""" Re-aligner small RNA sequence from SAM/BAM file (miRBase annotation)"""

import traceback
import os.path as op
import os
import re
import shutil
import pandas as pd
import pysam

from mirtop.libs import do
from mirtop.libs.utils import file_exists
import mirtop.libs.logger as mylog
from mirtop.install import _get_miraligner
from mirtop.mirna.snps import create_vcf
from mirtop.mirna.realign import *

logger = mylog.getLogger(__name__)

def _download_mirbase(args, version="CURRENT"):
    """
    Download files from mirbase
    """
    if not args.hairpin or not args.mirna:
        logger.info("Working with version %s" % version)
        hairpin_fn = op.join(op.abspath(args.out), "hairpin.fa.gz")
        mirna_fn = op.join(op.abspath(args.out), "miRNA.str.gz")
        if not file_exists(hairpin_fn):
            cmd_h = "wget ftp://mirbase.org/pub/mirbase/%s/hairpin.fa.gz -O %s &&  gunzip -f !$" % (version, hairpin_fn)
            do.run(cmd_h, "download hairpin")
        if not file_exists(mirna_fn):
            cmd_m = "wget ftp://mirbase.org/pub/mirbase/%s/miRNA.str.gz -O %s && gunzip -f !$" % (version, mirna_fn)
            do.run(cmd_m, "download mirna")
    else:
        return args.hairpin, args.mirna

def _make_unique(name, idx):
    """Make name unique in case only counts there"""
    seqname = re.compile(".[aA-zZ]+_x[0-9]+")
    if seqname.match(name):
        tags = name[1:].split("_x")
        return ">%s_%s_x%s" % (tags[0], idx, tags[1])
    return name.replace("@", ">")

def _filter_seqs(fn):
    """Convert names of sequences to unique ids"""
    out_file = op.splitext(fn)[0] + "_unique.fa"
    idx = 0
    if not file_exists(out_file):
        with open(out_file, 'w') as out_handle:
            with open(fn) as in_handle:
                for line in in_handle:
                    if line.startswith("@") or line.startswith(">"):
                        fixed_name = _make_unique(line.strip(), idx)
                        seq = in_handle.next().strip()
                        counts = _get_freq(fixed_name)
                        if len(seq) < 26 and (counts > 1 or counts == 0):
                            idx += 1
                            print >>out_handle, fixed_name
                            print >>out_handle, seq
                        try:
                            if line.startswith("@"):
                                in_handle.next()
                                in_handle.next()
                        except:
                            pass

    return out_file

def _convert_to_fasta(fn):
    out_file = op.splitext(fn)[0] + ".fa"
    with open(out_file, 'w') as out_handle:
        with open(fn) as in_handle:
            for line in in_handle:
                if line.startswith("@"):
                    seq = in_handle.next()
                    _ = in_handle.next()
                    qual = in_handle.next()
                elif line.startswith(">"):
                    seq = in_handle.next()
                count = 2
                if line.find("_x"):
                    count = int(line.strip().split("_x")[1])
                if count > 1:
                    print >>out_handle, ">%s" % line.strip()[1:]
                    print >>out_handle, seq.strip()
    return out_file

def _get_pos(string):
    name = string.split(":")[0][1:]
    pos = string.split(":")[1][:-1].split("-")
    return name, map(int, pos)

def _read_mature(matures, sps):
    mature = defaultdict(dict)
    with open(matures) as in_handle:
        for line in in_handle:
            if line.startswith(">") and line.find(sps) > -1:
                name = line.strip().replace(">", " ").split()
                mir5p = _get_pos(name[2])
                mature[name[0]] = {mir5p[0]: mir5p[1]}
                if len(name) > 3:
                    mir3p = _get_pos(name[3])
                    mature[name[0]].update({mir3p[0]: mir3p[1]})
    return mature

def _read_precursor(precursor, sps):
    """
    Load precursor file for that species
    """
    hairpin = defaultdict(str)
    name = None
    with open(precursor) as in_handle:
        for line in in_handle:
            if line.startswith(">"):
                if hairpin[name]:
                    hairpin[name] = hairpin[name] + "NNNNNNNNNNNN"
                name = line.strip().replace(">", " ").split()[0]
            else:
                hairpin[name] += line.strip()
        hairpin[name] = hairpin[name] + "NNNNNNNNNNNN"
    return hairpin

def _read_gtf(gtf):
    """
    Load GTF file with precursor positions on genome
    """
    if not gtf:
        return gtf
    db = defaultdict(list)
    with open(gtf) as in_handle:
        for line in in_handle:
            if line.startswith("#"):
                continue
            cols = line.strip().split("\t")
            name = [n.split("=")[1] for n in cols[-1].split(";") if n.startswith("Name")]
            chrom, start, end, strand = cols[0], cols[3], cols[4], cols[6]
            if cols[2] == "miRNA_primary_transcript":
                db[name[0]].append([chrom, int(start), int(end), strand])
    return db

