//(请在BactPG2程序路径下运行) BactPG2.exe  seqDir(绝对路径)  makeblastdb.exe   blastp.exe   identity(ex.70)  lenthpor(ex.0.7)
//注意：代码中如有注释代码要解除
package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
)

var (
	filesnums     []int
	rawLen        int
	files         []os.FileInfo
	combineLen    int
	PGfile        *os.File
	PGfirstline   []string
	UID           int
	seqdeletemaps = make(map[string]string)
	//等待组
	wg sync.WaitGroup

	//并发数控制管道,
	//限制并发
	chanSEM = make(chan string, 3)
)

func main() {

	if len(os.Args) < 6 {
		fmt.Println("lack of files")
		return
	}
	fmt.Println("请确保blastout、db、filter1、filter2文件夹全部删除再开始运行程序。\n请确保序列文件后缀名为.fasta。不能为.faa  fastq等格式")
	PGfile, _ = os.Create("./PG.tab.txt")
	defer PGfile.Close()
	err0 := os.MkdirAll("./db", os.ModePerm) //生成BactPG2下的db文件夹
	if err0 != nil {
		fmt.Printf("Error: %s\n", err0)
	}

	files, _ = ioutil.ReadDir(os.Args[1]) //遍历BactPG下的seq文件夹下所有文件
	fmt.Println("开始生成blast数据库")
	for i, _ := range files { //makeblastdb数据库
		dbname := `db/` + strings.Replace(files[i].Name(), ".fasta", "", -1)
		PGfirstline = append(PGfirstline, strings.Replace(files[i].Name(), ".fasta", "", -1))

		seqname := os.Args[1] + `/` + files[i].Name()
		wg.Add(1)
		go func(seqnameg string, dbnameg string) {

			makeblastdb(seqnameg, dbnameg, chanSEM)
			wg.Done()
		}(seqname, dbname)
	}
	wg.Wait()

	PGfile.WriteString("UID\t")
	fmt.Println("开始组合以下菌：", PGfirstline)
	for i, _ := range PGfirstline {
		PGfile.WriteString(PGfirstline[i] + "\t")
	}
	PGfile.WriteString("\n")
	fmt.Println("开始比对blastp(耗时较久)")
	err0 = os.MkdirAll("./blastout", os.ModePerm) //生成BactPG2下的blastout文件夹（最耗时的一步）
	if err0 != nil {
		fmt.Printf("Error: %s\n", err0)
	}

	for i := 0; i < len(files); i++ {
		for j := 0; j < len(files); j++ {
			if j == i {
				continue
			}
			dbname := `db/` + strings.Replace(files[j].Name(), ".fasta", "", -1)
			seqname := os.Args[1] + `/` + files[i].Name()
			//命名：在前是queryname，在后是dbname
			blastoutname := `blastout/` + strings.Replace(files[i].Name(), ".fasta", "", -1) + `.vs.` + strings.Replace(files[j].Name(), ".fasta", "", -1) + `.out`
			wg.Add(1)
			go func() {

				blastp(dbname, seqname, blastoutname, chanSEM)
				wg.Done()
			}()

		}
	}
	wg.Wait()

	//计算每条序列长度

	lenmaps := make(map[string]int)
	for i, _ := range files {
		faaname := os.Args[1] + `/` + files[i].Name()
		faafile, openfaaErr := os.Open(faaname)
		if openfaaErr != nil {
			fmt.Println("openfastafileErr:", openfaaErr)
			return
		}
		defer faafile.Close()
		faaReader := bufio.NewReader(faafile)
		var seqname string
		for {
			fastaraw, readfastaErr := faaReader.ReadString('\n')
			fastaraw = seqTrim(fastaraw)
			if strings.Contains(fastaraw, ">") {
				seqname = strings.Split(fastaraw[1:], " ")[0]

			} else if len(fastaraw) > 0 {
				lenmaps[seqname] = lenmaps[seqname] + len(fastaraw)
			}
			if readfastaErr == io.EOF {
				break
			}
		}

	}

	//第一次过滤
	blastoutfiles, _ := ioutil.ReadDir(`./blastout`)
	//创建第一次过滤文件夹
	err0 = os.MkdirAll("./filter1", os.ModePerm) //生成BactPG2下的filter1文件夹
	if err0 != nil {
		fmt.Printf("Error: %s\n", err0)
	}
	for i, _ := range blastoutfiles {
		wg.Add(1)
		go func(a int) {
			filer1(blastoutfiles[a].Name(), lenmaps, chanSEM)
			wg.Done()
		}(i)

	}
	wg.Wait()

	//开始第二步过滤
	err0 = os.MkdirAll("./filter2", os.ModePerm) //生成BactPG2下的filter2文件夹
	if err0 != nil {
		fmt.Printf("Error: %s\n", err0)
	}
	filter1files, _ := ioutil.ReadDir(`./filter1`)
	filter1maps := make(map[string]string)
	for i, _ := range filter1files {
		_, ok := filter1maps[filter1files[i].Name()]
		if ok {
			continue
		}
		file1name := `./filter1/` + filter1files[i].Name()
		Strainname := strings.Replace(filter1files[i].Name(), ".out", "", -1)
		nameblocks := strings.Split(Strainname, ".vs.")
		file2name := `./filter1/` + nameblocks[1] + `.vs.` + nameblocks[0] + `.out`
		filter1maps[filter1files[i].Name()] = "1"
		filter1maps[nameblocks[1]+`.vs.`+nameblocks[0]+`.out`] = "1"
		filter2(file1name, file2name, nameblocks[0], nameblocks[1])
	}
	//过滤结束
	for i := 0; i < len(files); i++ {
		filesnums = append(filesnums, i)
	}
	rawLen = len(filesnums)
	for i := len(filesnums); i >= 1; i-- { //开始全组合
		UID = 0
		if i == 1 { //如果C(n,1)的情况
			for j, _ := range files {
				seqname := files[j].Name()
				deletepro2(seqname)

			}

		} else {
			combineLen = i
			combine()
		}
	}
	PGfile.Close()
	//对已生成的PG.tab.txt进一步过滤
	PG2file, _ := os.Create("./PG_" + os.Args[4] + "_" + os.Args[5] + ".txt")
	openPGfile, openPGfileErr := os.Open("./PG.tab.txt")
	if openPGfileErr != nil {
		fmt.Println("openPGfileErr:", openPGfileErr)
	}
	defer openPGfile.Close()
	PGfileReader := bufio.NewReader(openPGfile)
	PGprotmaps := make(map[string]string)
	UIDmaps := make(map[int]int)
	for i := 1; true; i++ {
		UID := 0

		PGfileRaw, readPGfileErr := PGfileReader.ReadString('\n')
		if len(PGfileRaw) > 1 {
			PGfileRaw = seqTrim(PGfileRaw)
			if i == 1 {
				PG2file.WriteString(PGfileRaw + "\n")
			} else {
				PGlineblocks := strings.Split(PGfileRaw, "\t")
				for i, _ := range PGlineblocks {
					if i == 0 || len(PGlineblocks[i]) < 1 {
						continue
					} else if PGlineblocks[i] != "-" {
						UID++
						_, ok := PGprotmaps[PGlineblocks[i]]
						if ok && strings.Split(PGprotmaps[PGlineblocks[i]], "CG")[0] != strings.Split(PGlineblocks[0], "CG")[0] {
							goto outer
						}
					}
				}
				UIDmaps[UID]++
				//确定都没有之后再输入maps

				PG2file.WriteString(strconv.Itoa(UID) + "CG" + fmt.Sprintf("%04d", UIDmaps[UID]) + "\t")
				for i, _ := range PGlineblocks {
					if i == 0 || len(PGlineblocks[i]) < 1 {
						continue
					} else if PGlineblocks[i] != "-" {
						PGprotmaps[PGlineblocks[i]] = PGlineblocks[0]
						PG2file.WriteString(PGlineblocks[i] + "\t")
					} else if PGlineblocks[i] == "-" {
						PG2file.WriteString("-\t")
					}
				}
				PG2file.WriteString("\n")
			outer:
			}

		}
		if readPGfileErr == io.EOF {
			break
		}
	}

	fmt.Println("main over")

}

