text:
#=
Copyright (c) 2016, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

#using Debug

using CompilerTools.LivenessAnalysis
using CompilerTools.TransitiveDependence

function getArrayDistributionInfo(ast, state)
    set_user_partitionings(state)
    before_arr_partitionings = [state.arrs_dist_info[arr].partitioning for arr in keys(state.arrs_dist_info)]

    while true
        @dprintln(3,"DistPass state before array info walk: ",state)
        ParallelIR.AstWalk(ast, get_arr_dist_info, state)
        @dprintln(3,"DistPass state after array info walk: ",state)
        set_user_partitionings(state)
        new_arr_partitionings = [state.arrs_dist_info[arr].partitioning for arr in keys(state.arrs_dist_info)]
        # break if no new sequential array discovered
        if new_arr_partitionings==before_arr_partitionings
            break
        end
        before_arr_partitionings = new_arr_partitionings
    end

    # if all parfors are sequential
    if HPAT.force_parallel && all([state.parfor_partitioning[parfor_id]==SEQ for parfor_id in keys(state.parfor_partitioning)]) &&
          all([state.arrs_dist_info[arr].partitioning==SEQ for arr in keys(state.arrs_dist_info)])
        error("HPAT failed to parallelize! Fix the parallelism problem or use \"HPAT.setForceParallel(false)\" to run sequentially.")
    end
    # for debugging purposes
    save_array_partitionings(state)
end

function set_user_partitionings(state)
    for (arr,part) in state.user_partitionings
      state.arrs_dist_info[arr].partitioning = part
    end
end

function save_array_partitionings(state)
    d = Dict{LHSVar,Partitioning}()
    for (k,v) in state.arrs_dist_info
        d[k] = state.arrs_dist_info[k].partitioning
    end
    HPAT.set_saved_array_partitionings(d)
end

"""
mark sequential arrays
"""
function get_arr_dist_info(node::Expr, state::DistPassState, top_level_number, is_top_level, read)
    head = node.head
    # arrays written in parfors are ok for now

    @dprintln(3,"DistPass arr info walk Expr head: ", head)
    if head==:(=)
        @dprintln(3,"DistPass arr info walk assignment: ", node)
        lhs = toLHSVar(node.args[1])
        rhs = node.args[2]
        return get_arr_dist_info_assignment(node, state, top_level_number, lhs, rhs)
    elseif head==:parfor
        @dprintln(3,"DistPass arr info walk parfor: ", node)
        parfor = getParforNode(node)
        return get_arr_dist_info_parfor(node, state, top_level_number, parfor)
        # functions dist_ir_funcs are either handled here or do not make arrays sequential
    elseif head==:call && (isa(node.args[1],GlobalRef) || isa(node.args[1],TopNode)) && in(node.args[1].name, dist_ir_funcs)
        func = node.args[1].name
        if func==:__hpat_data_source_HDF5_read || func==:__hpat_data_source_TXT_read
            @dprintln(2,"DistPass arr info walk data source read ", node)
            # will be parallel IO, intentionally do nothing
        elseif func==:Kmeans
            @dprintln(2,"DistPass arr info walk kmeans ", node)
            # first array is cluster output and is sequential
            # second array is input matrix and is parallel
            setSEQ(toLHSVar(node.args[2]),state)
        elseif func==:LinearRegression || func==:NaiveBayes
            @dprintln(2,"DistPass arr info walk LinearRegression/NaiveBayes ", node)
            # first array is cluster output and is sequential
            # second array is input matrix and is parallel
            # third array is responses and is parallel
            setSEQ(toLHSVar(node.args[2]),state)
        elseif func==:__hpat_join
            # output columns of join can have variable length on different processors
            # set partitioning to ONE_D_VAR
            # node.args[3] gives number of output columns
            start_col_ind = 6
            end_col_ind = start_col_ind + node.args[3] - 1
            for col_ind in start_col_ind:end_col_ind
                col = node.args[col_ind]
                # table columns are 1D
                @assert length(state.arrs_dist_info[col].dim_sizes)==1
                state.arrs_dist_info[col].partitioning = ONE_D_VAR
            end
        elseif func==:__hpat_filter
            # output columns of filter can have variable length on different processors
            # set partitioning to ONE_D_VAR
            # node.args[4] gives number of output columns
            start_col_ind = 5
            end_col_ind = start_col_ind + node.args[4] - 1
            for col_ind in start_col_ind:end_col_ind
                col = node.args[col_ind]
                # table columns are 1D
                @assert length(state.arrs_dist_info[col].dim_sizes)==1
                state.arrs_dist_info[col].partitioning = ONE_D_VAR
            end
        elseif func==:__hpat_aggregate
            # output columns of aggregate can have variable length on different processors
            # set partitioning to ONE_D_VAR
            # node.args[4] gives number of output columns
            # start index is 2*node.args[4] because we need to skips expression columns and functions
            start_col_ind = (2*node.args[4]) + 5
            end_col_ind = start_col_ind + node.args[4] - 1
            # one extra output to account for output key column
            end_col_ind += 1
            for col_ind in start_col_ind:end_col_ind
                col = node.args[col_ind]
                # table columns are 1D
                @assert length(state.arrs_dist_info[col].dim_sizes)==1
                state.arrs_dist_info[col].partitioning = ONE_D_VAR
            end
        end
        return node
    elseif head==:gotoifnot
        @dprintln(3,"DistPass arr info gotoifnot: ", node)
        return CompilerTools.AstWalker.ASTWALK_RECURSE
        # arrays written in serial code are not distributed
    elseif head!=:body && head!=:block && head!=:lambda
        @dprintln(3,"DistPass arr info walk serial code: ", node)
        return get_arr_dist_info_serial_code(node, state, top_level_number)
    end
    return CompilerTools.AstWalker.ASTWALK_RECURSE
end

function get_arr_dist_info_serial_code(node, state, top_level_number)
    live_info = CompilerTools.LivenessAnalysis.from_lambda(state.LambdaVarInfo,
                                                           TypedExpr(nothing, :body, node), ParallelIR.pir_live_cb, state.LambdaVarInfo)
    # @dprintln(3, "body_lives = ", body_lives)
    # live_info = CompilerTools.LivenessAnalysis.find_top_number(top_level_number, state.lives)
    # all_vars = union(live_info.def, live_info.use)
    all_vars = []

    for bb in collect(values(live_info.basic_blocks))
        for stmt in bb.statements
            append!(all_vars, collect(union(stmt.def, stmt.use)))
        end
    end

    @dprintln(3,"DistPass arr info walk serial code vars: ", all_vars)
    # ReadWriteSet is not robust enough now
    #rws = CompilerTools.ReadWriteSet.from_exprs([node], ParallelIR.pir_live_cb, state.LambdaVarInfo)
    #readArrs = collect(keys(rws.readSet.arrays))
    #writeArrs = collect(keys(rws.writeSet.arrays))
    #allArrs = [readArrs;writeArrs]

    for var in all_vars
        if haskey(state.arrs_dist_info, toLHSVar(var))
            @dprintln(2,"DistPass arr info walk array sequential since in serial code: ", var, "  ",lookupVariableName(var, state.LambdaVarInfo), " ", node)
            setSEQ(toLHSVar(var),state)
        end
    end
    return node
end

function get_arr_dist_info_parfor(node, state, top_level_number, parfor)
    rws = CompilerTools.ReadWriteSet.from_exprs(parfor.body, ParallelAccelerator.ParallelIR.pir_rws_cb, state.LambdaVarInfo)
    partitioning = ONE_D

    if length(parfor.arrays_read_past_index)!=0 || length(parfor.arrays_written_past_index)!=0
        @dprintln(2,"DistPass arr info walk parfor sequential: ", node)
        partitioning = SEQ
    end

    indexVariable = toLHSVar(parfor.loopNests[1].indexVariable)

    allArrAccesses = merge(rws.readSet.arrays,rws.writeSet.arrays)
    myArrs = LHSVar[]

    body_lives = CompilerTools.LivenessAnalysis.from_lambda(state.LambdaVarInfo,
                                                            parfor.body, ParallelIR.pir_live_cb, state.LambdaVarInfo)
    # @dprintln(3, "body_lives = ", body_lives)
    @dprintln(2,"DistPass arr info walk parfor indexVariable: ", indexVariable, " accesses: ", allArrAccesses)
    @dprintln(2,"parfor = ", parfor)
    # If an array is accessed with a Parfor's index variable, the parfor and array should have same partitioning
    for arr in keys(allArrAccesses)
        # an array can be accessed multiple times in Pafor
        # for each access:
        for access_indices in allArrAccesses[arr]
            @dprintln(3, "access_indices = ", access_indices)
            indices = map(toLHSVar,access_indices)
            # if array would be accessed in parallel in this Parfor
            if indices[end]==indexVariable
                push!(myArrs, arr)
            end
            # An array access index can be dependent on parfor's
            # index variable as in nested comprehension case of K-Means.
            # Parfor can't be parallelized in general cases since array can't be partitioned properly.
            # ParallelIR should optimize out the trivial cases where indices are essentially equal (i=1+1*index-1 in k-means)
            if isAccessIndexDependent(indices, indexVariable, body_lives, state)
                #push!(myArrs, arr)
                @dprintln(2,"DistPass arr info walk arr index dependent: ",arr," ", indices, " ", indexVariable)
                partitioning = SEQ
            end
            # sequential if not accessed column major (last dimension)
            # TODO: generalize?
            if in(indexVariable, indices[1:end-1])
                @dprintln(2,"DistPass arr info walk arr index sequential: ",arr," ", indices, " ", indexVariable)
                partitioning = SEQ
            end
        end
    end
    # keep mapping from parfors to arrays
    # state.parfor_info[parfor.unique_id] = myArrs
    @dprintln(3,"DistPass arr info walk parfor arrays: ", myArrs)

    for arr in myArrs
        partitioning = min(partitioning,getArrayPartitioning(arr,state))
        if isSEQ(arr,state)
            # no need to check size for parallel arrays since ParallelIR already used equivalence class info
            # || !eqSize(state.arrs_dist_info[arr].dim_sizes[end], state.arrs_dist_info[myArrs[1]].dim_sizes[end])
            # last dimension of all parfor arrays should be equal since they are partitioned
            @dprintln(2,"DistPass parfor check array: ", arr," sequential: ", isSEQ(arr,state))
        end
    end
    # parfor and all its arrays have same partitioning
    state.parfor_partitioning[parfor.unique_id] = partitioning
    state.parfor_arrays[parfor.unique_id] = myArrs
    for arr in myArrs
        setArrayPartitioning(arr,partitioning,state)
    end
    return CompilerTools.AstWalker.ASTWALK_RECURSE
end

#function isAccessIndexDependent(indices::Vector{Any}, indexVariable::LHSVar, body_lives::BlockLiveness, state)
#    deps = CompilerTools.TransitiveDependence.computeDependencies(body_lives)
function isAccessIndexDependent(indices::Vector{Any}, indexVariable::LHSVar, deps, state)
    # if any index is dependent
    return reduce(|, [ isAccessIndexDependent(indices[i], indexVariable, deps, state) for i in 1:length(indices)] )
end

function isAccessIndexDependent(index::LHSVar, indexVariable::LHSVar, deps::Dict{LHSVar, Set{LHSVar}}, state)
    @dprintln(3,"isAccessIndexDependent for index ", index, " and variable ", indexVariable, " with deps ", deps)
    if index == indexVariable
        return false
    end

    if haskey(deps, index)
        ret = in(indexVariable, deps[index])
        @dprintln(3,"isAccessIndexDependent returned ", ret)
        return ret
    else
        @dprintln(2, "isAccessIndexDependent could not find index ", index, " in deps.")
        return false
    end
end

function isAccessIndexDependent(index::Int, indexVariable::LHSVar, deps::Dict{LHSVar, Set{LHSVar}}, state)
    return false
end

"""
For each statement of parfor's body, see if array access index is in Def set and parfor's indexVariable is in Use set.
This is a hack to work around not having dependence analysis in CompilerTools.
"""
function isAccessIndexDependent(index::LHSVar, indexVariable::LHSVar, body_lives::BlockLiveness, state)
    for bb in collect(values(body_lives.basic_blocks))
        for stmt in bb.statements
            if isBareParfor(stmt.tls.expr)
                return isAccessIndexDependent(index, indexVariable, state.deps, state)
            elseif in(index,stmt.def) && in(indexVariable, stmt.use)
                @dprintln(2,"DistPass arr info walk dependent index found: ", index," ",indexVariable,"  ",stmt)
                return true
            end
        end
    end
    return false
end

# TODO: is this general?
function isAccessIndexDependent(index::Int, indexVariable::LHSVar, body_lives::BlockLiveness, state)
    return false
end

function get_arr_dist_info(ast::Any, state::DistPassState, top_level_number, is_top_level, read)
    return CompilerTools.AstWalker.ASTWALK_RECURSE
end

""" return LHSVar if arg is RHSVar, otherwise no change
    used for allocation sizes which could be LHSVar or Int or TypedVar or Expr
"""
replaceAllocTypedVar(a::TypedVar) = toLHSVar(a)
replaceAllocTypedVar(a::Union{Int,LHSVar,Expr}) = a

function get_arr_dist_info_assignment(node::Expr, state::DistPassState, top_level_number, lhs::LHSVar, rhs::RHSVar)
    rhs = toLHSVar(rhs)
    if haskey(state.arrs_dist_info, rhs)
        state.arrs_dist_info[lhs].dim_sizes = state.arrs_dist_info[rhs].dim_sizes
        state.arrs_dist_info[lhs].starts = state.arrs_dist_info[rhs].starts
        state.arrs_dist_info[lhs].counts = state.arrs_dist_info[rhs].counts
        state.arrs_dist_info[lhs].strides = state.arrs_dist_info[rhs].strides
        state.arrs_dist_info[lhs].blocks = state.arrs_dist_info[rhs].blocks
        state.arrs_dist_info[lhs].local_sizes = state.arrs_dist_info[rhs].local_sizes
        state.arrs_dist_info[lhs].leftovers = state.arrs_dist_info[rhs].leftovers

        # lhs and rhs are sequential if either is sequential
        # partitioning based on precedence, SEQ has highest precedence
        partitioning = min(state.arrs_dist_info[lhs].partitioning, state.arrs_dist_info[rhs].partitioning)
        state.arrs_dist_info[lhs].partitioning = state.arrs_dist_info[rhs].partitioning = partitioning
        @dprintln(3,"DistPass arr info dim_sizes update: ", state.arrs_dist_info[lhs].dim_sizes)
        return node
        # lhs is sequential if rhs is unknown
    elseif haskey(state.arrs_dist_info, lhs)
        @dprintln(3,"DistPass assignment unknown rhs:", rhs,", sequentlial lhs: ",lhs)
        state.arrs_dist_info[lhs].partitioning = SEQ
    end
    return CompilerTools.AstWalker.ASTWALK_RECURSE
