import numpy as np
import xarray as xr
import matplotlib
matplotlib.use('Agg') # Must be before importing matplotlib.pyplot or pylab!
import matplotlib.pyplot as plt
import pandas as pd
import os
#import seaborn as sns
import matplotlib as mpl
import timeit

start_time = timeit.default_timer()

#%% read in the data
pr = 'OUTPUT/ches_scratch2_pr.000.nc'
pr_ds = xr.open_dataset(pr)
#print(pr_ds['time'])

##plot first theta plot

#pull in the z variables
ztheta_simulated = np.asarray(pr_ds['ztheta'])

#start plotting
plt.rc('font', family='serif')
plt.rc('xtick', labelsize='large')
plt.rc('ytick', labelsize='large')

fig = plt.figure(figsize=(8,8))

ax = fig.subplots(1,1)

ax.plot(np.asarray(pr_ds.theta.sel(time='604927000000')),ztheta_simulated,'-b')
ax.set_ylim([0., 3000.])
fig.savefig("Plot.png")
elapsed = timeit.default_timer() - start_time
print('Time elapsed ',elapsed, 'seconds')