import os
import re
import click
import warnings
from os.path import join

# torch stuff
from torchsummary import summary
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Subset

# mpi4py
import mpi4py
mpi4py.rc.initialize = False
from mpi4py import MPI

# molecules stuff
from molecules.ml.datasets import ContactMapDataset
from molecules.ml.hyperparams import OptimizerHyperparams
from molecules.ml.callbacks import (LossCallback, CheckpointCallback,
                                    SaveEmbeddingsCallback, TSNEPlotCallback)
from molecules.ml.unsupervised.vae import VAE, SymmetricVAEHyperparams, ResnetVAEHyperparams

def parse_dict(ctx, param, value):
    if value is not None:
        token = value.split(",")
        result = {}
        for item in token:
            k, v = item.split("=")
            result[k] = v
        return result

@click.command()
@click.option('-i', '--input', 'input_path', required=True,
              type=click.Path(exists=True),
              help='Path to file containing preprocessed contact matrix data')

@click.option('-dn', '--dataset_name', default='contact_map',
              help='Name of the dataset in the HDF5 file.')

@click.option('-rn', '--rmsd_name', default='rmsd',
              help='Name of the RMSD data in the HDF5 file.')

@click.option('-fn', '--fnc_name', default='fnc',
              help='Name of the fraction of contacts data in the HDF5 file.')

@click.option('-o', '--out', 'out_path', required=True,
              type=click.Path(exists=True),
              help='Output directory for model data')

@click.option('-m', '--model_prefix', required=True,
              help='Model prefix for file naming.')

@click.option('-h', '--dim1', required=True, type=int,
              help='H of (H,W) shaped contact matrix')

@click.option('-w', '--dim2', required=True, type=int,
              help='W of (H,W) shaped contact matrix')

@click.option('-c', '--checkpoint',
             type=click.Path(exists=True),
             help='Model checkpoint file to resume training. ' \
                  'Checkpoint files saved as .pt by CheckpointCallback.')

@click.option('-r', '--resume',is_flag=True,
              help='Resume from latest checkpoint')

@click.option('-iw', '--init_weights',
             type=click.Path(exists=True),
             help='Model checkpoint file to load model weights fromg. ' \
                  'Checkpoint files saved as .pt by CheckpointCallback.')

@click.option('-f', '--cm_format', default='sparse-concat',
              help='Format of contact map files. Options ' \
                   '[full, sparse-concat, sparse-rowcol]')

@click.option('-E', '--encoder_gpu', default=None, type=int,
              help='Encoder GPU id')

@click.option('-D', '--decoder_gpu', default=None, type=int,
              help='Decoder GPU id')

@click.option('-e', '--epochs', default=10, type=int,
              help='Number of epochs to train for')

@click.option('-b', '--batch_size', default=128, type=int,
              help='Batch size for training')

@click.option('-opt', '--optimizer', callback=parse_dict,
              help='Optimizer parameters')

@click.option('-lw', '--loss_weights', callback=parse_dict, default="lambda_rec=1.0",
              help='Loss weight components')

@click.option('-t', '--model_type', default='resnet',
              help='Model architecture option: [resnet, symmetric]')

@click.option('-d', '--latent_dim', default=10, type=int,
              help='Number of dimensions in latent space')

@click.option('-erl', '--encoder_resnet_layers', default=None, type=int,
              help='Number of resnet layers to target')

@click.option('-sf', '--scale_factor', default=2, type=int,
              help='Scale factor hparam for resnet VAE')

@click.option('-ei', '--embed_interval', default=1, type=int,
              help="Saves embeddings every interval'th point")

@click.option('-ti', '--tsne_interval', default=1, type=int,
              help='Saves model checkpoints, embeddings, tsne plots every ' \
                   "interval'th point")

@click.option('-S', '--sample_interval', default=20, type=int,
              help="For embedding plots. Plots every sample_interval'th point")

@click.option('-wp', '--wandb_project_name', default=None, type=str,
              help='Project name for wandb logging')

@click.option('--local_rank', default=None, type=int,
              help='Local rank on the machine, required for DDP')

@click.option('-a', '--amp', is_flag=True,
              help='Specify if we want to enable automatic mixed precision (AMP)')

@click.option('--distributed', is_flag=True,
              help='Enable distributed training')

@click.option('-ndw', '--num_data_workers', default=0, type=int, 
              help='Number of data loaders for training')

