"""Time-domain postprocessing for CloudScat impulse responses (NumPy only).

Units and conventions
---------------------
``dt`` is in seconds; ``h`` is a causal sampled response density, so its mass
is ``h * dt``.  Array element n is assigned delay n*dt.  Supply *excess arrival
delay*, after subtracting the direct source-to-observer travel time.  If a
histogram uses bin centres, account for that origin before calling this module.
No missing/truncated part of h is reconstructed.  Timing quantiles are
conditional on the response supplied; integrated response weight is retained.

The source is exp(-t/3 ms)*(1-exp(-t/tau_r))*H(t), normalized to unit peak.
tau_r is calibrated to the requested rising-edge 10--90% time, relative to
the source's peak.  This is deliberately distinct from onset-to-peak time.
The optional electronic response is an ideal first-order low-pass, not a
complete antialias filter. Its discrete kernel is exact for a linearly
interpolated input, preserving grid times without a half-bin leading bias.
"""

from functools import lru_cache
import json
import math
import numpy as np


SOURCE_DECAY_S = 3.0e-3
SOURCE_CALIBRATION_DT_S = 0.2e-6


def _validate_dt(dt):
    dt = float(dt)
    if not np.isfinite(dt) or dt <= 0:
        raise ValueError("dt must be finite and positive, in seconds")
    return dt


def _response(h, dt):
    dt = _validate_dt(dt)
    h = np.asarray(h, dtype=float)
    if h.ndim != 1 or h.size == 0 or not np.all(np.isfinite(h)):
        raise ValueError("h must be a nonempty, finite one-dimensional array")
    scale = float(np.max(np.abs(h)))
    if np.any(h < -max(scale * 1e-12, 1e-300)):
        raise ValueError("h must be nonnegative")
    return np.maximum(h, 0.0), dt


def _crossing(t, y, level, stop_index):
    """Earliest rising crossing before the selected peak, interpolated."""
    indexes = np.flatnonzero(y[:stop_index + 1] >= level)
    if indexes.size == 0:
        return None
    j = int(indexes[0])
    if j == 0:
        # Source and electronic outputs start at zero; this branch also
        # handles an impulse/step waveform whose first point is nonzero.
        return float(t[0])
    y0, y1 = float(y[j - 1]), float(y[j])
    if y1 == y0:
        return float(t[j])
    return float(t[j - 1] + (level - y0) / (y1 - y0) * (t[j] - t[j - 1]))


def waveform_metrics(t_s, y, onset_s=0.0):
    """Return rising-edge metrics; onset is supplied, never noise-thresholded."""
    t_s = np.asarray(t_s, dtype=float)
    y = np.asarray(y, dtype=float)
    if t_s.ndim != 1 or t_s.size < 2 or y.shape != t_s.shape:
        raise ValueError("t_s and y must be matching one-dimensional arrays")
    if not np.all(np.isfinite(y)) or not np.all(np.isfinite(t_s)):
        raise ValueError("waveform values must be finite")
    if not np.all(np.diff(t_s) > 0):
        raise ValueError("waveform times must be strictly increasing")
    peak_index = int(np.argmax(y))
    peak = float(y[peak_index])
    if peak <= 0:
        return {"rise_10_90_us": None, "onset_to_peak_us": None,
                "t10_us": None, "t90_us": None, "peak_time_us": None,
                "peak_amplitude": 0.0, "peak_at_window_end": False}
    t10 = _crossing(t_s, y, 0.1 * peak, peak_index)
    t90 = _crossing(t_s, y, 0.9 * peak, peak_index)
    return {
        "rise_10_90_us": None if t10 is None or t90 is None else (t90 - t10) * 1e6,
        "onset_to_peak_us": float((t_s[peak_index] - onset_s) * 1e6),
        "t10_us": None if t10 is None else (t10 - onset_s) * 1e6,
        "t90_us": None if t90 is None else (t90 - onset_s) * 1e6,
        "peak_time_us": float((t_s[peak_index] - onset_s) * 1e6),
        "peak_amplitude": peak,
        "peak_at_window_end": bool(peak_index == y.size - 1),
    }


