"""
Base extractor — abstract interface for extracting knowledge from models.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Optional

import numpy as np

from model_osmosis.knowledge.representation import (
    FunctionalKnowledge,
    MetaKnowledge,
    StructuralKnowledge,
    UniversalKnowledge,
)


class BaseExtractor(ABC):
    """Abstract base class for model knowledge extractors.

    Subclasses implement framework-specific extraction logic for PyTorch,
    scikit-learn, TensorFlow, etc.
    """

    @abstractmethod
    def can_handle(self, model: Any) -> bool:
        """Return True if this extractor knows how to handle the given model.

        Args:
            model: The trained model to inspect.
        """

    def extract(
        self,
        model: Any,
        sample_data: Optional[np.ndarray] = None,
        num_samples: int = 1000,
    ) -> UniversalKnowledge:
        """Extract full knowledge from a trained model.

        Args:
            model: The trained source model.
            sample_data: Optional sample input data for functional extraction.
                If None, synthetic data will be generated from the model's
                expected input shape.
            num_samples: Number of synthetic samples to generate if
                sample_data is not provided.

        Returns:
            A UniversalKnowledge object containing all extracted knowledge.
        """
        structural = self._extract_structural(model)
        meta = self._extract_meta(model)

        if sample_data is None:
            sample_data = self._generate_sample_data(model, meta, num_samples)

        functional = self._extract_functional(model, sample_data)
        return UniversalKnowledge(
            structural=structural,
            functional=functional,
            meta=meta,
        )

    @abstractmethod
    def _extract_structural(self, model: Any) -> StructuralKnowledge:
        """Extract architecture-level knowledge (weights, structure)."""

    @abstractmethod
    def _extract_functional(
        self, model: Any, sample_data: np.ndarray
    ) -> FunctionalKnowledge:
        """Extract behavioral knowledge (I/O pairs, activations)."""

    @abstractmethod
    def _extract_meta(self, model: Any) -> MetaKnowledge:
        """Extract metadata about the model."""

    def _generate_sample_data(
        self, model: Any, meta: MetaKnowledge, num_samples: int
    ) -> np.ndarray:
        """Generate synthetic sample data for functional extraction.

        Uses uniform random data in [0, 1] with the model's expected input shape.
        Subclasses can override for smarter sampling.
        """
        if meta.input_shape and len(meta.input_shape) > 0:
            shape = (num_samples, *meta.input_shape)
        else:
            # Fallback: try to infer from model or use a small default
            shape = (num_samples, 10)
        return np.random.rand(*shape).astype(np.float32)
