text:
###############################################################################
#    Model expression parsing
###############################################################################

# naming conventions
const ACC_SYM = :__acc       # name of accumulator variable
const PARAM_SYM = :__beta    # name of parameter vector
const TEMP_NAME = "tmp"      # prefix of temporary variables in log-likelihood function
const DERIV_PREFIX = "d"     # prefix of gradient variables in log-likelihood function

##########  creates a parameterized type to ease AST exploration  ############
type ExprH{H}
	head::Symbol
	args::Vector
	typ::Any
end
toExprH(ex::Expr) = ExprH{ex.head}(ex.head, ex.args, ex.typ)
toExpr(ex::ExprH) = Expr(ex.head, ex.args...)

typealias Exprequal    ExprH{:(=)}
typealias Exprdcolon   ExprH{:(::)}
typealias Exprpequal   ExprH{:(+=)}
typealias Exprmequal   ExprH{:(-=)}
typealias Exprtequal   ExprH{:(*=)}
typealias Exprtrans    ExprH{symbol("'")}
typealias Exprcall     ExprH{:call}
typealias Exprblock	   ExprH{:block}
typealias Exprline     ExprH{:line}
typealias Exprvcat     ExprH{:vcat}
typealias Exprref      ExprH{:ref}
typealias Exprif       ExprH{:if}
typealias Exprcomp     ExprH{:comparison}

## variable symbol polling functions
getSymbols(ex::Any) =        Set{Symbol}()
getSymbols(ex::Symbol) =     Set{Symbol}(ex)
getSymbols(ex::Array) =      mapreduce(getSymbols, union, ex)
getSymbols(ex::Expr) =       getSymbols(toExprH(ex))
getSymbols(ex::ExprH) =      mapreduce(getSymbols, union, ex.args)
getSymbols(ex::Exprcall) =   mapreduce(getSymbols, union, ex.args[2:end])  # skip function name
getSymbols(ex::Exprref) =    setdiff(mapreduce(getSymbols, union, ex.args), Set(:(:), symbol("end")) )# ':'' and 'end' do not count
getSymbols(ex::Exprcomp) =   setdiff(mapreduce(getSymbols, union, ex.args),
	Set(:(>), :(<), :(>=), :(<=), :(.>), :(.<), :(.<=), :(.>=), :(==)) )

## variable symbol subsitution functions
substSymbols(ex::Expr, smap::Dict) =          substSymbols(toExprH(ex), smap::Dict)
substSymbols(ex::Exprcall, smap::Dict) =      Expr(:call, ex.args[1], map(e -> substSymbols(e, smap), ex.args[2:end])...)
substSymbols(ex::ExprH, smap::Dict) =         Expr(ex.head, map(e -> substSymbols(e, smap), ex.args)...)
substSymbols(ex::Symbol, smap::Dict) =        haskey(smap, ex) ? smap[ex] : ex
substSymbols(ex::Vector{Expr}, smap::Dict) =  map(e -> substSymbols(e, smap), ex)
substSymbols(ex::Any, smap::Dict) =           ex

######### parameters structure  ############
# type MCMCParams
# 	sym::Symbol
# 	size::Vector{Integer}
# 	map::Union(Integer, Range1)
# end

######### structure for parsing model  ##############
type ParsingStruct
	bsize::Int                # length of beta, the parameter vector
	pars::Dict{Symbol, PDims} # parameters with their mapping to the beta real vector
	init::Vector{Float64}     # initial values of beta
	source::Expr              # model source, after first pass
	exprs::Vector{Expr}       # vector of assigments that make the model
	dexprs::Vector{Expr}      # vector of assigments that make the gradient
	finalacc::Symbol          # last symbol of loglik accumulator after renaming
	varsset::Set{Symbol}      # all the vars set in the model
	pardesc::Set{Symbol}      # all the vars set in the model that depend on model parameters
	accanc::Set{Symbol}       # all the vars (possibly external) that influence the accumulator
end
ParsingStruct() = ParsingStruct(0, Dict{Symbol, PDims}(), Float64[], :(), Expr[], Expr[], ACC_SYM,
	Set{Symbol}(), Set{Symbol}(), Set{Symbol}())

######### first pass on the model
#  - extracts parameters definition
#  - rewrite ~ operators  as acc += logpdf..(=)
#  - translates x += y into x = x + y, same for -= and *=
function parseModel!(m::ParsingStruct, source::Expr)
	local distribFound::Bool = false

	explore(ex::Expr) =       explore(toExprH(ex))
	explore(ex::ExprH) =      error("[parseModel] unmanaged expr type $(ex.head)")
	explore(ex::Exprline) =   nothing  # remove #line statements
	explore(ex::Exprref) =    toExpr(ex) # no processing
	explore(ex::Exprequal) =  toExpr(ex) # no processing
	explore(ex::Exprvcat) =   toExpr(ex) # no processing

	explore(ex::Exprpequal) = (args = ex.args ; Expr(:(=), args[1], Expr(:call, :+, args...)) )
	explore(ex::Exprmequal) = (args = ex.args ; Expr(:(=), args[1], Expr(:call, :-, args...)) )
	explore(ex::Exprtequal) = (args = ex.args ; Expr(:(=), args[1], Expr(:call, :*, args...)) )

	function explore(ex::Exprblock)
		al = {}
		for ex2 in ex.args
			if isa(ex2, Expr)
				ex3 = explore(ex2)
				ex3==nothing ? nothing : push!(al, ex3)
			else
				push!(al, ex2)
			end
		end
		Expr(ex.head, al...)
	end

	function explore(ex::Exprcall)
		ex.args[1] != :~ && return toExpr(ex)

		distribFound = true
		fn = symbol("logpdf$(ex.args[3].args[1])")
		return :($ACC_SYM = $ACC_SYM + $(Expr(:call, fn, ex.args[3].args[2:end]..., ex.args[2])))
	end

	@assert source.head==:block && length(source.args)>=1 "model should be a block with at least 1 statement"

	m.source = explore(source)

	# if no distribution expression '~' was found consider that last expr is the variable to be maximized
	if !distribFound
		m.source.args[end] = Expr(:(=) , ACC_SYM, m.source.args[end] )
	end
end

######## unfolds expressions to prepare derivation ###################
function unfold!(m::ParsingStruct)

	explore(ex::Expr) =       explore(toExprH(ex))
	explore(ex::ExprH) =      error("[unfold] unmanaged expr type $(ex.head)")
	explore(ex::Exprline) =   nothing     # remove line info
	explore(ex::Exprref) =    toExpr(ex)   # unchanged
	explore(ex::Exprcomp) =   toExpr(ex)  # unchanged
	explore(ex::Exprvcat) =   explore(Expr(:call, :vcat, ex.args...) )  # translate to vcat(), and explore
	explore(ex::Exprtrans) =  explore(Expr(:call, :transpose, ex.args[1]) )  # translate to transpose() and explore
	explore(ex::Any) =        ex

	explore(ex::Exprblock) =  mapreduce(explore, (a,b)->b, ex.args)  # process, and return last evaluated

	function explore(ex::Exprequal)
		lhs = ex.args[1]
		@assert typeof(lhs) == Symbol ||  (typeof(lhs) == Expr && lhs.head == :ref)
			"[unfold] not a symbol on LHS of assigment $ex"

		rhs = ex.args[2]
		if isa(rhs, Symbol) || isa(rhs, Real)
			push!(m.exprs, Expr(:(=), lhs, rhs))
		elseif isa(rhs, Expr) # only refs and calls will work
				ue = explore(toExprH(rhs)) # explore will return something in this case
				push!(m.exprs, Expr(:(=), lhs, ue))
		else  # unmanaged kind of lhs
		 	error("[unfold] can't handle RHS of assignment $ex")
		end
		lhs
	end

	function explore(ex::Exprcall)
		na = {ex.args[1]}   # function name
		args = ex.args[2:end]  # arguments

		# if more than 2 arguments, +, sum and * are converted  to nested expressions
		#  (easier for derivation)
		# TODO : apply to other n-ary (n>2) operators ?
		if in(na[1], [:+, :*, :sum])
			while length(args) > 2
				a2 = pop!(args)
				a1 = pop!(args)
				push!(args, Expr(:call, ex.args[1], a1, a2))
			end
		end

		for e2 in args
			if isa(e2, Expr) # only refs and calls will work
				ue = explore(e2)
				nv = gensym(TEMP_NAME)
				push!(m.exprs, :($nv = $ue))
				push!(na, nv)
			else
				push!(na, e2)
			end
		end

		Expr(ex.head, na...)
	end

	explore(m.source)
end

######### renames variables set several times to make them unique  #############
# FIXME : algo doesn't work when a variable sets individual elements, x = .. then x[3] = ...;
# FIXME 2 : external variables redefined within model are not renamed
function uniqueVars!(m::ParsingStruct)
	el = m.exprs
    subst = Dict{Symbol, Symbol}()
    used = Set(ACC_SYM)

    for idx in 1:length(el) # idx=4
        # first, substitute in the rhs the variables names that have been renamed
        el[idx].args[2] = substSymbols(el[idx].args[2], subst)

        # second, rename lhs symbol if set before
        lhs = collect(getSymbols(el[idx].args[1]))[1]  # there should be only one
        if in(lhs, used) # if var already set once => create a new one
            subst[lhs] = gensym("$lhs") # generate new name, add it to substitution list for following statements
            el[idx].args[1] = substSymbols(el[idx].args[1], subst)
        else # var set for the first time
            union!(used, Set(lhs))
        end
    end

	m.finalacc = haskey(subst, ACC_SYM) ? subst[ACC_SYM] : ACC_SYM  # keep reference of potentially renamed accumulator
end

######### identifies vars #############
# - lists variables that depend on model parameters
# - lists variables that influence the accumulator
# - lists variables defined
# In order to
#   1) restrict gradient code to the strictly necessary variables
#   2) move parameter independant variables definition out the function (but within closure)
#   3) TODO : remove unnecessary variables (with warning)
#   4) identify external vars
function categorizeVars!(m::ParsingStruct)
	lhsSymbol(ex) = Set(isa(ex.args[1], Symbol) ? ex.args[1] : ex.args[1].args[1])

    m.varsset = mapreduce(lhsSymbol, union, m.exprs)

    local parset = Set{Symbol}(collect(keys(m.pars))...)  # [p.sym for p in m.pars]...)
    m.pardesc = copy(parset)  # start with parameter symbols
    for ex2 in m.exprs
        lhs = lhsSymbol(ex2)
        rhs = getSymbols(ex2.args[2])

        !isempty(intersect(rhs, m.pardesc)) && union!(m.pardesc, lhs)
    end

    m.accanc = Set{Symbol}(m.finalacc)
    for ex2 in reverse(m.exprs) # proceed backwards ex2 = reverse(m.exprs)[3]
        lhs = lhsSymbol(ex2)
        rhs = setdiff(getSymbols(ex2), lhs) # to pickup potential index on lhs as an ancestor
        # isa(ex2.args[1], Expr) && ex2.args[1].head == :ref && union!(rhs, getSymbols(ex2.args[1].args[2]))

        !isempty(intersect(lhs, m.accanc)) && union!(m.accanc, rhs)
    end

    @assert in(m.finalacc, m.pardesc) "Model parameters do not seem to influence model outcome"

    local parset2 = setdiff(parset, m.accanc)
    @assert isempty(parset2) "Model parameter(s) $(collect(parset2)) do not seem to influence model outcome"

end

######### builds the gradient expression from unfolded expression ##############
function backwardSweep!(m::ParsingStruct)

	explore(ex::Expr) = explore(toExprH(ex))
	explore(ex::ExprH) = error("[backwardSweep] unmanaged expr type $(ex.head)")
	explore(ex::Exprline) = nothing

	function explore(ex::Exprequal)
		lhs = ex.args[1]
		if isa(lhs,Symbol) # simple var case
			dsym = lhs
			dsym2 = symbol("$(DERIV_PREFIX)$lhs")
		elseif isa(lhs,Expr) && lhs.head == :ref  # vars with []
			dsym = lhs
			dsym2 = Expr(:ref, symbol("$(DERIV_PREFIX)$(lhs.args[1])"), lhs.args[2:end]...)
		else
			error("[backwardSweep] not a symbol on LHS of assigment $(ex)")
		end

		rhs = ex.args[2]
		if !isa(rhs,Symbol) && !isa(rhs,Expr) # some kind of number, nothing to do

		elseif isa(rhs,Symbol)
			if in(rhs, avars)
				vsym2 = symbol("$(DERIV_PREFIX)$rhs")
				push!(m.dexprs, :( $vsym2 = $dsym2))
			end

		elseif isa(toExprH(rhs), Exprref)
			if in(rhs.args[1], avars)
				vsym2 = Expr(:ref, symbol("$(DERIV_PREFIX)$(rhs.args[1])"), rhs.args[2:end]...)
				push!(m.dexprs, :( $vsym2 = $dsym2))
			end

		elseif isa(toExprH(rhs), Exprcall)
			for i in 2:length(rhs.args)
				vsym = rhs.args[i]
				if isa(vsym, Symbol) && in(vsym, avars)
					m.dexprs = vcat(m.dexprs, derive(rhs, i-1, dsym))
				end
			end
		else
			error("[backwardSweep] can't derive $rhs")
		end
	end

	avars = intersect(m.accanc, m.pardesc)
	for ex2 in reverse(m.exprs)  # proceed backwards
		@assert isa(ex2, Expr) "[backwardSweep] not an expression : $ex2"
		explore(ex2)
	end
end

