import numpy as np
from netCDF4 import Dataset
import os
import sys

# path to input 4D file    
input_file = "/afast/pelletie/nemo_outputs/so025_noisf/pped/so025_noisf_1m_tempsal_1999.nc"
# path to output file
out_file = "/afast/pelletie/nemo_outputs/so025_noisf/pped/so025_noisf_1m_tempsal_layers_1999.nc"

# Upper and lower bounds of average ocean temperature layers.
upper_bounds = [0.,   50., 150., 300.]
lower_bounds = [50., 150., 300., 500.]

# path to 'mesh mask' file
mesh_file = "/afast/pelletie/nemo_outputs/so025_noisf/mesh/mesh_mask.nc"


# end input

tmpdat = Dataset(mesh_file, mode='r')
e3t = tmpdat.variables['e3t_0'][:]
gdepth_3d = tmpdat.variables['gdept_0'][:]
tmask = (tmpdat.variables['tmask'][:]).astype(bool)
tmpdat.close()


nz, ny, nx = np.shape(tmask)

tmask = np.where((np.logical_not(tmask[0,:,:]))[None,:,:] * (np.ones(shape=[nz], dtype=bool))[:,None,None], \
                 False, tmask)


n_layers = np.size(upper_bounds)
thick_intg = np.nan * np.ndarray(shape=[n_layers,nz, ny, nx], dtype=float)
# thick_intg = np.zeros(shape=[n_layers,nz, ny, nx], dtype=float)


upper_mesh = gdepth_3d - .5 * e3t
lower_mesh = gdepth_3d + .5 * e3t

for m in range(0,n_layers):

    fillidx = np.nonzero(np.logical_and( upper_bounds[m] <= upper_mesh, lower_bounds[m] >= lower_mesh))
    (thick_intg[m,:,:,:])[fillidx] = e3t[fillidx]
    
    fillidx = np.nonzero(np.logical_or( upper_bounds[m] > lower_mesh, lower_bounds[m] < upper_mesh))
    (thick_intg[m,:,:,:])[fillidx] = 0.

    fillidx = np.nonzero(np.logical_and( upper_bounds[m] > upper_mesh, upper_bounds[m] <= lower_mesh))
    (thick_intg[m,:,:,:])[fillidx] = (lower_mesh - upper_bounds[m])[fillidx]

    fillidx = np.nonzero(np.logical_and( lower_bounds[m] < lower_mesh, lower_bounds[m] >= upper_mesh))
    (thick_intg[m,:,:,:])[fillidx] = (lower_bounds[m] - upper_mesh)[fillidx]

    fillidx = np.nonzero(np.logical_not(tmask))
    (thick_intg[m,:,:,:])[fillidx] = 0.

    # print('nan count', np.count_nonzero(np.isnan(thick_intg[m,:,:,:])), nz*ny*nx)


#     print(m,thick_intg[m,:,300,500], upper_mesh[:,300,500], lower_mesh[:,300,500], \
#           tmask[:,300,500])

# sys.exit()

    

    
del upper_mesh,lower_mesh,e3t, gdepth_3d


fillval = -1.e20



tmpdat = Dataset(input_file, mode='r')

in_temp = tmpdat.variables['thetao'][:]
in_sal  = tmpdat.variables['so'][:]

lat = tmpdat.variables['nav_lat'][:]
lon = tmpdat.variables['nav_lon'][:]
depth_bounds = tmpdat.variables['deptht_bounds'][:]
depth = tmpdat.variables['deptht'][:]

out_nc = Dataset(out_file, mode='w')

out_nc.createDimension('y', ny)
out_nc.createDimension('x', nx)
out_nc.createDimension('depth', n_layers)
out_nc.createDimension('time_counter', 0)
out_nc.createDimension('axis_nbounds', 2)

tmp = out_nc.createVariable(varname='time_counter', datatype='f8', dimensions=['time_counter'])
tmp[:] = tmpdat.variables['time_counter'][:]
nt = np.size(tmp[:])
tmp.setncatts( tmpdat.variables['time_counter'].__dict__)

