from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from xarray import CFTimeIndex, DataArray, Dataset, infer_freq
from xarray.coding.calendar_ops import convert_calendar, interp_calendar
from xarray.coding.cftime_offsets import date_range
from xarray.testing import assert_identical
from xarray.tests import requires_cftime

cftime = pytest.importorskip("cftime")


@pytest.mark.parametrize(
    "source, target, use_cftime, freq",
    [
        ("standard", "noleap", None, "D"),
        ("noleap", "proleptic_gregorian", True, "D"),
        ("noleap", "all_leap", None, "D"),
        ("all_leap", "proleptic_gregorian", False, "4h"),
    ],
)
def test_convert_calendar(source, target, use_cftime, freq):
    src = DataArray(
        date_range("2004-01-01", "2004-12-31", freq=freq, calendar=source),
        dims=("time",),
        name="time",
    )
    da_src = DataArray(
        np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
    )

    conv = convert_calendar(da_src, target, use_cftime=use_cftime)

    assert conv.time.dt.calendar == target

    if source != "noleap":
        expected_times = date_range(
            "2004-01-01",
            "2004-12-31",
            freq=freq,
            use_cftime=use_cftime,
            calendar=target,
        )
    else:
        expected_times_pre_leap = date_range(
            "2004-01-01",
            "2004-02-28",
            freq=freq,
            use_cftime=use_cftime,
            calendar=target,
        )
        expected_times_post_leap = date_range(
            "2004-03-01",
            "2004-12-31",
            freq=freq,
            use_cftime=use_cftime,
            calendar=target,
        )
        expected_times = expected_times_pre_leap.append(expected_times_post_leap)
    np.testing.assert_array_equal(conv.time, expected_times)


def test_convert_calendar_dataset():
    # Check that variables without a time dimension are not modified
    src = DataArray(
        date_range("2004-01-01", "2004-12-31", freq="D", calendar="standard"),
        dims=("time",),
        name="time",
    )
    da_src = DataArray(
        np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
    ).expand_dims(lat=[0, 1])
    ds_src = Dataset({"hastime": da_src, "notime": (("lat",), [0, 1])})

    conv = convert_calendar(ds_src, "360_day", align_on="date")

    assert conv.time.dt.calendar == "360_day"
    assert_identical(ds_src.notime, conv.notime)


@pytest.mark.parametrize(
    "source,target,freq",
    [
        ("standard", "360_day", "D"),
        ("360_day", "proleptic_gregorian", "D"),
        ("proleptic_gregorian", "360_day", "4h"),
    ],
)
@pytest.mark.parametrize("align_on", ["date", "year"])
def test_convert_calendar_360_days(source, target, freq, align_on):
    src = DataArray(
        date_range("2004-01-01", "2004-12-30", freq=freq, calendar=source),
        dims=("time",),
        name="time",
    )
    da_src = DataArray(
        np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
    )

    conv = convert_calendar(da_src, target, align_on=align_on)

    assert conv.time.dt.calendar == target

    if align_on == "date":
        np.testing.assert_array_equal(
            conv.time.resample(time="ME").last().dt.day,
            [30, 29, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30],
        )
    elif target == "360_day":
        np.testing.assert_array_equal(
            conv.time.resample(time="ME").last().dt.day,
            [30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 29],
        )
    else:
        np.testing.assert_array_equal(
            conv.time.resample(time="ME").last().dt.day,
            [30, 29, 30, 30, 31, 30, 30, 31, 30, 31, 29, 31],
        )
    if source == "360_day" and align_on == "year":
        assert conv.size == 360 if freq == "D" else 360 * 4
    else:
        assert conv.size == 359 if freq == "D" else 359 * 4


def test_convert_calendar_360_days_random():
    da_std = DataArray(
        np.linspace(0, 1, 366),
        dims=("time",),
        coords={
            "time": date_range(
                "2004-01-01",
                "2004-12-31",
                freq="D",
                calendar="standard",
                use_cftime=False,
            )
        },
    )
    da_360 = DataArray(
        np.linspace(0, 1, 360),
        dims=("time",),
        coords={
            "time": date_range("2004-01-01", "2004-12-30", freq="D", calendar="360_day")
        },
    )

    conv = convert_calendar(da_std, "360_day", align_on="random")
    conv2 = convert_calendar(da_std, "360_day", align_on="random")
    assert (conv != conv2).any()

    conv = convert_calendar(da_360, "standard", use_cftime=False, align_on="random")
    assert np.datetime64("2004-02-29") not in conv.time
    conv2 = convert_calendar(da_360, "standard", use_cftime=False, align_on="random")
    assert (conv2 != conv).any()

    # Ensure that added days are evenly distributed in the 5 fifths of each year
    conv = convert_calendar(da_360, "noleap", align_on="random", missing=np.nan)
    conv = conv.where(conv.isnull(), drop=True)
    nandoys = conv.time.dt.dayofyear[:366]
    assert all(nandoys < np.array([74, 147, 220, 293, 366]))
    assert all(nandoys > np.array([0, 73, 146, 219, 292]))


