import os

import warnings

from collections.abc import Iterable

from pathlib import Path

import dask.array as da

import numpy as np

import pytest

from anndata import AnnData

from scipy import sparse

from sklearn.exceptions import ConvergenceWarning

from ehrapy.anndata.anndata_ext import _are_ndarrays_equal, _is_val_missing, _to_dense_matrix

from ehrapy.preprocessing._imputation import (

    _warn_imputation_threshold,

    explicit_impute,

    knn_impute,

    mice_forest_impute,

    miss_forest_impute,

    simple_impute,

)

from tests.conftest import ARRAY_TYPES, TEST_DATA_PATH

CURRENT_DIR = Path(__file__).parent

_TEST_PATH = f"{TEST_DATA_PATH}/imputation"

def _base_check_imputation(

    adata_before_imputation: AnnData,

    adata_after_imputation: AnnData,

    before_imputation_layer: str | None = None,

    after_imputation_layer: str | None = None,

    imputed_var_names: Iterable[str] | None = None,

):

    """Provides a base check for all imputations:

    - Imputation doesn't leave any NaN behind

    - Imputation doesn't modify anything in non-imputated columns (if the imputation on a subset was requested)

    - Imputation doesn't modify any data that wasn't NaN

    Args:

        adata_before_imputation: AnnData before imputation

        adata_after_imputation: AnnData after imputation

        before_imputation_layer: Layer to consider in the original ``AnnData``, ``X`` if not specified

        after_imputation_layer: Layer to consider in the imputated ``AnnData``, ``X`` if not specified

        imputed_var_names: Names of the features that were imputated, will consider all of them if not specified

    Raises:

        AssertionError: If any of the checks fail.

    """

    # Convert dask arrays to numpy arrays

    if isinstance(adata_before_imputation.X, da.Array):

        adata_before_imputation.X = adata_before_imputation.X.compute()

    if isinstance(adata_after_imputation.X, da.Array):

        adata_after_imputation.X = adata_after_imputation.X.compute()

    layer_before = _to_dense_matrix(adata_before_imputation, before_imputation_layer)

    layer_after = _to_dense_matrix(adata_after_imputation, after_imputation_layer)

    if layer_before.shape != layer_after.shape:

        raise AssertionError("The shapes of the two layers do not match")

    var_indices = (

        np.arange(layer_before.shape[1])

        if imputed_var_names is None

        else [

            adata_before_imputation.var_names.get_loc(var_name)

            for var_name in imputed_var_names

            if var_name in imputed_var_names

        ]

    )

    before_nan_mask = _is_val_missing(layer_before)

    imputed_mask = np.zeros(layer_before.shape[1], dtype=bool)

    imputed_mask[var_indices] = True

    # Ensure no NaN remains in the imputed columns of layer_after

    if np.any(before_nan_mask[:, imputed_mask] & _is_val_missing(layer_after[:, imputed_mask])):

        raise AssertionError("NaN found in imputed columns of layer_after.")

    # Ensure unchanged values outside imputed columns

    unchanged_mask = ~imputed_mask

    if not _are_ndarrays_equal(layer_before[:, unchanged_mask], layer_after[:, unchanged_mask]):

        raise AssertionError("Values outside imputed columns were modified.")

    # Ensure imputation does not alter non-NaN values in the imputed columns

    imputed_non_nan_mask = (~before_nan_mask) & imputed_mask

    if not _are_ndarrays_equal(layer_before[imputed_non_nan_mask], layer_after[imputed_non_nan_mask]):

        raise AssertionError("Non-NaN values in imputed columns were modified.")

    # If reaching here: all checks passed

    return

def test_base_check_imputation_incompatible_shapes(impute_num_adata):

    adata_imputed = knn_impute(impute_num_adata, copy=True)

    with pytest.raises(AssertionError):

        _base_check_imputation(impute_num_adata, adata_imputed[1:, :])

    with pytest.raises(AssertionError):

        _base_check_imputation(impute_num_adata, adata_imputed[:, 1:])

def test_base_check_imputation_nan_detected_after_complete_imputation(impute_num_adata):

    adata_imputed = knn_impute(impute_num_adata, copy=True)

    adata_imputed.X[0, 2] = np.nan

    with pytest.raises(AssertionError):

        _base_check_imputation(impute_num_adata, adata_imputed)

def test_base_check_imputation_nan_detected_after_partial_imputation(impute_num_adata):

    var_names = ("col2", "col3")

    adata_imputed = knn_impute(impute_num_adata, var_names=var_names, copy=True)

    adata_imputed.X[0, 2] = np.nan

    with pytest.raises(AssertionError):

        _base_check_imputation(impute_num_adata, adata_imputed, imputed_var_names=var_names)

