# File: SourceTracker.r
# Author: Dan Knights
# Contact: danknights@gmail.com
# License: GPL
# Copyright: Copyright 2011, Dan Knights
# Version: 0.9.1 (Beta)

# Function "sourcetracker"
#
# Gets prior counts of taxa in source environments for use by SourceTracker.
#
# Params:
# train - a sample x feature observation count matrix
# envs - a factor or character vector indicating the sample environments
# rarefaction_depth - if not NULL, all samples with > rarefaction_depth sequences are rarified;
#            This decreases the influence of high-coverage source samples.
#
# Value:
# Returns an object of class "sourcetracker". This is a list containing:
# sources - matrix of the total counts of each taxon in each environment
# train - a copy of the input training data
# envs - a copy of the input environment vector
"sourcetracker" <- function(train, envs, rarefaction_depth=1000){
    train <- as.matrix(train)

    # enforce integer data
    if(sum(as.integer(train) != as.numeric(train)) > 0){
        stop('Data must be integral. Consider using "ceiling(datatable)" or ceiling(1000*datatable) to convert floating-point data to integers.')
    }
    envs <- factor(envs)
    train.envs <- sort(unique(levels(envs)))

    # rarefy samples above maxdepth if requested
    if(!is.null(rarefaction_depth) && rarefaction_depth > 0) train <- rarefy(train, rarefaction_depth)
    # if(!is.null(rarefaction_depth) && rarefaction_depth > 0) cat(sprintf('Rarefying training data at %d\n',rarefaction_depth))

    # get source environment counts
    # sources is nenvs X ntaxa
    sources <- t(sapply(split(data.frame(train), envs), colSums))

    # add an empty row for "Unknown"
    sources <- rbind(sources, rep(0,ncol(train)))
    rownames(sources) <- c(train.envs,"Unknown")
    colnames(sources) <- colnames(train)
    sources <- as.matrix(sources)

    ret <- list(sources=sources, train=train, envs=envs)
    class(ret) <- "sourcetracker"

    return(invisible(ret))
}

# Function "predict.sourcetracker"
#
# S3-level function to estimate source proportions using a sourcetracker object.
# Returns an object of class "sourcetracker.fit"
#
# Params:
# stobj - output from function "sourcetracker"
# test - test data, a matrix of sample x feature counts
#   if test is NULL, performs leave-one-out predictions of the training
#   samples
# burnin - number of "burn-in" passes for Gibbs sampling
# nrestarts - number of times to restart the Gibbs sampling process
# n.draws.per.restart - number of Gibbs draws to collect per restart
# delay - number of passes between draws (ignored if n.draws.per.restart is 1)
# alpha1 - prior counts of each species in the training environments,
#   Higher values decrease the trust in the training data, and make the
#   source environment distributions over taxa smoother. By default, this is
#   set to 1e-3, which indicates reasonably high trust in all source environments, even
#   those with few training sequences. This is useful when only a small number
#   of biological samples are available from a source environment.
#   A more conservative value would be 0.001 or 0.01.
# alpha2 - prior counts of each species in the Unknown environment,
#   Higher values make the Unknown environment smoother and less prone to
#   over-fitting a given training sample. Default is 1e-3.
# beta - prior counts of test sequences in each environment.
#   Higher values cause a smoother distribution over source environments.
#   Default is 1e-2.
# rarefaction_depth - if not NULL, all test samples with > maxdepth sequences are rarified
# verbosity - if > 0, print progress updates while running
# full.results - return full draws from gibbs sampling (all assignments of all taxa)
#
# Value: A list containing:
# draws - an array of dimension (ndraws X nenvironments X nsamples),
#     containing all draws from gibbs sampling
# proportions - the mean proportions over all Gibbs draws for each sample
# proportions_sd - standard deviation of the mean proportions for each sample
# train.envs - the names of the source environments
# samplenames - the names of the test samples
"predict.sourcetracker" <- function(stobj, test=NULL,
            burnin=100, nrestarts=10, ndraws.per.restart=1, delay=10,
            alpha1=1e-3, alpha2=1e-1, beta=10, rarefaction_depth=1000,
            verbosity=1, full.results=FALSE){

    if(!is.null(test)){
        # if test is numeric, cast as a row matrix
        if(class(test) == "numeric" || class(test) == "integer"){
            test <- matrix(test, nrow=1)
        } else {
            test <- as.matrix(test)
        }
        if(sum(as.integer(test) != as.numeric(test)) > 0){
            stop('Data must be integral. Consider using "ceiling(datatable)" or ceiling(1000*datatable) to convert floating-point data to integers.')
        }
        sources <- stobj$sources
        if(verbosity>=1) {
            cat(rep(' ',nrestarts * ndraws.per.restart+1),sep='')
            cat(rep(' ',27),sep='')
            cat(sprintf('%11s',substr(rownames(sources),1,10)),sep='\t'); cat('\n')
        }
        T <- ncol(sources) # number of taxa
        V <- nrow(sources) # number of source envs
        if(is.null(dim(test))) N <- 1
        else N <- nrow(test) # number of sink samples
        samplenames <- rownames(test)
        draws <- run.gibbs(sources, test, V, T, N,
                burnin=burnin, nrestarts=nrestarts,
                ndraws.per.restart=ndraws.per.restart, delay=delay,
                alpha1=alpha1, alpha2=alpha2, beta=beta, maxdepth=rarefaction_depth,
                verbosity=verbosity, full.results=full.results)
        if(full.results) {
            full.draws <- draws$full.draws
            draws <- draws$draws
        }
    } else {  # leave-one-out
        samplenames <- rownames(stobj$train)
        envs <- stobj$envs

        T <- ncol(stobj$train) # number of taxa
        V <- nrow(stobj$sources) # number of source envs
        N <- nrow(stobj$train) # number of sink samples
        ndraws <- nrestarts * ndraws.per.restart # total number of draws
        draws <- array(0,dim=c(ndraws, V, N))
        cat(sprintf('ndraws=%d, V=%d, T=%d, N=%d\n',ndraws, V, T, N))
        if(full.results){
            full.draws <- array(0,dim=c(ndraws, V, T, N))
        }
        for(i in (1:N)){
            stobj.i <- sourcetracker(stobj$train[-i,], envs[-i], rarefaction_depth=rarefaction_depth)
            sources <- stobj.i$sources
            V.i <- nrow(sources) # number of source envs (might be missing one if there's only one sample from this env)
            draws.i <- run.gibbs(sources, stobj$train[i,], V.i, T, 1,
                burnin=burnin, nrestarts=nrestarts,
                ndraws.per.restart=ndraws.per.restart, delay=delay,
                alpha1=alpha1, alpha2=alpha2, beta=beta, maxdepth=rarefaction_depth,
                verbosity=verbosity, printing.index=i, printing.total=N, full.results=full.results)
            if(full.results){
                full.draws.i <- draws.i$full.draws
                draws.i <- draws.i$draws
            }
            # if(verbosity >= 1) cat(sprintf('%3d of %d: ',i,N))
            # handle case where there are no other samples from this env
            if(sum(envs[-i] == envs[i])==0){
                draws[,-which(rownames(stobj$sources)==envs[i]),i] <- drop(draws.i)
                if(full.results){
                    full.draws[,-which(rownames(stobj$sources)==envs[i]),,i] <- drop(full.draws.i)
                }
            } else {
                draws[,,i] <- drop(draws.i)
                if(full.results){
                    full.draws[,,,i] <- drop(full.draws.i)
                }
            }
        }
    }

    proportions <- matrix(nrow=N, ncol=V)
    proportions_sd <- matrix(nrow=N, ncol=V)
    for(i in 1:N){
        proportions[i,] <- apply(matrix(draws[,,i], ncol=V),2,mean)
        proportions_sd[i,] <- apply(matrix(draws[,,i], ncol=V),2,sd)
    }
    rownames(proportions) <- samplenames
    colnames(proportions) <- rownames(stobj$sources)
    rownames(proportions_sd) <- samplenames
    colnames(proportions_sd) <- rownames(stobj$sources)

    res <- list(draws=draws, proportions=proportions,
                proportions_sd=proportions_sd,
                train.envs=rownames(sources), samplenames=samplenames)

    if(full.results){
    	res$full.results <- full.draws
   	    dimnames(full.draws) <- list(
			sprintf('draw%05d',1:dim(full.draws)[1]),
			rownames(sources),
			colnames(sources),
			samplenames
    	)
	}
    class(res) <- "sourcetracker.fit"
    return(invisible(res))
}

