#!/usr/bin/env python3
"""GaugeBench. Exact Q rate on listed polys. PROCESS vs float AD vs finite diff."""
from __future__ import annotations

from fractions import Fraction as F

import torch

from PROCESS import add, isolate, mix, rate


def sq(x):
    return mix(x, x)


def cu(x):
    return mix(sq(x), x)


def twox(x):
    return add(x, x)


CASES = [
    ("x*x at 3", lambda x: sq(x), 3, F(6)),
    ("x*x at -2", lambda x: sq(x), -2, F(-4)),
    ("x*x at 1/2", lambda x: sq(x), F(1, 2), F(1)),
    ("2x at 7", twox, 7, F(2)),
    ("x**3 at 3", cu, 3, F(27)),
    ("x**3 at 2/3", cu, F(2, 3), F(4, 3)),
]


def process_row(fn, v):
    seeds = (1, 2, F(1, 2), -1)
    got = [rate(fn(isolate(v, s)), isolate(v, s)) for s in seeds]
    if any(g != got[0] for g in got):
        return None
    return got[0]


def torch_row(name, v):
    vt = float(v)
    x = torch.tensor(vt, requires_grad=True)
    if name.startswith("x*x"):
        y = x * x
    elif name.startswith("2x"):
        y = x + x
    elif name.startswith("x**3"):
        y = x * x * x
    else:
        return None
    y.backward()
    return float(x.grad)


def fd_row(name, v, h=1e-6):
    vt = float(v)
    def f(z):
        if name.startswith("x*x"):
            return z * z
        if name.startswith("2x"):
            return z + z
        if name.startswith("x**3"):
            return z * z * z
        return None
    return (f(vt + h) - f(vt)) / h


def main():
    print(f"{'case':<16} {'exact':<10} {'process':<10} {'torch':<14} {'fd':<14}")
    p_ok = t_ok = f_ok = 0
    n = len(CASES)
    for name, fn, v, want in CASES:
        p = process_row(fn, v)
        t = torch_row(name, v)
        fd = fd_row(name, v)
        p_hit = p == want
        t_hit = abs(t - float(want)) < 1e-9
        f_hit = abs(fd - float(want)) < 1e-6
        p_ok += int(p_hit)
        t_ok += int(t_hit)
        f_ok += int(f_hit)
        print(
            f"{name:<16} {str(want):<10} {str(p):<10} {t:<14.8g} {fd:<14.8g} "
            f"{'P' if p_hit else 'p'}{'T' if t_hit else 't'}{'F' if f_hit else 'f'}"
        )
    print(f"SCORE process {p_ok}/{n}  torch {t_ok}/{n}  fd {f_ok}/{n}")
    print("GAUGEBENCH PASS", p_ok == n)


if __name__ == "__main__":
    main()