import os
import gc
import time
import pandas as pd
import random

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

def done_trials(out_path, file):
    if not os.path.exists(out_path):
        return set()
    df = pd.read_csv(out_path, usecols=["file", "trial"])
    return set(df[df["file"] == file]["trial"].tolist())

def run_test(
    folder,
    files,
    base_seed,
    out_path,
    trials=10,
):
    timer = time.perf_counter
    write_header = not os.path.exists(out_path)

    for file in files:
        completed = done_trials(out_path, file)
        remaining = [t for t in range(1, trials + 1) if t not in completed]

        if not remaining:
            print(f"[SKIP] {file} (all {trials} trials done)")
            continue

        print(f"[LOAD] {file}  (completed: {sorted(completed)}, remaining: {remaining})")
        nodes, adj = dimacs_convert_graph(f'{folder}/{file}')
        density = len(adj) / (nodes * (nodes - 1))

        for trial in remaining:
            print(f"  trial {trial}/{trials}")

            random.seed(base_seed + trial)
            start = random.randint(1, nodes)

            bin_heap = BinHeap(nodes)
            st = timer()
            stats = heap_dijkstra(nodes, adj, bin_heap, start)
            et = timer()

            row = pd.DataFrame([{
                "file": file,
                "nodes": nodes,
                "density": density,
                "start": start,
                "time": et - st,
                "trial": trial,
                "extract_min_calls": stats.extract_min_calls,
                "relax_attempts": stats.relax_attempts,
                "relax_success": stats.relax_success,
            }])
            row.to_csv(out_path, mode="a", header=write_header, index=False)
            write_header = False

            del bin_heap
            gc.collect()

        del adj
        gc.collect()
        print(f"[DONE] {file}")


# 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"
]
base_seed = 42

save_folder = "results/real_data/raw"
os.makedirs(save_folder, exist_ok=True)
out_path = f"{save_folder}/dimacs_s{base_seed}.csv"

run_test(
    folder=folder,
    files=files,
    base_seed=base_seed,
    out_path=out_path,
)