"""
MultiModelMerger — the main API for merging multiple models into one.

Usage::

    from model_osmosis.merging import MultiModelMerger

    merger = MultiModelMerger(strategy="auto")
    result = merger.merge(
        models=[model_a, model_b, model_c],
        eval_data=X_eval,
    )
"""

from __future__ import annotations

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

import numpy as np

from model_osmosis.core.config import OsmosisConfig
from model_osmosis.core.exceptions import OsmosisError
from model_osmosis.merging.evaluator import MergeEvaluator
from model_osmosis.merging.streaming import StreamingParameterLoader
from model_osmosis.merging.strategies import (
    MergeStrategy,
    WeightedAverageStrategy,
    SLERPStrategy,
    TIESStrategy,
    get_strategy,
)


class MultiModelMerger:
    """Merge multiple models into one, efficiently and automatically.

    Features:
    - RAM-efficient: streams parameters layer-by-layer
    - Auto-optimize: finds best per-model weights via evaluation
    - Multiple strategies: weighted average, SLERP, TIES, or auto

    Args:
        strategy: Merge strategy name ('weighted_average', 'slerp', 'ties',
            'auto') or a MergeStrategy instance.
        model_weights: Optional per-model importance weights. If None,
            uses equal weighting (or auto-optimized if strategy='auto').
        config: Optional OsmosisConfig for growth/absorption settings.
        verbose: Print progress information.
    """

    def __init__(
        self,
        strategy: Union[str, MergeStrategy] = "weighted_average",
        model_weights: Optional[List[float]] = None,
        config: Optional[OsmosisConfig] = None,
        verbose: bool = True,
    ) -> None:
        if isinstance(strategy, str):
            self._strategy_name = strategy
            if strategy != "auto":
                self._strategy = get_strategy(strategy)
            else:
                self._strategy = None  # resolved during merge
        else:
            self._strategy_name = strategy.name
            self._strategy = strategy

        self.model_weights = model_weights
        self.config = config or OsmosisConfig(verbose=verbose)
        self.verbose = verbose
        self._loader = StreamingParameterLoader()
        self._evaluator = MergeEvaluator()

    def merge(
        self,
        models: List[Any],
        eval_data: Optional[np.ndarray] = None,
        base_model: Optional[Any] = None,
        inplace: bool = False,
    ) -> Any:
        """Merge multiple models into one.

        Args:
            models: List of models to merge. Can be:
                - Live PyTorch nn.Module objects
                - File paths to saved models
                - Dicts (state_dicts)
            eval_data: Optional evaluation data for auto-optimization.
                Required if strategy='auto'.
            base_model: Optional base model to use as the merge target.
                If None, uses a deep copy of the first model.
            inplace: If True, modify the first model (or base_model) in-place
                instead of deep-copying. Saves memory for large models.

        Returns:
            The merged model.
        """
        n = len(models)
        if n == 0:
            raise ValueError("Must provide at least one model")
        if n == 1:
            if inplace:
                return models[0]
            return copy.deepcopy(models[0]) if not isinstance(models[0], (str, dict)) else models[0]

        if self.verbose:
            print(f"[Merger] Merging {n} models with strategy: {self._strategy_name}")

        # Determine weights
        weights = self.model_weights
        strategy = self._strategy

        if self._strategy_name == "auto" and eval_data is not None:
            # Auto mode: try all strategies and find best weights
            strategy, weights = self._auto_optimize(models, eval_data)
        elif weights is None:
            weights = [1.0 / n] * n

        if strategy is None:
            strategy = WeightedAverageStrategy()

        if self.verbose:
            print(f"[Merger] Weights: {[f'{w:.3f}' for w in weights]}")
            print(f"[Merger] Strategy: {strategy.name}")

        # Perform the merge
        merged = self._merge_parameters(models, weights, strategy, base_model, inplace)

        if self.verbose:
            if hasattr(merged, "parameters"):
                total = sum(p.numel() for p in merged.parameters())
                print(f"[Merger] ✓ Merged model: {total:,} parameters")

        return merged

    def _auto_optimize(
        self,
        models: List[Any],
        eval_data: np.ndarray,
    ) -> Tuple[MergeStrategy, List[float]]:
        """Try multiple strategies and pick the best one."""
        if self.verbose:
            print("[Merger] Auto-optimizing: trying all strategies...")

        best_strategy = WeightedAverageStrategy()
        best_weights = [1.0 / len(models)] * len(models)
        best_score = float("-inf")

        strategies_to_try = [
            WeightedAverageStrategy(),
            SLERPStrategy(),
            TIESStrategy(density=0.5),
        ]

        for strat in strategies_to_try:
            weights, score = self._evaluator.find_best_weights(
                models=models,
                eval_data=eval_data,
                strategy=strat,
                n_trials=8,
                verbose=False,
            )
            if self.verbose:
                print(
                    f"  Strategy '{strat.name}': best score = {score:.6f}"
                )
            if score > best_score:
                best_score = score
                best_strategy = strat
                best_weights = weights

        if self.verbose:
            print(f"  → Winner: '{best_strategy.name}' (score: {best_score:.6f})")

        return best_strategy, best_weights

    def _merge_parameters(
        self,
        models: List[Any],
        weights: List[float],
        strategy: MergeStrategy,
        base_model: Optional[Any],
        inplace: bool = False,
    ) -> Any:
        """Merge parameters from multiple models layer-by-layer.

        RAM-efficient: builds per-model parameter dicts once, then
        merges by parameter name with O(1) lookup.
        """
        import torch
        import torch.nn as nn

        # Build parameter dicts for all models (O(n) per model, done once)
        all_params: List[Dict[str, np.ndarray]] = []
        for model in models:
            param_dict = dict(self._loader.stream_parameters(model))
            all_params.append(param_dict)

        param_names = list(all_params[0].keys())

        if self.verbose:
            print(f"[Merger] Merging {len(param_names)} parameter tensors...")

        # For each parameter, collect from all models and merge
        merged_state: Dict[str, np.ndarray] = {}

        for name in param_names:
            tensors = []
            for params in all_params:
                if name in params:
                    tensors.append(params[name])
                else:
                    # Fallback: use first model's value
                    tensors.append(all_params[0][name])

            # Skip merge if all tensors are identical (performance optimization)
            if all(np.array_equal(t, tensors[0]) for t in tensors[1:]):
                merged_state[name] = tensors[0].copy()
                continue

            # Merge this parameter
            if all(t.shape == tensors[0].shape for t in tensors):
                merged_state[name] = strategy.merge_tensors(tensors, weights)
            else:
                # Shape mismatch — use the first model's weights as fallback
                merged_state[name] = tensors[0].copy()
                if self.verbose:
                    print(
                        f"  [Merger] Shape mismatch for '{name}': "
                        f"{[t.shape for t in tensors]} — using first model"
                    )

        # Create the output model
        if base_model is not None:
            result = base_model if inplace else copy.deepcopy(base_model)
        else:
            if isinstance(models[0], nn.Module):
                result = models[0] if inplace else copy.deepcopy(models[0])
            else:
                raise ValueError(
                    "Cannot infer output model architecture. "
                    "Provide a base_model argument."
                )

        # Load merged parameters into the output model
        if isinstance(result, nn.Module):
            state_dict = result.state_dict()
            for name, arr in merged_state.items():
                if name in state_dict:
                    state_dict[name] = torch.from_numpy(arr)
            result.load_state_dict(state_dict)
            result.eval()

        return result
