#!/usr/bin/env python3
"""
Audio cross-correlation for movie_sync — torch GPU auto (ROCm/CUDA) with CPU fallback.

Content-to-content sync (Option A): ignores container start_time/edit-list mux delay.
If audio is delayed in container, we pad with silence on mux instead of shifting presentation.
Extracted wavs are pcm_f32le 48k mono (extract.py:41), so correlation is phase-accurate.

Usage:
  from sync_audio import load_wav, correlate, find_lag
  a, sr = load_wav("extracted/movie_a_audio.wav")
  lag_samples, peak, corr = correlate(a, b, sr=48000, method="phat", device="auto")
"""

from __future__ import annotations
import os
import sys
from typing import Tuple, Literal

import numpy as np

try:
    import soundfile as sf  # preferred for f32 wav
    HAS_SF = True
except ImportError:
    HAS_SF = False

try:
    import torch
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False


def _get_device(pref: str = "auto") -> str:
    if pref == "cpu":
        return "cpu"
    if pref in ("cuda", "auto") and HAS_TORCH and torch.cuda.is_available():
        return "cuda"
    return "cpu"


def load_wav(path: str, target_sr: int = 48000, mono: bool = True) -> Tuple[np.ndarray, int]:
    """Load wav as float32 mono. Resample if needed (cheap linear warn)."""
    if not os.path.exists(path):
        raise FileNotFoundError(path)
    if HAS_SF:
        data, sr = sf.read(path, always_2d=False, dtype="float32")
    else:
        # fallback scipy
        from scipy.io import wavfile
        sr, data = wavfile.read(path)
        if data.dtype == np.int16:
            data = data.astype(np.float32) / 32768.0
        elif data.dtype == np.int32:
            data = data.astype(np.float32) / 2147483648.0
        elif data.dtype != np.float32:
            data = data.astype(np.float32)
    if data.ndim > 1:
        if mono:
            data = data.mean(axis=1).astype(np.float32)
    if sr != target_sr:
        # simple resample via scipy if available, else warn
        try:
            from scipy.signal import resample_poly
            import math as _m
            g = _m.gcd(sr, target_sr)
            data = resample_poly(data, target_sr // g, sr // g).astype(np.float32)
            sr = target_sr
        except Exception:
            print(f"WARNING: {path} sr={sr} != {target_sr} and resample failed, using as-is", file=sys.stderr)
    return data.astype(np.float32), sr


def _next_pow2(n: int) -> int:
    return 1 << (n - 1).bit_length()


def _normalize(a: np.ndarray) -> np.ndarray:
    a = a - a.mean()
    s = a.std()
    return a / s if s > 1e-9 else a


def correlate(
    a: np.ndarray,
    b: np.ndarray,
    sr: int = 48000,
    method: Literal["phat", "standard"] = "phat",
    device: str = "auto",
    normalize: bool = True,
) -> Tuple[int, float, np.ndarray]:
    """
    Return (lag_samples, peak_value, corr_array).
    lag >0 means b delayed vs a (b starts later), need to shift b earlier by lag.
    Convention: corr = ifft(FFT(a) * conj(FFT(b))). peak at lag = argmax.
    GPU via torch if available, else numpy.
    """
    if normalize:
        a = _normalize(a)
        b = _normalize(b)

    n = len(a) + len(b) - 1
    N = _next_pow2(n)

    dev = _get_device(device)

    if dev == "cuda" and HAS_TORCH:
        # torch path — ROCm and CUDA both via cuda device
        try:
            ta = torch.from_numpy(a).to("cuda", dtype=torch.float32)
            tb = torch.from_numpy(b).to("cuda", dtype=torch.float32)
            Fa = torch.fft.rfft(ta, n=N)
            Fb = torch.fft.rfft(tb, n=N)
            R = Fa * torch.conj(Fb)
            if method == "phat":
                R = R / (torch.abs(R) + 1e-10)
            corr_t = torch.fft.irfft(R, n=N)
            # fftshift to center 0 lag
            corr_t = torch.roll(corr_t, shifts=N // 2)
            corr = corr_t.cpu().numpy()
            # compute lag from unshifted correlation for correct indexing
            corr_raw = torch.fft.irfft(R, n=N).cpu().numpy()
            peak_raw = int(np.argmax(np.abs(corr_raw)))
            if peak_raw > N // 2:
                lag = peak_raw - N
            else:
                lag = peak_raw
            peak_val = float(np.abs(corr_raw[peak_raw]))
            return lag, peak_val, corr
        except Exception as e:
            print(f"torch cuda failed ({e}), falling back cpu", file=sys.stderr)
            dev = "cpu"

    # numpy/scipy cpu path
    Fa = np.fft.rfft(a, n=N)
    Fb = np.fft.rfft(b, n=N)
    R = Fa * np.conj(Fb)
    if method == "phat":
        R = R / (np.abs(R) + 1e-10)
    corr_raw = np.fft.irfft(R, n=N)
    peak_raw = int(np.argmax(np.abs(corr_raw)))
    if peak_raw > N // 2:
        lag = peak_raw - N
    else:
        lag = peak_raw
    peak_val = float(np.abs(corr_raw[peak_raw]))
    corr_shifted = np.fft.fftshift(corr_raw)
    return lag, peak_val, corr_shifted


def lag_to_seconds(lag_samples: int, sr: int = 48000) -> float:
    return lag_samples / float(sr)


def apply_lag_with_silence(wav: np.ndarray, lag_samples: int) -> np.ndarray:
    """
    Apply lag by padding silence at start (no sound for delay).
    lag > 0: b delayed vs a. lag < 0: a delayed vs b.
    Caller decides which file to pad.
    """
    if lag_samples == 0:
        return wav
    if lag_samples > 0:
        # b delayed => pad b front with lag silence to align to a
        return np.concatenate([np.zeros(lag_samples, dtype=np.float32), wav])
    else:
        # b earlier => trim or pad a; this pads a-side: when applying to b, negative lag means b earlier, so we trim? Provide padded version for mux: pad with silence then trim alignment
        # For generic helper: pad with silence at start if lag negative we pad other file, but this is applied to lagging file only
        return np.concatenate([np.zeros(-lag_samples, dtype=np.float32), wav])


# torch device info for Makefile / GUI
def device_info() -> str:
    if not HAS_TORCH:
        return "torch not installed (cpu numpy)"
    if torch.cuda.is_available():
        try:
            name = torch.cuda.get_device_name(0)
        except Exception:
            name = "cuda"
        return f"torch {torch.__version__} + {name} (cuda/ROCm)"
    return f"torch {torch.__version__} (cpu only)"


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("a", help="wav A")
    ap.add_argument("b", help="wav B")
    ap.add_argument("--phat", action="store_true", help="use GCC-PHAT")
    ap.add_argument("--device", default="auto", choices=["auto","cpu","cuda"])
    ap.add_argument("--sr", type=int, default=48000)
    args = ap.parse_args()
    a, sa = load_wav(args.a, target_sr=args.sr)
    b, sb = load_wav(args.b, target_sr=args.sr)
    method = "phat" if args.phat else "standard"
    lag, peak, _ = correlate(a, b, sr=args.sr, method=method, device=args.device)
    print(f"device: {device_info()} method={method}")
    print(f"lag_samples={lag} lag_sec={lag/args.sr:.6f}s peak={peak:.4f}")
    print(f"apply: ffmpeg -i video.mkv -i audio.wav -filter_complex \"[1:a]adelay={max(0, lag/args.sr*1000):.0f}|{max(0, lag/args.sr*1000):.0f}[ad]\" -map 0:v -map \"[ad]\" out.mkv  (pad silence for delay)")