# -*- coding: utf-8 -*-
from __future__ import division
import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib
import os,sys
from matplotlib import rc
plt.rc('text', usetex=True)
plt.rc('font', family='serif',serif='Times',size=54)
I=0
paths=['fig6.eps','fig7.eps']
shading='flat'
cmap=plt.get_cmap('hot')
fontsizela=64
fsize=(20,14)
for folder in ['gvout_hr','gvouth_hr']:
    fig=plt.figure(figsize=fsize)
    if folder=='gvout_hr':
        harmonic_trap=False
    if folder=='gvouth_hr':
        harmonic_trap=True
    if not harmonic_trap:
        vs=np.loadtxt('%s/vs.txt'%(folder))
    else:
        wxs=np.loadtxt('%s/wxs.txt'%(folder))
        x0=np.loadtxt('%s/x0.txt'%(folder))
        vs=wxs*x0
    nv=len(vs)
    if os.path.exists('%s/params.txt'%(folder)):
        params=np.loadtxt('%s/params.txt'%(folder))
    gs=np.loadtxt('%s/gs.txt'%(folder))
    nj=len(gs)
    if len(params)==2:
        e,f=params
    if len(params)==3:
        e,f,s=params
    Ts1=np.loadtxt('%s/T1.txt'%(folder))
    Ts2=np.loadtxt('%s/T2.txt'%(folder))
    ax1=fig.add_subplot(311)
    ax2=fig.add_subplot(312)
    ax3=fig.add_subplot(313)
    im1=ax1.pcolormesh(vs,gs,np.transpose(Ts1),cmap=cmap,shading=shading,vmin=0,vmax=1)
    im2=ax2.pcolormesh(vs,gs,np.transpose(Ts2),cmap=cmap,shading=shading,vmin=0,vmax=1)
    im3=ax3.pcolormesh(vs,gs,np.transpose(Ts1-Ts2),cmap=cmap,shading=shading,vmin=0,vmax=1)
    fig.subplots_adjust(right=0.8)
    cbar_ax = fig.add_axes([0.83, 0.11, 0.03, 0.77])
    cbr1=fig.colorbar(im1,ticks=[0,0.2,0.4,0.6,0.8,1.0],cax=cbar_ax)
    if harmonic_trap:
        ax3.set_xlabel('$\\omega_x x_0$', labelpad=10, fontsize=fontsizela)
    else:
        ax3.set_xlabel('$v$', labelpad=10, fontsize=fontsizela)
    for i in range(1,4):
        exec('ax%d.set_ylabel("$g$",rotation=0,fontsize=fontsizela)'%(i))
        exec('ax%d.yaxis.set_label_coords(-0.12,0.35)'%(i))
    for i in range(1,4):
        exec('ax%d.set_xlim(np.min(vs),np.max(vs))'%(i))
        exec('ax%d.set_ylim(np.min(gs),np.max(gs))'%(i))
    fig.text(0.14,0.82,'(a)  $T_1$',fontsize=54,color='white')
    fig.text(0.14,0.55,'(b)  $T_2$',fontsize=54,color='white')
    fig.text(0.14,0.28,'(c)  $T_1-T_2$',fontsize=54,color='white')
    ax1.set_xticklabels([])
    ax2.set_xticklabels([])
    for i in range(1,4):
        exec('ax%d.plot(vs,np.zeros(len(vs)),color="gray",linewidth=1)'%(i))
    plt.savefig('%s'%paths[I],format='eps')
    I=I+1
