from netCDF4 import Dataset
import sys
import time
import numpy as np
from collections import OrderedDict
from nn_interp import nn_interp_1d

# Script for generating NEMO BDY data files from pre-interpolated data, and a NEMO BDY geometry file.
# Author: C. Pelletier

# Call:
# python generate_new_bdy_dta_onefile.py <input_file> <grid> <output_file>
# where:
# <input_file>: is an input dataset which has been preinterpolated to the NEMO grid, and contains all the fields you want to generate data to.
#              This script will "detect" any data fields in <input_file>.
# <grid> is the type of NEMO grid points (T, U or V) that your input data is defined on.
#             All fields within one <input_data> file MUST share the same type of points (e.g., don't mix velocities and tracers).
#             If the input file has 3D fields, then it should contain a depth variable AND a depth dimension, both called "depth<grid>" (e.g. "depthv" for V nodes)
# <output_file> is the path to an output file.

# BDY geometry file, coming from generate_new_bdy_geom.py
geom_file = "/storepelican/pelletie/nemo_bdy/grids/eORCA025_SO_BDY-COORDS_rim1.nc"

if(len(sys.argv) == 4):
    input_file = sys.argv[1]
    grid = sys.argv[2]
    output_file = sys.argv[3]
else:
    print('Please provide three input arguments: python generate_new_bdy_dta_onefile.py <input_file> <grid> <output_file>')
    sys.exit()

if(grid not in ['t', 'u', 'v']):
    print('grid should be t, u or v')
    sys.exit()

### END OF INPUT SPECIFICATION.


geom_nc = Dataset(geom_file, mode='r')


nbis = geom_nc.variables['nbi'+grid][0,:] - 1
nbjs = geom_nc.variables['nbj'+grid][0,:] - 1
nbrs = geom_nc.variables['nbr'+grid][0,:]
nbdys = len(geom_nc.dimensions['xb'+grid.upper()])
lons = geom_nc.variables['glam'+grid][0,:]
lats = geom_nc.variables['gphi'+grid][0,:]

geom_nc.close()

        
in_nc = Dataset(input_file, mode='r')
out_nc = Dataset(output_file, mode='w')

nt = len(in_nc.dimensions['time_counter'])

ny = len(in_nc.dimensions['y'])
nx = len(in_nc.dimensions['x'])

has_3d = False
nxb = nbdys


remap_vars = []

for name, var in in_nc.variables.items():
    if(( 'time_counter' in var.dimensions) \
       and ('x' in var.dimensions)):
        remap_vars.append(name)
        
        if('depth'+grid in var.dimensions):
            has_3d = True
            nz = len( in_nc.dimensions['depth'+grid] )

out_nc.createDimension('yb', 1)
out_nc.createDimension('xb'+grid.upper(), nxb)
out_nc.createDimension('time_counter', 0)
if(has_3d):
    out_nc.createDimension('depth'+grid, nz)

idx_flat_2d = np.repeat( np.arange(nt) * ny*nx, nxb ) \
              + np.tile(nbjs * nx, nt) \
              + np.tile(nbis, nt)

time_in = in_nc.variables['time_counter']
time_out = out_nc.createVariable(varname = 'time_counter', datatype = time_in.datatype, dimensions = time_in.dimensions)
time_out.setncatts(time_in.__dict__)
time_out[:] = time_in[:]

nbi_out = out_nc.createVariable(varname='nbi'+grid, datatype='i4', dimensions=['yb', 'xb'+grid.upper()])
nbi_out[0,:] = nbis[:] + 1

nbj_out = out_nc.createVariable(varname='nbj'+grid, datatype='i4', dimensions=['yb', 'xb'+grid.upper()])
nbj_out[0,:] = nbjs[:] + 1

nbr_out = out_nc.createVariable(varname='nbr'+grid, datatype='i4', dimensions=['yb', 'xb'+grid.upper()])
nbr_out[0,:] = nbrs[:]

if(has_3d):
    depth_in = in_nc.variables['depth'+grid]
    depth_out = out_nc.createVariable(varname = 'depth'+grid, datatype = depth_in.datatype, dimensions = depth_in.dimensions)
    depth_out.setncatts(depth_in.__dict__)
    depth_out[:] = depth_in[:]

    idx_flat_3d = np.repeat( np.arange(nt*nz) * ny*nx, nxb ) \
                  + np.tile(nbjs * nx, nt*nz) \
                  + np.tile(nbis, nt*nz)


lon_out = out_nc.createVariable(varname='nav_lon', datatype='f8', dimensions=['yb', 'xb'+grid.upper()])
lon_out.standard_name = "longitude"
lon_out.units = "degrees East"
lon_out[0,:] = lons[:]

lat_out = out_nc.createVariable(varname='nav_lat', datatype='f8', dimensions=['yb', 'xb'+grid.upper()])
lat_out.standard_name = "latitude"
lat_out.units = "degrees North"
lat_out[0,:] = lats[:]


for var in remap_vars:
    in_var = in_nc.variables[var]

    if( 'depth'+grid in in_var.dimensions ):

        out_var = out_nc.createVariable(var, datatype='f4', \
                                        dimensions=['time_counter', 'depth'+grid, 'yb', 'xb'+grid.upper()], fill_value = 1.e20)


        out_var[:,:,0,:] = np.reshape( (np.ndarray.flatten(in_var[:]))[idx_flat_3d], \
                                       (nt, nz, nxb) )


        if(np.ma.is_masked(out_var[:])):
            for t in range(0,nt):
                for k in range(0,nz):
                    if(np.all( out_var[:].mask[t,k,0,:] ) ):
                        # All masked nodes: take from upper level.
                        out_var[t,k,0,:] = out_var[t,k-1,0,:]
                    elif(np.any(out_var[:].mask[t,k,0,:])):
                        # Some masked nodes: drown at same level.
                        out_var[t,k,0,:] = nn_interp_1d(out_var[t,k,0,:])

    else:
        out_var = out_nc.createVariable(var, datatype='f4', \
                                        dimensions=['time_counter', 'yb', 'xb'+grid.upper()], \
                                        fill_value = 1.e20)
        out_var[:,0,:] =  np.reshape( (np.ndarray.flatten(in_var[:]))[idx_flat_2d], \
                                       (nt, nxb) )

        if(np.ma.is_masked(out_var[:])):
            for t in range(0,nt):
                if(np.any(out_var[:].mask[t,0,:])):
                    out_var[t,0,:] = nn_interp_1d(out_var[t,0,:])

in_nc.close()
out_nc.close()