# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/06_inference.ipynb (unless otherwise specified).

__all__ = ['bin_preds', 'get_shape', 'get_paper_syst_shapes', 'get_likelihood_width', 'interp_shape', 'calc_nll',
           'jacobian', 'calc_grad_hesse', 'calc_profile']

# Cell
from .model_wrapper import ModelWrapper
from .callback import PaperSystMod, PredHandler

import pandas as pd
import numpy as np
from typing import *
from collections import OrderedDict
from scipy.interpolate import InterpolatedUnivariateSpline
import itertools
from fastcore.all import partialler
from fastprogress import progress_bar
import math

from torch import Tensor, autograd
import torch
from torch.distributions import Distribution

# Cell
def bin_preds(df:pd.DataFrame, bins:np.ndarray=np.linspace(0.,10.,11), pred_name='pred') -> None:
    '''Bins predictions over specified range'''
    df[f'{pred_name}_bin'] = np.digitize(df[pred_name], bins)-1

# Cell
def get_shape(df:pd.DataFrame, targ:int, bins:np.ndarray=np.linspace(0.,10.,11), pred_name:str='pred_bin') -> Tensor:
    r'''Extracts normalised shape of class from binned predictions. Empty bins are filled with a small quantity to avoid zeros.'''
    f = df.loc[df.gen_target == targ, pred_name].value_counts(bins=bins-(1/len(bins)))
    f.sort_index(inplace=True)
    f += 1e-7
    f /= f.sum()
    return Tensor(f.values)

# Cell
def get_paper_syst_shapes(bkg_data:np.ndarray, df:pd.DataFrame, model:ModelWrapper, bins:np.ndarray=np.linspace(0.,10.,11), pred_cb:PredHandler=PredHandler(),
                          r_vals:Tuple[float,float,float]=[-0.2,0,0.2], l_vals:Tuple[float]=[2.5,3,3.5]) -> OrderedDict:
    r'''Pass background data through trained model in order to get up/down shape variations.'''
    def _get_shape(r,l):
        bp = model.predict(bkg_data, pred_cb=pred_cb, cbs=PaperSystMod(r=r,l=l))
        n = f'pred_{r}_{l}'
        df[n] = df.pred
        df.loc[df.gen_target == 0, n] = bp
        bin_preds(df, pred_name=n, bins=bins)
        return get_shape(df=df, targ=0, bins=np.linspace(0.,len(bins)-1,len(bins)), pred_name=f'{n}_bin')

    shapes = OrderedDict()
    for i,r in enumerate(r_vals):
        print(f'Running: r={r}')
        shapes[f'{i}_{1}'] = _get_shape(r,l_vals[1])
    for i,l in enumerate(l_vals):
        print(f'Running: l={l}')
        shapes[f'{1}_{i}'] = _get_shape(r_vals[1],l)
    return OrderedDict((('f_b_nom',shapes['1_1']),
                        ('f_b_up', torch.stack((shapes['2_1'],shapes['1_2']))),
                        ('f_b_dw', torch.stack((shapes['0_1'],shapes['1_0'])))))

# Cell
def get_likelihood_width(nll:np.ndarray, mu_scan:np.ndarray, val:float=0.5) -> float:
    r'''Compute width of likelihood at 95% confidence-level'''
    m = nll == nll
    r = InterpolatedUnivariateSpline(mu_scan[m], nll[m]-val-nll[m].min()).roots()
    if len(r) != 2: raise ValueError(f'No roots found at {val}, set val to a smaller value.')
    return (r[1]-r[0])/2

# Cell
def interp_shape(alpha:Tensor, f_nom:Tensor, f_up:Tensor, f_dw:Tensor):
    r'''Use quadratic interpolation between up/down systematic shapes and nominal in order to estimate shapes at arbitrary nuisance values.
    Linear extrapolation for absolute nuisances values greater than 1 (outside up/down shape range).
    Does not account for co-dependence of nuisances.
    Adapted from https://github.com/pablodecm/paper-inferno/blob/master/code/template_model.py under BSD 3-clause licence Copyright (c) 2018, Pablo de Castro, Tommaso Dorigo'''
    alpha_t = torch.repeat_interleave(alpha.unsqueeze(-1), repeats=f_nom.shape[-1], dim=-1)
    a = 0.5*(f_up+f_dw)[None,:]-f_nom
    b = 0.5*(f_up-f_dw)[None,:]

    switch = torch.where(alpha_t < 0., f_dw-f_nom, f_up-f_nom)
    abs_var = torch.where(torch.abs(alpha_t) > 1.,
                          (b+(torch.sign(alpha_t)*a))*(alpha_t-torch.sign(alpha_t))+switch,
                          a*torch.pow(alpha_t, 2)+ b * alpha_t)
    return (f_nom + abs_var.sum(1, keepdim=True)).squeeze(1)