end

function get_arr_dist_info_assignment(node::Expr, state::DistPassState, top_level_number, lhs::LHSVar, rhs::Expr)
    if isAllocation(rhs)
        # if allocated already, can't parallelize since we can't replace
        #   the arraysize() calls to global value statically. needs runtime support
        alloc_sizes = map(replaceAllocTypedVar, get_alloc_shape(rhs.args[2:end]))
        if state.arrs_dist_info[lhs].dim_sizes[1]!=0 && state.arrs_dist_info[lhs].dim_sizes!=alloc_sizes
            setSEQ(lhs,state)
            @dprintln(3,"DistPass arr info non-constant allocation found sequential: ", lhs," ",rhs.args[2:end])
        end
        state.arrs_dist_info[lhs].dim_sizes = alloc_sizes
        @dprintln(3,"DistPass arr info dim_sizes update: ", state.arrs_dist_info[lhs].dim_sizes)
    elseif isa(rhs,Expr) && rhs.head==:call && (isa(rhs.args[1],GlobalRef) || isa(rhs.args[1],TopNode)) && in(rhs.args[1].name, dist_ir_funcs)
        func = rhs.args[1]
        if isBaseFunc(func,:reshape)
            # only reshape() with constant tuples handled
            if haskey(state.tuple_table, rhs.args[3])
                state.arrs_dist_info[lhs].dim_sizes = state.tuple_table[rhs.args[3]]
                # TODO: are these required?
                #=
                state.arrs_dist_info[lhs].starts = state.arrs_dist_info[rhs.args[2]].starts
                state.arrs_dist_info[lhs].counts = state.arrs_dist_info[rhs.args[2]].counts
                state.arrs_dist_info[lhs].strides = state.arrs_dist_info[rhs.args[2]].strides
                state.arrs_dist_info[lhs].blocks = state.arrs_dist_info[rhs.args[2]].blocks
                state.arrs_dist_info[lhs].local_sizes = state.arrs_dist_info[rhs.args[2]].local_sizes
                state.arrs_dist_info[lhs].leftovers = state.arrs_dist_info[rhs.args[2]].leftovers
                =#

                @dprintln(3,"DistPass arr info dim_sizes update: ", state.arrs_dist_info[lhs].dim_sizes)
                # lhs and rhs are sequential if either is sequential
                # partitioning based on precedence, SEQ has highest precedence
                partitioning = min(state.arrs_dist_info[lhs].partitioning, state.arrs_dist_info[toLHSVar(rhs.args[2])].partitioning)
                state.arrs_dist_info[lhs].partitioning = state.arrs_dist_info[toLHSVar(rhs.args[2])].partitioning = partitioning
            else
                @dprintln(3,"DistPass arr info reshape tuple not found: ", rhs.args[3]," therefore sequential: ",lhs," ",toLHSVar(rhs.args[2]))
                setSEQ(lhs,state)
                setSEQ(toLHSVar(rhs.args[2]),state)
            end
        elseif isBaseFunc(rhs.args[1],:tuple)
            state.tuple_table[lhs] = [  toLHSVarOrNum(s) for s in rhs.args[2:end] ]
            # TODO: do tuples need to be constant?
            #=
            ok = true
            for s in rhs.args[2:end]
            if !(isa(s,SymbolNode) || isa(s,Int))
            ok = false
            end
            end
            if ok
            state.tuple_table[lhs] = [  toLHSVarOrNum(s) for s in rhs.args[2:end] ]
            @dprintln(3,"DistPass arr info tuple constant: ", lhs," ",rhs.args[2:end])
            else
            @dprintln(3,"DistPass arr info tuple not constant: ", lhs," ",rhs.args[2:end])
            end
            =#
        elseif isBaseFunc(func,:gemm_wrapper!)
            return get_arr_dist_info_gemm(node, state, top_level_number, lhs, rhs)
        elseif isBaseFunc(func,:gemv!)
            return get_arr_dist_info_gemv(node, state, top_level_number, lhs, rhs)
            # TODO: Check why toLHSVar(rhs.args[2])].dim_sizes[1] is zero for all arrays
        elseif func.name == :hcat
            @dprintln(3,"DistPass arr info handling hcat: ", rhs)
            state.arrs_dist_info[lhs].dim_sizes[1] = state.arrs_dist_info[toLHSVar(rhs.args[2])].dim_sizes[1]
            state.arrs_dist_info[lhs].dim_sizes[2] = length(rhs.args) - 1
            min_partitioning = state.arrs_dist_info[lhs].partitioning
            for curr_array_index in 2:length(rhs.args)
                min_partitioning = min(min_partitioning, state.arrs_dist_info[toLHSVar(rhs.args[curr_array_index])].partitioning)
            end
            for curr_array_index in 2:length(rhs.args)
                state.arrs_dist_info[toLHSVar(rhs.args[curr_array_index])].partitioning = min_partitioning
            end
            state.arrs_dist_info[lhs].partitioning = min_partitioning
        elseif func.name == :typed_hcat
            # lhs = Base.typed_hcat(Int64, arr1,...)
            @dprintln(3,"DistPass arr info handling typed_hcat: ", rhs)
            state.arrs_dist_info[lhs].dim_sizes[1] = state.arrs_dist_info[toLHSVar(rhs.args[3])].dim_sizes[1]
            state.arrs_dist_info[lhs].dim_sizes[2] = length(rhs.args) - 2
            min_partitioning = state.arrs_dist_info[lhs].partitioning
            for curr_array_index in 3:length(rhs.args)
                min_partitioning = min(min_partitioning, state.arrs_dist_info[toLHSVar(rhs.args[curr_array_index])].partitioning)
            end
            for curr_array_index in 3:length(rhs.args)
                state.arrs_dist_info[toLHSVar(rhs.args[curr_array_index])].partitioning = min_partitioning
            end
            state.arrs_dist_info[lhs].partitioning = min_partitioning
        elseif func.name == :vcat
            @dprintln(3,"DistPass arr info handling vcat: ", rhs)
            # output and all input arrays have same dim_sizes except first dimension
            state.arrs_dist_info[lhs].dim_sizes = state.arrs_dist_info[toLHSVar(rhs.args[2])].dim_sizes
            # HACK: a call expression summing dimension 1 of all input arrays
            state.arrs_dist_info[lhs].dim_sizes[1] =
               Expr(:call,:(+), map(x->state.arrs_dist_info[x].dim_sizes[1], rhs.args[2:end])...)
            min_partitioning = state.arrs_dist_info[lhs].partitioning
            for curr_array_index in 2:length(rhs.args)
                min_partitioning = min(min_partitioning, state.arrs_dist_info[toLHSVar(rhs.args[curr_array_index])].partitioning)
            end
            for curr_array_index in 2:length(rhs.args)
                state.arrs_dist_info[toLHSVar(rhs.args[curr_array_index])].partitioning = min_partitioning
            end
            state.arrs_dist_info[lhs].partitioning = min_partitioning
        elseif isBaseFunc(func,:transpose!)
            # arg1 of transpose!() is also output and needs to be updated
            out_arg = toLHSVar(rhs.args[2])
            state.arrs_dist_info[lhs].dim_sizes[2] = state.arrs_dist_info[toLHSVar(rhs.args[3])].dim_sizes[1]
            state.arrs_dist_info[out_arg].dim_sizes[2] = state.arrs_dist_info[toLHSVar(rhs.args[3])].dim_sizes[1]
            state.arrs_dist_info[lhs].dim_sizes[1] = state.arrs_dist_info[toLHSVar(rhs.args[3])].dim_sizes[2]
            state.arrs_dist_info[out_arg].dim_sizes[1] = state.arrs_dist_info[toLHSVar(rhs.args[3])].dim_sizes[2]
            partitioning = min(state.arrs_dist_info[lhs].partitioning, state.arrs_dist_info[toLHSVar(rhs.args[3])].partitioning)
            state.arrs_dist_info[lhs].partitioning = state.arrs_dist_info[toLHSVar(rhs.args[3])].partitioning = partitioning
            state.arrs_dist_info[out_arg].partitioning = partitioning
        elseif isBaseFunc(func,:transpose)
            in_arr = toLHSVar(rhs.args[2])
            state.arrs_dist_info[lhs].dim_sizes[2] = state.arrs_dist_info[in_arr].dim_sizes[1]
            state.arrs_dist_info[lhs].dim_sizes[1] = state.arrs_dist_info[in_arr].dim_sizes[2]
            partitioning = min(state.arrs_dist_info[lhs].partitioning, state.arrs_dist_info[in_arr].partitioning)
            state.arrs_dist_info[lhs].partitioning = state.arrs_dist_info[in_arr].partitioning = partitioning
        end
    else
        # lhs is sequential if rhs is unknown
        if haskey(state.arrs_dist_info, lhs)
            @dprintln(3,"DistPass assignment unknown rhs:", rhs,", sequentlial lhs: ",lhs)
            state.arrs_dist_info[lhs].partitioning = SEQ
        end
        return CompilerTools.AstWalker.ASTWALK_RECURSE
    end
    # Dead return
    return node
end

get_arr_dist_info_assignment(node::Expr, state::DistPassState, top_level_number, lhs::ANY, rhs::ANY) = CompilerTools.AstWalker.ASTWALK_RECURSE

function get_arr_dist_info_gemm(node::Expr, state::DistPassState, top_level_number, lhs::LHSVar, rhs::Expr)
    # determine output dimensions
    state.arrs_dist_info[lhs].dim_sizes = state.arrs_dist_info[toLHSVar(rhs.args[2])].dim_sizes
    state.arrs_dist_info[lhs].starts = state.arrs_dist_info[toLHSVar(rhs.args[2])].starts
    state.arrs_dist_info[lhs].counts = state.arrs_dist_info[toLHSVar(rhs.args[2])].counts
    state.arrs_dist_info[lhs].strides = state.arrs_dist_info[toLHSVar(rhs.args[2])].strides
    state.arrs_dist_info[lhs].blocks = state.arrs_dist_info[toLHSVar(rhs.args[2])].blocks
    state.arrs_dist_info[lhs].local_sizes = state.arrs_dist_info[toLHSVar(rhs.args[2])].local_sizes
    state.arrs_dist_info[lhs].leftovers = state.arrs_dist_info[toLHSVar(rhs.args[2])].leftovers

    arr1 = toLHSVar(rhs.args[5])
    t1 = (rhs.args[3]=='T')
    arr2 = toLHSVar(rhs.args[6])
    t2 = (rhs.args[4]=='T')

    partitioning=ONE_D

    # result is sequential if both inputs are sequential
    if isSEQ(arr1,state) && isSEQ(arr2,state)
        partitioning=SEQ
        # result is sequential but with reduction if both inputs are partitioned and second one is transposed
        # e.g. labels*points'
    elseif isONE_D(arr1,state) && isONE_D(arr2,state) && t2 && !t1
        partitioning=SEQ
        # first input is sequential but output is parallel if the second input is partitioned but not transposed
        # e.g. w*points
    elseif isONE_D(arr2,state) && !t2 && !isTWO_D(arr1,state)
        @dprintln(3,"DistPass arr info gemm first input is sequential: ", arr1)
        setSEQ(arr1,state)
        # if no array is sequential and any array is 2D, then all are TWO_D
    elseif !isSEQ(arr1,state) && !isSEQ(arr2,state) && !isSEQ(lhs,state) &&
        ( isTWO_D(arr1,state) || isTWO_D(arr2,state) || isTWO_D(lhs,state) )
        @dprintln(3,"DistPass arr info gemm 2D: ", arr1)
        setTWO_D(arr1,state)
        setTWO_D(arr2,state)
        partitioning=TWO_D
        # otherwise, no known pattern found, every array is sequential
    else
        @dprintln(3,"DistPass arr info gemm all sequential: ", arr1," ", arr2)
        setSEQ(arr1,state)
        setSEQ(arr2,state)
        partitioning=SEQ
    end

    if partitioning==SEQ
        @dprintln(3,"DistPass arr info gemm output is sequential: ", lhs," ",rhs.args[2])
    end
    setArrayPartitioning(lhs,partitioning,state)
    setArrayPartitioning(toLHSVar(rhs.args[2]),partitioning,state)
    return node
end

function get_arr_dist_info_gemv(node::Expr, state::DistPassState, top_level_number, lhs::LHSVar, rhs::Expr)
    # determine output dimensions
    state.arrs_dist_info[lhs].dim_sizes = state.arrs_dist_info[toLHSVar(rhs.args[2])].dim_sizes
    state.arrs_dist_info[lhs].starts = state.arrs_dist_info[toLHSVar(rhs.args[2])].starts
    state.arrs_dist_info[lhs].counts = state.arrs_dist_info[toLHSVar(rhs.args[2])].counts
    state.arrs_dist_info[lhs].strides = state.arrs_dist_info[toLHSVar(rhs.args[2])].strides
    state.arrs_dist_info[lhs].blocks = state.arrs_dist_info[toLHSVar(rhs.args[2])].blocks
    state.arrs_dist_info[lhs].local_sizes = state.arrs_dist_info[toLHSVar(rhs.args[2])].local_sizes
    state.arrs_dist_info[lhs].leftovers = state.arrs_dist_info[toLHSVar(rhs.args[2])].leftovers

    arr1 = toLHSVar(rhs.args[4])
    t1 = (rhs.args[3]=='T')
    arr2 = toLHSVar(rhs.args[5])

    partitioning = ONE_D

    # result is sequential if both inputs are sequential
    if isSEQ(arr1,state) && isSEQ(arr2,state)
        partitioning = SEQ
        # result is sequential but with reduction if both inputs are partitioned and matrix is not transposed (X*y)
    elseif !isSEQ(arr1,state) && !isSEQ(arr2,state) && !t1
        partitioning = SEQ
        # result is parallel if matrix is parallel and transposed (X'*x)
    elseif !isSEQ(arr1,state) && t1
        setSEQ(arr2,state)
        #seq = true
        # otherwise, no known pattern found, every array is sequential
    else
        @dprintln(3,"DistPass arr info gemv all sequential: ", arr1," ", arr2)
        setSEQ(arr1,state)
        setSEQ(arr2,state)
        partitioning = SEQ
    end

    if partitioning==SEQ
        @dprintln(3,"DistPass arr info gemv output is sequential: ", lhs," ",rhs.args[2])
    end
    setArrayPartitioning(lhs,partitioning,state)
    setArrayPartitioning(toLHSVar(rhs.args[2]),partitioning,state)
    return node
