from __future__ import annotations

import numpy as np
import pandas as pd
import pytest
import xarray as xr

import xclim.indices as xci
from xclim import atmos, convert, core, set_options
from xclim.core.calendar import build_climatology_bounds, percentile_doy
from xclim.core.units import convert_units_to

K2C = 273.15


class TestRainOnFrozenGround:
    @pytest.mark.parametrize("chunks", [{"time": 366}, None])
    def test_3d_data_with_nans(self, open_dataset, chunks):
        ds = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc")

        pr = ds.pr.copy()
        pr.values[1, 10] = np.nan

        if chunks:
            ds = ds.chunk(chunks)

        out = atmos.rain_on_frozen_ground_days(pr, ds.tas, freq="YS")
        np.testing.assert_array_equal(out.sel(location="Montréal"), [np.nan, 4, 5, 3])


class TestRainSeason:
    # @pytest.mark.parametrize("chunks", [{"time": 366}, None])
    def test_3d_data_with_nans(self, open_dataset):
        ds = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc")

        pr = ds.pr.isel(location=0).copy()
        pr[{"time": [0, 10, 100]}] = np.nan
        out = {}
        out["start"], out["end"], out["length"] = atmos.rain_season(
            pr,
            freq="YS-JAN",
            window_dry_end=5,
            date_min_start="01-01",
            date_min_end="01-01",
        )
        out_arr = np.array([out[var].values for var in ["start", "end", "length"]])
        out_exp = np.array(
            [
                [np.nan, 12.0, 6.0, 27.0],
                [np.nan, np.nan, 141.0, np.nan],
                [np.nan, 354.0, 135.0, 339.0],
            ]
        )
        np.testing.assert_array_equal(out_arr, out_exp)


class TestPrecipAccumulation:
    nc_pr = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"
    nc_tasmin = "NRCANdaily/nrcan_canada_daily_tasmin_1990.nc"

    def test_3d_data_with_nans(self, open_dataset):
        # test with 3d data
        pr = open_dataset(self.nc_pr).pr  # mm/s
        prMM = open_dataset(self.nc_pr).pr
        prMM *= 86400
        prMM.attrs["units"] = "mm/day"
        # put a nan somewhere
        prMM.values[10, 1, 0] = np.nan
        pr.values[10, 1, 0] = np.nan

        out1 = atmos.precip_accumulation(pr, freq="MS")
        out2 = atmos.precip_accumulation(prMM, freq="MS")

        # test kg m-2 s-1
        pr.attrs["units"] = "kg m-2 s-1"
        out3 = atmos.precip_accumulation(pr, freq="MS")

        np.testing.assert_allclose(out1, out2, rtol=1e-7, atol=1e-4)
        np.testing.assert_allclose(out1, out3, rtol=1e-7, atol=1e-4)

        # check some vector with and without a nan
        x1 = prMM[:31, 0, 0].values

        pr_tot = x1.sum()

        np.testing.assert_almost_equal(pr_tot, out1.values[0, 0, 0], 4)

        assert np.isnan(out1.values[0, 1, 0])
        assert np.isnan(out1.values[0, -1, -1])

    def test_with_different_phases(self, open_dataset):
        # test with different phases
        pr = open_dataset(self.nc_pr).pr  # mm/s
        tasmin = open_dataset(self.nc_tasmin).tasmin  # K

        out_tot = atmos.precip_accumulation(pr, freq="MS")
        out_sol = atmos.solid_precip_accumulation(pr, tas=tasmin, freq="MS")
        out_liq = atmos.liquid_precip_accumulation(pr, tas=tasmin, freq="MS")

        np.testing.assert_array_almost_equal(out_liq + out_sol, out_tot, 4)

        assert "solid" in out_sol.description
        assert "liquid" in out_liq.description
        assert out_sol.standard_name == "lwe_thickness_of_snowfall_amount"

        # With a non-default threshold
        out_sol = atmos.solid_precip_accumulation(pr, tas=tasmin, thresh="40 degF", freq="MS")
        out_liq = atmos.liquid_precip_accumulation(pr, tas=tasmin, thresh="40 degF", freq="MS")

        np.testing.assert_array_almost_equal(out_liq + out_sol, out_tot, 4)


