#' ---
#' title: "Prior probabilities in the interpretation of 'some': analysis of uniform prior wonky world model predictions"
#' author: "Judith Degen"
#' date: "January 26, 2014"
#' ---

library(ggplot2)
library(wesanderson)
theme_set(theme_bw(18))
setwd("/Users/titlis/cogsci/projects/stanford/projects/thegricean_sinking-marbles/writing/_2015/_journal_cognition/")
source("rscripts/helpers.r")

## load priors for generating plots
priorprobs = read.table(file="~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/12_sinking-marbles-prior15/results/data/smoothed_15marbles_priors_withnames.txt",sep="\t", header=T, quote="")
row.names(priorprobs) = priorprobs$Item
nrow(priorprobs)

# prior all-state probability for each item
prior_allprobs = priorprobs[,c("Item","X15")]
row.names(prior_allprobs) = prior_allprobs$Item

# prior expectation for each item
gathered_probs <- priorprobs %>%
  gather(State,Probability,X0:X15)
gathered_probs$State = as.numeric(as.character(gsub("X","",gathered_probs$State)))
prior_exps <- gathered_probs %>%
  group_by(Item) %>%
  summarise(exp.val=sum(State*Probability))
prior_exps = as.data.frame(prior_exps)
row.names(prior_exps) = prior_exps$Item
summary(prior_exps)

# histogram of expectations
exps = ggplot(prior_exps, aes(x=exp.val)) +
  geom_histogram() +
  scale_x_continuous(name="Expected value of prior distribution",breaks=seq(1,15, by=2)) +
  scale_y_continuous(name="Number of cases",breaks=seq(0,8, by=2))
ggsave("pics/priorexpectations-histogram.pdf",width=5,height=3.7)

# histogram of allstate-probs
allprobs = ggplot(prior_allprobs, aes(x=X15)) +
  geom_histogram() +
  scale_x_continuous(name="Prior all-state probability") +
  scale_y_continuous(name="Number of cases")
ggsave("priorallprobs-histogram.pdf")

pdf("pics/priordistributions.pdf",width=10,height=3.5)
grid.arrange(exps, allprobs, nrow=1,widths=unit.c(unit(.5, "npc"), unit(.5, "npc")))
dev.off()

# load inferred priors (with guessing parameter)
exps_inferred_phi = read.table(file="~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/bayesian_model_comparison/priors/results/summary_priorBDA-binomials-phi_incrMH100000burn50000.csv",sep=",", header=T, quote="")
exps_inferred_phi$X.Item. = gsub("\"","",exps_inferred_phi$X.Item.)
row.names(exps_inferred_phi) = exps_inferred_phi$X.Item.
nrow(exps_inferred_phi)
head(exps_inferred_phi)

# histogram of expectations
exps = ggplot(exps_inferred_phi, aes(x=X.expectation.*15)) +
  geom_histogram() +
  scale_x_continuous(name="Expected value of prior distribution",breaks=seq(1,15, by=2)) +
  scale_y_continuous(name="Number of cases")
ggsave("pics/priorexpectations-inferred-phi-histogram.pdf",width=5,height=3.7)

# load inferred priors (without guessing parameter)
exps_inferred = read.table(file="~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/bayesian_model_comparison/priors/results/summary_priorBDA-binomials_incrMH100000burn50000.csv",sep=",", header=T, quote="")
exps_inferred$X.Item. = gsub("\"","",exps_inferred$X.Item.)
row.names(exps_inferred) = exps_inferred$X.Item.
nrow(exps_inferred)
head(exps_inferred)
row.names(exps_inferred) = exps_inferred$X.Item.

# histogram of expectations
exps = ggplot(exps_inferred, aes(x=X.expectation.*15)) +
  geom_histogram() +
  scale_x_continuous(name="Expected value of prior distribution",breaks=seq(1,15, by=2)) +
  scale_y_continuous(name="Number of cases")
ggsave("pics/priorexpectations-inferred-histogram.pdf",width=5,height=3.7)

### PLOT COMPREHENSION RESULTS
# expectations
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/13_sinking-marbles-priordv-15/results/data/r.RData")

r$PriorExpectationProportion = prior_exps[as.character(r$Item),]$exp.val
r$PriorExpectationProportionInferred = exps_inferred[as.character(r$Item),]$X.expectation.
r$PriorExpectationProportionInferredPhi = exps_inferred_phi[as.character(r$Item),]$X.expectation.

cor(r$PriorExpectationProportionInferred*15,r$PriorExpectationProportion)
cor(r$PriorExpectationProportionInferredPhi*15,r$PriorExpectationProportion)
# exclude people who weren't literal above some threshold
# lit = droplevels(r[r$quantifier %in% c("All","None"),]) %>%
#   group_by(workerid,quantifier) %>%
#   summarise(Mean=mean(response))
# lit = as.data.frame(lit)
# lit = lit %>% spread(quantifier,Mean)
# lit$Literal = as.factor(as.character(ifelse(lit$All > .80*15 & lit$None < .20*15,"literal","non-literal")))
# table(lit$Literal)
# row.names(lit) = as.character(lit$workerid)
# r$Literal = lit[as.character(r$workerid),]$Literal

agr = aggregate(ProportionResponse ~ PriorExpectationProportion + quantifier + Item,data=r,FUN=mean)

min(agr[agr$quantifier == "Some",]$ProportionResponse)*15
max(agr[agr$quantifier == "Some",]$ProportionResponse)*15
agr$Quantifier = factor(x=agr$quantifier, levels=c("Some","All","None","long_filler","short_filler"))

p_eexps = ggplot(agr, aes(x=PriorExpectationProportion, y=ProportionResponse*15, color=Quantifier, shape=Quantifier)) +
  geom_point() +
  #geom_errorbar(aes(ymin=YMin,ymax=YMax)) +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +#values=c("#F8766D", "black", "#00BF7D", "gray30", "#00B0F6"),breaks=levels(agr$quantifier),labels=c("all","long filler","none","short filler","some")) +
  scale_y_continuous(breaks=seq(1,15,by=2),name="Posterior mean number of objects") +
  #  geom_abline(intercept=0,slope=1,color="gray70") +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects")
p_eexps
ggsave(file="empirical_exps.pdf")#,width=5,height=3.7)

# plot with inferred prior expectation on x axis
agr = aggregate(ProportionResponse ~ PriorExpectationProportionInferred + quantifier + Item,data=r,FUN=mean)

min(agr[agr$quantifier == "Some",]$ProportionResponse)*15
max(agr[agr$quantifier == "Some",]$ProportionResponse)*15
agr$Quantifier = factor(x=agr$quantifier, levels=c("Some","All","None","long_filler","short_filler"))

p_eexps = ggplot(agr, aes(x=PriorExpectationProportionInferred*15, y=ProportionResponse*15, color=Quantifier, shape=Quantifier)) +
  geom_point() +
  #geom_errorbar(aes(ymin=YMin,ymax=YMax)) +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +#values=c("#F8766D", "black", "#00BF7D", "gray30", "#00B0F6"),breaks=levels(agr$quantifier),labels=c("all","long filler","none","short filler","some")) +
  scale_y_continuous(breaks=seq(1,15,by=2),name="Posterior mean number of objects") +
  #  geom_abline(intercept=0,slope=1,color="gray70") +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects")
p_eexps
ggsave(file="pics/empirical_exps_inferredprior.pdf",width=7)#,width=5,height=3.7)

# plot with inferred prior (phi) expectation on x axis
agr = aggregate(ProportionResponse ~ PriorExpectationProportionInferredPhi + quantifier + Item,data=r,FUN=mean)

min(agr[agr$quantifier == "Some",]$ProportionResponse)*15
max(agr[agr$quantifier == "Some",]$ProportionResponse)*15
agr$Quantifier = factor(x=agr$quantifier, levels=c("Some","All","None","long_filler","short_filler"))

p_eexps = ggplot(agr, aes(x=PriorExpectationProportionInferredPhi*15, y=ProportionResponse*15, color=Quantifier, shape=Quantifier)) +
  geom_point() +
  #geom_errorbar(aes(ymin=YMin,ymax=YMax)) +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +#values=c("#F8766D", "black", "#00BF7D", "gray30", "#00B0F6"),breaks=levels(agr$quantifier),labels=c("all","long filler","none","short filler","some")) +
  scale_y_continuous(breaks=seq(1,15,by=2),name="Posterior mean number of objects") +
  #  geom_abline(intercept=0,slope=1,color="gray70") +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects")
