#!/usr/bin/env python
# -*- coding: utf-8 -*-
# written by Christoph Federrath

import argparse
from cfpack.defaults import * # define global constants, print, plotting styles, etc.
import numpy as np
from mytools import get_pdf

# ===== the following applies in case we are running this in script mode =====
if __name__ == "__main__":

    # create new parser and add arguments
    parser = argparse.ArgumentParser(description='Demo for MC error propagation.')
    args = parser.parse_args()

    # define measurement variables a and b
    a_mean = 10.0; a_std = 1.0
    # a_mean = 1.0; a_std = 0.5 # 2nd example for a
    b_mean = 1.0; b_std = 0.1

    # compute derived variable c based on best values and their errors
    c_mean = a_mean**2 / b_mean**2
    c_std = c_mean * np.sqrt( (2*a_std/a_mean)**2 + (2*b_std/b_mean)**2 )
    print("Based on normal error propagation, we find c = ", c_mean, "+/-", c_std)

    # === now do Monte Carlo error propagation ===
    n = int(1e6)

    # create Gaussian distributed measurement variables a and b
    a = np.random.normal(loc=a_mean, scale=a_std, size=n)
    b = np.random.normal(loc=b_mean, scale=b_std, size=n)

    # plot PDFs of a and b
    a_pdf_obj = get_pdf(a, xlabel=r"$a$", ylog=False, save="PDF_a.pdf")
    b_pdf_obj = get_pdf(b, xlabel=r"$b$", ylog=False, save="PDF_b.pdf")

    # compute derived quantity c and its PDF
    c = a**2 / b**2
    c_pdf_obj = get_pdf(c, xlabel=r"$c$", xlog=True, ylog=False, bins=2000, save="PDF_c.pdf")
    print("Based on MC error propagation (data), we find c (mean, std) = ",
          c.mean(), "+/-", c.std())
    print("Based on MC error propagation (PDF), we find c (mean, std)= ",
          c_pdf_obj.mean, "+/-", c_pdf_obj.std)
    print("Based on MC error propagation (PDF), we find c (median, 16th-84th) = ",
          c_pdf_obj.median,
          "-",c_pdf_obj.median-c_pdf_obj.sixteenth,
          "+", c_pdf_obj.eightyfourth-c_pdf_obj.median)

    stop()