import unittest

from useckit.v2.dataset import Dataset
from useckit.v2.keras_model_builder import ModelBuilder
from useckit.v2.evaluator import Paradigm
from useckit.v2.model_descriptions import MODEL_DESCRIPTIONS
from useckit.v2.utils.example_dataset import get_chi2021_dataset


class TestKerasModelBuilder(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Use feature vectors so that sklearn models can be trained easily
        cls.ds: Dataset = get_chi2021_dataset(feature_vector=True)

    def _build_validate_train(self, p: Paradigm, model_desc_key: str):
        print("Testing paradigm", p, "/", model_desc_key)

        mb = ModelBuilder(MODEL_DESCRIPTIONS)

        input_shape = self.ds.get_shape()
        num_classes = self.ds.get_unique_labels().shape[0]

        # Build steps
        if p == Paradigm.CSC:
            build_result = mb.build_multiclass_classification(
                arch_key=model_desc_key,
                input_shape=input_shape,
                num_classes=num_classes,
            )
        elif p == Paradigm.DL:
            build_result1 = mb.build_metriclearning_triplet(model_desc_key, input_shape=input_shape)
            build_result2 = mb.build_metriclearning_contrastive(model_desc_key, input_shape=input_shape)
        elif p == Paradigm.BV:
            build_result = mb.build_binary_verification_classifier(model_desc_key, input_shape=input_shape)
        elif p == Paradigm.AD:
            build_result = mb.build_autoencoder_oneclass(model_desc_key, input_shape=input_shape)
        else:
            raise ValueError(f"Unknown paradigm: {p}")

        # Validation Steps
        if p == Paradigm.CSC:
            build_result["model"].summary()
            assert build_result["model"].layers[-1].units == len(self.ds.get_unique_enrollment_labels()), \
                f"Expected {len(self.ds.get_unique_enrollment_labels())} units in the last layer, but got {build_result['model'].layers[-1].units} units."
            assert build_result["model"].loss == "categorical_crossentropy", \
                f"Expected categorical_crossentropy loss, but got {build_result['model'].loss}."
        elif p == Paradigm.DL:
            for build_result in [build_result1, build_result2]:
                build_result["model"].summary()
                from keras.src.layers import Lambda
                assert isinstance(build_result["model"].layers[-1], Lambda), \
                    f"Expected Lambda in the last layer, but got {type(build_result['model'].layers[-1])} instead."
        elif p == Paradigm.BV:
            build_result["model"].summary()
            assert build_result["model"].layers[-1].units == 1, \
                f"Expected 1 unit in the last layer, but got {build_result['model'].layers[-1].units} units."
            assert build_result["model"].loss == "binary_crossentropy", \
                f"Expected binary_crossentropy loss, but got {build_result['model'].loss}."
        elif p == Paradigm.AD:
            build_result["model"].summary()
            from keras.src.layers import Conv1D
            assert isinstance(build_result["model"].layers[-1], Conv1D), \
                f"Expected Conv1D in the last layer, but got {type(build_result['model'].layers[-1])} instead."

        # Training Steps
        verbose=1
        epochs=10
        if p == Paradigm.CSC:
            model, callbacks, metadata = build_result["model"], build_result["callbacks"], build_result["metadata"]
            model.fit(x=self.ds.get_trainset_data(),
                      y=self.ds.get_trainset_labels(),
                      callbacks=callbacks,
                      verbose=verbose,
                      epochs=epochs)
        elif p == Paradigm.DL:
            for build_result in [build_result1, build_result2]:
                model, callbacks, metadata = build_result["model"], build_result["callbacks"], build_result["metadata"]
                model.fit(x=self.ds.get_trainset_data(),
                          y=self.ds.get_trainset_labels(),
                          callbacks=callbacks,
                          verbose=verbose,
                          epochs=epochs)
        elif p == Paradigm.BV:
            model, callbacks, metadata = build_result["model"], build_result["callbacks"], build_result["metadata"]
            model.fit(x=self.ds.get_trainset_data(),
                      y=self.ds.get_trainset_labels(),
                      callbacks=callbacks,
                      verbose=verbose,
                      epochs=epochs)

        elif p == Paradigm.AD:
            model, callbacks, metadata = build_result["model"], build_result["callbacks"], build_result["metadata"]
            model.fit(x=self.ds.get_trainset_data(),
                      y=self.ds.get_trainset_labels(),
                      callbacks=callbacks,
                      verbose=verbose,
                      epochs=epochs)

        print("End.")


    # ---- Hardcoded test cases: CSC ----
    def test_csc_cnn_valid(self):
        self._build_validate_train(Paradigm.CSC, 'CNN_VALID')

    def test_csc_cnn_same(self):
        self._build_validate_train(Paradigm.CSC, 'CNN_SAME')

    def test_csc_fcn(self):
        self._build_validate_train(Paradigm.CSC, 'FCN')

    def test_csc_mlp(self):
        self._build_validate_train(Paradigm.CSC, 'MLP')

    def test_csc_resnet(self):
        self._build_validate_train(Paradigm.CSC, 'RESNET')

    def test_csc_encoder(self):
        self._build_validate_train(Paradigm.CSC, 'ENCODER')

    def test_csc_inception(self):
        self._build_validate_train(Paradigm.CSC, 'INCEPTION')

    def test_csc_tlenet(self):
        self._build_validate_train(Paradigm.CSC, 'TLENET')

    def test_csc_mcdcnn(self):
        self._build_validate_train(Paradigm.CSC, 'MCDCNN')

    def test_csc_mcnn(self):
        self._build_validate_train(Paradigm.CSC, 'MCNN')

    def test_csc_twiesn(self):
        self._build_validate_train(Paradigm.CSC, 'TWIESN')

    # ---- Hardcoded test cases: DL ----
    def test_dl_cnn_valid(self):
        self._build_validate_train(Paradigm.DL, 'CNN_VALID')

    def test_dl_cnn_same(self):
        self._build_validate_train(Paradigm.DL, 'CNN_SAME')

    def test_dl_fcn(self):
        self._build_validate_train(Paradigm.DL, 'FCN')

    def test_dl_mlp(self):
        self._build_validate_train(Paradigm.DL, 'MLP')

    def test_dl_resnet(self):
        self._build_validate_train(Paradigm.DL, 'RESNET')

    def test_dl_encoder(self):
        self._build_validate_train(Paradigm.DL, 'ENCODER')

    def test_dl_inception(self):
        self._build_validate_train(Paradigm.DL, 'INCEPTION')

    def test_dl_tlenet(self):
        self._build_validate_train(Paradigm.DL, 'TLENET')

    def test_dl_mcdcnn(self):
        self._build_validate_train(Paradigm.DL, 'MCDCNN')

    def test_dl_mcnn(self):
        self._build_validate_train(Paradigm.DL, 'MCNN')

    def test_dl_twiesn(self):
        self._build_validate_train(Paradigm.DL, 'TWIESN')

    # ---- Hardcoded test cases: BV ----
    def test_bv_cnn_valid(self):
        self._build_validate_train(Paradigm.BV, 'CNN_VALID')

    def test_bv_cnn_same(self):
        self._build_validate_train(Paradigm.BV, 'CNN_SAME')

    def test_bv_fcn(self):
        self._build_validate_train(Paradigm.BV, 'FCN')

    def test_bv_mlp(self):
        self._build_validate_train(Paradigm.BV, 'MLP')

    def test_bv_resnet(self):
        self._build_validate_train(Paradigm.BV, 'RESNET')

    def test_bv_encoder(self):
        self._build_validate_train(Paradigm.BV, 'ENCODER')

    def test_bv_inception(self):
        self._build_validate_train(Paradigm.BV, 'INCEPTION')

    def test_bv_tlenet(self):
        self._build_validate_train(Paradigm.BV, 'TLENET')

    def test_bv_mcdcnn(self):
        self._build_validate_train(Paradigm.BV, 'MCDCNN')

    def test_bv_mcnn(self):
        self._build_validate_train(Paradigm.BV, 'MCNN')

    def test_bv_twiesn(self):
        self._build_validate_train(Paradigm.BV, 'TWIESN')

    # ---- Hardcoded test cases: AD ----
    def test_ad_cnn_valid(self):
        self._build_validate_train(Paradigm.AD, 'CNN_VALID')

    def test_ad_cnn_same(self):
        self._build_validate_train(Paradigm.AD, 'CNN_SAME')

    def test_ad_fcn(self):
        self._build_validate_train(Paradigm.AD, 'FCN')

    def test_ad_mlp(self):
        self._build_validate_train(Paradigm.AD, 'MLP')

    def test_ad_resnet(self):
        self._build_validate_train(Paradigm.AD, 'RESNET')

    def test_ad_encoder(self):
        self._build_validate_train(Paradigm.AD, 'ENCODER')

    def test_ad_inception(self):
        self._build_validate_train(Paradigm.AD, 'INCEPTION')

    def test_ad_tlenet(self):
        self._build_validate_train(Paradigm.AD, 'TLENET')

    def test_ad_mcdcnn(self):
        self._build_validate_train(Paradigm.AD, 'MCDCNN')

    def test_ad_mcnn(self):
        self._build_validate_train(Paradigm.AD, 'MCNN')

    def test_ad_twiesn(self):
        self._build_validate_train(Paradigm.AD, 'TWIESN')



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