import math

import torch
import torch.nn.functional as F
from einops import rearrange
from torch import Tensor, einsum, nn
from torch.distributions import Categorical
from torch.nn.utils.rnn import pad_sequence
from tqdm import trange


def _create_mask(l, device):
    """1 is valid region and 0 is invalid."""
    seq = torch.arange(max(l), device=device).unsqueeze(0)  # (1 t)
    stop = torch.tensor(l, device=device).unsqueeze(1)  # (b 1)
    return (seq < stop).float()  # (b t)


def _list_to_tensor(x_list: list[Tensor], pattern="t b c -> b t c"):
    """
    Args:
        x_list: [(t d)]
    Returns:
        x: (? ? ?)
        m: (? ? ?), same as x
    """
    l = list(map(len, x_list))
    x = rearrange(pad_sequence(x_list), pattern)
    m = _create_mask(l, x_list[0].device)
    m = m.t().unsqueeze(-1)  # (t b 1)
    m = rearrange(m, pattern)
    return x, m


class SinusodialEmbedding(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        exponent = torch.arange(self.d_half, dtype=torch.float32)
        exponent = exponent / self.d_half
        omega = torch.exp(-math.log(1e4) * exponent)
        self.omega: torch.Tensor
        self.register_buffer("omega", omega, persistent=False)

    @property
    def d_half(self):
        assert self.d_model % 2 == 0, "Only support even d_model."
        return self.d_model // 2

    def forward(self, x):
        """
        Args:
            x: (...)
        Returns:
            pe: (... d)
        """
        omega = self.omega
        while omega.dim() <= x.dim():
            omega = omega.unsqueeze(0)  # (... d)

        x = x.unsqueeze(-1)  # (... 1)
        x = omega * x
        x = torch.cat([x.sin(), x.cos()], dim=-1)

        return x

    def get_pe(self, n: int):
        """
        Args:
            n: int
        Returns:
            pe: (n d)
        """
        device = self.omega.device
        return self.forward(torch.arange(n, device=device))

    def add_pe(self, x):
        """
        Args:
            x: (b t c)
        """
        e = self.get_pe(x.shape[1])  # t d
        e = e[None]  # b t d
        x = x + e
        return x


class CasualAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        dim_head = d_model // num_heads
        self.num_heads = num_heads
        self.scale = dim_head**-0.5
        self.to_qkv = nn.Linear(d_model, d_model * 3, bias=False)
        self.to_out = nn.Linear(d_model, d_model)

    def forward(self, x, m):
        """
        Args:
            x: (b t c)
            m: (b t c), 1 is data, 0 is padding
        Returns:
            x: (b t c)
        """
        h = self.num_heads

        q, k, v = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, "b t (h d) -> b t h d", h=h), (q, k, v))

        e = einsum("b i h d, b j h d -> b i j h", q, k)
        e = e * self.scale

        kpm = m.unsqueeze(1) * m.unsqueeze(2)  # b i j 1
        kpm = kpm.squeeze(-1).tril().unsqueeze(-1)  # b i j 1

        e = e.masked_fill(kpm == 0, -torch.finfo(e.dtype).max)
        a = e.softmax(dim=2)  # Normalize on j, i.e. key

        o = einsum("b i j h, b j h d -> b i h d", a, v)
        o = o.flatten(-2)
        o = self.to_out(o)  # b t c

        o = o * m

        return o


class PrenormResidual(nn.Module):
    def __init__(self, block, d_model, dropout, requires_mask=False):
        super().__init__()
        self.block = block
        self.requires_mask = requires_mask
        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, m):
        opts = {"m": m} if self.requires_mask else {}
        x = x + self.dropout(self.block(self.norm(x) * m, **opts))
        return x * m


class Block(nn.Sequential):
    def __init__(self, d_model, num_heads, dropout):
        super().__init__()
        self.attn = PrenormResidual(
            CasualAttention(d_model, num_heads),
            d_model=d_model,
            dropout=dropout,
            requires_mask=True,
        )
        self.ffn = PrenormResidual(
            nn.Sequential(
                nn.Linear(d_model, d_model * 4),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(d_model * 4, d_model),
            ),
            d_model=d_model,
            dropout=dropout,
        )

    def forward(self, x, m):
        """
        Args:
            x: (b t c)
            m: (b t 1)
        """
        x = self.attn(x, m)
        x = self.ffn(x, m)
        return x


class ListEmbedding(nn.Embedding):
    def forward(self, x: list[Tensor]) -> list[Tensor]:
        if len(x) == 0:
            return []
        return super().forward(torch.cat(x)).split([*map(len, x)])


def _join(x: tuple[Tensor], sep: Tensor):
    """
    Args:
        x: (k t d)
        sep: (d)
    """
    ret = x[0]
    for i in range(1, len(x)):
        ret = torch.cat((ret, sep[None], x[i]), dim=0)
    return ret


