// Copyright © 2016 Wei Shen <shenwei356@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"regexp"
	"runtime"
	"sort"
	"strconv"
	"strings"

	"github.com/shenwei356/bio/featio/gtf"
	"github.com/shenwei356/bio/seq"
	"github.com/shenwei356/bio/seqio/fai"
	"github.com/shenwei356/bio/seqio/fastx"
	"github.com/shenwei356/bwt"
	"github.com/shenwei356/bwt/fmi"
	"github.com/shenwei356/xopen"
	"github.com/spf13/cobra"
)

// ampliconCmd represents the amplicon command
var ampliconCmd = &cobra.Command{
	Use:   "amplicon",
	Short: "retrieve amplicon (or specific region around it) via primer(s)",
	Long: `retrieve amplicon (or specific region around it) via primer(s).

Attentions:
  1. Only one (the longest) matching location is returned for every primer pair.
  2. Mismatch is allowed, but the mismatch location (5' or 3') is not controled. 
  3. Degenerate bases/residues like "RYMM.." are also supported.
     But do not use degenerate bases/residues in regular expression, you need
     convert them to regular expression, e.g., change "N" or "X"  to ".".

Examples:
  0. no region given.
  
                    F
        -----===============-----
             F             R
        -----=====-----=====-----
             
             ===============         amplicon

  1. inner region (-r x:y).

                    F
        -----===============-----
             1 3 5                    x/y
                      -5-3-1          x/y
             F             R
        -----=====-----=====-----     x:y
        
             ===============          1:-1
             =======                  1:7
               =====                  3:7
                  =====               6:10
                  =====             -10:-6
                     =====           -7:-3
                                     -x:y (invalid)
                    
  2. flanking region (-r x:y -f)
        
                    F
        -----===============-----
         -3-1                        x/y
                            1 3 5    x/y
             F             R
        -----=====-----=====-----
        
        =====                        -5:-1
        ===                          -5:-3
                            =====     1:5
                              ===     3:5
            =================        -1:1
        =========================    -5:5
                                      x:-y (invalid)

`,
	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)
		alphabet := config.Alphabet
		idRegexp := config.IDRegexp
		lineWidth := config.LineWidth
		outFile := config.OutFile
		seq.AlphabetGuessSeqLengthThreshold = config.AlphabetGuessSeqLength
		seq.ValidateSeq = false
		gtf.Threads = config.Threads
		fai.MapWholeFile = false
		Threads = config.Threads
		runtime.GOMAXPROCS(config.Threads)
		bwt.CheckEndSymbol = false

		files := getFileListFromArgsAndFile(cmd, args, true, "infile-list", true)

		outfh, err := xopen.Wopen(outFile)
		checkError(err)
		defer outfh.Close()

		region := getFlagString(cmd, "region")
		fregion := getFlagBool(cmd, "flanking-region")

		forward0 := getFlagString(cmd, "forward")
		reverse0 := getFlagString(cmd, "reverse")
		primerFile := getFlagString(cmd, "primer-file")
		maxMismatch := getFlagNonNegativeInt(cmd, "max-mismatch")
		strict := getFlagBool(cmd, "strict-mode")
		onlyPositiveStrand := getFlagBool(cmd, "only-positive-strand")
		outFmtBED := getFlagBool(cmd, "bed")

		var list [][3]string
		var primers [][3][]byte

		if primerFile != "" {
			list, err = loadPrimers(primerFile)
			checkError(err)
		} else {
			list = [][3]string{{".", forward0, reverse0}}
		}

		primers, err = parsePrimers(list)
		checkError(err)

		if !config.Quiet {
			log.Infof("%d primer pair loaded", len(primers))
		}

		var begin, end int

		var usingRegion bool
		if region != "" {
			if !reRegion.MatchString(region) {
				checkError(fmt.Errorf(`invalid region: %s. type "seqkit amplicon -h" for more examples`, region))
			}
			r := strings.Split(region, ":")
			begin, err = strconv.Atoi(r[0])
			checkError(err)
			end, err = strconv.Atoi(r[1])
			checkError(err)

			if begin == 0 || end == 0 {
				checkError(fmt.Errorf("both begin and end in region should not be 0"))
			}

			if fregion {
				if begin > 0 && end < 0 {
					checkError(fmt.Errorf("invalid flanking region (x:-y): %d:%d", begin, end))
				}
				if begin > end {
					checkError(fmt.Errorf("invalid flanking region (x should be smaller than y): %d:%d", begin, end))
				}
			} else {
				if begin < 0 && end > 0 {
					checkError(fmt.Errorf("invalid inner region (-x:y): %d:%d", begin, end))
				}
			}
			usingRegion = true
		}

		var record *fastx.Record
		var fastxReader *fastx.Reader

		var finder *AmpliconFinder
		var loc []int

		strands := []string{"+", "-"}
		var strand string
		var tmpSeq *seq.Seq
		var primer [3][]byte

		for _, file := range files {
			fastxReader, err = fastx.NewReader(alphabet, file, idRegexp)
			checkError(err)

			for {
				record, err = fastxReader.Read()
				if err != nil {
					if err == io.EOF {
						break
					}
					checkError(err)
					break
				}
				if fastxReader.IsFastq {
					config.LineWidth = 0
					fastx.ForcelyOutputFastq = true
				}

				for _, strand = range strands {
					if strand == "-" {
						if onlyPositiveStrand {
							continue
						}
						record.Seq.RevComInplace()
					}

					for _, primer = range primers {
						finder, err = NewAmpliconFinder(record.Seq.Seq, primer[1], primer[2], maxMismatch)
						checkError(err)

						if usingRegion {
							loc, err = finder.LocateRange(begin, end, fregion, strict)
						} else {
							loc, err = finder.Locate()
						}
						checkError(err)

						if loc == nil {
							continue
						}

						if outFmtBED {
							outfh.WriteString(fmt.Sprintf("%s\t%d\t%d\t%s\t%d\t%s\t%s\n",
								record.ID,
								loc[0]-1,
								loc[1],
								primer[0],
								0,
								strand,
								record.Seq.SubSeq(loc[0], loc[1]).Seq))

							continue
						}
						tmpSeq = record.Seq

						record.Seq = record.Seq.SubSeq(loc[0], loc[1])
						record.FormatToWriter(outfh, config.LineWidth)

						record.Seq = tmpSeq
					}
				}
			}

			config.LineWidth = lineWidth
		}
	},
}

