"""Decode the command-verification tones on Jodrell Bank's Zond 6 tape #42 (November 1968).

Audio: https://www.jb.man.ac.uk/multimedia/audio/zond6.html
(CC BY-NC-ND 4.0, Jodrell Bank Centre for Astrophysics, University of Manchester). Not included here:
download zond6.mp3 from that page.

usage:    python zond6.py zond6.mp3 [outdir]
requires: python 3, numpy, scipy, ffmpeg on PATH
writes:   outdir/words.csv, outdir/pulses.csv, outdir/summary.json

Model being tested: every "second" of the signal starts with a marker tone (~1632 Hz = idle, ~2451 Hz = data).
A data second carries 8 bit cells after the marker; a ~978 Hz pulse in a cell = 1, the ~700 Hz base tone = 0.
"""
import csv
import json
import os
import subprocess
import sys

import numpy as np
from scipy import signal
from scipy.ndimage import percentile_filter

FS = 22050
HOP = 44                    # samples per track step (1.9955 ms)
DT = HOP / FS
TONES = {'base': 700.0, 'bit': 977.0, 'sync': 1632.5, 'mark': 2449.0}
SECTIONS = [(5, 63, 'before the Moon (~01:51 UT)'), (125, 211, 'after the Moon (~03:25 UT)')]  # voice-free stretches
rng = np.random.default_rng(7)


def load(path):
    raw = subprocess.run(['ffmpeg', '-v', 'quiet', '-i', path, '-f', 'f32le', '-ac', '1', '-ar', str(FS), '-'],
                         capture_output=True, check=True).stdout
    return np.frombuffer(raw, dtype=np.float32).astype(np.float64)


def tone_tracks(x):
    """log10 of each tone's envelope divided by its running noise floor (20th percentile over 10 s)."""
    out = {}
    for k, fc in TONES.items():
        sos = signal.butter(4, [fc - 60, fc + 60], btype='band', fs=FS, output='sos')
        y = signal.sosfiltfilt(sos, x)
        e = np.sqrt(np.maximum(signal.sosfiltfilt(signal.butter(2, 40, fs=FS, output='sos'), y * y), 0))[::HOP]
        out[k] = np.log10(e / (percentile_filter(e, 20, size=int(10 / DT)) + 1e-12) + 1e-3)
    return out


def tone_frequencies(x):
    seg = np.concatenate([x[int(a * FS):int(b * FS)] for a, b, _ in SECTIONS])
    n = 1 << 21
    S = np.abs(np.fft.rfft(seg * np.hanning(len(seg)), n))
    fr = np.fft.rfftfreq(n, 1 / FS)
    res = {}
    for k, f0 in TONES.items():
        m = (fr > f0 - 25) & (fr < f0 + 25)
        i = np.flatnonzero(m)[np.argmax(S[m])]
        y0, y1, y2 = np.log(S[i - 1:i + 2])
        res[k] = float(fr[i] + (y0 - y2) / (2 * (y0 - 2 * y1 + y2)) * (fr[1] - fr[0]))
    return res


def runs(mask, i0, i1, min_s):
    out, i = [], i0
    while i < i1:
        if mask[i]:
            j = i
            while j < i1 and mask[j]:
                j += 1
            if (j - i) * DT >= min_s:
                out.append((i * DT, j * DT))
            i = j
        else:
            i += 1
    return out


def hysteresis(a, hi=0.9, lo=0.6):
    on = np.zeros(len(a), bool)
    s = False
    for i, v in enumerate(a):
        s = (v > hi) if not s else (v > lo)
        on[i] = s
    return on


def best_grid(starts):
    best = None
    for T in np.arange(0.100, 0.145, 0.0002):
        for a0 in np.arange(0.02, 0.11, 0.0005):
            r = starts - (a0 + np.round((starts - a0) / T) * T)
            s = np.median(np.abs(r))
            if best is None or s < best[0]:
                best = (s, T, a0)
    return best