def test_base_check_imputation_nan_ignored_if_not_in_imputed_column(impute_num_adata):

    var_names = ("col2", "col3")

    adata_imputed = knn_impute(impute_num_adata, var_names=var_names, copy=True)

    # col1 has a NaN at row 2, should get ignored

    _base_check_imputation(impute_num_adata, adata_imputed, imputed_var_names=var_names)

def test_base_check_imputation_change_detected_in_non_imputed_column(impute_num_adata):

    var_names = ("col2", "col3")

    adata_imputed = knn_impute(impute_num_adata, var_names=var_names, copy=True)

    # col1 has a NaN at row 2, let's simulate it has been imputed by mistake

    adata_imputed.X[2, 0] = 42.0

    with pytest.raises(AssertionError):

        _base_check_imputation(impute_num_adata, adata_imputed, imputed_var_names=var_names)

def test_base_check_imputation_change_detected_in_imputed_column(impute_num_adata):

    adata_imputed = knn_impute(impute_num_adata, copy=True)

    # col3 didn't have a NaN at row 1, let's simulate it has been modified by mistake

    adata_imputed.X[1, 2] = 42.0

    with pytest.raises(AssertionError):

        _base_check_imputation(impute_num_adata, adata_imputed)

def test_mean_impute_no_copy(impute_num_adata):

    adata_not_imputed = impute_num_adata.copy()

    simple_impute(impute_num_adata)

    _base_check_imputation(adata_not_imputed, impute_num_adata)

def test_mean_impute_copy(impute_num_adata):

    adata_imputed = simple_impute(impute_num_adata, copy=True)

    assert id(impute_num_adata) != id(adata_imputed)

    _base_check_imputation(impute_num_adata, adata_imputed)

def test_mean_impute_throws_error_non_numerical(impute_adata):

    with pytest.raises(ValueError):

        simple_impute(impute_adata)

def test_mean_impute_subset(impute_adata):

    var_names = ("intcol", "indexcol")

    adata_imputed = simple_impute(impute_adata, var_names=var_names, copy=True)

    _base_check_imputation(impute_adata, adata_imputed, imputed_var_names=var_names)

    assert np.any([item != item for item in adata_imputed.X[::, 3:4]])

def test_median_impute_no_copy(impute_num_adata):

    adata_not_imputed = impute_num_adata.copy()

    simple_impute(impute_num_adata, strategy="median")

    _base_check_imputation(adata_not_imputed, impute_num_adata)

def test_median_impute_copy(impute_num_adata):

    adata_imputed = simple_impute(impute_num_adata, strategy="median", copy=True)

    _base_check_imputation(impute_num_adata, adata_imputed)

    assert id(impute_num_adata) != id(adata_imputed)

def test_median_impute_throws_error_non_numerical(impute_adata):

    with pytest.raises(ValueError):

        simple_impute(impute_adata, strategy="median")

def test_median_impute_subset(impute_adata):

    var_names = ("intcol", "indexcol")

    adata_imputed = simple_impute(impute_adata, var_names=var_names, strategy="median", copy=True)

    _base_check_imputation(impute_adata, adata_imputed, imputed_var_names=var_names)

def test_most_frequent_impute_no_copy(impute_adata):

    adata_not_imputed = impute_adata.copy()

    simple_impute(impute_adata, strategy="most_frequent")

    _base_check_imputation(adata_not_imputed, impute_adata)

def test_most_frequent_impute_copy(impute_adata):

    adata_imputed = simple_impute(impute_adata, strategy="most_frequent", copy=True)

    _base_check_imputation(impute_adata, adata_imputed)

    assert id(impute_adata) != id(adata_imputed)

def test_unknown_simple_imputation_strategy(impute_adata):

    with pytest.raises(ValueError):

        simple_impute(impute_adata, strategy="invalid_strategy", copy=True)  # type: ignore

def test_most_frequent_impute_subset(impute_adata):

    var_names = ("intcol", "strcol")

    adata_imputed = simple_impute(impute_adata, var_names=var_names, strategy="most_frequent", copy=True)

    _base_check_imputation(impute_adata, adata_imputed, imputed_var_names=var_names)

def test_knn_impute_check_backend(impute_num_adata):

    knn_impute(impute_num_adata, backend="faiss", copy=True)

    knn_impute(impute_num_adata, backend="scikit-learn", copy=True)

    with pytest.raises(

        ValueError,

        match="Unknown backend 'invalid_backend' for KNN imputation. Choose between 'scikit-learn' and 'faiss'.",

    ):

        knn_impute(impute_num_adata, backend="invalid_backend")  # type: ignore

def test_knn_impute_no_copy(impute_num_adata):

    adata_not_imputed = impute_num_adata.copy()

    knn_impute(impute_num_adata)

    _base_check_imputation(adata_not_imputed, impute_num_adata)

def test_knn_impute_copy(impute_num_adata):

    adata_imputed = knn_impute(impute_num_adata, n_neighbors=3, copy=True)

    _base_check_imputation(impute_num_adata, adata_imputed)

    assert id(impute_num_adata) != id(adata_imputed)