func init() {
	RootCmd.AddCommand(ampliconCmd)

	ampliconCmd.Flags().StringP("forward", "F", "", "forward primer (5'-primer-3'), degenerate bases allowed")
	ampliconCmd.Flags().StringP("reverse", "R", "", "reverse primer (5'-primer-3'), degenerate bases allowed")
	ampliconCmd.Flags().IntP("max-mismatch", "m", 0, "max mismatch when matching primers, no degenerate bases allowed")
	ampliconCmd.Flags().StringP("primer-file", "p", "", "3- or 2-column tabular primer file, with first column as primer name")

	ampliconCmd.Flags().StringP("region", "r", "", `specify region to return. type "seqkit amplicon -h" for detail`)
	ampliconCmd.Flags().BoolP("flanking-region", "f", false, "region is flanking region")
	ampliconCmd.Flags().BoolP("strict-mode", "s", false, "strict mode, i.e., discarding seqs not fully matching (shorter) given region range")
	ampliconCmd.Flags().BoolP("only-positive-strand", "P", false, "only search on positive strand")
	ampliconCmd.Flags().BoolP("bed", "", false, "output in BED6+1 format with amplicon as 7th columns")
}

func loadPrimers(file string) ([][3]string, error) {
	fh, err := os.Open(file)
	if err != nil {
		return nil, fmt.Errorf("load primers from '%s': %s", file, err)
	}

	var text string
	var items []string
	lists := make([][3]string, 0, 100)
	scanner := bufio.NewScanner(fh)
	for scanner.Scan() {
		text = strings.TrimSpace(scanner.Text())
		if text == "" || text[0] == '#' {
			continue
		}

		items = strings.Split(text, "\t")
		switch len(items) {
		case 3:
			lists = append(lists, [3]string{items[0], items[1], items[2]})
		case 2:
			lists = append(lists, [3]string{items[0], items[1], ""})
		default:
			continue
		}
	}
	if err = scanner.Err(); err != nil {
		return nil, fmt.Errorf("load primers from '%s': %s", file, err)
	}

	return lists, nil
}

