######################################################################
#######Series of scripts to determine adaptive introgression#########
####### Script used for the approach detailed in Appendix S2
###### Two parts to this: AFdiff with Envdif & LPdiff with Envdiff

########This script uses the following as input:
#### (1): List of outliers obtained from GEA (in this case we used bayenv)
#### (2): Minor allele frequency (MAF) by population
#### (3): env values for each population
#### (4): MAF for each loci (within hybrid zone and within each parental group)
#### (5): Raw SNP data file
#### (6): Fstat estimation file obtained from hierfstat
#### (7): Output of loci specific ancestry estimate as obtained from INTROGRESS
#### (8): List of SNPs used in INTROGRESS

#######Date: 14-Mar-2019



#####################################################################
#########FUNCTION & LIBRARIES#########################################

library(data.table)

afDiffBOOT<-function(bootDF,mafParents,mafHZ){
  
  #***************************
  #Function takes three inputs and generate the mean diff in AF for each pop from parent
  #data frame which contains columns labelled (loci,MAF &binMAF), 
  #dataframe with mafBypops and 
  #dataframe with mafByparents
  #***************************

  
  mafHZID_boot<-mafHZ[ ,colnames(mafHZ)%in%bootDF$loci]
  mafHZID_boot<-cbind(pops=mafHZ$population,mafHZID_boot)
  rownames(mafHZID_boot)<-mafHZID_boot$pops
  mafHZID_boot<-as.matrix(mafHZID_boot[ ,-1])

  
  mafParentsID_boot<-mafParents[mafParents$loci%in%bootDF$loci, ]
  mafParentT<-(t(mafParentsID_boot[ ,-1]))
  mafParentT<-apply(mafParentT,MARGIN = 2, function(Z) return(as.numeric(Z)))
  rownames(mafParentT)<-c("LP","SWWP")
  colnames(mafParentT)<-mafParentsID_boot$loci
  mafParentT<-mafParentT[ ,match(colnames(mafHZID_boot),colnames(mafParentT))]
  
  dfVar<-matrix(nrow=98,ncol=nrow(mafHZID_boot))
  for (i in 1:ncol(mafHZID_boot)){
    parental<- rep(mafParentT["LP",i],nrow(mafHZID_boot))
    dfVar[ ,i]<-parental-as.numeric(mafHZID_boot[ ,i])
  }
  
  dfVarMean<-apply(dfVar,1,function(x) median(abs(x),na.rm=TRUE))#one vector of 98 for each set of SNPs
  names(dfVarMean)<-rownames(mafHZID_boot)
  return(dfVarMean)
  
}



####################END OF FUNCTION##################################
############################################################
#LOAD DATASETS#

mafParents<-read.table("MAF_parentals.txt",sep="\t",header=T,stringsAsFactors = F)
maf98<-read.table("MAF98Pops_biallelicSNPs.txt",header=T,sep="\t")
popID<-read.table("InData/98PopIDs")
minor012<-fread("InData/SNPdataset.txt",header=T,sep="\t",data.table=F)
periphery_AF<-read.table("MAFbyPop",sep="\t",header=T)

pureSWWP_clim<-read.csv("InData/pureSWWP_Normal_1981_2010MSY.csv")
pureLP_clim<-read.csv("InData/pureLP_Normal_1981_2010MSY.csv")
pureSWWP_soil<-read.csv("InData/pureSWWP_soil1km.csv")
pureLP_soil<-read.csv("InData/pureLP_soil1km.csv")
periphery_clim<-read.table("InData/periphery_climate1981-2010.txt",header = T,sep="\t")
periphery_soil<-read.table("InData/SoilGrid1km.txt",header = T,sep="\t")

LP<-fread("../snpData/summaryStats_3489/introgress/allLoci/HomoP2_para.txt",sep="\t",data.table=F,stringsAsFactors = F) 
IntroLoci<-read.table("../snpData/summaryStats_3489/introgress/IntroLoci.txt",header=T,stringsAsFactors = F)



######################################################################################
#*****************START OF AFdiff -- Envdiff approach******************************
#*********************************************************************************
######################################################################################


#basic file manipulations 
pureSWWP<-merge(pureSWWP_clim,pureSWWP_soil,by=c("Site","Latitude","Longitude"))
pureLP<-merge(pureLP_clim,pureLP_soil,by=c("Site","Latitude","Longitude"))
periphery_env<-merge(periphery_clim,periphery_soil,by=c("Site","Latitude","Longitude"))
periphery_env<-periphery_env[periphery_env$Site%in%popID$V1, ]

periphery<-minor012[minor012$Pop%in%popID$V1, ]

pureLP<-as.matrix(pureLP[ ,-c(1:4)])
pureSWWP<-pureSWWP[ ,-c(1:4)]
periphery_env2<-periphery_env[ ,-c(1:4)]
pureLPmed<-apply(pureLP,2,function(X) return(median(X,na.rm=TRUE)))


