import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class DraftBlock(nn.Module):
    def __init__(self, D, H):
        super().__init__()
        self.fc1 = nn.Linear(D, H, bias=False)
        self.fc2 = nn.Linear(H, D, bias=False)
        self.alpha = nn.Parameter(torch.zeros(()))
        self.sqrtD = math.sqrt(D)

    def forward(self, s):
        ns = s * (self.sqrtD / (s.norm(dim=-1, keepdim=True) + 1e-12))
        return ns + self.fc2(F.silu(self.fc1(ns))) * self.alpha


class DraftDecoderLayer(nn.Module):
    def __init__(self, D, H):
        super().__init__()
        self.block = DraftBlock(D, H)

    def forward(self, hidden_states, *a, **k):
        return self.block(hidden_states)


def load_merged(model_dir, dtype=torch.bfloat16, device="cuda", chain_hidden=2048):
    from transformers import AutoConfig, AutoModelForCausalLM
    from safetensors.torch import load_file
    cfg = AutoConfig.from_pretrained(model_dir)
    m = AutoModelForCausalLM.from_config(cfg).to(device).to(dtype).eval()
    n = cfg.num_hidden_layers
    for j in range(4):
        m.model.layers[n - 4 + j] = DraftDecoderLayer(cfg.hidden_size, chain_hidden).to(device, dtype=dtype)
    m.model.norm = nn.Identity()
    sd = load_file(f"{model_dir}/model.safetensors", device=device)
    sd = {k.replace("model.language_model.", "model.", 1) if k.startswith("model.language_model.") else k: v
          for k, v in sd.items()}
    if "lm_head.weight" not in sd:
        sd["lm_head.weight"] = sd["model.embed_tokens.weight"]  
    m.load_state_dict(sd, strict=True)
    return m