"""
Phase 1: Emotion & Prosody Extraction Pipeline
===============================================
Extracts per-word F0/energy/emphasis + per-segment emotion from audio.

Models used (all cached in models/ folder):
  - Silero VAD          → sentence segmentation (CPU)
  - WhisperX large-v3   → transcription + word timestamps (GPU ~3GB)
  - parselmouth (Praat) → F0 contour + energy envelope (CPU)
  - emotion2vec+ base   → emotion classification (GPU ~0.4GB)

Usage:
  conda activate s2st
  python analyze_audio.py path/to/audio.wav
  python analyze_audio.py path/to/audio.wav --language ja --output results.json
"""

import argparse
import gc
import json
import os
import sys
import time
import warnings
from pathlib import Path

import librosa
import numpy as np
import parselmouth
import torch
import torchaudio

warnings.filterwarnings("ignore")

# ── Paths ────────────────────────────────────────────────────────────────────

PROJECT_DIR = Path(__file__).resolve().parent
MODELS_DIR = PROJECT_DIR / "models"

# Point HuggingFace cache to our local models/ folder
os.environ["HF_HOME"] = str(MODELS_DIR)
os.environ["TORCH_HOME"] = str(MODELS_DIR / "torch")

# ── Config ───────────────────────────────────────────────────────────────────

SAMPLE_RATE = 16000
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Emphasis thresholds (word vs segment mean)
F0_RATIO_THRESHOLD = 1.3
ENERGY_RATIO_THRESHOLD = 1.5
DURATION_RATIO_THRESHOLD = 1.4
STRONG_MULTIPLIER = 2.0  # single feature > threshold * 2 = emphasized
MIN_FEATURES_FOR_EMPHASIS = 2  # need 2+ features above threshold


def free_gpu():
    """Free VRAM between model stages."""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


# ── Stage 1: Load & Resample Audio ──────────────────────────────────────────

def load_audio(path: str) -> np.ndarray:
    """Load audio file, resample to 16kHz mono."""
    audio, sr = librosa.load(path, sr=SAMPLE_RATE, mono=True)
    print(f"  Loaded: {len(audio)/SAMPLE_RATE:.1f}s @ {SAMPLE_RATE}Hz")
    return audio


# ── Stage 2: VAD Segmentation ───────────────────────────────────────────────

def run_vad(audio: np.ndarray) -> list[dict]:
    """Silero VAD → speech segments with start/end times."""
    model, utils = torch.hub.load(
        repo_or_dir=str(MODELS_DIR / "torch" / "hub" / "snakers4_silero-vad_master"),
        model="silero_vad",
        source="local",
        onnx=False,
    )
    get_speech_timestamps = utils[0]

    audio_tensor = torch.from_numpy(audio).float()
    timestamps = get_speech_timestamps(audio_tensor, model, sampling_rate=SAMPLE_RATE)

    segments = []
    for ts in timestamps:
        segments.append({
            "start": ts["start"] / SAMPLE_RATE,
            "end": ts["end"] / SAMPLE_RATE,
        })

    # Merge short gaps (< 0.3s) into single segments
    merged = []
    for seg in segments:
        if merged and (seg["start"] - merged[-1]["end"]) < 0.3:
            merged[-1]["end"] = seg["end"]
        else:
            merged.append(seg)

    print(f"  VAD: {len(merged)} speech segments")
    del model
    free_gpu()
    return merged


# ── Stage 3: WhisperX Transcription + Word Alignment ────────────────────────

def run_whisperx(audio: np.ndarray, language: str | None = None) -> dict:
    """WhisperX → text + word-level timestamps."""
    import whisperx

    model = whisperx.load_model(
        "large-v3-turbo",
        device=DEVICE,
        compute_type="int8_float16" if DEVICE == "cuda" else "int8",
        language=language,
        download_root=str(MODELS_DIR / "hub"),
    )

    result = model.transcribe(audio, batch_size=8, task="transcribe")
    detected_lang = result.get("language", language or "unknown")
    print(f"  WhisperX: detected language = {detected_lang}")

    # Word-level forced alignment
    align_model, align_meta = whisperx.load_align_model(
        language_code=detected_lang, device=DEVICE
    )
    aligned = whisperx.align(
        result["segments"], align_model, align_meta, audio, DEVICE,
        return_char_alignments=False,
    )

    del model, align_model
    free_gpu()

    return {"segments": aligned["segments"], "language": detected_lang}


# ── Stage 4: Prosody Extraction (F0 + Energy) ───────────────────────────────

