import numpy
import matplotlib.pyplot as pyplot
from matplotlib.gridspec import GridSpec as GS
from jqc import jqc_plot
import os
import scipy.interpolate as interpolate

cwd = os.path.dirname(os.path.abspath(__file__))

grid = GS(1,2,width_ratios=[1.5,1])
jqc_plot.plot_style('normal')
fig = pyplot.figure("Higher N ")

ax1 = fig.add_subplot(grid[0])

x1,x2 = 0.6,0.9
y1,y2 = 0.55,1.0

positions = ax1.transAxes.transform([(x1,y1),(x2,y2)])
invert = fig.transFigure.inverted()
positions = invert.transform(positions)
print(positions)

x1,y1 = positions[0]
w,h = positions[1]-positions[0]

ax2 = fig.add_axes([x1,y1,w,h])

ax3 = fig.add_subplot(grid[1])

colours = [jqc_plot.colours['red'],jqc_plot.colours['blue']]
highlights = [jqc_plot.colours['reddish'],jqc_plot.colours['grayblue']]

for i,f in enumerate(["Sig"]):
    data_pol = numpy.genfromtxt(cwd+"\\Data\\pol\\"+f+"_ladder.csv",delimiter=',')
    data_d = numpy.genfromtxt(cwd+"\\Data\\d\\"+f+"_ladder.csv",delimiter=',')
    #delta_data_pol = [data_pol[i+1,1]-data_pol[i,1]for i in range(len(data_pol)-1)]
    delta_data_pol = [2* (data_pol[n,1]-data_pol[n+1,1])/(data_pol[n,1]+data_pol[n+1,1]) for n in range(len(data_pol)-1)]


    ax1.plot(data_pol[:,0],data_pol[:,1],color=highlights[i],
                mfc=colours[i],mec=colours[i],marker='o')
    ax2.plot(data_pol[:-1,0],numpy.abs(delta_data_pol),
                color=highlights[i],mfc=colours[i],mec=colours[i],marker='o')

    delta_data_d = [data_d[i+1,1]-data_d[i,1]for i in range(len(data_pol)-1)]
    ax3.plot(data_d[:,0],numpy.abs(data_d[:,1]),color=highlights[i+1],
    mfc=colours[i+1],mec=colours[i+1],marker='o')


err_pol_0 = 10
err_pol_2 = 1

ax1.errorbar(data_pol[0:2,0],data_pol[0:2,1],yerr=[err_pol_0,numpy.sqrt(err_pol_0**2+(0.2*err_pol_2)**2)],color=jqc_plot.colours['red'],capsize = 3.5, fmt='o',zorder = 2.5)


ax1.set_ylabel("Polarisability,$\\alpha$ (a$_0^3$)")
ax1.set_xlabel("Rotational Quantum \n Number, $N$")
ax1.set_ylim(650,900)
ax2.set_xlim(0,10)

ax2.set_yscale('log')

ax2.set_xlabel("$N$")
ax2.set_ylabel("$|\\Delta \\alpha|~(\\bar{\\alpha})$")


ax3.set_xlabel("Rotational Quantum \n Number, $N$")
ax3.set_ylabel("Transition Dipole Moment ($d_0$)")

ax1.text(0.01,0.03,"(a)",fontsize=20,transform=ax1.transAxes)
ax3.text(0.77,0.03,"(b)",fontsize=20,transform=ax3.transAxes)
fig.tight_layout()

pyplot.savefig(cwd+"\\HigherN.pdf")
pyplot.savefig(cwd+"\\HigherN.png")
pyplot.show()
