import copy
import json
import os
import time
from collections import OrderedDict
from typing import Optional

import gorilla
import numpy as np
import ray
import torch
from hyperopt import hp
from ray.tune import register_trainable, run_experiments
from ray.tune.suggest import HyperOptSearch
from ray.tune.trial import Trial
from tqdm import tqdm

from archai.common.common import expdir_abspath
from archai.common.config import Config
from archai.common.ordered_dict_logger import get_global_logger
from archai.common.stopwatch import StopWatch
from archai.supergraph.datasets.augmentation import (
    augment_list,
    policy_decoder,
    remove_deplicates,
)
from archai.supergraph.datasets.data import get_dataloaders
from archai.supergraph.models import get_model, num_class
from archai.supergraph.utils.augmented_trainer import train_and_eval
from archai.supergraph.utils.metrics import Accumulator

logger = get_global_logger()


# this method is overriden version of ray.tune.trial_runner.TrialRunner.step using monkey patching
def _step_w_log(self):
    original = gorilla.get_original_attribute(ray.tune.trial_runner.TrialRunner, "step")

    # collect counts by status for all trials
    cnts = OrderedDict()
    for status in [Trial.RUNNING, Trial.TERMINATED, Trial.PENDING, Trial.PAUSED, Trial.ERROR]:
        cnt = len(list(filter(lambda x: x.status == status, self._trials)))
        cnts[status] = cnt

    # get the best top1 accuracy from all finished trials so far
    best_top1_acc = 0.0
    for trial in filter(lambda x: x.status == Trial.TERMINATED, self._trials):
        if not trial.last_result:  # TODO: why would this happen?
            continue
        best_top1_acc = max(best_top1_acc, trial.last_result["top1_valid"])

    # display best accuracy from all finished trial
    logger.info("iter", self._iteration, "top1_acc=%.3f" % best_top1_acc, cnts, end="\r")

    # call original step method
    return original(self)


# override ray.tune.trial_runner.TrialRunner.step method so we can print best accuracy at each step
patch = gorilla.Patch(ray.tune.trial_runner.TrialRunner, "step", _step_w_log, settings=gorilla.Settings(allow_hit=True))
gorilla.apply(patch)


@ray.remote(num_gpus=torch.cuda.device_count(), max_calls=1)
def _train_model(conf, dataroot, augment, val_ratio, val_fold, save_path=None, only_eval=False):
    Config.set_inst(conf)
    conf["autoaug"]["loader"]["aug"] = augment
    model_type = conf["autoaug"]["model"]["type"]

    result = train_and_eval(conf, val_ratio=val_ratio, val_fold=val_fold, save_path=save_path, only_eval=only_eval)
    return model_type, val_fold, result


def _get_model_filepath(dataset, model, tag) -> Optional[str]:
    filename = "%s_%s_%s.model" % (dataset, model, tag)
    return expdir_abspath(filename)


def _train_no_aug(conf):
    sw = StopWatch.get()

    # region conf vars
    conf_dataset = conf["dataset"]
    dataroot = conf["dataroot"]
    conf_loader = conf["autoaug"]["loader"]
    conf_model = conf["autoaug"]["model"]
    model_type = conf_model["type"]
    ds_name = conf_dataset["name"]
    aug = conf_loader["aug"]
    val_ratio = conf_loader["val_ratio"]
    epochs = conf_loader["epochs"]
    cv_num = conf_loader["cv_num"]
    # endregion

    logger.info("----- Train without Augmentations cv=%d ratio(test)=%.1f -----" % (cv_num, val_ratio))
    sw.start(tag="train_no_aug")

    # for each fold, we will save model
    save_paths = [_get_model_filepath(ds_name, model_type, "ratio%.1f_fold%d" % (val_ratio, i)) for i in range(cv_num)]

    # Train model for each fold, save model in specified path, put result
    # in reqs list. These models are trained with aug specified in config.
    # TODO: configuration will be changed ('aug' key),
    #   but do we really need deepcopy everywhere?
    reqs = [
        # TODO: eliminate need for deep copy as only aug key is changed
        _train_model.remote(
            copy.deepcopy(copy.deepcopy(conf)), dataroot, aug, val_ratio, i, save_path=save_paths[i], only_eval=True
        )
        for i in range(cv_num)
    ]

    # we now probe saved models for each fold to check the epoch number
    # they are on. When every fold crosses an epoch number, we update
    # the progress.
    tqdm_epoch = tqdm(range(epochs))
    is_done = False
    for epoch in tqdm_epoch:
        while True:
            epochs_per_cv = OrderedDict()
            for cv_idx in range(cv_num):
                try:
                    if os.path.exists(save_paths[cv_idx]):
                        latest_ckpt = torch.load(save_paths[cv_idx])
                        if "epoch" not in latest_ckpt:
                            epochs_per_cv["cv%d" % (cv_idx + 1)] = epochs
                            continue
                    else:
                        continue
                    epochs_per_cv["cv%d" % (cv_idx + 1)] = latest_ckpt["epoch"]
                except Exception:
                    continue
            tqdm_epoch.set_postfix(epochs_per_cv)
            if len(epochs_per_cv) == cv_num and min(epochs_per_cv.values()) >= epochs:
                is_done = True
            if len(epochs_per_cv) == cv_num and min(epochs_per_cv.values()) >= epoch:
                break
            time.sleep(10)
        if is_done:
            break

    logger.info("getting results...")
    pretrain_results = ray.get(reqs)
    for r_model, r_cv, r_dict in pretrain_results:
        logger.info(
            "model=%s cv=%d top1_train=%.4f top1_valid=%.4f"
            % (r_model, r_cv + 1, r_dict["top1_train"], r_dict["top1_valid"])
        )
    logger.info("processed in %.4f secs" % sw.pause("train_no_aug"))