end

#=
function isEqualDimSize(sizes1::Array{Union{RHSVar,Int,Expr},1} , sizes2::Array{Union{RHSVar,Int,Expr},1})
    if length(sizes1)!=length(sizes2)
        return false
    end
    for i in 1:length(sizes1)
        if !eqSize(sizes1[i],sizes2[i])
            return false
        end
    end
    return true
end

function eqSize(a::Expr, b::Expr)
    if a.head!=b.head || length(a.args)!=length(b.args)
        return false
    end
    for i in 1:length(a.args)
        if !eqSize(a.args[i],b.args[i])
            return false
        end
    end
    return true
end

function eqSize(a::RHSVar, b::RHSVar)
    return toLHSVar(a) == toLHSVar(b)
end

function eqSize(a::Any, b::Any)
    return a==b
end
=#


text:
# ==============================================
# First-order approximation model with zonotopes
# ==============================================

"""
    FirstOrderZonotope{EM} <: AbstractApproximationModel

First order approximation model that works with zonotopes.

### Fields

- `exp` -- (optional, default: `BaseExp`) exponentiation method

### Algorithm

The transformations are:

- ``Φ ← \\exp(Aδ)``,
- ``Ω_0 ← bloat(zono(CH(\\mathcal{X}_0, Φ\\mathcal{X}_0)), α + β)``

where ``bloat(\\mathcal{X}, ε)`` bloats the set ``\\mathcal{X}`` with the value
``ε``, ``zono(·)`` overapproximates its argument with a zonotope, and ``α`` and
``β`` are factors computed for the homogeneous system and the inputs,
respectively.

### Reference

A. Girard: Reachability of uncertain linear systems using zonotopes. HSCC 2005.
"""
struct FirstOrderZonotope{EM} <: AbstractApproximationModel
    exp::EM
end

# convenience constructor
function FirstOrderZonotope(; exp=BaseExp)
    return FirstOrderZonotope(exp)
end

function Base.show(io::IO, alg::FirstOrderZonotope)
    print(io, "`FirstOrderZonotope` approximation model with: \n")
    print(io, "    - exponentiation method: $(alg.exp) \n")
end

Base.show(io::IO, m::MIME"text/plain", alg::FirstOrderZonotope) = print(io, alg)

# -----------------------------------------------
# FirstOrderZonotope approximation: Homogeneous case
# -----------------------------------------------

function discretize(ivp::IVP{<:CLCS, <:AbstractZonotope}, δ, alg::FirstOrderZonotope)
    A = state_matrix(ivp)
    X0 = initial_state(ivp)
    n = dim(X0)
    norm_A = opnorm(A, Inf)

    # matrix exponential
    Φ = _exp(A, δ, alg.exp)

    # shared code
    Z, α = _discretize_zonotope(Φ, X0, alg, δ, norm_A)

    # compute bloating of Z
    Ω0 = minkowski_sum(Z, BallInf(zeros(n), α))

    # create result
    X = stateset(ivp)
    Sdis = ConstrainedLinearDiscreteSystem(Φ, X)
    return InitialValueProblem(Sdis, Ω0)
end

# -------------------------------------------------
# FirstOrderZonotope approximation: Inhomogeneous case
# -------------------------------------------------

function discretize(ivp::IVP{<:CLCCS, <:AbstractZonotope}, δ, alg::FirstOrderZonotope)
    A = state_matrix(ivp)
    X0 = initial_state(ivp)
    n = dim(X0)
    norm_A = opnorm(A, Inf)

    # matrix exponential
    Φ = _exp(A, δ, alg.exp)

    # shared code
    Z, α = _discretize_zonotope(Φ, X0, alg, δ, norm_A)

    # compute bloating factor β
    U = next_set(inputset(ivp), 1)
    μ = norm(U, Inf)
    β = ((exp(norm_A * δ) - 1) * μ) / norm_A

    # compute bloating of Z
    Ω0 = minkowski_sum(Z, BallInf(zeros(n), α + β))

    # discretize inputs
    V = BallInf(zeros(n), β)

    # create result
    B = IdentityMultiple(one(eltype(A)), size(A, 1))
    X = stateset(ivp)
    Sdis = ConstrainedLinearControlDiscreteSystem(Φ, B, X, V)
    return InitialValueProblem(Sdis, Ω0)
end

# -----------
# Common code
# -----------

function _discretize_zonotope(Φ, X0, alg::FirstOrderZonotope, δ, norm_A)
    c = center(X0)
    G = genmat(X0)
    n, p = size(G)

    # compute X(δ)
    Xδ = linear_map(Φ, X0)

    # compute zonotope covering convex hull CH(X(0) ∪ X(δ))
    Z = overapproximate(CH(X0, Xδ), Zonotope)

    # compute bloating factor α
    α = (exp(norm_A * δ) - 1 - norm_A * δ) * norm(X0, Inf)

    return Z, α
end

ext: jl


text:
"""
    row_as_names!(table::CleanTable, i::Int; remove::Bool=true)

Renames the `table` using row `i` as new names and removes in-place all the rows above row `i` if
`remove=true`.

Default behavior is to remove rows above row `i`.
"""
function row_as_names!(table::CleanTable, i::Int; remove::Bool=true)
    columns = cols(table)
    length(columns[1]) < i && error("`i` can't be greater than the amount of rows")
    new_names = [Symbol(col[i]) for col in columns]
    to_delete = 1:i

    rename!(table, new_names)

    if remove
        for col in columns
            for i in to_delete
                deleteat!(col, 1)
            end
        end
    end

    return table
end

"""
    row_as_names(table, i::Int; remove::Bool=true)

Creates a `CleanTable` with copied columns and renames the table using row `i` as new names and
removes in-place all the rows above row `i` if `remove=true`.

Default behavior is to remove rows above row `i`.
"""
function row_as_names(table, i::Int; remove::Bool=true)
    return row_as_names!(CleanTable(table), i; remove=remove)
end

max_stars_repo_name: TheRoniOne/Cleaner

max_issues_repo_name: TheRoniOne/Cleaner

max_forks_repo_name: TheRoniOne/Cleaner


text:
module TestOptim
using GaussianProcesses, StatsFuns
using Test, Random

Random.seed!(1)

@testset "Optim" begin
    d, n = 2, 20
    X = 2π * rand(d, n)

    mean = MeanLin(zeros(d))
    kern = SE(1.0, 1.0)

    @testset "Without likelihood" begin
        y = X'rand(d) + 0.1*randn(n)
        noise = 0.0

        # Just checks that it doesn't crash
        # and that the final mll is better that the initial value
        @testset "Basic" begin
	          gp = GPE(X, y, mean, kern, -3.0)
	          init_target = gp.target
	          optimize!(gp)
	          @test gp.target > init_target
        end

        @testset "Fixed kernel" begin
            init_param = GaussianProcesses.get_params(kern)[1]
            fixed = fix(kern, GaussianProcesses.get_param_names(kern)[1])
            gp = GP(X, y, MeanZero(), fixed, -1.0)
            init_target = gp.target
            optimize!(gp)
            @test gp.target > init_target
            @test GaussianProcesses.get_params(kern)[1] == init_param
        end

        @testset "Keyword arguments" begin
            gp = GPE(X, y, mean, kern, noise)
            init_params = GaussianProcesses.get_params(gp; domean=true, kern=true,
                                                       noise=true)

            # Check mean fixed
            mean_params = GaussianProcesses.get_params(gp; domean=true, kern=false,
                                                       noise=false)
            optimize!(gp; domean=false, kern=true, noise=true)
            @test mean_params == GaussianProcesses.get_params(gp; domean=true,
                                                              kern=false, noise=false)

            set_params!(gp, init_params; domean=true, kern=true, noise=true)

            # Check kern fixed
            kern_params = GaussianProcesses.get_params(gp; domean=false, kern=true,
                                                       noise=false)
            optimize!(gp; domean=true, kern=false, noise=true)
            @test kern_params == GaussianProcesses.get_params(gp; domean=false, kern=true,
                                                              noise=false)

            set_params!(gp, init_params; domean=true, kern=true, noise=true)

            # Check noise fixed
            noise_params = GaussianProcesses.get_params(gp; domean=false, kern=false,
                                                        noise=true)
            optimize!(gp; domean=true, kern=true, noise=false)
            @test noise_params == GaussianProcesses.get_params(gp; domean=false, kern=false,
                                                               noise=true)
        end
    end

    @testset "With likelihood" begin
        f = X'rand(d) + 0.1*randn(n)
        y = collect(rand(n) .< normcdf.(f)) # Binary data
        lik = BernLik()  # Bernoulli likelihood for binary data {0,1}

        # Just checks that it doesn't crash
        # and that the final mll is better that the initial value
        @testset "Basic" begin
            gp = GPMC(X, y, mean, kern, lik)
            init_target = gp.target
            optimize!(gp)
            @test gp.target > init_target
        end

        @testset "Keyword arguments" begin
            gp = GPMC(X, y, mean, kern, lik)
            init_params = GaussianProcesses.get_params(gp; domean=true, kern=true, lik=true)

            # Check mean fixed
            mean_params = GaussianProcesses.get_params(gp.mean)
            optimize!(gp; domean=false, kern=true, lik=true)
            @test mean_params == GaussianProcesses.get_params(mean)

            set_params!(gp, init_params; domean=true, kern=true, lik=true)

            # Check kern fixed
            kern_params = GaussianProcesses.get_params(gp.kernel)
            optimize!(gp; domean=true, kern=false, lik=true)
            @test kern_params == GaussianProcesses.get_params(kern)

            set_params!(gp, init_params; domean=true, kern=true, lik=true)

            # Check lik fixed
            lik_params = GaussianProcesses.get_params(gp.lik)
            optimize!(gp; domean=true, kern=true, lik=false)
            @test lik_params == GaussianProcesses.get_params(lik)
        end
    end
end
end


text:
# Default parameter values
const DEFAULT_K = 256
const DEFAULT_M = 1
const DEFAULT_METHOD = :sample
const DEFAULT_DISTANCE = Distances.SqEuclidean()
const MAX_BITS = 128
const BITS_TO_TYPE = Dict(round(Int, log2(typemax(typ))) => typ
                          for typ in subtypes(Unsigned))
const TYPE_TO_BITS = Dict(t=>b for (b,t) in BITS_TO_TYPE)
const DEFAULT_PQ_MAXITER = 25
const DEFAULT_OPQ_MAXITER = 25
const DEFAULT_RVQ_MAXITER = 25


text:
"""
    Diagnostic(calc, prob; freq=1, nsteps=100, N=floor(Int, (nsteps+1)/freq))

Construct a diagnostic which stores the result of `calc(prob)` with frequency `freq`
"""
mutable struct Diagnostic{T,N} <: AbstractDiagnostic
   calc :: Function
   prob :: Problem
   data :: Vector{T}
      t :: Vector{Float64}
  steps :: Vector{Int}
   freq :: Int
      i :: Int
end

function Diagnostic(calc, prob; freq=1, nsteps=100, ndata=ceil(Int, (nsteps+1)/freq), N=ndata)
  firstvalue = calc(prob)
  T = typeof(firstvalue)

  data = Vector{T}(undef, N)
  t = Vector{Float64}(undef, N)
  steps = Vector{Int}(undef, N)

  data[1] = firstvalue
  t[1] = prob.clock.t
  steps[1] = prob.clock.step

  return Diagnostic{T,N}(calc, prob, data, t, steps, freq, 1)
end

"""
    extend!(diag::AbstractDiagnostic, n)

Extend the `data`, `time`, and `steps` vectors of `diag` by `n`.
"""
function extend!(d, n)
  resize!(d.data, length(d.steps)+n)
  resize!(d.t, length(d.steps)+n)
  resize!(d.steps, length(d.steps)+n)
  return nothing
end

extend!(diag::Diagnostic{T,N}) where {T,N} = extend!(diag, N)

"""
    update!(diag)

Update `diag`.
"""
function update!(d, i)
  i <= length(d.steps) || extend!(d)
   d.data[i] = d() #d.calc(prob)
      d.t[i] = d.prob.clock.t
  d.steps[i] = d.prob.clock.step
         d.i = i
  return nothing
end

"""
    increment!(diag)

Increment the Diagnostic diag, or an array of Diagnostics diags.
"""
function increment!(d)
  d.prob.clock.step % d.freq == 0 && update!(d, d.i+1)
  return nothing
end

function increment!(diags::AbstractVector)
  for d in diags
    increment!(d)
  end
  return nothing
end

"""
e.g. `plot(energydiag[:t], energydiag[:data])`.
"""
getindex(d::Diagnostic, s::Union{AbstractString,Symbol}) = getfield(d, Symbol(s))[1:d.i]
getindex(d::Diagnostic, idx...) = getindex(d.data, idx...)

lastindex(d::Diagnostic) = d.i

# d() = current value
# e = Diagnostic(energy, ...)
# e() returns energy at the current time.
# sensible?
(d::Diagnostic)() = d.calc(d.prob)


text:
#nb # %% A slide [markdown] {"slideshow": {"slide_type": "subslide"}}
# ## Square Loop Absorber
# This example is from Figure 7 of Costa and Monorchio: "A frequency selective
# radome with wideband absorbing properties", *IEEE Trans. AP-S*,
# Vol. 60, no. 6, June 2012, pp. 2740--2747.  It shows how one can use the `polyring`
# function to model square loop elements.  Three different designs are examined
# that employ different loop thicknesses and different values of sheet resistance.
# We compare the reflection coefficient magnitudes computed by PSSFSS with those digitized
# from the cited figure when the sheet is suspended
# 5 mm above a ground plane, hence we will also make use of the `pecsheet` function.

