"""CodeBharat-100M decoder-only Transformer.

The default configuration is deliberately matched to the prepared 100M
dataset:

* 49,152-token CodeBharat byte-level BPE vocabulary
* 1,024-token training sequences
* 100,679,424 trainable parameters with tied input/output embeddings

Architecture choices follow the conservative Llama/Qwen-style decoder recipe:
pre-norm RMSNorm, RoPE, SwiGLU, and grouped-query attention (GQA). The
implementation keeps full causal attention over all 1,024 positions because
code benefits from global context within a packed sequence.
"""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

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


BASE_DIR = Path(__file__).resolve().parent.parent


@dataclass
class ModelConfig:
    """Configuration for the default CodeBharat-100M model.

    The defaults are a single intentional architecture, not loose suggestions.
    They produce 100,679,424 trainable parameters when embeddings are tied.
    """

    vocab_size: int = 49_152
    hidden_size: int = 768
    num_layers: int = 10
    num_attention_heads: int = 12
    num_key_value_heads: int = 4
    intermediate_size: int = 2_048
    max_seq_len: int = 1_024
    rope_theta: float = 10_000.0
    rms_norm_eps: float = 1e-6
    initializer_range: float = 0.02
    attention_dropout: float = 0.0
    tie_word_embeddings: bool = True

    def __post_init__(self) -> None:
        positive_fields = {
            "vocab_size": self.vocab_size,
            "hidden_size": self.hidden_size,
            "num_layers": self.num_layers,
            "num_attention_heads": self.num_attention_heads,
            "num_key_value_heads": self.num_key_value_heads,
            "intermediate_size": self.intermediate_size,
            "max_seq_len": self.max_seq_len,
        }
        for name, value in positive_fields.items():
            if value <= 0:
                raise ValueError(f"{name} must be positive, got {value}")
        if self.hidden_size % self.num_attention_heads != 0:
            raise ValueError(
                "hidden_size must be divisible by num_attention_heads "
                f"({self.hidden_size} / {self.num_attention_heads})"
            )
        if self.num_attention_heads % self.num_key_value_heads != 0:
            raise ValueError(
                "num_attention_heads must be divisible by num_key_value_heads "
                f"({self.num_attention_heads} / {self.num_key_value_heads})"
            )
        if self.head_dim % 2:
            raise ValueError("head_dim must be even so RoPE can rotate pairs")
        if self.max_seq_len <= 1:
            raise ValueError("max_seq_len must be greater than one")
        if self.rope_theta <= 1.0:
            raise ValueError("rope_theta must be greater than one")
        if self.rms_norm_eps <= 0.0:
            raise ValueError("rms_norm_eps must be positive")
        if self.initializer_range <= 0.0:
            raise ValueError("initializer_range must be positive")
        if not 0.0 <= self.attention_dropout < 1.0:
            raise ValueError("attention_dropout must be in [0, 1)")

    @property
    def head_dim(self) -> int:
        """The dimensionality of each attention head."""
        return self.hidden_size // self.num_attention_heads

    @property
    def num_key_value_groups(self) -> int:
        """Number of query heads that share each key/value head."""
        return self.num_attention_heads // self.num_key_value_heads

    @property
    def estimated_parameter_count(self) -> int:
        """Return the exact count implied by this no-bias architecture."""
        embedding = self.vocab_size * self.hidden_size
        key_value_dim = self.num_key_value_heads * self.head_dim
        attention = self.hidden_size * (
            self.hidden_size + 2 * key_value_dim + self.hidden_size
        )
        swiglu = 3 * self.hidden_size * self.intermediate_size
        layer_norms = 2 * self.hidden_size
        transformer = self.num_layers * (attention + swiglu + layer_norms)
        final_norm = self.hidden_size
        output_head = 0 if self.tie_word_embeddings else embedding
        return embedding + transformer + final_norm + output_head

    def to_dict(self) -> dict[str, Any]:
        """Produce checkpoint-friendly, JSON-serializable configuration data."""
        return asdict(self)


