text:
#=
Handles models defined as named tuples.

* Developer: Hilding Elmqvist, Mogram AB
* First version: January 2021
* License: MIT (expat)

=#
export mergeModels, recursiveMerge, Redeclare, showModel, @showModel, drawModel, Model, Map, Par, Var, setLogMerge,
    constant, parameter, input, output, potential, flow, interval, @info_str, Boolean, Integ, @define

using Base.Meta: isexpr
using OrderedCollections: OrderedDict
using Unitful
using ModiaBase.Symbolic: removeBlock, prepend

function stringifyDefinition(v)
    if typeof(v) in [Symbol, Expr]
        v = removeBlock(v)
        v = ":(" * string(v) * ")"
    end
    return v
end

"Quote an expression; don't quote if the meaning in the AST is the same anyway."
selquote(x) = x
selquote(sym::Symbol) = Meta.quot(sym)
function selquote(ex::Expr)
    if isexpr(ex, :kw) || isexpr(ex, :call)
        Expr(ex.head, [ex.args[1], [selquote(e) for e in ex.args[2:end]]...]...)
    else
        Expr(ex.head, [selquote(e) for e in ex.args]...)
    end
end

macro define(modelDef)
    modelName = modelDef.args[1]
    println(modelName)
    model = modelDef.args[2]
    model = selquote(model)
    dump(model)
    res = :($modelName = Meta.quot($model))
    return esc(res)
end

function showModel(m, level=0)
    level += 1
#    print("  "^(level-1))
    if typeof(m) <: NamedTuple && haskey(m, :class)
        print(m.class)
    end
    println("(")
    for (k, v) in zip(keys(m), m)
        if typeof(v) <: NamedTuple
            print("  "^level, k, " = ")
            showModel(v, level)
        elseif k != :class
            println("  "^level, k, " = ", stringifyDefinition(v), ",")
        end
    end
    println("  "^(level-1), "),")
end

macro showModel(model)
    modelName = string(model)
    mod = :( print($modelName, " = "); showModel($model); println() )
    return esc(mod)
end

global logMerge = false

function setLogMerge(val)
    global logMerge
    logMerge = val
end

id = 1

function drawModel(name, m, level=0)
    this = Dict()
    children = []
    ports = []
    edges = []
    level += 1
    global id
    if typeof(m) <: NamedTuple && haskey(m, :class)
#        print(m.class)
    end
    for (k, v) in zip(keys(m), m)
        if typeof(v) <: NamedTuple && level <= 1
            child = drawModel(k, v, level)
            push!(children, child)
        elseif k in [:class, :equations, :input, :output, :potential, :flow, :info]
        elseif k == :connect
            for conn in v.args
                conn = conn.args
                source = conn[1]
                target = conn[2:end]
                id += 1
                edge = (; id, sources = [string(source)], targets=[string(target[1])])
                push!(edges, edge)
            end
        elseif level <= 2
            port = (; id=string(name) * "." * string(k), height=10, width=10)
            push!(ports, port)
        end
    end

    if length(children) > 0 && length(ports) > 0
        this = (;id=name, height=100, width=100, labels=[(;text=name)], ports, children, edges)
    elseif length(children) > 0
        this = (;id=name, height=100, width=100, labels=[(;text=name)], children, edges)
    elseif length(ports) > 0
        this = (;id=name, height=100, width=100, labels=[(;text=name)], ports)
    else
        this = (;id=name, height=100, width=100, labels=[(;text=name)])
    end
    return this
end

function mergeModels(m1::NamedTuple, m2::NamedTuple, env=Symbol())
#    println("mergedModels")
    mergedModels = OrderedDict{Symbol,Any}(pairs(m1))
    for (k,v) in collect(pairs(m2))
        if typeof(v) <: NamedTuple
            if k in keys(mergedModels) && ! (:_redeclare in keys(v))
                if logMerge; print("In $k: ") end
                m = mergeModels(mergedModels[k], v, k)
                mergedModels[k] = m
            elseif :_redeclare in keys(v)
                if logMerge; println("Redeclaring: $k = $v") end
                mergedModels[k] = v
            elseif nothing in values(v) # TODO: Refine

            else
                if !(:_redeclare in keys(mergedModels))
                    if logMerge; print("Adding: $k = "); showModel(v, 2) end
                end
                mergedModels[k] = v
            end
        elseif v === nothing
            if logMerge; println("Deleting: $k") end
            delete!(mergedModels, k)
        elseif k in keys(mergedModels) && k == :equations
            equa = copy(mergedModels[k])
            push!(equa.args, v.args...)
            mergedModels[k] = equa
            if logMerge
                println("Adding equations: ", v)
            end
        else
            if logMerge
                if k in keys(mergedModels)
                    if mergedModels[k] != v
                        println("Changing: $k = $(stringifyDefinition(mergedModels[k])) to $k = $(stringifyDefinition(v))")
                    end
                elseif !(:_redeclare in keys(mergedModels))
                    println("Adding: $k = $(stringifyDefinition(v))")
                end
            end
            mergedModels[k] = v
        end
    end
#    delete!(mergedModels, :class)
    return (; mergedModels...) # Transform OrderedDict to named tuple
end

Model(; kwargs...) = (; class = :Model, kwargs...)

#Map(; kwargs...) = (; class = :Map, kwargs...)
Map(; kwargs...) = (; kwargs...)

Par(; kwargs...) = Map(; class = :Par, kwargs...)

#Var(;kwargs...) = (; class = :Var, kwargs...)
function Var(values...; kwargs...)
    res = (; class = :Var, kwargs...)
    for v in values
        res = res | v
    end
    res
end

Integ(;kwargs...) = (; class = :Var, kwargs...)

Boolean(;kwargs...) = (; class = :Var, kwargs...)

Redeclare = ( _redeclare = true, )

constant = Var(constant = true)
parameter = Var(parameter = true)
input = Var(input = true)
output = Var(output = true)
potential = Var(potential = true)
flow = Var(flow = true)

interval(min, max) = Var(min=min, max=max)

macro info_str(text)
    Var(info=text)
end

Base.:|(m::NamedTuple, n::NamedTuple) =  mergeModels(m, n) # TODO: Chane to updated recursiveMerge
Base.:|(m, n) = begin if !(typeof(n) <: NamedTuple); recursiveMerge(m, (; value=n)) else recursiveMerge(n, (value=m,)) end end

recursiveMerge(x, ::Nothing) = x
recursiveMerge(x, y) = y
#recursiveMerge(x::Expr, y::Expr) = begin dump(x); dump(y); Expr(x.head, x.args..., y.args...) end
recursiveMerge(x::Expr, y::Tuple) = begin x = copy(x); xargs = x.args; xargs[y[2]] = y[3]; Expr(x.head, xargs...) end

#=
Base.:|(x::Symbol, y::NamedTuple) = begin
    :(eval($x) | $y)
end
=#

function recursiveMerge(nt1::NamedTuple, nt2::NamedTuple)
#    println("recursiveMerge")
#    @show nt1 nt2
    all_keys = union(keys(nt1), keys(nt2))
#    all_keys = setdiff(all_keys, [:class])
    gen = Base.Generator(all_keys) do key
        v1 = get(nt1, key, nothing)
        v2 = get(nt2, key, nothing)
        key => recursiveMerge(v1, v2)
    end
#    @show gen
    return (; gen...)
end

function unpackPath(path, sequence)
    if typeof(path) == Symbol
        push!(sequence, path)
    elseif isexpr(path, :.)
        unpackPath(path.args[1], sequence)
        push!(sequence, path.args[2].value)
    elseif isexpr(path, :ref)
        unpackPath(path.args[1], sequence)
        push!(sequence, path.args[2:end]...)
    end
end

function collectConnector(model)
    potentials = []
    flows = []
    for (k,v) in zip(keys(model), model)
        if typeof(v) <: NamedTuple && :class in keys(v) && v.class == :Var ||
            typeof(v) <: NamedTuple && :variable in keys(v) && v.variable
            if :potential in keys(v) && v.potential
                push!(potentials, k)
            elseif :flow in keys(v) && v.flow
                push!(flows, k)
            else
                push!(potentials, k)
            end
        end
    end
    return potentials, flows
