# -*- coding: utf-8 -*-
"""Measures in REW the same reference signals caa-validate measures in CAA.

Both tools are handed identical impulse-response WAV files whose answer is known
from theory, so the comparison is three-way: theory / CAA / REW. Needs REW 5.40
running with its REST API (roomeqwizard.exe -api).

    python rew_compare.py            # import, measure, print and save the table
"""
import base64, json, os, struct, time, urllib.error, urllib.request
import numpy as np
import scipy.signal as ss

API = "http://127.0.0.1:4735"
HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "refwav")
SR = 48000.0
os.makedirs(OUT, exist_ok=True)


def req(method, path, body=None, timeout=120):
    data = json.dumps(body).encode() if body is not None else None
    r = urllib.request.Request(API + path, data=data, method=method,
                               headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(r, timeout=timeout) as f:
            raw = f.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as e:
        raise RuntimeError("%s %s -> %d %s" % (method, path, e.code,
                                               e.read().decode("utf-8", "replace")[:200]))
    try:
        return json.loads(raw)
    except ValueError:
        return raw


def write_wav(path, x):
    """32-bit float WAV -- keeps the 90 dB of range these references need."""
    x = np.asarray(x, dtype=np.float32)
    with open(path, "wb") as f:
        n = len(x) * 4
        f.write(b"RIFF" + struct.pack("<I", 36 + n) + b"WAVE")
        f.write(b"fmt " + struct.pack("<IHHIIHH", 16, 3, 1, int(SR), int(SR) * 4, 4, 32))
        f.write(b"data" + struct.pack("<I", n) + x.tobytes())


def lr4_coeffs(fc):
    """The same RBJ section CAA builds: Q = 0.7071, cascaded twice."""
    w0 = 2 * np.pi * fc / SR
    al = np.sin(w0) / (2 * 0.70710678)
    c = np.cos(w0)
    b = np.array([(1 - c) / 2, (1 - c), (1 - c) / 2]) / (1 + al)
    a = np.array([1.0, -2 * c / (1 + al), (1 - al) / (1 + al)])
    return b, a


def lr4_ir(fc, n=32768, pre=240):
    b, a = lr4_coeffs(fc)
    x = np.zeros(n)
    x[pre] = 1.0
    return ss.lfilter(b, a, ss.lfilter(b, a, x))


def lr4_theory(f, fc=2500.0):
    b, a = lr4_coeffs(fc)
    z = np.exp(-1j * 2 * np.pi * np.asarray(f, float) / SR)
    return 40 * np.log10(np.abs((b[0] + b[1] * z + b[2] * z * z)
                                / (a[0] + a[1] * z + a[2] * z * z)))


def decay_ir(t60=0.5, secs=2.0, seed=12345):
    rng = np.random.RandomState(seed)
    t = np.arange(int(SR * secs)) / SR
    return rng.uniform(-1, 1, t.size) * 10 ** (-3 * t / t60)


def decode(b64):
    return np.frombuffer(base64.b64decode(b64), dtype=">f4")


def import_ir(path):
    before = set(req("GET", "/measurements"))
    req("POST", "/import/impulse-response", {"path": path.replace("/", "\\")})
    for _ in range(200):
        ms = req("GET", "/measurements")
        new = [k for k in ms if k not in before]
        if new:
            return new[-1]
    raise RuntimeError("REW did not import " + path)


def main():
    # REW keeps every import for the session; start from a clean list so the
    # ids in this run mean what they say.
    req("DELETE", "/measurements")
    lr4_path = os.path.join(OUT, "ref_lr4_2500.wav")
    dec_path = os.path.join(OUT, "ref_decay_500ms.wav")
    write_wav(lr4_path, lr4_ir(2500.0))
    write_wav(dec_path, decay_ir())

    rows = []

    # ---- crossover magnitude -------------------------------------------------
    mid = import_ir(lr4_path)
    fr = req("GET", "/measurements/%s/frequency-response?unit=dBFS&smoothing=None&ppo=96" % mid)
    mag = decode(fr["magnitude"])
    freqs = float(fr["startFreq"]) * 2.0 ** (np.arange(mag.size) / float(fr["ppo"]))

    def rew_at(f):
        return float(np.interp(f, freqs, mag))

    # REW normalises the imported IR's level; compare the SHAPE, referenced to
    # 200 Hz, exactly as caa-validate does.
    off = rew_at(200.0) - float(lr4_theory(200.0))
    for f in (1000.0, 2500.0, 5000.0, 10000.0):
        rows.append(("LR4 2.5 kHz @ %5d Hz" % f, "dB", float(lr4_theory(f)), rew_at(f) - off))

    # ---- decay ---------------------------------------------------------------
    did = import_ir(dec_path)
    # REW computes RT60 on demand; the octave filters have to be asked for
    # explicitly or the read-back 400s with "no RT60 data".
    req("POST", "/measurements/%s/command" % did,
        {"command": "Generate RT60", "parameters": {"filterOrder": 6, "octaveFrac": 1}})
    rt = None
    for _ in range(30):
        time.sleep(2)
        try:
            rt = req("GET", "/measurements/%s/rt60?octaveFrac=1" % did)
            break
        except Exception:
            continue
    band = (rt or {}).get("0.0", {})          # "0.0" is REW's broadband result
    t60 = 0.5
    q50 = 10 ** (-6 * 0.050 / t60)
    q80 = 10 ** (-6 * 0.080 / t60)
    if band:
        rows.append(("T20 of a 500 ms decay", "ms", 500.0, band["T20"] * 1000.0))
        rows.append(("T30 of a 500 ms decay", "ms", 500.0, band["T30"] * 1000.0))
        rows.append(("C50", "dB", 10 * np.log10((1 - q50) / q50), band["C50"]))
        rows.append(("C80", "dB", 10 * np.log10((1 - q80) / q80), band["C80"]))
    json.dump(rt, open(os.path.join(OUT, "rew_rt60.json"), "w"), indent=1)

    print("\n%-24s %-4s %10s %10s %8s" % ("check", "unit", "theory", "REW", "diff"))
    for name, unit, th, got in rows:
        print("%-24s %-4s %10.2f %10.2f %8.2f" % (name, unit, th, got, got - th))

    json.dump([{"check": n, "unit": u, "theory": t, "rew": g} for n, u, t, g in rows],
              open(os.path.join(OUT, "rew_results.json"), "w"), indent=1)
    print("\nsaved", os.path.join(OUT, "rew_results.json"))


if __name__ == "__main__":
    main()