def _coord(sequence, start, mirna, precursor, iso):
    """
    Define t5 and t3 isomirs
    """
    insertion = 0
    if iso.subs:
        insertion = 1 if iso.subs[0][-1] == "-" else 0
    end = start + (iso.end - len(iso.add) - insertion) - 1
    logger.debug("coor:s:%s len:%s flen:%s end:%s mirna:%s iso:%s" % (start, len(sequence), iso.end, end, mirna, iso.format()))
    dif = abs(mirna[0] - start)
    if start < mirna[0]:
        iso.t5 = sequence[:dif].upper()
    elif start > mirna[0]:
        iso.t5 = precursor[mirna[0] - 1:mirna[0] - 1 + dif].lower()
    elif start == mirna[0]:
        iso.t5 = "NA"
    if dif > 4:
        logger.debug("start > 3 %s %s %s %s %s" % (start, len(sequence), dif, mirna, iso.format()))
        return None

    dif = abs(mirna[1] - end)
    if iso.add:
        iso.add = iso.add.replace("-", "")
        sequence = sequence[:-len(iso.add)]
    # if dif > 3:
    #    return None
    if end > mirna[1]:
        iso.t3 = sequence[-dif:].upper()
    elif end < mirna[1]:
        iso.t3 = precursor[mirna[1] - dif:mirna[1]].lower()
    elif end == mirna[1]:
        iso.t3 = "NA"
    if dif > 4:
        logger.debug("end > 3 %s %s %s %s %s" % (len(sequence), end, dif, mirna, iso.format()))
        return None
    # logger.debug("coor end:%s %s %s %s %s iso:%s" % (start, len(sequence), end, dif, mirna, iso.format()))
    return True

def _annotate(reads, mirbase_ref, precursors):
    """
    Using SAM/BAM coordinates, mismatches and realign to annotate isomiRs
    """
    for r in reads:
        for p in reads[r].precursors:
            start = reads[r].precursors[p].start + 1  # convert to 1base
            end = start + len(reads[r].sequence)
            for mature in mirbase_ref[p]:
                mi = mirbase_ref[p][mature]
                is_iso = _coord(reads[r].sequence, start, mi, precursors[p], reads[r].precursors[p])
                logger.debug(("read:{s} pre:{p} start:{start} is_iso:{is_iso} mir:{mature} mir_p:{mi} mir_s:{mature_s}").format(s=reads[r].sequence, mature_s=precursors[p][mi[0]-1:mi[1]], **locals()))
                logger.debug("annotation:%s iso:%s" % (r, reads[r].precursors[p].format()))
                if is_iso:
                    reads[r].precursors[p].mirna = mature
                    break
    return reads

def _realign(seq, precursor, start, cigar):
    """
    The actual fn that will realign the sequence
    """
    seq, mature = cigar_correction(cigar, seq, precursor[start:])
    if seq.startswith("-"):
        seq = seq[1:]
    # print [cigar, seq, mature]
    logger.debug("Alignment: %s %s %s" % (cigar, seq, mature))
    error = set()
    pattern_addition = [[1, 1, 0], [1, 0, 1], [0, 1, 0], [0, 1, 1], [0, 0, 1], [1, 1, 1]]
    for pos in range(0, len(seq)):
        if seq[pos] != mature[pos]:
            error.add(pos)

    subs, add = [], []
    for e in error:
        if e < len(seq) - 3:
            subs.append([e, seq[e], mature[e]])

    pattern, error_add = [], []
    for e in range(len(seq) - 3, len(seq)):
        if e in error:
            pattern.append(1)
            error_add.append(e)
        else:
            pattern.append(0)
    for p in pattern_addition:
        if pattern == p:
            add = seq[error_add[0]:]
            break
    if not add and error_add:
        for e in error_add:
            subs.append([e, seq[e], mature[e]])

    return subs, add, len(mature)

def _clean_hits(reads):
    """
    Select only best matches
    """
    new_reads = defaultdict(realign)
    for r in reads:
        world = {}
        sc = 0
        for p in reads[r].precursors:
            world[p] = reads[r].precursors[p].get_score(len(reads[r].sequence))
            if sc < world[p]:
                sc = world[p]
        new_reads[r] = reads[r]
        for p in world:
            logger.debug("score %s %s %s" % (r, p, world[p]))
            if sc != world[p]:
                logger.debug("remove %s %s %s" % (r, p, world[p]))
                new_reads[r].remove_precursor(p)

    return new_reads

