#!/usr/bin/env python

import sys, os
import argparse
import ast
import ESMP
import netCDF4 as nc
import numpy as np

"""
Regridding tool, based on ESMP tutorial at:
http://www.earthsystemmodeling.org/python_releases/last_esmp/python_doc/html/index.html
"""

class Regridder:

    def __init__(self, grid_def, src_mask, dest_mask):

        ESMP.ESMP_Initialize()
        ESMP.ESMP_LogSet(True)

        # Create source and destination grids.
        self.src_grid = make_grid(grid_def, 'src_', src_mask)
        self.dest_grid = make_grid(grid_def, 'dest_', dest_mask)

        # Fields to regrid
        self.src_field = ESMP.ESMP_FieldCreateGrid(src_grid, 'src_field')
        self.dest_field = ESMP.ESMP_FieldCreateGrid(dest_grid, 'dest_field')

        # Area and fraction fields.
        self.src_frac = ESMP.ESMP_FieldCreateGrid(src_grid, 'src_frac')
        self.dest_frac = ESMP.ESMP_FieldCreateGrid(dest_grid, 'dest_frac')

        self.src_area = ESMP.ESMP_FieldCreateGrid(src_grid, 'src_area')
        self.dest_area = ESMP.ESMP_FieldCreateGrid(dest_grid, 'dest_area')

        self.src_shape = grid_def.variables['src_lons'].shape
        self.dest_shape = grid_def.variables['dest_lons'].shape


    def __del__(self):
        ESMP.ESMP_FieldDestroy(self.src_field)
        ESMP.ESMP_FieldDestroy(self.dest_field)

        ESMP.ESMP_FieldDestroy(self.src_frac)
        ESMP.ESMP_FieldDestroy(self.dest_frac)

        ESMP.ESMP_FieldDestroy(self.src_area)
        ESMP.ESMP_FieldDestroy(self.dest_area)

        ESMP.ESMP_GridDestroy(self.src_grid)
        ESMP.ESMP_GridDestroy(self.dest_grid)

        ESMP.ESMP_Finalize()


    def make_grid(self, grid_def, prefix, mask=None):
        """
        Define a grid
        """
        lons = np.copy(grid_def.variables[prefix + 'lons'])
        lats = np.copy(grid_def.variables[prefix + 'lats'])
        clons = np.copy(grid_def.variables[prefix + 'clons'])
        clats = np.copy(grid_def.variables[prefix + 'clats'])

        # Create a grid with 1 periodic dimension.
        grid = ESMP.ESMP_GridCreate1PeriDim(np.array((lons.shape[1], lons.shape[0]), dtype=np.int32))

        # Add centres to grid.
        ESMP.ESMP_GridAddCoord(grid, staggerloc=ESMP.ESMP_STAGGERLOC_CENTER)
        lb_center, ub_center = ESMP.ESMP_GridGetCoord(grid, ESMP.ESMP_STAGGERLOC_CENTER)

        [x, y] = [0, 1]
        x_center = ESMP.ESMP_GridGetCoordPtr(grid, x, ESMP.ESMP_STAGGERLOC_CENTER)
        y_center = ESMP.ESMP_GridGetCoordPtr(grid, y, ESMP.ESMP_STAGGERLOC_CENTER)
        x_center[:] = lons.flatten()
        y_center[:] = lats.flatten()

        # Add corners to grid. 
        ESMP.ESMP_GridAddCoord(grid, staggerloc=ESMP.ESMP_STAGGERLOC_CORNER)
        lb_corner, ub_corner = ESMP.ESMP_GridGetCoord(grid, ESMP.ESMP_STAGGERLOC_CORNER)
        nx, ny = ub_corner[:] - lb_corner[:]

        x_corner = ESMP.ESMP_GridGetCoordPtr(grid, x, ESMP.ESMP_STAGGERLOC_CORNER)
        x_corner[:] = np.NAN
        x_corner = x_corner.reshape(ny, nx)
        y_corner = ESMP.ESMP_GridGetCoordPtr(grid, y, ESMP.ESMP_STAGGERLOC_CORNER)
        y_corner[:] = np.NAN
        y_corner = y_corner.reshape(ny, nx)

        # Need to copy over the corners which are shape (4, 300, 360) to flat a (301, 360) 
        # array (for example). This is done by copying over all the corner 0 coordinates,
        # then fixing up the top row. 

        x_corner[:ny-1,:] = clons[0,:,:] 
        y_corner[:ny-1,:] = clats[0,:,:]

        assert (np.isnan(np.sum(x_corner)) and np.isnan(np.sum(y_corner)))

        x_corner[ny-1:,:] = clons[3,-1,:]
        y_corner[ny-1:,:] = clats[3,-1,:]

        x_corner = x_corner.reshape(ny*nx)
        y_corner = y_corner.reshape(ny*nx)

        x_corner = (x_corner + 360.) % 360.

        # Check that there are no NANs left.
        assert not (np.isnan(np.sum(x_corner)) or np.isnan(np.sum(y_corner)))

        if mask is not None:
           ESMP.ESMP_GridAddItem(grid, item=ESMP.ESMP_GRIDITEM_MASK)
           m = ESMP.ESMP_GridGetItem(grid, item=ESMP.ESMP_GRIDITEM_MASK)
           mask = mask.reshape(m.shape[0])
           m[mask == True] = 1
           m[mask == False] = 0

        return grid


    def regrid(self, input_field):
        '''
        '''

        # Load the field to be regridded.
        src_field_ptr = ESMP.ESMP_FieldGetPtr(self.src_field) 
        src_field_ptr[:] = (input_field.reshape(self.src_field.size))[:]

        # Call the regridding functions
        routehandle = ESMP.ESMP_FieldRegridStore(self.src_field, self.dest_field,
                                                 srcMaskValues=np.array([1], dtype=np.int32),
                                                 dstMaskValues=np.array([1], dtype=np.int32),
                                                 regridmethod=ESMP.ESMP_REGRIDMETHOD_CONSERVE,
                                                 unmappedaction=ESMP.ESMP_UNMAPPEDACTION_ERROR,
                                                 srcFracField=self.src_frac, dstFracField=self.dest_frac)
        ESMP.ESMP_FieldRegrid(self.src_field, self.dest_field, routehandle)
        ESMP.ESMP_FieldRegridRelease(routehandle)

        # Unload the result.
        dest_field_ptr = ESMP.ESMP_FieldGetPtr(self.dest_field)
        dest_field_ptr = dest_field_ptr.reshape(self.dest_shape)

        return dest_field_ptr


