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 = 'DATA_3D_AV_NETCDF_N03'


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



x_list = [2217,3981,6159,8619,9225,9669,723,3249,5391,6801,8091,9537]
y_list = [9147,9951,9309,7335,9465,8019,2799,3441,5499,5199,4161,3333]
site_name = ['nw1', 'nw4', 'ne1', 'ne2', 'ne3', 'ne4', 'sw1', 'sw3','sw4', 'se2', 'se3', 'se6']

#list to club all vm data

site_wtheta_list = []
site_wq_list = []

#looping through each site first
for i, site in enumerate(site_name):
    
    
         
    os.chdir(source)
    #get list of all the folders
    folder_list = os.listdir(source)
    
    
    
    #list to combine the data for one site
    wtheta_list = []
    wq_list = []
    #first collect data at all times for one site
    for j, folder in enumerate(folder_list):
        #for time> 5 hours, i > 1 
        if j > 1 :            
              #open data file inside the folder.
              ds = xr.open_dataset(source + '/' + folder + '/' + data)
              ds.close()
              #subset for variables
              ds = ds[['w','theta','wtheta','q','wq']]#.sel(zw_3d = height_level)
              ds = ds.sel(x=x_list[i],y=y_list[i])
              
              #subset for non zero time
              ds = ds.where(ds.time.notnull(),drop=True)
              #read in variables to data arrays, renaming the coords for theta and q to interpolate
              w = ds.w
              wtheta = ds.wtheta
              wq = ds.wq
              
              theta = ds.theta.rename({'zu_3d': 'zw_3d'})
              theta_interpolated = theta.interp_like(w)
              q = ds.q.rename({'zu_3d': 'zw_3d'})
              q_interpolated = q.interp_like(w)
              
              #calculate the fluxes
              
              #wtheta_turb = wtheta.sel(zw_3d=32) - w.sel(zw_3d=32)*theta_interpolated.sel(zw_3d=32)
              #wq_turb = wq.sel(zw_3d=32) - w.sel(zw_3d=32)*q_interpolated.sel(zw_3d=32)
              wtheta_turb = wtheta - w*theta_interpolated
              wq_turb = wq - w*q_interpolated
              wtheta_list.append(wtheta_turb)
              wq_list.append(wq_turb)
              
              
              print('Done!' + folder)
    
    #for each tower location combine the flux data        
    wtheta_ds = xr.concat(wtheta_list, dim='time')
    
    wq_ds = xr.concat(wq_list, dim='time')

    site_wtheta_list.append(wtheta_ds)
    site_wq_list.append(wq_ds)   
            
    print('done!' + site)
    
# #outside of the loop
# #concatenate all the flux data data

wtheta_all_sites_ds = xr.concat(site_wtheta_list, pd.Index(site_name,name='tower'))
wq_all_sites_ds = xr.concat(site_wq_list, pd.Index(site_name,name='tower'))

# #save the subset
wtheta_all_sites_ds.to_netcdf('/bog/incoming/CHEESEHEAD/palm/realistic_runs/ches_IOP03/OUTPUT/wtheta_all_sites_profiles.nc')
wq_all_sites_ds.to_netcdf('/bog/incoming/CHEESEHEAD/palm/realistic_runs/ches_IOP03/OUTPUT/wq_all_sites_profiles.nc')


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

print('DONE!')