import os
import torch.nn as nn 

from yolox_microbt.model.yolox_detector import YOLOXDetector, YOLOXTrainer
from yolox_microbt.model.postprocess.yolox_postprocess import YOLOXPostProcess
from yolox_microbt.data.datasets.coco import COCODataset_custom, TrainTransform_yuv444, ValTransform_yuv444
from yolox_microbt.utils.utils import init_yolo, _set_module, load_checkpoint

from mediatools.prepare_by_platform import BackendType, prepare_by_platform
from mediatools.observer import MinMaxObserver
from mediatools.fake_quantize import FixedFakeQuantize, LearnableFakeQuantize, PseudoFakeQuantize
from mediatools.scheme import QuantizeScheme

from mediatools.observer import PseudoObserver

from yolox.exp import Exp as MyExp

x_quantizer = FixedFakeQuantize.with_args(observer=PseudoObserver, 
                                          **QuantizeScheme(
                                              symmetry=True, 
                                              per_channel=False, 
                                              pot_scale=False, 
                                              half_range=True,
                                              bit=8).to_observer_params())
PseudoObserver_params = QuantizeScheme(symmetry=True, 
                                        per_channel=False,  
                                        pot_scale=False, 
                                        half_range=False,
                                        bit=8).to_observer_params()
PseudoObserver_params.update(
    {
        'limit_max': 8,
        'limit_min':-8,
    }
)
fixed_quantizer = PseudoFakeQuantize.with_args(observer=PseudoObserver, 
                                              **PseudoObserver_params)
