"""
Core geometry extraction from decision tree path structure.

Implements the DirectRS algorithm from Sudjianto's "Linear Models Hiding
Inside Your Decision Tree", adapted for both classification and regression.
"""

from __future__ import annotations

import numpy as np
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor

# Sentinel value in sklearn's tree arrays for "no node"
_TREE_LEAF = -1
_TREE_UNDEFINED = -2


def build_parent_map(tree_) -> np.ndarray:
    """Return array parent[i] = parent node index of node i, -1 for root."""
    n_nodes = tree_.node_count
    parent = np.full(n_nodes, -1, dtype=np.intp)
    children_left = tree_.children_left
    children_right = tree_.children_right
    for node in range(n_nodes):
        l, r = children_left[node], children_right[node]
        if l != _TREE_LEAF:
            parent[l] = node
        if r != _TREE_LEAF:
            parent[r] = node
    return parent


def get_leaf_indices(tree_) -> list[int]:
    """Return sorted list of leaf node indices."""
    return [
        i
        for i in range(tree_.node_count)
        if tree_.children_left[i] == _TREE_LEAF
    ]


def get_path_features(tree_, parent_map: np.ndarray, leaf_idx: int) -> set[int]:
    """
    Walk root → leaf and collect the set of feature indices split on.

    Returns a set of integer feature indices (0-based).
    """
    features: set[int] = set()
    node = leaf_idx
    while parent_map[node] != -1:
        par = parent_map[node]
        feat = tree_.feature[par]
        if feat >= 0:  # skip undefined/-2 (shouldn't happen on internal nodes)
            features.add(int(feat))
        node = par
    return features


def _leaf_values(tree_, task: str) -> np.ndarray:
    """
    Return a 1-D array of leaf values.

    Classification → class label (argmax of node value counts)
    Regression     → mean target stored in tree node
    """
    leaf_idxs = get_leaf_indices(tree_)
    values = []
    for i in leaf_idxs:
        v = tree_.value[i]  # shape (1, n_outputs, max_n_classes) or similar
        if task == "classification":
            # v shape: (1, n_classes) for a single-output tree
            values.append(float(np.argmax(v[0])))
        else:
            # Regression: mean is stored directly
            values.append(float(v[0, 0]))
    return np.array(values)


def compute_G_matrix(
    tree_,
    task: str = "classification",
) -> np.ndarray:
    """
    Compute the d×d geometry matrix G = Σ_ℓ w_ℓ A_ℓ.

    A_ℓ = (v_ℓ - v̄)² × outer(p_ℓ, p_ℓ)

    where:
      - p_ℓ ∈ {0,1}^d  binary path indicator (features split on root→leaf ℓ)
      - v_ℓ            leaf value (class label or regression mean)
      - v̄              cover-weighted mean leaf value
      - w_ℓ            fraction of training samples reaching leaf ℓ

    Returns
    -------
    G : ndarray of shape (d, d)
        Symmetric positive semi-definite geometry matrix.
    """
    n_features = tree_.n_features
    leaf_idxs = get_leaf_indices(tree_)
    parent_map = build_parent_map(tree_)
    n_total = tree_.n_node_samples[0]  # root sample count

    leaf_values = _leaf_values(tree_, task)
    leaf_weights = np.array(
        [tree_.n_node_samples[i] / n_total for i in leaf_idxs]
    )

    # Cover-weighted mean leaf value
    v_bar = float(np.dot(leaf_weights, leaf_values))

    G = np.zeros((n_features, n_features), dtype=float)

    for k, (leaf, v_l, w_l) in enumerate(zip(leaf_idxs, leaf_values, leaf_weights)):
        feat_set = get_path_features(tree_, parent_map, leaf)
        p_l = np.zeros(n_features, dtype=float)
        for f in feat_set:
            p_l[f] = 1.0
        A_l = (v_l - v_bar) ** 2 * np.outer(p_l, p_l)
        G += w_l * A_l

    return G


def compute_stretch_matrix(G: np.ndarray) -> np.ndarray:
    """
    Compute the stretch matrix S' = G^(1/2) via eigendecomposition.

    S' = V diag(sqrt(max(λ, 0))) Vᵀ

    Returns
    -------
    S_prime : ndarray of shape (d, d)
    """
    # Symmetrize to guard against tiny numerical asymmetry
    G_sym = (G + G.T) / 2.0
    eigenvalues, eigenvectors = np.linalg.eigh(G_sym)
    # Clamp to non-negative before sqrt
    sqrt_eigenvalues = np.sqrt(np.maximum(eigenvalues, 0.0))
    S_prime = eigenvectors @ np.diag(sqrt_eigenvalues) @ eigenvectors.T
    return S_prime


def build_embedding(
    X: np.ndarray,
    tree: DecisionTreeClassifier | DecisionTreeRegressor,
    S_prime: np.ndarray,
    leaf_indices: list[int],
) -> np.ndarray:
    """
    Build the piecewise-linear embedding matrix Φ of shape (n, L×(d+1)).

    For sample x landing in the pos-th leaf:
      - Block at column pos*(d+1) : (d+1) → [1, S'x]
      - All other L-1 blocks       → zeros

    Parameters
    ----------
    X : ndarray (n, d)
    tree : fitted sklearn decision tree
    S_prime : ndarray (d, d)  stretch matrix
    leaf_indices : list[int]  ordered list of leaf node indices

    Returns
    -------
    Phi : ndarray (n, L*(d+1))
    """
    n, d = X.shape
    L = len(leaf_indices)
    block_size = d + 1
    Phi = np.zeros((n, L * block_size), dtype=float)

    # Map leaf node index → position in leaf_indices list
    leaf_to_pos = {leaf: pos for pos, leaf in enumerate(leaf_indices)}

    leaf_node_ids = tree.apply(X)  # shape (n,)

    # Precompute S'x for all samples: shape (n, d)
    S_prime_X = X @ S_prime.T  # equivalent to (S'x)ᵀ

    for i in range(n):
        leaf = leaf_node_ids[i]
        pos = leaf_to_pos[leaf]
        start = pos * block_size
        Phi[i, start] = 1.0  # intercept
        Phi[i, start + 1 : start + block_size] = S_prime_X[i]

    return Phi
