import warnings

import numpy as np
import pytest
from numpy.testing import assert_allclose, assert_equal

import openamundsen as oa

from .compare import compare_datasets
from .conftest import base_config

single_point_results_all = [
    "single_point_results_multilayer",
    "single_point_results_cryolayers",
]


def base_config_snow():
    config = base_config()
    config.start_date = "2019-10-01"
    config.end_date = "2020-05-31"
    config.output_data.timeseries.variables = [
        {"var": "snow.num_layers", "name": "num_snow_layers"},
        {"var": "snow.albedo", "name": "snow_albedo"},
    ]
    return config


@pytest.fixture(scope="session")
def multilayer_run():
    config = base_config_snow()
    config.snow.model = "multilayer"
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run()
    return model


@pytest.fixture(scope="session")
def cryolayer_run():
    config = base_config_snow()
    config.snow.model = "cryolayers"
    config.output_data.timeseries.variables.append({"var": "snow.cold_content"})
    config.output_data.timeseries.variables.append({"var": "snow.layer_albedo"})
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run()
    return model


@pytest.fixture(scope="module")
def point_results_cryolayers():
    config = base_config()
    config.snow.model = "cryolayers"
    config.output_data.timeseries.variables.append({"var": "snow.cold_content"})
    config.output_data.timeseries.variables.append({"var": "snow.layer_albedo"})
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run()
    return model.point_output.data


@pytest.fixture(scope="function")
def single_point_results_multilayer(multilayer_run):
    return multilayer_run.point_output.data.sel(point="proviantdepot")


@pytest.fixture(scope="function")
def single_point_results_cryolayers(cryolayer_run):
    return cryolayer_run.point_output.data.sel(point="proviantdepot")


