library(adegenet);library(SNPRelate);library(tidyverse)
setwd("~/Desktop/ddRAD_epiRAD_powan/demultiplexed_reads/EpiRAD/Intermediate_files3/05.Stacks/batch_effect")

# SNPRelate ----

# Functions ----

# Consistent plot aesthetics for PCA
theme.pca <- function() {
  theme_bw() +
    theme(panel.grid.minor = element_blank(),
          panel.grid.major = element_blank(),
          panel.background = element_rect(colour="black",fill="white", size=1),
          axis.text = element_text(size=16, color = "black"),
          axis.ticks = element_line(size = 0.5, colour = "black"),
          axis.ticks.length = unit(3, "mm"),
          axis.title.y = element_text(size = 30),
          axis.title.x = element_text(size = 30),
          axis.text.x = element_text(size=20),
          axis.text.y = element_text(size=20),
          legend.title = element_text(size = 20),
          legend.text = element_text(size = 20))
}


# Import vcf file and popdata file ----
vcf.fn <- "populations.snps.vcf"
snpgdsVCF2GDS(vcf.fn, "test.gds", method="biallelic.only")
snpgdsSummary("test.gds")
genofile <- snpgdsOpen("test.gds")


pop_data <- read.table("~/Desktop/ddRAD_epiRAD_powan/demultiplexed_reads/EpiRAD/Intermediate_files3/popmap_ddRAD_epiRAD.txt",header=FALSE)


# Start analysis ----
set.seed(1000)  # for reproducibility

# pca with all SNPs
pca <- snpgdsPCA(genofile, num.thread=2, autosome.only = FALSE)  # 

# variance proportion (%)
pc.percent <- pca$varprop*100
pc.percent <- head(round(pc.percent, 2))

# Manipulate results ----
tab <- data.frame(sample.id = pca$sample.id,
                  EV1 = pca$eigenvect[,1],  # the first eigenvector
                  EV2 = pca$eigenvect[,2],  # the second eigenvector
                  EV3 = pca$eigenvect[,3],  # the third eigenvector
                  EV4 = pca$eigenvect[,4],
                  EV5 = pca$eigenvect[,5],
                  stringsAsFactors = FALSE)
head(tab)


# add population data
tab[,7] <- pop_data$V2
tab[,8] <- pop_data$Library
colnames(tab)[7] <- "Lake"
colnames(tab)[8] <- "Library"

# recode lake names
lakes_renamed <- recode(tab$Lake, gla = "Glashan", tar = "Tarsan", eck = "Eck", lai = "Allt na lairige", 
                        shi = "Shira", car = "Carron", sloy = "Sloy", lom = "Lomond")
tab[,7] <- lakes_renamed

# Plot results
cols <- c("Allt na lairige"="#a4dede","Carron"="#008396","Eck"="gray60","Glashan"="#9b2915","Lomond"="gray60","Shira"="#287c71",
          "Sloy"="#6c966f","Tarsan"="#fe7c73")
tab$Lake <- factor(tab$Lake,levels = c("Glashan", "Tarsan","Eck","Allt na lairige", "Carron","Shira","Sloy","Lomond"))


# PC1 and PC2
pdf(file = "~/Dropbox/Marco_Crotti/Evolutionary genomics of whitefish/Translocation project/Genomic analyses/figures/batch_effect_12_new.pdf", width = 11.69, height = 8.27)
ggplot(tab, aes(x = EV1, y = EV2, colour = Library)) + geom_point(size=6, alpha=0.85) +
  theme.pca() + labs(x=paste("EV1",round(pc.percent, 2)[1],"%"), y=paste("EV2", round(pc.percent, 2)[2], "%")) +  
  geom_vline(xintercept = 0, linetype = 'dashed') + geom_hline(yintercept = 0, linetype = 'dashed') 
dev.off()

# PC1 and PC3
pdf(file = "~/Dropbox/Marco_Crotti/Evolutionary genomics of whitefish/Translocation project/Genomic analyses/figures/batch_effect_13_new.pdf", width = 11.69, height = 8.27)
ggplot(tab, aes(x = EV1, y = EV3,col = Library)) + geom_point(size=6, alpha=0.85) +
  theme.pca() + labs(x=paste("EV1",round(pc.percent, 2)[1],"%"), y=paste("EV3", round(pc.percent, 2)[3], "%")) +
  geom_vline(xintercept = 0, linetype = 'dashed') + geom_hline(yintercept = 0, linetype = 'dashed') 
dev.off()

# To calculate the SNP correlations between eigenvactors and SNP genotypes: ----
# Get chromosome index
chr <- read.gdsn(index.gdsn(genofile, "snp.chromosome"))
chr2 <- parse_number(chr)
CORR <- snpgdsPCACorr(pca, genofile, eig.which=1:4)

savepar <- par(mfrow=c(3,1), mai=c(0.3, 0.55, 0.1, 0.25))
for (i in 1:3)
{
  plot(abs(CORR$snpcorr[i,]), ylim=c(0,1), xlab="", ylab=paste("PC", i),
       col=chr2, pch="+")
}

locus_id <- read.table("locus_name.txt",header=TRUE)
corr_table <- data.frame(t(CORR$snpcorr))
corr_table[,5] <- locus_id$ID
outliers <- filter(corr_table, X2 >= 0.3)

write.table(outliers$V5, "blacklist_snprelate.txt", row.names = FALSE, col.names = FALSE)

# Adegenet analysis ----
genind1 <- read.structure("populations.str", n.ind = 213, n.loc = 16349, 
                          onerowperind = FALSE, col.lab = 1, 
                          NA.char = "0", ask = FALSE, 
                          row.marknames = 1, quiet = FALSE) 


X <- scaleGen(genind1, NA.method="mean")
class(X)
pca1 <- dudi.pca(X,cent=TRUE,scale=TRUE,scannf=FALSE,nf=5)
s.label(pca1$li,xax=1, yax=2)
pca1$eig


contrib1 <- loadingplot(pca1$co, axis=2,
                        thres=.3, lab.jitter=1)
contrib1_table <- data.frame(unique(gsub("_.*","", contrib1$var.names)))
colnames(contrib1_table) <- "Locus"

write.table(contrib1_table, "blacklist_adegenet.txt", row.names = FALSE, col.names = FALSE, quote = FALSE)