import argparse
import time
from pathlib import Path

import numpy as np
import soundfile as sf

from flowhigh_sdk import audio


def decode_mel(voc, mel, chunk, overlap):
    hop = chunk - overlap
    wav = None
    for s in range(0, mel.shape[-1], hop):
        m = mel[..., s:s + chunk]
        if m.shape[-1] < chunk:
            m = np.pad(m, ((0, 0), (0, 0), (0, chunk - m.shape[-1])), mode="edge")
        c = voc.run(None, {"mel": np.ascontiguousarray(m)})[0][0, 0]
        if wav is None:
            wav = c
        else:
            n = min(overlap * 480, len(wav), len(c))
            if n > 0:
                fade = np.linspace(0.0, 1.0, n, dtype=np.float32)
                wav[-n:] = wav[-n:] * (1.0 - fade) + c[:n] * fade
            wav = np.concatenate([wav, c[n:]])
    return wav.astype(np.float32)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--input", required=True)
    p.add_argument("--output", required=True)
    p.add_argument("--models", default="models")
    p.add_argument("--steps", type=int, default=1)
    p.add_argument("--std", type=float, default=1e-4)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--chunk", type=int, default=100)
    p.add_argument("--overlap", type=int, default=20)
    p.add_argument("--flow-overlap", type=int, default=50)
    p.add_argument("--time-emb", default="assets/time_emb_4.npy")
    args = p.parse_args()

    x, sr = sf.read(args.input, always_2d=False)
    if x.ndim > 1:
        x = x.mean(axis=1)
    x = np.asarray(x, dtype=np.float32)
    target = int(round(audio.SR / sr * x.size))
    x = audio.resample(x, sr)
    peak = float(np.max(np.abs(x)))
    if peak > 0:
        x = x / peak
    cond_wave = x.astype(np.float32)
    mel = audio.log_mel(cond_wave)[0].T[None]
    frames = 500
    if mel.shape[-1] < frames:
        mel = np.pad(mel, ((0, 0), (0, 0), (0, frames - mel.shape[-1])), mode="edge")
    total = mel.shape[-1]
    if total > frames:
        n_chunks = (total - frames + (frames - args.flow_overlap) - 1) // (frames - args.flow_overlap) + 1
        padded = frames + (n_chunks - 1) * (frames - args.flow_overlap)
        if padded > total:
            mel = np.pad(mel, ((0, 0), (0, 0), (0, padded - total)), mode="edge")
    padded_total = mel.shape[-1]

    import axengine

    flow = axengine.InferenceSession(str(Path(args.models) / "flowhigh.axmodel"))
    voc = axengine.InferenceSession(str(Path(args.models) / "bigvgan.axmodel"))
    emb = np.load(args.time_emb).astype(np.float32)
    rng = np.random.RandomState(args.seed)
    hop = frames - args.flow_overlap
    starts = list(range(0, max(1, padded_total - frames + 1), hop))

    t0 = time.time()
    flow_ms = []
    out = None
    for s in starts:
        c = mel[..., s:s + frames]
        y = c + args.std * rng.randn(*c.shape).astype(np.float32)
        for i in range(args.steps):
            t1 = time.time()
            v = flow.run(None, {
                "x": np.ascontiguousarray(y),
                "cond": np.ascontiguousarray(c),
                "time_emb": np.ascontiguousarray(emb[i]),
            })[0]
            flow_ms.append((time.time() - t1) * 1000)
            y = y + (1.0 / args.steps) * v
        wav = decode_mel(voc, y, args.chunk, args.overlap)
        if out is None:
            out = wav
        else:
            n = min(args.flow_overlap * 480, len(out), len(wav))
            if n > 0:
                fade = np.linspace(0.0, 1.0, n, dtype=np.float32)
                out[-n:] = out[-n:] * (1.0 - fade) + wav[:n] * fade
            out = np.concatenate([out, wav[n:]])
    out = out[:int(total * 480)]
    out = audio.postprocess(out, cond_wave[:len(out)], len(out))[:target]
    sf.write(args.output, out, audio.SR, subtype="PCM_16")
    dt = time.time() - t0
    print(f"chunks={len(starts)} flow_ms={flow_ms} total_ms={dt*1000:.1f} rtf={dt/(x.size/audio.SR):.4f}")


if __name__ == "__main__":
    main()