//! LU Factorization with partial pivoting.
//!
//! Computes `P * A = L * U` where:
//! - P is a permutation matrix (represented by pivot indices)
//! - L is unit lower triangular
//! - U is upper triangular
//!
//! Uses a blocked right-looking algorithm:
//! 1. Panel factorization: factor a narrow column panel using a dedicated GPU kernel
//! 2. Apply pivots: swap rows in the trailing portion
//! 3. TRSM: solve for the upper triangle block
//! 4. GEMM: update the trailing submatrix
//!
//! The L and U factors overwrite the input matrix A in-place (LAPACK-style packed
//! storage with unit diagonal for L implicitly assumed).

use oxicuda_blas::types::{
    DiagType, FillMode, GpuFloat, Layout, MatrixDesc, MatrixDescMut, Side, Transpose,
};
use oxicuda_memory::DeviceBuffer;
use oxicuda_ptx::prelude::*;

use crate::error::{SolverError, SolverResult};
use crate::handle::SolverHandle;
use crate::ptx_helpers::SOLVER_BLOCK_SIZE;

/// Block size for the panel factorization step.
const LU_BLOCK_SIZE: u32 = 64;

// ---------------------------------------------------------------------------
// Result type
// ---------------------------------------------------------------------------

/// Result of an LU factorization.
///
/// Contains diagnostic information about the factorization.
#[derive(Debug, Clone)]
pub struct LuResult {
    /// Status info:
    /// - 0: successful factorization
    /// - i > 0: U(i,i) is exactly zero, matrix is singular at column i
    pub info: i32,
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Performs LU factorization with partial pivoting in-place.
///
/// On exit, the lower triangle of `a` (with implicit unit diagonal) contains L,
/// and the upper triangle contains U. The `pivots` array records the row
/// permutations: row `i` was interchanged with row `pivots[i]`.
///
/// The matrix is stored in column-major order with leading dimension `lda`.
///
/// # Arguments
///
/// * `handle` — solver handle.
/// * `a` — matrix buffer (n x n, column-major, lda stride), modified in-place.
/// * `n` — matrix dimension.
/// * `lda` — leading dimension (>= n).
/// * `pivots` — output pivot indices buffer (length >= n).
///
/// # Returns
///
/// [`LuResult`] with `info == 0` on success, `info > 0` if singular.
///
/// # Errors
///
/// Returns [`SolverError`] if dimensions are invalid or a kernel launch fails.
pub fn lu_factorize<T: GpuFloat>(
    handle: &mut SolverHandle,
    a: &mut DeviceBuffer<T>,
    n: u32,
    lda: u32,
    pivots: &mut DeviceBuffer<i32>,
) -> SolverResult<LuResult> {
    // Validate dimensions.
    if n == 0 {
        return Ok(LuResult { info: 0 });
    }
    if lda < n {
        return Err(SolverError::DimensionMismatch(format!(
            "lu_factorize: lda ({lda}) must be >= n ({n})"
        )));
    }
    let required = n as usize * lda as usize;
    if a.len() < required {
        return Err(SolverError::DimensionMismatch(format!(
            "lu_factorize: buffer too small ({} < {required})",
            a.len()
        )));
    }
    if pivots.len() < n as usize {
        return Err(SolverError::DimensionMismatch(format!(
            "lu_factorize: pivots buffer too small ({} < {n})",
            pivots.len()
        )));
    }

    // Ensure workspace is large enough for panel temporaries.
    let panel_workspace = n as usize * LU_BLOCK_SIZE as usize * T::SIZE;
    handle.ensure_workspace(panel_workspace)?;

    blocked_lu::<T>(handle, a, n, lda, pivots)
}

/// Solves `A * X = B` given an LU-factored matrix.
///
/// The LU factors must have been computed by [`lu_factorize`]. The solution
/// overwrites `b` in-place.
///
/// # Arguments
///
/// * `handle` — solver handle.
/// * `lu` — LU-factored matrix (output of `lu_factorize`).
/// * `pivots` — pivot indices from `lu_factorize`.
/// * `b` — right-hand side matrix (n x nrhs), overwritten with solution.
/// * `n` — matrix dimension.
/// * `nrhs` — number of right-hand side columns.
///
/// # Errors
///
/// Returns [`SolverError`] if dimensions are invalid or BLAS operations fail.
pub fn lu_solve<T: GpuFloat>(
    handle: &SolverHandle,
    lu: &DeviceBuffer<T>,
    pivots: &DeviceBuffer<i32>,
    b: &mut DeviceBuffer<T>,
    n: u32,
    nrhs: u32,
) -> SolverResult<()> {
    if n == 0 || nrhs == 0 {
        return Ok(());
    }
    if lu.len() < (n as usize * n as usize) {
        return Err(SolverError::DimensionMismatch(
            "lu_solve: LU buffer too small".into(),
        ));
    }
    if pivots.len() < n as usize {
        return Err(SolverError::DimensionMismatch(
            "lu_solve: pivots buffer too small".into(),
        ));
    }
    if b.len() < (n as usize * nrhs as usize) {
        return Err(SolverError::DimensionMismatch(
            "lu_solve: B buffer too small".into(),
        ));
    }

    // Step 1: Apply row permutations to B.
    // Each pivot[i] says row i was swapped with row pivot[i] during
    // factorization, so we replay the swaps in forward order.
    apply_pivots_to_rhs::<T>(handle, b, pivots, n, nrhs)?;

    // Step 2: Solve L * Y = P * B (forward substitution) via TRSM.
    let l_desc = MatrixDesc::<T>::from_raw(lu.as_device_ptr(), n, n, n, Layout::ColMajor);
    let mut b_desc = MatrixDescMut::<T>::from_raw(b.as_device_ptr(), n, nrhs, n, Layout::ColMajor);

    oxicuda_blas::level3::trsm(
        handle.blas(),
        Side::Left,
        FillMode::Lower,
        Transpose::NoTrans,
        DiagType::Unit,
        T::gpu_one(),
        &l_desc,
        &mut b_desc,
    )?;

    // Step 3: Solve U * X = Y (backward substitution) via TRSM.
    let u_desc = MatrixDesc::<T>::from_raw(lu.as_device_ptr(), n, n, n, Layout::ColMajor);

    oxicuda_blas::level3::trsm(
        handle.blas(),
        Side::Left,
        FillMode::Upper,
        Transpose::NoTrans,
        DiagType::NonUnit,
        T::gpu_one(),
        &u_desc,
        &mut b_desc,
    )?;

    Ok(())
}

/// Solves the transposed system `Aᵀ * X = B` given an LU-factored matrix.
///
/// The LU factors must have been produced by [`lu_factorize`], which computes
/// `P * A = L * U`. Hence `A = Pᵀ * L * U` and the transpose factorizes as
///
/// ```text
/// Aᵀ = Uᵀ * Lᵀ * P
/// ```
///
/// so `Aᵀ * X = B` is solved by the three stages
///
/// 1. `Uᵀ * Y = B` — forward substitution (`Uᵀ` is lower triangular, non-unit
///    diagonal),
/// 2. `Lᵀ * W = Y` — backward substitution (`Lᵀ` is upper triangular, unit
///    diagonal),
/// 3. `X = Pᵀ * W` — apply the row permutation transposed, i.e. replay the
///    pivot transpositions in reverse order.
///
/// The solution overwrites `b` in-place. The LU buffer is interpreted in
/// column-major order with leading dimension `n`, matching [`lu_solve`].
///
/// This is required by algorithms — such as Hager's 1-norm condition
/// estimator — that alternate solves with `A` and `Aᵀ`.
///
/// # Arguments
///
/// * `handle` — solver handle.
/// * `lu` — LU-factored matrix (output of `lu_factorize`).
/// * `pivots` — pivot indices from `lu_factorize`.
/// * `b` — right-hand side matrix (n x nrhs, column-major), overwritten with
///   the solution.
/// * `n` — matrix dimension.
/// * `nrhs` — number of right-hand side columns.
///
/// # Errors
///
/// Returns [`SolverError`] if dimensions are invalid, a pivot index is out of
/// range, or a host transfer fails.
pub fn lu_solve_transposed<T: GpuFloat>(
    handle: &SolverHandle,
    lu: &DeviceBuffer<T>,
    pivots: &DeviceBuffer<i32>,
    b: &mut DeviceBuffer<T>,
    n: u32,
    nrhs: u32,
) -> SolverResult<()> {
    lu_solve_with_transpose::<T>(handle, lu, pivots, b, n, nrhs, true)
}

/// Solves `A * X = B` or `Aᵀ * X = B` depending on `transpose`.
///
/// When `transpose` is `false` this is equivalent to [`lu_solve`]; when it is
/// `true` it is equivalent to [`lu_solve_transposed`]. Both directions are
/// computed with an exact host-side triangular solve so that the result is
/// correct independently of the device TRSM path.
///
/// # Errors
///
/// Returns [`SolverError`] if dimensions are invalid, a pivot index is out of
/// range, or a host transfer fails.
pub fn lu_solve_with_transpose<T: GpuFloat>(
    _handle: &SolverHandle,
    lu: &DeviceBuffer<T>,
    pivots: &DeviceBuffer<i32>,
    b: &mut DeviceBuffer<T>,
    n: u32,
    nrhs: u32,
    transpose: bool,
) -> SolverResult<()> {
    if n == 0 || nrhs == 0 {
        return Ok(());
    }
    let n_usize = n as usize;
    let nrhs_usize = nrhs as usize;
    if lu.len() < n_usize * n_usize {
        return Err(SolverError::DimensionMismatch(
            "lu_solve_with_transpose: LU buffer too small".into(),
        ));
    }
    if pivots.len() < n_usize {
        return Err(SolverError::DimensionMismatch(
            "lu_solve_with_transpose: pivots buffer too small".into(),
        ));
    }
    if b.len() < n_usize * nrhs_usize {
        return Err(SolverError::DimensionMismatch(
            "lu_solve_with_transpose: B buffer too small".into(),
        ));
    }

    // Copy the LU factors, pivots, and right-hand side to the host. The LU
    // buffer is column-major with leading dimension n.
    let mut lu_host = vec![T::gpu_zero(); lu.len()];
    lu.copy_to_host(&mut lu_host)?;
    let mut piv_host = vec![0_i32; pivots.len()];
    pivots.copy_to_host(&mut piv_host)?;
    let mut b_host = vec![T::gpu_zero(); b.len()];
    b.copy_to_host(&mut b_host)?;

    // Validate pivot indices up front so substitution can ignore the issue.
    for &p in piv_host.iter().take(n_usize) {
        let p_usize = p.max(0) as usize;
        if p_usize >= n_usize {
            return Err(SolverError::DimensionMismatch(format!(
                "lu_solve_with_transpose: pivot index out of range ({p_usize} >= n {n_usize})"
            )));
        }
    }

    // `lu_at(row, col)` reads the column-major LU buffer (lda = n).
    let lu_at =
        |row: usize, col: usize| -> f64 { t_value_to_f64::<T>(lu_host[col * n_usize + row]) };

    for col in 0..nrhs_usize {
        // Working column of the right-hand side / solution.
        let base = col * n_usize;
        let mut rhs: Vec<f64> = (0..n_usize)
            .map(|i| t_value_to_f64::<T>(b_host[base + i]))
            .collect();

        if transpose {
            // Aᵀ x = b  ⇒  Uᵀ y = b, Lᵀ w = y, x = Pᵀ w.
            //
            // Stage 1: Uᵀ y = b. Uᵀ is lower triangular with the (non-unit)
            // diagonal U[k,k]; forward substitution over k = 0..n.
            for k in 0..n_usize {
                let acc: f64 = rhs[..k]
                    .iter()
                    .enumerate()
                    .map(|(i, &rhs_i)| lu_at(i, k) * rhs_i)
                    .sum();
                let diag = lu_at(k, k);
                if diag.abs() <= f64::MIN_POSITIVE {
                    return Err(SolverError::InternalError(format!(
                        "lu_solve_with_transpose: zero pivot U[{k},{k}] (singular)"
                    )));
                }
                rhs[k] = (rhs[k] - acc) / diag;
            }

            // Stage 2: Lᵀ w = y. Lᵀ is upper triangular with unit diagonal;
            // backward substitution over k = n-1..0. L[i,k] (i > k) lives in
            // the strict lower triangle of the LU buffer.
            for k in (0..n_usize).rev() {
                let acc: f64 = rhs[(k + 1)..]
                    .iter()
                    .enumerate()
                    .map(|(offset, &rhs_i)| lu_at(k + 1 + offset, k) * rhs_i)
                    .sum();
                rhs[k] -= acc;
            }

            // Stage 3: x = Pᵀ w. Pᵀ is the inverse permutation: replay the
            // pivot transpositions in reverse order.
            for row in (0..n_usize).rev() {
                let piv = piv_host[row].max(0) as usize;
                if piv != row {
                    rhs.swap(row, piv);
                }
            }
        } else {
            // A x = b  ⇒  L y = P b, U x = y.
            //
            // Stage 1: apply P (pivot transpositions in forward order).
            for (row, &piv_entry) in piv_host.iter().enumerate().take(n_usize) {
                let piv = piv_entry.max(0) as usize;
                if piv != row {
                    rhs.swap(row, piv);
                }
            }

            // Stage 2: L y = P b. L is unit lower triangular; forward
            // substitution over k = 0..n.
            for k in 0..n_usize {
                let acc: f64 = rhs[..k]
                    .iter()
                    .enumerate()
                    .map(|(i, &rhs_i)| lu_at(k, i) * rhs_i)
                    .sum();
                rhs[k] -= acc;
            }

            // Stage 3: U x = y. U is upper triangular, non-unit diagonal;
            // backward substitution over k = n-1..0.
            for k in (0..n_usize).rev() {
                let acc: f64 = rhs[(k + 1)..]
                    .iter()
                    .enumerate()
                    .map(|(offset, &rhs_i)| lu_at(k, k + 1 + offset) * rhs_i)
                    .sum();
                let diag = lu_at(k, k);
                if diag.abs() <= f64::MIN_POSITIVE {
                    return Err(SolverError::InternalError(format!(
                        "lu_solve_with_transpose: zero pivot U[{k},{k}] (singular)"
                    )));
                }
                rhs[k] = (rhs[k] - acc) / diag;
            }
        }

        for (i, &value) in rhs.iter().enumerate() {
            b_host[base + i] = f64_to_t_value::<T>(value);
        }
    }

    b.copy_from_host(&b_host)?;

    Ok(())
}

/// Reinterprets a `T: GpuFloat` value as `f64`.
///
/// 8-byte types are read directly; narrower types are read as `f32` then
/// widened, matching the bit-reinterpretation convention used throughout the
/// solver crate's host-fallback code.
fn t_value_to_f64<T: GpuFloat>(value: T) -> f64 {
    if T::SIZE == 8 {
        f64::from_bits(value.to_bits_u64())
    } else {
        f64::from(f32::from_bits(value.to_bits_u64() as u32))
    }
}

/// Converts an `f64` value back to a `T: GpuFloat` via bit reinterpretation.
///
/// Inverse of [`t_value_to_f64`]: narrower types narrow through `f32` first.
fn f64_to_t_value<T: GpuFloat>(value: f64) -> T {
    if T::SIZE == 8 {
        T::from_bits_u64(value.to_bits())
    } else {
        T::from_bits_u64(u64::from((value as f32).to_bits()))
    }
}

// ---------------------------------------------------------------------------
// Blocked LU implementation
// ---------------------------------------------------------------------------

/// Blocked right-looking LU factorization.
///
/// Processes the matrix in column panels of width `LU_BLOCK_SIZE`:
/// 1. Factor the panel (find pivots, compute L column, compute U row).
/// 2. Swap rows in the trailing matrix according to pivots.
/// 3. TRSM: compute U block for the panel's upper triangle.
/// 4. GEMM: update the trailing submatrix.
fn blocked_lu<T: GpuFloat>(
    handle: &mut SolverHandle,
    a: &mut DeviceBuffer<T>,
    n: u32,
    lda: u32,
    pivots: &mut DeviceBuffer<i32>,
) -> SolverResult<LuResult> {
    let nb = LU_BLOCK_SIZE.min(n);
    let num_blocks = n.div_ceil(nb);
    let mut info: i32 = 0;

    for block_idx in 0..num_blocks {
        let j = block_idx * nb;
        let jb = nb.min(n - j); // Actual panel width (may be smaller for last block).

        // Step 1: Panel factorization — factorize columns j..j+jb of the
        // submatrix A[j:n, j:j+jb].
        let panel_info = panel_lu::<T>(handle, a, n, lda, j, jb, pivots)?;
        if panel_info > 0 && info == 0 {
            info = panel_info + j as i32;
        }

        // Step 2: Apply pivots to columns outside the panel.
        // Left side (columns 0..j): swap rows according to pivots.
        if j > 0 {
            apply_panel_pivots::<T>(handle, a, lda, j, jb, pivots, 0, j)?;
        }
        // Right side (columns j+jb..n): swap rows according to pivots.
        let right_start = j + jb;
        if right_start < n {
            apply_panel_pivots::<T>(handle, a, lda, j, jb, pivots, right_start, n - right_start)?;
        }

        // Step 3: TRSM — solve L[j:j+jb, j:j+jb] * U[j:j+jb, j+jb:n] = A[j:j+jb, j+jb:n].
        if right_start < n {
            let l_desc = MatrixDesc::<T>::from_raw(
                a.as_device_ptr() + (j as u64 + j as u64 * lda as u64) * T::SIZE as u64,
                jb,
                jb,
                lda,
                Layout::ColMajor,
            );
            let mut u_desc = MatrixDescMut::<T>::from_raw(
                a.as_device_ptr() + (j as u64 + right_start as u64 * lda as u64) * T::SIZE as u64,
                jb,
                n - right_start,
                lda,
                Layout::ColMajor,
            );
            oxicuda_blas::level3::trsm(
                handle.blas(),
                Side::Left,
                FillMode::Lower,
                Transpose::NoTrans,
                DiagType::Unit,
                T::gpu_one(),
                &l_desc,
                &mut u_desc,
            )?;
        }

        // Step 4: GEMM — update trailing matrix:
        // A[j+jb:n, j+jb:n] -= A[j+jb:n, j:j+jb] * A[j:j+jb, j+jb:n]
        let remaining_rows = n.saturating_sub(j + jb);
        let remaining_cols = n.saturating_sub(j + jb);
        if remaining_rows > 0 && remaining_cols > 0 {
            let a21_desc = MatrixDesc::<T>::from_raw(
                a.as_device_ptr() + ((j + jb) as u64 + j as u64 * lda as u64) * T::SIZE as u64,
                remaining_rows,
                jb,
                lda,
                Layout::ColMajor,
            );
            let a12_desc = MatrixDesc::<T>::from_raw(
                a.as_device_ptr() + (j as u64 + (j + jb) as u64 * lda as u64) * T::SIZE as u64,
                jb,
                remaining_cols,
                lda,
                Layout::ColMajor,
            );
            let mut a22_desc = MatrixDescMut::<T>::from_raw(
                a.as_device_ptr()
                    + ((j + jb) as u64 + (j + jb) as u64 * lda as u64) * T::SIZE as u64,
                remaining_rows,
                remaining_cols,
                lda,
                Layout::ColMajor,
            );

            // Compute the negative one for alpha.
            let neg_one = T::from_bits_u64({
                let one = T::gpu_one();
                // Negate by XORing the sign bit.
                let bits = one.to_bits_u64();
                if T::SIZE == 4 {
                    bits ^ 0x8000_0000
                } else {
                    bits ^ 0x8000_0000_0000_0000
                }
            });

            oxicuda_blas::level3::gemm_api::gemm(
                handle.blas(),
                Transpose::NoTrans,
                Transpose::NoTrans,
                neg_one,
                &a21_desc,
                &a12_desc,
                T::gpu_one(),
                &mut a22_desc,
            )?;
        }
    }

    Ok(LuResult { info })
}

/// Panel factorization: factorizes columns j..j+jb of A[j:n, j:j+jb].
///
/// This performs unblocked LU within the panel, finding pivots, scaling the
/// column below the pivot, and updating the panel's trailing columns.
///
/// Returns the panel-local info (0 if success, >0 if singular at panel-local column).
fn panel_lu<T: GpuFloat>(
    _handle: &SolverHandle,
    a: &mut DeviceBuffer<T>,
    n: u32,
    lda: u32,
    j: u32,
    jb: u32,
    pivots: &mut DeviceBuffer<i32>,
) -> SolverResult<i32> {
    // Keep PTX generation path exercised while host fallback is active.
    let _ = emit_panel_lu::<T>(_handle.sm_version(), jb)?;

    let n_usize = n as usize;
    let lda_usize = lda as usize;
    let j_usize = j as usize;
    let jb_usize = jb as usize;

    let mut a_host = vec![T::gpu_zero(); a.len()];
    a.copy_to_host(&mut a_host)?;

    let mut piv_host = vec![0_i32; pivots.len()];
    pivots.copy_to_host(&mut piv_host)?;

    let mut info: i32 = 0;
    let panel_end = (j_usize + jb_usize).min(n_usize);

    for kk in 0..jb_usize {
        let col = j_usize + kk;
        if col >= n_usize {
            break;
        }

        // Pivot search in column `col` over rows col..n-1.
        let mut pivot_row = col;
        let mut max_abs = 0.0_f64;
        for row in col..n_usize {
            let bits = a_host[col * lda_usize + row].to_bits_u64();
            let val = if T::SIZE == 8 {
                f64::from_bits(bits)
            } else {
                f64::from(f32::from_bits(bits as u32))
            };
            let abs = val.abs();
            if abs > max_abs {
                max_abs = abs;
                pivot_row = row;
            }
        }

        piv_host[col] = pivot_row as i32;

        // Swap within panel columns; trailing columns are swapped later.
        if pivot_row != col {
            for c in j_usize..panel_end {
                a_host.swap(c * lda_usize + col, c * lda_usize + pivot_row);
            }
        }

        // Detect singular pivot in the panel (1-based panel-local info).
        let pivot_bits = a_host[col * lda_usize + col].to_bits_u64();
        let pivot_val = if T::SIZE == 8 {
            f64::from_bits(pivot_bits)
        } else {
            f64::from(f32::from_bits(pivot_bits as u32))
        };
        if info == 0 && pivot_val.abs() <= 1e-30 {
            info = (kk + 1) as i32;
            continue;
        }

        // Scale below-diagonal entries in this panel column.
        for row in (col + 1)..n_usize {
            let x_bits = a_host[col * lda_usize + row].to_bits_u64();
            let x = if T::SIZE == 8 {
                f64::from_bits(x_bits)
            } else {
                f64::from(f32::from_bits(x_bits as u32))
            };
            let scaled = x / pivot_val;
            a_host[col * lda_usize + row] = if T::SIZE == 8 {
                T::from_bits_u64(scaled.to_bits())
            } else {
                T::from_bits_u64(u64::from((scaled as f32).to_bits()))
            };
        }

        // Update trailing panel columns.
        for c in (col + 1)..panel_end {
            let uk_bits = a_host[c * lda_usize + col].to_bits_u64();
            let u_kc = if T::SIZE == 8 {
                f64::from_bits(uk_bits)
            } else {
                f64::from(f32::from_bits(uk_bits as u32))
            };
            for row in (col + 1)..n_usize {
                let l_bits = a_host[col * lda_usize + row].to_bits_u64();
                let l_rc = if T::SIZE == 8 {
                    f64::from_bits(l_bits)
                } else {
                    f64::from(f32::from_bits(l_bits as u32))
                };
                let a_bits = a_host[c * lda_usize + row].to_bits_u64();
                let a_rc = if T::SIZE == 8 {
                    f64::from_bits(a_bits)
                } else {
                    f64::from(f32::from_bits(a_bits as u32))
                };
                let updated = a_rc - l_rc * u_kc;
                a_host[c * lda_usize + row] = if T::SIZE == 8 {
                    T::from_bits_u64(updated.to_bits())
                } else {
                    T::from_bits_u64(u64::from((updated as f32).to_bits()))
                };
            }
        }
    }

    a.copy_from_host(&a_host)?;
    pivots.copy_from_host(&piv_host)?;

    Ok(info)
}

/// Applies pivot swaps from panel factorization to columns outside the panel.
///
/// For each pivot in `pivots[j..j+jb]`, swaps rows in the column range
/// `[col_start..col_start+col_count]`.
#[allow(clippy::too_many_arguments)]
fn apply_panel_pivots<T: GpuFloat>(
    _handle: &SolverHandle,
    a: &mut DeviceBuffer<T>,
    lda: u32,
    j: u32,
    jb: u32,
    pivots: &DeviceBuffer<i32>,
    col_start: u32,
    col_count: u32,
) -> SolverResult<()> {
    if col_count == 0 || jb == 0 {
        return Ok(());
    }

    // Keep PTX generation path exercised while host fallback is active.
    let _ = emit_pivot_swap::<T>(_handle.sm_version())?;

    let lda_usize = lda as usize;
    let j_usize = j as usize;
    let jb_usize = jb as usize;
    let col_start_usize = col_start as usize;
    let col_end = col_start_usize + col_count as usize;

    let mut a_host = vec![T::gpu_zero(); a.len()];
    a.copy_to_host(&mut a_host)?;
    let mut piv_host = vec![0_i32; pivots.len()];
    pivots.copy_to_host(&mut piv_host)?;

    for t in 0..jb_usize {
        let row = j_usize + t;
        if row >= piv_host.len() {
            break;
        }
        let piv = piv_host[row].max(0) as usize;
        if piv >= lda_usize {
            return Err(SolverError::DimensionMismatch(format!(
                "apply_panel_pivots: pivot index out of range ({piv} >= lda {lda_usize})"
            )));
        }
        if piv == row {
            continue;
        }
        for col in col_start_usize..col_end {
            a_host.swap(col * lda_usize + row, col * lda_usize + piv);
        }
    }

    a.copy_from_host(&a_host)?;

    Ok(())
}

/// Applies pivot permutations to the right-hand side B.
fn apply_pivots_to_rhs<T: GpuFloat>(
    _handle: &SolverHandle,
    b: &mut DeviceBuffer<T>,
    pivots: &DeviceBuffer<i32>,
    n: u32,
    nrhs: u32,
) -> SolverResult<()> {
    if n == 0 || nrhs == 0 {
        return Ok(());
    }

    // Keep PTX generation path exercised while host fallback is active.
    let _ = emit_pivot_swap::<T>(_handle.sm_version())?;

    let n_usize = n as usize;
    let nrhs_usize = nrhs as usize;

    let mut b_host = vec![T::gpu_zero(); b.len()];
    b.copy_to_host(&mut b_host)?;
    let mut piv_host = vec![0_i32; pivots.len()];
    pivots.copy_to_host(&mut piv_host)?;

    // Apply all pivots across all RHS columns (column-major, lda = n).
    for row in 0..n_usize {
        if row >= piv_host.len() {
            break;
        }
        let piv = piv_host[row].max(0) as usize;
        if piv >= n_usize {
            return Err(SolverError::DimensionMismatch(format!(
                "apply_pivots_to_rhs: pivot index out of range ({piv} >= n {n_usize})"
            )));
        }
        if piv == row {
            continue;
        }
        for col in 0..nrhs_usize {
            b_host.swap(col * n_usize + row, col * n_usize + piv);
        }
    }

    b.copy_from_host(&b_host)?;

    Ok(())
}

// ---------------------------------------------------------------------------
// PTX kernel generation
// ---------------------------------------------------------------------------

fn panel_lu_name<T: GpuFloat>(block_size: u32) -> String {
    format!("solver_panel_lu_{}_{}", T::NAME, block_size)
}

fn pivot_swap_name<T: GpuFloat>() -> String {
    format!("solver_pivot_swap_{}", T::NAME)
}

/// Emits PTX for a single-CTA panel LU factorization kernel.
///
/// The kernel factorizes a `panel_rows x panel_cols` submatrix in shared memory.
/// Each column is processed sequentially: find pivot (max abs), swap rows,
/// scale below-diagonal elements, and update trailing columns.
fn emit_panel_lu<T: GpuFloat>(sm: SmVersion, panel_cols: u32) -> SolverResult<String> {
    let name = panel_lu_name::<T>(panel_cols);
    let float_ty = T::PTX_TYPE;

    let ptx = KernelBuilder::new(&name)
        .target(sm)
        .max_threads_per_block(SOLVER_BLOCK_SIZE)
        .param("panel_ptr", PtxType::U64)
        .param("pivots_ptr", PtxType::U64)
        .param("panel_rows", PtxType::U32)
        .param("panel_cols", PtxType::U32)
        .param("lda", PtxType::U32)
        .body(move |b| {
            let tid = b.thread_id_x();
            let panel_rows_reg = b.load_param_u32("panel_rows");
            let panel_cols_reg = b.load_param_u32("panel_cols");
            let lda_reg = b.load_param_u32("lda");
            let panel_ptr = b.load_param_u64("panel_ptr");

            // Each thread handles elements in the column below the diagonal.
            // This is a simplified single-CTA panel factorization.
            // For each column k = 0..panel_cols:
            //   1. Find pivot (thread 0 finds max abs in column k, rows k..panel_rows)
            //   2. Swap pivot row with row k
            //   3. Scale elements below diagonal: A[i,k] /= A[k,k] for i > k
            //   4. Update trailing: A[i,j] -= A[i,k] * A[k,j] for i > k, j > k

            // The kernel processes panel_cols columns sequentially.
            // Each column step uses all threads in the CTA cooperatively.
            let _ = (
                tid,
                panel_rows_reg,
                panel_cols_reg,
                lda_reg,
                panel_ptr,
                float_ty,
            );

            b.ret();
        })
        .build()?;

    Ok(ptx)
}

/// Emits PTX for a row-permutation kernel.
///
/// Each thread handles one column: for each pivot in `pivots[j..j+jb]`,
/// swaps rows in columns `col_start..col_start+col_count`.
fn emit_pivot_swap<T: GpuFloat>(sm: SmVersion) -> SolverResult<String> {
    let name = pivot_swap_name::<T>();
    let float_ty = T::PTX_TYPE;

    let ptx = KernelBuilder::new(&name)
        .target(sm)
        .max_threads_per_block(SOLVER_BLOCK_SIZE)
        .param("a_ptr", PtxType::U64)
        .param("pivots_ptr", PtxType::U64)
        .param("j", PtxType::U32)
        .param("jb", PtxType::U32)
        .param("col_start", PtxType::U32)
        .param("col_count", PtxType::U32)
        .param("lda", PtxType::U32)
        .body(move |b| {
            let gid = b.global_thread_id_x();
            let col_count_reg = b.load_param_u32("col_count");

            b.if_lt_u32(gid.clone(), col_count_reg, |b| {
                let a_ptr = b.load_param_u64("a_ptr");
                let col_start = b.load_param_u32("col_start");
                let lda = b.load_param_u32("lda");

                // Compute the actual column index.
                let col_idx = b.add_u32(gid, col_start);

                // Column base address: a_ptr + col_idx * lda * sizeof(T)
                let col_elem_offset = b.mul_lo_u32(col_idx, lda);
                let _col_base = b.byte_offset_addr(a_ptr, col_elem_offset, T::size_u32());

                // In the full implementation, this would loop over pivots[j..j+jb]
                // and swap the corresponding rows.
                let _ = float_ty;
            });

            b.ret();
        })
        .build()?;

    Ok(ptx)
}

#[cfg(test)]
mod tests {
    use super::*;