config = {
    "model_type": "det",
    "model_config":{
        "backbone_params": {
            "arch_def":dict(steam=dict(layer0='cn_k3_s2_i3_o8'),
                            stage0=dict(layer0='ds_k3_s1_i8_o8'),
                            stage1=dict(layer0='ir_k3_s2_e4_i8_o8',
                                        layer1='ir_k3_s1_e4_i8_o8'),
                            stage2=dict(layer0='ir_k3_s2_e4_i8_o8',
                                        layer1='ir_k3_s1_e4_i8_o8'),
                            stage3=dict(layer0='ir_k3_s2_e4_i8_o16',
                                        layer1='ir_k3_s1_e4_i16_o16',
                                        layer2='ir_k3_s1_e4_i16_o16'),
                            stage4=dict(layer0='ir_k3_s1_e4_i16_o24_noskip'),
                            stage5=dict(layer0='ir_k3_s1_e4_i24_o024'),
                            stage6=dict(layer0='ir_g2_k3_s2_e4_i24_o40',
                                        layer1='ir_g2_k3_s1_e4_i40_o40'),
                            stage7=dict(layer0='ir_g2_k3_s1_e4_i40_o80_noskip'))
        },
        "head_params": {
            'num_classes': 3
        },
        "loss_params": {
            "num_classes": 3,
            "feature_strides": [8, 16, 32],
            "center_radius": 2.0
        }
    },
    "dataset_config": {
        "class_to_inds": {
            'person': 0,
            'car': 1,
            'dog': 2
        }, # coco cats 0: person, 2:car, 16: dog
    },
    "postprocess_config":{
        "classes": [3, 5, 8],
        "num_classes": 3,
        "strides": [8, 16, 32],
        "cal_max_score": True,
        'nms_params': {
            'nms_thresh': 0.35,
            'score_thresh': 0.01,
            'post_nms': 100,
        }
    },
    "quant_config": {
        "deploy_backend": BackendType.AOFullA8W4,
    },
    "custom_quant_config": {
        "x_post_act_fake_quantizer": x_quantizer,
        "head0_layer_8_obj_conv0_conv0_post_act_fake_quantizer": fixed_quantizer,
        "head0_layer_8_cls_conv0_conv0_post_act_fake_quantizer": fixed_quantizer,
        "head0_layer_16_obj_conv0_conv0_post_act_fake_quantizer": fixed_quantizer,
        "head0_layer_16_cls_conv0_conv0_post_act_fake_quantizer": fixed_quantizer,
        "head0_layer_32_obj_conv0_conv0_post_act_fake_quantizer": fixed_quantizer,
        "head0_layer_32_cls_conv0_conv0_post_act_fake_quantizer": fixed_quantizer,
        },
    "fixpoint_params": {
        'quantizer_params': 
            {
                'per_channel_shift': True,
                'scale_bit': 15,
             },
        'bias_params':
            {
                'bias_bit': 32
            }
    },
    'float_pretrain': 'docs/training_sample_data/sa6921_vnne_3classes_voc_float_600e_best_ckpt.pth'
}


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.input_size = (416, 416)
        self.test_size = (416, 416)
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        self.data_num_workers = 2
        self.num_classes = config['model_config']['head_params']['num_classes']

        # --------------- QAT config ----------------- #
        # QAT warmup epoch, keep float training epoch
        self.qat_warmup_epoch = 0
        # number of epochs for calibration model, after qat_warmup and qat_calib, QAT starts
        self.qat_clib_epoch = 2

        # --------------- transform config ----------------- #
        # prob of applying mosaic aug
        self.mosaic_prob = 1.0
        # prob of applying mixup aug
        self.mixup_prob = 0.5
        # prob of applying hsv aug
        self.hsv_prob = 0.5
        # prob of applying flip aug
        self.flip_prob = 0.5
        # rotation angle range, for example, if set to 2, the true range is (-2, 2)
        self.degrees = 10.0
        # translate range, for example, if set to 0.1, the true range is (-0.1, 0.1)
        self.translate = 0.1
        self.mosaic_scale = (0.1, 2)
        # apply mixup aug or not
        self.enable_mixup = False
        self.mixup_scale = (0.5, 1.5)
        # shear angle range, for example, if set to 2, the true range is (-2, 2)
        self.shear = 2.0

        # --------------  training config --------------------- #
        # epoch number used for warmup
        self.warmup_epochs = 10
        # max training epoch
        self.max_epoch = 300
        # minimum learning rate during warmup
        self.warmup_lr = 0
        self.min_lr_ratio = 0.05
        # learning rate for one image. During training, lr will multiply batchsize.
        self.basic_lr_per_img = 0.002 / 64.0
        # name of LRScheduler
        self.scheduler = "warmcos"
        # last #epoch to close augmention like mosaic
        self.no_aug_epochs = 200
        # apply EMA during training
        self.ema = False

        # weight decay of optimizer
        self.weight_decay = 5e-4
        # momentum of optimizer
        self.momentum = 0.9
        # log period in iter, for example,
        # if set to 1, user could see log every iteration.
        self.print_interval = 20
        # eval period in epoch, for example,
        # if set to 1, model will be evaluate after every epoch.
        self.eval_interval = 1
        # save history checkpoint or not.
        # If set to False, yolox will only save latest and best ckpt.
        self.save_history_ckpt = True
        self.all_reduce_norm = True
        # name of experiment
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

    def get_model(self):
        model = YOLOXTrainer(**config['model_config'])
        model.apply(init_yolo)
        model.yolox.head0.initialized_params(1e-2)
        load_checkpoint(model, config['float_pretrain'])
        model.yolox = prepare_by_platform(model.yolox, **config['quant_config'])
        if "custom_quant_config"in config and config["custom_quant_config"] is not None:
            custom_modules = list(config["custom_quant_config"].keys())
            module_names = dict(model.yolox.named_modules()).keys()
            for module_name, module in model.yolox.named_modules():
                if module_name in custom_modules:
                    _set_module(model.yolox, module_name, config["custom_quant_config"][module_name]())
        self.model = model
        self.model.train()
        return self.model
    
    def get_postprocess(self):
        return YOLOXPostProcess(**config['postprocess_config'])
    
    def get_dataset(self, cache: bool = False, cache_type: str = "ram"):
        """
        Get dataset according to cache and cache_type parameters.
        Args:
            cache (bool): Whether to cache imgs to ram or disk.
            cache_type (str, optional): Defaults to "ram".
                "ram" : Caching imgs to ram for fast training.
                "disk": Caching imgs to disk for fast training.
        """
        from yolox_microbt.data.datasets.coco import COCODataset_custom, TrainTransform_yuv444
        from yolox_microbt.data.datasets.voc import VOCDetection_custom
        from yolox.data import VOCDetection, get_yolox_datadir

        return VOCDetection_custom(
            data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
            class_to_ind=config['dataset_config']['class_to_inds'],
            image_sets=[('2007', 'train'), ('2012', 'train')],
            img_size=self.input_size,
            preproc=TrainTransform_yuv444(
                max_labels=50,
                flip_prob=self.flip_prob,
                hsv_prob=self.hsv_prob),
            cache=cache,
            cache_type=cache_type,
            # max_imgs=640,
        )

    def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img: str = None):
        """
        Get dataloader according to cache_img parameter.
        Args:
            no_aug (bool, optional): Whether to turn off mosaic data enhancement. Defaults to False.
            cache_img (str, optional): cache_img is equivalent to cache_type. Defaults to None.
                "ram" : Caching imgs to ram for fast training.
                "disk": Caching imgs to disk for fast training.
                None: Do not use cache, in this case cache_data is also None.
        """
        import torch.distributed as dist
        from yolox_microbt.data.datasets.coco import TrainTransform_yuv444
        from yolox.data import (
            TrainTransform,
            YoloBatchSampler,
            DataLoader,
            InfiniteSampler,
            MosaicDetection,
            worker_init_reset_seed,
        )
        from yolox.utils import wait_for_the_master


        # if cache is True, we will create self.dataset before launch
        # else we will create self.dataset after launch
        if self.dataset is None:
            with wait_for_the_master():
                assert cache_img is None, \
                    "cache_img must be None if you didn't create self.dataset before launch"
                self.dataset = self.get_dataset(cache=False, cache_type=cache_img)

        self.dataset = MosaicDetection(
            dataset=self.dataset,
            mosaic=not no_aug,
            img_size=self.input_size,
            preproc=TrainTransform_yuv444(
                max_labels=120,
                flip_prob=self.flip_prob,
                hsv_prob=self.hsv_prob),
            degrees=self.degrees,
            translate=self.translate,
            mosaic_scale=self.mosaic_scale,
            mixup_scale=self.mixup_scale,
            shear=self.shear,
            enable_mixup=self.enable_mixup,
            mosaic_prob=self.mosaic_prob,
            mixup_prob=self.mixup_prob,
        )

        if is_distributed:
            batch_size = batch_size // dist.get_world_size()

        sampler = InfiniteSampler(len(self.dataset), seed=self.seed if self.seed else 0)

        batch_sampler = YoloBatchSampler(
            sampler=sampler,
            batch_size=batch_size,
            drop_last=False,
            mosaic=not no_aug,
        )

        dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True}
        dataloader_kwargs["batch_sampler"] = batch_sampler

        # Make sure each process has different random seed, especially for 'fork' method.
        # Check https://github.com/pytorch/pytorch/issues/63311 for more details.
        dataloader_kwargs["worker_init_fn"] = worker_init_reset_seed

        train_loader = DataLoader(self.dataset, **dataloader_kwargs)

        return train_loader

    def get_trainer(self, args):
        from yolox_microbt.core.trainer import Trainer_custom
        trainer = Trainer_custom(self, args)
        # NOTE: trainer shouldn't be an attribute of exp object
        return trainer
    
    def get_eval_dataset(self, **kwargs):
        from yolox.data import VOCDetection, get_yolox_datadir
        from yolox_microbt.data.datasets.voc import VOCDetection_custom
        testdev = kwargs.get("testdev", False)
        legacy = kwargs.get("legacy", False)

        return VOCDetection_custom(
            data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
            class_to_ind=config['dataset_config']['class_to_inds'],
            image_sets=[('2007', 'val')],
            img_size=self.test_size,
            preproc=ValTransform_yuv444(legacy=legacy),
            # max_imgs=200
        )
    
    def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
        from yolox.evaluators import VOCEvaluator

        return VOCEvaluator(
            dataloader=self.get_eval_loader(batch_size, is_distributed,
                                            testdev=testdev, legacy=legacy),
            img_size=self.test_size,
            confthre=self.test_conf,
            nmsthre=self.nmsthre,
            num_classes=self.num_classes,
        )