@requires_cftime
@pytest.mark.parametrize(
    "source,target,freq",
    [
        ("standard", "noleap", "D"),
        ("noleap", "proleptic_gregorian", "4h"),
        ("noleap", "all_leap", "ME"),
        ("360_day", "noleap", "D"),
        ("noleap", "360_day", "D"),
    ],
)
def test_convert_calendar_missing(source, target, freq):
    src = DataArray(
        date_range(
            "2004-01-01",
            "2004-12-31" if source != "360_day" else "2004-12-30",
            freq=freq,
            calendar=source,
        ),
        dims=("time",),
        name="time",
    )
    da_src = DataArray(
        np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
    )
    out = convert_calendar(da_src, target, missing=np.nan, align_on="date")

    expected_freq = freq
    assert infer_freq(out.time) == expected_freq

    expected = date_range(
        "2004-01-01",
        "2004-12-31" if target != "360_day" else "2004-12-30",
        freq=freq,
        calendar=target,
    )
    np.testing.assert_array_equal(out.time, expected)

    if freq != "ME":
        out_without_missing = convert_calendar(da_src, target, align_on="date")
        expected_nan = out.isel(time=~out.time.isin(out_without_missing.time))
        assert expected_nan.isnull().all()

        expected_not_nan = out.sel(time=out_without_missing.time)
        assert_identical(expected_not_nan, out_without_missing)


@requires_cftime
def test_convert_calendar_errors():
    src_nl = DataArray(
        date_range("0000-01-01", "0000-12-31", freq="D", calendar="noleap"),
        dims=("time",),
        name="time",
    )
    # no align_on for conversion to 360_day
    with pytest.raises(ValueError, match="Argument `align_on` must be specified"):
        convert_calendar(src_nl, "360_day")

    # Standard doesn't support year 0
    with pytest.raises(
        ValueError, match="Source time coordinate contains dates with year 0"
    ):
        convert_calendar(src_nl, "standard")

    # no align_on for conversion from 360 day
    src_360 = convert_calendar(src_nl, "360_day", align_on="year")
    with pytest.raises(ValueError, match="Argument `align_on` must be specified"):
        convert_calendar(src_360, "noleap")

    # Datetime objects
    da = DataArray([0, 1, 2], dims=("x",), name="x")
    with pytest.raises(ValueError, match="Coordinate x must contain datetime objects."):
        convert_calendar(da, "standard", dim="x")


def test_convert_calendar_dimension_name():
    src = DataArray(
        date_range("2004-01-01", "2004-01-31", freq="D", calendar="noleap"),
        dims=("date",),
        name="date",
    )

    out = convert_calendar(src, "proleptic_gregorian", dim="date")

    np.testing.assert_array_equal(src, out)


def test_convert_calendar_same_calendar():
    src = DataArray(
        date_range("2000-01-01", periods=12, freq="6h", use_cftime=False),
        dims=("time",),
        name="time",
    )
    out = convert_calendar(src, "proleptic_gregorian")
    assert src is out


@pytest.mark.parametrize(
    "source,target",
    [
        ("standard", "noleap"),
        ("noleap", "proleptic_gregorian"),
        ("standard", "360_day"),
        ("360_day", "proleptic_gregorian"),
        ("noleap", "all_leap"),
        ("360_day", "noleap"),
    ],
)
def test_interp_calendar(source, target):
    src = DataArray(
        date_range("2004-01-01", "2004-07-30", freq="D", calendar=source),
        dims=("time",),
        name="time",
    )
    tgt = DataArray(
        date_range("2004-01-01", "2004-07-30", freq="D", calendar=target),
        dims=("time",),
        name="time",
    )
    da_src = DataArray(
        np.linspace(0, 1, src.size), dims=("time",), coords={"time": src}
    )
    conv = interp_calendar(da_src, tgt)

    assert_identical(tgt.time, conv.time)

    np.testing.assert_almost_equal(conv.max(), 1, 2)
    assert conv.min() == 0


@requires_cftime
def test_interp_calendar_errors():
    src_nl = DataArray(
        [1] * 100,
        dims=("time",),
        coords={
            "time": date_range("0000-01-01", periods=100, freq="MS", calendar="noleap")
        },
    )
    tgt_360 = date_range("0001-01-01", "0001-12-30", freq="MS", calendar="standard")

    with pytest.raises(
        ValueError, match="Source time coordinate contains dates with year 0"
    ):
        interp_calendar(src_nl, tgt_360)

    da1 = DataArray([0, 1, 2], dims=("x",), name="x")
    da2 = da1 + 1

    with pytest.raises(
        ValueError, match="Both 'source.x' and 'target' must contain datetime objects."
    ):
        interp_calendar(da1, da2, dim="x")


@requires_cftime
@pytest.mark.parametrize(
    ("source_calendar", "target_calendar", "expected_index"),
    [("standard", "noleap", CFTimeIndex), ("all_leap", "standard", pd.DatetimeIndex)],
)
def test_convert_calendar_produces_time_index(
    source_calendar, target_calendar, expected_index
):
    # https://github.com/pydata/xarray/issues/9138
    time = date_range("2000-01-01", "2002-01-01", freq="D", calendar=source_calendar)
    temperature = np.ones(len(time))
    da = DataArray(
        data=temperature,
        dims=["time"],
        coords=dict(
            time=time,
        ),
    )
    converted = da.convert_calendar(target_calendar)
    assert isinstance(converted.indexes["time"], expected_index)