class AdditiveAttention(nn.Module):
    def __init__(self, query_dim, key_dim, hidden_dim):
        super(AdditiveAttention, self).__init__()
        self.query_proj = nn.Linear(query_dim, hidden_dim)
        self.key_proj = nn.Linear(key_dim, hidden_dim)
        self.v = nn.Linear(hidden_dim, 1, bias=False)

    def forward(self, query, key, mask=None):
        # query: (batch_size, query_len, query_dim)
        # key: (batch_size, key_len, key_dim)
        # mask: (batch_size, query_len, key_len)

        proj_query = self.query_proj(query)  # (batch_size, query_len, hidden_dim)
        proj_key = self.key_proj(key)        # (batch_size, key_len, hidden_dim)

        # Expand dimensions to perform addition
        proj_query = proj_query.unsqueeze(2)  # (batch_size, query_len, 1, hidden_dim)
        proj_key = proj_key.unsqueeze(1)      # (batch_size, 1, key_len, hidden_dim)

        # Add and apply non-linearity
        energy = torch.tanh(proj_query + proj_key)  # (batch_size, query_len, key_len, hidden_dim)

        # Compute attention scores
        scores = self.v(energy).squeeze(-1)  # (batch_size, query_len, key_len)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attention_weights = F.softmax(scores, dim=-1)  # (batch_size, query_len, key_len)

        return attention_weights

class TrainableMemory(nn.Module):
    def __init__(self, memory_size, memory_dim):
        super(TrainableMemory, self).__init__()
        self.memory = nn.Parameter(torch.randn(memory_size, memory_dim))

    def forward(self, batch_size):
        # Expand memory for the batch
        return self.memory.unsqueeze(0).expand(batch_size, -1, -1)  # (batch_size, memory_size, memory_dim)

class ImprovedAttention(nn.Module):
    def __init__(self, input_dim, num_heads, memory_size, memory_dim, hidden_dim):
        super(ImprovedAttention, self).__init__()
        assert input_dim % num_heads == 0, "input_dim must be divisible by num_heads"
        self.input_dim = input_dim
        self.num_heads = num_heads
        self.head_dim = input_dim // num_heads

        # Additive Attention
        self.additive_attention = AdditiveAttention(
            query_dim=self.head_dim,
            key_dim=self.head_dim,
            hidden_dim=hidden_dim
        )

        # Linear projections for multi-head
        self.query_linear = nn.Linear(input_dim, input_dim)
        self.key_linear = nn.Linear(input_dim, input_dim)
        self.value_linear = nn.Linear(input_dim, input_dim)

        # Parallel Feed-Forward Network for guidance
        self.parallel_ffn = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim)
        )

        # Trainable Memory
        self.memory_module = TrainableMemory(memory_size, memory_dim)
        self.memory_proj = nn.Linear(memory_dim, input_dim)

        # Final linear layer
        self.output_linear = nn.Linear(input_dim, input_dim)

        # Gating mechanism with learnable sigmoid gate
        self.gate_transform = nn.Linear(input_dim, input_dim)
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_seq, mask=None):
        """
        input_seq: (batch_size, seq_len, input_dim)
        mask: (batch_size, seq_len, seq_len)
        """

        batch_size, seq_len, _ = input_seq.size()

        # Linear projections
        queries = self.query_linear(input_seq)  # (batch_size, seq_len, input_dim)
        keys = self.key_linear(input_seq)       # (batch_size, seq_len, input_dim)
        values = self.value_linear(input_seq)   # (batch_size, seq_len, input_dim)

        # Reshape for multi-head attention
        queries = queries.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        keys = keys.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        values = values.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        # Now queries, keys, values are (batch_size, num_heads, seq_len, head_dim)

        # Additive Attention for each head
        attention_weights = self.additive_attention(
            query=queries.contiguous().view(batch_size * self.num_heads, seq_len, self.head_dim),
            key=keys.contiguous().view(batch_size * self.num_heads, seq_len, self.head_dim),
            mask=mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1).contiguous().view(batch_size * self.num_heads, seq_len, seq_len) if mask is not None else None
        )  # (batch_size * num_heads, seq_len, seq_len)

        attention_weights = attention_weights.view(batch_size, self.num_heads, seq_len, seq_len)
        # (batch_size, num_heads, seq_len, seq_len)

        # Weighted sum of values
        attention_output = torch.matmul(attention_weights, values)  # (batch_size, num_heads, seq_len, head_dim)

        # Concatenate heads
        attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.input_dim)
        # (batch_size, seq_len, input_dim)

        # Incorporate Trainable Memory
        memory = self.memory_module(batch_size)  # (batch_size, memory_size, memory_dim)
        memory_proj = self.memory_proj(memory)    # (batch_size, memory_size, input_dim)

        # Combine attention output with memory
        combined = torch.cat([attention_output, memory_proj], dim=1)  # (batch_size, seq_len + memory_size, input_dim)

        # Apply Parallel Feed-Forward Network
        guided = self.parallel_ffn(combined)  # (batch_size, seq_len + memory_size, input_dim)

        # Gating mechanism
        gate = self.sigmoid(self.gate_transform(guided))  # (batch_size, seq_len + memory_size, input_dim)
        output = gate * guided  # Element-wise multiplication

        # Final linear layer
        output = self.output_linear(output)  # (batch_size, seq_len + memory_size, input_dim)

        return output, attention_weights