"""Notebook-formula reference for the C++ `gpr_timing_check` tool.

Recomputes the same synthetic 4x4 scenario with the Multi_antenna_3D notebook
formulas (compute_pair_timestamps_for_snapshot + compute_frequency_sample_times
+ apply_intra_sweep_phase_correction, 'segmented_interleaved' branch) and
prints the motion phase at the same probe indices in the same format. Run:

    make gpr_timing_check && build/bin/gpr_timing_check > /tmp/cpp.txt
    python python_app/scripts/gpr_timing_check_reference.py > /tmp/ref.txt

The `int_minus` / `int_focus` lines of both outputs must match to ~1e-9 rad
(phases are compared wrapped to (-pi, pi]).
"""

from __future__ import annotations

import numpy as np

SPEED_OF_LIGHT = 299_792_458.0  # C++ kSpeedOfLightMetersPerSec
SPEED_M_S = 1.42
LOOK_ANGLE_DEG = 19.5
DIRECTION_SIGN = 1.0
F_START = 1.5e9
F_STOP = 6.0e9
N_FREQ = 11

# Trace windows of the synthetic collection (absolute seconds): segment A is
# Tx0/Tx1, segment B is Tx2/Tx3 (260 ms, stretched).
WINDOWS = {0: (1.000, 1.223), 1: (1.000, 1.223), 2: (1.223, 1.483), 3: (1.223, 1.483)}
TX_ORDER = [0, 1, 2, 3]
PROBE_INDICES = (0, 5, 10)


def wrap(phase: np.ndarray) -> np.ndarray:
    """Wrap phases to (-pi, pi] the way -arg(exp(-1j*phi)) reports them."""
    return -np.angle(np.exp(-1j * phase))


def main() -> None:
    freqs = np.linspace(F_START, F_STOP, N_FREQ)
    first = min(start for start, _ in WINDOWS.values())
    last = max(end for _, end in WINDOWS.values())
    acq_s = last - first
    t_ref = 0.5 * acq_s

    # Group Tx with identical windows into segments (snapshot_acquisition_timing
    # + compute_pair_timestamps_for_snapshot).
    groups: dict[tuple[float, float], list[int]] = {}
    for tx in TX_ORDER:
        groups.setdefault(WINDOWS[tx], []).append(tx)
    seg_of = {}
    for window in sorted(groups):
        for slot, tx in enumerate(groups[window]):
            seg_of[tx] = (window, slot, len(groups[window]))

    cos_theta = np.cos(np.radians(LOOK_ANGLE_DEG))
    motion_factor = DIRECTION_SIGN * SPEED_M_S * cos_theta

    print(f"reference: frame_center_s={t_ref:.9f} acq_s={acq_s:.9f}")
    phi_by_tx = {}
    for tx in TX_ORDER:
        (t0_abs, t1_abs), slot, n_slots = seg_of[tx]
        seg_start = t0_abs - first
        seg_end = t1_abs - first
        t_seg = seg_end - seg_start
        t_center = 0.5 * (seg_start + seg_end)

        # compute_frequency_sample_times, segmented branch.
        dt_event = t_seg / (n_slots * N_FREQ)
        t_abs = seg_start + (n_slots * np.arange(N_FREQ) + slot) * dt_event

        dt_intra = t_abs - t_center
        dtau_intra = 2.0 * motion_factor * dt_intra / SPEED_OF_LIGHT
        dt_pair = t_center - t_ref
        dtau_pair = 2.0 * motion_factor * dt_pair / SPEED_OF_LIGHT
        phi = 2.0 * np.pi * freqs * (dtau_intra + dtau_pair)
        phi_by_tx[tx] = phi

        wrapped = wrap(phi)
        print(
            f"int_minus tx o{tx}: "
            + " ".join(f"phi(k={k})={wrapped[k]:.12e}" for k in PROBE_INDICES)
        )

    for tx in TX_ORDER:
        phi = phi_by_tx[tx]
        # int_focus: remove the constant + linear-in-frequency component
        # (least squares on [1, f_rel], notebook intra_phase_focus_only).
        f_rel = freqs - freqs.mean()
        design = np.column_stack([np.ones_like(f_rel), f_rel])
        beta, *_ = np.linalg.lstsq(design, phi, rcond=None)
        residual = phi - design @ beta
        wrapped = wrap(residual)
        print(
            f"int_focus tx o{tx}: "
            + " ".join(f"phi(k={k})={wrapped[k]:.12e}" for k in PROBE_INDICES)
        )


if __name__ == "__main__":
    main()