import numpy as np
from pytest_cases import parametrize_with_cases
import xarray as xr
from pytest import raises
from pandas import DataFrame

from gval.utils.loading_datasets import (
    adjust_memory_strategy,
    get_current_memory_strategy,
    _handle_xarray_memory,
)


@parametrize_with_cases(
    "candidate_map, benchmark_map, positive_categories, negative_categories, rasterize_attributes, memory_strategies",
    glob="data_array_accessor_success",
)
def test_data_array_accessor_success(
    candidate_map,
    benchmark_map,
    positive_categories,
    negative_categories,
    rasterize_attributes,
    memory_strategies,
):
    adjust_memory_strategy(memory_strategies)

    assert get_current_memory_strategy() == memory_strategies

    candidate_map = _handle_xarray_memory(candidate_map)

    data = candidate_map.gval.categorical_compare(
        benchmark_map=benchmark_map,
        positive_categories=positive_categories,
        negative_categories=negative_categories,
        rasterize_attributes=rasterize_attributes,
    )

    if rasterize_attributes is None:
        assert isinstance(data[0], xr.DataArray)
    else:
        assert isinstance(data[0], DataFrame)
    assert isinstance(data[1], DataFrame)
    assert isinstance(data[2], DataFrame)


@parametrize_with_cases(
    "candidate_map, benchmark_map, positive_categories, negative_categories, memory_strategies, exception",
    glob="data_array_accessor_fail",
)
def test_data_array_accessor_fail(
    candidate_map,
    benchmark_map,
    positive_categories,
    negative_categories,
    memory_strategies,
    exception,
):
    with raises(exception):
        adjust_memory_strategy(memory_strategies)

        if exception == OSError:
            candidate_map.encoding["source"] = "arb"
            _handle_xarray_memory(candidate_map)

        _ = candidate_map.gval.categorical_compare(
            benchmark_map=benchmark_map,
            positive_categories=positive_categories,
            negative_categories=negative_categories,
        )


@parametrize_with_cases(
    "candidate_map, benchmark_map, rasterize_attributes",
    glob="data_array_accessor_homogenize",
)
def test_data_array_accessor_homogenize(
    candidate_map, benchmark_map, rasterize_attributes
):
    data = candidate_map.gval.homogenize(
        benchmark_map=benchmark_map, rasterize_attributes=rasterize_attributes
    )

    assert isinstance(data[0], xr.DataArray)
    assert isinstance(data[1], xr.DataArray)


@parametrize_with_cases(
    "candidate_map, benchmark_map, vectorized",
    glob="data_array_accessor_compute_agreement",
)
def test_data_array_accessor_compute_agreement_success(
    candidate_map, benchmark_map, vectorized
):
    aligned_cand, aligned_bench = candidate_map.gval.homogenize(
        benchmark_map=benchmark_map
    )

    if vectorized:
        aligned_cand.gval.agreement_map_format = "vector"

    data = aligned_cand.gval.compute_agreement_map(benchmark_map=aligned_bench)

    if vectorized:
        assert isinstance(data, DataFrame)
    else:
        assert isinstance(data, xr.DataArray)


@parametrize_with_cases(
    "candidate_map, benchmark_map", glob="data_array_accessor_compute_agreement"
)
def test_data_array_accessor_compute_agreement_fail(candidate_map, benchmark_map):
    with raises(ValueError):
        _ = candidate_map.gval.compute_agreement_map(benchmark_map=benchmark_map)


@parametrize_with_cases(
    "candidate_map, benchmark_map",
    glob="data_array_accessor_crosstab_table_success",
)
def test_data_array_accessor_crosstab_table_success(candidate_map, benchmark_map):
    aligned_cand, aligned_bench = candidate_map.gval.homogenize(
        benchmark_map=benchmark_map
    )
    data = aligned_cand.gval.compute_crosstab(benchmark_map=aligned_bench)

    assert isinstance(data, DataFrame)


@parametrize_with_cases(
    "candidate_map, benchmark_map, exception",
    glob="data_array_accessor_crosstab_table_fail",
)
def test_data_array_accessor_crosstab_table_fail(
    candidate_map, benchmark_map, exception
):
    with raises(exception):
        _ = candidate_map.gval.compute_crosstab(benchmark_map=benchmark_map)


@parametrize_with_cases(
    "candidate_map, benchmark_map, positive_categories, negative_categories, rasterize_attributes",
    glob="data_set_accessor_success",
)
def test_data_set_accessor_success(
    candidate_map,
    benchmark_map,
    positive_categories,
    negative_categories,
    rasterize_attributes,
):
    data = candidate_map.gval.categorical_compare(
        benchmark_map=benchmark_map,
        positive_categories=positive_categories,
        negative_categories=negative_categories,
        rasterize_attributes=rasterize_attributes,
    )

    if rasterize_attributes is None:
        assert isinstance(data[0], xr.Dataset)
    else:
        assert isinstance(data[0], DataFrame)
    assert isinstance(data[1], DataFrame)
    assert isinstance(data[2], DataFrame)