using Plots, PSSFSS, DelimitedFiles
D = 11 # Period of square lattice (mm)
r_outer = √2/2 * D/8 * [5,6,7] # radii of square outer vertices
thickness = D/16 * [1,2,3]
r_inner = r_outer - √2 * thickness  # radii of square inner vertices
Rs = [15,40,70] # Sheet resistance (Ω/□)
labels = ["Thin", "Medium", "Thick"]
colors = [:green, :blue, :red]
p = plot(title="Costa Absorber", xlim=(0,25),ylim=(-35,0),xtick=0:5:25,ytick=-35:5:0,
         xlabel="Frequency (GHz)", ylabel="Reflection Magnitude (dB)", legend=:bottomleft)
ps = []
for (i,(ri, ro, label, color, R)) in enumerate(zip(r_inner, r_outer, labels, colors, Rs))
    sheet = polyring(sides=4, s1=[D, 0], s2=[0, D], ntri=700, orient=45,
                     a=[ri], b=[ro], Rsheet=R, units=mm)
    push!(ps, plot(sheet, unitcell=true, title=label, lc=color))
    strata = [Layer()
              sheet
              Layer(width=5mm)
              pecsheet() # Perfectly conducting ground plane
              Layer()]
    results = analyze(strata, 1:0.2:25, (ϕ=0, θ=0), showprogress=false,
                      resultfile=devnull, logfile=devnull)
    data = extract_result(results, @outputs FGHz s11dB(h,h))
    plot!(p, data[:,1], data[:,2], label="PSSFSS "*label, lc=color)
    dat = readdlm("../src/assets/costa_2014_" * lowercase(label) * "_reflection.csv", ',')
    plot!(p, dat[:,1], dat[:,2], label="Costa "*label, ls=:dash, lc=color)
end
plot(ps..., layout=(1,3))
#md savefig("sqloop1.png"); nothing  # hide
#-
#md # ![](sqloop1.png)
#-
# This run takes about 85 seconds on my machine.
p
#md savefig(p,"sqloop2.png"); nothing  # hide
#-
#md # ![](sqloop2.png)

# It is useful to take a look at the log file created by PSSFSS for the last run above:
# ```
# Starting PSSFSS analysis on 2021-05-05 at 15:22:09.955
#
#
# Dielectric layer information...
#
# Layer  Width  units  epsr   tandel   mur  mtandel modes  beta1x  beta1y  beta2x  beta2y
# ----- ------------- ------- ------ ------- ------ ----- ------- ------- ------- -------
#     1    0.0000  mm    1.00 0.0000    1.00 0.0000     2   571.2     0.0     0.0   571.2
# ==================  Sheet   1  ========================   571.2     0.0     0.0   571.2
#     2    5.0000  mm    1.00 0.0000    1.00 0.0000    42   571.2     0.0     0.0   571.2
# ==================  Sheet   2  ========================     0.0     0.0     0.0     0.0
#     3    0.0000  mm    1.00 0.0000    1.00 0.0000     2   571.2     0.0     0.0   571.2
#
#
#
# PSS/FSS sheet information...
#
# Sheet  Loc         Style      Rot  J/M Faces Edges Nodes Unknowns  NUFP
# -----  ---  ---------------- ----- --- ----- ----- ----- -------- ------
#   1     1          polyring   0.0  J    753  1201   448    1058  567009
#   2     2              NULL   0.0  E      0     0     0       0       0
#
# ...
# ```
#-
# Note from the dielectric layer report that there are 42 modes defined in the region between the
# ground plane and the FSS sheet.  This is the number of modes selected by the code to include
# in the generalized scattering matrix formulation to properly account for electromagnetic coupling
# between the two surfaces. If the 5 mm spacing were increased to, say, 7 mm then fewer modes
# would be needed.  Also note in the FSS sheet information that `NUFP` (the number of unique face pairs)
# is exactly equal to the number of faces squared (``567009 = 753^2``), a consequence of the unstructured
# triangulation used for a `polyring`.

#-
# ### Conclusion
# PSSFSS results agree very well with those of the paper, except for the medium
# width loop, where the agreement is not quite as good.  The reason for this is
# not known.


text:
using Test
using LinearAlgebra
using PolynomialBasis
# using Revise
using CutCellDG
include("../useful_routines.jl")

function polar_to_cartesian(radialposition, angularposition)
    x = radialposition .* cos.(angularposition)
    y = radialposition .* sin.(angularposition)
    return vcat(x', y')
end

L, W = 1.0, 1.0
nelmts = [4, 4]
center = [0.5, 0.5]
radius = 0.4
polyorder = 2

levelsetbasis = LagrangeTensorProductBasis(2, polyorder)
dim,nf = size(interpolation_points(levelsetbasis))
mesh = CutCellDG.CGMesh([0.0, 0.0], [L, W], nelmts, nf)
levelset = CutCellDG.LevelSet(
    x -> circle_distance_function(x, center, radius),
    mesh,
    levelsetbasis,
)
cutmesh = CutCellDG.CutMesh(mesh, levelset)

refseedpoints, seedcellids = CutCellDG.seed_zero_levelset(2, levelset, cutmesh)
seedpoints = CutCellDG.map_to_spatial(refseedpoints, seedcellids, mesh)

numquerypts = 100
queryangles = range(0, stop = 2pi, length = numquerypts)
queryradius = 0.1 .+ abs.(sin.(queryangles))
querypoints = center .+ polar_to_cartesian(queryradius, queryangles)

dx = CutCellDG.element_size(mesh)
tol = 1e3eps()
boundingradius = 1.5 * maximum(dx)
closestpoints, closestcellids, flags =
    CutCellDG.closest_points_on_zero_levelset(
        querypoints,
        seedpoints,
        seedcellids,
        levelset,
        tol,
        boundingradius,
    )

testcp = closest_point_on_arc(querypoints, center, radius)
err = norm(testcp - closestpoints, Inf)
# err = vec(mapslices(x->norm(x,Inf),closestpoints - testcp,dims=1))
# maxidx = argmax(err)
# maxerr = err[maxidx]

@test all(flags)
@test err < 0.02
################################################################################

################################################################################
levelsetbasis = HermiteTensorProductBasis(2)

dim,nf = size(interpolation_points(levelsetbasis))
mesh = CutCellDG.CGMesh([0.0, 0.0], [L, W], nelmts, nf)
levelset = CutCellDG.LevelSet(
    x -> circle_distance_function(x, center, radius)[1],
    mesh,
    levelsetbasis,
)
cutmesh = CutCellDG.CutMesh(mesh, levelset)

refseedpoints, seedcellids = CutCellDG.seed_zero_levelset(2, levelset, cutmesh)
seedpoints = CutCellDG.map_to_spatial(refseedpoints, seedcellids, mesh)

numquerypts = 100
queryangles = range(0, stop = 2pi, length = numquerypts)
queryradius = 0.1 .+ abs.(sin.(queryangles))
querypoints = center .+ polar_to_cartesian(queryradius, queryangles)

dx = CutCellDG.element_size(mesh)
tol = 1e3eps()
boundingradius = 1.5 * maximum(dx)
closestpoints, closestcellids, flags =
    CutCellDG.closest_points_on_zero_levelset(
        querypoints,
        seedpoints,
        seedcellids,
        levelset,
        tol,
        boundingradius,
    )

testcp = closest_point_on_arc(querypoints, center, radius)
err = norm(testcp - closestpoints, Inf)

@test all(flags)
@test err < 0.002


text:
using ChainRulesCore

include("SimpleMPSTensor.jl")
include("GeneralMPSTensor.jl")

MPSTensor = Union{GeneralMPSTensor,SimpleMPSTensor}

function ITensors.inds(tensor::MPSTensor)
  return siteinds(tensor.mps) == [nothing] ? () : tuple(noncommoninds(tensor.mps...)...)
end

ITensors.sum(tensors::Vector{<:MPSTensor}; kwargs...) = ITensors.sum(tensors...; kwargs...)

function ITensors.sum(t1::MPSTensor, t2::MPSTensor...; kwargs...)
  return +(t1, sum(t2...; kwargs...); kwargs...)
end

ITensors.sum(t::MPSTensor; kwargs...) = t

"""
contract
"""
# contract into one ITensor
ITensors.ITensor(t::MPSTensor; kwargs...) = contract(collect(t.mps)...; kwargs...)

ITensors.contract(t1::MPSTensor; kwargs...) = t1

function ITensors.contract(t1::MPSTensor, t2::MPSTensor...; kwargs...)
  return contract(t1, contract(t2...; kwargs...); kwargs...)
end

ITensors.contract(t_list::Vector{<:MPSTensor}; kwargs...) = contract(t_list...; kwargs...)


text:
addprocs(8)

include("../../daggen.jl")
using DagScheduler
using BenchmarkTools

isdir(".mempool") && rm(".mempool"; recursive=true)
@everywhere begin
    DagScheduler.META_IMPL[:node] = ENV["NODE_META_IMPL"]
    DagScheduler.META_IMPL[:cluster] = ENV["CLUSTER_META_IMPL"]
    DagScheduler.META_IMPL[:map_num_entries] = 1024*5
    DagScheduler.META_IMPL[:map_entry_sz] = 256
end

node1 = NodeEnv(2, getipaddr(), [3,4,5])
node2 = NodeEnv(6, getipaddr(), [7,8,9])
runenv = DagScheduler.Plugin.setrunenv(RunEnv(; nodes=[node1,node2]))

const L = 6^4
const dag2 = gen_straight_dag(ones(Int, L));

result = collect(rundag(runenv, dag2));
#@time result = rundag(runenv, dag2);
@btime collect(rundag(runenv, dag2))

DagScheduler.cleanup(runenv)


text:
# using Dists: ValueSupport, VariateForm

# struct MixedSupport <: ValueSupport end
# struct MixedVariate <: VariateForm end

"""
    AbstractModel{A,B,M,Args,Obs}

Gives an abstract type for all Soss models

Type variables ending in T are type-level representations used to reconstruct

N gives the Names of arguments (each a Symbol)
B gives the Body, as an Expr
M gives the Module where the model is defined
"""
abstract type AbstractModel{A,B,M,Args,Obs} <: AbstractMeasure end

argstype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = AT
argstype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = AT

bodytype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = BT
bodytype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = BT

