i_create_fastqc_figure<-function(path,samples,out){
  require(ggplot2)
  require(Nozzle.R1)
  ############QUALITY PER NT###############
  tfqc<-data.frame()
  for (f in samples){
    src<-paste(path,f,"/qc/fastqc/fastqc_data.txt",sep="")
    endnt<-system(paste("grep -n END ",src," | head -2 | tail -1| sed 's/:/\t/' | cut -f 1"),
                  intern = TRUE)
    endnt<-as.integer(endnt)
    fqc<-suppressMessages(read.table(src,nrows=endnt-14,skip=12,sep="\t"))
    names(fqc)<-c("nt","mean","median","lowQ","upQ","10Q","90Q")
    fqc$sample<-substr(f,1,21)
    tfqc<-rbind(tfqc,fqc)

  }
  tfqcFile<-paste(out,"tfqc.txt",sep="")
  write.table(tfqc,tfqcFile,row.names=F,quote=F,sep="\t")

  p<-ggplot(tfqc,aes(factor(nt),mean)) +
    #geom_boxplot() +
    geom_jitter(position=position_jitter(width=0.3),aes(factor(nt), mean,colour=factor(sample)))+
    geom_abline(intercept = 30,slope=0)+
    ylim(25,42)+
    theme_bw(base_size = 12) +
    theme(axis.text.x = element_blank())+
    labs(list(x="nucleotide",y="score",colour="Samples"))

  ffqcFile="ffqc.jpg"
  ffqcHFile="ffqcH.pdf"
  jpeg(paste(out, ffqcFile, sep="" ) ,width=600,height=400,quality=100 );
  print(p)
  dev.off();
  pdf(paste( out, ffqcHFile, sep="" )  );
  print(p)
  dev.off();

  # create a figure and make it available for exporting
  FRfqc <- newFigure( ffqcFile, fileHighRes=ffqcHFile, exportId="FIGURE_FASTQC",
                      "This figure shows the quality per nucleotide for each sample.
                      Any value above 30 is quite good quality." );
  ################################################
  return(FRfqc)
}

i_create_rnaseqc_figure<-function(path,samples,out){
  require(ggplot2)
  require(Nozzle.R1)

  ############COVERAGE###############

  trqccov<-data.frame()
  for (f in samples){
    print (f)
    src<-paste(path,f,"/qc/rnaseqc/meanCoverageNorm_low.txt",sep="")
    temp<-read.table(src,skip=1)
    temp$type<-"low"
    temp$pos<-1:nrow(temp)
    temp$sample<-substr(f,1,21)
    trqccov<-rbind(trqccov,temp)
    src<-paste(path,f,"/qc/rnaseqc/meanCoverageNorm_medium.txt",sep="")
    temp<-read.table(src,skip=1)
    temp$type<-"medium"
    temp$pos<-1:nrow(temp)
    temp$sample<-substr(f,1,21)
    trqccov<-rbind(trqccov,temp)
    src<-paste(path,f,"/qc/rnaseqc/meanCoverageNorm_high.txt",sep="")
    temp<-read.table(src,skip=1)
    temp$type<-"high"
    temp$pos<-1:nrow(temp)
    temp$sample<-substr(f,1,21)
    trqccov<-rbind(trqccov,temp)

  }

  trqccovFile<-paste(out,"trqccov.txt",sep="")
  write.table(trqccov,trqccovFile,row.names=F,quote=F,sep="\t")

  trqccov$type<-factor(trqccov$type,levels=c("low","medium","high"))
  p<-ggplot(trqccov,aes(pos,log2(V1+0.5),colour=sample)) +
    geom_line()+
    theme_bw(base_size = 12) +
    theme(axis.text.x = element_blank())+
    labs(list(x="gene",y="coverage",colour="Samples"))+
    facet_wrap(~type,nrow=3)

  stdFile="frqccov.jpg"
  highFile= "frqccovH.jpg"
  jpeg( paste( out,stdFile , sep="" ),width=600,height=400,quality=100 );
  print(p)
  dev.off(frqccovFile);
  pdf( paste( out, highFile, sep="" ));
  print(p)
  dev.off();

  # create a figure and make it available for exporting
  FIG <- newFigure( stdFile, fileHighRes=highFile,exportId="FIGURE_COV",
                         "This shows the gene coverage by reads.
                    A good signal is an equal coverage along gene." );
  ################################################
  return(FIG)

}

i_create_gene_coverage<-function(path,samples,out){
  require(ggplot2)
  require(Nozzle.R1)

  ############COVERAGE###############
  trqccov<-data.frame()
  for (f in samples){
    src<-paste(path,f,"/qc/rnaseqc/meanCoverageNorm_low.txt",sep="")
    temp<-read.table(src,skip=1)
    temp$type<-"low"
    temp$pos<-1:nrow(temp)
    temp$sample<-substr(f,1,21)
    trqccov<-rbind(trqccov,temp)
    src<-paste(path,f,"/qc/rnaseqc/meanCoverageNorm_medium.txt",sep="")
    temp<-read.table(src,skip=1)
    temp$type<-"medium"
    temp$pos<-1:nrow(temp)
    temp$sample<-substr(f,1,21)
    trqccov<-rbind(trqccov,temp)
    src<-paste(path,f,"/qc/rnaseqc/meanCoverageNorm_high.txt",sep="")
    temp<-read.table(src,skip=1)
    temp$type<-"high"
    temp$pos<-1:nrow(temp)
    temp$sample<-substr(f,1,21)
    trqccov<-rbind(trqccov,temp)

  }

  trqccovFile<-paste(out,"trqccov.txt",sep="")
  write.table(trqccov,trqccovFile,row.names=F,quote=F,sep="\t")

  trqccov$type<-factor(trqccov$type,levels=c("low","medium","high"))
  p<-ggplot(trqccov,aes(pos,log2(V1+0.5),colour=sample)) +
    geom_line()+
    theme_bw(base_size = 12) +
    theme(axis.text.x = element_blank())+
    labs(list(x="gene",y="coverage",colour="Samples"))+
    facet_wrap(~type,nrow=3)

  stdFile="frqccov.jpg"
  highFile="frqccovH.pdf"
  jpeg( paste( out ,stdFile , sep="" ) ,width=600,height=400,quality=100 );
  print(p)
  dev.off();
  pdf( paste( out , highFile, sep="" ) );
  print(p)
  dev.off();

  # create a figure and make it available for exporting
  FIG <- newFigure( stdFile, fileHighRes=highFile,exportId="FIGURE_COV",
                         "This shows the gene coverage by reads.
                    A good signal is an equal coverage along the gene." );
  ################################################

  return(FIG)
}

i_create_count_top<-function(path,ssamples,out){
  require(Nozzle.R1)

  mainfolder<-dir(path,pattern="project")
  ############COUNTS###############
  counts<-read.table(paste(sep="",path,mainfolder,"/annotated_combined.counts"),header=T,sep="\t")
  names(counts)<-substr(names(counts),1,21)
  counts$mean<-rowMeans(counts[,2:(ncol(counts)-1)])
  counts<-counts[order(-counts$mean),]
  counts<-counts[counts$mean>0,]

  countsFile<-"counts.txt"
  write.table(counts,paste(out,countsFile,sep=""),row.names=F,quote=F,sep="\t")
  short<-counts[1:30,]
  TAB <- newTable(short , file=countsFile, exportId="TABLE_COUNTS",
                      "Top genes" );

  return(TAB)
}

i_create_distribution_counts<-function(path,samples,out,condition){
  require(ggplot2)
  require(Nozzle.R1)
  require(reshape)
  require(DESeq2)

  mainfolder<-dir(path,pattern="project")

  counts<-read.table(paste(sep="",path,mainfolder,"/combined.counts"),header=T,row.names=1,sep="\t")
  names(counts)<-substr(names(counts),1,21)
  keep<-rowSums(counts>0)>=1
  counts<-counts[keep,]
  cshape<-suppressMessages(melt(counts,by=1))
  p<-ggplot(cshape,aes(factor(variable),log2(value+0.1))) +
    geom_boxplot() +
    theme_bw(base_size = 12) +
    theme(axis.text.x = element_blank()) +
    labs(list(x="",y="log2(raw counts)"))

  stdFile= "cbox.jpg"
  jpeg( paste( out,stdFile, sep="" ),width=600,height=400,quality=100 );
  print(p)
  dev.off();

  # create a figure and make it available for exporting
  FIG1 <- newFigure( stdFile,exportId="FIGURE_COUNTS",
                    "Distribution of raw counts" );

  ma<-counts
  design<-data.frame(row.names=names(ma),condition=condition)
  dse<-DESeqDataSetFromMatrix(countData = ma,
                              design=~condition,
                              colData=design)
  rld<-rlogTransformation(dse)
  p<-ggplot(suppressMessages(melt(assay(rld))),aes(factor(X2),value)) +
    geom_boxplot() +
    theme_bw(base_size = 12) +
    theme(axis.text.x = element_blank())+
    labs(list(x="",y="log2(normalized counts)"))

  stdFile= "cboxnorm.jpg"
  jpeg( paste( out,stdFile, sep="" ),width=600,height=400,quality=100 );
  print(p)
  dev.off();

  # create a figure and make it available for exporting
  FIG2 <- newFigure( stdFile,exportId="FIGURE_COUNTSNORM",
                           "Distribution of normalized counts" );

  cum<-assay(rld)
  cum<-as.data.frame(apply(cum,2,sort))
  cum$pos<-nrow(cum):1
  cum<- suppressMessages(melt(cum,id.vars="pos"))
  p<-ggplot(cum,aes(pos,value,colour=variable)) +
    geom_point() +
    theme_bw(base_size = 12) +
    labs(list(x="sorted genes",y="log2(normalized counts)"))

  sdtFile="fcum.jpg"
  highFile="fcumH.pdf"
  jpeg(paste( out,sdtFile , sep="" ) ,width=600,height=400,quality=100 );
  print(p)
  dev.off()
  pdf( paste( out, highFile, sep="" ) )
  print(p)
  dev.off()

  # create a figure and make it available for exporting
  FIG3<- newFigure( sdtFile, fileHighRes=highFile,exportId="FIGURE_CUM",
                       "Cumulative expression by sample." )

  p<-mds(assay(rld),condition)
  p<-p + theme_bw(base_size = 12) +
    geom_text(aes(one,two,label=label))+
    scale_color_brewer(palette="Set1")

  stdFile= "mds.jpg"
  jpeg( paste( out,stdFile, sep="" ),width=600,height=400,quality=100 )
  print(p)
  dev.off()

  # create a figure and make it available for exporting
  FIG4 <- newFigure( stdFile,exportId="FIGURE_MDS",
                       "Multidimensional scaling plot" );

  return(list(FIG1,FIG2,FIG3,FIG4))

}

i_create_mapping_stats<-function(path,samples,out){
  require(Nozzle.R1)

  ############ALIGNED#################
  trqc<-data.frame()
  for (f in samples){
    src<-paste(path,f,"/qc/rnaseqc/metrics.tsv",sep="")
    rqc<-read.table(src,sep="\t",skip=1)
    rqc$sample<-substr(f,1,21)
    trqc<-rbind(trqc,rqc)
  }

  trqc<-trqc[,c(1,19,17,8,20,23,7,38)]
  names(trqc)<-c("Sample","Total","Mapped","Mapping rate","rRNA","transcripts","Exonic rate","Bases MM rate")

  tFile<-"trqc.txt"
  write.table(trqc,paste(out,tFile,sep=""),row.names=F,quote=F,sep="\t")

  # create a table and make it available for exporting
  TAB <- newTable( trqc, file=tFile, exportId="TABLE_RNASEQC",
                    "Stats of mapping reads " );

  ################################################

  return(TAB)

}


#!/usr/bin/env Rscript
# Load required libraries
#library(xcms)

# Set options
options(stringAsfactors=FALSE, useFancyQuotes=FALSE)
stack_height <- 0.67
max_ppm_shift <- 1
plot_type <- "simple"

# Read command line arguments
args <- commandArgs(trailingOnly=TRUE)

if (length(args) == 0) stop("No parameters are given!\n")

# Process command line arguments
for (arg in args) {
	command <- strsplit(x=arg, split="=")[[1]][1]
	value <- strsplit(x=arg, split="=")[[1]][2]

	if (command == "data_matrix_file") {
		data_matrix_file <- as.character(value)
		print(paste("Data matrix file is:", data_matrix_file))
	}

	if (command == "pdf_output") {
		pdf_output <- as.character(value)
		print(paste("PDF output set to:", pdf_output))
	}

	if (command == "stack_height") {
		stack_height <- as.numeric(value)
		print(paste("Setting graphical height of stack to:", stack_height))
	}

	if (command == "max_ppm_shift") {
		max_ppm_shift <- as.numeric(value)
		print(paste("Setting maximum ppm shift to:", max_ppm_shift))
	}

	if (command == "plot_type") {
		plot_type <- as.character(value)
		print(paste("Setting plot type to:", plot_type))
	}
}

# Check required command line arguments
if (is.na(data_matrix_file) | is.na(pdf_output)) stop("Please specify at least the command line options data_matrix_file and pdf_output!\n")

# MAIN
# Following code is adapted from Daniel Jacob

matfile <- data_matrix_file

specmat <- read.table(matfile, header=T, sep="\t", stringsAsFactors=FALSE)

ppm <- as.numeric(gsub('_', '.', gsub('B', '', colnames(specmat)[-1])))
ppm_range <- c(min(ppm), max(ppm))

cols <- rainbow(dim(specmat)[1])

# Prepare stacked Plot
Ymax <- max(specmat[,-1])        # intensity limit
ppm_lim <- ppm_range             # c(ppm min, ppm max)
K <- stack_height                # Graphical height of the stack (0 .. 1)
ppm_shift_max <- max_ppm_shift   # Maximum ppm shift

selspec <- seq(from=dim(specmat)[1], to=2, by=-1)

# Actual stacked plot
pdf(pdf_output, encoding="ISOLatin1", pointsize=10, width=5, height=5, family="Helvetica")

plot(cbind(ppm, t(specmat[1,-1])), xlim=rev(ppm_lim), ylim=c(0,Ymax), xlab="ppm", ylab="Intensity", main="NMR FID stacked plot", type="h", col="blue")

if (plot_type=="perspective") {
	for (ind in c(1:length(selspec))) {
		i <- selspec[ind]
		dppm <- ppm_shift_max*(i-1)/dim(specmat)[1]
		y_offset <- (i-1)*K*Ymax/dim(specmat)[1]
		segments(ppm + dppm, y_offset, ppm + dppm, t(specmat[i,-1]) + y_offset, col=cols[(ind-1) %% length(cols) + 1])
	}
} else {
	for (ind in c(1:length(selspec))) {
		i <- selspec[ind]
		segments( ppm, (i-1)*K*Ymax/dim(specmat)[1], ppm, t(specmat[i,-1]) + (i-1)*K*Ymax/dim(specmat)[1], col=cols[(ind-1) %% length(cols) + 1] )
	}
}

dev.off()


#' torr2p
#'
#' Convert pressure in torr in other unit.
#'
#' @param vp numeric pressure in torr.
#' @param vpunits character  Units available are : "atm","bar","Pa","mbar" (default),"dyne/cm2" and "psi".
#' @return pressure fro torr
#'
#'
#' @author Istituto per la Bioeconomia CNR Firenze Italy  Alfonso Crisci \email{alfonso.crisci@@ibe.cnr.it}
#' @keywords  torr2p
#'
#' @export
#'
#'
#'
#'

torr2p=function(vp,vpunits="mbar") {
                         if ( length(vp)>1) {vp=rep(vpunits,length(vp))}
                         ct$assign("vp", as.array(vp))
                         ct$assign("vpunits", as.array(vpunits))
                         ct$eval("var res=[]; for(var i=0, len=vp.length; i < len; i++){ res[i]=torr2p(vp[i],vpunits[i])};")
                         res=ct$get("res")
                         return(ifelse(res==9999,NA,res))
}


#' Compute the coordinates of the edges of an eddy
#'
#' @param x An "eddy" object as obtained using \code{\link{eddies}}.
#' @param ... arguments to be passed to \code{\link{ell.coords}}. Setting \code{n}
#' to a lower level speeds up computing time.
#' @keywords internal
#' @export
eddies_edges <- function(x, ...) {
  pix_size <- options()$eddies.pixel_size / 2
  out <- Map(ell.coords,
             x0 = x$Lon, y0 = x$Lat,
             major = unlist(x$MajorAxisLength) * pix_size * unlist(x$Extent),
             minor = unlist(x$MinorAxisLength) * pix_size * unlist(x$Extent),
             agl = unlist(x$Orientation) * pi /180, ...)
  out <- lapply(out, as.data.frame)
  nr <- median(sapply(out, nrow))
  out <- rblist(out)
  out$no_eddy <- rep(seq_along(x$Lat), each = nr)
  out
}

#' Match a location to the edges of an eddy
#'
#' @param lon location longitude
#' @param lat location latitude
#' @param edges_coords the coordinates of the edges of the eddies in \code{x}.
#' This arguments allows to speed up calculation when this function is called
#' from \code{\link{match_track2eddies}} and the same eddies are considered again
#' and again. See \code{\link{eddies_edges}}.
#' @inheritParams eddies_edges
#' @seealso \code{\link{match_track2eddies}}
#' @export
match_loc2eddies <- function(x, lon, lat, edges_coords = NULL, ...) {
  stopifnot(require("fields"))
  df <- na.omit(data.frame(lon, lat))
  if (nrow(df) == 0) return(NULL)

  # Get coordinates of the eddies' edges
  if (is.null(edges_coords)) {
    edges_coords <- eddies_edges(x, ...)
  }

  # Find closest match to location
  dist_mat <- as.vector(rdist.earth(df, edges_coords, miles = FALSE))
  mtch_idx <- which.min(dist_mat) %else% return(NULL)

  df$edd.dist <- dist_mat[mtch_idx]
  df$edd.edge_lon <- edges_coords[mtch_idx, 1]
  df$edd.edge_lat <- edges_coords[mtch_idx, 2]
  df$no_eddy <- edges_coords[mtch_idx, 3]
  df[ , paste0("edd.", names(x))] <- x[df$no_eddy, ]

  df
}

#' Match a track to the edges of an eddy
#'
#' @param track a set of locations. First column is longitude and second is latitude.
#' @param lon locations longitude
#' @param lat locations latitude
#' @inheritParams eddies_edges
#' @details If \code{lon} and \code{lat} are non-null then \code{track} is ignored.
#' @seealso \code{\link{match_loc2eddies}}
#' @export
match_track2eddies <- function(x, track, lon = NULL, lat = NULL, ...) {
  if (!missing(track) & is.null(lon) & is.null(lat)) {
    lon <- track[ , 1]
    lat <- track[ , 2]
  }
  edgs <- eddies_edges(x, ...)
  out <- Map(match_loc2eddies, list(x), lon, lat, list(edges_coords = edgs))
  out <- rblist(out)
  merge(data.frame(lon = lon, lat = lat), out, by = c("lon", "lat"), all = TRUE)
}


#===============================================================================
#  File:    08-treatment-group-comparisons.R
#  Date:    Feb 3, 2021
#  Purpose: replicate appendix analyses: Section 7
#  Data In:
#           ./data/survey_data.csv
#           ./data/daily_pulse_data.csv
#===============================================================================

# PACKAGES
#===============================================================================
library(tidyverse)
library(estimatr)
library(glmnet)
library(powerLATE)
library(haven)
source('code/functions.r')

# DATA
#===============================================================================
pulse <- read_csv("data/daily_pulse_data.csv")
svy <- read_csv("data/survey_data.csv")

# Dropping observations where treatment is missing
svy <- svy[!is.na(svy$W3_PATA306_treatment_w3),]

# Analysis ----------------------------------------------------------------

# https://declaredesign.org/blog/biased-fixed-effects.html

vars <- c("party7", "age", "agesq", "female", "raceeth", "educ",
          "ideo", "income", "employ", "state", "polint", "freq_tv", "freq_np",
          "freq_rad", "freq_net", "freq_disc", "log_news_pre", "diet_mean_pre")

# For all following analysis, HuffPost: treatment, FoxNews:control

########################################################
### Issue H1 & H2: issue opinions
########################################################

trt <- "HuffPost"
dv <- "issue_scale"
dv_pre <- "issue_scale_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv= "H1 & H2: Conservatism")

########################################################
### Issue H3 & H4: immigration attitudes
########################################################

trt <- "HuffPost"
dv <- "imm_issue_scale"
dv_pre <- "issue_scale_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control, more_vars = "policy14_imm_pre")
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv = "H3 & H4: Pro-immigration")

########################################################
### Issue RQ1: heterogeneous treatment effects
########################################################

# H1 & H2
heterogeneous_effect2(dv = "issue_scale", dv_pre = "issue_scale_pre",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "issue_scale", dv_pre = "issue_scale_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")
# H3 & H4
heterogeneous_effect2(dv = "imm_issue_scale", dv_pre = "issue_scale_pre",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "imm_issue_scale", dv_pre = "issue_scale_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### News H1: news browsing, 1 week
########################################################

trt <- "HuffPost"
dv <- "log_cons_1w"
dv_pre <- "log_cons_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control, more_vars = "log_total_pre")
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv = "H1: Cons. news (1 week)")

########################################################
### News H2: news browsing, 1 week
########################################################

trt <- "HuffPost"
dv <- "log_lib_1w"
dv_pre <- "log_lib_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control, more_vars = "log_total_pre")
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv = "H2: Lib. news (1 week)")

########################################################
### News H1: news browsing, 4 weeks
########################################################

trt <- "HuffPost"
dv <- "log_cons_4w"
dv_pre <- "log_cons_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt,  control = control, more_vars = "log_total_pre")
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv = "H1: Cons. news (4 weeks)")