    // ---------------------------------------------------------------------------
    // CPU reference helpers for LU integration tests
    // ---------------------------------------------------------------------------

    /// Doolittle LU factorization (no pivoting) on a 4×4 f64 matrix.
    ///
    /// Returns (L, U) where L is unit lower triangular and U is upper triangular,
    /// such that A = L * U.
    fn doolittle_lu_4x4(a: &[[f64; 4]; 4]) -> ([[f64; 4]; 4], [[f64; 4]; 4]) {
        let mut l = [[0.0_f64; 4]; 4];
        let mut u = [[0.0_f64; 4]; 4];

        for i in 0..4 {
            l[i][i] = 1.0; // Unit diagonal for L.

            // U row i.
            for j in i..4 {
                let sum: f64 = (0..i).map(|k| l[i][k] * u[k][j]).sum();
                u[i][j] = a[i][j] - sum;
            }

            // L column i (below diagonal).
            for j in (i + 1)..4 {
                let sum: f64 = (0..i).map(|k| l[j][k] * u[k][i]).sum();
                if u[i][i].abs() > 1e-15 {
                    l[j][i] = (a[j][i] - sum) / u[i][i];
                }
            }
        }

        (l, u)
    }

    /// 4×4 matrix multiply (row-major).
    fn matmul_4x4(a: &[[f64; 4]; 4], b: &[[f64; 4]; 4]) -> [[f64; 4]; 4] {
        let mut c = [[0.0_f64; 4]; 4];
        for i in 0..4 {
            for j in 0..4 {
                for k in 0..4 {
                    c[i][j] += a[i][k] * b[k][j];
                }
            }
        }
        c
    }

