import math

import torch
import torch.nn as nn

from giant.constants import COND_DIM, X_DIM


class SinusoidalEmbedding(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        assert dim % 2 == 0, "dim must be even"
        half = dim // 2
        freqs = torch.exp(
            -math.log(10000)
            * torch.arange(half, dtype=torch.float32)
            / max(half - 1, 1)
        )
        self.register_buffer("freqs", freqs)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        t = t.reshape(-1, 1).float()
        args = t * self.freqs.unsqueeze(0)  # (B, half)
        return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)


class ConditionEncoder(nn.Module):
    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        cont_dim: int = COND_DIM,
        emb_dim: int = 16,
        out_dim: int = 128,
    ) -> None:
        super().__init__()
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        in_dim = cont_dim + 2 * emb_dim
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, out_dim),
            nn.SiLU(),
            nn.Linear(out_dim, out_dim),
        )

    def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        pdg_e = self.pdg_emb(cond_cat[:, 0])
        mat_e = self.mat_emb(cond_cat[:, 1])
        x = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
        return self.mlp(x)


class ResBlock(nn.Module):
    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.linear1 = nn.Linear(dim, dim)
        self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
        self.act = nn.SiLU()
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        h = self.norm(x)
        h = self.linear1(h) + self.cond_proj(cond)
        h = self.act(h)
        h = self.dropout(h)
        h = self.linear2(h)
        return x + h


class DenoisingMLP(nn.Module):
    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(x_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, x_dim)

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)  # (B, time_dim)
        c_emb = self.cond_enc(cond_cont, cond_cat)  # (B, cond_out_dim)
        cond = torch.cat([t_emb, c_emb], dim=-1)  # (B, time_dim+cond_out_dim)
        x = self.input_proj(x_t)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)
