import math
import random
from mpmath import mp

def ellipse_S(alpha, beta):
    a = mp.sin(alpha)
    b = mp.sin(beta)
    a2 = a*a
    b2 = b*b
    a2m = 1.0 - a2
    b2m = 1.0 - b2
    c = (b * a2m) / (a * mp.sqrt(b2m))
    m = (a2 - b2) / b2m
    n = m / a2
    return float(mp.pi/2 - c * mp.ellippi(n, m))

def ellipse_S_fit(p, A=0.988812, B=0.328265, C=1.64709):
    x, y = p
    return A * (x*x*y) * (C + B*math.cos(x*math.pi/2))


####################################################
# Precompute residual factor table
####################################################

N = 32
table = [[None] * N for _ in range(N)]
# Perform precomputation in high precision
with mp.workdps(100):
    for yi in range(N):
        for xi in range(N):
            # Apply coordinate mapping
            x = mp.sqrt(xi / (N-1))
            y = mp.sqrt(yi / (N-1))

            # Convert to angles
            alpha = max(1e-10, min(1.0 - 1e-10, x) * math.pi/2)
            beta = max(1e-10, y*alpha)

            # Convert back for fit evaluation to account for the clamping
            x = alpha * 2 / math.pi
            y = beta / alpha

            # Tabulate residual factor
            table[yi][xi] = ellipse_S(alpha, beta) / ellipse_S_fit((x, y))

# Print table in C syntax
print(f"static const float table_ellipse_S[{N}*{N}] = {{")
for row in table:
    print("  " + " ".join(f"{float(p):.8f}f," for p in row))
print("};")


####################################################
# Evaluate fit quality
####################################################

# Emulate single precision
mp.prec = 24
table = [[mp.mpf(v) for v in row] for row in table]

def lookupGrid(x, y):
    # Coordinate encoding
    x = (x * x) * (N-1)
    y = (y * y) * (N-1)

    xi = int(x)
    yi = int(y)
    xf = math.fmod(x, 1.0)
    yf = math.fmod(y, 1.0)

    def mix(a, b, t):
        return a + (b - a)*t

    low = mix(table[yi][xi], table[yi][xi+1], xf)
    high = mix(table[yi+1][xi], table[yi+1][xi+1], xf)
    return mix(low, high, yf)

random.seed(42)
diff = []
for _ in range(10000):
    # Pick random angles
    alpha = random.random() * mp.pi/2
    beta = random.random() * mp.pi/2

    # Per definition, alpha >= beta
    if beta > alpha:
        alpha, beta = beta, alpha

    # Map to parametrization
    x = alpha * 2 / mp.pi
    y = beta / alpha

    lookup = lookupGrid(x, y) * ellipse_S_fit((x, y))
    with mp.workdps(100):
        correct = ellipse_S(alpha, beta)
    diff.append((correct - lookup)/correct)
print(f"Avg: {sum(diff)/len(diff)}, Max: {max(diff)}")