"""Tests for the attribute .X"""

from __future__ import annotations

from contextlib import nullcontext

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

import anndata as ad
from anndata import AnnData
from anndata._warnings import ImplicitModificationWarning
from anndata.tests.helpers import GEN_ADATA_NO_XARRAY_ARGS, assert_equal, gen_adata
from anndata.utils import asarray

UNLABELLED_ARRAY_TYPES = [
    pytest.param(sparse.csr_matrix, id="csr"),
    pytest.param(sparse.csc_matrix, id="csc"),
    pytest.param(sparse.csr_array, id="csr_array"),
    pytest.param(sparse.csc_array, id="csc_array"),
    pytest.param(asarray, id="ndarray"),
]
SINGULAR_SHAPES = [
    pytest.param(shape, id=str(shape)) for shape in [(1, 10), (10, 1), (1, 1)]
]


@pytest.mark.parametrize("shape", SINGULAR_SHAPES)
@pytest.mark.parametrize("orig_array_type", UNLABELLED_ARRAY_TYPES)
@pytest.mark.parametrize("new_array_type", UNLABELLED_ARRAY_TYPES)
def test_setter_singular_dim(shape, orig_array_type, new_array_type):
    # https://github.com/scverse/anndata/issues/500
    adata = gen_adata(shape, X_type=orig_array_type)
    to_assign = new_array_type(np.ones(shape))
    adata.X = to_assign
    np.testing.assert_equal(asarray(adata.X), 1)
    assert isinstance(adata.X, type(to_assign))


def test_repeat_indices_view():
    adata = gen_adata((10, 10), X_type=np.asarray)
    subset = adata[[0, 0, 1, 1], :]
    mat = np.array([np.ones(adata.shape[1]) * i for i in range(4)])
    with pytest.warns(
        FutureWarning,
        match=r"You are attempting to set `X` to a matrix on a view which has non-unique indices",
    ):
        subset.X = mat


@pytest.mark.parametrize("orig_array_type", UNLABELLED_ARRAY_TYPES)
@pytest.mark.parametrize("new_array_type", UNLABELLED_ARRAY_TYPES)
@pytest.mark.parametrize("copy_on_write_X", [True, False], ids=["CoW", "update"])
def test_setter_view(orig_array_type, new_array_type, *, copy_on_write_X: bool):
    ad.settings.copy_on_write_X = copy_on_write_X
    adata = gen_adata((10, 10), X_type=orig_array_type)
    orig_X = adata.X
    expected_X = asarray(orig_X.copy())
    to_assign = new_array_type(np.ones((9, 9)))
    if not copy_on_write_X:
        expected_X[:9, :9] = asarray(to_assign)
    view = adata[:9, :9]
    with (
        pytest.warns(
            ImplicitModificationWarning if copy_on_write_X else FutureWarning,
            match=r"initializing view as actual"
            if copy_on_write_X
            else r"will obey copy-on-write semantics",
        ),
        pytest.warns(UserWarning, match=r"Trying to set a dense array")
        if sparse.issparse(to_assign)
        and isinstance(orig_X, np.ndarray)
        and not copy_on_write_X
        else nullcontext(),
    ):
        view.X = to_assign
    assert_equal(view.X, to_assign)
    assert isinstance(view.X, type(to_assign) if copy_on_write_X else type(orig_X))
    assert_equal(adata.X, expected_X)
    # If cow, then not a view and if not cow, it is a view
    assert view.is_view != copy_on_write_X


###############################
# Tests for `adata.X is None` #
###############################


def test_set_x_is_none():
    # test setter and getter
    adata = AnnData(np.array([[1, 2, 3], [4, 5, 6]]), dict(o1=[1, 2], o2=[3, 4]))
    adata.X = None
    assert adata.X is None


def test_del_set_equiv_X():
    """Tests that `del adata.X` is equivalent to `adata.X = None`"""
    # test setter and deleter
    orig = gen_adata((10, 10))
    copy = orig.copy()

    del orig.X
    copy.X = None

    assert orig.X is None
    assert_equal(orig, copy)

    # Check that deleting again is still fine
    del orig.X
    assert orig.X is None


@pytest.mark.parametrize(
    ("obs", "var", "shape_expected"),
    [
        pytest.param(dict(obs_names=["1", "2"]), None, (2, 0), id="obs"),
        pytest.param(None, dict(var_names=["a", "b"]), (0, 2), id="var"),
        pytest.param(
            dict(obs_names=["1", "2", "3"]),
            dict(var_names=["a", "b"]),
            (3, 2),
            id="both",
        ),
    ],
)
def test_init_x_as_none_shape_from_obs_var(obs, var, shape_expected):
    adata = AnnData(None, obs, var)
    assert adata.X is None
    assert adata.shape == shape_expected


def test_init_x_as_none_explicit_shape():
    shape = (3, 5)
    adata = AnnData(None, uns=dict(test=np.array((3, 3))), shape=shape)
    assert adata.X is None
    assert adata.shape == shape


@pytest.mark.parametrize("shape", [*SINGULAR_SHAPES, pytest.param((5, 3), id="(5, 3)")])
def test_transpose_with_X_as_none(shape):
    adata = gen_adata(shape, X_type=lambda x: None)
    adataT = adata.transpose()
    assert_equal(adataT.shape, shape[::-1])
    assert_equal(adataT.obsp.keys(), adata.varp.keys())
    assert_equal(adataT.T, adata)


def test_copy():
    adata = AnnData(
        None,
        obs=pd.DataFrame(index=[f"cell{i:03}" for i in range(100)]),
        var=pd.DataFrame(index=[f"gene{i:03}" for i in range(200)]),
    )
    assert_equal(adata.copy(), adata)


def test_copy_view():
    adata = AnnData(
        None,
        obs=pd.DataFrame(index=[f"cell{i:03}" for i in range(100)]),
        var=pd.DataFrame(index=[f"gene{i:03}" for i in range(200)]),
    )
    v = adata[::-2, ::-2]
    assert_equal(v.copy(), v)


############
# IO tests #
############


def test_io_missing_X(tmp_path, diskfmt):
    file_pth = tmp_path / f"x_none_adata.{diskfmt}"
    write = lambda obj, pth: getattr(obj, f"write_{diskfmt}")(pth)
    read = lambda pth: getattr(ad, f"read_{diskfmt}")(pth)

    adata = gen_adata((20, 30), **GEN_ADATA_NO_XARRAY_ARGS)
    del adata.X

    write(adata, file_pth)
    from_disk = read(file_pth)

    assert_equal(from_disk, adata)


def test_set_dense_x_view_from_sparse():
    x = np.zeros((100, 30))
    x1 = np.ones((100, 30))
    orig = ad.AnnData(x)
    view = orig[:30]
    with (
        pytest.warns(
            UserWarning,
            match=r"Trying to set a dense array with a sparse array on a view",
        ),
        pytest.warns(
            ImplicitModificationWarning, match=r"Modifying `X` on a view results"
        ),
    ):
        view.X = sparse.csr_matrix(x1[:30])
    assert_equal(view.X, x1[:30])
    assert_equal(orig.X[:30], x1[:30])  # change propagates through
    assert_equal(orig.X[30:], x[30:])  # change propagates through


def test_fail_on_non_csr_csc_matrix():
    X = sparse.eye(100, format="coo")
    with pytest.raises(
        ValueError,
        match=r"Only CSR and CSC.*",
    ):
        ad.AnnData(X=X)
