import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

class IdentityMoELayer(nn.Module):
    def __init__(self, original_mlp, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        in_d = original_mlp.gate_proj.in_features
        hid_d = original_mlp.gate_proj.out_features

        # Router
        self.router = nn.Linear(in_d, num_experts, bias=False)

        # Experts
        self.experts = nn.ModuleList()
        for i in range(num_experts):
            expert = nn.Module()
            expert.gate_proj = nn.Linear(in_d, hid_d, bias=False)
            expert.up_proj = nn.Linear(in_d, hid_d, bias=False)
            expert.down_proj = nn.Linear(hid_d, in_d, bias=False)
            self.experts.append(expert)

    def forward(self, x):
        b, s, d = x.shape
        x_flat = x.view(-1, d)
        logits = self.router(x_flat)
        probs = F.softmax(logits, dim=-1)
        top_probs, top_indices = torch.topk(probs, self.top_k, dim=-1)
        top_probs = top_probs / top_probs.sum(dim=-1, keepdim=True)

        final_output = torch.zeros_like(x_flat)
        for i in range(self.num_experts):
            mask = (top_indices == i)
            if mask.any():
                for k in range(self.top_k):
                    idx = mask[:, k]
                    if idx.any():
                        exp_out = self.experts[i].down_proj(
                            F.silu(self.experts[i].gate_proj(x_flat[idx])) * 
                            self.experts[i].up_proj(x_flat[idx])
                        )
                        final_output[idx] += exp_out * top_probs[idx, k].unsqueeze(-1)
        return final_output.view(b, s, d)

def load_nanosota(repo_id, device="cuda"):
    print(f"Loading NanoSOTA MoE from {repo_id}...")
    tokenizer = AutoTokenizer.from_pretrained(repo_id)

    # 1. Load Base Shell
    model = AutoModelForCausalLM.from_pretrained(
        "Qwen/Qwen2.5-0.5B-Instruct", 
        torch_dtype=torch.bfloat16, 
        device_map=device,
        trust_remote_code=True
    )

    # 2. Apply Surgery
    for i in range(len(model.model.layers)):
        original_mlp = model.model.layers[i].mlp
        model.model.layers[i].mlp = IdentityMoELayer(original_mlp, num_experts=8, top_k=2).to(dtype=torch.bfloat16, device=device)

    # 3. Load Weights
    from safetensors.torch import load_file
    from huggingface_hub import hf_hub_download

    file_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")
    state_dict = load_file(file_path)

    # Load with strict=False to handle metadata differences
    model.load_state_dict(state_dict, strict=False)
    print("✅ Model loaded successfully!")
    return model, tokenizer