"""GLEAN-pruned OLMoE: HF-loadable model with ragged (variable-width) experts.

Pattern follows hbfreed/variable-flex-olmo's PrunedFlexOlmoForCausalLM
(docs/recon/prior-work-hbfreed.md), generalized from one scalar width to a
per-(layer, expert) width table: ``super().__init__`` builds the uniform
architecture from the config, then every MoE block is rebuilt to its pruned
shape — surviving experts only, each at its own width, router sliced to
match — so the state dict aligns exactly with what
``glean.prune.prune_channels_global`` leaves behind.

Caveat: ``output_router_logits=True`` (the load-balancing aux loss) assumes a
uniform ``config.num_experts`` and is unsupported on ragged models.
"""

import torch.nn as nn
from transformers.activations import ACT2FN
from transformers.models.olmoe.modeling_olmoe import OlmoeForCausalLM

from .configuration_pruned_olmoe import PrunedOlmoeConfig


class RaggedOlmoeMLP(nn.Module):
    """OlmoeMLP with an explicit intermediate width (SwiGLU, no biases)."""

    def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str):
        super().__init__()
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        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.act_fn = ACT2FN[hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class PrunedOlmoeForCausalLM(OlmoeForCausalLM):
    """OLMoE with per-layer surviving-expert lists at per-expert widths."""

    config_class = PrunedOlmoeConfig

    def __init__(self, config: PrunedOlmoeConfig):
        super().__init__(config)
        widths_table = getattr(config, "expert_widths", None)
        if widths_table is None:
            return  # unpruned: plain OLMoE
        if len(widths_table) != len(self.model.layers):
            raise ValueError(
                f"expert_widths has {len(widths_table)} rows but the model has "
                f"{len(self.model.layers)} decoder layers"
            )
        for layer, widths in zip(self.model.layers, widths_table):
            if any(w <= 0 for w in widths):
                raise ValueError("expert_widths must list surviving experts only (>0)")
            block = layer.mlp
            if len(widths) < block.top_k:
                raise ValueError(
                    f"a layer keeps {len(widths)} experts < top_k={block.top_k}"
                )
            block.num_experts = len(widths)
            block.gate = nn.Linear(config.hidden_size, len(widths), bias=False)
            block.experts = nn.ModuleList(
                RaggedOlmoeMLP(config.hidden_size, w, config.hidden_act)
                for w in widths
            )