p_eexps
ggsave(file="pics/empirical_exps_inferredprior_phi.pdf",width=7)#,width=5,height=3.7)

# allstate-probs
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/16_sinking-marbles-sliders-certain/results/data/r.RData")

# exclude people who are doing some sort of bullshit and not responding reasonably to all/none (see subject-variability.pdf for behavior on zero-slider)
#tmp = subset(r,!workerid %in% c(0,22,43,98,100,103,117,118))
r$AllPriorProbability = prior_allprobs[as.character(r$Item),]$X15
tmp=r
#tmpub = droplevels(subset(tmp, Proportion == "100"))
#tmpub$Quantifier = factor(x=tmpub$quantifier, levels=c("Some","All","None","long_filler","short_filler"))
agrr = aggregate(normresponse ~ AllPriorProbability + Proportion + quantifier + Item,data=r,FUN=mean)
ub = subset(agrr, Proportion == "100")
ub = droplevels(ub)
ub$Quantifier = factor(x=ub$quantifier, levels=c("Some","All","None","long_filler","short_filler"))

p_eprobs = ggplot(ub, aes(x=AllPriorProbability, y = normresponse, color=Quantifier, shape=Quantifier)) +
  geom_point() +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +#values=c("#F8766D", "black", "#00BF7D", "gray30", "#00B0F6"),breaks=levels(agr$quantifier),labels=c("all","long filler","none","short filler","some")) +
  scale_y_continuous(limits=c(0,1),breaks=seq(0,1,by=.2),name="Posterior probability of all-state ") +
  #  geom_abline(intercept=0,slope=1,color="gray70") +
  scale_x_continuous(name="Prior probability of all-state")
p_eprobs
ggsave(file="empirical-allprobs.pdf")#,width=5,height=3.7)

library(gridExtra)
#  share a legend between multiple plots
g <- ggplotGrob(p_eexps + theme(legend.position="right"))$grobs
legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
p_eexps_nolegend = p_eexps + theme(legend.position="none")
p_eprobs_nolegend = p_eprobs + theme(legend.position="none")

pdf("pics/empirical-results.pdf",width=13,height=5)
grid.arrange(p_eexps_nolegend, p_eprobs_nolegend, legend,nrow=1,widths=unit.c(unit(.45, "npc"), unit(.45, "npc"), unit(.1, "npc")))
dev.off()

### PLOT WONKY RESULTS
load("/Users/titlis/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/11_sinking-marbles-normal/results/data/r.RData")

r$PriorExpectation = prior_exps[as.character(r$Item),]$exp.val
ans = r %>%
  filter(quantifier %in% c("All","None","Some")) %>%
  group_by(quantifier, Item, PriorExpectation) %>%
  summarise(mean.prop=mean(numWonky))

p_wempirical =ggplot(ans,aes(x=PriorExpectation,y=mean.prop,color=quantifier)) +
  geom_point() +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3])) +
  scale_y_continuous(limits=c(-0.1,1.2),breaks=seq(0,1,by=.2),name="Proportion of `wonky' judgments") +
  scale_x_continuous(limits=c(0,15),breaks=seq(1,15,by=2),name="Prior mean number of objects")

# get wonky model predictions
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/wr-uniform.RData")

# plot expectations for best basic model:
toplot = droplevels(subset(wr, WonkyWorldPrior == .5 & Wonky == "true"))
nrow(toplot)

toplot = droplevels(subset(toplot, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)
#toplot$quantifier = capitalize(as.character(toplot$Quantifier))
toplot$Quantifier = factor(toplot$Quantifier, levels=c("all","none","some"))

p_wmodel = ggplot(toplot, aes(x=PriorExpectation, y=PosteriorProbability,color=Quantifier,shape=Quantifier)) +
  geom_point() +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +
  #scale_color_manual(values=c("#F8766D","#00BF7D","#00B0F6")) +
  scale_x_continuous(limits=c(0,15),breaks=seq(1,15,by=2),name="Prior mean number of objects") +
  scale_y_continuous(limits=c(-0.1,1.2),breaks=seq(0,1,by=.25),name="Predicted posterior wonkiness probability") +
  theme(plot.margin=unit(c(0,0,0,0),units="cm"))
p_wmodel

pdf("pics/wonkiness-results.pdf",width=10,height=4)
#grid.arrange(p_wmodel, p_wempirical, legend,nrow=1,widths=unit.c(unit(.45, "npc"), unit(.45, "npc"), unit(.1, "npc")))
grid.arrange(p_wmodel, p_wempirical, nrow=1,widths=unit.c(unit(.5, "npc"), unit(.5, "npc"), unit(.1, "npc")))
dev.off()

## PLOT SPEAKER RELIABILITY RESULTS
load("/Users/titlis/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/19_speaker_reliability_withins/results/data/r.RData")
r = r[!is.na(r$normresponse),]
r$AllPriorProbability = prior_allprobs[as.character(r$Item),]$X15

some100 = droplevels(subset(r, Proportion == "100" & quantifier == "Some"))
agr = aggregate(normresponse ~ AllPriorProbability  + trial_type + Item,data=some100,FUN=mean)
agr$CILow = aggregate(normresponse ~ AllPriorProbability + trial_type + Item,data=some100, FUN=ci.low)$normresponse
agr$CIHigh = aggregate(normresponse ~ AllPriorProbability + trial_type + Item,data=some100,FUN=ci.high)$normresponse
agr$YMin = agr$normresponse - agr$CILow
agr$YMax = agr$normresponse + agr$CIHigh

p=ggplot(agr, aes(x=AllPriorProbability,y=normresponse, color=trial_type)) +
  geom_point() +
  geom_errorbar(aes(ymin=YMin,ymax=YMax)) +
#  geom_smooth(method="lm") +
  scale_y_continuous(name="Posterior probability of all-state") +
  scale_x_continuous(name="Prior probability of all-state") +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3]),name="Trial type",breaks=levels(agr$trial_type),labels=c("uninformative (court)", "unreliable (drunk)", "cooperative (sober)"))
p
ggsave("pics/speakerreliabilityresults.pdf",width=7,height=4)

some100$Reliability = as.factor(ifelse(some100$trial_type == "sober","reliable","unreliable"))
centered = cbind(some100, myCenter(some100[,c("AllPriorProbability","Reliability","Trial")]))
m = lmer(normresponse ~ cAllPriorProbability * cReliability + (cAllPriorProbability * cReliability | Item) + (cAllPriorProbability * cReliability | workerid), data=centered)
summary(m)
save(m, file="data/model.RData")

m = lmer(normresponse ~ cAllPriorProbability * trial_type + (1 | Item) + (1 | workerid), data=centered)
summary(m)

m.trial = lmer(normresponse ~ cAllPriorProbability * cReliability * cTrial + (cAllPriorProbability * cReliability| Item) + (cAllPriorProbability * cReliability| workerid), data=centered)
summary(m.trial)

m.trial.nopr = lmer(normresponse ~ cAllPriorProbability + cReliability + cTrial + cReliability:cTrial + cAllPriorProbability:cTrial + cAllPriorProbability:cReliability:cTrial + (cAllPriorProbability * cReliability| Item) + (cAllPriorProbability * cReliability| workerid), data=centered)
summary(m.trial.nopr)

anova(m.trial.nopr,m.trial) #

m.trial.nopt = lmer(normresponse ~ cAllPriorProbability + cReliability + cTrial + cAllPriorProbability:cReliability + cReliability:cTrial + cAllPriorProbability:cReliability:cTrial + (cAllPriorProbability * cReliability| Item) + (cAllPriorProbability * cReliability| workerid), data=centered)
summary(m.trial.nopr)

anova(m.trial.nopt,m.trial)

# spell out all the interactions
#m.trial = lmer(normresponse ~ cAllPriorProbability + cReliability + cTrial + cAllPriorProbability:cReliability + cReliability:cTrial + cAllPriorProbability:cTrial + cAllPriorProbability:cReliability:cTrial + (cAllPriorProbability * cReliability| Item) + (cAllPriorProbability * cReliability| workerid), data=centered)

