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 pyro
import pyro.distributions as dist
from pyro.infer import Predictive, SVI, Trace_ELBO
from pyro.infer.autoguide import AutoDiagonalNormal
from pyro.optim import ClippedAdam
import torch
from torch.utils.data import DataLoader, TensorDataset


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 / "pyro_logistic_normal_report.html"
JSON_PATH = MODEL_DIR / "pyro_logistic_normal_metrics.json"
MODEL_PATH = MODEL_DIR / "pyro_logistic_normal_bundle.pt"
ATTACK_PRED_PATH = MODEL_DIR / "attack_pyro_logistic_normal_test_predictions.parquet"
PV_PRED_PATH = MODEL_DIR / "pv_pyro_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
BATCH_SIZE = 512
TRAIN_STEPS = 2200
LEARNING_RATE = 0.01
POSTERIOR_SAMPLES = 40


@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)
    torch.manual_seed(seed)
    pyro.set_rng_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 _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 _loader(split: SplitData, shuffle: bool) -> DataLoader:
    ds = TensorDataset(
        torch.from_numpy(split.x),
        torch.from_numpy(split.y),
        torch.from_numpy(split.prior_long),
        torch.from_numpy(split.prior_short),
    )
    return DataLoader(ds, batch_size=BATCH_SIZE, shuffle=shuffle)


def _load_attack_task() -> tuple[SplitData, SplitData, SplitData, list[str], str]:
    df = base._load_dataset()
    df = df[df["usable_for_model"]].copy().reset_index(drop=True)
    features, _numeric_cols, _attack_targets, _pv_targets, _dummy_cols = base._feature_matrix(df)
    attack_targets = [f"target_attack_share__{zone}" for zone in attack_mod.ZONE_ORDER]
    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(features, train_idx)
    y = _smooth_simplex(df[attack_targets].to_numpy(dtype=np.float32))
    prior_long = base._normalize_rows(df[[f"long_mean__actual_attack_share__{zone}" for zone in attack_mod.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 attack_mod.ZONE_ORDER]].to_numpy(dtype=float)).astype(np.float32)
    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),
        attack_targets,
        val_start_date,
    )


def _load_pv_task() -> tuple[SplitData, SplitData, SplitData, str]:
    df = base._load_dataset()
    df = df[df["usable_for_model"]].copy().reset_index(drop=True)
    y_dist, prior_long, prior_short = pv_mod._build_distributions(df)
    keep_mask = prior_long.sum(axis=1) > 0
    df = df.loc[keep_mask].copy().reset_index(drop=True)
    y_dist, prior_long, prior_short = pv_mod._build_distributions(df)
    features, _numeric_cols, _attack_targets, _pv_targets, _dummy_cols = base._feature_matrix(df)
    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(features, train_idx)
    y = _smooth_simplex(y_dist.astype(np.float32))
    prior_long = base._normalize_rows(np.clip(prior_long, 0.0, None)).astype(np.float32)
    prior_short = base._normalize_rows(np.clip(prior_short, 0.0, None)).astype(np.float32)
    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),
        val_start_date,
    )


def _pyro_model(x: torch.Tensor, prior: torch.Tensor, y: torch.Tensor | None = None):
    n_features = x.shape[1]
    n_outputs = prior.shape[1]
    weight = pyro.sample("weight", dist.Normal(0.0, 0.15).expand([n_outputs, n_features]).to_event(2))
    bias = pyro.sample("bias", dist.Normal(0.0, 0.10).expand([n_outputs]).to_event(1))
    log_concentration = pyro.sample("log_concentration", dist.Normal(np.log(30.0), 0.4))
    logits = torch.log(torch.clamp(prior, min=1e-6)) + x.matmul(weight.T) + bias
    mean = torch.softmax(logits, dim=1)
    pyro.deterministic("mean", mean)
    concentration = torch.exp(log_concentration)
    alpha = mean * concentration + 1e-4
    with pyro.plate("data", x.shape[0]):
        pyro.sample("obs", dist.Dirichlet(alpha), obs=y)