end

function mergeConnections!(connections)
    for i in 1:length(connections)
        for j in 1:i-1
            con1 = connections[i]
            con2 = connections[j]
            if length(con1.args) > 0 && length(con2.args) > 0 && length(intersect(Set(con1.args), Set(con2.args))) > 0
                connections[i] = Expr(:tuple, union(Set(con1.args), Set(con2.args))...)
                connections[j] = Expr(:tuple) # Empty tuple
                # if :(OpI.n2) in con1.args || :(OpI.n2) in con2.args # For bug testing
                #     @show i j con1.args con2.args
                #     @show connections[i] connections[j]
                # end
            end
        end
    end
end

function convertConnections!(connections, model, modelName, logging=false)
#    println("\nconvertConnections")
#    showModel(model)
    mergeConnections!(connections)

    connectEquations = []
    alreadyConnected = []
    for i in 1:length(connections)
        c = connections[i]
        if c.head == :tuple
            connected = c.args

            inflow = 0
            outflow = 0
            signalFlow1 = nothing
            connectedOutput = nothing
            potentials1 = nothing
            fullPotentials1 = []
            flows1 = []
            for con in connected
                if con in alreadyConnected
                    error("Already connected: $con, found in connection: $connected")
                end
                push!(alreadyConnected, con)
#                @show connected con

                sequence = []
                unpackPath(con, sequence)
                mod = model
                for s in sequence[1:end-1]
                    if s in keys(mod)
                        mod = mod[s]
                    else
                        error("Invalid path $con: $s not found in $(keys(mod))")
                    end
                end

                if sequence[end] in keys(mod)
                    mod = mod[sequence[end]]
#                    @show mod[:class]
                    if :input in keys(mod) && mod.input || :output in keys(mod) && mod.output || :class in keys(mod) && mod[:class] == :Var
                        signalFlow = con
                        if signalFlow1 !== nothing
                            push!(connectEquations, :($signalFlow1 = $signalFlow))
                        end
                        signalFlow1 = signalFlow

                        if :output in keys(mod) && mod.output
                            if connectedOutput != nothing # TODO: refine logic concerning inner and outer inputs and outputs
#                                error("It is not allowed to connect two outputs: ", connectedOutput, ", ", con)
                            end
                            connectedOutput = con
                        end
                    elseif mod[:class] == :Var
#                        println("\nConnect vars: ", connected)
#                        dump(connected)
                        if length(fullPotentials1) > 0
                            push!(connectEquations, :($(fullPotentials1[1]) = $con))
#                            println(:($(fullPotentials1[1]) = $con))
                        end
                        push!(fullPotentials1, con)
                    else
#                        @show mod typeof(mod)
                        potentials, flows = collectConnector(mod)
                        # Deprecated
                        if :potentials in keys(mod)
                            potentials = vcat(potentials, mod.potentials.args)
                        end
                        if :flows in keys(mod)
                            flows = vcat(flows, mod.flows.args)
                        end
                        if potentials1 != nothing && potentials != potentials1
                            error("Not compatible potential variables: $potentials1 != $potentials, found in connection: $connected")
                        end
                        fullPotentials = []
                        for i in 1:length(potentials)
                            p = potentials[i]
                            potential = append(con, p)
                            push!(fullPotentials, potential)
                            if potentials1 != nothing
                                push!(connectEquations, :($(fullPotentials1[i]) = $potential))
                            end
                        end
                        potentials1 = potentials
                        fullPotentials1 = fullPotentials

                        if length(flows1) > 0 && flows != flows1
                            error("Not compatible flow variables: $flows1 != $flows, found in connection: $connected")
                        end
                        for f in flows
                            flowVar = append(con, f)
                            if length(sequence) == 1
                                if inflow == 0
                                    inflow = flowVar
                                else
                                    inflow = :($inflow + $flowVar)
                                end
                            else
                                if outflow == 0
                                    outflow = flowVar
                                else
                                    outflow = :($outflow + $flowVar)
                                end
                            end
                        end
                        flows1 = flows
                    end
                else
                    # Deprecated
                    signalFlow = con
                    if signalFlow1 !== nothing
                        push!(connectEquations, :($signalFlow1 = $signalFlow))
                    end
                    signalFlow1 = signalFlow
                end
            end
            if inflow != 0 || outflow != 0
                push!(connectEquations, :($inflow = $outflow))
            end
        end
    end
    if length(connectEquations) > 0 && logging
        println("Connect equations in $modelName:")
        for e in connectEquations
            println("  ", e)
        end
    end
    return connectEquations
end

function flattenModelTuple!(model, modelStructure, modelName; unitless = false, log=false)
#    @show model
    connections = []
    extendedModel = merge(model, NamedTuple())
    for (k,v) in zip(keys(model), model)
#        @show k v typeof(v)
        # Deprecated
        if k in [:inputs, :outputs, :potentials, :flows, :class]
            if k in [:inputs, :outputs, :potentials, :flows]
                printstyled("  Deprecated construct in model $modelName: ", color=:yellow)
                println("$k = $v\n    Use: ... = $(string(k)[1:end-1]) ...")
            end
        elseif k == :init
            printstyled("  Deprecated construct in model $modelName: ", color=:yellow)
            println("$k = $v\n    Use: ... = Var(init=...) ...")
            for (x,x0) in zip(keys(v), v)
                if x != :class
                    if unitless && typeof(x0) != Expr
                        x0 = ustrip(x0)
                    end
                    modelStructure.init[x] = x0
                    modelStructure.mappedParameters = (;modelStructure.mappedParameters..., x => x0)
                end
            end
        elseif k == :start
            printstyled("  Deprecated construct in model $modelName: ", color=:yellow)
            println("$k = $v\n    Use: ... = Var(start=...) ...")
            for (s,s0) in zip(keys(v), v)
                if s != :class
                    if unitless
                        s0 = ustrip(s0)
                    end
                    modelStructure.start[s] = s0
                    modelStructure.mappedParameters = (;modelStructure.mappedParameters...,  s => s0)
                end
            end

        elseif typeof(v) in [Int64, Float64] || typeof(v) <: Unitful.Quantity ||
                typeof(v) in [Array{Int64,1}, Array{Int64,2}, Array{Float64,1}, Array{Float64,2}] ||
				typeof(v) <: NamedTuple && :class in keys(v) && v.class == :Par ||
				typeof(v) <: NamedTuple && :parameter in keys(v) && v.parameter
            if unitless && !(typeof(v) <: NamedTuple)
                v = ustrip(v)
            end
            modelStructure.parameters[k] = v
            modelStructure.mappedParameters = (;modelStructure.mappedParameters..., k => v)
        elseif (typeof(v) <: NamedTuple && :class in keys(v) && v.class in [:Par, :Var] ||
            typeof(v) <: NamedTuple && :parameter in keys(v) && v.parameter) &&
            :value in keys(v)
            v = v.value
            if typeof(v) in [Expr, Symbol]
                push!(modelStructure.equations, removeBlock(:($k = $v)))
            else
                modelStructure.parameters[k] = v
                modelStructure.mappedParameters = (;modelStructure.mappedParameters..., k => v)
            end
         elseif typeof(v) <: NamedTuple && :class in keys(v) && v.class in [:Var] ||
            typeof(v) <: NamedTuple && :variable in keys(v) && v.variable
            if :init in keys(v)
                x0 = v.init
                if unitless && typeof(x0) != Expr
                    x0 = ustrip(x0)
                end
                modelStructure.init[k] = x0
                modelStructure.mappedParameters = (;modelStructure.mappedParameters..., k => x0)
            end
            if :start in keys(v)
                s0 = v.start
                if unitless && typeof(s0) != Expr
                    s0 = ustrip(s0)
                end
                modelStructure.start[k] = s0
                modelStructure.mappedParameters = (;modelStructure.mappedParameters..., k => s0)
            end
            if :input in keys(v) && v[:input]
                modelStructure.inputs[k] = v
            end
            if :output in keys(v) && v[:output]
                modelStructure.outputs[k] = v
            end
        elseif typeof(v) <: NamedTuple # || typeof(v) == Symbol # instantiate
                if typeof(v) == Symbol
                    v = eval(eval(v))
                end
                subModelStructure = ModelStructure()
                flattenModelTuple!(v, subModelStructure, k; unitless, log)