# Function "plot.sourcetracker.fit"
#
# S3-level function to plot the SourceTracker output.
#
# Params:
# stresult - output from function "predict.sourcetracker"
# labels - Labels for samples; if NULL, uses rownames of data table
# type - One of 'pie', 'bar', or 'dist' (distribution plots). Default is 'pie'.
# gridsize - number of samples to plot per row; if NULL, will be estimated
# ... - Additional graphical parameters
"plot.sourcetracker.fit" <- function(stresult, labels=NULL,
        type=c('pie','bar','dist')[1], gridsize=NULL, env.colors=NULL,
        titlesize=NULL, indices=NULL, include.legend=FALSE, ...){
    if(is.null(env.colors)){
        env.colors <- std.env.colors
        # always set 'Unknown' to grey
        env.colors[stresult$train.envs=='Unknown'] <- std.env.colors[length(std.env.colors)]
    }

    if(is.null(indices)) indices <- 1:dim(stresult$draws)[3]
    N <- length(indices)
    V <- dim(stresult$draws)[2]

    if(include.legend) N <- N + 1

    if(!is.null(gridsize) && gridsize**2 < N)
        stop(sprintf('Please choose a gridsize of at least %d.',ceiling(sqrt(N))))

    if(is.null(labels)) labels <- stresult[['samplenames']]
    if(is.null(gridsize)) gridsize <- ceiling(sqrt(N))
    if(is.null(titlesize)){
        if(gridsize > 1){
            titlesize <- .3/log10(gridsize)
        } else {
            titlesize=1
        }
    }

    ngridrows <- ceiling(N / gridsize)
    par(mfrow=c(ngridrows,gridsize))
    par(oma=c(1,1,1,1), mar=c(0,0,titlesize,0))

    # legend will occupy one full plot in the upper left
    if(include.legend){
        plot(0,0,xlim=c(0,1), ylim=c(0,1), type='n', axes=FALSE)
        leg.cex <- 0.1
        leg <- legend('topleft',stresult$train.envs, fill=env.colors, bg='white', cex=leg.cex, plot=FALSE)
        maxdim <- max(leg$rect$w, leg$rect$h)

        # resize legend to be just big enough to fill the plot (80%), or to have maximum text size 2
        while(maxdim <.8 && leg.cex <= 2){
            leg.cex <- leg.cex + 0.01
            leg <- legend('topleft',stresult$train.envs, fill=env.colors, bg='white', cex=leg.cex, plot=FALSE)
            maxdim <- max(leg$rect$w, leg$rect$h)
        }
        leg <- legend('topleft',stresult$train.envs, fill=env.colors, bg='white', cex=leg.cex, border=NA)
    }

    if(type=='pie') plot.sourcetracker.pie(stresult, labels, gridsize, env.colors, titlesize, indices=indices, ...)
    if(type=='bar') plot.sourcetracker.bar(stresult, labels, gridsize, env.colors, titlesize, indices=indices, ...)
    if(type=='dist') plot.sourcetracker.dist(stresult, labels, gridsize, env.colors, titlesize, indices=indices, ...)
}

######### Internal functions below ####################

# Internal SourceTracker function to run Gibbs sampling
# total.n is used to supply the total number of samples in leave-one-out
# predictions for printing status updates
# full.results returns the source env. x taxon counts for every draw
"run.gibbs" <- function(sources, test, V, T, N,
        burnin=100, nrestarts=10, ndraws.per.restart=10, delay=10,
        alpha1=1e-3, alpha2=.1, beta=10, maxdepth=NULL,
        verbosity=1, printing.index=NULL, printing.total=NULL,
        full.results=FALSE){

    if(is.null(printing.total)) printing.total <- N

    train.envs <- rownames(sources)
    ndraws <- nrestarts * ndraws.per.restart # total number of draws
    npasses <- burnin + (ndraws.per.restart-1) * delay + 1 # passes per restart

    # draws will hold all draws (ndraws x V x N)
    draws <- array(dim=c(ndraws, V, N))
    if(full.results){
        full.draws <- array(dim=c(ndraws, V, T, N))
    }

    # rarefy samples above maxdepth if requested
    if(!is.null(maxdepth))    test <- rarefy(test, maxdepth)

    # sink samples must have integer counts
    if(is.null(dim(test))) test <- matrix(test,ncol=T)
    test <- round(test)

    # store original prior counts for "Unknown"
    unknown.prior <- sources[V,]
    # sources[-V,] <- sweep(sources[-V,],1,alpha1 * rowSums(sources[-V,]),'+')  # add relative alpha prior counts
    sources[-V,] <- sources[-V,] + alpha1 # add absolute alpha prior counts

    # for each sink sample
    for(i in 1:N){
        sink <- test[i,]
        D <- sum(sink) # sink sample depth
        incProgress(1,detail= paste0('Passes: ', i, ' on ', N))
        # precalculate denominator for Pr(env in sample)
        p_v_denominator = max(1,(D-1) + V*beta)

        # get taxon index for each sequence
        tax.cumsum <- cumsum(sink)
        tax.ix <- sapply(1:D,function(x) min(which(x<=tax.cumsum)))

        drawcount <- 1 # keeps running count of draws for this sample
        # for each restart
        for(j in 1:nrestarts){
            if(verbosity>=1) cat('.')
        	prev.warn <- options()$warn
            options(warn=-1)
            z <- sample(V,D,replace=TRUE) # random env assignments
            options(warn=prev.warn)

            sources[V,] <- unknown.prior # prior counts of taxa in Unknown
            sources[V,] <- sources[V,] + alpha2 * D # add relative alpha prior counts

            # tally counts in envs
            # count all assignments to the "other" environment
            # other environments don't get incremented because they are fixed from training data
            envcounts <- rep(beta, V)
            for(ix in 1:D){
                if(z[ix] == V)    sources[V,tax.ix[ix]] <- sources[V,tax.ix[ix]] + 1
                envcounts[z[ix]] <- envcounts[z[ix]] + 1
            }

            for(rep in 1:npasses){
                rand.ix <- sample(D) # random order for traversing sequence
                # temporary: not random
                # rand.ix <- 1:D

                cnt <- 0
                for(ix in rand.ix){
                    taxon <- tax.ix[ix]
                     # remove this sequence from all counts
                    envcounts[z[ix]] <- envcounts[z[ix]] - 1
                    if(z[ix] == V)    sources[V,taxon] <- sources[V,taxon] - 1

                    # get relative PDF over env assignments
                    p_tv <- sources[,taxon] / rowSums(sources) # Pr(taxon | env)
                    p_v <- envcounts/p_v_denominator# Pr(env in sample)

                    # re-sample this sequence's env assignment
                    z[ix] <- sample(1:V, prob=p_tv * p_v, size=1)

                    # replace this sequence in all counts
                    envcounts[z[ix]] <- envcounts[z[ix]] + 1

                    # if this sequence is assigned to "other", increase count
                    if(z[ix] == V)    sources[V,taxon] <- sources[V,taxon] + 1
                }

                # take sample
                if(rep > burnin && (((rep-burnin) %% delay)==1 || delay<=1)){

                    # save current mixing proportions
                    draws[drawcount,,i] <- round((envcounts - beta) / D,7)
                    draws[drawcount,,i] <- draws[drawcount,,i] / sum(draws[drawcount,,i])
                    # save full taxon-source assignments if requested
                    if(full.results){
                        # for each environment, save taxon counts
                        for(j in 1:V){
                            full.draws[drawcount,j,,i] <- sapply(1:T,function(x) sum(tax.ix[z==j]==x))
                        }
                    }
                    drawcount <- drawcount + 1
                }
            }
        }

        if(verbosity>=1){
            if(is.null(printing.index)){
                cat(sprintf('%4d of %4d, depth=%5d: ', i, printing.total, D))
            } else {
                cat(sprintf('%4d of %4d, depth=%5d: ', printing.index, printing.total, D))
            }
            props <- colMeans(matrix(draws[,,i],ncol=V))
            prop_devs <- apply(matrix(draws[,,i], ncol=V), 2, sd)
            cat(' ')
            cat(sprintf('%.2f (%.2f)', props, prop_devs),sep='\t')
            cat('\n')
        }
    }
    if(full.results){
        return(list(draws=draws, full.draws=full.draws))
    } else {
        return(draws=draws)
    }
}

# tries all values of alpha1 and alpha2 for best RMSE
# if individual.samples, tries to predict mixtures of single samples
# instead of mixtures of the environment means
# ntrials is the number of simulated samples per fit.
# nrepeats is the number of times to repeat the entire experiment at each alpha value
# verbosity > 1 means inner loop will print
"tune.st" <- function(otus, envs, individual.samples=TRUE, ntrials=25,
            rarefaction_depth=1000, alpha1=10**(-3),
            alpha2=10**(-3:0), beta=10, verbosity=0, ...){
    results <- list()
    alphas <- expand.grid(rev(alpha1), alpha2)
    colnames(alphas) <- c('alpha1','alpha2')
    rmse <- numeric(nrow(alphas))
    rmse.sem <- numeric(nrow(alphas))
    for(i in 1:nrow(alphas)){
        cat(sprintf('Loop %d of %d, alpha1=%f, alpha2=%f ',i,nrow(alphas),alphas[i,1], alphas[i,2]))
        if(verbosity > 2) cat('\n')
        results[[i]] <- eval.fit(otus, envs, individual.samples=individual.samples,
                                ntrials=ntrials, rarefaction_depth=rarefaction_depth,
                                alpha1=alphas[i,1], alpha2=alphas[i,2], beta=beta, verbosity=verbosity-1, ...)
        rmse[i] <- results[[i]]$rmse
        rmse.sem[i] <- results[[i]]$rmse.sem
        if(verbosity > 0) cat(sprintf('RMSE = %.3f +/- %.3f\n',rmse[i], rmse.sem[i]))
    }
    # choose alpha as most conservative value of alpha2 (smallest)
    # then most conservative value of alpha1 (largest)
    # that gives pseudo-r2 within 1 sem of the max.
    # best.ix <- min(which(rmse <= min(rmse + rmse.sem)))
    # Alternative: simply choose the lowest rmse
    best.ix <- which.min(rmse)
    best.rmse <- rmse[best.ix]
    best.alpha1 <- alphas[best.ix,1]
    best.alpha2 <- alphas[best.ix,2]
    return(list(alphas=alphas, rmse=rmse, rmse.sem=rmse.sem, best.rmse=best.rmse,
                best.alpha1=best.alpha1, best.alpha2=best.alpha2,
                results=results))
}

