import matplotlib.pyplot as plt
import numpy as np


no_of_data_points = 1000
tau_c= no_of_data_points/16
tau_phi = 2*tau_c
timing_delay = 5*tau_c


time_vector = np.arange(-no_of_data_points/2,no_of_data_points/2+1,1)
print(time_vector)

### DEFINE FUNCTIONS ####

def g2_chaotic(time_vector,tau_c):
	return 1 + np.exp(-2*np.abs(time_vector)/tau_c)


def g2x_int_chaotic(time_vector,tau_c,timing_delay):
	return np.exp(-2*timing_delay/tau_c) + np.exp(-2*np.abs(time_vector)/tau_c)

def g2_coherent(time_vector,tau_c):
	timing_delay_array = np.zeros(np.shape(time_vector))
	timing_delay_array.fill(1)
	return timing_delay_array 

def g2x_int_coherent(time_vector,tau_phi,timing_delay):
	timing_delay_array = np.zeros(np.shape(time_vector))
	timing_delay_array.fill(timing_delay)
	compared_time_vector = np.minimum(timing_delay_array,np.abs(time_vector))
	return np.exp(-2*np.abs(compared_time_vector)/tau_phi)


def g2_terms(g2_function,time_vector,tau_c,timing_delay):
	return 0.5*g2_function(time_vector,tau_c)+0.25*g2_function(time_vector+timing_delay,tau_c)+0.25*g2_function(time_vector-timing_delay,tau_c)

def g2x_int_term(g2x_int_function,time_vector,tau_c,timing_delay):
	return -0.5*g2x_int_function(time_vector,tau_c,timing_delay)


def g2x(g2_function,g2x_int_function,time_vector,tau_c,timing_delay):
	return 0.5*g2_function(time_vector,tau_c)+0.25*g2_function(time_vector+timing_delay,tau_c)+0.25*g2_function(time_vector-timing_delay,tau_c)-0.5*g2x_int_function(time_vector,tau_c,timing_delay)


plt.clf()
fig = plt.figure()
gs = fig.add_gridspec(3, 3, hspace=0.1,wspace=0.3)
axs = gs.subplots(sharex=True,)

# arrow config
arrow_length = 2
arrow_width=0.04
arrow_head_length= arrow_width*10

# CHAOTIC
axs[0,0].plot(time_vector/tau_c,g2_terms(g2_chaotic,time_vector,tau_c,timing_delay),"r-")
axs[0,1].plot(time_vector/tau_c,g2x_int_term(g2x_int_chaotic,time_vector,tau_c,timing_delay),"r-")
axs[0,2].plot(time_vector/tau_c,g2x(g2_chaotic,g2x_int_chaotic,time_vector,tau_c,timing_delay),"r-")
axs[0,0].arrow(-3,1.18,arrow_length,0,length_includes_head=True,width=arrow_width,head_length=arrow_head_length,facecolor="k")
axs[0,0].arrow(3,1.18,-1*arrow_length,0,length_includes_head=True,width=arrow_width,head_length=arrow_head_length,facecolor="k")
axs[0,0].text(1,1.3,r"$\tau_{c}$")

#COHERENT
axs[1,0].plot(time_vector/tau_c,g2_terms(g2_coherent,time_vector,tau_c,timing_delay),"g-")
axs[1,1].plot(time_vector/tau_c,g2x_int_term(g2x_int_coherent,time_vector,tau_c,timing_delay),"g-")
axs[1,2].plot(time_vector/tau_c,g2x(g2_coherent,g2x_int_coherent,time_vector,tau_c,timing_delay),"g-")
axs[1,1].arrow(-3.3,-0.18,arrow_length,0,length_includes_head=True,width=arrow_width,head_length=arrow_head_length,facecolor="k")
axs[1,1].arrow(3.3,-0.18,-1*arrow_length,0,length_includes_head=True,width=arrow_width,head_length=arrow_head_length,facecolor="k")
axs[1,1].text(1.5,-0.5,r"$\tau_{\phi}=2\tau_{c}$")


# Mixture
axs[2,0].plot(time_vector/tau_c,0.5*g2_terms(g2_coherent,time_vector,tau_c,timing_delay)+0.5*g2_terms(g2_chaotic,time_vector,tau_c,timing_delay),"b-")
axs[2,1].plot(time_vector/tau_c,0.5*g2x_int_term(g2x_int_coherent,time_vector,tau_c,timing_delay)+0.5*g2x_int_term(g2x_int_chaotic,time_vector,tau_c,timing_delay),"b-")
axs[2,2].plot(time_vector/tau_c,0.5*g2x(g2_coherent,g2x_int_coherent,time_vector,tau_phi,timing_delay)+0.5*g2x(g2_chaotic,g2x_int_chaotic,time_vector,tau_c,timing_delay),"b-")

# delta lines
axs[2,0].axvline(-5,color=(0,0,0,1),linestyle="dotted")
axs[2,0].axvline(5,color=(0,0,0,1),linestyle="dotted")
axs[2,0].text(-4.5,1.3,r"$-\Delta$")
axs[2,0].text(5.5,1.3,r"$\Delta$")


axs[0,0].axvline(-5,color=(0,0,0,1),linestyle="dotted")
axs[0,0].axvline(5,color=(0,0,0,1),linestyle="dotted")
axs[0,0].text(-4.5,1.3,r"$-\Delta$")
axs[0,0].text(5.5,1.3,r"$\Delta$")

axs[0,2].axvline(-5,color=(0,0,0,1),linestyle="dotted")
axs[0,2].axvline(5,color=(0,0,0,1),linestyle="dotted")
axs[0,2].text(-4.5,1.3,r"$-\Delta$")
axs[0,2].text(5.5,1.3,r"$\Delta$")

axs[2,2].axvline(-5,color=(0,0,0,1),linestyle="dotted")
axs[2,2].axvline(5,color=(0,0,0,1),linestyle="dotted")
axs[2,2].text(-4.5,1.3,r"$-\Delta$")
axs[2,2].text(5.5,1.3,r"$\Delta$")




axs[0,0].set_title(r"$A: g^{(2)}$ terms")
axs[0,1].set_title(r"$B: g^{(2X)}_{int}$ term")
axs[0,2].set_title(r"$g^{(2X)}= A + B$")

# axs[0,0].set_ylabel("Chaotic")
# axs[1,0].set_ylabel("Coherent")
# axs[2,0].set_ylabel("50:50 mixture")

axs[0,1].text(-8,0.4,"chaotic")
axs[1,1].text(-8,0.4,"coherent")
axs[2,1].text(-8,0.4,"50:50 mixture")

axs[2,0].set_xlabel(r"$(t_{2}-t_{1})/\tau_{c}$")
axs[2,1].set_xlabel(r"$(t_{2}-t_{1})/\tau_{c}$")
axs[2,2].set_xlabel(r"$(t_{2}-t_{1})/\tau_{c}$")


for i in range(3):
	for j in range(3):
		if j == 0:
			axs[i,j].set_ylim([0.45,1.55])
		elif j != 2:
			axs[i,j].set_ylim([-0.55,0.55])
		else:
			axs[i,j].set_ylim([0.45,1.55])


plt.tight_layout()
plt.savefig("theoretical_simulation.png",format="png",dpi=600)
# plt.show()