"""
Tests for knowledge extractors.
"""

import numpy as np
import pytest
import torch
import torch.nn as nn
from sklearn.tree import DecisionTreeClassifier
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestClassifier

from model_osmosis.extractors import PyTorchExtractor, SklearnExtractor, GenericExtractor
from model_osmosis.knowledge.representation import UniversalKnowledge


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def simple_pytorch_model():
    """A small PyTorch model for testing."""
    model = nn.Sequential(
        nn.Linear(10, 32),
        nn.ReLU(),
        nn.Linear(32, 16),
        nn.ReLU(),
        nn.Linear(16, 3),
    )
    # Simulate training by randomizing weights
    for p in model.parameters():
        nn.init.normal_(p)
    model.eval()
    return model


@pytest.fixture
def trained_tree():
    """A trained decision tree classifier."""
    rng = np.random.RandomState(42)
    X = rng.rand(200, 5)
    y = (X[:, 0] + X[:, 1] > 1.0).astype(int)
    clf = DecisionTreeClassifier(max_depth=5, random_state=42)
    clf.fit(X, y)
    return clf


@pytest.fixture
def trained_linear():
    """A trained linear regression model."""
    rng = np.random.RandomState(42)
    X = rng.rand(100, 4)
    y = X @ np.array([1.0, 2.0, -1.0, 0.5]) + 0.1
    model = LinearRegression()
    model.fit(X, y)
    return model


@pytest.fixture
def trained_forest():
    """A trained random forest."""
    rng = np.random.RandomState(42)
    X = rng.rand(200, 5)
    y = (X[:, 0] + X[:, 1] > 1.0).astype(int)
    clf = RandomForestClassifier(n_estimators=10, max_depth=4, random_state=42)
    clf.fit(X, y)
    return clf


@pytest.fixture
def sample_data_10d():
    """Sample data with 10 features."""
    return np.random.rand(50, 10).astype(np.float32)


@pytest.fixture
def sample_data_5d():
    """Sample data with 5 features."""
    return np.random.rand(50, 5).astype(np.float32)


@pytest.fixture
def sample_data_4d():
    """Sample data with 4 features."""
    return np.random.rand(50, 4).astype(np.float32)


# ---------------------------------------------------------------------------
# PyTorch Extractor Tests
# ---------------------------------------------------------------------------


class TestPyTorchExtractor:
    def test_can_handle_pytorch(self, simple_pytorch_model):
        ext = PyTorchExtractor()
        assert ext.can_handle(simple_pytorch_model) is True

    def test_can_handle_non_pytorch(self, trained_tree):
        ext = PyTorchExtractor()
        assert ext.can_handle(trained_tree) is False

    def test_extract_produces_valid_knowledge(
        self, simple_pytorch_model, sample_data_10d
    ):
        ext = PyTorchExtractor()
        knowledge = ext.extract(simple_pytorch_model, sample_data=sample_data_10d)

        assert isinstance(knowledge, UniversalKnowledge)
        assert knowledge.structural.model_type == "pytorch"
        assert knowledge.structural.total_params > 0
        assert len(knowledge.structural.parameters) > 0
        assert len(knowledge.structural.layer_info) > 0

    def test_extract_functional_has_matching_shapes(
        self, simple_pytorch_model, sample_data_10d
    ):
        ext = PyTorchExtractor()
        knowledge = ext.extract(simple_pytorch_model, sample_data=sample_data_10d)

        assert knowledge.functional.inputs.shape[0] == 50
        assert knowledge.functional.outputs.shape[0] == 50
        assert knowledge.functional.outputs.shape[1] == 3  # 3 output classes

    def test_extract_captures_activations(
        self, simple_pytorch_model, sample_data_10d
    ):
        ext = PyTorchExtractor()
        knowledge = ext.extract(simple_pytorch_model, sample_data=sample_data_10d)

        # Should have activations for each named module
        assert len(knowledge.functional.intermediate_activations) > 0

    def test_meta_captures_shapes(self, simple_pytorch_model, sample_data_10d):
        ext = PyTorchExtractor()
        knowledge = ext.extract(simple_pytorch_model, sample_data=sample_data_10d)

        assert knowledge.meta.source_framework == "pytorch"
        assert knowledge.meta.input_shape == (10,)
        assert knowledge.meta.output_shape == (3,)


# ---------------------------------------------------------------------------
# Sklearn Extractor Tests
# ---------------------------------------------------------------------------


class TestSklearnExtractor:
    def test_can_handle_sklearn(self, trained_tree):
        ext = SklearnExtractor()
        assert ext.can_handle(trained_tree) is True

    def test_can_handle_non_sklearn(self, simple_pytorch_model):
        ext = SklearnExtractor()
        assert ext.can_handle(simple_pytorch_model) is False

    def test_extract_tree(self, trained_tree, sample_data_5d):
        ext = SklearnExtractor()
        knowledge = ext.extract(trained_tree, sample_data=sample_data_5d)

        assert isinstance(knowledge, UniversalKnowledge)
        assert knowledge.structural.model_type == "sklearn_tree"
        assert knowledge.structural.total_params > 0
        assert "threshold" in knowledge.structural.parameters
        assert knowledge.meta.task_type == "classification"

    def test_extract_linear(self, trained_linear, sample_data_4d):
        ext = SklearnExtractor()
        knowledge = ext.extract(trained_linear, sample_data=sample_data_4d)

        assert knowledge.structural.model_type == "sklearn_linear"
        assert "coef" in knowledge.structural.parameters
        assert "intercept" in knowledge.structural.parameters

    def test_extract_forest(self, trained_forest, sample_data_5d):
        ext = SklearnExtractor()
        knowledge = ext.extract(trained_forest, sample_data=sample_data_5d)

        assert knowledge.structural.model_type == "sklearn_ensemble"
        assert knowledge.structural.total_params > 0
        assert knowledge.meta.feature_importances is not None

    def test_functional_outputs_shape(self, trained_tree, sample_data_5d):
        ext = SklearnExtractor()
        knowledge = ext.extract(trained_tree, sample_data=sample_data_5d)

        assert knowledge.functional.inputs.shape == (50, 5)
        assert knowledge.functional.outputs.shape[0] == 50
        assert knowledge.functional.is_classifier is True


# ---------------------------------------------------------------------------
# Generic Extractor Tests
# ---------------------------------------------------------------------------


class TestGenericExtractor:
    def test_can_handle_with_predict(self):
        class FakeModel:
            def predict(self, X):
                return np.zeros(len(X))

        ext = GenericExtractor()
        assert ext.can_handle(FakeModel()) is True

    def test_can_handle_with_call(self):
        class CallableModel:
            def __call__(self, X):
                return np.ones(len(X))

        ext = GenericExtractor()
        assert ext.can_handle(CallableModel()) is True

    def test_extract_generic(self):
        class SimpleModel:
            def predict(self, X):
                return np.sum(X, axis=1)

        ext = GenericExtractor()
        sample = np.random.rand(30, 7).astype(np.float32)
        knowledge = ext.extract(SimpleModel(), sample_data=sample)

        assert isinstance(knowledge, UniversalKnowledge)
        assert knowledge.structural.model_type == "generic"
        assert knowledge.functional.inputs.shape == (30, 7)
        assert knowledge.functional.outputs.shape[0] == 30