def _sam_to_bam(bam_fn):
    if not bam_fn.endswith("bam"):
        bam_out = "%s.bam" % os.path.splitext(bam_fn)[0]
        cmd = "samtools view -Sbh {bam_fn} -o {bam_out}"
        do.run(cmd.format(**locals()))
        return bam_out
    return bam_fn

def _bam_sort(bam_fn):
    bam_sort_by_n = op.splitext(bam_fn)[0] + "_sort.bam"
    if not file_exists(bam_sort_by_n):
        do.run(("samtools sort -n -o {bam_sort_by_n} {bam_fn}").format(**locals()))
    return bam_sort_by_n

def _read_bam(bam_fn, precursors):
    """
    read bam file and perform realignment of hits
    """
    mode = "r" if bam_fn.endswith("sam") else "rb"
    handle = pysam.Samfile(bam_fn, mode)
    reads = defaultdict(realign)
    for line in handle:
        if line.reference_id < 0:
            logger.debug("Sequence not mapped: %s" % line.reference_id)
            continue
        query_name = line.query_name
        if query_name not in reads:
            reads[query_name].sequence = line.query_sequence
        if line.is_reverse:
            logger.debug("Sequence is reverse: %s" % line.query_name)
            continue
        chrom = handle.getrname(line.reference_id)
        #  print "%s %s %s %s" % (line.query_name, line.reference_start, line.query_sequence, chrom)
        cigar = line.cigartuples
        iso = isomir()
        iso.align = line
        iso.start = line.reference_start
        if len(precursors[chrom]) < line.reference_start + len(reads[query_name].sequence):
            continue
        iso.subs, iso.add, iso.end = _realign(reads[query_name].sequence, precursors[chrom], line.reference_start, cigar)
        if len(iso.subs) < 2:
            reads[query_name].set_precursor(chrom, iso)

    reads = _clean_hits(reads)
    return reads

def _parse_mut(subs):
    """
    Parse mutation tag from miraligner output
    """
    if subs!="0":
        subs = [[subs.replace(subs[-2:], ""),subs[-2], subs[-1]]]
    return subs

def _read_miraligner(fn):
    """Read ouput of miraligner and create compatible output."""
    reads = defaultdict(realign)
    with open(fn) as in_handle:
        in_handle.next()
        for line in in_handle:
            cols = line.strip().split("\t")
            iso = isomir()
            query_name, seq = cols[1], cols[0]
            chrom, reference_start = cols[-2], cols[3]
            iso.mirna = cols[3]
            subs, add, iso.t5, iso.t3 = cols[6:10]
            if query_name not in reads:
                reads[query_name].sequence = seq
            iso.align = line
            iso.start = reference_start
            iso.subs, iso.add = _parse_mut(subs), add
            logger.debug("%s %s %s %s %s" % (query_name, reference_start, chrom, iso.subs, iso.add))
            reads[query_name].set_precursor(chrom, iso)
    return reads

def _cmd_miraligner(fn, out_file, species, hairpin):
    """
    Run miraligner for miRNA annotation
    """
    tool = _get_miraligner()
    path_db = op.dirname(op.abspath(hairpin))
    opts = "-Xms750m -Xmx4g"
    cmd = "{tool} -freq -i {fn} -o {out_file} -s {species} -db {path_db} -sub 1 -trim 3 -add 3"
    if not file_exists(out_file):
        do.run(cmd.format(**locals()), "miraligner with %s" % fn)
        shutil.move(out_file + ".mirna", out_file)
    return out_file

def _get_freq(name):
    """
    Check if name read contains counts (_xNumber)
    """
    try:
        counts = int(name.split("_x")[1])
    except:
        return 0
    return counts

