import os
import sys
import pencil as pc
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import rc
plt.rc('text', usetex=True)
plt.ion()
#
#-----Define parameters
gg = 1.e3
#
run5 = '../SW16condensP_coag0_buoyancyMean_G1p17em10/data/'
run6 = '../SW16condensP_coag0_buoyancyMean/data/'
#
ts5=pc.read_ts(datadir=run5)
ts6=pc.read_ts(datadir=run6)
#
row = 3
col = 2
siz_text = 25
siz_tit = 30
siz_leg = 20
xr = [0, 15.]
yr_B = [2.05, 2.15]
yr_T = [0, 2]
yr_q = [0.0, 0.3]
yr_s = [0., 0.1]
yr_urms = [0., 0.15]
yr_epsk = [0., 0.05]
yr_Cd = [1.e-6, 2.e-3]
plt.close('all')
plt.figure(figsize=(12,8))
plt.rc('font', family='Times', size=siz_text)
plt.rc('xtick', labelsize=siz_text)
plt.rc('ytick', labelsize=siz_text)
plt.tight_layout(pad=0)
plt.subplots_adjust(left=0.11, right=.96, top=.98, bottom=0.11, wspace=0.3)
#
plt.subplot(row, col, 1)
#
plt.plot(ts5.t, ts5.urms, '-r', label=r'$G=1.17\times10^{-10}$')
plt.plot(ts6.t, ts6.urms, '-k', label=r'$G=1.2\times10^{-12}$')
plt.xlim(xr)
plt.ylim(yr_urms)
#plt.ylabel(r'$\bar{\epsilon}\,[\rm{m}^2\rm{s}^{-3}]$')
plt.ylabel(r'$u_{\rm rms}\,[\rm{m}\rm{s}^{-2}]$', size=siz_tit)
plt.tick_params(top='on', right='on', direction='in', which='both')
#leg = plt.legend(fontsize=siz_text, loc=2, bbox_to_anchor=(0.01,1.03), handletextpad=1., frameon=False)
leg = plt.legend(fontsize=siz_leg, loc='best', handlelength=1., frameon=False)
def color_legend_texts(leg):
    """Color legend texts based on color of corresponding lines"""
    for line, txt in zip(leg.get_lines(), leg.get_texts()):
        txt.set_color(line.get_color())  
color_legend_texts(leg)
#
plt.subplot(row, col, 2)
#
plt.plot(ts5.t, ts5.epsK, '-r')
plt.plot(ts6.t, ts6.epsK, '-k')
plt.xlim(xr)
plt.ylim(yr_epsk)
plt.ylabel(r'$\bar{\epsilon}\,[\rm{m}^2\rm{s}^{-3}]$')
plt.tick_params(top='on', right='on', direction='in', which='both')
plt.subplot(row, col, 3)
#
plt.plot(ts5.t, np.sqrt(ts5.ttcrms**2-ts5.ttcm**2), '-r')
plt.plot(ts6.t, np.sqrt(ts6.ttcrms**2-ts6.ttcm**2), '-k')
plt.xlim(xr)
plt.ylim(yr_T)
plt.ylabel(r'$T_{\rm rms}\,[K]$', size=siz_tit)
plt.tick_params(top='on', right='on', direction='in', which='both')
#
plt.subplot(row, col, 4)
plt.plot(ts5.t,np.sqrt(np.abs(ts5.accrms**2-ts5.accm**2))*gg, '-r')
plt.plot(ts6.t,np.sqrt(np.abs(ts6.accrms**2-ts6.accm**2))*gg, '-k')
plt.xlim(xr)
plt.ylim(yr_q)
#plt.xlabel(r'$t\,[{\rm s}]$', size=siz_tit)
plt.ylabel(r'$q_{v,\rm rms}\,[\rm{g \cdot kg}^{-1}]$', size=siz_tit)
plt.tick_params(top='on', right='on', direction='in', which='both')
#
plt.subplot(row, col, 5)
plt.semilogy(ts5.t,np.sqrt(np.abs(ts5.condensationRaterms**2-ts5.condensationRatem**2)), '-r')
plt.semilogy(ts6.t,np.sqrt(np.abs(ts6.condensationRaterms**2-ts6.condensationRatem**2)), '-k')
plt.xlim(xr)
plt.ylim(yr_Cd)
plt.xlabel(r'$t\,[{\rm s}]$', size=siz_tit)
plt.ylabel(r'$C_{d,\rm rms}$', size=siz_tit)
plt.tick_params(top='on', right='on', direction='in', which='both')
#
plt.subplot(row, col, 6)
#plt.plot(ts1.t, np.sqrt(ts1.ssatrms**2-ts1.ssatm**2), '-r')
#plt.plot(ts2.t, np.sqrt(ts2.ssatrms**2-ts2.ssatm**2), '-k')
#plt.plot(ts3.t, np.sqrt(ts3.ssatrms**2-ts3.ssatm**2), '-b')
#plt.plot(ts4.t, np.sqrt(ts4.ssatrms**2-ts4.ssatm**2), '-g')
plt.plot(ts5.t, np.sqrt(ts5.ssatrms**2-ts5.ssatm**2), '-r')
plt.plot(ts6.t, np.sqrt(ts6.ssatrms**2-ts6.ssatm**2), '-k')
plt.xlim(xr)
plt.ylim(yr_s)
plt.xlabel(r'$t\,[\rm s]$', size=siz_tit)
plt.ylabel(r'$s_{\rm rms}$', size=siz_tit)
plt.tick_params(top='on', right='on', direction='in', which='both')
#-----mean ssat------------
plt.twinx()
#plt.plot(ts1.t, ts1.ssatm, ':r')
#plt.plot(ts2.t, ts2.ssatm, ':k')
#plt.plot(ts3.t, ts3.ssatm, ':b')
#plt.plot(ts4.t, ts4.ssatm, ':g')
plt.plot(ts5.t, ts5.ssatm, ':r')
plt.plot(ts6.t, ts6.ssatm, ':k')
plt.xlim(xr)
plt.ylim(yr_s)
plt.ylabel(r'$\bar{s}$', size=siz_tit)
#plt.yticks([])
plt.tick_params(labelright='off', direction='in')
#
filename=sys.argv[0]
print(filename)
basename=os.path.basename(filename)
base=basename.split('.')[0]
plt.savefig('python.eps',dpi=300)
print('!mv'+' '+'python.eps'+' '+'~/tex/xiangyu/supersat/fig/'+base+'.eps')