"""
Needleman-Wunsch global alignment with affine gap penalties (Gotoh algorithm).

Drop-in replacement for Bio.pairwise2.align.globalms used by Oncosplice.
"""

from collections import namedtuple
import numpy as np

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

# Traceback pointer constants
_DIAG = 0
_UP = 1    # gap in seq_b (consume seq_a)
_LEFT = 2  # gap in seq_a (consume seq_b)


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.

    Parameters
    ----------
    seq_a, seq_b : str
        Sequences to align.
    match : float
        Score for a matching pair.
    mismatch : float
        Score for a mismatching pair.
    gap_open : float
        Penalty for opening a gap (negative).
    gap_extend : float
        Penalty for extending a gap (negative or zero).
    penalize_end_gaps : bool
        If True, leading/trailing gaps are penalized (same as interior gaps).

    Returns
    -------
    list[Alignment]
        Single-element list with the optimal alignment.
    """
    n, m = len(seq_a), len(seq_b)

    NEG_INF = -1e9

    # M[i][j]  = best score ending with seq_a[i-1] aligned to seq_b[j-1]
    # Ix[i][j] = best score ending with a gap in seq_b (consuming seq_a[i-1])
    # Iy[i][j] = best score ending with a gap in seq_a (consuming seq_b[j-1])
    M = np.full((n + 1, m + 1), NEG_INF)
    Ix = np.full((n + 1, m + 1), NEG_INF)
    Iy = np.full((n + 1, m + 1), NEG_INF)

    # Traceback matrices: store which state we came from
    tb_M = np.zeros((n + 1, m + 1), dtype=np.int8)
    tb_Ix = np.zeros((n + 1, m + 1), dtype=np.int8)
    tb_Iy = np.zeros((n + 1, m + 1), dtype=np.int8)

    # Initialization
    M[0, 0] = 0.0

    gap_penalty = gap_open if penalize_end_gaps else 0.0
    ext_penalty = gap_extend if penalize_end_gaps else 0.0

    for i in range(1, n + 1):
        Ix[i, 0] = gap_penalty + ext_penalty * (i - 1)
    for j in range(1, m + 1):
        Iy[0, j] = gap_penalty + ext_penalty * (j - 1)

    # Fill
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            s = match if seq_a[i - 1] == seq_b[j - 1] else mismatch

            # M[i,j]: align seq_a[i-1] with seq_b[j-1]
            candidates_m = (M[i - 1, j - 1] + s,
                            Ix[i - 1, j - 1] + s,
                            Iy[i - 1, j - 1] + s)
            best_m = int(np.argmax(candidates_m))
            M[i, j] = candidates_m[best_m]
            tb_M[i, j] = best_m  # 0=from M, 1=from Ix, 2=from Iy

            # Ix[i,j]: gap in seq_b (consume seq_a[i-1], put '-' in seq_b)
            candidates_ix = (M[i - 1, j] + gap_open,
                             Ix[i - 1, j] + gap_extend)
            best_ix = int(np.argmax(candidates_ix))
            Ix[i, j] = candidates_ix[best_ix]
            tb_Ix[i, j] = best_ix  # 0=opened from M, 1=extended from Ix

            # Iy[i,j]: gap in seq_a (consume seq_b[j-1], put '-' in seq_a)
            candidates_iy = (M[i, j - 1] + gap_open,
                             Iy[i, j - 1] + gap_extend)
            best_iy = int(np.argmax(candidates_iy))
            Iy[i, j] = candidates_iy[best_iy]
            tb_Iy[i, j] = best_iy  # 0=opened from M, 1=extended from Iy

    # Find best score in bottom-right cell across all states
    final_scores = [M[n, m], Ix[n, m], Iy[n, m]]
    best_state = int(np.argmax(final_scores))
    best_score = final_scores[best_state]

    # Traceback
    aligned_a = []
    aligned_b = []
    i, j = n, m
    state = best_state  # 0=M, 1=Ix, 2=Iy

    while i > 0 or j > 0:
        if state == 0:  # M state: diagonal move
            if i == 0 or j == 0:
                break
            aligned_a.append(seq_a[i - 1])
            aligned_b.append(seq_b[j - 1])
            prev = tb_M[i, j]
            i -= 1
            j -= 1
            state = prev  # 0=M, 1=Ix, 2=Iy
        elif state == 1:  # Ix state: gap in seq_b
            if i == 0:
                break
            aligned_a.append(seq_a[i - 1])
            aligned_b.append('-')
            prev = tb_Ix[i, j]
            i -= 1
            state = 0 if prev == 0 else 1  # 0=came from M, 1=stayed in Ix
        elif state == 2:  # Iy state: gap in seq_a
            if j == 0:
                break
            aligned_a.append('-')
            aligned_b.append(seq_b[j - 1])
            prev = tb_Iy[i, j]
            j -= 1
            state = 0 if prev == 0 else 2  # 0=came from M, 1=stayed in Iy

    # Handle remaining characters (leading gaps)
    while i > 0:
        aligned_a.append(seq_a[i - 1])
        aligned_b.append('-')
        i -= 1
    while j > 0:
        aligned_a.append('-')
        aligned_b.append(seq_b[j - 1])
        j -= 1

    aligned_a.reverse()
    aligned_b.reverse()

    return [Alignment(''.join(aligned_a), ''.join(aligned_b), best_score)]
