import matplotlib.pyplot as pyplot
import numpy
from matplotlib.colors import LinearSegmentedColormap,ListedColormap
from matplotlib import rcParams, lines
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap, BoundaryNorm,LogNorm
from matplotlib.ticker import MultipleLocator, ScalarFormatter,FuncFormatter,FormatStrFormatter
from scipy.optimize import fsolve as fsolve
from matplotlib import gridspec
import scipy.interpolate as interp
from matplotlib.patches import ConnectionPatch,FancyBboxPatch,Rectangle
import os
directory = os.getcwd()

pyplot.style.use(directory+r'\Style Files\Wide_Origin_style.mplstyle')

File_path = directory+"\\Jesus_data\\"

colour_dict_twk_blue = {
    "red" : [(0.0,244.0/255.0,244.0/255.0),
            (0.33,124.0/255.0,124.0/255.0),
            (0.66,0.0,0.0),
            (1.0,0.0,0.0)] ,
    "green" : [(0.0,234.0/255.0,234.0/255.0),
            (0.33,154.0/255.0,154.0/255.0),
            (0.66,70.0/255.0,70.0/255.0),
            (1.0,32.0/255.0,32.0/255.0)]  ,
    "blue" : [(0.0,168.0/255.0,168.0/255.0),
            (0.33,148.0/255.0,148/255.0),
            (0.66,127.0/255.0,127.0/255.0),
            (1.0,58.0/255.0,58.0/255.0)]
}

colour_dict_twk_red = {
    "red" : [(0.0,244.0/255.0,244.0/255.0),
            (0.33,229.0/255.0,229.0/255.0),
            (0.66,214.0/255.0,214.0/255.0),
            (1.0,170.0/255.0,170.0/255.0)] ,
    "green" : [(0.0,234.0/255.0,234.0/255.0),
            (0.33,177.0/255.0,177.0/255.0),
            (0.66,120.0/255.0,120.0/255.0),
            (1.0,43.0/255.0,43.0/255.0)]  ,
    "blue" : [(0.0,168.0/255.0,168.0/255.0),
            (0.33,145.0/255.0,145/255.0),
            (0.66,122.0/255.0,122.0/255.0),
            (1.0,74.0/255.0,74.0/255.0)]

}

JQC = {'red'     :(198.0/255.0, 62.0/255.0, 98.0/255.0), \
       'blue'    :(0.0/255.0, 70.0/255.0, 127.0/255.0), \
       'purple'  :(126.0/255.0, 29.0/255.0, 123.0/255.0), \
       'sand'  :(244./255.0, 234./255.0, 168./255.0), \
       'grayblue'  :(212./255.0, 213./255.0, 220./255.0), \
       'green'   :(45.0/255.0, 159.0/255.0, 60.0/255.0)}

colour_dict_twk_blue_alpha = colour_dict_twk_blue.copy()
colour_dict_twk_blue_alpha['alpha'] = ((0.0, 0.0,0.0),
                #   (0.25,1.0, 1.0),
                   (0.5, 1.0, 1.0),
                #   (0.75,1.0, 1.0),
                   (1.0, 1.0, 1.0))


RbCs_map_twk_blue = LinearSegmentedColormap("RbCs_map_tweak_blue",colour_dict_twk_blue_alpha)
pyplot.register_cmap(cmap=RbCs_map_twk_blue)

colour_dict_twk_red_alpha = colour_dict_twk_red.copy()
colour_dict_twk_red_alpha['alpha'] = ((0.0, 0.0,0.0),
                #   (0.25,1.0, 1.0),
                   (0.5, 1.0, 1.0),
                #   (0.75,1.0, 1.0),
                   (1.0, 1.0, 1.0))


RbCs_map_twk_red = LinearSegmentedColormap("RbCs_map_tweak_red",colour_dict_twk_red_alpha)
pyplot.register_cmap(cmap=RbCs_map_twk_red)


def make_segments(x, y):
    '''
    Create list of line segments from x and y coordinates, in the correct format for LineCollection:
    an array of the form   numlines x (points per line) x 2 (x and y) array
    '''

    points = numpy.array([x, y]).T.reshape(-1, 1, 2)
    segments = numpy.concatenate([points[:-1], points[1:]], axis=1)

    return segments

def colorline(x, y, z=None, cmap=pyplot.get_cmap('copper'), norm=pyplot.Normalize(0.0, 1.0), linewidth=3, alpha=1.0,legend=False):
    '''
    Plot a colored line with coordinates x and y
    Optionally specify colors in the array z
    Optionally specify a colormap, a norm function and a line width
    '''

    # Default colors equally spaced on [0,1]:
    if z is None:
        z = numpy.linspace(0.0, 1.0, len(x))

    # Special case if a single number:
    if not hasattr(z, "__iter__"):  # to check for numerical input -- this is a hack
        z = numpy.array([z])

    z = numpy.asarray(z)

    segments = make_segments(x, y)
    lc = LineCollection(segments, array=z, cmap=cmap, norm=norm, linewidth=linewidth)

    ax = pyplot.gca()
    ax.add_collection(lc)

    return lc


