"""
description: offline analysis of kiap data
author: Ioannis Vlachos
date:   16.04.2019

Copyright (c) 2019 Ioannis Vlachos.
All rights reserved."""

import re
import os
import glob
import importlib
import urllib
import time

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from scipy import signal
from scipy import stats
# from cachetools import cached, LRUCache, TTLCache
from neo.io.blackrockio import BlackrockIO

import aux
from analytics import analytics1
importlib.reload(analytics1)


# @cached(cache={})
def read_data():
    '''returns first array1 and then array2 channels'''

    data = reader.get_analogsignal_chunk(i_stop=eval(str(params.lfp.i_stop)),
                                channel_indexes=params.lfp.array1 + params.lfp.array2)

    return data


params = aux.load_config()
ch_nr1 = len(params.lfp.array1)
ch_nr2 = len(params.lfp.array2)
psth_win = params.lfp.psth_win

# file_name = '/home/vlachos/data_kiap/20190321-135311/20190321-135311-001'
path = '/kiap/data/neural/k01/20190414-123429/'
file_name0 = '20190414-123429-001.ns2'
file_name = path + file_name0

plt.ioff()

# if not 'data0' in locals():
for fid in range(1, 76):
# for fid in range(62,63):

    file_name0 = f'20190414-123429-{fid:03d}.ns2'
    file_name = path + file_name0
    print(file_name)

    reader = BlackrockIO(filename=file_name)
    data0 = read_data()  # (samples, channels)

    data0 = data0.copy().astype(np.float64)
    # trials_msk = get_valid_trials(triggers_no_ns2, triggers_yes_ns2)

    array1 = params.lfp.array1
    array2 = params.lfp.array2
    # data1[data1 > params.lfp.artifact_thr] = 0
    # data1[data1 < -params.lfp.artifact_thr] = 0

    # data0 = data0 - np.mean(data0, 1)[:, np.newaxis]

    # reader.nev_data['Comments'][0]

    data0 = stats.zscore(data0)
    data01 = data0[:, :len(array1)]
    data02 = data0[:, len(array1):len(array1) + len(array2)]

    data1 = np.zeros((data01.shape[0], data01.shape[1], 3))     # (samples, channels, 3)
    data2 = np.zeros((data02.shape[0], data02.shape[1], 3))     # (samples, channels, 3)
    # Y1 = np.zeros((data01.shape[0], data01.shape[1], 3))

    # data1[:, :, 0], _, _ = analytics1.bw_filter(data01.T, params.lfp.filter_fc_lb, params.lfp.fs, params.lfp.filter_order_lb, plot=params.lfp.plot.filters)
    # data1[:, :, 1], _, _ = analytics1.bw_filter(data01.T, params.lfp.filter_fc_mb, params.lfp.fs, params.lfp.filter_order_mb, plot=params.lfp.plot.filters)
    # data1[:, :, 2], _, _ = analytics1.bw_filter(data01.T, params.lfp.filter_fc_hb, params.lfp.fs, params.lfp.filter_order_hb, plot=params.lfp.plot.filters)

    # data2[:, :, 0], _, _ = analytics1.bw_filter(data02.T, params.lfp.filter_fc_lb, params.lfp.fs, params.lfp.filter_order_lb, plot=params.lfp.plot.filters)
    # data2[:, :, 1], _, _ = analytics1.bw_filter(data02.T, params.lfp.filter_fc_mb, params.lfp.fs, params.lfp.filter_order_mb, plot=params.lfp.plot.filters)
    # data2[:, :, 2], _, _ = analytics1.bw_filter(data02.T, params.lfp.filter_fc_hb, params.lfp.fs, params.lfp.filter_order_hb, plot=params.lfp.plot.filters)

    # Y01, ff01 = analytics1.calc_fft(data01, params.lfp.fs, i_stop=1e5, axis=0)
    # Y02, ff02 = analytics1.calc_fft(data02, params.lfp.fs, i_stop=1e5, axis=0)
    # Y1, ff1 = analytics1.calc_fft(data1, params.lfp.fs, i_stop=1e5, axis=0)
    # Y2, ff2 = analytics1.calc_fft(data2, params.lfp.fs, i_stop=1e5, axis=0)

    #  SPECTROGRAMS
    # ----------------------------------

    nperseg = 20000

    ff, tt, Sxx1 = signal.spectrogram(data01, params.lfp.fs, axis=0, nperseg=nperseg, noverlap=int(nperseg * 0.5))
    ff, tt, Sxx2 = signal.spectrogram(data02, params.lfp.fs, axis=0, nperseg=nperseg, noverlap=int(nperseg * 0.5))

    f_idx = (ff > 0) & (ff < 200)

    fig1 = plt.figure(1, figsize=(14, 10))
    plt.clf()

    for ii in range(array1.__len__()):
        # plt.plot()
        ax = plt.subplot(8, 8, ii+1)
        ax.title.set_text(reader.header['signal_channels'][array1[ii] + 1][0])
        # plt.plot(ff01[fidx], Y01[fidx])
        plt.pcolormesh(tt, ff[f_idx], np.log10(Sxx1[f_idx, ii, :]))
        # plt.pcolormesh(tt, ff[f_idx], np.log10(Sxx2[f_idx, ii, :]))
        # plt.pcolormesh(tt, ff[f_idx], np.log10(Sxx1[f_idx, ii, :]/Sxx1[f_idx,:,:].mean(1)))
        print('plotting array1 channel: ', ii)
        if ii < 56:
            ax.set_xticks([])
        if np.mod(ii, 8) != 0:
            ax.set_yticks([])

    fname1 = path + os.path.splitext(file_name0)[0] + '-arr1.png'
    fig1.savefig(fname1)
    print(fname1, 'Saved')

    fig2 = plt.figure(1, figsize=(14, 10))
    plt.clf()

    for ii in range(array2.__len__()):
        # plt.plot()
        ax = plt.subplot(8, 8, ii+1)
        ax.title.set_text(reader.header['signal_channels'][array2[ii] + 1][0])
        # plt.plot(ff01[fidx], Y01[fidx])
        # plt.pcolormesh(tt, ff[f_idx], np.log10(Sxx1[f_idx, ii, :]))
        plt.pcolormesh(tt, ff[f_idx], np.log10(Sxx2[f_idx, ii, :]))
        # plt.pcolormesh(tt, ff[f_idx], np.log10(Sxx1[f_idx, ii, :]/Sxx1[f_idx,:,:].mean(1)))
        print('plotting array2 channel: ', ii)
        if ii < 56:
            ax.set_xticks([])
        if np.mod(ii, 8) != 0:
            ax.set_yticks([])

    fname2 = path + os.path.splitext(file_name0)[0]+'-arr2.png'
    fig2.savefig(fname2)
    print(fname2, 'Saved')

print('lfp_analysis2 script finished!')
# plt.show()