import torch
import torch.nn as nn
import numpy as np
from math import sqrt

class FullAttention(nn.Module):
    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
        super(FullAttention, self).__init__()
        self.scale = scale
        self.mask_flag = mask_flag
        self.output_attention = output_attention
        self.dropout = nn.Dropout(attention_dropout)

    def forward(self, queries, keys, values, instruct_mask, ts_mask):
        B, L, H, E = queries.shape
        _, S, _, D = values.shape
        scale = self.scale or 1. / sqrt(E)

        scores = torch.einsum("blhe,bshe->bhls", queries, keys)

        if self.mask_flag:
            ts_mask  = ts_mask[..., 0]
            
            instruct_mask_expanded = instruct_mask.unsqueeze(1).unsqueeze(3)  
            ts_mask_expanded = ts_mask.unsqueeze(1).unsqueeze(2)          
            mask = ~torch.logical_and(instruct_mask_expanded, ts_mask_expanded)
            scores.masked_fill_(mask, -np.inf)
        
        A = self.dropout(torch.softmax(scale * scores, dim=-1))
        A = torch.nan_to_num(A, nan=0.0, posinf=1.0, neginf=0.0) 
        V = torch.einsum("bhls,bshd->blhd", A, values)

        if self.output_attention:
            return V.contiguous(), A
        else:
            return V.contiguous(), None
        
class AttentionLayer(nn.Module):
    def __init__(self, attention, d_model, n_heads, d_keys=None,
                 d_values=None):
        super(AttentionLayer, self).__init__()

        d_keys = d_keys or (d_model // n_heads)
        d_values = d_values or (d_model // n_heads)

        self.inner_attention = attention
        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
        self.value_projection = nn.Linear(d_model, d_values * n_heads)
        self.out_projection = nn.Linear(d_values * n_heads, d_model)
        self.n_heads = n_heads

    def forward(self, queries, keys, values, instruct_mask, ts_mask):
        B, L, _ = queries.shape
        _, S, _ = keys.shape
        H = self.n_heads

        queries = self.query_projection(queries).view(B, L, H, -1)
        keys = self.key_projection(keys).view(B, S, H, -1)
        values = self.value_projection(values).view(B, S, H, -1)

        out, attn = self.inner_attention(
            queries,
            keys,
            values,
            instruct_mask, 
            ts_mask
        )
        out = out.view(B, L, -1)

        return self.out_projection(out), attn
