"""
Tests for ModalityBridge — cross-modal model merging.
"""

import numpy as np
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F

from model_osmosis.merging.modality_bridge import (
    ModalityBridge,
    BridgedModel,
    ModalityEncoder,
)


def _make_model(in_dim, hidden, out_dim, seed=0):
    torch.manual_seed(seed)
    model = nn.Sequential(
        nn.Linear(in_dim, hidden),
        nn.ReLU(),
        nn.Linear(hidden, out_dim),
    )
    model.eval()
    return model


class TestModalityBridge:
    def test_bridge_two_modalities(self):
        """Bridge two models with different input/output shapes."""
        # "Text" model: 10-dim input → 8-dim output
        text_model = _make_model(10, 32, 8, seed=0)
        # "Image" model: 20-dim input → 12-dim output
        image_model = _make_model(20, 64, 12, seed=1)

        # Paired data: 50 samples, each pair describes the same concept
        text_data = np.random.rand(50, 10).astype(np.float32)
        image_data = np.random.rand(50, 20).astype(np.float32)

        bridge = ModalityBridge(common_dim=64, verbose=False)
        result = bridge.merge(
            models=[text_model, image_model],
            paired_data=[text_data, image_data],
            modality_names=["text", "image"],
            epochs=30,
        )

        assert isinstance(result, BridgedModel)

        # Can encode text
        x_text = torch.rand(5, 8)
        common = result.encode(x_text, modality=0)
        assert common.shape == (5, 64)

        # Can encode image
        x_image = torch.rand(5, 12)
        common = result.encode(x_image, modality=1)
        assert common.shape == (5, 64)

    def test_bridge_three_modalities(self):
        """Bridge three different modalities."""
        model_a = _make_model(10, 32, 8, seed=0)
        model_b = _make_model(15, 48, 6, seed=1)
        model_c = _make_model(20, 64, 10, seed=2)

        n = 40
        data_a = np.random.rand(n, 10).astype(np.float32)
        data_b = np.random.rand(n, 15).astype(np.float32)
        data_c = np.random.rand(n, 20).astype(np.float32)

        bridge = ModalityBridge(common_dim=32, verbose=False)
        result = bridge.merge(
            models=[model_a, model_b, model_c],
            paired_data=[data_a, data_b, data_c],
            modality_names=["audio", "text", "image"],
            epochs=20,
        )

        assert len(result.modality_names) == 3
        assert result.common_dim == 32

    def test_forward_produces_output(self):
        """Full forward pass (encode → decode) should work."""
        model_a = _make_model(10, 32, 8, seed=0)
        model_b = _make_model(15, 48, 6, seed=1)

        n = 40
        data_a = np.random.rand(n, 10).astype(np.float32)
        data_b = np.random.rand(n, 15).astype(np.float32)

        bridge = ModalityBridge(common_dim=64, verbose=False)
        result = bridge.merge(
            models=[model_a, model_b],
            paired_data=[data_a, data_b],
            epochs=20,
        )

        # Forward with modality 0
        x = torch.rand(5, 8)
        with torch.no_grad():
            out = result(x, modality=0)
        assert out.shape[0] == 5

    def test_translate_between_modalities(self):
        """Translate from one modality to another via common space."""
        model_a = _make_model(10, 32, 8, seed=0)
        model_b = _make_model(15, 48, 6, seed=1)

        n = 40
        data_a = np.random.rand(n, 10).astype(np.float32)
        data_b = np.random.rand(n, 15).astype(np.float32)

        bridge = ModalityBridge(common_dim=64, verbose=False)
        result = bridge.merge(
            models=[model_a, model_b],
            paired_data=[data_a, data_b],
            epochs=20,
        )

        # Translate: encode as modality 0, get common representation
        x = torch.rand(5, 8)
        with torch.no_grad():
            common = result.translate(x, from_modality=0, to_modality=1)
        assert common.shape == (5, 64)

    def test_paired_alignment(self):
        """Paired samples should be closer in common space than random ones."""
        # Create models where paired data actually correlates
        model_a = _make_model(10, 32, 8, seed=0)
        model_b = _make_model(10, 48, 8, seed=1)

        # Use the SAME input data for both (perfect pairing)
        shared = np.random.rand(60, 10).astype(np.float32)

        bridge = ModalityBridge(common_dim=64, verbose=False)
        result = bridge.merge(
            models=[model_a, model_b],
            paired_data=[shared, shared],
            epochs=50,
        )

        # Encode same data through both modalities
        with torch.no_grad():
            out_a = model_a(torch.from_numpy(shared[:10]).float())
            out_b = model_b(torch.from_numpy(shared[:10]).float())
            enc_a = result.encode(out_a, modality=0)
            enc_b = result.encode(out_b, modality=1)

        # Paired should have higher similarity than random
        paired_sim = F.cosine_similarity(enc_a, enc_b).mean().item()
        # Shift one by 5 positions (mismatched pairs)
        random_sim = F.cosine_similarity(enc_a, enc_b.roll(5, 0)).mean().item()

        assert paired_sim > random_sim, (
            f"Paired similarity ({paired_sim:.3f}) should be higher "
            f"than random ({random_sim:.3f})"
        )

    def test_summary(self):
        """Summary should include modality info."""
        model_a = _make_model(10, 32, 8, seed=0)
        model_b = _make_model(15, 48, 6, seed=1)

        data_a = np.random.rand(30, 10).astype(np.float32)
        data_b = np.random.rand(30, 15).astype(np.float32)

        bridge = ModalityBridge(verbose=False)
        result = bridge.merge(
            models=[model_a, model_b],
            paired_data=[data_a, data_b],
            modality_names=["text", "protein"],
            epochs=10,
        )

        summary = result.summary()
        assert "text" in summary
        assert "protein" in summary
        assert "BridgedModel" in summary

    def test_validation_errors(self):
        """Should raise on invalid input."""
        bridge = ModalityBridge(verbose=False)

        # Single model
        with pytest.raises(ValueError, match="at least 2"):
            bridge.merge(
                models=[_make_model(10, 32, 8)],
                paired_data=[np.random.rand(10, 10).astype(np.float32)],
            )

        # Mismatched number of data arrays
        with pytest.raises(ValueError, match="paired data"):
            bridge.merge(
                models=[_make_model(10, 32, 8), _make_model(15, 48, 6)],
                paired_data=[np.random.rand(10, 10).astype(np.float32)],
            )

        # Different number of samples
        with pytest.raises(ValueError, match="same number of samples"):
            bridge.merge(
                models=[_make_model(10, 32, 8), _make_model(15, 48, 6)],
                paired_data=[
                    np.random.rand(10, 10).astype(np.float32),
                    np.random.rand(20, 15).astype(np.float32),
                ],
            )


