"""
Independent rigorous computation of OEIS A198683: number of distinct values of
i^i^...^i (n copies, all parenthesizations), principal branch.

Representation (see METHOD.md):
  Every value v != 0 is represented by its principal logarithm Lambda(v),
  Im in (-pi, pi], which is unique for v.  Lambda is held as an exact
  symbolic "pi-exponential sum"
      X = sum_j c_j * pi * exp(E_j),      c_j in Q(i),  E_j again such a sum,
  interned in a hash-consing table (id 0 = the empty sum = 0).
  v^w = exp(exp(Lambda_w) * Lambda_v) so
      Lambda(v^w) = sum_j c_j pi exp(E_j + Lambda_w)  +  2 pi i k
  where k is the unique integer bringing Im into (-pi, pi]; k is determined
  either exactly (when the sum is a pure constant c*pi) or from a certified
  Arb enclosure that lies strictly inside some ((2m-1)pi, (2m+1)pi).
  Canonicalisation of exp(E): the constant coefficient c0 = r + i s of E
  (the coefficient of pi*exp(0)) has s reduced to [0, 1/2) by pulling out
  exp(i pi m/2) = i^m, an exact identity.
Equal ids  => values provably equal (all rewrites are exact identities).
Disjoint certified enclosures of Lambda => values provably distinct.
"""
import sys
from fractions import Fraction as F
from math import floor
import flint
from flint import acb, arb, ctx

sys.setrecursionlimit(100000)

ZERO = F(0)
# ---------------- interning ----------------
KEYS = [()]          # id -> key (tuple of (eid, re, im)) ; id 0 = zero sum
TABLE = {(): 0}


def intern(d):
    key = tuple(sorted((e, c[0], c[1]) for e, c in d.items()
                       if c[0] != 0 or c[1] != 0))
    i = TABLE.get(key)
    if i is None:
        i = len(KEYS)
        KEYS.append(key)
        TABLE[key] = i
    return i


def terms(i):
    return {e: (a, b) for e, a, b in KEYS[i]}


def addto(d, e, c):
    a, b = d.get(e, (ZERO, ZERO))
    a += c[0]
    b += c[1]
    if a == 0 and b == 0:
        d.pop(e, None)
    else:
        d[e] = (a, b)


IPOW = [(F(1), F(0)), (F(0), F(1)), (F(-1), F(0)), (F(0), F(-1))]


def cmul(x, y):
    return (x[0] * y[0] - x[1] * y[1], x[0] * y[1] + x[1] * y[0])


def canon_exp(d):
    """exp(pi-sum d) = i^m * exp(E) with E canonical. returns (m mod 4, id(E))."""
    d = dict(d)
    r, s = d.get(0, (ZERO, ZERO))
    m = floor(2 * s)
    t = s - F(m, 2)
    if r == 0 and t == 0:
        d.pop(0, None)
    else:
        d[0] = (r, t)
    return m % 4, intern(d)


# ---------------- numerics ----------------
_cache = {}
PI = {}


def pi_at(p):
    if p not in PI:
        ctx.prec = p
        PI[p] = arb.pi()
    return PI[p]


def enclose(i, p):
    """Arb enclosure of the pi-sum with id i, computed at working precision p."""
    k = (i, p)
    r = _cache.get(k)
    if r is not None:
        return r
    if i == 0:
        r = acb(0)
    else:
        # evaluate sub-expressions first (they may change ctx.prec? no: we set it)
        subs = [(enclose(e, p) if e != 0 else None, a, b) for e, a, b in KEYS[i]]
        ctx.prec = p
        pi = pi_at(p)
        s = acb(0)
        for ev, a, b in subs:
            c = acb(arb(a.numerator) / a.denominator, arb(b.numerator) / b.denominator)
            s += c * pi * (ev.exp() if ev is not None else acb(1))
        r = s
    _cache[k] = r
    return r


PRECS = [64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]

STATS = {"exact_k": 0, "numeric_k": 0, "max_prec_k": 0, "max_absL": 0.0}


def is_constant(i):
    key = KEYS[i]
    return all(e == 0 for e, _, _ in key)


class Unresolved(Exception):
    pass


