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.
=#

module DistributedPass

#using Debug

import Base.show

using CompilerTools
import CompilerTools.DebugMsg
DebugMsg.init()

using CompilerTools.AstWalker
import CompilerTools.ReadWriteSet
using CompilerTools.LambdaHandling
using CompilerTools.Helper

import HPAT
using HPAT.Partitioning
using HPAT.SEQ
using HPAT.TWO_D
using HPAT.ONE_D
using HPAT.ONE_D_VAR

using ParallelAccelerator
import ParallelAccelerator.ParallelIR
import ParallelAccelerator.ParallelIR.isArrayType
import ParallelAccelerator.ParallelIR.getParforNode
import ParallelAccelerator.ParallelIR.isBareParfor
import ParallelAccelerator.ParallelIR.isAllocation
import ParallelAccelerator.ParallelIR.isAllocationAssignment
import ParallelAccelerator.ParallelIR.TypedExpr
import ParallelAccelerator.ParallelIR.get_alloc_shape
import ParallelAccelerator.ParallelIR.computeLiveness
import ParallelAccelerator.ParallelIR.PIRParForAst
import ParallelAccelerator.ParallelIR.isBareParfor

import ParallelAccelerator.ParallelIR.ISCAPTURED
import ParallelAccelerator.ParallelIR.ISASSIGNED
import ParallelAccelerator.ParallelIR.ISASSIGNEDBYINNERFUNCTION
import ParallelAccelerator.ParallelIR.ISCONST
import ParallelAccelerator.ParallelIR.ISASSIGNEDONCE
import ParallelAccelerator.ParallelIR.ISPRIVATEPARFORLOOP
import ParallelAccelerator.ParallelIR.PIRReduction

mk_call(fun,args) = Expr(:call, fun, args...)

function mk_mult_int_expr(args::Array)
    if length(args)==0
        return 1
    elseif length(args)==1
        return args[1]
    end
    next = 2
    prev_expr = args[1]

    while next<=length(args)
        m_call = mk_call(GlobalRef(Base,:mul_int),[prev_expr,args[next]])
        prev_expr  = mk_call(GlobalRef(Base,:box),[Int64,m_call])
        next += 1
    end
    return prev_expr
end

function mk_add_int_expr(args::Array)
    if length(args)==0
        return 0
    elseif length(args)==1
        return args[1]
    end
    next = 2
    prev_expr = args[1]

    while next<=length(args)
        m_call = mk_call(GlobalRef(Base,:add_int),[prev_expr,args[next]])
        prev_expr  = mk_call(GlobalRef(Base,:box),[Int64,m_call])
        next += 1
    end
    return prev_expr
end

mk_add_int_expr(a,b) = mk_call(GlobalRef(Base,:box),[Int64, mk_call(GlobalRef(Base,:add_int),[a,b])])
mk_sub_int_expr(a,b) = mk_call(GlobalRef(Base,:box),[Int64, mk_call(GlobalRef(Base,:sub_int),[a,b])])
mk_div_int_expr(a,b) = mk_call(GlobalRef(Base,:box),[Int64, mk_call(GlobalRef(Base,:sdiv_int),[a,b])])

mk_add_float_expr(a,b) = mk_call(GlobalRef(Base,:box),[Float64, mk_call(GlobalRef(Base,:add_float),[a,b])])
mk_mult_float_expr(a,b) = mk_call(GlobalRef(Base,:box),[Float64, mk_call(GlobalRef(Base,:mul_float),[a,b])])

dist_ir_funcs = Set([   :unsafe_arrayref,
                        :unsafe_arrayset,
                        :getindex,
                        :setindex,
                        :__hpat_data_source_HDF5_open,
                        :__hpat_data_sink_HDF5_open,
                        :__hpat_data_sink_HDF5_create,
                        :__hpat_data_sink_TXT_open,
                        :__hpat_data_source_HDF5_size,
                        :__hpat_get_H5_dim_size,
                        :__hpat_data_source_HDF5_read,
                        :__hpat_data_sink_HDF5_write,
                        :__hpat_data_sink_TXT_write,
                        :__hpat_data_source_TXT_open,
                        :__hpat_data_source_TXT_size,
                        :__hpat_get_TXT_dim_size,
                        :__hpat_data_source_TXT_read,
                        :__hpat_filter,
                        :__hpat_join,
                        :__hpat_aggregate,
                        :Kmeans,
                        :LinearRegression,
                        :NaiveBayes,
                        :arraylen, :arraysize, :reshape, :tuple, :hcat, :typed_hcat, :vcat,
                        :convert,
                        :transpose!, :transpose, :gemm_wrapper!,
                        :__hpat_transpose_hcat,
                        :gemv!, :cumsum!])

dist_opt = false
"""
Controls whether distributed pass tries to optimize (e.g. expand matrix multiply).
"""
function DistOptimize(x :: Bool)
   global dist_opt = x
end

# ENTRY to distributedIR
function from_root(function_name, ast::Tuple)
    @dprintln(1,"Starting main DistributedPass.from_root.  function = ", function_name, " ast = ", ast)

    (linfo, body) = ast
    lives = computeLiveness(body, linfo)
    user_partitionings = get_user_partitionings(body)
    deps = CompilerTools.TransitiveDependence.computeDependenciesAST(body, ParallelAccelerator.ParallelIR.dependenceCB, nothing)
    state::DistPassState = initDistState(linfo,lives,user_partitionings,deps)

    # find if an array should be partitioned, sequential, or shared
    getArrayDistributionInfo(body, state)
    # increase unique id to max of parfor ids to avoid conflicts for new parfors
    state.uniqueId = 1
    if length(keys(state.parfor_partitioning))>0
        state.uniqueId += maximum(keys(state.parfor_partitioning)) + 1
    end
    # perform domain-specific optimizations like expanding matrix multiply
    if dist_opt
        body = dist_optimize(body, state)
        @dprintln(3,"DistributedPass after dist_optimize = ", body)
        @dprintln(3,"DistributedPass after dist_optimize = ", state.LambdaVarInfo)
    end

    # transform body
    body.args = from_toplevel_body(body.args, state)
    @dprintln(1,"DistributedPass.from_root returns function = ", function_name, " ast = ", state.LambdaVarInfo, body)
    return state.LambdaVarInfo, body
end

type ArrDistInfo
    partitioning::Partitioning      # partitioning of array (SEQ, ONE_D_VAR, TWO_D, ONE_D)
    dim_sizes::Array{Union{RHSVar,Int,Expr},1}      # sizes of array dimensions
    # assuming only last dimension is partitioned
    arr_id::Int # assign ID to distributed array to access partitioning info later

    # array partitioning information similar to HDF5 hyperslab interface
    # for 1D, only start and count of last dimension are set
    # for 2D, these are set to mirror ScaLAPACK's block cyclic distribution
    starts::Array{LHSVar}
    counts::Array{LHSVar}
    strides::Array{LHSVar}
    blocks::Array{LHSVar}
    local_sizes::Array{LHSVar}
    leftovers::Array{LHSVar}

    function ArrDistInfo(num_dims::Int)
        # one dimensional partitioning is default
        new(ONE_D, zeros(Int64,num_dims),0,
        Array(LHSVar,num_dims), Array(LHSVar,num_dims),
        Array(LHSVar,num_dims), Array(LHSVar,num_dims),
        Array(LHSVar,num_dims), Array(LHSVar,num_dims))
    end
end

function show(io::IO, pnode::HPAT.DistributedPass.ArrDistInfo)
    print(io,"partitioning: ",pnode.partitioning," sizes: [")
    for i in 1:length(pnode.dim_sizes)
        print(io,pnode.dim_sizes[i])
        if i!=length(pnode.dim_sizes)
            print(io,",")
        end
    end
    print(io,"]")
end

# information about AST gathered and used in DistributedPass
type DistPassState
    # information about all arrays
    arrs_dist_info::Dict{LHSVar,ArrDistInfo}
    parfor_partitioning::Dict{Int,Partitioning}
    parfor_arrays::Dict{Int,Vector{LHSVar}}
    LambdaVarInfo::LambdaVarInfo
    uniqueId::Int
    lives  :: CompilerTools.LivenessAnalysis.BlockLiveness
    # keep values for constant tuples. They are often used for allocating and reshaping arrays.
    tuple_table              :: Dict{LHSVar,Array{Union{LHSVar,Int},1}}
    max_label :: Int # holds the max number of all LabelNodes
    user_partitionings::Dict{Symbol,Partitioning}
    dist_vars::Dict{Symbol,LHSVar} # a dictionary for distributed variables such as __hpat_num_pes
    deps
    parfor_stencils::Dict{Int,Vector{Int}}
    adjust_alloc::Dict{LHSVar,Vector{Union{LHSVar,Int,Expr}}} # adjust allocation for array after transformation
    function DistPassState(linfo, lives, user_partitionings, deps)
        new(Dict{LHSVar, Array{ArrDistInfo,1}}(), Dict{Int,Partitioning}(), Dict{Int,Vector{LHSVar}}(), linfo,0, lives,
             Dict{LHSVar,Array{Union{LHSVar,Int},1}}(),0,user_partitionings, Dict{Symbol,LHSVar}(), deps, Dict{Int,Vector{Int}}(), Dict{LHSVar,Vector{Int}}())
    end
end

isSEQ(arr,state) = (state.arrs_dist_info[arr].partitioning==SEQ)
isONE_D_VAR(arr,state) = (state.arrs_dist_info[arr].partitioning==ONE_D_VAR)
isONE_D(arr,state) = (state.arrs_dist_info[arr].partitioning==ONE_D)
isTWO_D(arr,state) = (state.arrs_dist_info[arr].partitioning==TWO_D)

function setSEQ(arr,state)
  state.arrs_dist_info[arr].partitioning=SEQ
end

function setTWO_D(arr,state)
  state.arrs_dist_info[arr].partitioning=TWO_D
end

getArrayPartitioning(arr,state) = state.arrs_dist_info[arr].partitioning

function setArrayPartitioning(arr,part,state)
  state.arrs_dist_info[arr].partitioning=part
end

include("distributed-pass-ast-walk.jl")
include("distributed-pass-optimize.jl")

function show(io::IO, pnode::HPAT.DistributedPass.DistPassState)
    println(io,"DistPassState arrs_dist_info:")
    for (k,v) in pnode.arrs_dist_info
        println(io,"  ", k,"  ",CompilerTools.LambdaHandling.lookupVariableName(k,pnode.LambdaVarInfo)," => ",v)
    end
    println(io,"DistPassState parfor_partitioning: ",pnode.parfor_partitioning)
    #= println(io,"DistPassState parfor_info:")
    for i in pnode.parfor_info
        println(io,"  ", i)
    end
    println(io,"DistPassState seq_parfors:")
    for i in pnode.seq_parfors
        print(io," ", i)
    end=#
    println(io,"")
end

function get_user_partitionings(body)
    first_arg = body.args[1]
    if isa(first_arg, Expr) && first_arg.head==:meta
        for meta in first_arg.args
            if meta.head==:hpat_partitioning
                @dprintln(3, "hpat partitionings found: ", meta)
                return meta.args[1]
            end
        end
    end
    return Dict{Symbol,Symbol}()
end

function initDistState(linfo::LambdaVarInfo, lives, user_partitionings, deps)
    state = DistPassState(linfo, lives, user_partitionings, deps)

    vars = getLocalVariables(linfo)
    # Populate the symbol table
    for var in vars
        typ = CompilerTools.LambdaHandling.getType(var, linfo)
        if isArrayType(typ)
            arrInfo = ArrDistInfo(ndims(typ))
            state.arrs_dist_info[var] = arrInfo
        end
    end

    return state
end

# nodes are :body of AST
function from_toplevel_body(nodes::Array{Any,1}, state::DistPassState)
    state.max_label = ParallelIR.getMaxLabel(state.max_label, nodes)
    res::Array{Any,1} = genDistributedInit(state)
    for node in nodes
        new_exprs = from_expr(node, state)
        append!(res, new_exprs)
    end
    #ParallelAccelerator.ParallelIR.hoistAllocation
    return res
end

# nodes are :body of Parfor
function from_nested_body(nodes::Array{Any,1}, state::DistPassState)
    res = Any[]
    for node in nodes
        new_exprs = from_expr(node, state)
        append!(res, new_exprs)
    end
    #ParallelAccelerator.ParallelIR.hoistAllocation
    return res
end

function from_expr(node::Expr, state::DistPassState)
    head = node.head
    if head==:(=)
        return from_assignment(node, state, toLHSVar(node.args[1]), node.args[2])
    elseif head==:parfor
        return from_parfor(node, state)
    #elseif head==:block
    elseif head==:call
        return from_call(node, state)
    else
        return [node]
    end
end

function from_expr(node::Any, state::DistPassState)
    return [node]
end

