library(plyr)
library(reshape2)
library(xtable)
options(xtable.floating=FALSE)
options(xtable.sanitize.text.function=function(x) x)
options(xtable.include.rownames=FALSE)
library(lattice)
library(latticeExtra)
sTheme<-standard.theme(color = FALSE)
sTheme$par.main.text$cex<-1
lattice.options(default.theme = sTheme)
library(lme4)
library(boot)
library(ks)
#
library(runjags)
#library(rjags)
library(coda)
runjags.options(silent.jags=TRUE, silent.runjags=TRUE) # ``progress'' is only shown in PDF, anyway
library(parallel)
#library(benchmarkme)

initJags <-function(chain) {
    chain<-as.numeric(chain)
    x<-as.list(my.init[[chain %% length(my.init) + 1]])
    x[[".RNG.seed"]]<-50+chain
    rngs<-c("base::Mersenne-Twister","base::Super-Duper","base::Wichmann-Hill","base::Marsaglia-Multicarry","lecuyer::RngStream")
    x[[".RNG.name"]]<-rngs[chain %% length(rngs) + 1]
    x
}
##
ssize<-10000;tthin<-1;chains<-8;

loadData <- function () {
    library(zTree)
    options(zTree.silent=FALSE) # <- less chatty
    setwd("raw/")
    allFiles<-list.files(".","*.xls$")
    dataZ <- zTreeTables(allFiles,tables=c("globals","subjects",'expectations','games'))
    setwd("..")
    save(file='dataZ0.Rdata',dataZ)
    load("dataZ0.Rdata")
    dataZZ<-join(dataZ$subjects,dataZ$globals,by=c("Date","Treatment","Period"))
    bidTreatments<-subset(ddply(dataZZ,~Date+Treatment,summarise,mPeriod=max(Period)),mPeriod==12)
    dataZZ<-match_df(dataZZ,bidTreatments)
    dataZZ<-subset(dataZZ,is.na(exp_x0)) # no expectation-treatments
    dataZZ<-subset(dataZZ,is.na(n_games)) # no risk-treatments
    dataZZ<-subset(dataZZ,is.na(oo) | oo==0) # no outside-option treatments
    oldbids<-is.na(dataZZ[["bid_x0"]])
    for(v in seq(0,50,10)) dataZZ[oldbids,sprintf("bid_x%d",v)]<-dataZZ[oldbids,sprintf("bid_%d",v+50)]
    dataZZ<-dataZZ[,c("Date","Period","Subject","Group","bid_x0","bid_x10","bid_x20","bid_x30","bid_x40","bid_x50","x","min_eingabe","max_eingabe","b_min","b_max","v_min","v_max")]
    dataZZ<-within(dataZZ,{
        Type<-as.factor(ifelse(Date %in% c("041130R0","041201NE","041201P8"),"2nd","1st"))
        sid<-as.factor(sprintf("%s-%02d",Date,Subject))
        min_eingabe[is.na(min_eingabe)]<-0
        x[is.na(x)]<-50
    })
    xxSid<-ddply(dataZZ,~sid+Date,summarise,minG=min(Group),maxG=max(Group))
    xxSid<-join(xxSid,ddply(xxSid,~Date,summarise,q1=quantile(maxG,.25),q2=quantile(minG,.75),critG=if(q1<q2) mean(q1,q2) else max(maxG))[,c(1,4)])
    dataZZ<-join(dataZZ,within(xxSid,gid<-as.factor(sprintf("%s-%d",Date,maxG<=critG)))[,c("Date","sid","gid")])
                                        #
    tlist<-ddply(dataZZ,~gid,summarise,omega=median(v_min),b=median(min_eingabe),type=Type[1],participants=length(unique(sid)))
    tlist[with(tlist,order(type,omega,b)),]
                                        #
    slist<-ddply(tlist,~omega+b+type,summarise,omega=as.character(median(omega)),b=as.character(median(b)),type=type[1],indep=length(gid),participants=sum(participants))
                                        #
    id<-c("Date","Period","sid",'gid','Type')
    dataL<-reshape(dataZZ,varying=grep('bid_x',names(dataZZ)),idvar=id,timevar='value',sep='_x',direction='long')
    dataL<-within(dataL,{
        bid<-bid-x
        over<-bid-value/ifelse(Type=='1st',2,1)
        treat<-as.factor(sprintf("%d%s%s",x,ifelse(Type=="2nd","II",""),ifelse(min_eingabe==0,'+','')))
        treat<-reorder(treat,x+2*(Type=="2nd")+(min_eingabe>0))
    })
                                        #
    xyplot(over~value,group=treat,data=dataL,type='smooth',auto.key=list(points=FALSE,lines=TRUE,columns=7))
    save(file='dataZ.Rdata',dataZZ,dataL,tlist,slist)
}