########################################################
### News H2: news browsing, 4 week
########################################################

trt <- "HuffPost"
dv <- "log_lib_4w"
dv_pre <- "log_lib_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control, more_vars = "log_total_pre")
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv = "H2: Lib. news (4 weeks)")

########################################################
### News H1: news browsing, 6 weeks
########################################################

trt <- "HuffPost"
dv <- "log_cons_6w"
dv_pre <- "log_cons_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control, more_vars = "log_total_pre")
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv = "H1: Cons. news (6 weeks)")

########################################################
### News H2: news browsing, 6 week
########################################################

trt <- "HuffPost"
dv <- "log_lib_6w"
dv_pre <- "log_lib_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control, more_vars = "log_total_pre")
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv = "H2: Lib. news (6 weeks)")

########################################################
### News RQ1: heterogeneous treatment effects by party ID / ideology
########################################################

# H1: n/s
heterogeneous_effect2(dv = "log_cons_1w", dv_pre = "log_cons_pre",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "log_cons_1w", dv_pre = "log_cons_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")
# H2: n/s
heterogeneous_effect2(dv = "log_lib_1w", dv_pre = "log_lib_pre",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "log_lib_1w", dv_pre = "log_lib_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### News RQ2: heterogeneous treatment effects by pre-treatment news consumption
########################################################

# H1: n/s
heterogeneous_effect2(dv = "log_cons_1w", dv_pre = "log_hp_pre",
                      moderator = "log_fn_pre", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "log_cons_1w", dv_pre = "log_fn_pre",
                      moderator = "log_fn_pre", trt = "HuffPost", control = "FoxNews")

# H2: +
heterogeneous_effect2(dv = "log_lib_1w", dv_pre = "log_hp_pre",
                      moderator = "log_hp_pre", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "log_lib_1w", dv_pre = "log_fn_pre",
                      moderator = "log_hp_pre", trt = "HuffPost", control = "FoxNews")

########################################################
### Twitter H1: conservative link shares
########################################################

# creating Twitter variables
svy$log_cons_links <- log(svy$cons_links+1)
svy$log_cons_links_pre <- log(svy$cons_links_pre+1)
svy$log_total_links_pre <- log(svy$total_links_pre+1)

trt <- "HuffPost"
dv <- "log_cons_links"
dv_pre <- "log_cons_links_pre"
control <- "FoxNews"

# using chrome vs others as blocks here
svy$chrome <- ifelse(svy$W3_Browser_treatment_w3=="Chrome", 1, 0)

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = chrome,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control, blocks="chrome", more_vars = "log_total_links_pre")
compute_proportion_missing_covars(itt)

