# -*- coding: utf-8 -*-
"""
DWTS Fan Vote Estimation (Q1) — Corrected Model + Weekly Viewers Vote Pool
========================================================================

Reads weekly viewers data from:
  dwts_total_viewers_s1_s34_fixed.xlsx
Columns used:
  season, episode_in_season, total_viewers_millions

Interpretation:
  total_viewers_millions = 12.08 means 12.08M viewers.
We convert to weekly total votes (vote pool):
  week_total_votes = total_viewers_millions * 1e6

Missing viewers filling rules:
  1) within-season: linear interpolate over week, then ffill/bfill
  2) across-season (same week): nearest season with data for that week
  3) fallback: default_week_total_votes (converted to millions)

Core model:
  - We estimate weekly fan share x_{i,w} (sum to 1 over active couples).
  - Elimination constraints enforced (soft) under percent scheme:
        C_i = wJ * p_i + wF * x_i
    eliminated must be in bottom-m of C (m = number eliminated that week).
  - We decompose x into star-fixed / pro-fixed / performance parts for interpretation.
  - Viewer-based vote pool only affects fan_votes_total_est output (share x unchanged).

Outputs:
  - dwts_fan_vote_estimates_components.csv
  - dwts_fan_vote_bounds_total.csv
  - dwts_star_pro_season_end_fans.csv
  - dwts_season_consistency_summary.csv

Dependencies:
  pip install pandas numpy scipy openpyxl
"""

import re
import numpy as np
import pandas as pd
from scipy.optimize import minimize, Bounds, LinearConstraint
from scipy.optimize import linprog

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

PARAMS = dict(
    # show combine weights
    wJ=1.0,
    wF=1.0,

    # fixed votes: pro-stock conversion multiplier
    alpha_pro_vote=0.70,

    # performance votes (weekly, from 0 baseline)
    perf_scale=0.40,
    perf_power_excellence=1.40,
    perf_danger_weight=0.60,         # near-bottom mobilization for weekly perf votes
    danger_bottom_k=2,               # bottom-k zone
    perf_danger_power=1.20,

    # fan-stock update dynamics (log space)
    star_gain_lr=0.18,
    star_loss_lr=0.22,
    star_decay=0.02,

    pro_gain_lr=0.16,
    pro_loss_lr=0.18,
    pro_decay=0.02,

    # nonlinear gain/loss shapes
    gain_power_excellence=1.40,
    danger_gain_weight=0.80,         # near-bottom mobilization for gaining fans
    danger_gain_power=1.20,
    loss_power=1.60,

    # tightness mapping (spread -> tight)
    tight_k=18.0,
    tight_threshold=0.20,
    eps=1e-12,

    # weekly inverse (x) solver
    lam_prior_x=1.0,
    lam_slack=120.0,
    maxiter_week=3600,

    # returning-star/pro shrinkage (outer loop)
    outer_iters=20,
    shrink_star=0.70,
    shrink_pro=0.70,

    # initial stock scales (absolute scale only affects reporting)
    init_star_total=1.0e6,
    init_pro_total=4.0e5,

    # fallback if viewers missing even after filling
    default_week_total_votes=10_000_000,
)

