"""Fuse-2 model: Qwen3 host + DeepSeek V4 Flash coding experts.

Architecture: per-layer expert augmentation (Option B from the master plan).
At each augmented host layer, coding experts from DeepSeek V4 Flash are added
alongside the host's native FFN. A learned router decides which experts fire.

Key design principles (from fuse1 lessons):
  - bridge_out zero-init → model starts as exact Qwen3-4B
  - repair_up zero-init → no residual correction initially
  - Router initialized to low activation → coding path fires rarely at first
  - Frozen experts, frozen host → only bridges + routers + repair train
  - KV cache supported — host attention uses DynamicCache, expert path is
    per-token (no cross-token attention) so cache works transparently
"""
from __future__ import annotations

import math
from copy import deepcopy
from typing import Iterator

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Qwen3Config, Qwen3ForCausalLM


class Fuse2Config(Qwen3Config):
    """Qwen3 config extended with Fuse-2 MoE coding expert parameters."""

    model_type = "fuse2"

    def __init__(
        self,
        # Expert configuration
        expert_hidden_size: int = 4096,       # DeepSeek V4 hidden
        expert_intermediate_size: int = 2048,  # DeepSeek V4 expert intermediate
        experts_per_layer: dict | None = None, # layer_idx -> list of expert IDs
        num_augmented_layers: int = 0,
        top_k_experts: int = 2,
        # Bridge configuration
        bridge_rank: int = 7,
        coding_enabled: bool = True,
        # Router configuration
        router_init_scale: float = -2.0,  # low initial activation
        load_balance_coef: float = 0.01,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.expert_hidden_size = expert_hidden_size
        self.expert_intermediate_size = expert_intermediate_size
        self.experts_per_layer = experts_per_layer or {}
        self.num_augmented_layers = num_augmented_layers
        self.top_k_experts = top_k_experts
        self.bridge_rank = bridge_rank
        self.coding_enabled = coding_enabled
        self.router_init_scale = router_init_scale
        self.load_balance_coef = load_balance_coef
        self.architectures = ["Fuse2ForCausalLM"]


class SwiGLUExpert(nn.Module):
    """A single DeepSeek V4 Flash expert (SwiGLU FFN).

    gate_proj: (intermediate, hidden)
    up_proj:   (intermediate, hidden)
    down_proj: (hidden, intermediate)
    """

    # DeepSeek V4 Flash uses swiglu_limit=10.0 to clamp intermediate
    # activations. Without this, outlier values grow exponentially across
    # 36 layers and produce NaN by layer 6.
    SWIGLU_LIMIT = 10.0

    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up = F.silu(self.gate_proj(x)) * self.up_proj(x)
        gate_up = gate_up.clamp(-self.SWIGLU_LIMIT, self.SWIGLU_LIMIT)
        return self.down_proj(gate_up)


class Fuse2Router(nn.Module):
    """Per-layer router for coding experts.

    Uses sqrtsoftplus scoring (matching DeepSeek V4's approach) with
    top-k selection and optional load balancing.
    """

    def __init__(
        self,
        input_dim: int,
        num_experts: int,
        top_k: int = 2,
        init_scale: float = -2.0,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = min(top_k, num_experts)
        self.gate = nn.Linear(input_dim, num_experts, bias=False)
        # Initialize to low activation so coding path fires rarely at start
        # Skip init on meta device (used by low_cpu_mem_usage / init_empty_weights)
        if self.gate.weight.device.type != 'meta':
            nn.init.normal_(self.gate.weight, mean=0.0, std=0.01)
        self.init_scale = init_scale

    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Route tokens to experts.

        Args:
            hidden_states: (batch*seq, expert_hidden) — already bridged

        Returns:
            router_weights: (batch*seq, top_k) — softmax weights for selected experts
            expert_indices: (batch*seq, top_k) — which experts were selected
            router_logits: (batch*seq, num_experts) — raw logits for load balancing
        """
        # sqrtsoftplus scoring (from DeepSeek V4)
        # Clamp softplus to min=1e-6 before sqrt to prevent NaN gradients
        # when logits are very negative (softplus → 0 → sqrt(0) = 0, but
        # gradient sqrt'(0) = inf).
        logits = self.gate(hidden_states)  # (tokens, num_experts)
        scores = F.softplus(logits).clamp(min=1e-6).sqrt()

        # Top-k selection
        topk_weights, topk_indices = scores.topk(self.top_k, dim=-1)

        # Normalize weights
        topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8)

        return topk_weights, topk_indices, logits


class Fuse2AugmentedLayer(nn.Module):
    """One Qwen3 layer augmented with DeepSeek V4 coding experts.

    Forward flow (v2 — with residual-safe architecture):
      1. Standard Qwen3 attention + FFN (frozen)
      2. bridge_in: host_hidden -> expert_hidden
      3. router: select top-k coding experts
      4. experts: parallel SwiGLU computation
      5. expert_output normalized to match expert_input scale
      6. bridge_out: expert_hidden -> host_hidden
      7. coding_norm: RMSNorm on coding_delta (learned scale)
      8. coding_gate: sigmoid gate (learned, init ~0.12)
      9. residual clamp: bound delta to max 10% of residual norm
     10. repair: rank-r residual correction (zero-init)
     11. hidden += gated coding_delta + repair_delta

    v2 safeguards (prevent residual stream degeneration):
      - Expert output normalization (consistent magnitude)
      - RMSNorm on coding_delta (bounded scale)
      - Learnable sigmoid gate (controls contribution strength)
      - Residual-safe clamping (prevents any layer from overwhelming)
    """

    def __init__(
        self,
        host_layer: nn.Module,
        host_hidden: int,
        expert_hidden: int,
        expert_intermediate: int,
        num_experts: int,
        top_k: int = 2,
        bridge_rank: int = 7,
        router_init_scale: float = -2.0,
        coding_enabled: bool = True,
        # v2 safeguards
        max_delta_ratio: float = 0.1,  # coding_delta <= 10% of residual norm
    ):
        super().__init__()
        self.host_layer = host_layer
        self.coding_enabled = coding_enabled
        self.num_experts = num_experts
        self.top_k = top_k
        self.max_delta_ratio = max_delta_ratio

        # Expose host layer attributes needed by the Qwen3 model forward pass
        self.attention_type = getattr(host_layer, "attention_type", "full_attention")

        # Bridge: host space <-> expert space
        self.bridge_in = nn.Linear(host_hidden, expert_hidden, bias=False)
        self.bridge_out = nn.Linear(expert_hidden, host_hidden, bias=False)

        # Router
        self.router = Fuse2Router(
            expert_hidden, num_experts, top_k, router_init_scale
        )

        # Experts (frozen, loaded from DeepSeek V4 Flash)
        self.experts = nn.ModuleList([
            SwiGLUExpert(expert_hidden, expert_intermediate)
            for _ in range(num_experts)
        ])

        # v2: RMSNorm on coding_delta (learned scale, init=1.0)
        # This normalizes the coding contribution to unit variance before
        # the gate scales it down. The learned scale allows the model to
        # adjust per-layer contribution magnitude during training.
        self.coding_norm = nn.RMSNorm(host_hidden, eps=1e-6)

        # v2: Learnable sigmoid gate (init=-2.0 -> sigmoid(-2) ~ 0.12)
        # This lets the model learn how strongly to incorporate coding at
        # each layer. Starting at 12% ensures coding contributes but can't
        # dominate. The gate is a single scalar per layer.
        self.coding_gate = nn.Parameter(torch.tensor(-2.0))

        # Residual repair (low-rank)
        self.repair_down = nn.Linear(host_hidden, bridge_rank, bias=False)
        self.repair_up = nn.Linear(bridge_rank, host_hidden, bias=False)

        # Initialize for preservation: zero-init bridge_out and repair_up
        # Skip init on meta device (used by low_cpu_mem_usage / init_empty_weights)
        if self.bridge_in.weight.device.type != 'meta':
            nn.init.normal_(self.bridge_in.weight, mean=0.0, std=0.02)
            nn.init.zeros_(self.bridge_out.weight)
            nn.init.normal_(self.repair_down.weight, mean=0.0, std=0.02)
            nn.init.zeros_(self.repair_up.weight)

    def _expert_computation(self, expert_input: torch.Tensor) -> torch.Tensor:
        """Compute expert output from expert_input.

        This is the checkpointed part of the forward pass. It includes:
        - Router scoring and top-k selection
        - Expert SwiGLU computation (reads weights from disk via monkey-patched forward)
        - Weighted accumulation of expert outputs
        - v2: Expert output normalization (consistent magnitude)

        When wrapped in torch.utils.checkpoint, only expert_input is saved
        during forward. The entire computation (including disk reads for expert
        weights) is recomputed during backward, keeping VRAM bounded to one
        layer's worth of expert weights at a time.
        """
        # Route to experts
        topk_weights, expert_indices, router_logits = self.router(expert_input)

        # Compute expert outputs (sparse — only selected experts)
        # Use index_add for autograd-safe accumulation (no in-place ops).
        expert_output = torch.zeros_like(expert_input)

        for k in range(self.top_k):
            indices = expert_indices[:, k]  # (tokens,)
            weights = topk_weights[:, k]    # (tokens,)

            # Group tokens by expert for efficient computation
            for eid in range(self.num_experts):
                mask = indices == eid
                if not mask.any():
                    continue
                expert_in = expert_input[mask]
                expert_out = self.experts[eid](expert_in)
                weighted = weights[mask].unsqueeze(-1) * expert_out
                token_idx = torch.where(mask)[0]
                expert_output = expert_output.index_add(
                    0, token_idx, weighted.to(expert_output.dtype))

        # v2: Expert output normalization
        # Normalize expert_output to match expert_input's per-token norm.
        # This ensures consistent output magnitude regardless of which experts
        # were selected or their internal scale differences. Without this,
        # different expert combinations produce wildly different output scales,
        # causing inconsistent coding_delta magnitudes that corrupt the residual.
        out_norm = expert_output.norm(dim=-1, keepdim=True) + 1e-6
        in_norm = expert_input.norm(dim=-1, keepdim=True) + 1e-6
        expert_output = expert_output * (in_norm / out_norm)

        return expert_output

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values=None,
        use_cache: bool | None = False,
        position_embeddings=None,
        **kwargs,
    ) -> torch.Tensor:
        # 1. Run the host layer (attention + FFN)
        hidden_states = self.host_layer(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            position_embeddings=position_embeddings,
            **kwargs,
        )

        if not self.coding_enabled or self.num_experts == 0:
            return hidden_states

        # v2-phase: Check for per-token coding mask
        # If coding_mask is set on the model, only apply coding to tokens
        # where mask=True. This enables phase-based coding:
        #   - Reasoning tokens (before </thinking>): coding OFF (pure Qwen)
        #   - Answer tokens (after </thinking>): coding ON (DeepSeek contributes)
        coding_mask = getattr(self, '_coding_mask', None)

        original_shape = hidden_states.shape
        h_flat = hidden_states.reshape(-1, original_shape[-1])

        # If we have a coding_mask, select only the tokens that need coding
        if coding_mask is not None:
            # coding_mask shape: (batch, seq_len) -> (tokens,)
            mask_flat = coding_mask.reshape(-1)
            if not mask_flat.any():
                # No tokens need coding — skip entirely
                return hidden_states
            # Only process tokens where mask is True
            h_coding = h_flat[mask_flat]  # (n_coding_tokens, hidden)
        else:
            mask_flat = None
            h_coding = h_flat  # Process all tokens

        # 2. Bridge to expert space
        expert_input = self.bridge_in(h_coding)  # (n_tokens, expert_hidden)

        # 3-4. Expert computation (router + experts + accumulation)
        # Use gradient checkpointing during training to avoid storing all
        # expert weights in VRAM. The checkpoint saves only expert_input
        # during forward and recomputes the expert computation (re-reading
        # weights from disk) during backward. This keeps peak VRAM bounded
        # to one layer's worth of expert weights at a time.
        if self.training and expert_input.requires_grad:
            expert_output = torch.utils.checkpoint.checkpoint(
                self._expert_computation, expert_input, use_reentrant=False,
            )
        else:
            expert_output = self._expert_computation(expert_input)

        # 5. Bridge back to host space
        coding_delta = self.bridge_out(expert_output)

        # v2-6a: RMSNorm on coding_delta
        coding_delta = self.coding_norm(coding_delta)

        # v2-6b: Learnable sigmoid gate
        gate = torch.sigmoid(self.coding_gate)
        coding_delta = coding_delta * gate

        # v2-6c: Residual-safe clamping
        h_for_norm = h_coding if coding_mask is not None else h_flat
        h_norm = h_for_norm.norm(dim=-1, keepdim=True) + 1e-6
        delta_norm = coding_delta.norm(dim=-1, keepdim=True) + 1e-6
        max_delta = h_norm * self.max_delta_ratio
        scale = (max_delta / delta_norm).clamp(max=1.0)
        coding_delta = coding_delta * scale

        # 7. Repair (only for coding tokens)
        repair_delta = self.repair_up(self.repair_down(h_for_norm))

        # 8. Residual addition
        # If using a mask, scatter coding_delta back to the right positions
        if coding_mask is not None:
            result = h_flat.clone()
            result[mask_flat] = h_coding + coding_delta + repair_delta
            return result.reshape(original_shape)
        else:
            result = h_flat + coding_delta + repair_delta
            return result.reshape(original_shape)

    def get_router_logits(self) -> torch.Tensor | None:
        """Return last router logits for load balancing loss."""
        return getattr(self, "_last_router_logits", None)


class Fuse2ForCausalLM(Qwen3ForCausalLM):
    """Qwen3-4B host + DeepSeek V4 Flash coding experts.

    The model starts as an exact Qwen3-4B (zero-init bridges) and learns
    to incorporate coding experts through bridge and router training.
    """

    config_class = Fuse2Config
    _no_split_modules = ["Qwen3DecoderLayer", "Fuse2AugmentedLayer"]

    def __init__(self, config: Fuse2Config):
        super().__init__(config)

        # Replace specified layers with augmented versions
        experts_per_layer = config.experts_per_layer or {}
        augmented_count = 0

        for layer_idx_str, expert_ids in experts_per_layer.items():
            layer_idx = int(layer_idx_str)
            if layer_idx >= len(self.model.layers):
                raise ValueError(
                    f"Layer {layer_idx} out of range "
                    f"(model has {len(self.model.layers)} layers)"
                )

            num_experts = len(expert_ids)
            if num_experts == 0:
                continue

            original_layer = self.model.layers[layer_idx]
            self.model.layers[layer_idx] = Fuse2AugmentedLayer(
                host_layer=original_layer,
                host_hidden=config.hidden_size,
                expert_hidden=config.expert_hidden_size,
                expert_intermediate=config.expert_intermediate_size,
                num_experts=num_experts,
                top_k=min(config.top_k_experts, num_experts),
                bridge_rank=config.bridge_rank,
                router_init_scale=config.router_init_scale,
                coding_enabled=config.coding_enabled,
            )
            augmented_count += 1

        config.num_augmented_layers = augmented_count

    def set_coding_enabled(self, enabled: bool) -> None:
        """Toggle the coding expert path."""
        for layer in self.model.layers:
            if isinstance(layer, Fuse2AugmentedLayer):
                layer.coding_enabled = enabled

    def set_coding_mask(self, mask: torch.Tensor | None) -> None:
        """Set per-token coding mask for phase-based coding.

        When set, only tokens where mask=True will have the coding path
        applied. Tokens where mask=False get pure host (Qwen3) output.

        This enables the "Qwen reasons, DeepSeek codes" architecture:
        - Set mask=False for reasoning tokens (before </thinking>)
        - Set mask=True for answer tokens (after </thinking>)

        Pass None to disable masking (coding applies to all tokens).
        """
        for layer in self.model.layers:
            if isinstance(layer, Fuse2AugmentedLayer):
                layer._coding_mask = mask

    def get_augmented_layers(self) -> list[tuple[int, Fuse2AugmentedLayer]]:
        """Return (index, layer) pairs for all augmented layers."""
        return [
            (i, layer)
            for i, layer in enumerate(self.model.layers)
            if isinstance(layer, Fuse2AugmentedLayer)
        ]

    def get_trainable_params(self) -> dict[str, nn.Parameter]:
        """Return only the trainable parameters (bridges, routers, repair, v2 safeguards)."""
        trainable = {}
        for name, param in self.named_parameters():
            if any(
                key in name
                for key in ("bridge_in", "bridge_out", "router",
                           "repair_down", "repair_up",
                           "coding_norm", "coding_gate")
            ):
                trainable[name] = param
        return trainable

    def freeze_host_and_experts(self) -> None:
        """Freeze everything except bridges, routers, repair, and v2 safeguards."""
        for name, param in self.named_parameters():
            if any(
                key in name
                for key in ("bridge_in", "bridge_out", "router",
                           "repair_down", "repair_up",
                           "coding_norm", "coding_gate")
            ):
                param.requires_grad = True
            else:
                param.requires_grad = False

    def count_parameters(self) -> dict[str, int]:
        """Count parameters by category."""
        counts = {
            "host": 0,
            "experts": 0,
            "bridges": 0,
            "routers": 0,
            "repair": 0,
            "total": 0,
            "trainable": 0,
        }
        for name, param in self.named_parameters():
            n = param.numel()
            counts["total"] += n
            if param.requires_grad:
                counts["trainable"] += n

            if "bridge_in" in name or "bridge_out" in name:
                counts["bridges"] += n
            elif "router" in name:
                counts["routers"] += n
            elif "repair" in name:
                counts["repair"] += n
            elif "experts" in name:
                counts["experts"] += n
            else:
                counts["host"] += n

        return counts

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values=None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        **kwargs,
    ):
        return super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            labels=labels,
            use_cache=use_cache,
            **kwargs,
        )

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        """Load from HuggingFace Hub with automatic runtime fixes.

        This overrides the default from_pretrained to apply three critical
        fixes after weight loading:

        1. Initialize coding_gate and coding_norm if they're still on meta
           device (these params are not in the safetensors checkpoint).
        2. Cast all RMSNorm/LayerNorm weights from float32 to bfloat16 to
           enable fused SDPA kernel dispatch (otherwise falls back to slow
           Python loops).
        3. Ensure coding_enabled is True (config default).

        With these fixes, from_pretrained produces a working model without
        any manual post-load patching.
        """
        model = super().from_pretrained(*args, **kwargs)
        model._apply_runtime_fixes()
        return model

    def _apply_runtime_fixes(self):
        """Apply runtime fixes after weight loading.

        Called automatically by from_pretrained. Can also be called manually
        if the model was loaded via a custom path (e.g., init_empty_weights +
        manual safetensors loading).
        """
        device = next(self.parameters()).device
        fixed_meta = 0
        fixed_norms = 0

        for layer in self.model.layers:
            if not isinstance(layer, Fuse2AugmentedLayer):
                continue

            # Fix 1: coding_gate on meta device → initialize to -2.0
            if hasattr(layer, 'coding_gate'):
                if layer.coding_gate.device.type == 'meta':
                    layer.coding_gate = nn.Parameter(
                        torch.tensor(-2.0, device=device))
                    fixed_meta += 1

            # Fix 2: coding_norm on meta device → create fresh RMSNorm
            if hasattr(layer, 'coding_norm'):
                if hasattr(layer.coding_norm, 'weight') and \
                   layer.coding_norm.weight.device.type == 'meta':
                    layer.coding_norm = nn.RMSNorm(
                        layer.coding_norm.weight.shape[0], eps=1e-6).to(device)
                    fixed_meta += 1

        # Fix 3: Cast float32 norm weights to bfloat16 for fused kernels
        for module in self.modules():
            if hasattr(module, 'weight') and hasattr(module, 'eps'):
                if module.weight.dtype == torch.float32:
                    module.weight.data = module.weight.data.to(torch.bfloat16)
                    fixed_norms += 1

        # Ensure coding is enabled
        self.set_coding_enabled(True)

        return {"meta_params_fixed": fixed_meta, "norms_cast_to_bf16": fixed_norms}


def load_expert_weights(
    model: Fuse2ForCausalLM,
    expert_dir: str,
    expert_mapping: dict[int, list[int]],
) -> dict:
    """Load extracted DeepSeek V4 expert weights into the Fuse2 model.

    Args:
        model: Fuse2 model with augmented layers
        expert_dir: directory containing expert safetensors
        expert_mapping: layer_idx -> list of expert IDs (matching selection order)

    Returns:
        Manifest of loaded tensors with hash verification
    """
    from safetensors.torch import load_file
    import glob

    # Load all shards
    shard_files = sorted(glob.glob(f"{expert_dir}/experts-*.safetensors"))
    if not shard_files:
        raise FileNotFoundError(f"No expert shards found in {expert_dir}")

    all_tensors = {}
    for shard in shard_files:
        all_tensors.update(load_file(shard))

    loaded = {}
    for layer_idx, expert_ids in expert_mapping.items():
        augmented = model.model.layers[layer_idx]
        if not isinstance(augmented, Fuse2AugmentedLayer):
            raise ValueError(f"Layer {layer_idx} is not augmented")

        for local_idx, global_eid in enumerate(expert_ids):
            prefix = f"layer{layer_idx:02d}_expert{global_eid:03d}"

            for pname in ("gate_proj.weight", "up_proj.weight", "down_proj.weight"):
                key = f"{prefix}.{pname}"
                if key not in all_tensors:
                    raise KeyError(f"Missing expert tensor: {key}")

                tensor = all_tensors[key]
                target_name = pname.replace(".", "_").replace("_weight", "")
                # Map to expert module
                parts = pname.split(".")
                module = augmented.experts[local_idx]
                for part in parts[:-1]:
                    module = getattr(module, part)
                param = getattr(module, parts[-1])
                param.data.copy_(tensor.to(param.dtype))

                loaded[key] = {
                    "shape": list(tensor.shape),
                    "destination": f"layers.{layer_idx}.experts.{local_idx}.{pname}",
                }

    return loaded