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

import ROOT
import math

def makeCppVector(myList):
    """ turn a python list into a C++ vector of doubles """
    retVal = ROOT.std.vector("double")()
    retVal.reserve(len(myList))
    for i in range(0, len(myList)):
        retVal.push_back(myList[i])
    return retVal


def buildLinSystemForFit(modelfuncs, xs, ys, sigmays, x0):
    n = len(modelfuncs)
    rhs = n * [0]
    mat = int(n * (n + 1) / 2) * [0]
    for x, y, sigma in zip(xs, ys, sigmays):
        w = 1 / (sigma * sigma)
        dx = x - x0
        terms = [modelfuncs[i](dx) for i in range(0, n)]
        for i in range(0, n):
            rhs[i] += w * y * terms[i]
            for j in range(0, i + 1):
                mat[int(i * (i + 1) / 2) + j] += w * terms[i] * terms[j]

    return rhs, mat


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


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 runPolynomialFit(n, xs, ys, sigmays, x0):
    modelfuncs = [lambda x, ii = i: x ** ii for i in range(0, n)]
    rhs, mat = buildLinSystemForFit(modelfuncs, xs, ys, sigmays, x0)
    
    crhs = makeCppVector(rhs)
    cmat = makeCppVector(mat)
    decomp = ROOT.Math.CholeskyDecompGenDim("double")(n, cmat.data())
    if not decomp:
        raise "Not Invertible"
    decomp.Solve(crhs)
    decomp.Invert(cmat.data())

    ndf = len(xs) - len(rhs)
    chi2 = calcChi2(modelfuncs, xs, ys, sigmays, x0, crhs)

    return {'params': crhs, 'cov': MatAdapter(cmat),
            'chi2': chi2, 'ndf': ndf}


def makeCorrel(n, cov):
    correl = []
    for i in range(0, n):
        correl.append([])
        for j in range(0, i + 1):
            correl[i].append(cov.at(i, j) /
                             math.sqrt(cov.at(i, i) * cov.at(j, j)))
    return correl


npoints = 50

xs = [2.5/0.4 * i / (npoints - 1) for i in range(0, npoints)]
ys = [math.tanh(x) for x in xs]
sigmas = npoints * [1]
x0 = 0.5 * 2.5 / 0.4

poly = runPolynomialFit(5, xs, ys, sigmas, x0)

print("polynomial fit deg. 5:")
print()
print("params: {}".format(poly['params']))
print("correlations")
correl = makeCorrel(5, poly['cov'])
for row in correl:
    print(row)