# -*- coding: utf-8 -*-
"""
NO-ARGS DWTS IG validation (Spyder/runfile friendly)
---------------------------------------------------
Inputs:
  - dwts_star_pro_season_end_fans.csv  (model output)
  - dwts_s34_exact_followers.xlsx      (IG followers long table)
Outputs:
  - ig_s34_merged.csv
  - ig_s34_metrics.json
  - ig_s34_outliers.csv
  - ig_s34_star_fit.png
  - ig_s34_pro_fit.png
"""

import re, unicodedata, json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import pearsonr, spearmanr

# ============ EDIT IF NEEDED ============
MODEL_CSV = "dwts_star_pro_season_end_fans.csv"
IG_XLSX   = "dwts_s34_exact_followers.xlsx"
SEASON    = 34
OUT_PREFIX = "ig_s34_"
TOP_OUTLIERS = 10
EPS = 1e-12
# =======================================


# --------- small helpers ----------
def _strip_accents(s):
    s = unicodedata.normalize("NFKD", str(s))
    return "".join(ch for ch in s if not unicodedata.combining(ch))

def norm_name(s):
    s = _strip_accents(s).strip()
    s = re.sub(r"\s+", " ", s)
    return s.casefold()

def norm_partner(s):
    s = _strip_accents(s).strip()
    s = re.sub(r"\([^)]*\)", "", s).strip()     # remove (...) notes
    s = re.sub(r"\s+", " ", s)
    if "/" in s:
        s = s.split("/")[0].strip()             # keep first before slash
    return s.casefold()

def safe_pos(x):
    x = np.asarray(x, float)
    return np.clip(x, EPS, None)

def fit_loglog(y_stock, x_ig):
    """
    Fit log(y) = a + b log(x).
    Returns: a, b, R2 (in log space), logRMSE, med_mult_err, residuals
    residual = log(y) - (a + b log(x))
    """
    x = np.log(safe_pos(x_ig))
    y = np.log(safe_pos(y_stock))
    A = np.vstack([np.ones_like(x), x]).T
    (a, b), *_ = np.linalg.lstsq(A, y, rcond=None)
    yhat = a + b * x
    resid = y - yhat

    ss_res = float(np.sum(resid**2))
    ss_tot = float(np.sum((y - y.mean())**2))
    r2 = 1 - ss_res / ss_tot if ss_tot > 0 else np.nan
    logrmse = float(np.sqrt(np.mean(resid**2)))
    med_mult = float(np.exp(np.median(np.abs(resid))))  # typical multiplicative error

    return float(a), float(b), float(r2), logrmse, med_mult, resid

def mult_stats(mult):
    mult = np.asarray(mult, float)
    mult = mult[np.isfinite(mult) & (mult > 0)]
    if len(mult) == 0:
        return {}
    q = lambda p: float(np.quantile(mult, p))
    return dict(
        n=int(len(mult)),
        median=float(np.median(mult)),
        p25=q(0.25), p75=q(0.75),
        p10=q(0.10), p90=q(0.90),
        min=float(mult.min()),
        max=float(mult.max()),
        mean=float(mult.mean()),
    )

def plot_scatter_with_fit(x, y, title, outpng):
    x = np.asarray(x, float); y = np.asarray(y, float)
    m = np.isfinite(x) & np.isfinite(y) & (x > 0) & (y > 0)
    x = x[m]; y = y[m]
    if len(x) < 3:
        return

    a, b, r2, *_ = fit_loglog(y, x)
    xs = np.logspace(np.log10(x.min()), np.log10(x.max()), 200)
    ys = np.exp(a) * (xs ** b)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_xscale("log"); ax.set_yscale("log")
    ax.scatter(x, y, s=28)
    ax.plot(xs, ys, linewidth=2.0, label=f"fit: y=exp(a)x^b, b={b:.2f}, R²={r2:.2f}")
    ax.set_title(title)
    ax.set_xlabel("Instagram followers (real)")
    ax.set_ylabel("Model fan stock (predicted)")
    ax.legend(fontsize=9)
    fig.tight_layout()
    fig.savefig(outpng, dpi=180)
    plt.close(fig)


