"""Unittest for the rayleigh correction utilities."""
import contextlib
import os
import unittest
from pathlib import Path
from typing import Iterator
from unittest.mock import patch

import dask.array as da
import numpy as np
import pytest

from pyspectral import rayleigh
from pyspectral.tests.data import (
    TEST_RAYLEIGH_AZID_COORD,
    TEST_RAYLEIGH_LUT,
    TEST_RAYLEIGH_SATZ_COORD,
    TEST_RAYLEIGH_SUNZ_COORD,
    TEST_RAYLEIGH_WVL_COORD,
)

TEST_RAYLEIGH_RESULT1 = np.array([10.339923,    8.64748], dtype='float32')
TEST_RAYLEIGH_RESULT2 = np.array([9.653559, 8.464997], dtype='float32')
TEST_RAYLEIGH_RESULT3 = np.array([5.488729, 8.533125], dtype='float32')
TEST_RAYLEIGH_RESULT4 = np.array([0.0,    8.64748], dtype='float32')
TEST_RAYLEIGH_RESULT5 = np.array([9.653559, np.nan], dtype='float32')
TEST_RAYLEIGH_RESULT_R1 = np.array([16.66666667, 20.83333333, 25.], dtype='float32')
TEST_RAYLEIGH_RESULT_R2 = np.array([0., 6.25, 12.5], dtype='float32')

TEST_ZENITH_ANGLES_RESULTS = np.array([68.67631374, 68.67631374, 32., 0.])


class RelativeSpectralResponseTestData:
    """Create the class instance to hold the RSR test data."""

    def __init__(self):
        """Make a test data set of relative spectral responses."""
        self.rsr = {}
        channel_names = ['ch12', 'ch13', 'ch10', 'ch11', 'ch16', 'ch14',
                         'ch15', 'ch1', 'ch2', 'ch3', 'ch4', 'ch5', 'ch6',
                         'ch7', 'ch8', 'ch9']
        wvl = [9.6372744012936646, 10.407492196078628, 7.3468642293967275,
               8.5926867614178715, 13.280724258676756, 11.239642285822033,
               12.380741429961382, 0.47063607733748003, 0.5099976405799187,
               0.63914891611559055, 0.85668832355426627, 1.6100814361999056,
               2.2568056299864101, 3.8853663735353847, 6.2428987228916233,
               6.9411756334211789]
        ch3_wvl = np.array([0.55518544, 0.56779468, 0.58099002, 0.59481323, 0.60931027,
                            0.62453163, 0.64053291, 0.65737575, 0.67512828, 0.69386619,
                            0.71367401])
        ch3_resp = np.array([2.61000005e-05, 1.07899999e-04, 3.26119992e-03,
                             2.90650606e-01, 9.02460396e-01, 9.60878074e-01,
                             9.97266889e-01, 9.94823873e-01, 7.18220174e-01,
                             8.31819978e-03, 9.34999989e-05])

        idx = 0
        for chname in channel_names:
            self.rsr[chname] = {'det-1': {}}
            self.rsr[chname]['det-1']['central_wavelength'] = wvl[idx]
            idx = idx + 1

        chname = 'ch3'
        self.rsr[chname]['det-1']['wavelength'] = ch3_wvl
        self.rsr[chname]['det-1']['response'] = ch3_resp


@pytest.fixture(scope="module")
def rayleigh_lut_dir(tmp_path_factory) -> Iterator[Path]:
    """Get directory for fake LUTs that will be deleted after usage."""
    import shutil

    rayleigh_dir = tmp_path_factory.mktemp("pyspectral_fake_rayleigh_luts_")
    yield rayleigh_dir
    shutil.rmtree(rayleigh_dir, ignore_errors=True)


@pytest.fixture(scope="module")
def fake_lut_hdf5(rayleigh_lut_dir):
    """Create a fake LUT HDF5 file and necessary mocks to use it."""
    from pyspectral.testing import mock_rayleigh_luts

    rayl_data = {
        "reflectance": TEST_RAYLEIGH_LUT,
        "wavelengths": TEST_RAYLEIGH_WVL_COORD,
        "azimuth_difference": TEST_RAYLEIGH_AZID_COORD,
        "sun_zenith_secant": TEST_RAYLEIGH_SUNZ_COORD,
        "satellite_zenith_secant": TEST_RAYLEIGH_SATZ_COORD,
    }
    with mock_rayleigh_luts(rayleigh_dir=rayleigh_lut_dir,
                            lut_data=rayl_data,
                            aerosol_types=("marine_clean_aerosol",),
                            atmospheres=("midlatitude_summer", "subarctic_winter", "tropical", "us-standard"),
                            ):
        yield


