import numpy as np
import math


class MelSpectrogram:
    def __init__(self, n_filt, freq_low, freq_high, frame_size, fs, n_fft):
        self.n_filt = n_filt
        self.n_fft = n_fft
        self.freq_low = freq_low
        self.freq_high = freq_high
        self.frame_size = frame_size
        self.fs = fs

        self.filt_bank = self.build_filterbank(freq_low, freq_high, n_filt, n_fft, fs)
        self.mel_power = np.zeros(n_filt, dtype=np.float32)
        self.mel_spectrogram_image = np.zeros((n_filt, n_filt), dtype=np.float32)

    def build_filterbank(self, l, h, n_filt, n_fft, fs):
        lower_mel = 1125 * math.log(1 + l / 700)
        higher_mel = 1125 * math.log(1 + h / 700)

        rows = math.floor(n_fft / 2)
        columns = n_filt
        filterbank = np.zeros((rows, columns), dtype=np.float32)

        diff = (higher_mel - lower_mel) / (n_filt + 1)
        mel_band = lower_mel
        freq_band = 700 * (math.exp(mel_band / 1125) - 1)

        f = np.zeros(n_filt + 2, dtype=np.int32)
        f[0] = math.floor((n_fft + 1) * freq_band / fs)

        for i in range(1, n_filt + 2):
            mel_band = mel_band + diff
            freq_band = 700 * (math.exp(mel_band / 1125) - 1)
            f[i] = math.floor((n_fft + 1) * freq_band / fs)

        for k in range(1, rows + 1):
            for m in range(1, columns + 1):
                if (k > f[m - 1]) and (k <= f[m]):
                    filterbank[k - 1][m - 1] = float(k - f[m - 1]) / float(
                        f[m] - f[m - 1]
                    )
                elif (k > f[m]) and (k <= f[m + 1]):
                    filterbank[k - 1][m - 1] = float(f[m + 1] - k) / float(
                        f[m + 1] - f[m]
                    )

        return filterbank

    def mel_calculate(self, fft_data):
        for i in range(self.n_filt):
            sum_val = 0
            for j in range(math.floor(self.n_fft / 2)):
                sum_val += self.filt_bank[j][i] * fft_data[j]

            self.mel_power[i] = math.log(sum_val + 1e-8)

    def mel_image_create(self):
        # Shift the 2-d image up
        for i in range(self.n_filt - 1):
            self.mel_spectrogram_image[i] = np.copy(self.mel_spectrogram_image[i + 1])

        # Add the current mel-spectrogram power to the bottom row
        self.mel_spectrogram_image[self.n_filt - 1] = np.copy(self.mel_power)

    def update_image(self, fft_data):
        self.mel_calculate(fft_data)
        self.mel_image_create()