#!/usr/bin/env python3
"""A198683: rigorous computation of the number of distinct values of i^i^...^i.

Usage: rigorous.py N PREC_BITS

Each value v is tracked two ways:
  * L : an Arb ball (rigorous enclosure) of the principal Log v (Im in (-pi, pi]).
  * P : an exact symbolic expression equal to Log v.

Symbolic expressions ("Expr") are finite sums of terms  c * pi^k * E(Q)  where
c is a Gaussian rational, k >= 0, and E(Q) = exp(Q) for a nested Expr Q (or 1).
Rules:
    Log(i)   = (i/2) pi
    Log(a^b) = b * Log(a) - 2 pi i k        (integer k proven from the ball)
    b        = E(P(b))                       (exactly, since P(b) is a log of b)
    E(Q) * E(R) = E(Q + R)
    E(Q + (m/2) pi i) = i^m E(Q)             (pull out exact powers of i; this also reduces Q mod 2 pi i)

Identical canonical Exprs  => values PROVEN equal.
Disjoint balls             => values PROVEN distinct.
Overlapping balls but different Exprs => UNRESOLVED (reported).
"""
import sys
import time
import math
from fractions import Fraction as Fr
from flint import acb, arb, ctx

N = int(sys.argv[1]) if len(sys.argv) > 1 else 12
PREC = int(sys.argv[2]) if len(sys.argv) > 2 else 2000
ctx.prec = PREC
PI = arb.pi()
TWO_PI = 2 * PI
BIG = 230          # |b| > e^BIG -> b treated as astronomically large (level-2)

ZERO = (Fr(0), Fr(0))
I_POW = [(Fr(1), Fr(0)), (Fr(0), Fr(1)), (Fr(-1), Fr(0)), (Fr(0), Fr(-1))]


def cmul(a, b):
    return (a[0] * b[0] - a[1] * b[1], a[0] * b[1] + a[1] * b[0])


# Expr: dict {(pi_pow, ekey): coeff}; ekey = () for no exponential, else frozen canonical Expr tuple
def freeze(P):
    return tuple(sorted(P.items()))


def e_add(P, Q):
    out = dict(P)
    for k, c in Q.items():
        a = out.get(k, ZERO)
        s = (a[0] + c[0], a[1] + c[1])
        if s == ZERO:
            out.pop(k, None)
        else:
            out[k] = s
    return out


def canon_E(Q):
    """E(Q) -> (gaussian factor, ekey). Pulls exact powers of i out of the (pi^1, no-E) term."""
    Q = dict(Q)
    key = (1, ())
    m = 0
    if key in Q:
        re, im = Q[key]
        m = math.floor(2 * im)              # im = m/2 + r, 0 <= r < 1/2
        im2 = im - Fr(m, 2)
        if re == 0 and im2 == 0:
            del Q[key]
        else:
            Q[key] = (re, im2)
    return I_POW[m % 4], (freeze(Q) if Q else ())


def unfreeze(ekey):
    return dict(ekey)


def times_b(P, Pb):
    """b * P where b = E(Pb)."""
    out = {}
    for (k, ek), c in P.items():
        g, ek2 = canon_E(e_add(unfreeze(ek), Pb) if ek else Pb)
        c2 = cmul(c, g)
        kk = (k, ek2)
        a = out.get(kk, ZERO)
        s = (a[0] + c2[0], a[1] + c2[1])
        if s == ZERO:
            out.pop(kk, None)
        else:
            out[kk] = s
    return out


def wrap_k(t):
    """Integer k with t - 2 pi k in (-pi, pi), proven from ball t; None if the ball touches the branch cut."""
    q = (t + PI) / TWO_PI
    k = q.mid().floor().unique_fmpz()          # exact integer, no float rounding
    if (q > arb(k)) and (q < arb(k + 1)):
        return int(k)
    return None