getmodule(::Type{AbstractModel{A,B,M,Args,Obs}}) where  {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = from_type(M)
getmodule(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = from_type(M)

getmoduletypencoding(::Type{AbstractModel{A,B,M,Args,Obs}}) where  {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = M
getmoduletypencoding(::AbstractModel{A,B,M,Args,Obs}) where  {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = M

argvalstype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = Args
argvalstype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = Args

obstype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = Obs
obstype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = Obs


text:
#############################################
## Set parameters
#############################################

using NIfTI, PyPlot, HDF5, MRIReco, LinearAlgebra, Dates

# parameters
@info "Setting Parameters"

verbosity_level = 1 # 0 = no Visualization, 1 = recon only, 2 = maps, 3 = iterations
do_recalc_sensitivity = false
do_b0_correction = true
b0_map_type = "smoothed" # "raw" or "smoothed" or "smoothed_resized"

do_plot_recon 		= verbosity_level >= 1
do_plot_maps 		= verbosity_level >= 2
do_plot_traj 		= do_plot_maps
do_plot_iterations 	= verbosity_level >= 3

w0_offset = 0 # rad/s, only needed for simulating f0 effects
dt = 1.8e-6 # acquisition dwell time [s]
TE = 0.02 # s
Nx = 240 # 300 # 240
Ny = 292 # 365 # 292
fov = [190, 230, 0.9+0.1]/1000 # [m]

idx_slice = 18; # slice(s) to reconstruct
do_select_slice =  !isempty(idx_slice)
# Paths
path_data = "D:\\SPIFI\\ExportETHResearchCollection"

# Raw Coil/Trajectory data: ISMRMRD file (.nii); units: rad/m
filename_data = joinpath(path_data,"SPIFI_0007_RawData_RotatedToSliceGeometry3D_spiralOut_singleVolume100.h5")

# Maps are NIFTI files (.nii); units: Hz
filename_sense_magnitude = joinpath(path_data, "SPIFI_0007_MapsForReconstruction", "coilSensitivityMaps_magnitude.nii")
filename_sense_phase = joinpath(path_data, "SPIFI_0007_MapsForReconstruction", "coilSensitivityMaps_phase.nii")

if b0_map_type == "smoothed_resized"
	# should give proper results
	filename_b0 = joinpath(path_data, "SPIFI_0007_MapsForReconstruction", "b0MapSmoothedResizedSpiralGeometry1_Hz.nii")
elseif b0_map_type == "smoothed"
	filename_b0 = joinpath(path_data, "SPIFI_0007_MapsForReconstruction", "b0MapSmoothed1_Hz.nii")
elseif b0_map_type == "raw"
	# use raw map instead:
	filename_b0 = joinpath(path_data, "SPIFI_0007_MapsForReconstruction", "b0MapRaw1_Hz.nii")
else
	error("Requested B0 map type not implemented or saved")
end

# rawdata: ch, il, read was order
# traj: il, read, kdim(x,y,z)
# Julia order: read, il, ch and kdim, read, il
#data = permutedims(h5read(filename_data, "rawdata"),[3,2,1,4])

#############################################
## Load ISMRMRD (coil/traj) data and maps (SENSE/B0)
#############################################

@info "Load Data"
dataFile = ISMRMRDFile(filename_data)
rawData = RawAcquisitionData(dataFile)

# select data from selected slices only
if do_select_slice
	indices = 1:length(rawData.profiles)
	#ic = [x for x ∈ indices if x ∉ idx_slice]
	ic = setdiff(indices, idx_slice)
	deleteat!(rawData.profiles,ic)

	n_slices =  size(idx_slice,1);
	# make sinlge slice parameter set
	rawData.profiles[1].head.idx.slice = 0;
	rawData.profiles[1].head.idx.repetition=0;
	rawData.params["enc_lim_slice"] = Limit(0, 0, 0)
	rawData.params["enc_lim_repetition"] = Limit(0, 0 ,0)
	rawData.params["reconSize"][3] = n_slices
	rawData.params["encodedSize"][3] = n_slices
	rawData.params["reconFOV"][3] = rawData.params["reconFOV"][3]/size(indices,1)*n_slices
	rawData.params["encodedFOV"][3] = rawData.params["encodedFOV"][3]/size(indices,1)*n_slices
	# print function for all parameters
	 for (key, value) in rawData.params; print(key); print(": ");print(value);print("\n"); end
end

# alternative read-in of data, not relying on mutable struct AcquisitionHeader in MRIReco.jl
# mytr = Trajectory(reshape(traj[1:2,:,:],2,:), n_interleaves, n_samples,
# TE=0.02, AQ=tAQ, circular=true)
# dat = Array{Array{Complex{Float64},2},3}(undef,1,1,1)
# dat[1,1,1] = reshape(rawdata,:,n_channels)
# acqData = AcquisitionData(mytr, dat, encodingSize=[Nx,Ny,1], fov=fov)

@info "Converting RawAcquisitionData to AcquisitionData"
# acqData = AcquisitionData(rawData,estimateProfileCenter=false)
acqData = AcquisitionData(rawData)

# adjust TE and TAQ after read-in, is not taken from file
n_samples = rawData.params["encodedSize"][1]
tAQ = (n_samples-1) * dt
acqData.traj[1].AQ=tAQ # important for B0 correction
acqData.traj[1].TE= TE
acqData.traj[1].times = TE .+ collect(0:dt:tAQ)

#############################################
## Load and convert traj from rad/m to -0.5 -> 0.5
#############################################

if do_plot_traj
	# Plot Raw Trajectory
	mytr = acqData.traj[1]
	figure(1);
	subplot(1,2,1)
	plot(mytr.nodes[1,:], mytr.nodes[2,:]);
	title("Raw Trajectory")
	axis(:square)
	gcf()
end

@info "Normalize trajectory rad/m -> [-1/2 1/2] FOV"
traj_node_normalized = acqData.traj[1].nodes
traj_node_normalized /= (2*π)
traj_node_normalized[2,:,:] /= Nx/fov[1]
traj_node_normalized[1,:,:] /= Ny/fov[2]

# Remove 3rd dimension for recon, otherwise MRIReco.jl has 3D assumption
traj_node_normalized = traj_node_normalized[1:2,:]

acqData.traj[1].nodes = traj_node_normalized

#############################################
## Plot Normalized Trajectory
##############################################

if do_plot_traj
	@info "Plot Normalized Trajectory"

	mytr = acqData.traj[1]
	figure(1);
	subplot(1,2,2)
	plot(mytr.nodes[1,:], mytr.nodes[2,:]);
	title("Normalized Trajectory")
	axis(:square)
	gcf()
end

################################
## Generate coil sensitivity maps
################################

if do_recalc_sensitivity == true
    @info "Calc Coil Sensitivity maps with Espirit"
    acqDataCart = regrid2d(acqData, (Nx,Ny); cgnr_iter=3)
    sense = espirit(acqDataCart,(6,6),30,eigThresh_1=0.02, eigThresh_2=0.98)
    sensitivity = sense;
#    sensitivity[1,1,1] = reshape(sense,:,n_channels)
else
    @info "Load Coil Sensitivity maps"
    # from previous Matlab calculation

	#loadImage(filename_sense_magnitude)
	sense = niread(filename_sense_magnitude).*exp.(1im.*niread(filename_sense_phase))

	n_channels = size(sense,4);

	sense = sense[:,:,idx_slice,:]
	# keep channels in 4th dim, even if singleton slice dim
	if ndims(sense)==3
		size_sense = size(sense)
		sense = reshape(sense, size_sense[1], size_sense[2], 1, size_sense[3])
	end

	n_slices = size(sense,3);

    sensitivity = Array{Array{Complex{Float64},2},4}(undef,1,1,1,1)
    sensitivity = sense;
end

sensitivity = mapslices(x ->imresize(x, (Nx, Ny)), sensitivity, dims=[1,2])
# conversion to ComplexF64 needed for reconstruction_2d solver
sensitivity = convert(Array{ComplexF64, 4}, reshape(sensitivity, Nx, Ny, 1, n_channels))

if do_plot_maps
	figure(2); clf(); for ch in 1:n_channels; subplot(8,4,ch); imshow(rotl90(abs.(sense[:,:,1,ch]))); end;
	subplots_adjust(wspace=0.05,hspace=0.05,left=0.05,bottom=0.0,right=1.0,top=0.95)
	gcf()
	figure(3); clf(); for ch in 1:n_channels; subplot(8,4,ch); imshow(rotl90(angle.(sense[:,:,1,ch]))); end;
	subplots_adjust(wspace=0.05,hspace=0.05,left=0.05,bottom=0.0,right=1.0,top=0.95)
	gcf()
end

##########################
## Load B0 map, adapt geometry
##########################

@info "Load B0 map"
b0 = niread(filename_b0).*2π .- w0_offset

#select recon slice
b0 = b0[:,:,idx_slice]

resizedB0 = mapslices(x->imresize(x,(Nx, Ny)), b0, dims=[1,2])

cmap = 1im.*resizedB0;

if do_plot_maps
	# Plot B0 map before and after resizing
	figure(4); cla();
	subplot(1,2,1)
	imshow(rotl90(b0), cmap="gray");
	colorbar()
	title("Loaded B0 map")
	subplot(1,2,2)
	imshow(rotl90(resizedB0), cmap="gray");
	title("B0 map, resized to Recon Geometry")
	subplots_adjust(wspace=0.05,hspace=0.05,left=0.05,bottom=0.0,right=1.0,top=0.95)
	colorbar()
	gcf()
end

##########################
## Perform reference reconstruction
##########################
@info "Reference cg-SENSE recon"
params = Dict{Symbol, Any}()
params[:reco] = "multiCoil"
params[:reconSize] = (Nx,Ny)
params[:regularization] = "L2"
params[:λ] = 1.e-2
params[:iterations] = 10
params[:solver] = "cgnr"
params[:solverInfo] = SolverInfo(ComplexF64,store_solutions=do_plot_iterations)
params[:senseMaps] = sensitivity

# params[:reco] = "direct"
#@time
begin
    if do_b0_correction
        params[:correctionMap] = cmap
        #params[:alpha] = 1.75 # oversampling factor for interpolation
        #params[:m] = 4.0 # truncation size of interpolation kernel
        #params[:K] = 28  # number of translates for LeastSquares approaches (not NFFT-approach) to time-segmentation
    end

    img_ref = reconstruction(acqData, params).data
end

# h5write(filename_recon, "/img_ref", img_ref)

##########################
## Plot resulting recon and iterations
##########################
if do_plot_recon
	figure(5); cla();
	subplot(1,2,1); imshow(rotl90(abs.(img_ref[:,:,1,1,1])), cmap="gray", vmax=0.8, aspect="equal");
	subplot(1,2,1); imshow(rotl90(abs.(img_ref[:,:,1,1,1])), cmap="gray", vmax=0.6, aspect="equal");
	subplot(1,2,2); imshow(rotl90(angle.(img_ref[:,:,1,1,1])), cmap="gray");
	subplots_adjust(wspace=0.05,hspace=0.05,left=0.05,bottom=0.0,right=1.0,top=0.95)
	gcf()

	figure(6); cla();
	subplot(1,2,1); plot(params[:solverInfo].convMeas); title("Linear")
	subplot(1,2,2); semilogy(params[:solverInfo].convMeas);  title("Log")
	gcf().suptitle("Convergence over Iterations")
	xlabel("Iterations")
	gcf()

	# plot with colorbar
	# begin
	# 	fig, (ax1, ax2) = subplots(figsize=(9, 3), ncols=2)
	# 	hp1 = ax1.imshow(rotl90(abs.(img_ref[:,:,1,1,1])), cmap="gray", vmax=0.8, aspect="equal")
	# 	hp2 = ax2.imshow(rotl90(angle.(img_ref[:,:,1,1,1].-π/10.0)), cmap="gray")
	# 	fig.colorbar(hp1, ax=ax1)
	# 	fig.colorbar(hp2, ax=ax2)
	# 	fig
	# end
end

# Plot Iterations
if do_plot_iterations
 	img_iter = Vector{Array{ComplexF64,5}}(undef,11)
 	for iter = 1:11
     	img_iter[iter] = reshape(params[:solverInfo].x_iter[iter],Nx,Ny,1,1,1)
     	figure(1000+iter);cla;imshow(abs.(img_ref[:,:,1,1,1])); gcf()
	end
end

##########################
## Save final recon result with a time stamp
##########################
if do_plot_recon
	figure(5);gcf()
	figSaveName = string("FinalRecon_", Dates.format(Dates.now(), "yyyy-mm-dd_HH_MM_SS"), ".png")
	savefig(figSaveName,dpi=300)
end


text:
function solve()
    n = parse(Int, readline())
    as = [parse(Int, x) for x in split(readline())]

    f1 = maximum(as[1:2^(n-1)])
    f2 = maximum(as[2^(n-1)+1:end])
    r = min(f1, f2)
    findfirst(x -> x==r, as)
end

println(solve())

max_stars_repo_name: aishikawa/atcoder-julia

max_issues_repo_name: aishikawa/atcoder-julia

max_forks_repo_name: aishikawa/atcoder-julia


text:
using Rasters, Test, Dates, DiskArrays
using Rasters.LookupArrays, Rasters.Dimensions
using Rasters: isdisk, ismem, filename

data1 = cumsum(cumsum(ones(10, 11); dims=1); dims=2)
data2 = 2cumsum(cumsum(ones(10, 11, 1); dims=1); dims=2)
dims1 = X(10:10:100), Y(-50:10:50)
dims2 = (dims1..., Ti([DateTime(2019)]))
refdimz = ()
mval = -9999.0
meta = NoMetadata()
nme = :test

ga2 = Raster(data2, dims2)
ga1 = Raster(data1; dims=dims1, refdims=refdimz, name=nme, metadata=meta, missingval=mval)

@test ga1 == data1
@test ga2 == data2

@testset "array properties" begin
    @test name(ga1) == :test
    @test missingval(ga1) == -9999.0
    @test isdisk(ga1) == false
    @test ismem(ga1) == true
    @test filename(ga1) === nothing
    @test DiskArrays.haschunks(ga1) == DiskArrays.Unchunked()
    @test DiskArrays.eachchunk(ga1) == DiskArrays.GridChunks((10, 11), (10, 11))
end

@testset "crs" begin
    crs(ga1) == nothing
    mappedcrs(ga1) == nothing
    gapr = setcrs(ga1, EPSG(4326))
    gampr = setmappedcrs(gapr, EPSG(4326))
    @test crs(gapr) == crs(gapr[X(1)]) == EPSG(4326)
    @test mappedcrs(gampr) == mappedcrs(gampr[X(1)]) == EPSG(4326)
end

@testset "arary dims have been formatted" begin
    @test index(ga2) == (10.0:10:100, -50.0:10:50.0, [DateTime(2019)])
    @test dims(ga1)[1:2] == dims(ga2)[1:2]
end

@testset "constructors" begin
    x = X(Projected(1.0:1.0:2.0; crs=ProjString("+proj=")))
    y = Y(Projected(1.0:1.0:2.0; crs=ProjString("+proj=")))

    A1 = rand(Float32, x, y)
    @test A1 isa Raster{Float32,2}
    @test size(A1) == (2, 2)

    A2 = zeros(Int, x, y)
    @test A2 isa Raster{Int,2}
    @test A2 == [0 0; 0 0]

    A3 = ones(x, y)
    @test A3 isa Raster{Float64,2}
    @test A3 == [1.0 1.0; 1.0 1.0]

    A4 = fill(99.0, x, y)
    @test A4 isa Raster{Float64,2}
    @test A4 == [99.0 99.0; 99.0 99.0]

    A5 = trues(x, y)
    @test A5 isa Raster{Bool,2}
    @test A5 == [true true; true true]

    A6 = falses(x, y)
    @test A6 isa Raster{Bool,2}
    @test A6 == [false false; false false]
end

@testset "show" begin
    sh = sprint(show, MIME("text/plain"), ga1)
    # Test but don't lock this down too much
    @test occursin("Raster", sh)
    @test occursin("Y", sh)
    @test occursin("X", sh)
end

@testset "skipmissing uses missingval" begin
    raster = Raster([NaN 1.0; 2.0 NaN], (X, Y); missingval=NaN)
    @test collect(skipmissing(raster)) == [2.0, 1.0]
    @test collect(keys(skipmissing(raster))) == [CartesianIndex(2, 1), CartesianIndex(1, 2)]
    @test collect(eachindex(skipmissing(raster))) == [2, 3]
    @test_throws MissingException skipmissing(raster)[1]
    @test skipmissing(raster)[2] == 2.0
    # It skips actual missing values as well
    mraster = Raster([NaN 1.0; missing NaN], (X, Y); missingval=NaN)
    @test collect(skipmissing(mraster)) == [1.0]
    @test collect(keys(skipmissing(mraster))) == [CartesianIndex(1, 2)]
    @test collect(eachindex(skipmissing(mraster))) == [3]
    @test skipmissing(mraster)[3] == 1.0
    @test_throws MissingException skipmissing(mraster)[2]

    r = Raster(ones(Int16, 8, 8), (X,Y); missingval = Int16(-9999))
    r[1:4, 1:4] .= -9999
    r = float.(r)
    @test !(missingval(r) in skipmissing(r))
    @test length(collect(skipmissing(r))) == 48
end

@testset "table" begin
    table = DimTable(ga1)
    tra = Raster(table, dims(ga1))
    @test tra == ga1
    @test name(tra) == :test
    @test_throws ArgumentError Raster(table, dims(ga1); name=:x)
end

@testset "from vector" begin
    vra = Raster(vec(ga1), dims(ga1); name=:from_vector)
    @test vra == ga1
    @test name(vra) == :from_vector
end


text:
include("utility.jl")
include("compressors.jl")
include("filters.jl")


text:
module Robotics

using LinearAlgebra

include("utils.jl")
include("geometry.jl")
include("dynamics.jl")
include("linalgutils.jl")

include("codegen.jl")

end # module


text:
abstract type AbstractReservesFormulation <: AbstractServiceFormulation end
struct RangeReserve <: AbstractReservesFormulation end
struct StepwiseCostReserve <: AbstractReservesFormulation end
############################### Reserve Variables` #########################################
"""
This function add the variables for reserves to the model
"""
function AddVariableSpec(
    ::Type{ActiveServiceVariable},
    ::PSIContainer,
    service::T,
) where {T <: PSY.Reserve}
    return AddVariableSpec(;
        variable_name = make_variable_name(PSY.get_name(service), T),
        binary = false,
        lb_value_func = x -> 0,
        devices_filter_func = x -> PSY.get_available(x),
    )
end

function AddVariableSpec(
    ::Type{T},
    ::Type{U},
    ::PSIContainer,
) where {T <: ServiceRequirementVariable, U <: PSY.ReserveDemandCurve}
    return AddVariableSpec(;
        variable_name = make_variable_name(T, U),
        binary = false,
        lb_value_func = x -> 0.0,
    )
end

################################## Reserve Requirement Constraint ##########################
function service_requirement_constraint!(
    psi_container::PSIContainer,
    service::SR,
    ::ServiceModel{SR, T},
) where {SR <: PSY.Reserve, T <: AbstractReservesFormulation}
    parameters = model_has_parameters(psi_container)
    use_forecast_data = model_uses_forecasts(psi_container)
    initial_time = model_initial_time(psi_container)
    @debug initial_time
    time_steps = model_time_steps(psi_container)
    name = PSY.get_name(service)
    constraint = get_constraint(psi_container, make_constraint_name(REQUIREMENT, SR))
    reserve_variable = get_variable(psi_container, name, SR)
    use_slacks = get_services_slack_variables(psi_container.settings)

    ts_vector = get_time_series(psi_container, service, "requirement")

    use_slacks && (slack_vars = reserve_slacks(psi_container, name))

    requirement = PSY.get_requirement(service)
    if parameters
        param = get_parameter_array(
            psi_container,
            UpdateRef{SR}(SERVICE_REQUIREMENT, "requirement"),
        )
        for t in time_steps
            param[name, t] = PJ.add_parameter(psi_container.JuMPmodel, ts_vector[t])
            if use_slacks
                resource_expression = sum(reserve_variable[:, t]) + slack_vars[t]
            else
                resource_expression = sum(reserve_variable[:, t])
            end
            constraint[name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                resource_expression >= param[name, t] * requirement
            )
        end
    else
        for t in time_steps
            constraint[name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                sum(reserve_variable[:, t]) >= ts_vector[t] * requirement
            )
        end
    end
    return
end

function service_requirement_constraint!(
    psi_container::PSIContainer,
    service::SR,
    ::ServiceModel{SR, T},
) where {SR <: PSY.StaticReserve, T <: AbstractReservesFormulation}
    parameters = model_has_parameters(psi_container)
    initial_time = model_initial_time(psi_container)
    @debug initial_time
    time_steps = model_time_steps(psi_container)
    name = PSY.get_name(service)
    constraint = get_constraint(psi_container, make_constraint_name(REQUIREMENT, SR))
    reserve_variable = get_variable(psi_container, name, SR)
    use_slacks = get_services_slack_variables(psi_container.settings)

    use_slacks && (slack_vars = reserve_slacks(psi_container, name))

    requirement = PSY.get_requirement(service)
    for t in time_steps
        resource_expression = JuMP.GenericAffExpr{Float64, JuMP.VariableRef}()
        JuMP.add_to_expression!(resource_expression, sum(reserve_variable[:, t]))
        if use_slacks
            resource_expression += slack_vars[t]
        end
        constraint[name, t] =
            JuMP.@constraint(psi_container.JuMPmodel, resource_expression >= requirement)
    end

    return
end

function cost_function!(
    psi_container::PSIContainer,
    service::SR,
    ::ServiceModel{SR, T},
) where {SR <: PSY.Reserve, T <: AbstractReservesFormulation}
    reserve = get_variable(psi_container, PSY.get_name(service), SR)
    for r in reserve
        JuMP.add_to_expression!(psi_container.cost_function, r, DEFAULT_RESERVE_COST)
    end
    return
end

function service_requirement_constraint!(
    psi_container::PSIContainer,
    service::SR,
    ::ServiceModel{SR, StepwiseCostReserve},
) where {SR <: PSY.Reserve}
    initial_time = model_initial_time(psi_container)
    @debug initial_time
    time_steps = model_time_steps(psi_container)
    name = PSY.get_name(service)
    constraint = get_constraint(psi_container, make_constraint_name(REQUIREMENT, SR))
    reserve_variable = get_variable(psi_container, name, SR)
    requirement_variable = get_variable(psi_container, SERVICE_REQUIREMENT, SR)

    for t in time_steps
        constraint[name, t] = JuMP.@constraint(
            psi_container.JuMPmodel,
            sum(reserve_variable[:, t]) >= requirement_variable[name, t]
        )
    end

    return
end

function AddCostSpec(
    ::Type{T},
    ::Type{StepwiseCostReserve},
    psi_container::PSIContainer,
) where {T <: PSY.Reserve}
    return AddCostSpec(;
        variable_type = ServiceRequirementVariable,
        component_type = T,
        has_status_variable = false,
        has_status_parameter = false,
        variable_cost = PSY.get_variable,
        start_up_cost = nothing,
        shut_down_cost = nothing,
        fixed_cost = nothing,
        sos_status = NO_VARIABLE,
    )
end

function add_to_cost!(
    psi_container::PSIContainer,
    spec::AddCostSpec,
    service::SR,
    component_name::String,
) where {SR <: PSY.Reserve}
    time_steps = model_time_steps(psi_container)
    use_forecast_data = model_uses_forecasts(psi_container)
    if !use_forecast_data
        error("StepwiseCostReserve is only supported with forecast")
    end
    variable_cost_forecast = get_time_series(psi_container, service, "variable_cost")
    variable_cost_forecast = map(PSY.VariableCost, variable_cost_forecast)
    for t in time_steps
        variable_cost!(psi_container, spec, component_name, variable_cost_forecast[t], t)
    end
    return
end

function cost_function!(
    psi_container::PSIContainer,
    service::SR,
    model::ServiceModel{SR, StepwiseCostReserve},
) where {SR <: PSY.ReserveDemandCurve}
    spec = AddCostSpec(SR, model.formulation, psi_container)
    @debug SR, spec
    add_to_cost!(psi_container, spec, service, PSY.get_name(service))
    return
end

function modify_device_model!(
    devices_template::Dict{Symbol, DeviceModel},
    service_model::ServiceModel{<:PSY.Reserve, <:AbstractReservesFormulation},
    contributing_devices::Vector{<:PSY.Device},
)
    device_types = unique(typeof.(contributing_devices))
    for dt in device_types
        for (device_model_name, device_model) in devices_template
            # add message here when it exists
            device_model.device_type != dt && continue
            service_model in device_model.services && continue
            push!(device_model.services, service_model)
        end
    end

    return
end

function include_service!(
    constraint_info::T,
    services,
    ::ServiceModel{SR, <:AbstractReservesFormulation},
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    SR <: PSY.Reserve{PSY.ReserveUp},
}
    for (ix, service) in enumerate(services)
        push!(
            constraint_info.additional_terms_ub,
            make_constraint_name(PSY.get_name(service), SR),
        )
    end
    return
end

function include_service!(
    constraint_info::T,
    services,
    ::ServiceModel{SR, <:AbstractReservesFormulation},
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    SR <: PSY.Reserve{PSY.ReserveDown},
}
    for (ix, service) in enumerate(services)
        push!(
            constraint_info.additional_terms_lb,
            make_constraint_name(PSY.get_name(service), SR),
        )
    end
    return
end

function add_device_services!(
    constraint_info::T,
    device::D,
    model::DeviceModel,
) where {
    T <: Union{AbstractRangeConstraintInfo, AbstractRampConstraintInfo},
    D <: PSY.Device,
}
    for service_model in get_services(model)
        if PSY.has_service(device, service_model.service_type)
            services =
                (s for s in PSY.get_services(device) if isa(s, service_model.service_type))
            @assert !isempty(services)
            include_service!(constraint_info, services, service_model)
        end
    end
    return
end

function add_device_services!(
    constraint_data_in::AbstractRangeConstraintInfo,
    constraint_data_out::AbstractRangeConstraintInfo,
    device::D,
    model::DeviceModel{D, <:AbstractStorageFormulation},
) where {D <: PSY.Storage}
    for service_model in get_services(model)
        if PSY.has_service(device, service_model.service_type)
            services =
                (s for s in PSY.get_services(device) if isa(s, service_model.service_type))
            @assert !isempty(services)
            if service_model.service_type <: PSY.Reserve{PSY.ReserveDown}
                for service in services
                    push!(
                        constraint_data_in.additional_terms_ub,
                        make_constraint_name(
                            PSY.get_name(service),
                            service_model.service_type,
                        ),
                    )
                end
            elseif service_model.service_type <: PSY.Reserve{PSY.ReserveUp}
                for service in services
                    push!(
                        constraint_data_out.additional_terms_ub,
                        make_constraint_name(
                            PSY.get_name(service),
                            service_model.service_type,
                        ),
                    )
                end
            end
        end
    end
    return
end


text:
module test_julia_reduce

using Test

@test mapfoldl((x -> [x, x+1]), vcat, [5, 8]) == [5, 6, 8, 9]

end # module test_julia_reduce


text:
using Revise
	using DiffEqOperators, ForwardDiff, IncompleteLU
	using PseudoArcLengthContinuation, LinearAlgebra, Plots, SparseArrays, Parameters, Setfield
	const PALC = PseudoArcLengthContinuation

norminf = x -> norm(x, Inf)

function Laplacian2D(Nx, Ny, lx, ly, bc = :Dirichlet)
	hx = 2lx/Nx
	hy = 2ly/Ny
	D2x = CenteredDifference(2, 2, hx, Nx)
	D2y = CenteredDifference(2, 2, hy, Ny)
	if bc == :Neumann
		Qx = Neumann0BC(hx)
		Qy = Neumann0BC(hy)
	elseif  bc == :Dirichlet
		Qx = Dirichlet0BC(typeof(hx))
		Qy = Dirichlet0BC(typeof(hy))
	else
		Qx = Dirichlet0BC(typeof(hx))
		Qy = Dirichlet0BC(typeof(hy))
	end
	D2xsp = sparse(D2x * Qx)[1]
	D2ysp = sparse(D2y * Qy)[1]
	if bc == :Periodic
		D2xsp[1,end] = D2xsp[1,2]
		D2xsp[end,1] = D2xsp[1,2]

		D2ysp[1,end] = D2ysp[1,2]
		D2ysp[end,1] = D2ysp[1,2]
	end
	A = kron(sparse(I, Ny, Ny), D2xsp) + kron(D2ysp, sparse(I, Nx, Nx))
	return A, D2x
end

function NL(u, p)
	@unpack r, μ, ν, c3, c5 = p
	n = div(length(u), 2)
	u1 = @view u[1:n]
	u2 = @view u[n+1:2n]

	ua = u1.^2 .+ u2.^2

	f = similar(u)
	f1 = @view f[1:n]
	f2 = @view f[n+1:2n]

	@. f1 .= r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1
	@. f2 .= r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2

	return f
end

function Fcgl!(f, u, p)
	mul!(f, p.Δ, u)
	f .= f .+ NL(u, p)
end

function Fcgl(u, p)
	f = similar(u)
	Fcgl!(f, u, p)
end

# computation of the first derivative
d1Fcgl(x, p, dx) = ForwardDiff.derivative(t -> Fcgl(x .+ t .* dx, p), 0.)

d1NL(x, p, dx) = ForwardDiff.derivative(t -> NL(x .+ t .* dx, p), 0.)

function dFcgl(x, p, dx)
	f = similar(dx)
	mul!(f, p.Δ, dx)
	nl = d1NL(x, p, dx)
	f .= f .+ nl
end

# computation of the second derivative
d2Fcgl(x, p, dx1, dx2) = ForwardDiff.derivative(t2 -> ForwardDiff.derivative( t1 -> Fcgl(x .+ t1 .* dx1 .+ t2 .* dx2, p), 0.), 0.)

# remark: I checked this against finite differences
function Jcgl(u, p)
	@unpack r, μ, ν, c3, c5, Δ = p

	n = div(length(u), 2)
	u1 = @view u[1:n]
	u2 = @view u[n+1:2n]

	ua = u1.^2 .+ u2.^2

	f1u = zero(u1)
	f2u = zero(u1)
	f1v = zero(u1)
	f2v = zero(u1)

	@. f1u =  r - 2 * u1 * (c3 * u1 - μ * u2) - c3 * ua - 4 * c5 * ua * u1^2 - c5 * ua^2
	@. f1v = -ν - 2 * u2 * (c3 * u1 - μ * u2)  + μ * ua - 4 * c5 * ua * u1 * u2
	@. f2u =  ν - 2 * u1 * (c3 * u2 + μ * u1)  - μ * ua - 4 * c5 * ua * u1 * u2
	@. f2v =  r - 2 * u2 * (c3 * u2 + μ * u1) - c3 * ua - 4 * c5 * ua * u2 ^2 - c5 * ua^2

	jacdiag = vcat(f1u, f2v)

	Δ + spdiagm(0 => jacdiag, n => f1v, -n => f2u)
end

####################################################################################################
Nx = 41*1
	Ny = 21*1
	n = Nx*Ny
	lx = pi
	ly = pi/2

	Δ = Laplacian2D(Nx, Ny, lx, ly)[1]
	par_cgl = (r = 0.5, μ = 0.1, ν = 1.0, c3 = -1.0, c5 = 1.0, Δ = blockdiag(Δ, Δ))
	sol0 = zeros(2Nx, Ny)

eigls = EigArpack(1.0, :LM)
	# eigls = eig_MF_KrylovKit(tol = 1e-8, dim = 60, x₀ = rand(ComplexF64, Nx*Ny), verbose = 1)
	opt_newton = PALC.NewtonPar(tol = 1e-9, verbose = true, eigsolver = eigls, maxIter = 20)
	out, hist, flag = @time PALC.newton(
		x ->  Fcgl(x, par_cgl),
		x ->  Jcgl(x, par_cgl),
		vec(sol0), opt_newton, normN = norminf)
####################################################################################################
# test for the Jacobian expression
# sol0 = rand(2Nx*Ny)
# J0 = ForwardDiff.jacobian(x-> Fcgl(x, par_cgl), sol0) |> sparse
# J1 = Jcgl(sol0, par_cgl)
# norm(J0 - J1, Inf)
####################################################################################################
opts_br = ContinuationPar(dsmin = 0.001, dsmax = 0.005, ds = 0.001, pMax = 2.5, detectBifurcation = 1, nev = 5, plotEveryNsteps = 50, newtonOptions = (@set opt_newton.verbose = false), maxSteps = 1060)

	br, u1 = @time PALC.continuation(
		(x, p) -> Fcgl(x, @set par_cgl.r = p),
		(x, p) -> Jcgl(x, @set par_cgl.r = p),
		vec(sol0), par_cgl.r,
		opts_br, verbosity = 0)
####################################################################################################
ind_hopf = 2
# number of time slices
M = 30
r_hopf, Th, orbitguess2, hopfpt, vec_hopf = PALC.guessFromHopf(br, ind_hopf, opt_newton.eigsolver, M, 22*sqrt(0.1); phase = 0.25)

orbitguess_f2 = reduce(hcat, orbitguess2)
orbitguess_f = vcat(vec(orbitguess_f2), Th) |> vec

poTrap = p -> PeriodicOrbitTrapProblem(
	x ->  Fcgl(x, p),
	x ->  Jcgl(x, p),
	real.(vec_hopf),
	hopfpt.u,
	M)

ls0 = GMRESIterativeSolvers(N = 2Nx*Ny, tol = 1e-9)#, Pl = lu(I + par_cgl.Δ))
poTrapMF = p -> PeriodicOrbitTrapProblem(
	x ->  Fcgl(x, p),
	x ->  (dx -> d1Fcgl(x, p, dx)),
	real.(vec_hopf),
	hopfpt.u,
	M, ls0)

poTrap(@set par_cgl.r = r_hopf - 0.1)(orbitguess_f) |> plot
poTrapMF(@set par_cgl.r = r_hopf - 0.1)(orbitguess_f) |> plot

plot();PALC.plotPeriodicPOTrap(orbitguess_f, M, Nx, Ny; ratio = 2);title!("")
deflationOp = DeflationOperator(2.0,(x,y) -> dot(x[1:end-1],y[1:end-1]),1.0,[zero(orbitguess_f)])

####################################################################################################
# circulant pre-conditioner
# Jpo = poTrap(@set par_cgl.r = r_hopf - 0.1)(Val(:JacFullSparse), orbitguess_f)[1:2n*M-2n,1:2n*M-2n]
# kΔ = spdiagm(0 => ones(M-1), -1 => ones(M-2), M-2 => [1])
# 	kI = spdiagm(0 => ones(M-1), -1 => -ones(M-2), M-2 => [-1])
#
# 	#diagonal precond
# 	# kΔ = spdiagm(0 => ones(M-1))
# 	# kI = spdiagm(0 => ones(M-1))
#
# 	h = orbitguess_f[end] / M
# 	Precs2 = kron(kI, spdiagm(0 => ones(2n))) ./1 -  h/2 * kron(kΔ, par_cgl.Δ)
# 	ls = GMRESIterativeSolvers(verbose = false, tol = 1e-4, N = size(Precs2,1), restart = 20, maxiter = 40, Pl = lu(Precs2), log=true)
# 	ls(Jpo, rand(ls.N))
####################################################################################################
# slow version DO NOT RUN!!!
# opt_po = (@set opt_po.eigsolver = eig_MF_KrylovKit(tol = 1e-4, x₀ = rand(2Nx*Ny), verbose = 2, dim = 20))
opt_po = (@set opt_po.eigsolver = DefaultEig())
opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.03, ds= 0.001, pMax = 2.5, maxSteps = 250, plotEveryNsteps = 3, newtonOptions = (@set opt_po.linsolver = DefaultLS()), computeEigenValues = true, nev = 5, precisionStability = 1e-7, detectBifurcation = 0)
	@assert 1==0 "Too much memory used!"
	br_pok2, upo , _= @time PALC.continuationPOTrap(
			p -> poTrap(@set par_cgl.r = p),
			orbitguess_f, r_hopf - 0.01,
			opts_po_cont, :FullLU;
			verbosity = 2,	plot = true,
			plotSolution = (x ;kwargs...) -> plotPeriodicPOTrap(x, M, Nx, Ny; kwargs...),
			printSolution = (u,p) -> PALC.amplitude(u, Nx*Ny, M), normC = norminf)
####################################################################################################
# we use an ILU based preconditioner for the newton method at the level of the full Jacobian of the PO functional
Jpo = poTrap(@set par_cgl.r = r_hopf - 0.01)(Val(:JacFullSparse), orbitguess_f)
Precilu = @time ilu(Jpo, τ = 0.003)
ls = GMRESIterativeSolvers(verbose = false, tol = 1e-3, N = size(Jpo,1), restart = 40, maxiter = 50, Pl = Precilu, log=true)
	ls(Jpo, rand(ls.N))

# ls = GMRESKrylovKit(verbose = 0, Pl = Precilu, rtol = 1e-3)
	ls(Jpo, rand(size(Jpo,1)))

opt_po = @set opt_newton.verbose = true
	outpo_f, _, flag = @time PALC.newton(poTrapMF(@set par_cgl.r = r_hopf - 0.01),
			orbitguess_f,
			(@set opt_po.linsolver = ls),
			:FullMatrixFree;
			normN = norminf,
			# callback = (x, f, J, res, iteration, options) -> (println("--> amplitude = ", PALC.amplitude(x, Nx*Ny, M; ratio = 2));true)
			)
	flag && printstyled(color=:red, "--> T = ", outpo_f[end], ", amplitude = ", PALC.amplitude(outpo_f, Nx*Ny, M; ratio = 2),"\n")
plot();PALC.plotPeriodicPOTrap(outpo_f, M, Nx, Ny; ratio = 2);title!("")

opt_po = @set opt_po.eigsolver = EigKrylovKit(tol = 1e-3, x₀ = rand(2n), verbose = 2, dim = 25)
# opt_po = @set opt_po.eigsolver = DefaultEig()
opt_po = @set opt_po.eigsolver = EigArpack(; tol = 1e-3, v0 = rand(2n))
opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.03, ds = 0.001, pMax = 2.2, maxSteps = 250, plotEveryNsteps = 3, newtonOptions = (@set opt_po.linsolver = ls), nev = 5, precisionStability = 1e-5, detectBifurcation = 2, dsminBisection =1e-7)
	br_po, _ , _= @time PALC.continuationPOTrap(
			p -> poTrapMF(@set par_cgl.r = p),
			outpo_f, r_hopf - 0.01,
			opts_po_cont, :FullMatrixFree;
			verbosity = 3,	plot = true,
			plotSolution = (x ;kwargs...) -> PALC.plotPeriodicPOTrap(x, M, Nx, Ny; ratio = 2, kwargs...),
			printSolution = (u, p) -> PALC.amplitude(u, Nx*Ny, M; ratio = 2), normC = norminf)

# branches = Any[br_po]
# push!(branches, br_po)
plotBranch([branches[1], br_po]; putbifptlegend = false, label="", xlabel="r",ylabel="Amplitude", legend = :bottomright);title!("")

###################################################################################################
# preconditioner taking into account the constraint
# Jpo = poTrap(@set par_cgl.r = r_hopf - 0.1)(Val(:JacFullSparse), orbitguess_f)[1:2n*M-2n,1:2n*M-2n]
# 	Jpo = blockdiag(Jpo, spdiagm(0 => ones(2Nx*Ny+1)) )
# Precilu = @time ilu(Jpo, τ = 0.004)
# ls = GMRESIterativeSolvers(verbose = false, tol = 1e-3, N = size(Jpo,1), restart = 30, maxiter = 50, Pl = Precilu, log=true)
# 	ls(Jpo, rand(ls.N))
####################################################################################################
# this preconditioner does not work very well here
Jpo = poTrap(@set par_cgl.r = r_hopf - 0.1)(Val(:JacCyclicSparse), orbitguess_f)
Precilu = @time ilu(Jpo, τ = 0.003)
ls = GMRESIterativeSolvers(verbose = false, tol = 1e-3, N = size(Jpo,1), restart = 30, maxiter = 50, Pl = Precilu, log=true)
	ls(Jpo, rand(ls.N))

opt_po = @set opt_newton.verbose = true
	outpo_f, hist, flag = @time PALC.newton(
			poTrapMF(@set par_cgl.r = r_hopf - 0.01),
			orbitguess_f,
			(@set opt_po.linsolver = ls), :BorderedMatrixFree;
			normN = norminf,
			# callback = (x, f, J, res, iteration, options; kwargs...) -> (println("--> amplitude = ", PALC.amplitude(x, Nx*Ny, M));@show kwargs;true)
			)
	plot();plotPeriodic(outpo_f);title!("")

function callbackPO(x, f, J, res, iteration, linsolver = ls, prob = poTrap, p = par_cgl; kwargs...)
	@show kwargs ls.N
	# we update the preconditioner every 10 continuation steps
	if mod(kwargs[:iterationC], 10) == 9 && iteration == 1
		@info "update Preconditioner"
		Jpo = poTrap(@set p.r = kwargs[:p])(Val(:JacCyclicSparse), x)
		Precilu = @time ilu(Jpo, τ = 0.003)
		ls.Pl = Precilu
	end
	true
end

opt_po = (@set opt_po.eigsolver = EigKrylovKit(tol = 1e-4, x₀ = rand(2n), verbose = 2, dim = 20))
opt_po = (@set opt_po.eigsolver = DefaultEig())
opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.03, ds= 0.001, pMax = 2.15, maxSteps = 450, plotEveryNsteps = 3, newtonOptions = (@set opt_po.linsolver = ls), computeEigenValues = true, nev = 5, precisionStability = 1e-7, detectBifurcation = 1)
# opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.01, ds= -0.001, pMax = 1.5, maxSteps = 400, plotEveryNsteps = 3, newtonOptions = (@set opt_po.linsolver = ls))
	br_pok2, upo , _= @time PALC.continuationPOTrap(
			p -> poTrap(@set par_cgl.r = p),
			outpo_f, r_hopf - 0.01,
			opts_po_cont, :BorderedMatrixFree;
			verbosity = 2,	plot = true,
			plotSolution = (x ;kwargs...) -> PALC.plotPeriodicPOTrap(x, M, Nx, Ny; kwargs...),
			callbackN = callbackPO,
			printSolution = (u, p) -> PALC.amplitude(u, Nx*Ny, M), normC = norminf)

# branches = Any[br_pok2]
# push!(branches, br_pok2)
# plotBranch(branches,label="", xlabel="r",ylabel="Amplitude");title!("")
####################################################################################################
# Jpo = poTrap(@set par_cgl.r = r_hopf - 0.1)(Val(:JacFullSparse), orbitguess_f)
# rhs = rand(size(Jpo,1))
# @time Jpo \ rhs
#
# using IterativeSolvers
# 	ls = GMRESIterativeSolvers(verbose = true, tol = 1e-3, N = size(Jpo,1), restart = 10, maxiter = 1000)
# 	ls(Jpo, rand(ls.N))
#
# n = Nx*Ny
# Jpo = @time  poTrap(@set par_cgl.r = r_hopf - 0.1)(Val(:JacFullSparse), orbitguess_f)[1:2n*M-2n,1:2n*M-2n]
# rhs = rand(size(Jpo,1))
# @time Jpo \ rhs
# ls = GMRESIterativeSolvers(verbose = true, tol = 1e-3, N = size(Jpo,1), restart = 10, maxiter = 1000)
# ls(Jpo, rand(ls.N))
#
# kΔ = spdiagm(0 => ones(M-1), -1 => ones(M-2), M-2 => [1])
# 	kI = spdiagm(0 => ones(M-1), -1 => -ones(M-2), M-2 => [-1])
# 	h = orbitguess_f[end] / M
# 	Jcglsp = Jcgl(orbitguess_f[1:2n], par_cgl)
# 	Precs2 = kron(kI, spdiagm(0 => ones(2n))) -  h/2 * kron(kΔ, par_cgl.Δ)
# 	ls = GMRESIterativeSolvers(verbose = true, tol = 1e-3, N = 2n*M-2n, restart = 20, maxiter = 1000, Pl = lu(Precs2), log=true)
#
# ls(Jpo, rand(ls.N))
####################################################################################################
####################################################################################################
# Experimental, full Inplace
function NL!(f, u, p, t = 0.)
	@unpack r, μ, ν, c3, c5 = p
	n = div(length(u), 2)
	u1v = @view u[1:n]
	u2v = @view u[n+1:2n]

	f1 = @view f[1:n]
	f2 = @view f[n+1:2n]

	@inbounds for ii = 1:n
		u1 = u1v[ii]
		u2 = u2v[ii]
		ua = u1^2+u2^2
		f1[ii] = r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1
		f2[ii] = r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2
	end
	return f
end

function dNL!(f, u, p, du)
	@unpack r, μ, ν, c3, c5 = p
	n = div(length(u), 2)
	u1v = @view u[1:n]
	u2v = @view u[n+1:2n]

	du1v = @view du[1:n]
	du2v = @view du[n+1:2n]

	f1 = @view f[1:n]
	f2 = @view f[n+1:2n]

	@inbounds for ii = 1:n
		u1 = u1v[ii]
		u2 = u2v[ii]
		du1 = du1v[ii]
		du2 = du2v[ii]
		ua = u1^2+u2^2
		f1[ii] = (-5*c5*u1^4 + (-6*c5*u2^2 - 3*c3)*u1^2 + 2*μ*u1*u2 - c5*u2^4 - c3*u2^2 + r) * du1 +
		(-4*c5*u2*u1^3 + μ*u1^2 + (-4*c5*u2^3 - 2*c3*u2)*u1 + 3*u2^2*μ - ν) * du2

		f2[ii] = (-4*c5*u2*u1^3 - 3*μ*u1^2 + (-4*c5*u2^3 - 2*c3*u2)*u1 - u2^2*μ + ν) * du1 + (-c5*u1^4 + (-6*c5*u2^2 - c3)*u1^2 - 2*μ*u1*u2 - 5*c5*u2^4 - 3*c3*u2^2 + r) * du2
	end

	return f
end

function Fcgl!(f, u, p, t = 0.)
	NL!(f, u, p)
	mul!(f, p.Δ, u, 1., 1.)
end

function dFcgl!(f, x, p, dx)
	# 19.869 μs (0 allocations: 0 bytes)
	dNL!(f, x, p, dx)
	mul!(f, p.Δ, dx, 1., 1.)
end

sol0f = vec(sol0)
out_ = similar(sol0f)
@time Fcgl!(out_, sol0f, par_cgl)
@time dFcgl!(out_, sol0f, par_cgl, sol0f)

ls0 = GMRESIterativeSolvers(N = 2Nx*Ny, tol = 1e-9)#, Pl = lu(I + par_cgl.Δ))
poTrapMFi = p -> PeriodicOrbitTrapProblem(
	(o, x) ->  Fcgl!(o, x, p),
	(o, x, dx) -> dFcgl!(o, x, p, dx),
	real.(vec_hopf),
	hopfpt.u,
	M, ls0; isinplace = true)

pb = poTrapMFi(@set par_cgl.r = r_hopf - 0.01)
outpo_f, _, flag = @time PALC.newton(pb,
	orbitguess_f, (@set opt_po.linsolver = ls),
	:FullMatrixFree; normN = norminf)

@assert 1==0 "tester map inplace dans gmresIS et voir si allocate"
@assert 1==0 "tester code_warntype dans POTrapFunctionalJac! st POTrapFunctional!"

lsi = GMRESIterativeSolvers!(verbose = false, N = length(orbitguess_f), tol = 1e-3, restart = 40, maxiter = 50, Pl = Precilu, log=true)

outpo_f, _, flag = @time PALC.newton(pb,
	orbitguess_f, (@set opt_po.linsolver = lsi),
	:FullMatrixFreeInplace; normN = norminf)

res = copy(orbitguess_f)
@code_warntype PALC.POTrapFunctional!(pb, res, orbitguess_f)
@code_warntype PALC.POTrapFunctionalJac!(pb, res, orbitguess_f, orbitguess_f)
using LinearMaps
Jmap! = LinearMap{Float64}((o, v) -> PALC.POTrapFunctionalJac!(pb, o, orbitguess_f, v), lsi.N, lsi.N ; ismutating = true)

@time mul!(res, Jmap!, orbitguess_f)

using IterativeSolvers

@time IterativeSolvers.gmres(Jmap!, orbitguess_f; tol = 1e-3, log = true, verbose = true, Pl = Precilu)

Jmap = LinearMap{Float64}((v) -> pb(orbitguess_f, v), lsi.N, lsi.N ; ismutating = false)
@time mul!(res, Jmap, orbitguess_f)
@time IterativeSolvers.gmres(Jmap, orbitguess_f; tol = 1e-3, log = true, verbose = true, Pl = Precilu)

@time IterativeSolvers.gmres(Jpo, orbitguess_f; tol = 1e-3, log = true, verbose = false, Pl = Precilu)

####################################################################################################
# Computation of Fold of limit cycle
function d2Fcglpb(f, x, dx1, dx2)
	return ForwardDiff.derivative(t2 -> ForwardDiff.derivative( t1-> f(x .+ t1 .* dx1 .+ t2 .* dx2,), 0.), 0.)
end

function JacT(F, x, v)
	# ForwardDiff.gradient(x -> dot(v, F(x)), x)
	ReverseDiff.gradient(x -> dot(v, F(x)), x)
end

function JacT2(F, x, v)
	# ForwardDiff.gradient(x -> dot(v, F(x)), x)
	Tracker.gradient(x -> dot(v, F(x)), x)
end

# dsol0 = rand(2n)
# 	dsol1 = rand(2n)
# 	sol0v = vec(sol0)
# 	@time d2Fcgl(sol0v, par_cgl, dsol0, dsol1)

# we look at the second fold point
indfold = 2
foldpt = PALC.FoldPoint(br_po, indfold)

Jpo = poTrap(@set par_cgl.r = r_hopf - 0.1)(Val(:JacFullSparse), orbitguess_f)
Precilu = @time ilu(Jpo, τ = 0.005)
ls = GMRESIterativeSolvers(verbose = false, tol = 1e-5, N = size(Jpo, 1), restart = 40, maxiter = 60, Pl = Precilu, log=true)
	ls(Jpo, rand(ls.N))

outfold, hist, flag = @time PALC.newtonFold(
		(x, p) -> poTrap(@set par_cgl.r = p)(x),
		(x, p) -> poTrap(@set par_cgl.r = p)(Val(:JacFullSparse), x),
		(x, p) -> transpose(poTrap(@set par_cgl.r = p)(Val(:JacFullSparse), x)),
		(x, p, dx1, dx2) -> d2Fcglpb(poTrap(@set par_cgl.r = p), x, dx1, dx2),
		br_pok2, indfold, #index of the fold point
		@set opt_po.linsolver = ls)
		# opt_po)
	flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold.p," from ", br_pok2.bifpoint[indfold][3],"\n")

