import os
import tempfile
import unittest

import netCDF4
import numpy as np
from netCDF4 import Dataset, EnumType
from numpy.testing import assert_array_equal

FILE_NAME = tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name
ENUM_NAME = 'cloud_t'
ENUM_BASETYPE = np.int8
VAR_NAME = 'primary_cloud'
ENUM_DICT = {'Altocumulus': 7, 'Missing': 127, 'Stratus': 2, 'Clear': 0,
'Nimbostratus': 6, 'Cumulus': 4, 'Altostratus': 5, 'Cumulonimbus': 1,
'Stratocumulus': 3}
datain = np.array([ENUM_DICT['Clear'],ENUM_DICT['Stratus'],ENUM_DICT['Cumulus'],\
                   ENUM_DICT['Missing'],ENUM_DICT['Cumulonimbus']],dtype=ENUM_BASETYPE)
datain_masked = np.ma.masked_values(datain,ENUM_DICT['Missing'])


class EnumTestCase(unittest.TestCase):

    def setUp(self):
        self.file = FILE_NAME
        f = Dataset(self.file,'w')
        cloud_type = f.createEnumType(ENUM_BASETYPE,ENUM_NAME,ENUM_DICT)
        # make sure KeyError raised if non-integer basetype used.
        try:
            cloud_typ2 = f.createEnumType(np.float32,ENUM_NAME,ENUM_DICT)  # type: ignore[arg-type]  # mypy correctly doesn't like float32
        except KeyError:
            pass
        f.createDimension('time',None)
        cloud_var =\
        f.createVariable(VAR_NAME,cloud_type,'time',\
                         fill_value=ENUM_DICT['Missing'])
        cloud_var[:] = datain_masked
        # make sure ValueError raised if illegal value assigned to Enum var.
        try:
            cloud_var[cloud_var.shape[0]] = 99
        except ValueError:
            pass
        f.close()

    def tearDown(self):
        # Remove the temporary files
        os.remove(self.file)

    def runTest(self):
        """testing enum data type"""
        f = Dataset(self.file, 'r')
        v = f.variables[VAR_NAME]
        assert isinstance(v.datatype, EnumType)
        assert v.datatype.enum_dict == ENUM_DICT
        assert list(f.enumtypes.keys()) == [ENUM_NAME]
        assert f.enumtypes[ENUM_NAME].name == ENUM_NAME # issue 775
        assert f.enumtypes[ENUM_NAME].dtype == ENUM_BASETYPE
        assert v._FillValue == ENUM_DICT['Missing']
        v.set_auto_mask(False)
        data = v[:]
        assert_array_equal(data, datain)
        v.set_auto_mask(True) # check to see if auto masking works
        data = v[:]
        assert_array_equal(data, datain_masked)
        assert_array_equal(data.mask, datain_masked.mask)
        f.close()

class EnumDictTestCase(unittest.TestCase):

    # issue 1128
    def setUp(self):
        DT = np.int16; BITS = 8
        self.STORED_VAL = DT(2**BITS)
        self.VAL_MAP = {f'bits_{n}': DT(2**n) for n in range(1,BITS+1)}
        self.VAL_MAP['invalid'] = DT(0)
        self.file = tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name
        with netCDF4.Dataset(self.file, 'w') as nc:
            # The enum is created with dtype=int16, so it will allow BITS values up to 15
            et = nc.createEnumType(DT, 'etype', self.VAL_MAP)
            ev = nc.createVariable('evar', et)
            # Succeeds because the created EnumType does keep the correct dict
            ev[...] = self.STORED_VAL

    def tearDown(self):
        os.remove(self.file)

    def runTest(self):
        with netCDF4.Dataset(self.file, 'r') as nc:
            read_var = nc['evar']
            read_et = nc.enumtypes["etype"]
            assert read_var[...] == self.STORED_VAL
            assert read_et.enum_dict == self.VAL_MAP

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