import json
import numpy as np
from cross_ratio_lattice import *


def stats(x):
    x = np.asarray(x, float)
    return {'max': float(np.nanmax(x)), 'median': float(np.nanmedian(x)), 'mean': float(np.nanmean(x))}

rng = np.random.default_rng(123)
# Direct formula and inverse formula on independent nonsingular triples.
a = rng.normal(size=10000) + 1j*rng.normal(size=10000)
b = rng.normal(size=10000) + 1j*rng.normal(size=10000)
c = rng.normal(size=10000) + 1j*rng.normal(size=10000)
d, bad = complete_d(a, b, c)
finite = np.isfinite(d)
formula_res = np.abs(cross_ratio(a[finite], b[finite], c[finite], d[finite]) + 1)
c2, bad2 = complete_c(a, b, d)
reverse_err = np.abs(c2[finite] - c[finite])
results = {'formula': {'singular': int(bad.sum()), 'residual': stats(formula_res),
                       'inverse_abs_error': stats(reverse_err)}}

for dtype in (np.complex128, np.complex64):
    z, singular = generate_lattice(18, seed=8, dtype=dtype)
    # Cast noise to the lattice precision: do not accidentally promote complex64.
    real_dtype = np.float32 if dtype == np.complex64 else np.float64
    noise = (rng.normal(size=z.shape[0]).astype(real_dtype) +
             1j * rng.normal(size=z.shape[0]).astype(real_dtype)).astype(dtype)
    amp = dtype.type(1e-4) if hasattr(dtype, 'type') else 1e-4
    top = (z[0, :] + amp * noise).astype(dtype)
    left = (z[:, 0] + amp * noise).astype(dtype)
    left[0] = top[0]
    constrained = constrained_from_boundary(top, left)
    additive = additive_from_boundary(top, left)
    results[str(dtype)] = {
        'generation_singular': singular,
        'exact_residual': stats(plaquette_residuals(z)),
        'constrained_perturbed_residual': stats(plaquette_residuals(constrained)),
        'additive_control_residual': stats(plaquette_residuals(additive)),
        'constrained_finite': bool(np.isfinite(constrained).all())}

# Completion followed by inverse completion measures the claimed reversibility.
for dtype in (np.complex128, np.complex64):
    rd = np.float32 if dtype == np.complex64 else np.float64
    aa = (rng.normal(size=5000).astype(rd) + 1j*rng.normal(size=5000).astype(rd)).astype(dtype)
    bb = (rng.normal(size=5000).astype(rd) + 1j*rng.normal(size=5000).astype(rd)).astype(dtype)
    cc = (rng.normal(size=5000).astype(rd) + 1j*rng.normal(size=5000).astype(rd)).astype(dtype)
    dd, bad = complete_d(aa, bb, cc)
    cb, _ = complete_c(aa, bb, dd)
    ok = np.isfinite(dd)
    results['roundtrip_' + str(dtype)] = {
        'finite': int(ok.sum()),
        'relative_error': float(np.linalg.norm((cb-cc)[ok]) / np.linalg.norm(cc[ok]))}

with open('results.json', 'w') as f:
    json.dump(results, f, indent=2)
print(json.dumps(results, indent=2))