from matplotlib import pyplot as plt
import numpy as np
import matplotlib.gridspec as gridspec

path = "data/BH_phase-diagrams-and-loss/data_rike/"

x = np.load(path + "x.npy")
y = np.load(path + "y.npy")
es = np.load(path + "es.npy")
dd = np.load(path + "dd.npy")
cost0 = np.load(path + "cost_bh_L12_trash2_d48_v5_seed17.npy")
cost1 = np.load(path + "cost_bh_L12_trash4_d25_v47_seed16.npy")
cost2 = np.load(path + "cost_bh_L12_trash6_d10_v10_seed22.npy")

vs = np.logspace(-2,2,50)
ds = np.linspace(-0.95,0.95,50)

cbar = np.zeros((2,2), dtype="object")

fig,axs = plt.subplots(ncols=2, nrows=2,figsize=(10,6),sharex=True, sharey=True)
cmap = plt.get_cmap("plasma") #'viridis'
cbar_labelsize = 14
cbar = np.zeros((2,2), dtype="object")

ax = axs[0,0]
im = ax.pcolormesh(x,y,dd.T, cmap=cmap, shading="auto",rasterized=True) #,rasterized=True necessary for pdf export
ax.set_xscale("log")
cbar[0,0] = fig.colorbar(im, ax=ax)
cbar[0,0].ax.tick_params(labelsize=cbar_labelsize)



ax = axs[0,1]
im = ax.pcolormesh(x,y,es.T, cmap=cmap, shading="auto",rasterized=True)
ax.set_xscale("log")
cbar[0,1] = fig.colorbar(im, ax=ax)
cbar[0,1].ax.tick_params(labelsize=cbar_labelsize)



ax = axs[1,0]
im = ax.pcolormesh(x,y,cost2.T,vmin=0, cmap=cmap, shading="auto",rasterized=True)
ax.set_xscale("log")
cbar[1,0] = fig.colorbar(im, ax=ax)
cbar[1,0].ax.tick_params(labelsize=cbar_labelsize)
ax.plot(vs[10],ds[10],"x", color="magenta")

ax = axs[1,1]
im = ax.pcolormesh(x,y,cost1.T,vmin=0, cmap=cmap, shading="auto",rasterized=True)
ax.set_xscale("log")
cbar[1,1] = fig.colorbar(im, ax=ax)
cbar[1,1].ax.tick_params(labelsize=cbar_labelsize)
ax.plot(vs[47],ds[25],"x", color="magenta")


for ax in axs.flatten():
    ax.tick_params(labelsize=14)
axs[1,0].set_xlabel(r"$V$",fontsize=18)
axs[1,1].set_xlabel(r"$V$",fontsize=18)
axs[0,0].set_ylabel(r"$\delta t$",fontsize=18)
axs[1,0].set_ylabel(r"$\delta t$",fontsize=18)
plt.tight_layout()

# has to happen after tight_layout()
axs[0,0].text(-0.25,0.9,"a", fontweight="bold", size=18, transform = axs[0,0].transAxes)
axs[1,0].text(-0.25,0.9,"c", fontweight="bold", size=18, transform = axs[1,0].transAxes)
axs[0,1].text(-0.1,0.9,"b", fontweight="bold", size=18, transform = axs[0,1].transAxes)
axs[1,1].text(-0.1,0.9,"d", fontweight="bold", size=18, transform = axs[1,1].transAxes)


cbar[0,0].ax.set_title("$\\langle n_i n_j \\rangle$")
axs[0,0].set_title("density-density correlations", fontsize=14)

cbar[0,1].ax.set_title("$\Delta \\lambda$")
axs[0,1].set_title("Entanglement spectrum degeneracy", fontsize=14)

cbar[1,0].ax.set_title("cost")
cbar[1,1].ax.set_title("cost")

plt.savefig("plots/BH_replot.pdf", bbox_inches='tight')

fig,aa = plt.subplots(figsize=(15,5.7),sharey="row",dpi=220)
ax00 = plt.subplot2grid(shape=(2,5), loc=(0, 0), rowspan=1, colspan=1, xticklabels=[])
ax10 = plt.subplot2grid(shape=(2,5), loc=(1, 0), rowspan=1, colspan=1, xticklabels=[])
ax01 = plt.subplot2grid(shape=(2,5), loc=(0, 1), rowspan=2, colspan=2, xticklabels=[], yticklabels=[])
ax02 = plt.subplot2grid(shape=(2,5), loc=(0, 3), rowspan=2, colspan=2, xticklabels=[], yticklabels=[])

