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

# Inputs:
# a) a set of processor-specific NEMO restart files defined with a certain domain decomposition, that you'll want to run at (2*n_proc, n_proc for oce and n_proc for ice). You can get this by doing a dummy run with the number of processors you'll eventually want.
# b) a pair of rebuilt (processor-wise) NEMO restart files, which you'll want to get the actual restart data from

# Outputs:
# a set of processor-specific NEMO restart files, using the a) domain decomposition, with the data coming from b)


# input a). The script expects the files to be <geom_in_stencil><oce,ice>_XXXX.nc, where XXXX are four numbers.
geom_in_stencil = "/afast/pelletie/nemo_rs/geom_in/mpp_209_fET171-h40/so025_paraso_temoin_00000096_restart_"
# Number of processors for destination geom (geom_in_stencil)
n_newdisp = 209

# input b) The script expects the files to be <rebuilt_rs_stencil><oce,ice>.nc
rebuilt_rs_stencil = "/afast/pelletie/nemo_rs/rs_data/h40/org/so025_fET171_spin_h40_00035040_restart_"
## WARNING: be sure that your b) files contain all the fields you'll need.
##          Some NEMO options (melt ponds, iscpl) create/require new restart fields.

# Output files. The script will write <output_stencil><oce,ice>_XXXX.nc
output_stencil = "/afast/pelletie/nemo_rs/rs_data/h40/redisp/mpp_209_fET171-h40/so025_fET171_spin_h40_00035040_remap209-mpp_restart_"

nz = 75
media = ['ice', 'oce']

for medium in media:

    big_rs = Dataset(rebuilt_rs_stencil+medium+'.nc', mode='r')
        
    for m in range(0,n_newdisp):

        geom_nc = Dataset(geom_in_stencil+medium+'_%04d'%m+'.nc', mode='r')

        nx_nd, ny_nd = geom_nc.DOMAIN_size_local
        sx_beg, sy_beg = geom_nc.DOMAIN_position_first - 1
        sx_end, sy_end = geom_nc.DOMAIN_position_last
        nx_nd, ny_nd = geom_nc.DOMAIN_size_local

        if( sx_end - sx_beg != nx_nd or sy_end - sy_beg != ny_nd):
            print('error nx, beg end')
            sys.exit()
        
        fmt = "%04d"%m
        outfile = output_stencil+medium+'_'+fmt+'.nc'

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

        out_nc.createDimension('x', nx_nd)
        out_nc.createDimension('y', ny_nd)
        out_nc.createDimension('z', nz)
        out_nc.createDimension('t', 0)

        print(sx_beg,sx_end,sy_beg,sy_end)
        for name, variable in big_rs.variables.items():
            varout = out_nc.createVariable(name, variable.datatype, variable.dimensions)
            len_tmp = len(variable.dimensions)
            
            if( 'x' in variable.dimensions):
                if(len_tmp == 2):
                    varout[:] = big_rs.variables[name][sy_beg:sy_end,sx_beg:sx_end]
                elif(len_tmp == 3):
                    varout[:] = big_rs.variables[name][:,sy_beg:sy_end,sx_beg:sx_end]
                elif(len_tmp == 4):
                    varout[:] = big_rs.variables[name][:,:,sy_beg:sy_end,sx_beg:sx_end]
                    
            else:
                varin = big_rs.variables[name][:]
                if(len_tmp > 0 ):
                    varout[:] = varin[:]
                else:
                    varout[:] = varin

        out_nc.setncatts(geom_nc.__dict__)
                
        out_nc.close()
        print('generated '+outfile)

        geom_nc.close()
        
    big_rs.close()