# ---------------------------
# Utilities: parse & features
# ---------------------------

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 load_dwts_wide_from_clean_panel(panel_xlsx_path: str, sheet_name: str = "panel_by_week") -> pd.DataFrame:
    """
    Load the cleaned long-format weekly panel (one row per pair-week) and pivot it back to the
    wide DWTS format expected by this solver (one row per pair-season with week{w}_judge{j}_score cols).

    Expected columns in the panel:
      celebrity_name, ballroom_partner, season, week,
      judge1_score..judgeK_score (typically 1..4),
      plus optionally: results, placement, celebrity_* metadata.

    Returns a DataFrame compatible with the original 2026_MCM_Problem_C_Data.csv schema.
    """
    panel = pd.read_excel(panel_xlsx_path, sheet_name=sheet_name)
    panel.columns = [str(c).strip() for c in panel.columns]

    # normalize a known naming difference
    if "celebrity_homecountry_region" in panel.columns and "celebrity_homecountry/region" not in panel.columns:
        panel = panel.rename(columns={"celebrity_homecountry_region": "celebrity_homecountry/region"})

    required = ["celebrity_name", "ballroom_partner", "season", "week"]
    missing = [c for c in required if c not in panel.columns]
    if missing:
        raise ValueError(f"clean panel missing required columns: {missing}. Found: {list(panel.columns)}")

    # numeric sanitize
    panel["season"] = pd.to_numeric(panel["season"], errors="coerce").astype("Int64")
    panel["week"] = pd.to_numeric(panel["week"], errors="coerce")
    panel = panel.dropna(subset=["season", "week"]).copy()
    panel["season"] = panel["season"].astype(int)
    panel["week"] = panel["week"].round().astype(int)

    # judge columns like judge1_score, judge2_score, ...
    judge_cols = [c for c in panel.columns if re.fullmatch(r"judge\d+_score", str(c), flags=re.I)]
    if not judge_cols:
        raise ValueError(
            "No judge*_score columns found in clean panel. "
            "Expected columns like judge1_score, judge2_score, ..."
        )

    base_cols = ["season", "celebrity_name", "ballroom_partner", "week"]
    long = panel[base_cols + judge_cols].melt(
        id_vars=base_cols, value_vars=judge_cols,
        var_name="judge_col", value_name="score"
    )
    long["judge"] = long["judge_col"].str.extract(r"judge(\d+)_score", flags=re.I).astype(int)
    long["wk_judge_col"] = long.apply(
        lambda r: f"week{int(r.week)}_judge{int(r.judge)}_score", axis=1
    )

    wide_scores = long.pivot_table(
        index=["celebrity_name", "ballroom_partner", "season"],
        columns="wk_judge_col",
        values="score",
        aggfunc="first",
    )
    wide_scores.columns.name = None
    wide_scores = wide_scores.reset_index()

    # carry static/meta columns (first occurrence per pair-season)
    meta_candidates = [
        "celebrity_industry",
        "celebrity_homestate",
        "celebrity_homecountry/region",
        "celebrity_age_during_season",
        "results",
        "placement",
    ]
    meta_cols = [c for c in meta_candidates if c in panel.columns]
    if meta_cols:
        meta = (
            panel.sort_values(["season", "celebrity_name", "ballroom_partner", "week"])
                 .groupby(["celebrity_name", "ballroom_partner", "season"], as_index=False)[meta_cols]
                 .first()
        )
    else:
        meta = wide_scores[["celebrity_name", "ballroom_partner", "season"]].copy()

    out = meta.merge(wide_scores, on=["celebrity_name", "ballroom_partner", "season"], how="left")
    return out


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 judge_rank_and_q(jw):
    n = len(jw)
    order = (-jw).argsort()
    ranks = np.empty(n, float)
    sj = jw[order]
    i = 0
    while i < n:
        j = i
        while j + 1 < n and np.isclose(sj[j + 1], sj[i], atol=1e-9, rtol=1e-9):
            j += 1
        ranks[order[i:j + 1]] = (i + 1 + j + 1) / 2.0
        i = j + 1
    q = (n - ranks + 1) / n
    return ranks, q

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))

def tightness_from_spread(jw, params):
    eps = params["eps"]
    spread = (np.max(jw) - np.min(jw)) / (np.mean(jw) + eps)
    tight = sigmoid(params["tight_k"] * (params["tight_threshold"] - spread))
    return float(spread), float(tight)

def excellence_score(q, params):
    return q ** params["gain_power_excellence"]

def danger_score(q, ranks, tight, n, params):
    bottom_k = int(params["danger_bottom_k"])
    danger_mask = (ranks >= (n - bottom_k + 1)).astype(float)
    return tight * danger_mask * ((1.0 - q) ** params["danger_gain_power"])

def loss_score(q, params):
    return (1.0 - q) ** params["loss_power"]

