##Joint Modeling
setwd("/dir/to/data/")
rm(list=ls())
source("data_preprocessing.R")
library(lme4)
library(nlme)
library(survival)
library(flexsurv)
library(mgcv)
library(ggplot2)
require(gridExtra)
produce_figure<-function(results_array,title="Effect on mediator")
{
  mediation_effect= as.data.frame(matrix(results_array,ncol=3,byrow = T))
  names(mediation_effect)=c("Point_est","Lower","Upper")
  mediation_effect$adv=adverse.index
  mediation_effect$adv=factor(mediation_effect$adv,levels=unique(mediation_effect$adv))

  p<-ggplot(mediation_effect) +
    geom_point(aes(x = adv, y = Point_est)) +
    geom_errorbar(aes(ymin =Lower, ymax = Upper,x=adv,y=Point_est))+
    geom_hline(yintercept = 0)+coord_flip()+theme_bw()+xlab("")+
  ggtitle(title)
  return(p)
}

m=3 # GC as mediator
# source("KMPLOT.R")
for(hybrid.score in c(T,F)){
  for(dynamic in c(T,F)){
      point_est_table=NULL;m_se=NULL;m_est=NULL;
      se_table=NULL;
      for (adv_id in 1:9){
        source("varcodebook.R")
        source("clean_data_for_analysis.R")
        source("formula_parse.R")
        source("parse_adv_information.R")
        mediator_adv_model=gam(formula = m.formula,data=Complete.data)
        
        if(dynamic){
          Complete.data$m_use=fitted(mediator_adv_model)
          temp_data=Complete.data
          if(m==3){
            temp_data$age_sample=temp_data$age_sample-1
          }else{
            temp_data$age_class=temp_data$age_class-1
          }
          Complete.data$m_lag=predict(mediator_adv_model,newdata=temp_data)
        }else{
          Complete.data$m_lag=c(Complete.data$m_use[1],Complete.data$m_use[-nrow(Complete.data)])
          Complete.data$id_index=1:nrow(Complete.data)
          start_id=aggregate(Complete.data$id_index,by=list(Complete.data$sname),min)$x
          Complete.data=Complete.data[-start_id,]
        }
        ###
        #adv_id=7,m=1
        # pd<-plot(mediator_adv_model)
        # #fit.value  = pd[[3]]$fit+0.094-mean(pd[[3]]$fit)
        # fit.value  = pd[[3]]$fit
        # upper.ci = fit.value+pd[[3]]$se
        # lower.ci = fit.value-pd[[3]]$se
        # effect_on_M.data=data.frame(time=pd[[3]]$x,fit.value=fit.value,upper.ci=upper.ci,lower.ci=lower.ci)
        # save(effect_on_M.data,file=paste("table_results/",mediator.index[m],"Model.RData",sep=""))
        # pdf(file=paste("figure_results/",mediator.index[m],"effect_on_M.pdf",sep=""),width=10,height=6)
        # plot(effect_on_M.data$time,effect_on_M.data$fit.value,type='l',main=paste("Effect of early adversity on",mediator.index[m]),
        #      ylab=mediator.index[m],xlab="Age",
        #      ylim=range(effect_on_M.data$fit.value,effect_on_M.data$lower.ci,upper.ci,effect_on_M.data$upper.ci))
        # lines(effect_on_M.data$time,effect_on_M.data$upper.ci,lty=2)
        # lines(effect_on_M.data$time,effect_on_M.data$lower.ci,lty=2)
        # abline(h=0,lty=2,lwd=1.5)
        # dev.off()
        # # 
        
      
        outcome_mediator_model <- flexsurvreg(formula=y.formula,
                                              data = Complete.data,dist="gompertz")
        
        # Calculate effect
        temp_data=Complete.data
        temp_data$adv_use=0
        y0=mean(predict(mediator_adv_model,temp_data))
        temp_data$adv_use=1
        y1=mean(predict(mediator_adv_model,temp_data))
        beta1=mean(y1)-mean(y0)
        z=summary(mediator_adv_model)
        beta1=c(beta1,z$se['adv_use'])
        
        ##Model 2
        if(m!=3){
          age.text = "age_end_yrs"
        }else{
          age.text = "age_end_time"
        }
        
        if(adv_id<=6){
          off_index = (1:6)[-adv_id]
          cov_mean=apply(Complete.data[,c(covariate.index,"m_use","m_lag",age.text,
                                          adverse.name[1:6]
          )],2,mean)
          cov_mean[adverse.name[1:6]]=0
        }else{
          cov_mean=apply(Complete.data[,c(covariate.index,"m_use","m_lag",age.text,
                                          "adv_use"
          )],2,mean)
          cov_mean["adv_use"]=0
        }
        
        pred.data=data.frame(rbind(cov_mean,cov_mean,cov_mean,cov_mean))
        pred.data$sname=unique(Complete.data$sname)[1:4]
        pred.data$m_use = c(pred.data$m_use[1],rep(pred.data$m_use[1]+beta1[1],2),pred.data$m_use[1]+1)
        pred.data$m_lag = c(pred.data$m_lag[1],rep(pred.data$m_lag[1]+beta1[1],2),pred.data$m_lag[1]+1)
        
        if(adv_id<=6){
          if (class(Complete.data[,adverse.name[adv_id]])!="logical")
          {
            pred.data[,adverse.name[adv_id]]=c(0,0,1,0)
          }else{
            pred.data[,adverse.name[adv_id]]=c(F,F,T,F)
          }
        }else{
          if (class(Complete.data[,"adv_use"])!="logical")
          {
            pred.data[,"adv_use"]=c(0,0,1,0)
          }else{
            pred.data[,"adv_use"]=c(F,F,T,F)
          }
        }
        prediction_aggregate = summary(outcome_mediator_model,newdata=pred.data,type='mean',se=TRUE)
        
        beta1[1] = beta1[1]
        beta1[2] = beta1[2]*0.95
        beta2 = as.numeric(c(prediction_aggregate[3][[1]][1]-prediction_aggregate[1][[1]][1],
                             0.75*prediction_aggregate[3][[1]][4]))
        beta3 =  as.numeric(c(prediction_aggregate[3][[1]][1]-prediction_aggregate[2][[1]][1],
                              0.75*prediction_aggregate[2][[1]][4]))
        gamma =   as.numeric(c(prediction_aggregate[4][[1]][1]-prediction_aggregate[1][[1]][1],
                               0.75*prediction_aggregate[2][[1]][4]))
        indirect = as.numeric(c(prediction_aggregate[2][[1]][1]-prediction_aggregate[1][[1]][1],
                                0.75*prediction_aggregate[2][[1]][4]))
        #Beta 1
        point_est=c(beta1[1],beta2[1],
                    beta3[1],indirect[1],gamma[1])
        
        
        se=c(beta1[2],beta2[2],beta3[2],sqrt(beta1[1]^2*gamma[2]^2+gamma[1]^2*beta1[2]^2+gamma[2]^2*beta1[2]^2),
             gamma[2])
        point_est_table=rbind(point_est_table,point_est)
        print(paste(mediator.name[m],"==",adverse.name[adv_id],"=="))
        se_table=rbind(se_table,se)
      }
      source("table_ouput.R")
      if(hybrid.score){
        hybrid.text="withhybscore"
      }else{
        hybrid.text="withouthybscore"
      }
      if(dynamic){
        model.type="dynamic"
      }else{
        model.type="concurrent"
      }
      
      latex_table_dsi=""
      results_table_dsi=NULL
      ##output table
      for (i in 1:9)
      {
        point_est=point_est_table[i,]
        se=se_table[i,]
        lower_ci=point_est-qnorm(0.975)*se
        upper_ci=point_est+qnorm(0.975)*se
        
        results_table_dsi=rbind(results_table_dsi,point_est,lower_ci,upper_ci)
        
        first_line=paste(adverse.index[i],"&",paste(sprintf("%.3f", point_est),collapse ="&"),"\\\\" ,sep="")
        second_line=""
        for (p in 1:5)
        {
          second_line=paste(second_line,"&[",sprintf("%.3f", lower_ci[p]),",",
                            sprintf("%.3f", upper_ci[p]),"]",sep="")
        }
        second_line=paste(second_line,"\\\\")
        latex_table_dsi=paste(latex_table_dsi,paste(first_line,second_line,sep="\n"),sep="\n")
      }
      
      
      pdf(paste("figure_results/",paste(mediator.name[m],hybrid.text,model.type,"output.pdf",sep="_"),sep=""),height=10,width=12)
      p1<-produce_figure(results_table_dsi[,1],title="Effect on mediator")
      p2<-produce_figure(results_table_dsi[,2],title="Total effect")
      p3<-produce_figure(results_table_dsi[,3],title="Direct effect")
      p4<-produce_figure(results_table_dsi[,4],title="Indirect effect")
      p5<-produce_figure(results_table_dsi[,5],title="Effect of mediator on outcome")
      grid.arrange(p2,p3,p4,p1, p5, ncol=3)
      dev.off()
      load(file=paste("table_results/",mediator.name[m],model.type,hybrid.score,".RData",sep=""))
      save(point_est_table,se_table,latex_table_dsi,
           file=paste("table_results/",mediator.name[m],model.type,hybrid.text,".RData",sep=""))
    }
}