library(tidyr)
library(reticulate)
library(ggbeeswarm)
source('config.r')
np <- import("numpy")

data <- np$load(
  "./outputs/all_scores_models_tuh_mae_shuffle-split.npy",
  # "./outputs/all_scores_mag_models_mnecommonsubjects_interval_rep-kfold.npy",
# "all_scores_mag_models_mnecommonsubjects.npy",
  allow_pickle = T)[[1]] %>%
  as.data.frame()

data_comp_scores <- read.csv(
  "outputs/tuh_component_scores.csv",
  stringsAsFactor = T
)

data_comp_scores_long <- data.frame(
  score = c(data_comp_scores$spoc, data_comp_scores$riemann),
  estimator = factor(rep(c("SPoC", "Riemann"),
                         each = nrow(data_comp_scores))),
  n_components = rep(data_comp_scores$n_components, times = 2),
  fold_idx = factor(rep(data_comp_scores$fold_idx, times = 2))
)

agg_scores <- aggregate(cbind(spoc, riemann) ~ n_components,
                        data = data_comp_scores, FUN = mean)

data_ <- data[, (!names(data) %in% c("dummy"))]
n_splits <- nrow(data_)

data_long <- data_ %>% gather(key = "estimator", value = "score")
# move to long format
data_long$estimator <- factor(data_long$estimator)

est_types <- c(
  "naive",
  "diag",
  "SPoC",
  "Riemann",
  "SPoC",
  "Riemann"
)

est_names <- c(
  "upper",
  "diag",
  "SPoC",
  "Riemann",
  sprintf("SPoC[%d]", which.min(agg_scores$spoc)),
  sprintf("Riemann[%d]", which.min(agg_scores$riemann))
)

est_labels <- setNames(
  c("upper", est_types[c(-1, -5, -6)]),
   est_types[c(-5, -6)]
)

# categorical colors based on: https://jfly.uni-koeln.de/color/
# beef up long data
data_long$est_type <- factor(rep(est_types, each = n_splits))

data_long$fold <- rep(1:n_splits, times = length(est_types))

# prepare properly sorted x labels
sort_idx <- order(apply(data_, 2, mean))
levels_est <- est_names[sort_idx]

my_color_cats <- setNames(
  with(
    color_cats,
    c(`sky blue`, `blueish green`, vermillon, orange)),
  c("naive", "diag", "SPoC", "Riemann"))

ggplot(data = subset(data_long, estimator != "dummy"),
       mapping = aes(y = score, x = reorder(estimator, I(score)))) +
  geom_beeswarm(
    priority = 'density',
    mapping = aes(color = est_type),
    size = 2.5,
    alpha = 0.2,
    show.legend = T, cex = 1) +
  scale_size_continuous(range = c(0.5, 2)) +
  scale_alpha_continuous(range = c(0.4, 0.7)) +
  geom_boxplot(mapping = aes(fill = est_type, color = est_type),
               alpha = 0.3,
               outlier.fill = NA, outlier.colour = NA) +
  geom_jitter(width = 0.2)+
  stat_summary(geom = 'text',
               mapping = aes(label = sprintf("%1.2f",
                                              ..y..)),
               fun.y = mean, size = 3.2, show.legend = FALSE,
               position = position_nudge(x = -0.49)) +
  my_theme +
  labs(y = "MAE", x = NULL, parse = T) +
  guides(size = F, alpha = F) +
  theme(legend.position = c(0.8, 0.25)) +
        # legend.position = "top", legend.text = element_text(size = 18)) +
  coord_flip(ylim=c(5,15)) +
  scale_fill_manual(values = my_color_cats,
                    breaks = names(my_color_cats),
                    labels = est_labels,
                    name = NULL) +
  scale_color_manual(
    values = my_color_cats,
    breaks = names(my_color_cats),
    labels = est_labels,
    name = NULL) +
  scale_x_discrete(labels = parse(text = levels_est)) +
  geom_hline(yintercept = mean(data$dummy), linetype = 'dashed') +
  annotate(geom = "text",
           y = mean(data$dummy) + 0.2, x = 3, label = 'predicting~bar(age)',
           size = annotate_text_size,
           parse = T, angle = 270)

fname <- "./outputs/fig_tuh_model_comp"
ggsave(paste0(fname, ".png"),
       width = 8, height = 5, dpi = 300)

ggsave(paste0(fname, ".pdf"),
       useDingbats = F,
       width = 8, height = 5, dpi = 300)
embedFonts(file = paste0(fname, ".pdf"), outfile = paste0(fname, ".pdf"))

component_labels <- setNames(
      rev(parse(text = est_names[-c(1:4, 7)])),
      c("Riemann", "SPoC"))

fig_components <- ggplot(data = data_comp_scores_long,
       mapping = aes(
         x = n_components, y = score,
  #  group = interaction(estimator, fold_idx),
         color = estimator, fill = estimator)) +
  stat_summary(inherit.aes = F,
               mapping = aes(fill = estimator, x = n_components,
                             y = score),
               fun.ymin = function(x) mean(x) - sd(x),
               fun.ymax = function(x) mean(x) + sd(x),
               geom = 'ribbon', alpha = 0.2) +
  stat_summary(fun.y = mean, geom = 'line', size = 1.5) +
  my_theme +
  theme(legend.position = c(0.8, 0.5)) +
  scale_color_manual(
    values = my_color_cats[c("Riemann", "SPoC")],
    breaks = c("Riemann", "SPoC"),
    labels = component_labels,
    name = NULL) +
  scale_fill_manual(
    values = my_color_cats[c("Riemann", "SPoC")],
    breaks = c("Riemann", "SPoC"),
    labels = component_labels,
    name = NULL) +
  labs(x='#components', y = "MAE") +
  geom_vline(
    xintercept = which.min(agg_scores$spoc),
    size = 0.6,
    color = my_color_cats['SPoC'], linetype = 'dashed') +
  geom_vline(
    xintercept = which.min(agg_scores$riemann),
    size = 0.6,
    color = my_color_cats['Riemann'], linetype = 'dashed') +
  scale_x_continuous(breaks = seq(0, 100, 10))
  # scale_y_continuous(breaks = seq(0, .8, .1)) +
  # coord_cartesian(ylim = c(0, 0.8))

fname <- "./outputs/fig_tuh_component_selection"
ggsave(paste0(fname, ".png"),
       width = 8, height = 5, dpi = 300)

ggsave(paste0(fname, ".pdf"),
       useDingbats = F,
       width = 8, height = 5, dpi = 300)
embedFonts(file = paste0(fname, ".pdf"), outfile = paste0(fname, ".pdf"))


#!/usr/bin/Rscript

# week olera kolera cholera epidemi pest farsot
# olera is the some of kolera and cholera

series <- read.table(file = "weekly_counts.txt",
                     col.names = c("week","olera","kolera","cholera","epidemi","pest","farsot"),
                     sep = "",
                     quote = "\"",
                     dec = ".",
                     fill = TRUE,
                     comment.char = "#")

pdf( "weekly_mentions.pdf", width = 7, height = 5 )

# cairo_ps("image.eps")

plot(x = series$week,
     y = series$olera,
     type = "l",
     ylab = "Occurence in corpus",
     xlab = "Week")


library(filesampler)

system.time(x <- wc_l("big.csv"))

x


library(data.table)

source('../../functions.r')

load('../generated/model/mcmc-results.rdata', v=F)

spplist <- fread('../acap-species-list.csv')[base_species == T]
setkeyv(spplist, 'code.sra')
### APF by genus
mcsel_genus <- mcmc[variable %in% mc_attributes[role == 'quarter0_spatial1' & vartype == 'apf_t', variable]]
mcsel_genus <- merge(mcsel_genus, mc_attributes, by = 'variable', all.x=T, all.y=F)
apf_samples_genus <- merge(mcsel_genus, spplist[, .(species_code=code.sra, genus=vulgroup)], by='species_code', all.x=T)
apf_genus <- apf_samples_genus[, .(apf=sum(value)), .(sample, genus)][, .(amean=mean(apf), ali=quantile(apf, 0.025), aui=quantile(apf, 0.975)), genus]
write.csv(apf_genus, row.names=FALSE, file='../generated/apf-genus.csv')


cl.out.path <- "/outputs"
setwd(cl.out.path)

library(scrattch.hicat)
library(Matrix)
library(scrattch.io)
library(mfishtools)
library(matrixStats)
library(ggplot2)
library(tidyr)
library(dplyr)
library(scrattch.vis)
library(RANN)
library(WGCNA)
library(gridExtra)
library(data.table)
library(R.utils)
library(Seurat)
library(cowplot)

options(stringsAsFactors=FALSE)

############################################################################
## Read in 10x data and additional information

load("data/hippocampus_10x.RData")
norm.dat.10X    <- logCPM(datExpr)
row.names(anno) <- colnames(norm.dat.10X)
samp.dat.10X    <- anno

exclude.genes   <- scan("data/exclusiongenes_mito_sex_tissue.txt", what="character")
use.genes       <- sort(setdiff(row.names(norm.dat.10X), exclude.genes))
save(use.genes, exclude.genes, file="use.genes.rda")

ionChannels     <- scan("data/ion_channels.txt",what="character")
epState         <- setNames(c("WG1","WG1","WG4","WG4"), c("H16.06.008","H17.06.015","H16.06.009","H16.06.010"))

############################################################################
## Assign classes using k-means clustering on all cells based on marker genes from HBA and Cell types database

# Read in predetermined genes from Allen Brain Map
highDG   <- intersect(scan("data/topDG_HBA.txt",what="character",sep="\n"),rownames(norm.dat.10X))[1:100]
highGlut <- intersect(scan("data/topGlut_CT.txt",what="character",sep="\n"),rownames(norm.dat.10X))
highGABA <- intersect(scan("data/topGABA_CT.txt",what="character",sep="\n"),rownames(norm.dat.10X))
highGlia <- intersect(scan("data/topGlia_CT.txt",what="character",sep="\n"),rownames(norm.dat.10X))
dexGenes <- unique(c(highDG,highGlut,highGABA,highGlia))