def extract_prosody(audio: np.ndarray, start: float, end: float) -> dict:
    """Extract F0 contour and energy envelope for a time range using parselmouth."""
    # Slice audio for this segment
    start_sample = int(start * SAMPLE_RATE)
    end_sample = int(end * SAMPLE_RATE)
    segment_audio = audio[start_sample:end_sample]

    if len(segment_audio) < 160:  # too short
        return {"f0_mean": 0, "f0_std": 0, "energy_mean": 0, "energy_std": 0,
                "f0_contour": [], "energy_contour": []}

    # Create parselmouth Sound object
    snd = parselmouth.Sound(segment_audio, sampling_frequency=SAMPLE_RATE)

    # F0 via autocorrelation
    pitch = snd.to_pitch_ac(
        time_step=0.01,
        pitch_floor=75,
        pitch_ceiling=600,
    )
    f0_values = pitch.selected_array["frequency"]
    f0_voiced = f0_values[f0_values > 0]

    # Energy (RMS in 10ms frames)
    frame_len = int(0.01 * SAMPLE_RATE)  # 160 samples
    n_frames = len(segment_audio) // frame_len
    energy_values = np.array([
        np.sqrt(np.mean(segment_audio[i * frame_len:(i + 1) * frame_len] ** 2))
        for i in range(n_frames)
    ])

    return {
        "f0_mean": float(np.mean(f0_voiced)) if len(f0_voiced) > 0 else 0,
        "f0_std": float(np.std(f0_voiced)) if len(f0_voiced) > 0 else 0,
        "energy_mean": float(np.mean(energy_values)) if len(energy_values) > 0 else 0,
        "energy_std": float(np.std(energy_values)) if len(energy_values) > 0 else 0,
        "f0_contour": f0_values.tolist(),
        "energy_contour": energy_values.tolist(),
    }


def extract_word_prosody(audio: np.ndarray, word_start: float, word_end: float) -> dict:
    """Extract F0 and energy for a single word's time span."""
    start_sample = int(word_start * SAMPLE_RATE)
    end_sample = int(word_end * SAMPLE_RATE)
    word_audio = audio[start_sample:end_sample]

    if len(word_audio) < 80:
        return {"f0": 0, "energy": 0, "duration": word_end - word_start}

    snd = parselmouth.Sound(word_audio, sampling_frequency=SAMPLE_RATE)

    # Praat requires pitch_floor > 3/duration, adapt for short words
    min_pitch_floor = 3.0 / snd.duration + 1.0
    pitch_floor = max(75, min_pitch_floor)
    pitch = snd.to_pitch_ac(time_step=0.005, pitch_floor=pitch_floor, pitch_ceiling=600)
    f0 = pitch.selected_array["frequency"]
    f0_voiced = f0[f0 > 0]

    energy = float(np.sqrt(np.mean(word_audio ** 2)))

    return {
        "f0": float(np.mean(f0_voiced)) if len(f0_voiced) > 0 else 0,
        "energy": energy,
        "duration": word_end - word_start,
    }


# ── Stage 5: Emphasis Detection ─────────────────────────────────────────────

def detect_emphasis(words: list[dict], segment_prosody: dict, segment_duration: float) -> list[dict]:
    """
    Score each word for emphasis based on F0, energy, and duration
    relative to segment averages.

    A word is emphasized if:
      - 2+ features exceed threshold ratios vs segment mean, OR
      - Any single feature exceeds 2x the threshold
    """
    n_words = len(words)
    if n_words == 0:
        return words

    seg_f0 = segment_prosody["f0_mean"]
    seg_energy = segment_prosody["energy_mean"]
    expected_duration = segment_duration / n_words if n_words > 0 else 1.0

    for word in words:
        wp = word.get("prosody", {})
        word_f0 = wp.get("f0", 0)
        word_energy = wp.get("energy", 0)
        word_dur = wp.get("duration", 0)

        # Compute ratios
        f0_ratio = word_f0 / seg_f0 if seg_f0 > 0 else 0
        energy_ratio = word_energy / seg_energy if seg_energy > 0 else 0
        duration_ratio = word_dur / expected_duration if expected_duration > 0 else 0

        # Count features above threshold
        features_above = 0
        strong_hit = False

        if f0_ratio > F0_RATIO_THRESHOLD:
            features_above += 1
        if f0_ratio > F0_RATIO_THRESHOLD * STRONG_MULTIPLIER:
            strong_hit = True

        if energy_ratio > ENERGY_RATIO_THRESHOLD:
            features_above += 1
        if energy_ratio > ENERGY_RATIO_THRESHOLD * STRONG_MULTIPLIER:
            strong_hit = True

        if duration_ratio > DURATION_RATIO_THRESHOLD:
            features_above += 1
        if duration_ratio > DURATION_RATIO_THRESHOLD * STRONG_MULTIPLIER:
            strong_hit = True

        emphasized = features_above >= MIN_FEATURES_FOR_EMPHASIS or strong_hit

        word["emphasis"] = {
            "is_emphasized": emphasized,
            "f0_ratio": round(f0_ratio, 2),
            "energy_ratio": round(energy_ratio, 2),
            "duration_ratio": round(duration_ratio, 2),
            "features_above_threshold": features_above,
        }

    return words


# ── Stage 6: Emotion Classification ─────────────────────────────────────────

