import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, AutoModelForCausalLM
from .configuration_parallel_mlp import ParallelMLPConfig


class ParallelMLPAdapter(nn.Module):
    """Simple bottleneck MLP (norm → up → SiLU → down) parallel to a transformer block.

    down_proj is zero-initialized so the adapter is a strict pass-through at the
    start of training, avoiding representation shock.
    """

    def __init__(self, hidden_size: int, intermediate_size: int, rms_norm_eps: float = 1e-5):
        super().__init__()
        self.norm      = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
        self.up_proj   = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
        nn.init.zeros_(self.down_proj.weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.down_proj(F.silu(self.up_proj(self.norm(x))))


class ParallelMLPBlockWrapper(nn.Module):
    """Wraps a transformer decoder layer, adding a parallel MLP adapter.

    The adapter sees the same pre-block hidden states as the original block.
    Its delta (adapter output minus input) is added to the original block's output,
    so neither path depends on the other — true parallel execution.
    """

    def __init__(self, original_block: nn.Module, mlp_adapter: ParallelMLPAdapter):
        super().__init__()
        self.original_block = original_block
        self.mlp_adapter = mlp_adapter  # name used by freeze filter: "mlp_adapter"

    def __getattr__(self, name: str):
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.original_block, name)

    def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
        block_outputs = self.original_block(hidden_states, *args, **kwargs)
        out_hidden = block_outputs[0] if isinstance(block_outputs, tuple) else block_outputs

        # adapter(x) returns x + delta; extract only the delta to add to block output
        adapter_delta = self.mlp_adapter(hidden_states) - hidden_states
        out_hidden = out_hidden + adapter_delta

        if isinstance(block_outputs, tuple):
            return (out_hidden,) + block_outputs[1:]
        return out_hidden


class UnifiedParallelMLPForCausalLM(PreTrainedModel):
    config_class = ParallelMLPConfig

    def __init__(self, config: ParallelMLPConfig, **kwargs):
        super().__init__(config)
        self.backbone = AutoModelForCausalLM.from_pretrained(
            config.base_model_name_or_path,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            attn_implementation="flash_attention_2",
        )

        backbone_dtype = next(self.backbone.parameters()).dtype
        layers = self.backbone.model.layers
        for pos in config.mlp_positions:
            if 0 < pos <= len(layers):
                original_layer = layers[pos - 1]
                adapter = ParallelMLPAdapter(
                    config.hidden_size,
                    config.mlp_intermediate_size,
                    config.rms_norm_eps,
                ).to(dtype=backbone_dtype)
                layers[pos - 1] = ParallelMLPBlockWrapper(original_layer, adapter)

    def forward(self, *args, **kwargs):
        return self.backbone(*args, **kwargs)

    def generate(self, *args, **kwargs):
        return self.backbone.generate(*args, **kwargs)