import torch.nn as nn
from models.positional_embeddings import FourierEmbedding, PositionalEmbedding
from models.networks.transformers import FusedMLP
import torch
import torch.nn.functional as F
import numpy as np
from einops import rearrange


class TimeEmbedder(nn.Module):
    def __init__(
        self,
        noise_embedding_type: str,
        dim: int,
        time_scaling: float,
        expansion: int = 4,
    ):
        super().__init__()
        self.encode_time = (
            PositionalEmbedding(num_channels=dim, endpoint=True)
            if noise_embedding_type == "positional"
            else FourierEmbedding(num_channels=dim)
        )
        self.time_scaling = time_scaling
        self.map_time = nn.Sequential(
            nn.Linear(dim, dim * expansion),
            nn.SiLU(),
            nn.Linear(dim * expansion, dim * expansion),
        )

    def forward(self, t):
        time = self.encode_time(t * self.time_scaling)
        time_mean = time.mean(dim=-1, keepdim=True)
        time_std = time.std(dim=-1, keepdim=True)
        time = (time - time_mean) / time_std
        return self.map_time(time)


def get_timestep_embedding(timesteps, embedding_dim, dtype=torch.float32):
    assert len(timesteps.shape) == 1
    timesteps = timesteps * 1000.0

    half_dim = embedding_dim // 2
    emb = np.log(10000) / (half_dim - 1)
    emb = (torch.arange(half_dim, dtype=dtype, device=timesteps.device) * -emb).exp()
    emb = timesteps.to(dtype)[:, None] * emb[None, :]
    emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
    if embedding_dim % 2 == 1:  # zero pad
        emb = F.pad(emb, (0, 1))
    assert emb.shape == (timesteps.shape[0], embedding_dim)
    return emb


class AdaLNMLPBlock(nn.Module):
    def __init__(self, dim, expansion):
        super().__init__()
        self.mlp = FusedMLP(
            dim, dropout=0.0, hidden_layer_multiplier=expansion, activation=nn.GELU
        )
        self.ada_map = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 3))
        self.ln = nn.LayerNorm(dim, elementwise_affine=False)

        nn.init.zeros_(self.mlp[-1].weight)
        nn.init.zeros_(self.mlp[-1].bias)

    def forward(self, x, y):
        gamma, mu, sigma = self.ada_map(y).chunk(3, dim=-1)
        x_res = (1 + gamma) * self.ln(x) + mu
        x = x + self.mlp(x_res) * sigma
        return x


