"""Object-based MIMO S21 corrections estimated after standard preprocessing.

Archives retain the standalone script's cal_*.npy/cal_meta.json layout. Pair
corrections multiply S21 after through calibration, reference subtraction and
notch filtering; they never alter S11 or the original measurement files.
"""

from __future__ import annotations

from contextlib import contextmanager
from collections.abc import Callable, Iterable, Sequence
from dataclasses import dataclass
import datetime
import json
import logging
from pathlib import Path
import re
import shutil
import tempfile
from typing import Any

import numpy as np
from numpy.typing import ArrayLike

from python_app.file_io import atomic_write_json
from python_app.models.dataset_model import SweepCollection

C_LIGHT = 299_792_458.0
_CAL_ARRAYS = ('freq_hz', 'pair', 'tx', 'rx', 'valid', 'tx_output_pos', 'rx_input_pos', 'antenna_params')
logger = logging.getLogger(__name__)


def _ports(values, label):
    array = np.asarray(values)
    if array.ndim != 1 or not array.size or array.dtype.kind not in 'iu':
        raise ValueError(f'{label} must contain distinct nonnegative integer ports')
    if np.any(array < 0) or np.any(array > np.iinfo(np.int32).max) or np.unique(array).size != array.size:
        raise ValueError(f'{label} must contain distinct nonnegative integer ports')
    return tuple(int(value) for value in array)


def _finite_array(value, label, shape=None, *, complex_values=False):
    array = np.asarray(value, dtype=complex if complex_values else float)
    if (shape is not None and array.shape != shape) or not np.all(np.isfinite(array)):
        raise ValueError(f'{label} has an invalid shape or non-finite values')
    return array


def _frequency(value, *, uniform=False):
    freq = _finite_array(value, 'frequency axis')
    if freq.ndim != 1 or freq.size < 2 or freq[0] < 0 or np.any(np.diff(freq) <= 0):
        raise ValueError('frequency axis must be increasing with at least two samples')
    # Dataset frequency axes are float32; account for their quantization at GHz.
    tolerance = max(1e-3, 2 * np.finfo(np.float32).eps * float(freq[-1]))
    if uniform and not np.allclose(np.diff(freq), np.median(np.diff(freq)), rtol=1e-5, atol=tolerance):
        raise ValueError('calibration requires a uniform frequency grid')
    return freq


@dataclass
class Settings:
    """Target geometry and estimation options, independent of example globals."""
    sphere_radius_m: float = .16
    input_idx: tuple[int, ...] = (0, 1, 2, 3)
    output_idx: tuple[int, ...] = (0, 1, 2, 3)
    antenna_height_m: float = .92
    pitch_deg: float = 10.
    f_min: float = 2e9
    f_max: float = 6e9
    gate_half_m: float = .15
    gate_taper_m: float = .03
    coarse_range_search_m: float = .6
    refine_axes: tuple[str, ...] = ('x', 'y', 'z')
    refine_range_m: float = .10
    apply_amplitude: bool = False
    phase_model: str = 'linear'

    def __post_init__(self):
        self.input_idx = tuple(sorted(_ports(self.input_idx, 'input_idx')))
        self.output_idx = tuple(sorted(_ports(self.output_idx, 'output_idx')))
        for name in ('sphere_radius_m', 'antenna_height_m', 'pitch_deg', 'f_min', 'f_max',
                     'gate_half_m', 'gate_taper_m', 'coarse_range_search_m', 'refine_range_m'):
            value = getattr(self, name)
            if isinstance(value, bool) or not np.isfinite(value):
                raise ValueError(f'{name} must be finite')
        if min(self.sphere_radius_m, self.antenna_height_m, self.coarse_range_search_m, self.refine_range_m, self.f_min) < 0:
            raise ValueError('radius, height, frequencies and search ranges must be nonnegative')
        if self.f_max <= self.f_min or min(self.gate_half_m, self.gate_taper_m) <= 0:
            raise ValueError('calibration band and gate must have positive widths')
        if self.phase_model not in ('const', 'linear', 'poly3', 'raw'):
            raise ValueError('phase_model must be const, linear, poly3 or raw')
        if len(set(self.refine_axes)) != len(self.refine_axes) or any(axis not in 'xyz' or len(axis) != 1 for axis in self.refine_axes):
            raise ValueError('refine_axes must contain distinct x/y/z axes')
        if not isinstance(self.apply_amplitude, bool):
            raise ValueError('apply_amplitude must be boolean')


def find_config_profile(start: str | Path, max_levels: int = 8) -> Path:
    """Locate the archived acquisition profile above a stage or snapshot folder."""
    path = Path(start)
    for _ in range(max_levels):
        candidate = path / 'config_profile.json'
        if candidate.is_file():
            return candidate
        path = path.parent
    raise FileNotFoundError(f'config_profile.json not found above {start}')


def _profile(path):
    config = json.loads(Path(path).read_text(encoding='utf-8'))
    if not isinstance(config, dict):
        raise ValueError('config_profile.json must contain an object')
    return config


