from . import model
import numpy as np
from scipy import special, stats

class RoyleNicholsModel(model.UnmarkedModel):
    def __init__(self, det_formula, abun_formula,  data):
        self.response = model.Response(data.y)
        abun = model.Submodel("Abundance", "abun", abun_formula, np.exp, data.site_covs)
        det = model.Submodel("Detection", "det", det_formula, special.expit, data.obs_covs)
        self.submodels = model.SubmodelDict(abun=abun, det=det)
    
    def negloglik(self, x, mod, K):
        x = np.array(x)
        beta_abun = x[mod["abun"].index]
        beta_det = x[mod["det"].index]
        y = mod.response.y
        N, J = y.shape
        lam = mod["abun"].predict(beta=beta_abun, interval=False)
        r = mod["det"].predict(beta=beta_det, interval=False).reshape(N, J)
        q = 1 - r
        nll = 0.0
        for i in range(N):
            kvals = range(int(mod.response.Kmin[i]), int(K)+1)
            f = stats.poisson.pmf(kvals, lam[i])
            ymat = np.tile(y[i,], (len(kvals), 1)) 
            qmat = np.tile(q[i,], (len(kvals), 1))
            kmat = np.tile(kvals, (J, 1)).transpose()
            pmat = 1 - qmat**kmat
            g = stats.binom.logpmf(ymat, 1, pmat).sum(axis=1)
            fg = f * np.exp(g)
            nll -= np.log(fg.sum())
        
        return nll
    
    def simulate(self):
        N, J = self.response.y.shape
        lam = self.predict("abun", interval=False)
        q = 1 - self.predict("det", interval=False).reshape(N, J)
        z = np.random.poisson(lam, N)
        zrep = np.tile(z, (J,1)).transpose()
        p = 1 - q**zrep
        y = np.empty((N, J))
        for i in range(N):
            y[i,] = np.random.binomial(1, p[i,], J)
        return y

[((13, 12, 13, 23), 'numpy.array', 'np.array', ({(13, 21, 13, 22): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((39, 12, 39, 37), 'numpy.random.poisson', 'np.random.poisson', ({(39, 30, 39, 33): 'lam', (39, 35, 39, 36): 'N'}, {}), '(lam, N)', True, 'import numpy as np\n'), ((42, 12, 42, 28), 'numpy.empty', 'np.empty', ({(42, 21, 42, 27): '(N, J)'}, {}), '((N, J))', True, 'import numpy as np\n'), ((24, 16, 24, 48), 'scipy.stats.poisson.pmf', 'stats.poisson.pmf', ({(24, 34, 24, 39): 'kvals', (24, 41, 24, 47): 'lam[i]'}, {}), '(kvals, lam[i])', False, 'from scipy import special, stats\n'), ((44, 20, 44, 51), 'numpy.random.binomial', 'np.random.binomial', ({(44, 39, 44, 40): '1', (44, 42, 44, 47): 'p[i,]', (44, 49, 44, 50): 'J'}, {}), '(1, p[i,], J)', True, 'import numpy as np\n'), ((30, 21, 30, 30), 'numpy.exp', 'np.exp', ({(30, 28, 30, 29): 'g'}, {}), '(g)', True, 'import numpy as np\n'), ((40, 15, 40, 32), 'numpy.tile', 'np.tile', ({(40, 23, 40, 24): 'z', (40, 26, 40, 31): '(J, 1)'}, {}), '(z, (J, 1))', True, 'import numpy as np\n'), ((27, 19, 27, 41), 'numpy.tile', 'np.tile', ({(27, 27, 27, 32): 'kvals', (27, 34, 27, 40): '(J, 1)'}, {}), '(kvals, (J, 1))', True, 'import numpy as np\n'), ((29, 16, 29, 49), 'scipy.stats.binom.logpmf', 'stats.binom.logpmf', ({(29, 35, 29, 39): 'ymat', (29, 41, 29, 42): '1', (29, 44, 29, 48): 'pmat'}, {}), '(ymat, 1, pmat)', False, 'from scipy import special, stats\n')]