import matplotlib.pyplot as plt
import numpy as np
from IPython import embed
import matplotlib as matplotlib
import math
import scipy.integrate as integrate
from scipy import signal
from scipy.interpolate import interp1d
from scipy.interpolate import CubicSpline
import scipy as sp
import pickle
from scipy.spatial import distance
from myfunctions import *
import time
from matplotlib import gridspec
from matplotlib_scalebar.scalebar import ScaleBar
import matplotlib.mlab as ml
import scipy.integrate as si
import pandas as pd

def plot_beat_corr(eod_fr,eod_fe = np.arange(1, 1500, 5)):
    beats = eod_fe - eod_fr
    beat_corr = eod_fe % eod_fr
    beat_corr[beat_corr > eod_fr / 2] = eod_fr - beat_corr[beat_corr > eod_fr / 2]
    gs0 = gridspec.GridSpec(3, 1, height_ratios=[4, 1, 1], hspace=0.7)

    plt.figure(figsize=(4.5, 6))
    style = 'dotted'
    color_v = 'black'
    color_b = 'silver'
    # plt.subplot(3,1,1)
    plt.subplot(gs0[0])
    np.max(beats) / eod_fr
    plt.plot(beats, beats, color=color_b)

    # plt.axvline(x = -250, Linestyle = style,color = color_v)
    # plt.axvline(x = 250, Linestyle = style,color = color_v)
    # plt.axvline(x = 750, Linestyle = style,color = color_v)
    # plt.axvline(x = 1500, Linestyle = style)
    # plt.subplot(3,1,2)
    plt.xlabel('Beats [Hz]')
    plt.ylabel('Difference frequency [Hz]')
    #plt.subplot(gs0[1])
    plt.plot(beats, beat_corr, color=color_b)
    plt.axvline(x=-250, Linestyle=style, color=color_v)
    plt.axvline(x=250, Linestyle=style, color=color_v)
    plt.axvline(x=750, Linestyle=style, color=color_v)
    plt.xlabel('EOD adjusted beat [Hz]')

    # plt.axvline(x = 1250, Linestyle = style,color = color_v)
    # plt.axvline(x = 1500, Linestyle = style,color = color_v)
    mult = np.array(beats) / eod_fr + 1
    # plt.subplot(3,1,3)
    plt.xlabel('Beats [Hz]')
    plt.ylabel('EOD adj. beat [Hz]', fontsize = 10)
    #plt.subplot(gs0[2])
    #plt.plot(mult, beat_corr, color=color_b)
    # plt.axvline(x = 0, Linestyle = style)
    #plt.axvline(x=0.5, Linestyle=style, color=color_v)
    # plt.axvline(x = 1, Linestyle = style)
    #plt.axvline(x=1.5, Linestyle=style, color=color_v)
    #plt.axvline(x=2.5, Linestyle=style, color=color_v)
    #plt.xlabel('EOD multiples')
    #plt.ylabel('EOD adj. beat [Hz]', fontsize = 10)
    plt.tight_layout()
    plt.subplots_adjust(left = 0.2,top = 0.97)
    plt.savefig(fname='../pictures_highbeats/' + 'beat_corr_illustration')
    plt.show()
    embed()


eod_fr = 500
eod_fe = np.arange(1, eod_fr*3, 5)
plot_beat_corr(eod_fr,eod_fe = eod_fe)