import sys, time, random, argparse
from copy import deepcopy
import torch
from pathlib import Path
import numpy as np
import torch.nn as nn
import math
import os
import torch.nn.init as init
import torch.nn.functional as F
import warnings
warnings.filterwarnings("ignore")

lib_dir = (Path(__file__).parent / ".." / ".." / "lib").resolve()
if str(lib_dir) not in sys.path:
    sys.path.insert(0, str(lib_dir))
from config_utils import load_config, dict2config
from datasets import get_datasets, get_nas_search_loaders
from procedures import (
    prepare_seed,
    prepare_logger,
    save_checkpoint,
    copy_checkpoint,
    get_optim_scheduler,
)
from utils import get_model_infos, obtain_accuracy
from log_utils import AverageMeter, time_string, convert_secs2time
from models import get_cell_based_tiny_net, get_search_spaces
from nas_201_api import NASBench201API as API
from models import create_cnn_model, count_parameters_in_MB


def search_func_modified(xloader,
    teacher_model,
    network,
    student_model,
    criterion,
    scheduler,
    w_optimizer,
    a_optimizer,
    epoch_str,
    print_freq,
    logger,
):
    teacher_model.eval()
    student_model.eval()
    data_time, batch_time = AverageMeter(), AverageMeter()
    base_losses, base_top1, base_top5 = AverageMeter(), AverageMeter(), AverageMeter()
    arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(), AverageMeter()
    end = time.time()
    T = 5
    lambda_ = 0.5

    for step, (base_inputs, base_targets, arch_inputs, arch_targets) in enumerate(
            xloader
    ):
        scheduler.update(None, 1.0 * step / len(xloader))
        base_targets = base_targets.cuda(non_blocking=True)
        arch_targets = arch_targets.cuda(non_blocking=True)
        # measure data loading time
        data_time.update(time.time() - end)

        network.train()
        w_optimizer.zero_grad()
        _, logits = network(base_inputs.cuda())
        logits_teacher = teacher_model(base_inputs.cuda())
        logits_student = student_model(base_inputs.cuda())
        loss_KD_TA = T * T * nn.KLDivLoss()(F.log_softmax(logits / T, dim=1),
                                            F.softmax(logits_teacher / T, dim=1))
        loss_KD_TA+= T * T * nn.KLDivLoss()(F.log_softmax(logits / T, dim=1),
                                            F.softmax(logits_student / T, dim=1))

        base_loss = (1 - lambda_) * criterion(logits, base_targets) + lambda_ * loss_KD_TA
        base_loss.backward()
        torch.nn.utils.clip_grad_norm_(network.parameters(), 5)
        w_optimizer.step()

        # record
        base_prec1, base_prec5 = obtain_accuracy(
            logits.data, base_targets.data, topk=(1, 5)
        )
        base_losses.update(base_loss.item(), base_inputs.size(0))
        base_top1.update(base_prec1.item(), base_inputs.size(0))
        base_top5.update(base_prec5.item(), base_inputs.size(0))
        # update the architecture-weight
        a_optimizer.zero_grad()
        _, logits = network(arch_inputs.cuda())

        logits_teacher = teacher_model(arch_inputs.cuda())
        logits_student = student_model(arch_inputs.cuda())

        loss_KD_TA = T * T * nn.KLDivLoss()(F.log_softmax(logits / T, dim=1),
                                            F.softmax(logits_teacher / T, dim=1))
        loss_KD_TA += T * T * nn.KLDivLoss()(F.log_softmax(logits / T, dim=1),
                                             F.softmax(logits_student / T, dim=1))
        arch_loss = (1 - lambda_) * criterion(logits, arch_targets) + lambda_ * loss_KD_TA
        arch_loss.backward()
        a_optimizer.step()
        # record
        arch_prec1, arch_prec5 = obtain_accuracy(
            logits.data, arch_targets.data, topk=(1, 5)
        )
        arch_losses.update(arch_loss.item(), arch_inputs.size(0))
        arch_top1.update(arch_prec1.item(), arch_inputs.size(0))
        arch_top5.update(arch_prec5.item(), arch_inputs.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if step % print_freq == 0 or step + 1 == len(xloader):
            Sstr = (
                    "*SEARCH* "
                    + time_string()
                    + " [{:}][{:03d}/{:03d}]".format(epoch_str, step, len(xloader))
            )
            Tstr = "Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})".format(
                batch_time=batch_time, data_time=data_time
            )
            Wstr = "Base [Loss {loss.val:.3f} ({loss.avg:.3f})  Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
                loss=base_losses, top1=base_top1, top5=base_top5
            )
            Astr = "Arch [Loss {loss.val:.3f} ({loss.avg:.3f})  Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Prec@5 {top5.val:.2f} ({top5.avg:.2f})]".format(
                loss=arch_losses, top1=arch_top1, top5=arch_top5
            )

            logger.log(Sstr + " " + Tstr + " " + Wstr + " " + Astr)

    return base_losses.avg, base_top1.avg, base_top5.avg, arch_losses.avg, arch_top1.avg, arch_top5.avg

