#! /usr/bin/python

import os
SCRIPT_PATH = os.path.dirname(os.path.realpath(__file__))
import sys

from collections import defaultdict


# To generate random reads
import random

# For manipulating CIGAR strings
import re

# For working wit gzip files
# Compressing profile to reduce size and make it small enough for GitHub
import gzip

# Mainly for callculating a reverse complement of a sequence
# and other utility stuff
import fastqparser

# To be able to generate random reads in SAM format
import utility_sam


CIGAR_OPERATI
ONS_ALL = ['M', 'I', 'D', 'N', 'S', 'H', 'P', '=', 'X'];
CIGAR_OPERATIONS_BASIC = ['M', 'I', 'D', 'S', 'H'];
CIGAR_OPERATIONS_EXTENDED = ['M', 'I', 'D', 'S', 'H', '=', 'X'];

# CIGAR profile version number
# Used to check if loaded profile is compatible with the current code
# Should be changed only when new code makes old profile unusable
PROFILE_VERSION = '1.1'


# CIGARProfile: Class containing information about a dataset
# The infromatio
n is stored as read CIGAR strings, plus some additional stuff
# It is used to generate simulated reads applying stored CIGAR strings to a random part of a given reference
# TODO:
# - Decide how to generate quals
# - Decide how to generate Ns (lots of them in nanopore reads)
# - (least important) Decide what to do about short tandem repeats (STR), inserts and deletes both
# - decide what to do with clipping
#       - randomly generate bases
#       - take part of the reference

# CIGARLine: Class containg information about one CIGAR string
class CIGARLine:
    def __init__(self, cigar = '', position = -1, qname = '', rname = '', quality = -1, GCcontent = 0.0, quals = None):
        self.cigar = cigar
        self.position = position
        self.qname = qname
        self.rname = rname
        self.quality = quality
        self.GCcontent = GCcontent          # GC content of a read that generated a CIGAR string
        self.quals = quals                  # Quality values of a read that generated a CIGAR string

    def toTSVstring(self):
        tsvstring = ('%s\t%d\t%s\t%s\t%d\t%0.4f' % (self.cigar, self.position, self.qname, self.rname, self.quality, self.GCcontent))

        # Add quals if they are defined
        if self.quals is not None:
            tsvstring += '\t%s' % self.quals

        return tsvstring

    def fromTSVstring(self, line):
        # If present, removing \n from the end of line
        if line[-1] == '\n':
            line = line[:-1]
        elements = line.split('\t')
        self.cigar = elements[0]
        self.position = int(elements[1])
        self.qname = elements[2]
        self.rname  = elements[3]
        self.quality = int(elements[4])
        self.GCcontent = float(elements[5])

        # getting quals string, if it exists
        if len(elements) > 6:
            self.quals = elements[6]

    # Splits the CIGAR string into individual operations, in the
    # same format as the original CIGAR string is in.
    # The result is returned as an array of a 2-element tuples, e.g.
    # [(12, 'M'), (3, 'D'), (4, 'M')].
    def splitCigar(self):
        # Using regular expressions to find repeating digit and skipping one character after that
        pattern = '(\d+)(.)'
        operations = re.findall(pattern, self.cigar)

        return operations

    # Returns an aproximate length of the read that generated the cigar string
    def length(self):
        return sum([int(x[0]) for x in self.splitCigar()])


class CIGARProfile:
    def __init__(self, name, GCcontent = 0.0):
        self.name = name
        self.clines = []

        # GC content for the whole profile
        # For GC content to be relevan't, the profile should be constructed for a single genome
        self.GCcontent = GCcontent

        # Table containing count of different mutations
        # Used to generate different mutations with different probability
        self.mutCntTable = None

        # Placeholder that allows the definition of profile-wide distribution for quals and Ns
        self.qualDistrib = None
        self.NDistrib = None

        # Profile version, used to determine if loaded profile is compatible with the current code
        self.version = PROFILE_VERSION

    def appendCLine(self, cline):
        if cline is None:
            raise Exception('Trying to add nonexisting CIGAR line! (%d)' % len(self.clines))
        self.clines.append(cline)

    def getRandomCLine(self):
        rand = random.randint(0, len(self.clines)-1)
        return self.clines[rand]

    # ATM using all cigar lines to generate reads at random positions
    # TODO: Consider the situation where reference is too short for a randomly chosen CIGAR
    def generateRandomRead(self, reference):
        cline = self.getRandomCLine()
        reflen = len(reference)
        readlen = cline.length()

        # Generate random position
        randpos = random.randint(0, reflen-readlen)
        startread = reference[randpos:randpos+readlen]

        # 50% of the time generate reverse complement sequence
        complement = random.choice((True, False))
        # if complement:
        #     startread = fastqparser.revcomp_seq(startread)

        read = ''
        pos = 0         # Where in the starting read are we currently positioned

        operations = cline.splitCigar()
        # applying cigar operations
        for op in operations:
            opsize = int(op[0])
            optype = op[1]
            if optype in ('H', 'S'):
                # Clipping: Generate random bases, atm just putting Hs
                bases = generate_randomseq(opsize)
                read += bases
                # read += 'H'*opsize
            elif optype == '=':
                # Copy from the starting read / reference
                read += startread[pos:pos+opsize]
                pos += opsize
            elif optype == 'X':
                for i in xrange(opsize):
                    read += mutate(startread[pos+i], self.mutCntTable)
                # read += 'X'*opsize
                pos += opsize
            elif optype == 'I':
                # Insert random bases, atm just putting Is
                # Could be joined together with clipping
                bases = generate_randomseq(opsize)
                read += bases
                # read += 'I'*opsize
            elif optype == 'D':
                # Skip bases in starting read
                pos += opsize
            elif optype == 'N':
                read += 'N'*opsize
                pos += opsize
            else:
                sys.stderr.write('\nInvalid CIGAR operation for generating random reads: %s\n' % optype)
                exit(1)

        # Return generated read in a form of a SAMline
        sline = utility_sam.SAMLine()
        sline.qname = ''
        sline.flag = 16 if complement else 0        # Setting forward/reverse flag
                                                    # Ignoring other possibilities, no secondary alignments and paired reads
        sline.rname = ''
        sline.pos = randpos
        sline.mapq = 255
        sline.cigar = cline.cigar
        sline.mrnm = ''
        sline.mpos = 0
        sline.isize = 0
        sline.seq = read
        sline.qual = cline.quals
        sline.original_line = ''

        return sline


    # Deprecated:
    # Separate function to generate a whole dataset faster and with less memory
    # (instead of calling generateRandomRead many times)
    def generateRandomReadsByNumber(self, reference, numreads):
        reads = []
        reflen = len(reference)

        # Generating reads
        for i in xrange(numreads):
            cline = self.getRandomCLine()
            readlen = cline.length()
            # Generate random position
            randpos = random.randint(0, reflen-readlen)
            startread = reference[randpos:randpos+readlen]

            # 50% of the time generate reverse complement sequence
            complement = random.choice((True, False))
            if complement:
                startread = fastqparser.revcomp_seq(startread)

            read = ''
            pos = 0         # Where in the starting read are we currently positioned

         