require(mgcv)
require(methods)
source("projectMarginal.R")

# Load data
load("DSdata.rda")

# Remove outliner:
data = data[data$word != "Beziehungen",]

# Recode covariates to different metric
# ... word length in terms of log2 
data$ll  <- log2(1/data$l)
data$ll1 <- log2(1/data$l1)
data$ll2 <- log2(1/data$l2)
data$id  <- as.factor(data$id);
data$wid  <- as.factor(data$wid);

# ... probs of predictabilities 
data$prb <- exp(data$p)/(1+exp(data$p))
data$prb <- ifelse(data$prb==min(data$prb), 0, data$prb) # No perfectly predicted word was fixated 

data$prb1 <- exp(data$p1)/(1+exp(data$p1))
data$prb1 <- ifelse(data$prb1==min(data$prb1), 0, ifelse(data$prb1==max(data$prb1), 1, data$prb1))

data$prb2 <- exp(data$p2)/(1+exp(data$p2))
data$prb2 <- ifelse(data$prb2==min(data$prb2), 0, ifelse(data$prb2==max(data$prb2), 1, data$prb2))

data$s1s2 <- factor(paste(data$s1, data$s2, sep=""), labels=c("FF", "FT", "TF", "TT"))
data$x0x1 <- factor(paste(data$x, data$x1, sep=""), levels=c("00", "10", "01", "11"), labels=c("CC", "FC", "CF", "FF"))

# Select only FF:
all_data <- data[data$s1s2 == "FF",];
#subset_data <- all_data[1:1000,]
# Define dimension of basis
k = 10; k2 = k**2; n_cores=4;

# Fit simple spline model to all data
model_all_post <- bam(ldur ~ s(ls, k=k) + s(ll1,k=k) + s(ll,k=k) + s(ll2, k=k) 
                             + s(f1, f, k=k2) + s(f, f2, k=k2) + s(id, bs="re") + s(wid, bs="re"), data=all_data)
# Fit SS-ANOVA to all data
#model_all_ten <- bam(ldur ~ s(ls, k=k) + s(ll1,k=k) + s(ll,k=k) + s(ll2, k=k)
#                            + ti(f1, bs="cr", k=k) + ti(f, bs="cr", k=k) + ti(f2, bs="cr", k=k)
#                            + ti(f1,f, bs="cr", k=c(k,k)) + ti(f2,f, bs="cr", k=c(k,k)) + s(id, bs="re"),
#                     data=all_data, cluster=n_cores);

model_all_ten <- bam(ldur ~ s(ls, k=k) + s(ll1,k=k) + s(ll,k=k) + s(ll2, k=k)
                            + ti(f1, bs="cr", k=k) + ti(f, bs="cr", k=k) + ti(f2, bs="cr", k=k)
                            + ti(f1,f, bs="cr", k=c(k,k)) + ti(f,f2, bs="cr", k=c(k,k)) + s(id, bs="re") + s(wid, bs="re"),
                     data=all_data, cluster=n_cores);

#model_all_ten_word_sn <- bam(ldur ~ s(ls, k=k) + s(ll1,k=k) + s(ll,k=k) + s(ll2, k=k)
#                          + ti(f1, bs="cr", k=k) + ti(f, bs="cr", k=k) + ti(f2, bs="cr", k=k)
#                          + ti(f1,f, bs="cr", k=c(k,k)) + ti(f2,f, bs="cr", k=c(k,k)) + s(id, bs="re") + s(wid, bs="re", by=sn),
#                          data=all_data, cluster=n_cores);

newdata = list(ls=0, ll1=0, ll=0, ll2=0, prb1=0, prb=0, prb2=0, f1=0, f=0, f2=0, id=all_data$id[[1]], wid=all_data$wid[[1]]);
N_points = 64;
P_f1f <- projectMarginals(model_all_post, "s(f1,f)", newdata, N_points);
P_ff2 <- projectMarginals(model_all_post, "s(f,f2)", newdata, N_points);
post_m_f1 <- P_f1f$Px;            post_s_f1 <- sqrt(diag(P_f1f$Vx))
post_m_f  <- P_f1f$Py + P_ff2$Px; post_s_f  <- sqrt(diag(P_f1f$Vy)) + sqrt(diag(P_ff2$Vx))
post_m_f2 <- P_ff2$Py;            post_s_f2 <- sqrt(diag(P_ff2$Vy))
post_m_f1f <- P_f1f$Pxy; dim(post_m_f1f) <- c(N_points, N_points)
post_m_ff2 <- P_ff2$Pxy; dim(post_m_ff2) <- c(N_points, N_points)
post_s_f1f <- sqrt(diag(P_f1f$Vxy)); dim(post_s_f1f) <- c(N_points, N_points)
post_s_ff2 <- sqrt(diag(P_ff2$Vxy)); dim(post_s_ff2) <- c(N_points, N_points)
post_f1f <- predict.gam(model_all_post, sample_f1f, type="terms", se.fit=F)[, "s(f1,f)"]
post_ff2 <- predict.gam(model_all_post, sample_ff2, type="terms", se.fit=F)[, "s(f,f2)"]
dim(post_f1f) <- c(N_points, N_points); dim(post_ff2) <- c(N_points, N_points)