# generates initialization code for distributed execution
function genDistributedInit(state::DistPassState)
    initCall = Expr(:call, GlobalRef(HPAT.API,:hpat_dist_init))
    numPesCall = Expr(:call, GlobalRef(HPAT.API,:hpat_dist_num_pes))
    nodeIdCall = Expr(:call, GlobalRef(HPAT.API,:hpat_dist_node_id))

    pes_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        Symbol("__hpat_num_pes"), Int32, ISASSIGNEDONCE | ISASSIGNED,state.LambdaVarInfo))
    node_id_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        Symbol("__hpat_node_id"), Int32, ISASSIGNEDONCE | ISASSIGNED, state.LambdaVarInfo))

    state.dist_vars[:num_pes] = pes_var
    state.dist_vars[:node_id] = node_id_var

    num_pes_assign = Expr(:(=), pes_var, numPesCall)
    node_id_assign = Expr(:(=), node_id_var, nodeIdCall)
    res = Any[initCall; num_pes_assign; node_id_assign]
    if haskey(ENV, "ENABLE_GAAS")
        # Just to make things working for now
        # Add variables like above so that it does not break in future
        initCallGAAS = Expr(:call, GlobalRef(HPAT.API,:hpat_dist_init_gaas))
        append!(res,[initCallGAAS])
    end

    # generate 2D init if there is any 2D array or parfor
    if any([state.arrs_dist_info[arr].partitioning==TWO_D for arr in keys(state.arrs_dist_info)]) ||
        any([state.parfor_partitioning[parfor_id]==TWO_D for parfor_id in keys(state.parfor_partitioning)])
      @dprintln(3,"DistPass generating 2D init")
      extra_2D_includes = """#include <mkl_blacs.h>
                            #include <mkl_scalapack.h>
                            #include <mkl_pblas.h>
                            extern "C" {
                            void descinit_(MKL_INT*,MKL_INT*,MKL_INT*,MKL_INT*,MKL_INT*,MKL_INT*,MKL_INT*,MKL_INT*,MKL_INT*,MKL_INT*);
                            MKL_INT numroc_(const MKL_INT*,const MKL_INT*,const MKL_INT*,const MKL_INT*,const MKL_INT*);
                            }
                            """
      HPAT.addHpatInclude(extra_2D_includes,"","")
      initCall2d = Expr(:call, GlobalRef(HPAT.API,:hpat_dist_2d_init))
      num_pes_x_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          Symbol("__hpat_num_pes_x"), Int32, ISASSIGNEDONCE | ISASSIGNED, state.LambdaVarInfo))
      num_pes_y_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          Symbol("__hpat_num_pes_y"), Int32, ISASSIGNEDONCE | ISASSIGNED, state.LambdaVarInfo))
      node_id_x_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          Symbol("__hpat_node_id_x"), Int32, ISASSIGNEDONCE | ISASSIGNED, state.LambdaVarInfo))
      node_id_y_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          Symbol("__hpat_node_id_y"), Int32, ISASSIGNEDONCE | ISASSIGNED, state.LambdaVarInfo))

      state.dist_vars[:num_pes_x] = num_pes_x_var
      state.dist_vars[:num_pes_y] = num_pes_y_var
      state.dist_vars[:node_id_x] = node_id_x_var
      state.dist_vars[:node_id_y] = node_id_y_var

      res2 = Any[initCall2d]
      res = [res;res2]
    end

    return res
end

function from_assignment(node::Expr, state::DistPassState, lhs::LHSVar, rhs::Expr)
    @assert node.head==:(=) "DistributedPass invalid assignment head"
    if isAllocation(rhs)
      return from_assignment_alloc(node,state,lhs,rhs)
    elseif rhs.head==:call && (isBaseFunc(rhs.args[1],:reshape) ||
           rhs.args[1]==GlobalRef(ParallelAccelerator.API,:reshape))
      return from_assignment_reshape(node,state,lhs,rhs)
    elseif rhs.head==:call && isBaseFunc(rhs.args[1],:vcat)
        return from_assignment_vcat(node,state,lhs,rhs)
    elseif rhs.head==:call && isBaseFunc(rhs.args[1],:gemm_wrapper!)
        return from_assignment_gemm(node,state,lhs,rhs)
    elseif rhs.head==:call && isBaseFunc(rhs.args[1],:gemv!)
        return from_assignment_gemv(node,state,lhs,rhs)
    else
        node.args[2] = from_expr(rhs,state)[1]
    end
    return [node]
end

function from_assignment(node::Expr, state::DistPassState, lhs::LHSVar, rhs::ANY)
  node.args[2] = from_expr(rhs,state)[1]
  return [node]
end

function from_assignment_alloc(node::Expr, state::DistPassState, arr::LHSVar, rhs::Expr)
  @dprintln(3,"from assingment alloc: ", node)
  if isONE_D(arr,state)
      @dprintln(3,"DistPass allocation array: ", arr)
      #shape = get_alloc_shape(node.args[2].args[2:end])
      #old_size = shape[end]
      dim_sizes = state.arrs_dist_info[arr].dim_sizes
      # generate array division
      # simple 1D partitioning of last dimension, more general partitioning needed
      # match common big data matrix reperesentation
      arr_tot_size = dim_sizes[end]

      # Not sure if need this in future
      # if arr_tot_size == -1
      #     return [node]
      # end

      arr_id = getDistNewID(state)
      state.arrs_dist_info[arr].arr_id = arr_id
      darr_start_var_name = Symbol("__hpat_dist_arr_start_"*string(arr_id))
      darr_div_var_name = Symbol("__hpat_dist_arr_div_"*string(arr_id))
      darr_count_var_name = Symbol("__hpat_dist_arr_count_"*string(arr_id))

      darr_start_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          darr_start_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
      darr_div_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          darr_div_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
      darr_count_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          darr_count_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

      state.arrs_dist_info[arr].starts[end] = darr_start_var
      state.arrs_dist_info[arr].counts[end] = darr_count_var

      darr_div_expr = Expr(:(=),darr_div_var, mk_div_int_expr(arr_tot_size, state.dist_vars[:num_pes]))
      # zero-based index to match C interface of HDF5
      darr_start_expr = Expr(:(=), darr_start_var, mk_mult_int_expr([state.dist_vars[:node_id], darr_div_var]))
      # darr_count_expr = :($darr_count_var = __hpat_node_id==__hpat_num_pes-1 ? $arr_tot_size-__hpat_node_id*$darr_div_var : $darr_div_var)
      darr_count_expr = Expr(:(=), darr_count_var, mk_call(
          GlobalRef(HPAT.API,:__hpat_get_node_portion),
              [arr_tot_size, darr_div_var, state.dist_vars[:num_pes], state.dist_vars[:node_id]]))

      # set new divided allocation size
      rhs.args[end-1] = darr_count_var

      res = [darr_div_expr; darr_start_expr; darr_count_expr; node]
      #debug_size_print = :(println("size ",$darr_count_var))
      #push!(res,debug_size_print)
      return res
  elseif isTWO_D(arr,state)
    return from_assignment_alloc_2d(node, state, arr, rhs)
  end
  return [node]
end

""" 2D block cyclic distribution to match ScaLAPACK's interface.
Keep start, stride, and count to match HDF5's interface.
Assign extra blocks and leftover rows as well.
"""
function from_assignment_alloc_2d(node::Expr, state::DistPassState, arr::LHSVar, rhs::Expr)
  arr_id = getDistNewID(state)
  state.arrs_dist_info[arr].arr_id = arr_id
  dim_sizes = state.arrs_dist_info[arr].dim_sizes
  arr_tot_size_x = dim_sizes[end-1]
  arr_tot_size_y = dim_sizes[end]

  # constant block size for block-cyclic partitioning
  BLOCK_SIZE = HPAT.BLOCK_SIZE
  block_size_var_name = Symbol("__hpat_dist_arr_2d_block_size_"*string(arr_id))
  block_size_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      block_size_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  # same block size for both dimensions
  # block size should be less than smaller dimension
  # block_size_expr = Expr(:(=), block_size_var, mk_call(GlobalRef(HPAT.API,:__hpat_min),[:($BLOCK_SIZE), dim_sizes[end], dim_sizes[end-1]])
  block_size_expr = Expr(:(=), block_size_var, :($BLOCK_SIZE))
  state.arrs_dist_info[arr].blocks[end-1] = state.arrs_dist_info[arr].blocks[end] = block_size_var

  # start = id*block_size
  start_x_var_name = Symbol("__hpat_dist_arr_2d_start_x_"*string(arr_id))
  start_y_var_name = Symbol("__hpat_dist_arr_2d_start_y_"*string(arr_id))
  start_x_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      start_x_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  start_y_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      start_y_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

  start_x_expr = Expr(:(=), start_x_var, mk_mult_int_expr(
      [state.dist_vars[:node_id_x],block_size_var]))
  start_y_expr = Expr(:(=), start_y_var, mk_mult_int_expr(
      [state.dist_vars[:node_id_y],block_size_var]))
  state.arrs_dist_info[arr].starts[end-1] = start_x_var
  state.arrs_dist_info[arr].starts[end] = start_y_var

  # stride = num_pes*block_size
  stride_x_var_name = Symbol("__hpat_dist_arr_2d_stride_x_"*string(arr_id))
  stride_y_var_name = Symbol("__hpat_dist_arr_2d_stride_y_"*string(arr_id))
  stride_x_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      stride_x_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  stride_y_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      stride_y_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

  stride_x_expr = Expr(:(=), stride_x_var, mk_mult_int_expr([state.dist_vars[:num_pes_x],block_size_var]))
  stride_y_expr = Expr(:(=), stride_y_var, mk_mult_int_expr([state.dist_vars[:num_pes_y],block_size_var]))
  state.arrs_dist_info[arr].strides[end-1] = stride_x_var
  state.arrs_dist_info[arr].strides[end] = stride_y_var

  # calculate number of blocks in each dimension (excluding leftover rows/columns)
  # total_size/block_size
  num_blocks_x_var_name = Symbol("__hpat_dist_arr_2d_num_blocks_x_"*string(arr_id))
  num_blocks_y_var_name = Symbol("__hpat_dist_arr_2d_num_blocks_y_"*string(arr_id))
  num_blocks_x_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      num_blocks_x_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  num_blocks_y_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      num_blocks_y_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

  nb_x_div_expr = Expr(:(=),num_blocks_x_var, mk_div_int_expr(arr_tot_size_x, block_size_var))
  nb_y_div_expr = Expr(:(=),num_blocks_y_var, mk_div_int_expr(arr_tot_size_y, block_size_var))

  # number of blocks per local PE in each dimension
  # num_blocks/num_pes + possible extra block
  blocks_per_pe_x_var_name = Symbol("__hpat_dist_arr_2d_blocks_per_pe_x_"*string(arr_id))
  blocks_per_pe_y_var_name = Symbol("__hpat_dist_arr_2d_blocks_per_pe_y_"*string(arr_id))
  blocks_per_pe_x_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      blocks_per_pe_x_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  blocks_per_pe_y_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      blocks_per_pe_y_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

  bppx_div_expr = Expr(:(=),blocks_per_pe_x_var, mk_div_int_expr(num_blocks_x_var,state.dist_vars[:num_pes_x]))
  bppy_div_expr = Expr(:(=),blocks_per_pe_y_var, mk_div_int_expr(num_blocks_y_var,state.dist_vars[:num_pes_y]))
  extra_block_call_x = mk_call(GlobalRef(HPAT.API,:__hpat_add_extra_block),
      [blocks_per_pe_x_var, num_blocks_x_var,state.dist_vars[:node_id_x],state.dist_vars[:num_pes_x]])
  extra_block_call_y = mk_call(GlobalRef(HPAT.API,:__hpat_add_extra_block),
      [blocks_per_pe_y_var, num_blocks_y_var,state.dist_vars[:node_id_y],state.dist_vars[:num_pes_y]])
  state.arrs_dist_info[arr].counts[end-1] = blocks_per_pe_x_var
  state.arrs_dist_info[arr].counts[end] = blocks_per_pe_y_var

  leftovers_x_var_name = Symbol("__hpat_dist_arr_2d_leftovers_x_"*string(arr_id))
  leftovers_y_var_name = Symbol("__hpat_dist_arr_2d_leftovers_y_"*string(arr_id))
  leftovers_x_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      leftovers_x_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  leftovers_y_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      leftovers_y_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

  leftovers_x_expr = Expr(:(=), leftovers_x_var, mk_call(GlobalRef(HPAT.API,:__hpat_get_leftovers),
      [num_blocks_x_var,state.dist_vars[:node_id_x],state.dist_vars[:num_pes_x],arr_tot_size_x,block_size_var]))
  leftovers_y_expr = Expr(:(=), leftovers_y_var, mk_call(GlobalRef(HPAT.API,:__hpat_get_leftovers),
      [num_blocks_y_var,state.dist_vars[:node_id_y],state.dist_vars[:num_pes_y],arr_tot_size_y,block_size_var]))
  state.arrs_dist_info[arr].leftovers[end-1] = leftovers_x_var
  state.arrs_dist_info[arr].leftovers[end] = leftovers_y_var

  # local sizes
  loc_size_x_var_name = Symbol("__hpat_dist_arr_2d_loc_size_x_"*string(arr_id))
  loc_size_y_var_name = Symbol("__hpat_dist_arr_2d_loc_size_y_"*string(arr_id))
  loc_size_x_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      loc_size_x_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  loc_size_y_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      loc_size_y_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

  loc_size_x_expr = Expr(:(=), loc_size_x_var, mk_add_int_expr(mk_mult_int_expr([blocks_per_pe_x_var,block_size_var]),leftovers_x_var))
  loc_size_y_expr = Expr(:(=), loc_size_y_var, mk_add_int_expr(mk_mult_int_expr([blocks_per_pe_y_var,block_size_var]),leftovers_y_var))
  state.arrs_dist_info[arr].local_sizes[end-1] = loc_size_x_var
  state.arrs_dist_info[arr].local_sizes[end] = loc_size_y_var

  # set new divided allocation size
  rhs.args[end-1] = loc_size_y_var
  rhs.args[end-3] = loc_size_x_var

  # TODO: handle extra blocks and partial block
  res = [block_size_expr; start_x_expr; start_y_expr; stride_x_expr; stride_y_expr;
            nb_x_div_expr; nb_y_div_expr; bppx_div_expr; bppy_div_expr;
            extra_block_call_x; extra_block_call_y;
            leftovers_x_expr; leftovers_y_expr;
            loc_size_x_expr; loc_size_y_expr; node]
  #debug_size_print = :(println("size ",$darr_count_var))
  #push!(res,debug_size_print)
  return res
