import nixio as nix
import os
from IPython import embed
#from utility import *
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import matplotlib.mlab as ml
import scipy.integrate as si
from scipy.ndimage import gaussian_filter
from IPython import embed
from myfunctions import *
from myfunctions import auto_rows
from myfunctions import remove_tick_ymarks
import string
from myfunctions import default_settings

def plot_amp(ax, mean1, dev,lw = 0.8,nrs = [3,4,5],nr_size = 12,name = 'amp',nr = 1, wide = 'gainsboro', middle = 'darkgrey', narrow = 'black'):
    np.unique(mean1['type'])
    all_means = mean1[mean1['type'] == name +' mean']
    original = all_means[all_means['dev'] == 'original']
    #m005 = all_means[all_means['dev'] == '005']
    m05 = all_means[all_means['dev'] == '05']
    m2 = all_means[all_means['dev'] == '2']
    # fig, ax = plt.subplots(nrows=4, ncols = 3, sharex=True)
    versions = [original, m05, m2] #m005,
    lim = [[]]*len(versions)
    for i in range(len(versions)):
        keys = [k for k in versions[i]][2::]
        try:
            data = np.array(versions[i][keys])[0]
        except:
            break
        axis = np.arange(0, len(data), 1)
        axis_new = axis * 1
        similarity = [keys, data]
        sim = np.argsort(similarity[0])
        # similarity[sim]
        all_means = mean1[mean1['type'] == name+' std']
        std = all_means[all_means['dev'] == dev[i]]
        std = np.array(std[keys])[0]
        #ax[1, 1].set_ylabel('Modulation depth')
        #ax[nr,i].set_title(dev[i] + ' ms')
        all_means = mean1[mean1['type'] == name+' 95']
        std95 = all_means[all_means['dev'] == dev[i]]
        std95 = np.array(std95[keys])[0]
        all_means = mean1[mean1['type'] == name+' 05']
        std05 = all_means[all_means['dev'] == dev[i]]
        std05 = np.array(std05[keys])[0]
        ax[nr, i].text(-0.1, 1.1, string.ascii_uppercase[nrs[i]], transform=ax[nr,i].transAxes,
                   size=nr_size, weight='bold')

        ax[nr,i].fill_between(np.array(keys)[sim], list(std95[sim]), list(std05[sim]),
                              color=wide)
        ax[nr,i].fill_between(np.array(keys)[sim], list(data[sim] + std[sim]), list(data[sim] - std[sim]),
                              color=middle)
        if i != 0:
            ax[nr, i] = remove_tick_ymarks(ax[nr,i])
        # ax[i].plot(data_tob.ff, data_tob.fe, color='grey', linestyle='--', label='AMf')
        ax[nr,i].plot(np.array(keys)[sim], data[sim],linewidth = lw,  color=narrow)
        lim[i] = np.nanmax(std95[sim])
        # ax[0].plot(data1.x, data1.freq20, color=colors[1], label='20 %')
    for i in range(len(versions)):
        ax[nr, i].set_ylim(0,np.nanmax(lim))
    #embed()
    return ax

def create_beat_corr(hz_range, eod_fr):
    beat_corr = hz_range%eod_fr
    beat_corr[beat_corr>eod_fr/2] = eod_fr[beat_corr>eod_fr/2] - beat_corr[beat_corr>eod_fr/2]
    return beat_corr


