"""Public API: SurrogateDecisionTreeClassifier and SurrogateDecisionTreeRegressor."""

import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
from sklearn.utils.multiclass import check_classification_targets
from sklearn.utils.validation import check_is_fitted, validate_data

from ._tree import TreeBuilder
from ._validation import validate_sample_weight


class SurrogateDecisionTreeClassifier(ClassifierMixin, BaseEstimator):
    """Decision tree classifier with surrogate splits for missing data.

    Handles NaN values natively using surrogate splits (CART-style) and
    learned missing directions (XGBoost-style). Supports three growth modes:
    depth-first, best-first (leaf-wise), and symmetric (oblivious).
    """

    def __init__(
        self,
        criterion="gini",
        max_depth=None,
        min_samples_split=2,
        min_samples_leaf=1,
        min_weight_fraction_leaf=0.0,
        max_features=None,
        max_leaf_nodes=None,
        min_impurity_decrease=0.0,
        max_surrogates=5,
        use_histogram=True,
        max_bins=255,
        symmetric_tree=False,
        random_state=None,
    ):
        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

    def __sklearn_tags__(self):
        tags = super().__sklearn_tags__()
        tags.input_tags.allow_nan = True
        return tags

    def fit(self, X, y, sample_weight=None):
        X, y = validate_data(self, X, y, accept_sparse=False,
                             ensure_all_finite="allow-nan", dtype=np.float64,
                             multi_output=False)
        sample_weight = validate_sample_weight(sample_weight, X.shape[0])
        check_classification_targets(y)

        # Remove zero-weight samples for consistency with duplicated samples
        nonzero = sample_weight > 0
        if not nonzero.all():
            X = X[nonzero]
            y = y[nonzero]
            sample_weight = sample_weight[nonzero]

        self.classes_ = np.unique(y)
        self.n_classes_ = len(self.classes_)

        # Remap labels to 0..n_classes-1
        label_map = {c: i for i, c in enumerate(self.classes_)}
        y_mapped = np.array([label_map[v] for v in y], dtype=np.intp)

        builder = TreeBuilder(
            is_classifier=True,
            n_classes=self.n_classes_,
            criterion=self.criterion,
            max_depth=self.max_depth,
            min_samples_split=self.min_samples_split,
            min_samples_leaf=self.min_samples_leaf,
            min_weight_fraction_leaf=self.min_weight_fraction_leaf,
            max_features=self.max_features,
            max_leaf_nodes=self.max_leaf_nodes,
            min_impurity_decrease=self.min_impurity_decrease,
            max_surrogates=self.max_surrogates,
            use_histogram=self.use_histogram,
            max_bins=self.max_bins,
            symmetric_tree=self.symmetric_tree,
            random_state=self.random_state,
        )

        self.tree_ = builder.build(X, y_mapped, sample_weight)
        self.max_depth_ = builder.max_depth_reached_
        self.node_count_ = builder.node_count_

        return self

    def predict(self, X):
        check_is_fitted(self)
        X = validate_data(self, X, accept_sparse=False,
                          ensure_all_finite="allow-nan", dtype=np.float64,
                          reset=False)
        predictions = np.array([self.tree_.predict_single(x) for x in X])
        class_indices = np.argmax(predictions, axis=1)
        return self.classes_[class_indices]

    def predict_proba(self, X):
        check_is_fitted(self)
        X = validate_data(self, X, accept_sparse=False,
                          ensure_all_finite="allow-nan", dtype=np.float64,
                          reset=False)
        raw = np.array([self.tree_.predict_single(x) for x in X])
        # Normalize to probabilities
        row_sums = raw.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1.0
        return raw / row_sums


class SurrogateDecisionTreeRegressor(RegressorMixin, BaseEstimator):
    """Decision tree regressor with surrogate splits for missing data."""

    def __init__(
        self,
        criterion="squared_error",
        max_depth=None,
        min_samples_split=2,
        min_samples_leaf=1,
        min_weight_fraction_leaf=0.0,
        max_features=None,
        max_leaf_nodes=None,
        min_impurity_decrease=0.0,
        max_surrogates=5,
        use_histogram=True,
        max_bins=255,
        symmetric_tree=False,
        random_state=None,
    ):
        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

    def __sklearn_tags__(self):
        tags = super().__sklearn_tags__()
        tags.input_tags.allow_nan = True
        return tags

    def fit(self, X, y, sample_weight=None):
        X, y = validate_data(self, X, y, accept_sparse=False,
                             ensure_all_finite="allow-nan", dtype=np.float64,
                             multi_output=False)
        sample_weight = validate_sample_weight(sample_weight, X.shape[0])
        y = y.astype(np.float64)

        # Remove zero-weight samples for consistency with duplicated samples
        nonzero = sample_weight > 0
        if not nonzero.all():
            X = X[nonzero]
            y = y[nonzero]
            sample_weight = sample_weight[nonzero]

        builder = TreeBuilder(
            is_classifier=False,
            criterion=self.criterion,
            max_depth=self.max_depth,
            min_samples_split=self.min_samples_split,
            min_samples_leaf=self.min_samples_leaf,
            min_weight_fraction_leaf=self.min_weight_fraction_leaf,
            max_features=self.max_features,
            max_leaf_nodes=self.max_leaf_nodes,
            min_impurity_decrease=self.min_impurity_decrease,
            max_surrogates=self.max_surrogates,
            use_histogram=self.use_histogram,
            max_bins=self.max_bins,
            symmetric_tree=self.symmetric_tree,
            random_state=self.random_state,
        )

        self.tree_ = builder.build(X, y, sample_weight)
        self.max_depth_ = builder.max_depth_reached_
        self.node_count_ = builder.node_count_

        return self

    def predict(self, X):
        check_is_fitted(self)
        X = validate_data(self, X, accept_sparse=False,
                          ensure_all_finite="allow-nan", dtype=np.float64,
                          reset=False)
        predictions = np.array([self.tree_.predict_single(x) for x in X])
        # value is [mean, weight_total] for regression
        return predictions[:, 0]
