from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import html
import json
import random
import sys

import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt


SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent
MODEL_DIR = PROJECT_ROOT / "data" / "modeling"
REPORTS_DIR = PROJECT_ROOT / "reports"
REPORT_PATH = REPORTS_DIR / "pymc_logistic_normal_report.html"
JSON_PATH = MODEL_DIR / "pymc_logistic_normal_metrics.json"
MODEL_PATH = MODEL_DIR / "pymc_logistic_normal_bundle.pt"
ATTACK_PRED_PATH = MODEL_DIR / "attack_pymc_logistic_normal_test_predictions.parquet"
PV_PRED_PATH = MODEL_DIR / "pv_pymc_logistic_normal_test_predictions.parquet"

if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

import experiment_attack_distribution_gnn as attack_mod  # noqa: E402
import experiment_pv_distribution_gnn as pv_mod  # noqa: E402
import train_attack_prediction_ffn as base  # noqa: E402


RANDOM_SEED = 42
ADVI_STEPS = 3500
POSTERIOR_DRAWS = 200

ZONE_ORDER = attack_mod.ZONE_ORDER
PRESS_CELLS = [f"x{x}_y{y}" for x in range(4) for y in range(3)]


@dataclass
class SplitData:
    x: np.ndarray
    y: np.ndarray
    prior_long: np.ndarray
    prior_short: np.ndarray
    metadata: pd.DataFrame


def _set_seed(seed: int = RANDOM_SEED) -> None:
    random.seed(seed)
    np.random.seed(seed)


def _smooth_simplex(y: np.ndarray, eps: float = 1e-4) -> np.ndarray:
    k = y.shape[1]
    out = (y + eps) / (1.0 + eps * k)
    return base._normalize_rows(out).astype(np.float32)


def _alr(x: np.ndarray) -> np.ndarray:
    x = np.clip(x, 1e-6, None)
    ref = x[:, [-1]]
    return np.log(x[:, :-1] / ref).astype(np.float32)


def _inverse_alr(z: np.ndarray) -> np.ndarray:
    exp_z = np.exp(z)
    denom = 1.0 + exp_z.sum(axis=1, keepdims=True)
    body = exp_z / denom
    last = 1.0 / denom
    return np.concatenate([body, last], axis=1).astype(np.float32)


def _make_split(
    df: pd.DataFrame,
    idx: pd.Index | np.ndarray,
    x_scaled: pd.DataFrame,
    y: np.ndarray,
    prior_long: np.ndarray,
    prior_short: np.ndarray,
) -> SplitData:
    arr_idx = np.asarray(idx)
    return SplitData(
        x=x_scaled.iloc[arr_idx].to_numpy(dtype=np.float32),
        y=y[arr_idx].astype(np.float32),
        prior_long=prior_long[arr_idx].astype(np.float32),
        prior_short=prior_short[arr_idx].astype(np.float32),
        metadata=df.iloc[arr_idx].copy().reset_index(drop=True),
    )


