text:
using DataStructures
using Random

#############
# Trimmed SDD types and structs
#############

"Root of the trimmed SDD manager node hierarchy"
abstract type TrimMgrNode <: SddMgrNode end

# alias structured logical nodes with a trimmed sdd manager vtree
const TrimTrue = SddTrueNode{TrimMgrNode}
const trimtrue = TrimTrue()
const TrimFalse = SddFalseNode{TrimMgrNode}
const trimfalse = TrimFalse()

const TrimConstant = SddConstantNode{TrimMgrNode}
const Trim⋁ = Sdd⋁Node{TrimMgrNode}
const Trim⋀ = Sdd⋀Node{TrimMgrNode}
const TrimNode = SddNode{<:TrimMgrNode} # would this be better?: Union{TrimTrue,TrimFalse,TrimConstant,Trim⋁,Trim⋀}

# alias SDD terminology
const Element = Tuple{TrimNode,TrimNode}
Element(prime::TrimNode, sub::TrimNode)::Element = (prime, sub)
const XYPartition = Set{Element}

const Unique⋁Cache = Dict{XYPartition,Trim⋁}
const ApplyCache = Dict{Tuple{TrimNode,TrimNode},TrimNode}

mutable struct TrimSddMgrInnerNode <: TrimMgrNode

    left::TrimMgrNode
    right::TrimMgrNode

    parent::Union{TrimSddMgrInnerNode, Nothing}

    variables::BitSet
    unique⋁cache::Unique⋁Cache

    conjoin_cache::ApplyCache
    disjoin_cache::ApplyCache

    TrimSddMgrInnerNode(left::TrimMgrNode, right::TrimMgrNode) = begin
        # @assert isempty(intersect(variables(left), variables(right)))
        this = new(left, right,
            nothing,
            union(variables(left), variables(right)),
            Unique⋁Cache(), ApplyCache(), ApplyCache())
        left.parent = this
        right.parent = this
        this
    end

end

mutable struct TrimSddMgrLeafNode <: TrimMgrNode

    var::Var
    parent::Union{TrimSddMgrInnerNode, Nothing}

    positive_literal::SddLiteralNode{TrimSddMgrLeafNode} # aka TrimLiteral
    negative_literal::SddLiteralNode{TrimSddMgrLeafNode} # aka TrimLiteral

    TrimSddMgrLeafNode(v::Var) = begin
        this = new(v, nothing)
        this.positive_literal = SddLiteralNode(var2lit(v), this)
        this.negative_literal = SddLiteralNode(-var2lit(v), this)
        this
    end

end

const TrimLiteral = SddLiteralNode{TrimSddMgrLeafNode}

const TrimSddMgr = AbstractVector{<:TrimMgrNode}
const TrimSdd = AbstractVector{<:TrimNode}

Base.eltype(::Type{TrimSdd}) = TrimNode

#####################
# Constructor
#####################

TrimMgrNode(v::Var) = TrimSddMgrLeafNode(v)
TrimMgrNode(left::TrimMgrNode, right::TrimMgrNode) = TrimSddMgrInnerNode(left, right)

#####################
# Traits
#####################

@inline NodeType(::TrimSddMgrLeafNode) = Leaf()
@inline NodeType(::TrimSddMgrInnerNode) = Inner()

#####################
# Methods
#####################

@inline children(n::TrimSddMgrInnerNode) = [n.left, n.right]

@inline variables(n::TrimSddMgrLeafNode)::BitSet = BitSet(n.var)
@inline variables(n::TrimSddMgrInnerNode)::BitSet = n.variables

import ..Utils: parent, descends_from, lca # make available for extension

@inline prime(e::Element) = e[1]
@inline sub(e::Element) = e[2]

@inline parent(n::TrimMgrNode)::Union{TrimSddMgrInnerNode, Nothing} = n.parent

@inline pointer_sort(s,t) = (hash(s) <= hash(t)) ? (s,t) : (t,s)

#TODO replace this by a bitset subset check on the set of variables
@inline descends_from(n::TrimNode, m::TrimNode) = descends_from(vtree(n), vtree(m))
@inline descends_from(::TrimMgrNode, ::TrimSddMgrLeafNode) = false
@inline descends_from(n::TrimMgrNode, m::TrimSddMgrInnerNode) = variables(n) ⊆ variables(m)

@inline descends_left_from(n::TrimNode, m::TrimNode)::Bool = descends_left_from(vtree(n), vtree(m))
@inline descends_left_from(n::TrimMgrNode, m::TrimSddMgrInnerNode)::Bool = (n === m.left) || descends_from(n, m.left)
@inline descends_left_from(::TrimMgrNode, ::TrimSddMgrLeafNode)::Bool = false

@inline descends_right_from(n::TrimNode, m::TrimNode)::Bool = descends_right_from(vtree(n), vtree(m))
@inline descends_right_from(n::TrimMgrNode, m::TrimSddMgrInnerNode)::Bool = (n === m.right) || descends_from(n,m.right)
@inline descends_right_from(::TrimMgrNode, ::TrimSddMgrLeafNode)::Bool = false

function lca(xy::XYPartition)::TrimSddMgrInnerNode
    # @assert !isempty(xy)
    # @assert all(e -> (prime(e) !== trimfalse), xy)
    element_vtrees = [parentlca(prime(e),sub(e)) for e in xy]
    return lca(element_vtrees...)
end

parentlca(p::TrimNode, s::TrimNode)::TrimSddMgrInnerNode =
    lca(parent(vtree(p)), parent(vtree(s)))
parentlca(p::TrimNode, ::TrimConstant)::TrimSddMgrInnerNode =
    parent(vtree(p))
parentlca(::TrimConstant, s::TrimNode)::TrimSddMgrInnerNode =
    parent(vtree(s))
parentlca(p::TrimConstant, s::TrimConstant)::TrimSddMgrInnerNode =
    error("This XY partition should have been trimmed to remove ($p,$s)!")

"""
Get the canonical compilation of the given XY Partition
"""
function canonicalize(xy::XYPartition)::TrimNode
    # @assert !isempty(xy)
    return canonicalize_compressed(compress(xy))
end

"""
Compress a given XY Partition (merge elements with identical subs)
"""
function compress(xy::XYPartition)::XYPartition
    # @assert !isempty(xy)
    sub2elems = groupby(e -> sub(e), xy)
    #TODO avoid making a new partition if existing one is unchanged
    compressed_elements = XYPartition()
    for (subnode,elements) in sub2elems
        primenode = mapreduce(e -> prime(e), (p1, p2) -> disjoin(p1, p2), elements)
        push!(compressed_elements, (primenode, subnode))
    end
    return compressed_elements
end

"""
Get the canonical compilation of the given compressed XY Partition
"""
function canonicalize_compressed(xy::XYPartition)::TrimNode
    # @assert !isempty(xy)
    # trim
    if length(xy) == 1 && (prime(first(xy)) === trimtrue)
        return sub(first(xy))
    elseif length(xy) == 2
        l = [xy...]
        if (sub(l[1]) === trimtrue) && (sub(l[2]) === trimfalse)
            return prime(l[1])
        elseif (sub(l[2]) === trimtrue) && (sub(l[1]) === trimfalse)
            return prime(l[2])
        end
    end
    # get unique node representation
    return unique⋁(xy)
end

@inline function remove_false_primes(xy)
    return filter(e -> (prime(e) !== trimfalse), xy)
end

"""
Construct a unique decision gate for the given vtree
"""
function unique⋁(xy::XYPartition, mgr::TrimSddMgrInnerNode = lca(xy))::Trim⋁
    #TODO add finalization trigger to remove from the cache when the node is gc'ed + weak value reference
    get!(mgr.unique⋁cache, xy) do
        node = Trim⋁(xy2ands(xy, mgr), mgr)
        not_xy = negate(xy)
        not_node = Trim⋁(xy2ands(not_xy, mgr), mgr, node)
        node.negation = not_node
        mgr.unique⋁cache[not_xy] = not_node
        node
    end
end

@inline negate(xy::XYPartition) = XYPartition([Element(prime(e), !sub(e)) for e in xy])
@inline xy2ands(xy::XYPartition, mgr::TrimSddMgrInnerNode) = [Trim⋀(prime(e), sub(e), mgr) for e in xy]

"""
Compile a given variable, literal, or constant
"""
compile(mgr::TrimSddMgr, x::Union{Var,Lit})::TrimLiteral = compile(mgr[end], x)

function compile(n::TrimMgrNode, v::Var)::TrimLiteral
    compile(n,var2lit(v))
end

function compile(n::TrimSddMgrLeafNode, l::Lit)::TrimLiteral
    # @assert n.var == lit2var(l)
    if l>0 # positive literal
        n.positive_literal
    else
        n.negative_literal
    end
end

function compile(n::TrimSddMgrInnerNode, l::Lit)::TrimLiteral
    if lit2var(l) in variables(n.left)
        compile(n.left, l)
    elseif lit2var(l) in variables(n.right)
        compile(n.right, l)
    else
        error("$v is not contained in this vtree")
    end
end

function compile(constant::Bool)::TrimConstant
    if constant == true
        trimtrue
    else
        trimfalse
    end
end

"""
Conjoin two SDDs
"""
@inline conjoin(::TrimFalse, ::TrimTrue)::TrimFalse = trimfalse
@inline conjoin(::TrimTrue, ::TrimFalse)::TrimFalse = trimfalse
@inline conjoin(s::TrimNode, ::TrimTrue)::TrimNode = s
@inline conjoin(::TrimNode, ::TrimFalse)::TrimFalse = trimfalse
@inline conjoin(::TrimTrue, s::TrimNode)::TrimNode = s
@inline conjoin(::TrimFalse, ::TrimNode)::TrimFalse = trimfalse
@inline conjoin(::TrimTrue, s::TrimTrue)::TrimNode = trimtrue
@inline conjoin(::TrimFalse, s::TrimFalse)::TrimNode = trimfalse

function conjoin(s::TrimLiteral, t::TrimLiteral)::TrimNode
    if vtree(s) == vtree(t)
        (s === t) ? s : trimfalse
    else
        conjoin_indep(s,t)
    end
end

function conjoin(s::TrimNode, t::TrimNode)::TrimNode
    if vtree(s) == vtree(t)
        conjoin_cartesian(t,s)
    elseif descends_from(s,t)
        conjoin_descendent(s,t)
    elseif descends_from(t,s)
        conjoin_descendent(t,s)
    else
        conjoin_indep(s,t)
    end
end

"""
Conjoin two SDDs when they respect the same vtree node
"""
function conjoin_cartesian(n1::TrimNode, n2::TrimNode)::TrimNode
    if n1 === n2
        return n1
    elseif n1 == !n2
        return trimfalse
    end
    (n1,n2) = pointer_sort(n1,n2)

    get!(vtree(n1).conjoin_cache, (n1,n2)) do
        elems_prod = Vector{Element}()
        elems1 = copy(children(n1))
        elems2 = copy(children(n2))
        for e1 in elems1
            for e2 in elems2
                if prime(e1) === prime(e2)
                    push!(elems_prod, Element(prime(e1), conjoin(sub(e1),sub(e2))))
                    filter!(e -> prime(e) !== prime(e1), elems1)
                    filter!(e -> prime(e) !== prime(e1), elems2)
                    break # go to next e1
                elseif prime(e1) === !prime(e2)
                    for e3 in elems2
                        if e3 !== e2
                            push!(elems_prod, Element(prime(e3), conjoin(sub(e3),sub(e1))))
                        end
                    end
                    for e4 in elems1
                        if e4 !== e1
                            push!(elems_prod, Element(prime(e4), conjoin(sub(e4),sub(e2))))
                        end
                    end
                    filter!(e -> prime(e) !== prime(e1), elems1)
                    filter!(e -> prime(e) !== prime(e2), elems2)
                    break # go to next e1
                end
            end
        end
        product = vec([(e1,e2) for e1 in elems1, e2 in elems2])
        #TODO sort product by probability of subsumes
        while !isempty(product)
            (e1, e2) = pop!(product)
            newprime = conjoin(prime(e1),prime(e2))
            if newprime != trimfalse
                push!(elems_prod, Element(newprime, conjoin(sub(e1),sub(e2))))
            end
            if newprime === prime(e1)
                # p1 |= p2 and therefore p1 will be mutex with all other p2-primes
                filter!(p -> prime(p[1]) !== prime(e1), product)
            elseif newprime === prime(e2)
                # p2 |= p1 and therefore p2 will be mutex with all other p1-primes
                filter!(p -> prime(p[2]) !== prime(e2), product)
            end
        end
        canonicalize(XYPartition(elems_prod))
    end
end

"""
Conjoin two SDDs when one descends from the other
"""
function conjoin_descendent(d::TrimNode, n::TrimNode)::TrimNode
    get!(vtree(n).conjoin_cache, (d,n)) do
        if descends_left_from(d, n)
            elements = Element[Element(conjoin(prime(e),d), sub(e)) for e in children(n)]
            elements = remove_false_primes(elements)
            push!(elements, Element(!d, trimfalse))
        else
            # @assert descends_right_from(d, n)
            elements = Element[Element(prime(e),conjoin(sub(e),d)) for e in children(n)]
        end
        #TODO are there cases where we don't need all of compress-trim-unique?
        canonicalize(XYPartition(elements))
    end
