import sys
from collections import Counter
import pysam
import math
import click

"""
Jia Jinbu, 2020.09.15.

Extract polyA position from PacBio ccs seqeunce without adapters and has the same strand with mRNA.
"""

@click.command()
@click.option('-i', '--inbam', help='Input bam file generated by minimap2', 
                    required=True, type=click.Path(exists=True))
@click.option('-s', '--inseq', help='Input fasta seqeunce used for minimap22', required=True)
@click.option('-o', '--out', help='Out poly results. Remove chimerical reads.', required=True)
def main(inbam, inseq, out):
    """
    Extract polyA position from PacBio ccs seqeunce without adapters and has the same strand with mRNA.
    """
    SEP = ","
    PAD_LENGTH = 20
    MATCH_SCORE = 1
    MISMATCH_SCORE = -1.5

    with open(out, 'w') as o:
        o.write('read_core_id\tpolya_start_base\tpolya_end_base\tpolya_length\tpolya_score\n')
        for read, (read_strand, seq_length, read_name, left_name, left_clip_seq, right_name, right_clip_seq) in iter_bam_clip_seq(inbam, inseq, PAD_LENGTH):
            #read_length = read.infer_read_length()
            seq = left_clip_seq if read_strand == "-" else right_clip_seq
            polyA_start, polyA_end, polyA_score, polyA_seq = polyA_finder(seq, match = MATCH_SCORE, mis = MISMATCH_SCORE)
            if polyA_start == 0: #0 indicate not find polyA
                polyA_length = 0
                read_polyA_end = 0
                read_polyA_start = 0
                #read_polyA_start = ""
                #after_polyA_seq = ""
            else:
                polyA_length = polyA_end - polyA_start + 1
                #read_polyA_start = seq[:polyA_start]
                #after_polyA_seq = seq[polyA_end:]
                read_polyA_end = seq_length - len(seq) + polyA_end
                read_polyA_start = seq_length - len(seq) + polyA_start
            read_core_id = ",".join(read_name.split(",")[:4])
            o.write(f'{read_core_id}\t{read_polyA_start}\t{read_polyA_end}\t{polyA_length}\t{polyA_score}\n')
            
def read_fasta_to_dict(filein):
    """
    !!!The function in included in both adapterFinder.py and 
    pacbio_find_polyA.py. They are same function, but haven't
    be put in a module to keep each script can be run independently.
    If you want to modify one of them, please modify them at the 
    same time.
    
    Input: fasta files
    
    Output: dict
    key: sequence name
         #>seq1 npr1
         #run: l[1:].split()[0]
         #will get `seq1`
    value: sequence
    """
    
    id2seq = {}
    
    seq_id, seq = "", ""
    for line_num, l in enumerate(open(filein)):
        l = l.strip()
        if l.startswith(">"):
            if seq_id:
                id2seq[seq_id] = seq
            seq = ""
            seq_id = l[1:].split()[0]
        else:
            seq += l
    if seq_id:
        id2seq[seq_id] = seq
        
    return(id2seq)

def revcom(seq):
    """
    !!!The function in included in both adapterFinder.py and 
    pacbio_find_polyA.py and extract_read_info.py. 
    They are same function, but haven't be put in a module to 
    keep each script can be run independently. If you want to 
    modify one of them, please modify them at the same time.
    
    Return the reverse complement sequence of origin sequence.
    
    The origin sequence can be in uppercase and lowercase letters.
    But the output is in uppercase letters.
    All letters not in `ATCGatcg` will be converted to `N`. 
    """
    def complement(seq):
        seq = seq.upper()
        basecomplement = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A','N':'N'}
        def _com(base):
            try:
                return basecomplement[base]
            except:
                return "N"
        letters = list(seq)
        letters = [_com(base) for base in letters]
        return ''.join(letters)
            
    return complement(seq[::-1])
    
