from scipy import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib as mpl
import sys

mpl.rc('font',**{'family':'Times New Roman','sans-serif':['Helvetica']})
mpl.rc('ytick', labelsize=12)
mpl.rc('xtick', labelsize=12)
mpl.rc('xtick.major', pad=8)
mpl.rc('ytick.major', pad=8)
mpl.rc('text', usetex=True)

def color_y_axis(ax, color):
    """Color your axes."""
    for t in ax.get_yticklabels():
        t.set_color(color)

    plt.setp(ax.spines.values(), color=color)
    plt.setp(ax.get_yticklines(), color=color)
    return None

fig = plt.figure(figsize=(8.27,11.69))
gs = gridspec.GridSpec(5, 3)
gs.update(wspace=0.1, hspace=0.3)

cell = ['-2', '0', '+2', '+4', '+6']
modes = ['TA1', 'TA2', 'LA']
Temps = ['10', '50', '100', '200', '300']

muev_to_psinv = 1.5192674606 / 1000
mev = 8.06554

for i in range(len(cell)):
    entry = cell[i]
    for j in range(len(modes)):
        mode = modes[j]
        ax = plt.subplot(gs[i,j])
        ax2 = ax.twinx()
        ax3 = ax.twinx()

        if i != len(cell) - 1:
            ax.set_xticklabels([])

        if j != 0:
            ax.set_yticklabels([])

        if j != 2:
            ax2.set_yticklabels([])

        #ax3.set_xticklabels([])
        ax3.set_yticklabels([])
        #ax3.set_xticks([])
        ax3.set_yticks([])
        try:
            target_file = ("cell%s/linewidth/TA2_Gam_ISC_300K.dat" % entry)
            data = np.loadtxt(target_file, usecols=[1,])
            ymax = np.max(data)
        except:
            data = np.zeros((len(data),2))
            pass
        ax.spines["left"].set_color('blue')
        ax2.spines["left"].set_visible(False)
        ax3.spines["left"].set_visible(False)
        ax3.spines["right"].set_visible(False)

        for T in Temps:
            #target_file = ("cell%s/linewidth/%s_Gam_harm_%sK.dat" % (entry, mode, T))
            target_file = ("cell%s/linewidth/%s_Gam_ISC_%sK.dat" % (entry, mode, T))
            try:
                data = np.loadtxt(target_file)
            except:
                data = np.zeros((len(data),2))
                pass
            ax.plot(data[:,0]/mev, data[:,1]/mev*1.0e+3, linestyle='--', marker='o', ms=4, label=("%sK" % T))
            if i == 1:
                ax.set_title(r"Cell: ``opt.'', Mode: %s" % mode)
            else:
                tmp = r"``$" + entry + "\%" + "$''"
                ax.set_title(r"Cell: %s, Mode: %s" % (tmp, mode))
            #ax.set_yscale('log')

            if i == 0 and j == 0:
                ax.legend(loc='best', fontsize=14, bbox_to_anchor=(3.25, 1.55), ncol=6)

        ax2.axis([0, 5.5, 0, ymax*1.5/mev*1.0e+3*muev_to_psinv])
        ax.axis([0, 5.5, 0, ymax*1.5/mev*1.0e+3])
        color_y_axis(ax, 'red')
        color_y_axis(ax2, 'blue')

fig.text(0.5, 0.04, r'Frequency (meV)', ha='center', size=20)
fig.text(0.03, 0.5, r'2$\Gamma^{(\mathrm{B})}$ ($\mu$eV)', va='center', rotation='vertical', size=20, color='red')
fig.text(0.97, 0.5, r'$1/\tau_{\mathrm{anh}}$ (ps$^{-1}$)', va='center', rotation='vertical', size=20, color='blue')

#plt.legend(loc='best')
plt.savefig("tau_acoustic_SCP.eps", format='eps', bbox_inches='tight')

plt.show()
