#load libraries
library(rstan)
library(mvtnorm)
library(ggplot2)
library(bayesplot)

#generate data
n<-1000
beta1<-0.25 #effect of instrument on independent variable
beta2<-.75 #effect of independent variable on dependent variable
gamma<-0.05 # size of exclusion restriction violation

#make variance-covariance matrix to make X and Y endogenous
covmat<-matrix(data=NA,nrow=2,ncol=2)
diag(covmat)<-1
covmat[1,2]<-0.25 #strength of endogeneity between X and Y
covmat[2,1]<-covmat[1,2]
eps<-rmvnorm(n,mean=c(0,0),sigma=covmat)

#make variables
z<-rnorm(n,.50,1) #instrument

x<-beta1*z+rnorm(n)+eps[,1] #endogenous independent variable

y<-beta2*x+rnorm(n)+eps[,2] #dependent variable

#now violate exclusion restriction
y<-y+gamma*z

####STAN PART
#now start with code for our stan model
full_bayes_code<-"
data{
  int n; // number of observations
  vector[n] x;
  vector[n] y;
  vector[n] z;
  real gamma_sig; //standard deviation of exclusion restriction violation
}
parameters{
  real beta;
  real pi;
  real<lower=0> gamma; //instrument's direct effect on second stage
}
model{
  y ~ normal(z * pi * beta + z*gamma , 1);
  x ~ normal(z * pi, 1);
  gamma ~ normal( 0 , gamma_sig);
}
"


full_bayes_df<-data.frame(x=x,y=y,z=z)

full_bayes_list<-list(
  n<-nrow(full_bayes_df),
  x<-full_bayes_df$x,
  y<-full_bayes_df$y,
  z<-full_bayes_df$z,
  gamma_sig<-0.1
)

#now run stan model
full_bayes_model<-stan(model_code = full_bayes_code,data=full_bayes_list,iter=2000,chains=4)
print(full_bayes_model)

#compare to frequentist IV to no uncertainty about exclusion restriction
summary(ivreg(y~x|z),diagnostics=TRUE,data=full_bayes_df)


#now try with a function
workshop_fun<-function(sig,iter=2000,chains=1,cores=1){
  mymodel<-stan(model_code = full_bayes_code,
                data=list(
                  n=nrow(full_bayes_df),
                  x=full_bayes_df$x,
                  y=full_bayes_df$y,
                  z=full_bayes_df$z,
                  gamma_sig=sig),
                iter=iter,chains=chains,cores=cores)
  
  if (bayesplot::rhat(mymodel)[1]<1.05) {
    output<-(rstan::extract(mymodel,pars=c("beta")))
    output<-c(mean(output$beta),sd(output$beta))
    print(mymodel)
    return(output)
  }
  else{
    print("Not converged!")
    workshop_fun(sig,iter=iter*2)

  }

}

workshop_fun(.1)

#now run over a range of values and store results
myseq<-seq(0.01,0.25,by=0.05)

mymatrix<-matrix(data=NA,nrow=length(myseq),ncol=2)

for(i in 1:length(myseq)){
  mymatrix[i,]<-workshop_fun(myseq[i])
}

#now plot
mydf<-as.data.frame(mymatrix)

g<-ggplot(data=mydf)+geom_line(aes(y=V1,x=myseq))+geom_line(aes(y=V1+1.96*V2,x=myseq))+geom_line(aes(y=V1-1.96*V2,x=myseq))+xlab("SD of Gamma")+geom_hline(yintercept=0,color="red")+ylab("Marginal Effect on Y")
g