def iter_bam_clip_seq(filein_bam, filein_seq, pad_length=20):

    '''
    !!!The function in included in both adapterFinder.py and 
    pacbio_find_polyA.py. They are same function, but haven't
    be put in a module to keep each script can be run independently.
    If you want to modify one of them, please modify them at the 
    same time. 
    
    genome 5'------------------------------------------------3'
    mapping region     |||||||||||||||||||||||||||
    read         ---------------------------------------
      5' clip    <<<<<<---                     --->>>>>> 3' clip
                       pad                     pad
    The 5' or 3' of clip sequence is based on the genome fasta 
    direction, not the original read direction.

    Extract 5' and 3' clip sequence, due to the clip sequence is removed 
    from bam file when alignment contain hard clip, thus you also need 
    to provide original sequence file.

    Input: 
    filein_bam: aligned file
    filein_seq: The origin seq file used to generate filein_bam.

    Output:
    iterator, each element is :
    [read, [read_strand, seq_length, read_name, \
            left_clip_name, left_clip_seq, \
            right_clip_name, right_clip_seq]]        
    
    read_id:        00000449-e191-4eda-92e7-c4b2843daba2
    read_core_id:   00000449-e191-4eda-92e7-c4b2843daba2,chr4,8523083,8523682
    read_name:      00000449-e191-4eda-92e7-c4b2843daba2,chr4,8523083,8523682,+,660,58,55,20
    left_clip_name: 00000449-e191-4eda-92e7-c4b2843daba2,chr4,8523083,8523682,+,660,58,55,20,5
    right_clip_name:00000449-e191-4eda-92e7-c4b2843daba2,chr4,8523083,8523682,+,660,58,55,20,3
    

    read:           the object of each line in bam file, generated by pysam
    read_id:        query_name in bam file, e.g.
    read_core_id:
        Due to one read may be mapped to multiple positions, thus only read
        name cannot uniquly represent the alignment, so we combine the read
        name (the first column in bam file) and the mapping position as the 
        read_core_id by `,`: `read_name,chr_name,start,end`.

        Note: Logically, the read_core_id may be also not uniquely because: 
        read contain duplicated regions, all regions can mapped to same 
        positions in genome. And `read_name,chr_name,start,end` didn't 
        consider strand information. The output didn't remove duplicated 
        read_core_id row, thusbefore generating the final result in this 
        pipeline, we need to removed them in specific step (now is in merge 
        step).
    read_strand:    read alignment direction `+` or `-`
    seq_length:     original sequence length
    read_name:      final read_name
        `read_name,chr_name,start,end,strand,seq_length,left_clip_length,
        right_clip_length,pad_length`.

        Note: This is used to record the clip information, but the final 
        read_name has been used in the downstream pipeline, thus you can 
        just replace it by read_core_id, which is exactly used in the 
        downstream pipeline.
    left_clip_name:      `read_name,5`
    left_clip_seq:  5' clip sequence contain adjecent pad_length mapping bases
    right_clip_name:     `read_name,3`
    right_clip_seq: 3' clip sequence contain adjecent pad_length mapping bases

    For developer:
    Note: The tenth column of bam file store query sequence, you can use
    read.query_sequence (pysam package) to extract it. In minimap2 alignment
    file, this column store the reverse complement of read sequence if 
    the alignment direction is minus (read.is_reverse is True). That means
    the sequence is in same direction with genome sequence, but it is not 
    genome sequence if mismatch or indel exist. If you want to extract genome 
    sequence, you need to use original genome fasta files. If the alignment 
    contain soft cliped bases, but didn't contain hard cliped bases, thus if 
    the alignment contain hard cliped bases (extract from read.cigartuples), 
    you need to extract original read sequence from fastq files, and reverse 
    complement it if the alignment direction is minus.
    
    left_clip_seq and right_clip_seq may be empty string.
    '''

    def get_clip_seq(read, origin_seqs, pad_length=20):
        #supposed that the first and the last type in cigartuples is in [0, 4, 5]
        read_strand = "-" if read.is_reverse else "+"
        left_clip_type, left_clip_length = read.cigartuples[0]
        right_clip_type, right_clip_length = read.cigartuples[-1]
        if left_clip_type == 0:
            left_clip_length = 0
        if right_clip_type == 0:
            right_clip_length = 0
        left_clip_length += pad_length
        right_clip_length += pad_length
    
        seq = read.query_sequence
        if left_clip_type == 5 or right_clip_type == 5:
            #if Hard clip, need origin_seqs to extract origin seq
            #In this case, if alignment direction is minus, 
            #need to reverse the origin seq
            seq = origin_seqs[read.query_name]
            if read_strand == "-": seq = revcom(seq)
        seq_length = len(seq)

        left_clip_seq = seq[:left_clip_length]
        left_clip_seq = revcom(left_clip_seq)
        right_clip_seq = seq[(-right_clip_length):]
    
        read_name = ",".join([read.query_name, 
                              read.reference_name,
                              str(read.reference_start + 1), 
                              str(read.reference_end),
                              read_strand,
                              str(seq_length),
                              str(left_clip_length),
                              str(right_clip_length),
                              str(pad_length)])
        left_clip_name, right_clip_name = [read_name + ",5", read_name + ",3"]
        return [read_strand, seq_length, read_name, left_clip_name, left_clip_seq, right_clip_name, right_clip_seq]

    origin_seqs = read_fasta_to_dict(filein_seq)

    with pysam.AlignmentFile(filein_bam, "rb") as bam_obj:
        for read in bam_obj.fetch():
            #remove unampped read
            if read.is_unmapped:
                continue
            yield([read, get_clip_seq(read, origin_seqs, pad_length)])

def max_subarray(A):
    """
    !!!The function in included in both polyCaller.py and 
    pacbio_find_polyA.py. They are same function, but haven't
    be put in a module to keep each script can be run independently.
    If you want to modify one of them, please modify them at the 
    same time.
    
    Maximum subarray problem: select subarray with maxmium sum
    modified Kadane's algorithm (not greedy)
    
    return (index is 0-based), you can get the subarray by A[start_index:(end_index+1)]:
    [start_index, end_index, sum]
    
    if the maxmium sum is <= 0, then return [-1, -1, 0]
    """
        
    max_ending_here = max_so_far = 0
    max_start_index = startIndex = 0
    max_end_index = -1
    for i, x in enumerate(A):
        if 0 >= max_ending_here + x:
        #For greedy at left side : if 0 > max_ending_here + x:
            startIndex = i+1
            max_ending_here = 0
        else:
            max_ending_here += x
        if max_ending_here > max_so_far:
        #For greedy at right side : if max_ending_here >= max_so_far:
            max_so_far = max_ending_here
            max_start_index = startIndex
            max_end_index = i
    
    if max_so_far <= 0 or (max_start_index > max_end_index):
        return ((-1, -1, 0))
    else:
        return (max_start_index, max_end_index, max_so_far)
    
def polyA_finder(seq, base="A", match = 1, mis = -1.5):
    """
    !!!The function in included in both polyCaller.py and 
    pacbio_find_polyA.py. They are same function, but haven't
    be put in a module to keep each script can be run independently.
    If you want to modify one of them, please modify them at the 
    same time.
    """
    scores = [match if base == s else mis for s in seq]
    start_index, end_index, max_score = max_subarray(scores)
    return (start_index+1, end_index+1, max_score, seq[start_index:(end_index+1)])
            
if __name__ == "__main__":
    main()