# -*- coding: utf-8 -*-
# This repo is licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import argparse
import bisect
import copy
import os
import time

import megengine as mge
import megengine.distributed as dist
from megengine.autodiff import GradManager
from megengine.data import DataLoader, Infinite, RandomSampler
from megengine.data import transform as T
from megengine.optimizer import AdamW, SGD
import math

from megengine.core._imperative_rt.utils import _set_defrag
_set_defrag(True)

from layers.tools.data_mapper import data_mapper
from layers.tools.utils import (
    AverageMeter,
    DetectionPadCollator,
    GroupedRandomSampler,
    get_config_info,
    import_from_file
)

logger = mge.get_logger(__name__)
logger.setLevel("INFO")
mge.device.set_prealloc_config(1024, 1024, 256 * 1024 * 1024, 4.0)


def make_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-f", "--file", default="net.py", type=str, help="net description file"
    )

    parser.add_argument(
        "-df", "--distill_file", default="distill_configs/ICD.py", type=str, help="distill description file"
    )

    parser.add_argument(
        "-tf", "--teacher_file", default="net.py", type=str, help="net description file"
    )

    parser.add_argument(
        "-w", "--weight_file", default=None, type=str, help="weights file",
    )

    parser.add_argument(
        "-tw", "--teacher_weight_file", default=None, type=str, help="weights file",
    )

    parser.add_argument(
        "-n", "--devices", default=1, type=int, help="total number of gpus for training",
    )
    parser.add_argument(
        "-b", "--batch_size", default=2, type=int, help="batch size for training",
    )
    parser.add_argument(
        "-d", "--dataset_dir", default="/data/datasets", type=str,
    )

    parser.add_argument(
        "-l", "--load_head", action='store_true'
    )

    parser.add_argument(
        "-sp", "--save_path", type=str,
    )


    return parser


def main():
    parser = make_parser()
    args = parser.parse_args()

    # ------------------------ begin training -------------------------- #
    logger.info("Device Count = %d", args.devices)

    from datetime import datetime
    now = datetime.now()

    log_dir = "log-of-ICD-{}-{}".format(os.path.basename(args.file).split(".")[0], now.strftime('%H:%M') + ('-Inherent' if args.load_head else ''))
    if not os.path.isdir(log_dir):
        os.makedirs(log_dir)

    args.save_path = log_dir

    if args.devices > 1:
        trainer = dist.launcher(worker, n_gpus=args.devices)
        trainer(args)
    else:
        worker(args)


def worker(args):
    mge.dtr.eviction_threshold = "6GB" # set the eviction threshold to 5 GB
    mge.dtr.enable()                   # enable the DTR optimization

    current_network = import_from_file(args.file)

    model = current_network.Net(current_network.Cfg())
    model.train()

    ########### BUILD TEACHER MODEL ############
    current_teacher = import_from_file(args.teacher_file)
    cfg_tea = current_teacher.Cfg()
    cfg_tea.backbone_pretrained = False
    model_tea = current_teacher.Net(cfg_tea)
    model_tea.train()
    # model_tea.training = True # run in train mod

    state_dict_tea = mge.load(args.teacher_weight_file)
    if "state_dict" in state_dict_tea:
        state_dict_tea = state_dict_tea["state_dict"]
    model_tea.load_state_dict(state_dict_tea)

    ############### LOADED ####################

    ################## DISTILLER ##############
    distiller_cfg = import_from_file(args.distill_file)
    distiller = distiller_cfg.Net()
    ############### END ####################

    if dist.get_rank() == 0:
        logger.info(get_config_info(model.cfg))
        logger.info(repr(model))

    params_with_grad = []
    for name, param in model.named_parameters():
        if "bottom_up.conv1" in name and model.cfg.backbone_freeze_at >= 1:
            continue
        if "bottom_up.layer1" in name and model.cfg.backbone_freeze_at >= 2:
            continue
        params_with_grad.append(param)

    opt = SGD(
        params_with_grad,
        lr=model.cfg.basic_lr * args.batch_size * dist.get_world_size(),
        momentum=model.cfg.momentum,
        weight_decay=model.cfg.weight_decay,
    )

    opt_d = AdamW(
        distiller.parameters(),
        lr=1e-4,
        weight_decay=1e-4,
    )
    params_with_grad.extend(distiller.parameters())
    gm = 
GradManager()