from .conftest import base_config
import numpy as np
from numpy.testing import assert_allclose, assert_array_equal
import openamundsen as oa
import openamundsen.errors as errors
import pandas as pd
from pathlib import Path
import pytest
import xarray as xr


def compare_tsp(tsp, **kwargs):
    d = dict(
        first_of_run=False,
        strict_first_of_year=False,
        strict_first_of_month=False,
        strict_first_of_day=False,
        first_of_year=False,
        first_of_month=False,
        first_of_day=False,
        last_of_run=False,
        strict_last_of_year=False,
        strict_last_of_month=False,
        strict_last_of_day=False,
        last_of_year=False,
        last_of_month=False,
        last_of_day=False,
    )
    d.update(kwargs)

    for key in d.keys():
        assert getattr(tsp, key) == d[key], f'{key} should be {d[key]}'


def test_timestep_properties():
    config = base_config()
    config.start_date = '2015-07-28 00:00'
    config.end_date = '2015-12-31'
    config.timestep = 'H'

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

    model.run_single()
    compare_tsp(
        model.timestep_props,
        first_of_run=True,
        first_of_year=True,
        first_of_month=True,
        first_of_day=True,
        strict_first_of_day=True,
    )

    model.run_single()
    compare_tsp(model.timestep_props)

    while model.date < pd.Timestamp('2015-07-29 00:00'):
        model.run_single()
    compare_tsp(
        model.timestep_props,
        first_of_day=True,
        strict_first_of_day=True,
    )

    while model.date < pd.Timestamp('2015-08-01 00:00'):
        model.run_single()
    compare_tsp(
        model.timestep_props,
        first_of_month=True,
        first_of_day=True,
        strict_first_of_month=True,
        strict_first_of_day=True,
    )

    config.start_date = '2015-07-28 23:00'
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run_single()
    compare_tsp(
        model.timestep_props,
        first_of_run=True,
        first_of_year=True,
        first_of_month=True,
        first_of_day=True,
        last_of_day=True,
        strict_last_of_day=True,
    )

    config.start_date = '2015-08-01 00:00'
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run_single()
    compare_tsp(
        model.timestep_props,
        first_of_run=True,
        first_of_year=True,
        first_of_month=True,
        first_of_day=True,
        strict_first_of_month=True,
        strict_first_of_day=True,
    )

    config.start_date = '2014-12-31 23:00'
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run_single()
    compare_tsp(
        model.timestep_props,
        first_of_run=True,
        first_of_year=True,
        first_of_month=True,
        first_of_day=True,
        last_of_year=True,
        last_of_month=True,
        last_of_day=True,
        strict_last_of_year=True,
        strict_last_of_month=True,
        strict_last_of_day=True,
    )

    model.run_single()
    compare_tsp(
        model.timestep_props,
        first_of_year=True,
        first_of_month=True,
        first_of_day=True,
        strict_first_of_year=True,
        strict_first_of_month=True,
        strict_first_of_day=True,
    )

    config.end_date = '2015-12-15 12:00'
    config.start_date = config.end_date
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run_single()
    compare_tsp(
        model.timestep_props,
        first_of_run=True,
        last_of_run=True,
        first_of_year=True,
        first_of_month=True,
        first_of_day=True,
        last_of_year=True,
        last_of_month=True,
        last_of_day=True,
    )


@pytest.mark.slow
def test_state_variable_reset():
    config = base_config()
    config.start_date = '2019-10-01'
    config.end_date = '2020-05-31'
    config.timestep = '3H'
    config.evapotranspiration.enabled = True
    config.canopy.enabled = True

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

    for category in model.state.categories:
        for var_name in model.state[category]._meta.keys():
            full_var_name = f'{category}.{var_name}'
            config.output_data.grids.variables.append({
                'var': full_var_name,
                'name': full_var_name,
            })

    config.reset_state_variables = False
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run()
    ds0 = model.gridded_output.data

    config.reset_state_variables = True
    model = oa.OpenAmundsen(config)
    model.initialize()
    model.run()
    ds1 = model.gridded_output.data

    compare_vars = [v['name'] for v in config.output_data.grids.variables]
    success = True

    for v in compare_vars:
        arr0 = ds0[v].values
        arr1 = ds1[v].values

        if not np.allclose(arr0, arr1, equal_nan=True):
            print(f'Mismatch in variable {v}\n'
                  f'Without reset:\n{arr0}\n'
                  f'With reset:\n{arr1}')
            success = False

    assert success
