text:
"""
    StateVecs

Provides a state vector that is includes information
of the domain decomposition into a grid mean,
environment and updrafts.
"""
module StateVecs

using ..FiniteDifferenceGrids
export StateVec, over_sub_domains, Cut, Dual
export var_names, var_string, var_suffix
export assign!, assign_real!, assign_ghost!, extrap!, extrap_0th_order!
export compare

include("DomainDecomp.jl")
include("DomainSubSet.jl")
include("DomainIdx.jl")
include("VarMapper.jl")

struct FieldsPerElement
  vals::Vector
end

"""
    StateVec{VM}

A state vector containing the number of subdomains,
`n_subdomains`, a `NamedTuple` variable mapper, a tuple
of the variable names, and a vector of vectors, containing
the values for all of the variables.
"""
struct StateVec{VM}
  """
  A `DomainIdx`, containing the complete set of sub-domain indexes
  """
  idx::DomainIdx
  """
  A `NamedTuple` variable mapper, containing indexes per variable
  """
  var_mapper::VM
  """
  A `Dict` index per sub-domain per variable
  """
  idx_ss_per_var::Dict
  """
  A mapper, containing mapping from global sub-domain indexes to sub-domain indexes per variable
  """
  a_map::Dict
  """
  A `NamedTuple` sub-domain mapper, containing indexes per variable which are sub-domain specific
  """
  sd_unmapped::Dict
  """
  A `DomainDecomp`, which specifies the complete set of `DomainDecomp` across all variables
  """
  dd::DomainDecomp
  """
  A `DomainDecomp` per variable
  """
  dss_per_var::NamedTuple
  """
  A vector of vectors containing all values for all variables.
  """
  fields::Vector{FieldsPerElement}
end

"""
    StateVec(vars::Tuple{Vararg{Tuple{Symbol, DD}}}, grid::Grid{T}) where {DD, T}

Return a state vector, given a tuple of tuples of variable
names and the number of their subdomains.
"""
function StateVec(vars, grid::Grid{T}, dd::DomainDecomp) where {T}
  n_subdomains = sum(dd)
  n_vars = sum([sum(dd, dss) for (v, dss) in vars])
  var_mapper, dss_per_var, var_names = get_var_mapper(vars, dd)
  idx = DomainIdx(dd)
  sd_unmapped = get_sd_unmapped(vars, idx, dd)

  idx_ss_per_var = Dict(ϕ => DomainIdx(dd, dss_per_var[ϕ]) for ϕ in var_names)
  a_map = Dict(ϕ => get_sv_a_map(idx, idx_ss_per_var[ϕ]) for ϕ in var_names)

  all_vars = Any[eltype(grid.Δz)(0) for v in 1:n_vars]
  fields = [FieldsPerElement(deepcopy(all_vars)) for k in over_elems(grid)]
  return StateVec{typeof(var_mapper)}(idx,
                                      var_mapper,
                                      idx_ss_per_var,
                                      a_map,
                                      sd_unmapped,
                                      dd,
                                      dss_per_var,
                                      fields)
end

"""
    var_names(sv::StateVec)

Get tuple of variable names.
"""
var_names(sv::StateVec{VM}) where {VM} = fieldnames(VM)

function Base.show(io::IO, sv::StateVec)
  println(io, "\n----------------------- State Vector")
  println(io, "DomainIdx")
  println(io, sv.idx)
  println(io, "var_mapper")
  println(io, sv.var_mapper)
  vn = var_names(sv)
  headers = [length(over_sub_domains(sv, ϕ))==1 ? string(ϕ) : string(ϕ)*"_"*string(i)
             for ϕ in vn for i in over_sub_domains(sv, ϕ)]
  n_vars = length(headers)
  n_elem = length(sv.fields)
  data = reshape([sv[ϕ, k, i] for ϕ in vn for
    i in over_sub_domains(sv, ϕ) for k in 1:n_elem], n_elem, n_vars)
  Base.print_matrix(io, [reshape(Text.(headers), 1, n_vars); data])
  println(io, "\n-----------------------")
end

DomainIdx(sv::StateVec) = sv.idx

"""
    over_sub_domains(sv::StateVec, ϕ::Symbol)

Get list of indexes over all subdomains for variable `ϕ`.
"""
function over_sub_domains(sv::StateVec, ϕ::Symbol)
  return [x for x in sv.sd_unmapped[ϕ] if !(x==0)]
end

function Base.getindex(sv::StateVec, ϕ::Symbol, k, i = 0)
  i==0 && (i = gridmean(sv.idx))
  i_sv = get_i_state_vec(sv.var_mapper, sv.a_map[ϕ], ϕ, i)
  @inbounds sv.fields[k].vals[i_sv]
end
function Base.setindex!(sv::StateVec, val, ϕ::Symbol, k, i = 0)
  i==0 && (i = gridmean(sv.idx))
  @inbounds i_sv = get_i_state_vec(sv.var_mapper, sv.a_map[ϕ], ϕ, i)
  @inbounds sv.fields[k].vals[i_sv] = val
end

abstract type AbstractCut{I} end

"""
    Cut{I} <: AbstractCut{I}

A Cut struct used to slice the state
vector along the grid-element dimension.
This is used as an API to pass Cuts
into local derivative/interpolation routines.
"""
struct Cut{I} <: AbstractCut{I}
  e::I
end
function Base.getindex(sv::StateVec, ϕ::Symbol, cut::Cut, i=0)
  @inbounds [sv[ϕ, k, i] for k in cut.e-1:cut.e+1]
end

"""
    Dual{I} <: AbstractCut{I}

A Dual struct used to slice the state
vector along the grid-element dimension.
In addition, this interpolates fields
[(f[k-1]+f[k])/2, (f[k]+f[k+1])/2]
"""
struct Dual{I} <: AbstractCut{I}
  e::I
end
function Base.getindex(sv::StateVec, ϕ::Symbol, dual::Dual, i=0)
  @inbounds [(sv[ϕ, k, i]+sv[ϕ, k+1, i])/2 for k in dual.e-1:dual.e]
end

Base.isnan(sv::StateVec) = any([any(isnan.(fpe.vals)) for fpe in sv.fields])
Base.isinf(sv::StateVec) = any([any(isinf.(fpe.vals)) for fpe in sv.fields])

"""
    var_suffix(sv::StateVec, ϕ::Symbol, i=0)

A suffix string for variable `ϕ` indicating its sub-domain.
"""
function var_suffix(sv::StateVec, ϕ::Symbol, i=0)
  i==0 && (i = gridmean(DomainIdx(sv)))
  var_suffix(sv.var_mapper, sv.idx, sv.idx_ss_per_var[ϕ], ϕ, i)
end

"""
    var_string(sv::StateVec, ϕ::Symbol, i=0)

A string of the variable `ϕ` with the appropriate suffix indicating its sub-domain.
"""
function var_string(sv::StateVec, ϕ::Symbol, i=0)
  i==0 && (i = gridmean(DomainIdx(sv)))
  var_string(sv.var_mapper, sv.idx, sv.idx_ss_per_var[ϕ], ϕ, i)
end

"""
    assign!(sv::StateVec, var_names, grid::Grid, val, i=1)

Assign value `val` to variable `ϕ` for all ghost points.
"""
function assign!(sv::StateVec, var_names, grid::Grid{FT}, val::FT) where FT
  gm, en, ud, sd, al = allcombinations(DomainIdx(sv))
  !(var_names isa Tuple) && (var_names = (var_names,))
  @inbounds for k in over_elems(grid), ϕ in var_names, i in over_sub_domains(sv, ϕ)
    sv[ϕ, k, i] = val
  end
end

"""
    assign!(sv::StateVec, grid::Grid, val)

Assign value `val` to all variables in state vector.
"""
function assign!(sv::StateVec, grid::Grid{FT}, val::FT) where FT
  gm, en, ud, sd, al = allcombinations(DomainIdx(sv))
  @inbounds for k in over_elems(grid), ϕ in var_names(sv), i in over_sub_domains(sv, ϕ)
    sv[ϕ, k, i] = val
  end
end

"""
    assign_real!(sv::StateVec, ϕ::Symbol, grid::Grid, vec_real, i=0)

Assign the real elements to the given vector `vec_real`.
"""
function assign_real!(sv::StateVec, ϕ::Symbol, grid::Grid, vec_real, i=0)
  i==0 && (i = gridmean(sv.idx))
  k_real = 1
  @assert length(vec_real)==length(over_elems_real(grid))
  @inbounds for k in over_elems_real(grid)
    sv[ϕ, k, i] = vec_real[k_real]
    k_real+=1
  end
end

"""
    assign_ghost!(sv::StateVec, ϕ::Symbol, grid::Grid, val, i=0)

Assign value `val` to variable `ϕ` for all ghost points.
"""
function assign_ghost!(sv::StateVec, ϕ::Symbol, grid::Grid, val, i=0)
  i==0 && (i = gridmean(sv.idx))
  @inbounds for k in over_elems_ghost(grid)
    sv[ϕ, k, i] = val
  end
end

"""
    assign_ghost!(sv::StateVec, ϕ::Symbol, grid::Grid, val, b::ZBoundary, i=1)

Assign value `val` to variable `ϕ` for all ghost points.
"""
function assign_ghost!(sv::StateVec, ϕ::Symbol, grid::Grid, val, b::ZBoundary, i=0)
  i==0 && (i = gridmean(sv.idx))
  @inbounds for k in over_elems_ghost(grid, b)
    sv[ϕ, k, i] = val
  end
end

"""
    compare(sv::StateVec, sv_expected::StateVec, grid::Grid, tol)

A dictionary, with keys in `var_names(vs)`, containing a `Bool` indicating
that `sv` ≈ `sv_expected` for all of their sub-domains for all elements.
"""
function compare(sv::StateVec, sv_expected::StateVec, grid::Grid{FT}, tol) where FT
  D = Dict(ϕ => [true for i in over_sub_domains(sv, ϕ)] for ϕ in var_names(sv))
  @inbounds for k in over_elems(grid)
    @inbounds for ϕ in var_names(sv)
      @inbounds for i in over_sub_domains(sv, ϕ)
        if abs(sv[ϕ, k, i] - sv_expected[ϕ, k, i]) > tol
          i_var = get_i_var(sv.a_map[ϕ], i)
          D[ϕ][i_var] = false
        end
      end
    end
  end
  return D
end

"""
    extrap!(sv::StateVec, ϕ::Symbol, grid::Grid, i=1)

Extrapolate variable `ϕ` to the first ghost point.
"""
function extrap!(sv::StateVec, ϕ::Symbol, grid::Grid, i=0)
  i==0 && (i = gridmean(sv.idx))
  @inbounds for b in (Zmin(), Zmax())
    kg, ki, kii = boundary_points(grid, b)
    sv[ϕ, kg, i] = 2*sv[ϕ, ki, i] - sv[ϕ, kii, i]
  end
end

"""
    extrap!(sv::StateVec, ϕ::Symbol, grid::Grid, ::ZBoundary, i=1)

Extrapolate variable `ϕ` to the first ghost point.
"""
function extrap!(sv::StateVec, ϕ::Symbol, grid::Grid, b::ZBoundary, i=0)
  i==0 && (i = gridmean(sv.idx))
  kg, ki, kii = boundary_points(grid, b)
  @inbounds sv[ϕ, kg, i] = 2*sv[ϕ, ki, i] - sv[ϕ, kii, i]
end

"""
    extrap!(sv::StateVec, ϕ::Symbol, grid::Grid, dual_val, b::ZBoundary, i=1)

Extrapolate variable `ϕ` to the first ghost point.
"""
function extrap!(sv::StateVec, ϕ::Symbol, grid::Grid, dual_val, b::ZBoundary, i=0)
  i==0 && (i = gridmean(sv.idx))
  ki = first_interior(grid, b)
  kg = first_ghost(grid, b)
  @inbounds sv[ϕ, kg, i] = 2*dual_val - sv[ϕ, ki, i]
end

"""
    extrap_0th_order!(sv::StateVec, ϕ::Symbol, grid::Grid, i=1)

Extrapolate variable `ϕ` to the first ghost point using zeroth order approximation.
"""
function extrap_0th_order!(sv::StateVec, ϕ::Symbol, grid::Grid, i=0)
  i==0 && (i = gridmean(sv.idx))
  @inbounds for b in (Zmin(), Zmax())
    kg, ki, kii = boundary_points(grid, b)
    sv[ϕ, kg, i] = sv[ϕ, ki, i]
  end
end

include("BoundaryConditions.jl")
include("ExportFuncs.jl")
include("PlotFuncs.jl")
include("DomainDecompFuncs.jl")

end

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
using Base: llvmcall
import Base.(+)

using ..Random123: R123Generator1x, R123Generator4x
import ..Random123: random123_r, set_counter!

const __m128i = NTuple{2, VecElement{UInt64}}
Base.convert(::Type{__m128i}, x::UInt128) = unsafe_load(Ptr{__m128i}(pointer_from_objref(Ref(x))))
Base.convert(::Type{UInt128}, x::__m128i) = unsafe_load(Ptr{UInt128}(pointer_from_objref(Ref(x))))
UInt128(x::__m128i) = convert(UInt128, x)
__m128i(x::UInt128) = convert(__m128i, x)
Base.convert(::Type{__m128i}, x::Union{Signed, Unsigned}) = convert(__m128i, UInt128(x))
Base.convert(::Type{T}, x::__m128i) where T <: Union{Signed, Unsigned} = convert(T, UInt128(x))

const LITTLE_ENDIAN = ENDIAN_BOM ≡ 0x04030201
__m128i(hi::UInt64, lo::UInt64) = LITTLE_ENDIAN ? (VecElement(lo), VecElement(hi)) : (VecElement(hi), VecElement(lo))

