import os

import time

import random

import pickle

import argparse

import os.path as osp

import torch

import torch.utils.data

from torch import nn

from torch_geometric.loader import DataLoader

import wandb

from rdkit import RDLogger

torch.set_num_threads(5)

RDLogger.DisableLog('rdApp.*')

from src.util.utils import *

from src.model.models import Generator, Discriminator, simple_disc

from src.data.dataset import DruggenDataset

from src.data.utils import get_encoders_decoders, load_molecules

from src.model.loss import discriminator_loss, generator_loss

class Train(object):

    """Trainer for DrugGEN."""

    def __init__(self, config):

        if config.set_seed:

            np.random.seed(config.seed)

            random.seed(config.seed)

            torch.manual_seed(config.seed)

            torch.cuda.manual_seed_all(config.seed)

            torch.backends.cudnn.deterministic = True

            torch.backends.cudnn.benchmark = False

            os.environ["PYTHONHASHSEED"] = str(config.seed)

            print(f'Using seed {config.seed}')

        self.device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')

        # Initialize configurations

        self.submodel = config.submodel

        # Data loader.

        self.raw_file = config.raw_file  # SMILES containing text file for dataset. 

                                         # Write the full path to file.

        self.drug_raw_file = config.drug_raw_file  # SMILES containing text file for second dataset. 

                                                   # Write the full path to file.

        # Automatically infer dataset file names from raw file names

        raw_file_basename = osp.basename(self.raw_file)

        drug_raw_file_basename = osp.basename(self.drug_raw_file)

        # Get the base name without extension and add max_atom to it

        self.max_atom = config.max_atom  # Model is based on one-shot generation.

        raw_file_base = os.path.splitext(raw_file_basename)[0]

        drug_raw_file_base = os.path.splitext(drug_raw_file_basename)[0]

        # Change extension from .smi to .pt and add max_atom to the filename

        self.dataset_file = f"{raw_file_base}{self.max_atom}.pt"

        self.drugs_dataset_file = f"{drug_raw_file_base}{self.max_atom}.pt"

        self.mol_data_dir = config.mol_data_dir  # Directory where the dataset files are stored.

        self.drug_data_dir = config.drug_data_dir  # Directory where the drug dataset files are stored.

        self.dataset_name = self.dataset_file.split(".")[0]

        self.drugs_dataset_name = self.drugs_dataset_file.split(".")[0]

        self.features = config.features  # Small model uses atom types as node features. (Boolean, False uses atom types only.)

                                         # Additional node features can be added. Please check new_dataloarder.py Line 102.

        self.batch_size = config.batch_size  # Batch size for training.

        self.parallel = config.parallel

        # Get atom and bond encoders/decoders

        atom_encoder, atom_decoder, bond_encoder, bond_decoder = get_encoders_decoders(

            self.raw_file,

            self.drug_raw_file,

            self.max_atom

        )

        self.atom_encoder = atom_encoder

        self.atom_decoder = atom_decoder

        self.bond_encoder = bond_encoder

        self.bond_decoder = bond_decoder

        self.dataset = DruggenDataset(self.mol_data_dir,

                                     self.dataset_file,

                                     self.raw_file,

                                     self.max_atom,

                                     self.features,

                                     atom_encoder=atom_encoder,

                                     atom_decoder=atom_decoder,

                                     bond_encoder=bond_encoder,

                                     bond_decoder=bond_decoder)

        self.loader = DataLoader(self.dataset,

                                 shuffle=True,

                                 batch_size=self.batch_size,

                                 drop_last=True)  # PyG dataloader for the GAN.

        self.drugs = DruggenDataset(self.drug_data_dir, 

                                 self.drugs_dataset_file, 

                                 self.drug_raw_file, 

                                 self.max_atom, 

                                 self.features,

                                 atom_encoder=atom_encoder,

                                 atom_decoder=atom_decoder,

                                 bond_encoder=bond_encoder,

                                 bond_decoder=bond_decoder)

        self.drugs_loader = DataLoader(self.drugs, 

                                       shuffle=True,

                                       batch_size=self.batch_size, 

                                       drop_last=True)  # PyG dataloader for the second GAN.

        self.m_dim = len(self.atom_decoder) if not self.features else int(self.loader.dataset[0].x.shape[1]) # Atom type dimension.

        self.b_dim = len(self.bond_decoder) # Bond type dimension.

        self.vertexes = int(self.loader.dataset[0].x.shape[0]) # Number of nodes in the graph.

        # Model configurations.

        self.act = config.act

        self.lambda_gp = config.lambda_gp

        self.dim = config.dim

        self.depth = config.depth

        self.heads = config.heads

        self.mlp_ratio = config.mlp_ratio

        self.ddepth = config.ddepth

        self.ddropout = config.ddropout

        # Training configurations.

        self.epoch = config.epoch

        self.g_lr = config.g_lr

        self.d_lr = config.d_lr

        self.dropout = config.dropout

        self.beta1 = config.beta1

        self.beta2 = config.beta2

        # Directories.

        self.log_dir = config.log_dir

        self.sample_dir = config.sample_dir

        self.model_save_dir = config.model_save_dir

        # Step size.

        self.log_step = config.log_sample_step

        # resume training

        self.resume = config.resume

        self.resume_epoch = config.resume_epoch

        self.resume_iter = config.resume_iter

        self.resume_directory = config.resume_directory

        # wandb configuration

        self.use_wandb = config.use_wandb

        self.online = config.online

        self.exp_name = config.exp_name

        # Arguments for the model.

        self.arguments = "{}_{}_glr{}_dlr{}_dim{}_depth{}_heads{}_batch{}_epoch{}_dataset{}_dropout{}".format(self.exp_name, self.submodel, self.g_lr, self.d_lr, self.dim, self.depth, self.heads, self.batch_size, self.epoch, self.dataset_name, self.dropout)

        self.build_model(self.model_save_dir, self.arguments)

    def build_model(self, model_save_dir, arguments):

        """Create generators and discriminators."""

        ''' Generator is based on Transformer Encoder: 

            @ g_conv_dim: Dimensions for MLP layers before Transformer Encoder

            @ vertexes: maximum length of generated molecules (atom length)

            @ b_dim: number of bond types

            @ m_dim: number of atom types (or number of features used)

            @ dropout: dropout possibility

            @ dim: Hidden dimension of Transformer Encoder

            @ depth: Transformer layer number

            @ heads: Number of multihead-attention heads

            @ mlp_ratio: Read-out layer dimension of Transformer

            @ drop_rate: depricated  

            @ tra_conv: Whether module creates output for TransformerConv discriminator

            '''

        self.G = Generator(self.act,

                           self.vertexes,

                           self.b_dim,

                           self.m_dim,

                           self.dropout,

                           dim=self.dim,

                           depth=self.depth,

                           heads=self.heads,

                           mlp_ratio=self.mlp_ratio)

        ''' Discriminator implementation with Transformer Encoder:

            @ act: Activation function for MLP

            @ vertexes: maximum length of generated molecules (molecule length)

            @ b_dim: number of bond types

            @ m_dim: number of atom types (or number of features used)

            @ dropout: dropout possibility

            @ dim: Hidden dimension of Transformer Encoder

            @ depth: Transformer layer number

            @ heads: Number of multihead-attention heads

            @ mlp_ratio: Read-out layer dimension of Transformer'''

        self.D = Discriminator(self.act,

                                self.vertexes,

                                self.b_dim,

                                self.m_dim,

                                self.ddropout,

                                dim=self.dim,

                                depth=self.ddepth,

                                heads=self.heads,

                                mlp_ratio=self.mlp_ratio)

        self.g_optimizer = torch.optim.AdamW(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])

        self.d_optimizer = torch.optim.AdamW(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])

        network_path = os.path.join(model_save_dir, arguments)

        self.print_network(self.G, 'G', network_path)

        self.print_network(self.D, 'D', network_path)

        if self.parallel and torch.cuda.device_count() > 1:

            print(f"Using {torch.cuda.device_count()} GPUs!")

            self.G = nn.DataParallel(self.G)

            self.D = nn.DataParallel(self.D)

        self.G.to(self.device)

        self.D.to(self.device)

    def print_network(self, model, name, save_dir):

        """Print out the network information."""

        num_params = 0

        for p in model.parameters():

            num_params += p.numel()

        if not os.path.exists(save_dir):

            os.makedirs(save_dir)

        network_path = os.path.join(save_dir, "{}_modules.txt".format(name))

        with open(network_path, "w+") as file:

            for module in model.modules():

                file.write(f"{module.__class__.__name__}:\n")

                print(module.__class__.__name__)

                for n, param in module.named_parameters():

                    if param is not None:

                        file.write(f"  - {n}: {param.size()}\n")

                        print(f"  - {n}: {param.size()}")

                break

            file.write(f"Total number of parameters: {num_params}\n")

            print(f"Total number of parameters: {num_params}\n\n")

    def restore_model(self, epoch, iteration, model_directory):

        """Restore the trained generator and discriminator."""

        print('Loading the trained models from epoch / iteration {}-{}...'.format(epoch, iteration))

        G_path = os.path.join(model_directory, '{}-{}-G.ckpt'.format(epoch, iteration))

        D_path = os.path.join(model_directory, '{}-{}-D.ckpt'.format(epoch, iteration))

        self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))

        self.D.load_state_dict(torch.load(D_path, map_location=lambda storage, loc: storage))

    def save_model(self, model_directory, idx,i):

        G_path = os.path.join(model_directory, '{}-{}-G.ckpt'.format(idx+1,i+1))

        D_path = os.path.join(model_directory, '{}-{}-D.ckpt'.format(idx+1,i+1))

        torch.save(self.G.state_dict(), G_path)

        torch.save(self.D.state_dict(), D_path)

    def reset_grad(self):

        """Reset the gradient buffers."""

        self.g_optimizer.zero_grad()

        self.d_optimizer.zero_grad()

    def train(self, config):

        ''' Training Script starts from here'''

        if self.use_wandb:

            mode = 'online' if self.online else 'offline'

        else:

            mode = 'disabled'

        kwargs = {'name': self.exp_name, 'project': 'druggen', 'config': config,

                'settings': wandb.Settings(_disable_stats=True), 'reinit': True, 'mode': mode, 'save_code': True}

        wandb.init(**kwargs)

        wandb.save(os.path.join(self.model_save_dir, self.arguments, "G_modules.txt"))

        wandb.save(os.path.join(self.model_save_dir, self.arguments, "D_modules.txt"))

        self.model_directory = os.path.join(self.model_save_dir, self.arguments)

        self.sample_directory = os.path.join(self.sample_dir, self.arguments)

        self.log_path = os.path.join(self.log_dir, "{}.txt".format(self.arguments))

        if not os.path.exists(self.model_directory):

            os.makedirs(self.model_directory)

        if not os.path.exists(self.sample_directory):

            os.makedirs(self.sample_directory)

        # smiles data for metrics calculation.

        drug_smiles = [line for line in open(self.drug_raw_file, 'r').read().splitlines()]

        drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]

        drug_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in drug_mols if x is not None]

        if self.resume:

            self.restore_model(self.resume_epoch, self.resume_iter, self.resume_directory)

        # Start training.

        print('Start training...')

        self.start_time = time.time()

        for idx in range(self.epoch):

            # =================================================================================== #

            #                             1. Preprocess input data                                #

            # =================================================================================== #

            # Load the data

            dataloader_iterator = iter(self.drugs_loader)

            wandb.log({"epoch": idx})

            for i, data in enumerate(self.loader):

                try:

                    drugs = next(dataloader_iterator)

                except StopIteration:

                    dataloader_iterator = iter(self.drugs_loader)

                    drugs = next(dataloader_iterator)

                wandb.log({"iter": i})

                # Preprocess both dataset

                real_graphs, a_tensor, x_tensor = load_molecules(

                    data=data,

                    batch_size=self.batch_size,

                    device=self.device,

                    b_dim=self.b_dim,

                    m_dim=self.m_dim,

                )

                drug_graphs, drugs_a_tensor, drugs_x_tensor = load_molecules(

                    data=drugs,

                    batch_size=self.batch_size,

                    device=self.device,

                    b_dim=self.b_dim,

                    m_dim=self.m_dim,

                )

                # Training configuration.

                GEN_node = x_tensor             # Generator input node features (annotation matrix of real molecules)

                GEN_edge = a_tensor             # Generator input edge features (adjacency matrix of real molecules)

                if self.submodel == "DrugGEN":

                    DISC_node = drugs_x_tensor  # Discriminator input node features (annotation matrix of drug molecules)

                    DISC_edge = drugs_a_tensor  # Discriminator input edge features (adjacency matrix of drug molecules)

                elif self.submodel == "NoTarget":

                    DISC_node = x_tensor      # Discriminator input node features (annotation matrix of real molecules)

                    DISC_edge = a_tensor      # Discriminator input edge features (adjacency matrix of real molecules)

                # =================================================================================== #

                #                                     2. Train the GAN                                #

                # =================================================================================== #

                loss = {}

                self.reset_grad()

                # Compute discriminator loss.

                node, edge, d_loss = discriminator_loss(self.G,

                                            self.D,

                                            DISC_edge,

                                            DISC_node,

                                            GEN_edge,

                                            GEN_node,

                                            self.batch_size,

                                            self.device,

                                            self.lambda_gp)

                d_total = d_loss

                wandb.log({"d_loss": d_total.item()})

                loss["d_total"] = d_total.item()

                d_total.backward()

                self.d_optimizer.step()

                self.reset_grad()

                # Compute generator loss.

                generator_output = generator_loss(self.G,

                                                    self.D,

                                                    GEN_edge,

                                                    GEN_node,

                                                    self.batch_size)

                g_loss, node, edge, node_sample, edge_sample = generator_output

                g_total = g_loss

                wandb.log({"g_loss": g_total.item()})

                loss["g_total"] = g_total.item()

                g_total.backward()

                self.g_optimizer.step()

                # Logging.

                if (i+1) % self.log_step == 0:

                    logging(self.log_path, self.start_time, i, idx, loss, self.sample_directory,

                            drug_smiles,edge_sample, node_sample, self.dataset.matrices2mol,

                            self.dataset_name, a_tensor, x_tensor, drug_vecs)

                    mol_sample(self.sample_directory, edge_sample.detach(), node_sample.detach(),

                               idx, i, self.dataset.matrices2mol, self.dataset_name)

                    print("samples saved at epoch {} and iteration {}".format(idx,i))

                    self.save_model(self.model_directory, idx, i)

                    print("model saved at epoch {} and iteration {}".format(idx,i))

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    # Data configuration.

    parser.add_argument('--raw_file', type=str, required=True)

    parser.add_argument('--drug_raw_file', type=str, required=False, help='Required for DrugGEN model, optional for NoTarget')

    parser.add_argument('--drug_data_dir', type=str, default='data')

    parser.add_argument('--mol_data_dir', type=str, default='data')

    parser.add_argument('--features', action='store_true', help='features dimension for nodes')

    # Model configuration.

    parser.add_argument('--submodel', type=str, default="DrugGEN", help="Chose model subtype: DrugGEN, NoTarget", choices=['DrugGEN', 'NoTarget'])

    parser.add_argument('--act', type=str, default="relu", help="Activation function for the model.", choices=['relu', 'tanh', 'leaky', 'sigmoid'])

    parser.add_argument('--max_atom', type=int, default=45, help='Max atom number for molecules must be specified.')

    parser.add_argument('--dim', type=int, default=128, help='Dimension of the Transformer Encoder model for the GAN.')

    parser.add_argument('--depth', type=int, default=1, help='Depth of the Transformer model from the GAN.')

    parser.add_argument('--ddepth', type=int, default=1, help='Depth of the Transformer model from the discriminator.')

    parser.add_argument('--heads', type=int, default=8, help='Number of heads for the MultiHeadAttention module from the GAN.')

    parser.add_argument('--mlp_ratio', type=int, default=3, help='MLP ratio for the Transformer.')

    parser.add_argument('--dropout', type=float, default=0., help='dropout rate')

    parser.add_argument('--ddropout', type=float, default=0., help='dropout rate for the discriminator')

    parser.add_argument('--lambda_gp', type=float, default=10, help='Gradient penalty lambda multiplier for the GAN.')

    # Training configuration.

    parser.add_argument('--batch_size', type=int, default=128, help='Batch size for the training.')

    parser.add_argument('--epoch', type=int, default=10, help='Epoch number for Training.')

    parser.add_argument('--g_lr', type=float, default=0.00001, help='learning rate for G')

    parser.add_argument('--d_lr', type=float, default=0.00001, help='learning rate for D')

    parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for Adam optimizer')

    parser.add_argument('--beta2', type=float, default=0.999, help='beta2 for Adam optimizer')

    parser.add_argument('--log_dir', type=str, default='experiments/logs')

    parser.add_argument('--sample_dir', type=str, default='experiments/samples')

    parser.add_argument('--model_save_dir', type=str, default='experiments/models')

    parser.add_argument('--log_sample_step', type=int, default=1000, help='step size for sampling during training')

    # Resume training.

    parser.add_argument('--resume', type=bool, default=False, help='resume training')

    parser.add_argument('--resume_epoch', type=int, default=None, help='resume training from this epoch')

    parser.add_argument('--resume_iter', type=int, default=None, help='resume training from this step')

    parser.add_argument('--resume_directory', type=str, default=None, help='load pretrained weights from this directory')

    # Seed configuration.

    parser.add_argument('--set_seed', action='store_true', help='set seed for reproducibility')

    parser.add_argument('--seed', type=int, default=1, help='seed for reproducibility')

    # wandb configuration.

    parser.add_argument('--use_wandb', action='store_true', help='use wandb for logging')

    parser.add_argument('--online', action='store_true', help='use wandb online')

    parser.add_argument('--exp_name', type=str, default='druggen', help='experiment name')

    parser.add_argument('--parallel', action='store_true', help='Parallelize training')

    config = parser.parse_args()

    # Check if drug_raw_file is provided when using DrugGEN model

    if config.submodel == "DrugGEN" and not config.drug_raw_file:

        parser.error("--drug_raw_file is required when using DrugGEN model")

    # If using NoTarget model and drug_raw_file is not provided, use a dummy file

    if config.submodel == "NoTarget" and not config.drug_raw_file:

        config.drug_raw_file = "data/akt_train.smi"  # Use a reference file for NoTarget model (AKT) (not used for training for ease of use and encoder/decoder's)

    trainer = Train(config)

    trainer.train(config)