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

__all__ = ['generate_times', 'sec_per_day', 'WeightFunction', 'make_exposure', 'Simulation']

# Cell
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats

from .config import Config
from .weights import PointSource

# Cell
import numbers

class _Sampler():
    """ Sample an arbitrary function or histogram

    - func -- a function, a histogram, or a fixed value<br>
        If a function, must be positive definite.<br>
        Assume histogram bins are 0 to 1.
    - a,b  -- limits (default 0,1)
    - n    -- table size (ignored if a histogram or value)
    - rng --  a numpy.random.Generator object, or a seed
    """

    def __init__(self, func, limits=(0,1), n=100, rng=None):

        a,b = limits
        self.x = np.linspace(a,b,n+1) # bin edges
        dx = (b-a)/(n)/2
        self.deltafun=None
        self.rng=rng

        if callable(func):
            # A function
            # evaluate at bin centers
            y = np.array([func(t-dx) for t in self.x])
            if np.any(y<0) or np.sum(y)==0:
                raise ValueError('Function is not positive definite')
        elif isinstance(func, numbers.Number):
            # a single value, or delta function
            self.deltafun = func
            if  func<0 or func>1:
                raise ValueError('Value not in range [0,1]')
            self.mean=func
            return
        else:
            n = len(func)
            self.x = np.linspace(a,b,n)
            y = func
        cy = np.cumsum(y)
        d = cy[-1]-cy[0]
        self.sy = (cy-cy[0])/d

        self.mean = np.sum( (self.x-dx) * y) / d

    def _evaluate(self, r):
        """evaluate inverse integral. expect 0<r<1 """
        return np.interp(r, self.sy, self.x)

    def __call__(self, size):
        """Generate `size` values
        """
        if self.deltafun: return np.full(size, self.deltafun)

        return self._evaluate(stats.uniform.rvs(size=size, random_state=self.rng))

# Cell
sec_per_day = 24*3600

def generate_times(start, stop, count, rng=None):
    """ Generate a list of times, distributed randomly

    - start, stop: times
    - count : expected number to generate with rate=count/(stop-start)
    - rng : random state generator

    returns : list of times between start and stop. Note that the actual number is Poisson-distributed
    """
    # note: can speed this up by making groups of random calls

    tt =[]
    t = start
    scale = (stop-start)/count
    if rng is None: rng = np.random.default_rng()
    while True:
        t += rng.exponential(scale =scale)
        if t>stop: break
        tt.append(t)
    return tt

# Cell
class WeightFunction(object):

    def __init__(self, s=1, b=1, wt_signif=0.1, rng=None):
        self.s = s
        self.b = b
        self.lam = wt_signif
        self.rng = rng

    def psf(self, rsq, sig=0.2, gamma=2):
        """ The PSF as a function of a radius**2
        """
        #return 1/(2*np.pi* sig**2) * (1-1/gamma) * (1 + rsq/sig**2/(2*gamma))**-gamma
        return np.exp(-rsq/self.lam)/(self.lam*(1-np.exp(-1/self.lam)))

    def __call__(self, rsq):
        """ The rsq distribution function
        """
        return self.s * self.psf(rsq) + self.b

    def sample(self, s, n):
        self.s = s
        return _Sampler(self, rng=self.rng)(n);

    def weights(self, s, n):
        h = self.sample(s,  n)
        p = self.s * self.psf(h)
        return 1/(1+self.b/p)

# Cell
def make_exposure(fexp, start, stop, interval=300):
    """
    - fexp -- exposure in cm^2, a value or a function of time in day units
    - start, stop -- range of time in day units
    - interval [300] -- 5-min interval (fermi data is 30 s)

    Returns: a DataFrame with start, stop, exp
    """
    def check_scalar( f):
        if np.isscalar(f):
            fval = f
            return lambda t: fval
        return f
    fexp = check_scalar(fexp)

    nbins = int((stop-start)*sec_per_day / interval)
    edges = np.linspace(start, start+nbins*interval/sec_per_day, nbins+1)
    starts, stops = edges[:-1], edges[1:]
    exp = fexp(starts) * interval
    return pd.DataFrame.from_dict(dict(start=starts, stop=stops,
                                       cos_theta=np.ones(nbins,np.float32),
                                       exp=exp,
                                      ))

# exp  = make_exposure(500, 0, 1 )
# days  = np.sum(exp.stop-exp.start); secs = days*24*3600
# exptot=np.sum(exp.exp)
# exp_text = f' average {exptot/secs:.0f} cm^2 for {secs/1e6:.1f} Ms'
# print(exp_text)
# exp.head()

# Cell
class Simulation(object):
    """
    - src_flux : source flux, scalar or function of days, typically around 1e-7
    - tstart, tstop :(days)
    - bkg_rate : background flux, scalar or function of day, typicaly 1e-6 for 4-deg cone
    - efun : scalar | function (of time in days) of the exposure/s. Typically 3000 cm^2 for fermi

    - wt_signif : the width of the PSF in (r/rmax)**2 coordinates
    - rng : random generator instance, or integer seed

    """

    def __init__(self, name, src_flux, tstart, tstop,
                 bkg_rate=1e-6,  efun=3000, wt_signif=0.1,
                 debug=False, rng=None, config=None):

        def check_scalar( f):
            if np.isscalar(f):
                fval = f
                return lambda t: fval
            return f
        self.config = config or Config()
        self.name = name
        self.src_fun = check_scalar(src_flux)
        self.bkg_fun = check_scalar(bkg_rate)
        self.flux_fun = lambda t: src_fun(t)+bkg_fun(t)
        self.wt_signif=wt_signif

        self.exposure = make_exposure(efun, tstart, tstop)
        if isinstance(rng, np.random.Generator):
            self.rng = rng
        else:
            self.rng = np.random.default_rng(rng)
        self.filename = None #flag that not a regular source
        self.debug=debug


    def run(self):
        times = []
        weights = []
        for start, stop, costh, exp  in self.exposure.itertuples(index=False,name=None):

            src = self.src_fun((start+stop)/2)
            bkg = self.bkg_fun((start+stop)/2)
            delta_t = (stop-start)*sec_per_day # tolal tim
            counts = (src+bkg) * exp
            new_times = generate_times(start, stop, counts, rng=self.rng)
            wfun = WeightFunction(wt_signif=self.wt_signif, s=src, b=bkg, rng=self.rng)
            new_wts = wfun.weights(s=src, n=len(new_times))
            if self.debug>0:
                print(f'time: {start:.3f} - {stop:.3f}, source {src}, exposure/s {exp/delta_t:.0f}, expected/generated counts {counts:.0f}/{len(new_wts)}')
            if self.debug>0:
                print(f'\t weights {np.array(new_wts).round(2)}')
                self.debug-=1

            assert len(new_times)==len(new_wts)
            times = np.append(times, new_times)
            weights = np.append(weights, new_wts)

        print(f'generated {len(times)} photons')
        self.photons=pd.DataFrame(dict(
            time=times,
            weight=weights.astype(np.float32),
        ))