#!/usr/bin/env python3
#
# day4-ex4.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


class ChebyshevIterator:
    def __init__(self, x, stop = 10):
        self._prev = 1
        self._cur = x
        self._twox = 2 * x
        self._stop = stop
    def __iter__(self):
        return self
    def __next__(self):
        tmp = self._prev
        self._prev = self._cur
        self._cur = self._twox * self._prev - tmp
        if 0 == self._stop:
            raise StopIteration
        self._stop = self._stop - 1
        return tmp


def buildLinSystemForFit(n, termbuilder, xs, ys, sigmays, x0):
    rhs = n * [0]
    mat = int(n * (n + 1) / 2) * [0]
    for x, y, sigma in zip(xs, ys, sigmays):
        w = 1 / (sigma * sigma)
        terms = [t for t in termbuilder(x, 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(n, termbuilder, xs, ys, sigmas, x0, param):
    retVal = 0
    for x, y, s in zip(xs, ys, sigmas):
        w = 1 / (s * s)
        terms = [t for t in termbuilder(x, 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 runChebyshevFit(n, xs, ys, sigmays, xa, xb):
    x0 = .5 * (xa + xb)
    ixw = 2 / (xb - xa)
    termbuilder = lambda x, n: ChebyshevIterator(ixw * (x - x0), n)
    rhs, mat = buildLinSystemForFit(n, termbuilder, 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(n, termbuilder, 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]
xa, xb = 0, 2.5 / 0.4

poly = runChebyshevFit(5, xs, ys, sigmas, xa, xb)

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