def _train_pyro(train: SplitData, val: SplitData) -> tuple[AutoDiagonalNormal, list[dict[str, float]]]:
    pyro.clear_param_store()
    guide = AutoDiagonalNormal(_pyro_model)
    optimizer = ClippedAdam({"lr": LEARNING_RATE, "clip_norm": 10.0})
    svi = SVI(_pyro_model, guide, optimizer, loss=Trace_ELBO())
    history: list[dict[str, float]] = []
    loader = _loader(train, shuffle=True)

    best_val = float("inf")
    best_params = None
    patience = 220
    patience_left = patience

    for step in range(1, TRAIN_STEPS + 1):
        running = 0.0
        n_batches = 0
        for x, y, prior_long, _prior_short in loader:
            loss = svi.step(x, prior_long, y)
            running += float(loss) / len(x)
            n_batches += 1

        if step % 25 == 0:
            pred_val = _predict_pyro(guide, val)
            val_mae = base._mean_abs_error(val.y, pred_val)
            history.append({"step": step, "train_loss": running / max(n_batches, 1), "val_mae": val_mae})
            if val_mae < best_val - 1e-6:
                best_val = val_mae
                best_params = {k: v.detach().cpu().clone() for k, v in pyro.get_param_store().items()}
                patience_left = patience
            else:
                patience_left -= 25
            if patience_left <= 0:
                break

    if best_params is not None:
        pyro.clear_param_store()
        for k, v in best_params.items():
            pyro.get_param_store()[k] = v.clone()
    return guide, history


