# 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 librosa
import librosa.display
from scipy import spatial
from typing import Optional, Tuple

# Ensure compatibility with numpy versions that do not have np.long
if not hasattr(np, 'long'):
    np.long = int

#### 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 Tonal_Fragment:
    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]:
        # Define profiles
        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]),
            # Add other profiles here
        }
        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}")

        # Compute correlations
        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',
    interpolation: str = 'linear',
) -> 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
        interpolation: type of interpolation for time codes ('linear' or 'none')

    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 = Tonal_Fragment(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]

    if interpolation == 'linear':
        df_segments['time'] = np.linspace(0, len(df_segments) * hoplen, len(df_segments))

    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

