# -*- Mode: python -*-
### -
### - Correct rmp file computed by OASIS MCT
## -
## -
## - Olivier Marti (olivier.marti@lsce.ipsl.fr)
## - Institut Pierre Simon Laplace
## - Laboratoire des Sciences du Climat et de l'Environnment
## -
#
# This program overcomes an interpolation bug in Scrip, uses by OASIS-MCT
#
# In IPSLCM6 at resolution ORCA1 x LMD 144x144, the bilinear interpolation of the wind stress is bugged.
# Two NEMO points around 0N/180W takes there wind stress values from points at 0N/0E.
#
# The bilinear interpolation compute weights and adresses that are written in a file :
# rmp_tlmd_to_torc_BILINEAR_ComputedByOasis.nc.
# The program copies the file to rmp_tlmd_to_torc_BILINEAR_Corrected.nc and makes a
# few corrections of adresses and weigths.
#
## See : https://forge.ipsl.jussieu.fr/igcmg/wiki/IPSLCM6/Coupling/BugBilinear
#

## - Needed Python modules
import cdms2, MV2, numpy as np
import os, time, shutil
from cdms2.coord import TransientVirtualAxis 

## - NetCDF parameters
# - If you want to turn that off or set different values of compression use the functions:
cdms2.setNetcdfShuffleFlag     (0)
cdms2.setNetcdfDeflateFlag     (0)
cdms2.setNetcdfDeflateLevelFlag(0)

#
rad = np.rad2deg (1.0) 

## Copy initial rmp file with bilinear weights
shutil.copyfile ( 'rmp_tlmd_to_torc_BILINEAR_ComputedByOasis.nc' , 'rmp_tlmd_to_torc_BILINEAR_Corrected.nc' )

## Read the new file
rmp_file             = cdms2.open   ( 'rmp_tlmd_to_torc_BILINEAR_Corrected.nc' , 'r+')

src_grid_center_lat  = rmp_file ( 'src_grid_center_lat' )
src_grid_center_lon  = rmp_file ( 'src_grid_center_lon' )
dst_grid_center_lat  = rmp_file ( 'dst_grid_center_lat' )
dst_grid_center_lon  = rmp_file ( 'dst_grid_center_lon' )
src_address          = rmp_file ( 'src_address'         )
dst_address          = rmp_file ( 'dst_address'         )
remap_matrix         = rmp_file ( 'remap_matrix'        )
src_grid_dims        = rmp_file ( 'src_grid_dims'       )
dst_grid_dims        = rmp_file ( 'dst_grid_dims'       )

src_grid_corner_lat  = rmp_file ( 'src_grid_corner_lat' )
src_grid_corner_lon  = rmp_file ( 'src_grid_corner_lon' )
src_grid_area        = rmp_file ( 'src_grid_area'       )
src_grid_imask       = rmp_file ( 'src_grid_imask'      )

dst_grid_corner_lat  = rmp_file ( 'dst_grid_corner_lat' )
dst_grid_corner_lon  = rmp_file ( 'dst_grid_corner_lon' )
dst_grid_area        = rmp_file ( 'dst_grid_area'       )
dst_grid_imask       = rmp_file ( 'dst_grid_imask'      )

# Get dimensions
( src_grid_size )       = src_grid_center_lat.shape
( dst_grid_size )       = dst_grid_center_lat.shape
( num_links, num_wgts)  = remap_matrix.shape
( src_grid_rank )       = src_grid_dims.shape
( dst_grid_rank )       = dst_grid_dims.shape

print 'Dimensions '
print 'src_grid_size : ', src_grid_size, '  src_grid_dims : ', src_grid_dims
print 'dst_grid_size : ', dst_grid_size, '  dst_grid_dims : ', dst_grid_dims
print ' '

## Get list of wrong weights
list_ind = np.where ( np.abs(remap_matrix) > 1.0 )[0]
print 'List of indexes of buggy points : ', list_ind
print ' '

for ind in list_ind[:]:
    print 'Working on point : ', ind, remap_matrix[ind]
    ## Compute 2D indexes
    ## warning : ind is a 'C' index
    src_grid_ind_2 = (src_address[ind]+1)/src_grid_dims[0] + 1
    src_grid_ind_1 = (src_address[ind]+1) - ( src_grid_ind_2 - 1 ) * src_grid_dims[0]

    dst_grid_ind_2 = dst_address[ind]/dst_grid_dims[0] + 1
    dst_grid_ind_1 = dst_address[ind] - ( dst_grid_ind_2 - 1 ) * dst_grid_dims[0]
    
    print src_address[ind], src_grid_ind_1, src_grid_ind_2, src_grid_center_lon[src_address[ind]]*rad, src_grid_center_lat[src_address[ind]]*rad
    print dst_address[ind], dst_grid_ind_1, dst_grid_ind_2, dst_grid_center_lon[dst_address[ind]]*rad, dst_grid_center_lat[dst_address[ind]]*rad
    print dst_grid_center_lon[dst_address[ind]]*rad - src_grid_center_lon[src_address[ind]]*rad, 
    print dst_grid_center_lat[dst_address[ind]]*rad - src_grid_center_lat[src_address[ind]]*rad

    #
    new_src_grid_ind_1 = np.mod (  (src_grid_ind_1 - src_grid_dims[0]/2)-1, src_grid_dims[0]) + 1 # Decalage de 180 degres ...
    new_src_grid_ind   = ( src_grid_ind_2 - 1 ) * src_grid_dims[0] + new_src_grid_ind_1 - 1
        
    print src_address[ind], new_src_grid_ind, new_src_grid_ind_1, src_grid_ind_2
    src_address [ind] = new_src_grid_ind

    dist = dst_grid_center_lon[dst_address[ind]-1]*rad - src_grid_center_lon[src_address[ind]-1]*rad

    if dist == 0.0:
        remap_matrix [ind] = 1.0
        other_src_grid_ind = 0 ; delta_x = 0.0
    else:
        if dist > 0.0:
            other_src_grid_ind_1 = np.mod ( (new_src_grid_ind_1 - 1 )-1, src_grid_dims[0]) + 1
        if dist < 0.0:
            other_src_grid_ind_1 = np.mod ( (new_src_grid_ind_1 + 1 )-1, src_grid_dims[0]) + 1
        
        other_src_grid_ind   = ( src_grid_ind_2 - 1 ) * src_grid_dims[0] + other_src_grid_ind_1 - 1
        delta_x = src_grid_center_lon[new_src_grid_ind]*rad- src_grid_center_lon[other_src_grid_ind]*rad
        remap_matrix [ind] = 1.0 - dist / delta_x
  
    print 'New ', new_src_grid_ind, other_src_grid_ind, dist, delta_x, remap_matrix[ind]
    print src_address[ind], new_src_grid_ind_1, src_grid_ind_2, src_grid_center_lon[src_address[ind]-1]*rad, src_grid_center_lat[src_address[ind]-1]*rad
    print dst_address[ind],     dst_grid_ind_1, dst_grid_ind_2, dst_grid_center_lon[dst_address[ind]-1]*rad, dst_grid_center_lat[dst_address[ind]-1]*rad
    print dst_grid_center_lon[dst_address[ind]-1]*rad - src_grid_center_lon[src_address[ind]-1]*rad, 
    print dst_grid_center_lat[dst_address[ind]-1]*rad - src_grid_center_lat[src_address[ind]-1]*rad
        
## Write corrected fields
rmp_file.write ( remap_matrix )
rmp_file.write ( src_address  )
 
rmp_file.close ()

## That's all folks !!!