sample_f1 <- sampleData(model_all_ten, c("f1"), newdata, N_points);
sample_f  <- sampleData(model_all_ten, c("f"), newdata, N_points);
sample_f2 <- sampleData(model_all_ten, c("f2"), newdata, N_points);
sample_f1f <- expand.grid(f1=sample_f1$f1, f=sample_f$f, f2=0, ls=0, ll1=0, ll=0, ll2=0, prb1=0, prb=0, prb2=0, id=all_data$id[[1]], wid=all_data$wid[[1]])
sample_ff2 <- expand.grid(f1=0, f=sample_f$f, f2=sample_f2$f2, ls=0, ll1=0, ll=0, ll2=0, prb1=0, prb=0, prb2=0, id=all_data$id[[1]], wid=all_data$wid[[1]])
ten_f1 <- predict.gam(model_all_ten, sample_f1, type="terms", se.fit=T, terms=c("ti(f1)"));
ten_f  <- predict.gam(model_all_ten, sample_f, type="terms", se.fit=T, terms=c("ti(f)"));
ten_f2 <- predict.gam(model_all_ten, sample_f2, type="terms", se.fit=T, terms=c("ti(f2)"));
ten_f1f <- predict.gam(model_all_ten, sample_f1f, type="terms", se.fit=T, unconditional=T)
ten_m_f1f <- ten_f1f$fit[, "ti(f1,f)"]; ten_s_f1f <- ten_f1f$se.fit[, "ti(f1,f)"]; 
dim(ten_m_f1f) <- c(N_points, N_points); dim(ten_s_f1f) <- c(N_points, N_points)
ten_ff2 <- predict.gam(model_all_ten, sample_ff2, type="terms", se.fit=T, unconditional=T)
ten_m_ff2 <- ten_ff2$fit[, "ti(f,f2)"]; ten_s_ff2 <- ten_ff2$se.fit[, "ti(f,f2)"]; 
dim(ten_m_ff2) <- c(N_points, N_points); dim(ten_s_ff2) <- c(N_points, N_points)


pdf(file="example_compare_all.pdf", width=7, height=4)
ylim = c(-0.15,0.15)
layout(matrix(c(1,2,3,4,5,6), 2,3, byrow=T),
       widths=c(1.3,1.1,1), heights=c(1,1.3))
par(oma=c(1,1,1,1))
#
# Plot tensor on all data:
#
par(mar=c(1,4.1,0,1))
plot(sample_f$f, ten_f1$fit, ylim=ylim, type="l", lty=1, xlab="", ylab="log fix. dur.", xaxt="n"); 
lines(sample_f$f, ten_f1$fit+ten_f1$se.fit, lty=2); lines(sample_f$f, ten_f1$fit-ten_f1$se.fit, lty=2)
par(mar=c(1,0,0,1))
plot(sample_f$f, ten_f$fit, ylim=ylim, type="l", lty=1, xaxt="n", yaxt="n"); 
lines(sample_f$f, ten_f$fit+ten_f$se.fit, lty=2); lines(sample_f$f, ten_f$fit-ten_f$se.fit, lty=2)
par(mar=c(1,0,0,0))
plot(sample_f$f, ten_f2$fit, ylim=ylim, type="l", lty=1, xaxt="n", yaxt="n"); 
lines(sample_f$f, ten_f2$fit+ten_f2$se.fit, lty=2); lines(sample_f$f, ten_f2$fit-ten_f2$se.fit, lty=2)


par(mar=c(4.6,4.1,0,1))
plot(sample_f$f, post_m_f1, ylim=ylim, type="l", lty=1, xlab="log freq. word N-1", ylab="log fix. dur.", xaxt="s"); 
lines(sample_f$f, post_m_f1 + post_s_f1, lty=2); lines(sample_f$f, post_m_f1 - post_s_f1, lty=2)
par(mar=c(4.6,0,0,1))
plot(sample_f$f, post_m_f, ylim=ylim, type="l", lty=1, xlab="log freq. word N", ylab="", xaxt="s", yaxt="n"); 
lines(sample_f$f, post_m_f + post_s_f, lty=2); lines(sample_f$f, post_m_f - post_s_f, lty=2)
par(mar=c(4.6,0,0,0))
plot(sample_f$f, post_m_f2, ylim=ylim, type="l", lty=1, xlab="log freq. word N+1", ylab="", xaxt="s", yaxt="n"); 
lines(sample_f$f, post_m_f2 + post_s_f2, lty=2); lines(sample_f$f, post_m_f2 - post_s_f2, lty=2)

dev.off()