func deletepro2(seqname string) { //C(n,1)的情况

	seqfile, openseqfileErr := os.Open(os.Args[1] + `/` + seqname)
	if openseqfileErr != nil {
		fmt.Println("openseqfileErr:", openseqfileErr)
	}
	seqfileReader := bufio.NewReader(seqfile)
	for {
		seqfileRaw, readseqfileErr := seqfileReader.ReadString('\n')
		seqfileRaw = seqTrim(seqfileRaw)
		if len(seqfileRaw) > 0 && strings.Contains(seqfileRaw, ">") {
			protseqname := strings.Split(seqfileRaw[1:], " ")
			_, ok := seqdeletemaps[protseqname[0]]
			if !ok {
				UID++
				UIDs := fmt.Sprintf("%04d", UID)
				//输出UID:
				PGfile.WriteString("1CG" + UIDs + "\t")
				for _, strainname := range PGfirstline {
					if strings.Contains(seqname, strainname) {
						PGfile.WriteString(protseqname[0] + "\t")
					} else {
						PGfile.WriteString("-\t")
					}
				}
				PGfile.WriteString("\n")
			}
		}

		if readseqfileErr == io.EOF {
			break
		}
	}

	seqfile.Close()

}

func combine() {
	//fmt.Println("start combine")
	arrLen := len(filesnums) - combineLen
	for i := 0; i <= arrLen; i++ {
		result := make([]int, combineLen)
		result[0] = filesnums[i]
		//fmt.Println(i, arrLen)
		doProcess(result, i, 1)
	}
}