def plot_mean_cells(type = '',tob_lw = 0.8):
    mean1 = pd.read_pickle('mean'+type+'.pkl')

    whole_page_width = 6.28
    intermediate_length = 4.7

    default_settings([0],intermediate_width = 6.29,intermediate_length = 3.7, ts = 10, ls = 9, fs = 9)

    x = np.arange(0, 2550, 50)
    corr = create_beat_corr(x, np.array([500] * len(x)))

    np.unique(mean1['type'])
    all_means = mean1[mean1['type'] == 'max mean']

    #versions = [[]]*len(dev)
    #for i in range(len(dev)):
    original = all_means[all_means['dev'] == 'original']
    m005 = all_means[all_means['dev'] == '005']
    m05 = all_means[all_means['dev'] == '05']
    m2 = all_means[all_means['dev'] == '2']
    dev = ['original',  '05', '2']#'005',
    titels = ['binary','0.5 ms','2 ms']
    fig, ax = plt.subplots(nrows=3, ncols=3, sharex=True)
    plt.suptitle(type)
    #embed()
    versions = [original,  m05, m2]#m005,
    nrs = [0,1,2]
    nr_size = 12
    lw = 1.2
    for i in range(len(versions)):
        keys = [k for k in versions[i]][2::]
        try:
            data = np.array(versions[i][keys])[0]
        except:
            #embed()
            break
        axis = np.arange(0, len(data), 1)
        axis_new = axis * 1
        similarity = [keys, data]
        sim = np.argsort(similarity[0])
        # similarity[sim]
        all_means = mean1[mean1['type'] == 'max std']
        std = all_means[all_means['dev'] == dev[i]]
        std = np.array(std[keys])[0]
        #ax[0,i].set_title(dev[i] +' ms')
        ax[0, i].set_title(titels[i])
        ax[ 0,0].set_ylabel('MPF [EODf]')
        all_means = mean1[mean1['type'] == 'max 95']
        std95 = all_means[all_means['dev'] == dev[i]]
        std95 = np.array(std95[keys])[0]
        all_means = mean1[mean1['type'] == 'max 05']
        std05 = all_means[all_means['dev'] == dev[i]]
        std05 = np.array(std05[keys])[0]
        #std[np.where(list(data[sim] + std[sim])>std95[sim])] = std95[sim][np.where(list(data[sim] + std[sim])>std95[sim])]
        #embed()
        middle  = 'pink'#'LightSalmon'
        wide='lightpink'#mistyrose'
        wide = 'gainsboro'
        narrow = 'crimson'
        tob_color = 'black'
        ax[0,i].fill_between(np.array(keys)[sim], list(std95[sim]), list(std05[sim]),
                              color=wide)# alpha = 0.45
        #embed()

        ax[0,i].fill_between(np.array(keys)[sim], list(data[sim] + std[sim]), list(data[sim] - std[sim]), color=middle)

        ax[0,i].plot(x / 500, corr / 500 + 1, color=tob_color, linestyle = '--', linewidth = tob_lw,  label='AMf')
        ax[0,i].plot(np.array(keys)[sim], data[sim], color=narrow, linewidth = lw)
        # plt.fill_between(np.array([0,1]),np.array([0,0]),np.array([1,1]))
        ax[0, i].text(-0.1, 1.1, string.ascii_uppercase[nrs[i]], transform=ax[0,i].transAxes,
                   size=nr_size, weight='bold')
        #embed()
        if i != 0:
            ax[0, i] = remove_tick_ymarks(ax[0, i])
        # ax[0].plot(data1.x, data1.freq20, color=colors[1], label='20 %')
    #ax[0, 2].legend(bbox_to_anchor=(0.7, 1, 0.7, .1), loc='lower left',
    #                ncol=3, mode="expand", borderaxespad=0.)

    #ax = plot_amp(ax, mean1, dev,name = 'amp',nr = 2)
    #ax[ 1,0].legend(bbox_to_anchor=(0.3, 1, 0.7, .1), loc='lower left',
    #                ncol=3, mode="expand", borderaxespad=0.)

    ax = plot_amp(ax,mean1, dev,lw = lw,nrs = [3,4,5],nr_size = nr_size,name='amp max', nr=1,wide = 'gainsboro', middle = 'lightblue', narrow = 'steelblue')
    ax = plot_amp(ax,mean1, dev,lw = lw, nrs = [6,7,8],nr_size = nr_size,name='tob max', nr=2,wide = 'gainsboro', middle = 'grey', narrow = 'black')

    ax[ 0,0].set_ylabel('EOD multiples')
    #ax[1, 0].set_ylabel('Modulation depth')
    ax[1, 0].set_ylabel('Modulation')
    #ax[2, 0].set_ylabel('Whole modulation ')
    plt.subplots_adjust(hspace = 0.35)
    for i in range(3):

        ax[i,0].spines['right'].set_visible(False)
        ax[i,0].spines['top'].set_visible(False)
        ax[i,1].spines['right'].set_visible(False)
        ax[i,1].spines['top'].set_visible(False)
        ax[i,2].spines['right'].set_visible(False)
        ax[i,2].spines['top'].set_visible(False)
        #ax[i,3].spines['right'].set_visible(False)
        #ax[i,3].spines['top'].set_visible(False)
    #for i in range(3):
    ax[2, 1].set_xlabel('EOD multiples')
        #ax[2, i].set_ylim([0, 370])
        #ax[2,i].set_ylim([0, 4700])
    #ax[2, 1].set_xlabel('stimulus frequency [EODf]')
    plt.subplots_adjust(bottom = 0.13, right = 0.96)
    #fig.tight_layout()f
    # fig.label_axes()
    return fig

if __name__ == "__main__":

    #type = ''
    type = ''
    type = 'simulation' #'simulation',
    nrs = ['simulation','sinz','sinz3','sinz5','sinz7','sinz9','' ,]# 'sinz11'
    for n in range(len(nrs)):
        type = str(nrs[n])
        print(type)
        fig = plot_mean_cells(type = type)
        #fig.savefig()
        plt.savefig('MPFmodulation_tob'+type+'.pdf')
        plt.savefig('../highbeats_pdf/MPFmodulation_tob'+type+'.pdf')
    plt.show()
    # plt.close()