def _geometry(config, input_idx, output_idx):
    result = []
    for key, identity, ports in (('tx_geometry', 'output_pos', output_idx), ('rx_geometry', 'input_pos', input_idx)):
        entries = config.get('gpr', {}).get(key, [])
        mapping = {}
        for entry in entries:
            port = _ports([entry[identity]], identity)[0]
            if port in mapping:
                raise ValueError(f'duplicate antenna geometry for {identity}={port}')
            mapping[port] = _finite_array([entry['x_m'], entry.get('y_m', 0.), entry.get('z_m', 0.)], key, (3,))
        missing = set(ports) - mapping.keys()
        if missing:
            raise ValueError(f'{key} is missing physical ports {sorted(missing)}')
        result.append(np.array([mapping[port] for port in ports]))
    return tuple(result)


def load_geometry(config_path: str | Path, input_idx: Sequence[int], output_idx: Sequence[int]) -> tuple[np.ndarray, np.ndarray]:
    """Read TX/RX radar coordinates in sorted physical port order."""
    return _geometry(_profile(config_path), sorted(_ports(input_idx, 'input_idx')), sorted(_ports(output_idx, 'output_idx')))


def list_snapshots(root: str | Path, prefixes: Sequence[str] | None = None) -> list[Path]:
    """Select complete snapshot directories, requiring unambiguous prefixes."""
    directories = sorted(path for path in Path(root).iterdir() if path.is_dir() and (path / 'meta.json').is_file())
    if prefixes is None:
        selected = directories
    else:
        selected = []
        for prefix in prefixes:
            matches = [path for path in directories if path.name.startswith(str(prefix))]
            if len(matches) != 1 or matches[0] in selected:
                raise ValueError(f'snapshot prefix {prefix!r} must select exactly one distinct snapshot in {root}')
            selected.extend(matches)
    if not selected:
        raise ValueError(f'no snapshots selected in {root}')
    return selected


def load_s21(snapshot_dir: str | Path, input_idx: Sequence[int], output_idx: Sequence[int]) -> tuple[np.ndarray, np.ndarray]:
    """Read a finite TX×RX×frequency matrix on one shared uniform grid."""
    inputs = sorted(_ports(input_idx, 'input_idx'))
    outputs = sorted(_ports(output_idx, 'output_idx'))
    rows, freq = [], None
    for output in outputs:
        row = []
        for input_ in inputs:
            suffix = f'i{input_}_o{output}'
            samples = np.load(Path(snapshot_dir) / f'{suffix}_s21.npy', allow_pickle=False)
            axis = _frequency(np.load(Path(snapshot_dir) / f'{suffix}_freq.npy', allow_pickle=False), uniform=True)
            if freq is not None and not np.array_equal(axis, freq):
                raise ValueError(f'frequency grids differ between pairs in {snapshot_dir}')
            freq = axis
            row.append(_finite_array(samples, suffix, freq.shape, complex_values=True))
        rows.append(row)
    assert freq is not None
    return np.array(rows), freq


def _source_signature(config) -> dict[str, Any]:
    radar = config.get('radar', {})
    s21 = config.get('preprocess', {}).get('s21', {})
    return {
        'radar': {key: radar[key] for key in ('model', 'serial') if key in radar},
        'slave_serials': radar.get('multi_device', {}).get('slave_serials', []),
        's21_calibration_set': s21.get('calibration', {}).get('set_name', ''),
        's21_reference_set': s21.get('reference', {}).get('set_name', ''),
        'notch': config.get('preprocess', {}).get('notch', {}),
    }


def validate_calibration(cal: dict[str, Any]) -> None:
    """Validate an archive before saving, applying or exporting runtime factors."""
    freq = _frequency(cal['freq_hz'], uniform=True)
    tx = _ports(cal['tx_output_pos'], 'tx_output_pos')
    rx = _ports(cal['rx_input_pos'], 'rx_input_pos')
    if tx != tuple(sorted(tx)) or rx != tuple(sorted(rx)):
        raise ValueError('calibration physical ports must be sorted')
    T = _finite_array(cal['tx'], 'tx corrections', (len(tx), freq.size), complex_values=True)
    R = _finite_array(cal['rx'], 'rx corrections', (len(rx), freq.size), complex_values=True)
    pair = _finite_array(cal['pair'], 'pair corrections', (len(tx), len(rx), freq.size), complex_values=True)
    if np.any(np.abs(T) == 0) or np.any(np.abs(R) == 0) or np.any(np.abs(pair) == 0):
        raise ValueError('calibration corrections must be nonzero')
    if not np.allclose(pair, T[:, None, :] * R[None, :, :], rtol=2e-5, atol=1e-7):
        raise ValueError('pair corrections disagree with tx/rx corrections')
    valid = np.asarray(cal['valid'])
    if valid.dtype.kind != 'b' or valid.shape != freq.shape or valid.sum() < 2:
        raise ValueError('calibration valid mask must select at least two frequencies')
    _finite_array(cal['antenna_params'], 'antenna parameters', (len(tx) + len(rx), 3))
    geometry = cal['meta']['geometry']
    _finite_array(geometry['tx_xyz_m'], 'TX geometry', (len(tx), 3))
    _finite_array(geometry['rx_xyz_m'], 'RX geometry', (len(rx), 3))