Base.zero(::Type{__m128i}) = __m128i(zero(UInt64), zero(UInt64))
Base.one(::Type{__m128i}) = __m128i(zero(UInt64), one(UInt64))
Base.xor(a::__m128i, b::__m128i) = llvmcall(
    """%3 = xor <2 x i64> %1, %0
    ret <2 x i64> %3""",
    __m128i, Tuple{__m128i, __m128i},
    a, b
)
(+)(a::__m128i, b::__m128i) = llvmcall(
    """%3 = add <2 x i64> %1, %0
    ret <2 x i64> %3""",
    __m128i, Tuple{__m128i, __m128i},
    a, b
)
(+)(a::__m128i, b::Integer) = a + __m128i(UInt128(b))

_aes_enc(a::__m128i, round_key::__m128i) = ccall(
    "llvm.x86.aesni.aesenc",
    llvmcall,
    __m128i, (__m128i, __m128i),
    a, round_key
)
_aes_enc_last(a::__m128i, round_key::__m128i) = ccall(
    "llvm.x86.aesni.aesenclast",
    llvmcall,
    __m128i, (__m128i, __m128i),
    a, round_key
)
_aes_key_gen_assist(a::__m128i, ::Val{R}) where R = ccall(
    "llvm.x86.aesni.aeskeygenassist",
    llvmcall,
    __m128i, (__m128i, UInt8),
    a, R
)

"Abstract RNG that generates one number at a time and is based on AESNI."
abstract type AbstractAESNI1x <: R123Generator1x{UInt128} end
"Abstract RNG that generates four numbers at a time and is based on AESNI."
abstract type AbstractAESNI4x <: R123Generator4x{UInt32} end

@inline function set_counter!(
    r::AbstractAESNI4x,
    ctr::NTuple{4, Integer}
)
    r.p = 0
    r.ctr1 = union_uint(Tuple(x % UInt32 for x in ctr))
    random123_r(r)
    r
end

@inline inc_counter!(r::AbstractAESNI4x) = (r.ctr1 += one(__m128i); r)

ext: jl


text:
module DaggerTimeDag

using Dagger
import TimeDag
import TimeDag: Node
import Dates: DateTime

function evaluate_internal(node::Node, time_start::DateTime, time_end::DateTime, args...)
    new_node = TimeDag.obtain_node(node.parents, node.op)
    state = TimeDag.start_at([node], time_start)
    TimeDag.evaluate_until!(state, time_end)
    return vcat(state.evaluated_node_to_blocks[node]...)
end
function evaluate(node::Node, time_start::DateTime, time_end::DateTime)
    to_transform = Node[node]
    node_to_task = IdDict{Node,Dagger.EagerThunk}()

    # Populate visit list
    while !isempty(to_transform)
        cur = popfirst!(to_transform)
        haskey(node_to_task, cur) && continue
        parent_tasks = Union{Node,Dagger.EagerThunk}[]
        for p in cur.parents
            push!(parent_tasks, get(node_to_task, p, p))
        end
        if !any(p->p isa Node, parent_tasks)
            # Ready to execute self
            task = Dagger.@spawn evaluate_internal(cur, time_start, time_end, parent_tasks...)
            node_to_task[cur] = task
        else
            # Need to revisit later
            # Push unprocessed parents first
            for p in filter(p->p isa Node, parent_tasks)
                push!(to_transform, p)
            end
            # Push self last
            push!(to_transform, cur)
        end
    end

    return node_to_task[node]
end

end # module


text:
#=
function set_upstream(sys::PSY.System, mapping::DataFrames.DataFrame)
    make_array(x) = ismissing(x) ? [] : split(strip(x, ['(', ')']), ",")
    for device in PSY.get_components(HydroEnergyCascade, sys)
        for upstream_device in
            make_array(mapping[findall(in([device.name]), mapping.hydrounit), :].upstream[1])
            if !isnothing(PSY.get_component(HydroEnergyCascade, sys, upstream_device))
                push!(
                    device.upstream,
                    PSY.get_component(HydroEnergyCascade, sys, upstream_device),
                )
            end
        end
    end
end
=#


text:
################################################################################
# Augmented Gradient Search ARS
################################################################################
using LinearAlgebra
using Statistics

import LinearAlgebra.normalize
import GeometryBasics.update

