src/chess_tutor/analysis/aggregates.py

"""Aggregate engine vs move-counterfactual heuristic vectors across games."""

from __future__ import annotations

from typing import Any, Dict, List, Sequence

import numpy as np

from chess_tutor.analysis.types import MoveHeuristicVectors


def mean_abs_heuristic_delta(vectors: MoveHeuristicVectors) -> float:
    """Mean |move_heuristic − engine_heuristic| over MLP outputs (0–100 units)."""
    ref = np.asarray(vectors.reference_0_100, dtype=np.float64)
    inv = np.asarray(vectors.inverted_0_100, dtype=np.float64)
    return float(np.mean(np.abs(inv - ref)))


def mean_abs_heuristic_delta_chained(chained: Any) -> float:
    """Adapter for ``ChainedMoveConceptResult`` (lazy import avoids cycles in tests)."""
    from chess_tutor.analysis.adapters import vectors_from_chained

    return mean_abs_heuristic_delta(vectors_from_chained(chained))


def vec_dict_add(store: Dict[str, np.ndarray], key: str, vec: np.ndarray) -> None:
    v = np.asarray(vec, dtype=np.float64).reshape(-1)
    if key not in store:
        store[key] = np.zeros_like(v)
    store[key] = store[key] + v


def per_heuristic_breakdown_rows(
    names: Sequence[str],
    n_moves: int,
    ref_sum: np.ndarray,
    inv_sum: np.ndarray,
    abs_diff_sum: np.ndarray,
    signed_diff_sum: np.ndarray,
) -> List[Dict[str, Any]]:
    """One row per heuristic dimension for a game or for all moves pooled."""
    if n_moves <= 0:
        return []
    n_moves_f = float(n_moves)
    ref_m = ref_sum / n_moves_f
    inv_m = inv_sum / n_moves_f
    mean_abs = abs_diff_sum / n_moves_f
    mean_signed = signed_diff_sum / n_moves_f
    out: List[Dict[str, Any]] = []
    for i, name in enumerate(names):
        out.append(
            {
                "heuristic": str(name),
                "engine_mean_0_100": round(float(ref_m[i]), 4),
                "move_counterfactual_mean_0_100": round(float(inv_m[i]), 4),
                "mean_signed_difference_move_minus_engine": round(float(mean_signed[i]), 4),
                "mean_abs_difference_across_moves": round(float(mean_abs[i]), 4),
            }
        )
    return out


def per_heuristic_equal_weight_across_games(
    names: Sequence[str],
    per_game_ref_sum: Dict[str, np.ndarray],
    per_game_inv_sum: Dict[str, np.ndarray],
    per_game_abs_sum: Dict[str, np.ndarray],
    per_game_signed_sum: Dict[str, np.ndarray],
    per_game_succ: Dict[str, int],
) -> List[Dict[str, Any]]:
    """Average per-game means so each game counts equally."""
    games = [gk for gk, n in per_game_succ.items() if n > 0 and gk in per_game_ref_sum]
    if not games:
        return []
    n_dim = len(names)
    rows: List[Dict[str, Any]] = []
    for i in range(n_dim):
        eng_g = [per_game_ref_sum[gk][i] / float(per_game_succ[gk]) for gk in games]
        mv_g = [per_game_inv_sum[gk][i] / float(per_game_succ[gk]) for gk in games]
        abs_g = [per_game_abs_sum[gk][i] / float(per_game_succ[gk]) for gk in games]
        sg_g = [per_game_signed_sum[gk][i] / float(per_game_succ[gk]) for gk in games]
        eng_m = float(np.mean(np.asarray(eng_g, dtype=np.float64)))
        mv_m = float(np.mean(np.asarray(mv_g, dtype=np.float64)))
        rows.append(
            {
                "heuristic": str(names[i]),
                "engine_mean_across_games_equal_weight": round(eng_m, 4),
                "move_counterfactual_mean_across_games_equal_weight": round(mv_m, 4),
                "mean_signed_difference_move_minus_engine": round(mv_m - eng_m, 4),
                "mean_abs_difference_averaged_per_game_then_across_games": round(
                    float(np.mean(np.asarray(abs_g, dtype=np.float64))), 4
                ),
                "mean_signed_difference_averaged_per_game_then_across_games": round(
                    float(np.mean(np.asarray(sg_g, dtype=np.float64))), 4
                ),
            }
        )
    return rows


def summary_across_heuristics(per_heuristic_rows: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
    """Mean of per-dimension scalars (twelve values → one summary number)."""
    if not per_heuristic_rows:
        return {
            "mean_abs_difference_averaged_over_heuristics": None,
            "mean_signed_difference_averaged_over_heuristics": None,
        }
    abs_v = np.asarray(
        [r["mean_abs_difference_across_moves"] for r in per_heuristic_rows], dtype=np.float64
    )
    sg_v = np.asarray(
        [r["mean_signed_difference_move_minus_engine"] for r in per_heuristic_rows],
        dtype=np.float64,
    )
    return {
        "mean_abs_difference_averaged_over_heuristics": round(float(np.mean(abs_v)), 4),
        "mean_signed_difference_averaged_over_heuristics": round(float(np.mean(sg_v)), 4),
    }


def summary_across_heuristics_equal_game(
    rows: Sequence[Dict[str, Any]],
) -> Dict[str, Any]:
    if not rows:
        return {
            "mean_abs_difference_averaged_over_heuristics": None,
            "mean_signed_difference_averaged_over_heuristics": None,
        }
    abs_v = np.asarray(
        [r["mean_abs_difference_averaged_per_game_then_across_games"] for r in rows],
        dtype=np.float64,
    )
    sg_v = np.asarray(
        [r["mean_signed_difference_move_minus_engine"] for r in rows],
        dtype=np.float64,
    )
    return {
        "mean_abs_difference_averaged_over_heuristics": round(float(np.mean(abs_v)), 4),
        "mean_signed_difference_averaged_over_heuristics": round(float(np.mean(sg_v)), 4),
    }