# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "librosa>=0.10.0",
#     "numpy",
#     "scipy",
#     "soundfile",
# ]
# ///

import librosa
import numpy as np
import sys
import json
import os
from pathlib import Path

def analyze_track(audio_path):
    """Analyze a single track and return features."""
    print(f"  Analyzing: {os.path.basename(audio_path)}")
    
    y, sr = librosa.load(audio_path, sr=22050, duration=180)
    
    # Tempo
    tempo, _ = librosa.beat.beat_track(y=y, sr=sr)
    tempo = float(tempo) if not hasattr(tempo, '__len__') else float(tempo[0])
    
    # Energy
    rms = librosa.feature.rms(y=y)[0]
    energy = float(np.mean(rms))
    energy_std = float(np.std(rms))
    
    # Spectral features
    spectral_centroid = librosa.feature.spectral_centroid(y=y, sr=sr)[0]
    brightness = float(np.mean(spectral_centroid))
    
    # Harmonic/Percussive
    y_harmonic, y_percussive = librosa.effects.hpss(y)
    harmonic_ratio = float(np.sum(np.abs(y_harmonic)) / (np.sum(np.abs(y)) + 1e-6))
    
    # Chroma
    chroma = librosa.feature.chroma_stft(y=y, sr=sr)
    chroma_mean = float(np.mean(chroma))
    
    # Onset strength
    onset_env = librosa.onset.onset_strength(y=y, sr=sr)
    rhythmic_density = float(np.mean(onset_env))
    
    # Zero crossing rate
    zcr = librosa.feature.zero_crossing_rate(y)[0]
    percussiveness = float(np.mean(zcr))
    
    return {
        "tempo_bpm": tempo,
        "energy_rms": energy,
        "energy_std": energy_std,
        "brightness_hz": brightness,
        "harmonic_ratio": harmonic_ratio,
        "chroma_mean": chroma_mean,
        "rhythmic_density": rhythmic_density,
        "percussiveness": percussiveness,
    }

def normalize_features(raw):
    """Convert raw features to 0-1 Genre Universe scales."""
    # Tempo: 60-180 BPM
    tempo_norm = np.clip((raw["tempo_bpm"] - 60) / 120, 0, 1)
    
    # Energy: RMS 0.01-0.25
    energy_norm = np.clip(raw["energy_rms"] / 0.2, 0, 1)
    
    # Brightness: 1000-4000 Hz
    brightness_norm = np.clip((raw["brightness_hz"] - 1000) / 3000, 0, 1)
    
    # Organic score
    organic_norm = np.clip(raw["harmonic_ratio"] * 1.2 - brightness_norm * 0.2, 0, 1)
    
    # Valence estimate (rough)
    valence_norm = np.clip(0.25 + brightness_norm * 0.25 + raw["chroma_mean"] * 0.3, 0, 1)
    
    # Danceability
    dance_tempo = 1 - abs(raw["tempo_bpm"] - 120) / 60
    danceability = np.clip(dance_tempo * 0.4 + raw["rhythmic_density"] * 0.35 + energy_norm * 0.25, 0, 1)
    
    # Acousticness (related to organic but different)
    acousticness = np.clip(raw["harmonic_ratio"] * 0.7 + (1 - brightness_norm) * 0.3, 0, 1)
    
    # Production density
    prod_density = np.clip((1 - raw["harmonic_ratio"]) * 0.5 + raw["energy_std"] * 8 + energy_norm * 0.3, 0, 1)
    
    return {
        "position": {
            "valence": round(float(valence_norm), 2),
            "tempo": round(float(tempo_norm), 2),
            "organic": round(float(organic_norm), 2),
        },
        "spikes": {
            "energy": round(float(energy_norm), 2),
            "acousticness": round(float(acousticness), 2),
            "danceability": round(float(danceability), 2),
            "production_density": round(float(prod_density), 2),
        }
    }

if __name__ == "__main__":
    tracks_dir = Path("surya_tracks")
    mp3_files = list(tracks_dir.glob("*.mp3"))
    
    if not mp3_files:
        print("No MP3 files found!")
        sys.exit(1)
    
    print(f"\n🎵 ANALYZING {len(mp3_files)} TRACKS FROM SURYA ALBUM\n")
    print("="*60)
    
    all_raw = []
    track_results = []
    
    for mp3 in sorted(mp3_files):
        raw = analyze_track(str(mp3))
        all_raw.append(raw)
        normalized = normalize_features(raw)
        track_results.append({
            "track": mp3.stem,
            "raw": raw,
            "normalized": normalized
        })
        print(f"    Tempo: {raw['tempo_bpm']:.1f} BPM | Energy: {raw['energy_rms']:.3f} | Harmonic: {raw['harmonic_ratio']:.2f}")
    
    # Average all tracks
    print("\n" + "="*60)
    print("📊 ALBUM AVERAGES")
    print("="*60)
    
    avg_raw = {
        key: np.mean([t[key] for t in all_raw])
        for key in all_raw[0].keys()
    }
    
    print(f"\n  Average Tempo: {avg_raw['tempo_bpm']:.1f} BPM")
    print(f"  Average Energy: {avg_raw['energy_rms']:.4f}")
    print(f"  Average Brightness: {avg_raw['brightness_hz']:.1f} Hz")
    print(f"  Average Harmonic Ratio: {avg_raw['harmonic_ratio']:.3f}")
    
    avg_normalized = normalize_features(avg_raw)
    
    print("\n" + "="*60)
    print("🎯 RECOMMENDED DAS POSITION (Album Average)")
    print("="*60)
    
    pos = avg_normalized["position"]
    spikes = avg_normalized["spikes"]
    
    print(f"\n  Position in Genre Universe:")
    print(f"    valence: {pos['valence']} (mood: {'sad' if pos['valence'] < 0.4 else 'happy' if pos['valence'] > 0.6 else 'bittersweet'})")
    print(f"    tempo:   {pos['tempo']} ({'slow' if pos['tempo'] < 0.4 else 'fast' if pos['tempo'] > 0.6 else 'medium'})")
    print(f"    organic: {pos['organic']} ({'electronic' if pos['organic'] < 0.4 else 'organic' if pos['organic'] > 0.6 else 'balanced'})")
    
    print(f"\n  Spike Values:")
    for key, val in spikes.items():
        bar = "█" * int(val * 20) + "░" * (20 - int(val * 20))
        print(f"    {key:20} [{bar}] {val}")
    
    # Output JSON for programmatic use
    result = {
        "tracks_analyzed": len(mp3_files),
        "track_names": [mp3.stem for mp3 in sorted(mp3_files)],
        "raw_averages": {k: round(v, 4) for k, v in avg_raw.items()},
        "recommended_position": pos,
        "recommended_spikes": spikes,
    }
    
    print("\n" + "="*60)
    print("📄 JSON OUTPUT")
    print("="*60)
    print(json.dumps(result, indent=2))