func parsePrimers(primers [][3]string) ([][3][]byte, error) {
	list := make([][3][]byte, 0, len(primers))

	for _, items := range primers {
		name := []byte(items[0])

		forward := []byte(items[1])
		if seq.DNAredundant.IsValid(forward) != nil {
			return nil, fmt.Errorf("invalid primer sequence: %s", forward)
		}

		reverse := []byte(items[2])
		if seq.DNAredundant.IsValid(reverse) != nil {
			return nil, fmt.Errorf("invalid primer sequence: %s", reverse)
		}

		// compute revcom of reverse
		s, _ := seq.NewSeq(seq.DNAredundant, reverse)
		reverse = s.RevCom().Seq

		list = append(list, [3][]byte{name, forward, reverse})
	}
	return list, nil
}

// AmpliconFinder is a struct for locating amplicon via primer(s).
type AmpliconFinder struct {
	Seq []byte
	F   []byte // Forward primer
	R   []byte // R should be reverse complementary sequence of reverse primer

	MaxMismatch int
	FMindex     *fmi.FMIndex

	searched, found bool
	iBegin, iEnd    int // 0-based

	rF, rR *regexp.Regexp
}

// NewAmpliconFinder returns a AmpliconFinder struct.
func NewAmpliconFinder(sequence, forwardPrimer, reversePrimerRC []byte, maxMismatch int) (*AmpliconFinder, error) {
	if len(sequence) == 0 {
		return nil, fmt.Errorf("non-blank sequence needed")
	}
	if len(forwardPrimer) == 0 && len(reversePrimerRC) == 0 {
		return nil, fmt.Errorf("at least one primer needed")
	}

	if len(forwardPrimer) == 0 { // F = R.revcom()
		forwardPrimer = reversePrimerRC
		reversePrimerRC = nil
	}

	finder := &AmpliconFinder{
		Seq: bytes.ToUpper(sequence), // to upper case
		F:   bytes.ToUpper(forwardPrimer),
		R:   bytes.ToUpper(reversePrimerRC),
	}

	if maxMismatch > 0 { // using FM-index
		index := fmi.NewFMIndex()
		_, err := index.Transform(finder.Seq)
		if err != nil {
			return nil, err
		}
		finder.MaxMismatch = maxMismatch
		finder.FMindex = index
	} else {
		if seq.DNA.IsValid(finder.F) != nil { // containing degenerate base
			if maxMismatch > 0 {
				checkError(fmt.Errorf("it does not support both degenerate base and mismatch"))
			}
			s, _ := seq.NewSeq(seq.DNA, finder.F)
			rF, err := regexp.Compile(s.Degenerate2Regexp())
			if err != nil {
				return nil, fmt.Errorf("fail to parse primer containing degenerate base: %s", finder.F)
			}
			finder.rF = rF
		}

		if seq.DNA.IsValid(finder.R) != nil { // containing degenerate base
			if maxMismatch > 0 {
				checkError(fmt.Errorf("it does not support both degenerate base and mismatch"))
			}
			s, _ := seq.NewSeq(seq.DNA, finder.R)
			rR, err := regexp.Compile(s.Degenerate2Regexp())
			if err != nil {
				return nil, fmt.Errorf("fail to parse primer containing degenerate base: %s", finder.R)
			}
			finder.rR = rR
		}
	}
	return finder, nil
}

