"""
SimpleStories Hindi - Custom HuggingFace Model

A small LLaMA-style decoder-only Transformer trained from scratch on Hindi stories.
Architecture: RoPE + SwiGLU + RMSNorm + FlashAttention (via scaled_dot_product_attention)

Usage:
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained("your-repo", trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained("your-repo", trust_remote_code=True)

    inputs = tokenizer("एक समय की बात है", return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.8)
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast


# ──────────────────────────────────────────────────────────────────────────────
# Config
# ──────────────────────────────────────────────────────────────────────────────
class SimpleStoriesConfig(PretrainedConfig):
    """Configuration for the SimpleStories Hindi model."""
    model_type = "simple_stories"

    def __init__(
        self,
        vocab_size=6000,
        max_seq_len=512,
        d_model=384,
        n_layers=10,
        n_heads=8,
        multiple_of=64,
        norm_eps=1e-5,
        dropout=0.0,
        pad_token_id=0,
        unk_token_id=1,
        bos_token_id=2,
        eos_token_id=3,
        **kwargs
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            unk_token_id=unk_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs
        )
        self.vocab_size = vocab_size
        self.max_seq_len = max_seq_len
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.multiple_of = multiple_of
        self.norm_eps = norm_eps
        self.dropout = dropout

        # HuggingFace standard aliases (needed by DynamicCache, generation, etc.)
        self.num_hidden_layers = n_layers
        self.hidden_size = d_model
        self.num_attention_heads = n_heads


# ──────────────────────────────────────────────────────────────────────────────
# RoPE helpers
# ──────────────────────────────────────────────────────────────────────────────
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def _precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(end, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)
    # duplicate so shape is [end, dim] instead of [end, dim//2]
    cos = torch.cat([cos, cos], dim=-1)
    sin = torch.cat([sin, sin], dim=-1)
    return cos, sin


def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    T = x.shape[2]
    cos_t = cos[:T, :].unsqueeze(0).unsqueeze(1)  # [1, 1, T, head_dim]
    sin_t = sin[:T, :].unsqueeze(0).unsqueeze(1)
    return (x * cos_t) + (_rotate_half(x) * sin_t)


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.pow(2).mean(-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight


class FeedForward(nn.Module):
    """SwiGLU feed-forward block."""
    def __init__(self, dim: int, hidden_dim: int = None, multiple_of: int = 32, dropout: float = 0.0):
        super().__init__()
        if hidden_dim is None:
            hidden_dim = int(2 * 4 * dim / 3)
            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(dim, hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout) if dropout > 0.0 else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.silu(self.w1(x)) * self.w2(x)
        out = self.w3(out)
        if self.dropout is not None:
            out = self.dropout(out)
        return out


class Attention(nn.Module):
    def __init__(self, config: SimpleStoriesConfig):
        super().__init__()
        self.n_heads = config.n_heads
        self.d_model = config.d_model
        self.head_dim = config.d_model // config.n_heads

        self.wq = nn.Linear(config.d_model, config.d_model, bias=False)
        self.wk = nn.Linear(config.d_model, config.d_model, bias=False)
        self.wv = nn.Linear(config.d_model, config.d_model, bias=False)
        self.wo = nn.Linear(config.d_model, config.d_model, bias=False)
        self.dropout_p = config.dropout

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        q = _apply_rotary_emb(q, cos, sin)
        k = _apply_rotary_emb(k, cos, sin)

        dropout_p = self.dropout_p if self.training else 0.0
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=dropout_p, is_causal=True)

        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.wo(out)


class TransformerBlock(nn.Module):
    def __init__(self, config: SimpleStoriesConfig):
        super().__init__()
        self.attention = Attention(config)
        self.feed_forward = FeedForward(dim=config.d_model, multiple_of=config.multiple_of, dropout=config.dropout)
        self.attention_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        x = x + self.attention(self.attention_norm(x), cos, sin)
        x = x + self.feed_forward(self.ffn_norm(x))
        return x


# ──────────────────────────────────────────────────────────────────────────────
# Main model
# ──────────────────────────────────────────────────────────────────────────────
class SimpleStoriesForCausalLM(PreTrainedModel, GenerationMixin):
    """LLaMA-style causal language model for Hindi story generation."""
    config_class = SimpleStoriesConfig
    _tied_weights_keys = ["output.weight"]

    def __init__(self, config: SimpleStoriesConfig):
        super().__init__(config)
        self.config = config
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.d_model)
        self.dropout = nn.Dropout(config.dropout) if config.dropout > 0.0 else None
        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
        self.norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.output = nn.Linear(config.d_model, config.vocab_size, bias=False)
        # weight tying
        self.tok_embeddings.weight = self.output.weight

        self.post_init()

    # Required by PreTrainedModel — we skip re-init since we load trained weights.
    def _init_weights(self, module):
        pass

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: torch.Tensor = None,
        labels: torch.LongTensor = None,
        past_key_values=None,
        use_cache: bool = False,
        return_dict: bool = True,
        **kwargs,
    ):
        B, T = input_ids.shape

        x = self.tok_embeddings(input_ids)
        if self.dropout is not None:
            x = self.dropout(x)

        # Compute RoPE embeddings on the fly (avoids meta-device buffer issues)
        head_dim = self.config.d_model // self.config.n_heads
        cos, sin = _precompute_freqs_cis(dim=head_dim, end=T)
        cos = cos.to(x.device, dtype=x.dtype)
        sin = sin.to(x.device, dtype=x.dtype)

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

        x = self.norm(x)
        logits = self.output(x)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                labels.view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)

        return CausalLMOutputWithPast(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        return {"input_ids": input_ids, "use_cache": False}