#!/usr/bin/env python3
"""
clone-voice.py — High-fidelity PersonaPlex voice cloning.

Applies LuxTTS-inspired preprocessing before embedding extraction:
  1. Resample to 24kHz mono
  2. Noise reduction (spectral gating)
  3. Silence trimming (energy-based VAD)
  4. LUFS normalization to -20 dBFS (tuned for PersonaPlex)
  5. Duration optimization (trim to 4-8s sweet spot)
  6. Multi-segment embedding averaging for long clips

Usage:
  python clone-voice.py --input voice.wav --name MyVoice [--device cuda]
  python clone-voice.py --input voice.wav --name MyVoice --segments 3  # multi-segment averaging
"""

import argparse
import os
import sys
import logging
import numpy as np

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Audio preprocessing (LuxTTS-inspired)
# ---------------------------------------------------------------------------

def preprocess_audio(input_path: str, target_sr: int = 24000,
                     target_lufs: float = -20.0,
                     min_duration: float = 2.0,
                     max_duration: float = 8.0,
                     denoise: bool = True) -> "torch.Tensor":
    """
    Full preprocessing pipeline:
      1. Load + resample to 24kHz mono
      2. Noise reduction via spectral gating
      3. Silence trimming (leading/trailing)
      4. LUFS normalization
      5. Duration clipping to sweet spot
    Returns: [1, T] tensor at target_sr
    """
    import torch
    import torchaudio

    log.info(f"  Loading: {input_path}")
    wav, sr = torchaudio.load(input_path)

    # Stereo → mono
    if wav.shape[0] > 1:
        wav = wav.mean(dim=0, keepdim=True)
        log.info(f"  Converted stereo → mono")

    # Resample
    if sr != target_sr:
        wav = torchaudio.transforms.Resample(sr, target_sr)(wav)
        sr = target_sr
        log.info(f"  Resampled to {target_sr}Hz")

    audio_np = wav.squeeze().numpy()
    orig_duration = len(audio_np) / sr

    # Step 1: Noise reduction
    if denoise:
        try:
            import noisereduce as nr
            log.info(f"  Denoising (spectral gating)...")
            audio_np = nr.reduce_noise(
                y=audio_np,
                sr=sr,
                prop_decrease=0.7,  # moderate — preserve voice character
                n_fft=2048,
                hop_length=512,
            )
            log.info(f"  Denoised")
        except ImportError:
            log.info(f"  noisereduce not available, skipping denoise")

    # Step 2: Silence trimming (energy-based)
    frame_length = int(0.025 * sr)  # 25ms frames
    hop = int(0.010 * sr)           # 10ms hop
    energy = []
    for i in range(0, len(audio_np) - frame_length, hop):
        frame = audio_np[i:i + frame_length]
        energy.append(np.sqrt(np.mean(frame ** 2)))
    energy = np.array(energy)

    if len(energy) > 0:
        # Threshold: 5% of peak energy (aggressive silence removal)
        threshold = max(np.percentile(energy, 10), np.max(energy) * 0.05)
        voiced = np.where(energy > threshold)[0]
        if len(voiced) > 0:
            start_sample = max(0, voiced[0] * hop - int(0.1 * sr))   # 100ms margin
            end_sample = min(len(audio_np), (voiced[-1] + 1) * hop + int(0.1 * sr))
            audio_np = audio_np[start_sample:end_sample]
            trimmed_dur = len(audio_np) / sr
            if trimmed_dur < orig_duration - 0.2:
                log.info(f"  Trimmed silence: {orig_duration:.1f}s → {trimmed_dur:.1f}s")

    # Step 3: LUFS normalization
    try:
        import pyloudnorm as pyln
        meter = pyln.Meter(sr)
        current_lufs = meter.integrated_loudness(audio_np)
        if not np.isinf(current_lufs) and not np.isnan(current_lufs):
            audio_np = pyln.normalize.loudness(audio_np, current_lufs, target_lufs)
            log.info(f"  Normalized: {current_lufs:.1f} → {target_lufs:.1f} LUFS")
    except Exception as e:
        # Fallback: simple RMS normalization
        rms = np.sqrt(np.mean(audio_np ** 2))
        if rms > 0:
            target_rms = 10 ** (target_lufs / 20)  # approximate
            audio_np = audio_np * (target_rms / rms)
            log.info(f"  RMS-normalized (fallback)")

    # Step 4: Duration clipping
    duration = len(audio_np) / sr
    if duration > max_duration:
        # Take the most energetic segment, not just the start
        segment_samples = int(max_duration * sr)
        # Sliding window energy to find the best segment
        window_energy = []
        step = int(0.5 * sr)  # 500ms steps
        for i in range(0, len(audio_np) - segment_samples, step):
            seg = audio_np[i:i + segment_samples]
            window_energy.append((i, np.sqrt(np.mean(seg ** 2))))
        if window_energy:
            best_start = max(window_energy, key=lambda x: x[1])[0]
            audio_np = audio_np[best_start:best_start + segment_samples]
            log.info(f"  Selected best {max_duration:.0f}s segment (from {duration:.1f}s)")
    elif duration < min_duration:
        log.warning(f"  Audio too short ({duration:.1f}s < {min_duration:.0f}s minimum)")

    final_duration = len(audio_np) / sr
    log.info(f"  Final: {final_duration:.1f}s, {sr}Hz mono")

    wav_out = torch.from_numpy(audio_np).float().unsqueeze(0)  # [1, T]
    return wav_out


