import shutil
import os
import time
import suffixtree


def scan_kmer(read):
    global K
    offset, seq1, seq2 = read.split("\t")
    ret = []
    for base in seq1:
        if base not in "ACGTN": return ret
    for base in seq2:
        if base not in "ACGTN": return ret
    for idx in range(len(seq1) - K + 1):
        ret += [(seq1[idx:idx + K], (offset, idx + 1))]
    for idx in range(len(seq2) - K + 1):
        ret += [(seq2[idx:idx + K], (offset, -idx - 1))]
    return filter(lambda x: "N" not in x[0], ret)


def align_by_anchor(kv):
    global ref_file_path, K, MATCH_RATE, JUDGE_THRESHOLD

    class anchored_read(object):
        __slots__ = ('id', 'kmeridx', 'reversed', 'offset1', 'offset2', 'seq1', 'seq2')

        def __init__(self, id, kmeridx, readFile1, readFile2):
            self.id = int(id)
            self.kmeridx = kmeridx
            self.reversed = False
            readFile1.seek(self.id)
            readFile2.seek(self.id)
            self.seq1 = readFile1.readline().strip()
            self.seq2 = readFile2.readline().strip()

        def reverse(self):
            self.reversed = not self.reversed
            self.kmeridx = -self.kmeridx
            self.seq1, self.seq2 = self.seq2, self.seq1

        def match_one_seq(self, seq1, seq2, offset, heuristic):
            """ to compute match score for same end with given offset
                seq1 and seq2 are one end of reads to compare,
                offset==3 means seq1[0] match seq2[3] (r2[offset])
                    0123---
                 0123456
            """
            if offset < 0:
                seq1, seq2 = seq2, seq1
                offset = -offset
            overlap = min(len(seq1), len(seq2) - offset)
            match_count = 0
            for i in range(overlap):
                match_count += (seq1[i] == seq2[i + offset])
                if heuristic and i == 4 and match_count < 4: return
            if float(match_count) / overlap < MATCH_RATE or match_count < K: return
            return match_count - (overlap - match_count)  # * (overlap - match_count)

        def match(self, read):
            """ "read" is the r1 on top, kmer locates at "a";
                "self" is the r2(anchor) on bottom, kmer locates at "b";
            """
            x = self.kmeridx - read.kmeridx
            anchored_score = self.match_one_seq(read.seq1, self.seq1, x, False)
            if not anchored_score: return
            best_match = None
            for y in range(-len(read.seq2) + 1, len(self.seq2)):
                match_result = self.match_one_seq(read.seq2, self.seq2, y, True)
                if not match_result: continue
                score = anchored_score + match_result  # - 0.05 * (x - y) * (x - y)
                if not best_match or best_match[0] < score:
                    best_match = (score, x, y)
            return best_match

    class anchor_aligment(list):

        def __init__(self, read):
            super(anchor_aligment, self).__init__()
            self.anchor(read, 0, 0)

        def anchor(self, read, left_offset, right_offset):
            self += [read]
            read.offset1, read.offset2 = (left_offset, right_offset)

        def match(self, read):
            return self[0].match(read)

        def report_doubt(self, kmer, col_id):
            doubts = []
            seq1_left, seq1_right = 0, len(self[0].seq1)
            seq2_left, seq2_right = 0, len(self[0].seq2)
            for read in self:
                seq1_left = min(seq1_left, read.offset1)
                seq1_right = max(seq1_right, read.offset1 + len(read.seq1))
                seq2_left = min(seq2_left, read.offset2)
                seq2_right = max(seq2_right, read.offset2 + len(read.seq2))
            for col_idx in range(seq1_left, seq1_right):
                basecount = {i: 0 for i in "ACGTN"}
                for read in self:
                    if 0 <= col_idx - read.offset1 < len(read.seq1):
                        basecount[read.seq1[col_idx - read.offset1]] += 1
                if basecount.values().count(0) == 4: continue
                for read in self:
                    if 0 <= col_idx - read.offset1 < len(read.seq1) and basecount[read.seq1[col_idx - read.offset1]] <= JUDGE_THRESHOLD:
                        doubts += [(kmer + str(col_id),
                                    (read.seq1[col_idx - read.offset1],
                                        (read.id + col_idx - read.offset1) * [1, -1][read.reversed],
                                     )
                                    )]
                for i in basecount:
                    if basecount[i] > JUDGE_THRESHOLD:
                        doubts += [(kmer + str(col_id), (i.lower(), basecount[i]))]
                col_id += 1
            for col_idx in range(seq2_left, seq2_right):
                basecount = {i: 0 for i in "ACGTN"}
                for read in self:
                    if 0 <= col_idx - read.offset2 < len(read.seq2):
                        basecount[read.seq2[col_idx - read.offset2]] += 1
                if basecount.values().count(0) == 4: continue
                for read in self:
                    if 0 <= col_idx - read.offset2 < len(read.seq2) and basecount[read.seq2[col_idx - read.offset2]] <= JUDGE_THRESHOLD:
                        doubts += [(kmer + str(col_id),
                                    (read.seq2[col_idx - read.offset2],
                                        (read.id + col_idx - read.offset2) * [-1, 1][read.reversed],
                                     )
                                    )]
                for i in basecount:
                    if basecount[i] > JUDGE_THRESHOLD:
                        doubts += [(kmer + str(col_id), (i.lower(), basecount[i]))]
                col_id += 1
            return doubts, col_id

        def print_pileup(self):
            pileup = ""
            seq1_left, seq1_right = 0, len(self[0].seq1)
            seq2_left, seq2_right = 0, len(self[0].seq2)
            for read in self:
                seq1_left = min(seq1_left, read.offset1)
                seq1_right = max(seq1_right, read.offset1 + len(read.seq1))
                seq2_left = min(seq2_left, read.offset2)
                seq2_right = max(seq2_right, read.offset2 + len(read.seq2))
            for read in self:
                pileup += str(read.id).center(10) + "." * (read.offset1 - seq1_left)
                pileup += read.seq1 + "." * (seq1_right - read.offset1 - len(read.seq1))
                pileup += "..." + "." * (read.offset2 - seq2_left) + read.seq2
                pileup += "." * (seq2_right - read.offset2 - len(read.seq2)) + "\n"
            return pileup

    # load reads from disk
    readFile1 = open(ref_file_path[0])
    readFile2 = open(ref_file_path[1])
    Ast = suffixtree.SuffixTree()
    Ust = suffixtree.SuffixTree()
    alignment_group = []
    for value in kv[1]:
        read = anchored_read(value[0], value[1], readFile1, readFile2)
        best, bestidx = None, -1
        if read.kmeridx < 0: read.reverse()
        Avote = Ast.tid_vote(read.seq1)
        Uvote = Ust.tid_vote(read.seq2)
        vote = [(tid, Avote[tid], Uvote[tid]) for tid in range(len(alignment_group))]
        if vote:
            bestidx, Abest, Ubest = max(vote, key=lambda x: x[1][1] + x[2][1])
            if read.match_one_seq(read.seq1, alignment_group[bestidx][0].seq1, Abest[0], False) and \
                    read.match_one_seq(read.seq2, alignment_group[bestidx][0].seq2, Ubest[0], False):
                best = (bestidx, Abest[0], Ubest[0], read.reversed)
        # for i, alignment in enumerate(alignment_group):
        #     match_result = alignment.match(read)
        #     if match_result and (not best or best[0] < match_result[0]):
        #         best, bestidx = match_result, i
        if not best:
            Ast.append(read.seq1, len(alignment_group))
            Ust.append(read.seq2, len(alignment_group))
            alignment_group += [anchor_aligment(read)]
        else:
            alignment_group[bestidx].anchor(read, best[1], best[2])
    report, col = [], 0
    log = "===%s===\n" % (kv[0])
    for alignment in alignment_group:
        doubts, col = alignment.report_doubt(kv[0], col)
        report += doubts
        log += alignment.print_pileup()
    # logfile = open("/temp/refread/log" + kv[0], "w")
    # logfile.write(log)
    # logfile.close()
    return report


