from __future__ import annotations

from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd
import pytest
from scipy import sparse

from anndata import AnnData, concat
from anndata._core.merge import _resolve_axis
from anndata.experimental.merge import as_group, concat_on_disk
from anndata.io import read_elem, write_elem
from anndata.tests.helpers import assert_equal, gen_adata
from anndata.utils import asarray

if TYPE_CHECKING:
    from pathlib import Path
    from typing import Literal


GEN_ADATA_OOC_CONCAT_ARGS = dict(
    obsm_types=(
        sparse.csr_matrix,
        np.ndarray,
        pd.DataFrame,
    ),
    varm_types=(sparse.csr_matrix, np.ndarray, pd.DataFrame),
    layers_types=(sparse.csr_matrix, np.ndarray, pd.DataFrame),
)


@pytest.fixture(params=[0, 1])
def axis(request) -> Literal[0, 1]:
    return request.param


@pytest.fixture(params=["array", "sparse", "sparse_array"])
def array_type(request) -> Literal["array", "sparse", "sparse_array"]:
    return request.param


@pytest.fixture(params=["inner", "outer"])
def join_type(request) -> Literal["inner", "outer"]:
    return request.param


@pytest.fixture(params=["zarr", "h5ad"])
def file_format(request) -> Literal["zarr", "h5ad"]:
    return request.param


# 1000 is enough to guarantee that the feature is being used
@pytest.fixture(params=[1_000, 100_000_000])
def max_loaded_elems(request) -> int:
    return request.param


def _adatas_to_paths(adatas, tmp_path, file_format):
    """
    Gets list of adatas, writes them and returns their paths as zarr
    """
    paths = None

    if isinstance(adatas, Mapping):
        paths = {}
        for k, v in adatas.items():
            p = tmp_path / (f"{k}." + file_format)
            write_elem(as_group(p, mode="a"), "", v)
            paths[k] = p
    else:
        paths = []
        for i, a in enumerate(adatas):
            p = tmp_path / (f"{i}." + file_format)
            write_elem(as_group(p, mode="a"), "", a)
            paths += [p]
    return paths


def assert_eq_concat_on_disk(
    adatas,
    tmp_path: Path,
    file_format: Literal["zarr", "h5ad"],
    max_loaded_elems: int | None = None,
    *args,
    **kwargs,
):
    # create one from the concat function
    res1 = concat(adatas, *args, **kwargs)
    # create one from the on disk concat function
    paths = _adatas_to_paths(adatas, tmp_path, file_format)
    out_name = tmp_path / f"out.{file_format}"
    if max_loaded_elems is not None:
        kwargs["max_loaded_elems"] = max_loaded_elems
    concat_on_disk(paths, out_name, *args, **kwargs)
    res2 = read_elem(as_group(out_name, mode="r"))
    assert_equal(res1, res2, exact=False)


def get_array_type(array_type, axis):
    if array_type == "sparse":
        return sparse.csr_matrix if axis == 0 else sparse.csc_matrix
    if array_type == "sparse_array":
        return sparse.csr_array if axis == 0 else sparse.csc_array
    if array_type == "array":
        return asarray
    msg = f"array_type {array_type} not implemented"
    raise NotImplementedError(msg)


