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

__all__ = ['PseudoData', 'paper_sig', 'paper_bkg']

# Cell
from fastcore.all import store_attr
from typing import Callable, Tuple, Union, List
import numpy as np
import pandas as pd

# Cell
class PseudoData():
    r'''Generic class for constructing pseudodata via sampling of random functions'''
    def __init__(self, func:Callable[[int],np.ndarray], targ:Union[float,int]): store_attr()
    def __getitem__(self, i:int) -> Tuple[np.ndarray,np.ndarray]: self.sample(1)
    def sample(self, n:int) -> Tuple[np.ndarray,np.ndarray,None]: return self.func(n).astype('float32'),self.targ*np.ones((n,1)).astype('float32'),None
    def __call__(self, n:int) -> Tuple[np.ndarray,np.ndarray,None]: return self.sample(n)
    def get_df(self, n:int) -> pd.DataFrame:
        d = self.sample(n)
        if d[2] is None: d = d[:2]
        df = pd.DataFrame(np.hstack(d))
        df.rename(columns={df.columns[-1]:'gen_target'}, inplace=True)
        return df

# Cell
class _PaperData():
    r'''Callable class generating pseudodata from Inferno paper'''
    def __init__(self, mu:List[float], conv:List[List[float]], r:float, l:float):
        store_attr(but=['mu', 'r'])
        self.r = np.array([r,0])
        self.mu = np.array(mu)

    def sample(self, n:int) -> np.ndarray:
        return np.hstack((np.random.multivariate_normal(self.mu+self.r, self.conv, n),
                          np.random.exponential(1/self.l, size=n)[:,None]))

    def __call__(self, n:int) -> np.ndarray: return self.sample(n)

# Cell
paper_sig = _PaperData(mu=[0,0], conv=[[1,0],[0,1]], r=0, l=2)

# Cell
paper_bkg = _PaperData(mu=[2,0], conv=[[5,0],[0,9]], r=0, l=3)