import torch
import torch.nn as nn

class DualStreamTransformer(nn.Module):
    def __init__(self, n_feat1=24, n_feat2=10, d_model=32, num_classes=2):
        super(DualStreamTransformer, self).__init__()
        
        dim_ff = 128 
        
        self.feat_tokenizers_1 = nn.ModuleList([nn.Linear(1, d_model) for _ in range(n_feat1)])
        self.cls_token_1 = nn.Parameter(torch.zeros(1, 1, d_model))
        encoder_layer_1 = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=dim_ff, batch_first=True)
        self.encoder_1 = nn.TransformerEncoder(encoder_layer_1, num_layers=2)
        
        self.feat_tokenizers_2 = nn.ModuleList([nn.Linear(1, d_model) for _ in range(n_feat2)])
        self.cls_token_2 = nn.Parameter(torch.zeros(1, 1, d_model))
        encoder_layer_2 = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=dim_ff, batch_first=True)
        self.encoder_2 = nn.TransformerEncoder(encoder_layer_2, num_layers=2)
        
    
        self.fusion = nn.Sequential(
            nn.Linear(d_model * 2, d_model), 
            nn.ReLU()                        
        )

    def forward(self, x1, x2):
        tokens1 = [layer(x1[:, i].unsqueeze(1)) for i, layer in enumerate(self.feat_tokenizers_1)]
        x1_emb = torch.stack(tokens1, dim=1)
        x1_emb = torch.cat((self.cls_token_1.expand(x1.size(0), -1, -1), x1_emb), dim=1)
        feat1 = self.encoder_1(x1_emb)[:, 0, :]
        
        
        tokens2 = [layer(x2[:, i].unsqueeze(1)) for i, layer in enumerate(self.feat_tokenizers_2)]
        x2_emb = torch.stack(tokens2, dim=1)
        x2_emb = torch.cat((self.cls_token_2.expand(x2.size(0), -1, -1), x2_emb), dim=1)
        feat2 = self.encoder_2(x2_emb)[:, 0, :]
        
        combined = torch.cat((feat1, feat2), dim=1)
        return self.fusion(combined) 

class ArcMarginProduct(nn.Module):
    def __init__(self, in_features, out_features, s=30.0, m=0.5):
        super(ArcMarginProduct, self).__init__()
        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)
        
    def predict(self, x):
        cosine = torch.matmul(nn.functional.normalize(x), nn.functional.normalize(self.weight).t())
        return cosine