#!/usr/bin/env python3
#
# day2-ex4.py - build linear system for parabola fit, and visualise

import math
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()

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

# we use a Cholesky decomposition class from ROOT
#
# it expects data in C++ vectors, so repack in makeCppVector
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)

ROOT.gROOT.SetBatch()
c = ROOT.TCanvas("c", "c", 800, 600)

# build graph for fit result (and uncertainty)
gfit = ROOT.TGraphErrors(120)
gfit.SetNameTitle("gfit", "fit result;V_{true} [V];V_{meas} [V]")
for i in range(0, 121):
    x = xs[len(xs) - 1] * (i - 10) / 100.  # range to extend a bit beyond data points
    dx = x - x0
    gfit.SetPoint(i, x, (crhs[2] * dx + crhs[1]) * dx + crhs[0])
    tmp = [cmat[0] + dx * (cmat[1] + dx * cmat[3]),
           cmat[1] + dx * (cmat[2] + dx * cmat[4]),
           cmat[3] + dx * (cmat[4] + dx * cmat[5])]
    err2 = tmp[0] + dx * (tmp[1] + dx * tmp[2])
    gfit.SetPointError(i, 0., math.sqrt(err2))
gfit.SetLineColor(ROOT.kGreen + 2)
gfit.SetFillColor(ROOT.kGreen)
gfit.SetMarkerColor(ROOT.kGreen + 2)
gfit.SetMarkerSize(0)
gfit.SetLineWidth(2)

# build graph with data
gdata = ROOT.TGraphErrors(len(xs))
gdata.SetNameTitle("gdata", "data")
for i in range(0, len(xs)):
    gdata.SetPoint(i, xs[i], ys[i])
    gdata.SetPointError(i, 0., sigmas[i])
gdata.SetLineWidth(2)
gdata.SetMarkerStyle(7)
gdata.SetMarkerSize(2)

gfit.Draw("A3")  # draw axes, and error band
gfit.Draw("LX")  # draw fitted line (without error bars/bands)
gdata.Draw("P")  # draw data

l = ROOT.TLegend(0.2, 0.7, 0.4, 0.9)
l.AddEntry(gdata)
l.AddEntry(gfit)
l.Draw()

c.SaveAs("day2-ex4.pdf")
c.SaveAs("day2-ex4.png")