src/chess_tutor/web/embedding.py

"""Shared cohort playstyle embeddings (PCA / t-SNE) for web scatter plots."""

from __future__ import annotations

import warnings
from dataclasses import dataclass
from typing import Any, Literal

import joblib
import numpy as np
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.neighbors import NearestNeighbors
from sklearn.preprocessing import StandardScaler

from chess_tutor.inference.pipeline import InferenceResult
from chess_tutor.paths import (
    default_cohort_db_path,
    default_web_pca_artifacts_path,
    default_web_tsne_artifacts_path,
)
from chess_tutor.playstyle_coaching.data import load_snapshots_df
from chess_tutor.playstyle_coaching.teacher import (
    TeacherCoachingBundle,
    bracket_mid_elo,
    default_teacher_bundle_path,
    load_teacher_bundle,
)

ELO_COLUMN = "meta_avg_player_elo_snapshot"
BRACKET_COLUMN = "elo_bracket"


def _impute_column_means(x: np.ndarray) -> np.ndarray:
    col_means = np.nanmean(x, axis=0)
    out = x.copy()
    nan_mask = np.isnan(out)
    if nan_mask.any():
        out[nan_mask] = np.take(col_means, np.where(nan_mask)[1])
    return out


def _variance_column_mask(x: np.ndarray, min_col_std: float = 1e-12) -> np.ndarray:
    return np.std(x, axis=0) > min_col_std


def _pca_axis_label(component_index: int, variance_ratio: float) -> str:
    pct = round(float(variance_ratio) * 100, 1)
    return f"PC{component_index + 1} ({pct}% variance)"


def _tsne_axis_label(component_index: int) -> str:
    return f"t-SNE {component_index + 1}"


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 _user_feature_vector(
    result: InferenceResult,
    kept_cols: list[str],
    cohort_col_means: np.ndarray,
) -> np.ndarray | None:
    user_row = result.snapshot.to_series()
    user_x = np.array([float(user_row.get(col, np.nan)) for col in kept_cols], dtype=float)
    for j in range(len(user_x)):
        if np.isnan(user_x[j]):
            user_x[j] = float(cohort_col_means[j])
    if not np.all(np.isfinite(user_x)):
        return None
    return user_x


def _subsample_indices(n_total: int, max_points: int, seed: int) -> np.ndarray:
    if n_total <= max_points:
        return np.arange(n_total)
    rng = np.random.default_rng(seed)
    return np.sort(rng.choice(n_total, size=max_points, replace=False))


def _scatter_points(
    indices: np.ndarray,
    x_values: np.ndarray,
    y_values: np.ndarray,
    brackets: np.ndarray,
    usernames: np.ndarray,
) -> list[dict[str, Any]]:
    return [
        {
            "x": round(float(x_values[j]), 4),
            "y": round(float(y_values[j]), 4),
            "bracket": str(brackets[indices[j]]),
            "username": str(usernames[indices[j]]),
        }
        for j in range(len(indices))
    ]


@dataclass(frozen=True)
class PreparedCohort:
    cohort_n_total: int
    kept_cols: list[str]
    cohort_scaled: np.ndarray
    cohort_col_means: np.ndarray
    scaler: StandardScaler
    brackets: np.ndarray
    usernames: np.ndarray
    elos: np.ndarray


@dataclass
class EmbeddingState:
    prepared: PreparedCohort
    username: str
    user_bracket: str
    user_elo: float
    user_scaled: np.ndarray
    pca: PCA
    pca_xy: np.ndarray
    user_pca_xy: np.ndarray
    tsne_indices: np.ndarray | None = None
    tsne_xy: np.ndarray | None = None
    user_tsne_xy: np.ndarray | None = None


def _try_load_pca_artifacts(corpus_id: str) -> dict[str, Any] | None:
    path = default_web_pca_artifacts_path(corpus_id)
    if not path.is_file():
        return None
    try:
        obj = joblib.load(path)
    except Exception:
        return None
    if not isinstance(obj, dict):
        return None
    return obj