def _select_feature_columns(df: pd.DataFrame, task: str) -> list[str]:
    cols: list[str] = []
    zone_prefixes = [
        "long_mean__actual_attack_share__",
        "short_mean__actual_attack_share__",
        "opp__long_mean__actual_conceded_attack_share__",
        "opp__short_mean__actual_conceded_attack_share__",
        "long_mean__actual_pv_share__",
        "short_mean__actual_pv_share__",
        "opp__long_mean__actual_conceded_pv_share__",
        "opp__short_mean__actual_conceded_pv_share__",
    ]
    for prefix in zone_prefixes:
        cols.extend(sorted(c for c in df.columns if c.startswith(prefix)))

    press_prefixes = [
        "short_mean__press_cnt__",
        "long_mean__press_cnt__",
        "opp__short_mean__press_cnt__",
        "opp__long_mean__press_cnt__",
        "short_mean__mf__press_cnt__",
        "long_mean__mf__press_cnt__",
        "opp__short_mean__mf__press_cnt__",
        "opp__long_mean__mf__press_cnt__",
    ]
    for prefix in press_prefixes:
        cols.extend(sorted(c for c in df.columns if c.startswith(prefix)))

    scalar_candidates = [
        "short_mean__recoveries_total",
        "long_mean__recoveries_total",
        "opp__short_mean__recoveries_total",
        "opp__long_mean__recoveries_total",
        "short_mean__shots_total",
        "long_mean__shots_total",
        "short_mean__passes_long",
        "long_mean__passes_long",
        "opp__short_mean__passes_long",
        "opp__long_mean__passes_long",
        "short_mean__goals_for",
        "long_mean__goals_for",
        "short_mean__goals_against",
        "long_mean__goals_against",
        "opp__short_mean__goals_for",
        "opp__long_mean__goals_for",
        "opp__short_mean__goals_against",
        "opp__long_mean__goals_against",
        "short_mean__mf__shots",
        "long_mean__mf__shots",
        "opp__short_mean__mf__shots",
        "opp__long_mean__mf__shots",
        "short_mean__mf__shots_on_target",
        "long_mean__mf__shots_on_target",
        "opp__short_mean__mf__shots_on_target",
        "opp__long_mean__mf__shots_on_target",
    ]
    cols.extend(c for c in scalar_candidates if c in df.columns)

    if task == "pv":
        pv_extra = [c for c in df.columns if c.startswith("short_mean__zone_pvAdded_shots__") or c.startswith("long_mean__zone_pvAdded_shots__")]
        opp_pv_extra = [c for c in df.columns if c.startswith("opp__short_mean__zone_pvAdded_shots__") or c.startswith("opp__long_mean__zone_pvAdded_shots__")]
        cols.extend(sorted(pv_extra))
        cols.extend(sorted(opp_pv_extra))

    return sorted(dict.fromkeys(cols))


def _load_task(task: str) -> tuple[SplitData, SplitData, SplitData, dict[str, object]]:
    df = base._load_dataset()
    df = df[df["usable_for_model"]].copy().reset_index(drop=True)

    if task == "attack":
        target_cols = [f"target_attack_share__{zone}" for zone in ZONE_ORDER]
        y = _smooth_simplex(df[target_cols].to_numpy(dtype=np.float32))
        prior_long = base._normalize_rows(
            df[[f"long_mean__actual_attack_share__{zone}" for zone in ZONE_ORDER]].to_numpy(dtype=float)
        ).astype(np.float32)
        prior_short = base._normalize_rows(
            df[[f"short_mean__actual_attack_share__{zone}" for zone in ZONE_ORDER]].to_numpy(dtype=float)
        ).astype(np.float32)
    elif task == "pv":
        y_dist, prior_long_raw, prior_short_raw = pv_mod._build_distributions(df)
        keep_mask = prior_long_raw.sum(axis=1) > 0
        df = df.loc[keep_mask].copy().reset_index(drop=True)
        y_dist, prior_long_raw, prior_short_raw = pv_mod._build_distributions(df)
        y = _smooth_simplex(y_dist.astype(np.float32))
        prior_long = base._normalize_rows(np.clip(prior_long_raw, 0.0, None)).astype(np.float32)
        prior_short = base._normalize_rows(np.clip(prior_short_raw, 0.0, None)).astype(np.float32)
        target_cols = [f"target_pv_dist__{zone}" for zone in ZONE_ORDER]
    else:
        raise ValueError(task)

    feature_cols = _select_feature_columns(df, task)
    x = df[feature_cols].fillna(0.0)
    train_df = df[df["split"] == "train"].copy()
    train_idx, val_idx, val_start_date = base._build_temporal_validation(train_df)
    test_idx = df.index[df["split"] == "test"]
    x_scaled, scaler_bundle = base._standardize_features(x, train_idx)

    info = {
        "target_cols": target_cols,
        "feature_cols": feature_cols,
        "val_start_date": val_start_date,
        "scaler": scaler_bundle,
    }
    return (
        _make_split(df, train_idx, x_scaled, y, prior_long, prior_short),
        _make_split(df, val_idx, x_scaled, y, prior_long, prior_short),
        _make_split(df, test_idx, x_scaled, y, prior_long, prior_short),
        info,
    )


