## TODO: MLP model，用来降低维度

import torch.nn as nn
import torch
import warnings
import math
import os 

MLP_NAME = "mlp_model.pt"
MLP_OPTIMIZER_NAME = "mlp_optimizer.pt"
MLP_SCHEDULER_NAME = "mlp_scheduler.pt"

def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor

def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)

class LLM2EmbedProjector(nn.Module):
    def __init__(self, in_dim, out_dim, fp16=False):
        super().__init__()
        self.fp16 = fp16
        self.mlp1 = nn.Sequential(
                nn.Linear(in_dim, out_dim),
            )
            
    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        with torch.amp.autocast(device_type='cuda', enabled=self.fp16):
            x = self.mlp1(x)
            x = nn.functional.normalize(x, dim=-1, p=2)
        return x

    def maybe_load_state_dict_from_path(self, ckpt_dir):
        state_dict = None
        ckpt_file = os.path.join(ckpt_dir, MLP_NAME)
        if os.path.exists(ckpt_file):
            state_dict = torch.load(ckpt_file, map_location='cpu')
        if state_dict:
            self.load_state_dict(state_dict)

class LLM2EmbedAddProjector(nn.Module):
    def __init__(self, in_dim, out_dim, fp16=False):
        super().__init__()
        self.fp16 = fp16
        assert in_dim % out_dim == 0, "in_dim must be divisible by out_dim"
        self.out_dim = out_dim
        self.in_dim = in_dim
        self.n = in_dim // out_dim

    def forward(self, x):
        # x shape: (batch_size, in_dim)
        batch_size = x.size(0)
        x_reshaped = x.view(batch_size, self.out_dim, self.n)  # (batch_size, out_dim, n)
        with torch.amp.autocast(device_type='cuda', enabled=self.fp16):
            output = x_reshaped.sum(dim=2)  # sum over the grouped dimensions
        return output  # shape: (batch_size, out_dim)

    def maybe_load_state_dict_from_path(self, ckpt_dir):
        print("No weights for LLM2EmbedAddProjector")