# -*- coding: utf-8 -*-
"""
DWTS Output Analyzer — Reasonableness + Consistency Metrics
==========================================================

Reads:
- dwts_fan_vote_estimates_components.csv (estimator output)
- 2026_MCM_Problem_C_Data.csv            (original DWTS data)

Computes sanity checks per (season, week):
- share_sum_error: |sum(fan_share_total) - 1|
- component_sum_error: |(star_fixed + pro_fixed + perf) - total|
- vote_pool_rel_error: |sum(fan_votes_total_est) - week_total_votes| / week_total_votes

Computes elimination-week consistency:
Let C_i = wJ * p_i + wF * x_i, where p_i is judge percent, x_i is estimated fan share.
For each elimination week with m eliminated:
- bottom_m_hit: eliminated set subset of bottom-m C among active
- pairwise_satisfaction_rate: fraction of (elim, survivor) pairs with C_survivor > C_elim
- margin: min(C_survivor) - max(C_elim) (positive is good)
- slack_week: taken from estimator output

Outputs:
- dwts_check_sanity_by_week.csv
- dwts_check_week_level_consistency.csv
- dwts_check_season_summary.csv
"""

import re
import numpy as np
import pandas as pd

WEEK_JUDGE_RE = re.compile(r"week(\d+)_judge(\d+)_score", re.IGNORECASE)

def find_week_judge_columns(df: pd.DataFrame):
    week_cols = {}
    for c in df.columns:
        m = WEEK_JUDGE_RE.fullmatch(c)
        if m:
            w = int(m.group(1))
            week_cols.setdefault(w, []).append(c)
    weeks = sorted(week_cols)
    for w in weeks:
        week_cols[w] = sorted(week_cols[w])
    return weeks, week_cols

def parse_elim_week(results):
    if not isinstance(results, str):
        return None
    m = re.search(r"Eliminated\s*Week\s*(\d+)", results, re.I)
    return int(m.group(1)) if m else None

def season_length(season_df, weeks, week_cols):
    last = 0
    for w in weeks:
        arr = season_df[week_cols[w]].to_numpy(float)
        if np.isfinite(arr).any():
            last = w
    return last

def build_J_totals(season_df, W, week_cols):
    N = len(season_df)
    J = np.zeros((N, W), float)
    for wi in range(1, W + 1):
        arr = season_df[week_cols[wi]].to_numpy(float)
        J[:, wi - 1] = np.nansum(arr, axis=1)
    return J

def compute_p_and_active(season_df):
    weeks, week_cols = find_week_judge_columns(season_df)
    W = season_length(season_df, weeks, week_cols)
    J = build_J_totals(season_df, W, week_cols)
    active = (J > 0)
    p = np.zeros((len(season_df), W), float)
    for w in range(1, W + 1):
        idx = np.where(active[:, w-1])[0]
        if len(idx) == 0:
            continue
        jw = J[idx, w-1]
        denom = jw.sum()
        p[idx, w-1] = jw/denom if denom > 0 else 1.0/len(idx)
    return W, J, active, p

