import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import nn
from einops import rearrange
import math
from einops_exts import check_shape, rearrange_many
from torch import Size, Tensor, nn

class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = x[:, None] * emb[None, :]
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb


def map_positional_encoding(v: Tensor, freq_bands: Tensor) -> Tensor:
    """Map v to positional encoding representation phi(v)

    Arguments:
        v (Tensor): input features (B, IFeatures)
        freq_bands (Tensor): frequency bands (N_freqs, )

    Returns:
        phi(v) (Tensor): fourrier features (B, 3 + (2 * N_freqs) * 3)
    """
    pe = [v]
    for freq in freq_bands:
        fv = freq * v
        pe += [torch.sin(fv), torch.cos(fv)]
    return torch.cat(pe, dim=-1)

class FeatureMapping(nn.Module):
    """FeatureMapping nn.Module

    Maps v to features following transformation phi(v)

    Arguments:
        i_dim (int): input dimensions
        o_dim (int): output dimensions
    """

    def __init__(self, i_dim: int, o_dim: int) -> None:
        super().__init__()
        self.i_dim = i_dim
        self.o_dim = o_dim

    def forward(self, v: Tensor) -> Tensor:
        """FeratureMapping forward pass

        Arguments:
            v (Tensor): input features (B, IFeatures)

        Returns:
            phi(v) (Tensor): mapped features (B, OFeatures)
        """
        raise NotImplementedError("Forward pass not implemented yet!")

class PositionalEncoding(FeatureMapping):
    """PositionalEncoding module

    Maps v to positional encoding representation phi(v)

    Arguments:
        i_dim (int): input dimension for v
        N_freqs (int): #frequency to sample (default: 10)
    """

    def __init__(
        self,
        i_dim: int,
        N_freqs: int = 10,
    ) -> None:
        super().__init__(i_dim, 3 + (2 * N_freqs) * 3)
        self.N_freqs = N_freqs

        a, b = 1, self.N_freqs - 1
        freq_bands = 2 ** torch.linspace(a, b, self.N_freqs)
        self.register_buffer("freq_bands", freq_bands)

    def forward(self, v: Tensor) -> Tensor:
        """Map v to positional encoding representation phi(v)

        Arguments:
            v (Tensor): input features (B, IFeatures)

        Returns:
            phi(v) (Tensor): fourrier features (B, 3 + (2 * N_freqs) * 3)
        """
        return map_positional_encoding(v, self.freq_bands)

class BaseTemperalPointModel(nn.Module):
    """ A base class providing useful methods for point cloud processing. """

    def __init__(
        self,
        *,
        num_classes,
        embed_dim,
        extra_feature_channels,
        dim: int = 768,
        num_layers: int = 6
    ):
        super().__init__()

        self.extra_feature_channels = extra_feature_channels
        self.timestep_embed_dim = 256
        self.output_dim = num_classes
        self.dim = dim
        self.num_layers = num_layers


        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(dim),
            nn.Linear(dim, self.timestep_embed_dim ),
            nn.SiLU(),
            nn.Linear(self.timestep_embed_dim , self.timestep_embed_dim )
        )

        self.positional_encoding = PositionalEncoding(i_dim=3, N_freqs=10)
        positional_encoding_d_out = 3 + (2 * 10) * 3

        # Input projection (point coords, point coord encodings, other features, and timestep embeddings)

        self.input_projection = nn.Linear(
            in_features=(3 + positional_encoding_d_out),
            out_features=self.dim
        )#b f p c

        # Transformer layers
        self.layers = self.get_layers()

        # Output projection
        self.output_projection = nn.Linear(self.dim, self.output_dim)
    def get_layers(self):
        raise NotImplementedError('This method should be implemented by subclasses')

    def forward(self, inputs: torch.Tensor, t: torch.Tensor):
        raise NotImplementedError('This method should be implemented by subclasses')