// Functional Group model

data {
  int<lower=0> N;        // Number of observations
  int<lower=0> T;        // Number of years
  int<lower=0> sites;    // number of sites
  int<lower=0> S;        // Number of species
  int<lower=0> N_func;   // Number of functional groups should be 5
  int functional[S];     // Vector of functional group assignments
  int y[N,S];     // Counts
  int<lower=0> year[N];  // Vector of years
  int<lower=0> site[N]; //vector of sites
  vector[S] meanalpha;
  int<lower=0> monum[N];
  int<lower=0> yearzoop;
  real zoop[yearzoop];
}

parameters {
  vector<lower=0>[S] sigma_obs;    // Observation different across species
  vector<lower=0>[N_func] sigma_proc;  // process error
  vector<lower=0>[S] sigma_site;
  vector<lower=0>[S] sigma_month;
  row_vector[S] rand[N];
  row_vector[S] sitere_raw[sites];
  row_vector[S] month_raw[3];
  row_vector[S] r_raw[T-1];
  vector<upper=20>[S] init_log_alpha;  // initial log_alpha
  vector[N_func] meanr[2];
  real ZooBeta[N_func];
  
}

transformed parameters {
  matrix[N,S] mu;
  matrix<upper=60>[T,S] log_alpha;
  matrix[sites,S] site_RE;
  matrix[3,S] month_RE;
  row_vector[S] r[T-1];

  for(j in 1:S){
   //initial population size
    log_alpha[1,j] = init_log_alpha[j];
    
    //Loop through first 13 years prior to zooplankton numbers            
    for(k in 1:T-1){     
      if(k<14){
        r[k,j] = meanr[1,functional[j]] + r_raw[k,j] * sigma_proc[functional[j]]; //mean r will be linked at the functional group level
      }
      else{
        r[k,j] = (meanr[2,functional[j]] + ZooBeta[functional[j]] * zoop[k-13] )+ r_raw[k,j] * sigma_proc[functional[j]]; //mean r will be linked at the functional group level
      }
      
      log_alpha[k+1,j] = log_alpha[k,j] + r[k,j];
  
    }
    
    //observation random effect for site
    for(k in 1:sites){
      site_RE[k,j] = sitere_raw[k,j] * sigma_site[j] ;
    }
  
    for(k in 1:3){
      month_RE[k,j] = month_raw[k,j] * sigma_month[j];
    }
  }

    for(i in 1:N){
      mu[i,] = exp(log_alpha[year[i],] + month_RE[monum[i],] + site_RE[site[i],] + rand[i,]);
    }  
  
}
model {
  init_log_alpha ~ normal(meanalpha,3);
  sigma_proc ~ normal(0,1); //probbly too restrictive. More variation than this would allow.
  sigma_site ~ normal(0,1);
  sigma_month ~ normal(0,1);
  ZooBeta ~ normal(0,1);
  sigma_obs ~ cauchy(0,1);
  meanr[1] ~ normal(0,2);
  meanr[2] ~ normal(0,2);

  for(k in 1:T-1){
    r_raw[k] ~ normal(0,1); // implies r[t,j] ~ normal(meanr[functional[j]],sigma_proc[functional[j]]);
  }
  
  for(k in 1:sites){
    sitere_raw[k] ~ normal(0,1);
  }
  
  for(k in 1:3){
  month_raw[k] ~ normal(0,1);
  }
  // Likelihood
  //observation process for one species
  for(x in 1:S){
    rand[,x] ~ normal(0,sigma_obs[x]);

    y[,x] ~ poisson(mu[,x]);
    
  }
}


generated quantities{
  row_vector[S] N_est[T];
  matrix[N,S] log_lik;

  for(x in 1:S){
    for(i in 1:N){;
      log_lik[i,x] = poisson_lpmf(y[i,x] | mu[i,x]);
    }
  }
    
  for(k in 1:T){
      N_est[k] = exp(log_alpha[k,]);
  }
}