def perf_vote_component(q, ranks, tight, n, params):
    exc = q ** params["perf_power_excellence"]
    bottom_k = int(params["danger_bottom_k"])
    danger_mask = (ranks >= (n - bottom_k + 1)).astype(float)
    dang = tight * danger_mask * ((1.0 - q) ** params["perf_danger_power"])
    return params["perf_scale"] * (exc + params["perf_danger_weight"] * dang)

def project_to_simplex(v):
    v = np.asarray(v, float)
    n = len(v)
    u = np.sort(v)[::-1]
    cssv = np.cumsum(u)
    rho = np.where(u + (1 - cssv) / (np.arange(n) + 1) > 0)[0]
    if len(rho) == 0:
        return np.ones(n) / n
    rho = rho[-1]
    theta = (cssv[rho] - 1) / (rho + 1)
    w = np.maximum(v - theta, 0)
    s = w.sum()
    return w / s if s > 0 else np.ones(n) / n

# ---------------------------
# Weekly solver: infer x with soft elimination constraints (percent scheme)
# ---------------------------

def solve_week_x_percent(p, x_hat, elim_local, params):
    """
    Variables: x (n), z (m>=0)
    Constraints:
      sum x = 1, x>=0
      For each eliminated e and survivor j:
        wF*(x_e - x_j) - z_k <= wJ*(p_j - p_e)
    Objective:
      lam_prior ||x - x_hat||^2 + lam_slack * sum(z)
    """
    n = len(p)
    wJ, wF = params["wJ"], params["wF"]
    lam_prior = params["lam_prior_x"]
    lam_slack = params["lam_slack"]

    elim_local = list(elim_local) if elim_local is not None else []
    survivors = [j for j in range(n) if j not in elim_local]
    pairs = []
    for e in elim_local:
        for j in survivors:
            rhs = wJ * (p[j] - p[e])
            pairs.append((e, j, rhs))
    m = len(pairs)

    if m == 0:
        return project_to_simplex(np.maximum(x_hat, 0)), 0.0

    dim = n + m
    lb = np.zeros(dim)
    ub = np.ones(dim)
    ub[n:] = np.inf
    bnds = Bounds(lb, ub)

    Aeq = np.zeros((1, dim))
    Aeq[0, :n] = 1.0
    ceq = LinearConstraint(Aeq, [1.0], [1.0])

    A = np.zeros((m, dim))
    ubA = np.zeros(m)
    for k, (e, j, rhs) in enumerate(pairs):
        A[k, e] = wF
        A[k, j] = -wF
        A[k, n + k] = -1.0
        ubA[k] = rhs
    cin = LinearConstraint(A, -np.inf * np.ones(m), ubA)

    x0 = project_to_simplex(np.maximum(x_hat, 1e-12))
    u0 = np.concatenate([x0, np.zeros(m)])

    def fun(u):
        x = u[:n]
        z = u[n:]
        return lam_prior * np.sum((x - x_hat) ** 2) + lam_slack * np.sum(z)

    def jac(u):
        x = u[:n]
        gx = 2.0 * lam_prior * (x - x_hat)
        gz = np.full(m, lam_slack)
        return np.concatenate([gx, gz])

    def hess(_u):
        H = np.zeros((dim, dim))
        H[:n, :n] = 2.0 * lam_prior * np.eye(n)
        return H

    res = minimize(
        fun, u0,
        method="trust-constr",
        jac=jac,
        hess=hess,
        constraints=[ceq, cin],
        bounds=bnds,
        options={"maxiter": int(params["maxiter_week"]), "verbose": 0}
    )

    if not res.success:
        # fallback to x_hat projection
        return project_to_simplex(np.maximum(x_hat, 0)), float("nan")

    x = project_to_simplex(np.maximum(res.x[:n], 0))
    slack_sum = float(np.sum(res.x[n:]))
    return x, slack_sum

# ---------------------------
# Uncertainty bounds for total x (hard constraints)
# ---------------------------

