#!/usr/bin/env python3
"""
Audio Analyzer v3 for Lux
Adds melodic contour tracking (note flow, pitch movement, melodic shape)
on top of v2's structural/harmonic/textural layers.

Pipeline: audio → features → holistic description → Lux

New in v3:
- F0 tracking with librosa.pyin (probabilistic YIN)
- Note segmentation (onset + pitch change detection)
- Melodic contour: range, step/leap ratio, direction trends
- Pitch histogram fine-grained by time
"""

import subprocess
import sys
import os
import tempfile
import json
import numpy as np

def convert_to_wav(input_path):
    """Convert any audio format to WAV using ffmpeg"""
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp:
        tmp_path = tmp.name

    cmd = [
        'ffmpeg', '-i', input_path,
        '-acodec', 'pcm_s16le', '-ar', '11025', '-ac', '1',
        '-y', tmp_path
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"Error converting audio: {result.stderr}", file=sys.stderr)
        return None
    return tmp_path

def extract_features(wav_path):
    """Extract comprehensive audio features including melodic contour"""
    import librosa

    # Load audio
    y, sr = librosa.load(wav_path, sr=11025)
    duration = librosa.get_duration(y=y, sr=sr)

    # === STRUCTURAL/TEMPORAL LAYER ===
    tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sr)
    tempo = float(tempo) if np.isscalar(tempo) else float(tempo[0])
    beat_times = librosa.frames_to_time(beat_frames, sr=sr)

    onset_env = librosa.onset.onset_strength(y=y, sr=sr)

    hop_length = 512
    tempogram = librosa.feature.tempogram(onset_envelope=onset_env, sr=sr, hop_length=hop_length)
    tempo_over_time = np.mean(tempogram, axis=0)

    # === HARMONIC/EMOTIONAL LAYER ===
    chroma = librosa.feature.chroma_stft(y=y, sr=sr)
    chroma_mean = np.mean(chroma, axis=1)

    key_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
    key_idx = np.argmax(chroma_mean)
    key = key_names[key_idx]

    major_profile = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88])
    minor_profile = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17])

    major_corr = np.corrcoef(chroma_mean, major_profile)[0, 1]
    minor_corr = np.corrcoef(chroma_mean, minor_profile)[0, 1]
    mode = "major" if major_corr > minor_corr else "minor"

    spectral_centroid = librosa.feature.spectral_centroid(y=y, sr=sr)
    centroid_mean = float(np.mean(spectral_centroid))

    spectral_contrast = librosa.feature.spectral_contrast(y=y, sr=sr, n_bands=4)
    contrast_mean = float(np.mean(spectral_contrast))

    # === MELODIC CONTOUR LAYER (NEW) ===
    # F0 tracking using probabilistic YIN
    f0, voiced_flag, voiced_probs = librosa.pyin(
        y, fmin=librosa.note_to_hz('C2'),
        fmax=librosa.note_to_hz('C7'),
        sr=sr
    )
    f0_times = librosa.frames_to_time(np.arange(len(f0)), sr=sr)

    # Filter to voiced frames only
    f0_voiced = f0[voiced_flag] if np.any(voiced_flag) else np.array([])
    voiced_times = f0_times[voiced_flag] if np.any(voiced_flag) else np.array([])

    # Convert F0 to MIDI notes for easier contour analysis
    f0_midi = librosa.hz_to_midi(f0_voiced) if len(f0_voiced) > 0 else np.array([])

    melodic_stats = {}
    if len(f0_midi) > 0:
        # Range
        melodic_stats['min_midi'] = float(np.min(f0_midi))
        melodic_stats['max_midi'] = float(np.max(f0_midi))
        melodic_stats['range_semitones'] = float(np.max(f0_midi) - np.min(f0_midi))
        melodic_stats['mean_midi'] = float(np.mean(f0_midi))
        melodic_stats['std_midi'] = float(np.std(f0_midi))

        # Find closest note names
        note_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        lowest_note = note_names[int(round(melodic_stats['min_midi'])) % 12]
        highest_note = note_names[int(round(melodic_stats['max_midi'])) % 12]
        melodic_stats['lowest_note'] = lowest_note
        melodic_stats['highest_note'] = highest_note

        # Contour direction: diffs between consecutive voiced frames
        midi_diffs = np.diff(f0_midi)
        midi_diffs = midi_diffs[np.abs(midi_diffs) > 0.5]  # Filter tiny changes

        if len(midi_diffs) > 0:
            ascending = np.sum(midi_diffs > 0)
            descending = np.sum(midi_diffs < 0)
            steps = np.sum(np.abs(midi_diffs) <= 2)  # Step = 2 semitones or less
            leaps = np.sum(np.abs(midi_diffs) > 2)   # Leap = more than 2 semitones

            melodic_stats['direction_ratio'] = float(ascending / max(descending, 1))
            melodic_stats['step_ratio'] = float(steps / max(steps + leaps, 1))
            melodic_stats['leap_ratio'] = float(leaps / max(steps + leaps, 1))

            # Detect contour shape per time segment
            segment_size = max(1, len(midi_diffs) // 8)  # 8 segments
            segment_directions = []
            for i in range(0, len(midi_diffs), segment_size):
                seg = midi_diffs[i:i+segment_size]
                if len(seg) > 0:
                    net = np.sum(seg)
                    if net > 2:
                        segment_directions.append('rising')
                    elif net < -2:
                        segment_directions.append('falling')
                    elif np.std(seg) < 1:
                        segment_directions.append('stable')
                    else:
                        segment_directions.append('wavering')
            melodic_stats['contour_segments'] = segment_directions

            # Overall trend
            net_direction = np.sum(midi_diffs)
            if net_direction > 10:
                melodic_stats['overall_trend'] = 'ascending'
            elif net_direction < -10:
                melodic_stats['overall_trend'] = 'descending'
            else:
                melodic_stats['overall_trend'] = 'balanced'

            # Maximum leap
            melodic_stats['max_leap_semitones'] = float(np.max(np.abs(midi_diffs))) if len(midi_diffs) > 0 else 0

        # Voiced ratio (how much of the track has detectable pitch)
        melodic_stats['voiced_ratio'] = float(np.sum(voiced_flag) / len(voiced_flag)) if len(voiced_flag) > 0 else 0

        # Onset detection for note segmentation
        onset_frames = librosa.onset.onset_detect(
            y=y, sr=sr, onset_envelope=onset_env,
            backtrack=True
        )
        onset_times = librosa.frames_to_time(onset_frames, sr=sr)
        melodic_stats['estimated_note_count'] = len(onset_times)

        # Pitch at onset points
        onset_pitches = []
        for onset_frame in onset_frames:
            if onset_frame < len(f0) and voiced_flag[onset_frame]:
                onset_pitches.append(float(f0[onset_frame]))
        melodic_stats['onset_pitches'] = onset_pitches[:20]  # First 20 notes

    # === TEXTURAL LAYER ===
    rms = librosa.feature.rms(y=y)
    rms_over_time = rms[0]

    zcr = librosa.feature.zero_crossing_rate(y)
    zcr_mean = float(np.mean(zcr))

    rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr)
    rolloff_mean = float(np.mean(rolloff))

    # === ENERGY ANALYSIS ===
    window_size = int(sr * 0.5)
    num_windows = len(y) // window_size
    energy_windows = []
    for i in range(num_windows):
        window = y[i*window_size:(i+1)*window_size]
        energy = float(np.sqrt(np.mean(window**2)))
        energy_windows.append(energy)

    energy_arr = np.array(energy_windows)
    energy_mean = np.mean(energy_arr)
    energy_std = np.std(energy_arr)

    quiet_threshold = energy_mean * 0.3
    loud_threshold = energy_mean * 1.5

    quiet_sections = []
    loud_sections = []
    for i, e in enumerate(energy_windows):
        time = i * 0.5
        if e < quiet_threshold:
            quiet_sections.append(time)
        elif e > loud_threshold:
            loud_sections.append(time)

    silence_threshold = 0.01
    silence_ratio = float(np.sum(np.abs(y) < silence_threshold) / len(y))

    features = {
        'duration': duration,
        'sample_rate': sr,
        'tempo': tempo,
        'beat_count': len(beat_times),
        'key': key,
        'mode': mode,
        'major_corr': float(major_corr),
        'minor_corr': float(minor_corr),
        'spectral_centroid': centroid_mean,
        'spectral_contrast': contrast_mean,
        'rms_mean': float(np.mean(rms_over_time)),
        'rms_std': float(np.std(rms_over_time)),
        'zcr_mean': zcr_mean,
        'rolloff_mean': rolloff_mean,
        'silence_ratio': silence_ratio,
        'energy_mean': float(energy_mean),
        'energy_std': float(energy_std),
        'quiet_sections': quiet_sections[:10],
        'loud_sections': loud_sections[:10],
        'energy_windows': energy_windows[:50],
        'chroma_mean': chroma_mean.tolist(),
        'tempo_over_time': tempo_over_time.tolist()[:20],
        'melodic': melodic_stats,
    }

    return features