def _save_pca_artifacts(corpus_id: str, artifacts: dict[str, Any]) -> None:
    path = default_web_pca_artifacts_path(corpus_id)
    try:
        joblib.dump(artifacts, path)
    except Exception:
        # Artifacts are purely for speed; do not fail inference if caching fails.
        pass


def _try_load_tsne_artifacts(corpus_id: str) -> dict[str, Any] | None:
    path = default_web_tsne_artifacts_path(corpus_id)
    if not path.is_file():
        return None
    try:
        obj = joblib.load(path)
    except Exception:
        return None
    if not isinstance(obj, dict):
        return None
    return obj


def _save_tsne_artifacts(corpus_id: str, artifacts: dict[str, Any]) -> None:
    path = default_web_tsne_artifacts_path(corpus_id)
    try:
        joblib.dump(artifacts, path)
    except Exception:
        pass


def _build_tsne_artifacts(
    cohort_scaled: np.ndarray,
    *,
    tsne_seed: int = 42,
) -> dict[str, Any] | None:
    n = len(cohort_scaled)
    if n < 20:
        return None

    perplexity = min(30.0, max(5.0, (n - 1) / 3.0))
    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            tsne = TSNE(
                n_components=2,
                perplexity=perplexity,
                init="pca",
                learning_rate="auto",
                random_state=tsne_seed,
            )
            tsne_xy = tsne.fit_transform(cohort_scaled)
    except (ValueError, TypeError):
        return None

    return {
        "cohort_n_total": n,
        "tsne_xy": tsne_xy.astype(np.float32, copy=False),
        "tsne_seed": tsne_seed,
    }


def _place_user_tsne(
    user_scaled: np.ndarray,
    cohort_scaled: np.ndarray,
    tsne_xy: np.ndarray,
    *,
    k_neighbors: int = 12,
) -> np.ndarray:
    """Place a new player on a fixed cohort t-SNE map via weighted kNN in feature space."""
    n_neighbors = min(k_neighbors, len(cohort_scaled))
    nbrs = NearestNeighbors(n_neighbors=n_neighbors)
    nbrs.fit(cohort_scaled)
    dists, indices = nbrs.kneighbors(user_scaled.reshape(1, -1))
    dists = dists[0]
    indices = indices[0]
    if dists[0] < 1e-12:
        return tsne_xy[indices[0]].astype(float)
    weights = 1.0 / (dists + 1e-12)
    weights /= weights.sum()
    return np.average(tsne_xy[indices], axis=0, weights=weights)


def _ensure_tsne_artifacts(
    corpus_id: str,
    cohort_scaled: np.ndarray,
    *,
    tsne_seed: int = 42,
) -> dict[str, Any] | None:
    tsne_artifacts = _try_load_tsne_artifacts(corpus_id)
    if (
        tsne_artifacts is not None
        and int(tsne_artifacts.get("cohort_n_total", -1)) == len(cohort_scaled)
    ):
        return tsne_artifacts

    tsne_artifacts = _build_tsne_artifacts(cohort_scaled, tsne_seed=tsne_seed)
    if tsne_artifacts is None:
        return None
    _save_tsne_artifacts(corpus_id, tsne_artifacts)
    return tsne_artifacts


