"""Base Transformer models for working with mass spectra and peptides"""
import re

import torch
import torch.nn as nn
import numpy as np
from .encoders import MassEncoder, PeakEncoder, PositionalEncoder
from ..masses import PeptideMass
from .. import utils
import torch.nn.functional as F
from transfusion.model  import timestep_embedding


class WindowedSelfAttention(nn.Module):
    def __init__(self, input_dim, window_size):
        super(WindowedSelfAttention, self).__init__()
        self.input_dim = input_dim
        self.window_size = window_size

    def forward(self, query, key, value):
        batch_size, query_seq_len, _ = query.size()
        _, key_seq_len, _ = key.size()

        # Initialize an empty tensor for the output
        output = torch.zeros(batch_size, query_seq_len, self.input_dim).to(query.device)

        # Calculate attention scores for each query position
        for i in range(query_seq_len):
            start = max(0, i - self.window_size // 2)
            end = min(key_seq_len, i + self.window_size // 2 + 1)
            attention_scores = torch.matmul(query[:, i, :].unsqueeze(1), key[:, start:end, :].transpose(1, 2))
            attention_scores = attention_scores / (key[:, start:end, :].size(-1) ** 0.5)
            attention_weights = F.softmax(attention_scores, dim=-1)
            weighted_sum = torch.matmul(attention_weights, value[:, start:end, :])
            output[:, i, :] = weighted_sum.squeeze(1)

        return output



class SelfAttentionConcatenation(nn.Module):
    def __init__(self, input_dim):
        super(SelfAttentionConcatenation, self).__init__()

        self.attention = nn.MultiheadAttention(input_dim, num_heads=1)

    def forward(self, *tensors):
        attn_tensors = []

        for tensor in tensors:
            attn_tensor, _ = self.attention(tensor, tensor, tensor)
            attn_tensors.append(attn_tensor)

        # Transpose and reshape the attention tensors to match the desired output shape
        output_dim = attn_tensors[0].shape[-1]
        attn_tensors = [attn_tensor.transpose(0, 1).reshape(-1, attn_tensor.shape[-2], output_dim) for attn_tensor in
                        attn_tensors]

        # Concatenate the tensors along the second dimension
        combined_tensor = torch.cat(attn_tensors, dim=1)

        return combined_tensor

class AttentionConcatenation(nn.Module):
    def __init__(self, input_dim):
        super(AttentionConcatenation, self).__init__()

        self.attention = nn.Sequential(
            nn.Linear(input_dim, 1),
            nn.Tanh(),
            nn.Softmax(dim=1)  # Apply softmax along the second dimension
        )

    def forward(self, tensor1, tensor2, tensor3):
        attn_weights1 = self.attention(tensor1)
        attn_weights2 = self.attention(tensor2)
        attn_weights3 = self.attention(tensor3)

        weighted_tensor1 = torch.mul(tensor1, attn_weights1)
        weighted_tensor2 = torch.mul(tensor2, attn_weights2)
        weighted_tensor3 = torch.mul(tensor3, attn_weights3)

        concatenated_tensor = torch.cat([weighted_tensor1, weighted_tensor2, weighted_tensor3], dim=1)

        return  concatenated_tensor #weighted_tensor1 #concatenated_tensor


class SpectrumEncoder(torch.nn.Module):
    """A Transformer encoder for input mass spectra.

    Parameters
    ----------
    dim_model : int, optional
        The latent dimensionality to represent peaks in the mass spectrum.
    n_head : int, optional
        The number of attention heads in each layer. ``dim_model`` must be
        divisible by ``n_head``.
    dim_feedforward : int, optional
        The dimensionality of the fully connected layers in the Transformer
        layers of the model.
    n_layers : int, optional
        The number of Transformer layers.
    dropout : float, optional
        The dropout probability for all layers.
    peak_encoder : bool, optional
        Use positional encodings m/z values of each peak.
    dim_intensity: int or None, optional
        The number of features to use for encoding peak intensity.
        The remaining (``dim_model - dim_intensity``) are reserved for
        encoding the m/z value.
    """

    def __init__(
            self,
            dim_model=128,
            vocab_size=26,
            win_size=10,
            num_ions=8,
            num_neighbor=5,
            n_head=8,
            dim_feedforward=1024,
            n_layers=1,
            dropout=0,
            peak_encoder=True,
            dim_intensity=None,
    ):
        """Initialize a SpectrumEncoder"""
        super().__init__()

        self.latent_spectrum = torch.nn.Parameter(torch.randn(1, 1, dim_model))
        window_size = 5  # the ms1 profile shape
        # self.attention_layer = AttentionConcatenation(dim_model) #SelfAttentionConcatenation(dim_model)  #WindowedSelfAttention(dim_model, window_size)
        # SelfAttentionConcatenation(dim_model) #AttentionConcatenation(dim_model)

        if peak_encoder:
            self.peak_encoder = PeakEncoder(
                dim_model,
                dim_intensity=dim_intensity,
            )
        else:
            self.peak_encoder = torch.nn.Linear(2, dim_model)

        self.ms1_encoder = torch.nn.Linear(1, dim_model)
        # multiple by two because of b ion and y ion
        self.ms2_encoder = torch.nn.Linear(vocab_size * win_size * num_ions * num_neighbor, dim_model)

        # The Transformer layers:
        layer = torch.nn.TransformerEncoderLayer(
            d_model=dim_model,
            nhead=n_head,
            dim_feedforward=dim_feedforward,
            batch_first=True,
            dropout=dropout,
        )

        self.transformer_encoder = torch.nn.TransformerEncoder(
            layer,
            num_layers=n_layers,
        )

    def mask1(self, X):
        zeros = ~X.sum(dim=2).bool()

        return zeros

    def calculate_mask(self, X):

        zeros = ~X.sum(dim=2).bool()
        mask = [
            torch.tensor([[False]] * X.shape[0]).type_as(zeros),
            zeros,
        ]

        mask = torch.cat(mask, dim=1)

        return mask

    def forward(self, spectra, ms1, ms2_b_ion, ms2_y_ion):
        """The forward pass.

        Parameters
        ----------
        spectra : torch.Tensor of shape (n_spectra, n_peaks, 2)
            The spectra to embed. Axis 0 represents a mass spectrum, axis 1
            contains the peaks in the mass spectrum, and axis 2 is essentially
            a 2-tuple specifying the m/z-intensity pair for each peak. These
            should be zero-padded, such that all of the spectra in the batch
            are the same length.
         ms1: torch.Tensor of shape (n_spectra, n_neghibors)
         ms2_b_ion: torch.Tensor of shape (n_spectra, seq_len, vocab_size, n_ion, n_neighbors, window_size)
         ms2_y_ion: torch.Tensor of shape (n_spectra, seq_len, vocab_size, n_ions, n_neighbors, window_size)
        Returns
        -------
        latent : torch.Tensor of shape (n_spectra, n_peaks + 1, dim_model)
            The latent representations for the spectrum and each of its
            peaks.
        mem_mask : torch.Tensor
            The memory mask specifying which elements were padding in X.
        """
        zeros = ~spectra.sum(dim=2).bool()
        mask = [
            torch.tensor([[False]] * spectra.shape[0]).type_as(zeros),
            zeros,
        ]
        spectra_mask = torch.cat(mask, dim=1)
        peaks = self.peak_encoder(spectra)

        # Add the spectrum representation to each input:
        latent_spectra = self.latent_spectrum.expand(peaks.shape[0], -1, -1)

        ms1 = ms1[:, :, None]
        ms1_mask = self.calculate_mask(ms1)
        ms1 = self.ms1_encoder(ms1.to(self.ms1_encoder.weight.dtype))

        ms2_b_ion = ms2_b_ion.view(ms2_b_ion.size(0), ms2_b_ion.size(1), -1)
        ms2_b_ion_mask = self.calculate_mask(ms2_b_ion)
        ms2_b_ion = self.ms2_encoder(ms2_b_ion.to(self.ms2_encoder.weight.dtype))

        ms2_y_ion = ms2_y_ion.view(ms2_y_ion.size(0), ms2_y_ion.size(1), -1)
        ms2_y_ion_mask = self.calculate_mask(ms2_y_ion)
        ms2_y_ion = self.ms2_encoder(ms2_y_ion.to(self.ms2_encoder.weight.dtype))

        ms2_b_ion = torch.cat([latent_spectra, ms2_b_ion], dim=1)
        ms2_y_ion = torch.cat([latent_spectra, ms2_y_ion], dim=1)
        ms2 = torch.cat([ms2_b_ion, ms2_y_ion], dim=1)

        # ms2 = torch.cat([ms2_b_ion, ms2_y_ion], dim =1)
        # ms2 = self.ms2_encoder(ms2.to(self.ms2_encoder.weight.dtype))

        ms2_mask = torch.cat([ms2_b_ion_mask, ms2_y_ion_mask], dim=1)

        ms1 = torch.cat([latent_spectra, ms1], dim=1)
        # ms2 = torch.cat([latent_spectra, ms2], dim=1)
        mask = torch.cat([spectra_mask, ms2_mask, ms1_mask], dim=1)

        peaks = torch.cat([latent_spectra, peaks], dim=1)
        peaks = torch.cat([peaks, ms2, ms1], dim=1)
        # peaks = self.windowed_attention(peaks, ms1, ms2)
        # peaks = self.attention_layer(peaks, ms2, ms1)

        return self.transformer_encoder(peaks, src_key_padding_mask=mask), mask

    @property
    def device(self):
        """The current device for the model"""
        return next(self.parameters()).device


class _PeptideTransformer(torch.nn.Module):
    """A transformer base class for peptide sequences.

    Parameters
    ----------
    dim_model : int
        The latent dimensionality to represent the amino acids in a peptide
        sequence.
    pos_encoder : bool
        Use positional encodings for the amino acid sequence.
    residues: Dict or str {"massivekb", "canonical"}, optional
        The amino acid dictionary and their masses. By default this is only
        the 20 canonical amino acids, with cysteine carbamidomethylated. If
        "massivekb", this dictionary will include the modifications found in
        MassIVE-KB. Additionally, a dictionary can be used to specify a custom
        collection of amino acids and masses.
    max_charge : int
        The maximum charge to embed.
    """

    def __init__(
        self,
        dim_model,
        pos_encoder,
        residues,
        max_charge,
    ):
        super().__init__()
        self.reverse = False
        self._peptide_mass = PeptideMass(residues=residues)
        self._amino_acids = list(self._peptide_mass.masses.keys()) + ["$"]
        self._idx2aa = {i + 1: aa for i, aa in enumerate(self._amino_acids)}
        self._aa2idx = {aa: i for i, aa in self._idx2aa.items()}

        if pos_encoder:
            self.pos_encoder = PositionalEncoder(dim_model)
        else:
            self.pos_encoder = torch.nn.Identity()

        self.charge_encoder = torch.nn.Embedding(max_charge, dim_model)
        self.aa_encoder = torch.nn.Embedding(
            len(self._amino_acids) + 1,
            dim_model,
            padding_idx=0,
        )

    def tokenize(self, sequence, partial=False):
        """Transform a peptide sequence into tokens

        Parameters
        ----------
        sequence : str
            A peptide sequence.

        Returns
        -------
        torch.Tensor
            The token for each amino acid in the peptide sequence.
        """
        if not isinstance(sequence, str):
            return sequence  # Assume it is already tokenized.

        sequence = sequence.replace("I", "L")
        sequence = re.split(r"(?<=.)(?=[A-Z])", sequence)
        if self.reverse:
            sequence = list(reversed(sequence))

        if not partial:
            sequence += ["$"]

        tokens = [self._aa2idx[aa] for aa in sequence]
        tokens = torch.tensor(tokens, device=self.device)
        return tokens

    def detokenize(self, tokens):
        """Transform tokens back into a peptide sequence.

        Parameters
        ----------
        tokens : torch.Tensor of shape (n_amino_acids,)
            The token for each amino acid in the peptide sequence.

        Returns
        -------
        list of str
            The amino acids in the peptide sequence.
        """
        sequence = [self._idx2aa.get(i.item(), "") for i in tokens]
        if "$" in sequence:
            idx = sequence.index("$")
            sequence = sequence[: idx + 1]

        if self.reverse:
            sequence = list(reversed(sequence))

        return sequence

    @property
    def vocab_size(self):
        """Return the number of amino acids"""
        return len(self._aa2idx)

    @property
    def device(self):
        """The current device for the model"""
        return next(self.parameters()).device


class PeptideEncoder(_PeptideTransformer):
    """A transformer encoder for peptide sequences.

    Parameters
    ----------
    dim_model : int
        The latent dimensionality to represent the amino acids in a peptide
        sequence.
    n_head : int, optional
        The number of attention heads in each layer. ``dim_model`` must be
        divisible by ``n_head``.
    dim_feedforward : int, optional
        The dimensionality of the fully connected layers in the Transformer
        layers of the model.
    n_layers : int, optional
        The number of Transformer layers.
    dropout : float, optional
        The dropout probability for all layers.
    pos_encoder : bool, optional
        Use positional encodings for the amino acid sequence.
    residues: Dict or str {"massivekb", "canonical"}, optional
        The amino acid dictionary and their masses. By default this is only
        the 20 canonical amino acids, with cysteine carbamidomethylated. If
        "massivekb", this dictionary will include the modifications found in
        MassIVE-KB. Additionally, a dictionary can be used to specify a custom
        collection of amino acids and masses.
    max_charge : int, optional
        The maximum charge state for peptide sequences.
    """

    def __init__(
        self,
        dim_model=128,
        n_head=8,
        dim_feedforward=1024,
        n_layers=1,
        dropout=0,
        pos_encoder=True,
        residues="canonical",
        max_charge=5,
    ):
        """Initialize a PeptideEncoder"""
        super().__init__(
            dim_model=dim_model,
            pos_encoder=pos_encoder,
            residues=residues,
            max_charge=max_charge,
        )

        # The Transformer layers:
        layer = torch.nn.TransformerEncoderLayer(
            d_model=dim_model,
            nhead=n_head,
            dim_feedforward=dim_feedforward,
            batch_first=True,
            dropout=dropout,
        )

        self.transformer_encoder = torch.nn.TransformerEncoder(
            layer,
            num_layers=n_layers,
        )

    def forward(self, sequences, charges):
        """Predict the next amino acid for a collection of sequences.

        Parameters
        ----------
        sequences : list of str or list of torch.Tensor of length batch_size
            The partial peptide sequences for which to predict the next
            amino acid. Optionally, these may be the token indices instead
            of a string.
        charges : torch.Tensor of size (batch_size,)
            The charge state of the peptide

        Returns
        -------
        latent : torch.Tensor of shape (n_sequences, len_sequence, dim_model)
            The latent representations for the spectrum and each of its
            peaks.
        mem_mask : torch.Tensor
            The memory mask specifying which elements were padding in X.
        """
        sequences = utils.listify(sequences)
        tokens = [self.tokenize(s) for s in sequences]
        tokens = torch.nn.utils.rnn.pad_sequence(tokens, batch_first=True)
        encoded = self.aa_encoder(tokens)

        # Encode charges
        charges = self.charge_encoder(charges - 1)[:, None]
        encoded = torch.cat([charges, encoded], dim=1)

        # Create mask
        mask = ~encoded.sum(dim=2).bool()

        # Add positional encodings
        encoded = self.pos_encoder(encoded)

        # Run through the model:
        latent = self.transformer_encoder(encoded, src_key_padding_mask=mask)
        return latent, mask


class PeptideDecoder(_PeptideTransformer):
    """A transformer decoder for peptide sequences.

    Parameters
    ----------
    dim_model : int, optional
        The latent dimensionality to represent peaks in the mass spectrum.
    n_head : int, optional
        The number of attention heads in each layer. ``dim_model`` must be
        divisible by ``n_head``.
    dim_feedforward : int, optional
        The dimensionality of the fully connected layers in the Transformer
        layers of the model.
    n_layers : int, optional
        The number of Transformer layers.
    dropout : float, optional
        The dropout probability for all layers.
    pos_encoder : bool, optional
        Use positional encodings for the amino acid sequence.
    reverse : bool, optional
        Sequence peptides from c-terminus to n-terminus.
    residues: Dict or str {"massivekb", "canonical"}, optional
        The amino acid dictionary and their masses. By default this is only
        the 20 canonical amino acids, with cysteine carbamidomethylated. If
        "massivekb", this dictionary will include the modifications found in
        MassIVE-KB. Additionally, a dictionary can be used to specify a custom
        collection of amino acids and masses.
    """

    def __init__(
        self,
        dim_model=128,
        n_head=8,
        dim_feedforward=1024,
        n_layers=1,
        dropout=0,
        pos_encoder=True,
        reverse=True,
        residues="canonical",
        max_charge=5,
        t_emb_dim = 128,
        t_emb_max_period: int = 10000
    ):
        """Initialize a PeptideDecoder"""
        super().__init__(
            dim_model=dim_model,
            pos_encoder=pos_encoder,
            residues=residues,
            max_charge=max_charge,
        )
        self.reverse = reverse
        self.t_emb_max_period = t_emb_max_period
        self.t_emb_dim = t_emb_dim

        # Additional model components
        self.mass_encoder = MassEncoder(dim_model)
        layer = torch.nn.TransformerDecoderLayer(
            d_model=dim_model,
            nhead=n_head,
            dim_feedforward=dim_feedforward,
            batch_first=True,
            dropout=dropout,
        )
         
        self.transformer_decoder = torch.nn.TransformerDecoder(
            layer,
            num_layers=n_layers,
        )

        self.t_emb_max_period = t_emb_max_period
        
        self.t_emb_dim = t_emb_dim
        
        self.t_layers = nn.Sequential(
            nn.Linear(t_emb_dim, dim_model),
            nn.SiLU(),
            nn.Linear(dim_model, dim_model)
        )
        self.final = torch.nn.Linear(dim_model, len(self._amino_acids) + 1)

    def forward(self, sequences, time_, precursors, memory, memory_key_padding_mask):
        """Predict the next amino acid for a collection of sequences.

        Parameters
        ----------
        sequences : list of str or list of torch.Tensor
            The partial peptide sequences for which to predict the next
            amino acid. Optionally, these may be the token indices instead
            of a string.
        time_: long tesnor of timestep indices
        precursors : torch.Tensor of size (batch_size, 2)
            The measured precursor mass (axis 0) and charge (axis 1) of each
            tandem mass spectrum
        memory : torch.Tensor of shape (batch_size, n_peaks, dim_model)
            The representations from a ``TransformerEncoder``, such as a
           ``SpectrumEncoder``.
        memory_key_padding_mask : torch.Tensor of shape (batch_size, n_peaks)
            The mask that indicates which elements of ``memory`` are padding.

        Returns
        -------
        scores : torch.Tensor of size (batch_size, len_sequence, n_amino_acids)
            The raw output for the final linear layer. These can be Softmax
            transformed to yield the probability of each amino acid for the
            prediction.
        tokens : torch.Tensor of size (batch_size, len_sequence)
            The input padded tokens.

        """
        # Prepare sequences
        if sequences is not None:
            sequences = utils.listify(sequences)
            tokens = [self.tokenize(s) for s in sequences]
            tokens = torch.nn.utils.rnn.pad_sequence(tokens, batch_first=True)
        else:
            tokens = torch.tensor([[]]).to(self.device)

        t = 0 
        if time_ is not None:
            self.t_emb = timestep_embedding(
                time_, self.t_emb_dim, self.t_emb_max_period, dtype=memory.dtype
            )
            t = self.t_layers(self.t_emb)[:, None]  # (bs, 1, dim)
        # Prepare mass and charge
        masses = self.mass_encoder(precursors[:, None, [0]])
        charges = self.charge_encoder(precursors[:, 1].int() - 1)
        precursors = masses + charges[:, None, :]

        # Feed through model:
        if sequences is None:
            tgt = precursors
        else:
            tgt = torch.cat([precursors, self.aa_encoder(tokens)], dim=1)

        tgt_key_padding_mask = tgt.sum(axis=2) == 0
        tgt = self.pos_encoder(tgt)
        tgt_mask = generate_tgt_mask(tgt.shape[1]).type_as(precursors)
        tgt = t + tgt
        preds = self.transformer_decoder(
            tgt=tgt,
            memory=memory,
            tgt_mask=tgt_mask,
            tgt_key_padding_mask=tgt_key_padding_mask,
            memory_key_padding_mask=memory_key_padding_mask.to(self.device),
        )
        return self.final(preds), tokens


def generate_tgt_mask(sz):
    """Generate a square mask for the sequence. The masked positions
    are filled with float('-inf'). Unmasked positions are filled with
    float(0.0).

    This function is a slight modification of the version in the PyTorch
    repository.

    Parameters
    ----------
    sz : int
        The length of the target sequence.
    """
    mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
    mask = (
        mask.float()
        .masked_fill(mask == 0, float("-inf"))
        .masked_fill(mask == 1, float(0.0))
    )
    return mask
