text:
__precompile__()
module FragmentGrammars

export Analysis, BaseDistribution, BaseRule, FragmentRule, AbstractRule, Fragment, Pointer, FragmentGrammar
export sample, add_obs!, rm_obs!
export run_chartparser, sample_tree
export category_type, terminal_type, category_rule_type, terminal_rule_type, startstate, startsymbols, score_type, state_type, completions, prob
export iterate, eltype, IteratorSize
# export show

import Base: iterate, eltype, IteratorSize#, length, show

include("GeneralizedChartparsing\\src\\GeneralizedChartparsing.jl")
using .GeneralizedChartparsing
using .GeneralizedChartparsing.Trees
using .GeneralizedChartparsing: sample_tree, Constituent
import .GeneralizedChartparsing: lhs, rhs, category_type, terminal_type, category_rule_type, terminal_rule_type, startstate, add_rule!, startsymbols, score_type, state_type, completions, prob
import .GeneralizedChartparsing: run_chartparser, categorical_sample

include("CompoundDists.jl");
using .CompoundDists
import .CompoundDists: sample, add_obs!, rm_obs!, logscore

# TODO: ALL OF THIS BADLY NEEDS TO BE REFACTORED.
# TODO: So much code can be turned into broadcasts
# TODO: Coins are flipped for preterminals??? Sheesh. Fix that.

###########################
# Helper structs and functions #
###########################
struct Fragment
    tree :: Tree
    variables :: Vector{Tree}
    leaves :: Vector{Tree}
end

struct Pointer
    fragment :: Fragment
    children :: Dict{Tree, Pointer}
end

function clone(ptr::Pointer)
    function clone(ptr::Pointer, tree::Tree, variables::Vector{Tree}, leaves::Vector{Tree}, children::Dict{Tree, Pointer})
        newtree = TreeNode(tree.data)
        if !isempty(tree.children)
            for child in tree.children
                insert_child!(newtree, clone(ptr, child, variables, leaves, children))
            end
        else
            if tree in ptr.fragment.variables
                if get(ptr.children, tree, nothing) != nothing
                    push!(children, newtree => ptr.children[tree])
                end
                push!(variables, newtree)
                push!(leaves, newtree)
            elseif tree in ptr.fragment.leaves
                push!(leaves, newtree)
            end
        end
        return newtree
    end
    variables = Tree[]
    leaves = Tree[]
    children = Dict{Tree, Pointer}()
    return Pointer(Fragment(clone(ptr, ptr.fragment.tree, variables, leaves, children), variables, leaves), children)
end

clone(frag::Fragment) = clone(Pointer(frag, Dict{Tree,Pointer}())).fragment

struct Analysis{C, CR} # Is this struct even needed? It's basically a Tuple of what comes out of FG sample
    pointer :: Pointer
    dm_obs :: Vector{Tuple{C,CR}}
    bb_obs :: Vector{Pair{Tuple{C, CR, C},Bool}}
    crp_obs :: Vector{Fragment}
end

abstract type AbstractRule{C1,C2} end

struct BaseRule{C1,C2} <: AbstractRule{C1,C2}
    # rule :: ContextFreeRule{C1,C2}
    lhs :: C1
    rhs :: Tuple{Vararg{C2}}
