# -*- coding: utf-8 -*-
import numpy as np
from scipy import signal
import mne
from mne.io.constants import FIFF
from matplotlib import pyplot as plt
from matplotlib import patches
import time
import csv
import sys

class wave_class(object):
    def __init__(self):
        self.spike_pn = []
        self.sharp_pn = []
        self.biphasic_pn = []
        self.spike_np = []
        self.sharp_np = []
        self.biphasic_np = []
        self.triphasic_npn = []
        self.triphasic_pnp = []
        self.polyphasic = []

def wave_analyze(raw):
    correct_kind(raw)
    picks_EEG = mne.pick_types(raw.info, meg=False, eeg=True, eog=False, stim=False, exclude='bads')
    raw.filter(1/(2*np.pi*0.3)/2, 70.0, picks=picks_EEG, method='fir')
    raw_pattern, picks_pattern = apply_pattern(raw, picks_EEG)
#    raw_pattern.plot()
    wave_all, wave_dict_all, discharge_time, peak_dics_slow, stds_all = waveform_analyze(raw_pattern, picks_pattern)
    return wave_all, wave_dict_all, discharge_time, raw_pattern, picks_pattern, peak_dics_slow, stds_all

def waveform_analyze(raw, picks):
    sfreq = raw.info['sfreq']
    channel_names = np.array(raw.info['ch_names'])[picks]
    raw_slow = raw.copy()
    raw_slow.filter(0.1, 15.0, picks=picks, method='fir')
    raw_fast = raw.copy()
    raw_fast.filter(14.0, 70.0, picks=picks, method='fir')

    data_slow = raw_slow.get_data(picks)
    data_slow *= 1e6
    data_fast = raw_fast.get_data(picks)
    data_fast *= 1e6

    peak_inds_slow = data_maxima_ind(data_slow)
    peak_inds_fast = data_maxima_ind(data_fast)

    peak_dics_slow, stds_all = peak_classify(data_slow, sfreq, peak_inds_slow, std_factor=0.6, smooth_count=1)
    wave_all = find_wave(peak_dics_slow, stds_all, data_fast, peak_inds_fast, sfreq)
    wave_dict_all = generate_dict(wave_all, channel_names)
    discharge_time=generate_times(wave_all)

#    large_inds_slow = generate_large_peak_ind(peak_dics_slow)
#    plot_data(raw_slow, picks, plot_names=['F8-AV'], gap=150, start=84, stop=94, peak_ind=peak_inds_slow, large_peak_ind=large_inds_slow, show_detail=True, line_colors=['k'])
    return wave_all, wave_dict_all, discharge_time, peak_dics_slow, stds_all

def correct_kind(raw):
    chs = raw.info['chs']
    for i in range(len(chs)):
        kind = chs[i]['kind']
        name = chs[i]['ch_name']
        if kind is FIFF.FIFFV_EEG_CH:
            if (u'STI' in name) or (u'sti' in name):
                chs[i]['kind'] = FIFF.FIFFV_STIM_CH
                continue
            elif (u'EOG' in name) or (u'eog' in name):
                chs[i]['kind'] = FIFF.FIFFV_EOG_CH
                continue
            elif (u'EMG' in name) or (u'emg' in name):
                chs[i]['kind'] = FIFF.FIFFV_EMG_CH
                continue
            elif (u'ECG' in name) or (u'ecg' in name):
                chs[i]['kind'] = FIFF.FIFFV_ECG_CH
                continue
            elif (u'MISC' in name) or (u'misc' in name):
                chs[i]['kind'] = FIFF.FIFFV_MISC_CH
                continue
            elif (u'DC' in name) or (u'dc' in name) or (name==u'POL E') or (name==u'EEG E') or (name==u'E') or (name==u'e'):
                chs[i]['kind'] = FIFF.FIFFV_SYST_CH
            elif (u'Event' in name) or (u'event' in name) or (u'Marker' in name) or (u'marker' in name):
                chs[i]['kind'] = FIFF.FIFFV_MISC_CH
                continue
    raw.info['chs'] = chs
    return

def apply_pattern(raw, picks_EEG):
    name_list_AV = ['Fp1','Fp2','F3','F4','C3','C4','P3','P4','O1','O2','F7','F8','T3','T4','T5','T6','Fz','Cz','Pz','T1','T2','A1','A2']
    further_pick_to_AV = further_pick_by_name(raw, picks_EEG, name_list_AV)
    picks_AV = picks_EEG[further_pick_to_AV]
    name_list_move_AV = ['Fp1','Fp2','F3','F4','C3','C4','P3','P4','O1','O2','F7','F8','T3','T4','T5','T6','Fz','Cz','Pz','A1','A2']
    further_pick_to_move_AV = further_pick_by_name(raw, picks_EEG, name_list_move_AV)
    picks_move_AV = picks_EEG[further_pick_to_move_AV]

    '''
    pick_C3 = picks_EEG[further_pick_by_name(raw, picks_EEG, ['C3'])]
    pick_C4 = picks_EEG[further_pick_by_name(raw, picks_EEG, ['C4'])]
    pick_T1 = picks_EEG[further_pick_by_name(raw, picks_EEG, ['T1'])]
    pick_T2 = picks_EEG[further_pick_by_name(raw, picks_EEG, ['T2'])]
    '''

    data_AV = np.mean(raw.get_data(picks_AV),axis=0)
    data = raw.get_data(picks_move_AV)
    data_mat_AV = np.dot(np.array([data_AV]).T, np.ones((1,data.shape[0]))).T
    raw_new = raw.copy()
    raw_new._data[picks_move_AV] = -(data-data_mat_AV)
    add_name_tail(raw_new, picks_move_AV, '-AV')

    '''
    raw_new._data[pick_T1] = -(raw._data[pick_T1]-raw._data[pick_C3])
    add_name_tail(raw_new, pick_T1, '-C3')
    raw_new._data[pick_T2] = -(raw._data[pick_T2]-raw._data[pick_C4])
    add_name_tail(raw_new, pick_T2, '-C4')
    '''
    picks_new = picks_AV
    return raw_new, picks_new

def add_name_tail(raw, picks, tail):
    for i in picks:
        raw.info['ch_names'][i] += tail
        raw.info['chs'][i]['ch_name'] += tail
    return

def further_pick_by_name(raw, picks, channel_names):
    further_picks = []
    all_names = np.array(raw.info['ch_names'])[picks]
    for j in range(len(all_names)):
        for this_name in channel_names:
            if '*' in this_name:
                this_name_key = this_name.split('*')[0]
                if (this_name_key in all_names[j]) and (all_names[j].split(this_name_key)[0]=='') and len(this_name_key)<len(all_names[j]):
                    further_picks.append(j)
                    break
            else:
                if this_name == all_names[j]:
                    further_picks.append(j)
                    break
    return np.array(further_picks)

def data_maxima_ind(data):
    n_picks = data.shape[0]
    maxima_inds = []
    for i in range(n_picks):
        maxima_P_ind = np.array(signal.argrelmax(data[i,:]))[0]
        maxima_N_ind = np.array(signal.argrelmin(data[i,:]))[0]
        maxima_ind = sorted(maxima_P_ind.tolist() + maxima_N_ind.tolist())
        if len(maxima_ind) == 0:
            maxima_ind.append(0)
        maxima_inds.append(np.array(maxima_ind))
    return maxima_inds

def peak_classify(data, freq, peak_inds, std_factor=1., smooth_count=1):
    n_channel = int(data.shape[0])
    peak_dics_all = []
    stds_all = []

    for i in range(n_channel):
        inds = np.array(peak_inds[i])
        peaks = data[i,inds]
        inds = inds.tolist()
        peaks = peaks.tolist()

        for count in range(smooth_count):
            inds_to_pop = []
            dones = [j<-1 for j in range(len(peaks))]
            for ii in range(len(peaks)):
                if dones[ii]:
                    pass
                elif ii<len(peaks)-3:
                    p0 = peaks[ii]
                    p1 = peaks[ii+1]
                    p2 = peaks[ii+2]
                    p3 = peaks[ii+3]
                    p1_line = p0 + float(inds[ii+1]-inds[ii])*(p3-p0)/float(inds[ii+3]-inds[ii])
                    p2_line = p0 + float(inds[ii+2]-inds[ii])*(p3-p0)/float(inds[ii+3]-inds[ii])
                    if max(abs(p1-p1_line),abs(p2-p2_line))<0.15*abs(p3-p0) and min(p1,p2)>min(p0,p3) and max(p1,p2)<max(p0,p3):
                        inds_to_pop += [ii+1,ii+2]
                        dones[ii+1:ii+3] = [True,True]
            for ii in range(len(inds_to_pop)):
                inds.pop(inds_to_pop[ii]-ii)
                peaks.pop(inds_to_pop[ii]-ii)

        inds = np.array(inds)
        peaks = np.array(peaks)
        means, stds = mean_std_around(data[i,:], inds, freq)
        peak_dics = []
        for ii in range(len(peaks)):
            mean = means[ii]
            std = stds[ii]
            peak_dic = {}
            peak_dic['type'] = 3
            peak_dic['done'] = False
            peak_dic['ind'] = inds[ii]
            peak_dic['time'] = inds[ii]/float(freq)
            peak_dic['peak'] = peaks[ii]
            if peak_dic['peak'] > mean+std_factor*std:
                peak_dic['position'] = 1
            elif peak_dic['peak'] < mean-std_factor*std:
                peak_dic['position'] = 2
            else:
                peak_dic['position'] = 0
            if ii > 1: #and max([abs(peak_dics[-2]['peak']),abs(peak_dics[-1]['peak']),abs(peak_dic['peak'])])>100:
                if peak_dics[-2]['peak'] > peak_dics[-1]['peak']:
                    peak_dics[-1]['type'] = which_v(peak_dics[-2],peak_dics[-1],peak_dic)
                else:
                    peak_dics[-1]['type'] = which_n(peak_dics[-2],peak_dics[-1],peak_dic)
            peak_dics.append(peak_dic)
        peak_dics_all.append(peak_dics)
        stds_all.append(stds)
    return peak_dics_all, stds_all

def generate_large_peak_ind(peak_dics_all):
    large_peak_inds = []
    for i in range(len(peak_dics_all)):
        peak_dics = peak_dics_all[i]
        this_large_peak_inds = []
        for j in range(len(peak_dics)):
            if peak_dics[j]['position'] > 0:
                this_large_peak_inds.append(peak_dics[j]['ind'])
        large_peak_inds.append(np.array(this_large_peak_inds))
    return large_peak_inds

def mean_std_around(data_array, inds, freq, mean_range=1, std_range=2):
    around_range_mean, around_range_std = mean_range*int(freq), std_range*int(freq)
    around_range = max([around_range_mean,around_range_std])
    inds_head = list(range(around_range))
    inds_head.reverse()
    data_head = data_array[inds_head]
    inds_tail = np.array(inds_head)+len(data_array)-around_range
    data_tail = data_array[inds_tail]
    data_extend = np.append(np.append(data_head,data_array),data_tail)
    datas_to_mean = []
    datas_to_std = []
    for ind in inds:
        data_to_mean = data_extend[(ind-around_range_mean+around_range):(ind+around_range_mean+around_range+1)]
        data_to_std = data_extend[(ind-around_range_std+around_range):(ind+around_range_std+around_range+1)]
        datas_to_mean.append(data_to_mean)
        datas_to_std.append(data_to_std)
    means = np.mean(datas_to_mean,axis=1)
    stds = np.std(datas_to_std,axis=1)
    return means, stds

def which_v(left,v,right):
    vtype = 0#'v_unknown'
    if v['position'] == 2:
        vtype = -1#'v_x2x'
    elif v['position'] == 1:
        vtype = -2#'v_101'
    else:
        if left['position'] == 0:
            if right['position'] == 0:
                vtype = -3#'v_000'
            elif right['position'] == 1:
                vtype = -4#'v_001'
        elif left['position'] == 1:
            if right['position'] == 0:
                vtype = -5#'v_100'
            elif right['position'] == 1:
                vtype = -2#'v_101'
    return vtype