######## sets inital values from 'init' given as parameter  ##########
function setInit!(m::ParsingStruct, init)
    @assert length(init)>=1 "There should be at leat one parameter specified, none found"

    for p in init  # p = collect(init)[1]
        par = p[1]  # param symbol defined here
        def = p[2]

        @assert typeof(par) == Symbol "[setInit] not a symbol in init param : $(par)"

        if isa(def, Real)  #  single param declaration
            # push!(m.pars, MCMCParams(par, Integer[], m.bsize+1))
            m.pars[par] = PDims(m.bsize+1, ())
            m.bsize += 1
            push!(m.init, def)

        elseif isa(def, Array) && ndims(def) == 1
            nb = size(def,1)
            # push!(m.pars, MCMCParams(par, Integer[nb], (m.bsize+1):(m.bsize+nb)))
            m.pars[par] = PDims(m.bsize+1, (nb,))
            m.bsize += nb
            m.init = [m.init, def...]

        elseif isa(def, Array) && ndims(def) == 2
            nb1, nb2 = size(def)
            # push!(m.pars, MCMCParams(par, Integer[nb1, nb2], (m.bsize+1):(m.bsize+nb1*nb2)))
            m.pars[par] = PDims(m.bsize+1, (nb1,nb2))
            m.bsize += nb1*nb2
            m.init = [m.init, vec(def)...]

        else
            error("[setInit] unsupported parameter type for $(par)")
        end
    end

end

######### returns an array of expr assigning parameters from the beta vector  ############
# function betaAssign(m::ParsingStruct)
# 	pmap = m.pars
# 	assigns = Expr[]
# 	for p in pmap
# 		if length(p.size) <= 1  # scalar or vector
# 			push!(assigns, :($(p.sym) = $PARAM_SYM[ $(Expr(:quote,p.map)) ]) )
# 		else # matrix case  (needs a reshape)
# 			push!(assigns, :($(p.sym) = reshape($PARAM_SYM[ $(Expr(:quote,p.map)) ], $(p.size[1]), $(p.size[2]))) )
# 		end
# 	end
# 	assigns
# end
function betaAssign(m::ParsingStruct)
	assigns = Expr[]
	for p in keys(m.pars)
		v = m.pars[p]

		if length(v.dims) == 0  # scalar
			push!(assigns, :($p = $PARAM_SYM[ $(v.pos) ]) )

		elseif length(v.dims) == 1  # vector
			r = v.pos:(v.pos+prod(v.dims)-1)
			push!(assigns, :($p = $PARAM_SYM[ $(Expr(:quote,r)) ]) )

		else # matrix case  (needs a reshape)
			r = v.pos:(v.pos+prod(v.dims)-1)
			push!(assigns, :($p = reshape($PARAM_SYM[ $(Expr(:quote,r)) ], $(v.dims[1]), $(v.dims[2]))) )
		end
	end
	assigns
end

######### evaluates once all variables to give type hints for derivation ############
#  most gradient calculation statements depend on the type of variables (Scalar or Array)
#  this is where they are evaluated (with values stored in global Dict 'vhint' )
function preCalculate(m::ParsingStruct)
    global vhint = Dict()

    body = Expr[ betaAssign(m)...,
                 :(local $ACC_SYM = 0.),
                 m.exprs...]

    vl = getSymbols(body)  # list of all vars (external, parameters, set by model, and accumulator)
    body = vcat(body, [ :(vhint[$(Expr(:quote, v))] = $v) for v in vl ], :(return $(m.finalacc)))

	# enclose in a try block to catch zero likelihoods (-Inf log likelihood)
	body = Expr(:try, Expr(:block, body...),
			          :e,
			          Expr(:block, :(if e == "give up eval"; return(-Inf); else; throw(e); end)))

	# identify external vars and add definitions x = Main.x
	ev = setdiff(m.accanc, union(m.varsset, Set(ACC_SYM, collect(keys(m.pars))...))) # vars that are external to the model
	vhooks = Expr(:block, [ :( local $v = $(Expr(:., :Main, Expr(:quote, v))) ) for v in ev]...) # assigment block

	# build and evaluate the let block containing the function and external vars hooks
	fn = gensym()
	body = Expr(:function, Expr(:call, fn, :($PARAM_SYM::Vector{Float64})),	Expr(:block, body) )
	body = :(let; global $fn; $vhooks; $body; end)
	eval(body)
	fn = eval(fn)

	# now evaluate vhint (or throw error if model does not evaluate for given initial values)
	res = fn(m.init)
	!isa(res, Real) && error("Model outcome should be a scalar, $(typeof(res)) found")
	res == -Inf && error("Initial values out of model support, try other values")
end

######### builds the model function ##############
# 'init' contains the dictionary of model params and their initial value
#    ! Initial values are used for the precalculate run that will allow
#    to know all variables types.
# If 'debug' is set to true, the function prints out the model function
#  that would have been created
#
function generateModelFunction(model::Expr; gradient=false, debug=false, init...)
	m = ParsingStruct()

	## checks initial values
	setInit!(m, init)

	## rewrites ~ , do some formatting ... on the model expression
	parseModel!(m, model)

	## process model
	unfold!(m)
	uniqueVars!(m)
	categorizeVars!(m)

	## build function expression
	if gradient  # case with gradient
		preCalculate(m)
		backwardSweep!(m)

		body = Expr[] # list of = expr making the model
		dsym(v::Symbol) = symbol("$DERIV_PREFIX$(v)")

		# initialization statements
		body = [ betaAssign(m)...,        # assigments beta vector -> model parameter vars
		         :($ACC_SYM = 0.),        # initialize accumulator
		         :($(dsym(m.finalacc)) = 1.0)] # initialize accumulator gradient accumulator

		avars = setdiff(intersect(m.accanc, m.pardesc), Set(m.finalacc)) # active vars without accumulator, treated above
		for v in avars
			vh = vhint[v]
			if isa(vh, Real)
				push!(body, :($(dsym(v)) = 0.))
			else
				# push!(body, :($(dsym(v)) = zeros(Float64, $(Expr(:quote,size(vh))))) )
				push!(body, :($(dsym(v)) = zeros(Float64, $(Expr(:tuple,size(vh)...)))) )
			end
		end

		# build function statements, and move to let block constant statements for optimization
		header = Expr[]  # let block statements
		fvars = union(Set([e.args[1] for e in body]...), Set(PARAM_SYM)) # vars that are re-evaluated at each function call
		for ex in [m.exprs..., m.dexprs...]
			if length(intersect(getSymbols(ex.args[2]), fvars)) > 0
				push!(body, ex)
				fvars = union(fvars, getSymbols(ex.args[1]))
			else
				push!(header, ex)
			end
		end

		# prefix statements with 'local' at first occurence
		vars  = Set(PARAM_SYM)
		for i in 1:length(header)
			if length(intersect(getSymbols(header[i].args[1]), vars)) == 0
				header[i] = :(local $(header[i]))
				vars = union(vars, getSymbols(header[i].args[1]))
			end
		end
		for i in 1:length(body)
			if length(intersect(getSymbols(body[i].args[1]), vars)) == 0
				body[i] = :(local $(body[i]))
				vars = union(vars, getSymbols(body[i].args[1]))
			end
		end

		# return statement (note : gradient vec should match beta variable mapping)
		# dexp = { :( vec([$(dsym(p.sym))]) ) for p in m.pars}
		gsym = dsym(PARAM_SYM)
		push!(body, :( local $gsym = similar($PARAM_SYM)))
		for p in keys(m.pars)
			v = m.pars[p]

			if length(v.dims) == 0  # scalar
				push!(body, :( $gsym[ $(v.pos) ] = $(dsym(p)) ) )
			elseif length(v.dims) == 1  # vector
				r = v.pos:(v.pos+prod(v.dims)-1)
				push!(body, :( $gsym[ $(Expr(:quote,r)) ] = $(dsym(p)) ) )
			else # matrix case  (needs a reshape)
				r = v.pos:(v.pos+prod(v.dims)-1)
				push!(body, :( $gsym[ $(Expr(:quote,r)) ] = vec($(dsym(p))) ))
			end
		end
		push!(body, :(($(m.finalacc), $gsym)))

		# dexp = { :( vec([$(dsym(p))]) ) for p in keys(m.pars)}
		# dexp = length(m.pars) > 1 ? Expr(:call, :vcat, dexp...) : dexp[1]
		# push!(body, :(($(m.finalacc), $dexp)))

		# enclose in a try block
		body = Expr(:try, Expr(:block, body...),
				          :e,
				          Expr(:block, :(if e == "give up eval"; return(-Inf, zero($PARAM_SYM)); else; throw(e); end)))

	else  # case without gradient
		body = [ betaAssign(m)...,        # assigments beta vector -> model parameter vars
		         :(local $ACC_SYM = 0.),  # initialize accumulator
                 m.source.args...,        # model statements
                 :(return($ACC_SYM)) ]

		# enclose in a try block
		body = Expr(:try, Expr(:block, body...),
				          :e,
				          Expr(:block, :(if e == "give up eval"; return(-Inf); else; throw(e); end)))

		header = Expr[]
	end

	# identify external vars and add definitions x = Main.x
	ev = setdiff(m.accanc, union(m.varsset, Set(ACC_SYM), Set(collect(keys(m.pars))...))) # vars that are external to the model
	header = [[ :( local $v = $(Expr(:., :Main, Expr(:quote, v))) ) for v in ev]..., header...] # assigment block

	# build and evaluate the let block containing the function and external vars hooks
	fn = gensym("ll")
	body = Expr(:function, Expr(:call, fn, :($PARAM_SYM::Vector{Float64})),	Expr(:block, body) )
	body = Expr(:let, Expr(:block, :(global $fn), header..., body))

	debug ? body : (eval(body) ; (eval(fn), m.bsize, m.pars, m.init) )
end


text:
# This file is part of the StateMachines.jl Julia package
#
# SMaul Bayer
# MIT License
# Github: https://github.com/pbayer/StateMachines.jl
#

c = Clock()

struct B1 <: Behavior end

SM1 = StateMachine{B1}("SM1", B1())
SM2 = StateMachine("SM2", B1())

@test SM1 isa StateMachine{B1}
@test SM2 isa StateMachine{B1}

ext: jl


text:
function a_res(model::PCSAFTFamily, z, v, T)
    return a_hc(model,z,v,T) + a_disp(model,z,v,T) + a_assoc(model,z,v,T)
end

function a_hc(model::PCSAFTFamily, z, v, T)
    x = z/sum(z[i] for i in model.components)
    m = model.params.segment
    m̄ = sum(x[i]*m[i] for i in model.components)
    return m̄*a_hs(model,z,v,T) - sum(x[i]*(m[i]-1)*log(g_hsij(model,z,v,T, i, i)) for i in model.components)
end

function a_disp(model::PCSAFTFamily, z, v, T)
    x = z/sum(z[i] for i in model.components)
    m = model.params.segment
    m̄ = sum(x[i]*m[i] for i in model.components)
    return -2*π*N_A*sum(z[i] for i in model.components)/v*I_n(model,z,v,T, 1)*m2ϵσ3(model,z,v,T, 1) - π*m̄*N_A*sum(z[i] for i in model.components)/v*C1(model,z,v,T)*I_n(model,z,v,T, 2)*m2ϵσ3(model,z,v,T, 2)
end

function d(model::PCSAFTFamily, z, v, T, component)
    ϵ = model.params.epsilon[component]
    σ = model.params.sigma[component]
    return σ * (1 - 0.12exp(-3ϵ/T))
end

function ζn(model::PCSAFTFamily, z, v, T, n)
    x = z/sum(z[i] for i in model.components)
    m = model.params.segment
    return N_A*sum(z[i] for i in model.components)*π/6/v * sum(x[i]*m[i]*d(model,z,v,T, i)^n for i in model.components)
end

function g_hsij(model::PCSAFTFamily, z, v, T, i, j)
    di = d(model,z,v,T, i)
    dj = d(model,z,v,T, j)
    ζ2 = ζn(model,z,v,T, 2)
    ζ3 = ζn(model,z,v,T, 3)
    return 1/(1-ζ3) + di*dj/(di+dj)*3ζ2/(1-ζ3)^2 + (di*dj/(di+dj))^2*2ζ2^2/(1-ζ3)^3
end

function a_hs(model::PCSAFTFamily, z, v, T)
    ζ0 = ζn(model,z,v,T, 0)
    ζ1 = ζn(model,z,v,T, 1)
    ζ2 = ζn(model,z,v,T, 2)
    ζ3 = ζn(model,z,v,T, 3)
    return 1/ζ0 * (3ζ1*ζ2/(1-ζ3) + ζ2^3/(ζ3*(1-ζ3)^2) + (ζ2^3/ζ3^2-ζ0)*log(1-ζ3))
end

function C1(model::PCSAFTFamily, z, v, T)
    x = z/sum(z[i] for i in model.components)
    η = ζn(model,z,v,T, 3)
    m = model.params.segment
    m̄ = sum(x[i]*m[i] for i in model.components)
    return (1 + m̄*(8η-2η^2)/(1-η)^4 + (1-m̄)*(20η-27η^2+12η^3-2η^4)/((1-η)*(2-η))^2)^-1
end

function m2ϵσ3(model::PCSAFTFamily, z, v, T, ϵ_power = 1)
    x = z/sum(z[i] for i in model.components)
    m = model.params.segment
    σ = model.params.sigma
    ϵ = model.params.epsilon
    #= return sum(x[i]*x[j]*m[i]*m[j] * (sqrt(ϵ[i]*ϵ[j])*(1-k[union(i,j)])/T)^ϵ_power * (0.5*(σ[i]+σ[j]))^3 for i in model.components, j in model.components) =#
    return sum(x[i]*x[j]*m[i]*m[j] * (ϵ[union(i,j)]*(1)/T)^ϵ_power * σ[union(i,j)]^3 for i in model.components, j in model.components)
end

function I_n(model::PCSAFTFamily, z, v, T, n)
    x = z/sum(z[i] for i in model.components)
    m = model.params.segment
    m̄ = sum(x[i]*m[i] for i in model.components)
    η = ζn(model,z,v,T, 3)
    if n == 1
        corr = [0.9105631445 -0.3084016918 -0.0906148351;
                0.6361281449 0.1860531159 0.4527842806;
                2.6861347891 -2.5030047259 0.5962700728;
                -26.547362491 21.419793629 -1.7241829131;
                97.759208784 -65.255885330 -4.1302112531;
                -159.59154087 83.318680481 13.776631870;
                91.297774084 -33.746922930 -8.6728470368]
    elseif n == 2
        corr = [0.7240946941 -0.5755498075 0.0976883116;
                2.2382791861 0.6995095521 -0.2557574982;
                -4.0025849485 3.8925673390 -9.1558561530;
                -21.003576815 -17.215471648 20.642075974;
                26.855641363 192.67226447 -38.804430052;
                206.55133841 -161.82646165 93.626774077;
                -355.60235612 -165.20769346 -29.666905585]
    end
    return sum((corr[i+1,1] + (m̄-1)/m̄*corr[i+1,2] + (m̄-1)/m̄*(m̄-2)/m̄*corr[i+1,3]) * η^i for i = 0:6)