# PLOT SPEAKER RELIABILITY RESULTS BY BLOCK
agr = aggregate(normresponse ~ AllPriorProbability  + trial_type + Item + block,data=some100,FUN=mean)
agr$CILow = aggregate(normresponse ~ AllPriorProbability + trial_type + Item + block,data=some100, FUN=ci.low)$normresponse
agr$CIHigh = aggregate(normresponse ~ AllPriorProbability + trial_type + Item + block,data=some100,FUN=ci.high)$normresponse
agr$YMin = agr$normresponse - agr$CILow
agr$YMax = agr$normresponse + agr$CIHigh

p=ggplot(agr, aes(x=AllPriorProbability,y=normresponse, color=trial_type)) +
  geom_point() +
  #geom_errorbar(aes(ymin=YMin,ymax=YMax)) +
  geom_smooth(method="lm") +
  scale_y_continuous(name="Posterior probability of all-state") +
  scale_x_continuous(name="Prior probability of all-state") +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3]),name="Trial type",breaks=levels(agr$trial_type),labels=c("uninformative (court)", "unreliable (drunk)", "cooperative (sober)")) +
  facet_wrap(~block)
p
ggsave("pics/speakerreliabilityresults-byblock.pdf",width=12,height=4)

#####################################
# plot model predictions: expectations
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/mp-uniform.RData")

library(dplyr)
library(plyr)
# plot expectations for best basic model:
toplot = droplevels(subset(mp, Quantifier == "some" & WonkyWorldPrior == .5))
nrow(toplot)
pexpectations = ddply(toplot, .(Item), summarise, PosteriorExpectation_predicted=sum(State*PosteriorProbability)/15)
head(pexpectations)
some = pexpectations#droplevels(subset(pexpectations, Quantifier == "some"))

toplot = some
nrow(toplot)
head(toplot)
toplot$Mode = priormodes[as.character(toplot$Item),]$Mode

ggplot(toplot, aes(x=PriorExpectation_smoothed, y=PosteriorExpectation_predicted)) +
  geom_point(color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth(color="#00B0F6") +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,1), name="Prior expectation") +
  scale_y_continuous(limits=c(0,1), name="Model predicted posterior expectation")

load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/mp-uniform.RData")

# plot expectations for best basic model:
toplot = droplevels(subset(mp, Quantifier == "some" & WonkyWorldPrior == .5 & State == 15))
nrow(toplot)

# adjust speaker optimality at will
toplot = droplevels(subset(toplot, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)

toplot_w = toplot
# get rRSA predictions for qud=how-many, alts=0_basic, spopt=2
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/complex_prior/smoothed_unbinned15/results/data/mp.RData")
summary(mp)
toplot = droplevels(subset(mp, QUD == "how-many" & Alternatives == "0_basic" &  State == 15))
nrow(toplot)

# adjust speaker optimality at will
toplot = droplevels(subset(toplot, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)

toplot_r = toplot
toplot_r$Model = "RSA"
toplot_w$Model = "wRSA"

# plot both rRSA and uniform wRSA expectation predictions in same plot
toplot = merge(toplot_r,toplot_w, all=T)
head(toplot)
nrow(toplot)
p_probs = ggplot(toplot, aes(x=PriorProbability, y=PosteriorProbability, color=Model)) +
  geom_point() + #color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth() + #color="#00B0F6") +
  scale_color_manual(values=c("darkred",wes_palette("Darjeeling")[1])) +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,1), name="Prior mean all-state probability") +
  scale_y_continuous(limits=c(0,1), name="Predicted posterior all-state probability")
p_probs
ggsave("pred_probs.pdf")

p_probs = ggplot(toplot_r, aes(x=PriorProbability, y=PosteriorProbability, color=Model)) +
  geom_point() + #color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth() + #color="#00B0F6") +
  scale_color_manual(values=c("darkred")) +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,1), name="Prior all-state probability") +
  scale_y_continuous(limits=c(0,1), name="Posterior all-state probability")
p_probs
ggsave("pred_prob_rsa.pdf")

# plot model predictions: expectations
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/mp-uniform.RData")

# plot expectations for best basic model:
toplot = droplevels(subset(mp, Quantifier == "some" & WonkyWorldPrior == .5))
nrow(toplot)

pexpectations = ddply(toplot, .(Item, SpeakerOptimality,PriorExpectation), summarise, PosteriorExpectation_predicted=sum(State*PosteriorProbability)/15)
head(pexpectations)
some = pexpectations#droplevels(subset(pexpectations, Quantifier == "some"))