# Run K-means clustering on cells with enough genes detected
set.seed(1)  # For reproducibility
kpCells  <- samp.dat.10X$gene.counts.0 > 1000
dat      <- t(norm.dat.10X[dexGenes,kpCells])
rcl      <- kmeans(as.matrix(as.data.frame(dat)),4)
rcl      <- kmeans(as.matrix(as.data.frame(dat)),4)
rcl      <- kmeans(as.matrix(as.data.frame(dat)),4)
# NOTE: this above line of code needs to be run at least 3 times to reproduce the results
# There is no reason why this should be the case, but it is.
cluster  <- norm.dat.10X[1,]*0
cluster[kpCells] <- rcl$cluster

# Assign clusters to appropriate classes based on marker gene expression
outVal   <- NULL
for(i in 1:4) {
  outVal <- cbind(outVal,Matrix::rowMeans(norm.dat.10X[c("PROX1","GAD1","SLC17A7","SLC1A3"),cluster==i]>0))
}
val <- 1
for (i in 1:4) {
  v   <- which.max(outVal[i,])
  val <- c(val, v+1)
  outVal[,v] = 0
}
classes  <- c("Low quality", "Dentate Gyrus", "Inhibitory", "Excitatory", "Non-neuronal")[order(val)]
classes  <- classes[cluster+1]
samp.dat.10X$class_label <- classes
print(table(classes[kpCells]))
#Dentate Gyrus    Excitatory    Inhibitory  Non-neuronal
#        12011          6386          5302          3543 # These are the values you should get.

############################################################################
## Plot relevant ion channels - This is figure 4B
ionch   <- c("KCNJ2","KCNA1","KCNB1","KCNC4","KCND2","KCNQ2","CACNA1D","CACNA1B","CACNA1H","KCNMA1","KCNN2","SCN1A","HCN1")
ionDat  <- 2^as.matrix(t(norm.dat.10X[ionch,]))-1  # Convert to linear CPM
plotDat <- data.frame(sample_name = colnames(norm.dat.10X),ionDat)
annoDat <- data.frame(sample_name = colnames(norm.dat.10X),cluster = samp.dat.10X$class)
annoDat <- annotate_cat(annoDat,"cluster")

p <- sample_fire_plot(plotDat,annoDat,ionch,"cluster",log_scale=TRUE, max_width=25,label_height=15)
p
ggsave("class_markers_new.pdf",p,device="pdf",height=5,width=4)

############################################################################
## Subset only the cells in dentate gyrus
kpDG   <- samp.dat.10X$class_label=="Dentate Gyrus"
datDG  <- as.matrix(norm.dat.10X[,kpDG])
annoDG <- samp.dat.10X[kpDG,]

# Reorder donor by WG
annoDG$Donor = factor(annoDG$Donor, levels=c("H16.06.008","H17.06.015","H16.06.009","H16.06.010"))
annoDG <- annotate_factor(annoDG,"Donor")

print(table(annoDG$Donor_label))
#  H16.06.008 H16.06.009 H16.06.010 H17.06.015
#        5035       1398       1853       3725

print(table(annoDG$Donor_label)/table(samp.dat.10X$Donor)[names(table(annoDG$Donor_label))])
#  H16.06.008 H16.06.009 H16.06.010 H17.06.015
#   0.5678998  0.2992935  0.2318569  0.4848367
# There is a significantly higher fraction of DG cells in the WG1 cases vs. the WG4 cases.

kpGn <- rowSums(datDG>1) >= (dim(datDG)[2]/20)  # Keep genes expressed in at least 5% of DG cells
useGenes <- sort(rownames(datDG)[kpGn])
datDG <- datDG[useGenes,]

fwrite(signif(datDG,3), file = "DG_data_new.csv", row.names=TRUE)
file.remove('DG_data_new.csv.gz')
gzip('DG_data_new.csv',destname='DG_data_new.csv.gz')

############################################################################
## Compute statistics per donor and find potentially significant genes and output
cl <- setNames(annoDG$Donor_label,colnames(datDG))
meansDG <- get_cl_means(datDG,cl)
e1 <- epState[colnames(meansDG)] == "WG1"
e4 <- epState[colnames(meansDG)] == "WG4"
dexDG <- rowMeans(meansDG[,e4]) - rowMeans(meansDG[,e1])
sigDG <- dexDG / (1+apply(meansDG[,e4],1,sd)+apply(meansDG[,e1],1,sd))

stats <- signif(data.frame(meansDG,dexDG,sigDG)[order(sigDG),],3)
colnames(stats) <- c(paste(epState[colnames(meansDG)],names(epState[colnames(meansDG)])),"Fold Change", "Score")

fwrite(stats, file = "DG_stats_new.csv", row.names=TRUE)

############################################################################
## Plot all the ion channels

# The first three genes in this plot is part of Figure 4E
ionChannels <- c("KCNMA1","CACNA1B","KCNJ2",intersect(names(sort(sigDG)), ionChannels))

plotDat <- data.frame(donor = as.factor(paste(epState[annoDG$Donor_label],annoDG$Donor_label)),
                      WG = epState[annoDG$Donor_label])

pdf("ionChannels_new.pdf", onefile = TRUE)
for (i in ionChannels){
  plotDat$log2CPM_UMI <- datDG[i,]
  main <- paste0(i,"; mean(WG4-WG1) = ",signif(dexDG[i],3),"; score = ",signif(sigDG[i],3))
  p <- ggplot(plotDat, aes(x=donor, y=log2CPM_UMI, fill=WG)) +
	   geom_boxplot(width=0.1, fill="white") +
	   labs(title=main) +
	   ylim(0,15)
  grid.arrange(p)
}
dev.off()

############################################################################
## Plot the best ion channels

# This plot is part of Figure 1B
ionch2 <- c(intersect(rownames(stats),ionChannels)[1:10],
            intersect(rownames(stats)[dim(stats)[1]:1],ionChannels)[10:1])
ionDG  <- 2^as.matrix(t(datDG[ionch2,]))-1  # Convert to linear CPM
plotDG <- data.frame(sample_name = colnames(datDG),ionDG)
annoDG$sample_name = colnames(datDG)
annoDG <- annoDG[,c("sample_name",setdiff(colnames(annoDG),"sample_name"))]

p = sample_fire_plot(plotDG,annoDG,ionch2,"Donor",log_scale=TRUE, max_width=35,label_height=15)
p
ggsave("ionChannel_DG_markers_new.pdf",p,device="pdf",height=5,width=2.5)

############################################################################
## Plot the best neuron projection genes

# This plot is part of Figure 1B
npGenes <- scan("data/neuron_projection_genes.txt", what="character")
np      <- c(intersect(rownames(stats),npGenes)[1:10],
            intersect(rownames(stats)[dim(stats)[1]:1],npGenes)[10:1])
npDG    <- 2^as.matrix(t(datDG[np,]))-1  # Convert to linear CPM
plotDG  <- data.frame(sample_name = colnames(datDG),npDG)

p = sample_fire_plot(plotDG,annoDG,np,"Donor",log_scale=TRUE, max_width=35,label_height=15)
p
ggsave("neuronProjection_DG_markers_new.pdf",p,device="pdf",height=5,width=2.5)

############################################################################
## Compare WG4/WG1 DEX genes to other combos.

dexR1 <- rowMeans(meansDG[,c(3,4)]) - rowMeans(meansDG[,c(1,2)])
dexR2 <- rowMeans(meansDG[,c(2,4)]) - rowMeans(meansDG[,c(1,3)])

out_dex <- data.frame(
    real_dex  = c(sum(dexDG >= 1),sum(dexDG <= -1)),
	rand_dex1 = c(sum(dexR1 >= 1),sum(dexR1 <= -1)),
	rand_dex2 = c(sum(dexR2 >= 1),sum(dexR2 <= -1)))
rownames(out_dex) <- c("Higher in WG4", "Lower in WG4")
print(out_dex)
#               real_dex rand_dex1 rand_dex2
# Higher in WG4      397       634       303
# Lower in WG4       341       176       289
# The number of genes is no different from what we'd expect by chance.  What if we put all the cells on a PC vector?

############################################################################
## Run Seurat on the data and output plots for PCs and UMAP

human_metadata <- data.frame(donor = annoDG$Donor_label,
                             state = epState[annoDG$Donor_label])
rownames(human_metadata) <- colnames(datDG)
human <- CreateSeuratObject(datDG, meta.data = human_metadata)
human <- NormalizeData(human, verbose = FALSE)
human <- FindVariableFeatures(human, selection.method = "vst", nfeatures = 2000, verbose = FALSE) # 750
human <- ScaleData(object = human, verbose = FALSE)
human <- RunPCA(object = human, npcs = 30, verbose = FALSE)
human <- RunUMAP(object = human, reduction = "pca", dims = 1:30, verbose = FALSE)

p1 <- DimPlot(object = human, group.by = "donor", reduction = "pca", do.return = TRUE,
      pt.size = 0.5, label=TRUE) + NoLegend()
p2 <- DimPlot(object = human, group.by = "state", reduction = "pca", do.return = TRUE,
      pt.size = 0.5, label=TRUE, cols=c("blue","red")) + NoLegend()
p3 <- DimPlot(object = human, group.by = "donor", reduction = "umap", do.return = TRUE,
      pt.size = 0.5, label=TRUE) + NoLegend()
p4 <- DimPlot(object = human, group.by = "state", reduction = "umap", do.return = TRUE,
      pt.size = 0.5, label=TRUE, cols=c("blue","red")) + NoLegend()
plot_grid(p1, p2, p3, p4, ncol = 2)
ggsave("umap_pca_plots.pdf",device="pdf",height=9,width=9)

# This plot is part of Figure 1B
DimPlot(object = human, group.by = "state", reduction = "pca", do.return = TRUE,
      pt.size = 0.4, label=FALSE, cols=c("blue","red")) + NoLegend()
ggsave("umap_pca_plots_PC_byState.pdf",device="pdf",height=5,width=5)


