from scipy import special
import numpy as np


def tabul_lnI(nmax, mmax, ygrid, psi1_S, psi1_H):
    """
    tabulate ln(I_beta).
    
    Parameters
    ----------
    nmax: upper bound of soft-band counts.
    mmax: upper bound of hard-band counts.
    ygrid: grid of y(HR, es, eH).
    psi1_S
    psi1_H
    
    Return
    ------
    ln(I_beta) on 3D grids of n (ranging 0-nmax), m (ranging 0-mmax), and y (i.e., ygrid).
    """
    nn, mm, yy = np.meshgrid(np.arange(nmax + 1), np.arange(mmax + 1), ygrid)
    return np.log(special.betainc(mm + psi1_H, nn + psi1_S, yy))

def generate_lnI(gridsize, stdygridpath, stdlnIpath):
    """
    generate a small lnI table from a massive lnI table.
    
    Parameters
    ----------
    gridsize: the desired gridsize of y(HR, es, eH).
    stdygridpath: the .npy file path saving the massive ygrid.
    stdlnIpath: the .npy file path saving the massive lnI table.
        The massive lnI table should be generated by the tabul_lnI function.
    
    Return
    ------
    out1: new y grid.
    out2: new lnI table.
    
    Notes
    -----
    the grid size must be at least 2 and do not exceed the original massive ygrid size.
    """
    stdygrid = np.load(stdygridpath)
    if not ((gridsize >= 2) and (gridsize <= stdygrid.size)):
        raise ValueError("Wrong grid size!")
    stdlnI = np.load(stdlnIpath)
    inds = (np.arange(0, gridsize) * (stdygrid.size - 1) / (gridsize - 1)).astype(int)
    if stdygrid.size - 1 not in inds:
        raise ValueError("grid not set correctly...")
    return stdygrid[inds], stdlnI[:, :, inds]

class fasthr:
    """
    class for HR calculations.
    """
    def __init__(self, S, H, e_S, e_H, psi1_S, psi2_S, psi1_H, psi2_H,
                 B_S = None, B_H = None, psi3_S = None, psi4_S = None, psi3_H = None, psi4_H = None, r_S = None, r_H = None,
                 xi_S = None, xi_H = None):
        """
        instantiate the class.
        
        Parameters
        ----------
        S: total soft-band counts in the source region.
        H: total hard-band counts in the source region.
        e_S: soft-band exposure time.
        e_H: hard-band exposure time.
        psi1_S, psi2_S: parameters of the sources' soft-band count-rate priors.
        psi1_H, psi2_H: parameters of the sources' hard-band count-rate priors.
        
        ### If you do not assume the background intensity to be exactly known, please specify:
        B_S: soft-band counts in the background region.
        B_H: hard-band counts in the background region.
        psi3_S, psi4_S: parameters of the background soft-band count-rate priors.
        psi3_H, psi4_H: parameters of the background hard-band count-rate priors.
        r_S: the soft-band ratio between the background-region area and the source-region area.
        r_H: the hard-band ratio between the background-region area and the source-region area.
        
        ### If you assume the background intensity to be exactly known, please specify:
        xi_S: the expected soft-band background count rate in the source region.
        xi_H: the expected hard-band background count rate in the source region.
        """
        self.S = S
        self.H = H
        self.e_S = e_S
        self.e_H = e_H
        self.psi1_S = psi1_S
        self.psi2_S = psi2_S
        self.psi1_H = psi1_H
        self.psi2_H = psi2_H
        self.nn, self.mm = np.meshgrid(np.arange(S + 1), np.arange(H + 1))
        # background parameters; if unknown bkg
        self.B_S = B_S
        self.B_H = B_H
        self.psi3_S = psi3_S
        self.psi4_S = psi4_S
        self.psi3_H = psi3_H
        self.psi4_H = psi4_H
        self.r_S = r_S
        self.r_H = r_H
        # background parameters; if known bkg
        self.xi_S = xi_S
        self.xi_H = xi_H
    
    def init_hrcdf(self, bkgtype):
        """
        initialize the calculation for HR CDF.
        
        Parmaeters
        ----------
        bkgtype: "fixed" or "unfixed"
            "fixed": assume the background intensity to be exactly known.
            "unfixed": do not assume the background intensity to be exactly known.
        """
        if bkgtype == "fixed":
            self.theta = special.gammaln(self.nn + self.psi1_S) + special.gammaln(self.mm + self.psi1_H)\
                - special.gammaln(self.nn + 1) - special.gammaln(self.mm + 1)\
                - special.gammaln(self.S - self.nn + 1) - special.gammaln(self.H - self.mm + 1)\
                - self.nn * np.log((self.e_S + self.psi2_S) * self.xi_S) - self.mm * np.log((self.e_H + self.psi2_H) * self.xi_H)
        elif bkgtype == "unfixed":
            self.theta = special.gammaln(self.S + self.B_S + self.psi3_S - self.nn) + special.gammaln(self.H + self.B_H + self.psi3_H - self.mm)\
                + special.gammaln(self.nn + self.psi1_S) + special.gammaln(self.mm + self.psi1_H)\
                - special.gammaln(self.nn + 1) - special.gammaln(self.mm + 1)\
                - special.gammaln(self.S - self.nn + 1) - special.gammaln(self.H - self.mm + 1)\
                + self.nn * np.log((self.psi4_S + (1 + self.r_S) * self.e_S) / (self.psi2_S + self.e_S))\
                + self.mm * np.log((self.psi4_H + (1 + self.r_H) * self.e_H) / (self.psi2_H + self.e_H))
        else:
            raise ValueError("bkgtype not supported.")
        # normalize theta
        self.theta -= self.theta.max()
    
    def calc_hrcdf(self, ygrid, lnI):
        """
        calculate the HR CDF.
        
        Parameters
        ----------
        ygrid: grid of y(HR, es, eH).
        lnI: lnI table.
        
        Return
        ------
        hrgrid: grid of HR, on which the CDF is evaluated.
        cdfgrid: the CDF values on the HR grid.
        
        Notes
        -----
        The input parameters should be obtained from the generate_lnI function.
        """
        hrgrid = (ygrid * (self.e_S + self.psi2_S) - (1. - ygrid) * (self.e_H + self.psi2_H)) / (ygrid * (self.e_S + self.psi2_S) + (1. - ygrid) * (self.e_H + self.psi2_H))
        cdfgrid = []
        for iy in range(len(ygrid)):
            cdfgrid.append(np.sum(np.exp(self.theta + lnI[self.mm.astype(int), self.nn.astype(int), iy])))
        cdfgrid = np.array(cdfgrid) / np.exp(self.theta).sum() # normalize so that CDF(1)=1
        return hrgrid, cdfgrid
