from utils_suseptibility import *

def model_full():

    plot_style()
    default_figsize(column=2, length=2.3)
    grid = gridspec.GridSpec(1, 1, wspace=0.6, bottom = 0.1,hspace=0.15, top=0.95, left=0.075, right=0.87)

    axes = []
    ##################################
    # model part
    ls = '--'
    lw = 0.5
    ax = plt.subplot(grid[0])

    axes.append(ax)

    perc,im,stack_final  = plt_model_big(ax, ls = ls, lw = 0.5)
    fr_waves = 139
    fr_noise = 120
    f1 = 33
    f2 = 139

    #embed()
    ax.plot(fr_noise * f1/fr_waves, fr_noise*f2/fr_waves, 'o', ms = 5, markeredgecolor = 'orange', markerfacecolor="None")
    ax.plot(-fr_noise * f1 / fr_waves, fr_noise * f2 / fr_waves, 'o', ms = 5, markeredgecolor='pink', markerfacecolor="None")

    # if len(cbar) > 0:
    ###############################
    # data part


    data_extra = False
    if data_extra:
        ax = plt.subplot(grid[0])
        axes.append(ax)
        cell = '2012-07-03-ak-invivo-1'
        mat_rev,stack_final_rev = load_stack_data_susept(cell, save_name = version_final(), end = '_revQuadrant_')
        mat, stack = load_stack_data_susept(cell, save_name=version_final(), end = '')
        #embed()
        #try:
        full_matrix = create_full_matrix2(np.array(mat),np.array(mat_rev))
        #except:
        #    print('full matrix something')
        #    embed()

        stack_final = get_axis_on_full_matrix(full_matrix, mat)

        abs_matrix = np.abs(stack_final)
        #embed()
        #if np.
        abs_matrix, add_nonlin_title, resize_val = rescale_colorbar_and_values(abs_matrix)


        ax.axhline(0, color = 'white', linestyle = ls, linewidth = lw)
        ax.axvline(0, color='white', linestyle = ls, linewidth = lw)
        im = plt_RAM_perc(ax, perc, abs_matrix)
        cbar, left, bottom, width, height = colorbar_outside(ax, im, add=5, width=0.01)

        set_clim_same_here([im], mats=[abs_matrix], lim_type='up', nr_clim='perc', clims='', percnr=95)

        #clim = im.get_clim()
        #if clim[1]> 1000:
        #todo: change clim values with different Hz values
        #embed()
        cbar.set_label(nonlin_title(add_nonlin_title = ' ['+add_nonlin_title), rotation=90, labelpad=8)

        set_ylabel_arrow(ax, xpos = -0.07, ypos = 0.97)
        set_xlabel_arrow(ax, xpos=1, ypos=-0.07)
        '''    eod_fr, stack_spikes = plt_data_suscept_single(ax, cbar_label, cell, cells, f, fig, file_names_exclude, lp, title,
                                                       width)'''
        cbar, left, bottom, width, height = colorbar_outside(ax, im, add=5, width=0.01)

    #print('finished model_full')
    fig = plt.gcf()
    #axes = plt.gca()
    #fig.tag(axes[::-1], xoffs=-4.5, yoffs=0.4)  # ax_ams[3],

    save_visualization()


def load_stack_data_susept(cell, save_name, end = ''):
    load_name = load_folder_name('calc_RAM') + '/' + save_name+end
    add = '_cell' + cell  +end# str(f)  # + '_amp_' + str(amp)
    #embed()
    stack_cell = load_data_susept(load_name + '_' + cell + '.pkl', load_name + '_' + cell, add=add,
                                  load_version='csv')
    file_names_exclude = get_file_names_exclude()
    stack_cell = stack_cell[~stack_cell['file_name'].isin(file_names_exclude)]
    # if len(stack_cell):
    file_names = stack_cell.file_name.unique()
    #embed()
    file_names = exclude_file_name_short(file_names)
    cut_off_nr = get_cutoffs_nr(file_names)
    try:
        maxs = list(map(float, cut_off_nr))
    except:
        embed()
    file_names = file_names[np.argmax(maxs)]
    #embed()
    stack_file = stack_cell[stack_cell['file_name'] == file_names]
    amps = [np.min(stack_file.amp.unique())]
    amps = restrict_punits(cell, amps)
    amp = np.min(amps)#[0]
    # for amp in amps:
    stack_amps = stack_file[stack_file['amp'] == amp]

    lengths = stack_amps.stimulus_length.unique()
    trial_nr_double = stack_amps.trial_nr.unique()
    trial_nr = np.max(trial_nr_double)
    stack_final = stack_amps[
        (stack_amps['stimulus_length'] == np.max(lengths)) & (stack_amps.trial_nr == trial_nr)]
    mat, new_keys = get_mat_susept(stack_final)
    return mat,stack_final


if __name__ == '__main__':
    model_full()