def _build_pca_artifacts(
    *,
    corpus_id: str,
    teacher_bundle: TeacherCoachingBundle,
    pca_seed: int,
) -> dict[str, Any] | 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()
    playstyle_cols = list(teacher_bundle.playstyle_cols)
    if len(playstyle_cols) < 2:
        return None

    missing_cols = [c for c in playstyle_cols if c not in cohort.columns]
    if missing_cols:
        return None

    cohort_x = cohort[playstyle_cols].astype(float).values
    cohort_x = _impute_column_means(cohort_x)
    keep = _variance_column_mask(cohort_x)
    if int(np.sum(keep)) < 2:
        return None

    kept_cols = [c for c, ok in zip(playstyle_cols, keep, strict=True) if ok]
    cohort_x = cohort_x[:, keep]
    if len(cohort_x) < 20:
        return None

    cohort_col_means = np.nanmean(cohort_x, axis=0)
    scaler = StandardScaler()
    cohort_scaled = scaler.fit_transform(cohort_x).astype(np.float32, copy=False)
    cohort_col_means = cohort_col_means.astype(np.float32, copy=False)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        pca = PCA(n_components=2, random_state=pca_seed)
        pca_xy = pca.fit_transform(cohort_scaled)

    brackets = cohort[BRACKET_COLUMN].astype(str).to_numpy()
    usernames = cohort["username"].astype(str).to_numpy()
    elos = np.array(
        [_listed_elo(row) or np.nan for _, row in cohort.iterrows()],
        dtype=float,
    )

    return {
        "cohort_n_total": len(cohort),
        "kept_cols": kept_cols,
        "cohort_col_means": cohort_col_means,
        "scaler": scaler,
        "cohort_scaled": cohort_scaled,
        "brackets": brackets,
        "usernames": usernames,
        "elos": elos,
        "pca": pca,
        "pca_xy": pca_xy,
    }


def build_cohort_pca_artifacts(
    corpus_id: str,
    *,
    pca_seed: int = 42,
    teacher_bundle: TeacherCoachingBundle | None = None,
) -> bool:
    """
    Precompute cohort PCA + scaling artifacts for fast web graph generation.
    """
    if teacher_bundle is None:
        teacher_path = default_teacher_bundle_path(corpus_id)
        if not teacher_path.is_file():
            return False
        teacher_bundle = load_teacher_bundle(teacher_path)

    artifacts = _build_pca_artifacts(
        corpus_id=corpus_id,
        teacher_bundle=teacher_bundle,
        pca_seed=pca_seed,
    )
    if artifacts is None:
        return False
    _save_pca_artifacts(corpus_id, artifacts)
    return True


def build_cohort_tsne_artifacts(
    corpus_id: str,
    *,
    tsne_seed: int = 42,
) -> bool:
    """Precompute cohort t-SNE coordinates for fast web graph generation."""
    pca_artifacts = _try_load_pca_artifacts(corpus_id)
    if pca_artifacts is None:
        if not build_cohort_pca_artifacts(corpus_id, pca_seed=tsne_seed):
            return False
        pca_artifacts = _try_load_pca_artifacts(corpus_id)
    if pca_artifacts is None:
        return False

    tsne_artifacts = _build_tsne_artifacts(
        pca_artifacts["cohort_scaled"],
        tsne_seed=tsne_seed,
    )
    if tsne_artifacts is None:
        return False
    _save_tsne_artifacts(corpus_id, tsne_artifacts)
    return True


