"""Export the interval covers for the RS closed-form rate to Lean modules.

Usage (from the repository root):
    python formal/scripts/export_rs_cover.py [outdir]

Writes `<outdir>/RSCover<k>.lean` (default `formal/MaxLeafFormal`), k = 0, 1, ...  Each module holds
a chain of consecutive intervals [a, b] with exact rational endpoints, given by their right ends, and
a kernel check (`decide +kernel`) of `MaxLeafFormal.RSRate.checkCover`: on each interval the
monotone upper bound of F*,

    ln 2 * F*(s) <= b (-ln a) + (1 - a)(-ln(1 - b)) + (1 - b) ln(1 - 2^(-1/a))   (s in [a, b]),

evaluated with the rational enclosures of `MaxLeafFormal.Enclose` (`lnQ`, `expQ`, `ln2Q`), lies below
the target times a lower bound of ln 2.  This script evaluates exactly the same rational functions
(a mirror of the Lean definitions in `Enclose.lean` and `RSCheck.lean`, in Python `Fraction`s), so
every exported interval passes the Lean check; a wrong export can only make a check fail.

The covers are produced by the adaptive bisection of
`ancillary/rs_closed_form_interval.py` (same ranges, same targets apart
from cover 1, same rule: bisect at the midpoint until accepted).  Cover 1 uses the paper's
intermediate value 0.8030906143 (proof of `prop:rs-rate`(c)) as its target; `RSRate.lean` proves
0.8030906143 < log2 1.744835.  The interval counts are not part of any claim.

    cover 1: [1/5, 4/5],       target 8030906143/10^10
    cover 2: [15/100, 31/100], target 7807/10^4
    cover 3: [54/100, 80/100], target 7802/10^4
"""
import hashlib
import math
import os
import sys
from fractions import Fraction as Fr

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SRC_REL = "ancillary/rs_closed_form_interval.py"
SRC = os.path.join(ROOT, *SRC_REL.split("/"))
CHUNK = 500

# Parameters of the enclosures; must equal `nL`, `nE`, `pr` in RSCheck.lean.
NL = 32
NE = 12
PR = 40

# ---------------------------------------------------------------------------------------------
# Mirror of MaxLeafFormal.Enclose (exact rational arithmetic)
# ---------------------------------------------------------------------------------------------


def floor_q(q):
    return q.numerator // q.denominator


