import os
import time
import pandas as pd
import numpy as np

from experiments.real_data.graph_converters.dimacs_graph_converter import dimacs_convert_graph
from core.dijkstra.heap_dijkstra import heap_dijkstra
from core.heaps.binary_heap import BinHeap
from core.heaps.fibonacci_heap import FiboHeap

def run_test(
    folder,
    files,
    trials,
    base_seed
):
    rows = []
    timer = time.perf_counter
    rng = np.random.default_rng(base_seed)
    INF = float('inf')

    for file in files:
        nodes, adj = dimacs_convert_graph(f'{folder}/{file}')
        density = len(adj) / (nodes * (nodes - 1))

        for trial in range(1, trials + 1):
            start, end = rng.choice(nodes, size=2, replace=False) + 1

            bin_heap = BinHeap(nodes)
            st = timer()
            dist, stats = heap_dijkstra(nodes, adj, bin_heap, start, end)
            et = timer()
            rows.append(
                {
                    "file": file,
                    "nodes": nodes,
                    "density": density,
                    "trial": trial,
                    "start": start,
                    "end": end,
                    "time": et - st,
                    "algorithm": "binary",
                    "reached": dist != INF,
                    "extract_min_calls": stats.extract_min_calls,
                    "relax_attempts": stats.relax_attempts,
                    "relax_success": stats.relax_success,
                }
            )

            fibo_heap = FiboHeap(nodes)
            st = timer()
            dist, stats = heap_dijkstra(nodes, adj, fibo_heap, start, end)
            et = timer()
            rows.append(
                {
                    "file": file,
                    "nodes": nodes,
                    "density": density,
                    "trial": trial,
                    "start": start,
                    "end": end,
                    "time": et - st,
                    "algorithm": "fibonacci",
                    "reached": dist != INF,
                    "extract_min_calls": stats.extract_min_calls,
                    "relax_attempts": stats.relax_attempts,
                    "relax_success": stats.relax_success,
                }
            )

    df = pd.DataFrame(rows)
    return df

# Settings
folder = "experiments/real_data/data/dimacs_data"
files = [
"USA-road-d.BAY.gr",
"USA-road-d.CAL.gr",
"USA-road-d.COL.gr",
"USA-road-d.CTR.gr",
"USA-road-d.E.gr",
"USA-road-d.FLA.gr",
"USA-road-d.LKS.gr",
"USA-road-d.NE.gr",
"USA-road-d.NW.gr",
"USA-road-d.NY.gr",
"USA-road-d.USA.gr",
"USA-road-d.W.gr"
]
trials = 200
base_seed = 42

df = run_test(
    folder=folder,
    files=files,
    trials=trials,
    base_seed=base_seed
)

save_folder = "results/real_data/raw"
os.makedirs(save_folder, exist_ok=True)
df.to_csv(f"{save_folder}/dimacs_t{trials}_s{base_seed}.csv", index=False)