backend/model_engine.py

"""V3RSUS prediction engine — v2.

Loads:
  - data/model_v2.pkl              ensemble + scaler + imputer + feature names
  - data-pipeline/features/fighter_snapshots.parquet   per-fighter ML features
  - data-pipeline/processed/fights.parquet             historical matchup lookup
  - data/ufc-master.csv                                career-aggregate stats for display

For arbitrary matchups: pull snapshots, compute diff vector, scale, predict.
Falls back gracefully if v2 artifact is missing (re-trains from scratch).
"""
import os
from pathlib import Path

import joblib
import numpy as np
import pandas as pd


EXCLUDED_FROM_DIFF = frozenset()

FEATURE_DISPLAY = {
    "rating": "Glicko Rating",
    "rd": "Rating Uncertainty",
    "height": "Height",
    "reach": "Reach",
    "weight": "Weight",
    "age": "Age",
    "layoff": "Layoff (days)",
    "career_fights": "Career Fights",
    "career_wins": "Career Wins",
    "career_losses": "Career Losses",
    "win_streak": "Win Streak",
    "lose_streak": "Lose Streak",
    "sig_str_landed_last3": "Sig. Strikes (last 3)",
    "sig_str_attempted_last3": "Sig. Strikes Attempted (last 3)",
    "td_landed_last3": "Takedowns (last 3)",
    "td_attempted_last3": "Takedowns Attempted (last 3)",
    "head_landed_last3": "Head Strikes (last 3)",
    "body_landed_last3": "Body Strikes (last 3)",
    "leg_landed_last3": "Leg Strikes (last 3)",
    "distance_landed_last3": "Distance Strikes (last 3)",
    "clinch_landed_last3": "Clinch Strikes (last 3)",
    "ground_landed_last3": "Ground Strikes (last 3)",
    "sub_att_last3": "Submission Attempts (last 3)",
    "kd_last3": "Knockdowns (last 3)",
    "ctrl_seconds_last3": "Control Time (last 3)",
    "sig_str_landed_last5": "Sig. Strikes (last 5)",
    "sig_str_attempted_last5": "Sig. Strikes Attempted (last 5)",
    "td_landed_last5": "Takedowns (last 5)",
    "td_attempted_last5": "Takedowns Attempted (last 5)",
    "head_landed_last5": "Head Strikes (last 5)",
    "body_landed_last5": "Body Strikes (last 5)",
    "leg_landed_last5": "Leg Strikes (last 5)",
    "distance_landed_last5": "Distance Strikes (last 5)",
    "clinch_landed_last5": "Clinch Strikes (last 5)",
    "ground_landed_last5": "Ground Strikes (last 5)",
    "sub_att_last5": "Submission Attempts (last 5)",
    "kd_last5": "Knockdowns (last 5)",
    "ctrl_seconds_last5": "Control Time (last 5)",
}

FEATURE_CATEGORY = {
    "rating": "experience", "rd": "experience",
    "height": "physical", "reach": "physical", "weight": "physical", "age": "physical",
    "layoff": "experience",
    "career_fights": "experience", "career_wins": "experience", "career_losses": "experience",
    "win_streak": "experience", "lose_streak": "experience",
}
for base in ("sig_str_landed", "sig_str_attempted", "head_landed", "body_landed", "leg_landed",
             "distance_landed", "clinch_landed", "kd"):
    for n in (3, 5):
        FEATURE_CATEGORY[f"{base}_last{n}"] = "striking"
for base in ("td_landed", "td_attempted", "ground_landed", "sub_att", "ctrl_seconds"):
    for n in (3, 5):
        FEATURE_CATEGORY[f"{base}_last{n}"] = "grappling"

CATEGORY_LABELS = {
    "striking": "Striking",
    "grappling": "Grappling",
    "physical": "Physical",
    "experience": "Experience",
}


