"""
Integration tests — full pipeline end-to-end.
"""

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 model_osmosis import OsmosisPipeline, OsmosisConfig
from model_osmosis.core.config import GrowthStrategyType


class TestPipelineNeuralToNeural:
    """Test full pipeline: PyTorch → PyTorch transfer."""

    def test_full_transfer_same_arch(self):
        """Source and target have same architecture."""
        config = OsmosisConfig(
            distillation_epochs=30,
            verbose=False,
        )
        pipeline = OsmosisPipeline(config)

        # Source
        source = nn.Sequential(
            nn.Linear(6, 24), nn.ReLU(),
            nn.Linear(24, 12), nn.ReLU(),
            nn.Linear(12, 3),
        )
        for p in source.parameters():
            if p.dim() > 1:
                nn.init.xavier_normal_(p)
        source.eval()

        # Target (same architecture, different weights)
        target = nn.Sequential(
            nn.Linear(6, 24), nn.ReLU(),
            nn.Linear(24, 12), nn.ReLU(),
            nn.Linear(12, 3),
        )

        sample = np.random.rand(200, 6).astype(np.float32)
        result = pipeline.absorb(source, target, sample_data=sample, auto_grow=False)

        # Verify: target should now produce similar outputs to source
        x = torch.from_numpy(sample[:20]).float()
        source.eval()
        result.eval()
        with torch.no_grad():
            src_out = source(x).numpy()
            tgt_out = result(x).numpy()

        # Correlation should be high after distillation
        for i in range(3):
            corr = np.corrcoef(src_out[:, i], tgt_out[:, i])[0, 1]
            assert corr > 0.8, f"Output {i} correlation too low: {corr:.3f}"

    def test_full_transfer_with_growth(self):
        """Smaller target grows to absorb larger source."""
        config = OsmosisConfig(
            distillation_epochs=30,
            growth_strategy=GrowthStrategyType.WIDER,
            verbose=False,
        )
        pipeline = OsmosisPipeline(config)

        # Large source
        source = nn.Sequential(
            nn.Linear(8, 64), nn.ReLU(),
            nn.Linear(64, 32), nn.ReLU(),
            nn.Linear(32, 4),
        )
        for p in source.parameters():
            if p.dim() > 1:
                nn.init.xavier_normal_(p)

        # Small target
        target = nn.Sequential(
            nn.Linear(8, 8), nn.ReLU(),
            nn.Linear(8, 4),
        )

        sample = np.random.rand(200, 8).astype(np.float32)

        old_params = sum(p.numel() for p in target.parameters())
        result = pipeline.absorb(source, target, sample_data=sample, auto_grow=True)
        new_params = sum(p.numel() for p in result.parameters())

        # Target should have grown
        assert new_params > old_params

        # Should still produce valid outputs
        x = torch.from_numpy(sample[:5]).float()
        with torch.no_grad():
            out = result(x)
        assert out.shape == (5, 4)


class TestPipelineSklearnToNeural:
    """Test full pipeline: sklearn → PyTorch transfer."""

    def test_tree_to_neural(self):
        """Transfer decision tree knowledge into a neural network."""
        config = OsmosisConfig(
            distillation_epochs=50,
            verbose=False,
        )
        pipeline = OsmosisPipeline(config)

        # Train a decision tree
        rng = np.random.RandomState(42)
        X = rng.rand(500, 5).astype(np.float32)
        y = (X[:, 0] * 2 + X[:, 1] > 1.5).astype(int)

        tree = DecisionTreeClassifier(max_depth=5, random_state=42)
        tree.fit(X, y)
        tree_accuracy = (tree.predict(X) == y).mean()

        # Neural target
        target = nn.Sequential(
            nn.Linear(5, 32), nn.ReLU(),
            nn.Linear(32, 16), nn.ReLU(),
            nn.Linear(16, 2),
        )

        result = pipeline.absorb(tree, target, sample_data=X, auto_grow=False)

        # Verify the neural net learned the tree's behavior
        x_test = torch.from_numpy(X[:100]).float()
        with torch.no_grad():
            logits = result(x_test)
            nn_preds = torch.argmax(logits, dim=1).numpy()

        tree_preds = tree.predict(X[:100])
        agreement = (nn_preds == tree_preds).mean()
        assert agreement > 0.7, f"Agreement too low: {agreement:.3f}"


class TestPipelineSklearnToSklearn:
    """Test full pipeline: sklearn → sklearn transfer."""

    def test_linear_to_linear(self):
        """Transfer one linear model's knowledge to another."""
        config = OsmosisConfig(verbose=False)
        pipeline = OsmosisPipeline(config)

        rng = np.random.RandomState(42)
        X = rng.rand(200, 4).astype(np.float32)
        y = X @ np.array([1.5, -2.0, 0.5, 3.0]) + 0.3

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

        target = LinearRegression()

        result = pipeline.absorb(source, target, sample_data=X, auto_grow=False)

        # Verify predictions match
        src_preds = source.predict(X[:50])
        tgt_preds = result.predict(X[:50])
        np.testing.assert_allclose(tgt_preds, src_preds, rtol=0.1)


class TestPipelineStepByStep:
    """Test using the pipeline step-by-step (extract → grow → inject)."""

    def test_step_by_step_api(self):
        config = OsmosisConfig(
            distillation_epochs=20,
            verbose=False,
        )
        pipeline = OsmosisPipeline(config)

        source = nn.Sequential(
            nn.Linear(6, 32), nn.ReLU(),
            nn.Linear(32, 3),
        )
        for p in source.parameters():
            if p.dim() > 1:
                nn.init.xavier_normal_(p)

        target = nn.Sequential(
            nn.Linear(6, 8), nn.ReLU(),
            nn.Linear(8, 3),
        )

        sample = np.random.rand(100, 6).astype(np.float32)

        # Step 1: Extract
        knowledge = pipeline.extract(source, sample_data=sample)
        assert knowledge.source_type == "pytorch"
        assert knowledge.param_count > 0

        # Step 2: Grow
        grown = pipeline.grow(target, knowledge)
        grown_params = sum(p.numel() for p in grown.parameters())
        original_params = sum(p.numel() for p in target.parameters())
        assert grown_params >= original_params

        # Step 3: Inject
        final = pipeline.inject(grown, knowledge)

        # Verify
        x = torch.from_numpy(sample[:5]).float()
        with torch.no_grad():
            out = final(x)
        assert out.shape == (5, 3)