end

"""
Conjoin two SDDs in separate parts of the vtree
"""
function conjoin_indep(s::TrimNode, t::TrimNode)::Trim⋁
    # @assert GateType(s)!=ConstantGate() && GateType(t)!=ConstantGate()
    mgr = parentlca(s,t)
    # @assert vtree(s) != mgr && vtree(t) != mgr
    (s,t) = pointer_sort(s,t)
    get!(mgr.conjoin_cache, (s,t)) do
        if descends_left_from(vtree(s), mgr)
            # @assert descends_right_from(vtree(t), mgr)
            elements = Element[Element(s,t),Element(!s,trimfalse)]
        else
            # @assert descends_left_from(vtree(t), mgr)
            # @assert descends_right_from(vtree(s), mgr)
            elements = Element[Element(t,s),Element(!t,trimfalse)]
        end
        # TODO: the XY partition must already be compressed and trimmed
        unique⋁(XYPartition(elements), mgr)
    end
end

@inline Base.:&(s,t) = conjoin(s,t)

"""
Disjoin two SDDs
"""
@inline disjoin(::TrimFalse, ::TrimTrue)::TrimTrue = trimtrue
@inline disjoin(::TrimTrue, ::TrimFalse)::TrimTrue = trimtrue
@inline disjoin(::TrimNode, ::TrimTrue)::TrimTrue = trimtrue
@inline disjoin(s::TrimNode, ::TrimFalse)::TrimNode = s
@inline disjoin(::TrimTrue, ::TrimNode)::TrimTrue = trimtrue
@inline disjoin(::TrimFalse, s::TrimNode)::TrimNode = s
@inline disjoin(::TrimTrue, s::TrimTrue)::TrimNode = trimtrue
@inline disjoin(::TrimFalse, s::TrimFalse)::TrimNode = trimfalse

function disjoin(s::TrimLiteral, t::TrimLiteral)::TrimNode
    if vtree(s) == vtree(t)
        (s === t) ? s : trimtrue
    else
        disjoin_indep(s,t)
    end
end

function disjoin(s::TrimNode, t::TrimNode)::TrimNode
    if vtree(s) == vtree(t)
        disjoin_cartesian(t,s)
    elseif descends_from(s,t)
        disjoin_descendent(s,t)
    elseif descends_from(t,s)
        disjoin_descendent(t,s)
    else
        disjoin_indep(s,t)
    end
end

"""
Disjoin two SDDs when they respect the same vtree node
"""
function disjoin_cartesian(n1::TrimNode, n2::TrimNode)::TrimNode
    if n1 === n2
        return n1
    elseif n1 == !n2
        return trimtrue
    end
    (n1,n2) = pointer_sort(n1,n2)

    get!(vtree(n1).disjoin_cache, (n1,n2)) do
        elems_prod = Vector{Element}()
        elems1 = copy(children(n1))
        elems2 = copy(children(n2))
        for e1 in elems1
            for e2 in elems2
                if prime(e1) === prime(e2)
                    push!(elems_prod, Element(prime(e1), disjoin(sub(e1),sub(e2))))
                    filter!(e -> prime(e) !== prime(e1), elems1)
                    filter!(e -> prime(e) !== prime(e2), elems2)
                    break # go to next e1
                elseif prime(e1) === !prime(e2)
                    for e3 in elems2
                        if e3 !== e2
                            push!(elems_prod, Element(prime(e3), disjoin(sub(e3),sub(e1))))
                        end
                    end
                    for e4 in elems1
                        if e4 !== e1
                            push!(elems_prod, Element(prime(e4), disjoin(sub(e4),sub(e2))))
                        end
                    end
                    filter!(e -> prime(e) !== prime(e1), elems1)
                    filter!(e -> prime(e) !== prime(e2), elems2)
                    break # go to next e1
                end
            end
        end
        product = vec([(e1,e2) for e1 in elems1, e2 in elems2])
        while !isempty(product)
            (e1, e2) = pop!(product)
            newprime = conjoin(prime(e1),prime(e2))
            if newprime != trimfalse
                push!(elems_prod, Element(newprime, disjoin(sub(e1),sub(e2))))
            end
            if newprime === prime(e1)
                # p1 |= p2 and therefore p1 will be mutex with all other p2-primes
                filter!(p -> prime(p[1]) !== prime(e1), product)
            elseif newprime === prime(e2)
                # p2 |= p1 and therefore p2 will be mutex with all other p1-primes
                filter!(p -> prime(p[2]) !== prime(e2), product)
            end
        end
        canonicalize(XYPartition(elems_prod))
    end
end

"""
Disjoin two SDDs when one descends from the other
"""
function disjoin_descendent(d::TrimNode, n::TrimNode)::TrimNode
    get!(vtree(n).disjoin_cache, (d,n)) do
        if descends_left_from(d, n)
            not_d = !d
            elements = Element[Element(conjoin(prime(e),not_d), sub(e)) for e in children(n)]
            elements = remove_false_primes(elements)
            push!(elements,Element(d, trimtrue))
        else
            # @assert descends_right_from(d, n)
            elements = Element[Element(prime(e),disjoin(sub(e),d)) for e in children(n)]
        end
        #TODO are there cases where we don't need all of compress-trim-unique?
        canonicalize(XYPartition(elements))
    end
end

"""
Disjoin two SDDs in separate parts of the vtree
"""
function disjoin_indep(s::TrimNode, t::TrimNode)::Trim⋁
    # @assert GateType(s)!=ConstantGate() && GateType(t)!=ConstantGate()
    mgr = parentlca(s,t)
    # @assert vtree(s) != mgr && vtree(t) != mgr
    (s,t) = pointer_sort(s,t)
    get!(mgr.disjoin_cache, (s,t)) do
        if descends_left_from(vtree(s), mgr)
            # @assert descends_right_from(vtree(t), mgr)
            elements = Element[Element(s,trimtrue),Element(!s,t)]
        else
            # @assert descends_left_from(vtree(t), mgr)
            # @assert descends_right_from(vtree(s), mgr)
            elements = Element[Element(t,trimtrue),Element(!t,s)]
        end
        # TODO: the XY partition must already be compressed and trimmed
        unique⋁(XYPartition(elements), mgr)
    end
end

@inline Base.:|(s,t) = disjoin(s,t)

"""
Negate an SDD
"""
@inline negate(::TrimFalse)::TrimTrue = trimtrue
@inline negate(::TrimTrue)::TrimFalse = trimfalse

function negate(s::TrimLiteral)::TrimLiteral
    if positive(s)
        vtree(s).negative_literal
    else
        vtree(s).positive_literal
    end
end

negate(node::Trim⋁)::Trim⋁ = node.negation

@inline Base.:!(s) = negate(s)


text:
type PSMMALAMCTune <: MCTunerState
  smmalatune::BasicMCTune
  malatune::BasicMCTune
  totaltune::BasicMCTune
  smmalafrequency::Real
  malafrequency::Real
end

PSMMALAMCTune(smmalatune::BasicMCTune, malatune::BasicMCTune, totaltune::BasicMCTune) =
  PSMMALAMCTune(smmalatune, malatune, totaltune, NaN, NaN)

PSMMALAMCTune(smmalastep::Real=1., malastep::Real=1.) =
  PSMMALAMCTune(BasicMCTune(smmalastep), BasicMCTune(malastep), BasicMCTune(1.), NaN, NaN)

immutable PSMMALAMCTuner <: MCTuner
  smmalatuner::VanillaMCTuner
  malatuner::VanillaMCTuner
  totaltuner::Union{VanillaMCTuner, AcceptanceRateMCTuner}
end

Base.show(io::IO, tuner::PSMMALAMCTuner) = print(io, "PSMMALAMCTuner")

Base.show(io::IO, ::MIME"text/plain", tuner::PSMMALAMCTuner) = show(io, tuner)

ext: jl


text:
using DataFrames, CSV, DataFramesMeta, Dates, BenchmarkTools, Cthulhu
using Revise
using AbnormalReturns

##

df_firm = CSV.File(joinpath("data", "firm_ret.csv")) |> DataFrame
df_mkt = CSV.File(joinpath("data", "mkt_ret.csv")) |> DataFrame
df_events = CSV.File(joinpath("data", "event_dates.csv")) |> DataFrame

##
@time data = MarketData(df_mkt, df_firm; id_col=:firm_id, valuecols_firms=[:ret])
# First run R5 3600: 29.455642 seconds (32.05 M allocations: 13.431 GiB, 7.98% gc time, 38.13% compilation time)
# Second run R5 3600: 17.942238 seconds (669.24 k allocations: 11.774 GiB, 11.31% gc time)
# First run i7 6700: 37.776454 seconds (32.06 M allocations: 13.430 GiB, 30.23% gc time, 32.16% compilation time)
# Second run i7 6700: 18.069534 seconds (669.95 k allocations: 11.774 GiB, 15.21% gc time)

##
df_firm = nothing
df_mkt = nothing
GC.gc()

##
@time df_temp = @chain df_events begin
    @transform(:reg_mkt = quick_reg(data[:firm_id, :est_window_start .. :est_window_end], @formula(ret ~ mkt)),)
    @transform(:reg_ffm = quick_reg(data[:firm_id, :est_window_start .. :est_window_end], @formula(ret ~ mkt + smb + hml + umd)),)
    @transform(
        :bhar_mkt = AbnormalReturns.bhar(data[:firm_id, :est_window_start .. :est_window_end], :reg_mkt),
        :bhar_ffm = AbnormalReturns.bhar(data[:firm_id, :est_window_start .. :est_window_end], :reg_ffm),
        # :bhar_simple = bhar(data[:firm_id, :est_window_start .. :est_window_end], ^(:ret), ^(:mkt)),
        # :car_simple = car(data[:firm_id, :est_window_start .. :est_window_end], ^(:ret), ^(:mkt))
    )
    @rtransform(
        :var_mkt = var(:reg_mkt),
        :var_ffm = var(:reg_ffm),
        :alpha = alpha(:reg_mkt),
        :beta = beta(:reg_mkt),
    )
end
# First run R5 3600: 19.867720 seconds (63.84 M allocations: 5.861 GiB, 13.96% gc time, 52.57% compilation time)
# Second run R5 3600: 8.928954 seconds (37.73 M allocations: 4.539 GiB, 17.06% gc time, 0.94% compilation time)
# First run i7 6700: 33.919134 seconds (64.13 M allocations: 5.936 GiB, 45.39% gc time, 65.25% compilation time)
# Second run i7 6700: 19.343996 seconds (37.73 M allocations: 4.602 GiB, 62.63% gc time, 0.44% compilation time)

##

cols = TimelineColumn.([:ret, :mkt, :smb, :hml, :umd])
@time @chain df_events[1:1000000, :] begin
    @rtransform(:reg = quick_reg(data[:firm_id, :est_window_start .. :est_window_end, cols], @formula(ret ~ mkt + smb + hml + umd)),)
end
# Run R5 3600: 64.389570 seconds (370.13 M allocations: 88.340 GiB, 15.36% gc time, 0.14% compilation time)

##

@time @chain df_events[1:1000000, :] begin
    @transform(:reg = quick_reg(data[:firm_id, :est_window_start .. :est_window_end], @formula(ret ~ mkt + smb + hml + umd)),)
end
# Run R5 3600: 2.824239 seconds (8.57 M allocations: 1.336 GiB, 1.89% compilation time)
##
@time @chain df_events[1:1000000, :] begin
    @transform(:bhar = bhar(data[:firm_id, :est_window_start .. :est_window_end]),)
end
##
quick_reg(data[1, Date(2017) .. Date(2018), cols], @formula(ret ~ mkt + smb + hml + umd))
##

@benchmark $data[i, TimelineColumn.([:ret, :mkt, :smb, :hml, :umd])] setup=(i=rand(1:10000))
##

# i = 151465
# ids = df_events.firm_id[1:100]
# date_mins = df_events.est_window_start[1:100]
# date_maxs = df_events.est_window_end[1:100]
# f = @formula(ret ~ mkt + smb + hml + umd)
# cols = AbnormalReturns.internal_termvars(f)
# data_dict = AbnormalReturns.construct_id_dict(ids)
# sch = apply_schema(f, schema(f, data))
# cache = AbnormalReturns.create_pred_matrix(data[1], sch)
# out = fill(BasicReg(0, f), length(ids))

# int_data = data[1]

b = @benchmarkable AbnormalReturns.vector_reg(
    $data,
    $ids,
    $date_mins,
    $date_maxs,
    $f,
)
b.params.evals=100
b.params.seconds=100
run(b)
##


text:
include("axelrodnet.jl")

agentcount::Int64=100,
n_iter::Int64=1000,
socialbotfrac::Float64=0.00,
m0::Int64=5,
rndseed::Int64=1,
repcount::Int64=1,
completemixing::Bool=false

test = Axelrodnet.run!()
test_2 = Axelrodnet.run!(sqrt_agentcount=1000, n_iter=1000, repcount=3)
using CSV
fullmix = Axelrodnet.run!(completemixing = true)

print(test[1])
print(fullmix[1])

grid([3,3])

trunc(Int64, 5.5)

test2 = convert2weightedGraph(test[1], test[2])

