from dataclasses import dataclass
import math
import numpy as np
from typing import Optional

import torch
import torch.nn as nn
from torch.nn import functional as F
from einops import rearrange

from .base import (
    Block,
    NormFunc,
    configure_optimizers,
    estimate_mfu,
    get_init_fn,
    init_weights_simple,
)


TIE_WEIGHTS = False
SIMPLE_INIT = True
Z_LOSS = True


class Projection(nn.Module):
    def __init__(self, input_dim, output_dim, dropout=0.5):
        super(Projection, self).__init__()
        self.linear_1 = nn.Linear(input_dim, output_dim, bias=False)
        self.linear_2 = nn.Linear(output_dim, output_dim, bias=False)
        self.layer_norm = nn.LayerNorm(output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        emb1 = self.linear_1(x)
        emb2 = self.dropout(self.linear_2(F.gelu(emb1)))
        return self.layer_norm(emb1 + emb2)


@dataclass
class GPTConfig:
    n_layer: int = 24
    n_head: int = 16  # query heads
    n_kv_head: Optional[int] = None
    d_head: int = 64
    block_size: int = 4288
    bias: bool = False
    dropout: float = 0.0
    text_vocab_size: int = 60_032
    text_codebook_size: int = 60_001
    text_pad_token: int = 1
    text_infer_token: int = 60_001
    text_end_token: int = 60_002
    text_cls_token: int = 60_013
    semantic_vocab_size: int = 4032
    semantic_codebook_size: int = 4000
    semantic_n_codebooks: int = 1
    semantic_pad_token: int = 4000
    semantic_infer_token: int = 4001
    semantic_cls_token: int = 4002
    semantic_rate_hz: int = 25
    semantic_shift_factor: int = 50
    coarse_vocab_size: int = 2112
    coarse_codebook_size: int = 2048
    coarse_n_codebooks: int = 12
    coarse_pad_token: int = 2048
    coarse_infer_token: int = 2049
    coarse_cls_token: int = 2050
    coarse_rate_hz: int = 25
    coarse_shift_factor: int = 5
    t_text: int = 150
    t_text_crop: int = 150
    t_audio: int = 3008
    t_memmap: int = 3008
    use_rotary_pos_emb: bool = False
    attention_type: str = "torch"  # "torch", "tao", "xformers"
    attention_sliding_window_size: int = -1

    def __post_init__(self):
        # default to multi head attention
        if self.n_kv_head is None:
            self.n_kv_head = self.n_head

    @property
    def n_embd(self):
        """The width of the residual stream"""
        return self.n_head * self.d_head

    @property
    def use_learned_pos_emb(self):
        return not self.use_rotary_pos_emb


class GPT(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.config = config

        model_dict = dict(
            wte_text=nn.Embedding(config.text_vocab_size, config.n_embd),
            ln_text=NormFunc(config.n_embd),
            wte_semantic=nn.ModuleList(
                [
                    nn.Embedding(config.semantic_vocab_size, config.n_embd)
                    for _ in range(config.semantic_n_codebooks)
                ]
            ),
            ln_semantic=NormFunc(config.n_embd),
            wte_coarse=nn.ModuleList(
                [
                    nn.Embedding(config.coarse_vocab_size, config.n_embd)
                    for _ in range(config.coarse_n_codebooks)
                ]
            ),
            ln_coarse=NormFunc(config.n_embd),
            drop=nn.Dropout(config.dropout),
            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f=NormFunc(config.n_embd),
        )
        if self.config.use_learned_pos_emb:
            model_dict["wpe"] = nn.Embedding(config.block_size + 2, config.n_embd)
        self.transformer = nn.ModuleDict(model_dict)
        self.text_prj = Projection(config.n_embd, 512)
        self.audio_prj = Projection(config.n_embd, 512)
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        self.contrastive_loss = nn.CrossEntropyLoss()
        self.lm_heads = nn.ModuleList(
            [nn.Linear(config.n_embd, config.text_vocab_size, bias=False) for _ in range(1)]
        )
        if TIE_WEIGHTS:
            for n in range(config.semantic_n_codebooks):
                self.transformer.wte_semantic[n].weight = self.lm_heads[n].weight
            for n in range(config.coarse_n_codebooks):
                n2 = config.semantic_n_codebooks + n
                self.transformer.wte_coarse[n].weight = self.lm_heads[n2].weight

        # init all weights
        if SIMPLE_INIT:
            self.apply(self._init_weights_simple)
            for pn, p in self.named_parameters():
                if pn.endswith("c_proj.weight"):
                    torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))
        else:
            self._init_weights()

        print(f"number of parameters: {self.get_num_params()/1e6:.0f}M")

    def forward(self, x, y=None, audio_offset=0, return_logits=False, last_only=True):
        device = x.device
        b, ns, t = x.size()
        assert ns == 1 + self.config.semantic_n_codebooks + self.config.coarse_n_codebooks

        if y is not None:
            assert t == self.config.block_size
            _, _, t2 = y.size()
            assert t2 == self.config.t_text

        # split text and audio tokens
        x_text = x[:, :, audio_offset:]
        x_audio = x[:, :, :audio_offset]

        # append [CLS] token
        x_text = nn.functional.pad(x_text, (0, 1), "constant", 0)
        x_text[:, :, -1] = x_text[:, :, -2]
        x_text[:, 0, self.config.t_text_crop] = self.config.text_cls_token
        x_audio = nn.functional.pad(x_audio, (0, 1), "constant", 0)
        x_audio[:, :, -1] = x_audio[:, :, -2]
        x_audio[:, 1, -1] = self.config.semantic_cls_token
        x_audio[:, 2:, -1] = self.config.coarse_cls_token

        # embed text
        x_text_emb = self.transformer.wte_text(x_text[:, 0, :])
        x_text_emb = self.transformer.ln_text(x_text_emb)

        # embed audio
        x_audio_emb = self.transformer.wte_semantic[0](x_audio[:, 1, :])
        x_audio_emb = self.transformer.ln_semantic(x_audio_emb)
        for n in range(self.config.coarse_n_codebooks):
            n2 = 1 + n + self.config.semantic_n_codebooks
            x_audio_emb += self.transformer.ln_coarse(self.transformer.wte_coarse[n](x_audio[:, n2, :]))

        # positional embedding
        if self.config.use_learned_pos_emb:
            pos = torch.arange(t, dtype=torch.long, device=device).unsqueeze(0)  # shape (1, t)
            pos_emb = self.transformer.wpe(pos)  # (1, t, n_embd)
            x_text_emb[:, :-1, :] += pos_emb[:, audio_offset:, :]
            x_audio_emb[:, :-1, :] += pos_emb[:, :audio_offset, :]

        # unimodal GPT
        x_text_emb = self.transformer.drop(x_text_emb)
        x_audio_emb = self.transformer.drop(x_audio_emb)
        x_cat_emb = rearrange([x_text_emb, x_audio_emb], "n b t c -> (n b) t c")

        for block in self.transformer.h[:-4]:
            x_cat_emb = block(x_cat_emb)
        x_text_emb, x_audio_emb = rearrange(x_cat_emb, "(n b) t c -> n b t c", n=2)
        x_text_emb = x_text_emb[:, : self.config.t_text_crop + 1, :]

        # contrastive loss
        cls_text = F.normalize(self.text_prj(x_text_emb[:, -1, :]), dim=-1)
        cls_audio = F.normalize(self.audio_prj(x_audio_emb[:, -1, :]), dim=-1)
        logits_per_text = self.logit_scale * cls_text @ cls_audio.t()
        logits_per_audio = logits_per_text.t()
        labels = torch.arange(cls_text.shape[0]).long().to(device)
        loss_dict = {}
        loss_dict["contrastive_loss"] = (
            self.contrastive_loss(logits_per_text, labels)
            + self.contrastive_loss(logits_per_audio, labels)
        ) / 2

        # concatenate
        x = torch.cat((x_audio_emb[:, :-1, :], x_text_emb[:, :-1, :]), 1)

        # multimodal GPT
        for block in self.transformer.h[-4:]:
            x = block(x)

        # x_emb (b, t, n_embd)

        x = self.transformer.ln_f(x)

        x = x[:, audio_offset : audio_offset + self.config.t_text, :]
        if return_logits:
            if last_only:
                x = x[:, -1, :]
            text_logits_list = []
            for n in range(1):
                text_logits_list.append(self.lm_heads[n](x))
            text_logits = torch.stack(text_logits_list).swapaxes(0, 1)
            return text_logits

        if Z_LOSS:
            loss_dict["z_loss"] = 0

        # projection
        logits = self.lm_heads[0](x)

        # flatten
        flat_logits = logits.reshape(-1, logits.size(-1))
        flat_y = y[:, 0, : self.config.t_text_crop].reshape(-1)
        mask = flat_y != 1
        loss_dict["text_0"] = F.cross_entropy(flat_logits[mask], flat_y[mask])
        # loss_dict["text_0"] = F.cross_entropy(
        #     logits.reshape(-1, logits.size(-1)), y[:, 0, :].reshape(-1)
        # )

        if Z_LOSS:
            loss_dict["z_loss"] += (torch.logsumexp(logits, dim=-1) ** 2).mean()
        return loss_dict, logits

    def get_num_params(self, non_embedding=True):
        n_params = sum(p.numel() for p in self.parameters())
        if non_embedding:
            for m in self.transformer.wte_semantic:
                n_params -= m.weight.numel()
            for m in self.transformer.wte_coarse:
                n_params -= m.weight.numel()
            if self.config.use_learned_pos_emb:
                n_params -= self.transformer.wpe.weight.numel()
        return n_params

    def _init_weights_simple(self, module):
        init_weights_simple(self, module)

    def _init_weights(self):
        # embeddings
        get_init_fn(self.config.n_embd, init_depth=None)(self.transformer.wte_text.weight)
        for module in self.transformer.wte_semantic:
            get_init_fn(self.config.n_embd, init_depth=None)(module.weight)
        for module in self.transformer.wte_coarse:
            get_init_fn(self.config.n_embd, init_depth=None)(module.weight)
        if self.config.use_learned_pos_emb:
            get_init_fn(self.config.n_embd, init_depth=None)(self.transformer.wpe.weight)
        # heads
        for module in self.lm_heads:
            get_init_fn(self.config.n_embd, init_depth=None)(module.weight)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        # attention blocks
        for layer_idx, block in enumerate(self.transformer.h):
            # mlp
            module = block.mlp.c_fc
            get_init_fn(self.config.n_embd, init_depth=layer_idx + 1)(module.weight)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
            module = block.mlp.c_proj
            get_init_fn(block.mlp.embd_inner, init_depth=layer_idx + 1)(module.weight)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
            # attention
            for module in [block.attn.c_attn, block.attn.c_proj]:
                get_init_fn(self.config.n_embd, init_depth=layer_idx + 1)(module.weight)
                if module.bias is not None:
                    torch.nn.init.zeros_(module.bias)

    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
        return configure_optimizers(self, weight_decay, learning_rate, betas, device_type)

    def estimate_mfu(self, fwdbwd_per_iter, dt):
        return estimate_mfu(self, fwdbwd_per_iter, dt)