"""
Global pairwise alignment with affine gap penalties using BioPython's
C-backed PairwiseAligner.

Parameters match the original Oncosplice convention:
    match=1, mismatch=-1, gap_open=-3, gap_extend=0

This means opening a gap is expensive but extending is free, so the
aligner prefers fewer, longer gaps over many short ones.
"""

from collections import namedtuple
from Bio.Align import PairwiseAligner

Alignment = namedtuple('Alignment', ['seqA', 'seqB', 'score'])


def _extract_aligned_strings(alignment) -> tuple[str, str]:
    """Extract gapped alignment strings from Bio.Align.Alignment coordinates."""
    coords = alignment.coordinates
    target_seq = str(alignment.target)
    query_seq = str(alignment.query)

    aligned_a = []
    aligned_b = []

    for k in range(coords.shape[1] - 1):
        t_start, t_end = coords[0, k], coords[0, k + 1]
        q_start, q_end = coords[1, k], coords[1, k + 1]

        t_len = t_end - t_start
        q_len = q_end - q_start

        if t_len > 0 and q_len > 0:
            aligned_a.append(target_seq[t_start:t_end])
            aligned_b.append(query_seq[q_start:q_end])
        elif t_len > 0:
            aligned_a.append(target_seq[t_start:t_end])
            aligned_b.append('-' * t_len)
        elif q_len > 0:
            aligned_a.append('-' * q_len)
            aligned_b.append(query_seq[q_start:q_end])

    return ''.join(aligned_a), ''.join(aligned_b)


def _trim_common(seq_a: str, seq_b: str) -> tuple[str, str, str, str, str]:
    """Strip identical prefix and suffix shared by two sequences.

    Returns (prefix, mid_a, mid_b, suffix) where:
        seq_a == prefix + mid_a + suffix
        seq_b == prefix + mid_b + suffix
    """
    min_len = min(len(seq_a), len(seq_b))

    # common prefix
    pre = 0
    while pre < min_len and seq_a[pre] == seq_b[pre]:
        pre += 1

    # common suffix (don't overlap with prefix)
    suf = 0
    while suf < (min_len - pre) and seq_a[-(suf + 1)] == seq_b[-(suf + 1)]:
        suf += 1

    end_a = len(seq_a) - suf if suf else len(seq_a)
    end_b = len(seq_b) - suf if suf else len(seq_b)

    return (
        seq_a[:pre],
        seq_a[pre:end_a],
        seq_b[pre:end_b],
        seq_a[end_a:] if suf else "",
    )


def global_align(seq_a: str, seq_b: str,
                 match: float = 1, mismatch: float = -1,
                 gap_open: float = -3, gap_extend: float = 0,
                 penalize_end_gaps: bool = True) -> list[Alignment]:
    """
    Global pairwise alignment with affine gap penalties.

    Identical prefix/suffix regions are trimmed before alignment and
    stitched back, reducing O(n²) to O(d²) where d is the length of the
    differing region.

    Returns
    -------
    list[Alignment]
        Single-element list with the optimal alignment.
    """
    # Fast path: identical sequences
    if seq_a == seq_b:
        score = len(seq_a) * match
        return [Alignment(seq_a, seq_b, score)]

    # Trim shared prefix/suffix so we only align the differing middle
    prefix, mid_a, mid_b, suffix = _trim_common(seq_a, seq_b)

    if not mid_a and not mid_b:
        # Shouldn't happen (caught by == above), but be safe
        score = len(seq_a) * match
        return [Alignment(seq_a, seq_b, score)]

    # Pure insertion or deletion: one middle is empty, no alignment needed
    if not mid_a or not mid_b:
        gap_len = len(mid_a) or len(mid_b)
        al_a = mid_a if mid_a else '-' * len(mid_b)
        al_b = mid_b if mid_b else '-' * len(mid_a)
        mid_score = gap_open + gap_extend * (gap_len - 1) if gap_len > 0 else 0
        flanking_score = (len(prefix) + len(suffix)) * match
        return [Alignment(prefix + al_a + suffix, prefix + al_b + suffix, flanking_score + mid_score)]

    aligner = PairwiseAligner()
    aligner.mode = 'global'
    aligner.match_score = match
    aligner.mismatch_score = mismatch
    aligner.open_gap_score = gap_open
    aligner.extend_gap_score = gap_extend

    if not penalize_end_gaps:
        aligner.target_end_gap_score = 0.0
        aligner.query_end_gap_score = 0.0

    al = aligner.align(mid_a, mid_b)[0]
    al_a, al_b = _extract_aligned_strings(al)

    # Stitch: prefix (matched 1:1) + aligned middle + suffix (matched 1:1)
    full_a = prefix + al_a + suffix
    full_b = prefix + al_b + suffix

    # Score: prefix matches + middle alignment score + suffix matches
    flanking_score = (len(prefix) + len(suffix)) * match
    total_score = flanking_score + al.score

    return [Alignment(full_a, full_b, total_score)]
