import numpy as np
import xarray as xr
import matplotlib as mpl
mpl.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 timeit
start_time = timeit.default_timer()

#%% plot parameters
'''
mpl.rcParams['axes.linewidth'] = 2.0 #set the value globally

plt.rc('font', family='serif',size = 16)  # controls default text layout
plt.rc('axes', titlesize=16)     # fontsize of the axes title
plt.rc('axes', labelsize=16)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=10)    # fontsize of the tick labels
plt.rc('ytick', labelsize=16)    # fontsize of the tick labels
plt.rc('legend', fontsize=16)    # legend fontsize
plt.rc('figure', titlesize=16)  # fontsize of the figure title

#plt.minorticks_on()
plt.rc('text', usetex=False)
sns.set(font = "Cambria",font_scale=1.5,style="whitegrid")
#sns.set(font = "Cambria",font_scale=1.5,style="white")

#mpl.rc_file_defaults() #undos seaborn effects

#custom fontsize
font = 12
'''
#%% data folder
source = '/bog/incoming/CHEESEHEAD/palm/realistic_runs/ches_IOP03/OUTPUT/ensemble.member.1'

data = 'DATA_2D_XY_AV_NETCDF_N03'
#variable of interest
var = 'tsurf*_xy'
#height_level = 82





os.chdir(source)
#get list of all the folders
folder_list = os.listdir(source)

#list to combine the data
ds_list = []
for i, folder in enumerate(folder_list):
     #open data file inside the folder.
     ds = xr.open_dataset(source+ '/' + folder + '/' + data)
     ds.close()
     #subset for height level, if needed
     ds = ds[var]#.sel(zu_3d = height_level)
     #drop the NaT times
     ds = ds.where(ds.time.notnull(),drop=True)
     #append the data file name to the list

     ds_list.append(ds)

#concatenate all data files in the list
data_ds = xr.concat(ds_list, dim='time')

#print(data_ds)
'''
#save the subset
#ds2.to_netcdf('/bog/incoming/CHEESEHEAD/palm/realistic_runs/ches_IOP03/OUTPUT/check.nc')

#make plots and save them if needed.
for i in range(0,len(data_ds.time),4): 
    fig = plt.figure(figsize=(8,8))
    data_ds.sel(time=data_ds.time[i]).plot( cmap = 'RdBu_r')
    #fname = '/bog/incoming/CHEESEHEAD/palm/realistic_runs/ches_IOP03/OUTPUT/plots' + '/' + var + str(height_level) + str(i) + '.png'
    fname = '/bog/incoming/CHEESEHEAD/palm/realistic_runs/ches_IOP03/OUTPUT/plots' + '/' + var + str(i) + '.png'
    fig.savefig(fname)
    plt.close(fig)
'''

#write to netcdf
data_ds.to_netcdf('surface_temperature.nc')

elapsed = timeit.default_timer() - start_time
print('Time elapsed ',elapsed, 'seconds')
print('done!')