class GeoAdaLNMLP(nn.Module):
    def __init__(self, input_dim, dim, depth, expansion, cond_dim):
        super().__init__()
        self.time_embedder = TimeEmbedder("positional", dim // 4, 1000, expansion=4)
        self.cond_mapper = nn.Linear(cond_dim, dim)
        self.initial_mapper = nn.Linear(input_dim, dim)
        self.blocks = nn.ModuleList(
            [AdaLNMLPBlock(dim, expansion) for _ in range(depth)]
        )
        self.final_adaln = nn.Sequential(
            nn.SiLU(),
            nn.Linear(dim, dim * 2),
        )
        self.final_ln = nn.LayerNorm(dim, elementwise_affine=False)
        self.final_linear = nn.Linear(dim, input_dim)

    def forward(self, batch):
        x = batch["y"]
        x = self.initial_mapper(x)
        gamma = batch["gamma"]
        cond = batch["emb"]
        t = self.time_embedder(gamma)
        cond = self.cond_mapper(cond)
        cond = cond + t
        for block in self.blocks:
            x = block(x, cond)
        gamma_last, mu_last = self.final_adaln(cond).chunk(2, dim=-1)
        x = (1 + gamma_last) * self.final_ln(x) + mu_last
        x = self.final_linear(x)
        return x


class GeoAdaLNMLPVonFisher(nn.Module):
    def __init__(self, input_dim, dim, depth, expansion, cond_dim):
        super().__init__()
        self.cond_mapper = nn.Linear(cond_dim, dim)
        self.blocks = nn.ModuleList(
            [AdaLNMLPBlock(dim, expansion) for _ in range(depth)]
        )
        self.final_adaln = nn.Sequential(
            nn.SiLU(),
            nn.Linear(dim, dim * 2),
        )
        self.final_ln = nn.LayerNorm(dim, elementwise_affine=False)
        self.mu_predictor = nn.Sequential(
            FusedMLP(dim, dropout=0.0, hidden_layer_multiplier=2, activation=nn.GELU),
            nn.Linear(dim, input_dim),
        )
        self.kappa_predictor = nn.Sequential(
            FusedMLP(dim, dropout=0.0, hidden_layer_multiplier=2, activation=nn.GELU),
            nn.Linear(dim, 1),
            torch.nn.Softplus(),
        )
        self.init_registers = torch.nn.Parameter(torch.randn(dim), requires_grad=True)
        torch.nn.init.trunc_normal_(
            self.init_registers, std=0.02, a=-2 * 0.02, b=2 * 0.02
        )

    def forward(self, batch):
        cond = batch["emb"]
        cond = self.cond_mapper(cond)
        x = self.init_registers.unsqueeze(0).repeat(cond.shape[0], 1)
        for block in self.blocks:
            x = block(x, cond)
        gamma_last, mu_last = self.final_adaln(cond).chunk(2, dim=-1)
        x = (1 + gamma_last) * self.final_ln(x) + mu_last
        mu = self.mu_predictor(x)
        mu = mu / mu.norm(dim=-1, keepdim=True)
        kappa = self.kappa_predictor(x)
        return mu, kappa


class GeoAdaLNMLPVonFisherMixture(nn.Module):
    def __init__(self, input_dim, dim, depth, expansion, cond_dim, num_mixtures=3):
        super().__init__()
        self.cond_mapper = nn.Linear(cond_dim, dim)
        self.blocks = nn.ModuleList(
            [AdaLNMLPBlock(dim, expansion) for _ in range(depth)]
        )
        self.final_adaln = nn.Sequential(
            nn.SiLU(),
            nn.Linear(dim, dim * 2),
        )
        self.final_ln = nn.LayerNorm(dim, elementwise_affine=False)
        self.mu_predictor = nn.Sequential(
            FusedMLP(dim, dropout=0.0, hidden_layer_multiplier=2, activation=nn.GELU),
            nn.Linear(dim, input_dim * num_mixtures),
        )
        self.kappa_predictor = nn.Sequential(
            FusedMLP(dim, dropout=0.0, hidden_layer_multiplier=2, activation=nn.GELU),
            nn.Linear(dim, num_mixtures),
            torch.nn.Softplus(),
        )
        self.mixture_weights = nn.Sequential(
            FusedMLP(dim, dropout=0.0, hidden_layer_multiplier=2, activation=nn.GELU),
            nn.Linear(dim, num_mixtures),
            torch.nn.Softmax(dim=-1),
        )
        self.num_mixtures = num_mixtures
        self.init_registers = torch.nn.Parameter(torch.randn(dim), requires_grad=True)
        torch.nn.init.trunc_normal_(
            self.init_registers, std=0.02, a=-2 * 0.02, b=2 * 0.02
        )

    def forward(self, batch):
        cond = batch["emb"]
        cond = self.cond_mapper(cond)
        x = self.init_registers.unsqueeze(0).repeat(cond.shape[0], 1)
        for block in self.blocks:
            x = block(x, cond)
        gamma_last, mu_last = self.final_adaln(cond).chunk(2, dim=-1)
        x = (1 + gamma_last) * self.final_ln(x) + mu_last
        mu = self.mu_predictor(x)
        mu = rearrange(mu, "b (n d) -> b n d", n=self.num_mixtures)
        mu = mu / mu.norm(dim=-1, keepdim=True)
        kappa = self.kappa_predictor(x)
        weights = self.mixture_weights(x)
        return mu, kappa, weights