"""The length bound of the completeness lemma is optimal for the token format.

A literal token covering c bytes occupies c + 1 bytes and a repeat token
covering a run of c equal bytes occupies 2, so the shortest recipe for a given
output is a shortest path over the positions of that output. The exact optimum
is computed here by that recurrence and compared with |f| + ceil(|f|/127) + 1
on strings with no two equal adjacent bytes, where no repeat token can cover
more than one byte and the bound is therefore attained by every shortest
recipe, and with the encoders of the artifact on arbitrary strings.
"""
import itertools
import json
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from selfrep import describe, describe_literal, decode

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 bound(n: int) -> int:
    return n + -(-n // 127) + 1


def shortest(f: bytes) -> int:
    """Length of a shortest r in R with delta(r) = f, by dynamic programming.

    cost[i] is the least number of bytes needed to cover f[i:]; the END byte is
    added once at the end. A literal token may carry 1..127 bytes at a cost of
    one more than it carries; a repeat token carries 1..127 copies of one byte
    at a cost of two.
    """
    n = len(f)
    INF = float("inf")
    cost = [INF] * (n + 1)
    cost[n] = 0
    for i in range(n - 1, -1, -1):
        best = INF
        for c in range(1, min(127, n - i) + 1):
            v = cost[i + c] + c + 1                       # literal token
            if v < best:
                best = v
        run = 0
        while run < min(127, n - i) and f[i + run] == f[i]:
            run += 1
            v = cost[i + run] + 2                         # repeat token
            if v < best:
                best = v
        cost[i] = best
    return cost[0] + 1


def run_free(f: bytes) -> bool:
    return all(f[i] != f[i + 1] for i in range(len(f) - 1))


def main() -> int:
    out = {}
    ok = True

    # Exhaustive over ternary strings with no two equal adjacent bytes.
    tested = 0
    bad = []
    for n in range(0, 11):
        for t in itertools.product((0, 1, 2), repeat=n):
            f = bytes(t)
            if not run_free(f):
                continue
            tested += 1
            s = shortest(f)
            if s != bound(n) or len(describe_literal(f)) != bound(n):
                bad.append(n)
    print(f"  every ternary string of length at most 10 with no two equal "
          f"adjacent bytes ({tested:,} of them): the shortest recipe has "
          f"exactly |f| + ceil(|f|/127) + 1 bytes: "
          f"{'yes' if not bad else f'NO at lengths {sorted(set(bad))[:5]}'}")
    ok &= not bad
    out["run_free_exhaustive"] = tested
    out["run_free_failures"] = len(bad)

    # Longer run-free strings over the full alphabet, including the lengths at
    # which a further literal token becomes necessary.
    rng = random.Random(4)
    lens = [0, 1, 126, 127, 128, 253, 254, 255, 380, 381, 508, 509, 700]
    lens += [rng.randrange(0, 900) for _ in range(120)]
    lbad = []
    for n in lens:
        f = bytearray()
        prev = -1
        for _ in range(n):
            c = rng.randrange(256)
            while c == prev:
                c = rng.randrange(256)
            f.append(c)
            prev = c
        f = bytes(f)
        if shortest(f) != bound(n) or len(describe_literal(f)) != bound(n):
            lbad.append(n)
    print(f"  {len(lens)} further run-free strings, lengths to 900 and at every "
          f"multiple of 127: shortest recipe equals the bound: "
          f"{'yes' if not lbad else f'NO at {lbad[:5]}'}")
    ok &= not lbad
    out["run_free_sampled"] = len(lens)
    out["run_free_sampled_failures"] = len(lbad)

    # On arbitrary strings the bound is an upper bound only; record how far the
    # encoder of the artifact is from the optimum.
    gaps = []
    abad = 0
    for _ in range(300):
        n = rng.randrange(0, 600)
        alpha = rng.choice([1, 2, 3, 9, 256])
        f = bytes(rng.randrange(alpha) for _ in range(n))
        s = shortest(f)
        d = len(describe(f))
        if decode(describe(f)) != f or d < s or d > bound(n):
            abad += 1
        gaps.append(d - s)
    print(f"  300 arbitrary strings: describe is never shorter than the optimum "
          f"and never longer than the bound: {'yes' if abad == 0 else 'NO'}; "
          f"excess over the optimum: max {max(gaps)}, mean {sum(gaps)/len(gaps):.2f}")
    ok &= abad == 0
    out["arbitrary_strings"] = len(gaps)
    out["arbitrary_failures"] = abad
    out["describe_excess_max"] = max(gaps)
    out["describe_excess_mean"] = sum(gaps) / len(gaps)

    # The serialization of the host is not run-free, so its optimum is shorter.
    from selfrep import HOST_PATH
    sigma = open(HOST_PATH, "rb").read()
    s = shortest(sigma)
    print(f"  sigma(N_host): {len(sigma):,} bytes; bound {bound(len(sigma)):,}, "
          f"describe {len(describe(sigma)):,}, shortest recipe {s:,}")
    out["sigma_bytes"] = len(sigma)
    out["sigma_bound"] = bound(len(sigma))
    out["sigma_describe"] = len(describe(sigma))
    out["sigma_shortest"] = s
    ok &= s <= len(describe(sigma)) <= bound(len(sigma))

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


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