// This program collects gene alignments from multiple FASTA files.
// Created by Mingzhi Lin (mingzhi9@gmail.com).
package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"runtime"

	"io/ioutil"

	"github.com/alecthomas/kingpin"
	"github.com/boltdb/bolt"
	"github.com/cheggaaa/pb"
	"github.com/kussell-lab/biogo/feat/gff"
	"github.com/kussell-lab/biogo/seq"
)

// BucketName is the name of bucket.
const BucketName = "col"

var (
	// ShowProgress indicates that progress can be showed.
	ShowProgress bool
)

// Genome stores genome meta and sequence data.
type Genome struct {
	Sample string
	Seq    *seq.Sequence
}

func main() {
	app := kingpin.New("CollectGeneAlignments", "collect gene alignments.")
	app.Version("0.4.1")
	sampleFile := app.Arg("list_file", "file contains a list of genomes").String()
	gffFile := app.Arg("gff_file", "gff file").String()
	dataDir := app.Arg("data_dir", "genome sequence folder").String()
	outFile := app.Arg("out_file", "output file").String()
	bufSize := app.Flag("buf_size", "gene load buf size (memory usage), default 1000 genes").Default("1000").Int()
	numGenome := app.Flag("num_genome", "number of samples per db").Default("1000").Int()
	appendix := app.Flag("appendix", "appendix (default .fasta)").Default(".fasta").String()
	tmpDir := app.Flag("temp_dir", "temp dir").Default(".").String()
	showProgress := app.Flag("progress", "showing progress").Default("false").Bool()
	ncpu := app.Flag("num_cpu", "number of threads").Default("0").Int()
	kingpin.MustParse(app.Parse(os.Args[1:]))

	ShowProgress = *showProgress
	if *ncpu <= 0 || *ncpu > runtime.NumCPU() {
		*ncpu = runtime.NumCPU()
	}
	runtime.GOMAXPROCS(*ncpu)

	// prepare gene splitter.
	gffs := readGff(*gffFile)
	splitter := NewGeneSplitter(gffs)

	samples := readSamples(*sampleFile)

	sampleChan := make(chan string)
	go func() {
		defer close(sampleChan)
		var pbar *pb.ProgressBar
		if ShowProgress {
			pbar = pb.StartNew(len(samples))
			defer pbar.FinishPrint("Finish collecting alignments")
		}
		for _, sample := range samples {
			sampleChan <- sample
			if ShowProgress {
				pbar.Increment()
			}
		}
	}()

	jobChan := make(chan *LoadingJob)
	done := make(chan bool)
	for i := 0; i < *ncpu; i++ {
		go func() {
			for {
				dbFile, err := ioutil.TempFile(*tmpDir, fmt.Sprintf("boltdb_%d", i))
				if err != nil {
					panic(err)
				}
				job := &LoadingJob{
					dbFile:   dbFile.Name(),
					dataDir:  *dataDir,
					appendix: *appendix,
					bufSize:  *bufSize,
					splitter: splitter,
				}
				job.Load(sampleChan, *numGenome)
				jobChan <- job
				if len(job.samples) < *numGenome {
					break
				}
			}
			done <- true
		}()
	}

	go func() {
		defer close(jobChan)
		for i := 0; i < *ncpu; i++ {
			<-done
		}
	}()

	geneSet := make(map[string]bool)

	var jobs []*LoadingJob
	for job := range jobChan {
		for gene := range job.geneSet {
			geneSet[gene] = true
		}
		jobs = append(jobs, job)
		defer os.Remove(job.dbFile)
	}

	geneIDChan := make(chan string)
	go func() {
		defer close(geneIDChan)
		for id := range geneSet {
			geneIDChan <- id
		}
	}()

	w, err := os.Create(*outFile)
	if err != nil {
		log.Fatal(err)
	}
	defer w.Close()
	wb := bufio.NewWriter(w)
	defer wb.Flush()
	// pull gene alignments, and write to file.
	var pbar *pb.ProgressBar
	if ShowProgress {
		pbar = pb.StartNew(len(geneSet))
		defer pbar.FinishPrint("Finish collecting alignments")
	}
	for id := range geneIDChan {
		alignChan := pullGenes(id, jobs, *ncpu)
		numGenes := 0
		alnLength := -1
		for aln := range alignChan {
			for _, gene := range aln {
				if alnLength == -1 {
					if len(gene.Seq) > 0 {
						alnLength = len(gene.Seq)
					}
				}
				if len(gene.Seq) == alnLength && len(strings.TrimSpace(gene.Id)) > 0 {
					wb.WriteString(">" + gene.Id + "\n")
					wb.Write(gene.Seq)
					wb.WriteString("\n")
					numGenes++
				}
			}
		}
		if numGenes > 0 {
			wb.WriteString("=\n")
		}
		if ShowProgress {
			pbar.Increment()
		}
	}
}

