#  Copyright © 2025 Emmi AI GmbH. All rights reserved.

from pathlib import Path
from unittest.mock import Mock

import pytest
import torch.nn as nn

from noether.core.initializers.checkpoint import CheckpointInitializer
from noether.core.initializers.resume import ResumeInitializer
from noether.core.models import Model
from noether.core.providers import PathProvider
from noether.core.schemas.initializers import CheckpointInitializerConfig, ResumeInitializerConfig
from noether.core.schemas.models.base import ModelBaseConfig
from noether.core.schemas.optimizers import OptimizerConfig
from noether.core.utils.model import compute_model_norm
from noether.core.utils.training.training_iteration import TrainingIteration


class DummyModel(Model):
    """
    A simple dummy model for testing, stored checkpoint weights are generated by using torch.manual_seed(0), which have norm 2.4621.
    """

    def __init__(self, model_config: ModelBaseConfig):
        super().__init__(model_config=model_config)
        self.name = "dummy_model"
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)


@pytest.fixture
def mock_path_provider():
    """Fixture for a mocked PathProvider."""
    provider = Mock()
    provider.stage_name = "provider_stage_name"
    return provider


@pytest.fixture
def path_provider_with_stage_name():
    """Fixture for a mocked PathProvider with a specific stage name."""
    return PathProvider(
        output_root_path=Path("tests/test_files"),
        run_id="test_run_id",
        stage_name="train",
        debug=False,
    )


@pytest.fixture
def base_config_dict():
    """Fixture for a base configuration dictionary."""
    return {
        "run_id": "test_run_id",
        "model_name": "dummy_model",
        "load_optim": True,
        "model_info": "ema",
        "pop_ckpt_kwargs_keys": ["key1"],
        "stage_name": "train",
        "checkpoint": "latest",
    }


@pytest.fixture
def dummy_model():
    return DummyModel(
        ModelBaseConfig(
            name="dummy_model", kind="...", optimizer_config=OptimizerConfig(kind="torch.optim.SGD", lr=0.01)
        )
    )


class DummyInitializer(CheckpointInitializer):
    def __init__(self, initializer_config, **kwargs):
        super().__init__(initializer_config=initializer_config, **kwargs)

    def init_weights(self, model):
        pass

    def init_optimizer(self, model):
        pass


def test_checkpoint_initializer_init_with_string_checkpoint(base_config_dict, mock_path_provider):
    """Test initialization with a string checkpoint."""
    config = CheckpointInitializerConfig(**base_config_dict)
    initializer = DummyInitializer(initializer_config=config, path_provider=mock_path_provider)

    assert initializer.run_id == "test_run_id"
    assert initializer.model_name == "dummy_model"
    assert initializer.load_optim is True
    assert initializer.model_info == "ema"
    assert initializer.pop_ckpt_kwargs_keys == ["key1"]
    assert initializer.stage_name == "train"
    assert initializer.checkpoint == "latest"


def test_checkpoint_initializer_init_with_dict_checkpoint(base_config_dict, mock_path_provider):
    """Test initialization with a dictionary checkpoint."""
    base_config_dict["checkpoint"] = "E10_U1000_S5000"
    config = CheckpointInitializerConfig(**base_config_dict)
    initializer = DummyInitializer(initializer_config=config, path_provider=mock_path_provider)

    assert isinstance(initializer.checkpoint, str)
    assert initializer.checkpoint == "E10_U1000_S5000"

    base_config_dict["checkpoint"] = {"epoch": 10, "update": 1000, "sample": 5000}
    config = CheckpointInitializerConfig(**base_config_dict)
    initializer = DummyInitializer(initializer_config=config, path_provider=mock_path_provider)

    assert isinstance(initializer.checkpoint, TrainingIteration)
    assert initializer.checkpoint.epoch == 10
    assert initializer.checkpoint.update == 1000
    assert initializer.checkpoint.sample == 5000
    assert initializer.checkpoint.is_fully_specified

    base_config_dict["checkpoint"] = {"epoch": 10}

    config = CheckpointInitializerConfig(**base_config_dict)
    initializer = DummyInitializer(initializer_config=config, path_provider=mock_path_provider)

    assert isinstance(initializer.checkpoint, TrainingIteration)
    assert initializer.checkpoint.epoch == 10
    assert initializer.checkpoint.is_minimally_specified