def week_bounds_total_x(p_local, elim_local, params):
    """
    Hard constraints:
      wJ*p_e + wF*x_e <= wJ*p_j + wF*x_j
    => wF*(x_e - x_j) <= wJ*(p_j - p_e)
    """
    n = len(p_local)
    wJ, wF = params["wJ"], params["wF"]
    E = list(elim_local)
    survivors = [j for j in range(n) if j not in E]

    A_ub, b_ub = [], []
    for e in E:
        for j in survivors:
            row = np.zeros(n)
            row[e] = wF
            row[j] = -wF
            A_ub.append(row)
            b_ub.append(wJ * (p_local[j] - p_local[e]))

    A_ub = np.array(A_ub) if A_ub else None
    b_ub = np.array(b_ub) if b_ub else None
    A_eq = np.ones((1, n))
    b_eq = np.array([1.0])
    bounds = [(0, None)] * n

    mins = np.full(n, np.nan)
    maxs = np.full(n, np.nan)
    for i in range(n):
        c = np.zeros(n); c[i] = 1.0
        r1 = linprog(c, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=b_eq, bounds=bounds, method="highs")
        if r1.success:
            mins[i] = r1.fun
        r2 = linprog(-c, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=b_eq, bounds=bounds, method="highs")
        if r2.success:
            maxs[i] = -r2.fun
    return mins, maxs

# ---------------------------
# Viewers (weekly vote pool): read xlsx + fill
# ---------------------------

def build_viewers_lookup(dwts_df, viewers_xlsx_path, params):
    """
    Returns dict:
      (season, week) -> (viewers_millions, week_total_votes, filled_flag, fill_method)

    Only builds season-week grid for weeks that exist in DWTS data (per season_length).
    """
    viewers = pd.read_excel(viewers_xlsx_path)
    colmap = {c.strip().lower(): c for c in viewers.columns}

    required = ["season", "episode_in_season", "total_viewers_millions"]
    for k in required:
        if k not in colmap:
            raise ValueError(f"Viewers xlsx missing column '{k}'. Found: {list(viewers.columns)}")

    viewers = viewers[[colmap["season"], colmap["episode_in_season"], colmap["total_viewers_millions"]]].copy()
    viewers.columns = ["season", "week", "viewers_millions"]

    viewers["season"] = pd.to_numeric(viewers["season"], errors="coerce").astype("Int64")
    viewers["week"] = pd.to_numeric(viewers["week"], errors="coerce")
    viewers["viewers_millions"] = pd.to_numeric(viewers["viewers_millions"], errors="coerce")

    viewers = viewers.dropna(subset=["season", "week"]).copy()
    viewers["week"] = viewers["week"].round().astype(int)
    viewers = viewers[viewers["week"] > 0].copy()

    # duplicates -> average
    viewers = viewers.groupby(["season", "week"], as_index=False)["viewers_millions"].mean()

    seasons = sorted(dwts_df["season"].dropna().unique().astype(int).tolist())
    weeks_all, week_cols_all = find_week_judge_columns(dwts_df)

    grid_rows = []
    for s in seasons:
        s_df = dwts_df[dwts_df["season"] == s].copy()
        W = season_length(s_df, weeks_all, week_cols_all)
        for w in range(1, W + 1):
            grid_rows.append((s, w))

    grid = pd.DataFrame(grid_rows, columns=["season", "week"])
    merged = grid.merge(viewers, on=["season", "week"], how="left")

    merged["filled"] = False
    merged["fill_method"] = ""
    merged["viewers_millions_filled"] = merged["viewers_millions"].astype(float)

    # mark observed
    obs_mask = merged["viewers_millions_filled"].notna()
    merged.loc[obs_mask, "fill_method"] = "observed"

    # Step 1: within-season interpolation + ffill/bfill
    merged = merged.sort_values(["season", "week"]).reset_index(drop=True)
    for s in seasons:
        mask = (merged["season"] == s)
        sub = merged.loc[mask, ["week", "viewers_millions_filled"]].copy()
        before_na = sub["viewers_millions_filled"].isna()

        sub["viewers_millions_filled"] = sub["viewers_millions_filled"].interpolate(
            method="linear", limit_direction="both"
        )
        sub["viewers_millions_filled"] = sub["viewers_millions_filled"].ffill().bfill()

        merged.loc[mask, "viewers_millions_filled"] = sub["viewers_millions_filled"].values
        after_na = sub["viewers_millions_filled"].isna()
        filled_now = before_na & (~after_na)

        if filled_now.any():
            idxs = merged.loc[mask].index[filled_now.values]
            merged.loc[idxs, "filled"] = True
            merged.loc[idxs, "fill_method"] = "within_season_interp_ffill_bfill"

    # Step 2: remaining missing -> nearest season with same week
    remaining = merged["viewers_millions_filled"].isna()
    if remaining.any():
        avail = merged.dropna(subset=["viewers_millions_filled"]).copy()
        for ridx in merged[remaining].index:
            s = int(merged.at[ridx, "season"])
            w = int(merged.at[ridx, "week"])
            cand = avail[avail["week"] == w]
            if cand.empty:
                continue
            cand = cand.assign(dist=(cand["season"] - s).abs()).sort_values("dist")
            merged.at[ridx, "viewers_millions_filled"] = float(cand.iloc[0]["viewers_millions_filled"])
            merged.at[ridx, "filled"] = True
            merged.at[ridx, "fill_method"] = "nearest_season_same_week"

    # Step 3: fallback
    remaining = merged["viewers_millions_filled"].isna()
    if remaining.any():
        fallback_m = float(params["default_week_total_votes"]) / 1e6
        merged.loc[remaining, "viewers_millions_filled"] = fallback_m
        merged.loc[remaining, "filled"] = True
        merged.loc[remaining, "fill_method"] = "default_fallback"

    merged["week_total_votes"] = merged["viewers_millions_filled"].astype(float) * 1e6

    lookup = {}
    for r in merged.itertuples(index=False):
        lookup[(int(r.season), int(r.week))] = (
            float(r.viewers_millions_filled),
            float(r.week_total_votes),
            bool(r.filled),
            str(r.fill_method)
        )
    return lookup