end

function from_assignment_reshape(node::Expr, state::DistPassState, arr::LHSVar, rhs::Expr)
  if isONE_D(arr,state)
      @dprintln(3,"DistPass reshape array: ", arr)
      dim_sizes = state.arrs_dist_info[arr].dim_sizes
      # generate array division
      # simple 1D partitioning of last dimension, more general partitioning needed
      # match common big data matrix reperesentation
      arr_tot_size = dim_sizes[end]

      arr_id = getDistNewID(state)
      state.arrs_dist_info[arr].arr_id = arr_id
      darr_start_var_name = Symbol("__hpat_dist_arr_start_"*string(arr_id))
      darr_div_var_name = Symbol("__hpat_dist_arr_div_"*string(arr_id))
      darr_count_var_name = Symbol("__hpat_dist_arr_count_"*string(arr_id))

      darr_start_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          darr_start_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
      darr_div_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          darr_div_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
      darr_count_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          darr_count_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

      state.arrs_dist_info[arr].starts[end] = darr_start_var
      state.arrs_dist_info[arr].counts[end] = darr_count_var

      darr_div_expr = Expr(:(=), darr_div_var, mk_div_int_expr(arr_tot_size, state.dist_vars[:num_pes]))
      # zero-based index to match C interface of HDF5
      darr_start_expr = Expr(:(=),darr_start_var, mk_mult_int_expr(
          [state.dist_vars[:node_id], darr_div_var]))
      #darr_count_expr = :($darr_count_var = __hpat_node_id==__hpat_num_pes-1 ? $arr_tot_size-__hpat_node_id*$darr_div_var : $darr_div_var)
      darr_count_expr = Expr(:(=), darr_count_var, mk_call(
          GlobalRef(HPAT.API,:__hpat_get_node_portion),
              [arr_tot_size, darr_div_var, state.dist_vars[:num_pes], state.dist_vars[:node_id]]))

      # create a new tuple for reshape
      tup_call = Expr(:call, GlobalRef(Core,:tuple), dim_sizes[1:end-1]... , darr_count_var)
      reshape_tup_var_name = Symbol("__hpat_dist_tup_var_"*string(arr_id))
      tup_typ = CompilerTools.LambdaHandling.getType(rhs.args[3], state.LambdaVarInfo)
      reshape_tup_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
          reshape_tup_var_name, tup_typ, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
      tup_expr = Expr(:(=),reshape_tup_var,tup_call)
      rhs.args[3] = reshape_tup_var
      res = [darr_div_expr; darr_start_expr; darr_count_expr; tup_expr; node]
      #debug_size_print = :(println("size ",$darr_count_var))
      #push!(res,debug_size_print)
      return res
  end
  return [node]
end

function from_assignment_vcat(node::Expr, state::DistPassState, arr::LHSVar, rhs::Expr)
    @dprintln(3,"DistPass vcat array: ", arr)
    dim_sizes = state.arrs_dist_info[arr].dim_sizes
    # generate array division
    # simple 1D partitioning of last dimension, more general partitioning needed
    # match common big data matrix reperesentation
    arr_tot_size = dim_sizes[end]

    arr_id = getDistNewID(state)
    state.arrs_dist_info[arr].arr_id = arr_id
    darr_start_var_name = Symbol("__hpat_dist_arr_start_"*string(arr_id))
    darr_div_var_name = Symbol("__hpat_dist_arr_div_"*string(arr_id))
    darr_count_var_name = Symbol("__hpat_dist_arr_count_"*string(arr_id))

    darr_start_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        darr_start_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
    darr_div_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        darr_div_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
    darr_count_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        darr_count_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

    state.arrs_dist_info[arr].starts[end] = darr_start_var
    state.arrs_dist_info[arr].counts[end] = darr_count_var

    darr_div_expr = Expr(:(=), darr_div_var, mk_div_int_expr(arr_tot_size, state.dist_vars[:num_pes]))
    # zero-based index to match C interface of HDF5
    darr_start_expr = Expr(:(=),darr_start_var, mk_mult_int_expr(
        [state.dist_vars[:node_id], darr_div_var]))
    #darr_count_expr = :($darr_count_var = __hpat_node_id==__hpat_num_pes-1 ? $arr_tot_size-__hpat_node_id*$darr_div_var : $darr_div_var)
    darr_count_expr = Expr(:(=), darr_count_var, mk_call(
        GlobalRef(HPAT.API,:__hpat_get_node_portion),
            [arr_tot_size, darr_div_var, state.dist_vars[:num_pes], state.dist_vars[:node_id]]))
    res = [darr_div_expr; darr_start_expr; darr_count_expr; node]
    return res
end

function from_assignment_gemm(node::Expr, state::DistPassState, lhs::LHSVar, rhs::Expr)
  arr1 = toLHSVar(rhs.args[5])
  t1 = (rhs.args[3]=='T')
  arr2 = toLHSVar(rhs.args[6])
  t2 = (rhs.args[4]=='T')

  # result is sequential but with reduction if both inputs are partitioned and second one is transposed
  # e.g. labels*points'
  if !isSEQ(arr1,state) && !isSEQ(arr2,state) && t2 && !t1 && isSEQ(lhs,state)
    @dprintln(3,"DistPass translating gemm reduce: ", node)
    # rhs.args[1] = :__hpat_gemm_reduce
    # allocate temporary array for local gemm values
    alloc_args = Array(Any,2)
    out_typ = CompilerTools.LambdaHandling.getType(lhs, state.LambdaVarInfo)
    alloc_args[1] = eltype(out_typ)
    out_dim_sizes = state.arrs_dist_info[lhs].dim_sizes
    alloc_args[2] = out_dim_sizes
    alloc_call = ParallelIR.from_alloc(alloc_args)
    reduce_num = getDistNewID(state)
    reduce_var_name = Symbol("__hpat_gemm_reduce_"*string(reduce_num))
    reduce_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        reduce_var_name, out_typ, ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
    reduce_var_init = Expr(:(=), reduce_var, Expr(:call,alloc_call...))
    # TODO: deallocate temporary array
    # reduce_var_dealloc = Expr(:call, TopNode(:ccall), QuoteNode(:jl_dealloc_array), reduce_var)

    # get reduction size
    reduce_size_var_name = Symbol("__hpat_gemm_reduce_size_"*string(reduce_num))
    reduce_size_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        reduce_size_var_name, Int, ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
    size_expr = Expr(:(=), reduce_size_var, mk_mult_int_expr(out_dim_sizes))

    # add allreduce call
    allreduceCall = Expr(:call, GlobalRef(HPAT.API,:hpat_dist_allreduce), reduce_var, GlobalRef(Base,:add_float), rhs.args[2], reduce_size_var)
    res_copy = Expr(:(=), lhs, rhs.args[2])
    # replace gemm output with local var
    #node.args[1] = reduce_var
    rhs.args[2] = reduce_var

    return [reduce_var_init; node; size_expr; allreduceCall; res_copy]
    # first input is sequential but output is parallel if the second input is partitioned but not transposed
    # e.g. w*points
  elseif isSEQ(arr1,state) && !isSEQ(arr2,state) && !t2 && !isSEQ(lhs,state)
    @dprintln(3,"DistPass arr info gemm first input is sequential: ", arr1)
    #rhs.args[1] = :__hpat_gemm_broadcast
    @dprintln(3,"DistPass translating gemm broadcast: ", node)
    # otherwise, no known pattern found
  # if any array is 2D, then all arrays should be 2d
  elseif any([isTWO_D(arr1,state), isTWO_D(arr2,state), isTWO_D(lhs,state)])
    @assert all([isTWO_D(arr1,state), isTWO_D(arr2,state), isTWO_D(lhs,state)]) "invalid 2d gemm"
    @dprintln(3,"DistPass translating gemm 2d: ", node)
    rhs.args[1] = GlobalRef(HPAT.API,:__hpat_gemm_2d)

    # result in rhs has partitioning set in alloc, assign to lhs
    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

    push!(rhs.args, state.arrs_dist_info[lhs].dim_sizes[end-1], state.arrs_dist_info[lhs].dim_sizes[end],
    state.arrs_dist_info[lhs].blocks[end-1], state.arrs_dist_info[lhs].blocks[end],
    state.arrs_dist_info[lhs].local_sizes[end-1], state.arrs_dist_info[lhs].local_sizes[end])

    push!(rhs.args, state.arrs_dist_info[arr1].dim_sizes[end-1], state.arrs_dist_info[arr1].dim_sizes[end],
    state.arrs_dist_info[arr1].blocks[end-1], state.arrs_dist_info[arr1].blocks[end],
    state.arrs_dist_info[arr1].local_sizes[end-1], state.arrs_dist_info[arr1].local_sizes[end])

    push!(rhs.args, state.arrs_dist_info[arr2].dim_sizes[end-1], state.arrs_dist_info[arr2].dim_sizes[end],
    state.arrs_dist_info[arr2].blocks[end-1], state.arrs_dist_info[arr2].blocks[end],
    state.arrs_dist_info[arr2].local_sizes[end-1], state.arrs_dist_info[arr2].local_sizes[end])
  end

  return [node]
end

function from_assignment_gemv(node::Expr, state::DistPassState, lhs::LHSVar, rhs::Expr)

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

  # result is sequential but with reduction if both inputs are partitioned and matrix is not transposed (X*y)
  # result is sequential but with reduction if matrix partitioned (X'*y)
  if !isSEQ(arr1,state) && isSEQ(lhs,state) && !isSEQ(arr2,state) && !t1
    @dprintln(3,"DistPass translating gemv reduce: ", node)
    # rhs.args[1] = :__hpat_gemm_reduce
    # allocate temporary array for local gemm values
    alloc_args = Array(Any,2)
    out_typ = CompilerTools.LambdaHandling.getType(lhs, state.LambdaVarInfo)
    alloc_args[1] = eltype(out_typ)
    out_dim_sizes = state.arrs_dist_info[lhs].dim_sizes
    alloc_args[2] = out_dim_sizes
    alloc_call = ParallelIR.from_alloc(alloc_args)
    reduce_num = getDistNewID(state)
    reduce_var_name = Symbol("__hpat_gemv_reduce_"*string(reduce_num))
    reduce_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        reduce_var_name, out_typ, ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
    reduce_var_init = Expr(:(=), reduce_var, Expr(:call,alloc_call...))
    # TODO: deallocate temporary array
    # reduce_var_dealloc = Expr(:call, TopNode(:ccall), QuoteNode(:jl_dealloc_array), reduce_var)

    # get reduction size
    reduce_size_var_name = Symbol("__hpat_gemv_reduce_size_"*string(reduce_num))
    reduce_size_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        reduce_size_var_name, Int, ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
    size_expr = Expr(:(=), reduce_size_var, mk_mult_int_expr(out_dim_sizes))

    # add allreduce call
    allreduceCall = Expr(:call, GlobalRef(HPAT.API,:hpat_dist_allreduce), reduce_var, GlobalRef(Base,:add_float), rhs.args[2], reduce_size_var)
    res_copy = Expr(:(=), lhs, rhs.args[2])
    # replace gemm output with local var
    #node.args[1] = reduce_var
    rhs.args[2] = reduce_var

    return [reduce_var_init; node; size_expr; allreduceCall; res_copy]
    # result and vector are sequential if matrix is parallel and transposed
  end
  return [node]
end