def merge_in_partition(kv_iter):
    ufset = {}
    posdict = {}
    basecount = {}

    def find(k):
        if ufset[k] != k: ufset[k] = find(ufset[k])
        return ufset[k]

    for v, k in kv_iter:
        if v not in ufset:
            ufset[v] = v
            basecount[v] = [0, 0, 0, 0]
        else:
            v = find(v)
        if k[0] == "a": basecount[v][0] += k[1]
        elif k[0] == "c": basecount[v][1] += k[1]
        elif k[0] == "g": basecount[v][2] += k[1]
        elif k[0] == "t": basecount[v][3] += k[1]
        elif k not in posdict:
            posdict[k] = v
            if k[0] == "A": basecount[v][0] += 1
            if k[0] == "C": basecount[v][1] += 1
            if k[0] == "G": basecount[v][2] += 1
            if k[0] == "T": basecount[v][3] += 1
        else:
            u = find(posdict[k])
            ufset[v] = u
            basecount[u] = [basecount[u][i] + basecount[v][i] for i in range(4)]
    for k in posdict:
        u = find(posdict[k])
        if k[0] == "A" and basecount[u][0] > JUDGE_THRESHOLD: continue
        if k[0] == "C" and basecount[u][1] > JUDGE_THRESHOLD: continue
        if k[0] == "G" and basecount[u][2] > JUDGE_THRESHOLD: continue
        if k[0] == "T" and basecount[u][3] > JUDGE_THRESHOLD: continue
        yield (u, k)
    for v in ufset:
        if ufset[v] == v:
            count = basecount[v]
            if count[0] > JUDGE_THRESHOLD: yield (v, ('a', count[0]))
            if count[1] > JUDGE_THRESHOLD: yield (v, ('c', count[1]))
            if count[2] > JUDGE_THRESHOLD: yield (v, ('g', count[2]))
            if count[3] > JUDGE_THRESHOLD: yield (v, ('t', count[3]))