def test_get_checkpoint_from_uri(base_config_dict, path_provider_with_stage_name):
    """Test getting checkpoint from URI."""

    config = CheckpointInitializerConfig(**base_config_dict)
    initializer = DummyInitializer(initializer_config=config, path_provider=path_provider_with_stage_name)

    prefix = "a"
    postfix = "b"

    path = (
        path_provider_with_stage_name.output_root
        / base_config_dict["run_id"]
        / base_config_dict["stage_name"]
        / "checkpoints"
        / f"{prefix}{base_config_dict['checkpoint']}{postfix}"
    )

    assert initializer._get_checkpoint_uri("a", "b") == path


def test_get_model_state_dict(base_config_dict, path_provider_with_stage_name, dummy_model):
    """Test getting model state dict from checkpoint."""

    config = CheckpointInitializerConfig(**base_config_dict)
    initializer = DummyInitializer(initializer_config=config, path_provider=path_provider_with_stage_name)

    state_dict, model_name, ckpt_uri = initializer._get_model_state_dict(dummy_model)

    assert model_name == "dummy_model"
    assert (
        ckpt_uri
        == path_provider_with_stage_name.output_root
        / base_config_dict["run_id"]
        / base_config_dict["stage_name"]
        / "checkpoints"
        / "dummy_model_cp=latest_ema_model.th"
    )
    assert pytest.approx(sum([p.norm() for p in state_dict.values()]), 0.0001) == 2.4621

    dummy_model.load_state_dict(state_dict)
    norm_after_load = sum([p.norm() for p in dummy_model.state_dict().values()])
    assert pytest.approx(norm_after_load, 0.0001) == 2.4621


def test_get_modelname_and_checkpoint_uri(base_config_dict, path_provider_with_stage_name, dummy_model):
    """Test errors in getting model name and checkpoint URI."""

    config = CheckpointInitializerConfig(**base_config_dict)
    initializer = DummyInitializer(initializer_config=config, path_provider=path_provider_with_stage_name)

    with pytest.raises(ValueError):
        initializer._get_modelname_and_checkpoint_uri(model=None, model_name=None, file_type="model")

    with pytest.raises(ValueError):
        initializer._get_modelname_and_checkpoint_uri(model=dummy_model, model_name=None, file_type="invalid_type")

    model_name, checkpoint_uri = initializer._get_modelname_and_checkpoint_uri(
        model=dummy_model, model_name=None, file_type="model"
    )
    assert model_name == "dummy_model"
    assert (
        checkpoint_uri
        == path_provider_with_stage_name.output_root
        / base_config_dict["run_id"]
        / base_config_dict["stage_name"]
        / "checkpoints"
        / "dummy_model_cp=latest_ema_model.th"
    )

    model_name, checkpoint_uri = initializer._get_modelname_and_checkpoint_uri(
        model=None, model_name="dummy_model", file_type="model"
    )
    assert model_name == "dummy_model"
    assert (
        checkpoint_uri
        == path_provider_with_stage_name.output_root
        / base_config_dict["run_id"]
        / base_config_dict["stage_name"]
        / "checkpoints"
        / "dummy_model_cp=latest_ema_model.th"
    )


def test_resume_initializer(base_config_dict, path_provider_with_stage_name, dummy_model):
    """Test ResumeInitializer specific behavior."""
    config = ResumeInitializerConfig(**base_config_dict)
    initializer = ResumeInitializer(initializer_config=config, path_provider=path_provider_with_stage_name)

    model_norm = compute_model_norm(dummy_model)
    initializer.init_weights(dummy_model)
    assert compute_model_norm(dummy_model) != model_norm
    assert pytest.approx(compute_model_norm(dummy_model).detach().numpy(), 0.0001) == 2.4621
    assert initializer.model_name == "dummy_model"
    assert initializer.load_optim is True