# Cell
def calc_nll(s_true:float, b_true:float, mu:Tensor, f_s_nom:Tensor, f_b_nom:Tensor,
             shape_alpha:Optional[Tensor]=None, s_norm_alpha:Optional[Tensor]=None, b_norm_alpha:Optional[Tensor]=None,
             f_s_up:Optional[Tensor]=None, f_s_dw:Optional[Tensor]=None,
             f_b_up:Optional[Tensor]=None, f_b_dw:Optional[Tensor]=None,
             s_norm_aux:Optional[Distribution]=None, b_norm_aux:Optional[Distribution]=None, shape_aux:Optional[List[Distribution]]=None) -> Tensor:
    r'''Compute negative log-likelihood for specified parameters.'''
    #  Adjust expectation by nuisances
    f_s = interp_shape(shape_alpha, f_s_nom, f_s_up, f_s_dw) if shape_alpha is not None and f_s_up is not None else f_s_nom
    f_b = interp_shape(shape_alpha, f_b_nom, f_b_up, f_b_dw) if shape_alpha is not None and f_b_up is not None  else f_b_nom
    s_exp = mu    +s_norm_alpha.sum() if s_norm_alpha is not None else mu
    b_exp = b_true+b_norm_alpha.sum() if b_norm_alpha is not None else b_true
    #  Compute NLL
    t_exp = (s_exp*f_s)+(b_exp*f_b)
    asimov = (s_true*f_s_nom)+(b_true*f_b_nom)
    nll = -torch.distributions.Poisson(t_exp, False).log_prob(asimov).sum()
    # Constrain nuisances
    if shape_aux is not None:
        if len(shape_aux) != len(shape_alpha): raise ValueError("Number of auxillary measurements must match the number of nuisance parameters.\
                                                           Pass `None`s for unconstrained nuisances.")
        for a,x in zip(shape_alpha, shape_aux):
            if x is not None: nll = nll-x.log_prob(a)
    if b_norm_alpha is not None:
        for a,x in zip(b_norm_alpha, b_norm_aux): nll = nll-x.log_prob(a)
    if s_norm_alpha is not None:
        for a,x in zip(s_norm_alpha, s_norm_aux): nll = nll-x.log_prob(a)
    return nll

# Cell
def jacobian(y:Tensor, x:Tensor, create_graph=False):
    r'''Compute full jacobian matrix for single tensor. Call twice for hessian.
    Copied from https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7 credits: Adam Paszke
    TODO: Fix this to work batch-wise (maybe https://gist.github.com/sbarratt/37356c46ad1350d4c30aefbd488a4faa)'''
    jac = []
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    for i in range(len(flat_y)):
        grad_y[i] = 1.
        grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
        jac.append(grad_x.reshape(x.shape))
        grad_y[i] = 0.
    return torch.stack(jac).reshape(y.shape + x.shape)

# Cell
def calc_grad_hesse(nll:Tensor, alpha:Tensor, create_graph:bool=False) -> Tuple[Tensor,Tensor]:
    r'''Compute full hessian and jacobian for single tensor'''
    grad = jacobian(nll, alpha, create_graph=True)
    hesse = jacobian(grad, alpha, create_graph=create_graph)
    return grad, hesse