format_latex3(dim, itt, dv = "H1: Tweets w/cons. links")

########################################################
### Twitter H2: liberal link shares
########################################################

# creating Twitter variables
svy$log_lib_links <- log(svy$lib_links+1)
svy$log_lib_links_pre <- log(svy$lib_links_pre+1)
svy$log_total_links_pre <- log(svy$total_links_pre+1)

trt <- "HuffPost"
dv <- "log_lib_links"
dv_pre <- "log_lib_links_pre"
control <- "FoxNews"

# using chrome vs others as blocks here
svy$chrome <- ifelse(svy$W3_Browser_treatment_w3=="Chrome", 1, 0)
## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = chrome,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control, blocks="chrome", more_vars = "log_total_links_pre")

format_latex3(dim, itt, dv = "H2: Tweets w/lib. links")

########################################################
### Twitter H3: follow conservative news sources
########################################################

# creating Twitter variables
svy$log_follows_cons_media <- log(svy$follows_cons_media+1)
svy$log_follows_elites <- log(svy$follows_elites+1)

trt <- "HuffPost"
dv <- "log_follows_cons_media"
control <- "FoxNews"

# using chrome vs others as blocks here
svy$chrome <- ifelse(svy$W3_Browser_treatment_w3=="Chrome", 1, 0)
## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = chrome,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control, blocks="chrome", more_vars = "log_follows_elites")

format_latex3(dim, itt, dv = "H3: Follow cons. media")

########################################################
### Twitter H4: follow liberal news sources
########################################################

# creating Twitter variables
svy$log_follows_lib_media <- log(svy$follows_lib_media+1)
svy$log_follows_elites <- log(svy$follows_elites+1)

trt <- "HuffPost"
dv <- "log_follows_lib_media"
control <- "FoxNews"

# using chrome vs others as blocks here
svy$chrome <- ifelse(svy$W3_Browser_treatment_w3=="Chrome", 1, 0)
## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = chrome,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control, blocks="chrome", more_vars = "log_follows_elites")

format_latex3(dim, itt, dv = "H4: Follow lib. media")

########################################################
### Affective H1a & H2a: Democrats, HuffPost treatment
########################################################

# creating new variables
svy$W4_PATA450a_new <- zap_labels(svy$W4_PATA450a)
svy$W2_PATA2_3_Dem_new <- zap_labels(svy$W2_PATA2_3_Dem)

trt <- "HuffPost"
dv <- "W4_PATA450a_new" # FT Democrats
dv_pre <- "W2_PATA2_3_Dem_new"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1a \\& H2a: Dem. thermometer")

########################################################
### Affective H1b & H2b: Republicans, HuffPost treatment
########################################################

# creating new variables
svy$W4_PATA450b_new <- zap_labels(svy$W4_PATA450b)
svy$W2_PATA2_3_Rep_new <- zap_labels(svy$W2_PATA2_3_Rep)

trt <- "HuffPost"
dv <- "W4_PATA450b_new" # FT Republicans
dv_pre <- "W2_PATA2_3_Rep_new"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1b \\& H2b: Rep. thermometer")

########################################################
### Affective H1c & H2c: Trump supporters, HuffPost treatment
########################################################

# creating new variables
svy$W4_PATA450e_new <- zap_labels(svy$W4_PATA450e)
svy$W2_PATA2_6_new <- zap_labels(svy$W2_PATA2_6)

trt <- "HuffPost"
dv <- "W4_PATA450e_new" # Trump social distance
dv_pre <- "W2_PATA2_6_new"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1c \\& H2c: Trump distance")

########################################################
### Polarization RQ1: heterogeneous treatment effects by party ID / ideology
########################################################