// LocateRange returns location of the range (begin:end, 1-based).
func (finder *AmpliconFinder) LocateRange(begin, end int, flanking bool, strictMode bool) ([]int, error) {
	if begin == 0 || end == 0 {
		checkError(fmt.Errorf("both begin and end in region should not be 0"))
	}

	if !finder.searched {
		_, err := finder.Locate()
		if err != nil {
			return nil, err
		}
	}
	if !finder.found {
		return nil, nil
	}

	var b, e int
	var ok bool
	if flanking {
		b, e, ok = SubLocationFlanking(len(finder.Seq), finder.iBegin, finder.iEnd, begin, end, strictMode)
	} else {
		b, e, ok = SubLocationInner(len(finder.Seq), finder.iBegin, finder.iEnd, begin, end, strictMode)
	}

	if ok {
		return []int{b, e}, nil
	}

	return nil, nil
}

// SubLocationInner returns location of a range (begin:end, relative to amplicon).
// B/E: 0-based, location of amplicon.
// begin/end: 1-based, begin: relative location to 5' end of amplicon,
// end: relative location to 3' end of amplicon.
// Returned locations are 1-based.
//
//                     F
//         -----===============-----
//              1 3 5                    x/y
//                       -5-3-1          x/y
//              F             R
//         -----=====-----=====-----     x:y
//
//              ===============          1:-1
//              =======                  1:7
//                =====                  3:7
//                   =====               6:10
//                   =====             -10:-6
//                      =====           -7:-3
//                                      -x:y (invalid)
//
func SubLocationInner(length, B, E, begin, end int, strictMode bool) (int, int, bool) {
	if begin == 0 || end == 0 {
		return 0, 0, false
	}

	if begin < 0 && end > 0 {
		checkError(fmt.Errorf("invalid inner region (-x:y): %d:%d", begin, end))
	}

	if length == 0 || B < 0 || B > length-1 || E < 0 || E > length-1 {
		return 0, 0, false
	}

	var b, e int

	if begin > 0 {
		b = B + begin
	} else {
		b = E + begin + 2
	}
	if b > length {
		if strictMode {
			return 0, 0, false
		}
		b = length
	} else if b < 1 {
		if strictMode {
			return 0, 0, false
		}
		b = 1
	}

	if end > 0 {
		e = B + end
	} else {
		e = E + end + 2
	}
	if e > length {
		if strictMode {
			return 0, 0, false
		}
		e = length
	} else if e < 1 {
		if strictMode {
			return 0, 0, false
		}
		e = 1
	}

	if b > e {
		return b, e, false
	}

	return b, e, true
}

// SubLocationFlanking returns location of a flanking range (begin:end, relative to amplicon).
// B/E: 0-based, location of amplicon.
// begin/end: 1-based, begin: relative location to 5' end of amplicon,
// end: relative location to 3' end of amplicon.
// Returned locations are 1-based.
//
//                     F
//         -----===============-----
//          -3-1                        x/y
//                             1 3 5    x/y
//              F             R
//         -----=====-----=====-----
//         =====                        -5:-1
//         ===                          -5:-3
//                             =====     1:5
//                               ===     3:5
//             =================        -1:1
//         =========================    -5:5
//                                       x:-y (invalid)
//
func SubLocationFlanking(length, B, E, begin, end int, strictMode bool) (int, int, bool) {
	if begin == 0 || end == 0 {
		return 0, 0, false
	}

	if begin > 0 && end < 0 {
		checkError(fmt.Errorf("invalid flanking region (x:-y): %d:%d", begin, end))
	}

	if length == 0 || B < 0 || B > length-1 || E < 0 || E > length-1 {
		return 0, 0, false
	}

	var b, e int
	var flag bool // 5' flanking is shorter than -begin

	if begin > 0 {
		b = E + begin + 1
	} else {
		b = B + begin + 1
	}
	if b > length {
		// b = length
		return 0, 0, false
	} else if b < 1 {
		if strictMode {
			return 0, 0, false
		}
		b = 1
		flag = true
		// return 0, 0, false
	}

	if end > 0 {
		e = E + end + 1
	} else {
		e = B + end + 1
	}
	if e > length {
		if strictMode {
			return 0, 0, false
		}
		e = length
	} else if e < 1 {
		if strictMode {
			return 0, 0, false
		}
		if flag {
			return 0, 0, false
		}
		e = 1
	}

	if b > e {
		return b, e, false
	}

	return b, e, true
}

