library(meffil)

########################################################
#### Step 1 - Methylation Data QC and normalization ####
####          according to meffil default pipeline  ####
########################################################
## load data into meffil and run QC
qc.object <- meffil.qc(samplesheet, cell.type.reference="blood gse35069", verbose=TRUE, featureset = "epic") 
## compute QC summary
qc.summary <- meffil.qc.summary(qc.object, parameters = meffil.qc.parameters(sex.outlier.sd = 6, meth.unmeth.outlier.sd = 4, detectionp.samples.threshold = 0.01, colour.code ="batch")) 
meffil.qc.report(qc.summary, output.file="qc/report2.html") 
## Remove samples that fail QC
qc.object.nobadsample <- meffil.remove.samples(qc.object, qc.summary$bad.samples$sample.name) 
## Normalization taking into account column, row and slide information
norm.objects <- meffil.normalize.quantiles(qc.object.nobadsample, number.pcs=10, random.effects = c("Slide","sentrix_row", "sentrix_col"))
## Remove probes that fail QC   
norm.beta <- meffil.normalize.samples(norm.objects, cpglist.remove=qc.summary$bad.cpgs$name) 

########################################################
#### Step 2 - Add age, sex and cell composition     ####
####          predictions to samplesheet            ####
########################################################
#### Extract sex and cell counts with meffil package
sex <- sapply(norm.objects, function(object) object$predicted.sex)
counts <- t(meffil.cell.count.estimates(norm.objects))
sexcounts.df <-data.frame(Sample_Name=rownames(counts),counts,sex)
#### Extract various age predictions with methylclock package
library(methylclock)
missCpGs <- checkClocks(norm.beta)
norm.beta.clockready<-data.frame(probe=rownames(norm.beta),norm.beta)
predictedAges<-DNAmAge(norm.beta.clockready,min.perc=.4)
predictedAges<-data.frame(predictedAges)
predictedAges$Sample_Name<-gsub("X","",predictedAges$id)
samplesheet<-merge(merge(samplesheet,sexcounts.df,by="Sample_Name"),predictedAges,by="Sample_Name")

########################################################
#### Step 3 - Compute residuals after ajustment for ####
####          predicted age, sex and cell composition ##
########################################################
residuals<-t(apply(norm.beta,1,function(probe){
    tmp<-data.frame(y=probe,samplesheet[,c("skinHorvath","predictedSex","CD8T","CD4T","NK","Bcell","Mono","Gran")])
    fit<-lm(y~.,data=tmp)
    return(residuals(fit))
}))

########################################################
#### Step 4 - Predict test dataset from train dataset ##
####          Relies on class::knn function         ####
####          Can be used either to                 ####
####            1) loop over train samples to       ####
####               compute leave one out Se and Sp  ####
####               estimates  						####
####            2) predict VUS predicted status     ####
####         RQ: best performances are guaranteed   ####
####         by multiclass knn, namely providing more ##
####         training cases covering many syndroms  ####
####         rather than just 1 syndrom vs controls ####
########################################################
library(class)
#### kNN prediction for 1 episignature
getSingleKNNpred<-function(residuals,
							### matrix of methylation residuals, cgs in rows, individuals in columns
							cgs,
							### vector of cg IDs included in the signature, as retrieved from supplementary materials
							pheno,
							### phenotype data.frame, minimum columns: Sample_Name, status ("geneABC Pathogenic", "geneABC VUS" or "Control")
							trainsamples,
							### samples with known pathogenic or negative status used to train the kNN algorithm
							testsamples,
							### samples to classify
							k=NULL,
							### number of nearest neighbours used for classification
							decision.thresh=1,
							### minimum percent of vote in favor of majority vote required for definite decision
							gene=NULL
							### target gene of interest, for visual representations
							){
	### prepare output dataframe
	out<-pheno[pheno$Sample_Name%in%testsamples,]
	out<-out[order(out$Sample_Name),]
	out$pred <- NA
	out$prob <- NA
	### separate dataset into training set and testing set, reorder
	rtrain<-residuals[rownames(residuals) %in% cgs, colnames(residuals) %in% trainsamples]
	rtrain<-rtrain[,order(colnames(rtrain))]
	ptrain<-pheno[pheno$Sample_Name %in% colnames(rtrain),]
	ptrain<-ptrain[order(ptrain$Sample_Name),]
	ntot<-sum(grepl("Pathogenic",ptrain$status))
	rtest<-residuals[rownames(residuals) %in% cgs, colnames(residuals) %in% testsamples]
	if (!is.null(dim(rtest))){rtest<-rtest[,order(colnames(rtest))]}
	### whatever the gene considered, make sure controls are used as reference
	ptrain$status<-relevel(factor(ptrain$status),ref="Control")
	### if no parameter k is given, require completely unanimous choice between neighbours
	if (is.null(k)){k<-ntot}
	### predictions using knn function of class R package
	predictions<-knn(t(rtrain),t(rtest),cl=ptrain$status,k,prob=T)
	
	ptest<-pheno[pheno$Sample_Name %in% colnames(rtest),]
	ptot<-rbind(ptrain, ptest)
	ptot$status<-relevel(factor(ptot$status),ref="Control")
	out$pred <- as.character(predictions)
	out$prob <- attr(predictions,"prob")
	out$prob[out$pred=="Control"] <- 1-out$prob[out$pred=="Control"]
	if (decision.thresh>.5 & sum(out$prob<decision.thresh)>0){
		out$pred[out$prob<decision.thresh]<-"Control"
		out$pred[out$prob<decision.thresh&out$prob>.5]<-"Uncertain"
	}
	if (!is.null(gene)){
		genesamples <- pheno$Sample_Name[grepl(gene,pheno$status)|grepl("Control",pheno$status)]
		short<-residuals[rownames(residuals)%in%cgs,colnames(residuals)%in%intersect(genesamples,c(ptrain$Sample_Name,ptest$Sample_Name))]
		ptot<-ptot[ptot$Sample_Name%in%colnames(short),]
		ptot$group<-paste(ptot$status,ptot$predictedSex)
		######### PCA plot
		fullpcs <- meffil.methylation.pcs(short,full.obj=T)
		pcs<-fullpcs$x
		eigs <- fullpcs$sdev^2
   	    perct <- eigs/sum(eigs)
		ptot<-ptot[order(ptot$Sample_Name),]
		pcs<-pcs[order(rownames(pcs)),]
		plotinfo<-data.frame(pc1=pcs[,1],pc2=pcs[,2],pc3=pcs[,3],pc4=pcs[,4],group=ptot$group,sample=ptot$Sample_Name,ID=ptot$ID)
		pdf(file=paste0("PCA_",gene,".pdf"))
			g<-ggplot(plotinfo,
				aes(x=pc1,y=pc2,col=group,shape=group,label=ID)) + geom_point() + geom_text_repel(size=4,show.legend=F) + theme_minimal()
			print(g)
		dev.off()
		######### Heatmap		
		colnames(short)<-ptot$ID	
		pdf(file=paste0("Heatmap_",gene,".pdf"))
		heatmap(shortforHeat,legend="col",cexCol=.6)
		dev.off()
	}
	return(out)
}


########################################################
#### Step 5 - Compute Se/Sp 						####
########################################################

out <-getSingleKNNpred(residuals,cgs,pheno,trainsamples,testsamples,k,decision.thresh,gene)							)
Se <- mean(out$pred[out$status==paste(gene,"Pathogenic")]==paste(gene,"Pathogenic"))
Sp <- mean(out$pred[out$status=="Control"]%in%c("Control","Uncertain"))

