#!/usr/bin/env python3
#
# day2-ex2.py - build linear system for parabola fit

import ROOT

def buildLinSystemForParabolaFit(xs, ys, sigmays, x0):
    rhs = [0, 0, 0]
    mat = [0, 0, 0, 0, 0, 0]
    for x, y, sigma in zip(xs, ys, sigmays):
        w = 1 / (sigma * sigma)
        dx = x - x0
        dx2 = dx * dx;
        rhs[0] += w * y
        rhs[1] += w * y * dx
        rhs[2] += w * y * dx2
        mat[0] += w
        mat[1] += w * dx
        mat[2] += w * dx2
        mat[4] += w * dx * dx2
        mat[5] += w * dx2 * dx2

    mat[3] = mat[2]

    return rhs, mat

xs = [0, 1, 2, 3, 4, 5]
ys = [-.10, .84, 1.91, 2.88, 4.06, 4.83]
sigmas = [.1, .1, .1, .1, .1, .1]

class MatAdapter:
    def __init__(self, m):
        self._m = m
    
    def at(self, i, j):
        if j > i:
            tmp = i
            i = j
            j = tmp
        idx = int((i * (i + 1) / 2) + j)
        return self._m[idx]

def prettyPrintMatrix(m):
    mm = MatAdapter(m)
    for i in range(0, 3):
        for j in range(0, 3):
            print(" {}".format(mm.at(i, j)), end = "")
        print()

rhs, mat = buildLinSystemForParabolaFit(xs, ys, sigmas, 2.5)
print(rhs)
prettyPrintMatrix(mat)

def makeCppVector(l):
    retVal = ROOT.std.vector("double")()
    retVal.reserve(len(l))
    for i in range(0, len(l)):
        retVal.push_back(l[i])
    return retVal

crhs = makeCppVector(rhs)
cmat = makeCppVector(mat)

decomp = ROOT.Math.CholeskyDecompGenDim("double")(3, cmat.data())
if not decomp:
    raise "Not Invertible"
decomp.Solve(crhs)
decomp.Invert(cmat.data())
print(crhs)
print(cmat)

def calcChi2(xs, ys, sigmas, x0, param):
    retVal = 0
    for x, y, s in zip(xs, ys, sigmas):
        w = 1 / (s * s)
        dx = x - x0
        ypred = param[0] + dx * (param[1] + dx * param[2])
        retVal += (y - ypred) * (y - ypred) * w
    return retVal

chi2 = calcChi2(xs, ys, sigmas, 2.5, crhs)
prob = ROOT.TMath.Prob(chi2, 3)

print("chi={} prob={}".format(chi2, prob))