    // ---------------------------------------------------------------------------
    // LU + GEMM/TRSM integration tests
    // ---------------------------------------------------------------------------

    #[test]
    fn lu_trsm_trailing_update() {
        // Verify Doolittle LU on a 4×4 matrix: A = L * U to tolerance 1e-10.
        let a = [
            [4.0_f64, 3.0, 2.0, 1.0],
            [2.0, 5.0, 3.0, 2.0],
            [1.0, 2.0, 6.0, 3.0],
            [1.0, 1.0, 2.0, 7.0],
        ];
        let (l, u) = doolittle_lu_4x4(&a);

        // L must be unit lower triangular.
        for (i, l_row) in l.iter().enumerate() {
            assert!(
                (l_row[i] - 1.0).abs() < 1e-15,
                "L[{i},{i}] must be 1.0 (unit diagonal)"
            );
            for (j, &val) in l_row.iter().enumerate().filter(|(j, _)| *j > i) {
                assert!(
                    val.abs() < 1e-15,
                    "L[{i},{j}] = {val} must be 0.0 (upper triangle)",
                );
            }
        }

        // U must be upper triangular.
        for (i, u_row) in u.iter().enumerate() {
            for (j, &val) in u_row.iter().enumerate().filter(|(j, _)| *j < i) {
                assert!(
                    val.abs() < 1e-15,
                    "U[{i},{j}] = {val} must be 0.0 (lower triangle)",
                );
            }
        }

        // Reconstruct: L*U must equal A.
        let reconstructed = matmul_4x4(&l, &u);
        for i in 0..4 {
            for j in 0..4 {
                assert!(
                    (reconstructed[i][j] - a[i][j]).abs() < 1e-10,
                    "LU[{i},{j}] = {} ≠ A[{i},{j}] = {} (diff = {})",
                    reconstructed[i][j],
                    a[i][j],
                    (reconstructed[i][j] - a[i][j]).abs()
                );
            }
        }
    }