@pytest.fixture(scope="module")
def fake_config_path(tmp_path_factory, rayleigh_lut_dir):
    """Create a custom config YAML on disk for this module's tests."""
    from pyspectral.testing import create_pyspectral_config_file

    config_path = tmp_path_factory.mktemp("pyspectral_fake_configs_") / "pyspectral.yaml"
    fake_config = {
        "rsr_dir": str(rayleigh_lut_dir),
        "rayleigh_dir": str(rayleigh_lut_dir),
        "download_from_internet": False,
    }
    create_pyspectral_config_file(config_path, fake_config)
    yield config_path


@pytest.fixture(scope="function")
def override_rayleigh_luts(fake_config_path, fake_lut_hdf5):
    """Mock pyspectral config to point to fake LUT directory.

    This fixture is function scoped so the overridden config doesn't have
    side effects for other tests. The rayleigh creation fixture is module
    scoped to avoid re-creating the HDF5 files unnecessarily.

    """
    from pyspectral.testing import override_config

    with override_config(config_path=fake_config_path):
        yield


def _create_rayleigh(platform='NOAA-20', sensor='VIIRS'):
    return rayleigh.Rayleigh(platform, sensor, atmosphere='midlatitude summer')


@contextlib.contextmanager
def mocked_rsr():
    """Mock the RSR class used by the Rayleigh class with fake data."""
    with patch('pyspectral.rayleigh.RelativeSpectralResponse') as mymock:
        instance = mymock.return_value
        instance.rsr = RelativeSpectralResponseTestData().rsr
        instance.unit = '1e-6 m'
        instance.si_scale = 1e-6
        yield mymock


def _create_dask_array(input_data, dtype):
    return da.from_array(np.array(input_data, dtype=dtype))


class TestRayleighDask:
    """Class for testing pyspectral.rayleigh - with dask-arrays as input."""

    def test_dask_clipping_angles_with_nans(self):
        """Test the clipping of angles outside coordinate range - with nan's in input."""
        from pyspectral.rayleigh import _clip_angles_inside_coordinate_range
        zenith_angle = da.array([79., 69., 32., np.nan])
        result = _clip_angles_inside_coordinate_range(zenith_angle, 2.75)
        np.testing.assert_allclose(result, TEST_ZENITH_ANGLES_RESULTS)

    def test_get_reflectance_dask_redband_outside_clip(self, override_rayleigh_luts):
        """Test getting the reflectance correction with dask inputs - using red band reflections outside 20-100."""
        sun_zenith = da.array([67., 32.])
        sat_zenith = da.array([45., 18.])
        azidiff = da.array([150., 110.])
        redband_refl = da.array([108., -0.5])
        viirs_rayl = _create_rayleigh()
        with mocked_rsr():
            refl_corr = viirs_rayl.get_reflectance(sun_zenith, sat_zenith, azidiff, 'ch3', redband_refl)
        np.testing.assert_allclose(refl_corr, TEST_RAYLEIGH_RESULT4)
        assert isinstance(refl_corr, da.Array)

    @pytest.mark.parametrize("dtype", [np.float32, np.float64])
    @pytest.mark.parametrize("use_dask", [False, True])
    @pytest.mark.parametrize(
        ("input_data", "exp_result"),
        [
            (([67.0, 32.0], [45.0, 18.0], [150.0, 110.0], [14.0, 5.0]), TEST_RAYLEIGH_RESULT1),
            (([60.0, 20.0], [49.0, 26.0], [140.0, 130.0], [12.0, 8.0]), TEST_RAYLEIGH_RESULT2),
        ]
    )
    def test_get_reflectance(self, override_rayleigh_luts, dtype, use_dask, input_data, exp_result):
        """Test getting the reflectance correction with dask inputs."""
        array_func = np.array if not use_dask else _create_dask_array
        sun_zenith = array_func(input_data[0], dtype=dtype)
        sat_zenith = array_func(input_data[1], dtype=dtype)
        azidiff = array_func(input_data[2], dtype=dtype)
        redband_refl = array_func(input_data[3], dtype=dtype)
        rayl = _create_rayleigh()
        with mocked_rsr():
            refl_corr = rayl.get_reflectance(sun_zenith, sat_zenith, azidiff, 'ch3', redband_refl)

        if use_dask:
            assert isinstance(refl_corr, da.Array)
            refl_corr_np = refl_corr.compute()
            assert refl_corr_np.dtype == refl_corr.dtype  # check that the final numpy array's dtype is equal
            refl_corr = refl_corr_np

        assert isinstance(refl_corr, np.ndarray)
        np.testing.assert_allclose(refl_corr, exp_result.astype(dtype), atol=4.0e-06)
        assert refl_corr.dtype == dtype  # check that the dask array's dtype is equal

    def test_get_reflectance_wvl_outside_range(self, override_rayleigh_luts):
        """Test getting the reflectance correction with wavelength outside correction range."""
        with mocked_rsr() as rsr_obj:
            rsr_obj.side_effect = IOError("No rsr data in pyspectral for this platform and sensor")
            sun_zenith = da.from_array([50., 10.])
            sat_zenith = da.from_array([39., 16.])
            azidiff = da.from_array([170., 110.])
            redband_refl = da.from_array([29., 12.])
            rayl = rayleigh.Rayleigh('UFO', 'unknown', atmosphere='midlatitude summer')

            # we gave a direct wavelength so the RSR object shouldn't be needed
            refl_corr = rayl.get_reflectance(sun_zenith, sat_zenith, azidiff, 1.2, redband_refl)
            np.testing.assert_allclose(refl_corr, 0)
            assert isinstance(refl_corr, da.Array)
            rsr_obj.assert_not_called()


