# -*- coding: utf-8 -*-
"""Compares a real CAA measurement against REW measuring the same loudspeaker.

`caa-validate` and `rew_compare.py` prove the arithmetic on signals whose answer
is known. They cannot prove anything about a microphone, an interface or a car.
This closes that gap: measure the same driver from the same seat with the same
microphone, once in CAA and once in REW, and put the two curves side by side.

    # 1. measure in CAA, save the session (Save & Report -> .caa)
    # 2. measure the same thing in REW, leave it open
    # 3. start REW with its API on:  roomeqwizard.exe -api
    python field_compare.py --caa run.caa

Writes field-compare.json (numbers), field-compare.svg (overlay + difference)
and prints a third-octave table. Everything is normalised by a single level
offset over the compared band -- what is being compared is the SHAPE, because
the two tools have no reason to agree on absolute level unless both were
calibrated against the same acoustic reference.

Options:
    --caa PATH        CAA session file
    --channel NAME    which channel in it (default: the first)
    --rew ID|NAME     which REW measurement (default: the selected one)
    --measure         ask REW to run its own sweep first -- THIS MAKES SOUND
    --band LO HI      compared range in Hz (default 30 16000)
    --out DIR         where to write the outputs (default: alongside the .caa)
"""
import argparse, base64, json, math, os, sys, time, urllib.error, urllib.request

import numpy as np

API = "http://127.0.0.1:4735"

THIRD_OCTAVE = [20, 25, 31.5, 40, 50, 63, 80, 100, 125, 160, 200, 250, 315, 400,
                500, 630, 800, 1000, 1250, 1600, 2000, 2500, 3150, 4000, 5000,
                6300, 8000, 10000, 12500, 16000, 20000]


# ---------------------------------------------------------------- REW ------
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 SystemExit("REW %s %s -> %d %s"
                         % (method, path, e.code, e.read().decode("utf-8", "replace")[:300]))
    except Exception as e:
        raise SystemExit("REW is not answering on %s (%s).\n"
                         "Start it with:  roomeqwizard.exe -api" % (API, e))
    try:
        return json.loads(raw)
    except ValueError:
        return raw


def rew_response(which):
    """(freqs, magDb) of a REW measurement, plus its title."""
    ms = req("GET", "/measurements")
    if not ms:
        raise SystemExit("REW has no measurements open.")

    key = None
    if which is None:
        sel = req("GET", "/measurements/selected")
        key = str(sel) if str(sel) in ms else None
    elif str(which) in ms:
        key = str(which)
    else:
        for k, v in ms.items():
            if which.lower() in str(v.get("title", "")).lower():
                key = k
                break
    if key is None:
        key = sorted(ms, key=lambda k: int(k))[-1]

    fr = req("GET", "/measurements/%s/frequency-response?unit=dBFS&smoothing=None&ppo=96" % key)
    mag = np.frombuffer(base64.b64decode(fr["magnitude"]), dtype=">f4").astype(float)
    freqs = float(fr["startFreq"]) * 2.0 ** (np.arange(mag.size) / float(fr["ppo"]))
    return freqs, mag, str(ms[key].get("title", key))


def rew_measure():
    print("asking REW to sweep -- this makes sound")
    req("POST", "/measure/command", {"command": "SPL"})
    for _ in range(300):
        time.sleep(2)
        ms = req("GET", "/measurements")
        if ms:
            return
    raise SystemExit("REW did not finish a measurement")


# ---------------------------------------------------------------- CAA ------
def caa_channels(path):
    """Every channel in a .caa session. Arrays are base64 float32, native
    (little-endian) order -- Session.cpp memcpy's them straight out."""
    d = json.load(open(path, encoding="utf-8-sig"))
    out = []
    for c in d.get("channels", []):
        mag = np.frombuffer(base64.b64decode(c.get("magDb") or ""), dtype="<f4").astype(float)
        if mag.size < 8:
            continue
        sr = float(c.get("sampleRate") or 48000.0)
        out.append(dict(name=c.get("name") or c.get("id") or "?",
                        sr=sr,
                        delayMs=float(c.get("delayMs") or 0.0),
                        freqs=np.arange(mag.size) * (sr * 0.5 / (mag.size - 1)),
                        mag=mag))
    return out, d


