
import torch
import torch.nn.quantized as nnq
from torch.quantization import QuantStub, DeQuantStub
import torchvision
import unittest
import os
from neural_compressor.adaptor import FRAMEWORKS
from neural_compressor.model import MODELS
import neural_compressor.adaptor.pytorch as nc_torch
from neural_compressor.adaptor.pytorch import PyTorchVersionMode
from neural_compressor.experimental import Quantization, common
import shutil
import copy
import numpy as np

try:
    try:
        import intel_pytorch_extension as ipex
    except:
        import intel_extension_for_pytorch as ipex
    TEST_IPEX = True
except:
    TEST_IPEX = False

PT_VERSION = nc_torch.get_torch_version()
if PT_VERSION >= PyTorchVersionMode.PT18.value:
    FX_MODE = True
else:
    FX_MODE = False

torch.manual_seed(1)

fake_ptq_yaml = '''
    model:
      name: imagenet
      framework: pytorch

    evaluation:
      accuracy:
        metric:
          MSE:
            compare_label: False
      performance:
        warmup: 5
        iteration: 10

    tuning:
      accuracy_criterion:
        absolute:  100.0
        higher_is_better: False
      exit_policy:
        timeout: 0
      random_seed: 9527
      workspace:
        path: saved
    '''

fake_dynamic_yaml = '''
    model:
      name: imagenet
      framework: pytorch

    quantization:
      approach: post_training_dynamic_quant
    evaluation:
      accuracy:
        metric:
          MSE:
            compare_label: False
      performance:
        warmup: 5
        iteration: 10

    tuning:
      accuracy_criterion:
        absolute:  100.0
        higher_is_better: False
      exit_policy:
        timeout: 0
      random_seed: 9527
      workspace:
        path: saved
    '''


def build_ptq_yaml():
    with open('ptq_yaml.yaml', 'w', encoding="utf-8") as f:
        f.write(fake_ptq_yaml)


def build_dynamic_yaml():
    with open('dynamic_yaml.yaml', 'w', encoding="utf-8") as f:
        f.write(fake_dynamic_yaml)


def build_fx_ptq_yaml():
    fake_fx_ptq_yaml = fake_ptq_yaml.replace('pytorch', 'pytorch_fx')
    with open('fx_ptq_yaml.yaml', 'w', encoding="utf-8") as f:
        f.write(fake_fx_ptq_yaml)


def build_fx_dynamic_yaml():
    fake_fx_dynamic_yaml = fake_dynamic_yaml.replace('pytorch', 'pytorch_fx')
    with open('fx_dynamic_yaml.yaml', 'w', encoding="utf-8") as f:
        f.write(fake_fx_dynamic_yaml)


def build_ipex_yaml():
    fake_yaml = '''
        model:
          name: imagenet
          framework: pytorch_ipex

        evaluation:
          accuracy:
            metric:
              MSE:
                compare_label: False
          performance:
            warmup: 5
            iteration: 10

        tuning:
          accuracy_criterion:
            relative:  0.01
          exit_policy:
            timeout: 0
          random_seed: 9527
          workspace:
            path: saved
        '''
    with open('ipex_yaml.yaml', 'w', encoding="utf-8") as f:
        f.write(fake_yaml)


@unittest.skipIf(TEST_IPEX, "TODO: Please wait to IPEX + PyTorch1.7 release")
class TestPytorchAdaptor(unittest.TestCase):
    framework_specific_info = {"device": "cpu",
                               "approach": "post_training_static_quant",
                               "random_seed": 1234,
                               "q_dataloader": None,
                               "workspace_path": './'}
    framework = "pytorch"
    adaptor = FRAMEWORKS[framework](framework_specific_info)
    model = torchvision.models.quantization.resnet18()
    nc_model = MODELS['pytorch'](model)

    @classmethod
    def setUpClass(self):
        build_ptq_yaml()
        build_dynamic_yaml()

    @classmethod
    def tearDownClass(self):
        os.remove('ptq_yaml.yaml')
        os.remove('dynamic_yaml.yaml')
        shutil.rmtree('./saved', ignore_errors=True)
        shutil.rmtree('runs', ignore_errors=True)

    def test_quantization_saved(self):
        from neural_compressor.utils.pytorch import load

        for fake_yaml in ['dynamic_yaml.yaml', 'ptq_yaml.yaml']:
            if fake_yaml in ['dynamic_yaml.yaml']:
                model = torchvision.models.quantization.resnet18()
            else:
                model = copy.deepcopy(self.model)
            if fake_yaml in ['ptq_yaml.yaml']:
                model.eval().fuse_model()
            quantizer = Quantization(fake_yaml)
            dataset = quantizer.dataset('dummy', (100, 3, 256, 256), label=True)
            quantizer.model = model
            quantizer.calib_dataloader = common.DataLoader(dataset)
            quantizer.eval_dataloader = common.DataLoader(dataset)
            q_model = quantizer.fit()
        self.assertTrue(bool(q_model))


