"""
Full definition of a LunarisCodex Language Model, all of it in this single file.
FIXED VERSION addressing numerical stability issues for training.

References:
1) the official GPT-2 TensorFlow implementation released by OpenAI:
https://github.com/openai/gpt-2/blob/master/src/model.py
2) huggingface/transformers PyTorch implementation:
https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
"""

import math
import inspect
from dataclasses import dataclass

import torch
import torch.nn as nn
from torch.nn import functional as F
import torch.utils.checkpoint as checkpoint

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

def apply_rotary_emb(xq, xk, freqs_cis):
    # Guarda o tipo original para converter de volta no final
    original_dtype = xq.dtype

    # Converte para float32 APENAS para a operação complexa
    xq_float = xq.float()
    xk_float = xk.float()

    # Agora, opera nos tensores float32
    xq_complex = torch.view_as_complex(xq_float.reshape(*xq_float.shape[:-1], -1, 2))
    xk_complex = torch.view_as_complex(xk_float.reshape(*xk_float.shape[:-1], -1, 2))

    # A multiplicação acontece em float32 complexo
    xq_out_complex = xq_complex * freqs_cis
    xk_out_complex = xk_complex * freqs_cis

    # Converte de volta para real (ainda em float32)
    xq_out = torch.view_as_real(xq_out_complex).flatten(3)
    xk_out = torch.view_as_real(xk_out_complex).flatten(3)

    # Converte o resultado final de volta para o dtype original (bfloat16)
    return xq_out.to(original_dtype), xk_out.to(original_dtype)

class RMSNorm(nn.Module):
    """ Fixed Root Mean Square Layer Normalization """
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
        self.bias = nn.Parameter(torch.zeros(dim))

    def _norm(self, x):
        # Upcast completo para float32, depois downcast do resultado
        x_f32 = x.to(torch.float32)
        variance = x_f32.pow(2).mean(-1, keepdim=True)
        return (x_f32 * torch.rsqrt(variance + self.eps)).to(x.dtype)

    def forward(self, x):
        # A função _norm agora lida com a lógica de dtype de forma robusta
        output = self._norm(x)
        return output * self.weight + self.bias

class CausalSelfAttention(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.d_model % config.n_heads == 0

        self.d_model = config.d_model
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads if config.n_kv_heads is not None else config.n_heads
        assert self.n_heads % self.n_kv_heads == 0, "Number of heads must be divisible by number of key-value heads"
        self.head_dim = self.d_model // self.n_heads

        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(self.d_model, self.d_model + 2 * self.n_kv_heads * self.head_dim, bias=config.bias)
        # output projection
        self.c_proj = nn.Linear(self.d_model, self.d_model, bias=config.bias)
        # regularization
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.dropout = config.dropout
        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
            # causal mask to ensure that attention is only applied to the left in the input sequence
            self.register_buffer("bias", torch.tril(torch.ones(config.max_seq_len, config.max_seq_len))
                                        .view(1, 1, config.max_seq_len, config.max_seq_len))

    def forward(self, x, freqs_cis):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (d_model)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v  = self.c_attn(x).split([self.d_model, self.n_kv_heads * self.head_dim, self.n_kv_heads * self.head_dim], dim=2)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) # (B, nh, T, hs)
        k = k.view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) # (B, n_kv_h, T, hs)
        v = v.view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) # (B, n_kv_h, T, hs)

        # apply rotary embeddings
        q, k = apply_rotary_emb(q, k, freqs_cis)

        # for GQA, repeat k and v heads to match q heads
        if self.n_kv_heads < self.n_heads:
            num_repeats = self.n_heads // self.n_kv_heads
            k = k.repeat_interleave(num_repeats, dim=1)
            v = v.repeat_interleave(num_repeats, dim=1)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        if self.flash:
            # efficient attention using Flash Attention CUDA kernels
            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        else:
            # manual implementation of attention
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

