import warnings
warnings.filterwarnings('ignore')
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim.lr_scheduler import ReduceLROnPlateau, MultiStepLR
from torch import optim
from torch.utils.data import Dataset, DataLoader
import os
import numpy as np
import re
import pickle
import pandas as pd
import random
import csv
from sklearn.preprocessing import LabelEncoder
from transformers import ( 
        RobertaTokenizer, RoFormerTokenizer,
        RoFormerModel,
        pipeline
    )

def seed_torch(seed=0):
	random.seed(seed)
	os.environ['PYTHONHASHSEED'] = str(seed) 
	np.random.seed(seed)
	torch.manual_seed(seed)
	torch.cuda.manual_seed(seed)
	torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
	torch.backends.cudnn.benchmark = False
	torch.backends.cudnn.deterministic = True
	os.environ['CUBLAS_WORKSPACE_CONFIG']=':16:8'
	torch.use_deterministic_algorithms(True, warn_only=True)


def format_bcr_for_tokenizer(df, igh_col='IGH', igl_col='IGL'):
    """
    Convert B cell sequences to tokenizer format.
    For each B cell, combine IGH and IGL sequences with [SEP] token.
    
    Parameters:
    -----------
    df : pandas.DataFrame
        DataFrame containing B cell data with IGH and IGL columns
    igh_col : str
        Column name for heavy chain sequences (default: 'IGH')
    igl_col : str
        Column name for light chain sequences (default: 'IGL')
    
    Returns:
    --------
    pandas.DataFrame
        DataFrame with additional 'tokenized_seq' column
    """
    
    def format_sequence(row):
        """Format a single B cell sequence"""
        igh_seq = row[igh_col]
        igl_seq = row[igl_col]
        
        # Check if sequences are valid
        if pd.isna(igh_seq) or pd.isna(igl_seq):
            return None
        
        # Convert sequences to space-separated amino acids
        igh_formatted = ' '.join(list(igh_seq))
        igl_formatted = ' '.join(list(igl_seq))
        
        # Combine with [SEP] token
        tokenized_seq = f"Ḣ {igh_formatted} [SEP] Ḷ {igl_formatted}"
        
        return tokenized_seq
    
    tokenized_seqs = df.apply(format_sequence, axis=1)
    
    return list(tokenized_seqs.values)

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

class Logger(object):

    def __init__(self, path, header):
        self.log_file = open(path, 'w')
        self.logger = csv.writer(self.log_file, delimiter='\t')

        self.logger.writerow(header)
        self.header = header

    def __del(self):
        self.log_file.close()

    def log(self, values):
        write_values = []
        for col in self.header:
            assert col in values
            write_values.append(values[col])

        self.logger.writerow(write_values)
        self.log_file.flush()
        
class EarlyStopping:
    """Early stops the training if validation loss doesn't improve after a given patience."""
    def __init__(self, save_path, checkpoint, patience=7, verbose=False, delta=0, save_epoch=None):
        """
        Args:
            save_path: the path for saved model
            checkpoint: the name for the saved model
            patience (int): How long to wait after last time validation loss improved.
                            Default: 7
            verbose (bool): If True, prints a message for each validation loss improvement. 
                            Default: False
            delta (float): Minimum change in the monitored quantity to qualify as an improvement.
                            Default: 0
            save_epoch (int, optional): If specified, save model at this specific epoch and stop training.
                            If None, use early stopping to save best model. Default: None
        """
        self.save_path = save_path
        self.checkpoint = checkpoint
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.val_loss_min = np.Inf
        self.delta = delta
        self.save_epoch = save_epoch
        self.current_epoch = 0
        self.should_stop = False  # New flag for save_epoch mode

    def __call__(self, val_loss, model, epoch):
        self.current_epoch = epoch
        
        # If save_epoch is specified, only save at that epoch and stop
        if self.save_epoch is not None:
            if epoch == self.save_epoch:
                self.save_checkpoint(val_loss, model)
                print(f'Saved model at specified epoch {self.save_epoch}. Stopping training.')
                self.should_stop = True  # Signal to stop training
            return
        
        # Otherwise, use early stopping logic
        score = -val_loss

        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
        elif score < self.best_score * (1-self.delta):
            self.counter += 1
            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
            self.counter = 0

    def save_checkpoint(self, val_loss, model):
        '''Saves model when validation loss decrease.'''
        if self.verbose or self.save_epoch is not None:
            if self.save_epoch is not None:
                print(f'Saving model at epoch {self.save_epoch}')
            else:
                print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')
        path = os.path.join(self.save_path, self.checkpoint)
        torch.save(model.state_dict(), path)
        self.val_loss_min = val_loss
        
