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

# Runtime = Add + Extract-min + Relax-attempts + Relax-success
#         = VlogV * k1 + VlogV * k2 + E * k3 + \alpha * logV * k4 (Binary)
#         = V * k1 + VlogV * k2 + E * k3 + \alpha * k4 (Fibonacci)
# (\alpha = relax_success)
#
# V increases -> Add, Extract-min, Relax-attempts, Relax-success all increases.
# Therfore, critical multicollinearlity occurs.
#
# To solve, run regression per each V.
# Then V becomes constant.
#
# Runtime(Binary) = VlogV * k1 + VlogV * k2 + E * k3 + \alpha * logV * k4 (Binary)
#                 = intercept + E * k3 + \alpha * k4
# Runtime(Fibonacci) = V * k1 + VlogV * k2 + E * k3 + \alpha * k4 (Fibonacci)
#                    = intercept + E * k3 + \alpha * k4
#
# Final equation
# Runtime of specific V = intercept + E * k3 + \alpha * k4
# 
# intercept absorbs V*(k1+k2) which is constant within each N.
# k3: cost per relax_attempt  (should be equal across heaps)
# k4: cost per decrease-key   (binary: should scale with log2N; fibonacci: should be constant)

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]}/processing_time_analysis"
os.makedirs(save_folder, exist_ok=True)

df = pd.read_csv(csv_file)
nodes = sorted(df["nodes"].unique())
algos = ["binary", "fibonacci"]

records = []

for algo in algos:
    sub = df[df["algorithm"] == algo]
    for n in nodes:
        g = sub[sub["nodes"] == n]

        X = g[["relax_attempts", "relax_success"]].to_numpy()
        y = g["time"].to_numpy()

        model = LinearRegression().fit(X, y)
        r2 = r2_score(y, model.predict(X))

        records.append({
            "algorithm":      algo,
            "N":         n,
            "log2N":     round(np.log2(n), 4),
            "intercept": model.intercept_,
            "k3 (E)":    model.coef_[0],
            "k4 (alpha)":model.coef_[1],
            "r2":        r2,
            "n":         len(g),
        })

results = pd.DataFrame(records)

for algo in algos:
    print(f"=== {algo} ===")
    r = results[results["algorithm"] == algo][["N","log2N","intercept","k3 (E)","k4 (alpha)","r2","n"]]
    print(r.to_string(index=False))
    print()

grouped = results.set_index(["N", "algorithm"])

def get_param(param):
    return (
        grouped.loc[(nodes, "binary"), param].values,
        grouped.loc[(nodes, "fibonacci"), param].values
    )

k3_bin, k3_fib = get_param("k3 (E)")
k4_bin, k4_fib = get_param("k4 (alpha)")
int_bin, int_fib = get_param("intercept")

# binary: k4 ~ log₂N 회귀
log_n = np.log2(np.array(nodes)).reshape(-1, 1)
bin_reg = LinearRegression().fit(log_n, k4_bin)
r2_bin = r2_score(k4_bin, bin_reg.predict(log_n))

# fibonacci: k4 ~ log₂N 회귀
fib_reg = LinearRegression().fit(log_n, k4_fib)
r2_fib = r2_score(k4_fib, fib_reg.predict(log_n))

n_line = np.linspace(min(nodes), max(nodes), 300)
log_n_line = np.log2(n_line).reshape(-1, 1)

print(f"\nbinary log₂N regression: coef={bin_reg.coef_[0]:.4e}, intercept={bin_reg.intercept_:.4e}, R²={r2_bin:.4f}")

# Plot
fig, ax = plt.subplots(figsize=(7, 5))
ax.plot(nodes, k4_bin, marker="o", color="steelblue", label="binary k4")
ax.plot(nodes, k4_fib, marker="s", color="darkorange", label="fibonacci k4")
ax.plot(n_line, bin_reg.predict(log_n_line), color="steelblue", linestyle="--", linewidth=1.5,
        label=f"binary fit (∝ log₂N)  R²={r2_bin:.3f}")
ax.plot(n_line, fib_reg.predict(log_n_line), color="steelblue", linestyle="--", linewidth=1.5,
        label=f"fibonacci fit (∝ log₂N)  R²={r2_fib:.3f}")
