/*
 * SPDX-License-Identifier: BSD-3-Clause
 * Copyright (c) 1996-2025, The SLICOT Team (original Fortran77 code)
 * Copyright (c) 2025, slicot.c contributors (C11 translation)
 */

#include "slicot.h"
#include "slicot_blas.h"
#include <math.h>
#include <ctype.h>
#include <stddef.h>

/**
 * @brief Solve system of linear equations R*x = b or R'*x = b in least squares sense.
 *
 * Solves R*x = b or R'*x = b where R is an n-by-n block upper triangular matrix.
 *
 * @param[in] cond Condition estimation mode ('E', 'N', 'U').
 * @param[in] uplo Storage scheme ('U', 'L').
 * @param[in] trans Form of system ('N', 'T', 'C').
 * @param[in] n Order of matrix R.
 * @param[in] ipar Integer parameters (st, bn, bsm, bsn).
 * @param[in] lipar Length of ipar.
 * @param[in,out] r Matrix R (ldr x nc).
 * @param[in] ldr Leading dimension of R.
 * @param[in] sdiag Diagonal elements of blocks (if uplo='L').
 * @param[in] s Transpose of last block column (if uplo='L').
 * @param[in] lds Leading dimension of S.
 * @param[in,out] b Right hand side vector b. On exit, solution x.
 * @param[in,out] ranks Numerical ranks of submatrices.
 * @param[in] tol Tolerance for rank determination.
 * @param[out] dwork Workspace.
 * @param[in] ldwork Length of dwork.
 * @param[out] info Exit code.
 */
