library(ClusterCall)
library(fitTetra)
library(viridis)

setwd("~") ### change accordingly
setwd("~/Google Drive/Manuscripts/FitTetra 2/supplemetary files")

if(!dir.exists("output")) dir.create("output")

### this assumes that Supplementary_file_2.zip was unpacked as a Supplementary_file_2 directory
source("Supplementary_file_2/helper_functions.R")

as   <- read.pop(theta.file = "Supplementary_file_2/122_2016_2845_MOESM2_ESM.csv",
                  r.file="Supplementary_file_2/122_2016_2845_MOESM3_ESM.csv")
AS   <- CC.bipop(as, parent1 = "Atlantic", parent2 = "Superior")

rgp  <- read.pop("Supplementary_file_2/122_2016_2845_MOESM8_ESM.csv", 
                r.file = "Supplementary_file_2/122_2016_2845_MOESM6_ESM.csv", error.checking = F)
RGP  <- CC.bipop(rgp, parent1 = "RioGrandeRusset", parent2 = "PremierRusset")

wl   <- read.pop("Supplementary_file_2/122_2016_2845_MOESM4_ESM.csv", 
                r.file = "Supplementary_file_2/122_2016_2845_MOESM5_ESM.csv",error.checking = F)
WL   <- CC.bipop(wl, parent1 = "Wauseon", parent2 = "Lenape")

fams <- CC.anypop(train = list(AS, RGP, WL), impute = F)

diversity <- read.pop(theta.file = "Supplementary_file_2/122_2016_2845_MOESM7_ESM.csv", 
                      error.checking = T, thresh = 2, cex = 0.5)

train <- list(AS, RGP, WL)

#prediction <- CC.anypop(train = train, predict = diversity[[1]])


### 3. clusterCall results
#We are checking the results with a checkResults helper function. We check per cross.

#### 3.1 Atlantic x Superior
resultsCC         <- NULL
ASGeno            <- fams@geno[,1:162] # Atlantic x Superior
curResults        <- checkResults(ASGeno, colnames(ASGeno)[3:162], 
                                  "Atlantic", "Superior", n.geno = nrow(ASGeno))
resultsCC         <- rbind(resultsCC,curResults$genotypeCategories)

#### 3.2 RioGrandeRusset x PremierRusset
RGPGeno           <- fams@geno[,163:326] # RioGrandeRusset x PremierRusset
curResults        <- checkResults(RGPGeno, colnames(RGPGeno)[3:164], 
                                  "RioGrandeRusset", "PremierRusset", n.geno = nrow(RGPGeno))
resultsCC         <- rbind(resultsCC,curResults$genotypeCategories)

#### 3.3  Wauseon x Lenape
WLGeno           <- fams@geno[,327:519] # Wauseon x Lenape
curResults       <- checkResults(WLGeno, colnames(WLGeno)[3:193], 
                                  "Wauseon", "Lenape", n.geno = nrow(WLGeno))
resultsCC        <- rbind(resultsCC,curResults$genotypeCategories)
colnames(resultsCC) <- c("no_model","too_many_NAs","monomorphic",
  "not_matching", "matching")
rownames(resultsCC) <- c("Atlantic x Superior", "RioGrandeRusset x PremierRusset", "Wauseon x Lenape")
save(resultsCC,file="output/resultsCC.RData")

### 4. fitTetra 2.0 analysis
#First, we transform the data back into cartesian coordinates. 
channelA <- AS@r*cos(AS@theta)*pi/2
channelA <- cbind(channelA,RGP@r*cos(RGP@theta)*pi/2)
channelA <- cbind(channelA,WL@r*cos(WL@theta)*pi/2)

channelB <- AS@r*sin(AS@theta)*pi/2
channelB <- cbind(channelB,RGP@r*sin(RGP@theta)*pi/2)
channelB <- cbind(channelB,WL@r*sin(WL@theta)*pi/2)

ratios   <- cbind(AS@theta,RGP@theta,WL@theta)

## population structure
popsT    <- c(4,5,rep(1,160),6,7,rep(2,162),8,9,rep(3,191)) 