class TestPrecipAverage:
    nc_pr = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"
    nc_tasmin = "NRCANdaily/nrcan_canada_daily_tasmin_1990.nc"

    def test_3d_data_with_nans(self, open_dataset):
        # test with 3d data
        pr = open_dataset(self.nc_pr).pr  # mm/s
        prMM = open_dataset(self.nc_pr).pr
        prMM *= 86400
        prMM.attrs["units"] = "mm/day"
        # put a nan somewhere
        prMM.values[10, 1, 0] = np.nan
        pr.values[10, 1, 0] = np.nan

        out1 = atmos.precip_average(pr, freq="MS")
        out2 = atmos.precip_average(prMM, freq="MS")

        # test kg m-2 s-1
        pr.attrs["units"] = "kg m-2 s-1"
        out3 = atmos.precip_average(pr, freq="MS")

        np.testing.assert_array_almost_equal(out1, out2, 3)
        np.testing.assert_array_almost_equal(out1, out3, 5)

        # check some vector with and without a nan
        x1 = prMM[:31, 0, 0].values

        pr_mean = x1.mean()

        np.testing.assert_almost_equal(pr_mean, out1.values[0, 0, 0], 4)

        assert np.isnan(out1.values[0, 1, 0])

        assert np.isnan(out1.values[0, -1, -1])

    def test_with_different_phases(self, open_dataset):
        # test with different phases
        pr = open_dataset(self.nc_pr).pr  # mm/s
        tasmin = open_dataset(self.nc_tasmin).tasmin  # K

        out_tot = atmos.precip_average(pr, freq="MS")
        out_sol = atmos.solid_precip_average(pr, tas=tasmin, freq="MS")
        out_liq = atmos.liquid_precip_average(pr, tas=tasmin, freq="MS")

        np.testing.assert_array_almost_equal(out_liq + out_sol, out_tot, 4)

        assert "solid" in out_sol.description
        assert "liquid" in out_liq.description
        assert out_sol.standard_name == "lwe_average_of_snowfall_amount"

        # With a non-default threshold
        out_sol = atmos.solid_precip_average(pr, tas=tasmin, thresh="40 degF", freq="MS")
        out_liq = atmos.liquid_precip_average(pr, tas=tasmin, thresh="40 degF", freq="MS")

        np.testing.assert_array_almost_equal(out_liq + out_sol, out_tot, 4)


class TestStandardizedPrecip:
    nc_ds = "sdba/CanESM2_1950-2100.nc"

    @pytest.mark.slow
    def test_3d_data_with_nans(self, open_dataset):
        # test with data
        ds = open_dataset(self.nc_ds)
        pr = ds.pr.sel(time=slice("2000"))  # kg m-2 s-1
        prMM = convert_units_to(pr, "mm/day", context="hydro")  # noqa
        # put a nan somewhere
        prMM.values[10] = np.nan
        pr.values[10] = np.nan

        out1 = atmos.standardized_precipitation_index(
            pr, freq="MS", window=1, dist="gamma", method="APP", fitkwargs={"floc": 0}
        )
        out2 = atmos.standardized_precipitation_index(
            prMM, freq="MS", window=1, dist="gamma", method="APP", fitkwargs={"floc": 0}
        )
        np.testing.assert_array_almost_equal(out1, out2, 3)

        # preparing water_budget for SPEI test
        with xr.set_options(keep_attrs=True):
            tasmax = ds.tasmax
            tas = tasmax - 2.5
            tasmin = tasmax - 5
            wb = xci.water_budget(pr, None, tasmin, tasmax, tas, None)
            wbMM = convert_units_to(wb, "mm/day", context="hydro")  # noqa

        out3 = atmos.standardized_precipitation_evapotranspiration_index(
            wb,
            freq="MS",
            window=1,
            dist="gamma",
            method="APP",
            fitkwargs={"floc": 0},
        )
        out4 = atmos.standardized_precipitation_evapotranspiration_index(
            wbMM,
            freq="MS",
            window=1,
            dist="gamma",
            method="APP",
            fitkwargs={"floc": 0},
        )

        np.testing.assert_array_almost_equal(out3, out4, 3)