# train SourceTracker object on training data
# returns RMSE
# ... are additional params to pass to sourcetracker predict object
"eval.fit" <- function(otus, envs, individual.samples=TRUE,
            ntrials=25, rarefaction_depth=1000, verbosity=1, ...){
    train.envs <- sort(unique(envs))
    V <- length(train.envs)
    env.sizes <- table(envs)

    # make sure each pair of envs gets picked
    # build up all pairs of samples, each column is a pair
    # each source env gets to be first and second sample once
    # pairs <- expand.grid(1:V,1:V)
    # pairs <- pairs[pairs[,1]!=pairs[,2],]
    # make nreps pairs randomly
    pairs <- NULL
    for(i in 1:ntrials){
        pairs <- rbind(pairs, sample(V,size=2))
    }

    mixtures <- runif(ntrials)
    y <- matrix(0,nrow=ntrials, ncol=V+1)
    yhat <- matrix(0,nrow=ntrials, ncol=V+1)
    yhat.sd <- matrix(0,nrow=ntrials, ncol=V+1)
    colnames(y) <- c(as.character(train.envs),'Unknown')
    colnames(yhat) <- c(as.character(train.envs),'Unknown')
    newsamples <- NULL
    allenvs <- NULL
    for(i in 1:ntrials){
        env1 <- pairs[i,1]
        env2 <- pairs[i,2]
        allenvs <- rbind(allenvs, c(env1, env2))
        if(verbosity > 1){
           cat(sprintf('%d of %d: %.2f*%s + %.2f*%s: \n',i,ntrials,mixtures[i], train.envs[env1],1-mixtures[i], train.envs[env2]))
        } else if(verbosity > 0){
            cat('.')
        }

        # all indices of each environment
        env1.ix.all <- which(envs == train.envs[env1])
        env2.ix.all <- which(envs == train.envs[env2])

        if(individual.samples){
            # get one sample from each env
            # cast as list so that sample doesn't misinterpret a length-1 vector
            env1.ix <- sample(as.list(env1.ix.all),size=1)[[1]]
            env2.ix <- sample(as.list(env2.ix.all),size=1)[[1]]

            # train sourcetracker, hold out entire second env. and first env. sample
            # note: don't hold out first sample if that env has only one sample
            if(length(env1.ix.all) == 1){
                st <- sourcetracker(otus[-env2.ix.all,], envs[-env2.ix.all])
            } else {
                st <- sourcetracker(otus[-c(env1.ix,env2.ix.all),], envs[-c(env1.ix,env2.ix.all)])
            }

            # make fake sample, weighted mixture of two source samples
            s1 <- otus[env1.ix,]
            s2 <- otus[env2.ix,]

        } else {
            # train sourcetracker, hold out entire second env.
            st <- sourcetracker(otus[-env2.ix.all,], envs[-env2.ix.all])

            # make fake sample as mixture of _environment_ means
            s1 <- colSums(rarefy(otus[env1.ix.all,], maxdepth=rarefaction_depth))
            s2 <- colSums(rarefy(otus[env2.ix.all,], maxdepth=rarefaction_depth))
        }

        newsample <- mixtures[i] * s1/sum(s1) + (1-mixtures[i]) * s2/sum(s2)
        newsample <- round(100000 * newsample)
        newsample <- matrix(newsample, nrow=1)
        newsample <- rarefy(newsample,maxdepth=ceiling(sum(s1+s2)/2))
        newsamples <- rbind(newsamples, newsample)
        y[i,env1] <- mixtures[i]
        y[i,V+1] <- 1-mixtures[i]

        # test on fake sample
        results <- predict(st, newsample, rarefaction_depth=rarefaction_depth, verbosity=verbosity-1, ...)
        for(j in 1:ncol(results$proportions)){
            whichenv <- which(colnames(yhat) == colnames(results$proportions)[j])
            yhat[i,whichenv] <- results$proportions[,j]
            yhat.sd[i,whichenv] <- results$proportions_sd[,j]
        }
    }

    # calculate RMSE
    se <- as.numeric((y[,-V] - yhat[,-V])**2)
    mse <- mean(se)
    se.sem <- sd(se)/sqrt(length(se))
    rmse <- mse**.5
    rmse.sem <- se.sem**.5

    return(list(y=y,yhat=yhat,yhat.sd=yhat.sd,newsamples=newsamples,
            env.pairs=allenvs, train.envs=train.envs, rmse=rmse, rmse.sem=rmse.sem))
}

# Internal SourceTracker function to perform rarefaction analysis
"rarefy" <- function(x,maxdepth){
    if(is.null(maxdepth)) return(x)

    if(!is.element(class(x), c('matrix', 'data.frame','array')))
        x <- matrix(x,nrow=1)
    nr <- nrow(x)
    nc <- ncol(x)

    for(i in 1:nrow(x)){
        if(sum(x[i,]) > maxdepth){
        	prev.warn <- options()$warn
            options(warn=-1)
            s <- sample(nc, size=maxdepth, prob=x[i,], replace=T)
            options(warn=prev.warn)
            x[i,] <- hist(s,breaks=seq(.5,nc+.5,1), plot=FALSE)$counts
        }
    }
    return(x)
}

# Internal SourceTracker function to plot pie charts
"plot.sourcetracker.pie" <- function(stresult, labels,
        gridsize, env.colors, titlesize, indices, ...){
    V <- length(stresult$train.envs)
    for(i in indices){
        props <- apply(matrix(stresult$draws[,,i], ncol=V),2,mean)
        pie(props, labels=NA, col=env.colors, main=labels[i],cex.main=titlesize, ...)
    }
}

# Internal SourceTracker function to plot bar plots
"plot.sourcetracker.bar" <- function(stresult, labels    ,
        gridsize, env.colors, titlesize, indices, ...){
    V <- length(stresult$train.envs)
    # add extra space at top for title
    new.margins <- par('mar')
    new.margins[3] <- max(.5, new.margins[3] * 1.5)
    new.margins[2] <- .25 * titlesize
    par(mar=new.margins)
    for(i in indices){
        props <- apply(matrix(stresult$draws[,,i], ncol=V),2,mean)
        prop_devs <- apply(matrix(stresult$draws[,,i], ncol=V),2,sd)
        centers <- barplot(props, col=env.colors, main=labels[i],
                cex.main=titlesize, axes=FALSE, axisnames=FALSE,
                ylim=c(0,1.5), ...)
        sourcetracker.error.bars(centers, props, prop_devs)
        for(j in 1:4)  axis(j, at=c(-100,100),labels=FALSE)
    }
}

# Internal SourceTracker function to plot distribution plots
"plot.sourcetracker.dist" <- function(stresult, labels,
        gridsize, env.colors, titlesize, indices, sortmethod=c('divergence', 'multilevel')[1], ...){
    # stop conditions
    if(dim(stresult$draws)[1] < 2)
        stop('Distribution plots require more than one draw.')
    V <- length(stresult$train.envs)
    for(i in indices){
        x <- stresult$draws[,,i]
        rownames(x) <- 1:nrow(x)
        if(sortmethod=='multilevel'){
            # sort by size of column
            sortby.ix <- sort(colMeans(x),index=T, dec=T)$ix
            x <- matrix(x, ncol=V)
            ix <- sortmatrix(x[,sortby.ix])
            x <- x[ix,]
        } else {
            d <- jsdmatrix(x)
            if(sum(d) > 0){
                # ensure that all mixtures were not identical (e.g. 100% unknown)
                ix <- cmdscale(jsdmatrix(x),k=1)
                ix <- sort(ix,index=T)$ix
                x <- x[ix,]
            }
        }
        centers <- barplot(t(x), beside=FALSE, col=env.colors,
                space=-1/ncol(x), border=NA, axes=FALSE, axisnames=FALSE,
                main=labels[i],cex.main=titlesize,
                ylim=c(-.05,1.05), ...)
        bounds <- c(0, min(centers)-.5, 1, max(centers)+.5)

        lines(c(bounds[2], bounds[4]), c(bounds[1], bounds[1]), lty=1, lwd=1)
        lines(c(bounds[2], bounds[4]), c(bounds[3], bounds[3]), lty=1, lwd=1)
        lines(c(bounds[2], bounds[2]), c(bounds[1], bounds[3]), lty=1, lwd=1)
        lines(c(bounds[4], bounds[4]), c(bounds[1], bounds[3]), lty=1, lwd=1)
    }
}

# Internal SourceTracker function to plot error bars on a bar chart
"sourcetracker.error.bars" <- function(x,centers,spread,...){
    width = .01

    xlim <- range(x)
    barw <- diff(xlim) * width

    upper <- centers + spread
    lower <- centers - spread
    keepix <- which(spread > 0)
    x <- x[keepix]
    upper <- upper[keepix]
    lower <- lower[keepix]

    segments(x, upper, x, lower, lwd=1.5, ...)
    segments(x - barw, upper, x + barw, upper, lwd=1.5, ...)
    segments(x - barw, lower, x + barw, lower, lwd=1.5, ...)
}

# plot type 1: predicted fit verses actual proportion of that site (when source is included)
# plot type 2: predicted fit versus actual proportion of every other site, when other source is hidden
"plot.eval" <- function(res, plot.all=FALSE,plot.type=c(1,2)[1], filename=NULL){

    V <- length(res$train.envs)
    if(plot.type==1){
        pch <- 16
        mycolors <- c('#E41A1C', '#377EB8', '#4DAF4A', '#984EA3', '#FF7F00', '#A65628', '#F781BF', '#999999', '#8DD3C7', '#FFFFB3', '#BEBADA', '#FB8072', '#80B1D3', '#FDB462', '#B3DE69', '#FCCDE5', '#D9D9D9', '#BC80BD', '#CCEBC5')
        mycolors <- sprintf('%sBB',mycolors)
        if(!is.null(filename)) pdf(filename,width=6,height=6)
        plot(0,0,xlim=0:1,ylim=0:1,type='n',xlab='True proportion', ylab='Estimated proportion')
        abline(0,1)
        for(i in 1:(V+1)){
            # sourcetracker.error.bars(res$y[,i], res$yhat[,i], res$yhat.sd[,i],col='#33333399')
            points(res$y[,i],res$yhat[,i],col=mycolors[i], pch=pch)
        }
        legend('topleft',legend=c(as.character(res$train.envs),'Unknown'),col=mycolors,pch=pch, cex=.75)
        if(!is.null(filename)) dev.off()
    } else {
        if(!is.null(filename)) pdf(filename,width=10,height=10)
        par(mfrow=c(V,V),mar=rep(.5,4))
        par(mfrow=c(V,V),mar=c(4.5,4.5,.5,.5))
        for(i in 1:V){
            for(j in 1:V){
                # keep only those points where j is the hidden source and i
                # is not present
                keepix <- which(res$env.pairs[,1] != i & res$env.pairs[,2] == j)
                # don't plot same-same plots
                if(i==j) {
                    plot(0,0,xlim=c(-1,1), ylim=c(-1,1),type='n',xaxt='none',yaxt='none',xlab='', ylab='')
                    text(0,0,res$train.envs[i], cex=2)
                } else if(length(keepix) < 3){
                    plot(0,0,type='n',xaxt='none', yaxt='none',xlab='', ylab='')
                    text(0,0,'Too few samples')
                } else if (i!=j){
                    a <- res$y[keepix,V+1] # actual hidden proportion
                    b <- res$yhat[keepix,i]; # recovered proportion for this other source
                    rmse <- mean(b**2)**.5
                    plot.color <- 'black'
                    # plot red if RMSE > .1
                    if(rmse > .1) plot.color='red'
                    plot(a,b,type='p',xaxt='none',yaxt='none',
                         xlim=0:1,ylim=0:1,col=plot.color,
                         xlab='', ylab='')
                    # draw axis labels close to the axes
                    cex = .65
                    # mtext(sprintf('Prop. %s',res$train.envs[j]), side=1, line=1, cex=cex)
                    # mtext(sprintf('Prop. %s when %s absent',res$train.envs[i],res$train.envs[j]), side=2, line=1, cex=cex)
                    mtext(sprintf('Prop. Env 1',res$train.envs[j]), side=1, line=1, cex=cex)
                    mtext(sprintf('Prop. Env 2 when 1 absent',res$train.envs[i],res$train.envs[j]), side=2, line=1, cex=cex)
                    cex=1
                    text(0.05,.90,sprintf('RMSE = %.3f',rmse),adj=c(0,0),cex=cex,col=plot.color)
                }
            }
        }
        if(!is.null(filename)) dev.off()
    }
}