#=
                println("subModelStructure")
                @show subModelStructure
                printModelStructure(subModelStructure, label=k)
=#
                mergeModelStructures(modelStructure, subModelStructure, k)
        elseif typeof(v) <:Array && length(v) > 0 && typeof(v[1]) <: NamedTuple # array of instances
            i = 0
            for a in v
                i += 1
                subModelStructure = ModelStructure()
                flattenModelTuple!(a, subModelStructure, k; unitless, log)
                mergeModelStructures(modelStructure, subModelStructure, Symbol(string(k)*"_"*string(i)) )
            end
        elseif isexpr(v, :vect) || isexpr(v, :vcat) || isexpr(v, :hcat)
            arrayEquation = false
            for e in v.args
                if isexpr(e, :(=))
                    if unitless
                        e = removeUnits(e)
                    end
                    push!(modelStructure.equations, removeBlock(e))
                elseif isexpr(e, :tuple)
                    push!(connections, e)
                elseif isexpr(e, :call) && e.args[1] == :connect
                    con = :( ( $(e.args[2]), $(e.args[3]) ) )
                    push!(connections, con)
                else
                    arrayEquation = true
                end
            end
            if arrayEquation
                push!(modelStructure.equations, removeBlock(:($k = $(prepend(v, :up)))))
#                 push!(modelStructure.equations, removeBlock(:($k = $v)))
            end
        elseif isexpr(v, :(=)) # Single equation
            if unitless
                v = removeUnits(v)
            end
            push!(modelStructure.equations, removeBlock(v))
        elseif v !== nothing # Binding expression
#            println("Binding expression")
#            @show k v typeof(v)

            if unitless
                v = removeUnits(v)
            end
            push!(modelStructure.equations, :($k = $(prepend(v, :up))))
#            push!(modelStructure.equations, :($k = $v))  # To fix Modelica_Electrical_Analog_Interfaces_ConditionalHeatPort
#            @show modelStructure.equations
        end
    end
#    printModelStructure(modelStructure, "flattened")
    connectEquations = convertConnections!(connections, extendedModel, modelName, log)
    push!(modelStructure.equations, connectEquations...)
end


text:
# calculate the volume of a sphere with radius r
# Arguments:  r = radius of the sphere
# Returns:    Volume

function sphereVolume(r)
  return 4/3 * pi * r^3
end

ext: jl


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

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 iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ d5559dba-9fe4-11ec-3744-ebd1408e7dc4
using LegibleLambdas, AbstractTrees, PlutoUI, HypertextLiteral, PlutoTest, LinearAlgebra, Plots

# ╔═╡ 20571e1a-9687-4d42-98b3-b3bc0b207b3b
md"## Let's write our own reverse-mode AD!"

# ╔═╡ 5e02b2c0-10cf-4746-840d-7017648f89f8
md"""
We will use Julia's dispatch system for simplicity. This means we create a type `Tracked` for keeping track of our input variables and everything we'll need to calculate the gradient later.
"""

# ╔═╡ 99b6ab91-a022-449c-988c-0e5c5719c910
begin
	struct Tracked{T} <: Number
		# The numerical result when doing the forward pass
		val::T
		name::Symbol
		# The pullback map for the reverse pass
		df
		# All the other variables this variable directly depends on
		deps::Vector{Tracked}
	end
	Tracked{T}(x, name=gensym()) where {T} = Tracked{T}(x, name, nothing, Tracked[])
	Base.convert(T::Type{Tracked{S}}, x::Tracked) where {S} = T(convert(S, x.val), x.name, x.df, x.deps)
	# This tells Julia to convert any number added to a `Tracked` to a `Tracked` first
	Base.promote_rule(::Type{Tracked{S}}, ::Type{T}) where {S<:Number, T<:Number} = Tracked{promote_type(S, T)}
end

# ╔═╡ 885cd51d-895f-4996-a23b-780498b5b810
md"""
All overloads will do the operation (e.g. sum `x` and `y`), but also remember the pullback map and input variables for the reverse pass.

`@λ` (from the *LegibleLambdas.jl* package) is just for the nicer printing, we could have replaced `@λ(Δ -> (Δ, Δ))` with `Δ -> (Δ, Δ)` if we didn't care about that
"""

# ╔═╡ 13487e65-5e48-4a37-9bea-f262dd7b6d56
function Base.:+(x::Tracked, y::Tracked)
	Tracked(x.val + y.val, :+, @λ(Δ -> (Δ, Δ)), Tracked[x, y])
end

# ╔═╡ b0cc4665-eb45-48ea-9a33-5acf56d2a283
function Base.:-(x::Tracked, y::Tracked)
	Tracked(x.val - y.val, :-, @λ(Δ -> (Δ, -Δ)), Tracked[x, y])
end

# ╔═╡ 73d638bf-30c1-4694-b3a8-4b29c5e3fa65
function Base.:*(x::Tracked, y::Tracked)
	Tracked(x.val * y.val, :*, @λ(Δ -> (Δ * y.val', x.val' * Δ)), Tracked[x, y])
end

# ╔═╡ ac097299-0a31-474c-ab26-a4fb24bb9046
function Base.:^(x::Tracked, n::Int)
	Tracked(x.val^n, Symbol("^$n"), @λ(Δ -> (Δ * n * x.val^(n-1),)), Tracked[x,])
end

# ╔═╡ 2141849b-675e-406c-8df4-34b2706507af
function Base.:/(x::Tracked, y::Tracked)
	Tracked(x.val / y.val, :/, @λ(Δ -> (Δ / y.val, -Δ * x.val / y.val^2)), Tracked[x, y])
end

# ╔═╡ 8ab0f55d-a393-4a8a-a48c-9ced26033f57
function Base.sin(x::Tracked)
	Tracked(sin(x.val), :sin, @λ(Δ -> (Δ * cos(x.val),)), Tracked[x,])
end

# ╔═╡ 4bfc2f7d-a5b0-44c7-8bb6-f1b834c1cc51
md"""
`Tracked` is a tree -- We just need to tell *AbstractTrees.jl* how to get the children for each node and we get tree printing and iteration over all nodes for free.
"""

# ╔═╡ 2188a663-5a85-4ce4-bc8d-20383481e59b
AbstractTrees.children(x::Tracked) = x.deps

# ╔═╡ 00da514b-c6be-4d95-a0de-aed486615f3a
md"""
Let's also overload `show` for nicer output:
"""

# ╔═╡ 7429ffcb-dcee-4090-972e-ffde8393a37a
begin
	# All this is just for nicer printing
	function Base.show(io::IO, x::Tracked)
		if x.df === nothing
			print(io, Base.isgensym(x.name) ? x.val : "$(x.name)=$(x.val)")
		else
			print(io, "Tracked(")
			show(io, x.val)
			print(io, ", ")
			print(io, x.name)
			print(io, ")")
		end
	end
	Base.show(io::IO, ::MIME"text/plain", x::Tracked) = print_tree(io, x)
	Base.:(==)(x::Tracked, y::Tracked) = x === y
end

# ╔═╡ 727b2de3-1ee7-4f14-897f-46d263fa12ee
md"""
Create some variables we want to eventually differentiate with respect to.
"""

# ╔═╡ 0b5e6560-81fd-4182-bba5-aca702fb3048
begin
   x = Tracked{Int}(3, :x)
   y = Tracked{Int}(5, :y)
end

# ╔═╡ ccafd0b9-95aa-4e58-8afc-26cd3ee61cc9
md"""
Straight away we get the primal result of our calculation:
"""

# ╔═╡ 81eb8a2d-a3a9-45af-a5a5-b96aefd48712
(2x*y + (x-1)^2).val # The result of `2x*y + (x-1)^2`

# ╔═╡ 01eacbd4-ef37-4524-aecc-2ef9a1044cf8
md"""
To also get the gradient, we'll use `PreOrderDFS` to traverse the tree we just created from the top down.
"""

