import warnings
warnings.filterwarnings("ignore")

import torch
import torch.nn as nn
import numpy as np

# Early stopping class for scAdam and scEve models
class EarlyStopping:
    def __init__(
        self, patience = 10, delta = 1e-4, mode = 'min', verbose = True
    ):
        """
        Early stopping class for tracking validation metrics
        
        Parameters
        ----------
        patience: int (default: 10)
            Number of consecutive epochs without improvement before performing early stopping.
            If patience is set to 0, then no early stopping will be performed.
            Note that if patience is enabled, then best weights from best epoch will automatically be loaded at the end of the training.
        delta: float (default: 1e-4)
            Minimal change in metric to use patients.
        mode: str
            Mode for early stopping trigger.
            'max' default for scAdam (accuracy metrics), 'min' default for scEve (error metrics).
        verbose: bool (default: True)
            Show progress bar for each epoch during training. 
            
        Internal abbreviations:
            bs: best score of metric
            es: early stopping
            bm: best model
        """
        self.patience = patience
        self.delta = delta
        self.mode = mode
        self.verbose = verbose
        self.bs = None
        self.es = False
        self.bm = None
        self.patience_count = 0

    def __call__(self, score, model) -> bool:
        if self.bs is None:
            self.bs = score
            self.save_model(model)
            return False

        if self.mode == 'max': # 'max' mode for scAdam metrics: accuracy, balanced accuracy, f1 score
            improved = score > (self.bs + self.delta)
        else: # 'min' mode for error metrics in scEve model training
            improved = score < (self.bs - self.delta)

        if improved: # if metric improved set patience_count to 0
            self.bs = score
            self.patience_count = 0
            self.save_model(model)
        else: # if not improved increase patience_count by 1
            self.patience_count += 1
            if self.patience_count >= self.patience: # Check if patience_count greater than patients
                self.es = True
                if self.verbose:
                    print(f'Early stopping triggered! Best score: {self.bs:.4f}')
        return self.es

    # Save best model
    def save_model(self, model): # Saves best model state to get best model before patience
        self.bm = {k: v.cpu().clone() for k, v in model.state_dict().items()}

    # Load best model
    def load_bm(self, model): # Loads best model state trained before patience
        model.load_state_dict(self.bm)


# Class for scRNA-seq data augmentation
class Augmentation:
    def __init__(
        self,
        prob = 0.15,
        noise_std = 0.1,
        dropout_prob = 0.1,
        alpha = 0.2
    ):
        """
        scRNA-seq data augmentation module. 
        Uses different augmentation methods:
            1) masking_augmentation - set a fraction of features (genes) to zero using gene masking probability (prob).
            2) noise_augmentation - add Gaussian noise using noize standard deviation (noise_std).
            3) dropout_augmentation - randomly zeros out each feature (gene) with probability (dropout_prob). 
            4) Generate new cell by combining two cells and their cell types using a weighted linear interpolation between them.
        Fourth augmentation only for scEve model training to address cell type imbalance.
        
        Parameters
        ----------
        prob: float (default: 0.15)
            Gene masking probability.
        noise_std: float (default: 0.1)
            Gaussian noise standard deviation.
        dropout_prob: float (default: 0.1)
            Dropout probability for simulating technical noise.
        alpha: float (default: 0.2)
            Alpha parameter for mixup augmentation.
        """
        # Save parameters
        self.prob = prob
        self.noise_std = noise_std
        self.dropout_prob = dropout_prob
        self.alpha = alpha

    # Set a fraction of features (genes) to zero with probability
    def masking_augmentation(self, x):
        mask = torch.rand_like(x) > self.prob
        return x * mask.float()

    # Add Gaussian noise with standard deviation
    def noise_augmentation(self, x):
        noise = torch.randn_like(x) * self.noise_std
        return torch.clamp(x + noise, min=0.0, max=1.0) # clip values to stay within [0,1]. 

    # Randomly zeros out each feature (gene) with probability 
    def dropout_augmentation(self, x):
        mask = torch.rand_like(x) > self.dropout_prob
        return x * mask.float()

    # Generate new cell by combining two cells and their cell types using a weighted linear interpolation between them.
    # Reduce overfitting by exposing the model to a broader range of input variations
    # Used only in scEve model training
    def mix_cells(self, cell_1, cell_2, celltype_1, celltype_2): # Addresses cell type imbalance in scEve model training
        weight = np.random.beta(self.alpha, self.alpha)
        # Generate new cell by linearly mix two input cells with weight coefficient
        x_mixed = weight * cell_1 + (1 - weight) * cell_2
        # Assign the cell type from the cell that has the larger contribution in mixed cell generation
        y_mixed = celltype_1 if weight > 0.5 else celltype_2
        return x_mixed, y_mixed
    
    def __call__(self, x):
        # Use only one augmentation at a time
        augmentation_type = np.random.randint(0, 3)
        if augmentation_type == 0:
            x = self.masking_augmentation(x)
        elif augmentation_type == 1:
            x = self.noise_augmentation(x)
        else:
            x = self.dropout_augmentation(x)
        
        return x