void nf01br(const char *cond, const char *uplo, const char *trans, i32 n, 
            const i32 *ipar, i32 lipar, f64 *r, i32 ldr, f64 *sdiag, 
            f64 *s, i32 lds, f64 *b, i32 *ranks, f64 tol, f64 *dwork, 
            i32 ldwork, i32 *info)
{
    /* Local variables */
    i32 bn, bsm, bsn, i, i1, j, k, l, nc, nths, rank, st, jwork;
    f64 toldef, dum[3];
    bool econd, ncond, lower, tranr, full;
    char transl[2] = {0, 0}, uplol[2] = {0, 0};
    f64 zero = 0.0, one = 1.0, svlmax = 0.0;
    i32 inc_1 = 1, inc_0 = 0;

    /* Check scalar input parameters */
    econd = (cond[0] == 'E' || cond[0] == 'e');
    ncond = (cond[0] == 'N' || cond[0] == 'n');
    lower = (uplo[0] == 'L' || uplo[0] == 'l');
    tranr = (trans[0] == 'T' || trans[0] == 't' || trans[0] == 'C' || trans[0] == 'c');

    *info = 0;
    if (!(econd || ncond || cond[0] == 'U' || cond[0] == 'u')) *info = -1;
    else if (!(lower || uplo[0] == 'U' || uplo[0] == 'u')) *info = -2;
    else if (!(tranr || trans[0] == 'N' || trans[0] == 'n')) *info = -3;
    else if (n < 0) *info = -4;
    else if (lipar < 4) *info = -6;
    else {
        st = ipar[0];
        bn = ipar[1];
        bsm = ipar[2];
        bsn = ipar[3];
        nths = bn * bsn;
        full = (bn <= 1 || bsn == 0);
        
        if (st < 0 || bn < 0 || bsm < 0 || bsn < 0) *info = -5;
        else if (n != nths + st) *info = -4;
        else if (ldr < 1 || (n > 0 && ldr < n)) *info = -8; /* Fixed: ldr >= max(1,n) */
        /* Fix: lds check logic */
        else {
            bool lds_ok = (lds >= 1);
            if (lower && !full && lds < st) lds_ok = false;
            if (!lds_ok) *info = -11;
            else {
                if (econd) {
                    if (full) l = 2 * n;
                    else l = 2 * ((bsn > st) ? bsn : st);
                } else {
                    l = 0;
                }
                if (ldwork < l) *info = -16;
            }
        }
    }

    if (*info != 0) {
        i32 err_code = -(*info);
        SLC_XERBLA("NF01BR", &err_code);
        return;
    }

    if (n == 0) return;

    if (econd) {
        toldef = tol;
        if (toldef <= zero) {
            toldef = (f64)n * SLC_DLAMCH("Epsilon");
        }
    }

    nc = bsn + st;
    
    if (full) {
        /* Special case: R is just upper triangular */
        
        if (lower) {
            /* Swap diagonal elements of R and SDIAG */
            i32 stride_diag = ldr + 1;
            SLC_DSWAP(&n, r, &stride_diag, sdiag, &inc_1);
            
            if (econd) {
                uplol[0] = 'U';
                transl[0] = trans[0];
                /* Swap upper and lower triangles */
                /* Not easy with BLAS swap. Done loop wise. */
                for (j = 0; j < n; j++) {
                    /* Swap R(J, J:N) with R(J:N, J) -- wait, lower triangle stores transpose of upper?
                       "strict lower triangle contains transpose of strict upper triangle"
                       So R[i, j] (lower) corresponds to R[j, i] (upper).
                       We want to move lower to upper.
                       Actually just swapping R[i, j] with R[j, i].
                    */
                    i32 len = n - 1 - j;
                    if (len > 0) {
                        /* Swap row j (starting col j+1) with col j (starting row j+1) */
                        /* Row j: &r[j + (j+1)*ldr]. Stride ldr. */
                        /* Col j: &r[(j+1) + j*ldr]. Stride 1. */
                        SLC_DSWAP(&len, &r[j + (j + 1) * ldr], &ldr, &r[(j + 1) + j * ldr], &inc_1);
                    }
                }
            } else {
                uplol[0] = uplo[0];
                transl[0] = tranr ? 'N' : 'T';
            }
        } else {
            uplol[0] = uplo[0];
            transl[0] = trans[0];
        }
        
        if (econd) {
            /* Call MB03OD for condition estimation */
            i32 rank_out;
            mb03od("No QR", n, n, r, ldr, NULL, toldef, svlmax, 
                   dwork, &rank_out, dum, dwork, ldwork, info);
            ranks[0] = rank_out;
            rank = rank_out;
        } else if (ncond) {
            rank = n;
            for (j = 0; j < n; j++) {
                if (r[j + j * ldr] == zero && rank == n) {
                    rank = j; /* 0-based index = number of non-zeros before it? */
                    /* Fortran: RANK = J - 1. If J=1 (first elt), rank=0. Correct. */
                }
            }
            ranks[0] = rank;
        } else {
            rank = ranks[0];
        }
        
        /* Solve R*x = b or R'*x = b */
        dum[0] = zero;
        if (rank < n) {
            i32 len = n - rank;
            SLC_DCOPY(&len, dum, &inc_0, &b[rank], &inc_1);
        }
        
        SLC_DTRSV(uplol, transl, "NonUnit", &rank, r, &ldr, b, &inc_1);
        
        if (lower) {
            /* Restore diagonal */
            i32 stride_diag = ldr + 1;
            SLC_DSWAP(&n, r, &stride_diag, sdiag, &inc_1);
            
            if (econd) {
                /* Restore triangles */
                for (j = 0; j < n; j++) {
                    i32 len = n - 1 - j;
                    if (len > 0) {
                        SLC_DSWAP(&len, &r[j + (j + 1) * ldr], &ldr, &r[(j + 1) + j * ldr], &inc_1);
                    }
                }
            }
        }
        return;
    }
    
    /* General case: l > 1 and BSN > 0 */
    /* Not implemented yet - sticking to Batch 2 scope and ensuring full functionality later if needed.
       But the test case uses BN=1, so we covered the test.
       Wait, if I don't implement the full logic, I can't claim full support.
       I should implement the general case too.
    */
    
    /* General case implementation */
    i = 0;
    l = bn;
    
    if (econd) {
        if (lower) {
            /* Swap loop */
            for (k = 0; k < bn; k++) {
                i32 stride_diag = ldr + 1;
                /* Swap BSN diagonal elements */
                SLC_DSWAP(&bsn, &r[i + 0 * ldr], &stride_diag, &sdiag[i], &inc_1); 
                /* Wait, R(I, 1). I increases.
                   Fortran: R(I, 1). Column 1? No, R is LDR x NC.
                   R is stored compressed. 
                   "leading (N-ST)-by-BSN part of this array"
                   R has BSN columns (if ST=0 and BN>1? No NC=BSN+ST).
                   
                   Structure Rc:
                   R_1 | L_1
                   R_2 | L_2
                   
                   R is (N) x (BSN+ST).
                   R_k is BSN x BSN.
                   Stored at R[ (k-1)*BSN : k*BSN-1, 0:BSN-1 ].
                   
                   Fortran: R(I, 1). I increases by BSN each block.
                   So R(1, 1) -> R(BSN+1, 1) -> ...
                   Diagonal of R_k starts at R((k-1)*BSN+1, 1).
                   Stride is LDR+1.
                */
                
                /* Swap upper/lower of R_k */
                /* R_k starts at row i, col 0. */
                for (j = 0; j < bsn; j++) {
                    /* Swap R_k(j, j+1:bsn) with R_k(j+1:bsn, j) */
                    /* Row i+j, col j. */
                    /* Strict upper part of R_k is stored in strict lower part of R_k?
                       "transpose of strict upper triangle of R_k is stored in strict lower triangle of R_k"
                       So we need to swap R_k[r, c] with R_k[c, r]? 
                       No, if we want to put it in upper part for MB03OD.
                       MB03OD expects upper triangular.
                       Lower part has transpose of upper.
                       So if we swap, we put the upper part back in upper.
                    */
                    /* Implementation detail: similar to full case but block wise */
                    i32 len = bsn - 1 - j;
                    if (len > 0) {
                        /* Swap row j of block (starting col j+1) with col j of block (starting row j+1) */
                        /* Block offset row: i. Col: 0. */
                        /* &r[(i+j) + (j+1)*ldr] vs &r[(i+j+1) + j*ldr] */
                        /* Stride ldr vs 1. */
                        SLC_DSWAP(&len, &r[(i + j) + (j + 1) * ldr], &ldr, &r[(i + j + 1) + j * ldr], &inc_1);
                    }
                }
                i += bsn;
            }
            
            if (st > 0) {
                /* Swap R_{l+1} diagonal */
                i32 stride_diag = ldr + 1;
                SLC_DSWAP(&st, &r[i + bsn * ldr], &stride_diag, &sdiag[i], &inc_1);
                
                /* Swap S and L_k^T */
                /* S is ST x (N-ST).
                   Last block column of R is [L_1; ...; L_l]. N-ST rows. ST cols?
                   No, "submatrix L_k".
                   Rc: [ ... | L_k ].
                   L_k is BSN x ST.
                   Last block column of R has ST columns. Indices BSN to BSN+ST-1.
                   
                   S stores transpose of last block column.
                   S: ST x (N-ST).
                   Last block column of R (without R_{l+1}): (N-ST) x ST.
                   So dimensions match for transpose.
                   
                   Swap R[0:N-ST, BSN:BSN+ST] with S.
                   R is col major. S is col major.
                   R block: rows 0..N-ST-1, cols BSN..BSN+ST-1.
                   S block: rows 0..ST-1, cols 0..N-ST-1.
                   
                   Fortran: CALL DSWAP( NTHS, R(1,J), 1, S(J-BSN,1), LDS )
                   Loop J = BSN+1 to NC.
                   R(1, J) is column J of R. Length NTHS.
                   S(J-BSN, 1) is row (J-BSN) of S? No, S(row, col).
                   S(J-BSN, 1). J-BSN goes from 1 to ST.
                   So row J-BSN of S. Stride LDS.
                   This swaps column of R with row of S.
                */
                for (j = bsn; j < nc; j++) {
                    /* Col j of R. &r[0 + j*ldr]. Stride 1. Length nths. */
                    /* Row j-bsn of S. &s[(j-bsn) + 0*lds]. Stride lds. Length nths? 
                       S has N-ST columns? No, S is ST x (N-ST).
                       Wait. S(input) dimension (LDS, N-ST).
                       "leading ST-by-(N-ST) part ... contains transpose of ... [ L_1' ... ]"
                       [L_1' ... L_l'] is ST x (N-ST).
                       So S stores ST rows, N-ST columns.
                       
                       Wait. Fortran: S(J-BSN, 1).
                       J goes BSN+1 to NC(=BSN+ST).
                       So J-BSN goes 1 to ST.
                       So we access row k of S.
                       Elements S(k, 1), S(k, 2)... S(k, N-ST).
                       We swap with R(1, J), R(2, J)... R(NTHS, J).
                       NTHS = N - ST.
                       So lengths match.
                    */
                    i32 s_row = j - bsn;
                    SLC_DSWAP(&nths, &r[0 + j * ldr], &inc_1, &s[s_row + 0 * lds], &lds);
                    
                    /* Swap upper/lower of R_{l+1} */
                    /* R_{l+1} is ST x ST. Starts at row i, col bsn. */
                    /* Loop to swap upper/lower of this block */
                    /* Similar to before, but block starts at col bsn */
                    /* Using previous loop approach, but now R_k logic for R_{l+1} */
                    /* Actually the Fortran code does:
                       CALL DSWAP( NC-J+1, R(I,J), LDR, R(I,J), 1 )
                       Here I is start of R_{l+1} rows. J iterates cols of R.
                       But J goes BSN+1 to NC.
                       This swaps lower triangle of R_{l+1} with upper.
                       Specifically, col J of R_{l+1} (lower part) with row J (upper part).
                    */
                    /* Let's reconstruct:
                       For j_local = 0 to st-1:
                         Swap R[i+j_local, bsn+j_local+1 ... ] with R[i+j_local+1..., bsn+j_local]
                    */
                }
                
                /* Doing the R_{l+1} swap properly: */
                for (j = 0; j < st; j++) {
                    i32 len = st - 1 - j;
                    if (len > 0) {
                        /* Swap row j of block (starting col j+1) with col j of block (starting row j+1) */
                        /* Block starts at row i, col bsn. */
                        /* Row j (of block): &r[(i+j) + (bsn+j+1)*ldr]. Stride ldr. */
                        /* Col j (of block): &r[(i+j+1) + (bsn+j)*ldr]. Stride 1. */
                        SLC_DSWAP(&len, &r[(i + j) + (bsn + j + 1) * ldr], &ldr, 
                                  &r[(i + j + 1) + (bsn + j) * ldr], &inc_1);
                    }
                }
                i += st;
            }
        }
        
        /* Incremental condition estimation */
        i1 = 0;
        for (k = 0; k < bn; k++) {
            i32 rank_out;
            mb03od("No QR", bsn, bsn, &r[i1], ldr, NULL, toldef, svlmax,
                   dwork, &rank_out, dum, dwork, ldwork, info);
            ranks[k] = rank_out;
            i1 += bsn; /* Next block starts bsn rows down */
        }
        
        if (st > 0) {
            i32 rank_out;
            /* R_{l+1} starts at row i1, col bsn */
            mb03od("No QR", st, st, &r[i1 + bsn * ldr], ldr, NULL, toldef, svlmax,
                   dwork, &rank_out, dum, dwork, ldwork, info);
            ranks[bn] = rank_out;
        }
    } else if (ncond) {
        i = 0;
        if (lower) {
            for (k = 0; k < bn; k++) {
                rank = bsn;
                for (j = 0; j < bsn; j++) {
                    if (sdiag[i] == zero && rank == bsn) rank = j;
                    i++;
                }
                ranks[k] = rank;
            }
            if (st > 0) {
                l++;
                rank = st;
                for (j = 0; j < st; j++) {
                    if (sdiag[i] == zero && rank == st) rank = j;
                    i++;
                }
                ranks[l - 1] = rank;
            }
        } else {
            for (k = 0; k < bn; k++) {
                rank = bsn;
                for (j = 0; j < bsn; j++) {
                    if (r[i + j * ldr] == zero && rank == bsn) rank = j;
                    i++;
                }
                ranks[k] = rank;
            }
            if (st > 0) {
                l++;
                rank = st;
                for (j = bsn; j < nc; j++) {
                    if (r[i + j * ldr] == zero && rank == st) rank = j - bsn;
                    i++;
                }
                ranks[l - 1] = rank;
            }
        }
    } else {
        if (st > 0) l++;
    }

    /* Solve */
    dum[0] = zero;
    if (lower && !econd) {
        if (!tranr) {
            i32 i1 = nths;
            if (st > 0) {
                rank = ranks[l - 1];
                if (rank < st) {
                    i32 len = st - rank;
                    SLC_DCOPY(&len, dum, &inc_0, &b[i1 + rank], &inc_1);
                }
                i32 stride_diag = ldr + 1;
                SLC_DSWAP(&st, &r[i1 + bsn * ldr], &stride_diag, &sdiag[i1], &inc_1);
                SLC_DTRSV("Lower", "Transpose", "NonUnit", &rank, &r[i1 + bsn * ldr], &ldr, &b[i1], &inc_1);
                SLC_DSWAP(&st, &r[i1 + bsn * ldr], &stride_diag, &sdiag[i1], &inc_1);

                f64 neg_one = -1.0;
                SLC_DGEMV("Transpose", &st, &nths, &neg_one, s, &lds, &b[nths], &inc_1, &one, b, &inc_1);
            }

            for (k = bn - 1; k >= 0; k--) {
                i1 -= bsn;
                rank = ranks[k];
                if (rank < bsn) {
                    i32 len = bsn - rank;
                    SLC_DCOPY(&len, dum, &inc_0, &b[i1 + rank], &inc_1);
                }
                i32 stride_diag = ldr + 1;
                SLC_DSWAP(&bsn, &r[i1], &stride_diag, &sdiag[i1], &inc_1);
                SLC_DTRSV("Lower", "Transpose", "NonUnit", &rank, &r[i1], &ldr, &b[i1], &inc_1);
                SLC_DSWAP(&bsn, &r[i1], &stride_diag, &sdiag[i1], &inc_1);
            }
        } else {
            i32 i1 = 0;
            char transl[2];
            if (tranr) transl[0] = 'N', transl[1] = 0;
            else transl[0] = 'T', transl[1] = 0;

            for (k = 0; k < bn; k++) {
                rank = ranks[k];
                if (rank < bsn) {
                    i32 len = bsn - rank;
                    SLC_DCOPY(&len, dum, &inc_0, &b[i1 + rank], &inc_1);
                }
                i32 stride_diag = ldr + 1;
                SLC_DSWAP(&bsn, &r[i1], &stride_diag, &sdiag[i1], &inc_1);
                SLC_DTRSV("Lower", transl, "NonUnit", &rank, &r[i1], &ldr, &b[i1], &inc_1);
                SLC_DSWAP(&bsn, &r[i1], &stride_diag, &sdiag[i1], &inc_1);
                i1 += bsn;
            }

            if (st > 0) {
                rank = ranks[l - 1];
                if (rank < st) {
                    i32 len = st - rank;
                    SLC_DCOPY(&len, dum, &inc_0, &b[i1 + rank], &inc_1);
                }
                f64 neg_one = -1.0;
                SLC_DGEMV("NoTranspose", &st, &nths, &neg_one, s, &lds, b, &inc_1, &one, &b[i1], &inc_1);

                i32 stride_diag = ldr + 1;
                SLC_DSWAP(&st, &r[i1 + bsn * ldr], &stride_diag, &sdiag[i1], &inc_1);
                SLC_DTRSV("Lower", transl, "NonUnit", &rank, &r[i1 + bsn * ldr], &ldr, &b[i1], &inc_1);
                SLC_DSWAP(&st, &r[i1 + bsn * ldr], &stride_diag, &sdiag[i1], &inc_1);
            }
        }
        return;
    }
    
    /* Default simple solve if BN>1 and UPLO=U? */
    /* The Fortran code handles BN>1 UPLO=U (Back substitution) */
    /* "Solve R*x = b using back substitution" */
    
    if (!tranr) {
        /* R*x = b */
        /* Start with R_{l+1} if ST > 0 */
        i1 = nths; /* 0-based index of last part */
        if (st > 0) {
            /* Solve R_{l+1} x_{l+1} = b_{l+1} */
            /* R_{l+1} at r[i1 + bsn*ldr] */
            rank = ranks[bn]; // Last rank
            if (rank < st) {
                i32 len = st - rank;
                SLC_DCOPY(&len, dum, &inc_0, &b[i1 + rank], &inc_1);
            }
            SLC_DTRSV("Upper", trans, "NonUnit", &rank, &r[i1 + bsn * ldr], &ldr, &b[i1], &inc_1);
            
            /* Update b: b_prev = b_prev - L_k * x_{l+1} */
            /* L_k is part of last block column of R.
               Rows 0 to nths-1. Cols bsn to bsn+st-1.
               Call GEMV: y = alpha*A*x + beta*y
               A = R(0:nths-1, bsn:bsn+st-1).
               x = b(i1:i1+st-1) (solution x_{l+1})
               y = b(0:nths-1)

               Formula: R x = b.
               [ R_prev  L ] [ x_prev ]   [ b_prev ]
               [   0    R_last ] [ x_last ] = [ b_last ]

               x_last solved, so: R_prev x_prev = b_prev - L x_last
            */
            f64 neg_one = -1.0;
            SLC_DGEMV(trans, &nths, &st, &neg_one, &r[0 + bsn * ldr], &ldr,
                      &b[i1], &inc_1, &one, b, &inc_1);
        }
        
        /* Solve remaining diagonal blocks R_k */
        /* They are decoupled now? 
           "R is n-by-n block upper triangular ... with upper triangular submatrices R_k ... first l of same order BSN"
           Wait, structure:
           R1 0 ... 0 | L1
           0 R2 ... 0 | L2
           
           Yes, R_k are diagonal blocks.
           Since we removed L_k parts, we just solve R_k x_k = b_k.
        */
        
        for (k = bn - 1; k >= 0; k--) {
            /* Block k */
            i1 = k * bsn;
            rank = ranks[k];
            /* R_k at r[i1 + 0*ldr] */
            /* B_k at b[i1] */
            
            if (rank < bsn) {
                i32 len = bsn - rank;
                SLC_DCOPY(&len, dum, &inc_0, &b[i1 + rank], &inc_1);
            }
            SLC_DTRSV("Upper", trans, "NonUnit", &rank, &r[i1], &ldr, &b[i1], &inc_1);
        }
    } else {
        /* R'*x = b */
        /* [ R_prev'  0 ] [ x_prev ] = [ b_prev ]
           [ L'    R_last' ] [ x_last ] = [ b_last ]
           
           R_prev' x_prev = b_prev. (Decoupled)
           L' x_prev + R_last' x_last = b_last
           R_last' x_last = b_last - L' x_prev
        */
        
        /* Solve x_prev blocks */
        for (k = 0; k < bn; k++) {
            i1 = k * bsn;
            rank = ranks[k];
            if (rank < bsn) {
                i32 len = bsn - rank;
                SLC_DCOPY(&len, dum, &inc_0, &b[i1 + rank], &inc_1);
            }
            SLC_DTRSV("Upper", trans, "NonUnit", &rank, &r[i1], &ldr, &b[i1], &inc_1);
        }
        
        if (st > 0) {
            i1 = nths;
            rank = ranks[bn];
            
            /* Update b_last: b_last = b_last - L' x_prev */
            /* L is NTHS x ST. L' is ST x NTHS.
               x_prev is b(0:nths-1).
               target is b(i1:i1+st-1).
            */
            f64 neg_one = -1.0;
            SLC_DGEMV("Transpose", &nths, &st, &neg_one, &r[0 + bsn * ldr], &ldr, 
                      b, &inc_1, &one, &b[i1], &inc_1);
                      
            /* Solve R_last' x_last = b_last */
            if (rank < st) {
                i32 len = st - rank;
                SLC_DCOPY(&len, dum, &inc_0, &b[i1 + rank], &inc_1);
            }
            SLC_DTRSV("Upper", trans, "NonUnit", &rank, &r[i1 + bsn * ldr], &ldr, &b[i1], &inc_1);
        }
    }
}
