"""key / tempo / meter detection. pure librosa dsp — no ML."""
import librosa
import numpy as np
import soundfile as sf

# krumhansl-schmuckler key profiles
# major and minor correlation vectors for pitch class distribution
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])

PITCH_CLASSES = ['C', 'C#', 'D', 'D#', 'E', 'F',
                 'F#', 'G', 'G#', 'A', 'A#', 'B']

# all the analysis here works on envelopes and chroma; 22k mono is plenty
ANALYSIS_SR = 22050


def _key_from_audio(track, sr):
    """chroma-based key detection using krumhansl-schmuckler profiles"""
    chroma = librosa.feature.chroma_cqt(y=track, sr=sr)
    pitch_dist = np.mean(chroma, axis=1)

    # normalize
    pitch_dist = (pitch_dist - pitch_dist.mean()) / (pitch_dist.std() + 1e-8)

    best_corr = -2
    best_key = 'C major'

    for shift in range(12):
        rolled = np.roll(pitch_dist, -shift)

        # check major
        major_norm = (MAJOR_PROFILE - MAJOR_PROFILE.mean()) / MAJOR_PROFILE.std()
        corr_major = np.corrcoef(rolled, major_norm)[0, 1]
        if corr_major > best_corr:
            best_corr = corr_major
            best_key = f'{PITCH_CLASSES[shift]} major'

        # check minor
        minor_norm = (MINOR_PROFILE - MINOR_PROFILE.mean()) / MINOR_PROFILE.std()
        corr_minor = np.corrcoef(rolled, minor_norm)[0, 1]
        if corr_minor > best_corr:
            best_corr = corr_minor
            best_key = f'{PITCH_CLASSES[shift]} minor'

    return best_key


def _scalar_tempo(tempo):
    # librosa sometimes returns an array
    if hasattr(tempo, '__len__'):
        return float(tempo[0])
    return float(tempo)


def _time_signature_from_beats(onset_env, beats):
    """
    estimate the meter from where the accents land: take the onset strength
    at each tracked beat and score how well a "downbeat every N beats" grid
    lines up with the loud ones, for N = 3 and 4. waltz time has to win
    clearly — ambiguous material is called 4/4, like most music.
    """
    if len(beats) < 12:
        return "4/4"

    strengths = onset_env[beats].astype(float)
    spread = strengths.std()
    if spread < 1e-8:
        return "4/4"
    strengths = (strengths - strengths.mean()) / spread

    def accent_score(meter):
        # best phase alignment of the downbeat grid
        return max(float(np.mean(strengths[offset::meter]))
                   for offset in range(meter))

    three, four = accent_score(3), accent_score(4)
    if three > 0 and three > four * 1.25:
        return "3/4"
    return "4/4"


def find_key(path):
    track, sr = librosa.load(path, sr=ANALYSIS_SR, mono=True)
    return _key_from_audio(track, sr)


def get_tempo(path):
    track, sr = librosa.load(path, sr=ANALYSIS_SR, mono=True)
    tempo, _ = librosa.beat.beat_track(y=track, sr=sr)
    return _scalar_tempo(tempo)


def get_time_signature(path):
    track, sr = librosa.load(path, sr=ANALYSIS_SR, mono=True)
    onset_env = librosa.onset.onset_strength(y=track, sr=sr)
    _, beats = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr)
    return _time_signature_from_beats(onset_env, beats)


def get_duration(path):
    return librosa.get_duration(path=path)


def fingerprint(path):
    """one-stop analysis: key, bpm, meter, duration — single decode pass."""
    track, sr = librosa.load(path, sr=ANALYSIS_SR, mono=True)

    onset_env = librosa.onset.onset_strength(y=track, sr=sr)
    tempo, beats = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr)

    # source-file metadata without a second full decode
    try:
        meta = sf.info(path)
        native_sr, channels = meta.samplerate, meta.channels
    except Exception:
        native_sr, channels = sr, 1

    return {
        'key': _key_from_audio(track, sr),
        'bpm': round(_scalar_tempo(tempo), 1),
        'time_signature': _time_signature_from_beats(onset_env, beats),
        'duration': round(len(track) / sr, 2),
        'sample_rate': native_sr,
        'channels': channels,
    }


if __name__ == '__main__':
    import sys
    if len(sys.argv) < 2:
        print('usage: python analyze.py <audio_file>')
        sys.exit(1)

    info = fingerprint(sys.argv[1])
    print(f"key: {info['key']}")
    print(f"bpm: {info['bpm']}")
    print(f"time signature: {info['time_signature']} (estimated)")
    print(f"duration: {info['duration']}s")
    print(f"sample rate: {info['sample_rate']}Hz")
    print(f"channels: {info['channels']}")