def clear_frame(ax=None):
    # Taken from a post by Tony S Yu
    if ax is None:
        ax = pyplot.gca()
    ax.xaxis.set_visible(False)
    ax.yaxis.set_visible(False)
    for spine in ax.spines.itervalues():
        spine.set_visible(False)

vmin=1e-3

#scaling = 0.82

fig = pyplot.figure()

IntensityScaling = 1.00

grid=gridspec.GridSpec(2,2,width_ratios=[1,1.3])

ax_Low = fig.add_subplot(grid[1,0])
ax_High = fig.add_subplot(grid[0,0],sharex=ax_Low)

for mF in [1,2,3,4,5,6,7]:
	Dat_m = numpy.genfromtxt(File_path+"MFm0"+str(mF)+"\\results.dat")
	Dat_p = numpy.genfromtxt(File_path+"MFp0"+str(mF)+"\\results.dat")
	pyplot.figure(1)

	for i in range(2,len(Dat_m[0,:])):
		if mF== 5:
			if i ==2:
				ax_Low.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['red'],zorder=2)
			elif i ==3:
				ax_High.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['purple'],zorder=2)
			else:
				ax_Low.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)
				ax_High.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)
		elif mF == 6:
			if i ==2:
				ax_High.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['blue'],zorder=2)
			else:
				ax_Low.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)
				ax_High.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)

		elif mF == 4:
			if i ==5:
				ax_High.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['green'],zorder=2)
				ax_Low.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)
			else:
				ax_Low.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)
				ax_High.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)
		else:
			ax_Low.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)
			ax_High.plot(Dat_p[:,1]*1e4,Dat_p[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)


		ax_Low.plot(Dat_m[:,1]*1e4,Dat_m[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)
		ax_High.plot(Dat_m[:,1]*1e4,Dat_m[:,i]*1e-3,color=JQC['grayblue'],zorder=0.5)

#ax.set_rasterization_zorder(0.9)
d = .03  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax_Low.transAxes, color='k', clip_on=False)
ax_Low.plot((1-d, 1+d), (1-d, 1+d), **kwargs)        # top-right diagonal
ax_Low.plot((- d, + d), (1-d, 1+d), **kwargs)  # bottom-right diagonal

ax_Low.spines['top'].set_visible(False)
ax_High.tick_params(axis='x',which='both',labelbottom='off',bottom='off')

ax_High.spines['bottom'].set_visible(False)
kwargs.update(transform=ax_High.transAxes)  # switch to the bottom axes
ax_High.plot((-d, +d), (- d, + d), **kwargs)  # bottom-left diagonal
ax_High.plot((1-d, 1+d), (- d,+ d), **kwargs)  # bottom-right diagonal

ax_Low.set_yticks([-1,-0.5,0,0.5,1])
ax_High.set_yticks([979,979.5,980,980.5,981])

ax_Low.set_ylim(-0.75,0.75)
ax_Low.set_xlim(0,250)

#ax_Low.add_patch(Rectangle((0,-0.2),100,0.5,fill=False,ec ='k', lw= 2.0, transform=ax_Low.transData,zorder=10))
ax_High.set_ylim(979.25,980.75)

ax_Low.set_xlabel("Magnetic Field (G)")
ax_Low.text(-0.3,1.3,"Energy$/h$ (MHz)",rotation ='vertical',transform = ax_Low.transAxes,clip_on=False)
#fig.subplots_adjust(left = 0.15, bottom = 0.14, right=0.89,top=0.97,hspace=0.05,wspace=0.4)


ax_Low.text(20,-0.23,"$+5$",color=JQC['red'],clip_on=False,transform=ax_Low.transData,fontsize=14)
ax_Low.text(120,-0.1,"$N=0$",color='k',clip_on=False,transform=ax_Low.transData,fontsize=14)
ax_High.text(20,979.95,"$+5$",color=JQC['purple'],clip_on=False,transform=ax_High.transData,fontsize=14)
ax_High.text(20,980.195,"$+4$",color=JQC['green'],clip_on=False,transform=ax_High.transData,fontsize=14)
ax_High.text(20,980.4,"$+6$",color=JQC['blue'],clip_on=False,transform=ax_High.transData,fontsize=14)
ax_High.text(120,980.15,"$N=1$",color='k',clip_on=False,transform=ax_High.transData,fontsize=14)

#ax_Low.axvline(181.5,linestyle='--',color='k')
#ax_High.axvline(181.5,linestyle='--',color='k')
ax_High.add_artist(ConnectionPatch(xyB=(181.5,-0.75),xyA=(181.5,980.75),coordsB='data',coordsA='data',axesA=ax_High,axesB=ax_Low,color='k',ls='--',lw=2,clip_on=False,zorder=15))


########################### AC-Stark Shift Figure ########################
sub = fig.add_subplot(grid[:,1])


Data = numpy.genfromtxt(r"ACStark\lines.dat",delimiter=" ")
DT_int = Data[:,0].copy()*IntensityScaling
Lines = Data[:,1:].copy()

for i in range(0,len(Lines[0,:])):
	#1e3 is for kHz, in which the frequencies are recorded.
	pyplot.plot(DT_int*1e-3,Lines[:,i]*1e-3,color=(244.0/255.0, 234.0/255.0, 168.0/255.0), zorder=0)





Data = numpy.genfromtxt(r"ACStark\sigma+.dat",delimiter=" ")
Lines_int1 = Data[:,1:].copy()

Data = numpy.genfromtxt(r"ACStark\sigma-.dat",delimiter=" ")
Lines_int2 = Data[:,1:].copy()
Lines_int = Lines_int1+Lines_int2



min_int = 0.0 # intensity is >0
max_int = 0.1 #set the normalisation to be maxed at the maximum of int (for colour)

for i in range(0,len(Lines[0,:])):
	#1e3 is for kHz, in which the frequencies are recorded.
	cl=colorline(DT_int*1e-3,Lines[:,i]*1e-3,Lines_int[:,i],cmap='RbCs_map_tweak_red',norm=LogNorm(vmin,vmax=1),linewidth=2.0)
cbar1 = pyplot.colorbar(cl, pad=-.08) #only include a colourbar on the first or last line else you get overlapping
cbar1.set_label('Relative Transition Strength')
cbar1.set_ticks([1e-3, 1e-2, 1e-1, 1e-0])

Data = numpy.genfromtxt(r"ACStark\pi.dat",delimiter=" ")
Lines_int = Data[:,1:].copy()

min_int = 0.0 # intensity is >0
max_int = 0.1 #set the normalisation to be maxed at the maximum of int (for colour)

for i in range(0,len(Lines[0,:])):
	#1e3 is for kHz, in which the frequencies are recorded.
	cl=colorline(DT_int*1e-3,Lines[:,i]*1e-3,Lines_int[:,i],cmap='RbCs_map_tweak_blue',norm=LogNorm(vmin,vmax=1),linewidth=2.0)

cbar2 = pyplot.colorbar(cl) #only include a colourbar on the first or last line else you get overlapping
cbar2.set_ticks([])




#sub.axvline(10.4,linestyle='--',color='k')
#sub.axvline(20.7,linestyle='--',color='k')
#sub.axvline(30,linestyle='--',color='k')
sub.set_xlabel("Intensity (kW/cm$^{2}$)")
sub.set_xticks([0, 10, 20, 30, 40])

sub.set_xlim(0,40)
sub.set_ylim(979.8,980.6)
sub.set_ylabel("Transition Frequency (MHz)")
Offset = 980
sub.ticklabel_format(axis='y',useOffset=Offset,style='plain')
sub.yaxis.offsetText.set_visible(False)
sub.annotate("+"+str("%.0f"%Offset)+"MHz",xy=(0.01,1.01),xycoords="axes fraction")
sub.text(10,980.07,"$+5$",color=JQC['purple'],clip_on=False,transform=sub.transData,fontsize=14)
sub.text(10,980.29,"$+4$",color=JQC['green'],clip_on=False,transform=sub.transData,fontsize=14)
sub.text(10,980.4,"$+6$",color=JQC['blue'],clip_on=False,transform=sub.transData,fontsize=14)


sub.annotate("$\Delta M_F =$",xy=(0.82,1.015),xycoords="axes fraction", color=(0.0/255.0, 0.0/255.0, 0.0/255.0))
sub.annotate("$0$",xy=(1.06,1.015),xycoords="axes fraction", color=(0.0/255.0, 70.0/255.0, 127.0/255.0))
sub.annotate("$\pm1$",xy=(1.12,1.015),xycoords="axes fraction", color=(198.0/255.0, 62.0/255.0, 98.0/255.0))


ax_Low.text(0.017,0.065,"(a)",transform=ax_Low.transAxes,fontsize=20)
sub.text(0.02,0.035,"(b)",transform=sub.transAxes,fontsize=20)
pyplot.subplots_adjust(wspace = 0.32 , hspace =0.05,right =0.94,top=0.95,left=0.11,bottom=0.16)
pyplot.savefig("Stark+HF.eps")
pyplot.show()