# ---------------------------
# Season solver (sequential)
# ---------------------------

def solve_season_corrected(season_df, star_global_log, pro_global_log, params):
    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)
    N = len(season_df)

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

    active = (J > 0)

    # 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(N)], int)
    elim_week = []
    for i, resu in enumerate(season_df["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)

    # init stocks
    init_star_total = params["init_star_total"]
    init_pro_total = params["init_pro_total"]
    shrink_star = params["shrink_star"]
    shrink_pro = params["shrink_pro"]

    idx1 = np.where(active[:, 0])[0]
    if len(idx1) > 0:
        jw1 = J[idx1, 0]
        z = (jw1 - jw1.mean()) / (jw1.std() + 1e-9)
        wgt = np.exp(0.8 * z)
        wgt = wgt / wgt.sum()
        star_guess = np.zeros(N, float)
        star_guess[idx1] = init_star_total * wgt
    else:
        star_guess = np.full(N, init_star_total / max(N, 1), float)

    pro_guess = np.full(N, init_pro_total / max(N, 1), float)

    logF_star = np.full((N, W + 1), np.nan, float)
    logF_pro  = np.full((N, W + 1), np.nan, float)

    for i in range(N):
        c = names[i]
        p = pros[i]
        gS = star_global_log.get(c, np.log(star_guess[i] + 1.0))
        gP = pro_global_log.get(p, np.log(pro_guess[i] + 1.0))
        logF_star[i, 0] = shrink_star * gS + (1 - shrink_star) * np.log(star_guess[i] + 1.0)
        logF_pro[i, 0]  = shrink_pro  * gP + (1 - shrink_pro)  * np.log(pro_guess[i] + 1.0)

    x_total = np.zeros((N, W), float)
    x_star_fixed = np.zeros((N, W), float)
    x_pro_fixed = np.zeros((N, W), float)
    x_perf = np.zeros((N, W), float)

    slack_w = np.zeros(W, float)
    pmat = np.zeros((N, W), float)
    rankmat = np.full((N, W), np.nan, float)
    qmat = np.zeros((N, W), float)
    spread_w = np.full(W, np.nan, float)
    tight_w = np.full(W, np.nan, float)

    end_star_fans = np.full(N, np.nan, float)
    end_pro_fans = np.full(N, np.nan, float)
    end_week = np.full(N, np.nan, 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_local = jw / denom if denom > 0 else np.ones(len(idx)) / len(idx)

        ranks, q = judge_rank_and_q(jw)
        spread, tight = tightness_from_spread(jw, params)

        pmat[idx, w - 1] = p_local
        rankmat[idx, w - 1] = ranks
        qmat[idx, w - 1] = q
        spread_w[w - 1] = spread
        tight_w[w - 1] = tight

        # elimination set at end of week w
        E = [i for i, ew in enumerate(elim_week) if ew == w]
        idx_set = set(idx.tolist())
        local_map = {i: k for k, i in enumerate(idx)}
        elim_local = [local_map[i] for i in E if i in idx_set]

        # build x_hat from predicted votes (arbitrary scale)
        F_star = np.exp(logF_star[idx, w - 1])
        F_pro  = np.exp(logF_pro[idx, w - 1])

        v_star_fixed = F_star
        v_pro_fixed = params["alpha_pro_vote"] * F_pro
        v_perf = perf_vote_component(q, ranks, tight, len(idx), params)

        v_pred = v_star_fixed + v_pro_fixed + v_perf + params["eps"]
        x_hat = v_pred / v_pred.sum()

        x, slack = solve_week_x_percent(p_local, x_hat, elim_local, params)
        slack_w[w - 1] = 0.0 if np.isnan(slack) else slack

        # decompose by predicted proportions
        frac_star = v_star_fixed / v_pred
        frac_pro  = v_pro_fixed / v_pred
        frac_perf = v_perf / v_pred

        x_total[idx, w - 1] = x
        x_star_fixed[idx, w - 1] = x * frac_star
        x_pro_fixed[idx, w - 1]  = x * frac_pro
        x_perf[idx, w - 1]       = x * frac_perf

        # update fan stocks AFTER week w
        exc = excellence_score(q, params)
        dang = danger_score(q, ranks, tight, len(idx), params)
        gain_term = exc + params["danger_gain_weight"] * dang
        loss_term = loss_score(q, params)

        logF_star[idx, w] = logF_star[idx, w - 1] + params["star_gain_lr"] * gain_term - params["star_loss_lr"] * loss_term - params["star_decay"]
        logF_pro[idx, w]  = logF_pro[idx, w - 1]  + params["pro_gain_lr"]  * gain_term - params["pro_loss_lr"]  * loss_term - params["pro_decay"]

        # carry forward for non-updated
        for i in range(N):
            if np.isnan(logF_star[i, w]) and np.isfinite(logF_star[i, w - 1]):
                logF_star[i, w] = logF_star[i, w - 1]
            if np.isnan(logF_pro[i, w]) and np.isfinite(logF_pro[i, w - 1]):
                logF_pro[i, w] = logF_pro[i, w - 1]

        # record end stocks after update
        for i in idx:
            if elim_week[i] == w or (elim_week[i] is None and w == W):
                end_star_fans[i] = float(np.exp(logF_star[i, w]))
                end_pro_fans[i]  = float(np.exp(logF_pro[i, w]))
                end_week[i] = float(w)

    # season consistency: eliminated subset of bottom-m of C
    ok, tot = 0, 0
    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
        C = params["wJ"] * pmat[:, w - 1] + params["wF"] * x_total[:, w - 1]
        m = len(E)
        C_local = C[idx]
        bottom_m_local = idx[np.argpartition(C_local, m - 1)[:m]]
        tot += 1
        ok += int(set(E).issubset(set(bottom_m_local.tolist())))
    consistency = ok / tot if tot else np.nan

    return dict(
        W=W, J=J, active=active, elim_week=elim_week,
        names=names, pros=pros,
        x_total=x_total,
        x_star_fixed=x_star_fixed,
        x_pro_fixed=x_pro_fixed,
        x_perf=x_perf,
        p=pmat, rank=rankmat, q=qmat,
        spread=spread_w, tight=tight_w,
        slack_week=slack_w,
        end_star_fans=end_star_fans,
        end_pro_fans=end_pro_fans,
        end_week=end_week,
        consistency=consistency,
        slack_total=float(np.nansum(slack_w)),
        start_logF_star=logF_star[:, 0],
        start_logF_pro=logF_pro[:, 0],
    )

# ---------------------------
# Outer loop across seasons
# ---------------------------

def run_all_seasons(dwts_csv_path, viewers_xlsx_path, out_prefix="dwts_", params=None):
    if params is None:
        params = PARAMS

        # DWTS main table can be provided either as the original wide CSV
    # OR as the cleaned long weekly panel xlsx (we pivot it back to wide here).
    if str(dwts_csv_path).lower().endswith((".xlsx", ".xls")):
        df = load_dwts_wide_from_clean_panel(dwts_csv_path, sheet_name=params.get("panel_sheet", "panel_by_week"))
    else:
        df = pd.read_csv(dwts_csv_path)


    viewers_lookup = build_viewers_lookup(df, viewers_xlsx_path, params)
    seasons = sorted(df["season"].dropna().unique().astype(int).tolist())

    # init global baselines (log)
    star_global_log = {}
    pro_global_log = {}
    all_names = df["celebrity_name"].astype(str).unique().tolist()
    all_pros  = df["ballroom_partner"].astype(str).unique().tolist()
    init_star = np.log(params["init_star_total"] / max(1, len(all_names)) + 1.0)
    init_pro  = np.log(params["init_pro_total"] / max(1, len(all_pros)) + 1.0)
    for n in all_names:
        star_global_log[n] = init_star
    for p in all_pros:
        pro_global_log[p] = init_pro

    est_rows, bound_rows, end_rows, summary_rows = [], [], [], []

    outer_iters = int(params["outer_iters"])
    for it in range(1, outer_iters + 1):
        star_starts = {k: [] for k in star_global_log.keys()}
        pro_starts = {k: [] for k in pro_global_log.keys()}

        keep_outputs = (it == outer_iters)
        if keep_outputs:
            est_rows.clear(); bound_rows.clear(); end_rows.clear(); summary_rows.clear()

        for season in seasons:
            season_df = df[df["season"] == season].reset_index(drop=True)
            sol = solve_season_corrected(season_df, star_global_log, pro_global_log, params)

            # collect season-start stocks for empirical-bayes update
            for i, c in enumerate(sol["names"]):
                if np.isfinite(sol["start_logF_star"][i]):
                    star_starts[c].append(float(sol["start_logF_star"][i]))
            for i, p in enumerate(sol["pros"]):
                if np.isfinite(sol["start_logF_pro"][i]):
                    pro_starts[p].append(float(sol["start_logF_pro"][i]))

            if not keep_outputs:
                continue

            summary_rows.append({
                "season": int(season),
                "weeks": int(sol["W"]),
                "consistency_rate_bottom_m": float(sol["consistency"]),
                "total_slack": float(sol["slack_total"]),
            })

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

                viewers_m, week_total_votes, filled_flag, fill_method = viewers_lookup.get(
                    (int(season), int(w)),
                    (float(params["default_week_total_votes"]) / 1e6,
                     float(params["default_week_total_votes"]),
                     True,
                     "default_fallback")
                )

                for i in idx:
                    x = sol["x_total"][i, w - 1]
                    est_rows.append({
                        "season": int(season),
                        "week": int(w),

                        "viewers_millions": float(viewers_m),
                        "week_total_votes": float(week_total_votes),
                        "viewers_filled": bool(filled_flag),
                        "viewers_fill_method": str(fill_method),

                        "celebrity_name": str(sol["names"][i]),
                        "ballroom_partner": str(sol["pros"][i]),

                        "judge_total": float(sol["J"][i, w - 1]),
                        "judge_percent": float(sol["p"][i, w - 1]),
                        "judge_rank": float(sol["rank"][i, w - 1]) if np.isfinite(sol["rank"][i, w - 1]) else np.nan,
                        "perf_q": float(sol["q"][i, w - 1]),
                        "spread_week": float(sol["spread"][w - 1]) if np.isfinite(sol["spread"][w - 1]) else np.nan,
                        "tight_week": float(sol["tight"][w - 1]) if np.isfinite(sol["tight"][w - 1]) else np.nan,

                        "fan_share_total": float(x),
                        "fan_share_star_fixed": float(sol["x_star_fixed"][i, w - 1]),
                        "fan_share_pro_fixed": float(sol["x_pro_fixed"][i, w - 1]),
                        "fan_share_perf": float(sol["x_perf"][i, w - 1]),

                        "fan_votes_total_est": float(x * week_total_votes),

                        "slack_week": float(sol["slack_week"][w - 1]),
                    })

                # bounds (certainty proxy) on elimination weeks
                E = [i for i, ew in enumerate(sol["elim_week"]) if ew == w]
                if E:
                    local_map = {i: k for k, i in enumerate(idx)}
                    elim_local = [local_map[i] for i in E if i in local_map]
                    p_local = sol["p"][idx, w - 1]
                    mins, maxs = week_bounds_total_x(p_local, elim_local, params)
                    for k, i in enumerate(idx):
                        mn, mx = mins[k], maxs[k]
                        bound_rows.append({
                            "season": int(season),
                            "week": int(w),
                            "celebrity_name": str(sol["names"][i]),
                            "ballroom_partner": str(sol["pros"][i]),
                            "fan_share_total_est": float(sol["x_total"][i, w - 1]),
                            "fan_share_total_min": float(mn) if np.isfinite(mn) else np.nan,
                            "fan_share_total_max": float(mx) if np.isfinite(mx) else np.nan,
                            "ci_width_proxy": float(mx - mn) if np.isfinite(mn) and np.isfinite(mx) else np.nan,
                        })

            # season-end stock table
            for i in range(len(sol["names"])):
                if np.isfinite(sol["end_week"][i]):
                    end_rows.append({
                        "season": int(season),
                        "celebrity_name": str(sol["names"][i]),
                        "ballroom_partner": str(sol["pros"][i]),
                        "end_week": int(sol["end_week"][i]),
                        "star_fans_after_end": float(sol["end_star_fans"][i]),
                        "pro_fans_after_end": float(sol["end_pro_fans"][i]),
                    })

        # update global baselines (mean of starts, smoothed)
        rho = 0.6
        for c, vals in star_starts.items():
            if len(vals) > 0:
                new_mu = float(np.mean(vals))
                star_global_log[c] = (1 - rho) * star_global_log[c] + rho * new_mu
        for p, vals in pro_starts.items():
            if len(vals) > 0:
                new_mu = float(np.mean(vals))
                pro_global_log[p] = (1 - rho) * pro_global_log[p] + rho * new_mu

        print(f"[Outer {it}/{outer_iters}] global baselines updated.")

    # write outputs
    pd.DataFrame(est_rows).to_csv(out_prefix + "fan_vote_estimates_components.csv", index=False)
    pd.DataFrame(bound_rows).to_csv(out_prefix + "fan_vote_bounds_total.csv", index=False)
    pd.DataFrame(end_rows).to_csv(out_prefix + "star_pro_season_end_fans.csv", index=False)
    pd.DataFrame(summary_rows).to_csv(out_prefix + "season_consistency_summary.csv", index=False)

    print("Saved:")
    print(" -", out_prefix + "fan_vote_estimates_components.csv")
    print(" -", out_prefix + "fan_vote_bounds_total.csv")
    print(" -", out_prefix + "star_pro_season_end_fans.csv")
    print(" -", out_prefix + "season_consistency_summary.csv")


if __name__ == "__main__":
    run_all_seasons(
        dwts_csv_path="2026_MCM_C_clean_panel_week.xlsx",
        viewers_xlsx_path="dwts_total_viewers_s1_s34_fixed.xlsx",
        out_prefix="dwts_",
        params=PARAMS
    )