zlim = c(-0.25,0.25)
pdf(file="example_compare_sum_f1f.pdf", width=4.6, height=4)
filled.contour(sample_f$f, sample_f1$f1, t(post_f1f), color.palette=topo.colors, zlim=zlim)
dev.off()
pdf(file="example_compare_sum_ff2.pdf", width=4.6, height=4)
filled.contour(sample_f2$f2, sample_f$f, t(post_ff2), color.palette=topo.colors, zlim=zlim)
dev.off()
pdf(file="example_compare_interaction_f1f.pdf", width=4.6, height=4)
filled.contour(sample_f$f, sample_f1$f1, t(post_m_f1f), color.palette=topo.colors, zlim=zlim)
dev.off()
pdf(file="example_compare_interaction_ff2.pdf", width=4.6, height=4)
filled.contour(sample_f2$f2, sample_f$f, t(post_m_ff2), color.palette=topo.colors, zlim=zlim)
dev.off()
post_f1f_sig <- post_m_f1f; post_f1f_sig[post_f1f_sig < 2*post_s_f1f] <- NA; 
post_ff2_sig <- post_m_ff2; post_ff2_sig[post_ff2_sig < 2*post_s_ff2] <- NA; 
pdf(file="example_compare_sig_f1f.pdf", width=4.6, height=4)
filled.contour(sample_f$f, sample_f1$f1, t(post_f1f_sig), color.palette=topo.colors, zlim=zlim)
dev.off()
pdf(file="example_compare_sig_ff2.pdf", width=4.6, height=4)
filled.contour(sample_f2$f2, sample_f$f, t(post_ff2_sig), color.palette=topo.colors, zlim=zlim)
dev.off()

filled.contour(sample_f$f, sample_f1$f1, t(ten_m_f1f), color.palette=topo.colors, zlim=zlim)
filled.contour(sample_f2$f2, sample_f$f, t(ten_m_ff2), color.palette=topo.colors, zlim=zlim)
ten_sig_f1f <- ten_m_f1f; ten_sig_f1f[ten_sig_f1f < 2*ten_s_f1f] <- NA;
ten_sig_ff2 <- ten_m_ff2; ten_sig_ff2[ten_sig_ff2 < 2*ten_s_ff2] <- NA; 
filled.contour(sample_f$f, sample_f1$f1, t(ten_sig_f1f), color.palette=topo.colors, zlim=zlim)
filled.contour(sample_f2$f2, sample_f$f, t(ten_sig_ff2), color.palette=topo.colors, zlim=zlim)


# Plot other main effects of TPS
sample_ls <- sampleData(model_all_ten, c("ls"), newdata, N_points);
sample_l1 <- sampleData(model_all_ten, c("ll1"), newdata, N_points);
sample_l  <- sampleData(model_all_ten, c("ll"), newdata, N_points);
sample_l2 <- sampleData(model_all_ten, c("ll2"), newdata, N_points);
post_ls <- predict.gam(model_all_post, sample_ls, type="terms", se.fit=T, terms=c("s(ls)"));
post_l1 <- predict.gam(model_all_post, sample_l1, type="terms", se.fit=T, terms=c("s(ll1)"));
post_l  <- predict.gam(model_all_post, sample_l, type="terms", se.fit=T, terms=c("s(ll)"));
post_l2 <- predict.gam(model_all_post, sample_l2, type="terms", se.fit=T, terms=c("s(ll2)"));

old_par <- par(lwd=3)
pdf(file="example_compare_all_ls.pdf", width=4, height=3, pointsize=10)
plot(sample_ls$ls, post_ls$fit, ylim=c(-0.2, 0.2), xlab="saccade amplitude", ylab="log. fix. dur.", type="l")
lines(sample_ls$ls, post_ls$fit+post_ls$se.fit, lty=2); lines(sample_ls$ls, post_ls$fit-post_ls$se.fit, lty=2)
dev.off()

pdf(file="example_compare_all_l1.pdf", width=4, height=3, pointsize=10)
plot(sample_l1$ll1, post_l1$fit, ylim=c(-0.2, 0.2), xlab="log. word length (N-1)", ylab="log. fix. dur.", type="l")
lines(sample_l1$ll1, post_l1$fit+post_l1$se.fit, lty=2); lines(sample_l1$ll1, post_l1$fit-post_l1$se.fit, lty=2)
dev.off()

pdf(file="example_compare_all_l.pdf", width=4, height=3, pointsize=10)
plot(sample_l$ll, post_l$fit, ylim=c(-0.2, 0.2), xlab="log. word length (N)", ylab="log. fix. dur.", type="l")
lines(sample_l$ll, post_l$fit+post_l$se.fit, lty=2); lines(sample_l$ll, post_l$fit-post_l$se.fit, lty=2)
dev.off()

pdf(file="example_compare_all_l2.pdf", width=4, height=3, pointsize=10)
plot(sample_l2$ll2, post_l2$fit, ylim=c(-0.3, 0.3), xlab="log. word length (N+1)", ylab="log. fix. dur.", type="l")
lines(sample_l2$ll2, post_l2$fit+post_l2$se.fit, lty=2); lines(sample_l2$ll2, post_l2$fit-post_l2$se.fit, lty=2)
dev.off()

par(old_par)