def branch_k(i):
    """integer k with Im(X_i) + 2 pi k in (-pi, pi]."""
    if is_constant(i):
        STATS["exact_k"] += 1
        s = terms(i).get(0, (ZERO, ZERO))[1]   # Im = s*pi exactly
        # need s + 2k in (-1, 1]  ->  k = ceil((s-1)/2)... compute directly
        k = -floor((s + 1) / 2)
        if s + 2 * k <= -1:
            k += 1
        assert -1 < s + 2 * k <= 1
        return k
    for p in PRECS:
        ctx.prec = p
        x = enclose(i, p).imag / pi_at(p)       # Im/pi
        if not x.is_finite():
            continue
        # want integer m with x strictly in (2m-1, 2m+1)
        f = ((x + 1) / 2).floor().unique_fmpz()
        if f is None:
            continue
        m = int(f)
        if x > 2 * m - 1 and x < 2 * m + 1:
            STATS["numeric_k"] += 1
            STATS["max_prec_k"] = max(STATS["max_prec_k"], p)
            return -m
    raise Unresolved("branch undetermined for id %d" % i)


FLAGGED = set()   # ids of RAW (branch-unreduced) logs, only allowed at final level


def power(lv, lw, final=False):
    """Lambda(v^w) given ids of Lambda(v), Lambda(w).
    If final=True and the branch cannot be certified, returns the raw sum
    (a logarithm of v^w, not nec. principal) and records it in FLAGGED."""
    W = terms(lw)
    out = {}
    for e, a, b in KEYS[lv]:
        d = terms(e)
        for e2, c in W.items():
            addto(d, e2, c)
        m, eid = canon_exp(d)
        addto(out, eid, cmul((a, b), IPOW[m]))
    raw = intern(out)
    try:
        k = branch_k(raw)
    except Unresolved:
        if not final:
            raise
        FLAGGED.add(raw)
        return raw
    if k:
        addto(out, 0, (ZERO, F(2 * k)))
        return intern(out)
    return raw


# ---------------- distinctness ----------------
def boxes(ids, p):
    ctx.prec = p
    res = {}
    for i in ids:
        z = enclose(i, p)
        res[i] = (z.real.lower(), z.real.upper(), z.imag.lower(), z.imag.upper(), z)
    return res


def overlapping_pairs(ids, p):
    B = boxes(ids, p)
    ctx.prec = p
    order = sorted(ids, key=lambda i: B[i][0])
    pairs = []
    active = []
    for i in order:
        rl, rh, il, ih, _ = B[i]
        active = [j for j in active if B[j][1] >= rl]
        for j in active:
            if i in FLAGGED or j in FLAGGED or not (B[j][3] < il or B[j][2] > ih):
                pairs.append((j, i))
        active.append(i)
    return pairs


def certify_distinct(ids):
    """returns list of pairs that could not be separated."""
    ids = list(ids)
    pairs = overlapping_pairs(ids, PRECS[0])
    bad = []
    for (a, b) in pairs:
        ok = False
        for p in PRECS:
            ctx.prec = p
            za, zb = enclose(a, p), enclose(b, p)
            if a in FLAGGED or b in FLAGGED:
                # Re(log) = log|value| is branch independent; Im is not.
                sep = not za.real.overlaps(zb.real)
            else:
                sep = not (za.real.overlaps(zb.real) and za.imag.overlaps(zb.imag))
            if sep:
                ok = True
                break
        if not ok:
            bad.append((a, b))
    return bad, len(pairs)


def main(N):
    LI = intern({0: (ZERO, F(1, 2))})        # Log i = i pi/2
    V = {1: {LI}}
    print("a(1) = 1", flush=True)
    for n in range(2, N + 1):
        S = set()
        for k in range(1, n):
            for lv in V[k]:
                for lw in V[n - k]:
                    S.add(power(lv, lw, final=(n == N)))
        bad, npairs = certify_distinct(S)
        V[n] = S
        maxabs = 0
        for i in S:
            ctx.prec = 64
            z = enclose(i, 64)
            if i not in FLAGGED:
                maxabs = max(maxabs, float(abs(z).upper().log()) / 2.302585)
        print("a(%d) = %d   (distinct ids; %d close pairs re-checked; unresolved %d; "
              "max log10|Lambda| %.2f; interned %d; stats %s)"
              % (n, len(S), npairs, len(bad), maxabs, len(KEYS), STATS), flush=True)
        for i in FLAGGED & S:
            ctx.prec = 256
            z = enclose(i, 256)
            print("  FLAGGED (branch not certifiable) raw log =", z.str(8), flush=True)
        if bad:
            for a, b in bad:
                print("  UNRESOLVED", a, b, KEYS[a], KEYS[b])


if __name__ == "__main__":
    main(int(sys.argv[1]) if len(sys.argv) > 1 else 12)
