from typing import Callable, Union

from .transformer import Reconstructor, Tokenizer, Transformer
import torch
import torch.nn as nn
import torch.optim
import math
from einops import rearrange

ModuleType = Union[str, Callable[..., nn.Module]]

# class PositionalEmbedding(torch.nn.Module):
#     def __init__(self, num_channels, max_positions=10000, endpoint=False):
#         super().__init__()
#         self.num_channels = num_channels
#         self.max_positions = max_positions
#         self.endpoint = endpoint

#     def forward(self, x):
#         freqs = torch.arange(start=0, end=self.num_channels//2, dtype=torch.float32, device=x.device)
#         freqs = freqs / (self.num_channels // 2 - (1 if self.endpoint else 0))
#         freqs = (1 / self.max_positions) ** freqs
#         x = x.ger(freqs.to(x.dtype))
#         x = torch.cat([x.cos(), x.sin()], dim=1)
#         return x



class TimeStepEmbedding(nn.Module):
    """
    Layer that embeds diffusion timesteps.

     Args:
        - dim (int): the dimension of the output.
        - max_period (int): controls the minimum frequency of the embeddings.
        - n_layers (int): number of dense layers
        - fourer (bool): whether to use random fourier features as embeddings
    """

    def __init__(
        self,
        dim: int,
        max_period: int = 10000,
        n_layers: int = 2,
        fourier: bool = False,
        scale=16,
    ):
        super().__init__()
        self.dim = dim
        self.max_period = max_period
        self.n_layers = n_layers
        self.fourier = fourier

        if dim % 2 != 0:
            raise ValueError(f"embedding dim must be even, got {dim}")

        if fourier:
            self.register_buffer("freqs", torch.randn(dim // 2) * scale)

        layers = []
        for i in range(n_layers - 1):
            layers.append(nn.Linear(dim, dim))
            layers.append(nn.SiLU())
        self.fc = nn.Sequential(*layers, nn.Linear(dim, dim))

    def forward(self, timesteps):
        if not self.fourier:
            d, T = self.dim, self.max_period
            mid = d // 2
            fs = torch.exp(-math.log(T) / mid * torch.arange(mid, dtype=torch.float32))
            fs = fs.to(timesteps.device)
            args = timesteps[:, None].float() * fs[None]
            emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        else:
            x = timesteps.ger((2 * torch.pi * self.freqs).to(timesteps.dtype))
            emb = torch.cat([x.cos(), x.sin()], dim=1)

        return self.fc(emb)
    
    

# class MLPDiffusion(nn.Module):
#     def __init__(self, d_in, dim_t = 512, h_dim = 2*512, use_mlp=True):
#         super().__init__()
#         self.dim_t = dim_t

#         self.proj = nn.Linear(d_in, dim_t)

#         self.mlp = nn.Sequential(
#             nn.Linear(dim_t, h_dim),
#             nn.SiLU(),
#             nn.Linear(h_dim, h_dim),
#             nn.SiLU(),
#             nn.Linear(h_dim, dim_t),
#             nn.SiLU(),
#             nn.Linear(dim_t, d_in),
#         ) if use_mlp else nn.Linear(dim_t, d_in)

#         # self.map_noise = PositionalEmbedding(num_channels=dim_t)
#         # self.time_embed = nn.Sequential(
#         #     nn.Linear(dim_t, dim_t),
#         #     nn.SiLU(),
#         #     nn.Linear(dim_t, dim_t)
#         # )
#         self.time_emb = TimeStepEmbedding(dim_t)
        
#         self.use_mlp = use_mlp
    
#     def forward(self, x, timesteps):
#         t_emb = self.time_emb(timesteps)    
#         x = self.proj(x) + t_emb
#         return self.mlp(x)


class MLPDiffusion(nn.Module):
    def __init__(self, d_in, n_layers, n_units, emb_dim):
        super().__init__()

        self.proj = nn.Linear(d_in, emb_dim)

        in_dims = [emb_dim] + (n_layers - 1) * [n_units]
        out_dims = n_layers * [n_units]
        layers = nn.ModuleList()
        for i in range(len(in_dims)):
            layers.append(nn.Linear(in_dims[i], out_dims[i]))
            layers.append(nn.ReLU())
        # add final layer
        layers.append(nn.Linear(n_units, d_in)) 
        self.mlp = nn.Sequential(*layers)
        self.time_emb = TimeStepEmbedding(emb_dim)

    def forward(self, x, timesteps):
        t_emb = self.time_emb(timesteps)    
        x = self.proj(x) + t_emb
        return self.mlp(x)
    
    
    
class UniModMLP(nn.Module):
    """
        Input:
            x_num: [bs, d_numerical]
            x_cat: [bs, len(categories)]
        Output:
            x_num_pred: [bs, d_numerical], the predicted mean for numerical data
            x_cat_pred: [bs, sum(categories)], the predicted UNORMALIZED logits for categorical data
    """
    def __init__(
            self, d_numerical, categories, num_layers, d_token,
            n_head = 1, factor = 4, bias = True, dim_t=512, h_dim=2*512, use_mlp=True, 
            n_layers=5, n_units=512, **kwargs
        ):
        super().__init__()
        self.d_numerical = d_numerical
        self.categories = categories
        self.d_token = d_token

        self.tokenizer = Tokenizer(d_numerical, categories, d_token, bias = bias)
        # self.encoder = Transformer(num_layers, d_token, n_head, d_token, factor)
        d_in = d_token * (d_numerical + len(categories))
        self.mlp = MLPDiffusion(d_in, n_layers, n_units, dim_t)
        # self.decoder = Transformer(num_layers, d_token, n_head, d_token, factor)
        self.detokenizer = Reconstructor(d_numerical, categories, d_token)
        
        self.model = nn.ModuleList([self.mlp])

    def forward(self, x_num, x_cat, timesteps):
        
        # x_num has shape (batch, num features)
        # x_cat_soft has shape (batch, sum(categories))? as its is onehot encoded
        
        e = self.tokenizer(x_num, x_cat)
        mlp_input = e[:, 1:, :] # ignore the first CLS token. 
        
        y = self.mlp(mlp_input.flatten(1), timesteps)
        pred_e = rearrange(y, 'b (n d) -> b n d', d=self.d_token)
        
        x_num_pred, x_cat_pred = self.detokenizer(pred_e)
        x_cat_pred = torch.cat(x_cat_pred, dim=-1) if len(x_cat_pred)>0 else torch.zeros_like(x_cat).to(x_num_pred.dtype)

        return x_num_pred, x_cat_pred


class Precond(nn.Module):
    def __init__(self,
        denoise_fn,
        sigma_data = 0.5,              # Expected standard deviation of the training data.
        net_conditioning = "sigma",
    ):
        super().__init__()
        self.sigma_data = sigma_data
        self.net_conditioning = net_conditioning
        self.denoise_fn_F = denoise_fn

    def forward(self, x_num, x_cat, t, sigma):

        x_num = x_num.to(torch.float32)

        sigma = sigma.to(torch.float32)
        assert sigma.ndim == 2
        if sigma.dim() > 1: # if learnable column-wise noise schedule, sigma conditioning is set to the defaults schedule of rho=7
            sigma_cond = (0.002 ** (1/7) + t * (80 ** (1/7) - 0.002 ** (1/7))).pow(7)
        else:
            sigma_cond = sigma 
        dtype = torch.float32

        c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)
        c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2).sqrt()
        c_in = 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt()
        c_noise = sigma_cond.log() / 4

        x_in = c_in * x_num
        if self.net_conditioning == "sigma":
            F_x, x_cat_pred = self.denoise_fn_F(x_in, x_cat, c_noise.flatten())
        elif self.net_conditioning == "t":
            F_x, x_cat_pred = self.denoise_fn_F(x_in, x_cat, t)

        assert F_x.dtype == dtype
        D_x = c_skip * x_num + c_out * F_x.to(torch.float32)
        
        return D_x, x_cat_pred
    

class Model(nn.Module):
    def __init__(
            self, denoise_fn,
            sigma_data=0.5, 
            precond=False, 
            net_conditioning="sigma",
            **kwargs
        ):
        super().__init__()
        self.precond = precond
        if precond:
            self.denoise_fn_D = Precond(
                denoise_fn,
                sigma_data=sigma_data,
                net_conditioning=net_conditioning
            )
        else:
            self.denoise_fn_D = denoise_fn

    def forward(self, x_num, x_cat, t, sigma=None):
        if self.precond:
            return self.denoise_fn_D(x_num, x_cat, t, sigma)
        else:
            return self.denoise_fn_D(x_num, x_cat, t)