class VALLEAR(nn.Module):
    def __init__(
        self,
        num_tokens: int,
        d_model=256,
        num_heads=8,
        dropout=0.1,
        num_layers=12,
    ):
        super().__init__()
        # Here, simply use num_tokens := max(num_text_tokens, num_prompt_tokens, num_output_tokens)
        self.text_emb = ListEmbedding(num_tokens, d_model)
        self.prompt_emb = ListEmbedding(num_tokens, d_model)
        self.output_emb = ListEmbedding(num_tokens, d_model)
        self.sin_emb = SinusodialEmbedding(d_model)
        self.sep = nn.Parameter(torch.randn(d_model))  # start of sequence token
        self.blocks = nn.ModuleList(
            [Block(d_model, num_heads, dropout) for _ in range(num_layers)]
        )
        self.fc = nn.Linear(d_model, num_tokens)

    @property
    def _ignore_index(self):
        return -100

    @staticmethod
    def _elementwise_merge_tensors(*l, sep):
        return [*map(lambda ts: _join(ts, sep), zip(*l))]

    def forward(
        self,
        text_list: list[Tensor],
        prompt_list: list[Tensor],
        output_list: list[Tensor],
        compute_loss: bool = True,
    ) -> Tensor:
        """
        Args:
            text_list: b t d
            prompt_list: b t d
            output_list: b t d
        Returns:
            y: logits of last output, b k
        """
        device = text_list[0].device

        x_list = self._elementwise_merge_tensors(
            self.text_emb(text_list),
            self.prompt_emb(prompt_list),
            self.output_emb(output_list),
            sep=self.sep,
        )

        x, m = _list_to_tensor(x_list)
        x = self.sin_emb.add_pe(x)

        for block in self.blocks:
            x = block(x, m)

        h = self.fc(x) * m

        h_list = [hi[:li] for hi, li in zip(h, map(len, x_list))]

        if compute_loss and len(output_list) > 0:
            y_list = self._elementwise_merge_tensors(
                text_list,
                prompt_list,
                output_list,
                sep=torch.tensor(self._ignore_index, device=device),
            )

            # make y_list earlier as it is future that is unknown
            for i in range(len(y_list)):
                y_list[i] = y_list[i].roll(-1, dims=0)
                y_list[i][-1] = self._ignore_index

            self.loss = dict(
                nll=F.cross_entropy(
                    torch.cat(h_list),
                    torch.cat(y_list),
                    ignore_index=self._ignore_index,
                )
            )

        logits = torch.stack([hi[-1] for hi in h_list])

        return logits

    @staticmethod
    def _prune(l: Tensor, stop_token: int | None):
        if stop_token is None:
            return l

        indices = (l == stop_token).nonzero()

        if len(indices) == 0:
            return l

        return l[: indices[0].item()]

    def generate(
        self,
        text_list: list[Tensor],
        prompt_list: list[Tensor],
        max_steps: int = 1000,
        stop_token: int | None = None,
    ):
        device = text_list[0].device
        output_list: list[Tensor] = [
            torch.zeros(0, device=device).long() for _ in text_list
        ]
        stopped = [False] * len(text_list)
        for _ in trange(max_steps):
            logits = self.forward(
                text_list,
                prompt_list,
                output_list,
                compute_loss=False,
            )
            o = Categorical(logits=logits).sample()
            for i, oi in enumerate(o):
                if oi.item() == stop_token:
                    stopped[i] = True
                output_list[i] = torch.cat([output_list[i], oi[None]])
            if all(stopped):
                break
        pruned = [self._prune(o, stop_token) for o in output_list]
        return pruned


def example_usage():
    import soundfile

    device = "cuda"

    test_qnt = torch.load("data/test/test.qnt.pt")[0, 0].to(device)
    num_qnts = 1024 + 1
    eoq = num_qnts - 1

    model = VALLEAR(num_qnts).to(device)

    text_list = [
        torch.tensor([1, 2, 3], device=device),
        torch.tensor([2, 3], device=device),
    ]

    prompt_list = [
        torch.tensor([1, 2, 3], device=device),
        torch.tensor([2, 3], device=device),
    ]

    output_list = [
        torch.tensor([1, 2, 3, eoq], device=device),
        torch.tensor([*test_qnt, eoq], device=device),
    ]

    out = model.generate(
        text_list,
        prompt_list,
        max_steps=200,
        stop_token=eoq,
    )

    print(out)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    for i in range(100):
        optimizer.zero_grad()
        _ = model(text_list, prompt_list, output_list)

        losses = model.loss
        sum(losses.values()).backward()
        optimizer.step()

        if i % 20 == 0:
            print(f"iter={i}, {losses}.")

    out = model.generate(
        text_list,
        prompt_list,
        max_steps=200,
        stop_token=eoq,
    )

    print(out)

    from ..emb.qnt import decode

    codes = rearrange(out[1], "t -> 1 1 t")
    wavs, sr = decode(codes)
    soundfile.write("data/test/test.ar.recon.wav", wavs.cpu()[0, 0], sr)


if __name__ == "__main__":
    example_usage()