#' Function that corrects chronologies for sudden jumps in time
#'
#' Some occurrences in the model results can lead the CumDY function
#' to detect extra year transitions, resulting in sudden jumps in
#' the shell chronology or a start of the chronology at an age
#' beyond 1 year. This function removes these sharp transitions
#' and late onset by adding or subtracting whole years to the age
#' result.
#' @param resultarray Array containing the full results of
#' the optimized growth model
#' @param T_per The period length of one year (in days)
#' @param agecorrection Correct for jumps in age (/code{TRUE}) or
#' only for starting time (/code{FALSE})
#' @param plot Should the results be plotted? (/code{TRUE/FALSE})
#' @return An updated and corrected version of \code{resultarray}
#' @references package dependencies: ggplot2 3.2.1
#' @examples
#' testarray <- array(NA, dim = c(20, 16, 9)) # Create empty array
#' # with correct third dimension
#' windowfill <- seq(10, 100, 10) # Create dummy simulation data
#' # (ages) to copy through the array
#' for(i in 6:length(testarray[1, , 1])){
#'     testarray[, i, 3] <- c(windowfill,
#'         rep(NA, length(testarray[, 1, 3]) - length(windowfill)))
#'     windowfill <- c(NA, windowfill + 31)
#' }
#' testarray2 <- age_corr(testarray, 365, FALSE, FALSE) # Apply function on
#' array
#' @export
age_corr <- function(resultarray,
    T_per = 365,
    plot = TRUE,
    agecorrection = TRUE
    ){
    Age <- Age_corr <- NULL # Predefine variables to circumvent global variable binding error
    # Check on glitches where consecutive windows are placed (almost) 1 year apart.
    # These are repaired by adding one T_per to the cumulative Day values of those windows in
    # resultarray[, , 3]
    mean_window_age <- data.frame(
        1:(length(resultarray[1,,1]) - 5),
        unname(colMeans(resultarray[, 6:length(resultarray[1, , 1]), 3], na.rm = TRUE)))
    colnames(mean_window_age) <- c("window", "Age")

    # Plot mean window ages
    dev.new() # Plot Window age to check for strange jumps in age
    ageplot <- ggplot2::ggplot(mean_window_age, ggplot2::aes(window, Age)) +
        ggplot2::geom_line() +
        ggplot2::geom_point() +
        ggplot2::ggtitle("Plot of average ages of modelling windows") +
        ggplot2::xlab("Window #") +
        ggplot2::scale_y_continuous("Age (days)", seq(0, 365 * ceiling(max(unname(colMeans(
            resultarray[, 6:length(resultarray[1,,1]),3], na.rm = TRUE)),
            na.rm = TRUE) / 365), 365))
    plot(ageplot)

    if(agecorrection == TRUE){
        agecorr <- rep(0, length(mean_window_age[, 2])) # Create vector to store corrections
        # (in days) to the mean ages of modelling windows
        for(i in 1:(length(mean_window_age[,2]) - 1)){ # Loop through all windows
            if(mean_window_age[i + 1, 2] - mean_window_age[i, 2] > T_per / 2){ # If there is a
            # large (> half a year) positive step between consecutive windows...
                # ...subtract one year from the subsequent windows
                agecorr[(i + 1):length(agecorr)] <- agecorr[(i + 1):length(agecorr)] - T_per
            }else if(mean_window_age[i + 1, 2] - mean_window_age[i, 2] < T_per / -2){ # If there is
                # a large (> half a year) negative drop between consecutive windows, add one year
                # from the subsequent windows
                agecorr[(i + 1):length(agecorr)] <- agecorr[(i + 1):length(agecorr)] + T_per
            }
        }

        if(min(mean_window_age$Age + agecorr) > T_per){
            agecorr <- agecorr - 365 # Subtract whole number of years from correction if the
            # smallest value is more than one year old
        }

        mean_window_age$Age_corr <- mean_window_age$Age + agecorr
        mean_window_age$correction <- agecorr

        if(plot == TRUE){ # Plot updated result
            ageplot <- ggplot2::ggplot(mean_window_age, ggplot2::aes(window, Age)) + # Plot Window
            # age to check for strange jumps in age
            ggplot2::geom_line() +
            ggplot2::geom_point() +
            ggplot2::ggtitle("Plot of average ages of modelling windows") +
            ggplot2::xlab("Window #") +
            ggplot2::scale_y_continuous("Age (days)", seq(0, 365 * ceiling(max(unname(colMeans(
                resultarray[, 6:length(resultarray[1,,1]),3], na.rm = TRUE)),
                na.rm = TRUE) / 365), 365)) +
            ggplot2::geom_line(ggplot2::aes(window, Age_corr), col = "red") # Add corrected window
            # age to plot
            plot(ageplot)
        }

        resultarray[, 6:length(resultarray[1, , 1]), 3] <- resultarray[, 6:length(
            resultarray[1, , 1]), 3] + matrix(agecorr, nrow = length(resultarray[, 1, 1]),
            ncol = length(agecorr), byrow = TRUE) # Apply correction on the results of
            # age modelling
    }else{
        resultarray[, 6:length(resultarray[1, , 1]), 3] <- resultarray[, 6:length(
            resultarray[1, , 1]), 3] - floor(min(mean_window_age$Age) / 365) * 365 # Correct age if
            # all windows occur past year one
    }
    return(resultarray)
}


#assert existence of
#commonfile
source(commonfile)
#bamfile
load(bamfile)
#pwmfile
#load(paste0(pwmdir,pwmid,'.pwmout.RData'))
#tmpdir

bamnames = names(allreads[[1]])

obschrnames=names(allreads)
preads=allreads[[obschrnames[1]]][[1]]$plus
cutat=10

#stablize the variance (helps when there are few high coverage sites).
tfun <- function(x){
    y = x
    x[x>cutat]=cutat
    y[x>0] = sqrt(x[x>0])
    y
}

unlink(paste0(tmpdir,'/',pwmid,'/','*tf',pwmid,'*'))

use.w=!is.null(whitelist)
if(use.w){
    wtable=read.table(whitelist)
    white.list=lapply(levels(wtable[,1]),function(i){
        wtchr=wtable[wtable[,1]==i,]
        ir=IRanges(wtchr[,2],wtchr[,3])
    })
    names(white.list)=levels(wtable[,1])
}

dir.create(paste0(tmpdir,'/',pwmid),recursive=T)
makeTFmatrix <- function(coords,prefix='',offset=0){
    cwidth = width(coords[[1]][1])
    obschrnames=names(allreads)
    validchr = obschrnames[which(obschrnames%in%ncoords)]
    readcov=lapply(1:length(bamnames),function(j){
        sapply(validchr,function(i){length(allreads[[i]][[j]]$plus)+length(allreads[[i]][[j]]$minus)})/seqlengths(genome)[validchr]
    })
    readfact = lapply(readcov,function(i){i/i[1]})
    slen = seqlengths(genome)
    scrd =sapply(coords,length)
    minbgs=floor(max(10000,sum(scrd))*(scrd/sum(scrd)));
    for(chr in validchr){
        chrlen = slen[chr]
        print(chr)
        chrcoord=coords[[chr]]
        pos.mat = do.call(rBind,lapply(1:length(bamnames),function(i){
            pluscoord=allreads[[chr]][[i]]$plus
            if(length(pluscoord)>0){
                rre = rle(sort(pluscoord))
                irp=IRanges(start=rre$values,width=1)
                fos=findOverlaps(chrcoord,irp)
                uquery=queryHits(fos)
                querycoord=rre$values[subjectHits(fos)]
                uoffset = querycoord-start(chrcoord)[uquery]+1
                rval= rre$lengths[subjectHits(fos)] / readfact[[i]][chr]
                pos.triple = cbind(round(uquery),round(uoffset),tfun(rval))
                pos.mat=sparseMatrix(i=round(uoffset),j=round(uquery),x=tfun(rval),dims=c(2*wsize,length(chrcoord)),giveCsparse=T)
            }else{
                pos.triple=cbind(1,1,0)
                pos.mat=Matrix(0,nrow=2*wsize,ncol=length(chrcoord))
            }
            pos.mat
        }))
    #
        neg.mat = do.call(rBind,lapply(1:length(bamnames),function(i){
            minuscoord=allreads[[chr]][[i]]$minus
            if(length(minuscoord)>0){
                rre = rle(sort(minuscoord))
                irp=IRanges(start=rre$values,width=1)
                fos=findOverlaps(chrcoord,irp)
                uquery=queryHits(fos)
                querycoord=rre$values[subjectHits(fos)]
                uoffset = querycoord-start(chrcoord)[uquery]+1
                rval= rre$lengths[subjectHits(fos)] / readfact[[i]][chr]
                neg.triple = cbind(round(uquery),round(uoffset),tfun(rval))
                neg.mat=sparseMatrix(i=round(uoffset),j=round(uquery),x=tfun(rval),dims=c(2*wsize,length(chrcoord)),giveCsparse=T)
            }else{
                neg.triple=cbind(1,1,0)
                neg.mat=Matrix(0,nrow=2*wsize,ncol=length(chrcoord))
            }
            neg.mat
        }))
#
        save(pos.mat,neg.mat,file=paste0(tmpdir,'/',pwmid,'/',prefix,'tf',pwmid,'-',chr,'.RData'))
	gc()
    }
}

load(paste0(pwmdir,pwmid,'.pwmout.RData'))
coords2=sapply(coords.short,flank,width=wsize,both=T)
makeTFmatrix(coords2,'positive.')
load(paste0(pwmdir.bg,pwmid,'.pwmout.RData'))
coords2=sapply(coords.short,flank,width=wsize,both=T)
makeTFmatrix(coords2,'background.',10000)

load(paste0(pwmdir,pwmid,'.pwmout.RData'))

#
#####


source('~/project/main.r')
raw_tmp <-arrange(filter(row_refine,車種=='B03'))
raw_tmp <-arrange(filter(raw_tmp,當事者行動狀態==9))
raw_tmp <-arrange(filter(raw_tmp, 事故類型及型態==15)) #事故類型及型態意義=車與車-路口交岔撞 - 15

raw_tmp <- select(raw_tmp,區,死,受傷)
raw_tmp$count <- 1
raw_tmp[is.na(raw_tmp)] <- 0
raw_tmp_summary <- (raw_tmp %>% group_by(區) %>%
                      summarise_all(sum)
)

raw_tmp_summary$死比率 <- raw_tmp_summary$死 / sum(raw_tmp_summary$count) * 100
raw_tmp_summary$受傷比率 <- raw_tmp_summary$受傷 / sum(raw_tmp_summary$count) * 100
raw_tmp_summary$意外比率 <- raw_tmp_summary$count / sum(raw_tmp_summary$count) * 100