"save.mapping.file" <- function(results, map, filename='map.txt',
        include.contamination.predictions=FALSE){
    # also append results to a new version of the mapping file in output directory
    # add columns filled with NA to the mapping file, add values where available
    for(j in 1:ncol(results$proportions)){
        colname <- sprintf('Proportion_%s',colnames(results$proportions)[j])
        colname_sd <- sprintf('Proportion_SD_%s',colnames(results$proportions)[j])
        map[,colname] <- rep(NA,nrow(map))
        map[rownames(results$proportions), colname] <- results$proportions[,j]
        map[rownames(results$proportions), colname_sd] <- results$proportions_sd[,j]
    }

    # add this environment and its s.d. to mapping file
    thisenv <- rep(NA,nrow(results$proportions))
    thisenv_sd <- rep(NA,nrow(results$proportions))
    for(i in 1:nrow(results$proportions)){
        sampleid <- rownames(results$proportions)[i]
        whichenv <- which(results$train.envs==map[sampleid,'Env'])
        if(length(whichenv) > 0){
            thisenv[i] <-  results$proportions[i,whichenv]
            thisenv_sd[i] <- results$proportions_sd[i,whichenv]
        }
    }
    map[,'Proportion_This_Env'] <- rep(NA, nrow(map))
    map[,'Proportion_SD_This_Env'] <- rep(NA, nrow(map))
    map[rownames(results$proportions), 'Proportion_This_Env'] <- thisenv
    map[rownames(results$proportions), 'Proportion_SD_This_Env'] <- thisenv_sd

    if(include.contamination.predictions){
        for(threshold in seq(0.05, 0.95, .05)){
            header <- sprintf('contaminated_at_%.02f',threshold)
            map[,header] <- rep(NA, nrow(map))
            map[rownames(results$proportions), header] <- as.character(thisenv < threshold)
        }
    }

    # write new mapping file
    sink(filename)
    cat('#SampleID\t')
    write.table(map,sep='\t',quote=F)
    sink(NULL)
}

# sorts a matrix by first column, breaks ties by the 2nd, 3rd, etc. columns
# returns row indices
"sortmatrix" <- function(x){
    # sort by last column, then 2nd-to-last, etc.
    ix <- 1:nrow(x)
    for(j in ncol(x):1){
        ixj <- sort(x[ix,j], index=T)$ix
        ix <- ix[ixj]
    }
    return(ix)
}

"jsdmatrix" <- function(x){
    d <- matrix(0,nrow=nrow(x),ncol=nrow(x))
    for(i in 1:(nrow(x)-1)){
        for(j in (i+1):nrow(x)){
            d[i,j] <- jsd(x[i,], x[j,])
            d[j,i] <- d[i,j]
        }
    }
    return(d)
}

"jsd" <- function(p,q){
    m <- (p + q)/2
    return((kld(p,m) + kld(q,m))/2)
}

"kld" <- function(p,q){
    nonzero <- p>0 & q>0
    return(sum(p[nonzero] * log2(p[nonzero]/q[nonzero])))
}

# global definition of standard env colors
std.env.colors <- c(
'#885588',
'#CC6666',
'#47697E',
'#5B7444',
'#79BEDB',
'#663333',
'#3F52A2',
'#128244',
'#e93E4A',
'#1DBDBC',
'#A43995',
'#FFCC33',
'#B1BDCD',
'#A3C586',
'#6B78B4',
'#266A2E',
'#FCF1D1',
'#660F57',
'#272B20',
'#003366',
'#656565'
)


require(nza)
require(nzr)
require(xgboost)
require(caTools)
require(caret)
require(pROC)

nzConnectDSN('DBM_SANDBOX')

