import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

MAX_SLATE_SIZE = 160
EMBED_DIM = 256

class TriModalTwoTower(nn.Module):
    def __init__(self, num_items, plot_embs, gnn_embs, scalars):
        super().__init__()
        self.register_buffer("plot_feats", torch.tensor(np.vstack([np.zeros((1, 1024)), plot_embs]), dtype=torch.float32))
        self.register_buffer("gnn_feats", torch.tensor(np.vstack([np.zeros((1, 256)), gnn_embs]), dtype=torch.float32))
        self.register_buffer("scalar_feats", torch.tensor(np.vstack([np.zeros((1, 10)), scalars]), dtype=torch.float32))
        
        self.plot_proj = nn.Linear(1024, 256)
        self.gnn_proj = nn.Linear(256, 128)
        self.scalar_proj = nn.Linear(10, 64)
        
        self.content_fusion = nn.Sequential(
            nn.Linear(256 + 128 + 64, 512), nn.LayerNorm(512), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(512, EMBED_DIM), nn.LayerNorm(EMBED_DIM) 
        )
        self.cf_emb = nn.Embedding(num_items + 1, EMBED_DIM, padding_idx=0)
        self.cf_norm = nn.LayerNorm(EMBED_DIM) 
        self.id_dropout = nn.Dropout(0.3)
        self.rating_emb = nn.Embedding(6, EMBED_DIM, padding_idx=0) 
        
        self.long_attention = nn.MultiheadAttention(embed_dim=EMBED_DIM, num_heads=4, batch_first=True)
        self.long_proj = nn.Sequential(nn.Linear(EMBED_DIM, EMBED_DIM), nn.LayerNorm(EMBED_DIM))
        self.short_gru = nn.GRU(EMBED_DIM, EMBED_DIM, batch_first=True)
        self.short_proj = nn.Sequential(nn.Linear(EMBED_DIM, EMBED_DIM), nn.LayerNorm(EMBED_DIM))
        self.gate = nn.Sequential(nn.Linear(EMBED_DIM * 2, EMBED_DIM), nn.GELU(), nn.Linear(EMBED_DIM, 1), nn.Sigmoid())

    def get_item_representations(self, item_ids, is_target=False):
        p = F.relu(self.plot_proj(self.plot_feats[item_ids]))
        g = F.relu(self.gnn_proj(self.gnn_feats[item_ids]))
        s = F.relu(self.scalar_proj(self.scalar_feats[item_ids]))
        content_vec = self.content_fusion(torch.cat([p, g, s], dim=-1))
        cf_vec = self.cf_norm(self.cf_emb(item_ids)) 
        if not is_target: cf_vec = self.id_dropout(cf_vec) 
        return content_vec + cf_vec

    def forward_user(self, long_items, long_ratings, short_items, short_ratings):
        long_embs = self.get_item_representations(long_items, is_target=False) + self.rating_emb(long_ratings)
        l_pad_mask = (long_items == 0)
        attn_out, _ = self.long_attention(long_embs, long_embs, long_embs, key_padding_mask=l_pad_mask)
        mask_f = (~l_pad_mask).unsqueeze(-1).float()
        long_vec = self.long_proj((attn_out * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp(min=1e-6))
        
        short_embs = self.get_item_representations(short_items, is_target=False) + self.rating_emb(short_ratings)
        lengths = (~(short_items == 0)).sum(dim=1).clamp(min=1).cpu()
        packed = nn.utils.rnn.pack_padded_sequence(short_embs, lengths, batch_first=True, enforce_sorted=False)
        _, hidden = self.short_gru(packed)
        short_vec = self.short_proj(hidden.squeeze(0))
        
        alpha = self.gate(torch.cat([long_vec, short_vec], dim=-1))
        return F.normalize(alpha * short_vec + (1.0 - alpha) * long_vec, p=2, dim=1)

    def forward_item(self, item_ids):
        return F.normalize(self.get_item_representations(item_ids, is_target=True), p=2, dim=1)

class SemanticSlateRanker(nn.Module):
    def __init__(self, continuous_dim, num_items, num_franchises, pretrained_text_embs, d_model=128, dropout=0.1):
        super().__init__()
        self.item_emb = nn.Embedding(num_items, pretrained_text_embs.shape[1], padding_idx=0)
        self.item_emb.weight.data.copy_(torch.from_numpy(pretrained_text_embs).float())
        self.item_emb.weight.requires_grad = False 
        
        self.franchise_emb = nn.Embedding(num_franchises, 16, padding_idx=0)
        self.pos_encoder = nn.Embedding(MAX_SLATE_SIZE, d_model)
        
        self.item_proj = nn.Sequential(nn.Linear(pretrained_text_embs.shape[1], d_model), nn.LayerNorm(d_model), nn.GELU())
        # continuous_dim (11) + franchise_emb (16) = 27 
        self.tab_proj = nn.Sequential(
            nn.Linear(continuous_dim + 16, d_model), nn.LayerNorm(d_model), nn.GELU(),
            nn.Linear(d_model, d_model), nn.LayerNorm(d_model), nn.GELU()
        )
        self.fusion_proj = nn.Sequential(nn.Linear(d_model * 2, d_model), nn.LayerNorm(d_model), nn.GELU())
        
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=d_model*2, dropout=dropout, batch_first=True, norm_first=True)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2, enable_nested_tensor=False)
        self.scorer = nn.Sequential(nn.Linear(d_model, 64), nn.GELU(), nn.Linear(64, 1))

    def forward(self, tab, f_ids, item_ids, valid_mask):
        f_embs = self.franchise_emb(f_ids)
        tab_encoded = self.tab_proj(torch.cat([tab, f_embs], dim=-1))          
        c_emb = self.item_proj(self.item_emb(item_ids))              
        fused = self.fusion_proj(torch.cat([tab_encoded, c_emb], dim=-1))   
        
        seq_length = tab.size(1)
        positions = torch.arange(seq_length, device=tab.device).unsqueeze(0).expand(tab.size(0), -1)
        contextualized = self.transformer(fused + self.pos_encoder(positions), src_key_padding_mask=~valid_mask)
        return self.scorer(contextualized).squeeze(-1)