#!/home/jyhs/cse251b/.venv/bin/python3
# run  := python3 test.py
# dir  := .
# kid  :=

from dataclasses import dataclass
import inspect
import math

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


class CombinedOptimizer:
    def __init__(self, *optimizers):
        self.optimizers = [opt for opt in optimizers if opt is not None]
        self.param_groups = []
        for opt in self.optimizers:
            self.param_groups.extend(opt.param_groups)

    def step(self, *args, **kwargs):
        out = None
        for opt in self.optimizers:
            out = opt.step(*args, **kwargs)
        return out

    def zero_grad(self, *args, **kwargs):
        for opt in self.optimizers:
            opt.zero_grad(*args, **kwargs)

    def state_dict(self):
        return {"optimizers": [opt.state_dict() for opt in self.optimizers]}

    def load_state_dict(self, state_dict):
        states = state_dict["optimizers"]
        if len(states) != len(self.optimizers):
            raise ValueError(f"optimizer count mismatch: {len(states)} != {len(self.optimizers)}")
        for opt, state in zip(self.optimizers, states):
            opt.load_state_dict(state)


@dataclass
class LyraConfig:
    block_size: int = 1024
    vocab_size: int = 50304
    n_layer: int = 12
    n_head: int = 10
    n_embd: int = 640
    rope_base: float = 10000.0
    use_qk_norm: bool = True
    logit_softcap: float = 0.0
    mlp_hidden: int = 2048
    dropout: float = 0.0
    bias: bool = False


class ChannelRMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        x_fp = x.float()
        rms = x_fp.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x_fp * rms).to(x.dtype) * self.weight


def build_rotary_cache(head_dim: int, max_seq_len: int, base: float = 10000.0):
    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    t = torch.arange(max_seq_len, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)
    return freqs.cos(), freqs.sin()


def apply_rotary_embedding(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    cos = cos[None, None, :, :].to(x.dtype)
    sin = sin[None, None, :, :].to(x.dtype)
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class RotarySelfAttention(nn.Module):
    def __init__(self, config: LyraConfig):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        assert self.head_dim % 2 == 0, "head_dim must be even for RoPE"

        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.c_proj.NANOGPT_SCALE_INIT = 1

        self.use_qk_norm = config.use_qk_norm
        if self.use_qk_norm:
            self.q_norm = ChannelRMSNorm(self.head_dim)
            self.k_norm = ChannelRMSNorm(self.head_dim)

    def forward(self, x, cos, sin):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        if self.use_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        q = apply_rotary_embedding(q, cos[:T], sin[:T])
        k = apply_rotary_embedding(k, cos[:T], sin[:T])

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)


class GatedFeedForward(nn.Module):
    def __init__(self, config: LyraConfig):
        super().__init__()
        hidden = config.mlp_hidden
        self.c_gate = nn.Linear(config.n_embd, hidden, bias=False)
        self.c_up = nn.Linear(config.n_embd, hidden, bias=False)
        self.c_proj = nn.Linear(hidden, config.n_embd, bias=False)
        self.c_proj.NANOGPT_SCALE_INIT = 1

    def forward(self, x):
        return self.c_proj(F.silu(self.c_gate(x)) * self.c_up(x))


class DecoderLayer(nn.Module):
    def __init__(self, config: LyraConfig):
        super().__init__()
        self.ln_1 = ChannelRMSNorm(config.n_embd)
        self.attn = RotarySelfAttention(config)
        self.ln_2 = ChannelRMSNorm(config.n_embd)
        self.mlp = GatedFeedForward(config)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.ln_1(x), cos, sin)
        x = x + self.mlp(self.ln_2(x))
        return x