class RMSNorm(nn.Module):
    """Root mean square normalization, computed safely in float32."""

    def __init__(self, hidden_size: int, eps: float) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.float()
        variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
        normalized = hidden_states * torch.rsqrt(variance + self.eps)
        return self.weight * normalized.to(input_dtype)


class RotaryEmbedding(nn.Module):
    """Rotary position embeddings with interleaved real/imaginary pairs."""

    def __init__(self, head_dim: int, theta: float) -> None:
        super().__init__()
        inv_freq = 1.0 / (
            theta
            ** (
                torch.arange(0, head_dim, 2, dtype=torch.float32)
                / head_dim
            )
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(
        self,
        position_ids: torch.Tensor,
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return cos/sin values shaped as (batch, sequence, head_dim / 2)."""
        positions = position_ids.to(dtype=torch.float32)
        inv_freq = self.inv_freq.to(device=position_ids.device)
        angles = positions.unsqueeze(-1) * inv_freq
        return angles.cos().to(dtype=dtype), angles.sin().to(dtype=dtype)


def apply_rotary_embedding(
    query_states: torch.Tensor,
    key_states: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply interleaved RoPE to Q and K tensors.

    Query and key have shape (batch, heads, sequence, head_dim). Cosine and
    sine have shape (batch, sequence, head_dim / 2).
    """

    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)

    def rotate(x: torch.Tensor) -> torch.Tensor:
        x_even = x[..., ::2]
        x_odd = x[..., 1::2]
        rotated = torch.stack(
            (x_even * cos - x_odd * sin, x_even * sin + x_odd * cos),
            dim=-1,
        )
        return rotated.flatten(start_dim=-2)

    return rotate(query_states), rotate(key_states)


class GroupedQueryAttention(nn.Module):
    """Causal self-attention with GQA and PyTorch SDPA kernels."""

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.num_attention_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = config.num_key_value_groups
        self.head_dim = config.head_dim
        self.attention_dropout = config.attention_dropout

        key_value_dim = self.num_key_value_heads * self.head_dim
        self.q_proj = nn.Linear(
            config.hidden_size,
            config.hidden_size,
            bias=False,
        )
        self.k_proj = nn.Linear(config.hidden_size, key_value_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, key_value_dim, bias=False)
        self.o_proj = nn.Linear(
            config.hidden_size,
            config.hidden_size,
            bias=False,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.shape

        query_states = self.q_proj(hidden_states).view(
            batch_size,
            seq_len,
            self.num_attention_heads,
            self.head_dim,
        )
        key_states = self.k_proj(hidden_states).view(
            batch_size,
            seq_len,
            self.num_key_value_heads,
            self.head_dim,
        )
        value_states = self.v_proj(hidden_states).view(
            batch_size,
            seq_len,
            self.num_key_value_heads,
            self.head_dim,
        )

        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)
        value_states = value_states.transpose(1, 2)
        query_states, key_states = apply_rotary_embedding(
            query_states,
            key_states,
            cos,
            sin,
        )

        # Manual expansion is portable across CPU, CUDA, and Apple MPS. It
        # avoids relying on device-specific SDPA GQA support.
        if self.num_key_value_groups > 1:
            key_states = key_states.repeat_interleave(
                self.num_key_value_groups,
                dim=1,
            )
            value_states = value_states.repeat_interleave(
                self.num_key_value_groups,
                dim=1,
            )

        dropout_p = self.attention_dropout if self.training else 0.0
        attn_output = F.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            attn_mask=None,
            dropout_p=dropout_p,
            is_causal=True,
        )
        attn_output = attn_output.transpose(1, 2).contiguous().view(
            batch_size,
            seq_len,
            -1,
        )
        return self.o_proj(attn_output)


