#!/usr/bin/env python3
"""Batch sonic/spectral analysis of all audio files in a directory.
Outputs a summary and saves per-track JSON results.

Usage: python3 batch_analyze.py "/path/to/music/folder"
"""

import os, sys, json
import numpy as np
import librosa
from collections import Counter

def analyze_folder(music_dir: str):
    exts = ('.mp3', '.flac', '.m4a', '.ogg', '.wav')
    files = sorted(f for f in os.listdir(music_dir) if f.lower().endswith(exts))
    if not files:
        print("No audio files found.")
        return

    results = []
    for i, f in enumerate(files):
        path = os.path.join(music_dir, f)
        try:
            y, sr = librosa.load(path, sr=22050, duration=45)
            tempo = librosa.feature.rhythm.tempo(y=y, sr=sr)[0]
            rms = float(np.nanmean(librosa.feature.rms(y=y)[0]))
            centroid = float(np.nanmean(librosa.feature.spectral_centroid(y=y, sr=sr)[0]))
            zcr = float(np.nanmean(librosa.feature.zero_crossing_rate(y=y)[0]))
            dur = float(librosa.get_duration(y=y, sr=sr))
            chroma = librosa.feature.chroma_cqt(y=y, sr=sr).mean(axis=1)
            keys = ['C','C#','D','D#','E','F','F#','G','G#','A','A#','B']
            key = keys[int(np.argmax(chroma))]
            results.append({
                'file': f,
                'tempo': round(tempo, 1),
                'rms': round(rms, 4),
                'centroid': round(centroid, 1),
                'zcr': round(zcr, 4),
                'duration': round(dur, 1),
                'key': key,
            })
        except Exception as e:
            results.append({'file': f, 'error': str(e)})

    good = [r for r in results if 'tempo' in r]
    bad = [r for r in results if 'error' in r]

    tempos = [r['tempo'] for r in good]
    energies = [r['rms'] for r in good]
    centroids = [r['centroid'] for r in good]
    durations = [r['duration'] for r in good]
    keys = [r['key'] for r in good]
    kc = Counter(keys)

    folder_name = os.path.basename(music_dir.rstrip('/'))
    print(f"=== {folder_name} — SONIC ANALYSIS ===")
    print(f"Total files: {len(files)} | Analyzed: {len(good)} | Errors: {len(bad)}")
    if bad:
        for r in bad:
            print(f"  ✗ {r['file']}: {r['error']}")

    print(f"\n--- COLLECTIVE STATS ---")
    print(f"Tempo:     mean={np.mean(tempos):.1f}  std={np.std(tempos):.1f}  min={min(tempos):.1f}  max={max(tempos):.1f}")
    print(f"Energy:    mean={np.mean(energies):.4f}  std={np.std(energies):.4f}")
    print(f"Brightness (centroid): mean={np.mean(centroids):.0f} Hz  std={np.std(centroids):.0f} Hz")
    print(f"Duration:  mean={np.mean(durations):.1f}s  total={sum(durations)/60:.1f} min")

    print(f"\n--- KEY DISTRIBUTION ---")
    for k, _ in kc.most_common():
        print(f"  {k}: {kc[k]}")

    print(f"\n--- FASTEST ---")
    for r in sorted(good, key=lambda x: x['tempo'], reverse=True)[:5]:
        print(f"  {r['tempo']:.0f} BPM — {r['file']}")

    print(f"\n--- SLOWEST ---")
    for r in sorted(good, key=lambda x: x['tempo'])[:5]:
        print(f"  {r['tempo']:.0f} BPM — {r['file']}")

    print(f"\n--- HIGHEST ENERGY ---")
    for r in sorted(good, key=lambda x: x['rms'], reverse=True)[:5]:
        print(f"  {r['rms']:.4f} — {r['file']}")

    print(f"\n--- BRIGHTEST (highest centroid) ---")
    for r in sorted(good, key=lambda x: x['centroid'], reverse=True)[:5]:
        print(f"  {r['centroid']:.0f} Hz — {r['file']}")

    print(f"\n--- DARKEST (lowest centroid) ---")
    for r in sorted(good, key=lambda x: x['centroid'])[:5]:
        print(f"  {r['centroid']:.0f} Hz — {r['file']}")

    out_path = f"/tmp/{folder_name.replace(' ', '_')}_sonic_results.json"
    with open(out_path, "w") as fh:
        json.dump(results, fh, indent=2, ensure_ascii=False)
    print(f"\n\nFull results saved to {out_path}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} /path/to/music/folder")
        sys.exit(1)
    analyze_folder(sys.argv[1])