highchart() %>%
  hc_chart(type ="column") %>%
  hc_xAxis(categories = raw_tmp_summary$區) %>%
  hc_add_series(data = raw_tmp_summary$死比率, name = "死/總意外比率") %>%
  hc_add_series(data = raw_tmp_summary$受傷比率, name = "受傷/總意外比率") %>%
  hc_add_series(data = raw_tmp_summary$意外比率, name = "意外/總意外比率")


library(rrBLUP)
library(data.table)
args = commandArgs(trailingOnly=TRUE)

X_file <- args[1] # geno matrix
Y_file <- args[2] # pheno matrix
test_file <- args[3] # individuals in test set

geno <- fread(X_file)
test <- scan(test_file, what='character')
training <- geno[geno$ID %in% setdiff(geno$ID,test),]
write.csv(training,'geno_training.csv',quote=F,row.names=F)

Y <- read.csv(Y_file, row.names=1)
Y <- Y[!rownames(Y) %in% test,]
write.csv(Y,'pheno_training.csv',row.names=T,quote=F)


# $ r --file=analyze.r

library(lme4)

tasks <- read.csv("tasks.csv", header = TRUE)

# 'data.frame':  208 obs. of  27 variables:
#  $ participantNumber                   : int  1 1 1 1 1 1 1 1 1 1 ...
#  $ usedMouse                           : Factor w/ 2 levels "no","yes": 1 1 1 1 1 1 1 1 1 1 ...
#  $ taskNumber                          : int  1 2 3 4 5 6 7 8 9 10 ...
#  $ secondEncounter                     : Factor w/ 2 levels "no","yes": 1 1 1 1 2 2 2 2 1 1 ...
#  $ task                                : Factor w/ 6 levels "Four Squares",..: 5 3 4 6 5 4 6 3 1 2 ...
#  $ treatment                           : Factor w/ 3 levels "BoxSelectOnly",..: 3 1 1 1 1 3 3 3 2 2 ...
#  $ allEventsLogged                     : Factor w/ 2 levels "no","yes": 2 2 2 2 2 2 2 2 2 2 ...
#  $ completed                           : Factor w/ 2 levels "no","yes": 2 2 2 2 2 2 2 1 2 1 ...
#  $ timedOut                            : Factor w/ 2 levels "no","yes": 1 1 1 1 1 1 1 2 1 2 ...
#  $ grossTime                           : num  141.6 201.2 118.1 72 59.9 ...
#  $ interactionTime                     : num  91.4 178.8 94.6 52.2 29.2 ...
#  $ textSelectInteractionsStartedCount  : int  3 0 0 0 0 7 4 19 4 10 ...
#  $ textSelectInteractingTime           : num  71.6 0 0 0 0 ...
#  $ textSelectRefactoringsCount         : int  2 0 0 0 0 6 3 5 3 6 ...
#  $ singleArgTextSelectRefactoringsCount: int  1 0 0 0 0 5 1 0 3 6 ...
#  $ multiArgTextSelectRefactoringsCount : int  1 0 0 0 0 1 2 5 0 0 ...
#  $ boxSelectInteractionsStartedCount   : int  0 9 4 3 2 0 0 0 12 54 ...
#  $ boxSelectInteractingTime            : num  0 96.2 68.2 38.6 22.8 ...
#  $ boxSelectRefactoringsCount          : int  0 9 4 2 2 0 0 0 10 33 ...
#  $ singleArgBoxSelectRefactoringsCount : int  0 2 2 0 1 0 0 0 3 6 ...
#  $ multiArgBoxSelectRefactoringsCount  : int  0 7 2 2 1 0 0 0 7 27 ...
#  $ interactionsStartedCount            : int  3 9 4 3 2 7 4 19 16 64 ...
#  $ interactingTime                     : num  71.6 96.2 68.2 38.6 22.8 ...
#  $ refactoringsCount                   : int  2 9 4 2 2 6 3 5 13 39 ...
#  $ undoCount                           : int  0 4 0 0 0 0 0 5 2 39 ...
#  $ redoCount                           : int  0 0 0 0 0 0 0 0 0 3 ...
#  $ invocations                         : Factor w/ 155 levels "","Abstract (rect \"yellowgree... over its constants; Rename 'rect2' to 'oneCorner'; Undo; Undo; Merge 4 rects by "| __truncated__,..: 72 147 111 15 72 109 21 135 81 61 ...

tasks[ , "completionProbability"] <- sapply(tasks$completed, function (yn) { as.double(yn == "yes") })
tasks[ , "firstEncounter"]        <- sapply(tasks$secondEncounter, function (yn) { if(yn=="yes") "no" else "yes" })
tasks[ , "logInteractionTime"]    <- sapply(tasks$interactionTime, log)

headToHeadTasks <- subset(tasks, treatment != "CodeToolsOnly")

completedHeadToHeadTasks <- subset(headToHeadTasks, completed == "yes")

interactionTimeModel <- lmer(interactionTime ~ treatment + (1 | participantNumber) + (1 + secondEncounter | task) + secondEncounter + treatment:secondEncounter + taskNumber + usedMouse + usedOwnComputer, completedHeadToHeadTasks)
summary(interactionTimeModel)

logInteractionTimeModel <- lmer(logInteractionTime ~ treatment + (1 | participantNumber) + (1 + secondEncounter | task) + secondEncounter + treatment:secondEncounter + taskNumber + usedMouse + usedOwnComputer, completedHeadToHeadTasks)
summary(logInteractionTimeModel)

# Swap firstEncounter for secondEncounter to get an easy to interpret p value for treatment on second encounter.
interactionTimeModel <- lmer(interactionTime ~ treatment + (1 | participantNumber) + (1 + firstEncounter | task) + firstEncounter + treatment:firstEncounter + taskNumber + usedMouse + usedOwnComputer, completedHeadToHeadTasks)
summary(interactionTimeModel)

# Swap firstEncounter for secondEncounter to get an easy to interpret p value for treatment on second encounter.
logInteractionTimeModel <- lmer(logInteractionTime ~ treatment + (1 | participantNumber) + (1 + firstEncounter | task) + firstEncounter + treatment:firstEncounter + taskNumber + usedMouse + usedOwnComputer, completedHeadToHeadTasks)
summary(logInteractionTimeModel)

completionProbabilityModel <- glmer(completionProbability ~ treatment + (1 | participantNumber) + (1 + secondEncounter | task) + secondEncounter + treatment:secondEncounter + taskNumber + usedMouse + usedOwnComputer, headToHeadTasks, family=binomial(link="logit"), control=glmerControl(optimizer="bobyqa"))
summary(completionProbabilityModel)

# Swap firstEncounter for secondEncounter to get an easy to interpret p value for treatment on second encounter.
completionProbabilityModel <- glmer(completionProbability ~ treatment + (1 | participantNumber) + (1 + firstEncounter | task) + firstEncounter + treatment:firstEncounter + taskNumber + usedMouse + usedOwnComputer, headToHeadTasks, family=binomial(link="logit"), control=glmerControl(optimizer="bobyqa"))
summary(completionProbabilityModel)

invocationsModel <- lmer(refactoringsCount ~ treatment + (1 | participantNumber) + (1 + secondEncounter | task) + secondEncounter + treatment:secondEncounter + taskNumber + usedMouse + usedOwnComputer, completedHeadToHeadTasks)
summary(invocationsModel)

# Swap firstEncounter for secondEncounter to get an easy to interpret p value for treatment on second encounter.
invocationsModel <- lmer(refactoringsCount ~ treatment + (1 | participantNumber) + (1 + firstEncounter | task) + firstEncounter + treatment:firstEncounter + taskNumber + usedMouse + usedOwnComputer, completedHeadToHeadTasks)
summary(invocationsModel)

interactingTimeModel <- lmer(interactingTime ~ treatment + (1 | participantNumber) + (1 + secondEncounter | task) + secondEncounter + treatment:secondEncounter + taskNumber + usedMouse + usedOwnComputer, completedHeadToHeadTasks)
summary(interactingTimeModel)

# Swap firstEncounter for secondEncounter to get an easy to interpret p value for treatment on second encounter.
interactingTimeModel <- lmer(interactingTime ~ treatment + (1 | participantNumber) + (1 + firstEncounter | task) + firstEncounter + treatment:firstEncounter + taskNumber + usedMouse + usedOwnComputer, completedHeadToHeadTasks)
summary(interactingTimeModel)

# summaryToBootstrap <- function(.) { fixef(.) }
#
# set.seed(101)
# ## 3.8s (on a 5600 MIPS 64bit fast(year 2009) desktop "AMD Phenom(tm) II X4 925"):
# boot01 <- bootMer(interactionTimeModel, summaryToBootstrap, nsim = 100)
#
# ## to "look" at it
# require("boot") ## a recommended package, i.e. *must* be there
# boot01


clump<-function(dat, window = 5000, thres1 = 1e-7, thres2 = 5e-5, listProbes = TRUE){
#' takes all significant association less than thres1 in order of significance as an index site and identifies all other associations within specified window less than thres2. If a less significant probe is clumped with a more significant probe it is excluded for consideration as an index site.
#'
#' @param dat A matrix with the columns ids, pval, chr and pos,
#' @param window Distance in bp of region around site to look for other sinificant associations.
#' @param thres1 P value threshold to select significant sites to clump.
#' @param thres2 P value threshold to select additional sites to clump with site under consideration
#' @return a ma

	tmp<-dat[order(dat$pval),]
	out<-NULL
	while(tmp[1,2] < thres1){
		chrom<-tmp[1,3]
		start<-tmp[1,4]-window
		end<-tmp[1,4]+window
		index<-which(tmp$chr == chrom & tmp$pos <= end & tmp$pos >= start & tmp$pos != tmp[1,4])	## remove self from consideration

		n.sig<-length(which(tmp[index,2] < thres2))
		if(n.sig > 0){
			probes<-tmp[index,1][which(tmp[index,2] < thres2)]
			if(n.sig > 1){
				probes<-paste(probes, collapse = ";")
				}
		} else {
			probes<-NA
		}
		out<-rbind(out, c(unlist(tmp[1,]), length(index), n.sig, probes))
		## remove these probes from furture consideration
		tmp<-tmp[-c(1,index),]
	}
	colnames(out)<-c("Index", "Pvalue", "Chr", "Position", "nSitesinWindow", "nSigSites", "SigDNAmSites")
	return(out)

}


