import math

from einops import rearrange
import torch
from torch import nn
import numpy as np
import torch.distributed

from base import ContinuousTransformer, ScaledSinusoidalEmbedding


class FourierFeatures(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        assert out_features % 2 == 0
        self.weight = nn.Parameter(torch.zeros(out_features // 2, in_features))

    def forward(self, input):
        f = 2 * math.pi * input @ self.weight.T
        return torch.cat([f.cos(), f.sin()], dim=-1)


class SemanticConditioner(nn.Module):
    def __init__(
        self,
        vocab_size: int = 4001,
        input_dim: int = 768,
        output_dim: int = 1536,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.embedding = torch.nn.Embedding(self.vocab_size, input_dim)
        self.proj_out = nn.Linear(input_dim, output_dim)
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)

    def _initialize(self, centroid_path):
        centroids = np.load(centroid_path)[0]
        # create a new tensor with an extra row for the pad token
        centroids = np.concatenate([centroids, np.random.randn(1, centroids.shape[1])], axis=0)
        assert self.embedding.weight.shape == centroids.shape
        self.embedding.weight.data.copy_(torch.tensor(centroids, dtype=torch.float32))

    def forward(self, x):
        x = self.embedding(x)
        x = self.proj_out(x)
        x = x + self.pos_embedding(x)
        return x  # [batch, sequence, channels]


class TextConditioner(nn.Module):
    def __init__(
        self,
        vocab_size: int = 60001,
        output_dim: int = 1536,
    ):
        super().__init__()
        self.embedding = torch.nn.Embedding(vocab_size, output_dim)
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)

    def forward(self, x):
        x = self.embedding(x)
        x = x + self.pos_embedding(x)
        return x  # [batch, sequence, channels]


class LatentContextConditioner(nn.Module):
    def __init__(
        self,
        io_channels: int = 128,
        output_dim: int = 1536,
    ):
        super().__init__()
        self.proj_out = nn.Linear(io_channels, output_dim)
        self.pos_embedding = ScaledSinusoidalEmbedding(output_dim)

    def forward(self, x):
        x = self.proj_out(x)
        x = x + self.pos_embedding(x)
        return x


class DiffusionTransformer(nn.Module):
    def __init__(
        self,
        io_hz=100,
        io_channels=128,
        embed_dim=1536,
        depth=24,
        n_heads=24,
        qk_norm=True,
        block_size=3000,
        cond_semantic_n_vocab=4001,
        cond_semantic_len=750,
        cond_text_n_vocab=60001,
        cond_text_len=2560,
        ctx_len=None,
    ):
        super().__init__()
        self.io_hz = io_hz
        self.io_channels = io_channels
        self.embed_dim = embed_dim
        self.depth = depth
        self.n_heads = n_heads
        self.qk_norm = qk_norm
        self.block_size = block_size
        self.cond_semantic_n_vocab = cond_semantic_n_vocab
        self.cond_semantic_len = cond_semantic_len
        self.cond_text_n_vocab = cond_text_n_vocab
        self.cond_text_len = cond_text_len

        self.semantic_conditioner = SemanticConditioner(
            vocab_size=cond_semantic_n_vocab,
            output_dim=embed_dim,
        )
        self.text_conditioner = TextConditioner(vocab_size=cond_text_n_vocab, output_dim=embed_dim)
        self.ctx_len = ctx_len
        if ctx_len is not None:
            self.ctx_conditioner = LatentContextConditioner(
                io_channels=io_channels, output_dim=embed_dim
            )
            self.vae_pad_embed = nn.Parameter(torch.zeros(io_channels))

        # Timestep embeddings
        timestep_features_dim = 256
        self.timestep_features = FourierFeatures(1, timestep_features_dim)
        self.to_timestep_embed = nn.Sequential(
            nn.Linear(timestep_features_dim, embed_dim, bias=True),
            nn.SiLU(),
            nn.Linear(embed_dim, embed_dim, bias=True),
        )

        # Conditioning tokens
        self.to_cond_embed = nn.Sequential(
            nn.Linear(embed_dim, embed_dim, bias=False),
            nn.SiLU(),
            nn.Linear(embed_dim, embed_dim, bias=False),
        )

        # Transformer
        assert embed_dim % n_heads == 0
        dim_heads = embed_dim // n_heads
        self.transformer = ContinuousTransformer(
            embed_dim,
            depth,
            dim_heads=dim_heads,
            dim_in=io_channels,
            dim_out=io_channels,
            causal=False,
            qk_norm=qk_norm,
        )
        self.preprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)
        self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)

    @property
    def device(self):
        return next(self.parameters()).device

    @property
    def dtype(self):
        return next(self.parameters()).dtype

    def _initialize(self, semantic_centroid_path=None):
        # Note: this is simple, but so far just as good as other inits
        nn.init.zeros_(self.preprocess_conv.weight)
        nn.init.zeros_(self.postprocess_conv.weight)
        nn.init.normal_(self.timestep_features.weight)
        for layer in self.transformer.layers:
            nn.init.zeros_(layer.ff.ff[1].weight)
            nn.init.zeros_(layer.ff.ff[1].bias)
            nn.init.zeros_(layer.self_attn.to_out.weight)
            nn.init.zeros_(layer.cross_attn.to_out.weight)
        with torch.no_grad():
            for param in self.parameters():
                param *= 0.5
        if semantic_centroid_path is not None:
            self.semantic_conditioner._initialize(semantic_centroid_path)

    def _forward(
        self,
        x,
        t,
        cross_attn_cond=None,
    ):
        if cross_attn_cond is not None:
            cross_attn_cond = self.to_cond_embed(cross_attn_cond)

        # Get the batch of timestep embeddings
        timestep_embed = self.to_timestep_embed(self.timestep_features(t[:, None]))  # (b, embed_dim)
        prepend_inputs = timestep_embed.unsqueeze(1)
        prepend_length = prepend_inputs.shape[1]

        x = self.preprocess_conv(x) + x

        x = rearrange(x, "b c t -> b t c")

        output = self.transformer(
            x,
            context=cross_attn_cond,
            prepend_embeds=prepend_inputs,
        )

        output = rearrange(output, "b t c -> b c t")[:, :, prepend_length:]

        output = self.postprocess_conv(output) + output

        return output

    def forward(
        self,
        x,
        t,
        text_codes=None,
        semantic_codes=None,
        ctx_vae=None,
        ctx_mask=None,
        future_ctx_vae=None,
        future_ctx_mask=None,
    ):
        assert text_codes is not None
        assert semantic_codes is not None
        assert torch.all(text_codes < self.cond_text_n_vocab), text_codes.max()
        assert torch.all(semantic_codes < self.cond_semantic_n_vocab)
        text_embeds = self.text_conditioner(text_codes)
        semantic_embeds = self.semantic_conditioner(semantic_codes)
        cross_attn_cond = torch.concat([text_embeds, semantic_embeds], dim=1)
        if self.ctx_len is not None:
            # ctx_vae: [batch, channels, seq_len]
            # ctx_mask: [batch, seq_len]
            # Expand vae_pad_embed to match input shape
            expanded_vae_pad = self.vae_pad_embed.view(1, -1, 1).expand(
                ctx_vae.shape[0], -1, ctx_vae.shape[2]
            )
            # Invert the mask (1 where we want to replace, 0 elsewhere)
            inverted_ctx_mask = (~ctx_mask).float().unsqueeze(1)
            # Use the mask to blend the original input with the pad embedding
            ctx_vae_input = (
                ctx_vae * ctx_mask.float().unsqueeze(1) + expanded_vae_pad * inverted_ctx_mask
            )
            ctx_vae_input = ctx_vae_input.permute(0, 2, 1).to(x.dtype)
            ctx_embeds = self.ctx_conditioner(ctx_vae_input)

            cross_attn_cond = torch.concat([cross_attn_cond, ctx_embeds], dim=1)

        return self._forward(x, t, cross_attn_cond=cross_attn_cond)

    def forward_inference(
        self,
        x,
        t,
        cross_attn_cond=None,
        no_text_cross_attn_cond=None,
        no_tag_cross_attn_cond=None,
        no_sem_cross_attn_cond=None,
        no_ctx_cross_attn_cond=None,
        text_cfg_scale=1.0,
        tag_cfg_scale=1.0,
        sem_cfg_scale=1.0,
        ctx_cfg_scale=1.0,
    ):
        assert cross_attn_cond is not None
        assert no_text_cross_attn_cond is not None
        has_text_cfg = text_cfg_scale != 1.0 and no_text_cross_attn_cond is not None
        has_tag_cfg = tag_cfg_scale != 1.0 and no_tag_cross_attn_cond is not None
        has_sem_cfg = sem_cfg_scale != 1.0 and no_sem_cross_attn_cond is not None
        has_ctx_cfg = ctx_cfg_scale != 1.0 and no_ctx_cross_attn_cond is not None
        n_rep = 1 + int(has_text_cfg) + int(has_tag_cfg) + int(has_sem_cfg) + int(has_ctx_cfg)

        # Classifier-free guidance

        # Concatenate conditioned and unconditioned inputs on the batch dimension
        batch_inputs = torch.cat([x] * n_rep, dim=0)
        batch_timestep = torch.cat([t] * n_rep, dim=0)

        # Handle CFG for cross-attention conditioning
        if has_text_cfg:
            cross_attn_cond = torch.cat(
                [
                    cross_attn_cond,
                    no_text_cross_attn_cond,
                ],
                dim=0,
            )
        if has_tag_cfg:
            cross_attn_cond = torch.cat([cross_attn_cond, no_tag_cross_attn_cond], dim=0)
        if has_sem_cfg:
            cross_attn_cond = torch.cat([cross_attn_cond, no_sem_cross_attn_cond], dim=0)
        if has_ctx_cfg:
            cross_attn_cond = torch.cat([cross_attn_cond, no_ctx_cross_attn_cond], dim=0)

        batch_output = self._forward(
            batch_inputs,
            batch_timestep,
            cross_attn_cond=cross_attn_cond,
        )

        cfg_list = []
        if has_text_cfg:
            cfg_list.append(text_cfg_scale)
        if has_tag_cfg:
            cfg_list.append(tag_cfg_scale)
        if has_sem_cfg:
            cfg_list.append(sem_cfg_scale)
        if has_ctx_cfg:
            cfg_list.append(ctx_cfg_scale)
        if len(cfg_list) == 4:
            cond_output, no_out_1, no_out_2, no_out_3, no_out_4 = torch.chunk(batch_output, 5, dim=0)
            batch_output = (
                cond_output
                + (cond_output - no_out_1) * (cfg_list[0] - 1.0)
                + (cond_output - no_out_2) * (cfg_list[1] - 1.0)
                + (cond_output - no_out_3) * (cfg_list[2] - 1.0)
                + (cond_output - no_out_4) * (cfg_list[3] - 1.0)
            )
        elif len(cfg_list) == 3:
            cond_output, no_out_1, no_out_2, no_out_3 = torch.chunk(batch_output, 4, dim=0)
            batch_output = (
                cond_output
                + (cond_output - no_out_1) * (cfg_list[0] - 1.0)
                + (cond_output - no_out_2) * (cfg_list[1] - 1.0)
                + (cond_output - no_out_3) * (cfg_list[2] - 1.0)
            )
        elif len(cfg_list) == 2:
            cond_output, no_out_1, no_out_2 = torch.chunk(batch_output, 3, dim=0)
            batch_output = (
                cond_output
                + (cond_output - no_out_1) * (cfg_list[0] - 1.0)
                + (cond_output - no_out_2) * (cfg_list[1] - 1.0)
            )
        elif len(cfg_list) == 1:
            cond_output, no_out_1 = torch.chunk(batch_output, 2, dim=0)
            batch_output = cond_output + (cond_output - no_out_1) * (cfg_list[0] - 1.0)
        else:
            batch_output = batch_output

        return batch_output