def validate_for_config(cal: dict[str, Any], config: dict[str, Any]) -> None:
    """Require physical geometry, combo coverage and a compatible phase base."""
    validate_calibration(cal)
    tx, rx = _geometry(config, cal['rx_input_pos'], cal['tx_output_pos'])
    geometry = cal['meta']['geometry']
    if not (np.allclose(tx, geometry['tx_xyz_m'], rtol=0, atol=1e-9) and
            np.allclose(rx, geometry['rx_xyz_m'], rtol=0, atol=1e-9)):
        raise ValueError('antenna geometry differs from the object calibration')
    for combo in config.get('run', {}).get('combos', []):
        input_, output = _ports([combo['input']], 'combo input')[0], _ports([combo['output']], 'combo output')[0]
        if input_ not in cal['rx_input_pos'] or output not in cal['tx_output_pos']:
            raise ValueError(f'object calibration does not cover combo {combo}')
    source = cal['meta'].get('source_preprocess', {})
    current = _source_signature(config)
    for key in ('radar', 'slave_serials', 's21_calibration_set'):
        if source.get(key) and current.get(key) and source[key] != current[key]:
            raise ValueError(f'object calibration source differs: {key}')


@contextmanager
def _staged_directory(destination, overwrite, *, preserve_existing=False):
    destination = Path(destination)
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.exists() and (not overwrite or not destination.is_dir() or destination.is_symlink()):
        raise FileExistsError(f'{destination} already exists; choose another name or overwrite=True')
    stage = Path(tempfile.mkdtemp(prefix=f'.{destination.name}.new-', dir=destination.parent))
    backup = None
    retain = False
    try:
        if preserve_existing and destination.exists():
            shutil.copytree(destination, stage, dirs_exist_ok=True, symlinks=True)
        yield stage
        if destination.exists():
            if not overwrite:
                raise FileExistsError(f'{destination} appeared during save; choose another name or overwrite=True')
            backup = Path(tempfile.mkdtemp(prefix=f'.{destination.name}.old-', dir=destination.parent))
            backup.rmdir()
            destination.replace(backup)
        try:
            stage.replace(destination)
        except BaseException:
            if backup is not None:
                try:
                    backup.replace(destination)
                except OSError as error:
                    retain = True
                    raise RuntimeError(f'calibration rollback failed; originals retained at {backup}, new files at {stage}') from error
            raise
    finally:
        if not retain:
            for temporary in (stage, backup):
                if temporary is not None and temporary.exists():
                    try:
                        shutil.rmtree(temporary)
                    except OSError:
                        logger.exception('Could not remove calibration temporary directory %s', temporary)


def save_calibration(cal: dict[str, Any], cal_dir: str | Path, overwrite: bool = False) -> Path:
    """Commit a complete legacy archive, restoring its predecessor on failure."""
    validate_calibration(cal)
    destination = Path(cal_dir)
    if destination.is_dir() and any(destination.glob('cal_*')) and not overwrite:
        raise FileExistsError(f'{destination} already contains calibration files; use overwrite=True')
    meta: dict[str, Any] = dict(cal['meta'], name=destination.name)
    # Older one-target archives used nonstandard JSON NaN for unavailable LOO.
    quality = dict(meta.get('quality', {}))
    if 'coherence_leave_one_out' in quality:
        quality['coherence_leave_one_out'] = [value if value is None or np.isfinite(value) else None for value in quality['coherence_leave_one_out']]
        meta['quality'] = quality
    with _staged_directory(destination, overwrite or destination.is_dir(), preserve_existing=True) as stage:
        for key in _CAL_ARRAYS:
            path = stage / f'cal_{key}.npy'
            path.unlink(missing_ok=True)
            np.save(path, cal[key], allow_pickle=False)
        (stage / 'cal_meta.json').unlink(missing_ok=True)
        atomic_write_json(stage / 'cal_meta.json', meta)
    cal['meta'], cal['dir'] = meta, destination
    return destination


def load_calibration(cal_dir: str | Path) -> dict[str, Any]:
    """Load and validate cal_*.npy arrays plus cal_meta.json without pickle."""
    directory = Path(cal_dir)
    cal: dict[str, Any] = {key: np.load(directory / f'cal_{key}.npy', allow_pickle=False) for key in _CAL_ARRAYS}
    cal.update(meta=json.loads((directory / 'cal_meta.json').read_text(encoding='utf-8')), dir=directory)
    validate_calibration(cal)
    return cal