# ╔═╡ f0814e23-6f75-4db8-b277-d21d4926f876
z = (2x*y + (x-1)^2)

# ╔═╡ e52aa672-69a9-419b-a992-e7a3d1364fb6
# `PreOrderDFS` traverses this tree from the top down
Text.(collect(PreOrderDFS(z)))

# ╔═╡ d5565717-239b-41ab-b311-d59b383130ed
md"""
Ok, let's create our function `grad` which will accumulate all intermediate gradients into a dictionary:
"""

# ╔═╡ 99a3507b-ca03-429f-acde-e2d1ebb32054
function grad(f::Tracked)
	d = Dict{Any, Any}(f => 1)
	for x in PreOrderDFS(f) # recursively traverse all dependents
		x.df === nothing && continue # ignore untracked variables like constants
		dy = x.df(d[x]) # evaluate pullback
		for (yᵢ, dyᵢ) in zip(x.deps, dy)
			# store the gradient in d
			# if we have already stored a gradient for this variable, we need to add them
			d[yᵢ] = get(d, yᵢ, 0) .+ dyᵢ
		end
	end
	return d
end

# ╔═╡ d4e9b202-242e-4420-986b-12d2ab57af93
grad(f::Tracked, x::Tracked) = grad(f)[x]

# ╔═╡ 7fcccb65-4a7b-4527-97be-d25f481f6eaf
md"""
We can verify that it does the right thing:
"""

# ╔═╡ 18b1c55d-a6b5-44f6-b0b3-50bdb0aa9d96
w = x*y + x

# ╔═╡ 506d408e-dc2b-4e12-b917-286e3f4079a2
grad(w)

# ╔═╡ a7c8cb6a-6e17-4d8f-8958-fe3527c5b8e7
grad(w, x), grad(w, y)

# ╔═╡ e55dfad2-db50-459f-ab54-fa7637fc3638
md"""
## How can we visualize both the forward and the reverse pass?

We can further visualize each steps we just took. First we do the forwards calculation, where we also build up our tree, then we go down the tree in the opposite direction to accumulate our gradient.
"""

# ╔═╡ d9304d8d-9a34-46f2-908d-42d5cc5f5c5f
👽 = Tracked{Int}(12, :👽)

# ╔═╡ 27b39d7d-fc08-4ccc-aea4-b64f8a4f5726
#ex = :(3y*x + 2(x-1)*x)
ex = :(x^3 + y + sin(👽))

# ╔═╡ a6c2d3a2-326a-41dd-864d-aa3662466222
x,y

# ╔═╡ 5696b588-21c8-41cf-a28b-0b148a13dfa4
html"Move me!"

# ╔═╡ bcff2aa8-2387-44c7-a28f-39cd505a7adf
md"""
We can also visualize what Julia does in the forward pass on the code itself:
"""

# ╔═╡ 53ae4070-8818-4d21-8648-19df9319918a
md"""
## Let's write our own neural network!
"""

# ╔═╡ 8a70e835-dae6-4727-8204-95c87d5c23da
md"""
We need some more overloads
"""

# ╔═╡ 520f280d-c78e-433b-a0a2-8ef05b04f7cc
function Base.broadcasted(::typeof(tanh), x::Tracked)
	res = tanh.(x.val)
	Tracked(res, :tanh, @λ(Δ -> (Δ .* (1 .- res.^2),)), Tracked[x,])
end

# ╔═╡ a5671f2f-48ca-4896-9428-147dc671d2b9
function Base.broadcasted(::typeof(+), x::Tracked{<:VecOrMat}, y::Tracked{<:Vector})
	Tracked(x.val .+ y.val, :.+, @λ(Δ -> (Δ, sum(Δ; dims=2))), Tracked[x, y])
end

# ╔═╡ e01114fa-a847-424a-931f-8f42e623109f
function LinearAlgebra.norm(x::Tracked)
	res = norm(x.val)
	Tracked(res, :norm, @λ(Δ -> (Δ/res .* x.val,)), Tracked[x,])
end

# ╔═╡ e6ee1921-57aa-4194-b2eb-8e32fa4a6c44
macro t(s)
	:(Tracked{typeof($(esc(s)))}($(esc(s)), $(QuoteNode(s))))
end

# ╔═╡ 266bf2ed-d625-47be-8a87-21e12052b27f
md"""
Let's use three dense layers, all of the form

$x \mapsto \tanh.(W \cdot x + b)$
"""

# ╔═╡ 788903f0-a372-4265-bbf3-5b507a705100
NN(x, (W1, W2, W3, b1, b2, b3)) = W3 * tanh.(W2 * tanh.(W1 * @t(x) .+ b1) .+ b2) .+ b3

# ╔═╡ 337bc837-174f-4ed6-bee1-c9e80dc03b52
input = collect(range(0, 2π; length=100))'

# ╔═╡ 8f1ec8d7-89ab-4bc5-9eef-2b0a69019854
md"""
We'll try to approximate a sine curve:
"""

# ╔═╡ d2d21d41-9a4e-428b-b696-6878901849ba
ŷ = sin.(input) #.+ .2input .^ 2 .- cos.(5input)

# ╔═╡ 98807f0d-f8a5-4fe0-9983-2b95290347d9
md"""
### Tada! 🎉

(It's even reasonably fast for neural nets)
"""

# ╔═╡ e5d3f23e-3295-4d9b-a9aa-17270e5ca67a
function Base.show(io::IO, x::Tracked{<:AbstractArray})
	io = IOContext(io, :compact=>true)
		print(io, Base.isgensym(x.name) ? x.val : string(x.name))
end

# ╔═╡ d82adc20-4c8c-4f2c-9839-d03ad7e7f581
begin
	struct EX
		x::Any
		function EX(ex)
			if Meta.isexpr(ex, :call) && ex.args[1] === :+ && length(ex.args) > 3
				new(Expr(:call, :+, Expr(:call, :+, ex.args[2:end-1]...), ex.args[end]))
			else
				new(ex)
			end
		end
	end
	show_tree(ex::Expr) = show_tree(EX(ex))
	function Base.show(io::IO, ex::EX)
		Base.show_unquoted(io, Meta.isexpr(ex.x, :call) ? ex.x.args[1] : ex.x)
		if Meta.isexpr(ex.x, :call) && ex.x.args[1] === :^
			print(io, ex.x.args[3])
		end
	end
	function AbstractTrees.children(ex::EX)
		if Meta.isexpr(ex.x, :call)
			ex.x.args[1] === :^ ? [EX(ex.x.args[2])] : EX.(ex.x.args[2:end])
		else
			EX[]
		end
	end
	Base.:(==)(ex1::EX, ex2::EX) = ex1.x == ex2.x
	Base.hash(ex::EX, i::UInt) = hash(ex.x, i)
end

# ╔═╡ 96286b65-1a22-4458-a399-46579248cce4
begin
	_rand(s...) = randn(s...) / √prod(s)
	W1, W2, W3 = _rand(32, 1), _rand(32, 32), _rand(1, 32)
	b1, b2, b3 = _rand(32), _rand(32), _rand(1)
	params = [@t(W1), @t(W2), @t(W3), @t(b1), @t(b2), @t(b3)]
	p = @animate for i in 1:10000
		loss = norm(NN(input, params) - @t(ŷ))
		∇ = grad(loss)
		for p in params
			p.val .-= 1e-3 .* ∇[p]
		end
		plot(input', [NN(input, params).val; ŷ]'; label=["prediction" "training data"])
	end every 200
	gif(p; fps=5)
end

# ╔═╡ 8110f306-a7bb-43a2-bb36-6182c59b4b2e
begin
	struct TTREE
		x
		d::Dict{Any, Any}
	end
	function Base.show(io::IO, x::TTREE)
		show(io, x.x)
		print(io, " ")
		show(io, MIME("text/html"), get(x.d, x.x, @htl("")))
	end
	AbstractTrees.children(x::TTREE) = (TTREE(i, x.d) for i in children(x.x))
end

