SonicForgeStudio/tools/p0_baseline.py

"""P0 baseline metrics - SonicForgeStudio (do truoc khi port Rust).

Do: kich thuoc bundle, thoi gian cold start + RAM cua daw_engine.exe,
thoi gian import app.main, thoi gian mixdown 4 track 30s.

Chay: python tools/p0_baseline.py            # in ra stdout + ghi md/BASELINE_P0.md
      python tools/p0_baseline.py --print     # chi in
"""
import argparse
import io
import json
import os
import shutil
import statistics
import subprocess
import sys
import tempfile
import time
import urllib.request

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
ENGINE_EXE = os.path.join(ROOT, "dist", "daw_engine", "daw_engine.exe")
ENGINE_DIRS = ["dist/daw_engine", "src-tauri/resources/daw_engine"]
OUT_MD = os.path.join(ROOT, "md", "BASELINE_P0.md")
SR = 48000


def dir_size_mb(rel):
    p = os.path.join(ROOT, rel)
    if not os.path.exists(p):
        return None
    total = files = 0
    for r, _d, fs in os.walk(p):
        for f in fs:
            try:
                total += os.path.getsize(os.path.join(r, f))
                files += 1
            except OSError:
                pass
    return {"mb": round(total / 1048576.0, 1), "files": files}


def measure_bundle():
    out = {}
    for d in ENGINE_DIRS:
        s = dir_size_mb(d)
        if s:
            out[d] = s
    for sub in ("torch", "torchvision", "sympy", "networkx", "tensorflow"):
        p = os.path.join(ROOT, "dist", "daw_engine", "_internal", sub)
        out["has:" + sub] = os.path.exists(p)
    return out


def _tasklist_mem_kb(pid):
    try:
        r = subprocess.run(["tasklist", "/FI", "PID eq %d" % pid, "/FO", "CSV", "/NH"],
                           capture_output=True, text=True, timeout=15)
        row = r.stdout.strip().splitlines()[0]
        cells = [c.strip('"') for c in row.split('","')]
        return int(cells[-1].replace(",", "").replace(" K", "").strip())
    except Exception:
        return None


def measure_cold_start(timeout_s=90):
    if not os.path.exists(ENGINE_EXE):
        return {"error": "engine exe khong ton tai: " + ENGINE_EXE}
    env = dict(os.environ)
    env["SF_PORT"] = "8000"
    t0 = time.perf_counter()
    proc = subprocess.Popen([ENGINE_EXE], cwd=os.path.dirname(ENGINE_EXE), env=env,
                            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    port = None
    try:
        while time.perf_counter() - t0 < timeout_s:
            if proc.poll() is not None:
                return {"error": "engine thoat som rc=%s" % proc.returncode}
            for p in range(8000, 8011):
                try:
                    with urllib.request.urlopen("http://127.0.0.1:%d/health" % p, timeout=0.5) as r:
                        if r.status == 200:
                            port = p
                            break
                except Exception:
                    continue
            if port:
                break
            time.sleep(0.05)
        ms = round((time.perf_counter() - t0) * 1000.0, 1)
        mem = _tasklist_mem_kb(proc.pid)
        result = {"startup_ms": ms, "port": port, "rss_mb": round(mem / 1024.0, 1) if mem else None}
    finally:
        subprocess.run(["taskkill", "/PID", str(proc.pid), "/T", "/F"],
                       capture_output=True, text=True)
        try:
            proc.wait(timeout=10)
        except Exception:
            pass
    return result


def measure_import_time():
    code = ("import time,sys;t=time.perf_counter();import app.main;"
            "print(round((time.perf_counter()-t)*1000,1))")
    t0 = time.perf_counter()
    r = subprocess.run([sys.executable, "-c", code], cwd=ROOT, capture_output=True, text=True)
    wall = round((time.perf_counter() - t0) * 1000.0, 1)
    try:
        imp = float(r.stdout.strip().splitlines()[-1])
    except Exception:
        imp = None
    return {"import_app_main_ms": imp, "subprocess_wall_ms": wall}


def measure_render(reps=3):
    import numpy as np
    sys.path.insert(0, ROOT)
    from app.core import audio_editor
    tmp = tempfile.mkdtemp(prefix="sf_bench_")
    dur = 30.0
    n = int(SR * dur)
    t = np.arange(n) / SR
    paths = []
    for i, f in enumerate((110.0, 220.0, 330.0, 440.0)):
        y = (0.3 * np.sin(2 * np.pi * f * t)).astype(np.float32)
        p = os.path.join(tmp, "t%d.wav" % i)
        import soundfile as sf
        sf.write(p, np.column_stack([y, y]), SR, subtype="FLOAT")
        paths.append(p)
    meta = [{"file_path": p, "volume": 0.8, "muted": False} for p in paths]
    times = []
    for i in range(reps):
        out = os.path.join(tmp, "mix%d.wav" % i)
        t0 = time.perf_counter()
        audio_editor.mix_multitrack_session(meta, out, sample_rate=SR, bit_depth=32)
        times.append(round((time.perf_counter() - t0) * 1000.0, 1))
    shutil.rmtree(tmp, ignore_errors=True)
    return {"tracks": 4, "seconds_each": dur, "sr": SR,
            "runs_ms": times, "min_ms": min(times), "median_ms": round(statistics.median(times), 1),
            "realtime_factor": round(min(times) / (dur * 1000.0), 4)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--print", action="store_true")
    args = ap.parse_args()
    metrics = {
        "date": time.strftime("%Y-%m-%d %H:%M"),
        "bundle": measure_bundle(),
        "cold_start": measure_cold_start(),
        "import": measure_import_time(),
        "render_mixdown": measure_render(),
    }
    text = "# Baseline P0 - SonicForgeStudio (Python engine)\n\n" \
           "Sinh boi `tools/p0_baseline.py` ngay %s. Dung lam moc so sanh khi port Rust.\n\n```json\n%s\n```\n" % (
               metrics["date"], json.dumps(metrics, indent=1, ensure_ascii=False))
    print(text)
    if not args.print:
        os.makedirs(os.path.dirname(OUT_MD), exist_ok=True)
        with io.open(OUT_MD, "w", encoding="utf-8", newline="\n") as f:
            f.write(text)
        print("WROTE", OUT_MD)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())