# Cell
def calc_profile(f_s_nom:Tensor, f_b_nom:Tensor, n_obs:int, mu_scan:Tensor, mu_true:int,
                 f_s_up:Optional[Tensor]=None, f_s_dw:Optional[Tensor]=None,
                 f_b_up:Optional[Tensor]=None, f_b_dw:Optional[Tensor]=None,
                 shape_aux:Optional[List[Distribution]]=None,
                 s_norm_aux:Optional[List[Distribution]]=None, b_norm_aux:Optional[List[Distribution]]=None, nonaux_b_norm:bool=False,
                 n_steps:int=100, lr:float=0.1,  verbose:bool=True) -> Tensor:
    r'''Compute profile likelihoods for range of mu values, optimising on full hessian.
    Ideally mu-values should be computed in parallel, but batch-wise hessian in PyTorch is difficult.'''
    for f in [f_s_nom, f_s_up, f_s_dw, f_b_nom, f_b_up, f_b_dw]:  # Ensure correct dimensions
        if f is not None and len(f.shape) < 2: f.unsqueeze_(0)
    # Cases where nuisance only causes up xor down variation
    if (f_s_up is None and f_s_dw is not None): f_s_up = torch.repeat_interleave(f_s_nom, repeats=len(f_s_dw), dim=0)
    if (f_s_dw is None and f_s_up is not None): f_s_dw = torch.repeat_interleave(f_s_nom, repeats=len(f_s_up), dim=0)
    if (f_b_up is None and f_b_dw is not None): f_b_up = torch.repeat_interleave(f_s_nom, repeats=len(f_b_dw), dim=0)
    if (f_b_dw is None and f_b_up is not None): f_b_dw = torch.repeat_interleave(f_s_nom, repeats=len(f_b_up), dim=0)
    if f_s_up is not None and f_b_up is not None and len(f_s_up) != len(f_b_up):
        raise ValueError("Shape variations for signal & background must have the same number of variations. \
                          Please enter the nominal templates for nuisances that only affect either signal of background.")
    # Norm uncertainties
    if s_norm_aux is None: s_norm_aux = []
    if b_norm_aux is None: b_norm_aux = []
    # Compute nuisance indeces
    n_alpha = np.max((len(f_b_up) if f_b_up is not None else 0, len(f_s_up) if f_s_up is not None else 0))
    shape_idxs = list(range(n_alpha))
    s_norm_idxs = list(range(n_alpha, n_alpha+len(s_norm_aux)))
    n_alpha += len(s_norm_aux)
    b_norm_idxs = list(range(n_alpha, n_alpha+len(b_norm_aux)+nonaux_b_norm))
    n_alpha += len(b_norm_aux)+nonaux_b_norm

    b_true = n_obs-mu_true
    if n_alpha > 0:
        nlls = []
        get_nll = partialler(calc_nll, s_true=mu_true, b_true=b_true,
                         f_s_nom=f_s_nom, f_s_up=f_s_up, f_s_dw=f_s_dw,
                         f_b_nom=f_b_nom, f_b_up=f_b_up, f_b_dw=f_b_dw,
                         s_norm_aux=s_norm_aux, b_norm_aux=b_norm_aux, shape_aux=shape_aux)
        for mu in progress_bar(mu_scan, display=verbose):  # TODO: Fix this to run mu-scan in parallel
                alpha = torch.zeros((n_alpha), requires_grad=True, device=f_b_nom.device)
                for i in range(n_steps):  # Newton optimise nuisances
                    nll = get_nll(shape_alpha=alpha[shape_idxs], mu=mu, s_norm_alpha=alpha[s_norm_idxs], b_norm_alpha=alpha[b_norm_idxs])
                    grad, hesse = calc_grad_hesse(nll, alpha, create_graph=False)
                    step = lr*grad.detach()@torch.inverse(hesse)
                    step = torch.clamp(step, -100, 100)
                    alpha = alpha-step
                nlls.append(get_nll(shape_alpha=alpha[shape_idxs], mu=mu, s_norm_alpha=alpha[s_norm_idxs], b_norm_alpha=alpha[b_norm_idxs]).detach())
                if alpha[shape_idxs].abs().max() > 1: print(f'Linear regime: Mu {mu.data.item()}, shape nuisances {alpha[shape_idxs].data}')
        nlls = torch.stack(nlls)
    else:
        nlls = -torch.distributions.Poisson((mu_scan.reshape((-1,1))*f_s_nom)+(b_true*f_b_nom), False).log_prob((mu_true*f_s_nom)+(b_true*f_b_nom)).sum(1)
    return nlls