nzdata = nzQuery("SELECT HIST.*,
	   NVL(RESP.TARGET_RESPONSE, 0) AS TARGET_RESPSONSE
                 FROM (
                 SELECT DISTINCT DETAIL.LAST_NAME,
                 DETAIL.ADDRESS_LINE_1,
                 DETAIL.CITY_NAME,
                 DETAIL.STATE_CODE,
                 DETAIL.ZIP,
                 1 AS TARGET_RESPONSE
                 FROM (
                 SELECT IND.LAST_NAME,
                 ADDR.ADDRESS_LINE_1,
                 ADDR.CITY_NAME,
                 ADDR.STATE_CODE,
                 ADDR.ZIP,
                 IND.REWARDS_ACTIVE_CODE,
                 IND.VISA_STATUS,
                 DT.DATE_VALUE,
                 SH.INBOUND_INTERACTION_KEY,
                 ST.CHANNEL_GROUP,
                 PD.DEPARTMENT_MEMBER_NUMBER,
                 PD.SUB_DEPARTMENT_MEMBER_NUMBER,
                 SUM(SD.SALE_QUANTITY) AS UNITS,
                 SUM(SD.SALES_PRICE) AS SALES
                 FROM EDW_SPOKE..FACT_BPS_SALES_HEADER SH
                 JOIN EDW_SPOKE..FACT_BPS_SALES_DETAIL SD ON SH.INBOUND_INTERACTION_KEY = SD.SH_INBOUND_INTERACTION_KEY
                 JOIN EDW_SPOKE..DIM_STORE ST ON SD.STORE_MEMBER_KEY = ST.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_DATE DT ON SD.SALE_DATE_KEY = DT.DATE_KEY
                 JOIN EDW_SPOKE..DIM_PRODUCT_BPS PD ON SD.SALES_PRODUCT_MEMBER_KEY = PD.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_PERSONA PERS ON SH.CUSTOMER_MEMBER_KEY = PERS.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_INDIVIDUAL IND ON PERS.INDIVIDUAL_MEMBER_KEY = IND.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_ADDRESS ADDR ON IND.ADDRESS_KEY = ADDR.MEMBER_KEY
                 WHERE ST.CHANNEL_GROUP IN ('DIRECT','RETAIL')
                 AND DT.DATE_VALUE- 1 >= '2017-10-13'::DATE
                 AND SD.ACTUAL_SALE_FLAG = 'Y'
                 AND SD.RETURN_FLAG <> 'Y'
                 AND PD.DEPARTMENT_MEMBER_NUMBER NOT IN (875,999)
                 GROUP BY IND.LAST_NAME, ADDR.ADDRESS_LINE_1, ADDR.CITY_NAME, ADDR.STATE_CODE, ADDR.ZIP, IND.REWARDS_ACTIVE_CODE, IND.VISA_STATUS, DT.DATE_VALUE,
                 SH.INBOUND_INTERACTION_KEY, ST.CHANNEL_GROUP, PD.DEPARTMENT_MEMBER_NUMBER, PD.SUB_DEPARTMENT_MEMBER_NUMBER
                 ORDER BY IND.LAST_NAME, ADDR.ADDRESS_LINE_1, ADDR.CITY_NAME, ADDR.STATE_CODE, ADDR.ZIP
                 ) AS DETAIL
                 ORDER BY DETAIL.LAST_NAME, DETAIL.ADDRESS_LINE_1, DETAIL.CITY_NAME, DETAIL.STATE_CODE, DETAIL.ZIP
                 ) RESP
                 RIGHT JOIN (
                 SELECT DETAIL.LAST_NAME,
                 DETAIL.ADDRESS_LINE_1,
                 DETAIL.CITY_NAME,
                 DETAIL.STATE_CODE,
                 DETAIL.ZIP,
                 MAX(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-(365) AND '2017-10-13'::DATE THEN 1 ELSE 0 END) AS PURCH_6MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-(365) AND '2017-10-13'::DATE THEN DETAIL.SALES ELSE 0 END) AS SLS_6MO,
                 MAX(CASE WHEN DETAIL.REWARDS_ACTIVE_CODE IN ('A','Y') THEN 'Y' ELSE 'N' END) AS REWARDS,
                 MAX(CASE WHEN DETAIL.VISA_STATUS = 'Y' THEN 'Y' ELSE 'N' END) AS VISA,
                 MIN(DETAIL.DATE_VALUE) AS FIRST_DATE,
                 MAX(DETAIL.DATE_VALUE) AS LAST_DATE,
                 MAX(DETAIL.DATE_VALUE)-MIN(DETAIL.DATE_VALUE) AS TIME_AS_CUSTOMER,
                 COUNT(DISTINCT DETAIL.INBOUND_INTERACTION_KEY) AS ORDERS_LT,
                 (MAX(DETAIL.DATE_VALUE) - MIN(DETAIL.DATE_VALUE)) / NULLIF(COUNT(DISTINCT DETAIL.INBOUND_INTERACTION_KEY), 0)::DOUBLE AS AVG_DAYS_BTW_PURCH,
                 NVL(SUM(DETAIL.SALES)/NULLIF(COUNT(DISTINCT DETAIL.INBOUND_INTERACTION_KEY),0)::DOUBLE,0) AS AVG_TICKET_LT,
                 NVL(SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE THEN DETAIL.SALES ELSE 0 END) /
                 NULLIF(COUNT(DISTINCT CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE THEN DETAIL.INBOUND_INTERACTION_KEY ELSE NULL END),0),0) AS AVG_TICKET_24MO,
                 NVL(SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE THEN DETAIL.SALES ELSE 0 END) /
                 NULLIF(COUNT(DISTINCT CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE THEN DETAIL.INBOUND_INTERACTION_KEY ELSE NULL END),0),0) AS AVG_TICKET_48MO,
                 SUM(CASE WHEN DETAIL.CHANNEL_GROUP = 'RETAIL' THEN DETAIL.SALES ELSE 0 END)/ NULLIF(SUM(DETAIL.SALES),0) AS PERCENT_RETAIL_SALES,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE THEN DETAIL.SALES ELSE 0 END) AS SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-365 AND '2017-10-13'::DATE THEN DETAIL.SALES ELSE 0 END) AS SALES_12MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER IN (100, 151, 200, 400, 450) THEN DETAIL.SALES ELSE 0 END) AS CORE_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 100 THEN DETAIL.SALES ELSE 0 END) AS D100_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 151 THEN DETAIL.SALES ELSE 0 END) AS D151_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 175 THEN DETAIL.SALES ELSE 0 END) AS D175_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 200 THEN DETAIL.SALES ELSE 0 END) AS D200_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 300 THEN DETAIL.SALES ELSE 0 END) AS D300_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 350 THEN DETAIL.SALES ELSE 0 END) AS D350_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 400 THEN DETAIL.SALES ELSE 0 END) AS D400_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 450 THEN DETAIL.SALES ELSE 0 END) AS D450_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 475 THEN DETAIL.SALES ELSE 0 END) AS D475_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 500 THEN DETAIL.SALES ELSE 0 END) AS D500_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 600 THEN DETAIL.SALES ELSE 0 END) AS D600_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 650 THEN DETAIL.SALES ELSE 0 END) AS D650_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-729 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 675 THEN DETAIL.SALES ELSE 0 END) AS D675_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER IN (100, 151, 200, 400, 450) THEN DETAIL.SALES ELSE 0 END) AS CORE_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 100 THEN DETAIL.SALES ELSE 0 END) AS D100_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 151 THEN DETAIL.SALES ELSE 0 END) AS D151_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 175 THEN DETAIL.SALES ELSE 0 END) AS D175_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 200 THEN DETAIL.SALES ELSE 0 END) AS D200_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 300 THEN DETAIL.SALES ELSE 0 END) AS D300_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 350 THEN DETAIL.SALES ELSE 0 END) AS D350_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 400 THEN DETAIL.SALES ELSE 0 END) AS D400_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 450 THEN DETAIL.SALES ELSE 0 END) AS D450_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 475 THEN DETAIL.SALES ELSE 0 END) AS D475_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 500 THEN DETAIL.SALES ELSE 0 END) AS D500_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 600 THEN DETAIL.SALES ELSE 0 END) AS D600_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 650 THEN DETAIL.SALES ELSE 0 END) AS D650_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN '2017-10-13'::DATE-1458 AND '2017-10-13'::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 675 THEN DETAIL.SALES ELSE 0 END) AS D675_SALES_48MO
                 FROM (
                 SELECT IND.LAST_NAME,
                 ADDR.ADDRESS_LINE_1,
                 ADDR.CITY_NAME,
                 ADDR.STATE_CODE,
                 ADDR.ZIP,
                 IND.REWARDS_ACTIVE_CODE,
                 IND.VISA_STATUS,
                 DT.DATE_VALUE,
                 SH.INBOUND_INTERACTION_KEY,
                 ST.CHANNEL_GROUP,
                 PD.DEPARTMENT_MEMBER_NUMBER,
                 PD.SUB_DEPARTMENT_MEMBER_NUMBER,
                 SUM(SD.SALE_QUANTITY) AS UNITS,
                 SUM(SD.SALES_PRICE) AS SALES
                 FROM EDW_SPOKE..FACT_BPS_SALES_HEADER SH
                 JOIN EDW_SPOKE..FACT_BPS_SALES_DETAIL SD ON SH.INBOUND_INTERACTION_KEY = SD.SH_INBOUND_INTERACTION_KEY
                 JOIN EDW_SPOKE..DIM_STORE ST ON SD.STORE_MEMBER_KEY = ST.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_DATE DT ON SD.SALE_DATE_KEY = DT.DATE_KEY
                 JOIN EDW_SPOKE..DIM_PRODUCT_BPS PD ON SD.SALES_PRODUCT_MEMBER_KEY = PD.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_PERSONA PERS ON SH.CUSTOMER_MEMBER_KEY = PERS.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_INDIVIDUAL IND ON PERS.INDIVIDUAL_MEMBER_KEY = IND.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_ADDRESS ADDR ON IND.ADDRESS_KEY = ADDR.MEMBER_KEY
                 WHERE ST.CHANNEL_GROUP IN ('DIRECT','RETAIL')
                 AND DT.DATE_VALUE BETWEEN '2017-10-13'::DATE - (365*5) AND '2017-10-13'::DATE
                 AND SD.ACTUAL_SALE_FLAG = 'Y'
                 AND SD.RETURN_FLAG <> 'Y'
                 AND PD.DEPARTMENT_MEMBER_NUMBER NOT IN (875,999)
                 GROUP BY IND.LAST_NAME, ADDR.ADDRESS_LINE_1, ADDR.CITY_NAME, ADDR.STATE_CODE, ADDR.ZIP, IND.REWARDS_ACTIVE_CODE, IND.VISA_STATUS, DT.DATE_VALUE,
                 SH.INBOUND_INTERACTION_KEY, ST.CHANNEL_GROUP, PD.DEPARTMENT_MEMBER_NUMBER, PD.SUB_DEPARTMENT_MEMBER_NUMBER
                 ORDER BY IND.LAST_NAME, ADDR.ADDRESS_LINE_1, ADDR.CITY_NAME, ADDR.STATE_CODE, ADDR.ZIP, DT.DATE_VALUE, SH.INBOUND_INTERACTION_KEY,
                 PD.DEPARTMENT_MEMBER_NUMBER, PD.SUB_DEPARTMENT_MEMBER_NUMBER
                 ) AS DETAIL
                 GROUP BY DETAIL.LAST_NAME, DETAIL.ADDRESS_LINE_1, DETAIL.CITY_NAME, DETAIL.STATE_CODE, DETAIL.ZIP
                 ) HIST ON RESP.LAST_NAME = HIST.LAST_NAME
                 AND RESP.ADDRESS_LINE_1 = HIST.ADDRESS_LINE_1
                 AND RESP.CITY_NAME = HIST.CITY_NAME
                 AND RESP.STATE_CODE = HIST.STATE_CODE
                 AND RESP.ZIP = HIST.ZIP
                 WHERE HIST.AVG_TICKET_48MO > 0
                 ORDER BY RANDOM()
                 LIMIT 100000")

nzDisconnect()

modelData <- nzdata
nzdata <- NULL

#modelData <- CHURNDATA
#CHURNDATA <- NULL

modelData$LAST_NAME <- NULL
modelData$ADDRESS_LINE_1 <- NULL
modelData$CITY_NAME <- NULL
modelData$STATE_CODE <- NULL
modelData$ZIP <- NULL

summary(modelData)

modelData$PURCH_6MO <- factor(modelData$PURCH_6MO, levels=c(0,1))
modelData$SLS_6MO <- as.numeric(modelData$SLS_6MO)
modelData$REWARDS <- factor( ifelse(modelData$REWARDS=='Y',1,0), levels=c(1,0) )
modelData$VISA <- factor(ifelse(modelData$VISA=='Y', 1,0), levels=c(1,0))
modelData$FIRST_DATE <- as.Date(modelData$FIRST_DATE)
modelData$LAST_DATE <- as.Date(modelData$LAST_DATE)
modelData$ORDERS_LT <- as.numeric(modelData$ORDERS_LT)
modelData$AVG_DAYS_BTW_PURCH <- as.numeric(modelData$AVG_DAYS_BTW_PURCH)
modelData$TIME_AS_CUSTOMER <- as.numeric(modelData$TIME_AS_CUSTOMER)
modelData$AVG_TICKET_24MO <- as.numeric(modelData$AVG_TICKET_24MO)
modelData$AVG_TICKET_48MO <- as.numeric(modelData$AVG_TICKET_48MO)
modelData$AVG_TICKET_LT <- as.numeric(modelData$AVG_TICKET_LT)
modelData$PERCENT_RETAIL_SALES <- as.numeric(modelData$PERCENT_RETAIL_SALES)
modelData$SALES_24MO <- as.numeric(modelData$SALES_24MO)
modelData$CORE_SALES_24MO <- as.numeric(modelData$SALES_24MO)
modelData$CORE_SALES_48MO <- as.numeric(modelData$CORE_SALES_48MO)
modelData$SALES_12MO <- - as.numeric(modelData$SALES_12MO)