def main(input_path, dataset_name, rmsd_name, fnc_name, out_path, checkpoint, resume, model_prefix,
         dim1, dim2, init_weights, cm_format, encoder_gpu, decoder_gpu, epochs, batch_size, optimizer, loss_weights, model_type,
         latent_dim, encoder_resnet_layers, scale_factor, embed_interval, tsne_interval, sample_interval, 
         wandb_project_name, local_rank, amp, distributed, num_data_workers):

    """Example for training Fs-peptide with either Symmetric or Resnet VAE."""
    
    if tsne_interval < embed_interval:
        warnings.warn('Found tsne_interval < embed_interval. Will result in ' \
                      'duplicated t-SNE plots.')

    # do some scaffolding for DDP
    comm_rank = 0
    comm_size = 1
    comm_local_rank = 0
    comm = None
    if distributed and dist.is_available():

        # init mpi4py:
        MPI.Init_thread()

        # get communicator: duplicate from comm world
        comm = MPI.COMM_WORLD.Dup()

        # now match ranks between the mpi comm and the nccl comm
        os.environ["WORLD_SIZE"] = str(comm.Get_size())
        os.environ["RANK"] = str(comm.Get_rank())

        # init torch distributed
        dist.init_process_group(backend='nccl',
                                init_method='env://')
        comm_rank = dist.get_rank()
        comm_size = dist.get_world_size()

        if local_rank is not None:
            comm_local_rank = local_rank
        else:
            comm_local_rank = int(os.getenv("LOCAL_RANK", 0))
            
        if comm_rank == 0:
            print("Distributed setup complete")

    assert model_type in ['symmetric', 'resnet']

    # Note: See SymmetricVAEHyperparams, ResnetVAEHyperparams class definitions
    #       for hyperparameter options. 

    if model_type == 'symmetric':
        # Optimal Fs-peptide params
        fs_peptide_hparams ={'filters': [100, 100, 100, 100],
                             'kernels': [5, 5, 5, 5],
                             'strides': [1, 2, 1, 1],
                             'affine_widths': [64],
                             'affine_dropouts': [0],
                             'latent_dim': latent_dim,
                             'output_activation': 'None'}

        input_shape = (1, dim1, dim2)
        hparams = SymmetricVAEHyperparams(**fs_peptide_hparams)

    elif model_type == 'resnet':
        resnet_hparams = {'max_len': dim1,
                          'nchars': dim2,
                          'latent_dim': latent_dim,
                          'dec_filters': dim1,
                          'enc_reslayers': encoder_resnet_layers,
                          'scale_factor': scale_factor,
                          'lambda_rec': float(loss_weights["lambda_rec"]),
                          'output_activation': 'None'}

        input_shape = (dim1, dim1)
        hparams = ResnetVAEHyperparams(**resnet_hparams)

    optimizer_hparams = OptimizerHyperparams(name=optimizer["name"], hparams={'lr': float(optimizer["lr"])})

    # For ease of training multiple models
    model_path = join(out_path, f'model-{model_prefix}')

    # Save hparams to disk 
    if comm_rank == 0:
        os.makedirs(model_path, exist_ok=True)
        hparams.save(join(model_path, 'model-hparams.json'))
        optimizer_hparams.save(join(model_path, 'optimizer-hparams.json'))

    # create model
    vae = VAE(input_shape, hparams, optimizer_hparams,
              gpu=(encoder_gpu, decoder_gpu),
              enable_amp=amp,
              init_weights=init_weights)

    enc_device = torch.device(f'cuda:{encoder_gpu}')
    dec_device = torch.device(f'cuda:{decoder_gpu}')
    if comm_size > 1:
        if (encoder_gpu == decoder_gpu):
            vae.model = DDP(vae.model, device_ids = [enc_device], output_device = enc_device)
        else:
            vae.model = DDP(vae.model, device_ids = None, output_device = None)

    # set global default device
    torch.cuda.set_device(enc_device.index)

    # Diplay model
    if comm_rank == 0:
        print(vae)
        # Only print summary when encoder_gpu is None or 0
        #summary(vae.model, input_shape)

    # Load training and validation data
    # training
    train_dataset = ContactMapDataset(input_path,
                                      dataset_name,
                                      rmsd_name,
                                      fnc_name,
                                      input_shape,
                                      split='train',
                                      cm_format=cm_format)

    # split across nodes
    if comm_size > 1:
        chunksize = len(train_dataset) // comm_size
        train_dataset = Subset(train_dataset,
                               list(range(chunksize * comm_rank, chunksize * (comm_rank + 1))))
    
    train_loader = DataLoader(train_dataset,
                              batch_size = batch_size,
                              drop_last = True,
                              shuffle = True,
                              pin_memory = True,
                              num_workers = num_data_workers)

    # validation
    valid_dataset = ContactMapDataset(input_path,
                                      dataset_name,
                                      rmsd_name,
                                      fnc_name,
                                      input_shape,
                                      split='valid',
                                      cm_format=cm_format)

    # split across nodes
    if comm_size > 1:
        chunksize = len(valid_dataset) // comm_size
        valid_dataset = Subset(valid_dataset,
                               list(range(chunksize * comm_rank, chunksize * (comm_rank + 1))))
    
    valid_loader = DataLoader(valid_dataset,
                              batch_size = batch_size,
                              drop_last = True,
                              shuffle = True,
                              pin_memory = True,
                              num_workers = num_data_workers)

    ## we call next once here to make sure the data is pinned to the right GPU
    #with torch.cuda.device(enc_device.index):
    #    _ = next(train_loader)
    #    _ = valid_loader.next()

    # For ease of training multiple models
    model_path = join(out_path, f'model-{model_prefix}')
    os.makedirs(model_path, exist_ok=True)

    # do we want wandb
    wandb_config = None
    if (comm_rank == 0) and (wandb_project_name is not None):
        import wandb
        wandb.init(project = wandb_project_name,
                   name = model_prefix,
                   id = model_prefix,
                   dir = model_path,
                   resume = False)
        wandb_config = wandb.config
        
        # log HP
        wandb_config.dim1 = dim1
        wandb_config.dim2 = dim2
        wandb_config.latent_dim = latent_dim
        
        # optimizer
        wandb_config.optimizer_name = optimizer_hparams.name
        for param in optimizer_hparams.hparams:
            wandb_config['optimizer_' + param] = optimizer_hparams.hparams[param]

        # loss weight
        wandb_config.lambda_rec = hparams.lambda_rec
            
        # watch model
        wandb.watch(vae.model)
    
    # Optional callbacks
    loss_callback = LossCallback(join(model_path, 'loss.json'),
                                 wandb_config=wandb_config,
                                 mpi_comm=comm)
    
    checkpoint_callback = CheckpointCallback(out_dir=join(model_path, 'checkpoint'),
                                             mpi_comm=comm)

    save_callback = SaveEmbeddingsCallback(out_dir=join(model_path, 'embeddings'),
                                           interval=embed_interval,
                                           sample_interval=sample_interval,
                                           mpi_comm=comm)

    # TSNEPlotCallback requires SaveEmbeddingsCallback to run first
    tsne_callback = TSNEPlotCallback(out_dir=join(model_path, 'embeddings'),
                                     projection_type='3d',
                                     target_perplexity=100,
                                     colors=['rmsd', 'fnc'],
                                     interval=tsne_interval,
                                     wandb_config=wandb_config,
                                     mpi_comm=comm)

    # Train model with callbacks
    callbacks = [loss_callback, checkpoint_callback, save_callback, tsne_callback]

    # see if resume is set
    if resume and (checkpoint is None):
        clist = [x for x in os.listdir(join(model_path, 'checkpoint')) if x.endswith(".pt")]
        checkpoints = sorted(clist, key=lambda x: re.match("epoch-\d*?-(\d*?-\d*?).pt", x).groups()[0])
        if checkpoints:
            checkpoint = join(model_path, 'checkpoint', checkpoints[-1])
            if comm_rank == 0:
                print(f"Resuming from checkpoint {checkpoint}.")
        else:
            if comm_rank == 0:
                print(f"No checkpoint files in directory {join(model_path, 'checkpoint')}, \
                       cannot resume training, will start from scratch.")
    
    # create model
    vae.train(train_loader, valid_loader, epochs,
              checkpoint=checkpoint, callbacks=callbacks)

    if comm_rank == 0:
        # Save loss history to disk.
        loss_callback.save(join(model_path, 'loss.json'))

        # Save final model weights to disk
        vae.save_weights(join(model_path, 'encoder-weights.pt'),
                         join(model_path, 'decoder-weights.pt'))

    # Output directory structure
    #  out_path
    # ├── model_path
    # │   ├── checkpoint
    # │   │   ├── epoch-1-20200606-125334.pt
    # │   │   └── epoch-2-20200606-125338.pt
    # │   ├── decoder-weights.pt
    # │   ├── encoder-weights.pt
    # │   ├── loss.json
    # │   ├── model-hparams.json
    # │   └── optimizer-hparams.json
    # |   |__ wandb_cache/

if __name__ == '__main__':
    # set forkserver
    torch.multiprocessing.set_start_method('forkserver', force = True)
    
    # call main
    main()