# ╔═╡ 1f1b384a-6588-45a5-9dd3-6de3face8bfb
function ad_steps(x::Expr; color_fwd="red", color_bwd="green", font_size=".8em")
	x = EX(x)
	repr(x) = sprint(show, x; context=:compact=>true)
	span_fwd = @htl ""
	span_bwd = @htl ""

    d1 = Dict(
		let e = eval(i.x)
			i => @htl " $(span_fwd)$(repr(e isa Tracked ?  e.val : e))"
		end
		for i in PostOrderDFS(x) if isempty(children(i))
	)

	res = accumulate(Iterators.filter(x -> !isempty(children(x)), PostOrderDFS(x)); init=d1) do d,i
		d = copy(d)
		e = eval(i.x)
		d[i] = @htl " $(span_fwd)$(repr(e isa Tracked ?  e.val : e))"
		d
	end

	pushfirst!(res, d1)

	f = eval(x.x)
	d = Dict{Any, Any}(f => 1)
	let d1 = copy(res[end])
		d1[x] = @htl "$(d1[x]) $(span_bwd)1"
		push!(res, d1)
	end
	for (x, e) in zip(PreOrderDFS.((f, x))...)
		x.df === nothing && continue
		dy = x.df(d[x])
		for (yᵢ, dyᵢ, e) in zip(x.deps, dy, children(e))
			d1 = copy(res[end])
			if haskey(d, yᵢ)
				d1[e] = @htl "$(get(d1, e, ""))$(span_bwd) + $(repr(dyᵢ))"
			else
				d1[e] = @htl "$(get(d1, e, "")) $(span_bwd)$(repr(dyᵢ))"
			end
			push!(res, d1)
			d[yᵢ] = get(d, yᵢ, 0) + dyᵢ
		end
	end
	res
end

# ╔═╡ 5585e9bb-7160-4cbf-b072-eb482edb8771
steps = ad_steps(ex);

# ╔═╡ 419842ed-fc24-420b-84eb-c9f9e575b860
i = @bind i Slider(1:length(steps))

# ╔═╡ 1a154bb7-93a3-4973-8908-788db77ac294
s2 = @htl """
<link rel="stylesheet" href="https://fperucic.github.io/treant-js/Treant.css"/>

"""

# ╔═╡ 6b1fb808-e993-4c2b-b81b-6710f8206de7
function to_json(x)
	d = Dict{Symbol, Any}(
		:innerHTML => sprint(AbstractTrees.printnode, x),
		:children => Any[to_json(c) for c in children(x)],
		#:collapsed => !isempty(children(x)),
	)
end

# ╔═╡ 437285d4-ec53-4bb7-9966-fcfb5352e205
function show_tree(x; height=400)
	s2
	id = gensym()
	@htl """

	"""
end

# ╔═╡ bb9bf66f-5ac8-4836-9d33-646a5c6f9015
show_tree(TTREE(EX(ex), steps[i]))

# ╔═╡ 86aa821b-a373-4814-953e-535f3a33c002
show_tree(norm(NN(input, params) - @t(ŷ)); height = 1000)

# ╔═╡ f6ce8448-d9ce-4453-9e47-dc6443d50f55
s1 = html"""

max_stars_repo_name: simeonschaub/ReverseModePluto

max_issues_repo_name: simeonschaub/ReverseModePluto

max_forks_repo_name: simeonschaub/ReverseModePluto


text:
volume_hypersphere(n, r=1) = π^(n / 2) / gamma(1 + n / 2) * r^n

"""
    log_vol_hypersphere(n, r=1)

Computes the log volume of a hypersphere

# Arguments

- `n`: the number of dimensions

# Keywords

- `r=1`: the radius
"""
log_vol_hypersphere(n; r=1) = (n / 2) * log(π) - loggamma(1 + n / 2) + n * log(r)

"""
    log_vol_ellipsiod(n, cov_mat; r = 1)

Computes the log volume of a ellipsoid

# Arguments
- `n`: the number of dimensions
- `cov_mat`: covariance matrix

# Keywords

- `r=1`: the radius
"""
function log_vol_ellipsiod(n, cov_mat; r = 1)
    vol = log_vol_hypersphere(n; r)
    eig_val,_ = eigen(cov_mat)
    return vol + .5 * sum(log(n + 2) .+ log.(eig_val))
end

"""
    volume_ellipsoid(n, cov_mat; r=1)

Computes the volume of an ellipsoid

# Arguments
- `n`: the number of dimensions
- `cov_mat`: covariance matrix

# Keywords

- `r=1`: the radius
"""
function volume_ellipsoid(n, cov_mat; r = 1)
    return exp(log_vol_ellipsiod(n, cov_mat; r))
end

volume_ellipsoid(cov_mat; r=1) = volume_ellipsoid(size(cov_mat, 2), cov_mat; r)

log_vol_ellipsiod(cov_mat; r=1) = log_vol_ellipsiod(size(cov_mat, 2), cov_mat; r)

"""
    sample_ellipsoid(μ, n_dims, cov_mat)

Samples a point uniformly from the interior of an ellipsoid

# Arguments

- `μ`: mean of ellipsoid on each dimension
- `n_dims`: the number of dimensions
- `cov_mat`: covariance matrix
"""
function sample_ellipsoid(μ, n_dims, cov_mat)
    x = random_position(rand(), n_dims)
    return μ .+ sqrt((n_dims + 2) * cov_mat) * x
end

"""
    sample_ellipsoid_surface(μ, n_dims, cov_mat; r=1)

Samples a point uniformly from the surface of an ellipsoid

# Arguments

- `μ`: mean of ellipsoid on each dimension
- `n_dims`: the number of dimensions
- `cov_mat`: covariance matrix

# Keywords

- `r=1`: radius of underlying unit hypersphere
"""
function sample_ellipsoid_surface(μ, n_dims, cov_mat; r=1)
    x = random_position(r, n_dims)
    return μ .+ sqrt((n_dims + 2) * cov_mat) * x
end

"""
    bias_correction(
        model,
        p_fun,
        points,
        target,
        bounds,
        args...;
        parm_names = Symbol.("p", 1:size(points,2)),
        n_sim=10_000,
        kwargs...
    )

Uses the hit or miss technique to adjust volume estimates based on ellipsoids.

# Arguments

- `model`: a model function that returns predictions given a vector of parameters
- `p_fun`: a function that that evaluates the qualitative data pattern
- `points`: a p x n matrix of sampled points where p is the number of parameters and n is the number of samples
- `bounds`: a vector of tuples representing (lowerbound, upperbound) for each dimension in
- `args...`: arguments passed to `model` and `p_fun`

# Keywords

- `parm_names`: parameter names with default `p1`, `p2,`... `pn`
- `n_sim=10_000`: number of simulations for hit or miss adjustment
- `kwargs...`: optional keyword arguments for `model` and `p_fun`
"""
function bias_correction(
    model,
    p_fun,
    points,
    target,
    bounds,
    args...;
    parm_names = Symbol.("p", 1:size(points,2)),
    n_sim=10_000,
    kwargs...
    )

    _μ = mean(points, dims=1)[:]
    μ = ComponentArray(_μ, make_axis(parm_names))
    cov_mat = cov(points)
    n = length(μ)

    _model = x -> model(x, args...; kwargs...)
    _p_fun = x -> p_fun(x, args...; kwargs...)

    hits = 0
    for i in 1:n_sim
        parms = sample_ellipsoid(μ, n, cov_mat)
        if in_bounds(parms, bounds)
            pattern = _p_fun(_model(parms))
            hits += hit_or_miss(pattern, target, parms, bounds)
        end
    end
    return hits / n_sim
end

function hit_or_miss(target, pattern, parms, bounds)
    if target == pattern
        return 1
    end
    return 0
end

"""
    estimate_volume(
        model,
        p_fun,
        points,
        target,
        bounds,
        args...;
        n_sim = 10_000,
        kwargs...
    )

Estimate volume of region with an eillipsoid and hit or miss bias adjustment.

# Arguments
- `model`: a model function that returns predictions given a vector of parameters
- `p_fun`: a function that that evaluates the qualitative data pattern
- `points`: a p x n matrix of sampled points where p is the number of parameters and n is the number of samples
- `target`: the target pattern associated with the `points`
- `bounds`: a vector of tuples representing (lowerbound, upperbound) for each dimension in
- `args...`: arguments passed to `model` and `p_fun`

