import itertools
import math
import os
from collections import OrderedDict

import torch
from torch import nn
from torch.nn.parallel.data_parallel import DataParallel
from tqdm import tqdm

from archai.common import ml_utils, utils
from archai.common.common import get_tb_writer
from archai.common.ordered_dict_logger import get_global_logger
from archai.supergraph.datasets.data import get_dataloaders
from archai.supergraph.models import get_model, num_class
from archai.supergraph.utils.metrics import Accumulator

logger = get_global_logger()


# TODO: remove scheduler parameter?
def run_epoch(
    conf, logger, model: nn.Module, loader, loss_fn, optimizer, split_type: str, epoch=0, verbose=1, scheduler=None
):
    """Runs epoch for given dataloader and model. If optimizer is supplied
    then backprop and model update is done as well. This can be called from
    test to train modes.
    """

    writer = get_tb_writer()

    # region conf vars
    conf_loader = conf["autoaug"]["loader"]
    epochs = conf_loader["epochs"]
    conf_opt = conf["autoaug"]["optimizer"]
    grad_clip = conf_opt["clip"]
    # endregion

    tqdm_disable = bool(os.environ.get("TASK_NAME", ""))  # TODO: remove?
    if verbose:
        loader = tqdm(loader, disable=tqdm_disable)
        loader.set_description("[%s %04d/%04d]" % (split_type, epoch, epochs))

    metrics = Accumulator()
    cnt = 0
    steps = 0
    for data, label in loader:
        steps += 1
        data, label = data.cuda(), label.cuda()

        if optimizer:
            optimizer.zero_grad()

        preds = model(data)
        loss = loss_fn(preds, label)

        if optimizer:
            loss.backward()
            if getattr(optimizer, "synchronize", None):
                optimizer.synchronize()  # for horovod
            # grad clipping defaults to 5 (same as Darts)
            if grad_clip > 0.0:
                nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            optimizer.step()

        top1, top5 = ml_utils.accuracy(preds, label, (1, 5))
        metrics.add_dict(
            {
                "loss": loss.item() * len(data),
                "top1": top1.item() * len(data),
                "top5": top5.item() * len(data),
            }
        )
        cnt += len(data)
        if verbose:
            postfix = metrics / cnt
            if optimizer:
                if "lr" in optimizer.param_groups[0]:
                    postfix["lr"] = optimizer.param_groups[0]["lr"]
            loader.set_postfix(postfix)

        # below changes LR for every batch in epoch
        # TODO: should we do LR step at epoch start only?
        # if scheduler is not None:
        #     scheduler.step(epoch - 1 + float(steps) / total_steps)

        del preds, loss, top1, top5, data, label

    if tqdm_disable:
        if optimizer:
            logger.info(
                "[%s %03d/%03d] %s lr=%.6f", split_type, epoch, epochs, metrics / cnt, optimizer.param_groups[0]["lr"]
            )
        else:
            logger.info("[%s %03d/%03d] %s", split_type, epoch, epochs, metrics / cnt)

    metrics /= cnt
    if optimizer:
        if "lr" in optimizer.param_groups[0]:
            metrics.metrics["lr"] = optimizer.param_groups[0]["lr"]
    if verbose:
        for key, value in metrics.items():
            writer.add_scalar("{}/{}".format(key, split_type), value, epoch)
    return metrics


# NOTE that 'eval' is overloaded in this code base. 'eval' here means
# taking a trained model and running it on val or test sets. In NAS 'eval'
# often means taking a found model and training it fully (often termed 'final training').