toplot = droplevels(subset(some, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)
toplot_w = toplot

# get rRSA predictions for qud=how-many, alts=0_basic, spopt=2
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/complex_prior/smoothed_unbinned15/results/data/toplot-expectations.RData")
toplot_r = toplot
head(toplot_w)
summary(toplot_r)
toplot_r$Model = "RSA"
toplot_r$PriorExpectation = toplot_r$PriorExpectation_smoothed*15
toplot_w$Model = "wRSA"

# plot both rRSA and uniform wRSA expectation predictions in same plot
toplot = merge(toplot_r,toplot_w, all=T)
head(toplot)
nrow(toplot)

p_exps = ggplot(toplot, aes(x=PriorExpectation, y=PosteriorExpectation_predicted*15, color=Model)) +
  geom_point() + #color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth() + #color="#00B0F6") +
  scale_color_manual(values=c("darkred",wes_palette("Darjeeling")[1])) +#values=c("#007fb1", "#4ecdff")) +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,15), breaks=seq(1,15,by=2), name="Prior mean number of objects") +
  scale_y_continuous(limits=c(0,15), breaks=seq(1,15,by=2), name="Predicted posterior number of objects")
p_exps
ggsave("pred_exps.pdf")

p_exps = ggplot(toplot_r, aes(x=PriorExpectation, y=PosteriorExpectation_predicted*15, color=Model)) +
  geom_point() + #color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth() + #color="#00B0F6") +
  scale_color_manual(values=c("darkred")) +#values=c("#007fb1", "#4ecdff")) +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,15), breaks=seq(1,15,by=2), name="Prior expectation") +
  scale_y_continuous(limits=c(0,15), breaks=seq(1,15,by=2), name="Posterior expectation")
p_exps
ggsave("pred_exps_rsa.pdf")

library(gridExtra)
#  share a legend between multiple plots
g <- ggplotGrob(p_exps + theme(legend.position="right"))$grobs
legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
p_exps_nolegend = p_exps + theme(legend.position="none")
p_probs_nolegend = p_probs + theme(legend.position="none")

pdf("rsa-predictions.pdf",width=12,height=4.9)
grid.arrange(p_exps_nolegend, p_probs_nolegend, legend,nrow=1,widths=unit.c(unit(.45, "npc"), unit(.45, "npc"), unit(.1, "npc")))
dev.off()

# expectations
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/13_sinking-marbles-priordv-15/results/data/r.RData")

agr = aggregate(ProportionResponse ~ PriorExpectationProportion + quantifier + Item,data=r,FUN=mean)

min(agr[agr$quantifier == "Some",]$ProportionResponse)
max(agr[agr$quantifier == "Some",]$ProportionResponse)
agr$Quantifier = factor(x=agr$quantifier, levels=c("Some","All","None","long_filler","short_filler"))

p_eexps = ggplot(agr, aes(x=PriorExpectationProportion*15, y=ProportionResponse*15, color=Quantifier, shape=Quantifier)) +
  geom_point() +
  #geom_errorbar(aes(ymin=YMin,ymax=YMax)) +
  geom_smooth(method="lm") +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +#values=c("#F8766D", "black", "#00BF7D", "gray30", "#00B0F6"),breaks=levels(agr$quantifier),labels=c("all","long filler","none","short filler","some")) +
  scale_y_continuous(breaks=seq(1,15,by=2),name="Posterior mean number of objects") +
  #  geom_abline(intercept=0,slope=1,color="gray70") +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects")
p_eexps
ggsave(file="empirical_exps.pdf")#,width=5,height=3.7)

# allstate-probs
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/16_sinking-marbles-sliders-certain/results/data/r.RData")

# exclude people who are doing some sort of bullshit and not responding reasonably to all/none (see subject-variability.pdf for behavior on zero-slider)
tmp = subset(r,!workerid %in% c(0,22,43,98,100,103,117,118))
agrr = aggregate(normresponse ~ AllPriorProbability + Proportion + quantifier + Item,data=tmp,FUN=mean)
agrr = aggregate(normresponse ~ AllPriorProbability + Proportion + quantifier + Item,data=r,FUN=mean)
ub = subset(agrr, Proportion == "100")
ub = droplevels(ub)
ub$Quantifier = factor(x=ub$quantifier, levels=c("Some","All","None","long_filler","short_filler"))

p_eprobs = ggplot(ub, aes(x=AllPriorProbability, y = normresponse, color=Quantifier, shape=Quantifier)) +
  geom_point() +
  geom_smooth(method="lm") +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +#values=c("#F8766D", "black", "#00BF7D", "gray30", "#00B0F6"),breaks=levels(agr$quantifier),labels=c("all","long filler","none","short filler","some")) +
  scale_y_continuous(limits=c(0,1),name="Posterior probability of all-state ") +
  #  geom_abline(intercept=0,slope=1,color="gray70") +
  scale_x_continuous(limits=c(0,1),name="Prior probability of all-state")
p_eprobs
ggsave(file="empirical-allprobs.pdf")#,width=5,height=3.7)

library(gridExtra)
#  share a legend between multiple plots
g <- ggplotGrob(p_eexps + theme(legend.position="right"))$grobs
legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
p_eexps_nolegend = p_eexps + theme(legend.position="none")
p_eprobs_nolegend = p_eprobs + theme(legend.position="none")

pdf("empirical-results.pdf",width=12,height=4.9)
grid.arrange(p_eexps_nolegend, p_eprobs_nolegend, legend,nrow=1,widths=unit.c(unit(.45, "npc"), unit(.45, "npc"), unit(.1, "npc")))
dev.off()

load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/wr-uniform.RData")

# plot expectations for best basic model:
toplot = droplevels(subset(wr, WonkyWorldPrior == .5 & Wonky == "true"))
nrow(toplot)
toplot$Mode = priormodes[as.character(toplot$Item),]$Mode

toplot = droplevels(subset(toplot, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)
#toplot$quantifier = capitalize(as.character(toplot$Quantifier))
toplot$Quantifier = factor(toplot$Quantifier, levels=c("some","all","none"))

p_wmodel = ggplot(toplot, aes(x=PriorExpectation, y=PosteriorProbability,color=Quantifier,shape=Quantifier)) +
  geom_point() +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +
  #scale_color_manual(values=c("#F8766D","#00BF7D","#00B0F6")) +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects") +
  scale_y_continuous(breaks=seq(0,1,by=.25),name="Predicted posterior wonkiness probability") +
  theme(plot.margin=unit(c(0,0,0,0),units="cm"))
p_wmodel
ggsave("model-wonkiness-uniform.pdf",width=6.8,height=5)#,width=30,height=10)

############
# empirical wonkiness posteriors
############
load("/Users/titlis/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/17_sinking-marbles-normal-sliders/results/data/r.RData")
head(r)
nrow(r)

toplot = aggregate(response ~ quantifier + Item + PriorExpectation, FUN="mean", data=r)
toplot = droplevels(subset(toplot, quantifier %in% c("All","Some","None")))
toplot$Quantifier = factor(tolower(toplot$quantifier), levels=c("all", "none","some"))
toplot$Mode = priormodes[as.character(toplot$Item),]$Mode
toplot$quantifier = capitalize(as.character(toplot$Quantifier))
toplot$Quantifier = factor(toplot$quantifier, levels=c("Some","All","None"))

p_wempirical = ggplot(toplot, aes(x=PriorExpectation, y=response, color=Quantifier,shape=Quantifier)) +
  geom_point() +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects") +
  scale_y_continuous(breaks=seq(0,1,by=.25),name="Mean empirical wonkiness probability")  +
  theme(plot.margin=unit(c(0,0,0,0),units="cm"))
p_wempirical
ggsave(file="empirical-wonkiness.pdf",width=6.8,height=5)

library(gridExtra)
#  share a legend between multiple plots
g <- ggplotGrob(p_wmodel + theme(legend.position="right"))$grobs
legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
p_wmodel_nolegend = p_wmodel + theme(legend.position="none")
p_wempirical_nolegend = p_wempirical + theme(legend.position="none")

#pdf("rsa-predictions-uniform.pdf",width=10,height=4)
pdf("wonkiness-fullplot.pdf",width=12,height=5)
grid.arrange(p_wmodel_nolegend, p_wempirical_nolegend, legend,nrow=1,widths=unit.c(unit(.45, "npc"), unit(.45, "npc"), unit(.1, "npc")))
dev.off()

## load four-step priors for generating appendix plots
priorprobs = read.table(file="~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/24_sinking-marbles-prior-fourstep/results/data/smoothed_15marbles_priors_withnames.txt",sep="\t", header=T, quote="")
row.names(priorprobs) = priorprobs$Item
nrow(priorprobs)

# prior all-state probability for each item
prior_allprobs = priorprobs[,c("Item","X15")]
row.names(prior_allprobs) = prior_allprobs$Item

# prior expectation for each item
gathered_probs <- priorprobs %>%
  gather(State,Probability,X0:X15)
gathered_probs$State = as.numeric(as.character(gsub("X","",gathered_probs$State)))
prior_exps <- gathered_probs %>%
  group_by(Item) %>%
  summarise(exp.val=sum(State*Probability))
prior_exps = as.data.frame(prior_exps)
row.names(prior_exps) = prior_exps$Item
summary(prior_exps)

# histogram of expectations
exps = ggplot(prior_exps, aes(x=exp.val)) +
  geom_histogram() +
  scale_x_continuous(name="Expected value of prior distribution",breaks=seq(1,15, by=2)) +
  scale_y_continuous(name="Number of cases",breaks=seq(0,8, by=2))
ggsave("pics/priorexpectations-histogram-fourstep.pdf",width=5,height=3.7)

# histogram of allstate-probs
allprobs = ggplot(prior_allprobs, aes(x=X15)) +
  geom_histogram() +
  scale_x_continuous(name="Prior all-state probability") +
  scale_y_continuous(name="Number of cases")
ggsave("priorallprobs-histogram-fourstep.pdf")

pdf("pics/priordistributions-fourstep.pdf",width=10,height=3.5)
grid.arrange(exps, allprobs, nrow=1,widths=unit.c(unit(.5, "npc"), unit(.5, "npc")))
dev.off()


compare <- function(a, b)
{
  cat(paste(a, "is of type", class(a), "and", b, "is of type", class(b), "\n"))

  if (a < b) cat(paste(a, "is strictly less than", b, "\n"))
  if (a <= b) cat(paste(a, "is less than or equal to", b, "\n"))
  if (a > b) cat(paste(a, "is strictly greater than", b, "\n"))
  if (a >= b) cat(paste(a, "is greater than or equal to", b, "\n"))
  if (a == b) cat(paste(a, "is equal to", b, "\n"))
  if (a != b) cat(paste(a, "is not equal to", b, "\n"))

  invisible()
}

compare('YUP', 'YUP')
compare('BALL', 'BELL')
compare('24', '123')
compare(24, 123)
compare(5.0, 5)


# Dual Luciferase Analysis Pipeline
# Elias Brandorff, SILS Amsterdam Version 1
# import excel data from GloMax luminescence machine
# level 1 calculations: normalize expression to internal control
# plot normalized expression values
# produce a summary table
# level 2 calculations: normalize expression of experimental to control condition
# plot relative expression

#set working path
setwd("C:/DualReporterPipeline/")

# load the relevant libraries
library(readxl)
library(tidyr)
library(ggplot2)
library(dplyr)
library(gridExtra)
library(magrittr)

# import complete excel output GloMax with readxl into a dataframe.
# either point to number or name of the correct sheet
all_data_from_excel <- read_excel("Dual Luciferase Reporter Assay System ELIAS 2020.03.10 05_34_44.xlsx", sheet = "Results")

#Subset firefly data
firefly <- all_data_from_excel[19:26,6:17] %>% unlist()

#Subset renilla data
renilla <- all_data_from_excel[40:47,6:17] %>% unlist()

#Read conditions from a csv file - empty cells are interpreted as NA
df_conditions <- read.csv("conditions_exp1.csv", stringsAsFactors = FALSE, na.strings = c("","NA","na"))

#Subset the dataframe, to select only conditions (get rid of row names)
df_conditions <- df_conditions[1:8,2:13]

#Convert the dataframe with conditions to a vector
condition <- df_conditions %>% unlist(use.names = FALSE)

#Combine all vectors in a (tidy) dataframe, remove data that has no condition associated with it
#Note: By using filtering for NA, the 'condition' vector defines the data you want to analyze/display
df <- data.frame(condition,firefly,renilla) %>% filter(!is.na(condition))

#add a column that in which data is normalized to internal control
df <- df %>% mutate(FR=firefly/renilla) %>% na.omit() #Remove NA from the table (i.e. cells without data)

#Ensure that the order of the conditions is kept (otherwise ordering is alphabetical)
df$condition <- factor(df$condition, levels=unique(df$condition))

#Use this line of code when you want to sort the data acoording to median of FR
# df$condition <- reorder(df$condition, df$FR, median, na.rm = TRUE)

FR_tidy <- df

#plotting FR ratios for data overview
p1 <- ggplot(FR_tidy, aes(x = condition, y = FR)) +
  geom_jitter(position=position_jitter(0.1), cex=2, color="grey40") +
  stat_summary(fun = median, fun.min = median, fun.max = median,
               geom = "crossbar", width = 0.2)

p1

#Save p1 in png format
png("FR_summary_exp1.png", width = 300, height = 300)
plot(p1)
dev.off()

#Attemp to spread the dataframe.
FR_tidy5 <- df %>%
  select(condition, FR) %>%
  group_by(condition) %>%
  mutate(grouped_id = row_number()) %>%
  pivot_wider(id_cols = NULL, names_from = condition, values_from = FR) %>%
  summarize_all(funs(median)) %>%
  select(-grouped_id)
FR_tidy5

#The table is in the correct orientation now. But we want to ad the mean and median as the 2 bottom rows to a the FR-values

#produce a table in PDF for publication containing the mean FR ratio per condition
pdf("FRtable_exp1.pdf", height=11, width=10, pivot = TRUE)
grid.table(FR_tidy5)
dev.off()

#FR summary
FR2_tidy2 <- FR_tidy %>%
  group_by(condition) %>%
  summarise(median = median(FR, na.rm = TRUE), mean = mean(FR, na.rm = TRUE))
FR2_tidy2

#doing the relative calculation i.e. normalize to expression of the empty vector "Ev"
#Calculate the average of the EV control
EV_mean <- FR_tidy %>% filter(condition=="control") %>% summarise(mean_EV_FR=mean(FR))  %>% unlist(use.names = FALSE)

#Divide FR by EV control value
FR_tidy <- FR_tidy %>% mutate(FC = FR/EV_mean)

#### Here a next step in the normalization can be made. I tested some data generated by my student. In this case we have an empty-vector control
#### and two reporter constructs that should be compared to each other, after normalization to the empty reporter vector.
#### eventually, the "control" conditions can be omitted from the second plot.

p2 <- ggplot(FR_tidy, aes(condition, FC)) +
  geom_point() +
  stat_summary(fun = median, fun.min = median, fun.max = median,
               geom = "crossbar", width = 0.2) +
  ylab("FC luciferase expression / EV normalized to renilla") +
  ylim(0,1)

p2

#Save p2 in png format
png("FC_lucexpression1.png", width = 300, height = 300)
plot(p2)
dev.off()

tbl_median <- tableGrob(t(FR2_tidy2), theme = ttheme_default(8))
grid.arrange(p1, p2, tbl_median, ncol=2, nrow=2, as.table=TRUE, heights=c(3,1))


#' Traceplot of the occupancy index rate
#'
#' @description Traceplot of the index occupancy rate, with
#' effective sample size.
#'
#' @param modelResults Output of the function \code{runModel}
#' @param index_year Index of the year to plot. Indexes go from 1 to the number of years.
#'
#' @importFrom magrittr %>%
#'
#' @export
#'
#' @return The traceplot with the estimate of the effective sample size.
#'
#' @examples
#'
#' tracePlot_OccupancyIndexRate(sampleResults, 1)
#'
tracePlot_OccupancyIndexRate <- function(modelResults, index_year) {

  psi_mean_output <- modelResults$modelOutput$psi_mean_output

  nchain <- nrow(psi_mean_output)
  niter <- ncol(psi_mean_output)

  psi_mean_output <-
    matrix(psi_mean_output[, , index_year], nrow = nchain, ncol = niter)

  psi_mean_output_long <- reshape2::melt(psi_mean_output)

  diagnosticsPlot <-
    ggplot2::ggplot(data = psi_mean_output_long, ggplot2::aes(
      x = Var2,
      y = value,
      group = Var1,
      color = factor(Var1)
    )) + ggplot2::geom_line() +
    ggplot2::xlab("Iterations") + ggplot2::ylab("Value") +
    ggplot2::theme(
      plot.title = ggplot2::element_text(hjust = 0.5, size = 17),
      axis.title = ggplot2::element_text(size = 16, face = "bold"),
      axis.text.y = ggplot2::element_text(size = 11, face = "bold"),
      axis.text.x = ggplot2::element_text(
        size = 11,
        face = "bold",
        hjust = 1
      ),
      axis.line = ggplot2::element_line(colour = "black", size = 0.15),
      # panel.grid.minor = element_line(colour="grey", size=0.15),
      panel.grid.major = ggplot2::element_line(colour = "grey", size = 0.15),
      panel.background = ggplot2::element_rect(fill = "white", color = "black"),
      legend.position = "none"
    )

  plotTitle <- createPlotTitle(psi_mean_output, nchain)

  diagnosticsPlot <- diagnosticsPlot + ggplot2::ggtitle(plotTitle)

  diagnosticsPlot
}

#' Traceplot of the year-specific random effects
#'
#' @description Traceplot of the year-specific random effects, with
#' effective sample size.
#'
#' @param modelResults Output of the function \code{runModel}
#' @param index_year Index of the year of the random effect to plot. Indexes go from 1 to the number of years.
#'
#' @importFrom magrittr %>%
#'
#' @export
#'
#' @return The traceplot with the estimate of the effective sample size.
#'
#' @examples
#'
#' tracePlot_OccupancyYearEffect(sampleResults, 1)
#'
tracePlot_OccupancyYearEffect <- function(modelResults, index_year) {
  beta_psi_output <- modelResults$modelOutput$beta_psi_output

  nchain <- nrow(beta_psi_output)
  niter <- ncol(beta_psi_output)

  beta_psi_output <-
    matrix(beta_psi_output[, , index_year], nrow = nchain, ncol = niter)

  beta_psi_output_long <- reshape2::melt(beta_psi_output)

  diagnosticsPlot <-
    ggplot2::ggplot(data = beta_psi_output_long, ggplot2::aes(
      x = Var2,
      y = value,
      group = Var1,
      color = factor(Var1)
    )) + ggplot2::geom_line() +
    ggplot2::xlab("Iterations") + ggplot2::ylab("Value") +
    ggplot2::theme(
      plot.title = ggplot2::element_text(hjust = 0.5, size = 17),
      axis.title = ggplot2::element_text(size = 16, face = "bold"),
      axis.text.y = ggplot2::element_text(size = 11, face = "bold"),
      axis.text.x = ggplot2::element_text(
        size = 11,
        face = "bold",
        hjust = 1
      ),
      axis.line = ggplot2::element_line(colour = "black", size = 0.15),
      # panel.grid.minor = element_line(colour="grey", size=0.15),
      panel.grid.major = ggplot2::element_line(colour = "grey", size = 0.15),
      panel.background = ggplot2::element_rect(fill = "white", color = "black"),
      legend.position = "none"
    )

  plotTitle <- createPlotTitle(beta_psi_output, nchain)

  diagnosticsPlot <- diagnosticsPlot + ggplot2::ggtitle(plotTitle)

  diagnosticsPlot
}

#' Traceplot of the site-specific autocorrelated random effects
#'
#' @description Traceplot of the site-specific random effects, with
#' effective sample size.
#'
#' @param modelResults Output of the function \code{runModel}
#' @param index_site Index of the random effect to plot. Indexes go from 1 to the number of sites in the approximation.
#'
#' @importFrom magrittr %>%
#'
#' @export
#'
#' @return The traceplot with the estimate of the effective sample size.
#'
#' @examples
#'
#' tracePlot_OccupancySiteEffect(sampleResults, 1)
#'
tracePlot_OccupancySiteEffect <- function(modelResults, index_site) {
  usingSpatial <- modelResults$dataCharacteristics$usingSpatial

  if (!usingSpatial) {
    print("No spatial effect")
  } else {
    years <- modelResults$dataCharacteristics$Years
    Y <- length(years)

    beta_psi_output <- modelResults$modelOutput$beta_psi_output

    nchain <- nrow(beta_psi_output)
    niter <- ncol(beta_psi_output)

    beta_psi_output <-
      matrix(beta_psi_output[, , Y + index_site], nrow = nchain, ncol = niter)

    beta_psi_output_long <- reshape2::melt(beta_psi_output)

    diagnosticsPlot <-
      ggplot2::ggplot(data = beta_psi_output_long, ggplot2::aes(
        x = Var2,
        y = value,
        group = Var1,
        color = factor(Var1)
      )) + ggplot2::geom_line() +
      ggplot2::xlab("Iterations") + ggplot2::ylab("Value") +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, size = 17),
        axis.title = ggplot2::element_text(size = 16, face = "bold"),
        axis.text.y = ggplot2::element_text(size = 11, face = "bold"),
        axis.text.x = ggplot2::element_text(
          size = 11,
          face = "bold",
          hjust = 1
        ),
        axis.line = ggplot2::element_line(colour = "black", size = 0.15),
        # panel.grid.minor = element_line(colour="grey", size=0.15),
        panel.grid.major = ggplot2::element_line(colour = "grey", size = 0.15),
        panel.background = ggplot2::element_rect(fill = "white", color = "black"),
        legend.position = "none"
      )

    plotTitle <- createPlotTitle(beta_psi_output, nchain)

    diagnosticsPlot <- diagnosticsPlot + ggplot2::ggtitle(plotTitle)

    return(diagnosticsPlot)
  }

}