# Keywords
- `n_sim=10_000`: the number of samples for hit or miss bias adjustment
- `kwargs...`: additional keyword arguments passed to `model` or `p_fun`
"""
function estimate_volume(
    model,
    p_fun,
    points,
    target,
    bounds,
    args...;
    n_sim = 10_000,
    parm_names = Symbol.("p", 1:size(points,2)),
    kwargs...
    )

    cf = bias_correction(
        model,
        p_fun,
        points,
        target,
        bounds,
        args...;
        n_sim,
        parm_names,
        kwargs...
    )

    cov_mat = cov(points)
    volume = volume_ellipsoid(cov_mat)
    volume *= cf
    return volume
end

make_axis(symbols) = Axis(NamedTuple(symbols .=> eachindex(symbols)))

max_stars_repo_name: itsdfish/ParameterSpacePartitions

max_forks_repo_name: itsdfish/ParameterSpacePartitions


text:
abstract type AbstractODMNode end

struct ODMNodeType{Symbol}
    function ODMNodeType(s::Symbol)
        new{s}()
    end
    function ODMNodeType(s::Nothing)
        new{:TextNode}()
    end
end

struct ODMRoot <: AbstractODMNode
    name::Symbol
    attr::Dict{Symbol, String}
    el::Vector{AbstractODMNode}
    function ODMRoot(attr)
        new(:ODM, attr, AbstractODMNode[])
    end
end

struct ODMNode <: AbstractODMNode
    name::Symbol
    attr::Dict{Symbol, String}
    el::Vector{AbstractODMNode}
    function ODMNode(name, attr, el)
        new(name, attr, el)
    end
    function ODMNode(name, attr)
        ODMNode(name, attr, AbstractODMNode[])
    end
end
struct ODMTextNode <: AbstractODMNode
    content::String
end
struct StudyMetaData <: AbstractODMNode
    metadata::ODMNode
    el::Vector
end

function Base.show(io::IO, n::T) where T <: AbstractODMNode
    print(io, "$(n.name)  ($(:OID in keys(n.attr) ? "OID:$(attribute(n, :OID)))" : "")$(:StudyOID in keys(n.attr) ? " StudyOID:$(attribute(n, :StudyOID)))" : ""))")
end
function Base.show(io::IO, n::ODMRoot)
    print(io, "ODM root node")
end
function Base.show(io::IO, n::StudyMetaData)
    print(io, "Completed Study MetaData ($(length(n.el)) elements), OID: $(attribute(n.metadata, :OID)), Name: $(attribute(n.metadata, :Name))")
end
function Base.show(io::IO, n::ODMTextNode)
    print(io, "Text Node")
end

function AbstractTrees.children(x::T) where T <: AbstractODMNode
    x.el
end
function AbstractTrees.children(i::ODMTextNode)
    []
end
    #AbstractTrees.nodetype(::IntTree) = IntTree

function makenode(str, attr)
    symb = Symbol(str)
    return ODMNode(symb, attr)
end
function makenode(content::String)
    return ODMTextNode(content)
end
function attributes_dict(n)
    d = Dict{Symbol, String}()
    for i in eachattribute(n)
        d[Symbol(i.name)] = i.content
    end
    d
end
"""
    importxml(file::AbstractString)

Import odm.xml file.
"""
function importxml(file::AbstractString)
    isfile(file) || error("File not found!")
    xdoc  = readxml(file)
    xodm  = root(xdoc)
    odm   = ODMRoot(attributes_dict(xodm))
    importxml_(odm, xodm)
    odm
end
    function importxml_(parent, root)
        if haselement(root)
            chld = EzXML.eachelement(root)
            for c in chld
                if iselement(c)
                    attr = attributes_dict(c)
                    odmn = makenode(nodename(c), attr)
                    push!(parent.el, odmn)
                    if hasnode(c) importxml_(odmn, c) end
                end
            end
        else
            chld = EzXML.eachnode(root)
            for c in chld
                if istext(c)
                    odmn = makenode(nodecontent(c))
                    push!(parent.el, odmn)
                end
            end
        end
        parent
    end
################################################################################
# SUPPORT FUNCTIONS
################################################################################
function attribute(n::AbstractODMNode, attr::Symbol)
    if ht_keyindex(n.attr, attr) > 0 return n.attr[attr] else return "" end
end
function attribute(n::AbstractODMNode, attr::String)
    attribute(n, Symbol(attr))
end
function attributes(n::AbstractODMNode, attrs::AbstractVector{T}) where T <: StrOrSym
    broadcast(x-> attribute(n, x), attrs)
end

function name(n::ODMRoot)
    :ODM
end
function name(n::ODMNode)
    n.name
end
function name(n::ODMTextNode)
    nothing
end
function have_attr(n::AbstractODMNode, attr::Symbol)
    if ht_keyindex(n.attr, attr) > 0 return true else return false end
end
function have_oid(n::AbstractODMNode)
    have_attr(n, :OID)
end
function appendelements!(inds::AbstractVector, n::AbstractODMNode, name::Symbol)
    for i in n.el
        if i.name == name
            push!(inds, i)
        end
    end
    inds
end
function content(n::ODMTextNode)
    n.content
end
function content(n::AbstractODMNode)
    for i in n.el
        if isa(i, ODMTextNode)
            return content(i)
        end
    end
    nothing
end
################################################################################
# BASIC FUNCTIONS
################################################################################
"""
    findelement(n::AbstractODMNode, name::Symbol, oid::AbstractString)

Find first element by name and oid.
"""
function findelement(n::AbstractODMNode, name::Symbol, oid::AbstractString)
    for i in n.el
        if i.name == name
            if have_oid(i)
                if attribute(i, :OID) == oid return i end
            end
        end
    end
end
"""
    findelement(n::AbstractODMNode, name)

Find first element by name.
"""
function findelement(n::AbstractODMNode, name::Symbol)
    for i in n.el
        if i.name == name
            return i
        end
    end
end
"""
    findelements(n::AbstractODMNode, name::Symbol)

Find all elements by name.
"""
function findelements(n::AbstractODMNode, name::Symbol)
    inds = AbstractODMNode[]
    for i in n.el
        if i.name == name
            push!(inds, i)
        end
    end
    inds
end
"""
    countelements(n::AbstractODMNode, name::Symbol)

Count elements by name.
"""
function countelements(n::AbstractODMNode, name::Symbol)
    inds = 0
    for i in n.el
        if i.name == name
            inds += 1
        end
    end
    inds
end
################################################################################
# SEARCH FUNCTIONS (find*)
################################################################################
"""
    metadatalist(odm::ODMRoot)

Returm table of MetaDataVersion elements.
"""
function metadatalist(odm::ODMRoot)
    df = DataFrame(StudyOID = String[], OID = String[], Name = String[])
    for i in odm.el
        if i.name == :Study
            for j in i.el
                if j.name == :MetaDataVersion
                    push!(df, (attribute(i, "OID"), attribute(j, "OID"), attribute(j, "Name")))
                end
            end
        end
    end
    df
end
"""
    studylist(odm::ODMRoot)

Returm table of Study elements.
"""
function studylist(odm::ODMRoot)
    df = DataFrame(StudyOID = String[])
    for i in odm.el
        if i.name == :Study
            push!(df, (attribute(i, "OID"),))
        end
    end
    df
end
"""
    clinicaldatalist(odm::ODMRoot)

Returm table of ClinicalData elements.
"""
function clinicaldatalist(odm::ODMRoot)
    df = DataFrame(StudyOID = String[], MetaDataVersionOID = String[])
    for i in odm.el
        if i.name == :ClinicalData
            push!(df, (attribute(i, "StudyOID"), attribute(i, "MetaDataVersionOID")))
        end
    end
    df
end
"""
    findclinicaldata(odm::ODMRoot, soid::AbstractString)

