from mpi4py import MPI
import numpy as np
import sys

import PyFVCOM as pf

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

if rank == 0:
    fvcom_file = sys.argv[1]
    lower_left_ll = np.asarray(sys.argv[2].split(','), dtype=float)
    upper_right_ll = np.asarray(sys.argv[3].split(','), dtype=float)
    grid_res = float(sys.argv[4])
    depth_layers = np.asarray(sys.argv[5].split(','), dtype=float)
    var_list_raw = sys.argv[6].split(',')
    varmatch = {'temp':'nodes', 'salinity':'nodes', 'u':'elements', 'v':'elements', 'zeta':'surface', 'h':'static_nodes'}
    varlist = {}
    for this_var in var_list_raw:
        varlist[this_var] = varmatch[this_var]

    fvcom = pf.read.FileReader(fvcom_file)
    nt = fvcom.dims.time
    start, stop = 0, nt
    global_time_indices = np.array_split(np.arange(start, stop), size)

else:
    fvcom_file = None
    lower_left_ll = None
    upper_right_ll = None
    grid_res = None
    depth_layers = None
    varlist = None
    global_time_indices = None
# get the fvcom points within the extended grid and land mask the grid mesh

fvcom_file = comm.bcast(fvcom_file, root=0)
lower_left_ll = comm.bcast(lower_left_ll, root=0)
upper_right_ll = comm.bcast(upper_right_ll, root=0)
grid_res = comm.bcast(grid_res, root=0)
depth_layers = comm.bcast(depth_layers, root=0)
varlist = comm.bcast(varlist, root=0)

time_indices = comm.scatter(global_time_indices, root=0)

interped_data = {}

worker = pf.interpolate.MPIRegularInterpolateWorker(fvcom_file, time_indices, comm=comm, verbose=True)
worker.InitialiseGrid(lower_left_ll, upper_right_ll, grid_res, depth_layers, time_varying_depth=True)

for var, var_mode in varlist.items(): 
    interped_data[var] = worker.InterpolateRegular(var, mode=var_mode)

all_interped_data = comm.gather(interped_data, root=0)

# resort the data by timestep
if rank == 0:
    collected_interp_data = {}

    for this_var in varlist.keys():
        output_list = []
        for this_dict in all_interped_data:
            output_list.append(this_dict[this_var])
        collected_interp_data[this_var] = np.ma.masked_invalid(np.vstack(output_list))

    # write to cmems format
    fvcom = pf.read.FileReader(fvcom_file)

    dims = {'time':len(fvcom.time.datetime), 'depth':len(worker.regular_grid.dep_lays), 'lon':len(worker.regular_grid.lons),
                                'lat':len(worker.regular_grid.lats), 'DateStrLen': 26} 

    output_file_name = 'output.nc'

    all_attributes = {'lon':{'standard_name':'longitude', 'units':'degrees_east', 'long_name':'longitude'},
            'lat':{'standard_name':'latitude', 'units':'degrees_north', 'long_name':'latitude'},
            'depth':{'standard_name':'depth', 'units':'m', 'long_name':'depth, measured downwards from free surface', 'axis':'Z', 'positive':'down'},
            'h':{'standard_name':'h', 'units':'m', 'long_name':'model bathymetry depth, measured downwards from geoid', 'axis':'Z', 'positive':'down'},
            'temp':{'standard_name':'sea_water_potential_temperature','units':'C','missing_value':-32768, 'long_name':'Sea Water Potential Temperature'} ,
            'salinity':{'standard_name':'sea_water_salinity','units':'psu', 'missing_value':-32768, 'long_name':'Sea Water Salinity'},
            'u':{'standard_name':'eastward_sea_water_velocity','units':'m s-1', 'missing_value':-32768, 'long_name':'Eastward Current Velocity'},
            'v':{'standard_name':'northward_sea_water_velocity','units':'m s-1', 'missing_value':-32768, 'long_name':'Northward Current Velocity'},
            'zeta':{'standard_name':'sea_surface_height_above_geoid','units':'m', 'missing_value':-32768,'long_name':'Sea surface height above geoid'}}

    globals = {'type': 'OPERATIONAL MODEL REGULAR GRID OUTPUT',
                    'title': 'Regularly gridded data interpolated from FVCOM output',
                    'history': 'File created using PyFVCOM',
                    'filename': str(output_file_name),
                    'Conventions': 'CF-1.0',
                    '_FillValue':-32768}

    for this_val in collected_interp_data.values():
        np.ma.set_fill_value(this_val, -32768)
    ncopts = {}

    #for this_var, this_mode in varlist.items():
    #    this_output_file_name = 'output_{}.nc'.format(this_var)
    
    with pf.preproc.WriteForcing(output_file_name, dims, global_attributes=globals, clobber=True, format='NETCDF4') as outfile:
        # Add the variables.
        outfile.write_fvcom_time(fvcom.time.datetime)
        outfile.add_variable('lon', worker.regular_grid.lons, ['lon'], attributes=all_attributes['lon'])
        outfile.add_variable('lat', worker.regular_grid.lats, ['lat'], attributes=all_attributes['lat'])
        outfile.add_variable('depth', worker.regular_grid.dep_lays, ['depth'], attributes=all_attributes['depth'])
    
        for this_var, this_mode in varlist.items():
            if this_mode == 'surface':
                outfile.add_variable(this_var, collected_interp_data[this_var], ['time', 'lon', 'lat'],
                                         attributes=all_attributes[this_var], ncopts=ncopts)
            elif this_mode == 'static_nodes':
                outfile.add_variable(this_var, collected_interp_data[this_var], ['lon', 'lat'],
                                         attributes=all_attributes[this_var], ncopts=ncopts)
            else:
                outfile.add_variable(this_var, collected_interp_data[this_var], ['time', 'lon', 'lat', 'depth'],
                                         attributes=all_attributes[this_var], ncopts=ncopts)
            
else:
    print('Rank {} process finished'.format(rank))