"""Fuse-3 model: LFM2-2.6B host + Qwen3.6-35B-A3B coding experts.

Architecture: per-layer FFN augmentation with Qwen3.6 MoE experts.
At each augmented host layer, coding experts from Qwen3.6 are added
alongside the host's native SwiGLU FFN. A learned router decides which
experts fire.

Key design differences from Fuse-2:
  - NO bridges needed: LFM2 hidden_size (2048) == Qwen3.6 expert input (2048).
    Experts are 2048->512->2048, directly operating on host hidden states.
  - LFM2 is NOT a Qwen model. Lfm2ForCausalLM has a hybrid conv+attention
    architecture with 22 short-conv layers + 8 GQA layers.
  - The augmented layer splits LFM2's forward into attention/conv + FFN,
    then adds expert output as a parallel path to the dense FFN.
  - expert_scale zero-init -> model starts as exact LFM2-2.6B
  - Router initialized to low activation -> coding path fires rarely at first
  - Frozen host, frozen experts -> only router + scale train
  - use_cache = False initially (correctness first)
  - Optional SwiGLU clamping (limit=10.0) for Qwen3.6 expert stability
"""
from __future__ import annotations

import math
from typing import Iterator

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Lfm2Config, Lfm2ForCausalLM


class Fuse3Config(Lfm2Config):
    """LFM2 config extended with Fuse-3 MoE coding expert parameters."""

    model_type = "fuse3"

    def __init__(
        self,
        # Expert configuration
        expert_intermediate_size: int = 512,   # Qwen3.6 expert intermediate
        experts_per_layer: dict | None = None,  # layer_idx -> list of expert IDs
        num_augmented_layers: int = 0,
        top_k_experts: int = 8,                 # Qwen3.6 uses 8 routed
        # Router configuration
        router_init_scale: float = -2.0,        # low initial activation
        load_balance_coef: float = 0.001,
        # Expert stability
        swiglu_limit: float = 10.0,             # clamp expert activations (from fuse-2 lessons)
        coding_enabled: bool = True,
        # Scale
        expert_scale_init: float = -5.0,         # softplus(-5)≈0.007, small but active
        **kwargs,
    ):
        super().__init__(**kwargs)
        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.router_init_scale = router_init_scale
        self.load_balance_coef = load_balance_coef
        self.swiglu_limit = swiglu_limit
        self.coding_enabled = coding_enabled
        self.expert_scale_init = expert_scale_init
        self.use_cache = False
        self.architectures = ["Fuse3ForCausalLM"]