def calibration_on_grid(cal: dict[str, Any], freq: ArrayLike) -> np.ndarray:
    """Interpolate separable corrections and continue phase smoothly at edges."""
    validate_calibration(cal)
    freq = _frequency(freq)
    fc = np.asarray(cal['freq_hz'])
    if np.array_equal(freq, fc):
        return cal['pair']

    def interp(correction, params):
        phase = np.unwrap(np.angle(correction))
        out = np.interp(freq, fc, phase)
        slope = 2 * np.pi * params[1] / C_LIGHT
        for selected, edge in ((freq < fc[0], 0), (freq > fc[-1], -1)):
            out[selected] = phase[edge] + slope * (freq[selected] - fc[edge])
        return np.interp(freq, fc, np.abs(correction)) * np.exp(1j * out)

    n_tx = len(cal['tx_output_pos'])
    tx = np.array([interp(value, cal['antenna_params'][index]) for index, value in enumerate(cal['tx'])])
    rx = np.array([interp(value, cal['antenna_params'][n_tx + index]) for index, value in enumerate(cal['rx'])])
    corrections = (tx[:, None, :] * rx[None, :, :]).astype(np.asarray(cal['pair']).dtype)
    _finite_array(corrections, 'interpolated corrections', complex_values=True)
    return corrections


def apply_calibration(s21: ArrayLike, freq: ArrayLike, cal: dict[str, Any]) -> np.ndarray:
    """Multiply a TX×RX×frequency S21 matrix by validated pair corrections."""
    corrections = calibration_on_grid(cal, freq)
    samples = _finite_array(s21, 'S21', corrections.shape, complex_values=True)
    return _finite_array(samples * corrections, 'corrected S21', corrections.shape, complex_values=True)


def apply_to_measurement(meas_dir: str | Path, cal: dict[str, Any], overwrite: bool = False) -> tuple[Path, int, int, int]:
    """Stage corrected S21 and copied companion files in postprocessed safely."""
    directory = Path(meas_dir)
    pre, post = directory / 'preprocessed', directory / 'postprocessed'
    if not pre.is_dir():
        raise FileNotFoundError(f'no {pre}')
    config = _profile(find_config_profile(pre))
    _reject_corrected_source(config)
    validate_for_config(cal, config)
    if post.exists() and overwrite and not (post / 'calibration_applied.json').is_file():
        raise FileExistsError(f'{post} has no calibration_applied.json; refusing to replace it')
    tx = {int(port): index for index, port in enumerate(cal['tx_output_pos'])}
    rx = {int(port): index for index, port in enumerate(cal['rx_input_pos'])}
    pattern = re.compile(r'^i(\d+)_o(\d+)_s21\.npy$')
    n_snap = n_fix = n_copy = 0
    cache = {}
    with _staged_directory(post, overwrite) as stage:
        for item in sorted(pre.iterdir()):
            if item.is_file():
                shutil.copy2(item, stage / item.name)
                n_copy += 1
                continue
            out = stage / item.name
            out.mkdir()
            n_snap += 1
            for path in sorted(item.iterdir()):
                match = pattern.match(path.name)
                if path.is_dir() or match is None:
                    if path.is_dir():
                        shutil.copytree(path, out / path.name)
                    else:
                        shutil.copy2(path, out / path.name)
                    n_copy += 1
                    continue
                input_, output = int(match[1]), int(match[2])
                if input_ not in rx or output not in tx:
                    raise ValueError(f'object calibration does not cover {path.name}')
                freq = _frequency(np.load(item / f'i{input_}_o{output}_freq.npy', allow_pickle=False))
                key = freq.tobytes()
                if key not in cache:
                    cache[key] = calibration_on_grid(cal, freq)
                samples = np.load(path, allow_pickle=False)
                dtype = samples.dtype
                samples = _finite_array(samples, path.name, freq.shape, complex_values=True)
                corrected = (samples * cache[key][tx[output], rx[input_]]).astype(dtype)
                _finite_array(corrected, 'corrected S21', complex_values=True)
                np.save(out / path.name, corrected, allow_pickle=False)
                n_fix += 1
        atomic_write_json(stage / 'calibration_applied.json', dict(
            calibration=cal['meta'].get('name'), calibration_dir=str(cal.get('dir') or ''),
            created=datetime.datetime.now().isoformat(timespec='seconds'), source=str(pre),
            rule='postprocessed S21 = preprocessed S21 * cal_pair[tx(output_pos), rx(input_pos)]',
            frequency_grids=[np.frombuffer(key, dtype=np.float64).tolist() for key in cache]))
    return post, n_snap, n_fix, n_copy


def _reject_corrected_source(config):
    object_calibration = config.get('preprocess', {}).get('s21', {}).get('object_calibration', {})
    if object_calibration.get('set_name') or object_calibration.get('bundle_path'):
        raise ValueError('calibration sources must have no object correction; object_calibration is already active')


def _best_phases(Z):
    """Maximize pair _coherence by alternating antenna phases, anchored at TX0/RX0."""
    t = np.zeros(Z.shape[0])
    r = np.zeros(Z.shape[1])
    for _ in range(100):
        t = np.angle((Z * np.exp(-1j * r)[None, :]).sum(axis=1))
        r = np.angle((Z * np.exp(-1j * t)[:, None]).sum(axis=0))
        t = t - t[0]
        r = r - r[0]
    return t, r


