"""
evaluate_sequence_block_model_calibrated.py
------------------------------------------
Compara el modelo OneVsRest actual vs una version calibrada (sigmoid)
usando el cache de features v3 y el mismo split holdout por partido.
"""

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

from sklearn.calibration import CalibratedClassifierCV
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import mean_absolute_error, roc_auc_score
from sklearn.model_selection import GroupShuffleSplit
from sklearn.multiclass import OneVsRestClassifier


PREPROCESSED = Path("/Users/pagrois/Documents/Racing/preprocessed_SSD_25-26.csv")
FEATURE_CACHE = Path("/tmp/sequence_block_model_cache/league_sequence_features_v3.pkl")
REPORTS_DIR = Path("/Users/pagrois/Racing/reports")
REPORTS_DIR.mkdir(parents=True, exist_ok=True)

BLOCKS = ["High", "Medium", "Low"]
BLOCK_TO_IDX = {b: i for i, b in enumerate(BLOCKS)}
PHASE_TO_BLOCK = {
    "Build Up against High Block": "High",
    "Build Up against Medium Block": "Medium",
    "Build Up against Low Block": "Low",
}
RANDOM_SEED = 42

FEATURE_COLS = [
    "feat_n_pressure_own_half",
    "feat_avg_pass_options_rival",
    "feat_top2_avg_x",
    "feat_team_avg_x",
    "feat_back4_avg_x",
    "feat_lowest_def_no_gk_avg_x",
    "feat_convex_hull_area_no_gk",
    "feat_team_width_no_gk",
    "feat_team_depth_no_gk",
    "feat_n_players_rival_half_avg",
    "feat_n_pressure_high",
    "feat_n_pressure_medium",
    "feat_n_pressure_low",
    "feat_avg_event_x",
    "feat_pct_actions_rival",
    "feat_n_actions_rival",
]


def resolve_blocks(grp):
    max_n = grp["n_ev"].max()
    tied = set(grp.loc[grp["n_ev"] == max_n, "block_raw"])
    if len(tied) == 1:
        return list(tied)
    valid = set(grp.loc[(grp["n_ev"] >= 3) & (grp["block_raw"].isin(tied)), "block_raw"])
    blocks = list(valid) if valid else list(tied)
    if len(blocks) > 1 and "Medium" in blocks:
        nm = [b for b in blocks if b != "Medium"]
        if nm:
            blocks = nm
    return blocks


def normalize_rows(arr):
    arr = np.asarray(arr, dtype=float)
    row_sum = arr.sum(axis=1, keepdims=True)
    out = arr.copy()
    mask = row_sum.squeeze() > 0
    out[mask] = out[mask] / row_sum[mask]
    out[~mask] = 1.0 / arr.shape[1]
    return out


def one_hot(labels_idx, n_classes=3):
    arr = np.zeros((len(labels_idx), n_classes), dtype=float)
    arr[np.arange(len(labels_idx)), labels_idx] = 1.0
    return arr


def probs_from_ovr(model, X):
    prob_list = model.predict_proba(X)
    if isinstance(prob_list, list):
        probs = np.column_stack([p[:, 1] for p in prob_list])
    elif isinstance(prob_list, np.ndarray) and prob_list.ndim == 3:
        probs = np.column_stack([prob_list[i][:, 1] for i in range(prob_list.shape[0])])
    else:
        probs = np.asarray(prob_list)
        if probs.ndim == 2 and probs.shape[1] == 3:
            return probs
    return probs


def topk_prediction_sets(probs_norm, n_passes, prob_threshold=0.28, gap_threshold=0.08):
    pred_sets = []
    for i in range(len(probs_norm)):
        order = np.argsort(probs_norm[i])[::-1]
        top1, top2 = order[0], order[1]
        if (
            n_passes[i] >= 7
            and probs_norm[i, top1] >= prob_threshold
            and probs_norm[i, top2] >= prob_threshold
            and abs(probs_norm[i, top1] - probs_norm[i, top2]) <= gap_threshold
        ):
            pred_sets.append({top1, top2})
        else:
            pred_sets.append({top1})
    return pred_sets


def set_metrics(y_bin, pred_sets):
    true_sets = [set(np.where(row > 0)[0]) for row in y_bin]
    overlap_hit = np.mean([len(t & p) > 0 for t, p in zip(true_sets, pred_sets)])
    exact_hit = np.mean([t == p for t, p in zip(true_sets, pred_sets)])
    predicted_multi = np.mean([len(p) > 1 for p in pred_sets])
    multi_rows = [(t, p) for t, p in zip(true_sets, pred_sets) if len(p) > 1]
    multi_overlap = np.mean([len(t & p) > 0 for t, p in multi_rows]) if multi_rows else np.nan
    return {
        "overlap_hit": overlap_hit,
        "exact_hit": exact_hit,
        "predicted_multi_rate": predicted_multi,
        "predicted_multi_overlap": multi_overlap,
    }


def evaluate(y_bin, y_weight, probs, n_passes):
    probs_norm = normalize_rows(probs)
    pred_idx = probs_norm.argmax(axis=1)
    top1_hit = y_bin[np.arange(len(y_bin)), pred_idx].mean()
    true_mult = y_bin.sum(axis=1) > 1
    multi_top1 = y_bin[np.arange(len(y_bin)), pred_idx][true_mult].mean() if true_mult.any() else np.nan
    multi_mae = mean_absolute_error(y_weight[true_mult], probs_norm[true_mult]) if true_mult.any() else np.nan
    pred_sets = topk_prediction_sets(probs_norm, n_passes)
    set_stats = set_metrics(y_bin, pred_sets)
    return {
        "auc_macro": roc_auc_score(y_bin, probs, average="macro"),
        "auc_micro": roc_auc_score(y_bin, probs, average="micro"),
        "mae": mean_absolute_error(y_weight, probs_norm),
        "top1_hit": top1_hit,
        "multi_top1_hit": multi_top1,
        "multi_mae": multi_mae,
        **set_stats,
    }