cmap = plt.get_cmap("plasma") #'viridis'
cbar_labelsize = 10
cbar = np.zeros((2,2), dtype="object")

axs=np.array([[ax00,ax10],[ax01,ax02]])
print(axs[0,0])

ax = axs[0,0]
im = ax.pcolormesh(x,y,dd.T, cmap=cmap, shading="auto",rasterized=True) #,rasterized=True necessary for pdf export
#im.ax.tick_params(labelsize=cbar_labelsize)
ax.set_xscale("log")
cbar[0,0] = fig.colorbar(im, ax=ax)
cbar[0,0].ax.tick_params(labelsize=cbar_labelsize)

ax = axs[0,1]
im = ax.pcolormesh(x,y,es.T, cmap=cmap, shading="auto",rasterized=True)
ax.set_xscale("log")
cbar[0,1] = fig.colorbar(im, ax=ax)
cbar[0,1].ax.tick_params(labelsize=cbar_labelsize)

ax = axs[1,0]
im = ax.pcolormesh(x,y,cost2.T,vmin=0, cmap=cmap, shading="auto",rasterized=True)
ax.set_xscale("log")
cbar[1,0] = fig.colorbar(im, ax=ax)
cbar[1,0].ax.tick_params(labelsize=cbar_labelsize)
ax.plot(vs[10],ds[10],"X", color="magenta",markersize=15)

ax = axs[1,1]
im = ax.pcolormesh(x,y,cost1.T,vmin=0, cmap=cmap, shading="auto",rasterized=True)
ax.set_xscale("log")
cbar[1,1] = fig.colorbar(im, ax=ax)
cbar[1,1].ax.tick_params(labelsize=cbar_labelsize)
ax.plot(vs[47],ds[25],"X", color="magenta",markersize=15)

for ax in axs.flatten():
    ax.tick_params(labelsize=10)
axs[1,0].set_xlabel(r"$V$",fontsize=12)
axs[0,1].set_ylabel(r"$\delta J$",fontsize=12)
axs[1,1].set_xlabel(r"$V$",fontsize=12)
axs[0,0].set_ylabel(r"$\delta J$",fontsize=12)
axs[1,0].set_ylabel(r"$\delta J$",fontsize=12)
#plt.tight_layout()

# has to happen after tight_layout()
axs[0,0].text(0.01,0.89,"a", fontweight="bold", size=18, transform = axs[0,0].transAxes, color="white")
axs[1,0].text(0.02,0.95,"c", fontweight="bold", size=18, transform = axs[1,0].transAxes)
axs[0,1].text(0.01,0.88,"b", fontweight="bold", size=18, transform = axs[0,1].transAxes)
axs[1,1].text(0.02,0.94,"d", fontweight="bold", size=18, transform = axs[1,1].transAxes)

axs[1,0].text(0.23,0.05,"MI", fontweight="bold", size=18, transform = axs[1,0].transAxes, color="white")
axs[1,0].text(0.23,0.8,"TMI", fontweight="bold", size=18, transform = axs[1,0].transAxes)
axs[1,0].text(0.77,0.5,"CDW", fontweight="bold", size=18, transform = axs[1,0].transAxes)

#cbar[0,0].ax.set_title("$\\langle n_i n_j \\rangle$")
cbar[0,0].ax.set_title("$O_{CDW}$")
axs[0,0].set_title("DMRG", fontsize=14)

cbar[0,1].ax.set_title("ES")
#axs[0,1].set_title("Entanglement spectrum degeneracy", fontsize=14)

cbar[1,0].ax.set_title("cost")
cbar[1,1].ax.set_title("cost")

plt.savefig("plots/BH_replot2.pdf", bbox_inches='tight')

fig, ax = plt.subplots(figsize=(6,4))
im = ax.pcolormesh(x,y,cost2.T,vmin=0, cmap=cmap, shading="auto",rasterized=True)
ax.set_xscale("log")
cbar = fig.colorbar(im, ax=ax)
cbar.ax.tick_params(labelsize=cbar_labelsize)
ax.set_title("Anomaly Syndrome", fontsize=14)
cbar.ax.set_title("cost")
ax.plot(vs[10],ds[10],"x", markersize=15, color="magenta")
ax.text(vs[10],ds[10]-0.25,"Training example", color="magenta")
ax.tick_params(labelsize=14)
#ax.set_xlabel(r"$V$",fontsize=18)
#ax.set_ylabel(r"$\delta t$",fontsize=18)
plt.tight_layout()
plt.savefig("plots/BH_ex_for_presentation.png", dpi=500, bbox_inches='tight')