# H1a
heterogeneous_effect2(dv = "W4_PATA450a_new", dv_pre = "W2_PATA2_3_Dem_new",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "W4_PATA450a_new", dv_pre = "W2_PATA2_3_Dem_new",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")
# H1b
heterogeneous_effect2(dv = "W4_PATA450b_new", dv_pre = "W2_PATA2_3_Rep_new",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "W4_PATA450b_new", dv_pre = "W2_PATA2_3_Rep_new",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

# H1c
heterogeneous_effect2(dv = "W4_PATA450e_new", dv_pre = "W2_PATA2_6_new",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "W4_PATA450e_new", dv_pre = "W2_PATA2_6_new",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### Perceived H1a & H1b: HuffPost treatment
########################################################

trt <- "HuffPost"
dv <- "ppol"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1a \\& H1b: Perceived polarization")

########################################################
### Perceived RQ1: heterogeneous treatment effects by party ID / ideology
########################################################

# H1a & H1b
heterogeneous_effect2(dv = "ppol",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "ppol",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### Agenda H1 & H2: HuffPost treatment
########################################################

trt <- "HuffPost"
dv <- "agenda_lean"
dv_pre <- "agenda_lean_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1 \\& H2: Agenda setting")

########################################################
### Agenda RQ1: heterogeneous treatment effects by party ID / ideology
########################################################

# H1 & H2
heterogeneous_effect2(dv = "agenda_lean", dv_pre = "agenda_lean_pre",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "agenda_lean", dv_pre = "agenda_lean_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### Agenda RQ2: heterogeneous treatment effects by pre-treatment habits
########################################################

# H1
heterogeneous_effect2(dv = "agenda_lean", dv_pre = "log_fn_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")
# H2
heterogeneous_effect2(dv = "agenda_lean", dv_pre = "log_hp_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### Approval H1a & H2a: President
########################################################

trt <- "HuffPost"
dv <- "trump_approve"
dv_pre <- "trump_approve_pre"
control = "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1a \\& H2a: Pres. approval")

########################################################
### Approval H1b & H2b: Republicans in Congress [Congress control preference]
########################################################

trt <- "HuffPost"
dv <- "cong_rep_approve"
dv_pre <- "cong_rep_approve_pre"
control = "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1b \\& H2b: Rep. Cong. pref.")

########################################################
### Approval H1c & H2c: Democrats in Congress [Congress control preference]
########################################################

trt <- "HuffPost"
dv <- "cong_dem_approve"
dv_pre <- "cong_dem_approve_pre"
control = "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1c \\& H2c: Dem. Cong. pref.")

########################################################
### Approval RQ1: heterogeneous treatment effects by party ID / ideology
########################################################

# H1a & H2a
heterogeneous_effect2(dv = "trump_approve", dv_pre = "trump_approve_pre",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "trump_approve", dv_pre = "trump_approve_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")
# H1b & H2b
heterogeneous_effect2(dv = "cong_rep_approve",
                      dv_pre = "cong_rep_approve_pre",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "cong_rep_approve",
                      dv_pre = "cong_rep_approve_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")
# H1c & H2c
heterogeneous_effect2(dv = "cong_dem_approve",
                      dv_pre = "cong_dem_approve_pre",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "cong_dem_approve",
                      dv_pre = "cong_dem_approve_pre",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### Turnout RQ1a: Huff Post
########################################################

trt <- "HuffPost"
dv <- "vote_after"
dv_pre <- "vote_pre"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, dv_pre = dv_pre, trt = trt, control = control)

format_latex3(dim, itt, dv = "RQ1a: Turnout")

########################################################
### Media Trust H1a & H2a: Trust in Fox News for HuffPost treatment
########################################################

trt <- "HuffPost"
dv <- "trust_fox"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1a \\& H2a: Trust in Fox News")

########################################################
### Media Trust H1b & H2b: Trust in HuffPost for FN treatment
########################################################

trt <- "HuffPost"
dv <- "trust_hp"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1b \\& H2b: Trust in HuffPost")

########################################################
### Media Trust H3a & H3b: Media liberal bias for HP treatment
########################################################

trt <- "HuffPost"
dv <- "bias"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H3a \\& H3b: Liberal media bias")

########################################################
### Media Trust RQ1
########################################################

# H1a & H2a
heterogeneous_effect2(dv = "trust_fox",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "trust_fox",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")
# H1b & H2b
heterogeneous_effect2(dv = "trust_hp",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "trust_hp",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")
# H3a & H3b
heterogeneous_effect2(dv = "bias",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "bias",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### Media Trust H4a: Media trust for HP group
########################################################

trt <- "HuffPost"
dv <- "trust"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H4a: Media trust")

########################################################
### Media Trust H4a: Media trust for HP group (w7)
########################################################

trt <- "HuffPost"
dv <- "trust_w7"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H4a: Media trust (w7)")

########################################################
### Factual knowledge H1: % foreign born
########################################################

trt <- "HuffPost"
dv <- "W4_PATA430a"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1: \\% foreign born")

########################################################
### Factual knowledge H2: % unemployment rate
########################################################

trt <- "HuffPost"
dv <- "W4_PATA430b"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H2: \\% unemployment")

########################################################
### Factual knowledge RQ1a & RQ1b: news reception for HP group
########################################################

trt <- "HuffPost"
dv <- "event_mokken"
dv_pre <- "event_pre_mokken"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))
## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, dv_pre = dv_pre, control = control)

format_latex3(dim, itt, dv = "RQ1a \\& RQ1b: Event knowledge")

########################################################
### Knowledge RQ2
########################################################

# H1
heterogeneous_effect2(dv = "W4_PATA430a",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "W4_PATA430a",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

# H2
heterogeneous_effect2(dv = "W4_PATA430b",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "W4_PATA430b",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

# RQ1a & RQ1b
heterogeneous_effect2(dv = "event_mokken", dv_pre = "event_pre_mokken",
                      moderator = "party7", trt = "HuffPost", control = "FoxNews")
heterogeneous_effect2(dv = "event_mokken", dv_pre = "event_pre_mokken",
                      moderator = "ideo", trt = "HuffPost", control = "FoxNews")

########################################################
### Factual knowledge RQ1b: immigration knowledge for HP group
########################################################

svy$know_immig_w2_TRUE <- svy$W2_PATA2_10_m_3 == 1
svy$know_immig_w4_TRUE <- svy$W4_PATA409_1 == 1

trt <- "HuffPost"
dv <- "know_immig_w4_TRUE"
dv_pre <- "know_immig_w2_TRUE"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, dv_pre = dv_pre, control = control)

format_latex3(dim, itt, dv = "RQ1b: Immigration event knowledge")

########################################################
### Election prediction H1
########################################################

# Subjects assigned to the Fox News treatment...

# H1a & H2a: ... will be more likely to predict the Republican Party as winning the control of the House of Representatives
trt <- "HuffPost"
dv <- "predicted_house_winner_gop"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1a \\& H2a: \\ GOP House winner")

# H1b & H2b: ... will predict a higher House vote share for the Republican Party

trt <- "HuffPost"
dv <- "predicted_vote_share_gop"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1b \\& H2b: \\ GOP House vote share")

# H1c & H2c: ... will be more likely to predict the Republican candidate to win the race for the House seat in their Congressional district

trt <- "HuffPost"
dv <- "predicted_district_winner_gop"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1c \\& H2c: \\ GOP House district winner")

########################################################
### Election prediction H3
########################################################

# (a) Pre-treatment Republican partisans in the Fox News treatment group and (b) pre-treatment Democratic partisans in the HuffPost treatment group will report a higher average certainty in their predictions than subjects assigned to the control group.
## We revised this hypothesis to be able to compare the HuffPost and the FoxNews groups directly
## Pre-treatment Democratic partisans in the HuffPost treatment group will report a higher average certainty in their predictions than pre-treatment Republican partisans in the Fox News treatment group.

svy_sub <- filter(svy, (W3_PATA306_treatment_w3 == "FoxNews" & partylean == "Republican") | (W3_PATA306_treatment_w3 == "HuffPost" & partylean == "Democrat"))

trt <- "HuffPost"
dv <- "prediction_certainty"
control <- "FoxNews"

## DIM WITHOUT BLOCKING (data sparsity issue)
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            #blocks = W3_Browser_treatment_w3,
                            data = svy_sub,
                            condition1 = control,
                            condition2 = trt))

## ITT WITHOUT Lin's covariate adjustment (data sparsity issue)
form <- as.formula("prediction_certainty ~ W3_PATA306_treatment_w3")
# itt <- lm_robust(form, data = dplyr::filter(svy_sub, W3_PATA306_treatment_w3 != "Control"))
itt <- lm_robust(form, svy_sub)

format_latex3(dim, list(itt), dv = "H3a \\& H3b: \\ Prediction certainty")

########################################################
### Election prediction H4
########################################################

# (a) # Pre-treatment Republican partisans in the HuffPost treatment group and (b) pre-treatment Democratic partisans in the Fox News treatment group will report a lower average certainty in their predictions than subjects assigned to the control group.
## We revised this hypothesis to be able to compare the HuffPost and the FoxNews groups directly
## Pre-treatment Republican partisans in the HuffPost treatment group will report a higher average certainty in their predictions than pre-treatment Democratic partisans in the Fox News treatment group.

svy_sub <- filter(svy, (W3_PATA306_treatment_w3 == "FoxNews" & partylean == "Democrat") | (W3_PATA306_treatment_w3 == "HuffPost" & partylean == "Republican"))

trt <- "HuffPost"
dv <- "prediction_certainty"
control <- "FoxNews"

## DIM WITHOUT BLOCKING (data sparsity issue)
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            #blocks = W3_Browser_treatment_w3,
                            data = svy_sub,
                            condition1 = control,
                            condition2 = trt))

## ITT WITHOUT Lin's covariate adjustment (data sparsity issue)
form <- as.formula("prediction_certainty ~ W3_PATA306_treatment_w3")
# itt <- lm_robust(form, data = dplyr::filter(svy_sub, W3_PATA306_treatment_w3 != "FoxNews"))

itt <- lm_robust(form, svy_sub)

format_latex3(dim, list(itt), dv = "H4a \\& H4b: \\ Prediction certainty")

########################################################
### Bombing H1
########################################################

# Bombing H1a: HuffPost vs. Fox News treatment

trt <- "HuffPost"
dv <- "W4_PATA462_a"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H1a \\& H1b: Mail-bombing is a false flag")

########################################################
### Bombing H2
########################################################

# Bombing H2a: HuffPost vs. Fox News treatment

trt <- "HuffPost"
dv <- "W4_PATA462_b"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H2a \\& H2b: Media is accountable")

########################################################
### Bombing H3a
########################################################

# Bombing H3a: HuffPost vs. Fox News treatment

trt <- "HuffPost"
dv <- "W4_PATA462_c"
control <- "FoxNews"

## DIM
(dim <- difference_in_means(formula(paste0(dv, " ~ W3_PATA306_treatment_w3")),
                            blocks = W3_Browser_treatment_w3,
                            data = svy,
                            condition1 = control,
                            condition2 = trt))

## ITT, with Lin's covariate adjustment
itt <- run_model(dv = dv, trt = trt, control = control)

format_latex3(dim, itt, dv = "H3a \\& H3b: Trump is accountable")


