#!/usr/bin/env python3
import argparse
from npcaller.validator import AlignmentStats, randomize_fasta, align_to_reference, sam_to_bam, filter_fasta
from npcaller.fasta import FastaReader
import sys


def align(fasta_file, ref_file, sam_basename, graphmap_bin, ncores):
    samfile_name = sam_basename + ".sam"
    bamfile_name = sam_basename  # without .bam!
    align_to_reference(fasta_file, ref_file, samfile_name, graphmap_bin, ncores)
    sam_to_bam(samfile_name, bamfile_name)


def stats(ref_file, bam_file, ncores, tsv_file):
    fr = FastaReader(ref_file)
    _, ref = next(fr.get_entries())
    al_stats = AlignmentStats(bam_file, ref, ncores)
    al_stats.to_tsv(tsv_file)


if __name__ == "__main__":
    argp = argparse.ArgumentParser("Tools to validate the performance of a basecaller.")
    subparsers = argp.add_subparsers()
    randp = subparsers.add_parser("randomize", help="randomize a fasta file by mutating all nucleotides assuming a "
                                                    "uniform distribution. The sequence lengths of the "
                                                    "input file are kept.")
    randp.add_argument("-f", "--fasta", required=True, type=argparse.FileType('r'),
                       help="fasta file")
    randp.add_argument("-o", "--output", required=False, type=argparse.FileType('w'),
                       help="randomized fasta file output", default=sys.stdout)
    randp.set_defaults(which="randomize")

    statp = subparsers.add_parser("align", help="Align fasta file to reference ('mapping')")
    statp.add_argument("-f", "--fasta", required=True, type=argparse.FileType('r'),
                       help="fasta file with base-called reads. ")
    statp.add_argument("-r", "--reference", required=True, type=argparse.FileType('r'),
                       help="fasta file with reference sequence")
    statp.add_argument("-s", "--sam_basename", required=True, type=str,
                       help="Name of sam/bam files (without file extension)")
    statp.add_argument("-g", "--graphmap", required=False, type=str, default="",
                       help="Path to graphmap alignment tool")
    statp.add_argument("-n", "--ncores", required=False, type=int,
                       help="#CPU cores", default=-1)
    statp.set_defaults(which="align")

    statp = subparsers.add_parser("stats", help="Calculate statistics from Alignment")
    statp.add_argument("-r", "--reference", required=True, type=argparse.FileType('r'),
                       help="fasta file with reference sequence")
    statp.add_argument("-b", "--bam", required=True, type=argparse.FileType('rb'),
                       help="path to BAM-file that was generated by the 'align' command.")
    statp.add_argument("-t", "--tsv", required=False, type=argparse.FileType('w'), default=sys.stdout,
                       help="Write summary to this file")
    statp.add_argument("-n", "--ncores", required=False, type=int,
                       help="#CPU cores", default=None)
    statp.set_defaults(which="stats")

    filterp = subparsers.add_parser("filter", help="filter fasta file for keywords, such as 'template'. "
                                                   "All reads containing the keyword in the header are kept. ")
    filterp.add_argument("-k", "--keyword", required=True, type=str,
                       help="the keyword.")
    filterp.add_argument("-f", "--fasta", required=True, type=argparse.FileType('r'),
                       help="fasta file")
    filterp.add_argument("-o", "--output", required=False, type=argparse.FileType('w'),
                       help="filtered fasta file output", default=sys.stdout)
    filterp.set_defaults(which="filter")

    args = argp.parse_args()
    try:
        if args.which == "filter":
            filter_fasta(args.fasta, args.output, args.keyword)
        elif args.which == "align":
            align(args.fasta, args.reference, args.sam_basename, args.graphmap, args.ncores)
        elif args.which == "stats":
            stats(args.reference, args.bam, args.ncores, args.tsv)
        elif args.which == "randomize":
            randomize_fasta(args.fasta, args.output)
    except AttributeError:
        # no option specified.
        argp.print_usage()
