from __future__ import annotations

import anndata as ad
import cupy as cp
import numpy as np
import pandas as pd
import pooch
import pytest
import scanpy as sc
from scipy.stats import pearsonr

import rapids_singlecell as rsc
from rapids_singlecell.preprocessing._harmony._helper import (
    _choose_colsum_algo_benchmark,
    _choose_colsum_algo_heuristic,
    _colsum_heuristic,
    _scatter_add_cp,
)


def _get_measure(x, base, norm):
    assert norm in ["r", "L2"]

    if norm == "r":
        corr, _ = pearsonr(x, base)
        return corr
    else:
        return np.linalg.norm(x - base) / np.linalg.norm(base)


@pytest.fixture
def adata_reference():
    X_pca_file = pooch.retrieve(
        "https://github.com/slowkow/harmonypy/raw/refs/heads/master/data/pbmc_3500_pcs.tsv.gz",
        known_hash="md5:27e319b3ddcc0c00d98e70aa8e677b10",
    )
    X_pca = pd.read_csv(X_pca_file, delimiter="\t")
    X_pca_harmony_file = pooch.retrieve(
        "https://github.com/slowkow/harmonypy/raw/refs/heads/master/data/pbmc_3500_pcs_harmonized.tsv.gz",
        known_hash="md5:a7c4ce4b98c390997c66d63d48e09221",
    )
    X_pca_harmony = pd.read_csv(X_pca_harmony_file, delimiter="\t")
    meta_file = pooch.retrieve(
        "https://github.com/slowkow/harmonypy/raw/refs/heads/master/data/pbmc_3500_meta.tsv.gz",
        known_hash="md5:8c7ca20e926513da7cf0def1211baecb",
    )
    meta = pd.read_csv(meta_file, delimiter="\t")
    adata = ad.AnnData(
        X=None,
        obs=meta,
        obsm={"X_pca": X_pca.values, "harmony_org": X_pca_harmony.values},
    )
    return adata


@pytest.mark.parametrize("correction_method", ["fast", "original", "batched"])
@pytest.mark.parametrize("dtype", [cp.float32, cp.float64])
def test_harmony_integrate(correction_method, dtype):
    """
    Test that Harmony integrate works.

    This is a very simple test that just checks to see if the Harmony
    integrate wrapper successfully added a new field to ``adata.obsm``
    and makes sure it has the same dimensions as the original PCA table.
    """
    adata = sc.datasets.pbmc68k_reduced()
    rsc.pp.harmony_integrate(
        adata,
        "bulk_labels",
        correction_method=correction_method,
        dtype=dtype,
    )
    assert adata.obsm["X_pca_harmony"].shape == adata.obsm["X_pca"].shape


@pytest.mark.parametrize("dtype", [cp.float32, cp.float64])
def test_harmony_integrate_algos(dtype):
    """
    Test that Harmony integrate works.

    This is a very simple test that just checks to see if the Harmony
    integrate wrapper successfully added a new field to ``adata.obsm``
    and makes sure it has the same dimensions as the original PCA table.
    """
    adata = sc.datasets.pbmc68k_reduced()
    rsc.pp.harmony_integrate(
        adata, "bulk_labels", correction_method="fast", dtype=dtype
    )
    fast = adata.obsm["X_pca_harmony"].copy()
    rsc.pp.harmony_integrate(
        adata, "bulk_labels", correction_method="original", dtype=dtype
    )
    slow = adata.obsm["X_pca_harmony"].copy()
    assert _get_measure(fast, slow, "r").min() > 0.99
    assert _get_measure(fast, slow, "L2").max() < 0.1


@pytest.mark.parametrize("algo", ["columns", "atomics", "gemm"])
@pytest.mark.parametrize("dtype", [cp.float32, cp.float64, cp.int32])
def test_colsum_algo(algo, dtype):
    # Int32 testing for correctness of the algorithm
    if dtype == cp.int32:
        X = cp.random.randint(0, 10, size=(20, 10), dtype=dtype)
    else:
        X = cp.random.randn(20, 10, dtype=dtype)
    algo_func = _choose_colsum_algo_heuristic(X.shape[0], X.shape[1], algo)
    algo_out = algo_func(X)
    cupy_out = X.sum(axis=0)
    if dtype == cp.int32:
        cp.testing.assert_array_equal(algo_out, cupy_out)
    elif dtype == cp.float32:
        cp.testing.assert_allclose(algo_out, cupy_out, atol=1e-5)
    else:
        cp.testing.assert_allclose(algo_out, cupy_out)