class TestRayleigh:
    """Class for testing pyspectral.rayleigh."""

    def test_get_effective_wavelength_and_band_name_with_floats(self, override_rayleigh_luts):
        """Test getting the effective wavelength."""
        this = rayleigh.Rayleigh('Himawari-8', 'ahi')
        # Only ch3 (~0.63) testdata implemented yet...
        ewl, band_name = this._get_effective_wavelength_and_band_name(0.65)
        np.testing.assert_allclose(ewl, 650)  # 635.6167)
        assert isinstance(band_name, str)

        this = rayleigh.Rayleigh('Himawari-8', 'ahi')
        ewl, band_name = this._get_effective_wavelength_and_band_name(0.7)
        assert ewl == 700.0
        assert isinstance(band_name, str)
        ewl, band_name = this._get_effective_wavelength_and_band_name(0.9)
        assert ewl == 900.0
        assert isinstance(band_name, str)
        ewl, band_name = this._get_effective_wavelength_and_band_name(0.455)
        assert ewl == 455.0
        assert isinstance(band_name, str)

    def test_rayleigh_init(self, override_rayleigh_luts):
        """Test creating the Rayleigh object."""
        with patch('pyspectral.rayleigh.RelativeSpectralResponse') as mymock:
            instance = mymock.return_value
            instance.rsr = RelativeSpectralResponseTestData().rsr

            with pytest.raises(AttributeError):
                rayleigh.Rayleigh('Himawari-8', 'ahi', atmosphere='unknown')

            with pytest.raises(AttributeError):
                rayleigh.Rayleigh('Himawari-8', 'ahi', aerosol_type='unknown')

            this = rayleigh.Rayleigh('Himawari-8', 'ahi', atmosphere='subarctic winter')
            assert os.path.basename(this.reflectance_lut_filename).endswith('subarctic_winter.h5')
            assert this.sensor == 'ahi'

            this = rayleigh.Rayleigh('NOAA-19', 'avhrr/3', atmosphere='tropical')
            assert this.sensor == 'avhrr3'

    def test_clipping_angles_with_nans(self):
        """Test the clipping of angles outside coordinate range - with nan's in input."""
        from pyspectral.rayleigh import _clip_angles_inside_coordinate_range
        zenith_angle = np.array([79., 69., 32., np.nan])
        result = _clip_angles_inside_coordinate_range(zenith_angle, 2.75)
        np.testing.assert_allclose(result, TEST_ZENITH_ANGLES_RESULTS)

    def test_rayleigh_reduction(self, override_rayleigh_luts):
        """Test the code that reduces Rayleigh correction for high zenith angles."""
        # Test the Rayleigh reduction code
        sun_zenith = np.array([70., 65., 60.])
        in_rayleigh = [50, 50, 50]
        # Test case where no reduction is done.
        retv = rayleigh.Rayleigh.reduce_rayleigh_highzenith(sun_zenith, in_rayleigh, 70., 90., 1)
        np.testing.assert_allclose(retv, in_rayleigh)
        # Test case where moderate reduction is performed.
        retv = rayleigh.Rayleigh.reduce_rayleigh_highzenith(sun_zenith, in_rayleigh, 30., 90., 1)
        np.testing.assert_allclose(retv, TEST_RAYLEIGH_RESULT_R1)
        # Test case where extreme reduction is performed.
        retv = rayleigh.Rayleigh.reduce_rayleigh_highzenith(sun_zenith, in_rayleigh, 30., 90., 1.5)
        np.testing.assert_allclose(retv, TEST_RAYLEIGH_RESULT_R2)

    @pytest.mark.parametrize(
        ("platform_name", "sensor"),
        [
            ("FY-4B", "VIIRS"),
            ("FY-4B", "nosensor"),
        ]
    )
    def test_rayleigh_getname_failures(self, platform_name, sensor):
        """Test Rayleigh instrument selection arguments that should fail."""
        from pyspectral.rayleigh import normalize_sensor

        with pytest.raises(ValueError):
            normalize_sensor(platform_name, sensor)

    @pytest.mark.parametrize(
        ("platform_name", "sensor", "exp_name"),
        [
            ("FY-4A", "agri", "agri"),
            ("FY-4B", "agri", "agri"),
            ("FY-4B", "ghi", "ghi"),
            ("Metop-C", "avhrr/3", "avhrr3"),
        ]
    )
    def test_rayleigh_getname(self, platform_name, sensor, exp_name):
        """Test logic for Rayleigh instrument selection."""
        from pyspectral.rayleigh import normalize_sensor

        assert normalize_sensor(platform_name, sensor) == exp_name

    @patch('pyspectral.rayleigh.da', None)
    def test_get_reflectance_redband_outside_clip(self, override_rayleigh_luts):
        """Test getting the reflectance correction - using red band reflections outside 20 to 100."""
        sun_zenith = np.array([67., 32.])
        sat_zenith = np.array([45., 18.])
        azidiff = np.array([150., 110.])
        redband_refl = np.array([100., 20.])
        rayl = _create_rayleigh()
        with mocked_rsr():
            refl_corr1 = rayl.get_reflectance(
                sun_zenith, sat_zenith, azidiff, 'ch3', redband_refl)

        np.testing.assert_allclose(refl_corr1, TEST_RAYLEIGH_RESULT4)

        rng = np.random.default_rng(12345)
        rints_low = rng.integers(low=-10, high=20, size=2).astype('float')
        rints_high = rng.integers(low=100, high=200, size=2).astype('float')

        redband_refl = np.array([rints_high[0], rints_low[0]])
        with mocked_rsr():
            refl_corr2 = rayl.get_reflectance(
                sun_zenith, sat_zenith, azidiff, 'ch3', redband_refl)

        redband_refl = np.array([rints_high[1], rints_low[1]])
        with mocked_rsr():
            refl_corr3 = rayl.get_reflectance(
                sun_zenith, sat_zenith, azidiff, 'ch3', redband_refl)

        assert isinstance(refl_corr1, np.ndarray)
        assert isinstance(refl_corr2, np.ndarray)
        np.testing.assert_allclose(refl_corr1, refl_corr2)
        np.testing.assert_allclose(refl_corr2, refl_corr3)

    @patch('pyspectral.rayleigh.da', None)
    @pytest.mark.parametrize(
        ("sun_zenith", "sat_zenith", "azidiff", "redband_refl", "exp_result"),
        [
            (np.array([67., 32.]), np.array([45., 18.]), np.array([150., 110.]), np.array([14., 5.]),
             TEST_RAYLEIGH_RESULT1),
            (np.array([60., 20.]), np.array([49., 26.]), np.array([140., 130.]), np.array([12., 8.]),
             TEST_RAYLEIGH_RESULT2),
            (np.array([60., 20.]), np.array([49., 26.]), np.array([140., 130.]), np.array([12., np.nan]),
             TEST_RAYLEIGH_RESULT5),
        ]
    )
    @pytest.mark.parametrize("dtype", [np.float32, np.float64])
    def test_get_reflectance(self, override_rayleigh_luts, sun_zenith, sat_zenith, azidiff,
                             redband_refl, exp_result, dtype):
        """Test getting the reflectance correction."""
        rayl = _create_rayleigh()
        with mocked_rsr():
            refl_corr = rayl.get_reflectance(
                sun_zenith.astype(dtype),
                sat_zenith.astype(dtype),
                azidiff.astype(dtype),
                'ch3',
                redband_refl.astype(dtype))
        assert isinstance(refl_corr, np.ndarray)
        np.testing.assert_allclose(refl_corr, exp_result.astype(dtype), atol=4.0e-06)

    @patch('pyspectral.rayleigh.da', None)
    def test_get_reflectance_no_rsr(self, override_rayleigh_luts):
        """Test getting the reflectance correction, simulating that we have no RSR data."""
        with mocked_rsr() as rsr_obj:
            rsr_obj.side_effect = IOError("No rsr data in pyspectral for this platform and sensor")
            sun_zenith = np.array([50., 10.])
            sat_zenith = np.array([39., 16.])
            azidiff = np.array([170., 110.])
            redband_refl = np.array([29., 12.])
            ufo = rayleigh.Rayleigh('UFO', 'unknown', atmosphere='midlatitude summer')
            with pytest.raises(KeyError):
                ufo.get_reflectance(sun_zenith, sat_zenith, azidiff, 'ch3', redband_refl)

    @patch('pyspectral.rayleigh.da', None)
    def test_get_reflectance_float_wavelength(self, override_rayleigh_luts):
        """Test getting the reflectance correction."""
        with mocked_rsr() as rsr_obj:
            rsr_obj.side_effect = IOError("No rsr data in pyspectral for this platform and sensor")
            sun_zenith = np.array([50., 10.])
            sat_zenith = np.array([39., 16.])
            azidiff = np.array([170., 110.])
            redband_refl = np.array([29., 12.])
            rayl = rayleigh.Rayleigh('UFO', 'unknown', atmosphere='midlatitude summer')

            # we gave a direct wavelength so the RSR object shouldn't be needed
            refl_corr = rayl.get_reflectance(sun_zenith, sat_zenith, azidiff, 0.634, redband_refl)
            np.testing.assert_allclose(refl_corr, TEST_RAYLEIGH_RESULT3)
            assert isinstance(refl_corr, np.ndarray)
            rsr_obj.assert_not_called()

    @patch('pyspectral.rayleigh.da', None)
    def test_get_reflectance_wvl_outside_range(self, override_rayleigh_luts):
        """Test getting the reflectance correction with wavelength outside correction range."""
        with mocked_rsr() as rsr_obj:
            rsr_obj.side_effect = IOError("No rsr data in pyspectral for this platform and sensor")
            sun_zenith = np.array([50., 10.])
            sat_zenith = np.array([39., 16.])
            azidiff = np.array([170., 110.])
            redband_refl = np.array([29., 12.])
            rayl = rayleigh.Rayleigh('UFO', 'unknown', atmosphere='midlatitude summer')

            # we gave a direct wavelength so the RSR object shouldn't be needed
            refl_corr = rayl.get_reflectance(sun_zenith, sat_zenith, azidiff, 1.2, redband_refl)
            np.testing.assert_allclose(refl_corr, 0)
            assert isinstance(refl_corr, np.ndarray)
            rsr_obj.assert_not_called()

    def test_get_reflectance_no_lut(self, override_rayleigh_luts):
        """Test that missing a LUT causes an exception.."""
        ryl = rayleigh.Rayleigh('UFO', 'unknown', atmosphere='subarctic summer')

        with pytest.raises(FileNotFoundError):
            ryl.get_reflectance(
                np.zeros((5, 5), dtype=np.float64),
                np.zeros((5, 5), dtype=np.float64),
                np.zeros((5, 5), dtype=np.float64),
                0.64)


@pytest.mark.parametrize(
    ("version", "exp_download"),
    [
        (True, False),
        ("v0.0.0", True),
        (False, True),
    ],
)
def test_check_and_download(tmp_path, version, exp_download):
    """Test that check_and_download only downloads when necessary."""
    from pyspectral.rayleigh import check_and_download
    from pyspectral.testing import mock_rayleigh

    with mock_rayleigh(rayleigh_dir=tmp_path, existing_version=version), \
            unittest.mock.patch("pyspectral.rayleigh.download_luts") as download:
        check_and_download()
        if exp_download:
            download.assert_called()
        else:
            download.assert_not_called()