class SwiGLUMLP(nn.Module):
    """Bias-free SwiGLU feed-forward network."""

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(
            config.hidden_size,
            config.intermediate_size,
            bias=False,
        )
        self.up_proj = nn.Linear(
            config.hidden_size,
            config.intermediate_size,
            bias=False,
        )
        self.down_proj = nn.Linear(
            config.intermediate_size,
            config.hidden_size,
            bias=False,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        gate = F.silu(self.gate_proj(hidden_states))
        return self.down_proj(gate * self.up_proj(hidden_states))


class DecoderLayer(nn.Module):
    """Pre-norm decoder block: attention residual, then SwiGLU residual."""

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.input_layernorm = RMSNorm(
            config.hidden_size,
            config.rms_norm_eps,
        )
        self.self_attn = GroupedQueryAttention(config)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size,
            config.rms_norm_eps,
        )
        self.mlp = SwiGLUMLP(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, cos, sin)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        return residual + hidden_states


class CodeBharat(nn.Module):
    """Dense, causal CodeBharat-100M language model.

    Inputs must be integer token IDs with shape (batch, sequence). Packed corpus
    shards are uint16 on disk; the data loader must convert each batch to int64
    or int32 before calling this model.
    """

    def __init__(self, config: ModelConfig | None = None) -> None:
        super().__init__()
        self.config = config if config is not None else ModelConfig()

        self.token_embeddings = nn.Embedding(
            self.config.vocab_size,
            self.config.hidden_size,
        )
        self.layers = nn.ModuleList(
            DecoderLayer(self.config) for _ in range(self.config.num_layers)
        )
        self.final_norm = RMSNorm(
            self.config.hidden_size,
            self.config.rms_norm_eps,
        )
        self.rotary_emb = RotaryEmbedding(
            self.config.head_dim,
            self.config.rope_theta,
        )
        self.lm_head: nn.Linear | None
        if self.config.tie_word_embeddings:
            self.lm_head = None
        else:
            self.lm_head = nn.Linear(
                self.config.hidden_size,
                self.config.vocab_size,
                bias=False,
            )

        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=self.config.initializer_range,
            )

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Return next-token logits with shape (batch, sequence, vocab_size)."""
        if input_ids.ndim != 2:
            raise ValueError(
                "input_ids must have shape (batch, sequence), "
                f"got {tuple(input_ids.shape)}"
            )
        if input_ids.dtype not in (torch.int32, torch.int64):
            raise TypeError(
                "input_ids must be torch.int32 or torch.int64; "
                f"got {input_ids.dtype}. Cast packed uint16 batches first."
            )

        batch_size, seq_len = input_ids.shape
        if seq_len > self.config.max_seq_len:
            raise ValueError(
                f"sequence length {seq_len} exceeds configured maximum "
                f"{self.config.max_seq_len}"
            )

        if position_ids is None:
            position_ids = torch.arange(
                seq_len,
                device=input_ids.device,
                dtype=torch.long,
            ).unsqueeze(0).expand(batch_size, -1)
        elif position_ids.shape != input_ids.shape:
            raise ValueError(
                "position_ids must have the same shape as input_ids, "
                f"got {tuple(position_ids.shape)} and {tuple(input_ids.shape)}"
            )
        elif position_ids.dtype not in (torch.int32, torch.int64):
            raise TypeError("position_ids must be torch.int32 or torch.int64")
        elif position_ids.numel() and position_ids.max().item() >= self.config.max_seq_len:
            raise ValueError(
                "position_ids contains a position outside the configured "
                f"maximum of {self.config.max_seq_len}"
            )

        hidden_states = self.token_embeddings(input_ids)
        cos, sin = self.rotary_emb(position_ids, hidden_states.dtype)

        for layer in self.layers:
            hidden_states = layer(hidden_states, cos, sin)
        hidden_states = self.final_norm(hidden_states)

        if self.lm_head is None:
            return F.linear(hidden_states, self.token_embeddings.weight)
        return self.lm_head(hidden_states)

    def count_parameters(self) -> int:
        """Return trainable parameters, respecting weight tying."""
        return sum(
            parameter.numel()
            for parameter in self.parameters()
            if parameter.requires_grad
        )


def _verify_packed_data_contract(config: ModelConfig) -> None:
    """Fail early if model defaults drift from the packed-data contract."""
    metadata_path = BASE_DIR / "data" / "tokenized" / "meta.json"
    if not metadata_path.exists():
        print(f"[smoke] packed-data metadata not found: {metadata_path}")
        return

    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    expected = {
        "vocab_size": config.vocab_size,
        "seq_len": config.max_seq_len,
        "dtype": "uint16",
    }
    actual = {name: metadata.get(name) for name in expected}
    if actual != expected:
        raise AssertionError(
            f"Packed-data contract mismatch: expected {expected}, got {actual}"
        )
    print("[smoke] packed-data contract OK")


def run_smoke_test() -> None:
    """Check parameter budget, data contract, causality, and gradients."""
    torch.manual_seed(7)

    config = ModelConfig()
    model = CodeBharat(config).eval()
    parameter_count = model.count_parameters()
    if parameter_count != config.estimated_parameter_count:
        raise AssertionError(
            "Parameter estimate mismatch: "
            f"{parameter_count:,} actual vs {config.estimated_parameter_count:,} expected"
        )
    if not 100_000_000 <= parameter_count <= 101_000_000:
        raise AssertionError(
            f"Default model is outside the 100M target: {parameter_count:,}"
        )

    with torch.inference_mode():
        input_ids = torch.randint(
            0,
            config.vocab_size,
            (1, 16),
            dtype=torch.long,
        )
        logits = model(input_ids)
    expected_shape = (1, 16, config.vocab_size)
    if logits.shape != expected_shape:
        raise AssertionError(
            f"Unexpected default-model logits shape: {tuple(logits.shape)}"
        )
    if not torch.isfinite(logits).all():
        raise AssertionError("Default-model logits contain non-finite values")
    _verify_packed_data_contract(config)

    # A small model makes causal and backward checks fast while using the same
    # components as the 100M model.
    tiny_config = ModelConfig(
        vocab_size=128,
        hidden_size=64,
        num_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        intermediate_size=192,
        max_seq_len=32,
    )
    tiny_model = CodeBharat(tiny_config).eval()
    tiny_input = torch.randint(0, tiny_config.vocab_size, (2, 12))
    altered_input = tiny_input.clone()
    altered_input[:, -1] = (altered_input[:, -1] + 1) % tiny_config.vocab_size

    with torch.inference_mode():
        original_logits = tiny_model(tiny_input)
        altered_logits = tiny_model(altered_input)
    torch.testing.assert_close(
        original_logits[:, :-1],
        altered_logits[:, :-1],
        rtol=0.0,
        atol=1e-6,
        msg="A future token changed an earlier causal prediction",
    )

    tiny_model.train()
    train_logits = tiny_model(tiny_input)
    loss = F.cross_entropy(
        train_logits[:, :-1].reshape(-1, tiny_config.vocab_size),
        tiny_input[:, 1:].reshape(-1),
    )
    loss.backward()
    if tiny_model.token_embeddings.weight.grad is None:
        raise AssertionError("Backward pass did not produce embedding gradients")

    print(f"[smoke] parameters: {parameter_count:,} ({parameter_count / 1e6:.2f}M)")
    print(f"[smoke] default forward: {tuple(logits.shape)}")
    print(f"[smoke] tiny causal + backward checks: OK (loss={loss.item():.4f})")
    print("[smoke] CodeBharat-100M model: PASS")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--smoke-test",
        action="store_true",
        help="run architecture and packed-data contract checks",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.smoke_test:
        run_smoke_test()
        return

    config = ModelConfig()
    print("CodeBharat-100M architecture")
    print(json.dumps(config.to_dict(), indent=2))
    print(f"Estimated parameters: {config.estimated_parameter_count:,}")
    print("Run with --smoke-test to execute forward, causal, and gradient checks.")


if __name__ == "__main__":
    main()