# -*- coding: utf-8 -*-
# MegEngine 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 os
import time
import pickle
import numpy as np
import megengine.amp as amp

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 SGD
import megengine.jit as jit

from tools.utils import (
    AverageMeter,
    DetectionPadCollator,
    GroupedRandomSampler,
    PseudoDetectionDataset,
    get_config_info,
    import_from_file
)
logger = mge.get_logger(__name__)
logger.setLevel("INFO")
mge.device.set_prealloc_config(1024, 1024, 512 * 1024 * 1024, 2.0)


def make_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-a", "--arch", type=str, help="model architecture",
    )
    parser.add_argument(
        "-n", "--ngpus", default=1, type=int, help="total number of gpus for training",
    )
    parser.add_argument(
        "-b", "--batch-size", default=2, type=int, help="batchsize for training",
    )
    parser.add_argument(
        "-s", "--steps", default=100, type=int, help="number of train steps (default: 100)",
    )
    parser.add_argument(
        "--trace",
        action='store_true',
        default=False,
        help="whether use trace or not (default: False)",
    )
    parser.add_argument(
        "--preloader",
        action='store_true',
        default=False,
        help="whether use preloader or not (default: False)",
    )
    parser.add_argument(
        "--symbolic",
        action='store_true',
        default=False,
        help="whether use symbolic trace or not (default: False)",
    )
    parser.add_argument(
        "-d", "--loader", default=False, action="store_true", help="use pseudo detection dataset loader",
    )
    parser.add_argument(
        "-p", "--print-freq", default=1, type=int, help="print frequency (default: 1)",
    )
    parser.add_argument(
        "-m",
        "--mode",
        default="normal",
        type=str,
        choices=["normal", "mp"],
        help="Quantization Mode\n"
        "normal: no quantization, using float32\n"
        "mp: input type is fp16",
    )
    parser.add_argument("--preload", default=False, action="store_true", help="whether use preload")
    return parser


def main():
    parser = make_parser()
    args = parser.parse_args()
    # ------------------------ begin training -------------------------- #
    if args.ngpus > 1:
        train_func = dist.launcher(worker, n_gpus=args.ngpus)
        train_func(args)
    else:
        worker(args)


def worker(args):
    config_file = {
        "faster_rcnn": "configs/faster_rcnn_res50_coco_1x_800size.py",
        "atss": "configs/atss_res50_coco_1x_800size.py",
        "retinanet": "configs/retinanet_res50_coco_1x_800size.py",
    }[args.arch]

    current_network = import_from_file(config_file)

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

    model.train()

    # 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,
        momentum=model.cfg.momentum,
        weight_decay=model.cfg.weight_decay * dist.get_world_size(),
    )

    gm = GradManager()
    if dist.get_world_size() > 1:
        gm.attach(
            params_with_grad,
            callbacks=[dist.make_allreduce_cb("SUM", dist.WORLD)]
        )
    else:
        gm.attach(params_with_grad)

    if dist.get_world_size() > 1:
        dist.bcast_list_(model.parameters(), dist.WORLD)  # sync parameters

    train_loader = None

    for epoch in range(1):
        train_one_epoch(model, train_loader, opt, gm, epoch, args)

def train_one_epoch(model, data_queue, opt, gm, epoch, args):
    @amp.autocast(enabled=args.mode == "mp")
    def train_func(image, im_info, gt_boxes):
        with gm:
            loss_dict = model(image=image, im_info=im_info, gt_boxes=gt_boxes)
            gm.backward(loss_dict["total_loss"])
            loss_list = list(loss_dict.values())
        opt.step().clear_grad()
        return loss_list

    if args.trace:
        if args.symbolic:
            train_func = jit.trace(train_func, symbolic=True, sublinear_memory_config=
jit.SublinearMemoryConfig(genetic_nr_iter=50)