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

__all__ = ['VariableSoftmax', 'AbsInferno', 'PaperInferno', 'InfernoPred', 'AbsApproxInferno', 'ApproxPaperInferno']

# Cell
from .callback import AbsCallback, PredHandler
from .inference import calc_grad_hesse

import numpy as np
from abc import abstractmethod
from fastcore.all import store_attr, delegates, is_close
from typing import Optional, List

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

# Cell
class VariableSoftmax(nn.Softmax):
    r'''Softmax with temperature'''
    def __init__(self, temp:float=1, dim:int=-1):
        super().__init__(dim=dim)
        self.temp = temp

    def forward(self, x:Tensor) -> Tensor: return super().forward(x/self.temp)

# Cell
class AbsInferno(AbsCallback):
    r'''Attempted reproduction of TF1 & TF2 INFERNO with exact effect of nuisances being passed through model'''
    def __init__(self, b_true:float, mu_true:float, n_shape_alphas:int=0, s_shape_alpha:bool=False, b_shape_alpha:bool=False, nonaux_b_norm:bool=False,
                 shape_aux:Optional[List[Distribution]]=None, b_norm_aux:Optional[List[Distribution]]=None, s_norm_aux:Optional[List[Distribution]]=None):
        super().__init__()
        store_attr()
        if self.shape_aux is not None and len(self.shape_aux) != self.n_shape_alphas: raise ValueError("Number of auxillary measurements on shape nuisances must match the number of shape nuisance parameters")
        self.n=self.mu_true+self.b_true
        # Norm uncertainties
        if self.s_norm_aux is None: self.s_norm_aux = []
        if self.b_norm_aux is None: self.b_norm_aux = []
        # Compute nuisance indeces
        self.poi_idx = [0]
        self.shape_idxs = list(range(1,self.n_shape_alphas+1))
        self.n_alpha = 1+self.n_shape_alphas
        self.s_norm_idxs = list(range(self.n_alpha, self.n_alpha+len(self.s_norm_aux)))
        self.n_alpha += len(self.s_norm_aux)
        self.b_norm_idxs = list(range(self.n_alpha, self.n_alpha+len(self.b_norm_aux)+self.nonaux_b_norm))
        self.n_alpha += len(self.b_norm_aux)+self.nonaux_b_norm

    def on_train_begin(self) -> None:
        self.wrapper.loss_func = None  # Ensure loss function is skipped, callback computes loss value in `on_forwards_end`
        for c in self.wrapper.cbs:
            if hasattr(c, 'loss_is_meaned'): c.loss_is_meaned = False  # Ensure that average losses are correct
        self.alpha = torch.zeros((self.n_alpha), requires_grad=True, device=self.wrapper.device)  # Nuisances set to zero (true values)
        with torch.no_grad(): self.alpha[self.poi_idx] = self.mu_true  # POI set to true value

    def on_batch_begin(self) -> None:
        self.b_mask = self.wrapper.y.squeeze() == 0
        self._aug_data(self.wrapper.x)

    def on_batch_end(self) -> None:
        self.alpha.grad.data.zero_()

    @abstractmethod
    def _aug_data(self, x:Tensor) -> Tensor:
        r'''Include nuisances in input data. Overide this for specific problem.'''
        pass

    def get_inv_ikk(self, f_s:Tensor, f_b:Tensor, f_s_asimov:Tensor, f_b_asimov:Tensor) -> Tensor:
        r'''Compute full hessian at true param values'''
        # Compute nll
        s_exp = self.alpha[self.poi_idx]+self.alpha[self.s_norm_idxs].sum() if len(self.s_norm_idxs) > 0 else self.alpha[self.poi_idx]
        b_exp = self.b_true             +self.alpha[self.b_norm_idxs].sum() if len(self.b_norm_idxs) > 0 else self.b_true
        t_exp  = (s_exp*f_s)+(b_exp*f_b)
        asimov = (self.mu_true*f_s)+(self.b_true*f_b_asimov)
        nll = -torch.distributions.Poisson(t_exp, False).log_prob(asimov).sum()
        # Constrain nll
        if self.shape_aux is not None:
            for a,x in zip(self.alpha[self.shape_idxs], self.shape_aux):
                if x is not None: nll = nll-x.log_prob(a)
        if len(self.b_norm_idxs):
            for i,x in zip(self.b_norm_idxs, self.b_norm_aux): nll = nll-x.log_prob(self.alpha[i])
        if len(self.s_norm_idxs):
            for i,x in zip(self.s_norm_idxs, self.s_norm_aux): nll = nll-x.log_prob(self.alpha[i])
        _,h = calc_grad_hesse(nll, self.alpha, create_graph=True)
        return torch.inverse(h)[self.poi_idx,self.poi_idx]

    @staticmethod
    def to_shape(p:Tensor, w:Optional[Tensor]=None) -> Tensor:
        f = (p*w).sum(0)+1e-7 if w is not None else p.sum(0)+1e-7
        return f/f.sum()

    def on_forwards_end(self) -> None:
        r'''Compute loss and replace wrapper loss value'''

        w_s = self.wrapper.w[~self.b_mask] if self.wrapper.w is not None else None
        w_b = self.wrapper.w[self.b_mask] if self.wrapper.w is not None else None

        # Shapes with derivatives w.r.t. nuisances
        f_s = self.to_shape(self.wrapper.y_pred[~self.b_mask], w_s)
        f_b = self.to_shape(self.wrapper.y_pred[self.b_mask], w_b)

        # Shapes without derivatives w.r.t. nuisances
        f_s_asimov = self.to_shape(self.wrapper.model(self.wrapper.x[~self.b_mask].detach()), w_s) if self.s_shape_alpha else f_s
        f_b_asimov = self.to_shape(self.wrapper.model(self.wrapper.x[self.b_mask].detach()), w_b)  if self.b_shape_alpha else f_b

        self.wrapper.loss_val = self.get_inv_ikk(f_s=f_s, f_b=f_b, f_s_asimov=f_s_asimov, f_b_asimov=f_b_asimov)