def which_n(left,n,right):
    ntype = 0
    if n['position'] == 1:
        ntype = 1#'n_x1x'
    elif n['position'] == 2:
        ntype = 2#'n_202'
    else:
        if left['position'] == 0:
            if right['position'] == 0:
                ntype = 3#'n_000'
            elif right['position'] == 2:
                ntype = 4#'n_002'
        elif left['position'] == 2:
            if right['position'] == 0:
                ntype = 5#'n_200'
            elif right['position'] == 2:
                ntype = 2#'n_202'
    return ntype

def find_wave(peak_dics_all, stds_all, data_spike, inds_spike, freq=256.):
    n_channel = len(peak_dics_all)
    spike_limit = 70e-3*freq
    slow_limit = 200e-3*freq
    wave_all  = []
    for i in range(n_channel):
        wave_inst = wave_class()
        peaks_spike = data_spike[i,inds_spike[i]]
        ind_spike = inds_spike[i]
        dones = [j<-1 for j in range(len(peaks_spike))]
        for ii in range(len(peaks_spike)):
            amp_threshold = 100.
            try:
                if dones[ii]:
                    pass
                elif abs(peaks_spike[ii+1]-peaks_spike[ii])>amp_threshold and ii>0 and ii<len(peaks_spike)-1:
                    wave_inst.spike_pn.append(np.array((ind_spike[ii-1]/freq,ind_spike[ii+2]/freq))+analysis_start_time)
                    dones[ii] = True
                    dones[ii+1] = True
            except:
                Exception

        peak_dics = peak_dics_all[i]
        stds = stds_all[i]
        for ii in range(len(peak_dics)):
            amp_threshold = max([2*2*stds[ii],0.5*100.])
            try:
                if peak_dics[ii]['done']:
                    pass
                else:
                    if (peak_dics[ii]['type']==-4 or peak_dics[ii]['type']==-2 or peak_dics[ii]['type']==-1) \
                    and (peak_dics[ii+1]['type']==1) \
                    and (peak_dics[ii+2]['type']==-1) \
                    and (peak_dics[ii+3]['type']==5 or peak_dics[ii+3]['type']==2 or (peak_dics[ii+3]['type']==1 and (abs(peak_dics[ii+3]['peak']-peak_dics[ii+2]['peak'])<0.7*abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak']) or abs(peak_dics[ii+3]['peak']-peak_dics[ii]['peak'])<0.4*abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak']) ) )) \
                    and (peak_dics[ii+3]['ind']-peak_dics[ii]['ind']<60*slow_limit):
                        if peak_dics[ii+3]['ind']-peak_dics[ii]['ind']<spike_limit and abs(peak_dics[ii+1]['peak']-peak_dics[ii]['peak'])>=amp_threshold:
                            wave_inst.spike_pn.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+3]['ind']/freq))+analysis_start_time)
                            peak_dics[ii]['done'] = True
                            peak_dics[ii+1]['done'] = True
                            peak_dics[ii+2]['done'] = True
                            continue
                        elif (peak_dics[ii+1]['ind']-peak_dics[ii]['ind']<slow_limit) and (peak_slope(peak_dics[ii],peak_dics[ii+1])/peak_slope(peak_dics[ii+2],peak_dics[ii+3])>2 or peak_slope(peak_dics[ii],peak_dics[ii+1])/peak_slope(peak_dics[ii+1],peak_dics[ii+2])<-2) and (peak_dics[ii+1]['peak']-peak_dics[ii]['peak'])/(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])>0.6*0 and abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])>=0.75*amp_threshold:#float((peak_dics[ii+1]['ind']-peak_dics[ii]['ind']))/float((peak_dics[ii+3]['ind']-peak_dics[ii]['ind']))<0.28 and (peak_dics[ii+1]['peak']-peak_dics[ii]['peak'])/(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])>0.6:
                            wave_inst.sharp_pn.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+3]['ind']/freq))+analysis_start_time)
                            peak_dics[ii]['done'] = True
                            peak_dics[ii+1]['done'] = True
                            peak_dics[ii+2]['done'] = True
                            continue
                        elif (peak_dics[ii+1]['ind']-peak_dics[ii]['ind']<slow_limit) and abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])>=0.75*amp_threshold and peak_dics[ii+3]['ind']-peak_dics[ii]['ind']>slow_limit and abs(peak_dics[ii]['peak']-peak_dics[ii+3]['peak'])/abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])<0.3 and 0.5*(peak_dics[ii]['peak']+peak_dics[ii+3]['peak'])>min(peak_dics[ii+1]['peak'],peak_dics[ii+2]['peak'])+0.3*abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak']) and 0.5*(peak_dics[ii]['peak']+peak_dics[ii+3]['peak'])<min(peak_dics[ii+1]['peak'],peak_dics[ii+2]['peak'])+0.7*abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak']):
                            wave_inst.biphasic_pn.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+3]['ind']/freq))+analysis_start_time)
                            peak_dics[ii]['done'] = True
                            peak_dics[ii+1]['done'] = True
                            peak_dics[ii+2]['done'] = True
                            continue

                    if (peak_dics[ii]['type']==4 or peak_dics[ii]['type']==2 or peak_dics[ii]['type']==1) and (peak_dics[ii+1]['type']==-1) and (peak_dics[ii+2]['type']==1 and abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>75*2) and (peak_dics[ii+3]['type']==-1):
                        if (((peak_dics[ii+4]['type']==5) or peak_dics[ii+4]['type']==3 or (peak_dics[ii+4]['type']==2)) or (peak_dics[ii+4]['type']==1 and abs(peak_dics[ii+4]['peak']-peak_dics[ii+3]['peak'])<0.8*abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak']))) and abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>=0.75*amp_threshold and abs(peak_dics[ii]['peak']-peak_dics[ii+1]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>0.25 and abs(peak_dics[ii]['peak']-peak_dics[ii+1]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])<0.75 and abs(peak_dics[ii+3]['peak']-peak_dics[ii+4]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>0.25 and abs(peak_dics[ii+3]['peak']-peak_dics[ii+4]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])<0.75 and peak_dics[ii+3]['ind']-peak_dics[ii]['ind']>slow_limit and peak_dics[ii+1]['ind']-peak_dics[ii]['ind']>0.15*(peak_dics[ii+4]['ind']-peak_dics[ii]['ind']):
                            wave_inst.triphasic_npn.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+4]['ind']/freq))+analysis_start_time)
                            peak_dics[ii]['done'] = True
                            peak_dics[ii+1]['done'] = True
                            peak_dics[ii+2]['done'] = True
                            peak_dics[ii+3]['done'] = True
                            continue
                        elif (peak_dics[ii+1]['ind']-peak_dics[ii]['ind']<slow_limit) and peak_dics[ii+4]['type']==1 and abs(peak_dics[ii+4]['peak']-peak_dics[ii+3]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>0.6*0 and abs(peak_dics[ii+4]['peak']-peak_dics[ii+3]['peak'])>=0.75*amp_threshold and peak_dics[ii+3]['ind']-peak_dics[ii]['ind']>slow_limit:
                            peak_dics[ii]['done'] = True
                            iii = 0
                            while ((peak_dics[ii+iii+1]['type']==-1) or (peak_dics[ii+iii+1]['type']==1)) and (iii<2 or abs(peak_dics[ii+iii+1]['peak']-peak_dics[ii+iii]['peak'])>0.6*0*abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])):
                                peak_dics[ii+iii+1]['done'] = True
                                iii = iii+1
                            wave_inst.polyphasic.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+iii+1]['ind']/freq))+analysis_start_time)
                            continue

                    if (peak_dics[ii]['type']==4 or peak_dics[ii]['type']==2 or peak_dics[ii]['type']==1) and (peak_dics[ii+1]['type']==-1) and (peak_dics[ii+2]['type']==1) and (peak_dics[ii+3]['type']==-5 or peak_dics[ii+3]['type']==-2 or (peak_dics[ii+3]['type']==-1 and (abs(peak_dics[ii+3]['peak']-peak_dics[ii+2]['peak'])<0.7*abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak']) or abs(peak_dics[ii+3]['peak']-peak_dics[ii]['peak'])<0.4*abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak']) ) )) and (peak_dics[ii+3]['ind']-peak_dics[ii]['ind']<60*slow_limit):
                        if peak_dics[ii+3]['ind']-peak_dics[ii]['ind']<spike_limit and abs(peak_dics[ii+1]['peak']-peak_dics[ii]['peak'])>=amp_threshold:
                            wave_inst.spike_np.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+3]['ind']/freq))+analysis_start_time)
                            peak_dics[ii]['done'] = True
                            peak_dics[ii+1]['done'] = True
                            peak_dics[ii+2]['done'] = True
                            continue
                        elif (peak_dics[ii+1]['ind']-peak_dics[ii]['ind']<slow_limit) and (peak_slope(peak_dics[ii],peak_dics[ii+1])/peak_slope(peak_dics[ii+2],peak_dics[ii+3])>2 or peak_slope(peak_dics[ii],peak_dics[ii+1])/peak_slope(peak_dics[ii+1],peak_dics[ii+2])<-2) and (peak_dics[ii+1]['peak']-peak_dics[ii]['peak'])/(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])>0.6*0 and abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])>=0.75*amp_threshold:#float((peak_dics[ii+1]['ind']-peak_dics[ii]['ind']))/float((peak_dics[ii+3]['ind']-peak_dics[ii]['ind']))<0.28 and (peak_dics[ii+1]['peak']-peak_dics[ii]['peak'])/(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])>0.6:
                            wave_inst.sharp_np.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+3]['ind']/freq))+analysis_start_time)
                            peak_dics[ii]['done'] = True
                            peak_dics[ii+1]['done'] = True
                            peak_dics[ii+2]['done'] = True
                            continue
                        elif (peak_dics[ii+1]['ind']-peak_dics[ii]['ind']<slow_limit) and abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])>=0.75*amp_threshold and peak_dics[ii+3]['ind']-peak_dics[ii]['ind']>slow_limit and abs(peak_dics[ii]['peak']-peak_dics[ii+3]['peak'])/abs(peak_dics[ii+1]['peak']-peak_dics[ii+2]['peak'])<0.3 and 0.5*(peak_dics[ii]['peak']+peak_dics[ii+3]['peak'])>min(peak_dics[ii+1]['peak'],peak_dics[ii+2]['peak'])+0.3*abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak']) and 0.5*(peak_dics[ii]['peak']+peak_dics[ii+3]['peak'])<min(peak_dics[ii+1]['peak'],peak_dics[ii+2]['peak'])+0.7*abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak']):
                            wave_inst.biphasic_np.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+3]['ind']/freq))+analysis_start_time)
                            peak_dics[ii]['done'] = True
                            peak_dics[ii+1]['done'] = True
                            peak_dics[ii+2]['done'] = True
                            continue

                    if (peak_dics[ii]['type']==-4 or peak_dics[ii]['type']==-2 or peak_dics[ii]['type']==-1) and (peak_dics[ii+1]['type']==1) and (peak_dics[ii+2]['type']==-1 and abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>75*2) and (peak_dics[ii+3]['type']==1):
                        if (((peak_dics[ii+4]['type']==-5) or peak_dics[ii+4]['type']==-3 or (peak_dics[ii+4]['type']==-2)) or (peak_dics[ii+4]['type']==-1 and abs(peak_dics[ii+4]['peak']-peak_dics[ii+3]['peak'])<0.8*abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak']))) and abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>=0.75*amp_threshold and abs(peak_dics[ii]['peak']-peak_dics[ii+1]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>0.25 and abs(peak_dics[ii]['peak']-peak_dics[ii+1]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])<0.75 and abs(peak_dics[ii+3]['peak']-peak_dics[ii+4]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>0.25 and abs(peak_dics[ii+3]['peak']-peak_dics[ii+4]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])<0.75 and peak_dics[ii+3]['ind']-peak_dics[ii]['ind']>slow_limit and peak_dics[ii+1]['ind']-peak_dics[ii]['ind']>0.15*(peak_dics[ii+4]['ind']-peak_dics[ii]['ind']):
                            wave_inst.triphasic_pnp.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+4]['ind']/freq))+analysis_start_time)
                            peak_dics[ii]['done'] = True
                            peak_dics[ii+1]['done'] = True
                            peak_dics[ii+2]['done'] = True
                            peak_dics[ii+3]['done'] = True
                            continue
                        elif (peak_dics[ii+1]['ind']-peak_dics[ii]['ind']<slow_limit) and peak_dics[ii+4]['type']==-1 and abs(peak_dics[ii+4]['peak']-peak_dics[ii+3]['peak'])/abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])>0.6*0 and abs(peak_dics[ii+4]['peak']-peak_dics[ii+3]['peak'])>=0.75*amp_threshold and peak_dics[ii+3]['ind']-peak_dics[ii]['ind']>slow_limit:
                            peak_dics[ii]['done'] = True
                            iii = 0
                            while ((peak_dics[ii+iii+1]['type']==1) or (peak_dics[ii+iii+1]['type']==-1)) and (iii<2 or abs(peak_dics[ii+iii+1]['peak']-peak_dics[ii+iii]['peak'])>0.6*0*abs(peak_dics[ii+2]['peak']-peak_dics[ii+1]['peak'])):
                                peak_dics[ii+iii+1]['done'] = True
                                iii = iii+1
                            wave_inst.polyphasic.append(np.array((peak_dics[ii]['ind']/freq,peak_dics[ii+iii+1]['ind']/freq))+analysis_start_time)
                            continue
            except:
                Exception
        wave_all.append(wave_inst)
    return wave_all