def run_analyzer(
    est_csv="dwts_fan_vote_estimates_components.csv",
    dwts_csv="2026_MCM_Problem_C_Data.csv",
    out_prefix="dwts_check_",
    wJ=1.0,
    wF=1.0
):
    est = pd.read_csv(est_csv)
    dwts = pd.read_csv(dwts_csv)

    # normalize key strings (防止空格导致 join 失败)
    for c in ["celebrity_name", "ballroom_partner"]:
        if c in est.columns:
            est[c] = est[c].astype(str).str.strip()
        if c in dwts.columns:
            dwts[c] = dwts[c].astype(str).str.strip()

    # --------- Sanity checks by (season, week)
    grp = est.groupby(["season", "week"], as_index=False)

    sum_total = grp["fan_share_total"].sum().rename(columns={"fan_share_total": "sum_fan_share_total"})
    sum_star  = grp["fan_share_star_fixed"].sum().rename(columns={"fan_share_star_fixed": "sum_star_fixed"})
    sum_pro   = grp["fan_share_pro_fixed"].sum().rename(columns={"fan_share_pro_fixed": "sum_pro_fixed"})
    sum_perf  = grp["fan_share_perf"].sum().rename(columns={"fan_share_perf": "sum_perf"})
    sum_votes = grp["fan_votes_total_est"].sum().rename(columns={"fan_votes_total_est": "sum_fan_votes_est"})
    week_votes = grp["week_total_votes"].first().rename(columns={"week_total_votes": "week_total_votes"})
    slack_week = grp["slack_week"].first().rename(columns={"slack_week": "slack_week"})
    fill_flag = grp["viewers_filled"].first().rename(columns={"viewers_filled": "viewers_filled"})
    fill_method = grp["viewers_fill_method"].first().rename(columns={"viewers_fill_method": "viewers_fill_method"})

    sanity = sum_total.merge(sum_star, on=["season","week"])\
                      .merge(sum_pro, on=["season","week"])\
                      .merge(sum_perf, on=["season","week"])\
                      .merge(sum_votes, on=["season","week"])\
                      .merge(week_votes, on=["season","week"])\
                      .merge(slack_week, on=["season","week"])\
                      .merge(fill_flag, on=["season","week"])\
                      .merge(fill_method, on=["season","week"])

    sanity["share_sum_error"] = (sanity["sum_fan_share_total"] - 1.0).abs()
    sanity["component_sum_error"] = (sanity["sum_star_fixed"] + sanity["sum_pro_fixed"] + sanity["sum_perf"] - sanity["sum_fan_share_total"]).abs()
    sanity["vote_pool_rel_error"] = (sanity["sum_fan_votes_est"] - sanity["week_total_votes"]).abs() / sanity["week_total_votes"].replace(0, np.nan)

    # --------- Consistency checks on elimination weeks
    week_rows = []
    season_rows = []

    seasons = sorted(est["season"].dropna().unique().astype(int).tolist())

    # build quick lookup for x by (season, week, celeb, pro)
    est_key = {}
    for r in est.itertuples(index=False):
        key = (int(r.season), int(r.week), str(r.celebrity_name).strip(), str(r.ballroom_partner).strip())
        est_key[key] = float(r.fan_share_total)

    for s in seasons:
        dwts_s = dwts[dwts["season"] == s].reset_index(drop=True)
        if dwts_s.empty:
            continue

        W, J, active, p = compute_p_and_active(dwts_s)

        # elimination week inference
        last_pos = np.array([(np.where(J[i, :] > 0)[0].max() + 1) if (J[i, :] > 0).any() else 0
                             for i in range(len(dwts_s))], int)
        elim_week = []
        for i, resu in enumerate(dwts_s["results"].tolist()):
            ew = parse_elim_week(resu)
            lp = last_pos[i]
            if ew is not None:
                elim_week.append(min(ew, lp) if lp > 0 else ew)
            else:
                elim_week.append(lp if (lp > 0 and lp < W) else None)

        elim_weeks = 0
        hit_weeks = 0
        pairwise_rates = []
        margins = []
        slack_sum = 0.0

        names = dwts_s["celebrity_name"].astype(str).str.strip().tolist()
        pros  = dwts_s["ballroom_partner"].astype(str).str.strip().tolist()

        for w in range(1, W + 1):
            idx = np.where(active[:, w-1])[0]
            if len(idx) == 0:
                continue

            E = [i for i, ew in enumerate(elim_week) if ew == w]
            if not E:
                continue

            # build x_local from estimator output
            x_local = []
            ok_join = True
            for i in idx:
                key = (int(s), int(w), names[i], pros[i])
                if key not in est_key:
                    ok_join = False
                    break
                x_local.append(est_key[key])
            if not ok_join:
                continue
            x_local = np.array(x_local, float)

            p_local = p[idx, w-1]
            C_local = wJ * p_local + wF * x_local

            m = len(E)
            elim_local = [np.where(idx == i)[0][0] for i in E if i in set(idx)]
            if len(elim_local) == 0:
                continue

            bottom_m = np.argpartition(C_local, m-1)[:m]
            hit = set(elim_local).issubset(set(bottom_m.tolist()))

            survivors = [j for j in range(len(idx)) if j not in elim_local]
            pair_ok = []
            for e in elim_local:
                for j in survivors:
                    pair_ok.append(C_local[j] > C_local[e])
            pair_rate = float(np.mean(pair_ok)) if pair_ok else np.nan

            margin = float(np.min(C_local[survivors]) - np.max(C_local[elim_local])) if survivors else np.nan

            sl = sanity[(sanity["season"] == s) & (sanity["week"] == w)]["slack_week"]
            sl = float(sl.iloc[0]) if len(sl) else np.nan
            if np.isfinite(sl):
                slack_sum += sl

            elim_weeks += 1
            hit_weeks += int(hit)
            if np.isfinite(pair_rate):
                pairwise_rates.append(pair_rate)
            if np.isfinite(margin):
                margins.append(margin)

            week_rows.append({
                "season": int(s),
                "week": int(w),
                "n_active": int(len(idx)),
                "n_elim": int(m),
                "bottom_m_hit": bool(hit),
                "pairwise_satisfaction_rate": pair_rate,
                "margin": margin,
                "slack_week": sl,
            })

        season_rows.append({
            "season": int(s),
            "elim_weeks_evaluated": int(elim_weeks),
            "bottom_m_hit_rate": float(hit_weeks / elim_weeks) if elim_weeks else np.nan,
            "avg_pairwise_satisfaction": float(np.mean(pairwise_rates)) if pairwise_rates else np.nan,
            "avg_margin": float(np.mean(margins)) if margins else np.nan,
            "total_slack_sum": float(slack_sum),

            "max_share_sum_error": float(sanity[sanity["season"] == s]["share_sum_error"].max()) if not sanity[sanity["season"] == s].empty else np.nan,
            "max_vote_pool_rel_error": float(sanity[sanity["season"] == s]["vote_pool_rel_error"].max()) if not sanity[sanity["season"] == s].empty else np.nan,
            "weeks_with_filled_viewers": int(sanity[(sanity["season"] == s) & (sanity["viewers_filled"] == True)].shape[0]),
        })

    week_df = pd.DataFrame(week_rows)
    season_df = pd.DataFrame(season_rows)

    sanity.to_csv(out_prefix + "sanity_by_week.csv", index=False)
    week_df.to_csv(out_prefix + "week_level_consistency.csv", index=False)
    season_df.to_csv(out_prefix + "season_summary.csv", index=False)

    print("Saved:")
    print(" -", out_prefix + "sanity_by_week.csv")
    print(" -", out_prefix + "week_level_consistency.csv")
    print(" -", out_prefix + "season_summary.csv")


if __name__ == "__main__":
    run_analyzer(
        est_csv="dwts_fan_vote_estimates_components.csv",
        dwts_csv="2026_MCM_Problem_C_Data.csv",
        out_prefix="dwts_check_",
        wJ=1.0, wF=1.0
    )
