"""
StackOptimizer — post-stacking optimization to make the stacked model
genuinely better than any individual expert.

Features:
- Expert contribution scoring: measure each expert's value
- Auto-pruning: remove experts that barely contribute
- Conflict resolution tuning: train confidence heads
- Expert quantization: compress to int8 for 4x less RAM
"""

from __future__ import annotations

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

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

from model_osmosis.stacking.stacked_model import StackedModel


class StackOptimizer:
    """Optimize a StackedModel after stacking.

    Makes the stacked model genuinely better by:
    1. Scoring each expert's contribution
    2. Removing weak experts (pruning)
    3. Training confidence heads for conflict resolution
    4. Quantizing expert weights for RAM efficiency

    Args:
        verbose: Print progress.
    """

    def __init__(self, verbose: bool = True) -> None:
        self.verbose = verbose

    def optimize(
        self,
        model: StackedModel,
        eval_data: np.ndarray,
        prune_threshold: float = 0.05,
        train_confidence: bool = True,
        confidence_epochs: int = 30,
        quantize: bool = False,
    ) -> StackedModel:
        """Full optimization pipeline.

        Args:
            model: The StackedModel to optimize.
            eval_data: Evaluation data for measuring quality.
            prune_threshold: Remove experts contributing less than this %.
            train_confidence: Whether to train confidence heads.
            confidence_epochs: Epochs for confidence training.
            quantize: Whether to quantize expert weights to int8.

        Returns:
            The optimized StackedModel.
        """
        if self.verbose:
            print("[Optimizer] Starting optimization...")
            print(f"[Optimizer] Initial: {model.num_experts} experts, "
                  f"{model.total_params():,} params")

        # Step 1: Score experts
        scores = self.score_experts(model, eval_data)
        if self.verbose:
            for name, score in scores.items():
                print(f"[Optimizer]   {name}: contribution = {score:.4f}")

        # Step 2: Prune weak experts
        model = self.prune_experts(model, scores, threshold=prune_threshold)

        # Step 3: Train confidence heads
        if train_confidence and model.use_confidence:
            self.train_confidence_heads(model, eval_data, epochs=confidence_epochs)

        # Step 4: Quantize (optional)
        if quantize:
            model = self.quantize_experts(model)

        if self.verbose:
            print(f"[Optimizer] Final: {model.num_experts} experts, "
                  f"{model.total_params():,} params")
            print("[Optimizer] ✓ Optimization complete!")

        return model

    def score_experts(
        self,
        model: StackedModel,
        eval_data: np.ndarray,
    ) -> Dict[str, float]:
        """Measure each expert's contribution to output quality.

        Method: Compare the full model's output to the output without
        each expert (leave-one-out). The expert whose removal causes
        the biggest quality drop has the highest contribution score.

        Returns:
            Dict mapping expert name to contribution score (0.0 = useless, 1.0 = critical).
        """
        x = torch.from_numpy(eval_data).float()

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

        scores = {}
        for i, name in enumerate(model.expert_names):
            # Temporarily disable this expert by zeroing its gate weight
            degraded_output = self._forward_without_expert(model, x, exclude_idx=i)
            # Measure quality drop
            mse_full = 0.0  # baseline
            mse_degraded = torch.mean(
                (full_output - degraded_output) ** 2
            ).item()
            # Higher MSE = this expert matters more
            scores[name] = mse_degraded

        # Normalize to [0, 1]
        max_score = max(scores.values()) if scores else 1.0
        if max_score > 0:
            scores = {k: v / max_score for k, v in scores.items()}

        return scores

    def _forward_without_expert(
        self,
        model: StackedModel,
        x: torch.Tensor,
        exclude_idx: int,
    ) -> torch.Tensor:
        """Forward pass with one expert excluded."""
        batch_size = x.size(0)

        # Get gate logits manually
        if x.dim() > 2:
            flat = x.view(x.size(0), -1)
        else:
            flat = x

        logits = model.gate.gate(flat)
        # Mask out the excluded expert
        logits[:, exclude_idx] = float("-inf")

        # Top-K from remaining experts
        top_k = min(model.top_k, model.num_experts - 1)
        top_k_values, top_k_indices = torch.topk(logits, top_k, dim=-1)
        top_k_weights = torch.softmax(top_k_values, dim=-1)

        # Compute output
        output = torch.zeros(batch_size, model.output_dim, device=x.device)

        for batch_idx in range(batch_size):
            weighted_sum = torch.zeros(model.output_dim, device=x.device)
            total_weight = 0.0

            for k in range(top_k):
                expert_idx = top_k_indices[batch_idx, k].item()
                gate_w = top_k_weights[batch_idx, k]

                expert_input = model.input_projectors[expert_idx](
                    x[batch_idx:batch_idx+1]
                )
                expert_out = model.experts[expert_idx](expert_input)
                projected = model.output_projectors[expert_idx](expert_out)

                weighted_sum += gate_w * projected.squeeze(0)
                total_weight += gate_w.item()

            if total_weight > 1e-8:
                output[batch_idx] = weighted_sum / total_weight

        return output

    def prune_experts(
        self,
        model: StackedModel,
        scores: Dict[str, float],
        threshold: float = 0.05,
    ) -> StackedModel:
        """Remove experts with contribution below threshold.

        Args:
            model: The StackedModel.
            scores: Expert contribution scores from score_experts().
            threshold: Minimum contribution to keep (0.0-1.0).

        Returns:
            A new StackedModel with weak experts removed.
        """
        # Find experts to keep (always keep at least 1)
        keep_indices = []
        keep_names = []
        pruned_names = []

        for i, name in enumerate(model.expert_names):
            if scores.get(name, 0.0) >= threshold:
                keep_indices.append(i)
                keep_names.append(name)
            else:
                pruned_names.append(name)

        # Always keep at least 1 expert
        if not keep_indices:
            best_idx = max(range(model.num_experts),
                          key=lambda i: scores.get(model.expert_names[i], 0.0))
            keep_indices = [best_idx]
            keep_names = [model.expert_names[best_idx]]
            pruned_names = [n for n in model.expert_names if n != keep_names[0]]

        if len(pruned_names) == 0:
            if self.verbose:
                print("[Optimizer] No experts to prune — all above threshold")
            return model

        if self.verbose:
            print(f"[Optimizer] Pruning {len(pruned_names)} weak experts: "
                  f"{pruned_names}")
            print(f"[Optimizer] Keeping {len(keep_indices)} experts: {keep_names}")

        # Build new model with only kept experts
        new_experts = [copy.deepcopy(model.experts[i]) for i in keep_indices]
        new_input_projs = [model.input_projectors[i] for i in keep_indices]
        new_output_projs = [model.output_projectors[i] for i in keep_indices]

        # Detect expert dims
        expert_input_dims = []
        expert_output_dims = []
        for i in keep_indices:
            in_proj = model.input_projectors[i]
            out_proj = model.output_projectors[i]
            # Get dims from the projector layers
            for m in in_proj.modules():
                if isinstance(m, nn.Linear):
                    expert_input_dims.append(m.out_features)
                    break
            else:
                expert_input_dims.append(model.input_dim)

            for m in out_proj.modules():
                if isinstance(m, nn.Linear):
                    expert_output_dims.append(m.in_features)
                    break
            else:
                expert_output_dims.append(model.output_dim)

        pruned = StackedModel(
            experts=new_experts,
            input_dim=model.input_dim,
            output_dim=model.output_dim,
            top_k=min(model.top_k, len(new_experts)),
            expert_input_dims=expert_input_dims,
            expert_output_dims=expert_output_dims,
            expert_names=keep_names,
            use_confidence=model.use_confidence,
        )

        return pruned

    def train_confidence_heads(
        self,
        model: StackedModel,
        eval_data: np.ndarray,
        epochs: int = 30,
        lr: float = 1e-3,
    ) -> None:
        """Train confidence heads so experts learn when they're reliable.

        Each expert's confidence head learns to output high confidence
        when that expert produces a good output, and low confidence
        when it produces a bad one (compared to the consensus).
        """
        if model.confidence_heads is None:
            return

        x = torch.from_numpy(eval_data).float()

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

        # Consensus = mean of all expert outputs
        consensus = torch.stack(expert_outputs).mean(dim=0)

        # Train each confidence head
        conf_params = list(model.confidence_heads.parameters())
        optimizer = torch.optim.Adam(conf_params, lr=lr)

        model.confidence_heads.train()
        for epoch in range(epochs):
            optimizer.zero_grad()
            total_loss = torch.tensor(0.0)

            for i in range(model.num_experts):
                # Target confidence: how close is this expert to consensus?
                with torch.no_grad():
                    error = torch.mean(
                        (expert_outputs[i] - consensus) ** 2, dim=-1, keepdim=True
                    )
                    # Low error = high confidence target
                    max_error = error.max().clamp(min=1e-8)
                    target_conf = 1.0 - (error / max_error)

                pred_conf = model.confidence_heads[i](expert_outputs[i].detach())
                loss = nn.functional.mse_loss(pred_conf, target_conf)
                total_loss = total_loss + loss

            total_loss.backward()
            optimizer.step()

        model.confidence_heads.eval()
        if self.verbose:
            print(f"[Optimizer] Confidence heads trained ({epochs} epochs, "
                  f"loss={total_loss.item():.6f})")

    def quantize_experts(self, model: StackedModel) -> StackedModel:
        """Quantize expert weights to int8 for ~4x RAM reduction.

        Uses PyTorch's dynamic quantization on expert Linear layers.
        Gate and confidence heads remain in fp32 for accuracy.
        """
        try:
            for i, expert in enumerate(model.experts):
                model.experts[i] = torch.ao.quantization.quantize_dynamic(
                    expert,
                    {nn.Linear},
                    dtype=torch.qint8,
                )

            if self.verbose:
                print("[Optimizer] Experts quantized to int8")
        except Exception as e:
            if self.verbose:
                print(f"[Optimizer] Quantization skipped: {e}")

        return model