# ---------------------------------------------------------------------------
# Multi-segment embedding averaging
# ---------------------------------------------------------------------------

def clone_voice_multiseg(input_wav: str, output_name: str, device: str = "cuda",
                         hf_repo: str = "nvidia/personaplex-7b-v1",
                         cpu_offload: bool = False,
                         n_segments: int = 1,
                         target_lufs: float = -20.0,
                         max_seg_duration: float = 8.0):
    """
    Clone voice with optional multi-segment averaging.

    For n_segments > 1 on long clips:
      - Split preprocessed audio into overlapping segments
      - Extract embeddings from each segment independently
      - Average embeddings in latent space
      - Use final segment's KV-cache (most voice information)
    """
    import torch
    import torchaudio
    from huggingface_hub import hf_hub_download
    from moshi.models import loaders
    from moshi.models.lm import LMGen

    voices_dir = os.path.join(os.path.dirname(__file__), "custom_voices")
    os.makedirs(voices_dir, exist_ok=True)
    output_pt = os.path.join(voices_dir, f"{output_name}.pt")

    if os.path.exists(output_pt):
        log.info(f"Voice '{output_name}' already exists at {output_pt}")
        log.info("Delete it first if you want to re-clone.")
        return output_pt

    # ── Preprocessing ────────────────────────────────────────────────────
    log.info("Step 1: Preprocessing audio...")
    # For multi-segment, preprocess with longer max duration
    effective_max = max_seg_duration if n_segments == 1 else max_seg_duration * n_segments
    preprocessed = preprocess_audio(
        input_wav,
        target_lufs=target_lufs,
        max_duration=effective_max,
        denoise=True,
    )

    # Save preprocessed audio for inspection
    prep_path = os.path.join(voices_dir, f"{output_name}_preprocessed.wav")
    torchaudio.save(prep_path, preprocessed, 24000)
    log.info(f"  Saved preprocessed audio: {prep_path}")

    # ── Split into segments if requested ─────────────────────────────────
    sr = 24000
    total_samples = preprocessed.shape[1]
    total_duration = total_samples / sr

    if n_segments > 1 and total_duration > max_seg_duration:
        seg_samples = int(max_seg_duration * sr)
        overlap = int(1.0 * sr)  # 1s overlap
        stride = max(seg_samples - overlap, int(2.0 * sr))
        segments = []
        for i in range(0, total_samples - seg_samples + 1, stride):
            seg = preprocessed[:, i:i + seg_samples]
            segments.append(seg)
            if len(segments) >= n_segments:
                break
        # If we didn't get enough, just use what we have
        if len(segments) == 0:
            segments = [preprocessed]
        log.info(f"  Split into {len(segments)} segments ({max_seg_duration:.0f}s each, 1s overlap)")
    else:
        segments = [preprocessed]

    # ── Load models ──────────────────────────────────────────────────────
    log.info("\nStep 2: Loading models...")
    mimi_weight = hf_hub_download(hf_repo, loaders.MIMI_NAME)
    mimi = loaders.get_mimi(mimi_weight, device)
    mimi.streaming_forever(1)

    moshi_weight = hf_hub_download(hf_repo, loaders.MOSHI_NAME)
    lm = loaders.get_moshi_lm(moshi_weight, device=device, cpu_offload=cpu_offload)
    lm.eval()

    frame_size = int(mimi.sample_rate / mimi.frame_rate)
    other_mimi = loaders.get_mimi(mimi_weight, device)
    other_mimi.streaming_forever(1)

    log.info("  Warming up...")
    from moshi.offline import warmup
    lm_gen = LMGen(
        lm,
        audio_silence_frame_cnt=int(0.5 * mimi.frame_rate),
        sample_rate=mimi.sample_rate,
        device=device,
        frame_rate=mimi.frame_rate,
        save_voice_prompt_embeddings=True,
        use_sampling=False,
        temp=0.8,
        temp_text=0.7,
        top_k=250,
        top_k_text=25,
    )
    lm_gen.streaming_forever(1)
    warmup(mimi, other_mimi, lm_gen, device, frame_size)

    # ── Extract embeddings per segment ───────────────────────────────────
    all_embeddings = []
    final_cache = None

    for seg_idx, seg_audio in enumerate(segments):
        seg_dur = seg_audio.shape[1] / sr
        log.info(f"\nStep 3.{seg_idx + 1}: Extracting embeddings from segment {seg_idx + 1}/{len(segments)} ({seg_dur:.1f}s)...")

        # Write temp WAV for this segment
        tmp_seg = os.path.join(voices_dir, f"_tmp_seg_{seg_idx}.wav")
        torchaudio.save(tmp_seg, seg_audio, sr)

        # Reset state for each segment
        mimi.reset_streaming()
        other_mimi.reset_streaming()
        lm_gen.reset_streaming()

        # Load and process
        lm_gen.load_voice_prompt(tmp_seg)

        # Trick save path so .pt goes where we want
        lm_gen.voice_prompt = os.path.join(voices_dir, f"_tmp_seg_{seg_idx}.wav")
        lm_gen._step_voice_prompt(mimi)

        # Collect embeddings
        auto_saved = os.path.join(voices_dir, f"_tmp_seg_{seg_idx}.pt")
        if os.path.exists(auto_saved):
            state = torch.load(auto_saved, map_location="cpu", weights_only=False)
            all_embeddings.append(state["embeddings"])
            final_cache = state["cache"]
            os.remove(auto_saved)
            log.info(f"    Extracted {state['embeddings'].shape[0]} frames")
        else:
            log.warning(f"    Segment {seg_idx + 1} failed to produce embeddings")

        # Clean up temp
        os.remove(tmp_seg)

    if not all_embeddings:
        log.error("No embeddings extracted!")
        return None

    # ── Average embeddings across segments ───────────────────────────────
    if len(all_embeddings) > 1:
        log.info(f"\nStep 4: Averaging {len(all_embeddings)} segment embeddings...")
        # Pad/truncate to same length, then average
        min_frames = min(e.shape[0] for e in all_embeddings)
        truncated = [e[:min_frames] for e in all_embeddings]
        stacked = torch.stack(truncated, dim=0)  # [N_segs, min_frames, 1, 1, 4096]
        averaged = stacked.mean(dim=0)            # [min_frames, 1, 1, 4096]
        log.info(f"  Averaged: {averaged.shape[0]} frames from {len(all_embeddings)} segments")
    else:
        averaged = all_embeddings[0]

    # ── Save final .pt ───────────────────────────────────────────────────
    torch.save({
        "embeddings": averaged.detach().cpu(),
        "cache": final_cache.detach().cpu() if final_cache is not None else torch.zeros(1, 17, 4, dtype=torch.int64),
    }, output_pt)

    # Verify
    state = torch.load(output_pt, map_location="cpu", weights_only=False)
    emb_shape = state["embeddings"].shape
    log.info(f"\nVoice cloned successfully!")
    log.info(f"  Output: {output_pt}")
    log.info(f"  Embeddings: {emb_shape} ({emb_shape[0]} frames, ~{emb_shape[0] / 12.5:.1f}s)")
    log.info(f"  Segments averaged: {len(all_embeddings)}")
    log.info(f"  Preprocessing: denoise + silence trim + {target_lufs:.0f} LUFS normalize")

    return output_pt


