## This code is written by Davide Albanese, <davide.albanese@gmail.com>
## Copyright (C) 2015-2016 Fondazione Edmund Mach
## Copyright (C) 2015-2016 Davide Albanese

## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.

## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.

## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.



import subprocess
import os
import os.path
import re

import numpy as np
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord

from strainest import mummer_path
import strainest.mummer


nucmer_path = 'nucmer'
show_aligns_path = 'show-aligns'


def mapgenome(genome_fn, reference_fn, delta_prefix):

    NO_ALIGN_STR = "ERROR: Could not find any alignments for"

    # run nucmer
    cmd = [nucmer_path, '--prefix', delta_prefix, reference_fn, genome_fn]
    devnull = open(os.devnull, 'w')
    proc = subprocess.Popen(cmd, stdout=devnull, stderr=subprocess.PIPE)
    _, out_stderr = proc.communicate()
    devnull.close()
    if proc.returncode:
        raise Exception(out_stderr)

    delta_fn = delta_prefix + '.delta'

    # get the first reference sequence only
    ref_record = next(SeqIO.parse(reference_fn, 'fasta'))
    ref_seqid = ref_record.id
    ref_len = len(ref_record)

    genome_seq_mapped = np.asarray(['.'] * ref_len)
    genome_count_mapped = np.zeros(ref_len, dtype=np.int)
    for record in SeqIO.parse(genome_fn, 'fasta'):
        genome_seqid = record.id

        # run show-aligns
        align_fn = delta_prefix + '.align'
        with open(align_fn, 'wb') as align_handler:
            cmd = [show_aligns_path, delta_fn, ref_seqid, genome_seqid]
            proc = subprocess.Popen(cmd, stdout=align_handler,
                                    stderr=subprocess.PIPE)
            _, out_stderr = proc.communicate()

        if proc.returncode:
            if out_stderr.startswith(NO_ALIGN_STR):
                os.remove(align_fn)
                continue
            else:
                raise Exception(out_stderr)

        # load alignments
        with open(align_fn, 'rb') as align_handler:
            mar = strainest.mummer.MummerAlignmentReader(align_handler)
            for al in mar:
                ref_seq = np.asarray(list(al.seq1))
                genome_seq = np.asarray(list(al.seq2))
                genome_seq = genome_seq[np.where(ref_seq != '.')]
                genome_seq_mapped[al.start1-1:al.end1] = genome_seq
                genome_count_mapped[al.start1-1:al.end1] += 1

        os.remove(align_fn)

    os.remove(delta_fn)

    # remove repeats
    genome_seq_mapped[genome_count_mapped > 1] = '.'

    # sub '.' with '-'
    genome_seq_mapped[genome_seq_mapped == '.']  = '-'

    # write mapped genome
    out_seq = Seq(''.join(genome_seq_mapped)).upper()
    out_id = re.sub('\s+', '_', os.path.basename(genome_fn))
    out_record = SeqRecord(out_seq, id=out_id, description='')

    return out_record
