import gstlearn as gl
import numpy as np

class LocalKriging :
    def __init__(self,driftFunctions,db,model):  
        
        # Input Data Bae
        self.ndim = db.getNDim()
        self.dbc = gl.Db(db)
        namesCoords = db.getNamesByLocator(gl.ELoc.X)
        coords = db[namesCoords]  
        for i,f in enumerate(driftFunctions):
            self.dbc["addedDrift" + str(i)] = f(coords)
        self.dbc.setLocators(["addedDrift*"],gl.ELoc.F)
        
        # Output Data Base
        self.dbout = gl.Db.createFromOnePoint([0 for i in range(self.ndim)])
        iptrEst = self.dbout.addColumnsByConstant(1,radix="z_estim",locatorType = gl.ELoc.Z)
        coords = self.dbout.getAllCoordinates().transpose()[0]
        for i,f in enumerate(driftFunctions):
            self.dbout.addColumns(f(np.atleast_2d(np.array(coords))), "addedDrift" + str(i))
        self.dbout.setLocators(["addedDrift*"],gl.ELoc.F)
        
        # Model
        self.modelc = gl.Model(model)
        self.modelc.setDriftIRF(0,len(driftFunctions))
        
        # Neighborhood (Unique)
        self.neigh = gl.NeighUnique()
        
        # For later use
        self.drifts = driftFunctions

        # Instantiate the KrigingSystem
        self.ks = gl.KrigingSystem(self.dbc, self.dbout, self.modelc, self.neigh)
        self.ks.updKrigOptEstim(iptrEst, -1, -1)
        self.ks.isReady()
            
    def eval(self,coordsTarget):
        for i in range(self.ndim):
            self.dbout.setCoordinate(0, i, coordsTarget[i])
        
        for i in range(len(self.drifts)):
            self.dbout.setLocVariable(gl.ELoc.F,0, i, self.drifts[i](np.atleast_2d(coordsTarget))[0])
        self.ks.estimate(0)
        return self.dbout["z_estim"][0]

db = gl.Db.create()
np.random.seed(123)
db["x"] = np.random.uniform(size=100)
db["y"] = np.random.uniform(size=100)
db["z"] = np.random.normal(size=100)
db.setLocators(["x","y"],gl.ELoc.X)
db.setLocators(["z"],gl.ELoc.Z)

def driftFunction1(coords):
    return coords[:,0]**2

def driftFunction2(coords):
    return coords[:,1]**2

driftFunctions = [driftFunction1,driftFunction2]

model = gl.Model.createFromParam(gl.ECov.MATERN,param=1)