import sys

import numpy as np

import tables as tb
from tables.tests import common


class LargeTable(tb.IsDescription):
    time = tb.Int32Col()


class BasicTestCase(common.TempFileMixin, common.PyTablesTestCase):
    # file  = "test.h5"
    open_mode = "w"
    title = "This is the table title"
    dim1, dim2, dim3 = 24, 721, 1440
    nrows = dim1 * dim2 * dim3  # rows for a day
    chunkshape = nrows
    complib = "blosc2"  # default

    def setUp(self):
        super().setUp()

        # Create an instance of an HDF5 Table
        self.populateFile()
        self.h5file.close()

    def populateFile(self):
        group = self.h5file.root
        table = self.h5file.create_table(
            group,
            "table",
            LargeTable,
            "Large table",
            tb.Filters(complevel=1, complib=self.complib),
            chunkshape=self.chunkshape,
        )

        # Structured NumPy buffer for every day
        self.day_block = day_block = np.empty(self.nrows, dtype=table.dtype)
        day_block["time"] = np.arange(self.nrows)

        # Append groups of rows ("days") so that we have more than 2**31
        # (see https://github.com/PyTables/PyTables/issues/995)
        self.ndays = ndays = 90
        self.assertTrue(ndays * self.nrows > 2**31)
        if common.verbose:
            print(f"Writing {ndays} days...")
        for day in range(ndays):
            table.append(day_block)
        table.flush()

    def test00_values(self):
        """Check that written values are correct."""

        self.h5file = tb.open_file(self.h5fname)
        table = self.h5file.root.table
        nrows = self.nrows
        day_block = self.day_block
        if common.verbose:
            print(f"Checking {self.ndays} days...")
        for nday in range(self.ndays):
            day_block2 = table[nday * nrows : (nday + 1) * nrows]
            self.assertEqual(
                np.sum(day_block2["time"] == day_block["time"]),
                nrows,
                f"Values differ in day {nday}",
            )


@common.unittest.skipIf(
    not common.blosc_avail, "BLOSC compression library not available"
)
class BloscTestCase(BasicTestCase):
    title = "Blosc table"
    complib = "blosc"


@common.unittest.skipIf(
    not common.blosc2_avail, "BLOSC2 compression library not available"
)
class Blosc2TestCase(BasicTestCase):
    title = "Blosc2 table"
    complib = "blosc2"


class ZlibTestCase(BasicTestCase):
    title = "Zlib table"
    complib = "zlib"


def suite():
    theSuite = common.unittest.TestSuite()
    niter = 1
    # common.heavy = 1  # Uncomment this only for testing purposes

    for n in range(niter):
        theSuite.addTest(common.make_suite(BloscTestCase))
        theSuite.addTest(common.make_suite(Blosc2TestCase))
        if common.heavy:
            theSuite.addTest(common.make_suite(ZlibTestCase))

    return theSuite


if __name__ == "__main__":
    common.parse_argv(sys.argv)
    common.print_versions()
    common.unittest.main(defaultTest="suite")