"""
Tests for knowledge representation and storage.
"""

import tempfile
from pathlib import Path

import numpy as np
import pytest

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


@pytest.fixture
def sample_knowledge():
    """Create a sample UniversalKnowledge object."""
    structural = StructuralKnowledge(
        model_type="pytorch",
        architecture_desc="Test model",
        layer_info=[
            {"name": "layer1", "type": "Linear", "in_features": 10, "out_features": 32},
            {"name": "layer2", "type": "Linear", "in_features": 32, "out_features": 5},
        ],
        parameters={
            "layer1.weight": np.random.rand(32, 10).astype(np.float32),
            "layer1.bias": np.random.rand(32).astype(np.float32),
            "layer2.weight": np.random.rand(5, 32).astype(np.float32),
            "layer2.bias": np.random.rand(5).astype(np.float32),
        },
        total_params=32 * 10 + 32 + 5 * 32 + 5,
    )

    inputs = np.random.rand(100, 10).astype(np.float32)
    outputs = np.random.rand(100, 5).astype(np.float32)

    functional = FunctionalKnowledge(
        inputs=inputs,
        outputs=outputs,
        is_classifier=True,
        num_classes=5,
    )

    meta = MetaKnowledge(
        source_framework="pytorch",
        source_class_name="TestModel",
        input_shape=(10,),
        output_shape=(5,),
        task_type="classification",
    )

    return UniversalKnowledge(
        structural=structural,
        functional=functional,
        meta=meta,
    )


class TestUniversalKnowledge:
    def test_creation(self, sample_knowledge):
        assert sample_knowledge.source_type == "pytorch"
        assert sample_knowledge.param_count == 32 * 10 + 32 + 5 * 32 + 5
        assert sample_knowledge.is_classifier is True

    def test_summary(self, sample_knowledge):
        summary = sample_knowledge.summary()
        assert "pytorch" in summary
        assert "classification" in summary
        assert "TestModel" in summary

    def test_output_shape_auto_inferred(self):
        func = FunctionalKnowledge(
            inputs=np.random.rand(10, 4).astype(np.float32),
            outputs=np.random.rand(10, 3).astype(np.float32),
        )
        assert func.output_shape == (3,)


class TestStorage:
    def test_save_and_load_roundtrip(self, sample_knowledge):
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "test_knowledge.osm"
            save_knowledge(sample_knowledge, path)

            loaded = load_knowledge(path)

            assert isinstance(loaded, UniversalKnowledge)
            assert loaded.source_type == sample_knowledge.source_type
            assert loaded.param_count == sample_knowledge.param_count
            assert loaded.meta.source_class_name == "TestModel"
            np.testing.assert_array_equal(
                loaded.functional.inputs, sample_knowledge.functional.inputs
            )
            np.testing.assert_array_equal(
                loaded.functional.outputs, sample_knowledge.functional.outputs
            )

    def test_load_nonexistent_raises(self):
        with pytest.raises(FileNotFoundError):
            load_knowledge("/nonexistent/path/knowledge.osm")

    def test_load_wrong_type_raises(self):
        import pickle

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "bad.osm"
            with open(path, "wb") as f:
                pickle.dump({"not": "knowledge"}, f)

            with pytest.raises(TypeError):
                load_knowledge(path)

    def test_save_creates_parents(self, sample_knowledge):
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "sub" / "dir" / "knowledge.osm"
            result = save_knowledge(sample_knowledge, path)
            assert result.exists()