outfold, hist, flag = @time PALC.newtonFold(
		(x, p) -> poTrapMF(@set par_cgl.r = p)(x),
		(x, p) -> (dx -> poTrapMF(@set par_cgl.r = p)(x, dx)),
		(x, p) -> (dx -> JacT(x -> poTrapMF(@set par_cgl.r = p)(x), x, dx)),
		(x, p, dx1, dx2) -> d2Fcglpb(poTrap(@set par_cgl.r = p), x, dx1, dx2),
		br_pok2, indfold, #index of the fold point
		@set opt_po.linsolver = ls)
		# opt_po)
	flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold.p," from ", br_pok2.bifpoint[indfold][3],"\n")

@time JacT( poTrapMF(@set par_cgl.r = 0.95), orbitguess_f, orbitguess_f)

@time JacT2( poTrapMF(@set par_cgl.r = 0.95), orbitguess_f, orbitguess_f)

optcontfold = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 40.1, pMin = -10., newtonOptions = (@set opt_po.linsolver = ls), maxSteps = 20)

# optcontfold = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 40.1, pMin = -10., newtonOptions = opt_po, maxSteps = 10)

outfoldco, hist, flag = @time PALC.continuationFold(
	(x, r, p) -> poTrap(setproperties(par_cgl, (r=r, c5=p)))(x),
	(x, r, p) -> poTrap(setproperties(par_cgl, (r=r, c5=p)))(Val(:JacFullSparse), x),
	(x, r, p) -> transpose(poTrap(setproperties(par_cgl, (r=r, c5=p)))(Val(:JacFullSparse), x)),
	p -> ((x, r, dx1, dx2) -> d2Fcglpb(poTrap(setproperties(par_cgl, (r=r, c5=p))), x, dx1, dx2)),
	br_pok2, indfold,
	par_cgl.c5, plot = true, verbosity = 2,
	optcontfold)

