"""
sklearn-compatible DirectRS models.

DirectRSClassifier  — piecewise-linear logistic regression on tree geometry
DirectRSRegressor   — piecewise-linear ridge regression on tree geometry
"""

from __future__ import annotations

import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.utils.validation import check_is_fitted

from .geometry import (
    build_embedding,
    compute_G_matrix,
    compute_stretch_matrix,
    get_leaf_indices,
)


class DirectRSClassifier(BaseEstimator, ClassifierMixin):
    """
    DirectRS classification model.

    Extracts path-based geometry from a decision tree and fits a multinomial
    logistic regression on the resulting piecewise-linear embedding.

    Parameters
    ----------
    max_depth : int, default=4
    min_samples_split : int, default=2
    min_samples_leaf : int, default=1
    C : float, default=1.0
        Inverse regularization strength for LogisticRegression.
    random_state : int or None, default=None
    """

    def __init__(
        self,
        max_depth: int = 4,
        min_samples_split: int = 2,
        min_samples_leaf: int = 1,
        C: float = 1.0,
        random_state: int | None = None,
    ):
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.min_samples_leaf = min_samples_leaf
        self.C = C
        self.random_state = random_state

    def fit(self, X: np.ndarray, y: np.ndarray) -> "DirectRSClassifier":
        X = np.asarray(X, dtype=float)
        y = np.asarray(y)

        # 1. Fit the base decision tree
        self.tree_ = DecisionTreeClassifier(
            max_depth=self.max_depth,
            min_samples_split=self.min_samples_split,
            min_samples_leaf=self.min_samples_leaf,
            random_state=self.random_state,
        )
        self.tree_.fit(X, y)
        self.classes_ = self.tree_.classes_
        self.n_features_in_ = X.shape[1]

        # 2. Extract geometry
        self.G_ = compute_G_matrix(self.tree_.tree_, task="classification")
        self.S_prime_ = compute_stretch_matrix(self.G_)
        self.leaf_indices_ = get_leaf_indices(self.tree_.tree_)

        # 3. Build embedding and fit logistic regression
        Phi = build_embedding(X, self.tree_, self.S_prime_, self.leaf_indices_)
        self.lr_ = LogisticRegression(
            C=self.C,
            max_iter=1000,
            solver="lbfgs",
            random_state=self.random_state,
        )
        self.lr_.fit(Phi, y)

        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        check_is_fitted(self)
        X = np.asarray(X, dtype=float)
        Phi = build_embedding(X, self.tree_, self.S_prime_, self.leaf_indices_)
        return self.lr_.predict(Phi)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        check_is_fitted(self)
        X = np.asarray(X, dtype=float)
        Phi = build_embedding(X, self.tree_, self.S_prime_, self.leaf_indices_)
        return self.lr_.predict_proba(Phi)

    def score(self, X: np.ndarray, y: np.ndarray) -> float:
        """Return mean accuracy on (X, y)."""
        return float(np.mean(self.predict(X) == y))

    @property
    def feature_importances_(self) -> dict[str, np.ndarray]:
        """
        Return dict with 'tree' (gain-based) and 'geometry' (diag S') importances.
        Both arrays are normalized to sum to 1.
        """
        check_is_fitted(self)
        tree_imp = self.tree_.feature_importances_

        diag = np.abs(np.diag(self.S_prime_))
        total = diag.sum()
        geom_imp = diag / total if total > 0 else diag

        return {"tree": tree_imp, "geometry": geom_imp}


class DirectRSRegressor(BaseEstimator, RegressorMixin):
    """
    DirectRS regression model.

    Extracts path-based geometry from a decision tree and fits a ridge
    regression on the resulting piecewise-linear embedding.

    Parameters
    ----------
    max_depth : int, default=4
    min_samples_split : int, default=2
    min_samples_leaf : int, default=1
    alpha : float, default=1.0
        Ridge regularization strength.
    random_state : int or None, default=None
    """

    def __init__(
        self,
        max_depth: int = 4,
        min_samples_split: int = 2,
        min_samples_leaf: int = 1,
        alpha: float = 1.0,
        random_state: int | None = None,
    ):
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.min_samples_leaf = min_samples_leaf
        self.alpha = alpha
        self.random_state = random_state

    def fit(self, X: np.ndarray, y: np.ndarray) -> "DirectRSRegressor":
        X = np.asarray(X, dtype=float)
        y = np.asarray(y, dtype=float)

        # 1. Fit the base decision tree
        self.tree_ = DecisionTreeRegressor(
            max_depth=self.max_depth,
            min_samples_split=self.min_samples_split,
            min_samples_leaf=self.min_samples_leaf,
            random_state=self.random_state,
        )
        self.tree_.fit(X, y)
        self.n_features_in_ = X.shape[1]

        # 2. Extract geometry
        self.G_ = compute_G_matrix(self.tree_.tree_, task="regression")
        self.S_prime_ = compute_stretch_matrix(self.G_)
        self.leaf_indices_ = get_leaf_indices(self.tree_.tree_)

        # 3. Build embedding and fit ridge regression
        Phi = build_embedding(X, self.tree_, self.S_prime_, self.leaf_indices_)
        self.ridge_ = Ridge(alpha=self.alpha)
        self.ridge_.fit(Phi, y)

        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        check_is_fitted(self)
        X = np.asarray(X, dtype=float)
        Phi = build_embedding(X, self.tree_, self.S_prime_, self.leaf_indices_)
        return self.ridge_.predict(Phi)

    def score(self, X: np.ndarray, y: np.ndarray) -> float:
        """Return R² on (X, y)."""
        from sklearn.metrics import r2_score
        return float(r2_score(y, self.predict(X)))

    @property
    def feature_importances_(self) -> dict[str, np.ndarray]:
        """
        Return dict with 'tree' (gain-based) and 'geometry' (diag S') importances.
        Both arrays are normalized to sum to 1.
        """
        check_is_fitted(self)
        tree_imp = self.tree_.feature_importances_

        diag = np.abs(np.diag(self.S_prime_))
        total = diag.sum()
        geom_imp = diag / total if total > 0 else diag

        return {"tree": tree_imp, "geometry": geom_imp}