def peak_slope(peak_dic_1, peak_dic_2):
    x_1 = float(peak_dic_1['ind'])
    x_2 = float(peak_dic_2['ind'])
    y_1 = peak_dic_1['peak']
    y_2 = peak_dic_2['peak']
    slope = (y_2-y_1)/(x_2-x_1)
    return slope

def generate_dict(wave_all, channel_names):
    wave_dict_all = {}
    for i in range(len(wave_all)):
        key = channel_names[i]
        wave_inst = wave_all[i]
        wave_dict = {}
        wave_dict['spike'] = wave_inst.spike_pn + wave_inst.spike_np
        wave_dict['sharp'] = wave_inst.sharp_pn + wave_inst.sharp_np
        wave_dict['biphasic'] = wave_inst.biphasic_pn + wave_inst.biphasic_np
        wave_dict['triphasic'] = wave_inst.triphasic_npn + wave_inst.triphasic_pnp
        wave_dict['polyphasic'] = wave_inst.polyphasic
        wave_dict_all[key] = wave_dict
    return wave_dict_all

def generate_times(wave_all):
    discharge_time = []
    for i in range(len(wave_all)):
        this_wave = wave_all[i]
        this_time = this_wave.sharp_pn + this_wave.sharp_np# + this_wave.spike_pn + this_wave.spike_np
        discharge_time.append(this_time)
    return discharge_time

def plot_data(raw, picks, plot_names, gap=1000, start=0., stop=None, peak_ind=None, threshold=None, large_peak_ind=None, HFO_peak_ind=None, show_detail=False, HFO_time=None, line_colors=['b']):
    if raw._times[-1] < original_total_time:
        need_shift = True
    else:
        need_shift = False
    further_picks = further_pick_by_name(raw, picks, plot_names)
    temp_raw = raw.copy()
    if need_shift:
        try:
            temp_raw.crop(start-analysis_start_time, stop-analysis_start_time)
        except Exception:
            print('\033[0;31mData (or part of data) at [%.2fs, %.2fs] has not been filtered.\033[0m'%(start,stop))
            return
        data, times = temp_raw[picks[further_picks]]
    else:
        temp_raw.crop(start, stop)
        data, times = temp_raw[picks[further_picks]]
    sfreq = raw.info['sfreq']
    names = np.array(raw.info['ch_names'])[picks[further_picks]]
    data = data*1e6
    n_plots = len(further_picks)
    if show_detail and (peak_ind is not None):
        mark_inds = []
        for i in range(n_plots):
            mark_ind = peak_ind[further_picks[i]]-(temp_raw.first_samp-raw.first_samp)
            mark_ind = mark_ind[np.where(mark_ind>=0)]
            mark_ind = mark_ind[np.where(mark_ind<data.shape[1])]
            mark_inds.append(mark_ind)
        if large_peak_ind is not None:
            large_mark_inds = []
            for i in range(n_plots):
                large_mark_ind = large_peak_ind[further_picks[i]]-(temp_raw.first_samp-raw.first_samp)
                large_mark_ind = large_mark_ind[np.where(large_mark_ind>=0)]
                large_mark_ind = large_mark_ind[np.where(large_mark_ind<data.shape[1])]
                large_mark_inds.append(large_mark_ind)
        if HFO_peak_ind is not None:
            HFO_mark_inds = []
            for i in range(n_plots):
                temp_HFO_ind = np.hstack(HFO_peak_ind[further_picks[i]]) if len(HFO_peak_ind[further_picks[i]])>0 else []
                HFO_mark_ind = peak_ind[further_picks[i]][temp_HFO_ind]-(temp_raw.first_samp-raw.first_samp)
                HFO_mark_ind = HFO_mark_ind[np.where(HFO_mark_ind>=0)]
                HFO_mark_ind = HFO_mark_ind[np.where(HFO_mark_ind<data.shape[1])]
                HFO_mark_inds.append(HFO_mark_ind)
    if (not show_detail) and (HFO_time is not None):
        HFO_mark_inds = []
        for i in range(n_plots):
            if need_shift:
                temp_HFO_ind = (sfreq*(np.hstack(HFO_time[further_picks[i]])-analysis_start_time)).astype(int) if len(HFO_time[further_picks[i]])>0 else []
            else:
                temp_HFO_ind = (sfreq*np.hstack(HFO_time[further_picks[i]])).astype(int) if len(HFO_time[further_picks[i]])>0 else []
            HFO_mark_ind = temp_HFO_ind - (temp_raw.first_samp-raw.first_samp)
            HFO_mark_ind = HFO_mark_ind[np.where(HFO_mark_ind>=0)]
            HFO_mark_ind = HFO_mark_ind[np.where(HFO_mark_ind<data.shape[1])]
            HFO_mark_inds.append(HFO_mark_ind)

    plt.figure(figsize=[14, 1.*n_plots])
    data_plot = plt.gca()
    if stop is None:
        stop = start+times[-1]
    data_plot.axis(xmin=start, xmax=stop, ymin=-0.5*gap, ymax=(n_plots-0.5)*gap)
    data_plot.set_yticks([])
    for i in range(n_plots):
        line_color = line_colors[i%len(line_colors)]
        data_plot.plot(times+start, data[i]+gap*(n_plots-1-i), color=line_color)
        data_plot.plot(times+start, 0*data[i]+gap*(n_plots-1-i-0.5), color='k', linewidth=0.2)
        data_plot.text(start, gap*(n_plots-1-i), names[i], horizontalalignment='right', verticalalignment='center', color=line_color)
        if (peak_ind is not None) and show_detail:
            data_plot.plot(times[mark_inds[i]]+start, data[i][mark_inds[i]]+gap*(n_plots-1-i),'c.')
        if (threshold is not None) and show_detail:
            data_plot.plot(times+start, 0*data[i]+gap*(n_plots-1-i)+threshold[further_picks[i]]*1e6, 'g--', linewidth=0.5)
            data_plot.plot(times+start, 0*data[i]+gap*(n_plots-1-i)-threshold[further_picks[i]]*1e6, 'g--', linewidth=0.5)
        if (large_peak_ind is not None) and show_detail:
            data_plot.plot(times[large_mark_inds[i]]+start, data[i][large_mark_inds[i]]+gap*(n_plots-1-i),'r.')
        if (HFO_peak_ind is not None) or (HFO_time is not None):
            for j in range(len(HFO_mark_inds[i])):
                data_plot.plot(np.ones(10)*times[HFO_mark_inds[i][j]]+start, data[i][HFO_mark_inds[i][j]]+np.linspace(gap*(n_plots-1.3-i),gap*(n_plots-0.7-i),10),'m', linewidth=2)
    data_plot.set_xlabel('time [s]')
    data_plot.text(start, -gap, 'spacing between channels: %.0f$\mu V$'%(gap), horizontalalignment='left', verticalalignment='center')
    plt.show()
    return data_plot

def key_for_names(pick_names, i):
    name = pick_names[i].split(' ')[-1].split('-')[0]
    tail = pick_names[i].split(' ')[-1].split('-')[-1] if len(pick_names[i].split(' ')[-1].split('-'))>1 else ""
    if "'" in name:
        alpha, number = name[:2], int(name[2:])
    else:
        alpha, number = name[:1], int(name[1:])
    key = alpha + chr(number) + tail
    return key

def get_hit_ratio(discharge_time, labels, record_start, tolerance=0.1):
    hit_sum = 0
    for i in np.arange(1,len(labels)):
        label_i = 60.*float(labels[i][0]) + float(labels[i][1]) - record_start
        hit = False
        for j in range(len(discharge_time)):
            this_discharge_time = discharge_time[j]
            for k in range(len(this_discharge_time)):
                if label_i >= this_discharge_time[k][0]-tolerance and label_i <= this_discharge_time[k][1]+tolerance:
                    hit = True
                    break
            if hit:
                hit_sum += 1
                break
    hit_ratio = float(hit_sum)/(len(labels)-1)
    return hit_ratio, hit_sum

if __name__ == "__main__":
    print("Welcome to the WAVE analysis system (beta)!")
    mne.utils.logger.setLevel(mne.utils.logging.CRITICAL)

    DATA_LOADED = False
    TIME_SPECIFIED = False
    START_TIME_SPECIFIED = False
    END_TIME_SPECIFIED = False
    data_path = input("\033[0;32mEnter the full path of the .edf file:\n\033[0m")
#    data_path = r'D:\Brain\source\case\wangshuqi_1\wangshuqi_1.edf'

    while not DATA_LOADED:
        try:
            if data_path == 'exit':
                sys.exit("WAVE analyis terminated.")
            print("\nLoading data...")
            raw = mne.io.read_raw_edf(data_path, preload=True)
        except Exception as e:
            data_path = input("\033[0;32mNot a valid .edf file, please re-enter the full path of the .edf file:\n\033[0m")
        else:
            DATA_LOADED = True

    raw_to_analyze = raw.copy()
    start_time_str = input("\033[0;32mSpecify the start time in seconds (from [0.00, %.2f]) for analysis:\n\033[0m"%raw._times[-1])
    while not START_TIME_SPECIFIED:
        try:
            if start_time_str == 'exit':
                sys.exit("WAVE analyis terminated.")
            start_time = float(start_time_str)
            raw_to_analyze.copy().crop(start_time, raw._times[-1])
        except Exception as e:
            start_time_str = input("\033[0;32mNot valid start time, please re-enter the start time in seconds (from [0.00, %.2f]) for analysis:\n\033[0m"%raw._times[-1])
        else:
            START_TIME_SPECIFIED = True

    end_time_str = input("\033[0;32mSpecify the end time in seconds (from [%.2f, %.2f]) for analysis:\n\033[0m"%(start_time,raw._times[-1]))
    while not END_TIME_SPECIFIED:
        try:
            if end_time_str == 'exit':
                sys.exit("WAVE analyis terminated.")
            end_time = float(end_time_str)
            raw_to_analyze.crop(start_time, end_time)
        except Exception as e:
            end_time_str = input("\033[0;32mNot valid end time, please re-enter the start time in seconds (from [%.2f, %.2f]) for analysis:\n\033[0m"%(start_time,raw._times[-1]))
        else:
            END_TIME_SPECIFIED = True

    global analysis_start_time, analysis_end_time, original_total_time
    analysis_start_time = start_time
    analysis_end_time = end_time
    original_total_time = raw._times[-1]

    print("\nStarting to analyze...")
    time_start_find = time.time()