def search(conf):
    sw = StopWatch.get()

    # region conf vars
    conf_dataset = conf["dataset"]
    dataroot = conf["dataroot"]
    redis_ip = conf["redis"]
    conf_loader = conf["autoaug"]["loader"]
    conf_model = conf["autoaug"]["model"]
    model_type = conf_model["type"]
    ds_name = conf_dataset["name"]
    aug = conf_loader["aug"]
    val_ratio = conf_loader["val_ratio"]
    epochs = conf_loader["epochs"]
    val_fold = conf_loader["val_fold"]
    cv_num = conf_loader["cv_num"]
    num_policy = conf["autoaug"]["num_policy"]
    num_op = conf["autoaug"]["num_op"]
    num_search = conf["autoaug"]["num_search"]
    num_result_per_cv = conf["autoaug"]["num_result_per_cv"]
    smoke_test = conf["smoke_test"]
    resume = conf["resume"]
    # endregion

    ray.init(
        redis_address=redis_ip,
        # allocate all GPUs on local node if cluster is not specified
        num_gpus=torch.cuda.device_count() if not redis_ip else None,
    )

    # first train with no aug
    _train_no_aug(conf)

    # get values from config
    num_samples = 4 if smoke_test else num_search

    logger.info("----- Search Test-Time Augmentation Policies -----")
    sw.start(tag="search")

    save_paths = [_get_model_filepath(ds_name, model_type, "ratio%.1f_fold%d" % (val_ratio, i)) for i in range(cv_num)]

    copied_c = copy.deepcopy(conf)
    ops = augment_list(False)
    space = {}
    for i in range(num_policy):
        for j in range(num_op):
            space["policy_%d_%d" % (i, j)] = hp.choice("policy_%d_%d" % (i, j), list(range(0, len(ops))))
            space["prob_%d_%d" % (i, j)] = hp.uniform("prob_%d_ %d" % (i, j), 0.0, 1.0)
            space["level_%d_%d" % (i, j)] = hp.uniform("level_%d_ %d" % (i, j), 0.0, 1.0)

    final_policy_set = []
    total_computation = 0
    reward_attr = "top1_valid"  # top1_valid or minus_loss
    for _ in range(1):  # run multiple times.
        for val_fold in range(cv_num):
            name = "search_%s_%s_fold%d_ratio%.1f" % (ds_name, model_type, val_fold, val_ratio)
            # logger.info(name)
            register_trainable(name, (lambda augs, rpt: _eval_tta(copy.deepcopy(copied_c), augs, rpt)))
            algo = HyperOptSearch(space, max_concurrent=4 * 20, reward_attr=reward_attr)

            exp_config = {
                name: {
                    "run": name,
                    "num_samples": num_samples,
                    "resources_per_trial": {"gpu": 1},
                    "stop": {"training_iteration": num_policy},
                    "config": {
                        "dataroot": dataroot,
                        "save_path": save_paths[val_fold],
                        "val_ratio": val_ratio,
                        "val_fold": val_fold,
                        "num_op": num_op,
                        "num_policy": num_policy,
                    },
                }
            }
            results = run_experiments(
                exp_config,
                search_alg=algo,
                scheduler=None,
                verbose=0,
                queue_trials=True,
                resume=resume,
                raise_on_failed_trial=False,
            )

            results = [x for x in results if x.last_result is not None]
            results = sorted(results, key=lambda x: x.last_result[reward_attr], reverse=True)

            # calculate computation usage
            for result in results:
                total_computation += result.last_result["elapsed_time"]

            for result in results[:num_result_per_cv]:
                final_policy = policy_decoder(result.config, num_policy, num_op)
                logger.info(
                    "loss=%.12f top1_valid=%.4f %s"
                    % (result.last_result["minus_loss"], result.last_result["top1_valid"], final_policy)
                )

                final_policy = remove_deplicates(final_policy)
                final_policy_set.extend(final_policy)

    logger.info(json.dumps(final_policy_set))
    logger.info("final_policy=%d" % len(final_policy_set))
    logger.info("processed in %.4f secs, gpu hours=%.4f" % (sw.pause("search"), total_computation / 3600.0))
    logger.info(
        "----- Train with Augmentations model=%s dataset=%s aug=%s ratio(test)=%.1f -----"
        % (model_type, ds_name, aug, val_ratio)
    )
    sw.start(tag="train_aug")

    num_experiments = 5
    default_path = [
        _get_model_filepath(ds_name, model_type, "ratio%.1f_default%d" % (val_ratio, _)) for _ in range(num_experiments)
    ]
    augment_path = [
        _get_model_filepath(ds_name, model_type, "ratio%.1f_augment%d" % (val_ratio, _)) for _ in range(num_experiments)
    ]
    reqs = [
        _train_model.remote(copy.deepcopy(copied_c), dataroot, aug, 0.0, 0, save_path=default_path[_], only_eval=True)
        for _ in range(num_experiments)
    ] + [
        _train_model.remote(copy.deepcopy(copied_c), dataroot, final_policy_set, 0.0, 0, save_path=augment_path[_])
        for _ in range(num_experiments)
    ]

    tqdm_epoch = tqdm(range(epochs))
    is_done = False
    for epoch in tqdm_epoch:
        while True:
            epochs = OrderedDict()
            for exp_idx in range(num_experiments):
                try:
                    if os.path.exists(default_path[exp_idx]):
                        latest_ckpt = torch.load(default_path[exp_idx])
                        epochs["default_exp%d" % (exp_idx + 1)] = latest_ckpt["epoch"]
                except:
                    pass
                try:
                    if os.path.exists(augment_path[exp_idx]):
                        latest_ckpt = torch.load(augment_path[exp_idx])
                        epochs["augment_exp%d" % (exp_idx + 1)] = latest_ckpt["epoch"]
                except:
                    pass

            tqdm_epoch.set_postfix(epochs)
            if len(epochs) == num_experiments * 2 and min(epochs.values()) >= epochs:
                is_done = True
            if len(epochs) == num_experiments * 2 and min(epochs.values()) >= epoch:
                break
            time.sleep(10)
        if is_done:
            break

    logger.info("getting results...")
    final_results = ray.get(reqs)

    for train_mode in ["default", "augment"]:
        avg = 0.0
        for _ in range(num_experiments):
            r_model, r_cv, r_dict = final_results.pop(0)
            logger.info("[%s] top1_train=%.4f top1_test=%.4f" % (train_mode, r_dict["top1_train"], r_dict["top1_test"]))
            avg += r_dict["top1_test"]
        avg /= num_experiments
        logger.info("[%s] top1_test average=%.4f (#experiments=%d)" % (train_mode, avg, num_experiments))
    logger.info("processed in %.4f secs" % sw.pause("train_aug"))

    logger.info(sw)


