#!/usr/bin/env python3
#
# day2-ex2.py - build linear system for parabola 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 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


def calcChi2(xs, ys, sigmas, x0, param):
    retVal = 0
    for x, y, s in zip(xs, ys, sigmas):
        w = 1 / (s * s)
        dx = x - x0
        ypred = param[0] + dx * (param[1] + dx * param[2])
        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 runParabolicFit(xs, ys, sigmays, x0):
    rhs, mat = buildLinSystemForParabolaFit(xs, ys, sigmays, x0)
    
    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())

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

    return crhs, MatAdapter(cmat), chi2, ndf


def generateMeasurements(param_true, n, xlo, xhi, x0, sigmay, rng):
    xs = [xlo * i / (n - 1) + xhi * (n - 1 - i) / (n - 1) for i in range(0, n)]
    ys = []
    sigmas = []
    for x in xs:
        dx = x - x0
        ytrue = param_true[0] + dx * (param_true[1] + dx * param_true[2])
        y = rng.Gaus(ytrue, sigmay)
        ys.append(y)
        sigmas.append(sigmay)

    return xs, ys, sigmas

def runToys(ntoys, param_true, nmeas, xlo, xhi, x0, sigmay, rng):
    toys = []
    for i in range(0, ntoys):
        xs, ys, sigmas = generateMeasurements(param_true, nmeas, xlo, xhi, x0,
                                              sigmay, rng)
        rhs, cov, chi2, ndf = runParabolicFit(xs, ys, sigmas, x0)
        toys.append({'chi2' : chi2, 'ndf': ndf, 'param': rhs, 'param_true':
                     param_true, 'cov': cov})
    return toys


rng = ROOT.TRandom()
toys = runToys(10000, [2.4, 0.9, 0.07], 6, 0, 5, 2.5, 0.1, rng)

histos = []

hchi2ndf = ROOT.TH1D("hchi2ndf", "#chi^2/ndf;#chi^2/ndf;entries", 100, 0., 10.)
histos.append(hchi2ndf)
hchi2prob = ROOT.TH1D("hchi2prob", "#chi^2 probability;Prob(#chi^2,ndf);entries", 100, 0., 1.)
histos.append(hchi2prob)
hp0 = ROOT.TH1D("hp0", "p_{0};p_{0};entries", 100, 2., 3.)
histos.append(hp0)
hp1 = ROOT.TH1D("hp1", "p_{1};p_{1};entries", 100, 0.7, 1.1)
histos.append(hp1)
hp2 = ROOT.TH1D("hp2", "p_{2};p_{2};entries", 100, 0., 0.15)
histos.append(hp2)
hp0resid = ROOT.TH1D("hp0resid", "p_{0}-p_{0}^{true};p_{0}-p_{0}^{true};entries", 100,
                     -.3, .3)
histos.append(hp0resid)
hp1resid = ROOT.TH1D("hp1resid", "p_{1}-p_{1}^{true};p_{1}-p_{1}^{true};entries", 100,
                     -.1, .1)
histos.append(hp1resid)
hp2resid = ROOT.TH1D("hp2resid", "p_{2}-p_{2}^{true};p_{2}-p_{2}^{true};entries", 100,
                     -.075, .075)
histos.append(hp2resid)
hp0pull = ROOT.TH1D("hp0pull", "p_{0} pull;#frac{p_{0}-p_{0}^{true}}{#sigma_{0}};entries",
                    100, -5, 5.)
histos.append(hp0pull)
hp1pull = ROOT.TH1D("hp1pull", "p_{1} pull;#frac{p_{1}-p_{1}^{true}}{#sigma_{1}};entries",
                    100, -5, 5.)
histos.append(hp1pull)
hp2pull = ROOT.TH1D("hp2pull", "p_{2} pull;#frac{p_{2}-p_{2}^{true}}{#sigma_{2}};entries",
                    100, -5, 5.)
histos.append(hp2pull)

for toy in toys:
    hchi2ndf.Fill(toy['chi2'] / toy['ndf'])
    hchi2prob.Fill(ROOT.TMath.Prob(toy['chi2'], toy['ndf']))
    hp0.Fill(toy['param'][0])
    hp1.Fill(toy['param'][1])
    hp2.Fill(toy['param'][2])
    hp0resid.Fill(toy['param'][0] - toy['param_true'][0])
    hp1resid.Fill(toy['param'][1] - toy['param_true'][1])
    hp2resid.Fill(toy['param'][2] - toy['param_true'][2])
    hp0pull.Fill((toy['param'][0] - toy['param_true'][0]) /
                 math.sqrt(toy['cov'].at(0, 0)))
    hp1pull.Fill((toy['param'][1] - toy['param_true'][1]) /
                 math.sqrt(toy['cov'].at(1, 1)))
    hp2pull.Fill((toy['param'][2] - toy['param_true'][2]) /
                 math.sqrt(toy['cov'].at(2, 2)))

c = ROOT.TCanvas("c", "c", 800, 600)
for h in histos:
    h.SetMinimum(0)
    h.Draw()
    c.SaveAs("day3-ex1-{}.pdf".format(h.GetName()))
    c.SaveAs("day3-ex1-{}.png".format(h.GetName()))