def ceil_q(q):
    return -((-q.numerator) // q.denominator)


def rdn(q, prec):
    return Fr(floor_q(q * 2 ** prec), 2 ** prec)


def rup(q, prec):
    return Fr(ceil_q(q * 2 ** prec), 2 ** prec)


def ilog2(n):
    """`ilog2 1024 n` of Enclose.lean: floor(log2 n) for n >= 2, else 0."""
    return n.bit_length() - 1 if n >= 2 else 0


def log_series(x, n):
    return sum((x ** (i + 1) / (i + 1) for i in range(n)), Fr(0))


def log_err(x, n):
    return abs(x) ** (n + 1) / (1 - abs(x))


def neg_log1m(x, n, prec):
    xl, xh = rdn(x, prec), rup(x, prec)
    if abs(xl) < 1 and abs(xh) < 1:
        return (rdn(log_series(xl, n) - log_err(xl, n), prec),
                rup(log_series(xh, n) + log_err(xh, n), prec))
    return (x, x / (1 - x) if x < 1 else Fr(0))


_ln2 = {}


def ln2q(n, prec):
    if (n, prec) not in _ln2:
        _ln2[(n, prec)] = neg_log1m(Fr(1, 2), n, prec)
    return _ln2[(n, prec)]


def pow2z(k):
    return Fr(2) ** k


def red_exp(q):
    k = ilog2(abs(q.numerator)) - ilog2(q.denominator)
    return k - 1 if q / pow2z(k) < 1 else k


def ln_q(q, n, prec):
    k = red_exp(q)
    m = q / pow2z(k)
    x = 1 - 1 / m
    if 0 <= x < 1:
        lo2, hi2 = ln2q(n, prec)
        sm = (k * lo2, k * hi2) if k >= 0 else (k * hi2, k * lo2)
        nl = neg_log1m(x, n, prec)
        return (rdn(sm[0] + nl[0], prec), rup(sm[1] + nl[1], prec))
    return (1 - 1 / q, q - 1)


def exp_series(t, n):
    return sum((t ** i / math.factorial(i) for i in range(n)), Fr(0))


def exp_err(t, n):
    return abs(t) ** n * (Fr(n + 1) / (math.factorial(n) * n))


def halvings(x):
    return ilog2(max(ceil_q(abs(x)), 0)) + 1


def exp_q(x, n, prec):
    k = halvings(x)
    t = x / 2 ** k
    N = n + 1
    if abs(t) <= 1:
        T = exp_series(t, N)
        e = exp_err(t, N)
        lo, hi = max(Fr(0), rdn(T - e, prec)), rup(T + e, prec)
        for _ in range(k):
            lo, hi = rdn(lo * lo, prec), rup(hi * hi, prec)
        return (lo, hi)
    return (Fr(0), Fr(3) ** max(ceil_q(x), 0))


# ---------------------------------------------------------------------------------------------
# Mirror of MaxLeafFormal.RSRate.cellOK (RSCheck.lean)
# ---------------------------------------------------------------------------------------------


def cell_ok(T, a, b):
    if not (0 < a < b < 1):
        return False
    L2lo, L2hi = ln2q(NL, PR)
    LA = ln_q(a, NL, PR)[0]
    LB = ln_q(1 - b, NL, PR)[0]
    PL = exp_q(-L2hi / a, NE, PR)[0]
    if not PL < 1:
        return False
    U = ln_q(1 - PL, NL, PR)[1]
    if not U <= 0:
        return False
    return b * (-LA) + (1 - a) * (-LB) + (1 - b) * U < T * L2lo


def float_bound(a, b):
    a, b = float(a), float(b)
    L2 = math.log(2)
    return (b * -math.log(a) + (1 - a) * -math.log(1 - b) + (1 - b) * math.log(1 - 2 ** (-1 / a))) / L2


def cover(lo, hi, target):
    stack = [(Fr(lo), Fr(hi))]
    out = []
    while stack:
        a, b = stack.pop()
        # floating-point pre-filter (only skips the exact check where it would clearly fail)
        if float_bound(a, b) < float(target) + 1e-7 and cell_ok(target, a, b):
            out.append((a, b))
            continue
        assert b - a > Fr(1, 10 ** 12), (a, b)
        m = (a + b) / 2
        stack.append((m, b))
        stack.append((a, m))
    for (a, b), (c, d) in zip(out, out[1:]):
        assert b == c
    assert out[0][0] == lo and out[-1][1] == hi
    return out


def q(x):
    x = Fr(x)
    assert x >= 0
    return f"q {x.numerator} {x.denominator}"


def src_sha():
    # hash of the file with LF line endings (as committed), independent of a CRLF checkout
    data = open(SRC, "rb").read().replace(b"\r\n", b"\n")
    return hashlib.sha256(data).hexdigest()


SPECS = [
    ("1", Fr(1, 5), Fr(4, 5), Fr(8030906143, 10 ** 10)),
    ("2", Fr(15, 100), Fr(31, 100), Fr(7807, 10 ** 4)),
    ("3", Fr(54, 100), Fr(80, 100), Fr(7802, 10 ** 4)),
]


def main(outdir):
    sha = src_sha()
    k = 0
    index = []
    for name, lo, hi, target in SPECS:
        cells = cover(lo, hi, target)
        print(f"cover {name}: {len(cells)} intervals", file=sys.stderr)
        chunks = [cells[i:i + CHUNK] for i in range(0, len(cells), CHUNK)]
        for ch in chunks:
            a0, b1 = ch[0][0], ch[-1][1]
            body = [
                "import MaxLeafFormal.RSCheck",
                "",
                "/-!",
                f"# RS closed-form rate: cover {name}, chunk {k} (generated)",
                "",
                "Generated by `formal/scripts/export_rs_cover.py` (bisection as in",
                f"`{SRC_REL}`,",
                f"SHA-256 `{sha}`).",
                f"{len(ch)} intervals covering [{a0}, {b1}] for the target {target}.",
                "Do not edit by hand. The data are not trusted: `decide +kernel` re-checks every interval.",
                "-/",
                "",
                "set_option maxRecDepth 100000",
                "",
                "namespace MaxLeafFormal.RSRate",
                "",
                f"/-- The right ends of the intervals, starting from `{a0}`. -/",
                f"def cover{k} : List ℚ := [",
                ",\n".join(f"  {q(b)}" for _, b in ch),
                "]",
                "",
                f"theorem cover{k}_ok :",
                f"    checkCover ({q(target)}) ({q(b1)}) ({q(a0)}) cover{k} = true := by",
                "  decide +kernel",
                "",
                f"theorem cover{k}_sound : ∀ σ ∈ Set.Icc ((({q(a0)}) : ℚ) : ℝ) ((({q(b1)}) : ℚ) : ℝ),",
                f"    rsF σ ≤ ((({q(target)}) : ℚ) : ℝ) :=",
                f"  checkCover_sound cover{k}_ok",
                "",
                "end MaxLeafFormal.RSRate",
            ]
            path = os.path.join(outdir, f"RSCover{k}.lean")
            with open(path, "w", encoding="utf-8", newline="\n") as fh:
                fh.write("\n".join(body) + "\n")
            index.append((k, name, a0, b1, target, len(ch)))
            k += 1
    for row in index:
        print(*row)


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else os.path.join(ROOT, "formal", "MaxLeafFormal"))