function from_parfor(node::Expr, state)
    @assert node.head==:parfor "DistributedPass invalid parfor head"

    parfor = node.args[1]

    parfor.preParFor = from_nested_body(parfor.preParFor, state)
    parfor.hoisted = from_nested_body(parfor.hoisted, state)
    parfor.body = from_nested_body(parfor.body, state)
    parfor.postParFor = from_nested_body(parfor.postParFor, state)

    parfor_rws = CompilerTools.ReadWriteSet.from_exprs(parfor.body, ParallelAccelerator.ParallelIR.pir_rws_cb, state.LambdaVarInfo, state.LambdaVarInfo)

    if parfor.unique_id in keys(state.parfor_stencils)
        return from_parfor_stencil_1d(node, state, parfor, state.parfor_stencils[parfor.unique_id])
    elseif state.parfor_partitioning[parfor.unique_id]==ONE_D
      return from_parfor_1d(node, state, parfor)
    elseif state.parfor_partitioning[parfor.unique_id]==TWO_D
      return from_parfor_2d(node, state, parfor)
    else
        # broadcast results of sequential parfors if rand() is used
        has_rand = false
        for stmt in parfor.body
            if isa(stmt,Expr) && stmt.head==:(=) && isa(stmt.args[2],Expr) && stmt.args[2].head==:call &&
                 (isBaseFunc(stmt.args[2].args[1],:rand!) || isBaseFunc(stmt.args[2].args[1],:rand) ||
                   isBaseFunc(stmt.args[2].args[1],:randn!) || isBaseFunc(stmt.args[2].args[1],:randn) )
                has_rand = true
                break
            end
        end
        if has_rand
            # only rank 0 executes rand(), then broadcasts results
            writeArrs = collect(keys(parfor_rws.writeSet.arrays))
            @assert length(writeArrs)==1 "Only one parfor output supported now"
            write_arr = toLHSVar(writeArrs[1])
            # generate new label
            label = next_label(state)
            label_node = LabelNode(label)
            goto_node = Expr(:gotoifnot, mk_call(GlobalRef(Base,:(===)),
                [state.dist_vars[:node_id], 0]), label)
            # get broadcast size
            bcast_size_var_name = Symbol("__hpat_bcast_size_"*string(label))
            bcast_size_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
                bcast_size_var_name, Int, ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
            size_expr = Expr(:(=), bcast_size_var, mk_mult_int_expr(state.arrs_dist_info[write_arr].dim_sizes))
            bcast_expr = Expr(:call,GlobalRef(HPAT.API,:__hpat_dist_broadcast), write_arr, bcast_size_var)
            push!(parfor.hoisted, goto_node)
            p_out = pop!(parfor.postParFor)
            push!(parfor.postParFor, label_node)
            push!(parfor.postParFor, size_expr)
            push!(parfor.postParFor, bcast_expr)
            push!(parfor.postParFor, p_out)

            @dprintln(3,"DistPass rand() in sequential parfor ", parfor)
            #return [goto_node; node; label_node; size_expr; bcast_expr]
            return [node]
        end
    end
    return [node]
end

function from_parfor_stencil_1d(node::Expr, state, parfor, stencil_inds::Vector{Int})
    @dprintln(3,"DistPass translating stencil 1d parfor: ", parfor.unique_id)

    in_arr = state.parfor_arrays[parfor.unique_id][1]
    out_arr = state.parfor_arrays[parfor.unique_id][end]
    dprintln(3, "stencil input arr: ", in_arr," output_arr: ", out_arr, " parfor node: ", node)
    data_typ = eltype(CompilerTools.LambdaHandling.getType(out_arr, state.LambdaVarInfo))

    # indices on left and right side of stencil e.g. -2, -1 , 1, 2
    left_inds = filter(x->x<0, stencil_inds)
    right_inds = filter(x->x>0, stencil_inds)
    @assert left_inds==[-1] && right_inds==[1] "only simple 3-point stencil supported for now"

    # no change to the loopnest start since index variable is not used in actual stencil computation
    # upper bound should be replaced with local size of array
    loopnest = parfor.loopNests[1]
    dprintln(3, "stencil loopnest before change: ", loopnest)
    loopnest.upper.args[3].args[2] = Expr(:call, GlobalRef(Base, :arraysize), in_arr,1)
    dprintln(3, "stencil loopnest after change: ", loopnest)

    recv_left_tmp = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        Symbol("recv_left_tmp"), data_typ, ISASSIGNED, state.LambdaVarInfo))
    recv_right_tmp = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        Symbol("recv_right_tmp"), data_typ, ISASSIGNED, state.LambdaVarInfo))

    # create MPI request objects for Isend/Irecv
    init_stencil_reqs = mk_call(GlobalRef(HPAT.API,:__hpat_init_stencil_reqs),[])

    # send/recv on left
    # node 0 should just set output to input
    label_left = next_label(state)
    label_left_node = LabelNode(label_left)
    # goto label if not node_id==0
    goto_left_node = Expr(:gotoifnot, mk_call(GlobalRef(Base,:(===)),
        [state.dist_vars[:node_id], 0]), label_left)
    # out_arr[1] = in_arr[1] on first node
    assign_leftmost = Expr(:call, GlobalRef(Base,:unsafe_arrayset), out_arr,
        Expr(:call, GlobalRef(Base,:unsafe_arrayref), in_arr, 1), 1)
    send_left = Expr(:call,GlobalRef(HPAT.API,:__hpat_send_left), in_arr)
    recv_left = Expr(:call,GlobalRef(HPAT.API,:__hpat_recv_left), recv_left_tmp)
    label_after_left = next_label(state)
    label_after_left_node = LabelNode(label_after_left)
    goto_after_left_node = GotoNode(label_after_left)

    # send/recv on right
    # node pes-1 should just set output to input
    label_right = next_label(state)
    label_right_node = LabelNode(label_right)
    # goto label if not node_id==num_pes-1
    goto_right_node = Expr(:gotoifnot, mk_call(GlobalRef(Base,:(===)),
        [state.dist_vars[:node_id], mk_call(GlobalRef(Core.Intrinsics,:sub_int), [state.dist_vars[:num_pes],1])]), label_right)
    # out_arr[n] = in_arr[n] on last node
    assign_rightmost = Expr(:call, GlobalRef(Base,:unsafe_arrayset), out_arr,
        Expr(:call, GlobalRef(Base,:unsafe_arrayref), in_arr,
        mk_call(GlobalRef(Base, :arraysize), [in_arr,1])),
        mk_call(GlobalRef(Base, :arraysize), [out_arr,1]) )
    send_right = Expr(:call, GlobalRef(HPAT.API,:__hpat_send_right), in_arr)
    recv_right = Expr(:call, GlobalRef(HPAT.API,:__hpat_recv_right), recv_right_tmp)
    label_after_right = next_label(state)
    label_after_right_node = LabelNode(label_after_right)
    goto_after_right_node = GotoNode(label_after_right)

    # node 0 doesn't wait left
    label_wleft = next_label(state)
    label_wleft_node = LabelNode(label_wleft)
    # goto label if not node_id!=0
    goto_wleft_node = Expr(:gotoifnot, mk_call(GlobalRef(Core.Intrinsics,:ne_int),
        [state.dist_vars[:node_id], 0]), label_wleft)
    wait_left = Expr(:call, GlobalRef(HPAT.API,:__hpat_wait_left))

    # node pes-1 doesn't wait right
    label_wright = next_label(state)
    label_wright_node = LabelNode(label_wright)
    # goto label if not node_id==num_pes-1
    goto_wright_node = Expr(:gotoifnot, mk_call(GlobalRef(Core.Intrinsics,:ne_int),
        [state.dist_vars[:node_id], mk_call(GlobalRef(Core.Intrinsics,:sub_int), [state.dist_vars[:num_pes],1])]), label_wright)
    wait_right = Expr(:call, GlobalRef(HPAT.API,:__hpat_wait_right))

    border_left_body = get_stencil_border_compute(parfor.body, recv_left_tmp, left_inds[1], 1)
    border_right_body = get_stencil_border_compute(parfor.body, recv_right_tmp, right_inds[1],
        mk_call(GlobalRef(Base, :arraysize), [in_arr,1]))

    # if ! node==0 goto 1, out[1]=in[1] goto 2, 1: send/rev 2: ...
    # if ! node==pes-1 goto 1, out[n]=in[n] goto 2, 1: send/rev 2: ...
    res = Any[init_stencil_reqs, goto_left_node, assign_leftmost, goto_after_left_node, label_left_node, send_left, recv_left, label_after_left_node,
        goto_right_node, assign_rightmost, goto_after_right_node, label_right_node, send_right, recv_right, label_after_right_node, node,
        goto_wleft_node, wait_left, border_left_body..., label_wleft_node, goto_wright_node, wait_right, border_right_body..., label_wright_node]
    dprintln(3, "stencil returns: ", res)
  return res
end

function get_stencil_border_compute(body, tmp_var, stencil_ind, border_ind)
    new_body = deepcopy(body)
    for node in new_body
        if isAssignmentNode(node) && isCall(node.args[2]) &&
              node.args[2].args[1]==GlobalRef(Base,:unsafe_arrayref)
            rhs = node.args[2]
            index = rhs.args[3]
            # if access is stencil
            if isa(index,Expr) && index.head==:call && index.args[1]==GlobalRef(Base,:box) && isa(index.args[3],Expr) &&
                  index.args[3].head==:call && index.args[3].args[1]==GlobalRef(Core.Intrinsics,:add_int) &&
                  isa(index.args[3].args[3],Int)
                # if this is the border we are replacing, replace arrayref call
                if index.args[3].args[3]==stencil_ind
                    node.args[2] = tmp_var
                else
                    # replace index variable with 1
                    index.args[3].args[2] = border_ind
                end
            else
                @assert isa(index, LHSVar) "unexpected index $index"
                # replace index variable with 1
                rhs.args[3] = border_ind
            end
        elseif isCall(node) && node.args[1]==GlobalRef(Base,:unsafe_arrayset)
            node.args[4] = border_ind
        end
    end
    return new_body
end

function from_parfor_1d(node::Expr, state, parfor)
  @dprintln(3,"DistPass translating 1d parfor: ", parfor.unique_id)
  @dprintln(3, parfor)

  # TODO: assuming 1st loop nest is the last dimension
  loopnest = parfor.loopNests[1]
  parfor.force_simd = true
  # TODO: build a constant table and check the loop variables at this stage
  # @assert loopnest.lower==1 && loopnest.step==1 "DistPass only simple PIR loops supported now"

  loop_start_var_name = Symbol("__hpat_loop_start_"*string(parfor.unique_id))
  loop_end_var_name = Symbol("__hpat_loop_end_"*string(parfor.unique_id))
  loop_div_var_name = Symbol("__hpat_loop_div_"*string(parfor.unique_id))

  loop_start_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      loop_start_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  loop_end_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      loop_end_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  loop_div_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      loop_div_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

  #first_arr = state.parfor_info[parfor.unique_id][1];
  #@dprintln(3,"DistPass parfor first array ", first_arr)
  #global_size = state.arrs_dist_info[first_arr].dim_sizes[1]

  # some parfors have no arrays
  global_size = loopnest.upper

  loop_div_expr = Expr(:(=),loop_div_var, mk_div_int_expr(
      global_size, state.dist_vars[:num_pes]))
  loop_start_expr = Expr(:(=), loop_start_var, mk_add_int_expr(
      mk_mult_int_expr([state.dist_vars[:node_id],loop_div_var]),1))
  #loop_end_expr = :($loop_end_var = __hpat_node_id==__hpat_num_pes-1 ?$(global_size):(__hpat_node_id+1)*$loop_div_var)
  loop_end_expr = Expr(:(=), loop_end_var, mk_call(
      GlobalRef(HPAT.API,:__hpat_get_node_end),
          [global_size, loop_div_var, state.dist_vars[:num_pes], state.dist_vars[:node_id]]))

  loopnest.lower = loop_start_var
  loopnest.upper = loop_end_var

  for stmt in parfor.body
      #adjust_arrayrefs(stmt, loop_start_var)
      ParallelIR.AstWalk(stmt, adjust_arrayrefs, loopnest)
  end
  res = [loop_div_expr; loop_start_expr; loop_end_expr; node]

  dist_reductions = gen_dist_reductions(parfor.reductions, state)
  append!(res, dist_reductions)

  #debug_start_print = :(println("parfor start", $loop_start_var))
  #debug_end_print = :(println("parfor end", $loop_end_var))
  #push!(res,debug_start_print)
  #push!(res,debug_end_print)

  #debug_div_print = :(println("parfor div ", $loop_div_var))
  #push!(res,debug_div_print)
  #debug_pes_print = :(println("parfor pes ", __hpat_num_pes))
  #push!(res,debug_pes_print)
  #debug_rank_print = :(println("parfor rank ", __hpat_node_id))
  #push!(res,debug_rank_print)
  return res
end

function from_parfor_2d(node::Expr, state, parfor)
  @dprintln(3,"DistPass translating 2d parfor: ", parfor.unique_id)

  # TODO: assuming parfor has an array
  # TODO: using parfor index (e.g. comprehension) is not supported yet
  first_arr = state.parfor_arrays[parfor.unique_id][1]
  # assuming 1st loop nest is the last dimension
  parfor.loopNests[1].upper = state.arrs_dist_info[first_arr].local_sizes[end]
  parfor.loopNests[2].upper = state.arrs_dist_info[first_arr].local_sizes[end-1]

  res = [node]

  dist_reductions = gen_dist_reductions(parfor.reductions, state)
  append!(res, dist_reductions)

  return res
end

function next_label(state)
    state.max_label = state.max_label + 1
    return state.max_label
end