plotBranch(outfoldco, label="", xlabel="c5", ylabel="r");title!("")

####################################################################################################
# Continuation of periodic orbits on the GPU
using CuArrays
CuArrays.allowscalar(false)
import LinearAlgebra: mul!, axpby!
mul!(x::CuArray, y::CuArray, α::T) where {T <: Number} = (x .= α .* y)
mul!(x::CuArray, α::T, y::CuArray) where {T <: Number} = (x .= α .* y)
axpby!(a::T, X::CuArray, b::T, Y::CuArray) where {T <: Number} = (Y .= a .* X .+ b .* Y)

par_cgl_gpu = @set par_cgl.Δ = CuArrays.CUSPARSE.CuSparseMatrixCSC(par_cgl.Δ);
Jpo = poTrap(@set par_cgl.r = r_hopf - 0.01)(Val(:JacFullSparse), orbitguess_f)
Precilu = @time ilu(Jpo, τ = 0.003)

struct LUperso
	L
	Ut	# transpose of U in LU decomposition
end

import Base: ldiv!
# https://github.com/JuliaDiffEq/DiffEqBase.jl/blob/master/src/init.jl#L146-L150
function LinearAlgebra.ldiv!(_lu::LUperso, rhs::Array)
	@show "bla"
	_x = (_lu.Ut) \ ((_lu.L) \ rhs)
	rhs .= vec(_x)
	# CuArrays.unsafe_free!(_x)
	rhs