func doProcess(result []int, rawIndex int, curIndex int) {
	var choice int = rawLen - rawIndex + curIndex - combineLen
	//fmt.Printf("Choice: %d, rawLen: %d, rawIndex : %d, curIndex : %d \r\n", choice, rawLen, rawIndex, curIndex)
	var tResult []int
	for i := 0; i < choice; i++ {
		tResult = result
		//fmt.Println(curIndex)
		tResult[curIndex] = filesnums[i+1+rawIndex]

		if curIndex+1 == combineLen {
			var referfilename string
			linkmaps := make(map[string][]string)
			strainnamemaps := make(map[string]int)
			var flag int = -1
			for j, fileindex := range tResult {
				if j == 0 { //把第一个文件设置为参照
					referfilename = strings.Replace(files[fileindex].Name(), ".fasta", "", -1)
				} else if j > 0 {
					filter2file, openfilter2Err := os.Open(`./filter2/` + referfilename + `_` + strings.Replace(files[fileindex].Name(), ".fasta", "", -1) + `.out`)
					if openfilter2Err != nil {
						fmt.Println("openfilter2Err:", openfilter2Err)
						return
					}

					filter2fileReader := bufio.NewReader(filter2file)
					for x := 1; true; x++ {
						filter2fileRaw, readfilter2fileErr := filter2fileReader.ReadString('\n')
						filter2fileRaw = seqTrim(filter2fileRaw)
						lineblocks := strings.Split(filter2fileRaw, "\t")
						if x == 1 { //将每个filter2file的首行与位置对应
							flag++
							if j == 1 {
								strainnamemaps[lineblocks[0]] = flag - 1
								strainnamemaps[lineblocks[1]] = flag
							} else {
								strainnamemaps[lineblocks[1]] = flag
							}

							continue
						}

						if len(filter2fileRaw) > 0 {

							if j == 1 { //把1_2.out所有存入maps
								linkmaps[lineblocks[0]] = append(linkmaps[lineblocks[0]], lineblocks[1]) //如果这里1_2.out首列有重复，就有问题
							} else {
								_, ok := linkmaps[lineblocks[0]]
								if ok {
									linkmaps[lineblocks[0]] = append(linkmaps[lineblocks[0]], lineblocks[1])
								} else if !ok {

								}
							}
						}
						if readfilter2fileErr == io.EOF {
							break
						}
					}

					filter2file.Close()
				}
			}

			//剔除不合格的蛋白
			for y, _ := range linkmaps {
				if len(linkmaps[y]) < len(tResult)-1 {
					delete(linkmaps, y)
				}
			}

			for y, _ := range linkmaps {
				UID++
				UIDs := fmt.Sprintf("%04d", UID)
				//输出UID:
				PGfile.WriteString(strconv.Itoa(combineLen) + "CG" + UIDs)
				PGfile.WriteString("\t")
				// for _, D := range linkmaps[y] { //注意注意，用完删除
				// 	seqdeletemaps[D] = "1"
				// }
				for _, strainname := range PGfirstline {
					_, ok := strainnamemaps[strainname]
					if !ok {
						PGfile.WriteString("-\t")
					} else if ok {
						seqdeletemaps[y] = "1"
						if strainnamemaps[strainname] == -1 {
							PGfile.WriteString(y + "\t")
						} else {
							seqdeletemaps[linkmaps[y][strainnamemaps[strainname]]] = "1"
							PGfile.WriteString(linkmaps[y][strainnamemaps[strainname]] + "\t")
						}
					}
				}
				PGfile.WriteString("\n")

			}
			//开始删除filter2文件中的蛋白
			for j, fileindex := range tResult {
				if j == 0 { //把第一个文件设置为参照
					referfilename = strings.Replace(files[fileindex].Name(), ".fasta", "", -1)
				} else if j > 0 {
					//改成并发
					wg.Add(1)
					go func(a int) {
						deletepro(`./filter2/`+referfilename+`_`+strings.Replace(files[a].Name(), ".fasta", "", -1)+`.out`, linkmaps)
						wg.Done()
					}(fileindex)
					wg.Wait()
				}
			}

			continue
		} else {
			doProcess(tResult, rawIndex+i+1, curIndex+1)
		}

	}
}

