src/chess_tutor/web/style_elo_plot.py

"""Elo vs tactical/positional playstyle scatter and 0–100 style scale for the web UI."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import joblib
import numpy as np
import pandas as pd

from chess_tutor.inference.pipeline import InferenceResult
from chess_tutor.paths import (
    default_cohort_db_path,
    default_web_style_artifacts_path,
)
from chess_tutor.playstyle_coaching.data import load_snapshots_df
from chess_tutor.playstyle_coaching.teacher import bracket_mid_elo

TACTICAL_FEATURES: tuple[str, ...] = (
    "profile_agg_cctx_capture_under_sharp_rate_mean",
    "profile_agg_cctx_king_unsafe_and_sharp_rate_mean",
    "profile_agg_move_capture_given_pband_tension_sharp_rate_mean",
    "profile_agg_ten_adds_tension_rate_mean",
    "profile_agg_cxp_sharpen_given_pband_tension_sharp_rate_mean",
)

POSITIONAL_FEATURES: tuple[str, ...] = (
    "profile_agg_human_quiet_in_sharp_position_rate_mean",
    "profile_agg_cxp_simplify_given_pband_tension_sharp_rate_mean",
    "profile_agg_human_pawn_break_rate_mean",
    "profile_agg_cctx_simplify_when_complex_rate_mean",
    "profile_agg_pband_tension_quiet_rate_mean",
    "profile_agg_theme_quietMove_rate_mean",
)

STYLE_ZONES: tuple[dict[str, Any], ...] = (
    {
        "min": 0,
        "max": 20,
        "id": "extremely_strategic",
        "label": "Extremely strategic",
        "short_label": "Ext. strategic",
        "color": "#4c78a8",
    },
    {
        "min": 20,
        "max": 40,
        "id": "strategic",
        "label": "Strategic",
        "short_label": "Strategic",
        "color": "#72b7b2",
    },
    {
        "min": 40,
        "max": 60,
        "id": "balanced",
        "label": "Balanced",
        "short_label": "Balanced",
        "color": "#54a24b",
    },
    {
        "min": 60,
        "max": 80,
        "id": "aggressive",
        "label": "Aggressive",
        "short_label": "Aggressive",
        "color": "#f58518",
    },
    {
        "min": 80,
        "max": 100,
        "id": "extremely_aggressive",
        "label": "Extremely aggressive",
        "short_label": "Ext. aggressive",
        "color": "#e45756",
    },
)

ELO_COLUMN = "meta_avg_player_elo_snapshot"
BRACKET_COLUMN = "elo_bracket"


@dataclass(frozen=True)
class StyleCohortData:
    elos: list[float]
    styles: list[float]
    brackets: list[str]
    usernames: list[str]
    user_style: float
    user_elo: float
    username: str
    user_bracket: str
    n_tactical_features: int
    n_positional_features: int


@dataclass(frozen=True)
class StyleCohortArtifacts:
    tactical_cols: list[str]
    positional_cols: list[str]
    tactical_means: dict[str, float]
    tactical_stds: dict[str, float]
    positional_means: dict[str, float]
    positional_stds: dict[str, float]
    elos: list[float]
    styles: list[float]
    brackets: list[str]
    usernames: list[str]
    n_tactical_features: int
    n_positional_features: int


_STYLE_COHORT_ARTIFACTS_CACHE: dict[str, StyleCohortArtifacts] = {}


def _try_load_style_artifacts(corpus_id: str) -> StyleCohortArtifacts | None:
    if corpus_id in _STYLE_COHORT_ARTIFACTS_CACHE:
        return _STYLE_COHORT_ARTIFACTS_CACHE[corpus_id]
    path = default_web_style_artifacts_path(corpus_id)
    if not path.is_file():
        return None
    try:
        obj = joblib.load(path)
    except Exception:
        return None
    try:
        artifacts = StyleCohortArtifacts(**obj)
    except Exception:
        return None
    _STYLE_COHORT_ARTIFACTS_CACHE[corpus_id] = artifacts
    return artifacts


def _save_style_artifacts(corpus_id: str, artifacts: StyleCohortArtifacts) -> None:
    path = default_web_style_artifacts_path(corpus_id)
    try:
        joblib.dump(
            {
                "tactical_cols": artifacts.tactical_cols,
                "positional_cols": artifacts.positional_cols,
                "tactical_means": artifacts.tactical_means,
                "tactical_stds": artifacts.tactical_stds,
                "positional_means": artifacts.positional_means,
                "positional_stds": artifacts.positional_stds,
                "elos": artifacts.elos,
                "styles": artifacts.styles,
                "brackets": artifacts.brackets,
                "usernames": artifacts.usernames,
                "n_tactical_features": artifacts.n_tactical_features,
                "n_positional_features": artifacts.n_positional_features,
            },
            path,
        )
    except Exception:
        # Cache is best-effort only.
        pass


def _build_style_cohort_artifacts(
    corpus_id: str,
) -> StyleCohortArtifacts | None:
    db_path = default_cohort_db_path(corpus_id)
    if not db_path.is_file():
        return None

    try:
        cohort = load_snapshots_df(db_path, corpus_id)
    except (ValueError, OSError):
        return None

    if cohort.empty or BRACKET_COLUMN not in cohort.columns:
        return None

    cohort = cohort.dropna(subset=[BRACKET_COLUMN]).copy()
    tactical_cols = _available_features(cohort, TACTICAL_FEATURES)
    positional_cols = _available_features(cohort, POSITIONAL_FEATURES)
    if not tactical_cols or not positional_cols:
        return None

    tactical_frame = cohort[tactical_cols].astype(float)
    positional_frame = cohort[positional_cols].astype(float)

    tactical_means = {col: float(tactical_frame[col].mean()) for col in tactical_cols}
    tactical_stds = {
        col: float(tactical_frame[col].std(ddof=0)) or 1.0 for col in tactical_cols
    }
    positional_means = {
        col: float(positional_frame[col].mean()) for col in positional_cols
    }
    positional_stds = {
        col: float(positional_frame[col].std(ddof=0)) or 1.0
        for col in positional_cols
    }

    for col in tactical_cols:
        if tactical_stds[col] < 1e-12:
            tactical_stds[col] = 1.0
    for col in positional_cols:
        if positional_stds[col] < 1e-12:
            positional_stds[col] = 1.0

    elos: list[float] = []
    styles: list[float] = []
    brackets: list[str] = []
    usernames: list[str] = []

    for _, row in cohort.iterrows():
        elo = _listed_elo(row)
        if elo is None:
            continue
        style = _style_axis(
            row,
            tactical_cols,
            positional_cols,
            tactical_means,
            tactical_stds,
            positional_means,
            positional_stds,
        )
        if not np.isfinite(style):
            continue
        elos.append(elo)
        styles.append(style)
        brackets.append(str(row[BRACKET_COLUMN]))
        usernames.append(str(row.get("username", "")))

    if len(elos) < 20:
        return None

    return StyleCohortArtifacts(
        tactical_cols=tactical_cols,
        positional_cols=positional_cols,
        tactical_means=tactical_means,
        tactical_stds=tactical_stds,
        positional_means=positional_means,
        positional_stds=positional_stds,
        elos=elos,
        styles=styles,
        brackets=brackets,
        usernames=usernames,
        n_tactical_features=len(tactical_cols),
        n_positional_features=len(positional_cols),
    )


def build_style_cohort_artifacts(corpus_id: str = "standard_600") -> bool:
    """Precompute style cohort artifacts for the web UI."""
    artifacts = _build_style_cohort_artifacts(corpus_id)
    if artifacts is None:
        return False
    _STYLE_COHORT_ARTIFACTS_CACHE[corpus_id] = artifacts
    _save_style_artifacts(corpus_id, artifacts)
    return True


def _available_features(
    cohort: pd.DataFrame,
    features: tuple[str, ...],
) -> list[str]:
    return [col for col in features if col in cohort.columns]


def _listed_elo(row: pd.Series) -> float | None:
    if ELO_COLUMN in row.index:
        val = row.get(ELO_COLUMN)
        if val is not None and not (isinstance(val, float) and np.isnan(val)):
            elo = float(val)
            if elo > 0.0:
                return elo
    bracket = row.get(BRACKET_COLUMN)
    if bracket is not None and str(bracket).strip():
        return bracket_mid_elo(str(bracket))
    return None


def _group_zscore_mean(
    row: pd.Series,
    cols: list[str],
    means: dict[str, float],
    stds: dict[str, float],
) -> float:
    zs: list[float] = []
    for col in cols:
        if col not in row.index:
            continue
        val = row.get(col)
        if val is None or (isinstance(val, float) and np.isnan(val)):
            continue
        zs.append((float(val) - means[col]) / stds[col])
    if not zs:
        return 0.0
    return float(np.mean(zs))


def _style_axis(
    row: pd.Series,
    tactical_cols: list[str],
    positional_cols: list[str],
    tactical_means: dict[str, float],
    tactical_stds: dict[str, float],
    positional_means: dict[str, float],
    positional_stds: dict[str, float],
) -> float:
    tactical = _group_zscore_mean(row, tactical_cols, tactical_means, tactical_stds)
    positional = _group_zscore_mean(row, positional_cols, positional_means, positional_stds)
    return tactical - positional


def _percentile_score(styles: list[float], user_style: float) -> float:
    arr = np.asarray(styles, dtype=float)
    below = float(np.sum(arr < user_style))
    equal = float(np.sum(arr == user_style))
    pct = (below + 0.5 * equal) / len(arr) * 100.0
    return round(min(100.0, max(0.0, pct)), 1)


def _zone_for_score(score: float) -> dict[str, Any]:
    value = min(100.0, max(0.0, score))
    for zone in STYLE_ZONES:
        if zone["min"] <= value < zone["max"]:
            return zone
    return STYLE_ZONES[-1]


def _load_style_cohort_data(
    result: InferenceResult,
    *,
    corpus_id: str = "standard_600",
) -> StyleCohortData | None:
    artifacts = _try_load_style_artifacts(corpus_id)
    if artifacts is None:
        artifacts = _build_style_cohort_artifacts(corpus_id)
        if artifacts is None:
            return None
        _STYLE_COHORT_ARTIFACTS_CACHE[corpus_id] = artifacts
        _save_style_artifacts(corpus_id, artifacts)

    user_row = result.snapshot.to_series()
    user_elo = _listed_elo(user_row)
    if user_elo is None:
        user_elo = float(result.player_elo)
    user_style = _style_axis(
        user_row,
        artifacts.tactical_cols,
        artifacts.positional_cols,
        artifacts.tactical_means,
        artifacts.tactical_stds,
        artifacts.positional_means,
        artifacts.positional_stds,
    )
    if not np.isfinite(user_style):
        return None

    return StyleCohortData(
        elos=artifacts.elos,
        styles=artifacts.styles,
        brackets=artifacts.brackets,
        usernames=artifacts.usernames,
        user_style=user_style,
        user_elo=user_elo,
        username=result.username,
        user_bracket=str(result.snapshot.elo_bracket),
        n_tactical_features=artifacts.n_tactical_features,
        n_positional_features=artifacts.n_positional_features,
    )


def build_style_scale(
    result: InferenceResult,
    *,
    corpus_id: str = "standard_600",
) -> dict[str, Any] | None:
    """
    Map the player onto a 0–100 scale (0 = positional/strategic, 100 = tactical/aggressive)
    using cohort percentiles.
    """
    data = _load_style_cohort_data(result, corpus_id=corpus_id)
    if data is None:
        return None

    score = _percentile_score(data.styles, data.user_style)
    zone = _zone_for_score(score)
    return {
        "score": score,
        "zone_id": zone["id"],
        "zone_label": zone["label"],
        "zones": list(STYLE_ZONES),
        "axis_left": "Positional / strategic",
        "axis_right": "Tactical / aggressive",
        "label": data.username,
        "bracket": data.user_bracket,
        "n_total": len(data.styles),
    }


def build_style_elo_plot(
    result: InferenceResult,
    *,
    corpus_id: str = "standard_600",
    max_points: int = 6000,
    seed: int = 42,
) -> dict[str, Any] | None:
    """Place cohort snapshots and the analyzed player on listed Elo (x) vs style (y)."""
    data = _load_style_cohort_data(result, corpus_id=corpus_id)
    if data is None:
        return None

    n_total = len(data.elos)
    indices = np.arange(n_total)
    if n_total > max_points:
        rng = np.random.default_rng(seed)
        indices = np.sort(rng.choice(n_total, size=max_points, replace=False))

    points = [
        {
            "x": round(float(data.elos[i]), 1),
            "y": round(float(data.styles[i]), 4),
            "bracket": data.brackets[i],
            "username": data.usernames[i],
        }
        for i in indices
    ]

    return {
        "axis_labels": ["Listed Elo", "Tactical ← → Positional"],
        "n_total": n_total,
        "n_shown": len(points),
        "n_tactical_features": data.n_tactical_features,
        "n_positional_features": data.n_positional_features,
        "points": points,
        "user": {
            "x": round(float(data.user_elo), 1),
            "y": round(float(data.user_style), 4),
            "label": data.username,
            "bracket": data.user_bracket,
        },
    }