from ibllib.io import spikeglx
import numpy as np
import ibllib.dsp as dsp
from scipy import signal
from ibllib.misc import print_progress
from pathlib import Path
import alf.io as aio
import logging
import ibllib.ephys.ephysqc as ephysqc
from phylib.io import alf

_logger = logging.getLogger('ibllib')


RMS_WIN_LENGTH_SECS = 3
WELCH_WIN_LENGTH_SAMPLES = 1024


def rmsmap(fbin, spectra=True):
    """
    Computes RMS map in time domain and spectra for each channel of Neuropixel probe

    :param fbin: binary file in spike glx format (will look for attached metatdata)
    :type fbin: str or pathlib.Path
    :param spectra: whether to compute the power spectrum (only need for lfp data)
    :type: bool
    :return: a dictionary with amplitudes in channeltime space, channelfrequency space, time
     and frequency scales
    """
    if not isinstance(fbin, spikeglx.Reader):
        sglx = spikeglx.Reader(fbin)
    rms_win_length_samples = 2 ** np.ceil(np.log2(sglx.fs * RMS_WIN_LENGTH_SECS))
    # the window generator will generates window indices
    wingen = dsp.WindowGenerator(ns=sglx.ns, nswin=rms_win_length_samples, overlap=0)
    # pre-allocate output dictionary of numpy arrays
    win = {'TRMS': np.zeros((wingen.nwin, sglx.nc)),
           'nsamples': np.zeros((wingen.nwin,)),
           'fscale': dsp.fscale(WELCH_WIN_LENGTH_SAMPLES, 1 / sglx.fs, one_sided=True),
           'tscale': wingen.tscale(fs=sglx.fs)}
    win['spectral_density'] = np.zeros((len(win['fscale']), sglx.nc))
    # loop through the whole session
    for first, last in wingen.firstlast:
        D = sglx.read_samples(first_sample=first, last_sample=last)[0].transpose()
        # remove low frequency noise below 1 Hz
        D = dsp.hp(D, 1 / sglx.fs, [0, 1])
        iw = wingen.iw
        win['TRMS'][iw, :] = dsp.rms(D)
        win['nsamples'][iw] = D.shape[1]
        if spectra:
            # the last window may be smaller than what is needed for welch
            if last - first < WELCH_WIN_LENGTH_SAMPLES:
                continue
            # compute a smoothed spectrum using welch method
            _, w = signal.welch(D, fs=sglx.fs, window='hanning', nperseg=WELCH_WIN_LENGTH_SAMPLES,
                                detrend='constant', return_onesided=True, scaling='density',
                                axis=-1)
            win['spectral_density'] += w.T
        # print at least every 20 windows
        if (iw % min(20, max(int(np.floor(wingen.nwin / 75)), 1))) == 0:
            print_progress(iw, wingen.nwin)
    return win


def extract_rmsmap(fbin, out_folder=None, spectra=True):
    """
    Wrapper for rmsmap that outputs _ibl_ephysRmsMap and _ibl_ephysSpectra ALF files

    :param fbin: binary file in spike glx format (will look for attached metatdata)
    :param out_folder: folder in which to store output ALF files. Default uses the folder in which
     the `fbin` file lives.
    :param spectra: whether to compute the power spectrum (only need for lfp data)
    :type: bool
    :return: None
    """
    _logger.info(f"Computing QC for {fbin}")
    sglx = spikeglx.Reader(fbin)
    # check if output ALF files exist already:
    if out_folder is None:
        out_folder = Path(fbin).parent
    else:
        out_folder = Path(out_folder)
    alf_object_time = f'_iblqc_ephysTimeRms{sglx.type.upper()}'
    alf_object_freq = f'_iblqc_ephysSpectralDensity{sglx.type.upper()}'

    # crunch numbers
    rms = rmsmap(fbin, spectra=spectra)
    # output ALF files, single precision with the optional label as suffix before extension
    if not out_folder.exists():
        out_folder.mkdir()
    tdict = {'rms': rms['TRMS'].astype(np.single), 'timestamps': rms['tscale'].astype(np.single)}
    aio.save_object_npy(out_folder, object=alf_object_time, dico=tdict)
    if spectra:
        fdict = {'power': rms['spectral_density'].astype(np.single),
                 'freqs': rms['fscale'].astype(np.single)}
        aio.save_object_npy(out_folder, object=alf_object_freq, dico=fdict)


def _sample2v(ap_file):
    """
    Convert raw ephys data to Volts
    """
    md = spikeglx.read_meta_data(ap_file.with_suffix('.meta'))
    s2v = spikeglx._conversion_sample2v_from_meta(md)
    return s2v['ap'][0]


def ks2_to_alf(ks_path, bin_path, out_path, bin_file=None, ampfactor=1, label=None, force=True):
    """
    Convert Kilosort 2 output to ALF dataset for single probe data
    :param ks_path:
    :param bin_path: path of raw data
    :param out_path:
    :return:
    """
    m = ephysqc.phy_model_from_ks2_path(ks2_path=ks_path, bin_path=bin_path, bin_file=bin_file)
    ephysqc.spike_sorting_metrics_ks2(ks_path, m, save=True)
    ac = alf.EphysAlfCreator(m)
    ac.convert(out_path, label=label, force=force, ampfactor=ampfactor)


def extract_data(ks_path, ephys_path, out_path):
    efiles = spikeglx.glob_ephys_files(ephys_path)

    for efile in efiles:
        if efile.get('ap') and efile.ap.exists():
            ks2_to_alf(ks_path, ephys_path, out_path, bin_file=efile.ap,
                       ampfactor=_sample2v(efile.ap), label=None, force=True)

            extract_rmsmap(efile.ap, out_folder=out_path, spectra=False)
        if efile.get('lf') and efile.lf.exists():
            extract_rmsmap(efile.lf, out_folder=out_path)