#    wave_dict_all = find_waveform.waveform_analysis_direct(raw_to_analyze)
    wave_all, wave_dict_all, discharge_time, raw_pattern, picks_pattern, peak_dics_slow, stds_all = wave_analyze(raw_to_analyze)
    time_end_find = time.time()

    labels = []
    try:
        with open(data_path[:-4]+'_labels.csv', 'rb') as csvfile:
            spamreader = csv.reader(csvfile,dialect='excel')
            for row in spamreader:
                labels.append(row)
    except:
        print ("\033[0;31mFailed to open " + data_path[:-4] + "_labels.csv.\033[0m")

    record_start = 8*60.+8
    hit_ratio, hit_sum = get_hit_ratio(discharge_time, labels, record_start)
    print("WAVE analysis finished in %.3f seconds for data at [%.2f, %.2f]s. Hit ratio is %.2f%% (%d out of %d)."%(time_end_find-time_start_find, analysis_start_time, analysis_end_time, 100*hit_ratio, hit_sum, len(labels)-1))

    print ("\033[0;36m\nNow the analyzed data can be plotted with the following code:\nplot_data(raw_pattern, picks_pattern, plot_names=np.array(raw_pattern.info['ch_names'])[picks_pattern], gap=100, start=80, stop=90, HFO_time=discharge_time, line_colors=['c','r'])\033[0m")

#D:\Brain\source\case\wangshuqi_1\wangshuqi_1.edf
#D:\Brain\source\case\wangshuqi_2\wangshuqi.edf

#i = 11
#ii = 1828
#peak_dics = peak_dics_slow[i]
#amp_threshold = max([2*2.3*stds_all[i][ii],0.5*100.])


from PIL import Image, ImageOps
import numpy as np
import skimage.io as io
from src.models.class_patcher import patcher
from src.utils.imgproc import *
from skimage.color import rgb2hsv, hsv2rgb

class patcher(patcher):
    def __init__(self, body='./body/body_kyoko.png', **options):
        super().__init__('京狐', body=body, pantie_position=[718, 1464], **options)
        self.mask = io.imread('./mask/mask_kyoko.png')
        try:
            self.with_garter = self.options['with_garter']
        except:
            self.with_garter = self.ask(question='With garter belt?', default=True)
        if self.with_garter:
            self.garter_position = [701, 1272]
            self.garter = np.float32(io.imread('./material/garter_kyoko.png') / 255)
            self.garter_shade = np.float32(io.imread('./material/garter_kyoko_shade.png') / 255)
            self.garter_shade_alpha = self.garter_shade[:, :, -1]

        self.bra_position = [700, 1008]
        self.bra = np.float32(io.imread('./mask/bra_kyoko.png') / 255)
        self.bra_center = np.float32(io.imread('./mask/bra_kyoko_center.png') / 255)
        self.bra_shade = np.float32(io.imread('./material/bra_kyoko_shade.png') / 255)
        self.bra_lace = np.float32(io.imread('./material/bra_kyoko_lace.png') / 255)
        self.bra_shade_alpha = self.bra_shade[:, :, -1]
        self.bra_lace_mask = self.bra_lace[:, :, -1] > 0.3

        self.pantie_ribbon_position = [745, 1528]
        self.bra_ribbon_position = [800, 1173]
        self.ribbon = np.float32(io.imread('./material/ribbon_kyoko.png') / 255)
        self.ribbon_shade = np.float32(io.imread('./material/ribbon_kyoko_shade.png') / 255)
        self.ribbon_shade_alpha = self.ribbon_shade[:, :, -1]

    def pick_color(self, arr):
        return np.mean(np.mean(arr, axis=0), axis=0)

    def extract_base_color(self, pantie):
        front = pantie[20:100, 30:80, :3] / 255.0
        front_shade = pantie[130:150, 0:40, :3] / 255.0
        front_color = self.pick_color(front)
        front_shade_color = self.pick_color(front_shade)
        front_shade_color = rgb2hsv(front_shade_color[None, None])
        front_shade_color[0, 0, 1] *= front_shade_color[0, 0, 2] / 0.3
        if front_shade_color[0, 0, 1] > 0.7:
            front_shade_color[0, 0, 1] *= 0.7
        front_shade_color[0, 0, 2] *= front_shade_color[0, 0, 2] / 0.4
        front_shade_color = np.clip(hsv2rgb(front_shade_color)[0, 0], 0, 1)
        return front_color, front_shade_color

    def gen_ribbon(self, image):
        pantie = np.array(image)
        ribbon = pantie[24:32, 15:27, :3] / 255.0
        ribbon_color = self.pick_color(ribbon)
        ribbon_shade = pantie[26:30, 12:15, :3] / 255.0
        ribbon_shade_color = self.pick_color(ribbon_shade)
        ribbon = self.ribbon[:, :, :3] * ribbon_color
        ribbon_shade = (self.ribbon_shade[:, :, -1])[:, :, None] * ribbon_shade_color
        ribbon = alpha_brend(ribbon_shade, ribbon, self.ribbon_shade_alpha)
        ribbon = np.dstack((ribbon, self.ribbon[:, :, -1] > 0.5))
        return Image.fromarray(np.uint8(np.clip(ribbon, 0, 1) * 255))

    def gen_garter(self, image):
        pantie = np.array(image)
        front_color, front_shade_color = self.extract_base_color(pantie)
        garter = self.garter[:, :, :3] * front_color
        garter_shade = (self.garter_shade[:, :, -1])[:, :, None] * front_shade_color
        garter = alpha_brend(garter_shade, garter, self.garter_shade_alpha)
        garter = np.dstack((garter, self.garter[:, :, -1] > 0.5))
        return Image.fromarray(np.uint8(np.clip(garter, 0, 1) * 255))

    def gen_bra(self, image):
        pantie = np.array(image)
        front_color, front_shade_color = self.extract_base_color(pantie)
        ribbon = pantie[24:32, 15:27, :3] / 255.0
        ribbon_color = self.pick_color(ribbon)
        center = np.float32(pantie[20:170, -200:-15, :3][:, ::-1]) / 255
        bra_center = np.copy(self.bra_center)
        bra_center[80:80 + center.shape[0], 30:30 + center.shape[1], :3] = center * np.float32(bra_center[80:80 + center.shape[0], 30:30 + center.shape[1], :3] > 0)
        bra = self.bra[:, :, :3] * front_color
        bra_shade = (self.bra_shade[:, :, -1])[:, :, None] * front_shade_color
        bra_lace = self.bra_lace[:, :, :3] * ribbon_color
        bra = alpha_brend(bra_center[:, :, :3], bra[:, :, :3], bra_center[:, :, 0] > 0.1)
        bra = alpha_brend(bra_lace, bra, self.bra_lace_mask)
        bra = alpha_brend(bra_shade, bra, self.bra_shade_alpha)
        bra = np.dstack((bra, self.bra[:, :, 0] > 0.8))
        return Image.fromarray(np.uint8(np.clip(bra, 0, 1) * 255))

    def convert(self, image):
        pantie = np.array(image)
        pantie = ribbon_inpaint(pantie)
        patch = np.copy(pantie[-140:-5, 546:, :])
        [pr, pc, d] = patch.shape
        pantie[125:125 + pr, :pc, :] = patch
        pantie[-140:, 546:, :] = 0
        pantie = np.uint8(resize(pantie, [0.7, 0.7]) * 255)[:170]
        pantie = np.bitwise_and(pantie, self.mask)
        return Image.fromarray(pantie)

    def patch(self, image, transparent=False):
        pantie = self.convert(image)
        if transparent:
            patched = Image.new("RGBA", (4096, 4096))
        else:
            patched = self.body.copy()
        patched = self.paste(patched, pantie, self.pantie_position)
        patched = self.paste(patched, self.gen_bra(image), self.bra_position)
        ribbon = self.gen_ribbon(image)
        patched = self.paste(patched, ribbon, self.pantie_ribbon_position)
        patched = self.paste(patched, ribbon.resize((int(ribbon.width * 0.62), int(ribbon.height * 0.62))), self.bra_ribbon_position)
        if self.with_garter:
            patched = self.paste(patched, self.gen_garter(image), self.garter_position)
        return patched


# Support for python2
from __future__ import print_function
# Import ctypes
from ctypes import *
# Importing Numpy (math, arrays, etc...)
import numpy as np
# Import platform to detect OS
from sys import platform, exit
# Import os utils
from os import path
# Import thermo
from . import thermo

c_len_type = thermo.c_len_type