func pullGenes(id string, jobList []*LoadingJob, ncpu int) chan []seq.Sequence {
	jobChan := make(chan *LoadingJob)
	go func() {
		defer close(jobChan)
		for _, job := range jobList {
			jobChan <- job
		}
	}()

	resChan := make(chan []seq.Sequence, ncpu)
	done := make(chan bool)
	for i := 0; i < ncpu; i++ {
		go func() {
			for job := range jobChan {
				db := job.db
				var geneIDs []string
				for j := 0; j < len(job.samples); j++ {
					geneID := fmt.Sprintf("%s_%d", id, j)
					geneIDs = append(geneIDs, geneID)
				}
				genes := getGene(db, BucketName, geneIDs)
				for j, gene := range genes {
					gene.Id = id + " " + job.samples[j]
					genes[j] = gene
				}
				resChan <- genes
			}
			done <- true
		}()
	}

	go func() {
		defer close(resChan)
		for i := 0; i < ncpu; i++ {
			<-done
		}
	}()
	return resChan
}

func splitSamples(samples []string, perDB int) (subSamples [][]string) {
	numSubSamples := len(samples)/perDB + 1
	for i := 0; i < numSubSamples; i++ {
		start := i * perDB
		end := (i + 1) * perDB
		if end > len(samples) {
			end = len(samples)
		}
		subSamples = append(subSamples, samples[start:end])
	}

	return
}

// LoadingJob load genomes into a db.
type LoadingJob struct {
	samples  []string
	db       *bolt.DB
	dbFile   string
	dataDir  string
	appendix string
	bufSize  int
	splitter Splitter
	geneSet  map[string]bool
}

// Load loads genomes into a db.
func (job *LoadingJob) Load(sampleChan chan string, numGenome int) {

	// create bolt db.
	job.db = createDB(job.dbFile)
	createBucket(job.db, BucketName)

	// geneCount stores number of genes of the same id.
	job.geneSet = make(map[string]bool)

	// splitting genes and push them into a channel.
	geneChan := make(chan *seq.Sequence)
	go func() {
		defer close(geneChan)
		for sample := range sampleChan {
			genomes := readGenome(sample, job.dataDir, job.appendix)
			for _, genome := range genomes {
				genes := job.splitter.Split(genome)
				for _, gene := range genes {
					job.geneSet[gene.Id] = true
					gene.Id = fmt.Sprintf("%s_%d", gene.Id, len(job.samples))
					geneChan <- gene
				}
			}
			job.samples = append(job.samples, sample)

			if len(job.samples) > numGenome {
				break
			}
		}
	}()

	// push genes into bolt db.
	putGene(job.db, BucketName, geneChan, job.bufSize)
	return
}

// readSamples return a list of samples from a sample file.
func readSamples(filename string) (samples []string) {
	f, err := os.Open(filename)
	if err != nil {
		log.Fatalf("Error when reading file %s:%v", filename, err)
	}
	defer f.Close()

	r := bufio.NewReader(f)
	for {
		line, err := r.ReadString('\n')

		if err != nil {
			if err != io.EOF {
				log.Fatalf("Error when reading file %s: %v", filename, err)
			}
			break
		}
		samples = append(samples, strings.TrimSpace(line))
	}
	return
}

func readGenome(sample, workDir, appendix string) (genomes []*seq.Sequence) {
	filename := filepath.Join(workDir, sample+appendix)

	f, err := os.Open(filename)
	if err != nil {
		log.Fatalf("Error when opening file %s: %v", filename, err)
	}
	defer f.Close()

	rd := seq.NewFastaReader(f)
	genomes, err = rd.ReadAll()
	if err != nil {
		log.Panic(err)
	}

	return genomes
}

// readGff return gene features.
func readGff(filename string) []*gff.Record {
	f, err := os.Open(filename)
	if err != nil {
		log.Fatalf("Error when reading file %s:%v", filename, err)
	}
	defer f.Close()

	r := gff.NewReader(f)
	records, err := r.ReadAll()
	if err != nil {
		log.Fatalf("Error when reading file %s:%v", filename, err)
	}
	return records
}

// createDB creates a bolt db.
func createDB(dbFile string) *bolt.DB {
	db, err := bolt.Open(dbFile, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}
	return db
}

// createBucket creates a bucket.
func createBucket(db *bolt.DB, bucketName string) {
	fn := func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(bucketName))
		return err
	}

	err := db.Update(fn)
	if err != nil {
		log.Fatal(err)
	}
}

func putGene(db *bolt.DB, bucketName string, geneChan chan *seq.Sequence, bufSize int) {
	buf := make([]*seq.Sequence, bufSize)
	count := 0
	for gene := range geneChan {
		if count >= len(buf) {
			loadGenes(db, bucketName, buf)
			count = 0
		}
		buf[count] = gene
		count++
	}

	loadGenes(db, bucketName, buf[:count])
}

func loadGenes(db *bolt.DB, bucketName string, genes []*seq.Sequence) {
	fn := func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		for _, gene := range genes {
			err := b.Put([]byte(gene.Id), gene.Seq)
			if err != nil {
				return err
			}
		}

		return nil
	}

	err := db.Update(fn)
	if err != nil {
		log.Fatal(err)
	}
}

func getGene(db *bolt.DB, bucketName string, ids []string) (genes []seq.Sequence) {
	fn := func(tx *bolt.Tx) error {
		for _, id := range ids {
			b := tx.Bucket([]byte(bucketName))
			v := b.Get([]byte(id))
			var gene seq.Sequence
			gene.Id = id
			gene.Seq = v
			genes = append(genes, gene)
		}

		return nil
	}

	db.View(fn)

	return
}