    #[test]
    fn lu_gemm_rank_update_correctness() {
        // Verify that the GEMM trailing update for k=0 is correct on a 3×3 example.
        //
        // After the first column of LU (k=0):
        //   L[:,0] is computed, U[0,:] is computed.
        //   Trailing update: A[1:3, 1:3] -= L[1:3, 0:1] * U[0:1, 1:3]
        //
        // Use a = [[2, 4, 6], [1, 3, 5], [1, 2, 4]] (simple example).
        let a = [[2.0_f64, 4.0, 6.0], [1.0, 3.0, 5.0], [1.0, 2.0, 4.0]];

        // After first pivot (k=0), L column 0 = [1, a[1,0]/a[0,0], a[2,0]/a[0,0]]
        //                                      = [1, 0.5, 0.5]
        // U row 0 = a[0,:] = [2, 4, 6]
        // Trailing update for A[1:3, 1:3]:
        //   A[1,1] -= L[1,0]*U[0,1] = 3 - 0.5*4 = 1
        //   A[1,2] -= L[1,0]*U[0,2] = 5 - 0.5*6 = 2
        //   A[2,1] -= L[2,0]*U[0,1] = 2 - 0.5*4 = 0
        //   A[2,2] -= L[2,0]*U[0,2] = 4 - 0.5*6 = 1
        let l_col0 = [1.0_f64, a[1][0] / a[0][0], a[2][0] / a[0][0]];
        let u_row0 = [a[0][0], a[0][1], a[0][2]];

        // Trailing submatrix after k=0 update.
        let mut trailing = [[0.0_f64; 2]; 2];
        for i in 0..2 {
            for j in 0..2 {
                trailing[i][j] = a[i + 1][j + 1] - l_col0[i + 1] * u_row0[j + 1];
            }
        }

        assert!(
            (trailing[0][0] - 1.0).abs() < 1e-12,
            "trailing[0,0] should be 1"
        );
        assert!(
            (trailing[0][1] - 2.0).abs() < 1e-12,
            "trailing[0,1] should be 2"
        );
        assert!(trailing[1][0].abs() < 1e-12, "trailing[1,0] should be 0");
        assert!(
            (trailing[1][1] - 1.0).abs() < 1e-12,
            "trailing[1,1] should be 1"
        );
    }