class cubic(thermo.thermopack):
    """
    Interface to cubic
    """
    def __init__(self):
        """
        Initialize cubic specific function pointers
        """
        # Load dll/so
        super(cubic, self).__init__()

        # Init methods
        self.eoslibinit_init_cubic = getattr(self.tp, self.get_export_name("eoslibinit", "init_cubic"))
        # Tuning methods
        self.s_get_kij = getattr(self.tp, self.get_export_name("", "thermopack_getkij"))
        self.s_set_kij = getattr(self.tp, self.get_export_name("", "thermopack_setkijandji"))

        self.s_get_hv_param = getattr(self.tp, self.get_export_name("", "thermopack_gethvparam"))
        self.s_set_hv_param = getattr(self.tp, self.get_export_name("", "thermopack_sethvparam"))

        self.s_get_ci = getattr(self.tp, self.get_export_name("", "thermopack_get_volume_shift_parameters"))
        self.s_set_ci = getattr(self.tp, self.get_export_name("", "thermopack_set_volume_shift_parameters"))

    #################################
    # Init
    #################################

    def init(self, comps, eos, mixing="vdW", alpha="Classic",
             parameter_reference="Default", volume_shift=False):
        """Initialize cubic model in thermopack

        Args:
            comps (str): Comma separated list of component names
            eos (str): Equation of state (SRK, PR, ...)
            mixing (str, optional): Mixture model. Defaults to "vdW".
            alpha (str, optional): Alpha model. Defaults to "Classic".
            parameter_reference (str, optional): Which parameters to use?. Defaults to "Default".
        """
        self.activate()
        eos_c = c_char_p(eos.encode('ascii'))
        eos_len = c_len_type(len(eos))
        mixing_c = c_char_p(mixing.encode('ascii'))
        mixing_len = c_len_type(len(mixing))
        alpha_c = c_char_p(alpha.encode('ascii'))
        alpha_len = c_len_type(len(alpha))
        comp_string_c = c_char_p(comps.encode('ascii'))
        comp_string_len = c_len_type(len(comps))
        ref_string_c = c_char_p(parameter_reference.encode('ascii'))
        ref_string_len = c_len_type(len(parameter_reference))

        if volume_shift:
            vol_shift_c = c_int(1)
        else:
            vol_shift_c = c_int(0)

        self.eoslibinit_init_cubic.argtypes = [c_char_p,
                                               c_char_p,
                                               c_char_p,
                                               c_char_p,
                                               c_char_p,
                                               POINTER (c_int),
                                               c_len_type,
                                               c_len_type,
                                               c_len_type,
                                               c_len_type,
                                               c_len_type]

        self.eoslibinit_init_cubic.restype = None

        self.eoslibinit_init_cubic(comp_string_c,
                                   eos_c,
                                   mixing_c,
                                   alpha_c,
                                   ref_string_c,
                                   byref(vol_shift_c),
                                   comp_string_len,
                                   eos_len,
                                   mixing_len,
                                   alpha_len,
                                   ref_string_len)
        self.nc = max(len(comps.split(" ")),len(comps.split(",")))

    def get_kij(self, c1, c2):
        """Get attractive energy interaction parameter

        Args:
            c1 (int): Component one
            c2 (int): Component two

        Returns:
            kij (float): i-j interaction parameter
        """
        self.activate()
        c1_c = c_int(c1)
        c2_c = c_int(c2)
        kij_c = c_double(0.0)
        self.s_get_kij.argtypes = [POINTER(c_int),
                                   POINTER(c_int),
                                   POINTER(c_double)]

        self.s_get_kij.restype = None

        self.s_get_kij(byref(c1_c),
                       byref(c2_c),
                       byref(kij_c))

        return kij_c.value

    def set_kij(self, c1, c2, kij):
        """Set attractive energy interaction parameter

        Args:
            c1 (int): Component one
            c2 (int): Component two
            kij (float): i-j interaction parameter
        """
        self.activate()
        c1_c = c_int(c1)
        c2_c = c_int(c2)
        kij_c = c_double(kij)
        self.s_set_kij.argtypes = [POINTER(c_int),
                                   POINTER(c_int),
                                   POINTER(c_double)]

        self.s_set_kij.restype = None

        self.s_set_kij(byref(c1_c),
                       byref(c2_c),
                       byref(kij_c))

    def get_lij(self, c1, c2):
        """Get co-volume interaction

        Args:
            c1 (int): Component one
            c2 (int): Component two

        Returns:
            lij (float): i-j interaction parameter
        """
        return 1.0

    def set_lij(self, c1, c2, lij):
        """Set co-volume interaction

        Args:
            c1 (int): Component one
            c2 (int): Component two
            lij ([type]): [description]
        """
        print("Setting lij")

    def get_hv_param(self, c1, c2):
        """Get Huron-Vidal parameters

        Args:
            c1 (int): Component one
            c2 (int): Component two

        Returns:
            alpha_ij (float): alpha param i-j
            alpha_ji (float): alpha param j-i
            a_ij (float): a param i-j
            a_ji (float): a param j-i
            b_ij (float): b param i-j
            b_ji (float): b param j-i
            c_ij (float): c param i-j
            c_ji (float): c param j-i
        """
        self.activate()
        c1_c = c_int(c1)
        c2_c = c_int(c2)
        alpha_ij_c = c_double(0.0)
        alpha_ji_c = c_double(0.0)
        a_ij_c = c_double(0.0)
        a_ji_c = c_double(0.0)
        b_ij_c = c_double(0.0)
        b_ji_c = c_double(0.0)
        c_ij_c = c_double(0.0)
        c_ji_c = c_double(0.0)

        self.s_get_hv_param.argtypes = [POINTER(c_int),
                                        POINTER(c_int),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double)]

        self.s_get_hv_param.restype = None

        self.s_get_hv_param(byref(c1_c),
                            byref(c2_c),
                            byref(alpha_ij_c),
                            byref(alpha_ji_c),
                            byref(a_ij_c),
                            byref(a_ji_c),
                            byref(b_ij_c),
                            byref(b_ji_c),
                            byref(c_ij_c),
                            byref(c_ji_c))
        return alpha_ij_c.value, alpha_ji_c.value, a_ij_c.value, a_ji_c.value, b_ij_c.value, b_ji_c.value, c_ij_c.value, c_ji_c.value

    def set_hv_param(self, c1, c2, alpha_ij, alpha_ji, a_ij, a_ji, b_ij, b_ji, c_ij, c_ji):
        """Set Huron-Vidal parameters

        Args:
            c1 (int): Component one
            c2 (int): Component two
            alpha_ij (float): alpha param i-j
            alpha_ji (float): alpha param j-i
            a_ij (float): a param i-j
            a_ji (float): a param j-i
            b_ij (float): b param i-j
            b_ji (float): b param j-i
            c_ij (float): c param i-j
            c_ji (float): c param j-i
        """
        self.activate()
        c1_c = c_int(c1)
        c2_c = c_int(c2)
        alpha_ij_c = c_double(alpha_ij)
        alpha_ji_c = c_double(alpha_ji)
        a_ij_c = c_double(a_ij)
        a_ji_c = c_double(a_ji)
        b_ij_c = c_double(b_ij)
        b_ji_c = c_double(b_ji)
        c_ij_c = c_double(c_ij)
        c_ji_c = c_double(c_ji)

        self.s_set_hv_param.argtypes = [POINTER(c_int),
                                        POINTER(c_int),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double),
                                        POINTER(c_double)]

        self.s_set_hv_param.restype = None

        self.s_set_hv_param(byref(c1_c),
                            byref(c2_c),
                            byref(alpha_ij_c),
                            byref(alpha_ji_c),
                            byref(a_ij_c),
                            byref(a_ji_c),
                            byref(b_ij_c),
                            byref(b_ji_c),
                            byref(c_ij_c),
                            byref(c_ji_c))

    def get_ci(self, cidx):
        """Get volume correction parameters

        Args:
            cidx (int): Component index

        Returns:
            ciA (float): Volume shift param of component cidx (m3/mol)
            ciB (float): Volume shift param of component cidx (m3/mol/K)
            ciC (float): Volume shift param of component cidx (m3/mol/K^2)
            ci_type (int): Volume shift type (CONSTANT=1, LINEAR=2, QUADRATIC=3)
        """
        cidx_c = c_int(cidx)
        ciA_c = c_double(0.0)
        ciB_c = c_double(0.0)
        ciC_c = c_double(0.0)
        ci_type_c = c_int(0)
        self.s_get_ci.argtypes = [POINTER(c_int),
                                  POINTER(c_double),
                                  POINTER(c_double),
                                  POINTER(c_double),
                                  POINTER(c_int)]

        self.s_get_ci.restype = None

        self.s_get_ci(byref(cidx_c),
                      byref(ciA_c),
                      byref(ciB_c),
                      byref(ciC_c),
                      byref(ci_type_c))

        return ciA_c.value, ciB_c.value, ciC_c.value, ci_type_c.value

    def set_ci(self, cidx, ciA, ciB=0.0, ciC=0.0, ci_type=1):
        """Set volume correction parametrs

        Args:
            cidx (int): Component index
            ciA (float): Volume shift param of component cidx (m3/mol)
            ciB (float): Volume shift param of component cidx (m3/mol/K)
            ciC (float): Volume shift param of component cidx (m3/mol/K^2)
            ci_type (int): Volume shift type (CONSTANT=1, LINEAR=2, QUADRATIC=3)
        """
        cidx_c = c_int(cidx)
        ciA_c = c_double(ciA)
        ciB_c = c_double(ciB)
        ciC_c = c_double(ciC)
        ci_type_c = c_int(ci_type)
        self.s_set_ci.argtypes = [POINTER(c_int),
                                  POINTER(c_double),
                                  POINTER(c_double),
                                  POINTER(c_double),
                                  POINTER(c_int)]

        self.s_set_ci.restype = None

        self.s_set_ci(byref(cidx_c),
                      byref(ciA_c),
                      byref(ciB_c),
                      byref(ciC_c),
                      byref(ci_type_c))


'''WIP --- Doesn't work well yet'''
import plac
import random
import six

import pathlib
import cPickle as pickle
from itertools import izip

import spacy

import cytoolz
import cupy as xp
import cupy.cuda
import chainer.cuda

import chainer.links as L
import chainer.functions as F
from chainer import Chain, Variable, report
import chainer.training
import chainer.optimizers
from chainer.training import extensions
from chainer.iterators import SerialIterator
from chainer.datasets import TupleDataset

class SentimentAnalyser(object):
    @classmethod
    def load(cls, path, nlp, max_length=100):
        raise NotImplementedError
        #with (path / 'config.json').open() as file_:
        #    model = model_from_json(file_.read())
        #with (path / 'model').open('rb') as file_:
        #    lstm_weights = pickle.load(file_)
        #embeddings = get_embeddings(nlp.vocab)
        #model.set_weights([embeddings] + lstm_weights)
        #return cls(model, max_length=max_length)

    def __init__(self, model, max_length=100):
        self._model = model
        self.max_length = max_length

    def __call__(self, doc):
        X = get_features([doc], self.max_length)
        y = self._model.predict(X)
        self.set_sentiment(doc, y)

    def pipe(self, docs, batch_size=1000, n_threads=2):
        for minibatch in cytoolz.partition_all(batch_size, docs):
            minibatch = list(minibatch)
            sentences = []
            for doc in minibatch:
                sentences.extend(doc.sents)
            Xs = get_features(sentences, self.max_length)
            ys = self._model.predict(Xs)
            for sent, label in zip(sentences, ys):
                sent.doc.sentiment += label - 0.5
            for doc in minibatch:
                yield doc

    def set_sentiment(self, doc, y):
        doc.sentiment = float(y[0])
        # Sentiment has a native slot for a single float.
        # For arbitrary data storage, there's:
        # doc.user_data['my_data'] = y

class Classifier(Chain):
    def __init__(self, predictor):
        super(Classifier, self).__init__(predictor=predictor)

    def __call__(self, x, t):
        y = self.predictor(x)
        loss = F.softmax_cross_entropy(y, t)
        accuracy = F.accuracy(y, t)
        report({'loss': loss, 'accuracy': accuracy}, self)
        return loss

class SentimentModel(Chain):
    def __init__(self, nlp, shape, **settings):
        Chain.__init__(self,
            embed=_Embed(shape['nr_vector'], shape['nr_dim'], shape['nr_hidden'],
                initialW=lambda arr: set_vectors(arr, nlp.vocab)),
            encode=_Encode(shape['nr_hidden'], shape['nr_hidden']),
            attend=_Attend(shape['nr_hidden'], shape['nr_hidden']),
            predict=_Predict(shape['nr_hidden'], shape['nr_class']))
        self.to_gpu(0)

    def __call__(self, sentence):
        return self.predict(
                  self.attend(
                      self.encode(
                          self.embed(sentence))))

class _Embed(Chain):
    def __init__(self, nr_vector, nr_dim, nr_out):
        Chain.__init__(self,
            embed=L.EmbedID(nr_vector, nr_dim),
            project=L.Linear(None, nr_out, nobias=True))
        #self.embed.unchain_backward()

    def __call__(self, sentence):
        return [self.project(self.embed(ts)) for ts in F.transpose(sentence)]

class _Encode(Chain):
    def __init__(self, nr_in, nr_out):
        Chain.__init__(self,
            fwd=L.LSTM(nr_in, nr_out),
            bwd=L.LSTM(nr_in, nr_out),
            mix=L.Bilinear(nr_out, nr_out, nr_out))

    def __call__(self, sentence):
        self.fwd.reset_state()
        fwds = map(self.fwd, sentence)
        self.bwd.reset_state()
        bwds = reversed(map(self.bwd, reversed(sentence)))
        return [F.elu(self.mix(f, b)) for f, b in zip(fwds, bwds)]

class _Attend(Chain):
    def __init__(self, nr_in, nr_out):
        Chain.__init__(self)

    def __call__(self, sentence):
        sent = sum(sentence)
        return sent

class _Predict(Chain):
    def __init__(self, nr_in, nr_out):
        Chain.__init__(self,
            l1=L.Linear(nr_in, nr_in),
            l2=L.Linear(nr_in, nr_out))

    def __call__(self, vector):
        vector = self.l1(vector)
        vector = F.elu(vector)
        vector = self.l2(vector)
        return vector

class SentenceDataset(TupleDataset):
    def __init__(self, nlp, texts, labels, max_length):
        self.max_length = max_length
        sents, labels = self._get_labelled_sentences(
            nlp.pipe(texts, batch_size=5000, n_threads=3),
            labels)
        TupleDataset.__init__(self,
            get_features(sents, max_length),
            labels)

    def __getitem__(self, index):
        batches = [dataset[index] for dataset in self._datasets]
        if isinstance(index, slice):
            length = len(batches[0])
            returns = [tuple([batch[i] for batch in batches])
                       for i in six.moves.range(length)]
            return returns
        else:
            return tuple(batches)

    def _get_labelled_sentences(self, docs, doc_labels):
        labels = []
        sentences = []
        for doc, y in izip(docs, doc_labels):
            for sent in doc.sents:
                sentences.append(sent)
                labels.append(y)
        return sentences, xp.asarray(labels, dtype='i')

