import pandas as pd
import os
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error, mean_absolute_error

raw_path  = "results/real_data/raw/dimacs_s42.csv"
pred_path = "results/real_data/derived/distribution/dimacs_prediction_result/prediction.csv"
save_folder = "results/real_data/derived/comparison"
os.makedirs(save_folder, exist_ok=True)

raw  = pd.read_csv(raw_path)
pred = pd.read_csv(pred_path)

# real: trial 평균
agg = (
    raw.groupby("file")
    .agg(
        nodes        = ("nodes",        "first"),
        density      = ("density",      "first"),
        real_time    = ("time",         "mean"),
        real_relax   = ("relax_success","mean"),
    )
    .reset_index()
)

# prediction: file 기준으로 merge
pred_cols = pred[["file", "predict_time", "relax_success"]].rename(
    columns={"relax_success": "pred_relax"}
)

merged = agg.merge(pred_cols, on="file", how="inner")

# 오차 계산
for real_col, pred_col, prefix in [
    ("real_time",  "predict_time", "time"),
    ("real_relax", "pred_relax",   "relax"),
]:
    merged[f"{prefix}_abs_err"] = (merged[pred_col] - merged[real_col]).abs()
    merged[f"{prefix}_rel_err"] = merged[f"{prefix}_abs_err"] / merged[real_col]

cols = [
    "file", "nodes", "density",
    "real_time",  "predict_time",  "time_abs_err",  "time_rel_err",
    "real_relax", "pred_relax",    "relax_abs_err", "relax_rel_err",
]
merged[cols].to_csv(os.path.join(save_folder, "comparison.csv"), index=False)
print(merged[cols].to_string(index=False))
print(f"\nSaved → {save_folder}/comparison.csv")

for real_col, pred_col, label in [
    ("real_time",  "predict_time", "time"),
    ("real_relax", "pred_relax",   "relax_success"),
]:
    y     = merged[real_col].to_numpy()
    y_hat = merged[pred_col].to_numpy()
    ss_res = ((y - y_hat) ** 2).sum()
    ss_tot = ((y - y.mean()) ** 2).sum()
    r2   = 1 - ss_res / ss_tot
    rmse = mean_squared_error(y, y_hat) ** 0.5
    mae  = mean_absolute_error(y, y_hat)
    print(f"\n[{label}]  R²={r2:.4f}  RMSE={rmse:.4g}  MAE={mae:.4g}")


# ── Plot ───────────────────────────────────────────────────────────────────────
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

for ax, real_col, pred_col, label in [
    (axes[0], "real_time",  "predict_time", "Runtime [s]"),
    (axes[1], "real_relax", "pred_relax",   "Relax Success"),
]:
    x = merged[real_col].to_numpy()
    y = merged[pred_col].to_numpy()
    mn, mx = min(x.min(), y.min()), max(x.max(), y.max())

    ax.scatter(x, y, s=60, zorder=3)
    for _, row in merged.iterrows():
        ax.annotate(
            row["file"].replace("USA-road-d.", "").replace(".gr", ""),
            (row[real_col], row[pred_col]),
            fontsize=6, textcoords="offset points", xytext=(4, 2)
        )
    ax.plot([mn, mx], [mn, mx], color="red", linestyle="--", linewidth=1.2, label="y = x")
    ax.set_xlabel(f"Real {label}")
    ax.set_ylabel(f"Predicted {label}")
    ax.set_title(f"Predicted vs Real — {label}")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.4)

fig.tight_layout()
fig.savefig(os.path.join(save_folder, "comparison_plot.png"), dpi=300)
plt.close(fig)
print(f"Saved → {save_folder}/comparison_plot.png")