"""
StackedModel — Optimized Mixture-of-Experts wrapper.

Features:
- Top-K sparse routing: only activate the best K experts per input
- Confidence heads: each expert reports how confident it is
- Sparse forward: skip non-selected experts entirely (saves RAM + compute)
- Supports quantized experts for reduced memory
"""

from __future__ import annotations

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

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class TopKGating(nn.Module):
    """Sparse Top-K gating network.

    Instead of softmax over ALL experts (dense), this selects only the
    top K experts and zeros out the rest. This means non-selected experts
    don't need to run at all.

    Args:
        input_dim: Input feature dimension.
        num_experts: Total number of experts.
        top_k: Number of experts to activate per input.
        noise_std: Noise added during training for load balancing.
    """

    def __init__(
        self,
        input_dim: int,
        num_experts: int,
        top_k: int = 2,
        noise_std: float = 0.1,
    ) -> None:
        super().__init__()
        self.num_experts = num_experts
        self.top_k = min(top_k, num_experts)
        self.noise_std = noise_std

        hidden = max(64, input_dim // 2)
        self.gate = nn.Sequential(
            nn.Linear(input_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, num_experts),
        )

    def forward(
        self, x: torch.Tensor, training: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute sparse gating weights.

        Args:
            x: Input tensor (batch, features).
            training: Whether to add noise for load balancing.

        Returns:
            (weights, indices):
                weights: (batch, top_k) — normalized weights for selected experts
                indices: (batch, top_k) — indices of selected experts
        """
        if x.dim() > 2:
            flat = x.view(x.size(0), -1)
        else:
            flat = x

        logits = self.gate(flat)

        # Add noise during training for load balancing
        if training and self.noise_std > 0:
            noise = torch.randn_like(logits) * self.noise_std
            logits = logits + noise

        # Select top-K experts
        top_k_values, top_k_indices = torch.topk(logits, self.top_k, dim=-1)

        # Softmax only over selected experts (sparse)
        top_k_weights = F.softmax(top_k_values, dim=-1)

        return top_k_weights, top_k_indices


class ConfidenceHead(nn.Module):
    """Per-expert confidence estimator.

    Each expert outputs not just a prediction but also a confidence score.
    When experts disagree, the more confident one gets more weight.
    """

    def __init__(self, output_dim: int) -> None:
        super().__init__()
        self.head = nn.Sequential(
            nn.Linear(output_dim, max(16, output_dim // 2)),
            nn.ReLU(),
            nn.Linear(max(16, output_dim // 2), 1),
            nn.Sigmoid(),
        )

    def forward(self, expert_output: torch.Tensor) -> torch.Tensor:
        """Returns confidence score in [0, 1] for this expert's output."""
        return self.head(expert_output)


class OutputProjector(nn.Module):
    """Projects an expert's output to a common dimension."""

    def __init__(self, in_dim: int, out_dim: int) -> None:
        super().__init__()
        if in_dim == out_dim:
            self.proj = nn.Identity()
        else:
            self.proj = nn.Linear(in_dim, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.proj(x)


class InputProjector(nn.Module):
    """Projects common input to an expert's expected input dimension."""

    def __init__(self, in_dim: int, out_dim: int) -> None:
        super().__init__()
        if in_dim == out_dim:
            self.proj = nn.Identity()
        else:
            self.proj = nn.Linear(in_dim, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.proj(x)


class StackedModel(nn.Module):
    """Optimized Mixture-of-Experts model.

    Every expert model is kept in full — nothing is averaged or discarded.
    A Top-K sparse gating network routes each input to only the best K
    experts, skipping the rest entirely.

    Args:
        experts: List of nn.Module models to stack.
        input_dim: Input feature dimension.
        output_dim: Common output dimension.
        top_k: Number of experts to activate per input (default: 2).
        expert_input_dims: Per-expert input dimensions for projectors.
        expert_output_dims: Per-expert output dimensions for projectors.
        expert_names: Optional names for each expert.
        use_confidence: Whether to use confidence heads for conflict resolution.
    """

    def __init__(
        self,
        experts: List[nn.Module],
        input_dim: int,
        output_dim: int,
        top_k: int = 2,
        expert_input_dims: Optional[List[int]] = None,
        expert_output_dims: Optional[List[int]] = None,
        expert_names: Optional[List[str]] = None,
        use_confidence: bool = True,
    ) -> None:
        super().__init__()

        self.num_experts = len(experts)
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.top_k = min(top_k, self.num_experts)
        self.use_confidence = use_confidence
        self.expert_names = expert_names or [
            f"expert_{i}" for i in range(self.num_experts)
        ]

        # Store experts as ModuleList (preserves all parameters)
        self.experts = nn.ModuleList(experts)

        # Input projectors
        self.input_projectors = nn.ModuleList()
        if expert_input_dims:
            for edim in expert_input_dims:
                self.input_projectors.append(InputProjector(input_dim, edim))
        else:
            for _ in experts:
                self.input_projectors.append(InputProjector(input_dim, input_dim))

        # Output projectors
        self.output_projectors = nn.ModuleList()
        if expert_output_dims:
            for edim in expert_output_dims:
                self.output_projectors.append(OutputProjector(edim, output_dim))
        else:
            for _ in experts:
                self.output_projectors.append(OutputProjector(output_dim, output_dim))

        # Sparse Top-K gating
        self.gate = TopKGating(
            input_dim=input_dim,
            num_experts=self.num_experts,
            top_k=self.top_k,
        )

        # Per-expert confidence heads
        if use_confidence:
            self.confidence_heads = nn.ModuleList(
                [ConfidenceHead(output_dim) for _ in range(self.num_experts)]
            )
        else:
            self.confidence_heads = None

        # Track expert usage statistics
        self.register_buffer(
            "_expert_usage", torch.zeros(self.num_experts),
        )
        self._forward_count = 0

        # Expert offloader (disabled by default)
        self._offloader = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Sparse forward pass — only top-K experts run.

        1. Gate selects top-K experts per input
        2. Only those K experts are loaded to GPU (if offloading enabled)
        3. Confidence weighting resolves conflicts
        4. Weighted sum produces final output
        """
        batch_size = x.size(0)

        # Get sparse gating: (batch, top_k) weights and indices
        gate_weights, gate_indices = self.gate(x, training=self.training)

        # Initialize output accumulator
        output = torch.zeros(batch_size, self.output_dim, device=x.device)

        # Track which experts are selected (for usage statistics)
        for idx in gate_indices.reshape(-1):
            self._expert_usage[idx.item()] += 1
        self._forward_count += batch_size

        # Run ONLY selected experts (sparse — this is the speed optimization)
        unique_experts = gate_indices.unique().tolist()

        # Offload: load only selected experts to compute device
        if self._offloader is not None:
            self._offloader.load_experts(self.experts, unique_experts)

        expert_outputs_cache: Dict[int, torch.Tensor] = {}
        expert_confidence_cache: Dict[int, torch.Tensor] = {}

        for expert_idx in unique_experts:
            # Only compute for inputs that selected this expert
            mask = (gate_indices == expert_idx).any(dim=-1)
            if not mask.any():
                continue

            # Get the appropriate device for this expert
            try:
                expert_device = next(self.experts[expert_idx].parameters()).device
            except StopIteration:
                # Quantized experts may have no regular parameters
                expert_device = x.device
            expert_x = x[mask].to(expert_device)

            expert_input = self.input_projectors[expert_idx](expert_x)
            expert_out = self.experts[expert_idx](expert_input)
            projected = self.output_projectors[expert_idx](expert_out)

            # Move result back to main device if needed
            projected = projected.to(x.device)

            expert_outputs_cache[expert_idx] = (projected, mask)

            if self.use_confidence and self.confidence_heads is not None:
                conf = self.confidence_heads[expert_idx](projected)
                expert_confidence_cache[expert_idx] = conf

        # Combine outputs using gate weights × confidence
        for batch_idx in range(batch_size):
            weighted_sum = torch.zeros(self.output_dim, device=x.device)
            total_weight = 0.0

            for k in range(self.top_k):
                expert_idx = gate_indices[batch_idx, k].item()
                gate_w = gate_weights[batch_idx, k]

                if expert_idx in expert_outputs_cache:
                    out_tensor, mask = expert_outputs_cache[expert_idx]
                    # Find this batch item's position in the masked output
                    mask_indices = mask.nonzero(as_tuple=True)[0]
                    local_idx = (mask_indices == batch_idx).nonzero(as_tuple=True)[0]

                    if len(local_idx) > 0:
                        expert_out = out_tensor[local_idx[0]]

                        # Apply confidence weighting
                        if (
                            self.use_confidence
                            and expert_idx in expert_confidence_cache
                        ):
                            conf = expert_confidence_cache[expert_idx][local_idx[0]]
                            w = gate_w * conf.squeeze()
                        else:
                            w = gate_w

                        weighted_sum += w * expert_out
                        total_weight += w.item()

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

        return output

    def enable_offloading(
        self,
        device: str = "cuda",
        cache_size: int = 2,
        verbose: bool = False,
    ) -> None:
        """Enable expert offloading for RAM savings.

        Moves all experts to CPU and only loads the selected top-K
        to the compute device (GPU) on demand.

        Args:
            device: Compute device for active experts (e.g., 'cuda').
            cache_size: Number of experts to keep on GPU (LRU cache).
            verbose: Print offloading activity.
        """
        from model_osmosis.stacking.offloader import ExpertOffloader

        self._offloader = ExpertOffloader(
            device=device,
            cache_size=max(cache_size, self.top_k),
            verbose=verbose,
        )
        # Move all experts to CPU
        self._offloader.offload_all_to_cpu(self.experts)
        # Keep gate + projectors + confidence on compute device
        self.gate.to(device)
        self.input_projectors.to(device)
        self.output_projectors.to(device)
        if self.confidence_heads is not None:
            self.confidence_heads.to(device)

    def disable_offloading(self) -> None:
        """Disable offloading — move all experts back to the current device."""
        self._offloader = None

    @property
    def offloading_enabled(self) -> bool:
        """Whether expert offloading is active."""
        return self._offloader is not None

    def memory_stats(self) -> Dict[str, Any]:
        """Get memory usage statistics."""
        if self._offloader is not None:
            return self._offloader.memory_stats(self.experts)
        total_params = self.total_params()
        return {
            "total_experts": self.num_experts,
            "cached_experts": self.num_experts,
            "total_params": total_params,
            "cached_params": total_params,
            "total_memory_mb": total_params * 4 / (1024 * 1024),
            "gpu_memory_mb": total_params * 4 / (1024 * 1024),
            "ram_savings_pct": 0.0,
        }

    def total_params(self) -> int:
        """Total parameters across all experts + gate + confidence."""
        return sum(p.numel() for p in self.parameters())

    def expert_params(self) -> List[Tuple[str, int]]:
        """Parameter count per expert."""
        result = []
        for name, expert in zip(self.expert_names, self.experts):
            count = sum(p.numel() for p in expert.parameters())
            result.append((name, count))
        return result

    def usage_stats(self) -> Dict[str, float]:
        """How often each expert is selected (after inference)."""
        if self._forward_count == 0:
            return {n: 0.0 for n in self.expert_names}
        return {
            name: (usage.item() / self._forward_count)
            for name, usage in zip(self.expert_names, self._expert_usage)
        }

    def summary(self) -> str:
        """Human-readable summary."""
        lines = [
            f"StackedModel (Sparse MoE) — {self.num_experts} experts, top-{self.top_k} routing:",
            f"  Input dim: {self.input_dim}",
            f"  Output dim: {self.output_dim}",
            f"  Confidence heads: {'Yes' if self.use_confidence else 'No'}",
        ]
        total_expert = 0
        for name, count in self.expert_params():
            lines.append(f"  {name}: {count:,} params")
            total_expert += count
        gate_params = sum(p.numel() for p in self.gate.parameters())
        lines.append(f"  Gate: {gate_params:,} params")
        lines.append(f"  Total: {self.total_params():,} params")
        lines.append(f"  Active per input: {self.top_k}/{self.num_experts} experts")
        return "\n".join(lines)

    def save(self, path: str, name: Optional[str] = None) -> str:
        """Save the stacked model to a folder.

        Files are saved directly in the folder (flat, no subfolders):
        - {name}.pt — model weights
        - {name}_config.json — architecture info
        - {name}_readme.txt — summary

        Args:
            path: Directory to save into (created if it doesn't exist).
            name: Model name for file prefixes. Defaults to 'model'.

        Returns:
            The full path to the saved .pt file.
        """
        import json
        import os

        name = name or "model"
        os.makedirs(path, exist_ok=True)

        # Gather expert layer shapes for reconstruction
        expert_architectures = []
        for expert in self.experts:
            layers_info = []
            for m in expert.modules():
                if isinstance(m, nn.Linear):
                    layers_info.append({
                        "type": "Linear",
                        "in": m.in_features,
                        "out": m.out_features,
                    })
                elif isinstance(m, nn.ReLU):
                    layers_info.append({"type": "ReLU"})
            expert_architectures.append(layers_info)

        # Gather projector dims
        expert_input_dims = []
        expert_output_dims = []
        for i in range(self.num_experts):
            for m in self.input_projectors[i].modules():
                if isinstance(m, nn.Linear):
                    expert_input_dims.append(m.out_features)
                    break
            else:
                expert_input_dims.append(self.input_dim)
            for m in self.output_projectors[i].modules():
                if isinstance(m, nn.Linear):
                    expert_output_dims.append(m.in_features)
                    break
            else:
                expert_output_dims.append(self.output_dim)

        config = {
            "num_experts": self.num_experts,
            "input_dim": self.input_dim,
            "output_dim": self.output_dim,
            "top_k": self.top_k,
            "use_confidence": self.use_confidence,
            "expert_names": self.expert_names,
            "expert_input_dims": expert_input_dims,
            "expert_output_dims": expert_output_dims,
            "expert_architectures": expert_architectures,
            "total_params": self.total_params(),
        }

        with open(os.path.join(path, f"{name}_config.json"), "w") as f:
            json.dump(config, f, indent=2)

        pt_path = os.path.join(path, f"{name}.pt")
        torch.save(self.state_dict(), pt_path)

        with open(os.path.join(path, f"{name}_readme.txt"), "w") as f:
            f.write(self.summary())
            f.write(f"\n\nSaved to: {os.path.abspath(path)}\n")

        return pt_path

    @classmethod
    def load(cls, path: str, name: str = "model") -> "StackedModel":
        """Load a stacked model from a folder.

        Args:
            path: Directory containing the saved files.
            name: Model name (file prefix used during save).

        Returns:
            The loaded StackedModel.
        """
        import json
        import os

        with open(os.path.join(path, f"{name}_config.json"), "r") as f:
            config = json.load(f)

        state_dict = torch.load(
            os.path.join(path, f"{name}.pt"),
            map_location="cpu",
            weights_only=True,
        )

        # Reconstruct experts from saved architecture info
        experts = []
        for arch in config["expert_architectures"]:
            layers = []
            for layer_info in arch:
                if layer_info["type"] == "Linear":
                    layers.append(nn.Linear(layer_info["in"], layer_info["out"]))
                elif layer_info["type"] == "ReLU":
                    layers.append(nn.ReLU())
            experts.append(nn.Sequential(*layers))

        expert_input_dims = config.get("expert_input_dims")
        expert_output_dims = config.get("expert_output_dims")

        model = cls(
            experts=experts,
            input_dim=config["input_dim"],
            output_dim=config["output_dim"],
            top_k=config["top_k"],
            expert_input_dims=expert_input_dims,
            expert_output_dims=expert_output_dims,
            expert_names=config.get("expert_names"),
            use_confidence=config.get("use_confidence", True),
        )

        model.load_state_dict(state_dict)
        model.eval()
        return model

    def push_to_hub(
        self,
        repo_id: str,
        commit_message: str = "Upload stacked model",
        private: bool = False,
        token: Optional[str] = None,
        create_model_card: bool = True,
        safe_serialization: bool = True,
        base_models: Optional[List[str]] = None,
    ) -> str:
        """Push the stacked model to HuggingFace Hub.

        Args:
            repo_id: Repository ID (e.g., "username/model-name").
            commit_message: Git commit message.
            private: Whether the repo should be private.
            token: HuggingFace API token (uses cached if None).
            create_model_card: Auto-generate README.md.
            safe_serialization: Use safetensors format (recommended).
            base_models: List of base model IDs for the model card.

        Returns:
            URL of the uploaded model.

        Example::

            stacked = pipeline.stack([model_a, model_b], sample_data=X)
            url = stacked.push_to_hub(
                "username/my-stacked-model",
                base_models=["org/model-a", "org/model-b"],
            )
            print(f"Model uploaded to: {url}")
        """
        try:
            from huggingface_hub import HfApi, create_repo
        except ImportError:
            raise ImportError(
                "HuggingFace Hub integration requires `huggingface_hub`. "
                "Install with: pip install huggingface_hub"
            )

        import json
        import tempfile
        import os

        api = HfApi(token=token)

        # Create repo if needed
        try:
            create_repo(repo_id, private=private, exist_ok=True, token=token)
        except Exception as e:
            # Repo might already exist
            pass

        with tempfile.TemporaryDirectory() as tmpdir:
            # Save weights
            if safe_serialization:
                try:
                    from model_osmosis.merging.streaming import StreamingParameterLoader
                    loader = StreamingParameterLoader()
                    loader.save_to_safetensors(
                        self.state_dict(),
                        os.path.join(tmpdir, "model.safetensors"),
                        metadata={
                            "format": "pt",
                            "library": "model-osmosis",
                            "model_type": "stacked_moe",
                        }
                    )
                except ImportError:
                    # Fallback to PyTorch format
                    torch.save(self.state_dict(), os.path.join(tmpdir, "pytorch_model.bin"))
            else:
                torch.save(self.state_dict(), os.path.join(tmpdir, "pytorch_model.bin"))

            # Save config
            expert_input_dims = []
            expert_output_dims = []
            expert_architectures = []

            for i in range(self.num_experts):
                # Get input dims
                for m in self.input_projectors[i].modules():
                    if isinstance(m, nn.Linear):
                        expert_input_dims.append(m.out_features)
                        break
                else:
                    expert_input_dims.append(self.input_dim)

                # Get output dims
                for m in self.output_projectors[i].modules():
                    if isinstance(m, nn.Linear):
                        expert_output_dims.append(m.in_features)
                        break
                else:
                    expert_output_dims.append(self.output_dim)

                # Get architecture
                layers_info = []
                for m in self.experts[i].modules():
                    if isinstance(m, nn.Linear):
                        layers_info.append({
                            "type": "Linear",
                            "in": m.in_features,
                            "out": m.out_features,
                        })
                    elif isinstance(m, nn.ReLU):
                        layers_info.append({"type": "ReLU"})
                expert_architectures.append(layers_info)

            config = {
                "architectures": ["StackedModel"],
                "model_type": "stacked_moe",
                "num_experts": self.num_experts,
                "input_dim": self.input_dim,
                "output_dim": self.output_dim,
                "top_k": self.top_k,
                "use_confidence": self.use_confidence,
                "expert_names": self.expert_names,
                "expert_input_dims": expert_input_dims,
                "expert_output_dims": expert_output_dims,
                "expert_architectures": expert_architectures,
                "total_params": self.total_params(),
                "library_name": "model-osmosis",
            }

            with open(os.path.join(tmpdir, "config.json"), "w") as f:
                json.dump(config, f, indent=2)

            # Generate model card
            if create_model_card:
                try:
                    from model_osmosis.utils.model_card import (
                        ModelCardGenerator,
                        StackCardConfig,
                    )
                    generator = ModelCardGenerator()
                    card_config = StackCardConfig(
                        model_name=repo_id.split("/")[-1],
                        expert_models=base_models or self.expert_names,
                        expert_names=self.expert_names,
                        top_k=self.top_k,
                        description=f"A Mixture-of-Experts model with {self.num_experts} experts using top-{self.top_k} routing.",
                    )
                    generator.save(card_config, os.path.join(tmpdir, "README.md"))
                except Exception as e:
                    # Fallback: simple README
                    with open(os.path.join(tmpdir, "README.md"), "w") as f:
                        f.write(f"# {repo_id.split('/')[-1]}\n\n")
                        f.write(self.summary())
                        f.write("\n\n---\n*Created with ModelOsmosis*\n")

            # Upload all files
            api.upload_folder(
                folder_path=tmpdir,
                repo_id=repo_id,
                commit_message=commit_message,
            )

        return f"https://huggingface.co/{repo_id}"

    @classmethod
    def from_pretrained(
        cls,
        repo_id: str,
        revision: Optional[str] = None,
        token: Optional[str] = None,
    ) -> "StackedModel":
        """Load a stacked model from HuggingFace Hub.

        Args:
            repo_id: Repository ID (e.g., "username/model-name").
            revision: Specific revision to load.
            token: HuggingFace API token.

        Returns:
            The loaded StackedModel.

        Example::

            model = StackedModel.from_pretrained("username/my-stacked-model")
        """
        try:
            from huggingface_hub import snapshot_download
        except ImportError:
            raise ImportError(
                "HuggingFace Hub integration requires `huggingface_hub`. "
                "Install with: pip install huggingface_hub"
            )

        import json
        import tempfile
        import os

        # Download all files
        local_dir = snapshot_download(
            repo_id,
            revision=revision,
            token=token,
        )

        # Load config
        config_path = os.path.join(local_dir, "config.json")
        with open(config_path, "r") as f:
            config = json.load(f)

        # Load weights
        safetensors_path = os.path.join(local_dir, "model.safetensors")
        pytorch_path = os.path.join(local_dir, "pytorch_model.bin")

        if os.path.exists(safetensors_path):
            from safetensors.torch import load_file
            state_dict = load_file(safetensors_path)
        elif os.path.exists(pytorch_path):
            state_dict = torch.load(pytorch_path, map_location="cpu", weights_only=True)
        else:
            raise FileNotFoundError(
                f"No model weights found in {repo_id}. "
                "Expected model.safetensors or pytorch_model.bin"
            )

        # Reconstruct experts
        experts = []
        for arch in config.get("expert_architectures", []):
            layers = []
            for layer_info in arch:
                if layer_info["type"] == "Linear":
                    layers.append(nn.Linear(layer_info["in"], layer_info["out"]))
                elif layer_info["type"] == "ReLU":
                    layers.append(nn.ReLU())
            experts.append(nn.Sequential(*layers))

        # Create model
        model = cls(
            experts=experts,
            input_dim=config["input_dim"],
            output_dim=config["output_dim"],
            top_k=config.get("top_k", 2),
            expert_input_dims=config.get("expert_input_dims"),
            expert_output_dims=config.get("expert_output_dims"),
            expert_names=config.get("expert_names"),
            use_confidence=config.get("use_confidence", True),
        )

        model.load_state_dict(state_dict)
        model.eval()
        return model