def judge(key_values):
    values = list(key_values[1])
    ACGT = {i: 0 for i in "ACGT"}
    for value in values:
        if value[0] != "N":
            if value[0] in "ACGT": ACGT[value[0]] += 1
            if value[0] in "acgt": ACGT[value[0].upper()] += value[1]
    ref, refcnt = max(ACGT.items(), key=lambda x: x[1])
    if refcnt <= JUDGE_THRESHOLD: return
    for value in values:
        if value[0] in "acgt" or ACGT[value[0]] > JUDGE_THRESHOLD and value[0] != 'N': continue
        yield str(value[1]) + "\t" + ref


def run(sc, inputfile="/home/x/xieluyu/reads/lineno_seq1_seq2.txt", outputdir="/home/x/xieluyu/output"):
    global P1, P2, P3, P4
    if os.path.exists(outputdir): shutil.rmtree(outputdir)
    start = time.time()
    reads_file = sc.textFile(inputfile, 64)
    P1 = reads_file.flatMap(scan_kmer)
    P2 = P1.groupByKey(1024).filter(lambda kv: len(kv[1]) > 1).flatMap(align_by_anchor)
    P3 = P2.mapPartitions(merge_in_partition)
    num_partitions = P2.getNumPartitions()
    while num_partitions != 1:
        num_partitions = (num_partitions - 1) / 4 + 1
        P3 = P3.partitionBy(num_partitions).mapPartitions(merge_in_partition)
    P4 = P3.groupByKey().flatMap(judge)
    P4.saveAsTextFile("file://" + outputdir)
    end = time.time()
    print "elapse:", end - start, "seconds"


# ref_file_path = ["/temp/refread/Phix1.fq", "/temp/refread/Phix2.fq"]
ref_file_path = ["/temp/refread/Ecoli1.fq", "/temp/refread/Ecoli2.fq"]
K = 14
MATCH_RATE = 0.9
JUDGE_THRESHOLD = 2
if __name__ == '__main__':
    from pyspark import SparkContext, SparkConf
    conf = SparkConf().setAppName("dec")  # .setMaster("local[4]")
    sc = SparkContext(conf=conf)
    run(sc)