@pytest.mark.parametrize("compute_capability", ["100", "80"])
def test_choose_colsum_algo(compute_capability):
    # Test that the choose_colsum_algo function returns the correct algorithm
    # for the given shape of the matrix
    for rows in np.arange(1000, 300000, 1000):
        for columns in np.arange(10, 5000, 50):
            algo = _colsum_heuristic(rows, columns, compute_capability)
            assert algo in ["columns", "atomics", "gemm"]


@pytest.mark.parametrize("dtype", [cp.float32, cp.float64])
def test_benchmark_colsum_algorithms(dtype):
    # Test that the benchmark_colsum_algorithms function returns the correct algorithm
    # for the given shape of the matrix
    test_shape = (1000, 100)
    algo_func = _choose_colsum_algo_benchmark(test_shape[0], test_shape[1], dtype)
    assert callable(algo_func)


@pytest.mark.parametrize("dtype", [cp.float32, cp.float64])
@pytest.mark.parametrize("column", ["gemm", "columns", "atomics"])
@pytest.mark.parametrize("correction_method", ["fast", "original", "batched"])
def test_harmony_integrate_reference(
    adata_reference, *, dtype, column, correction_method
):
    """
    Test that Harmony integrate works.
    """
    rsc.pp.harmony_integrate(
        adata_reference,
        "donor",
        correction_method=correction_method,
        dtype=dtype,
        colsum_algo=column,
        max_iter_harmony=20,
    )

    assert (
        _get_measure(
            adata_reference.obsm["harmony_org"],
            adata_reference.obsm["X_pca_harmony"],
            "L2",
        ).max()
        < 0.05
    )
    assert (
        _get_measure(
            adata_reference.obsm["harmony_org"],
            adata_reference.obsm["X_pca_harmony"],
            "r",
        ).min()
        > 0.95
    )


@pytest.mark.parametrize("n_cells", [1000, 60000])
@pytest.mark.parametrize("n_pcs", [20, 50])
@pytest.mark.parametrize("n_batches", [3, 10])
@pytest.mark.parametrize("switcher", [0, 1])
def test_scatter_add_shared_vs_optimized(n_cells, n_pcs, n_batches, switcher):
    """
    Test that shared memory and non-shared scatter add kernels produce identical results.

    Uses small integer values (as float32) for exact verification of correctness.
    """
    rng = np.random.default_rng(42)
    X_np = rng.integers(1, 10, size=(n_cells, n_pcs)).astype(np.float32)
    cats_np = rng.integers(0, n_batches, size=n_cells, dtype=np.int32)

    X = cp.asarray(X_np)
    cats = cp.asarray(cats_np)

    # Compute expected result using numpy
    expected_np = np.zeros((n_batches, n_pcs), dtype=np.float32)
    for i in range(n_cells):
        cat = cats_np[i]
        if switcher == 1:
            expected_np[cat] += X_np[i]
        else:
            expected_np[cat] -= X_np[i]
    expected = cp.asarray(expected_np)

    # Run optimized (non-shared) kernel via _scatter_add_cp
    out_optimized = cp.zeros((n_batches, n_pcs), dtype=cp.float32)
    _scatter_add_cp(X, out_optimized, cats, switcher, n_batches, use_shared=False)

    # Run shared memory kernel via _scatter_add_cp
    out_shared = cp.zeros((n_batches, n_pcs), dtype=cp.float32)
    _scatter_add_cp(X, out_shared, cats, switcher, n_batches, use_shared=True)

    # Both kernels should produce identical results
    cp.testing.assert_array_equal(out_optimized, expected)
    cp.testing.assert_array_equal(out_shared, expected)
    cp.testing.assert_array_equal(out_optimized, out_shared)
