functions {
  
  real impute_missing(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;
  }

  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 0:(size(y_i)-1)) {
      if (y_i[size(y_i) - k]) {
        return size(y_i) - k;
      }
    }
    return -1;
  }

  int find_block_start(int[] y_i, int start_ind, int i) {
    int kback;
    for (k in 1:start_ind) {
      kback = start_ind - k;
      if (y_i[kback]) {
        return kback;
      }
    }
  return 0;
  }
  
}
data {
  int<lower = 0> T_years; // years in dataset
  int<lower = 0> N_plant; // number of plants in dataset
  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 was alive in its final census
    // (implied: there will be no records for "dead", i.e., 3)
    if (y[i, last_year] < 3) {

      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
        } 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
    } else {
      
      int block_start = find_block_start(y[i], last_year, i);
//      print("i", i);
//      print("y_i", y[i]);
//      print("last year", last_year);
//      print("block_star", block_start);
      
      if (y[i,last_year-1]) {

        1 ~ bernoulli(p[2,last_year]);
        // if plant was flowering in prior year
        if (y[i, last_year-1] == 2) {
          // flower as flowering plant in year last-1 NO
          0 ~ bernoulli(f[2, last_year-1]);
          // flpl survived from year t-1 YES
          1 ~ bernoulli(s[2, last_year-1]);
        // if plant was basal in prior year
        } else {
          // flower as basal plant in year t-1 NO
          0 ~ bernoulli(f[1, last_year-1]);
          // survived from year t-1 YES
          1 ~ bernoulli(s[1, last_year-1]);
        }
      } else {
        int k_ind = y[i, block_start];
        real prob_combined = impute_missing(k_ind, block_start, last_year, p, s, f);
        print(prob_combined)
        // 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]);
          }
        }
      }
    }
  }
}