def prepare_cohort_embedding(
    result: InferenceResult,
    *,
    corpus_id: str = "standard_600",
    teacher_bundle: TeacherCoachingBundle | None = None,
    tsne_seed: int = 42,
) -> EmbeddingState | None:
    artifacts = _try_load_pca_artifacts(corpus_id)

    if artifacts is None:
        teacher_path = default_teacher_bundle_path(corpus_id)
        if teacher_bundle is None:
            if not teacher_path.is_file():
                return None
            teacher_bundle = load_teacher_bundle(teacher_path)

        artifacts = _build_pca_artifacts(
            corpus_id=corpus_id,
            teacher_bundle=teacher_bundle,
            pca_seed=tsne_seed,
        )
        if artifacts is None:
            return None
        _save_pca_artifacts(corpus_id, artifacts)

    kept_cols = artifacts["kept_cols"]
    cohort_col_means = artifacts["cohort_col_means"]
    scaler = artifacts["scaler"]
    cohort_scaled = artifacts["cohort_scaled"]
    brackets = artifacts["brackets"]
    usernames = artifacts["usernames"]
    elos = artifacts["elos"]
    pca = artifacts["pca"]
    pca_xy = artifacts["pca_xy"]

    user_x = _user_feature_vector(result, kept_cols, cohort_col_means)
    if user_x is None:
        return None

    user_scaled = scaler.transform(user_x.reshape(1, -1))[0]
    user_pca_xy = pca.transform(user_scaled.reshape(1, -1))[0]

    prepared = PreparedCohort(
        cohort_n_total=int(artifacts["cohort_n_total"]),
        kept_cols=kept_cols,
        cohort_scaled=cohort_scaled,
        cohort_col_means=cohort_col_means,
        scaler=scaler,
        brackets=brackets,
        usernames=usernames,
        elos=elos,
    )

    user_row = result.snapshot.to_series()
    user_elo = _listed_elo(user_row)
    if user_elo is None:
        user_elo = float(result.player_elo)

    state = EmbeddingState(
        prepared=prepared,
        username=result.username,
        user_bracket=str(result.snapshot.elo_bracket),
        user_elo=user_elo,
        user_scaled=user_scaled,
        pca=pca,
        pca_xy=pca_xy,
        user_pca_xy=user_pca_xy,
    )

    n_total = prepared.cohort_n_total
    tsne_artifacts = _ensure_tsne_artifacts(
        corpus_id,
        cohort_scaled,
        tsne_seed=tsne_seed,
    )
    if tsne_artifacts is not None:
        state.tsne_indices = np.arange(n_total)
        state.tsne_xy = tsne_artifacts["tsne_xy"]
        state.user_tsne_xy = _place_user_tsne(user_scaled, cohort_scaled, state.tsne_xy)
    return state


def build_embedding_2d_plot(
    state: EmbeddingState,
    method: Literal["pca", "tsne"],
    *,
    max_points: int = 6000,
    seed: int = 42,
) -> dict[str, Any] | None:
    prepared = state.prepared
    n_total = prepared.cohort_n_total

    if method == "pca":
        indices = _subsample_indices(n_total, max_points, seed)
        points = _scatter_points(
            indices,
            state.pca_xy[indices, 0],
            state.pca_xy[indices, 1],
            prepared.brackets,
            prepared.usernames,
        )
        evr = np.nan_to_num(state.pca.explained_variance_ratio_, nan=0.0)
        return {
            "method": "pca",
            "axis_labels": [_pca_axis_label(0, evr[0]), _pca_axis_label(1, evr[1])],
            "variance_explained": [round(float(evr[0]), 4), round(float(evr[1]), 4)],
            "n_total": n_total,
            "n_shown": len(points),
            "n_features": len(prepared.kept_cols),
            "points": points,
            "user": {
                "x": round(float(state.user_pca_xy[0]), 4),
                "y": round(float(state.user_pca_xy[1]), 4),
                "label": state.username,
                "bracket": state.user_bracket,
            },
        }

    if state.tsne_indices is None or state.tsne_xy is None or state.user_tsne_xy is None:
        return None

    fit_indices = state.tsne_indices
    tsne_xy = state.tsne_xy

    points = [
        {
            "x": round(float(tsne_xy[j, 0]), 4),
            "y": round(float(tsne_xy[j, 1]), 4),
            "bracket": str(prepared.brackets[i]),
            "username": str(prepared.usernames[i]),
        }
        for j, i in enumerate(fit_indices)
    ]
    return {
        "method": "tsne",
        "axis_labels": [_tsne_axis_label(0), _tsne_axis_label(1)],
        "n_total": n_total,
        "n_shown": len(points),
        "n_features": len(prepared.kept_cols),
        "points": points,
        "user": {
            "x": round(float(state.user_tsne_xy[0]), 4),
            "y": round(float(state.user_tsne_xy[1]), 4),
            "label": state.username,
            "bracket": state.user_bracket,
        },
    }


