"""
ModelStacker — easy API to stack multiple models into one.

Usage::

    from model_osmosis.stacking import ModelStacker

    stacker = ModelStacker()
    result = stacker.stack(
        [model_a, model_b, model_c],
        sample_data=X,
        optimize=True,      # auto-prune + train confidence
    )
"""

from __future__ import annotations

import copy
from typing import Any, 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.utils.determinism import set_global_seed


class ModelStacker:
    """Stack multiple models into one mega-model.

    Every model is preserved in full as an "expert". A sparse Top-K
    gating network routes inputs to only the best K experts.

    Args:
        output_dim: Output dimension. Auto-detected if None.
        verbose: Print progress.
    """

    def __init__(
        self,
        output_dim: Optional[int] = None,
        verbose: bool = True,
    ) -> None:
        self.output_dim = output_dim
        self.verbose = verbose

    def stack(
        self,
        models: List[nn.Module],
        input_dim: Optional[int] = None,
        output_dim: Optional[int] = None,
        sample_data: Optional[np.ndarray] = None,
        expert_names: Optional[List[str]] = None,
        top_k: int = 2,
        train_gate: bool = True,
        gate_epochs: int = 50,
        gate_lr: float = 1e-3,
        optimize: bool = False,
        prune_threshold: float = 0.05,
        save_to: Optional[str] = None,
        model_name: Optional[str] = None,
        seed: int = 42,
    ) -> StackedModel:
        """Stack multiple models into a single StackedModel.

        Args:
            models: List of PyTorch nn.Module models to stack.
            input_dim: Input feature dimension. Auto-detected if not provided.
            output_dim: Output dimension. Auto-detected if not provided.
            sample_data: Sample input data for auto-detection and training.
            expert_names: Optional names for each expert.
            top_k: Number of experts to activate per input (sparse routing).
            train_gate: Whether to train the gating network.
            gate_epochs: Epochs to train the gate.
            gate_lr: Learning rate for gate training.
            optimize: Auto-optimize after stacking (prune + confidence training).
            prune_threshold: Min contribution to keep an expert (0.0-1.0).
            save_to: Folder to auto-save the model to (e.g., 'models').
            model_name: Name prefix for saved files.

        Returns:
            A StackedModel containing all experts + a learned gate.
        """
        n = len(models)
        if n == 0:
            raise ValueError("Must provide at least one model")
        set_global_seed(seed)

        if self.verbose:
            print(f"[Stacker] Stacking {n} models (top-{min(top_k, n)} routing)...")

        # Auto-detect dimensions
        expert_input_dims, expert_output_dims = self._detect_dims(
            models, sample_data, input_dim, output_dim
        )

        # Resolve common input/output dims
        resolved_input_dim = input_dim or (
            sample_data.shape[-1] if sample_data is not None else expert_input_dims[0]
        )
        resolved_output_dim = output_dim or self.output_dim or max(expert_output_dims)

        if self.verbose:
            print(f"[Stacker] Common input dim: {resolved_input_dim}")
            print(f"[Stacker] Common output dim: {resolved_output_dim}")
            for i, model in enumerate(models):
                name = expert_names[i] if expert_names else f"expert_{i}"
                params = sum(p.numel() for p in model.parameters())
                in_dim = expert_input_dims[i]
                out_dim = expert_output_dims[i]
                print(f"[Stacker]   {name}: {params:,} params (in={in_dim}, out={out_dim})")

        # Deep copy all models to avoid mutating originals
        experts = [copy.deepcopy(m) for m in models]
        for e in experts:
            e.eval()

        # Create the stacked model
        stacked = StackedModel(
            experts=experts,
            input_dim=resolved_input_dim,
            output_dim=resolved_output_dim,
            top_k=min(top_k, n),
            expert_input_dims=expert_input_dims,
            expert_output_dims=expert_output_dims,
            expert_names=expert_names,
            use_confidence=True,
        )

        if self.verbose:
            print(f"[Stacker] Created StackedModel:")
            print(f"[Stacker]   Total parameters: {stacked.total_params():,}")
            print(f"[Stacker]   Active experts per input: {stacked.top_k}/{n}")

        # Train the gating network
        if train_gate and sample_data is not None and n > 1:
            if self.verbose:
                print(f"[Stacker] Training gate + confidence ({gate_epochs} epochs)...")
            self._train_gate(stacked, sample_data, gate_epochs, gate_lr)

        # Post-stacking optimization
        if optimize and sample_data is not None:
            from model_osmosis.stacking.optimizer import StackOptimizer
            optimizer = StackOptimizer(verbose=self.verbose)
            stacked = optimizer.optimize(
                stacked,
                eval_data=sample_data,
                prune_threshold=prune_threshold,
            )

        # Auto-save
        if save_to:
            name = model_name or "model"
            saved_path = stacked.save(save_to, name=name)
            if self.verbose:
                print(f"[Stacker] Saved to: {saved_path}")

        if self.verbose:
            print("[Stacker] ✓ Stacking complete!")

        return stacked

    def _detect_dims(
        self,
        models: List[nn.Module],
        sample_data: Optional[np.ndarray],
        input_dim: Optional[int],
        output_dim: Optional[int],
    ) -> Tuple[List[int], List[int]]:
        """Auto-detect input/output dimensions for each expert."""
        expert_input_dims = []
        expert_output_dims = []

        for model in models:
            in_dim, out_dim = self._probe_model_dims(model, sample_data, input_dim)
            expert_input_dims.append(in_dim)
            expert_output_dims.append(out_dim)

        return expert_input_dims, expert_output_dims

    def _probe_model_dims(
        self,
        model: nn.Module,
        sample_data: Optional[np.ndarray],
        input_dim: Optional[int],
    ) -> Tuple[int, int]:
        """Probe a model to find its input/output dimensions."""
        in_dim = input_dim
        out_dim = None

        linear_layers = []
        for m in model.modules():
            if isinstance(m, nn.Linear):
                linear_layers.append(m)

        if linear_layers:
            if in_dim is None:
                in_dim = linear_layers[0].in_features
            out_dim = linear_layers[-1].out_features

        if sample_data is not None and in_dim is not None:
            try:
                probe_input = torch.from_numpy(sample_data[:2]).float()
                if probe_input.shape[-1] == in_dim:
                    model.eval()
                    with torch.no_grad():
                        probe_out = model(probe_input)
                    out_dim = probe_out.shape[-1]
            except Exception:
                pass

        if in_dim is None:
            raise ValueError(
                f"Cannot auto-detect input dim for {type(model).__name__}. "
                f"Provide input_dim explicitly."
            )
        if out_dim is None:
            raise ValueError(
                f"Cannot auto-detect output dim for {type(model).__name__}. "
                f"Provide output_dim explicitly."
            )

        return in_dim, out_dim

    def _train_gate(
        self,
        stacked: StackedModel,
        sample_data: np.ndarray,
        epochs: int,
        lr: float,
    ) -> None:
        """Train the gating network and confidence heads.

        Strategy: Train the gate to produce outputs that match the
        consensus of all experts. This teaches the gate to find
        the optimal routing per input.
        """
        x = torch.from_numpy(sample_data).float()

        # Get each expert's output (frozen)
        expert_outputs = []
        for i in range(stacked.num_experts):
            stacked.experts[i].eval()
            with torch.no_grad():
                expert_in = stacked.input_projectors[i](x)
                out = stacked.experts[i](expert_in)
                projected = stacked.output_projectors[i](out)
                expert_outputs.append(projected)

        # Target: consensus of all experts
        target = torch.stack(expert_outputs).mean(dim=0)

        # Train gate + projectors + confidence heads
        trainable = (
            list(stacked.gate.parameters())
            + list(stacked.input_projectors.parameters())
            + list(stacked.output_projectors.parameters())
        )
        if stacked.confidence_heads is not None:
            trainable += list(stacked.confidence_heads.parameters())

        optimizer = torch.optim.Adam(trainable, lr=lr)

        stacked.train()
        for epoch in range(epochs):
            optimizer.zero_grad()
            output = stacked(x)
            loss = nn.functional.mse_loss(output, target)
            loss.backward()
            optimizer.step()

        stacked.eval()
        if self.verbose:
            print(f"[Stacker]   Gate training final loss: {loss.item():.6f}")
