import pandas as pd
import numpy as np
import re
from pathlib import Path

# ============================================================
# Config (edit these paths)
# ============================================================
INPUT_PATH = Path(r"2026_MCM_Problem_C_Data.xlsx")
OUTPUT_PATH = Path(r"2026_MCM_C_clean_panel_week.xlsx")

# ============================================================
# Helpers
# ============================================================
def find_score_columns(cols):
    """Find columns like weekX_judgeY_score."""
    pattern = re.compile(r"^week(\d+)_judge(\d+)_score$")
    return [c for c in cols if pattern.match(str(c))]

# ============================================================
# 1) Load
# ============================================================
if not INPUT_PATH.exists():
    raise FileNotFoundError(f"Input file not found: {INPUT_PATH.resolve()}")

raw = pd.read_excel(INPUT_PATH)

# Rename the awkward column name to a clean English name (if present)
raw = raw.rename(columns={"celebrity_homecountry/region": "celebrity_homecountry_region"})

# ============================================================
# 2) Identify columns
# ============================================================
score_cols = find_score_columns(raw.columns)
static_cols = [c for c in raw.columns if c not in score_cols]

# Minimal keys (DO NOT include NaN-prone static fields here)
id_cols = ["season", "celebrity_name"]

# ============================================================
# 3) Static table (one row per contestant-season)
# ============================================================
static = (
    raw[static_cols]
    .drop_duplicates(subset=id_cols)
    .copy()
)

# ============================================================
# 4) Wide -> long (one row per judge score)
# ============================================================
long_scores = raw.melt(
    id_vars=id_cols,
    value_vars=score_cols,
    var_name="week_judge",
    value_name="raw_score",
)

wk_j = long_scores["week_judge"].astype(str).str.extract(r"^week(\d+)_judge(\d+)_score$")
long_scores["week"] = wk_j[0].astype(int)
long_scores["judge_id"] = wk_j[1].astype(int)

long_scores["raw_score"] = pd.to_numeric(long_scores["raw_score"], errors="coerce")

# Dataset rule: 0 indicates "not performed after elimination" -> NaN
long_scores["score"] = long_scores["raw_score"].where(long_scores["raw_score"] > 0, np.nan)

# ============================================================
# 5) Drop padded weeks that do not exist per season
# ============================================================
week_exists = (
    long_scores.groupby(["season", "week"])["score"]
    .apply(lambda s: s.notna().any())
    .rename("week_exists")
    .reset_index()
)

long_scores = long_scores.merge(week_exists, on=["season", "week"], how="left")
long_scores = long_scores[long_scores["week_exists"]].copy()

# Season max existing week
season_max_week = (
    week_exists[week_exists["week_exists"]]
    .groupby("season")["week"]
    .max()
    .rename("season_max_week")
    .reset_index()
)

# ============================================================
# 6) Long -> panel per (season, celebrity_name, week)
# ============================================================
panel = (
    long_scores.pivot_table(
        index=id_cols + ["week"],
        columns="judge_id",
        values="score",
        aggfunc="first",
    )
    .reset_index()
)

# Rename judge columns to judge1_score..judge4_score
panel.columns = [
    f"judge{int(c)}_score" if isinstance(c, (int, np.integer)) else c
    for c in panel.columns
]

# Ensure judge1-4 columns exist
for j in range(1, 5):
    col = f"judge{j}_score"
    if col not in panel.columns:
        panel[col] = np.nan

judge_cols = [f"judge{j}_score" for j in range(1, 5)]

# ============================================================
# 7) Derived judge metrics
# ============================================================
panel["num_judges"] = panel[judge_cols].notna().sum(axis=1)

# min_count=1 avoids turning all-NaN into 0
panel["judge_total"] = panel[judge_cols].sum(axis=1, skipna=True, min_count=1)

# Internal active flag (will NOT be exported)
panel["__active"] = panel["judge_total"].fillna(0) > 0

# Merge static info and season max week
panel = panel.merge(static, on=id_cols, how="left")
panel = panel.merge(season_max_week, on="season", how="left")