def build_elo_embedding_plot(
    state: EmbeddingState,
    method: Literal["pca", "tsne"],
    *,
    component: int = 0,
    max_points: int = 6000,
    seed: int = 42,
) -> dict[str, Any] | None:
    prepared = state.prepared
    valid = np.isfinite(prepared.elos)

    if method == "pca":
        y_all = state.pca_xy[:, component]
        valid_indices = np.where(valid)[0]
        if len(valid_indices) > max_points:
            pick = _subsample_indices(len(valid_indices), max_points, seed)
            indices = valid_indices[pick]
        else:
            indices = valid_indices
        points = _scatter_points(
            indices,
            prepared.elos[indices],
            y_all[indices],
            prepared.brackets,
            prepared.usernames,
        )
        for p in points:
            p["x"] = round(float(p["x"]), 1)
        evr = np.nan_to_num(state.pca.explained_variance_ratio_, nan=0.0)
        y_label = _pca_axis_label(component, evr[component])
        user_y = float(state.user_pca_xy[component])
    else:
        if state.tsne_indices is None or state.tsne_xy is None or state.user_tsne_xy is None:
            return None
        fit_indices = state.tsne_indices
        fit_elos = prepared.elos[fit_indices]
        valid_fit = np.isfinite(fit_elos)
        fit_indices = fit_indices[valid_fit]
        tsne_y = state.tsne_xy[valid_fit, component]
        fit_elos = fit_elos[valid_fit]

        points = [
            {
                "x": round(float(fit_elos[j]), 1),
                "y": round(float(tsne_y[j]), 4),
                "bracket": str(prepared.brackets[i]),
                "username": str(prepared.usernames[i]),
            }
            for j, i in enumerate(fit_indices)
        ]
        y_label = _tsne_axis_label(component)
        user_y = float(state.user_tsne_xy[component])

    n_total = int(np.sum(valid)) if method == "pca" else len(points)
    return {
        "method": method,
        "axis_labels": ["Listed Elo", y_label],
        "n_total": n_total,
        "n_shown": len(points),
        "points": points,
        "user": {
            "x": round(float(state.user_elo), 1),
            "y": round(user_y, 4),
            "label": state.username,
            "bracket": state.user_bracket,
        },
    }


def build_tsne_web_plots(
    result: InferenceResult,
    *,
    corpus_id: str = "standard_600",
    teacher_bundle: TeacherCoachingBundle | None = None,
    max_points: int = 6000,
    seed: int = 42,
) -> dict[str, dict[str, Any] | None]:
    """Build t-SNE playstyle map and rating vs t-SNE plots for the web UI."""
    state = prepare_cohort_embedding(
        result,
        corpus_id=corpus_id,
        teacher_bundle=teacher_bundle,
        tsne_seed=seed,
    )
    if state is None:
        return {"tsne_plot": None, "elo_tsne_plot": None}
    return {
        "tsne_plot": build_embedding_2d_plot(state, "tsne", max_points=max_points, seed=seed),
        "elo_tsne_plot": build_elo_embedding_plot(
            state, "tsne", max_points=max_points, seed=seed
        ),
    }


def build_all_embedding_plots(
    result: InferenceResult,
    *,
    corpus_id: str = "standard_600",
    teacher_bundle: TeacherCoachingBundle | None = None,
    max_points: int = 6000,
    seed: int = 42,
) -> dict[str, dict[str, Any] | None]:
    state = prepare_cohort_embedding(
        result,
        corpus_id=corpus_id,
        teacher_bundle=teacher_bundle,
        tsne_seed=seed,
    )
    if state is None:
        return {
            "pca_plot": None,
            "tsne_plot": None,
            "elo_pca_plot": None,
            "elo_tsne_plot": None,
        }
    return {
        "pca_plot": build_embedding_2d_plot(state, "pca", max_points=max_points, seed=seed),
        "tsne_plot": build_embedding_2d_plot(state, "tsne", max_points=max_points, seed=seed),
        "elo_pca_plot": build_elo_embedding_plot(
            state, "pca", max_points=max_points, seed=seed
        ),
        "elo_tsne_plot": build_elo_embedding_plot(
            state, "tsne", max_points=max_points, seed=seed
        ),
    }