"""Evaluator — computes structured metrics from a FitResult."""

from __future__ import annotations

from typing import Any, Literal

import numpy as np
import numpy.typing as npt
import pandas as pd

from lizyml.core.types.fit_result import FitResult
from lizyml.metrics.base import BaseMetric
from lizyml.metrics.registry import MetricEntry, get_metrics_for_task
from lizyml.training.oof_assembly import compute_oof_valid_mask

TaskType = Literal["regression", "binary", "multiclass"]


def _normalize_multiclass_proba(
    pred: npt.NDArray[np.float64],
) -> npt.NDArray[np.float64]:
    """Row-wise normalize multiclass predictions to sum to 1.0.

    Required when the objective is ``multiclassova`` (independent sigmoid
    per class).  For ``multiclass`` (softmax), predictions already sum to
    1.0 and this operation is idempotent.
    """
    row_sums = pred.sum(axis=1, keepdims=True)
    # Guard against all-zero rows (degenerate edge case).
    # Positive values always sum to a positive number, so exact == 0.0
    # only fires when every element in the row is 0.0.
    row_sums = np.where(row_sums == 0.0, 1.0, row_sums)
    normalized: npt.NDArray[np.float64] = pred / row_sums
    return normalized


def _pred_for_metric(
    metric: BaseMetric,
    raw_pred: npt.NDArray[np.float64],
    task: TaskType,
) -> npt.NDArray[Any]:
    """Return the appropriate prediction array for *metric*.

    - ``needs_proba=True`` and ``needs_simplex=True``:
      - multiclass 2-D predictions are row-normalised so that metrics
        receiving probabilities always see a valid distribution (sum = 1).
        This is necessary for ``multiclassova`` (independent sigmoid),
        and idempotent for ``multiclass`` (softmax).
    - ``needs_proba=True`` and ``needs_simplex=False``:
      - Per-class OvR metrics (e.g. AUCPR, Brier) receive raw predictions.
    - ``needs_proba=False``: binarise (binary) or argmax (multiclass).
    """
    if metric.needs_proba:
        if task == "multiclass" and raw_pred.ndim == 2 and metric.needs_simplex:
            return _normalize_multiclass_proba(raw_pred)
        return raw_pred
    if task == "binary":
        return (raw_pred >= 0.5).astype(int)
    if task == "multiclass":
        result: npt.NDArray[np.intp] = raw_pred.argmax(axis=1)
        return result
    return raw_pred  # regression


def _compute_metrics(
    metrics: list[BaseMetric],
    y_true: npt.NDArray[Any],
    y_pred: npt.NDArray[np.float64],
    task: TaskType,
) -> dict[str, float]:
    result: dict[str, float] = {}
    for m in metrics:
        pred = _pred_for_metric(m, y_pred, task)
        result[m.name] = m(y_true, pred)
    return result


class Evaluator:
    """Compute structured evaluation metrics from a :class:`FitResult`.

    The output structure is fixed::

        {
            "raw": {
                "oof":          {metric_name: float, ...},
                "oof_per_fold": [{metric_name: float}, ...],
                "if_mean":      {metric_name: float, ...},
                "if_per_fold":  [{metric_name: float}, ...],
            },
            "calibrated": { ... }   # populated only when calibrator is set
        }

    Metric computation is centralised here; no metric calculation is done
    outside this class.

    Args:
        task: ML task type; used for task-compatibility validation.
    """

    def __init__(self, task: TaskType = "regression") -> None:
        self.task = task

    def evaluate(
        self,
        fit_result: FitResult,
        y: pd.Series | npt.NDArray[Any],
        metric_entries: list[MetricEntry],
    ) -> dict[str, Any]:
        """Compute OOF, IF-per-fold, and IF-mean metrics.

        Args:
            fit_result: Completed CV training output.
            y: Ground-truth target for the full dataset (same order as
                ``fit_result.oof_pred``).
            metric_entries: Metric names or parameterised ``MetricEntry``
                dicts to compute.  Must be compatible with the task this
                :class:`Evaluator` was constructed for.

        Returns:
            Nested dict with ``"raw"`` (and ``"calibrated"`` when applicable).
        """
        metrics = get_metrics_for_task(metric_entries, self.task)
        y_arr = np.asarray(y)

        # --- OOF coverage mask (H-0057) --------------------------------------
        oof_mask = compute_oof_valid_mask(fit_result.splits.outer, len(y_arr))
        covered_pred = fit_result.oof_pred[oof_mask]

        # Assert no NaN in covered rows — NaN here indicates a pipeline bug,
        # not a structurally uncovered row (e.g. TimeSeriesCV first period).
        _nan_flags = np.isnan(covered_pred)
        # For 2-D predictions (multiclass), collapse to per-row bool.
        nan_in_covered = (
            _nan_flags if _nan_flags.ndim == 1 else np.asarray(_nan_flags.any(axis=1))
        )
        if nan_in_covered.any():
            nan_count = int(nan_in_covered.sum())
            nan_indices = np.where(oof_mask)[0][nan_in_covered]
            raise ValueError(
                f"OOF predictions contain NaN in {nan_count} row(s) "
                f"covered by validation folds (indices: "
                f"{nan_indices[:5].tolist()}{'...' if nan_count > 5 else ''}). "
                f"This indicates a bug in the training pipeline."
            )

        oof_coverage = float(oof_mask.sum()) / len(y_arr) if len(y_arr) > 0 else 0.0

        # --- OOF metrics (covered rows only) ----------------------------------
        oof_scores = _compute_metrics(metrics, y_arr[oof_mask], covered_pred, self.task)

        # --- Per-fold OOF metrics (valid_idx) --------------------------------
        oof_per_fold: list[dict[str, float]] = []
        for _, valid_idx in fit_result.splits.outer:
            y_valid = y_arr[valid_idx]
            oof_valid_pred = fit_result.oof_pred[valid_idx]
            oof_per_fold.append(
                _compute_metrics(metrics, y_valid, oof_valid_pred, self.task)
            )

        # --- Per-fold IF metrics ---------------------------------------------
        if_per_fold: list[dict[str, float]] = []
        for k, (train_idx, _) in enumerate(fit_result.splits.outer):
            y_train = y_arr[train_idx]
            fold_pred = fit_result.if_pred_per_fold[k]
            fold_scores = _compute_metrics(metrics, y_train, fold_pred, self.task)
            if_per_fold.append(fold_scores)

        # --- IF mean ---------------------------------------------------------
        if_mean: dict[str, float] = {
            m.name: float(np.mean([fold[m.name] for fold in if_per_fold]))
            for m in metrics
        }

        return {
            "raw": {
                "oof": oof_scores,
                "oof_per_fold": oof_per_fold,
                "if_mean": if_mean,
                "if_per_fold": if_per_fold,
                "oof_coverage": oof_coverage,
            }
        }