# metric could be 'last', 'test', 'val', 'train'.
def train_and_eval(conf, val_ratio, val_fold, save_path, only_eval, reporter=None, metric="test"):

    writer = get_tb_writer()

    # region conf vars
    conf_dataset = conf["dataset"]
    dataroot = utils.full_path(conf_dataset["dataroot"])
    horovod = conf["common"]["horovod"]
    checkpoint_freq = conf["common"]["checkpoint"]["freq"]
    conf_loader = conf["autoaug"]["loader"]
    conf_model = conf["autoaug"]["model"]
    ds_name = conf_dataset["name"]
    aug = conf_loader["aug"]
    cutout = conf_loader["cutout"]
    batch_size = conf_loader["batch"]
    max_batches = conf_dataset["max_batches"]
    epochs = conf_loader["epochs"]
    conf_model = conf["autoaug"]["model"]
    conf_opt = conf["autoaug"]["optimizer"]
    conf_lr_sched = conf["autoaug"]["lr_schedule"]
    n_workers = conf_loader["n_workers"]
    # endregion

    # initialize horovod
    # TODO: move to common init
    if horovod:
        import horovod.torch as hvd

        hvd.init()
        device = torch.device("cuda", hvd.local_rank())
        torch.cuda.set_device(device)

    if not reporter:

        def reporter(**kwargs):
            return 0

    # get dataloaders with transformations and splits applied
    train_dl, valid_dl, test_dl = get_dataloaders(
        ds_name,
        batch_size,
        dataroot,
        aug,
        cutout,
        load_train=True,
        load_test=True,
        val_ratio=val_ratio,
        val_fold=val_fold,
        horovod=horovod,
        n_workers=n_workers,
        max_batches=max_batches,
    )

    # create a model & an optimizer
    model = get_model(conf_model, num_class(ds_name), data_parallel=(not horovod))

    # select loss function and optimizer
    lossfn = nn.CrossEntropyLoss()
    optimizer = ml_utils.create_optimizer(conf_opt, model.parameters())

    # distributed optimizer if horovod is used
    is_master = True
    if horovod:
        optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
        # issue : https://github.com/horovod/horovod/issues/1099
        optimizer._requires_update = set()
        hvd.broadcast_parameters(model.state_dict(), root_rank=0)
        hvd.broadcast_optimizer_state(optimizer, root_rank=0)
        if hvd.rank() != 0:
            is_master = False
    logger.debug("is_master=%s" % is_master)

    # select LR schedule
    scheduler = ml_utils.create_lr_scheduler(conf_lr_sched, epochs, optimizer, len(train_dl))

    result = OrderedDict()
    epoch_start = 1
    # if model available from previous checkpount then load it
    if save_path and os.path.exists(save_path):
        logger.info("%s checkpoint found. loading..." % save_path)
        data = torch.load(save_path)

        # when checkpointing we do add 'model' key so other cases are special cases
        if "model" in data or "state_dict" in data:
            key = "model" if "model" in data else "state_dict"
            logger.info("checkpoint epoch@%d" % data["epoch"])

            # TODO: do we need change here?
            if not isinstance(model, DataParallel):
                # for non-dataparallel models, remove default 'module.' prefix
                model.load_state_dict({k.replace("module.", ""): v for k, v in data[key].items()})
            else:
                # for dataparallel models, make sure 'module.' prefix exist
                model.load_state_dict({k if "module." in k else "module." + k: v for k, v in data[key].items()})

            # load optimizer
            optimizer.load_state_dict(data["optimizer"])

            # restore epoch count
            if data["epoch"] < epochs:
                epoch_start = data["epoch"]
            else:
                # epochs finished, switch to eval mode
                only_eval = False

        else:
            model.load_state_dict({k: v for k, v in data.items()})
        del data
    else:
        logger.info(
            'model checkpoint does not exist at "%s". skip \
            to pretrain weights...'
            % save_path
        )
        only_eval = False  # we made attempt to load checkpt but as it does not exist, switch to train mode

    # if eval only then run model on train, test and val sets
    if only_eval:
        logger.info("evaluation only+")
        model.eval()
        rs = dict()  # stores metrics for each set
        rs["train"] = run_epoch(conf, logger, model, train_dl, lossfn, None, split_type="train", epoch=0)
        if valid_dl:
            rs["valid"] = run_epoch(conf, logger, model, valid_dl, lossfn, None, split_type="valid", epoch=0)
        rs["test"] = run_epoch(conf, logger, model, test_dl, lossfn, None, split_type="test", epoch=0)

        for key, setname in itertools.product(["loss", "top1", "top5"], ["train", "valid", "test"]):
            result["%s_%s" % (key, setname)] = rs[setname][key]

        result["epoch"] = 0
        return result

    # train loop
    best_top1, best_valid_loss = 0, 10.0e10
    max_epoch = epochs
    for epoch in range(epoch_start, max_epoch + 1):
        # if horovod:
        #     trainsampler.set_epoch(epoch)

        # run train epoch and update the model
        model.train()
        rs = dict()
        rs["train"] = run_epoch(
            conf,
            logger,
            model,
            train_dl,
            lossfn,
            optimizer,
            split_type="train",
            epoch=epoch,
            verbose=is_master,
            scheduler=scheduler,
        )
        if scheduler[0]:
            scheduler[0].step()

        model.eval()

        # check for nan loss
        if math.isnan(rs["train"]["loss"]):
            raise Exception("train loss is NaN.")

        # collect metrics on val and test set, checkpoint
        if epoch % checkpoint_freq == 0 or epoch == max_epoch:
            if valid_dl:
                rs["valid"] = run_epoch(
                    conf, logger, model, valid_dl, lossfn, None, split_type="valid", epoch=epoch, verbose=is_master
                )
            rs["test"] = run_epoch(
                conf, logger, model, test_dl, lossfn, None, split_type="test", epoch=epoch, verbose=is_master
            )

            # TODO: is this good enough condition?
            if rs[metric]["loss"] < best_valid_loss or rs[metric]["top1"] > best_top1:
                best_top1 = rs[metric]["top1"]
                best_valid_loss = rs[metric]["loss"]
                for key, setname in itertools.product(["loss", "top1", "top5"], ["train", "valid", "test"]):
                    result["%s_%s" % (key, setname)] = rs[setname][key]
                result["epoch"] = epoch

                writer.add_scalar("best_top1/valid", rs["valid"]["top1"], epoch)
                writer.add_scalar("best_top1/test", rs["test"]["top1"], epoch)

                reporter(
                    loss_valid=rs["valid"]["loss"],
                    top1_valid=rs["valid"]["top1"],
                    loss_test=rs["test"]["loss"],
                    top1_test=rs["test"]["top1"],
                )

                # save checkpoint
                if is_master and save_path:
                    logger.info("save model@%d to %s" % (epoch, save_path))
                    torch.save(
                        {
                            "epoch": epoch,
                            "log": {
                                "train": rs["train"].get_dict(),
                                "valid": rs["valid"].get_dict(),
                                "test": rs["test"].get_dict(),
                            },
                            "optimizer": optimizer.state_dict(),
                            "model": model.state_dict(),
                        },
                        save_path,
                    )

    del model

    result["top1_test"] = best_top1
    return result