def _fit_model(train: SplitData) -> tuple[pm.Model, object]:
    x_train = train.x.astype(np.float32)
    y_train = _alr(train.y)
    prior_train = _alr(train.prior_long)
    n_features = x_train.shape[1]
    n_outputs = y_train.shape[1]

    with pm.Model() as model:
        X = pm.Data("X", x_train)
        prior_alr = pm.Data("prior_alr", prior_train)
        beta = pm.Normal("beta", mu=0.0, sigma=0.12, shape=(n_features, n_outputs))
        bias = pm.Normal("bias", mu=0.0, sigma=0.08, shape=(n_outputs,))
        sigma = pm.HalfNormal("sigma", sigma=0.20, shape=(n_outputs,))
        mu = prior_alr + pt.dot(X, beta) + bias
        pm.Normal("obs", mu=mu, sigma=sigma, observed=y_train)
        approx = pm.fit(
            n=ADVI_STEPS,
            method="advi",
            obj_optimizer=pm.adam(learning_rate=0.01),
            progressbar=False,
            random_seed=RANDOM_SEED,
        )
    return model, approx


def _predict(model: pm.Model, approx: object, split: SplitData) -> np.ndarray:
    x_arr = split.x.astype(np.float32)
    prior_arr = _alr(split.prior_long)
    with model:
        pm.set_data({"X": x_arr, "prior_alr": prior_arr})
        posterior = approx.sample(draws=POSTERIOR_DRAWS, return_inferencedata=True, random_seed=RANDOM_SEED)
    beta = posterior.posterior["beta"].mean(dim=("chain", "draw")).values
    bias = posterior.posterior["bias"].mean(dim=("chain", "draw")).values
    mu = prior_arr + x_arr @ beta + bias
    pred = _inverse_alr(mu)
    return base._normalize_rows(pred).astype(np.float32)


def _metrics(y_true: np.ndarray, pred_model: np.ndarray, pred_long: np.ndarray, pred_short: np.ndarray) -> dict[str, float]:
    return {
        "model_mae": base._mean_abs_error(y_true, pred_model),
        "season_baseline_mae": base._mean_abs_error(y_true, pred_long),
        "short8_baseline_mae": base._mean_abs_error(y_true, pred_short),
        "model_jsd": base._jsd_mean(y_true, pred_model),
        "season_baseline_jsd": base._jsd_mean(y_true, pred_long),
        "short8_baseline_jsd": base._jsd_mean(y_true, pred_short),
        "model_kl_proxy": float(np.mean(np.sum(y_true * (np.log(y_true + 1e-12) - np.log(pred_model + 1e-12)), axis=1))),
        "season_baseline_kl_proxy": float(np.mean(np.sum(y_true * (np.log(y_true + 1e-12) - np.log(pred_long + 1e-12)), axis=1))),
        "short8_baseline_kl_proxy": float(np.mean(np.sum(y_true * (np.log(y_true + 1e-12) - np.log(pred_short + 1e-12)), axis=1))),
    }


def _mae_rows(df: pd.DataFrame, target_prefix: str) -> pd.DataFrame:
    target_cols = [c for c in df.columns if c.startswith(target_prefix)]
    model_cols = [c.replace(target_prefix, f"pred_model__{target_prefix}") for c in target_cols]
    season_cols = [c.replace(target_prefix, f"pred_season__{target_prefix}") for c in target_cols]
    short8_cols = [c.replace(target_prefix, f"pred_short8__{target_prefix}") for c in target_cols]
    out = df[["fecha", "team_name", "opponent_name"]].copy()
    out["model_mae"] = np.abs(df[target_cols].to_numpy() - df[model_cols].to_numpy()).mean(axis=1)
    out["season_mae"] = np.abs(df[target_cols].to_numpy() - df[season_cols].to_numpy()).mean(axis=1)
    out["short8_mae"] = np.abs(df[target_cols].to_numpy() - df[short8_cols].to_numpy()).mean(axis=1)
    return out


def _last3_rows(df: pd.DataFrame) -> list[dict[str, object]]:
    sub = df[df["team_name"].eq("Racing de Santander")].sort_values("fecha").tail(3)
    rows = sub.to_dict(orient="records")
    for row in rows:
        row["fecha"] = pd.Timestamp(row["fecha"]).strftime("%Y-%m-%d")
    return rows


