"""Histogram-based binning for fast split finding (LightGBM-style)."""

import numpy as np

# Sentinel bin index for NaN values
NAN_BIN = 255


class HistogramBinner:
    """Quantile-based feature binning for histogram split finding.

    Bins each feature into at most `max_bins` quantile-based bins.
    NaN values get the sentinel bin index NAN_BIN (255).
    """

    def __init__(self, max_bins: int = 255):
        self.max_bins = min(max_bins, 254)  # Reserve 255 for NaN
        self.bin_edges_ = None  # list of arrays, one per feature
        self.n_features_ = 0

    def fit(self, X: np.ndarray, sample_weight: np.ndarray = None) -> "HistogramBinner":
        """Learn bin edges from training data.

        When sample_weight is provided, computes weighted quantiles so that
        bin edges are consistent between weighted fitting and fitting with
        duplicated/removed samples.
        """
        n_samples, n_features = X.shape
        self.n_features_ = n_features
        self.bin_edges_ = []

        for f in range(n_features):
            col = X[:, f]
            nan_mask = np.isnan(col)
            valid = col[~nan_mask]

            if len(valid) == 0:
                self.bin_edges_.append(np.array([]))
                continue

            unique_vals = np.unique(valid)
            n_unique = len(unique_vals)

            if n_unique <= 1:
                self.bin_edges_.append(np.array([valid.max()]))
                continue

            if n_unique <= self.max_bins:
                # Use midpoints between unique values as edges
                edges = (unique_vals[:-1] + unique_vals[1:]) / 2.0
            elif sample_weight is not None:
                # Weighted quantiles for consistency with duplicated samples
                valid_w = sample_weight[~nan_mask]
                sorted_idx = np.argsort(valid)
                sorted_vals = valid[sorted_idx]
                sorted_w = valid_w[sorted_idx]
                cum_w = np.cumsum(sorted_w)
                total_w = cum_w[-1]

                n_bins = min(self.max_bins, n_unique)
                percentiles = np.linspace(0, 1, n_bins + 1)[1:-1]
                edge_list = []
                for p in percentiles:
                    idx = np.searchsorted(cum_w, p * total_w)
                    idx = min(idx, len(sorted_vals) - 1)
                    edge_list.append(sorted_vals[idx])
                edges = np.unique(edge_list)
            else:
                n_bins = min(self.max_bins, n_unique)
                percentiles = np.linspace(0, 100, n_bins + 1)[1:-1]
                edges = np.unique(np.percentile(valid, percentiles))

            self.bin_edges_.append(edges)

        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        """Bin features into uint8 indices. NaN → NAN_BIN."""
        n_samples, n_features = X.shape
        X_binned = np.empty((n_samples, n_features), dtype=np.uint8)

        for f in range(n_features):
            col = X[:, f]
            nan_mask = np.isnan(col)
            edges = self.bin_edges_[f]

            if nan_mask.all():
                X_binned[:, f] = NAN_BIN
            elif len(edges) == 0:
                X_binned[:, f] = 0
                X_binned[nan_mask, f] = NAN_BIN
            elif nan_mask.any():
                binned = np.empty(n_samples, dtype=np.uint8)
                binned[~nan_mask] = np.searchsorted(edges, col[~nan_mask], side="right").astype(np.uint8)
                binned[nan_mask] = NAN_BIN
                X_binned[:, f] = binned
            else:
                X_binned[:, f] = np.searchsorted(edges, col, side="right").astype(np.uint8)

        return X_binned

    def fit_transform(self, X: np.ndarray, sample_weight: np.ndarray = None) -> np.ndarray:
        return self.fit(X, sample_weight=sample_weight).transform(X)

    def n_bins_per_feature(self) -> np.ndarray:
        """Return the number of non-NaN bins for each feature."""
        return np.array([len(edges) + 1 for edges in self.bin_edges_])


def build_histogram_classification(
    X_binned_col: np.ndarray,
    y: np.ndarray,
    sample_weight: np.ndarray,
    n_bins: int,
    n_classes: int,
    sample_indices: np.ndarray = None,
) -> np.ndarray:
    """Build histogram of class counts per bin.

    Returns shape (n_bins + 1, n_classes) — last row is NaN bin.
    """
    hist = np.zeros((n_bins + 1, n_classes), dtype=np.float64)

    if sample_indices is not None:
        bins = X_binned_col[sample_indices]
        labels = y[sample_indices]
        weights = sample_weight[sample_indices]
    else:
        bins = X_binned_col
        labels = y
        weights = sample_weight

    for i in range(len(bins)):
        b = bins[i]
        if b == NAN_BIN:
            row = n_bins  # Last row for NaN
        else:
            row = b
        hist[row, int(labels[i])] += weights[i]

    return hist


def build_histogram_regression(
    X_binned_col: np.ndarray,
    y: np.ndarray,
    sample_weight: np.ndarray,
    n_bins: int,
    sample_indices: np.ndarray = None,
) -> np.ndarray:
    """Build histogram of (sum_y, sum_y2, count) per bin.

    Returns shape (n_bins + 1, 3) — last row is NaN bin.
    """
    hist = np.zeros((n_bins + 1, 3), dtype=np.float64)

    if sample_indices is not None:
        bins = X_binned_col[sample_indices]
        targets = y[sample_indices]
        weights = sample_weight[sample_indices]
    else:
        bins = X_binned_col
        targets = y
        weights = sample_weight

    for i in range(len(bins)):
        b = bins[i]
        if b == NAN_BIN:
            row = n_bins
        else:
            row = b
        w = weights[i]
        val = targets[i]
        hist[row, 0] += w * val       # sum_y
        hist[row, 1] += w * val * val  # sum_y2
        hist[row, 2] += w             # count

    return hist


def histogram_subtract(parent_hist: np.ndarray, left_hist: np.ndarray) -> np.ndarray:
    """Compute right histogram via subtraction: right = parent - left."""
    return parent_hist - left_hist