class DocDataset(TupleDataset):
    def __init__(self, nlp, texts, labels):
        self.max_length = max_length
        DatasetMixin.__init__(self,
            get_features(
                nlp.pipe(texts, batch_size=5000, n_threads=3), self.max_length),
            labels)

def read_data(data_dir, limit=0):
    examples = []
    for subdir, label in (('pos', 1), ('neg', 0)):
        for filename in (data_dir / subdir).iterdir():
            with filename.open() as file_:
                text = file_.read()
            examples.append((text, label))
    random.shuffle(examples)
    if limit >= 1:
        examples = examples[:limit]
    return zip(*examples) # Unzips into two lists

def get_features(docs, max_length):
    docs = list(docs)
    Xs = xp.zeros((len(docs), max_length), dtype='i')
    for i, doc in enumerate(docs):
        j = 0
        for token in doc:
            if token.has_vector and not token.is_punct and not token.is_space:
                Xs[i, j] = token.norm
                j += 1
                if j >= max_length:
                    break
    return Xs

def set_vectors(vectors, vocab):
    for lex in vocab:
        if lex.has_vector and (lex.rank+1) < vectors.shape[0]:
            lex.norm = lex.rank+1
            vectors[lex.rank + 1] = lex.vector
        else:
            lex.norm = 0
    vectors.unchain_backwards()
    return vectors

def train(train_texts, train_labels, dev_texts, dev_labels,
        lstm_shape, lstm_settings, lstm_optimizer, batch_size=100, nb_epoch=5,
        by_sentence=True):
    nlp = spacy.load('en', entity=False)
    if 'nr_vector' not in lstm_shape:
        lstm_shape['nr_vector'] = max(lex.rank+1 for lex in vocab if lex.has_vector)
    print("Make model")
    model = Classifier(SentimentModel(nlp, lstm_shape, **lstm_settings))
    print("Parsing texts...")
    if by_sentence:
        train_data = SentenceDataset(nlp, train_texts, train_labels, lstm_shape['max_length'])
        dev_data = SentenceDataset(nlp, dev_texts, dev_labels, lstm_shape['max_length'])
    else:
        train_data = DocDataset(nlp, train_texts, train_labels)
        dev_data = DocDataset(nlp, dev_texts, dev_labels)
    train_iter = SerialIterator(train_data, batch_size=batch_size,
                                shuffle=True, repeat=True)
    dev_iter = SerialIterator(dev_data, batch_size=batch_size,
                              shuffle=False, repeat=False)
    optimizer = chainer.optimizers.Adam()
    optimizer.setup(model)
    updater = chainer.training.StandardUpdater(train_iter, optimizer, device=0)
    trainer = chainer.training.Trainer(updater, (20, 'epoch'), out='result')

    trainer.extend(extensions.Evaluator(dev_iter, model, device=0))
    trainer.extend(extensions.LogReport())
    trainer.extend(extensions.PrintReport([
        'epoch', 'main/accuracy', 'validation/main/accuracy']))
    trainer.extend(extensions.ProgressBar())

    trainer.run()

def evaluate(model_dir, texts, labels, max_length=100):
    def create_pipeline(nlp):
        '''
        This could be a lambda, but named functions are easier to read in Python.
        '''
        return [nlp.tagger, nlp.parser, SentimentAnalyser.load(model_dir, nlp,
                                                               max_length=max_length)]

    nlp = spacy.load('en')
    nlp.pipeline = create_pipeline(nlp)

    correct = 0
    i = 0
    for doc in nlp.pipe(texts, batch_size=1000, n_threads=4):
        correct += bool(doc.sentiment >= 0.5) == bool(labels[i])
        i += 1
    return float(correct) / i

@plac.annotations(
    train_dir=("Location of training file or directory"),
    dev_dir=("Location of development file or directory"),
    model_dir=("Location of output model directory",),
    is_runtime=("Demonstrate run-time usage", "flag", "r", bool),
    nr_hidden=("Number of hidden units", "option", "H", int),
    max_length=("Maximum sentence length", "option", "L", int),
    dropout=("Dropout", "option", "d", float),
    learn_rate=("Learn rate", "option", "e", float),
    nb_epoch=("Number of training epochs", "option", "i", int),
    batch_size=("Size of minibatches for training LSTM", "option", "b", int),
    nr_examples=("Limit to N examples", "option", "n", int)
)
def main(model_dir, train_dir, dev_dir,
         is_runtime=False,
         nr_hidden=64, max_length=100, # Shape
         dropout=0.5, learn_rate=0.001, # General NN config
         nb_epoch=5, batch_size=32, nr_examples=-1):  # Training params
    model_dir = pathlib.Path(model_dir)
    train_dir = pathlib.Path(train_dir)
    dev_dir = pathlib.Path(dev_dir)
    if is_runtime:
        dev_texts, dev_labels = read_data(dev_dir)
        acc = evaluate(model_dir, dev_texts, dev_labels, max_length=max_length)
        print(acc)
    else:
        print("Read data")
        train_texts, train_labels = read_data(train_dir, limit=nr_examples)
        dev_texts, dev_labels = read_data(dev_dir, limit=nr_examples)
        print("Using GPU 0")
        #chainer.cuda.get_device(0).use()
        train_labels = xp.asarray(train_labels, dtype='i')
        dev_labels = xp.asarray(dev_labels, dtype='i')
        lstm = train(train_texts, train_labels, dev_texts, dev_labels,
                     {'nr_hidden': nr_hidden, 'max_length': max_length, 'nr_class': 2,
                      'nr_vector': 2000, 'nr_dim': 32},
                      {'dropout': 0.5, 'lr': learn_rate},
                      {},
                      nb_epoch=nb_epoch, batch_size=batch_size)

if __name__ == '__main__':
    plac.call(main)


# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

#

#

#

#
ObsPy Tutorial

#
Handling Waveform Data

#

#

#

# image: User:Abbaszade656 / Wikimedia Commons / CC-BY-SA-4.0

# ## Workshop for the "Training in Network Management Systems and Analytical Tools for Seismic"
# ### Baku, October 2018
#
# Seismo-Live: http://seismo-live.org
#
# ##### Authors:
# * Tobias Megies ([@megies](https://github.com/megies))
# * Lion Krischer ([@krischer](https://github.com/krischer))
# ---

# ![](images/obspy_logo_full_524x179px.png)

# %matplotlib inline
import matplotlib.pyplot as plt
plt.style.use('ggplot')
plt.rcParams['figure.figsize'] = 12, 8

#
#
# * read waveform data is returned as a **`Stream`** object.

from obspy import read
st = read("./data/waveform_PFO.mseed", format="mseed")
print(st)

# - UNIX wildcards can be used to read multiple files simultaneously
# - automatic file format detection, no need to worry about file formats
#
#   - currently supported: **MiniSEED, SAC, SEG-Y, SEG-2, GSE1/2, Seisan, SH, DataMark, CSS, wav, y, Q, ASCII, Guralp Compressed Format, Kinemetrics EVT, K-NET&KiK formats, PDAS, Reftek 130, WIN (keeps growing...)**
#   - more file formats are included whenever a basic reading routine is provided (or e.g. sufficient documentation on data compression etc.)
#   - custom user-specific file formats can be added (through plug-ins) to be auto-discovered in a local ObsPy installation

from obspy import read
st = read("./data/waveform_*")
print(st)

# - for MiniSEED files, only reading short portions of files without all of the file getting read into memory is supported (saves time and memory when working on large collections of big files)

# +
from obspy import UTCDateTime

t = UTCDateTime("2011-03-11T05:46:23.015400Z")
st = read("./data/waveform_*", starttime=t + 10 * 60, endtime=t + 12 * 60)
print(st)
# -

# #### The Stream Object
#
#  - A Stream object is a collection of Trace objects

from obspy import read
st = read("./data/waveform_PFO.mseed")
print(type(st))

print(st)

print(st.traces)

print(st[0])

# - More traces can be assembled using **`+`** operator (or using the `.append()` and `.extend()` methods)

# +
st1 = read("./data/waveform_PFO.mseed")
st2 = read("./data/waveform_PFO_synthetics.mseed")

st = st1 + st2
print(st)

st3 = read("./data/waveform_BFO_BHE.sac")

st += st3
print(st)
# -

#  - convenient (and nicely readable) looping over traces

for tr in st:
    print(tr.id)

#  - Stream is useful for applying the same processing to a larger number of different waveforms or to group Traces for processing (e.g. three components of one station in one Stream)

# #### The Trace Object
#
# - a Trace object is a single, contiguous waveform data block (i.e. regularly spaced time series, no gaps)
# - a Trace object contains a limited amount of metadata in a dictionary-like object (as **`Trace.stats`**) that fully describes the time series by specifying..
#   * recording location/instrument (network, station, location and channel code)
#   * start time
#   * sampling rate

st = read("./data/waveform_PFO.mseed")
tr = st[0]  # get the first Trace in the Stream
print(tr)

print(tr.stats)

# - For custom applications it is sometimes necessary to directly manipulate the metadata of a Trace.
# - The metadata of the Trace will **stay consistent**, as all values are derived from the starttime, the data and the sampling rate and are **updated automatically**

print(tr.stats.delta, "|", tr.stats.endtime)

tr.stats.sampling_rate = 5.0
print(tr.stats.delta, "|", tr.stats.endtime)

print(tr.stats.npts)

tr.data = tr.data[:100]
print(tr.stats.npts, "|", tr.stats.endtime)

# - convenience methods make basic manipulations simple

tr = read("./data/waveform_PFO.mseed")[0]
tr.plot()

print(tr)
tr.resample(sampling_rate=100.0)
print(tr)

print(tr)
tr.trim(tr.stats.starttime + 12 * 60, tr.stats.starttime + 14 * 60)
print(tr)
tr.plot()

tr.detrend("linear")
tr.taper(max_percentage=0.05, type='cosine')
tr.filter("lowpass", freq=0.1)
tr.plot()

# try tr.<Tab> for other methods defined for Trace
tr.detrend?

# - Raw data available as a [**`numpy.ndarray`**](http://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.html) (as **`Trace.data`**)

print(tr.data[:20])

# - Data can be directly modified e.g. ..
#
# ..by doing arithmetic operations (fast, handled in C by NumPy)

print(tr.data ** 2 + 0.5)

# ..by using [**`numpy.ndarray`** builtin methods](http://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.html) (also done in C by NumPy)

print(tr.data.max())
print(tr.data.mean())
print(tr.data.ptp())
# try tr.data.<Tab> for a list of numpy methods defined on ndarray

# ..by using **`numpy`** functions (also done in C by NumPy)

import numpy as np
print(np.abs(tr.data))
# you can try np.<Tab> but there is a lot in there
# try np.a<Tab>

# ..by feeding pointers to existing C/Fortran routines from inside Python!
#
# This is done internally in several places, e.g. for cross correlations, beamforming or in third-party filetype libraries like e.g. libmseed.
#
# - Trace objects can also be manually generated with data in a [**`numpy.ndarray`**](http://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.html) (e.g. when needing to parse waveforms from non-standard ascii files)

# +
from obspy import Trace

x = np.random.randint(-100, 100, 500)
tr = Trace(data=x)
tr.stats.station = "XYZ"
tr.stats.starttime = UTCDateTime()

tr.plot()
# -

# - Stream objects can be assembled from manually generated Traces

# +
from obspy import Stream

tr2 = Trace(data=np.random.randint(-300, 100, 1000))
tr2.stats.starttime = UTCDateTime()
tr2.stats.sampling_rate = 10.0
st = Stream([tr, tr2])

st.plot()
# -