#### continuous_rqa_parameters-pac.r: Part of `politeness_and_coordination.Rmd` ####
#
# This script explores the parameters for the continuous cross-recurrence analysis
# that we'll run over the movement data. Because this is a lengthy process,
# we create a series of files along the way that can be re-run in pieces if needed.
# This allows us to keep this file commented by default.
#
# Written by: A. Paxton (University of California, Berkeley)
# Date last modified: 18 July 2018
#####################################################################################

#### 1. Preliminaries ####

# read in libraries and functions
source('./scripts/03-data_analysis/libraries_and_functions-pac.r')

# prep workspace and libraries
invisible(lapply(c('tseriesChaos',
                   'nonlinearTseries',
                   'crqa',
                   # 'beepr',
                   'quantmod'),
                 require,
                 character.only = TRUE))

# read in conversation_df dataset
conversation_df = read.table('./data/pac-filtered_movement_data.csv',
                             sep=',',header=TRUE)

# spread data to wide form
conversation_df = spread(conversation_df, interlocutor, movement) %>%
  rename(movement_0 = '0',
         movement_1 = '1')

#### 2. Determine delay with average mutual information (AMI) ####

# set maximum AMI
ami.lag.max = 200

# get AMI lag for both participants in each conversation
conversation_df = conversation_df %>% ungroup() %>%
  group_by(participant_id, task, partner_type) %>%
  mutate(ami.loc0 = first_local_minimum(
    mutual(movement_0, lag.max = ami.lag.max, plot = FALSE))
    ) %>%
  mutate(ami.loc1 = first_local_minimum(
    mutual(movement_1, lag.max = ami.lag.max, plot = FALSE))
    ) %>%
  group_by(participant_id, task, partner_type, ami.loc0, ami.loc1) %>%
  distinct() %>%
  mutate(ami.selected = min(ami.loc1, ami.loc0, na.rm = TRUE)) %>%
  ungroup()

# write AMI information to file
amis = conversation_df %>%
  select(participant_id, task, partner_type, contains('ami')) %>%
  distinct()
write.table(amis,'./data/crqa_parameters/ami_calculations-pac.csv', sep=',',row.names=FALSE,col.names=TRUE)

##### 2b. Merging (if we already have amis) #####

# if we've already run it, load it in
amis = read.table('./data/crqa_parameters/ami_calculations-pac.csv', sep=',',header=TRUE)
conversation_df = full_join(conversation_df, amis,
                            by=c('participant_id',
                                 'partner_type',
                                 'task'))

#### 3. Determing embedding dimension with FNN ####

# set maximum percentage of false nearest neighbors
fnnpercent = 10

# create empty dataframe
fnn_dataframe = data.frame()

# cycle through each conversation
convo.dfs = split(conversation_df,
                  list(conversation_df$participant_id,
                       conversation_df$task,
                       conversation_df$partner_type))
for (next_conv in convo.dfs){
  if (dim(data.frame(next_conv))[1] > 0){

    # clean up conversation data
    next_conv_df = data.frame(next_conv)

    # print update
    print(paste0("Beginning FNN calculations for Participant ", unique(next_conv_df$participant_id),
                 ", Task `", unique(next_conv_df$task),"`, Partner `",unique(next_conv_df$partner_type),"`"))

    # get the AMI value for that conversation
    next_ami = unique(next_conv_df$ami.selected)

    # calculate false nearest neighbors for participant 0
    fnn_0 = false.nearest(next_conv_df$movement_0,
                          m = fnnpercent,
                          d = next_ami,
                          t = 0,
                          rt = 10,
                          eps = sd(next_conv_df$movement_0) / 10)
    fnn_0 = fnn_0[1,][complete.cases(fnn_0[1,])]
    threshold_0 = as.numeric(fnn_0[1]/fnnpercent)

    # calculate false nearest neighbors for participant 1
    fnn_1 = false.nearest(next_conv_df$movement_1,
                          m = fnnpercent,
                          d = next_ami,
                          t = 0,
                          rt = 10,
                          eps = sd(next_conv_df$movement_1) / 10)
    fnn_1 = fnn_1[1,][complete.cases(fnn_1[1,])]
    threshold_1 = as.numeric(fnn_1[1]/fnnpercent)

    # identify the largest dimension after a large drop for each participant
    embed_0 = max(as.numeric(which(diff(fnn_0) < -threshold_0))) + 1
    embed_1 = max(as.numeric(which(diff(fnn_1) < -threshold_1))) + 1

    # bind everything to data frame
    next_conv_df = next_conv_df %>% ungroup() %>%
      mutate(embed.0 = embed_0) %>%
      mutate(embed.1 = embed_1) %>%
      mutate(embed.selected = max(embed_0, embed_1))
    fnn_dataframe = rbind.data.frame(fnn_dataframe, next_conv_df)
}}

# rename it to save over the old conversation_df
conversation_df = fnn_dataframe

# save false nearest neighbor calculations to file
fnn.merged = conversation_df %>% ungroup() %>%
  select(participant_id, task, partner_type, contains('embed')) %>%
  distinct()
write.table(fnn.merged,'./data/crqa_parameters/fnn_calculations-pac.csv', sep=',',row.names=FALSE,col.names=TRUE)

##### 3b. Merging (if we already have fnns) #####

# if we've already run it, load it in
fnn.merged = read.table('./data/crqa_parameters/fnn_calculations-pac.csv', sep=',',header=TRUE)
conversation_df = full_join(conversation_df, fnn.merged,
                            by=c('participant_id',
                                 'partner_type',
                                 'task'))

#### 4. Determine optimal radius ####

# rescale by mean distance
conversation_df_crqa = conversation_df %>% ungroup() %>%
  dplyr::select(participant_id, task, partner_type,
                movement_0, movement_1, ami.selected, embed.selected) %>%
  group_by(participant_id, task, partner_type) %>%
  mutate(rescale.movement_0 = movement_0/mean(movement_0)) %>%
  mutate(rescale.movement_1 = movement_1/mean(movement_1))

# cycle through all conversations of all dyads
crqa.data = split(conversation_df_crqa,
                  list(conversation_df$participant_id,
                       conversation_df$task,
                       conversation_df$partner_type))
for (next.conv in crqa.data){

  # make sure we only proceed if we have data for the conversation
  if (dim(next.conv)[1] > 0){

    # reset `target` variables for new radius (above what RR can be)
    from.target = 101
    last.from.target = 102

    # identify parameters for the conversation
    chosen.delay = unique(next.conv$ami.selected)
    chosen.embed = unique(next.conv$embed.selected)

    # start the radius
    chosen.radius = .00

    # arbitrary starting point for `rr`
    rr = 100

    # if we're still improving, keep going
    while ((from.target < last.from.target) | rr == 0){

      # keep the previous iteration's performance
      last.from.target = from.target

      # don't let it continue forever
      if (chosen.radius > 2) {
        break
      }

      # increment radius size
      chosen.radius = chosen.radius + .01

      # print update
      print(paste("Participant ", unique(next.conv$participant_id),
                  ", partner ",unique(next.conv$partner_type),
                  ", task ", unique(next.conv$task),
                  ": radius ",chosen.radius,sep=""))

      # run CRQA and grab recurrence rate (RR)
      rec_analysis = crqa(next.conv$rescale.movement_0,
                          next.conv$rescale.movement_1,
                          delay = chosen.delay,
                          embed = chosen.embed,
                          r = chosen.radius,
                          normalize = 0,
                          rescale = 0,
                          mindiagline = 2,
                          minvertline = 2,
                          tw = 0,
                          whiteline = FALSE,
                          recpt=FALSE)
      rr = rec_analysis$RR

      # clear it so we don't take up too much memory (optional)
      rm(rec_analysis)

      # identify how far off the RR is from our target (5%)
      from.target = abs(rr - 5)

      # save individual radius calculations
      participant_id = unique(next.conv$participant_id)
      task = unique(next.conv$task)
      partner_type = unique(next.conv$partner_type)
      write.table(cbind.data.frame(participant_id,
                                   task,
                                   partner_type,
                                   chosen.delay,
                                   chosen.embed,
                                   chosen.radius,
                                   rr,
                                   from.target),
                  paste('./data/crqa_parameters/radius_calculations/radius_calculations-mean_scaled-r',
                        chosen.radius, '-', participant_id ,'_', task, '_', partner_type, '-pac.csv', sep=''),
                  sep=',',row.names=FALSE,col.names=TRUE)
    }}}

# # let us know when it's finished
# beepr::beep("fanfare")

# concatenate the radius files
radius_files = list.files('./data/crqa_parameters/radius_calculations',
                          pattern='radius_calculations',
                          full.names = TRUE)
radius_selection = data.frame()
for (next_file in radius_files){
  radius_selection = rbind.data.frame(radius_selection,
                                      read.table(next_file, header = TRUE, sep = ','))
}
write.table(radius_selection,'./data/crqa_parameters/radius_calculations-mean_scaled-pac.csv',
            sep=',', row.names=FALSE, col.names=TRUE)

#### 4b. Reload datasets (if we've already run things) ####

# rescale by mean distance
conversation_df_crqa = conversation_df %>% ungroup() %>%
  dplyr::select(participant_id, task, partner_type,
                movement_0, movement_1, ami.selected, embed.selected) %>%
  group_by(participant_id, task, partner_type) %>%
  mutate(rescale.movement_0 = movement_0/mean(movement_0)) %>%
  mutate(rescale.movement_1 = movement_1/mean(movement_1))

#### 5. Expand radius ####

# load back in the data
radius_selection = read.table('./data/crqa_parameters/radius_calculations-mean_scaled-pac.csv',
                              sep=',',header=TRUE)

# get radius yielding closest to 5% (only the largest radius, if more than one)
radius_stats = radius_selection %>% ungroup() %>%
  group_by(participant_id, partner_type, task) %>%
  dplyr::filter(from.target==min(from.target)) %>%
  dplyr::arrange(participant_id, partner_type, task) %>%
  dplyr::arrange(desc(from.target)) %>%
  slice(1) %>%
  distinct()

# link conversation numbers to types
checking_numbers = conversation_df_crqa %>%
  ungroup() %>%
  dplyr::select(participant_id, task, partner_type) %>%
  distinct()