@unittest.skipIf(not FX_MODE, "Unsupport Fx Mode with PyTorch Version Below 1.8")
class TestPytorchFXAdaptor(unittest.TestCase):
    framework_specific_info = {"device": "cpu",
                               "approach": "post_training_static_quant",
                               "random_seed": 1234,
                               "q_dataloader": None,
                               "workspace_path": './'}
    framework = "pytorch_fx"
    adaptor = FRAMEWORKS[framework](framework_specific_info)
    model = torchvision.models.quantization.resnet18()
    nc_model = MODELS['pytorch_fx'](model)

    @classmethod
    def setUpClass(self):
        build_fx_ptq_yaml()
        build_fx_dynamic_yaml()

    @classmethod
    def tearDownClass(self):
        os.remove('fx_ptq_yaml.yaml')
        os.remove('fx_dynamic_yaml.yaml')
        shutil.rmtree('./saved', ignore_errors=True)
        shutil.rmtree('runs', ignore_errors=True)

    def test_fx_static_quantization_saved(self):
        from neural_compressor.utils.pytorch import load
        fake_yaml = 'fx_ptq_yaml.yaml'
        model = copy.deepcopy(self.model)
        model.eval().fuse_model()
        quantizer = Quantization(fake_yaml)
        dataset = quantizer.dataset('dummy', (100, 3, 256, 256), label=True)
        quantizer.model = model
        quantizer.calib_dataloader = common.DataLoader(dataset)
        quantizer.eval_dataloader = common.DataLoader(dataset)
        q_model = quantizer.fit()
        self.assertTrue(bool(q_model))
    
    @unittest.skipIf(PT_VERSION < PyTorchVersionMode.PT19.value,
      "Please use PyTroch 1.9 or higher version for dynamic quantization with pytorch_fx backend")
    def test_fx_dynamic_quantization_saved(self):
        from neural_compressor.utils.pytorch import load
        fake_yaml = 'fx_dynamic_yaml.yaml'
        model = torchvision.models.resnet18()
        quantizer = Quantization(fake_yaml)
        dataset = quantizer.dataset('dummy', (100, 3, 256, 256), label=True)
        quantizer.model = model
        quantizer.calib_dataloader = common.DataLoader(dataset)
        quantizer.eval_dataloader = common.DataLoader(dataset)
        q_model = quantizer.fit()
        self.assertTrue(bool(q_model))


@unittest.skipIf(not TEST_IPEX, "Unsupport Intel PyTorch Extension")
class TestPytorchIPEXAdaptor(unittest.TestCase):
    @classmethod
    def setUpClass(self):
        build_ipex_yaml()

    @classmethod
    def tearDownClass(self):
        os.remove('ipex_yaml.yaml')
        shutil.rmtree('./saved', ignore_errors=True)
        shutil.rmtree('runs', ignore_errors=True)
    def test_tuning_ipex(self):
        from neural_compressor.experimental import Quantization
        model = torchvision.models.resnet18()
        quantizer = Quantization('ipex_yaml.yaml')
        dataset = quantizer.dataset('dummy', (100, 3, 256, 256), label=True)
        quantizer.model = model
        quantizer.calib_dataloader = common.DataLoader(dataset)
        quantizer.eval_dataloader = common.DataLoader(dataset)
        nc_model = quantizer.fit()
        nc_model.save("./saved")
        try:
            script_model = torch.jit.script(model.to(ipex.DEVICE))
        except:
            script_model = torch.jit.trace(model.to(ipex.DEVICE), torch.randn(10, 3, 224, 224).to(ipex.DEVICE))
        from neural_compressor.experimental import Benchmark
        evaluator = Benchmark('ipex_yaml.yaml')
        evaluator.model = script_model
        evaluator.b_dataloader = common.DataLoader(dataset)
        results = evaluator()


if __name__ == "__main__":
    unittest.main()
