# coding: utf-8

from __future__ import absolute_import
import os, shutil, tempfile

from pandas import read_hdf
from pandas.testing import assert_frame_equal

from jass.models.worktable import create_worktable_file

from . import JassTestCase


class TestWorkTable(object):
    __test__ = False

    @classmethod
    def setUpClass(cls):
        # Create a temporary directory
        cls.test_dir = tempfile.mkdtemp()
        init_file_path = cls.get_file_path_fn("initTable.hdf5")
        cls.result_hdf_path = os.path.join(cls.test_dir, "test_worktable.hdf5")
        cls.expected_hdf_path = cls.get_file_path_fn("worktable.hdf5")
        create_worktable_file(
            cls.phenotypes_sel, init_file_path, cls.result_hdf_path, cls.remove_nan
        )
        cls.expected_sumstatjosttab = read_hdf(cls.expected_hdf_path, "SumStatTab")
        cls.expected_sumstatjosttab.reset_index(drop=True, inplace=True)
        cls.result_sumstatjosttab = read_hdf(cls.result_hdf_path, "SumStatTab")
        cls.result_sumstatjosttab.reset_index(drop=True, inplace=True)
        cls.expected_regionsubtable = read_hdf(cls.expected_hdf_path, "Regions")
        cls.result_regionsubtable = read_hdf(cls.result_hdf_path, "Regions")
        cls.expected_summarytable = read_hdf(cls.expected_hdf_path, "summaryTable")
        cls.result_summarytable = read_hdf(cls.result_hdf_path, "summaryTable")
        cls.expected_subcov = read_hdf(cls.expected_hdf_path, "COV")
        cls.result_subcov = read_hdf(cls.result_hdf_path, "COV")
        cls.expected_phenolist = read_hdf(cls.expected_hdf_path, "PhenoList")
        cls.result_phenolist = read_hdf(cls.result_hdf_path, "PhenoList")

    @classmethod
    def tearDownClass(cls):
        # Remove the directory after the test
        shutil.rmtree(cls.test_dir)
        pass

    def test_compare_sumstatjosttab(self):
        """
        Compare result and expected SumStatJostTab
        """
        assert_frame_equal(self.expected_sumstatjosttab, self.result_sumstatjosttab, check_like=True)

    def test_compare_regionsubtable(self):
        """
        Compare result and expected RegionSubTable
        """
        assert_frame_equal(self.expected_regionsubtable, self.result_regionsubtable, check_like=True)

    def test_compare_summarytable(self):
        """
        Compare result and expected SummaryTable
        """
        assert_frame_equal(self.expected_summarytable, self.result_summarytable, check_like=True)

    def test_compare_subcov(self):
        """
        Compare result and expected subCOV
        """
        assert_frame_equal(self.expected_subcov, self.result_subcov, check_like=True)

    def test_compare_phenolist(self):
        """
        Compare result and expected PhenoList
        """
        assert_frame_equal(self.expected_phenolist, self.result_phenolist, check_like=True)



phenotypes_disney = ["z_DISNEY_POCAHONT", "z_DISNEY_RATATOUY"]
phenotypes_cars = [
    "z_BMW_ISETTA",
    "z_BMW_MINI",
    "z_FIAT_CINQCENT",
    "z_FIAT_CINQUECENTO",
    "z_MERCO_SMART",
    "z_TATA_TATANANO",
]

params = [
    ("nonans", "data_test1", phenotypes_disney, True),
    ("withnans", "data_test2", phenotypes_cars, False),
]

for name, param, phenotypes_sel, remove_nan in params:
    cls_name = "TestWorkTable_%s" % (name,)
    globals()[cls_name] = type(
        cls_name,
        (TestWorkTable, JassTestCase),
        {
            "test_folder": param,
            "phenotypes_sel": phenotypes_sel,
            "remove_nan": remove_nan,
            "__test__": True,
        },
    )

if __name__ == "__main__":
    import unittest

    unittest.main()
