"""Clearline-NC-GPU on a recording: background noise out, every voice kept.

    pip install -r requirements.txt
    python example.py noisy.wav enhanced.wav [--live] [--device cuda]

The model is Voice Isolator's released network (its code is in vi/) fine-tuned to keep every talker, and it is never
given an enrollment. It is causal (it looks 10 ms ahead) and remembers a whole call, so a recording is never cut into
pieces (60 s pieces with 10 s of context match it to only 20.6 dB): up to two minutes it runs in one pass, longer ones
through Voice Isolator's live engine in float16 (memory stays at a few GB; over 600 s it matches an exact float32 run
to 55.1 dB, three times faster than real time). --live writes what a live call would hear: the same engine, 20 ms at a
time, the output 10 ms behind the input.
"""
import argparse, copy, importlib, importlib.util, os, sys
from pathlib import Path

import numpy as np
import torch

HERE = Path(os.path.abspath(__file__)).parent  # not resolve(): a Hub snapshot's files are links into its blob store
SR = 16000
_PKG = "clearline_nc_gpu_vi"  # vi/ under its own name, so a `vi` package imported from elsewhere cannot shadow it


def _vi(sub):
    if _PKG not in sys.modules:
        spec = importlib.util.spec_from_file_location(_PKG, HERE / "vi" / "__init__.py",
                                                      submodule_search_locations=[str(HERE / "vi")])
        mod = importlib.util.module_from_spec(spec)
        sys.modules[_PKG] = mod
        spec.loader.exec_module(mod)
    return importlib.import_module(f"{_PKG}.{sub}")


def load(path=HERE, device="cuda"):
    """The model from a copy of this repository (config.json and model.safetensors)."""
    return _vi("infer").load_model(str(path), device=device)


@torch.no_grad()
def enhance(model, x, one_pass_s=120.0):
    """x: float32 mono at 16 kHz -> the enhanced signal, the same length and aligned with x. One pass up to one_pass_s
    (its causal WavLM's attention mask grows with the square of the length), else the live engine on a float16 copy."""
    dev = next(model.parameters()).device
    x = np.ascontiguousarray(x, dtype=np.float32)
    if len(x) <= one_pass_s * SR:
        return model(torch.from_numpy(x)[None].to(dev))[0].float().cpu().numpy()
    out = stream(copy.deepcopy(model), x, graphs=dev.type == "cuda")
    k = out["delay"]
    y = np.zeros_like(x)
    y[: len(out["y"]) - k] = out["y"][k:]
    return y


def stream(model, x, chunk_hops=2, dtype=torch.float16, graphs=True):
    """x as one live call through Voice Isolator's engine (vi/stream.py), self-enrollment off: the model keeps its null
    enrollment, as it was trained. Converts `model` to `dtype` in place (load a second copy for offline use). Returns
    dict(y=the output of every whole 20 ms chunk, delay=the samples it lags the input)."""
    dev = next(model.parameters()).device
    s = _vi("stream").Stream(model, chunk_hops=chunk_hops, device=str(dev), dtype=dtype, graphs=graphs)
    s.srv.self_enroll = False
    return dict(y=s.push(np.asarray(x, np.float32)), delay=s.delay)


def main():
    import soundfile as sf
    ap = argparse.ArgumentParser()
    ap.add_argument("inp")
    ap.add_argument("out")
    ap.add_argument("--live", action="store_true")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    a = ap.parse_args()
    x, sr = sf.read(a.inp, dtype="float32", always_2d=True)
    x = x.mean(1)
    if sr != SR:
        import soxr
        x = soxr.resample(x, sr, SR).astype(np.float32)
    m = load(device=a.device)
    y = stream(m, x, graphs=a.device != "cpu")["y"] if a.live else enhance(m, x)
    sf.write(a.out, y, SR)


if __name__ == "__main__":
    main()