# ============================================================
# 8) Elimination week (single final column)
#    A) parsed from results text: "Eliminated Week X"
#    B) inferred from last active week (covers "2nd Place", "3rd Place", etc.)
# ============================================================
if "results" in panel.columns:
    elim_txt = panel["results"].astype(str).str.extract(r"Eliminated Week\s*(\d+)", expand=False)
    elim_parsed = pd.to_numeric(elim_txt, errors="coerce")
else:
    elim_parsed = pd.Series(np.nan, index=panel.index)

last_active_week = (
    panel.loc[panel["__active"]]
    .groupby(id_cols)["week"]
    .max()
    .rename("last_active_week")
    .reset_index()
)
panel = panel.merge(last_active_week, on=id_cols, how="left")

elim_inferred = np.where(
    panel["last_active_week"].notna() & (panel["last_active_week"] < panel["season_max_week"]),
    panel["last_active_week"],
    np.nan,
)

panel["elimination_week"] = elim_parsed.fillna(pd.Series(elim_inferred, index=panel.index))

panel["eliminated_this_week"] = (
    panel["elimination_week"].notna() & (panel["week"] == panel["elimination_week"])
).astype(int)

# ============================================================
# 9) Weekly judge_rank and judge_percent (NO groupby.apply)
# ============================================================
panel["judge_rank"] = np.nan
panel["judge_percent"] = np.nan

active_mask = panel["__active"]

panel.loc[active_mask, "judge_rank"] = (
    panel.loc[active_mask]
    .groupby(["season", "week"])["judge_total"]
    .rank(method="min", ascending=False)
)

denom = (
    panel.loc[active_mask]
    .groupby(["season", "week"])["judge_total"]
    .transform("sum")
)
panel.loc[active_mask, "judge_percent"] = panel.loc[active_mask, "judge_total"] / denom

# ============================================================
# 10) Drop internal / intermediate columns for MAIN output
#     (Also remove fan_votes/fan_rank/fan_percent/pair_key as requested)
# ============================================================
drop_cols = ["__active", "season_max_week", "last_active_week",
             "fan_votes", "fan_rank", "fan_percent", "pair_key"]
panel_main = panel.drop(columns=[c for c in drop_cols if c in panel.columns], errors="ignore")

# ============================================================
# 11) Reorder columns for output (your exact requirement)
#     1) celebrity_name
#     2) ballroom_partner
#     3) season
#     then the rest shifted accordingly
# ============================================================
desired_order = [
    "celebrity_name",
    "ballroom_partner",
    "season",
    "week",
    "judge1_score",
    "judge2_score",
    "judge3_score",
    "judge4_score",
    "num_judges",
    "judge_total",
    "results",
    "placement",
    "elimination_week",
    "eliminated_this_week",
    "judge_rank",
    "judge_percent",
    "celebrity_industry",
    "celebrity_homestate",
    "celebrity_homecountry_region",
    "celebrity_age_during_season",
]

# Keep only existing columns, append any remaining columns at the end (safety)
existing = [c for c in desired_order if c in panel_main.columns]
remaining = [c for c in panel_main.columns if c not in existing]
panel_main = panel_main[existing + remaining]

# ============================================================
# 12) Build two sorted sheets
#     Sheet1: panel_by_week (first)
#     Sheet2: panel_by_pairkey (second)
# ============================================================
panel_by_week = panel_main.sort_values(
    ["season", "week", "judge_rank", "celebrity_name"],
    na_position="last"
).reset_index(drop=True)

panel_by_pairkey = panel_main.sort_values(
    ["season", "celebrity_name", "week"],
    na_position="last"
).reset_index(drop=True)

# ============================================================
# 13) Export (sheet order matters)
# ============================================================
with pd.ExcelWriter(OUTPUT_PATH, engine="openpyxl") as writer:
    panel_by_week.to_excel(writer, index=False, sheet_name="panel_by_week")          # 1st sheet
    panel_by_pairkey.to_excel(writer, index=False, sheet_name="panel_by_pairkey")    # 2nd sheet

print(f"Saved: {OUTPUT_PATH.resolve()}")
