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 = 'VM_OUTPUT_N03/site0005_N03'
#variable of interest
var = 'ta'
#height_level = 82

#folder = '01.ches_IOP3.14592'
#ds = xr.open_dataset(source+ '/' + folder + '/' + data)
#ds.close()


#print(ds[var].time)
#print(ds)
     
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):
     if i > 0 :            
          #open data file inside the folder.
          ds = xr.open_dataset(source+ '/' + folder + '/' + data)
          ds.close()
          #subset for height level, if needed
          ds = ds.sel(station=14)
          #swap dimensions from ntime to the actual timestep
          ds = ds.swap_dims({'ntime' : 'time'})
          #drop the NaT times
          #ds = ds.where(ds.time.notnull(),drop=True)
          #drop 0 values for vm data
          ds = ds.where(ds[var]!=0,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)
data_ds = data_ds.resample(time="30Min").mean()


fig = plt.figure(figsize=(8,8))
data_ds.ta.plot()
fig.savefig('plot.png')

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


#make plots and save them.

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 = 'Blues')
    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)

'''

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