#!/usr/bin/env python
"""
TODO
1. decrease memory usage
2. multi-fasta fastq file, ex. 454
3. split reads into small chuncks?

SHRiMP wrapper

Inputs:
1. reference seq
2. reads

Outputs:
1. table of 8 columns:
         chrom   ref_loc     read_id     read_loc    ref_nuc     read_nuc    quality     coverage
2. SHRiMP output

Parameters:
    -s    Spaced Seed                             (default: 111111011111)
    -n    Seed Matches per Window                 (default: 2)
    -t    Seed Hit Taboo Length                   (default: 4)
    -9    Seed Generation Taboo Length            (default: 0)
    -w    Seed Window Length                      (default: 115.00%)
    -o    Maximum Hits per Read                   (default: 100)
    -r    Maximum Read Length                     (default: 1000)
    -d    Kmer Std. Deviation Limit               (default: -1 [None])

    -m    S-W Match Value                         (default: 100)
    -i    S-W Mismatch Value                      (default: -150)
    -g    S-W Gap Open Penalty (Reference)        (default: -400)
    -q    S-W Gap Open Penalty (Query)            (default: -400)
    -e    S-W Gap Extend Penalty (Reference)      (default: -70)
    -f    S-W Gap Extend Penalty (Query)          (default: -70)
    -h    S-W Hit Threshold                       (default: 68.00%)

Command:
%rmapper -s spaced_seed -n seed_matches_per_window -t seed_hit_taboo_length -9 seed_generation_taboo_length -w seed_window_length -o max_hits_per_read -r max_read_length -d kmer -m sw_match_value -i sw_mismatch_value -g sw_gap_open_ref -q sw_gap_open_query -e sw_gap_ext_ref -f sw_gap_ext_query -h sw_hit_threshold <query> <target> > <output> 2> <log>

SHRiMP output:
>7:2:1147:982/1 chr3    +   36586562    36586595    2   35  36  2900    3G16G13
>7:2:1147:982/1 chr3    +   95338194    95338225    4   35  36  2700    9T7C14
>7:2:587:93/1   chr3    +   14913541    14913577    1   35  36  2960    19--16
"""

from __future__ import print_function

import os
import os.path
import re
import sys
import tempfile

assert sys.version_info[:2] >= (2, 6)


def stop_err(msg):
    sys.stderr.write("%s\n" % msg)
    sys.exit()


def reverse_complement(s):
    complement_dna = {
        "A": "T",
        "T": "A",
        "C": "G",
        "G": "C",
        "a": "t",
        "t": "a",
        "c": "g",
        "g": "c",
        "N": "N",
        "n": "n",
        ".": ".",
        "-": "-",
    }
    reversed_s = []
    for i in s:
        reversed_s.append(complement_dna[i])
    reversed_s.reverse()
    return "".join(reversed_s)


