#!/usr/bin/env python3
"""
E4-H160 — Range resolution in wavelength-steered optical arrays.

Three separable calculations:
  T1  geometry: illuminated bandwidth B_eff and the naive range floor vs aperture
  T2  information bound: Cramer-Rao bound on delay for a point and an extended target
  T3  Monte-Carlo check that an estimator attains the bound (FFT correlator)

No external data. All parameters are stated in the text and printed.

Run:  python3 h160_investigation_v2.py
Out:  ../results/*.csv
"""
import numpy as np
import csv, os

C = 2.99792458e8
OUT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "results"))
rng = np.random.default_rng(20261101)

PITCH = 2.0e-6
THETA0 = np.deg2rad(20.0)
LAM0 = PITCH * np.sin(THETA0)
COS0 = np.cos(THETA0)
DTH_DLAM = 1.0 / (PITCH * COS0)


def beam_halfwidth(N):
    return LAM0 / (N * PITCH * COS0)


def b_eff(N):
    return beam_halfwidth(N) * C / (LAM0 ** 2 * DTH_DLAM)


def envelope(f, N):
    sigma_theta = beam_halfwidth(N) / 2.0
    dtheta_df = LAM0 ** 2 * DTH_DLAM / C
    return np.exp(-0.5 * (f * dtheta_df / sigma_theta) ** 2)


def crb(A, f, sigma):
    J = (2 * np.pi) ** 2 * np.sum((f ** 2) * np.abs(A) ** 2) / (sigma ** 2)
    return 1.0 / np.sqrt(J) if J > 0 else np.inf


def corrupt(A, f, tau, snr_peak, n):
    p = np.max(np.abs(A) ** 2)
    sigma = np.sqrt(p / snr_peak)
    y = A * np.exp(-2j * np.pi * f * tau)
    return y + sigma * (rng.normal(size=n) + 1j * rng.normal(size=n)) / np.sqrt(2), sigma


def estimate(y, f, weight, pad=1 << 20):
    Y = np.fft.fftshift(np.fft.ifft(y * weight, n=pad))
    mag = np.abs(Y)
    k = int(np.argmax(mag))
    if 0 < k < pad - 1:
        a, b, c = mag[k - 1], mag[k], mag[k + 1]
        den = (a - 2 * b + c)
        delta = 0.5 * (a - c) / den if den != 0 else 0.0
    else:
        delta = 0.0
    df = f[1] - f[0]
    m = k - pad // 2 + delta        # fftshift moved the zero delay to pad//2
    return m / (pad * df)


def main():
    os.makedirs(OUT, exist_ok=True)
    rows = []
    print("=" * 78)
    print("E4-H160  independent investigation")
    print("=" * 78)
    print(f"pitch = {PITCH*1e6:.2f} um, theta0 = {np.rad2deg(THETA0):.1f} deg, "
          f"lambda0 = {LAM0*1e9:.1f} nm")

    print("\nT1  illuminated bandwidth and naive floor vs aperture")
    print(f"    {'N':>6s} {'D (mm)':>8s} {'B_eff (GHz)':>12s} {'floor (um)':>11s} {'floor/D':>8s}")
    for N in (256, 512, 1024, 2048, 4096, 8192):
        D = N * PITCH
        be = b_eff(N)
        floor = C / (2 * be)
        print(f"    {N:6d} {D*1e3:8.2f} {be/1e9:12.1f} {floor*1e6:11.2f} {floor/D:8.3f}")
        rows.append(dict(table="T1", N=N, D_m=D, B_eff_GHz=be/1e9, floor_um=floor*1e6))

    N = 2048
    n = 4096
    snr = 1e2          # 20 dB peak SNR: keeps the CRB above the delay grid
    print(f"\nT2  Cramer-Rao delay bound, N={N} (B_eff={b_eff(N)/1e9:.0f} GHz), peak SNR={snr:.0e}")
    print(f"    {'B (GHz)':>9s} {'B/B_eff':>8s} {'CRB point (ps)':>15s} {'CRB extended (ps)':>18s} {'point/ext':>10s}")
    for B in (10e9, 50e9, 100e9, 214e9, 500e9, 1000e9, 2000e9):
        df = B / n
        f = (np.arange(n) - n / 2) * df
        sigma = np.sqrt(1.0 / snr)
        c_pt = crb(envelope(f, N), f, sigma)
        c_ex = crb(np.ones_like(f), f, sigma)
        print(f"    {B/1e9:9.1f} {B/b_eff(N):8.2f} {c_pt*1e12:15.4f} {c_ex*1e12:18.4f} {c_pt/c_ex:10.2f}")
        rows.append(dict(table="T2", B_GHz=B/1e9, ratio=B/b_eff(N),
                         CRB_point_ps=c_pt*1e12, CRB_extended_ps=c_ex*1e12,
                         ratio_point_ext=c_pt/c_ex))

    print(f"\nT3  Monte-Carlo attainment (N={N}, 60 trials, peak SNR={snr:.0e})")
    print(f"    {'B (GHz)':>9s} {'estimator':>12s} {'RMSE (ps)':>11s} {'CRB (ps)':>10s} {'RMSE/CRB':>9s}")
    tau_true = 2 * 1.0 / C          # 1 m: far inside c/(2*df) at every B tested
    for B in (50e9, 100e9):  # B>100 GHz wraps on the delay grid; see RESULTS.md
        df = B / n
        f = (np.arange(n) - n / 2) * df
        A = envelope(f, N)
        sigma = np.sqrt(1.0 / snr)
        c_pt = crb(A, f, sigma)
        for name, w in (("matched", A), ("unweighted", np.ones_like(f))):
            errs = []
            for _ in range(120):
                y, _ = corrupt(A, f, tau_true, snr, n)
                t = estimate(y, f, w)
                errs.append(t - tau_true)
            rmse = float(np.sqrt(np.mean(np.square(errs))))
            print(f"    {B/1e9:9.1f} {name:>12s} {rmse*1e12:11.4f} {c_pt*1e12:10.4f} {rmse/c_pt:9.2f}")
            rows.append(dict(table="T3", B_GHz=B/1e9, estimator=name,
                             RMSE_ps=rmse*1e12, CRB_ps=c_pt*1e12, ratio=rmse/c_pt))

    path = os.path.join(OUT, "h160_results.csv")
    keys = sorted({k for r in rows for k in r})
    with open(path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=keys)
        w.writeheader()
        for r in rows:
            w.writerow(r)
    print(f"\nwrote {path}  ({len(rows)} rows)")


if __name__ == "__main__":
    main()