def _report_html(summary: dict[str, object]) -> str:
    def rows_html(rows: list[dict[str, object]]) -> str:
        return "".join(
            f"<tr><td>{html.escape(str(r['fecha']))}</td><td>{html.escape(str(r['opponent_name']))}</td>"
            f"<td>{r['model_mae']:.4f}</td><td>{r['season_mae']:.4f}</td><td>{r['short8_mae']:.4f}</td><td>{r['previous_model_mae']:.4f}</td></tr>"
            for r in rows
        )

    return f"""<!DOCTYPE html>
<html lang="es">
<head>
  <meta charset="utf-8" />
  <title>PyMC Logistic-Normal</title>
  <style>
    body {{ margin: 0; background: #f1f4ef; color: #14342B; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }}
    .wrap {{ max-width: 1200px; margin: 0 auto; padding: 30px 24px 48px; }}
    h1 {{ margin: 0 0 10px; font-size: 38px; }}
    .lead {{ margin: 0 0 22px; font-size: 18px; color: #35574D; }}
    .grid2 {{ display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 16px; margin-bottom: 18px; }}
    .card {{ background: white; border-radius: 20px; padding: 18px 20px; box-shadow: 0 8px 24px rgba(12, 36, 28, 0.08); }}
    table {{ width: 100%; border-collapse: collapse; font-size: 14px; }}
    th, td {{ padding: 10px 8px; border-bottom: 1px solid #E5ECE6; text-align: left; }}
    th {{ color: #587468; text-transform: uppercase; font-size: 12px; letter-spacing: .06em; }}
    @media (max-width: 980px) {{ .grid2 {{ grid-template-columns: 1fr; }} }}
  </style>
</head>
<body>
  <div class="wrap">
    <h1>Bayesiano logistic-normal con PyMC</h1>
    <p class="lead">Regresion bayesiana en espacio log-ratio, con prior del partido basado en temporada y correccion por features de ataque, concesion rival, pases, presiones y recuperaciones.</p>
    <div class="grid2">
      <section class="card">
        <h2>Ataque - test</h2>
        <table>
          <tr><th>Metrica</th><th>PyMC</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          <tr><td>MAE</td><td>{summary['attack']['test_metrics']['model_mae']:.4f}</td><td>{summary['attack']['test_metrics']['season_baseline_mae']:.4f}</td><td>{summary['attack']['test_metrics']['short8_baseline_mae']:.4f}</td><td>{summary['attack']['previous_test_metrics']['model_mae']:.4f}</td></tr>
          <tr><td>JSD</td><td>{summary['attack']['test_metrics']['model_jsd']:.4f}</td><td>{summary['attack']['test_metrics']['season_baseline_jsd']:.4f}</td><td>{summary['attack']['test_metrics']['short8_baseline_jsd']:.4f}</td><td>{summary['attack']['previous_test_metrics']['model_jsd']:.4f}</td></tr>
          <tr><td>KL</td><td>{summary['attack']['test_metrics']['model_kl_proxy']:.4f}</td><td>{summary['attack']['test_metrics']['season_baseline_kl_proxy']:.4f}</td><td>{summary['attack']['test_metrics']['short8_baseline_kl_proxy']:.4f}</td><td>{summary['attack']['previous_test_metrics']['model_kl_proxy']:.4f}</td></tr>
        </table>
      </section>
      <section class="card">
        <h2>PV - test</h2>
        <table>
          <tr><th>Metrica</th><th>PyMC</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          <tr><td>MAE</td><td>{summary['pv']['test_metrics']['model_mae']:.4f}</td><td>{summary['pv']['test_metrics']['season_baseline_mae']:.4f}</td><td>{summary['pv']['test_metrics']['short8_baseline_mae']:.4f}</td><td>{summary['pv']['previous_test_metrics']['model_mae']:.4f}</td></tr>
          <tr><td>JSD</td><td>{summary['pv']['test_metrics']['model_jsd']:.4f}</td><td>{summary['pv']['test_metrics']['season_baseline_jsd']:.4f}</td><td>{summary['pv']['test_metrics']['short8_baseline_jsd']:.4f}</td><td>{summary['pv']['previous_test_metrics']['model_jsd']:.4f}</td></tr>
          <tr><td>KL</td><td>{summary['pv']['test_metrics']['model_kl_proxy']:.4f}</td><td>{summary['pv']['test_metrics']['season_baseline_kl_proxy']:.4f}</td><td>{summary['pv']['test_metrics']['short8_baseline_kl_proxy']:.4f}</td><td>{summary['pv']['previous_test_metrics']['model_kl_proxy']:.4f}</td></tr>
        </table>
      </section>
    </div>
    <div class="grid2">
      <section class="card">
        <h2>Ataque - ultimos 3 de Racing</h2>
        <table>
          <tr><th>Fecha</th><th>Rival</th><th>PyMC</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          {rows_html(summary['attack']['last3_racing'])}
        </table>
      </section>
      <section class="card">
        <h2>PV - ultimos 3 de Racing</h2>
        <table>
          <tr><th>Fecha</th><th>Rival</th><th>PyMC</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          {rows_html(summary['pv']['last3_racing'])}
        </table>
      </section>
    </div>
  </div>
</body>
</html>"""


