from dataclasses import dataclass
import math
from typing import Optional

import torch
import torch.nn as nn
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


@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 = 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_rate_hz: int = 25
    coarse_shift_factor: int = 5
    t_text: int = 1152
    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

    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, config.n_embd)
        self.transformer = nn.ModuleDict(model_dict)
        self.lm_heads = nn.ModuleList(
            [
                nn.Linear(config.n_embd, config.semantic_vocab_size, bias=False)
                for _ in range(config.semantic_n_codebooks)
            ]
            + [
                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.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, score_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 <= t - score_offset

        # embed text
        x_emb = self.transformer.wte_text(x[:, 0, :])
        x_emb = self.transformer.ln_text(x_emb)
        # embed semantic
        for n in range(self.config.semantic_n_codebooks):
            x_emb += self.transformer.ln_semantic(self.transformer.wte_semantic[n](x[:, 1 + n, :]))
        # embed coarse
        for n in range(self.config.coarse_n_codebooks):
            n2 = 1 + n + self.config.semantic_n_codebooks
            x_emb += self.transformer.ln_coarse(self.transformer.wte_coarse[n](x[:, n2, :]))

        # x_emb (b, t, n_embd)

        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_emb += pos_emb

        x = self.transformer.drop(x_emb)

        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)

        x = x[:, score_offset : score_offset + t2, :]
        if return_logits:
            if last_only:
                x = x[:, -1, :]
            semantic_logits_list = []
            for n in range(self.config.semantic_n_codebooks):
                semantic_logits_list.append(self.lm_heads[n](x))
            semantic_logits = torch.stack(semantic_logits_list).swapaxes(0, 1)
            coarse_logits_list = []
            for n in range(self.config.coarse_n_codebooks):
                n2 = n + self.config.semantic_n_codebooks
                coarse_logits_list.append(self.lm_heads[n2](x))
            coarse_logits = torch.stack(coarse_logits_list).swapaxes(0, 1)
            return semantic_logits, coarse_logits

        loss_dict = {}
        if Z_LOSS:
            loss_dict["z_loss"] = 0
        for n in range(self.config.semantic_n_codebooks):
            logits = self.lm_heads[n](x)
            loss_dict[f"semantic_{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()
        for n in range(self.config.coarse_n_codebooks):
            n2 = n + self.config.semantic_n_codebooks
            logits = self.lm_heads[n2](x)
            loss_dict[f"coarse_{n}"] = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                y[:, n2, :].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_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)