end

function a_assoc(model::PCSAFTFamily, z, v, T)
    x = z/sum(z[i] for i in model.components)
    X_iA = X_assoc(model,z,v,T)
    return sum(x[i]*sum(log(X_iA[i,a])-X_iA[i,a]/2 + model.params.n_sites[i][a]/2 for a in keys(model.params.n_sites[i])) for i in model.components)
end

function X_assoc(model::PCSAFTFamily, z, v, T)
    x = z/sum(z[i] for i in model.components)
    ρ = N_A*sum(z[i] for i in model.components)/v
    X_iA = Dict()
    X_iA_old = Dict()
    tol = 1.
    iter = 1
    while tol > 1e-12 && iter < 100
        for i in model.components
            for a in keys(model.params.n_sites[i])
                A = 0.
                for j in model.components
                    B = 0
                    for b in keys(model.params.n_sites[j])
                        if haskey(model.params.epsilon_assoc,Set([(i,a),(j,b)]))
                            if iter!=1
                                B+=X_iA_old[j,b]*Δ(model,z,v,T,i,j,a,b)
                            else
                                B+=Δ(model,z,v,T,i,j,a,b)
                            end
                        end
                    end
                    A += ρ*x[j]*B
                end
                if iter == 1
                    X_iA[i,a] =0.5+0.5*(1+A)^-1
                else
                    X_iA[i,a] =0.5*X_iA_old[i,a]+0.5*(1+A)^-1
                end
            end
        end
        if iter == 1
            tol = sqrt(sum(sum((1. -X_iA[i,a])^2 for a in keys(model.params.n_sites[i])) for i in model.components))
        else
            tol = sqrt(sum(sum((X_iA_old[i,a] -X_iA[i,a])^2 for a in keys(model.params.n_sites[i])) for i in model.components))
        end
        X_iA_old = deepcopy(X_iA)
        iter += 1
    end

    return X_iA
end

function Δ(model::PCSAFTFamily, z, v, T, i, j, a, b)
    ϵ_assoc = model.params.epsilon_assoc[Set([(i,a),(j,b)])]
    κ = model.params.bond_vol[Set([(i,a),(j,b)])]
    σ = model.params.sigma[union(i,j)]
    g = g_hsij(model,z,v,T,i,j)
    return g*σ^3*(exp(ϵ_assoc/T)-1)*κ
end


text:
using  LinearAlgebra #
using  WAV
using  FFTW
using  DSP
using Statistics
include("adaa.jl")
#for the more imaginative distortions:
#rolling average like
function applyDistortion(signal::AbstractVector{T}, increment::Integer, distortion::Function)where
  {T<:Real}
    0 < increment || throw(DomainError(increment, "increment must be greater than or equal to 1"))
  #signal to return
  newSignal = T[]
  signalLength = length(signal)
  for index in 1:increment:signalLength
      if index + increment  > signalLength
          increment = signalLength - index +1
      end
      block = signal[index:index+increment-1]
      append!(newSignal, distortion(block))
      if index == signalLength  && length(newSignal) < signalLength
          push!(newSignal, 80)
      end
    end
    return newSignal
end
#clip a signal of a vector
function hardClip(signal::AbstractVector{T}, limit::T) where
    {T<:Real}
       0.0 <= limit <= 1.0 || throw(DomainError(limit, "limit must be between 0 and 1"))
       clip(i) = (abs(i[1]) < limit && return i[1])  || return i[1] > 0 ? limit : -limit
        return clip.(signal)
end

#hardClipping implented with gain/more idiomatically
function hardC(signal::Vector{Float64}, gain)::Vector{Float64}
  c(x,y) = (abs(x*y) < 1 && return x*y)  || return sign(x)
  return c.(signal,gain)
end

function hardClipAD1(signal::Vector{Float64}, gain)::Vector{Float64}
  f(x) = (abs(x) < 1 ? x : sign(x) )
  f1(x) = (abs(x) < 1 ? x^2/2 : x*sign(x) - .5)
  return processAD1(gain*signal, f,f1)
end

function hardClipAD2(signal::Vector{Float64}, gain)::Vector{Float64}
  f(x) = (abs(x) < 1 ? x : sign(x) )
  f1(x) = (abs(x) < 1 ? x^2/2 : x*sign(x) - .5)
  f2(x) = (abs(x) < 1 ? x^3/6 : (x^2/2 + 1/6) * sign(x) - x/2)
  return processAD2(gain*signal, f,f1,f2)
end

#softclipping of a vector. maybe latter on i'll add
#more sigmoids
function softClip(signal::AbstractVector{T}, limit)where
    {T<:Real}
      #as limit becomes really large function behaves like hard cliping.
      #might need to check for zed
      -1300.0 < limit < 13300.0 || throw(DomainError(limit, "limit must be between -1300 and 13000"))
      soft(i) = tanh(limit * i)
      return soft.(signal)
end
#weird movable wave shapping that i found
function chebby(signal::AbstractVector{T}, limit)where
    {T<:Real}
    0.0 <= limit <= 4.0 || throw(DomainError(limit, "limit must be between 0 and 4"))
    cheb(x)=limit*x[1]^3+(1-limit)x[1]
    return cheb.(signal)
end
function chebyshev(order,x)
  if order == 0
      return 0
  elseif  order == 1
      return x
  elseif order < 21
      Tn_0 = Float64(1)
      Tn_1 = x
      Tn = Float64(0)
      for n in 2:1:order
          Tn = 2*x*Tn_1-Tn_0
          Tn_0 = Tn_1
          Tn_1= Tn
      end
      return Tn
  else
      return cos(order*acos(x))
  end

end

#=gonna figure this out latter:
#supposed to return a function that represents
a weighted sum of chebbysehv polynomials
right now it will output the a lambda function in string form
=#
function poly(weights)
  scale = sum(weights[2])
  nomials = AbstractString["(1)", "(x[1])", "(2x[1]^2-1)",
                              "(4x[1]^3-3x[1])",
                              "(8x[1]^4[1]-8x[1]^2+1)",
                              "(16x[1]^5-20x[1]^3+5x[1])",
                             "(32x[1]^6-48x[1]^4+18x[1]^2-1))",
                             "(64x[1]^7-112x[1]^5+56x[1]^3-7x[1])"]

    func = "x-> "
    for i in 1:length(weights[1]) func *= nomials[i]  *= "*$(weights[2][i])/$scale + "
    end
    func = func[1:length(func)-2]
    return func

end
#first four chebby polynomials weighted
function cheb(signal::AbstractVector)
  ch(val)= ((4 * val^3 - 3 * val) + (2 * val^2 - 1) + val + 1)/4
  return ch.(signal)
end
#=moving average filter but worse:
# perhaps interested for percussive elements?
=#
function fakebitcrush(signal::AbstractVector{T}, rate::Integer) where
    {T<:Real}
      1 < rate <= length(signal)/rate ||
      throw(DomainError(rate, "rate must be between 0 and length(signal)/100"))
        bit(block) = mean(block)*ones(length(block))
      return applyDistortion(signal, rate, bit)
end

function linear(block)#linear regression
    linreg(x,y) = hcat(ones(size(y, 1)), y) \ x
    timeVector = range(0.0, step=inv(41000), length=length(block))
    return linreg(block,timeVector)
end
#=replaces partition amount of samples with linear approximation
sounds like a combo of bitcrush and lpf applied to random noise.
pretty cool
=#
function linearDistortion(signal::AbstractVector{T}, partition::Integer) where
    {T<:Real}
      0 < partition || throw(DomainError(iterations, "partition must be greater than or equal to 1"))
    #applies best fit apprx
    f(y,x) = y[2] .* x .+ y[1]
    line(block) = f(linear(block),range(0.0, step=inv(41000), length=length(block)))
    applyDistortion(signal, partition,line)
end
#= replace partition samples with the sample samples but sorted
ascending or decending based on the linear regression of said partition()
similar to linearDistortion, but more intense
=#
function sortDistortion(signal::AbstractVector{T}, partition::Integer) where
    {T<:Real}
      0 < partition || throw(DomainError(iterations, "partition must be greater than or equal to 1"))
    sorting(block) =  linear(block)[2] > 0 ? sort(block) : sort(block,rev=true)
    return applyDistortion(signal, partition, sorting)
end

max_stars_repo_name: justinbroce/julia-audio-stuffs

max_issues_repo_name: justinbroce/julia-audio-stuffs

max_forks_repo_name: justinbroce/julia-audio-stuffs


text:
using Router, App, MapGen, Agent, JSON

route("/", named = :root) do
    Router.serve_static_file("/main.html")
end

route("/rmap", method=POST, named = :randommap) do
    height = parse(Int, @params(:h))
    width = parse(Int, @params(:w))
    global navimap = MapGen.random_map(height, width)
    respond(Dict(:json => JSON.json(Dict(:name => navimap.name, :nodes => navimap.nodes, :edges => navimap.edges))))
end

route("/execute", method=POST, named = :predict) do
    global navimap
    result = Agent.predict(@params(:instruction), eval(parse(@params(:initial))), navimap)
    respond(Dict(:json => JSON.json(Dict(:path => result))))
end

max_stars_repo_name: ozanarkancan/navi-app

max_issues_repo_name: ozanarkancan/navi-app

max_forks_repo_name: ozanarkancan/navi-app


text:
# Implements Bayesian networks

module BayesianNetworks

using SPNetworks
import SPNetworks: Node, SumNode, ProductNode, LeafNode, IndicatorFunction

"Tree-shaped Bayesian network. A tree is conneted digraph where each node has at most one parent."
abstract type BayesianTree end
struct BTRoot <: BayesianTree # root node
    variable::Int
    factor::Vector{Float64} # array of unconditional probabilities
    children::Vector{BayesianTree}
end
struct BTNode <: BayesianTree # inner node
    variable::Int
    factor::Matrix{Float64} # square matrix where columns are conditional probability distributions given each value of the parent
    children::Vector{BayesianTree}
end
struct BTLeaf <: BayesianTree # leaf node
    variable::Int
    factor::Matrix{Float64}
end

"Compile Bayesian Tree into SPN."
function compile(root::BTRoot)
    # init cache of conditional SPNs
    cache = Dict{Tuple{Int,Int,Int},Int}() # (node.var,parent.var,value) -> node id
    # init cache of indicators
    icache = Dict{Tuple{Int,Int},Int}()    # (node.var,value) -> node id
    # root node
    sumnode = SumNode(Vector{UInt}(undef,size(root.factor,1)),root.factor)
    # list of nodes
    nodelist = Node[ sumnode ]
    # apply Shannon-like decomposition recursively
    for value=1:size(root.factor,1)
        # Decompose on root.variable = value
        prodnode = ProductNode([])
        push!(nodelist, prodnode)
        sumnode.children[value] = length(nodelist) # set child id
        push!(nodelist, IndicatorFunction(root.variable,value))
        push!(prodnode.children,length(nodelist))
        for child in root.children
            # recur on child
            nodeid = compile!(cache,icache,nodelist,child,root,value)
            push!(prodnode.children,nodeid)
        end
    end
    # return compiled SPN
    spn = SumProductNetwork(nodelist)
    # node ids are not topologically orders - sort in-place nodes in bfs-order
    SPNetworks.sort!(spn)
    # return sorted network
    spn
end
"
    compile!(cache::Dict{Tuple{Int,Int,Int},Int},icache::Dict{Tuple{Int,Int}},nodelist::Vector{Node},node::BayesianTree,parent,pavalue)

Compile subtree rooted at node and adds respective nodes to nodelist. Returns id of root node of subSPN.
Caches the node id given the parent id and value to avoid duplication.
"
function compile!(cache::Dict{Tuple{Int,Int,Int},Int},icache::Dict{Tuple{Int,Int}},nodelist::Vector{Node},node::BTLeaf,parent,pavalue)
    if haskey(cache,(node.variable,parent.variable,pavalue))
        return cache[(node.variable,parent.variable,pavalue)]
    end
    sumnode = SumNode(Vector{UInt}(undef,size(node.factor,1)),node.factor[:,pavalue])
    push!(nodelist, sumnode)
    cache[(node.variable,parent.variable,pavalue)] = length(nodelist)
    for value = 1:size(node.factor,1)
        push!(nodelist,IndicatorFunction(node.variable,value))
        sumnode.children[value] = length(nodelist) # set child id
    end
    # return root node id
    cache[(node.variable,parent.variable,pavalue)]
end
function compile!(cache::Dict{Tuple{Int,Int,Int},Int},icache::Dict{Tuple{Int,Int}},nodelist::Vector{Node},node::BTNode,parent,pavalue)
    if haskey(cache,(node.variable,parent.variable,pavalue))
        return cache[(node.variable,parent.variable,pavalue)]
    end
    sumnode = SumNode(Vector{UInt}(undef,size(node.factor,1)),node.factor[:,pavalue])
    push!(nodelist, sumnode)
    cache[(node.variable,parent.variable,pavalue)] = length(nodelist)
    for value=1:size(node.factor,1)
        prod = ProductNode([])
        push!(nodelist, prod)
        sumnode.children[value] = length(nodelist) # set child id
        if haskey(icache,(node.variable,value)) # recover indicator if already created
            push!(prod.children, icache[(node.variable,value)])
        else # new indicator
            push!(nodelist, IndicatorFunction(node.variable,value))
            icache[(node.variable,value)] = length(nodelist)
            push!(prod.children, length(nodelist))
        end
        for child in node.children
            # recur
            nodeid = compile!(cache,icache,nodelist,child,node,value)
            push!(prod.children, nodeid)
        end
    end
    # return root node id
    cache[(node.variable,parent.variable,pavalue)]
