"""Applies a trained model to recordings.

    python3 ml/predict.py --model RUN_DIR/model.pt [--logf] REAL.wav [more.wav ...]

Each recording is converted to mono at 48 kHz, cut to 3 s from its onset (see
match.py), peak-normalized, turned into the spectrogram the model was trained
on, and passed through the network. To show the sensitivity to the alignment,
the prediction is repeated with the onset shifted by a few milliseconds and
the spread is reported alongside the value at the detected onset.
"""

import argparse
import math
import os
import sys

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from features import spectrogram, spectrogram_logf  # noqa: E402
from match import read_wav_any, onset, segment, FS  # noqa: E402
from train import Net, SpecNet, TARGETS, UNITS  # noqa: E402

torch.backends.cudnn.enabled = False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--logf", action="store_true")
    ap.add_argument("--shifts", default="0,-5,5,10,20", help="onset shifts in ms for the sensitivity check")
    ap.add_argument("wavs", nargs="+")
    args = ap.parse_args()
    ck = torch.load(args.model, map_location="cpu", weights_only=False)
    net = SpecNet(len(ck["targets"]), n_freq=ck["n_freq"]) if ck.get("arch") == "spec1d" else Net(len(ck["targets"]))
    net.load_state_dict(ck["state"])
    net.eval()
    mu, sd = ck["mu"], ck["sd"]
    Tinv = np.linalg.inv(ck["T"]) if "T" in ck else np.eye(len(ck["targets"]))
    logf = ck["input"] == "logf" if "input" in ck else args.logf
    spec = spectrogram_logf if logf else spectrogram
    shifts = [float(s) for s in args.shifts.split(",")]
    for path in args.wavs:
        x = read_wav_any(path)
        st = onset(x)
        preds = []
        for sh in shifts:
            s0 = max(0, st + int(sh * FS / 1000))
            X = torch.from_numpy(spec(segment(x, s0)).astype(np.float32))[None, None]
            with torch.no_grad():
                z = net(X)[0].numpy()
            preds.append((z * sd + mu) @ Tinv.T)
        preds = np.array(preds)
        p0, lo, hi = preds[0], preds.min(0), preds.max(0)
        print(f"{path}: {len(x) / FS:.2f} s, onset at {st / FS:.3f} s")
        for i, t in enumerate(TARGETS):
            unit, scale = UNITS[t]
            print(f"  {t:>4}: {math.exp(p0[i]) * scale:8.2f} {unit:<7} (onset shifts give {math.exp(lo[i]) * scale:.2f} to {math.exp(hi[i]) * scale:.2f})")
        d = {t: p0[i] for i, t in enumerate(TARGETS)}
        kappa = math.exp(d["h"] + 0.5 * (d["E"] - d["rho"])) / math.sqrt(12 * (1 - 0.29 ** 2))
        print(f"  E/rho {math.exp(d['E'] - d['rho']) / 1e6:.1f} MPa m^3/kg, kappa {kappa:.3f} m^2/s, area {math.exp(d['Lx'] + d['Ly']) / 0.0254 ** 2:.0f} in^2, aspect {math.exp(d['Lx'] - d['Ly']):.2f}")


if __name__ == "__main__":
    main()