#! python

from collections import defaultdict
from .aln.utils import get_query_start_end_in_forward_read

"""
Find and tag inversion alignments and flanking alignments before
and after inversions.
    iv:A:B --> flank alignments before an inversion
    iv:A:I --> inversion alignment
    iv:A:A --> flank alignments after an inversion
"""

# Define coords
IDX, ISREVERSE, REF_NAME, REF_START, REF_END, READ_START, READ_END = range(0,7) # COORD COLUMNS

# inversion tag
INVERSION_TAG_NAME = 'iv'
INVERSION_TAG_TYPE = 'A'
FLANK_BEFORE_TAG_VALUE = 'B'
INVERSION_TAG_VALUE = 'I'
FLANK_AFTER_TAG_VALUE = 'A'


def tag_inversions(alns, inversion_poswiggle):
    """
    Find and tag inversions and flanking alignments
    """
    # require a minimum of three alignments to detect inversion
    has_inversion = False
    for sorted_alns in yield_sorted_alnobjs_of_a_reference(alns, min_groupsize=3):
        has_inversion = has_inversion or tag_inversions_of_sorted_alns(sorted_alns, inversion_poswiggle)
    return has_inversion


def tag_inversions_of_sorted_alns(sorted_alns, inversion_poswiggle):
    """
    Input sorted_alns are a list of sorted mapped alignments with the same
    query name and reference id, sorted by reference start.
    Find inversions and flanking alignments in inputs and set tags to
    inversion alignments as well as their flanking alignments.
    """
    assert sorted_alns >= 3 # assert minimum grouped alignments size >= 3, OK to disable
    assert_alns_mapped_grouped_sorted(sorted_alns) # sanity check, OK to disable

    # convert alignments to coordinates of alignments, attach aln index to coord for back tracking
    coords = [aln_to_coord_tuple(aln, idx) for (idx, aln) in enumerate(sorted_alns)]

    def _set_tags(coord_before, coord_middle, coord_after):
        set_tags_for_inv_flank_alns(sorted_alns[coord_tuple_to_index(coord_before)],
                                    sorted_alns[coord_tuple_to_index(coord_middle)],
                                    sorted_alns[coord_tuple_to_index(coord_after)])

    return find_inversion(coords, _set_tags, inversion_poswiggle)


def yield_sorted_alnobjs_of_a_reference(alnobjs, min_groupsize):
    """
    All alnobjs must be mapped or unmapped alignments of a read.
    Each time yield a list of sorted alignments of a group, where
        * unmmapped alignments are dropped
        * mapped alignments are grouped by reference_id
        * grouped alignments are sorted by reference start.
        * group which contain less than min_groupsize alignments will be droped
    """
    reftoreads = defaultdict(list)
    for aln in alnobjs:
        if not aln.is_unmapped and not aln.is_secondary:
            reftoreads[aln.reference_id].append(aln)

    for reference_id in reftoreads:
        # get alignments mapped to this reference, sort by reference start positions
        sorted_alns = sorted(reftoreads[reference_id], key=lambda r:r.reference_start)
        if not all([aln.reference_id==reference_id for aln in sorted_alns]):
            raise ValueError("Not filtered")
        if len(sorted_alns) >= min_groupsize:
            yield sorted_alns


def assert_alns_mapped_grouped_sorted(sorted_alns):
    """Assert sorted_alns are all mapped alignments, have the same query and reference
    and sorted by reference start."""
    if len(sorted_alns) == 0:
        return
    query_name, reference_id = sorted_alns[0].query_name, sorted_alns[0].reference_id
    for idx, aln in enumerate(sorted_alns):
        if idx == 0:
            continue
        if aln.is_unmapped:
            raise ValueError("Input sorted alignments must all be mapped.")
        if not (aln.reference_id == reference_id and
                aln.query_name == query_name and
                aln.reference_start >= sorted_alns[idx-1].reference_start):
            raise ValueError("Input sorted alignments must have the same query and reference name, sorted by reference start.")


def is_inv_aln(aln):
    """Return if an alignment is mapped and is an inversion."""
    return ((not aln.is_unmapped) and aln.has_tag(INVERSION_TAG_NAME) and
            aln.get_tag(tag=INVERSION_TAG_NAME) == INVERSION_TAG_VALUE)


def is_inv_or_flank_aln(aln):
    """Return True if aln is an inversion or flanking alignments of an inversion."""
    return ((not aln.is_unmapped) and aln.has_tag(INVERSION_TAG_NAME) and
            aln.get_tag(INVERSION_TAG_NAME) in [FLANK_BEFORE_TAG_VALUE, INVERSION_TAG_VALUE, FLANK_AFTER_TAG_VALUE])


def are_both_inv_or_flank_alns(a0, a1):
    """Return True if both a0 and a1 are either inversion or flanking alignments."""
    return is_inv_or_flank_aln(a0) and is_inv_or_flank_aln(a1)