end

end


text:
pgfplots()
srand(3)
a1 = 1.0
H = 50
for f in 1:3
    a = 5*[2.5, 3.0, 3.0][f]
    aζ = 5*[3.0, 2.5, 3.0][f]
    lim = [22.0, 1.05, 2.0][f]
    name = "output/volatility/figprior$f.pdf"

    v = zeros(H)
    ζ = zeros(H)

    v[1] = rand(InverseGamma(a1, a1))

    for i in 2:H
        ζ[i] = rand(InverseGamma(aζ, inv(v[i-1]/aζ)))
        v[i] = rand(InverseGamma(a, inv(ζ[i]/a)))
    end

    display(plot(v, color=colorant"#0044FF", size = (300, 300/1.4), ylim=(0.0, lim), label = LaTeXString("\$\\alpha = $a, \\alpha_\\zeta = $aζ\$")))
    savefig(name)
    println(v)
end


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 0532232e-5017-11eb-2ee5-bd82db649d1c
using DataFrames

# ╔═╡ ab5fe00a-74cb-11eb-1b63-8358d1a19718
using RDatasets

# ╔═╡ 02e69f0a-5030-11eb-376c-2909777b59ef
using Random

# ╔═╡ 5e0be22e-5031-11eb-02b1-679548ac5754
begin
	using Plots
	pyplot()
end

# ╔═╡ 346e3170-5035-11eb-1d53-f3717ff82404
using PlutoUI

# ╔═╡ 878ae3d0-50ab-11eb-265d-d52ac266e8a4
using TableIO

# ╔═╡ ea900542-50ac-11eb-2869-e5694b94f504
using FreqTables

# ╔═╡ 15f6a92c-50ae-11eb-0a64-4de1f7bb490a
using Statistics

# ╔═╡ f30185a6-5015-11eb-1a8d-25603607e724
md"# Module 3: _Dataframes in Julia_

**CS6741: Statistical Foundations of Data Science**

In this module, you will get to learn one of the important packages required to process structured data - `dataframes`. We will use this throughout the course to input data, perform basic data cleaning operations, and then run statistical operations. Thus, it forms an important base for the rest of the course.

"

# ╔═╡ b66b26ee-5016-11eb-01e6-694b2e39bf5d
md"
#### Installing and using Dataframes.jl
"

# ╔═╡ 4ae89f8e-5016-11eb-16df-13d654d299d5
md"We will be using the `Dataframes.jl` package. To install this package in your local Julia instance do the following:"

# ╔═╡ 5f4fdd02-5016-11eb-201d-1d9426646956
md"
```
using Pkg
Pkg.add(\"DataFrames\")
```
"

# ╔═╡ 16b7eac0-5017-11eb-15e5-b9f201470cbc
md"Notice that you do not have to quit a Pluto session to install a package. You can do it in another terminal and once the package is installed it would be available to the active Pluto session."

# ╔═╡ c8529ea0-5016-11eb-38d9-c1ecc5550d2d
md"Like any other package, we need to add it to the current Julia runtime:"

# ╔═╡ b3d48e04-5016-11eb-28f2-c71e4540a0dc
md"#### Creating dataframes"

# ╔═╡ 35fabd48-5017-11eb-1a67-e90302279f9e
md"Throughout this course and in most practical applications too, you will encounter relational data which can be laid out in a table where each row is a data item and columns represent different fields or attributes of a data item."

# ╔═╡ 4efdd714-5017-11eb-05d3-6dcec7617db1
md"This can be represented in the `DataFrame` data type. We create below our first dataframe."

# ╔═╡ 5ec1b85a-5017-11eb-1652-877f43479ca0
df = DataFrame(Number = 1:9, IsPrime = ["F", "T", "T", "F", "T", "F", "T", "F", "F"])

# ╔═╡ b31b85fe-5017-11eb-22fa-13762810aa85
md"There are many other ways of creating a dataframe. The following does it column wise."

# ╔═╡ c0d7126a-5017-11eb-2af4-89db7d01e57a
begin
	df2 = DataFrame()
	df2.Number = 1:9
	df2.IsPrime = ["F", "T", "T", "F", "T", "F", "T", "F", "F"]
	df2
end

# ╔═╡ 083152f2-74cc-11eb-3882-8b79170cd7aa
md"👉 Create a dataframe of the details of the planets in the solar system such as size, weight, and periods of rotation and revolution."

# ╔═╡ ec081484-5017-11eb-05a9-75fdf117f284
md"The following does it row-wise."

# ╔═╡ f2cdcade-5017-11eb-3ecf-59e0476e76e7
begin
	df3 = DataFrame(Number = Int[], IsPrime = String[])
	push!(df3, (1, "F")) # as a tuple
	push!(df3, [2, "T"]) # as a vector
	push!(df3, Dict(:Number => 3, :IsPrime => "T")) # as a dictionary
	push!(df3, (4, "F"))
	push!(df3, (5, "T"))
	push!(df3, (6, "F"))
	push!(df3, (7, "T"))
	push!(df3, (8, "F"))
	push!(df3, (9, "F"))
	df3
end

# ╔═╡ 1feb0404-5019-11eb-17c6-cdf37432a2b0
md"Certain cells may have missing data. These can be explicitly stated."

# ╔═╡ 2dcaacca-5019-11eb-3691-b99785b2dce9
df4 = DataFrame(
	RollNo     = [6, 28, 496],
	Name       = ["Arun", "Arunima", "Jacob"],
	FavNo      = [π, exp(1), missing],
	State      = ["TN", "J&K", "TN"],
	Attendance = [missing, 18/21, 14/24]
)

# ╔═╡ a43be892-5017-11eb-1962-d30a900d1bae
md"We have so far seen the usual data types - integers, floats, strings, boolean. But often in statistics we are also interested in _categorical data_, i.e., data which belongs on one of a countable number of categories. An example of a categorical data is the `State` field in the above table."

# ╔═╡ c42539c2-5019-11eb-2521-f978d2dc23a6
df5 = DataFrame(
	RollNo     = [6, 28, 496],
	Name       = ["Arun", "Arunima", "Jacob"],
	FavNo      = [π, exp(1), missing],
	State      = categorical(["TN", "J&K", "TN"]),
	Attendance = [missing, 18/21, 14/24]
)

# ╔═╡ 5f569306-74cc-11eb-1efe-671de09ad499
md"👉 Many datasets are available in the [RDatasets package](https://github.com/JuliaStats/RDatasets.jl). Use it to read the [diamonds data](https://bookdown.org/yih_huynh/Guide-to-R-Book/diamonds.html) into a dataframe. Which of these columns are categorical?"

# ╔═╡ 1bd6c8c0-7969-11eb-3f1c-17cbe159a26d
diamonds = dataset("ggplot2", "diamonds")

# ╔═╡ fd412c90-797a-11eb-1f71-075ec205a39f
describe(diamonds)

# ╔═╡ f005c3d0-501a-11eb-20e4-1f42902ba037
md"#### Selecting parts of a dataframe"

# ╔═╡ f9b80a8c-501a-11eb-2b51-8b7ca90b99fc
md"One of the basic functions on a dataframe is to select a subset of rows and columns based on specific selection criteria."

# ╔═╡ 13bac0d6-501c-11eb-0c0e-1d91358a121b
md"In the following we discuss simple row and column range selectors."

# ╔═╡ 15885816-501b-11eb-3070-ff59092748e8
df[1:3, :]

# ╔═╡ 75939448-501b-11eb-2a93-b1679d178b36
df[[1, 3, 5, 7], :]

# ╔═╡ 7c44c238-501b-11eb-2ce0-81e458713d63
df[:, [:IsPrime]]

# ╔═╡ 8be7f980-501b-11eb-3f71-879d6490e0a6
df[1:4, [:IsPrime]]

# ╔═╡ bd25434a-501b-11eb-2251-b92ecfd1e7b7
df[!, :IsPrime]

# ╔═╡ d84c0596-501b-11eb-0a10-a768b0538384
df[4, :IsPrime]

# ╔═╡ 1c6ed2bc-501c-11eb-37a0-b7cecc35c6fe
md"We can also use a _regular expressions_ to select columns. For instance to select all columns with `No` in them."

# ╔═╡ 53458446-501c-11eb-10df-85efdd2c8e78
df5

# ╔═╡ 25de35c2-501c-11eb-2666-5f73b1e2b004
df5[!, r"te"]

# ╔═╡ a02b1e7e-74cc-11eb-1400-7f702301e663
md"👉 In the diamonds dataset, filter the diamonds based on regular expressions on the cut."

# ╔═╡ aefc5360-7962-11eb-3e23-fbb46cdf92dd
diamonds[!,r"Cut"]

# ╔═╡ add1dd20-7962-11eb-1e18-7d4de9e5f0c7

# ╔═╡ 78f9d638-5028-11eb-0378-f52a1d091d11
md"We can also select rows based on conditions on the data fields"

# ╔═╡ a2cfe2ce-5028-11eb-1341-892384a89e52
df[df.Number .> 4, :]

# ╔═╡ 5d86370e-501c-11eb-15bb-5f0095fc5e41
df5[df5.State .== "TN", :]

# ╔═╡ be611380-5028-11eb-2e4e-a956b61e501f
df5[(df5.State .== "TN") .& (df5.RollNo .> 100), :]

# ╔═╡ bae24224-74cc-11eb-1235-29eec31f5fe0
md"👉 In the diamonds dataset, filter the diamonds based on a certain size lower bound and certain price upper-bound."

# ╔═╡ 41e0254e-7ab5-11eb-1794-0f11687dc655
describe(diamonds)

# ╔═╡ 41376230-7ab5-11eb-3a37-19897d26a0e8
diamonds[(diamonds.Price.<3993) .& (diamonds.Price.>3991),:]

# ╔═╡ 0dc5a0f2-502f-11eb-23b0-89669cb10fc5
md"If we want to create a _new_ data frame from a part of it, then we can use the `select` keyword (think of `select` statements in SQL). Syntax is quite similar."

# ╔═╡ 36703c42-502f-11eb-01fe-d1b6e32454db
select(df5, r"No")

# ╔═╡ 58cfd450-502f-11eb-01e9-e77d7d31dcf5
md"We can transform different aspects including column names with the `select` statement."

# ╔═╡ a6416e74-502f-11eb-3763-d71db19aa35d
df

# ╔═╡ 6819c894-502f-11eb-1502-0b8abafed9a5
select(df, :Number, :IsPrime => (x -> x .== "T"))

# ╔═╡ d36bfbaa-74cc-11eb-370d-f387239ee7ca
md"👉 In the diamonds dataset, convert the price column to a categorical variable depending with three categories: Low, Medium, High, based on your chosen thresholds of price."

# ╔═╡ 61945f60-7b3c-11eb-1574-7da8ce075837
select(diamonds,:Price,:Price=>(x->(x.<1000) .& (x.>500)) =>:Threshold)

# ╔═╡ 53a82460-7b3a-11eb-3461-c90925f5a6be

# ╔═╡ 4c025570-7ab6-11eb-193a-99c726f6eac9
md"select(diamonds,:Price,categorical!(diamonds,:Price=>(x-> x.<1000 ? Low : x.>3000 ? High : Medium)=>:Range))"

# ╔═╡ 37970c30-7b78-11eb-2783-b7e60a8dc32a
begin
	a = [2,4,10,4,8,8]
	b = [5,9,7,2,8,7]
	c = [2,9,7,6,8,1]

	dfd = DataFrame(A = a, B = b, C = c)
end

# ╔═╡ 4516dbb0-7b78-11eb-052f-0d0582e70faa
select(dfd,:A,:B,:C,AsTable(:) => ByRow(extrema)=> [:MaxElem,:MinElem])

# ╔═╡ f85c27da-502f-11eb-3628-d9e4fba3a36a
md"More complex transformations are also possible and are often used."

# ╔═╡ 09dacec6-5030-11eb-1420-2fa4b4384a7b
Random.seed!(1)

# ╔═╡ 0e9da3e8-5030-11eb-2595-09e52e9dc5bb
df7 = DataFrame(rand(5, 3), [:a, :b, :c])

# ╔═╡ 6dd0ed84-5030-11eb-28ec-416ed6ef4e00
select(df7, :a, :b, :c => (x -> x .- mean(x)) => :c_offset)

# ╔═╡ 92673a22-5030-11eb-0657-7583d5dfda62
md"Doing the same thing with the `transform` keyword retains the existing columns"

# ╔═╡ f701ade2-74cc-11eb-27cf-cfeb48f63662
md"👉 For the diamonds dataset, the price column is in US dollars, convert them into INR."

# ╔═╡ d3fdf860-7ab8-11eb-3bc0-f1e6b2e52eec
select(diamonds,:Price,:Price=>(x->x.*75)=>:INR)

# ╔═╡ a2dba6b8-5030-11eb-1bdb-656ac06ecd58
md"The above functions are per column. For row-wise transformations we have the `ByRow` keyword."

# ╔═╡ 8807526a-5030-11eb-2476-411d7af8938d
select(df7, :a, :a => ByRow(sqrt))

# ╔═╡ d48515da-5030-11eb-2e6d-b51cc9f54c40
select(df7, :a, :a => ByRow(x -> 1/(1 + exp(-x))))

# ╔═╡ 8b74e2f8-74ce-11eb-31c2-6f4a42953d89
md"👉 Do you recognise the function above?"

# ╔═╡ 386e6730-7ab9-11eb-256b-9f879495ed47
md"Ans:Sigmoid"

# ╔═╡ 6c7039f8-5031-11eb-1049-1d1bef7cbf3e
plot(x->x, x->1/(1 + exp(-x)), -5, 5, legend = false, line=5)

# ╔═╡ bffd0eea-74ce-11eb-18c1-43dc5a6eec6f
md"We can also use the select command to create multiple columns in a go, for instance finding extrema row-wise."

# ╔═╡ a8bc02ca-5031-11eb-016a-6974fbe0a879
df7

# ╔═╡ b96cb2e0-5031-11eb-1be8-f90e124225f6
select(df7, AsTable(:) => ByRow(extrema) => [:lo, :hi])

