import argparse
import torch
import datetime
import os
import numpy as np
from dateutil import tz
from omegaconf import OmegaConf
from pytorch_lightning import seed_everything
from pytorch_lightning import loggers as pl_loggers
from pytorch_lightning.trainer import Trainer
from pytorch_lightning.callbacks import (
    ModelCheckpoint,
    EarlyStopping,
    LearningRateMonitor,
)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-c",
        "--config",
        metavar="base_config.yaml",
        help="paths to base config",
        required=True,
    )
    parser.add_argument(
        "--train", action="store_true", default=False, help="specify to train model"
    )
    parser.add_argument(
        "--test",
        action="store_true",
        default=False,
        help="specify to test model"
        "By default run.py trains a model based on config file",
    )
    parser.add_argument(
        "--ckpt_path", type=str, default=None, help="Checkpoint path for the save model"
    )
    parser.add_argument("--random_seed", type=int, default=23, help="Random seed")
    parser.add_argument(
        "--train_pct", type=float, default=1.0, help="Percent of training data"
    )
    parser.add_argument(
        "--splits",
        type=int,
        default=1,
        help="Train on n number of splits used for training. Defaults to 1",
    )
    parser = Trainer.add_argparse_args(parser)

    return parser


def main(cfg, args):

    # get datamodule
    dm = MLRL.builder.build_data_module(cfg)
    # define lightning module
    model = MLRL.builder.build_lightning_model(cfg, dm)
    # callbacks
    callbacks = [LearningRateMonitor(logging_interval="step")]
    if "checkpoint_callback" in cfg.lightning:
        checkpoint_callback = ModelCheckpoint(**cfg.lightning.checkpoint_callback)
        callbacks.append(checkpoint_callback)
    if "early_stopping_callback" in cfg.lightning:
        early_stopping_callback = EarlyStopping(**cfg.lightning.early_stopping_callback)
        callbacks.append(early_stopping_callback)
    if cfg.train.scheduler is not None:
        lr_monitor = LearningRateMonitor(logging_interval="step")
        callbacks.append(lr_monitor)

    # logging
    if "logger" in cfg.lightning:
        logger_type = cfg.lightning.logger.pop("logger_type")
        logger_class = getattr(pl_loggers, logger_type)
        cfg.lightning.logger.name = f"{cfg.experiment_name}_{cfg.extension}"
        logger = logger_class(**cfg.lightning.logger)
        cfg.lightning.logger.logger_type = logger_type
    else:
        logger = None

    # setup pytorch-lightning trainer
    cfg.lightning.trainer.val_check_interval = args.val_check_interval
    cfg.lightning.trainer.auto_lr_find = args.auto_lr_find
    trainer_args = argparse.Namespace(**cfg.lightning.trainer)
    if cfg.lightning.ckpt!='None':
        trainer = Trainer.from_argparse_args(
            args=trainer_args, deterministic=True, callbacks=callbacks, logger=logger,  resume_from_checkpoint = cfg.lightning.ckpt
        )
    else:
        trainer = Trainer.from_argparse_args(
            args=trainer_args, deterministic=True, callbacks=callbacks, logger=logger
        )

    # learning rate finder
    if trainer_args.auto_lr_find is not False:
        lr_finder = trainer.tuner.lr_find(model, datamodule=dm)
        new_lr = lr_finder.suggestion()
        model.lr = new_lr
        print("=" * 80 + f"\nLearning rate updated to {new_lr}\n" + "=" * 80)

    if args.train:
        trainer.fit(model, dm)
    if args.test:
        trainer.test(model=model, datamodule=dm)

    # save top weights paths to yaml
    if "checkpoint_callback" in cfg.lightning:
        ckpt_paths = os.path.join(
            cfg.lightning.checkpoint_callback.dirpath, "best_ckpts.yaml"
        )
        checkpoint_callback.to_yaml(filepath=ckpt_paths)


if __name__ == "__main__":

    # parse arguments
    parser = get_parser()
    args = parser.parse_args()
    cfg = OmegaConf.load(args.config)

    # edit experiment name
    cfg.data.frac = args.train_pct
    # if cfg.trial_name is not None:
    #     cfg.experiment_name = f"{cfg.experiment_name}_{cfg.trial_name}"
    if args.splits is not None:
        cfg.experiment_name = f"{cfg.experiment_name}_{args.train_pct}"  # indicate % data used in trial name

    # loop over the number of independent training splits, defaults to 1 split
    for split in np.arange(args.splits):

        # get current time
        now = datetime.datetime.now(tz.tzlocal())
        timestamp = now.strftime("%Y_%m_%d_%H_%M_%S")

        # random seed
        args.random_seed = split + 1
        seed_everything(args.random_seed)

        # set directory names
        cfg.extension = str(args.random_seed) if args.splits != 1 else timestamp
        cfg.output_dir = cfg.lightning.logger.save_dir
        cfg.lightning.checkpoint_callback.dirpath = os.path.join(
            cfg.lightning.checkpoint_callback.dirpath,
            f"{cfg.experiment_name}/{cfg.extension}",
        )

        # create directories
        if not os.path.exists(cfg.lightning.logger.save_dir):
            os.makedirs(cfg.lightning.logger.save_dir)
        if not os.path.exists(cfg.lightning.checkpoint_callback.dirpath):
            os.makedirs(cfg.lightning.checkpoint_callback.dirpath)
        if not os.path.exists(cfg.output_dir):
            os.makedirs(cfg.output_dir)

        # save config
        config_path = os.path.join(cfg.output_dir, "config.yaml")
        with open(config_path, "w") as fp:
            OmegaConf.save(config=cfg, f=fp.name)

        main(cfg, args)