# I/O
import os
import sys
from pathlib import Path
import argparse
from datetime import datetime

# Data handling
import pandas as pd
import numpy as np
from pytorch_lightning.core import datamodule

# Modeling
import torch
from pytorch_lightning import Trainer, seed_everything
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.plugins import DDPPlugin
from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
from torchdrug import datasets, utils
from torchdrug.data import MoleculeDataset, DataLoader
from torchdrug.data.constant import ATOM_SYMBOL

# Custom
from datasets import ChemGraphDataModule
from models import ChemGraphEmbeddingNetwork


class ZINC5K(datasets.ZINC250k):
    """
    Testing dataset
    Subset of ZINC compound database for virtual screening.
    Statistics:
        - #Molecule: 5,000
        - #Regression task: 2
    Parameters:
        path (str): path to store the dataset
        verbose (int, optional): output verbose level
        **kwargs
    """

    url = "https://raw.githubusercontent.com/aspuru-guzik-group/chemical_vae/master/models/zinc_properties/" \
          "250k_rndm_zinc_drugs_clean_3.csv"
    md5 = "b59078b2b04c6e9431280e3dc42048d5"
    target_fields = ["logP", "qed"]

    def __init__(self, path, verbose=1, **kwargs):
        path = os.path.expanduser(path)
        if not os.path.exists(path):
            os.makedirs(path)
        self.path = path

        file_name = utils.download(self.url, path, md5=self.md5)

        sub_file_name = os.path.join(self.path, "5k_header_zinc_drugs_clean_3.csv")
        if not os.path.exists(sub_file_name):
            zinc5k = pd.read_csv(file_name, nrows=5001)
            zinc5k.to_csv(sub_file_name, index=False)
        file_name = sub_file_name

        self.load_csv(file_name, smiles_field="smiles", target_fields=self.target_fields,
                      verbose=verbose, **kwargs)


def process(args):
    torch.multiprocessing.set_sharing_strategy('file_system')
    # Pretraining data
    if args.dataset == "ZINC5k":
        dataset = ZINC5K("../data/torchdrug/molecule-datasets/", 
                         node_feature="pretrain",
                         edge_feature="pretrain",
                         lazy=True)
    elif args.dataset == "ZINC250k":
        dataset = datasets.ZINC250k("../data/torchdrug/molecule-datasets/", 
                                    node_feature="pretrain",
                                    edge_feature="pretrain",
                                    lazy=True)
    elif args.dataset == "ZINC2m":
        # defaults to lazy load
        dataset = datasets.ZINC2m("../data/torchdrug/molecule-datasets/", 
                                node_feature="pretrain",
                                edge_feature="pretrain")

    # CTRP smiles to embed
    ctrp = pd.read_csv("../data/drug_screens/CTRP/v20.meta.per_compound.txt", sep="\t")
    ctrp_ds = MoleculeDataset()
    ctrp_ds.load_smiles(smiles_list=ctrp['cpd_smiles'], 
                        targets=dict(), 
                        node_feature='pretrain',
                        edge_feature='pretrain')
    
    # Self-supervised pretraining
    dm = ChemGraphDataModule.from_argparse_args(args, train=dataset, predict=ctrp_ds)
    model = ChemGraphEmbeddingNetwork(task=args.task,
                                      input_dim=dataset.node_feature_dim,
                                      hidden_dims=[512]*5,
                                      edge_input_dim=dataset.edge_feature_dim,
                                      batch_norm=True,
                                      readout="mean",
                                      mask_rate=0.15)
    # Callbacks
    fname = f"{args.name}_{args.task}_{args.dataset}"
    logger = TensorBoardLogger(save_dir=args.default_root_dir,
                               version=fname,
                               name='lightning_logs')
    early_stop = EarlyStopping(monitor='accuracy',
                               min_delta=0.001,
                               patience=5,
                               verbose=False,
                               mode='max')
    checkpoint_callback = ModelCheckpoint(monitor='accuracy',
                                          mode='max')
    trainer = Trainer.from_argparse_args(args,
                                         default_root_dir=logger.log_dir,
                                         logger=logger,
                                         callbacks=[early_stop, checkpoint_callback],
                                         strategy=DDPPlugin(find_unused_parameters=False),
                                         profiler='simple')
    trainer.fit(model, dm)

    # Generate CTRP embeddings 
    model.to('cpu')
    model.eval()
    dl = DataLoader(ctrp_ds, batch_size=len(ctrp_ds))
    graph_embeds = []
    node_embeds = []
    for batch in dl:
        graph_feature, node_feature = model(batch)
        graph_embeds.append(graph_feature.detach())
        node_embeds.append(node_feature.detach())
    graph_embeds = torch.cat(graph_embeds).numpy()
    node_embeds = torch.cat(node_embeds).numpy()

    # Write out
    node_cpd_ids = [np.repeat(cpd_id, n['graph'].num_node) for n,cpd_id in zip(ctrp_ds, ctrp['broad_cpd_id'])]
    node_cpd_ids = np.concatenate(node_cpd_ids)
    node_embeds = pd.DataFrame(node_embeds, index=node_cpd_ids)
    node_embeds['atom_type'] = np.concatenate([[ATOM_SYMBOL[a] for a in n['graph'].atom_type] for n in ctrp_ds])
    graph_embeds = pd.DataFrame(graph_embeds, index=ctrp['broad_cpd_id'])
    node_embeds.to_csv(f"../data/torchdrug/molecule-datasets/{fname}_ctrp_node_embeds.csv", sep=",")
    graph_embeds.to_csv(f"../data/torchdrug/molecule-datasets/{fname}_ctrp_graph_embeds.csv", sep=",")


def main():
    """Parse Arguments"""
    desc = "Pretrain graph molecule representations."
    parser = argparse.ArgumentParser(description=desc,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--name", type=str,
                        help="Prepended name of experiment")
    parser.add_argument("--dataset", type=str, choices=["ZINC5k", "ZINC250k", "ZINC2m"], default="ZINC5k")
    parser = ChemGraphEmbeddingNetwork.add_model_specific_args(parser)
    parser = ChemGraphDataModule.add_argparse_args(parser)
    parser = Trainer.add_argparse_args(parser)                       
    args = parser.parse_args()
    
    process(args)


if __name__ == "__main__":
    sys.exit(main())