# ------------------------------------------------------------- compare -----
def compare(caa, rew, lo, hi):
    """Both curves on one log axis, levelled by a single offset over the band."""
    grid = np.geomspace(max(lo, 1.0), hi, 480)
    a = np.interp(grid, caa["freqs"], caa["mag"])
    b = np.interp(grid, rew[0], rew[1])

    # One constant, chosen so the mean difference is zero. Not a fit: it removes
    # a level offset and nothing else, so every shape error survives it.
    offset = float(np.mean(a - b))
    b = b + offset
    return grid, a, b, offset


def svg(path, grid, a, b, lo, hi, caa_name, rew_name):
    W, H, PAD_L, PAD_R, PAD_T = 980, 520, 56, 16, 18
    top_h, gap, bot_h = 300, 34, 120

    def x(f):
        return PAD_L + (math.log10(f) - math.log10(lo)) / (math.log10(hi) - math.log10(lo)) \
               * (W - PAD_L - PAD_R)

    span = max(1.0, float(np.nanmax([a.max(), b.max()]) - np.nanmin([a.min(), b.min()])))
    ymax = float(np.nanmax([a.max(), b.max()])) + span * 0.08
    ymin = float(np.nanmin([a.min(), b.min()])) - span * 0.08

    def y(v):
        return PAD_T + (ymax - v) / (ymax - ymin) * top_h

    diff = a - b
    dmax = max(1.0, float(np.nanmax(np.abs(diff))) * 1.25)

    def yd(v):
        return PAD_T + top_h + gap + (dmax - v) / (2 * dmax) * bot_h

    p = ['<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 %d %d" width="%d" height="%d" '
         'font-family="system-ui,-apple-system,Segoe UI,sans-serif" font-size="11">' % (W, H, W, H),
         '<rect width="%d" height="%d" fill="#0b0b0d"/>' % (W, H)]

    for f in THIRD_OCTAVE:
        if not (lo <= f <= hi):
            continue
        nice = f in (20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000)
        p.append('<line x1="%.1f" y1="%d" x2="%.1f" y2="%.1f" stroke="#2a2a30" stroke-width="%s"/>'
                 % (x(f), PAD_T, x(f), PAD_T + top_h, "1" if nice else "0.5"))
        if nice:
            lab = "%gk" % (f / 1000) if f >= 1000 else "%g" % f
            p.append('<text x="%.1f" y="%.1f" fill="#8a8a92" text-anchor="middle">%s</text>'
                     % (x(f), PAD_T + top_h + 14, lab))

    step = 5 if (ymax - ymin) < 45 else 10
    v = math.ceil(ymin / step) * step
    while v <= ymax:
        p.append('<line x1="%d" y1="%.1f" x2="%d" y2="%.1f" stroke="#2a2a30" stroke-width="0.5"/>'
                 % (PAD_L, y(v), W - PAD_R, y(v)))
        p.append('<text x="%d" y="%.1f" fill="#8a8a92" text-anchor="end">%d dB</text>'
                 % (PAD_L - 6, y(v) + 3, v))
        v += step

    def poly(vals, fy, colour, width="1.8"):
        pts = " ".join("%.1f,%.1f" % (x(f), fy(val)) for f, val in zip(grid, vals))
        return ('<polyline points="%s" fill="none" stroke="%s" stroke-width="%s" '
                'stroke-linejoin="round"/>' % (pts, colour, width))

    p.append(poly(b, y, "#FFD60A"))
    p.append(poly(a, y, "#64D2FF"))

    p.append('<rect x="%d" y="%.1f" width="%d" height="%d" fill="#141418"/>'
             % (PAD_L, PAD_T + top_h + gap, W - PAD_L - PAD_R, bot_h))
    p.append('<line x1="%d" y1="%.1f" x2="%d" y2="%.1f" stroke="#3a3a44"/>'
             % (PAD_L, yd(0), W - PAD_R, yd(0)))
    for lvl in (1.0, -1.0):
        if abs(lvl) < dmax:
            p.append('<line x1="%d" y1="%.1f" x2="%d" y2="%.1f" stroke="#2a2a30" '
                     'stroke-dasharray="3 3"/>' % (PAD_L, yd(lvl), W - PAD_R, yd(lvl)))
            p.append('<text x="%d" y="%.1f" fill="#8a8a92" text-anchor="end">%+g dB</text>'
                     % (PAD_L - 6, yd(lvl) + 3, lvl))
    p.append(poly(diff, yd, "#30D158", "1.4"))

    p.append('<rect x="%d" y="%d" width="10" height="3" fill="#64D2FF"/>' % (PAD_L + 8, PAD_T + 10))
    p.append('<text x="%d" y="%d" fill="#c9c9d1">CAA — %s</text>'
             % (PAD_L + 24, PAD_T + 14, caa_name))
    p.append('<rect x="%d" y="%d" width="10" height="3" fill="#FFD60A"/>' % (PAD_L + 8, PAD_T + 26))
    p.append('<text x="%d" y="%d" fill="#c9c9d1">REW — %s</text>'
             % (PAD_L + 24, PAD_T + 30, rew_name))
    p.append('<text x="%d" y="%.1f" fill="#8a8a92">CAA − REW</text>'
             % (PAD_L + 8, PAD_T + top_h + gap + 14))
    p.append("</svg>")
    open(path, "w", encoding="utf-8").write("\n".join(p))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--caa", required=True)
    ap.add_argument("--channel")
    ap.add_argument("--rew")
    ap.add_argument("--measure", action="store_true")
    ap.add_argument("--band", nargs=2, type=float, default=[30.0, 16000.0])
    ap.add_argument("--out")
    args = ap.parse_args()

    chans, session = caa_channels(args.caa)
    if not chans:
        raise SystemExit("no channel in %s carries a frequency response" % args.caa)
    ch = chans[0]
    if args.channel:
        hit = [c for c in chans if args.channel.lower() in c["name"].lower()]
        if not hit:
            raise SystemExit("no channel matching %r -- have: %s"
                             % (args.channel, ", ".join(c["name"] for c in chans)))
        ch = hit[0]

    if args.measure:
        rew_measure()
    rew = rew_response(args.rew)

    lo, hi = args.band
    grid, a, b, offset = compare(ch, rew, lo, hi)
    diff = a - b

    out_dir = args.out or os.path.dirname(os.path.abspath(args.caa))
    os.makedirs(out_dir, exist_ok=True)
    svg(os.path.join(out_dir, "field-compare.svg"), grid, a, b, lo, hi, ch["name"], rew[2])

    print("CAA  %s   (%s)" % (ch["name"], os.path.basename(args.caa)))
    print("REW  %s" % rew[2])
    print("levelled by %+.2f dB over %g-%g Hz\n" % (offset, lo, hi))
    print("%9s %9s %9s %9s" % ("Hz", "CAA", "REW", "diff"))
    rows = []
    for f in THIRD_OCTAVE:
        if not (lo <= f <= hi):
            continue
        ca = float(np.interp(f, grid, a))
        rw = float(np.interp(f, grid, b))
        print("%9g %9.2f %9.2f %+9.2f" % (f, ca, rw, ca - rw))
        rows.append(dict(hz=f, caa=round(ca, 3), rew=round(rw, 3), diff=round(ca - rw, 3)))

    summary = dict(
        mean_abs=float(np.mean(np.abs(diff))),
        max_abs=float(np.max(np.abs(diff))),
        max_at_hz=float(grid[int(np.argmax(np.abs(diff)))]),
        std=float(np.std(diff)),
        level_offset=offset,
        band=[lo, hi],
        caa_channel=ch["name"],
        caa_delay_ms=ch["delayMs"],
        rew_measurement=rew[2],
        caa_file=os.path.abspath(args.caa),
        app_version=session.get("app"),
    )
    print("\nmean |diff| %.2f dB   max %.2f dB at %.0f Hz   sd %.2f dB"
          % (summary["mean_abs"], summary["max_abs"], summary["max_at_hz"], summary["std"]))

    json.dump(dict(summary=summary, thirdOctave=rows),
              open(os.path.join(out_dir, "field-compare.json"), "w"), indent=1)
    print("wrote field-compare.json and field-compare.svg in", out_dir)


if __name__ == "__main__":
    main()