import os
import numpy as np
import pandas as pd
import re
import matplotlib.pyplot as plt
from scipy.spatial.distance import squareform
from statsmodels.tsa.stattools import acf, pacf


def get_stimtimes(sfn):
    """Reads in stimulus times
    
    :param str sfn: stimulation file name ()
    """

    header = ['EventType', 'TimeStamp', 'Duration', 'Channels']

    stim_data = pd.read_csv(sfn, skiprows=12, comment=';', header=None, names=header, index_col=False,
                            usecols=[1, 2, 3, 4])

    stimtimes = stim_data[stim_data['EventType'] == 'S  1']['TimeStamp'].as_matrix()

    return stimtimes


# Read in and file utilities
def read_eeg(fname_without_extension, path='', n_channels=27):
    """Reads in EEG data files from *.eeg files

    :param str fname_without_extension: file name without extension
    :param str path: path to datafile location
    :param int n_channels: number of recording-channels in the file
    :return: [n, n_channels] numpy array of data
    :type return: numpy.ndarray
    """
    data =\
        np.fromfile(path + fname_without_extension + '.eeg', dtype='int16')
    n = len(data) / n_channels
    data = data.reshape([n, n_channels])
    return data


def read_vhdr_lines(fname_without_extension, path=''):
    """Reads in .vhdr file

    :param str fname_without_extension: file name without extension
    :param str path: path to file location
    :return: list with file rows as elements (str)
    :type return: list of str
    """
    with open(path + fname_without_extension + '.vhdr', 'r') as f:
        lines = [line.replace('\r\n', '') for line in f]
    return lines


def get_channel_dict(lines):
    """Gets the channel names and numbers from a list of the .vhdr rows

    :param list[str] lines: .vhdr file contents (generated with read_vhdr_lines function)
    :return: dictionary with channel numbers as keys and channel names as values
    :type return: dict
    """
    channels = [line.split(',')[0].split('=') for line in lines if line[:2] == 'Ch']
    channel_dict = {int(ch[0][2:]): ch[1] for ch in channels if re.match("Ch([1-9])", ch[0])}
    return channel_dict


def get_channel_name_dict(fname, path):
    '''Returns a dictionary with channel names as keys and numbers as values

    :param str fname: string with filenames
    :param str path: string with path to file ('/' in the end)
    :return:
    '''
    lines = read_vhdr_lines(fname, path)
    inv_dict = get_channel_dict(lines)
    ch_dict = {inv_dict[i]: i - 1 for i in inv_dict.keys()}
    return ch_dict


def get_channel_number(fname_without_extension, path):
    """Returns the channel number for a given file

    :param str fname_without_extension: file name without extension
    :param str path: path to file location
    :return: number of channels for the file
    :type return: list[int]
    """
    lines = read_vhdr_lines(fname_without_extension, path)
    ch_dict = get_channel_dict(lines)
    return len(ch_dict.keys())


def get_stim_filenames(path="./"):
    """Returns a list of eeg file names (with header and stimulation files) without file extension.

    :params str path: folder path to data files ("/" in the end needed)
    :return: list with eeg file names in the folder given by path (has also .vhdr and .mkdr counterpart)
    :type return: list[str]
    """
    file_list = os.listdir(path)
    eeg_file_list = [f[:-4] for f in file_list if f[-3:] == 'eeg']
    eeg_files_with_description = [f for f in eeg_file_list if (f + '.vhdr' in file_list and f + '.vmrk' in file_list)]
    stimlens = {filename: len(get_stimtimes(path + filename + '.vmrk')) for filename in eeg_files_with_description}
    eeg_files_with_stimuli = [s for s in stimlens.keys() if stimlens[s] > 0]
    return eeg_files_with_stimuli


def get_stimperiods(fname_without_extension, path='', data=None, filter_lb=0):
    """Cuts and concatenates signal from stimulation periods

    :param str fname_without_extension: filename without extension
    :param str path:                    path to file (with '/'-s and also '/' in the end of the path)
    :param numpy.ndarray data:          data array, if None then it tries to read data from folder
                                        with given filename
    :param float filter_lb: only keeps stimulation-periods with stimulus_freq > filter_lb
    :return: concatenated signal from stimulation periods (the interstimulation parts are dropped)
            and chunk endtimes (unusually indexing starts from one!!!)
    """
    if data is None:
        # Read in time-series data from file
        eeg_data = read_in_data(fname_without_extension, path,
                     get_channel_number(fname_without_extension, path))
    else:
        eeg_data = data

    # Get stimulus start and end times
    stims = get_stimtimes(path + fname_without_extension + '.vmrk')
    stims = _hp_filter_stims(stims, filter_lb)
    diffstim = np.diff(stims)
    ends = np.where(diffstim > 1000)[0]
    starts = ends + 1

    ends = ends.tolist() + [-1]
    starts = [0] + starts.tolist()

    chunks = [eeg_data[stims[starts[i]]:stims[ends[i]]+1, :] for i in range(len(ends))]
    endpoints = np.cumsum([len(i) for i in chunks])
    stim_data = np.concatenate(chunks, axis=0)
    return stim_data, endpoints


def _hp_filter_stims(stims, l_bound):
    """Filters out low frequency stimuli which is smaller than l_bound

    :param numpy.ndarray stims: stimuli times (1D array of int-s)
    :param float l_bound: lower bound on stimulus frequency
    :return: stimuli higher frequency than l_bound
    :rtype: numpy.ndarray
    """
    freqs = convert2freq(stims)
    valids = np.where(freqs > l_bound)[0]
    return stims[range(valids[0] - 1, valids[-1] + 2)]