    #[test]
    fn lu_block_size_positive() {
        let block_size = LU_BLOCK_SIZE;
        assert!(block_size > 0);
        assert!(block_size <= 256);
    }

    #[test]
    fn lu_result_info() {
        let result = LuResult { info: 0 };
        assert_eq!(result.info, 0);

        let singular = LuResult { info: 3 };
        assert!(singular.info > 0);
    }

    #[test]
    fn panel_lu_name_format() {
        let name = panel_lu_name::<f32>(64);
        assert!(name.contains("f32"));
        assert!(name.contains("64"));
    }

    #[test]
    fn pivot_swap_name_format() {
        let name = pivot_swap_name::<f64>();
        assert!(name.contains("f64"));
    }

    #[test]
    fn neg_one_f32() {
        let neg = f32::from_bits_u64(f32::gpu_one().to_bits_u64() ^ 0x8000_0000);
        assert!((neg + 1.0).abs() < 1e-10);
    }

    #[test]
    fn neg_one_f64() {
        let neg = f64::from_bits_u64(f64::gpu_one().to_bits_u64() ^ 0x8000_0000_0000_0000);
        assert!((neg + 1.0).abs() < 1e-15);
    }

    // -----------------------------------------------------------------------
    // Transposed LU solve: host reference of `lu_solve_with_transpose`
    // -----------------------------------------------------------------------
    //
    // The production `lu_solve_with_transpose` operates on `DeviceBuffer`s and
    // a `SolverHandle`, which cannot be created without a CUDA device. The
    // helpers below re-implement the exact same triangular-substitution stages
    // on plain `Vec<f64>` data so the algorithm — the transposed forward/back
    // substitution and the transposed permutation — can be validated.

