import numpy as np
import matplotlib.pyplot as plt

lsize = 20

fig,ax=plt.subplots(nrows=2,ncols=2,figsize=(20,10))
plt.rc('text', usetex=True)

# Loading data

data = np.load("data.npz")

L = 9 # length of domain [0,L mm]
Lrho = 3 # length of domain for rho. We shortened it, in order to speed up the simulations
I = 9000 # dx=1micron
Irho = 3000 # dx=1micron

X=np.linspace(0,Lrho,Irho)

# Subfigure A

axA = ax[0,0]

#N_data = data["N_data_AC"] # saved frames of the evolution of N. shape : 50x9000
rho_data = data["rho_data_AC"] # saved frames of the evolution of rho. shape : 50x3000
neutralfraction_data = data["neutralfraction_data_AC"] # saved frames of the evolution of the neutral fraction. shape : 50x3000

rho = rho_data[0,:] # taking the 0th saved frame
neutralfraction = neutralfraction_data[0,:] # taking the 0th saved frame
i_thresh = 850 # cells left to index i_thresh have been labeled

color = 'black'
axA.set_xlabel(r'\textit{Distance from center (mm})',size=lsize)
axA.set_ylabel(r'\textit{Cell density $\rho$ (cell/mm$^{-2}$)}', color=color,size=lsize)
axA.fill_between(X,0,neutralfraction, color='yellow', label=r'Cells initially left of the peak')
axA.fill_between(X,neutralfraction,rho, color='cyan', label=r'Cells initially right of the peak')
axA.plot(X, rho, ls='-', color=color)
axA.set_ylim(0,200)
axA.set_xlim(0,Lrho)
axA.tick_params(axis='y', labelcolor=color,labelsize=lsize)
axA.tick_params(axis='x',labelsize=lsize)
axA.legend(fontsize=lsize, loc=1)
axAbis = axA.twinx()
color = 'blue'
Y = 100*np.ones(Irho)
Y[i_thresh:] *= 0
axAbis.plot(X, Y, color=color)
axAbis.set_ylim(0,150)
axAbis.set_xlim(0,Lrho)
axAbis.tick_params(axis='y', labelcolor=color,labelsize=lsize)
labels = [0, 25, 50, 75, 100]
axAbis.set_yticks(labels)

# Subfigure C

axC = ax[1,0]

rho = rho_data[49,:] # taking the 49th saved frame
neutralfraction = neutralfraction_data[49,:] # taking the 49th saved frame

color = 'black'
axC.set_xlabel(r'\textit{Distance from center (mm})',size=lsize)
axC.set_ylabel(r'\textit{Cell density $\rho$ (cell/mm$^{-2}$)}', color=color,size=lsize)
axC.fill_between(X,0,neutralfraction, color='yellow', label=r'Cells initially left of the peak')
axC.fill_between(X,neutralfraction,rho, color='cyan', label=r'Cells initially right of the peak')
axC.plot(X, rho, ls='-', color=color)
axC.set_ylim(0,200)
axC.set_xlim(0,Lrho)
axC.tick_params(axis='y', labelcolor=color,labelsize=lsize)
axC.tick_params(axis='x',labelsize=lsize)
axC.legend(fontsize=lsize, loc=1)
axCbis = axC.twinx()
color = 'blue'
Y = 100*np.ones(I)
Y[i_thresh:] *= 0
axCbis.plot(X, 100*neutralfraction/rho, color=color)
axCbis.set_ylim(0,150)
axCbis.set_xlim(0,Lrho)
axCbis.tick_params(axis='y', labelcolor=color,labelsize=lsize)
labels = [0, 25, 50, 75, 100]
axCbis.set_yticks(labels)

# Subfigure B

axB = ax[0,1]

#N_data = data["N_data_BD"] # saved frames of the evolution of N. shape : 50x9000
rho_data = data["rho_data_BD"] # saved frames of the evolution of rho. shape : 50x3000
neutralfraction_data = data["neutralfraction_data_BD"] # saved frames of the evolution of the neutral fraction. shape : 50x3000

rho = rho_data[0,:] # taking the 0th saved frame
neutralfraction = neutralfraction_data[0,:] # taking the 0th saved frame
i_thresh = 560 # cells left to index i_thresh have been labeled

color = 'black'
axB.set_xlabel(r'\textit{Distance from center (mm})',size=lsize)
axB.set_ylabel(r'\textit{Cell density $\rho$ (cell/mm$^{-2}$)}', color=color,size=lsize)
axB.fill_between(X,0,neutralfraction, color='yellow', label=r'Cells initially left of the peak')
axB.fill_between(X,neutralfraction,rho, color='cyan', label=r'Cells initially right of the peak')
axB.plot(X, rho, ls='-', color=color)
axB.set_ylim(0,100)
axB.set_xlim(0,Lrho)
axB.tick_params(axis='y', labelcolor=color,labelsize=lsize)
axB.tick_params(axis='x',labelsize=lsize)
axB.legend(fontsize=lsize, loc=1)
axBbis = axB.twinx()
color = 'blue'
Y = 100*np.ones(Irho)
Y[i_thresh:] *= 0
axBbis.plot(X, Y, color=color)
axBbis.set_ylim(0,150)
axBbis.set_xlim(0,Lrho)
axBbis.tick_params(axis='y', labelcolor=color,labelsize=lsize)
labels = [0, 25, 50, 75, 100]
axBbis.set_yticks(labels)

# Subfigure D

axD = ax[1,1]

rho = rho_data[49,:] # taking the 49th saved frame
neutralfraction = neutralfraction_data[49,:] # taking the 49th saved frame

color = 'black'
axD.set_xlabel(r'\textit{Distance from center (mm})',size=lsize)
axD.set_ylabel(r'\textit{Cell density $\rho$ (cell/mm$^{-2}$)}', color=color,size=lsize)
axD.fill_between(X,0,neutralfraction, color='yellow', label=r'Cells initially left of the peak')
axD.fill_between(X,neutralfraction,rho, color='cyan', label=r'Cells initially right of the peak')
axD.plot(X, rho, ls='-', color=color)
axD.set_ylim(0,100)
axD.set_xlim(0,Lrho)
axD.tick_params(axis='y', labelcolor=color,labelsize=lsize)
axD.tick_params(axis='x',labelsize=lsize)
axD.legend(fontsize=lsize, loc=1)
axDbis = axD.twinx()
color = 'blue'
Y = 100*np.ones(I)
Y[i_thresh:] *= 0
axDbis.plot(X, 100*neutralfraction/rho, color=color)
axDbis.set_ylim(0,150)
axDbis.set_xlim(0,Lrho)
axDbis.tick_params(axis='y', labelcolor=color,labelsize=lsize)
labels = [0, 25, 50, 75, 100]
axDbis.set_yticks(labels)

# Saving figure

fig.tight_layout(pad=3)

fig.savefig('FigSI20.png',dpi=200)