class TestWetDays:
    nc_file = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"

    def test_3d_data_with_nans(self, open_dataset):
        # test with 3d data
        pr = open_dataset(self.nc_file).pr
        prMM = open_dataset(self.nc_file).pr
        prMM.values *= 86400.0
        prMM.attrs["units"] = "mm/day"
        # put a nan somewhere
        prMM.values[10, 1, 0] = np.nan
        pr.values[10, 1, 0] = np.nan
        pr_min = "5 mm/d"
        out1 = atmos.wetdays(pr, thresh=pr_min, freq="MS")
        out2 = atmos.wetdays(prMM, thresh=pr_min, freq="MS")

        # test kg m-2 s-1
        pr.attrs["units"] = "kg m-2 s-1"
        out3 = atmos.wetdays(pr, thresh=pr_min, freq="MS")

        np.testing.assert_array_equal(out1, out2)
        np.testing.assert_array_equal(out1, out3)

        # check some vector with and without a nan
        x1 = prMM[:31, 0, 0].values

        wd1 = (x1 >= int(pr_min.split(" ")[0])).sum()

        assert wd1 == out1.values[0, 0, 0]

        assert np.isnan(out1.values[0, 1, 0])

        # make sure that vector with all nans gives nans whatever skipna
        assert np.isnan(out1.values[0, -1, -1])
        # assert (np.isnan(wds.values[0, -1, -1]))


class TestWetPrcptot:
    """Testing of prcptot with wet days"""

    def test_simple(self, atmosds):
        pr = atmosds.pr

        thresh = "1 mm/day"
        out = atmos.wet_precip_accumulation(pr, thresh=thresh)

        # Reference value
        t = core.units.convert_units_to(thresh, pr, context="hydro")
        pa = atmos.precip_accumulation(pr.where(pr >= t, 0))
        np.testing.assert_array_equal(out, pa)


class TestDailyIntensity:
    # testing of wet_day and daily_pr_intensity, both are related
    nc_file = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"

    def test_3d_data_with_nans(self, open_dataset):
        # test with 3d data
        pr = open_dataset(self.nc_file).pr
        prMM = open_dataset(self.nc_file).pr
        prMM.values *= 86400.0
        prMM.attrs["units"] = "mm/day"
        # put a nan somewhere
        prMM.values[10, 1, 0] = np.nan
        pr.values[10, 1, 0] = np.nan

        # compute with both skipna options
        pr_min = "2 mm/d"
        # dis = daily_pr_intensity(pr, pr_min=pr_min, freq='MS', skipna=True)

        out1 = atmos.daily_pr_intensity(pr, thresh=pr_min, freq="MS")
        out2 = atmos.daily_pr_intensity(prMM, thresh=pr_min, freq="MS")

        # test kg m-2 s-1
        pr.attrs["units"] = "kg m-2 s-1"
        out3 = atmos.daily_pr_intensity(pr, thresh=pr_min, freq="MS")

        np.testing.assert_array_almost_equal(out1, out2, 3)
        np.testing.assert_array_almost_equal(out1, out3, 3)

        x1 = prMM[:31, 0, 0].values

        di1 = x1[x1 >= int(pr_min.split(" ")[0])].mean()
        # buffer = np.ma.masked_invalid(x2)
        # di2 = buffer[buffer >= pr_min].mean()

        assert np.allclose(di1, out1.values[0, 0, 0])
        # assert (np.allclose(di1, dis.values[0, 0, 0]))
        assert np.isnan(out1.values[0, 1, 0])
        # assert (np.allclose(di2, dis.values[0, 1, 0]))
        assert np.isnan(out1.values[0, -1, -1])
        # assert (np.isnan(dis.values[0, -1, -1]))