tmp = out_nc.createVariable(varname='time_counter_bounds', datatype='f8', dimensions=['time_counter', 'axis_nbounds'])
tmp[:] = tmpdat.variables['time_counter_bounds'][:]
tmp.setncatts( tmpdat.variables['time_counter_bounds'].__dict__)

tmp = out_nc.createVariable(varname='nav_lat', datatype='f8', dimensions=['y', 'x'])
tmp[:] = tmpdat.variables['nav_lat'][:]
tmp.setncatts( tmpdat.variables['nav_lat'].__dict__)

tmp = out_nc.createVariable(varname='nav_lon', datatype='f8', dimensions=['y', 'x'])
tmp[:] = tmpdat.variables['nav_lon'][:]
tmp.setncatts( tmpdat.variables['nav_lon'].__dict__)

tmp = out_nc.createVariable(varname='depth', datatype='f8', dimensions=['depth'])
tmp[:] = np.mean(np.vstack((upper_bounds, lower_bounds)), axis=0)
tmp.standard_name = "depth"
tmp.units="m"
tmp.axis="Z"
tmp.positive="down"
tmp.bounds="depth_bounds"

tmp = out_nc.createVariable(varname='depth_bounds', datatype='f8', dimensions=['depth', 'axis_nbounds'])
tmp[:,0] = upper_bounds
tmp[:,1] = lower_bounds

out_temp = out_nc.createVariable(varname='thetao', datatype='f4', dimensions=['time_counter', 'depth', 'y', 'x'], fill_value=fillval)
out_temp.standard_name = "conservative_temperature_mean_on_sublayers"
out_temp.units = "degC"
out_temp.coordinates = "time_counter depth nav_lat nav_lon"

out_sal = out_nc.createVariable(varname='so', datatype='f4', dimensions=['time_counter', 'depth', 'y', 'x'], fill_value=fillval)
out_sal.standard_name = "absolute_salinity_mean_on_sublayers"
out_sal.units = "g/kg"
out_sal.coordinates = "time_counter depth nav_lat nav_lon"



for m in range(0,n_layers):

    out_temp[:,m,:,:] = np.where( (np.all(thick_intg[m,:,:,:] == 0., axis=0))[None,:,:] * (np.ones(shape=[nt], dtype=bool))[:,None,None], \
                                 fillval, \
                                 np.sum(in_temp *   (thick_intg[m,:,:,:])[None,:,:,:] * (np.ones(shape=[nt], dtype=float))[:,None,None,None], axis=1) / np.sum((thick_intg[m,:,:,:])[None,:,:,:] * (np.ones(shape=[nt], dtype=float))[:,None,None,None], axis=1))


    out_sal[:,m,:,:] = np.where( (np.all(thick_intg[m,:,:,:] == 0., axis=0))[None,:,:] * (np.ones(shape=[nt], dtype=bool))[:,None,None], \
                                 fillval, \
                                 np.sum(in_sal *   (thick_intg[m,:,:,:])[None,:,:,:] * (np.ones(shape=[nt], dtype=float))[:,None,None,None], axis=1) / np.sum((thick_intg[m,:,:,:])[None,:,:,:] * (np.ones(shape=[nt], dtype=float))[:,None,None,None], axis=1))

    
    # print('dbg thick_intg', m, thick_intg[m,:,301,490])
    # print('dbg out temp', m, out_temp[m,:,301,490])

    # out_temp[:,m,:,:] = np.sum( (thick_intg[m,:,:,:])[None,:,:,:] * in_temp, axis=1)
    # print('isnan out', np.count_nonzero(np.isnan(out_temp[0,m,:,:])))


# out_nc.createDimension('z', nz)
# tmp = out_nc.createVariable(varname='thick_intg', datatype='f8', dimensions=['layer', 'z', 'y', 'x'])
# tmp[:] = thick_intg[:]

tmpdat.close()
out_nc.close()