import unittest

import torch
from model_zoo import CIFAR, VGG16Cifar10
from predtuner import TorchApp, accuracy, config_pylogger, get_knobs_from_file
from torch.nn import Conv2d, Linear
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.dataset import Subset

msg_logger = config_pylogger(output_dir="/tmp", verbose=True)


class TestTorchApp(unittest.TestCase):
    def setUp(self):
        dataset = CIFAR.from_file(
            "model_data/cifar10/input.bin", "model_data/cifar10/labels.bin"
        )
        self.dataset = Subset(dataset, range(100))
        self.module = VGG16Cifar10()
        self.module.load_state_dict(torch.load("model_data/vgg16_cifar10.pth.tar"))

    def get_app(self):
        return TorchApp(
            "TestTorchApp",
            self.module,
            DataLoader(self.dataset, batch_size=500),
            DataLoader(self.dataset, batch_size=500),
            get_knobs_from_file(),
            accuracy,
        )

    def test_init(self):
        app = self.get_app()
        n_knobs = {op: len(ks) for op, ks in app.op_knobs.items()}
        self.assertEqual(len(n_knobs), 34)
        for op_name, op in app.midx.name_to_module.items():
            if isinstance(op, Conv2d):
                nknob = 56
            elif isinstance(op, Linear):
                nknob = 2
            else:
                nknob = 1
            self.assertEqual(n_knobs[op_name], nknob)

    def test_baseline_qos(self):
        app = self.get_app()
        qos, _ = app.measure_qos_perf({}, False)
        self.assertAlmostEqual(qos, 88.0)

    def test_tuning(self):
        app = self.get_app()
        baseline, _ = app.measure_qos_perf({}, False)
        tuner = app.get_tuner()
        tuner.tune(100, baseline - 3.0)
        configs = tuner.get_all_configs()
        for conf in configs:
            self.assertTrue(conf.qos > baseline - 3.0)