"""
Merge evaluator — auto-optimize merge ratios by evaluating output quality.
"""

from __future__ import annotations

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

import numpy as np

from model_osmosis.merging.strategies import MergeStrategy, get_strategy


class MergeEvaluator:
    """Evaluate merged model quality and find optimal merge weights.

    Given evaluation data, tries different weight configurations and
    picks the one that produces the best results.
    """

    def __init__(
        self,
        eval_fn: Optional[Callable[[Any, np.ndarray], float]] = None,
    ) -> None:
        """
        Args:
            eval_fn: Custom evaluation function that takes (model, eval_data)
                and returns a quality score (higher = better).
                If None, uses default MSE-based evaluation.
        """
        self.eval_fn = eval_fn

    def find_best_weights(
        self,
        models: List[Any],
        eval_data: np.ndarray,
        strategy: MergeStrategy,
        n_trials: int = 10,
        verbose: bool = False,
    ) -> Tuple[List[float], float]:
        """Find the best per-model weights via grid search.

        Args:
            models: List of source models.
            eval_data: Evaluation input data.
            strategy: The merge strategy to use.
            n_trials: Number of weight configurations to try.
            verbose: Print progress.

        Returns:
            (best_weights, best_score) — the optimal weight configuration
            and its quality score.
        """
        n_models = len(models)

        if n_models <= 1:
            return [1.0], 0.0

        # Generate candidate weight configurations
        candidates = self._generate_weight_candidates(n_models, n_trials)

        # Get reference outputs from each model
        model_outputs = []
        for model in models:
            out = self._get_outputs(model, eval_data)
            model_outputs.append(out)

        best_weights = candidates[0]
        best_score = float("-inf")

        for i, weights in enumerate(candidates):
            # Blend outputs according to weights
            blended = np.zeros_like(model_outputs[0], dtype=np.float64)
            for output, w in zip(model_outputs, weights):
                blended += w * output.astype(np.float64)
            blended = blended.astype(model_outputs[0].dtype)

            # Evaluate quality
            score = self._evaluate_blend(blended, model_outputs, weights)

            if score > best_score:
                best_score = score
                best_weights = weights

            if verbose:
                print(
                    f"  [Evaluator] Trial {i + 1}/{len(candidates)}: "
                    f"weights={[f'{w:.2f}' for w in weights]} "
                    f"score={score:.4f}"
                )

        if verbose:
            print(
                f"  [Evaluator] Best: weights={[f'{w:.2f}' for w in best_weights]} "
                f"score={best_score:.4f}"
            )

        return best_weights, best_score

    def _generate_weight_candidates(
        self, n_models: int, n_trials: int
    ) -> List[List[float]]:
        """Generate diverse weight configurations to try."""
        candidates = []

        # Always include uniform weights
        candidates.append([1.0 / n_models] * n_models)

        # Include "one model dominates" configurations
        for i in range(n_models):
            w = [0.1 / max(n_models - 1, 1)] * n_models
            w[i] = 0.9
            candidates.append(w)

        # Random diverse configurations
        rng = np.random.RandomState(42)
        for _ in range(max(0, n_trials - len(candidates))):
            raw = rng.dirichlet(np.ones(n_models))
            candidates.append(raw.tolist())

        return candidates[:n_trials]

    def _get_outputs(self, model: Any, data: np.ndarray) -> np.ndarray:
        """Get model outputs on evaluation data."""
        try:
            import torch
            import torch.nn as nn

            if isinstance(model, nn.Module):
                model.eval()
                x = torch.from_numpy(data).float()
                device = next(model.parameters()).device
                x = x.to(device)
                with torch.no_grad():
                    out = model(x)
                if isinstance(out, torch.Tensor):
                    return out.cpu().numpy()
                return out[0].cpu().numpy()
        except (ImportError, StopIteration):
            pass

        if hasattr(model, "predict_proba"):
            return model.predict_proba(data).astype(np.float32)
        if hasattr(model, "predict"):
            out = model.predict(data)
            if out.ndim == 1:
                out = out.reshape(-1, 1)
            return out.astype(np.float32)

        raise ValueError(f"Cannot get outputs from {type(model).__name__}")

    def _evaluate_blend(
        self,
        blended: np.ndarray,
        model_outputs: List[np.ndarray],
        weights: List[float],
    ) -> float:
        """Evaluate the quality of a blended output.

        Uses diversity-weighted consensus: higher score when the blend
        captures agreement between models while incorporating unique strengths.
        """
        # Measure agreement: how well does the blend correlate with each source?
        agreement_score = 0.0
        for output, w in zip(model_outputs, weights):
            # Use negative MSE as score (higher = better = less error)
            mse = np.mean((blended - output.astype(np.float64)) ** 2)
            agreement_score -= w * mse

        # Diversity bonus: reward blends that aren't just copying one model
        entropy = -sum(
            w * np.log(max(w, 1e-10)) for w in weights
        )
        diversity_bonus = entropy * 0.01

        return agreement_score + diversity_bonus