def _coherence(Z, t=None, r=None):
    """Return |sum(pair contributions)| / sum(|pair contributions|)."""
    if t is not None:
        if r is None:
            raise ValueError('RX phases must accompany TX phases')
        Z = Z * np.exp(-1j * (t[:, None] + r[None, :]))
    return float(np.abs(Z.sum()) / (np.abs(Z).sum() + 1e-30))


def _circular_median(phases):
    """Choose the phase minimizing circular distances along the first axis."""
    d = np.abs(np.angle(np.exp(1j * (phases[:, None] - phases[None, :]))))
    best = np.argmin(d.sum(axis=1), axis=0)
    return np.take_along_axis(phases, best[None], axis=0)[0]


def _combine(sets):
    """Combine targets using circular phase and logarithmic gain medians."""
    out = []
    for key in ('T', 'R'):
        Z = np.array([s[key] for s in sets])                      # (targets, antennas, frequencies)
        out.append(np.exp(np.median(np.log(np.abs(Z)), axis=0)) * np.exp(1j * _circular_median(np.angle(Z))))
    return tuple(out)


class _Estimator:
    """Geometry, gated target echoes and the original alternating estimator."""

    def __init__(self, st, tx, rx, freq, check=lambda: None):
        self.st, self.tx, self.rx, self.freq = st, tx, rx, freq
        self.check = check
        self.n_tx, self.n_rx = len(tx), len(rx)
        self.band = (freq >= st.f_min) & (freq <= st.f_max)
        self.fb = freq[self.band]
        if self.fb.size < 4:
            raise ValueError('calibration band must contain at least four frequency samples')
        self.win = np.hanning(self.fb.size)
        self.ok = self.win > 0.2                                  # Avoid unreliable window division at band edges.
        if self.ok.sum() < (4 if st.phase_model == 'poly3' else 2):
            raise ValueError('too few reliable frequencies for the requested phase model')
        self.nfft = 1 << int(np.ceil(np.log2(self.fb.size * 8)))
        self.dpath = np.fft.fftfreq(self.nfft, self.fb[1] - self.fb[0]) * C_LIGHT     # Excess propagation path in metres.
        self.gate = np.clip(np.minimum((self.dpath + st.gate_half_m) / st.gate_taper_m,
                                       (st.gate_half_m - self.dpath) / st.gate_taper_m) + 0.5, 0, 1)

    def scene_to_radar(self, q):
        """Convert ground-relative scene coordinates into radar coordinates."""
        p = np.radians(self.st.pitch_deg)
        y0 = q[1] - self.st.antenna_height_m
        return np.array([q[0], y0 * np.cos(p) + q[2] * np.sin(p), -y0 * np.sin(p) + q[2] * np.cos(p)])

    def pair_paths(self, q_scene):
        """Return bistatic paths to a point, or sphere surface paths shortened by
        2 * radius * cos(beta/2), using the measured sphere center."""
        a = self.st.sphere_radius_m
        pr = self.scene_to_radar(np.asarray(q_scene, float))
        vt = self.tx - pr
        vr = self.rx - pr
        dt = np.linalg.norm(vt, axis=1)
        dr = np.linalg.norm(vr, axis=1)
        if np.any(dt <= a) or np.any(dr <= a):
            raise ValueError('reflector intersects an antenna or has zero range')
        paths = dt[:, None] + dr[None, :]
        if a:
            cos_beta = (vt @ vr.T) / (dt[:, None] * dr[None, :])
            paths = paths - 2 * a * np.sqrt(np.clip((1 + cos_beta) / 2, 0, 1))
        return paths

    def _unwound(self, s21_band, paths):
        return s21_band * np.exp(2j * np.pi * self.fb[None, None, :] * paths[:, :, None] / C_LIGHT)

    def target_echo(self, s21_band, paths):
        """Isolate the unwound target echo with the tapered excess-path gate."""
        h = np.fft.ifft(self._unwound(s21_band, paths) * self.win, self.nfft, axis=2)
        return np.fft.fft(h * self.gate, axis=2)[:, :, :self.fb.size] / np.maximum(self.win, 1e-3)

    def broadband(self, E):
        """Sum reliable-band echo contributions at the target for each pair."""
        return E[:, :, self.ok].sum(axis=2)

    def coarse_range(self, s21, q_nominal):
        """Move nominal Z toward the strongest envelope within the path search range."""
        q = q_nominal.copy()
        if not self.st.coarse_range_search_m:
            return q
        s = s21[:, :, self.band]
        sel = np.abs(self.dpath) <= self.st.coarse_range_search_m
        for _ in range(3):
            self.check()
            env = np.abs(np.fft.ifft(self._unwound(s, self.pair_paths(q)) * self.win, self.nfft, axis=2)).sum(axis=(0, 1))
            q[2] += self.dpath[sel][np.argmax(env[sel])] / 2
        return q

    def refine_target(self, s21, q0):
        """Refine target coordinates and antenna phases by alternating Nelder-Mead fits."""
        P0 = self.pair_paths(q0)
        E0 = self.target_echo(s21[:, :, self.band], P0)
        axes = [('xyz').index(a) for a in self.st.refine_axes]

        def E_at(v):
            q = q0.copy()
            q[axes] += v
            dP = self.pair_paths(q) - P0
            return E0 * np.exp(2j * np.pi * self.fb[None, None, :] * dP[:, :, None] / C_LIGHT), q

        def score(v, t, r):
            self.check()
            if np.any(np.abs(v) > self.st.refine_range_m):
                return 0.0
            return _coherence(self.broadband(E_at(v)[0]), t, r)

        v = np.zeros(len(axes))
        if axes:
            from scipy.optimize import minimize

            for _ in range(6):
                self.check()
                t, r = _best_phases(self.broadband(E_at(v)[0]))
                simplex = np.vstack([v] + [v + 0.03 * np.eye(len(axes))[k] for k in range(len(axes))])
                v = minimize(lambda w: -score(w, t, r), v, method='Nelder-Mead',
                             options=dict(xatol=5e-4, fatol=1e-7, initial_simplex=simplex)).x
        return E_at(v)

    def factorize(self, E):
        """Factor E(tx,rx,f) into antenna phases and least-squares log-amplitudes."""
        nb = E.shape[2]
        T = np.ones((self.n_tx, nb), complex)
        R = np.ones((self.n_rx, nb), complex)
        L = np.log(np.abs(E) + 1e-30)
        for k in range(nb):
            self.check()
            t, r = _best_phases(E[:, :, k])
            tt = L[:, :, k].mean(axis=1) - L[:, :, k].mean()
            rr = L[:, :, k].mean(axis=0) - L[:, :, k].mean()
            T[:, k] = np.exp(tt + 1j * t)
            R[:, k] = np.exp(rr + 1j * r)
        return T, R

    def antenna_params(self, Z):
        """Fit zero-frequency phase, excess delay path and gain for one antenna."""
        ph = np.unwrap(np.angle(Z[self.ok]))
        f = self.fb[self.ok]
        slope, ph0 = np.polyfit(f, ph, 1)
        return dict(phase0_deg=float(np.degrees(ph0)), delay_path_m=float(-slope * C_LIGHT / (2 * np.pi)),
                    gain_db=float(20 * np.log10(np.mean(np.abs(Z[self.ok])))),
                    phase_rms_after_linear_deg=float(np.degrees(np.std(ph - (slope * f + ph0)))))

    def phase_model(self, Z):
        """Fit the selected antenna phase model over reliable band samples."""
        model = self.st.phase_model
        if model == 'raw':
            return Z
        ph = np.unwrap(np.angle(Z))
        f = (self.fb - self.fb.mean()) / 1e9
        if model == 'const':
            u = np.exp(1j * np.angle(Z[self.ok])).mean()
            return np.abs(Z) * np.exp(1j * np.angle(u))
        deg = {'linear': 1, 'poly3': 3}[model]
        return np.abs(Z) * np.exp(1j * np.polyval(np.polyfit(f[self.ok], ph[self.ok], deg), f))

    def pair_correction(self, T, R):
        """Invert antenna factors, retaining measured amplitudes in phase-only mode."""
        Z = T[:, None, :] * R[None, :, :]
        return np.conj(Z) / np.abs(Z) if not self.st.apply_amplitude else 1.0 / Z

    def to_full_grid(self, Z_band, params):
        """Extend antenna factors over the full grid with continuous edge-anchored
        linear phase and constant edge gain outside the reliable band."""
        idx = np.flatnonzero(self.band)[self.ok]
        ph_in = np.unwrap(np.angle(Z_band[self.ok]))
        mag_in = np.abs(Z_band[self.ok])
        slope = -2 * np.pi * params['delay_path_m'] / C_LIGHT               # Radians per Hz.
        ph = np.empty(self.freq.size)
        mag = np.empty(self.freq.size)
        ph[idx] = ph_in
        mag[idx] = mag_in
        lo, hi = idx[0], idx[-1]
        ph[:lo] = ph_in[0] + slope * (self.freq[:lo] - self.freq[lo])
        mag[:lo] = mag_in[0]
        ph[hi + 1:] = ph_in[-1] + slope * (self.freq[hi + 1:] - self.freq[hi])
        mag[hi + 1:] = mag_in[-1]
        return mag * np.exp(1j * ph)