#' Traceplot of the covariates of the occupancy probability
#'
#' @description Traceplot of the covariates coefficient, with
#' effective sample size.
#'
#' @param modelResults Output of the function \code{runModel}
#' @param index_cov Index of the covariate to plot. Indexes go from 1 to the number of covariate.
#'
#' @importFrom magrittr %>%
#'
#' @export
#'
#' @return The traceplot with the estimate of the effective sample size.
#'
#' @examples
#'
#' tracePlot_OccupancyCovariate(sampleResults, 1)
#'
tracePlot_OccupancyCovariate <- function(modelResults, index_cov) {
  beta_psi_output <- modelResults$modelOutput$beta_psi_output

  usingSpatial <- modelResults$dataCharacteristics$usingSpatial

  Y <- length(modelResults$dataCharacteristics$Years)

  if (usingSpatial) {
    X_centers <- nrow(modelResults$dataCharacteristics$X_tilde)
  } else {
    X_centers <- 0
  }

  numTimeSpaceCov <-
    modelResults$dataCharacteristics$numTimeSpaceCov

  if ((Y + X_centers + numTimeSpaceCov + index_cov) <= dim(beta_psi_output)[3]) {
    nchain <- nrow(beta_psi_output)
    niter <- ncol(beta_psi_output)

    beta_psi_output <-
      matrix(beta_psi_output[, , Y + X_centers + numTimeSpaceCov + index_cov], nrow = nchain, ncol = niter)

    beta_psi_output_long <- reshape2::melt(beta_psi_output)

    diagnosticsPlot <-
      ggplot2::ggplot(data = beta_psi_output_long, ggplot2::aes(
        x = Var2,
        y = value,
        group = Var1,
        color = factor(Var1)
      )) + ggplot2::geom_line() +
      ggplot2::xlab("Iterations") + ggplot2::ylab("Value") +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, size = 17),
        axis.title = ggplot2::element_text(size = 16, face = "bold"),
        axis.text.y = ggplot2::element_text(size = 11, face = "bold"),
        axis.text.x = ggplot2::element_text(
          size = 11,
          face = "bold",
          hjust = 1
        ),
        axis.line = ggplot2::element_line(colour = "black", size = 0.15),
        # panel.grid.minor = element_line(colour="grey", size=0.15),
        panel.grid.major = ggplot2::element_line(colour = "grey", size = 0.15),
        panel.background = ggplot2::element_rect(fill = "white", color = "black"),
        legend.position = "none"
      )

    plotTitle <- createPlotTitle(beta_psi_output, nchain)

    diagnosticsPlot <- diagnosticsPlot + ggplot2::ggtitle(plotTitle)

    return(diagnosticsPlot)

  } else {
    print("Index above the number of covariates")

  }

}