function from_call(node::Expr, state)
    @assert node.head==:call "Invalid call node"
    @dprintln(2,"DistPass from_call ", node)

    func = node.args[1]
    if (func==GlobalRef(HPAT.API,:__hpat_data_source_HDF5_read) || func==GlobalRef(HPAT.API,:__hpat_data_source_TXT_read))
        arr = toLHSVar(node.args[3])
        @dprintln(3,"DistPass data source for array: ", arr)
        if isONE_D(arr,state)
          # 1D read, add start and count indices of last dimension
          push!(node.args, state.arrs_dist_info[arr].starts[end], state.arrs_dist_info[arr].counts[end])
        elseif isTWO_D(arr,state)
          push!(node.args, state.arrs_dist_info[arr].starts[end-1], state.arrs_dist_info[arr].starts[end],
           state.arrs_dist_info[arr].strides[end-1], state.arrs_dist_info[arr].strides[end],
           state.arrs_dist_info[arr].counts[end-1], state.arrs_dist_info[arr].counts[end],
           state.arrs_dist_info[arr].blocks[end-1], state.arrs_dist_info[arr].blocks[end],
           state.arrs_dist_info[arr].local_sizes[end-1], state.arrs_dist_info[arr].local_sizes[end],
           state.arrs_dist_info[arr].leftovers[end-1], state.arrs_dist_info[arr].leftovers[end])
        end
        return [node]
    elseif func==GlobalRef(HPAT.API,:__hpat_data_sink_HDF5_write)
        arr = toLHSVar(node.args[4])
        @dprintln(3,"DistPass data sink HDF5 for array: ", arr)
        if isONE_D(arr,state)
          # 1D write, add start and count indices of last dimension, total sizes
          push!(node.args, state.arrs_dist_info[arr].starts[end], state.arrs_dist_info[arr].counts[end],
                    state.arrs_dist_info[arr].dim_sizes)
        elseif isTWO_D(arr,state)
          push!(node.args, state.arrs_dist_info[arr].starts[end-1], state.arrs_dist_info[arr].starts[end],
           state.arrs_dist_info[arr].strides[end-1], state.arrs_dist_info[arr].strides[end],
           state.arrs_dist_info[arr].counts[end-1], state.arrs_dist_info[arr].counts[end],
           state.arrs_dist_info[arr].blocks[end-1], state.arrs_dist_info[arr].blocks[end],
           state.arrs_dist_info[arr].local_sizes[end-1], state.arrs_dist_info[arr].local_sizes[end],
           state.arrs_dist_info[arr].leftovers[end-1], state.arrs_dist_info[arr].leftovers[end],
           state.arrs_dist_info[arr].dim_sizes)
        end

        return [node]
    elseif func==GlobalRef(HPAT.API,:__hpat_data_sink_TXT_write)
        arr = toLHSVar(node.args[3])
        @dprintln(3,"DistPass data sink TXT for array: ", arr)
        if isONE_D(arr,state)
          # 1D write, add start and count indices of last dimension, total sizes
          push!(node.args, state.arrs_dist_info[arr].starts[end], state.arrs_dist_info[arr].counts[end],
                    state.arrs_dist_info[arr].dim_sizes)
        end
        return [node]
    elseif func==GlobalRef(HPAT.API,:Kmeans) && (isONE_D(toLHSVar(node.args[3]), state) || isONE_D_VAR(toLHSVar(node.args[3]), state))
        # 1st call argument (node.args[2]) is coeffs (lhs of original source code)
        arr = toLHSVar(node.args[3])
        @dprintln(3,"DistPass kmeans call for array: ", arr)
        node.args[1] = GlobalRef(HPAT.API,:Kmeans_dist)

        extra_daal_includes = """ #include "daal.h"
        using namespace daal;
        using namespace daal::data_management;
        using namespace daal::algorithms;
        """
        HPAT.addHpatInclude(extra_daal_includes,"-daal","-daal")
        # rebalance array if necessary
        # table operations like filter produce irregular chunk sizes on different processors
        rebalance_out = gen_rebalance_array(arr, state)

        push!(node.args, state.arrs_dist_info[arr].starts[end], state.arrs_dist_info[arr].counts[end],
                state.arrs_dist_info[arr].dim_sizes[1], state.arrs_dist_info[arr].dim_sizes[end])
        return [rebalance_out; node]
    elseif func==GlobalRef(HPAT.API,:LinearRegression) || func==GlobalRef(HPAT.API,:NaiveBayes)
        # 1st call argument (node.args[2]) is coeffs (lhs of original source code)
        arr1 = toLHSVar(node.args[3])
        arr2 = toLHSVar(node.args[4])
        @dprintln(3,"DistPass LinearRegression/NaiveBayes call for arrays: ", arr1," ", arr2)
        # both arrays can be 1D or variable length 1D
        if (isONE_D(arr1,state) || isONE_D_VAR(arr1,state)) &&
            (isONE_D(arr2,state) || isONE_D_VAR(arr2,state))

            node.args[1] = GlobalRef(HPAT.API, Symbol("$(func.name)_dist"))

            extra_daal_includes = """ #include "daal.h"
            using namespace daal;
            using namespace daal::data_management;
            using namespace daal::algorithms;
            """
            HPAT.addHpatInclude(extra_daal_includes,"-daal", "-daal")
            # rebalance array if necessary
            # table operations like filter produce irregular chunk sizes on different processors
            rebalance_out1 = gen_rebalance_array(arr1, state)
            rebalance_out2 = gen_rebalance_array(arr2, state)

            # response array could be 1D or 2D array
            # number of responses for each feature is 1 if 1D array to use same 2D DAAL table interface
            response_col_size = state.arrs_dist_info[arr2].dim_sizes[1]
            if length(state.arrs_dist_info[arr2].dim_sizes)==1
                response_col_size = 1
            end

            push!(node.args, state.arrs_dist_info[arr1].starts[end], state.arrs_dist_info[arr1].counts[end],
                    state.arrs_dist_info[arr1].dim_sizes[1], state.arrs_dist_info[arr1].dim_sizes[end])
            push!(node.args, state.arrs_dist_info[arr2].starts[end], state.arrs_dist_info[arr2].counts[end],
                    response_col_size, state.arrs_dist_info[arr2].dim_sizes[end])
            return [rebalance_out1; rebalance_out2; node]
        end
    elseif isBaseFunc(func, :arraysize) && (isONE_D(toLHSVar(node.args[2]),state) || isTWO_D(toLHSVar(node.args[2]),state))
        arr = toLHSVar(node.args[2])
        # don't replace if it is variable length (ONE_D_VAR)
        # can be 2D, like hcat-transpose of variable length arrays
        @dprintln(3,"found arraysize on dist array: ",node," ",arr)

        # replace last dimension size queries since it is partitioned
        #if node.args[3]==length(state.arrs_dist_info[arr].dim_sizes)
        #    return [state.arrs_dist_info[arr].dim_sizes[end]]
        #end
        # Parfor array length does not get zero value
        @assert state.arrs_dist_info[arr].dim_sizes[node.args[3]]!=0 "$arr array size could not be zero"
        return [state.arrs_dist_info[arr].dim_sizes[node.args[3]]]
    elseif isBaseFunc(func,:arraylen) && (isONE_D(toLHSVar(node.args[2]), state) || isTWO_D(toLHSVar(node.args[2]),state))
        arr = toLHSVar(node.args[2])
        # don't replace if it is variable length (ONE_D_VAR)
        # can be 2D, like hcat-transpose of variable length arrays

        #len = parse(foldl((a,b)->"$a*$b", "1",state.arrs_dist_info[arr].dim_sizes))
        len = mk_mult_int_expr(state.arrs_dist_info[arr].dim_sizes)
        @dprintln(3,"found arraylen on dist array: ",node," ",arr," len: ",len)
        @dprintln(3,"found arraylen on dist array: ",node," ",arr)
        return [len]
    elseif isBaseFunc(func,:cumsum!) && isONE_D(toLHSVar(node.args[2]), state)
        new_func = GlobalRef(HPAT.API, :dist_cumsum!)
        node.args[1] = new_func
    end
    return [node]
end

function getDistNewID(state)
    state.uniqueId+=1
    return state.uniqueId
end

function adjust_arrayrefs(stmt::Expr, loopnest, top_level_number, is_top_level, read)

    if isCall(stmt)
        topCall = stmt.args[1]
        #ref_args = stmt.args[2:end]
        if isBaseFunc(topCall,:unsafe_arrayref) || isBaseFunc(topCall,:unsafe_arrayset) ||
            topCall==GlobalRef(ParallelAccelerator.API,:getindex) || topCall==GlobalRef(ParallelAccelerator.API,:setindex)
            # TODO: simply divide the last dimension, more general partitioning needed
            index_arg = toLHSVar(stmt.args[end])
            if isa(index_arg, LHSVar) && index_arg==toLHSVar(loopnest.indexVariable)
                stmt.args[end] = mk_add_int_expr(mk_sub_int_expr(toLHSVar(index_arg),loopnest.lower),1)
                return stmt
            end
        end
    end
    CompilerTools.AstWalker.ASTWALK_RECURSE
end

function adjust_arrayrefs(stmt::Any, loopnest,top_level_number, is_top_level, read)
    CompilerTools.AstWalker.ASTWALK_RECURSE
end

function gen_dist_reductions(reductions::Array{PIRReduction,1}, state)
    @dprintln(3,"DistPass generating reductions: ", reductions)
    res = Any[]
    for reduce in reductions
        reduce_var_name = Symbol("__hpat_reduce_"*string(getDistNewID(state)))
        typ = CompilerTools.LambdaHandling.getType(reduce.reductionVar, state.LambdaVarInfo)
        reduce_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
            reduce_var_name, typ, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))

        reduce_var_init = Expr(:(=), reduce_var, 0)
        size = 1
        if isArrayType(typ)
            alloc_args = Array(Any,2)
            alloc_args[1] = eltype(typ)
            out_dim_sizes = state.arrs_dist_info[reduce.reductionVar].dim_sizes
            alloc_args[2] = out_dim_sizes
            alloc_call = ParallelIR.from_alloc(alloc_args)
            reduce_var_init = Expr(:(=), reduce_var, Expr(:call,alloc_call...))
            size = mk_mult_int_expr(out_dim_sizes)
        end
        reduceCall = Expr(:call,GlobalRef(HPAT.API,:hpat_dist_allreduce),
            reduce.reductionVar,reduce.reductionFunc, reduce_var, size)
        rootCopy = Expr(:(=), reduce.reductionVar, reduce_var)
        append!(res,[reduce_var_init; reduceCall; rootCopy])
    end
    return res
end

"""
    generate code to rebalance 1D or 2D arrays resulting from table operations
    like filter and join.
"""
function gen_rebalance_array(arr::LHSVar, state)
  # no need to rebalance if not variable chunk length
  if !isONE_D_VAR(arr, state)
    return Any[]
  end

  num_dims = length(state.arrs_dist_info[arr].dim_sizes)
  out = []
  arr_id = getDistNewID(state)
  state.arrs_dist_info[arr].arr_id = arr_id

  # get total array size with allreduce
  darr_size_var_name = Symbol("_glob_arr_size_"*string(arr_id))
  darr_loc_size_var_name = Symbol("_loc_arr_size_"*string(arr_id))
  darr_size_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      darr_size_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  darr_loc_size_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      darr_loc_size_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  size_var_init = Expr(:(=), darr_size_var, -1)
  # get size of last dimension, assume 1D partitioning
  loc_size_var_init = Expr(:(=), darr_loc_size_var, Expr(:call, GlobalRef(Base, :arraysize), arr, num_dims))
  push!(out, size_var_init)
  push!(out, loc_size_var_init)
  reduceCall = Expr(:call, GlobalRef(HPAT.API,:hpat_dist_allreduce),
    darr_loc_size_var, GlobalRef(Base, :add_int), darr_size_var, 1)
  push!(out, reduceCall)
  state.arrs_dist_info[arr].dim_sizes[end] = darr_size_var

  darr_start_var_name = Symbol("__hpat_dist_arr_start_"*string(arr_id))
  darr_div_var_name = Symbol("__hpat_dist_arr_div_"*string(arr_id))
  darr_count_var_name = Symbol("__hpat_dist_arr_count_"*string(arr_id))

  darr_start_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      darr_start_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  darr_div_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      darr_div_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  darr_count_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
      darr_count_var_name, Int, ISASSIGNEDONCE | ISASSIGNED | ISPRIVATEPARFORLOOP, state.LambdaVarInfo))
  state.arrs_dist_info[arr].starts[end] = darr_start_var
  state.arrs_dist_info[arr].counts[end] = darr_count_var

  darr_div_expr = Expr(:(=),darr_div_var, mk_div_int_expr(darr_size_var, state.dist_vars[:num_pes]))
  # zero-based index to match C interface of HDF5
  darr_start_expr = Expr(:(=), darr_start_var, mk_mult_int_expr(
      [state.dist_vars[:node_id], darr_div_var]))
  darr_count_expr = Expr(:(=), darr_count_var, mk_call(
      GlobalRef(HPAT.API,:__hpat_get_node_portion),
          [darr_size_var, darr_div_var, state.dist_vars[:num_pes], state.dist_vars[:node_id]]))
  push!(out, darr_div_expr, darr_start_expr, darr_count_expr)

  rebalance_call = mk_call(GlobalRef(HPAT.API,:__hpat_arr_rebalance),[arr, darr_count_var])
  push!(out, rebalance_call)

  return out
end

end # DistributedPass


text:
# Asset constructors

"""
    @asset asset_type code_symb name

Create a dimension and a reference unit for an asset.

The macros `@dimension` and `@refunit` are called with arguments derived
from `asset_type`, `code_symb` and `name`.

# Examples

```jldoctest
julia> @asset Cash USD "USD"
USD

julia> @asset Metal GOLD "GOLD"
GOLD

julia> @asset Stock DIS "DIS"
DIS
```
"""
macro asset(asset_type, code_symb, name)
    asset_str = titlecase(string(asset_type), strict=true)
    code_abbr = string(code_symb)

    code_abbr_bold = boldify(code_abbr)
    dimension = Symbol(code_abbr_bold)
    dim_abbr = asset_str * "{" * code_abbr_bold * "}"
    dim_name = Symbol(dim_abbr)

    esc(quote
        Unitful.@dimension($dimension, $dim_abbr, $dim_name)
        Unitful.@refunit($code_symb, $code_abbr, $name, $dimension, true)
    end)