def get_masks():

    return (None, None)

    src_mask = None
    dest_mask = None
    # Check whether the field has missing values, in which case a src mask is needed.
    if hasattr(src_file.variables[field_name], 'missing_value'):
        assert(args.src_mask_file is None and args.src_mask_var is None)
        src_mask = np.ma.getmask(src_file.variables[field_name][:])
        src_mask = src_mask[0,:,:]

        # FIXME: HACK to get a particular field regridded. Put 0 everywhere where there's a missing value.
        src_mask = None
        src_field_tmp[src_field_tmp == src_file.variables[field_name].missing_value] = 0.0

    if args.src_mask_file:
        # If no missing values, then use --src_mask_var
        assert(args.src_mask_var is not None)

        with nc.Dataset(args.src_mask_file) as f:
            src_mask = f.variables[args.src_mask_var][:]
            # In case src_mask is ints or reals, convert to proper masked array. 
            src_mask = np.ma.make_mask(src_mask)

    else:
        sys.stdout.write("INFO: no mask present, regridding without mask.\n")
        # If neither of the above, then give a message/warning but proceed anyway. 

    if src_mask is not None:
        if (args.dest_mask_file is None) or (args.dest_mask_var is None):
            sys.stderr.write("Source field has mask, a dest mask and var must be provided.\n") 
            return 1

        with nc.Dataset(args.dest_mask_file) as f:
            dest_mask = f.variables[args.dest_mask_var][:]
            # In case src_mask is ints or reals, convert to proper masked array. 
            dest_mask = np.ma.make_mask(dest_mask)