#' Traceplot of the intercept of the detection probability
#'
#' @description Traceplot of the intercept, with
#' effective sample size.
#'
#' @param modelResults Output of the function \code{runModel}
#' @param index_year Index of the year to plot, in case \code{usingYearDetProb} is
#' set to \code{TRUE}. Indexes go from 1 to the number of years.
#'
#' @importFrom magrittr %>%
#'
#' @export
#'
#' @return The traceplot with the estimate of the effective sample size.
#'
#' @examples
#'
#' tracePlot_DetectionIntercept(sampleResults, 1)
#'
tracePlot_DetectionIntercept <- function(modelResults, index_year) {

  usingYearDetProb <- modelResults$dataCharacteristics$usingYearDetProb

  years <- modelResults$dataCharacteristics$Years
  Y <- length(years)

  if(!usingYearDetProb){
    index_year <- 1
  } else if(index_year > Y){

    print("Index above the number of years")
    return(NULL)

  }

  beta_p_output <- modelResults$modelOutput$beta_p_output

  nchain <- nrow(beta_p_output)
  niter <- ncol(beta_p_output)

  beta_p_output <-
    matrix(beta_p_output[, , index_year], nrow = nchain, ncol = niter)

  beta_psi_output_long <- reshape2::melt(beta_p_output)

  diagnosticsPlot <-
    ggplot2::ggplot(data = beta_psi_output_long, ggplot2::aes(
      x = Var2,
      y = value,
      group = Var1,
      color = factor(Var1)
    )) + ggplot2::geom_line() +
    ggplot2::xlab("Iterations") + ggplot2::ylab("Value") +
    ggplot2::theme(
      plot.title = ggplot2::element_text(hjust = 0.5, size = 17),
      axis.title = ggplot2::element_text(size = 16, face = "bold"),
      axis.text.y = ggplot2::element_text(size = 11, face = "bold"),
      axis.text.x = ggplot2::element_text(
        size = 11,
        face = "bold",
        hjust = 1
      ),
      axis.line = ggplot2::element_line(colour = "black", size = 0.15),
      # panel.grid.minor = element_line(colour="grey", size=0.15),
      panel.grid.major = ggplot2::element_line(colour = "grey", size = 0.15),
      panel.background = ggplot2::element_rect(fill = "white", color = "black"),
      legend.position = "none"
    )

  plotTitle <- createPlotTitle(beta_p_output, nchain)

  diagnosticsPlot <- diagnosticsPlot + ggplot2::ggtitle(plotTitle)

  diagnosticsPlot
}

