"""Surrogate split discovery using CART predictive measure of association.

After a primary split is chosen, this module finds backup splits on other
features that best mimic the primary's partitioning. When a sample is
missing the primary split feature, surrogates route it through the tree.
"""

import numpy as np

from ._node import SurrogateRecord


def find_surrogates(
    X: np.ndarray,
    sample_indices: np.ndarray,
    primary_feature: int,
    primary_threshold: float,
    primary_missing_direction: int,
    max_surrogates: int = 5,
    all_feature_indices: np.ndarray = None,
    sample_weight: np.ndarray = None,
) -> list:
    """Find surrogate splits for a primary split.

    For each candidate feature, finds the threshold that maximizes
    weighted concordance with the primary split's left/right assignment.
    Computes the predictive measure of association (lambda) and
    supports mirrored surrogates for negatively correlated features.
    """
    if max_surrogates <= 0:
        return []

    n_features = X.shape[1]
    if all_feature_indices is None:
        all_feature_indices = np.arange(n_features)

    # Determine primary's left/right for non-NaN samples
    primary_vals = X[sample_indices, primary_feature]
    primary_valid = ~np.isnan(primary_vals)

    if primary_valid.sum() < 2:
        return []

    valid_node_idx = sample_indices[primary_valid]
    primary_goes_left = X[valid_node_idx, primary_feature] <= primary_threshold

    # Use weights if provided
    if sample_weight is not None:
        weights = sample_weight[valid_node_idx]
    else:
        weights = np.ones(len(valid_node_idx), dtype=np.float64)

    w_total = weights.sum()
    w_left_primary = weights[primary_goes_left].sum()
    w_right_primary = w_total - w_left_primary

    # Naive error rate
    p_majority = max(w_left_primary, w_right_primary) / w_total
    p_naive_error = 1.0 - p_majority

    if p_naive_error < 1e-10:
        return []

    surrogates = []

    for feat in all_feature_indices:
        if feat == primary_feature:
            continue

        surr_vals = X[valid_node_idx, feat]
        surr_valid = ~np.isnan(surr_vals)

        if surr_valid.sum() < 2:
            continue

        both_idx = np.where(surr_valid)[0]
        if len(both_idx) < 2:
            continue

        surr_col = surr_vals[both_idx]
        primary_left_here = primary_goes_left[both_idx]
        weights_here = weights[both_idx]
        w_both = weights_here.sum()

        # Sort by surrogate feature
        sort_order = np.argsort(surr_col)
        sorted_surr = surr_col[sort_order]
        sorted_primary_left = primary_left_here[sort_order]
        sorted_weights = weights_here[sort_order]

        best_lambda = -np.inf
        best_threshold = 0.0
        best_mirror = False

        # Weighted cumulative concordance
        cum_w_prim_left = 0.0  # weighted count of primary_left in surrogate's left
        cum_w_prim_right = 0.0  # weighted count of primary_right in surrogate's left
        total_w_prim_left = (sorted_weights * sorted_primary_left).sum()
        total_w_prim_right = w_both - total_w_prim_left

        for i in range(len(both_idx) - 1):
            w = sorted_weights[i]
            if sorted_primary_left[i]:
                cum_w_prim_left += w
            else:
                cum_w_prim_right += w

            if sorted_surr[i] == sorted_surr[i + 1]:
                continue

            # Non-mirrored weighted concordance
            agree_normal = cum_w_prim_left + (total_w_prim_right - cum_w_prim_right)
            disagree_normal = w_both - agree_normal

            # Mirrored weighted concordance
            agree_mirror = cum_w_prim_right + (total_w_prim_left - cum_w_prim_left)
            disagree_mirror = w_both - agree_mirror

            if agree_normal >= agree_mirror:
                disagreement_rate = disagree_normal / w_both
                mirror = False
            else:
                disagreement_rate = disagree_mirror / w_both
                mirror = True

            lam = (p_naive_error - disagreement_rate) / p_naive_error

            if lam > best_lambda:
                best_lambda = lam
                best_threshold = (sorted_surr[i] + sorted_surr[i + 1]) / 2.0
                best_mirror = mirror

        if best_lambda > 0:
            surr_record = SurrogateRecord(
                feature=feat,
                threshold=best_threshold,
                mirror=best_mirror,
                lambda_=best_lambda,
                missing_direction=primary_missing_direction,
            )
            surrogates.append(surr_record)

    surrogates.sort(key=lambda s: -s.lambda_)
    return surrogates[:max_surrogates]
