#!/usr/bin/env python3
"""Cubic windows for the Appendix D maps that were still cleartext-only.

Sollya builds a longest-feasible partition whose fixed-point cubic stays
within half an ulp. Mathematica (Remez), Maple (numapprox[minimax]), and
MATLAB/Chebfun (minimax) each fit the same pieces. The shipped cubic is
the one with the smallest absolute error after rounding coefficients to
2^{-(k+16)}.
"""

from __future__ import annotations

import json
import math
import re
import subprocess
import sys
from dataclasses import dataclass
from decimal import Decimal, ROUND_HALF_UP
from pathlib import Path

PREC = (8, 12, 16, 20, 24, 28, 32)
SOLLYA = "sollya"
MATH = "/home/dawkteraitch/Wolfram/15.0/Executables/math"
MAPLE = "/home/dawkteraitch/maple2025/bin/maple"
MATLAB = "/home/dawkteraitch/MATLAB/R2026a/bin/matlab"
WORK = Path("/tmp/cas-lic/window")
EXPR = {
    "lecun": "(17159/10000) * tanh((2 * x) / 3)",
    "hard": "expm1(x) * (x + 1) / 2",
}
NUMBER = re.compile(r"[+-]?(?:\d+\.\d*|\d+\.?\d*)(?:[eE][+-]?\d+)?")


@dataclass
class Piece:
    left: int
    right: int
    coeffs: list[int]  # c0..c3 at scale 2^q
    err: float
    source: str


class Sollya:
    def __init__(self) -> None:
        self.proc: subprocess.Popen[str] | None = None
        self.start()

    def start(self) -> None:
        if self.proc is not None:
            self.proc.kill()
            self.proc.wait()
        self.proc = subprocess.Popen(
            ["stdbuf", "-o0", "-e0", SOLLYA],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )
        self._run("verbosity = 0;\nprec = 120;\n", "READY")

    def close(self) -> None:
        if self.proc is not None:
            self.proc.kill()
            self.proc.wait()
            self.proc = None

    def _run(self, script: str, sentinel: str) -> list[str]:
        assert self.proc is not None and self.proc.stdin and self.proc.stdout
        self.proc.stdin.write(script)
        self.proc.stdin.write(f'print("{sentinel}");\n')
        self.proc.stdin.flush()
        lines: list[str] = []
        while True:
            line = self.proc.stdout.readline()
            if line == "":
                raise RuntimeError("sollya exited")
            if sentinel in line:
                return lines
            lines.append(line.rstrip("\n"))

    def fit(
        self, name: str, left: int, right: int, k: int, q: int, tol: float
    ) -> tuple[float, list[int]] | None:
        # fpminimax is ill-conditioned on a tiny interval at 0 and can report
        # an error above half an ulp for a one-ulp step. Taylor is the
        # fallback, and the search keeps the smaller certified error.
        mini = self._polynomial(name, left, right, k, q, taylor=False)
        if mini is not None and mini[0] <= tol:
            return mini
        tay = self._polynomial(name, left, right, k, q, taylor=True)
        cands = [item for item in (mini, tay) if item is not None]
        if not cands:
            return None
        return min(cands, key=lambda item: item[0])

    def _polynomial(
        self, name: str, left: int, right: int, k: int, q: int, taylor: bool
    ) -> tuple[float, list[int]] | None:
        expr = EXPR[name]
        if taylor:
            # taylor(f, n, t) returns p such that p(x-t) is the expansion at t.
            build = (
                f"t0 = (({left})+({right}))/2 * 2^(-{k});\n"
                f"pshift = taylor({expr}, 3, t0);\n"
                f"p = pshift(x - t0);"
            )
        else:
            build = (
                f"p = fpminimax({expr}, 3, [|80,80,80,80|], "
                f"[{left}*2^(-{k}), {right}*2^(-{k})], absolute);"
            )
        script = f"""
{build}
qbits = {q};
i0 = nearestint(coeff(p, 0) * 2^qbits);
i1 = nearestint(coeff(p, 1) * 2^qbits);
i2 = nearestint(coeff(p, 2) * 2^qbits);
i3 = nearestint(coeff(p, 3) * 2^qbits);
pr = i0*2^(-qbits) + i1*2^(-qbits)*x + i2*2^(-qbits)*x^2 + i3*2^(-qbits)*x^3;
e = dirtyinfnorm(({expr}) - pr, [{left}*2^(-{k}), {right}*2^(-{k})]);
print(i0);
print(i1);
print(i2);
print(i3);
print(e);
"""
        try:
            lines = self._run(script, "FITDONE")
        except RuntimeError:
            self.start()
            return None
        if any("Error" in line or "error" in line for line in lines):
            self.start()
            return None
        nums: list[str] = []
        for line in lines:
            nums.extend(NUMBER.findall(line))
        if len(nums) < 5:
            self.start()
            return None
        coeffs = [int(Decimal(nums[i]).to_integral_value(ROUND_HALF_UP)) for i in range(4)]
        err = max(float(n) for n in nums[4:])
        if not math.isfinite(err):
            return None
        return err, coeffs

    def score(self, name: str, left: int, right: int, k: int, q: int, coeffs: list[int]) -> float | None:
        expr = EXPR[name]
        c0, c1, c2, c3 = coeffs
        script = f"""
pr = ({c0})*2^(-{q}) + ({c1})*2^(-{q})*x + ({c2})*2^(-{q})*x^2 + ({c3})*2^(-{q})*x^3;
e = dirtyinfnorm(({expr}) - pr, [{left}*2^(-{k}), {right}*2^(-{k})]);
print(e);
"""
        try:
            lines = self._run(script, "SCOREDONE")
        except RuntimeError:
            self.start()
            return None
        if any("Error" in line or "error" in line for line in lines):
            self.start()
            return None
        nums: list[str] = []
        for line in lines:
            nums.extend(NUMBER.findall(line))
        if not nums:
            return None
        err = max(float(n) for n in nums)
        return err if math.isfinite(err) else None


