#!/usr/bin/env python3
#
# day1-ex2.py - straight line fit

import math

def invert2x2(m):
    """ invert 2x2 matrix - we should use something better here! """
    det = m[0][0] * m[1][1] - m[0][1] * m[1][0]
    return [[m[1][1] / det, -m[1][0] / det],
            [-m[0][1] / det, m[0][0] / det]]


def mulM2x2V(m, v):
    """ multiply a 2x2 matrix with a two element vector """
    return [m[0][0] * v[0] + m[0][1] * v[1], m[1][0] * v[0] + m[1][1] * v[1]]


def fitStraightLine(xs, ys, sigmas):
    """ fit a straight line to given data points """
    mat = [[0, 0], [0, 0]]
    rhs = [0, 0]
    for x, y, sigma in zip(xs, ys, sigmas):
        w = 1 / (sigma * sigma)
        rhs[0] += w * y * x
        rhs[1] += w * y
        mat[0][0] += w * x * x
        mat[0][1] += w * x
        mat[1][1] += w
    mat[1][0] = mat[0][1]
    cov = invert2x2(mat)
    sol = mulM2x2V(cov, rhs)
    return sol, cov


Vtrue = [0, 1, 2, 5]
V = [.09, 1.05, 1.72, 4.15]
sigmas = [.1, .1, .1, .1]

sol, cov = fitStraightLine(Vtrue, V, sigmas)

print("Solution: fit parameters a={} +/- {}".format(sol[0],
                                                    math.sqrt(cov[0][0])))
print("Solution: fit parameters b={} +/- {}".format(sol[1],
                                                    math.sqrt(cov[1][1])))
print("Correlation matrix:")
print("[ {}, {} ]".format(1, cov[0][1] / math.sqrt(cov[0][0]*cov[1][1])))
print("[ {}, {} ]".format(cov[1][0] / math.sqrt(cov[0][0]*cov[1][1]), 1))