def main():
    t0 = time.time()
    P_i = {(1, ()): (Fr(0), Fr(1, 2))}
    V = {1: [dict(L=acb(0, PI / 2), P=P_i, key=freeze(P_i), expr="i", idx=0)]}
    mapping = {1: {}}          # mapping[n][(k, class_a, class_b)] = class at n
    level2, ambiguous = [], []
    report = {}
    for n in range(2, N + 1):
        cands = []
        cands_l2 = []
        for k in range(1, n):
            for A in V[k]:
                for B in V[n - k]:
                    if B["L"].real > BIG:
                        level2.append((n, A, B))
                        cands_l2.append((k, A["idx"], B["idx"]))
                        continue
                    raw = B["L"].exp() * A["L"]
                    kk = wrap_k(raw.imag)
                    expr = f"({A['expr']})^({B['expr']})"
                    if kk is None:
                        ambiguous.append((n, expr, raw.imag))
                        continue
                    L = acb(raw.real, raw.imag - TWO_PI * kk)
                    P = times_b(A["P"], B["P"])
                    if kk:
                        P = e_add(P, {(1, ()): (Fr(0), Fr(-2 * kk))})
                    cands.append(dict(L=L, P=P, key=freeze(P), expr=expr, src=(k, A["idx"], B["idx"])))

        cands.sort(key=lambda c: float(c["L"].real.mid()))
        reps, by_key = [], {}
        merged = 0
        unresolved = []
        for c in cands:
            r = by_key.get(c["key"])
            if r is not None:
                if not r["L"].overlaps(c["L"]):
                    raise RuntimeError("identical Expr but disjoint balls: " + c["expr"])
                merged += 1
                c["rep"] = r
                continue
            clash = None
            for rr in reversed(reps):
                if float(rr["L"].real.mid()) < float(c["L"].real.mid()) - 1.0 and rr["L"].real.upper() < c["L"].real.lower():
                    break
                if rr["L"].overlaps(c["L"]):
                    clash = rr
                    break
            if clash is not None:
                unresolved.append((c, clash))
                by_key[c["key"]] = clash        # assume equal for now (recorded)
                c["rep"] = clash
                continue
            c["rep"] = c
            reps.append(c)
            by_key[c["key"]] = c
        for j, c in enumerate(reps):
            c["idx"] = j
        mapping[n] = {c["src"]: c["rep"]["idx"] for c in cands}
        # level-2 values (exponent astronomically large): each gets its own class after the level-1 ones;
        # main() proves below that they are distinct from all level-1 values (only one occurs for n <= 12)
        for j, src in enumerate(cands_l2):
            mapping[n][src] = len(reps) + j
        V[n] = reps
        report[n] = (len(reps), merged, unresolved)
        amb_n = sum(1 for a in ambiguous if a[0] == n)
        print(f"n={n:2d}: distinct={len(reps):5d}  proven-equal merges={merged:5d}  "
              f"UNRESOLVED={len(unresolved)}  branch-ambiguous={amb_n}  [{time.time()-t0:.1f}s]", flush=True)

    # exhaustive disjointness check of representatives (sweep on real-part intervals)
    for n in range(2, N + 1):
        reps = sorted(V[n], key=lambda c: float(c["L"].real.lower()))
        active, bad = [], 0
        for c in reps:
            lo = c["L"].real.lower()
            active = [a for a in active if not (a["L"].real.upper() < lo)]
            bad += sum(1 for a in active if a["L"].overlaps(c["L"]))
            active.append(c)
        if bad:
            print(f"n={n}: WARNING {bad} overlapping representative pairs")

    for n, expr, t in ambiguous:
        print(f"AMBIGUOUS n={n}: {expr}\n    Im(raw) = {t.mid().str(25)} +/- {t.rad().str(3)}")
    for n, A, B in level2:
        LB = B["L"]
        s = LB.imag.sin()
        print(f"LEVEL-2 n={n}: ({A['expr']})^(<huge>) with Re Log b = {LB.real.mid().str(15)}, "
              f"sin(Im Log b) = {s.str(10)}")
    # level-2 distinctness proof: v = exp(b * Log a) cannot be enclosed, but Re Log v = Re(b * Log a) can
    # (Arb exponents are unbounded). Prove it lies outside the real-part range of every level-1 value at n
    # and that level-2 values at the same n are pairwise separated.
    for n in range(2, N + 1):
        l2 = [(A, B) for m, A, B in level2 if m == n]
        if not l2:
            continue
        lo = min((c["L"].real.lower() for c in V[n]), key=float)
        hi = max((c["L"].real.upper() for c in V[n]), key=float)
        res = [(B["L"].exp() * A["L"]).real for A, B in l2]
        for X in res:
            if not (X < lo or X > hi):
                raise RuntimeError(f"n={n}: level-2 value not separated from level-1 values")
        for i in range(len(res)):
            for j in range(i + 1, len(res)):
                if res[i].overlaps(res[j]):
                    raise RuntimeError(f"n={n}: level-2 values {i},{j} not separated")
        print(f"LEVEL-2 PROOF n={n}: {len(res)} value(s), Re Log v in {[x.str(6) for x in res]}; "
              f"level-1 Re Log v in [{lo.str(6)}, {hi.str(6)}] -> all distinct")
    # closest pairs at n=N (relative and absolute distance between log-balls)
    reps = sorted(V[N], key=lambda c: float(c["L"].real.mid()))
    pairs = []
    for i in range(len(reps)):
        for j in range(i + 1, min(i + 40, len(reps))):
            d = abs(reps[i]["L"] - reps[j]["L"])
            sc = max(abs(reps[i]["L"]), arb(1))
            pairs.append((float((d / sc).mid().log()) / math.log(10), float(d.mid().log()) / math.log(10), reps[i], reps[j]))
    pairs.sort(key=lambda x: x[0])
    for rel, ab, a, b in pairs[:6]:
        print(f"CLOSE n={N}: rel dist 1e{rel:.1f}, abs dist 1e{ab:.1f}, |L| ~ {abs(a['L']).mid().str(5)}")
        print(f"    {a['expr']}\n    {b['expr']}")
    if len(sys.argv) > 3 and sys.argv[3] == "--tables":
        write_tables(V, mapping, level2, N)
    print("\ncounts (level-1 distinct):", [1] + [report[n][0] for n in range(2, N + 1)])
    for n in report:
        for c, r in report[n][2][:15]:
            print(f"  UNRESOLVED n={n}:\n    {c['expr']}\n    {r['expr']}")