class TestMaxPrIntensity:
    def test_simple(self, pr_hr_series):
        pr1 = pr_hr_series(np.zeros(366 * 24))
        pr1[10:20] += np.arange(10)
        pr2 = pr_hr_series(np.ones(366 * 24))

        pr = xr.concat([pr1, pr2], dim="site")
        out = atmos.max_pr_intensity(pr, window=2, freq="YS")
        np.testing.assert_array_almost_equal(out.isel(time=0), [8.5 * 3600, 3600])

    def test_indexing(self, pr_hr_series):
        pr = pr_hr_series(np.zeros(366 * 24))
        pr[20:30] += np.arange(10)

        out1 = atmos.max_pr_intensity(pr, window=2, freq="YS")
        out2 = atmos.max_pr_intensity(pr, window=2, freq="YS", doy_bounds=[200, 1])
        np.testing.assert_array_almost_equal(out1.isel(time=0), [8.5 * 3600])
        np.testing.assert_array_almost_equal(out2.isel(time=0), [2.5 * 3600])


class TestMax1Day:
    # testing of wet_day and daily_pr_intensity, both are related

    nc_file = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"

    def test_3d_data_with_nans(self, open_dataset):
        # test with 3d data
        pr = open_dataset(self.nc_file).pr
        prMM = open_dataset(self.nc_file).pr
        prMM.values *= 86400.0
        prMM.attrs["units"] = "mm/day"
        # put a nan somewhere
        prMM.values[10, 1, 0] = np.nan
        pr.values[10, 1, 0] = np.nan

        out1 = atmos.max_1day_precipitation_amount(pr, freq="MS")
        out2 = atmos.max_1day_precipitation_amount(prMM, freq="MS")

        # test kg m-2 s-1
        pr.attrs["units"] = "kg m-2 s-1"
        out3 = atmos.max_1day_precipitation_amount(pr, freq="MS")

        np.testing.assert_array_almost_equal(out1, out2, 3)
        np.testing.assert_array_almost_equal(out1, out3, 3)

        x1 = prMM[:31, 0, 0].values
        rx1 = x1.max()

        assert np.allclose(rx1, out1.values[0, 0, 0])
        assert np.isnan(out1.values[0, 1, 0])
        assert np.isnan(out1.values[0, -1, -1])


class TestMaxNDay:
    # testing of wet_day and daily_pr_intensity, both are related

    nc_file = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"

    @pytest.mark.parametrize(
        "units,factor,chunks",
        [
            ("mm/day", 86400.0, None),
            ("kg m-2 s-1", 1, None),
            ("mm/s", 1, {"time": 73.0}),
        ],
    )
    def test_3d_data_with_nans(self, open_dataset, units, factor, chunks):
        # test with 3d data
        pr1 = open_dataset(self.nc_file).pr
        pr2 = open_dataset(self.nc_file, chunks=chunks).pr
        pr2.values *= factor
        pr2.attrs["units"] = units
        # put a nan somewhere
        pr2.values[10, 1, 0] = np.nan
        pr1.values[10, 1, 0] = np.nan
        wind = 3
        out1 = atmos.max_n_day_precipitation_amount(pr1, window=wind, freq="MS")
        out2 = atmos.max_n_day_precipitation_amount(pr2, window=wind, freq="MS")

        np.testing.assert_array_almost_equal(out1, out2, 3)

        x1 = pr1[:31, 0, 0].values * 86400
        df = pd.DataFrame({"pr": x1})
        rx3 = df.rolling(wind).sum().max()

        assert np.allclose(rx3, out1.values[0, 0, 0])
        assert np.isnan(out1.values[0, 1, 0])
        assert np.isnan(out1.values[0, -1, -1])

    def test_indexing(self, open_dataset):
        pr = open_dataset(self.nc_file).pr
        out1 = atmos.max_n_day_precipitation_amount(pr, window=7, freq="YS")
        out2 = atmos.max_n_day_precipitation_amount(pr, window=7, freq="YS", month=[6, 7, 8])

        np.testing.assert_allclose(out1.values[0, 0, 0], 55.48, atol=1e-2)
        np.testing.assert_allclose(out2.values[0, 0, 0], 50.57, atol=1e-2)