def _callbacks(progress, cancelled):
    def check(message=None):
        if cancelled is not None and cancelled():
            raise InterruptedError('object calibration cancelled')
        if message is not None and progress is not None:
            progress(message)
    return check


def _mean_samples(samples, check, expected_freq=None):
    mean, count, freq = None, 0, expected_freq
    for data, axis in samples:
        check()
        axis = _frequency(axis, uniform=True)
        if freq is not None and not np.array_equal(axis, freq):
            raise ValueError('frequency grids differ between calibration measurements')
        freq = axis
        if mean is None:
            mean = np.zeros_like(data, dtype=np.complex128)
        data = _finite_array(data, 'calibration measurement', mean.shape, complex_values=True)
        count += 1
        mean += (data - mean) / count
    if mean is None:
        raise ValueError('calibration measurements must contain at least one collection')
    return mean, freq


def _settings(settings):
    settings = settings or Settings()
    settings.__post_init__()
    return settings


def _compute(sets, empty, freq, settings, config, config_path, check) -> dict[str, Any]:
    tx, rx = _geometry(config, settings.input_idx, settings.output_idx)
    est = _Estimator(settings, tx, rx, freq, check)
    for index, record in enumerate(sets):
        check(f'Estimating target {index + 1}/{len(sets)}')
        data = record.pop('s21')
        if empty is not None:
            data = data - empty
        nominal = _finite_array(record['target_scene_m'], 'target_scene_m', (3,))
        echo, refined = est.refine_target(data, est.coarse_range(data, nominal))
        if not np.all(np.any(np.abs(echo[:, :, est.ok]) > 0, axis=2)):
            raise ValueError(f'target {index + 1} has no usable echo for one or more antenna pairs')
        _finite_array(echo, 'target echo', complex_values=True)
        record.update(E=echo, q=refined, q_nominal=nominal)
        record['T'], record['R'] = est.factorize(echo)

    T_all, R_all = _combine(sets)
    T_all = np.array([est.phase_model(value) for value in T_all])
    R_all = np.array([est.phase_model(value) for value in R_all])
    antennas = {f'Tx{index}': est.antenna_params(value) for index, value in enumerate(T_all)}
    antennas.update({f'Rx{index}': est.antenna_params(value) for index, value in enumerate(R_all)})
    T_full = np.array([est.to_full_grid(value, antennas[f'Tx{index}']) for index, value in enumerate(T_all)])
    R_full = np.array([est.to_full_grid(value, antennas[f'Rx{index}']) for index, value in enumerate(R_all)])
    unit = (lambda value: np.conj(value) / np.abs(value)) if not settings.apply_amplitude else (lambda value: 1 / value)
    saved_pair = unit(T_full[:, None, :] * R_full[None, :, :]).astype(np.complex64)
    quality = dict(coherence_before=[], coherence_own=[], coherence_leave_one_out=[], coherence_final=[])
    diagnostics = []
    final_correction = saved_pair[:, :, est.band]
    for record in sets:
        check()
        before = est.broadband(record['E'])
        after = est.broadband(record['E'] * final_correction)
        quality['coherence_before'].append(_coherence(before))
        quality['coherence_own'].append(_coherence(est.broadband(record['E'] * est.pair_correction(record['T'], record['R']))))
        others = [other for other in sets if other is not record]
        quality['coherence_leave_one_out'].append(
            _coherence(est.broadband(record['E'] * est.pair_correction(*_combine(others)))) if others else None)
        quality['coherence_final'].append(_coherence(after))
        diagnostics.append(dict(phase_before_deg=np.degrees(np.angle(before)).tolist(),
                                phase_after_deg=np.degrees(np.angle(after)).tolist(),
                                coherence_before=_coherence(before), coherence_after=_coherence(after)))

    valid = np.zeros(freq.size, bool)
    valid[np.flatnonzero(est.band)[est.ok]] = True
    meta = dict(
        created=datetime.datetime.now().isoformat(timespec='seconds'), name=None,
        convention=dict(apply='S21_corrected[tx, rx, f] = S21[tx, rx, f] * cal_pair[tx, rx, f]',
                        index_order='sorted physical tx_output_pos x rx_input_pos',
                        phase_reference='relative to Tx0 and Rx0; common phase/delay is retained',
                        amplitude='equalized' if settings.apply_amplitude else 'phase only',
                        outside_band='linear antenna phase continuation anchored at the band edge'),
        geometry=dict(config=str(config_path), tx_xyz_m=tx.tolist(), rx_xyz_m=rx.tolist(),
                      antenna_height_m=settings.antenna_height_m, pitch_deg=settings.pitch_deg),
        estimation=dict(band_hz=[settings.f_min, settings.f_max], gate_half_m=settings.gate_half_m,
                        gate_taper_m=settings.gate_taper_m, coarse_range_search_m=settings.coarse_range_search_m,
                        refine_axes=list(settings.refine_axes), refine_range_m=settings.refine_range_m,
                        phase_model=settings.phase_model, sphere_radius_m=settings.sphere_radius_m),
        source_preprocess=_source_signature(config),
        targets=[dict(path=record.get('path', ''), snapshots=record.get('snapshots', []),
                      nominal_scene_m=record['q_nominal'].tolist(), refined_scene_m=record['q'].tolist()) for record in sets],
        antennas=antennas, quality=quality, diagnostics=diagnostics)
    cal = dict(freq_hz=freq.astype(np.float64), pair=saved_pair,
               tx=unit(T_full).astype(np.complex64), rx=unit(R_full).astype(np.complex64), valid=valid,
               tx_output_pos=np.array(settings.output_idx, np.int32), rx_input_pos=np.array(settings.input_idx, np.int32),
               antenna_params=np.array([[np.radians(value['phase0_deg']), value['delay_path_m'], value['gain_db']]
                                        for value in antennas.values()]), meta=meta, dir=None)
    validate_calibration(cal)
    check('Object calibration complete')
    return cal