modelData$D100_SALES_24MO <- as.numeric(modelData$D100_SALES_24MO)
modelData$D151_SALES_24MO <- as.numeric(modelData$D151_SALES_24MO)
modelData$D175_SALES_24MO <- as.numeric(modelData$D175_SALES_24MO)
modelData$D200_SALES_24MO <- as.numeric(modelData$D200_SALES_24MO)
modelData$D300_SALES_24MO <- as.numeric(modelData$D300_SALES_24MO)
modelData$D350_SALES_24MO <- as.numeric(modelData$D350_SALES_24MO)
modelData$D400_SALES_24MO <- as.numeric(modelData$D400_SALES_24MO)
modelData$D450_SALES_24MO <- as.numeric(modelData$D450_SALES_24MO)
modelData$D475_SALES_24MO <- as.numeric(modelData$D475_SALES_24MO)
modelData$D500_SALES_24MO <- as.numeric(modelData$D500_SALES_24MO)
modelData$D600_SALES_24MO <- as.numeric(modelData$D600_SALES_24MO)
modelData$D650_SALES_24MO <- as.numeric(modelData$D650_SALES_24MO)
modelData$D675_SALES_24MO <- as.numeric(modelData$D675_SALES_24MO)

modelData$D100_SALES_48MO <- as.numeric(modelData$D100_SALES_48MO)
modelData$D151_SALES_48MO <- as.numeric(modelData$D151_SALES_48MO)
modelData$D175_SALES_48MO <- as.numeric(modelData$D175_SALES_48MO)
modelData$D200_SALES_48MO <- as.numeric(modelData$D200_SALES_48MO)
modelData$D300_SALES_48MO <- as.numeric(modelData$D300_SALES_48MO)
modelData$D350_SALES_48MO <- as.numeric(modelData$D350_SALES_48MO)
modelData$D400_SALES_48MO <- as.numeric(modelData$D400_SALES_48MO)
modelData$D450_SALES_48MO <- as.numeric(modelData$D450_SALES_48MO)
modelData$D475_SALES_48MO <- as.numeric(modelData$D475_SALES_48MO)
modelData$D500_SALES_48MO <- as.numeric(modelData$D500_SALES_48MO)
modelData$D600_SALES_48MO <- as.numeric(modelData$D600_SALES_48MO)
modelData$D650_SALES_48MO <- as.numeric(modelData$D650_SALES_48MO)
modelData$D675_SALES_48MO <- as.numeric(modelData$D675_SALES_48MO)
modelData$TARGET_RESPSONSE <- factor(modelData$TARGET_RESPSONSE, levels=c(0,1))

summary(modelData)

table(modelData$TARGET_RESPSONSE)
table(modelData$PURCH_6MO)

noMissingTarget <- modelData[!is.na(modelData$TARGET_RESPSONSE), ]
noMissingTarget <- NULL

columnList <- as.list(colnames(noMissingTarget))

vars <- c('PURCH_6MO','REWARDS', 'VISA', 'TIME_AS_CUSTOMER',
           'ORDERS_LT', 'AVG_DAYS_BTW_PURCH', 'AVG_TICKET_LT', 'AVG_TICKET_24MO',
           'PERCENT_RETAIL_SALES', 'SALES_12MO', 'CORE_SALES_24MO',
          'D100_SALES_24MO', 'D151_SALES_24MO', 'D175_SALES_24MO', 'D200_SALES_24MO', 'D300_SALES_24MO',
          'D350_SALES_24MO', 'D400_SALES_24MO', 'D450_SALES_24MO', 'D475_SALES_24MO', 'D500_SALES_24MO',
          'D600_SALES_24MO', 'D650_SALES_24MO', 'D675_SALES_24MO',  'TARGET_RESPSONSE')

vars.x.num <- c('ORDERS_LT', 'AVG_DAYS_BTW_PURCH', 'AVG_TICKET_LT', 'AVG_TICKET_24MO', 'TIME_AS_CUSTOMER',
                'PERCENT_RETAIL_SALES', 'SALES_12MO', 'CORE_SALES_24MO',
                'D100_SALES_24MO', 'D151_SALES_24MO', 'D175_SALES_24MO', 'D200_SALES_24MO', 'D300_SALES_24MO',
                'D350_SALES_24MO', 'D400_SALES_24MO', 'D450_SALES_24MO', 'D475_SALES_24MO', 'D500_SALES_24MO',
                'D600_SALES_24MO', 'D650_SALES_24MO', 'D675_SALES_24MO')
vars.x.nonum <- c('PURCH_6MO','REWARDS','VISA')

vars.y <- c('TARGET_RESPSONSE')

modelData2 <- modelData[ , vars]

var_diagnostics <- function(data, xvars1, xvars2){
  for(i in xvars1){
    print(paste(i, 'contains', sum(is.na(data[,i])), 'empty values.'))
  } #end inner for loop

  for(j in xvars2){
    print(j)
    print(table(data[,j]))
  }

  for(p in xvars1){
    for(q in xvars1){
      cc <- cor(data[,p], data[,q], use='pairwise.complete.obs', method=c('pearson'))
      if(cc > .8 & p != q){
        print(paste(p, ' - ', q, ' = ', cc))
      }

    }
  }

} #end function

var_diagnostics(modelData2, vars.x.num, vars.x.nonum)

scaled = scale(modelData2[, vars.x.num], center = TRUE)

mahal <- mahalanobis(scaled[ , vars.x.num],
                     colMeans(scaled[ , vars.x.num], na.rm = TRUE),
                     cov(scaled[ , vars.x.num],
                         use = "pairwise.complete.obs"),
                     tol=1e-20)

cutoff <- qchisq(1-.001, length(vars.x.num))

summary(mahal < cutoff)

modelData3 <- cbind(as.data.frame(scaled),
                    as.data.frame(modelData2[, vars.x.nonum]))
modelData3$TARGET_RESPONSE <- modelData2[, vars.y]

nooutliers <- as.data.frame(subset(modelData3, mahal < cutoff))

train <- nooutliers[1:50000,]
test <- nooutliers[50001:93667,]

form <- 'TARGET_RESPONSE ~ ORDERS_LT + AVG_DAYS_BTW_PURCH + AVG_TICKET_LT +
         TIME_AS_CUSTOMER + SALES_12MO +
        CORE_SALES_24MO + D200_SALES_24MO + D350_SALES_24MO + D400_SALES_24MO +
        D450_SALES_24MO + PURCH_6MO + REWARDS + VISA'
logmod <- glm(formula = form,
              family = binomial(link='logit'),
              data=train)

save(logmod, file = "C:\\Users\\pairwin\\Documents\\GitHub\\RAnalyses\\CHURN_LOGISTIC.rda")

summary(logmod)
anova(logmod, test='Chisq')
#plot(logmod)

target <-c('TARGET_RESPONSE')
finalvars <- c('ORDERS_LT','AVG_DAYS_BTW_PURCH','AVG_TICKET_LT', 'TIME_AS_CUSTOMER',
'SALES_12MO','CORE_SALES_24MO', 'D200_SALES_24MO', 'D350_SALES_24MO',
'D400_SALES_24MO','D450_SALES_24MO','PURCH_6MO','REWARDS', 'VISA')

train$PURCH_6MO <- as.numeric(as.character(train$PURCH_6MO))
train$REWARDS <- as.numeric(as.character(train$REWARDS))
train$VISA <- as.numeric(as.character(train$VISA))
train$TARGET_RESPONSE <- as.numeric(as.character(train$TARGET_RESPONSE))

test$PURCH_6MO <- as.numeric(as.character(test$PURCH_6MO))
test$REWARDS <- as.numeric(as.character(test$REWARDS))
test$VISA <- as.numeric(as.character(test$VISA))
test$TARGET_RESPONSE <- as.numeric(as.character(test$TARGET_RESPONSE))

train.data <- as.matrix(train[,finalvars])
train.label <- as.array(train[,target])

test.data <- as.matrix(test[,finalvars])
test.label <- as.array(test[,target])

#----------------Advanced features --------------
# to use advanced features, we need to put data in xgb.DMatrix
dtrain <- xgb.DMatrix(data = train.data, label=train.label)
dtest <- xgb.DMatrix(data = test.data, label=test.label)
#---------------Using watchlist----------------
# watchlist is a list of xgb.DMatrix, each of them is tagged with name
watchlist <- list(train=dtrain, test=dtest)
# to train with watchlist, use xgb.train, which contains more advanced features
# watchlist allows us to monitor the evaluation result on all data in the list
print("Train xgboost using xgb.train with watchlist")
bst <- xgb.train(data=dtrain,
                 max_depth=length(finalvars),
                 eta=0.3,
                 nrounds=53,
                 subsample = 0.8,
                 min_child_weight = 5,
                 colsample_bytree = .5,
                 watchlist=watchlist,
                 eval_metric = 'error',
                 nthread = 8,
                 objective = "binary:logistic",
                 early_stopping_rounds = 53)
# Finally, you can check which features are the most important.
print("Most important features (look at column Gain):")
imp_matrix <- xgb.importance(feature_names = colnames(train.data), model = bst)
print(imp_matrix)

# Feature importance bar plot by gain
print("Feature importance Plot : ")
print(xgb.plot.importance(importance_matrix = imp_matrix))

label = getinfo(dtest, "label")
pred <- predict(bst, dtest)
err <- as.numeric(sum(as.integer(pred > 0.5) != label))/length(label)
print(paste("test-error=", err))

predfactor <- factor(ifelse(pred >.5, 1,0), levels=c(1,0))

confusionMatrix(predfactor, factor(test.label, levels=c(1,0)))

rc <- roc(test.label, pred, plot=TRUE, smooth=TRUE)

#-------------------save and load models-------------------------
# save model to binary local file
xgb.save(bst, 'C:\\Users\\pairwin\\Documents\\GitHub\\RAnalyses\\xgboost_churn.model')
# load binary model to R
bst2 <- xgb.load('C:\\Users\\pairwin\\Documents\\GitHub\\RAnalyses\\xgboost_churn.model')


library(httr, quietly = T)
library(graph, quietly = T)
library(jsonlite, quietly = T)
library(stringi, quietly = T)
library(infotheo, quietly = T)
options(show.error.messages = FALSE)
options(error = function() {
    err <- stri_replace_all_regex(geterrmessage(), '\n', paste0('\n', ANSI_RED))
    colorize_log(ANSI_RED, err)
    colorize_log(ANSI_RED, 'Execution halted')
    quit(save='no', status=1, runLast=FALSE)
    })