using DataFrames

d = DataFrame(A = Int64[], B = Any[])

append!(
    d,
    DataFrame(
        A = [1, 2, 3],
        B = [3, "b", 3]
    )
)

max_stars_repo_name: JohannesNakayama/axelrodnet

max_issues_repo_name: JohannesNakayama/axelrodnet

max_forks_repo_name: JohannesNakayama/axelrodnet


text:
using Revise, YaoGraphRegister, LightGraphs, GraphPlot, Compose

using Colors

g = SimpleGraph(10, 10)

nodecolor = [colorant"lightseagreen", colorant"orange"]
membership = [i == 3 ? 2 : 1 for i in 1:10]
nodefillc = nodecolor[membership]

gplot(g; nodelabel=1:10, nodefillc=nodefillc, layout=circular_layout)

max_stars_repo_name: QuantumBFS/YaoGraphs

max_issues_repo_name: QuantumBFS/YaoGraphs

max_forks_repo_name: QuantumBFS/YaoGraphs


text:
# Scoped Options

ContextVariablesX.@contextvar options_context::NamedTuple = NamedTuple()

function with_options(f, options::NamedTuple)
    ContextVariablesX.with_context(options_context => options) do
        f()
    end
end
with_options(f; options...) = with_options(f, NamedTuple(options))

get_options() = options_context[]
get_options(key::Symbol) = getproperty(get_options(), key)
function get_options(key::Symbol, default)
    opts = get_options()
    return haskey(opts, key) ? getproperty(opts, key) : default
end


text:
include("linearcode.jl")

function repetitioncode(n::Integer)
    F, _ = FiniteField(2, 1, "α")
    G = matrix(F, 1, n, [1 for i in 1:n])
    M2 = MatrixSpace(F, n - 1, 1)
    M3 = MatrixSpace(F, n - 1, n - 1)
    H = hcat(M2([1 for i in 1:(n - 1)]), M3(1))
    return LinearCode(F, n, 1, n, G, G, H, H, G, H, missing)
end

# function Hammingcode(p::Integer, r::Integer)
#
# end

# function tetracode()
#     return Hammingcode(3, 2)
# end

# both of these are actually the extended Golay codes
function Golay(p::Integer)
    if p == 2
        F, _ = FiniteField(2, 1, "α")
        M = MatrixSpace(F, 12 , 12)
        A = M([0 1 1 1 1 1 1 1 1 1 1 1;
             1 1 1 0 1 1 1 0 0 0 1 0;
             1 1 0 1 1 1 0 0 0 1 0 1;
             1 0 1 1 1 0 0 0 1 0 1 1;
             1 1 1 1 0 0 0 1 0 1 1 0;
             1 1 1 0 0 0 1 0 1 1 0 1;
             1 1 0 0 0 1 0 1 1 0 1 1;
             1 0 0 0 1 0 1 1 0 1 1 1;
             1 0 0 1 0 1 1 0 1 1 1 0;
             1 0 1 0 1 1 0 1 1 1 0 0;
             1 1 0 1 1 0 1 1 1 0 0 0;
             1 0 1 1 0 1 1 1 0 0 0 1])
        G = hcat(M(1), A)
        H = hcat(-A', M(1))
        return LinearCode(F, 24, 12, 8, G, G, H, H, G, H, missing)
    elseif p == 3
        F, _ = FiniteField(3, 1, "α")
        M = MatrixSpace(F, 6 , 6)
        A = M([0 1 1 1 1 1;
               1 0 1 -1 -1 1;
               1 1 0 1 -1 -1;
               1 -1 1 0 1 -1;
               1 -1 -1 1 0 1;
               1 1 -1 -1 1 0])
        G = hcat(M(1), A)
        H = hcat(-A', M(1))
        return LinearCode(F, 12, 6, 6, G, G, H, H, G, H, missing)
    else
        error("Golay code not implemented for q = $q.")
    end
end

max_stars_repo_name: esabo/CodingTheory

max_issues_repo_name: esabo/CodingTheory

max_forks_repo_name: esabo/CodingTheory


text:
using TrafficAssignment
using Test, LinearAlgebra
using DelimitedFiles

include("test_functions.jl")

test_tntp()

summarize_ta_data()

data_time = time()
ta_data = load_ta_network("SiouxFalls")
println("Data Loading Completed, time:", time() - data_time, " seconds")

@testset "Various FW Methods" begin
    ta_data = load_ta_network("SiouxFalls")

    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:bfw, step=:exact, log=:off, tol=1e-3, max_iter_no=5)
    @test abs( objective - 4.963799502172654e6 ) < 1e6

    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:cfw, step=:newton, log=:off, tol=1e-3, max_iter_no=5)
    @test abs( objective - 4.963799502172654e6 ) < 1e6

    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:fw, step=:exact, log=:off, tol=1e-3, max_iter_no=5)
    @test abs( objective - 4.963799502172654e6 ) < 1e6
end

@testset "Testing Sioux Falls" begin
    ta_data = load_ta_network("SiouxFalls")
    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:cfw, step=:newton, log=:off, tol=1e-5, max_iter_no=50000)
    solution, header = readdlm("SiouxFalls_flow.csv", ',', header=true)
    solution_flow = Float64.(solution[:, 3])
    @test norm(link_volume - solution_flow) / norm(solution_flow) < 0.01
end

@testset "Testing Anaheim" begin
    ta_data = load_ta_network("Anaheim")
    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:cfw, step=:newton, log=:off, tol=1e-5, max_iter_no=50000)
    solution, header = readdlm("Anaheim_flow.csv", ',', header=true)
    solution_flow = Float64.(solution[:, 3])
    @test norm(link_volume - solution_flow) / norm(solution_flow) < 0.01
end

summarize_ta_data(markdown=true)


text:
using IterTools
import IteratorInterfaceExtensions
import TableTraits
using Random
using ProgressMeter
using Serialization
using CSVFiles
using FileIO

function print_nonempty(name, vector)
    if length(vector) > 0
        println("  $name:")
        for obj in vector
            println("    ", obj)
        end
    end
end

function Base.show(io::IO, sim_def::SimulationDef{T}) where T <: AbstractSimulationData
    println("SimulationDef{$T}")

    println("  rvdict:")
    for (key, value) in sim_def.rvdict
        println("    $key: $(typeof(value))")
    end

    print_nonempty("translist", sim_def.translist)
    print_nonempty("savelist",  sim_def.savelist)
    println("  nt_type: $(sim_def.nt_type)")

    Base.show(io, sim_def.data)  # note: data::T
end