def _tab_output(reads, out_file, sample):
    """Create info matrix"""
    seen = set()
    lines = []
    lines_pre = []
    seen_ann = {}
    dt = None
    with open(out_file, 'w') as out_handle:
        print >>out_handle, "seq\tname\tfreq\tchrom\tstart\tend\tmism\tadd\tt5\tt3\ts5\ts3\tDB\tprecursor\tambiguity\tName"
        for r, read in reads.iteritems():
            hits = set()
            [hits.add(mature.mirna) for mature in read.precursors.values() if mature.mirna]
            hits = len(hits)
            for p, iso in read.precursors.iteritems():
                if len(iso.subs) > 3 or not iso.mirna:
                    continue
                if (r, iso.mirna) not in seen:
                    seen.add((r, iso.mirna))
                    chrom = iso.mirna
                    if not chrom:
                        chrom = p
                    count = _get_freq(r)
                    seq = reads[r].sequence
                    if iso.get_score(len(seq)) < 1:
                        continue
                    if iso.subs:
                        iso.subs = [] if "N" in iso.subs[0] else iso.subs
                    annotation = "%s.%s" % (chrom, iso.format_id(sep="."))
                    idname = ("{0}.{1}").format(chrom, iso.format_id(sep="."))
                    res = ("{seq}\t{r}\t{count}\t{chrom}\tNA\tNA\t{format}\tNA\tNA\tmiRNA\t{p}\t{hits}\t{idname}").format(format=iso.format().replace("NA", "0"), **locals())
                    if annotation in seen_ann and seq.find("N") < 0 and seen_ann[annotation].split("\t")[0].find("N") < 0:
                        logger.warning("Same isomir %s from different sequence: \n%s and \n%s" % (annotation, res, seen_ann[annotation]))
                    seen_ann[annotation] = res
                    lines.append([annotation, chrom, count, sample, hits])
                    lines_pre.append([annotation, chrom, p, count, sample, hits])
                    print >>out_handle, res

    if lines:
        dt = pd.DataFrame(lines)
        dt.columns = ["isomir", "chrom", "counts", "sample", "hits"]
        dt = dt[dt['hits']>0]
        dt = dt.loc[:, "isomir":"sample"]
        dt = dt.groupby(['isomir', 'chrom', 'sample'], as_index=False).sum()
        dt.to_csv(out_file + "_summary")
        dt_pre = pd.DataFrame(lines_pre)
        dt_pre.columns = ["isomir", "mature", "chrom", "counts", "sample", "hits"]
        dt_pre = dt_pre[dt_pre['hits']==1]
        dt_pre = dt_pre.loc[:, "isomir":"sample"]
        dt_pre = dt_pre.groupby(['isomir', 'chrom', 'mature', 'sample'], as_index=False).sum()
        return out_file, dt, dt_pre
    return None

def _merge(dts):
    """
    merge multiple samples in one matrix
    """
    df= pd.concat(dts)

    ma = df.pivot(index='isomir', columns='sample', values='counts')
    ma_mirna = ma
    ma = ma.fillna(0)
    ma_mirna['mirna'] = [m.split(":")[0] for m in ma.index.values]
    ma_mirna = ma_mirna.groupby(['mirna']).sum()
    ma_mirna = ma_mirna.fillna(0)
    return ma, ma_mirna

def _create_counts(out_dts, out_dir):
    """Summarize results into single files."""
    ma, ma_mirna = _merge(out_dts)
    out_ma = op.join(out_dir, "counts.tsv")
    out_ma_mirna = op.join(out_dir, "counts_mirna.tsv")
    ma.to_csv(out_ma, sep="\t")
    ma_mirna.to_csv(out_ma_mirna, sep="\t")
    return out_ma_mirna, out_ma

def miraligner(args):
    """
    Realign BAM hits to miRBAse to get better accuracy and annotation
    """
    hairpin, mirna = _download_mirbase(args)
    precursors = _read_precursor(args.hairpin, args.sps)
    matures = _read_mature(args.mirna, args.sps)
    gtf = _read_gtf(args.gtf)
    out_dts = []
    for bam_fn in args.files:
        sample = op.splitext(op.basename(bam_fn))[0]
        if bam_fn.endswith("bam") or bam_fn.endswith("sam"):
            logger.info("Reading %s" % bam_fn)
            bam_fn = _sam_to_bam(bam_fn)
            bam_sort_by_n = _bam_sort(bam_fn)
            reads = _read_bam(bam_sort_by_n, precursors)
        elif bam_fn.endswith("fasta") or bam_fn.endswith("fa") or bam_fn.endswith("fastq"):
            out_file = op.join(args.out, sample + ".premirna")
            bam_fn = _filter_seqs(bam_fn)
            if args.miraligner:
                _cmd_miraligner(bam_fn, out_file, args.sps, args.hairpin)
                reads = _read_miraligner(out_file)
            else:
                logger.error("No other aligner supported")
                raise ValueError("Specify --miraligner or give BAM files")
        else:
            raise ValueError("Format not recognized.")

        if not args.miraligner:
            reads = _annotate(reads, matures, precursors)
        out_file = op.join(args.out, sample + ".mirna")
        out_file, dt, dt_pre= _tab_output(reads, out_file, sample)
        try:
            vcf_file = op.join(args.out, sample + ".vcf")
            if not file_exists(vcf_file):
            # if True:
                create_vcf(dt_pre, matures, gtf, vcf_file)
            try:
                import vcf
                vcf.Reader(filename=vcf_file)
            except Exception as e:
                logger.warning(e.__doc__)
                logger.warning(e.message)
        except Exception as e:
            # traceback.print_exc()
            logger.warning(e.__doc__)
            logger.warning(e.message)
        if isinstance(dt, pd.DataFrame):
            out_dts.append(dt)

    if out_dts:
        _create_counts(out_dts, args.out)
        # _summarize(out_dts)
    else:
        print "No files analyzed!"