ax.set_xlabel("N (nodes)")
ax.set_ylabel("k4 — decrease-key cost per call [s]")
ax.set_title("k4 (decrease-key unit cost) vs N")
ax.legend(fontsize=8)
ax.grid(True, alpha=0.4)
fig.tight_layout()
fig.savefig(f"{save_folder}/k4_vs_N.png", dpi=300)
plt.close(fig)
print(f"Saved: {save_folder}/k4_vs_N.png")


# k3
fig, ax = plt.subplots(figsize=(7, 5))
ax.plot(nodes, k3_bin, marker="o", color="steelblue", label="binary k3")
ax.plot(nodes, k3_fib, marker="s", color="darkorange", label="fibonacci k3")
ax.set_xlabel("N (nodes)")
ax.set_ylabel("k3 — iteration cost per call [s]")
ax.set_title("k3 (iteration unit cost) vs N")
ax.legend(fontsize=8)
ax.grid(True, alpha=0.4)
fig.tight_layout()
fig.savefig(f"{save_folder}/k3_vs_N.png", dpi=300)
plt.close(fig)
print(f"Saved: {save_folder}/k3_vs_N.png")


# Intercept
fig, ax = plt.subplots(figsize=(7, 5))
ax.plot(nodes, int_bin, marker="o", color="steelblue", label="binary k4")
ax.plot(nodes, int_fib, marker="s", color="darkorange", label="fibonacci k4")
ax.set_xlabel("N (nodes)")
# ax.set_ylabel("k4 — decrease-key cost per call [s]")
# ax.set_title("k4 (decrease-key unit cost) vs N")
ax.legend(fontsize=8)
ax.grid(True, alpha=0.4)
fig.tight_layout()
fig.savefig(f"{save_folder}/int_vs_N.png", dpi=300)
plt.close(fig)
print(f"Saved: {save_folder}/int_vs_N.png")


# ── VIF Check ──────────────────────────────────────────────────────────────────
# relax_success = ratio × relax_attempts → 두 변수가 선형 상관될 수 있음
# VIF = 1 / (1 - R²), where R² is from regressing X1 on X2 (2-predictor case)
# VIF > 10: severe multicollinearity, coefficients become unstable

vif_records = []

for algo in algos:
    sub = df[df["algorithm"] == algo]
    for n in nodes:
        g = sub[sub["nodes"] == n]
        X = g[["relax_attempts", "relax_success"]].to_numpy()

        # Regress relax_attempts on relax_success
        reg_vif = LinearRegression().fit(X[:, 1].reshape(-1, 1), X[:, 0])
        r2_vif = r2_score(X[:, 0], reg_vif.predict(X[:, 1].reshape(-1, 1)))
        vif = 1 / (1 - r2_vif) if r2_vif < 1.0 else float("inf")

        # Pearson correlation (simpler diagnostic)
        corr = np.corrcoef(X[:, 0], X[:, 1])[0, 1]

        vif_records.append({
            "algorithm": algo,
            "N": n,
            "corr(E, alpha)": round(corr, 6),
            "R2_vif": round(r2_vif, 6),
            "VIF": round(vif, 2),
        })

vif_df = pd.DataFrame(vif_records)

for algo in algos:
    print(f"\n=== VIF — {algo} ===")
    v = vif_df[vif_df["algorithm"] == algo][["N", "corr(E, alpha)", "R2_vif", "VIF"]]
    print(v.to_string(index=False))

# Plot VIF vs N
fig, ax = plt.subplots(figsize=(7, 4))
for algo, color, marker in [("binary", "steelblue", "o"), ("fibonacci", "darkorange", "s")]:
    v = vif_df[vif_df["algorithm"] == algo]
    ax.plot(v["N"], v["VIF"], marker=marker, color=color, label=algo)
ax.axhline(10, color="red", linestyle="--", linewidth=1.2, label="VIF = 10 (threshold)")
ax.set_xlabel("N (nodes)")
ax.set_ylabel("VIF")
ax.set_title("VIF (relax_attempts vs relax_success) per N")
ax.legend(fontsize=8)
ax.grid(True, alpha=0.4)
fig.tight_layout()
fig.savefig(f"{save_folder}/vif_vs_N.png", dpi=300)
plt.close(fig)
print(f"\nSaved: {save_folder}/vif_vs_N.png")