function Base.show(io::IO, sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData
    println("SimulationInstance{$T}")
    print_nonempty("translist for model params", sim_inst.translist_modelparams)

    Base.show(io, sim_inst.sim_def)

    println("  trials: $(sim_inst.trials)")
    println("  current_trial: $(sim_inst.current_trial)")
    sim_inst.current_trial > 0 && println("  current_data: $(sim_inst.current_data)")

    println("  $(length(sim_inst.models)) models")
    println("  $(length(sim_inst.results)) results dicts")
end

function Base.show(obj::T) where T <: AbstractSimulationData
    nothing
end

"""
    _store_param_results(m::AbstractModel, datum_key::Tuple{Symbol, Symbol},
                        trialnum::Int, scen_name::Union{Nothing, String},
                        results::Dict{Tuple, DataFrame})

Store `results` for a single parameter `datum_key` in model `m` and return the
dataframe for this particular `trial_num`/`scen_name` combination.
"""
function _store_param_results(m::AbstractModel, datum_key::Tuple{Symbol, Symbol},
                            trialnum::Int, scen_name::Union{Nothing, String},
                            results::Dict{Tuple, DataFrame})
    @debug "\nStoring trial results for $datum_key"

    (comp_name, datum_name) = datum_key
    dims = dim_names(m, comp_name, datum_name)
    has_scen = ! (scen_name === nothing)

    if length(dims) == 0        # scalar value
        value = m[comp_name, datum_name]
        # println("Scalar: $value")

        if haskey(results, datum_key)
            results_df = results[datum_key]
        else
            cols = [[], []]
            names = [datum_name, :trialnum]
            if has_scen
                push!(cols, [])
                push!(names, :scen)
            end
            results_df = DataFrame(cols, names)
            results[datum_key] = results_df
        end

        trial_df = DataFrame(datum_name => value, :trialnum => trialnum)
        has_scen ? trial_df[!, :scen] .= scen_name : nothing
        append!(results_df, trial_df)
        # println("results_df: $results_df")

    else
        trial_df = getdataframe(m, comp_name, datum_name)
        trial_df[!, :trialnum] .= trialnum
        has_scen ? trial_df[!, :scen] .= scen_name : nothing
        # println("size of trial_df: $(size(trial_df))")

        if haskey(results, datum_key)
            results_df = results[datum_key]
            # println("Appending trial_df $(size(trial_df)) to results_df $(size(results_df))")
            append!(results_df, trial_df)
        else
            # println("Setting results[$datum_key] = trial_df $(size(trial_df))")
            results[datum_key] = trial_df
        end
    end

    return trial_df
end

"""
    _store_trial_results(sim_inst::SimulationInstance{T}, trialnum::Int,
                        scen_name::Union{Nothing, String}, output_dir::Union{Nothing, String},
                        streams::Dict{String, CSVFiles.CSVFileSaveStream{IOStream}}) where T <: AbstractSimulationData

Save the stored simulation results ` from trial `trialnum` and scenario `scen_name`
to files in the directory `output_dir`
"""
function _store_trial_results(sim_inst::SimulationInstance{T}, trialnum::Int,
                                scen_name::Union{Nothing, String}, output_dir::Union{Nothing, String},
                                streams::Dict{String, CSVFiles.CSVFileSaveStream{IOStream}}) where T <: AbstractSimulationData
    savelist = sim_inst.sim_def.savelist

    model_index = 1
    for (m, results) in zip(sim_inst.models, sim_inst.results)
        for datum_key in savelist
            trial_df = _store_param_results(m, datum_key, trialnum, scen_name, results)
            if output_dir !== nothing

                # get sub_dir, which is different from output_dir if there are multiple models
                if (length(sim_inst.results) > 1)
                    sub_dir = joinpath(output_dir, "model_$(model_index)")
                else
                    sub_dir = output_dir
                end
                mkpath(sub_dir, mode=0o750)

                # get filtered trial_df, which is different from trial_df if there are multiple scenarios
                if scen_name !== nothing
                    trial_df_filtered = filter(row -> row[:scen] .== scen_name, trial_df)[:, 1:end-1] # remove scen field
                else
                    trial_df_filtered = trial_df
                end

                datum_name = join(map(string, datum_key), "_")
                _save_trial_results(trial_df_filtered, datum_name, sub_dir, streams)
            end
        end
        model_index += 1
    end
end

"""
    _save_trial_results(trial_df::DataFrame, datum_name::String, output_dir::String,
                        streams::Dict{String, CSVFiles.CSVFileSaveStream{IOStream}})

Save the stored simulation results in `trial_df` from trial `trialnum` to files
in the directory `output_dir`
"""
function _save_trial_results(trial_df::DataFrame, datum_name::String, output_dir::AbstractString, streams::Dict{String, CSVFiles.CSVFileSaveStream{IOStream}}) where T <: AbstractSimulationData
    filename = joinpath(output_dir, "$datum_name.csv")
    if haskey(streams, filename)
        write(streams[filename], trial_df)
    else
        streams[filename] = savestreaming(filename, trial_df)
    end
end

"""
    save_trial_inputs(sim_inst::SimulationInstance, filename::String)

Save the trial inputs for `sim_inst` to `filename`.
"""
function save_trial_inputs(sim_inst::SimulationInstance, filename::String)
    mkpath(dirname(filename), mode=0o750)   # ensure that the specified path exists
    save(filename, sim_inst)
    return nothing
end

# TBD: Modify lhs() to return an array of SampleStore{T} instances?
"""
    get_trial(sim_inst::SimulationInstance, trialnum::Int)

Return a NamedTuple with the data for next trial. Note that the `trialnum`
parameter is used only to support a 1-deep data cache that allows this
function to be called successively with the same `trialnum` to retrieve
the same NamedTuple. If `trialnum` does not match the current trial number,
the argument is ignored.
"""
function get_trial(sim_inst::SimulationInstance, trialnum::Int)

    if sim_inst.current_trial == trialnum
        return sim_inst.current_data
    end

    sim_def = sim_inst.sim_def

    vals = [rand(rv.dist) for rv in values(sim_def.rvdict)]
    sim_inst.current_data = sim_def.nt_type((vals...,))
    sim_inst.current_trial = trialnum

    return sim_inst.current_data
end

"""
    generate_trials!(sim_inst::SimulationInstance{T}, samples::Int; filename::Union{String, Nothing}=nothing)

Generate trials for the given `SimulationInstance` using the defined `samplesize.
Call this before running the sim to pre-generate data to be used by all scenarios.
Also saves inputs if a filename is given.
"""
function generate_trials!(sim_inst::SimulationInstance{T}, samplesize::Int;
                        filename::Union{String, Nothing}=nothing) where T <: AbstractSimulationData
    sample!(sim_inst, samplesize)

    # TBD: If user asks for trial data to be saved, generate it up-front, or
    # open a file that can be written to for each trialnum/scenario set?
    if filename != nothing
        save_trial_inputs(sim_inst, filename)
    end
end

function sample!(sim_inst::MonteCarloSimulationInstance, samplesize::Int)
    sim_inst.trials = samplesize
    rand!(sim_inst)
end

"""
    Random.rand!(sim_inst::SimulationInstance{T})

Replace all RVs originally of type Distribution with SampleStores with
values drawn from that original distribution.
"""
function Random.rand!(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData
    sim_def = sim_inst.sim_def
    rvdict = sim_def.rvdict
    trials = sim_inst.trials

    for rv in values(sim_def.rvdict)
        # use underlying distribution, if known
        orig_dist = (rv.dist isa SampleStore ? rv.dist.dist : rv.dist)
        dist = (orig_dist === nothing ? rv.dist : orig_dist)
        values = rand(dist, trials)
        rvdict[rv.name] = RandomVariable(rv.name, SampleStore(values, orig_dist))
    end
end

"""
    _copy_sim_params(sim_inst::SimulationInstance{T})

Copy the parameters that are perturbed so we can restore them after each trial. This
is necessary when we are applying distributions by adding or multiplying original values.
"""
function _copy_sim_params(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData

    # If there is a MarginalModel, need to copy the params for both the base and marginal modeldefs separately
    flat_model_list = _get_flat_model_list(sim_inst)
    param_vec = Vector{Dict{Symbol, ModelParameter}}(undef, length(flat_model_list))

    for (i, m) in enumerate(flat_model_list)
        md = modelinstance_def(m)
        param_vec[i] = Dict{Symbol, ModelParameter}(trans.paramnames[i] => copy(model_param(md, trans.paramnames[i])) for trans in sim_inst.translist_modelparams)
    end

    return param_vec
end

function _restore_sim_params!(sim_inst::SimulationInstance{T},
                              param_vec::Vector{Dict{Symbol, ModelParameter}}) where T <: AbstractSimulationData

    # Need to flatten the list of models so that if there is a MarginalModel,
    # both its base and marginal models will have their separate params restored
    flat_model_list = _get_flat_model_list(sim_inst)

    for (i, m) in enumerate(flat_model_list)
        params = param_vec[i]
        md = m.mi.md
        for trans in sim_inst.translist_modelparams
            name = trans.paramnames[i]
            param = params[name]
            _restore_param!(param, name, md, i, trans)
        end
    end

    return nothing
end

function _restore_param!(param::ScalarModelParameter{T}, name::Symbol, md::ModelDef, i::Int, trans::TransformSpec_ModelParams) where T
    md_param = model_param(md, name)
    md_param.value = param.value
end

function _restore_param!(param::ArrayModelParameter{T}, name::Symbol, md::ModelDef, i::Int, trans::TransformSpec_ModelParams) where T
    md_param = model_param(md, name)
    indices = _param_indices(param, md, i, trans)
    md_param.values[indices...] = param.values[indices...]
end

function _param_indices(param::ArrayModelParameter{T}, md::ModelDef, i::Int, trans::TransformSpec_ModelParams) where T
    pdims = dim_names(param)   # returns [] for scalar parameters
    num_pdims = length(pdims)

    tdims  = trans.dims
    num_dims = length(tdims)

    # special case for handling reshaped data where a single draw returns a matrix of values
    if num_dims == 0
        indices = repeat([Colon()], num_pdims)
        return indices
    end

    if num_pdims != num_dims
        pname = trans.paramnames[i]
        error("Dimension mismatch: model parameter :$pname has $num_pdims dimensions ($pdims); Sim has $num_dims")
    end

    indices = Vector()
    for (dim_name, dim_values) in zip(pdims, tdims)
        dim = dimension(md, dim_name)
        dim_indices = dim[dim_values]
        dim_name == :time ? dim_indices = TimestepIndex.(dim_indices) : nothing
        push!(indices, dim_indices)
    end

    return indices
end

function _perturb_param!(param::ScalarModelParameter{T}, md::ModelDef, i::Int, trans::TransformSpec_ModelParams, rvalue::Number) where T
    op = trans.op

    if op == :(=)
        param.value = T(rvalue)

    elseif op == :(*=)
        param.value *= rvalue

    else
        param.value += rvalue
    end
end

# rvalue is an Array so we expect the dims to match and don't need to worry about
# broadcasting
function _perturb_param!(param::ArrayModelParameter{T}, md::ModelDef, i::Int,
    trans::TransformSpec_ModelParams, rvalue::Array{<: Number, N}) where {T, N}

    op = trans.op
    pvalue = value(param)
    indices = _param_indices(param, md, i, trans)

    if op == :(=)
        pvalue[indices...] = rvalue

    elseif op == :(*=)
        pvalue[indices...] *= rvalue

    else
        pvalue[indices...] += rvalue

    end
end

# rvalue is a Number so we might need to deal with broadcasting
function _perturb_param!(param::ArrayModelParameter{T}, md::ModelDef, i::Int,
                         trans::TransformSpec_ModelParams, rvalue::Number) where {T, N}
    op = trans.op
    pvalue = value(param)
    indices = _param_indices(param, md, i, trans)

    if op == :(=)

        # first we check for a time index
        ti = get_time_index_position(param)

        # If there is no time index we have all methods needed to broadcast normally
        if isnothing(ti)
            broadcast_flag = sum(map(x -> length(x) > 1, indices)) > 0
            broadcast_flag ? pvalue[indices...] .= rvalue : pvalue[indices...] = rvalue

        else
            indices1, ts, indices2 = split_indices(indices, ti)
            non_ts_indices = [indices1..., indices2...]
            broadcast_flag = isempty(non_ts_indices) ? false : sum(map(x -> length(x) > 1, non_ts_indices)) > 0

            # Loop over the Array of TimestepIndex
            if isa(ts, Array)
                for el in ts
                    broadcast_flag ? pvalue[indices1..., el, indices2...] .= rvalue : pvalue[indices1..., el, indices2...] = rvalue
                end

            # The time is just a single TimestepIndex and we can proceed with broadcast
            else
                broadcast_flag ? pvalue[indices...] .= rvalue : pvalue[indices...] = rvalue
            end
        end

    elseif op == :(*=)
        pvalue[indices...] *= rvalue

    else
        pvalue[indices...] += rvalue
    end
end

"""
    _perturb_params!(sim_inst::SimulationInstance{T}, trialnum::Int)

Modify the stochastic parameters for all models in `sim_inst`, using the
values drawn for trial `trialnum`.
"""
function _perturb_params!(sim_inst::SimulationInstance{T}, trialnum::Int) where T <: AbstractSimulationData
    if trialnum > sim_inst.trials
        error("Attempted to run trial $trialnum, but only $(sim_inst.trials) trials are defined")
    end

    trialdata = get_trial(sim_inst, trialnum)

    # If it's a MarginalModel, need to perturb the params in both the base and marginal modeldefs
    flat_model_list = _get_flat_model_list(sim_inst)
    for (i, m) in enumerate(flat_model_list)
        for trans in sim_inst.translist_modelparams
            param = model_param(m.mi.md, trans.paramnames[i])
            rvalue = getfield(trialdata, trans.rvname)
            _perturb_param!(param, m.mi.md, i, trans, rvalue)
        end
    end
    return nothing
end

function _reset_rvs!(sim_def::SimulationDef{T}) where T <: AbstractSimulationData
    for rv in values(sim_def.rvdict)
        if rv.dist isa SampleStore
            reset(rv.dist)
        end
    end
end

"""
    _reset_results!(sim_inst::SimulationInstance{T})

Reset all simulation results storage to a vector of empty dicts
"""
function _reset_results!(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData
    sim_inst.results = [Dict{Tuple, DataFrame}() for m in sim_inst.models]
end

# Append a string representation of the tuple args to the given directory name
function _compute_output_dir(orig_output_dir, tup)
    if orig_output_dir === nothing
        output_dir = nothing
    else
        output_dir = joinpath(orig_output_dir, join(map(string, tup), "_"))
        mkpath(output_dir, mode=0o750)
    end
    return output_dir
end

"""
    Base.run(sim_def::SimulationDef{T},
            models::Union{Vector{M}, AbstractModel},
            samplesize::Int;
            ntimesteps::Int=typemax(Int),
            trials_output_filename::Union{Nothing, AbstractString}=nothing,
            results_output_dir::Union{Nothing, AbstractString}=nothing,
            pre_trial_func::Union{Nothing, Function}=nothing,
            post_trial_func::Union{Nothing, Function}=nothing,
            scenario_func::Union{Nothing, Function}=nothing,
            scenario_placement::ScenarioLoopPlacement=OUTER,
            scenario_args=nothing,
            results_in_memory::Bool=true) where {T <: AbstractSimulationData, M <: AbstractModel}

Run the simulation definition `sim_def` for the `models` using `samplesize` samples.

Optionally run the `models` for `ntimesteps`, if specified,
else to the maximum defined time period. Note that trial data are applied to all the
associated models even when running only a portion of them.

If provided, the generated trials and results will be saved in the indicated
`trials_output_filename` and `results_output_dir` respectively. If `results_in_memory` is set
to false, then results will be cleared from memory and only stored in the
`results_output_dir`.

If `pre_trial_func` or `post_trial_func` are defined, the designated functions are called
just before or after (respectively) running a trial. The functions must have the signature:

    fn(sim_inst::SimulationInstance, trialnum::Int, ntimesteps::Int, tup::Tuple)

where `tup` is a tuple of scenario arguments representing one element in the cross-product
of all scenario value vectors. In situations in which you want the simulation loop to run only
some of the models, the remainder of the runs can be handled using a `pre_trial_func` or
`post_trial_func`.

If provided, `scenario_args` must be a `Vector{Pair}`, where each `Pair` is a symbol and a
`Vector` of arbitrary values that will be meaningful to `scenario_func`, which must have
the signature:

    scenario_func(sim_inst::SimulationInstance, tup::Tuple)

By default, the scenario loop encloses the simulation loop, but the scenario loop can be
placed inside the simulation loop by specifying `scenario_placement=INNER`. When `INNER`
is specified, the `scenario_func` is called after any `pre_trial_func` but before the model
is run.

Returns the type `SimulationInstance` that contains a copy of the original `SimulationDef`,
along with mutated information about trials, in addition to the model list and
results information.
"""
function Base.run(sim_def::SimulationDef{T},
                models::Union{Vector{M}, AbstractModel},
                samplesize::Int;
                ntimesteps::Int=typemax(Int),
                trials_output_filename::Union{Nothing, AbstractString}=nothing,
                results_output_dir::Union{Nothing, AbstractString}=nothing,
                pre_trial_func::Union{Nothing, Function}=nothing,
                post_trial_func::Union{Nothing, Function}=nothing,
                scenario_func::Union{Nothing, Function}=nothing,
                scenario_placement::ScenarioLoopPlacement=OUTER,
                scenario_args=nothing,
                results_in_memory::Bool=true) where {T <: AbstractSimulationData, M <: AbstractModel}

    # If the provided models list has both a Model and a MarginalModel, it will be a Vector{Any}, and needs to be converted
    if models isa Vector{Any}
        models = convert(Vector{AbstractModel}, models)
    end

    # Quick check for results saving
    if (!results_in_memory) && (results_output_dir === nothing)
        error("The results_in_memory keyword arg is set to ($results_in_memory) and
        results_output_dir keyword arg is set to ($results_output_dir), thus
        results will not be saved either in memory or in a file.")
    end

    # Initiate the SimulationInstance and set the models and trials for the copied
    # sim held within sim_inst
    sim_inst = SimulationInstance{typeof(sim_def.data)}(sim_def)
    set_models!(sim_inst, models)
    generate_trials!(sim_inst, samplesize; filename=trials_output_filename)
    set_translist_modelparams!(sim_inst) # should this use m.md or m.mi.md (after building below)?

    if (scenario_func === nothing) != (scenario_args === nothing)
        error("run: scenario_func and scenario_arg must both be nothing or both set to non-nothing values")
    end

    for m in sim_inst.models
        is_built(m) || build!(m)
    end

    trials = 1:sim_inst.trials

    # Save the original dir since we modify the output_dir to store scenario results
    orig_results_output_dir = results_output_dir

    # booleans vars to simplify the repeated tests in the loop below
    has_results_output_dir  = (orig_results_output_dir !== nothing)
    has_scenario_func  = (scenario_func !== nothing)
    has_outer_scenario = (has_scenario_func && scenario_placement == OUTER)
    has_inner_scenario = (has_scenario_func && scenario_placement == INNER)

    if has_scenario_func
        scen_names  = [arg.first  for arg in scenario_args]
        scen_values = [arg.second for arg in scenario_args]

        # precompute all combinations of scenario arguments so we can run
        # a single loop regardless of the number of scenario arguments.
        arg_tuples = Iterators.product(scen_values...)

        if has_outer_scenario
            arg_tuples_outer = arg_tuples
            arg_tuples_inner = (nothing,)   # allows one iteration when no scenario loop specified
        else
            arg_tuples_outer = (nothing,)   # as above
            arg_tuples_inner = arg_tuples
        end
    else
        arg_tuples = arg_tuples_outer = arg_tuples_inner = (nothing,)
        scen_name = nothing
    end

    # Set up progress bar
    nscenarios = length(arg_tuples)
    ntrials = length(trials)
    total_runs = nscenarios * ntrials
    counter = 1
    p = Progress(total_runs, counter, "Running $ntrials trials for $nscenarios scenarios...")

    for outer_tup in arg_tuples_outer
        if has_outer_scenario
            @debug "Calling outer scenario_func with $outer_tup"
            scenario_func(sim_inst, outer_tup)

            # we'll store the results of each in a subdir composed of tuple values
            results_output_dir = _compute_output_dir(orig_results_output_dir, outer_tup)

            # we'll need a scenario name for the DataFrame
            scen_name = join(map(string, outer_tup), "_")
        end

        # Save the params to be perturbed so we can reset them after each trial
        original_values = _copy_sim_params(sim_inst)

        # Reset internal index to 1 for all stored parameters to reuse the data
        _reset_rvs!(sim_inst.sim_def)

        # Create a Dictionary of streams
        streams = Dict{String, CSVFiles.CSVFileSaveStream{IOStream}}()

        try
            for (i, trialnum) in enumerate(trials)
                @debug "Running trial $trialnum"

                for inner_tup in arg_tuples_inner
                    tup = has_inner_scenario ? inner_tup : outer_tup

                    _perturb_params!(sim_inst, trialnum)

                    if pre_trial_func !== nothing
                        @debug "Calling pre_trial_func($trialnum, $tup)"
                        pre_trial_func(sim_inst, trialnum, ntimesteps, tup)
                    end

                    if has_inner_scenario
                        @debug "Calling inner scenario_func with $inner_tup"
                        scenario_func(sim_inst, inner_tup)

                        results_output_dir = _compute_output_dir(orig_results_output_dir, inner_tup)

                        # we'll need a scenario name for the DataFrame
                        scen_name = join(map(string, inner_tup), "_")
                    end

                    for m in sim_inst.models   # note that list of models may be changed in scenario_func
                        @debug "Running model"
                        run(m, ntimesteps=ntimesteps)
                    end

                    if post_trial_func !== nothing
                        @debug "Calling post_trial_func($trialnum, $tup)"
                        post_trial_func(sim_inst, trialnum, ntimesteps, tup)
                    end

                    _store_trial_results(sim_inst, trialnum, scen_name, results_output_dir, streams)
                    _restore_sim_params!(sim_inst, original_values)

                    counter += 1
                    ProgressMeter.update!(p, counter)

                    if has_results_output_dir && ! results_in_memory
                        _reset_results!(sim_inst)
                    end
                end
            end
        finally
            close.(values(streams))   # use broadcasting to close all stream
        end
    end

    return sim_inst
end

"""
    _get_flat_model_list(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData

Return a flattened vector of models, splatting out the base and modified models of
a MarginalModel.
"""
function _get_flat_model_list(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData

    flat_model_list = []
    for m in sim_inst.models
        if m isa MarginalModel
            push!(flat_model_list, m.base)
            push!(flat_model_list, m.modified)
        else
            push!(flat_model_list, m)
        end
    end
    return flat_model_list
end

"""
    _get_flat_model_list_names(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData

Return a vector of names referring to a flattened vector of models, splatting out
the base and modified models of a MarginalModel.
"""
function _get_flat_model_list_names(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData

    flat_model_list_names = [] # use for errors
    for (i, m) in enumerate(sim_inst.models)
        if m isa MarginalModel
            push!(flat_model_list_names, Symbol("Model$(i)_Base"))
            push!(flat_model_list_names, Symbol("Model$(i)_Modified"))
        else
            push!(flat_model_list_names, Symbol("Model$(i)"))
        end
    end
    return flat_model_list_names

end

# Set models
"""
    set_models!(sim_inst::SimulationInstance{T}, models::Vector{M}) where {T <: AbstractSimulationData, M <: AbstractModel}

Set the `models` to be used by the SimulationDef held by `sim_inst`.
"""
function set_models!(sim_inst::SimulationInstance{T}, models::Vector{M}) where {T <: AbstractSimulationData, M <: AbstractModel}
    sim_inst.models = models
    _reset_results!(sim_inst)    # sets results vector to same length
end

"""
    set_models!(sim_inst::SimulationInstance{T}, m::AbstractModel)  where T <: AbstractSimulationData

Set the model `m` to be used by the Simulation held by `sim_inst`.
"""
set_models!(sim_inst::SimulationInstance{T}, m::AbstractModel)  where T <: AbstractSimulationData = set_models!(sim_inst, [m])

"""
    set_translist_modelparams!(sim_inst::SimulationInstance{T})

Create the transform spec list for the simulation instance, finding the matching
model parameter names for each transform spec parameter for each model.
"""
function set_translist_modelparams!(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData

    # build flat model list that splats out the base and modified models of MarginalModel
    flat_model_list = _get_flat_model_list(sim_inst)
    flat_model_list_names = _get_flat_model_list_names(sim_inst)

    # allocate simulation instance translist
    sim_inst.translist_modelparams = Vector{TransformSpec_ModelParams}(undef, length(sim_inst.sim_def.translist))

    for (trans_idx, trans) in enumerate(sim_inst.sim_def.translist)

        # initialize the vector of model parameters
        model_parameters_vec = Vector{Symbol}(undef, length(flat_model_list))

        # handling an unshared parameter specific to a component/parameter pair
        compname = trans.compname
        if !isnothing(compname)
            for (model_idx, m) in enumerate(flat_model_list)

                # check for component in the model
                compname in keys(components(m.md)) || error("Component $compname does not exist in $(flat_model_list_names[model_idx]).")

                model_param_name = get_model_param_name(m.md, compname, trans.paramname)

                # if this is a shared parameter the user should use syntax without
                # compname in it, although this could warn or error
                if is_shared(model_param(m.md, model_param_name))
                    @warn string("Parameter indicated in `defsim` as $compname.$(trans.paramname) ",
                            "is connected to a SHARED parameter $model_param_name. Thus the ",
                            "value will be varied in all component parameters connected to ",
                            "that shared model parameter.  We suggest using $model_param_name = distribution ",
                            "syntax to be transparent about this.")
                end
                model_parameters_vec[model_idx] = model_param_name
            end

        # no component, so this should be referring to a shared parameter ... but
        # historically might not have done so and been using one set by default etc.
        else
            paramname = trans.paramname
            suggestion_string = "use the `ComponentName.ParameterName` syntax in your SimulationDefinition to explicitly define this transform ie. `ComponentName.$paramname = RandomVariable`"

            for (model_idx, m) in enumerate(flat_model_list)
                model_name = flat_model_list_names[model_idx]

                # found the shared parameter
                if has_parameter(m.md, paramname)
                    model_parameters_vec[model_idx] = paramname

                # didn't find the shared parameter, will try to resolve
                else
                    @warn "Parameter name $paramname not found in $model_name's shared parameter list, will attempt to resolve."
                    unshared_paramname = nothing
                    unshared_compname = nothing

                    for (compname, compdef) in components(m.md)
                        if has_parameter(compdef, paramname)
                            if isnothing(unshared_paramname) # first time the parameter was found in a component
                                unshared_paramname = get_model_param_name(m.md, compname, paramname) # NB might not need to use m.mi.md here could be m.md
                                unshared_compname = compname
                            else # already found in a previous component
                                error("Cannot resolve because parameter name $paramname found in more than one component of $model_name, including $unshared_compname and $compname. Please $suggestion_string.")
                            end
                        end
                    end
                    if isnothing(unshared_paramname)
                        error("Cannot resolve because $paramname not found in any of the components of $model_name.  Please $suggestion_string.")
                    else
                        @warn("Found $paramname in $unshared_compname with model parameter name $unshared_paramname. Will use this model parameter, but in the future we suggest you $suggestion_string")
                        model_parameters_vec[model_idx] = unshared_paramname
                    end
                end
            end
        end
        new_trans = TransformSpec_ModelParams(model_parameters_vec, trans.op, trans.rvname, trans.dims)
        sim_inst.translist_modelparams[trans_idx] = new_trans
    end
end

#
# Iterator functions for Simulation instance directly, and for use as an IterableTable.
#
function Base.iterate(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData
    _reset_rvs!(sim_inst.sim_def)
    trialnum = 1
    return get_trial(sim_inst, trialnum), trialnum + 1
end

function Base.iterate(sim_inst::SimulationInstance{T}, trialnum) where T <: AbstractSimulationData
    if trialnum > sim_inst.trials
        return nothing
    else
        return get_trial(sim_inst, trialnum), trialnum + 1
    end
end

IteratorInterfaceExtensions.isiterable(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData = true
TableTraits.isiterabletable(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData = true

IteratorInterfaceExtensions.getiterator(sim_inst::SimulationInstance{T}) where T = SimIterator{sim_inst.sim_def.nt_type, T}(sim_inst)

column_names(sim_def::SimulationDef{T}) where T <: AbstractSimulationData = fieldnames(sim_def.nt_type)
column_types(sim_def::SimulationDef{T}) where T <: AbstractSimulationData = [eltype(fld) for fld in values(sim_def.rvdict)]

column_names(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData = column_names(sim_inst.sim_def)
column_types(sim_inst::SimulationInstance{T}) where T <: AbstractSimulationData = column_types(sim_inst.sim_def)

#
# Iteration support (which in turn supports the "save" method)
#
column_names(iter::SimIterator) = column_names(iter.sim_inst)
column_types(iter::SimIterator) = error("Not implemented") # Used to be `IterableTables.column_types(iter.sim_def)`

function Base.iterate(iter::SimIterator)
    _reset_rvs!(iter.sim_inst.sim_def)
    idx = 1
    return get_trial(iter.sim_inst, idx), idx + 1
end

function Base.iterate(iter::SimIterator, idx)
    if idx > iter.sim_inst.trials
        return nothing
    else
        return get_trial(iter.sim_inst, idx), idx + 1
    end
end

Base.length(iter::SimIterator) = iter.sim_inst.trials

Base.eltype(::Type{SimIterator{NT, T}}) where {NT, T} = NT


text:
import .MCMCChains: Chains

# A basic chains constructor that works with the Transition struct we defined.
function AbstractMCMC.bundle_samples(
    ts::Vector{<:Transition{<:AbstractArray}},
    model::DensityModel,
    sampler::MHSampler,
    state,
    chain_type::Type{Chains};
    param_names=missing,
    kwargs...
)
    # Turn all the transitions into a vector-of-vectors.
    vals = [vcat(t.params, t.lp) for t in ts]

    # Check if we received any parameter names.
    if ismissing(param_names)
        param_names = [Symbol(:param_, i) for i in 1:length(keys(ts[1].params))]
    else
        # Generate new array to be thread safe.
        param_names = Symbol.(param_names)
    end

    # Add the log density field to the parameter names.
    push!(param_names, :lp)

    # Bundle everything up and return a Chains struct.
    return Chains(vals, param_names, (internals = [:lp],))
end

function AbstractMCMC.bundle_samples(
    ts::Vector{<:Transition{<:NamedTuple}},
    model::DensityModel,
    sampler::MHSampler,
    state,
    chain_type::Type{Chains};
    param_names=missing,
    kwargs...
)
    # Convert to a Vector{NamedTuple} first
    nts = AbstractMCMC.bundle_samples(ts, model, sampler, state, Vector{NamedTuple}; param_names=param_names, kwargs...)

    # Get all the keys
    all_keys = unique(mapreduce(collect∘keys, vcat, nts))

    # Preallocate array
    # vals = []

    # Push linearized draws onto array
    trygetproperty(thing, key) = key in keys(thing) ? getproperty(thing, key) : missing
    vals = map(nt -> [trygetproperty(nt, k) for k in all_keys], nts)

    # Check if we received any parameter names.
    if ismissing(param_names)
        param_names = all_keys
    else
        # Generate new array to be thread safe.
        param_names = Symbol.(param_names)
    end

    # Bundle everything up and return a Chains struct.
    return Chains(vals, param_names, (internals = [:lp],))
end

function AbstractMCMC.bundle_samples(
    ts::Vector{<:Vector{<:Transition}},
    model::DensityModel,
    sampler::Ensemble,
    state,
    chain_type::Type{Chains};
    param_names=missing,
    kwargs...
)
    # Preallocate return array
    # NOTE: requires constant dimensionality.
    n_params = length(ts[1][1].params)
    nsamples = length(ts)
    # add 1 parameter for lp
    vals = Array{Float64, 3}(undef, nsamples, n_params + 1, sampler.n_walkers)

    for n in 1:nsamples
        for i in 1:sampler.n_walkers
            walker = ts[n][i]
            for j in 1:n_params
                vals[n, j, i] = walker.params[j]
            end
            vals[n, n_params + 1, i] = walker.lp
        end
    end

    # Check if we received any parameter names.
    if ismissing(param_names)
        param_names = [Symbol(:param_, i) for i in 1:length(keys(ts[1][1].params))]
    else
        # Generate new array to be thread safe.
        param_names = Symbol.(param_names)
    end

    # Add the log density field to the parameter names.
    push!(param_names, :lp)

    # Bundle everything up and return a Chains struct.
    return Chains(vals, param_names, (internals=[:lp],))
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule LibSSH2_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("LibSSH2")
JLLWrappers.@generate_main_file("LibSSH2", UUID("29816b5a-b9ab-546f-933c-edad1886dfa8"))
end  # module LibSSH2_jll


text:
module TKEMassFlux

using OceanTurb
using OceanTurb: nan2inf, inf2zero

using Printf

import ..OceanTurb: oncell, onface, maxsqrt, minuszero
import .KPP: ∂B∂z, u★, w★, isunstable
import .ModularKPP: AbstractModularKPPModel

const nsol = 5
@solution U V T S e

"Returns √ϕ if ϕ is positive and not NaN. Otherwise returns 0."
@inline function zeroed_sqrt(ϕ)
    ϕ *= 1 - isnan(ϕ)
    return maxsqrt(ϕ)
end

@inline zeroed_sqrt(ϕ, i) = @inbounds zeroed_sqrt(ϕ[i])

@inline sqrt_e(m, i) = @inbounds maxsqrt(m.solution.e[i])

"Returns (∂z U)^2 + (∂z V)^2 at cell interfaces."
@inline shear_squared(m, i) = ∂z(m.solution.U, i)^2 + ∂z(m.solution.V, i)^2

"Returns ∂z B at cell interfaces."
@inline ∂B∂z(m, i) = ∂B∂z(m.solution.T, m.solution.S,
                          m.constants.g, m.constants.α, m.constants.β, i)

"Returns ∂z B at cell centers."
@inline oncell_∂B∂z(m, i) = oncell(∂B∂z, m, i) # Fallback valid for linear equations of state

@inline sqrt_∂B∂z(m, i) = maxsqrt(∂B∂z(m, i))

"Returns a velocity scale associated with convection across grid cell i."
@inline wΔ³(m, i=m.grid.N) = max(zero(eltype(m.grid)), m.state.Qb) * Δc(m.grid, i)

mutable struct Model{L, K, A, W, N, E, H, K0, C, ST, G, TS, S, BC, T} <: AbstractModel{TS, G, T}

                       clock :: Clock{T}
                        grid :: G
                 timestepper :: TS
                    solution :: S
                         bcs :: BC
               mixing_length :: L
          eddy_diffusivities :: K
       convective_adjustment :: A
              tke_wall_model :: W
                tke_equation :: E
        boundary_layer_depth :: H
               nonlocal_flux :: N
    background_diffusivities :: K0
                   constants :: C
                       state :: ST

end

include("state.jl")
include("nonlocal_flux.jl")
include("mixing_length.jl")
include("tke_equation.jl")
include("wall_models.jl")
include("diffusivities.jl")
include("convective_adjustment.jl")

function ModelBoundaryConditions(FT=Float64; U = DefaultBoundaryConditions(FT),
                                             V = DefaultBoundaryConditions(FT),
                                             T = DefaultBoundaryConditions(FT),
                                             S = DefaultBoundaryConditions(FT),
                                 )

    return (U=U, V=V, T=T, S=S)
end

function Model(;
                                   grid = UniformGrid(N, H),
                              constants = Constants(),
                          mixing_length = SimpleMixingLength(),
                     eddy_diffusivities = IndependentDiffusivities(),
                  convective_adjustment = nothing,
                         tke_wall_model = PrescribedSurfaceTKEFlux(),
                           tke_equation = TKEParameters(),
                   boundary_layer_depth = nothing,
                          nonlocal_flux = nothing,
               background_diffusivities = BackgroundDiffusivities(),
                                stepper = :BackwardEuler,
                                    bcs = ModelBoundaryConditions(eltype(grid))
              )

    bcs = merge(bcs, (e = TKEBoundaryConditions(eltype(grid), tke_wall_model),))

    solution = Solution((CellField(grid) for i=1:nsol)...)

    Kϕ = (U=KU, V=KV, T=KT, S=KS, e=Ke)
    Rϕ = (U=RU, V=RV, T=RT, S=RS, e=Re)
    Lϕ = (U=minuszero, V=minuszero, T=minuszero, S=minuszero, e=Le)
    eq = Equation(K=Kϕ, R=Rϕ, L=Lϕ, update=update_state!)
    lhs = OceanTurb.build_lhs(solution)

    timestepper = Timestepper(stepper, eq, solution, lhs)

    return Model(Clock(),
                 grid,
                 timestepper,
                 solution,
                 bcs,
                 mixing_length,
                 eddy_diffusivities,
                 convective_adjustment,
                 tke_wall_model,
                 tke_equation,
                 boundary_layer_depth,
                 nonlocal_flux,
                 background_diffusivities,
                 constants,
                 State(grid, mixing_length, boundary_layer_depth, nonlocal_flux)
                )
end

@inline RU(m, i) = @inbounds   m.constants.f * m.solution.V[i] - ∂z_NLᵁ(m, i)
@inline RV(m, i) = @inbounds - m.constants.f * m.solution.U[i] - ∂z_NLⱽ(m, i)
@inline RT(m, i) = - ∂z_NLᵀ(m, i)
@inline RS(m, i) = - ∂z_NLˢ(m, i)

@inline Re(m, i) = production(m, i) + buoyancy_flux(m, i) - dissipation(m, i) - ∂z_NLᵉ(m, i)
@inline Le(m, i) = 0

end # module


text:
function node(abaco::Context, ne, tag)
    # if tag is unknow then fallback to the default settings
    ##settings = get(abaco.cfg, tag, abaco.cfg[DEFAULT_TYPE])
    if abaco.interval == -1
        el = Node(ne, tag)
    else
        el = Node(ne, tag, abaco.ages)
    end

    if haskey(abaco.cfg, tag)
            for formula in values(abaco.cfg[tag].formula)
                for snap in values(el.snap)
                    snap.outputs[formula.output] = FormulaState(false, formula.output)
                end
            end
    end

    abaco.node[ne] = el
    el
end

function node(abaco::Context, ne::String, tag::String, parent::String)
    if parent !== ""
        container = abaco.node[parent]
        node(abaco, container, ne, tag)
    else
        node(abaco, ne, tag)
    end
end

function node(abaco::Context, target::Node, ne::String, tag::String)
    elem = node(abaco, ne, tag)

    if haskey(abaco.origins, target.ne)
        push!(abaco.origins[target.ne], elem)
    else
        abaco.origins[target.ne] = Set([elem])
    end
    abaco.target[ne] = (tag, target.ne)
    return elem
end

function delete_node(abaco, ne)
    delete!(abaco.node, ne)
end


text:
function one_hot(y)
    n = length(unique(y))
    if minimum(y) == 0
        return eye(n)[(y+1),:]
    elseif minimum(y) == -1 && n == 2
        y = trunc(Int64,(y + 1)/2+1)
        return eye(2)[y,:]
    end
    return eye(n)[y,:]
end

function softmax(x::Vector)
    x = exp(x)
    pos = x./sum(x)
    return indmax(pos)
end

function softmax(X::Matrix)
    n_sample = size(X,1)
    res = zeros(n_sample)
    for i = 1:n_sample
        res[i] = softmax(X[i,:])
    end
    return res
end

function calc_variance(X)
    n_sample = size(X,1)
    mean_ = repmat(mean(X, 1),n_sample,1)
    de_mean = X - mean_
    return 1/n_sample * diag(de_mean' * de_mean)
end

function get_random_subsets(X, y , n_subsets;replacement = true)
    n_sample, n_feature = size(X)

    X_y = hcat(X, y)
    X_y = X_y[shuffle(1:n_sample), :]
    subsample_size = trunc(Int64,n_sample / 2)
    if replacement == true
        subsample_size = n_sample
    end
    sets = zeros(n_subsets,subsample_size, n_feature + 1)
    for i = 1:n_subsets
        idx = sample(1:n_sample, subsample_size, replace = replacement)
        temp = X_y[idx,:]
        sets[i,:,:] = temp
    end
    return sets
end

function train_test_split(X,y;train_size=0.75,rand_seed=2135)
    srand(rand_seed)
    rand_indx = shuffle(1:size(X,1))
    train_num = Int(floor(size(X,1) * train_size))
    X_train = X[rand_indx[1:train_num],:]
    X_test = X[rand_indx[(train_num+1):end],:]
    y_train = y[rand_indx[1:train_num]]
    y_test = y[rand_indx[(train_num+1):end]]
    return  X_train, X_test, y_train, y_test
end

function calc_entropy(y)
    if size(y,2) > 1
        y = unhot(y)
    end
    feature_unique = unique(y)
    num_sample = length(y)
    entro = 0
    for i in feature_unique
        num_feature = sum(y .== i)
        p = num_feature / num_sample
        entro += - p * log2(p)
    end
    return entro
end

function batch_iter(X, batch_size = 64)
    n_sample = size(X,1)
    n_batch = floor(n_sample / (batch_size))
    batch_end = 0

    batch = []
    for i = 1n_batch
        batch_begin = (i-1)*n_batch
        batch_end = i * n_batch
        if i < n_batch
            push!(batch,X[batch_beginbatch_end, ])
        else
            push!(batch,X[batch_beginend, ])
        end
    end
    return batch
end

function euc_dist(x,y)
    return norm(x-y)
end

function l2_dist(X)
    sum_X = sum(X .* X, 1)
    return (-2 * X * X' + sum_X)' + sum_X
end

function unhot(predicted)
    """Convert one-hot representation into one column."""
    actual = zeros(size(predicted,1))
    for i = 1:size(predicted, 1)
        predicted_data = predicted[i,:]
        actual[i] = indmax(predicted_data)
    end
    return actual
end

function normalize_(X::Matrix)
    std_ = std(X, 1)
    mean_ = mean(X, 1)
    for i = 1:size(X,2)
        if std_[i] != 0
            X[:,i] = (X[:, i] - mean_[i])/std_[i]
        end
    end
    return X
end
function absolute_error(actual, predicted)
    return abs(actual - predicted)
end

function classification_error(actual, predicted)
    if size(actual,2) > 1 && length(actual) > 1
        actual = unhot(actual)
        predicted = unhot(predicted)
    end
    return sum(actual .!= predicted) / size(actual,1)
end

function accuracy(actual, predicted)
    return 1.0 - classification_error(actual, predicted)
end

function mean_absolute_error(actual, predicted)
    return mean(absolute_error(actual, predicted))
end

function squared_error(actual, predicted)
    return (actual - predicted) .^ 2
end

function squared_log_error(actual, predicted)
    return (log(actual + 1) - log(predicted + 1)) .^ 2
end

function mean_squared_log_error(actual, predicted)
    return mean(squared_log_error(actual, predicted))
end

function mean_squared_error(actual, predicted)
    return mean(squared_error(actual, predicted))
end

function root_mean_squared_error(actual, predicted)
    return sqrt(mean_squared_error(actual, predicted))
end

function root_mean_squared_log_error(actual, predicted)
    return sqrt(mean_squared_log_error(actual, predicted))
end

function logloss(actual, predicted)
    predicted = clamp(predicted, 0, 1 - 1e-15)
    loss = -sum(actual .* log(predicted))
    return loss / size(actual,1)
end

function hinge(actual, predicted)
    return mean(max(1. - actual .* predicted, 0.))
end

function binary_crossentropy(actual, predicted)
    predicted = clamp(predicted, 1e-15, 1 - 1e-15)
    return mean(-sum(actual .* log(predicted) -
                           (1 - actual) .* log(1 - predicted)))
end

function sigmoid_tanh(x)
    return 0.5 * (tanh(x) + 1)
end

function sigmoid(x)
    return 1./(1+exp(-x))
end

function sigmoid_prime(x)
    return sigmoid(x).*(1-sigmoid(x))
end

function make_cla(;n_samples = 1200, n_features = 10, n_classes = 2)
    X, y = dat.make_classification(n_samples=n_samples, n_features=n_features,
                               random_state=1111, n_classes= n_classes)
    # Convert y to {-1, 1}
    y = (y * 2) - 1
    X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8,
                                                        rand_seed=1111)
    X_train, X_test, y_train, y_test
end

function make_reg(;n_samples = 200,
                   n_features = 10)
    X, y = dat.make_regression(n_samples=n_samples, n_features=n_features, n_targets=1, noise=0.05,
                           random_state=1111, bias=0.5)
    X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8,
                                                        rand_seed=1111)
    X_train, X_test, y_train, y_test
end

function make_iris()
    data= dat.load_iris()
    X = data["data"]
    y = data["target"]
    X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8,
                                                        rand_seed=1111)
    X_train, X_test, y_train, y_test
end

function make_blo()
    X, y = dat.make_blobs(centers=4, n_samples=500, n_features=2,
                          random_state=42)
    return  X, y
end

function make_digits()
    data = dat.load_digits()
    X = data["data"]
    y = data["target"]
    X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8,
                                                         rand_seed=1111)
    X_train, X_test, y_train, y_test
end


text:
abstract type AbstractIFS <: Function end

abstract type AbstractIFSComplex <: AbstractIFS end

"""
An Iterated Function System formed by general Complex Contractions.
"""
struct IFSComplex{F <: Function} <: AbstractIFSComplex
  contractions::Vector{F}

  function IFSComplex{F}(funs::Vector{F}) where F <: Function
    new(deepcopy(funs))
  end
end

IFSComplex(funs::Vector{F}) where F <: Function = IFSComplex{F}(funs)

function (ifs::IFSComplex)(z::Number)
  n = rand(1:size(ifs))
  ifs.contractions[n](z), n
end

"""
An Iterated Function System formed by Complex Affine Contractions.
"""
struct IFSComplexAffine{F <: AbstractAffineTransformation} <: AbstractIFSComplex
  contractions::Vector{F}

  function IFSComplexAffine{F}(funs::Vector{F}) where F <: AbstractAffineTransformation
    new(deepcopy(funs))
  end
end

IFSComplexAffine(funs::Vector{F}) where F <: AbstractAffineTransformation =
IFSComplexAffine{F}(funs)

function (ifs::IFSComplexAffine)(z::Number)
  n = rand(1:size(ifs))
  ifs.contractions[n](z), n
end

contractions(ifs::AbstractIFS) = ifs.contractions
size(ifs::AbstractIFS) = length(ifs.contractions)


text:
function a_ideal(model::SAFT, z, v, T)
    x = z/sum(z)
    return sum(x[i]*log(z[i]/v) for i in model.components)-1
end


text:
# setup project paths -
const _PATH_TO_SRC = dirname(pathof(@__MODULE__))

# load external packages -
using Test
using HTTP
using DataFrames
using BioSequences
using ProgressMeter
using BSON

# load my codes -
include(joinpath(_PATH_TO_SRC, "Types.jl"))
include(joinpath(_PATH_TO_SRC, "Base.jl"))
include(joinpath(_PATH_TO_SRC, "Factory.jl"))
include(joinpath(_PATH_TO_SRC, "Network.jl"))
include(joinpath(_PATH_TO_SRC, "Reactions.jl"))
include(joinpath(_PATH_TO_SRC, "Compounds.jl"))
include(joinpath(_PATH_TO_SRC, "Sequence.jl"))
include(joinpath(_PATH_TO_SRC, "Organism.jl"))
include(joinpath(_PATH_TO_SRC, "ECNumbers.jl"))
include(joinpath(_PATH_TO_SRC, "Files.jl"))

# URL string -
const _KEGG_LINKAGE_REACTION_URL = "http://rest.kegg.jp/link/rn"
const _KEGG_LINKAGE_EC_URL = "http://rest.kegg.jp/link/ec"
const _KEGG_LIST_COMPOUND_URL = "http://rest.kegg.jp/list/compound"
const _KEGG_GET_URL = "http://rest.kegg.jp/get"
const _KEGG_LINK_URL = "http://rest.kegg.jp/link"
const _KEGG_LIST_URL = "http://rest.kegg.jp/list"


text:
function _convert(ctx::ConverterContext, decl::CLFunctionDecl)
	name = _convertName(ctx, decl)
	csym = repr(Symbol(spelling(decl)))

	if Clang.calling_conv(type(decl)) == LibClang.CXCallingConv_C
		convention = "CBinding.CDECL"
	else
		error("Support for function calling convention `$(Clang.calling_conv(type(decl)))` is not yet implemented")
	end

	ret = _convertType(ctx, result_type(decl))
	args = [CLCursor(LibClang.clang_Cursor_getArgument(decl.cursor, UInt(i-1))) for i in 1:argnum(decl)]
	names = map(arg -> _convertName(ctx, arg), args)
	types = map(arg -> _convertType(ctx, type(arg)), args)
	names = map(enumerate(names)) do (i, aname)
		if isempty(aname)
			aname = _gensym(ctx, string(i))
		end
		return aname
	end
	if isvariadic(decl)
		push!(names, "$(_gensym(ctx, string(length(names)+1)))...")
		push!(types, "Base.Vararg")
	end
	sig  = "CBinding.Cfunction{$(ret), Base.Tuple{$(join(types, ", "))}}"

	body = ""
	if isinlined(decl)
		todo"need inlined functions to be saved/stored differently so the closure lasts forever"
		@warn "Unable to convert inline function:  $(name)"
	elseif !(name in ctx.oneofs)
		push!(ctx.oneofs, name)

		_export(ctx, name)
		def = "$(name)($(join(names, ", "))) = $(_gensym(ctx, name))[]($(join([names..., "convention = $(convention)"], ", ")))"
		push!(ctx.converted, JuliaizedC(
			decl,
			def,
			:atcompile,
		))

		ptr = "const $(_gensym(ctx, name)) = Base.Ref(Base.Ptr{$(sig)}(Base.C_NULL))"
		push!(ctx.converted, JuliaizedC(
			decl,
			ptr,
			:atcompile_bindings,
		))

		load = "CBinding.bind($(_gensym(ctx, name)), $(csym), $(_gensym(ctx, "libraries"))...)"
		push!(ctx.converted, JuliaizedC(
			decl,
			load,
			:atload,
		))
	end
end


text: struct OutageStatus <: PSI.InitialConditionType end


text:
module ATools

include("math.jl")
include("util.jl")
include("histos.jl")
include("fits.jl")
include("raytracing.jl")
include("io.jl")
include("reco.jl")

export in_range,gline2p, rxy,phixy,fphi,dxyz,wstd,mean_std,fit,
       fit_pol1, fit_pol2,fit_pol3, fit_gauss,fit_gauss_fm,
       fitg1, fitg2, fit_profile, plot_fit_gauss  #math
export select_event,select_by_index,select_by_column_value,
       select_by_column_value_lt, select_by_column_value_le,
       select_by_column_value_gt, select_by_column_value_ge,
       select_by_column_value_interval,select_by_column_value_closed_interval,
       select_by_column_value_closed_left_interval,select_by_column_value_closed_right_interval,
       select_by_index,find_max_xy   #util
export hist1d, edges, centers, hist2d, p1df #histos
export readdf, read_abc, readh5_dset, readh5_todf, read_evtpar
#export lfit, fit_pol1, fit_pol2,fit_pol3, fit_gauss,fit_gauss_fm,
#       fitg1, fitg2, fit_profile #fits

end

max_stars_repo_name: jjgomezcadenas/ATools

max_issues_repo_name: jjgomezcadenas/ATools

max_forks_repo_name: jjgomezcadenas/ATools


text:
using Random

###############################################################################
#
#  Direct product decomposition
#
###############################################################################

function direct_product_decomposition(G::Main.ForeignGAP.MPtr, ab::Tuple{Int, Int})

  decompositions = Tuple{Main.ForeignGAP.MPtr, Main.ForeignGAP.MPtr}[]
  if GAP.Globals.IsAbelian(G)
    return ab, (1, 1), 1, 1
  end
  n = ab[1]
  subs = GAP.Globals.NormalSubgroups(G)
  #First, I collect all the possible decompositions
  decompositions = Tuple{Main.ForeignGAP.MPtr, Main.ForeignGAP.MPtr}[]
  for i = 1:length(subs)
    g1 = subs[i]
    if isone(GAP.Globals.Size(g1))
      continue
    end
    o1 = GAP.Globals.Size(g1)
    for j = 1:length(subs)
      g2 = subs[j]
      if isone(GAP.Globals.Size(g2))
        continue
      end
      o2 = GAP.Globals.Size(g2)
      if o1*o2 != n
        continue
      end
      if isone(GAP.Globals.Size(GAP.Globals.Intersection(g1, g2)))
        push!(decompositions, (g1, g2))
      end
    end
  end
  if isempty(decompositions)
    return ab, (1, 1) , 1, 1
  end

  #We pass to the list with the group ID
  grp_id_list = Array{Tuple{Main.ForeignGAP.MPtr, Main.ForeignGAP.MPtr}, 1}(undef, length(decompositions))
  for i = 1:length(grp_id_list)
    grp_id_list[i] = (GAP.Globals.IdGroup(decompositions[i][1]), GAP.Globals.IdGroup(decompositions[i][2]))
  end

  l1 = grp_id_list[1][1]
  a1 = l1[1]
  b1 = l1[2]
  l2 = grp_id_list[1][2]
  a2 = l2[1]
  b2 = l2[2]
  if length(grp_id_list) == 1
    return (a1, b1), (a2, b2), 1, 1
  end
  #I count the redundancy, i.e. the number of possible decompositions of the same type.
  red = 1
  for i = 2:length(grp_id_list)
    l1 = grp_id_list[i][1]
    a11 = l1[1]
    b11 = l1[2]
    l2 = grp_id_list[i][2]
    a21 = l2[1]
    b21 = l2[2]
    if min(a11, a21) > min(a1, a2)
      red = 1
      a1 = a11
      b1 = b11
      a2 = a21
      b2 = b21
    else
      if (a11, b11) == (a1, b1) && (a21, b21) == (a2, b2)
        red += 1
      end
    end
  end
  #Finally, I count the numer of times a single subgroup appears in the lists.
  redfirst = 1
  GID_first = [a1, b1]
  GID_second = [a2, b2]
  ind = 1
  for i = 1:length(decompositions)
    if GAP.Globals.IdGroup(decompositions[i][1]) == GID_first && GAP.Globals.IdGroup(decompositions[i][2]) == GID_second
      ind = i
      break
    end
  end
  for i = ind+1:length(decompositions)
    if decompositions[i][1] == decompositions[ind][1]
      redfirst += 1
    end
  end
  return (a1, b1), (a2, b2), red, redfirst

end

###############################################################################
#
#  Merging fields
#
###############################################################################

function _to_composite(x::FieldsTower, y::FieldsTower)

  Kns, mx, my = NumberField(x.field, y.field, cached = false, check = false)
  K, mK = simple_extension(Kns, check = false)

  #First, compute the maximal order
  O1 = maximal_order(x.field)
  B1 = basis(O1, copy = false)
  O2 = maximal_order(y.field)
  B2 = basis(O2, copy = false)
  imagesx = Array{nf_elem, 1}(undef, length(B1))
  imagesy = Array{nf_elem, 1}(undef, length(B2))
  for i = 1:degree(x.field)
    imagesx[i] = mK\(mx(B1[i].elem_in_nf))
  end
  for i = 1:degree(y.field)
    imagesy[i] = mK\(my(B2[i].elem_in_nf))
  end
  prod_bas = Hecke.product_basis(imagesx, imagesy)
  fac = factor(gcd(discriminant(O1), discriminant(O2)))
  lp = collect(keys(fac.fac))
  MatOrd = basis_matrix(prod_bas, FakeFmpqMat)
  discr = fmpz(1)
  if istotally_real(x.field) && istotally_real(y.field)
    discr = abs(discriminant(O1)^degree(O2) * discriminant(O2)^degree(O1))
  else
    discr = (-1)^divexact(degree(K), 2) * abs(discriminant(O1)^degree(O2) * discriminant(O2)^degree(O1))
  end
  inde = root(divexact(numerator(discriminant(K)), discr), 2)
  MatOrdNum = MatOrd.num
  Hecke.hnf_modular_eldiv!(MatOrdNum, MatOrd.den, :lowerleft)
  OK = NfAbsOrd(K, FakeFmpqMat(MatOrdNum, MatOrd.den))
  #Careful: we need to compute also the sign of the discriminant!
  OK.disc = discr
  OK.index = inde
  OK.gen_index = fmpq(OK.index)
  if !isempty(lp)
    OK = Hecke.pmaximal_overorder_at(OK, lp)
  end
  OK.ismaximal = 1
  Hecke._set_maximal_order_of_nf(K, OK)

  # Now, I have to translate the automorphisms.
  # Easy thing: first, I write the automorphisms of the non simple extension
  # Then translating them is straightforward.
  autK = Array{NfToNfMor, 1}(undef, length(x.generators_of_automorphisms)+ length(y.generators_of_automorphisms))
  el = mK.prim_img
  for i = 1:length(x.generators_of_automorphisms)
    ima = mx(x.generators_of_automorphisms[i].prim_img)
    autns = Hecke.NfAbsNSToNfAbsNS(Kns, Kns, NfAbsNSElem[ima, gens(Kns)[2]])
    ima = mK\(autns(el))
    autK[i] = NfToNfMor(K, K, ima)
  end
  for j = 1:length(y.generators_of_automorphisms)
    ima = my(y.generators_of_automorphisms[j].prim_img)
    autns = Hecke.NfAbsNSToNfAbsNS(Kns, Kns, NfAbsNSElem[gens(Kns)[1], ima])
    ima = mK\(autns(el))
    autK[j+length(x.generators_of_automorphisms)] = NfToNfMor(K, K, ima)
  end

  #Computing closure of the automorphisms
  all_autos = Vector{NfToNfMor}(undef, degree(K))
  ind = 1
  aut1 = automorphisms(x.field, copy = false)
  aut2 = automorphisms(y.field, copy = false)
  for a1 in aut1
    for a2 in aut2
      ima1 = mx(a1.prim_img)
      ima2 = my(a2.prim_img)
      autns = Hecke.NfAbsNSToNfAbsNS(Kns, Kns, NfAbsNSElem[ima1, ima2])
      ima = mK\(autns(el))
      all_autos[ind] = NfToNfMor(K, K, ima)
      ind += 1
    end
  end
  Hecke._set_automorphisms_nf(K, all_autos)

  #Last thing: I have to add the maps of the subfields!
  emb1 = NfToNfMor(x.field, K, mK\(mx.prim_img))
  emb2 = NfToNfMor(y.field, K, mK\(my.prim_img))
  l1 = append!(NfToNfMor[emb1, emb2], x.subfields)
  embs = append!(l1, y.subfields)
  return FieldsTower(K, autK, embs)

end

#merge function when all the fields are automatically linearly disjoint
function _easy_merge(list1, list2, absolute_bound::fmpz)

  res = FieldsTower[]
  for x in list1
    for y in list2
      check_bound_disc(x, y, absolute_bound) || continue
      composite = _to_composite(x, y)
      if abs(discriminant(maximal_order(composite.field))) <= absolute_bound
        push!(res, composite)
      end
    end
  end
  return res

end

function _disjoint_ab_subs(list1::Vector{FieldsTower}, list2::Vector{FieldsTower})

  deg_ab_sub1 = 1
  deg_ab_sub2 = 1
  for i = 1:length(list1[1].subfields)
    if degree(domain(list1[1].subfields[i])) == 1
      deg_ab_sub1 *= degree(codomain(list1[1].subfields[i]))
    end
  end
  for i = 1:length(list2[1].subfields)
    if degree(domain(list2[1].subfields[i])) == 1
      deg_ab_sub2 *= degree(codomain(list2[1].subfields[i]))
    end
  end
  return deg_ab_sub1, deg_ab_sub2, gcd(deg_ab_sub1, deg_ab_sub2) == 1

end

function check_bound_disc(K::FieldsTower, L::FieldsTower, bound::fmpz)
  dK = degree(K)
  dL = degree(L)
  discK = abs(discriminant(maximal_order(K)))
  discL = abs(discriminant(maximal_order(L)))
  d = gcd(discK, discL)
  if issquare(d)
    d = root(d, 2)
  end
  d = d^(2*dK*dL - 2)
  d1 = gcd(discK^dL, discL^dK)
  d = gcd(d, d1)
  n = discK^dL * discL^dK
  return div(n, d) <= bound
end

function maximal_abelian_subextension(F::FieldsTower)
  fields = AnticNumberField[]
  for x in F.subfields
    if degree(domain(x)) == 1
      push!(fields, codomain(x))
    end
  end
  return fields
end

function check_norm_group_and_disc(lfieldsK::Array{AnticNumberField, 1}, lfieldsL::Array{AnticNumberField, 1}, bound::fmpz)

  target_deg = prod(degree(x) for x in lfieldsK) * prod(degree(x) for x in lfieldsL)
  discK = lcm([discriminant(maximal_order(x)) for x in lfieldsK])
  discL = lcm([discriminant(maximal_order(x)) for x in lfieldsL])
  modulo = Int(lcm(discK, discL))
  y = PolynomialRing(QQ, "y", cached = false)[2]
  K = NumberField(y-1, cached = false)[1]
  O = maximal_order(K)
  n_quo1 = lcm([degree(x) for x in lfieldsK])
  n_quo2 = lcm([degree(x) for x in lfieldsL])
  r, mr = Hecke.ray_class_groupQQ(O, modulo, true, lcm(n_quo1, n_quo2))
  Kt = PolynomialRing(K, "t", cached = false)[1]
  h = change_ring(lfieldsK[1].pol, Kt)
  S = norm_group(h, mr)[1]
  for i = 2:length(lfieldsK)
    h = change_ring(lfieldsK[i].pol, Kt)
    s = norm_group(h, mr)[1]
    S = intersect(s, S)
  end
  for i = 1:length(lfieldsL)
    h = change_ring(lfieldsL[i].pol, Kt)
    s = norm_group(h, mr)[1]
    S = intersect(s, S)
  end
  Q, mQ = quo(r, S)
  if order(Q) != target_deg
    return false
  else
    C = ray_class_field(mr, mQ)
    return Hecke.discriminant_conductorQQ(O, C, modulo, bound, Int(order(Q)))
  end

end

function _first_sieve(list1::Vector{FieldsTower}, list2::Vector{FieldsTower}, absolute_bound::fmpz, redfirst::Int)
  ab1, ab2, fl = _disjoint_ab_subs(list1, list2)
  bound_max_ab_sub = root(absolute_bound, divexact(degree(list1[1])*degree(list2[1]), ab1*ab2))
  D = Dict{Tuple{Set{fmpz}, Set{fmpz}, Bool}, Array{Tuple{Int, Int}, 1}}() #The boolean true means real
  for i1 = 1:length(list1)
    @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())Combinations with field $(i1)/$(length(list1)) of the first list"
    @vprint :FieldsNonFancy 1 "Fields $(i1)/$(length(list1))\n"
    K = list1[i1]
    DK = Dict{Tuple{Set{fmpz}, Set{fmpz}}, Array{Int, 1}}()
    rK = Hecke.ramified_primes(maximal_order(K))
    lfieldsK = maximal_abelian_subextension(K)
    rK1 = Set(fmpz[])
    for x in lfieldsK
      rK1 = union(rK1, Hecke.ramified_primes(maximal_order(x)))
    end
    #First, I only check for disjointness
    for i2 = 1:length(list2)
      L = list2[i2]
      check_bound_disc(K, L, absolute_bound) || continue
      #First, I check if the discriminants are coprime
      rL = Hecke.ramified_primes(maximal_order(L))
      lfieldsL = maximal_abelian_subextension(L)
      rL1 = Set(fmpz[])
      for x in lfieldsL
        rL1 = union(rL1, Hecke.ramified_primes(maximal_order(x)))
      end
      if length(lfieldsL) == 1 && length(lfieldsK) == 1
        if degree(lfieldsL[1]) == 2 && degree(lfieldsK[1]) == 2
          if isempty(intersect(rL1, rK1)) || discriminant(maximal_order(lfieldsK[1])) != discriminant(maximal_order(lfieldsL[1]))
            k =  Set(vcat(rK, rL))
            k1 = union(rK1, rL1)
            if haskey(DK, (k, k1))
              push!(DK[(k, k1)],  i2)
            else
              DK[(k, k1)] = Int[i2]
            end
          end
        else
          if isempty(intersect(rL1, rK1)) || Hecke.islinearly_disjoint(lfieldsK[1], lfieldsL[1])
            k =  Set(vcat(rK, rL))
            k1 = union(rK1, rL1)
            if haskey(DK, (k, k1))
              push!(DK[(k, k1)],  i2)
            else
              DK[(k, k1)] = Int[i2]
            end
          end
        end
      else
        if isempty(intersect(rL1, rK1)) || check_norm_group_and_disc(lfieldsK, lfieldsL, bound_max_ab_sub)
          k =  Set(vcat(rK, rL))
          k1 = union(rK1, rL1)
          if haskey(DK, (k, k1))
            push!(DK[(k, k1)],  i2)
          else
            DK[(k, k1)] = Int[i2]
          end
        end
      end
    end
    #Now, I look at the length of the lists. I know that this must be divisible by redfirst...
    for (k, v) in DK
      if length(v) < redfirst
        continue
      else
        if !istotally_real(K.field)
          ar = Array{Tuple{Int, Int}, 1}(undef, length(v))
          for j = 1:length(v)
            ar[j] = (i1, v[j])
          end
          if haskey(D, (k[1], k[2], false))
            append!(ar, D[(k[1], k[2], false)])
          end
          D[(k[1], k[2], false)] = ar
        else
          ar_real = Array{Tuple{Int, Int}, 1}()
          ar_not_real = Array{Tuple{Int, Int}, 1}()
          for j = 1:length(v)
            if istotally_real(list2[v[j]].field)
              push!(ar_real, (i1, v[j]))
            else
              push!(ar_not_real, (i1, v[j]))
            end
          end
          if !isempty(ar_real)
            if haskey(D, (k[1], k[2], true))
              append!(ar_real, D[(k[1], k[2], true)])
            end
            D[(k[1], k[2], true)] = ar_real
          end
          if !isempty(ar_not_real)
            if haskey(D, (k[1], k[2], false))
              append!(ar_not_real, D[(k[1], k[2], false)])
            end
            D[(k[1], k[2], false)] = ar_not_real
          end
        end
      end
    end
  end
  return D
end

function _second_sieve(list1::Vector{FieldsTower}, list2::Vector{FieldsTower}, absolute_bound::fmpz, red::Int, redsecond::Int, D::Dict)
  redfirst = divexact(red, redsecond)
  res = FieldsTower[]
  indD = 0
  lD = length(D)
  for v in values(D)
    indD += 1
    @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())Cluster $(indD)/$(lD) of length $(length(v))"
    @vprint :FieldsNonFancy 1 "Cluster $(indD)/$(lD) of length $(length(v))\n"
    if length(v) < red
      continue
    elseif length(v) == red
      @vtime :Fields 2 composite = _to_composite(list1[v[1][1]], list2[v[1][2]])
      disccomp = abs(discriminant(maximal_order(composite)))
      if disccomp <= absolute_bound
        push!(res, composite)
      end
    else
      number_expected_fields = div(length(v), red)
      to_discard = 0
      small_res = FieldsTower[]
      small_res_ind = Tuple{Int, Int}[]
      permu = randperm(length(v))
      for idx = 1:length(v)
        if length(small_res) == number_expected_fields || idx + to_discard > length(v)
          break
        end
        i = permu[idx]
        #I compute the field only if at least redfirst first fields appear in the list
        # and at least red/redfirst second fields appear
        ind1 = v[i][1]
        cnt = 1
        for j = i+1:length(v)
          if v[j][1] == ind1
            cnt += 1
          end
          if cnt == redfirst
            break
          end
        end
        if cnt != redfirst
          to_discard -= 1
          continue
        end
        ind2 = v[i][2]
        cnt = 1
        for j = i+1:length(v)
          if v[j][2] == ind2
            cnt += 1
          end
          if cnt == redsecond
            break
          end
        end
        if cnt != redsecond
          to_discard -= 1
          continue
        end
        @vtime :Fields 2 composite = _to_composite(list1[v[i][1]], list2[v[i][2]])
        disccomp = abs(discriminant(maximal_order(composite)))
        if disccomp <= absolute_bound
          #Check if I already have that field.
          found = false
          for j = 1:length(small_res)
            fj = small_res[j].field
            if first_check_isom(fj, composite.field)
              if v[i][1] != small_res_ind[j][1] && !Hecke.issubfield_normal(list1[v[i][1]].field, fj)[1]
                continue
              elseif v[i][2] != small_res_ind[j][2] && !Hecke.issubfield_normal(list2[v[i][2]].field, fj)[1]
                continue
              else
                found = true
                break
              end
            end
          end
          if found
            to_discard -= 1
            continue
          end
          to_discard += redfirst - 1
          push!(small_res, composite)
          push!(small_res_ind, (v[i][1], v[i][2]))
        end
      end
      append!(res, small_res)
    end
  end
  return res

end

function _merge(list1::Vector{FieldsTower}, list2::Vector{FieldsTower}, absolute_bound::fmpz, red::Int, redsecond::Int)

  if red == 1
    #All the fields are automatically linearly disjoint
    @vprint :Fields 1 "All the fields are linearly disjoint, easy case \n"
    @vprint :FieldsNonFancy 1 "All the fields are linearly disjoint, easy case \n"
    return _easy_merge(list1, list2, absolute_bound)
  end

  redfirst = divexact(red, redsecond)
  if redfirst < redsecond
    #switch the lists
    return _merge(list2, list1, absolute_bound, red, redfirst)
  end
  #Bad case: a mess.
  #I check that the maximal abelian subextensions are linearly disjoint.
  #Working with polynomials may be more expensive.
  #Since I am working with abelian extensions, I can compute the norm groups and see there if the norm groups are correct.
  @vprint :Fields 1 "Merging fields\n"
  @vprint :Fields 1 "Number of fields: $(length(list1)*length(list2))\n"
  @vprint :FieldsNonFancy 1 "Merging fields\n"
  @vprint :FieldsNonFancy 1 "Number of fields: $(length(list1)*length(list2))\n"
  @vprint :Fields 1 "Redundancy: $(red), $(redfirst), $(redsecond)\n"
  @vprint :FieldsNonFancy 1 "Redundancy: $(red), $(redfirst), $(div(red, redfirst))\n"

  res = FieldsTower[]
  D = _first_sieve(list1, list2, absolute_bound, redfirst)
  if isempty(D)
    return res
  end
  @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())"

  nfields_after = sum(length(v) for v in values(D))
  @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())Fields after first part of sieving: $(nfields_after)\n"
  @vprint :FieldsNonFancy 1 "Fields after first part of sieving: $(nfields_after)\n"

  res = _second_sieve(list1, list2, absolute_bound, red, redsecond, D)

  @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())Fields found: $(length(res))\n"
  @vprint :FieldsNonFancy 1 "Fields found: $(length(res))\n"
  return res

end

function first_check_isom(K::AnticNumberField, L::AnticNumberField)
  OK = maximal_order(K)
  OL = maximal_order(L)
  if discriminant(OK) != discriminant(OL)
    return false
  end

  p = 200
  cnt = 1
  while cnt < 20
    cnt += 1
    p = next_prime(p)
    lp1 = prime_decomposition_type(OK, p)
    lp2 = prime_decomposition_type(OL, p)
    if length(lp1) != length(lp2)
      return false
    end
    if lp1[1] != lp2[1]
      return false
    end
  end
  return true
end