import math

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

from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from configuration_supernova import SupernovaConfig


# ================================================================================
# EXACT CUSTOM SELF-ATTENTION
#
# Checkpoint keys:
#
# attn.qkv.weight
# attn.qkv.bias
# attn.out.weight
# attn.out.bias
#
# DO NOT replace this with nn.MultiheadAttention.
# ================================================================================

class SupernovaSelfAttention(nn.Module):

    def __init__(self, d_model, n_heads, dropout=0.0):

        super().__init__()

        assert d_model % n_heads == 0

        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        # EXACT checkpoint parameter names.
        self.qkv = nn.Linear(
            d_model,
            3 * d_model,
            bias=True,
        )

        self.out = nn.Linear(
            d_model,
            d_model,
            bias=True,
        )

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, causal_mask=None):

        B, T, C = x.shape

        qkv = self.qkv(x)

        q, k, v = qkv.chunk(3, dim=-1)

        q = q.view(
            B,
            T,
            self.n_heads,
            self.head_dim,
        ).transpose(1, 2)

        k = k.view(
            B,
            T,
            self.n_heads,
            self.head_dim,
        ).transpose(1, 2)

        v = v.view(
            B,
            T,
            self.n_heads,
            self.head_dim,
        ).transpose(1, 2)

        # [B, H, T, T]
        scores = torch.matmul(
            q,
            k.transpose(-2, -1),
        ) / math.sqrt(self.head_dim)

        if causal_mask is not None:

            scores = scores.masked_fill(
                causal_mask.view(1, 1, T, T),
                float("-inf"),
            )

        weights = F.softmax(
            scores,
            dim=-1,
        )

        weights = self.dropout(weights)

        y = torch.matmul(
            weights,
            v,
        )

        y = y.transpose(
            1,
            2,
        ).contiguous().view(
            B,
            T,
            C,
        )

        return self.out(y)


# ================================================================================
# TRANSFORMER BLOCK
# ================================================================================

class SupernovaBlock(nn.Module):

    def __init__(self, config):

        super().__init__()

        self.ln1 = nn.LayerNorm(
            config.d_model
        )

        self.attn = SupernovaSelfAttention(
            d_model=config.d_model,
            n_heads=config.n_heads,
            dropout=config.dropout,
        )

        self.ln2 = nn.LayerNorm(
            config.d_model
        )

        # EXACT checkpoint numbering:
        #
        # ff.0 = Linear
        # ff.1 = GELU
        # ff.2 = Linear
        #
        self.ff = nn.Sequential()

        self.ff.add_module(
            "0",
            nn.Linear(
                config.d_model,
                config.ffn_dim,
            )
        )

        self.ff.add_module(
            "1",
            nn.GELU(),
        )

        self.ff.add_module(
            "2",
            nn.Linear(
                config.ffn_dim,
                config.d_model,
            )
        )

    def forward(
        self,
        x,
        causal_mask=None,
    ):

        x = x + self.attn(
            self.ln1(x),
            causal_mask=causal_mask,
        )

        x = x + self.ff(
            self.ln2(x)
        )

        return x


# ================================================================================
# MODEL
# ================================================================================

class SupernovaForCausalLM(PreTrainedModel):

    config_class = SupernovaConfig
    base_model_prefix = "supernova"

    def __init__(self, config):

        super().__init__(config)

        self.token_embedding = nn.Embedding(
            config.vocab_size,
            config.d_model,
            padding_idx=config.pad_token_id,
        )

        self.position_embedding = nn.Embedding(
            config.max_seq_len,
            config.d_model,
        )

        self.blocks = nn.ModuleList([
            SupernovaBlock(config)
            for _ in range(config.n_layers)
        ])

        self.ln_final = nn.LayerNorm(
            config.d_model
        )

        self.lm_head = nn.Linear(
            config.d_model,
            config.vocab_size,
            bias=False,
        )

        # IMPORTANT:
        # Do not call post_init().
        # We need to preserve the original checkpoint exactly.

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        **kwargs,
    ):

        if input_ids is None:
            raise ValueError(
                "input_ids must be provided"
            )

        B, T = input_ids.shape

        if T > self.config.max_seq_len:
            raise ValueError(
                f"Sequence length {T} exceeds "
                f"max_seq_len {self.config.max_seq_len}"
            )

        positions = torch.arange(
            T,
            device=input_ids.device,
        ).unsqueeze(0)

        x = (
            self.token_embedding(input_ids)
            + self.position_embedding(positions)
        )

        causal_mask = torch.triu(
            torch.ones(
                T,
                T,
                device=input_ids.device,
                dtype=torch.bool,
            ),
            diagonal=1,
        )

        for block in self.blocks:

            x = block(
                x,
                causal_mask=causal_mask,
            )

        x = self.ln_final(x)

        logits = self.lm_head(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=self.config.pad_token_id,
            )

        return CausalLMOutput(
            loss=loss,
            logits=logits,
        )