end
BaseRule(rule) = BaseRule(#= rule, =# lhs(rule), rhs(rule))
# BaseRule(lhs::C1, rhs::Tuple{Vararg{C2}}) where {C1,C2} = BaseRule(ContextFreeRule{C1,C2}(lhs, rhs))

struct FragmentRule{C1,C2} <: AbstractRule{C1,C2}
    fragment :: Fragment
    # rule :: ContextFreeRule{C1,C2}
    lhs :: C1
    rhs :: Tuple{Vararg{C2}}
end
FragmentRule(rule, fragment) = FragmentRule(fragment, #= rule, =#lhs(rule), rhs(rule))
FragmentRule(fragment) = FragmentRule(fragment, fragment.tree.data, (getfield.(fragment.leaves, :data)...,))

mutable struct ApproxRule{C1,C2} <: AbstractRule{C1,C2}
    lhs :: C1
    rhs :: Tuple{Vararg{C2}}
    rules :: Vector{AbstractRule{C1,C2}}
    probs :: Vector{LogProb}
    prob :: LogProb
end

approxprob(ar::ApproxRule) = ar.prob
approxprob(br::BaseRule) = one(LogProb)

# Unzip rule-logprob pairs into two lists, cast the logprob list into floats, and call categorical_sample
sample(r::ApproxRule) = categorical_sample(r.rules, r.probs)
# sample(r::ApproxRule) = categorical_sample((x->(x[1], float.(x[2])))(collect(zip(r.rules...)))...)
categorical_sample(tokens::Vector{X}, weights::Vector{LogProb}) where X = categorical_sample(tokens, float.(weights))

lhs(r::AbstractRule) = r.lhs
rhs(r::AbstractRule) = r.rhs
(r::BaseRule)(cat) = rhs(r)
(r::BaseRule)() = rhs(r)
(r::FragmentRule)(cat) = rhs(r)
(r::FragmentRule)() = rhs(r)
(r::ApproxRule)(cat) = rhs(r)
(r::ApproxRule)() = rhs(r)

show(io::IO, r::AbstractRule{C1,C2}) where {C1,C2} =
print("AbstractRule($(lhs(r)),$(rhs(r)))")
show(io::IO, r::BaseRule{C1,C2}) where {C1,C2} =
println("BaseRule($(lhs(r)),$(rhs(r)))")
show(io::IO, r::FragmentRule{C1,C2}) where {C1,C2} =
print("FragmentRule($(lhs(r)),$(rhs(r)))")
show(io::IO, r::ApproxRule{C1,C2}) where {C1,C2} =
print("ApproxRule($(lhs(r)),$(rhs(r)))")

eltype(::Type{Pointer}) = Pointer
IteratorSize(::Type{Pointer}) = Base.SizeUnknown()
function iterate(frag_pointer::Pointer, state = [frag_pointer])
    if isempty(state)
        nothing
    else
        state[1], prepend!(state[2:end], collect(values(state[1].children)))
    end
end

eltype(::Type{State{C,CR}}) where {C,CR} = State{C,CR}
IteratorSize(::Type{State{C,CR}}) where {C,CR} = Base.SizeUnknown()
function iterate(fsa_state::State, state = [fsa_state])
    if isempty(state)
        nothing
    else
        state[1], prepend!(state[2:end], collect(values(state[1].trans)))
    end
end

get_idx(A::AbstractVector{T}, i::T) where T = (
    for (j,k) in enumerate(A)
        if i == k
            return j
        end
    end; error("element $i not found in $A")
)

# Basically copied from add_rule! in GeneralizedChartparsing/src/Grammars.jl
function add_rule!(state::State{C,AbstractRule{C,C}}, rule::AbstractRule{C,C}, head, cats) where {C} #, CR<:ApproxRule{C,C}}
    s = state
    for c in cats
        if is_possible_transition(s, c)
            s = transition(s, c)
        else
            s.isfinal = false
            s = s.trans[c] = State(C, AbstractRule{C,C})
        end
    end
    # All unique lhs-rhs pairs are collapsed to include base rules and fragments, and transitions between states determine RHS, therefore we only need to check LHS here.
    aridx = findall((comp -> comp[1] == head), s.comp)
    # @show aridx, head, cats, rule
    if isempty(aridx) # if no approx rule yet
        ar = ApproxRule(head, cats, AbstractRule{C,C}[rule], LogProb[zero(LogProb)], zero(LogProb))
        push!(s.comp, (head, ar))
    else
        ar = s.comp[aridx[1]][2] # there can be only one! [2] because [1] is lhs and [2] is approxrule
        if !(rule in ar.rules) push!(ar.rules, rule); push!(ar.probs, zero(LogProb)) end
    end
    # push!(s.comp, (head, rule))
end

function rm_rule!(state::State{C,AbstractRule{C,C}}, rule::AbstractRule{C,C}, head, cats) where {C} #, CR<:ApproxRule{C,C}}
    s = state
    for c in cats
        if is_possible_transition(s, c)
            s = transition(s, c)
        else
            s.isfinal = false
            s = s.trans[c] = State(C, CR)
        end
    end
    aridx = findall((comp -> comp[1] == head), s.comp)
    # @show aridx, head, cats, rule
    # if isempty(aridx) # if no approx rule yet
        # Nothing needs to happen here, right?
    # else
    if !isempty(aridx) # Shouldn't this never be the case?
        ar = s.comp[aridx[1]][2] # there can be only one! [2] because [1] is lhs and [2] is approxrule
        # @show ar.rules
        ridx = findall(r -> r===rule, ar.rules)
        if !isempty(ridx)
            ridx = ridx[1] # there can be only one!
            deleteat!(ar.rules, ridx)
            deleteat!(ar.probs, ridx)
            # filter!(r -> r[1]≠rule, ar.rules) # Remove "rule" from the approx rule
            if isempty(ar.rules) && isempty(ar.probs)
                deleteat!(s.comp, aridx) # Remove empty approx rule from completions
            end
        end
    end
end

################################
# Fragment Grammar definitions #
################################
"""
    FragmentGrammar(categories, startcategories, category_rules, terminals, terminal_rules, startstate, terminal_dict, CRP, DM, BB)
"""
mutable struct FragmentGrammar{C, CR, T, TR}
    categories :: Vector{C}
    startcategories :: Vector{C}
    category_rules :: Vector{CR}
    terminals :: Vector{T}
    terminal_rules :: Vector{TR}
    preterminals :: Vector{C}
    startstate :: State{C, AbstractRule{C,C}}
    terminal_dict :: Dict{T, Vector{Tuple{C, TR}}}
    CRP :: Dict{C, ChineseRest{Fragment}}
    DM :: Dict{C, DirCat{CR, Float64}}
    BB :: Dict{Tuple{C, CR, C}, BetaBern{Bool, Int}}
end

show(io::IO, fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} =
    print("FragmentGrammar{$C, $(CR), $T, $(TR)}()")

show(io::IO, crp::ChineseRest) =
    print("ChineseRest($(crp.a), $(crp.b), $(crp.num_tables), $(crp.num_customers), $(crp.tables))")

category_type(fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} = C
terminal_type(fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} = T
category_rule_type(fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} = AbstractRule{C,C}
terminal_rule_type(fg::FragmentGrammar{C, CR, T, TR}) where {C, CR, T, TR} = AbstractRule{C,T}
categories(fg::FragmentGrammar) = fg.categories
terminals(fg::FragmentGrammar) = fg.terminals
preterminals(fg::FragmentGrammar) = fg.preterminals
startstate(fg::FragmentGrammar) = fg.startstate
startsymbols(fg::FragmentGrammar) = fg.startcategories
score_type(fg::FragmentGrammar) = LogProb
state_type(fg::FragmentGrammar{C}) where C = State{C,AbstractRule{C,C}} # typeof(startstate(fg))
completions(fg::FragmentGrammar{C}, state::State{C,AbstractRule{C,C}}) where C =
    ((cat, r, prob(fg, cat, r)) for (cat, r) in completions(state))
completions(fg::FragmentGrammar{C, CR, T, TR}, t::T) where {C, T, CR, TR} =
    ((cat, rule, prob(fg, cat, rule)) for (cat, rule) in fg.terminal_dict[t])
prob(fg::FragmentGrammar{C}, cat::C, rule::BaseRule{C,C}) where C = one(score_type(fg))
prob(fg::FragmentGrammar{C}, cat::C, rule::R) where {C, R<:ApproxRule{C,C}} = rule.prob

function update_approx_probs!(rule::ApproxRule{C,C}, fg::FragmentGrammar{C}) where C
    # Should really use broadcasting or something here...
    cat = lhs(rule)
    rule.prob = zero(LogProb)
    for (i,r) in enumerate(rule.rules)
        if r isa BaseRule
            new_p = logscore(fg.DM[cat], r) * new_table_logscore(fg.CRP[cat])
            rule.prob += new_p
            rule.probs[i] = new_p
        elseif r isa FragmentRule
            new_p = logscore(fg.CRP[cat], r.fragment, true)
            rule.prob += new_p
            rule.probs[i] = new_p
        else error("ApproxRule subrule is of type $(typeof(r)), should be BaseRule or FragmentRule") end
    end
end
update_approx_probs!(completion::Tuple{C,ApproxRule{C,C}}, fg::FragmentGrammar{C}) where C = update_approx_probs!(completion[2], fg)
update_approx_probs!(fg::FragmentGrammar) = for state in startstate(fg) update_approx_probs!.(state.comp, Ref(fg)) end

logscore(fg::FragmentGrammar) = reduce(*, logscore.(values(fg.DM))) * reduce(*, logscore.(values(fg.CRP))) * reduce(*, logscore.(values(fg.BB)))

"""
    FragmentGrammar(categories, startcategories, category_rules, terminals, terminal_rules[, a::Float64, b::Float64])
"""
function FragmentGrammar(cats::Vector{C}, starts::Vector{C}, cat_rules::Vector{CR}, terms::Vector{T}, term_rules::Vector{TR}, a=0.01, b=0.2, dm_pseudo=1.0, bb_alpha=1, bb_beta=1) where {C, CR<:BaseRule{C,C}, T, TR}
    # cat_rules = Vector{CR}(cat_rules)
    startstate = State(C, AbstractRule{C,C})# ApproxRule{C,C})
    for r in cat_rules
        add_rule!(startstate, r, lhs(r), rhs(r))
        # rm_rule!(startstate, r, r.lhs, r.rhs)
    end

    CRP = Dict{C,ChineseRest{Fragment}}()
    DM = Dict{C, DirCat{CR, Float64}}(cat => DirCat{CR,Float64}(Dict(x => dm_pseudo for x in CR[r for (i, r) in enumerate(cat_rules) if lhs(r) == cat]), dm_pseudo) for (j, cat) in enumerate(cats))
    BB = Dict{Tuple{C, CR, C}, BetaBern{Bool, Int}}((lhs(r), r, rhs) => BetaBern(bb_alpha, bb_beta) for r in cat_rules for rhs in rhs(r))
    preterminals = C[]
    terminal_dict = Dict{T, Vector{Tuple{C, TR}}}()
    for r in term_rules
        push!(preterminals, lhs(r))
        t = rhs(r)[1]
        if haskey(terminal_dict, t)
            push!(terminal_dict[t], (lhs(r), r))
        else
            terminal_dict[t] = [(lhs(r), r)]
        end
    end

    fg = FragmentGrammar{C, CR, T, TR}(cats, starts, cat_rules, terms, term_rules, preterminals, startstate, terminal_dict, CRP, DM, BB)

    fg.CRP = Dict{C,ChineseRest{Fragment}}(cat => ChineseRest(a, b, BaseDistribution(fg, cat)) for (i, cat) in enumerate(cats))

    return fg
end

function sample(fg :: FragmentGrammar{C,CR,T,TR}, cat :: C) where {C, CR, T, TR}
    crp_sample = sample(fg.CRP[cat])
    fragment, dm_counts, bb_counts, crp_counts = crp_sample isa Tuple ? crp_sample : (crp_sample, Tuple{C,CR}[], Pair{Tuple{C,CR,C}, Bool}[], Fragment[])
    children = Dict{Tree, Pointer}()
    for variable in fragment.variables
        ptr, dm_counts_child, bb_counts_child, crp_counts_child = sample(fg, variable.data)
        append!(dm_counts, dm_counts_child)
        append!(bb_counts, bb_counts_child)
        append!(crp_counts, crp_counts_child)
        push!(children, variable => ptr)
    end
    return Pointer(fragment, children), dm_counts, bb_counts, crp_counts
end

struct BaseDistribution{C} <: Distribution{Fragment}
    fg :: FragmentGrammar{C}
    category :: C
end

show(io::IO, bd::BaseDistribution{C}) where C = print("BaseDistribution{$C}()")

function sample(basedist :: BaseDistribution)
    C, CR, T, TR = category_type(basedist.fg), category_rule_type(basedist.fg), terminal_type(basedist.fg), terminal_rule_type(basedist.fg)

    variables = Tree{C}[]
    leaves = Tree{C}[]
    tree = TreeNode(basedist.category)

    dm_sample = sample(basedist.fg.DM[basedist.category])
    dm_counts = Tuple{C,CR}[(basedist.category, dm_sample)]
    # push!(dm_counts, (basedist.category, dm_sample))
    bb_counts = Pair{Tuple{C, CR, C}, Bool}[]
    crp_counts = Fragment[]

    r = dm_sample
    children = r(basedist.category) # or r.rhs or rhs(r)

    for child in children
        if child in basedist.fg.preterminals
            # Make a leaf node (non-variable)
            leaf_tree = TreeNode(child)
            insert_child!(tree, leaf_tree)
            push!(leaves, leaf_tree)
        else
            bbidx = (basedist.category, dm_sample, child)
            extend = sample(basedist.fg.BB[bbidx])
            push!(bb_counts, bbidx => extend)
            if extend # if we extend the fragment
                # Get (recursive) Pointer, DM counts, BB counts, and CRP counts (fragments)
                ptr, dm_counts_child, bb_counts_child, crp_counts_child = sample(basedist.fg, child)
                fragchild = clone(ptr.fragment) # ptr.fragment # We only need the fragment for the base distribution
                # Merge BB counts
                append!(bb_counts, bb_counts_child)
                # Merge DM counts
                append!(dm_counts, dm_counts_child)
                # Merge CRP counts (recursively counts only if we're *extending* a non-terminal with another fragment)
                # If it's a new fragment, it may contain a fragment in crp_counts_child. If it's already stored, then it contains nothing.
                append!(crp_counts, crp_counts_child)
                # Add to tree
                insert_child!(tree, fragchild.tree) # deepcopy(fragchild.tree))   # Better than modifying the underlying tree if taken from CRP preexisting fragment
                # Merge variables (non-terminal leaves) from recursive calls to FG
                append!(variables, fragchild.variables)
                append!(leaves, fragchild.leaves)
            else
                # Make a leaf node (variable)
                variable_tree = TreeNode(child)
                push!(variables, variable_tree)
                push!(leaves, variable_tree)
                insert_child!(tree, variable_tree)
            end
        end
    end

    frag = Fragment(tree, variables, leaves)
    push!(crp_counts, frag)

    return frag, dm_counts, bb_counts, crp_counts
end

# TODO: Write tests for add/rm_obs!
function add_obs!(fg :: FragmentGrammar, analysis :: Analysis)
    # Add DM counts
    for dm_obs in analysis.dm_obs
        add_obs!(fg.DM[dm_obs[1]], dm_obs[2])
    end
    # Add BB counts
    for bb_obs in analysis.bb_obs
        add_obs!(fg.BB[bb_obs[1]], bb_obs[2])
    end
    # Add fragments to CRP
    for frag in analysis.crp_obs
        add_obs!(fg.CRP[frag.tree.data], frag)
        fr = FragmentRule(frag)
        add_rule!(startstate(fg), fr, lhs(fr), rhs(fr))
    end
    for frag_ptr in analysis.pointer
        add_obs!(fg.CRP[frag_ptr.fragment.tree.data], frag_ptr.fragment)
        # Add completions to finite state machine (for approx. PCFG)
        fr = FragmentRule(frag_ptr.fragment)
        add_rule!(startstate(fg), fr, lhs(fr), rhs(fr))
    end
end

function rm_obs!(fg :: FragmentGrammar, analysis :: Analysis)
    # rm DM counts
    for dm_obs in analysis.dm_obs
        rm_obs!(fg.DM[dm_obs[1]], dm_obs[2])
    end
    # rm BB counts
    for bb_obs in analysis.bb_obs
        rm_obs!(fg.BB[bb_obs[1]], bb_obs[2])
    end
    # rm fragments from CRP
    for frag in analysis.crp_obs
        rm_obs!(fg.CRP[frag.tree.data], frag)
        fr = FragmentRule(frag)
        rm_rule!(startstate(fg), fr, lhs(fr), rhs(fr))
    end
    for frag_ptr in analysis.pointer
        rm_obs!(fg.CRP[frag_ptr.fragment.tree.data], frag_ptr.fragment)
        # rm completions to finite state machine (for approx. PCFG)
        fr = FragmentRule(frag_ptr.fragment)
        rm_rule!(startstate(fg), fr, lhs(fr), rhs(fr))
    end
end

# TODO: Refactor this method so it isn't spaghetti
function sample(parsetree :: Tree{Tuple{Cons,AbstractRule{C,C}}}, fg::FragmentGrammar{C,CR,T,TR}) where {Cons<:Constituent, C, CR, T, TR}
    cons, rule = parsetree.data

    local tree, frag
    variables = Tree{C}[]
    leaves = Tree{C}[]
    children = Dict{Tree, Pointer}()

    dm_obs = Tuple{C,CR}[]
    bb_obs = Pair{Tuple{C, CR, C}, Bool}[]
    crp_obs = Fragment[]

    if rule isa BaseRule{C,C} # If terminal rule, just return nothing and the caller will make a leaf from the preterminal.
        return nothing, dm_obs, bb_obs, crp_obs

    elseif rule isa ApproxRule{C,C} # If it's ApproxRule, we care about it, i.e. it's not a preterminal. We may still sample a BaseRule from this.
        sampled_rule = sample(rule)

        if sampled_rule isa BaseRule{C,C}
            # TODO: -Make a fragment of current contiguous tree fragment and bookkeep (crp_obs vector or something)
            #       -Flip a coin to decide if this fragment's tree will continue being the current contiguous tree fragment
            push!(dm_obs, (lhs(rule), sampled_rule))
            tree = TreeNode(lhs(rule))

            for child in parsetree.children
                child_cat = lhs(child.data[2])
                ptr_child, dm_obs_child, bb_obs_child, crp_obs_child = sample(child, fg)

                if ptr_child != nothing
                    append!(dm_obs, dm_obs_child)
                    append!(bb_obs, bb_obs_child)
                    append!(crp_obs, crp_obs_child)

                    bbidx = (lhs(rule), sampled_rule, child_cat)
                    extend = sample(fg.BB[bbidx])
                    push!(bb_obs, bbidx => extend)

                    if extend
                        ptr_child_clone = clone(ptr_child) # Keeps references to the same other Pointer objects from variable nodes, but creates a clone of the underlying Fragment tree and variables/leaves Vectors.
                        frag_child = ptr_child.fragment
                        tree_child = frag_child.tree
                        insert_child!(tree, tree_child)
                        append!(variables, frag_child.variables)
                        append!(leaves, frag_child.leaves)
                        merge!(children, ptr_child_clone.children)
                    else
                        variable_tree = TreeNode(child_cat)
                        insert_child!(tree, variable_tree)
                        push!(variables, variable_tree)
                        push!(leaves, variable_tree)
                        push!(children, variable_tree => ptr_child)
                    end
                else
                    # TODO: We're at a terminal rule child and are making a new fragment, so add preterminal tree node as leaf of fragment.
                    tree_child = TreeNode(child_cat)
                    insert_child!(tree, tree_child)
                    push!(leaves, tree_child)
                end
            end

            frag = Fragment(tree, variables, leaves)
            push!(crp_obs, frag)
            return Pointer(frag, children), dm_obs, bb_obs, crp_obs

        elseif sampled_rule isa FragmentRule{C,C}
            push!(crp_obs, sampled_rule.fragment)
            for (i, child) in enumerate(parsetree.children)
                leaf_node = sampled_rule.fragment.leaves[i]
                if leaf_node in sampled_rule.fragment.variables # For each of these recurse. If return != nothing, do the logical stuff, etc.
                    ptr_child, dm_obs_child, bb_obs_child, crp_obs_child = sample(child, fg)
                    append!(dm_obs, dm_obs_child)
                    append!(bb_obs, bb_obs_child)
                    append!(crp_obs, crp_obs_child)
                    push!(children, leaf_node => ptr_child)
                end
            end
            return Pointer(sampled_rule.fragment, children), dm_obs, bb_obs, crp_obs
        end
    end
end

approx_rule(r::Tuple{Cons,AbstractRule}) where Cons <: Constituent = r[2]

function logscore(approx_tree::Tree{Tuple{Cons,AbstractRule{C,C}}}, fg::FragmentGrammar{C,CR,T,TR}) where {Cons<:Constituent, C, CR, T, TR}
    reduce(*, approxprob.(approx_rule.(data.(approx_tree))))
end

run_chartparser(input::Tree, fg::FragmentGrammar) = run_chartparser(leaf_data(input), fg, boolean_dependency_dict(input))

function boolean_dependency_dict(tree::Tree{C}) where C
    dependency_dict = Dict{Tuple{C,Int,Int},Bool}()
    function helper(tree, dependency_dict::Dict{Tuple{C,Int,Int}, Bool}, dot::Ref{Int})
        start = dot[]
        if isempty(tree.children)
            dot[] += 1
        else
            for child in tree.children
                helper(child, dependency_dict, dot)
            end
        end
        push!(dependency_dict, (tree.data, start, dot[])=>true)
    end
    helper(tree, dependency_dict, Ref(1))
    # @show dependency_dict
    return dependency_dict
end

function run_mcmc!(fg::FragmentGrammar, inputs::Vector{Tree}, sweeps=1)
    update_approx_probs!(fg)
    analyses = init_mcmc(fg, inputs)
    add_obs!.([fg], getindex.(values(analyses),1))
    update_approx_probs!(fg)

    # local curr_prob::LogProb

    for sweep in 1:sweeps
        for input_tree in keys(analyses)
            curr_prob = logscore(fg)
            curr_analysis, curr_approx_tree = analyses[input_tree]
            rm_obs!(fg, curr_analysis)

            update_approx_probs!(fg) # F-f(i), needed for logscore(::Tree{Tuple{Cons,AbstractRule{C,C}}}, fg::FragmentGrammar) as well.

            proposal_approx_tree = sample_tree(run_chartparser(input_tree, fg))
            proposal_analysis = Analysis(sample(proposal_approx_tree, fg)...)
            forward_trans_prob = logscore(proposal_approx_tree, fg)
            reverse_trans_prob = logscore(curr_approx_tree, fg)

            add_obs!(fg, proposal_analysis) # No need to update approx probs, not going to use them here.

            new_prob = logscore(fg)

            accept_prob = min(one(LogProb), new_prob/curr_prob * reverse_trans_prob/forward_trans_prob)
            if flip(float(accept_prob)) == 0 # If we reject.
                @show "Reject", accept_prob
                rm_obs!(fg, proposal_analysis)
                add_obs!(fg, curr_analysis)
                @show curr_prob.log
            else # If we accept.
                @show "Accept", accept_prob
                analyses[input_tree] = (proposal_analysis, proposal_approx_tree)
                @show new_prob.log
            end
            update_approx_probs!(fg)
        end
    end
    return analyses
end

function init_mcmc(fg::FragmentGrammar, inputs::Vector{Tree})
    analyses = Dict{Tree,Tuple{Analysis,Tree}}()
    for input_tree in inputs
        analysis = Analysis(initialize_fragments(input_tree, fg)...)
        approx_tree = sample_tree(run_chartparser(input_tree, fg)) # This should always be BaseRules
        analyses[input_tree] = (analysis, approx_tree)
    end
    return analyses
end

function initialize_fragments(tree::Tree, fg::FragmentGrammar{C,CR}) where {C,CR}
    variables = Tree{C}[]
    leaves = Tree{C}[]
    children = Dict{Tree, Pointer}()

    dm_obs = Tuple{C,CR}[]
    bb_obs = Pair{Tuple{C, CR, C}, Bool}[]
    crp_obs = Fragment[]

    l = tree.data

    working_tree = TreeNode(l)

    # FIND BASERULE USED
    st = startstate(fg)
    for child in tree.children
        # I can assume that the FG only has base rules in its states right now.
        st = transition(st, child.data)
        insert_child!(working_tree, TreeNode(child.data))
    end
    ar = completions(st)[1][2] # Completion is Tuple{C, CR}
    # ar had sure as hell better be a BaseRule at this point.
    r = ar.rules[1] # There should seriously only be one.
    push!(dm_obs, (lhs(r), r))

    # RECURSE FOR FRAGMENTS, CHECK IF TERMINAL, ETC.
    for (i, child) in enumerate(working_tree.children)
        variable_tree = child
        push!(leaves, variable_tree)

        if !(child.data in preterminals(fg))
            push!(variables, variable_tree)
            bbidx = (lhs(r), r, child.data)
            push!(bb_obs, bbidx => false)

            ptr_child, dm_obs_child, bb_obs_child, crp_obs_child = initialize_fragments(tree.children[i], fg)
            append!(dm_obs, dm_obs_child)
            append!(bb_obs, bb_obs_child)
            append!(crp_obs, crp_obs_child)

            push!(children, variable_tree => ptr_child)
        end
    end

    frag = Fragment(working_tree, variables, leaves)
    push!(crp_obs, frag)

    return Pointer(frag, children), dm_obs, bb_obs, crp_obs
end

end


text:
"""
    coordinates(geometry)
Returns the edges/vertices/coordinates of a geometry. Is allowed to return lazy iterators!
Use `decompose(ConcretePointType, geometry)` to get `Vector{ConcretePointType}` with
`ConcretePointType` to be something like `Point{3, Float32}`.
"""
function coordinates(points::AbstractVector{<:AbstractPoint})
    return points
end

"""
    faces(geometry)
Returns the face connections of a geometry. Is allowed to return lazy iterators!
Use `decompose(ConcreteFaceType, geometry)` to get `Vector{ConcreteFaceType}` with
`ConcreteFaceType` to be something like `TriangleFace{Int}`.
"""
function faces(f::AbstractVector{<:AbstractFace})
    return f
end

function normals(primitive, nvertices=nothing)
    # doesn't have any specific algorithm to generate normals
    # so will be generated from faces + positions
    # which we indicate by returning nothing!
    # Overload normals(primitive::YourPrimitive), to calcalute the normals
    # differently
    return nothing
end

function faces(primitive, nvertices=nothing)
    # doesn't have any specific algorithm to generate faces
    # so will try to triangulate the coordinates!
    return nothing
end

"""
    Tesselation(primitive, nvertices)
For abstract geometries, when we generate
a mesh from them, we need to decide how fine grained we want to mesh them.
To transport this information to the various decompose methods, you can wrap it
in the Tesselation object e.g. like this:

```julia
sphere = Sphere(Point3f0(0), 1)
m1 = mesh(sphere) # uses a default value for tesselation
m2 = mesh(Tesselation(sphere, 64)) # uses 64 for tesselation
length(coordinates(m1)) != length(coordinates(m2))
```
For grid based tesselation, you can also use a tuple:
```julia
rect = Rect2D(0, 0, 1, 1)
Tesselation(rect, (5, 5))
"""
struct Tesselation{Dim, T, Primitive, NGrid}
    primitive::Primitive
    nvertices::NTuple{NGrid, Int}
end

function Tesselation(primitive::GeometryPrimitive{Dim, T}, nvertices::NTuple{N, <:Integer}) where {Dim, T, N}
    return Tesselation{Dim, T, typeof(primitive), N}(primitive, Int.(nvertices))
end

Tesselation(primitive, nvertices::Integer) = Tesselation(primitive, (nvertices,))

# This is a bit lazy, I guess we should just refactor these methods
# to directly work on Tesselation - but this way it's backward compatible and less
# refactor work :D
nvertices(tesselation::Tesselation) = tesselation.nvertices
nvertices(tesselation::Tesselation{T, N, P, 1}) where {T, N, P} = tesselation.nvertices[1]

coordinates(tesselation::Tesselation) = coordinates(tesselation.primitive, nvertices(tesselation))
faces(tesselation::Tesselation) = faces(tesselation.primitive, nvertices(tesselation))
normals(tesselation::Tesselation) = normals(tesselation.primitive, nvertices(tesselation))
texturecoordinates(tesselation::Tesselation) = texturecoordinates(tesselation.primitive, nvertices(tesselation))

## Decompose methods
# Dispatch type to make `decompose(UV{Vec2f0}, priomitive)` work
# and to pass through tesselation information

# Types that can be converted to a mesh via the functions below
const Meshable{Dim, T} = Union{Tesselation{Dim, T}, Mesh{Dim, T},
                               GeometryPrimitive{Dim, T}, AbstractVector{<: AbstractPoint{Dim, T}}}

struct UV{T} end
UV(::Type{T}) where T = UV{T}()
struct UVW{T} end
UVW(::Type{T}) where T = UVW{T}()
struct Normal{T} end
Normal(::Type{T}) where T = Normal{T}()

function decompose(::Type{F}, primitive) where {F<:AbstractFace}
    f = faces(primitive)
    f === nothing && return nothing
    return collect_with_eltype(F, f)
end

function decompose(::Type{P}, primitive) where {P<:AbstractPoint}
    return collect_with_eltype(P, metafree(coordinates(primitive)))
end

function decompose(::Type{Point}, primitive::Meshable{Dim, T}) where {Dim, T}
    return collect_with_eltype(Point{Dim, T}, metafree(coordinates(primitive)))
end

function decompose(::Type{T}, primitive) where {T}
    return collect_with_eltype(T, primitive)
end

decompose_uv(primitive) = decompose(UV(Vec2f0), primitive)
decompose_uvw(primitive) = decompose(UVW(Vec3f0), primitive)
decompose_normals(primitive) = decompose(Normal(Vec3f0), primitive)

function decompose(::Normal{T}, primitive) where T
    n = normals(primitive)
    n === nothing && return nothing
    return collect_with_eltype(T, n)
end

function decompose(::Union{UV{T}, UVW{T}}, primitive) where T
    return collect_with_eltype(T, texturecoordinates(primitive))
end

ext: jl


text:
# Jonas August, 2/2022

#=
This problem was asked by Microsoft.

A number is considered perfect if its digits sum up to exactly 10.

Given a positive integer n, return the n-th perfect number.

For example, given 1, you should return 19. Given 2, you should return 28.
=#

" Returns n-th perfect number. "
function perfect10(n::Int)::Int
    @assert n > 0
    i = 0   # number being checked
    m = 0   # current count of perfect numbers ≤ i CHECK
    while m < n
        i += 1
        # Convert number to digits by converting to string and then converting each digit to Int individually.
        s = 0  # sum of the converted digits
        for digit in string(i)
            s += parse(Int, digit)
        end
        if s == 10
            m += 1
        end
    end
    i
end

using Test
@test perfect10(1) == 19
@test perfect10(2) == 28

for n in 1:100
    @show n, perfect10(n)
end

max_stars_repo_path: Perfect number - Problem 70.jl

max_issues_repo_path: Perfect number - Problem 70.jl

max_forks_repo_path: Perfect number - Problem 70.jl


text:
using DataFrames

export Mack, claims2cum, cum2claims, cum2futureclaims
"""
A reserve triangle with chain ladder information
"""
type Mack
  "Cummulative claims (upper triangle)"
  cum::Matrix{Real}
  "Time vector of development factors"
  f::Vector{Real}
  "Time vector of future claims payments"
  futureclaims::Vector{Real}
  "Vector of reseves"
  res::Vector{Real}
  """
  Total reserves: ``\mathit{tot\_res} =\sum_{i=1}^I \mathit{res}_i``
  """
  tot_res::Real
  "Vector of mean square errors"
  mse::Vector{Real}
  "Total mean square error"
  tot_mse::Real
end

function Mack(df_triang::DataFrame; cum::Bool = false)
Mack(convert(DataArray{Real, 2}, df_triang); cum = cum)
end

function Mack(da_triang::DataArray{Real, 2}; cum::Bool = false)
  Mack(convert(Array{Real,2}, da_triang, 0); cum = cum)
end

function Mack(triang::Array{Real,2}; cum::Bool = false)
  if cum
    C = deepcopy(triang)
  else
    C = claims2cum(triang)
  end
  # Number of accident / development years
  I = size(C)[1]
  # Development factors
  f = [sum(C[1:(I-k),k+1])/sum(C[1:(I-k),k]) for k in 1:(I-1)]
  # Ultimate claim amounts
  for k  in 2:I
    for l in (I+1-k+1):I
      C[k,l] = C[k,I+1-k] * prod(f[(I+1-k):(l-1)])
    end
  end
  futclaims = cum2futureclaims(C)
  # Reserves
  R = [C[k,I]-C[k,I+1-k] for k in 1:I]
  σ² =
    [1/(I-k-1) *
      C[1:(I-k),k] ⋅ (C[1:(I-k),k+1] ./ C[1:(I-k),k]-f[k]).^2
    for k in 1:(I-1)]
  σ²[I-1] = min(σ²[I-2]^2 / σ²[I-3], min(σ²[I-3],σ²[I-2]))
  # mean square error
  mse = zeros(Real, I)
  for i in 2:I
    mse[i] = 0
    for k in (I+1-i):(I-1)
      mse[i] += σ²[k]/f[k]^2 * (1/C[i,k] + 1/ sum(C[1:I-k,k]))
    end
    mse[i] *= C[i,I]^2
  end
  # total mean square error
  mse_total = 0.0
  for i in 2:I
    tmp = 0.0
    for k in collect(I+1-i : I-1)
      tmp += (2σ²[k] / f[k]^2) / sum(C[1 : I-k, k])
    end
    tmp *= C[i,I] * sum(C[i+1 : I, I])
    tmp += mse[i]
    mse_total += tmp
  end
  return Mack(C, f, futclaims, R, sum(R), mse, mse_total)
end

"""
    claims2cum(c::Array{Real,2})

    Convert a quadratic upper triangular reserve matrix to
    the correponding cumulative triangular reserve matrix.
    The (strictly) lower triangular part is set to 0
"""
function claims2cum(c::Array{Real,2})
  if size(c,1) ≠ size(c,2)
    error("claims2cum: Matrix is not quadratic")
  end
  cum = zeros(Real,size(c))
  for i in 1:size(c, 1)
    cum[i,1:size(c, 1)-i+1] = cumsum(c[i,1:size(c, 1)-i+1], 2)
  end
  cum
end

"""
    cum2claims(c::Array{Real,2})

    Convert a quadratic cumulative claims matrix into a
    claims matrix.
"""
function cum2claims(c::Array{Real,2})
  I = size(c,2)
  if size(c,1) ≠ I
    error("cum2claims: Matrix is not quadratic")
  end
  claims = zeros(Real,size(c))
  for i in 1:I
    claims[i,1] = c[i,1]
    for k in 2:I
     claims[i,k] = c[i,k]-c[i,k-1]
    end
  end
  claims
end

"""
    cum2futclaims(c::Array{Real,2})

    Calculate the future claims vector from the quadratic,
    (strictly) lower triangle of a projected cumulative claims
    matrix
"""
function cum2futureclaims(c::Array{Real,2})
  I = size(c,2)
  if size(c,1) ≠ I
    error("cumclaims: Matrix is not quadratic")
  end
  claims = cum2claims(c)
  futureclaims = zeros(Real,I-1)
  for k in 2:I
    futureclaims[k-1] = sum(claims[I-k+2:I,k])
  end
  futureclaims
end


text:
#nb # %% A slide [markdown] {"slideshow": {"slide_type": "Slide"}}
# ## Meanderline-Based CPSS
# A "CPSS" is a circular polarization selective structure, i.e., a structure that reacts differently
# to the two senses of circular polarization.
# We'll first look at analyzing a design presented in the literature, and then proceed to optimize another
# design using PSSFSS as the analysis engine inside the optimization objective function.
# ### Sjöberg and Ericsson Design
# This example comes from the paper D. Sjöberg and A. Ericsson, "A multi layer meander line circular
# polarization selective structure (MLML-CPSS)," The 8th European Conference on Antennas and Propagation
# (EuCAP 2014), 2014, pp. 464-468, doi: 10.1109/EuCAP.2014.6901792.
#
# The authors describe an ingenious structure consisting of 5 progressively rotated meanderline sheets, which
# acts as a circular polarization selective surface: it passes LHCP (almost) without attenuation or
# reflection, and reflects RHCP (without changing its sense!) almost without attenuation or transmission.
#

# Here is the script that analyzes their design:

using PSSFSS
## Define convenience functions for sheets:
outer(rot) = meander(a=3.97, b=3.97, w1=0.13, w2=0.13, h=2.53+0.13, units=mm, ntri=600, rot=rot)
inner(rot) = meander(a=3.97*√2, b=3.97/√2, w1=0.1, w2=0.1, h=0.14+0.1, units=mm, ntri=600, rot=rot)
center(rot) = meander(a=3.97, b=3.97, w1=0.34, w2=0.34, h=2.51+0.34, units=mm, ntri=600, rot=rot)
## Note our definition of `h` differs from that of the reference by the width of the strip.
t1 = 4mm # Outer layers thickness
t2 = 2.45mm # Inner layers thickness
substrate = Layer(width=0.1mm, epsr=2.6)
foam(w) = Layer(width=w, epsr=1.05) # Foam layer convenience function
rot0 = 0 # rotation of first sheet
strata = [
    Layer()
    outer(rot0)
    substrate
    foam(t1)
    inner(rot0 - 45)
    substrate
    foam(t2)
    center(rot0 - 2*45)
    substrate
    foam(t2)
    inner(rot0 - 3*45)
    substrate
    foam(t1)
    outer(rot0 - 4*45)
    substrate
    Layer() ]
steering = (θ=0, ϕ=0)
flist = 10:0.1:20
##
results = analyze(strata, flist, steering, showprogress=false,
                  resultfile=devnull, logfile=devnull);

# Here are plots of the five meanderline sheets:

using Plots
plot(outer(rot0), unitcell=true, title="Sheet1")
#-
plot(inner(rot0-45), unitcell=true, title="Sheet2")
#-
plot(center(rot0-2*45), unitcell=true, title="Sheet3 (Center)")
#-
plot(inner(rot0-3*45), unitcell=true, title="Sheet4")
#-
plot(outer(rot0-4*45), unitcell=true, title="Sheet5")

# Notice that not only are the meanders rotated, but so too are the unit cell rectangles.
# This is because we used the generic `rot` keyword argument that rotates the entire unit
# cell and its contents. `rot` can be used for any FSS or PSS element type.  As a consequence of
# the different rotations applied to each unit cell, interactions between sheets due to higher-order
# modes cannot be accounted for; only the dominant ``m=n=0`` TE and TM modes are used in cascading
# the individual sheet scattering matrices.  This approximation is adequate for sheets that are
# sufficiently separated.  We can see from the log file (saved from a previous run where it was not
# disabled) that only 2 modes are used to model the interactions between sheets:
#
# ```
# Starting PSSFSS analysis on 2021-05-09 at 04:16:07.871
#
#
# ******************* Warning ***********************
#    Unequal unit cells in sheets 1 and 2
#    Setting #modes in dividing layer 3 to 2
# ******************* Warning ***********************
#
#
# ******************* Warning ***********************
#    Unequal unit cells in sheets 2 and 3
#    Setting #modes in dividing layer 5 to 2
# ******************* Warning ***********************
#
#
# ******************* Warning ***********************
#    Unequal unit cells in sheets 3 and 4
#    Setting #modes in dividing layer 7 to 2
# ******************* Warning ***********************
#
#
# ******************* Warning ***********************
#    Unequal unit cells in sheets 4 and 5
#    Setting #modes in dividing layer 9 to 2
# ******************* Warning ***********************
#
#
# Dielectric layer information...
#
#  Layer  Width  units  epsr   tandel   mur  mtandel modes  beta1x  beta1y  beta2x  beta2y
#  ----- ------------- ------- ------ ------- ------ ----- ------- ------- ------- -------
#      1    0.0000  mm    1.00 0.0000    1.00 0.0000     2  1582.7     0.0     0.0  1582.7
#  ==================  Sheet   1  ========================  1582.7     0.0     0.0  1582.7
#      2    0.1000  mm    2.60 0.0000    1.00 0.0000     0     0.0     0.0     0.0     0.0
#      3    4.0000  mm    1.05 0.0000    1.00 0.0000     2  1582.7     0.0     0.0  1582.7
#  ==================  Sheet   2  ========================   791.3  -791.3  1582.7  1582.7
#      4    0.1000  mm    2.60 0.0000    1.00 0.0000     0     0.0     0.0     0.0     0.0
#      5    2.4500  mm    1.05 0.0000    1.00 0.0000     2   791.3  -791.3  1582.7  1582.7
#  ==================  Sheet   3  ========================     0.0 -1582.7  1582.7     0.0
#      6    0.1000  mm    2.60 0.0000    1.00 0.0000     0     0.0     0.0     0.0     0.0
#      7    2.4500  mm    1.05 0.0000    1.00 0.0000     2     0.0 -1582.7  1582.7     0.0
#  ==================  Sheet   4  ========================  -791.3  -791.3  1582.7 -1582.7
#      8    0.1000  mm    2.60 0.0000    1.00 0.0000     0     0.0     0.0     0.0     0.0
#      9    4.0000  mm    1.05 0.0000    1.00 0.0000     2  -791.3  -791.3  1582.7 -1582.7
#  ==================  Sheet   5  ======================== -1582.7    -0.0     0.0 -1582.7
#     10    0.1000  mm    2.60 0.0000    1.00 0.0000     0     0.0     0.0     0.0     0.0
#     11    0.0000  mm    1.00 0.0000    1.00 0.0000     2 -1582.7    -0.0     0.0 -1582.7
# ...
# ```
#
# Note that PSSFSS prints warnings to the log file where it is forced to set the number of layer
# modes to 2 because of unequal unit cells.  Also, in the dielectric layer list it can be seen
# that these layers are assigned 2 modes each.  The thin layers adjacent to sheets are assigned 0
# modes because these sheets are incorporated into so-called "GSM blocks" or "Gblocks" wherein
# the presence of the thin layer is accounted for using the stratified medium Green's functions.

# Here is the script that compares PSSFSS predicted performance with very
# high accuracy predictions from CST and COMSOL that were digitized from figures in the paper.

using Plots, DelimitedFiles
RL11rr = -extract_result(results, @outputs s11db(r,r))
AR11r = extract_result(results, @outputs ar11db(r))
IL21L = -extract_result(results, @outputs s21db(L,L))
AR21L = extract_result(results, @outputs ar21db(L))

default(lw=2, xlim=(10,20), xtick=10:20, ylim=(0,3), ytick=0:0.5:3, gridalpha=0.3,
        framestyle=:box)
p = plot(flist,RL11rr,title="RHCP → RHCP Return Loss",
         xlabel="Frequency (GHz)", ylabel="Return Loss (dB)", label="PSSFSS")
cst = readdlm("../src/assets/cpss_cst_fine_digitized_rl.csv", ',')
plot!(p, cst[:,1], cst[:,2], label="CST")
comsol = readdlm("../src/assets/cpss_comsol_fine_digitized_rl.csv", ',')
plot!(p, comsol[:,1], comsol[:,2], label="COMSOL")
#md savefig("cpssa1.png"); nothing  # hide
#-
#md # ![](cpssa1.png)
#-
p = plot(flist,AR11r,title="RHCP → RHCP Reflected Axial Ratio",
         xlabel="Frequency (GHz)", ylabel="Axial Ratio (dB)", label="PSSFSS")
cst = readdlm("../src/assets/cpss_cst_fine_digitized_ar_reflected.csv", ',')
plot!(p, cst[:,1], cst[:,2], label="CST")
comsol = readdlm("../src/assets/cpss_comsol_fine_digitized_ar_reflected.csv", ',')
plot!(p, comsol[:,1], comsol[:,2], label="COMSOL")
#md savefig("cpssa2.png"); nothing  # hide
#-
#md # ![](cpssa2.png)
#-
p = plot(flist,IL21L,title="LHCP → LHCP Insertion Loss",
         xlabel="Frequency (GHz)", ylabel="Insertion Loss (dB)", label="PSSFSS")
cst = readdlm("../src/assets/cpss_cst_fine_digitized_il.csv", ',')
plot!(p, cst[:,1], cst[:,2], label="CST")
comsol = readdlm("../src/assets/cpss_comsol_fine_digitized_il.csv", ',')
plot!(p, comsol[:,1], comsol[:,2], label="COMSOL")
#md savefig("cpssa3.png"); nothing  # hide
#-
#md # ![](cpssa3.png)
#-
p = plot(flist,AR21L,title="LHCP → LHCP Transmitted Axial Ratio",
         xlabel="Frequency (GHz)", ylabel="Axial Ratio (dB)", label="PSSFSS")
cst = readdlm("../src/assets/cpss_cst_fine_digitized_ar_transmitted.csv", ',')
plot!(p, cst[:,1], cst[:,2], label="CST")
comsol = readdlm("../src/assets/cpss_comsol_fine_digitized_ar_transmitted.csv", ',')
plot!(p, comsol[:,1], comsol[:,2], label="COMSOL")
#md savefig("cpssa4.png"); nothing  # hide
#-
#md # ![](cpssa4.png)

# The PSSFSS results generally track well with the high-accuracy solutions, but are less accurate
# especially at the high end of the band, presumably because cascading is performed in PSSFSS
# for this structure using only the two principal Floquet modes.  As previosly discussed,
# this is necessary because the
# rotated meanderlines are achieved by rotating the entire unit cell, and the unit cell for sheets
# 2 and 4 are not square.  Since the periodicity of the sheets in the structure varies from sheet
# to sheet, higher order Floquet modes common to neighboring sheets cannot be defined, so we are forced
# to use only the dominant (0,0) modes which are independent of the periodicity.  This limitation is removed
# in a later example.
# Meanwhile, it is of interest to note that their high-accuracy runs
# required 10 hours for CST and 19 hours for COMSOL on large engineering workstations.  The PSSFSS run
# took about 60 seconds on my desktop machine.

# ### Design Based on PSSFSS Optimization with CMAES
# Here we use PSSFSS in conjunction with the CMAES optimizer from the
# [CMAEvolutionStrategy](https://github.com/jbrea/CMAEvolutionStrategy.jl) package.  I've used CMAES
# in the past with good success on some tough optimization problems.  Here is the code that defines
# the objective function:

# ```julia
# using PSSFSS
# using Dates: now
#
# let bestf = typemax(Float64)
#     function objective(x)
#         period, wo, ho, wi, hi, wc, hc, t1, t2 = x
#         ao = bo = ai = bi = ac = bc = period
#         ai *= √2
#         bi /= √2
#         # Ensure physically realizable or return large value:
#         (bo > ho > 2.1*wo && bi > hi > 2.1*wi && bc > hc > 2.1*wc) || (return 5000.0)
#
#         outer(rot) = meander(a=ao, b=bo, w1=wo, w2=wo, h=ho, units=mm, ntri=400, rot=rot)
#         inner(rot) = meander(a=ai, b=bi, w1=wi, w2=wi, h=hi, units=mm, ntri=400, rot=rot)
#         center(rot) = meander(a=ac, b=bc, w1=wc, w2=wc, h=hc, units=mm, ntri=400, rot=rot)
#
#         substrate = Layer(width=0.1mm, epsr=2.6)
#         foam(w) = Layer(width=w, epsr=1.05)
#         rot0 = 0
#
#         strata = [
#                 Layer()
#                 outer(rot0)
#                 substrate
#                 foam(t1*1mm)
#                 inner(rot0 - 45)
#                 substrate
#                 foam(t2*1mm)
#                 center(rot0 - 2*45)
#                 substrate
#                 foam(t2*1mm)
#                 inner(rot0 - 3*45)
#                 substrate
#                 foam(t1*1mm)
#                 outer(rot0 - 4*45)
#                 substrate
#                 Layer() ]
#         steering = (θ=0, ϕ=0)
#         flist = 11:0.25:19
#         results = analyze(strata, flist, steering, showprogress=false)
#         s11rr, s21ll, ar11db, ar21db = eachcol(extract_result(results,
#                        @outputs s11db(R,R) s21db(L,L) ar11db(R) ar21db(L)))
#         RL = -s11rr
#         IL = -s21ll
#         obj = maximum(vcat(RL,IL,ar11db,ar21db))
#         if obj < bestf
#             bestf = obj
#             open("optimization_best.log", "a") do fid
#                 xround = round.(x, digits=4)
#                 println(fid, round(obj,digits=4), " at x = ", xround, "  #", now())
#             end
#         end
#         return obj
#     end
# end
# ```

# We optimize at 33 frequencies between 11 and 19 GHz.  The actual frequency range of interest is
# 12 to 18 GHz; the wider optimization band provides some safety margin.  Each objective function
# evaluation takes about 24  seconds on my machine.
# As you can see from the code above, each successive sheet in the structure is rotated an additional
# 45 degrees relative to its predecessor.
# The objective is defined as the largest value of RHCP reflected return loss, LHCP insertion loss, or
# reflected or transmitted axial ratio that
# occurs at any of the analysis frequencies (i.e. we are setting up for "minimax" optimization). Also,
# the `let` block allows the objective function to maintain persistent state in the
# variable `bestf` which is initialized to the largest 64-bit floating point value. Each time a set
# of inputs results in a lower objective function value, `bestf` is updated with this value and
# the inputs and objective function value are
# written to the file "optimization_best.log", along with a date/time stamp.  This allows the user
# to monitor the optimization and to terminate the it prematurely, if desired, without losing the
# best result achieved so far.

# Here is the code for running the optimization:

# ```julia
# using CMAEvolutionStrategy
# #  x = [period,  wo,  ho,  wi,   hi,  wc,   hc,  t1,  t2]
# xmin = [3.0,     0.1, 0.1, 0.1,  0.1, 0.1,  0.1, 1.5, 1.5]
# xmax = [5.5,     0.35,4.0, 0.35, 4.0, 0.35, 4.0, 6.0, 6.0]
# x0 = 0.5 * (xmin + xmax)
# popsize = 2*(4 + floor(Int, 3*log(length(x0))))
# isfile("optimization_best.log") && rm("optimization_best.log")
# opt = minimize(objective, x0, 1.0;
#            lower = xmin,
#            upper = xmax,
#            maxfevals = 9000,
#            xtol = 1e-4,
#            ftol = 1e-5,
#            popsize=popsize)
# ```

# Note that I set the population size to twice the normal default value.  Based
# on previous experience, using 2 to 3 times the default population size helps the
# optimizer to do better on tough objective functions like the present one.
# I let the optimizer run for 6 hours, during which time it reduced the objective function
# value from 11.88 dB to 0.86 dB.  It was then interrupted due to a file system error.  I
# restarted it after setting the starting value to the current best and reducing the
# "sigma" value (the third argument to `minimize`, which controls the algorithms
#  exploratory tendency) to 0.2 (slightly greater than the value it had achieved during the aborted run)
# from its default value of 1.  After about 13 hours the program terminated normally due to
# insufficient changes in the `x` variable.  The final value of objective function was 0.65 dB.

# Here is a look at the final portion of the file "optimization_best.log":
# ```
# 0.6535 at x = [3.0968, 0.1025, 2.1601, 0.1003, 0.95, 0.3377, 2.3584, 4.3813, 2.2974]  #2021-05-09T23:02:40.230
# 0.6533 at x = [3.0991, 0.1028, 2.162, 0.1002, 0.944, 0.3372, 2.3562, 4.3801, 2.3007]  #2021-05-09T23:27:38.132
# 0.6532 at x = [3.0985, 0.1029, 2.1652, 0.1003, 0.9414, 0.337, 2.3547, 4.3766, 2.3036]  #2021-05-09T23:46:19.068
# 0.6531 at x = [3.0998, 0.1028, 2.164, 0.1001, 0.9443, 0.3378, 2.3558, 4.3783, 2.3028]  #2021-05-09T23:49:13.202
# 0.6529 at x = [3.0988, 0.1028, 2.1652, 0.1002, 0.9422, 0.3372, 2.3545, 4.3765, 2.3039]  #2021-05-09T23:50:58.064
# 0.6529 at x = [3.0985, 0.1029, 2.1687, 0.1002, 0.9407, 0.3368, 2.3526, 4.3714, 2.3073]  #2021-05-10T00:11:36.314
# 0.6529 at x = [3.0976, 0.1028, 2.1688, 0.1004, 0.9443, 0.3374, 2.3532, 4.3716, 2.3062]  #2021-05-10T00:16:16.205
# 0.6528 at x = [3.0991, 0.1028, 2.1679, 0.1002, 0.9444, 0.3376, 2.3539, 4.3735, 2.3056]  #2021-05-10T00:20:04.181
# 0.6527 at x = [3.0987, 0.1027, 2.1664, 0.1001, 0.9443, 0.3375, 2.3544, 4.3751, 2.3045]  #2021-05-10T00:33:08.779
# 0.6527 at x = [3.0972, 0.1029, 2.1692, 0.1002, 0.9362, 0.336, 2.3512, 4.3707, 2.3083]  #2021-05-10T00:39:14.306
# 0.6527 at x = [3.0985, 0.1028, 2.1663, 0.1001, 0.9426, 0.3372, 2.3539, 4.3749, 2.3049]  #2021-05-10T00:40:59.223
# 0.6527 at x = [3.0974, 0.1028, 2.1688, 0.1002, 0.9424, 0.3371, 2.3527, 4.3718, 2.3069]  #2021-05-10T00:51:08.383
# 0.6526 at x = [3.0982, 0.1028, 2.1678, 0.1001, 0.9422, 0.3371, 2.3531, 4.373, 2.3061]  #2021-05-10T00:56:58.562
# 0.6526 at x = [3.0979, 0.1028, 2.1682, 0.1001, 0.9406, 0.3369, 2.3529, 4.3725, 2.3066]  #2021-05-10T00:59:51.875
# 0.6525 at x = [3.0974, 0.1028, 2.1675, 0.1001, 0.9411, 0.3369, 2.353, 4.3735, 2.306]  #2021-05-10T01:03:39.402
# ```

# The performance of this design is shown below:

# ![](./assets/cpss_cmaesopt_ar_refl.png)

# ![](./assets/cpss_cmaesopt_ar_trans.png)

# ![](./assets/cpss_cmaesopt_il_trans.png)

# ![](./assets/cpss_cmaesopt_rl_refl.png)

# It would probably be possible to do improve the performance somewhat over the 12-18 GHz band by weighting the
# various contributions to the objective function and/or tapering the
# requirements of the objective function at the band edges.  Also, in a serious design effort, several additional
# runs of the optimizer should be attempted, since results vary for stochastic algorithms like CMAES.


text:
function cfun2(a, A)
    cf = mean(abs.(a-A))
    return cf
end

function  pfun(R)
    Rs = copy(R)
    ts = sample(1:length(R), 2)
    Rs[reverse(ts)] = R[ts]
    return Rs
end

function efun6(a,b,c,b2,R,nlag=50,n2lag=25)
    A = autocor(R, 1:nlag)
    B = autocor(abs.(R), 1:n2lag)
    C = crosscor(R, abs.(R), -nlag:nlag)
    B2 = autocor(R.*R, 1:n2lag)
    Z = cfun2(a,A)+cfun2(b,B)+cfun2(c,C)+cfun2(b2,B2)
    return Z
end

max_stars_repo_name: silvaac/SMCJulia

max_issues_repo_name: silvaac/SMCJulia

max_forks_repo_name: silvaac/SMCJulia


text:
@testset "NegativeBinomialLikelihood" begin
    for args in ((), (logistic,), (LogisticLink(),)), kwargs in ((), (; successes=1))
        lik = NegativeBinomialLikelihood(args...; kwargs...)
        @test lik isa NegativeBinomialLikelihood{LogisticLink,Int}
        @test lik.successes == 1
    end

    for args in ((normcdf,), (NormalCDFLink(),)), kwargs in ((; successes=2.0),)
        lik = NegativeBinomialLikelihood(args...; kwargs...)
        @test lik isa NegativeBinomialLikelihood{NormalCDFLink,Float64}
        @test lik.successes == 2
    end

    lik = NegativeBinomialLikelihood()
    test_interface(lik, NegativeBinomial; functor_args=(:successes, :invlink))
end


text:
"""
DIRCLUSTERCOEFFS
----------------
    compute clustering coefficients for a directed graph.
    cc = dirclustercoeffs(A) returns the directed clustering coefficients
    (which generalize the clustering coefficients of an undirected graph,
    and so calling this function on an undirected graph will produce the same
    answer as clustercoeffs, but less efficiently.)

    This function implements the algorithm from Fagiolo, Phys Rev. E.
    76026107 (doi:10:1103/PhysRevE.76.026107).

    (cc,cccyc,ccmid,ccin,ccout,nf) = dirclusteringcoeffs(A) returns different
    components of the clustering coefficients corresponding to cycles,
    middles, in triangles, and out triangles.  See the manuscript for a
    description of the various types of triangles counted in the above metrics.

Functions
---------
- (cc,cccyc,ccmid,ccin,ccout,nf) = dirclustercoeffs{T}(A::SparseMatrixCSC{T,Int64},weighted::Bool)
- (cc,cccyc,ccmid,ccin,ccout,nf) = dirclustercoeffs{T}(A::SparseMatrixCSC{T,Int64})
- (cc,cccyc,ccmid,ccin,ccout,nf) = dirclustercoeffs{T}(A::SparseMatrixCSC{T,Int64},weighted::Bool,normalized::Bool)

Example
-------
~~~
(A,xy,labels) = load_matrix_network_metadata("celegans")
(cc, cccyc, ccmid, ccin, ccout, nf) = dirclustercoeffs(A, true, true)
(maxval, maxind) = findmax(cc)
labels[maxind]
~~~
"""

###########################
###########################

function dirclustercoeffs{T}(A::SparseMatrixCSC{T,Int64},weighted::Bool)
    return dirclustercoeffs(A,weighted,true)
end

function dirclustercoeffs{T}(A::SparseMatrixCSC{T,Int64})
    return dirclustercoeffs(A,true,true)
end

function dirclustercoeffs{T}(A::SparseMatrixCSC{T,Int64},weighted::Bool,normalized::Bool)
    donorm = true
    usew = true

    if !normalized
        donorm = false
    end
    if !weighted
        usew = false
    end

    if usew
        (rp,ci,ai) = sparse_to_csr(A)
        (col_ptr,ri,ati) = sparse_to_csr(A')
    else
        (rp,ci) = sparse_to_csr(A)
        (col_ptr,ri) = sparse_to_csr(A')
    end

    if any(ai.<0)
        error("only positive edge weights allowed")
    end

    n = length(rp)-1
    # initialize all the variables
    cc = zeros(Float64,n)
    ind = zeros(Bool,n)
    cache = zeros(Float64,n)
    degs = zeros(Int64,n)
    cccyc = zeros(Float64,n)
    ccmid = zeros(Float64,n)
    ccin = zeros(Float64,n)
    ccout = zeros(Float64,n)
    nf = zeros(Int64,n)

    # precompute degrees
    for v = 1:n
        for rpi = rp[v]:rp[v+1]-1
            w = ci[rpi]
            if v == w
                continue
            else
                degs[w] = degs[w] + 1
                degs[v] = degs[v] + 1
            end
        end
    end

    ew = 1
    ew2 = 1
    for v = 1:n
        # setup counts for the different cycle types
        bilatedges = 0
        curcccyc = 0
        curccmid = 0
        curccin = 0
        curccout = 0
        # 1.
        # find triangles with out links as last step, so precompute the inlinks
        # back to node v
        for cpi = col_ptr[v]:col_ptr[v+1]-1
            w = ri[cpi]
            if usew
                ew = ati[cpi]
            end
            if v != w
                ind[w] = 1
                cache[w] = ew^(1/3)
            end
        end
        # count cycles (cycles are out->out->out)
        for rpi = rp[v]:rp[v+1]-1
            w = ci[rpi]
            if v == w
                continue
            end # discount self-loop
            if usew
                ew = ai[rpi]^(1/3)
            end
            for rpi2 = rp[w]:rp[w+1]-1
                x = ci[rpi2]
                if x == w
                    continue
                end
                if x == v
                    bilatedges = bilatedges + 1
                    continue
                end
                if ind[x] == 1
                    if usew
                        ew2 = ai[rpi2]
                    end
                    curcccyc = curcccyc + ew*ew2^(1/3) * cache[x]
                end
            end
        end
        # count middle-man circuits (out->in->out)
        for rpi = rp[v]:rp[v+1]-1
            w = ci[rpi]
            if v == w
                continue
            end # discount self-loop
            if usew
                ew = ai[rpi]^(1/3)
            end
            for cpi = col_ptr[w]:col_ptr[w+1]-1
                x=ri[cpi]
                if x==w
                    continue
                end
                if ind[x] == 1
                    if usew
                        ew2 = ati[cpi]
                    end
                    curccmid=curccmid+ew*ew2^(1/3)*cache[x]
                end
            end
        end
        # count in-link circuits (in->out->out)
        for cpi = col_ptr[v]:col_ptr[v+1]-1
            w = ri[cpi]
            if v == w
                continue
            end # discount self-loop
            if usew
                ew = ati[cpi]^(1/3)
            end
            for rpi2 = rp[w]:rp[w+1]-1
                x = ci[rpi2]
                if x == w
                    continue
                end
                if ind[x] == 1
                    if usew
                        ew2 = ai[rpi2]
                    end
                    curccin = curccin + ew*ew2^(1/3) * cache[x]
                end
            end
        end
        # reset and reinit the cache for outlinks
        for cpi = col_ptr[v]:col_ptr[v+1]-1
            w = ri[cpi]
            ind[w]=0
        end # reset indicator
        for rpi = rp[v]:rp[v+1]-1
            w = ci[rpi]
            if usew
                ew = ai[rpi]
            end
            if v != w
                ind[w] = 1
                cache[w] = ew^(1/3)
            end
        end
        # count out-link circuits (out->out->in)
        for rpi = rp[v]:rp[v+1]-1
            w = ci[rpi]
            if v == w
                continue
            end # discount self-loop
            if usew
                ew = ai[rpi]^(1/3)
            end
            for rpi2 = rp[w]:rp[w+1]-1
                x = ci[rpi2]
                if x == w
                    continue
                end
                if ind[x] == 1
                    if usew
                        ew2 = ai[rpi2]
                    end
                    curccout = curccout+ew*ew2^(1/3) * cache[x]
                end
            end
        end
        for rpi = rp[v]:rp[v+1]-1
            w = ci[rpi]
            ind[w] = 0
        end # reset indicator
        # store the values
        curnf = degs[v]*(degs[v]-1) - 2*bilatedges
        curcc = curcccyc + curccmid + curccin + curccout
        if curnf>0 && donorm
            curcc = curcc/curnf
        end
        cc[v] = curcc
        cccyc[v] = curcccyc
        ccmid[v] = curccmid
        ccin[v] = curccin
        ccout[v] = curccout
        nf[v] = curnf

    end
    return (cc,cccyc,ccmid,ccin,ccout,nf)
end


text:
module IZNeuron
export IZStates, IZParams, evolve

struct IZStates{T <: AbstractFloat}
    v::T
    u::T
end

struct IZParams{T <: AbstractFloat}
    a::T
    b::T
    c::T
    d::T
    bc::T
    lazy_threshold::T

    IZParams(a, b, c, d, lazy_threshold) = new{T}(a, b, c, d, lazy_threshold, b * c)
end

function evolve(t, dt, states::IZStates, params::IZParams, inject_fn, update, fire_fn, put_task)

    i_syn, delta_v = inject_fn()
    new_v = states.v + dt * (0.04 * states.v^2 + 5 * states.v + 140 - states.u + i_syn) + delta_v
    new_u = states.u + dt * params.a * (params.b * states.v - states.u)

    if new_v >= 30.0
        new_v = params.c
        new_u += params.d
        fire_fn(t, dt)
    end

    update(IZStates(new_v, new_u))

    if abs(new_v - params.c) > params.lazy_threshold || abs(new_u - params.bc) > params.lazy_threshold
        put_task(t + dt)
    end

end

end # Module end


text:
module Oulu

# Write your package code here.

function addNumbers(a, b)
    a + b
end

end


text:
"""
unfold_particle_filter(num_particles::Int, objects_observed::Matrix{Array{Detection2D}}, camera_trajectories::Matrix{Camera_Params}, num_receptive_fields::Int64)

Performs inference procedure.

# Arguments
- v_matrix::Union{Matrix{Float64}, Nothing} V matrix to condition on. If nothing, do inference over V matrix.
- num_particles::Int
- objects_observed: indexed by ``(\\mathrm{scene, frame, receptive field, detection})``
- camera_trajectories: Indexed by ``(\\mathrm{scene, frame})``
- num_receptive_fields::Int64

array for each frame and array for each receptive field and array for those detections

    # Returns
    - A trace containing the inferred visual system and world state.

    """
    function unfold_particle_filter(v_matrix::Union{Matrix{Float64}, Nothing},
        num_particles::Int64,
        mcmc_steps_outer::Int64,
        mcmc_steps_inner::Int64,
        objects_observed::Matrix{Array{Detection2D}},
        camera_trajectories::Matrix{Camera_Params},
        params::Video_Params,
        V_file::IOStream, ws_file::IOStream, order::Vector{Int64})

        lesioned = !isnothing(v_matrix)

        init_obs = Gen.choicemap()

        if lesioned
            for j = 1:params.n_possible_objects
                init_obs[:init_v_matrix => :lambda_fa => j => :fa] = v_matrix[j,1]
                init_obs[:init_v_matrix => :miss_rate => j => :miss] = v_matrix[j,2]
            end
        end

        #no video, no frames
        println("initialize pf")
        state = initialize_particle_filter(main, (lesioned, 0, 0, params), init_obs, num_particles)

        num_videos, num_frames = size(objects_observed)

        inferred_realities = Array{Any}(undef, num_videos)
        avg_v = zeros(params.n_possible_objects, 2)

        global acceptance_counter = 0
        global total = 0

        for v=1:num_videos
            @time begin

                println("v ", v)

                maybe_resample!(state)
                ess = effective_sample_size(normalize_weights(state.log_weights)[2])
                println("ess after resample ", ess)

                #want to do this for every frame at once. put a : instead of f
                #rf is for receptive field
                obs = Gen.choicemap()

                for f = 1:num_frames
                    camera_params = camera_trajectories[v, f]
                    obs[:videos => v => :frame_chain => f => :camera => :camera_location_x] = camera_params.camera_location.x
                    obs[:videos => v => :frame_chain => f => :camera => :camera_location_y] = camera_params.camera_location.y
                    obs[:videos => v => :frame_chain => f => :camera => :camera_location_z] = camera_params.camera_location.z
                    obs[:videos => v => :frame_chain => f => :camera => :camera_focus_x] = camera_params.camera_focus.x
                    obs[:videos => v => :frame_chain => f => :camera => :camera_focus_y] = camera_params.camera_focus.y
                    obs[:videos => v => :frame_chain => f => :camera => :camera_focus_z] = camera_params.camera_focus.z

                    #for rf = 1:num_receptive_fields
                    #println("objects_observed[v, f][rf] ", objects_observed[v, f][rf])
                    #println("type ", typeof(objects_observed[v, f][rf]))
                    #println(typeof(objects_observed[v, f]))
                    #println(typeof(convert(Array{Any, 1}, objects_observed[v, f])))
                    obs[:videos => v => :frame_chain => f => :observations_2D] = convert(Array{Any, 1}, objects_observed[v, f])
                    #end
                end
                #def should be using map to replace for loops here
                #point is, condition on the camera trajectory and on the observations

                if lesioned
                    for j = 1:params.n_possible_objects
                        obs[:videos => v => :v_matrix => :lambda_fa => j => :fa] = v_matrix[j,1]
                        obs[:videos => v => :v_matrix => :miss_rate => j => :miss] = v_matrix[j,2]
                    end
                end

                particle_filter_step!(state, (lesioned, v, num_frames, params), (UnknownChange(),), obs)

                max_particle_weight = maximum(map(i -> get_score(state.traces[i]), collect(1:num_particles)))
                #println("max log score ", max_particle_weight)

                min_particle_weight = minimum(map(i -> get_score(state.traces[i]), collect(1:num_particles)))
                #println("min log score ", min_particle_weight)

                #println("particles")
                # for i = 1:num_particles
                #      println("particle ", i)
                #      trace = state.traces[i]
                #     choice = get_choices(trace)
                #     open("$i.txt", "w") do f
                #         Base.show(f, "text/plain", get_choices(trace))
                #     end
                #
                #      println("score ", get_score(trace))
                # #
                #      if v==2
                #          selection = select(:videos => 2 => :frame_chain => 111 => :observations_2D)
                #          args = get_args(trace)
                #          println("args ", args)
                #          trace2, weight = regenerate(trace, args, (), selection)
                #          println("after regenerate score ", get_score(trace2))
                #          #println("after regenerate score ", weight)
                #          println("after regenerate ", trace2[:videos => 2 => :frame_chain => 111 => :observations_2D])
                #      end
                #
                #
                #     #println("choices")
                #     #println(get_choices(trace))
                #     println("init v ", trace[:init_v_matrix])
                #     println("scene 1 ", trace[:videos => 1 => :init_scene])
                #     println("at v 1 ", trace[:videos => 1 => :v_matrix])
                #     if v==2
                #         println("scene 2 ", trace[:videos => 2 => :init_scene])
                #         println("at v 2 ", trace[:videos => 2 => :v_matrix])
                #     end
                #
                #     #println(trace)
                # end

                ess = effective_sample_size(normalize_weights(state.log_weights)[2])
                println("ess after pf step ", ess)

                #optional rejuvination
                (perturb_params, n_objects_per_category) = get_probs_categories(objects_observed, params, v, num_frames)
                #println("perturb_params ", perturb_params)
                line_segments_per_category = get_line_segments_per_category(params, objects_observed, camera_trajectories, v, num_frames)
                #line_segments = get_line_segments(objects_observed, camera_trajectories, params, v, num_frames, num_receptive_fields, total_n_objects)
                #to-do: threads here
                #seeds = collect(1:num_particles)
                Threads.@threads for i = 1:num_particles
                    #Random.seed!(seeds[i])
                    #for i = 1:num_particles
                    println("score particle i ", i, " before perturbation", get_score(state.traces[i]))
                    println("trace ", state.traces[i][:videos => v => :init_scene])
                    #println("perturb particle i ", i)
                    state.traces[i] = perturb(lesioned, state.traces[i], v, perturb_params, mcmc_steps_outer, mcmc_steps_inner, line_segments_per_category, params)
                    #println("done perturbing i ", i)
                    println("score particle i ", i, " after perturbation", get_score(state.traces[i]))
                    println("trace ", state.traces[i][:videos => v => :init_scene])
                    #println("trace ", state.traces[i][:videos => v => :init_scene])
                    #println("log score of this trace ", get_score(state.traces[i]))
                    #visualize_trace(state.traces, i, camera_trajectories, v, 1, params)

                    #visualize_trace(state.traces, i, camera_trajectories, v, 2, params)
                end

                #visualize_trace_with_heatmap(state.traces, camera_trajectories, v, 1, params )

                min_particle_weight = minimum(map(i -> get_score(state.traces[i]), collect(1:num_particles)))
                #println("min log score ", min_particle_weight)

                max_particle_weight = maximum(map(i -> get_score(state.traces[i]), collect(1:num_particles)))
                #println("max log score ", max_particle_weight)

                ess = effective_sample_size(normalize_weights(state.log_weights)[2])
                println("ess after rejuvination ", ess)

                inferred_realities[v], avg_v = print_Vs_and_Rs_to_file(V_file, ws_file, state.traces, num_particles, params, v, order, v==num_videos)
                #println("avg_v ", avg_v)
                #println("time of v ", v)
            end #end timer
        end

        println(acceptance_counter)
        println(total)

        return (sample_unweighted_traces(state, num_particles), inferred_realities, avg_v)

    end

    """
    perturb(trace, v::Int64, perturb_params::Perturb_Params, line_segments_per_category::Array{Array{Line_Segment,1},1})
    Does 500 MCMC steps (with different proposal functions) on the scene and on the v matrix.

    """

    @gen function perturb(lesioned::Bool,
        trace,
        v::Int64,
        perturb_params::Perturb_Params,
        mcmc_steps_outer::Int,
        mcmc_steps_inner::Int,
        line_segments_per_category::Array{Array{Line_Segment,1},1},
        params::Video_Params)
        #acceptance_counter = 0
        #proposal_counter = 0

        #println("before perturbations")
        #println("v ", v)
        #
        # println("alpha ", get_retval(trace)[end][2][5,2])
        # println("beta ", get_retval(trace)[end][3][5,2])
        #
        # #println("miss 2 ", trace[:init_v_matrix => :miss_rate => 2 => :miss])
        # println("init miss 5 ", trace[:init_v_matrix => :miss_rate => 5 => :miss])
        # #println("miss 2 ", trace[:videos => v => :v_matrix => :miss_rate => 2 => :miss])
        # println("miss rate 5 ", trace[:videos => v => :v_matrix => :miss_rate => 5 => :miss])
        #
        # println("scene at v ", trace[:videos => v => :init_scene])
        n = length(perturb_params.probs_possible_objects)

        for iter=1:mcmc_steps_outer #try 100 MH moves
            #println("iter ", iter)
            #println("trace ", trace[:videos => v => :init_scene])

            trace = perturb_scene(trace, v, perturb_params, line_segments_per_category, params)
            if lesioned == false
                for iter2=1:mcmc_steps_inner
                    trace = perturb_whole_v_matrix_mh(trace, v, n)
                end
            end
            # println("lambda_fa 2 ", trace[:v_matrix => (:lambda_fa, 2)])
            # println("miss 2 ", trace[:v_matrix => (:miss_rate, 2)])
            # println("lambda_fa 5 ", trace[:v_matrix => (:lambda_fa, 5)])
            # println("miss 5 ", trace[:v_matrix => (:miss_rate, 5)])
        end
        #println("acceptance_counter $(acceptance_counter/proposal_counter)")

        #println("after perturbations")
        #println("v ", v)

        # println("alpha ", get_retval(trace)[end][2][5,2])
        # println("beta ", get_retval(trace)[end][3][5,2])
        #
        # #println("miss 2 ", trace[:init_v_matrix => :miss_rate => 2 => :miss])
        # println("init miss 5 ", trace[:init_v_matrix => :miss_rate => 5 => :miss])
        # #println("miss 2 ", trace[:videos => v => :v_matrix => :miss_rate => 2 => :miss])
        # println("miss rate 5 ", trace[:videos => v => :v_matrix => :miss_rate => 5 => :miss])
        #
        # println("scene at v ", trace[:videos => v => :init_scene])
        return trace
    end

    """
    perturb_scene(trace, v::Int64, perturb_params::Perturb_Params, line_segments_per_category::Array{Array{Line_Segment,1},1})
    Does 3 MCMC steps (with different proposal functions) on the scene.

    """

    function perturb_scene(trace, v::Int64, perturb_params::Perturb_Params, line_segments_per_category::Array{Array{Line_Segment,1},1}, params::Video_Params)
        #println("trace ", trace[:videos => v => :init_scene])
        trace, accepted = add_remove_kernel(trace, v, line_segments_per_category, perturb_params)

        #println("accepted? ", accepted)
        #println("trace ", trace[:videos => v => :init_scene])

        #println("trace ", trace[:videos => v => :init_scene])

        #only try changing location or category if there's at least one object in the scene
        if length(trace[:videos => v => :init_scene]) > 0
            trace, accepted = change_location_kernel(trace, v, 0.1, params, line_segments_per_category)

            #global acceptance_counter = acceptance_counter + accepted
            #global total = total + 1

            #println("accepted? ", accepted)
            #println("trace ", trace[:videos => v => :init_scene])

            #trace, accepted = change_category_kernel(trace, v, perturb_params)
            #println("accepted? ", accepted)
            #println("trace ", trace[:videos => v => :init_scene])
        end
        return trace
    end

    #perturb the whole v matrix
    function perturb_whole_v_matrix_mh(trace, v::Int64, n::Int64)
        selection = get_selection_whole(v, n)
        trace, accepted = mh(trace, selection) #proposes new trace from the prior
        return trace
    end

    """
    perturb_v_matrix(trace, perturb_params::Perturb_Params)
    Picks one element of the v matrix to perturb using metropolis hastings.

    """
    #just pick an element of the matrix to perturb
    function perturb_individual_element_v_matrix_mh(trace, v::Int64, n::Int64)
        i = categorical([0.5, 0.5])
        j = categorical(fill(1/n, n))

        selection = get_selection(v, j, i)
        trace, accepted = mh(trace, selection) #proposes new trace from the prior

        # if i == 1 #if perturbing fa
        #     trace, accepted = mh(trace, proposal_for_v_matrix_fa, (v, j))
        # else
        #     trace, accepted = mh(trace, proposal_for_v_matrix_miss, (v, j))
        # end

        #println("accepted? ", accepted)
        #println("trace ", trace[:v_matrix])
        return trace
    end

    #return a selection for one element of the matrix
    function get_selection(v::Int64, j::Int64, i::Int64)
        if i == 1 #change lambda_fa
            if v == 1
                selection = select(:videos => v => :v_matrix => :lambda_fa => j => :fa,
                :init_v_matrix => :lambda_fa => j => :fa)
            else
                selection = select(:videos => v => :v_matrix => :lambda_fa => j => :fa,
                :videos => v-1 => :v_matrix => :lambda_fa => j => :fa)
            end
        else
            if v == 1
                selection = select(:videos => v => :v_matrix => :miss_rate => j => :miss,
                :init_v_matrix => :miss_rate => j => :miss)
            else
                selection = select(:videos => v => :v_matrix => :miss_rate => j => :miss,
                :videos => v-1 => :v_matrix => :miss_rate => j => :miss)
            end
        end
        return selection
    end

    #return a selection for the whole v matrix
    function get_selection_whole(v::Int64, n::Int64)
        selection = select()

        for j = 1:n
            if v == 1
                push!(selection, :videos => v => :v_matrix => :lambda_fa => j => :fa)
                push!(selection, :init_v_matrix => :lambda_fa => j => :fa)
                push!(selection, :v_matrix => :miss_rate => j => :miss)
                push!(selection, :init_v_matrix => :miss_rate => j => :miss)
            else
                push!(selection, :videos => v => :v_matrix => :lambda_fa => j => :fa)
                push!(selection, :videos => v-1 => :v_matrix => :lambda_fa => j => :fa)
                push!(selection, :videos => v => :v_matrix => :miss_rate => j => :miss)
                push!(selection, :videos => v-1 => :v_matrix => :miss_rate => j => :miss)
            end
        end

        return selection
    end

    """
    perturb_v_matrix_hmc(trace, perturb_params::Perturb_Params)
    Perturbs an element of the v_matrix using Hamiltonian MC

    """

    #just pick an element of the matrix to perturb
    function perturb_v_matrix_hmc(trace, v::Int64, perturb_params::Perturb_Params)
        n = length(perturb_params.probs_possible_objects)
        i = categorical([0.5, 0.5])
        j = categorical(fill(1/n, n))

        if i == 1 #if perturbing fa
            trace, accepted = hmc(trace, select(:videos => v => :v_matrix => :lambda_fa => j => :fa))
        else
            trace, accepted = hmc(trace, select(:videos => v => :v_matrix => :miss_rate => j => :miss))
        end

        #println("accepted? ", accepted)
        #println("trace ", trace[:v_matrix])

        return trace
    end

    """
    proposal_for_v_matrix_fa(trace, j::Int64)
    Performs on MH step on the false alarm rate for object of category j.

        """

        @gen function proposal_for_v_matrix_fa(trace, v::Int64, j::Int64)
            std = 0.01
            choices = get_choices(trace)
            #centered on previous value
            @trace(trunc_normal(choices[:videos => v => :v_matrix => :lambda_fa => j => :fa], std, 0.0, 10000.0), :videos => v => :v_matrix => :lambda_fa => j => :fa) #had to make trunc_normal because of error otherwise
        end

        """
        proposal_for_v_matrix_miss(trace, j::Int64)
        Performs on MH step on the miss rate for object of category j.

            """

            @gen function proposal_for_v_matrix_miss(trace, v::Int64, j::Int64)
                std = 0.1
                choices = get_choices(trace)
                #centered on previous value
                @trace(trunc_normal(choices[:videos => v => :v_matrix => :miss_rate => j => :miss], std, 0.0, 1.0), :videos => v => :v_matrix => :miss_rate => j => :miss)
            end

            """
            get_probs_categories(objects_observed::Matrix{Array{Detection2D}}, params::Video_Params, v::Int64, num_frames::Int64, num_receptive_fields::Int64)

            Returns a probability distribution over the object categories to be used
            in the proposal functions.

            Derived from the objects that the visual system observes.
            """

            function get_probs_categories(objects_observed::Matrix{Array{Detection2D}}, params::Video_Params, v::Int64, num_frames::Int64)
                track_categories = zeros(length(params.possible_objects)) #each element will be the number of times that category was detected. adding 1
                probabilities = zeros(length(params.possible_objects))
                for f = 1:num_frames
                    #for rf = 1:num_receptive_fields
                    for (index, value) in enumerate(objects_observed[v, f])
                        track_categories[value[3]] = track_categories[value[3]]+1#category
                    end
                    #end
                end
                #make it so that 10% of the weight gets divided evenly among all the categories
                other_ten_percent = sum(track_categories)/9
                each = other_ten_percent/length(params.possible_objects)
                probabilities = copy(track_categories)
                probabilities = probabilities .+ each

                #in case of no observations, make uniform
                if sum(probabilities) == 0
                    probabilities .= 1
                end

                return (Perturb_Params(probs_possible_objects = (probabilities)./sum(probabilities)), track_categories)
            end

            """
            get_line_segments_per_category(params::Video_Params, objects_observed::Matrix{Array{Detection2D}}, camera_trajectories::Matrix{Camera_Params}, v::Int64, num_frames::Int64, num_receptive_fields::Int64)

            Gets all line segments in 3d space that correspond to each 2D detection
            in every frame of a scene.
            """
            function get_line_segments_per_category(params::Video_Params, objects_observed::Matrix{Array{Detection2D}}, camera_trajectories::Matrix{Camera_Params}, v::Int64, num_frames::Int64)
                line_segments = Array{Array{Line_Segment, 1}}(undef, length(params.possible_objects))
                for j = 1:length(params.possible_objects)
                    line_segments[j] = []
                end
                for f = 1:num_frames
                    camera_params = camera_trajectories[v, f]
                    #for rf = 1:num_receptive_fields
                    for (index, value) in enumerate(objects_observed[v, f])
                        #println("value ", value)
                        line_segment = get_line_segment(camera_params, params, value)
                        #println("line_segment in pf_inference ", line_segment)
                        push!(line_segments[value[3]], line_segment)
                    end
                    #end
                end
                return line_segments
            end

            """Duplicated from Gen library"""
            function effective_sample_size(log_normalized_weights::Vector{Float64})
                log_ess = -logsumexp(2. * log_normalized_weights)
                return exp(log_ess)
            end

            """Duplicated from Gen Library"""
            function normalize_weights(log_weights::Vector{Float64})
                log_total_weight = logsumexp(log_weights)
                log_normalized_weights = log_weights .- log_total_weight
                return (log_total_weight, log_normalized_weights)
            end

            export unfold_particle_filter

max_stars_repo_name: zhangir-azerbayev/MetaGen

max_issues_repo_name: zhangir-azerbayev/MetaGen

max_forks_repo_name: zhangir-azerbayev/MetaGen


text:
module PlotMesh
using Luxor, Colors, HDF5
# Plotting mesh from JuliaFEM h5-file
# Currently deformed mesh plots only 1 mode of a modal solution.

function triangle(points, mesh_col, br)
	#Plotting a coloured triangle from points.
    jet = RGB{Float64}[
          RGB(clamp(min(4x - 1.5, -4x + 4.5) ,0.0,1.0*br),
              clamp(min(4x - 0.5, -4x + 3.5) ,0.0,1.0*br),
              clamp(min(4x + 0.5, -4x + 2.5) ,0.0,1.0*br))
	      for x in LinRange(0.0,1.0,12)]

    mesh_col = mesh(points, jet[mesh_col])
    setmesh(mesh_col)
    poly(points, :fill, close=true)
    setline(1)
    sethue("black")
    poly(points, :stroke, close=true)
end

function bounds3D(mesh)
  max_x = []
  min_x = []
  max_y = []
  min_y = []
  max_z = []
  min_z = []
  for key in keys(mesh.nodes)
    if isempty(max_x)
      max_x = mesh.nodes[key][1]
      min_x = mesh.nodes[key][1]
      max_y = mesh.nodes[key][2]
      min_y = mesh.nodes[key][2]
      max_z = mesh.nodes[key][3]
      min_z = mesh.nodes[key][3]
    end
    if mesh.nodes[key][1] > max_x
      max_x = mesh.nodes[key][1]
    end
    if mesh.nodes[key][1] < min_x
      min_x = mesh.nodes[key][1]
    end
    if mesh.nodes[key][1] > max_y
      max_y = mesh.nodes[key][2]
    end
    if mesh.nodes[key][1] < min_y
      min_y = mesh.nodes[key][2]
    end
    if mesh.nodes[key][1] > max_z
      max_z = mesh.nodes[key][3]
    end
    if mesh.nodes[key][1] < min_z
      min_z = mesh.nodes[key][3]
    end
  end
  return max_x, min_x, max_y, min_y, max_z, min_z
end

function triangle_normal(P1,P2,P3)
  #Calculating normal from three points.
  x1,y1,z1 = P1
  x2,y2,z2 = P2
  x3,y3,z3 = P3
  n = [(y2-y1)*(z3-z1)-(y3-y1)*(z2-z1)
       (z2-z1)*(x3-x1)-(x2-x1)*(z3-z1)
       (x2-x1)*(y3-y1)-(x3-x1)*(y2-y1)]
  return n
end

function sorted_tri(mesh)
  #Counting all the Tet10 elements
  nTet10 = count(i->(i==:Tet10),values(mesh.element_types))
  #Creating an array for to keep all nodal data and normal for the side triangles
  triangles = zeros(nTet10*4,9+3)
  i = 1
  #Collecting the element data
  for key in keys(mesh.elements)
    if mesh.element_types[key] == :Tet10
      nodes = mesh.elements[key]
      xyz1 = mesh.nodes[nodes[1]]
      xyz2 = mesh.nodes[nodes[2]]
      xyz3 = mesh.nodes[nodes[3]]
      xyz4 = mesh.nodes[nodes[4]]
      #1st triangle
      nodes = [xyz1[1],xyz1[2],xyz1[3],xyz2[1],xyz2[2],xyz2[3],xyz4[1],xyz4[2],xyz4[3]]
      nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
      triangles[i,:] = vcat(nodes,nTri)
      i += 1
      #2nd triangle
      nodes = [xyz2[1],xyz2[2],xyz2[3],xyz3[1],xyz3[2],xyz3[3],xyz4[1],xyz4[2],xyz4[3]]
      nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
      triangles[i,:] = vcat(nodes,nTri)
      i += 1
      #3rd triangle
      nodes = [xyz3[1],xyz3[2],xyz3[3],xyz1[1],xyz1[2],xyz1[3],xyz4[1],xyz4[2],xyz4[3]]
      nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
      triangles[i,:] = vcat(nodes,nTri)
      i += 1
      #4th triangle
      nodes = [xyz1[1],xyz1[2],xyz1[3],xyz3[1],xyz3[2],xyz3[3],xyz2[1],xyz2[2],xyz2[3]]
      nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
      triangles[i,:] = vcat(nodes,nTri)
      i += 1
     end
  end
  return triangles
end

function trias_h5(h5data)
  #Counting all the Tet10 elements
  nTet10 = length(h5data["Topology"]["Tet10"]["Element IDs"])
  #Creating an array for to keep all nodal data and normal for the side triangles
  triangles = zeros(nTet10*4,9+3)
  i = 1
  #Collecting the element data
  for key in range(1,stop=nTet10)
    nodes = h5data["Topology"]["Tet10"]["Connectivity"][:,key]

    xyz1 = h5data["Geometry"][:,nodes[1]+1]
    xyz2 = h5data["Geometry"][:,nodes[2]+1]
    xyz3 = h5data["Geometry"][:,nodes[3]+1]
    xyz4 = h5data["Geometry"][:,nodes[4]+1]

    #1st triangle
    nodes = [xyz1[1],xyz1[2],xyz1[3],xyz2[1],xyz2[2],xyz2[3],xyz4[1],xyz4[2],xyz4[3]]
    nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
    triangles[i,:] = vcat(nodes,nTri)
    i += 1
    #2nd triangle
    nodes = [xyz2[1],xyz2[2],xyz2[3],xyz3[1],xyz3[2],xyz3[3],xyz4[1],xyz4[2],xyz4[3]]
    nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
    triangles[i,:] = vcat(nodes,nTri)
    i += 1
    #3rd triangle
    nodes = [xyz3[1],xyz3[2],xyz3[3],xyz1[1],xyz1[2],xyz1[3],xyz4[1],xyz4[2],xyz4[3]]
    nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
    triangles[i,:] = vcat(nodes,nTri)
    i += 1
    #4th triangle
    nodes = [xyz1[1],xyz1[2],xyz1[3],xyz3[1],xyz3[2],xyz3[3],xyz2[1],xyz2[2],xyz2[3]]
    nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
    triangles[i,:] = vcat(nodes,nTri)
    i += 1

  end
  return triangles
end

function pproject(rP,r0,n,e1,e2)
  #Function projects 3D points to 2D plane with a dot product
  #Returns, 2D coordinates and distance from plane
  g = rP-r0
  s = n'g
  t1 = e1'g
  t2 = e2'g
  return [t1,t2,s]
end

function rangeInt(x,xMin,xMax,n)
  vn = collect(xMin:((xMax-xMin)/n):xMax)
  ff = searchsorted(vn,x)
  idx = ff.stop
  if idx > n
    idx = n
  end
  if idx < 1
    idx = 1
  end
  return idx
end

function draw_triangles(triangles,plane,imageSize, deform, h5data)
  if plane=="XY"
    r0 = [0.,0.,0.]
    e1 = [1.,0.,0.]
    e2 = [0.,1.,0.]
    n  = [0.,0.,1.]
  end
  if plane=="YZ"
    r0 = [0.,0.,0.]
    e1 = [0.,1.,0.]
    e2 = [0.,0.,1.]
    n  = [1.,0.,0.]
  end
  if plane=="XZ"
    r0 = [0.,0.,0.]
    e1 = [1.,0.,0.]
    e2 = [0.,0.,1.]
    n  = [0.,1.,0.]
  end
  if plane=="XYZ"
    r0 = [0.,0.,0.]
    e1 = [-1/sqrt(2),1/sqrt(2),0.]
    e2 = [0.,0.,1.]
    n  = [1/sqrt(3),1/sqrt(3),1/sqrt(3)]
  end

  if deform>-10000.0
    #Counting all the Tet10 elements
    nTet10 = length(h5data["Topology"]["Tet10"]["Element IDs"])
    #Creating an array for to keep all nodal data and normal for the side triangles
    triangles = zeros(nTet10*4,9+3)
    i = 1
    #Collecting the element data
    for key in range(1,stop=nTet10)
      nodes = h5data["Topology"]["Tet10"]["Connectivity"][:,key]

      modeNodes=h5data["Results"]["Natural Frequency Analysis"]["Displacement"]["Mode 1"]

      xyz1 = h5data["Geometry"][:,nodes[1]+1] + modeNodes[:,nodes[1]+1]*deform
      xyz2 = h5data["Geometry"][:,nodes[2]+1] + modeNodes[:,nodes[2]+1]*deform
      xyz3 = h5data["Geometry"][:,nodes[3]+1] + modeNodes[:,nodes[3]+1]*deform
      xyz4 = h5data["Geometry"][:,nodes[4]+1] + modeNodes[:,nodes[4]+1]*deform

      #1st triangle
      nodes = [xyz1[1],xyz1[2],xyz1[3],xyz2[1],xyz2[2],xyz2[3],xyz4[1],xyz4[2],xyz4[3]]
      nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
      triangles[i,:] = vcat(nodes,nTri)
      i += 1
      #2nd triangle
      nodes = [xyz2[1],xyz2[2],xyz2[3],xyz3[1],xyz3[2],xyz3[3],xyz4[1],xyz4[2],xyz4[3]]
      nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
      triangles[i,:] = vcat(nodes,nTri)
      i += 1
      #3rd triangle
      nodes = [xyz3[1],xyz3[2],xyz3[3],xyz1[1],xyz1[2],xyz1[3],xyz4[1],xyz4[2],xyz4[3]]
      nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
      triangles[i,:] = vcat(nodes,nTri)
      i += 1
      #4th triangle
      nodes = [xyz1[1],xyz1[2],xyz1[3],xyz3[1],xyz3[2],xyz3[3],xyz2[1],xyz2[2],xyz2[3]]
      nTri = triangle_normal(nodes[1:3],nodes[4:6],nodes[7:9])
      triangles[i,:] = vcat(nodes,nTri)
      i += 1
    end
  end

  #Projecting and calculating distance and angle from viewplane
  projTri = zeros(size(triangles,1),8)

  for  i=1:size(triangles,1)
    x1, y1, d1 = pproject(triangles[i,1:3],r0,n,e1,e2)
    x2, y2, d2 = pproject(triangles[i,4:6],r0,n,e1,e2)
    x3, y3, d3 = pproject(triangles[i,7:9],r0,n,e1,e2)
    d = (d1+d2+d3)/3
    nt = triangles[i,10:12]
    angle = acos((n'nt)/sqrt(nt'nt+n'n))
    projTri[i,:] = [x1,y1,x2,y2,x3,y3,d,angle]
  end

  #Centering xy-data
  xMax = maximum(projTri[:,[1,3,5]])
  xMin = minimum(projTri[:,[1,3,5]])
  yMax = maximum(projTri[:,[2,4,6]])
  yMin = minimum(projTri[:,[2,4,6]])

  projTri[:,1] = projTri[:,1].-(xMin+xMax)/2
  projTri[:,2] = projTri[:,2].-(yMin+yMax)/2
  projTri[:,3] = projTri[:,3].-(xMin+xMax)/2
  projTri[:,4] = projTri[:,4].-(yMin+yMax)/2
  projTri[:,5] = projTri[:,5].-(xMin+xMax)/2
  projTri[:,6] = projTri[:,6].-(yMin+yMax)/2

  #Scaling to figure size
  xMax = maximum(projTri[:,[1,3,5]])
  xMin = minimum(projTri[:,[1,3,5]])
  yMax = maximum(projTri[:,[2,4,6]])
  yMin = minimum(projTri[:,[2,4,6]])

  scaleX=imageSize[1]/(xMax-xMin)
  scaleY=imageSize[2]/(yMax-yMin)
  scale = minimum([scaleX,scaleY])*0.95

  projTri[:,1:6] = projTri[:,1:6]*scale

  projTri = sortslices(projTri, dims=1, by=x->x[7])

  #Color range index for plotting
  valueMax = maximum(projTri[:,[2,4,6]])
  valueMin = minimum(projTri[:,[2,4,6]])
  nValue = 12

  #Plotting the trianguladed surfaces with color scale
  for i=1:size(projTri,1)
    if projTri[i,8]<=(pi/2)
      nCol = [rangeInt(projTri[i,2],valueMin,valueMax,nValue),
	      rangeInt(projTri[i,4],valueMin,valueMax,nValue),
	      rangeInt(projTri[i,6],valueMin,valueMax,nValue)]
      br = (pi/2-projTri[i,8])/(pi/2)
      triangle([Point(projTri[i,1], projTri[i,2]),
		Point(projTri[i,3], projTri[i,4]),
		Point(projTri[i,5], projTri[i,6])],nCol,br)
    end
  end
end

function plotmesh(filename; imageSize = (400,400), backgroundcolor="white",
				  outputfn="plotmesh_output.png")
	function readh5data()
    	c = h5open(filename, "r") do file
        	h5data=read(file)
    	end
		return h5data
	end

    #Define image
    d = Drawing(imageSize[1], imageSize[2], outputfn)
    background(backgroundcolor)
    origin()

    #Draw mesh
	h5data = readh5data()
    triangles = trias_h5(h5data)
    draw_triangles(triangles,"XZ",imageSize, -Inf, h5data)

    finish()
    return d
end

#=
#This will open a h5-file and read Tri6 mesh and save plot
i = "test_geom.h5"
c = h5open(i, "r") do file
	global h5data=read(file)
end

#Define image
imageSize = (800,800)
Drawing(imageSize[1], imageSize[2], string(i[1:end-2],"png"))
background("white")
origin()

#Draw mesh
triangles = trias_h5(h5data)
draw_triangles(triangles,"XZ",imageSize, -Inf, h5data)

finish()
=#

export trias_h5, draw_triangles, plotmesh
end


text:
hit_chance(chance::Real) = rand() <= chance
hit_chance{T<:Real}(chance::Array{T}) = rand(size(chance)) .<= chance

function Base.endswith{T<:AbstractString}(filename::AbstractString, suffixes::AbstractVector{T})
    for suffix in suffixes
        if endswith(filename, suffix)
            return true
        elseif endswith(filename, uppercase(suffix))
            return true
        end
    end
    return false
end

"""
`listfiles(dir; nargs...)` → `Vector{UTF8String}`

Description
============

Returns the relative paths to all the visible files in the directory
`dir` whose file-endings are contained within `formats`. If `dir` is
not provided the current working path will be used.

Usage
======

    listfiles([dir, hidden=false, expand=false, recursive=true, formats=[".png", ".jpg", ".jpeg", ".bmp"])

Arguments
==========

- `hidden`: If `true`, then files starting with "." are also included.

- `recursive`: If `true`, then the function will also recursively step
through all the subdirectories and append the paths to their content
relative to `dir`.

- `expand`: If `true`, then all the paths will be expaneded to the
full absolute paths instead of being relpathtive to `dir`.

- `formats`: The allowed file endings. Files with a different suffix
will not be included in the return value.

Author(s)
==========

- Christof Stocker (Github: https://github.com/Evizero)

see also
=========

`DirImageSource`
"""
function listfiles(
        dir = ".";
        hidden = false,
        expand = false,
        recursive = true,
        formats = [".png", ".jpg", ".jpeg", ".bmp"])
    dircontent = convert(Vector{UTF8String}, readdir(dir))

    # exclude hidden files (i.e. the ones starting with ".")
    if !hidden
        dircontent = filter(_ -> !startswith(_, "."), dircontent)
    end

    # prepend the directory path to the files
    dircontent = map(_ -> joinpath(dir, _), dircontent)

    # fully expand the path for all files if requested
    if expand
        dircontent = convert(Vector{UTF8String}, map(abspath, dircontent))
    end

    # just filter out entries that are files (exclude dirs)
    files = filter(isfile, dircontent)

    # only include files that end with one of the allowed formats
    files = filter(_ -> endswith(_, formats), files)

    # filter out the dirs as well and recursively append
    # the content of the subdirs to the file list.
    if recursive
        dirs = filter(isdir, dircontent)
        for subdir in dirs
            append!(files,
                    listfiles(subdir;
                        hidden = hidden,
                        expand = expand,
                        recursive = recursive))
        end
    end
    files
end