end

ext: jl


text:
# ```@meta
# EditURL = "https://github.com/tisztamo/Catwalk.jl/blob/main/docs/src/usage.jl"
# ```
# # Usage

# Let's say you have a long-running calculation, organized into batches:

const NUM_BATCHES = 1000

function runbatches()
    for batchidx = 1:NUM_BATCHES
        hotloop()
        ## Log progress, etc.
    end
end

# The hot loop calls the type-unstable function `get_some_x()` and
# passes its result to a relatively cheap calculation `calc_with_x()`.

const NUM_ITERS_PER_BATCH = 1_000_000

function hotloop()
    for i = 1:NUM_ITERS_PER_BATCH
        x = get_some_x(i)
        calc_with_x(x)
    end
end

const xs = Any[1, 2.0, ComplexF64(3.0, 3.0)]
get_some_x(i) = xs[i % length(xs) + 1]

const result = Ref(ComplexF64(0.0, 0.0))

function calc_with_x(x)
    result[] += x
end

# As `get_some_x` is not type-stable, `calc_with_x` must be dynamically
# dispatched, which slows down the calculation.
#
# Sometimes it is not feasible to type-stabilize `get_some_x`.
# Catwalk.jl is here for those cases.
#
# You mark `hotloop`, the outer function
# with the `@jit` macro and provide the name of the dynamically
# dispatched function
# and the argument to operate on (the API will hopefully
# improve in the future). You also have to add an extra argument
# named `jitctx` to the jit-ed function:

using Catwalk

@jit calc_with_x x function hotloop_jit(jitctx)
    for i = 1:NUM_ITERS_PER_BATCH
        x = get_some_x(i)
        calc_with_x(x)
    end
end

# The Catwalk optimizer will provide you the `jitctx` context which you have to pass
# to the jit-ed function manually.
# Also, every batch needs a bit housekeeping to drive the Catwalk optimizer:

function runbatches_jit()
    jit = Catwalk.JIT() ## Also works inside a function (no eval used)
    for batch = 1:NUM_BATCHES
        Catwalk.step!(jit)
        hotloop_jit(Catwalk.ctx(jit)) # Dynamic dispatch is "pulled out" from the hot loop to this call
    end
end

# *Note that the body of the batch can be arbitrarily complex,
# not just a simple for loop like in this example. See [a more complex test case](https://github.com/tisztamo/Catwalk.jl/blob/main/test/scheduling.jl).*
#
# Yes, it is a bit complicated to integrate your code with Catwalk, but it may
# worth the effort:

result[] = ComplexF64(0, 0)
@time runbatches_jit()

## 4.608471 seconds (4.60 M allocations: 218.950 MiB, 0.56% gc time, 21.68% compilation time)

jit_result = result[]

result[] = ComplexF64(0, 0)
@time runbatches()

## 23.387341 seconds (1000.00 M allocations: 29.802 GiB, 7.71% gc time)

# And the results are the same:

jit_result == result[] || error("JIT must be a no-op!")

# Please note that the speedup depends on the portion of the runtime spent in dynamic dispatch,
# which is most likely smaller in your case than in this contrived example.
#
# You can find this example under [docs/src/usage.jl](https://github.com/tisztamo/Catwalk.jl/blob/main/docs/src/usage.jl) in the repo.


text:
abstract type AbstractSIMP <: TopOptAlgorithm end

# MMA wrapper
include("mma_optimizer.jl")

# Basic SIMP
include("basic_simp.jl")

## Continuation SIMP
include("continuation_schemes.jl")
include("continuation_simp.jl")


text:
using Sequent
using MLStyle
using HMRowUnification
using Test

@data Nat begin
    Z
    S(Nat)
end

rule_z = @sequent FOLD_Z begin
    _ ⊢ₛ Z => 0
end

rule_s = @sequent FOLD_S begin
    _ ⊢ₛ term => n
    _ ⊢ₛ S(term) => n + 1
end

@semantics [rule_z, rule_s]

@testset "Nat" begin

@test (nothing ⊢ₛ Z) == 0
@test (nothing ⊢ₛ S(S(Z))) == 2

end

@data ML begin
    LApp(ML, ML)
    LFun(Symbol, ML)
    LVar(Symbol)
    LVal(Int)
end

ImD = Base.ImmutableDict

struct Sigma{O}
    unbox :: ImD{Symbol, O}
end

Sigma{O}() where O = Sigma(ImD{Symbol, O}())

Base.getindex(s::Sigma{O}, x::Pair{Symbol, A}) where {O, A<:O} =
    Sigma(ImD(s.unbox, x))
Base.getindex(s::Sigma{O}, x::Symbol) where O =
    s.unbox[x]

(s::Sigma{O})(x::Pair{Symbol, A}) where {O, A<:O} = ImD(s, x)

r_app = @sequent APP begin
    (Γ, σ) ⊢ʷ a => ta
    (Γ, σ) ⊢ʷ f => Arrow(ta′, tr)
    true = Γ.unify(ta, ta′)
    (Γ, σ) ⊢ʷ LApp(f, a) => tr
end

r_var = @sequent VAR begin
    (Γ, σ) ⊢ʷ LVar(s) => σ[s]
end

r_fun = @sequent FUN begin
    a′ = Γ.new_tvar()
    (Γ, σ[a => a′]) ⊢ʷ r => tr
    (Γ, σ) ⊢ʷ LFun(a, r) => Arrow(a′, tr)
end

r_val = @sequent VAL begin
    _ ⊢ʷ LVal(_) => Nom(:int)
end

@semantics [r_app, r_fun, r_val, r_var]

@testset "ML" begin
    tctx = HMT[]
    Γ = mk_tcstate(tctx)
    σ = Sigma{HMT}()
    int_t = Nom(:int)
    σ = σ[:add => Arrow(int_t, Arrow(int_t, int_t))]
    term =
        LFun(:a,
            LApp(
                LApp(
                    LVar(:add),
                    LVar(:a)),
                LVar(:a)))
    infer_ty =  (Γ, σ) ⊢ʷ term
    @test Γ.prune(infer_ty) == Arrow(int_t, int_t)
end


text:
# Code lists begin on page 322 of the SUDS manual
const unsupported = Int16.((3, 4, 8, 9, 11, 12, 13, 15, 16, 17, 18, 19, 21, 22, 23, 24))
const disp_only = Int16.((25, 26, 27, 28, 29, 31))

const suds_codes = Dict{Int64, String}(
  1   => "STATIONID",
  2   => "STRUCTTAG",
  3   => "PC_TERMINATOR",
  4   => "EQUIPMENT",
  5   => "STATIONCOMP",
  6   => "MUXDATA",
  7   => "DESCRIPTRACE",
  8   => "LOCTRACE",
  9   => "PC_CALIBRATION",
  10  => "FEATURE",
  11  => "RESIDUAL",
  12  => "PC_EVENT",
  13  => "EVDESCR",
  14  => "ORIGIN",
  15  => "ERROR",
  16  => "FOCALMECH",
  17  => "MOMENT",
  18  => "VELMODEL",
  19  => "LAYERS",
  20  => "PC_COMMENT",
  21  => "PROFILE",
  22  => "SHOTGATHER",
  23  => "CALIB",
  24  => "COMPLEX",
  25  => "TRIGGERS",
  26  => "TRIGSETTING",
  27  => "EVENTSETTING",
  28  => "DETECTOR",
  29  => "ATODINFO",
  30  => "TIMECORRECTION",
  32  => "CHANSET",
  31  => "INSTRUMENT",
  33  => "CHANSETENTRY"
  )

const sensor_types = Dict{UInt8, String}(
  0x42 => "K"         , # 'B' = "bolometer"
  0x43 => "{unknown}" , # 'C' = "local clock"
  0x48 => "%"         , # 'H' = "humidity"
  0x50 => "Pa"        , # 'P' = "pressure sensor"
  0x52 => "{unknown}" , # 'R' = "rainfall"
  0x53 => "m/m"       , # 'S' = "linear strain meter"
  0x54 => "K"         , # 'T' = "temperature sensor"
  0x56 => "m3/m3"     , # 'V' = "volumetric strain meter"
  0x57 => "{unknown}" , # 'W' = "wind"
  0x61 => "m/s2"      , # 'a' = "accelerometer"
  0x63 => "m"         , # 'c' = "creep meter"
  0x64 => "m"         , # 'd' = "displacement sensor"
  0x67 => "cm/s2"     , # 'g' = "gravimeter"
  0x69 => "degrees"   , # 'i' = "tilt meter/inclinometer"
  0x6d => "T"         , # 'm' = "magnetic field"
  0x72 => "{radon}"   , # 'r' = "radon sensor"
  0x73 => "{unknown}" , # 's' = "satellite time code"
  0x74 => "m"         , # 't' = "tidal meter"
  0x76 => "m/s"       , # 'v' = "velocity seismometer"
  0x77 => "N.m2"      , # 'w' = "torsion"
  0x78 => "{unknown}" , # 'x' = "experimental"
  )

const data_types = Dict{Int64, String}(
  -33 => "PAD4", -32 => "PAD2", -31 => "PAD1 ", -30 => "INT3", -29 => "IDXPTR ",
  -28 => "YESNO", -27 => "UCHAR ", -26 => "UINT2", -25 => "UINT4 ",
  -24 => "MEMPTR", -23 => "CHRPTR", -22 => "GENPTR", -21 => "CODESTR ",
  -20 => "INT2TM", -19 => "LIST", -18 => "LATIT", -17 => "LONGIT ",
  -16 => "MS i TIME", -15 => "ST i TIME", -14 => "FLOAT8", -13 => "FLOAT4 ",
  -12 => "AUTHOR", -11 => "DOMAIN", -10 => "REFERS2", -9 => "LABEL",
  -8 => "CODE4", -7 => "CODE2 ", -6 => "CODE1", -5 => "FIXED", -4 => "INT4",
  -3 => "INT2 ", -2 => "STRING", -1 => "CHAR", 2 => "structtag",
  3 => "pc i terminator", 4 => "equipment", 5 => "stationcomp", 6 => "muxdata",
  7 => "descriptrace", 8 => "loctrace", 9 => "pc i calibration",
  10 => "feature", 11 => "residual", 12 => "pc i event", 13 => "evdescr",
  14 => "origin", 15 => "error", 16 => "focalmech", 17 => "moment",
  18 => "velmodel", 19 => "layers", 20 => "pc i comment", 23 => "calib",
  25 => "triggers", 26 => "trigsetting", 27 => "eventsetting", 28 => "detector",
  29 => "atodinfo", 30 => "timecorrection", 31 => "instrument", 32 => "chanset",
  33 => "chansetentry", 104 => "sig i path i cmp", 105 => "signal i path",
  106 => "mux i waveform", 107 => "waveform", 108 => "data i group",
  109 => "response", 110 => "pick", 111 => "pick i residual", 112 => "event",
  113 => "signif i event", 114 => "solution", 115 => "solution i err",
  116 => "focal i mech", 118 => "vel i model", 119 => "vel i layer i data",
  123 => "resp i pz i data", 125 => "lsa i detection", 126 => "lsa i setting",
  130 => "clock i rate", 131 => "recorder", 200 => "variable i info",
  201 => "structure i info", 202 => "member i info", 203 => "stream",
  204 => "file i index", 205 => "code i list", 206 => "gui i default",
  211 => "comment", 212 => "structure i tag", 213 => "terminator",
  214 => "code i data", 300 => "site", 301 => "spectra",
  302 => "sig i cmp i data", 303 => "resp i fap i data",
  304 => "lsa i set i data", 305 => "resp i cfs i data",
  306 => "sig i path i data", 307 => "magnitude", 308 => "processing",
  309 => "polarity", 310 => "ssam i setup", 311 => "ssam i output",
  312 => "map i element", 313 => "seismometer", 314 => "resp i fir i data",
  315 => "filter", 316 => "sig i path i ass", 317 => "ssam i band i data",
  318 => "beam i data", 319 => "resp i sen i data", 320 => "calibration",
  321 => "source", 322 => "user i vars", 323 => "service",
  324 => "recorder i ass", 325 => "seismo i ass", 326 => "coordinate i sys"
  )

const pick_types = Dict{Int16, String}(
  0 => "NOT GIVEN",
  1 => "WINDOW",
  2 => "F FINIS",
  3 => "X MAXIMUM AMPLITUDE",
  4 => "INCREASE GAIN STEP",
  5 => "DECREASE GAIN STEP",
  50 => "P",
  51 => "P",
  52 => "P∗",
  53 => "PP",
  54 => "PPP",
  55 => "PPPP",
  56 => "PPS",
  57 => "PG",
  58 => "PN",
  59 => "PDIFFRACTED",
  60 => "PCP",
  61 => "PCPPKP",
  62 => "PCS",
  63 => "PP",
  64 => "PPP",
  65 => "PKP",
  66 => "PKPPKP",
  67 => "PKPPKS",
  68 => "PKPSKS",
  69 => "PKS",
  70 => "PPKS",
  71 => "PKKP",
  72 => "PKKS",
  73 => "PCPPKP",
  74 => "PCSPKP",
  100 => "S",
  101 => "S",
  102 => "S∗",
  103 => "SS",
  104 => "SSS",
  105 => "SSSS",
  106 => "SG",
  107 => "SN",
  108 => "SCS",
  109 => "SPCS",
  110 => "SS",
  111 => "SSS",
  112 => "SSSS",
  113 => "SSCS",
  114 => "SCSPKP",
  115 => "SCP",
  116 => "SKS",
  117 => "SKKS",
  118 => "SKKKS",
  119 => "SKSSKS",
  120 => "SKP",
  121 => "SKKP",
  122 => "SKKKP",
  201 => "LG",
  202 => "LR",
  203 => "LR2",
  204 => "LR3",
  205 => "LR4",
  206 => "LQ",
  207 => "LQ2",
  208 => "LQ3",
  209 => "LQ4",
  301 => "T")

