import pytest

import numpy as np
import pandas as pd

import anndata as ad

from moscot.costs._costs import BarcodeDistance, _scaled_hamming_dist
from moscot.costs._utils import get_cost


class TestBarcodeDistance:
    RNG = np.random.RandomState(0)

    @staticmethod
    def test_barcode_distance_init():
        adata = ad.AnnData(TestBarcodeDistance.RNG.rand(3, 3), obsm={"barcodes": TestBarcodeDistance.RNG.rand(3, 3)})
        # initialization failure when no adata is provided
        with pytest.raises(TypeError):
            get_cost("barcode_distance", backend="moscot")
        # initialization failure when invalid key is provided
        with pytest.raises(KeyError):
            get_cost("barcode_distance", backend="moscot", adata=adata, key="invalid_key", attr="obsm")
        # initialization failure when invalid attr
        with pytest.raises(AttributeError):
            get_cost("barcode_distance", backend="moscot", adata=adata, key="barcodes", attr="invalid_attr")
        # check if not None
        cost_fn: BarcodeDistance = get_cost(
            "barcode_distance", backend="moscot", adata=adata, key="barcodes", attr="obsm"
        )
        assert cost_fn is not None

    @staticmethod
    def test_scaled_hamming_dist_with_sample_inputs():
        # Sample input arrays
        x = np.array([1, -1, 0, 1])
        y = np.array([0, 1, 1, 1])

        # Expected output
        expected_distance = 2.0 / 3

        # Compute the scaled Hamming distance
        computed_distance = _scaled_hamming_dist(x, y)

        # Check if the computed distance matches the expected distance
        np.testing.assert_almost_equal(computed_distance, expected_distance, decimal=4)

    @staticmethod
    def test_scaled_hamming_dist_if_nan():
        # Sample input arrays with no shared indices
        x = np.array([-1, -1, 0, 1])
        y = np.array([0, 1, -1, -1])

        with pytest.raises(ValueError, match="No shared indices."):
            _scaled_hamming_dist(x, y)

    @staticmethod
    def test_barcode_distance_with_sample_input():
        # Example barcodes
        barcodes = np.array([[1, 0, 1], [1, 1, 0], [0, 1, 1]])

        # Create a dummy AnnData object with the example barcodes
        adata = ad.AnnData(TestBarcodeDistance.RNG.rand(3, 3))
        adata.obsm["barcodes"] = barcodes

        # Initialize BarcodeDistance
        cost_fn: BarcodeDistance = get_cost(
            "barcode_distance", backend="moscot", adata=adata, key="barcodes", attr="obsm"
        )

        # Compute distances
        computed_distances = cost_fn()

        # Expected distance matrix
        expected_distances = np.array([[0.0, 2.0, 2.0], [2.0, 0.0, 2.0], [2.0, 2.0, 0.0]]) / 3.0

        # Check if the computed distances match the expected distances
        np.testing.assert_almost_equal(computed_distances, expected_distances, decimal=4)


class TestLeafDistance:
    @staticmethod
    def create_dummy_adata_leaf():
        import networkx as nx

        adata: ad.AnnData = ad.AnnData(
            X=np.ones((10, 10)),
            obs=pd.DataFrame(data={"day": [0, 0, 0, 1, 1, 1, 2, 2, 2, 2]}),
        )
        g: nx.DiGraph = nx.DiGraph()
        g.add_nodes_from([str(i) for i in range(3)] + ["root"])
        g.add_edges_from([("root", str(i)) for i in range(3)])
        adata.uns["tree"] = {0: g}
        return adata

    @staticmethod
    def test_leaf_distance_init():
        adata = TestLeafDistance.create_dummy_adata_leaf()
        # initialization failure when no adata is provided
        with pytest.raises(TypeError):
            get_cost("leaf_distance", backend="moscot")
        # initialization failure when invalid key is provided
        with pytest.raises(KeyError, match="Unable to find tree in"):
            get_cost("leaf_distance", backend="moscot", adata=adata, key="invalid_key", attr="uns", dist_key=0)
        # initialization failure when invalid dist_key is provided
        with pytest.raises(KeyError, match="Unable to find tree in"):
            get_cost("leaf_distance", backend="moscot", adata=adata, key="tree", attr="uns")
        # when leaves do not match adata.obs_names
        with pytest.raises(ValueError, match="Leaves do not match"):
            get_cost("leaf_distance", backend="moscot", adata=adata, key="tree", attr="uns", dist_key=0)()
        # now giving valid input
        adata0 = adata[adata.obs.day == 0]
        cost_fn = get_cost("leaf_distance", backend="moscot", adata=adata0, key="tree", attr="uns", dist_key=0)
        np.testing.assert_equal(cost_fn(), np.array([[0, 2, 2], [2, 0, 2], [2, 2, 0]]))
        # when tree is not a networkx.Graph
        adata0.uns["tree"] = {0: 1}
        with pytest.raises(TypeError, match="networkx.Graph"):
            get_cost("leaf_distance", backend="moscot", adata=adata0, key="tree", attr="uns", dist_key=0)