import pytest
import numpy as np
import pandas as pd
import xarray as xr
from xarray.tests import assert_allclose

from xskillscore.core.deterministic import (
    pearson_r, rmse, mae)
from xskillscore.core.probabilistic import (
    xr_crps_gaussian)


@pytest.fixture
def ds_dask():
    dates = pd.date_range('1/1/2000', '1/3/2000', freq='D')
    lats = np.arange(4)
    lons = np.arange(5)
    a = xr.DataArray(
        np.random.rand(len(dates), len(lats), len(lons)),
        coords=[dates, lats, lons],
        dims=['time', 'lat', 'lon']
    ).chunk()
    b = xr.DataArray(
        np.random.rand(len(dates), len(lats), len(lons)),
        coords=[dates, lats, lons],
        dims=['time', 'lat', 'lon']
    ).chunk()

    ds = xr.Dataset()
    ds['a'] = a
    ds['b'] = b
    return ds


@pytest.fixture
def mu():
    return xr.DataArray(5)


@pytest.fixture
def sigma():
    return xr.DataArray(6)


def test_pearson_r_accessor(ds_dask):
    ds = ds_dask.load()

    dim = 'time'
    actual = pearson_r(ds['a'], ds['b'], dim)
    expected = ds.xs.pearson_r('a', 'b', dim)
    assert_allclose(actual, expected)


def test_rmse_accessor_dask(ds_dask):
    dim = 'lon'
    actual = rmse(ds_dask['a'], ds_dask['b'], dim).compute()
    expected = ds_dask.xs.rmse('a', 'b', dim).compute()
    assert_allclose(actual, expected)


def test_mae_accessor_outer_array(ds_dask):
    ds = ds_dask.load()
    b = ds['b']
    ds = ds.drop('b')
    dim = 'lat'

    actual = mae(ds['a'], b, dim)
    expected = ds.xs.mae('a', b, dim)
    assert_allclose(actual, expected)


def test_crps_gaussian_accessor(ds_dask, mu, sigma):
    ds = ds_dask.load()
    ds['mu'] = mu
    ds['sigma'] = sigma

    actual = xr_crps_gaussian(ds['a'], mu, sigma)
    expected = ds.xs.crps_gaussian('a', 'mu', 'sigma')
    assert_allclose(actual, expected)


def test_crps_gaussian_accessor_outer_array(ds_dask, mu, sigma):
    ds = ds_dask.load()
    actual = xr_crps_gaussian(ds['a'], mu, sigma)
    expected = ds.xs.crps_gaussian('a', mu, sigma)
    assert_allclose(actual, expected)