#' Traceplot of the covariates of the detection probability
#'
#' @description Traceplot of the covariates coefficient, with
#' effective sample size.
#'
#' @param modelResults Output of the function \code{runModel}
#' @param index_cov Index of the covariate to plot. Indexes go from 1 to the number of covariate.
#'
#' @importFrom magrittr %>%
#'
#' @export
#'
#' @return The traceplot with the estimate of the effective sample size.
#'
#' @examples
#'
#' tracePlot_DetectionCovariate(sampleResults, 1)
#'
tracePlot_DetectionCovariate <- function(modelResults, index_cov) {
  beta_p_output <- modelResults$modelOutput$beta_p_output

  usingYearDetProb <- modelResults$dataCharacteristics$usingYearDetProb
  p_intercepts <- ifelse(usingYearDetProb, Y, 1)

  if ((p_intercepts + index_cov) <= dim(beta_p_output)[3]) {
    nchain <- nrow(beta_p_output)
    niter <- ncol(beta_p_output)

    beta_p_output <-
      matrix(beta_p_output[, , p_intercepts + index_cov], nrow = nchain, ncol = niter)

    beta_psi_output_long <- reshape2::melt(beta_p_output)

    diagnosticsPlot <-
      ggplot2::ggplot(data = beta_psi_output_long, ggplot2::aes(
        x = Var2,
        y = value,
        group = Var1,
        color = factor(Var1)
      )) + ggplot2::geom_line() +
      ggplot2::xlab("Iterations") + ggplot2::ylab("Value") +
      ggplot2::theme(
        plot.title = ggplot2::element_text(hjust = 0.5, size = 17),
        axis.title = ggplot2::element_text(size = 16, face = "bold"),
        axis.text.y = ggplot2::element_text(size = 11, face = "bold"),
        axis.text.x = ggplot2::element_text(
          size = 11,
          face = "bold",
          hjust = 1
        ),
        axis.line = ggplot2::element_line(colour = "black", size = 0.15),
        # panel.grid.minor = element_line(colour="grey", size=0.15),
        panel.grid.major = ggplot2::element_line(colour = "grey", size = 0.15),
        panel.background = ggplot2::element_rect(fill = "white", color = "black"),
        legend.position = "none"
      )

    plotTitle <- createPlotTitle(beta_p_output, nchain)

    diagnosticsPlot <- diagnosticsPlot + ggplot2::ggtitle(plotTitle)

    return(diagnosticsPlot)

  } else {
    print("Index above the number of covariates")

  }

}

# tracePlot_l_T <- function(modelResults) {
#   l_T_output <- modelResults$modelOutput$l_T_output
#
#   nchain <- nrow(l_T_output)
#
#   l_T_psi_output_long <- reshape2::melt(l_T_output)
#
#   diagnosticsPlot <-
#     ggplot(data = l_T_psi_output_long, aes(
#       x = Var2,
#       y = value,
#       group = Var1,
#       color = factor(Var1)
#     )) + geom_line() +
#     xlab("Iterations") + ylab("Value") +
#     theme(
#       plot.title = element_text(hjust = 0.5, size = 17),
#       axis.title = element_text(size = 16, face = "bold"),
#       axis.text.y = element_text(size = 11, face = "bold"),
#       axis.text.x = element_text(
#         size = 11,
#         face = "bold",
#         hjust = 1
#       ),
#       axis.line = element_line(colour = "black", size = 0.15),
#       # panel.grid.minor = element_line(colour="grey", size=0.15),
#       panel.grid.major = element_line(colour = "grey", size = 0.15),
#       panel.background = element_rect(fill = "white", color = "black"),
#       legend.position = "none"
#     )
#
#   plotTitle <- createPlotTitle(l_T_output, nchain)
#
#   diagnosticsPlot <- diagnosticsPlot + ggtitle(plotTitle)
#
#   diagnosticsPlot
#
# }
#
# tracePlot_sigma_T <- function(modelResults) {
#   sigma_T_output <- modelResults$modelOutput$sigma_T_output
#
#   nchain <- nrow(sigma_T_output)
#
#   sigma_T_output_long <- reshape2::melt(sigma_T_output)
#
#   diagnosticsPlot <-
#     ggplot(data = sigma_T_output_long, aes(
#       x = Var2,
#       y = value,
#       group = Var1,
#       color = factor(Var1)
#     )) + geom_line() +
#     xlab("Iterations") + ylab("Value") +
#     theme(
#       plot.title = element_text(hjust = 0.5, size = 17),
#       axis.title = element_text(size = 16, face = "bold"),
#       axis.text.y = element_text(size = 11, face = "bold"),
#       axis.text.x = element_text(
#         size = 11,
#         face = "bold",
#         hjust = 1
#       ),
#       axis.line = element_line(colour = "black", size = 0.15),
#       # panel.grid.minor = element_line(colour="grey", size=0.15),
#       panel.grid.major = element_line(colour = "grey", size = 0.15),
#       panel.background = element_rect(fill = "white", color = "black"),
#       legend.position = "none"
#     )
#
#   plotTitle <- createPlotTitle(sigma_T_output, nchain)
#
#   diagnosticsPlot <- diagnosticsPlot + ggtitle(plotTitle)
#
#   diagnosticsPlot
#
# }
#
# tracePlot_l_S <- function(modelResults) {
#   usingSpatial <- modelResults$dataCharacteristics$usingSpatial
#
#   if (usingSpatial) {
#     l_s_output <- modelResults$modelOutput$l_s_output
#
#     nchain <- nrow(l_s_output)
#
#     l_s_psi_output_long <- reshape2::melt(l_s_output)
#
#     diagnosticsPlot <-
#       ggplot(data = l_s_psi_output_long, aes(
#         x = Var2,
#         y = value,
#         group = Var1,
#         color = factor(Var1)
#       )) + geom_line() +
#       xlab("Iterations") + ylab("Value") +
#       theme(
#         plot.title = element_text(hjust = 0.5, size = 17),
#         axis.title = element_text(size = 16, face = "bold"),
#         axis.text.y = element_text(size = 11, face = "bold"),
#         axis.text.x = element_text(
#           size = 11,
#           face = "bold",
#           hjust = 1
#         ),
#         axis.line = element_line(colour = "black", size = 0.15),
#         # panel.grid.minor = element_line(colour="grey", size=0.15),
#         panel.grid.major = element_line(colour = "grey", size = 0.15),
#         panel.background = element_rect(fill = "white", color = "black"),
#         legend.position = "none"
#       )
#
#     plotTitle <- createPlotTitle(l_s_output, nchain)
#
#     diagnosticsPlot <- diagnosticsPlot + ggtitle(plotTitle)
#
#     return(diagnosticsPlot)
#   } else {
#     print("No spatial effect")
#   }
#
#
# }
#
# tracePlot_sigma_S <- function(modelResults) {
#   sigma_s_output <- modelResults$modelOutput$sigma_s_output
#
#   nchain <- nrow(sigma_s_output)
#
#   sigma_s_output_long <- reshape2::melt(sigma_s_output)
#
#   diagnosticsPlot <-
#     ggplot2::ggplot(data = sigma_s_output_long, aes(
#       x = Var2,
#       y = value,
#       group = Var1,
#       color = factor(Var1)
#     )) + ggplot2::geom_line() +
#     ggplot2::xlab("Iterations") + ggplot2::ylab("Value") +
#     ggplot2::theme(
#       plot.title = ggplot2::element_text(hjust = 0.5, size = 17),
#       axis.title = ggplot2::element_text(size = 16, face = "bold"),
#       axis.text.y = ggplot2::element_text(size = 11, face = "bold"),
#       axis.text.x = ggplot2::element_text(
#         size = 11,
#         face = "bold",
#         hjust = 1
#       ),
#       axis.line = ggplot2::element_line(colour = "black", size = 0.15),
#       # panel.grid.minor = element_line(colour="grey", size=0.15),
#       panel.grid.major = ggplot2::element_line(colour = "grey", size = 0.15),
#       panel.background = ggplot2::element_rect(fill = "white", color = "black"),
#       legend.position = "none"
#     )
#
#   plotTitle <- createPlotTitle(sigma_s_output, nchain)
#
#   diagnosticsPlot <- diagnosticsPlot + ggplot2::ggtitle(plotTitle)
#
#   diagnosticsPlot
#
# }
#
# tracePlot_sigma_eps <- function(modelResults) {
#   sigma_as_output <- modelResults$modelOutput$sigma_as_output
#
#   nchain <- nrow(sigma_as_output)
#
#   sigma_as_output_long <- reshape2::melt(sigma_as_output)
#
#   ggplot2::diagnosticsPlot <-
#     ggplot2::ggplot(data = sigma_as_output_long, aes(
#       x = Var2,
#       y = value,
#       group = Var1,
#       color = factor(Var1)
#     )) + ggplot2::geom_line() +
#     ggplot2::xlab("Iterations") + ggplot2::ylab("Value") +
#     ggplot2::theme(
#       plot.title = ggplot2::element_text(hjust = 0.5, size = 17),
#       axis.title = ggplot2::element_text(size = 16, face = "bold"),
#       axis.text.y = ggplot2::element_text(size = 11, face = "bold"),
#       axis.text.x = ggplot2::element_text(
#         size = 11,
#         face = "bold",
#         hjust = 1
#       ),
#       axis.line = ggplot2::element_line(colour = "black", size = 0.15),
#       # panel.grid.minor = element_line(colour="grey", size=0.15),
#       panel.grid.major = ggplot2::element_line(colour = "grey", size = 0.15),
#       panel.background = ggplot2::element_rect(fill = "white", color = "black"),
#       legend.position = "none"
#     )
#
#   plotTitle <- createPlotTitle(sigma_as_output, nchain)
#
#   diagnosticsPlot <- diagnosticsPlot + ggplot2::ggtitle(plotTitle)
#
#   diagnosticsPlot
#
# }

