1 files changed, 30 insertions, 24 deletions

diff --git a/src/nll_gdistsamp.cpp b/src/nll_gdistsamp.cpp
index 80dc2d4..baed007 100644
--- a/src/nll_gdistsamp.cpp
+++ b/src/nll_gdistsamp.cpp
@@ -16,8 +16,7 @@ double nll_gdistsamp(arma::vec beta, arma::uvec n_param, arma::vec y,
     int mixture, std::string keyfun, std::string survey,
     arma::mat Xlam, arma::vec Xlam_offset, arma::vec A, arma::mat Xphi,
     arma::vec Xphi_offset, arma::mat Xdet, arma::vec Xdet_offset, arma::vec db,
-    arma::mat a, arma::mat u, arma::vec w, arma::vec k, arma::vec lfac_k,
-    arma::cube lfac_kmyt, arma::cube kmyt, arma::uvec Kmin, int threads){
+    arma::mat a, arma::mat u, arma::vec w, int K, arma::uvec Kmin, int threads){
 
   #ifdef _OPENMP
     omp_set_num_threads(threads);
@@ -25,9 +24,7 @@ double nll_gdistsamp(arma::vec beta, arma::uvec n_param, arma::vec y,
 
   int M = Xlam.n_rows;
   int T = Xphi.n_rows / M;
-  int R = y.size() / M;
-  unsigned J = R / T;
-  int K = k.size() - 1;
+  unsigned J = db.size() - 1;
 
   //Abundance
   const vec lambda = exp(Xlam * beta_sub(beta, n_param, 0) + Xlam_offset) % A;
@@ -51,43 +48,52 @@ double nll_gdistsamp(arma::vec beta, arma::uvec n_param, arma::vec y,
   #pragma omp parallel for reduction(+: loglik) if(threads > 1)
   for (int i=0; i<M; i++){
 
+    vec f = zeros(K+1); // should this be ones(K+1);
+    for (int k=Kmin(i); k<(K+1); k++){
+      f(k) = N_density(mixture, k, lambda(i), log_alpha);
+    }
+
     int t_ind = i * T;
     int y_ind = i * T * J;
 
     vec y_sub(J);
+    vec y_all(J+1);
+
+    vec cp(J);
+    vec cp_all(J+1);
+    double ptotal;
+
+    vec g = zeros(K+1);
 
-    mat mn = zeros(K+1, T);
     for(int t=0; t<T; t++){
       int y_stop = y_ind + J - 1;
       y_sub = y.subvec(y_ind, y_stop);
-      uvec not_missing = find_finite(y_sub);
+      y_all.subvec(0, (J-1)) = y_sub;
 
-      if(not_missing.size() == J){
+      uvec nm = find_finite(y_sub);
 
-        vec p1 = lfac_kmyt.subcube(span(i),span(t),span());
-        vec p = distprob(keyfun, det_param(t_ind), scale, survey, db,
-                          w, a.row(i));
-        vec p3 = p % u.col(i) * phi(t_ind);
-        //the following line causes a segfault only in R CMD check,
-        //when kmyt contains NA values
-        vec p4 = kmyt.subcube(span(i),span(t),span());
+      if((nm.size() == J)){
 
-        double p5 = 1 - sum(p3);
+        cp = distprob(keyfun, det_param(t_ind), scale, survey, db,
+                      w, a.row(i));
+        cp = cp % u.col(i) * phi(t_ind);
+        ptotal = sum(cp);
+
+        cp_all.subvec(0, (J-1)) = cp;
+        cp_all(J) = 1 - ptotal;
+
+        for (int k=Kmin(i); k<(K+1); k++){
+          y_all(J) = k - sum(y_sub);
+          g(k) += dmultinom(y_all, cp_all);
+        }
 
-        mn.col(t) = lfac_k - p1 + sum(y_sub % log(p3)) + p4 * log(p5);
       }
 
       t_ind += 1;
       y_ind += J;
     }
 
-    double site_lp = 0.0;
-    for (int j=Kmin(i); j<(K+1); j++){
-      site_lp += N_density(mixture, j, lambda(i), log_alpha) *
-        exp(sum(mn.row(j)));
-    }
-
-    loglik += log(site_lp + DBL_MIN);
+    loglik += log(sum(f % exp(g)));
 
   }