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'


site_list = ['0005','0010','0011','0012','0013','0014','0015','0017','0018','0019','0020','0025']

site_name = ['nw1', 'nw4', 'ne1', 'ne2', 'ne3', 'ne4', 'sw1', 'sw3','sw4', 'se2', 'se3', 'se6']

#list to club all vm data

vm_list = []


for i, site in enumerate(site_list):
    
    data = 'VM_OUTPUT_N03/site' + site + '_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 for one site
    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 from one vm in the list
    data_ds = xr.concat(ds_list, dim='time')
    #print(data_ds)
    data_ds = data_ds.resample(time="30Min").mean()
    
    #append the vm data to the master list
    vm_list.append(data_ds)
    
#outside of the loop
#concatenate all the vm data
vm_ds = xr.concat(vm_list, pd.Index(site_name,name='tower'))

    
print(vm_ds)

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


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