import time
import pandas as pd
import numpy as np
from tqdm import tqdm
import datetime
import json
import os

from experiments.synthetic_data.graph_generators.outdegree_graph_generator import outdegree_generate_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(
    nodes_li,
    densities,
    stds,
    trials,
    mean_dist,
    start,
    base_seed
):
    timer = time.perf_counter
    file_exists = False
    INF = float('inf')

    total_jobs = (
        len(nodes_li) *
        len(densities) *
        len(stds) *
        trials
    )

    with tqdm(total=total_jobs) as pbar:
        for nodes in nodes_li:
            rows = []

            for density in densities:
                for std in stds:
                    sigma = np.sqrt(np.log((std / mean_dist)**2 + 1))
                    mu = np.log(mean_dist) - (sigma**2) / 2
                    distribution = ('lognormal', (mu, sigma))

                    for trial in range(1, trials + 1):
                        seed = (base_seed * 1_000_003) ^ (nodes * 9176) ^ int(density * 1e9) ^ int(sigma * 1e6) ^ trial
                        seed &= 0xFFFFFFFF

                        adj = outdegree_generate_graph(
                            NODES=nodes,
                            DENSITY=density,
                            DISTRIBUTION=distribution,
                            SEED=seed
                        )

                        bin_heap = BinHeap(nodes)
                        st = timer()
                        stats = heap_dijkstra(nodes, adj, bin_heap, start)
                        et = timer()
                        rows.append(
                            {
                                "nodes": nodes,
                                "density": density,
                                "std": std,
                                "sigma": sigma,
                                "trial": trial,
                                "seed": seed,
                                "time": et - st,
                                "algorithm": "binary",
                                "extract_min_calls": stats.extract_min_calls,
                                "relax_attempts": stats.relax_attempts,
                                "relax_success": stats.relax_success,
                            }
                        )

                        fibo_heap = FiboHeap(nodes)
                        st = timer()
                        stats = heap_dijkstra(nodes, adj, fibo_heap, start)
                        et = timer()
                        rows.append(
                            {
                                "nodes": nodes,
                                "density": density,
                                "std": std,
                                "sigma": sigma,
                                "trial": trial,
                                "seed": seed,
                                "time": et - st,
                                "algorithm": "fibonacci",
                                "extract_min_calls": stats.extract_min_calls,
                                "relax_attempts": stats.relax_attempts,
                                "relax_success": stats.relax_success,
                            }
                        )

                        pbar.update(1)

            df_nodes = pd.DataFrame(rows)
            df_nodes.to_csv(
                csv_file,
                mode='a',
                header=not file_exists,
                index=False
            )
            file_exists = True
            print(f"=== SAVED nodes = {nodes} ({len(df_nodes)} rows) ===")
    return

# Settings
nodes_li = [2000, 4000, 8000, 16000]
densities = [0.0000001, 0.0000003, 0.000001, 0.000003, 0.00001, 0.00003, 0.0001, 0.0003, 0.001, 0.003, 0.01, 0.03, 0.1, 0.3]
stds = [1000, 1500, 2000, 3000, 4000, 6000, 8000, 12000, 16000]
trials = 10
mean_dist = 3000
start = 0
base_seed = 42

save_folder = "results/synthetic_data/raw"
os.makedirs(save_folder, exist_ok=True)
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
csv_file = f"{save_folder}/{timestamp}.csv"
config_file = f"{save_folder}/{timestamp}.json"

config = {
    "nodes_li": nodes_li,
    "densities": densities,
    "stds": stds,
    "trials": trials,
    "mean_dist": mean_dist,
    "start": start,
    "base_seed": base_seed,
}
with open(config_file, "w") as f:
    json.dump(config, f, indent=4)

df = run_test(
    nodes_li = nodes_li,
    densities = densities,
    stds = stds,
    trials = trials,
    mean_dist = mean_dist,
    start = start,
    base_seed = base_seed
)