"""Tests for the certificate formalism.  Run with `python3 tests/test_core.py`
(or pytest if available)."""

import heapq
import math
import os
import random
import sys

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

from certlang.core import (ATOM_FRAME, RULE_FRAME, Certificate, IncrementalLearner,
                           Rule, build_library, cheapest_subset, choice,
                           exact_min_certificate, greedy_learner, popcount)
from certlang.domains import all_functions, tiny_domain


def forward_min_cost(labels, domain, library):
    """Brute-force Dijkstra over the FULL graph of extensional functions:
    nodes are all |Y|^|X| functions, edges are overwrite transformations
    T_{S,y} weighted 1 + L(S) + log2|Y|.  Source: any constant function
    (cost log2|Y| + 1 stop bit folded into totals identically to
    Certificate.cost)."""
    n, ny = domain.n, domain.ny
    target = tuple(labels)
    start_cost = choice(ny) + 1.0
    dist = {}
    heap = []
    for y in range(ny):
        g0 = tuple([y] * n)
        dist[g0] = start_cost
        heapq.heappush(heap, (start_cost, g0))
    while heap:
        g, f = heapq.heappop(heap)
        if g > dist.get(f, math.inf):
            continue
        if f == target:
            return g
        for mask, scost, _desc in library:
            for y in range(ny):
                nf = list(f)
                m = mask
                while m:
                    low = m & -m
                    nf[low.bit_length() - 1] = y
                    m ^= low
                nf = tuple(nf)
                ng = g + RULE_FRAME + scost + choice(ny)
                if ng < dist.get(nf, math.inf):
                    dist[nf] = ng
                    heapq.heappush(heap, (ng, nf))
    return math.inf


def test_backward_equals_forward():
    """On a micro domain (|X|=4, |Y|=2) the settled-set reduction must give
    exactly the same optimum as brute-force search over all 16 functions."""
    dom = tiny_domain([2], nvals=2)
    assert dom.n == 4
    lib = build_library(dom, max_atoms=2)
    for labels in all_functions(dom.n, dom.ny):
        labels = list(labels)
        cert = exact_min_certificate(labels, dom, lib)
        fwd = forward_min_cost(labels, dom, lib)
        assert cert.predict_all(dom.n) == labels, (labels, cert.explain(dom.ylabels))
        assert abs(cert.cost() - fwd) < 1e-9, (labels, cert.cost(), fwd)
    print("ok: backward settled-set search == forward function-graph Dijkstra (16/16)")


def test_certificate_semantics():
    c = Certificate(0, 2)
    c.rules.append(Rule(0b0110, 1, 1.0, "mid"))
    c.rules.append(Rule(0b0011, 0, 1.0, "low"))  # overwrites bit 1 back to 0
    assert c.predict_all(4) == [0, 0, 1, 0]
    assert c.predict(2) == 1 and c.predict(1) == 0
    obs = [0b0011, 0b0100]  # class0: {0,1}, class1: {2}
    assert c.consistent(obs, 0b1111)
    obs_bad = [0b0111, 0b1000]
    assert not c.consistent(obs_bad, 0b1111)
    print("ok: last-write-wins semantics and consistency check")


def test_exact_on_named_functions():
    dom = tiny_domain([3], nvals=2)
    lib = build_library(dom, max_atoms=3)
    n = dom.n

    def cost_of(labels):
        return exact_min_certificate(list(labels), dom, lib).cost()

    const = cost_of([0] * n)
    maj = cost_of([1 if bin(x).count("1") >= 2 else 0 for x in range(n)])
    # value at position p counts 1-bits of x in this encoding (v=1 digit)
    rnd = random.Random(1)
    randoms = [cost_of([rnd.randrange(2) for _ in range(n)]) for _ in range(20)]
    assert const < maj < min(randoms) + 1e-9 or maj < sum(randoms) / len(randoms), \
        (const, maj, randoms)
    assert const == choice(2) + 1.0
    print(f"ok: cost(const)={const:.2f} < cost(majority)={maj:.2f} "
          f"<= mean(random)={sum(randoms)/len(randoms):.2f}")


def test_incremental_consistency():
    dom = tiny_domain([3], nvals=2)
    rng = random.Random(7)
    for _ in range(50):
        labels = [rng.randrange(2) for _ in range(dom.n)]
        order = list(range(dom.n))
        rng.shuffle(order)
        learner = IncrementalLearner(dom, max_atoms=3)
        for x in order:
            learner.observe(x, labels[x])
            for xx, yy in learner.obs.items():
                assert learner.cert.predict(xx) == yy
        learner.recompress(rng=random.Random(0))
        assert learner.cert.predict_all(dom.n) == labels
    print("ok: incremental learner stays consistent through 50 random targets")


def test_cheapest_subset_universality():
    dom = tiny_domain([3], nvals=2)
    # forbid everything except x: search must fall back to (or beat) singleton
    for x in range(dom.n):
        forbidden = dom.universe & ~(1 << x)
        cost, mask, desc = cheapest_subset(dom, x, forbidden, max_atoms=3)
        assert mask == 1 << x, (x, desc)
    print("ok: singleton always reachable (universality)")


def test_kraft_sanity():
    """Costs must be strictly positive and per-family choice budgets exact:
    an atom family with k parameter combinations contributes k descriptions
    each of cost >= log2(k) + tag, so sum 2^-cost <= 1 per family holds by
    construction.  Spot-check: no atom is cheaper than its family's
    log2(#atoms-in-family) would allow after deduping."""
    dom = tiny_domain([3], nvals=2)
    assert all(a.cost > 0 for a in dom.atoms)
    total = sum(2 ** -a.cost for a in dom.atoms)
    assert total <= 2.0 + 1e-9, total  # 2 families, each sums to <= 1
    print(f"ok: kraft mass of atom library = {total:.3f} <= #families")


def test_greedy_learner():
    dom = tiny_domain([3], nvals=2)
    labels = [1 if bin(x).count("1") >= 2 else 0 for x in range(dom.n)]
    obs = {x: labels[x] for x in range(dom.n)}
    cert = greedy_learner(obs, dom, max_atoms=3)
    assert cert.predict_all(dom.n) == labels
    print(f"ok: greedy learner consistent, cost={cert.cost():.2f}")


if __name__ == "__main__":
    test_certificate_semantics()
    test_backward_equals_forward()
    test_exact_on_named_functions()
    test_incremental_consistency()
    test_cheapest_subset_universality()
    test_kraft_sanity()
    test_greedy_learner()
    print("all tests passed")