def regrid_this_field(field, src_file, grid_def):
    """
    Return whether or not a certain field should be regridded. 

    Regrid any 3D fields whose 2nd and 3rd dimensions match those of the grid definition.
    """

    return True


def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("grid_def", help="""
                        A grid definition file containing source and destination grid specifications.
                        The file must contain the following fields:
                            * src_lats, dest_lats - latitude points of the source and destination grids.
                            * src_longs, dest_longs - longitude points of the source and destination grids.
                            * src_clats, dest_clats - latitude of corners of the grids.
                            * src_clons, dest_clats - longitude of corners of the grids. 
                        Optionally it can also contain:
                            * src_mask, dest_mask - masks, all masked points are ingnored in the interpolation. 

                        There is an example grid definition file included with this program.
                        Any easy way to make a file like this is to use ncks. For example to copy and rename a 
                        field:
                        ncks -v latt_bonds ice_lgrid.nc grid_def.nc   # copies latt_bonds variable
                        ncrename -v latt_bonds,src_clats grid_def.nc  # renames latt_bonds variable
                        ncrename -d nx,nxl grid_def.nc                # renames nx dimension
                        """)
    parser.add_argument("source_file", help="""
                         The file containing fields to be regridded. They are must conform to the grid
                         definition given in the grid_def file described above. """)
    parser.add_argument("dest_file", help="""
                         The file that contains the regridded fields, this should not already exist.""")
    parser.add_argument("--masking", default="source", help="""
                        Specify how masking will be done. This can be one of:
                            * 'source' - use missing values from the source field as a mask.
                            * 'grid' - use the masks given in the definition file.
                            * 'none' - don't do any masking.
                        The default is 'source'.""")
    parser.add_argument("--skip_tests", action="store_true", help="Skip the runtime checks/tests.")

    args = parser.parse_args()

    grid_def = nc.Dataset(args.grid_def)
    src_file = nc.Dataset(args.src_file)
    assert(not os.path.exists(args.dest_file))
    dest_file = nc.Dataset(args.dest_file, 'w')

    # Make the destination file. 
    dest_file.createDimension('lons', len(grid_def.variables['dest_lons']))
    dest_file.createDimension('lats', len(grid_def.variables['dest_lons']))
    dest_file.createDimension('time', None)
    dest_file.createVariable('lons', 'f8', ('lons'))
    dest_file.createVariable('lats', 'f8', ('lats'))
    dest_file.createVariable('time', 'f8', ('time'))

    (src_mask, dest_mask) = get_masks(grid_def, src_file)

    # Open up the source and destination grids
    regridder = Regridder(grid_def, src_mask, dest_mask)

    # Iterate over all variables and time points, regridding each.

    for field_name in src_file.variables.keys():
        if not regrid_this_field(field_name, src_file, grid_def):
            continue

        dest_file.createVariable(field_name, 'f8', ('time', 'lons', 'lats'))

        for t in range(src_field_tmp.shape[0]):
            print 'Regridding at timestep %s' % t

            # Set the time value.
            dest_file.variables['time'][t] = src_file.variables['time'][t]
            # Do the regridding. 
            dest_file.variables[field_name][:] = regridder.regrid(src_file.variables[field_name][t,:,:])

            # Check some results. 
            #src_mass, src_area = compute_mass(src_field_cl, src_area_cl, src_frac_cl, True)
            #dest_mass, dest_area = compute_mass(dest_field, dest_area_cl, 0, False)
            #np.testing.assert_approx_equal(src_mass, dest_mass)

    src_file.close()
    dest_file.close()
    grid_def.close()

if __name__ == "__main__":
    sys.exit(main())