recheck_radii = radius_stats %>%
  dplyr::left_join(x=.,
                   y=checking_numbers,
                   by=c("participant_id"="participant_id",
                        "task"="task",
                        "partner_type"="partner_type")) %>%
  mutate(recheck.conv = paste(participant_id,
                              task,
                              partner_type,
                              sep='.')) %>%
  mutate(from.target = rr - 5)

# cycle through all conversations of all dyads to refine
recheck_conversation_df = dplyr::left_join(recheck_radii,
                                           conversation_df_crqa)

# split the data
crqa.data = split(recheck_conversation_df,
                  list(recheck_conversation_df$participant_id,
                       recheck_conversation_df$task,
                       recheck_conversation_df$partner_type))

# cycle through the conversations
recheck_radii = crqa.data[recheck_radii$recheck.conv]
for (next.conv in recheck_radii){

  # make sure we only proceed if we have data for the conversation
  if (dim(next.conv)[1] != 0){

    # figure out whether we need to have a larger or smaller radius
    original.from.target = unique(next.conv$from.target)
    chosen.radius = unique(next.conv$chosen.radius)
    rr = unique(next.conv$rr)

    # reset `target` variables for new radius (above what RR can be)
    from.target = abs(original.from.target)
    last.from.target = 102

    # identify parameters
    chosen.delay = unique(next.conv$ami.selected)
    chosen.embed = unique(next.conv$embed.selected)

    # if we haven't hit our target threshold or we're still improving, keep going
    while ((from.target > 1) | (from.target < last.from.target)){

      # set some bounds for our radius explorations
      if (chosen.radius < 0 | chosen.radius > 2) { break } # if we've already a bunch
      if (from.target > (last.from.target) + .5) { break } # if we're not improving

      # keep the previous iteration's performance
      last.from.target = from.target

      # update radius based on whether we overshot or undershot target RR
      if (original.from.target > 0){
        chosen.radius = chosen.radius - .001
      } else {
        chosen.radius = chosen.radius + .001
      }

      # print update
      print(paste("Participant ", unique(next.conv$participant_id),
                  ", partner ",unique(next.conv$partner_type),
                  ", task ", unique(next.conv$task),
                  ": radius ",chosen.radius,sep=""))

      # run CRQA and grab recurrence rate (RR)
      rec_analysis = crqa(next.conv$rescale.movement_0,
                          next.conv$rescale.movement_1,
                          delay = chosen.delay,
                          embed = chosen.embed,
                          r = chosen.radius,
                          normalize = 0,
                          rescale = 0,
                          mindiagline = 2,
                          minvertline = 2,
                          tw = 0,
                          whiteline = FALSE,
                          recpt=FALSE)
      rr = rec_analysis$RR

      # clear it so we don't take up too much memory (optional)
      rm(rec_analysis)

      # identify how far off the RR is from our target (5%)
      from.target = abs(rr - 5)

      # save individual radius calculations
      participant_id = unique(next.conv$participant_id)
      task = unique(next.conv$task)
      partner_type = unique(next.conv$partner_type)
      write.table(cbind.data.frame(participant_id,
                                   task,
                                   partner_type,
                                   chosen.delay,
                                   chosen.embed,
                                   chosen.radius,
                                   rr,
                                   from.target),
                  paste('./data/crqa_parameters/radius_calculations/radius_calculations-mean_scaled-r',
                        chosen.radius, '-', participant_id ,'_', task, '_', partner_type, '-pac.csv', sep=''),
                  sep=',',row.names=FALSE,col.names=TRUE)
    }}}

# # let us know when it's finished
# beepr::beep("fanfare")

# concatenate the radius files
radius_files = list.files('./data/crqa_parameters/radius_calculations',
                          pattern='radius_calculations',
                          full.names = TRUE)
radius_selection = data.frame()
for (next_file in radius_files){
  radius_selection = rbind.data.frame(radius_selection,
                                      read.table(next_file, header = TRUE, sep = ','))
}

# save the new version
write.table(radius_selection,'./data/crqa_parameters/radius_recheck_calculations-mean_scaled-pac.csv',
            sep=',', row.names=FALSE, col.names=TRUE)

#### 6. Export chosen radii for all conversations ####

# load in files
radius_selection = read.table('./data/crqa_parameters/radius_recheck_calculations-mean_scaled-pac.csv',
                              sep=',',header=TRUE)

# identify the target radii
radius_stats = radius_selection %>% ungroup() %>%
  group_by(participant_id, partner_type, task) %>%
  dplyr::filter(from.target==min(from.target)) %>%
  dplyr::arrange(participant_id, partner_type, task) %>%
  dplyr::arrange(desc(from.target)) %>%
  slice(1) %>%
  distinct()

# join the dataframes
conversation_df_crqa = full_join(radius_stats,
                                 conversation_df_crqa,
                                 by = c("participant_id",
                                        "task",
                                        "partner_type",
                                        "chosen.embed" = "embed.selected",
                                        "chosen.delay" = "ami.selected"))

# save to files
write.table(conversation_df_crqa,'./data/crqa_data_and_parameters-pac.csv',
            sep=',',row.names=FALSE,col.names=TRUE)
write.table(radius_stats, './data/crqa_parameters-pac.csv',
            sep=',',row.names=FALSE,col.names=TRUE)


library(lme4)
library(dplyr)
library(tidyr)
library(reshape2)
library(ggplot2)

dir_separator = '/'

# dedupe function
dedupe = function(data, verbose = TRUE){
  # Deduplicates columns in data
  # Parameters:
  # data: data you want deduplicated
  # Returns: deduplicated dataframe

  deduped <- data[!duplicated(as.list(data))]
  if (verbose) {
    print("Number of duplicate columns:")
    print((ncol(data) - ncol(deduped)))
  }
  return(deduped)
}

get_qq = function(x) {
  # Returns x values of each data point from a qq plt

  return(qqnorm(x)$x)
}
combine_category = function(f, local_title, save_path) {
  # Combines multiple samples into a single data point
  # For instances in which there are zero values,
  # runs a mixed effect model that imputes the actual value
  # based on the bias of the sample and of other values for that
  # period that are not zero.
  # Parameters:
  # f: filename of csv with samples
  # local_title: region name
  # save_path: directory to save results in

  # print basic setup information
  print(local_title)
  print(paste("Loading: ", f))

  # Loading
  data = read.table(f, header=TRUE, sep = ",", quote = "'")

  # reorder columns
  # API code places the period in a somewhat random spot
  # It also may save the index in the first column, not any actual data
  # So first check if the first column is an integer; if so, drop it
  if (class(data[,1]) == "integer") {
    data = data[, -1]
  }
  # And then reorder
  df = data %>% select(period, everything())
  # and dedupe just in case
  df = dedupe(df)

  # get the region and topic
  filesplit = strsplit(f, dir_separator)
  fn = filesplit[[1]][length(filesplit[[1]])]
  local_fn = substr(fn, 1, nchar(fn)-4)
  print(paste("Working on:", local_fn))

  # report the number of samples; excluding the period column
  print(paste("# samples:", ncol(df)-1 ))

  # Some search terms will simply not return any values
  # So after deduplication there will just be the date column and a single value column
  # No need to do any additional processing on that
  if (ncol(df) == 2) {
    print("Found all zeroes. Just saving this.")
    colnames(df) = c('period_date', 'value')
    write.csv(df, paste0(save_path,local_fn, '_combined.csv'))
  }
  else
  {
    # convert data to long format
    print("Converting to long format")
    print(head(df))
    long = df %>% gather(sample, value, -period)
    long = long %>% group_by(period) %>%
      mutate(qq = get_qq(value),
             value2 = na_if(value, 0),
             log_value = log(value2))
    long$sample = factor(long$sample)
    long$timeperiod = long$period

    # Fit mixed effects model allowing intercept and slope to vary by sample and period
    print("Running Mixed effects model")
    lm1 = lmer(log_value~qq + (1 + qq|period) + (1 + qq|sample), REML = F, data=long)

    # Get coefficients in order to calculate predicted values conditional on week and sample
    print("Determining period and sample effects")
    period_effects = as.data.frame(ranef(lm1)$period) %>% select(int_period=1, qq_period=2)
    period_effects$period = row.names(period_effects)
    sample_effects = as.data.frame(ranef(lm1)$sample) %>% select(int_sample=1, qq_sample=2)
    sample_effects$sample = row.names(sample_effects)

    print("Running Fixed effect model")
    fixed_effects = fixef(lm1)

    long$int_fixed = fixed_effects[1]
    long$qq_fixed = fixed_effects[2]

    print("Imputing..")
    long_betas = long %>% left_join(period_effects) %>% left_join(sample_effects) %>%
      mutate(int_full = int_fixed + int_sample + int_period,
             qq_full = qq_fixed + qq_sample + qq_period,
             pred = int_full + qq_full*qq,  # Predicted value on the log scale
             pred_exp = exp(pred),          # Exponentiated predicted value
             imputed_value = ifelse(is.na(value2), pred_exp, value)) # Use true value if available, else pred_exp
    # save this output
    write.csv(long_betas, paste0(save_path, local_fn, '_imputed_samples.csv'))

    # and convert and take the mean
    print("Taking the mean..")
    condensed = long_betas %>% group_by(period) %>%
      summarize(imp_value=mean(imputed_value),
                avg_value=mean(value)) # take the mean of these
    condensed$period_date = as.Date(condensed$period, format="%Y-%m-%d")
    condensed = arrange(condensed, period_date)

    final = condensed %>% select(period_date, imp_value, avg_value) %>%
      melt(id="period_date")
    print(save_path)
    save_title = paste(save_path, "-", local_title)
    plt = ggplot(final, aes(x = period_date, y=value, color=variable)) + geom_line() +
      ggtitle(local_title)
    plt
    ggsave(paste0(save_path, local_fn, '.png'),device='png')
    write.csv(final, paste0(save_path, local_fn, '_combined.csv'))
  } # else

}

