functions {
  matrix kron(matrix A, matrix B) { //kronecker product, stan  does not have built in
  matrix[rows(A) * rows(B), cols(A) * cols(B)] C;
  int m;
  int n;
  int p;
  int q;
  m = rows(A);
  n = cols(A);
  p = rows(B);
  q = cols(B);
  for (i in 1:m) {
    for (j in 1:n) {
      int row_start;
      int row_end;
      int col_start;
      int col_end;
      row_start = (i - 1) * p + 1;
      row_end = (i - 1) * p + p;
      col_start = (j - 1) * q + 1;
      col_end = (j - 1) * q + q;
      C[row_start:row_end, col_start:col_end] = A[i, j] * B;
    }
  }
  return C;
  }
}

data {
  int<lower=1> p; //lag order
  int<lower=p+1> N; //number of observations
  int<lower=2> k; //number of endogenous variables
  int<lower=1> q; //number of deterministic variables
  matrix[N, k] Y; //endogenous variables (y_t)
  matrix[N, q] D; //deterministic variables (d_t)
  matrix[N, k*p] W; //lagged endogenous variables
  matrix[N, q*p] Q; //lagged deterministic variables
  vector[k*p*k] theta_beta; //vec_beta prior mean
  matrix[k*p*k, k*p*k] Omega_beta; //vec_beta prior covariance matrix
  vector[k*q] theta_Psi; //vec_Psi prior mean
  matrix[k*q, k*q] Omega_Psi; //vec_Psi prior covariance matrix
  vector[k*(k-1)/2] theta_A; //a prior mean
  matrix[k*(k-1)/2, k*(k-1)/2] Omega_A; //a prior covariance matrix
  vector[k] mu_log_lambda_1;        // log lambda initial condition prior means
  vector<lower=0>[k] sigma2_log_lambda_1; // log lambda initial condition prior variances
  vector<lower=0>[k] alpha_phi;     // phi prior shapes
  vector<lower=0>[k] beta_phi;      // phi prior scales
  int<lower=1> H; // Forecast horizon
  matrix[H, q] d_pred; //future deterministic variables
}

transformed data {
    matrix[p, p] I_p = diag_matrix(rep_vector(1, p)); // Identity matrix
}

parameters {
  matrix[k*p, k] beta; //beta' = (Pi_1,...,Pi_p)
  matrix[k, q] Psi; //Psi * d_t = steady state
  matrix[N, k] z; //standard normals
  vector<lower=0>[k] phi; //log volatility innovation variances
  vector[k*(k-1)/2] a; //free parameters in A
}

transformed parameters {
  matrix[k,k] A;
  matrix[k,k] Ainv;
  matrix[k,k] Sigma_u[N]; //time varying covariance matrix of reduced form errors u_t
  matrix[N, k] log_lambda; //log volatilities
  // construct A (1's on diagonal, and then free parameters on lower triangular)
  A = diag_matrix(rep_vector(1,k));
  {
    int idx = 1;
    for (i in 2:k) {
      for (j in 1:(i-1)) {
        A[i,j] = a[idx];
        idx += 1;
      }
    }
  }
  Ainv = inverse(A);
  
  for (i in 1:k) {
    log_lambda[1, i] = mu_log_lambda_1[i] + sqrt(sigma2_log_lambda_1[i]) * z[1, i];
  }
  for (t in 2:N) {
    for (i in 1:k) {
      log_lambda[t, i] = log_lambda[t-1, i] + sqrt(phi[i]) * z[t, i];
    }
  }
  
  for (t in 1:N) {
    matrix[k,k] sqrt_Lambda_t = diag_matrix(sqrt(exp(log_lambda[t]')));
    Sigma_u[t] = tcrossprod(Ainv * sqrt_Lambda_t);
  }
}

model {
  
  to_vector(z) ~ std_normal();

  for(t in 1:N){
      vector[k] u_t = (Y[t] - (D[t]*Psi' + (W[t]-Q[t]*(kron(I_p,Psi')))*beta))';
      u_t ~ multi_normal(rep_vector(0,k), Sigma_u[t]);
  }

  to_vector(beta) ~ multi_normal(theta_beta, Omega_beta);
  to_vector(Psi) ~ multi_normal(theta_Psi, Omega_Psi);
  a  ~ multi_normal(theta_A, Omega_A);
  for (i in 1:k) {
  phi[i] ~ inv_gamma(alpha_phi[i], beta_phi[i]);
}
}

generated quantities {

  matrix[k, k] Pi[p];
  for (i in 1:p) {
    Pi[i] = (beta[((i - 1) * k + 1):(i * k), :])'; //extract Pi_1, ..., Pi_p
  }
  
  matrix[H,k] y_pred;
  matrix[H,k] log_lambda_pred;
  matrix[k,k] Sigma_u_pred[H];

  
  for (i in 1:k) {
    log_lambda_pred[1,i] = normal_rng(log_lambda[N,i], sqrt(phi[i]));
    }
  
  for (h in 2:H) {
    for (i in 1:k) {
      log_lambda_pred[h,i] = normal_rng(log_lambda_pred[h-1,i], sqrt(phi[i]));
      }
    }

  
  for (h in 1:H) {
    
    matrix[k,k] sqrt_Lambda;
    sqrt_Lambda = diag_matrix(sqrt(exp(log_lambda_pred[h]')));
    
    vector[k] epsilon = multi_normal_rng(rep_vector(0,k), diag_matrix(rep_vector(1, k)));
    vector[k] u_t = Ainv * sqrt_Lambda * epsilon;
    
    Sigma_u_pred[h] = tcrossprod(Ainv * sqrt_Lambda);
    
    vector[k] yhat_t = (d_pred[h]*Psi')';

    if (h > 1) {
      for (i in 1:min(h-1, p)) {
        yhat_t += to_vector((y_pred[h-i] - d_pred[h-i]*Psi') * Pi[i]');
      }
    }

    if (h <= p) {
      for (i in h:p) {
        yhat_t += to_vector((Y[N + h - i] - D[N + h - i]*Psi') * Pi[i]');
      }
    }
    y_pred[h] = (yhat_t + u_t)';
  }
}
