import os
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
from scipy.optimize import curve_fit
import numpy as np


# 1. E vs relax_attempts
def E_vs_relax_attempts(df):
    X = df[["E"]].to_numpy()
    y = df["relax_attempts"].to_numpy()

    reg = LinearRegression()
    reg.fit(X, y)
    y_hat = reg.predict(X)
    r2 = r2_score(y, y_hat)
    equation = f"y = {reg.coef_[0]:.4g}x + {reg.intercept_:.4g}\n$R^2$ = {r2:.4f}"

    plt.figure(figsize=(6, 4))
    plt.scatter(X, y, marker="o")
    plt.plot(X, y_hat, color="red", label=equation)
    plt.xlabel("E (edges)")
    plt.ylabel("Average relax_attempts")
    plt.title("Edges vs Relax Attempts")
    plt.legend(loc="upper right", fontsize=8)
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(f"{save_folder}/E_vs_relax_attempts.png", dpi=300)
    plt.close()

    print(f"Saved E vs relax_attempts plots to {save_folder}")


# 2. sigma vs relax_success_ratio
def sigma_vs_relax_success_ratio(df):
    rows = []
    sigma_save_folder = os.path.join(save_folder, "sigma_vs_relax_ratio_controlled")
    os.makedirs(sigma_save_folder, exist_ok=True)
    cnt = 1

    for pair, group in df.groupby(["nodes", "density"]):
        nodes, density = pair
        grouped = group.groupby("sigma")["relax_success_ratio"].mean()

        sigma_vals = grouped.index.to_numpy()
        ratio_vals = grouped.to_numpy()

        valid = np.isfinite(ratio_vals)  # When attempts = 0 -> ratio becomes infinite.
        sigma_vals = sigma_vals[valid]
        ratio_vals = ratio_vals[valid]

        if len(sigma_vals) < 2:  # When too many invalid values -> Can't do regression.
            plt.close()
            cnt += 1
            continue

        log_sigma = np.log(sigma_vals).reshape(-1, 1)
        reg = LinearRegression()
        reg.fit(log_sigma, ratio_vals)
        r2 = r2_score(ratio_vals, reg.predict(log_sigma))

        sigma_line = np.linspace(sigma_vals.min(), sigma_vals.max(), 200)
        ratio_line = reg.predict(np.log(sigma_line).reshape(-1, 1))
        equation = f"ratio = {reg.coef_[0]:.4g}$\\cdot \\ln{{\\sigma}}$ + {reg.intercept_:.4g}\n$R^2$ = {r2:.4f}"

        fig, ax = plt.subplots(figsize=(6, 4))
        ax.scatter(sigma_vals, ratio_vals, marker="o")
        ax.plot(sigma_line, ratio_line, color="red", label=equation)
        ax.set_xlabel("Sigma")
        ax.set_ylabel("Relax Success Ratio")
        ax.set_title(
            f"Sigma vs Relax Success Ratio\n" f"nodes={nodes}, density={density}"
        )
        plt.legend(loc="upper left", fontsize=8)
        ax.grid(True)
        fig.tight_layout()
        fname = (
            f"{cnt}. nodes{nodes}_density{density:.20f}".rstrip("0").rstrip(".")
            + ".png"
        )
        fig.savefig(os.path.join(sigma_save_folder, fname), dpi=300)
        plt.close(fig)

        rows.append(
            {
                "nodes": nodes,
                "density": density,
                "avg_degree": (nodes - 1) * density,
                "coef": reg.coef_[0],
                "intercept": reg.intercept_,
                "r2": r2,
            }
        )
        cnt += 1

    res = pd.DataFrame(rows)
    res.to_csv(os.path.join(sigma_save_folder, "r2.csv"))
    print(f"Saved controlled sigma vs relax_success_ratio plots to {sigma_save_folder}")