def half_ulp(k: int) -> float:
    return 2.0 ** (-(k + 1))


def q_of(k: int) -> int:
    return k + 16


def quantize(coeffs: list[str], q: int) -> list[int]:
    scale = Decimal(2) ** q
    out: list[int] = []
    for raw in coeffs:
        value = Decimal(raw) * scale
        if value >= 0:
            out.append(int(value.to_integral_value(ROUND_HALF_UP)))
        else:
            out.append(-int((-value).to_integral_value(ROUND_HALF_UP)))
    return out


def lecun_tail(k: int) -> int:
    num = 17159 * (1 << k)
    den = 10000
    return (num + den // 2) // den


def lecun_end(k: int) -> int:
    """First raw at which the constant tail is within 1/4 ulp of the asymptote gap."""
    tail = lecun_tail(k) / 2**k
    # f increases to 1.7159. Leave a quarter ulp for rounding the tail itself.
    gap = 0.25 * (2 ** (-k))
    target = 1.7159 - gap
    lo, hi = 0.0, 40.0
    for _ in range(80):
        mid = (lo + hi) / 2
        if 1.7159 * math.tanh(2 * mid / 3) >= target:
            hi = mid
        else:
            lo = mid
    return max(1, math.ceil(hi * 2**k - 1e-9))


def march(sollya: Sollya, name: str, start: int, end: int, k: int) -> list[Piece]:
    q = q_of(k)
    tol = half_ulp(k)
    pieces: list[Piece] = []
    left = start
    guess = 1
    while left < end:
        step = max(1, guess)
        last: tuple[int, float, list[int]] | None = None
        fail: int | None = None
        while left + step <= end:
            got = sollya.fit(name, left, left + step, k, q, tol)
            if got is not None and got[0] <= tol:
                last = (left + step, got[0], got[1])
                step *= 2
            else:
                fail = min(end, left + step)
                break
        else:
            if last is None or last[0] != end:
                got = sollya.fit(name, left, end, k, q, tol)
                if got is not None and got[0] <= tol:
                    last = (end, got[0], got[1])
                    fail = None
                else:
                    fail = end
        if fail is None and last is not None:
            right, err, coeffs = last
        else:
            lo = last[0] if last is not None else left
            hi = fail if fail is not None else end
            best = last
            while hi - lo > 1:
                mid = (lo + hi) // 2
                got = sollya.fit(name, left, mid, k, q, tol)
                if got is not None and got[0] <= tol:
                    best = (mid, got[0], got[1])
                    lo = mid
                else:
                    hi = mid
            if best is None:
                got = sollya.fit(name, left, left + 1, k, q, tol)
                if got is None:
                    raise RuntimeError(f"{name} k={k} no cubic on [{left}, {left+1}]")
                best = (left + 1, got[0], got[1])
            right, err, coeffs = best
        pieces.append(Piece(left, right, coeffs, err, "sollya"))
        guess = max(1, right - left)
        left = right
        print(f"  {name} k={k} pieces={len(pieces)} span={right - pieces[-1].left} err={err:.3e}", flush=True)
    return pieces


def write_intervals(pieces_by_key: dict[tuple[str, int], list[Piece]], path: Path) -> None:
    with path.open("w") as handle:
        for (name, k), pieces in pieces_by_key.items():
            q = q_of(k)
            for piece in pieces:
                handle.write(f"{name} {k} {piece.left} {piece.right} {q}\n")


def run_maple(interval_path: Path, out_path: Path) -> None:
    program = WORK / "maple_batch.mpl"
    program.write_text(
        f"""
Digits := 40:
with(numapprox):
fd := fopen("{out_path}", WRITE):
infd := fopen("{interval_path}", READ):
while true do
  line := readline(infd);
  if line = 0 then break; end if;
  parts := sscanf(line, "%s %d %d %d %d");
  gname := parts[1];
  a := parts[3] * 2^(-parts[2]);
  b := parts[4] * 2^(-parts[2]);
  try
    if gname = "lecun" then
      p := minimax(17159/10000*tanh(2*x/3), x = a .. b, [3, 0]);
    else
      p := minimax((exp(x)-1)*(x+1)/2, x = a .. b, [3, 0]);
    end if;
    p := expand(p);
    fprintf(fd, "%.25e %.25e %.25e %.25e\\n",
      evalf(coeff(p, x, 0)), evalf(coeff(p, x, 1)),
      evalf(coeff(p, x, 2)), evalf(coeff(p, x, 3)));
  catch:
    fprintf(fd, "FAIL\\n");
  end try;
end do:
fclose(fd):
fclose(infd):
quit:
"""
    )
    subprocess.run([MAPLE, "-q", str(program)], check=True, timeout=3600)


def run_math(interval_path: Path, out_path: Path) -> None:
    program = WORK / "mma_batch.wl"
    program.write_text(
        r"""
pow[base_, 0] := 1;
pow[base_, n_Integer] := base^n;
lecun[x_] := 17159/10000*Tanh[2*x/3];
hard[x_] := (Exp[x] - 1)*(x + 1)/2;
remez[ff_, a_, b_] := Module[
  {n = 3, m = 5, mid, half, nodes, coef, mat, rhs, sol, grid, err, idx, drop, tpoly, xpoly},
  mid = N[(a + b)/2, 40];
  half = N[(b - a)/2, 40];
  nodes = Sort[Table[N[Cos[Pi*(k + 1/2)/m], 40], {k, 0, m - 1}]];
  coef = {0, 0, 0, 0};
  Do[
    mat = N[Table[If[j <= n + 1, pow[nodes[[i]], j - 1], -(-1)^i], {i, m}, {j, m}], 40];
    rhs = N[Table[ff[mid + half*nodes[[i]]], {i, m}], 40];
    sol = Quiet[Check[LinearSolve[mat, rhs], $Failed]];
    If[sol === $Failed || Head[sol] =!= List, Break[]];
    coef = N[sol[[1 ;; n + 1]], 40];
    grid = N[Subdivide[-1, 1, 256], 40];
    err = Table[
      coef[[1]] + coef[[2]]*grid[[i]] + coef[[3]]*grid[[i]]^2 + coef[[4]]*grid[[i]]^3
        - ff[mid + half*grid[[i]]],
      {i, Length[grid]}];
    idx = {1, Length[grid]};
    Do[
      If[Abs[err[[i]]] >= Abs[err[[i - 1]]] && Abs[err[[i]]] >= Abs[err[[i + 1]]],
        AppendTo[idx, i]],
      {i, 2, Length[grid] - 1}];
    While[Length[idx] > m,
      drop = Ordering[Abs[err[[idx]]], 1][[1]];
      idx = Delete[idx, drop];
    ];
    If[Length[idx] == m, nodes = Sort[N[grid[[idx]], 40]]];
    , {8}];
  tpoly = coef[[1]] + coef[[2]]*((x - mid)/half) + coef[[3]]*((x - mid)/half)^2
    + coef[[4]]*((x - mid)/half)^3;
  xpoly = Expand[N[tpoly, 30]];
  Table[N[Coefficient[xpoly, x, k], 20], {k, 0, 3}]
];
lines = ReadList["""
        + f'"{interval_path}"'
        + r""", String];
out = OpenWrite["""
        + f'"{out_path}"'
        + r"""];
Do[
  parts = StringSplit[line];
  If[Length[parts] < 5, Continue[]];
  name = parts[[1]];
  kk = ToExpression[parts[[2]]];
  left = ToExpression[parts[[3]]];
  right = ToExpression[parts[[4]]];
  aa = N[left*2^(-kk), 30];
  bb = N[right*2^(-kk), 30];
  ff = If[name == "lecun", lecun, hard];
  cs = Quiet[Check[remez[ff, aa, bb], $Failed]];
  If[cs === $Failed || Head[cs] =!= List,
    WriteString[out, "FAIL\n"],
    WriteString[out, StringRiffle[Map[ToString[CForm[#]] &, cs], " "], "\n"]
  ];
  , {line, lines}];
Close[out];
Exit[];
"""
    )
    subprocess.run([MATH, "-script", str(program)], check=True, timeout=3600)


def run_matlab(interval_path: Path, out_path: Path) -> None:
    program = WORK / "matlab_batch.m"
    program.write_text(
        f"""
fid = fopen('{interval_path}', 'r');
out = fopen('{out_path}', 'w');
while true
  line = fgetl(fid);
  if ~ischar(line), break; end
  parts = strsplit(strtrim(line));
  if numel(parts) < 5, continue; end
  name = parts{{1}};
  kk = str2double(parts{{2}});
  left = str2double(parts{{3}});
  right = str2double(parts{{4}});
  aa = left * 2^(-kk);
  bb = right * 2^(-kk);
  try
    if strcmp(name, 'lecun')
      f = chebfun(@(x) 1.7159*tanh(2*x/3), [aa, bb]);
    else
      f = chebfun(@(x) expm1(x).*(x+1)/2, [aa, bb]);
    end
    p = minimax(f, 3);
    c = poly(p);
    c = c(:).';
    if numel(c) < 4
      c = [zeros(1, 4-numel(c)), c];
    end
    fprintf(out, '%.20e %.20e %.20e %.20e\\n', c(4), c(3), c(2), c(1));
  catch
    fprintf(out, 'FAIL\\n');
  end
end
fclose(fid);
fclose(out);
"""
    )
    subprocess.run(
        [MATLAB, "-batch", f"run('{program}')"],
        check=True,
        timeout=7200,
    )


def load_coeff_file(path: Path) -> list[list[str] | None]:
    rows: list[list[str] | None] = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or "FAIL" in line:
            rows.append(None)
            continue
        parts = NUMBER.findall(line.replace("D", "E"))
        if len(parts) < 4:
            rows.append(None)
            continue
        rows.append(parts[:4])
    return rows


def tournament(
    sollya: Sollya,
    pieces_by_key: dict[tuple[str, int], list[Piece]],
    others: dict[str, list[list[str] | None]],
) -> dict[str, int]:
    wins = {"sollya": 0, "mathematica": 0, "maple": 0, "matlab": 0}
    index = 0
    for key in pieces_by_key:
        name, k = key
        q = q_of(k)
        for piece in pieces_by_key[key]:
            best_err = piece.err
            best_coeffs = piece.coeffs
            best_source = "sollya"
            for source in ("mathematica", "maple", "matlab"):
                row = others[source][index] if index < len(others[source]) else None
                if row is None:
                    continue
                try:
                    coeffs = quantize(row, q)
                except Exception:
                    continue
                err = sollya.score(name, piece.left, piece.right, k, q, coeffs)
                if err is not None and err < best_err:
                    best_err = err
                    best_coeffs = coeffs
                    best_source = source
            piece.coeffs = best_coeffs
            piece.err = best_err
            piece.source = best_source
            wins[best_source] += 1
            index += 1
        print(f"tournament {name} k={k} done", flush=True)
    return wins


def split_i128(value: int) -> tuple[int, int]:
    raw = value + (1 << 128) if value < 0 else value
    lo = raw & ((1 << 64) - 1)
    hi = (raw >> 64) & ((1 << 64) - 1)
    if hi >= (1 << 63):
        hi -= 1 << 64
    return hi, lo


TAGS = {"lecun": "LECUN", "hard": "HARDELISH"}


def emit_table(name: str, k: int, pieces: list[Piece]) -> str:
    knots = [pieces[0].left] + [piece.right for piece in pieces]
    q = q_of(k)
    tag = TAGS[name]
    lines = [f"static constexpr std::int64_t {tag}_K{k}_KNOTS[] = {{"]
    lines.append("    " + ", ".join(str(knot) for knot in knots) + ",")
    lines.append("};")
    lines.append(f"static constexpr grotto::principal_detail::cubic_bits {tag}_K{k}_PIECES[] = {{")
    for piece in pieces:
        his, los = [], []
        for coeff in piece.coeffs:
            hi, lo = split_i128(coeff)
            his.append(str(hi))
            los.append(f"{lo}u")
        lines.append("    {{" + ", ".join(his) + "}, {" + ", ".join(los) + "}},")
    lines.append("};")
    lines.append(
        f"static constexpr window_table {tag}_K{k}_TABLE"
        f"{{{tag}_K{k}_KNOTS, {tag}_K{k}_PIECES, {len(pieces)}, {q}}};"
    )
    lines.append("")
    return "\n".join(lines)


def emit(pieces_by_key: dict[tuple[str, int], list[Piece]], wins: dict[str, int], path: Path) -> None:
    chunks = [
        "// Generated on mocha2.",
        "// Sollya fpminimax chose the knots (absolute error at most half an ulp).",
        "// Each cubic is the best of Sollya, Mathematica Remez, Maple numapprox[minimax],",
        "// and MATLAB/Chebfun minimax after rounding coefficients to 2^{-(k+16)}.",
        "// Wins: " + ", ".join(f"{name} {count}" for name, count in wins.items()) + ".",
        "// lecun is the positive half of an odd map; hard covers (-1, 0).",
        "// Do not edit.",
        "#pragma once",
        "",
    ]
    ks = [k for (name, k) in pieces_by_key if name == "lecun"]
    for name in ("lecun", "hard"):
        for k in ks:
            chunks.append(emit_table(name, k, pieces_by_key[(name, k)]))
        tag = TAGS[name]
        refs = ",\n    ".join(f"&{tag}_K{k}_TABLE" for k in ks)
        chunks.append(
            f"static constexpr window_table const* const {tag}[{len(ks)}] = {{\n    {refs},\n}};"
        )
        chunks.append("")
    tails = ", ".join(str(lecun_tail(k)) for k in ks)
    chunks.append(f"static constexpr std::int64_t LECUN_TAIL[{len(ks)}] = {{\n    {tails},\n}};")
    chunks.append("")
    path.write_text("\n".join(chunks))


def summarize(pieces_by_key: dict[tuple[str, int], list[Piece]], wins: dict[str, int]) -> None:
    summary = {"wins": wins, "functions": {}}
    for name in ("lecun", "hard"):
        summary["functions"][name] = {}
        for k in selected_precisions():
            pieces = pieces_by_key[(name, k)]
            summary["functions"][name][str(k)] = {
                "parts": len(pieces),
                "max_abs_error": max(piece.err for piece in pieces),
                "half_ulp": half_ulp(k),
            }
    (WORK / "summary.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2), flush=True)


def self_test() -> None:
    sollya = Sollya()
    got = sollya.fit("lecun", 0, 1 << 15, 16, q_of(16), half_ulp(16))
    print("sollya", got[0] if got else None)
    sollya.close()


def selected_precisions() -> tuple[int, ...]:
    if "--k" in sys.argv:
        return (int(sys.argv[sys.argv.index("--k") + 1]),)
    return PREC


def main() -> None:
    if len(sys.argv) > 1 and sys.argv[1] == "--self-test":
        self_test()
        return
    WORK.mkdir(parents=True, exist_ok=True)
    sollya = Sollya()
    pieces_by_key: dict[tuple[str, int], list[Piece]] = {}
    for k in selected_precisions():
        end = lecun_end(k)
        print(f"march lecun k={k} end={end} tail={lecun_tail(k)}", flush=True)
        pieces_by_key[("lecun", k)] = march(sollya, "lecun", 0, end, k)
        one = 1 << k
        print(f"march hard k={k}", flush=True)
        pieces_by_key[("hard", k)] = march(sollya, "hard", -one, 0, k)
    interval_path = WORK / "intervals.txt"
    write_intervals(pieces_by_key, interval_path)
    print("intervals", sum(len(v) for v in pieces_by_key.values()), flush=True)
    jobs = {
        "maple": (run_maple, WORK / "maple_coeffs.txt"),
        "mathematica": (run_math, WORK / "mma_coeffs.txt"),
        "matlab": (run_matlab, WORK / "matlab_coeffs.txt"),
    }
    for name, (fn, out) in jobs.items():
        # Run sequentially inside this process: MATLAB and Maple both want a license,
        # and a failure is easier to see one at a time. They do not share state.
        print("running", name, flush=True)
        if out.exists():
            out.unlink()
        fn(interval_path, out)
        print("finished", name, out.stat().st_size, flush=True)
    expected = sum(len(v) for v in pieces_by_key.values())
    others = {
        "maple": load_coeff_file(WORK / "maple_coeffs.txt"),
        "mathematica": load_coeff_file(WORK / "mma_coeffs.txt"),
        "matlab": load_coeff_file(WORK / "matlab_coeffs.txt"),
    }
    for source, rows in others.items():
        if len(rows) != expected:
            raise SystemExit(f"{source} wrote {len(rows)} rows, expected {expected}")
    wins = tournament(sollya, pieces_by_key, others)
    sollya.close()
    emit(pieces_by_key, wins, WORK / "window_extra_tables.inc")
    summarize(pieces_by_key, wins)


if __name__ == "__main__":
    main()