"""Primary split finding with NaN-direction learning.

Supports both exact (sort-scan) and histogram-based split search.
For each candidate split, evaluates placing NaN samples left vs right
and keeps the direction with better gain (XGBoost-style).
"""

import numpy as np

from ._impurity import (
    gini_impurity,
    entropy_impurity,
    mse_impurity,
    impurity_improvement,
    CLASSIFICATION_CRITERIA,
    REGRESSION_CRITERIA,
)
from ._node import SplitRecord
from ._histogram import NAN_BIN, histogram_subtract


def find_best_split_classification(
    X: np.ndarray,
    y: np.ndarray,
    sample_weight: np.ndarray,
    sample_indices: np.ndarray,
    feature_indices: np.ndarray,
    n_classes: int,
    criterion: str = "gini",
    min_samples_leaf: int = 1,
    min_weight_leaf: float = 0.0,
    min_impurity_decrease: float = 0.0,
) -> SplitRecord:
    """Find the best split across candidate features for classification."""
    criterion_fn = CLASSIFICATION_CRITERIA[criterion]
    best = SplitRecord()

    n_node = len(sample_indices)
    if n_node < 2 * min_samples_leaf:
        return best

    # Parent impurity
    y_node = y[sample_indices]
    w_node = sample_weight[sample_indices]
    parent_counts = np.bincount(y_node.astype(int), weights=w_node, minlength=n_classes).astype(np.float64)
    parent_impurity = criterion_fn(parent_counts)
    w_total = w_node.sum()

    for feat in feature_indices:
        col = X[sample_indices, feat]
        nan_mask = np.isnan(col)
        valid_mask = ~nan_mask

        n_valid = valid_mask.sum()
        if n_valid < 2 * min_samples_leaf:
            continue

        # Separate NaN and non-NaN
        valid_idx = sample_indices[valid_mask]
        nan_idx = sample_indices[nan_mask]
        n_nan = len(nan_idx)

        # NaN class counts
        if n_nan > 0:
            nan_counts = np.bincount(y[nan_idx].astype(int), weights=sample_weight[nan_idx], minlength=n_classes).astype(np.float64)
            w_nan = sample_weight[nan_idx].sum()
        else:
            nan_counts = np.zeros(n_classes, dtype=np.float64)
            w_nan = 0.0

        # Sort valid samples by feature value
        col_valid = X[valid_idx, feat]
        sort_order = np.argsort(col_valid)
        sorted_idx = valid_idx[sort_order]
        sorted_vals = col_valid[sort_order]

        # Cumulative scan
        left_counts = np.zeros(n_classes, dtype=np.float64)
        valid_counts = parent_counts - nan_counts
        right_counts = valid_counts.copy()
        w_left = 0.0
        w_right = sample_weight[sorted_idx].sum()

        for i in range(n_valid - 1):
            si = sorted_idx[i]
            cls = int(y[si])
            w = sample_weight[si]

            left_counts[cls] += w
            right_counts[cls] -= w
            w_left += w
            w_right -= w

            # Skip if same value as next
            if sorted_vals[i] == sorted_vals[i + 1]:
                continue

            # Check min_samples_leaf (approximate via weights)
            n_left_approx = i + 1
            n_right_approx = n_valid - i - 1

            # Evaluate NaN going left
            if n_nan > 0:
                left_plus_nan = left_counts + nan_counts
                w_left_nan = w_left + w_nan
                w_right_nan = w_right

                if w_left_nan > 0 and w_right_nan > 0:
                    imp_left = criterion_fn(left_plus_nan)
                    imp_right = criterion_fn(right_counts)
                    gain_nan_left = impurity_improvement(
                        w_total, parent_impurity,
                        w_left_nan, imp_left,
                        w_right_nan, imp_right,
                    )
                else:
                    gain_nan_left = -np.inf

                # NaN going right
                right_plus_nan = right_counts + nan_counts
                w_left_nonan = w_left
                w_right_nonan = w_right + w_nan

                if w_left_nonan > 0 and w_right_nonan > 0:
                    imp_left2 = criterion_fn(left_counts)
                    imp_right2 = criterion_fn(right_plus_nan)
                    gain_nan_right = impurity_improvement(
                        w_total, parent_impurity,
                        w_left_nonan, imp_left2,
                        w_right_nonan, imp_right2,
                    )
                else:
                    gain_nan_right = -np.inf

                if gain_nan_left >= gain_nan_right:
                    gain = gain_nan_left
                    missing_dir = -1
                    final_imp_left = imp_left if gain_nan_left > -np.inf else 0.0
                    final_imp_right = imp_right if gain_nan_left > -np.inf else 0.0
                    final_n_left = n_left_approx + n_nan
                    final_n_right = n_right_approx
                else:
                    gain = gain_nan_right
                    missing_dir = 1
                    final_imp_left = imp_left2 if gain_nan_right > -np.inf else 0.0
                    final_imp_right = imp_right2 if gain_nan_right > -np.inf else 0.0
                    final_n_left = n_left_approx
                    final_n_right = n_right_approx + n_nan
            else:
                # No NaN samples
                if w_left == 0 or w_right == 0:
                    continue
                imp_left = criterion_fn(left_counts)
                imp_right = criterion_fn(right_counts)
                gain = impurity_improvement(
                    w_total, parent_impurity,
                    w_left, imp_left,
                    w_right, imp_right,
                )
                missing_dir = -1
                final_imp_left = imp_left
                final_imp_right = imp_right
                final_n_left = n_left_approx
                final_n_right = n_right_approx

            if gain < min_impurity_decrease:
                continue

            if gain > best.gain:
                threshold = (sorted_vals[i] + sorted_vals[i + 1]) / 2.0
                best = SplitRecord(
                    feature=feat,
                    threshold=threshold,
                    gain=gain,
                    missing_direction=missing_dir,
                    n_left=final_n_left,
                    n_right=final_n_right,
                    impurity_left=final_imp_left,
                    impurity_right=final_imp_right,
                )

    return best