process_region = function(data_directory, region, results_directory="") {
  # Using a list of samples of Google Health calls contained in csv files in data_dictionary,
  # combine multiple samples into single data point and save results
  # Most of the combining work is actually done in combine_category
  # Parameters:
  # data_directory: parent directory above where data is stored
  # region: shorthand for region. DMA code for DMAs, state (e.g. US-CA) or country (e.g. US) code
  # results_directory: where to save the combined results

  current_wd <- getwd()

  # Scan directory for files
  setwd(data_directory)
  if (results_directory == "") { results_directory = paste0(region, '/combined/') }
  filenames = list.files(region, pattern='*.csv', full.names=TRUE)

  # run each file
  for (f in filenames) {
    if (startsWith(f, paste0(region, '/'))) {
      fn = substring(f, nchar(region)+2, nchar(f))
    }
    else { fn = f }
    local_title = substr(fn, 1, nchar(f)-4) # take the .csv off
    if (!dir.exists(results_directory)) { dir.create(results_directory) }
    combine_category(f, local_title, results_directory)
    print(f)
  } # for loop
  setwd(current_wd)
}

# run code
# Assumptions:
# - Starting directory (below) is the parent directory in which data is stored
# - for each region, data is stored in a folder named after the region - either the DMA or the state/country code
# - under that, each file is a series of samples run by the Search Sampler python package
starting_dir = '~/repos/flint_water/Search/DATA/'

process_region(starting_dir, '513')
process_region(starting_dir, 'US-MI')
process_region(starting_dir, 'US')


library(ggplot2)
qplot(factor(names(prob), levels = names(prob)), hebrew, geom = "histogram")


library(sf)
library(mapview)
library(gdtools)
options(stringsAsFactors = FALSE)
ap = "Lelystad Airport"
cd = "LEY"
class(cd)

lat = 52.460278
lon = 5.527222
class(lon)

airport = data.frame(ap, cd, lat, lon, stringsAsFactors = FALSE)

NL_Airports = read.csv(file = 'http://www.twiav.nl/files/NL_Airports.csv')

names(airport) = names(NL_Airports)
NL_Airports = rbind(NL_Airports, airport)
str(NL_Airports)

NL_Airports = st_as_sf(NL_Airports, coords = c('longitude', 'latitude'), crs = 4326)
class(NL_Airports)
NL_Airports

plot(st_geometry(NL_Airports), main = "Airports in the Netherlands", pch = 17)

st_write(NL_Airports, "NL_Airports.geojson", append = FALSE)

mapview(NL_Airports, color = "red", col.regions = "orange", alpha.regions = 1, label =NL_Airports$airport)

GR_Airports = read.csv(file = "C:\\Users\\kougi\\Downloads\\gr-airports.csv")
GR_Airports = GR_Airports[-1,]
GR_Airports_sm = GR_Airports[, c("name","iata_code", "longitude_deg","latitude_deg")]

names(GR_Airports_sm)[names(GR_Airports_sm) == 'longitude_deg'] <- 'longitude'
names(GR_Airports_sm)[names(GR_Airports_sm) == 'latitude_deg'] <- 'latitude'
str(GR_Airports_sm)

GR_Airports_sm = st_as_sf(GR_Airports_sm, coords = c('longitude', 'latitude'), crs = 4326)
GR_Airports_sm

mapview(GR_Airports_sm, color = "red", col.regions = "orange", alpha.regions = 1, label =GR_Airports_sm$name)


#===============================================================================
#  File:    06-additional-results.R
#  Date:    Feb 3, 2021
#  Purpose: replicate appendix analyses: Sections 5 and 6
#  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")

########################################################
### Issue H1: issue opinions, Fox News treatment
########################################################

trt <- "FoxNews"
dv <- "issue_scale_w5"
dv_pre <- "issue_scale_pre"
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H1: Conservatism")

########################################################
### Issue H2: issue opinions, HuffPost treatment
########################################################

trt <- "HuffPost"
dv <- "issue_scale_w5"
dv_pre <- "issue_scale_pre"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H2: Conservatism")

########################################################
### Issue H3: immigration attitudes, FoxNews treatment
########################################################

trt <- "FoxNews"
dv <- "imm_issue_scale_w5"
dv_pre <- "issue_scale_pre"
D <- "comp_fn"

## 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, more_vars = "policy14_imm_pre")
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H3: Pro-immigration")

########################################################
### Issue H4: immigration attitudes, HuffPost treatment
########################################################

trt <- "HuffPost"
dv <- "imm_issue_scale_w5"
dv_pre <- "issue_scale_pre"
D <- "comp_hp"

## 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, more_vars = "policy14_imm_pre")
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H4: Pro-immigration")

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

# H1
heterogeneous_effect(dv = "issue_scale_w5", dv_pre = "issue_scale_pre",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "issue_scale_w5", dv_pre = "issue_scale_pre",
                     moderator = "ideo", trt = "FoxNews")
# H2
heterogeneous_effect(dv = "issue_scale_w5", dv_pre = "issue_scale_pre",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "issue_scale_w5", dv_pre = "issue_scale_pre",
                     moderator = "ideo", trt = "HuffPost")
# H3
heterogeneous_effect(dv = "imm_issue_scale_w5", dv_pre = "issue_scale_pre",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "imm_issue_scale_w5", dv_pre = "issue_scale_pre",
                     moderator = "ideo", trt = "FoxNews")
# H4
heterogeneous_effect(dv = "imm_issue_scale_w5", dv_pre = "issue_scale_pre",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "imm_issue_scale_w5", dv_pre = "issue_scale_pre",
                     moderator = "ideo", trt = "HuffPost")

########################################################
### Agenda H1: FoxNews treatment
########################################################

trt <- "FoxNews"
dv <- "agenda_lean_w5"
dv_pre <- "agenda_lean_pre"
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H1: Agenda setting")

########################################################
### Agenda H2: HuffPost treatment
########################################################

trt <- "HuffPost"
dv <- "agenda_lean_w5"
dv_pre <- "agenda_lean_pre"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace,  trt, pwr,
              dv = "H2: Agenda setting")

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

# H1
heterogeneous_effect(dv = "agenda_lean_w5", dv_pre = "agenda_lean_pre",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "agenda_lean_w5", dv_pre = "agenda_lean_pre",
                     moderator = "ideo", trt = "FoxNews")

# H2
heterogeneous_effect(dv = "agenda_lean_w5", dv_pre = "agenda_lean_pre",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "agenda_lean_w5", dv_pre = "agenda_lean_pre",
                     moderator = "ideo", trt = "HuffPost")

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

# H1
heterogeneous_effect(dv = "agenda_lean_w5", dv_pre = "log_fn_pre",
                     moderator = "ideo", trt = "FoxNews")
# H2
heterogeneous_effect(dv = "agenda_lean_w5", dv_pre = "log_hp_pre",
                     moderator = "ideo", trt = "HuffPost")

########################################################
### Approval H1a: President
########################################################

trt <- "FoxNews"
dv <- "trump_approve_w5"
dv_pre <- "trump_approve_pre"
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H1a: Pres. approval")

########################################################
### Approval H2a: President
########################################################

trt <- "HuffPost"
dv <- "trump_approve_w5"
dv_pre <- "trump_approve_pre"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H2a: Pres. approval")

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

# H1a
heterogeneous_effect(dv = "trump_approve_w5", dv_pre = "trump_approve_pre",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "trump_approve_w5", dv_pre = "trump_approve_pre",
                     moderator = "ideo", trt = "FoxNews")

# H2a
heterogeneous_effect(dv = "trump_approve_w5", dv_pre = "trump_approve_pre",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "trump_approve_w5", dv_pre = "trump_approve_pre",
                     moderator = "ideo", trt = "HuffPost")

########################################################
### Turnout RQ1a: Fox News
########################################################

trt <- "FoxNews"
dv <- "vote_after"
dv_pre <- "vote_pre"
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "RQ1a: Turnout")

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

trt <- "HuffPost"
dv <- "vote_after"
dv_pre <- "vote_pre"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "RQ1b: Turnout")

########################################################
### Media Trust H1a: Trust in Fox News for FN treatment
########################################################

trt <- "FoxNews"
dv <- "trust_fox_w5"
dv_pre <- ""
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H1a: Trust in Fox News")

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

trt <- "FoxNews"
dv <- "trust_hp_w5"
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H1b: Trust in HuffPost")

########################################################
### Media Trust H2a: Trust in HuffPost for HP treatment
########################################################

trt <- "HuffPost"
dv <- "trust_hp_w5"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H2a: Trust in HuffPost")

########################################################
### Media Trust H2b: Trust in Fox News for HP treatment
########################################################

trt <- "HuffPost"
dv <- "trust_fox_w5"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H2b: Trust in Fox News")

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

# H1a
heterogeneous_effect(dv = "trust_fox_w5",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "trust_fox_w5",
                     moderator = "ideo", trt = "FoxNews")
# H1b
heterogeneous_effect(dv = "trust_hp_w5",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "trust_hp_w5",
                     moderator = "ideo", trt = "FoxNews")
# H2a
heterogeneous_effect(dv = "trust_hp_w5",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "trust_hp_w5",
                     moderator = "ideo", trt = "HuffPost")
# H2b
heterogeneous_effect(dv = "trust_fox_w5",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "trust_fox_w5",
                     moderator = "ideo", trt = "HuffPost")

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

trt <- "FoxNews"
dv <- "trust_w7"
dv_pre <- ""
D <- "comp_fn"

# svy$trust_w7 <- zap_labels(svy$trust_w7)

# resetting after het fx analysis
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")

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H4a: Media trust (W7)")

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

# svy$trust_w7 <- svy$trust_w7/sd(svy$trust_w7[which(svy$W3_PATA306_treatment_w3 == "Control")], na.rm = TRUE)

trt <- "HuffPost"
dv <- "trust_w7"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H4b: Media trust (W7)")

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

trt <- "HuffPost"
dv <- "event_mokken_w5"
dv_pre <- "event_pre_mokken"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "RQ1a: Event knowledge")

########################################################
### Factual knowledge RQ1b: news reception for FN group
########################################################

trt <- "FoxNews"
dv <- "event_mokken_w5"
dv_pre <- "event_pre_mokken"
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "RQ1b: Event knowledge")

# RQ1a
heterogeneous_effect(dv = "event_mokken_w5", dv_pre = "event_pre_mokken",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "event_mokken_w5", dv_pre = "event_pre_mokken",
                     moderator = "ideo", trt = "HuffPost")

# RQ1b
heterogeneous_effect(dv = "event_mokken_w5", dv_pre = "event_pre_mokken",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "event_mokken_w5", dv_pre = "event_pre_mokken",
                     moderator = "ideo", trt = "FoxNews")

#### WAVE 7 ####

########################################################
### Issue H1: issue opinions, Fox News treatment (w7)
########################################################