class TestMaxConsecWetDays:
    nc_file = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"

    def test_3d_data_with_nans(self, open_dataset):
        # test with 3d data
        pr = open_dataset(self.nc_file).pr
        prMM = open_dataset(self.nc_file).pr
        prMM.values *= 86400.0
        prMM.attrs["units"] = "mm/day"
        # put a nan somewhere
        prMM.values[10, 1, 0] = np.nan
        pr.values[10, 1, 0] = np.nan
        pr_min = "5 mm/d"
        out1 = atmos.maximum_consecutive_wet_days(pr, thresh=pr_min, freq="MS")
        out2 = atmos.maximum_consecutive_wet_days(prMM, thresh=pr_min, freq="MS")

        # test kg m-2 s-1
        pr.attrs["units"] = "kg m-2 s-1"
        out3 = atmos.maximum_consecutive_wet_days(pr, thresh=pr_min, freq="MS")

        np.testing.assert_array_equal(out1, out2)
        np.testing.assert_array_equal(out1, out3)

        # check some vector with and without a nan
        x1 = prMM[:31, 0, 0] * 0.0
        x1[5:10] = 10
        x1.attrs["units"] = "mm/day"
        cwd1 = atmos.maximum_consecutive_wet_days(x1, freq="MS")

        assert cwd1 == 5

        assert np.isnan(out1.values[0, 1, 0])

        # make sure that vector with all nans gives nans whatever skipna
        assert np.isnan(out1.values[0, -1, -1])


class TestMaxConsecDryDays:
    nc_file = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"

    def test_3d_data_with_nans(self, open_dataset):
        # test with 3d data
        pr = open_dataset(self.nc_file).pr
        prMM = open_dataset(self.nc_file).pr
        prMM.values *= 86400.0
        prMM.attrs["units"] = "mm/day"
        # put a nan somewhere
        prMM.values[10, 1, 0] = np.nan
        pr.values[10, 1, 0] = np.nan
        pr_min = "5 mm/d"
        out1 = atmos.maximum_consecutive_dry_days(pr, thresh=pr_min, freq="MS")
        out2 = atmos.maximum_consecutive_dry_days(prMM, thresh=pr_min, freq="MS")

        # test kg m-2 s-1
        pr.attrs["units"] = "kg m-2 s-1"
        out3 = atmos.maximum_consecutive_dry_days(pr, thresh=pr_min, freq="MS")

        np.testing.assert_array_equal(out1, out2)
        np.testing.assert_array_equal(out1, out3)

        # check some vector with and without a nan
        x1 = prMM[:31, 0, 0] * 0.0 + 50.0
        x1[5:10] = 0
        x1.attrs["units"] = "mm/day"
        cdd1 = atmos.maximum_consecutive_dry_days(x1, freq="MS")

        assert cdd1 == 5

        assert np.isnan(out1.values[0, 1, 0])

        # make sure that vector with all nans gives nans whatever skipna
        assert np.isnan(out1.values[0, -1, -1])


