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

[((522, 533), 'numpy.array', 'np.array', (['x'], {}), '(x)\n', (530, 533), True, 'import numpy as np\n'), ((1524, 1549), 'numpy.random.poisson', 'np.random.poisson', (['lam', 'N'], {}), '(lam, N)\n', (1541, 1549), True, 'import numpy as np\n'), ((1631, 1647), 'numpy.empty', 'np.empty', (['(N, J)'], {}), '((N, J))\n', (1639, 1647), True, 'import numpy as np\n'), ((948, 980), 'scipy.stats.poisson.pmf', 'stats.poisson.pmf', (['kvals', 'lam[i]'], {}), '(kvals, lam[i])\n', (965, 980), False, 'from scipy import special, stats\n'), ((1695, 1726), 'numpy.random.binomial', 'np.random.binomial', (['(1)', 'p[i,]', 'J'], {}), '(1, p[i,], J)\n', (1713, 1726), True, 'import numpy as np\n'), ((1255, 1264), 'numpy.exp', 'np.exp', (['g'], {}), '(g)\n', (1261, 1264), True, 'import numpy as np\n'), ((1565, 1583), 'numpy.tile', 'np.tile', (['z', '(J, 1)'], {}), '(z, (J, 1))\n', (1572, 1583), True, 'import numpy as np\n'), ((1103, 1125), 'numpy.tile', 'np.tile', (['kvals', '(J, 1)'], {}), '(kvals, (J, 1))\n', (1110, 1125), True, 'import numpy as np\n'), ((1188, 1221), 'scipy.stats.binom.logpmf', 'stats.binom.logpmf', (['ymat', '(1)', 'pmat'], {}), '(ymat, 1, pmat)\n', (1206, 1221), False, 'from scipy import special, stats\n')]