#!/usr/bin/env python3
"""Check tables/trees_nXX.tsv against tables/classes_nXX.tsv using only the table files.

For every parenthesization: parse the expression text, evaluate it from scratch with Arb
(principal power z^w = exp(w Log z), no shared intermediate results), take the principal Log,
and check that it lies inside the certified ball of its assigned class.
Trees whose direct evaluation is impossible (an exponent around e^(1e35)) must be assigned to
the level-2 class.

Usage: check_table.py N [PREC_BITS]
"""
import sys
from flint import acb, arb, ctx

N = int(sys.argv[1])
ctx.prec = int(sys.argv[2]) if len(sys.argv) > 2 else 3000
I = acb(0, 1)


def parse(s):
    """Grammar: E := 'i' | '(' E ')' '^' '(' E ')' | 'i^i'.  Returns nested tuples."""
    pos = 0

    def term():
        nonlocal pos
        if s.startswith("(", pos):
            depth = 0
            for j in range(pos, len(s)):
                depth += s[j] == "("
                depth -= s[j] == ")"
                if depth == 0:
                    inner = s[pos + 1:j]
                    pos = j + 1
                    return parse(inner)
            raise ValueError(s)
        if s.startswith("i", pos):
            pos += 1
            return "i"
        raise ValueError(s)

    left = term()
    if pos == len(s):
        return left
    assert s[pos] == "^", s
    pos += 1
    right = term()
    assert pos == len(s), s
    return (left, right)


def ev(t):
    if t == "i":
        return I
    a, b = ev(t[0]), ev(t[1])
    return (b * a.log()).exp()


def main():
    classes = {}
    level2 = set()
    for line in open(f"tables/classes_n{N:02d}.tsv"):
        if line.startswith("#"):
            continue
        cid, re_s, im_s, _ = line.rstrip("\n").split("\t")
        if re_s.startswith("Re Log v ="):
            level2.add(int(cid))
            continue
        classes[int(cid)] = acb(arb(re_s), arb(im_s))
    # 1) the listed class balls must be pairwise disjoint (so the classes are provably distinct values)
    items = sorted(classes.items(), key=lambda kv: float(kv[1].real.lower()))
    active, overl = [], 0
    for cid, B in items:
        lo = B.real.lower()
        active = [a for a in active if not (a[1].real.upper() < lo)]
        overl += sum(1 for a in active if a[1].overlaps(B))
        active.append((cid, B))
    print(f"n={N}: {len(classes)} level-1 class balls, overlapping pairs: {overl}")
    ok = bad = l2ok = 0
    seen = set()
    lo1 = min((B.real.lower() for B in classes.values()), key=float)
    l2re = {}
    for line in open(f"tables/trees_n{N:02d}.tsv"):
        if line.startswith("#"):
            continue
        idx, cid, expr = line.rstrip("\n").split("\t")
        cid = int(cid)
        seen.add(cid)
        v = ev(parse(expr))
        finite = v.real.is_finite() and v.imag.is_finite() and not v.contains(0)
        if cid in level2:
            # v itself cannot be enclosed, but Re Log v = Re(b * Log a) can: evaluate a and b from the text
            # and check that Re Log v lies far below every level-1 class ball
            a, b = (ev(x) for x in parse(expr))
            X = (b * a.log()).real
            if X < lo1:
                l2ok += 1
                l2re.setdefault(cid, []).append(X)
            else:
                bad += 1
                print(f"FAIL tree {idx}: level-2 Re Log v {X.str(10)} not below level-1 range: {expr}")
            continue
        if not finite:
            print(f"FAIL tree {idx}: direct evaluation not finite but class {cid} is level-1: {expr}")
            bad += 1
            continue
        L = v.log()
        if L.overlaps(classes[cid]):
            ok += 1
        else:
            bad += 1
            print(f"FAIL tree {idx}: Log value {L.str(15)} not in class {cid} ball: {expr}")
    l2 = sorted(l2re.items())
    for i in range(len(l2)):
        for j in range(i + 1, len(l2)):
            if any(x.overlaps(y) for x in l2[i][1] for y in l2[j][1]):
                bad += 1
                print(f"FAIL: level-2 classes {l2[i][0]} and {l2[j][0]} not separated")
    for cid, xs in l2:
        print(f"level-2 class {cid}: {len(xs)} trees, Re Log v = {xs[0].str(8)} < level-1 minimum {lo1.str(8)}")
    print(f"n={N}: {ok} trees consistent with their class ball, {l2ok} level-2 trees, {bad} failures; "
          f"{len(seen)} classes used (table lists {len(classes) + len(level2)})")


if __name__ == "__main__":
    main()