# Cell
class PaperInferno(AbsInferno):
    r'''Inheriting class for dealing with INFERNO paper synthetic problem'''
    @delegates(AbsInferno, but=['b_shape_alpha', 's_shape_alpha'])
    def __init__(self, float_r:bool, float_l:bool, l_init:float=3, b_true:float=1000, mu_true:float=50, **kwargs):
        super().__init__(b_true=b_true, mu_true=mu_true, n_shape_alphas=float_r+float_l, b_shape_alpha=True, **kwargs)
        self.float_r,self.float_l,self.l_init = float_r,float_l,l_init

    def _aug_data(self, x:Tensor) -> None:
        if self.float_r: x[self.b_mask,0] += self.alpha[self.shape_idxs[0]]
        if self.float_l: x[self.b_mask,2] *= (self.alpha[self.shape_idxs[-1]]+self.l_init)/self.l_init

# Cell
class InfernoPred(PredHandler):
    r'''Prediction handler for hard assignments'''
    def get_preds(self) -> np.ndarray: return np.argmax(self.preds, 1)

# Cell
from .inference import calc_nll

from fastcore.all import partialler
from typing import Tuple

# Cell
class AbsApproxInferno(AbsInferno):
    r'''Attempted reproduction INFERNO following paper description implementations with nuisances being approximated by creating up/down shapes and interpolating
    Includes option to randomise params per batch and converge to better values, which results in slightly better performance'''
    @delegates(AbsInferno)
    def __init__(self, aug_alpha:bool=False, n_steps:int=100, lr:float=0.1, **kwargs):
        super().__init__(**kwargs)
        store_attr('aug_alpha, n_steps, lr')

    def _aug_data(self): pass  # Override abs method
    def on_batch_begin(self) -> None: pass
    def on_batch_end(self) -> None: pass

    def on_train_begin(self) -> None:
        self.wrapper.loss_func = None  # Ensure loss function is skipped, callback computes loss value in `on_forwards_end`
        for c in self.wrapper.cbs:
            if hasattr(c, 'loss_is_meaned'): c.loss_is_meaned = False  # Ensure that average losses are correct

    @abstractmethod
    def _get_up_down(self, x_s:Tensor, x_b:Tensor, w_s:Optional[Tensor]=None, w_b:Optional[Tensor]=None) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:
        r'''Compute upd/down shapes for signal and background seperately. Overide this for specific problem.'''
        pass

    def get_ikk(self, f_s_nom:Tensor, f_b_nom:Tensor, f_s_up:Optional[Tensor], f_s_dw:Optional[Tensor], f_b_up:Optional[Tensor], f_b_dw:Optional[Tensor]) -> Tensor:
        r'''Compute full hessian at true param values, or at random starting values with Newton updates'''
        if self.aug_alpha: alpha = torch.randn((self.n_alpha), requires_grad=True, device=self.wrapper.device)/10
        else:              alpha = torch.zeros((self.n_alpha), requires_grad=True, device=self.wrapper.device)
        with torch.no_grad(): alpha[self.poi_idx] += self.mu_true
        get_nll = partialler(calc_nll, s_true=self.mu_true, b_true=self.b_true,
                             f_s_nom=f_s_nom, f_b_nom=f_b_nom, f_s_up=f_s_up, f_s_dw=f_s_dw,
                             f_b_up=f_b_up, f_b_dw=f_b_dw, shape_aux=self.shape_aux, s_norm_aux=self.s_norm_aux, b_norm_aux=self.b_norm_aux)
        if self.aug_alpha:  # Alphas carry noise, optimise via Newton
            for i in range(self.n_steps):  # Newton optimise nuisances & mu
                nll = get_nll(mu=alpha[self.poi_idx], s_norm_alpha=alpha[self.s_norm_idxs], b_norm_alpha=alpha[self.b_norm_idxs], shape_alpha=alpha[self.shape_idxs])
                g,h = calc_grad_hesse(nll, alpha)
                s = torch.clamp(self.lr*g.detach()@torch.inverse(h), -100, 100)
                alpha = alpha-s
        nll = get_nll(mu=alpha[self.poi_idx], s_norm_alpha=alpha[self.s_norm_idxs], b_norm_alpha=alpha[self.b_norm_idxs], shape_alpha=alpha[self.shape_idxs])
        _,h = calc_grad_hesse(nll, alpha, create_graph=True)
        return torch.inverse(h)[self.poi_idx,self.poi_idx]

    def on_forwards_end(self) -> None:
        r'''Compute loss and replace wrapper loss value'''
        b = self.wrapper.y.squeeze() == 0
        w_s = self.wrapper.w[~b] if self.wrapper.w is not None else None
        w_b = self.wrapper.w[b] if self.wrapper.w is not None else None
        f_s = self.to_shape(self.wrapper.y_pred[~b], w_s)
        f_b = self.to_shape(self.wrapper.y_pred[b], w_b)
        (f_s_up,f_s_dw),(f_b_up,f_b_dw)= self._get_up_down(self.wrapper.x[~b], self.wrapper.x[b])
        self.wrapper.loss_val = self.get_ikk(f_s_nom=f_s, f_b_nom=f_b, f_s_up=f_s_up, f_s_dw=f_s_dw, f_b_up=f_b_up, f_b_dw=f_b_dw)