# Multi-Head Axial Attention
class Attention(nn.Module):
    def __init__(self, gn, nh = 8, dropout = 0.3):
        """
        Axial attention.
        Linearly projects input genes into queries, keys, and values.

        Parameters
        ----------
        gn: int
            Number of input genes.
        nh: int (default: 8) 
            Number of heads in Attention.
        dropout: float (default: 0.1)
            Portion of neurons that temporarily ignored during training (prevents overfitting).
        """
        super().__init__()
        
        self.gn = gn
        self.nh = nh
        self.hd = gn // nh
        self.scale = self.hd ** -0.5

        # Single linear layer that produces Q, K, V in one shot.
        self.qkv = nn.Linear(gn, gn * 3, bias=False)
        # Output projection to merge all heads back into the original feature (gene) space.
        self.proj = nn.Linear(gn, gn)
        # The dropout makes the attention weights and output projections more robust and stable to overfitting.
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        # Get: B = batch size, N = number of features (genes), C = feature dimensionality
        B, N, C = x.shape

        # Compute Q, K, V in a single matmul, then reshape to separate heads and components
        qkv = self.qkv(x).reshape(B, N, 3, self.nh, self.hd)
        # Reorder dimensions to (3, B, nh, N, hd) 
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, nh, N, hd)
        q, k, v = qkv[0], qkv[1], qkv[2]
        
        # Attention
        # Scaled dot-product attention: for each head compute similarity between queries and keys
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.dropout(attn)
        
        # Combine heads
        # Apply attention weights to values, then bring heads back into the sequence dimension
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.dropout(x)
        
        return x, attn


# Feed-Forward neural network
class FFNN(nn.Module):
    def __init__(self, gn, ff_hd, dropout = 0.3):
        """
        Feed Forward. It expands each gene embedding into a higher-dimensional space, 
        applies a GELU activation and dropout, and then projects it back to the original gene dimension. 
        It increases the representational power of the transformer block (RibBlock) 
        by applying a non-linear transformation to each gene independently, while dropout helps prevent overfitting.
        
        Parameters
        ----------
        gn: int
            Number of input genes.
        ff_hd: int (default: 512) 
            Number of nodes in each layer in feed forward network.
        dropout: float (default: 0.1)
            Portion of neurons that temporarily ignored during training (prevents overfitting).
        """
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(gn, ff_hd),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ff_hd, gn),
            nn.Dropout(dropout)
        )
    
    def forward(self, x):
        return self.net(x)
        

# Main Block for scAdam and scEve transformer models
# "God created Eve by placing Adam in a deep sleep, removing one of his ribs (or "side"), 
#  and fashioning it into the first woman to be a partner for him"
class RibBlock(nn.Module): 
    def __init__(
        self, 
        input_dim, 
        nh = 8, 
        ff_hd = 512,
        dropout = 0.3
    ):
        """
        Transformer block with Attention and FeedForward.
        Main block for scAdam and scEve models.

        Parameters
        ----------
        input_dim: int
            Number of input features (ex. genes) in adata.
        nh: int (default: 8) 
            Number of heads in Attention.
        ff_hd: int (default: 512) 
            Number of nodes in each layer in feed forward network.
        dropout: float (default: 0.1)
            Portion of neurons that temporarily ignored during training (prevents overfitting).
        """
        super().__init__()
        self.norm1 = nn.LayerNorm(input_dim)
        self.attn = Attention(input_dim, nh, dropout) # Attention to mix information across genes
        self.norm2 = nn.LayerNorm(input_dim)
        self.ff = FFNN(input_dim, ff_hd, dropout) # FFNN to transform each gene embedding independently
        
    def forward(self, x):
        x_norm = self.norm1(x)
        attn_out, attn_weights = self.attn(x_norm)
        x = x + attn_out
        x = x + self.ff(self.norm2(x))
        
        return x, attn_weights


# Feature Embedding layer
class Embedding(nn.Module):
    def __init__(self, input_dim, ed, nc = 16, dropout = 0.1):
        """
        Feature (gene) embedding. Transforms feature expression vectors into a compact, 
        structured representation by dividing genes into fixed-size chunks, 
        projecting each chunk into an embedding space, and adding learnable positional encodings.
        Chunking reduces computational cost and memory usage.

        Parameters
        ----------
        input_dim: int 
            Number of input features (ex. genes) in adata.
        ed: int
            Embedding dimensionality.
        nc: int (default: 16)
            Number of chunks for features from adata.

        Internal abbreviations:
            gcs: number of genes in chunk.
            xc: feature (gene) expression matrix created from chunks.
            xe: feature (gene) expression matrix after embedding creation.
        """
        super().__init__()
        self.input_dim = input_dim
        self.ed = ed
        self.nc = nc
        
        # Convert genes to chunks
        self.gcs = input_dim // nc # Divide number of genes by number of chunks
        # gcs: number of genes in chunk
        if input_dim % nc != 0:
            self.gcs += 1
        
        # Chunk projection
        self.chunkembed = nn.Linear(self.gcs, ed)
        
        # Positional encoding
        self.posembed = nn.Parameter(torch.randn(1, nc, ed) * 0.02)
        self.norm = nn.LayerNorm(ed)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x): # x - feature (gene) matrix
        # Get batch size from gene matrix
        B = x.shape[0]

        # Split the gene expression into chunks
        chunks = []
        for i in range(self.nc):
            start_idx = i * self.gcs
            end_idx = min((i + 1) * self.gcs, self.input_dim)
            chunk = x[:, start_idx:end_idx]
            # Process last chunk
            # Adds zeros if it's smaller than could be
            if chunk.shape[1] < self.gcs: # Check if there are fewer genes in the last chunk than there should be in the chunk
                zeros = torch.zeros(B, self.gcs - chunk.shape[1], 
                                     device=x.device, dtype=x.dtype)
                chunk = torch.cat([chunk, zeros], dim=1)
            chunks.append(chunk)

        # Create new gene matrix embedding using chunks
        xc = torch.stack(chunks, dim=1)
        xe = self.chunkembed(xc)
        xe = xe + self.posembed
        xe = self.norm(xe)
        xe = self.dropout(xe)
        
        return xe