# ╔═╡ c62c16f4-5031-11eb-11e1-e181d156670e
extrema([3, 2, -1, 0])

# ╔═╡ d2d84480-74ce-11eb-0f13-f1bef9d197f5
md"Extrema is a generic function you can use with any list and returns tuples. So, in the select command you could have had any other custom function which returns tuples."

# ╔═╡ 06fe9d18-74cf-11eb-0723-1dd875c22fea
md"We can also transform dataframes by outputing columns. Note that column names (like variable names) are prefixed with `:`"

# ╔═╡ d8a720dc-5031-11eb-04ac-817b4512d51a
transform(df7, AsTable(:) => ByRow(argmax) => :prediction)

# ╔═╡ 31e4d220-5032-11eb-2a98-671624a40b07
string(transform(df7, AsTable(:) => ByRow(argmax) => :prediction)[1, :prediction])

# ╔═╡ e7f28678-74ce-11eb-260a-ef3b6576bcc2
md"👉 In the diamonds dataframe, write a select command to create two new columns: `Largest` - which calculates the largest dimension amongst x, y, z, and `Largest axis` which is either x, y, z along which the largest dimension was found."

# ╔═╡ 5380a4b8-5032-11eb-18b1-b5aa0ec65677
md"#### Joins"

# ╔═╡ 5944e882-5032-11eb-150b-6153c7bbccd5
md"A common action to be performed in data manipulation is to combine the data across multiple dataframes. The different ways of joining data forms an important part in any database course. For a quick overview see [this](https://en.wikipedia.org/wiki/Join_(SQL)) helpfully detailed page on Wikipedia"

# ╔═╡ 1c0c7f38-5033-11eb-2026-05c96439cd4b
players = DataFrame(Jersey = [10, 8, 19], Name = ["Sachin Tendulkar", "Anil Kumble", "Rahul Dravid"])

# ╔═╡ 75b13424-5034-11eb-2103-d7b7e0636164
roles = DataFrame(Jersey = [10, 8, 24], Job = ["Batsman", "Leg spinner", "Captain"])

# ╔═╡ a8efdb18-5034-11eb-1d61-335fd069e273
innerjoin(players, roles, on = :Jersey)

# ╔═╡ 886d70fe-5035-11eb-287e-4bdd2388c053
md"Let us use Pluto's interactivity and some cool Julia syntax to try out different other types of joins"

# ╔═╡ 49a45892-74cf-11eb-3487-9501aab49a7e
md"Try out each of the possible joins below in the dropdown and see its effect."

# ╔═╡ 39b697f6-50a8-11eb-301a-0d829c9ea14d
begin
	prompt = Dict("innerjoin" => "The output contains rows for values of the key that exist in all passed data frames", "leftjoin" => "The output contains rows for values of the key that exist in the first (left) argument, whether or not that value exists in the second (right) argument.", "rightjoin" => "The output contains rows for values of the key that exist in the second (right) argument, whether or not that value exists in the first (left) argument.", "outerjoin" => "The output contains rows for values of the key that exist in any of the passed data frames.", "semijoin" => "Like an inner join, but output is restricted to columns from the first (left) argument.", "antijoin" => "The output contains rows for values of the key that exist in the first (left) but not the second (right) argument. As with semijoin, output is restricted to columns from the first (left) argument.")
end

# ╔═╡ 38fa429a-5035-11eb-2c28-95752c218eb1
@bind jointype Select(["innerjoin", "leftjoin", "rightjoin", "outerjoin", "semijoin", "antijoin"])

# ╔═╡ 8d878060-50a9-11eb-1e5b-1d46abf93d19
begin
	output = prompt[jointype]
	md"$output"
end

# ╔═╡ 77b0aab0-5035-11eb-3203-61cf9603c09f
getfield(DataFrames, Symbol(jointype))(players, roles, on = :Jersey)

# ╔═╡ e83b5134-5037-11eb-31f9-93a2e545a174
md"#### Split-Apply-Combine"

# ╔═╡ f1cb42a4-5037-11eb-19b3-11e1cbc3898c
md"
Many data operations involve a sequence of the following three operations:
1. _Split_ data set into groups
2. _Apply_ some functions or transformations on each group separately
3. _Combine_ results from different groups
"

# ╔═╡ 4cc32816-5038-11eb-391f-651889126b66
md"This process is described in great detail (with graphics and examples) in [this](https://www.jstatsoft.org/index.php/jss/article/view/v040i01/v40i01.pdf) paper."

# ╔═╡ 28b04e72-50ab-11eb-39fa-01758aed2871
md"To use files in Pluto, we can use a file picker."

# ╔═╡ cafde6d2-74cb-11eb-2dda-2f895a328dd2
md"You can use Pluto to work with local files, by adding a filepicker as shown below"

# ╔═╡ 246ca6ee-50ab-11eb-3d58-81843829a32b
@bind f PlutoUI.FilePicker()

# ╔═╡ d6b2715a-74cb-11eb-2f8a-b9f6a20895f3
md"For this tutorial, we will be using the RDatasets package which has many of the interesting datasets for us"

# ╔═╡ b65971f6-74cb-11eb-3cb1-43a033b360d3
Iris = dataset("datasets", "iris")

# ╔═╡ cf2a0ff4-50ab-11eb-363d-b7aa7467876b
md"We can summarise dataframes to quickly get a sense of what the columns are like. It calculates statistics, but more on this later."

# ╔═╡ e9f09aa6-50ab-11eb-0973-bda264b6637e
describe(Iris)

# ╔═╡ 538609f6-50ac-11eb-3f36-dbc78ef740f5
md"This is not quite helpful for `string` fields. We can check the unique values with the `unique` command."

# ╔═╡ 2cf9f3d8-50ac-11eb-33ed-3fd3cfdd91c0
unique(Iris.Species)

# ╔═╡ fef98472-50ac-11eb-0992-e50602b5667c
md"To compute frequency counts, one of the good options (that works well with DataFrames) is the `FreqTables.jl` library."

# ╔═╡ a8b8646c-50ac-11eb-1b8c-03ee856da36c
freqtable(Iris.Species)

# ╔═╡ 67f23202-50ac-11eb-0f5c-53c2579f8b93
md"But we may also want to convert this column into a categorical column. We can do this inplace with `categorical!`"

# ╔═╡ 086fe27a-50ac-11eb-2c36-73cae74cd4ea
IrisCat = categorical(Iris, :Species, compress=true)

# ╔═╡ 507fc780-50be-11eb-3bf9-9b53314547fc
freqtable(IrisCat.Species)

# ╔═╡ 62ab88a6-50ad-11eb-14a1-4f13505743c5
md"Back to split-combine-apply. First we can `groupby` the data on a specific field."

# ╔═╡ 54a8be54-50ad-11eb-15f0-fb33b5e890f9
gdf = groupby(Iris, :Species)

# ╔═╡ e5d124e8-50ad-11eb-3929-ffbe4ad5a34d
md"Once we have the data split as per `Species`, we want to apply a function and then combine the results. One simple function is to compute the `mean`. We will use the `Statistics` package to do this. This is a package we will use a lot."

# ╔═╡ 9870c276-50ad-11eb-0bbc-69a75ca01ca6
combine(gdf, :PetalLength => mean)

# ╔═╡ 02a854b8-50b0-11eb-36b6-5391455c6067
md"We can also compute multiple columns using for instance the `extrema` function that we had seen earlier."

# ╔═╡ 113fc6d2-50b0-11eb-1bd1-83ad2189b1f3
combine(gdf, :PetalLength => (x -> [extrema(x)]) => [:min, :max])

# ╔═╡ 3af1fd76-50ae-11eb-233d-c1900d65f6bd
md"We can create multiple apply functions and then combine them together."

# ╔═╡ 56fd3198-50ae-11eb-3545-79baa8e18c80
combine(gdf, :SepalLength => mean, :SepalWidth => mean, :PetalLength => mean)

# ╔═╡ c570e282-50ae-11eb-2907-3d5bfc5fb6cb
md"We can also apply functions across columns"

# ╔═╡ ccb232da-50ae-11eb-369d-fdfbc59638c6
combine(gdf, [:PetalLength, :SepalLength] => ((p, s) -> (a=mean(p)/mean(s))) => :two_lengths)

# ╔═╡ 2703f45e-50b0-11eb-139f-dfc3bf1de490
md"We can also access different parts of a grouped structure individually"

# ╔═╡ f94d01dc-50b3-11eb-00d9-c3bcc9cba289
gdf[1]

# ╔═╡ 7b390ed8-75e3-11eb-343c-b38b5fd416cc
md"👉 Try out group by two different columns"

# ╔═╡ 2694f998-75c8-11eb-280d-61943feeb8c3
md"### Differently organising the dataframe"

# ╔═╡ 63c723b4-50b5-11eb-0d5a-651ed812b699
md"There are different ways in which data can formatted. The common one is what we saw where each row is a flower and columns are all its properties."

# ╔═╡ 61c7feb8-50b4-11eb-2d14-a5a09c2cbadf
Iris

# ╔═╡ eab1b56a-50b5-11eb-362f-f72b9a5b9ad5
md"Let us add to this table an index which denotes a unique number of each flower."

# ╔═╡ f74f21f4-50b5-11eb-01ea-f32f27e075c0
insertcols!(Iris, 1, :Index => 1:150, makeunique=true)

# ╔═╡ ccb7540c-50b5-11eb-0584-1da2140322ce
md"An alternative is to have each data-item in a single row"

# ╔═╡ 5c03b2c4-50b4-11eb-3870-13459725f829
IrisStacked = stack(Iris, 2:6, :Index)

# ╔═╡ 7de8c59e-50bb-11eb-16af-4b02ba7c577a
md"A better way to see this would be to sort the dataframe"

# ╔═╡ 5bfbe84c-50bb-11eb-17db-15ac85f02f03
sort(IrisStacked)

# ╔═╡ 88a488ba-50bb-11eb-1e10-a766f21a86d8
md"Now we can see that every row contains one data-point and the details of a single flower are spread across 5 rows."

# ╔═╡ ca822d50-50bb-11eb-391b-b18f27cbe6cd
md"So with the `stack` command, we have moved from a fat short matrix to a thin tall matrix. We can now move in the reverse direction with the `unstack` command."

# ╔═╡ b0141668-50bb-11eb-1948-574d89edf39c
unstack(IrisStacked, :Index, :variable, :value)

# ╔═╡ 3a6e3e66-75c8-11eb-15b3-656eba412155
md"👉 Think through the advantages and disadvantages of the _long_ and _wide_ formats"

# ╔═╡ d74c4188-75c8-11eb-2098-5154990736ba
md"For a good overview of these choices, it is recommended that you read [this paper titled 'Tidy Data'](https://vita.had.co.nz/papers/tidy-data.pdf)"