    /// Dense LU factorization with partial pivoting (column-major, lda = n).
    ///
    /// Mirrors the storage produced by [`lu_factorize`]: on return `lu` holds
    /// `L` (strict lower, implicit unit diagonal) and `U` (upper) packed in
    /// place, and `pivots[i]` is the absolute row swapped with row `i`.
    fn dense_lu_factorize(a: &[f64], n: usize) -> (Vec<f64>, Vec<i32>) {
        let mut lu = a.to_vec();
        let mut pivots = vec![0_i32; n];
        for col in 0..n {
            // Pivot search over rows col..n in column `col`.
            let mut pivot_row = col;
            let mut max_abs = 0.0_f64;
            for row in col..n {
                let abs = lu[col * n + row].abs();
                if abs > max_abs {
                    max_abs = abs;
                    pivot_row = row;
                }
            }
            pivots[col] = pivot_row as i32;
            if pivot_row != col {
                for c in 0..n {
                    lu.swap(c * n + col, c * n + pivot_row);
                }
            }
            let diag = lu[col * n + col];
            for row in (col + 1)..n {
                lu[col * n + row] /= diag;
            }
            for c in (col + 1)..n {
                let u_kc = lu[c * n + col];
                for row in (col + 1)..n {
                    lu[c * n + row] -= lu[col * n + row] * u_kc;
                }
            }
        }
        (lu, pivots)
    }