mer2init <- function(x.mer,preB=NULL,postB=NULL) {
    x.init<-NULL;i<-0
    while(TRUE) {
        i <- i+1;
        fixefInd<- (1:length(fixef(x.mer)))
        x.init<-rbind(x.init,bootMer(x.mer,function(x) c(fixef(x),sigma(x),sqrt(unlist(VarCorr(x)))),2,i)$t)
        x.init<-x.init[apply(t(t(abs(x.init)) > apply(x.init,2,mean)/4),1,min)>0,] # ← σ is not small!
        if(i>400) stop("can't finit init values with positive standard deviation for ME model")
        if(is.null(nrow(x.init))) next
        if(nrow(x.init)>=8) break
    }
    x<-NULL;
    for(i in 1:8) x[[i]] <- list(beta=c(preB,x.init[i,fixefInd],postB),tau=1/(x.init[i,-fixefInd])^2)
    x
}
#

slopeDD <- function(dd) {
    ## dd is the data with variables bid,value,compare
    ## compare defines the two conditions
    ## the function estimtes a single Bayes model
   dd<-within(dd,{value<-value-mean(value);compare<-compare-mean(compare)}) # ← normalise
   d1.mer<-lmer(bid~value*compare+(1|gid)+(1|sid),data=dd) # ← ME model to get init values for Bayes
   d1.init<- mer2init(d1.mer)
   d1.model<-'model {
       for (i in 1:length(bid)) {
         bid[i]~dnorm(beta[1]+beta[2]*v[i]+beta[3]*c[i]+beta[4]*c[i]*v[i]+uG[gid[i]]+uS[sid[i]],tau[1])
       }
       # group specific effects:
       for (s in 1:max(gid)) {uG[s]   ~ dnorm(0,tau[3])}
       for (s in 1:max(sid)) {uS[s]   ~ dnorm(0,tau[2])}
       for (i in 1:4) { beta[i] ~ dnorm(0,.0001) }
       for (i in 1:3) { tau[i]   ~ dgamma(m[i]*m[i]/(r[i]*r[i]),m[i]/(r[i]*r[i]))
                        m[i] ~ dgamma(1,1);r[i] ~ dgamma(1,1);}
   }'
   my.init<<-d1.init
   d1.data<-with(dd,list(bid=bid,c=as.numeric(compare),v=value,gid=as.numeric(factor(gid)),sid=as.numeric(factor(sid))))
   run.jags(model=d1.model,data=d1.data,monitor=c("beta","tau"),inits=initJags,method="parallel",n.chains=chains)
}