@parametrize_with_cases(
    "candidate_map, benchmark_map, rasterize_attributes",
    glob="data_set_accessor_homogenize",
)
def test_data_set_accessor_homogenize(
    candidate_map, benchmark_map, rasterize_attributes
):
    data = candidate_map.gval.homogenize(
        benchmark_map=benchmark_map, rasterize_attributes=rasterize_attributes
    )

    assert isinstance(data[0], xr.Dataset)
    assert isinstance(data[1], xr.Dataset)


@parametrize_with_cases(
    "candidate_map, benchmark_map", glob="data_set_accessor_compute_agreement"
)
def test_data_set_accessor_compute_agreement_success(candidate_map, benchmark_map):
    aligned_cand, aligned_bench = candidate_map.gval.homogenize(
        benchmark_map=benchmark_map
    )
    data = aligned_cand.gval.compute_agreement_map(benchmark_map=aligned_bench)

    assert isinstance(data, xr.Dataset)


@parametrize_with_cases(
    "candidate_map, benchmark_map", glob="data_set_accessor_compute_agreement"
)
def test_data_set_accessor_compute_agreement_fail(candidate_map, benchmark_map):
    with raises(ValueError):
        _ = candidate_map.gval.compute_agreement_map(benchmark_map=benchmark_map)


@parametrize_with_cases(
    "candidate_map, benchmark_map", glob="data_set_accessor_crosstab_table_success"
)
def test_data_set_accessor_crosstab_table_success(candidate_map, benchmark_map):
    aligned_cand, aligned_bench = candidate_map.gval.homogenize(
        benchmark_map=benchmark_map
    )
    data = aligned_cand.gval.compute_crosstab(benchmark_map=aligned_bench)

    assert isinstance(data, DataFrame)


@parametrize_with_cases(
    "candidate_map, benchmark_map",
    glob="data_set_accessor_crosstab_table_fail",
)
def test_data_set_accessor_crosstab_table_fail(candidate_map, benchmark_map):
    with raises(IndexError):
        _ = candidate_map.gval.compute_crosstab(benchmark_map=benchmark_map)


@parametrize_with_cases(
    "crosstab_df, positive_categories, negative_categories",
    glob="data_frame_accessor_compute_metrics",
)
def test_data_frame_accessor_compute_metrics(
    crosstab_df, positive_categories, negative_categories
):
    data = crosstab_df.gval.compute_categorical_metrics(
        positive_categories=positive_categories, negative_categories=negative_categories
    )

    assert isinstance(data, DataFrame)


@parametrize_with_cases(
    "candidate_map, crs, entries",
    glob="categorical_plot_success",
)
def test_categorical_plot_success(candidate_map, crs, entries):
    candidate_map.rio.set_crs(crs)
    viz_object = candidate_map.gval.cat_plot(basemap=None)
    assert len(viz_object.axes.get_legend().texts) == entries


@parametrize_with_cases(
    "candidate_map, legend_labels, num_classes",
    glob="categorical_plot_fail",
)
def test_categorical_plot_fail(candidate_map, legend_labels, num_classes):
    candidate_map.data = np.random.choice(np.arange(num_classes), candidate_map.shape)
    with raises(ValueError):
        _ = candidate_map.gval.cat_plot(legend_labels=legend_labels, basemap=None)


@parametrize_with_cases(
    "candidate_map, axes",
    glob="continuous_plot_success",
)
def test_continuous_plot_success(candidate_map, axes):
    viz_object = candidate_map.gval.cont_plot(basemap=None)
    assert len(viz_object.figure.axes) == axes


@parametrize_with_cases(
    "candidate_map",
    glob="continuous_plot_fail",
)
def test_continuous_plot_fail(candidate_map):
    with raises(ValueError):
        _ = candidate_map.gval.cont_plot()


@parametrize_with_cases(
    "candidate_map, benchmark_map",
    glob="data_array_accessor_continuous",
)
def test_data_array_accessor_continuous(candidate_map, benchmark_map):
    agreement_map, metrics_df = candidate_map.gval.continuous_compare(
        benchmark_map=benchmark_map, metrics="all"
    )

    assert isinstance(agreement_map, xr.DataArray)
    assert isinstance(metrics_df, DataFrame)


@parametrize_with_cases(
    "candidate_map, benchmark_map",
    glob="data_set_accessor_continuous",
)
def test_data_set_accessor_continuous(
    candidate_map,
    benchmark_map,
):
    agreement_map, metrics_df = candidate_map.gval.continuous_compare(
        benchmark_map=benchmark_map, metrics="all"
    )

    assert isinstance(agreement_map, xr.Dataset)
    assert isinstance(metrics_df, DataFrame)