def find_best_split_regression(
    X: np.ndarray,
    y: np.ndarray,
    sample_weight: np.ndarray,
    sample_indices: np.ndarray,
    feature_indices: np.ndarray,
    criterion: str = "squared_error",
    min_samples_leaf: int = 1,
    min_weight_leaf: float = 0.0,
    min_impurity_decrease: float = 0.0,
) -> SplitRecord:
    """Find the best split across candidate features for regression."""
    best = SplitRecord()
    use_mse = criterion == "squared_error"

    n_node = len(sample_indices)
    if n_node < 2 * min_samples_leaf:
        return best

    y_node = y[sample_indices]
    w_node = sample_weight[sample_indices]
    w_total = w_node.sum()

    if use_mse:
        parent_impurity = _weighted_variance(y_node, w_node)
    else:
        parent_impurity = _weighted_mae(y_node, w_node)

    for feat in feature_indices:
        col = X[sample_indices, feat]
        nan_mask = np.isnan(col)
        valid_mask = ~nan_mask

        n_valid = valid_mask.sum()
        if n_valid < 2 * min_samples_leaf:
            continue

        valid_idx = sample_indices[valid_mask]
        nan_idx = sample_indices[nan_mask]
        n_nan = len(nan_idx)

        if n_nan > 0:
            nan_y = y[nan_idx]
            nan_w = sample_weight[nan_idx]
            nan_sum_y = (nan_y * nan_w).sum()
            nan_sum_y2 = (nan_y ** 2 * nan_w).sum()
            w_nan = nan_w.sum()
        else:
            nan_sum_y = 0.0
            nan_sum_y2 = 0.0
            w_nan = 0.0

        col_valid = X[valid_idx, feat]
        sort_order = np.argsort(col_valid)
        sorted_idx = valid_idx[sort_order]
        sorted_vals = col_valid[sort_order]

        if use_mse:
            _scan_mse(
                sorted_idx, sorted_vals, y, sample_weight,
                n_valid, n_nan, w_total, w_nan,
                nan_sum_y, nan_sum_y2,
                parent_impurity, feat,
                min_samples_leaf, min_impurity_decrease,
                best,
            )
        else:
            _scan_mae(
                sorted_idx, sorted_vals, y, sample_weight,
                n_valid, n_nan, w_total, w_nan,
                nan_idx, parent_impurity, feat,
                min_samples_leaf, min_impurity_decrease,
                best,
            )

    return best


