//------------------------------------------------------------------------------
// GB_Vector_diag: extract a diagonal from a matrix, as a vector
//------------------------------------------------------------------------------

// SuiteSparse:GraphBLAS, Timothy A. Davis, (c) 2017-2025, All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//------------------------------------------------------------------------------

#define GB_FREE_WORKSPACE   \
{                           \
    GB_Matrix_free (&T) ;   \
}

#define GB_FREE_ALL         \
{                           \
    GB_FREE_WORKSPACE ;     \
    GB_phybix_free (V) ;    \
}

#include "diag/GB_diag.h"
#include "select/GB_select.h"
#include "scalar/GB_Scalar_wrap.h"

GrB_Info GB_Vector_diag     // extract a diagonal from a matrix, as a vector
(
    GrB_Matrix V,                   // output vector (as an n-by-1 matrix)
    const GrB_Matrix A,             // input matrix
    int64_t k,
    GB_Werk Werk
)
{

    //--------------------------------------------------------------------------
    // check inputs
    //--------------------------------------------------------------------------

    GrB_Info info ;
    ASSERT_MATRIX_OK (A, "A input for GB_Vector_diag", GB0) ;
    ASSERT_MATRIX_OK (V, "V input for GB_Vector_diag", GB0) ;
    ASSERT (GB_VECTOR_OK (V)) ;             // V is a vector on input
    ASSERT (!GB_any_aliased (A, V)) ;       // A and V cannot be aliased
    ASSERT (!GB_IS_HYPERSPARSE (V)) ;       // vectors cannot be hypersparse

    struct GB_Matrix_opaque T_header ;
    GrB_Matrix T = NULL ;

    GrB_Type atype = A->type ;
    GrB_Type vtype = V->type ;
    int64_t nrows = GB_NROWS (A) ;
    int64_t ncols = GB_NCOLS (A) ;
    int64_t n ;
    if (k >= ncols || k <= -nrows)
    { 
        // output vector V must have zero length
        n = 0 ;
    }
    else if (k >= 0)
    { 
        // if k is in range 0 to n-1, V must have length min (m,n-k)
        n = GB_IMIN (nrows, ncols - k) ;
    }
    else
    { 
        // if k is in range -1 to -m+1, V must have length min (m+k,n)
        n = GB_IMIN (nrows + k, ncols) ;
    }

    if (n != V->vlen)
    { 
        GB_ERROR (GrB_DIMENSION_MISMATCH,
            "Input vector must have size " GBd "\n", n) ;
    }

    if (!GB_Type_compatible (atype, vtype))
    { 
        GB_ERROR (GrB_DOMAIN_MISMATCH, "Input matrix of type [%s] "
            "cannot be typecast to output of type [%s]\n",
            atype->name, vtype->name) ;
    }

    //--------------------------------------------------------------------------
    // finish any pending work in A and clear the output vector V
    //--------------------------------------------------------------------------

    GB_MATRIX_WAIT (A) ;
    GB_phybix_free (V) ;

    //--------------------------------------------------------------------------
    // handle the CSR/CSC format of A
    //--------------------------------------------------------------------------

    bool csc = A->is_csc ;
    if (!csc)
    { 
        // The kth diagonal of a CSC matrix is the same as the (-k)th diagonal
        // of the CSR format, so if A is CSR, negate the value of k.  Then
        // treat A as if it were CSC in the rest of this method.
        k = -k ;
    }

    //--------------------------------------------------------------------------
    // extract the kth diagonal of A into the temporary hypersparse matrix T
    //--------------------------------------------------------------------------

    struct GB_Scalar_opaque scalar_header ;
    GrB_Scalar scalar = GB_Scalar_wrap (&scalar_header, GrB_INT64, &k) ;
    GB_CLEAR_MATRIX_HEADER (T, &T_header) ;
    GB_OK (GB_selector (T, GrB_DIAG, false, A, scalar, Werk)) ;
    GB_OK (GB_convert_any_to_hyper (T, Werk)) ;
    GB_MATRIX_WAIT (T) ;
    ASSERT_MATRIX_OK (T, "T = diag (A,k)", GB0) ;

    //--------------------------------------------------------------------------
    // transplant the pattern of T into the sparse vector V
    //--------------------------------------------------------------------------

    int64_t vnz = GB_nnz (T) ;
    float bitmap_switch = V->bitmap_switch ;
    int sparsity_control = V->sparsity_control ;

    bool Vp_is_32 = T->p_is_32 ;
    bool Vj_is_32 = T->j_is_32 ;
    bool Vi_is_32 = (k >= 0) ? T->i_is_32 : T->j_is_32 ;

    GB_OK (GB_new (&V, // existing header
        vtype, n, 1, GB_ph_malloc, true, GxB_SPARSE,
        GxB_NEVER_HYPER, 1, Vp_is_32, Vj_is_32, Vi_is_32)) ;

    V->sparsity_control = sparsity_control ;
    V->bitmap_switch = bitmap_switch ;
    V->iso = T->iso ;       // OK
    if (V->iso)
    { 
        GBURBLE ("(iso diag) ") ;
    }

    GB_Cp_DECLARE (Vp, ) ; GB_Cp_PTR (Vp, V) ;
    GB_ISET (Vp, 0, 0) ;    // Vp [0] = 0 ;
    GB_ISET (Vp, 1, vnz) ;  // Vp [1] = vnz ;

    V->nvals = vnz ;
    if (k >= 0)
    { 
        // transplant T->i into V->i
        V->i = T->i ;
        V->i_size = T->i_size ;
        T->i_shallow = true ;
    }
    else
    { 
        // transplant T->h into V->i
        V->i = T->h ;
        V->i_size = T->h_size ;
        T->h_shallow = true ;
    }

    //--------------------------------------------------------------------------
    // transplant or typecast the values from T to V
    //--------------------------------------------------------------------------

    GB_Type_code vcode = vtype->code ;
    GB_Type_code acode = atype->code ;
    if (vcode == acode)
    { 
        // transplant T->x into V->x
        V->x = T->x ;
        V->x_size = T->x_size ;
        T->x = NULL ;
    }
    else
    {
        // V->x = (vtype) T->x
        // V is sparse so malloc is OK
        V->x = GB_XALLOC_MEMORY (false, V->iso, vnz, vtype->size,
            &(V->x_size)) ;
        if (V->x == NULL)
        { 
            // out of memory
            GB_FREE_ALL ;
            return (GrB_OUT_OF_MEMORY) ;
        }
        GB_OK (GB_cast_matrix (V, T)) ;
    }

    //--------------------------------------------------------------------------
    // finalize the vector V
    //--------------------------------------------------------------------------

    V->jumbled = T->jumbled ;
//  V->nvec_nonempty = (vnz == 0) ? 0 : 1 ;
    GB_nvec_nonempty_set ((GrB_Matrix) V, (vnz == 0) ? 0 : 1) ;
    V->magic = GB_MAGIC ;

    //--------------------------------------------------------------------------
    // free workspace, conform V to its desired format, and return result
    //--------------------------------------------------------------------------

    GB_FREE_WORKSPACE ;
    GB_OK (GB_conform (V, Werk)) ;
    ASSERT_MATRIX_OK (V, "V output for GB_Vector_diag", GB0) ;
    return (GrB_SUCCESS) ;
}

