"""
Chess Transformer Model for the Chess Challenge.

This module provides a GPT-style transformer architecture
designed to fit within the 1M parameter constraint.

Key improvements for legal move generation:
- Optimized architecture for move-level tokenization
- Better parameter distribution
- Support for board-aware training
"""

from __future__ import annotations

import math
from typing import Optional, Tuple, Union

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


def calculate_parameters(
    vocab_size: int,
    n_embd: int,
    n_layer: int,
    n_head: int,
    n_ctx: int,
    n_inner: int,
    tie_weights: bool = True,
) -> int:
    """
    Calculate the total number of parameters for a given configuration.
    """
    # Token embeddings: vocab_size * n_embd
    token_emb = vocab_size * n_embd

    # Position embeddings: n_ctx * n_embd
    pos_emb = n_ctx * n_embd

    # Per transformer layer:
    ln1 = 2 * n_embd
    attn_qkv = n_embd * 3 * n_embd + 3 * n_embd
    attn_out = n_embd * n_embd + n_embd
    ln2 = 2 * n_embd
    ffn_up = n_embd * n_inner + n_inner
    ffn_down = n_inner * n_embd + n_embd

    per_layer = ln1 + attn_qkv + attn_out + ln2 + ffn_up + ffn_down
    all_layers = n_layer * per_layer

    # Final layer norm
    final_ln = 2 * n_embd

    # LM head (shared if tie_weights)
    lm_head = 0 if tie_weights else vocab_size * n_embd

    total = token_emb + pos_emb + all_layers + final_ln + lm_head
    return total


def find_optimal_config(
    vocab_size: int,
    target_params: int = 980_000,
    max_params: int = 999_999,
    n_ctx: int = 256,
    tie_weights: bool = True,
) -> dict:
    """
    Find optimal model configuration that fits within the parameter budget.

    Prioritizes deeper models with moderate width for better pattern learning.
    """
    best_config = None
    best_params = 0

    # Search configurations - prioritize depth for sequential pattern learning
    configs_to_try = [
        # (n_embd, n_layer, n_head, ffn_mult) - deeper is better for sequence modeling
        (128, 12, 8, 2.0),  # Deep and narrow
        (128, 10, 8, 2.5),
        (112, 12, 8, 2.0),
        (120, 10, 8, 2.0),
        (128, 8, 8, 3.0),
        (112, 10, 8, 2.5),
        (96, 12, 8, 2.5),
        (128, 8, 8, 2.5),
        (120, 8, 8, 2.5),
        (112, 8, 8, 3.0),
        (96, 10, 8, 3.0),
        (128, 6, 8, 3.0),
        (112, 8, 8, 2.5),
        (96, 8, 8, 3.0),
    ]

    for n_embd, n_layer, n_head, ffn_mult in configs_to_try:
        if n_embd % n_head != 0:
            continue

        n_inner = int(n_embd * ffn_mult)

        params = calculate_parameters(
            vocab_size=vocab_size,
            n_embd=n_embd,
            n_layer=n_layer,
            n_head=n_head,
            n_ctx=n_ctx,
            n_inner=n_inner,
            tie_weights=tie_weights,
        )

        if params <= max_params and params > best_params:
            best_params = params
            best_config = {
                "n_embd": n_embd,
                "n_layer": n_layer,
                "n_head": n_head,
                "n_ctx": n_ctx,
                "n_inner": n_inner,
                "params": params,
            }

            if params >= target_params:
                return best_config

    return best_config


class ChessConfig(PretrainedConfig):
    """
    Configuration class for the Chess Transformer model.
    """

    model_type = "chess_transformer"

    def __init__(
        self,
        vocab_size: int = 1200,
        n_embd: int = 128,
        n_layer: int = 8,
        n_head: int = 8,
        n_ctx: int = 256,
        n_inner: Optional[int] = None,
        dropout: float = 0.1,
        layer_norm_epsilon: float = 1e-5,
        tie_weights: bool = True,
        pad_token_id: int = 0,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        use_cache: bool = True,
        **kwargs,
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs,
        )

        self.vocab_size = vocab_size
        self.n_embd = n_embd
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_ctx = n_ctx
        self.n_inner = n_inner if n_inner is not None else 3 * n_embd
        self.dropout = dropout
        self.layer_norm_epsilon = layer_norm_epsilon
        self.tie_weights = tie_weights
        self.tie_word_embeddings = bool(tie_weights)
        self.use_cache = use_cache


