# relative mode: A estimation of relative mode from audio.
# part of the initial code has been derived from:
# https://github.com/jackmcarthur/musical-key-finder
#
# Full documentation of this model is available from
# [https://github.com/tuomaseerola/relative_mode](https://github.com/tuomaseerola/relative_mode)

#### Libraries ####
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import librosa
import librosa.display
from scipy import spatial
from scipy.interpolate import interp1d
import math
import heapq
from typing import Optional, Tuple

#### Helper Functions ####
def calculate_cosine_distance(a: np.ndarray, b: np.ndarray) -> float:
    return float(spatial.distance.cosine(a, b))

def calculate_cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    return 1 - calculate_cosine_distance(a, b)

def calculate_euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    return float(spatial.distance.euclidean(a, b))

def calculate_euclidean_similarity(a: np.ndarray, b: np.ndarray) -> float:
    return 1 - calculate_euclidean_distance(a, b)

#### Tonal Fragment Class ####
class TonalFragment:
    def __init__(
        self,
        waveform: np.ndarray,
        sr: int,
        tstart: Optional[float] = None,
        tend: Optional[float] = None,
        chromatype: str = 'CQT',
        profile: str = 'krumhansl',
        octaves: int = 7,
        thres: float = 0.0,
        bins_octave: int = 36,
        wlensmooth: int = 41,
        f_min: float = 65.4,
        distance: str = 'pearson',
    ):
        self.waveform = waveform
        self.sr = sr
        self.tstart = librosa.time_to_samples(tstart, sr=sr) if tstart else None
        self.tend = librosa.time_to_samples(tend, sr=sr) if tend else None
        self.chromatype = chromatype
        self.profile = profile
        self.octaves = octaves
        self.thres = thres
        self.bins_octave = bins_octave
        self.win_len_smooth = wlensmooth
        self.distance = distance
        self.y_segment = self.waveform[self.tstart:self.tend]

        self.chromograph = self._compute_chromograph(f_min)
        self.chroma_vals = [np.sum(self.chromograph[i]) for i in range(12)]
        self.keyfreqs = self._compute_key_frequencies()
        self.key_dict, self.key, self.bestcorr = self._compute_key_correlations()

    def _compute_chromograph(self, f_min: float) -> np.ndarray:
        if self.chromatype == 'CENS':
            return librosa.feature.chroma_cens(
                y=self.y_segment, sr=self.sr, fmin=f_min, n_octaves=self.octaves,
                bins_per_octave=36, hop_length=8192, win_len_smooth=self.win_len_smooth
            )
        elif self.chromatype == 'STFT':
            return librosa.feature.chroma_stft(y=self.y_segment, sr=self.sr)
        elif self.chromatype == 'CQT':
            return librosa.feature.chroma_cqt(
                y=self.y_segment, sr=self.sr, n_octaves=self.octaves, threshold=self.thres,
                fmin=f_min, bins_per_octave=36, cqt_mode='hybrid', hop_length=8192
            )
        else:
            raise ValueError(f"Unsupported chromatype: {self.chromatype}")

    def _compute_key_frequencies(self) -> dict:
        pitches = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        return {pitches[i]: self.chroma_vals[i] for i in range(12)}

    def _compute_key_correlations(self) -> Tuple[dict, str, float]:
        profiles = {
            'krumhansl': ([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88],
                          [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17]),
            'aarden': ([17.7661, 0.145624, 14.9265, 0.160186, 19.8049, 11.3587, 0.291248, 22.062, 0.145624, 8.15494, 0.232998, 4.95122],
                       [18.2648, 0.737619, 14.0499, 16.8599, 0.702494, 14.4362, 0.702494, 18.6161, 4.56621, 1.93186, 7.37619, 1.75623]),
        }
        maj_profile, min_profile = profiles.get(self.profile, (None, None))
        if not maj_profile or not min_profile:
            raise ValueError(f"Unsupported profile: {self.profile}")

        pitches = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        keys = [p + ' major' for p in pitches] + [p + ' minor' for p in pitches]
        maj_corrs, min_corrs = [], []
        for i in range(12):
            key_test = [self.keyfreqs.get(pitches[(i + m) % 12]) for m in range(12)]
            if self.distance == 'pearson':
                maj_corrs.append(np.corrcoef(maj_profile, key_test)[1, 0])
                min_corrs.append(np.corrcoef(min_profile, key_test)[1, 0])
            elif self.distance == 'cosine':
                maj_corrs.append(calculate_cosine_similarity(maj_profile, key_test))
                min_corrs.append(calculate_cosine_similarity(min_profile, key_test))
            else:
                raise ValueError(f"Unsupported distance: {self.distance}")

        key_dict = {**{keys[i]: maj_corrs[i] for i in range(12)},
                    **{keys[i + 12]: min_corrs[i] for i in range(12)}}
        best_key = max(key_dict, key=key_dict.get)
        best_corr = max(key_dict.values())
        return key_dict, best_key, best_corr

#### Main Functions ####
def relative_mode(
    y: np.ndarray,
    sr: int,
    winlen: int = 3,
    hoplen: int = 3,
    distance: str = 'cosine',
    profile: str = 'albrecht',
    chromatype: str = 'CENS',
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Run relative mode estimation across analysis windows.

    Arguments:
        y: audio
        sr: sampling rate
        winlen: analysis window length in seconds
        hoplen: analysis window hop length in seconds
        distance: distance measure (string), either cosine, pearson, or euclidean
        profile: key profile used (string), either krumhansl, albrecht, simple, aarden
        chromatype: type of chroma extraction (string), either CQT, CENS, HPCC

    Returns:
        df: relative mode scalar value (-1 to +1).
        df_segments: relative mode scalar value (-1 to +1) for each window with time codes.
    """
    t = librosa.estimate_tuning(y=y, sr=sr)
    y440 = librosa.effects.pitch_shift(y, sr=sr, n_steps=-t)
    df_segments = pd.DataFrame(columns=['onset', 'time', 'tonmaxmaj', 'tonmaxmin', 'tonkey', 'tondeltamax'])
    frames = librosa.util.frame(y440, frame_length=int(sr * winlen), hop_length=int(sr * hoplen))
    for ii in range(frames.shape[1]):
        ton = TonalFragment(frames[:, ii], sr, distance=distance, profile=profile, chromatype=chromatype)
        time_code = ii * hoplen
        df_segments.loc[len(df_segments)] = [ii, time_code, ton.max_maj, ton.max_min, ton.key, ton.maj_min_delta_max]

    df = pd.DataFrame({
        'tonmaxmaj': [np.mean(df_segments['tonmaxmaj'])],
        'tonmaxmin': [np.mean(df_segments['tonmaxmin'])],
        'tondeltamax': [np.mean(df_segments['tondeltamax'])],
        'tondeltamaxMd': [np.median(df_segments['tondeltamax'])],
        'tondeltamaxSi': [sum(np.sign(df_segments['tondeltamax']))],
    })

    return df, df_segments