class SwiGLUExpert(nn.Module):
    """A single Qwen3.6 MoE expert (SwiGLU FFN).

    gate_proj: (intermediate, hidden)  — w1
    up_proj:   (intermediate, hidden)  — w3
    down_proj: (hidden, intermediate)  — w2

    Input:  (..., hidden_size)
    Output: (..., hidden_size)
    """

    def __init__(self, hidden_size: int, intermediate_size: int, swiglu_limit: float = 0.0):
        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)
        self.swiglu_limit = swiglu_limit

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.gate_proj(x)
        up = self.up_proj(x)
        act = F.silu(gate) * up
        if self.swiglu_limit > 0:
            act = act.clamp(-self.swiglu_limit, self.swiglu_limit)
        return self.down_proj(act)


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

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

    def __init__(
        self,
        input_dim: int,
        num_experts: int,
        top_k: int = 8,
        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)
        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, hidden_size)

        Returns:
            router_weights: (batch*seq, top_k) — normalized 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
        """
        logits = self.gate(hidden_states)
        scores = F.softplus(logits).sqrt()
        topk_weights, topk_indices = scores.topk(self.top_k, dim=-1)
        topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8)
        return topk_weights, topk_indices, logits


class Fuse3AugmentedLayer(nn.Module):
    """One LFM2 layer augmented with Qwen3.6 coding experts.

    Forward flow (mirrors Lfm2DecoderLayer but with MoE added to FFN):
      1. Attention or ShortConv (frozen host, same as LFM2)
      2. Dense FFN (frozen host SwiGLU)
      3. Router: select top-k coding experts from FFN input
      4. Experts: parallel SwiGLU computation (frozen, from Qwen3.6)
      5. expert_scale * expert_output added to residual
      6. No bridge needed — hidden_size matches on both sides

    The model starts as exact LFM2 (expert_scale=0) and learns to
    incorporate coding experts through router + scale training.
    """

    def __init__(
        self,
        host_layer: nn.Module,
        hidden_size: int,
        expert_intermediate: int,
        num_experts: int,
        top_k: int = 8,
        swiglu_limit: float = 10.0,
        router_init_scale: float = -2.0,
        coding_enabled: bool = True,
        expert_scale_init: float = 0.0,
    ):
        super().__init__()
        self.host_layer = host_layer
        self.coding_enabled = coding_enabled
        self.num_experts = num_experts
        self.top_k = min(top_k, num_experts)
        self.hidden_size = hidden_size

        # Expose host layer attributes needed by the LFM2 model forward pass
        self.is_attention_layer = getattr(host_layer, "is_attention_layer", False)

        # Router (operates on hidden_size directly — no bridge)
        self.router = Fuse3Router(
            hidden_size, num_experts, top_k, router_init_scale
        )

        # Experts (frozen, loaded from Qwen3.6)
        self.experts = nn.ModuleList([
            SwiGLUExpert(hidden_size, expert_intermediate, swiglu_limit)
            for _ in range(num_experts)
        ])

        # Scale factor for expert output (zero-init = exact LFM2)
        self.expert_scale = nn.Parameter(torch.tensor(expert_scale_init))

        # Store last router logits for load balancing loss
        self._last_router_logits: torch.Tensor | None = None

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values=None,
        cache_position: torch.LongTensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        # ── 1. Run the original host layer (attention/conv + FFN) ──
        # This delegates to the original Lfm2DecoderLayer.forward, ensuring
        # perfect compatibility with the host model's conv/attention impl.
        layer_output = self.host_layer(
            hidden_states,
            position_embeddings=position_embeddings,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            cache_position=cache_position,
            **kwargs,
        )

        # If coding disabled or no experts, return original output unchanged
        if not self.coding_enabled or self.num_experts == 0:
            return layer_output

        # ── 2. Always run router (for load balancing gradients) ──
        original_shape = layer_output.shape
        h_flat = layer_output.reshape(-1, original_shape[-1])

        topk_weights, expert_indices, router_logits = self.router(h_flat)
        self._last_router_logits = router_logits

        # Skip expert computation if scale is effectively zero
        # (router still ran, so LB loss gradients flow)
        scale = F.softplus(self.expert_scale)
        if scale.item() < 1e-4:
            return layer_output

        # ── 3. Compute expert outputs (sparse) ──
        # Detach expert inputs/outputs — experts are frozen, so no gradients
        # need to flow through the expert weights. Gradients only flow through
        # topk_weights (router) and scale (expert_scale).
        # Use index_add (out-of-place) to avoid in-place op autograd issues.
        expert_output = torch.zeros_like(h_flat)

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

            for eid in range(self.num_experts):
                token_mask = indices == eid
                if not token_mask.any():
                    continue
                expert_in = h_flat[token_mask].detach()
                expert_out = self.experts[eid](expert_in).detach()
                weighted_out = weights[token_mask].unsqueeze(-1) * expert_out
                token_indices = torch.where(token_mask)[0]
                expert_output = expert_output.index_add(
                    0, token_indices, weighted_out
                )

        # ── 4. Normalize expert output to match host activation scale ──
        # The experts come from Qwen3.6 which has different activation
        # distributions. Rescale to match host std, but clamp the ratio
        # to prevent amplification of sparse outputs.
        host_std = layer_output.std().detach() + 1e-6
        expert_std = expert_output.std().detach() + 1e-6
        ratio = torch.clamp(host_std / expert_std, max=2.0)
        expert_output = expert_output * ratio

        # ── 5. Scale and add to residual ──
        # Clamp scale to prevent runaway
        scale = torch.clamp(scale, max=0.1)
        expert_delta = scale * expert_output
        expert_delta = expert_delta.reshape(original_shape)

        return layer_output + expert_delta

    def get_router_logits(self) -> torch.Tensor | None:
        return self._last_router_logits


class Fuse3ForCausalLM(Lfm2ForCausalLM):
    """LFM2-2.6B host + Qwen3.6-35B-A3B coding experts.

    The model starts as an exact LFM2-2.6B (expert_scale=0) and learns
    to incorporate coding experts through router and scale training.
    """

    config_class = Fuse3Config
    _no_split_modules = ["Lfm2DecoderLayer", "Fuse3AugmentedLayer"]

    def __init__(self, config: Fuse3Config):
        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] = Fuse3AugmentedLayer(
                host_layer=original_layer,
                hidden_size=config.hidden_size,
                expert_intermediate=config.expert_intermediate_size,
                num_experts=num_experts,
                top_k=min(config.top_k_experts, num_experts),
                swiglu_limit=config.swiglu_limit,
                router_init_scale=config.router_init_scale,
                coding_enabled=config.coding_enabled,
                expert_scale_init=config.expert_scale_init,
            )
            augmented_count += 1

        config.num_augmented_layers = augmented_count

    def set_coding_enabled(self, enabled: bool) -> None:
        for layer in self.model.layers:
            if isinstance(layer, Fuse3AugmentedLayer):
                layer.coding_enabled = enabled

    def get_augmented_layers(self) -> list[tuple[int, Fuse3AugmentedLayer]]:
        return [
            (i, layer)
            for i, layer in enumerate(self.model.layers)
            if isinstance(layer, Fuse3AugmentedLayer)
        ]

    def get_trainable_params(self) -> dict[str, nn.Parameter]:
        trainable = {}
        for name, param in self.named_parameters():
            if any(key in name for key in ("router", "expert_scale")):
                trainable[name] = param
        return trainable

    def freeze_host_and_experts(self) -> None:
        for name, param in self.named_parameters():
            if any(key in name for key in ("router", "expert_scale")):
                param.requires_grad = True
            else:
                param.requires_grad = False

    def get_all_router_logits(self) -> list[torch.Tensor]:
        """Collect router logits from all augmented layers after forward pass."""
        logits = []
        for layer in self.model.layers:
            if hasattr(layer, '_last_router_logits') and layer._last_router_logits is not None:
                logits.append(layer._last_router_logits)
        return logits

    def get_expert_scales(self) -> list[float]:
        """Get current effective expert scale (softplus) from all augmented layers."""
        import torch.nn.functional as F
        scales = []
        for layer in self.model.layers:
            if hasattr(layer, 'expert_scale'):
                scales.append(F.softplus(layer.expert_scale).item())
        return scales

    def reset_router_logits(self) -> None:
        """Clear stored router logits (call before each forward pass during training)."""
        for layer in self.model.layers:
            if hasattr(layer, '_last_router_logits'):
                layer._last_router_logits = None

    def count_parameters(self) -> dict[str, int]:
        counts = {
            "host": 0, "experts": 0, "routers": 0, "scale": 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 "router" in name:
                counts["routers"] += n
            elif "expert_scale" in name:
                counts["scale"] += 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,
    ):
        # Delegate to Lfm2ForCausalLM.forward with original args.
        # The augmented layers handle expert computation internally;
        # KV cache works because Fuse3AugmentedLayer delegates to the
        # original host layer's forward for attention/conv.
        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,
        )


def load_expert_weights(
    model: Fuse3ForCausalLM,
    expert_dir: str,
    expert_mapping: dict[int, list[int]],
) -> dict:
    """Load extracted Qwen3.6 expert weights into the Fuse3 model.

    Args:
        model: Fuse3 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

    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():
        if layer_idx >= len(model.model.layers):
            continue  # skip out-of-range layers
        augmented = model.model.layers[layer_idx]
        if not isinstance(augmented, Fuse3AugmentedLayer):
            continue  # skip non-augmented layers

        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]
                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


# ── AutoModel registration ─────────────────────────────────────────────
# This allows `AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)`
# to correctly load Fuse3 models, including bitsandbytes quantized versions.

from transformers import AutoConfig, AutoModelForCausalLM

AutoConfig.register("fuse3", Fuse3Config)
AutoModelForCausalLM.register(Fuse3Config, Fuse3ForCausalLM)