def synthesize_description(features):
    """Compose features into holistic, experiential description"""
    duration = features['duration']
    tempo = features['tempo']
    key = features['key']
    mode = features['mode']
    centroid = features['spectral_centroid']
    contrast = features['spectral_contrast']
    rms_mean = features['rms_mean']
    zcr = features['zcr_mean']
    rolloff = features['rolloff_mean']
    silence_ratio = features['silence_ratio']
    energy_mean = features['energy_mean']
    energy_std = features['energy_std']
    melodic = features.get('melodic', {})

    lines = []

    # === DURATION ===
    if duration < 60:
        lines.append(f"A {duration:.0f}-second piece.")
    else:
        mins = int(duration // 60)
        secs = duration % 60
        lines.append(f"A {mins}-minute, {secs:.0f}-second piece.")

    # === MOOD/CHARACTER ===
    if mode == "major":
        mood = "bright, open"
        if key in ['C', 'G', 'D']:
            mood = "bright, confident"
        elif key in ['A', 'E', 'B']:
            mood = "bright, warm"
        elif key in ['F', 'Bb', 'Eb']:
            mood = "bright, gentle"
    else:
        mood = "dark, introspective"
        if key in ['A', 'E', 'B']:
            mood = "dark, intense"
        elif key in ['D', 'G', 'C']:
            mood = "dark, melancholic"
        elif key in ['F', 'Bb', 'Eb']:
            mood = "dark, mysterious"

    lines.append(f"Key: {key} {mode}. The emotional color is {mood}.")

    # === BRIGHTNESS/TEXTURE ===
    if centroid < 1000:
        brightness = "dark and warm"
    elif centroid < 2000:
        brightness = "balanced, neither too bright nor too dark"
    elif centroid < 3000:
        brightness = "bright and present"
    else:
        brightness = "sharp and shimmering"

    if contrast > 30:
        richness = "rich and full"
    elif contrast > 20:
        richness = "moderately rich"
    else:
        richness = "lean and focused"

    lines.append(f"The sound is {brightness}, {richness}.")

    # === RHYTHM/TEMPO ===
    if tempo < 60:
        tempo_desc = "very slow, almost static"
    elif tempo < 80:
        tempo_desc = "slow, deliberate"
    elif tempo < 100:
        tempo_desc = "moderate, walking pace"
    elif tempo < 120:
        tempo_desc = "upbeat, forward-moving"
    elif tempo < 140:
        tempo_desc = "fast, energetic"
    elif tempo < 160:
        tempo_desc = "fast, driving"
    else:
        tempo_desc = "very fast, urgent"

    lines.append(f"Tempo: {tempo:.0f} BPM. {tempo_desc}.")

    # === DYNAMICS ===
    if rms_mean < 0.02:
        dynamics = "very quiet, intimate"
    elif rms_mean < 0.05:
        dynamics = "quiet, restrained"
    elif rms_mean < 0.10:
        dynamics = "moderate"
    elif rms_mean < 0.20:
        dynamics = "loud, present"
    else:
        dynamics = "very loud, intense"

    if energy_std > energy_mean * 0.3:
        dynamics += ", with significant variation"
    elif energy_std > energy_mean * 0.15:
        dynamics += ", with moderate variation"
    else:
        dynamics += ", fairly consistent"

    lines.append(f"Energy: {dynamics}.")

    # === TEXTURE ===
    if zcr > 0.1:
        texture = "noisy, textured"
    elif zcr > 0.05:
        texture = "mixed, with some grain"
    else:
        texture = "smooth, tonal"

    if rolloff < 2000:
        weight = "bass-heavy, grounded"
    elif rolloff < 4000:
        weight = "mid-focused, balanced"
    else:
        weight = "treble-focused, airy"

    lines.append(f"Texture: {texture}. Weight: {weight}.")

    # === SILENCE/SPACE ===
    if silence_ratio > 0.3:
        lines.append(f"Contains {silence_ratio*100:.0f}% silence. Lots of space and breathing room.")
    elif silence_ratio > 0.1:
        lines.append(f"Contains {silence_ratio*100:.0f}% silence. Some moments of rest.")
    else:
        lines.append("Continuous sound, minimal silence.")

    # === ENERGY SHAPES ===
    quiet = features['quiet_sections']
    loud = features['loud_sections']

    if quiet:
        times = [f"{int(t//60)}:{int(t%60):02d}" for t in quiet[:3]]
        lines.append(f"Quiet moments at {', '.join(times)}.")

    if loud:
        times = [f"{int(t//60)}:{int(t%60):02d}" for t in loud[:3]]
        lines.append(f"Peak energy at {', '.join(times)}.")

    # === MELODIC CONTOUR (NEW) ===
    if melodic:
        lines.append("")
        lines.append("--- Melodic Analysis ---")

        if melodic.get('voiced_ratio', 0) > 0.3:
            # Range
            rng = melodic.get('range_semitones', 0)
            lowest = melodic.get('lowest_note', '?')
            highest = melodic.get('highest_note', '?')
            if rng < 5:
                range_desc = "very narrow"
            elif rng < 12:
                range_desc = "narrow, staying within an octave"
            elif rng < 24:
                range_desc = "wide, spanning about two octaves"
            else:
                range_desc = "very wide, leaping across octaves"
            lines.append(f"The vocal/instrumental melody has a {range_desc} "
                        f"range ({rng:.0f} semitones, from {lowest} to {highest}).")

            # Direction
            trend = melodic.get('overall_trend', 'balanced')
            dir_ratio = melodic.get('direction_ratio', 1.0)
            if trend == 'ascending':
                lines.append("The melody tends to climb overall, reaching upward.")
            elif trend == 'descending':
                lines.append("The melody tends to descend overall, settling downward.")
            else:
                if dir_ratio > 1.3:
                    lines.append("More ascending than descending movement — the melody lifts.")
                elif dir_ratio < 0.7:
                    lines.append("More descending than ascending — the melody falls.")
                else:
                    lines.append("Movement is balanced between rising and falling.")

            # Step vs leap
            step_ratio = melodic.get('step_ratio', 0.5)
            leap_ratio = melodic.get('leap_ratio', 0.5)
            if step_ratio > 0.7:
                lines.append("The melody moves mostly in steps — smooth, connected, like walking.")
            elif leap_ratio > 0.3:
                leap_desc = "frequent" if leap_ratio > 0.5 else "occasional"
                max_leap = melodic.get('max_leap_semitones', 0)
                lines.append(f"The melody includes {leap_desc} leaps (up to {max_leap:.0f} semitones), "
                            f"creating moments of drama and surprise.")
            else:
                lines.append("The melody moves in small steps, rarely jumping far.")

            # Contour segments
            segments = melodic.get('contour_segments', [])
            if segments:
                seg_str = []
                for seg in segments:
                    if seg == 'rising':
                        seg_str.append('↑')
                    elif seg == 'falling':
                        seg_str.append('↓')
                    elif seg == 'stable':
                        seg_str.append('→')
                    else:
                        seg_str.append('↔')
                lines.append(f"Contour shape over time: {' '.join(seg_str)} "
                            f"(each arrow = one section of the song).")

            # Note density
            note_count = melodic.get('estimated_note_count', 0)
            if note_count > 0:
                notes_per_sec = note_count / max(duration, 1)
                if notes_per_sec < 1:
                    note_desc = "sparse, with space between each note"
                elif notes_per_sec < 3:
                    note_desc = "moderate, a comfortable pace"
                elif notes_per_sec < 6:
                    note_desc = "dense, with rapid note changes"
                else:
                    note_desc = "very dense, almost continuous motion"
                lines.append(f"Note density: {note_desc} "
                            f"({note_count} detected notes, {notes_per_sec:.0f}/second).")

        else:
            lines.append("No clear melodic line detected — this may be "
                        "instrumental texture, ambient sound, or the melody is "
                        "buried in the mix.")

    return '\n'.join(lines)

def main():
    if len(sys.argv) < 2:
        print("Usage: audio-analyzer-v3.py <audio_file>")
        sys.exit(1)

    input_path = sys.argv[1]
    if not os.path.exists(input_path):
        print(f"Error: File not found: {input_path}")
        sys.exit(1)

    wav_path = convert_to_wav(input_path)
    if not wav_path:
        sys.exit(1)

    try:
        features = extract_features(wav_path)
        description = synthesize_description(features)

        print(description)

        print("\n--- Raw Features (JSON) ---")
        json_features = {}
        for k, v in features.items():
            if isinstance(v, np.ndarray):
                json_features[k] = v.tolist()
            elif isinstance(v, (np.float32, np.float64)):
                json_features[k] = float(v)
            elif isinstance(v, (np.int32, np.int64)):
                json_features[k] = int(v)
            else:
                json_features[k] = v
        print(json.dumps(json_features, indent=2))

    finally:
        if os.path.exists(wav_path):
            os.unlink(wav_path)

if __name__ == '__main__':
    main()