# This file is a part of Julia. License is MIT: https://julialang.org/license

## getindex
const rep = 20
const reps = rep^2

function sparse_getindex_perf()
    ## create some sparse matrices (need to be square):
    seed = 1
    srand(seed)

    small = 10^3
    med = 10^4
    large = 10^5
    huge = 10^6
    # # 1 entry per line
    # ss = []
    # push!(ss, sprand(small, small, 1e-3))
    # #push!(ss, sprand(med, med, 1e-4))
    # push!(ss, sprand(large, large, 1e-5))

    # 10 entries per line
    ts = []
    push!(ts, sprand(small, small, 1e-2))
    #push!(ts, sprand(med, med, 1e-3))
    push!(ts, sprand(large, large, 1e-4))

    # 100 entries per line
    us = []
    push!(us, sprand(small, small, 1e-1))
    #push!(us, sprand(med, med, 1e-2))
    push!(us, sprand(large, large, 1e-3))
    #push!(us, sprand(huge, huge, 1e-4))

    # # 1000 entries per line
    # vs = []
    # push!(vs, sprand(small, small, 1.0))
    # #push!(vs, sprand(med, med, 1e-1))
    # push!(vs, sprand(large, large, 1e-2))
    # #push!(vs, sprand(huge, huge, 1e-3))

    ## using uint32 (works up to 10^9)
    uus = []
    for u in us
        push!(uus, SparseMatrixCSC(u.m, u.n, map(UInt32,u.colptr), map(UInt32,u.rowval), u.nzval))
    end

    # test performance with matrices in us, uus and ts for
    # - integer indexing
    # - range indexing
    # - ordered array indexing
    # - disordered array indexing
    # - 1d indexing with integers and booleans

    # setup:
    # - indices
    intinds = nothing
    logicalinds = nothing # needs to be generated for a specific matrix size.
    rangeinds = 121:237
    orderedinds = collect(rangeinds)
    disorderedinds = orderedinds[randperm(length(orderedinds))]

    inds = [(intinds, "integers"), (logicalinds, "logical array"), (rangeinds, "a range"),
            (orderedinds, "a ordered array"), (disorderedinds, "a disordered array")]

    # - matrix sizes
    sizes = [(1, "small", "Small sparse matrix"), (2, "medium", "Medium sparse matrix")]

    # - matrix types
    mattyp = [(ts, "10 entries/column"), (us, "100 entries/column"), (uus, "100 entries/column uint32")]

    # - functions
    funs = [(integer_indexing, 1, "indexing"), (one_arg_indexing, 1, "1d indexing"),
            (row_indexing, 2, "indexing rows"), (col_indexing, 2, "indexing columns"),
            (row_col_indexing, 3, "indexing rows & columns")]

    # performance tests:
    counters = [1,1] # for small and medium matrix
    # integer indexing
    for (sz,s1,s2) in sizes # size of the matrix
        for (mt, ms) in mattyp  # type of the matrix
            m = mt[sz]
            c = counters[sz]
            @timeit integer_indexing(m) "sparse_getindex_$s1$c" "$s2 with $ms, indexing with integers"
            counters[sz] += 1
        end
    end

    # range & array indexing
    for (sz,s1,s2) in sizes # size of the matrix
        for (mt, ms) in mattyp  # type of the matrix
            m = mt[sz]
            for (fun, nargs, funstr) in funs[3:5] # indexing test function
                for (ind,indstr) in inds[3:5] # type of indices
                    c = counters[sz]
                    if indstr=="logical array"
                        # make a logical array of the right size
                        ind = sprand(Bool, size(m,1)..., 1e-5)
                    end
                    if nargs==2
                        @timeit fun(m, ind)  "sparse_getindex_$s1$c" "Sparse matrix with $ms, $funstr with $indstr"
                    elseif nargs==3
                        @timeit fun(m, ind, ind) "sparse_getindex_$s1$c" "Sparse matrix with $ms, $funstr with $indstr"
                    else
                        error("Something is amiss here.")
                    end
                    counters[sz] += 1
                end
            end
        end
    end

    # linear indexing
    for (sz,s1,s2) in sizes # size of the matrix
        for (mt, ms) in mattyp  # type of the matrix
            m = mt[sz]
            for (ind,indstr) in inds[2:5] # type of indices
                if indstr=="logical array"
                    if sz==2
                        continue # logical indexing with medium size sparse matrix takes too long
                    end
                    # make a logical array of the right size
                    ind = sprand(Bool, size(m)..., 1e-5)
                    c = counters[sz]
                    @timeit one_arg_indexing(m, ind) "sparse_getindex_$s1$c" "$s2 with $ms, linear indexing with $indstr"
                    counters[sz] += 1
                else
                    c = counters[sz]
                    @timeit one_arg_indexing(m, ind) "sparse_getindex_$s1$c" "$s2 with $ms, linear indexing with $indstr"
                    counters[sz] += 1
                end
            end
        end
    end

end

function integer_indexing(A)
    # index with two random integers
    nI, nJ = size(A)
    rI = 1:nI
    rJ = 1:nJ
    tmp = zero(eltype(A))
    for i in rand(rI, reps)
        for j in rand(rJ, rep)
            tmp += A[i,j]
        end
    end
    tmp
end

function row_indexing(A, rowinds)
    # index rows with rowinds and columns with a random integer
    nI, nJ = size(A)
    rI = 1:nI
    rJ = 1:nJ
    tmp = zero(eltype(A))
    for j in rand(rJ, reps)
        tmp += sum(A[rowinds,j])
    end
    tmp
end

function col_indexing(A, colinds)
    # index rows with a random integer and columns with colinds
    nI, nJ = size(A)
    rI = 1:nI
    rJ = 1:nJ
    tmp = zero(eltype(A))
    for i in rand(rI, div(reps,10) )
        tmp += sum(A[i,colinds])
    end
    tmp
end

function row_col_indexing(A, rowinds, colinds)
    # index rows with rowinds and columns with colinds
    # we need:
    (maximum(rowinds)+rep < size(A,1) && maximum(colinds)+rep < size(A, 2)) || error("bad rowinds or colinds")
    nI, nJ = size(A)
    rI = 1:nI
    rJ = 1:nJ
    for i in 1:10
        for j in 1:10
            tmp2 = A[rowinds.+i, colinds.+j]
        end
    end
end

function one_arg_indexing(A, lininds)
    # This is for 1d-indexing and indexing with one array of logicals.
    # Both return a nx1 sparse matrix.
    tmp = zero(eltype(A))
    if isa(eltype(A), Bool)
        tmp = sum(A[lininds])
    else
        for i in 1:rep
            tmp += sum(A[lininds])
        end
    end
    tmp
end
