functions {
  
  real impute_observed(int k, int T0, int T1, real[,] p, real[,] s, real[,] f) {
        real prob;
    int T = T1 - T0;
    matrix[T, T] A;
    for (row in 1:T) {
      if (row > 1) {
      A[row,1] = s[k,T0]
                        * (1 - f[k,T0])
                        * (1 - p[1,T0+1]);
      } else { 
      A[row,1] = (1 - s[k,T0])
                        * (1 - p[2,T0+1]);
      }
      for (col in 2:T) {
        if (col < T) {
          if (col > row) {
          A[row,col] = A[row,col-1]
                        * (1 - p[2,T0+col]);
          } else if (col == row) {
          A[row,col] = A[row,col-1]
                        * (1 - s[1,(T0-1)+col])
                        * (1 - p[2,T0+col]);
          } else {
          A[row,col] = A[row,col-1]
                        * s[1,(T0-1)+col]
                        * (1 - f[1,(T0-1)+col])
                        * (1 - p[1,T0+col]);
          }
        } else {
          if (row < T) {
          A[row,T] = A[row,T-1]
                      * p[2,T1];
          } else {
          A[row,T] = A[row,T-1]
                      * (1 - s[1,T1-1])
                      * p[2,T1];
          }
        }
      }
    }
    
    prob = 0.0;
    for (i in 1:T)
      prob += A[i,T];
    return prob;
  }

  real impute_unobserved(int k, int T0, int T1, real[,] p, real[,] s, real[,] f) {
    real prob;
    int T = T1 - T0 + 1;
    matrix[T, T] A;
    for (row in 1:T) {
      if (row > 1) {
      A[row,1] = s[k,T0]
                        * (1 - f[k,T0])
                        * (1 - p[1,T0+1]);
      } else { 
      A[row,1] = (1 - s[k,T0])
                        * (1 - p[2,T0+1]);
      }
      for (col in 2:(T-1)) {
        if (col > row) {
          A[row,col] = A[row,col-1]
                        * (1 - p[2,T0+col]);
        } else if (col == row) {
          A[row,col] = A[row,col-1]
                        * (1 - s[1,T0+(col-1)])
                        * (1 - p[2,T0+col]);
        } else {
          A[row,col] = A[row,col-1]
                        * s[1,T0+(col-1)]
                        * (1 - f[1,T0+(col-1)])
                        * (1 - p[1,T0+col]);
        }
      }
      if (row < T) {
        A[row,T] = A[row,T-1]
                    * (1 - p[2,T1]);
        } else {
        A[row,T] = A[row,T-1]
                    * (1 - s[1,T1-1])
                    * (1 - p[2,T1]);
        }
      }
    
    prob = 0.0;
    for (i in 1:T)
      prob += A[i,T];
    return prob;
  }

  int first_obs(int[] y_i) {
    for (k in 1:size(y_i)) {
      if (y_i[k]) return k;
    }
    return 0;
  }

  int last_obs(int[] y_i, int i) {
    for (k in 1:(size(y_i))) {
      if (y_i[k] == 3) return k;
    }
  return size(y_i);
  }

  int find_block_start(int[] y_i, int start_ind, int i) {
    int k_back;
    for (k in 1:start_ind) {
      k_back = start_ind - k;
      if (y_i[k_back]) {
        return k_back;
      }
    }
  return 0;
  }
  
}
data {
  int<lower = 0> T_years;
  int<lower = 0> N_plant;
  int<lower = 0, upper = 3> y[N_plant, T_years];
}
parameters {
  real<lower = 0, upper = 1> p[2, T_years]; // detection probability in each year
  // (p[1,] is basal plants, p[2,] is dead plants
  // by design detection of adult plants is 1)
  real<lower = 0, upper = 1> s[2, T_years-1]; // survival probability in each year
  // (s[1,] is survival of basal plants, s[2,] is survival of flowering plants
  // by definition dead plants remain dead)
  real<lower = 0, upper = 1> f[2, T_years-1]; // flowering probability in each year
  // (f[1,] is the flowering probability for basal plants, f[2,] is same for
  // for flowering plants; dead plants don't flower
}
// y[i, t] is obs. where rows are individual plants and cols are years
// y = 0 for unobserved, 1 for basal, 2 for flowering, 3 for dead
// p[1:2, t] is detection probability for basal (1) and dead plants in each year (conditioned on state)
// s[1:2, t] is survival prob. for basal (1) and flowering (2) plants in each year (conditioned on state)
// f[1:2, t] is probability of basal plant transitioning to flowering plant in given year
model {
  
  for (i in 1:2)
     for (j in 1:T_years-1) {
       s[i,j] ~ beta(12, 2);
       f[i,j] ~ beta(2, 4);
       p[i,j] ~ beta(0.75, 0.8);
     }

  for (i in 1:N_plant) {
    int first_year;
    int last_year;
    first_year = first_obs(y[i]);
    last_year = last_obs(y[i], i);

    // CASE 1
    // if plant observed was alive in its final census)
    if (y[i, last_year] == 1 || y[i, last_year] == 2) {

      for (t in (first_year+1):last_year) {

        // if plant is basal in year t
        if (y[i, t] == 1) {
          // detected in year t YES
          1 ~ bernoulli(p[1, t]);
          // if plant was flowering
          if (y[i, t-1] == 2) {
            // flower as flowering plant in year t-1 NO
            0 ~ bernoulli(f[2, t-1]);
            // flpl survived from year t-1 YES
            1 ~ bernoulli(s[2, t-1]);
            // if plant was not flowering (observed or not)
          } else {
            // flower as basal plant in year t-1 NO
            0 ~ bernoulli(f[1, t-1]);
            // basal plant survived from year t-1 YES
            1 ~ bernoulli(s[1, t-1]);
          }
          // if plant was flowering
        } else if (y[i, t] == 2) {
          // detection probability for flpl is always 1
          // if plant was flowering the year before
          if (y[i, t-1] == 2) {
            // flower as flowering plant in year t-1 YES
            1 ~ bernoulli(f[2, t-1]);
            // flpl survived from year t-1 YES
            1 ~ bernoulli(s[2, t-1]);
            // if plant was not flowering (observed for not)
          } else {
            // flower as basal plant in year t-1 YES
            1 ~ bernoulli(f[1, t-1]);
            // basal survived from year t-1 YES
            1 ~ bernoulli(s[1, t-1]);
          }
          // if plant was not observed
          // note - implicit here is that because there is a later observation of a living plant,
          // this plant MUST be basal (not)
        } else {
          // detected in year t NO
          0 ~ bernoulli(p[1, t]);
          // if plant was flowering in prior year
          if (y[i, t-1] == 2) {
            // flower as flowering plant in year t-1 NO
            0 ~ bernoulli(f[2, t-1]);
            // flpl surived from year t-1 YES
            1 ~ bernoulli(s[2, t-1]);
            // if plant was not flowering (observed or not)
          } else {
            // flower as basal plant in year t-1 NO
            0 ~ bernoulli(f[1, t-1]);
            // basal surived from year t-1 YES
            1 ~ bernoulli(s[1, t-1]);
          }
        }
      }

    // CASE 2
    // if plant was not alive in its final census (0, unobserved, or 3, dead)
    } else {
      
      int block_start = find_block_start(y[i], last_year, i);
      
      int k_ind = y[i, block_start];
      real prob_combined;
      if (y[i, last_year]) {
        // last observation is a dead plant - impute a block ending with observed death
        prob_combined = impute_observed(k_ind, block_start, last_year, p, s, f);

      } else {
        // there is no observation of death - impute a block ending with no observation
        prob_combined = impute_unobserved(k_ind, block_start, last_year, p, s, f);
      }
      // targ += bernoulli_pmf(1 | prob_combined)
      1 ~ bernoulli(prob_combined);

      for (t in (first_year+1):block_start) {

        // if plant is basal in year t
        if (y[i, t] == 1) {
          // detected in year t YES
          1 ~ bernoulli(p[1, t]);
          // if plant was flowering
          if (y[i, t-1] == 2) {
            // flower as flowering plant in year t-1 NO
            0 ~ bernoulli(f[2, t-1]);
            // flpl survived from year t-1 YES
            1 ~ bernoulli(s[2, t-1]);
            // if plant was not flowering (observed or not)
          } else {
            // flower as basal plant in year t-1 NO
            0 ~ bernoulli(f[1, t-1]);
            // survived from year t-1 YES
            1 ~ bernoulli(s[1, t-1]);
          }
        // if plant was flowering
        } else if (y[i, t] == 2) {
          // obs. prob. for flpl is always 1
          // if plant was flowering
          if (y[i, t-1] == 2) {
            // flower as flowering plant in year t-1 YES
            1 ~ bernoulli(f[2, t-1]);
            // flpl survived from year t-1 YES
            1 ~ bernoulli(s[2, t-1]);
            // if plant was not flowering (observed for not)
          } else {
            // flower as basal plant in year t-1 YES
            1 ~ bernoulli(f[1, t-1]);
            // basal plant survived from year t-1 YES
            1 ~ bernoulli(s[1, t-1]);
          }
        // if plant was not observed
        } else {
          // detected in year t YES
          0 ~ bernoulli(p[1, t]);
          // if plant was flowering in prior year
          if (y[i, t-1] == 2) {
            // flower as flowering plant in year t-1 NO
            0 ~ bernoulli(f[2, t-1]);
            // flpl surived from year t-1 YES
            1 ~ bernoulli(s[2, t-1]);
            // if plant was not flowering (observed or not)
          } else {
            // flower as basal plant in year t-1 NO
            0 ~ bernoulli(f[1, t-1]);
            // basal surived from year t-1 NO
            1 ~ bernoulli(s[1, t-1]);
          }
        }
      }
    }
  }
}