import numba as nb
import numpy as np


@nb.njit()
def find_nearest_stft_bin(f0_, freqs):
    freqs = np.expand_dims(freqs, 0)
    f0_ = np.expand_dims(f0_, 1)
    return np.abs(freqs - f0_).argmin()


@nb.njit()
def get_med_curve(f0, step_size=20):
    v_begin = -1
    v_end = -1
    x_med_curve = []
    y_med_curve = []
    T = len(f0)

    for i in range(T):
        if f0[i] >= 50 and i < T - 1:
            if v_begin == -1:
                v_begin = i
            v_end = i
        else:
            if v_end != -1:
                if v_end - v_begin > 3:
                    for j in range(v_begin, v_end + 1 - step_size, step_size):
                        frag_med = np.median(f0[j:j + step_size])
                        x_med_curve.append(j)
                        y_med_curve.append(frag_med)
                    x_med_curve.append(v_end)
                    y_med_curve.append(np.median(f0[v_end - step_size:v_end + 1]))
            v_end = v_begin = -1
    x_med_curve = [0] + x_med_curve + [T]
    x_med_curve = np.array(x_med_curve)
    y_med_curve = [y_med_curve[0]] + y_med_curve + [y_med_curve[-1]]
    y_med_curve = np.array(y_med_curve)
    return x_med_curve, y_med_curve


@nb.njit()
def clean_short_v_frag(f0):
    v_begin = -1
    T = len(f0)

    uv = np.zeros_like(f0).astype(np.bool_)
    for i in range(T):
        if f0[i] >= 1e-4 and i < T - 1:
            if v_begin == -1:
                v_begin = i
        else:
            if v_begin != -1:
                v_end = i if f0[i] >= 1e-4 else i - 1
                if v_end - v_begin + 1 < 3:
                    uv[v_begin:v_end + 1] = 1
            v_begin = -1
    return uv


@nb.njit()
def find_best_f0_using_har_energy(spec, pitches, freqs, hars, hars_mhalf, f0_min, f0_max):
    re = np.zeros_like(spec)
    T = len(spec)
    for i in range(T):
        spec_i = spec[i]
        for j, f0_j in enumerate(pitches[i]):
            if f0_j == 0 or f0_j < f0_min[i] or f0_j > f0_max[i]:
                continue
            mask = np.zeros((10000,))
            mask_mhalf = np.zeros((10000,))
            for mul in hars:
                b = find_nearest_stft_bin(np.array((f0_j * mul,)), freqs)
                for delta in range(-1, 2):
                    mask[b + delta] = 1
            for mul in hars_mhalf:
                b_mhalf = find_nearest_stft_bin(np.array((f0_j * (mul - 0.5),)), freqs)
                for delta in range(-1, 2):
                    mask_mhalf[b_mhalf + delta] = 1
            mask = mask[:len(spec_i)]
            mask_mhalf = mask_mhalf[:len(spec_i)]
            energy = (np.exp(spec_i) * mask).sum() / mask.sum()
            energy_mhalf = (np.exp(spec_i) * mask_mhalf).sum() / mask_mhalf.sum()
            re[i, j] = energy / energy_mhalf
    f0_2d_mask = 10000 * (re > 2) + 20000 * (re > 3) + np.expand_dims(np.arange(re.shape[1])[::-1], 0)
    f0_idx = np.zeros((T,), dtype=np.int_)
    for i in range(T):
        f0_idx[i] = f0_2d_mask[i].argmax()
    uv = re.sum(-1) == 0

    f0 = np.zeros((T,))
    for i in range(T):
        f0[i] = pitches[i, f0_idx[i]]
    f0 = f0 * (1 - uv)
    uv = clean_short_v_frag(f0)
    f0[uv] = 0
    x_med_curve, y_med_curve = get_med_curve(f0)
    re = re * (re > 1.5)
    return re, f0, x_med_curve, y_med_curve