# ╔═╡ Cell order:
# ╟─f30185a6-5015-11eb-1a8d-25603607e724
# ╟─b66b26ee-5016-11eb-01e6-694b2e39bf5d
# ╟─4ae89f8e-5016-11eb-16df-13d654d299d5
# ╟─5f4fdd02-5016-11eb-201d-1d9426646956
# ╟─16b7eac0-5017-11eb-15e5-b9f201470cbc
# ╟─c8529ea0-5016-11eb-38d9-c1ecc5550d2d
# ╠═0532232e-5017-11eb-2ee5-bd82db649d1c
# ╟─b3d48e04-5016-11eb-28f2-c71e4540a0dc
# ╟─35fabd48-5017-11eb-1a67-e90302279f9e
# ╟─4efdd714-5017-11eb-05d3-6dcec7617db1
# ╠═5ec1b85a-5017-11eb-1652-877f43479ca0
# ╟─b31b85fe-5017-11eb-22fa-13762810aa85
# ╠═c0d7126a-5017-11eb-2af4-89db7d01e57a
# ╟─083152f2-74cc-11eb-3882-8b79170cd7aa
# ╟─ec081484-5017-11eb-05a9-75fdf117f284
# ╠═f2cdcade-5017-11eb-3ecf-59e0476e76e7
# ╟─1feb0404-5019-11eb-17c6-cdf37432a2b0
# ╠═2dcaacca-5019-11eb-3691-b99785b2dce9
# ╟─a43be892-5017-11eb-1962-d30a900d1bae
# ╠═c42539c2-5019-11eb-2521-f978d2dc23a6
# ╟─5f569306-74cc-11eb-1efe-671de09ad499
# ╠═ab5fe00a-74cb-11eb-1b63-8358d1a19718
# ╠═1bd6c8c0-7969-11eb-3f1c-17cbe159a26d
# ╠═fd412c90-797a-11eb-1f71-075ec205a39f
# ╟─f005c3d0-501a-11eb-20e4-1f42902ba037
# ╟─f9b80a8c-501a-11eb-2b51-8b7ca90b99fc
# ╟─13bac0d6-501c-11eb-0c0e-1d91358a121b
# ╠═15885816-501b-11eb-3070-ff59092748e8
# ╠═75939448-501b-11eb-2a93-b1679d178b36
# ╠═7c44c238-501b-11eb-2ce0-81e458713d63
# ╠═8be7f980-501b-11eb-3f71-879d6490e0a6
# ╠═bd25434a-501b-11eb-2251-b92ecfd1e7b7
# ╠═d84c0596-501b-11eb-0a10-a768b0538384
# ╟─1c6ed2bc-501c-11eb-37a0-b7cecc35c6fe
# ╠═53458446-501c-11eb-10df-85efdd2c8e78
# ╠═25de35c2-501c-11eb-2666-5f73b1e2b004
# ╟─a02b1e7e-74cc-11eb-1400-7f702301e663
# ╠═aefc5360-7962-11eb-3e23-fbb46cdf92dd
# ╠═add1dd20-7962-11eb-1e18-7d4de9e5f0c7
# ╟─78f9d638-5028-11eb-0378-f52a1d091d11
# ╠═a2cfe2ce-5028-11eb-1341-892384a89e52
# ╠═5d86370e-501c-11eb-15bb-5f0095fc5e41
# ╠═be611380-5028-11eb-2e4e-a956b61e501f
# ╟─bae24224-74cc-11eb-1235-29eec31f5fe0
# ╠═41e0254e-7ab5-11eb-1794-0f11687dc655
# ╠═41376230-7ab5-11eb-3a37-19897d26a0e8
# ╟─0dc5a0f2-502f-11eb-23b0-89669cb10fc5
# ╠═36703c42-502f-11eb-01fe-d1b6e32454db
# ╟─58cfd450-502f-11eb-01e9-e77d7d31dcf5
# ╠═a6416e74-502f-11eb-3763-d71db19aa35d
# ╠═6819c894-502f-11eb-1502-0b8abafed9a5
# ╟─d36bfbaa-74cc-11eb-370d-f387239ee7ca
# ╠═61945f60-7b3c-11eb-1574-7da8ce075837
# ╠═53a82460-7b3a-11eb-3461-c90925f5a6be
# ╟─4c025570-7ab6-11eb-193a-99c726f6eac9
# ╠═37970c30-7b78-11eb-2783-b7e60a8dc32a
# ╠═4516dbb0-7b78-11eb-052f-0d0582e70faa
# ╟─f85c27da-502f-11eb-3628-d9e4fba3a36a
# ╠═02e69f0a-5030-11eb-376c-2909777b59ef
# ╠═09dacec6-5030-11eb-1420-2fa4b4384a7b
# ╠═0e9da3e8-5030-11eb-2595-09e52e9dc5bb
# ╠═6dd0ed84-5030-11eb-28ec-416ed6ef4e00
# ╟─92673a22-5030-11eb-0657-7583d5dfda62
# ╠═f701ade2-74cc-11eb-27cf-cfeb48f63662
# ╠═d3fdf860-7ab8-11eb-3bc0-f1e6b2e52eec
# ╟─a2dba6b8-5030-11eb-1bdb-656ac06ecd58
# ╠═8807526a-5030-11eb-2476-411d7af8938d
# ╠═d48515da-5030-11eb-2e6d-b51cc9f54c40
# ╟─8b74e2f8-74ce-11eb-31c2-6f4a42953d89
# ╟─386e6730-7ab9-11eb-256b-9f879495ed47
# ╠═5e0be22e-5031-11eb-02b1-679548ac5754
# ╠═6c7039f8-5031-11eb-1049-1d1bef7cbf3e
# ╟─bffd0eea-74ce-11eb-18c1-43dc5a6eec6f
# ╠═a8bc02ca-5031-11eb-016a-6974fbe0a879
# ╠═b96cb2e0-5031-11eb-1be8-f90e124225f6
# ╠═c62c16f4-5031-11eb-11e1-e181d156670e
# ╟─d2d84480-74ce-11eb-0f13-f1bef9d197f5
# ╟─06fe9d18-74cf-11eb-0723-1dd875c22fea
# ╠═d8a720dc-5031-11eb-04ac-817b4512d51a
# ╠═31e4d220-5032-11eb-2a98-671624a40b07
# ╟─e7f28678-74ce-11eb-260a-ef3b6576bcc2
# ╟─5380a4b8-5032-11eb-18b1-b5aa0ec65677
# ╟─5944e882-5032-11eb-150b-6153c7bbccd5
# ╠═1c0c7f38-5033-11eb-2026-05c96439cd4b
# ╠═75b13424-5034-11eb-2103-d7b7e0636164
# ╠═a8efdb18-5034-11eb-1d61-335fd069e273
# ╟─886d70fe-5035-11eb-287e-4bdd2388c053
# ╠═346e3170-5035-11eb-1d53-f3717ff82404
# ╟─49a45892-74cf-11eb-3487-9501aab49a7e
# ╟─39b697f6-50a8-11eb-301a-0d829c9ea14d
# ╠═38fa429a-5035-11eb-2c28-95752c218eb1
# ╟─8d878060-50a9-11eb-1e5b-1d46abf93d19
# ╠═77b0aab0-5035-11eb-3203-61cf9603c09f
# ╟─e83b5134-5037-11eb-31f9-93a2e545a174
# ╟─f1cb42a4-5037-11eb-19b3-11e1cbc3898c
# ╟─4cc32816-5038-11eb-391f-651889126b66
# ╟─28b04e72-50ab-11eb-39fa-01758aed2871
# ╟─cafde6d2-74cb-11eb-2dda-2f895a328dd2
# ╠═246ca6ee-50ab-11eb-3d58-81843829a32b
# ╟─d6b2715a-74cb-11eb-2f8a-b9f6a20895f3
# ╠═b65971f6-74cb-11eb-3cb1-43a033b360d3
# ╠═878ae3d0-50ab-11eb-265d-d52ac266e8a4
# ╟─cf2a0ff4-50ab-11eb-363d-b7aa7467876b
# ╠═e9f09aa6-50ab-11eb-0973-bda264b6637e
# ╟─538609f6-50ac-11eb-3f36-dbc78ef740f5
# ╠═2cf9f3d8-50ac-11eb-33ed-3fd3cfdd91c0
# ╟─fef98472-50ac-11eb-0992-e50602b5667c
# ╠═ea900542-50ac-11eb-2869-e5694b94f504
# ╠═a8b8646c-50ac-11eb-1b8c-03ee856da36c
# ╟─67f23202-50ac-11eb-0f5c-53c2579f8b93
# ╠═086fe27a-50ac-11eb-2c36-73cae74cd4ea
# ╠═507fc780-50be-11eb-3bf9-9b53314547fc
# ╟─62ab88a6-50ad-11eb-14a1-4f13505743c5
# ╠═54a8be54-50ad-11eb-15f0-fb33b5e890f9
# ╟─e5d124e8-50ad-11eb-3929-ffbe4ad5a34d
# ╠═15f6a92c-50ae-11eb-0a64-4de1f7bb490a
# ╠═9870c276-50ad-11eb-0bbc-69a75ca01ca6
# ╟─02a854b8-50b0-11eb-36b6-5391455c6067
# ╠═113fc6d2-50b0-11eb-1bd1-83ad2189b1f3
# ╟─3af1fd76-50ae-11eb-233d-c1900d65f6bd
# ╠═56fd3198-50ae-11eb-3545-79baa8e18c80
# ╟─c570e282-50ae-11eb-2907-3d5bfc5fb6cb
# ╠═ccb232da-50ae-11eb-369d-fdfbc59638c6
# ╟─2703f45e-50b0-11eb-139f-dfc3bf1de490
# ╠═f94d01dc-50b3-11eb-00d9-c3bcc9cba289
# ╟─7b390ed8-75e3-11eb-343c-b38b5fd416cc
# ╟─2694f998-75c8-11eb-280d-61943feeb8c3
# ╟─63c723b4-50b5-11eb-0d5a-651ed812b699
# ╠═61c7feb8-50b4-11eb-2d14-a5a09c2cbadf
# ╟─eab1b56a-50b5-11eb-362f-f72b9a5b9ad5
# ╠═f74f21f4-50b5-11eb-01ea-f32f27e075c0
# ╟─ccb7540c-50b5-11eb-0584-1da2140322ce
# ╠═5c03b2c4-50b4-11eb-3870-13459725f829
# ╟─7de8c59e-50bb-11eb-16af-4b02ba7c577a
# ╠═5bfbe84c-50bb-11eb-17db-15ac85f02f03
# ╟─88a488ba-50bb-11eb-1e10-a766f21a86d8
# ╟─ca822d50-50bb-11eb-391b-b18f27cbe6cd
# ╠═b0141668-50bb-11eb-1948-574d89edf39c
# ╟─3a6e3e66-75c8-11eb-15b3-656eba412155
# ╟─d74c4188-75c8-11eb-2098-5154990736ba


text:
include("trajectories/trajectories.jl")
include("named_policy.jl")
include("agent.jl")
include("multi_agent.jl")


text:
immutable PlaneWaveHH3DTD{T,P,F} <: TDFunctional{T}
    direction::P
    speed_of_light::T
    signature::F
    amplitude::T
end

#speedoflight(f::PlaneWaveHH3DTD) = f.speed_of_light

function planewave(direction, speedoflight::Number, signature, amplitude=one(speedoflight)::Number)
    PlaneWaveHH3DTD(direction, speedoflight, signature, amplitude)
end

*(a, f::PlaneWaveHH3DTD) = PlaneWaveHH3DTD(f.direction, f.speed_of_light, f.signature, a*f.amplitude)

function(f::PlaneWaveHH3DTD)(r,t)
    k = f.direction
    u = sum(k[i]*r[i] for i in 1:length(k))
    #u = dot(f.direction, r)
    a = f.amplitude
    c = f.speed_of_light
    h = f.signature
    a * h(c*t - u)
end


text:
parent_type(::Type{AlgMatElem{T, S, Mat}}) where {T, S, Mat} = S

@doc Markdown.doc"""
    matrix(a::AlgMatElem; copy::Bool = true) -> MatElem

Returns the matrix which defines $a$.
"""
function matrix(a::AlgMatElem; copy::Bool = true)
  if copy
    return deepcopy(a.matrix)
  else
    return a.matrix
  end
end

################################################################################
#
#  Coefficients
#
################################################################################

function assure_has_coeffs(a::AlgMatElem)
  if a.has_coeffs
    return nothing
  end

  A = parent(a)
  Ma = matrix(a)
  b, c = _check_matrix_in_algebra(Ma, A)
  @assert b "The element is not in the algebra."
  a.coeffs = c
  a.has_coeffs = true
  return nothing
end

@doc Markdown.doc"""
    coefficients(a::AlgMatElem; copy::Bool = true) -> Vector{RingElem}

Returns the coefficients of $a$ in the basis of `algebra(a)`.
"""
function coefficients(a::AlgMatElem; copy::Bool = true)
  assure_has_coeffs(a)
  if copy
    return deepcopy(a.coeffs)
  else
    return a.coeffs
  end
end

################################################################################
#
#  String I/O
#
################################################################################

function show(io::IO, a::AlgMatElem)
  show(IOContext(io, :compact => true), matrix(a, copy = false))
end

################################################################################
#
#  Unary operations
#
################################################################################

@doc Markdown.doc"""
    -(a::AlgMatElem) -> AlgMatElem

Returns $-a$.
"""
function -(a::AlgMatElem)
  b = parent(a)(-matrix(a, copy = false))
  if a.has_coeffs
    b.coeffs = [ -coefficients(a, copy = false)[i] for i = 1:dim(parent(a)) ]
    b.has_coeffs = true
  end
  return b
end

################################################################################
#
#  Binary operations
#
################################################################################

@doc Markdown.doc"""
    +(a::AlgMatElem, b::AlgMatElem) -> AlgMatElem

Return $a + b$.
"""
function +(a::AlgMatElem{T, S, V}, b::AlgMatElem{T, S, V}) where {T, S, V}
  parent(a) != parent(b) && error("Parents don't match.")
  c = parent(a)(matrix(a, copy = false) + matrix(b, copy = false))
  if a.has_coeffs && b.has_coeffs
    c.coeffs = [ coefficients(a, copy = false)[i] + coefficients(b, copy = false)[i] for i = 1:dim(parent(a)) ]
    c.has_coeffs = true
  end
  return c
end

@doc Markdown.doc"""
    -(a::AlgMatElem, b::AlgMatElem) -> AlgMatElem

Return $a - b$.
"""
function -(a::AlgMatElem{T, S, V}, b::AlgMatElem{T, S, V}) where {T, S, V}
  parent(a) != parent(b) && error("Parents don't match.")
  c = parent(a)(matrix(a, copy = false) - matrix(b, copy = false))
  if a.has_coeffs && b.has_coeffs
    c.coeffs = [ coefficients(a, copy = false)[i] - coefficients(b, copy = false)[i] for i = 1:dim(parent(a)) ]
    c.has_coeffs = true
  end
  return c
end

@doc Markdown.doc"""
    *(a::AlgMatElem, b::AlgMatElem) -> AlgMatElem

Return $a \cdot b$.
"""
function *(a::AlgMatElem{T, S, V}, b::AlgMatElem{T, S, V}) where {T, S, V}
  parent(a) != parent(b) && error("Parents don't match.")
  return parent(a)(matrix(a, copy = false)*matrix(b, copy = false))
end

################################################################################
#
#  Unsafe operations
#
################################################################################

function zero!(a::AlgMatElem)
  a.matrix = zero!(a.matrix)
  return a
end

function add!(c::AlgMatElem{T, S, V}, a::AlgMatElem{T, S, V}, b::AlgMatElem{T, S, V}) where {T, S, V}
  parent(a) != parent(b) && error("Parents don't match.")
  parent(c) != parent(b) && error("Parents don't match.")
  A = parent(a)

  if c === a || c === b
    d = A()
    d = add!(d, a, b)
    return d
  end

  c.matrix = add!(c.matrix, a.matrix, b.matrix)
  c.has_coeffs = false

  return c
end

function mul!(c::AlgMatElem{T, S, V}, a::AlgMatElem{T, S, V}, b::AlgMatElem{T, S, V}) where {T, S, V}
  parent(a) != parent(b) && error("Parents don't match.")
  A = parent(a)

  if c === a || c === b
    d = parent(a)()
    mul!(d, a, b)
    return d
  end

  c.matrix = mul!(c.matrix, a.matrix, b.matrix)
  c.has_coeffs = false

  return c
end

################################################################################
#
#  Ad hoc operations
#
################################################################################

function *(a::AlgMatElem, b::T) where {T <: RingElem}
  A = parent(a)
  if parent(b) == base_ring(A)
    b = coefficient_ring(A)(b)
  end
  return A(matrix(a, copy = false)*b)::elem_type(A)
end

function *(b::T, a::AlgMatElem) where {T <: RingElem}
  A = parent(a)
  if parent(b) == base_ring(A)
    b = coefficient_ring(A)(b)
  end
  return A(b*matrix(a, copy = false))::elem_type(A)
end

function *(a::AlgMatElem{S, T, U}, b::U) where { S, T, U <: MatElem }
  A = parent(a)
  return A(matrix(a, copy = false)*b)
end

function *(b::U, a::AlgMatElem{S, T, U}) where { S, T, U <: MatElem }
  A = parent(a)
  return A(b*matrix(a, copy = false))
