#!/usr/bin/env python3
"""Independent, integer-only checking of terminal rejection certificates.

The search implementation is C++; this checker uses Python's standard library.
It performs no Gaussian elimination and does not trust the enumerator's pivots.
Run with one or more *.jsonl certificate streams. Empty streams are valid only
when the associated completed search log reports endpoint_pass=0.
"""
from __future__ import annotations
import argparse
from collections import Counter
from fractions import Fraction
from itertools import combinations
import json
from math import lcm
from pathlib import Path
import re
from typing import Any


def integers(xs: list[Any]) -> list[int]:
    return [int(x) for x in xs]


def check_certificate(c: dict[str, Any]) -> str:
    P = c['P']; kap = c['kappa']; k = c['k']; h = c['h']; t = len(P)
    assert len(kap) == t and all(len(row) == k for row in P)
    assert all(type(x) is int and 0 <= x <= 3 for row in P for x in row)
    assert all(type(x) is int and x > 0 for x in kap)
    assert all(sum(row[j] for row in P) == 3 for j in range(k))
    r = [sum(row)-a for row, a in zip(P,kap)]
    assert min(r) > 0 and sum(r) == 3*h and sum(kap) == 3*(k-h)
    kind = c['type']
    if kind == 'inconsistent':
        y=integers(c['y']); assert len(y)==t
        assert all(sum(y[u]*P[u][j] for u in range(t))==0 for j in range(k))
        assert sum(y[u]*kap[u] for u in range(t)) != 0
    elif kind == 'rank_deficient':
        z=integers(c['z']); assert len(z)==k and any(z)
        assert all(sum(P[u][j]*z[j] for j in range(k))==0 for u in range(t))
    elif kind == 'coefficient_bounds':
        mu=[Fraction(x) for x in c['mu']]; assert len(mu)==k
        D=lcm(*(x.denominator for x in mu)); z=[x.numerator*(D//x.denominator) for x in mu]
        assert all(sum(P[u][j]*z[j] for j in range(k))==D*kap[u] for u in range(t))
        assert any(not (0<x<D) for x in z)
        B=[integers(row) for row in c['left_rows']]; ds=integers(c['diagonal'])
        assert len(B)==len(ds)==k and all(len(row)==t for row in B) and all(x>0 for x in ds)
        assert all(sum(B[i][u]*P[u][j] for u in range(t))==(ds[i] if i==j else 0)
                   for i in range(k) for j in range(k))
    elif kind == 'forced_equality':
        u,v=c['pair']; assert 0<=u<v<t
        z=integers(c['z']); D=int(c['scale']); assert len(z)==k and D>0
        assert all(sum(P[i][j]*z[j] for j in range(k))==D*((i==u)-(i==v)) for i in range(t))
    else:
        raise ValueError(f'Not a rejection certificate: {kind}')
    return kind


def parse_footer(path: Path) -> tuple[dict[str,int], set[int]]:
    lines=path.read_text().splitlines()
    assert lines and lines[-1].startswith('COMPLETE '), f'Incomplete run: {path}'
    footer={key:int(value) for key,value in re.findall(r'(\w+)=(\d+)',lines[-1])}
    case_ids=[]
    for line in lines[:-1]:
        if line.startswith('case='):
            case_ids.append(int(re.match(r'case=(\d+)',line)[1]))
    assert len(case_ids)==len(set(case_ids))==footer['cases']
    assert footer['found']==0
    return footer,set(case_ids)


def check_stream(path: Path) -> dict[str,Any]:
    counts=Counter(); ids=set()
    with path.open() as f:
        for line_number,line in enumerate(f,1):
            if not line.strip():continue
            c=json.loads(line)
            try:kind=check_certificate(c)
            except Exception as exc:raise AssertionError(f'{path}:{line_number}: invalid certificate') from exc
            counts[kind]+=1;ids.add(c['case_id'])
    log=path.with_suffix('.txt')
    footer,case_ids=parse_footer(log)
    assert sum(counts.values())==footer['endpoint_pass']
    assert counts['inconsistent']==footer['inconsistent']
    assert counts['rank_deficient']==footer['rankbad']
    assert counts['coefficient_bounds']==footer['boundsbad']
    assert counts['forced_equality']==footer['equal']
    assert ids <= case_ids
    return {'file':path.name,'verified':True,'certificates':sum(counts.values()),
            'types':dict(counts),'completed_cases':footer['cases'],'footer':footer}


def main() -> None:
    if not __debug__:
        raise RuntimeError("Run without -O: assertion checks are required.")
    ap=argparse.ArgumentParser(description=__doc__)
    ap.add_argument('streams',type=Path,nargs='+')
    ap.add_argument('--output',type=Path)
    args=ap.parse_args()
    result=[check_stream(p) for p in args.streams]
    text=json.dumps({'all_verified':True,'streams':result},indent=2)+'\n'
    if args.output:args.output.write_text(text)
    print(text,end='')

if __name__=='__main__':main()