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

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

from mediatools.observer import PseudoObserver

from yolox.exp import Exp as MyExp
import torch

x_quantizer = FixedFakeQuantize.with_args(observer=MinMaxObserver, 
                                          **QuantizeScheme(
                                              symmetry=True, 
                                              per_channel=False, 
                                              pot_scale=False, 
                                              half_range=True,
                                              bit=8).to_observer_params())
fixed_quantizer = FixedFakeQuantize.with_args(observer=MinMaxObserver, 
                                              **QuantizeScheme(
                                                  symmetry=True, 
                                                  per_channel=False,  
                                                  pot_scale=False, 
                                                  half_range=False,
                                                  bit=8).to_observer_params())
custom_quantizer = LearnableFakeQuantize.with_args(observer=MinMaxObserver, 
                                                   **QuantizeScheme(
                                                       symmetry=True, 
                                                       per_channel=True, 
                                                       pot_scale=False, 
                                                       half_range=False,
                                                       bit=4).to_observer_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": {
        "catids": [1, 3, 17], # coco cats 0: person, 2:car, 16: dog
    },
    "postprocess_config":{
        "classes": [1, 2, 3],
        "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.AOFullA8W6,
    },
    "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,
        "head0.layer_8_obj.conv0.conv0.weight_fake_quant": custom_quantizer,
        "head0.layer_8_cls.conv0.conv0.weight_fake_quant": custom_quantizer,
        "head0.layer_16_obj.conv0.conv0.weight_fake_quant": custom_quantizer,
        "head0.layer_16_cls.conv0.conv0.weight_fake_quant": custom_quantizer,
        "head0.layer_32_obj.conv0.conv0.weight_fake_quant": custom_quantizer,
        "head0.layer_32_cls.conv0.conv0.weight_fake_quant": custom_quantizer,
        },
    "fixpoint_params": {
        'quantizer_params': 
            {
                'scale_bit': 15,
             },
        'bias_params':
            {
                'bias_bit': 32
            }
    },
    "exportor_params": {
        'output_mode': 0,
        'valid_classes': 3
    }
}


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.input_size = (416, 416)
        self.mosaic_scale = (0.5, 1.5)
        self.random_size = (10, 20)
        self.test_size = (416, 416)
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        self.enable_mixup = False
        self.data_num_workers = 2
        self.ema = False
        self.eval_interval = 10
        self.warmup_epochs = 5
        self.qat_warmup_epoch = 5
        self.qat_clib_epoch = 1
        self.max_epoch = 120
        self.no_aug_epochs = 60

        self.all_reduce_norm = True

    def get_model(self):

        model = YOLOXTrainer(**config['model_config'])
        model.apply(init_yolo)
        model.yolox.head0.initialized_params(1e-2)
        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

        return COCODataset_custom(
            **config['dataset_config'],
            data_dir=self.data_dir,
            json_file=self.train_ann,
            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=10240
        )

    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):
        testdev = kwargs.get("testdev", False)
        legacy = kwargs.get("legacy", False)

        return COCODataset_custom(
            **config['dataset_config'],
            data_dir=self.data_dir,
            json_file=self.val_ann if not testdev else self.test_ann,
            name="val2017" if not testdev else "test2017",
            img_size=self.test_size,
            preproc=ValTransform_yuv444(legacy=legacy),
            # max_imgs=200
        )