##ESTIMATE OBSERVED difference in env between LP and the hybrid pop#########

EnvDiff<-matrix(nrow=98,ncol=88)
for (i in 1:ncol(EnvDiff)){
  lpVar<-rep(pureLPmed[i],98)
  EnvDiff[ ,i]<-abs(lpVar-periphery_env2[ ,i])
}

colnames(EnvDiff)<-colnames(periphery_env2)
EnvDiff<-data.frame(EnvDiff)
EnvDiff<-cbind(pops=periphery_env$Site,EnvDiff)

########Now using the function to first get observed differences in AF between LP and the hybrid pop for each env variable using SNPs deemed as outliers in bayenv


#Load outliers first
clim<-list.files("../snpData/summaryStats_3489/bayenvOut/clim/convergence/",full.names = T,pattern="Overlap")
Soil<-list.files("../snpData/summaryStats_3489/bayenvOut/Soil/convergence/",full.names = T,pattern="Overlap")
outliers<-c(clim,Soil)


filesBF<-vector("list",length(outliers))

for (f in 1:length(filesBF)){
  
  df<-read.table(outliers[f],header=T,sep="\t")
  df<-df[ ,1:5]
  filesBF[[f]]<-df
}

#just getting variable names here
files<-list.files("../snpData/summaryStats_3489/bayenvOut/clim/convergence/",pattern="Overlap")
files<-c(files,list.files("../snpData/summaryStats_3489/bayenvOut/Soil/convergence/",pattern="Overlap"))
varID<-sapply(strsplit(files,"Overlap1_3chainConvg"),"[",2)
varID<-sapply(strsplit(varID,".txt"),"[",1)


names(filesBF)<-varID
remove<-c("Elevation" ,"Latitude","Longitude")
BFoutliers<-filesBF[!(names(filesBF)%in%remove)]

cat("GEA outlier file is loaded")

afVAR_bayOut<-lapply(BFoutliers,function(Z) return(afDiffBOOT(Z,mafParents = mafParents,mafHZ = periphery_AF)))
afVAR_bayOut<-lapply(afVAR_bayOut,function(V) return(V[match(EnvDiff$pops,names(V))]))


#get observed correlations
corBayEst<-NULL
for (i in 1:length(afVAR_bayOut)){
  ID<-names(afVAR_bayOut)[i]
  corBayEst[i]<-cor.test(afVAR_bayOut[[i]],EnvDiff[ ,ID],method="spearman")[[4]]
}
names(corBayEst)<-names(BFoutliers)


#########END OF ESTIMATION FOR OBSERVED CORRELATION##################################
#####################################################################################



#############EMPIRICAL DISTRIBUTION FOR EACH ENV VARIABLE STARTING################

#bin by maf and FCT#############################################
binMAF<-data.frame(matrix(nrow = nrow(maf98),ncol = 1))
colnames(binMAF)<-"binMAF"

test<-cbind(maf98,binMAF)

maf<-0
b<-1

while (maf<=0.60){
  for ( i in 1:nrow(test)){
    if (test[i,"MAF"]<=maf){
      #implemnting a nested if to change bin value only if it is NA and prevent overwriting
      if (is.na(test[i,"binMAF"])){ 
        test[i,"binMAF"]<-b
      }
    }
  }
  maf<-maf+0.02 
  b<-b+1
}

#read in Fst/Fct file
LocF<-read.table("../snpData/summaryStats_3489/LocF_98hybrids.txt",header=T,sep="\t")
LocF<-cbind(loci=rownames(LocF),LocF)
rownames(LocF)<-NULL
colnames(LocF)<-c("loci","var1","var2","var3","var4","Fst","Fct")

binFst<-data.frame(matrix(nrow = nrow(LocF),ncol = 1))
colnames(binFst)<-"bin2"
binFst<-cbind(loci=LocF$loci,binFst)

f<-0
b<-1

while (f<=1){
  for ( i in 1:nrow(binFst)){
    if (LocF[i,"Fct"]<=f){
      if (is.na(binFst[i,"bin2"])){
        binFst[i,"bin2"]<-b
      }
    }
  }
  f<-f+0.02 
  b<-b+1
}

test<-merge(test,binFst,by="loci")
bin2D<-paste(test$binMAF,test$bin2,sep=":")
test<-cbind(test,bin2D)


#generate a list based on bin numbers so we can sample from each individually for each envVariable 
binsGlobal<-split(test,f = test$bin2D) 


##Large loop to move through all env Variables, do bootstrapping and get correlations 

afVAR_boot<-vector("list",length(BFoutliers))

