from dataclasses import dataclass
import math
import diffdist
from typing import Optional

import torch
import torch.nn as nn
import numpy as np
from torch.nn import functional as F

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_cls_token: int = 60_001
    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_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_tags: int = 152
    t_text_lyrics: int = 1000
    t_audio: int = 3136
    t_memmap: int = 3008
    use_rotary_pos_emb: bool = False
    attention_type: str = "torch"  # "torch", "tao", "xformers"
    attention_sliding_window_size: int = -1
    n_unimodal: int = 6

    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_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, config.n_embd)
        self.transformer = nn.ModuleDict(model_dict)
        self.text_prj = Projection(config.n_embd, 128)
        self.audio_prj = Projection(config.n_embd, 128)
        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.coarse_vocab_size, bias=False)
                for _ in range(config.coarse_n_codebooks)
            ]
        )
        if TIE_WEIGHTS:
            for n in range(config.coarse_n_codebooks):
                self.transformer.wte_coarse[n].weight = self.lm_heads[n].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,
        text_offset=0,
        return_logits=False,
        last_only=True,
        tag_cls_ix=None,
    ):
        device = x.device
        b, ns, t = x.size()
        assert ns == 1 + self.config.coarse_n_codebooks

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

        # split text and audio
        x_text = x[:, 0, :text_offset]
        x_audio = x[:, 1:, text_offset:]

        # embed text
        x_text_emb = self.transformer.wte_text(x_text)
        x_text_emb = self.transformer.ln_text(x_text_emb)

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

        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 += pos_emb[:, :text_offset, :]
            x_audio_emb += pos_emb[:, text_offset:, :]

        # unimodal GPT
        x_text_emb = self.transformer.drop(x_text_emb)
        x_audio_emb = self.transformer.drop(x_audio_emb)
        for block in self.transformer.h[: self.config.n_unimodal]:
            x_text_emb = block(x_text_emb)
            x_audio_emb = block(x_audio_emb)

        # cls embeddings
        cls_text = F.normalize(self.text_prj(x_text_emb[:, tag_cls_ix, :]), dim=-1)
        cls_audio = F.normalize(self.audio_prj(x_audio_emb[:, self.config.t_audio, :]), dim=-1)

        # all gather for contrastive learning
        text_list = [torch.zeros_like(cls_text) for _ in range(torch.distributed.get_world_size())]
        text_list = diffdist.functional.all_gather(text_list, cls_text)
        cat_text = torch.cat(text_list, dim=0)
        audio_list = [torch.zeros_like(cls_audio) for _ in range(torch.distributed.get_world_size())]
        audio_list = diffdist.functional.all_gather(audio_list, cls_audio)
        cat_audio = torch.cat(audio_list, dim=0)

        logits_per_text = self.logit_scale * cat_text @ cat_audio.t()
        logits_per_audio = logits_per_text.t()
        labels = torch.arange(cat_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_text_emb, x_audio_emb[:, : self.config.t_audio, :]), 1)

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

        # x_emb (b, t, n_embd)

        x = self.transformer.ln_f(x)

        x = x[:, text_offset : text_offset + self.config.t_audio, :]
        if return_logits:
            if last_only:
                x = x[:, -1, :]
            coarse_logits_list = []
            for n in range(self.config.coarse_n_codebooks):
                coarse_logits_list.append(self.lm_heads[n](x))
            coarse_logits = torch.stack(coarse_logits_list).swapaxes(0, 1)
            return coarse_logits

        if Z_LOSS:
            loss_dict["z_loss"] = 0
        for n in range(self.config.coarse_n_codebooks):
            logits = self.lm_heads[n](x)
            loss_dict[f"coarse_{n}"] = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                y[:, n, :].reshape(-1),
                ignore_index=-1,
            )
            if Z_LOSS:
                loss_dict["z_loss"] += (torch.logsumexp(logits, dim=-1) ** 2).mean()

        return loss_dict

    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_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_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)