//! Rust-accelerated core for the Needleman-Wunsch DP fill and traceback.
//!
//! Only the performance-critical, type-agnostic numeric work lives here.
//! Everything that touches generic Python objects (score matrix construction
//! from a callable, alignment pair assembly from sequences) stays in Python
//! because it runs in O(n+m) and cannot avoid per-element Python overhead
//! anyway.
//!
//! One Python entry point is exposed from this module:
//! - `nw_traceback_indices`: DP fill + traceback, returns index arrays
//!
//! Safe `Array2` indexing is used throughout.  Benchmarking showed that
//! `unsafe get_unchecked` offers no measurable improvement -- LLVM eliminates
//! the bounds checks in release mode.

use numpy::ndarray::{Array2, ArrayView2};
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray2};
use pyo3::prelude::*;

#[pymodule(name = "_pynw_core")]
mod pynw_core {
    use super::*;

    const GAP: isize = -1;

    #[derive(Clone, Copy)]
    enum Direction {
        Diagonal,
        Up,
        Left,
    }

    type NwTracebackIndicesOutput<'py> = (
        f64,
        Bound<'py, PyArray1<isize>>,
        Bound<'py, PyArray1<isize>>,
    );

    fn fill_dp_and_traceback(
        score_matrix: &ArrayView2<f64>,
        gap_penalty_seq1: f64,
        gap_penalty_seq2: f64,
    ) -> (Array2<f64>, Array2<Direction>) {
        // TODO: We could optimize the space complexity of this method by
        // calculating only the traceback using a rolling band of 2 rows.
        let (n, m) = (score_matrix.nrows(), score_matrix.ncols());
        let mut dp = Array2::<f64>::zeros((n + 1, m + 1));
        let mut tb = Array2::from_elem((n + 1, m + 1), Direction::Diagonal);

        for i in 1..=n {
            dp[[i, 0]] = i as f64 * gap_penalty_seq2;
            tb[[i, 0]] = Direction::Up;
        }
        for j in 1..=m {
            dp[[0, j]] = j as f64 * gap_penalty_seq1;
            tb[[0, j]] = Direction::Left;
        }

        for i in 1..=n {
            for j in 1..=m {
                let diagonal = dp[[i - 1, j - 1]] + score_matrix[[i - 1, j - 1]];
                let up = dp[[i - 1, j]] + gap_penalty_seq2;
                let left = dp[[i, j - 1]] + gap_penalty_seq1;

                // Precedence for tie-breaking: diagnonal > up > left
                let (best, direction) = if left > up && left > diagonal {
                    (left, Direction::Left)
                } else if up > diagonal {
                    (up, Direction::Up)
                } else {
                    (diagonal, Direction::Diagonal)
                };

                dp[[i, j]] = best;
                tb[[i, j]] = direction;
            }
        }

        (dp, tb)
    }

    fn traceback_indices(traceback_matrix: &ArrayView2<Direction>) -> (Vec<isize>, Vec<isize>) {
        let (n, m) = (traceback_matrix.nrows() - 1, traceback_matrix.ncols() - 1);

        let mut s1 = Vec::with_capacity(n + m);
        let mut s2 = Vec::with_capacity(n + m);

        let mut i = n;
        let mut j = m;

        while i > 0 || j > 0 {
            let dir = &traceback_matrix[[i, j]];
            match dir {
                Direction::Diagonal => {
                    i -= 1;
                    j -= 1;
                    s1.push(i as isize);
                    s2.push(j as isize);
                }
                Direction::Up => {
                    i -= 1;
                    s1.push(i as isize);
                    s2.push(GAP);
                }
                Direction::Left => {
                    j -= 1;
                    s1.push(GAP);
                    s2.push(j as isize);
                }
            }
        }

        s1.reverse();
        s2.reverse();

        (s1, s2)
    }

    /// Combined DP fill + traceback returning index arrays.
    ///
    /// Returns `(optimal_score, seq1_idx, seq2_idx)` where each index
    /// array is a 1-D `intp` numpy array.  A value of `-1` indicates
    /// a gap at that alignment position.
    ///
    /// This is the fast path for the public `needleman_wunsch` API:
    /// no intermediate Python objects are created.
    #[pyfunction]
    fn nw_traceback_indices<'py>(
        py: Python<'py>,
        score_matrix: PyReadonlyArray2<'py, f64>,
        gap_penalty_seq1: f64,
        gap_penalty_seq2: f64,
    ) -> PyResult<NwTracebackIndicesOutput<'py>> {
        let sm = score_matrix.as_array();
        let (n, m) = (sm.nrows(), sm.ncols());

        let (dp, tb) = fill_dp_and_traceback(&sm, gap_penalty_seq1, gap_penalty_seq2);
        let (s1, s2) = traceback_indices(&tb.view());

        let best_score = dp[[n, m]];
        Ok((best_score, s1.into_pyarray(py), s2.into_pyarray(py)))
    }
}
