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

import math
import ROOT

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

fitparams = [0.8014, 0.1496]
fitcov = [[0.02673 * 0.02673, -0.7303 * 0.02673 * 0.07319],
          [-0.7303 * 0.02673 * 0.07319, 0.07319 * 0.07319]]

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

# build graph for fit result (and uncertainty)
gfit = ROOT.TGraphErrors(100)
gfit.SetNameTitle("gfit", "fit result;V_{true} [V];V_{meas} [V]")
for i in range(0, 121):
    x = Vtrue[3] * (i - 10) / 100.  # range to extend a bit beyond data points
    gfit.SetPoint(i, x, fitparams[0] * x + fitparams[1])
    tmp = [fitcov[0][0] * x + fitcov[0][1], fitcov[1][0] * x + fitcov[1][1]]
    err2 = tmp[0] * x + tmp[1]
    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(Vtrue))
gdata.SetNameTitle("gdata", "data")
for i in range(0, len(Vtrue)):
    gdata.SetPoint(i, Vtrue[i], V[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("day1-ex3.pdf")
c.SaveAs("day1-ex3.png")