import json
import cmor
import unittest
import numpy
from pathlib import Path

from netCDF4 import Dataset

CMIP7_TABLES_PATH = "cmip7-cmor-tables/tables"
CV_PATH = "TestTables/CMIP7_CV.json"

USER_INPUT = {
    "_AXIS_ENTRY_FILE": "CMIP7_coordinate.json",
    "_FORMULA_VAR_FILE": "CMIP7_formula_terms.json",
    "_cmip7_option": 1,
    "_controlled_vocabulary_file": CV_PATH,
    "activity_id": "CMIP",
    "calendar": "360_day",
    "cv_version": "6.2.19.0",
    "drs_specs": "MIP-DRS7",
    "experiment_id": "piControl",
    "forcing_index": "f30",
    "grid_label": "gn",
    "initialization_index": "i000001d",
    "institution_id": "PCMDI",
    "license_id": "CC BY 4.0",
    "nominal_resolution": "250 km",
    "outpath": ".",
    "physics_index": "p1",
    "realization_index": "r009",
    "source_id": "PCMDI-test-1-0",
    "tracking_prefix": "hdl:21.14100",
    "host_collection": "CMIP7",
    "frequency": "mon",
    "region": "glb",
    "archive_id": "WCRP",
    "output_path_template": "<activity_id><source_id><experiment_id><member_id><variable_id><branding_suffix><grid_label>",
    "output_file_template": "<variable_id><branding_suffix><frequency><region><grid_label><source_id><experiment_id><variant_label>",
}