# Instrument codes
# 0 = "not specified"
# 1 = "sp usgs"
# 2 = "sp wwssn"
# 3 = "lp wwssn"
# 4 = "sp dwwssn"
# 5 = "lp dwwssn"
# 6 = "hglp lamont"
# 7 = "lp hglp lamont"
# 8 = "sp sro"
# 9 = "lp sro"
# 10 = "sp asro"
# 11 = "lp asro"
# 12 = "sp rstn"
# 13 = "lp rstn"
# 1 4 = "sp uofa U of alaska"
# 15 = "STS-1/UVBB"
# 16 = "STS-1/VBB"
# 17 = "STS-2"
# 18 = "FBA-23"
# 19 = "Wilcoxin "
# 50 = "USGS cassette"
# 51 = "GEOS"
# 52 = "EDA"
# 53 = "Sprengnether refraction"
# 54 = "Teledyne refraction"
# 55 = "Kinemetrics refraction"
# 300 = "amplifier"
# 301 = "amp/vco"
# 302 = "filter"
# 303 = "summing amp"
# 304 = "transmitter"
# 305 = "receiver"
# 306 = "antenna"
# 307 = "battery"
# 308 = "solar cell"
# 309 = "discriminator"
# 310 = "discr. rack"
# 311 = "paper recorder"
# 312 = "film recorder"
# 313 = "smoked glass recorder"
# 314 = "atod converter"
# 315 = "computer"
# 316 = "clock"
# 317 = "time receiver"
# 318 = "magnetic tape"
# 319 = "magnetic disk"
# 320 = "optical disk"

# ampunits
# 'd' = "digital counts",
# 'm' = "millimeters on develocorder" ,
# 'n' = "nanometers (/sec or /sec/sec)",
# 'v' = "millivolts",

const auth = Dict{Int32, String}(
  0 => "none: not given" ,
  1 => "temp: Temporary, for testing purposes",
  2 => "suds: Internal to SUDS",
  101 => "calnet usgs menlo park, ca",
  102 => "alaska net usgs menlo park, ca" ,
  103 => "katmai net usgs menlo park, ca",
  104 => "scalnet usgs pasadena, ca.",
  120 => "shumagin net lamont palisades,ny",
  10000 => "gsmen: US Geological Survey, Menlo Park, CA",
  10001 => "suds: testing of suds at the USGS, Menlo Park, CA" ,
  10002 =>"calnt: network porcessing group, USGS, Menlo Park, CA",
  10005 => "5day: 5 day recorders US Geological Survey, Menlo Park, CA" ,
  10006 => "geos: GEOS recorders US Geological Survey, Menlo Park, CA",
  10007 => "cent: centipede recorders US Geological Survey, Menlo Park, CA" ,
  10008 => "citgs: CIT stations maintained by USGS, Menlo Park, CA",
  10009 => "lllgs: LLL stations maintained by USGS, Menlo Park, CA" ,
  10010 => "dwrgs: LLL stations maintained by USGS, Menlo Park, CA",
  10011 => "unrgs: UNR stations maintained by USGS, Menlo Park, CA",
  10012 => "yel: Yellowstone Park, Wyoming, maintained by USGS, Menlo Park, CA",
  10500 => "RTP: main rtp, USGS, Menlo Park",
  10501 => "PRTP: prototype rtp, USGS, Menlo Park",
  10502 => "MRTP: motorola rtp, USGS, Menlo Park" ,
  10503 => "TUST1: CUSP Tustin A/D #1",
  10504 => "TUST2: CUSP Tustin A/D #2" ,
  10505 => "ECLIP: CUSP Eclipse digitizer",
  10506 => "CVAX: CUSP-VAX/750 digitizer" ,
  10507 => "HPARK: Haliburton digital, Parkfield",
  10520 => "CITT1: Tustin #1, Pasadena",
  10521 => "CITT2: Tustin #2, Pasadena",
  10522 => "CITN3: 11/34 online, Pasadena" ,
  10523 => "CITS3: 11/34 online, Pasadena",
  10524 => "CITD1: Nova/Eclipse, Pasadena" ,
  10525 => "CITF: VAX, Pasadena",
  10526 => "CITH: hand timed in Pasadena" ,
  11000 => "Daiss, Charles, USGS, Menlo Park, CA",
  11001 => "Oppenheimer, Dave, USGS, Menlo Park, CA" ,
  11002 => "Eaton, Jerry, USGS, Menlo Park, CA",
  15000 => "gspas: US Geological Survey, Pasadena, CA" ,
  15001 => "tergp: TERRAscope, US Geological Survey, Pasadena, CA",
  20000 => "uofa: Geophysical Institute, University of Alaska, College, AK" ,
  30000 => "uofw: Geophysics, University of Washington, WA",
  40000 => "ldgo: Lamont Doherty Geological Observatory, Palisades, NY" ,
  50000 => "iris: IRIS Consortium, Seattle Data Center, WA",
  51000 => "gsn: Global Seismographic Network, USGS, Albuquerque, NM" ,
  52000 => "asro: Abbreviated Seismic Research Observatories",
  53000 => "passc: PASSCAL Program, IRIS",
  60000 => "lll: Lawrence Livermore Labs, Livermore, CA",
  70000 => "lbl: Lawrence Berkeley Labs, U. C. Berkeley, CA" ,
  80000 => "lanl: Los Alamos National Labs, Los Alamos, NM",
  90000 => "stl: St. Louis University, St. Louis, MO",
  100000 => "ucsd: University of California, San Diego and SCRIPPS",
  110000 => "ucb: University of California, Berkeley, CA",
  120000 => "ucsb: University of California, Santa Barbara, CA",
  130000 => "ucsc: University of California, Santa Cruz, CA",
  140000 => "usc: University of Southern California, Los Angeles, CA",
  150000 => "cit: California Institute of Technology, Pasadena, CA",
  150001 => "terct: TERRAscope network, California Institute of Technology, Pasadena, CA",
  160000 => "nnunr: Northern Nevada net, University of Nevada, Reno, NV",
  160001 => "snunr: Southern Nevada net, University of Nevada, Reno, NV",
  170000 => "utah: University of Utah, Salt Lake City, UT",
  180000 => "msu: Memphis State University, Memphis, TN",
  180000 => "msu: Memphis State University, Memphis, TN" ,
  181010 => "sanju: PANDA experiment in SAN JUAN, Argentina",
  181011 => "jujuy: PANDA experiment in JUJUY, Argentina",
  181020 => "newma: PANDA experiment in NEW MADRID, TN",
  181030 => "arken: PANDA experiment in AK",
  181040 => "hawii: PANDA experiment in HAWII",
  181050 => "palmn: PANDA experiment in PALMERSTON NORTH, New Zealand" ,
  181051 => "taran: PANDA2 experiment in mountain TARAMAKI, New Zealand",
  181060 => "taiwa: PANDA2 experiment in Taiwan",
  187000 => "archj: ARCH Johnston, professor, director of research",
  187001 => "jmch: Jer-Ming CHiu, professor",
  187002 => "wych: Wai-Ying CHung, associate research professor",
  187003 => "hjdo: H.James DOrman, executive director",
  187004 => "mell: Michael ELLis, associate professor",
  187005 => "Josep: JOSE Pujol, associate professor",
  187006 => "paulr: PAUL Rydelek, assistant research professor",
  187007 => "robsm: ROBert SMalley, assistant research professor" ,
  187008 => "paulb: PAUL Bodin, assistant professor",
  187009 => "eusc: EUgene SChweig, adjunct professor, USGS geologist" ,
  187010 => "johng: JOHN Geomberg, adjunct professor, USGS geophysicist",
  187011 => "scda: SCott DAvis, USGS guest researcher",
  187500 => "jimbo: JIM BOllwerk, seismic networks engineer",
  187501 => "stepb: STEPhen Brewer, ceri seismic networks director" ,
  187502 => "cchiu: Christy CHIU, research associate II",
  187503 => "michf: MICHael Frohme, director of computing",
  188000 => "zrli: ZhaoRen LI, graduate research assistant",
  188001 => "kcch: Kou-Cheng Chen, graduate research assistant" ,
  189000 => "group: data processing GROUP in ceri",
  190000 => "aftac: AFTAC Center for Seismic Studies, Alexandria, VA" ,
  200000 => "uhhil: University of Hawaii, Hilo, HA",
  210000 => "uhhon: University of Hawaii, Honolulu, HA" ,
  220000 => "mit: Massachusetts Institute of Technology, Cambridge, MA",
  230000 => "dtm: Department of Terrestrial Magnetism, Washington, DC" ,
  240000 => "vpi: Virginia Polytechnic Institute, Blacksburg, VA",
  250000 => "anu: Australian National University",
  260000 => "gsgol: US Geological Survey, Golden, CO",
  260001 => "nngsg: Northern Nevada network, US Geological Survey, Golden, CO" ,
  260002 => "sngsg: Southern Nevada network, US Geological Survey, Golden, CO",
  270000 => "bmr: Bureau of Mineral Resources",
  280000 => "cands: Canadian Digital Seismic Network",
  290000 => "cdsn: China Digital Seismic Network",
  300000 => "cdmg: California Division Mines-Geology, Sacramento, CA",
  310000 => "pge: Pacific Gas and Electric/Woodward-Clyde, CA",
  315001 => "unoiv: Union Oil, Imperial Valley, CA",
  315002 => "unoml: Union Oil, Medicine Lake",
  320000 => "terra: Terra Corporation, Mendocino, CA",
  330000 => "cadwr: California Division of Water Resources",
  340000 => "gikar: Geophysical Institute, Karlsruhe, Germany",
  350000 => "gfz: GeoForschungsZentrum, Potsdam, Germany",
  360000 => "cnrir: CNR-IRS, Milan, Italy",
  370000 => "gsc: Geological Survey of Canada, Ottawa, Canada" ,
  380000 => "ind: industry",
  385000 => "geot: Geotech, Garland, Texas",
  390000 => "nano: Nanometrics, Kanata, Ontario, Canada",
  395000 => "lenn: Lennartz Electronic, Tubingen, Germany",
  400000 => "kine: Kinemetrics, Pasadena, CA",
  405000 => "snl: Sandia National Laboratories, Albuquerque, NM",
  410000 => "cices: CICESE, Ensenada, Mexico",
  415000 => "nmt: New Mexico Inst Mining and Tech, Soccorro, NM"
)

const mag_scale = ("coda", "tau", "xmag", "ml", "mb", "ms", "mw")

# const mag_type = Dict{Char, String}(
# 'A' => "average coda and amplitude",
# 'S' => "Msz",
# 'a' => "amplitude",
# 'b' => "Mb",
# 'c' => "coda",
# 'l' => "Ml",
# 'm' => "moment",
# 's' => "Ms",
# 'w' = "Mw"
# )

const loc_prog = Dict{Char, String}(
'7' => "Hypo-71",
'l' => "HypoLayer",
'e' => "HypoEllipse",
'i' => "HypoInverse", # Also given as 'h' in SAC readsuds documentation
'r' => "relp",
'u' => "Uhrhammer",
'c' => "Centroid",
'h' => "HypoInverse",
'v' => "Velest"
)


text:
#md # ---

using DifferentialEquations, Plots

g = 9.79 # Gravitational constants
L = 1.00 # Length of the pendulum

#Initial Conditions
u₀ = [0, π / 60] # Initial speed and initial angle
tspan = (0.0, 6.3) # time domain

#Define the problem
function simplependulum(du, u, p, t)
    θ = u[1]
    dθ = u[2]
    du[1] = dθ
    du[2] = -(g/L)*θ
end

#Pass to solvers
prob = ODEProblem(simplependulum, u₀, tspan)
sol = solve(prob, Tsit5(), reltol = 1e-6)
#md nothing # hide

#md # ---

# Analytic solution
u = u₀[2] .* cos.(sqrt(g / L) .* sol.t)

plot(sol.t, getindex.(sol.u, 2), label = "Numerical")
plot!(sol.t, u, label = "Analytic")
#md savefig("pendulum1.svg"); nothing # hide

#md # ![](pendulum1.svg)

#md # ---

# [Numbers with Uncertainties](http://tutorials.juliadiffeq.org/html/type_handling/02-uncertainties.html)

using Measurements

g = 9.79 ± 0.02; # Gravitational constants
L = 1.00 ± 0.01; # Length of the pendulum

#Initial Conditions
u₀ = [0 ± 0, π / 60 ± 0.01] # Initial speed and initial angle

#Define the problem
function simplependulum(du, u, p, t)
    θ = u[1]
    dθ = u[2]
    du[1] = dθ
    du[2] = -(g/L)*θ
