"""Unit tests for KG edge co-occurrence learning (Phase 17 Plan 01)."""
import torch
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

from arbitor.components import MoEGraph


def _make_mg():
    return MoEGraph(cb_dim=64, trigram_dim=512, num_experts=4, core_rank=32, shared_inter=128,
                    max_iters=2, codebook_size=16, active_graph_max_nodes=32)

def test_int_score_cooccurrence():
    tg = _make_mg()
    vq1 = torch.tensor([[2, 4, 8, 15]])
    tg.update_kg_edges(vq1)
    score_after_1 = tg.edge_score.clone()

    vq2 = torch.tensor([[2, 4, 9, 15]])
    tg.update_kg_edges(vq2)
    score_after_2 = tg.edge_score.clone()

    assert not torch.equal(score_after_1, score_after_2), "second update should change int edge scores"
    print(" PASS test_int_score_cooccurrence")


def test_ternary_quantize():
    tg = _make_mg()
    tg.edge_score.fill_(0)
    tg.edge_score[:10] = 4
    tg.edge_score[10:20] = -4
    tg._steps_since_requant.fill_(50)
    tg.update_kg_edges(torch.tensor([[0, 1, 2, 3]]))
    assert (tg.edge_attr[:10] == 1).all(), "positive should be +1"
    assert (tg.edge_attr[10:20] == -1).all(), "negative should be -1"
    print(" PASS test_ternary_quantize")


def test_batch_detection():
    tg = _make_mg()
    vq_ids = torch.randint(0, 16, (2, 4))
    old_score = tg.edge_score.clone()
    tg.update_kg_edges(vq_ids)
    assert not torch.equal(old_score, tg.edge_score), "batch update should change edge scores"
    print(" PASS test_batch_detection")


def test_int_scores_stay_finite():
    tg = _make_mg()
    for _ in range(10):
        vq = torch.randint(0, 16, (2, 4))
        tg.update_kg_edges(vq)
    assert tg.edge_score.dtype == torch.int8
    print(" PASS test_int_scores_stay_finite")


def test_checkpoint_persistence():
    tg1 = _make_mg()
    vq = torch.randint(0, 16, (2, 4))
    tg1.update_kg_edges(vq)
    sd = tg1.state_dict()

    tg2 = _make_mg()
    tg2.load_state_dict(sd)
    assert torch.equal(tg1.edge_score, tg2.edge_score), "state_dict should preserve edge scores"
    print(" PASS test_checkpoint_persistence")


if __name__ == "__main__":
    test_int_score_cooccurrence()
    test_ternary_quantize()
    test_batch_detection()
    test_int_scores_stay_finite()
    test_checkpoint_persistence()
    print("\nAll KG edge tests PASS")