def _eval_tta(conf, augment, reporter):
    Config.set_inst(conf)

    # region conf vars
    conf_dataset = conf["dataset"]
    conf_loader = conf["autoaug"]["loader"]
    conf_model = conf["autoaug"]["model"]
    ds_name = conf_dataset["name"]
    cutout = conf_loader["cutout"]
    n_workers = conf_loader["n_workers"]
    # endregion

    val_ratio, val_fold, save_path = augment["val_ratio"], augment["val_fold"], augment["save_path"]

    # setup - provided augmentation rules
    aug = policy_decoder(augment, augment["num_policy"], augment["num_op"])

    # eval
    model = get_model(conf_model, num_class(ds_name))
    ckpt = torch.load(save_path)
    if "model" in ckpt:
        model.load_state_dict(ckpt["model"])
    else:
        model.load_state_dict(ckpt)
    model.eval()

    loaders = []
    for _ in range(augment["num_policy"]):
        tl, validloader, tl2 = get_dataloaders(
            augment["dataroot"],
            ds_name,
            aug,
            cutout,
            load_train=True,
            load_test=True,
            val_ratio=val_ratio,
            val_fold=val_fold,
            n_workers=n_workers,
        )
        loaders.append(iter(validloader))
        del tl, tl2  # TODO: why exclude validloader?

    start_t = time.time()
    metrics = Accumulator()
    loss_fn = torch.nn.CrossEntropyLoss(reduction="none")
    try:
        while True:
            losses = []
            corrects = []
            for loader in loaders:
                data, label = next(loader)
                data, label = data.cuda(), label.cuda()

                pred = model(data)

                loss = loss_fn(pred, label)
                losses.append(loss.detach().cpu().numpy())

                _, pred = pred.topk(1, 1, True, True)
                pred = pred.t()
                correct = pred.eq(label.view(1, -1).expand_as(pred)).detach().cpu().numpy()
                corrects.append(correct)
                del loss, correct, pred, data, label

            losses = np.concatenate(losses)
            losses_min = np.min(losses, axis=0).squeeze()

            corrects = np.concatenate(corrects)
            corrects_max = np.max(corrects, axis=0).squeeze()
            metrics.add_dict(
                {"minus_loss": -1 * np.sum(losses_min), "correct": np.sum(corrects_max), "cnt": len(corrects_max)}
            )
            del corrects, corrects_max
    except StopIteration:
        pass

    del model
    metrics = metrics / "cnt"
    gpu_secs = (time.time() - start_t) * torch.cuda.device_count()
    reporter(minus_loss=metrics["minus_loss"], top1_valid=metrics["correct"], elapsed_time=gpu_secs, done=True)
    return metrics["correct"]
