from __future__ import annotations

import math
from typing import Any

import torch
from torch import Tensor, nn
from torch.nn import functional as F
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutput

from .configuration_tiny_transformer import TinyTransformerConfig


class CausalSelfAttention(nn.Module):
    def __init__(self, config: TinyTransformerConfig):
        super().__init__()
        self.heads = config.heads
        self.head_dim = config.dim // config.heads
        self.dropout = config.dropout
        self.qkv = nn.Linear(config.dim, 3 * config.dim)
        self.output = nn.Linear(config.dim, config.dim)
        self.residual_dropout = nn.Dropout(config.dropout)

    def forward(self, x: Tensor, attention_mask: Tensor | None = None) -> Tensor:
        batch, steps, width = x.shape
        qkv = self.qkv(x).view(batch, steps, 3, self.heads, self.head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)

        if attention_mask is None:
            attended = F.scaled_dot_product_attention(
                q,
                k,
                v,
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=True,
            )
        else:
            causal = torch.ones(
                (steps, steps), dtype=torch.bool, device=x.device
            ).tril()
            valid_keys = attention_mask[:, None, None, :].to(torch.bool)
            combined_mask = causal[None, None, :, :] & valid_keys
            attended = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=combined_mask,
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=False,
            )

        attended = attended.transpose(1, 2).contiguous().view(batch, steps, width)
        return self.residual_dropout(self.output(attended))


class FeedForward(nn.Module):
    def __init__(self, config: TinyTransformerConfig):
        super().__init__()
        self.input = nn.Linear(config.dim, config.ff_dim)
        self.output = nn.Linear(config.ff_dim, config.dim)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: Tensor) -> Tensor:
        return self.dropout(self.output(F.gelu(self.input(x))))


class TransformerBlock(nn.Module):
    def __init__(self, config: TinyTransformerConfig):
        super().__init__()
        self.attention_norm = nn.LayerNorm(config.dim)
        self.attention = CausalSelfAttention(config)
        self.feed_forward_norm = nn.LayerNorm(config.dim)
        self.feed_forward = FeedForward(config)

    def forward(self, x: Tensor, attention_mask: Tensor | None = None) -> Tensor:
        x = x + self.attention(self.attention_norm(x), attention_mask)
        return x + self.feed_forward(self.feed_forward_norm(x))


class TinyTransformerForCausalLM(PreTrainedModel, GenerationMixin):
    """Decoder-only TinyStories model compatible with AutoModelForCausalLM."""

    config_class = TinyTransformerConfig
    model_type = "tiny_transformer"
    base_model_prefix = "tiny_transformer"
    main_input_name = "input_ids"
    _no_split_modules = ["TransformerBlock"]
    _supports_sdpa = True

    def __init__(self, config: TinyTransformerConfig):
        super().__init__(config)
        self.token_embedding = nn.Embedding(config.vocab_size, config.dim)
        self.position_embedding = nn.Embedding(config.context, config.dim)
        self.embedding_dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList(
            [TransformerBlock(config) for _ in range(config.layers)]
        )
        self.final_norm = nn.LayerNorm(config.dim)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        self.post_init()

        residual_std = 0.02 / math.sqrt(2 * config.layers)
        for block in self.blocks:
            nn.init.normal_(block.attention.output.weight, mean=0.0, std=residual_std)
            nn.init.normal_(block.feed_forward.output.weight, mean=0.0, std=residual_std)

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)

    def get_input_embeddings(self) -> nn.Embedding:
        return self.token_embedding

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.token_embedding = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_output_embeddings(self, value: nn.Linear) -> None:
        self.lm_head = value

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        labels: Tensor | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **_: Any,
    ) -> CausalLMOutput | tuple[Tensor, ...]:
        _, steps = input_ids.shape
        if steps > self.config.context:
            raise ValueError(
                f"sequence has {steps} tokens; context limit is {self.config.context}"
            )
        if position_ids is None:
            if attention_mask is None:
                position_ids = torch.arange(steps, device=input_ids.device)[None, :]
            else:
                position_ids = attention_mask.long().cumsum(dim=-1) - 1
                position_ids.masked_fill_(attention_mask == 0, 0)

        x = self.embedding_dropout(
            self.token_embedding(input_ids) + self.position_embedding(position_ids)
        )
        collect_hidden = bool(
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        hidden_states: list[Tensor] | None = [x] if collect_hidden else None
        for block in self.blocks:
            x = block(x, attention_mask)
            if hidden_states is not None:
                hidden_states.append(x)
        logits = self.lm_head(self.final_norm(x))

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

        use_return_dict = (
            return_dict if return_dict is not None else self.config.return_dict
        )
        if not use_return_dict:
            result: tuple[Tensor, ...] = (logits,)
            if hidden_states is not None:
                result += (tuple(hidden_states),)  # type: ignore[assignment]
            return ((loss,) + result) if loss is not None else result
        return CausalLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=tuple(hidden_states) if hidden_states is not None else None,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        **_: Any,
    ) -> dict[str, Tensor | bool | None]:
        if input_ids.size(1) > self.config.context:
            input_ids = input_ids[:, -self.config.context :]
            if attention_mask is not None:
                attention_mask = attention_mask[:, -self.config.context :]
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "use_cache": False,
        }