# 3. avg_degree vs relax_success_ratio
def avg_degree_vs_relax_success_ratio(df):
    avgdeg_scatter_folder = os.path.join(save_folder, "avg_deg_vs_ratio_controlled")
    os.makedirs(avgdeg_scatter_folder, exist_ok=True)
    cnt = 1

    for pair, group in df.groupby(["nodes", "sigma"]):
        nodes, sigma = pair
        grouped = group.groupby("avg_degree")["relax_success_ratio"].mean().reset_index()

        fig, ax = plt.subplots(figsize=(6, 4))
        ax.plot(grouped["avg_degree"], grouped["relax_success_ratio"], marker="o")
        ax.set_xlabel("Avg Degree = (N-1) × density")
        ax.set_ylabel("Relax Success Ratio")
        ax.set_title(
            f"Avg Degree vs Relax Success Ratio\n"
            f"nodes={nodes}, sigma={sigma:.4f}"
        )
        ax.grid(True)
        fig.tight_layout()
        fname = f"{cnt}. nodes{nodes}_sigma{sigma:.6f}.png"
        fig.savefig(os.path.join(avgdeg_scatter_folder, fname), dpi=300)
        plt.close(fig)
        cnt += 1

    print(f"Saved avg_degree vs ratio plots to {avgdeg_scatter_folder}")


# 4. log_avg_degree vs log_ratio
def log_avg_degree_vs_log_ratio(df):
    CEILING = 1
    avgdeg_loglog_folder = os.path.join(save_folder, "avgdeg_vs_ratio_loglog")
    os.makedirs(avgdeg_loglog_folder, exist_ok=True)
    cnt = 0.99

    for pair, group in df.groupby(["nodes", "sigma"]):
        nodes, sigma = pair
        grouped = group.groupby("avg_degree")["relax_success_ratio"].mean().reset_index()
        grouped = grouped[grouped["relax_success_ratio"] > 0]

        deg_all = grouped["avg_degree"].to_numpy(dtype=float)
        ratio_all = grouped["relax_success_ratio"].to_numpy(dtype=float)

        ceiling_mask = ratio_all >= CEILING
        decline_mask = ~ceiling_mask

        fig, ax = plt.subplots(figsize=(6, 4))

        if ceiling_mask.any():
            ax.scatter(
                deg_all[ceiling_mask],
                ratio_all[ceiling_mask],
                color="gray",
                label="ratio ≥ 0.99 (ceiling)",
            )

        if decline_mask.sum() >= 2:
            x_dec = deg_all[decline_mask]
            y_dec = ratio_all[decline_mask]
            log_x = np.log2(x_dec)
            log_y = np.log2(y_dec)

            reg = LinearRegression()
            reg.fit(log_x.reshape(-1, 1), log_y)
            r2 = r2_score(log_y, reg.predict(log_x.reshape(-1, 1)))
            slope = reg.coef_[0]
            intercept = reg.intercept_

            log_x_line = np.linspace(log_x.min(), log_x.max(), 200)
            x_line = 2 ** log_x_line
            y_line = 2 ** (slope * log_x_line + intercept)

            ax.scatter(
                x_dec, y_dec, color="blue", label="declining points"
            )
            ax.plot(
                x_line,
                y_line,
                color="blue",
                linestyle="--",
                label=f"fit (slope={slope:.3f}, R²={r2:.3f})",
            )
        elif decline_mask.any():
            ax.scatter(
                deg_all[decline_mask],
                ratio_all[decline_mask],
                color="blue",
                label="declining points",
            )

        ax.set_xscale("log")
        ax.set_yscale("log")
        ax.set_xlabel("Avg Degree (log scale)")
        ax.set_ylabel("Relax Success Ratio (log scale)")
        ax.set_title(
            f"Avg Degree vs Relax Success Ratio (log-log)\n"
            f"nodes={nodes}, sigma={sigma:.4f}"
        )
        ax.legend(fontsize=8)
        ax.grid(True)
        fig.tight_layout()
        fname = f"{cnt}. nodes{nodes}_sigma{sigma:.6f}.png"
        fig.savefig(os.path.join(avgdeg_loglog_folder, fname), dpi=300)
        plt.close(fig)
        cnt += 1

    print(f"Saved avg_degree vs ratio log-log plots to {avgdeg_loglog_folder}")