class FightPredictor:
    def __init__(self):
        self.models = {}
        self.scaler = None
        self.imputer = None
        self.features: list[str] = []
        self.model_metrics = {}
        self.lr_coefficients = None

        # Data sources
        self.snapshots: pd.DataFrame | None = None
        self.fights_df: pd.DataFrame | None = None   # ufc-master.csv format for display + fighters list
        self.norm_fights: pd.DataFrame | None = None  # normalized fights table for historical lookup

        # Stats
        self.fighter_count = 0
        self.fight_count = 0
        self.feature_count = 0

    # ------------------------------------------------------------------
    # Loading
    # ------------------------------------------------------------------

    def load_artifact(self, artifact_path: str, data_path: str) -> bool:
        artifact = joblib.load(artifact_path)
        self.models = artifact["models"]
        self.scaler = artifact["scaler"]
        self.imputer = artifact.get("imputer")
        self.features = artifact["features"]
        self.model_metrics = artifact["model_metrics"]
        self.lr_coefficients = artifact["lr_coefficients"]
        self.feature_count = len(self.features)

        self.fights_df = pd.read_csv(data_path)
        self.fight_count = len(self.fights_df)

        # data_dir = the same folder we just loaded the artifact from
        data_dir = Path(artifact_path).resolve().parent
        repo_root = Path(__file__).resolve().parents[1]
        self._load_v2_sources(repo_root, data_dir)
        return True

    def _load_v2_sources(self, repo_root: Path, data_dir: Path | None = None) -> None:
        d = data_dir or (repo_root / "data")
        snap_path = d / "fighter_snapshots.parquet"
        fights_path = d / "fights.parquet"
        events_path = d / "events.parquet"

        if snap_path.exists():
            self.snapshots = pd.read_parquet(snap_path)
            self.fighter_count = len(self.snapshots)
            print(f"  loaded {len(self.snapshots)} fighter snapshots")

        if fights_path.exists() and events_path.exists():
            f = pd.read_parquet(fights_path)
            ev = pd.read_parquet(events_path)
            self.norm_fights = f.merge(ev[["event_id", "date"]], on="event_id", how="left")
            print(f"  loaded {len(self.norm_fights)} normalized fights")

    # ------------------------------------------------------------------
    # Backwards-compat fall-back train (used if no v2 artifact)
    # ------------------------------------------------------------------

    def train(self, data_path: str = "data/ufc-master.csv") -> bool:
        """Legacy training on ufc-master.csv. Kept as fallback only."""
        from sklearn.linear_model import LogisticRegression
        from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
        from sklearn.preprocessing import StandardScaler
        from sklearn.model_selection import cross_val_score

        if not os.path.exists(data_path):
            print(f"data not found at {data_path}")
            return False

        self.fights_df = pd.read_csv(data_path)
        fights = self.fights_df.copy()
        self.fight_count = len(fights)
        fights["target"] = (fights["Winner"] == "Red").astype(int)

        diff_cols = [c for c in fights.columns if c.endswith("Dif") and pd.api.types.is_numeric_dtype(fights[c])]
        X = fights[diff_cols].fillna(0)
        y = fights["target"]
        self.features = diff_cols
        self.feature_count = len(diff_cols)

        self.scaler = StandardScaler()
        X_scaled = self.scaler.fit_transform(X)

        configs = [
            ("logistic_regression", LogisticRegression(max_iter=1000, random_state=42)),
            ("random_forest", RandomForestClassifier(n_estimators=100, max_depth=8, random_state=42, n_jobs=-1)),
            ("gradient_boosting", GradientBoostingClassifier(n_estimators=100, max_depth=4, random_state=42)),
        ]
        for name, m in configs:
            m.fit(X_scaled, y)
            self.models[name] = m
            scores = cross_val_score(m, X_scaled, y, cv=5, scoring="accuracy")
            self.model_metrics[name] = {"accuracy": round(float(scores.mean()), 4), "std": round(float(scores.std()), 4)}

        self.lr_coefficients = self.models["logistic_regression"].coef_[0]
        return True

    # ------------------------------------------------------------------
    # Lookups
    # ------------------------------------------------------------------

    def get_fighter_stats(self, name: str) -> dict | None:
        """Returns display stats (ufc-master.csv format). Used by build_profile in app.py."""
        if self.fights_df is None:
            return None
        n = name.lower().strip()
        red = self.fights_df[self.fights_df["RedFighter"].str.lower() == n]
        blue = self.fights_df[self.fights_df["BlueFighter"].str.lower() == n]
        if red.empty and blue.empty:
            return None

        latest_red = red.sort_values("Date", ascending=False).iloc[0] if not red.empty else None
        latest_blue = blue.sort_values("Date", ascending=False).iloc[0] if not blue.empty else None

        if latest_red is None:
            row, prefix = latest_blue, "Blue"
        elif latest_blue is None:
            row, prefix = latest_red, "Red"
        else:
            if latest_red["Date"] >= latest_blue["Date"]:
                row, prefix = latest_red, "Red"
            else:
                row, prefix = latest_blue, "Blue"

        stats = {}
        for col in self.fights_df.columns:
            if col.startswith(prefix):
                stats[col.replace(prefix, "", 1)] = row[col]
        stats["ActualName"] = row[prefix + "Fighter"]
        stats["TotalFights"] = len(red) + len(blue)
        return stats

    def _lookup_snapshot(self, name: str) -> pd.Series | None:
        if self.snapshots is None:
            return None
        mask = self.snapshots["name"].str.lower() == name.lower().strip()
        if not mask.any():
            return None
        return self.snapshots[mask].iloc[0]

    # ------------------------------------------------------------------
    # Prediction
    # ------------------------------------------------------------------

    @staticmethod
    def _safe(v):
        if v is None:
            return 0.0
        if isinstance(v, (int, np.integer)):
            return float(v)
        if isinstance(v, (float, np.floating)):
            return 0.0 if (np.isnan(v) or np.isinf(v)) else float(v)
        return 0.0

    def _build_diff_vector(self, s1: pd.Series, s2: pd.Series) -> np.ndarray:
        """Map each training feature to a diff(s1, s2) value."""
        # Map column suffix (e.g. "sig_str_landed_last3") to snapshot column
        # Snapshot column names match the rolling table: sig_str_landed_last3 etc.
        snap_aliases = {
            "rating": "post_rating",
            "rd": "post_rd",
            "height": "height_cm",
            "reach": "reach_cm",
            "weight": "weight_lbs",
            "age": "current_age",
            "layoff": "current_layoff_days",
            "career_fights": "career_fights_before",
            "career_wins": "career_wins_before",
            "career_losses": "career_losses_before",
            "win_streak": "win_streak_before",
            "lose_streak": "lose_streak_before",
        }

        out = np.zeros(len(self.features), dtype=float)
        for i, fname in enumerate(self.features):
            if fname.startswith("wc_"):
                continue  # weight class dummies not used at prediction (unknown until fight)
            base = fname.replace("diff_", "", 1)
            col = snap_aliases.get(base, base)
            v1 = self._safe(s1.get(col)) if col in s1.index else 0.0
            v2 = self._safe(s2.get(col)) if col in s2.index else 0.0
            out[i] = v1 - v2
        return out

    def predict_matchup(self, f1_name: str, f2_name: str) -> dict:
        s1 = self._lookup_snapshot(f1_name)
        s2 = self._lookup_snapshot(f2_name)
        if s1 is None:
            raise ValueError(f"Fighter not found in snapshot table: {f1_name}")
        if s2 is None:
            raise ValueError(f"Fighter not found in snapshot table: {f2_name}")

        diff_vec = self._build_diff_vector(s1, s2).reshape(1, -1)
        if self.imputer is not None:
            diff_vec = self.imputer.transform(diff_vec)
        scaled = self.scaler.transform(diff_vec)[0]

        model_breakdown = {}
        ensemble_prob = np.zeros(2)
        for name, model in self.models.items():
            prob = model.predict_proba(scaled.reshape(1, -1))[0]
            model_breakdown[name] = {
                "f1Prob": round(float(prob[1]), 4),
                "f2Prob": round(float(prob[0]), 4),
                "accuracy": self.model_metrics.get(name, {}).get("accuracy"),
            }
            ensemble_prob += prob
        ensemble_prob /= len(self.models)
        f1_prob = float(ensemble_prob[1])
        f2_prob = float(ensemble_prob[0])

        f1_display = self.get_fighter_stats(f1_name) or {"ActualName": s1["name"]}
        f2_display = self.get_fighter_stats(f2_name) or {"ActualName": s2["name"]}
        f1_actual = f1_display.get("ActualName", s1["name"])
        f2_actual = f2_display.get("ActualName", s2["name"])

        winner = f1_actual if f1_prob > 0.5 else f2_actual
        confidence = max(f1_prob, f2_prob)
        model_breakdown["ensemble"] = {"f1Prob": round(f1_prob, 4), "f2Prob": round(f2_prob, 4)}

        return {
            "winner": winner,
            "confidence": round(confidence, 4),
            "f1Prob": round(f1_prob, 4),
            "f2Prob": round(f2_prob, 4),
            "f1Name": f1_actual,
            "f2Name": f2_actual,
            "f1Stats": f1_display,
            "f2Stats": f2_display,
            "keyFactors": self._key_factors(scaled, f1_actual, f2_actual, s1, s2),
            "categoryAnalysis": self._category_scores(scaled),
            "modelBreakdown": model_breakdown,
            "historicalMatchups": self._historical(f1_actual, f2_actual),
        }

    # ------------------------------------------------------------------
    # Explanation
    # ------------------------------------------------------------------

    def _display(self, fname: str) -> str:
        base = fname.replace("diff_", "", 1)
        return FEATURE_DISPLAY.get(base, base.replace("_", " ").title())

    def _category(self, fname: str) -> str:
        base = fname.replace("diff_", "", 1)
        return FEATURE_CATEGORY.get(base, "other")

    def _key_factors(self, scaled: np.ndarray, f1_name: str, f2_name: str, s1: pd.Series, s2: pd.Series):
        if self.lr_coefficients is None:
            return []
        contribs = self.lr_coefficients * scaled
        raw = []
        for i, fname in enumerate(self.features):
            if fname.startswith("wc_"):
                continue
            c = contribs[i]
            if abs(c) < 0.005:
                continue
            base = fname.replace("diff_", "", 1)
            snap_aliases = {
                "rating": "post_rating", "rd": "post_rd",
                "height": "height_cm", "reach": "reach_cm", "weight": "weight_lbs",
                "age": "current_age", "layoff": "current_layoff_days",
                "career_fights": "career_fights_before",
                "career_wins": "career_wins_before",
                "career_losses": "career_losses_before",
                "win_streak": "win_streak_before",
                "lose_streak": "lose_streak_before",
            }
            col = snap_aliases.get(base, base)
            v1 = self._safe(s1.get(col)) if col in s1.index else 0.0
            v2 = self._safe(s2.get(col)) if col in s2.index else 0.0
            raw.append({
                "factor": self._display(fname),
                "category": self._category(fname),
                "advantage": f1_name if c > 0 else f2_name,
                "impact": round(abs(float(c)), 4),
                "f1Value": round(v1, 1) if isinstance(v1, float) else v1,
                "f2Value": round(v2, 1) if isinstance(v2, float) else v2,
            })

        seen = {}
        for r in raw:
            if r["factor"] not in seen or r["impact"] > seen[r["factor"]]["impact"]:
                seen[r["factor"]] = r
        return sorted(seen.values(), key=lambda x: x["impact"], reverse=True)[:8]

    def _category_scores(self, scaled: np.ndarray) -> dict:
        if self.lr_coefficients is None:
            return {c: {"score": 50, "advantage": "even", "label": l} for c, l in CATEGORY_LABELS.items()}
        contribs = self.lr_coefficients * scaled
        buckets: dict[str, list[float]] = {}
        for i, fname in enumerate(self.features):
            cat = self._category(fname)
            if cat in CATEGORY_LABELS:
                buckets.setdefault(cat, []).append(contribs[i])

        result = {}
        for cat, label in CATEGORY_LABELS.items():
            vals = buckets.get(cat, [])
            if vals:
                avg = float(np.mean(vals))
                score = 50 + 50 * float(np.tanh(avg * 3))
                score = max(0, min(100, score))
                if avg > 0.01:
                    adv = "fighter1"
                elif avg < -0.01:
                    adv = "fighter2"
                else:
                    adv = "even"
                result[cat] = {"score": round(score), "advantage": adv, "label": label}
            else:
                result[cat] = {"score": 50, "advantage": "even", "label": label}
        return result

    def _historical(self, f1_actual: str, f2_actual: str):
        if self.norm_fights is None:
            return None
        f1 = f1_actual.lower().strip()
        f2 = f2_actual.lower().strip()
        nf = self.norm_fights
        mask = (
            ((nf["red_name"].str.lower() == f1) & (nf["blue_name"].str.lower() == f2)) |
            ((nf["red_name"].str.lower() == f2) & (nf["blue_name"].str.lower() == f1))
        )
        matches = nf[mask].sort_values("date", ascending=False)
        if matches.empty:
            return None
        results = []
        for _, row in matches.iterrows():
            if row["winner"] == "red":
                winner = row["red_name"]
            elif row["winner"] == "blue":
                winner = row["blue_name"]
            elif row["winner"] == "draw":
                winner = "Draw"
            else:
                winner = "No Contest"
            results.append({
                "date": str(row["date"].date()) if pd.notna(row.get("date")) else "",
                "winner": winner,
                "method": str(row["method"]) if pd.notna(row.get("method")) else None,
                "round": int(row["finish_round"]) if pd.notna(row.get("finish_round")) else None,
                "time": str(row["finish_time"]) if pd.notna(row.get("finish_time")) else None,
            })
        return results

    # ------------------------------------------------------------------
    # Info
    # ------------------------------------------------------------------

    def get_model_info(self) -> dict:
        return {
            "fightCount": self.fight_count,
            "fighterCount": self.fighter_count,
            "featureCount": self.feature_count,
            "models": self.model_metrics,
            "categories": list(CATEGORY_LABELS.keys()),
        }

    def save_artifact(self, artifact_path: str) -> None:
        joblib.dump({
            "models": self.models,
            "scaler": self.scaler,
            "imputer": self.imputer,
            "features": self.features,
            "model_metrics": self.model_metrics,
            "lr_coefficients": self.lr_coefficients,
        }, artifact_path, compress=3)