"""
Tests for model save/load functionality.
"""

import os
import json
import numpy as np
import pytest
import torch
import torch.nn as nn

from model_osmosis.stacking import StackedModel, ModelStacker, SafeStacker


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),
    )
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_normal_(p)
    model.eval()
    return model


class TestSaveLoad:
    def test_save_creates_files(self, tmp_path):
        """Save should create {name}.pt, {name}_config.json, {name}_readme.txt."""
        experts = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        stacked = StackedModel(experts=experts, input_dim=6, output_dim=3)

        stacked.save(str(tmp_path), name="my_model")

        assert os.path.exists(tmp_path / "my_model.pt")
        assert os.path.exists(tmp_path / "my_model_config.json")
        assert os.path.exists(tmp_path / "my_model_readme.txt")

    def test_save_default_name(self, tmp_path):
        """Default name is 'model'."""
        experts = [_make_model(6, 16, 3, seed=i) for i in range(2)]
        stacked = StackedModel(experts=experts, input_dim=6, output_dim=3)

        stacked.save(str(tmp_path))

        assert os.path.exists(tmp_path / "model.pt")
        assert os.path.exists(tmp_path / "model_config.json")

    def test_save_config_is_valid(self, tmp_path):
        """Config should have all necessary fields."""
        experts = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        stacked = StackedModel(
            experts=experts, input_dim=6, output_dim=3,
            expert_names=["a", "b", "c"],
        )

        stacked.save(str(tmp_path), name="test")

        with open(tmp_path / "test_config.json") as f:
            config = json.load(f)

        assert config["num_experts"] == 3
        assert config["input_dim"] == 6
        assert config["output_dim"] == 3
        assert config["expert_names"] == ["a", "b", "c"]
        assert config["total_params"] > 0

    def test_load_roundtrip(self, tmp_path):
        """Save then load should produce same outputs."""
        experts = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        stacked = StackedModel(
            experts=experts, input_dim=6, output_dim=3,
            expert_names=["code", "math", "chat"],
        )
        stacked.eval()

        stacked.save(str(tmp_path), name="roundtrip")
        loaded = StackedModel.load(str(tmp_path), name="roundtrip")

        # State dicts should be identical
        for key in stacked.state_dict():
            torch.testing.assert_close(
                stacked.state_dict()[key], loaded.state_dict()[key],
            )

        # Same input should produce same output
        torch.manual_seed(42)
        x = torch.rand(10, 6)
        with torch.no_grad():
            out_original = stacked(x)
            out_loaded = loaded(x)
        torch.testing.assert_close(out_original, out_loaded)

    def test_load_preserves_config(self, tmp_path):
        """Loaded model should have same config."""
        experts = [_make_model(6, 16, 3, seed=i) for i in range(2)]
        stacked = StackedModel(
            experts=experts, input_dim=6, output_dim=3,
            top_k=1, expert_names=["alpha", "beta"],
        )

        stacked.save(str(tmp_path), name="cfg")
        loaded = StackedModel.load(str(tmp_path), name="cfg")

        assert loaded.num_experts == 2
        assert loaded.input_dim == 6
        assert loaded.output_dim == 3
        assert loaded.top_k == 1
        assert loaded.expert_names == ["alpha", "beta"]

    def test_multiple_models_same_folder(self, tmp_path):
        """Multiple models can be saved to the same folder."""
        m1 = StackedModel(
            experts=[_make_model(6, 16, 3, seed=0)],
            input_dim=6, output_dim=3,
        )
        m2 = StackedModel(
            experts=[_make_model(6, 32, 3, seed=1)],
            input_dim=6, output_dim=3,
        )

        m1.save(str(tmp_path), name="code_model")
        m2.save(str(tmp_path), name="math_model")

        assert os.path.exists(tmp_path / "code_model.pt")
        assert os.path.exists(tmp_path / "math_model.pt")

        loaded1 = StackedModel.load(str(tmp_path), name="code_model")
        loaded2 = StackedModel.load(str(tmp_path), name="math_model")

        assert loaded1.num_experts == 1
        assert loaded2.num_experts == 1

    def test_stacker_auto_save(self, tmp_path):
        """ModelStacker auto-saves when save_to is provided."""
        models_dir = str(tmp_path / "models")
        models = [_make_model(6, 16, 3, seed=i) for i in range(3)]
        sample = np.random.rand(30, 6).astype(np.float32)

        stacker = ModelStacker(verbose=False)
        result = stacker.stack(
            models,
            sample_data=sample,
            save_to=models_dir,
            model_name="mega",
        )

        assert os.path.exists(os.path.join(models_dir, "mega.pt"))
        assert os.path.exists(os.path.join(models_dir, "mega_config.json"))

        loaded = StackedModel.load(models_dir, name="mega")
        x = torch.rand(5, 6)
        with torch.no_grad():
            out = loaded(x)
        assert out.shape == (5, 3)

    def test_safe_stacker_save(self, tmp_path):
        """SafeStacker result can be saved and loaded."""
        eval_data = np.random.rand(50, 6).astype(np.float32)
        safe = SafeStacker(eval_data=eval_data, verbose=False)
        safe.add(_make_model(6, 16, 3, seed=0), name="code")
        safe.add(_make_model(6, 32, 3, seed=1), name="math")

        result = safe.build()
        result.save(str(tmp_path / "models"), name="safe_model")

        loaded = StackedModel.load(str(tmp_path / "models"), name="safe_model")
        x = torch.rand(5, 6)
        with torch.no_grad():
            out = loaded(x)
        assert out.shape == (5, 3)