end

################################################################################
#
#  Inverse
#
################################################################################

@doc Markdown.doc"""
    inv(a::AlgMatElem) -> AlgMatElem

Returns $a^{-1}$.
"""
function inv(a::AlgMatElem)
  return parent(a)(inv(matrix(a, copy = false)))
end

################################################################################
#
#  Exponentiation
#
################################################################################

@doc Markdown.doc"""
    ^(a::AlgMatElem, b::Int) -> AlgMatElem

Returns $a^b$.
"""
function ^(a::AlgMatElem, b::Int)
  return parent(a)(matrix(a, copy = false)^b)
end

################################################################################
#
#  Parent object overloading
#
################################################################################

function (A::AlgMat)()
  n = degree(A)
  return A(zero_matrix(coefficient_ring(A), n, n))
end

function (A::AlgMat{T, S})(M::S) where { T, S }
  @assert base_ring(M) === coefficient_ring(A)
  return AlgMatElem{T, typeof(A), S}(A, deepcopy(M))
end

#function (A::AlgMat{T, S})(a::T) where { T, S }
#  return a*one(A)
#end
#
#function (A::AlgMat)(x::fmpz)
#  return x * one(A)
#end
#
#function (A::AlgMat)(x::Integer)
#  return x * one(A)
#end

function (A::AlgMat{T, S})(v::Vector{T}) where { T, S }
  @assert length(v) == dim(A)
  R = coefficient_ring(A)
  M = zero_matrix(R, degree(A), degree(A))
  B = basis(A)
  for i = 1:dim(A)
    #M = add!(M, M, matrix(basis(A)[i], copy = false)*v[i])
    M += matrix(B[i], copy = false)*R(v[i])
  end
  a = A(M)
  a.coeffs = copy(v)
  a.has_coeffs = true
  return a
end

################################################################################
#
#  Equality
#
################################################################################

@doc Markdown.doc"""
    ==(a::AlgMatElem, b::AlgMatElem) -> Bool

Returns `true` if $a$ and $b$ are equal and `false` otherwise.
"""
function ==(a::AlgMatElem{T, S, V}, b::AlgMatElem{T, S, V}) where {T, S, V}
  parent(a) != parent(b) && return false
  return matrix(a, copy = false) == matrix(b, copy = false)
end

################################################################################
#
#  isone/iszero
#
################################################################################

isone(a::AlgMatElem) = isone(matrix(a, copy = false))

iszero(a::AlgMatElem) = iszero(matrix(a, copy = false))

################################################################################
#
#  elem_from_mat_row
#
################################################################################

function elem_from_mat_row(A::AlgMat{T, S}, M::MatElem{T}, i::Int) where { T, S }
  v = Vector{T}(undef, dim(A))
  for c = 1:ncols(M)
    v[c] = deepcopy(M[i, c])
  end
  return A(v)
end

function elem_from_mat_row(A::AlgMat, M::fmpz_mat, i::Int, d::fmpz = fmpz(1))
  v = Vector{fmpq}(undef, dim(A))
  for j in 1:ncols(M)
    v[j] = fmpq(M[i, j], d)
  end
  return A(v)
end

################################################################################
#
#  Hashing
#
################################################################################

function Base.hash(a::AlgMatElem, h::UInt)
  return Base.hash(matrix(a, copy = false), h)
end

################################################################################
#
#  getindex / setindex
#
################################################################################

function getindex(a::AlgMatElem, r::Int, c::Int)
  return matrix(a, copy = false)[r, c]
end

function setindex!(a::AlgMatElem, d::T, r::Int, c::Int) where T <: Union{ RingElem, Int }
  a.matrix[r, c] = coefficient_ring(parent(a))(d)
  a.has_coeffs = false
  return nothing
end


text:
"""
    Storage{T,N}

    contains maximal-representation trajectories

    x: position
    q: orientation (Quaternion)
    v: linear velocity (midpoint)
    ω: angular velocity (midpoint)
    px: linear momentum
    pq: angular momentum
    vl: linear velocity
    ωl: angular velocity
"""
struct Storage{T,N}
    x::Vector{Vector{SVector{3,T}}}
    q::Vector{Vector{Quaternion{T}}}
    v::Vector{Vector{SVector{3,T}}}
    ω::Vector{Vector{SVector{3,T}}}
    px::Vector{Vector{SVector{3,T}}}
    pq::Vector{Vector{SVector{3,T}}}
    vl::Vector{Vector{SVector{3,T}}}
    ωl::Vector{Vector{SVector{3,T}}}

    function Storage{T}(steps, nbodies) where T
        x = [[szeros(T, 3) for i = steps] for j = 1:nbodies]
        q = [[one(Quaternion{T}) for i = steps] for j = 1:nbodies]
        v = [[szeros(T, 3) for i = steps] for j = 1:nbodies]
        ω = [[szeros(T, 3) for i = steps] for j = 1:nbodies]
        px = [[szeros(T, 3) for i = steps] for j = 1:nbodies]
        pq = [[szeros(T, 3) for i = steps] for j = 1:nbodies]
        vl = [[szeros(T, 3) for i = steps] for j = 1:nbodies]
        ωl = [[szeros(T, 3) for i = steps] for j = 1:nbodies]
        new{T,length(steps)}(x, q, v, ω, px, pq, vl, ωl)
    end

    Storage(steps, nbodies) = Storage{Float64}(steps, nbodies)
    Storage{T}(nsteps::Integer, nbodies) where T = Storage{T}(Base.OneTo(nsteps), nbodies)
    Storage(nsteps::Integer, nbodies) = Storage{Float64}(nsteps, nbodies)

    Storage{T}() where T = Storage{T}(Base.OneTo(0),0)
end

Base.length(storage::Storage{T,N}) where {T,N} = N

function save_to_storage!(mechanism::Mechanism, storage::Storage, i::Int)
    for (ind, body) in enumerate(mechanism.bodies)
        state = body.state
        storage.x[ind][i] = state.x2 # x2
        storage.q[ind][i] = state.q2 # q2
        storage.v[ind][i] = state.v15 # v1.5
        storage.ω[ind][i] = state.ϕ15 # ω1.5
        q2 = state.q2
        p2 = momentum(mechanism, body) # p1 in world frame
        px2 = p2[SVector{3,Int}(1,2,3)] # px1 in world frame
        pq2 = p2[SVector{3,Int}(4,5,6)] # pq1 in world frame
        v2 = px2 ./ body.mass # in world frame
        ω2 = body.inertia \ (rotation_matrix(inv(q2)) * pq2) # in body frame, we rotate using the current quaternion q2 = state.q2
        storage.px[ind][i] = px2 # px2
        storage.pq[ind][i] = pq2 # pq2
        storage.vl[ind][i] = v2 # v2
        storage.ωl[ind][i] = ω2 # ω2
    end
    return
end

function generate_storage(mechanism::Mechanism, z)
    N = length(z)
    M = length(mechanism.bodies)
	storage = Storage{Float64}(N, M)

    for t = 1:N
        off = 0
        for i = 1:M
            storage.x[i][t] = z[t][off .+ (1:3)]
            storage.v[i][t] = z[t][off .+ (4:6)]
            storage.q[i][t] = Quaternion(z[t][off .+ (7:10)]...)
            storage.ω[i][t] = z[t][off .+ (11:13)]
            off += 13
        end
    end

    return storage
end

function get_maximal_state(storage::Storage{T,N}) where {T,N}
	z = [get_maximal_state(storage, i) for i=1:N]
	return z
end

function get_maximal_state(storage::Storage{T,N}, i::Int) where {T,N}
	Nb = length(storage.x)
	z = zeros(13 * Nb)
	for j = 1:Nb
		x2 = storage.x[j][i]
		q2 = storage.q[j][i]
		v15 = storage.v[j][i]
		ϕ15 = storage.ω[j][i]
		z[13 * (j-1) .+ (1:13)] = [x2; v15; vector(q2); ϕ15]
	end
	return z
end


text:
export AR, Autoregressive, ARsafe, ARunsafe, ARMeta

import StatsFuns: log2π

struct AR end

const Autoregressive = AR

struct ARsafe end
struct ARunsafe end

struct ARMeta{F <: VariateForm, S}
    order :: Int
    stype :: S
end

function ARMeta(::Type{ Univariate }, order, stype::S) where S
    order === 1 || @warn "ARMeta{Univariate} has been created with order equals to $(order). Order has been forced to be equal to 1."
    return ARMeta{Univariate, S}(1, stype)
end

function ARMeta(::Type{Multivariate}, order, stype::S) where S
    return ARMeta{Multivariate, S}(order, stype)
end

getvform(meta::ARMeta{F}) where { F } = F
getorder(meta::ARMeta)                = meta.order
getstype(meta::ARMeta)                = meta.stype

is_multivariate(meta::ARMeta) = getvform(meta) === Multivariate
is_univariate(meta::ARMeta)   = getvform(meta) === Univariate

is_safe(meta::ARMeta)   = getstype(meta) === ARsafe()
is_unsafe(meta::ARMeta) = getstype(meta) === ARunsafe()

@node AR Stochastic [ y, x, θ, γ ]

default_meta(::Type{ AR }) = error("Autoregressive node requires meta flag explicitly specified")

@average_energy AR (q_y_x::MultivariateNormalDistributionsFamily, q_θ::NormalDistributionsFamily, q_γ::GammaShapeRate, meta::ARMeta) = begin
    mθ, Vθ   = mean(q_θ), cov(q_θ)
    myx, Vyx = mean(q_y_x), cov(q_y_x)
    mγ       = mean(q_γ)

    order = getorder(meta)

    mx, Vx   = ar_slice(getvform(meta), myx, order+1:2order), ar_slice(getvform(meta), Vyx, order+1:2order, order+1:2order)
    my1, Vy1 = first(myx), first(Vyx)
    Vy1x     = ar_slice(getvform(meta), Vyx, 1, order+1:2order)

    AE = -0.5*(logmean(q_γ)) + 0.5log2π + 0.5*mγ*(Vy1+my1^2 - 2*mθ'*(Vy1x + mx*my1) + tr(Vθ*Vx) + mx'*Vθ*mx + mθ'*(Vx + mx*mx')*mθ)

    # correction
    if is_multivariate(meta)
        AE += entropy(q_y_x)
        idc = [1, order+1:2order...]
        myx_n = myx[idc]
        Vyx_n = Vyx[idc, idc]
        q_y_x = MvNormalMeanCovariance(myx_n, Vyx_n)
        AE -= entropy(q_y_x)
    end

    return AE
end

# Helpers for AR rules

"""
    ar_slice(::T, array, ranges...)

Returns `array[ranges...]` in case if T is Multivariate, and `first(array[ranges...])` in case if T is Univariate
"""
function ar_slice end

function ar_slice(::Type{Multivariate}, array, ranges...)
    return array[ranges...]
end

function ar_slice(::Type{Univariate}, array, ranges...)
    return first(array[ranges...])
end

"""
ar_uvector(::T, order)

Returns `[ 1.0, 0.0 ... 0.0 ]` with length equal to order in case if T is Multivariate, and `1.0` in case if T is Univariate
"""
function ar_unit end

function ar_unit(::Type{Multivariate}, order)
    c    = zeros(order)
    c[1] = 1.0
    return c
end

function ar_unit(::Type{Univariate}, order)
    return 1.0
end

function ar_precision(::Type{Multivariate}, order, γ)
    mw               = zeros(order, order)
    mw[diagind(mw)] .= huge
    mw[1, 1]         = γ
    return mw
end

function ar_precision(::Type{Univariate}, order, γ)
    return γ
end

function ar_transition(::Type{Multivariate}, order, γ)
    V = zeros(order, order)
    V[1] = 1/γ
    return V
end

function ar_transition(::Type{Univariate}, order, γ)
    return inv(γ)
end


text:
using Clang, Clang.LibClang
using CEnum
using DataStructures, OrderedCollections

include("clanghelpers.jl")

const VX_INCLUDE = joinpath(
    homedir(),
    "git",
    "OpenVX.jl",
    "deps",
    "include",
    "VX"
)

const VX_HEADERS = [joinpath(VX_INCLUDE, header) for header in readdir(VX_INCLUDE) if endswith(header, ".h")]

ctx = DefaultContext()

parse_headers!(ctx, VX_HEADERS,
              args=["-I", VX_INCLUDE, "-fparse-all-comments"],
              includes=vcat(VX_INCLUDE, CLANG_INCLUDE),
              flags = CXTranslationUnit_DetailedPreprocessingRecord)

ctx.libname = "libopenvx"
ctx.options["is_function_strictly_typed"] = true
ctx.options["is_struct_mutable"] = false

blacklist = ["unix", "linux"]
blackprefix = ["Brig", "_"]

functions = Vector{CLCursor}()
remaining = Vector{CLCursor}()
typealiases = OrderedDict()

for trans_unit in ctx.trans_units

    root_cursor = getcursor(trans_unit)
    push!(ctx.cursor_stack, root_cursor)
    header = spelling(root_cursor)
    ctx.children = children(root_cursor)

    for (i, child) in enumerate(ctx.children)
        # Cursor properties
        child_name = name(child)
        child_kind = kind(child)

        # Skip compiler constants/garbage and redundants
        child_kind == CXCursor_MacroInstantiation   && continue
        child_kind == CXCursor_LastPreprocessing    && continue
        child_name ∈ blacklist                      && continue
        if (any(startswith.(child_name, blackprefix)) && child_kind == CXCursor_MacroDefinition)
            continue
        end
        child_name ∈ keys(ctx.common_buffer)        && continue
        child_name ∈ name.(functions)               && continue
        child_kind == CXCursor_FunctionDecl     && (wrap_vx_fun!(ctx, child); push!(functions, child); continue)
        wrap!(ctx, child)
    end
end

api_file = joinpath(@__DIR__, "libopenvx_api.jl")
api_stream = open(api_file, "w")
print_buffer(api_stream, ctx.api_buffer)
empty!(ctx.api_buffer)
close(api_stream)

common_file = joinpath(@__DIR__, "libopenvx_common.jl")
open(common_file, "w") do f
    print_buffer(f, dump_to_buffer(ctx.common_buffer))
