library("aroma.affymetrix");
library("matrixStats"); # rowMedians()
verbose <- Arguments$getVerbose(-50, timestamp=TRUE);

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Setup
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
dataSet <- "GSE13372,testset";
chipType <- "GenomeWideSNP_6,Full";

csR <- AffymetrixCelSet$byName(dataSet, chipType=chipType);
print(csR);


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# AS-CRMAv2
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
res <- doASCRMAv2(csR, drop=FALSE, verbose=verbose);
print(res);


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Plot (TCN,BAF) on Chr2:75-90Mb in Array #1
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
cesN <- res$cesN;
ce <- cesN[[1]];
sampleName <- getName(ce);

# Calculate reference signals
ceR <- getAverage(cesN, verbose=verbose);
print(ceR);

# Get annotation data
cdf <- getCdf(cesN);
ugp <- getAromaUgpFile(cdf);
ufl <- getAromaUflFile(cdf);


chr <- 2;
chrTag <- sprintf("Chr%02d", chr);
units <- getUnitsOnChromosome(ugp, chromosome=2, region=c(75,90)*1e6);
pos <- getPositions(ugp, units=units) / 1e6;

# Get (TCN,BAF) signals
thetaR <- extractTotalAndFreqB(ceR, units=units)[,"total"];
data <- extractTotalAndFreqB(ce, units=units);
data[,"total"] <- 2*data[,"total"] / thetaR;
C <- data[,"total"];
B <- data[,"freqB"];

# Find genotype clusters in BAF
okB <- which(is.finite(B));
fit <- kmeans(B[okB], centers=c(0,1/2,1));

toPNG(getFullName(cesN), tags=c(sampleName, "PSCN,tracks"), {
  layout(matrix(1:2, ncol=1));
  par(mar=c(3,4,2,1)+0.1, pch=".");
  xlim <- range(pos, na.rm=TRUE);

  # Plot TCN track
  plot(NA, xlim=xlim, ylim=c(0,4), xlab="pos", ylab=expression(C));
  abline(h=0:4, lty=3, lwd=2, col="#999999");
  abline(h=median(C, na.rm=TRUE), lwd=2, col="red");
  points(pos,C, cex=3);
  lines(smooth.spline(pos,C), lwd=2, col="blue");
  abline(v=c(83.1,83.7));
  stext(side=3, pos=0, sampleName);
  stext(side=3, pos=1, chrTag);

  # Plot BAF track
  plot(NA, xlim=xlim, ylim=c(0,1), xlab="pos", ylab=expression(beta));
  abline(h=fit$centers, lwd=2, col="black");
  points(pos[okB], B[okB], cex=3, col=fit$cluster+1);
  abline(v=c(83.1,83.7));
});


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Plot (TCN,BAF) on Chr2 in Array #1
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
units <- getUnitsOnChromosome(ugp, chromosome=2);
pos <- getPositions(ugp, units=units) / 1e6;

# Get (TCN,BAF) signals
thetaR <- extractTotalAndFreqB(ceR, units=units)[,"total"];
data <- extractTotalAndFreqB(ce, units=units);
data[,"total"] <- 2*data[,"total"] / thetaR;
C <- data[,"total"];
B <- data[,"freqB"];

# Find genotype clusters in BAF
okB <- which(is.finite(B));
fit <- kmeans(B[okB], centers=c(0,1/2,1))

# Identify non-polymorphic loci
isCN <- (getUnitTypes(cdf, units=units) == 5);

toPNG(getFullName(cesN), tags=c(sampleName, "PSCN,tracks,SNPsAndNonSNPs"), {
  layout(matrix(1:3, ncol=1))
  par(mar=c(3,4,2,1)+0.1, pch=".")
  xlim <- range(pos, na.rm=TRUE);

  # Plot TCN track for SNPs
  plot(NA, xlim=xlim, ylim=c(0,4), xlab="pos", ylab=expression(C));
  abline(h=0:4, lty=3, lwd=2, col="#999999");
  abline(h=median(C, na.rm=TRUE), lwd=2, col="red");
  x <- pos[!isCN];
  y <- C[!isCN];
  points(x,y, cex=3);
  lines(smooth.spline(x,y), lwd=2, col="blue");
  abline(v=c(83.1,83.7));
  stext(side=3, pos=0, sampleName);
  stext(side=3, pos=1, chrTag);

  # Plot TCN track for non-polymorphic loci
  plot(NA, xlim=xlim, ylim=c(0,4), xlab="pos", ylab=expression(C));
  abline(h=0:4, lty=3, lwd=2, col="#999999");
  abline(h=median(C, na.rm=TRUE), lwd=2, col="red");
  x <- pos[isCN];
  y <- C[isCN];
  points(x,y, cex=3);
  lines(smooth.spline(x,y), lwd=2, col="blue");
  abline(v=c(83.1,83.7));
  stext(side=3, pos=0, sampleName);
  stext(side=3, pos=1, chrTag);

  # Plot BAF track
  plot(NA, xlim=xlim, ylim=c(0,1), xlab="pos", ylab=expression(beta));
  abline(h=fit$centers, lwd=2, col="black");
  points(pos[okB], B[okB], cex=3, col=fit$cluster+1);
  abline(v=c(83.1,83.7));
});


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Plot TCN vs fragment length stratified by SNP & non-polymorphic loci
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Identify non-polymorphic loci
isCN <- (getUnitTypes(cdf, units=units) == 5);
fl <- ufl[units,];
nbrOfEnzymes <- ncol(fl);

