import librosa
import numpy as np


# 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']


def find_key(path):
    """chroma-based key detection using krumhansl-schmuckler profiles"""
    track, sr = librosa.load(path, sr=None)

    # pull chroma energy, average across time
    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 get_tempo(path):
    track, sr = librosa.load(path, sr=None)
    tempo, _ = librosa.beat.beat_track(y=track, sr=sr)
    # librosa sometimes returns an array
    if hasattr(tempo, '__len__'):
        return float(tempo[0])
    return float(tempo)


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


def fingerprint(path):
    track, sr = librosa.load(path, sr=None, mono=False)

    channels = 1 if track.ndim == 1 else track.shape[0]

    # reload mono for analysis
    if channels > 1:
        mono = librosa.to_mono(track)
    else:
        mono = track

    key_sig = find_key(path)
    bpm = get_tempo(path)
    length = librosa.get_duration(y=mono, sr=sr)

    return {
        'key': key_sig,
        'bpm': round(bpm, 1),
        'duration': round(length, 2),
        'sample_rate': 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"duration: {info['duration']}s")
    print(f"sample rate: {info['sample_rate']}Hz")
    print(f"channels: {info['channels']}")