def main():
    parser = argparse.ArgumentParser(
        description="High-fidelity PersonaPlex voice cloning"
    )
    parser.add_argument(
        "--input", "-i", required=True,
        help="Input audio file (WAV, MP3, FLAC — any format torchaudio supports)"
    )
    parser.add_argument(
        "--name", "-n", required=True,
        help="Name for the cloned voice (e.g. 'MyVoice')"
    )
    parser.add_argument(
        "--device", "-d", default="cuda",
        help="Device to run on (default: cuda)"
    )
    parser.add_argument(
        "--cpu-offload", action="store_true",
        help="Offload to CPU if GPU memory is insufficient"
    )
    parser.add_argument(
        "--hf-repo", default="nvidia/personaplex-7b-v1",
        help="HuggingFace model repo"
    )
    parser.add_argument(
        "--segments", "-s", type=int, default=1,
        help="Number of segments to average (1 = single pass, 3+ = multi-segment averaging for longer clips)"
    )
    parser.add_argument(
        "--lufs", type=float, default=-20.0,
        help="Target LUFS normalization (default: -20, PersonaPlex default is -24)"
    )
    parser.add_argument(
        "--max-duration", type=float, default=8.0,
        help="Max duration per segment in seconds (default: 8, built-in voices use ~4)"
    )
    parser.add_argument(
        "--no-denoise", action="store_true",
        help="Skip noise reduction (if reference is already clean)"
    )

    args = parser.parse_args()

    if not os.path.exists(args.input):
        print(f"Error: Input file not found: {args.input}")
        sys.exit(1)

    import torch
    with torch.no_grad():
        result = clone_voice_multiseg(
            input_wav=args.input,
            output_name=args.name,
            device=args.device,
            hf_repo=args.hf_repo,
            cpu_offload=args.cpu_offload,
            n_segments=args.segments,
            target_lufs=args.lufs,
            max_seg_duration=args.max_duration,
        )

    sys.exit(0 if result else 1)


if __name__ == "__main__":
    main()