#!/usr/bin/env python3
"""
Music analysis script for level design.
Detects beats, onsets, tempo, and structural changes in level music.
Uses scipy directly to avoid librosa/numpy version conflicts.
"""

import json
import wave
import numpy as np
from scipy import signal
from scipy.ndimage import uniform_filter1d
from pathlib import Path


def load_wav(filepath: Path) -> tuple:
    """Load a WAV file and return (samples, sample_rate)."""
    with wave.open(str(filepath), 'rb') as wf:
        sr = wf.getframerate()
        n_channels = wf.getnchannels()
        n_frames = wf.getnframes()
        sampwidth = wf.getsampwidth()

        raw_data = wf.readframes(n_frames)

        if sampwidth == 2:
            dtype = np.int16
        elif sampwidth == 4:
            dtype = np.int32
        else:
            dtype = np.uint8

        samples = np.frombuffer(raw_data, dtype=dtype)

        # Convert to mono if stereo
        if n_channels == 2:
            samples = samples.reshape(-1, 2).mean(axis=1)

        # Normalize to float -1 to 1
        samples = samples.astype(np.float32) / np.iinfo(dtype).max

    return samples, sr


def compute_onset_envelope(y: np.ndarray, sr: int, hop_length: int = 512) -> tuple:
    """Compute onset strength envelope using spectral flux."""
    # Compute spectrogram
    nperseg = 2048
    noverlap = nperseg - hop_length
    f, t, Sxx = signal.spectrogram(y, sr, nperseg=nperseg, noverlap=noverlap)

    # Spectral flux (difference between consecutive frames)
    flux = np.zeros(Sxx.shape[1])
    for i in range(1, Sxx.shape[1]):
        diff = Sxx[:, i] - Sxx[:, i-1]
        flux[i] = np.sum(np.maximum(0, diff))  # Half-wave rectification

    # Smooth the envelope
    flux = uniform_filter1d(flux, size=3)

    return flux, t