trt <- "FoxNews"
dv <- "issue_scale_w7"
dv_pre <- "issue_scale_pre"
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H1: Conservatism")

########################################################
### Issue H2: issue opinions, HuffPost treatment (w7)
########################################################

trt <- "HuffPost"
dv <- "issue_scale_w7"
dv_pre <- "issue_scale_pre"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H2: Conservatism")

########################################################
### Issue H3: immigration attitudes, FoxNews treatment (w7)
########################################################

trt <- "FoxNews"
dv <- "imm_issue_scale_w7"
dv_pre <- "issue_scale_pre"
D <- "comp_fn"

## 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, more_vars = "policy14_imm_pre")
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H3: Pro-immigration")

########################################################
### Issue H4: immigration attitudes, HuffPost treatment (w7)
########################################################

trt <- "HuffPost"
dv <- "imm_issue_scale_w7"
dv_pre <- "issue_scale_pre"
D <- "comp_hp"

## 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, more_vars = "policy14_imm_pre")
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H4: Pro-immigration")

########################################################
### Issue RQ1: heterogeneous treatment effects (w7)
########################################################

# H1
heterogeneous_effect(dv = "issue_scale_w7", dv_pre = "issue_scale_pre",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "issue_scale_w7", dv_pre = "issue_scale_pre",
                     moderator = "ideo", trt = "FoxNews")
# H2
heterogeneous_effect(dv = "issue_scale_w7", dv_pre = "issue_scale_pre",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "issue_scale_w7", dv_pre = "issue_scale_pre",
                     moderator = "ideo", trt = "HuffPost")
# H3
heterogeneous_effect(dv = "imm_issue_scale_w7", dv_pre = "issue_scale_pre",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "imm_issue_scale_w7", dv_pre = "issue_scale_pre",
                     moderator = "ideo", trt = "FoxNews")
# H4
heterogeneous_effect(dv = "imm_issue_scale_w7", dv_pre = "issue_scale_pre",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "imm_issue_scale_w7", dv_pre = "issue_scale_pre",
                     moderator = "ideo", trt = "HuffPost")

########################################################
### Agenda H1: FoxNews treatment (w8)
########################################################

trt <- "FoxNews"
dv <- "agenda_lean_w8"
dv_pre <- "agenda_lean_pre"
D <- "comp_fn"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H1: Agenda setting")

########################################################
### Agenda H2: HuffPost treatment (w8)
########################################################

trt <- "HuffPost"
dv <- "agenda_lean_w8"
dv_pre <- "agenda_lean_pre"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace,  trt, pwr,
              dv = "H2: Agenda setting")

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

# H1
heterogeneous_effect(dv = "agenda_lean_w8", dv_pre = "agenda_lean_pre",
                     moderator = "party7", trt = "FoxNews")
heterogeneous_effect(dv = "agenda_lean_w8", dv_pre = "agenda_lean_pre",
                     moderator = "ideo", trt = "FoxNews")

# H2
heterogeneous_effect(dv = "agenda_lean_w8", dv_pre = "agenda_lean_pre",
                     moderator = "party7", trt = "HuffPost")
heterogeneous_effect(dv = "agenda_lean_w8", dv_pre = "agenda_lean_pre",
                     moderator = "ideo", trt = "HuffPost")

########################################################
### Agenda RQ2: heterogeneous treatment effects by pre-treatment habits (w8)
########################################################

# H1
heterogeneous_effect(dv = "agenda_lean_w8", dv_pre = "log_fn_pre",
                     moderator = "ideo", trt = "FoxNews")
# H2
heterogeneous_effect(dv = "agenda_lean_w8", dv_pre = "log_hp_pre",
                     moderator = "ideo", trt = "HuffPost")

########################################################
### Media Trust H4a: Media trust for FN group (w8)
########################################################

# svy$trust_w8 <- svy$trust_w8/sd(svy$trust_w8[which(svy$W3_PATA306_treatment_w3 == "Control")], na.rm = TRUE)

trt <- "FoxNews"
dv <- "trust_w8"
D <- "comp_fn"

# svy$trust_w8 <- zap_labels(svy$trust_w8)

# resetting after het fx analysis
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")

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H4a: Media trust (W8)")

########################################################
### Media Trust H4b: Media trust for HP group (w8)
########################################################

trt <- "HuffPost"
dv <- "trust_w8"
D <- "comp_hp"

## 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)
compute_proportion_missing_covars(itt)
## CACE
cace <- estimate_cace(Y=dv, D = D, Z = "W3_PATA306_treatment_w3",
                      X = extract_covariates(itt), trt=trt)

# power analysis
pwr <- power2(dv, dim, itt, cace, covariates = TRUE, D = D, trt = trt)

format_latex2(dim, itt, cace, trt, pwr,
              dv = "H4b: Media trust (W8)")


ANCOM <- function (real.data, sig = 0.05, multcorr = 3, tau = 0.02, theta = 0.1)
{
    #no_cores <- detectCores() - 1

    colnames(real.data)[dim(real.data)[2]] <- "Group"
    real.data$Group <- factor(real.data$Group)
    real.data <- data.frame(real.data[which(is.na(real.data$Group) ==
        F), ], row.names = NULL)
    num_OTU <- ncol(real.data) - 1
    W.detected <- ancom.detect(real.data, num_OTU, sig, multcorr,
        ncore = ncores)
    W_stat <- W.detected
    if (num_OTU < 10) {
        detected <- colnames(real.data)[which(W.detected > num_OTU -
            1)]
    }
    else {
        if (max(W.detected)/num_OTU >= theta) {
            c.start <- max(W.detected)/num_OTU
            cutoff <- c.start - c(0.05, 0.1, 0.15, 0.2, 0.25)
            prop_cut <- rep(0, length(cutoff))
            for (cut in 1:length(cutoff)) {
                prop_cut[cut] <- length(which(W.detected >= num_OTU *
                  cutoff[cut]))/length(W.detected)
            }
            del <- rep(0, length(cutoff) - 1)
            for (ii in 1:(length(cutoff) - 1)) {
                del[ii] <- abs(prop_cut[ii] - prop_cut[ii + 1])
            }
            if (del[1] < tau & del[2] < tau & del[3] < tau) {
                nu = cutoff[1]
            }
            else if (del[1] >= tau & del[2] < tau & del[3] <
                tau) {
                nu = cutoff[2]
            }
            else if (del[2] >= tau & del[3] < tau & del[4] <
                tau) {
                nu = cutoff[3]
            }
            else {
                nu = cutoff[4]
            }
            up_point <- min(W.detected[which(W.detected >= nu *
                num_OTU)])
            W.detected[W.detected >= up_point] <- 99999
            W.detected[W.detected < up_point] <- 0
            W.detected[W.detected == 99999] <- 1
            detected <- colnames(real.data)[which(W.detected ==
                1)]
        }
        else {
            W.detected <- 0
            detected <- "No significant OTUs detected"
        }
    }
    results <- list(W = W_stat, detected = detected, dframe = real.data)
    class(results) <- "ancom"
    return(results)
}


#Usage:
#Rscript cell_cluster.r cluster_number
library("monocle")
library("reshape")
library('ggplot2')
args <- commandArgs(trailingOnly=TRUE)
if(length(args)==0)
{
	print("Usage: Rscript cell_cluster.r cluster_number cluster_labels")
}
cluster_number <- args[1]
HSMM <- readRDS('post_quality_control_data.rds')
HSMM <- clusterCells(HSMM, num_clusters=as.numeric(cluster_number), tol = 1e-6,max_components = 10, verbose=TRUE,param.gamma=300)
saveRDS(HSMM,'post_quality_control_data.rds')
print(plot_cell_trajectory(HSMM, color="Cluster"))
ggsave("cell_cluster_plot.png")
ggsave("cell_cluster_plot.tiff", width = 7, height = 7,dpi=600)
pData(HSMM)$Size_Factor <- NULL
rownames(pData(HSMM)) <- NULL
pData(HSMM)$Total_mRNAs <- Matrix::colSums(exprs(HSMM))
write.table(pData(HSMM), 'cluster_table_after.txt', sep="\t",quote=FALSE)
dev.off()


# From Trevor Thompson tkt2@cdc.gov

require(rms)
if(FALSE) {
require(rms)

## simulate competing risk data with 2 predictors
set.seed(115)
dat1<- crisk.sim(n=2000, foltime=200, dist.ev=c("weibull","weibull"), anc.ev=c(0.8,0.9), beta0.ev=c(2,2), dist.cens="weibull",
                 anc.cens=1,beta0.cens=1, z=NULL, beta=list(c(-0.69,0), c(-0.35, -0.5)), x=list(c("bern", 0.3), c("bern", 0.6)), nsit=2)

dat1$event<-factor(ifelse(is.na(dat1$cause), 0, dat1$cause))
dat1$x<-factor(dat1$x)
dat1$x.1<-factor(dat1$x.1)

## set some predictor data to missing
miss1<-rbinom(2000, 1, .15)
miss2<-rbinom(2000, 1, .05)

dat1[miss1==1,]$x<-NA
dat1[miss2==1,]$x.1<-NA

describe(dat1)

## impute missing data
imp.obj <- aregImpute(~ x + x.1 + event + time, n.impute=20, data=dat1, x=TRUE)
imp.obj

## create Fine-Gray function for dtrans
dtrans.fg <-
  function(data) finegray(Surv(time, event) ~ ., data=data, etype=1)

chk <- dtrans.fg(dat1)
dim(chk)

## fit model with imputed data
mod.fg <-
  fit.mult.impute(Surv(fgstart, fgstop, fgstatus) ~ x + x.1,
                  cph, data=dat1, xtrans=imp.obj, dtrans=dtrans.fg,
                  weights=fgwt, fit.reps=TRUE, x=TRUE, y=TRUE, surv=TRUE)
# Problem: fgwt is in data created by dtrans, not in dat1
}

x <- runif(10)
y <- runif(10)
x[1] <- NA
ww = seq(0.01, 1, length=10)
d <- data.frame(x, y, ww)
a <- aregImpute(~ x + y, data=d, nk=0)

dt <- function(dat) cbind(dat, ww=1:10)

f <- fit.mult.impute(Surv(y) ~ x, cph, weights=ww,
                     data=d, xtrans=a, dtrans=dt)