func deletepro(filter2filename string, deleteseqname map[string][]string) {

	chanSEM <- "1"
	filter2file, openfilter2fileErr := os.Open(filter2filename)
	if openfilter2fileErr != nil {
		fmt.Println("openfastaErr:", openfilter2fileErr)
	}
	filter2fileReader := bufio.NewReader(filter2file)

	//创建生成文件
	testfile, createtestErr := os.Create(filter2filename + "replace")
	if createtestErr != nil {
		fmt.Println("createtestErr:", createtestErr)
	}

	for {

		filter2raw, readfilter2Err := filter2fileReader.ReadString('\n')
		filter2raw = seqTrim(filter2raw)
		if len(filter2raw) > 0 {
			filter2lineblocks := strings.Split(filter2raw, "\t")
			_, ok := deleteseqname[filter2lineblocks[0]]
			if ok {

			} else {
				testfile.WriteString(filter2raw + "\n")

			}
		}

		if readfilter2Err == io.EOF {
			break
		}

	}
	filter2file.Close()
	testfile.Close()
	err0 := os.RemoveAll(filter2filename) //删除原来faa文件
	if err0 != nil {
		fmt.Printf("Error: %s\n", err0)
	}
	renamerr := os.Rename(filter2filename+"replace", filter2filename) //重命名:把新的faa重命名为旧的faa
	if renamerr != nil {
		fmt.Printf("file rename Error:%s", renamerr)
	}
	<-chanSEM
}