# 5. sigma vs. relax_success_ratio in controlled node,density
def regime_distribution(df):
    r2_csv = os.path.join(save_folder, "sigma_vs_relax_ratio_controlled", "r2.csv")
    res = pd.read_csv(r2_csv)

    avg_deg   = res["avg_degree"].to_numpy()   # (nodes-1)*density
    r2        = res["r2"].to_numpy()
    coef      = res["coef"].to_numpy()
    intercept = res["intercept"].to_numpy()

    # ── Phase classification ──────────────────────────────────────────────────
    tree_mask  = (intercept == 1) & (np.abs(coef) == 0)
    giant_mask = (~tree_mask) & (r2 >= 0.9)
    trans_mask = (~tree_mask) & (~giant_mask)

    labels = np.empty(len(res), dtype=object)
    labels[tree_mask]  = "Tree"
    labels[giant_mask] = "Giant"
    labels[trans_mask] = "Transition"

    print("\n=== Phase counts ===")
    for phase in ["Tree", "Transition", "Giant"]:
        print(f"  {phase}: {(labels == phase).sum()}")

    regime_save_folder = os.path.join(save_folder, "regime_distribution")
    os.makedirs(regime_save_folder, exist_ok=True)

    colors = {"Tree": "gray", "Transition": "darkorange", "Giant": "steelblue"}
    bins = np.logspace(
        np.log10(avg_deg[avg_deg > 0].min()),
        np.log10(avg_deg.max()),
        40
    )

    # ── Histogram: avg_degree distribution per phase ──────────────────────────
    fig, ax = plt.subplots(figsize=(8, 5))
    for phase, color in colors.items():
        vals = avg_deg[labels == phase]
        ax.hist(vals, bins=bins, alpha=0.6, color=color, label=phase)
    ax.axvline(1.0, color="red", linestyle="--", linewidth=1.5, label="avg_degree = 1")
    ax.set_xscale("log")
    ax.set_xlabel("Avg Degree = (N-1) × density (log scale)")
    ax.set_ylabel("Count")
    ax.set_title("Phase Distribution by Avg Degree")
    ax.legend(fontsize=9)
    ax.grid(True, which="both", alpha=0.4)
    fig.tight_layout()
    fig.savefig(os.path.join(regime_save_folder, "phase_histogram.png"), dpi=300)
    plt.close(fig)

    # ── Scatter: avg_degree vs R², colored by phase ───────────────────────────
    fig, ax = plt.subplots(figsize=(8, 5))
    for phase, color in colors.items():
        mask = labels == phase
        ax.scatter(avg_deg[mask], r2[mask], s=15, alpha=0.6, color=color, label=phase)
    ax.axvline(1.0, color="red", linestyle="--", linewidth=1.5, label="avg_degree = 1")
    ax.set_xscale("log")
    ax.set_xlabel("Avg Degree (log scale)")
    ax.set_ylabel("R²")
    ax.set_title("R² vs Avg Degree, colored by phase")
    ax.legend(fontsize=9)
    ax.grid(True, which="both", alpha=0.4)
    fig.tight_layout()
    fig.savefig(os.path.join(regime_save_folder, "r2_by_phase.png"), dpi=300)
    plt.close(fig)

    print(f"Saved regime distribution plots to {regime_save_folder}")