end

#Pass to solvers
prob = ODEProblem(simplependulum, u₀, tspan)
sol = solve(prob, Tsit5(), reltol = 1e-6);
nothing # hide

#md # ---

# Analytic solution
u = u₀[2] .* cos.(sqrt(g / L) .* sol.t)

plot(sol.t, getindex.(sol.u, 2), label = "Numerical")
plot!(sol.t, u, label = "Analytic")
#md savefig("pendulum2.svg"); nothing # hide

#md # ![](pendulum2.svg)

max_stars_repo_name: pnavaro/DiffEq

max_issues_repo_name: pnavaro/DiffEq

max_forks_repo_name: pnavaro/DiffEq


text:
@testset "$TEST $G" begin

r1 = G(10,20)
r2 = DG(5,10)
@test typeof(r1) == G
@test typeof(r2) == DG
@test nv(r1) == 10
@test ne(r1) == 20
@test nv(r2) == 5
@test ne(r2) == 10

@test G(10,20,seed=3) == G(10, 20, seed=3)
@test DG(10,20,seed=3) == DG(10, 20, seed=3)
@test G(10,20,seed=3) == erdos_renyi(10, 20, G, seed=3)
@test ne(G(10,40,seed=3)) == 40
@test ne(DG(10,80,seed=3)) == 80

er = erdos_renyi(10, 0.5, G)
@test nv(er) == 10
@test ne(er) < 40
@test is_directed(er) == false

er = erdos_renyi(10, 0.5, DG)
@test nv(er) == 10
@test is_directed(er) == true

er = erdos_renyi(10, 0.5, G, seed=17)
@test nv(er) == 10
@test is_directed(er) == false

ws = watts_strogatz(10, 4, 0.2,  G)
@test nv(ws) == 10
@test ne(ws) == 20
@test is_directed(ws) == false

ws = watts_strogatz(10, 4, 0.2, DG)
@test nv(ws) == 10
@test ne(ws) == 20
@test is_directed(ws) == true

ba = barabasi_albert(10, 2,  G)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == false

ba = barabasi_albert(10, 2, 2,  G)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == false

ba = barabasi_albert(10, 4, 2, G)
@test nv(ba) == 10
@test ne(ba) == 12
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(2, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 17
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(2, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 17
@test is_directed(ba) == false

ba = barabasi_albert!(CompleteGraph(4, G), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == false

ba = barabasi_albert(10, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == true

ba = barabasi_albert(10, 2, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 16
@test is_directed(ba) == true

ba = barabasi_albert(10, 4, 2, DG)
@test nv(ba) == 10
@test ne(ba) == 12
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(2, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(2, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 18
@test is_directed(ba) == true

ba = barabasi_albert!(CompleteDiGraph(4, DG), 10, 2)
@test nv(ba) == 10
@test ne(ba) == 24
@test is_directed(ba) == true

fm = static_fitness_model(20, rand(10), G)
@test nv(fm) == 10
@test ne(fm) == 20
@test is_directed(fm) == false

fm = static_fitness_model(20, rand(10), rand(10), DG)
@test nv(fm) == 10
@test ne(fm) == 20
@test is_directed(fm) == true

sf = static_scale_free(10, 20, 2.0, G)
@test nv(sf) == 10
@test ne(sf) == 20
@test is_directed(sf) == false

sf = static_scale_free(10, 20, 2.0, 2.0, DG)
@test nv(sf) == 10
@test ne(sf) == 20
@test is_directed(sf) == true

rr = random_regular_graph(5, 0, G)
@test nv(rr) == 5
@test ne(rr) == 0
@test is_directed(rr) == false

rd = random_regular_digraph(10,0)
@test nv(rd) == 10
@test ne(rd) == 0
@test is_directed(rd)

rr = random_regular_graph(6, 3, G, seed=1)
@test nv(rr) == 6
@test ne(rr) == 9
@test is_directed(rr) == false

rr = random_regular_graph(100, 5, G)
@test nv(rr) == 100
@test ne(rr) == 250
@test is_directed(rr) == false
for v in vertices(rr)
    @test degree(rr, v) == 5
end

rr = random_configuration_model(10, repeat([2,4] ,5), G, seed=3)
@test nv(rr) == 10
@test ne(rr) == 15
@test is_directed(rr) == false
num2 = 0; num4 = 0
for v in vertices(rr)
    d = degree(rr, v)
    @test  d == 2 || d == 4
    d == 2 ? num2 += 1 : num4 += 1
end
@test num4 == 5
@test num2 == 5

rr = random_configuration_model(100, zeros(Int,100), G)
@test nv(rr) == 100
@test ne(rr) == 0
@test is_directed(rr) == false

rr = random_configuration_model(3, [2,2,2], G, check_graphical=true)
@test nv(rr) == 3
@test ne(rr) == 3
@test is_directed(rr) == false

rd = random_regular_digraph(100, 4, DG)
@test nv(rd) == 100
@test ne(rd) == 400
@test is_directed(rd)
@test std(out_degree(rd)) == 0

rd = random_regular_digraph(100, 4, DG, dir=:in)
@test nv(rd) == 100
@test ne(rd) == 400
@test is_directed(rd)
@test std(in_degree(rd)) == 0

rr = random_regular_graph(10, 8, G, seed=4)
@test nv(rr) == 10
@test ne(rr) == 40
@test is_directed(rr) == false
for v in vertices(rr)
    @test degree(rr, v) == 8
end

rd = random_regular_digraph(10, 8, DG, dir=:out, seed=4)
@test nv(rd) == 10
@test ne(rd) == 80
@test is_directed(rd)

g = stochastic_block_model(2., 3., [100,100], G)
@test  3.5 < mean(degree(g)) < 6.5
g = stochastic_block_model(3., 4., [100,100,100])
@test  9.5 < mean(degree(g)) < 12.5

end # testset


text:
module CSSUtil

using WebIO
using JSON
import WebIO: render

export style
export @md_str

using Measures
using Colors

function style(dict::Dict)
    Dict(:style=>dict)
end

function style(p::Pair...)
    style(Dict(p...))
end

"""
When a Node's `instanceof` field is set to `Fallthrough()`
it renders it only child and places it in place of itself
("splat"s its children with the Node's own siblings)

This is useful when you want to defer rendering an object
till its parents have been rendered
"""
struct Fallthrough
end

wrapnode(n::Node) = n
wrapnode(x) = node(Fallthrough(), x)

function render(n::Node{Fallthrough})
    inner = render(first(children(n)))
    inner isa String ? inner : inner(props(n))
end

JSON.lower(n::Node{Fallthrough}) = JSON.lower(render(n))

function style(elem, dict::Dict)
    wrapnode(elem)(style(dict))
end

function style(elem, p::Pair...)
    wrapnode(elem)(style(p...))
end

function style(::Nothing, arg::Pair...)
    style(arg...)
end

function style(::Nothing, arg::Dict)
    style(arg)
end

empty() = dom"div"()

export mm, em, cm, inch, pt, px, vw, vh, vmin, cent
"1mm"
const mm = Length(:mm, 1.0)
" 1em "
const em = Length(:em, 1.0)
"1cm"
const cm = Length(:cm, 1.0)
"1inch"
const inch = Length(:in, 1.0)
"1pt"
const pt = Length(:pt, 1.0)
"1px"
const px = Length(:px, 1.0)
"1vw"
const vw = Length(:vw, 1.0)
"1vh"
const vh = Length(:vh, 1.0)
"1vmin"
const vmin = Length(:vmin, 1.0)
"1% length"
const cent = Length(:cent, 1.0)

JSON.lower(l::Length{u}) where {u} = "$(l.value)$u"
JSON.lower(l::Length{:cent}) = "$(l.value)%"
JSON.lower(c::Color) = "#$(hex(c))"

function assertoneof(x, xs, name="argument")
    if !(string(x) in xs)
        throw(ArgumentError("$name should be one of $(join(xs, ",", "or")) as a string or symbol."))
    end
end

include("layout.jl")
include("theme.jl")
include("markdown.jl")

end # module

max_stars_repo_name: JuliaGizmos/CSSUtil


text:
### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, StatsPlots
	Random.seed!(2)
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing5.13"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	alpha  = 0.05
	L(obs) = obs - (1-sqrt(alpha))
	U(obs) = obs + (1-sqrt(alpha))

	mu = 5.57
	triDist = TriangularDist(mu-1,mu+1,mu)
end

# ╔═╡ 28756a6e-3af8-11eb-012f-93dc4ef1c3c9
begin
	N = 100
	hitBounds, missBounds = zeros(N, 2), zeros(N,2)
	for i in 1:N
		observation = rand(triDist)
		LL, UU = L(observation), U(observation)
		if LL <= mu && mu <= UU
			hitBounds[i,:] = [LL  UU-LL]
		else
			missBounds[i,:] = [LL  UU-LL]
		end
	end
end

# ╔═╡ 28759980-3af8-11eb-17c2-a3ee2e86a55d
begin
	groupedbar(hitBounds, bar_position=:stack,
		c=:blue, la=0, fa=[0 1], label="", ylims=(3,8))
	groupedbar!(missBounds, bar_position=:stack,
		c=:red, la=0, fa=[0 1], label="", ylims=(3,8))
	plot!([0,N+1],[mu,mu],
		c=:black, xlims=(0,N+1),
		ylims=(3,8), label="Parameter value", ylabel="Value Estimate")
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing5.13"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═28756a6e-3af8-11eb-012f-93dc4ef1c3c9
# ╠═28759980-3af8-11eb-17c2-a3ee2e86a55d
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


text:
# set up

include("set_up.jl")

using JLD
using HDF5

################################################################################
##      parameters for training data                                          ##
################################################################################

# set parameters for all jobs

# burn-in
burn_in = 10000 # 10000

# length training data
length_training_data = 5000 # 5000

# length test data
length_test_data = 5000 # 5000

# nbr iterations
nbr_iterations = burn_in+length_training_data + length_test_data

# set nbr parameters
set_nbr_params = 7

# set nbr cores
nbr_of_cores = 4

# log-scale prior
log_scale_prior = false

# algorithm
mcmc_alg = "MCWM"  # set MCWM or PMCMC

# type of job
job = ARGS[1] # set work to simdata or new_data

# set jod dep. parameters
if job == "simdata"

	# jobname
	jobname = "da_ada_training_data"*job

	# nbr particels
	nbr_particels = 3*400

	# use simulated data
	sim_data = true # set to true to use sim data

	# data set
	data_set = "old" # was "old"

	# dt
	dt = 0.035 # new = 0.35 old = 0.035

	# dt_U
	dt_U = 1. # new = 1 old = 1

elseif job == "new_data"

	# jobname
	jobname = "da_ada_training_data"*job

	# nbr particels
	nbr_particels = 2000 #3*500

	# use simulated data
	sim_data = false # set to true to use sim data

	# data set
	data_set = "new" # was "old"

	# dt
	dt = 0.35 # new = 0.35 old = 0.035

	# dt_U
	dt_U = 1. # new = 1 old = 1

end

################################################################################
##                         set model parameters  for training                 ##
################################################################################

problem_training = set_up_problem(nbr_of_unknown_parameters=set_nbr_params, use_sim_data = sim_data, data_set = data_set)
problem_training.alg_param.alg = mcmc_alg
problem_training.alg_param.R = nbr_iterations
problem_training.alg_param.N = nbr_particels
problem_training.alg_param.burn_in = burn_in
problem_training.alg_param.nbr_of_cores = nbr_of_cores
problem_training.alg_param.dt = dt
problem_training.alg_param.dt_U = dt_U
problem_training.alg_param.alg = mcmc_alg
problem_training.adaptive_update =  AMUpdate_gen(eye(set_nbr_params), 1/sqrt(set_nbr_params), 0.15, 1, 0.8, 25)

################################################################################
##                generate training data                                     ###
################################################################################

if !log_scale_prior
	tic()
    println("run Normal prior model")
	res_training, theta_training, loglik_training, cov_matrix = mcmc(problem_training, true, true)
	time_pre_er = toc()
	#export_parameters(res_problem_normal_prior_est_AM_gen[2],jobname)
else
	tic()
    println("run log-scale prior model")
	res_training, theta_training, loglik_training, cov_matrix  = @time mcmc(problem_training_nonlog, true, true)
	time_pre_er = toc()
end

################################################################################
##                store generated data                                       ###
################################################################################

# export generated data
export_data(problem_training, res_training[1],"gp_training_$(set_nbr_params)_par_lunarc_new_data_4_cores"*job)

# split training and test data

theta_test = theta_training[:,(end-length_test_data+1):end]
loglik_test = loglik_training[(end-length_test_data+1):end]

theta_training = theta_training[:,1:length_training_data]
loglik_training = loglik_training[1:length_training_data]

# save training data, test data, and covaraince matrix to a Julia workspace file
save("gp_training_$(set_nbr_params)_par_training_and_test"*job*"lunarc_new_new.jld", "res_training", res_training, "theta_training", theta_training, "loglik_training", loglik_training, "theta_test", theta_test, "loglik_test", loglik_test,"cov_matrix",cov_matrix)

max_stars_repo_name: SamuelWiqvist/adamcmcpaper

max_issues_repo_name: SamuelWiqvist/adamcmcpaper

max_forks_repo_name: SamuelWiqvist/adamcmcpaper