import json, time
import numpy as np


def fourier_weights(mask, bandwidth, clip=False):
    """Periodic Fourier projection of a sampled domain indicator."""
    n0, n1 = mask.shape
    coeff = np.fft.fft2(mask)
    k0 = np.fft.fftfreq(n0) * n0
    k1 = np.fft.fftfreq(n1) * n1
    keep = (np.abs(k0[:, None]) <= bandwidth[0]) & (np.abs(k1[None, :]) <= bandwidth[1])
    weights = np.fft.ifft2(coeff * keep).real
    if clip:
        weights = np.maximum(weights, 0.0)
    return weights


def domains(x, y):
    X, Y = np.meshgrid(x, y, indexing="ij")
    ellipse = ((X - .48) / .34) ** 2 + ((Y - .52) / .27) ** 2 <= 1
    hole = ellipse & (((X - .48) / .105) ** 2 + ((Y - .52) / .085) ** 2 > 1)
    # A cusp-like cardioid, plus a nonconvex boundary.
    xx, yy = X - .5, Y - .5
    r = np.sqrt(xx * xx + yy * yy)
    theta = np.arctan2(yy, xx)
    cusp = (r <= .27 * (1 + .55 * np.cos(theta)))
    polygon = (Y > .18 + .18 * X) & (Y < .84 - .26 * X) & (X > .16) & (X < .84)
    return {"ellipse": ellipse, "ellipse_hole": hole, "cusp": cusp, "polygon": polygon}


def fields(x, y):
    X, Y = np.meshgrid(x, y, indexing="ij")
    return np.stack([
        0.4 + 0.7 * X + 0.3 * Y,
        np.sin(2*np.pi*(1.3*X + .7*Y)) + .2*np.cos(2*np.pi*Y),
        np.exp(-35*((X-.3)**2 + (Y-.7)**2)),
        np.sin(2*np.pi*8*X) * np.cos(2*np.pi*5*Y),
    ])


def normalized_pool(f, w):
    return np.sum(f * w) / np.sum(w)


def run():
    np.random.seed(7)
    # Reference is much finer than the tested grid; all fields are evaluated analytically.
    nr = 768
    xr = (np.arange(nr) + .5) / nr
    ref_domains = domains(xr, xr)
    ref_fields = fields(xr, xr)
    references = {name: np.array([normalized_pool(f, m.astype(float)) for f in ref_fields])
                  for name, m in ref_domains.items()}

    rows = []
    for n in (24, 48, 96):
        x = (np.arange(n) + .5) / n
        ms = domains(x, x)
        fs = fields(x, x)
        for name, m in ms.items():
            m = m.astype(float)
            ref = references[name]
            base = np.array([normalized_pool(f, m) for f in fs])
            for frac in (.25, .5):
                w = fourier_weights(m, (int(frac*n), int(frac*n)), clip=False)
                fw = np.array([normalized_pool(f, w) for f in fs])
                wc = fourier_weights(m, (int(frac*n), int(frac*n)), clip=True)
                fc = np.array([normalized_pool(f, wc) for f in fs])
                rows.append({"n": n, "domain": name, "bandwidth_fraction": frac,
                             "masked_mean_mae": float(np.mean(np.abs(base-ref))),
                             "fourier_mae": float(np.mean(np.abs(fw-ref))),
                             "fourier_clipped_mae": float(np.mean(np.abs(fc-ref))),
                             "weight_min": float(w.min()), "weight_max": float(w.max()),
                             "mask_fraction": float(m.mean())})

    # Claimed reusable-weight property: pooling multiple channels is one matrix-weighted sum,
    # and geometry preprocessing is amortized over channels.
    n = 64; x = (np.arange(n)+.5)/n; m = domains(x,x)["cusp"].astype(float)
    t0 = time.perf_counter(); w = fourier_weights(m, (n//2,n//2)); prep = time.perf_counter()-t0
    channels = fields(x,x)
    t0 = time.perf_counter()
    for _ in range(300):
        _ = np.sum(channels*w, axis=(1,2))/np.sum(w)
    pool_ms = (time.perf_counter()-t0)*1000/300
    # Fourier projection's direct indicator approximation error (L2), for the representative cusp.
    projection_l2 = float(np.sqrt(np.mean((w-m)**2)))

    summary = {
        "rows": rows,
        "representative": {"preprocess_ms": prep*1000, "pool_4_channels_ms": pool_ms,
                            "projection_l2_n64_Fhalf": projection_l2},
        "mean_ratio_fourier_over_mask": float(np.mean([r["fourier_mae"]/max(r["masked_mean_mae"],1e-12) for r in rows])),
        "mean_ratio_clipped_over_mask": float(np.mean([r["fourier_clipped_mae"]/max(r["masked_mean_mae"],1e-12) for r in rows]))
    }
    with open("results.json", "w") as f: json.dump(summary, f, indent=2)
    print(json.dumps(summary, indent=2))

if __name__ == "__main__": run()