"""
Scikit-learn absorber — inject knowledge into sklearn estimators.

Strategy:
- For linear models: direct coefficient transfer if compatible, else refit.
- For all models: refit using functional knowledge (I/O pairs).
"""

from __future__ import annotations

from typing import Any

import numpy as np

from model_osmosis.absorbers.base import BaseAbsorber
from model_osmosis.core.config import OsmosisConfig
from model_osmosis.core.exceptions import AbsorptionError
from model_osmosis.knowledge.representation import UniversalKnowledge


class SklearnAbsorber(BaseAbsorber):
    """Absorb knowledge into a scikit-learn estimator."""

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

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

    def absorb(
        self,
        target_model: Any,
        knowledge: UniversalKnowledge,
        config: OsmosisConfig,
    ) -> Any:
        """Inject knowledge into a sklearn estimator.

        For linear models with compatible source: attempts direct coefficient
        transfer. For all models: refits on the source's functional knowledge.
        """
        # Try direct linear transfer first
        if _is_linear(target_model) and knowledge.source_type == "sklearn_linear":
            transferred = self._direct_linear_transfer(target_model, knowledge)
            if transferred:
                if config.verbose:
                    print("[Absorber] Direct linear coefficient transfer successful")
                return target_model

        # Universal fallback: refit on functional knowledge
        self._refit(target_model, knowledge, config)
        return target_model

    def _direct_linear_transfer(
        self,
        target_model: Any,
        knowledge: UniversalKnowledge,
    ) -> bool:
        """Copy coefficients directly into a linear model.

        Returns True if transfer was successful.
        """
        params = knowledge.structural.parameters

        if "coef" in params and hasattr(target_model, "coef_"):
            source_coef = params["coef"]
            try:
                # Ensure the target can accept these coefficients
                target_model.coef_ = source_coef.copy()
                if "intercept" in params and hasattr(target_model, "intercept_"):
                    target_model.intercept_ = params["intercept"].copy()
                return True
            except Exception:
                return False
        return False

    def _refit(
        self,
        target_model: Any,
        knowledge: UniversalKnowledge,
        config: OsmosisConfig,
    ) -> None:
        """Refit the target model using the source's I/O pairs.

        For classifiers, converts probability outputs to hard labels for fitting.
        For regressors, uses raw outputs.
        """
        X = knowledge.functional.inputs
        y = knowledge.functional.outputs

        if knowledge.functional.is_classifier:
            # Convert probabilities to class labels for sklearn classifiers
            if y.ndim > 1 and y.shape[1] > 1:
                y_labels = np.argmax(y, axis=1)
            else:
                y_labels = (y.ravel() > 0.5).astype(int)

            # Ensure the target has classes_ attribute setup
            if not hasattr(target_model, "classes_"):
                unique_labels = np.unique(y_labels)
                target_model.classes_ = unique_labels

            target_model.fit(X, y_labels)
        else:
            if y.ndim > 1 and y.shape[1] == 1:
                y = y.ravel()
            target_model.fit(X, y)

        if config.verbose:
            print(
                f"[Absorber] Refit {type(target_model).__name__} on "
                f"{len(X)} functional samples"
            )


def _is_linear(model: Any) -> bool:
    """Check if a model is a linear sklearn model."""
    class_name = type(model).__name__.lower()
    linear_names = {
        "linearregression", "logisticregression", "ridge", "lasso",
        "elasticnet", "sgdclassifier", "sgdregressor",
    }
    return any(name in class_name for name in linear_names) or hasattr(model, "coef_")