class TestSnowfallDate:
    tasmin_file = "NRCANdaily/nrcan_canada_daily_tasmin_1990.nc"
    pr_file = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"

    @classmethod
    def get_snowfall(cls, nimbus):
        dnr = xr.merge(
            (
                xr.open_dataset(nimbus.fetch(cls.pr_file), engine="h5netcdf"),
                xr.open_dataset(nimbus.fetch(cls.tasmin_file), engine="h5netcdf"),
            )
        )
        return convert.snowfall_approximation(dnr.pr, tas=dnr.tasmin, thresh="-0.5 degC", method="binary")

    def test_first_snowfall(self, nimbus):
        with set_options(check_missing="skip"):
            fs = atmos.first_snowfall(prsn=self.get_snowfall(nimbus), thresh="0.5 mm/day")

        np.testing.assert_array_equal(
            fs[:, [0, 45, 82], [10, 105, 155]],
            np.array(
                [
                    [[1, 1, 1], [1, 1, 1], [11, np.nan, np.nan]],
                    [[254, 256, 277], [274, 292, 275], [300, np.nan, np.nan]],
                ]
            ),
        )

    def test_last_snowfall(self, nimbus):
        with set_options(check_missing="skip"):
            ls = atmos.last_snowfall(prsn=self.get_snowfall(nimbus), thresh="0.5 mm/day")

        np.testing.assert_array_equal(
            ls[:, [0, 45, 82], [10, 105, 155]],
            np.array(
                [
                    [[155, 151, 129], [127, 157, 110], [106, np.nan, np.nan]],
                    [[365, 363, 363], [365.0, 364, 364], [362, np.nan, np.nan]],
                ]
            ),
        )


class TestDaysWithSnow:
    def test_simple(self, open_dataset, prsn_series):
        prsn = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc").prsn
        out = atmos.days_with_snow(prsn, low="0 kg m-2 s-1", high="1e6 kg m-2 s-1")
        np.testing.assert_array_equal(out[1], [np.nan, 162, 159, 126, np.nan])


def test_days_over_precip_doy_thresh(open_dataset):
    pr = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc").pr
    per = percentile_doy(pr, window=5, per=80)

    out1 = atmos.days_over_precip_doy_thresh(pr, per)
    np.testing.assert_array_equal(out1[1, :, 0], np.array([81, 60, 68, 78]))

    out2 = atmos.days_over_precip_doy_thresh(pr, per, thresh="2 mm/d")
    np.testing.assert_array_equal(out2[1, :, 0], np.array([80, 59, 66, 78]))

    assert "only days with at least 2 mm/d are counted." in out2.description
    assert "[80]th percentile" in out2.attrs["description"]
    assert "['1990-01-01', '1993-12-31'] period" in out2.attrs["description"]
    assert "5 day(s)" in out2.attrs["description"]


def test_days_over_precip_thresh(open_dataset):
    pr = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc").pr
    per = pr.quantile(0.8, "time", keep_attrs=True)
    per.attrs["climatology_bounds"] = build_climatology_bounds(pr)

    out = atmos.days_over_precip_thresh(pr, per)

    np.testing.assert_allclose(out[1, :], np.array([80.0, 64.0, 65.0, 83.0]), atol=0.001)
    assert "80.0th percentile" in out.attrs["description"]
    assert "['1990-01-01', '1993-12-31'] period" in out.attrs["description"]


def test_days_over_precip_thresh__seasonal_indexer(open_dataset):
    # GIVEN
    pr = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc").pr
    per = pr.quantile(0.8, "time", keep_attrs=True)
    # WHEN
    out = atmos.days_over_precip_thresh(pr, per, freq="YS", date_bounds=("01-10", "12-31"))
    # THEN
    np.testing.assert_almost_equal(out[0], np.array([81.0, 66.0, 66.0, 75.0]))


def test_fraction_over_precip_doy_thresh(open_dataset):
    pr = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc").pr
    per = percentile_doy(pr, window=5, per=80)

    out = atmos.fraction_over_precip_doy_thresh(pr, per)
    np.testing.assert_allclose(out[1, :, 0], np.array([0.803, 0.747, 0.745, 0.806]), atol=0.001)

    out = atmos.fraction_over_precip_doy_thresh(pr, per, thresh="0.002 m/d")
    np.testing.assert_allclose(out[1, :, 0], np.array([0.822, 0.780, 0.771, 0.829]), atol=0.001)

    assert "only days with at least 0.002 m/d are included" in out.description
    assert "[80]th percentile" in out.attrs["description"]
    assert "['1990-01-01', '1993-12-31'] period" in out.attrs["description"]
    assert "5 day(s)" in out.attrs["description"]


