"""
Core Layers & Transformer Encoder Block for Nova 1.0
Includes RMSNorm, SwiGLU, CausalSelfAttention with RoPE, and Post-Norm initialization.
Grounded in Hierarchical Reasoning Model (HRM) Paper specifications.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from .rope import RotaryEmbedding


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # RMS normalization over last dimension without mean subtraction or bias
        variance = x.pow(2).mean(-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight


class SwiGLU(nn.Module):
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        # Gate, Up, and Down linear projections (bias=False as in LLaMA/HRM paper)
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU(x) = (Swish(x W_g) * x W_u) W_d
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, max_seq_len: int = 2048):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        assert self.head_dim * n_heads == d_model, "d_model must be divisible by n_heads"

        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_model, bias=False)
        self.v_proj = nn.Linear(d_model, d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)

        self.rotary_emb = RotaryEmbedding(dim=self.head_dim, max_seq_len=max_seq_len)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        batch_size, seq_len, _ = x.shape

        # Linear projections & reshape to (batch, n_heads, seq_len, head_dim)
        q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)

        # Apply Rotary Position Embedding
        q, k = self.rotary_emb(q, k, seq_len=seq_len)

        # Scaled dot-product attention with causal mask
        if mask is None:
            # Create lower-triangular causal mask
            causal_mask = torch.tril(torch.ones((seq_len, seq_len), device=x.device, dtype=torch.bool))
            attn_mask = torch.zeros((seq_len, seq_len), device=x.device, dtype=q.dtype)
            attn_mask.masked_fill_(~causal_mask, float("-inf"))
        else:
            attn_mask = mask

        # PyTorch SDPA (uses ROCm optimized kernels when available)
        attn_out = F.scaled_dot_product_attention(
            q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False
        )

        # Reshape back to (batch, seq_len, d_model)
        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.out_proj(attn_out)


class TransformerBlock(nn.Module):
    """
    Encoder Transformer Layer implementing Post-Norm architecture as described in HRM paper.
    """
    def __init__(self, d_model: int, n_heads: int, d_ff: int, max_seq_len: int = 2048):
        super().__init__()
        self.attn = CausalSelfAttention(d_model=d_model, n_heads=n_heads, max_seq_len=max_seq_len)
        self.norm1 = RMSNorm(d_model)
        self.mlp = SwiGLU(d_model=d_model, d_ff=d_ff)
        self.norm2 = RMSNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Post-Norm residual connections (x = Norm(x + Sublayer(x)))
        x = self.norm1(x + self.attn(x))
        x = self.norm2(x + self.mlp(x))
        return x


def init_weights_lecun(module: nn.Module):
    """
    Truncated LeCun Normal Initialization as specified in HRM Paper (Wang et al., 2025).
    """
    if isinstance(module, nn.Linear):
        std = 1.0 / math.sqrt(module.in_features)
        nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-2.0 * std, b=2.0 * std)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)