class LyraLM(nn.Module):
    def __init__(self, config: LyraConfig):
        super().__init__()
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),
            h=nn.ModuleList([DecoderLayer(config) for _ in range(config.n_layer)]),
            ln_f=ChannelRMSNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.transformer.wte.weight = self.lm_head.weight

        head_dim = config.n_embd // config.n_head
        cos, sin = build_rotary_cache(head_dim, config.block_size, config.rope_base)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            if hasattr(module, "NANOGPT_SCALE_INIT"):
                std *= (2 * self.config.n_layer) ** -0.5
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.size()
        assert T <= self.config.block_size, (
            f"sequence length {T} > block_size {self.config.block_size}"
        )
        x = self.transformer.wte(idx)
        for block in self.transformer.h:
            x = block(x, self.rope_cos, self.rope_sin)
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)
        if self.config.logit_softcap > 0:
            cap = self.config.logit_softcap
            logits = cap * torch.tanh(logits / cap)
        logits = logits[:, :, :50257]
        if targets is None:
            return logits
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss

    def configure_optimizers(
        self,
        weight_decay,
        learning_rate,
        betas,
        device_type,
        use_muon=False,
        muon_lr=4e-4,
        muon_momentum=0.95,
        muon_ns_steps=5,
        muon_nesterov=True,
        muon_adjust_lr_fn="original",
        embed_lr=0.0,
        scalar_lr=0.0,
    ):
        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}
        if use_muon:
            embed_names = ("transformer.wte", "lm_head")
            muon_items = [
                (n, p) for n, p in param_dict.items()
                if p.dim() == 2 and not any(name in n for name in embed_names)
            ]
            muon_ids = {id(p) for _, p in muon_items}
            adamw_items = [(n, p) for n, p in param_dict.items() if id(p) not in muon_ids]
            embed_params = [p for n, p in adamw_items if p.dim() >= 2]
            scalar_params = [p for n, p in adamw_items if p.dim() < 2]

            print(
                f"Muon parameter tensors: {len(muon_items)}, "
                f"with {sum(p.numel() for _, p in muon_items):,} parameters"
            )
            print(
                f"AdamW embed tensors: {len(embed_params)}, "
                f"with {sum(p.numel() for p in embed_params):,} parameters"
            )
            print(
                f"AdamW scalar tensors: {len(scalar_params)}, "
                f"with {sum(p.numel() for p in scalar_params):,} parameters"
            )

            muon_optimizer = torch.optim.Muon(
                [{"params": [p for _, p in muon_items], "lr": muon_lr, "initial_lr": muon_lr}],
                lr=muon_lr,
                weight_decay=weight_decay,
                momentum=muon_momentum,
                nesterov=muon_nesterov,
                ns_steps=muon_ns_steps,
                adjust_lr_fn=muon_adjust_lr_fn,
            )
            adamw_groups = []
            if embed_params:
                elr = embed_lr if embed_lr > 0 else learning_rate
                adamw_groups.append({
                    "params": embed_params,
                    "lr": elr,
                    "initial_lr": elr,
                    "weight_decay": 0.0,
                })
            if scalar_params:
                slr = scalar_lr if scalar_lr > 0 else learning_rate
                adamw_groups.append({
                    "params": scalar_params,
                    "lr": slr,
                    "initial_lr": slr,
                    "weight_decay": 0.0,
                })
            fused_available = "fused" in inspect.signature(torch.optim.AdamW).parameters
            use_fused = fused_available and device_type == "cuda"
            adamw_optimizer = torch.optim.AdamW(
                adamw_groups,
                lr=learning_rate,
                betas=betas,
                **(dict(fused=True) if use_fused else {}),
            )
            print(
                "using Muon + AdamW: "
                f"muon_lr={muon_lr}, adamw_lr={learning_rate}, fused AdamW={use_fused}"
            )
            return CombinedOptimizer(muon_optimizer, adamw_optimizer)

        decay_params = [p for p in param_dict.values() if p.dim() >= 2]
        nodecay_params = [p for p in param_dict.values() if p.dim() < 2]
        optim_groups = [
            {"params": decay_params, "weight_decay": weight_decay},
            {"params": nodecay_params, "weight_decay": 0.0},
        ]
        num_decay_params = sum(p.numel() for p in decay_params)
        num_nodecay_params = sum(p.numel() for p in nodecay_params)
        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
        fused_available = "fused" in inspect.signature(torch.optim.AdamW).parameters
        use_fused = fused_available and device_type == "cuda"
        extra_args = dict(fused=True) if use_fused else dict()
        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
        print(f"using fused AdamW: {use_fused}")
        return optimizer

    def estimate_mfu(self, fwdbwd_per_iter, dt):
        N = num_params(self)
        cfg = self.config
        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd // cfg.n_head, cfg.block_size
        flops_per_token = 6 * N + 12 * L * H * Q * T
        flops_per_fwdbwd = flops_per_token * T
        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
        flops_achieved = flops_per_iter * (1.0 / dt)
        flops_promised = 312e12
        return flops_achieved / flops_promised


def num_params(model: LyraLM) -> int:
    seen = set()
    total = 0
    for p in model.parameters():
        if id(p) in seen:
            continue
        seen.add(id(p))
        total += p.numel()
    return total


def strip_runtime_prefixes(state_dict):
    clean = {}
    for key, value in state_dict.items():
        for prefix in ("_orig_mod.", "module."):
            if key.startswith(prefix):
                key = key[len(prefix):]
        clean[key] = value
    return clean


def load_model(checkpoint_path: str, device: str = "cuda") -> nn.Module:
    cfg = LyraConfig()
    model = LyraLM(cfg)
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    state_dict = checkpoint["model"] if isinstance(checkpoint, dict) and "model" in checkpoint else checkpoint
    model.load_state_dict(strip_runtime_prefixes(state_dict), strict=True)
    model.to(device)
    model.eval()
    return model


if __name__ == "__main__":
    cfg = LyraConfig()
    m = LyraLM(cfg)
    n = num_params(m)
    print(f"LyraConfig = {cfg}")
    print(f"params: {n:,} ({n/1e6:.2f} M)")
    assert n < 100_000_000, "OVER 100M PARAM CAP"
    x = torch.randint(0, 50257, (2, 64))
    logits, loss = m(x, x)
    print(f"in: {tuple(x.shape)} out: {tuple(logits.shape)} loss: {loss.item():.3f}")
    print("ok")