# 6. Nonlinear regression: r = (a·ln(sigma) + b) · avg_deg^c
def nonlinear_regression(compute_local=False):
    # Filter: giant component regime only
    sub = df[(df["avg_degree"] >= 1) & (df["relax_success_ratio"] < 0.99)].copy()
    sub = sub[sub["relax_success_ratio"] > 0].dropna(subset=["relax_success_ratio", "avg_degree", "sigma"])

    avg_deg = sub["avg_degree"].to_numpy()
    sigma   = sub["sigma"].to_numpy()
    r       = sub["relax_success_ratio"].to_numpy()

    def model(X, a, b, c):
        avg_deg_, sigma_ = X
        return (a * np.log(sigma_) + b) * avg_deg_ ** c

    # Initial guess
    p0 = [0.05, 0.5, -0.7]

    popt, pcov = curve_fit(model, (avg_deg, sigma), r, p0=p0, maxfev=10000)
    a, b, c = popt
    perr = np.sqrt(np.diag(pcov))

    r_pred = model((avg_deg, sigma), a, b, c)
    ss_res = np.sum((r - r_pred) ** 2)
    ss_tot = np.sum((r - r.mean()) ** 2)
    r2 = 1 - ss_res / ss_tot

    if compute_local == False:
        return (a, b, c)

    print("\n=== Nonlinear regression: r = (a·ln(σ) + b) · avg_deg^c ===")
    print(f"  a = {a:.6f} ± {perr[0]:.6f}")
    print(f"  b = {b:.6f} ± {perr[1]:.6f}")
    print(f"  c = {c:.6f} ± {perr[2]:.6f}")
    print(f"  R² = {r2:.4f}")

    # ── Predicted vs Actual ───────────────────────────────────────────────────
    nlr_save_folder = os.path.join(save_folder, "nonlinear_regression")
    os.makedirs(nlr_save_folder, exist_ok=True)

    fig, ax = plt.subplots(figsize=(6, 5))
    ax.scatter(r, r_pred, s=5, alpha=0.3, color="steelblue")
    mn, mx = min(r.min(), r_pred.min()), max(r.max(), r_pred.max())
    ax.plot([mn, mx], [mn, mx], color="red", linewidth=1.5, linestyle="--")
    ax.set_xlabel("Actual ratio")
    ax.set_ylabel("Predicted ratio")
    ax.set_title(f"Nonlinear fit: r = (a·ln(σ)+b)·avg_deg^c\nR²={r2:.4f}")
    ax.grid(True, alpha=0.4)
    fig.tight_layout()
    fig.savefig(os.path.join(nlr_save_folder, "predicted_vs_actual.png"), dpi=300)
    plt.close(fig)

    # ── Residuals by nodes ────────────────────────────────────────────────────
    residuals = r - r_pred
    fig, ax = plt.subplots(figsize=(7, 5))
    for n in sorted(sub["nodes"].unique()):
        mask = sub["nodes"].to_numpy() == n
        ax.scatter(r_pred[mask], residuals[mask], s=5, alpha=0.4, label=f"N={n}")
    ax.axhline(0, color="red", linewidth=1.2, linestyle="--")
    ax.set_xlabel("Predicted ratio")
    ax.set_ylabel("Residual")
    ax.set_title("Residuals by N")
    ax.legend(fontsize=6, ncol=3, markerscale=2)
    ax.grid(True, alpha=0.4)
    fig.tight_layout()
    fig.savefig(os.path.join(nlr_save_folder, "residuals_by_N.png"), dpi=300)
    plt.close(fig)

    print(f"Saved nonlinear regression plots to {nlr_save_folder}")



# Initial Setup
csv_file = "results/synthetic_data/raw/20260421_031740.csv"
save_folder = f"results/synthetic_data/derived/{os.path.splitext(os.path.basename(csv_file))[0]}/call_number_analysis"
os.makedirs(save_folder, exist_ok=True)

df = pd.read_csv(csv_file)
df = df[df["algorithm"] == "binary"].copy()
df["relax_success_ratio"] = df["relax_success"] / df["relax_attempts"]
df["E"] = df["nodes"] * (df["nodes"] - 1) * df["density"]
df["decrease_key"] = df["relax_success"]
df["avg_degree"] = (df["nodes"] - 1) * df["density"]

E_vs_relax_attempts(df) # 1
sigma_vs_relax_success_ratio(df) # 2
avg_degree_vs_relax_success_ratio(df) # 3
log_avg_degree_vs_log_ratio(df) # 4
regime_distribution(df) # 5
nonlinear_regression(compute_local=True) # 6