ANSI_RED <- '\033[31m'
ANSI_GREEN <- '\033[32m'
ANSI_RESET <- '\033[0m\n'

colorize_log <- function(color, string) {
    # ANSI coloring
    cat(paste0(color, string, ANSI_RESET))
}

check_request <- function(api_host, request, job_id) {
    if (http_error(request)) {
        save(request, file=paste0('/logs/job_', job_id, '_error.RData'))
        error_request <- RETRY("PUT", paste0('http://', api_host, '/api/job/', job_id), times = 5, quiet=FALSE)
        warn_for_status(error_request)
        stop_for_status(request)
    }
}

get_dataset <- function(api_host, dataset_id, job_id, sampling_factor=1) {
    url <- paste0('http://', api_host, '/api/dataset/', dataset_id, '/loadwithids')
    colorize_log(ANSI_GREEN, paste0('Load dataset from ', url))
    start_time <- Sys.time()
    df_request <- RETRY("GET", url, times = 5, quiet=FALSE)
    check_request(api_host, df_request, job_id)
    dataset_loading_time <- as.double(difftime(Sys.time(),start_time,unit="s"))
    colorize_log(ANSI_GREEN, paste('Successfully loaded dataset (size ', headers(df_request)$`x-content-length`,
                ' bytes) in', dataset_loading_time))

    df <- read.csv(text=content(df_request, 'text'), header=TRUE, check.names=FALSE)
    df <- df[sample(nrow(df), nrow(df) * sampling_factor), ]
    return(list(df,dataset_loading_time))
}

estimate_weight <- function(from_node, to_node, graph, df, continuous=FALSE) {
    disc_df <- if(continuous) discretize(df) else df
    from_node_name <- colnames(df)[strtoi(from_node)]
    to_node_name <- colnames(df)[strtoi(to_node)]

    # expression <- causal.effect(to_node_name, from_node_name, G=igraph.from.graphNEL(graph), expr=FALSE)
    # parents <- sapply(unlist(inEdges(from_node, graph)), function(x) colnames(df)[strtoi(x)])

    # Only mutual information for now
    mi <- round(mutinformation(disc_df[[from_node_name]], disc_df[[to_node_name]]), digits = 4)
    # cmi <- condinformation(disc_df[[from_node_name]], disc_df[[to_node_name]], disc_df[parents])
    return(mi)
}

store_graph_result <- function(api_host, graph, sepsets, df, job_id, independence_test, send_sepsets, meta_results, execution_time, dataset_loading_time) {
    edges <- edges(graph)
    edge_list <- list(from_node=c(), to_node=c())
    i <- 1
    for (node in names(edges)){
        for (edge in edges[[node]]){
            edge_list[['from_node']][[i]] <- as.numeric(colnames(df)[strtoi(node)])
            edge_list[['to_node']][[i]] <- as.numeric(colnames(df)[strtoi(edge)])

            weight <- estimate_weight(
                node, edge, graph, df,
                continuous=(independence_test != 'binCI' && independence_test != 'disCI')
            )
            edge_list[['weight']][[i]] <- weight

            i <- i + 1
        }
    }
    edge_list <- data.frame(edge_list)

    sepset_list <- list(from_node=c(), to_node=c(), statistic=c(), level=c())

    if (send_sepsets > 0) {
        ss_nodes_list <- list()
        i <- 1
        for(from_node in 1:length(sepsets)){
            for(to_node in 1:length(sepsets)){
                sepset <- sapply(sepsets[[from_node]][[to_node]], (function (x) colnames(df)[x]))
                if(length(sepset) > 0){
                    sepset_list[['from_node']][[i]] <- as.numeric(colnames(df)[from_node])
                    sepset_list[['to_node']][[i]] <- as.numeric(colnames(df)[to_node])
                    ss_nodes_list[[i]] <- if(length(sepset) > 1) sepset else list(sepset)
                    sepset_list[['statistic']][[i]] <- 0
                    sepset_list[['level']][[i]] <- length(sepset)
                    i <- i + 1
                }
            }
        }
    }
    sepset_list <- data.frame(sepset_list)

    if (send_sepsets > 0) {
        sepset_list$nodes <- ss_nodes_list
    }

    result_json <- jsonlite::toJSON(list(
        job_id=strtoi(job_id),
        edge_list=if(nrow(edge_list) == 0) list() else edge_list,
        meta_results=meta_results,
        sepset_list=if(nrow(sepset_list) == 0) list() else sepset_list,
        execution_time =execution_time,
        dataset_loading_time=dataset_loading_time
    ), auto_unbox=TRUE)

    graph_request <- RETRY("POST", paste0('http://', api_host, '/api/job/', job_id, '/result'),
                                 body=result_json,
                                 add_headers("Content-Type" = "application/json"), times = 5, quiet=FALSE)
    check_request(api_host, graph_request, job_id)
    colorize_log(ANSI_GREEN, paste0('Successfully executed job ', job_id))
    return(graph_request)
}

store_graph_result_bn <- function(api_host, bn_result, df, job_id, independence_test, meta_results, execution_time, dataset_loading_time) {
    edge_list <- list(from_node=c(), to_node=c())

    for (i in 1:(length(bn_result$'arcs')/2) ){
        edge_list[['from_node']][[i]] <- as.numeric(bn_result$'arcs'[i])
        edge_list[['to_node']][[i]] <- as.numeric(bn_result$'arcs'[i+(length(bn_result$'arcs')/2)])
        edge_list[['weight']][[i]] <- 0
    }
    edge_list <- data.frame(edge_list)

    result_json <- jsonlite::toJSON(list(
        job_id=strtoi(job_id),
        edge_list=if(nrow(edge_list) == 0) list() else edge_list,
        meta_results=meta_results,
        sepset_list=list(),
        execution_time =execution_time,
        dataset_loading_time=dataset_loading_time
    ), auto_unbox=TRUE)

    graph_request <- RETRY("POST", paste0('http://', api_host, '/api/job/', job_id, '/result'),
                                 body=result_json,
                                 add_headers("Content-Type" = "application/json"), times = 5, quiet=FALSE)
    check_request(api_host, graph_request, job_id)
    colorize_log(ANSI_GREEN, paste0('Successfully executed job ', job_id))
    return(graph_request)
}

store_graph_result_bnlearn_hc <- function(api_host, bn_result, df, job_id, meta_results, execution_time, dataset_loading_time){
    edge_list <- list(from_node=c(), to_node=c())

    for (i in 1:(length(bn_result$'arcs')/2) ){
        edge_list[['from_node']][[i]] <- as.numeric(bn_result$'arcs'[i])
        edge_list[['to_node']][[i]] <- as.numeric(bn_result$'arcs'[i+(length(bn_result$'arcs')/2)])
        edge_list[['weight']][[i]] <- 0
    }
    edge_list <- data.frame(edge_list)

    result_json <- jsonlite::toJSON(list(
        job_id=strtoi(job_id),
        edge_list=if(nrow(edge_list) == 0) list() else edge_list,
        meta_results=meta_results,
        sepset_list=list(),
        execution_time =execution_time,
        dataset_loading_time=dataset_loading_time
    ), auto_unbox=TRUE)

    graph_request <- RETRY("POST", paste0('http://', api_host, '/api/job/', job_id, '/result'),
                                 body=result_json,
                                 add_headers("Content-Type" = "application/json"), times = 5, quiet=FALSE)
    check_request(api_host, graph_request, job_id)
    colorize_log(ANSI_GREEN, paste0('Successfully executed job ', job_id))
    return(graph_request)
}


 figure.observed.size.freq = function(regions= c("cfanorth", "cfasouth", "cfa4x"), years=NULL, outdir=NULL ) {

    odb = observer.db( DS="odb")
    # Remove CW's outside norms
    ii = which( odb$cw > 50 & odb$cw < 170)
    odb = odb[ii,]

    if (years=="all") years = sort( unique( odb$fishyr ) )
#years=2014
    for (reg in regions) {
      r = polygon_inside(x=odb, region=aegis.polygons::polygon_internal_code(reg), planar=F)
      for (y in years) {
        # remove production (pre-sorted) samples in historical data
        #if(y < 2004) odb = odb[which(odb$prodcd_id=="0"),]
        out= NULL
        i = which( odb$fishyr==y )  # use fishing year and not year of the actual sample
        j = intersect (r, i)
        breaks = seq(44, 184, by=4)
        soft = which(odb$durometer<68)
        cc1 = intersect( which( odb$shell==1 ), j )
        cc2 = intersect( which( odb$shell==2 ), j )
        cc3 = intersect( which( odb$shell==3 ), j )
        cc4 = intersect( which( odb$shell==4 ), j )
        cc5 = intersect( which( odb$shell==5 ), j )

        histcc1= hist(odb$cw[cc1], breaks=breaks, freq=F, plot=F)
        histcc2= hist(odb$cw[cc2], breaks=breaks, freq=F, plot=F)
        histcc3= hist(odb$cw[cc3], breaks=breaks, freq=F, plot=F)
        histcc4= hist(odb$cw[cc4], breaks=breaks, freq=F, plot=F)
        histcc5= hist(odb$cw[cc5], breaks=breaks, freq=F, plot=F)

        totsum = sum(histcc1$counts + histcc2$counts + histcc3$counts + histcc4$counts + histcc5$counts )

        histcc1$density = histcc1$counts/totsum * 100
        histcc2$density = histcc2$counts/totsum * 100
        histcc3$density = histcc3$counts/totsum * 100
        histcc4$density = histcc4$counts/totsum * 100
        histcc5$density = histcc5$counts/totsum * 100

        toplot= rbind(histcc1$density, histcc2$density, histcc3$density, histcc4$density, histcc5$density)
        rowcounts = rbind(histcc1$counts, histcc2$counts, histcc3$counts, histcc4$counts, histcc5$counts)
        ntot = sum( toplot)
        rowtot = rowSums(toplot)
        row.count.tot = rowSums( rowcounts )
        ncc = round( rowtot / ntot * 100, 1)
        legend.text = c( paste("CC5 - ",ncc[5],"%", "(", row.count.tot[5], ")", sep=""),
                         paste("CC4 - ",ncc[4],"%", "(", row.count.tot[4], ")", sep=""),
                         paste("CC3 - ",ncc[3],"%", "(", row.count.tot[3], ")", sep=""),
                         paste("CC2 - ",ncc[2],"%", "(", row.count.tot[2], ")", sep=""),
                         paste("CC1 - ",ncc[1],"%", "(", row.count.tot[1], ")", sep="") )
        # maintitle = switch(reg,
        #   cfanorth = paste("Size frequency distribution in N-ENS --", y),
        #   cfasouth = paste("Size frequency distribution in S-ENS --", y),
        #   cfa4x = paste("Size frequency distribution in 4X --", y)
        # )
        maintitle = NULL

        ylim = c(0, 18)

        dir.create( outdir, recursive=T, showWarnings=F  )
        fn = file.path( outdir, paste("size.freq", reg, y, ".pdf", sep="") )
      #  Cairo( file=fn, type="png", bg="white",  units="in", width=6, height=4, pointsize=50, dpi=300 )
        #png( file=fn,units='in', width=15,height=12,pointsize=18, res=450,type='cairo')
        pdf(file=fn, width=6, height=5, bg='white')

           barplot( toplot[c(5:1),], space=0, names.arg=breaks[-1], lwd=3,
             main=maintitle, legend.text=legend.text, xlab="Carapace width (mm)", ylab="Percentage", ylim=ylim,
             args.legend=list(x = "topright", cex=0.8) )
           abline( v=12, lty="dashed", lwd=3)
        dev.off()

        #cmd( "convert -trim -frame 10x10 -mattecolor white ", fn, fn )

        print( "-------------------------------")
        print( paste( reg, y, ": N total =", ntot) )
        print( toplot )
        print(fn)
    }  }
    return("Done")
  }