def main(xargs):
    assert torch.cuda.is_available(), "CUDA is not available."
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.set_num_threads(xargs.workers)
    prepare_seed(xargs.rand_seed)
    logger = prepare_logger(args)

    train_data, valid_data, xshape, class_num = get_datasets(
        xargs.dataset, xargs.data_path, -1
    )
    # config_path = 'configs/nas-benchmark/algos/GDAS.config'
    config = load_config(
        xargs.config_path, {"class_num": class_num, "xshape": xshape}, logger
    )
    search_loader, _, valid_loader = get_nas_search_loaders(
        train_data,
        valid_data,
        xargs.dataset,
        "../../configs/nas-benchmark",
        config.batch_size,
        xargs.workers,
    )
    teacher_model = xargs.teacher_model
    TA = xargs.TA
    student = xargs.student_model
    teacher_checkpoint = xargs.teacher_checkpoint
    student_checkpoint = xargs.student_checkpoint
    epoch_online = xargs.epoch_online
    logger.log(
        "||||||| {:10s} ||||||| Search-Loader-Num={:}, batch size={:}".format(
            xargs.dataset, len(search_loader), config.batch_size
        )
    )
    logger.log("||||||| {:10s} ||||||| Config={:}".format(xargs.dataset, config))

    search_space = get_search_spaces("cell", xargs.search_space_name)
    if xargs.model_config is None:
        model_config = dict2config(
            {
                "name": "GDAS",
                "C": xargs.channel,
                "N": xargs.num_cells,
                "max_nodes": xargs.max_nodes,
                "num_classes": class_num,
                "space": search_space,
                "affine": False,
                "track_running_stats": bool(xargs.track_running_stats),
            },
            None,
        )
    else:
        model_config = load_config(
            xargs.model_config,
            {
                "num_classes": class_num,
                "space": search_space,
                "affine": False,
                "track_running_stats": bool(xargs.track_running_stats),
            },
            None,
        )
    search_model = get_cell_based_tiny_net(model_config)
    logger.log("search-model :\n{:}".format(search_model))
    logger.log("model-config : {:}".format(model_config))

    student_accuracy = {'best': -1}
    TA_accuracy = {'best': -1}

    w_optimizer, w_scheduler, criterion = get_optim_scheduler(
        search_model.get_weights(), config
    )
    a_optimizer = torch.optim.Adam(
        search_model.get_alphas(),
        lr=xargs.arch_learning_rate,
        betas=(0.5, 0.999),
        weight_decay=xargs.arch_weight_decay,
    )
    logger.log("w-optimizer : {:}".format(w_optimizer))
    logger.log("a-optimizer : {:}".format(a_optimizer))
    logger.log("w-scheduler : {:}".format(w_scheduler))
    logger.log("criterion   : {:}".format(criterion))
    flop, param = get_model_infos(search_model, xshape)
    logger.log("FLOP = {:.2f} M, Params = {:.2f} MB".format(flop, param))
    logger.log("search-space [{:} ops] : {:}".format(len(search_space), search_space))
    if xargs.arch_nas_dataset is None:
        api = None
    else:
        api = API(xargs.arch_nas_dataset)
    logger.log("{:} create API = {:} done".format(time_string(), api))

    last_info, model_base_path, model_best_path = (
        logger.path("info"),
        logger.path("model"),
        logger.path("best"),
    )
    network, criterion = torch.nn.DataParallel(search_model).cuda(), criterion.cuda()

    if last_info.exists():  # automatically resume from previous checkpoint
        logger.log(
            "=> loading checkpoint of the last-info '{:}' start".format(last_info)
        )
        last_info = torch.load(last_info)
        start_epoch = last_info["epoch"]
        checkpoint = torch.load(last_info["last_checkpoint"])
        genotypes = checkpoint["genotypes"]
        valid_accuracies = checkpoint["valid_accuracies"]
        search_model.load_state_dict(checkpoint["search_model"])
        w_scheduler.load_state_dict(checkpoint["w_scheduler"])
        w_optimizer.load_state_dict(checkpoint["w_optimizer"])
        a_optimizer.load_state_dict(checkpoint["a_optimizer"])
        logger.log(
            "=> loading checkpoint of the last-info '{:}' start with {:}-th epoch.".format(
                last_info, start_epoch
            )
        )
    else:
        logger.log("=> do not find the last-info file : {:}".format(last_info))
        start_epoch, valid_accuracies, genotypes = (
            0,
            {"best": -1},
            {-1: search_model.genotype()},
        )

    # start training
    start_time, search_time, epoch_time, total_epoch = (
        time.time(),
        AverageMeter(),
        AverageMeter(),
        config.epochs + config.warmup,
    )

    teacher_model, teacher_optimizer, teacher_scheduler = create_cnn_model(teacher_model, train_data, total_epoch, teacher_checkpoint, use_cuda=1)
    # if teacher_checkpoint:
    # teacher_model = load_checkpoint(teacher_model, teacher_checkpoint)
    # else:
    #     teacher_model = train_teacher(search_loader, teacher_model, criterion, teacher_optimizer, logger, total_epoch, xargs.teacher_model)

    # if TA:
    student_model, student_optimizer, student_scheduler = create_cnn_model(student, train_data, total_epoch, student_checkpoint, use_cuda=1)
    # if student_checkpoint:
    # student_model = load_checkpoint(student_model, student_checkpoint)
    # checkpoint = torch.load(student_checkpoint)
    # student_optimizer.load_state_dict(checkpoint['optimizer'])

    # if TA != 'GDAS':
    #     network, w_optimizer = create_cnn_model(TA, train_data, use_cuda=1)
    #     # w_optimizer = torch.optim.Adam(network .parameters(), lr=0.025, betas=(0.5, 0.999), weight_decay=1e-4)
    #     a_optimizer = None

    for epoch in range(start_epoch, total_epoch):
        w_scheduler.update(epoch, 0.0)
        need_time = "Time Left: {:}".format(
            convert_secs2time(epoch_time.val * (total_epoch - epoch), True)
        )
        epoch_str = "{:03d}-{:03d}".format(epoch, total_epoch)
        search_model.set_tau(
            xargs.tau_max - (xargs.tau_max - xargs.tau_min) * epoch / (total_epoch - 1)
        )
        logger.log(
            "\n[Search the {:}-th epoch] {:}, tau={:}, LR={:}".format(
                epoch_str, need_time, search_model.get_tau(), min(w_scheduler.get_lr())
            )
        )

        search_w_loss, search_w_top1, search_w_top5, valid_a_loss, valid_a_top1, valid_a_top5 = search_func_modified(
            search_loader,
            teacher_model,
            network,
            student_model,
            criterion,
            w_scheduler,
            w_optimizer,
            a_optimizer,
            epoch_str,
            xargs.print_freq,
            logger,
            )
        # else:
        #     Student_optimizer = None
        #     training_mode = 0
        #     search_w_loss, search_w_top1, search_w_top5, valid_a_loss, valid_a_top1, valid_a_top5, student_loss, student_top1, student_top5 = search_func(search_loader,
        #                                                                                                                                                     teacher_model,
        #                                                                                                                                                     network,
        #                                                                                                                                                     student_model,
        #                                                                                                                                                     criterion,
        #                                                                                                                                                     w_scheduler,
        #                                                                                                                                                     w_optimizer,
        #                                                                                                                                                     a_optimizer,
        #                                                                                                                                                     Student_optimizer,

        search_time.update(time.time() - start_time)
        logger.log(
            "[{:}] searching : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%, time-cost={:.1f} s".format(
                epoch_str, search_w_loss, search_w_top1, search_w_top5, search_time.sum
            )
        )
        logger.log(
            "[{:}] evaluate  : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%".format(
                epoch_str, valid_a_loss, valid_a_top1, valid_a_top5
            )
        )

        # logger.log(
        #     "[{:}] student  : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%".format(
        #         epoch_str, student_loss, student_top1, student_top5
        #     )
        # )

        # check the best accuracy

        # student_accuracy[epoch] = student_top1
        # if student_top1 > student_accuracy['best']:
        #     student_accuracy['best'] = student_top1

        TA_accuracy[epoch] = search_w_top1
        if search_w_top1 > TA_accuracy['best']:
            TA_accuracy['best'] = search_w_top1

        valid_accuracies[epoch] = valid_a_top1
        if valid_a_top1 > valid_accuracies["best"]:
            valid_accuracies["best"] = valid_a_top1
            genotypes["best"] = search_model.genotype()
            find_best = True
        else:
            find_best = False

        genotypes[epoch] = search_model.genotype()
        logger.log(
            "<<<--->>> The {:}-th epoch : {:}".format(epoch_str, genotypes[epoch])
        )
        # save checkpoint

        save_path = save_checkpoint(
            {
                "epoch": epoch + 1,
                "args": deepcopy(xargs),
                "search_model": search_model.state_dict(),
                "w_optimizer": w_optimizer.state_dict(),
                "a_optimizer": a_optimizer.state_dict(),
                "w_scheduler": w_scheduler.state_dict(),
                "genotypes": genotypes,
                "valid_accuracies": valid_accuracies,
            },
            model_base_path,
            logger,
        )
        last_info = save_checkpoint(
            {
                "epoch": epoch + 1,
                "args": deepcopy(args),
                "last_checkpoint": save_path,
            },
            logger.path("info"),
            logger,
        )
        if find_best:
            logger.log(
                "<<<--->>> The {:}-th epoch : find the highest validation accuracy : {:.2f}%.".format(
                    epoch_str, valid_a_top1
                )
            )
            copy_checkpoint(model_base_path, model_best_path, logger)
        with torch.no_grad():
            logger.log("{:}".format(search_model.show_alphas()))
        if api is not None:
            logger.log("{:}".format(api.query_by_arch(genotypes[epoch], "200")))
            # measure elapsed time
        epoch_time.update(time.time() - start_time)
        start_time = time.time()

        # if TA!='GDAS':
        #     student_model, student_optimizer = create_cnn_model(student, train_data, use_cuda=1)
        #
        # student_best = -1
        # for epoch in range(start_epoch, total_epoch):
        #     student_loss, student_top1, student_top5 = train_student(search_loader,
        #         network,
        #         student_model,
        #         criterion,
        #         student_optimizer,
        #         epoch_str,
        #         xargs.print_freq,
        #         logger,)
        #
        #     student_accuracy[epoch] = student_top1
        #     if student_top1 > student_accuracy['best']:
        #         student_accuracy['best'] = student_top1



        logger.log("\n" + "-" * 100)
        # check the performance from the architecture dataset
        logger.log(
            "GDAS : run {:} epochs, cost {:.1f} s, last-geno is {:}.".format(
                total_epoch, search_time.sum, genotypes[epoch - 1]
            )
        )
        if api is not None:
            logger.log("{:}".format(api.query_by_arch(genotypes[epoch - 1], "200")))


    logger.log('----------------')

    logger.log('we used {:} as our Teacher with param size {:}'.format(xargs.teacher_model, count_parameters_in_MB(teacher_model)))
    logger.log('we used {:} as our TA with param size {:}'.format(TA, count_parameters_in_MB(network)))
    logger.log('we used {:} as our Student with param size {:}'.format(xargs.student_model, count_parameters_in_MB(student_model)))

    logger.log('we used {:} online epochs out of total epochs of {:}'.format(xargs.epoch_online, total_epoch))
    logger.log('The best ACC of TA: {:.2f}%'.format(TA_accuracy['best']))
    logger.log('The best ACC of Student: {:.2f}%'.format(student_accuracy['best']))
    logger.log('----------------')

    logger.close()

    # else:
    #     if student:
    #         student_model, student_optimizer = create_cnn_model(student, train_data, use_cuda=1)
    #         student_best = -1
    #         for epoch in range(start_epoch, total_epoch):
    #             epoch_str = "{:03d}-{:03d}".format(epoch, total_epoch)
    #             student_loss, student_top1, student_top5 = train_student(search_loader,
    #                                                                      teacher_model,
    #                                                                      student_model,
    #                                                                      criterion,
    #                                                                      student_optimizer,
    #                                                                      epoch_str,
    #                                                                      xargs.print_freq,
    #                                                                      logger, )
    #
    #             student_accuracy[epoch] = student_top1
    #             if student_top1 > student_accuracy['best']:
    #                 student_accuracy['best'] = student_top1
    #
    #         logger.log('----------------')
    #         logger.log('we used {:} as our Teacher with param size {:}'.format(xargs.teacher_model, count_parameters_in_MB(teacher_model)))
    #         logger.log('we used {:} as our TA with param size {:}'.format(TA, count_parameters_in_MB(network)))
    #         logger.log('we used {:} as our Student with param size {:}'.format(xargs.student_model, count_parameters_in_MB(student_model)))
    #         logger.log('we used {:} online epochs out of total epochs of {:}'.format(xargs.epoch_online, total_epoch))
    #         logger.log('The best ACC of  : {:.2f}%'.format(TA_accuracy['best']))
    #         logger.log('The best ACC of Student: {:.2f}%'.format(student_accuracy['best']))
    #         logger.log('----------------')
    #         logger.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser("GDAS")
    parser.add_argument("--data_path", default='../../data', type=str, help="Path to dataset")
    parser.add_argument(
        "--dataset",
        type=str,
        default= 'cifar100',
        choices=["cifar10", "cifar100", "ImageNet16-120"],
        help="Choose between Cifar10/100 and ImageNet-16.",
    )
    # channels and number-of-cells
    parser.add_argument("--search_space_name", default='darts', type=str, help="The search space name.")
    parser.add_argument("--max_nodes", type=int, help="The maximum number of nodes.")
    parser.add_argument("--channel", type=int, help="The number"
                                                    " of channels.")
    parser.add_argument(
        "--num_cells", type=int, help="The number of cells in one "
                                      "。stage."
    )
    parser.add_argument(
        "--track_running_stats",
        type=int,
        default=1,
        choices=[0, 1],
        help="Whether use track_running_stats or not in the BN layer.",
    )
    parser.add_argument(
        "--config_path", default='../../configs/search-opts/GDAS-NASNet-CIFAR.config', type=str, help="The path of the configuration."
    )
    parser.add_argument(
        "--model_config",
        default='../../configs/search-archs/GDAS-NASNet-CIFAR.config',
        type=str,
        help="The path of the model configuration. When this arg is set, it will cover max_nodes / channels / num_cells.",
    )
    # architecture leraning rate
    parser.add_argument(
        "--arch_learning_rate",
        type=float,
        default=3e-4,
        help="learning rate for arch encoding",
    )
    parser.add_argument(
        "--arch_weight_decay",
        type=float,
        default=1e-3,
        help="weight decay for arch encoding",
    )
    parser.add_argument("--tau_min", default=10, type=float, help="The minimum tau for Gumbel")
    parser.add_argument("--tau_max", default=0.1, type=float, help="The maximum tau for Gumbel")
    # log
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="number of data loading workers (default: 2)",
    )
    parser.add_argument(
        "--save_dir", default='./output/search-cell-dar/GDAS-cifar10-BN1', type=str, help="Folder to save checkpoints and log."
    )
    parser.add_argument(


        "--arch_nas_dataset",
        type=str,


        help="The path to load the architecture dataset   (tiny-nas-benchmark).",
    )
    parser.add_argument("--print_freq", default=200, type=int, help="print frequency (default: 200)")
    parser.add_argument("--rand_seed", default= -1, type=int, help="manual seed")
    parser.add_argument("--teacher_model", default="vgg19", type=str, help="type of teacher mode")
    parser.add_argument("--TA", default='GDAS', type=str, help="type of TA")
    parser.add_argument("--student_model", default='lenet', type=str, help="type of student mode")
    parser.add_argument("--teacher_checkpoint", default='../output/nas-infer/cifar10-BS96-gdas_serached/checkpoint/seed-71956-bestNone_vgg19_68.08%_07-19,06.pth', type=str, help="teacher mode's check point")
    parser.add_argument("--student_checkpoint", default='../output/nas-infer/cifar10-BS96-gdas_serached/checkpoint/seed-62735-bestNone_lenet_33.10%_07-20,01.pth', type=str,
                        help="student mode's check point")
    parser.add_argument("--epoch_online", default=250, type=int, help="online training of TA and student")
    args = parser.parse_args()
    if args.rand_seed is None or args.rand_seed < 0:
        args.rand_seed = 99999
    # #
    # #
    # teacher_models = ['resnet110', 'resnet56', 'resnet44', 'resnet32', 'resnet26', 'resnet20', 'resnet14', 'resnet8', 'plane10', 'plane8', 'plane  6','plane4','plane2' ]
    # teacher_models = ['vgg16']
    # for one in teacher_models:
    #     args.teacher_model = one

    # TA_models = ['plane4', 'plane6', 'resnet26', 'resnet20']
    # for one in TA_models:
    #     args.TA = one
    main(args)