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()
#
run1 = '../SW16condensP_coag0_grav_buoyancy_n1e8/data'
run2 = '../SW16condensP_coag0_grav_buoyancy_n1e8_qveLow/data'
run3 = '../SW16condensP_coag0_grav_buoyancy_n1e8_T/data'
#
ts1=pc.read_ts(datadir=run1)
ts2=pc.read_ts(datadir=run2)
ts3=pc.read_ts(datadir=run3)
#
row = 2
col = 2
siz_text = 30
siz_leg = 20
xr = [0, 30.]
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.17, right=.985, top=.94, bottom=0.13, wspace=0.40)
#
plt.subplot(row, col, 1)
#
plt.plot(ts1.t, ts1.epsK, 'k', label=r'$s_0=-0.0168$, $q_{v,e}=0.0163$')
plt.plot(ts2.t, ts2.epsK, 'r', label=r'$s_0=-0.0168$, $q_{v,e}=0.0156$')
plt.plot(ts3.t, ts3.epsK, 'b', label=r'$s_0=0.014$')
plt.xlim(xr)
plt.ylabel(r'$\bar{\epsilon}\,[\rm{m}^2\rm{s}^{-3}]$')
#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(ts1.t, ts1.ttcrms, 'k')
plt.plot(ts2.t, ts2.ttcrms, 'r')
plt.plot(ts3.t, ts3.ttcrms, 'b')
plt.xlim(xr)
plt.ylabel(r'$T_{\rm rms}\,[K]$')
#
plt.subplot(row, col, 3)
plt.plot(ts1.t, ts1.accrms, 'k')
plt.plot(ts2.t, ts2.accrms, 'r')
plt.plot(ts3.t, ts3.accrms, 'b')
plt.xlim(xr)
plt.xlabel(r'$t\,[{\rm s}]$')
plt.ylabel(r'$qv_{\rm rms}\,[kg \cdot kg^{-1}]$')
#
plt.subplot(row, col, 4)
plt.plot(ts1.t, ts1.ssatrms, 'k')
plt.plot(ts2.t, ts2.ssatrms, 'r')
plt.plot(ts3.t, ts3.ssatrms, 'b')
plt.xlim(xr)
plt.xlabel(r'$t\,[\rm s]$')
plt.ylabel(r'$s_{\rm rms}\,[\%]$')
#
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')