toPNG(getFullName(cesN), tags=c(sampleName, "TCNvsFragmentLength"), {
  layout(matrix(1:4, ncol=1));
  par(mar=c(1,1,1,1)+0.1, pch=".")
  xlim <- c(0,2000);
  Clim <- c(0,4);

  # For SNPs and non-polymorphic loci
  for (type in c("SNP", "CN")) {
    if (type == "SNP") {
      idxs <- whichVector(!isCN);
    } else {
      idxs <- whichVector(isCN);
    }

    # For each enzyme
    for (ee in 1:nbrOfEnzymes) {
      x <- fl[idxs,ee];
      y <- C[idxs];
      ok <- whichVector(is.finite(x) & is.finite(y));
      x <- x[ok];
      y <- y[ok];
      plot(x,y, cex=2, xlim=xlim, ylim=Clim);
      lines(smooth.spline(x,y), lwd=2, col="red");
    } # for (ee ...)
  } # for (kk ...)
});


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Plot density of BAF stratified by TCN
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
breaks <- quantile(C, probs=seq(0,1,by=0.1), na.rm=TRUE);
cuts <- cut(C, breaks=c(0,1.4,1.68,10));
cuts <- as.integer(cuts);
ucuts <- unique(cuts);

toPNG(getFullName(cesN), tags=c(sampleName, "BAFDensityByTCN"), {
  cols <- terrain.colors(length(ucuts));
  plot(NA, xlim=c(0,1), ylim=c(0,3), xlab="BAF", ylab="density");

  # For each TCN bin...
  for (kk in seq_along(ucuts)) {
    idxs <- whichVector(cuts == kk);
    zkk <- na.omit(B[idxs]);
    if (length(zkk) > 1) {
      printf("Number of loci: %d\n", length(zkk));
      d <- density(zkk, adjust=0.6, kernel="biweight");
      lines(d, lwd=2, col=cols[kk]);
    }
  } # for (kk ...)
});


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Analyze autosomal chromosomes
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
units <- getUnitsOnChromosomes(ugp, 1:22);
theta <- extractTheta(cesN, units=units);
thetaR <- rowMedians(theta[,1,]+theta[,2,], na.rm=TRUE);

# (TCN,BAF) and (CA,CB) across all arrays
C <- 2*(theta[,1,]+theta[,2,])/thetaR;
CA <- 2*theta[,1,]/thetaR;
CB <- 2*theta[,2,]/thetaR;
B <- CB/C;

# Average TCN & (CA,CB) per array
mu <- colMedians(C, na.rm=TRUE);
muA <- colMedians(CA, na.rm=TRUE);
muB <- colMedians(CB, na.rm=TRUE);

Clim <- c(-0.2,3.2);
xlab <- "BAF";
Clab <- "TCN";
CAlab <- "CA";
CBlab <- "CB";

toPNG(getFullName(cesN), tags=c("ASCN"), {
  layout(matrix(1:9, nrow=3, byrow=TRUE));
  par(mar=c(4,4,0.5,0.5)+0.1);
  centers <- matrix(c(0,2, 1,1, 2,0), nrow=3, ncol=2, byrow=TRUE);

  # For each array...
  for (ii in 1:ncol(C)) {
    xx <- CA[,ii];
    yy <- CB[,ii];
    X <- cbind(xx,yy);
    ok <- (is.finite(X) & -1 < X & X < 30);
    ok <- ok[,1] & ok[,2];
    X <- X[ok,];
    smoothScatter(X, xlim=Clim, ylim=Clim, xlab=CAlab, ylab=CBlab);
    abline(h=0:3, lty=3, lwd=1);
    abline(v=0:3, lty=3, lwd=1);

    # Plot centers
    fit <- kmeans(X, centers=centers);
    print(fit$centers);
    points(fit$centers, pch=19, col="red");
  } # for (ii ...)
});


toPNG(getFullName(cesN), tags=c("TCNvsBAF"), {
  layout(matrix(1:9, nrow=3, byrow=TRUE));
  par(mar=c(4,4,0.5,0.5)+0.1);
  centers <- matrix(c(0,2, 1/2,2, 1,2), nrow=3, ncol=2, byrow=TRUE);

  # For each array...
  for (ii in 1:ncol(C)) {
    xx <- B[,ii];
    yy <- C[,ii];
    X <- cbind(xx,yy);
    ok <- (is.finite(X) & -1 < X & X < 30);
    ok <- ok[,1] & ok[,2];
    X <- X[ok,];
    smoothScatter(X, xlim=c(0,1), ylim=Clim, xlab=xlab, ylab=Clab);
    abline(h=0:3, lty=3, lwd=1);
    abline(v=0:2/2, lty=3, lwd=1);

    # Plot centers
    fit <- kmeans(X, centers=centers);
    print(fit$centers);
    points(fit$centers, pch=19, col="red");
  } # for (ii ...)
});