# Cell
class ApproxPaperInferno(AbsApproxInferno):
    r'''Inheriting class for dealing with INFERNO paper synthetic problem'''
    @delegates(AbsApproxInferno, but=['b_shape_alpha', 's_shape_alpha'])
    def __init__(self, r_mods:Optional[Tuple[float,float]]=(-0.2,0.2), l_mods:Optional[Tuple[float,float]]=(2.5,3.5), l_init:float=3,
                 b_true:float=1000, mu_true:float=50, **kwargs):
        super().__init__(b_true=b_true, mu_true=mu_true, n_shape_alphas=(r_mods is not None)+(l_mods is not None), b_shape_alpha=True, **kwargs)
        store_attr('r_mods, l_mods, l_init')

    def on_train_begin(self) -> None:
        super().on_train_begin()
        if self.r_mods is not None:
            self.r_mod_t = (torch.zeros(1,3, device=self.wrapper.device),torch.zeros(1,3, device=self.wrapper.device))
            self.r_mod_t[0][0,0] = self.r_mods[0]
            self.r_mod_t[1][0,0] = self.r_mods[1]
        if self.l_mods is not None:
            self.l_mod_t = (torch.ones(1,3, device=self.wrapper.device),torch.ones(1,3, device=self.wrapper.device))
            self.l_mod_t[0][0,2] = self.l_mods[0]/self.l_init
            self.l_mod_t[1][0,2] = self.l_mods[1]/self.l_init

    def _get_up_down(self, x_s:Tensor, x_b:Tensor, **kwargs) -> Tuple[Tuple[Optional[Tensor],Optional[Tensor]],Tuple[Optional[Tensor],Optional[Tensor]]]:
        if self.r_mods is None and self.l_mods is None: return (None,None),(None,None)
        u,d = [],[]
        if self.r_mods is not None:
            with torch.no_grad(): x_b = x_b+self.r_mod_t[0]
            d.append(self.to_shape(self.wrapper.model(x_b)))
            with torch.no_grad(): x_b = x_b+self.r_mod_t[1]-self.r_mod_t[0]
            u.append(self.to_shape(self.wrapper.model(x_b)))
            with torch.no_grad(): x_b = x_b-self.r_mod_t[1]
        if self.l_mods is not None:
            with torch.no_grad(): x_b = x_b*self.l_mod_t[0]
            d.append(self.to_shape(self.wrapper.model(x_b)))
            with torch.no_grad(): x_b = x_b*self.l_mod_t[1]/self.l_mod_t[0]
            u.append(self.to_shape(self.wrapper.model(x_b)))
            with torch.no_grad(): x_b = x_b/self.l_mod_t[1]
        return (None,None),(torch.stack(u),torch.stack(d))