def _run_task(task: str, pred_path: Path, prev_metrics_path: Path, prev_pred_path: Path) -> dict[str, object]:
    train, val, test, info = _load_task(task)
    model, approx = _fit_model(train)
    val_pred = _predict(model, approx, val)
    test_pred = _predict(model, approx, test)

    val_metrics = _metrics(val.y, val_pred, val.prior_long, val.prior_short)
    test_metrics = _metrics(test.y, test_pred, test.prior_long, test.prior_short)

    pred_df = test.metadata.copy().reset_index(drop=True)
    target_prefix = "target_attack_share__" if task == "attack" else "target_pv_dist__"
    for i, zone in enumerate(ZONE_ORDER):
        pred_df[f"{target_prefix}{zone}"] = test.y[:, i]
        pred_df[f"pred_model__{target_prefix}{zone}"] = test_pred[:, i]
        pred_df[f"pred_season__{target_prefix}{zone}"] = test.prior_long[:, i]
        pred_df[f"pred_short8__{target_prefix}{zone}"] = test.prior_short[:, i]
    pred_df.to_parquet(pred_path, index=False)

    prev_metrics = json.loads(prev_metrics_path.read_text(encoding="utf-8"))["test_metrics"]
    current_rows = _mae_rows(pred_df, target_prefix)
    prev_rows = _mae_rows(pd.read_parquet(prev_pred_path), target_prefix)
    current_rows = current_rows.merge(
        prev_rows.rename(columns={"model_mae": "previous_model_mae", "season_mae": "previous_season_mae", "short8_mae": "previous_short8_mae"}),
        on=["fecha", "team_name", "opponent_name"],
        how="left",
    )

    return {
        "train_rows": int(len(train.metadata)),
        "val_rows": int(len(val.metadata)),
        "test_rows": int(len(test.metadata)),
        "val_start_date": info["val_start_date"],
        "feature_count": len(info["feature_cols"]),
        "test_metrics": test_metrics,
        "val_metrics": val_metrics,
        "previous_test_metrics": prev_metrics,
        "last3_racing": _last3_rows(current_rows),
    }


def main() -> None:
    _set_seed()
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    REPORTS_DIR.mkdir(parents=True, exist_ok=True)

    attack_summary = _run_task(
        "attack",
        ATTACK_PRED_PATH,
        MODEL_DIR / "attack_distribution_gnn_metrics.json",
        MODEL_DIR / "attack_distribution_gnn_test_predictions.parquet",
    )
    pv_summary = _run_task(
        "pv",
        PV_PRED_PATH,
        MODEL_DIR / "pv_distribution_gnn_metrics.json",
        MODEL_DIR / "pv_distribution_gnn_test_predictions.parquet",
    )

    summary = {
        "model": "pymc_logistic_normal",
        "attack": attack_summary,
        "pv": pv_summary,
        "config": {
            "advi_steps": ADVI_STEPS,
            "posterior_draws": POSTERIOR_DRAWS,
        },
        "artifacts": {
            "attack_predictions": str(ATTACK_PRED_PATH),
            "pv_predictions": str(PV_PRED_PATH),
        },
    }

    JSON_PATH.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    REPORT_PATH.write_text(_report_html(summary), encoding="utf-8")
    MODEL_PATH.write_text(json.dumps({"model": "pymc_logistic_normal", "summary_path": str(JSON_PATH)}, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"Metricas guardadas en: {JSON_PATH}")
    print(f"Reporte guardado en: {REPORT_PATH}")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()