def _scan_mse(
    sorted_idx, sorted_vals, y, sample_weight,
    n_valid, n_nan, w_total, w_nan,
    nan_sum_y, nan_sum_y2,
    parent_impurity, feat,
    min_samples_leaf, min_impurity_decrease,
    best,
):
    """MSE-based split scan using cumulative sums."""
    sum_left = 0.0
    sum_left_sq = 0.0
    w_left = 0.0

    sorted_w = sample_weight[sorted_idx]
    sorted_y = y[sorted_idx]

    total_sum = (sorted_y * sorted_w).sum()
    total_sum_sq = (sorted_y ** 2 * sorted_w).sum()
    w_valid = sorted_w.sum()

    for i in range(n_valid - 1):
        w = sorted_w[i]
        val = sorted_y[i]
        sum_left += w * val
        sum_left_sq += w * val * val
        w_left += w

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

        w_right = w_valid - w_left
        sum_right = total_sum - sum_left
        sum_right_sq = total_sum_sq - sum_left_sq

        # Evaluate both NaN directions
        for nan_dir in ([-1, 1] if n_nan > 0 else [-1]):
            if nan_dir == -1:
                wl = w_left + w_nan
                sl = sum_left + nan_sum_y
                sl2 = sum_left_sq + nan_sum_y2
                wr = w_right
                sr = sum_right
                sr2 = sum_right_sq
            else:
                wl = w_left
                sl = sum_left
                sl2 = sum_left_sq
                wr = w_right + w_nan
                sr = sum_right + nan_sum_y
                sr2 = sum_right_sq + nan_sum_y2

            if wl <= 0 or wr <= 0:
                continue

            imp_l = sl2 / wl - (sl / wl) ** 2
            imp_r = sr2 / wr - (sr / wr) ** 2
            gain = impurity_improvement(
                w_total, parent_impurity, wl, imp_l, wr, imp_r
            )

            if gain < min_impurity_decrease:
                continue

            if gain > best.gain:
                threshold = (sorted_vals[i] + sorted_vals[i + 1]) / 2.0
                n_left = (i + 1) + (n_nan if nan_dir == -1 else 0)
                n_right = (n_valid - i - 1) + (n_nan if nan_dir == 1 else 0)
                best.feature = feat
                best.threshold = threshold
                best.gain = gain
                best.missing_direction = nan_dir
                best.n_left = n_left
                best.n_right = n_right
                best.impurity_left = max(imp_l, 0.0)
                best.impurity_right = max(imp_r, 0.0)


def _scan_mae(
    sorted_idx, sorted_vals, y, sample_weight,
    n_valid, n_nan, w_total, w_nan,
    nan_idx, parent_impurity, feat,
    min_samples_leaf, min_impurity_decrease,
    best,
):
    """MAE-based split scan (brute force — O(n^2) but correct)."""
    for i in range(n_valid - 1):
        if sorted_vals[i] == sorted_vals[i + 1]:
            continue

        left_idx = sorted_idx[: i + 1]
        right_idx = sorted_idx[i + 1 :]

        for nan_dir in ([-1, 1] if n_nan > 0 else [-1]):
            if nan_dir == -1:
                li = np.concatenate([left_idx, nan_idx]) if n_nan > 0 else left_idx
                ri = right_idx
            else:
                li = left_idx
                ri = np.concatenate([right_idx, nan_idx])

            wl = sample_weight[li].sum()
            wr = sample_weight[ri].sum()
            if wl <= 0 or wr <= 0:
                continue

            imp_l = _weighted_mae(y[li], sample_weight[li])
            imp_r = _weighted_mae(y[ri], sample_weight[ri])
            gain = impurity_improvement(
                w_total, parent_impurity, wl, imp_l, wr, imp_r
            )

            if gain < min_impurity_decrease:
                continue
            if gain > best.gain:
                threshold = (sorted_vals[i] + sorted_vals[i + 1]) / 2.0
                best.feature = feat
                best.threshold = threshold
                best.gain = gain
                best.missing_direction = nan_dir
                best.n_left = len(li)
                best.n_right = len(ri)
                best.impurity_left = imp_l
                best.impurity_right = imp_r


