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()

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

data = 'VM_OUTPUT_N03/site0010_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)


ds1 = xr.open_dataset('check1.nc')

ds2 = xr.open_dataset('check2.nc')

ds_concat = xr.concat([ds1,ds2],pd.Index(['nw1','nw4'], name='tower'))
print(ds_concat)


'''     
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()

print(data_ds)

#save the subset
data_ds.to_netcdf('/bog/incoming/CHEESEHEAD/palm/realistic_runs/ches_IOP03/OUTPUT/check2.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!')