"""
Generic fallback extractor.

Works with any model that has a `predict()` or `__call__()` method.
Only captures functional knowledge (I/O pairs) — no structural knowledge.
"""

from __future__ import annotations

from typing import Any, Optional, Tuple

import numpy as np

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


class GenericExtractor(BaseExtractor):
    """Fallback extractor for any model with a predict/call interface.

    This is used when no framework-specific extractor can handle the model.
    It treats the model as a black box and only captures I/O behavior.
    """

    def can_handle(self, model: Any) -> bool:
        return (
            hasattr(model, "predict")
            or hasattr(model, "__call__")
        )

    def _extract_structural(self, model: Any) -> StructuralKnowledge:
        # Opaque model — no structural knowledge available
        return StructuralKnowledge(
            model_type="generic",
            architecture_desc=f"Opaque model: {type(model).__name__}",
        )

    def _extract_functional(
        self, model: Any, sample_data: np.ndarray
    ) -> FunctionalKnowledge:
        # Use predict if available, else __call__
        if hasattr(model, "predict"):
            outputs = model.predict(sample_data)
        else:
            outputs = model(sample_data)

        outputs = np.asarray(outputs, dtype=np.float32)
        if outputs.ndim == 1:
            outputs = outputs.reshape(-1, 1)

        # Guess if it's a classifier
        is_classifier = False
        num_classes = 0
        if hasattr(model, "predict_proba"):
            try:
                prob_outputs = model.predict_proba(sample_data)
                outputs = np.asarray(prob_outputs, dtype=np.float32)
                is_classifier = True
                num_classes = outputs.shape[1] if outputs.ndim > 1 else 2
            except Exception:
                pass

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

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

        # Try to infer shapes from common attributes
        if hasattr(model, "n_features_in_"):
            input_shape = (model.n_features_in_,)
        if hasattr(model, "n_outputs_"):
            output_shape = (model.n_outputs_,)

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