def test_knn_impute_non_numerical_data(impute_adata):

    with pytest.raises(ValueError):

        knn_impute(impute_adata, n_neighbors=3, copy=True)

def test_knn_impute_numerical_data(impute_num_adata):

    adata_imputed = knn_impute(impute_num_adata, copy=True)

    _base_check_imputation(impute_num_adata, adata_imputed)

def test_missforest_impute_non_numerical_data(impute_adata):

    with pytest.raises(ValueError):

        miss_forest_impute(impute_adata, copy=True)

def test_missforest_impute_numerical_data(impute_num_adata):

    warnings.filterwarnings("ignore", category=ConvergenceWarning)

    adata_imputed = miss_forest_impute(impute_num_adata, copy=True)

    _base_check_imputation(impute_num_adata, adata_imputed)

def test_missforest_impute_subset(impute_num_adata):

    warnings.filterwarnings("ignore", category=ConvergenceWarning)

    var_names = ("col2", "col3")

    adata_imputed = miss_forest_impute(impute_num_adata, var_names=var_names, copy=True)

    _base_check_imputation(impute_num_adata, adata_imputed, imputed_var_names=var_names)

@pytest.mark.parametrize(

    "array_type,expected_error",

    [

        (np.array, None),

        (da.from_array, NotImplementedError),

        (sparse.csr_matrix, NotImplementedError),

    ],

)

def test_miceforest_array_types(impute_num_adata, array_type, expected_error):

    impute_num_adata.X = array_type(impute_num_adata.X)

    if expected_error:

        with pytest.raises(expected_error):

            mice_forest_impute(impute_num_adata, copy=True)

@pytest.mark.skipif(os.name == "Darwin", reason="miceforest Imputation not supported by MacOS.")

def test_miceforest_impute_no_copy(impute_iris_adata):

    adata_not_imputed = impute_iris_adata.copy()

    mice_forest_impute(impute_iris_adata)

    _base_check_imputation(adata_not_imputed, impute_iris_adata)

@pytest.mark.skipif(os.name == "Darwin", reason="miceforest Imputation not supported by MacOS.")

def test_miceforest_impute_copy(impute_iris_adata):

    adata_imputed = mice_forest_impute(impute_iris_adata, copy=True)

    _base_check_imputation(impute_iris_adata, adata_imputed)

    assert id(impute_iris_adata) != id(adata_imputed)

@pytest.mark.skipif(os.name == "Darwin", reason="miceforest Imputation not supported by MacOS.")

def test_miceforest_impute_non_numerical_data(impute_titanic_adata):

    with pytest.raises(ValueError):

        mice_forest_impute(impute_titanic_adata)

@pytest.mark.skipif(os.name == "Darwin", reason="miceforest Imputation not supported by MacOS.")

def test_miceforest_impute_numerical_data(impute_iris_adata):

    adata_not_imputed = impute_iris_adata.copy()

    mice_forest_impute(impute_iris_adata)

    _base_check_imputation(adata_not_imputed, impute_iris_adata)

@pytest.mark.parametrize(

    "array_type,expected_error",

    [

        (np.array, None),

        (da.from_array, None),

        (sparse.csr_matrix, NotImplementedError),

    ],

)

def test_explicit_impute_array_types(impute_num_adata, array_type, expected_error):

    impute_num_adata.X = array_type(impute_num_adata.X)

    if expected_error:

        with pytest.raises(expected_error):

            explicit_impute(impute_num_adata, replacement=1011, copy=True)

@pytest.mark.parametrize("array_type", ARRAY_TYPES)

def test_explicit_impute_all(array_type, impute_num_adata):

    impute_num_adata.X = array_type(impute_num_adata.X)

    warnings.filterwarnings("ignore", category=FutureWarning)

    adata_imputed = explicit_impute(impute_num_adata, replacement=1011, copy=True)

    _base_check_imputation(impute_num_adata, adata_imputed)

    assert np.sum([adata_imputed.X == 1011]) == 3

@pytest.mark.parametrize("array_type", ARRAY_TYPES)

def test_explicit_impute_subset(impute_adata, array_type):

    impute_adata.X = array_type(impute_adata.X)

    adata_imputed = explicit_impute(impute_adata, replacement={"strcol": "REPLACED", "intcol": 1011}, copy=True)

    _base_check_imputation(impute_adata, adata_imputed, imputed_var_names=("strcol", "intcol"))

    assert np.sum([adata_imputed.X == 1011]) == 1

    assert np.sum([adata_imputed.X == "REPLACED"]) == 1

def test_warning(impute_num_adata):

    warning_results = _warn_imputation_threshold(impute_num_adata, threshold=20, var_names=None)

    assert warning_results == {"col1": 25, "col3": 50}