"""
Arrange spike times into bins by stimulus frame times
"""
from pathlib import Path
import numpy as np
import json


__all__ = [
    'binspikes',
]


def binspikes(path):
    """
    Arrange spike times into bins by stimulus frame times

    Parameters
    ----------
    path : str or pathlib.Path
        Path to the data directory containing parameters, stimulus frame
        times, and spike times of all cells

    Returns
    -------
    runningbin : (k, r, c) numpy.ndarray
        Binned spikes for the non-repeating "running" stimulus parts,
        where k is the number of trials, r is the number of bins per
        part, and c is the number of cells

    frozenbin : (k, f, c) numpy.ndarray
        Binned spikes for the repeating "frozen" stimulus parts

    moviebin : (k, m, c) numpy.ndarray
        Binned spikes for the repeating "movie" stimulus parts. Only
        returned if the experiment contained the movie stimulation
    """
    # Load stimulus parameters from disk
    path = Path(path).expanduser()
    filepath = path/'parameters.json'
    p = json.loads(filepath.read_text())

    # Determine length (number of bins) of one trial
    len_run = p.setdefault('RunningFrames', 1500) * p['Nblinks']
    len_frz = p.setdefault('FrozenFrames', 0) * p['Nblinks']
    len_mov = p.setdefault('MovieFrames', 0) * p.setdefault('Nblinksmovie', 0)
    len_trial = len_run + len_frz + len_mov

    # Load stimulus frame times and spike times from disk
    ft_file = path/'frametimes.txt'
    sp_files = path.glob('spiketimes_cell*.txt')
    sp_files = sorted(sp_files, key=lambda x: int(x.name[15:-4]))  # Sort files
    frametimes = np.loadtxt(ft_file)
    spiketimes = [np.loadtxt(file) for file in sp_files]
    num_cells = len(spiketimes)

    # Add the monitor delay to the stimulus frame times
    filepath = path/'../expdata.json'
    e = json.loads(filepath.read_text())
    frametimes += e['projector']['delay']

    # Check frame time frequency with respect to monitor refresh rate
    ftimes_freq = 1000 / (np.diff(frametimes).mean() * 1000)
    pulses = int(round(e['projector']['refreshrate'] / ftimes_freq))

    # Bin the spike times by stimulus frame times
    spkbin = np.zeros((frametimes.size * pulses, num_cells), dtype='uint8')
    frametimes_ext = np.append(frametimes, frametimes[-1])  # Add last bin edge
    for i in range(num_cells):
        spkbin[::pulses, i] = np.histogram(spiketimes[i], frametimes_ext)[0]

    # Reshape binned spikes by trial
    num_trials = int(frametimes.size*pulses / len_trial)
    total_len = num_trials * len_trial
    spkbin = spkbin[:total_len]
    spkbin = spkbin.reshape(num_trials, len_trial, num_cells)

    # Separate running and frozen parts
    runningbin = spkbin[:, :len_run, :]
    frozenbin = spkbin[:, len_run:len_run+len_frz, :]
    moviebin = spkbin[:, len_run+len_frz:, :]

    # Adjust bins to blinks
    shape_r = (num_trials, p['RunningFrames'], p['Nblinks'], num_cells)
    shape_f = (num_trials, p['FrozenFrames'], p['Nblinks'], num_cells)
    shape_m = (num_trials, p['MovieFrames'], p['Nblinksmovie'], num_cells)
    runningbin = runningbin.reshape(shape_r).sum(axis=2, dtype='uint8')
    frozenbin = frozenbin.reshape(shape_f).sum(axis=2, dtype='uint8')
    moviebin = moviebin.reshape(shape_m).sum(axis=2, dtype='uint8')

    if p['MovieFrames']:
        return runningbin, frozenbin, moviebin
    elif p['FrozenFrames']:
        return runningbin, frozenbin
    else:
        return runningbin,  # Return tuple for consistency


if __name__ == '__main__':
    print('\nThis file is a python module. Import it to call the function:')
    print('  >>> from binspikes import binspikes')
    print('  >>> path = \'20220208_60pMEA_right_i1/FrozenNoise_1/\'')
    print('  >>> runningbin, frozenbin = binspikes(path)\n')
    print('Details:')
    print(binspikes.__doc__)