"""Unit tests for directrs.geometry."""

import numpy as np
import pytest
from sklearn.datasets import load_iris, load_diabetes
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor

from directrs.geometry import (
    build_parent_map,
    build_embedding,
    compute_G_matrix,
    compute_stretch_matrix,
    get_leaf_indices,
    get_path_features,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def small_clf_tree():
    """Tiny classification tree on iris (max_depth=2)."""
    X, y = load_iris(return_X_y=True)
    clf = DecisionTreeClassifier(max_depth=2, random_state=0)
    clf.fit(X, y)
    return clf, X, y


@pytest.fixture
def small_reg_tree():
    """Tiny regression tree on diabetes (max_depth=2)."""
    X, y = load_diabetes(return_X_y=True)
    reg = DecisionTreeRegressor(max_depth=2, random_state=0)
    reg.fit(X, y)
    return reg, X, y


# ---------------------------------------------------------------------------
# build_parent_map
# ---------------------------------------------------------------------------

def test_build_parent_map_root_has_no_parent(small_clf_tree):
    clf, _, _ = small_clf_tree
    parent = build_parent_map(clf.tree_)
    assert parent[0] == -1, "Root node should have no parent"


def test_build_parent_map_all_internal_nodes_are_parents(small_clf_tree):
    clf, _, _ = small_clf_tree
    tree_ = clf.tree_
    parent = build_parent_map(tree_)
    for node in range(tree_.node_count):
        left = tree_.children_left[node]
        right = tree_.children_right[node]
        if left != -1:
            assert parent[left] == node
        if right != -1:
            assert parent[right] == node


def test_build_parent_map_length(small_clf_tree):
    clf, _, _ = small_clf_tree
    parent = build_parent_map(clf.tree_)
    assert len(parent) == clf.tree_.node_count


# ---------------------------------------------------------------------------
# get_leaf_indices
# ---------------------------------------------------------------------------

def test_get_leaf_indices_are_leaves(small_clf_tree):
    clf, _, _ = small_clf_tree
    tree_ = clf.tree_
    leaves = get_leaf_indices(tree_)
    for leaf in leaves:
        assert tree_.children_left[leaf] == -1
        assert tree_.children_right[leaf] == -1


def test_get_leaf_indices_count_matches_depth2():
    """A complete binary tree of depth 2 has 4 leaves."""
    X, y = load_iris(return_X_y=True)
    clf = DecisionTreeClassifier(max_depth=2, random_state=0)
    clf.fit(X, y)
    leaves = get_leaf_indices(clf.tree_)
    # At depth 2, at most 4 leaves; iris may have fewer if some branches are pure
    assert 1 <= len(leaves) <= 4


# ---------------------------------------------------------------------------
# get_path_features
# ---------------------------------------------------------------------------

def test_get_path_features_subset_of_all_features(small_clf_tree):
    clf, _, _ = small_clf_tree
    tree_ = clf.tree_
    parent = build_parent_map(tree_)
    n_features = tree_.n_features
    for leaf in get_leaf_indices(tree_):
        feats = get_path_features(tree_, parent, leaf)
        assert feats.issubset(set(range(n_features)))


def test_get_path_features_returns_set(small_clf_tree):
    clf, _, _ = small_clf_tree
    tree_ = clf.tree_
    parent = build_parent_map(tree_)
    leaf = get_leaf_indices(tree_)[0]
    result = get_path_features(tree_, parent, leaf)
    assert isinstance(result, set)


# ---------------------------------------------------------------------------
# compute_G_matrix
# ---------------------------------------------------------------------------

def test_compute_G_matrix_shape_classification(small_clf_tree):
    clf, _, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    d = clf.tree_.n_features
    assert G.shape == (d, d)


def test_compute_G_matrix_psd_classification(small_clf_tree):
    clf, _, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    eigenvalues = np.linalg.eigvalsh(G)
    assert np.all(eigenvalues >= -1e-10), f"G is not PSD: min eigenvalue = {eigenvalues.min()}"


def test_compute_G_matrix_symmetric(small_clf_tree):
    clf, _, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    np.testing.assert_allclose(G, G.T, atol=1e-12)


def test_compute_G_matrix_shape_regression(small_reg_tree):
    reg, _, _ = small_reg_tree
    G = compute_G_matrix(reg.tree_, task="regression")
    d = reg.tree_.n_features
    assert G.shape == (d, d)


def test_compute_G_matrix_psd_regression(small_reg_tree):
    reg, _, _ = small_reg_tree
    G = compute_G_matrix(reg.tree_, task="regression")
    eigenvalues = np.linalg.eigvalsh(G)
    assert np.all(eigenvalues >= -1e-10)


# ---------------------------------------------------------------------------
# compute_stretch_matrix
# ---------------------------------------------------------------------------

def test_compute_stretch_matrix_shape(small_clf_tree):
    clf, _, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    S = compute_stretch_matrix(G)
    assert S.shape == G.shape


def test_compute_stretch_matrix_psd(small_clf_tree):
    clf, _, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    S = compute_stretch_matrix(G)
    eigenvalues = np.linalg.eigvalsh(S)
    assert np.all(eigenvalues >= -1e-10), f"S' not PSD: min = {eigenvalues.min()}"


def test_compute_stretch_matrix_symmetric(small_clf_tree):
    clf, _, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    S = compute_stretch_matrix(G)
    np.testing.assert_allclose(S, S.T, atol=1e-10)


def test_compute_stretch_matrix_squared_approx_G(small_clf_tree):
    """S' @ S' ≈ G (by construction of matrix square root)."""
    clf, _, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    S = compute_stretch_matrix(G)
    np.testing.assert_allclose(S @ S, G, atol=1e-8)


# ---------------------------------------------------------------------------
# build_embedding
# ---------------------------------------------------------------------------

def test_build_embedding_shape(small_clf_tree):
    clf, X, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    S = compute_stretch_matrix(G)
    leaves = get_leaf_indices(clf.tree_)
    Phi = build_embedding(X, clf, S, leaves)
    n, d = X.shape
    L = len(leaves)
    assert Phi.shape == (n, L * (d + 1))


def test_build_embedding_one_nonzero_block_per_row(small_clf_tree):
    """Each row of Phi should have exactly one nonzero block."""
    clf, X, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    S = compute_stretch_matrix(G)
    leaves = get_leaf_indices(clf.tree_)
    Phi = build_embedding(X, clf, S, leaves)
    d = X.shape[1]
    block_size = d + 1
    L = len(leaves)
    for i in range(len(X)):
        nonzero_blocks = 0
        for b in range(L):
            block = Phi[i, b * block_size : (b + 1) * block_size]
            if np.any(block != 0):
                nonzero_blocks += 1
        assert nonzero_blocks == 1, f"Row {i} has {nonzero_blocks} non-zero blocks"


def test_build_embedding_intercept_is_one(small_clf_tree):
    """The intercept (first element of the active block) should be 1."""
    clf, X, _ = small_clf_tree
    G = compute_G_matrix(clf.tree_, task="classification")
    S = compute_stretch_matrix(G)
    leaves = get_leaf_indices(clf.tree_)
    leaf_to_pos = {leaf: pos for pos, leaf in enumerate(leaves)}
    Phi = build_embedding(X, clf, S, leaves)
    d = X.shape[1]
    block_size = d + 1
    leaf_ids = clf.apply(X)
    for i in range(min(20, len(X))):
        pos = leaf_to_pos[leaf_ids[i]]
        assert Phi[i, pos * block_size] == 1.0
