from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd
import pytest
from anndata import AnnData
from matplotlib.testing.compare import compare_images

import ehrapy as ep
from ehrapy.io import read_csv

if TYPE_CHECKING:
    import os

    from matplotlib.figure import Figure

TEST_DATA_PATH = Path(__file__).parent / "data"


@pytest.fixture
def root_dir():
    return Path(__file__).resolve().parent


@pytest.fixture
def rng():
    return np.random.default_rng(seed=42)


@pytest.fixture
def obs_data():
    return {
        "disease": ["cancer", "tumor"],
        "country": ["Germany", "switzerland"],
        "sex": ["male", "female"],
    }


@pytest.fixture
def var_data():
    return {
        "alive": ["yes", "no", "maybe"],
        "hospital": ["hospital 1", "hospital 2", "hospital 1"],
        "crazy": ["yes", "yes", "yes"],
    }


@pytest.fixture
def missing_values_adata(obs_data, var_data):
    return AnnData(
        X=np.array([[0.21, np.nan, 41.42], [np.nan, np.nan, 7.234]], dtype=np.float32),
        obs=pd.DataFrame(data=obs_data),
        var=pd.DataFrame(data=var_data, index=["Acetaminophen", "hospital", "crazy"]),
    )


@pytest.fixture
def lab_measurements_simple_adata(obs_data, var_data):
    X = np.array([[73, 0.02, 1.00], [148, 0.25, 3.55]], dtype=np.float32)
    return AnnData(
        X=X,
        obs=pd.DataFrame(data=obs_data),
        var=pd.DataFrame(data=var_data, index=["Acetaminophen", "Acetoacetic acid", "Beryllium, toxic"]),
    )


@pytest.fixture
def lab_measurements_layer_adata(obs_data, var_data):
    X = np.array([[73, 0.02, 1.00], [148, 0.25, 3.55]], dtype=np.float32)
    return AnnData(
        X=X,
        obs=pd.DataFrame(data=obs_data),
        var=pd.DataFrame(data=var_data, index=["Acetaminophen", "Acetoacetic acid", "Beryllium, toxic"]),
        layers={"layer_copy": X},
    )


@pytest.fixture
def mimic_2():
    adata = ep.dt.mimic_2()
    return adata


@pytest.fixture
def mimic_2_encoded():
    adata = ep.dt.mimic_2(encoded=True)
    return adata


@pytest.fixture
def mimic_2_10():
    mimic_2_10 = ep.dt.mimic_2()[:10]

    return mimic_2_10


@pytest.fixture
def mar_adata(rng) -> AnnData:
    """Generate MAR data using dependent columns."""
    data = rng.random((100, 10))
    # Assume missingness in the last column depends on the values of the first column
    missing_indicator = data[:, 0] < np.percentile(data[:, 0], 0.1 * 100)
    data[missing_indicator, -1] = np.nan  # Only last column has missing values dependent on the first column

    return AnnData(data)


@pytest.fixture
def mcar_adata(rng) -> AnnData:
    """Generate MCAR data by randomly sampling."""
    data = rng.random((100, 10))
    missing_indices = rng.choice(a=[False, True], size=data.shape, p=[1 - 0.1, 0.1])
    data[missing_indices] = np.nan

    return AnnData(data)


@pytest.fixture
def adata_mini():
    return read_csv(f"{TEST_DATA_PATH}/dataset1.csv", columns_obs_only=["glucose", "weight", "disease", "station"])


@pytest.fixture
def adata_move_obs_num() -> AnnData:
    return read_csv(TEST_DATA_PATH / "io/dataset_move_obs_num.csv")


@pytest.fixture
def adata_move_obs_mix() -> AnnData:
    return read_csv(TEST_DATA_PATH / "io/dataset_move_obs_mix.csv")


@pytest.fixture
def impute_num_adata() -> AnnData:
    adata = read_csv(dataset_path=f"{TEST_DATA_PATH}/imputation/test_impute_num.csv")
    return adata


@pytest.fixture
def impute_adata() -> AnnData:
    adata = read_csv(dataset_path=f"{TEST_DATA_PATH}/imputation/test_impute.csv")
    return adata


@pytest.fixture
def impute_iris_adata() -> AnnData:
    adata = read_csv(dataset_path=f"{TEST_DATA_PATH}/imputation/test_impute_iris.csv")
    return adata


@pytest.fixture
def impute_titanic_adata():
    adata = read_csv(dataset_path=f"{TEST_DATA_PATH}/imputation/test_impute_titanic.csv")
    return adata


@pytest.fixture
def encode_ds_1_adata() -> AnnData:
    adata = read_csv(dataset_path=f"{TEST_DATA_PATH}/encode/dataset1.csv")
    return adata


@pytest.fixture
def encode_ds_2_adata() -> AnnData:
    adata = read_csv(dataset_path=f"{TEST_DATA_PATH}/encode/dataset2.csv")
    return adata


# simplified from https://github.com/scverse/scanpy/blob/main/scanpy/tests/conftest.py
@pytest.fixture
def check_same_image(tmp_path):
    def check_same_image(
        fig: Figure,
        base_path: Path | os.PathLike,
        *,
        tol: float,
    ) -> None:
        expected = Path(base_path).parent / (Path(base_path).name + "_expected.png")
        if not Path(expected).is_file():
            raise OSError(f"No expected output found at {expected}.")
        actual = tmp_path / "actual.png"

        fig.savefig(actual, dpi=80)

        result = compare_images(expected, actual, tol=tol, in_decorator=True)

        if result is None:
            return None

        raise AssertionError(result)

    return check_same_image


def asarray(a):
    import numpy as np

    return np.asarray(a)


def as_dense_dask_array(a, chunk_size=1000):
    import dask.array as da

    return da.from_array(a, chunks=chunk_size)


ARRAY_TYPES = (asarray, as_dense_dask_array)