def main():
    path = sys.argv[1]
    outdir = sys.argv[2] if len(sys.argv) > 2 else '.'
    os.makedirs(outdir, exist_ok=True)
    x = load(path)
    L = tone_tracks(x)
    freqs = tone_frequencies(x)

    # 1. data-second markers: >= 60 ms of the ~2451 Hz tone, >= 0.5 s after the previous one
    cand = []
    for t0, t1, label in SECTIONS:
        for a, b in runs(L['mark'] > 1.0, int(t0 / DT), int(t1 / DT), 0.06):
            if not cand or a - cand[-1][0] > 0.5:
                cand.append((a, b, label))
    on = np.array([c[0] for c in cand])
    gaps = np.diff(on)
    slot = float(np.median(gaps[(gaps > 0.95) & (gaps < 1.15)]))
    marks = []
    for i, (a, b, label) in enumerate(cand):          # keep markers that sit on the one-second rhythm
        d = np.abs(np.delete(on, i) - a)
        d = d[d < 12]
        if np.any(np.abs(d / slot - np.round(d / slot)) * slot < 0.06):
            marks.append((a, b, label))

    # 2. bit-tone pulses relative to each marker onset
    secs = []
    for a, b, label in marks:
        i0 = int(a / DT)
        p = runs(hysteresis(L['bit'][i0:i0 + int(1.0 / DT)]), 0, int(1.0 / DT), 0.04)
        secs.append(dict(t=a, marker_s=b - a, section=label, pulses=p))
    starts = np.array([s for sec in secs for s, e in sec['pulses']])
    med, T, a0 = best_grid(starts)
    resid = starts - (a0 + np.round((starts - a0) / T) * T)

    # 3. chance test: random pulse starts, same count and window
    null = np.array([best_grid(rng.uniform(starts.min(), starts.max(), len(starts)))[0] for _ in range(200)])

    # 4. words: method A = tone fills > p of the cell's first 80 ms; method B = 978 Hz vs 700 Hz power
    def word(t, method, p):
        bits, vals = '', []
        for k in range(8):
            c = t + a0 + k * T
            i0, i1 = int(c / DT), int((c + 0.08) / DT)
            v = float(np.mean(L['bit'][i0:i1] > 0.8)) if method == 'A' else float(np.mean(L['bit'][i0:i1] - L['base'][i0:i1]))
            vals.append(round(v, 2))
            bits += '1' if v > p else '0'
        return bits, vals

    ref = [word(s['t'], 'A', 0.5) for s in secs]
    agree = {f'A p={p}': sum(word(s['t'], 'A', p)[0] == r[0] for s, r in zip(secs, ref)) for p in (0.35, 0.65)}
    agree.update({f'B p={p}': sum(word(s['t'], 'B', p)[0] == r[0] for s, r in zip(secs, ref)) for p in (0.0, 0.5)})

    with open(os.path.join(outdir, 'words.csv'), 'w', newline='') as fh:
        w = csv.writer(fh)
        w.writerow(['n', 'tape_time_s', 'tape_mmss', 'section', 'word_b0_to_b7', 'off_grid_pulse'] +
                   [f'fill_b{k}' for k in range(8)])
        for n, (s, (bits, vals)) in enumerate(zip(secs, ref)):
            st = np.array([p[0] for p in s['pulses']])
            off = bool(len(st) and np.any(np.abs(st - (a0 + np.round((st - a0) / T) * T)) > 0.020))
            mm, ss = divmod(s['t'], 60)
            w.writerow([n, round(s['t'], 3), f'{int(mm)}:{ss:06.3f}', s['section'], bits, int(off)] + vals)
    with open(os.path.join(outdir, 'pulses.csv'), 'w', newline='') as fh:
        w = csv.writer(fh)
        w.writerow(['n', 'marker_onset_tape_s', 'marker_ms', 'pulse_start_ms', 'pulse_end_ms', 'cell', 'grid_residual_ms'])
        for n, s in enumerate(secs):
            for ps, pe in s['pulses']:
                c = int(np.round((ps - a0) / T))
                w.writerow([n, round(s['t'], 3), round(s['marker_s'] * 1000, 1), round(ps * 1000, 1), round(pe * 1000, 1),
                            c, round((ps - (a0 + c * T)) * 1000, 1)])
    summary = dict(
        audio='https://www.jb.man.ac.uk/multimedia/audio/zond6.html (CC BY-NC-ND 4.0, Jodrell Bank Centre for Astrophysics)',
        tones_hz={k: round(v, 2) for k, v in freqs.items()},
        master_hz={'base x7': round(freqs['base'] * 7, 1), 'bit x5': round(freqs['bit'] * 5, 1),
                   'sync x3': round(freqs['sync'] * 3, 1), 'mark x2': round(freqs['mark'] * 2, 1)},
        slot_s=round(slot, 4), data_seconds=len(secs), pulses=int(len(starts)),
        marker_ms_median=round(float(np.median([s['marker_s'] for s in secs])) * 1000, 1),
        grid=dict(cell_ms=round(T * 1000, 2), first_cell_ms=round(a0 * 1000, 1),
                  median_abs_residual_ms=round(med * 1000, 2), within_10ms=round(float(np.mean(np.abs(resid) < 0.010)), 3),
                  within_20ms=round(float(np.mean(np.abs(resid) < 0.020)), 3)),
        chance_test=dict(trials=len(null), as_good_or_better=int(np.sum(null <= med)),
                         random_median_abs_residual_ms=round(float(np.median(null)) * 1000, 2)),
        decoder_agreement_of=len(secs), decoder_agreement=agree,
    )
    json.dump(summary, open(os.path.join(outdir, 'summary.json'), 'w'), indent=1, ensure_ascii=False)
    print(json.dumps(summary, indent=1, ensure_ascii=False))
    for s, (bits, _) in zip(secs, ref):
        mm, ss = divmod(s['t'], 60)
        print(f"{int(mm)}:{ss:06.3f}  {bits}")


if __name__ == '__main__':
    main()