# # # This source code is subject to the license referenced at
# # # https://github.com/NRLMMD-GEOIPS.

"""Test script for representative product comparisons."""

import logging

from geoips.interfaces.class_based.output_checkers import BaseOutputCheckerPlugin
from geoips.commandline.log_setup import log_with_emphasis
from geoips.geoips_utils import get_numpy_seeded_random_generator

LOG = logging.getLogger(__name__)


class NetcdfOutputCheckerPlugin(BaseOutputCheckerPlugin):
    """NetCDF Output Checker Plugin Class.

    Inherits from BaseOutputCheckerPlugin, which contains shared functionality among
    output_checker plugins. Responsible for comparing two NetCDF files
    (produced, expected) within a given threshold.
    """

    interface = "output_checkers"
    family = "standard"
    name = "netcdf"

    def get_test_files(self, test_data_dir):
        """Return a Series of Netcdf paths modified from compare."""
        import xarray as xr
        from os import makedirs
        from os.path import exists, join

        savedir = join(test_data_dir, "scratch", "unit_tests", "test_netcdf")
        if not exists(savedir):
            makedirs(savedir)
        # Path for the "compare" NetCDF file
        compare_file = join(savedir, "compare.nc")

        # Generate random data for the "compare" file
        predictable_random = get_numpy_seeded_random_generator()
        compare_data = predictable_random.random((100, 100))
        compare_ds = xr.Dataset(data_vars={"data": (("x", "y"), compare_data)})
        compare_ds.to_netcdf(compare_file)

        # Paths for the other files
        matched_path = join(savedir, "matched.nc")
        close_mismatch_path = join(savedir, "close_mismatch.nc")
        bad_mismatch_path = join(savedir, "bad_mismatch.nc")

        # Randomly modify the "compare" data for "close_mismatch" and "bad_mismatch"
        close_mismatch_data = compare_data + predictable_random.normal(
            scale=0.05, size=compare_data.shape
        )
        bad_mismatch_data = compare_data + predictable_random.normal(
            scale=0.25, size=compare_data.shape
        )

        # Create DataArrays for the modified data
        close_mismatch_da = xr.DataArray(data=close_mismatch_data, dims=("x", "y"))
        bad_mismatch_da = xr.DataArray(data=bad_mismatch_data, dims=("x", "y"))

        # Create Datasets and save them to NetCDF files
        close_mismatch_ds = xr.Dataset(data_vars={"data": close_mismatch_da})
        bad_mismatch_ds = xr.Dataset(data_vars={"data": bad_mismatch_da})

        close_mismatch_ds.to_netcdf(close_mismatch_path)
        bad_mismatch_ds.to_netcdf(bad_mismatch_path)
        compare_ds.to_netcdf(matched_path)
        return compare_file, [matched_path, close_mismatch_path, bad_mismatch_path]

    def perform_test_comparisons(self, compare_file, test_files):
        """Test the comparison of two Netcdf files with the Netcdf Output Checker.

        Parameters
        ----------
        compare_file : str
            Path to the comparison (baseline) file.
        test_files : list[str]
            List of test file paths [matched, close_mismatch, bad_mismatch].

        Raises
        ------
        AssertionError
            If any comparison result doesn't match expected outcome.
        """
        for path_idx in range(len(test_files)):
            retval = self.outputs_match(
                test_files[path_idx],
                compare_file,
            )
            if path_idx == 0:
                assert retval is True
            else:
                assert retval is False

    def correct_file_format(self, fname):
        """Check if fname is a geoips formatted netcdf file.

        Parameters
        ----------
        fname : str
            Name of file to check.

        Returns
        -------
        bool
            True if it is a geoips netcdf file, False otherwise.
        """
        import xarray

        try:
            xobj = xarray.open_dataset(fname)
        except Exception:
            return False
        from geoips.geoips_utils import get_required_geoips_xarray_attrs

        return set(get_required_geoips_xarray_attrs()).issubset(set(xobj.attrs.keys()))

    def outputs_match(self, output_product, compare_product):
        """Check if two geoips formatted netcdf files match.

        Parameters
        ----------
        output_product : str
            Full path to current output product
        compare_product : str
            Full path to comparison product

        Returns
        -------
        bool
            Return True if products match, False if they differ
        """
        out_difftxt = self.get_out_diff_fname(compare_product, output_product)
        diffout = []
        retval = True
        import xarray

        out_xobj = xarray.open_dataset(output_product)
        compare_xobj = xarray.open_dataset(compare_product)

        if out_xobj.attrs != compare_xobj.attrs:
            message = "BAD GeoIPS NetCDF file attributes do NOT match exactly"
            log_with_emphasis(
                LOG.interactive,
                message,
                f"output_product: {output_product}",
                f"compare_product: {compare_product}",
            )
            for attr in out_xobj.attrs.keys():
                if attr not in compare_xobj.attrs:
                    diffstr = (
                        f"\nattr {attr}\n\n"
                        f"output\n{out_xobj.attrs[attr]}\n\n"
                        "not in comparison\n"
                    )
                    diffout += [diffstr]
                    LOG.interactive(diffstr)
                elif out_xobj.attrs[attr] != compare_xobj.attrs[attr]:
                    diffstr = (
                        f"\nattr {attr}\n\n"
                        f"output\n{out_xobj.attrs[attr]}\n\n"
                        f"comparison\n{compare_xobj.attrs[attr]}\n"
                    )
                    diffout += [diffstr]
                    LOG.interactive(diffstr)
            for attr in compare_xobj.attrs.keys():
                if attr not in out_xobj.attrs:
                    diffstr = (
                        f"\nattr {attr}\n\n"
                        f"not in output\n\n"
                        f"comparison\n{compare_xobj.attrs[attr]}\n"
                    )
                    diffout += [diffstr]
                    LOG.interactive(diffstr)
                elif out_xobj.attrs[attr] != compare_xobj.attrs[attr]:
                    diffstr = (
                        f"\nattr {attr}\n\n"
                        f"output\n{out_xobj.attrs[attr]}\n\n"
                        f"comparison\n{compare_xobj.attrs[attr]}\n"
                    )
                    diffout += [diffstr]
                    LOG.interactive(diffstr)
            diffout += ["\n"]
            retval = False

        if retval is False:
            with open(out_difftxt, "w") as fobj:
                fobj.writelines(diffout)
            return False

        try:
            xarray.testing.assert_allclose(compare_xobj, out_xobj)
        except AssertionError as resp:
            message = "BAD GeoIPS NetCDF files do not match within tolerance"
            log_with_emphasis(
                LOG.interactive,
                message,
                f"output_product: {output_product}",
                f"compare_product: {compare_product}",
            )
            log_with_emphasis(
                LOG.interactive, *[line for line in str(resp).split("\n")]
            )
            diffout += [
                "\nxarray objects do not match between current output and comparison\n"
            ]
            diffout += [f"\nOut: {out_xobj}\n"]
            diffout += [f"\nCompare: {compare_xobj}\n"]
            diffout += [f"\n{resp}\n"]
            for varname in compare_xobj.variables:
                logged_messages = self.log_object_diff_values(
                    out_xobj, compare_xobj, varname, log_function=LOG.warning
                )
                diffout.extend([[m] for m in logged_messages])
            retval = False

        try:
            xarray.testing.assert_identical(compare_xobj, out_xobj)
        except AssertionError as resp:
            lines = ["INFORMATIONAL ONLY assert_identical differences"]
            for line in str(resp).split("\n"):
                lines += [line]
            for varname in compare_xobj.variables:
                self.log_object_diff_values(
                    out_xobj, compare_xobj, varname, log_function=LOG.info
                )
            log_with_emphasis(LOG.info, *lines)

        if retval is False:
            with open(out_difftxt, "w") as fobj:
                fobj.writelines(diffout)
            return False

        return True

    def log_object_diff_values(
        self, object1, object2, compare_key, log_function=LOG.info
    ):
        r"""
        Log differences two objects via key and returns messages detailing differences.

        Computes the maximum, minimum, and mean differences for the values associated
        with a specified key in two objects. Generates messages for any non-zero
        differences and logs these messages. The messages are also returned as a list
        of strings for further use. Takes in an optional custom logging function,
        defaults to info logging.

        Objects must expose a min, max, and mean function after the difference
        between them is computed.

        Parameters
        ----------
        object1 : object
            The first object containing the key for comparison. Must support
            subscripting with the compare key and arithmetic operations.
        object2 : object
            The second object for comparison. Must also support subscripting with the
            compare key and arithmetic operations.
        compare_key : str
            The key for the values to be compared between `object1` and `object2`.
        log_func : func, optional
            The function to be used for logging, must take in a list of strings of min
            length 0.

        Returns
        -------
        messages : list of str
            A list of messages indicating non-zero min, max, and mean differences for
            the values associated with `compare_key`, formatted as strings.

        Notes
        -----
        Only logs and returns messages for non-zero differences. If there are no
        non-zero differences, an empty list is returned.

        Uses `log_with_emphasis` to log the messages.

        Examples
        --------
        Given two objects `obj1` and `obj2` with a key `temp`:

        >>> obj1 = {'temp': np.array([1, 2, 3])}
        >>> obj2 = {'temp': np.array([2, 3, 4])}
        >>> log_object_diff_values(obj1, obj2, 'temp')
        ['mindiff temp: -1\\n', 'maxdiff temp: -1\\n', 'meandiff temp: -1.0\\n']

        This will log and return messages about differences in `temp` values.
        """
        maxdiff = (object2[compare_key] - object1[compare_key]).max()
        mindiff = (object2[compare_key] - object1[compare_key]).min()
        meandiff = (object2[compare_key] - object1[compare_key]).mean()
        messages = [
            f"mindiff {compare_key}: {mindiff}\n",
            f"maxdiff {compare_key}: {maxdiff}\n",
            f"meandiff {compare_key}: {meandiff}\n",
        ]
        log_with_emphasis(log_function, *messages)
        return messages

    def call(self, compare_path, output_products):
        """Compare the "correct" NetCDF files found the list of current output_products.

        Compares files produced in the current processing run with the list of
        "correct" files contained in "compare_path".

        Parameters
        ----------
        compare_path : str
            Path to directory of "correct" products - filenames must match
            output_products
        output_products : list of str
            List of strings of current output products,
            to compare with products in compare_path

        Returns
        -------
        int
            Binary code: 0 if all comparisons were completed successfully.
        """
        retval = self.compare_outputs(
            compare_path,
            output_products,
        )
        return retval


PLUGIN_CLASS = NetcdfOutputCheckerPlugin
