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

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)