#!/apps/PYTHON/2.7.3/bin/python
# This function extrapolates horizontally each level of a (x,y,z)
# field using the nearest neighbour method.
#
# Usage   : python extrap.py <input file> <input variable name>
#    <grid description file> <2d longitude name in grid file>
#    <2d latitude name in grid file> <mask 1=extrapolate>
#    <mask variable name> <output file>
#
# History : Virginie Guemas - Initial version     -       2012
#           Virginie Guemas - Masking the outputs - March 2014
#           Francesco Benincasa - Optimizing      - Jan   2015
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

import sys
import cdms2 as cdms
from cdms2 import MV
import numpy as N
N.set_printoptions(threshold=N.inf)
from scipy.spatial import ckdtree

cdms.setNetcdfShuffleFlag(1) #returning 1 if shuffling is enabled, 0 otherwise
cdms.setNetcdfDeflateFlag(1) #returning 1 if deflate is used, 0 otherwise
cdms.setNetcdfDeflateLevelFlag(1)


class MainClass:

    def __init__(self):

        self.X      = None
        self.Y      = None
        self.lat    = None
        self.lon    = None
        self.Pfill  = None
        self.var0_a = None
        self.var4   = None
        self.mask3d = None
        self.KDT    = None


    #@profile
    def lon_lat_to_cartesian(self, lon, lat, R=1):
        """calculates lon, lat coordinates of a point on a sphere with
        radius R
        """
        lon_r = N.radians(lon)
        lat_r = N.radians(lat)
        x = R * N.cos(lat_r) * N.cos(lon_r)
        y = R * N.cos(lat_r) * N.sin(lon_r)
        z = R * N.sin(lat_r)
        return x, y, z


    #@profile
    def mainFunc(self, fileIN, varIN, fileM, varlon, varlat, fileF, varF, fileOUT):
        #
        # 2. Get the input files
        # =======================
        #
        f=cdms.open(fileIN)
        self.var0_a=f(varIN,squeeze=1)
        f.close()

        self.mask3d=self.var0_a.mask
        (lz1,ly0,lx0)=self.var0_a.shape

        f=cdms.open(fileM)
        lon=f(varlon,squeeze=1)
        lat=f(varlat,squeeze=1)
        f.close()

        self.lon = lon
        self.lat = lat

        f=cdms.open(fileF)
        self.Pfill=f(varF,squeeze=1)
        f.close()

        m3d_tmp = self.mask3d
        v0a_tmp = self.var0_a
        vr4_tmp = N.zeros((lz1,ly0,lx0), dtype='d')

        # var4 intitialization
        self.var4 = N.where(m3d_tmp==False,v0a_tmp,vr4_tmp)

        #2d to 1d coordinates
        nlon = lon.flatten().data
        nlat = lat.flatten().data

        #indexes of points to extrapolate
        indexes = N.where(self.var0_a.mask & (self.Pfill>=1))

        #each level has a different mask, loop over the levels
        for lev in N.arange(lz1-1):
            print('LEV', lev)
            #orig matrix at level "lev"
            newvar = self.var0_a[lev]
            #not masked coordinates for the present level
            nomlon = lon.data[~newvar.mask]
            nomlat = lat.data[~newvar.mask]

            #cartesian coordinates of not masked points
            levxs, levys, levzs = self.lon_lat_to_cartesian(nomlon, nomlat)
            #print(levxs.shape, levys.shape, levzs.shape)
            levkdt = ckdtree.cKDTree(zip(levxs, levys, levzs))

            #points to extrapolate at the current level
            newindexes = N.where(indexes[0]==lev)
            latidxs = indexes[1].data[newindexes]
            lonidxs = indexes[2].data[newindexes]
            exlon = nlon[lonidxs]
            exlat = nlat[latidxs]

            #cartesian coordinates of the points to extrapolate
            xt, yt, zt = self.lon_lat_to_cartesian(exlon, exlat)
            #calculates distances and indexes of the nearest neighbours
            dist, idx = levkdt.query(zip(xt, yt, zt), k=2)

            #assign to the result matrix the found values
            self.var4[lev, latidxs, lonidxs] = newvar.data[~newvar.mask].flatten()[idx.T[1]]

        maskout = MV.where(self.Pfill>0.5,False,self.mask3d)
        var4 = cdms.createVariable(self.var4,id=self.var0_a.id)
      #  var4 = MV.where(maskout<0.5,var4,1e20)
        var4.getAxis(0).id = 'z'
        var4.getAxis(0)[:] = self.var0_a.getAxis(0)[:]
        var4.getAxis(1).id = 'y'
        var4.id = self.var0_a.id

        h = cdms.open(fileOUT,'w')
        h.write(var4)
        h.close()


if __name__ == "__main__":
    #
    # 1. Input arguments
    # ===================
    #
    # Input file and var names :
    # ---------------------------
    fileIN=sys.argv[1]
    varIN=sys.argv[2]
    #
    # Meshmask file :
    # ----------------
    fileM=sys.argv[3]
    varlon=sys.argv[4]
    varlat=sys.argv[5]
    #
    # Location of points to fill :
    # ----------------------------
    fileF=sys.argv[6]
    varF=sys.argv[7]
    #
    # Output file name :
    # -------------------
    fileOUT=sys.argv[8]

    mc = MainClass()
    #execute
    mc.mainFunc(fileIN, varIN, fileM, varlon, varlat, fileF, varF, fileOUT)
