################################################################################
#                                                                              #
#                           UNCERTANTIES WRAPPER                               #
#                                                                              #
################################################################################

from uncertainties import *
import numpy as np
from uncertainties import unumpy as unp

__all__ = []

### funcs



def numericJacobian(f, x, vals, f_size = 1):
  J = np.zeros([len(x), f_size, len(vals)])
  for l in range(0,len(vals)):
    if vals[l]!= 0:    h = np.sqrt(np.finfo(float).eps)*vals[l];
    else:           h = 1e-14;
    vals1 = np.copy(vals); vals1[l] += +h
    vals2 = np.copy(vals); vals2[l] += -h;
    f1 = f(x,*vals1).astype(np.float64)
    f2 = f(x,*vals2).astype(np.float64)
    thisJ = ((f(x,*vals1) - f(x,*vals2))/(2*h)).astype(np.float64)
    J[:,0,l] = thisJ # nowadays only scalar f allowed
  return J.T

def propagate_term(der,unc):
  return der**2*unc**2

def wrap_unc(f,x,*args):

  # get parameters and uncertainties
  vals = np.array([args[k].nominal_value  for k in range(0,len(args))])
  uncs = np.array([args[k].std_dev        for k in range(0,len(args))])

  # compute f nominal_value
  f_val = f(x,*vals)
  if hasattr(f(np.array([1]),*vals), "__len__"):
    f_size = len(f(np.array([1]),*vals))
  else:
    f_size = 1

  # get numeric derivatives
  derivatives = numericJacobian(f, x, vals, f_size)
  f_unc = np.zeros([len(x)])

  # compute f std_dev
  for i in range(0,len(uncs)):
    f_unc[:] += propagate_term(derivatives[i],uncs[i])[0]
  f_unc = np.sqrt(f_unc)

  return unp.uarray(f_val,f_unc)

def get_confidence_bands(x,y,sigma=1):
  nom = unp.nominal_values(y)
  std = unp.std_devs(y)
  # uncertainty lines (sigma confidence)
  return nom+sigma*std, nom-sigma*std