def generate_sub_table(result_file, ref_file, score_files, table_outfile, hit_per_read, insertion_size):
    invalid_editstring_char = 0
    all_score_file = score_files.split(",")

    if len(all_score_file) != hit_per_read:
        stop_err("One or more query files is missing. Please check your dataset.")

    temp_table_name = tempfile.NamedTemporaryFile().name
    temp_table = open(temp_table_name, "w")

    outfile = open(table_outfile, "w")

    # reference seq: not a single fasta seq
    refseq = {}
    chrom_cov = {}
    seq = ""
    title = None

    for line in open(ref_file):
        line = line.rstrip()
        if not line or line.startswith("#"):
            continue

        if line.startswith(">"):
            if seq:
                if title in refseq:
                    pass
                else:
                    refseq[title] = seq
                    chrom_cov[title] = {}
                seq = ""
            title = line[1:]
        else:
            seq += line
    if seq:
        if title not in refseq:
            refseq[title] = seq
            chrom_cov[title] = {}

    # find hits: one end and/or the other
    hits = {}
    for line in open(result_file):
        line = line.rstrip()
        if not line or line.startswith("#"):
            continue

        # FORMAT: readname contigname strand contigstart contigend readstart readend readlength score editstring
        fields = line.split("\t")
        readname = fields[0][1:]
        chrom = fields[1]
        strand = fields[2]
        chrom_start = int(fields[3]) - 1
        chrom_end = int(fields[4])
        read_start = fields[5]
        score = fields[8]
        editstring = fields[9]

        if hit_per_read == 1:
            endindex = "1"
        else:
            readname, endindex = readname.split("/")

        if readname in hits:
            if endindex in hits[readname]:
                hits[readname][endindex].append([strand, editstring, chrom_start, chrom_end, read_start, chrom])
            else:
                hits[readname][endindex] = [[strand, editstring, chrom_start, chrom_end, read_start, chrom]]
        else:
            hits[readname] = {}
            hits[readname][endindex] = [[strand, editstring, chrom_start, chrom_end, read_start, chrom]]

    # find score: one end and the other end
    hits_score = {}
    readname = ""
    score = ""
    for num_score_file in range(len(all_score_file)):
        score_file = all_score_file[num_score_file]
        for line in open(score_file):
            line = line.rstrip()
            if not line or line.startswith("#"):
                continue

            if line.startswith(">"):
                if score:
                    if readname in hits:
                        if len(hits[readname]) == hit_per_read:
                            if readname in hits_score:
                                if endindex in hits_score[readname]:
                                    pass
                                else:
                                    hits_score[readname][endindex] = score
                            else:
                                hits_score[readname] = {}
                                hits_score[readname][endindex] = score
                    score = ""
                if hit_per_read == 1:
                    readname = line[1:]
                    endindex = "1"
                else:
                    readname, endindex = line[1:].split("/")
            else:
                score = line

        if score:  # the last one
            if readname in hits:
                if len(hits[readname]) == hit_per_read:
                    if readname in hits_score:
                        if endindex in hits_score[readname]:
                            pass
                        else:
                            hits_score[readname][endindex] = score
                    else:
                        hits_score[readname] = {}
                        hits_score[readname][endindex] = score

    # call to all mappings
    for readkey in hits.keys():
        if len(hits[readkey]) != hit_per_read:
            continue

        matches = []
        match_count = 0

        if hit_per_read == 1:
            if len(hits[readkey]["1"]) == 1:
                matches = [hits[readkey]["1"]]
                match_count = 1
        else:
            end1_data = hits[readkey]["1"]
            end2_data = hits[readkey]["2"]

            for end1_hit in end1_data:
                crin_strand = {"+": False, "-": False}
                crin_insertSize = {"+": False, "-": False}

                crin_strand[end1_hit[0]] = True
                crin_insertSize[end1_hit[0]] = int(end1_hit[2])

                for end2_hit in end2_data:
                    crin_strand[end2_hit[0]] = True
                    crin_insertSize[end2_hit[0]] = int(end2_hit[2])

                    if end1_hit[-1] != end2_hit[-1]:
                        continue

                    if crin_strand["+"] and crin_strand["-"]:
                        if (crin_insertSize["-"] - crin_insertSize["+"]) <= insertion_size:
                            matches.append([end1_hit, end2_hit])
                            match_count += 1

        if match_count == 1:
            for x, end_data in enumerate(matches[0]):
                end_strand, end_editstring, end_chr_start, end_chr_end, end_read_start, end_chrom = end_data
                end_read_start = int(end_read_start) - 1

                if end_strand == "-":
                    refsegment = reverse_complement(refseq[end_chrom][end_chr_start:end_chr_end])
                else:
                    refsegment = refseq[end_chrom][end_chr_start:end_chr_end]

                match_len = 0
                editindex = 0
                gap_read = 0

                while editindex < len(end_editstring):
                    editchr = end_editstring[editindex]
                    chrA = ""
                    chrB = ""

                    if editchr.isdigit():
                        editcode = ""

                        while editchr.isdigit() and editindex < len(end_editstring):
                            editcode += editchr
                            editindex += 1
                            if editindex < len(end_editstring):
                                editchr = end_editstring[editindex]

                        for baseIndex in range(int(editcode)):
                            chrA += refsegment[match_len + baseIndex]
                            chrB = chrA

                        match_len += int(editcode)

                    elif editchr == "x":
                        # crossover: inserted between the appropriate two bases
                        # Two sequencing errors: 4x15x6 (25 matches with 2 crossovers)
                        # Treated as errors in the reads; Do nothing.
                        editindex += 1

                    elif editchr.isalpha():
                        editcode = editchr
                        editindex += 1
                        chrA = refsegment[match_len]
                        chrB = editcode
                        match_len += len(editcode)

                    elif editchr == "-":
                        editcode = editchr
                        editindex += 1
                        chrA = refsegment[match_len]
                        chrB = editcode
                        match_len += len(editcode)
                        gap_read += 1

                    elif editchr == "(":
                        editcode = ""

                        while editchr != ")" and editindex < len(end_editstring):
                            if editindex < len(end_editstring):
                                editchr = end_editstring[editindex]
                            editcode += editchr
                            editindex += 1

                        editcode = editcode[1:-1]
                        chrA = "-" * len(editcode)
                        chrB = editcode

                    else:
                        invalid_editstring_char += 1

                    if end_strand == "-":
                        chrA = reverse_complement(chrA)
                        chrB = reverse_complement(chrB)

                    pos_line = ""
                    rev_line = ""

                    for mappingIndex in range(len(chrA)):
                        # reference
                        chrAx = chrA[mappingIndex]
                        # read
                        chrBx = chrB[mappingIndex]

                        if chrAx and chrBx and chrBx.upper() != "N":
                            if end_strand == "+":
                                chrom_loc = end_chr_start + match_len - len(chrA) + mappingIndex
                                read_loc = end_read_start + match_len - len(chrA) + mappingIndex - gap_read

                                if chrAx == "-":
                                    chrom_loc -= 1

                                if chrBx == "-":
                                    scoreBx = "-1"
                                else:
                                    scoreBx = hits_score[readkey][str(x + 1)].split()[read_loc]

                                # 1-based on chrom_loc and read_loc
                                pos_line = (
                                    pos_line
                                    + "\t".join(
                                        (
                                            end_chrom,
                                            str(chrom_loc + 1),
                                            readkey + "/" + str(x + 1),
                                            str(read_loc + 1),
                                            chrAx,
                                            chrBx,
                                            scoreBx,
                                        )
                                    )
                                    + "\n"
                                )

                            else:
                                chrom_loc = end_chr_end - match_len + mappingIndex
                                read_loc = end_read_start + match_len - 1 - mappingIndex - gap_read

                                if chrAx == "-":
                                    chrom_loc -= 1

                                if chrBx == "-":
                                    scoreBx = "-1"
                                else:
                                    scoreBx = hits_score[readkey][str(x + 1)].split()[read_loc]

                                # 1-based on chrom_loc and read_loc
                                rev_line = (
                                    "\t".join(
                                        (
                                            end_chrom,
                                            str(chrom_loc + 1),
                                            readkey + "/" + str(x + 1),
                                            str(read_loc + 1),
                                            chrAx,
                                            chrBx,
                                            scoreBx,
                                        )
                                    )
                                    + "\n"
                                    + rev_line
                                )

                            if end_chrom in chrom_cov:
                                if chrom_loc in chrom_cov[end_chrom]:
                                    chrom_cov[end_chrom][chrom_loc] += 1
                                else:
                                    chrom_cov[end_chrom][chrom_loc] = 1

                            else:
                                chrom_cov[end_chrom] = {}
                                chrom_cov[end_chrom][chrom_loc] = 1

                    if pos_line:
                        temp_table.write("%s\n" % (pos_line.rstrip("\r\n")))
                    if rev_line:
                        temp_table.write("%s\n" % (rev_line.rstrip("\r\n")))

    temp_table.close()

    # chrom-wide coverage
    for line in open(temp_table_name):
        line = line.rstrip()
        if not line or line.startswith("#"):
            continue

        fields = line.split()
        chrom = fields[0]
        eachBp = int(fields[1])
        readname = fields[2]

        if hit_per_read == 1:
            fields[2] = readname.split("/")[0]

        if eachBp in chrom_cov[chrom]:
            outfile.write("%s\t%d\n" % ("\t".join(fields), chrom_cov[chrom][eachBp]))
        else:
            outfile.write("%s\t%d\n" % ("\t".join(fields), 0))

    outfile.close()

    if os.path.exists(temp_table_name):
        os.remove(temp_table_name)

    if invalid_editstring_char:
        print("Skip ", invalid_editstring_char, " invalid characters in editstrings")

    return True


