text:
@given("entity(person) set owns attribute type: is-alive") do context
    @fail "Implement me"
end

@given("entity(person) set owns attribute type: age") do context
    @fail "Implement me"
end

@given("entity(person) set owns attribute type: score") do context
    @fail "Implement me"
end

@given("entity(person) set owns attribute type: name") do context
    @fail "Implement me"
end

@given("entity(person) set owns attribute type: email") do context
    @fail "Implement me"
end

@given("entity(person) set owns attribute type: birth-date") do context
    @fail "Implement me"
end

@when("delete attribute: \$x") do context
    @fail "Implement me"
end

#=
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

from behave import *
from hamcrest import *

from typedb.common.exception import TypeDBClientException
from typedb.common.label import Label
from tests.behaviour.config.parameters import parse_bool, parse_list, RootLabel, parse_label
from tests.behaviour.context import Context

@step("put {root_label:RootLabel} type: {type_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    if root_label == RootLabel.ENTITY:
        context.tx().concepts().put_entity_type(type_label)
    elif root_label == RootLabel.RELATION:
        context.tx().concepts().put_relation_type(type_label)
    else:
        raise ValueError("Unrecognised value")

@step("delete {root_label:RootLabel} type: {type_label}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).delete()
        assert False
    except TypeDBClientException:
        pass

@step("delete {root_label:RootLabel} type: {type_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).delete()

@step("{root_label:RootLabel}({type_label}) is null: {is_null}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, is_null):
    is_null = parse_bool(is_null)
    assert_that(context.get_thing_type(root_label, type_label) is None, is_(is_null))

@step("{root_label:RootLabel}({type_label}) set label: {new_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, new_label: str):
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_label(new_label)

@step("{root_label:RootLabel}({type_label}) get label: {get_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, get_label: str):
    assert_that(context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_label().name(), is_(get_label))

@step("{root_label:RootLabel}({type_label}) set abstract: {is_abstract}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, is_abstract):
    is_abstract = parse_bool(is_abstract)
    thing_type = context.get_thing_type(root_label, type_label)
    if is_abstract:
        thing_type.as_remote(context.tx()).set_abstract()
    else:
        thing_type.as_remote(context.tx()).unset_abstract()

@step("{root_label:RootLabel}({type_label}) is abstract: {is_abstract}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, is_abstract):
    is_abstract = parse_bool(is_abstract)
    assert_that(context.get_thing_type(root_label, type_label).as_remote(context.tx()).is_abstract(), is_(is_abstract))

@step("{root_label:RootLabel}({type_label}) set supertype: {super_label}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, super_label: str):
    if root_label == RootLabel.ENTITY:
        entity_supertype = context.tx().concepts().get_entity_type(super_label)
        try:
            context.tx().concepts().get_entity_type(type_label).as_remote(context.tx()).set_supertype(entity_supertype)
            assert False
        except TypeDBClientException:
            pass
    elif root_label == RootLabel.ATTRIBUTE:
        attribute_supertype = context.tx().concepts().get_attribute_type(super_label)
        try:
            context.tx().concepts().get_attribute_type(type_label).as_remote(context.tx()).set_supertype(attribute_supertype)
            assert False
        except TypeDBClientException:
            pass
    elif root_label == RootLabel.RELATION:
        relation_supertype = context.tx().concepts().get_relation_type(super_label)
        try:
            context.tx().concepts().get_relation_type(type_label).as_remote(context.tx()).set_supertype(relation_supertype)
            assert False
        except TypeDBClientException:
            pass
    else:
        raise ValueError("Unrecognised value")

@step("{root_label:RootLabel}({type_label}) set supertype: {super_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, super_label: str):
    if root_label == RootLabel.ENTITY:
        entity_supertype = context.tx().concepts().get_entity_type(super_label)
        context.tx().concepts().get_entity_type(type_label).as_remote(context.tx()).set_supertype(entity_supertype)
    elif root_label == RootLabel.ATTRIBUTE:
        attribute_supertype = context.tx().concepts().get_attribute_type(super_label)
        context.tx().concepts().get_attribute_type(type_label).as_remote(context.tx()).set_supertype(attribute_supertype)
    elif root_label == RootLabel.RELATION:
        relation_supertype = context.tx().concepts().get_relation_type(super_label)
        context.tx().concepts().get_relation_type(type_label).as_remote(context.tx()).set_supertype(relation_supertype)
    else:
        raise ValueError("Unrecognised value")

@step("{root_label:RootLabel}({type_label}) get supertype: {super_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, super_label: str):
    supertype = context.get_thing_type(root_label, super_label)
    assert_that(context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_supertype(), is_(supertype))

@step("{root_label:RootLabel}({type_label}) get supertypes contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    super_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_supertypes()]
    for super_label in super_labels:
        assert_that(actuals, has_item(super_label))

@step("{root_label:RootLabel}({type_label}) get supertypes do not contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    super_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_supertypes()]
    for super_label in super_labels:
        assert_that(actuals, not_(has_item(super_label)))

@step("{root_label:RootLabel}({type_label}) get subtypes contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    sub_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_subtypes()]
    for sub_label in sub_labels:
        assert_that(actuals, has_item(sub_label))

@step("{root_label:RootLabel}({type_label}) get subtypes do not contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    sub_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_subtypes()]
    for sub_label in sub_labels:
        assert_that(actuals, not_(has_item(sub_label)))

@step("{root_label:RootLabel}({type_label}) set owns key type: {att_type_label} as {overridden_label}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str, overridden_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    overridden_type = context.tx().concepts().get_attribute_type(overridden_label)
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_owns(attribute_type, overridden_type, is_key=True)
        assert False
    except TypeDBClientException:
        pass

@step("{root_label:RootLabel}({type_label}) set owns key type: {att_type_label}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_owns(attribute_type, is_key=True)
        assert False
    except TypeDBClientException:
        pass

@step("{root_label:RootLabel}({type_label}) set owns key type: {att_type_label} as {overridden_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str, overridden_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    overridden_type = context.tx().concepts().get_attribute_type(overridden_label)
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_owns(attribute_type, overridden_type, is_key=True)

@step("{root_label:RootLabel}({type_label}) set owns key type: {att_type_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_owns(attribute_type, is_key=True)

@step("{root_label:RootLabel}({type_label}) unset owns attribute type: {att_type_label}; throws exception")
@step("{root_label:RootLabel}({type_label}) unset owns key type: {att_type_label}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).unset_owns(attribute_type)
        assert False
    except TypeDBClientException:
        pass

@step("{root_label:RootLabel}({type_label}) unset owns attribute type: {att_type_label}")
@step("{root_label:RootLabel}({type_label}) unset owns key type: {att_type_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).unset_owns(attribute_type)

@step("{root_label:RootLabel}({type_label}) get owns key types contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    attribute_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_owns(keys_only=True)]
    for attribute_label in attribute_labels:
        assert_that(actuals, has_item(attribute_label))

@step("{root_label:RootLabel}({type_label}) get owns key types do not contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    attribute_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_owns(keys_only=True)]
    for attribute_label in attribute_labels:
        assert_that(actuals, not_(has_item(attribute_label)))

@step("{root_label:RootLabel}({type_label}) set owns attribute type: {att_type_label} as {overridden_label}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str, overridden_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    overridden_type = context.tx().concepts().get_attribute_type(overridden_label)
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_owns(attribute_type, overridden_type)
        assert False
    except TypeDBClientException:
        pass

@step("{root_label:RootLabel}({type_label}) set owns attribute type: {att_type_label}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_owns(attribute_type)
        assert False
    except TypeDBClientException:
        pass

@step("{root_label:RootLabel}({type_label}) set owns attribute type: {att_type_label} as {overridden_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str, overridden_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    overridden_type = context.tx().concepts().get_attribute_type(overridden_label)
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_owns(attribute_type, overridden_type)

@step("{root_label:RootLabel}({type_label}) set owns attribute type: {att_type_label}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, att_type_label: str):
    attribute_type = context.tx().concepts().get_attribute_type(att_type_label)
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_owns(attribute_type)

@step("{root_label:RootLabel}({type_label}) get owns attribute types contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    attribute_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_owns()]
    for attribute_label in attribute_labels:
        assert_that(actuals, has_item(attribute_label))

@step("{root_label:RootLabel}({type_label}) get owns attribute types do not contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    attribute_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_owns()]
    for attribute_label in attribute_labels:
        assert_that(actuals, not_(has_item(attribute_label)))

@step("{root_label:RootLabel}({type_label}) set plays role: {role_label:ScopedLabel} as {overridden_label:ScopedLabel}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, role_label: Label, overridden_label: Label):
    role_type = context.tx().concepts().get_relation_type(role_label.scope()).as_remote(context.tx()).get_relates(role_label.name())
    overridden_type = context.tx().concepts().get_relation_type(overridden_label.scope()).as_remote(context.tx()).get_relates(overridden_label.name())
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_plays(role_type, overridden_type)
        assert False
    except TypeDBClientException:
        pass

@step("{root_label:RootLabel}({type_label}) set plays role: {role_label:ScopedLabel} as {overridden_label:ScopedLabel}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, role_label: Label, overridden_label: Label):
    role_type = context.tx().concepts().get_relation_type(role_label.scope()).as_remote(context.tx()).get_relates(role_label.name())
    overridden_type = context.tx().concepts().get_relation_type(overridden_label.scope()).as_remote(context.tx()).get_relates(overridden_label.name())
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_plays(role_type, overridden_type)

@step("{root_label:RootLabel}({type_label}) set plays role: {role_label:ScopedLabel}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, role_label: Label):
    role_type = context.tx().concepts().get_relation_type(role_label.scope()).as_remote(context.tx()).get_relates(role_label.name())
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_plays(role_type)
        assert False
    except TypeDBClientException:
        pass

@step("{root_label:RootLabel}({type_label}) set plays role: {role_label:ScopedLabel}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, role_label: Label):
    role_type = context.tx().concepts().get_relation_type(role_label.scope()).as_remote(context.tx()).get_relates(role_label.name())
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).set_plays(role_type)

@step("{root_label:RootLabel}({type_label}) unset plays role: {role_label:ScopedLabel}; throws exception")
def step_impl(context: Context, root_label: RootLabel, type_label: str, role_label: Label):
    role_type = context.tx().concepts().get_relation_type(role_label.scope()).as_remote(context.tx()).get_relates(role_label.name())
    try:
        context.get_thing_type(root_label, type_label).as_remote(context.tx()).unset_plays(role_type)
        assert False
    except TypeDBClientException:
        pass

@step("{root_label:RootLabel}({type_label}) unset plays role: {role_label:ScopedLabel}")
def step_impl(context: Context, root_label: RootLabel, type_label: str, role_label: Label):
    role_type = context.tx().concepts().get_relation_type(role_label.scope()).as_remote(context.tx()).get_relates(role_label.name())
    context.get_thing_type(root_label, type_label).as_remote(context.tx()).unset_plays(role_type)

@step("{root_label:RootLabel}({type_label}) get playing roles contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    role_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_plays()]
    for role_label in role_labels:
        assert_that(role_label, is_in(actuals))

@step("{root_label:RootLabel}({type_label}) get playing roles do not contain")
def step_impl(context: Context, root_label: RootLabel, type_label: str):
    role_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.get_thing_type(root_label, type_label).as_remote(context.tx()).get_plays()]
    for role_label in role_labels:
        assert_that(role_label, not_(is_in(actuals)))

@step("thing type root get supertypes contain")
def step_impl(context: Context):
    super_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.tx().concepts().get_root_thing_type().as_remote(context.tx()).get_supertypes()]
    for super_label in super_labels:
        assert_that(super_label, is_in(actuals))

@step("thing type root get supertypes do not contain")
def step_impl(context: Context):
    super_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.tx().concepts().get_root_thing_type().as_remote(context.tx()).get_supertypes()]
    for super_label in super_labels:
        assert_that(super_label, not_(is_in(actuals)))

@step("thing type root get subtypes contain")
def step_impl(context: Context):
    sub_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.tx().concepts().get_root_thing_type().as_remote(context.tx()).get_subtypes()]
    for sub_label in sub_labels:
        assert_that(sub_label, is_in(actuals))

@step("thing type root get subtypes do not contain")
def step_impl(context: Context):
    sub_labels = [parse_label(s) for s in parse_list(context.table)]
    actuals = [t.get_label() for t in context.tx().concepts().get_root_thing_type().as_remote(context.tx()).get_subtypes()]
    for sub_label in sub_labels:
        assert_that(sub_label, not_(is_in(actuals)))
=#


text:
@testset "correctly detect stdlib objects" begin
    @test Lint.stdlibobject(:var) == nothing
    @test Lint.stdlibobject(:axes) !== nothing
end

@testset "Handle empty code block" begin
    let s = """
        a = begin end
        """,
        ex=Meta.parse(s)
        @test Lint.get_tag_per_condition(ex) == Lint.AssignTag
        @test Lint.guesstype(ex, LintContext()) == Nothing
        msgs = lintstr(s)
        @test isempty(msgs)
    end
end
@testset "Colon (no parameters)" begin
    @test Lint.iscolon(:(:))
    @test Lint.iscolon(:(Colon()))
end

ext: jl


text:
using DSGE
using HDF5

include("../util.jl")
path = dirname(@__FILE__)

# Test in generic case
# Global minimum at (a, a^2)
a = 1
b = 100
function rosenbrock(x::Vector)
    return (a-x[1])^2.0 + b*(x[2]-x[1]^2.0)^2.0
end

# this is the actual gradient of the Rosenbrock function
function rosenbrock_grad(x::Array)
    dr = similar(x)
    dr[1] = -2*(a-x[1]) - 2b*(2x[1])*(x[2]-x[1]^2.0)
    dr[2] = 2b*(x[2]-x[1]^2.0)
    badg = false
    return dr, badg
end

# A really bad guess
x_init = [10.0, -9.0]
x_expected = [a, a^2.0]

res_real_grad, _ = csminwel(rosenbrock, rosenbrock_grad, x_init)
@test_approx_eq x_expected res_real_grad.minimum
res_numeric_grad, _ = csminwel(rosenbrock, x_init)
@test_approx_eq_eps x_expected res_numeric_grad.minimum 1e-8

nothing


text:
function vecrotx(angle::Real)
    sn = sin(angle)
    cn = cos(angle)
    rx = [1 0 0 0; 0 cn -sn 0; 0 sn cn 0; 0 0 0 1]
    return rx
end

function vecroty(angle::Real)
    sn = sin(angle)
    cn = cos(angle)
    ry = [cn 0 sn 0; 0 1 0 0; -sn 0 cn 0; 0 0 0 1]
    return ry
end

function vecrotz(angle::Real)
    sn = sin(angle)
    sn = sin(angle)
    cn = cos(angle)
    rz = [ cn -sn 0 0; sn cn 0 0; 0 0 1 0; 0 0 0 1]
    return rz
end

function vecangle(num, den)
    if length(num)==1
        ang = [atan(num, den)]
    else
        ang = zero(num)
        for i=1:length(num)
            ang[i] = atan(num[i], den[i])
        end
    end
    index = findall(ang .< 0.0)
    ang[index] = 2*pi .+ ang[index]
    return ang
end

function vecmag(vec)
    mag  = sqrt.(sum(vec.^2, dims=1))
    return mag
end

function vecnorm(vec)
    nvec = vec./repeat(sqrt.(sum(vec.^2, dims=1)), outer=tuple(vcat(size(vec, 1), ones(Int, ndims(vec)-1))...))
    return nvec
end

function vecscale(vector::Vector)
    s = vcat(vector, [0, 0])
    ss = [s[1] 0 0 0; 0 s[2] 0 0; 0 0 s[3] 0; 0 0 0 1]
    return ss
end

function vectrans(vector::Vector)
    v = vcat(vector,[0, 0])
    dd = [1 0 0 v[1]; 0 1 0 v[2]; 0 0 1 v[3]; 0 0 0 1]
    return dd
end


text:
__precompile__()

module ExprRules

import TreeView: walk_tree
using StatsBase
using AbstractTrees
using DataStructures  #NodeRecycler

export
        Grammar,
        RuleNode,
        NodeLoc,

        ExpressionIterator,

        @grammar,
        max_arity,
        depth,
        node_depth,
        isterminal,
        iseval,
        return_type,
        contains_returntype,
        nchildren,
        child_types,
        nonterminals,
        get_executable,
        sample,
        root_node_loc,
        count_expressions,
        mindepth_map,
        mindepth,

        SymbolTable,
        interpret,

        NodeRecycler,
        recycle!

include("interpreter.jl")
using .Interpreter

"""
    isterminal(rule::Any, types::AbstractVector{Symbol})

Returns true if the rule is terminal, ie does not contain any of the types in the provided vector.
For example, :(x) is terminal, and :(1+1) is terminal, but :(Real + Real) is typically not.
"""
function isterminal(rule::Any, types::AbstractVector{Symbol})
    if isa(rule, Expr)
        for arg in rule.args
            if !isterminal(arg, types)
                return false
            end
        end
    end
    return rule ∉ types
end

"""
    iseval(rule::Any)

Returns true if the rule is the special evaluate immediately function, i.e., _()
"""
iseval(rule) = false
iseval(rule::Expr) = (rule.head == :call && rule.args[1] == :_)

"""
    get_childtypes(rule::Any, types::AbstractVector{Symbol})

Returns the child types of a production rule.
"""
function get_childtypes(rule::Any, types::AbstractVector{Symbol})
    retval = Symbol[]
    if isa(rule, Expr)
        for arg in rule.args
            append!(retval, get_childtypes(arg, types))
        end
    elseif rule ∈ types
        push!(retval, rule)
    end
    return retval
end

"""
    Grammar

Represents a grammar and its production rules.
Use the @grammar macro to create a Grammar object.
"""
struct Grammar
    rules::Vector{Any}    # list of RHS of rules (subexpressions)
    types::Vector{Symbol} # list of LHS of rules (types, all symbols)
    isterminal::BitVector # whether rule i is terminal
    iseval::BitVector     # whether rule i is an eval rule
    bytype::Dict{Symbol,Vector{Int}}   # maps type to all rules of said type
    childtypes::Vector{Vector{Symbol}} # list of types of the children for each rule. Empty if terminal
end
"""
    @grammar

Define a grammar and return it as a Grammar. For example:
```julia-repl
grammar = @grammar begin
    R = x
    R = 1 | 2
    R = R + R
end
```
"""
macro grammar(ex)
    rules = Any[]
    types = Symbol[]
    bytype = Dict{Symbol,Vector{Int}}()
    for e in ex.args
        if isa(e, Expr)
            if e.head == :(=)
                s = e.args[1] # name of return type
                rule = e.args[2] # expression?
                rvec = Any[]
                _parse_rule!(rvec, rule)
                for r in rvec
                    push!(rules, r)
                    push!(types, s)
                    bytype[s] = push!(get(bytype, s, Int[]), length(rules))
                end
            end
        end
    end
    alltypes = collect(keys(bytype))
    is_terminal = [isterminal(rule, alltypes) for rule in rules]
    is_eval = [iseval(rule) for rule in rules]
    childtypes = [get_childtypes(rule, alltypes) for rule in rules]
    return Grammar(rules, types, is_terminal, is_eval, bytype, childtypes)
end
_parse_rule!(v::Vector{Any}, r) = push!(v, r)
function _parse_rule!(v::Vector{Any}, ex::Expr)
    if ex.head == :call && ex.args[1] == :|
         terms = length(ex.args) == 2 ?
            collect(interpret(ex.args[2])) :    #|(a:c) case
            ex.args[2:end]                      #a|b|c case
        for t in terms
            _parse_rule!(v, t)
        end
    else
        push!(v, ex)
    end
end

Base.getindex(grammar::Grammar, typ::Symbol) = grammar.bytype[typ]

"""
   append!(grammar1::Grammar, grammar2::Grammar)

Appends the production rules of grammar2 to grammar1.
"""
function Base.append!(grammar1::Grammar, grammar2::Grammar)
    N = length(grammar1.rules)
    append!(grammar1.rules, grammar2.rules)
    append!(grammar1.types, grammar2.types)
    append!(grammar1.isterminal, grammar2.isterminal)
    append!(grammar1.iseval, grammar2.iseval)
    append!(grammar1.childtypes, copy.(grammar2.childtypes))
    for (s,v) in grammar2.bytype
        grammar1.bytype[s] = append!(get(grammar1.bytype, s, Int[]), N .+ v)
    end
    grammar1
end

"""
    nonterminals(grammar::Grammar)

Returns a list of nonterminals in the grammar.
"""
nonterminals(grammar::Grammar) = collect(keys(grammar.bytype))

"""
    return_type(grammar::Grammar, rule_index::Int)

Returns the type of the production rule at rule_index.
"""
return_type(grammar::Grammar, rule_index::Int) = grammar.types[rule_index]

"""
    child_types(grammar::Grammar, rule_index::Int)

Returns the types of the children (nonterminals) of the production rule at rule_index.
"""
child_types(grammar::Grammar, rule_index::Int) = grammar.childtypes[rule_index]

"""
    isterminal(grammar::Grammar, rule_index::Int)

Returns true if the production rule at rule_index is terminal, i.e., does not contain any nonterminal symbols.
"""
isterminal(grammar::Grammar, rule_index::Int) = grammar.isterminal[rule_index]

"""
    iseval(grammar::Grammar)

Returns true if any production rules in grammar contain the special _() eval function.
"""
iseval(grammar::Grammar) = any(grammar.iseval)

"""
    iseval(grammar::Grammar, rule_index::Int)

Returns true if the production rule at rule_index contains the special _() eval function.
"""
iseval(grammar::Grammar, index::Int) = grammar.iseval[index]

"""
    nchildren(grammar::Grammar, rule_index::Int)

Returns the number of children (nonterminals) of the production rule at rule_index.
"""
nchildren(grammar::Grammar, rule_index::Int) = length(grammar.childtypes[rule_index])

"""
    max_arity(grammar::Grammar)

Returns the maximum arity (number of children) over all production rules in the grammar.
"""
max_arity(grammar::Grammar) = maximum(length(cs) for cs in grammar.childtypes)

"""
    RuleNode

Type for representing nodes in an expression tree.
"""
mutable struct RuleNode
    ind::Int # index in grammar
    _val::Any  #value of _() evals
    children::Vector{RuleNode}
end
RuleNode(ind::Int) = RuleNode(ind, nothing, RuleNode[])
RuleNode(ind::Int, children::Vector{RuleNode}) = RuleNode(ind, nothing, children)
RuleNode(ind::Int, _val::Any) = RuleNode(ind, _val, RuleNode[])

include("recycler.jl")

"""
    return_types(grammar::Grammar, node::RuleNode)

Returns the return type in the production rule used by node.
"""
return_type(grammar::Grammar, node::RuleNode) = grammar.types[node.ind]

"""
    child_types(grammar::Grammar, node::RuleNode)

Returns the list of child types in the production rule used by node.
"""
child_types(grammar::Grammar, node::RuleNode) = grammar.childtypes[node.ind]

"""
    isterminal(grammar::Grammar, node::RuleNode)

Returns true if the production rule used by node is terminal, i.e., does not contain any nonterminal symbols.
"""
isterminal(grammar::Grammar, node::RuleNode) = grammar.isterminal[node.ind]

"""
    nchildren(grammar::Grammar, node::RuleNode)

Returns the number of children in the production rule used by node.
"""
nchildren(grammar::Grammar, node::RuleNode) = length(child_types(grammar, node))

"""
    contains_returntype(node::RuleNode, grammar::Grammar, sym::Symbol, maxdepth::Int=typemax(Int))

Returns true if the tree rooted at node contains at least one node at depth less than maxdepth
with the given return type.
"""
function contains_returntype(node::RuleNode, grammar::Grammar, sym::Symbol, maxdepth::Int=typemax(Int))
    maxdepth < 1 && return false
    if return_type(grammar, node) == sym
        return true
    end
    for c in node.children
        if contains_returntype(c, grammar, sym, maxdepth-1)
            return true
        end
    end
    return false
end
function Base.:(==)(A::RuleNode, B::RuleNode)
    (A.ind == B.ind) &&
        (A._val == B._val) &&
        (length(A.children) == length(B.children)) && #required because zip doesn't check lengths
        all(isequal(a,b) for (a,b) in zip(A.children, B.children))
end

function Base.hash(node::RuleNode, h::UInt=zero(UInt))
    retval = hash(node.ind, h)
    for child in node.children
        retval = hash(child, retval)
    end
    return retval
end

function Base.show(io::IO, grammar::Grammar)
    for i in eachindex(grammar.rules)
        println(io, i, ": ", grammar.types[i], " = ", grammar.rules[i])
    end
end

function Base.show(io::IO, node::RuleNode; separator=",", last_child::Bool=false)
    print(io, node.ind)
    if !isempty(node.children)
        print(io, "{")
        for (i,c) in enumerate(node.children)
            show(io, c, separator=separator, last_child=(i == length(node.children)))
        end
        print(io, "}")
    elseif !last_child
        print(io, separator)
    end
end

"""
Return the number of vertices in the tree rooted at root.
"""
function Base.length(root::RuleNode)
    retval = 1
    for c in root.children
        retval += length(c)
    end
    return retval
end

"""
    depth(root::RuleNode)

Return the depth of the expression tree rooted at root.
"""
function depth(root::RuleNode)
    retval = 1
    for c in root.children
        retval = max(retval, depth(c)+1)
    end
    return retval
end

"""
    node_depth(root::RuleNode, node::RuleNode)

Return the depth of node for an expression tree rooted at root.  Depth is 1 when root == node.
"""
function node_depth(root::RuleNode, node::RuleNode)
    root === node && return 1
    for c in root.children
        d = node_depth(c, node)
        d > 0 && (return d+1)
    end
    return 0
end

"""
    get_executable(rulenode::RuleNode, grammar::Grammar)

Returns the executable julia expression represented in the expression tree with root rulenode.  The returned expression can be evaluated using eval().
"""
function get_executable(rulenode::RuleNode, grammar::Grammar)
    root = (rulenode._val != nothing) ?
        rulenode._val : deepcopy(grammar.rules[rulenode.ind])
    if !grammar.isterminal[rulenode.ind] # not terminal
        root,j = _get_executable(root, rulenode, grammar)
    end
    return root
end
function _get_executable(expr::Expr, rulenode::RuleNode, grammar::Grammar, j=0)
    for (k,arg) in enumerate(expr.args)
        if isa(arg, Expr)
            expr.args[k],j = _get_executable(arg, rulenode, grammar, j)
        elseif haskey(grammar.bytype, arg)
            child = rulenode.children[j+=1]
            expr.args[k] = (child._val != nothing) ?
                child._val : deepcopy(grammar.rules[child.ind])
            if !isterminal(grammar, child)
                expr.args[k],_ = _get_executable(expr.args[k], child, grammar, 0)
            end
        end
    end
    return expr, j
end
function _get_executable(typ::Symbol, rulenode::RuleNode, grammar::Grammar, j=0)
    retval = typ
    if haskey(grammar.bytype, typ)
        child = rulenode.children[1]
        retval = (child._val != nothing) ?
            child._val : deepcopy(grammar.rules[child.ind])
        if !grammar.isterminal[child.ind]
            retval,_ = _get_executable(retval, child, grammar, 0)
        end
    end
    retval, j
end

"""
    Core.eval(rulenode::RuleNode, grammar::Grammar)

Evaluate the expression tree with root rulenode.
"""
Core.eval(rulenode::RuleNode, grammar::Grammar) = Core.eval(Main, get_executable(rulenode, grammar))
Core.eval(grammar::Grammar, index::Int) = Core.eval(Main, grammar.rules[index].args[2])
"""
    Core.eval(tab::SymbolTable, ex::Expr)

Evaluate the expression ex using symbol table tab
"""
Core.eval(tab::SymbolTable, ex) = interpret(tab, ex)
function Base.display(rulenode::RuleNode, grammar::Grammar)
    root = get_executable(rulenode, grammar)
    if isa(root, Expr)
        walk_tree(root)
    else
        root
    end
end

"""
    rand(::Type{RuleNode}, grammar::Grammar, typ::Symbol, max_depth::Int=10)

Generates a random RuleNode of return type typ and maximum depth max_depth.
"""
function Base.rand(::Type{RuleNode}, grammar::Grammar, typ::Symbol, max_depth::Int=10,
    bin::Union{NodeRecycler,Nothing}=nothing)
    rules = grammar[typ]

    if max_depth <= 1
        terminals = filter(r->isterminal(grammar,r), rules)
        rule_index = !isempty(terminals) ? StatsBase.sample(terminals) : StatsBase.sample(rules)
    else
        rule_index = StatsBase.sample(rules)
    end

    rulenode = iseval(grammar, rule_index) ?
        RuleNode(bin, rule_index, Core.eval(grammar, rule_index)) :
        RuleNode(bin, rule_index)

    if !grammar.isterminal[rule_index]
        for ch in child_types(grammar, rule_index)
            push!(rulenode.children, rand(RuleNode, grammar, ch, max_depth-1, bin))
        end
    end
    return rulenode
end
"""
    rand(::Type{RuleNode}, grammar::Grammar, typ::Symbol, dmap::AbstractVector{Int}, max_depth::Int=10)

Generates a random RuleNode of return type typ and maximum depth max_depth guided by a minimum depth map dmap.
"""
function Base.rand(::Type{RuleNode}, grammar::Grammar, typ::Symbol, dmap::AbstractVector{Int},
    max_depth::Int=10, bin::Union{NodeRecycler,Nothing}=nothing)
    rules = grammar[typ]
    rule_index = StatsBase.sample(filter(r->dmap[r] ≤ max_depth, rules))

    rulenode = iseval(grammar, rule_index) ?
        RuleNode(bin, rule_index, Core.eval(grammar, rule_index)) :
        RuleNode(bin, rule_index)

    if !grammar.isterminal[rule_index]
        for ch in child_types(grammar, rule_index)
            push!(rulenode.children, rand(RuleNode, grammar, ch, dmap, max_depth-1, bin))
        end
    end
    return rulenode
end

mutable struct RuleNodeAndCount
    node::RuleNode
    cnt::Int
end
"""
    sample(root::RuleNode, typ::Symbol, grammar::Grammar, maxdepth::Int=typemax(Int))

Selects a uniformly random node from the tree, limited to maxdepth.
"""
function StatsBase.sample(root::RuleNode, maxdepth::Int=typemax(Int))
    x = RuleNodeAndCount(root, 1)
    for child in root.children
        _sample(child, x, maxdepth-1)
    end
    x.node
end
function _sample(node::RuleNode, x::RuleNodeAndCount, maxdepth::Int)
    maxdepth < 1 && return
    x.cnt += 1
    if rand() <= 1/x.cnt
        x.node = node
    end
    for child in node.children
        _sample(child, x, maxdepth-1)
    end
end

"""
    sample(root::RuleNode, typ::Symbol, grammar::Grammar,
                          maxdepth::Int=typemax(Int))

Selects a uniformly random node of the given return type, typ, limited to maxdepth.
"""
function StatsBase.sample(root::RuleNode, typ::Symbol, grammar::Grammar,
                          maxdepth::Int=typemax(Int))
    x = RuleNodeAndCount(root, 0)
    if grammar.types[root.ind] == typ
        x.cnt += 1
    end
    for child in root.children
        _sample(child, typ, grammar, x, maxdepth-1)
    end
    grammar.types[x.node.ind] == typ || error("type $typ not found in RuleNode")
    x.node
end
function _sample(node::RuleNode, typ::Symbol, grammar::Grammar, x::RuleNodeAndCount,
                 maxdepth::Int)
    maxdepth < 1 && return
    if grammar.types[node.ind] == typ
        x.cnt += 1
        if rand() <= 1/x.cnt
            x.node = node
        end
    end
    for child in node.children
        _sample(child, typ, grammar, x, maxdepth-1)
    end
end

"""
    NodeLoc

A helper struct that points to a node in the tree via its parent such that the child can be easily
swapped out.
If i is 0 the node pointed to is the root node and parent is the node itself.
"""
struct NodeLoc
    parent::RuleNode
    i::Int
end

"""
    root_node_loc(root::RuleNode)

Returns a NodeLoc pointing to the root node.
"""
root_node_loc(root::RuleNode) = NodeLoc(root, 0)

"""
    get(root::RuleNode, loc::NodeLoc)

Obtain the node pointed to by loc.
"""
function Base.get(root::RuleNode, loc::NodeLoc)
    parent, i = loc.parent, loc.i
    if loc.i > 0
        return parent.children[i]
    else
        return root
    end
end

"""
    insert!(loc::NodeLoc, rulenode::RuleNode)

Replaces the subtree pointed to by loc with the given rulenode.
"""
function Base.insert!(root::RuleNode, loc::NodeLoc, rulenode::RuleNode)
    parent, i = loc.parent, loc.i
    if loc.i > 0
        parent.children[i] = rulenode
    else
        root.ind = rulenode.ind
        root._val = rulenode._val
        root.children = rulenode.children
    end
    return root
end

mutable struct NodeLocAndCount
    loc::NodeLoc
    cnt::Int
end
"""
    sample(::Type{NodeLoc}, root::RuleNode, maxdepth::Int=typemax(Int))

Selects a uniformly random node in the tree no deeper than maxdepth using reservoir sampling.
Returns a NodeLoc that specifies the location using its parent so that the subtree can be replaced.
"""
function StatsBase.sample(::Type{NodeLoc}, root::RuleNode, maxdepth::Int=typemax(Int))
    x = NodeLocAndCount(root_node_loc(root), 1)
    _sample(NodeLoc, root, x, maxdepth-1)
    x.loc
end
function _sample(::Type{NodeLoc}, node::RuleNode, x::NodeLocAndCount, maxdepth::Int)
    maxdepth < 1 && return
    for (j,child) in enumerate(node.children)
        x.cnt += 1
        if rand() <= 1/x.cnt
            x.loc = NodeLoc(node, j)
        end
        _sample(NodeLoc, child, x, maxdepth-1)
    end
end

"""
    sample(::Type{NodeLoc}, root::RuleNode, typ::Symbol, grammar::Grammar)

Selects a uniformly random node in the tree of a given type, specified using its parent such that the subtree can be replaced.
Returns a NodeLoc.
"""
function StatsBase.sample(::Type{NodeLoc}, root::RuleNode, typ::Symbol, grammar::Grammar,
                          maxdepth::Int=typemax(Int))
    x = NodeLocAndCount(root_node_loc(root), 0)
    if grammar.types[root.ind] == typ
        x.cnt += 1
    end
    _sample(NodeLoc, root, typ, grammar, x, maxdepth-1)
    grammar.types[get(root,x.loc).ind] == typ || error("type $typ not found in RuleNode")
    x.loc
end
function _sample(::Type{NodeLoc}, node::RuleNode, typ::Symbol, grammar::Grammar,
                 x::NodeLocAndCount, maxdepth::Int)
    maxdepth < 1 && return
    for (j,child) in enumerate(node.children)
        if grammar.types[child.ind] == typ
            x.cnt += 1
            if rand() <= 1/x.cnt
                x.loc = NodeLoc(node, j)
            end
        end
        _sample(NodeLoc, child, typ, grammar, x, maxdepth-1)
    end
end

###

function _next_state!(node::RuleNode, grammar::Grammar, max_depth::Int)

    if max_depth < 1
        return (node, false) # did not work
    elseif isterminal(grammar, node)
        # do nothing
        return (node, false) # cannot change leaves
    else # !isterminal
        if isempty(node.children)
            if max_depth ≤ 1
                return (node,false) # cannot expand
            end

            # build out the node
            for c in child_types(grammar, node)

                worked = false
                i = 0
                child = RuleNode(0)
                child_rules = grammar[c]
                while !worked && i < length(child_rules)
                    i += 1
                    child = RuleNode(child_rules[i])
                    worked = true
                    if !isterminal(grammar, child)
                        child, worked = _next_state!(child, grammar, max_depth-1)
                    end
                end
                if !worked
                    return (node, false) # did not work
                end
                push!(node.children, child)
            end

            return (node, true)
        else # not empty
            # make one change, starting with rightmost child
            worked = false
            child_index = length(node.children) + 1
            while !worked && child_index > 1
                child_index -= 1
                child = node.children[child_index]

                child, child_worked = _next_state!(child, grammar, max_depth-1)
                while !child_worked
                    child_type = return_type(grammar, child)
                    child_rules = grammar[child_type]
                    i = something(findfirst(isequal(child.ind), child_rules), 0)
                    if i < length(child_rules)
                        child_worked = true
                        child = RuleNode(child_rules[i+1])
                        if !isterminal(grammar, child)
                            child, child_worked = _next_state!(child, grammar, max_depth-1)
                        end
                        node.children[child_index] = child
                    else
                        break
                    end
                end

                if child_worked
                    worked = true

                    # reset remaining children
                    for child_index2 in child_index+1 : length(node.children)
                        c = child_types(grammar, node)[child_index2]
                        worked = false
                        i = 0
                        child = RuleNode(0)
                        child_rules = grammar[c]
                        while !worked && i < length(child_rules)
                            i += 1
                            child = RuleNode(child_rules[i])
                            worked = true
                            if !isterminal(grammar, child)
                                child, worked = _next_state!(child, grammar, max_depth-1)
                            end
                        end
                        if !worked
                            break
                        end
                        node.children[child_index2] = child
                    end
                end
            end

            return (node, worked)
        end
    end
end

"""
    ExpressionIterator(grammar::Grammar, max_depth::Int, sym::Symbol)

An iterator over all possible expressions of a grammar up to max_depth with start symbol sym.
"""
mutable struct ExpressionIterator
    grammar::Grammar
    max_depth::Int
    sym::Symbol
end
Base.IteratorSize(::ExpressionIterator) = Base.SizeUnknown()
Base.eltype(::ExpressionIterator) = RuleNode

function Base.iterate(iter::ExpressionIterator)
    grammar, sym, max_depth = iter.grammar, iter.sym, iter.max_depth
    node = RuleNode(grammar[sym][1])
    if isterminal(grammar, node)
        return (deepcopy(node), node)
    else
        node, worked =  _next_state!(node, grammar, max_depth)
        while !worked
            # increment root's rule
            rules = grammar[sym]
            i = something(findfirst(isequal(node.ind), rules), 0)
            if i < length(rules)
                node, worked = RuleNode(rules[i+1]), true
                if !isterminal(grammar, node)
                    node, worked = _next_state!(node, grammar, max_depth)
                end
            else
                break
            end
        end
        return worked ? (deepcopy(node), node) : nothing
    end
end

function Base.iterate(iter::ExpressionIterator, state::RuleNode)
    grammar, max_depth = iter.grammar, iter.max_depth
    node, worked = _next_state!(state, grammar, max_depth)

    while !worked
        # increment root's rule
        rules = grammar[iter.sym]
        i = something(findfirst(isequal(node.ind), rules), 0)
        if i < length(rules)
            node, worked = RuleNode(rules[i+1]), true
            if !isterminal(grammar, node)
                node, worked = _next_state!(node, grammar, max_depth)
            end
        else
            break
        end
    end
    return worked ? (deepcopy(node), node) : nothing
end

"""
    count_expressions(grammar::Grammar, max_depth::Int, sym::Symbol)

Count the number of possible expressions of a grammar up to max_depth with start symbol sym.
"""
function count_expressions(grammar::Grammar, max_depth::Int, sym::Symbol)
    retval = 0
    for root_rule in grammar[sym]
        node = RuleNode(root_rule)
        if isterminal(grammar, node)
            retval += 1
        else
            node, worked = _next_state!(node, grammar, max_depth)
            while worked
                retval += 1
                node, worked = _next_state!(node, grammar, max_depth)
            end
        end
    end
    return retval
end

"""
    count_expressions(iter::ExpressionIterator)

Count the number of possible expressions in the expression iterator.
"""
count_expressions(iter::ExpressionIterator) = count_expressions(iter.grammar, iter.max_depth, iter.sym)

# Interface to AbstractTrees.jl
AbstractTrees.children(node::RuleNode) = node.children
AbstractTrees.printnode(io::IO, node::RuleNode) = print(io, node.ind)

"""
    mindepth_map(grammar::Grammar)

Returns the minimum depth achievable for each production rule, dmap.
"""
function mindepth_map(grammar::Grammar)
    dmap0 = Int[isterminal(grammar,i) ? 0 : typemax(Int)/2 for i in eachindex(grammar.rules)]
    dmap1 = fill(-1, length(grammar.rules))
    while dmap0 != dmap1
        for i in eachindex(grammar.rules)
            dmap1[i] = _mindepth(grammar, i, dmap0)
        end
        dmap1, dmap0 = dmap0, dmap1
    end
    dmap0
end
function _mindepth(grammar::Grammar, rule_index::Int, dmap::AbstractVector{Int})
    isterminal(grammar, rule_index) && return 0
    return 1 + maximum([mindepth(grammar, ctyp, dmap) for ctyp in child_types(grammar, rule_index)])
end
"""
    mindepth(grammar::Grammar, typ::Symbol, dmap::AbstractVector{Int})

Returns the minimum depth achievable for a given nonterminal symbol
"""
function mindepth(grammar::Grammar, typ::Symbol, dmap::AbstractVector{Int})
    return minimum(dmap[grammar.bytype[typ]])
end

"""
    Interpreter.SymbolTable(grammar::Grammar, mod::Module=Main)

Returns a symbol table populated with mapping from symbols in grammar to
symbols in module mod or Main, if defined.
"""
function Interpreter.SymbolTable(grammar::Grammar, mod::Module=Main)
    tab = SymbolTable()
    for rule in grammar.rules
        _add_to_symboltable!(tab, rule, mod)
    end
    tab
end
_add_to_symboltable!(tab::SymbolTable, rule::Any, mod::Module) = true
function _add_to_symboltable!(tab::SymbolTable, rule::Expr, mod::Module)
    if rule.head == :call && !iseval(rule)
        s = rule.args[1]
        if !_add_to_symboltable!(tab, s, mod)
            @warn "Unable to add function $s to symbol table"
        end
        for s in rule.args[2:end]  #nested exprs
            _add_to_symboltable!(tab, s, mod)
        end
    end
    return true
end
function _add_to_symboltable!(tab::SymbolTable, s::Symbol, mod::Module)
    if isdefined(mod, s)
        tab[s] = getfield(mod, s)
        return true
    elseif isdefined(Base, s)
        tab[s] = getfield(Base, s)
        return true
    elseif isdefined(Main, s)
        tab[s] = getfield(Main, s)
        return true
    else
        return false
    end
end

end # module


text:
#
# SurfaceNets in Julia
#
# Ported from the Javascript implementation by Mikola Lysenko (C) 2012
# https://github.com/mikolalysenko/mikolalysenko.github.com/blob/master/Isosurface/js/surfacenets.js
# MIT License
#
# Based on: S.F. Gibson, "Constrained Elastic Surface Nets". (1998) MERL Tech Report.
#

#Look up Table
include("lut/sn.jl")

function isosurface(sdf::AbstractArray{T, 3}, method::NaiveSurfaceNets, ::Type{VertType}=SVector{3,Float64}, ::Type{FaceType}=SVector{4, Int};
                    origin=VertType(-1,-1,-1), widths=VertType(2,2,2)) where {T, VertType, FaceType}

    scale = widths ./ VertType(size(sdf) .- 1)  # subtract 1 because an SDF with N points per side has N-1 cells

    dims = size(sdf)

    vertices = VertType[]
    faces = FaceType[]

    sizehint!(vertices,ceil(Int,maximum(dims)^2))
    sizehint!(faces,ceil(Int,maximum(dims)^2))

    n = 0
    R = Array{Int}([1, (dims[1]+1), (dims[1]+1)*(dims[2]+1)])
    buf_no = 1

    buffer = fill(zero(Int),R[3]*2)

    #March over the voxel grid
    zi = 0
    @inbounds while zi<dims[3]-1

        # m is the pointer into the buffer we are going to use.
        # This is slightly obtuse because javascript does not have good support for packed data structures, so we must use typed arrays :(
        # The contents of the buffer will be the indices of the vertices on the previous x/y slice of the volume
        m = 1 + (dims[1]+1) * (1 + buf_no * (dims[2]+1))

        for yi = 0:dims[2]-2
            for xi = 0:dims[1]-2

                inds = (xi,yi,zi)

                # Read in 8 field values around this vertex and store them in an array
                # Also calculate 8-bit mask, like in marching cubes, so we can speed up sign checks later
                @inbounds grid = (sdf[xi+1,yi+1,zi+1],
                                  sdf[xi+2,yi+1,zi+1],
                                  sdf[xi+1,yi+2,zi+1],
                                  sdf[xi+2,yi+2,zi+1],
                                  sdf[xi+1,yi+1,zi+2],
                                  sdf[xi+2,yi+1,zi+2],
                                  sdf[xi+1,yi+2,zi+2],
                                  sdf[xi+2,yi+2,zi+2])

                mask = method.insidepositive ? _get_cubeindex_pos(grid, method.iso) : _get_cubeindex(grid, method.iso)

                # Check for early termination if cell does not intersect boundary
                if _no_triangles(mask)
                    m += 1
                    continue
                end

                # iso level correction
                grid = grid .- method.iso

                #Sum up edge intersections
                edge_mask = sn_edge_table[mask]

                _sn_add_verts!(inds, vertices, grid, edge_mask, buffer, m, scale, origin, method.eps, VertType)

                #Now we need to add faces together, to do this we just loop over 3 basis components
                _sn_add_faces!(inds, faces, edge_mask, mask, buffer, m, R, FaceType)

                m += 1
            end
            m += 2
        end
        zi += 1
        buf_no = xor(buf_no,1)
        R[3]=-R[3]
    end

    vertices, faces
end

function isosurface(f::Function, method::NaiveSurfaceNets,
                    ::Type{VertType}=SVector{3,Float64}, ::Type{FaceType}=SVector{4, Int};
                    origin=VertType(-1,-1,-1), widths=VertType(2,2,2),
                    samples::NTuple{3,T}=_DEFAULT_SAMPLES) where {T <: Integer, VertType, FaceType}

    scale = widths ./ VertType(samples .- 1)  # subtract 1 because an SDF with N points per side has N-1 cells

    vertices = VertType[]
    faces = FaceType[]

    sizehint!(vertices,ceil(Int,maximum(samples)^2))
    sizehint!(faces,ceil(Int,maximum(samples)^2))

    R = Array{Int}([1, (samples[1]+1), (samples[1]+1)*(samples[2]+1)])
    buf_no = 1

    buffer = fill(zero(Int),R[3]*2)

    zv = zero(eltype(VertType))
    zvt = zero(VertType)
    grid = (zv,zv,zv,zv,zv,zv,zv,zv)
    points = (zvt,zvt,zvt,zvt,zvt,zvt,zvt,zvt)

    #March over the voxel grid
    zi = 0
    @inbounds while zi<samples[3]-1

        # m is the pointer into the buffer we are going to use.
        # This is slightly obtuse because javascript does not have good support for packed data structures, so we must use typed arrays :(
        # The contents of the buffer will be the indices of the vertices on the previous x/y slice of the volume
        m = 1 + (samples[1]+1) * (1 + buf_no * (samples[2]+1))

        for yi = 0:samples[2]-2
            for xi = 0:samples[1]-2

                inds = (xi,yi,zi)

                if xi == 0
                    points = (VertType(xi,yi,zi).* scale .+ origin,
                              VertType(xi+1,yi,zi).* scale .+ origin,
                              VertType(xi,yi+1,zi).* scale .+ origin,
                              VertType(xi+1,yi+1,zi).* scale .+ origin,
                              VertType(xi,yi,zi+1).* scale .+ origin,
                              VertType(xi+1,yi,zi+1).* scale .+ origin,
                              VertType(xi,yi+1,zi+1).* scale .+ origin,
                              VertType(xi+1,yi+1,zi+1).* scale .+ origin)
                    grid = (f(points[1]),
                            f(points[2]),
                            f(points[3]),
                            f(points[4]),
                            f(points[5]),
                            f(points[6]),
                            f(points[7]),
                            f(points[8]))
                else
                    points = (points[2],
                              VertType(xi+1,yi,zi).* scale .+ origin,
                              points[4],
                              VertType(xi+1,yi+1,zi).* scale .+ origin,
                              points[6],
                              VertType(xi+1,yi,zi+1).* scale .+ origin,
                              points[8],
                              VertType(xi+1,yi+1,zi+1).* scale .+ origin)
                    grid = (grid[2],
                            f(points[2]),
                            grid[4],
                            f(points[4]),
                            grid[6],
                            f(points[6]),
                            grid[8],
                            f(points[8]))
                end

                # Also calculate 8-bit mask, like in marching cubes, so we can speed up sign checks later
                mask = method.insidepositive ? _get_cubeindex_pos(grid, method.iso) : _get_cubeindex(grid, method.iso)

                # Check for early termination if cell does not intersect boundary
                if _no_triangles(mask)
                    m += 1
                    continue
                end

                # iso level correction
                grid = grid .- method.iso

                #Sum up edge intersections
                edge_mask = sn_edge_table[mask]

                # add vertices
                _sn_add_verts!(inds, vertices, grid, edge_mask, buffer, m, scale, origin, method.eps, VertType)

                #Now we need to add faces together, to do this we just loop over 3 basis components
                _sn_add_faces!(inds, faces, edge_mask, mask, buffer, m, R, FaceType)

                m += 1
            end
            m += 2
        end
        zi += 1
        buf_no = xor(buf_no,1)
        R[3]=-R[3]
    end

    vertices, faces
end

@inline function _sn_add_verts!(inds, vertices, grid, edge_mask, buffer, m, scale, origin, eps, ::Type{VertType}) where {VertType}
    v = zero(VertType)
    T = eltype(VertType)
    e_count = 0

    #For every edge of the cube...
    @inbounds for i=0x00:0x0b

        #Use edge mask to check if it is crossed
        iszero(edge_mask & (0x0001<<i)) && continue

        #If it did, increment number of edge crossings
        e_count += 1

        #Now find the point of intersection
        e0 = cube_edges[(i<<0x01)+1]       #Unpack vertices
        e1 = cube_edges[(i<<0x01)+2]
        g0 = grid[e0+1]                 #Unpack grid values
        g1 = grid[e1+1]
        t  = g0 - g1                 #Compute point of intersection

        abs(t) <= eps && continue

        t = g0 / t

        #Interpolate vertices and add up intersections (this can be done without multiplying)
        a1, a2, a3 = 0x01 & e0, 0x02 & e0, 0x04 & e0
        b1, b2, b3 = 0x01 & e1, 0x02 & e1, 0x04 & e1
        xj = T(!iszero(a1))
        yj = T(!iszero(a2))
        zj = T(!iszero(a3))
        a1 != b1 && (xj += !iszero(a1) ? -t : t)
        a2 != b2 && (yj += !iszero(a2) ? -t : t)
        a3 != b3 && (zj += !iszero(a3) ? -t : t)
        v = v .+ VertType(xj,yj,zj)

    end # edge check

    #Now we just average the edge intersections and add them to coordinate
    s = one(T) / e_count
    v = (VertType(inds...)  .+ s .* v) .* scale .+ origin

    #Add vertex to buffer, store pointer to vertex index in buffer
    buffer[m+1] = length(vertices)
    push!(vertices, v)
end

function _sn_add_faces!(inds, faces, edge_mask, mask, buffer, m, R, ::Type{FaceType}) where {FaceType}
    for i = 0x00:0x02
        #The first three entries of the edge_mask count the crossings along the edge
        iszero(edge_mask & (0x0001<<i)) && continue

        # i = axes we are point along.  iu, iv = orthogonal axes
        iu = (i+0x01)%0x03
        iv = (i+0x02)%0x03

        #If we are on a boundary, skip it
        iszero(inds[iu+1]) || iszero(inds[iv+1]) && continue

        #Otherwise, look up adjacent edges in buffer
        du = R[iu+1]
        dv = R[iv+1]

        #Remember to flip orientation depending on the sign of the corner.
        if !iszero(mask & 0x01)
            if length(FaceType) == 4
                push!(faces,FaceType(buffer[m+1]+1, buffer[m-du+1]+1, buffer[m-du-dv+1]+1, buffer[m-dv+1]+1))
            elseif length(FaceType) == 3
                push!(faces,FaceType(buffer[m+1]+1, buffer[m-du+1]+1, buffer[m-du-dv+1]+1))
                push!(faces,FaceType(buffer[m-du-dv+1]+1, buffer[m-dv+1]+1, buffer[m+1]+1))
            end
        else
            if length(FaceType) == 4
                push!(faces,FaceType(buffer[m+1]+1, buffer[m-dv+1]+1, buffer[m-du-dv+1]+1, buffer[m-du+1]+1))
            elseif length(FaceType) == 3
                push!(faces,FaceType(buffer[m+1]+1, buffer[m-dv+1]+1, buffer[m-du-dv+1]+1))
                push!(faces,FaceType(buffer[m-du-dv+1]+1, buffer[m-du+1]+1, buffer[m+1]+1))
            end
        end
    end
end


text:
function trapezoidal_rule(n::Int, a = -1.0, b = 1.0, ::Type{T} = typeof((b-a)/n)) where {T}
    grid = EquispacedGrid(n, a, b)
    weights = step(grid) * ones(T, n)
    weights[1] /= 2
    weights[end] /= 2
    grid, weights
end

function rectangular_rule(n::Int, a = 0.0, b = 1.0, ::Type{T} = typeof((b-a)/n)) where {T}
    grid = PeriodicEquispacedGrid(n, a, b)
    weight = step(grid)
    grid, Fill(weight,length(grid))
end

# The implementation of the first and second Fejer rule and of Clenshaw-Curtis
# quadrature is based on the exposition in:
# J. Waldvogel, Fast construction of the Fejer and Clenshaw-Curtis quadrature
# rules, BIT 43(1):1-18, 2003.
function fejer_vector1!(v)
    n = length(v)
    T = eltype(v)
    v[1] = T(2)
    for k in 1:(n-1)>>1
        r = T(2)/(1-4k^2)*exp(im*k*T(pi)/n)
        v[k+1] = r
        v[n-k+1] = conj(r)
    end
    if iseven(n)
        v[n>>1+1] = 0
    end
    v
end

function fejer_first_rule(n::Int, ::Type{T} = Float64) where {T}
    v = zeros(Complex{T}, n)
    fejer_vector1!(v)
    ChebyshevNodes{T}(n), real(ifft(v))
end

function fejer_vector2!(v)
    n = length(v)
    T = eltype(v)
    v[1] = T(2)
    for k in 1:(n-1)>>1
        r = T(2)/(1-4k^2)
        v[k+1] = r
        v[n-k+1] = r
    end
    if iseven(n)
        k = n>>1-1
        v[k+1] = T(2)/(1-4k^2)
        v[n>>1+1] = T(n-3)/(2*(n>>1)-1) - 1
    end
    v
end

# TODO: verify, is this correct?
function fejer_second_rule(n::Int, ::Type{T} = Float64) where {T}
    v = zeros(T, n)
    weights = zeros(T, n+1)
    fejer_vector2!(v)
    weights[1:n] = real(ifft(v))
    weights[n+1] = weights[1]
    ChebyshevExtremae{T}(n+1), weights
end

function cc_vector_g!(g)
    n = length(g)
    T = eltype(g)
    w0cc = one(T) / (n^2-1+mod(n,2))
    for k in 1:n
        g[k] = -w0cc
    end
    if iseven(n)
        g[n>>1+1] = w0cc * ((2-mod(n,2))*n-1)
    end
    g
end

function clenshaw_curtis(n, ::Type{T} = Float64) where {T}
    v = zeros(T, n)
    g = zeros(T, n)
    weights = zeros(T, n+1)
    fejer_vector2!(v)
    cc_vector_g!(g)
    weights[1:n] = real(ifft(v+g))
    weights[n+1] = weights[1]
    ChebyshevExtremae{T}(n+1), weights
end

function rescale_cc_quad(x, w, a, b)
    y = (reverse(x).+1)/2 * (b-a) .+ a
    v = w .* (b-a)/2
    y,v
end

function rescale_fejer_quad(x, w, a, b)
    y = (x.+1)/2 * (b-a) .+ a
    v = w .* (b-a)/2
    y,v
end

function graded_rule(sigma, a, b, M, n, T = Float64, ndiff = 2)
    xg = zeros(T, 0)
    wg = zeros(T, 0)
    q1 = one(T)
    for m = M-1:-1:1
        x,w = BasisFunctions.fejer_first_rule(n, T)
        q0 = sigma*q1
        xs,ws = rescale_fejer_quad(x, w, q0, q1)
        xg = [xg; xs]
        wg = [wg; ws]
        q1 = q0
        n -= ndiff
    end
    x,w = BasisFunctions.fejer_first_rule(n, T)
    xs,ws = rescale_fejer_quad(x, w, 0, q1)
    xg = [xg; xs]
    wg = [wg; ws]
    a .+ xg*(b-a), wg * (b-a)
end


text:
using U2Net
using Test


text:
function lintblock(ex::Expr, ctx::LintContext)
    lastexpr = nothing
    similarexprs = Expr[]
    diffs = Float64[]
    checksimilarityflag = !(LintIgnore(:W651, "") in ctx.ignore)

    if checksimilarityflag
        checksimilarity = ()->begin
            if length(similarexprs) <= 2 # not much I can do
                diffs = Float64[]
                lastexpr = nothing
                similarexprs = Expr[]
            else
                # cyclic diffs, so now we have at least 3 similarity scores
                push!(diffs, expr_similar_score(similarexprs[1], similarexprs[end]))
                local n = length(diffs)
                local m = mean(diffs)
                local s = std(diffs)
                local m2 = mean([diffs[end-1], diffs[end]])
                # look for screw up at the end
                #println(diffs, "\nm=", m, " s=", s, " m2=", m2, " m-m2=", m-m2)
                if m2 < m && m-m2 > s/2.5 && s / m > 0.0001
                    msg(ctx, :W651, "the last of a $(n)-expr block looks different" *
                        ";  Avg similarity score: " * @sprintf("%8.2f", m) *
                        ";  Last part: " * @sprintf("%9.2f", m2))
                end
                diffs = Float64[]
                lastexpr = nothing
                similarexprs = Expr[]
            end
        end
    end

    for (i,sube) in enumerate(ex.args)
        if isa(sube, Expr)
            if sube.head == :line
                ctx.line = ctx.lineabs + sube.args[1]-1
                continue
            elseif sube.head == :return && i != length(ex.args)
                msg(ctx, :W641, "unreachable code after return")
                lintexpr(sube, ctx)
                break
            else
                if checksimilarityflag
                    if lastexpr != nothing
                        local dif = expr_similar_score(lastexpr, sube)
                        if dif > SIMILARITY_THRESHOLD
                            if !isempty(similarexprs)
                                append!(similarexprs, [lastexpr, sube])
                            else
                                push!(similarexprs, sube)
                            end
                            push!(diffs, dif)
                        else
                            checksimilarity()
                        end
                    end
                end
                lintexpr(sube, ctx)
                lastexpr = sube
            end
        elseif isa(sube, QuoteNode)
            lintexpr(sube,ctx)
        elseif isa(sube, LineNumberNode)
            ctx.line = ctx.lineabs + sube.line-1
            continue
        elseif isa(sube, Symbol)
            registersymboluse(sube, ctx)
            if checksimilarityflag
                checksimilarity()
            end
        end
    end

    if checksimilarityflag
        checksimilarity()
    end
end

function expr_similar_score(e1::Expr, e2::Expr, base::Float64 = 1.0)
    if e1.head != e2.head
        return -base
    end

    score = base - abs(length(e1.args) - length(e2.args)) * base * 2.0

    for i in 1:min(length(e1.args), length(e2.args))
        if typeof(e1.args[i]) == Expr && typeof(e2.args[i]) == Expr
            score += expr_similar_score(e1.args[i], e2.args[i], base * 1.1)
        elseif typeof(e1.args[i]) == typeof(e2.args[i])
            score += base * 0.3
            if e1.args[i] == e2.args[i]
                score += base * 0.8
            end
        else
            score -= base
        end
        if score < 0.0 # so early disagreement dominates and short-circuit
            break
        end
    end
    return score
end

function test_similarity_string(str::T) where T<:AbstractString
    i = start(str)
    firstexpr = nothing
    lastexpr = nothing
    diffs = Float64[]
    while !done(str,i)
        problem = false
        ex = nothing
        try
            (ex, i) = parse(str,i)
        catch
            problem = true
        end
        if !problem
            if firstexpr == nothing
                firstexpr = ex
            end
            if lastexpr != nothing
                push!(diffs, expr_similar_score(lastexpr, ex))
            end
            lastexpr = ex
        else
            break
        end
    end
    if lastexpr != nothing && length(diffs) >= 2
        push!(diffs, expr_similar_score(lastexpr, firstexpr))
    end
    return diffs
end

function linttry(ex::Expr, ctx::LintContext)
    # try
    withcontext(ctx, LocalContext(ctx.current)) do
        lintexpr(ex.args[1], ctx)
    end

    # catch
    if isa(ex.args[2], Symbol)
        withcontext(ctx, LocalContext(ctx.current)) do
            set!(ctx.current, ex.args[2], VarInfo(location(ctx), Exception))
            lintexpr(ex.args[3], ctx)
        end
    end

    # finally
    if length(ex.args) > 3
        @assert length(ex.args) == 4
        withcontext(ctx, LocalContext(ctx.current)) do
            lintexpr(ex.args[4], ctx)
        end
    end
end

function lintlet(ex::Expr, ctx::LintContext)
    withcontext(ctx, LocalContext(ctx.current)) do
        for arg in ex.args
            # it's always assignment, or the parser would have thrown at the very start
            if isexpr(arg, :(=)) && !isexpr(arg.args[1], :call)
                lintassignment(arg, ctx; islocal = true)
            end
        end
        blk = ex.args[1]
        @assert isexpr(blk, :block)
        for arg in blk.args
            if isexpr(arg, :(=)) && !isexpr(arg.args[1], :call)
                lintassignment(arg, ctx; islocal = true)
            else
                lintexpr(arg, ctx)
            end
        end
    end
end


text:
function Adjust(da_or_de, beta, alpha)
    return da_or_de^beta + alpha
end

function integrate(a0, e0, p)
    """
    General purpose integrator for Nyadzani & Razzaque eqns for a & e.

    Params:
        a0 - The initial semimajor axis, measured in solar radii
        e0 - The initial eccentricity, dimensionless
        p  - A vector of parameters:
                 p[1] = m1 (units: Solar Mass)
                 p[2] = m2 (units: Solar Mass)
                 p[3] = Beta for LT Adjustment (units: dimensionless)
                 p[4] = Alpha for LT adjustment (units: [unit]/s)
                 p[5] = Lifetime (evolution + rejuvenation)

    Returns:
        A  - An array of the semimajor axes of the binary system over time. (Solar Mass)
        E  - An array of the eccentricities of the binary system over time. (no dim.)
    """

    Solar_Mass   = 1.989e30  # kg
    Solar_Radius = 696340000.0 # m
               G = 6.67e-11  # m^3 kg^-1 s^-2
               c = 299792458.0 # m/s

               A = Float64[]
               E = Float64[]
               H = Float64[]

    m1, m2, expo, alph, evotime = p[1], p[2], p[3], p[4], p[5]

    # number of seconds in a year
    seconds_per_year = 60 * 60 * 24 * 365.25

    #######################
    ##     Unit Check     #
    #######################
    a = a0 * Solar_Radius # Meters
    e = e0                # Dimensionless
    m1 = m1 * Solar_Mass  # Kilogram
    m2 = m2 * Solar_Mass  # Kilogram
    #######################
    ##   End Unit Check   #
    #######################

    # Beta has units m^4 / s
    beta = ((64/5) * G^3 * m1 * m2 * (m1 + m2) / (c^5))

    A = push!(A, a)
    E = push!(E, e)
    H = push!(H, 0.0)

    total_time = 0

    # Integrate until past the end of the universe, or a 10km orbit
    while total_time/seconds_per_year + evotime < 1e11 && a > 1e4
        initial_da = (- beta / ( (a^3) * (1 - e^2)^(7/2) ))
        da = initial_da * (1 + (73/24) * e^2 + (37/96) * e^4)

        intial_de = (((-19/12) * beta) / (a^4*(1-e^2)^(5/2)))
        de = intial_de * (e + (121/304) * e^3) # Units: s^-1

        # da = Adjust(da, expo, alph)
        # de = Adjust(de, expo, alph)

        timeA = abs(1e-2 * a/da)
        if e > 1e-10
            timeE = abs(1e-2 * e/de)
        else
            de = 0
            e = 1e-10
            timeE = timeA * 10
        end

        dt2 = (evotime + total_time/seconds_per_year)*0.23076752*0.5*seconds_per_year

        dt = min(timeE, timeA, dt2)

        a = a + dt * da
        e = e + dt * de

        A = push!(A, a)
        E = push!(E, e)
        H = push!(H, dt)

        total_time = total_time + dt
    end

    stop_reason = "flag_not_set"

    if total_time/seconds_per_year + evotime >= 1e11
        stop_reason = "out_of_time"
    end

    if a <= 1e4
        stop_reason = "merged"
    end

    # Solar Radii, Dimensionless
    return A / Solar_Radius, E, H, stop_reason
end

max_stars_repo_name: Krytic/Takahe

max_issues_repo_name: Krytic/Takahe

max_forks_repo_name: Krytic/Takahe


text:
type FiveConstantArms <: BanditEnvironment
end

function arms(env::FiveConstantArms, trial::Int64)
  [
   BernoulliArm(0.1),
   BernoulliArm(0.3),
   BernoulliArm(0.5),
   BernoulliArm(0.4),
   BernoulliArm(0.2)
   ]
end

function n_arms(env::FiveConstantArms)
  5
end

max_stars_repo_name: rmminusrslash/BanditsBook

max_issues_repo_name: rmminusrslash/BanditsBook

max_forks_repo_name: rmminusrslash/BanditsBook


text:
using GLVisualize, GLAbstraction
using FileIO, GeometryTypes, Reactive

if !isdefined(:runtests)
    window = glscreen()
    timesignal = loop(linspace(0f0,1f0,360))
end
description = """
Introductory example, explaining Signals, file loading and
how to transform a visualization.
"""

# loadasset is defined in GLVisualize like this:
# loadasset(path_segments...) = FileIO.load(assetpath(path_segments...))
# where assetpath just looks up the file in the asset folder
# You can load these (file types)[https://github.com/JuliaIO/FileIO.jl/blob/master/docs/registry.md]
mesh = loadasset("cat.obj")

# GLAbstraction.const_lift is an alias for Reactive.map, which also works for non
# signal arguments.
# Reactive.map takes a function and signals like the one created via `loop` (or just Signal(x))
# as an argument, applies the function to the signals whenever they update and
# returns a new signal.

rotation_angle  = const_lift(*, timesignal, 2f0*pi)

# the cat needs some rotation on the x axis to stand straight
# so we start off with a rotation of 90 degrees
start_rotation  = Signal(rotationmatrix_x(deg2rad(90f0)))
rotation        = map(rotationmatrix_y, rotation_angle)
final_rotation  = map(*, start_rotation, rotation)

# now we visualize the mesh and pass the rotation via the model keyword argument,
# which is short for the modelmatrix, which allows you to transform the visualization
# with any arbitrary transformation matrix
# You can create the most common transformation matrix with `translationmatrix(::Vec3f0)`,
# `rotationmatrix_x`/`y`/`z` (rotation around axis x,y,z), and `scalematrix(::Vec3f0)`

# the visualize function always only takes one argument, plus an optional style
# argument and then visualization dependant many keywords to customize the visualization.
# for all parameters Signals can be used and thus the visualization becomes animated

robj = visualize(mesh, model = final_rotation)

_view(robj, window)

if !isdefined(:runtests)
    renderloop(window)
end


text:
module LaTeXUnicode

using CSV
using DataFrames

export translationary

function build()
    # From http://milde.users.sourceforge.net/LUCR/Math/data/unimathsymbols.txt
    path = joinpath(dirname(dirname(pathof(LaTeXUnicode))), "data", "unimathsymbols.txt")
    df = DataFrame(CSV.File(path; comment="#", header = false, delim = '^'))
    rename!(df,
        [
            :code_point, :character, :latex, :class,
            :category, :requirements, :conflicts, :description
        ]
    )

    return df
end

const translationary = build()

end


text:
using FITSIO
using FITSIO.Libcfitsio

# The low-level API is now part of a submodule Libcfitsio

f001 = fits_open_file("f001a066.fits")
n = fits_get_hdrspace(f001)[1]; # => 128

for i = 1:n
   println(fits_read_keyn(f001,i))
end

max_stars_repo_path: Module 3/Chapter10/main-fitsio.jl

max_issues_repo_path: Module 3/Chapter10/main-fitsio.jl

max_forks_repo_path: Module 3/Chapter10/main-fitsio.jl


text:
"""Walk an AST and find all unique symbols.
"""
function find_symbols_in_ast(ex)
    # accumulate sub-expressions for recursion
    expressions = [ex]
    # accumulate found symbols
    symbols = []

    # since no Tail Call Optimization in Julia,
    # we write recursion in a while loop
    while length(expressions) > 0
        ex = pop!(expressions)
        first, rest = _find_symbols_in_ast(ex)
        if typeof(first) == Symbol
            # got value, no more recursion
            push!(symbols, first)
        end
        if ~isnothing(rest)
            # recur
            expressions = vcat(expressions, rest)
        end
    end
    return unique(symbols)
end

# https://docs.julialang.org/en/v1/base/base/#Keywords
RESERVED = [
    :baremodule, :begin, :break, :catch, :const, :continue, :do, :else,
    :elseif, :end, :export, :false, :finally, :for, :function, :global, :if,
    :import, :let, :local, :macro, :module, :quote, :return, :struct, :true,
    :try, :using, :while
]
"Return true if symbol could be assigned to a Thunk."
function maybe_thunk(symbol, ignore_symbols)
    # println("maybe_thunk: $symbol")
    if symbol in RESERVED
        ret = false
    elseif symbol in ignore_symbols
        ret = false
    elseif string(symbol)[1] == '.'
        # ie .+ is not a possible thunk
        ret = false
    else
        ret = true
    end
    # println("maybe_thunk: $symbol, $ret")
    ret
end

"""
We find all "assigned symbols", as in, everything except:
- anonymous functions :((x,y,z) -> x*x).args[1]
- generators :((x for x in 1:10))
- Reserved keywords
- ... others that we should catch...?
"""
function find_assigned_symbols_in_ast(ex)
    # accumulate sub-expressions for recursion
    # println("===== intitial call ========")
    expressions = [ex]
    # paired list of symbols to ignore per expression
    ignore_sym_per_ex = [Vector{Symbol}()]

    # accumulate found symbols
    symbols = Vector{Symbol}()

    while length(expressions) > 0
        # @show expressions
        # @show ignore_sym_per_ex
        # @show symbols
        ex = pop!(expressions)
        ignore_symbols = pop!(ignore_sym_per_ex)
        if (typeof(ex) == Expr)
            if ex.head == :->
                # anonymous function, so need to ignore unassigned symbols
                ignore_symbols = vcat(ignore_symbols, find_symbols_in_ast(ex.args[1]))
                first, rest = _find_symbols_in_ast(ex.args[2])
            elseif ex.head == :generator
                # right hand of generator, first arg is the assignment
                ignore_symbols = vcat(ignore_symbols, ex.args[2].args[1])
                first, rest = _find_symbols_in_ast(ex)
            elseif ex.head == :kw
                # ignore keyword argument right side
                push!(ignore_symbols, ex.args[1])
                first, rest = _find_symbols_in_ast(ex.args[2])
            else
                first, rest = _find_symbols_in_ast(ex)
            end
        else
            first, rest = _find_symbols_in_ast(ex)
        end

        # 0 or 1 of next if blocks will be executed

        if typeof(first) == Symbol
            # got value, no more recursion
            if maybe_thunk(first, ignore_symbols)
                push!(symbols, first)
            end
        end

        if ~isnothing(rest)
            # recur
            expressions = vcat(expressions, rest)
            ignore_sym_per_ex = vcat(ignore_sym_per_ex,
                repeat([ignore_symbols], length(rest)))
        end
    end
    unique(symbols)
end

"""
Return either (nothing, List[Expr]), (Symbol, nothing), or (nothing, nothing).

Helper function following `cons` and `nil` pattern from Lisp.
"""
function _find_symbols_in_ast(ex)
    if typeof(ex) == Expr
        head, args = ex.head, ex.args
        return nothing, args
    elseif typeof(ex) == Symbol
        return ex, nothing
    else
        return nothing, nothing
    end
end

"""
Rewrite expression for lazy evaluation using thunks.
Not intended for public usage.
"""
function thunkify(ex, type=Thunk)
    if ex.head == :block
        t = thunkify_block(ex, type)
    else
        t = _thunkify(ex, type)
    end
    t
end

"Rewrite a begin...end block to be lazy."
function thunkify_block(ex, type=Thunk)
    @assert ex.head == :block "unexpected head ($(ex.head) != :block) for: $ex"
    args = map(t -> _thunkify(t,type), ex.args)
    Expr(:block, args...)
end

"Assign left of `=` to a thunk."
function _thunkify_eq(ex, type=Thunk)
    @assert ex.head == :(=) "unexpected head ($(ex.head) != :(=)) for: $ex"
    l = ex.args[1]
    if typeof(ex.args[2]) == Expr
        # recursion...
        r = thunkify(ex.args[2], type)
    else
        r = ex.args[2]
    end
    :($l = $r)
end

"""Actually create the Thunk struct.

Wrap expression in anonymous function, with each assigned symbol as an arg.
"""
function _thunkify_expr(ex, type=Thunk)
    symbols = find_assigned_symbols_in_ast(ex)
    vars = "("
    for sym in symbols
        vars *= "$sym,"
    end
    vars *= ")"
    vars = Meta.parse(vars)
    wrapped = :($vars -> $ex)
    if type == Thunk
        Expr(:call, Thunk, wrapped, vars)
    elseif type == Unevaluated
        vars = :(() -> $vars)
        Expr(:call, Unevaluated, wrapped, vars)
    elseif type == Reversible
        Expr(:call, Reversible, wrapped, vars)
    else
        error("_thunkify_expr not implemented for type: $type")
    end
end

"Helper function to dispatch based on Expr head."
function _thunkify(ex, type=Thunk)
    if typeof(ex) == Expr
        if ex.head == :(=)
            _thunkify_eq(ex, type)
        else
            _thunkify_expr(ex, type)
        end
    else
        # e.g. for LineNumberNode in a block
        ex
    end
end

"""
    @lazy <Expr>

Lazy evaluation of arbitrary expressions.
```julia
@lazy x+y
@lazy x = f(y)
@lazy begin
    a = 1
    b = 2
    c = 3
    abc = sum([a,b,c])
end
reify(abc)
# output
6
```
"""
macro lazy(ex)
    new = thunkify(ex, Thunk)
    esc(new)
end

"""
    @noeval <Expr>

Lazy evaluation of arbitrary expressions that does not check for variable
definition. This is safe to use if the assumption of referential transprency
holds, and should be used with extreme caution otherwise, e.g. if referenced
variables might be mutated. Note that evaluation occurs in the scope of where
`reify` is called.

```jldoctest
@noeval begin
    # will not error
    not_used = not_defined + 1
    a = 1
    b = 2
    c = 3
    abc = sum([a,b,c])
end
c = -1
reify(abc)
# output
2
```
"""
macro noeval(ex)
    new = thunkify(ex, Unevaluated)
    esc(new)
end

"""
    @reversible <Expr>

Lazy evaluation of arbitrary expressions that allows for computation to be
undone.

```jldoctest
@reversible begin
    a = 1
    b = 2
    c = 3
    abc = sum([a,b,c])
end
@assert abc.evaluated == false
reify(abc)
@assert abc.evaluated == true
undo(abc)
@assert abc.evaluated == false
reify(abc)
@assert abc.evaluated == true
# output

```
"""
macro reversible(ex)
    new = thunkify(ex, Reversible)
    esc(new)
end


text:
module Curvature

include("functions.jl")

end


text:
using Test, TestPackage

@test f(2) == 4

max_stars_repo_name: SciML/TestPackage

max_issues_repo_name: SciML/TestPackage

max_forks_repo_name: SciML/TestPackage


text:
## Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra
using StaticArrays
using HybridArrays

# %%code cell

function finklet(j, l)
    # -------- set filters
    jrad = 7-j
    dθ = π/8        # 8 angular bins hardwired
    θ_l = dθ*l
    # -------- define coordinates
    nx = 256
    xbox = LinRange(-nx/2, nx/2-1 , nx)
    # make a 256x256 grid of X
    sx = xbox' .* ones(nx)
    sy = ones(nx)' .* xbox
    r  = sqrt.((sx).^2 + (sy).^2)
    θ  = mod.(atan.(sy, sx).+π .-θ_l,2*π)
    nozeros = r .> 0
    logr = log2.(r[nozeros])
    r[nozeros] = logr
    # -------- in Fourier plane, envelope of psi_j,l
    mask = (abs.(θ.-π).<= dθ) .& (abs.(r.-jrad) .<= 1)
    # -------- angular part
    ang = cos.((θ.-π).*4)
    # -------- radial part
    rad = cos.((r.-jrad).*π./2)
    psi = mask.*ang.*rad             #mask times angular part times radial part
    return psi
end

function fink_filter_bank(J,L)
    fink_filter = Array{Float64, 4}(undef, 256, 256, J, L)
    for l = 1:L
        for j = 1:J
            fink_filter[:,:,j,l]=fftshift(finklet(j-1,l-1))
        end
    end
    return fink_filter
end

function had!(A,B)
    m,n = size(A)
    @assert (m,n) == size(B)
    for j in 1:n
       for i in 1:m
         @inbounds A[i,j] *= B[i,j]
       end
    end
    return A
end

fink_filter_set = fink_filter_bank(8,8)

test_img = zeros(256,256)
copyto!(test_img,fink_filter_set[:,:,1,3])

function DHC(image, filter_set; norm_on = 1, coeff_12_on =1, coeff_20_on = 1)
    FFTW.set_num_threads(2)
    (Nx, Ny) = size(image)

    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []

    ## Coeff at (1,0)
    #takes care of image normalization
    S0 = zeros(2)
    if norm_on == 1
        S0[1] = mean(image)
        norm_im = image.-S0[1]
        S0[2] = sum(abs2.(norm_im))/(Nx*Ny)
        norm_im = norm_im./sqrt(Nx*Ny*S0[2])
    else
        norm_im = image
    end

    append!(out_coeff,S0)

    ##Coeff at (1,1)
    #store the fft of images and im*filters for later
    im_fd_0 = fft(norm_im)

    S1 = zeros((J,L))
    im_fd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    @views for l = 1:L
        for j = 1:J
            im_fd_0_1[:,:,j,l] .= im_fd_0
        end
    end

    if coeff_20_on == 1
        im_rd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    end

    @views for l = 1:L
        for j = 1:J
            had!(im_fd_0_1[:,:,j,l],filter_set[:,:,j,l]) #wavelet already in fft domain not shifted
            S1[j,l]+=sum(abs2.(im_fd_0_1[:,:,j,l])) #normalization choice arb to make order unity
            if coeff_20_on == 1
                im_rd_0_1[:,:,j,l] .= ifft(im_fd_0_1[:,:,j,l])
            end
        end
    end
    append!(out_coeff,S1)

    ##Coeff at (2,0)
    if coeff_20_on == 1
        S20 = zeros((J, L, J, L))

        Atmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,Nx,Ny))
        Btmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,Nx,Ny))
        interm = 0
        @views for l2 = 1:L
            for j2 = 1:J
                copyto!(Btmp, im_rd_0_1[:,:,j2,l2])
                for l1 = 1:L
                    for j1  = 1:J
                        copyto!(Atmp, im_rd_0_1[:,:,j1,l1])
                        S20[j1,l1,j2,l2] += sum(abs.(Atmp.*Btmp))
                    end
                end
            end
        end

        append!(out_coeff,S20)
    end

    ##Coeff at (1,2) level
    if coeff_12_on == 1
        S12 = Array{Float64,4}(undef, J, J, L, L)
        for j1 = 1:J
            for l1 = 1:L
                not_rot_im = fftshift(im_fd_0_1[j1,l1,:,:])
                for l2 = 1:L
                    for j2 = 1:J
                        S12[j1,j2,l1,l2] += sum(abs.(not_rot_im.*fftshift(im_fd_0_1[j2,l2,:,:])))
                    end
                end
            end
        end
    append!(out_coeff,S12)
    end
    return out_coeff
end

function dual_conj(A,B)
    @. abs(A.*B)
end

temp = DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

@time DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Profile.clear()
@profile DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.samples = 5
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 120
@benchmark DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

##Trying out static speed ups

function dual_conj_pre(A,B)
    Atmp = SMatrix{256,256}(A)
    Btmp = SMatrix{256,256}(B)
    @. abs(Atmp.*Btmp)
end

@time dual_conj_pre(test_img,test_img)

test_img_s = @SMatrix rand(4,4)

function mult_wrapper(A,B)
    sum(abs.(A.*B))
end

@time mult_wrapper(test_img_s,test_img_s)

test_img_2 = rand(4,4)

@time mult_wrapper(test_img_2,test_img_2)

A = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(randn(256,256))

mult_wrapper(A,A)

@benchmark mult_wrapper(A,A)

A = randn(256,256)

@benchmark mult_wrapper(A,A)

function mult_wrapper(A,B)
    A = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(randn(256,256))
    B = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(randn(256,256))
    sum(abs.(A.*B))
end

@benchmark mult_wrapper(A,A)

Btmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
copyto!(Btmp, fink_filter_set[:,:,5,5])

## Suppose I was willing to reshape in DHC_ISO

function structure_test(J,L)
    S20 = zeros((J, L, J, L))
    @views for l2 = 1:L
        for j2 = 1:J
            for l1 = 1:L
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = rand()
                end
            end
        end
    end
end

@benchmark structure_test(8,8)

function structure_test(J,L)
    S20 = zeros(J*L*J*L)
    @views for l2 = 1:L
        for j2 = 1:J
            for l1 = 1:L
                for j1  = 1:J
                    S20[(j1-1)*J+(l1-1)*L+(j2-1)*J+l2] = rand()
                end
            end
        end
    end
end

@benchmark structure_test(8,8)

function structure_test(J,L)
    S20 = zeros(J*L*J*L)
    Atmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
    Btmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
    @views for l2 = 1:L
        for j2 = 1:J
            copyto!(Btmp, fink_filter_set[:,:,j2,l2])
            for l1 = 1:L
                for j1  = 1:J
                    copyto!(Atmp, fink_filter_set[:,:,j1,l1])
                    S20[(j1-1)*J+(l1-1)*L+(j2-1)*J+l2] = sum(abs.(Atmp.*Btmp))
                end
            end
        end
    end
end

structure_test(8,8)

@time structure_test(8,8)

@benchmark structure_test(8,8)

function structure_test(J,L)
    S20 = zeros(J,L,J,L)
    Atmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
    Btmp = HybridArray{Tuple{256,StaticArrays.Dynamic()}}(zeros(ComplexF64,256,256))
    @views for l2 = 1:L
        for j2 = 1:J
            copyto!(Btmp, fink_filter_set[:,:,j2,l2])
            for l1 = 1:L
                for j1  = 1:J
                    copyto!(Atmp, fink_filter_set[:,:,j1,l1])
                    S20[j1,l1,j2,l2] = sum(abs.(Atmp.*Btmp))
                end
            end
        end
    end
end

structure_test(8,8)

@benchmark structure_test(8,8)

## Using BLAS

function structure_test(J,L,test_img_1)
    S20 = zeros((J, L, J, L))
    @views for l2 = 1:L
        for j2 = 1:J
            for l1 = 1:L
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = sum(abs.(test_img_1))
                end
            end
        end
    end
end

structure_test(8,8,rand(256,256))

@benchmark structure_test(8,8,rand(256,256))

function structure_test(J,L,test_img_1)
    S20 = zeros((J, L, J, L))
    @views for l2 = 1:L
        for j2 = 1:J
            for l1 = 1:L
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = BLAS.asum(test_img_1)
                end
            end
        end
    end
end

structure_test(8,8,rand(256,256))

@benchmark structure_test(8,8,rand(256,256))

# %% codecell

function DHC(image, filter_set; norm_on = 1, coeff_12_on =1, coeff_20_on = 1)
    FFTW.set_num_threads(2)
    (Nx, Ny) = size(image)

    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []

    ## Coeff at (1,0)
    #takes care of image normalization
    S0 = zeros(2)
    if norm_on == 1
        S0[1] = mean(image)
        norm_im = image.-S0[1]
        S0[2] = sum(abs2.(norm_im))/(Nx*Ny)
        norm_im = norm_im./sqrt(Nx*Ny*S0[2])
    else
        norm_im = image
    end

    append!(out_coeff,S0)

    ##Coeff at (1,1)
    #store the fft of images and im*filters for later
    im_fd_0 = fft(norm_im)

    S1 = zeros((J,L))
    im_fd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    @views for l = 1:L
        for j = 1:J
            im_fd_0_1[:,:,j,l] .= im_fd_0
        end
    end

    if coeff_20_on == 1
        im_rd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    end

    @views for l = 1:L
        for j = 1:J
            had!(im_fd_0_1[:,:,j,l],filter_set[:,:,j,l]) #wavelet already in fft domain not shifted
            S1[j,l]+=sum(abs2.(im_fd_0_1[:,:,j,l])) #normalization choice arb to make order unity
            if coeff_20_on == 1
                im_rd_0_1[:,:,j,l] .= ifft(im_fd_0_1[:,:,j,l])
            end
        end
    end
    append!(out_coeff,S1)

    ##Coeff at (2,0)
    if coeff_20_on == 1
        S20 = zeros((J, L, J, L))

        Atmp = zeros(ComplexF64,Nx,Ny)
        Btmp = zeros(ComplexF64,Nx,Ny)
        interm = 0
        @views for l2 = 1:L
            for j2 = 1:J
                copyto!(Btmp, im_rd_0_1[:,:,j2,l2])
                for l1 = 1:L
                    for j1  = 1:J
                        copyto!(Atmp, im_rd_0_1[:,:,j1,l1])
                        S20[j1,l1,j2,l2] += BLAS.asum(Atmp.*Btmp)
                    end
                end
            end
        end

        append!(out_coeff,S20)
    end

    ##Coeff at (1,2) level
    if coeff_12_on == 1
        S12 = Array{Float64,4}(undef, J, J, L, L)
        for j1 = 1:J
            for l1 = 1:L
                not_rot_im = fftshift(im_fd_0_1[j1,l1,:,:])
                for l2 = 1:L
                    for j2 = 1:J
                        S12[j1,j2,l1,l2] += sum(abs.(not_rot_im.*fftshift(im_fd_0_1[j2,l2,:,:])))
                    end
                end
            end
        end
    append!(out_coeff,S12)
    end
    return out_coeff
end

temp = DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

@time DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Profile.clear()
@profile DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.samples = 5
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 120
@benchmark DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

## Factor of 2 from BLAS asum
# Need to do something for .*

BLAS.dot(test_img,test_img)

dot(test_img,test_img)

At = [1 2 3; 4 5 6; 7 8 9]
Bt = [1 2 3; 4 5 6; 7 8 9]
dot(At,Bt)

1^2+2^2+3^2+4^2+5^2+6^2+7^2+8^2+9^2

## Ok, so asum was a cool find, but dot(abs) is prob better

function structure_test1(J,L,test_img)
    S1 = zeros((J,L))
    @views for l = 1:L
        for j = 1:J
            S1[j,l]+=BLAS.dot(test_img,test_img) #normalization choice arb to make order unity
        end
    end
end

structure_test1(8,8,rand(256,256))

@time structure_test1(8,8,rand(256,256))

@benchmark structure_test1(8,8,rand(256,256))

function structure_test2(J,L,test_img)
    S1 = zeros((J,L))
    @views for l = 1:L
        for j = 1:J
            S1[j,l]+=sum(abs2.(test_img))#normalization choice arb to make order unity
        end
    end
end

structure_test2(8,8,rand(256,256))

@time structure_test2(8,8,rand(256,256))

@benchmark structure_test2(8,8,rand(256,256))

## Yeah, ok, so that is much better. Can we take abs as we allocate?

function structure_test(J,L,test_img_1)
    S20 = zeros((J, L, J, L))
    Atmp = zeros(Float64,256,256)
    Btmp = zeros(Float64,256,256)
    @views for l2 = 1:L
        copyto!(Btmp, abs.(test_img_1))
        for j2 = 1:J
            for l1  = 1:L
                copyto!(Atmp, abs.(test_img_1))
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = BLAS.dot(Atmp,Btmp)
                end
            end
        end
    end
end

structure_test(8,8,rand(256,256))

@time structure_test(8,8,rand(256,256))

@benchmark structure_test(8,8,rand(256,256))

function structure_test(J,L,test_img_1)
    S20 = zeros((J, L, J, L))
    Atmp = zeros(ComplexF64,256,256)
    Btmp = zeros(ComplexF64,256,256)
    @views for l2 = 1:L
        copyto!(Btmp, test_img_1)
        for j2 = 1:J
            for l1  = 1:L
                copyto!(Atmp, test_img_1)
                for j1  = 1:J
                    S20[j1,l1,j2,l2] = BLAS.asum(Atmp.*Btmp)
                end
            end
        end
    end
end

structure_test(8,8,rand(256,256))

@time structure_test(8,8,rand(256,256))

@benchmark structure_test(8,8,rand(256,256))

##
function DHC(image, filter_set; norm_on = 1, coeff_12_on =1, coeff_20_on = 1)
    FFTW.set_num_threads(2)
    (Nx, Ny) = size(image)

    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []

    ## Coeff at (1,0)
    #takes care of image normalization
    S0 = zeros(2)
    if norm_on == 1
        S0[1] = mean(image)
        norm_im = image.-S0[1]
        S0[2] = BLAS.dot(norm_im,norm_im)/(Nx*Ny)
        norm_im = norm_im./sqrt(Nx*Ny*S0[2])
    else
        norm_im = image
    end

    append!(out_coeff,S0)

    ##Coeff at (1,1)
    #store the fft of images and im*filters for later
    im_fd_0 = fft(norm_im)

    S1 = zeros((J,L))
    im_fd_0_1 = Array{ComplexF64,4}(undef, Nx, Ny, J, L)
    @views for l = 1:L
        for j = 1:J
            im_fd_0_1[:,:,j,l] .= im_fd_0
        end
    end

    if coeff_20_on == 1
        im_rd_0_1 = Array{Float64,4}(undef, Nx, Ny, J, L)
    end

    Atmp = zeros(ComplexF64,Nx,Ny)
    Btmp = zeros(Float64,Nx,Ny)
    Ctmp = zeros(Float64,Nx,Ny)
    @views for l = 1:L
        for j = 1:J
            had!(im_fd_0_1[:,:,j,l],filter_set[:,:,j,l]) #wavelet already in fft domain not shifted
            Btmp .= abs.(im_fd_0_1[:,:,j,l])
            S1[j,l]+=BLAS.dot(Btmp,Btmp) #normalization choice arb to make order unity
            if coeff_20_on == 1
                Atmp .= ifft(im_fd_0_1[:,:,j,l])
                im_rd_0_1[:,:,j,l] .= abs.(Atmp)
            end
        end
    end
    append!(out_coeff,S1)

    ##Coeff at (2,0)
    if coeff_20_on == 1
        S20 = zeros(Float64,J, L, J, L)
        @views for l2 = 1:L
            for j2 = 1:J
                copyto!(Btmp, im_rd_0_1[:,:,j2,l2])
                for l1 = 1:L
                    for j1  = 1:J
                        copyto!(Ctmp, im_rd_0_1[:,:,j1,l1])
                        S20[j1,l1,j2,l2] += BLAS.dot(Btmp,Ctmp)
                    end
                end
            end
        end

        append!(out_coeff,S20)
    end

    ##Coeff at (1,2) level
    if coeff_12_on == 1
        S12 = zeros(Float64,J, L, J, L)
        for l2 = 1:L
            for j2 = 1:J
                not_rot_im = fftshift(im_fd_0_1[:,:,j1,l1])
                for l1 = 1:L
                    for j1 = 1:J
                        S12[j1,l1,j2,l2] += sum(abs.(not_rot_im.*fftshift(im_fd_0_1[:,:,j2,l2])))
                    end
                end
            end
        end
    append!(out_coeff,S12)
    end
    return out_coeff
end

temp = DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

@time DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Profile.clear()
@profile DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.samples = 5
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 120
@benchmark DHC(test_img,fink_filter_set,coeff_12_on =0, coeff_20_on = 1)

## Starting to test the S12 coeff

function fftshift_wrap()
    A = rand(256,256)
    fftshift(A)
end

@time fftshift_wrap()

@benchmark fftshift_wrap()

function fftshift_wrap()
    A = rand(256,256)
end

@time fftshift_wrap()

@benchmark fftshift_wrap()

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC