import numpy as np
import xarray as xr
import timeit
import os 
#%%
#author: sreenath paleri

start_time = timeit.default_timer()

source = 'ensemble.member.0'

#get folder from user input later or choose to run inside folder
folder = '18.ches_IOP3.8982'


#create output directory for sliced data
outpath = source+ '/' + folder + '/OUTPUT' 
if not os.path.exists(outpath):
    os.makedirs(outpath)
    

#list all files in folder
files = os.listdir(source+ '/' + folder)
#list all topo files
rm = [k for k in files if 'DATA_TOPO' in k]

#delete all topo files
for i, file_name in enumerate(rm):
    print('Removing:' + source+ '/' + folder + '/' + file_name)
    os.remove(source+ '/' + folder + '/' + file_name)


#read in data from the folder    
for i, file_name in enumerate(os.listdir(source+ '/' + folder)):
    #subset for output files starting with 'DATA'
    if file_name.startswith('DATA'):
        #open and read those files
        print('Reading:' + source+ '/' + folder + '/' + file_name)
        ds = xr.open_dataset(source+ '/' + folder + '/' + file_name)
        ds.close()
        time_array = ds.time

        #subset for the non null time
        time_array = time_array.where(time_array.time.notnull(),drop=True)

        #subset dataset for the first and last available time steps 
        sample = ds.loc[dict(time=slice(time_array[0],time_array[-1]))]
        
        #name of the modiefied outfile, sliced for non-null times
        outfile = source+ '/' + folder + '/OUTPUT/' + file_name + 'slice'
        print('Writing NetCDF file',flush=True)
        print(outfile)
        
        #create an empty outfile
        nc_output = xr.Dataset()
        #copy all the input attributes
        nc_output.attrs = sample.attrs
        
        #write out the empty outfile with just the global attributes
        nc_output.to_netcdf(outfile)
        nc_output.close()
        
        #write out all variable individually
        for var in sample.data_vars:
            sample[var].to_netcdf(outfile, mode='a')
        
        #delete variables from memory    
        del sample, ds


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