error_check <- function(func_args, model_type){

  x <- all.vars(func_args$formula[[3]]) # covariate names
  y <- all.vars(func_args$formula[[2]]) # outcome name

  nparams <- length(x) + 1

  ## check init_k - initial number of clusters
  if(class(func_args$init_k)!='numeric' | length(func_args$init_k)!=1 ){
    stop("ERROR: init_k must be length 1 numeric integer such that 0 < init_k <= n .")
  }else if(func_args$init_k<0){
    stop("ERROR: init_k must be 0 < init_k <= n .")
  }else if( func_args$init_k %% 1 !=0 ){
    stop('ERROR: init_k must be an integer.')
  }else if(  func_args$init_k > nrow(func_args$d_train) ){
    stop('ERROR: init_k cannot be greater than sample size.')
  }

  ## Check burnin and iter settings
  if( ! (func_args$burnin %% 1 ==0 & func_args$iter %% 1 ==0) ){
    stop("ERROR: burnin and iter must be integers.")
  }else if( func_args$burnin < 0 | func_args$iter < 0){
    stop("ERROR: burnin and iter must be positive.")
  }else if(func_args$burnin==0){
    warning("Warning (non-fatal): burnin==0. Should specify 0 < burnin < iter.")
  }else if(func_args$iter==0){
    stop('ERROR: iter==0. Should specify 0 < burnin < iter.')
  }else if(! func_args$burnin < func_args$iter){
    stop("ERROR: iter must be greater than burnin.")
  }

  # check training data
  if( !is.data.frame(func_args$d_train) ){
    stop("ERROR: d_train is not a data.frame object.")
  }else if( nrow(func_args$d_train)==0 ){
    stop("ERROR: d_train has no rows.")
  }else if( ncol(func_args$d_train)==0 ){
    stop("ERROR: d_train has no columns.")
  }else if( nrow(func_args$d_train)<=2  ){
    stop('ERROR: d_train should contain at least 3 observations.')
  }else if( nrow(func_args$d_train) < func_args$init_k ){
    stop("ERROR: nrow(d_train) < init_k.")
  }else if( max( ! x %in% colnames(func_args$d_train) )==1 ){
    stop("ERROR: Some covariate specified in formula are not in d_train")
  }else if(  max(apply(func_args$d_train[,x, drop=F], 2, class)!="numeric") ){
    stop("ERROR: All columns of d_train must be numeric.")
  }

  for( var in x ){
    unvals <- unique(func_args$d_train[,var])
    nu <- length(unvals)
    if( nu==2 & sum(unvals %in% c(0,1)) != 2 ){
      stop(paste0('ERROR: ',var," in d_train is binary, but not coded as 0-1"))
    }
  }

  # check testing data set if supplied.
  if(! is.null(func_args$d_test) ){

    if( !is.data.frame(func_args$d_test) ){
      stop("ERROR: d_test is not a data.frame object.")
    }else if( nrow(func_args$d_test)==0 ){
      stop("ERROR: d_test has no rows.")
    }else if( ncol(func_args$d_test)==0 ){
      stop("ERROR: d_test has no columns.")
    }else if( max( ! x %in% colnames(func_args$d_test) )==1 ){
      stop("ERROR: Some covariate specified in formula are not in d_test.")
    }else if(  max(apply(func_args$d_test[,x, drop=F], 2, class)!="numeric") ){
      stop("ERROR: All columns of d_test must be numeric.")
    }

  }

  if(!is.null(func_args$beta_prior_mean)){
    if( class(func_args$beta_prior_mean) != 'numeric' ){
      stop('ERROR: beta_prior_mean should be a numeric vector.')
    }else if(length(func_args$beta_prior_mean) != (length(x)+1) ){
      stop("ERROR: beta_prior_mean should be length p+1, where p is number of covariates.")
    }
  }

  if(!is.null(func_args$beta_prior_var)){
    if( class(func_args$beta_prior_var) != 'numeric' ){
      stop('ERROR: beta_prior_var should be a numeric vector.')
    }else if(length(func_args$beta_prior_var) != (length(x)+1) ){
      stop("ERROR: beta_prior_var should be length p+1, where p is number of covariates.")
    }else if( max(func_args$beta_prior_var<0)==1 ){
      stop("ERROR: beta_prior_var should contain non-zero values since it is a vector of variances.")
    }
  }

  if(length(func_args$beta_var_scale)!=1){
    stop("ERROR: beta_var_scale should be a scalar, length 1 numeric.")
  }else if(class(func_args$beta_var_scale)!='numeric'){
    stop("ERROR: beta_var_scale should be numeric.")
  }else if(func_args$beta_var_scale<0){
    stop("ERROR: beta_var_scale should be >0.")
  }

  if(!is.numeric(func_args$mu_scale) | length(func_args$mu_scale)!=1 | func_args$mu_scale<=0 ){
    stop("ERROR: mu_scale should be length 1 numeric values greater than 0.")
  }

  if(model_type!='fDP'){
    if(!is.numeric(func_args$tau_scale) | length(func_args$tau_scale)!=1 | func_args$tau_scale<=0 ){
      stop("ERROR: tau_scale should be length 1 numeric values greater than 0.")
    }
  }

  if(! class(func_args$d_train[,y]) %in% c('numeric', 'integer') ){
    stop("ERROR: outcome is not numeric")
  }

  ## model-specific checks
  if(model_type=='ZDP'){

    if( length(func_args$phi_y)!=2 ){
      stop("ERROR: phi_y should be a length 2 vector of positive real numbers.")
    }else if( max(func_args$phi_y<0)==1 ){
      stop("ERROR: phi_y contains negative values. Should be positive real numbers")
    }

    if( sum(func_args$d_train[,y]==0)==0 ){
      stop("ERROR: outcome does not contain any zeros. ZDPMix() is inappropriate.")
    }

    if(!is.null(func_args$gamma_prior_mean)){
      if( class(func_args$gamma_prior_mean) != 'numeric' ){
        stop('ERROR: gamma_prior_mean should be a numeric vector.')
      }else if(length(func_args$gamma_prior_mean) != (length(x)+1) ){
        stop("ERROR: gamma_prior_mean should be length p+1, where p is number of covariates.")
      }
    }

    if(!is.null(func_args$gamma_prior_var)){

      if( class(func_args$gamma_prior_var) != 'numeric' ){
        stop('ERROR: gamma_prior_var should be a numeric vector.')
      }else if(length(func_args$gamma_prior_var) != (length(x)+1) ){
        stop("ERROR: gamma_prior_var should be length p+1, where p is number of covariates.")
      }else if( max(func_args$gamma_prior_var<0)==1 ){
        stop("ERROR: gamma_prior_var should contain non-zero values since it is a vector of variances.")
      }

    }

    if( is.null(func_args$prop_sigma_z) ){
      stop('ERROR: Proposal covariance is null.')
    }else if( class(func_args$prop_sigma_z)!='matrix' ){
      stop('ERROR: prop_sigma_z should be a (p+1) X (p+1) numeric covariance matrix')
    }else if(!is.numeric(func_args$prop_sigma_z)){
      stop("ERROR: prop_sigma_z is not numeric.")
    }else if( !isSymmetric(func_args$prop_sigma_z) ){
      stop("ERROR: prop_sigma_z is not symmetric.")
    }else if( det(func_args$prop_sigma_z)==0 ){
      stop("ERROR: prop_sigma_z is not positive definite.")
    }

  }else if(model_type=='NDP'){

    if( length(func_args$phi_y)!=2 ){
      stop("ERROR: phi_y should be a length 2 vector of positive real numbers.")
    }else if( max(func_args$phi_y<0)==1 ){
      stop("ERROR: phi_y contains negative values. Should be positive real numbers")
    }

  }else if(model_type=='fDP'){

    if( !is.null(func_args$phi_y) & length(func_args$phi_y)!=2 ){
      stop("ERROR: phi_y should be a length 2 vector of positive real numbers.")
    }else if( max(func_args$phi_y<0)==1 ){
      stop("ERROR: phi_y contains negative values. Should be positive real numbers")
    }

    if( !is.null(func_args$tau_x) & length(func_args$tau_x)!=2 ){
      stop("ERROR: tau_x should be a length 2 vector of positive real numbers.")
    }else if( max(func_args$phi_y<0)==1 ){
      stop("ERROR: tau_x contains negative values. Should be positive real numbers")
    }

  }else if(model_type=='PDP'){

    uy <- unique(func_args$d_train[,y])
    if(length(uy)!=2){
      stop("ERROR: outcome is not binary.")
    }else if( sum(uy %in% c(0,1))!=2  ){
      stop("ERROR: outcome must be coded as numeric vector with 1 and 0. 1 indicates events.")
    }

    if( is.null(func_args$prop_sigma_b) ){
      stop('ERROR: Proposal covariance is null.')
    }else if( class(func_args$prop_sigma_b)!='matrix' ){
      stop('ERROR: prop_sigma_b should be a (p+1) X (p+1) numeric covariance matrix')
    }else if(!is.numeric(func_args$prop_sigma_b)){
      stop("ERROR: prop_sigma_b is not numeric.")
    }else if( !isSymmetric(func_args$prop_sigma_b) ){
      stop("ERROR: prop_sigma_b is not symmetric.")
    }else if( det(func_args$prop_sigma_b)==0 ){
      stop("ERROR: prop_sigma_b is not positive definite.")
    }

  }

}

error_check_credible_gradient = function(func_args){

  if( class(func_args$x) != 'numeric' ){
    stop("ERROR: x is not numeric")
  }else if( class(func_args$post_draws) != "matrix" ){
    stop("ERROR: post_draws is not a matrix. Should be matrix with dimensions length(x) by #posterior draws ")
  }else if( !is.character(func_args$col_gradient) | !is.character(func_args$col_mean_line) ){
    stop("ERROR: either col_gradient or col_mean_line is not a character vector specifying a particular color")
  }

}


#' Find "closest" servers (geography-wise) from master server list
#'
#' @md
#' @param servers if not `NULL`, then the data frame from [spd_servers()]. If
#'        `NULL`, then [spd_servers()] will be called to retrieve the server list.
#' @param config client configuration retrieved via [spd_config()]. If `NULL` it
#'        will be retrieved
#' @return server list in order of geographic closeness
#' @export
#' @examples \dontrun{
#' spd_closest_servers()
#' }
spd_closest_servers <- function(servers=NULL, config=NULL) {

  if (is.null(config)) config <- spd_config()

  if (is.null(servers)) servers <- spd_servers(config)

  # we don't need great circle for this, just best effort
  idx <- order(sqrt((servers$lat - as.numeric(config$client$lat))^2 +
                      (servers$lng - as.numeric(config$client$lon))^2))

  servers[idx,]

}