def test_default_multilayer():
    config = base_config_snow()
    assert config.snow.model == "multilayer"


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_swe(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    swe = ds.swe.values
    ice_content = ds.ice_content.values
    liquid_water_content = ds.liquid_water_content.values
    assert np.all(np.isfinite(swe))
    assert swe.min() == 0.0
    assert_allclose(
        swe,
        (ice_content + liquid_water_content).sum(axis=1),
        atol=0.1,
    )
    assert np.all(ice_content >= 0.0)
    assert np.all(liquid_water_content >= 0.0)


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_density(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    swe3d = (ds.ice_content + ds.liquid_water_content).values
    density = ds.snow_density.values
    assert np.all(np.isnan(density[swe3d == 0.0]))
    assert np.all(np.isfinite(density[swe3d > 0.0]))
    assert not np.any(density < 0.0)
    assert not np.any(density > 1000.0)


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_depth(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    depth = ds.snow_depth.values
    swe = ds.swe.values
    assert np.all(depth[swe == 0.0] == 0.0)
    assert np.all(depth[swe > 0.0] > 0.0)
    assert not np.any(depth < 0.0)

    swe3d = (ds.ice_content + ds.liquid_water_content).values
    thickness = ds.snow_thickness.values
    assert_allclose(depth, thickness.sum(axis=1), atol=1e-3)

    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            "invalid value encountered in divide",
            RuntimeWarning,
        )
        density_calc = swe3d / thickness

    density = ds.snow_density.values
    assert_allclose(density_calc, density, atol=1.0)


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_num_layers(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    num_layers = ds.num_snow_layers.values
    thickness = ds.snow_thickness.values
    assert_equal(num_layers, (thickness > 0).sum(axis=1))


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_albedo(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    config = base_config()
    albedo = ds.snow_albedo.values
    swe = ds.swe.values
    min_albedo = config.snow.albedo.min
    max_albedo = config.snow.albedo.max
    pos_snow = swe > 0.0
    assert np.all(albedo[pos_snow] >= np.float32(min_albedo))
    assert np.all(albedo[pos_snow] <= np.float32(max_albedo))
    assert np.all(np.isnan(albedo[~pos_snow]))


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_melt(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    assert np.all(ds.melt >= 0)


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_runoff(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    assert np.all(ds.runoff >= 0)


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_refreezing(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    assert np.all(ds.refreezing >= 0)


@pytest.mark.slow
def test_cryolayers_layer_albedo(single_point_results_cryolayers):
    config = base_config()
    ds = single_point_results_cryolayers
    layer_albedo = ds.layer_albedo.values
    swe3d = (ds.ice_content + ds.liquid_water_content).values
    min_albedo = config.snow.albedo.min
    max_albedo = config.snow.albedo.max
    pos_snow = swe3d > 0.0
    assert np.all(layer_albedo[pos_snow] >= np.float32(min_albedo))
    assert np.all(layer_albedo[pos_snow] <= np.float32(max_albedo))
    assert np.all(np.isnan(layer_albedo[~pos_snow]))


@pytest.mark.slow
def test_cryolayers_cold_content(single_point_results_cryolayers):
    config = base_config()
    ds = single_point_results_cryolayers
    cc = ds.cold_content.values
    ic = ds.ice_content.values
    lwc = ds.liquid_water_content.values
    melt = ds.melt.values
    sublimation = ds.sublimation.values
    swe3d = (ds.ice_content + ds.liquid_water_content).values
    cold_holding_capacity = config.snow.cryolayers.cold_holding_capacity
    pos_snow = swe3d > 0
    assert np.all(cc >= 0.0)
    assert cc.max() > 0.0
    assert np.all(cc[~pos_snow] == 0.0)

    # Check if the CC is <= swe * cold_colding_capacity (add melt and sublimation because the
    # maximum CC is calculated before reducing SWE by melt and sublimation)
    assert np.all(
        (((ic + lwc).sum(axis=1) + melt + sublimation) * cold_holding_capacity - cc.sum(axis=1))
        >= -0.1
    )


@pytest.mark.slow
@pytest.mark.parametrize("method", ["temperature_index", "enhanced_temperature_index"])
def test_temperature_index(method):
    config = base_config_snow()
    config.snow.model = "cryolayers"
    config.timestep = "D"
    config.output_data.timeseries.variables = [{"var": "snow.cold_content"}]
    config.snow.melt.method = method

    if method == "temperature_index":
        config.snow.melt.degree_day_factor = 5.0
    elif method == "enhanced_temperature_index":
        config.snow.melt.degree_day_factor = 3.0
        config.snow.melt.albedo_factor = 0.1

    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run()

    ds = model.point_output.data.sel(point="proviantdepot")
    assert np.all(ds.melt >= 0)
    assert np.any(ds.melt > 0)
    assert np.all(ds.melt[ds.temp <= config.snow.melt.degree_day_factor] == 0.0)
    melt_potential = (
        ds.ice_content.sum("snow_layer").shift(time=1)
        + ds.cold_content.sum("snow_layer").shift(time=1)
        + ds.snowfall
    )
    assert np.all(ds.melt <= melt_potential.fillna(np.inf) + 1e-6)


@pytest.mark.slow
@pytest.mark.comparison
def test_compare_multilayer(multilayer_run):
    compare_datasets(
        "snow_multilayer_point",
        multilayer_run.point_output.data,
        point="proviantdepot",
    )


@pytest.mark.slow
@pytest.mark.comparison
def test_compare_cryolayers(cryolayer_run):
    compare_datasets(
        "snow_cryolayers_point",
        cryolayer_run.point_output.data,
        point="proviantdepot",
    )


@pytest.mark.slow
@pytest.mark.parametrize("fixture_name", single_point_results_all)
def test_mass_balance(fixture_name, request):
    ds = request.getfixturevalue(fixture_name)
    ds_prev = ds.shift(time=1)

    ic1 = ds.ice_content.sum("snow_layer")
    ic2 = (
        ds_prev.ice_content.sum("snow_layer")
        + ds.snowfall
        - ds.melt
        - ds.sublimation
        + ds.refreezing
    )
    assert_allclose(ic1, ic2, atol=1e-4)

    lwc1 = ds.liquid_water_content.sum("snow_layer")
    lwc2 = (
        ds_prev.liquid_water_content.sum("snow_layer")
        + ds.rainfall
        + ds.melt
        - ds.runoff
        - ds.refreezing
    )
    assert_allclose(lwc1, lwc2, atol=1e-4)