def run_emotion(audio: np.ndarray, segments: list[dict]) -> list[dict]:
    """emotion2vec+ base → per-segment emotion label."""
    from funasr import AutoModel

    model = AutoModel(
        model="emotion2vec/emotion2vec_plus_base",
        hub="hf",
        cache_dir=str(MODELS_DIR / "hub"),
        device=DEVICE,
    )

    EMOTION_LABELS = [
        "angry", "disgusted", "fearful", "happy", "neutral",
        "other", "sad", "surprised", "unknown",
    ]

    for seg in segments:
        start_sample = int(seg["start"] * SAMPLE_RATE)
        end_sample = int(seg["end"] * SAMPLE_RATE)
        seg_audio = audio[start_sample:end_sample]

        if len(seg_audio) < 400:
            seg["emotion"] = {"label": "unknown", "scores": {}}
            continue

        result = model.generate(
            input=seg_audio.astype(np.float32),
            output_dir=None,
            granularity="utterance",
        )

        scores = result[0]["scores"]
        if isinstance(scores, np.ndarray):
            scores = scores.tolist()

        top_idx = int(np.argmax(scores))
        label = EMOTION_LABELS[top_idx] if top_idx < len(EMOTION_LABELS) else "unknown"

        seg["emotion"] = {
            "label": label,
            "confidence": round(float(scores[top_idx]), 3),
            "scores": {EMOTION_LABELS[i]: round(float(s), 3) for i, s in enumerate(scores) if i < len(EMOTION_LABELS)},
        }

    del model
    free_gpu()
    return segments


# ── Main Pipeline ────────────────────────────────────────────────────────────

def analyze(audio_path: str, language: str | None = None) -> dict:
    """Run full Phase 1 analysis pipeline."""
    t0 = time.time()

    print("[1/6] Loading audio...")
    audio = load_audio(audio_path)

    print("[2/6] Running VAD...")
    vad_segments = run_vad(audio)

    print("[3/6] Running WhisperX transcription + alignment...")
    whisper_result = run_whisperx(audio, language)
    detected_lang = whisper_result["language"]

    # Map WhisperX segments to VAD segments
    print("[4/6] Extracting prosody (F0 + energy)...")
    segments = []
    for wx_seg in whisper_result["segments"]:
        seg_start = wx_seg.get("start", 0)
        seg_end = wx_seg.get("end", 0)
        seg_text = wx_seg.get("text", "")
        seg_duration = seg_end - seg_start

        # Segment-level prosody
        prosody = extract_prosody(audio, seg_start, seg_end)

        # Word-level prosody
        words = []
        for w in wx_seg.get("words", []):
            w_start = w.get("start")
            w_end = w.get("end")
            if w_start is None or w_end is None:
                words.append({"word": w.get("word", ""), "start": None, "end": None,
                              "prosody": {"f0": 0, "energy": 0, "duration": 0}})
                continue
            wp = extract_word_prosody(audio, w_start, w_end)
            words.append({
                "word": w.get("word", ""),
                "start": round(w_start, 3),
                "end": round(w_end, 3),
                "prosody": wp,
            })

        # Emphasis detection
        print(f"[5/6] Detecting emphasis for segment: \"{seg_text[:50]}...\"")
        words = detect_emphasis(words, prosody, seg_duration)

        segments.append({
            "start": round(seg_start, 3),
            "end": round(seg_end, 3),
            "text": seg_text,
            "prosody": {
                "f0_mean": round(prosody["f0_mean"], 1),
                "f0_std": round(prosody["f0_std"], 1),
                "energy_mean": round(prosody["energy_mean"], 6),
                "energy_std": round(prosody["energy_std"], 6),
            },
            "words": words,
        })

    print("[6/6] Running emotion classification...")
    segments = run_emotion(audio, segments)

    elapsed = time.time() - t0
    result = {
        "file": os.path.basename(audio_path),
        "duration_seconds": round(len(audio) / SAMPLE_RATE, 2),
        "language": detected_lang,
        "processing_time_seconds": round(elapsed, 1),
        "device": DEVICE,
        "segments": segments,
    }

    # Summary
    emphasized_words = sum(
        1 for seg in segments for w in seg.get("words", [])
        if w.get("emphasis", {}).get("is_emphasized", False)
    )
    total_words = sum(len(seg.get("words", [])) for seg in segments)
    emotions = [seg.get("emotion", {}).get("label", "?") for seg in segments]

    print(f"\n{'='*60}")
    print(f"  File:      {result['file']}")
    print(f"  Duration:  {result['duration_seconds']}s")
    print(f"  Language:  {detected_lang}")
    print(f"  Segments:  {len(segments)}")
    print(f"  Words:     {total_words} ({emphasized_words} emphasized)")
    print(f"  Emotions:  {', '.join(emotions)}")
    print(f"  Time:      {elapsed:.1f}s on {DEVICE}")
    print(f"{'='*60}")

    return result


# ── CLI ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Phase 1: Emotion & Prosody Extraction")
    parser.add_argument("audio", help="Path to audio file (wav/mp3/flac)")
    parser.add_argument("--language", "-l", default=None,
                        help="Language code (e.g., ja, en, hi). Auto-detected if omitted.")
    parser.add_argument("--output", "-o", default=None,
                        help="Output JSON path. Defaults to <audio_name>_analysis.json")
    args = parser.parse_args()

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

    result = analyze(args.audio, args.language)

    output_path = args.output or (Path(args.audio).stem + "_analysis.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)

    print(f"\nSaved to: {output_path}")