import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import random


def read_dimacs(file):
    idx = 0
    with open(f"{folder}/{file}", "r") as f:
        for line in f:
            line = line.strip()
            if line.startswith("c"):
                continue
            elif line.startswith("p"):
                _, _, n, e = line.split()
                edges = int(e)
                data = [0] * edges
            elif line.startswith("a"):
                _, u, v, w = line.split()
                data[idx] = int(w)
                idx += 1
    return data


def visualize_weights(file):
    data = read_dimacs(file)
    sample = np.array(data)

    save_folder = f"{res_folder}/{file}"
    os.makedirs(save_folder, exist_ok=True)

    plt.figure(figsize=(8, 5))
    plt.hist(sample, bins=100)
    plt.title("Edge Weight Distribution (Original Scale)")
    plt.xlabel("Weight")
    plt.ylabel("Frequency")
    plt.tight_layout()
    plt.savefig(f"{save_folder}/hist_original.png")
    plt.close()

    log_sample = np.log(sample[sample > 0])

    plt.figure(figsize=(8, 5))
    plt.hist(log_sample, bins=100)
    plt.title("Log(Weight) Distribution")
    plt.xlabel("log(Weight)")
    plt.ylabel("Frequency")
    plt.tight_layout()
    plt.savefig(f"{save_folder}/hist_log.png")
    plt.close()

    plt.figure(figsize=(8, 5))
    plt.hist(log_sample, bins=100, density=True)
    plt.title("Log(Weight) Density")
    plt.xlabel("log(Weight)")
    plt.ylabel("Density")
    plt.tight_layout()
    plt.savefig(f"{save_folder}/hist_log_density.png")
    plt.close()

    return {
        "file": file,
        "edges": len(sample),
        "mean": np.mean(sample),
        "std": np.std(sample),
        "min": np.min(sample),
        "max": np.max(sample)
    }

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"
]

res_folder = "results/real_data/derived/dimacs_graph_distribution"
os.makedirs(res_folder, exist_ok=True)

rows = []
for file in files:
    row = visualize_weights(file)
    rows.append(row)
df = pd.DataFrame(rows)
df.to_csv(f"{res_folder}/distribution.csv", index=False)