def _source_peak(tau_r_s, decay_s):
    t_peak = tau_r_s * math.log1p(decay_s / tau_r_s)
    peak = math.exp(-t_peak / decay_s) * (-math.expm1(-t_peak / tau_r_s))
    return t_peak, peak


def _source_rise_on_fine_grid(tau_r_s, decay_s):
    t_peak, peak = _source_peak(tau_r_s, decay_s)
    t = np.arange(0.0, t_peak, SOURCE_CALIBRATION_DT_S)
    t = np.append(t, t_peak)
    y = np.exp(-t / decay_s) * (-np.expm1(-t / tau_r_s)) / peak
    t10 = _crossing(t, y, 0.1, y.size - 1)
    t90 = _crossing(t, y, 0.9, y.size - 1)
    return t90 - t10


@lru_cache(maxsize=128)
def source_parameters(intrinsic_rise_us, decay_s=SOURCE_DECAY_S):
    """Calibrate tau_r by bisection with 0.2 us crossing interpolation."""
    target = float(intrinsic_rise_us) * 1e-6
    decay_s = float(decay_s)
    if not np.isfinite(target) or target <= 0 or decay_s <= 0:
        raise ValueError("source rise and decay time must be positive")
    lo = max(target * 1e-4, 1e-12)
    hi = max(target, 1e-9)
    for _ in range(30):
        if _source_rise_on_fine_grid(hi, decay_s) >= target:
            break
        hi *= 2.0
    else:
        raise ValueError("requested rise is too long for the selected decay time")
    for _ in range(64):
        mid = (lo + hi) / 2.0
        if _source_rise_on_fine_grid(mid, decay_s) < target:
            lo = mid
        else:
            hi = mid
    tau_r = (lo + hi) / 2.0
    t_peak, _ = _source_peak(tau_r, decay_s)
    return {"intrinsic_rise_us": float(intrinsic_rise_us),
            "tau_r_s": tau_r, "tau_r_us": tau_r * 1e6,
            "decay_s": decay_s, "decay_ms": decay_s * 1e3,
            "analytic_onset_to_peak_us": t_peak * 1e6,
            "calibration_dt_us": SOURCE_CALIBRATION_DT_S * 1e6}


def pulse_waveform(dt, intrinsic_rise_us, decay_s=SOURCE_DECAY_S,
                   tail_e_folds=12.0):
    """Return (time_seconds, unit-peak source); tail is finite and documented."""
    dt = _validate_dt(dt)
    p = source_parameters(float(intrinsic_rise_us), float(decay_s))
    t = np.arange(int(math.ceil(tail_e_folds * decay_s / dt)) + 1) * dt
    _, peak = _source_peak(p["tau_r_s"], decay_s)
    y = np.exp(-t / decay_s) * (-np.expm1(-t / p["tau_r_s"])) / peak
    return t, y


def _fft_convolve(x, z):
    size = len(x) + len(z) - 1
    nfft = 1 << (size - 1).bit_length()
    y = np.fft.irfft(np.fft.rfft(x, nfft) * np.fft.rfft(z, nfft), nfft)[:size]
    return np.maximum(y, 0.0)


def one_pole_kernel(dt, fc_hz, tail_tolerance=1e-12):
    """Exact first-order kernel for piecewise-linear input sample values.

    y[n] = a*y[n-1] + b0*x[n] + b1*x[n-1], zero prehistory.
    This is not a claim that a first-order filter provides adequate antialiasing.
    """
    dt = _validate_dt(dt)
    fc_hz = float(fc_hz)
    if not np.isfinite(fc_hz) or fc_hz <= 0:
        raise ValueError("fc_hz must be positive, or None to omit electronics")
    u = 2.0 * math.pi * fc_hz * dt
    a = math.exp(-u)
    q = -math.expm1(-u) / u
    b0, b1 = 1.0 - q, q - a
    n = max(2, int(math.ceil(-math.log(tail_tolerance) / u)) + 2)
    kernel = np.empty(n)
    kernel[0] = b0
    kernel[1:] = (a * b0 + b1) * np.exp(-u * np.arange(n - 1))
    return kernel