for (L in 1:length(outliers)){
  
  ID<-names(outliers)[L]
  cat("working on variable", ID,"\n")
  
  
  Outlier<-outliers[[L]]
  targetSeq<-as.character(Outlier$loci)
  binsNULL<-test[!(test$loci%in%targetSeq),] #remove the outliers for env Var
  binsNULL<-split(binsNULL,f = binsNULL$bin2D,drop=TRUE)
  df<-Outlier
  df<-merge(test,df,by="loci")
  
  binsEnv<-data.frame(table(df$bin2D))
  binsEnv$Var1<-as.character(binsEnv$Var1)
  binsNULL_var1<-binsNULL[binsEnv$Var1] #subset to keep only bins present in the envVariable outlier set
  
  boot<-vector("list",100) #list for holding the bootstrap replicates
  empList<-vector("list",nrow(binsEnv)) #list for the bins
  
  for (r in 1:length(boot)){
    
    for (i in 1:length(binsNULL_var1)){
      name<-names(binsNULL_var1)[i]
      
      set<-binsNULL_var1[[name]]
      empList[[i]]<-set[sample(nrow(set),size = binsEnv[i,"Freq"],replace = FALSE), ] #sample based on the number of outliers in each bin
    }
    boot[[r]]<-do.call(rbind,empList)
    
  }
  
  #use function to obtain diff in AF for each boot replicate from LP AND finally estimate correlation 
  afVAR_boot[[L]]<-lapply(boot,function(Z) return(afDiffBOOT(Z,mafParents = mafParents,mafHZ = periphery_AF)))
  afVAR_boot[[L]]<-lapply(afVAR_boot[[L]],function(V) return(V[match(EnvDiff$pops,names(V))]))
  afVAR_boot[[L]]<-unlist(lapply(afVAR_boot[[L]],function(X) return(cor.test(X,EnvDiff[ ,ID],method="spearman")[[4]])))
  
}

names(afVAR_boot)<-names(BFoutliers)


#################END OF EMPIRICAL ESTIMATION###############################################



#checking if variables are outisde the 95th/25th percentile of empirical distribution
for (i in 1:length(afVAR_boot)){
    ID<-names(afVAR_boot)[i]
    emp<-afVAR_boot[[i]]
    if (quantile(emp,probs = c(0.95))>corBayEst[[ID]]){
        cat("corr for",ID,"is NOT significant and the observed corr is", corBayEst[[ID]],"\n")
      }
    }



######################################################################################
#*****************START OF LPdiff -- Envdiff approach******************************
#*********************************************************************************
######################################################################################


#some file modifications

colnames(LP)[1]<-"loci"
CONTIG<-LP$loci 
LP<-LP[ ,-1]
popID<-sapply(strsplit(colnames(LP),"_"),"[",1)
LP_tr<-data.frame(t(LP))
LP_tr<-cbind(pop=popID,LP_tr)
colnames(LP_tr)[2:ncol(LP_tr)]<-CONTIG

LPancs<-split(LP_tr,LP_tr$pop)

#The env difference between LP and hybrids was estimated as part of the AFdiff-Envdiff script, so we use that below
EnvDiff<-EnvDiff[match(names(LPancs), EnvDiff$pops), ]


##now estimate observed correlation for the set of outliers that was also loaded as a part of the AFdiff-Envdiff script above

corBayEst<-NULL 

for (i in 1:length(BFoutliers)){
  ID<-names(BFoutliers)[i]
  bay<-BFoutliers[[i]]
  
  #obtain only the outliers for one env variable at a time to do the correlations
  #take median ancestry estimate of all bayenv outliers per pop
  IDloci<-lapply(LPancs,function(df) return(df[ ,colnames(df)%in%bay$loci]))
  IDloci<-unlist(lapply(IDloci,function(df) return(median(as.matrix(df))))) 
  corBayEst[i]<-cor.test(IDloci,EnvDiff[ ,ID],method="spearman")[[4]]
}

names(corBayEst)<-names(BFoutliers)

###Now obtain empirical null using a simple bootstrapping procedure
corAdptBoot<-vector("list",length(BFoutliers)) 

for (L in 1:length(BFoutliers)){
  ID<-names(BFoutliers)[L]
  bay<-outliers[[L]]
  null<-IntroLoci[!(IntroLoci$x%in%bay$loci), ]
  
  boot<-NULL
  for (i in 1:1000){
    
    set<-sample(null,nrow(bay),replace = F)
    IDloci<-lapply(LPancs,function(df) return(df[ ,colnames(df)%in%set]))
    IDloci<-unlist(lapply(IDloci,function(df) return(median(as.matrix(df)))))
    
    boot[i]<-cor.test(IDloci,EnvDiff[ ,ID],method="spearman")[[4]]
  }
  
  corAdptBoot[[L]]<-boot
  
}

names(corAdptBoot)<-names(outliers)


##simple for loop to see which variables are outliers
for (i in 1:length(corAdptBoot)){
  var<-names(corAdptBoot)[i]
  obs<-corBayEst[names(corBayEst)==var]
  
  if (obs<quantile(corAdptBoot[[i]],probs = c(0.025))){
    cat(var, "outlier snps are adaptively introgressed AND the obs rho is", obs, "\n")
  }
}