"""Warm the shared separation cache for the FIXED eval set, on a free GPU, ONCE.

Separation (htdemucs v4) is identical across every restorer/generator variant we score — same degraded input
-> same stems. Without a cache, each variant eval re-separates the same 1024+62 clips (the dominant cost). This
script pre-separates the exact eval set eval_exp builds (same seed/selection) using the SAME audio_separator
path eval uses (parity), storing int16 stems content-addressed (deduped). Future eval_exp runs then hit the
cache and skip separation entirely.

Cache key = hash(degraded-array) + separator-kind + version -> matches eval_exp's run_upload(separate_chunk)
key exactly (the FLOAT-subtype temp wav eval_exp writes round-trips losslessly, no resample/trim). A miss is
harmless (eval just separates live), so exact set parity is a speed concern, not a correctness one.

Run (free GPU, isolated from cuda:0/1 jobs):
  CUDA_VISIBLE_DEVICES=2 RESTOFLOW_DEVICE=cuda:0 RESTOFLOW_SEP_CACHE=<repo>/sep_cache \
      python -m restoflow.precache_sep --ours-n 1024 --musdb-n 62
"""
from __future__ import annotations
import argparse, glob, os, random, time
from pathlib import Path
import numpy as np
import soundfile as sf
import librosa

BASE = Path("/media/maindisk/ksoil/repos/Restoration/Dac_demucs_lyra")
SYNTH_VAL = Path("/media/maindisk/melkor169/MSRKit/xlance-msr/stereo_a2sb/synthetic_out/val")  # READ-ONLY


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--device", default="cuda:0")           # physical GPU chosen via CUDA_VISIBLE_DEVICES
    p.add_argument("--ours-n", type=int, default=1024); p.add_argument("--musdb-n", type=int, default=62)
    p.add_argument("--excerpt-s", type=float, default=9.0); p.add_argument("--lp-sr", type=int, default=8000)
    p.add_argument("--musdb-root", default="/media/maindisk/melkor169/drums_dereverb/data/gmd_musdb18hq_stereo")
    a = p.parse_args()
    os.environ["RESTOFLOW_DEVICE"] = a.device
    os.environ.setdefault("RESTOFLOW_SEP_CACHE", str(BASE / "sep_cache"))
    import restoflow.app as APP
    SR = APP.SR; random.seed(0)
    cdir = APP._sep_cache_dir()
    print(f"[precache] cache dir = {cdir}")

    # ---- build the SAME degraded arrays eval_exp builds (identical seed/selection) ----
    degs = []                                              # (tag, deg[2,T])
    cleans = sorted(glob.glob(str(SYNTH_VAL / "clean" / "*.wav"))); random.shuffle(cleans); n_ours = 0
    for cp in cleans[: a.ours_n * 2]:
        dp = cp.replace("/clean/", "/degraded/")
        if not os.path.exists(dp):
            continue
        d, _ = sf.read(dp, dtype="float32"); c, _ = sf.read(cp, dtype="float32")
        d = d.T if d.ndim == 2 else np.stack([d, d]); c = c.T if c.ndim == 2 else np.stack([c, c])
        n = min(c.shape[1], d.shape[1]); degs.append((f"ours/{Path(cp).stem}", np.ascontiguousarray(d[:, :n])))
        n_ours += 1
        if n_ours >= a.ours_n:
            break
    if Path(a.musdb_root).exists():
        exc = int(a.excerpt_s * SR)
        tracks = sorted([d for d in glob.glob(str(Path(a.musdb_root) / "*")) if Path(d).is_dir()])
        random.shuffle(tracks); n_m = 0
        for td in tracks[: a.musdb_n * 2]:
            mixp = Path(td) / "mixture.wav"
            if mixp.exists():
                clean, _ = sf.read(mixp, dtype="float32"); clean = clean.T
            else:
                sts = [sf.read(Path(td) / f"{s}.wav", dtype="float32")[0].T for s in APP.STEMS
                       if (Path(td) / f"{s}.wav").exists()]
                if not sts:
                    continue
                clean = sum(sts)
            if clean.ndim == 1:
                clean = np.stack([clean, clean])
            if clean.shape[1] < exc:
                continue
            st = (clean.shape[1] - exc) // 2; clean = np.ascontiguousarray(clean[:, st:st + exc])
            deg = np.stack([librosa.resample(clean[c], orig_sr=SR, target_sr=a.lp_sr) for c in range(2)])
            deg = np.ascontiguousarray(np.stack([librosa.resample(deg[c], orig_sr=a.lp_sr, target_sr=SR)
                                                 for c in range(2)])[:, :clean.shape[1]].astype(np.float32))
            if not np.isfinite(deg).all() or float(np.sqrt(np.mean(deg ** 2))) < 1e-5:
                continue
            degs.append((f"musdb/{Path(td).name}", deg)); n_m += 1
            if n_m >= a.musdb_n:
                break
    print(f"[precache] eval set: {len(degs)} unique degraded inputs to separate")

    sep = APP.get_separator()
    print(f"[precache] separator kind = {sep['kind']}")
    t0 = time.time(); hits = 0; done = 0
    for tag, deg in degs:
        key = APP._sep_cache_key(deg, sep["kind"])
        if (cdir / f"{key}.flac").exists():               # idempotent: skip already-cached
            hits += 1; continue
        APP.separate_chunk(sep, deg)                      # populates cache under the same key eval_exp will hit
        done += 1
        if done % 25 == 0:
            el = time.time() - t0
            print(f"  [{done+hits}/{len(degs)}] separated {done} (skipped {hits}) · {el/max(done,1):.2f}s/clip "
                  f"· ETA {el/max(done,1)*(len(degs)-done-hits)/60:.1f} min", flush=True)
    sz = sum(f.stat().st_size for f in cdir.glob("*.flac")) / 1e9 if cdir.exists() else 0.0
    print(f"[precache] DONE: separated {done}, already-cached {hits}, total {len(degs)} · cache {sz:.2f} GB on disk")
    print("PRECACHE_DONE")


if __name__ == "__main__":
    main()