# #### Builtin methods defined on **`Stream`** / **`Trace`**
#
# - Most methods that work on a Trace object also work on a Stream object. They are simply executed for every trace. [See ObsPy documentation for an overview of available methods](http://docs.obspy.org/packages/autogen/obspy.core.stream.Stream.html) (or try **`st.<Tab>`**).
#  - **`st.filter()`** - Filter all attached traces.
#  - **`st.trim()`** - Cut all traces.
#  - **`st.resample()`** / **`st.decimate()`** - Change the sampling rate.
#  - **`st.trigger()`** - Run triggering algorithms.
#  - **`st.plot()`** / **`st.spectrogram()`** - Visualize the data.
#  - **`st.attach_response()`**/**`st.remove_response()`**, **`st.simulate()`** - Instrument correction
#  - **`st.merge()`**, **`st.normalize()`**, **`st.detrend()`**, **`st.taper()`**, ...
# - A **`Stream`** object can also be exported to many formats, so ObsPy can be used to convert between different file formats.

st = read("./data/waveform_*.sac")
st.write("output_file.mseed", format="MSEED")
# !ls -l output_file*

#

# #### Trace Exercises
#  - Make an **`numpy.ndarray`** with zeros and (e.g. use **`numpy.zeros()`**) and put an ideal pulse somewhere in it
#  - initialize a **`Trace`** object with your data array
#  - Fill in some station information (e.g. network, station, ..)
#  - Print trace summary and plot the trace
#  - Change the sampling rate to 20 Hz
#  - Change the starttime of the trace to the start time of this session
#  - Print the trace summary and plot the trace again

x = np.zeros(300)
x[100] = 1.0
tr = Trace(data=x)
tr.stats.station = "ABC"
tr.stats.sampling_rate = 20.0
tr.stats.starttime = UTCDateTime(2014, 2, 24, 15, 0, 0)
print(tr)
tr.plot()

# - Use **`tr.filter(...)`** and apply a lowpass filter with a corner frequency of 1 Hertz.
# - Display the preview plot, there are a few seconds of zeros that we can cut off.

tr.filter("lowpass", freq=1)
tr.plot()

# - Use **`tr.trim(...)`** to remove some of the zeros at start and at the end
# - show the preview plot again

tr.trim(tr.stats.starttime + 3, tr.stats.endtime - 5)
tr.plot()

# - Scale up the amplitudes of the trace by a factor of 500
# - Add standard normal gaussian noise to the trace (use [**`np.random.randn()`**](http://docs.scipy.org/doc/numpy/reference/generated/numpy.random.randn.html))
# - Display the preview plot again

tr.data = tr.data * 500
tr.data = tr.data + np.random.randn(len(tr))
tr.plot()

# #### Stream Exercises
#
# - Read all Tohoku example earthquake data into a stream object ("./data/waveform\_\*")
# - Print the stream summary

st = read("./data/waveform_*")
print(st)

# - Use **`st.select()`** to only keep traces of station BFO in the stream. Show the preview plot.

st = st.select(station="BFO")
st.plot()

# - trim the data to a 10 minute time window around the first arrival (just roughly looking at the preview plot)
# - display the preview plot and spectrograms for the stream (with logarithmic frequency scale, use `wlen=50` for the spectrogram plot)

t1 = UTCDateTime(2011, 3, 11, 5, 55)
st.trim(t1, t1 + 10 * 60)
st.plot()
st.spectrogram(log=True, wlen=50);

# - remove the linear trend from the data, apply a tapering and a lowpass at 0.1 Hertz
# - show the preview plot again

st.detrend("linear")
st.filter("lowpass", freq=0.1)
st.plot()


# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

""" Implementation of external Vitis-AI accel operation """

import os
import tvm
import warnings
import numpy as np

try:
    from dnndk import n2cube, dputils
except:
    warnings.warn("Could not import dnndk n2cube")

@tvm.register_func("tvm.accel.accel_fused")
def accel_fused(kernel_name, input_name, output_name,
    layout, out, *ins):

    # Attach to DPU driver and prepare for running
    n2cube.dpuOpen()

    # Create DPU Kernels
    kernel = n2cube.dpuLoadKernel(kernel_name)

    # Create DPU Tasks for kernel
    task = n2cube.dpuCreateTask(kernel, 0)

    # Load image to DPU
    X = ins[0].asnumpy().reshape((-1))
    n2cube.dpuSetInputTensorInHWCFP32(task, input_name, X, len(X))

    # Model run on DPU """
    n2cube.dpuRunTask(task)

    # Get the output tensor size
    size = n2cube.dpuGetOutputTensorSize(task, output_name)
    address = n2cube.dpuGetOutputTensorAddress(task, output_name)

    value = [0 for i in range(size)]

    # Get the output tensor data
    n2cube.dpuGetTensorData(address, value, size)
    scale = n2cube.dpuGetOutputTensorScale(task, output_name, idx=0)

    value = np.array(value).astype(np.float32) * float(scale)

    value_shape = tuple(out.shape) if layout == 'NHWC' else  \
        (out.shape[0], out.shape[2], out.shape[3], out.shape[1])
    value = np.reshape(value, value_shape)

    # DPU output is in NHWC
    if layout == 'NCHW':
        value = np.transpose(value,(0,3,1,2))

    tvm.nd.array(value).copyto(out)


#!/usr/bin/python

import jetson.inference
import jetson.utils

import argparse
import sys
import pyrealsense2 as rs
import numpy as np
import cv2
from enum import IntEnum

class Preset(IntEnum):
    Custom = 0
    Default = 1
    Hand = 2
    HighAccuracy = 3
    HighDensity = 4
    MediumDensity = 5

# parse the command line
parser = argparse.ArgumentParser(description="Locate objects in a live camera stream using an object detection DNN.",
                                 formatter_class=argparse.RawTextHelpFormatter, epilog=jetson.inference.detectNet.Usage())

parser.add_argument("--network", type=str, default="ssd-mobilenet-v2",
                    help="pre-trained model to load (see below for options)")
parser.add_argument("--overlay", type=str, default="box,labels,conf",
                    help="detection overlay flags (e.g. --overlay=box,labels,conf)\nvalid combinations are:  'box', 'labels', 'conf', 'none'")
parser.add_argument("--threshold", type=float, default=0.4,
                    help="minimum detection threshold to use")
parser.add_argument("--camera", type=str, default="0",
                    help="index of the MIPI CSI camera to use (e.g. CSI camera 0)\nor for VL42 cameras, the /dev/video device to use.\nby default, MIPI CSI camera 0 will be used.")
parser.add_argument("--width", type=int, default=1280,
                    help="desired width of camera stream (default is 1280 pixels)")
parser.add_argument("--height", type=int, default=720,
                    help="desired height of camera stream (default is 720 pixels)")

try:
    opt = parser.parse_known_args()[0]
except:
    print("")
    parser.print_help()
    sys.exit(0)

pipeline = rs.pipeline()
config = rs.config()
config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)
config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
profile = pipeline.start(config)
colorizer = rs.colorizer()
depth_sensor = profile.get_device().first_depth_sensor()
depth_sensor.set_option(rs.option.visual_preset, 4)
# Getting the depth sensor's depth scale (see rs-align example for explanation)
depth_scale = depth_sensor.get_depth_scale()

# We will not display the background of objects more than
#  clipping_distance_in_meters meters away
clipping_distance_in_meters = 3  # 3 meter
clipping_distance = clipping_distance_in_meters / depth_scale

align_to = rs.stream.color
align = rs.align(align_to)

# load the object detection network
net = jetson.inference.detectNet(opt.network, sys.argv, opt.threshold)

# create the camera and display
# camera = jetson.utils.gstCamera(opt.width, opt.height, opt.camera)
display = jetson.utils.glDisplay()

# process frames until user exits
try:
   while True:
    # capture the image
    # img, width, height = camera.CaptureRGBA()

    # Get frameset of color and depth
    frames = pipeline.wait_for_frames()

    # Align the depth frame to color frame
    aligned_frames = align.process(frames)

    spatial = rs.spatial_filter()
    # spatial.set_option(rs.option.holes_fill, 5)

    # spatial = rs.spatial_filter()
    # spatial.set_option(rs.option.filter_magnitude, 5)
    spatial.set_option(rs.option.filter_smooth_alpha, .25)
    spatial.set_option(rs.option.filter_smooth_delta, 45)
    spatial.set_option(rs.option.holes_fill, 2)
    # Get aligned frames
    aligned_depth_frame = aligned_frames.get_depth_frame()

    aligned_depth_frame = spatial.process(aligned_depth_frame)
    color_frame = aligned_frames.get_color_frame()

    depth_image = np.asanyarray(aligned_depth_frame.get_data())
    color_image = np.asanyarray(color_frame.get_data())
    colorized_depth = np.asanyarray(colorizer.colorize(aligned_depth_frame).get_data())
    rgba = cv2.cvtColor(color_image,cv2.COLOR_BGR2RGBA)
    width = color_image.shape[1]
    height = color_image.shape[0]

    img = jetson.utils.cudaFromNumpy(rgba)
    # detect objects in the image (with overlay)
    detections = net.Detect(img, width, height, opt.overlay)

    # print the detections
    print("detected {:d} objects in image".format(len(detections)))

    for detection in detections:
        print(detection)
        (ox,oy)=detection.Center
        ox = int(ox)
        oy = int(oy)
        print("#####################################################")
        print("Mean center of the object detected")
        print(ox,oy)
        depthy = depth_image[oy,ox]

        print("#####################################################")
        print(depthy)

        print("#####################################################")
        print("depth of the object in meter scale ")
        print(depthy*depth_scale)
        print("#####################################################")

    # render the image
    display.RenderOnce(img, width, height)

    cv2.imshow('jetcolor depth', colorized_depth)

    # update the title bar
    display.SetTitle("{:s} | Network {:.0f} FPS".format(
        opt.network, net.GetNetworkFPS()))

    # print out performance info
    net.PrintProfilerTimes()

    '''
    grey_color = 10
    #depth image is 1 channel, color is 3 channels
    depth_image_3d = np.dstack((depth_image, depth_image, depth_image))
    bg_removed = np.where((depth_image_3d > clipping_distance) | \
    (depth_image_3d <= 0), grey_color, color_image)
    cv2.namedWindow('Recorder Realsense', cv2.WINDOW_AUTOSIZE)
    cv2.imshow('Recorder Realsense', bg_removed)
    '''
    key = cv2.waitKey(1)

    # if 'esc' button pressed, escape loop and exit program
    if key & 0xFF == ord('q') or key == 27:
        cv2.destroyAllWindows()
        break

finally:
    pipeline.stop()


"""Single Shot Object Detection model with pretrained ResNet34 backbone extending :class:`.ComposerModel`."""

import os
import tempfile
from typing import Any, Sequence, Tuple, Union

import numpy as np
import requests
from torch import Tensor
from torchmetrics import Metric, MetricCollection

from composer.models.base import ComposerModel
from composer.models.ssd.base_model import Loss
from composer.models.ssd.ssd300 import SSD300
from composer.models.ssd.utils import Encoder, SSDTransformer, dboxes300_coco
from composer.utils.import_helpers import MissingConditionalImportError

__all__ = ["SSD"]