def _collection_samples(collection, settings):
    from python_app.storage.npz.serialize import validate_trace_collection

    validate_trace_collection(collection, require_finite=True)
    traces = {(trace.combo.input, trace.combo.output): trace for trace in collection.traces}
    rows, freq = [], None
    for output in settings.output_idx:
        row = []
        for input_ in settings.input_idx:
            if (input_, output) not in traces:
                raise ValueError(f'calibration collection is missing input={input_}, output={output}')
            trace = traces[input_, output]
            axis = _frequency(trace.frequency_hz, uniform=True)
            if freq is not None and not np.array_equal(axis, freq):
                raise ValueError('frequency grids differ between collection pairs')
            freq = axis
            row.append(_finite_array(trace.s21, 'S21', freq.shape, complex_values=True))
        rows.append(row)
    return np.array(rows), freq


def compute_from_collections(
    cal_sets: Iterable[dict[str, Any]], config_dict: dict[str, Any],
    empty_collections: Iterable[SweepCollection] | None = None, settings: Settings | None = None,
    *, progress: Callable[[str], None] | None = None, cancelled: Callable[[], bool] | None = None,
) -> dict[str, Any]:
    """Estimate corrections from standard-preprocessed sweeps without object correction.

    cal_sets contains {collections: [SweepCollection, ...], target_scene_m: XYZ}.
    progress accepts a status string; cancelled returns bool and raises
    InterruptedError promptly in data, refinement and factorization loops.
    """
    settings = _settings(settings)
    _reject_corrected_source(config_dict)
    check = _callbacks(progress, cancelled)
    sets = list(cal_sets)
    if not sets:
        raise ValueError('at least one reflector target is required')
    empty, freq = None, None
    if empty_collections is not None:
        check('Averaging empty scene')
        empty, freq = _mean_samples((_collection_samples(value, settings) for value in empty_collections), check)
    prepared = []
    for index, record in enumerate(sets):
        check(f'Averaging target {index + 1}/{len(sets)}')
        collections = list(record['collections'])
        mean, freq = _mean_samples((_collection_samples(value, settings) for value in collections), check, freq)
        prepared.append(dict(s21=mean, target_scene_m=record['target_scene_m'],
                             snapshots=[str(value.collection_id) for value in collections]))
    return _compute(prepared, empty, freq, settings, config_dict, 'live capture', check)