def impulse_metrics(h, dt):
    """Timing of the supplied weighted impulse, with total mass unnormalized."""
    h, dt = _response(h, dt)
    mass = h * dt
    total = float(np.sum(mass))
    out = {"integrated_response_weight": total,
           "response_window_end_us": (h.size - 1) * dt * 1e6,
           "quantiles_conditional_on_supplied_response": True,
           "delay_grid_origin_us": 0.0,
           "delay_grid_convention": "mass at delay n*dt"}
    if total == 0:
        out.update({"status": "no_detected_weight", "q10_us": None,
                    "q50_us": None, "q90_us": None,
                    "spread_q90_q10_us": None,
                    "fraction_after_100us": None, "fraction_after_200us": None,
                    "fraction_after_3ms": None})
        return out
    t = np.arange(h.size) * dt
    positive = mass > 0
    # Invert the weighted empirical CDF without interpolating across gaps:
    # a direct-light atom plus a separate delayed atom must retain that gap.
    cdf = np.cumsum(mass)[positive] / total
    times = t[positive]
    indexes = np.searchsorted(cdf, [0.1, 0.5, 0.9], side="left")
    quantiles = times[np.minimum(indexes, times.size - 1)] * 1e6
    out.update({"status": "ok", "q10_us": float(quantiles[0]),
                "q50_us": float(quantiles[1]), "q90_us": float(quantiles[2]),
                "spread_q90_q10_us": float(quantiles[2] - quantiles[0])})
    for name, limit in [("100us", 100e-6), ("200us", 200e-6), ("3ms", 3e-3)]:
        out["fraction_after_" + name] = float(np.sum(mass[t > limit]) / total)
    return out


def propagate_waveform(h, dt, intrinsic_rise_us, fc_hz=None, return_metadata=False):
    """Return (time_seconds, output); output retains response weight.

    With return_metadata=True, return (time_seconds, output, metadata).
    """
    h, dt = _response(h, dt)
    _, source = pulse_waveform(dt, intrinsic_rise_us)
    y = _fft_convolve(source, h * dt)
    if fc_hz is not None:
        y = _fft_convolve(y, one_pole_kernel(dt, fc_hz))
    t = np.arange(y.size) * dt
    metadata = {"time_step_us": dt * 1e6,
                "intrinsic_rise_us": float(intrinsic_rise_us),
                "source": dict(source_parameters(float(intrinsic_rise_us))),
                "analog_bandwidth_hz": None if fc_hz is None else float(fc_hz),
                "source_tail_e_folds": 12.0,
                "electronic_kernel": None if fc_hz is None else "exact piecewise-linear input",
                "output_window_end_ms": float(t[-1] * 1e3)}
    return (t, y, metadata) if return_metadata else (t, y)


def metrics(h, dt, intrinsic_rise_us, fc_hz=None):
    """JSON-serializable result; use propagate_waveform for waveform arrays."""
    h, dt = _response(h, dt)
    im = impulse_metrics(h, dt)
    if im["integrated_response_weight"] == 0:
        return {"status": "no_detected_weight", "intrinsic_rise_us": float(intrinsic_rise_us),
                "analog_bandwidth_hz": None if fc_hz is None else float(fc_hz),
                "rise_10_90_us": None, "onset_to_peak_us": None, "impulse": im}
    t, y, result = propagate_waveform(h, dt, intrinsic_rise_us, fc_hz, return_metadata=True)
    result.update(waveform_metrics(t, y))
    result.update({"status": "ok", "impulse": im})
    return result


def _summary(values):
    a = np.asarray([v for v in values if v is not None], dtype=float)
    if a.size == 0:
        return {k: None for k in ["min", "p05", "median", "p95", "max"]}
    v = np.quantile(a, [0.0, 0.05, 0.5, 0.95, 1.0])
    return dict(zip(["min", "p05", "median", "p95", "max"], map(float, v)))


