"""Accuracy as a function of ensemble size.

For every size 1..M, samples distinct member subsets, scores each as an ensemble
(argmax of the mean member probability), and reports the mean accuracy with a
+/-1 std band across the sampled configurations. Answers "how many models do we
actually need?".
"""

import itertools
import math
from typing import Any, Dict, List, Tuple

import numpy as np
from matplotlib.ticker import MaxNLocator

from analysis.plotting import panel_grid, save_figure, save_json, style_axes
from analysis.sources import FAMILY_LABELS, ordered_families
from analysis.context import AnalysisContext, FAMILY_COLORS
from analysis.registry import register

#: Member-subset configurations sampled per ensemble size (all of them when the
#: number of distinct subsets is smaller).
TARGET_CONFIGS = 30


def _sample_subsets(
    n_models: int, size: int, target: int, rng: np.random.Generator
) -> List[List[int]]:
    """Distinct member-index subsets of ``size``.

    Enumerates exhaustively when ``C(n_models, size) <= target`` (e.g. size 1 or
    size M), otherwise draws ``target`` distinct random subsets.
    """
    if math.comb(n_models, size) <= target:
        return [list(combo) for combo in itertools.combinations(range(n_models), size)]

    seen: set[Tuple[int, ...]] = set()
    subsets: List[List[int]] = []
    while len(subsets) < target:
        pick = tuple(
            sorted(int(i) for i in rng.choice(n_models, size=size, replace=False))
        )
        if pick not in seen:
            seen.add(pick)
            subsets.append(list(pick))
    return subsets


def _size_curve(
    probs: np.ndarray, true_idx: np.ndarray, target: int, rng: np.random.Generator
) -> Dict[str, Any]:
    """Mean/std accuracy over sampled member subsets, for every size 1..M."""
    n_models = probs.shape[0]
    sizes: List[int] = []
    means: List[float] = []
    stds: List[float] = []
    n_configs: List[int] = []
    for size in range(1, n_models + 1):
        subsets = _sample_subsets(n_models, size, target, rng)
        accuracies = [
            float(np.mean(probs[subset].mean(axis=0).argmax(axis=1) == true_idx))
            for subset in subsets
        ]
        acc = np.asarray(accuracies, dtype=float)
        sizes.append(size)
        means.append(float(acc.mean()))
        stds.append(float(acc.std()))
        n_configs.append(len(subsets))
    return {"sizes": sizes, "mean": means, "std": stds, "n_configs": n_configs}


@register(
    "ensemble_size_accuracy",
    title="Accuracy vs. number of ensemble members",
)
def ensemble_size_accuracy(ctx: AnalysisContext) -> None:
    base_seed = ctx.base_seed()

    curves: Dict[str, Dict[str, Any]] = {}
    for index, family in enumerate(ctx.families):
        selected = ctx.clean_source(family)
        if selected is None:
            continue
        stem, ds = selected
        probs = ctx.probs(family, ds)
        if probs.shape[0] < 1:
            continue
        # Distinct, reproducible RNG stream per family for subset sampling.
        rng = np.random.default_rng([base_seed, index])
        curve = _size_curve(probs, ctx.true_index(ds), TARGET_CONFIGS, rng)
        curve["source"] = stem
        curves[family] = curve
        ctx.log.info(
            f"ensemble_size_accuracy: {family} - sizes 1..{probs.shape[0]}, "
            f"up to {TARGET_CONFIGS} configs each."
        )

    if not curves:
        ctx.log.error("ensemble_size_accuracy: no usable probabilities.")
        return

    fig, axes = panel_grid(1, n_cols=1, panel_size=(7.0, 4.3))
    ax = axes[0]
    for family in ordered_families(curves.keys()):
        curve = curves[family]
        color = FAMILY_COLORS.get(family, "#666666")
        x = np.asarray(curve["sizes"])
        mean = np.asarray(curve["mean"])
        std = np.asarray(curve["std"])
        ax.fill_between(x, mean - std, mean + std, color=color, alpha=0.18, linewidth=0)
        ax.plot(
            x,
            mean,
            color=color,
            marker="o",
            markersize=4,
            linewidth=2,
            label=FAMILY_LABELS.get(family, family),
        )

    ax.set_xlabel("Ensemble size (number of models)")
    ax.set_ylabel("Accuracy")
    ax.set_title("Ensemble accuracy vs. number of models")
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax.margins(x=0.02)
    style_axes(ax)
    ax.legend(frameon=False, loc="lower right")

    png = save_figure(fig, ctx.out_dir, "ensemble_size_accuracy")
    save_json(
        {"target_configs": TARGET_CONFIGS, "families": curves},
        ctx.out_dir,
        "ensemble_size_accuracy",
    )
    ctx.log.info(
        f"ensemble_size_accuracy: wrote {png.name} ({len(curves)} family/families)."
    )