Return ClinicalData element by study OID (`soid`), nothing if not found.
"""
function findclinicaldata(odm::ODMRoot, soid::AbstractString)
    for i in odm.el
        if i.name == :ClinicalData && attribute(i, "StudyOID") == soid return i end
    end
    nothing
end
"""
    findclinicaldata(odm::ODMRoot, soid::AbstractString, moid::AbstractString)

Find ClinicalData by StudyOID and MetaDataVersionOID.
"""
function findclinicaldata(odm::ODMRoot, soid::AbstractString, moid::AbstractString)
    for i in odm.el
        if i.name == :ClinicalData && attribute(i, "StudyOID") == soid && attribute(i, "MetaDataVersionOID") == moid return i end
    end
    nothing
end
"""
    findstudy(odm::ODMRoot, oid::AbstractString)

Find Study element by OID (`oid`), nothing if not found.
"""
function findstudy(odm::ODMRoot, oid::AbstractString)
    for i in odm.el
        if i.name == :Study && attribute(i, "OID") == oid return i end
    end
    nothing
end
"""
    findstudymetadata(odm::ODMRoot, soid::AbstractString, moid::AbstractString)

Find metadata for study with study OID soid and metadata OID moid.
"""
function findstudymetadata(odm::ODMRoot, soid::AbstractString, moid::AbstractString)
    study = findstudy(odm, soid)
    findelement(study, :MetaDataVersion, moid)
end
################################################################################
# List functions (return DataFrames)
################################################################################
"""
    eventlist(md::AbstractODMNode)

Return events (StudyEventDef).
"""
function eventlist(md::AbstractODMNode)
    df = DataFrame(OID = String[], Name = String[], Repeating= String[], Type = String[])
    for i in md.el
        if name(i) == :StudyEventDef
            push!(df, (attribute(i, "OID"), attribute(i, "Name"), attribute(i, "Repeating"), attribute(i, "Type")))
        end
    end
    df
end
"""
    formlist(md::AbstractODMNode)

Return forms (FormDef).
"""
function formlist(md::AbstractODMNode)
    formlist(md.el)
end
"""
    formlist(el::Vector{T}) where T <: AbstractODMNode

Return forms (FormDef).
"""
function formlist(el::Vector{T}) where T <: AbstractODMNode
    df = DataFrame(OID = String[], Name = String[], Repeating= String[])
    for i in el
        if name(i) == :FormDef
            push!(df, (attribute(i, "OID"), attribute(i, "Name"), attribute(i, "Repeating")))
        end
    end
    df
end

"""
    itemgrouplist(md::AbstractODMNode; optional = false, attrs::Union{AbstractVector, Nothing} = nothing)

Return item groups (ItemGroupDef).

If optional = true - return all optional attributes.
attrs - list of attributes.
"""
function itemgrouplist(md::AbstractODMNode; optional = false, attrs::Union{AbstractVector, Nothing} = nothing)
    itemgrouplist(md.el; optional = optional, attrs = attrs)
end
"""
    itemgrouplist(el::Vector{T}; optional = false, attrs::Union{AbstractVector, Nothing} = nothing) where T <: AbstractODMNode

Return item groups (ItemGroupDef).

If optional = true - return all optional attributes.
attrs - list of attributes.
"""
function itemgrouplist(el::Vector{T}; optional = false, attrs::Union{AbstractVector, Nothing} = nothing) where T <: AbstractODMNode
    if isnothing(attrs)
        if optional
            attrs = [:OID, :Name, :Repeating, :SASDatasetName, :Comment]
        else
            attrs = [:OID, :Name, :Repeating]
        end
    end
    df = DataFrame(Matrix{Union{Missing, String}}(undef, 0, length(attrs)), attrs)
    for i in el
        if name(i) == :ItemGroupDef
            push!(df, attributes(i, attrs))
        end
    end
    df
end

"""
    itemlist(md::AbstractODMNode; optional = false, attrs = nothing)

Return items (ItemDef).

If optional = true - return all optional attributes.
"""
function itemlist(md::AbstractODMNode; optional = false, attrs = nothing)
    itemlist(md.el; optional = optional, attrs = attrs)
end
"""
    itemlist(el::Vector{T}; optional = false, attrs = nothing) where T <: AbstractODMNode
"""
function itemlist(el::Vector{T}; optional = false, attrs = nothing) where T <: AbstractODMNode
    if isnothing(attrs)
        if optional
            attrs = [:OID, :Name, :DataType, :Length, :SignificantDigits, :SASFieldName, :SDSVarName, :Origin, :Comment]
        else
            attrs = [:OID, :Name, :DataType]
        end
    end
    df = DataFrame(Matrix{Union{Missing, String}}(undef, 0, length(attrs)), attrs)
    for i in el
        if name(i) == :ItemDef
            push!(df, attributes(i, attrs))
        end
    end
    df
end
################################################################################
# CONTENT
################################################################################
function formcontent_(md, oid)
    ig   = findelement(md, :FormDef, oid)
    inds = AbstractODMNode[]
    for i in ig.el
        if name(i) == :ItemGroupRef
            el = findelement(md, :ItemGroupDef, attribute(i, "ItemGroupOID"))
            if !(isnothing(el)) push!(inds, el) end
        end
    end
    inds
end

function itemgroupcontent_(md, oid)
    ig   = findelement(md, :ItemGroupDef, oid)
    inds = AbstractODMNode[]
    for i in ig.el
        if name(i) == :ItemRef
            el = findelement(md, :ItemDef, attribute(i, "ItemOID"))
            if !(isnothing(el)) push!(inds, el) end
        end
    end
    inds
end

"""
    formcontent(md, oid)

Return item groups (ItemGroupDef) for concrete form (FormDef) by OID.

"""
function formcontent(md, oid; optional = false, attrs::Union{AbstractVector, Nothing} = nothing)
    inds = formcontent_(md, oid)
    itemgrouplist(inds; optional = optional, attrs = attrs)
end
"""
    itemgroupcontent(md, oid; optional = false)

Return items (ItemDef) for concrete item group (ItemGroupDef) by OID.

If optional = true - return all optional attributes.
"""
function itemgroupcontent(md, oid; optional = false)
    inds = itemgroupcontent_(md, oid)
    itemlist(inds; optional = optional)
end

"""
    itemsformcontent(md, oid; optional = false)

"""
function itemsformcontent(md, oid; optional = false)
    inds   = AbstractODMNode[]
    frm    = findelement(md, :FormDef, oid)
    iginds = findelements(frm, :ItemGroupRef)
    for i in iginds
        igoid = attribute(i, "ItemGroupOID")
        append!(inds, itemgroupcontent_(md, igoid))
    end
    itemlist(inds; optional = optional)
end
################################################################################
# TOP LEVEL FUNCTIONS
# clinicaldatatable
# buildmetadata
# buildelementsdata
################################################################################
"""
    buildmetadata(odm::ODMRoot, soid::AbstractString, moid::AbstractString)

Build MetaData from MetaDataVersion.
"""
function buildmetadata(odm::ODMRoot, soid::AbstractString, moid::AbstractString)
    mdat   = findstudymetadata(odm, soid, moid)
    stmd   = StudyMetaData(mdat, AbstractODMNode[])
    fillstmd_(stmd.el, stmd.metadata, odm)
    stmd
end
"""
    buildmetadata(mdat::AbstractODMNode)

Build MetaData from MetaDataVersion mdat.
"""
function buildmetadata(mdat::AbstractODMNode)
    if name(cd) != :MetaDataVersion error("This is not MetaDataVersion") end
    stmd   = StudyMetaData(mdat, AbstractODMNode[])
    fillstmd_(stmd.el, stmd.metadata, odm)
    stmd
end
"""
    clinicaldatatable(cd::AbstractODMNode; itemgroup = nothing)

