"""Unit tests for pure weight-merge backend."""

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

from model_osmosis.format.weight_merge import merge_weights_to_osmo


def _make_model(seed: int, out_features: int = 4) -> nn.Module:
    torch.manual_seed(seed)
    model = nn.Sequential(
        nn.Linear(4, out_features),
        nn.ReLU(),
        nn.Linear(out_features, 3),
    )
    model.eval()
    return model


def test_weight_merge_is_deterministic_for_same_inputs():
    model_a = _make_model(0)
    model_b = _make_model(1)

    result_a = merge_weights_to_osmo(
        [model_a, model_b],
        source_refs=["pt:a.pt", "pt:b.pt"],
        strategy="auto",
        weights=[0.5, 0.5],
        coverage_floor=1.0,
        seed=42,
    )
    result_b = merge_weights_to_osmo(
        [model_a, model_b],
        source_refs=["pt:a.pt", "pt:b.pt"],
        strategy="auto",
        weights=[0.5, 0.5],
        coverage_floor=1.0,
        seed=42,
    )

    report_a = dict(result_a.report)
    report_b = dict(result_b.report)
    report_a.pop("selection_latency_ms", None)
    report_b.pop("selection_latency_ms", None)
    assert report_a == report_b
    assert result_a.report["strategy_used"] == result_b.report["strategy_used"]
    assert result_a.report["strategy_used"] in {
        "weighted_average",
        "slerp",
        "ties",
        "dare",
        "breadcrumbs",
        "consensus",
    }
    assert result_a.report["skip_reason_counts"] == {
        "dtype_mismatch": 0,
        "missing": 0,
        "shape_mismatch": 0,
    }

    keys_a = sorted(result_a.merged_osmo.parameters.keys())
    keys_b = sorted(result_b.merged_osmo.parameters.keys())
    assert keys_a == keys_b
    for key in keys_a:
        np.testing.assert_array_equal(
            result_a.merged_osmo.parameters[key],
            result_b.merged_osmo.parameters[key],
        )


def test_weight_merge_reports_skip_reasons_for_incompatible_tensors():
    model_a = _make_model(2, out_features=4)
    model_b = _make_model(3, out_features=5)

    result = merge_weights_to_osmo(
        [model_a, model_b],
        source_refs=["pt:a.pt", "pt:b.pt"],
        strategy="weighted_average",
        coverage_floor=0.0,
        seed=7,
    )

    assert result.report["compatible_tensor_count"] >= 0
    assert result.report["skipped_tensor_count"] > 0
    skip_counts = result.report["skip_reason_counts"]
    assert (skip_counts["shape_mismatch"] + skip_counts["missing"] + skip_counts["dtype_mismatch"]) > 0


def test_weight_merge_coverage_floor_failure():
    model_a = _make_model(4, out_features=4)
    model_b = _make_model(5, out_features=7)

    with pytest.raises(RuntimeError, match="coverage floor not met"):
        merge_weights_to_osmo(
            [model_a, model_b],
            source_refs=["pt:a.pt", "pt:b.pt"],
            strategy="weighted_average",
            coverage_floor=0.70,
            seed=42,
        )