def convert_fastqsolexa_to_fasta_qual(infile_name, query_fasta, query_qual):
    outfile_seq = open(query_fasta, "w")
    outfile_score = open(query_qual, "w")

    seq_title_startswith = ""
    qual_title_startswith = ""

    default_coding_value = 64  # Solexa ascii-code
    fastq_block_lines = 0

    for i, line in enumerate(open(infile_name)):
        line = line.rstrip()
        if not line or line.startswith("#"):
            continue

        fastq_block_lines = (fastq_block_lines + 1) % 4
        line_startswith = line[0:1]

        if fastq_block_lines == 1:
            # first line is @title_of_seq
            if not seq_title_startswith:
                seq_title_startswith = line_startswith

            if line_startswith != seq_title_startswith:
                outfile_seq.close()
                outfile_score.close()
                stop_err("Invalid fastqsolexa format at line %d: %s." % (i + 1, line))

            read_title = line[1:]
            outfile_seq.write(">%s\n" % line[1:])

        elif fastq_block_lines == 2:
            # second line is nucleotides
            read_length = len(line)
            outfile_seq.write("%s\n" % line)

        elif fastq_block_lines == 3:
            # third line is +title_of_qualityscore ( might be skipped )
            if not qual_title_startswith:
                qual_title_startswith = line_startswith

            if line_startswith != qual_title_startswith:
                outfile_seq.close()
                outfile_score.close()
                stop_err("Invalid fastqsolexa format at line %d: %s." % (i + 1, line))

            quality_title = line[1:]
            if quality_title and read_title != quality_title:
                outfile_seq.close()
                outfile_score.close()
                stop_err(
                    'Invalid fastqsolexa format at line %d: sequence title "%s" differes from score title "%s".'
                    % (i + 1, read_title, quality_title)
                )

            if not quality_title:
                outfile_score.write(">%s\n" % read_title)
            else:
                outfile_score.write(">%s\n" % line[1:])

        else:
            # fourth line is quality scores
            qual = ""
            fastq_integer = True
            # peek: ascii or digits?
            val = line.split()[0]
            fastq_integer = True
            try:
                int(val)
            except ValueError:
                fastq_integer = False

            if fastq_integer:
                # digits
                qual = line
            else:
                # ascii
                quality_score_length = len(line)
                if quality_score_length == read_length + 1:
                    # first char is qual_score_startswith
                    qual_score_startswith = ord(line[0:1])
                    line = line[1:]
                elif quality_score_length == read_length:
                    qual_score_startswith = default_coding_value
                else:
                    stop_err(
                        "Invalid fastqsolexa format at line %d: the number of quality scores ( %d ) is not the same as bases ( %d )."
                        % (i + 1, quality_score_length, read_length)
                    )

                for char in line:
                    score = ord(char) - qual_score_startswith  # 64
                    qual = "%s%s " % (qual, str(score))

            outfile_score.write("%s\n" % qual)

    outfile_seq.close()
    outfile_score.close()

    return True