print("Cargando targets y features...", flush=True)
pre = pd.read_csv(PREPROCESSED, usecols=["matchId", "sequenceId", "phaseLabel"], low_memory=False)
pre["matchId"] = pre["matchId"].astype(str)
pre["sequenceId"] = pd.to_numeric(pre["sequenceId"], errors="coerce")
pre = pre.dropna(subset=["sequenceId"]).copy()
pre["sequenceId"] = pre["sequenceId"].astype(int).astype(str)
pre["block_raw"] = pre["phaseLabel"].map(PHASE_TO_BLOCK)

feat = pd.read_pickle(FEATURE_CACHE)
pass_count_map = feat.set_index(["matchId", "sequenceId"])["n_passes"].to_dict()
block_ev = (
    pre.dropna(subset=["block_raw"])
    .groupby(["matchId", "sequenceId", "block_raw"])
    .size()
    .reset_index(name="n_ev")
)

rows = []
for (mid, sid), grp in block_ev.groupby(["matchId", "sequenceId"]):
    n_passes = int(pass_count_map.get((mid, sid), -1))
    if n_passes < 3:
        continue
    grp_multi = grp[grp["n_ev"] >= 3]
    if n_passes >= 7 and len(grp_multi) > 1:
        total = grp_multi["n_ev"].sum()
        weights = {b: 0.0 for b in BLOCKS}
        for r in grp_multi.itertuples(index=False):
            weights[r.block_raw] = r.n_ev / total
    else:
        blocks = resolve_blocks(grp)
        weights = {b: 0.0 for b in BLOCKS}
        w = 1.0 / len(blocks)
        for b in blocks:
            weights[b] = w
    rows.append({"matchId": mid, "sequenceId": sid, "y_high": weights["High"], "y_medium": weights["Medium"], "y_low": weights["Low"]})

target = pd.DataFrame(rows)
df_model = feat.merge(target, on=["matchId", "sequenceId"], how="inner")
df_model = df_model[df_model["n_passes"] >= 3].dropna(subset=FEATURE_COLS).copy()

X = df_model[FEATURE_COLS].to_numpy(dtype=float)
y_weight = df_model[["y_high", "y_medium", "y_low"]].to_numpy(dtype=float)
y_bin = (y_weight > 0).astype(int)
groups = df_model["matchId"].to_numpy()
n_passes = df_model["n_passes"].to_numpy()

train_idx, test_idx = next(GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RANDOM_SEED).split(X, y_bin, groups))
X_train, X_test = X[train_idx], X[test_idx]
y_train_bin, y_test_bin = y_bin[train_idx], y_bin[test_idx]
y_train_weight, y_test_weight = y_weight[train_idx], y_weight[test_idx]
n_passes_test = n_passes[test_idx]

base_rf = RandomForestClassifier(
    n_estimators=400,
    min_samples_leaf=5,
    random_state=RANDOM_SEED,
    n_jobs=1,
    class_weight="balanced_subsample",
)
uncal_model = OneVsRestClassifier(base_rf, n_jobs=1)
cal_model = OneVsRestClassifier(
    CalibratedClassifierCV(
        estimator=RandomForestClassifier(
            n_estimators=300,
            min_samples_leaf=5,
            random_state=RANDOM_SEED,
            n_jobs=1,
            class_weight="balanced_subsample",
        ),
        method="sigmoid",
        cv=3,
    ),
    n_jobs=1,
)

print("Entrenando modelo original...", flush=True)
uncal_model.fit(X_train, y_train_bin)
uncal_probs = probs_from_ovr(uncal_model, X_test)
uncal_metrics = evaluate(y_test_bin, y_test_weight, uncal_probs, n_passes_test)

print("Entrenando modelo calibrado...", flush=True)
cal_model.fit(X_train, y_train_bin)
cal_probs = probs_from_ovr(cal_model, X_test)
cal_metrics = evaluate(y_test_bin, y_test_weight, cal_probs, n_passes_test)

lines = []
lines.append("# Comparacion modelo calibrado")
lines.append("")
lines.append(f"- Train sequences: {len(train_idx):,}")
lines.append(f"- Test sequences: {len(test_idx):,}")
lines.append("")
lines.append("## Holdout test")
for name, metrics in [("Original", uncal_metrics), ("Calibrado", cal_metrics)]:
    lines.append(
        f"- {name}: "
        f"AUC macro={metrics['auc_macro']:.3f} | "
        f"AUC micro={metrics['auc_micro']:.3f} | "
        f"MAE={metrics['mae']:.3f} | "
        f"Top-1 hit={metrics['top1_hit']:.3f} | "
        f"Multi Top-1={metrics['multi_top1_hit']:.3f} | "
        f"Multi MAE={metrics['multi_mae']:.3f}"
    )
    lines.append(
        f"  Variante dual: overlap={metrics['overlap_hit']:.3f} | "
        f"exact={metrics['exact_hit']:.3f} | "
        f"% multi pred={100*metrics['predicted_multi_rate']:.1f}% | "
        f"overlap en multi pred={metrics['predicted_multi_overlap']:.3f}"
    )

out = REPORTS_DIR / "sequence_block_model_calibrated_report.md"
out.write_text("\n".join(lines), encoding="utf-8")
print(f"Reporte: {out}")
print("\n".join(lines))