class SSD(ComposerModel):
    """Single Shot Object detection Model with pretrained ResNet34 backbone extending :class:`.ComposerModel`.

    Args:
        input_size (int, optional): input image size. Default: ``300``.
        num_classes (int, optional): The number of classes to detect. Default: ``80``.
        overlap_threshold (float, optional): Minimum IOU threshold for NMS. Default: ``0.5``.
        nms_max_detections (int, optional): Max number of boxes after NMS. Default: ``200``.
        data (str, optional): path to coco dataset. Default: ``"/localdisk/coco"``.
    """

    def __init__(self, input_size: int, overlap_threshold: float, nms_max_detections: int, num_classes: int, data: str):
        super().__init__()

        self.input_size = input_size
        self.overlap_threshold = overlap_threshold
        self.nms_max_detections = nms_max_detections
        self.num_classes = num_classes
        url = "https://download.pytorch.org/models/resnet34-333f7ec4.pth"
        with tempfile.TemporaryDirectory() as tempdir:
            with requests.get(url, stream=True) as r:
                r.raise_for_status()
                pretrained_backbone = os.path.join(tempdir, "weights.pth")
                with open(pretrained_backbone, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        f.write(chunk)
                self.module = SSD300(self.num_classes, model_path=pretrained_backbone)

        dboxes = dboxes300_coco()
        self.loss_func = Loss(dboxes)

        self.encoder = Encoder(dboxes)
        self.data = data
        self.MAP = coco_map(self.data)
        val_annotate = os.path.join(self.data, "annotations/instances_val2017.json")
        val_coco_root = os.path.join(self.data, "val2017")
        input_size = self.input_size
        val_trans = SSDTransformer(dboxes, (input_size, input_size), val=True)
        from composer.datasets.coco import COCODetection
        self.val_coco = COCODetection(val_coco_root, val_annotate, val_trans)

    def loss(self, outputs: Any, batch: Any) -> Union[Tensor, Sequence[Tensor]]:

        (_, _, _, bbox, label) = batch  #type: ignore
        if not isinstance(bbox, Tensor):
            raise TypeError("bbox must be a singular tensor")
        trans_bbox = bbox.transpose(1, 2).contiguous()

        ploc, plabel = outputs
        gloc, glabel = trans_bbox, label

        loss = self.loss_func(ploc, plabel, gloc, glabel)
        return loss

    def metrics(self, train: bool = False) -> Union[Metric, MetricCollection]:
        return self.MAP

    def forward(self, batch: Any) -> Tensor:
        (img, _, _, _, _) = batch  #type: ignore
        ploc, plabel = self.module(img)
        return ploc, plabel  #type: ignore

    def validate(self, batch: Any) -> Tuple[Any, Any]:
        inv_map = {v: k for k, v in self.val_coco.label_map.items()}
        ret = []
        overlap_threshold = self.overlap_threshold
        nms_max_detections = self.nms_max_detections

        (img, img_id, img_size, _, _) = batch  #type: ignore
        ploc, plabel = self.module(img)

        results = []
        try:
            results = self.encoder.decode_batch(ploc,
                                                plabel,
                                                overlap_threshold,
                                                nms_max_detections,
                                                nms_valid_thresh=0.05)
        except:
            print("No object detected")

        (htot, wtot) = [d.cpu().numpy() for d in img_size]  #type: ignore
        img_id = img_id.cpu().numpy()  #type: ignore
        if len(results) > 0:
            # Iterate over batch elements
            for img_id_, wtot_, htot_, result in zip(img_id, wtot, htot, results):
                loc, label, prob = [r.cpu().numpy() for r in result]  #type: ignore
                # Iterate over image detections
                for loc_, label_, prob_ in zip(loc, label, prob):
                    ret.append([img_id_, loc_[0]*wtot_, \
                                loc_[1]*htot_,
                                (loc_[2] - loc_[0])*wtot_,
                                (loc_[3] - loc_[1])*htot_,
                                prob_,
                                inv_map[label_]])

        return ret, ret

class coco_map(Metric):

    def __init__(self, data):
        super().__init__()
        try:
            from pycocotools.coco import COCO
        except ImportError as e:
            raise MissingConditionalImportError(extra_deps_group="coco",
                                                conda_channel="conda-forge",
                                                conda_package="pycocotools") from e
        self.add_state("predictions", default=[])
        val_annotate = os.path.join(data, "annotations/instances_val2017.json")
        self.cocogt = COCO(annotation_file=val_annotate)

    def update(self, pred, target):
        self.predictions.append(pred)  #type: ignore
        np.squeeze(self.predictions)  #type: ignore

    def compute(self):
        try:
            from pycocotools.cocoeval import COCOeval
        except ImportError as e:
            raise MissingConditionalImportError(extra_deps_group="coco",
                                                conda_channel="conda-forge",
                                                conda_package="pycocotools") from e
        cocoDt = self.cocogt.loadRes(np.array(self.predictions))
        E = COCOeval(self.cocogt, cocoDt, iouType='bbox')
        E.evaluate()
        E.accumulate()
        E.summarize()
        return E.stats[0]


# remove onxx imports if using opencv backend
import onnx
import onnx.utils
import onnxruntime
import cv2
import datetime
import os
import numpy as np
import colorama
from colorama import Fore, Style
# cv2 import error fix
import sys

sys.path.append('/usr/local/lib/python2.7/site-packages')
sys.path.append('/usr/local/lib/python3.7/site-packages')
# Find file relative to the location of this code files
default_onnx_path = f"{os.path.dirname(__file__)}/centerface.onnx"

class CenterFace:
    def __init__(self, onnx_path=None, in_shape=None, backend="auto"):
        self.in_shape = in_shape
        self.onnx_input_name = "input.1"
        self.onnx_output_names = ["537", "538", "539", "540"]

        if onnx_path is None:
            onnx_path = default_onnx_path

        if backend == "auto":

            backend = "onnxrt"
            print(
                "attempting to use onnxrt, if its not compatible with your system, change 'backend' to opencv"
            )

        self.backend = backend

        if self.backend == "opencv":
            self.net = cv2.dnn.readNetFromONNX(onnx_path)
        elif self.backend == "onnxrt":
            import onnx
            import onnx.utils
            import onnxruntime

            print(Fore.BLUE + "self.backend is onnxrt")
            # Silence warnings about unnecessary bn initializers
            onnxruntime.set_default_logger_severity(3)

            static_model = onnx.load(onnx_path)
            dyn_model = self.dynamicize_shapes(static_model)
            dyn_model = onnx.utils.polish_model(dyn_model)
            self.sess = onnxruntime.InferenceSession(dyn_model.SerializeToString())

            preferred_provider = self.sess.get_providers()[0]
            print(Fore.BLUE + "preferred_provider ", preferred_provider)
            preferred_device = "GPU" if preferred_provider.startswith("CUDA") else "CPU"
            # print(f'Running on {preferred_device}.')

    @staticmethod
    def dynamicize_shapes(static_model):
        from onnx.tools.update_model_dims import update_inputs_outputs_dims

        input_dims, output_dims = {}, {}
        for node in static_model.graph.input:
            dims = [d.dim_value for d in node.type.tensor_type.shape.dim]
            input_dims[node.name] = dims
        for node in static_model.graph.output:
            dims = [d.dim_value for d in node.type.tensor_type.shape.dim]
            output_dims[node.name] = dims
        input_dims.update({"input.1": ["B", 3, "H", "W"]})  # RGB input image
        output_dims.update(
            {
                "537": ["B", 1, "h", "w"],  # heatmap
                "538": ["B", 2, "h", "w"],  # scale
                "539": ["B", 2, "h", "w"],  # offset
                "540": ["B", 10, "h", "w"],  # landmarks
            }
        )
        dyn_model = update_inputs_outputs_dims(static_model, input_dims, output_dims)
        return dyn_model

    def __call__(self, img, threshold=0.5):
        self.orig_shape = img.shape[:2]
        if self.in_shape is None:
            print("self.in_shape is none")
            self.in_shape = self.orig_shape[::-1]
            print("self.in_shape set to ", self.in_shape)
        if not hasattr(self, "h_new"):  # First call, need to compute sizes
            self.w_new, self.h_new, self.scale_w, self.scale_h = self.transform(
                self.in_shape
            )

        blob = cv2.dnn.blobFromImage(
            img,
            scalefactor=1.0,
            size=(self.w_new, self.h_new),
            mean=(0, 0, 0),
            swapRB=False,
            crop=False,
        )
        if self.backend == "opencv":
            self.net.setInput(blob)
            heatmap, scale, offset, lms = self.net.forward(self.onnx_output_names)
        elif self.backend == "onnxrt":
            heatmap, scale, offset, lms = self.sess.run(
                self.onnx_output_names, {self.onnx_input_name: blob}
            )
        else:
            raise RuntimeError(f"Unknown backend {self.backend}")
        dets, lms = self.decode(
            heatmap, scale, offset, lms, (self.h_new, self.w_new), threshold=threshold
        )
        if len(dets) > 0:
            dets[:, 0:4:2], dets[:, 1:4:2] = (
                dets[:, 0:4:2] / self.scale_w,
                dets[:, 1:4:2] / self.scale_h,
            )
            lms[:, 0:10:2], lms[:, 1:10:2] = (
                lms[:, 0:10:2] / self.scale_w,
                lms[:, 1:10:2] / self.scale_h,
            )
        else:
            dets = np.empty(shape=[0, 5], dtype=np.float32)
            lms = np.empty(shape=[0, 10], dtype=np.float32)

        return dets, lms

    def transform(self, in_shape):
        h_orig, w_orig = self.orig_shape

        if in_shape == None:
            return h_orig, w_orig, 1, 1
        else:
            w_new, h_new = in_shape
            # Make spatial dims divisible by 32
            w_new, h_new = int(np.ceil(w_new / 32) * 32), int(np.ceil(h_new / 32) * 32)
            scale_w, scale_h = w_new / w_orig, h_new / h_orig
            return w_new, h_new, scale_w, scale_h

    def decode(self, heatmap, scale, offset, landmark, size, threshold=0.1):
        heatmap = np.squeeze(heatmap)
        scale0, scale1 = scale[0, 0, :, :], scale[0, 1, :, :]
        offset0, offset1 = offset[0, 0, :, :], offset[0, 1, :, :]
        c0, c1 = np.where(heatmap > threshold)
        boxes, lms = [], []
        if len(c0) > 0:
            for i in range(len(c0)):
                s0, s1 = (
                    np.exp(scale0[c0[i], c1[i]]) * 4,
                    np.exp(scale1[c0[i], c1[i]]) * 4,
                )
                o0, o1 = offset0[c0[i], c1[i]], offset1[c0[i], c1[i]]
                s = heatmap[c0[i], c1[i]]
                x1, y1 = (
                    max(0, (c1[i] + o1 + 0.5) * 4 - s1 / 2),
                    max(0, (c0[i] + o0 + 0.5) * 4 - s0 / 2),
                )
                x1, y1 = min(x1, size[1]), min(y1, size[0])
                boxes.append([x1, y1, min(x1 + s1, size[1]), min(y1 + s0, size[0]), s])
                lm = []
                for j in range(5):
                    lm.append(landmark[0, j * 2 + 1, c0[i], c1[i]] * s1 + x1)
                    lm.append(landmark[0, j * 2, c0[i], c1[i]] * s0 + y1)
                lms.append(lm)
            boxes = np.asarray(boxes, dtype=np.float32)
            lms = np.asarray(lms, dtype=np.float32)
            keep = self.nms(boxes[:, :4], boxes[:, 4], 0.3)
            boxes = boxes[keep, :]
            lms = lms[keep, :]
        return boxes, lms

    @staticmethod
    def nms(boxes, scores, nms_thresh):
        x1 = boxes[:, 0]
        y1 = boxes[:, 1]
        x2 = boxes[:, 2]
        y2 = boxes[:, 3]
        areas = (x2 - x1) * (y2 - y1)
        order = np.argsort(scores)[::-1]
        num_detections = boxes.shape[0]
        suppressed = np.zeros((num_detections,), dtype=np.uint8)
        for _i in range(num_detections):
            i = order[_i]
            if suppressed[i]:
                continue
            ix1 = x1[i]
            iy1 = y1[i]
            ix2 = x2[i]
            iy2 = y2[i]
            iarea = areas[i]

            for _j in range(_i + 1, num_detections):
                j = order[_j]
                if suppressed[j]:
                    continue

                xx1 = max(ix1, x1[j])
                yy1 = max(iy1, y1[j])
                xx2 = min(ix2, x2[j])
                yy2 = min(iy2, y2[j])
                w = max(0, xx2 - xx1)
                h = max(0, yy2 - yy1)

                inter = w * h
                ovr = inter / (iarea + areas[j] - inter)
                if ovr >= nms_thresh:
                    suppressed[j] = True
        keep = np.nonzero(suppressed == 0)[0]
        return keep

if __name__ == "__main__":
    main()