F1par    <- matrix(c(4,5, 6,7, 8,9, NA,NA, NA,NA, NA,NA, NA,NA, NA,NA, NA,NA), nrow=9, byrow=T)
rownames(F1par)      <- c("F1_AS", "F1_RGP", "F1_WL", "P_A", "P_S", "P_R", "P_P","P_W","P_L")

# Then we run fitEtra 2.0 on the set. We use p.threshold of 0.9 as this directly corresponds to
# clusterCall min.posterior parameter. We used call.threshold = 0 and peak.threshold = 1 so
# that fitTetra returns all called SNPs. We will do quality check in the next step, outside of the 
# package.

mrkresult  <- list(NULL)
for(mrk in 1:nrow(AS@geno)){
    if(!mrk%%50) print(mrk)
    markeriT                <- data.frame(SampleName=colnames(ratios),
                                          MarkerName=rownames(AS@geno)[mrk], 
                                          ratio=ratios[mrk,],                                                                                    population=popsT)
    
    try(mrkresult[[mrk]] <- fitTetra(marker = 1, data = markeriT, pop.parents = F1par,
                                     maxiter = 50, p.threshold = 0.9, call.threshold = 0,
                                     peak.threshold = 1))
}
save(mrkresult, file=paste0("mrkresults_ft2_",format(Sys.time(), "%d-%m-%Y"),".RData"))

#We unpack the results into a matrix.
allGenoFT <- NULL
snps      <- NULL
for(i in 1:length(mrkresult)){
  try(allGenoFT <- rbind(allGenoFT,mrkresult[[i]]$scores[,13]))
  try(snps      <- c(snps,as.character(mrkresult[[i]]$scores[1,2])))
}
rownames(allGenoFT) <- snps
colnames(allGenoFT) <- mrkresult[[1]]$scores[,3]
save(allGenoFT, file=paste0("allGenoFT_ft2_",format(Sys.time(), "%d-%m-%Y"),".RData"))

### 5. fitTetra 2.0 results
# We check fitTetra 2.0 results analogous to clusterCall, per cross.

#### 5.1 Atlantic x Superior
resultsFT         <- NULL
ASGeno            <- allGenoFT[,colnames(AS@geno)]
curResults        <- checkResults(ASGeno, colnames(ASGeno)[3:162], 
                                  "Atlantic", "Superior", n.geno = nrow(ASGeno))
resultsFT <- rbind(resultsFT,curResults$genotypeCategories)

#### 5.2 RioGrandeRusset x PremierRusset
RGPGeno           <- allGenoFT[,colnames(RGP@geno)]
curResults        <- checkResults(RGPGeno, colnames(RGPGeno)[3:164], 
                                  "RioGrandeRusset", "PremierRusset", n.geno = nrow(RGPGeno))

resultsFT <- rbind(resultsFT,curResults$genotypeCategories)

#### 5.3  Wauseon x Lenape
WLGeno            <- allGenoFT[,colnames(WL@geno)] # Wauseon x Lenape
curResults        <- checkResults(WLGeno, colnames(WLGeno)[3:193], 
                                  "Wauseon", "Lenape", n.geno = nrow(WLGeno))

resultsFT <- rbind(resultsFT,curResults$genotypeCategories)
colnames(resultsFT) <- c("no_model","too_many_NAs","monomorphic",
  "not_matching", "matching")
rownames(resultsFT) <- c("Atlantic x Superior", "RioGrandeRusset x PremierRusset", "Wauseon x Lenape")
save(resultsFT, file="output/resultsFT.Rdata")

pdf("output/results_clusterCall_and_fitTetra_ClusterCall-data_plotted.pdf")
plotResults(resultsFT[1,],resultsFT[2,],resultsFT[3,], main="FitTetra 2.0")
plotResults(resultsCC[1,],resultsCC[2,],resultsCC[3,], main="ClusterCall")
dev.off()

write.csv(rbind(resultsCC,resultsFT),file="output/results_clusterCall_and_fitTetra_ClusterCall-data.csv",quote = F)