from mathutils import *
from plotutils import *

#Fitting the quadratic function f(x) = a x2 + b x + c
#Data created with [a = 8, b = 10, c = 16] and 0<noise<1

x = [ 1.0,  2.0,   3.0,   4.0,   5.0,   6.0,   7.0,   8.0,   9.0,  10.0]
y = [36.0, 66.0, 121.0, 183.0, 263.0, 365.0, 473.0, 603.0, 753.0, 917.0]
num_samples = len(x)
w = [ 1.0] * num_samples
W = MatrixUtils.createRealDiagonalMatrix(w)

p=plot(y, xdata=x)[0]


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()
print ""
for i in range (num_samples):
    print x[i], y[i], poly(x[i], [c,b,a])

 
plot_function(p, PolynomialFunction((c,b,a)), "Fit", x)