def test_fraction_over_precip_thresh(open_dataset):
    pr = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc").pr
    per = pr.quantile(0.8, "time", keep_attrs=True)
    per.attrs["climatology_bounds"] = build_climatology_bounds(pr)

    out = atmos.fraction_over_precip_thresh(pr, per)

    np.testing.assert_allclose(out[1, :], np.array([0.839, 0.812, 0.776, 0.864]), atol=0.001)

    assert "80.0th percentile" in out.attrs["description"]
    assert "['1990-01-01', '1993-12-31'] period" in out.attrs["description"]


def test_liquid_precip_ratio(open_dataset):
    ds = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc")

    out = atmos.liquid_precip_ratio(pr=ds.pr, tas=ds.tas, thresh="0 degC", freq="YS")
    np.testing.assert_allclose(out[:, 0], np.array([0.919, 0.805, 0.525, 0.740, 0.993]), atol=1e3)

    with set_options(cf_compliance="raise"):
        # Test if tasmax is allowed
        out = atmos.liquid_precip_ratio(pr=ds.pr, tas=ds.tasmax, thresh="33 degF", freq="YS")
        np.testing.assert_allclose(out[:, 0], np.array([0.975, 0.921, 0.547, 0.794, 0.999]), atol=1e3)
        assert "where temperature is above 33 degf." in out.description


def test_dry_spell(atmosds):
    pr = atmosds.pr

    events = atmos.dry_spell_frequency(pr, thresh="3 mm", window=7, freq="YS")
    total_d_sum = atmos.dry_spell_total_length(pr, thresh="3 mm", window=7, op="sum", freq="YS")
    total_d_max = atmos.dry_spell_total_length(pr, thresh="3 mm", window=7, op="max", freq="YS")
    max_d_sum = atmos.dry_spell_max_length(pr, thresh="3 mm", window=7, op="sum", freq="YS")
    max_d_max = atmos.dry_spell_max_length(pr, thresh="3 mm", window=7, op="max", freq="YS")
    total_d_sum = total_d_sum.sel(location="Halifax", drop=True).isel(time=slice(0, 2))
    total_d_max = total_d_max.sel(location="Halifax", drop=True).isel(time=slice(0, 2))
    max_d_sum = max_d_sum.sel(location="Halifax", drop=True).isel(time=slice(0, 2))
    max_d_max = max_d_max.sel(location="Halifax", drop=True).isel(time=slice(0, 2))

    np.testing.assert_allclose(events[0:2, 0], [5, 7], rtol=1e-1)
    np.testing.assert_allclose(total_d_sum, [50, 53], rtol=1e-1)
    np.testing.assert_allclose(total_d_max, [68, 97], rtol=1e-1)
    np.testing.assert_allclose(max_d_sum, [14, 10], rtol=1e-1)
    np.testing.assert_allclose(max_d_max, [14, 14], rtol=1e-1)

    assert (
        "The annual number of dry periods of 7 day(s) or more, "
        "during which the total precipitation on a window of 7 day(s) is below 3 mm."
    ) in events.description
    assert "The annual number of days in dry periods of 7 day(s) or more" in total_d_sum.description
    assert "The annual number of days in dry periods of 7 day(s) or more" in total_d_max.description
    assert "The maximum annual number of consecutive days in a dry period of 7 day(s) or more" in max_d_sum.description
    assert "The maximum annual number of consecutive days in a dry period of 7 day(s) or more" in max_d_max.description


def test_dry_spell_total_length_indexer(pr_series):
    pr = pr_series([np.nan] + [1] * 4 + [0] * 10 + [1] * 350, start="1900-01-01", units="mm/d")
    out = atmos.dry_spell_total_length(
        pr,
        window=7,
        op="sum",
        thresh="3.1 mm",
        freq="MS",
    )
    np.testing.assert_allclose(out, [np.nan] + [0] * 11)

    out = atmos.dry_spell_total_length(
        pr, window=7, op="sum", thresh="3.1 mm", freq="MS", date_bounds=("01-10", "12-31")
    )
    np.testing.assert_allclose(out, [9] + [0] * 11)


