"""Closed forms for the number of steps of the constructor and of the copier.

Reading a recipe token by token, the decoding loop of P spends a fixed number
of steps on the token's tag and a fixed number on each byte it contributes to
the output, so the whole run is a sum over the token census:

    steps(P, r) = 5 f_lit + 4 f_rep + 7 L + 9 R + 8,

where L and R count the literal and repeat tokens, f_lit and f_rep the bytes
each kind contributes, and the last term is the end token together with the
halt instruction. The copying phase of P* costs six steps per tape byte and
two to detect the end of tape. The formulas are compared here against the
reference run on the recipes of the paper and on random recipes.
"""
import json
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from selfrep import (HOST_PATH, M_P, M_STAR, describe, describe_literal, decode,
                     run_reference, tau_star, Tape, HALT_PC)

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def runs_path(name: str) -> str:
    d = os.path.join(REPO, "paper", "runs")
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, name)


def census(r: bytes):
    """(literal tokens, repeat tokens, bytes from each) of a recipe."""
    i = 0
    lit = rep = f_lit = f_rep = 0
    while True:
        t = r[i]
        i += 1
        if t == 0:
            return lit, rep, f_lit, f_rep
        if t <= 127:
            lit += 1
            f_lit += t
            i += t
        else:
            rep += 1
            f_rep += 256 - t
            i += 1


def steps_P(r: bytes) -> int:
    """Closed form for the run of the constructor P on a recipe."""
    lit, rep, f_lit, f_rep = census(r)
    return 5 * f_lit + 4 * f_rep + 7 * lit + 9 * rep + 8


def steps_P_star(tau: bytes):
    """Closed form for the three phases of P* on its own description tape."""
    lit, rep, f_lit, f_rep = census(tau)
    a = 5 * f_lit + 4 * f_rep + 7 * lit + 9 * rep + 7
    return {"A": a, "R": 1, "B": 6 * len(tau) + 2}


def upper_bound(r: bytes) -> int:
    """5|delta(r)| + 9k + 8 with k the number of tokens, k <= (|r|-1)/2."""
    lit, rep, f_lit, f_rep = census(r)
    return 5 * (f_lit + f_rep) + 9 * (lit + rep) + 8


def coarse_bound(r: bytes) -> int:
    return 5 * len(decode(r)) + (9 * (len(r) - 1)) // 2 + 8


def phase_split(mem0, tau: bytes):
    mem = list(mem0)
    dev = Tape(tau)
    pc = 0
    counts = {"A": 0, "R": 0, "B": 0}
    while pc != HALT_PC:
        counts["A" if pc < 3 * 20 else ("R" if pc == 3 * 20 else "B")] += 1
        A = mem[pc]
        B = mem[(pc + 1) & 0xFF]
        C = mem[(pc + 2) & 0xFF]
        r = (mem[B] - mem[A]) & 0xFF
        mem[B] = r
        pc = C if (r == 0 or r >= 0x80) else (pc + 3) & 0xFF
        dev.apply(mem)
    return counts


def main() -> int:
    out = {}
    ok = True
    sigma = open(HOST_PATH, "rb").read()

    named = [("r_H = describe(sigma(N_host))", describe(sigma)),
             ("the all-literal encoding of sigma(N_host)", describe_literal(sigma)),
             ("tau* (the self-reproducing tape)", tau_star(sigma, M_STAR))]
    rows = []
    for label, r in named:
        _, n = run_reference(M_P, r)
        pred = steps_P(r)
        rows.append({"tape": label, "recipe_bytes": len(r), "output_bytes":
                     len(decode(r)), "steps": n, "closed_form": pred,
                     "agrees": n == pred, "upper": upper_bound(r),
                     "coarse": coarse_bound(r)})
        print(f"  {label}: {n:,} steps, closed form {pred:,}, "
              f"{'agree' if n == pred else 'DISAGREE'}; "
              f"bound 5|f|+9k+8 = {upper_bound(r):,}", flush=True)
        ok &= n == pred and n <= upper_bound(r) <= coarse_bound(r)
    out["named"] = rows

    # random recipes: both encoders, and strings built to force each token kind
    rng = random.Random(1009)
    bad = 0
    tested = 0
    for _ in range(400):
        n = rng.randrange(0, 700)
        alpha = rng.choice([1, 2, 3, 17, 256])
        f = bytes(rng.randrange(alpha) for _ in range(n))
        for enc in (describe, describe_literal):
            r = enc(f)
            _, s = run_reference(M_P, r)
            tested += 1
            if s != steps_P(r) or s > upper_bound(r) or upper_bound(r) > coarse_bound(r):
                bad += 1
    print(f"  {tested:,} random runs of P: closed form exact and within the "
          f"bound on every one: {'yes' if bad == 0 else f'NO ({bad})'}", flush=True)
    ok &= bad == 0
    out["random_runs"] = tested
    out["random_failures"] = bad

    # the three phases of P*
    tau = tau_star(sigma, M_STAR)
    measured = phase_split(M_STAR, tau)
    pred = steps_P_star(tau)
    agree = measured == pred
    print(f"  P* phases measured {measured}, closed form {pred}: "
          f"{'agree' if agree else 'DISAGREE'}", flush=True)
    print(f"  total {sum(measured.values()):,} = {sum(pred.values()):,}", flush=True)
    ok &= agree
    out["phases_measured"] = measured
    out["phases_closed_form"] = pred
    out["total"] = sum(pred.values())

    # phase B is six steps per tape byte and two more, over recipes of every
    # shape; P* is run only on recipes, on which it halts
    bbad = 0
    probes = 0
    for _ in range(40):
        for enc in (describe, describe_literal):
            probe = enc(bytes(rng.randrange(rng.choice([2, 4, 256]))
                              for _ in range(rng.randrange(0, 200))))
            got = phase_split(M_STAR, probe)
            probes += 1
            if got != steps_P_star(probe):
                bbad += 1
    print(f"  phase B = 6|tau| + 2 and phase R = 1 on {probes} further recipes: "
          f"{'yes' if bbad == 0 else f'NO ({bbad})'}", flush=True)
    ok &= bbad == 0
    out["phase_b_failures"] = bbad

    json.dump(out, open(runs_path("paper_steps.json"), "w"), indent=1)
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())