"""
Tests for knowledge absorbers.
"""

import numpy as np
import pytest
import torch
import torch.nn as nn
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.tree import DecisionTreeClassifier

from model_osmosis.absorbers import PyTorchAbsorber, SklearnAbsorber
from model_osmosis.core.config import OsmosisConfig
from model_osmosis.extractors import PyTorchExtractor, SklearnExtractor
from model_osmosis.knowledge.representation import UniversalKnowledge


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


@pytest.fixture
def config():
    return OsmosisConfig(
        distillation_epochs=20,
        distillation_lr=1e-3,
        verbose=False,
    )


@pytest.fixture
def source_pytorch_model():
    """Trained PyTorch source model."""
    model = nn.Sequential(
        nn.Linear(8, 32),
        nn.ReLU(),
        nn.Linear(32, 16),
        nn.ReLU(),
        nn.Linear(16, 4),
    )
    # "Train" by setting non-random weights
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_normal_(p)
    model.eval()
    return model


@pytest.fixture
def target_pytorch_model():
    """Target model with same architecture as source."""
    model = nn.Sequential(
        nn.Linear(8, 32),
        nn.ReLU(),
        nn.Linear(32, 16),
        nn.ReLU(),
        nn.Linear(16, 4),
    )
    model.eval()
    return model


@pytest.fixture
def smaller_target_pytorch():
    """Smaller target model."""
    model = nn.Sequential(
        nn.Linear(8, 16),
        nn.ReLU(),
        nn.Linear(16, 4),
    )
    model.eval()
    return model


@pytest.fixture
def pytorch_knowledge(source_pytorch_model):
    """Knowledge extracted from the source PyTorch model."""
    ext = PyTorchExtractor()
    sample = np.random.rand(200, 8).astype(np.float32)
    return ext.extract(source_pytorch_model, sample_data=sample)


# ---------------------------------------------------------------------------
# PyTorch Absorber Tests
# ---------------------------------------------------------------------------


class TestPyTorchAbsorber:
    def test_can_handle(self, target_pytorch_model):
        absorber = PyTorchAbsorber()
        assert absorber.can_handle(target_pytorch_model) is True

    def test_absorb_same_architecture(
        self, target_pytorch_model, pytorch_knowledge, config
    ):
        """Absorbing into same-architecture model should work via direct transfer + distillation."""
        absorber = PyTorchAbsorber()
        result = absorber.absorb(target_pytorch_model, pytorch_knowledge, config)

        assert result is target_pytorch_model  # modified in-place

        # The model should now produce predictions
        x = torch.from_numpy(pytorch_knowledge.functional.inputs[:10]).float()
        with torch.no_grad():
            preds = result(x)
        assert preds.shape == (10, 4)

    def test_absorb_smaller_target(
        self, smaller_target_pytorch, pytorch_knowledge, config
    ):
        """Absorbing into a smaller model via distillation should still work."""
        absorber = PyTorchAbsorber()
        result = absorber.absorb(smaller_target_pytorch, pytorch_knowledge, config)

        x = torch.from_numpy(pytorch_knowledge.functional.inputs[:10]).float()
        with torch.no_grad():
            preds = result(x)
        assert preds.shape == (10, 4)

    def test_absorb_cross_framework(self, config):
        """Absorb sklearn knowledge into a PyTorch model."""
        # Train an sklearn model
        rng = np.random.RandomState(42)
        X = rng.rand(300, 5).astype(np.float32)
        y = (X[:, 0] + X[:, 1] > 1.0).astype(int)

        tree = DecisionTreeClassifier(max_depth=4, random_state=42)
        tree.fit(X, y)

        # Extract
        ext = SklearnExtractor()
        knowledge = ext.extract(tree, sample_data=X)

        # Create a PyTorch target
        target = nn.Sequential(
            nn.Linear(5, 16),
            nn.ReLU(),
            nn.Linear(16, 2),
        )

        # Absorb
        absorber = PyTorchAbsorber()
        result = absorber.absorb(target, knowledge, config)

        # Should produce valid outputs
        x_test = torch.from_numpy(X[:10]).float()
        with torch.no_grad():
            preds = result(x_test)
        assert preds.shape == (10, 2)


# ---------------------------------------------------------------------------
# Sklearn Absorber Tests
# ---------------------------------------------------------------------------


class TestSklearnAbsorber:
    def test_can_handle(self):
        absorber = SklearnAbsorber()
        assert absorber.can_handle(LinearRegression()) is True

    def test_absorb_linear_refit(self, config):
        """Test refitting a linear model with functional knowledge."""
        # Create source knowledge from a linear model
        rng = np.random.RandomState(42)
        X = rng.rand(100, 4).astype(np.float32)
        y = X @ np.array([1.0, 2.0, -1.0, 0.5]) + 0.1

        source = LinearRegression()
        source.fit(X, y)

        ext = SklearnExtractor()
        knowledge = ext.extract(source, sample_data=X)

        # Target: fresh linear model
        target = LinearRegression()

        absorber = SklearnAbsorber()
        result = absorber.absorb(target, knowledge, config)

        # Should now make similar predictions
        preds = result.predict(X[:10])
        assert preds.shape == (10,)

        # Predictions should be close to the source
        source_preds = source.predict(X[:10])
        np.testing.assert_allclose(preds, source_preds, rtol=0.1)

    def test_absorb_tree_from_functional(self, config):
        """Test refitting a tree from functional knowledge."""
        rng = np.random.RandomState(42)
        X = rng.rand(200, 5).astype(np.float32)
        y = (X[:, 0] + X[:, 1] > 1.0).astype(int)

        source = DecisionTreeClassifier(max_depth=5, random_state=42)
        source.fit(X, y)

        ext = SklearnExtractor()
        knowledge = ext.extract(source, sample_data=X)

        target = DecisionTreeClassifier(max_depth=5)
        absorber = SklearnAbsorber()
        result = absorber.absorb(target, knowledge, config)

        preds = result.predict(X[:10])
        assert len(preds) == 10
