import spikeextractors as se
from .filterrecording import FilterRecording
import numpy as np
from scipy import special
import scipy.fft
import math


class BandpassFilterRecording(FilterRecording):
    def __init__(self, *, recording: se.RecordingExtractor, freq_min, freq_max, freq_wid):
        self._padding = 3000
        target_ram = 100 * 1000 * 1000
        target_chunk_size = math.ceil(
            min(
                recording.get_sampling_frequency() * 30,
                target_ram / (recording.get_num_channels() * 4)
            )
        )
        # It's important that the fft's have size 2^x. So we prepare the chunk sizes to have size 2^x - 2*padding
        chunk_size = int(2 ** math.ceil(np.log2(target_chunk_size)) - self._padding * 2)
        FilterRecording.__init__(self, recording=recording, chunk_size=chunk_size)
        self._params = dict(
            name='bandpass_filter',
            freq_min=freq_min,
            freq_max=freq_max,
            freq_wid=freq_wid
        )
        self._recording = recording

    def paramsForHash(self):
        return self._params

    def filterChunk(self, *, start_frame, end_frame):
        i1 = start_frame - self._padding
        i2 = end_frame + self._padding
        padded_chunk = self._read_chunk(i1, i2)
        if i1 < 0:
            for m in range(padded_chunk.shape[0]):
                padded_chunk[m, :-i1] = padded_chunk[m, -i1]
        if i2 > self._recording.get_num_frames():
            for m in range(padded_chunk.shape[0]):
                aa = (i2 - self._recording.get_num_frames())
                padded_chunk[m, -aa:] = padded_chunk[m, aa - 1]
        filtered_padded_chunk = self._do_filter(padded_chunk)
        return filtered_padded_chunk[:, start_frame - i1:end_frame - i1]

    def _create_filter_kernel(self, N, samplerate, freq_min, freq_max, freq_wid=1000):
        # Matches ahb's code /matlab/processors/ms_bandpass_filter.m
        # improved ahb, changing tanh to erf, correct -3dB pts  6/14/16
        T = N / samplerate  # total time
        df = 1 / T  # frequency grid
        relwid = 3.0  # relative bottom-end roll-off width param, kills low freqs by factor 1e-5.

        k_inds = np.arange(0, N)
        k_inds = np.where(k_inds <= (N + 1) / 2, k_inds, k_inds - N)

        fgrid = df * k_inds
        absf = np.abs(fgrid)

        val = np.ones(fgrid.shape)
        if freq_min != 0:
            val = val * (1 + special.erf(relwid * (absf - freq_min) / freq_min)) / 2  # pylint: disable=no-member
            val = np.where(np.abs(k_inds) < 0.1, 0, val)  # kill DC part exactly
        if freq_max != 0:
            val = val * (1 - special.erf((absf - freq_max) / freq_wid)) / 2  # pylint: disable=no-member
        val = np.sqrt(val)  # note sqrt of filter func to apply to spectral intensity not ampl
        return val

    def _do_filter(self, chunk):
        samplerate = self._recording.get_sampling_frequency()
        M = chunk.shape[0]
        chunk2 = chunk
        # Subtract off the mean of each channel unless we are doing only a low-pass filter
        # if self._params['freq_min']!=0:
        #    for m in range(M):
        #        chunk2[m,:]=chunk2[m,:]-np.mean(chunk2[m,:])
        # Do the actual filtering with a DFT with real input
        chunk_fft = scipy.fft.rfft(chunk2, workers=1)
        kernel = self._create_filter_kernel(
            chunk2.shape[1],
            samplerate,
            self._params['freq_min'], self._params['freq_max'], self._params['freq_wid']
        )
        kernel = kernel[0:chunk_fft.shape[1]]  # because this is the DFT of real data
        chunk_fft = chunk_fft * np.tile(kernel, (M, 1))
        chunk_filtered = scipy.fft.irfft(chunk_fft, workers=1)
        return chunk_filtered

    def _read_chunk(self, i1, i2):
        M = len(self._recording.get_channel_ids())
        N = self._recording.get_num_frames()
        if i1 < 0:
            i1b = 0
        else:
            i1b = i1
        if i2 > N:
            i2b = N
        else:
            i2b = i2
        ret = np.zeros((M, i2 - i1))
        ret[:, i1b - i1:i2b - i1] = self._recording.get_traces(start_frame=i1b, end_frame=i2b)
        return ret


def bandpass_filter(recording, freq_min=300, freq_max=6000, freq_wid=1000, resample=None):
    return BandpassFilterRecording(
        recording=recording,
        freq_min=freq_min,
        freq_max=freq_max,
        freq_wid=freq_wid,
    )