def main():
    # 1) read model output
    m = pd.read_csv(MODEL_CSV)
    m = m[pd.to_numeric(m["season"], errors="coerce") == SEASON].copy()

    # normalize ids
    m["celebrity_norm"] = m["celebrity_name"].map(norm_name)
    m["pro_norm"] = m["ballroom_partner"].map(norm_partner)

    for c in ["star_fans_after_end", "pro_fans_after_end", "end_week"]:
        m[c] = pd.to_numeric(m[c], errors="coerce")

    # 2) read IG xlsx
    ig = pd.read_excel(IG_XLSX)
    ig["role"] = ig["role"].astype(str).str.strip().str.casefold()
    ig["name_norm"] = ig["name"].map(norm_name)
    ig["instagram_followers_exact"] = pd.to_numeric(ig["instagram_followers_exact"], errors="coerce")

    celeb_map = dict(zip(
        ig.loc[ig["role"] == "celebrity", "name_norm"],
        ig.loc[ig["role"] == "celebrity", "instagram_followers_exact"]
    ))
    pro_map = dict(zip(
        ig.loc[ig["role"] == "pro", "name_norm"],
        ig.loc[ig["role"] == "pro", "instagram_followers_exact"]
    ))

    # 3) merge
    df = m.copy()
    df["celeb_ig_followers"] = df["celebrity_norm"].map(celeb_map)
    df["pro_ig_followers"] = df["pro_norm"].map(pro_map)

    # 4) core derived columns
    df["mult_star_ig_over_stock"] = df["celeb_ig_followers"] / df["star_fans_after_end"]
    df["mult_pro_ig_over_stock"]  = df["pro_ig_followers"]  / df["pro_fans_after_end"]

    # 5) fit + residuals (star)
    df["star_log_resid"] = np.nan
    ms = df.dropna(subset=["celeb_ig_followers", "star_fans_after_end"])
    star_fit = {}
    if len(ms) >= 3:
        a, b, r2, logrmse, medmult, resid = fit_loglog(ms["star_fans_after_end"].values, ms["celeb_ig_followers"].values)
        df.loc[ms.index, "star_log_resid"] = resid
        star_fit = dict(a=a, b=b, r2=r2, logrmse=logrmse, typical_mult_err=medmult)

    # 6) fit + residuals (pro)
    df["pro_log_resid"] = np.nan
    mp = df.dropna(subset=["pro_ig_followers", "pro_fans_after_end"])
    pro_fit = {}
    if len(mp) >= 3:
        a, b, r2, logrmse, medmult, resid = fit_loglog(mp["pro_fans_after_end"].values, mp["pro_ig_followers"].values)
        df.loc[mp.index, "pro_log_resid"] = resid
        pro_fit = dict(a=a, b=b, r2=r2, logrmse=logrmse, typical_mult_err=medmult)

    # 7) correlations on log scale (ranking + linear)
    metrics = dict(
        season=SEASON,
        n_pairs=int(len(df)),
        star_match_rate=float(df["celeb_ig_followers"].notna().mean()),
        pro_match_rate=float(df["pro_ig_followers"].notna().mean()),
        mult_star_stats=mult_stats(df["mult_star_ig_over_stock"].values),
        mult_pro_stats=mult_stats(df["mult_pro_ig_over_stock"].values),
        star_fit=star_fit,
        pro_fit=pro_fit,
    )

    if len(ms) >= 3:
        x = np.log(safe_pos(ms["celeb_ig_followers"].values))
        y = np.log(safe_pos(ms["star_fans_after_end"].values))
        metrics["star_log_pearson_r"] = float(pearsonr(x, y)[0])
        metrics["star_log_spearman_rho"] = float(spearmanr(x, y)[0])

    if len(mp) >= 3:
        x = np.log(safe_pos(mp["pro_ig_followers"].values))
        y = np.log(safe_pos(mp["pro_fans_after_end"].values))
        metrics["pro_log_pearson_r"] = float(pearsonr(x, y)[0])
        metrics["pro_log_spearman_rho"] = float(spearmanr(x, y)[0])

    # ratio check: (star/pro) model vs IG
    both = df.dropna(subset=["celeb_ig_followers","pro_ig_followers","star_fans_after_end","pro_fans_after_end"])
    if len(both) >= 3:
        r_m = safe_pos(both["star_fans_after_end"].values) / safe_pos(both["pro_fans_after_end"].values)
        r_i = safe_pos(both["celeb_ig_followers"].values) / safe_pos(both["pro_ig_followers"].values)
        metrics["ratio_log_pearson_r"] = float(pearsonr(np.log(r_i), np.log(r_m))[0])
        metrics["ratio_log_spearman_rho"] = float(spearmanr(np.log(r_i), np.log(r_m))[0])

    # 8) outliers: large deviation from fitted relation (sum of abs residuals)
    df["abs_resid_sum"] = np.abs(df["star_log_resid"]) + np.abs(df["pro_log_resid"])
    outliers = df.sort_values("abs_resid_sum", ascending=False).head(TOP_OUTLIERS)

    # 9) save tables
    df.to_csv(OUT_PREFIX + "merged.csv", index=False)
    outliers.to_csv(OUT_PREFIX + "outliers.csv", index=False)
    with open(OUT_PREFIX + "metrics.json", "w", encoding="utf-8") as f:
        json.dump(metrics, f, ensure_ascii=False, indent=2)

    # 10) plots (scatter + regression line ONLY)
    plot_scatter_with_fit(df["celeb_ig_followers"].values, df["star_fans_after_end"].values,
                          "Celebrity: IG followers vs Model stock (log-log)", OUT_PREFIX + "star_fit.png")
    plot_scatter_with_fit(df["pro_ig_followers"].values, df["pro_fans_after_end"].values,
                          "Pro: IG followers vs Model stock (log-log)", OUT_PREFIX + "pro_fit.png")

    print("Saved:")
    print(" -", OUT_PREFIX + "merged.csv")
    print(" -", OUT_PREFIX + "outliers.csv")
    print(" -", OUT_PREFIX + "metrics.json")
    print(" -", OUT_PREFIX + "star_fit.png")
    print(" -", OUT_PREFIX + "pro_fit.png")


if __name__ == "__main__":
    main()