// Locate returns location of amplicon.
// Locations are 1-based, nil returns if not found.
func (finder *AmpliconFinder) Locate() ([]int, error) {
	if finder.searched {
		if finder.found {
			return []int{finder.iBegin, finder.iEnd}, nil
		}
		return nil, nil
	}

	if finder.MaxMismatch <= 0 { // exactly matching
		// search F
		var i int

		if finder.rF == nil {
			i = bytes.Index(finder.Seq, finder.F)
			if i < 0 { // not found
				finder.searched, finder.found = true, false
				return nil, nil
			}
		} else {
			loc := finder.rF.FindSubmatchIndex(finder.Seq)
			if len(loc) == 0 {
				finder.searched, finder.found = true, false
				return nil, nil
			}
			i = loc[0]
		}

		if len(finder.R) == 0 { // only forward primer, returns location of F
			finder.searched, finder.found = true, true
			finder.iBegin, finder.iEnd = i, i+len(finder.F)-1
			return []int{i + 1, i + len(finder.F)}, nil
		}

		// two primers given, need to search R
		var j int
		if finder.rR == nil {
			j = bytes.Index(finder.Seq, finder.R)
			if j < 0 {
				finder.searched, finder.found = true, false
				return nil, nil
			}

			for {
				if j+1 >= len(finder.Seq) {
					break
				}
				k := bytes.Index(finder.Seq[j+1:], finder.R)
				if k < 0 {
					break
				}
				j += k + 1
			}
		} else {
			loc := finder.rR.FindAllSubmatchIndex(finder.Seq, -1)
			if len(loc) == 0 {
				finder.searched, finder.found = true, false
				return nil, nil
			}
			j = loc[len(loc)-1][0]
		}

		if j < i { // wrong location of F and R:  5' ---R-----F---- 3'
			finder.searched, finder.found = true, false
			return nil, nil
		}
		finder.searched, finder.found = true, true
		finder.iBegin, finder.iEnd = i, j+len(finder.R)-1
		return []int{i + 1, j + len(finder.R)}, nil
	}

	// search F
	locsI, err := finder.FMindex.Locate(finder.F, finder.MaxMismatch)
	if err != nil {
		return nil, err
	}
	if len(locsI) == 0 { // F not found
		finder.searched, finder.found = true, false
		return nil, nil
	}
	if len(finder.R) == 0 { // returns location of F
		sort.Ints(locsI) // remain the first location
		finder.searched, finder.found = true, true
		finder.iBegin, finder.iEnd = locsI[0], locsI[0]+len(finder.F)-1
		return []int{locsI[0] + 1, locsI[0] + len(finder.F)}, nil
	}

	// search R
	locsJ, err := finder.FMindex.Locate(finder.R, finder.MaxMismatch)
	if err != nil {
		return nil, err
	}
	if len(locsJ) == 0 {
		finder.searched, finder.found = true, false
		return nil, nil
	}
	sort.Ints(locsI) // to remain the FIRST location
	sort.Ints(locsJ) // to remain the LAST location
	finder.searched, finder.found = true, true
	finder.iBegin, finder.iEnd = locsI[0], locsI[0]+len(finder.R)-1
	return []int{locsI[0] + 1, locsJ[len(locsJ)-1] + len(finder.R)}, nil
}

// Location returns location of amplicon.
// Locations are 1-based, nil returns if not found.
func (finder *AmpliconFinder) Location() ([]int, error) {
	if !finder.searched {
		_, err := finder.Locate()
		if err != nil {
			return nil, err
		}
	}
	if !finder.found {
		return nil, nil
	}

	return []int{finder.iBegin + 1, finder.iEnd + 1}, nil
}