def set_tags_for_inv_flank_alns(aln_before, aln_middle, aln_after):
    """Set tags for inversion alignment and flanking alignments before and after inversion.
    """
    aln_before.set_tag(INVERSION_TAG_NAME, FLANK_BEFORE_TAG_VALUE, value_type=INVERSION_TAG_TYPE)
    aln_middle.set_tag(INVERSION_TAG_NAME, INVERSION_TAG_VALUE, value_type=INVERSION_TAG_TYPE)
    aln_after.set_tag(INVERSION_TAG_NAME, FLANK_AFTER_TAG_VALUE, value_type=INVERSION_TAG_TYPE)


def coord_tuple_to_index(coord):
    return coord[IDX]


def aln_to_coord_tuple(aln, aln_id):
    """Convert an AlignedSegment object to a coordinate tuple,
    Each coordinate tuple has six fields:
    (IDX, ISREVERSE, REF_START, REF_END, READ_START, READ_END)
    """
    return ((aln_id, aln.is_reverse, aln.reference_name, aln.reference_start, aln.reference_end) +
            get_query_start_end_in_forward_read(aln))


def _is_middle_coord_tuple_inversion(coord_tuples, inversion_poswiggle):
    """
    Return True if c1 is an inversion compared with c0 and c2.

    coord_tuples = (c0, c1, c2) where each c* is a tuple
    (idx, aln.is_reverse, aln.reference_start, aln.reference_end, query_start, query_end)
    , where forward_read[query_start: query_end] is aligned with reference.
    """
    assert len(coord_tuples) == 3
    c0, c1, c2 = coord_tuples
    cp0, cp1, cp2 = list(sorted(coord_tuples, key=lambda c:c[READ_START]))
    return (c0[ISREVERSE] == c2[ISREVERSE] and c0[ISREVERSE] != c1[ISREVERSE] and
            max(abs(c0[REF_END] - c1[REF_START]), abs(c1[REF_END] - c2[REF_START])) <= inversion_poswiggle and
            max(abs(cp0[READ_END] - cp1[READ_START]), abs(cp1[READ_END] - cp2[READ_START])) <= inversion_poswiggle)


def assert_coords_sorted_by_ref_start(coords):
    """Assert coords are sorted by reference start."""
    for idx, c in enumerate(coords):
        if idx == 0:
            continue
        if c[REF_START] < coords[idx-1][REF_START]:
            raise ValueError("Coords must be sorted by reference start while coords[{}]={} violates the rule.".format(idx, c))


def find_inversion(coords, apply_to_coords_f, inversion_poswiggle):
    """
    Search flanking alignments of coord and look for inversions.

            left2 left1 coord right1 right2

    coords --- coordinate tuples of alignments sorted by reference start.
    apply_to_coords_f --- function apply_to_coords_f(coord_before, coord_middle, coord_after) takes coords of
    flanking alignments before an inversion, the inversion and flanking alignments after
    the inversion as input.
    inversion_poswiggle --- maximum position wiggle in base pairs between inversion alignment and
        flanking alignments on either ends.

   This function will search for simple inversions:
                      Inv
    REF  >>>>>>>>> >>>>>>>>> >>>>>>>>>
    READ >>>>>>>>> <<<<<<<<< >>>>>>>>>
    READ <<<<<<<<< >>>>>>>>> <<<<<<<<<

    This function will *ONLY* report *ONE* out of multiple conseutive inversions, *NOT* all inversions
                      Inv                 Inv
    REF  >>>>>>>>> >>>>>>>>> >>>>>>>>> >>>>>>>>> >>>>>>>>>
    READ >>>>>>>>> <<<<<<<<< >>>>>>>>> <<<<<<<<< >>>>>>>>>
    READ <<<<<<<<< >>>>>>>>> <<<<<<<<< >>>>>>>>> >>>>>>>>>

    ...doctest:
        >>> def f(c0, c1, c2): print (c1)
        >>> coords = [(1, False, 'chr', 0, 100, 0, 100), (1, True, 'chr', 100, 200, 90, 210), (2, False, 'chr', 190, 310, 210, 290)]
        >>> find_inversion(coords, f, 10)
        (1, True, 'chr', 100, 200, 90, 210)
        True
        >>> find_inversion(coords, f, 9)
        False
    """
    if len(coords) <= 2:
        return False
    assert_coords_sorted_by_ref_start(coords) # sanity check, can be disabled

    has_inversion = False
    _idx = 2
    while _idx < len(coords):
        if _is_middle_coord_tuple_inversion(coords[_idx-2:_idx+1], inversion_poswiggle):
            has_inversion = True
            apply_to_coords_f(coords[_idx-2], coords[_idx-1], coords[_idx])
            break
        _idx += 1
    return has_inversion