class TestCMIP7(unittest.TestCase):

    def _load_dataset(self, updated_fields=None, removed_fields=None):
        input_data = USER_INPUT.copy()

        if removed_fields is not None:
            for field in removed_fields:
                input_data.pop(field, None)

        if updated_fields is not None:
            input_data.update(updated_fields)

        with open(self.input_json, "w") as input_file_handle:
            json.dump(input_data, input_file_handle, sort_keys=True, indent=4)

        error_flag = cmor.dataset_json(str(self.input_json))
        if error_flag:
            raise RuntimeError("CMOR dataset_json call failed")

    def setUp(self):
        """
        Write out a simple file using CMOR
        """
        self.input_json = Path("Test/input_cmip7.json")

        # Set up CMOR
        cmor.setup(inpath=CMIP7_TABLES_PATH, netcdf_file_action=cmor.CMOR_REPLACE)

        # Define dataset using USER_INPUT
        with open(self.input_json, "w") as input_file_handle:
            json.dump(USER_INPUT, input_file_handle, sort_keys=True, indent=4)

    def tearDown(self):
        self.input_json.unlink(missing_ok=True)

    def test_cmip7(self):
        self._load_dataset()

        data = [27] * (2 * 3 * 4)
        tos = numpy.array(data)
        tos.shape = (2, 3, 4)
        lat = numpy.array([10, 20, 30])
        lat_bnds = numpy.array([5, 15, 25, 35])
        lon = numpy.array([0, 90, 180, 270])
        lon_bnds = numpy.array([-45, 45,
                                135,
                                225,
                                315
                                ])
        time = numpy.array([15.5, 45])
        time_bnds = numpy.array([0, 31, 60])
        cmor.load_table("CMIP7_ocean.json")
        cmorlat = cmor.axis("latitude",
                            coord_vals=lat,
                            cell_bounds=lat_bnds,
                            units="degrees_north")
        cmorlon = cmor.axis("longitude",
                            coord_vals=lon,
                            cell_bounds=lon_bnds,
                            units="degrees_east")
        cmortime = cmor.axis("time",
                             coord_vals=time,
                             cell_bounds=time_bnds,
                             units="days since 2018")
        axes = [cmortime, cmorlat, cmorlon]
        cmortos = cmor.variable("tos_tavg-u-hxy-sea", "degC", axes)
        self.assertEqual(cmor.write(cmortos, tos), 0)
        filename = cmor.close(cmortos, file_name=True)
        self.assertEqual(cmor.close(), 0)

        ds = Dataset(filename)
        attrs = ds.ncattrs()
        test_attrs = {
            'branded_variable': 'tos_tavg-u-hxy-sea',
            'branding_suffix': 'tavg-u-hxy-sea',
            'temporal_label': 'tavg',
            'vertical_label': 'u',
            'horizontal_label': 'hxy',
            'area_label': 'sea',
            'region': 'glb',
            'frequency': 'mon',
            'archive_id': 'WCRP',
            'mip_era': 'CMIP7',
            'data_specs_version': 'CMIP-7.0.0.0',
            'host_collection': 'CMIP7',
        }

        for attr, val in test_attrs.items():
            self.assertIn(attr, attrs)
            self.assertEqual(val, ds.getncattr(attr))
        institution_id = USER_INPUT["institution_id"]
        license_id = "CC BY 4.0"
        license_type = "Creative Commons Attribution 4.0 International"
        license_url = "https://creativecommons.org/licenses/by/4.0/"
        license = \
            (f"{license_id}; CMIP7 data produced by {institution_id} is "
             f"licensed under a {license_type} License ({license_url}). "
             "Consult https://wcrp-cmip.github.io/cmip7-guidance/docs/CMIP7/"
             "Guidance_for_users for terms of "
             "use governing CMIP7 output, including citation requirements and "
             "proper acknowledgment. The data producers and data providers "
             "make no warranty, either express or implied, including, but not "
             "limited to, warranties of merchantability and fitness for a "
             "particular purpose. All liabilities arising from the supply of "
             "the information (including any liability arising in negligence) "
             "are excluded to the fullest extent permitted by law."
             )
        self.assertTrue("license" in attrs)
        self.assertTrue("license_id" in attrs)
        self.assertTrue("license_type" not in attrs)
        self.assertTrue("license_url" not in attrs)
        self.assertEqual(license, ds.getncattr("license"))
        self.assertEqual(license_id, ds.getncattr("license_id"))

        ds.close()

    def test_secondary_modeling_realm(self):
        self._load_dataset()

        data = [27] * (2 * 3 * 4)
        pr = numpy.array(data)
        pr.shape = (2, 3, 4)
        lat = numpy.array([10, 20, 30])
        lat_bnds = numpy.array([5, 15, 25, 35])
        lon = numpy.array([0, 90, 180, 270])
        lon_bnds = numpy.array([-45, 45,
                                135,
                                225,
                                315
                                ])
        time = numpy.array([15.5, 45])
        time_bnds = numpy.array([0, 31, 60])
        cmor.load_table("CMIP7_atmos.json")
        cmorlat = cmor.axis("latitude",
                            coord_vals=lat,
                            cell_bounds=lat_bnds,
                            units="degrees_north")
        cmorlon = cmor.axis("longitude",
                            coord_vals=lon,
                            cell_bounds=lon_bnds,
                            units="degrees_east")
        cmortime = cmor.axis("time",
                             coord_vals=time,
                             cell_bounds=time_bnds,
                             units="days since 2018")
        axes = [cmortime, cmorlat, cmorlon]
        cmorpr = cmor.variable("prra_tavg-u-hxy-is", "kg m-2 s-1", axes)
        self.assertEqual(cmor.write(cmorpr, pr), 0)
        filename = cmor.close(cmorpr, file_name=True)
        self.assertEqual(cmor.close(), 0)

        ds = Dataset(filename)
        attrs = ds.ncattrs()
        test_attrs = {
            'branded_variable': 'prra_tavg-u-hxy-is',
            'branding_suffix': 'tavg-u-hxy-is',
            'temporal_label': 'tavg',
            'vertical_label': 'u',
            'horizontal_label': 'hxy',
            'area_label': 'is',
            'region': 'glb',
            'frequency': 'mon',
            'archive_id': 'WCRP',
            'mip_era': 'CMIP7',
            'data_specs_version': 'CMIP-7.0.0.0',
            'host_collection': 'CMIP7',
            'realm': 'atmos land landIce',
        }

        for attr, val in test_attrs.items():
            self.assertIn(attr, attrs)
            self.assertEqual(val, ds.getncattr(attr))

        ds.close()

    def test_climatology(self):
        self._load_dataset()

        data = numpy.array([27, 28])
        time = numpy.array([15, 45])
        time_bnds = numpy.array([[0, 31], [31, 60]])
        cmor.load_table("CMIP7_atmos.json")
        cmortime = cmor.axis("time2",
                             coord_vals=time,
                             cell_bounds=time_bnds,
                             units="days since 2018")
        cmorco2 = cmor.variable("co2_tclm-u-hm-u", "mol mol-1", [cmortime])
        self.assertEqual(cmor.write(cmorco2, data), 0)
        filename = cmor.close(cmorco2, file_name=True)
        self.assertEqual(cmor.close(), 0)

        basename = Path(filename).name
        self.assertTrue(basename.endswith("_201801-201802.nc"))

        ds = Dataset(filename)
        self.assertEqual("mon", ds.getncattr("frequency"))
        time_var = ds.variables["time"]
        self.assertEqual("climatology_bnds", time_var.getncattr("climatology"))
        self.assertEqual("Monthly Climatology", time_var.getncattr("long_name"))

        ds.close()

    def test_conventions_set_by_user(self):
        self._load_dataset(updated_fields={"Conventions": "CF-1.13"})

        data = numpy.array([27, 28])
        time = numpy.array([15, 45])
        time_bnds = numpy.array([[0, 31], [31, 60]])
        cmor.load_table("CMIP7_atmos.json")
        cmortime = cmor.axis("time2",
                             coord_vals=time,
                             cell_bounds=time_bnds,
                             units="days since 2018")
        cmorco2 = cmor.variable("co2_tclm-u-hm-u", "mol mol-1", [cmortime])
        self.assertEqual(cmor.write(cmorco2, data), 0)
        filename = cmor.close(cmorco2, file_name=True)
        self.assertEqual(cmor.close(), 0)

        ds = Dataset(filename)
        self.assertEqual("CF-1.13", ds.getncattr("Conventions"))

        ds.close()


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