//------------------------------------------------------------------------------
// GB_subref.h: definitions for GB_subref_* functions
//------------------------------------------------------------------------------

// SuiteSparse:GraphBLAS, Timothy A. Davis, (c) 2017-2020, All Rights Reserved.
// http://suitesparse.com   See GraphBLAS/Doc/License.txt for license.

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

#ifndef GB_SUBREF_H
#define GB_SUBREF_H
#include "GB_ij.h"

GB_PUBLIC   // accessed by the MATLAB tests in GraphBLAS/Test only
GrB_Info GB_subref              // C = A(I,J): either symbolic or numeric
(
    // output
    GrB_Matrix *Chandle,
    // input, not modified
    const bool C_is_csc,        // requested format of C
    const GrB_Matrix A,
    const GrB_Index *I,         // index list for C = A(I,J), or GrB_ALL, etc.
    const int64_t ni,           // length of I, or special
    const GrB_Index *J,         // index list for C = A(I,J), or GrB_ALL, etc.
    const int64_t nj,           // length of J, or special
    const bool symbolic,        // if true, construct Cx as symbolic
    const bool must_sort,       // if true, must return C sorted
    GB_Context Context
) ;

GrB_Info GB_subref_phase0
(
    // output
    int64_t *GB_RESTRICT *p_Ch,         // Ch = C->h hyperlist, or NULL standard
    int64_t *GB_RESTRICT *p_Ap_start,   // A(:,kA) starts at Ap_start [kC]
    int64_t *GB_RESTRICT *p_Ap_end,     // ... and ends at Ap_end [kC] - 1
    int64_t *p_Cnvec,       // # of vectors in C
    bool *p_need_qsort,     // true if C must be sorted
    int *p_Ikind,           // kind of I
    int64_t *p_nI,          // length of I
    int64_t Icolon [3],     // for GB_RANGE, GB_STRIDE
    int64_t *p_nJ,          // length of J
    // input, not modified
    const GrB_Matrix A,
    const GrB_Index *I,     // index list for C = A(I,J), or GrB_ALL, etc.
    const int64_t ni,       // length of I, or special
    const GrB_Index *J,     // index list for C = A(I,J), or GrB_ALL, etc.
    const int64_t nj,       // length of J, or special
    const bool must_sort,   // true if C must be returned sorted
    GB_Context Context
) ;

GrB_Info GB_subref_slice
(
    // output:
    GB_task_struct **p_TaskList,    // array of structs, of size max_ntasks
    int *p_max_ntasks,              // size of TaskList
    int *p_ntasks,                  // # of tasks constructed
    int *p_nthreads,                // # of threads for subref operation
    bool *p_post_sort,              // true if a final post-sort is needed
    int64_t *GB_RESTRICT *p_Mark,      // for I inverse, if needed; size avlen
    int64_t *GB_RESTRICT *p_Inext,     // for I inverse, if needed; size nI
    int64_t *p_nduplicates,         // # of duplicates, if I inverse computed
    // from phase0:
    const int64_t *GB_RESTRICT Ap_start,   // location of A(imin:imax,kA)
    const int64_t *GB_RESTRICT Ap_end,
    const int64_t Cnvec,            // # of vectors of C
    const bool need_qsort,          // true if C must be sorted
    const int Ikind,                // GB_ALL, GB_RANGE, GB_STRIDE or GB_LIST
    const int64_t nI,               // length of I
    const int64_t Icolon [3],       // for GB_RANGE and GB_STRIDE
    // original input:
    const int64_t avlen,            // A->vlen
    const int64_t anz,              // nnz (A)
    const GrB_Index *I,
    GB_Context Context
) ;

GrB_Info GB_subref_phase1               // count nnz in each C(:,j)
(
    int64_t *GB_RESTRICT *Cp_handle,       // output of size Cnvec+1
    int64_t *Cnvec_nonempty,            // # of non-empty vectors in C
    // tasks from phase0b:
    GB_task_struct *GB_RESTRICT TaskList,  // array of structs
    const int ntasks,                   // # of tasks
    const int nthreads,                 // # of threads to use
    const int64_t *Mark,                // for I inverse buckets, size A->vlen
    const int64_t *Inext,               // for I inverse buckets, size nI
    const int64_t nduplicates,          // # of duplicates, if I inverted
    // analysis from phase0:
    const int64_t *GB_RESTRICT Ap_start,
    const int64_t *GB_RESTRICT Ap_end,
    const int64_t Cnvec,
    const bool need_qsort,
    const int Ikind,
    const int64_t nI,
    const int64_t Icolon [3],
    // original input:
    const GrB_Matrix A,
    const GrB_Index *I,         // index list for C = A(I,J), or GrB_ALL, etc.
    const bool symbolic,
    GB_Context Context
) ;

GrB_Info GB_subref_phase2   // C=A(I,J)
(
    GrB_Matrix *Chandle,    // output matrix (unallocated on input)
    // from phase1:
    const int64_t *GB_RESTRICT Cp,         // vector pointers for C
    const int64_t Cnvec_nonempty,       // # of non-empty vectors in C
    // from phase0b:
    const GB_task_struct *GB_RESTRICT TaskList,    // array of structs
    const int ntasks,                           // # of tasks
    const int nthreads,                         // # of threads to use
    const bool post_sort,               // true if post-sort needed
    const int64_t *Mark,                // for I inverse buckets, size A->vlen
    const int64_t *Inext,               // for I inverse buckets, size nI
    const int64_t nduplicates,          // # of duplicates, if I inverted
    // from phase0:
    const int64_t *GB_RESTRICT Ch,
    const int64_t *GB_RESTRICT Ap_start,
    const int64_t *GB_RESTRICT Ap_end,
    const int64_t Cnvec,
    const bool need_qsort,
    const int Ikind,
    const int64_t nI,
    const int64_t Icolon [3],
    const int64_t nJ,
    // original input:
    const bool C_is_csc,        // format of output matrix C
    const GrB_Matrix A,
    const GrB_Index *I,
    const bool symbolic,
    GB_Context Context
) ;