createPlotTitle <- function(mcmc_output, nchain) {
  eff_samplesize <- ess(mcmc_output)

  plotTitle <-
    paste0("Effective sample size = ", round(eff_samplesize, 1))

  if (nchain > 1) {
    Rhat <- compute_rhat(mcmc_output)

    plotTitle <-
      paste0(plotTitle, paste0(" / R.hat = ", round(Rhat, 3)))
  }

  plotTitle
}

compute_rhat <- function(mcmc_output) {
  mcmc_output_list <- lapply(1:nrow(mcmc_output), function(i) {
    coda::mcmc(mcmc_output[i, ])
  })
  mcmc_output_list_2 <- coda::as.mcmc.list(mcmc_output_list)

  Rhat <- coda::gelman.diag(mcmc_output_list_2)

  Rhat$psrf[2]
}

ess <- function(mcmc_output) {
  mcmc_output_list <- lapply(1:nrow(mcmc_output), function(i) {
    coda::mcmc(mcmc_output[i, ])
  })
  mcmc_output_list_2 <- coda::as.mcmc.list(mcmc_output_list)

  coda::effectiveSize(mcmc_output_list_2)

}


dyn.load('/Library/Java/JavaVirtualMachines/jdk1.8.0_131.jdk/Contents/Home/jre/lib/server/libjvm.dylib')
library(rJava)

setwd("/Users/mengmengjiang/all datas/print")

library(xlsx)

# reading ux and sd

k1<-read.xlsx("dotx.xlsx",sheetName="600",header=TRUE)
k2<-read.xlsx("dotx.xlsx",sheetName="1khz",header=TRUE)
k3<-read.xlsx("dotx.xlsx",sheetName="2khz",header=TRUE)
k4<-read.xlsx("dotx.xlsx",sheetName="2.5khz",header=TRUE)

# errorbar
error.bar <- function(x, y, upper, coll,lower=upper, length=0.05,...){
if(length(x) != length(y) | length(y) !=length(lower) | length(lower) != length(upper))
stop("vectors must be same length")
arrows(x,y+upper, x, y-lower,col=coll, angle=90, code=3, length=length, ...)
}

# color setting

yan<-c("red","blue","black","green3")
pcc<-c(0,1,2,5)

# plot
par(fig=c(0,1,0,1),new=T)

plot(k1$ux,k1$sdeva, col=0,xlab = expression(italic(U["x"]) (mm/s)),
          ylab = expression(italic(S["d"]) (um)), mgp=c(1.1, 0, 0),tck=0.02,
               main = "", xlim = c(-50,300),ylim=c(0,400))

mtext("Sd+d and fitting",3,line=0.2,font=2,cex=1.2)

points(k1$ux,(k1$sdeva+k1$deva),col=yan[1],pch=pcc[1])
points(k2$ux,(k2$sdeva+k2$deva),col=yan[2],pch=pcc[2])
points(k3$ux,(k3$sdeva+k3$deva),col=yan[3],pch=pcc[3])
points(k4$ux,(k4$sdeva+k4$deva),col=yan[4],pch=pcc[4])

z1<-lm((k1$sdeva+k1$deva)~k1$ux)
z2<-lm((k2$sdeva+k2$deva)~k2$ux)
z3<-lm((k3$sdeva+k3$deva)~k3$ux)
z4<-lm((k4$deva+k4$sdeva)~k4$ux)

abline(z1,col=yan[1],lwd=1.5,lty=2)
abline(z2,col=yan[2],lwd=1.5,lty=2)
abline(z3,col=yan[3],lwd=1.5,lty=2)
abline(z4,col=yan[4],lwd=1.5,lty=2)

error.bar(k1$ux,(k1$deva+k1$sdeva),(k1$stdd+k1$sdstd)/2,col=yan[1])
error.bar(k2$ux,(k2$deva+k2$sdeva),(k2$stdd+k2$sdstd)/2,col=yan[2])
error.bar(k3$ux,(k3$deva+k3$sdeva),(k3$stdd+k3$sdstd)/2,col=yan[3])
error.bar(k4$ux,(k4$deva+k4$sdeva),(k4$stdd+k4$sdstd)/2,col=yan[4])

leg<-c("600Hz","1KHz","2KHz","2.5KHz")

legend("topright",legend=leg,col=yan,pch=pcc,lwd=1.5,lty=2,inset=.02,
bty="n")

# 斜率获取
b11<-z1$coefficients[2] #b11 = 2.113 ,b1=2.13
b2<-z2$coefficients[2]
b3<-z3$coefficients[2]
b4<-z4$coefficients[2]

x<-b1/b11

a1<-z1$coefficients[1]
a2<-z2$coefficients[1]
a3<-z3$coefficients[1]
a4<-z4$coefficients[1]

b<-c(b1,b2,b3,b4)

# 近似的斜率获取
b1<-2.13

kv2<-b1*0.6

kv3<-b1*0.3

kv4<-b1*0.24

kv<-c(b1,kv2,kv3,kv4) # kv 是基于sd下的斜率拟合

kv22<-kv2*x
kv33<-kv3*x
kv44<-kv4*x

kv1<-c(b1,kv22,kv33,kv44) # kv1是基于sd+d下的斜率拟合

# 画斜线
par(fig=c(0.08,0.6,0.4,0.98),new=T)

ux<-c(0.6,1,2,2.5)

plot(ux,b, col=0,bty="n",xlab = expression(italic(f["p"]) (KHz)),
          ylab = expression(italic(Slope)), mgp=c(1.1, 0, 0),tck=0.02,
               main = "", xlim = c(0.5,2.5),ylim=c(0,2.2),cex=0.8)

               lines(ux,b,col="red",lwd=1.5,lty=2,pch=pcc[1],type="b")
               lines(ux,kv,col="blue",lwd=1.5,lty=2,pch=pcc[2],type="b")
               lines(ux,kv1,col="black",lwd=1.5,lty=2,pch=pcc[3],type="b")

leg1<-c("fit","fit at Sd","fit at Sd+d")

legend("topright",legend=leg1,col=c("red","blue","black"),pch=c(0,1,2),lwd=1.5,lty=2,inset=.02,bty="n")


# There are some config parameters that depends on the data it-self. In this file we are going to create the functions
# that allow us to compute the best config parameter for Data Processing in the cellSizeDistribution step.

# CalculateBarcodeInflections calculates an adaptive inflection point ("knee")
# of the barcode distribution for each sample group. This is
# useful for determining a threshold for removing low-quality
# samples.
generate_default_values_cellSizeDistribution <- function(seurat_obj, config) {
  # Q: should we check for precalculated values? e.g.:
  # is.null(Tool(seurat_obj, slot = "CalculateBarcodeInflections"))
  # Returns Seurat object with a new list in the `tools` slot,
  # `CalculateBarcodeInflections` including inflection point calculatio

  seurat_obj_tmp <- CalculateBarcodeInflections(
                                        object = seurat_obj,
                                        barcode.column = "nCount_RNA",
                                        group.column = "orig.ident",
                                        # [HARDCODED]
                                        threshold.low = 1e2,
                                        threshold.high = NULL
  )
  # extracting the inflection point value which can serve as minCellSize threshold
  # all other side effects to the scdate object will be discarded
  # [TODO unittest]: this calculation assumes a single sample, i.e. only one group in orig.ident!
  # otherwise it will return multiple values for each group!
  # returned is both the rank(s) as well as inflection point
  # orig.ident          nCount_RNA  rank
  # SeuratProject       1106        10722
  tmp <- Tool(seurat_obj_tmp, slot = "CalculateBarcodeInflections")$inflection_points
  # extracting only inflection point(s)
  return(tmp$nCount_RNA)
}

cellSizeDistribution_config <- function(seurat_obj, config) {

    minCellSize <- generate_default_values_cellSizeDistribution(seurat_obj, config)
    # update config
    config$filterSettings$minCellSize <- minCellSize

    return(config)
}