def detect_beats(onset_env: np.ndarray, times: np.ndarray, sr: int) -> tuple:
    """Detect beats using autocorrelation-based tempo estimation."""
    # Estimate tempo via autocorrelation
    corr = np.correlate(onset_env, onset_env, mode='full')
    corr = corr[len(corr)//2:]  # Take positive lags only

    # Look for peaks in tempo range (60-200 BPM)
    hop_length = 512
    min_lag = int(sr / hop_length * 60 / 200)  # 200 BPM
    max_lag = int(sr / hop_length * 60 / 60)   # 60 BPM

    if max_lag >= len(corr):
        max_lag = len(corr) - 1

    # Find strongest peak in tempo range
    search_range = corr[min_lag:max_lag+1]
    if len(search_range) == 0:
        return 120.0, []

    peak_idx = np.argmax(search_range) + min_lag

    # Convert lag to BPM
    beat_period_samples = peak_idx * hop_length
    bpm = 60.0 * sr / beat_period_samples if beat_period_samples > 0 else 120.0

    # Detect beat positions using peak picking
    beat_interval_frames = int(sr / hop_length * 60 / bpm)

    # Find peaks in onset envelope
    peaks, _ = signal.find_peaks(onset_env, distance=beat_interval_frames * 0.8,
                                  prominence=np.std(onset_env) * 0.3)

    beat_times = times[peaks] if len(peaks) > 0 else []

    return bpm, list(beat_times)


def detect_energy_peaks(onset_env: np.ndarray, times: np.ndarray) -> list:
    """Find high-energy moments (where drums kick in, etc.)."""
    threshold = np.mean(onset_env) + 1.5 * np.std(onset_env)

    # Find significant peaks
    peaks, properties = signal.find_peaks(onset_env, height=threshold,
                                           distance=20, prominence=np.std(onset_env))

    return [float(times[p]) for p in peaks]


def detect_segments(y: np.ndarray, sr: int, n_segments: int = 8) -> list:
    """Detect structural changes in the music."""
    # Compute a simple feature: RMS energy over time
    hop = sr // 4  # 250ms windows
    rms = []
    for i in range(0, len(y) - hop, hop):
        rms.append(np.sqrt(np.mean(y[i:i+hop]**2)))
    rms = np.array(rms)

    if len(rms) < 2:
        return []

    # Detect changes in RMS (derivative)
    diff = np.abs(np.diff(rms))

    # Find top N-1 change points
    if len(diff) < n_segments:
        return []

    peaks, _ = signal.find_peaks(diff, distance=4)  # At least 1 second apart

    if len(peaks) == 0:
        return []

    # Sort by magnitude and take top ones
    sorted_peaks = sorted(peaks, key=lambda p: diff[p], reverse=True)
    top_peaks = sorted(sorted_peaks[:n_segments-1])

    # Convert to time
    segment_times = [float(p * hop / sr) for p in top_peaks]

    return segment_times


def analyze_track(filepath: Path) -> dict:
    """Analyze a music track and return timing data."""
    print(f"Analyzing {filepath.name}...")

    # Load audio
    y, sr = load_wav(filepath)
    duration_ms = int(len(y) / sr * 1000)

    # Compute onset envelope
    onset_env, onset_times = compute_onset_envelope(y, sr)

    # Detect tempo and beats
    bpm, beat_times = detect_beats(onset_env, onset_times, sr)
    beats_ms = [int(t * 1000) for t in beat_times]

    # Detect onsets (significant sound events)
    threshold = np.mean(onset_env) + 0.5 * np.std(onset_env)
    onset_peaks, _ = signal.find_peaks(onset_env, height=threshold, distance=5)
    onsets_ms = [int(onset_times[p] * 1000) for p in onset_peaks]

    # Find high-energy moments
    energy_peak_times = detect_energy_peaks(onset_env, onset_times)
    energy_peaks_ms = [int(t * 1000) for t in energy_peak_times]

    # Detect structural changes
    segment_times = detect_segments(y, sr)
    segments_ms = [int(t * 1000) for t in segment_times]

    # Downbeats (first beat of each measure - assuming 4/4)
    downbeat_indices = list(range(0, len(beats_ms), 4))
    downbeats_ms = [beats_ms[i] for i in downbeat_indices if i < len(beats_ms)]

    # Calculate BPM subdivisions for level design
    beat_interval_ms = int(60000 / bpm) if bpm > 0 else 500

    return {
        "file": filepath.name,
        "duration_ms": duration_ms,
        "bpm": round(bpm, 1),
        "beat_interval_ms": beat_interval_ms,
        "bpm_div_2": round(bpm / 2, 1),  # Half-time feel
        "bpm_div_4": round(bpm / 4, 1),  # Quarter-time
        "bpm_div_8": round(bpm / 8, 1),  # Eighth-time
        "beats_ms": beats_ms,
        "downbeats_ms": downbeats_ms,
        "onsets_ms": onsets_ms[:100],  # Limit to first 100 for readability
        "energy_peaks_ms": energy_peaks_ms,
        "segments_ms": segments_ms,
        "notes": {
            "beats_ms": "Regular beat timestamps - use for basic enemy spawns",
            "downbeats_ms": "First beat of each measure (4/4) - use for emphasis/bosses",
            "energy_peaks_ms": "High-energy moments - great for swift_cinder bursts!",
            "segments_ms": "Structural changes - new sections, drums kick in, etc.",
            "bpm_div_4": "Use this for flare_serpent/basalt_guardian wave frequency",
        }
    }


def main():
    music_dir = Path(__file__).parent / "fire_nation_attacked" / "audio" / "44k"
    output_file = Path(__file__).parent / "music_analysis.json"

    # Find all level music files
    level_files = sorted(music_dir.glob("lvl*.wav"))

    if not level_files:
        print(f"No level music found in {music_dir}")
        return

    results = {}
    for filepath in level_files:
        level_name = filepath.stem  # e.g., "lvl3"
        results[level_name] = analyze_track(filepath)

    # Write results
    with open(output_file, "w") as f:
        json.dump(results, f, indent=2)

    print(f"\nAnalysis complete! Results written to {output_file}")

    # Print summary
    print("\n" + "="*60)
    print("SUMMARY")
    print("="*60)
    for name, data in results.items():
        print(f"\n{name}.wav:")
        print(f"  Duration: {data['duration_ms']/1000:.1f}s | BPM: {data['bpm']}")
        print(f"  Recommended wave BPM: {data['bpm_div_4']} (BPM/4) or {data['bpm_div_8']} (BPM/8)")
        if data['energy_peaks_ms']:
            peaks_str = ", ".join([f"{p/1000:.1f}s" for p in data['energy_peaks_ms'][:5]])
            print(f"  Energy peaks: {peaks_str}{'...' if len(data['energy_peaks_ms']) > 5 else ''}")
        if data['segments_ms']:
            segs_str = ", ".join([f"{s/1000:.1f}s" for s in data['segments_ms'][:5]])
            print(f"  Segments: {segs_str}{'...' if len(data['segments_ms']) > 5 else ''}")


if __name__ == "__main__":
    main()