slopeComp <- function(dd) {
    ## dd is the data with variables bid,value,compare
    ## compare defines the two conditions
    ## the function provides a Bayes model comparison with and without conditioning on “compare”
   #
   dd<-within(dd,{value<-value-mean(value);compare<-compare-mean(compare)}) # ← normalise
   #
   d0.mer<-lmer(bid~value+(1|gid)+(1|sid),data=dd) # ← ME model to get init values for Bayes
   d0.init<- mer2init(d0.mer)
   #
   d1.mer<-lmer(bid~value*compare+(1|gid)+(1|sid),data=dd) # ← ME model to get init values for Bayes
   d1.init<- mer2init(d1.mer)
   # 
   d0.model<-'model {
       for (i in 1:length(bid)) {
         bid[i]~dnorm(beta[1]+beta[2]*v[i]+uG[gid[i]]+uS[sid[i]],tau[1])
       }
       # group specific effects:
       for (s in 1:max(gid)) {uG[s]   ~ dnorm(0,1/tau[3])}
       for (s in 1:max(sid)) {uS[s]   ~ dnorm(0,1/tau[2])}
       for (i in 1:2) { beta[i] ~ dnorm(0,.0001) }
       for (i in 1:3) { tau[i]   ~ dgamma(m[i]*m[i]/(r[i]*r[i]),m[i]/(r[i]*r[i]))
                        m[i] ~ dgamma(1,1);r[i] ~ dgamma(1,1);}
   }'
   #
   d1.model<-'model {
       for (i in 1:length(bid)) {
         bid[i]~dnorm(beta[1]+beta[2]*v[i]+beta[3]*c[i]+beta[4]*c[i]*v[i]+uG[gid[i]]+uS[sid[i]],tau[1]^2)
       }
       # group specific effects:
       for (s in 1:max(gid)) {uG[s]   ~ dnorm(0,tau[3]^2)}
       for (s in 1:max(sid)) {uS[s]   ~ dnorm(0,tau[2]^2)}
       for (i in 1:4) { beta[i] ~ dnorm(0,.0001) }
       for (i in 1:3) { tau[i]   ~ dgamma(m[i]*m[i]/(r[i]*r[i]),m[i]/(r[i]*r[i]))
                        m[i] ~ dgamma(1,1);r[i] ~ dgamma(1,1);}
   }'
   #
   my.init<<-d0.init
   d1.data<-with(dd,list(bid=bid,c=as.numeric(compare),v=value,gid=as.numeric(factor(gid)),sid=as.numeric(factor(sid))))
   system.time(d0.bayes<-run.jags(model=d0.model,data=d1.data,monitor=c("beta","tau"),inits=initJags,method="parallel",n.chains=chains))
   #
   my.init<<-d1.init
   system.time(d1.bayes<-run.jags(model=d1.model,data=d1.data,monitor=c("beta","tau"),inits=initJags,method="parallel",n.chains=chains))
   d1.bayes
   #
   #
   d01.model<-'model {
       for (i in 1:length(bid)) {
         mu[1,i]<-beta[1]+beta[2]*v[i]+uG0[gid[i]]+uS0[sid[i]]
         mu[2,i]<-beta[3]+beta[4]*c[i]+beta[5]*c[i]*v[i]+beta[6]*v[i]+uG1[gid[i]]+uS1[sid[i]]
         bid[i] ~dnorm(mu[case+1,i],preci[case+1])
       }
       for (s in 1:max(gid)) {     # group specific effects:
         uG0[s]   ~ dnorm(0,tau[1]))
         uG1[s]   ~ dnorm(0,tau[4]))
       }
       for (s in 1:max(sid)) {     # subject specific effects:
          uS1[s]   ~ dnorm(0,tau[2]))
          uS0[s]   ~ dnorm(0,tau[5]))
       }
       for (i in 1:length(mBeta)) { # fixed effects:
          beta[i]   ~ dnorm(mBeta[i],precBeta[i])
       }
       for (i in 1:length(Atau)) {   # all standard deviations:
          tau[i] ~ dgamma(Atau[i],Btau[i])
       }
       preci[1]<-tau[3])
       preci[2]<-tau[6])
       case~dbern(p1)
       p1 ~ dbeta(1,1)
   }'
   d01.stats<-rbind(summary(d0.bayes)[["statistics"]],summary(d1.bayes)[["statistics"]])
   d01.beta.stats<-d01.stats[grep("^beta",rownames(d01.stats)),]
   d01.tau.stats<-d01.stats[grep("^tau",rownames(d01.stats)),]
   d01.beta.stats
   d1.data[["mBeta"]]<-d01.beta.stats[,1]
   d1.data[["precBeta"]]<-1/(d01.beta.stats[,2])^2
   d1.data[["Atau"]]<-d01.tau.stats[,1]^2/d01.tau.stats[,2]^2
   d1.data[["Btau"]]<-d01.tau.stats[,1]  /d01.tau.stats[,2]^2
   #
   xx<-cbind(head(d0.bayes$mcmc[[1]]),head(d1.bayes$mcmc[[1]]))
   betas<-xx[,grep("^beta",colnames(xx))]
   taus<-xx[,grep("^tau",colnames(xx))]
   my.init<<-NULL
   for(i in 1:4) my.init[[i]]<<-list(beta=betas[i,],tau=taus[i,])
   #                                        #
   d01.comp<-run.jags(model=d01.model,data=d1.data,monitor=c("beta","tau","p1"),inits=initJags,method="parallel",n.chains=chains)
}