# if __name__ == '__main__':
#
#    ephys_path = Path('C:/Users/Mayo/Downloads/raw_ephys_data')
#    ks_path = Path('C:/Users/Mayo/Downloads/KS2')
#    out_path = Path('C:/Users/Mayo/Downloads/alf')
#    extract_data(ks_path, ephys_path, out_path)

[((12, 10, 12, 37), 'logging.getLogger', 'logging.getLogger', ({(12, 28, 12, 36): '"""ibllib"""'}, {}), "('ibllib')", False, 'import logging\n'), ((34, 13, 34, 85), 'ibllib.dsp.WindowGenerator', 'dsp.WindowGenerator', (), '', True, 'import ibllib.dsp as dsp\n'), ((76, 11, 76, 32), 'ibllib.io.spikeglx.Reader', 'spikeglx.Reader', ({(76, 27, 76, 31): 'fbin'}, {}), '(fbin)', False, 'from ibllib.io import spikeglx\n'), ((91, 4, 91, 71), 'alf.io.save_object_npy', 'aio.save_object_npy', (), '', True, 'import alf.io as aio\n'), ((103, 10, 103, 53), 'ibllib.io.spikeglx._conversion_sample2v_from_meta', 'spikeglx._conversion_sample2v_from_meta', ({(103, 50, 103, 52): 'md'}, {}), '(md)', False, 'from ibllib.io import spikeglx\n'), ((115, 8, 115, 95), 'ibllib.ephys.ephysqc.phy_model_from_ks2_path', 'ephysqc.phy_model_from_ks2_path', (), '', True, 'import ibllib.ephys.ephysqc as ephysqc\n'), ((116, 4, 116, 60), 'ibllib.ephys.ephysqc.spike_sorting_metrics_ks2', 'ephysqc.spike_sorting_metrics_ks2', (), '', True, 'import ibllib.ephys.ephysqc as ephysqc\n'), ((117, 9, 117, 31), 'phylib.io.alf.EphysAlfCreator', 'alf.EphysAlfCreator', ({(117, 29, 117, 30): 'm'}, {}), '(m)', False, 'from phylib.io import alf\n'), ((122, 13, 122, 50), 'ibllib.io.spikeglx.glob_ephys_files', 'spikeglx.glob_ephys_files', ({(122, 39, 122, 49): 'ephys_path'}, {}), '(ephys_path)', False, 'from ibllib.io import spikeglx\n'), ((31, 15, 31, 36), 'ibllib.io.spikeglx.Reader', 'spikeglx.Reader', ({(31, 31, 31, 35): 'fbin'}, {}), '(fbin)', False, 'from ibllib.io import spikeglx\n'), ((36, 19, 36, 51), 'numpy.zeros', 'np.zeros', ({(36, 28, 36, 50): '(wingen.nwin, sglx.nc)'}, {}), '((wingen.nwin, sglx.nc))', True, 'import numpy as np\n'), ((37, 23, 37, 47), 'numpy.zeros', 'np.zeros', ({(37, 32, 37, 46): '(wingen.nwin,)'}, {}), '((wingen.nwin,))', True, 'import numpy as np\n'), ((38, 21, 38, 86), 'ibllib.dsp.fscale', 'dsp.fscale', (), '', True, 'import ibllib.dsp as dsp\n'), ((45, 12, 45, 42), 'ibllib.dsp.hp', 'dsp.hp', ({(45, 19, 45, 20): 'D', (45, 22, 45, 33): '1 / sglx.fs', (45, 35, 45, 41): '[0, 1]'}, {}), '(D, 1 / sglx.fs, [0, 1])', True, 'import ibllib.dsp as dsp\n'), ((47, 29, 47, 39), 'ibllib.dsp.rms', 'dsp.rms', ({(47, 37, 47, 38): 'D'}, {}), '(D)', True, 'import ibllib.dsp as dsp\n'), ((81, 21, 81, 37), 'pathlib.Path', 'Path', ({(81, 26, 81, 36): 'out_folder'}, {}), '(out_folder)', False, 'from pathlib import Path\n'), ((95, 8, 95, 75), 'alf.io.save_object_npy', 'aio.save_object_npy', (), '', True, 'import alf.io as aio\n'), ((32, 42, 32, 80), 'numpy.log2', 'np.log2', ({(32, 50, 32, 79): '(sglx.fs * RMS_WIN_LENGTH_SECS)'}, {}), '(sglx.fs * RMS_WIN_LENGTH_SECS)', True, 'import numpy as np\n'), ((54, 19, 56, 40), 'scipy.signal.welch', 'signal.welch', (), '', False, 'from scipy import signal\n'), ((60, 12, 60, 43), 'ibllib.misc.print_progress', 'print_progress', ({(60, 27, 60, 29): 'iw', (60, 31, 60, 42): 'wingen.nwin'}, {}), '(iw, wingen.nwin)', False, 'from ibllib.misc import print_progress\n'), ((79, 21, 79, 31), 'pathlib.Path', 'Path', ({(79, 26, 79, 30): 'fbin'}, {}), '(fbin)', False, 'from pathlib import Path\n'), ((59, 33, 59, 59), 'numpy.floor', 'np.floor', ({(59, 42, 59, 58): '(wingen.nwin / 75)'}, {}), '(wingen.nwin / 75)', True, 'import numpy as np\n')]