@pytest.mark.parametrize("reindex", [True, False], ids=["reindex", "no_reindex"])
def test_anndatas(
    *,
    axis: Literal[0, 1],
    array_type: Literal["array", "sparse", "sparse_array"],
    join_type: Literal["inner", "outer"],
    tmp_path: Path,
    max_loaded_elems: int,
    file_format: Literal["zarr", "h5ad"],
    reindex: bool,
):
    _, off_axis_name = _resolve_axis(1 - axis)
    random_axes = {0, 1} if reindex else {axis}
    sparse_fmt = "csr" if axis == 0 else "csc"
    kw = (
        GEN_ADATA_OOC_CONCAT_ARGS
        if not reindex
        else dict(
            obsm_types=(get_array_type("sparse", 1 - axis), np.ndarray, pd.DataFrame),
            varm_types=(get_array_type("sparse", 1 - axis), np.ndarray, pd.DataFrame),
            layers_types=(get_array_type("sparse", axis), np.ndarray, pd.DataFrame),
        )
    )

    adatas = []
    for i in range(5):
        M, N = (np.random.randint(1, 100) if a in random_axes else 50 for a in (0, 1))
        a = gen_adata(
            (M, N), X_type=get_array_type(array_type, axis), sparse_fmt=sparse_fmt, **kw
        )
        # ensure some names overlap, others do not, for the off-axis so that inner/outer is properly tested
        off_names = getattr(a, f"{off_axis_name}_names").array
        off_names[1::2] = f"{i}-" + off_names[1::2]
        setattr(a, f"{off_axis_name}_names", off_names)
        adatas.append(a)

    assert_eq_concat_on_disk(
        adatas,
        tmp_path,
        file_format,
        max_loaded_elems,
        axis=axis,
        join=join_type,
    )


def test_concat_ordered_categoricals_retained(tmp_path, file_format):
    a = AnnData(
        X=np.ones((5, 1)),
        obs=pd.DataFrame(
            {
                "cat_ordered": pd.Categorical(list("aabcd"), ordered=True),
            },
            index=[f"cell{i:02}" for i in range(5)],
        ),
    )
    b = AnnData(
        X=np.ones((5, 1)),
        obs=pd.DataFrame(
            {
                "cat_ordered": pd.Categorical(list("abcdd"), ordered=True),
            },
            index=[f"cell{i:02}" for i in range(5, 10)],
        ),
    )

    adatas = [a, b]
    assert_eq_concat_on_disk(adatas, tmp_path, file_format)


@pytest.fixture
def xxxm_adatas():
    def gen_index(n):
        return [f"cell{i}" for i in range(n)]

    return [
        AnnData(
            X=sparse.csr_matrix((3, 5)),
            obs=pd.DataFrame(index=gen_index(3)),
            obsm={
                "dense": np.arange(6).reshape(3, 2),
                "sparse": sparse.csr_matrix(np.arange(6).reshape(3, 2)),
                "df": pd.DataFrame(
                    {
                        "a": np.arange(3),
                        "b": list("abc"),
                        "c": pd.Categorical(list("aab")),
                    },
                    index=gen_index(3),
                ),
            },
        ),
        AnnData(
            X=sparse.csr_matrix((4, 10)),
            obs=pd.DataFrame(index=gen_index(4)),
            obsm=dict(
                dense=np.arange(12).reshape(4, 3),
                df=pd.DataFrame(dict(a=np.arange(3, 7)), index=gen_index(4)),
            ),
        ),
        AnnData(
            X=sparse.csr_matrix((2, 100)),
            obs=pd.DataFrame(index=gen_index(2)),
            obsm={
                "sparse": np.arange(8).reshape(2, 4),
                "dense": np.arange(4, 8).reshape(2, 2),
                "df": pd.DataFrame(
                    {
                        "a": np.arange(7, 9),
                        "b": list("cd"),
                        "c": pd.Categorical(list("ab")),
                    },
                    index=gen_index(2),
                ),
            },
        ),
    ]


def test_concatenate_xxxm(xxxm_adatas, tmp_path, file_format, join_type):
    if join_type == "outer":
        for i in range(len(xxxm_adatas)):
            xxxm_adatas[i] = xxxm_adatas[i].T
            xxxm_adatas[i].X = sparse.csr_matrix(xxxm_adatas[i].X)
    assert_eq_concat_on_disk(xxxm_adatas, tmp_path, file_format, join=join_type)


def test_output_dir_exists(tmp_path):
    in_pth = tmp_path / "in.h5ad"
    out_pth = tmp_path / "does_not_exist" / "out.h5ad"

    AnnData(X=np.ones((5, 1))).write_h5ad(in_pth)

    with pytest.raises(FileNotFoundError, match=f"{out_pth}"):
        concat_on_disk([in_pth], out_pth)


def test_failure_w_no_args(tmp_path):
    with pytest.raises(ValueError, match=r"No objects to concatenate"):
        concat_on_disk([], tmp_path / "out.h5ad")