Return clinical data table in long formal. `cd` should be ClinicalData.
"""
function clinicaldatatable(cd::AbstractODMNode; itemgroup = nothing)
    if name(cd) != :ClinicalData error("This is not ClinicalData") end
    df = DataFrame(SubjectKey = String[], StudyEventOID = String[], FormOID = String[], ItemGroupOID = String[], ItemGroupRepeatKey = String[], ItemOID = String[], Value = String[])
    sdl = findelements(cd, :SubjectData)
    edl = AbstractODMNode[]
    fdl = AbstractODMNode[]
    gdl = AbstractODMNode[]
    idl = AbstractODMNode[]
    for s in sdl
        resize!(edl, 0)
        appendelements!(edl, s, :StudyEventData)
        for e in edl
            resize!(fdl, 0)
            appendelements!(fdl, e, :FormData)
            for f in fdl
                resize!(gdl, 0)
                appendelements!(gdl, f, :ItemGroupData)
                for g in gdl
                    if !isnothing(itemgroup)
                        if attribute(g, :ItemGroupOID) != itemgroup continue end
                    end
                    resize!(idl, 0)
                    appendelements!(idl, g, :ItemData)
                    for i in idl
                        push!(df, (attribute(s, :SubjectKey),
                        attribute(e, :StudyEventOID),
                        attribute(f, :FormOID),
                        attribute(g, :ItemGroupOID),
                        attribute(g, :ItemGroupRepeatKey),
                        attribute(i, :ItemOID),
                        attribute(i, :Value)))
                    end
                end
            end
        end
    end
    df
end
"""
    clinicaldatatable(odm::ODMRoot, soid::AbstractString, moid::AbstractString; itemgroup = nothing)

Return clinical data table in long formal.
"""
function clinicaldatatable(odm::ODMRoot, soid::AbstractString, moid::AbstractString; itemgroup = nothing)
    cld  = findclinicaldata(odm, soid, moid)
    isnothing(cld) && error("ClinicalData not found")
    clinicaldatatable(cld; itemgroup = itemgroup)
end
"""
    subjectdatatable(cld::AbstractODMNode; optional = false, attrs::Union{AbstractVector, Nothing} = nothing)

Subject information table
"""
function subjectdatatable(cld::AbstractODMNode; optional = false, attrs::Union{AbstractVector, Nothing} = nothing)
    if name(cld) != :ClinicalData error("This is not ClinicalData") end
    if isnothing(attrs)
        if optional
            attrs = [:SubjectKey, :TransactionType]
        else
            attrs = [:SubjectKey]
        end
    end
    df = DataFrame(Matrix{Union{String, Missing}}(undef, 0, length(attrs)), attrs)
    for s in cld.el
        if name(s) == :SubjectData
            push!(df, attributes(s, attrs))
        end
    end
    df
end
"""
    subjectdatatable(odm::ODMRoot; optional = false, attrs::Union{AbstractVector, Nothing} = nothing)
"""
function subjectdatatable(odm::ODMRoot; optional = false, attrs::Union{AbstractVector, Nothing} = nothing)
    cld = findelements(odm, :ClinicalData)
    if length(cld) > 0
        df = subjectdatatable(cld[1]; optional = optional, attrs = attrs)
        if length(cld) > 1
            for i = 2:length(cld)
                append!(df, subjectdatatable(cld[i]; optional = optional, attrs = attrs))
            end
        end
        return df
    end
    nothing
end

"""
    buildelementsdata(stmd::AbstractODMNode)

Build elements data.
"""
function buildelementsdata(stmd::AbstractODMNode)
    error("not implemented")
end

################################################################################
# BUILD MetaData functions
################################################################################
function inlist(n, dest)
    for i in dest
        if n.name == i.name
            if have_oid(n)
                if n.attr[:OID] == i.attr[:OID] return true end
            else
                return true
            end
        end
    end
    false
end
function fillstmd_(dest, source, odm)
    inds = AbstractODMNode[]
    for i in source.el
        if !inlist(i, dest)
            push!(inds, i)
        end
    end
    if length(inds) > 0
        append!(dest, inds)
    end
    incl = findelements(source, :Include)
    if length(incl) > 0
        for i in incl
            fillstmd_(dest, findstudymetadata(odm, attribute(i, :StudyOID), attribute(i, :MetaDataVersionOID)), odm)
        end
    end
    dest
end
################################################################################
# Information
################################################################################
"""
    studyinfo(odm::ODMRoot, oid::AbstractString;  io = stdout)

Study information.
"""
function studyinfo(odm::ODMRoot, oid::AbstractString;  io = stdout)
    for i in odm.el
        if name(i) == :Study && attribute(i, :OID) == oid
            studyinfo(i; io = io)
        end
    end
    print(io, "")
end
"""
    studyinfo(odm::ODMRoot;  io = stdout)

Study information.
"""
function studyinfo(odm::ODMRoot;  io = stdout)
    str = ""
    for i in odm.el
        if name(i) == :Study
            str *= "--------------------------------------\n"
            str *= studyinfo_(i)
        end
    end
    str *= "--------------------------------------\n"
    print(io, str)
end
"""
    studyinfo(st::AbstractODMNode;  io = stdout)

Study information.
"""
function studyinfo(st::AbstractODMNode;  io = stdout)
    str = studyinfo_(st)
    print(io, str)
end
function studyinfo_(st::AbstractODMNode)
    if name(st) != :Study error("This is not Study") end
    str = "Study OID: $(attribute(st, "OID"))\n"
    gv  = findelement(st, :GlobalVariables)
    sn  = findelement(gv, :StudyName)
    str *= "StudyName: $(content(sn))\n"
    sd  = findelement(gv, :StudyDescription)
    str *= "StudyDescription: $(content(sd))\n"
    pn  = findelement(gv, :ProtocolName)
    str *= "ProtocolName: $(content(pn))\n"
    str *= "MetaDataVersion: \n"
    mdl = findelements(st, :MetaDataVersion)
    if length(mdl) > 0
        for i in mdl
            str *= "    OID: $(attribute(i, :OID)); Name: $(attribute(i, :Name))\n"
        end
    else
        str *= "    No\n"
    end
    str
end


text:
module Experiments

include("data.jl")
include("neuralnetwork.jl")
include("plotlib.jl")
include("plotutils.jl")
include("structs.jl")

using DataStructures
using PlotlyJS

export run_experiments

plot_dir = "plots"

"""
	run_experiments(experiments, prefix = "exp_1", random_moons = true, show_plots = false, save_plots = true, plot_size = "medium", plot_type = "both", epochs = 500, batchsize = 42)

Runs all experiments as specified by an DataStructures.OrderedDict.
This includes the creation of moon data, the network, the training and the plotting of the results.
This function is the work horse to conduct any analyses for the notebook.
"""
function run_experiments(experiments::DataStructures.OrderedDict; prefix::String = "", random_moons::Bool = true, show_plots::Bool = true, save_plots::Bool = true, plot_size::String = "medium", plot_type::String = "both", args...)
	plot_type = lowercase(plot_type)
	@assert plot_type in ["boundary", "stats", "both"] "Requested plot type '$plot_type' is not supported."

    plot_width, plot_height = PlotUtils.get_plot_resolution(plot_size)

    for (name, def) in experiments
        # load data
		if random_moons
        	train_loader = Data.get_moon_data_loader(offset = def.offset)
		else
			train_loader = Data.get_moon_data_loader(offset = def.offset, seed = 1337)
		end

        # create and train the model
        model, stats = NeuralNetwork.train(train_loader, def.optimizer, def.sd; args...)

        # create plots
        p1 = PlotLib.plot_decision_boundary(train_loader, model, title = name)
        p2 = PlotLib.plot_loss_and_accuracy(stats["loss"], stats["accuracy"]; args...)

		if plot_type == "both"
			# combine subplots
			p = [p1 p2]
			# make background and plots transparent (looks better in dark-styled browsers)
			PlotlyJS.relayout!(p, paper_bgcolor = "rgba(0, 0, 0, 0)", plot_bgcolor = "rgba(0, 0, 0, 0)")
		elseif plot_type == "boundary"
			p = p1
		elseif plot_type == "stats"
			p = p2
		end

		# save requested plot
        if save_plots
			p_name = "$plot_dir/$prefix" * (prefix != "" ? "_" : "") * replace(name, " " => "_") * (plot_type == "both" ? "" : plot_type) * ".png"
			PlotlyJS.savefig(p, p_name, width = plot_width, height = plot_height)
		end

        if show_plots
			display(p)
		end
    end

	@info "Finished $(length(experiments)) experiments."
end

end