"""Tree builder supporting depth-first, best-first, and symmetric growth modes."""

import heapq

import numpy as np

from ._node import Node, SplitRecord
from ._impurity import CLASSIFICATION_CRITERIA, REGRESSION_CRITERIA, impurity_improvement
from ._splitter import (
    find_best_split_classification,
    find_best_split_regression,
    find_best_split_histogram_classification,
    find_best_split_histogram_regression,
)
from ._surrogate import find_surrogates
from ._histogram import HistogramBinner
from ._validation import compute_feature_indices


class TreeBuilder:
    """Builds a decision tree with surrogate splits.

    Supports three growth modes:
    - depth_first (default): Standard recursive CART
    - best_first (max_leaf_nodes set): LightGBM-style leaf-wise growth
    - symmetric (symmetric_tree=True): CatBoost-style oblivious trees
    """

    def __init__(
        self,
        is_classifier: bool,
        n_classes: int = 0,
        criterion: str = "gini",
        max_depth: int = None,
        min_samples_split: int = 2,
        min_samples_leaf: int = 1,
        min_weight_fraction_leaf: float = 0.0,
        max_features=None,
        max_leaf_nodes: int = None,
        min_impurity_decrease: float = 0.0,
        max_surrogates: int = 5,
        use_histogram: bool = True,
        max_bins: int = 255,
        symmetric_tree: bool = False,
        random_state=None,
    ):
        self.is_classifier = is_classifier
        self.n_classes = n_classes
        self.criterion = criterion
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.min_samples_leaf = min_samples_leaf
        self.min_weight_fraction_leaf = min_weight_fraction_leaf
        self.max_features = max_features
        self.max_leaf_nodes = max_leaf_nodes
        self.min_impurity_decrease = min_impurity_decrease
        self.max_surrogates = max_surrogates
        self.use_histogram = use_histogram
        self.max_bins = max_bins
        self.symmetric_tree = symmetric_tree
        self.random_state = random_state

        self.rng_ = np.random.RandomState(random_state)
        self.node_count_ = 0
        self.max_depth_reached_ = 0
        self.binner_ = None
        self.X_binned_ = None

    def build(self, X, y, sample_weight):
        """Build the tree and return the root Node."""
        n_samples, n_features = X.shape
        self.n_features_ = n_features
        self.n_max_features_ = compute_feature_indices(
            self.max_features, n_features, self.rng_
        )
        self.min_weight_leaf_ = self.min_weight_fraction_leaf * sample_weight.sum()

        # Histogram binning
        if self.use_histogram:
            self.binner_ = HistogramBinner(max_bins=self.max_bins)
            self.X_binned_ = self.binner_.fit_transform(X, sample_weight=sample_weight)
            self.n_bins_per_feature_ = self.binner_.n_bins_per_feature()

        sample_indices = np.arange(n_samples, dtype=np.intp)

        if self.symmetric_tree:
            return self._build_symmetric(X, y, sample_weight, sample_indices)
        elif self.max_leaf_nodes is not None:
            return self._build_best_first(X, y, sample_weight, sample_indices)
        else:
            return self._build_depth_first(X, y, sample_weight, sample_indices, depth=0)

    def _make_leaf(self, y, sample_weight, sample_indices, depth):
        """Create a leaf node."""
        node = Node(
            node_id=self._next_id(),
            depth=depth,
            is_leaf=True,
            n_samples=len(sample_indices),
            weighted_n_samples=sample_weight[sample_indices].sum(),
        )
        if self.is_classifier:
            counts = np.bincount(
                y[sample_indices].astype(int),
                weights=sample_weight[sample_indices],
                minlength=self.n_classes,
            ).astype(np.float64)
            node.value = counts
            node.impurity = CLASSIFICATION_CRITERIA[self.criterion](counts)
        else:
            ys = y[sample_indices]
            ws = sample_weight[sample_indices]
            w_total = ws.sum()
            if w_total > 0:
                mean = np.average(ys, weights=ws)
            else:
                mean = 0.0
            node.value = np.array([mean, w_total])
            if self.criterion == "squared_error":
                node.impurity = np.average((ys - mean) ** 2, weights=ws) if w_total > 0 else 0.0
            else:
                from ._splitter import _weighted_mae
                node.impurity = _weighted_mae(ys, ws) if w_total > 0 else 0.0
        self.max_depth_reached_ = max(self.max_depth_reached_, depth)
        return node

    def _should_stop(self, sample_indices, depth, sample_weight):
        """Check stopping criteria."""
        n = len(sample_indices)
        if n < self.min_samples_split:
            return True
        if self.max_depth is not None and depth >= self.max_depth:
            return True
        w = sample_weight[sample_indices].sum()
        if w <= 0:
            return True
        # Check if all targets are the same
        y_unique = np.unique(self.y_[sample_indices])
        if len(y_unique) <= 1:
            return True
        return False

    def _find_split(self, X, y, sample_weight, sample_indices):
        """Find the best split for a node."""
        feature_indices = self.rng_.choice(
            self.n_features_,
            size=self.n_max_features_,
            replace=False,
        )

        if self.use_histogram and self.is_classifier:
            split, _ = find_best_split_histogram_classification(
                self.X_binned_, y, sample_weight, sample_indices,
                feature_indices, self.n_classes,
                self.n_bins_per_feature_, self.binner_.bin_edges_,
                criterion=self.criterion,
                min_samples_leaf=self.min_samples_leaf,
                min_impurity_decrease=self.min_impurity_decrease,
            )
        elif self.use_histogram and not self.is_classifier:
            split, _ = find_best_split_histogram_regression(
                self.X_binned_, y, sample_weight, sample_indices,
                feature_indices,
                self.n_bins_per_feature_, self.binner_.bin_edges_,
                criterion=self.criterion,
                min_samples_leaf=self.min_samples_leaf,
                min_impurity_decrease=self.min_impurity_decrease,
            )
        elif self.is_classifier:
            split = find_best_split_classification(
                X, y, sample_weight, sample_indices,
                feature_indices, self.n_classes,
                criterion=self.criterion,
                min_samples_leaf=self.min_samples_leaf,
                min_impurity_decrease=self.min_impurity_decrease,
            )
        else:
            split = find_best_split_regression(
                X, y, sample_weight, sample_indices,
                feature_indices,
                criterion=self.criterion,
                min_samples_leaf=self.min_samples_leaf,
                min_impurity_decrease=self.min_impurity_decrease,
            )

        return split

    def _partition(self, X, sample_indices, split):
        """Partition samples into left/right based on split, handling NaN."""
        col = X[sample_indices, split.feature]
        nan_mask = np.isnan(col)
        valid_mask = ~nan_mask

        left_valid = col <= split.threshold
        right_valid = ~left_valid

        # Apply to valid only
        left_mask = np.zeros(len(sample_indices), dtype=bool)
        right_mask = np.zeros(len(sample_indices), dtype=bool)

        left_mask[valid_mask] = left_valid[valid_mask]
        right_mask[valid_mask] = right_valid[valid_mask]

        # NaN samples go to learned direction
        if split.missing_direction <= 0:
            left_mask[nan_mask] = True
        else:
            right_mask[nan_mask] = True

        return sample_indices[left_mask], sample_indices[right_mask]

    def _build_depth_first(self, X, y, sample_weight, sample_indices, depth):
        """Recursive depth-first tree building."""
        self.y_ = y  # Store for _should_stop

        if self._should_stop(sample_indices, depth, sample_weight):
            return self._make_leaf(y, sample_weight, sample_indices, depth)

        split = self._find_split(X, y, sample_weight, sample_indices)

        if split.gain <= 0 or split.feature < 0:
            return self._make_leaf(y, sample_weight, sample_indices, depth)

        # Find surrogates
        if self.max_surrogates > 0:
            split.surrogates = find_surrogates(
                X, sample_indices,
                split.feature, split.threshold, split.missing_direction,
                max_surrogates=self.max_surrogates,
                all_feature_indices=np.arange(self.n_features_),
                sample_weight=sample_weight,
            )

        left_idx, right_idx = self._partition(X, sample_indices, split)

        if len(left_idx) == 0 or len(right_idx) == 0:
            return self._make_leaf(y, sample_weight, sample_indices, depth)

        node = Node(
            node_id=self._next_id(),
            depth=depth,
            is_leaf=False,
            n_samples=len(sample_indices),
            weighted_n_samples=sample_weight[sample_indices].sum(),
            split=split,
        )

        if self.is_classifier:
            counts = np.bincount(
                y[sample_indices].astype(int),
                weights=sample_weight[sample_indices],
                minlength=self.n_classes,
            ).astype(np.float64)
            node.value = counts
            node.impurity = CLASSIFICATION_CRITERIA[self.criterion](counts)
        else:
            ys = y[sample_indices]
            ws = sample_weight[sample_indices]
            w_total = ws.sum()
            mean = np.average(ys, weights=ws) if w_total > 0 else 0.0
            node.value = np.array([mean, w_total])
            node.impurity = split.impurity_left  # Will be overridden

        node.left_child = self._build_depth_first(X, y, sample_weight, left_idx, depth + 1)
        node.right_child = self._build_depth_first(X, y, sample_weight, right_idx, depth + 1)

        self.max_depth_reached_ = max(self.max_depth_reached_, depth)
        return node

    def _build_best_first(self, X, y, sample_weight, sample_indices):
        """Best-first (leaf-wise) tree building with max_leaf_nodes constraint."""
        self.y_ = y

        root = self._make_leaf(y, sample_weight, sample_indices, depth=0)
        root.is_leaf = True

        # Priority queue: (-gain, counter, node, sample_indices, depth)
        counter = 0
        heap = []

        split = self._find_split(X, y, sample_weight, sample_indices)
        if split.gain > 0 and split.feature >= 0:
            heapq.heappush(heap, (-split.gain, counter, root, sample_indices, 0, split))
            counter += 1

        n_leaves = 1

        while heap and n_leaves < self.max_leaf_nodes:
            neg_gain, _, node, node_indices, depth, split = heapq.heappop(heap)

            if self.max_depth is not None and depth >= self.max_depth:
                continue

            # Find surrogates
            if self.max_surrogates > 0:
                split.surrogates = find_surrogates(
                    X, node_indices,
                    split.feature, split.threshold, split.missing_direction,
                    max_surrogates=self.max_surrogates,
                    all_feature_indices=np.arange(self.n_features_),
                )

            left_idx, right_idx = self._partition(X, node_indices, split)

            if len(left_idx) == 0 or len(right_idx) == 0:
                continue

            # Convert leaf to internal node
            node.is_leaf = False
            node.split = split

            left_child = self._make_leaf(y, sample_weight, left_idx, depth + 1)
            right_child = self._make_leaf(y, sample_weight, right_idx, depth + 1)
            node.left_child = left_child
            node.right_child = right_child

            n_leaves += 1  # Was 1 leaf, now 2 = net +1

            # Try splitting children
            for child, child_idx in [(left_child, left_idx), (right_child, right_idx)]:
                if len(child_idx) >= self.min_samples_split:
                    child_split = self._find_split(X, y, sample_weight, child_idx)
                    if child_split.gain > 0 and child_split.feature >= 0:
                        heapq.heappush(heap, (-child_split.gain, counter, child, child_idx, depth + 1, child_split))
                        counter += 1

        return root

    def _build_symmetric(self, X, y, sample_weight, sample_indices):
        """Symmetric (oblivious) tree building — same split at each depth level."""
        self.y_ = y

        root = self._make_leaf(y, sample_weight, sample_indices, depth=0)

        # Each level: list of (node, sample_indices)
        current_level = [(root, sample_indices)]
        depth = 0

        while True:
            if self.max_depth is not None and depth >= self.max_depth:
                break

            # Pool all samples at this level to find the best global split
            all_indices = np.concatenate([idx for _, idx in current_level])

            if len(all_indices) < self.min_samples_split:
                break

            # Check if any meaningful split exists
            y_unique = np.unique(y[all_indices])
            if len(y_unique) <= 1:
                break

            split = self._find_split(X, y, sample_weight, all_indices)

            if split.gain <= 0 or split.feature < 0:
                break

            next_level = []

            for node, node_indices in current_level:
                if len(node_indices) < self.min_samples_split:
                    continue

                # Use the global split for this node
                node_split = SplitRecord(
                    feature=split.feature,
                    threshold=split.threshold,
                    gain=split.gain,
                    missing_direction=split.missing_direction,
                )

                # Find surrogates per node
                if self.max_surrogates > 0:
                    node_split.surrogates = find_surrogates(
                        X, node_indices,
                        node_split.feature, node_split.threshold,
                        node_split.missing_direction,
                        max_surrogates=self.max_surrogates,
                        all_feature_indices=np.arange(self.n_features_),
                        sample_weight=sample_weight,
                    )

                left_idx, right_idx = self._partition(X, node_indices, node_split)

                if len(left_idx) == 0 or len(right_idx) == 0:
                    continue

                node.is_leaf = False
                node.split = node_split

                left_child = self._make_leaf(y, sample_weight, left_idx, depth + 1)
                right_child = self._make_leaf(y, sample_weight, right_idx, depth + 1)
                node.left_child = left_child
                node.right_child = right_child

                next_level.append((left_child, left_idx))
                next_level.append((right_child, right_idx))

            if not next_level:
                break

            current_level = next_level
            depth += 1

        return root

    def _next_id(self):
        nid = self.node_count_
        self.node_count_ += 1
        return nid
