"""
SafeStacker — guaranteed-better model stacking.

Every model added is validated against eval data. If it doesn't
improve quality, it's rejected. The stacked model NEVER gets worse.

Usage::

    from model_osmosis.stacking import SafeStacker

    safe = SafeStacker(eval_data=X, eval_fn=my_scorer)
    safe.add(code_model, name="code")       # ✅ accepted: +12% quality
    safe.add(math_model, name="math")       # ✅ accepted: +8% quality
    safe.add(bad_model, name="junk")        # ❌ rejected: -3% quality
    safe.add(chat_model, name="chat")       # ✅ accepted: +5% quality

    result = safe.build()  # → StackedModel with only the good experts
"""

from __future__ import annotations

import copy
from typing import Any, Callable, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn

from model_osmosis.stacking.stacked_model import StackedModel
from model_osmosis.stacking.stacker import ModelStacker


class SafeStacker:
    """Guaranteed-better stacking: every model added must improve quality.

    Uses a validation-gated approach:
    1. Measure current model quality on eval data
    2. Temporarily add the new expert
    3. Train gate with new expert included
    4. Measure quality again
    5. Keep only if quality improved (or stayed the same)

    This guarantees the stacked model NEVER gets worse.

    Args:
        eval_data: Required evaluation data for quality measurement.
        eval_fn: Custom evaluation function (model, data) → score.
            Higher = better. If None, uses default output-diversity score.
        min_improvement: Minimum quality improvement to accept a new expert.
            0.0 = accept if not worse, 0.01 = must be at least 1% better.
        top_k: Number of experts to activate per input.
        verbose: Print progress.
    """

    def __init__(
        self,
        eval_data: np.ndarray,
        eval_fn: Optional[Callable[[nn.Module, np.ndarray], float]] = None,
        min_improvement: float = 0.0,
        top_k: int = 2,
        verbose: bool = True,
    ) -> None:
        self.eval_data = eval_data
        self.eval_fn = eval_fn or self._default_eval
        self.min_improvement = min_improvement
        self.top_k = top_k
        self.verbose = verbose

        self._experts: List[nn.Module] = []
        self._names: List[str] = []
        self._scores: List[float] = []  # quality after each addition
        self._current_model: Optional[StackedModel] = None
        self._current_score: float = float("-inf")
        self._stacker = ModelStacker(verbose=False)

        # Track accepted and rejected models
        self.accepted: List[str] = []
        self.rejected: List[Tuple[str, float, float]] = []  # (name, attempted_score, current_score)

    def add(
        self,
        model: nn.Module,
        name: Optional[str] = None,
    ) -> bool:
        """Try to add a model. Only accepts if it improves quality.

        Args:
            model: The model to try adding.
            name: Optional name for the expert.

        Returns:
            True if the model was accepted, False if rejected.
        """
        name = name or f"expert_{len(self._experts)}"

        if self.verbose:
            print(f"\n[SafeStacker] Trying to add '{name}'...")

        # If this is the first model, always accept
        if len(self._experts) == 0:
            self._experts.append(copy.deepcopy(model))
            self._names.append(name)
            self._current_model = self._build_model()
            self._current_score = self._evaluate(self._current_model)
            self._scores.append(self._current_score)
            self.accepted.append(name)

            if self.verbose:
                print(f"[SafeStacker] ✅ Accepted '{name}' (first expert, "
                      f"score: {self._current_score:.6f})")
            return True

        # Try adding this expert
        trial_experts = self._experts + [copy.deepcopy(model)]
        trial_names = self._names + [name]

        trial_model = self._build_model(trial_experts, trial_names)
        trial_score = self._evaluate(trial_model)

        improvement = trial_score - self._current_score

        if self.verbose:
            print(f"[SafeStacker]   Current score: {self._current_score:.6f}")
            print(f"[SafeStacker]   With '{name}': {trial_score:.6f}")
            print(f"[SafeStacker]   Improvement: {improvement:+.6f}")

        if improvement >= self.min_improvement:
            # Accept!
            self._experts = trial_experts
            self._names = trial_names
            self._current_model = trial_model
            self._current_score = trial_score
            self._scores.append(self._current_score)
            self.accepted.append(name)

            if self.verbose:
                print(f"[SafeStacker] ✅ Accepted '{name}' "
                      f"({improvement:+.6f} improvement)")
            return True
        else:
            # Reject!
            self.rejected.append((name, trial_score, self._current_score))

            if self.verbose:
                print(f"[SafeStacker] ❌ Rejected '{name}' "
                      f"(would decrease quality by {-improvement:.6f})")
            return False

    def build(self) -> StackedModel:
        """Return the final stacked model with only accepted experts.

        Returns:
            The optimized StackedModel.
        """
        if self._current_model is None:
            raise ValueError("No models added yet. Call add() first.")

        if self.verbose:
            print(f"\n[SafeStacker] Final model: {len(self._experts)} experts")
            print(f"[SafeStacker]   Accepted: {self.accepted}")
            if self.rejected:
                print(f"[SafeStacker]   Rejected: {[r[0] for r in self.rejected]}")
            print(f"[SafeStacker]   Final score: {self._current_score:.6f}")

        return self._current_model

    def _build_model(
        self,
        experts: Optional[List[nn.Module]] = None,
        names: Optional[List[str]] = None,
    ) -> StackedModel:
        """Build a StackedModel from current experts."""
        experts = experts or self._experts
        names = names or self._names

        return self._stacker.stack(
            models=experts,
            sample_data=self.eval_data,
            expert_names=names,
            top_k=min(self.top_k, len(experts)),
            train_gate=True,
            gate_epochs=30,
        )

    def _evaluate(self, model: StackedModel) -> float:
        """Evaluate model quality."""
        return self.eval_fn(model, self.eval_data)

    @staticmethod
    def _default_eval(model: nn.Module, data: np.ndarray) -> float:
        """Default evaluation: output stability + diversity.

        Measures:
        1. Output consistency (low variance across similar inputs)
        2. Output range utilization (uses full output space)
        3. Gradient of confidence (higher confidence = better)
        """
        x = torch.from_numpy(data).float()

        model.eval()
        with torch.no_grad():
            output = model(x)

        # Score 1: Output stability — low variance is good
        stability = -torch.var(output).item()

        # Score 2: Output range — using more of the output space is good
        out_range = (output.max() - output.min()).item()

        # Score 3: Output differentiation — different inputs should give
        # different outputs (the model is actually processing inputs)
        if len(output) > 1:
            # Compare first half to second half
            half = len(output) // 2
            diff = torch.mean(
                torch.abs(output[:half] - output[half:half*2])
            ).item()
        else:
            diff = 0.0

        # Combined score (all positive = better)
        score = stability * 0.1 + out_range * 0.3 + diff * 0.6

        return score

    def summary(self) -> str:
        """Human-readable summary of the stacking process."""
        lines = [
            f"SafeStacker Summary:",
            f"  Models tried: {len(self.accepted) + len(self.rejected)}",
            f"  Accepted: {len(self.accepted)} → {self.accepted}",
            f"  Rejected: {len(self.rejected)} → {[r[0] for r in self.rejected]}",
        ]
        if self._scores:
            lines.append(f"  Score progression: {[f'{s:.4f}' for s in self._scores]}")
            lines.append(f"  Total improvement: {self._scores[-1] - self._scores[0]:+.6f}")
        return "\n".join(lines)
