"""
Scikit-learn model extractor.

Handles decision trees, random forests, gradient boosted trees, SVMs,
linear models, and any other sklearn estimator.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from model_osmosis.core.exceptions import ExtractionError
from model_osmosis.extractors.base import BaseExtractor
from model_osmosis.knowledge.representation import (
    FunctionalKnowledge,
    MetaKnowledge,
    StructuralKnowledge,
)


class SklearnExtractor(BaseExtractor):
    """Extract knowledge from scikit-learn estimators."""

    def can_handle(self, model: Any) -> bool:
        try:
            from sklearn.base import BaseEstimator

            return isinstance(model, BaseEstimator)
        except ImportError:
            return False

    def _extract_structural(self, model: Any) -> StructuralKnowledge:
        model_category = _categorize_sklearn_model(model)
        layer_info: List[Dict[str, Any]] = []
        parameters: Dict[str, np.ndarray] = {}
        total_params = 0

        if model_category == "tree":
            layer_info, parameters, total_params = _extract_tree_structure(model)
        elif model_category == "ensemble":
            layer_info, parameters, total_params = _extract_ensemble_structure(model)
        elif model_category == "linear":
            layer_info, parameters, total_params = _extract_linear_structure(model)
        elif model_category == "svm":
            layer_info, parameters, total_params = _extract_svm_structure(model)
        else:
            # Opaque model — no structural knowledge
            pass

        return StructuralKnowledge(
            model_type=f"sklearn_{model_category}",
            architecture_desc=f"sklearn {type(model).__name__}",
            layer_info=layer_info,
            parameters=parameters,
            total_params=total_params,
        )

    def _extract_functional(
        self, model: Any, sample_data: np.ndarray
    ) -> FunctionalKnowledge:
        is_classifier = hasattr(model, "predict_proba") or hasattr(model, "classes_")
        num_classes = 0

        if is_classifier and hasattr(model, "predict_proba"):
            outputs = model.predict_proba(sample_data)
            num_classes = outputs.shape[1] if outputs.ndim > 1 else 2
        else:
            outputs = model.predict(sample_data)
            if outputs.ndim == 1:
                outputs = outputs.reshape(-1, 1)

        return FunctionalKnowledge(
            inputs=sample_data.astype(np.float32),
            outputs=outputs.astype(np.float32),
            is_classifier=is_classifier,
            num_classes=num_classes,
        )

    def _extract_meta(self, model: Any) -> MetaKnowledge:
        input_shape: Tuple[int, ...] = ()
        output_shape: Tuple[int, ...] = ()
        feature_importances: Optional[np.ndarray] = None

        # Infer input shape
        if hasattr(model, "n_features_in_"):
            input_shape = (model.n_features_in_,)

        # Infer output shape
        if hasattr(model, "n_outputs_"):
            output_shape = (model.n_outputs_,)
        elif hasattr(model, "classes_"):
            output_shape = (len(model.classes_),)

        # Feature importances
        if hasattr(model, "feature_importances_"):
            feature_importances = model.feature_importances_.copy()
        elif hasattr(model, "coef_"):
            feature_importances = np.abs(model.coef_).mean(axis=0) if model.coef_.ndim > 1 else np.abs(model.coef_)

        # Task type
        is_classifier = hasattr(model, "classes_")
        task_type = "classification" if is_classifier else "regression"

        return MetaKnowledge(
            source_framework="sklearn",
            source_class_name=type(model).__name__,
            input_shape=input_shape,
            output_shape=output_shape,
            feature_importances=feature_importances,
            task_type=task_type,
        )

    def _generate_sample_data(
        self, model: Any, meta: MetaKnowledge, num_samples: int
    ) -> np.ndarray:
        if meta.input_shape and len(meta.input_shape) > 0:
            shape = (num_samples, *meta.input_shape)
        else:
            shape = (num_samples, 10)
        return np.random.rand(*shape).astype(np.float32)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _categorize_sklearn_model(model: Any) -> str:
    """Determine the category of an sklearn model."""
    class_name = type(model).__name__.lower()

    tree_names = {"decisiontree", "tree"}
    ensemble_names = {"randomforest", "gradientboosting", "adaboost", "bagging", "extratrees"}
    linear_names = {
        "linearregression", "logisticregression", "ridge", "lasso",
        "elasticnet", "sgd", "perceptron", "passiveaggressive",
    }
    svm_names = {"svc", "svr", "nusvc", "nusvr", "linearsvc", "linearsvr"}

    for name in tree_names:
        if name in class_name:
            return "tree"
    for name in ensemble_names:
        if name in class_name:
            return "ensemble"
    for name in svm_names:
        if name in class_name:
            return "svm"
    for name in linear_names:
        if name in class_name:
            return "linear"

    # Fallback: check for common attributes
    if hasattr(model, "tree_"):
        return "tree"
    if hasattr(model, "estimators_"):
        return "ensemble"
    if hasattr(model, "coef_"):
        return "linear"
    if hasattr(model, "support_vectors_"):
        return "svm"

    return "generic"


def _extract_tree_structure(model: Any):
    """Extract structure from a single decision tree."""
    tree = model.tree_
    parameters = {
        "feature": tree.feature.copy(),
        "threshold": tree.threshold.copy(),
        "value": tree.value.copy(),
        "children_left": tree.children_left.copy(),
        "children_right": tree.children_right.copy(),
    }
    total_params = sum(p.size for p in parameters.values())

    layer_info = [{
        "type": "DecisionTree",
        "n_nodes": tree.node_count,
        "max_depth": model.get_depth() if hasattr(model, "get_depth") else tree.max_depth,
        "n_features": tree.n_features,
        "n_outputs": tree.n_outputs,
    }]

    return layer_info, parameters, total_params


def _extract_ensemble_structure(model: Any):
    """Extract structure from an ensemble (forest, boosting, etc.)."""
    layer_info = []
    parameters = {}
    total_params = 0

    for i, estimator in enumerate(model.estimators_):
        # Handle nested arrays (e.g., RandomForest stores in flat list,
        # GradientBoosting stores arrays of arrays)
        if hasattr(estimator, "__iter__") and not hasattr(estimator, "tree_"):
            for j, sub_est in enumerate(estimator):
                if hasattr(sub_est, "tree_"):
                    tree = sub_est.tree_
                    prefix = f"estimator_{i}_{j}"
                    parameters[f"{prefix}_feature"] = tree.feature.copy()
                    parameters[f"{prefix}_threshold"] = tree.threshold.copy()
                    parameters[f"{prefix}_value"] = tree.value.copy()
                    total_params += tree.feature.size + tree.threshold.size + tree.value.size
                    layer_info.append({
                        "type": "DecisionTree",
                        "name": prefix,
                        "n_nodes": tree.node_count,
                    })
        elif hasattr(estimator, "tree_"):
            tree = estimator.tree_
            prefix = f"estimator_{i}"
            parameters[f"{prefix}_feature"] = tree.feature.copy()
            parameters[f"{prefix}_threshold"] = tree.threshold.copy()
            parameters[f"{prefix}_value"] = tree.value.copy()
            total_params += tree.feature.size + tree.threshold.size + tree.value.size
            layer_info.append({
                "type": "DecisionTree",
                "name": prefix,
                "n_nodes": tree.node_count,
            })

    return layer_info, parameters, total_params


def _extract_linear_structure(model: Any):
    """Extract structure from a linear model."""
    parameters = {}
    total_params = 0

    if hasattr(model, "coef_"):
        coef = np.atleast_2d(model.coef_)
        parameters["coef"] = coef.copy()
        total_params += coef.size

    if hasattr(model, "intercept_"):
        intercept = np.atleast_1d(model.intercept_)
        parameters["intercept"] = intercept.copy()
        total_params += intercept.size

    layer_info = [{
        "type": "Linear",
        "in_features": parameters.get("coef", np.array([])).shape[-1] if "coef" in parameters else 0,
        "out_features": parameters.get("coef", np.array([])).shape[0] if "coef" in parameters else 0,
    }]

    return layer_info, parameters, total_params


def _extract_svm_structure(model: Any):
    """Extract structure from an SVM."""
    parameters = {}
    total_params = 0

    if hasattr(model, "support_vectors_"):
        sv = model.support_vectors_
        parameters["support_vectors"] = sv.copy()
        total_params += sv.size

    if hasattr(model, "dual_coef_"):
        dc = model.dual_coef_
        parameters["dual_coef"] = dc.copy()
        total_params += dc.size

    if hasattr(model, "intercept_"):
        intercept = np.atleast_1d(model.intercept_)
        parameters["intercept"] = intercept.copy()
        total_params += intercept.size

    layer_info = [{
        "type": "SVM",
        "kernel": getattr(model, "kernel", "unknown"),
        "n_support_vectors": len(model.support_vectors_) if hasattr(model, "support_vectors_") else 0,
    }]

    return layer_info, parameters, total_params