def compute_calibration(
    cal_sets: Iterable[dict[str, Any]], empty_scene_path: str | Path | None = None, settings: Settings | None = None,
    *, progress: Callable[[str], None] | None = None, cancelled: Callable[[], bool] | None = None,
) -> dict[str, Any]:
    """Legacy archive API backed by the same estimator as live collection capture."""
    settings = _settings(settings)
    check = _callbacks(progress, cancelled)
    cal_sets = list(cal_sets)
    if not cal_sets:
        raise ValueError('at least one reflector target is required')
    config_path = find_config_profile(cal_sets[0]['path'])
    config = _profile(config_path)
    _reject_corrected_source(config)
    geometry = _geometry(config, settings.input_idx, settings.output_idx)
    signature = _source_signature(config)

    def samples(directories):
        for directory in directories:
            check()
            profile_path = find_config_profile(directory)
            path = Path(directory)
            while True:
                if (path / 'calibration_applied.json').is_file():
                    raise ValueError(f'object calibration is already applied to archive {directory}')
                if path == profile_path.parent:
                    break
                path = path.parent
            source = _profile(profile_path)
            _reject_corrected_source(source)
            current = _geometry(source, settings.input_idx, settings.output_idx)
            if any(not np.allclose(first, second, rtol=0, atol=1e-9) for first, second in zip(geometry, current)):
                raise ValueError(f'antenna geometry differs in {directory}')
            if signature != _source_signature(source):
                raise ValueError(f'preprocessing or hardware configuration differs in {directory}')
            yield load_s21(directory, settings.input_idx, settings.output_idx)

    empty, freq = None, None
    if empty_scene_path is not None:
        empty, freq = _mean_samples(samples(list_snapshots(empty_scene_path)), check)
    prepared = []
    for record in cal_sets:
        directories = list_snapshots(record['path'], record.get('snapshots'))
        mean, freq = _mean_samples(samples(directories), check, freq)
        prepared.append(dict(s21=mean, target_scene_m=record['target_scene_m'], path=str(record['path']),
                             snapshots=[directory.name for directory in directories]))
    cal = _compute(prepared, empty, freq, settings, config, config_path, check)
    cal['meta']['estimation']['empty_scene'] = str(empty_scene_path) if empty_scene_path is not None else None
    return cal