import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForCausalLM
from RBFLayer import RBFLayer

def l_norm(x, p=2):
    return torch.norm(x, p=p, dim=-1)


# Gaussian RBF
def rbf_gaussian(x):
    return (-x.pow(2)).exp()

class CustomRBFFeedForward(nn.Module):
    def __init__(self, in_features, out_features, num_kernels):
        super(CustomRBFFeedForward, self).__init__()
        # RBFLayer from the given implementation
        self.rbf_layer = RBFLayer(
            in_features_dim=in_features,  # Input size (e.g., 896)
            num_kernels=num_kernels,  # Number of kernels in the RBF layer (can be tuned)
            out_features_dim=out_features,  # Output size (e.g., 4864)
            radial_function=rbf_gaussian,  # Use the Gaussian RBF
            norm_function=l_norm  # Use Euclidean norm
        )

    def forward(self, x):
        # Apply the RBF layer to the input x
        return self.rbf_layer(x)