def _predict_pyro(guide: AutoDiagonalNormal, split: SplitData) -> np.ndarray:
    outputs = []
    loader = _loader(split, shuffle=False)
    for x, _y, prior_long, _prior_short in loader:
        posterior_median = guide.median(x, prior_long, None)
        conditioned = pyro.poutine.condition(_pyro_model, data=posterior_median)
        trace = pyro.poutine.trace(conditioned).get_trace(x, prior_long, None)
        mean_np = trace.nodes["mean"]["value"].detach().cpu().numpy()
        outputs.append(mean_np)
    return np.vstack(outputs)


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:
    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>Pyro 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 Pyro</h1>
    <p class="lead">Modelo composicional bayesiano: prior del partido en logit-space + corrección por features con pesos bajo prior gaussiano, inferidos por SVI.</p>
    <div class="grid2">
      <section class="card">
        <h2>Ataque - test</h2>
        <table>
          <tr><th>Metrica</th><th>Pyro</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>Pyro</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>Pyro</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>Pyro</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 main() -> None:
    _set_seed()
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    REPORTS_DIR.mkdir(parents=True, exist_ok=True)

    attack_train, attack_val, attack_test, attack_targets, attack_val_start = _load_attack_task()
    attack_guide, attack_history = _train_pyro(attack_train, attack_val)
    attack_val_pred = _predict_pyro(attack_guide, attack_val)
    attack_test_pred = _predict_pyro(attack_guide, attack_test)
    attack_val_metrics = _metrics(attack_val.y, attack_val_pred, attack_val.prior_long, attack_val.prior_short)
    attack_test_metrics = _metrics(attack_test.y, attack_test_pred, attack_test.prior_long, attack_test.prior_short)

    attack_pred_df = attack_test.metadata.copy().reset_index(drop=True)
    for i, zone in enumerate(attack_mod.ZONE_ORDER):
        attack_pred_df[f"target_attack_share__{zone}"] = attack_test.y[:, i]
        attack_pred_df[f"pred_model__target_attack_share__{zone}"] = attack_test_pred[:, i]
        attack_pred_df[f"pred_season__target_attack_share__{zone}"] = attack_test.prior_long[:, i]
        attack_pred_df[f"pred_short8__target_attack_share__{zone}"] = attack_test.prior_short[:, i]
    attack_pred_df.to_parquet(ATTACK_PRED_PATH, index=False)

    pv_train, pv_val, pv_test, pv_val_start = _load_pv_task()
    pv_guide, pv_history = _train_pyro(pv_train, pv_val)
    pv_val_pred = _predict_pyro(pv_guide, pv_val)
    pv_test_pred = _predict_pyro(pv_guide, pv_test)
    pv_val_metrics = _metrics(pv_val.y, pv_val_pred, pv_val.prior_long, pv_val.prior_short)
    pv_test_metrics = _metrics(pv_test.y, pv_test_pred, pv_test.prior_long, pv_test.prior_short)

    pv_pred_df = pv_test.metadata.copy().reset_index(drop=True)
    for i, zone in enumerate(pv_mod.ZONE_ORDER):
        pv_pred_df[f"target_pv_dist__{zone}"] = pv_test.y[:, i]
        pv_pred_df[f"pred_model__target_pv_dist__{zone}"] = pv_test_pred[:, i]
        pv_pred_df[f"pred_season__target_pv_dist__{zone}"] = pv_test.prior_long[:, i]
        pv_pred_df[f"pred_short8__target_pv_dist__{zone}"] = pv_test.prior_short[:, i]
    pv_pred_df.to_parquet(PV_PRED_PATH, index=False)

    attack_prev_metrics = json.loads((MODEL_DIR / "attack_distribution_gnn_metrics.json").read_text(encoding="utf-8"))["test_metrics"]
    pv_prev_metrics = json.loads((MODEL_DIR / "pv_distribution_gnn_metrics.json").read_text(encoding="utf-8"))["test_metrics"]

    attack_last3 = _mae_rows(attack_pred_df, "target_attack_share__")
    attack_prev_last3 = _mae_rows(pd.read_parquet(MODEL_DIR / "attack_distribution_gnn_test_predictions.parquet"), "target_attack_share__")
    attack_last3 = attack_last3.merge(
        attack_prev_last3.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",
    )
    pv_last3 = _mae_rows(pv_pred_df, "target_pv_dist__")
    pv_prev_last3 = _mae_rows(pd.read_parquet(MODEL_DIR / "pv_distribution_gnn_test_predictions.parquet"), "target_pv_dist__")
    pv_last3 = pv_last3.merge(
        pv_prev_last3.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",
    )

    summary = {
        "model": "pyro_logistic_normal",
        "attack": {
            "train_rows": int(len(attack_train.metadata)),
            "val_rows": int(len(attack_val.metadata)),
            "test_rows": int(len(attack_test.metadata)),
            "val_start_date": attack_val_start,
            "test_metrics": attack_test_metrics,
            "val_metrics": attack_val_metrics,
            "previous_test_metrics": attack_prev_metrics,
            "last3_racing": _last3_rows(attack_last3),
        },
        "pv": {
            "train_rows": int(len(pv_train.metadata)),
            "val_rows": int(len(pv_val.metadata)),
            "test_rows": int(len(pv_test.metadata)),
            "val_start_date": pv_val_start,
            "test_metrics": pv_test_metrics,
            "val_metrics": pv_val_metrics,
            "previous_test_metrics": pv_prev_metrics,
            "last3_racing": _last3_rows(pv_last3),
        },
        "config": {
            "train_steps": TRAIN_STEPS,
            "learning_rate": LEARNING_RATE,
            "batch_size": BATCH_SIZE,
            "posterior_samples": POSTERIOR_SAMPLES,
        },
        "artifacts": {
            "attack_predictions": str(ATTACK_PRED_PATH),
            "pv_predictions": str(PV_PRED_PATH),
        },
    }

    torch.save({"summary": summary, "pyro_param_store": {k: v.detach().cpu() for k, v in pyro.get_param_store().items()}}, MODEL_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")

    print(f"Modelo guardado en: {MODEL_PATH}")
    print(f"Predicciones ataque guardadas en: {ATTACK_PRED_PATH}")
    print(f"Predicciones PV guardadas en: {PV_PRED_PATH}")
    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()