# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 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.
"""Finetune a pretrained fp32 with int8 post train quantization(calibration)"""
import argparse
import collections
import numbers
import os
import time

# pylint: disable=import-error
import models

import megengine as mge
import megengine.data as data
import megengine.data.transform as T
import megengine.distributed as dist
import megengine.functional as F
import megengine.quantization as Q
from megengine.quantization.quantize import enable_observer, quantize, quantize_qat

logger = mge.get_logger(__name__)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--arch", default="resnet18", type=str)
    parser.add_argument("-d", "--data", default=None, type=str)
    parser.add_argument("-s", "--save", default="/data/models", type=str)
    parser.add_argument(
        "-c",
        "--checkpoint",
        default=None,
        type=str,
        help="pretrained model to finetune",
    )

    parser.add_argument("-n", "--ngpus", default=None, type=int)
    parser.add_argument("-w", "--workers", default=4, type=int)
    parser.add_argument("--report-freq", default=50, type=int)
    args = parser.parse_args()

    world_size = (
        dist.helper.get_device_count_by_fork("gpu")
        if args.ngpus is None
        else args.ngpus
    )
    world_size = 1 if world_size == 0 else world_size
    if world_size != 1:
        logger.warning(
            "Calibration only supports single GPU now, %d provided", world_size
        )
    proc_func = dist.launcher(worker) if world_size > 1 else worker
    proc_func(world_size, args)


def get_parameters(model, cfg):
    if isinstance(cfg.WEIGHT_DECAY, numbers.Number):
        return {
            "params": model.parameters(requires_grad=True),
            "weight_decay": cfg.WEIGHT_DECAY,
        }

    groups = collections.defaultdict(list)  # weight_decay -> List[param]
    for pname, p in model.named_parameters(requires_grad=True):
        wd = cfg.WEIGHT_DECAY(pname, p)
        groups[wd].append(p)
    groups = [
        {"params": params, "weight_decay": wd} for wd, params in groups.items()
    ]  # List[{param, weight_decay}]
    return groups


def worker(world_size, args):
    # pylint: disable=too-many-statements

    rank = dist.get_rank()
    if world_size > 1:
        # Initialize distributed process group
        logger.info("init distributed process group {} / {}".format(rank, world_size))

    save_dir = os.path.join(args.save, args.arch + "." + "calibration")
    if not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)
    mge.set_log_file(os.path.join(save_dir, "log.txt"))

    model = models.__dict__[args.arch]()

    # load calibration model
    assert args.checkpoint
    logger.info("Load pretrained weights from %s", args.checkpoint)
    ckpt = mge.load(args.checkpoint)
    ckpt = ckpt["state_dict"] if "state_dict" in ckpt else ckpt
    model.load_state_dict(ckpt, strict=False)

    # Build valid datasets
    valid_dataset = data.dataset.ImageNet(args.data, train=False)
    valid_sampler = data.SequentialSampler(
        valid_dataset, batch_size=100, drop_last=False
    )
    valid_queue = data.DataLoader(
        valid_dataset,
        sampler=valid_sampler,
        transform=T.Compose(
            [T.Resize(256), T.CenterCrop(224), T.Normalize(mean=128), 
T.ToMode("CHW")