func filter2(file1name1 string, file1name2 string, nameblock1 string, nameblock2 string) {
	filter1file, openfilter1fileErr := os.Open(file1name1)
	if openfilter1fileErr != nil {
		fmt.Println("openfilter1fileErr:", openfilter1fileErr)
		return
	}
	//创建文件
	filter2file, createfilter2Err := os.Create(`./filter2/` + nameblock1 + `_` + nameblock2 + `.out`)
	if createfilter2Err != nil {
		fmt.Println("createfilter2Err:", createfilter2Err)
	}
	defer filter2file.Close()

	filter2file.WriteString(nameblock1 + "\t" + nameblock2 + "\n")
	defer filter1file.Close()
	filter1Reader := bufio.NewReader(filter1file)
	querymaps := make(map[string]string)
	for {
		filterfile1Raw, readfilterfile1Err := filter1Reader.ReadString('\n')
		if len(filterfile1Raw) > 1 {
			lineblocks := strings.Split(filterfile1Raw, "\t")
			querymaps[lineblocks[0]+"\t"+lineblocks[1]] = "1"
		}
		if readfilterfile1Err == io.EOF {
			break
		}
	}
	filter1file2, openfilter1file2Err := os.Open(file1name2)
	if openfilter1file2Err != nil {
		fmt.Println("openfilter1file2Err:", openfilter1file2Err)
		return
	}
	defer filter1file2.Close()
	filter2Reader := bufio.NewReader(filter1file2)
	for {
		filterfile2Raw, readfilterfile2Err := filter2Reader.ReadString('\n')
		if len(filterfile2Raw) > 1 {
			file2blocks := strings.Split(filterfile2Raw, "\t")
			_, ok := querymaps[file2blocks[1]+"\t"+file2blocks[0]]
			if ok {
				filter2file.WriteString(file2blocks[1] + "\t" + file2blocks[0] + "\n")
			} /*else if !ok {
				fmt.Println(file2blocks[1] + "\t" + file2blocks[0])
			}*/
		}
		if readfilterfile2Err == io.EOF {
			break
		}
	}

}

func filer1(blastoutfilename string, lenmaps map[string]int, chSem chan string) {
	//控制并发
	chSem <- "1"
	blastoutfile, openblastoutERR := os.Open(`./blastout/` + blastoutfilename)
	if openblastoutERR != nil {
		fmt.Println("openblastoutERR:", openblastoutERR)
		return
	}

	filter1file, createfilter1Err := os.Create(`./filter1/` + blastoutfilename)
	if createfilter1Err != nil {
		fmt.Println("createfilter1Err:", createfilter1Err)
	}
	defer filter1file.Close()
	defer blastoutfile.Close()

	blastoutReader := bufio.NewReader(blastoutfile)
	seqmaps := make(map[string]string)
	for {
		blastoutraw, readblastoutErr := blastoutReader.ReadString('\n')
		blastoutraw = seqTrim(blastoutraw)
		if len(blastoutraw) > 0 {
			lineblocks := strings.Split(blastoutraw, "\t")
			_, ok := seqmaps[lineblocks[0]]
			if ok { //如果if ok continue,则会只保留最好比对结果那条
				continue
			}
			seqmaps[lineblocks[0]] = "1"
			identity, _ := strconv.ParseFloat(lineblocks[2], 64)
			identityArg, _ := strconv.ParseFloat(os.Args[4], 64)
			lenth, _ := strconv.Atoi(lineblocks[3])
			lenthpor := float64(lenth) / float64(lenmaps[lineblocks[0]])
			lenthporArg, _ := strconv.ParseFloat(os.Args[5], 64)
			lenthporArg = lenthporArg / 100
			// fmt.Println(identity, lenthpor, lenth, lenmaps[lineblocks[0]])
			if identity >= identityArg && lenthpor >= lenthporArg {
				filter1file.WriteString(blastoutraw + "\n")
			}
		}
		if readblastoutErr == io.EOF {
			break
		}
	}
	<-chSem

}

func makeblastdb(seq string, dbname string, chSem chan string) {
	chSem <- "1"
	cmd := exec.Command("cmd.exe", "/C", os.Args[2], `-in`, seq, `-dbtype`, `prot`, `-out`, dbname)
	err0 := cmd.Run()
	if err0 != nil {
		fmt.Printf("Error: %s\n", err0)
	}
	<-chSem
}

func blastp(dbname string, queryname string, blastoutname string, chSem chan string) {
	chSem <- "1"
	cmd := exec.Command("cmd.exe", "/C", os.Args[3], `-db`, dbname, `-query`, queryname, `-out`, blastoutname, `-outfmt`, `6`, `-evalue`, `1e-3`)
	err0 := cmd.Run()
	if err0 != nil {
		fmt.Printf("Error: %s\n", err0)
	}
	<-chSem
}

func seqTrim(s string) string {
	s = strings.Trim(s, "\r\n")
	s = strings.Trim(s, "\n")
	return s
}