def write_tables(V, mapping, level2, N):
    """tables/classes_nXX.tsv: every distinct value with a certified enclosure of its principal Log.
    tables/trees_nXX.tsv: every one of the Catalan(n-1) parenthesizations with its class id.
    Tree order: t(1) = [i]; t(n) = for k = 1..n-1, for L in t(k), for R in t(n-k): (L)^(R)."""
    import os
    os.makedirs("tables", exist_ok=True)
    trees = {1: [("i", 0)]}
    for n in range(2, N + 1):
        out = []
        for k in range(1, n):
            for le, lc in trees[k]:
                for re_, rc in trees[n - k]:
                    out.append((f"({le})^({re_})" if n > 2 else "i^i", mapping[n][(k, lc, rc)]))
        trees[n] = out
    for n in range(1, N + 1):
        reps = V[n]
        nl2 = sum(1 for x in level2 if x[0] == n)
        with open(f"tables/classes_n{n:02d}.tsv", "w") as f:
            f.write("# A198683 n=%d: %d distinct values. Columns: class, Re Log v (ball), Im Log v (ball), representative\n"
                    % (n, len(reps) + nl2))
            for c in reps:
                f.write(f"{c['idx']}\t{c['L'].real.str(30, radius=True)}\t{c['L'].imag.str(30, radius=True)}\t{c['expr']}\n")
            for j, (_, A, B) in enumerate(x for x in level2 if x[0] == n):
                LB = B["L"]
                f.write(f"{len(reps)+j}\tRe Log v = -(pi/2)*exp({LB.real.str(20, radius=True)})*sin({LB.imag.str(20, radius=True)})"
                        f"\t(level-2: exponent b has Log b = above; |v| ~ exp(-0.79e^(9.49e34)))\t({A['expr']})^({B['expr']})\n")
        with open(f"tables/trees_n{n:02d}.tsv", "w") as f:
            f.write(f"# A198683 n={n}: all {len(trees[n])} parenthesizations. Columns: index, class, expression\n")
            for t, (e, c) in enumerate(trees[n]):
                f.write(f"{t}\t{c}\t{e}\n")
        ncls = len(set(c for _, c in trees[n]))
        print(f"tables n={n}: {len(trees[n])} trees -> {ncls} classes")


if __name__ == "__main__":
    main()