end


text:
# Bradford, E., Imsland, L., Zhang, D., & del Rio Chanona, E. A. (2020).
# Stochastic data-driven model predictive control using Gaussian processes.
# Computers & Chemical Engineering, 139, 106844.

# objective: maximize C_qc

# state constraints
# C_N(t) - 150 ≤ 0              t = T
# C_N(t) − 800 ≤ 0              ∀t
# 0.011 C_X(t) - C_qc(t) ≤ 3f-2 ∀t

function bioreactor(; store_results=false::Bool)
    datadir = nothing
    if store_results
        datadir = generate_data_subdir(@__FILE__)
    end

    # initial conditions and timepoints
    t0 = 0.0f0
    tf = 240.0f0
    Δt = 10.0f0
    tspan = (t0, tf)
    C_X₀, C_N₀, C_qc₀ = 1.0f0, 150.0f0, 0.0f0
    u0 = [C_X₀, C_N₀, C_qc₀]
    control_ranges = [(120.0f0, 400.0f0), (0.0f0, 40.0f0)]

    system_params = (
        u_m=0.0572f0,
        u_d=0.0f0,
        K_N=393.1f0,
        Y_NX=504.5f0,
        k_m=0.00016f0,
        k_d=0.281f0,
        k_s=178.9f0,
        k_i=447.1f0,
        k_sq=23.51f0,
        k_iq=800.0f0,
        K_Np=16.89f0,
    )

    function system!(du, u, p, t, controller; input=:state)
        @argcheck input in (:state, :time)

        u_m, u_d, K_N, Y_NX, k_m, k_d, k_s, k_i, k_sq, k_iq, K_Np = values(system_params)

        # neural network outputs controls based on state
        C_X, C_N, C_qc = u
        if input == :state
            I, F_N = controller(u, p)
        elseif input == :time
            I, F_N = controller(t, p)
        end

        # auxiliary variables
        I_ksi = I / (I + k_s + I^2.0f0 / k_i)
        CN_KN = C_N / (C_N + K_N)

        I_kiq = I / (I + k_sq + I^2.0f0 / k_iq)
        Cqc_KNp = C_qc / (C_N + K_Np)

        # dynamics of the controlled system
        dC_X = u_m * I_ksi * C_X * CN_KN - u_d * C_X
        dC_N = -Y_NX * u_m * I_ksi * C_X * CN_KN + F_N
        dC_qc = k_m * I_kiq * C_X - k_d * Cqc_KNp

        # update in-place
        @inbounds begin
            du[1] = dC_X
            du[2] = dC_N
            du[3] = dC_qc
        end
    end

    # set arquitecture of neural network controller
    controller = FastChain(
        (x, p) -> [x[1], x[2] / 10.0f0, x[3] * 10.0f0],  # input scaling
        FastDense(3, 16, tanh_fast; initW=(x, y) -> Float32(5 / 3) * glorot_uniform(x, y)),
        FastDense(16, 16, tanh_fast; initW=(x, y) -> Float32(5 / 3) * glorot_uniform(x, y)),
        FastDense(16, 2; initW=(x, y) -> Float32(5 / 3) * glorot_uniform(x, y)),
        # I ∈ [120, 400] & F ∈ [0, 40] in Bradford 2020
        # (x, p) -> [280f0 * sigmoid(x[1]) + 120f0, 40f0 * sigmoid(x[2])],
        scaled_sigmoids(control_ranges),
    )

    controlODE = ControlODE(controller, system!, u0, tspan; Δt)

    function collocation(
        u0;
        num_supports::Integer=length(controlODE.tsteps),
        nodes_per_element::Integer=4,
        constrain_states::Bool=false,
    )
        optimizer = optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0)
        model = InfiniteModel(optimizer)
        method = OrthogonalCollocation(nodes_per_element)
        @infinite_parameter(
            model, t in [t0, tf], num_supports = num_supports, derivative_method = method
        )

        @variables(
            model,
            begin
                # state variables
                x[1:3], Infinite(t)
                # control variables
                c[1:2], Infinite(t)
            end
        )

        # fixed parameters
        u_m, u_d, K_N, Y_NX, k_m, k_d, k_s, k_i, k_sq, k_iq, K_Np = values(system_params)

        # initial conditions
        @constraint(model, [i = 1:3], x[i](0) == u0[i])

        # control range
        @constraints(
            model,
            begin
                control_ranges[1][1] <= c[1] <= control_ranges[1][2]
                control_ranges[2][1] <= c[2] <= control_ranges[2][2]
            end
        )

        if constrain_states
            @constraints(
                model,
                begin
                    x[2] <= 150
                    x[2] <= 800
                    1.1f-2 * x[1] - x[3] <= 3.0f-2
                end
            )
        end

        # dynamic equations
        @constraints(
            model,
            begin
                ∂(x[1], t) ==
                u_m *
                (c[1] / (c[1] + k_s + c[1]^2.0f0 / k_i)) *
                x[1] *
                (x[2] / (x[2] + K_N)) - u_d * x[1]
                ∂(x[2], t) ==
                -Y_NX *
                u_m *
                (c[1] / (c[1] + k_s + c[1]^2.0f0 / k_i)) *
                x[1] *
                (x[2] / (x[2] + K_N)) + c[2]
                ∂(x[3], t) ==
                k_m * (c[1] / (c[1] + k_sq + c[1]^2.0f0 / k_iq)) * x[1] -
                k_d * (x[3] / (x[2] + K_Np))
            end
        )

        @objective(model, Max, x[3](tf))

        optimize!(model)

        model |> optimizer_model |> solution_summary
        states = hcat(value.(x)...) |> permutedims
        controls = hcat(value.(c)...) |> permutedims
        return model, supports(t), states, controls
    end

    control_profile, infopt_model, times_collocation, states_collocation, controls_collocation = collocation_preconditioner(
        u0, collocation; plot=true
    )

    θ = preconditioner(
        controlODE,
        control_profile;
        progressbar=true,
        plot_progress=false,
        # control_range_scaling=[range[end] - range[1] for range in control_ranges],
    )

    plot_simulation(controlODE, θ; only=:states)
    plot_simulation(controlODE, θ; only=:controls)
    store_simulation("precondition", controlODE, θ; datadir)

    function state_penalty_functional(solution_array, time_intervals; δ=1.0f1)
        @argcheck size(solution_array, 2) == length(time_intervals) + 1

        ratio_X_N = 3.0f-2 / 800.0f0
        C_N_over = map(y -> relaxed_log_barrier(800.0f0 - y; δ), solution_array[2, 1:end])
        C_X_over = map(
            (x, z) -> relaxed_log_barrier(3.0f-2 - (1.1f-2 * x - z); δ=δ * ratio_X_N),
            solution_array[1, 1:end],
            solution_array[3, 1:end],
        )
        C_N_over_last = relaxed_log_barrier(150.0f0 - solution_array[2, end]; δ)

        return sum((C_N_over .+ C_X_over)[2:end] .* time_intervals) + C_N_over_last
    end

    # state constraints on control change
    # C_N(t) - 150 ≤ 0              t = T
    # C_N(t) − 800 ≤ 0              ∀t
    # 0.011 C_X(t) - C_qc(t) ≤ 3f-2 ∀t
    function losses(
        controlODE, params; δ=1.0f1, α=1.0f-5, μ=(3.125f-8, 3.125f-6), ρ=1.0f-1, tsteps=(), kwargs...
    )

        # integrate ODE system
        sol_raw = solve(controlODE, params; kwargs...)
        sol = Array(sol_raw)

        # approximate integral penalty
        # state_penalty = Δt * (sum(C_N_over) + sum(C_X_over)) + C_N_over_last  # for fixed timesteps
        Zygote.ignore() do
            global time_intervals = [
                sol_raw.t[i + 1] - sol_raw.t[i] for i in eachindex(sol_raw.t[1:(end - 1)])
            ]
        end

        state_penalty = α * state_penalty_functional(sol, time_intervals; δ)

        # penalty on change of controls
        control_penalty = 0.0f0
        for i in 1:(size(sol, 2) - 1)
            prev = controller(sol[:, i], params)
            post = controller(sol[:, i + 1], params)
            for j in 1:length(prev)
                control_penalty += μ[j] * (prev[j] - post[j])^2
            end
        end

        regularization = ρ * mean(abs2, θ)  # sum(abs2, θ)

        objective = -sol[3, end]  # maximize C_qc

        return objective, state_penalty, control_penalty, regularization
    end

    function plots_callback(controlODE, θ)
        plot_simulation(
            controlODE, θ; only=:states, vars=[2], yrefs=[800, 150]
        )
        plot_simulation(
            controlODE,
            θ;
            only=:states,
            fun=(x, y, z) -> 1.1f-2x - z,
            yrefs=[3.0f-2],
        )
        return plot_simulation(controlODE, θ; only=:controls)
    end

    # α: penalty coefficient
    # δ: barrier relaxation coefficient
    α0, δ0 = 1.0f-5, 100.0f0
    barrier_iterations = 0:20
    αs = [α0 for _ in barrier_iterations]
    δs = [δ0 * 0.8f0^i for i in barrier_iterations]
    θ = constrained_training(
        controlODE,
        losses;
        αs,
        δs,
        tsteps,
        starting_params=θ,
        show_progressbar=true,
        # plots_callback,
        datadir,
    )

    final_values = NamedTuple{(
        :objective, :state_penalty, :control_penalty, :regularization
    )}(
        loss(θ; δ=δs[end], α=αs[end], tsteps)
    )

    @info "Final states"
    # plot_simulation(controlODE, θ; only=:states, vars=[1], show=final_values)
    plot_simulation(
        controlODE,
        θ;
        only=:states,
        vars=[2],
        yrefs=[800, 150],
    )
    plot_simulation(controODE, θ; only=:states, vars=[3])
    plot_simulation(
        controlODE,
        θ;
        only=:states,
        fun=(x, y, z) -> 1.1f-2x - z,
        yrefs=[3.0f-2],
    )

    @info "Final controls"
    plot_simulation(
        controODE, θ; only=:controls, vars=[1]
    )
    plot_simulation(
        controODE, θ; only=:controls, vars=[2]
    )

    @info "Final loss" loss(θ; δ=δs[end], α=αs[end], tsteps)

    # initial conditions and timepoints
    # t0 = 0f0
    # tf = 240f0
    # Δt = 10f0
    # C_X₀, C_N₀, C_qc₀ = 1f0, 150f0, 0f0
    # u0 = [C_X₀, C_N₀, C_qc₀]
    # tspan = (t0, tf)
    # tsteps = t0:Δt:tf
    # control_ranges = [(120f0, 400f0), (0f0, 40f0)]
    perturbation_specs = [
        (variable=1, type=:centered, scale=10.0f0, samples=10, percentage=2.0f-2)
        (variable=2, type=:centered, scale=800.0f0, samples=10, percentage=2.0f-2)
        (variable=3, type=:positive, scale=5.0f-1, samples=10, percentage=2.0f-2)
    ]
    return plot_initial_perturbations(controlODE, θ, perturbation_specs)
end  # script wrapper


text:
export Measurement

struct Measurement
    name::String
    fields::Dict{String,Union{String,Int,Float64,Bool}}
    tags::Dict{String,String}
    timestamp::Float64

    function Measurement(name, fields, tags, timestamp = time())
        return new(
            String(name),
            Dict(String(k) => v for (k, v) in fields),
            Dict(String(k) => String(v) for (k, v) in tags),
            Float64(timestamp),
        )
    end
end

function payload(m::Measurement)
    buff = IOBuffer()
    payload!(buff, m)
    return String(take!(buff))
end

function check_allowed(x)
    contains(x, '\n') && error("'\\n' not allowed in " * repr(x))
    endswith(x, '\\') && error("'\\\\' not allowed at the end of " * repr(x))
    return true
end

function escape_measurement(x)
    check_allowed(x)
    x = replace(x, r"(\\[ ,=\"])" => s"\\\1")
    x = replace(x, r"([ ,=\"])" => s"\\\1")
    return x
end

escape_tag(x)           = (check_allowed(x) && replace(x, r"([ ,=])" => s"\\\1"))
escape_tag_key(x)       = (check_allowed(x) && replace(x, r"([ ,=])" => s"\\\1"))
escape_field_key(x)     = (check_allowed(x) && replace(x, r"([ ,=\"])" => s"\\\1"))
escape_field(x::String) = '"' * replace(x, r"([\"\\])" => s"\\\1") * '"'
escape_field(x::Number) = string(x)

function payload!(io::IO, m::Measurement)
    write(io, escape_measurement(m.name))
    if !isempty(m.tags)
        join(io, [",$(escape_tag_key(k))=$(escape_tag(v))" for (k, v) in m.tags])
    end
    write(io, " ")
    join(io, ["$(escape_field_key(k))=$(escape_field(v))" for (k, v) in m.fields], ",")
    write(io, " ")
    write(io, string(round(Int64, m.timestamp*1e9)))
end


text:
using KernelAbstractions
using KernelAbstractions: CPUEvent

function __wait(; dependencies=nothing)
  wait(CPU(), MultiEvent(dependencies), yield)
end

event = CPUEvent(@async(__nonexistent!()))
event = CPUEvent(@async(__wait(dependencies=event)))
wait(event)


text:
# This file was generated, do not modify it. # hide
describe(Xc, :mean, :std)


text:
const WAN_EXECS= ["wannier90.x"]
include(joinpath(depsdir, "wannier90flags.jl"))
const WAN_FLAGS = _WAN_FLAGS()
flagtype(::Type{Wannier90}, flag) = haskey(WAN_FLAGS, flag) ? WAN_FLAGS[flag] : Nothing
flagtype(::DFInput{Wannier90}, flag) = flagtype(Wannier90, flag)


text:
# Autogenerated wrapper script for Hwloc_jll for x86_64-apple-darwin
export libhwloc

JLLWrappers.@generate_wrapper_header("Hwloc")
JLLWrappers.@declare_library_product(libhwloc, "@rpath/libhwloc.15.dylib")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libhwloc,
        "lib/libhwloc.15.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()