suppressPackageStartupMessages({
    source("../../../.Rprofile", chdir = TRUE)
    library(Seurat)
    library(Signac)
    library(EnsDb.Hsapiens.v86)
    library(dplyr)
    library(ggplot2)
})

# the 10x hdf5 file contains both data types.
inputdata.10x <- Read10X_h5("pbmc_granulocyte_sorted_10k_filtered_feature_bc_matrix.h5")

# extract RNA and ATAC data
rna_counts <- inputdata.10x$`Gene Expression`
atac_counts <- inputdata.10x$Peaks

# Create Seurat object
pbmc <- CreateSeuratObject(counts = rna_counts)
pbmc[["percent.mt"]] <- PercentageFeatureSet(pbmc, pattern = "^MT-")

# Now add in the ATAC-seq data
# we'll only use peaks in standard chromosomes
grange.counts <- StringToGRanges(rownames(atac_counts), sep = c(":", "-"))
grange.use <- seqnames(grange.counts) %in% standardChromosomes(grange.counts)
atac_counts <- atac_counts[as.vector(grange.use), ]
annotations <- GetGRangesFromEnsDb(ensdb = EnsDb.Hsapiens.v86)
seqlevelsStyle(annotations) <- 'UCSC'
genome(annotations) <- "hg38"

frag.file <- "pbmc_granulocyte_sorted_10k_atac_fragments.tsv.gz"
chrom_assay <- CreateChromatinAssay(
   counts = atac_counts,
   sep = c(":", "-"),
   genome = 'hg38',
   fragments = frag.file,
   min.cells = 10,
   annotation = annotations
 )
pbmc[["ATAC"]] <- chrom_assay

VlnPlot(pbmc, features = c("nCount_ATAC", "nCount_RNA","percent.mt"), ncol = 3,
  log = TRUE, pt.size = 0) + NoLegend()

pbmc <- subset(
  x = pbmc,
  subset = nCount_ATAC < 7e4 &
    nCount_ATAC > 5e3 &
    nCount_RNA < 25000 &
    nCount_RNA > 1000 &
    percent.mt < 20
)

# RNA analysis
DefaultAssay(pbmc) <- "RNA"
pbmc <- SCTransform(pbmc, verbose = FALSE) %>% RunPCA() %>% RunUMAP(dims = 1:50, reduction.name = 'umap.rna', reduction.key = 'rnaUMAP_')

# ATAC analysis
# We exclude the first dimension as this is typically correlated with sequencing depth
DefaultAssay(pbmc) <- "ATAC"
pbmc <- RunTFIDF(pbmc)
pbmc <- FindTopFeatures(pbmc, min.cutoff = 'q0')
pbmc <- RunSVD(pbmc)
pbmc <- RunUMAP(pbmc, reduction = 'lsi', dims = 2:50, reduction.name = "umap.atac", reduction.key = "atacUMAP_")

pbmc <- FindMultiModalNeighbors(pbmc, reduction.list = list("pca", "lsi"), dims.list = list(1:50, 2:50))
pbmc <- RunUMAP(pbmc, nn.name = "weighted.nn", reduction.name = "wnn.umap", reduction.key = "wnnUMAP_")
pbmc <- FindClusters(pbmc, graph.name = "wsnn", algorithm = 3, verbose = FALSE)

# perform sub-clustering on cluster 6 to find additional structure
pbmc <- FindSubCluster(pbmc, cluster = 6, graph.name = "wsnn", algorithm = 3)
Idents(pbmc) <- "sub.cluster"

# add annotations
pbmc <- RenameIdents(pbmc, '19' = 'pDC','20' = 'HSPC','15' = 'cDC')
pbmc <- RenameIdents(pbmc, '0' = 'CD14 Mono', '9' ='CD14 Mono', '5' = 'CD16 Mono')
pbmc <- RenameIdents(pbmc, '17' = 'Naive B', '11' = 'Intermediate B', '10' = 'Memory B', '21' = 'Plasma')
pbmc <- RenameIdents(pbmc, '7' = 'NK')
pbmc <- RenameIdents(pbmc, '4' = 'CD4 TCM', '13'= "CD4 TEM", '3' = "CD4 TCM", '16' ="Treg", '1' ="CD4 Naive", '14' = "CD4 Naive")
pbmc <- RenameIdents(pbmc, '2' = 'CD8 Naive', '8'= "CD8 Naive", '12' = 'CD8 TEM_1', '6_0' = 'CD8 TEM_2', '6_1' ='CD8 TEM_2')
pbmc <- RenameIdents(pbmc, '18' = 'MAIT')
pbmc <- RenameIdents(pbmc, '6_2' ='gdT', '6_3' = 'gdT')
pbmc$celltype <- Idents(pbmc)

write.csv(pbmc@meta.data, "wnn_meta_data.csv", quote=FALSE)

p1 <- DimPlot(pbmc, reduction = "umap.rna", group.by = "celltype", label = TRUE, label.size = 2.5, repel = TRUE) + ggtitle("RNA")
p2 <- DimPlot(pbmc, reduction = "umap.atac", group.by = "celltype", label = TRUE, label.size = 2.5, repel = TRUE) + ggtitle("ATAC")
p3 <- DimPlot(pbmc, reduction = "wnn.umap", group.by = "celltype", label = TRUE, label.size = 2.5, repel = TRUE) + ggtitle("WNN")
p1 + p2 + p3 & NoLegend() & theme(plot.title = element_text(hjust = 0.5))

## to make the visualization easier, subset T cell clusters
celltype.names <- levels(pbmc)
tcell.names <- grep("CD4|CD8|Treg", celltype.names,value = TRUE)
tcells <- subset(pbmc, idents = tcell.names)
CoveragePlot(tcells, region = 'CD8A', features = 'CD8A', assay = 'ATAC', expression.assay = 'SCT', peaks = FALSE)


library(sf)
library(stars)
library(glue)
library(tictoc)
library(devtools)

ville <- "La Rochelle"
localdata <- "~/files/la rochelle/r5"
DVFdata <- "~/files/DVFdata"

# choix d'un petit secteur de test. Core plutôt que commuting
FUA <- st_read("{DVFdata}/sources/FUA/FRA core/FRA_core.shp" |> glue(),
               stringsAsFactors=FALSE) |> st_transform(3035)

fua <- FUA |> dplyr::filter(fuaname == ville) |> dplyr::pull(geometry)

# choix de la résolution
resol <- 200

jour_du_transit <- as.POSIXct("2022-03-09 08:00:00 UTC")

jMem <- "8G"
options(java.parameters = '-Xmx12G')

#moteur_r5 <- r5r::setup_r5(localdata, verbose = TRUE, overwrite = FALSE)

larochelle_c200 <- qs::qread("~/files/la rochelle/la_rochelle_c200.rda", nthreads = 4)
larochelle_emp <- qs::qread(file = "~/files/la rochelle/emp_LaRochelle.rda", nthreads = 4)

# Pour l'instant, test sur core et non commuting
larochelle_emp <- larochelle_emp[fua]
st_dimensions(larochelle_c200) <- st_dimensions(larochelle_emp)

opportunites_stars <- c(larochelle_emp, dplyr::select(larochelle_c200, Ind))
rm(larochelle_c200, larochelle_emp)

opportunites <- st_as_sf(opportunites_stars, as_points = TRUE) |>
  dplyr::transmute(emplois = ifelse(is.na(emplois_total), 0, emplois_total)) |>
  dplyr::filter(emplois != 0)

rJava::.jinit()

# load_all()
moteur_r5_wparam <- routing_setup_r5(path = localdata, date = jour_du_transit, jMem = jMem, n_threads = 4)

tic()
iso_transit_r5 <- iso_accessibilite(quoi = opportunites,
                                    ou = fua |> st_zm(drop=TRUE),
                                    resolution = resol,
                                    tmax = 30,
                                    pdt = 5,
                                    dir = "r5_temp_larochelle" |> glue(),
                                    routing = moteur_r5_wparam,
                                    future=TRUE)
toc()