##--------------------------------------------------------------------------------
##                               MIXTURE
##--------------------------------------------------------------------------------
##
#
TYPES<-c("RAND","MD","FLEX","BNE")
runMixture <- function() {
ssize<-5000;tthin<-1;chains<-4;
ssize<-10000;tthin<-10;chains<-8;
#
# for the Bayes Regression / Mixture Model below we use the following approximation, which is
# exact only for ρ∈{0,1}.
#
#xx<-within(bids,{rho<-as.numeric(as.character(rho));
#   bA<-x/(1+r)-(1-rho)*r^2/(1+r*(1-rho))/2;bD<-bA-bid})
#xyplot(bD ~ x | risk,group=rho,data=xx,auto.key=TRUE)
#
#--------------------
# interestingly: it does not work (well) to estimate the mix.model jointly for the different treatments.
# this forces all the priors to follow the same distribution, and should help stability, but it is too much:
# most bidders are booked into “flex”, practically nobody into “markdown” (except for 50II+).
# hence we go for treatment-wise estimation of the mix-model
# we use two steps model comparison. In a first step pseudopriors
# are constructed. Then in a second step the model is selected.
# Shrinkage for coefficients is here provided through mu,t (this is necessary)
# 
#--------------------------------------------------
mix.model<-'model {
       for (i in 1:length(bid)) {
         mmu[i] <- beta[kI[sid[i]],1,sid[i]] + v[i]*beta[kI[sid[i]],2,sid[i]]+nuS[kI[sid[i]],sid[i]]+nuG[kI[sid[i]],gid[sid[i]]]
         bid[i] ~ dnorm(mmu[i],tau[1])
       }
       for (i in 1:max(sid)) {
         beta[1,1,i] ~ dnorm(mu[1,1],t[1,1]) # random
         beta[1,2,i] <- 0                    # random
         beta[2,1,i] ~ dnorm(mu[2,1],t[2,1]) # markdown
         beta[2,2,i] <- 1                    # markdown
         beta[3,1,i] ~ dnorm(mu[3,1],t[3,1]) # flex
         beta[3,2,i] ~ dnorm(mu[3,2],t[3,2]) # flex
         beta[4,1,i] <- 0                    # BNE
         beta[4,2,i] ~ dnorm(mu[4,2],t[4,2]) # BNE
         kI[i] ~ dcat(g)
         for (k in 1:4) {
            nuS[k,i] ~ dnorm(0,tau[3])          # RE for subjects
         }
       }
       #
       # 
       for (j in 1:3) {                      # priors for RE: 
          tau[j]~ dgamma(mTau[j]^2/sTau[j]^2,mTau[j]/sTau[j]^2)
          mTau[j] ~ dexp(1);sTau[j] ~ dexp(1);
       }
       for (k in 1:4) {
          for (i in 1:max(gid)) {
             nuG[k,i] ~ dnorm(0,tau[2]) # RE for groups
          }
          for(j in 1:2) {
            mu[k,j] ~ dnorm(0,.0001)
            t[k,j]  ~ dgamma(mT[k,j]^2/sT[k,j]^2,mT[k,j]/sT[k,j]^2)
            mT[k,j] ~ dexp(1);sT[k,j] ~ dexp(1)
          }
          gamma[k]~dexp(1);
          g[k]<-gamma[k]/sum(gamma)
          kS[k]<-sum(kI==k) # measure actual frequency
       }
   }'
