"""Estimate object-based MIMO S21 factors after through/reference preprocessing."""

from __future__ import annotations

from collections.abc import Callable, Iterable
from dataclasses import asdict, dataclass
import copy
import datetime
from typing import Any

import numpy as np

from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.storage.npz.serialize import validate_trace_collection
from python_app.workflows.object_calibration import geometry_signature

C_LIGHT = 299_792_458.0


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:
    """Reflector geometry and phase-estimation settings."""
    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 _geometry(config, input_idx, output_idx):
    geometry, result = geometry_signature(config), []
    for key, identity, ports in (('tx_geometry', 'output_pos', output_idx), ('rx_geometry', 'input_pos', input_idx)):
        mapping = {entry[identity]: [entry['x_m'], entry['y_m'], entry['z_m']] for entry in geometry[key]}
        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 _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 _compute(sets, freq, settings, config, provenance, check) -> tuple[SweepCollection, 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')
        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)])
    saved_pair = est.pair_correction(T_full, R_full).astype(np.complex64)
    _finite_array(saved_pair, 'pair corrections', complex_values=True)
    diagnostics = []
    final_correction = saved_pair[:, :, est.band]
    for record in sets:
        check()
        before = est.broadband(record['E'])
        after = est.broadband(record['E'] * final_correction)
        # Notebook cell 5: each map uses its own summed contribution as phase zero.
        phase_before = (np.degrees(np.angle(before) - np.angle(before.sum())) + 180) % 360 - 180
        phase_after = (np.degrees(np.angle(after) - np.angle(after.sum())) + 180) % 360 - 180
        diagnostics.append(dict(phase_before_deg=phase_before.tolist(),
                                phase_after_deg=phase_after.tolist(),
                                phase_std_before_deg=float(np.std(phase_before)),
                                phase_std_after_deg=float(np.std(phase_after)),
                                coherence_before=_coherence(before), coherence_after=_coherence(after)))

    metadata = dict(
        created=datetime.datetime.now().isoformat(timespec='seconds'),
        provenance=copy.deepcopy(provenance),
        tx_output_pos=list(settings.output_idx), rx_input_pos=list(settings.input_idx),
        geometry=dict(tx_xyz_m=tx.tolist(), rx_xyz_m=rx.tolist()),
        settings=asdict(settings),
        source_preprocess=dict(reference_name=config.get('preprocess', {}).get('s21', {}).get('reference', {}).get('set_name', ''),
                               notch=copy.deepcopy(config.get('preprocess', {}).get('notch', {}))),
        targets=[dict(captures=record['captures'], nominal_scene_m=record['q_nominal'].tolist(),
                      refined_scene_m=record['q'].tolist()) for record in sets],
        antennas=antennas, diagnostics=diagnostics,
    )
    frequency = freq.astype(np.float32)
    collection = SweepCollection(0, 0, [
        TraceData(ComboKey(input_, output), frequency, saved_pair[t, r])
        for t, output in enumerate(settings.output_idx) for r, input_ in enumerate(settings.input_idx)
    ])
    validate_trace_collection(collection, require_finite=True)
    check('Object calibration complete')
    return collection, metadata


def _collection_samples(collection, settings):
    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], settings: Settings | None = None,
    *, provenance: dict[str, Any], progress: Callable[[str], None] | None = None,
    cancelled: Callable[[], bool] | None = None,
) -> tuple[SweepCollection, dict[str, Any]]:
    """Estimate a saved set from preprocessed sweeps without object correction.

    Each target provides ``collections`` and ground-relative ``target_scene_m``.
    Reference subtraction has already been applied by the live preprocessor.
    Cancellation is checked during averaging, refinement and factorization.
    """
    settings = Settings() if settings is None else settings
    check = _callbacks(progress, cancelled)
    sets = list(cal_sets)
    if not sets:
        raise ValueError('at least one reflector target is required')
    prepared, freq = [], None
    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'], captures=[
            dict(collection_id=value.collection_id, monotonic_ns=value.monotonic_ns,
                 capture_start_ns=value.capture_start_ns, capture_end_ns=value.capture_end_ns)
            for value in collections
        ]))
    collection, metadata = _compute(prepared, freq, settings, config_dict, provenance, check)
    # Multi-device acquisition keeps every stage for timing. Pairs outside the
    # configured geometry carry identity factors, explicitly marked in the set.
    corrected = {trace.combo for trace in collection.traces}
    unchanged = [trace for trace in collections[0].traces if trace.combo not in corrected]
    collection.traces.extend(TraceData(
        trace.combo, np.asarray(trace.frequency_hz, dtype=np.float32),
        np.ones(trace.frequency_hz.size, dtype=np.complex64),
    ) for trace in unchanged)
    metadata['unchanged_combos'] = [dict(input=trace.combo.input, output=trace.combo.output) for trace in unchanged]
    return collection, metadata