end

function LinearAlgebra.ldiv!(_lu::LUperso, rhs::CuArrays.CuArray)
	_x = UpperTriangular(_lu.Ut) \ (LowerTriangular(_lu.L) \ rhs)
	rhs .= vec(_x)
	CuArrays.unsafe_free!(_x)
	rhs
end

import PseudoArcLengthContinuation: extractPeriodFDTrap
extractPeriodFDTrap(x::CuArray) = x[end:end]

sol0_f = vec(sol0)
	sol0gpu = CuArray(sol0_f)
	_dxh = rand(length(sol0_f))
	_dxd = CuArray(_dxh)

	outh = Fcgl(sol0_f, par_cgl);
	outd = Fcgl(sol0gpu, par_cgl_gpu);
	@assert norm(outh-Array(outd), Inf) < 1e-12

	outh = dFcgl(sol0_f, par_cgl, _dxh);
	outd = dFcgl(sol0gpu, par_cgl_gpu, _dxd);
	@assert norm(outh-Array(outd), Inf) < 1e-12

orbitguess_cu = CuArray(orbitguess_f)
norm(orbitguess_f - Array(orbitguess_cu), Inf)

Precilu_gpu = LUperso(LowerTriangular(CuArrays.CUSPARSE.CuSparseMatrixCSR(I+Precilu.L)), UpperTriangular(CuArrays.CUSPARSE.CuSparseMatrixCSR(sparse(Precilu.U'))));

Precilu_host = LUperso((I+Precilu.L), (sparse(Precilu.U')));

rhs = rand(size(Jpo,1))
	sol_0 = Precilu \ rhs
	sol_1 = UpperTriangular(sparse(Precilu.U')) \ (LowerTriangular(I+Precilu.L)  \ (rhs))
	# sol_2 = LowerTriangular(Precilu.U') \ (LowerTriangular(sparse(I+Precilu.L))  \ (rhs))
	norm(sol_1-sol_0, Inf64)
	# norm(sol_2-sol_0, Inf64)

sol_0 = (I+Precilu.L) \ rhs
	sol_1 = LowerTriangular(CuArrays.CUSPARSE.CuSparseMatrixCSR(I+Precilu.L)) \ CuArray(rhs)
	@assert norm(sol_0-Array(sol_1), Inf64) < 1e-10

sol_0 = (Precilu.U)' \ rhs
	sol_1 = UpperTriangular(CuArrays.CUSPARSE.CuSparseMatrixCSR(sparse(Precilu.U'))) \ CuArray(rhs)
	norm(sol_0-Array(sol_1), Inf64)
	@assert norm(sol_0-Array(sol_1), Inf64) < 1e-10

sol_0 = Precilu \ rhs
	sol_1 = ldiv!(Precilu_host, copy(rhs))
	@assert norm(sol_1-sol_0, Inf64) < 1e-10

sol_0 = ldiv!(Precilu_gpu, copy(CuArray(rhs)));
	sol_1 = ldiv!(Precilu_host, copy(rhs))
	norm(sol_1-Array(sol_0), Inf64)
	@assert norm(sol_1-Array(sol_0), Inf64) < 1e-10

# matrix-free problem on the gpu
ls0gpu = GMRESKrylovKit(rtol = 1e-9)
poTrapMFGPU = p -> PeriodicOrbitTrapProblem(
	x ->  Fcgl(x, p),
	x ->  (dx -> dFcgl(x, p, dx)),
	CuArray(real.(vec_hopf)),
	CuArray(hopfpt.u),
	M, ls0gpu)

pb = poTrapMF(@set par_cgl.r = r_hopf - 0.1);
pbgpu = poTrapMFGPU(@set par_cgl_gpu.r = r_hopf - 0.1);

pbgpu(orbitguess_cu);
pbgpu(orbitguess_cu, orbitguess_cu);

ls = GMRESKrylovKit(verbose = 2, Pl = Precilu, rtol = 1e-3, dim  = 20)
	outh, _, _ = @time ls((Jpo), orbitguess_f) #0.4s

lsgpu = GMRESKrylovKit(verbose = 2, Pl = Precilu_gpu, rtol = 1e-3, dim  = 20)
	Jpo_gpu = CuArrays.CUSPARSE.CuSparseMatrixCSR(Jpo);
	outd, _, _ = @time lsgpu(Jpo_gpu, orbitguess_cu)

@assert norm(outh-Array(outd), Inf) < 1e-12

outh = @time pb(orbitguess_f);
	outd = @time pbgpu(orbitguess_cu);
	norm(outh-Array(outd), Inf)

_dxh = rand(length(orbitguess_f));
	_dxd = CuArray(_dxh);
	outh = @time pb(orbitguess_f, _dxh);
	outd = @time pbgpu(orbitguess_cu, _dxd);
	norm(outh-Array(outd), Inf)

outpo_f, hist, flag = @time PALC.newton(
		poTrapMF(@set par_cgl.r = r_hopf - 0.01),
		orbitguess_f,
		(@set opt_po.linsolver = ls), :FullMatrixFree;
		normN = x -> maximum(abs.(x)),
		# callback = (x, f, J, res, iteration, options) -> (println("--> amplitude = ", amplitude(x));true)
		) #14s
flag && printstyled(color=:red, "--> T = ", outpo_f[end], ", amplitude = ", amplitude(outpo_f, Nx*Ny, M),"\n")

opt_po = @set opt_newton.verbose = true
	outpo_f, hist, flag = @time PALC.newton(
			poTrapMFGPU(@set par_cgl_gpu.r = r_hopf - 0.01),
			orbitguess_cu,
			(@set opt_po.linsolver = lsgpu), :FullMatrixFree;
			normN = x -> maximum(abs.(x)),
			# callback = (x, f, J, res, iteration, options) -> (println("--> amplitude = ", PALC.amplitude(x, Nx*Ny, M));true)
			) #7s
	flag && printstyled(color=:red, "--> T = ", outpo_f[end:end], ", amplitude = ", amplitude(outpo_f, Nx*Ny, M),"\n")

opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.03, ds= 0.001, pMax = 2.2, maxSteps = 35, plotEveryNsteps = 3, newtonOptions = (@set opt_po.linsolver = lsgpu))
	br_pok2, upo , _= @time PALC.continuationPOTrap(
		p -> poTrapMFGPU(@set par_cgl_gpu.r = p),
		orbitguess_cu, r_hopf - 0.01,
		opts_po_cont, :FullMatrixFree;
		verbosity = 2,
		printSolution = (u,p) -> PALC.amplitude(u, Nx*Ny, M), normC = x -> maximum(abs.(x)))