#--------------------
mix1.jags<-function(x,kI) {
    cat("****************************************\n")
    cat("***  treat = ",x$treat[1]," kI=",kI,"*** \n")
    cat("****************************************\n")
    sidN<-length(levels(factor(x$sid)))
    mixP.data<-with(x,list(bid=bid,v=value,gid=as.numeric(factor(gid)),sid=as.numeric(factor(sid)),kI=rep(kI,sidN)))
    mixPi.jags<-run.jags(model=mix.model,data=mixP.data,monitor=c("mu","tau"),inits=initJags,method="parallel",n.chains=chains,sample=ssize,thin=tthin,summarise=TRUE)
    mixPi.df<-as.data.frame(as.mcmc(mixPi.jags))
    mu.df<-mixPi.df[,grep(sprintf("mu\\[%g,",kI),names(mixPi.df))]
    tau.df<-mixPi.df[,grep(sprintf("tau\\[%g",kI),names(mixPi.df))]
    mu<-unlist(lapply(mu.df,mean))
    t<-pmin(10^6,1/unlist(lapply(mu.df,var)))
    list(mu=mu,t=t,sseff=mixPi.jags[["mcse"]][["sseff"]],psrf=mixPi.jags[["psrf"]][["psrf"]][,"Point est."])
}
#--------------------
mix.jags<- function(x) {
    sidN<-length(levels(factor(x$sid)))
    ## first round: generate pseudopriors
    ## these are based on vague priors for beta and tau 
    ## now the pseudopriors for mu,t are constructed
    mu<-array(NA,c(4,2))
    t<-array(NA,c(4,2))
    psrf<-array(NA,c(4,11))
    sseff<-array(NA,c(4,11))
    mTau<-c(NA,NA,NA)
    sTau<-c(NA,NA,NA)
    for (k in 1:4) {
        x.jags<-mix1.jags(x,k)
        mu[k,]<-x.jags$mu
        t[k,]<-x.jags$t
        psrf[k,]<-x.jags$psrf
        sseff[k,]<-x.jags$sseff
    }
    mix.data<-with(x,list(bid=bid,v=value,gid=as.numeric(factor(gid)),sid=as.numeric(factor(sid)),gamma=c(NA,NA,NA,1),mu=mu,t=t))
    save(mu,t,psrf,sseff,file=sprintf("stats1-%s.Rdata",x$treat[1]))
    run.jags(model=mix.model,data=mix.data,monitor=c("g","kS"),inits=initJags,method="parallel",n.chains=chains,sample=ssize,thin=tthin)
}
#--------------------
my.init<<-NULL
dd<-subset(dataL,Period>6)
ddOther<-within(subset(dd,!(as.character(treat) %in% c("0+","50II+"))),treat<-"other 1st")
runjags.options(silent.jags=FALSE, silent.runjags=FALSE,modules="glm")
# Debug:
# runjags.options(silent.jags=FALSE, silent.runjags=FALSE);ssize=1000;tthin=10
# x<-subset(dd,treat=="0+")
# x.jags<-mix.jags(x) # take 1.5 minutes
mixG.time<-(system.time(mixG.j.all<-dlply(dd,~treat,mix.jags)))
mixG.other<-mix.jags(ddOther)
save(mixG.j.all,mixG.time,mixG.other,file="mixG.Rdata")
}

getStatsFile <- function(treat) {
    fn<-sprintf("stats1-%s.Rdata",treat)
    if(!file.exists(fn)) return(list())
    load(fn)
    BIG=10^20
    mu[1,2] <- 0                    # random
    mu[2,2] <- 1                    # markdown
    mu[4,1] <- 0                    # BNE
    t[1,2] <- BIG                    # random
    t[2,2] <- BIG                    # markdown
    t[4,1] <- BIG                    # BNE
    list(mu=mu,sd=1/sqrt(t),psrf=psrf,sseff=sseff)
}





##--------------------------------------------------------------------------------