def test_dry_spell_max_length_indexer(pr_series):
    pr = pr_series([np.nan] + [1] * 4 + [0] * 10 + [1] * 350, start="1900-01-01", units="mm/d")
    out = atmos.dry_spell_max_length(
        pr,
        window=7,
        op="sum",
        thresh="3.1 mm",
        freq="MS",
    )
    np.testing.assert_allclose(out, [np.nan] + [0] * 11)

    out = atmos.dry_spell_total_length(
        pr, window=7, op="sum", thresh="3.1 mm", freq="MS", date_bounds=("01-10", "12-31")
    )
    np.testing.assert_allclose(out, [9] + [0] * 11)


def test_dry_spell_frequency_op(open_dataset):
    pr = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc").pr
    test_sum = atmos.dry_spell_frequency(pr, thresh="3 mm", window=7, freq="MS", op="sum")
    test_max = atmos.dry_spell_frequency(pr, thresh="3 mm", window=7, freq="MS", op="max")

    np.testing.assert_allclose(test_sum[0, :14], [1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0], rtol=1e-1)
    np.testing.assert_allclose(test_max[0, :14], [1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 2, 1], rtol=1e-1)
    assert (
        "The monthly number of dry periods of 7 day(s) or more, "
        "during which the total precipitation on a window of 7 day(s) is below 3 mm."
    ) in test_sum.description
    assert (
        "The monthly number of dry periods of 7 day(s) or more, "
        "during which the maximal precipitation on a window of 7 day(s) is below 3 mm."
    ) in test_max.description


class TestSnowfallMeteoSwiss:
    tasmin_file = "NRCANdaily/nrcan_canada_daily_tasmin_1990.nc"
    pr_file = "NRCANdaily/nrcan_canada_daily_pr_1990.nc"

    @classmethod
    def get_snowfall(cls, nimbus):
        dnr = xr.merge(
            (
                xr.open_dataset(nimbus.fetch(cls.pr_file), engine="h5netcdf"),
                xr.open_dataset(nimbus.fetch(cls.tasmin_file), engine="h5netcdf"),
            )
        )
        return convert.snowfall_approximation(dnr.pr, tas=dnr.tasmin, thresh="-0.5 degC", method="binary")

    def test_snowfall_frequency(self, nimbus):
        prsn = self.get_snowfall(nimbus)
        with set_options(check_missing="skip"):
            sf = atmos.snowfall_frequency(prsn=prsn, thresh="1 mm/day")
        expected = np.array(
            [
                [
                    [27.624, 29.834, 25.414],
                    [22.652, 25.414, 22.652],
                    [12.155, 0.0, 0.0],
                ],
                [
                    [23.913, 23.370, 20.652],
                    [17.391, 15.761, 13.043],
                    [4.891, 0.0, 0.0],
                ],
            ]
        )
        np.testing.assert_allclose(
            sf[:, [0, 45, 82], [10, 105, 155]],
            expected,
            rtol=1e-3,
        )

    def test_snowfall_intensity(self, nimbus):
        prsn = self.get_snowfall(nimbus)
        with set_options(check_missing="skip"):
            si = atmos.snowfall_intensity(prsn=prsn, thresh="1 mm/day")
        expected = np.array(
            [
                [
                    [3.585, 3.839, 4.446],
                    [5.148, 5.884, 5.764],
                    [5.910, 0.0, 0.0],
                ],
                [
                    [5.093, 5.284, 5.539],
                    [5.946, 6.840, 9.702],
                    [9.522, 0.0, 0.0],
                ],
            ]
        )
        np.testing.assert_allclose(
            si[:, [0, 45, 82], [10, 105, 155]],
            expected,
            rtol=1e-3,
        )