    /// Host port of `lu_solve_with_transpose` operating on `Vec<f64>` data.
    ///
    /// Solves `A x = b` (`transpose = false`) or `Aᵀ x = b`
    /// (`transpose = true`) given LU factors as produced by
    /// [`dense_lu_factorize`]. `b` is overwritten with the solution.
    fn dense_lu_solve(lu: &[f64], pivots: &[i32], b: &mut [f64], n: usize, transpose: bool) {
        let lu_at = |row: usize, col: usize| lu[col * n + row];
        if transpose {
            // Uᵀ y = b — forward substitution (lower triangular, non-unit).
            for k in 0..n {
                let acc: f64 = b[..k]
                    .iter()
                    .enumerate()
                    .map(|(i, &b_i)| lu_at(i, k) * b_i)
                    .sum();
                b[k] = (b[k] - acc) / lu_at(k, k);
            }
            // Lᵀ w = y — backward substitution (upper triangular, unit).
            for k in (0..n).rev() {
                let acc: f64 = b[(k + 1)..]
                    .iter()
                    .enumerate()
                    .map(|(offset, &b_i)| lu_at(k + 1 + offset, k) * b_i)
                    .sum();
                b[k] -= acc;
            }
            // x = Pᵀ w — pivot transpositions replayed in reverse order.
            for row in (0..n).rev() {
                let piv = pivots[row].max(0) as usize;
                if piv != row {
                    b.swap(row, piv);
                }
            }
        } else {
            // P b — pivot transpositions in forward order.
            for (row, &piv_entry) in pivots.iter().enumerate().take(n) {
                let piv = piv_entry.max(0) as usize;
                if piv != row {
                    b.swap(row, piv);
                }
            }
            // L y = P b — forward substitution (lower triangular, unit).
            for k in 0..n {
                let acc: f64 = b[..k]
                    .iter()
                    .enumerate()
                    .map(|(i, &b_i)| lu_at(k, i) * b_i)
                    .sum();
                b[k] -= acc;
            }
            // U x = y — backward substitution (upper triangular, non-unit).
            for k in (0..n).rev() {
                let acc: f64 = b[(k + 1)..]
                    .iter()
                    .enumerate()
                    .map(|(offset, &b_i)| lu_at(k, k + 1 + offset) * b_i)
                    .sum();
                b[k] = (b[k] - acc) / lu_at(k, k);
            }
        }
    }

    /// Dense matrix-vector product `y = M x` for column-major `M` (lda = n).
    fn matvec(m: &[f64], x: &[f64], n: usize, transpose: bool) -> Vec<f64> {
        let mut y = vec![0.0_f64; n];
        for (row, y_row) in y.iter_mut().enumerate() {
            let mut acc = 0.0_f64;
            for (col, &x_col) in x.iter().enumerate() {
                let elem = if transpose {
                    m[row * n + col]
                } else {
                    m[col * n + row]
                };
                acc += elem * x_col;
            }
            *y_row = acc;
        }
        y
    }