class ProjectionHead(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(ProjectionHead, self).__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.layers(x)

class BCRencoder(nn.Module):
    def __init__(self):
        super(BCRencoder, self).__init__()
        # Find the directory of the current file (combcr.py)
        current_dir = os.path.dirname(os.path.abspath(__file__))
        self.berta = RoFormerModel.from_pretrained(os.path.join(current_dir, "BCRencoder"))
        #for p in self.parameters():
        #    p.requires_grad = False
        
    def forward(self, x):
        bertaoutput = self.berta(**x).last_hidden_state
        embs = []
        
        attention_mask = x['attention_mask']
        
        for i in range(len(x['input_ids'])):
            actual_length = attention_mask[i].sum().item()
            valid_embeddings = bertaoutput[i, 1:actual_length-1]
            embs.append(valid_embeddings.mean(0).reshape(1, -1))
        
        x = torch.concat(embs)
        #x = bertaoutput[:,0,:].reshape(-1, bertaoutput.shape[-1])
        return x
        

class Encoder_profile(nn.Module):
    """
    this is the profile encoder class, which contains the 3 linear transformation
    layers and 1 layer normalization
    
    Parameters:
        param in_dim: the input profile dimension
        param hid_dim: the hidden dimension of first transformation layer
        param hid_dim2: the hidden dimension of second transformation layer
        param out_dim: the output dimension
    
    Returns:
        new profile compressed embedding(bs, out_dim)
    """
    
    def __init__(self, in_dim = 5, hid_dim = 5, hid_dim2 = 5, out_dim = 5):
        super(Encoder_profile, self).__init__()
        self.layer1 = nn.Linear(in_dim, hid_dim)
        self.layer2 = nn.Linear(hid_dim, hid_dim2)
        self.layer3 = nn.Linear(hid_dim2, out_dim)
        
        # define the tanh as activation function
        self.activation = nn.ReLU()
        
        # define the layer normalization
        self.layer_norm = nn.LayerNorm(out_dim)
        
    def forward(self, profile):
        compressed_profile = self.activation(self.layer1(profile))
        compressed_profile = self.activation(self.layer2(compressed_profile))
        compressed_profile = self.activation(self.layer3(compressed_profile))
        compressed_profile = self.layer_norm(compressed_profile)
        return compressed_profile

class CustomDataset(Dataset):
    def __init__(self, bcrfile_path, rnafile_path, separatebatch=False, user_defined_cluster=False):
        super(CustomDataset, self).__init__()
        rnafile = pd.read_csv(rnafile_path, sep=",", index_col="barcode", low_memory=False)
        self.rnamatrix = torch.from_numpy(rnafile.values).float()
        
        bcrfile = pd.read_csv(bcrfile_path, sep=",", index_col="barcode", low_memory=False)
        le = LabelEncoder()
        
        # For paired-chain: concatenate IGH and IGL
        bcrfile["whole_seq"] = bcrfile['IGH'].str.cat(bcrfile['IGL'])
        bcrfile['sample'] = bcrfile['sample'].astype(str)
        
        # BCR identity (for loss_cmc)
        if separatebatch:
            print("keep invariant of a batch")
            bcrfile["new_whole_seq"] = bcrfile["sample"].str.cat(bcrfile["whole_seq"], sep=":")
            self.bcridentity = le.fit_transform(bcrfile.new_whole_seq.tolist())
        else:
            self.bcridentity = le.fit_transform(bcrfile.whole_seq.tolist())
        self.bcridentity = torch.from_numpy(self.bcridentity)
        
        # Format BCR sequences for tokenizer
        self.bcr = format_bcr_for_tokenizer(bcrfile)
        
        # Cluster label (for loss_p2p_inner)
        if user_defined_cluster and 'cluster_label' in bcrfile.columns:
            print("Using cluster_label for profile contrastive loss")
            le_cluster = LabelEncoder()
            if separatebatch:
                bcrfile["new_cluster_label"] = bcrfile["sample"].str.cat(bcrfile["cluster_label"].astype(str), sep=":")
                self.cluster_identity = le_cluster.fit_transform(bcrfile.new_cluster_label.tolist())
            else:
                self.cluster_identity = le_cluster.fit_transform(bcrfile.cluster_label.astype(str).tolist())
            self.cluster_identity = torch.from_numpy(self.cluster_identity)
            self.has_cluster = True
        else:
            if user_defined_cluster:
                print("Warning: cluster_label not found, using BCR identity for intra-modal contrastive loss")
            # using BCR identity for intra-modal contrastive loss
            self.cluster_identity = self.bcridentity.clone()
            self.has_cluster = False
        
        assert(len(rnafile)==len(bcrfile))
    
    def __len__(self):
        return self.rnamatrix.shape[0]
    
    def __getitem__(self, idx):
        identity = self.bcridentity[idx]
        cluster_id = self.cluster_identity[idx]
        gex = self.rnamatrix[idx]
        bcr = self.bcr[idx]
        
        return bcr, gex, identity, cluster_id, idx

class CoMBCR_model(nn.Module):
    def __init__(self, encoderBCR_out_dim,
                 encoderprofile_in_dim, encoderprofile_hid_dim, encoderprofile_hid_dim2, encoderprofle_out_dim):
        super(CoMBCR_model, self).__init__()
        
        # define the BCRencoder
        
        self.encoder_BCR = BCRencoder()
        self.BCR_project = ProjectionHead(1024, encoderBCR_out_dim, encoderBCR_out_dim)
        # define the temperature (default set as 1)
        self.encoder_profile = Encoder_profile(encoderprofile_in_dim, encoderprofile_hid_dim, 
                                              encoderprofile_hid_dim2, encoderprofle_out_dim)
        # define the project
        self.profile_proj = ProjectionHead(encoderprofle_out_dim, encoderprofle_out_dim, encoderprofle_out_dim)
        
        
        
    def forward(self, bcr_token, exp):
             
        # extract the features of BCR
        encoderBCR_feature = self.encoder_BCR(bcr_token)
        encoderBCR_embedding = self.BCR_project(encoderBCR_feature)
        # normalization of BCR sequence embedding
        encoderBCR_embedding = F.normalize(encoderBCR_embedding, dim = -1)
        
        # extract the profile features based on the encoder_profile
        encoderprofile_feature = self.encoder_profile(exp)
        encoderprofile_embedding = self.profile_proj(encoderprofile_feature)
        # normalization of expression embedding
        encoderprofile_embedding = F.normalize(encoderprofile_embedding, dim = -1)
        
        return encoderBCR_embedding, encoderprofile_embedding
    
class mydefine_loss(nn.Module):
    def __init__(self, temperature=0.2, lam=0.1):
        super().__init__()
        self.temperature = temperature
        self.lam = lam
        
    def forward(self, encoderBCR_embedding, encoderprofile_embedding, batch_BCRsimilar, batch_expdist, idx_BCR, idx_cluster):
        '''
        output: graphs.num x emb
        batch_BCRdist/batch_expdist: batch_size x batch_size
        idx_BCR: BCR identity for loss_cmc
        idx_cluster: cluster label for loss_p2p_inner
        '''
        
        # pMHC encoder and TCR encoder similarity score calculation
        sim_p2b = encoderprofile_embedding @ encoderBCR_embedding.T / self.temperature
        sim_b2p = encoderBCR_embedding @ encoderprofile_embedding.T / self.temperature
        
        # same BCR with multiple different profiles (for loss_cmc)
        with torch.no_grad():
            idx_BCR = idx_BCR.view(-1, 1)
            pos_idx_BCR = torch.eq(idx_BCR, idx_BCR.T)
            
            sim_targets_bcr = (pos_idx_BCR).float().to(encoderBCR_embedding.device)
            sim_targets_bcr = sim_targets_bcr / sim_targets_bcr.sum(1, keepdim=True)
        
        # p2b loss
        loss_p2t = -torch.sum(F.log_softmax(sim_p2b, dim=1) * sim_targets_bcr, dim=1).mean()
        loss_b2p = -torch.sum(F.log_softmax(sim_b2p, dim=1) * sim_targets_bcr, dim=1).mean()
        loss_cmc = (loss_p2t + loss_b2p) / 2
        
        # b2b loss
        sim_b2b = encoderBCR_embedding @ encoderBCR_embedding.T
        loss_b2b = ((sim_b2b - batch_BCRsimilar)**2).mean()
        
        # profile similarity score calculation
        sim_p2p = encoderprofile_embedding @ encoderprofile_embedding.T
        
        # same cluster profiles should be similar (for loss_p2p_inner)
        with torch.no_grad():
            idx_cluster = idx_cluster.view(-1, 1)
            pos_idx_cluster = torch.eq(idx_cluster, idx_cluster.T)
            
            sim_targets_cluster = (pos_idx_cluster).float().to(encoderBCR_embedding.device)
            sim_targets_cluster = sim_targets_cluster / sim_targets_cluster.sum(1, keepdim=True)
        
        loss_p2p_inner = -torch.sum(F.log_softmax(sim_p2p / self.temperature, dim=1) * sim_targets_cluster, dim=1).mean()
        
        # p2p loss
        loss_p2p = -torch.sum(F.log_softmax(1 - sim_p2p, dim=1) * torch.softmax(batch_expdist, dim=1), dim=1).mean() + self.lam * loss_p2p_inner
        
        return loss_cmc, loss_b2b, loss_p2p


def CoMBCR_main(bcrpath, rnapath, bcroriginal, outdir, checkpoint="best_network.pth", 
                lr=1e-6, lam=1e-1, batch_size=256, epochs=200, patience=15, lr_step=[40,100], 
                encoderprofile_in_dim=5000, separatebatch=False, user_defined_cluster=False, save_epoch=None):
    
    
    # Find the directory of the current file (combcr.py)
    current_dir = os.path.dirname(os.path.abspath(__file__))
    
    if not os.path.exists(os.path.join(outdir)):
        os.mkdir(os.path.join(outdir))
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    print('learning rate is ', lr)
    bcrfile_path = bcrpath
    rnafile_path = rnapath
    dataset = CustomDataset(bcrfile_path=bcrfile_path, rnafile_path=rnafile_path, 
                           separatebatch=separatebatch, user_defined_cluster=user_defined_cluster)
    
    # prepare intra-distance
    bcrsimilar = pd.read_csv(bcroriginal, index_col="barcode", low_memory=False).values
    assert(bcrsimilar.shape[0] == len(dataset))
    bcrsimilar = torch.from_numpy(bcrsimilar)
    bcrsimilar = F.normalize(bcrsimilar, dim = -1)
    bcrsimilar = bcrsimilar @ bcrsimilar.T
    rnavalues = pd.read_csv(rnafile_path, sep=",", index_col="barcode").values
    rnavalues = torch.from_numpy(rnavalues)
    rnadist = torch.cdist(rnavalues, rnavalues)
    
    seed_torch()
    tokenizer = RoFormerTokenizer.from_pretrained(os.path.join(current_dir, "BCRencoder"))
    early_stopping = EarlyStopping(outdir, patience=patience, checkpoint=checkpoint, save_epoch=save_epoch)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    runmodel = CoMBCR_model(encoderBCR_out_dim=256,
                     encoderprofile_in_dim=encoderprofile_in_dim, encoderprofile_hid_dim=1024, 
                     encoderprofile_hid_dim2=512, encoderprofle_out_dim=256)
    loss_cmc_epoch = AverageMeter()
    loss_b2b_epoch = AverageMeter()
    loss_p2p_epoch = AverageMeter()
    loss_all = AverageMeter()
    logger = Logger(os.path.join(outdir, "{}.log".format(checkpoint)), ['Epoch', 'loss', 'loss_cmc', 'loss_p2p', 'loss_b2b'])
    calculate_loss = mydefine_loss(temperature=0.2, lam=lam)
    calculate_loss.to(device)
    runmodel.to(device)
    #optimizer=optim.Adam(filter(lambda p: p.requires_grad, runmodel.parameters()), lr=lr)
    optimizer=optim.Adam(runmodel.parameters(), lr=lr)
    #scheduler = ReduceLROnPlateau(optimizer, patience=10, verbose=True)
    scheduler = MultiStepLR(optimizer, milestones=lr_step, verbose=True)
    epochs = epochs
    seed_torch()
    
    runmodel.train()
    for epoch in np.arange(epochs):
        loss_cmc_epoch.reset()
        loss_b2b_epoch.reset()
        loss_p2p_epoch.reset()
        loss_all.reset()
        for batch, (bcr, gex, identity, cluster_id, idx) in enumerate(loader):
            tokenized_input = tokenizer(list(bcr), return_tensors='pt', padding='max_length', truncation=True, max_length=255)
            tokenized_input = tokenized_input.to(device)
            gex = gex.to(device)
            identity_BCR = identity.to(device)
            cluster_identity = cluster_id.to(device)
            idxs = idx
            optimizer.zero_grad()
           
            encoderBCR_embedding, encoderprofile_embedding = runmodel(tokenized_input, gex)
            batch_bcr_simiar = bcrsimilar[:, idxs][idxs,:].to(device)
            batch_gex_dist = rnadist[:, idxs][idxs,:].to(device)
            loss_cmc, loss_b2b, loss_p2p = calculate_loss(encoderBCR_embedding, encoderprofile_embedding, 
                                                          batch_bcr_simiar, batch_gex_dist, 
                                                          identity_BCR, cluster_identity)
            loss = loss_cmc + loss_p2p/(loss_p2p.detach()/loss_cmc.detach()) + loss_b2b/(loss_b2b.detach()/loss_cmc.detach())
            #loss = loss_cmc/(loss_cmc.detach()) + loss_p2p/loss_p2p.detach() + loss_b2b/loss_b2b.detach()
            loss.backward()
            optimizer.step()
            
            loss_cmc_epoch.update(loss_cmc.item(), 1) 
            loss_b2b_epoch.update(loss_b2b.item(), 1)
            loss_p2p_epoch.update(loss_p2p.item(), 1)
            loss_all.update(loss_cmc.item()+loss_b2b.item()+loss_p2p.item(), 1)
            #loss_all.update(loss.item(), 1)

            # if batch %10 == 0:
            #     print('steps:{}\tloss:{:.5f}\tloss_cmc:{:.5f}\tloss_p2p:{:.5f}\tloss_b2b:{:.5f}'.format(batch, loss, loss_cmc, loss_p2p, loss_b2b))
            torch.cuda.empty_cache()
        # Log before checking stopping condition
        logger.log({'Epoch': epoch, 'loss': round(loss_all.avg, 6), 'loss_cmc': round(loss_cmc_epoch.avg, 6), 
                   'loss_p2p': round(loss_p2p_epoch.avg, 6), 'loss_b2b': round(loss_b2b_epoch.avg, 6)})
        if epoch % 1 == 0:
            print('Epoch:[{}/{}]\tloss:{:.5f}\tloss_cmc:{:.6f}\tloss_p2p:{:.6f}\tloss_b2b:{:.6f}'.format(
                epoch, epochs, loss_all.avg, loss_cmc_epoch.avg, loss_p2p_epoch.avg, loss_b2b_epoch.avg))
        
        # Pass epoch to early_stopping
        early_stopping(loss_all.avg, runmodel, epoch)
        scheduler.step()
        
        # Check stopping conditions
        if save_epoch is not None and early_stopping.should_stop:
            print(f"Training stopped after saving model at epoch {save_epoch}")
            break
        elif save_epoch is None and early_stopping.early_stop:
            print("Early stopping")
            break

    runmodel = CoMBCR_model(encoderBCR_out_dim=256,
                 encoderprofile_in_dim=encoderprofile_in_dim, encoderprofile_hid_dim=1024, 
                 encoderprofile_hid_dim2=512, encoderprofle_out_dim=256)
    loader = DataLoader(dataset, batch_size=256, shuffle=False)
    runmodel.load_state_dict(torch.load(os.path.join(outdir, checkpoint), weights_only=True))
    runmodel.eval()
    runmodel.to(device)
    BCRallembs = []
    gexallembs = []
    for epoch in np.arange(1):
        with torch.no_grad():
            for batch, (bcr, gex, identity, cluster_id, idx) in enumerate(loader):
                tokenized_input = tokenizer(list(bcr), return_tensors='pt', padding='max_length', truncation=True, max_length=255)
                tokenized_input = tokenized_input.to(device)
                encoderBCR_embedding = runmodel.encoder_BCR.eval()(tokenized_input)
                encoderBCR_embedding = runmodel.BCR_project.eval()(encoderBCR_embedding)
                # normalization of BCR sequence embedding
                encoderBCR_embedding = F.normalize(encoderBCR_embedding, dim = -1)
                BCRallembs.append(encoderBCR_embedding.cpu().numpy())
                
                gex = gex.to(device)
                encoderprofile_embedding = runmodel.encoder_profile.eval()(gex)
                encoderprofile_embedding = runmodel.profile_proj.eval()(encoderprofile_embedding)
                # normalization of BCR sequence embedding
                encoderprofile_embedding = F.normalize(encoderprofile_embedding, dim = -1)
                gexallembs.append(encoderprofile_embedding.cpu().numpy())
    BCRembeddings = np.concatenate(BCRallembs, axis=0)
    gexembeddings = np.concatenate(gexallembs, axis=0)
    
    if not os.path.exists(os.path.join(outdir, "Embeddings")):
        os.mkdir(os.path.join(outdir, "Embeddings"))
    
    # Read and save with barcode index
    bcrfile = pd.read_csv(bcrfile_path, sep=",", index_col="barcode", low_memory=False)
    
    # Save BCR embeddings with barcode index
    bcr_emb_df = pd.DataFrame(BCRembeddings, index=bcrfile.index)
    bcr_emb_df.to_csv(os.path.join(outdir, "Embeddings", "bcrembeddings.csv"), index=True)
    
    # Save GEX embeddings with barcode index
    gex_emb_df = pd.DataFrame(gexembeddings, index=bcrfile.index)
    gex_emb_df.to_csv(os.path.join(outdir, "Embeddings", "gexembeddings.csv"), index=True)
    
    return BCRembeddings, gexembeddings
 