def __main__():
    # SHRiMP path
    shrimp = "rmapper-ls"

    # I/O
    input_target_file = sys.argv[1]  # fasta
    shrimp_outfile = sys.argv[2]  # shrimp output
    table_outfile = sys.argv[3]  # table output
    single_or_paired = sys.argv[4].split(",")

    insertion_size = 600

    if len(single_or_paired) == 1:  # single or paired
        type_of_reads = "single"
        hit_per_read = 1
        input_query = single_or_paired[0]
        query_fasta = tempfile.NamedTemporaryFile().name
        query_qual = tempfile.NamedTemporaryFile().name

    else:  # paired-end
        type_of_reads = "paired"
        hit_per_read = 2
        input_query_end1 = single_or_paired[0]
        input_query_end2 = single_or_paired[1]
        insertion_size = int(single_or_paired[2])
        query_fasta_end1 = tempfile.NamedTemporaryFile().name
        query_fasta_end2 = tempfile.NamedTemporaryFile().name
        query_qual_end1 = tempfile.NamedTemporaryFile().name
        query_qual_end2 = tempfile.NamedTemporaryFile().name

    # SHRiMP parameters: total = 15, default values
    spaced_seed = "111111011111"
    seed_matches_per_window = "2"
    seed_hit_taboo_length = "4"
    seed_generation_taboo_length = "0"
    seed_window_length = "115.0"
    max_hits_per_read = "100"
    max_read_length = "1000"
    kmer = "-1"
    sw_match_value = "100"
    sw_mismatch_value = "-150"
    sw_gap_open_ref = "-400"
    sw_gap_open_query = "-400"
    sw_gap_ext_ref = "-70"
    sw_gap_ext_query = "-70"
    sw_hit_threshold = "68.0"

    # TODO: put the threshold on each of these parameters
    if len(sys.argv) > 5:
        try:
            if sys.argv[5].isdigit():
                spaced_seed = sys.argv[5]
            else:
                stop_err("Error in assigning parameter: Spaced seed.")
        except Exception:
            stop_err("Spaced seed must be a combination of 1s and 0s.")

        seed_matches_per_window = sys.argv[6]
        seed_hit_taboo_length = sys.argv[7]
        seed_generation_taboo_length = sys.argv[8]
        seed_window_length = sys.argv[9]
        max_hits_per_read = sys.argv[10]
        max_read_length = sys.argv[11]
        kmer = sys.argv[12]
        sw_match_value = sys.argv[13]
        sw_mismatch_value = sys.argv[14]
        sw_gap_open_ref = sys.argv[15]
        sw_gap_open_query = sys.argv[16]
        sw_gap_ext_ref = sys.argv[17]
        sw_gap_ext_query = sys.argv[18]
        sw_hit_threshold = sys.argv[19]

    # temp file for shrimp log file
    shrimp_log = tempfile.NamedTemporaryFile().name

    # convert fastq to fasta and quality score files
    if type_of_reads == "single":
        convert_fastqsolexa_to_fasta_qual(input_query, query_fasta, query_qual)
    else:
        convert_fastqsolexa_to_fasta_qual(input_query_end1, query_fasta_end1, query_qual_end1)
        convert_fastqsolexa_to_fasta_qual(input_query_end2, query_fasta_end2, query_qual_end2)

    # SHRiMP command
    if type_of_reads == "single":
        command = " ".join(
            (
                shrimp,
                "-s",
                spaced_seed,
                "-n",
                seed_matches_per_window,
                "-t",
                seed_hit_taboo_length,
                "-9",
                seed_generation_taboo_length,
                "-w",
                seed_window_length,
                "-o",
                max_hits_per_read,
                "-r",
                max_read_length,
                "-d",
                kmer,
                "-m",
                sw_match_value,
                "-i",
                sw_mismatch_value,
                "-g",
                sw_gap_open_ref,
                "-q",
                sw_gap_open_query,
                "-e",
                sw_gap_ext_ref,
                "-f",
                sw_gap_ext_query,
                "-h",
                sw_hit_threshold,
                query_fasta,
                input_target_file,
                ">",
                shrimp_outfile,
                "2>",
                shrimp_log,
            )
        )

        try:
            os.system(command)
        except Exception as e:
            if os.path.exists(query_fasta):
                os.remove(query_fasta)
            if os.path.exists(query_qual):
                os.remove(query_qual)
            stop_err(str(e))

    else:  # paired
        command_end1 = " ".join(
            (
                shrimp,
                "-s",
                spaced_seed,
                "-n",
                seed_matches_per_window,
                "-t",
                seed_hit_taboo_length,
                "-9",
                seed_generation_taboo_length,
                "-w",
                seed_window_length,
                "-o",
                max_hits_per_read,
                "-r",
                max_read_length,
                "-d",
                kmer,
                "-m",
                sw_match_value,
                "-i",
                sw_mismatch_value,
                "-g",
                sw_gap_open_ref,
                "-q",
                sw_gap_open_query,
                "-e",
                sw_gap_ext_ref,
                "-f",
                sw_gap_ext_query,
                "-h",
                sw_hit_threshold,
                query_fasta_end1,
                input_target_file,
                ">",
                shrimp_outfile,
                "2>",
                shrimp_log,
            )
        )
        command_end2 = " ".join(
            (
                shrimp,
                "-s",
                spaced_seed,
                "-n",
                seed_matches_per_window,
                "-t",
                seed_hit_taboo_length,
                "-9",
                seed_generation_taboo_length,
                "-w",
                seed_window_length,
                "-o",
                max_hits_per_read,
                "-r",
                max_read_length,
                "-d",
                kmer,
                "-m",
                sw_match_value,
                "-i",
                sw_mismatch_value,
                "-g",
                sw_gap_open_ref,
                "-q",
                sw_gap_open_query,
                "-e",
                sw_gap_ext_ref,
                "-f",
                sw_gap_ext_query,
                "-h",
                sw_hit_threshold,
                query_fasta_end2,
                input_target_file,
                ">>",
                shrimp_outfile,
                "2>>",
                shrimp_log,
            )
        )

        try:
            os.system(command_end1)
            os.system(command_end2)
        except Exception as e:
            if os.path.exists(query_fasta_end1):
                os.remove(query_fasta_end1)
            if os.path.exists(query_fasta_end2):
                os.remove(query_fasta_end2)
            if os.path.exists(query_qual_end1):
                os.remove(query_qual_end1)
            if os.path.exists(query_qual_end2):
                os.remove(query_qual_end2)
            stop_err(str(e))

    # check SHRiMP output: count number of lines
    num_hits = 0
    if shrimp_outfile:
        for line in open(shrimp_outfile):
            line = line.rstrip("\r\n")
            if not line or line.startswith("#"):
                continue
            try:
                line.split()
                num_hits += 1
            except Exception as e:
                stop_err(str(e))

    if num_hits == 0:  # no hits generated
        err_msg = ""
        if shrimp_log:
            for line in open(shrimp_log):
                if line.startswith("error"):  # deal with memory error:
                    err_msg += line  # error: realloc failed: Cannot allocate memory
                if re.search("Reads Matched", line):  # deal with zero hits
                    if int(line[8:].split()[2]) == 0:
                        err_msg = "Zero hits found.\n"
        stop_err("SHRiMP Failed due to:\n" + err_msg)

    # convert to table
    if type_of_reads == "single":
        generate_sub_table(shrimp_outfile, input_target_file, query_qual, table_outfile, hit_per_read, insertion_size)
    else:
        generate_sub_table(
            shrimp_outfile,
            input_target_file,
            query_qual_end1 + "," + query_qual_end2,
            table_outfile,
            hit_per_read,
            insertion_size,
        )

    # remove temp. files
    if type_of_reads == "single":
        if os.path.exists(query_fasta):
            os.remove(query_fasta)
        if os.path.exists(query_qual):
            os.remove(query_qual)
    else:
        if os.path.exists(query_fasta_end1):
            os.remove(query_fasta_end1)
        if os.path.exists(query_fasta_end2):
            os.remove(query_fasta_end2)
        if os.path.exists(query_qual_end1):
            os.remove(query_qual_end1)
        if os.path.exists(query_qual_end2):
            os.remove(query_qual_end2)

    if os.path.exists(shrimp_log):
        os.remove(shrimp_log)


if __name__ == "__main__":
    __main__()