function train(env::Environment, policy::Policy{T}, normalizer::Normalizer{T}, hp::HyperParameters{T}) where T
    envs = [deepcopy(env) for i = 1:Threads.nthreads()]

    output_size, input_size = size(policy.θ)
    nx = input_size
    nu = output_size
    nθ = output_size * input_size
    # Rollout data
    X = [zeros(nx) for k = 1:hp.horizon+1]
    A = [zeros(nu) for k = 1:hp.horizon]
    fx = [zeros(nx,nx) for k = 1:hp.horizon]
    fu = [zeros(nx,nu) for k = 1:hp.horizon]
    # Gradient data
    ∇θ = zeros(1,nθ)
    ∂x∂θ = [zeros(nx, nθ) for k = 1:hp.horizon]

    for episode = 1:hp.main_loop_size
        # Stem
        state = reset(env)
        X[1] .= copy(state)
        done = false
        sum_reward = 0
        num_plays = 0
        for k = 1:hp.horizon
            done && break
            observe(normalizer, state)
            state = normalize(normalizer, state)
            action = evaluate(policy, state)
            A[k] .= copy(action)
            state, reward, done, _ = step(env, action, diff = true)
            X[k+1] .= copy(state)
            fx[k] .= copy(env.dynamics_jacobian_state)
            fu[k] .= copy(env.dynamics_jacobian_input)
            reward = max(min(reward, 1), -1)
            sum_reward += reward
            num_plays += 1
        end

        ∇θ .= 0.0
        ∂u∂θ = [cat([X[k]' for i = 1:output_size]..., dims = (1,2)) for k = 1:num_plays]
        ∂r∂x = [FiniteDiff.finite_difference_jacobian(x -> [-cost(env, x, A[k])], X[k]) for k = 1:num_plays]
        ∂r∂u = [FiniteDiff.finite_difference_jacobian(u -> [-cost(env, X[k], u)], A[k]) for k = 1:num_plays]
        for k = 2:num_plays
            ∂x∂θ[k] .= fx[k-1] * ∂x∂θ[k-1] + fu[k-1] * ∂u∂θ[k-1] #TODO
        end
        for k = 1:num_plays
            if norm(∇θ, Inf) > 8e2
                println("grad up to step $k")
                break
            end
            ∇θ += ∂r∂x[k] * ∂x∂θ[k] + ∂r∂u[k] * ∂u∂θ[k]
        end
        ∇ = transpose(reshape(∇θ, (input_size, output_size)))
        (norm(∇, Inf) < 1e3) && gradient_update(policy, ∇)

        # finish, print:
        println("episode $episode ∇∞ $(scn(norm(∇, Inf))) r $sum_reward")
    end
    return nothing
end

function gradient_update(policy::Policy, ∇)
    policy.θ += policy.hp.step_size * ∇ ./ norm(∇, Inf)
    return nothing
end


text:
function FieldVector2D(arg0::Field, arg1::Vector2D)
    return FieldVector2D((Field, Vector2D), arg0, arg1)
end

function FieldVector2D(arg0::RealFieldElement, arg1::FieldVector2D)
    return FieldVector2D((RealFieldElement, FieldVector2D), arg0, arg1)
end

function FieldVector2D(arg0::RealFieldElement, arg1::FieldVector2D, arg2::RealFieldElement, arg3::FieldVector2D)
    return FieldVector2D((RealFieldElement, FieldVector2D, RealFieldElement, FieldVector2D), arg0, arg1, arg2, arg3)
end

function FieldVector2D(arg0::RealFieldElement, arg1::FieldVector2D, arg2::RealFieldElement, arg3::FieldVector2D, arg4::RealFieldElement, arg5::FieldVector2D)
    return FieldVector2D((RealFieldElement, FieldVector2D, RealFieldElement, FieldVector2D, RealFieldElement, FieldVector2D), arg0, arg1, arg2, arg3, arg4, arg5)
end

function FieldVector2D(arg0::RealFieldElement, arg1::FieldVector2D, arg2::RealFieldElement, arg3::FieldVector2D, arg4::RealFieldElement, arg5::FieldVector2D, arg6::RealFieldElement, arg7::FieldVector2D)
    return FieldVector2D((RealFieldElement, FieldVector2D, RealFieldElement, FieldVector2D, RealFieldElement, FieldVector2D, RealFieldElement, FieldVector2D), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function FieldVector2D(arg0::RealFieldElement, arg1::RealFieldElement)
    return FieldVector2D((RealFieldElement, RealFieldElement), arg0, arg1)
end

function FieldVector2D(arg0::RealFieldElement, arg1::Vector2D)
    return FieldVector2D((RealFieldElement, Vector2D), arg0, arg1)
end

function FieldVector2D(arg0::RealFieldElement, arg1::Vector2D, arg2::RealFieldElement, arg3::Vector2D)
    return FieldVector2D((RealFieldElement, Vector2D, RealFieldElement, Vector2D), arg0, arg1, arg2, arg3)
end

function FieldVector2D(arg0::RealFieldElement, arg1::Vector2D, arg2::RealFieldElement, arg3::Vector2D, arg4::RealFieldElement, arg5::Vector2D)
    return FieldVector2D((RealFieldElement, Vector2D, RealFieldElement, Vector2D, RealFieldElement, Vector2D), arg0, arg1, arg2, arg3, arg4, arg5)
end

function FieldVector2D(arg0::RealFieldElement, arg1::Vector2D, arg2::RealFieldElement, arg3::Vector2D, arg4::RealFieldElement, arg5::Vector2D, arg6::RealFieldElement, arg7::Vector2D)
    return FieldVector2D((RealFieldElement, Vector2D, RealFieldElement, Vector2D, RealFieldElement, Vector2D, RealFieldElement, Vector2D), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function FieldVector2D(arg0::Vector{RealFieldElement})
    return FieldVector2D((Vector{RealFieldElement},), arg0)
end

function FieldVector2D(arg0::jdouble, arg1::FieldVector2D)
    return FieldVector2D((jdouble, FieldVector2D), arg0, arg1)
end

function FieldVector2D(arg0::jdouble, arg1::FieldVector2D, arg2::jdouble, arg3::FieldVector2D)
    return FieldVector2D((jdouble, FieldVector2D, jdouble, FieldVector2D), arg0, arg1, arg2, arg3)
end

function FieldVector2D(arg0::jdouble, arg1::FieldVector2D, arg2::jdouble, arg3::FieldVector2D, arg4::jdouble, arg5::FieldVector2D)
    return FieldVector2D((jdouble, FieldVector2D, jdouble, FieldVector2D, jdouble, FieldVector2D), arg0, arg1, arg2, arg3, arg4, arg5)
end

function FieldVector2D(arg0::jdouble, arg1::FieldVector2D, arg2::jdouble, arg3::FieldVector2D, arg4::jdouble, arg5::FieldVector2D, arg6::jdouble, arg7::FieldVector2D)
    return FieldVector2D((jdouble, FieldVector2D, jdouble, FieldVector2D, jdouble, FieldVector2D, jdouble, FieldVector2D), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function add(obj::FieldVector2D, arg0::FieldVector2D)
    return jcall(obj, "add", FieldVector2D, (FieldVector2D,), arg0)
end

function add(obj::FieldVector2D, arg0::RealFieldElement, arg1::FieldVector2D)
    return jcall(obj, "add", FieldVector2D, (RealFieldElement, FieldVector2D), arg0, arg1)
end

function add(obj::FieldVector2D, arg0::RealFieldElement, arg1::Vector2D)
    return jcall(obj, "add", FieldVector2D, (RealFieldElement, Vector2D), arg0, arg1)
end

function add(obj::FieldVector2D, arg0::Vector2D)
    return jcall(obj, "add", FieldVector2D, (Vector2D,), arg0)
end

function add(obj::FieldVector2D, arg0::jdouble, arg1::FieldVector2D)
    return jcall(obj, "add", FieldVector2D, (jdouble, FieldVector2D), arg0, arg1)
end

function add(obj::FieldVector2D, arg0::jdouble, arg1::Vector2D)
    return jcall(obj, "add", FieldVector2D, (jdouble, Vector2D), arg0, arg1)
end

function angle(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "angle", RealFieldElement, (FieldVector2D, FieldVector2D), arg0, arg1)
end

function angle(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::Vector2D)
    return jcall(FieldVector2D, "angle", RealFieldElement, (FieldVector2D, Vector2D), arg0, arg1)
end

function angle(::Type{FieldVector2D}, arg0::Vector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "angle", RealFieldElement, (Vector2D, FieldVector2D), arg0, arg1)
end

function cross_product(obj::FieldVector2D, arg0::FieldVector2D, arg1::FieldVector2D)
    return jcall(obj, "crossProduct", RealFieldElement, (FieldVector2D, FieldVector2D), arg0, arg1)
end

function cross_product(obj::FieldVector2D, arg0::Vector2D, arg1::Vector2D)
    return jcall(obj, "crossProduct", RealFieldElement, (Vector2D, Vector2D), arg0, arg1)
end

function distance(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "distance", RealFieldElement, (FieldVector2D, FieldVector2D), arg0, arg1)
end

function distance(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::Vector2D)
    return jcall(FieldVector2D, "distance", RealFieldElement, (FieldVector2D, Vector2D), arg0, arg1)
end

function distance(::Type{FieldVector2D}, arg0::Vector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "distance", RealFieldElement, (Vector2D, FieldVector2D), arg0, arg1)
end

function distance(obj::FieldVector2D, arg0::FieldVector2D)
    return jcall(obj, "distance", RealFieldElement, (FieldVector2D,), arg0)
end

function distance(obj::FieldVector2D, arg0::Vector2D)
    return jcall(obj, "distance", RealFieldElement, (Vector2D,), arg0)
end

function distance1(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "distance1", RealFieldElement, (FieldVector2D, FieldVector2D), arg0, arg1)
end

function distance1(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::Vector2D)
    return jcall(FieldVector2D, "distance1", RealFieldElement, (FieldVector2D, Vector2D), arg0, arg1)
end

function distance1(::Type{FieldVector2D}, arg0::Vector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "distance1", RealFieldElement, (Vector2D, FieldVector2D), arg0, arg1)
end

function distance1(obj::FieldVector2D, arg0::FieldVector2D)
    return jcall(obj, "distance1", RealFieldElement, (FieldVector2D,), arg0)
end

function distance1(obj::FieldVector2D, arg0::Vector2D)
    return jcall(obj, "distance1", RealFieldElement, (Vector2D,), arg0)
end

function distance_inf(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "distanceInf", RealFieldElement, (FieldVector2D, FieldVector2D), arg0, arg1)
end

function distance_inf(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::Vector2D)
    return jcall(FieldVector2D, "distanceInf", RealFieldElement, (FieldVector2D, Vector2D), arg0, arg1)
end

function distance_inf(::Type{FieldVector2D}, arg0::Vector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "distanceInf", RealFieldElement, (Vector2D, FieldVector2D), arg0, arg1)
end

function distance_inf(obj::FieldVector2D, arg0::FieldVector2D)
    return jcall(obj, "distanceInf", RealFieldElement, (FieldVector2D,), arg0)
end

function distance_inf(obj::FieldVector2D, arg0::Vector2D)
    return jcall(obj, "distanceInf", RealFieldElement, (Vector2D,), arg0)
end

function distance_sq(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "distanceSq", RealFieldElement, (FieldVector2D, FieldVector2D), arg0, arg1)
end

function distance_sq(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::Vector2D)
    return jcall(FieldVector2D, "distanceSq", RealFieldElement, (FieldVector2D, Vector2D), arg0, arg1)
end

function distance_sq(::Type{FieldVector2D}, arg0::Vector2D, arg1::FieldVector2D)
    return jcall(FieldVector2D, "distanceSq", RealFieldElement, (Vector2D, FieldVector2D), arg0, arg1)
end

function distance_sq(obj::FieldVector2D, arg0::FieldVector2D)
    return jcall(obj, "distanceSq", RealFieldElement, (FieldVector2D,), arg0)
end

function distance_sq(obj::FieldVector2D, arg0::Vector2D)
    return jcall(obj, "distanceSq", RealFieldElement, (Vector2D,), arg0)
end

function dot_product(obj::FieldVector2D, arg0::FieldVector2D)
    return jcall(obj, "dotProduct", RealFieldElement, (FieldVector2D,), arg0)
end

function dot_product(obj::FieldVector2D, arg0::Vector2D)
    return jcall(obj, "dotProduct", RealFieldElement, (Vector2D,), arg0)
end

function equals(obj::FieldVector2D, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_minus_i(::Type{FieldVector2D}, arg0::Field)
    return jcall(FieldVector2D, "getMinusI", FieldVector2D, (Field,), arg0)
end

function get_minus_j(::Type{FieldVector2D}, arg0::Field)
    return jcall(FieldVector2D, "getMinusJ", FieldVector2D, (Field,), arg0)
end

function get_na_n(::Type{FieldVector2D}, arg0::Field)
    return jcall(FieldVector2D, "getNaN", FieldVector2D, (Field,), arg0)
end

function get_negative_infinity(::Type{FieldVector2D}, arg0::Field)
    return jcall(FieldVector2D, "getNegativeInfinity", FieldVector2D, (Field,), arg0)
end

function get_norm(obj::FieldVector2D)
    return jcall(obj, "getNorm", RealFieldElement, ())
end

function get_norm1(obj::FieldVector2D)
    return jcall(obj, "getNorm1", RealFieldElement, ())
end

function get_norm_inf(obj::FieldVector2D)
    return jcall(obj, "getNormInf", RealFieldElement, ())
end

function get_norm_sq(obj::FieldVector2D)
    return jcall(obj, "getNormSq", RealFieldElement, ())
end

function get_plus_i(::Type{FieldVector2D}, arg0::Field)
    return jcall(FieldVector2D, "getPlusI", FieldVector2D, (Field,), arg0)
end

function get_plus_j(::Type{FieldVector2D}, arg0::Field)
    return jcall(FieldVector2D, "getPlusJ", FieldVector2D, (Field,), arg0)
end

function get_positive_infinity(::Type{FieldVector2D}, arg0::Field)
    return jcall(FieldVector2D, "getPositiveInfinity", FieldVector2D, (Field,), arg0)
end

function get_x(obj::FieldVector2D)
    return jcall(obj, "getX", RealFieldElement, ())
end

function get_y(obj::FieldVector2D)
    return jcall(obj, "getY", RealFieldElement, ())
end

function get_zero(::Type{FieldVector2D}, arg0::Field)
    return jcall(FieldVector2D, "getZero", FieldVector2D, (Field,), arg0)
end

function hash_code(obj::FieldVector2D)
    return jcall(obj, "hashCode", jint, ())
end

function is_infinite(obj::FieldVector2D)
    return jcall(obj, "isInfinite", jboolean, ())
end

function is_na_n(obj::FieldVector2D)
    return jcall(obj, "isNaN", jboolean, ())
end

function negate(obj::FieldVector2D)
    return jcall(obj, "negate", FieldVector2D, ())
end

function normalize(obj::FieldVector2D)
    return jcall(obj, "normalize", FieldVector2D, ())
end

function orientation(::Type{FieldVector2D}, arg0::FieldVector2D, arg1::FieldVector2D, arg2::FieldVector2D)
    return jcall(FieldVector2D, "orientation", RealFieldElement, (FieldVector2D, FieldVector2D, FieldVector2D), arg0, arg1, arg2)
end

function scalar_multiply(obj::FieldVector2D, arg0::RealFieldElement)
    return jcall(obj, "scalarMultiply", FieldVector2D, (RealFieldElement,), arg0)
end

function scalar_multiply(obj::FieldVector2D, arg0::jdouble)
    return jcall(obj, "scalarMultiply", FieldVector2D, (jdouble,), arg0)
end

function subtract(obj::FieldVector2D, arg0::FieldVector2D)
    return jcall(obj, "subtract", FieldVector2D, (FieldVector2D,), arg0)
end

function subtract(obj::FieldVector2D, arg0::RealFieldElement, arg1::FieldVector2D)
    return jcall(obj, "subtract", FieldVector2D, (RealFieldElement, FieldVector2D), arg0, arg1)
end

function subtract(obj::FieldVector2D, arg0::RealFieldElement, arg1::Vector2D)
    return jcall(obj, "subtract", FieldVector2D, (RealFieldElement, Vector2D), arg0, arg1)
end

function subtract(obj::FieldVector2D, arg0::Vector2D)
    return jcall(obj, "subtract", FieldVector2D, (Vector2D,), arg0)
end

function subtract(obj::FieldVector2D, arg0::jdouble, arg1::FieldVector2D)
    return jcall(obj, "subtract", FieldVector2D, (jdouble, FieldVector2D), arg0, arg1)
end

function subtract(obj::FieldVector2D, arg0::jdouble, arg1::Vector2D)
    return jcall(obj, "subtract", FieldVector2D, (jdouble, Vector2D), arg0, arg1)
end

function to_array(obj::FieldVector2D)
    return jcall(obj, "toArray", Vector{RealFieldElement}, ())
end

function to_string(obj::FieldVector2D)
    return jcall(obj, "toString", JString, ())
end

function to_string(obj::FieldVector2D, arg0::NumberFormat)
    return jcall(obj, "toString", JString, (NumberFormat,), arg0)
end

function to_vector2d(obj::FieldVector2D)
    return jcall(obj, "toVector2D", Vector2D, ())
end


text:
using Revise
using ADCME
using ADCMEKit
using NNFEM
using PyPlot
using ProgressMeter
using Printf
using MAT
using Clustering
using JLD2

close("all")
include("load_domain_function.jl")

NT = 100
Δt = 30/NT

mesh = "crack_wider.msh"
c1 = (3.5, 0.0)
c2 = (4.0, 0.5)
slip_scale = 5.0 * sqrt(2)

#################### solve static elastic response  ####################
domain = load_crack_domain(mesh=mesh, c1=c1, c2=c2, slip_scale=slip_scale)
# domain = load_crack_domain()

H = domain.elements[1].mat[1].H
EBC_func = nothing
FBC_func = nothing
body_func = nothing
globaldata = GlobalData(missing, missing, missing, missing, domain.neqs, EBC_func, FBC_func, body_func)
Fext = compute_external_force(globaldata, domain)
d = LinearStaticSolver(globaldata, domain, domain.state, H, Fext)
sess = Session(); init(sess)
domain.state = run(sess, d)

#################### solve viscoelastic relaxation  ####################
strain = getStrain(domain)
stress = getStress(domain)
state = domain.state

domain = load_crack_domain(mesh=mesh, c1=c1, c2=c2, slip_scale=slip_scale, option="mixed")
# domain = load_crack_domain(option="mixed")

ts = GeneralizedAlphaSolverTime(Δt, NT)
ubd, abd = compute_boundary_info(domain, globaldata, ts)
Fext = compute_external_force(domain, globaldata, ts)

d0 = state
v0 = zeros(2domain.nnodes)
a0  = zeros(2domain.nnodes)
# σ0 = zeros(getNGauss(domain),3)
# ϵ0 = zeros(getNGauss(domain),3)
σ0 = stress
ϵ0 = strain

μ = zeros(getNGauss(domain))
λ = zeros(getNGauss(domain))
η = zeros(getNGauss(domain))
k = 0
for i = 1:domain.neles
  e = domain.elements[i]
  for mat in e.mat
    global k += 1
    μ[k] = mat.μ
    λ[k] = mat.λ
    η[k] = mat.η
  end
end

gnodes = getGaussPoints(domain)
η = 1e10 *( constant(3.0) - constant(1.5) * (gnodes[:,2])/maximum(gnodes[:,2]))

d, v, a, σ, ϵ = ViscoelasticitySolver(
  globaldata, domain, d0, v0, a0, σ0, ϵ0, Δt, NT, μ, λ, η, Fext, ubd, abd
)

sess = Session(); init(sess)
d_, σ_, η_ = run(sess, [d, σ, η])
matwrite("data/dippingfault_viscosity_inversion.mat", Dict(
  "sigma"=> σ_,
  "d" => d_,
  "eta" => η_,
  "d0" => d0,
  "v0" => v0,
  "a0" => a0,
  "epsilon0"=> ϵ0,
  "sigma0" => σ0,
  "NT" => NT,
  "dt" => Δt
))
@save "data/dippingfault_viscosity_inversion.jld2" domain

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
using Plots, Statistics, Random
using Flux, DiffEqFlux, StochasticDiffEq, CUDA, RegNeuralDE, Tracker
using BSON
using Tracker: data
using ColorBrewer
colors = ColorBrewer.palette("Set2", 8)

dataset = BSON.load("data/sde_demo.bson")
sde_data = dataset[:sde_data]
sde_data_vars = dataset[:sde_data_vars]

u0 = reshape(Float32[2.0; 0.0], :, 1)
datasize = 30
tspan = [0.0f0, 1.0f0]
tsteps = range(tspan[1], tspan[2], length = datasize) |> track

plots1, plots2 = [], []
ptime1, ptime2 = [], []
ttime1, ttime2 = [], []
loss1, loss2 = [], []
nfe1, nfe2 = [], []

neuralsde = nothing
REGULARIZE = false

function loss_function(u0, p, i, c = 0.2f0)
    sol, nfe1, nfe2, sv = neuralsde(u0, p)
    means = mean(sol; dims = 3)
    vars = var(sol; dims = 3, mean = means)
    l2_means = mean(abs2, sde_data .- means)
    l2_vars = mean(abs2, sde_data_vars .- vars)
    reg = REGULARIZE ? c * sum(sv.saveval) : 0.0f0
    loss = l2_means + l2_vars
    if i % 50 == 0
        @show i, data(loss), data(l2_means), data(l2_vars), data(reg), nfe1, nfe2
    end
    i == -1 && return (loss, nfe1)
    return loss + reg
end

for i in [5]#, 35, 1]
    for reg in [false, true]
        Random.seed!(i)

        drift_dudt = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2)) |> track
        diffusion_dudt = Dense(2, 2) |> track

        REGULARIZE = reg

        neuralsde = TrackedNeuralDSDE(
            drift_dudt,
            diffusion_dudt,
            [0.0f0, 1.0f0 + eps(Float32)],
            REGULARIZE,
            SOSRI(),
            saveat = tsteps,
            reltol = 3f-1,
            abstol = 3f-1,
        )

        u0_ = repeat(u0, 1, 100)
        loss_function(u0_ |> track, neuralsde.p, 0)
        Tracker.gradient(p -> loss_function(u0_ |> track, p, 0), neuralsde.p)

        opt = AdaBelief(0.01)
        ps = neuralsde.p
        ps_best = ps
        total_time = 0
        for iter = 1:250
            _t = time()
            l, back = Tracker.forward(p -> loss_function(u0_ |> track, p, iter), ps)
            gs = back(1.0f0)[1]
            total_time += time() - _t
            ps_best = ps
            update_parameters!((ps,), (gs,), opt)
        end

        u0_ = repeat(u0, 1, 100)

        begin
            neuralsde(u0_ |> track, ps_best)
            ptime = @belapsed neuralsde($u0_ |> track, $ps_best)[1]
            preds = neuralsde(u0_ |> track, ps_best)[1]
            preds = preds |> untrack
            means, vars = mean(preds; dims = 3), var(preds, dims = 3)

            plt = plot(
                tsteps |> untrack,
                means[1, :, 1],
                ribbon = vars[1, :, 1],
                color = colors[1],
                fillalpha = 0.2,
                label = "Pred: Dim 1",
                linewidth = 3,
                legend = :bottomleft,
            )
            plot!(
                plt,
                tsteps |> untrack,
                means[2, :, 1],
                ribbon = vars[2, :, 1],
                color = colors[2],
                fillalpha = 0.2,
                label = "Pred: Dim 2",
                linewidth = 3,
            )
            scatter!(
                plt,
                tsteps |> untrack,
                sde_data[1, :],
                yerror = sde_data_vars[1, :],
                label = "Data: Dim 1",
                color = colors[1],
                msc = colors[1],
            )
            scatter!(
                plt,
                tsteps |> untrack,
                sde_data[2, :],
                yerror = sde_data_vars[2, :],
                label = "Data: Dim 2",
                color = colors[2],
                msc = colors[2],
            )

            if REGULARIZE
                push!(plots1, plt)
            else
                push!(plots2, plt)
            end
            display(plt)
        end

        if REGULARIZE
            push!(ttime1, total_time)
            l, n = loss_function(u0_ |> track, ps_best, -1)
            push!(loss1, l)
            push!(ptime1, ptime)
            push!(nfe1, n)

            @show ttime1[end], l, n, ptime[end]
        else
            push!(ttime2, total_time)
            l, n = loss_function(u0_ |> track, ps_best, -1)
            push!(loss2, l)
            push!(ptime2, ptime)
            push!(nfe2, n)

            @show ttime2[end], l, n, ptime2[end]
        end

        GC.gc(true)
        GC.gc(true)
        GC.gc(true)
    end
end

@show mean(loss1), std(loss1), mean(loss2), std(loss2)
@show mean(nfe1), std(nfe1), mean(nfe2), std(nfe2)
@show mean(ptime1), std(ptime1), mean(ptime2), std(ptime2)
@show mean(ttime1), std(ttime1), mean(ttime2), std(ttime2)

title!(plots2[1], "Unregularized Neural SDE")
title!(plots1[1], "Regularized Neural SDE")
xlabel!(plots1[1], "t")
plt = plot(plots2[1], plots1[1], layout = (2, 1))

savefig(plt, "spiral_sde.pdf")


text:
# The neuron meshes are provided in a separate git repository. Make sure to clone
#
#     https://github.com/SpinDoctorMRI/RealNeuronMeshes.git
#
# in order to run SpinDoctor simulations on real neuron geometries. You may need to manually
# unzip the mesh files of interest, as they are stored in text format.
#
# The file name is given by `prefix * name * postfix` (uncomment the desired file parts).

neuron_dir = "RealNeuronMeshes/volume_meshes"
isdir(neuron_dir) || error("""Cannot find neuron meshes.
                           Run `git clone https://github.com/SpinDoctorMRI/RealNeuronMeshes.git`
                           to download the mesh files.""")

## Choose whether to use whole neuron or certain parts
# prefix, postfix = "separated_", "_soma"
# prefix, postfix = "separated_", "_dendrites"
prefix, postfix = "separated_", "_dendrites_1"
# prefix, postfix = "separated_", "_dendrites_2"
# prefix, postfix = "", ""

## Neuron name
# name = "spindles/03a_spindle2aFI"
# name = "spindles/03a_spindle6aFI"
# name = "spindles/03b_spindle4aACC"
# name = "spindles/03b_spindle5aACC"
# name = "spindles/03b_spindle6aACC"
# name = "spindles/03b_spindle7aACC"
# name = "spindles/04b_spindle3aFI"
# name = "spindles/05b_spindle5aFI"
name = "spindles/06b_spindle8aACC"
# name = "spindles/07b_spindle9aACC"
# name = "spindles/08a_spindle13aACC"
# name = "spindles/09o_spindle7aFI"
# name = "spindles/09o_spindle8aFI"
# name = "spindles/10a_spindle18aACC"
# name = "spindles/12a_spindle19aACC"
# name = "spindles/12o_spindle9aFI"
# name = "spindles/13o_spindle10aFI"
# name = "spindles/15o_spindle12aFI"
# name = "spindles/16o_spindle13aFI"
# name = "spindles/19o_spindle14aFI"
# name = "spindles/21o_spindle15aFI"
# name = "spindles/23o_spindle16aFI"
# name = "spindles/25o_spindle17aFI"
# name = "spindles/26o_spindle18aFI"
# name = "spindles/27o_spindle19aFI"
# name = "spindles/28o_spindle20aFI"
# name = "spindles/28o_spindle21aFI"
# name = "spindles/29o_spindle22aFI"
# name = "spindles/30o_spindle23aFI"
# name = "pyramidals/02a_pyramidal2aFI"
# name = "pyramidals/02b_pyramidal1aACC"
# name = "pyramidals/02b_pyramidal1aFI"
# name = "pyramidals/03a_pyramidal9aFI"
# name = "pyramidals/03b_pyramidal2aACC"
# name = "pyramidals/03b_pyramidal3aACC"
# name = "pyramidals/03b_pyramidal3aFI"
# name = "pyramidals/03b_pyramidal4aFI"
# name = "pyramidals/03b_pyramidal9aFI"
# name = "pyramidals/04a_pyramidal4aACC"
# name = "pyramidals/04a_pyramidal5aACC"
# name = "pyramidals/04b_pyramidal5aFI"
# name = "pyramidals/04b_pyramidal6aACC"
# name = "pyramidals/04b_pyramidal6aFI"
# name = "pyramidals/04b_pyramidal7aACC"
# name = "pyramidals/05a_pyramidal10aACC"
# name = "pyramidals/05a_pyramidal8aACC"
# name = "pyramidals/05b_pyramidal7aFI"
# name = "pyramidals/05b_pyramidal8aFI"
# name = "pyramidals/05b_pyramidal9aACC"
# name = "pyramidals/06a_pyramidal11aACC"
# name = "pyramidals/06b_pyramidal10aFI"
# name = "pyramidals/06b_pyramidal12aACC"
# name = "pyramidals/07a_pyramidal13aACC"
# name = "pyramidals/07b_pyramidal14aACC"
# name = "pyramidals/08o_pyramidal11aFI"
# name = "pyramidals/10a_pyramidal15aACC"
# name = "pyramidals/11a_pyramidal16aACC"
# name = "pyramidals/11o_pyramidal12aFI"
# name = "pyramidals/17o_pyramidal13aFI"
# name = "pyramidals/18o_pyramidal14aFI"
# name = "pyramidals/20o_pyramidal15aFI"
# name = "pyramidals/22o_pyramidal16aFI"
# name = "pyramidals/24o_pyramidal17aFI"
# name = "pyramidals/25o_pyramidal18aFI"
# name = "pyramidals/31o_pyramidal19aFI"

## Resulting neuron filename
name = prefix * name * postfix

# Floating point type for simulations
T = Float64

# Geometrical setup
setup = NeuronSetup{T}(;
    name,
    meshdir = neuron_dir,
    ecs_shape = :no_ecs,
    ecs_ratio = 0.3,
    # refinement = 0.5,
)

# Get compartimentalized coefficient vectors
coeffs = coefficients(
    setup;
    D = (; neuron = 0.002 * I(3), ecs = 0.002 * I(3)),
    T₂ = (; neuron = Inf, ecs = Inf),
    ρ = (; neuron = 1.0, ecs = 1.0),
    κ = (; neuron_ecs = 1e-4, in = 0.0, neuron = 0.0, ecs = 0.0),
    γ = 2.67513e-4,
)


text:
using SolidStateDetectors.ConstructiveSolidGeometry: CSGUnion, CSGIntersection, CSGDifference, CSG_dict,
        parse_r_of_primitive, parse_height_of_primitive, parse_translate_vector, parse_rotation_matrix
using SolidStateDetectors: construct_units
import SolidStateDetectors.ConstructiveSolidGeometry: Geometry, Dictionary
using SolidStateDetectors.ConstructiveSolidGeometry: AbstractGeometry, Transformations,
    CSGUnion, CSGDifference, CSGIntersection, Cone, CartesianVector, HexagonalPrism
using DataStructures: OrderedDict
using IntervalSets
using Rotations
using StaticArrays
using Unitful
using JSON
using YAML

#old implementation of Constructive Solid Geometries:
function Geometry(::Type{T}, t::Type{CSGUnion}, dict::Union{Dict, OrderedDict}, input_units::NamedTuple, transformations::Transformations{T}) where {T}
    @assert haskey(dict, "parts") "Please specify 'parts' of the '$(dict["type"])'."
    sum( map(x-> Geometry(T, x, input_units, transformations), dict["parts"]) )
end

function Geometry(::Type{T}, ::Type{CSGIntersection}, dict::Union{Dict, OrderedDict}, input_units::NamedTuple, transformations::Transformations{T}) where {T}
    @assert haskey(dict, "parts") "Please specify 'parts' of the '$(dict["type"])'."
    parts = map(x-> Geometry(T, x, input_units, transformations), dict["parts"])
    reduce(&, parts)
end

function Geometry(::Type{T}, ::Type{CSGDifference}, dict::Union{Dict, OrderedDict}, input_units::NamedTuple, transformations::Transformations{T}) where {T}
    @assert haskey(dict, "parts") "Please specify 'parts' of the '$(dict["type"])'."
    Geometry(T, dict["parts"][1], input_units, transformations) - sum( map(x-> Geometry(T, x, input_units, transformations), dict["parts"][2:end]) )
end

function Geometry(::Type{T}, dict::Union{Dict, OrderedDict}, input_units::NamedTuple, transformations::Transformations{T}) where {T}
    if haskey(dict, "translate")
        length_unit = input_units.length
        t::CartesianVector{T} = parse_translate_vector(T, dict["translate"], length_unit)
        gdict::Dict{Any,Any} = filter(p -> first(p) != "translate", dict)
        return Geometry(T, gdict, input_units, transformations) + t
    end
    Geometry(T, CSG_dict[dict["type"]], dict, input_units, transformations)
end

function Geometry(::Type{T}, filename::String, input_units::NamedTuple, transformations::Transformations{T}) where {T}
    @assert isfile(filename) "The given filename '$(filename)' does not lead to a valid file."
    dict = if endswith(filename, ".json")
        JSON.parsefile(filename)
    elseif endswith(filename, ".yaml")
        YAML.load_file(filename)
    else
        @error "Only JSON and YAML formats are supported at the moment."
    end
    Geometry(T, dict, input_units, transformations::Transformations{T})
end

# UPDATE DUE TO BREAKING CHANGES TO TUBE/CONE: CENTERED AROUND ZERO
# AND REDEFINITION OF HEXAGONALPRISM
function update_primitives!(dict::AbstractDict, input_units::NamedTuple)
    dict_keys = keys(dict)
    #Translate all Tubes and Cone that are defined via "h" and where h is not zero
    if "type" in dict_keys && dict["type"] == "HexagonalPrism"
        rotZ = haskey(dict, "rotZ") ? dict["rotZ"] : 0
        rotZ -= input_units.angle == u"°" ? 30 : π/6
        rotZ = string(rotZ * input_units.angle)
        delete!(dict, "rotZ")
        if rotZ != 0 dict["rotate"] = Dict("Z" => rotZ) end
    elseif "type" in dict_keys && dict["type"] in ["tube", "cone"] && "h" in dict_keys && dict["h"] != 0
        #add translate vector with z=0 if it does not exist
        if !("translate" in dict_keys) dict["translate"] = Dict{String,Any}("z" => 0) end
        dict["translate"]["z"] = dict["translate"]["z"] + Float64(dict["h"] / 2)
        if dict["translate"]["z"] == 0 delete!(dict["translate"], "z") end
        if length(keys(dict["translate"])) == 0 delete!(dict, "translate") end
    end
    for k in dict_keys
        if typeof(dict[k]) <: AbstractDict
            dict[k] = update_primitives!(dict[k], input_units)
        elseif typeof(dict[k]) <: Array
            for i in eachindex(dict[k])
                if typeof(dict[k][i]) <: AbstractDict
                    dict[k][i] = update_primitives!(dict[k][i], input_units)
                end
            end
        end
    end
    dict
end

function update_geometry!(g::AbstractDict, T::DataType, input_units = (length = NoUnits, angle = NoUnits))
    transformations::Transformations{T} = (rotation = one(SMatrix{3, 3, T, 9}), translation = zero(CartesianVector{T}))
    g["geometry"] = Dictionary(Geometry(T, g["geometry"], input_units, transformations))
    return g
end

function restructure_config_file_dict!(config_file_dict::AbstractDict, T::DataType = Float64)

    input_units = construct_units(config_file_dict)
    update_primitives!(config_file_dict, input_units)
    update_units = (length = NoUnits, angle = input_units.angle)

    t = broadcast(obj -> obj["type"], config_file_dict["objects"])

    semiconductors = config_file_dict["objects"][findall(t .== "semiconductor")]
    @assert length(semiconductors) == 1 "Config files need exactly one semiconductor per detector. This one has $(length(semiconductors)) semiconductors."
    semiconductor = semiconductors[1]

    semiconductor = begin
        delete!(semiconductor, "type")
        if "charge_density_model" in keys(semiconductor)
            semiconductor["impurity_density"] = semiconductor["charge_density_model"]
            delete!(semiconductor, "charge_density_model")
        end
        update_geometry!(semiconductor, T, update_units)
        semiconductor
    end

    contacts = [
        begin
        delete!(contact, "type")
        if "channel" in keys(contact)
            contact["id"] = contact["channel"]
            delete!(contact, "channel")
        end
        update_geometry!(contact, T, update_units)
        contact
        end
        for contact in config_file_dict["objects"][findall(t .== "contact")]
    ]

    passives = [
        begin
        delete!(passive, "type")
        update_geometry!(passive, T, update_units)
        passive
        end
        for passive in config_file_dict["objects"][findall(t .== "passive")]
    ]

    virtual_drift_volumes = [
        begin
        delete!(virtual_volume, "type")
        update_geometry!(virtual_volume, T, update_units)
        virtual_volume
        end
        for virtual_volume in config_file_dict["objects"][findall(t .== "virtual_drift_volume")]
    ]

    config_file_dict["detectors"] = []
    push!(config_file_dict["detectors"], OrderedDict{String, Any}("semiconductor" => semiconductor) )
    if length(contacts) > 0 config_file_dict["detectors"][1]["contacts"] = contacts end
    if length(virtual_drift_volumes) > 0 config_file_dict["detectors"][1]["virtual_drift_volumes"] = virtual_drift_volumes end
    if length(passives) > 0 config_file_dict["surroundings"] = passives end

    delete!(config_file_dict, "objects")
    config_file_dict

end

"""
    update_config_file(filename::String)::String

Updates the format of config files from v0.5.3 to v0.6.
Currently supported formats for the config file: .json, .yaml.
Returns the new filename (or an empty string if file could not be converted).
"""
function update_config_file(filename::String)::String
    @assert isfile(filename) "'$(filename)' does not exist!"

    if endswith(filename, ".json")
        data = JSON.parsefile(filename, dicttype = OrderedDict)
        if !haskey(data, "objects")
            @warn "The configuration file does not have the key 'objects'.\n"*
                  "Assuming that this configuration file has already been updated. Skipping..."
            return ""
        end
        restructure_config_file_dict!(data)
        new_filename = replace(filename, ".json" => "_updated.json")
        open(new_filename,"w") do f
            JSON.print(f, data, 4)
        end
    elseif endswith(filename, ".yaml")
        data = YAML.load_file(filename, dicttype = OrderedDict)
        if !haskey(data, "objects")
            @warn "The configuration file does not have the key 'objects'.\n"*
                  "Assuming that this configuration file has already been updated. Skipping..."
            return ""
        end
        restructure_config_file_dict!(data)
        new_filename = replace(filename, ".yaml" => "_updated.yaml")
        YAML.write_file(new_filename, data)
    else
        error("Currently only .json and .yaml files are supported.")
        new_filename = ""
    end
    new_filename
end


text:
# temperature correlation functions for PressureDrop package.

"""
`Shiu_Beggs_relaxationfactor(<arguments>)`

Generates the relaxation factor, A, needed for the Ramey method, for underspecified conditions.

This correlation assumes flow has stabilized and that the transient time component f(t) is not changing.

# Arguments
All arguments are in U.S. field units.

- `q_o`: oil rate in stb/d
- `q_w`: water rate in stb/d
- `APIoil`: API oil gravity
- `sg_water`: water specific gravity
- `GLR`: gas:liquid ratio in scf/stb
- `sg_gas`: gas specific gravity
- `id`: flow path inner diameter in inches
- `WHP`: wellhead/outlet absolute pressure in **psig**
"""
function Shiu_Beggs_relaxationfactor(q_o, q_w, GLR, APIoil, sg_water, sg_gas, id, WHP)

    sg_oil = 141.5/(APIoil + 131.5)
    q_g = (q_o + q_w) * GLR
    w = 1/86400 * (350*(q_o*sg_oil + q_w*sg_water) + 0.0764*q_g*sg_gas) #mass flow rate in lb/sec
    ρ_l_sc = 62.4 * mixture_properties_simple(q_o, q_w, sg_oil, sg_water)

    # Original Shiu-Beggs coefficients for known WHP:
    C_0 = 0.0063
    C_1 = 0.4882
    C_2 = 2.9150
    C_3 = -0.3476
    C_4 = 0.2219
    C_5 = 0.2519
    C_6 = 4.7240

    return C_0 * w^C_1 * ρ_l_sc^C_2 * id^C_3 * WHP^C_4 * APIoil^C_5 * sg_gas^C_6 #relaxation distance
end

"""
`Ramey_wellboretemp(z, inclination, T_bh, A, G_g = 1.0)`

Estimates wellbore temp using Ramey 1962 method.

# Arguments
- `z`: true vertical depth **from the bottom of the well**, ft
- `T_bh`: bottomhole temperature, °F
- `A`: relaxation factor
- `G_g = 1.0`: geothermal gradient in °F per 100 ft of true vertical depth
"""
function Ramey_temp(z, T_bh, A, G_g = 1.0)

    g_g = G_g / 100

    return T_bh - g_g * z + A * g_g * (1 - exp(-z / A))
end

"""
`linear_wellboretemp(;WHT, BHT, wellbore::Wellbore)`

Linear temperature profile from a wellhead temperature and bottomhole temperature in °F for a Wellbore object.

Interpolation is based on true vertical depth of the wellbore, not md.
"""
function linear_wellboretemp(;WHT, BHT, wellbore::Wellbore,
                            kwargs...) #catch extra arguments from a WellModel for convenience

    temp_slope = (BHT - WHT) / maximum(wellbore.tvd)

    return [WHT + depth * temp_slope for depth in wellbore.tvd]
end

"""
`Shiu_wellboretemp(<named arguments>)`

Wrapper to compute temperature profile for a Wellbore object using Ramey correlation with Shiu relaxation factor correlation.

# Arguments
- `BHT`: bottomhole temperature in °F
- `geothermal_gradient = 1.0`: geothermal gradient in °F per 100 feet
- `wellbore::Wellbore`: Wellbore object to use as reference for segmentation, inclination, and
- `q_o`: oil rate in stb/d
- `q_w`: water rate in stb/d
- `GLR`: gas:liquid ratio in scf/day
- `APIoil`: oil gravity
- `sg_water`: water specific gravity
- `sg_gas`: gas specific gravity
- `WHP`: wellhead/outlet absolute pressure in **psig**
"""
function Shiu_wellboretemp(;BHT, geothermal_gradient = 1.0, wellbore::Wellbore, q_o, q_w, GLR, APIoil, sg_water, sg_gas, WHP,
                            kwargs...) #catch extra arguments from a WellModel for convenience

    id_avg = sum(wellbore.id)/length(wellbore.id)
    A = Shiu_Beggs_relaxationfactor(q_o, q_w, GLR, APIoil, sg_water, sg_gas, id_avg, WHP) #use average inner diameter to calculate relaxation factor
    TD = maximum(wellbore.tvd)
    depths = TD .- wellbore.tvd
    temp_profile = [Ramey_temp(z, BHT, A, geothermal_gradient) for z in depths]

    return temp_profile
end


text:
using BenchmarkTools

include("RegionTutorialBenchmarks.jl")
include("getindex.jl")

const SUITE = BenchmarkGroup()
SUITE["one_region"] = @benchmarkable run_oneregion()
SUITE["two_regions"] = @benchmarkable run_tworegion()
SUITE["getindex"] = @benchmarkable run_getindex()


text:
# Spinnaker.jl: wrapper for FLIR/Point Grey Spinnaker SDK
# Copyright (C) 2019 Samuel Powell

# examples/exposure.jl: example to configure manual exposure
using Spinnaker

camlist = CameraList()
cam = camlist[0]
triggersource!(cam, "Software")
triggermode!(cam, "On")
gain!(cam, 0)
acquisitionmode!(cam, "Continuous")
start!(cam)
for expval in 0:1e5:2e6
    expact = exposure!(cam, expval)
    @info "Exposure set to $(expact/1e6)s"
    trigger!(cam)
    saveimage(cam, joinpath(@__DIR__, "exposure_$(expval/1e6)s.png"), spinImageFileFormat(6))
    @info "Image saved"
end
stop!(cam)


text:
module sampleproject

export foo

"""
    foo()

Return the foo.

# Example

```jldoctest
julia> using sampleproject

julia> foo()
"bar"
```
"""
foo() = "bar"

end # module

max_stars_repo_name: nluetts/sampleproject

max_issues_repo_name: nluetts/sampleproject

max_forks_repo_name: nluetts/sampleproject


text:
using Pointy
using Test
using BenchmarkTools

@testset "Pointy.jl" begin
    #ideal hex lattice, 60° between basal plane vectors
    u = [1,0,0]; v = [.5,√3/2,0]; w = [0,0,√(8/3)];
    @test length(pointgroup_basic(u,v,w))==24
    #ideal hex lattice, 120° between basal plane vectors
    a = -u + v; b = v; c = w
    @test length(pointgroup_basic(a,b,c))==24
    # simple cubic lattice in unique orientation
    d,e,f = threeDrotation([1, 0, 0],[0, 1, 0],[0, 0, 1],π/3,π/5,π/7)
    @test length(pointgroup_basic(d,e,f))==48
    # hexagonal lattice in unique orientation
    g,h,i = threeDrotation([1,0,0], [.5,√3/2,0],[0,0,√(8/3)],π/7,π/11,π/3)
    @test length(pointgroup_basic(g,h,i))==24
    # Add a timing check against basic and "fast" make sure fast is fast
    # hexagonal case
    slow = @belapsed pointgroup_basic($g,$h,$i)
    fast = @belapsed pointgroup($g,$h,$i)
    @test  slow/fast > 25
    # cubic case
    slow = @belapsed pointgroup_basic($d,$e,$f)
    fast = @belapsed pointgroup($d,$e,$f)
    @test  slow/fast > 54
    # tetragonal case
    d,e,f = threeDrotation([1.1, 0, 0],[0, 1, 0],[0, 0, 1],π/3,π/5,π/7)
    slow = @belapsed pointgroup_basic($d,$e,$f)
    fast = @belapsed pointgroup($d,$e,$f)
    @test  slow/fast > 60
    # orthorhombic case
    d,e,f = threeDrotation([1.1, 0, 0],[0, .9, 0],[0, 0, 1],π/3,π/5,π/7)
    slow = @belapsed pointgroup_basic($d,$e,$f)
    fast = @belapsed pointgroup($d,$e,$f)
    @test  slow/fast > 62
    u = [1,1,2]; v = [1,2,1]; w = [2,1,1];
    u,v,w = threeDrotation(u,v,w,π/3,π/5,π/7)
    @test length(pointgroup(u,v,w))==12
end


text:
# using Statistics
using StatsBase:sample
# using Base.Threads:@threads
using BenchmarkTools

n_obs = Int(1e6)
n_vars = 100
n_bins = 64
K = 3
𝑖 = collect(1:n_obs)
δ = rand(n_obs, K)
hist = zeros(K, n_bins, n_vars);
X_bin = sample(UInt8.(1:n_bins), n_obs * n_vars);
X_bin = reshape(X_bin, n_obs, n_vars);

function iter_1(X_bin, hist, δ, 𝑖)
    hist .= 0.0
    @inbounds for i in 𝑖
        @inbounds for k in 1:3
            hist[k, X_bin[i,1], 1] += δ[i,k]
        end
    end
end

𝑖_sample = sample(𝑖, Int(n_obs / 2), ordered=true)

@time iter_1(X_bin, hist, δ, 𝑖_sample)
@btime iter_1($X_bin, $hist, $δ, $𝑖_sample)

function iter_2(X_bin, hist, δ, 𝑖)
    hist .= 0.0
    @inbounds @simd for i in CartesianIndices(𝑖)
        @inbounds @simd for k in 1:3
            hist[k, X_bin[𝑖[i],1], 1] += δ[𝑖[i],k]
        end
    end
end

𝑖_sample = sample(𝑖, Int(n_obs / 2), ordered=true)

@time iter_2(X_bin, hist, δ, 𝑖_sample)
@btime iter_2($X_bin, $hist, $δ, $𝑖_sample)

# slower
δ = rand(K, n_obs)
hist = zeros(K, n_bins, n_vars);
X_bin = sample(UInt8.(1:n_bins), n_obs * n_vars);
X_bin = reshape(X_bin, n_obs, n_vars);

function iter_1(X_bin, hist, δ, 𝑖)
    hist .= 0.0
    @inbounds for i in 𝑖
        @inbounds for k in 1:3
            hist[k, X_bin[i,1], 1] += δ[k,i]
        end
    end
end

𝑖_sample = sample(𝑖, Int(n_obs / 2), ordered=true)

@time iter_1(X_bin, hist, δ, 𝑖_sample)
@btime iter_1($X_bin, $hist, $δ, $𝑖_sample)


text:
# benchmark for binomial samplers

using Sampling
import Sampling: BinomialRmathSampler

function bench_binomial_sampler(s, params::(Int, Float64), ns::Int)
    # warm up
    for i = 1:10; rand(s); end
    # measure
    et = @elapsed for i = 1:ns; rand(s); end
    mps = (ns * 1.0e-6) / et
    @printf("%-20s on (%4d, %4.2f)  =>  %9.4f MPS\n", string(typeof(s)), params[1], params[2], mps)
end

for n in 2.^(1:7), p in [0.3, 0.5, 0.9]
    bench_binomial_sampler(BinomialGeomSampler(n, p), (n, p), 10^5)
end
println()

for n in (20 * 2.^(1:8)), p in [0.3, 0.6]
    bench_binomial_sampler(BinomialTPESampler(n, p), (n, p), 10^5)
end
println()

for n in 2.^(1:12), p in [0.3, 0.5, 0.9]
    bench_binomial_sampler(BinomialPolySampler(n, p), (n, p), 10^5)
end
println()

for n in 2.^(1:12), p in [0.3, 0.5, 0.9]
    bench_binomial_sampler(BinomialRmathSampler(n, p), (n, p), 10^5)
end
println()

for n in 2.^(1:12), p in [0.3, 0.5, 0.9]
    bench_binomial_sampler(BinomialAliasSampler(n, p), (n, p), 10^5)
end
println()


text:
julia> macro do_while(condition, block)
           quote
               let
                   $block
                   while $condition
                       $block
                   end
               end
           end |> esc
       end
@do_while (macro with 1 method)

julia> i = 0
0

julia> @do_while i % 6 ≠ 0 begin
           @show i
           i += 1
       end
i = 0
i = 1
i = 2
i = 3
i = 4
i = 5

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
"""
    get_rel_freqs(sample::ScreenSample,
                       callable_bc::DataFrame)::DataFrame

Computes the relative frequencies of barcodes as compared to the median
frequency of occurrence for all the negative controls. The frequencies are
computed initially by dividing by the total number of reads for that sample.
"""
function get_rel_freqs(sample::ScreenSample,
                       callable_bc::DataFrame)::DataFrame
    # construct a DataFrame from the dictionary
    df = DataFrame(id=collect(keys(sample.counts)),
                   counts=float(collect(values(sample.counts))),
                   sample_id=sample.filename)

    total_reads = sum(df[:counts])

    # use all rows from barcodes even if they are missing
    tmp = join(callable_bc, df, on=:id, kind=:left)
    # set missing sgRNAs to the pseudocount value
    tmp[isna.(tmp[:counts]), :counts] = 0.5
    # compute frequencies by normalizing by the total number of reads
    tmp[:freqs] = tmp[:counts] ./ total_reads

    negcontrols = @from i in tmp begin
        @where contains(get(i.gene), "NonTargetingControlGuide") && i.freqs > 0
        @select i
        @collect DataFrame
    end
    (size(negcontrols)[1] < 100) && warn("Very few negative controls detected, this might be problem")

    negcontrol_med = median(negcontrols[:freqs])

    tmp[:rel_freqs] = tmp[:freqs] ./ negcontrol_med

    tmp
end

"""
    get_log2fc_pvals(exp::DataFrame, control::DataFrame)

Given a treatment dataframe `exp` and a control dataframe `control` calculates
the log 2 fold changes on a per guide level and then uses a Mann Whitney test
to compare this the negative control guides to determine the significance on
a per gene level.
"""
function get_log2fc_pvals(exp::DataFrame, control::DataFrame)

    # join controls dataframe with treatment and compute log2fc between them
    x = calc_log2fc(exp, control)

    # get all log2fc values for the negative controls
    negcontrols_fc = @from i in x begin
        @where contains(get(i.gene), "NonTargetingControlGuide") && isfinite(get(i.log2fc))
        @select get(i.log2fc)
        @collect
    end

    z = compute_gene_stats(x, negcontrols_fc)

    sort!(z, cols=:pval, rev=true)
    z[:ranking] = collect(1:size(z)[1])
    z

end

"""
    calc_log2fc(exp::DataFrame, control::DataFrame)

Combine `control` with `exp` and compute the log2 fold change between their
`:rel_freqs` columns
"""
function calc_log2fc(exp::DataFrame, control::DataFrame)
    @from i in control begin
        @join j in exp on i.id equals j.id
        @select {
                    i.id,
                    log2fc=Nullable(log2(get(j.rel_freqs)/get(i.rel_freqs))),
                    rel_freq1=i.rel_freqs,
                    rel_freq2=j.rel_freqs,
                    i.gene,
                    i.sample_id
                }
        @collect DataFrame
    end
end

"""
    compute_gene_stats(df::DataFrame, negcontrols_fc::Vector{Float64};
                       column::Symbol=:log2fc)

Calculates pvalues and medians on a per gene basis for the given DataFrame `df`
using the specified `column`. P-values are computed using a Mann-Whitney test of
the gene values against the values provided in `negcontrols_fc` for the negative
controls.
"""
function compute_gene_stats(df::DataFrame, negcontrols_fc::Vector{Float64};
                            column::Symbol=:log2fc)
    res = @from i in df begin
        @where isfinite(_try_get(getfield(i, column)))
        @group i by i.gene into g
        @let test = -log10(pvalue(MannWhitneyUTest(map(j->_try_get(getfield(j, column)), g),
                                                   negcontrols_fc)))
        @let med = median(map(j->_try_get(getfield(j, column)), g))
        @select {
            gene = _try_get(g.key),
            med_val = med,
            pval = test,
            prod = abs(med * test),
            n_guides = length(g)
        }
        @collect DataFrame
    end
    rename!(res, :med_val, Symbol("med_", column))
    res
end

"""
    fit_lines(t0::DataFrame, t1::DataFrame, t2::DataFrame)

Fit a line to the relative frequencies of each guide at `t0`, `t1`, and `t2`
by solving the linear least squares equation for an overdetermined system.
"""
fit_lines(t0::DataFrame, t1::DataFrame, t2::DataFrame) = fit_lines(join(join(t0, t1, on=:id), t2, on=:id))

function fit_lines(data::DataFrame)
    X = hcat(ones(3, 1), [0,1,2])
    tmp = @from i in data begin
        # Solve the system of linear equations for the betahats, which give you the slope
        # and the intercept
        # https://en.wikipedia.org/wiki/Linear_least_squares_(mathematics)#The_general_problem
        @let betahat = X \ [get(i.rel_freqs), get(i.rel_freqs_1), get(i.rel_freqs_2)]
        @select {
            id=get(i.id),
            gene=get(i.gene),
            intercept=betahat[1],
            slope=betahat[2],
            t0_count=get(i.counts),
            t2_count=get(i.counts_2),
            t0_rel_freqs=get(i.rel_freqs),
            t1_rel_freqs=get(i.rel_freqs_1),
            t2_rel_freqs=get(i.rel_freqs_2)
        }
        @collect DataFrame
    end
end


text:
using Random

# strings
x = "ppzz!#"
r = RObject(x)
@test isa(r,RObject{StrSxp})
@test length(r) == 1
@test rcopy(r) == x
@test isascii(r)
@test isascii(rcopy(r))
@test rcopy(r[1]) == x
@test isa(r[1],RObject{CharSxp})

x = "aaα₁"
r = RObject(x)
@test isa(r,RObject{StrSxp})
@test length(r) == 1
@test rcopy(r) == x
@test !isascii(rcopy(r))
@test rcopy(Symbol, r) == Symbol(x)

v = ["ap","xα⟩","pp"]
r = RObject(v)
@test isa(r,RObject{StrSxp})
@test length(r) == length(v)
@test rcopy(r) == v
@test rcopy(r[1]) == v[1]
@test rcopy(Array, r)[1] == v[1]
@test rcopy(Vector, r)[1] == v[1]
@test rcopy(Array{String}, r)[2] == v[2]
@test rcopy(Array{Symbol}, r)[2] == Symbol(v[2])
@test isa(RCall.sexp(RCall.RClass{:character}, :a), Ptr{StrSxp})

s = SubString{String}["a","b"]
r = RObject(s)
@test isa(r,RObject{StrSxp})
@test length(r) == length(s)
@test rcopy(r) == s
@test rcopy(r[1]) == s[1]

# logical
r = RObject(false)
@test isa(r,RObject{LglSxp})
@test length(r) == 1
@test rcopy(r) === false
@test r[1] === convert(Cint,0)

v = bitrand(10)
r = RObject(v)
@test isa(r,RObject{LglSxp})
@test length(r) == length(v)
@test size(r) == size(v)
@test rcopy(r) == v
@test rcopy(Vector, r) == v
@test rcopy(Array, r) == v
@test r[1] === convert(Cint,v[1])

# integer
@test rcopy(Int, R"TRUE") == 1
@test rcopy(Int, R"1") == 1
@test rcopy(Array{Int}, R"c(1,2,3)") == [1,2,3]

x = 7
r = RObject(x)
@test isa(r,RObject{IntSxp})
@test length(r) == 1
@test size(r) == (1,)
@test rcopy(r) === x
@test r[1] === convert(Cint,x)

v = -7:3
r = RObject(v)
@test isa(r,RObject{IntSxp})
@test length(r) == length(v)
@test size(r) == size(v)
@test isa(rcopy(r), Vector{Int})
@test rcopy(r) == collect(v)
@test rcopy(Vector, r) == collect(v)
@test r[1] === convert(Cint,v[1])
@test r[3] === convert(Cint,v[3])
r[2] = -100
@test r[2] === convert(Cint,-100)
@test r[1] === convert(Cint,v[1])
@test r[3] === convert(Cint,v[3])

m = Int[-5 2 9; 7 -8 3]
r = RObject(m)
@test isa(r,RObject{IntSxp})
@test length(r) == length(m)
@test size(r) == size(m)
@test isa(rcopy(r), Matrix{Int})
@test rcopy(r) == m
@test rcopy(Array, r) == m
@test r[1] === convert(Cint,m[1])
@test r[3] === convert(Cint,m[3])
@test r[2,2] === convert(Cint,m[2,2])
r[2] = -100
@test r[2] === convert(Cint,-100)
r[1,3] = -101
@test r[1,3] === convert(Cint,-101)

a = rand(-20:20,2,4,5)
r = RObject(a)
@test isa(r,RObject{IntSxp})
@test length(r) == length(a)
@test size(r) == size(a)
@test isa(rcopy(r), Array{Int,3})
@test rcopy(r) == a
@test r[1] === convert(Cint,a[1])
@test r[3] === convert(Cint,a[3])
@test r[2,3,2] === convert(Cint,a[2,3,2])

a = rand(-20:20,2,4,2,3)
r = RObject(a)
@test isa(r,RObject{IntSxp})
@test length(r) == length(a)
@test size(r) == size(a)
@test isa(rcopy(r), Array{Int,4})
@test rcopy(r) == a
@test r[1] === convert(Cint,a[1])
@test r[3] === convert(Cint,a[3])
@test r[2,3,1,2] === convert(Cint,a[2,3,1,2])

# real
@test rcopy(Float64, R"TRUE") == 1.0
@test rcopy(Float64, R"1L") == 1.0
@test rcopy(Array{Float64}, R"c(1L,2L,3L)") == [1.0,2.0,3.0]

x = 7.0
r = RObject(x)
@test isa(r,RObject{RealSxp})
@test length(r) == 1
@test size(r) == (1,)
@test rcopy(r) === convert(Float64,x)
@test r[1] === convert(Float64,x)

v = -7.0:3.0
r = RObject(v)
@test isa(r,RObject{RealSxp})
@test length(r) == length(v)
@test size(r) == size(v)
@test isa(rcopy(r), Vector{Float64})
@test rcopy(r) == collect(v)
@test rcopy(Vector, r) == collect(v)
@test r[1] === convert(Float64,v[1])
@test r[3] === convert(Float64,v[3])

m = Float64[-5 2 9; 7 -8 3]
r = RObject(m)
@test isa(r,RObject{RealSxp})
@test length(r) == length(m)
@test size(r) == size(m)
@test isa(rcopy(r), Matrix{Float64})
@test rcopy(r) == m
@test rcopy(Array, r) == m
@test r[1] === convert(Float64,m[1])
@test r[3] === convert(Float64,m[3])
@test r[2,2] === convert(Float64,m[2,2])

a = rand(2,4,5)
r = RObject(a)
@test isa(r,RObject{RealSxp})
@test length(r) == length(a)
@test size(r) == size(a)
@test isa(rcopy(r), Array{Float64,3})
@test rcopy(r) == a
@test r[1] === convert(Float64,a[1])
@test r[3] === convert(Float64,a[3])
@test r[2,3,2] === convert(Float64,a[2,3,2])

a = rand(2,4,2,3)
r = RObject(a)
@test isa(r,RObject{RealSxp})
@test length(r) == length(a)
@test size(r) == size(a)
@test isa(rcopy(r), Array{Float64,4})
@test rcopy(r) == a
@test r[1] === convert(Float64,a[1])
@test r[3] === convert(Float64,a[3])
@test r[2,3,1,2] === convert(Float64,a[2,3,1,2])

# complex
x = 7.0-2.0*im
r = RObject(x)
@test isa(r,RObject{CplxSxp})
@test length(r) == 1
@test size(r) == (1,)
@test rcopy(r) === convert(ComplexF64,x)
@test r[1] === convert(ComplexF64,x)

v = randn(7)+im*randn(7)
r = RObject(v)
@test isa(r,RObject{CplxSxp})
@test length(r) == length(v)
@test size(r) == size(v)
@test isa(rcopy(r), Vector{ComplexF64})
@test rcopy(r) == collect(v)
@test rcopy(Vector, r) == collect(v)
@test r[1] === convert(ComplexF64,v[1])
@test r[3] === convert(ComplexF64,v[3])

m = Float64[-5 2 9; 7 -8 3] + im*Float64[-5 2 9; 7 -8 3]
r = RObject(m)
@test isa(r,RObject{CplxSxp})
@test length(r) == length(m)
@test size(r) == size(m)
@test isa(rcopy(r), Matrix{ComplexF64})
@test rcopy(r) == m
@test rcopy(Array, r) == m
@test r[1] === convert(ComplexF64,m[1])
@test r[3] === convert(ComplexF64,m[3])
@test r[2,2] === convert(ComplexF64,m[2,2])

a = rand(2,4,5)+im*randn(2,4,5)
r = RObject(a)
@test isa(r,RObject{CplxSxp})
@test length(r) == length(a)
@test size(r) == size(a)
@test isa(rcopy(r), Array{ComplexF64,3})
@test rcopy(r) == a
@test r[1] === convert(ComplexF64,a[1])
@test r[3] === convert(ComplexF64,a[3])
@test r[2,3,2] === convert(ComplexF64,a[2,3,2])

a = rand(2,4,2,3)+im*randn(2,4,2,3)
r = RObject(a)
@test isa(r,RObject{CplxSxp})
@test length(r) == length(a)
@test size(r) == size(a)
@test isa(rcopy(r), Array{ComplexF64,4})
@test rcopy(r) == a
@test r[1] === convert(ComplexF64,a[1])
@test r[3] === convert(ComplexF64,a[3])
@test r[2,3,1,2] === convert(ComplexF64,a[2,3,1,2])

# dict
d = Dict(:a=>[1, 2, 4], :b=> ["e", "d", "f"])
r = RObject(d)
@test r[:a][3] == 4
@test rcopy(r[:b][2]) == "d"
l = rcopy(R"list(a=1,b=c(1,3,4))")
@test l[:a] == 1
@test l[:b][3] == 4
d = RObject(Dict("a"=>2))
@test Dict{Any,Any}("a" => 2) == rcopy(Dict, d)
@test Dict{String,Int}("a"=>2) == rcopy(Dict{String,Int}, d)

# list
a = Any[1, 1:10]
r = RObject(a)
@test isa(r, RObject{VecSxp})
@test isa(rcopy(r), Array{Any})
@test isa(rcopy(Array, r), Array{Any})

# raw
a = 0x01
r = RObject(a)
@test a == rcopy(r)
@test rcopy(UInt8, r) == 0x01

a = UInt8[0x01, 0x0c, 0xff]
r = RObject(a)
@test a == rcopy(r)
@test rcopy(Array, r) == a

# function
function funk(x,y)
    x+y
end
f1 = RObject(funk)
@test rcopy(Function, f1)(1,2) == 3
@test rcopy(Function, f1.p)(1,2) == 3

# misc
@test rcopy(Int32(1)) == 1
@test rcopy(Cint, Int32(1)) == 1
a = RObject(rand(10))
@test length(rcopy(Any, a)) == 10
# @test typeof(RCall.sexp(Cint, 1)) == Cint
# @test typeof(RCall.sexp(Float64, 1)) == Float64
# @test typeof(RCall.sexp(ComplexF64, 1)) == ComplexF64
@test typeof(rcopy(Vector{Float64}, a.p)) == Vector{Float64}
b = RObject(true)
@test rcopy(Cint, b.p) == 1
@test rcopy(Vector{Cint}, b.p) == [1]
@test rcopy(Array{Cint}, b.p) == [1]
@test rcopy(Vector{Bool}, b.p) == [true]
@test rcopy(BitVector, b.p) == [true]
@test isa(convert(Any, R"list(a=1,b=2)"), RObject)
@test isa(convert(RObject{}, R"list(a=1,b=2)"), RObject)

# issue 195
@test isa(rcopy(R"list(a=NULL)")[:a], Nothing)

# convert to Any
@test isa(rcopy(Any, R"1"), Float64)
@test isa(convert(Any, R"1"), RObject)

# s4

@test isa(rcopy(reval("""
   setClass("Foo", representation(x = "numeric"))
   foo <- new("Foo", x = 20)
""")), RObject)


text:
abstract type LassoRoot

end


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 DataTablePass

#using Debug

using CompilerTools
import CompilerTools.DebugMsg
DebugMsg.init()
using CompilerTools.Helper
using CompilerTools.LambdaHandling
import CompilerTools.ReadWriteSet
import CompilerTools.LambdaHandling.matchVarDef
import CompilerTools.LambdaHandling.lookupLHSVarByName

import HPAT
import HPAT.CaptureAPI.getColName
using HPAT.DomainPass.get_table_meta
using HPAT.DomainPass.get_col_name_from_arr

using ParallelAccelerator
using ParallelAccelerator.DomainIR.AstWalk
import ParallelAccelerator.ParallelIR
import ParallelAccelerator.ParallelIR.computeLiveness

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

"""
    get a unique id for renaming
"""
function get_unique_id(state)
  state.unique_id += 1
  return state.unique_id
end

type QueryTreeNode
    expr::Expr
    children::Vector{QueryTreeNode}
    # position in original AST
    ast_index::Int
    parent::Union{QueryTreeNode,Void}

    # empty node constructor
    function QueryTreeNode()
        new(Expr(:null), QueryTreeNode[], -1, nothing)
    end
    # Constructor
    function QueryTreeNode(expr::Expr, ind)
        new(expr, QueryTreeNode[], ind, nothing)
    end
end

isRoot(n::QueryTreeNode) = (n.parent == nothing)

function print_tree(root_qtn)
    @dprintln(3, string("     Tree::", string(root_qtn.expr.head)), " ", root_qtn.ast_index)
    for qn in root_qtn.children
        print_tree(qn)
    end
end

# information about AST gathered and used in DomainPass
type DataTableState
    linfo  :: LambdaVarInfo
    tableCols::Dict{Symbol,Vector{Symbol}}
    tableArrs::Dict{Symbol,Vector{LHSVar}}
    query_tree::QueryTreeNode
    unique_id::Int
end

# ENTRY to datatable-pass
function from_root(function_name, ast::Tuple)
    @dprintln(1, "Starting main DataTablePass.from_root.  function = ", function_name, " ast = ", ast)
    (linfo, body) = ast
    lives = computeLiveness(body, linfo)
    tableCols, tableTypes, tableIds = get_table_meta(body)
    # find variables for columns of each table
    tableArrs = get_table_arrs(tableCols, linfo)

    query_tree = make_query_tree(body.args)
    print_tree(query_tree)
    # HACK: initialize unique id to length(body) to avoid conflict with DomainPass
    state = DataTableState(linfo, tableCols, tableArrs, query_tree, length(body.args))

    # transform body
    body.args = from_toplevel_body(body.args, state)
    @dprintln(1,"DataTablePass.from_root returns function = ", function_name, " ast = ", body)
    return LambdaVarInfoToLambda(linfo, body.args, ParallelAccelerator.DomainIR.AstWalk)
end

# nodes are :body of AST
function from_toplevel_body(nodes::Array{Any,1}, state)
    res::Array{Any,1} = []
    # TODO Handle optimization; Need to replace all uses of filter output table after pushing up with new filter output table
    optimize_table_oprs(nodes, state.query_tree, state)
    print_tree(state.query_tree)
    @dprintln(3,"Datatable pass: Body after query optimizations ", Expr(:body,nodes...))

    # After optimizations make actuall call nodes for cgen
    for (index, node) in enumerate(nodes)
        if isa(node, Expr) && node.head==:filter
            append!(res, translate_hpat_filter(node, state))
        elseif isa(node, Expr) && node.head==:join
            append!(res, translate_hpat_join(node, state))
       elseif isa(node, Expr) && node.head==:aggregate
            append!(res, translate_hpat_aggregate(node, state))
        else
            append!(res, [node])
        end
    end
    return res
end

# find variables for columns of each table
function get_table_arrs(tableCols, linfo)
    tableArrs = Dict{Symbol,Vector{LHSVar}}()
    for (t,cols) in tableCols
        arrs = LHSVar[]
        for col in cols
            var = CompilerTools.LambdaHandling.lookupLHSVarByName(getColName(t,col), linfo)
            push!(arrs, var)
        end
        tableArrs[t] = arrs
    end
    return tableArrs
end

"""
    Build a query tree of relational operations.
"""
function make_query_tree(nodes::Array{Any,1})
    # node of trees where new nodes are added
    q_nodes = QueryTreeNode[]
    root = QueryTreeNode()

    for (index, node) in enumerate(reverse(nodes))
        # convert reverse index to normal index
        index = length(nodes)-index+1
        new_qt_node = QueryTreeNode()
        # ignore non-expression nodes
        # TODO: handle having different basic blocks
        if !isa(node, Expr) continue end
        if node.head==:filter
            # -2 because id and condition are above it and these are part of filter
            new_qt_node = QueryTreeNode(node, index)
        elseif node.head==:join
            # -1 because id is above it which part of join
            new_qt_node = QueryTreeNode(node, index)
        elseif node.head==:aggregate
            # node.args[4] gives length of expression nodes list and they are part of aggregate
            # * 2 because each expression nodes occupy two places
            # TODO Fix this. It might not work in future. Aggregate node indexes are variable
            new_qt_node = QueryTreeNode(node, index)
        else
            continue
        end
        # the first node is root of the tree
        # TODO: handle projections as root
        if length(q_nodes)==0
            root = new_qt_node
            push!(q_nodes, root)
        else
            add_child!(q_nodes, new_qt_node)
        end
    end
    return root
end

"""
    adds new node to tree and updates q_nodes accordingly
"""
function add_child!(q_nodes::Vector{QueryTreeNode}, new_qt_node::QueryTreeNode)
    for q_node in q_nodes
        # if any of q_node's inputs are output of new_qt_node
        if isDependent(q_node, new_qt_node)
            push!(q_node.children, new_qt_node)
            new_qt_node.parent = q_node
        end
    end
    push!(q_nodes, new_qt_node)
end

"""
    if any of t2's inputs are output of t1
"""
function isDependent(t2::QueryTreeNode, t1::QueryTreeNode)
    local t1_out_table::Symbol
    if t1.expr.head==:filter
        t1_out_table = t1.expr.args[2]
    elseif t1.expr.head==:aggregate || t1.expr.head==:join
        t1_out_table = t1.expr.args[1]
    else
        throw("invalid input query node: $t1")
    end
    local t2_in_tables::Vector{Symbol}
    if t2.expr.head==:filter
        t2_in_tables = [t2.expr.args[3]]
    elseif t2.expr.head==:join
        t2_in_tables = t2.expr.args[2:3]
    elseif t2.expr.head==:aggregate
        t2_in_tables = [t2.expr.args[2]]
    else
        throw("invalid input query node: $t2")
    end
    return t1_out_table in t2_in_tables
end

"""
    Translate join node so that backend can translate
"""
function translate_hpat_join(node, state)
    # args: id, length of output table columns, length of 1st input table columns,
    #       length of 2nd input table columns, output columns, input1 columns, input2 columns
    open_call = mk_call(GlobalRef(HPAT.API,:__hpat_join),
                        [node.args[10]; length(node.args[7]); length(node.args[8]); length(node.args[9]); node.args[7]; node.args[8]; node.args[9]])
    dprintln(3, "Datatable pass: join translated: ", open_call)
    return [open_call]
end

"""
    Make filter :call node with the following layout
    condition expression lhs, columns length, columns names(#t1#c1) ...
"""
function translate_hpat_filter(node, state)
    num_cols = length(node.args[4])
    # args: id, condition, number of columns, output table columns, input table columns
    open_call = mk_call(GlobalRef(HPAT.API,:__hpat_filter),
                        [node.args[7]; node.args[1]; num_cols; node.args[5]; node.args[6]])
    dprintln(3, "Datatable pass: filter translated: ", open_call)
    return [open_call]
end

"""
    Translate aggragte node so that backend can translate
"""
function translate_hpat_aggregate(node, state)
    # args: id, key, number of expressions, expression list
    open_call = mk_call(GlobalRef(HPAT.API,:__hpat_aggregate),
                        [node.args[7]; node.args[3]; length(node.args[4]); node.args[4]; node.args[5]; node.args[6]])
    dprintln(3, "Datatable pass: aggregate translated: ", open_call)
    return [open_call]
end

"""
    Starting point of all optimization performed on query tree
    Add rules to perform more optimizations
"""
function optimize_table_oprs(nodes, tree_node, state)
    # TODO Add more rules of optimizations
    for child in tree_node.children
        if tree_node.expr.head==:filter && child.expr.head==:join
            # try to push filter up join (filter earlier) if possible
            push_filter_up(nodes, tree_node, child, state)
        end
    end
    # call recursively on children
    for child in tree_node.children
        optimize_table_oprs(nodes, child, state)
    end

#=    curr_qtn = root_qtn
    while true
        isRoot(curr_qtn) && break
        # save child pointer because swap may happen otherwise concurrent modifiction error would be thrown
        child_qtn = curr_qtn.child
        # RULE 1: PUSH FILTER UP
        # if parent is join and child is filter then switch the order
        if curr_qtn.parent.data.head==:join && curr_qtn.data.head==:filter
            rename_map = push_filter_up(nodes, curr_qtn, tableCols, linfo)
            call_astwalk_down_tree(nodes, curr_qtn.end_pos, rename_map)
            swap_querytree_nodes(curr_qtn.parent, curr_qtn)
        end
        curr_qtn = child_qtn
    end
    =#
end

"""
    Swap the positions and data of two query tree nodes
"""
function swap_querytree_nodes(qtn1, qtn2)
    if qtn1.data.head==:join && qtn2.data.head==:filter
        # filter node is shoved above join and filter has 3 nodes
        qtn2.start_pos = qtn1.start_pos
        qtn2.end_pos = qtn2.start_pos + 2
        qtn1.start_pos = qtn1.start_pos + 3
        qtn1.end_pos = qtn1.end_pos + 3
    else
        @assert false "Write rules to handle swaping"
    end
    qtn2_child = qtn2.child
    qtn1_parent = qtn1.parent

    qtn1.parent.child = qtn2
    qtn1.child = qtn2_child

    qtn2.child = qtn1
    qtn2_child.parent = qtn1
    qtn2.parent = qtn1_parent
    qtn1.parent = qtn2

end

"""
    OPTIMIZATION: PUSH FILTER UP
        if there is a join before a filter then move that filter above join
        Assign new output to filter operation and use that in join node
        Old output of filter needs to be replaced with whatever table operation output right above filter
"""

function push_filter_up(nodes::Array{Any,1}, filter_qnode, join_qnode, state)

    join_node = join_qnode.expr
    filter_node = filter_qnode.expr
    @dprintln(3, "push filter up filter: ", filter_node, "\n       join: ", join_node)
    # args[2] and args[3] are join input tables (symbol names of tables)
    # Search in these input tables on which filter condition can be applied.
    join_input_tables = [join_node.args[2]; join_node.args[3]]
    filter_output_table = filter_node.args[2]

    # conditional expression is an mmap before filter
    cond_mmap = nodes[filter_qnode.ast_index-1].args[2]
    @dprintln(3, "filter conditional mmap: ", cond_mmap)
    cond_input_arrs = cond_mmap.args[1]
    # find out which input table of join neeeds to be filtered
    cond_column_names = map(x->get_col_name_from_arr(x, state.linfo), cond_input_arrs)
    @dprintln(3, "filtered columns: ", cond_column_names)

    new_cond_node = Expr(:null)
    new_filter_node = Expr(:null)
    # extra index if second input table is replaced
    # assuming elements of second input table in join are right after first one
    sec_ind = 0
    # if columns of first join input table are used in filtering
    if issubset(cond_column_names, state.tableCols[join_input_tables[1]])
        @dprintln(3, "first input table of join to replace")
        new_cond_node, new_filter_node = create_new_filter_node(join_input_tables[1], cond_mmap, state)
    elseif issubset(cond_column_names, state.tableCols[join_input_tables[2]])
        @dprintln(3, "second input table of join to replace")
        new_cond_node, new_filter_node = create_new_filter_node(join_input_tables[2], cond_mmap, state)
        sec_ind = 1
    else
        # return without changing
        # optimization not valid if columns of both tables involved
        return
    end

    # TODO: check for validity of optimization, e.g. make sure output of join
    #       is not used for other things using LivenessAnalysis

    @dprintln(3, "new cond node: ", new_cond_node,"\n    ** new filter node: ", new_filter_node)
    @dprintln(3, "updated lambda info: ", state.linfo)
    # update join node
    new_table_name = new_filter_node.args[2]
    join_node.args[2+sec_ind] = new_table_name
    # ordering of input arrs should be the same (key first)
    in_col_names = map(x->getColName(new_table_name, x), join_node.args[5+sec_ind])
    new_input_arrs = map(x->lookupLHSVarByName(x, state.linfo), in_col_names)
    join_node.args[8+sec_ind] = new_input_arrs
    @dprintln(3, "updated join node: ", join_node)

    old_filter_index = filter_qnode.ast_index
    join_index = join_qnode.ast_index

    # TODO: assuming there is at most one escaping variable in mmap that is assigned right before it
    dep_vars = cond_mmap.args[2].linfo.escaping_vars
    num_deps = length(dep_vars)
    @assert num_deps<=1 "too many escaping variables in filter mmap"
    if num_deps==1 && !in(dep_vars[1], state.linfo.input_params)
        dep_var = lookupLHSVarByName(dep_vars[1], state.linfo)
        @dprintln(3, "mmap escaping variable: ", dep_vars[1]," ",dep_var, " ", typeof(dep_var))
        dep_assign = nodes[old_filter_index-2]
        @assert dep_assign.head==:(=) && dep_assign.args[1]==dep_var "assignment to mmap escaping variable expected"
        deleteat!(nodes, old_filter_index-2)
        insert!(nodes, join_index, dep_assign)
        join_index += 1
    end
    # delete old filter node and its mmap
    deleteat!(nodes, (old_filter_index-1, old_filter_index))
    # insert new condition and filter nodes
    insert!(nodes, join_index, new_cond_node)
    insert!(nodes, join_index+1, new_filter_node)
    # update query tree
    # update tree node data before swapping
    filter_qnode.expr = new_filter_node
    filter_qnode.ast_index = join_index+1
    join_qnode.ast_index = join_index+2
    # filter->join->t1..t2 becomes join->(filter->t1)..t2

    old_filter_parent = filter_qnode.parent
    # if filter was root, update root in state
    if old_filter_parent==nothing
        state.query_tree = join_qnode
    else
        # update top parent's children
        for (i,c) in enumerate(old_filter_parent.children)
            if c==filter_qnode
                deleteat!(old_filter_parent.children, i)
                push!(old_filter_parent.children, join_qnode)
                break
            end
        end
    end
    join_qnode.parent = filter_qnode.parent
    filter_qnode.parent = join_qnode
    # join has two children, the one corresponding to transformed table is replaced
    jchild_found = false
    for (i,c) in enumerate(join_qnode.children)
        filter_in_table = new_filter_node.args[3]
        if getOutputTable(c.expr)==filter_in_table
            filter_qnode.children = [c]
            deleteat!(join_qnode.children, i)
            push!(join_qnode.children, filter_qnode)
            jchild_found = true
            break
        end
    end
    # if join input doesn't come from other relational operations, i.e. directly from I/O
    if jchild_found==false
        filter_qnode.children = []
        join_qnode.children = [filter_qnode]
    end

    # replace output of filter with output of join in rest of AST
    rename_map::Dict{LHSVar,LHSVar} = Dict{LHSVar,LHSVar}()
    old_filter_output = filter_node.args[5]
    old_filter_output_names = filter_node.args[4]
    corresponding_join_output_names = map(x->getColName(join_node.args[1],x), old_filter_output_names)
    corresponding_join_output = map(x->lookupLHSVarByName(x,state.linfo), corresponding_join_output_names)
    for i in 1:length(old_filter_output)
        rename_map[old_filter_output[i]] = corresponding_join_output[i]
    end
    # rename in all nodes after join
    AstWalk(Expr(:body, nodes[(join_qnode.ast_index+1):end]...), rename_symbols, rename_map)
    # filter output table symbol -> join output table symbol
    table_rename_tuple = (filter_output_table, join_node.args[1])
    AstWalk(Expr(:body, nodes[(join_qnode.ast_index+1):end]...), rename_input_table, table_rename_tuple)

    #=
    # start_pos is the cond pos
    new_cond_node = nodes[parent.start_pos]
    filter_id_node = nodes[parent.start_pos+1]
    # args[7] is the conditional
    cond_rhs = string(filter_node.args[7].args[2])
    # Extract column name from filter condition and check in join input tables on which it was applied
    # It will return table name and join input table index 1 or 2 which should be filtered
    table_name,j_ind = find_table_from_cond(tableCols,join_input_tables,cond_rhs)

    # Change condition rhs input table; Does not change output table
    replace_table_in_cond(new_cond_node,table_name)
    # return output filter table name which will be replaced in join
    out_filter_table, out_filter_table_cols, in_filter_table_cols = replace_table_in_filter_node(filter_node,table_name,tableCols)
    # As new table columns are made after pushing. We need to add to linfo
    for co in 1:length(out_filter_table_cols)
        CompilerTools.LambdaHandling.addLocalVariable(out_filter_table_cols[co],
                                                      CompilerTools.LambdaHandling.getType(in_filter_table_cols[co],linfo),
                                                      ISASSIGNEDONCE | ISASSIGNED, linfo)
    end
    # Adding mapping from filter table output to join output
    # In future change generalize it because any node output above filter node should be used instead of join output
    add_mapping(filter_output_table, join_node.args[1], tableCols, rename_map)
    # + 1 because we need right index in join expression node
    join_node.args[j_ind + 1] = out_filter_table
    # Also replace table columns list with appropriate name
    # + 7 because we need right index in join expression node
    join_node.args[j_ind + 7] = out_filter_table_cols
    # remove condition,id and filter node
    for ind in curr_qtn.start_pos:curr_qtn.end_pos
        nodes=deleteat!(nodes,ind)
    end
    # Move condition, id and filter node
    # curr_qtn parents start_pos give me place to put filter
    splice!(nodes, curr_qtn.parent.start_pos:1, [new_cond_node, filter_id_node, filter_node])

    return rename_map
    =#
end

function getOutputTable(expr::Expr)
    if expr.head==:filter
        return expr.args[2]
    elseif expr.head==:join
        return expr.args[1]
    elseif expr.head==:aggregate
        return expr.args[1]
    else
        throw("invalid relational operation: $expr")
    end
end

function create_new_filter_node(t1::Symbol, cond_mmap::Expr, state)
    cond_in_arrs = cond_mmap.args[1]
    new_cond_in_arrs = map(x->find_corresponding_column(x, t1, state), cond_in_arrs)
    new_mmap = deepcopy(cond_mmap)
    new_mmap.args[1] = new_cond_in_arrs
    new_cond_out = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
        Symbol("@$(t1)@cond_e"), Vector{Bool}, ISASSIGNEDONCE | ISASSIGNED, state.linfo))
    cond_assign = Expr(:(=), new_cond_out, new_mmap)
    tid = get_unique_id(state)
    new_table_name = Symbol("$(t1)_$(tid)")
    state.tableCols[new_table_name] = state.tableCols[t1]
    new_arrs = Any[]
    for (ind,col) in enumerate(state.tableCols[new_table_name])
        col_name = getColName(new_table_name, col)
        old_col_name = getColName(t1, col)
        old_var = state.tableArrs[t1][ind]
        typ = CompilerTools.LambdaHandling.getType(old_var, state.linfo)
        new_var = toLHSVar(CompilerTools.LambdaHandling.addLocalVariable(
            col_name, typ, ISASSIGNEDONCE | ISASSIGNED, state.linfo))
        push!(new_arrs, new_var)
    end
    state.tableArrs[new_table_name] = new_arrs
    id = get_unique_id(state)
    new_filter_node = Expr(:filter, new_cond_out, new_table_name, t1, state.tableCols[t1], new_arrs, state.tableArrs[t1], id)
    return cond_assign, new_filter_node
end

"""
    find variable of column in t1 that corresponds to col_var
"""
function find_corresponding_column(col_var::LHSVar, t1::Symbol, state)
    # e.g. _16 to :val2
    col_name = get_col_name_from_arr(col_var, state.linfo)
    # e.g. :val2 to :t1@val2
    full_col_name = getColName(t1, col_name)
    var = lookupLHSVarByName(full_col_name, state.linfo)
    return var
end

"""
    Utility function to perform renaming astwalk starting from the argument pos
"""
function call_astwalk_down_tree(nodes, pos, rename_map::Dict{LHSVar,LHSVar})
    for i in pos:length(nodes)
        AstWalk(nodes[i], rename_symbols, rename_map)
    end
end

"""
    rename input table of relational operations according to table_rename_tuple
"""
function rename_input_table(node::Expr, table_rename_tuple::Tuple{Symbol,Symbol}, top_level_number, is_top_level, read)
    if node.head==:filter
        if node.args[3]==table_rename_tuple[1]
            node.args[3] = table_rename_tuple[2]
        end
    elseif node.head==:join
        if node.args[2]==table_rename_tuple[1]
            node.args[2] = table_rename_tuple[2]
        end
        if node.args[3]==table_rename_tuple[1]
            node.args[3] = table_rename_tuple[2]
        end
    elseif node.head==:aggregate
        if node.args[2]==table_rename_tuple[1]
            node.args[2] = table_rename_tuple[2]
        end
    end
    return CompilerTools.AstWalker.ASTWALK_RECURSE
end

function rename_input_table(node::ANY, table_rename_tuple::Tuple{Symbol,Symbol}, top_level_number, is_top_level, read)
    return CompilerTools.AstWalker.ASTWALK_RECURSE
end

function rename_symbols(node::LHSVar, rename_map::Dict{LHSVar,LHSVar}, top_level_number, is_top_level, read)
    new_sym = node
    while haskey(rename_map, new_sym)
      new_sym = rename_map[new_sym]
    end
    return new_sym
end

function rename_symbols(node::TypedVar, rename_map::Dict{LHSVar,LHSVar}, top_level_number, is_top_level, read)
    typ = node.typ
    new_sym = toLHSVar(node)
    while haskey(rename_map, new_sym)
      new_sym = rename_map[new_sym]
    end
    return TypedVar(new_sym,typ)
end

function rename_symbols(node::ANY, rename_map::Dict{LHSVar,LHSVar}, top_level_number, is_top_level, read)
    return CompilerTools.AstWalker.ASTWALK_RECURSE
end

function add_mapping(old_t::Symbol, new_t::Symbol, tableCols::Dict{Symbol,Array{Symbol,1}}, mapping::Dict{LHSVar,LHSVar})
    mapping[old_t]=new_t
    old_t_cols = tableCols[old_t]
    new_t_cols = tableCols[new_t]
    for (ind,old_col) in enumerate(old_t_cols)
        new_col = new_t_cols[ind]
        @assert new_col==old_col "Both table cols should be same"
        mapping[getColName(old_t,old_col)] = getColName(new_t,new_col)
    end
    @dprintln(3,"Datatable pass: After adding new mappings = ",mapping)
end

"""
    OPTIMIZATION: Remove extra columns.
    Insert Project(select) above aggregate and join
"""
function prune_column(nodes::Array{Any,1})
end

"""
    OPTIMIZATION: Combine two or more filters using AND on condition which are on
    same table without any other inner query operations
"""
function combine_filters()
end

"""
    OPTIMIZATION: If select/projection operator only use one of the table columns after join then
    we can safely remove that join.
    TODO check with ParallelAccelerator
"""
function remove_unnessary_joins()
end

"""
    Find table name for the given filter condition
    This necessary to replace the table name if you move filter
    above join or aggregates
    TODO make it short
"""
function find_table_from_cond(tableCols,join_input_tables,cond)
    # I am assuming that second element has column name
    arr = split(cond,'@')
    col_name = arr[3]
    for (j_ind,k) in enumerate(join_input_tables)
        arr = tableCols[k]
        for (index, value) in enumerate(arr)
            curr_col = string(value)
            if curr_col == col_name
                return string(k),j_ind
            end
        end
    end
    @assert false "Could not find column in join input tables"
end

"""
    Replaces table name in the filter condition(mmap) rhs
    e.g
     table1#cond_e = table1#col1 > 1 => table1#cond_e = table2#col1 > 1
"""
function replace_table_in_cond(node,table_name)
    arr2 = split(string(node.args[2].args[1][1].name),'@')
    node.args[2].args[1][1] = Symbol(string("@",table_name,"@",arr2[3]))
end

"""
    Replaces table name and columns accordingly in the filter node after moving
"""
function replace_table_in_filter_node(node, table_name, tableCols)
    # node.args[2] gives output table
    out_table_name = string(table_name, "_fil")
    node.args[2] = Symbol(out_table_name)
    # New table must be added in tableCols which will have same columns as input
    tableCols[Symbol(out_table_name)] = tableCols[Symbol(table_name)]
    # node.args[3] gives input table
    node.args[3] = Symbol(table_name)
    # replace in condition expression (could be removed in future)
    # args[7] gives cond_lhs
    arr2 = split(string(node.args[7].args[2].name),'#')
    node.args[7].args[2] = getColName(Symbol(table_name), Symbol(arr2[3]))
    # Replace columns list with new columns
    node.args[4] = tableCols[Symbol(table_name)]
    # Replace output and input table columns array with new table columns
    node.args[5] = []
    node.args[6] = []
    for (index, col) in enumerate(tableCols[node.args[2]])
        append!(node.args[5], [getColName(Symbol(out_table_name), col)])
        append!(node.args[6], [getColName(Symbol(table_name), col)])
    end
    # Return output filter table and its columns list and input tables columns list
    return node.args[2], node.args[5], node.args[6]
end

end # DataTablePass