def find_best_split_histogram_classification(
    X_binned: np.ndarray,
    y: np.ndarray,
    sample_weight: np.ndarray,
    sample_indices: np.ndarray,
    feature_indices: np.ndarray,
    n_classes: int,
    n_bins_per_feature: np.ndarray,
    bin_edges: list,
    criterion: str = "gini",
    min_samples_leaf: int = 1,
    min_impurity_decrease: float = 0.0,
    parent_histograms: dict = None,
    sibling_histograms: dict = None,
) -> tuple:
    """Histogram-based best split for classification.

    Returns (SplitRecord, dict_of_histograms).
    """
    criterion_fn = CLASSIFICATION_CRITERIA[criterion]
    best = SplitRecord()
    node_histograms = {}

    n_node = len(sample_indices)
    if n_node < 2 * min_samples_leaf:
        return best, node_histograms

    y_node = y[sample_indices]
    w_node = sample_weight[sample_indices]
    parent_counts = np.bincount(y_node.astype(int), weights=w_node, minlength=n_classes).astype(np.float64)
    parent_impurity = criterion_fn(parent_counts)
    w_total = w_node.sum()

    for feat in feature_indices:
        n_bins = n_bins_per_feature[feat]

        # Build or derive histogram
        if sibling_histograms is not None and feat in sibling_histograms and parent_histograms is not None and feat in parent_histograms:
            hist = histogram_subtract(parent_histograms[feat], sibling_histograms[feat])
        else:
            hist = np.zeros((n_bins + 1, n_classes), dtype=np.float64)
            for si in sample_indices:
                b = X_binned[si, feat]
                row = n_bins if b == NAN_BIN else b
                hist[row, int(y[si])] += sample_weight[si]

        node_histograms[feat] = hist

        # NaN bin counts
        nan_counts = hist[n_bins]
        w_nan = nan_counts.sum()

        # Scan bins left to right
        left_counts = np.zeros(n_classes, dtype=np.float64)
        valid_counts = parent_counts - nan_counts
        w_left = 0.0

        for b in range(n_bins):
            left_counts += hist[b]
            w_left += hist[b].sum()
            w_right = valid_counts.sum() - w_left

            if w_left <= 0 or w_right <= 0:
                continue

            right_counts = valid_counts - left_counts

            # NaN left vs right
            for nan_dir in ([-1, 1] if w_nan > 0 else [-1]):
                if nan_dir == -1:
                    lc = left_counts + nan_counts
                    rc = right_counts
                    wl = w_left + w_nan
                    wr = w_right
                else:
                    lc = left_counts
                    rc = right_counts + nan_counts
                    wl = w_left
                    wr = w_right + w_nan

                if wl <= 0 or wr <= 0:
                    continue

                imp_l = criterion_fn(lc)
                imp_r = criterion_fn(rc)
                gain = impurity_improvement(w_total, parent_impurity, wl, imp_l, wr, imp_r)

                if gain < min_impurity_decrease:
                    continue

                if gain > best.gain:
                    edges = bin_edges[feat]
                    if b < len(edges):
                        threshold = edges[b]
                    elif len(edges) > 0:
                        threshold = edges[-1]
                    else:
                        continue

                    best = SplitRecord(
                        feature=feat,
                        threshold=threshold,
                        gain=gain,
                        missing_direction=nan_dir,
                        impurity_left=imp_l,
                        impurity_right=imp_r,
                    )

    return best, node_histograms


