#!/usr/bin/env python3
"""Process invariants of a /solve-maths run, independent of its mathematics.

usage: python3 check_invariants.py <run>/data.json [--json out.json]

Reads only the run record (data.json and returns/). Each invariant is a
property of the record that a correct run must satisfy whatever the
mathematics says. Prints one line per invariant: name, pass/fail, counts.
"""
import json, os, sys
from datetime import datetime

def ts(s):
    return datetime.fromisoformat(s.replace('Z', '+00:00'))

def main(path, out=None):
    root = os.path.dirname(os.path.abspath(path))
    D = json.load(open(path))
    run_agents = [a['id'] for a in D['agents'] if a.get('round') in (1, 2)]
    nodes = [n for n in D['nodes'] if n['id'].split('-')[0] in run_agents]
    roots = [n for n in D['root_nodes'] if n['id'] in ('R-1', 'R-2', 'R-3', 'R-4')]
    by = {n['id']: n for n in D['nodes'] + D['root_nodes']}
    lane = lambda i: i.split('-')[0]
    res = []

    # I1 independence: a round-1 step cites only steps of its own agent
    r1 = [(p, n['id']) for n in nodes if n.get('round') == 1 for p in n['parents']]
    bad = [e for e in r1 if lane(e[0]) != lane(e[1])]
    res.append(('I1 independence', not bad, f'{len(bad)} of {len(r1)} round-1 edges cross agents', bad))

    # I2 provenance: every parent exists; every agent step is in its verbatim return
    edges = [(p, n['id']) for n in nodes + roots for p in n['parents']]
    miss = [e for e in edges if e[0] not in by]
    ret = set()
    for f in os.listdir(os.path.join(root, 'returns')):
        if f.endswith('.json') and f != 'briefs.json':
            ret |= {m['id'] for m in json.load(open(os.path.join(root, 'returns', f))).get('nodes', [])}
    unret = [n['id'] for n in nodes if n['id'] not in ret]
    res.append(('I2 provenance', not miss and not unret,
                f'{len(miss)} dangling parents; {len(unret)} of {len(nodes)} steps missing from returns/', miss + unret))

    # I3 time order: a step is written after every step it cites
    late = [e for e in edges if e[0] in by and ts(by[e[0]]['t']) > ts(by[e[1]]['t'])]
    res.append(('I3 time order', not late, f'{len(late)} of {len(edges)} edges point backwards in time', late))

    # I4 budget: at most nine subagents across all rounds
    res.append(('I4 budget', len(run_agents) <= 9, f'{len(run_agents)} agents (limit 9)', []))

    # I5 audited before use: every step the root cites as "proved" was checked
    # in the run by an agent other than its author before the root cited it
    cited = sorted({(p, r['id']) for r in roots for p in r['parents'] if by[p].get('status') == 'proved'})
    def checked(p, t):
        return any(p in m['parents'] and lane(m['id']) != lane(p) and ts(m['t']) <= t
                   for m in nodes if m['kind'] in ('check', 'objection', 'verdict'))
    unch = [(p, r) for p, r in cited if not checked(p, ts(by[r]['t']))]
    res.append(('I5 audited before use', not unch,
                f'{len(unch)} of {len(cited)} proved inputs cited by the root before an independent check', unch))

    # I6 labels: every step carries a kind and a status from the schema
    K = {'idea', 'lemma', 'computation', 'check', 'objection', 'blocked', 'verdict'}
    S = {'proved', 'refuted', 'blocked', 'open', 'heuristic', 'superseded', 'partial'}
    unl = [n['id'] for n in nodes if n.get('kind') not in K or n.get('status') not in S]
    res.append(('I6 labels', not unl, f'{len(unl)} of {len(nodes)} steps with a missing or unknown label', unl))

    for name, ok, msg, ex in res:
        print(f"{name:24s} {'pass' if ok else 'FAIL'}  {msg}" + (f"  e.g. {ex[:4]}" if ex else ''))
    if out:
        json.dump([{'name': a, 'pass': b, 'msg': c, 'examples': d} for a, b, c, d in res], open(out, 'w'), indent=1)

if __name__ == '__main__':
    a = sys.argv[1:]
    main(a[0], a[a.index('--json') + 1] if '--json' in a else None)