import org.apache.commons.math3.geometry.euclidean.twod.Vector2D as Vector2D
import org.apache.commons.math3.linear.ArrayRealVector as ArrayRealVector
import org.apache.commons.math3.linear.Array2DRowRealMatrix as Array2DRowRealMatrix
import org.apache.commons.math3.linear.MatrixUtils as MatrixUtils
import org.apache.commons.math3.util.Pair as Pair
import org.apache.commons.math3.optimization.general.LevenbergMarquardtOptimizer as LevenbergMarquardtOptimizer
import org.apache.commons.math3.fitting.leastsquares.MultivariateJacobianFunction as MultivariateJacobianFunction
import org.apache.commons.math3.fitting.leastsquares.LeastSquaresBuilder as LeastSquaresBuilder
import org.apache.commons.math3.fitting.leastsquares.LevenbergMarquardtOptimizer as LevenbergMarquardtOptimizer



x = [12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0]
y = [30.0, 21.0, 17.0, 15.0, 16.0, 19.0, 22.0, 33.0]
v = [0.25, 0.5,  0.75, 0.33, 0.7.  0.35, 0.22, 0.35]

num_samples = len(x)
w = [1.0/i for i in v]

W = MatrixUtils.createRealDiagonalMatrix(w)




initial = [1.0, 1.0, 1.0]    #a, b, c


class Model(MultivariateJacobianFunction):
    def value(self, variables):
        value = ArrayRealVector(num_samples)
        jacobian = Array2DRowRealMatrix(num_samples, 3)        
        for i in range(num_samples):
            (a,b,c) = (variables.getEntry(0), variables.getEntry(1), variables.getEntry(2))   
            model = a*x[i]*x[i] + b*x[i] + c
            value.setEntry(i, model)  
                      
            # derivative with respect to a
            jacobian.setEntry(i, 0,  x[i]*x[i])
            # derivative with respect to b
            jacobian.setEntry(i, 1,  x[i])
            # derivative with respect to c
            jacobian.setEntry(i, 2,  1.0)            
        
        return Pair(value, jacobian)


model = Model()

# the target is to have all points at the specified radius from the center
target = [v for v in y]

problem = LeastSquaresBuilder().start(initial).model(model).target(target).lazyEvaluation(False).maxEvaluations(1000).maxIterations(1000).weight(W).build()

optimizer = LevenbergMarquardtOptimizer()
optimum = optimizer.optimize(problem)
optimalValues = optimum.getPoint()
a,b,c = optimum.getPoint().getEntry(0), optimum.getPoint().getEntry(1), optimum.getPoint().getEntry(2)

print "A: ",  a 
print "B: ",  b
print "C: ",  c
print ""
print "RMS: "          , optimum.getRMS()
print "evaluations: "  , optimum.getEvaluations()
print "iterations: "   , optimum.getIterations()