*** Reaching Definitions ***
*** InSet Map ***
Node[-1, 0, ifthen.c:2.0-10.1 "int main() {\n   ... }"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[31, 19, ifthen.c:9.2-8 "x = 42"]>
Node[10, 20, ifthen.c:2.11-10.1 "{\n  x = 7; ... }"]
Node[11, 5, ifthen.c:3.2-7 "x = 7"]
Node[17, 7, ifthen.c:4.6-11 "x > 0"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[11, 5, ifthen.c:3.2-7 "x = 7"]>
Node[22, 10, ifthen.c:5.4-7 "x++"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[11, 5, ifthen.c:3.2-7 "x = 7"]>
Node[27, 14, ifthen.c:7.4-7 "x--"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[11, 5, ifthen.c:3.2-7 "x = 7"]>
Node[31, 19, ifthen.c:9.2-8 "x = 42"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[22, 10, ifthen.c:5.4-7 "x++"]>  #  <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[27, 14, ifthen.c:7.4-7 "x--"]>
*** OutSet Map ***
Node[-1, 0, ifthen.c:2.0-10.1 "int main() {\n   ... }"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[31, 19, ifthen.c:9.2-8 "x = 42"]>
Node[10, 20, ifthen.c:2.11-10.1 "{\n  x = 7; ... }"]
Node[11, 5, ifthen.c:3.2-7 "x = 7"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[11, 5, ifthen.c:3.2-7 "x = 7"]>
Node[17, 7, ifthen.c:4.6-11 "x > 0"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[11, 5, ifthen.c:3.2-7 "x = 7"]>
Node[22, 10, ifthen.c:5.4-7 "x++"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[22, 10, ifthen.c:5.4-7 "x++"]>
Node[27, 14, ifthen.c:7.4-7 "x--"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[27, 14, ifthen.c:7.4-7 "x--"]>
Node[31, 19, ifthen.c:9.2-8 "x = 42"] ==> <Entity[kind=VARIABLE, name=x, linkage=EXTERNAL, type=0, definition=1]@Node[31, 19, ifthen.c:9.2-8 "x = 42"]>


# -------------------------------------
# --------- WATER SENSOR DATA ---------
# -------------------------------------

# get sensor data from Chicago Open Data Portal
water_sensors <- read.csv("https://data.cityofchicago.org/api/views/qmqz-2xku/rows.csv?accessType=DOWNLOAD", stringsAsFactors = F, header = T)
sensor_locations <- read.csv("https://data.cityofchicago.org/api/views/g3ip-u8rb/rows.csv?accessType=DOWNLOAD", stringsAsFactors = F, header = T)

# merge sensor data
water_sensors <- merge(water_sensors, sensor_locations, by.x="Beach.Name", by.y="Sensor.Name")

# drop unecessary columns
drop <- c("Sensor.Type","Location")
water_sensors <- water_sensors[,!(names(water_sensors) %in% drop)]

# normalize names
water_sensors$Beach.Name <- substr(water_sensors$Beach.Name,1,nchar(water_sensors$Beach.Name)-6)

# match beaches with water and weather stations
beach_mappings <- read.csv("data/ExternalData/Beach_Mappings.csv", stringsAsFactors=FALSE)
waternames <- setNames(beach_mappings$water, beach_mappings$beach)
beach_readings$Water.Station <- waternames[beach_readings$Client.ID]

# parse dates for water and weather
water_sensors$Date <- substr(water_sensors$Measurement.Timestamp,1,10)
water_sensors$Date<- as.Date(water_sensors$Date, format="%m/%d/%Y")

# Find average, max, and min values for all water and weather metrics per day per senser location
df_min_water_temp <- aggregate(water_sensors$Water.Temperature, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=min)
df_max_water_temp <- aggregate(water_sensors$Water.Temperature, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=max)
df_avg_water_temp <- aggregate(water_sensors$Water.Temperature, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=mean)
water_temp <- join_all(list(df_min_water_temp, df_max_water_temp, df_avg_water_temp), by=c("Group.1", "Group.2"))
names(water_temp) <- c("Water.Station", "Full_date", "Water.Temp.Min", "Water.Temp.Max", "Water.Temp.Mean")

df_min_turbidity <- aggregate(water_sensors$Turbidity, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=min)
df_max_turbidity <- aggregate(water_sensors$Turbidity, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=max)
df_avg_turbidity <- aggregate(water_sensors$Turbidity, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=mean)
turbidity <- join_all(list(df_min_turbidity, df_max_turbidity, df_avg_turbidity), by=c("Group.1", "Group.2"))
names(turbidity) <- c("Water.Station", "Full_date", "Turbidity.Min", "Trubidity.Max", "Turbidity.Mean")

df_min_transducer_depth <- aggregate(water_sensors$Transducer.Depth, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=min)
df_max_transducer_depth <- aggregate(water_sensors$Transducer.Depth, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=max)
df_avg_transducer_depth <- aggregate(water_sensors$Transducer.Depth, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=mean)
transducer_depth <- join_all(list(df_min_transducer_depth, df_max_transducer_depth, df_avg_transducer_depth), by=c("Group.1", "Group.2"))
names(transducer_depth) <- c("Water.Station", "Full_date", "Transducer.Depth.Min", "Transducer.Depth.Max", "Transducer.Depth.Mean")

df_min_wave_height <- aggregate(water_sensors$Wave.Height, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=min)
df_max_wave_height <- aggregate(water_sensors$Wave.Height, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=max)
df_avg_wave_height <- aggregate(water_sensors$Wave.Height, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=mean)
wave_height <- join_all(list(df_min_wave_height, df_max_wave_height, df_avg_wave_height), by=c("Group.1", "Group.2"))
names(wave_height) <- c("Water.Station", "Full_date", "Wave.Height.Min", "Wave.Height.Max", "Wave.Height.Mean")

df_min_wave_period <- aggregate(water_sensors$Wave.Period, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=min)
df_max_wave_period <- aggregate(water_sensors$Wave.Period, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=max)
df_avg_wave_period <- aggregate(water_sensors$Wave.Period, by=list(water_sensors$Beach.Name, water_sensors$Date), FUN=mean)
wave_period <- join_all(list(df_min_wave_period, df_max_wave_period, df_avg_wave_period), by=c("Group.1", "Group.2"))
names(wave_period) <- c("Water.Station", "Full_date", "Wave.Period.Min", "Wave.Period.Max", "Wave.Period.Mean")

water_sensors_daily <- join_all(list(water_temp, turbidity, transducer_depth, wave_height, wave_period), by=c("Water.Station", "Full_date"))

# merge with beach_readings
beach_readings <- merge(beach_readings, water_sensors_daily, all.x = TRUE)

# clean up workspace
rm(list = ls(pattern = "^df_min_"))
rm(list = ls(pattern = "^df_max_"))
rm(list = ls(pattern = "^df_avg_"))
rm(list=c("water_temp", "wave_period", "wave_height", "turbidity", "transducer_depth"))

# clean up order
beach_readings <- beach_readings[order(-as.numeric(beach_readings$Full_date)), c(
  "Client.ID",  "Full_date", "Year", "Date",
  "Laboratory.ID", "Reading.1", "Reading.2",
  "Escherichia.coli", "Units", "Sample.Collection.Time",
  "Weekday", "Month", "Day",
  "Drek_Reading", "Drek_Prediction", "Drek_Worst_Swim_Status",
  "e_coli_geomean_actual_calculated", "elevated_levels_actual_calculated", "Drek_elevated_levels_predicted_calculated",
  "Water.Station", "Water.Temp.Min", "Water.Temp.Max","Water.Temp.Mean",
  "Turbidity.Min","Trubidity.Max", "Turbidity.Mean",
  "Transducer.Depth.Min","Transducer.Depth.Max", "Transducer.Depth.Mean",
  "Wave.Height.Min","Wave.Height.Max", "Wave.Height.Mean",
  "Wave.Period.Min","Wave.Period.Max","Wave.Period.Mean"
)]


# Formatting the 2011 Diallel phenotype data
# Jinliang Yang
# Jan 17th, 2012
# CL, CD and CW: Diallel Cob Trait

#setwd("/Users/yangjl/Documents/workingSpace/pheno2011/2011rawdata")

setwd("/Users/yangjl/Documents/pheno2011")

#############################################################
# Check the data completeness and correct the notes
#Notes:
#Dairy 8651-9000 v
#Zumwalt 9251-9600 v
#Johnson J2601-J2950 v
#Ki11: kill@ => Ki11 v
#Ms71: MS71 => Ms71 v

file1 <- read.csv("LL_cob.test.csv")
dim(file1)
#[1] 8774   10| kernel [1] 8286   11 |KRN #8678 8

diallel <- subset(file1, Note.y=="Diallel" | Note.y=="NAM Filler" | Note.y=="B73");
dim(diallel)
#[1] 2966   10 | [1] 2819   11 | [1] 2850    8
summary(diallel)

###########################################################
### Delete duplicated input and missing data
# Get ride of 126 duplicated records
barcode <- diallel[duplicated(diallel$Barcode),]$Barcode
diallel[diallel$Barcode %in% barcode,]
diallel <- diallel[!duplicated(diallel$Barcode),]

# Get ride of the missing records
#diallel[is.na(diallel$TKW),]
#diallel <- diallel[!is.na(diallel$KRN),]
##########################################################
### checking for outlyers
diallel$CL <- as.numeric(as.character(diallel$CL));
diallel$CD <- as.numeric(as.character(diallel$CD));
diallel$CW <- as.numeric(as.character(diallel$CW));

hist(diallel$CL);
hist(diallel$CD);
hist(diallel$CW);
idx <- which.max(diallel$CW)
diallel[idx,]
diallel[idx, ]$CW <- 28.99
#######################################################
### get the LSmean
library(nlme)

#fit the model (it takes a while to run in R)
tem <- diallel[!is.na(diallel$CL),]
lmeout1 <- lme(CL~Genotype, data=tem, random=~1|Farm);
lmeout2 <- lme(CD~Genotype, data=diallel, random=~1|Farm);
lmeout3 <- lme(CW~Genotype, data=diallel, random=~1|Farm);

#extract the estimates of coefficents, which are the lsmeans
ped.hat1 <- lmeout1$coef$fixed
ped.hat2 <- lmeout2$coef$fixed
ped.hat3 <- lmeout3$coef$fixed

#reparameterize the estimates to be more interpretable
ped.hat1[-1] <- ped.hat1[-1]+ped.hat1[1];
ped.hat2[-1] <- ped.hat2[-1]+ped.hat2[1];
ped.hat3[-1] <- ped.hat3[-1]+ped.hat3[1];

#rename the first estimate
head(ped.hat1)
head(diallel[order(diallel$Genotype),], 30);

names(ped.hat1)[1]="B73 x B97";
names(ped.hat2)[1]="B73 x B97";
names(ped.hat3)[1]="B73 x B97";

#get rid of the prefix "Genotype" in the levels names so that they are more readable
names(ped.hat1) <- gsub("Genotype","",names(ped.hat1));
names(ped.hat2) <- gsub("Genotype","",names(ped.hat2));
names(ped.hat3) <- gsub("Genotype","",names(ped.hat3));
#transpose the results
cl <- data.frame(Genotype=names(ped.hat1), CL=ped.hat1);
cd <- data.frame(Genotype=names(ped.hat2), CD=ped.hat2);
cw <- data.frame(Genotype=names(ped.hat3), CW=ped.hat3);

############# H2 ###############################
fit1 <- aov(CL ~ Farm + Genotype, data=diallel)
#2954644/(2954644+39598+1261566)=69.4 pval=2.2e-16 ***
fit2 <- aov(CD ~ Farm + Genotype, data=diallel)
#28022/(28022+15250+174)=64.5 pval=2.2e-16 ***
fit3 <- aov(CW ~ Farm + Genotype, data=diallel)
#390226/(390226+108061+2492)=77.9 pval=2.2e-16 ***
###########################################################

### merge the trait of reciprocal F1
library(ggplot2)
########-------CL
cl$Genotype <- as.character(cl$Genotype);
cl$p1 <- NA;
cl$p2 <- NA;
for (i in 1:nrow(cl)){
  tem <- sort(unlist(strsplit(cl$Genotype[i], split=" x ")));
  cl$p1[i] <- tem[1];
  cl$p2[i] <- tem[2];
}

cl <- cl[order(cl$p1, cl$p2),]
cl$ped <- paste(cl$p1, cl$p2, sep="|")
lsmean_cl <- ddply(cl, .(ped), summarise,
                        CL=mean(CL))
dim(lsmean_cl)
#[1] 251   2

########-------CW
cw$Genotype <- as.character(cw$Genotype);
cw$p1 <- NA;
cw$p2 <- NA;
for (i in 1:nrow(cw)){
  tem <- sort(unlist(strsplit(cw$Genotype[i], split=" x ")));
  cw$p1[i] <- tem[1];
  cw$p2[i] <- tem[2];
}

cw <- cw[order(cw$p1, cw$p2),]
cw$ped <- paste(cw$p1, cw$p2, sep="|")
lsmean_cw <- ddply(cw, .(ped), summarise,
                   CW=mean(CW))
dim(lsmean_cw)
#[1] 251   2

########-------CD
cd$Genotype <- as.character(cd$Genotype);
cd$p1 <- NA;
cd$p2 <- NA;
for (i in 1:nrow(cd)){
  tem <- sort(unlist(strsplit(cd$Genotype[i], split=" x ")));
  cd$p1[i] <- tem[1];
  cd$p2[i] <- tem[2];
}

cd <- cd[order(cd$p1, cd$p2),]
cd$ped <- paste(cd$p1, cd$p2, sep="|")
lsmean_cd <- ddply(cd, .(ped), summarise,
                   CD=mean(CD))
dim(lsmean_cd)
#[1] 251   2
######################################################################
nathan <- read.csv("Nathan_p1_subset.csv")
nam <- read.table("nam_parents", header=TRUE)
p <- c("B73", as.character(nam$parent))
p <- toupper(p);
nathan$INBRED <- toupper(nathan$INBRED)

pcob <- subset(nathan, INBRED %in% p)
pcob[pcob=="."] <- NA;
pcob$CW_Inbred <- as.numeric(as.character(pcob$CW_Inbred))
pcob$CL_Inbred <- as.numeric(as.character(pcob$CL_Inbred))
pcob$CD_Inbred <- as.numeric(as.character(pcob$CD_Inbred))
pcobmean <- ddply(pcob, .(INBRED), summarise,
              CW=mean(CW_Inbred, na.rm=T),
              CL=mean(CL_Inbred, na.rm=T),
              CD=mean(CD_Inbred, na.rm=T)
              )
b73 <- data.frame(INBRED="B73", CW=15.9, CL=115.8867, CD=25.59333)
pcobmean <- rbind(b73, pcobmean)
########################################################
#------------------ CL F1 and parents lsmean:
idx1 <- grep("@", lsmean_cl$ped)
f1cl <- lsmean_cl[-idx1,]

pcl <- lsmean_cl[idx1,]
pcl$ped <- gsub('@\\|NA',"", pcl$ped)

pcl$ped <- toupper(pcl$ped);
pcl2 <- merge(pcl, pcobmean[, c("INBRED", "CL")], by.x="ped", by.y="INBRED", all=T)

cor(subset(pcl2, !is.na(CL.x) & !is.na(CL.y))[,2:3])
#0.578

mymean <- function(x) mean(x, na.rm=T);
pcl2$CL <- apply(pcl2[,2:3], 1, mymean)
pcl2 <- pcl2[, c(1,4)]
names(pcl2)[1] <- "ped"

master_cl <- rbind(f1cl, pcl2)

#------------------ CW F1 and parents lsmean:
idx2 <- grep("@", lsmean_cw$ped)
f1cw <- lsmean_cw[-idx2,]

pcw <- lsmean_cw[idx2,]
pcw$ped <- gsub('@\\|NA',"", pcw$ped)

pcw$ped <- toupper(pcw$ped);
pcw2 <- merge(pcw, pcobmean[, c("INBRED", "CW")], by.x="ped", by.y="INBRED", all=T)

cor(subset(pcw2, !is.na(CW.x) & !is.na(CW.y))[,2:3])
#0.45

mymean <- function(x) mean(x, na.rm=T);
pcw2$CW <- apply(pcw2[,2:3], 1, mymean)
pcw2 <- pcw2[, c(1,4)]
names(pcw2)[1] <- "ped"

master_cw <- rbind(f1cw, pcw2)

#------------------ CD F1 and parents lsmean:
idx3 <- grep("@", lsmean_cd$ped)
f1cd <- lsmean_cd[-idx3,]

pcd <- lsmean_cd[idx3,]
pcd$ped <- gsub('@\\|NA',"", pcd$ped)

pcd$ped <- toupper(pcd$ped);
pcd2 <- merge(pcd, pcobmean[, c("INBRED", "CD")], by.x="ped", by.y="INBRED", all=T)

cor(subset(pcd2, !is.na(CD.x) & !is.na(CD.y))[,2:3])
#0.64

mymean <- function(x) mean(x, na.rm=T);
pcd2$CD <- apply(pcd2[,2:3], 1, mymean)
pcd2 <- pcd2[, c(1,4)]
names(pcd2)[1] <- "ped"

master_cd <- rbind(f1cd, pcd2)
###########################################
#### Merge the data for output
master_cob <- merge(master_cl, master_cw, by="ped", all=T)
master_cob <- merge(master_cob, master_cd, by="ped", all=T)

write.table(master_cob, "diallel_lsmean_cob_020112.csv", sep=",", row.names=FALSE, quote=FALSE)
## plot and output lsmean
idx2 <- grep("\\|", master_cd$ped);
f1 <- master_cd[idx2,]
p <- master_cd[-idx2, ]
plot(density(f1$CD), main="CD of Diallel and parents", xlab="KRN", lwd=3, col="blue")
lines(density(p$CD), lwd=3, lty=2, col="red")


#' Process output from a trainnBayesODM model
#'
#' Summarizes output for all types of trainBayesODM model functions.
#' @param shape SpatialPolygonsDataFrame.
#' @param mcmc Object of class \code{mcmc.list}.
#' @param effort Character, name of field in \code{shape} with effort.
#' @param detect Character, name of field in \code{shape} with detections.
#' @param neighIndex Index of units with neighbors (from the spatial object). Set as \code{NULL} if there are none.
#' @param islandndex Index of units that are islands (from the spatial object). Set as \code{NULL} if there are none.
#' @param pByState Logical, if \code{FALSE} (default), it is assumed there is one value of p (global) or one value of p per county. If \code{TRUE}, then it is assumed that there is one value of p per state/province, and \code{stateProv} must be specified.
#' @param stateProv Name of field in \code{shape} with state/province names. Ignored if \code{pByState} is \code{FALSE}.
#' @return List.
#' @examples

.processBayesODM <- function(
	shape,
	mcmc,
	effort,
	detect,
	neighIndex = NULL,
	islandIndex = NULL,
	pByState = FALSE,
	stateProv = NULL
) {

	nchains <- coda::nchain(mcmc)
	mcmc <- list(samples=mcmc)

	### calculate summary
	#####################

		mcmc$summary$all.chains <- matrix(NA, nrow=ncol(mcmc$samples[[1]]), ncol=5)
		rownames(mcmc$summary$all.chains) <- colnames(mcmc$samples[[1]])
		colnames(mcmc$summary$all.chains) <- c('Mean', 'Median', 'St.Dev', '95%CI_low', '95%CI_upp')

		samps <- mcmc$samples[[1]]
		if (nchains > 1) {
			for (chain in 2:nchains) {
				samps <- rbind(samps, mcmc$samples[[chain]])
			}
		}

		mcmc$summary$all.chains[ , 'Mean'] <- colMeans(samps, na.rm=TRUE)
		mcmc$summary$all.chains[ , 'Median'] <- apply(samps, 2, median, na.rm=TRUE)
		mcmc$summary$all.chains[ , 'St.Dev'] <- apply(samps, 2, sd, na.rm=TRUE)
		mcmc$summary$all.chains[ , '95%CI_low'] <- apply(samps, 2, quantile, probs=0.025, na.rm=TRUE)
		mcmc$summary$all.chains[ , '95%CI_upp'] <- apply(samps, 2, quantile, probs=0.975, na.rm=TRUE)

		rm(samps); gc()

	### update shapefile
	####################

		# psi
		if (!is.null(neighIndex[1])) {

			psiNeighIndex <- which(grepl(rownames(mcmc$summary$all.chains), pattern='psi_neigh[[]'))
			psiNeigh <- mcmc$summary$all.chains[psiNeighIndex, 'Mean']
			psiUncerNeigh <- mcmc$summary$all.chains[psiNeighIndex, '95%CI_upp'] - mcmc$summary$all.chains[psiNeighIndex, '95%CI_low']

		}

		if (!is.null(islandIndex[1])) {

			psiIslandIndex <- which(grepl(rownames(mcmc$summary$all.chains), pattern='psi_island[[]'))
			psiIsland <- mcmc$summary$all.chains[psiIslandIndex, 'Mean']
			psiUncerIsland <- mcmc$summary$all.chains[psiIslandIndex, '95%CI_upp'] - mcmc$summary$all.chains[psiIslandIndex, '95%CI_low']

		}

		# p
		pIndex <- if (any(grepl(rownames(mcmc$summary$all.chains), pattern='p[[]'))) {
			which(grepl(rownames(mcmc$summary$all.chains), pattern='p[[]'))
		} else {
			which(rownames(mcmc$summary$all.chains) == 'p')
		}
		p <- mcmc$summary$all.chains[pIndex, 'Mean']
		pUncer <- mcmc$summary$all.chains[pIndex, '95%CI_upp'] - mcmc$summary$all.chains[pIndex, '95%CI_low']

		# if p values pertain to states, not counties
		if (pByState) {

			state <- as.numeric(as.factor(shape@data[ , stateProv]))
			p <- p[state]
			pUncer <- pUncer[state]

		}

		# assign
		shape@data$pUncer <- shape@data$p <- shape@data$psiUncer <- shape@data$psi <- NA

		shape@data$psiUncer <- shape@data$psi <- NA
		if (!is.null(neighIndex[1])) {
			shape@data$psi[neighIndex] <- psiNeigh
			shape@data$psiUncer[neighIndex] <- psiUncerNeigh
		}

		if (!is.null(islandIndex[1])) {
			shape@data$psi[islandIndex] <- psiIsland
			shape@data$psiUncer[islandIndex] <- psiUncerIsland
		}

		shape@data$p <- p
		shape@data$pUncer <- pUncer

		names(shape@data)[(ncol(shape@data) - 3):ncol(shape@data)] <- c('psi', 'psi90CI', 'p', 'p90CI')

	### done
	########

		list(mcmc=mcmc, shape=shape)

}


# Make arrow RT plot
library(ggplot2)
library(ggstance)
library(ggrepel)

make_rt_point_plot <- function(rt_data_long, label, JOBID, ext = ".png"){
  rt_levels_state <- unique(rt_data_long$state)
  rt_data_long$state <- factor(rt_data_long$state, levels = rt_levels_state)
  rt_levels_state_name <- unique(rt_data_long$state_name)
  rt_data_long$state_name <- factor(rt_data_long$state_name, levels = rt_levels_state_name)

  p1 <- ggplot(rt_data_long) +
    geom_point(aes(y = state_name, x = rt, shape = x, col = groupings),  stat="identity",
               position=position_dodge(width = -0.5)) +
    geom_pointrangeh(aes(x = rt, xmin=rt_min, xmax=rt_max, y=state_name, group = x, col = groupings),
                   position = position_dodge(width= -0.5), size = 0.05) +
    ylab("") + xlab(expression(R[t])) +
    scale_x_log10(expand = c(0, 0),breaks=0:6) +
    scale_colour_discrete(name = "") +
    scale_shape_discrete(name = "", labels = c("Initial", "Current")) +
    geom_vline(xintercept=0:6,colour="gray",alpha=.5) +
    geom_vline(aes(xintercept=1)) +
    theme(axis.line = element_line(colour = "black"),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          panel.background = element_blank(),
          legend.key = element_blank())
  ggsave(paste0("usa/figures/rt_point_", label, "_", JOBID, ext), p1, width = 8, height=7)
}


\name{BvCat.cv}
\alias{BvCat.cv}
\title{Function perform cross-validation and fit the penalized bivariate multinomial logistic regression.}
\usage{
BvCat.cv(X, Y, ngamma = 100, lambda.vec = seq(.01,1, length=10),
                    nfolds = 5, delta = .1,
										standardize = TRUE, tol = 1e-8, quiet = TRUE,
										inner.quiet= TRUE)
}
\arguments{
\item{X}{An \eqn{n \times p} matrix of predictors. This should not include an intercept, nor should it be standardized.}
\item{Y}{A list of length two: each list corresponds to one of the two response variables. The first component of the list should be made up of integers from 1 to J; the second component of integers from 1 to K. }
\item{ngamma}{The number of candidate tuning parameters \eqn{\gamma} to be considered.}
\item{lambda.vec}{A set of candidate tuning parameters for \eqn{\lambda}. In general, values between \eqn{10^{-4}} and \eqn{1} tend to be best, but this may need to be explored for each particular application.}
\item{nfolds}{The number of folds for cross-validation. Default is \code{5}. If cross-validation is not desired, set equal to NULL.}
\item{delta}{The ratio of maximum to minimum \eqn{\gamma} candidate values. Smaller values of \eqn{\delta} would allow for less sparse model fits, but at the expensive of computing time.}
\item{standardize}{Should the predictors be standardized (centered and scaled to have unit standard deviation) for model fitting? Default is TRUE.}
\item{tol}{Convergence tolerance parameter. Convergence is claimed when three consecutive objective function value changes are below this tolerance. Default is \eqn{10^{-8}}, but for some applications, this may need to be adjusted.}
\item{quiet}{Should any progress be printed?}
\item{inner.quiet}{Should much more progress be printed? Use with caution. }
}

\value{
\item{beta}{A sparse matrix containing regression coefficient estimates. Note that if \code{standardize=TRUE}, these will be on the scale of the standardized predictors, so use with caution. It is highly recommended to extract coefficients with \code{BvCat.coef}.}
\item{J}{Number of response categories for the first response.}
\item{K}{Number of response categories for the second response.}
\item{D}{The constraint matrix used in the penalty. See description for more details.}
\item{standardize}{Were predictors standardized for model fitting?}
\item{errs.joint}{An array of dimension \code{ngamma} \eqn{\times} \code{length(lambda.vec)} \eqn{\times} \code{nfolds} of misclassification errors for the candidate tuning parameters in each of the \code{nfolds}.}
\item{deviance}{An array of dimension \code{ngamma} \eqn{\times} \code{length(lambda.vec)} \eqn{\times} \code{nfolds} of multinomial deviances for the candidate tuning parameters in each of the \code{nfolds}. Note that misclassification error is the default performance metric.}
\item{lambda.min.joint}{Tuning parameter value for \eqn{\lambda} which minimized average misclassification error.}
\item{gamma.min.joint}{Tuning parameter value for \eqn{\gamma} which minimized average misclassification error.}
\item{X.mean}{The columnwise averages of the input matrix \code{X}. Used for extracting coefficients and prediction when model fit with \code{standardize=TRUE}.}
\item{X.sd}{The columnwise standard deviations input matrix \code{X}. Used for extracting coefficients and prediction when model fit with \code{standardize=TRUE}.}
\item{cv.index}{Which observations belonged to which cross-validation folds?}
\item{lambda.vec}{Exactly the input \code{lambda.vec}.}
\item{gamma.vec}{The candidate values of \eqn{\gamma} used for model fitting.}
}

\examples{
# ----------------------------------------
# Example for BvCategorical
# ----------------------------------------
set.seed(1)
p <- 50
n <- 300
J <- 3
K <- 2
Results <- NULL
SigmaX <- matrix(0, nrow=p, ncol=p)
for(jj in 1:p){
  for(kk in 1:p){
    SigmaX[jj,kk] <- .5^abs(jj-kk)
  }
}
eo <- eigen(SigmaX)
SigmaXsqrt <- tcrossprod(tcrossprod(eo$vec, diag(eo$val^.5)), eo$vec)

# generate data -------------
X <- tcrossprod(matrix(rnorm(n*p), nrow=n), SigmaXsqrt)

beta <- matrix(0, nrow=p, ncol=J*K)
temp <- sample(1:p, 10)

# predict affect joint probabilities -------------
for(j in 1:6){
  beta[temp[j],] <- runif(6, -3, 3)
}
# predictors affect marginal probabilities only ---
for(j in 7:10){
  temp1 <- runif(4, -3, 3)
  beta[temp[j],] <- c(-temp1[4] + temp1[3] + temp1[1], temp1[1], temp1[2],
                      temp1[3], temp1[4], -temp1[1]+ temp1[4] + temp1[2])
}

lleval <- exp(crossprod(t(X), beta))
Pmat <- lleval/tcrossprod(rowSums(lleval), rep(1, J*K))
Ymat <- matrix(0, nrow=n, ncol=J*K)
Y <- list(NA)
Y[[1]] <- rep(0, n)
Y[[2]] <- rep(0, n)
refClasses <- rbind(c(1:J, 1:J), c(rep(1, J), rep(2, J)))
for(k in 1:n){
  Ymat[k,] <- rmultinom(1, 1, Pmat[k,])
  Y[[1]][k] <- refClasses[1,which(Ymat[k,]==1)]
  Y[[2]][k] <- refClasses[2,which(Ymat[k,]==1)]
}

# ----------------------------------------
# Fit model with cross-validation
# ----------------------------------------
modfit <- BvCat.cv(X, Y, ngamma = 20,
                   lambda.vec = 10^seq(-4, 0, length=5),
                   nfolds = 5, delta = .01, standardize = TRUE,
                   tol = 1e-8,  quiet = FALSE, inner.quiet = TRUE)

# ---------------------------------------
# CV errors
# ---------------------------------------
apply(modfit$errs.joint, c(1,2), mean) # default
apply(modfit$deviance, c(1,2), mean)

# -----------------------------------------------------
# Coefficients at tuning parameter pair selected by CV
# -----------------------------------------------------
temp.coef <- BvCat.coef(modfit, type="matrix")
temp.coef.tensor <- BvCat.coef(modfit, type="tensor")
temp.coef.untuned <- BvCat.coef(modfit, lambda = modfit$lambda.vec[1], gamma = modfit$gamma.vec[10], type="matrix")

ntest <- 1
Xtest <- tcrossprod(matrix(rnorm(ntest*p), nrow=ntest),SigmaXsqrt)
temp.predict <- BvCat.predict(Xtest, modfit, type="class")
temp.predict.probs <- BvCat.predict(Xtest, modfit, type="probabilities")
# first dimension of temp.predict.probs$preds length ntest, second J, and third K

# ----------------------------------------
# Fit model with no cross-validation
# ----------------------------------------
modfit.noCV <- BvCat.cv(X, Y, ngamma = 20,
                   lambda.vec = 10^seq(-1, -4, length=5),
                   nfolds = NULL, delta = .01, standardize = TRUE,
                   tol = 1e-8,  quiet = TRUE, inner.quiet = TRUE)

# temp.coef <- BvCat.coef(modfit.noCV, type="matrix") # returns errror
temp.coef.noCV <- BvCat.coef(modfit.noCV, lambda = modfit$lambda.vec[1], gamma = modfit$gamma.vec[10], type="matrix")

}


# function to create a custom heatmap using polygons

#' Create heat map
#'
#' Function for a simple heat map with custom limits in base r plotting. The function is a wrapper for the image function
#'
#' @param mat matrix with values to create heatmap for
#' @param xrange the x range in which to plot the grid (vector length 2)
#' @param yrange the y range in which to plot the grid (vector length 2)
#' @param xlimit the x limits of the plot (vector length 2)
#' @param ylimit the y limits of the plot (vector length 2)
#' @param colramp the colors to use for plotting (vector)
#' @param maketicks draw simple x,y axes (logical).
#' @param xticks position of xticks, defaults to \code{pretty}
#' @param yticks posiition of yticks, defaults to \code{pretty}
#' @param makebox draw box line around plot (logical)
#' @param axislwd linewidth to use for ticks and box (numeric)
#' @param locut low threshold on raster values (numeric)
#' @param hicut high threshold on raster values (numeric)
#' @param nacut values outside cutrange are set to NA instead of cut value (logical)
#' @param axisvars list with additional arguments to \code{axis} (named list, not including \code{side},\code{at} or \code{lwd})
#' @param ... other parameters passed to \code{image} function
#'
#' @return plot
#' @export
heatGrid <- function(mat=NULL, xrange=NULL, yrange=NULL, xlimit=NULL, ylimit=NULL,
                     colramp=topo.colors(200), maketicks=T, xticks=NULL, yticks=NULL, makebox=T, axislwd=0.5,
                     locut=NULL, hicut=NULL, nacut=F, axisvars=NULL, ...){

  # setup default data
  testn    <- 25
  trigo    <- cos(seq(0,2*pi,length.out=testn))
  testdata <- do.call(cbind,lapply(1:length(trigo), function(i) abs(rev(sort(rnorm(testn,trigo[i],0.5))))))
  # testdata <- ifelse(testdata>0.2,NA,testdata)
  if (is.null(mat)){mat <- testdata}

  # seq axes sequencing
  if (is.null(xrange)){xrange <- c(1,ncol(mat))}
  if (is.null(yrange)){yrange <- c(1,nrow(mat))}
  xseq <- seq(xrange[1],xrange[2],length.out=ncol(mat))
  yseq <- seq(yrange[1],yrange[2],length.out=nrow(mat))

  # transpose matrix for image function
  pdat <- t(mat[nrow(mat):1,])

  # cut values outside range
  if (!is.null(locut)){
    if (nacut){
      pdat[pdat<=locut] <- NA
    }else{
      pdat[pdat<=locut] <- locut
    }
  }
  if (!is.null(hicut)){
    if (nacut){
        pdat[pdat>=hicut] <- NA
    }else{
      pdat[pdat>=hicut] <- hicut
    }
  }

  # set zlimit
  zlimit <- range(mat)
  if (!is.null(locut)){zlimit[1] <- locut}
  if (!is.null(hicut)){zlimit[2] <- hicut}

  # get exact border coordinates
  if (is.null(xlimit)){xlimit <- c(xseq[1]-0.5*diff(xseq[1:2]), tail(xseq,1)+0.5*diff(xseq[1:2]))}
  if (is.null(ylimit)){ylimit <- c(yseq[1]-0.5*diff(yseq[1:2]), tail(yseq,1)+0.5*diff(yseq[1:2]))}
  image(xseq,yseq, pdat, col=colramp, axes=F, xaxs='i', yaxs='i',
        xlim=xlimit, ylim=ylimit, zlim=zlimit,
        ...)

  if (maketicks){
    if (is.null(xticks)){
      at_x <- pretty(xseq)
    }else{
      at_x <- xticks
    }
    if (is.null(yticks)){
      at_y <- pretty(yseq)
    }else{
      at_y <- yticks
    }
    do.call(axis, c(list(side=1, at=at_x, lwd=axislwd), axisvars))
    do.call(axis, c(list(side=2, at=at_y, lwd=axislwd), axisvars))
  }
  if (makebox){
    box(lwd=axislwd)
  }

}

#' @rdname heatGrid
#' @export
pkHeatmap <- heatGrid