def find_best_split_histogram_regression(
    X_binned: np.ndarray,
    y: np.ndarray,
    sample_weight: np.ndarray,
    sample_indices: np.ndarray,
    feature_indices: np.ndarray,
    n_bins_per_feature: np.ndarray,
    bin_edges: list,
    criterion: str = "squared_error",
    min_samples_leaf: int = 1,
    min_impurity_decrease: float = 0.0,
    parent_histograms: dict = None,
    sibling_histograms: dict = None,
) -> tuple:
    """Histogram-based best split for regression."""
    best = SplitRecord()
    node_histograms = {}

    n_node = len(sample_indices)
    if n_node < 2 * min_samples_leaf:
        return best, node_histograms

    y_node = y[sample_indices]
    w_node = sample_weight[sample_indices]
    w_total = w_node.sum()
    parent_impurity = _weighted_variance(y_node, w_node) if criterion == "squared_error" else _weighted_mae(y_node, w_node)

    for feat in feature_indices:
        n_bins = n_bins_per_feature[feat]

        if sibling_histograms is not None and feat in sibling_histograms and parent_histograms is not None and feat in parent_histograms:
            hist = histogram_subtract(parent_histograms[feat], sibling_histograms[feat])
        else:
            # (sum_y, sum_y2, count) per bin + NaN bin at index n_bins
            hist = np.zeros((n_bins + 1, 3), dtype=np.float64)
            for si in sample_indices:
                b = X_binned[si, feat]
                row = n_bins if b == NAN_BIN else b
                w = sample_weight[si]
                v = y[si]
                hist[row, 0] += w * v
                hist[row, 1] += w * v * v
                hist[row, 2] += w

        node_histograms[feat] = hist

        nan_row = hist[n_bins]
        nan_sum_y, nan_sum_y2, w_nan = nan_row

        sum_left = 0.0
        sum_left_sq = 0.0
        w_left = 0.0

        total_valid_sum = hist[:n_bins, 0].sum()
        total_valid_sum_sq = hist[:n_bins, 1].sum()
        w_valid = hist[:n_bins, 2].sum()

        for b in range(n_bins):
            sum_left += hist[b, 0]
            sum_left_sq += hist[b, 1]
            w_left += hist[b, 2]

            w_right = w_valid - w_left
            if w_left <= 0 or w_right <= 0:
                continue

            sum_right = total_valid_sum - sum_left
            sum_right_sq = total_valid_sum_sq - sum_left_sq

            for nan_dir in ([-1, 1] if w_nan > 0 else [-1]):
                if nan_dir == -1:
                    wl = w_left + w_nan
                    sl = sum_left + nan_sum_y
                    sl2 = sum_left_sq + nan_sum_y2
                    wr = w_right
                    sr = sum_right
                    sr2 = sum_right_sq
                else:
                    wl = w_left
                    sl = sum_left
                    sl2 = sum_left_sq
                    wr = w_right + w_nan
                    sr = sum_right + nan_sum_y
                    sr2 = sum_right_sq + nan_sum_y2

                if wl <= 0 or wr <= 0:
                    continue

                imp_l = sl2 / wl - (sl / wl) ** 2
                imp_r = sr2 / wr - (sr / wr) ** 2
                gain = impurity_improvement(w_total, parent_impurity, wl, max(imp_l, 0), wr, max(imp_r, 0))

                if gain < min_impurity_decrease:
                    continue

                if gain > best.gain:
                    edges = bin_edges[feat]
                    if b < len(edges):
                        threshold = edges[b]
                    elif len(edges) > 0:
                        threshold = edges[-1]
                    else:
                        continue

                    best = SplitRecord(
                        feature=feat,
                        threshold=threshold,
                        gain=gain,
                        missing_direction=nan_dir,
                        impurity_left=max(imp_l, 0),
                        impurity_right=max(imp_r, 0),
                    )

    return best, node_histograms


def _weighted_variance(y, w):
    total_w = w.sum()
    if total_w == 0:
        return 0.0
    mean = np.average(y, weights=w)
    return np.average((y - mean) ** 2, weights=w)


def _weighted_mae(y, w):
    total_w = w.sum()
    if total_w == 0:
        return 0.0
    sorted_idx = np.argsort(y)
    cum_w = np.cumsum(w[sorted_idx])
    median_idx = np.searchsorted(cum_w, total_w / 2.0)
    median_idx = min(median_idx, len(y) - 1)
    median = y[sorted_idx[median_idx]]
    return np.average(np.abs(y - median), weights=w)