GrB_Info GB_I_inverse           // invert the I list for C=A(I,:)
(
    const GrB_Index *I,         // list of indices, duplicates OK
    int64_t nI,                 // length of I
    int64_t avlen,              // length of the vectors of A
    // outputs:
    int64_t *GB_RESTRICT *p_Mark,  // head pointers for buckets, size avlen
    int64_t *GB_RESTRICT *p_Inext, // next pointers for buckets, size nI
    int64_t *p_ndupl,           // number of duplicate entries in I
    GB_Context Context
) ;

//------------------------------------------------------------------------------
// GB_subref_method: select a method for C(:,kC) = A(I,kA), for one vector of C
//------------------------------------------------------------------------------

// Determines the method used for to construct C(:,kC) = A(I,kA) for a
// single vector of C and A.

static inline int GB_subref_method  // return the method to use (1 to 12)
(
    // output
    int64_t *p_work,                // work required
    bool *p_this_needs_I_inverse,   // true if I needs to be inverted
    // input:
    const int64_t ajnz,             // nnz (A (:,j))
    const int64_t avlen,            // A->vlen
    const int Ikind,                // GB_ALL, GB_RANGE, GB_STRIDE, or GB_LIST
    const int64_t nI,               // length of I
    const bool I_inverse_ok,        // true if I is invertable 
    const bool need_qsort,          // true if C(:,k) requires sorting
    const int64_t iinc,             // increment for GB_STRIDE
    const int64_t nduplicates       // # of duplicates in I (zero if not known) 
)
{

    //--------------------------------------------------------------------------
    // initialize return values
    //--------------------------------------------------------------------------

    int method ;            // determined below
    bool this_needs_I_inverse = false ; // most methods do not need I inverse
    int64_t work ;          // most methods require O(nnz(A(:,j))) work

    //--------------------------------------------------------------------------
    // determine the method to use for C(:,j) = A (I,j)
    //--------------------------------------------------------------------------

    if (ajnz == avlen)
    {
        // A(:,j) is dense
        if (Ikind == GB_ALL)
        { 
            // Case 1: C(:,k) = A(:,j) are both dense
            method = 1 ;
            work = nI ;   // ajnz == avlen == nI
        }
        else
        { 
            // Case 2: C(:,k) = A(I,j), where A(:,j) is dense,
            // for Ikind == GB_RANGE, GB_STRIDE, or GB_LIST
            method = 2 ;
            work = nI ;
        }
    }
    else if (nI == 1)
    { 
        // Case 3: one index
        method = 3 ;
        work = 1 ;
    }
    else if (Ikind == GB_ALL)
    { 
        // Case 4: I is ":"
        method = 4 ;
        work = ajnz ;
    }
    else if (Ikind == GB_RANGE)
    { 
        // Case 5: C (:,k) = A (ibegin:iend,j)
        method = 5 ;
        work = ajnz ;
    }
    else if ((Ikind == GB_LIST && !I_inverse_ok) ||  // must do Case 6
        (64 * nI < ajnz))    // Case 6 faster
    { 
        // Case 6: nI not large; binary search of A(:,j) for each i in I
        method = 6 ;
        work = nI * 64 ;
    }
    else if (Ikind == GB_STRIDE)
    { 
        if (iinc >= 0)
        { 
            // Case 7: I = ibegin:iinc:iend with iinc >= 0
            method = 7 ;
            work = ajnz ;
        }
        else if (iinc < -1)
        { 
            // Case 8: I = ibegin:iinc:iend with iinc < =1
            method = 8 ;
            work = ajnz ;
        }
        else // iinc == -1
        { 
            // Case 9: I = ibegin:(-1):iend
            method = 9 ;
            work = ajnz ;
        }
    }
    else // Ikind == GB_LIST, and I inverse buckets will be used
    {
        // construct the I inverse buckets
        this_needs_I_inverse = true ;
        if (need_qsort)
        { 
            // Case 10: nI large, need qsort
            // duplicates are possible so cjnz > ajnz can hold.  If fine tasks
            // use this method, a post sort is needed when all tasks are done.
            method = 10 ;
            work = ajnz * 32 ;
        }
        else if (nduplicates > 0)
        { 
            // Case 11: nI large, no qsort, with duplicates
            // duplicates are possible so cjnz > ajnz can hold.  Note that the
            // # of duplicates is only known after I is inverted, which might
            // not yet be done.  In that case, nuplicates is assumed to be
            // zero, and Case 11 is assumed to be used instead.  This is
            // revised after I is inverted.
            method = 11 ;
            work = ajnz * 2 ;
        }
        else
        { 
            // Case 12: nI large, no qsort, no dupl
            method = 12 ;
            work = ajnz ;
        }
    }

    //--------------------------------------------------------------------------
    // return result
    //--------------------------------------------------------------------------

    if (p_work != NULL)
    { 
        (*p_work) = work ;
    }
    if (p_this_needs_I_inverse != NULL)
    { 
        (*p_this_needs_I_inverse) = this_needs_I_inverse ;
    }
    return (method) ;
}

#endif