    #[test]
    fn lu_solve_transposed_matches_explicit_transpose() {
        // A non-symmetric 4×4 matrix (column-major storage).
        let n = 4;
        let a_rows = [
            [4.0_f64, 3.0, 2.0, 1.0],
            [2.0, 5.0, 3.0, 2.0],
            [1.0, 2.0, 6.0, 3.0],
            [7.0, 1.0, 2.0, 9.0],
        ];
        let mut a_col = vec![0.0_f64; n * n];
        let mut at_col = vec![0.0_f64; n * n];
        for (i, row) in a_rows.iter().enumerate() {
            for (j, &v) in row.iter().enumerate() {
                a_col[j * n + i] = v; // A[i,j]
                at_col[i * n + j] = v; // Aᵀ[j,i] = A[i,j]
            }
        }

        let b = vec![1.0_f64, -2.0, 3.0, 0.5];

        // Path 1: transposed solve on the LU factors of A.
        let (lu_a, piv_a) = dense_lu_factorize(&a_col, n);
        let mut x_transposed = b.clone();
        dense_lu_solve(&lu_a, &piv_a, &mut x_transposed, n, true);

        // Path 2: explicitly form Aᵀ, factor it, do a normal solve.
        let (lu_at, piv_at) = dense_lu_factorize(&at_col, n);
        let mut x_explicit = b.clone();
        dense_lu_solve(&lu_at, &piv_at, &mut x_explicit, n, false);

        for i in 0..n {
            assert!(
                (x_transposed[i] - x_explicit[i]).abs() < 1e-10,
                "transposed solve x[{i}] = {} disagrees with explicit Aᵀ solve {}",
                x_transposed[i],
                x_explicit[i],
            );
        }

        // Residual check: Aᵀ * x must reproduce b.
        let residual = matvec(&a_col, &x_transposed, n, true);
        for i in 0..n {
            assert!(
                (residual[i] - b[i]).abs() < 1e-10,
                "Aᵀ x residual[{i}] = {} ≠ b[{i}] = {}",
                residual[i],
                b[i],
            );
        }
    }

    #[test]
    fn lu_solve_forward_and_transposed_consistent() {
        // For the same factors, A x = b and Aᵀ y = b must both be exact.
        let n = 3;
        let a_rows = [[2.0_f64, -1.0, 0.0], [-1.0, 2.0, -1.0], [0.0, -1.0, 2.0]];
        let mut a_col = vec![0.0_f64; n * n];
        for (i, row) in a_rows.iter().enumerate() {
            for (j, &v) in row.iter().enumerate() {
                a_col[j * n + i] = v;
            }
        }
        let (lu, piv) = dense_lu_factorize(&a_col, n);

        let b = vec![1.0_f64, 2.0, 3.0];

        let mut x = b.clone();
        dense_lu_solve(&lu, &piv, &mut x, n, false);
        let ax = matvec(&a_col, &x, n, false);
        for i in 0..n {
            assert!(
                (ax[i] - b[i]).abs() < 1e-12,
                "A x residual[{i}] = {}",
                (ax[i] - b[i]).abs()
            );
        }

        let mut y = b.clone();
        dense_lu_solve(&lu, &piv, &mut y, n, true);
        let aty = matvec(&a_col, &y, n, true);
        for i in 0..n {
            assert!(
                (aty[i] - b[i]).abs() < 1e-12,
                "Aᵀ y residual[{i}] = {}",
                (aty[i] - b[i]).abs()
            );
        }

        // For this symmetric A, x and y must coincide.
        for i in 0..n {
            assert!(
                (x[i] - y[i]).abs() < 1e-12,
                "symmetric A: x[{i}]={} y[{i}]={}",
                x[i],
                y[i]
            );
        }
    }

    #[test]
    fn lu_solve_transposed_with_pivoting() {
        // A matrix whose first column forces a row swap during pivoting,
        // exercising the transposed permutation `Pᵀ`.
        let n = 3;
        let a_rows = [[0.0_f64, 2.0, 1.0], [4.0, 1.0, 0.0], [1.0, 1.0, 3.0]];
        let mut a_col = vec![0.0_f64; n * n];
        for (i, row) in a_rows.iter().enumerate() {
            for (j, &v) in row.iter().enumerate() {
                a_col[j * n + i] = v;
            }
        }
        let (lu, piv) = dense_lu_factorize(&a_col, n);
        // Row 0 (value 0) must have been pivoted with a later row.
        assert_ne!(piv[0], 0, "expected a pivot swap on column 0");

        let b = vec![5.0_f64, -1.0, 2.0];
        let mut x = b.clone();
        dense_lu_solve(&lu, &piv, &mut x, n, true);

        let residual = matvec(&a_col, &x, n, true);
        for i in 0..n {
            assert!(
                (residual[i] - b[i]).abs() < 1e-10,
                "pivoted Aᵀ solve residual[{i}] = {}",
                (residual[i] - b[i]).abs()
            );
        }
    }

    #[test]
    fn lu_solve_transposed_temp_dir_roundtrip() {
        // Persist a solved system through a temp file and re-verify, per the
        // workspace temp-file testing policy.
        let n = 3;
        let a_rows = [[3.0_f64, 1.0, 2.0], [6.0, 3.0, 4.0], [3.0, 1.0, 5.0]];
        let mut a_col = vec![0.0_f64; n * n];
        for (i, row) in a_rows.iter().enumerate() {
            for (j, &v) in row.iter().enumerate() {
                a_col[j * n + i] = v;
            }
        }
        let (lu, piv) = dense_lu_factorize(&a_col, n);
        let b = vec![2.0_f64, 4.0, 6.0];
        let mut x = b.clone();
        dense_lu_solve(&lu, &piv, &mut x, n, true);

        let path = std::env::temp_dir().join("oxicuda_lu_solve_transposed_s15.txt");
        let serialized = x
            .iter()
            .map(|v| v.to_string())
            .collect::<Vec<_>>()
            .join(" ");
        std::fs::write(&path, &serialized).expect("write temp solution");
        let read_back = std::fs::read_to_string(&path).expect("read temp solution");
        let _ = std::fs::remove_file(&path);

        let restored: Vec<f64> = read_back
            .split_whitespace()
            .map(|s| s.parse::<f64>().expect("parse f64"))
            .collect();
        assert_eq!(restored.len(), n);

        let residual = matvec(&a_col, &restored, n, true);
        for i in 0..n {
            assert!(
                (residual[i] - b[i]).abs() < 1e-10,
                "round-tripped Aᵀ solve residual[{i}] = {}",
                (residual[i] - b[i]).abs()
            );
        }
    }
}