def sampling_phase_metrics(y, dt, fs_hz, phase_count=128, fc_hz=None):
    """Noiseless point-sampling across clock phases, using sampled peak.

    Each sampled waveform's own peak defines its 10/90% levels. Negative-time
    samples are explicitly zero. Crossing times use linear interpolation.
    These spreads do not include shot noise or detector/ADC noise, and the
    optional fc check is only a necessary Nyquist check, not filter validation.
    """
    dt = _validate_dt(dt)
    y = np.asarray(y, dtype=float)
    if y.ndim != 1 or y.size < 2 or not np.all(np.isfinite(y)):
        raise ValueError("y must be a finite one-dimensional waveform")
    fs_hz = float(fs_hz)
    if not np.isfinite(fs_hz) or fs_hz <= 0 or phase_count < 2:
        raise ValueError("fs_hz must be positive and phase_count at least two")
    period = 1.0 / fs_hz
    t = np.arange(y.size) * dt
    rises, peaks = [], []
    for phase in np.arange(int(phase_count)) * period / phase_count:
        n = int(math.floor((t[-1] - phase) / period))
        sampled_t = phase + np.arange(-1, n + 1) * period
        sampled_y = np.interp(sampled_t, t, y, left=0.0, right=0.0)
        m = waveform_metrics(sampled_t, sampled_y)
        rises.append(m["rise_10_90_us"])
        peaks.append(m["onset_to_peak_us"])
    continuous = waveform_metrics(t, y)
    flags = []
    if fs_hz <= 1.0e4:
        resolution = "invalid_resolution_for_30us_rise"
        flags.append("100us or longer sample intervals cannot resolve a 30us rise")
    elif fs_hz < 1.0e5:
        resolution = "marginal_resolution_for_30us_rise"
        flags.append("fewer than three sample intervals across a 30us rise")
    else:
        resolution = "at_least_three_intervals_per_30us_rise"
    if period < dt:
        flags.append("sampling faster than computational grid: interpolation adds no information")
    if fc_hz is not None and float(fc_hz) >= 0.5 * fs_hz:
        flags.append("analog cutoff reaches/exceeds Nyquist; sampling result is not antialiased")
    return {"sampling_rate_hz": fs_hz, "sample_interval_us": period * 1e6,
            "phase_count": int(phase_count), "resolution_status": resolution,
            "flags": flags, "noise_included": False,
            "continuous_rise_10_90_us": continuous["rise_10_90_us"],
            "sampled_rise_10_90_us": _summary(rises),
            "sampled_onset_to_peak_us": _summary(peaks),
            "peak_reference": "each sampled waveform's own peak"}


def selfcheck():
    """Compare a delta cloud response with independent continuous expressions."""
    dt = 2e-6
    delta = np.array([1.0 / dt])
    # Independent continuous convolution of two source exponentials with RC.
    expected_rises = {(30, 10000): 48.03517444, (30, 50000): 31.31325808,
                      (80, 10000): 90.28856458, (80, 50000): 80.40709689,
                      (200, 10000): 204.42401224, (200, 50000): 200.03885856}
    rows = []
    for (rise, fc), expected in expected_rises.items():
        m = metrics(delta, dt, rise, fc)
        error = abs(m["rise_10_90_us"] - expected)
        if error > 0.6:
            raise AssertionError("RC/source analytical check failed: " + str((rise, fc, error)))
        rows.append({"intrinsic_rise_us": rise, "fc_hz": fc,
                     "computed_rise_us": m["rise_10_90_us"],
                     "continuous_reference_rise_us": expected,
                     "absolute_error_us": error,
                     "computed_onset_to_peak_us": m["onset_to_peak_us"]})
    im = impulse_metrics(delta, dt)
    assert im["integrated_response_weight"] == 1.0 and im["q90_us"] == 0.0
    t, y = propagate_waveform(delta, dt, 30, 10000)
    sampled = sampling_phase_metrics(y, dt, 10000, fc_hz=10000)
    assert sampled["resolution_status"] == "invalid_resolution_for_30us_rise"
    assert metrics(np.zeros(10), dt, 30)["status"] == "no_detected_weight"
    return {"status": "passed", "dt_us": dt * 1e6, "analytic_checks": rows,
            "example_sampling_10kSps": sampled}


if __name__ == "__main__":
    print(json.dumps(selfcheck(), indent=2))