import torch
import torch.nn as nn
import torch.nn.functional as F

class GatedUpdateCell(nn.Module):
    def __init__(self, input_dim, hidden_dim,TMRB_dropout):
        super(GatedUpdateCell, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.input_to_hidden = nn.Linear(input_dim, hidden_dim)
        
        self.W_r = nn.Linear(hidden_dim*2, hidden_dim)
        self.W_z = nn.Linear(hidden_dim*2, hidden_dim)
        self.W_t = nn.Linear(hidden_dim*2, hidden_dim)
        self.dropout = nn.Dropout(p=TMRB_dropout)

    def forward(self, x, h_prev):
        if x.size(-1) != self.hidden_dim:
            x = self.input_to_hidden(x)
        combined = self.dropout(torch.cat((x, h_prev), dim=-1))
        r_t = torch.sigmoid(self.W_r(combined))
        z_t = torch.sigmoid(self.W_z(combined))
        h_t = torch.tanh(self.W_t(torch.cat((x, h_prev * r_t), dim=-1)))
        h_next = z_t * h_t + (1 - z_t) * x
        return h_next

class TMRB(nn.Module):
    def __init__(self, input_dim, out_dim,top_k, TMRB_dropout,is_update,select_k,device="cuda:0"):
        super(TMRB,self).__init__()
        self.input_dim = input_dim
        self.out_dim = out_dim
        self.device = device
        self.TMRB_dropout = TMRB_dropout
        self.is_update = is_update 
        self.GatedUpdateCell = GatedUpdateCell(out_dim, out_dim,self.TMRB_dropout).to(self.device)
        self.top_k = top_k
        self.select_k = select_k
        self.mlp = nn.Linear(self.top_k * self.input_dim, self.out_dim)
        self.init_hidden = nn.Parameter(
    nn.init.xavier_uniform_(torch.empty(input_dim, 1)))
        self.dorpout = nn.Dropout(TMRB_dropout)
        
    def forward(self, tem_emb, year, hidden_states_per_year):
        B, N, D = tem_emb.shape
        tem_emb = tem_emb.transpose(1,-1)
        
        if year - 1 in hidden_states_per_year.keys():
            prev_hidden = hidden_states_per_year[year - 1]
            prev_hidden = prev_hidden.expand(size = (N,*prev_hidden.shape)).transpose(0,1)
        else:
            prev_hidden = self.init_hidden
            prev_hidden = prev_hidden.expand(size = (N,*prev_hidden.shape)).transpose(0,1)

        prev_hidden = prev_hidden.expand(size = (B,*prev_hidden.shape))
        prev_hidden = prev_hidden.squeeze(-1)
        
        if self.select_k:
            time_step_diff = torch.abs(tem_emb - prev_hidden)
            _, top_k_indices = torch.topk(time_step_diff, k=self.top_k, dim=-1, largest=True, sorted=False)
            top_k_features = torch.gather(tem_emb, dim=2, index=top_k_indices)
            

        else:
            top_k_indices = torch.randint(0, N, (B, self.top_k), device=tem_emb.device)
            top_k_features = torch.gather(tem_emb, dim=2, index=top_k_indices.unsqueeze(-1).expand(-1, -1, D))

        top_k_features = top_k_features.view(B, -1)
        time_step_input = self.mlp(top_k_features)
        
        if self.is_update:
            prev_hidden_avg = torch.mean(prev_hidden, dim=2)
            prev_hidden_avg = prev_hidden_avg.view(B, -1)
            updated_hidden_pool = self.GatedUpdateCell(time_step_input,prev_hidden_avg)
        else:
            updated_hidden_pool = time_step_input

        updated_hidden_pool = updated_hidden_pool.unsqueeze(2).expand(B, self.out_dim, N)
        return updated_hidden_pool