class MultiHeadAttention(nn.Module):
    """Multi-head self-attention with causal masking."""

    def __init__(self, config: ChessConfig):
        super().__init__()

        assert config.n_embd % config.n_head == 0

        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        self.scale = 1.0 / math.sqrt(self.head_dim)

        # Combined QKV projection
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)

        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        # Causal mask
        self.register_buffer(
            "bias",
            torch.tril(torch.ones(config.n_ctx, config.n_ctx)).view(
                1, 1, config.n_ctx, config.n_ctx
            ),
            persistent=False,
        )

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        B, T, C = x.size()

        # QKV projection
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=2)

        # Reshape for multi-head attention
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        # Attention scores
        att = (q @ k.transpose(-2, -1)) * self.scale

        # Causal mask
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))

        # Padding mask
        if attention_mask is not None:
            att = att.masked_fill(
                attention_mask.unsqueeze(1).unsqueeze(2) == 0, float("-inf")
            )

        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)

        # Apply attention
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)

        return self.resid_dropout(self.c_proj(y))


class FeedForward(nn.Module):
    """Feed-forward network with GELU activation."""

    def __init__(self, config: ChessConfig):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, config.n_inner)
        self.c_proj = nn.Linear(config.n_inner, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.c_fc(x)
        x = F.gelu(x)
        x = self.c_proj(x)
        return self.dropout(x)


class TransformerBlock(nn.Module):
    """Transformer block with pre-normalization."""

    def __init__(self, config: ChessConfig):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.attn = MultiHeadAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.mlp = FeedForward(config)

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        x = x + self.attn(self.ln_1(x), attention_mask=attention_mask)
        x = x + self.mlp(self.ln_2(x))
        return x


class ChessForCausalLM(PreTrainedModel):
    """
    Chess Transformer for next-move prediction.
    """

    config_class = ChessConfig
    base_model_prefix = "transformer"
    supports_gradient_checkpointing = True
    _keys_to_ignore_on_load_missing = ["lm_head.weight"]

    def __init__(self, config: ChessConfig):
        super().__init__(config)

        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.wpe = nn.Embedding(config.n_ctx, config.n_embd)
        self.drop = nn.Dropout(config.dropout)

        self.h = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.n_layer)
        ])

        self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        if config.tie_weights:
            self._tied_weights_keys = ["lm_head.weight"]

        self.post_init()

        if config.tie_weights:
            self.tie_weights()

    def get_input_embeddings(self) -> nn.Module:
        return self.wte

    def set_input_embeddings(self, new_embeddings: nn.Module):
        self.wte = new_embeddings
        if getattr(self.config, "tie_weights", False):
            self.tie_weights()

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

    def set_output_embeddings(self, new_embeddings: nn.Module):
        self.lm_head = new_embeddings

    def tie_weights(self):
        if getattr(self.config, "tie_weights", False) or getattr(
            self.config, "tie_word_embeddings", False
        ):
            self._tie_or_clone_weights(self.lm_head, self.wte)

    def _init_weights(self, module: nn.Module):
        """Initialize weights with small std for stability."""
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.LayerNorm):
            torch.nn.init.ones_(module.weight)
            torch.nn.init.zeros_(module.bias)

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        B, T = input_ids.size()
        device = input_ids.device

        if position_ids is None:
            position_ids = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)

        # Embeddings
        tok_emb = self.wte(input_ids)
        pos_emb = self.wpe(position_ids)
        x = self.drop(tok_emb + pos_emb)

        # Transformer blocks
        for block in self.h:
            x = block(x, attention_mask=attention_mask)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        # Loss computation
        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,
            )

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

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )

    @torch.no_grad()
    def generate_move(
        self,
        input_ids: torch.LongTensor,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
    ) -> int:
        """Generate the next move token."""
        self.eval()

        outputs = self(input_ids)
        logits = outputs.logits[:, -1, :]

        if temperature > 0:
            logits = logits / temperature

        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = float("-inf")

        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)

        return next_token.item()


# Register with Auto classes
from transformers import AutoConfig, AutoModelForCausalLM

AutoConfig.register("chess_transformer", ChessConfig)
AutoModelForCausalLM.register(ChessConfig, ChessForCausalLM)