class SwiGLU(nn.Module):
    """ SwiGLU Gated Linear Unit Feed-Forward Network """
    def __init__(self, config):
        super().__init__()
        hidden_dim = 4 * config.d_model
        hidden_dim = int(2 * hidden_dim / 3)
        # custom multiple_of for efficiency
        multiple_of = 256
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.gate_up_proj = nn.Linear(config.d_model, 2 * hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
        x = self.down_proj(F.silu(gate) * up)
        x = self.dropout(x)
        return x

class Block(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.ln_1 = RMSNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = RMSNorm(config.d_model)
        self.mlp = SwiGLU(config)
        self.use_gradient_checkpointing = config.use_gradient_checkpointing

    def _block_forward(self, x, freqs_cis):
        """Helper function for the forward pass, used by checkpointing."""
        x = x + self.attn(self.ln_1(x), freqs_cis)
        x = x + self.mlp(self.ln_2(x))
        return x

    def forward(self, x, freqs_cis):
        if self.training and self.use_gradient_checkpointing:
            return checkpoint.checkpoint(self._block_forward, x, freqs_cis, use_reentrant=False)
        else:
            return self._block_forward(x, freqs_cis)

@dataclass
class LunarisCodexConfig:
    max_seq_len: int = 1024
    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
    n_layers: int = 12
    n_heads: int = 12
    n_kv_heads: int = None # Number of key-value heads for GQA
    d_model: int = 768
    dropout: float = 0.0
    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
    rope_theta: float = 10000.0
    use_gradient_checkpointing: bool = False # Whether to use gradient checkpointing to save memory

class LunarisCodex(nn.Module):

    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None
        assert config.max_seq_len is not None
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.d_model),
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layers)]),
            ln_f = RMSNorm(config.d_model),
        ))
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        # with weight tying when using torch.compile() some warnings get generated:
        # "UserWarning: functional_call was passed multiple values for tied weights.
        # This behavior is deprecated and will be an error in future versions"
        # not 100% sure what this is, so far seems to be harmless. TODO investigate
        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying

        # pre-compute and register the freqs_cis as a buffer
        freqs_cis = precompute_freqs_cis(
            self.config.d_model // self.config.n_heads,
            self.config.max_seq_len,
            self.config.rope_theta
        )
        self.register_buffer("freqs_cis", freqs_cis)

        # init all weights
        self.apply(self._init_weights)

        # report number of parameters
        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))

    def get_num_params(self, non_embedding=True):
        """
        Return the number of parameters in the model.
        For non-embedding count (default), the position embeddings get subtracted.
        The token embeddings would too, except due to the parameter sharing these
        params are actually used as weights in the final layer, so we include them.
        """
        n_params = sum(p.numel() for p in self.parameters())
        if non_embedding:
            # With RoPE, there are no learned position embeddings
            pass
        return n_params

    def _init_weights(self, module):
        """
        Initializes weights with a consistent, modern scheme.
        This method is called by `self.apply(self._init_weights)`.
        """
        # Calculate the standard deviation for scaled initialization
        std = 0.02 / math.sqrt(2 * self.config.n_layers)

        # Handle container modules first to apply specific initializations to their children.
        # This avoids double-initialization and ensures the correct scheme is used.
        if isinstance(module, CausalSelfAttention):
            # Scaled initialization for all linear layers in the attention block
            torch.nn.init.normal_(module.c_attn.weight, mean=0.0, std=std)
            torch.nn.init.normal_(module.c_proj.weight, mean=0.0, std=std)
            if module.c_attn.bias is not None:
                torch.nn.init.zeros_(module.c_attn.bias)
            if module.c_proj.bias is not None:
                torch.nn.init.zeros_(module.c_proj.bias)

        elif isinstance(module, SwiGLU):
            # Scaled initialization for the output projection layer (down_proj)
            torch.nn.init.normal_(module.down_proj.weight, mean=0.0, std=std)
            # Standard initialization for the fused input gate layer (gate_up_proj)
            torch.nn.init.normal_(module.gate_up_proj.weight, mean=0.0, std=0.02)

        # Handle leaf modules that are not part of the handled containers.
        elif isinstance(module, nn.Embedding):
            # Standard initialization for token embeddings
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

        elif isinstance(module, RMSNorm):
            # Initialize RMSNorm weights to 1 and bias to 0
            torch.nn.init.ones_(module.weight)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)

        # Note: We do NOT provide a generic `elif isinstance(module, nn.Linear)` rule.
        # This is intentional. The linear layers within CausalSelfAttention and SwiGLU
        # are already handled above. The only remaining nn.Linear is the `lm_head`,
        # whose weights are tied to `wte` and should not be re-initialized. This
        # strategy correctly avoids corrupting the weight tying.

    def forward(self, idx, targets=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.max_seq_len, f"Cannot forward sequence of length {t}, block size is only {self.config.max_seq_len}"

        # pre-computed rotary embeddings for the sequence
        freqs_cis = self.freqs_cis[:t]

        # forward the LunarisCodex model itself
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, d_model)
        x = self.transformer.drop(tok_emb)
        for block in self.transformer.h:
            x = block(x, freqs_cis)
        x = self.transformer.ln_f(x)

        if targets is not None:
            # if we are given some desired targets also calculate the loss
            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # inference-time mini-optimization: only forward the lm_head on the very last position
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None

        return logits, loss

    @classmethod
    def from_pretrained(cls, model_type, override_args=None):
        raise NotImplementedError("from_pretrained is not supported in this architecture.")

    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
        # start with all of the candidate parameters
        param_dict = {pn: p for pn, p in self.named_parameters()}
        # filter out those that do not require grad
        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
        nodecay_params = [p for n, p in param_dict.items() 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")
        # Create AdamW optimizer and use the fused version if it is available
        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):
        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
        # first estimate the number of flops we do per iteration.
        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
        N = self.get_num_params()
        cfg = self.config
        L, H, Q, T = cfg.n_layers, cfg.n_heads, cfg.d_model//cfg.n_heads, cfg.max_seq_len
        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
        # express our flops throughput as ratio of A100 bfloat16 peak flops
        flops_achieved = flops_per_iter * (1.0/dt) # per second
        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
        mfu = flops_achieved / flops_promised
        return mfu

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
        the sequence max_new_tokens times, feeding the predictions back into the model each time.
        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
        """
        for _ in range(max_new_tokens):
            # if the sequence context is growing too long we must crop it at max_seq_len
            idx_cond = idx if idx.size(1) <= self.config.max_seq_len else idx[:, -self.config.max_seq_len:]
            # forward the model to get the logits for the index in the sequence
            logits, _ = self(idx_cond)
            # pluck the logits at the final step and scale by desired temperature
            logits = logits[:, -1, :] / temperature
            # optionally crop the logits to only the top k options
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            # apply softmax to convert logits to (normalized) probabilities
            probs = F.softmax(logits, dim=-1)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)
            # append sampled index to the running sequence and continue
            idx = torch.cat((idx, idx_next), dim=1)

        return idx