import torch
from torch import nn

from .base_path import BasePath
from .linear import LinearPath

class MLPpath(BasePath):
    """
    Multilayer Perceptron (MLP) path class for generating geometric paths.

    Args:
        n_embed (int, optional): Number of embedding dimensions. Defaults to 1.
        depth (int, optional): Depth of the MLP. Defaults to 2.
        activation (str, optional): Activation function to use. Defaults to "gelu".
        base: Path class to correct. Defaults to LinearPath.
    """
    def __init__(
        self,
        n_embed: int = 1,
        depth: int = 2,
        activation: str = "gelu",
        base: BasePath = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        activation_dict = {key.lower(): key for key in dir(nn) if not key.startswith('_')}
        name = activation_dict[activation.lower()]
        activation_class = getattr(nn, name)
        self.activation = activation_class()
        # input_sizes = [1] + [n_embed]*(depth - 1)
        input_sizes = [1] + [self.final_position.shape[-1] * n_embed]*(depth - 1)
        output_sizes = input_sizes[1:] + [self.final_position.shape[-1]]
        self.layers = [
            nn.Linear(
                input_sizes[i//2], output_sizes[i//2], dtype=self.dtype, bias=True
            ) if i%2 == 0 else self.activation\
            for i in range(depth*2 - 1)
        ]
       
        self.mlp = nn.Sequential(*self.layers)
        self.mlp.to(self.device)
        self.neval = 0

        self.base = base if base is not None else LinearPath(**kwargs)
        
        print("Number of trainable parameters in MLP:", sum(p.numel() for p in self.parameters() if p.requires_grad))
        print(self.mlp)

    def get_positions(self, time: float):
        mlp_out = self.mlp(time) * (1 - time) * time #* 4
        if self.fix_positions is not None:
            mlp_out[:, self.fix_positions] = 0.0
        base_out = self.base.get_positions(time) #* (1 - (1 - time) * time * 4)
        out = base_out + mlp_out
        return out

    def __get_positions(self, time: float):
        """
        Generates a geometric path using the MLP.

        Args:
            time (float): Time parameter for generating the path.
            *args: Additional arguments.

        Returns:
            torch.Tensor: The geometric path generated by the MLP.
        """
        # mlp_out = self.mlp(time) - (1 - time) * self.mlp(self.t_init) - time * self.mlp(self.t_final)
        # out = mlp_out
        mlp_out = self.mlp(time) * (1 - time) * time #* 4
        if self.fix_positions is not None:
            mlp_out[:, self.fix_positions.repeat_interleave(3)] = 0.0
        base_out = self.base.get_positions(time) #* (1 - (1 - time) * time * 4)
        out = base_out + mlp_out
        return out

# class ResNetLayer(nn.Module):
#     def __init__(
#         self,
#         output_size: int,
#         n_embed: int,
#         activation: str = "selu",
#     ):
#         super().__init__()
#         self.activation = activation_dict[activation]
#         self.layer = nn.Sequential(
#             nn.Linear(output_size, output_size * n_embed, dtype=torch.float64, bias=True),
#             self.activation,
#             nn.Linear(output_size * n_embed, output_size, dtype=torch.float64, bias=True),
#         )
#         # self.layer = SwiGLU(output_size, output_size)

#     def forward(self, x):
#         return x + self.layer(x)
    
# class SwiGLU(nn.Module):
#     def __init__(self, in_features, out_features):
#         super().__init__()
#         self.fc1 = nn.Linear(in_features, out_features)
#         self.fc2 = nn.Linear(in_features, out_features)
#         self.activation = nn.SiLU()

#     def forward(self, x):
#         return self.fc1(x) * self.activation(self.fc2(x))