class TestUnsupervisedBridge:
    def test_unsupervised_two_modalities(self):
        """Bridge two modalities WITHOUT paired data."""
        model_a = _make_model(10, 32, 8, seed=0)
        model_b = _make_model(20, 64, 12, seed=1)

        # Completely independent data — NOT paired
        data_a = np.random.rand(60, 10).astype(np.float32)
        data_b = np.random.rand(80, 20).astype(np.float32)  # different count!

        bridge = ModalityBridge(common_dim=64, verbose=False)
        result = bridge.unsupervised_merge(
            models=[model_a, model_b],
            unpaired_data=[data_a, data_b],
            modality_names=["text", "image"],
            epochs=30,
        )

        assert isinstance(result, BridgedModel)
        assert result.common_dim == 64

        # Can encode both modalities
        x_a = torch.rand(5, 8)
        x_b = torch.rand(5, 12)
        with torch.no_grad():
            enc_a = result.encode(x_a, modality=0)
            enc_b = result.encode(x_b, modality=1)
        assert enc_a.shape == (5, 64)
        assert enc_b.shape == (5, 64)

    def test_unsupervised_produces_output(self):
        """Full forward pass works with unsupervised bridge."""
        model_a = _make_model(10, 32, 8, seed=0)
        model_b = _make_model(15, 48, 6, seed=1)

        data_a = np.random.rand(50, 10).astype(np.float32)
        data_b = np.random.rand(70, 15).astype(np.float32)

        bridge = ModalityBridge(common_dim=32, verbose=False)
        result = bridge.unsupervised_merge(
            models=[model_a, model_b],
            unpaired_data=[data_a, data_b],
            epochs=20,
        )

        x = torch.rand(5, 8)
        with torch.no_grad():
            out = result(x, modality=0)
        assert out.shape[0] == 5

    def test_unsupervised_three_modalities(self):
        """Bridge three modalities without pairs."""
        models = [_make_model(d, 32, d+2, seed=i) for i, d in enumerate([10, 15, 20])]
        data = [np.random.rand(40 + i*10, d).astype(np.float32) for i, d in enumerate([10, 15, 20])]

        bridge = ModalityBridge(common_dim=32, verbose=False)
        result = bridge.unsupervised_merge(
            models=models,
            unpaired_data=data,
            epochs=20,
        )

        assert len(result.modality_names) == 3

