import os
from tqdm import tqdm
import csv


try: 
    import Tailer.TailerFunctions as tf #hmm, something is weird here, here's a workaround
    import Tailer.LocalAligner as la
    # if someone smarter knows what I'm doing wrong, let me know: timnicholsonshaw@gmail.com
except:
    import TailerFunctions as tf
    import LocalAligner as la

class miRNA_Tail(tf.Tail):
    def __init__(self, alignment, gene, totalReferenceLengthDict,rpad=0):
        self.readID = alignment.query_name 
        self.is_reverse = alignment.is_reverse # Was the sequence rev comp'd by aligner

        self.gene = gene

        self.threeEnd = alignment.reference_end - (totalReferenceLengthDict[gene] - rpad)

        if alignment.cigartuples[-1][0] == 4: # Interested in softclips at end of cigar, 4=softclipped
            self.tailLen = alignment.cigartuples[-1][1]
            self.tailSeq = alignment.query_sequence[alignment.query_length - self.tailLen:] # saves seq of tail
        else:
            self.tailLen = 0
            self.tailSeq = None
            


        
        



def tailedReadsToTailFilemiRNA(TailedReads, outLoc, threeEndThresh = 100, seq_out=False, mature_end=0):
    """Writes TailedReads object to file
    
    :param TailedReads          Previously generated TailedReads object
    :param outLoc               String indictating where it should be written
    :param threeEndThresh       Threshold for inclusion in write, absolute distance to mature end
    """

    out =[]

    for seq, tailedRead in TailedReads.items():
        bestTails = tailedRead.findBestAlignment()

        if not bestTails: continue

        # Joins ensIDs/geneNames together for reads with multiple best tails
        #ensIDs = "|".join([str(x.geneID) for x in bestTails])
        ensIDs = "NA"
        #geneNames = "|".join([str(x.gene) for x in bestTails])
        geneNames = sorted([str(x.gene) for x in bestTails])
        geneNames = "".join(geneNames)

        # If it's further than this threshold from the mature end, it's probably spurious
        if bestTails:
            if bestTails[0].threeEnd < -threeEndThresh: continue
            elif bestTails[0].threeEnd > threeEndThresh: continue

            if not seq_out: #Won't add sequence information without sequence flag
                out.append([ tailedRead.count,
                ensIDs, 
                geneNames, 
                bestTails[0].threeEnd + bestTails[0].tailLen - mature_end, 
                bestTails[0].tailLen, 
                bestTails[0].tailSeq])

            else:
                if bestTails[0].is_reverse:
                    seq = tf.reverse_complement(seq)
                else:
                    pass

                out.append([seq, 
                tailedRead.count,
                ensIDs, 
                geneNames, 
                bestTails[0].threeEnd + bestTails[0].tailLen - mature_end, 
                bestTails[0].tailLen, 
                bestTails[0].tailSeq])




    #Sort by count
    key_column = 0
    if seq_out: key_column = 1
    out = sorted(out, key=lambda x:x[key_column], reverse=True)

    # Add header
    if seq_out:
        out = [["Sequence", "Count", "EnsID", "Gene_Name", "End_Position", "Tail_Length", "Tail_Sequence" ]] + out
    else:
        out = [["Count", "EnsID", "Gene_Name", "End_Position", "Tail_Length", "Tail_Sequence" ]] + out

    
    with open(outLoc, 'w') as csvfile:
        writer = csv.writer(csvfile)

        for line in out:
            writer.writerow(line)


def makeTailedReadsDictmiRNA(alignment_dict, totalReferenceLengthDict=""):
    """Collapses reads by sequence and creates Tail objects for every alignment/gene

    :param alignment_dict       generated by above function
    :param gtf_db               SQL database generated by above function
    :return                     dictionary with sequence as keys and TailedRead objects as values
    """

    out_dict = {}
    for key, _ in tqdm(alignment_dict.items()):
        
        try: # If the sequence has been seen before, just increments the count
            out_dict[alignment_dict[key][0].query_sequence].count+=1
            continue
        except KeyError: # If it hasn't been seen before, creates a new TailedRead object
            out_dict[alignment_dict[key][0].query_sequence] = tf.TailedRead()
        
        # Iterates over every alignment and gene combination and creates a Tail object
        #This probably leads to some wrong combos, but it gets sorted out later by filters
        for aln in alignment_dict[key]:
            gene = aln.reference_name
            out_dict[alignment_dict[key][0].query_sequence].tails.append(miRNA_Tail(aln, gene, totalReferenceLengthDict))

    return out_dict

def run_miRNA_tailer(args):
    
    for file in args.files:
        pre, ext = os.path.splitext(file)

        alignment = tf.getHandleOnIndexedBam(file)
        totalReferenceLengthDict = {x:y for x, y in zip(alignment.references, alignment.lengths)}

        aln_dict = tf.makeAlignmentDict(alignment)

        print("Calculating tails...")
        # tails everything and creates a dictionary of tailed reads
        tailedReads = makeTailedReadsDictmiRNA(aln_dict, totalReferenceLengthDict=totalReferenceLengthDict)

        # write to file
        tailedReadsToTailFilemiRNA(tailedReads, pre + "_tail.csv", threeEndThresh=args.threshold, seq_out=args.sequence, mature_end=args.mature)

        print("Wrote " + pre + "_tail.csv " + " to disk.")




