module louisenlund
    using Random
    include("./fixed_parameters.jl")
    using .fixed_parameters
    include("./cohort.jl")
    using .cohort
    include("house.jl")
    using .house

    function get_num_ex_teachers(houses, class_info)
        num_ex_teachers = 0
        for info in class_info
            num_ex_teachers += info[2]
        end
        for house in houses
            num_ex_teachers -= length(house.crew[2])
        end
        return num_ex_teachers
    end

    function get_house_walls(index)
        if index==1
            return [([0,0], [lund_size[2]-hwh,lund_size[2]]), lundwalls[1]]
        elseif index==num_houses
            return [lundwalls[num_houses-1], ([lund_size[1],lund_size[1]], [lund_size[2]-hwh,lund_size[2]])]
        else
            return [lundwalls[index-1], lundwalls[index]]
        end
    end

    function get_classroom_walls(index)
        if index==1
            return [([0,0], [0,cwh]), lundwalls[num_houses]]
        elseif index==num_classrooms
            return [lundwalls[end], ([lund_size[1],lund_size[1]], [0,cwh])]
        else
            return [lundwalls[num_houses+index-2], lundwalls[num_houses+index-1]]
        end
    end

    function extend!(aTuple::Tuple)
        result = aTuple[1]
        for i in 2:length(aTuple)
            for j in aTuple[i]
                push!(result, j)
            end
        end
        return result
    end

    mutable struct Louisenlund
        intern_random_infection_rate::Float64
        extern_random_infection_rate::Float64
        teacher_break_rules::Float64
        student_break_rules::Float64
        num_cohorts::Int64

        cohorts::Array
        houses::Array
        classes::Array
        fakewalls::Array
        classrooms::Array
        chosen_classes::Array

        tick::Int64
        infected_list::Array
        healthy_list::Array
        regulation::Bool
        quarantine::Int64

        ticking::Any
        start_infection::Any

        function Louisenlund(num_cohorts::Int64, regulation::Bool, intern_random_infection_rate::Float64, extern_random_infection_rate::Float64, teacher_break_rules::Float64, student_break_rules::Float64)
            Weiden = House(get_house_walls(1), 16, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Kuh = House(get_house_walls(2), 16, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Meierei = House(get_house_walls(3), 20, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Allee = House(get_house_walls(4), 12, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Eschen = House(get_house_walls(5), 21, 2, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            House_E_9 = [Weiden, Kuh, Meierei, Allee, Eschen]
            class_info_E_9 = [(15,3), (10,2), (7,2), (15,3), (11,3), (8,2), (9,2), (15,3), (12,3)]
            num_ex_t_E_9 = get_num_ex_teachers(House_E_9, class_info_E_9)

            Fuchsbau_Q1 = House(get_house_walls(24), 4, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Förse_unten = House(get_house_walls(22), 14, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Ahorn = House(get_house_walls(15), 12, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Glocken = House(get_house_walls(10), 7, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Gärtnerei = House(get_house_walls(17), 6, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Eichen_oben = House(get_house_walls(18), 14, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Linden = House(get_house_walls(20), 17, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Förse_oben = House(get_house_walls(23), 21, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            House_Q1 = [Fuchsbau_Q1, Förse_unten, Ahorn, Glocken, Gärtnerei, Eichen_oben, Linden, Förse_oben]
            class_info_Q1 = [(12,3), (12,3), (16,3), (14,3), (17,3), (17,3), (7,2)]
            num_ex_t_Q1 = get_num_ex_teachers(House_Q1, class_info_Q1)

            Pibo = House(get_house_walls(16), 13, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Marstall = House(get_house_walls(8), 3, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Hausmeisterei = House(get_house_walls(7), 2, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Fuchsbau_Q2 = House(get_house_walls(25), 4, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Wald = House(get_house_walls(13), 11, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Buchen = House(get_house_walls(21), 7, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Gilden = House(get_house_walls(14), 16, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Schloss_rechts = House(get_house_walls(12), 7, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            House_Q2 = [Pibo, Marstall, Hausmeisterei, Fuchsbau_Q2, Wald, Buchen, Gilden, Schloss_rechts]
            class_info_Q2 = [(10,2), (17,3), (17,3), (8,2), (15,3)]
            num_ex_t_Q2 = get_num_ex_teachers(House_Q2, class_info_Q2)

            Birken = House(get_house_walls(6), 14, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Schloss_links = House(get_house_walls(11), 6, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Eichen_unten = House(get_house_walls(19), 11, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Fuchsbau_IB = House(get_house_walls(26), 4, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            Kava = House(get_house_walls(9), 23, 1, intern_random_infection_rate, student_break_rules, teacher_break_rules)
            House_IB = [Birken, Schloss_links, Eichen_unten, Fuchsbau_IB, Kava]
            class_info_IB = [(11,3), (29,6), (18,4)]
            num_ex_t_IB = get_num_ex_teachers(House_IB, class_info_IB)

            if num_cohorts==1
                my_cohorts = [Cohort(21, num_ex_t_E_9+num_ex_t_Q1+num_ex_t_Q2+num_ex_t_IB, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, extend!((House_E_9,House_Q1,House_Q2,House_IB)), extend!((class_info_E_9,class_info_Q1,class_info_Q2,class_info_IB)))]
            elseif num_cohorts==2
                E_9 = Cohort(17, num_ex_t_E_9, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, House_E_9, class_info_E_9)
                Q1_Q2_IB = Cohort(4, num_ex_t_Q1+num_ex_t_Q2+num_ex_t_IB, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, extend!((House_Q1,House_Q2,House_IB)), extend!((class_info_Q1,class_info_Q2,class_info_IB)))
                my_cohorts = [E_9, Q1_Q2_IB]
            elseif num_cohorts==3
                E_9 = Cohort(17, num_ex_t_E_9, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, House_E_9, class_info_E_9)
                Q1_Q2 = Cohort(4, num_ex_t_Q1+num_ex_t_Q2, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, extend!((House_Q1,House_Q2)), extend!((class_info_Q1,class_info_Q2)))
                IB = Cohort(0, num_ex_t_IB, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, House_IB, class_info_IB)
                my_cohorts = [E_9, Q1_Q2, IB]
            elseif num_cohorts==4
                E_9 = Cohort(17, num_ex_t_E_9, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, House_E_9, class_info_E_9)
                Q1 = Cohort(0, num_ex_t_Q1, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, House_Q1, class_info_Q1)
                Q2 = Cohort(4, num_ex_t_Q2, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, House_Q2, class_info_Q2)
                IB = Cohort(0, num_ex_t_IB, intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, House_IB, class_info_IB)
                my_cohorts = [E_9, Q1, Q2, IB]
            end

            my_classes = Any[]
            my_houses = Any[]
            for i in my_cohorts
                my_classes = extend!((my_classes, i.classes))
                my_houses = extend!((my_houses, i.houses))
            end
            my_fakewalls = Tuple[]
            for i in my_cohorts
                for j in i.houses
                    for k in i.houses
                        if j!=k
                            for w in j.walls
                                if w in k.walls && !(w in my_fakewalls)
                                    push!(my_fakewalls, w)
                                end
                            end
                        end
                    end
                end
            end

            new(intern_random_infection_rate, extern_random_infection_rate, teacher_break_rules, student_break_rules, num_cohorts, my_cohorts, my_houses, my_classes, my_fakewalls, Vector(1:num_classrooms), Any[], 0, Any[], Any[], regulation, 0, ticking, start_infection)
        end
    end

    function ticking(self::Louisenlund, move_intensity::Int64, free_from_quarantine::Float64, time_of_lessons::Float64, num_lessons::Int64, time_in_house::Float64)
        self.tick += 1

        if self.tick%(day_tick/move_intensity)<time_of_lessons*day_tick/move_intensity
            if (self.tick%(day_tick/move_intensity))%(time_of_lessons/num_lessons*day_tick)==0
                empty!(self.chosen_classes)
            end
            if self.chosen_classes==[]
                Random.shuffle!(self.classrooms)
                Random.shuffle!(self.classes)
                self.chosen_classes = self.classes[1:20]
            end
            classrooms = deepcopy(self.classrooms)
            for i in self.cohorts
                if !self.regulation || (self.regulation && !i.quarantine)
                    for c in i.classes
                        if c in self.chosen_classes
                            #println(classrooms[end])
                            c.lessoning(c, get_classroom_walls(pop!(classrooms)), self.fakewalls)
                        else
                            for j in c.crew
                                for k in j
                                    if k.intern
                                        k.destination = nothing
                                        k.move(k, self.fakewalls)
                                        k.step = move_speed
                                    else
                                        k.position = (nothing, nothing)
                                    end
                                end
                            end
                        end
                    end
                else
                    i.house_time(i, self.fakewalls)
                end
            end
        else
            for i in self.cohorts
                i.house_time(i, self.fakewalls)
            end
            empty!(self.chosen_classes)
        end

        for (i, j, k) in self.infected_list
            for (a, b, c) in self.healthy_list
                if self.tick%infection_period==0 && !self.cohorts[i].crew[j][k].in_quarantine && self.cohorts[i].crew[j][k].position!=(nothing,nothing) && self.cohorts[a].crew[b][c].position!=(nothing,nothing) && self.cohorts[i].crew[j][k].contact(self.cohorts[i].crew[j][k], self.cohorts[a].crew[b][c])
                    push!(self.infected_list, (a, b, c))
                    filter!(e->e!=(a, b, c), self.healthy_list)
                end
            end

            self.cohorts[i].crew[j][k].infected_time += 1
            if self.cohorts[i].crew[j][k].in_quarantine
                self.cohorts[i].crew[j][k].quarantine_time += 1
            end

            if !self.cohorts[i].crew[j][k].show_infection_sign
                if self.regulation && self.tick%day_tick==0
                    show_infection_sign = get_show_infection_sign_rate(self.cohorts[i].crew[j][k].infected_time/day_tick)
                    self.cohorts[i].crew[j][k].show_infection_sign = rand()<=show_infection_sign ? true : false
                else
                    self.cohorts[i].crew[j][k].show_infection_sign = false
                end
            end
            if !self.cohorts[i].crew[j][k].in_quarantine && self.cohorts[i].crew[j][k].show_infection_sign && (rand()<=1-free_from_quarantine || self.cohorts[i].quarantine)
                self.cohorts[i].crew[j][k].in_quarantine = true
                self.quarantine += 1
                if self.regulation
                    self.cohorts[i].set_Quarantine(self.cohorts[i])
                    for z in self.cohorts
                        z.warn(z)
                    end
                end
            end

            if self.cohorts[i].crew[j][k].infected_time > end_infection_time*day_tick
                self.cohorts[i].crew[j][k].infected = false
                self.cohorts[i].crew[j][k].show_infection_sign = false
                if self.cohorts[i].crew[j][k].in_quarantine
                    self.quarantine -= 1
                    self.cohorts[i].crew[j][k].in_quarantine = false
                    cure_rate = get_cure_rate((self.cohorts[i].crew[j][k].infected_time-self.cohorts[i].crew[j][k].quarantine_time)/day_tick)
                    if rand() <= cure_rate
                        self.cohorts[i].crew[j][k].immunized = true
                    else
                        self.cohorts[i].crew[j][k].alive = false
                    end
                else
                    if !self.cohorts[i].crew[j][k].show_infection_sign || rand()<=1-death_rate
                        self.cohorts[i].crew[j][k].immunized = true
                    else
                        self.cohorts[i].crew[j][k].alive = false
                    end
                end
                filter!(e->e!=(i, j, k), self.infected_list)
            end
        end

        for (i, c) in enumerate(self.cohorts)
            if c.quarantine
                c.quarantine_time += 1
                if c.quarantine_time%day_tick==0
                    for (x, y, z) in self.infected_list
                        if x==i && !self.cohorts[x].crew[y][z].in_quarantine && rand()<=1-fake_negative_test
                            self.cohorts[x].crew[y][z].in_quarantine = true
                            self.quarantine += 1
                        end
                    end
                    c.quarantine = false
                    c.quarantine_time = 0
                    for h in c.houses
                        h.quarantine = false
                    end
                end
            end
        end
    end

    function start_infection(self::Louisenlund)
        i = rand(1:self.num_cohorts)
        j = rand(1:2)
        k = rand(1:length(self.cohorts[i].crew[j]))
        self.cohorts[i].crew[j][k].infected = true

        push!(self.infected_list, (i, j, k))
        for a in 1:self.num_cohorts
            for b in 1:2
                for c in 1:length(self.cohorts[a].crew[b])
                    if (a, b, c) != (i, j, k)
                        push!(self.healthy_list, (a, b, c))
                    end
                end
            end
        end
    end

    export Louisenlund

end


using Base.Broadcast: BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted

struct StridedArrayStyle{N} <: AbstractArrayStyle{N}
end

Broadcast.BroadcastStyle(::Type{<:AbstractStridedView{<:Any,N}}) where {N} =
    StridedArrayStyle{N}()

StridedArrayStyle(::Val{N}) where {N} = StridedArrayStyle{N}()
StridedArrayStyle{M}(::Val{N}) where {M,N} = StridedArrayStyle{N}()

Broadcast.BroadcastStyle(a::StridedArrayStyle, ::DefaultArrayStyle{0}) = a
Broadcast.BroadcastStyle(::StridedArrayStyle{N}, a::DefaultArrayStyle) where {N} =
    BroadcastStyle(DefaultArrayStyle{N}(), a)
Broadcast.BroadcastStyle(::StridedArrayStyle{N}, ::Broadcast.Style{Tuple}) where {N} =
    DefaultArrayStyle{N}()

function Base.similar(bc::Broadcasted{<:StridedArrayStyle{N}}, ::Type{T}) where {N,T}
    StridedView(similar(convert(Broadcasted{DefaultArrayStyle{N}}, bc), T))
end

Base.dotview(a::AbstractStridedView{<:Any,N}, I::Vararg{SliceIndex,N}) where {N} =
    getindex(a, I...)

# Broadcasting implementation
@inline function Base.copyto!(dest::AbstractStridedView{<:Any,N},
                                bc::Broadcasted{StridedArrayStyle{N}}) where {N}
    # convert to map

    # flatten and only keep the AbstractStridedView arguments
    # promote AbstractStridedView to have same size, by giving artificial zero strides
    stridedargs = promoteshape(size(dest), capturestridedargs(bc)...)
    c = make_capture(bc)
    _mapreduce_fuse!(c, nothing, nothing, size(dest), (dest, stridedargs...))
    return dest
end

const WrappedScalarArgs = Union{AbstractArray{<:Any,0}, Ref{<:Any}}

@inline capturestridedargs(t::Broadcasted, rest...) =
    (capturestridedargs(t.args...)..., capturestridedargs(rest...)...)
@inline capturestridedargs(t::AbstractStridedView, rest...) =
    (t, capturestridedargs(rest...)...)
@inline capturestridedargs(t, rest...) = capturestridedargs(rest...)
@inline capturestridedargs() = ()

# broadcast by promoting size 1 dimensions to full size
# making the stride in that dimension zero
promoteshape(sz::Dims, a1::AbstractStridedView, As...) = (promoteshape1(sz, a1), promoteshape(sz, As...)...)
promoteshape(sz::Dims) = ()
function promoteshape1(sz::Dims{N}, a::StridedView) where {N}
    newstrides = ntuple(Val(N)) do d
        if size(a, d) == sz[d]
            stride(a, d)
        elseif size(a, d) == 1
            0
        else
            throw(DimensionMismatch("array could not be broadcasted to match destination"))
        end
    end
    return StridedView(a.parent, sz, newstrides, a.offset, a.op)
end
function promoteshape1(sz::Dims{N}, a::UnsafeStridedView) where {N}
    newstrides = ntuple(Val(N)) do d
        if size(a, d) == sz[d]
            stride(a, d)
        elseif size(a, d) == 1
            0
        else
            throw(DimensionMismatch("array could not be broadcast to match destination"))
        end
    end
    return UnsafeStridedView(a.ptr, sz, newstrides, a.offset, a.op)
end

struct CaptureArgs{F, Args<:Tuple}
    f::F
    args::Args
end
struct Arg
end

# construct CaptureArgs
@inline function make_capture(bc::Broadcasted)
    args = make_tcapture(bc.args)
    CaptureArgs(bc.f, args)
end
@inline make_tcapture(t::Tuple{}) = t
@inline make_tcapture(t::Tuple) = (make_capture(t[1]), make_tcapture(tail(t))...)
@inline make_capture(a::WrappedScalarArgs) = a[]
@inline make_capture(a::AbstractStridedView) = Arg()
@inline make_capture(a) = a

# Evaluate CaptureArgs
(c::CaptureArgs)(vals...) = consume(c, vals)[1]
@inline function consume(c::CaptureArgs{F,Args}, vals) where {F,Args}
    args, newvals = t_consume(c.args, vals)
    return c.f(args...), newvals
end
@inline consume(a::Arg, vals::Tuple) = vals[1], tail(vals)
@inline consume(a, vals) = a, vals
@inline t_consume(t::Tuple{}, vals) = t, vals
@inline function t_consume(t::Tuple, vals)
    t1, newvals1 = consume(t[1], vals)
    ttail, newvals = t_consume(tail(t), newvals1)
    return (t1, ttail...), newvals
end


module InformationMeasures

include("../src/Discretization.jl")
include("../src/Estimators.jl")
include("../src/Formulae.jl")
include("../src/Measures.jl")

end # module


module Indexing

using DocStringExtensions
using TriMatrices: TriLayout, TriLower, TriUpper, TriSymmetric,
                   nelems, hasdiag

include("math.jl")
include("conversion.jl")
include("tri_indices.jl")

end  # module


# Testing the LeafModels struct
A = Fvcb()
g0 = 0.03; g1 = 12.0
Gs = Medlyn(g0,g1) # Instance of a Medlyn type with g0 = 0.03 and g1 = 0.1

@testset "LeafModels()" begin
    leaf = LeafModels(photosynthesis = A, stomatal_conductance = Gs)
    @test typeof(leaf) == LeafModels{Missing,Missing,Fvcb{Float64},Medlyn{Float64},
        MutableNamedTuples.MutableNamedTuple{(:PPFD, :Tₗ, :Cₛ, :A, :Gₛ, :Cᵢ, :Dₗ),
        NTuple{7,Base.RefValue{Float64}}}}

    @test typeof(leaf.photosynthesis) == Fvcb{Float64}
    @test typeof(leaf.stomatal_conductance) == Medlyn{Float64}
    @test leaf.photosynthesis.Tᵣ == 25.0
    @test leaf.stomatal_conductance.g0 ≈ g0
    @test leaf.stomatal_conductance.g1 ≈ g1
end;

@testset "init_status!" begin
    leaf = LeafModels(photosynthesis = A, stomatal_conductance = Gs)
    @test leaf.status.Tₗ == -999.99

    init_status!(leaf, Tₗ = 25.0)
    @test leaf.status.Tₗ == 25.0

    model = read_model("inputs/models/plant_coffee.yml")
    init_status!(model, Tₗ = 25.0)
    @test model["Leaf"].status.Tₗ == 25.0
end;

@testset "Vars to initialise" begin
    leaf = LeafModels(photosynthesis = A, stomatal_conductance = Gs)
    @test to_initialise(leaf) == [:PPFD, :Tₗ, :Cₛ, :Dₗ]
    @test to_initialise(leaf) == to_initialise(A,Gs)
    @test to_initialise(A) == [:PPFD, :Tₗ, :Cₛ]

    @test leaf.status.Tₗ == -999.99
    @test is_initialised(leaf) == false

    leaf = LeafModels(photosynthesis = A, stomatal_conductance = Gs,
        Tₗ = 25.0, PPFD = 1000.0, Cₛ = 400.0, Dₗ = 1.2)

    @test is_initialised(leaf) == true
end;


type Timedata{T} <: AbstractTimedata
    vals::DataFrame
    idx::Array{T, 1}

    function Timedata(vals::DataFrame, idx::Array{T, 1})
        chkIdx(idx)
        if(size(vals, 1) != length(idx))
            if (length(idx) == 0) | (size(vals, 1) == 0)
                df = convert(DataFrame, rand(2, 2))
                return new(df[1:2, [false, false]], Array{T, 1}[])
            end
            error(length(idx), " idx entries, but ", size(vals, 1), " rows of data")
        end
        return new(vals, idx)
    end
end

function Timedata{T}(vals::DataFrame, idx::Array{T, 1})
    return Timedata{T}(vals, idx)
end

#######################
## find boolean true ##
#######################

import Base.find
function find(td::TimeData.Timedata)
    ## apply find to column DataArrays
    (allRowInds, allColInds) = find2sub(td)
    return sub2ind(size(td), allRowInds, allColInds)
end

function find2sub(td::TimeData.Timedata)
    ## apply find to column DataArrays and return tuple of subscripts
    allRowInds = Int[]
    allColInds = Int[]
    for (name, col) in eachcol(td.vals)
        rowInds = find(col)
        colInd = td.vals.colindex[name]
        colInds = repmat([colInd], length(rowInds))
        append!(allRowInds, [rowInds])
        append!(allColInds, [colInds])
    end
    return (allRowInds, allColInds)
end

function find(f::Function, td::TimeData.AbstractTimedata)
    ## apply function elementwise to td
    ## apply find to column DataArrays

    (allRowInds, allColInds) = find2sub(f, td)
    return sub2ind(size(td), allRowInds, allColInds)
end

function find2sub(f::Function, td::TimeData.AbstractTimedata)
    ## apply function elementwise to td
    ## apply find to column DataArrays

    (nObs, nVars) = size(td)

    allRowInds = Int[]
    allColInds = Int[]

    for ii=1:nObs
        for jj=1:nVars
            fval = f(get(td, ii, jj))
            if !isa(fval, Bool) & !isna(fval)
                error("function must return boolean value or NA")
            end
            if isequal(fval, true)
                push!(allRowInds, ii)
                push!(allColInds, jj)
            end
        end
    end
    return (allRowInds, allColInds)
end


function xwalkforward(x)
  2x
end


# level 1 - generic routines

function broutine2x2!(st::AbstractVecOrMat, U::AbstractMatrix, locs::Locations)
    U11 = U[1, 1]; U12 = U[1, 2];
    U21 = U[2, 1]; U22 = U[2, 2];
    for loc in locs
        _broutine2x2!(st, (U11,U12,U21,U22), loc)
    end
    return st
end

# NOTE:
# wrapping the following code seems to have a small overhead
# so we are not wrapping them in a function here
function broutine2x2_kernel_expr(idx_1, idx_2)
    return quote
        ST1 = U11 * st[$idx_1] + U12 * st[$idx_2]
        ST2 = U21 * st[$idx_1] + U22 * st[$idx_2]

        st[$idx_1] = ST1
        st[$idx_2] = ST2
    end
end

function broutine2x2_m_kernel_expr(idx_1, idx_2)
    return quote
        ST1 = U11 * st[b, $idx_1] + U12 * st[b, $idx_2]
        ST2 = U21 * st[b, $idx_1] + U22 * st[b, $idx_2]

        st[b, $idx_1] = ST1
        st[b, $idx_2] = ST2
    end
end

# ST1 = U11 * st[j+i] + U12 * st[j+i+step_1]
# ST2 = U21 * st[j+i] + U22 * st[j+i+step_1]
# st[j+i] = ST1
# st[j+i+step_1] = ST2
body_broutine2x2 = quote
    step_1 = 1 << (plain(locs) - 1)
    step_2 = 1 << plain(locs)

    @inbounds if step_1 == 1
        for j in 0:step_2:size(st, 1)-step_1
            $(broutine2x2_kernel_expr(:(j+1), :(j+2)))
        end
        return st
    end
end

body_broutine2x2_m = quote
    step_1 = 1 << (plain(locs) - 1)
    step_2 = 1 << plain(locs)

    @inbounds if step_1 == 1
        for j in 0:step_2:size(st, 2)-step_1, b in 1:size(st, 1)
            $(broutine2x2_m_kernel_expr(:(j+1), :(j+2)))
        end
        return st
    end
end

init_broutine2x2_ctrl = quote
    raw = first(plain(locs))
    raw_ctrl = plain(ctrl)
    step_1 = 1 << (raw - 1)
    step_2 = 1 << raw
    ctrl_mask = bmask(raw_ctrl)
    flag_mask = reduce(+, one(Int) << (raw_ctrl[i] - 1) for i in 1:length(ctrl) if ctrl.flags[i])
end

body_broutine2x2_ctrl = quote
    $init_broutine2x2_ctrl
    @inbounds if step_1 == 1
        for j in 0:step_2:size(st, 1)-step_1
            if ismatch(j, ctrl_mask, flag_mask)
                $(broutine2x2_kernel_expr(:(j+1), :(j+2)))
            end
        end
        return st
    end
end

body_broutine2x2_ctrl_m = quote
    $init_broutine2x2_ctrl
    @inbounds if step_1 == 1
        for j in 0:step_2:size(st, 2)-step_1, b in 1:size(st, 1)
            if ismatch(j, ctrl_mask, flag_mask)
                $(broutine2x2_m_kernel_expr(:(j+1), :(j+2)))
            end
        end
        return st
    end
end

# case 1: we expand the inner loop of step_1 if it's smaller than 16
for k in [2,4,8,16]
    push!(body_broutine2x2.args, quote
        @inbounds if step_1 == $k
            for j in 0:step_2:size(st, 1)-step_1
                Base.Cartesian.@nexprs $k i-> begin
                    $(broutine2x2_kernel_expr(:(j+i), :(j+i+step_1)))
                end
            end
            return st
        end
    end)

    push!(body_broutine2x2_m.args, quote
        @inbounds if step_1 == $k
            for j in 0:step_2:size(st, 2)-step_1, b in 1:size(st, 1)
                Base.Cartesian.@nexprs $k i-> begin
                    $(broutine2x2_m_kernel_expr(:(j+i), :(j+i+step_1)))
                end
            end
            return st
        end
    end)

    push!(body_broutine2x2_ctrl.args, quote
        @inbounds if step_1 == $k
            for j in 0:step_2:size(st, 1)-step_1
                Base.Cartesian.@nexprs $k i->begin
                    if ismatch(j+i-1, ctrl_mask, flag_mask)
                        $(broutine2x2_kernel_expr(:(j+i), :(j+i+step_1)))
                    end
                end
            end
            return st
        end
    end)

    push!(body_broutine2x2_ctrl_m.args, quote
        @inbounds if step_1 == $k
            for j in 0:step_2:size(st, 2)-step_1, b in 1:size(st, 1)
                Base.Cartesian.@nexprs $k i->begin
                    if ismatch(j+i-1, ctrl_mask, flag_mask)
                        $(broutine2x2_m_kernel_expr(:(j+i), :(j+i+step_1)))
                    end
                end
            end
            return st
        end
    end)
end

# case 2: we expand the inner loop by 16, if step_1 is larger than 16
push!(body_broutine2x2.args, quote
    for j in 0:step_2:size(st, 1)-step_1
        for i in j:16:j+step_1-1
            Base.Cartesian.@nexprs 16 k->begin
                $(broutine2x2_kernel_expr(:(i+k), :(i+step_1+k)))
            end
        end
    end
    return st
end)

push!(body_broutine2x2_m.args, quote
    for j in 0:step_2:size(st, 2)-step_1, b in 1:size(st, 1)
        for i in j:16:j+step_1-1
            Base.Cartesian.@nexprs 16 k->begin
                $(broutine2x2_m_kernel_expr(:(i+k), :(i+step_1+k)))
            end
        end
    end
    return st
end)

push!(body_broutine2x2_ctrl.args, quote
    for j in 0:step_2:size(st, 1)-step_1
        for i in j:16:j+step_1-1
            Base.Cartesian.@nexprs 16 k->begin
                if ismatch(i+k-1, ctrl_mask, flag_mask)
                    $(broutine2x2_kernel_expr(:(i+k), :(i+step_1+k)))
                end
            end
        end
    end
    return st
end)

push!(body_broutine2x2_ctrl_m.args, quote
    for j in 0:step_2:size(st, 2)-step_1, b in 1:size(st, 1)
        for i in j:16:j+step_1-1
            Base.Cartesian.@nexprs 16 k->begin
                if ismatch(i+k-1, ctrl_mask, flag_mask)
                    $(broutine2x2_m_kernel_expr(:(i+k), :(i+step_1+k)))
                end
            end
        end
    end
    return st
end)

# create broutine2x2 function expr
def_broutine2x2 = Dict{Symbol, Any}(
    :name => :_broutine2x2!,
    :args => [:(st::AbstractVector), :((U11,U12,U21,U22)), :(locs::Locations)],
    :body => body_broutine2x2
)

def_broutine2x2_m = Dict{Symbol, Any}(
    :name => :_broutine2x2!,
    :args => [:(st::AbstractMatrix), :((U11,U12,U21,U22)), :(locs::Locations)],
    :body => body_broutine2x2_m
)

def_broutine2x2_ctrl = Dict{Symbol, Any}(
    :name => :_broutine2x2!,
    :args => [:(st::AbstractVector), :((U11,U12,U21,U22)), :(locs::Locations), :(ctrl::CtrlLocations)],
    :body => body_broutine2x2_ctrl,
)

def_broutine2x2_ctrl_m = Dict{Symbol, Any}(
    :name => :_broutine2x2!,
    :args => [:(st::AbstractMatrix), :((U11,U12,U21,U22)), :(locs::Locations), :(ctrl::CtrlLocations)],
    :body => body_broutine2x2_ctrl_m,
)

@eval $(combinedef(def_broutine2x2))
@eval $(combinedef(def_broutine2x2_m))
@eval $(combinedef(def_broutine2x2_ctrl))
@eval $(combinedef(def_broutine2x2_ctrl_m))

function broutine2x2!(st::AbstractVecOrMat, U::AbstractMatrix, locs::Locations, ctrl::CtrlLocations)
    U11 = U[1, 1]; U12 = U[1, 2];
    U21 = U[2, 1]; U22 = U[2, 2];
    for loc in locs
        _broutine2x2!(st, (U11,U12,U21,U22), locs, ctrl)
    end
    return st
end

# level 2 - special matrix routines

# level 3 - intrinsic gate routines
@broutine X = [0 1;1 0]
@broutine Y = [0 -im;im 0]
@broutine Z = [1 0;0 -1]
@broutine H = [1/sqrt(2) 1/sqrt(2); 1/sqrt(2) -1/sqrt(2)]
@broutine T = [1 0;0 exp(im * π / 4)]
@broutine S = [1 0;0 im]
@broutine Sdag = [1 0;0 -im]
@broutine Tdag = [1 0;0 exp(-im * π / 4)]
@broutine Rx(θ::Real) = [cos(θ / 2) -im*sin(θ / 2); -im*sin(θ / 2) cos(θ / 2)]
@broutine Ry(θ::Real) = [cos(θ / 2) -sin(θ / 2); sin(θ / 2) cos(θ / 2)]
@broutine Rz(θ::Real) = [exp(-im * θ / 2) 0; 0 exp(im * θ / 2)]

@broutine PSWAP(θ) = [-im * sin(θ/2) + cos(θ/2) 0 0 0;
                0 cos(θ/2) -im * sin(θ/2) 0;
                0 -im * sin(θ/2) cos(θ/2) 0;
                0 0 0 -im * sin(θ/2) + cos(θ/2)]

# @broutine SWAP = [1 0 0 0;
#                 0 0 1 0;
#                 0 1 0 0;
#                 0 0 0 1]
#
# NOTE: the compiler seems to be happier without us optimizing the kernel function
#      don't remove T1 and T4, so let's just define this manually :(

const BSWAP = [1 0 0 0;
                0 0 1 0;
                0 1 0 0;
                0 0 0 1]

@inline function unsafe_swap_kernel(st, b3, step_2, step_4)
    @fastmath @inbounds begin
        i1 = b3+1
        i2 = b3+step_2+1
        i3 = b3+step_4+1
        i4 = b3+step_2+step_4+1

        T1 = st[i1]
        T2 = st[i3]
        T3 = st[i2]
        T4 = st[i4]

        st[i1] = T1
        st[i2] = T2
        st[i3] = T3
        st[i4] = T4
    end
end

@inline function unsafe_swap_kernel(st, b, b3, step_2, step_4)
    @fastmath @inbounds begin
        i1 = b3+1
        i2 = b3+step_2+1
        i3 = b3+step_4+1
        i4 = b3+step_2+step_4+1

        T1 = st[b, i1]
        T2 = st[b, i3]
        T3 = st[b, i2]
        T4 = st[b, i4]

        st[b, i1] = T1
        st[b, i2] = T2
        st[b, i3] = T3
        st[b, i4] = T4
    end
end

function _swap_step(locs)
    loc_1 = first(plain(locs))
    loc_2 = last(plain(locs))

    if loc_1 > loc_2
        loc_1 = last(plain(locs))
        loc_2 = first(plain(locs))
    end

    step_1 = 1 << loc_1
    step_2 = 1 << (loc_1 - 1)
    step_3 = 1 << loc_2
    step_4 = 1 << (loc_2 - 1)

    return step_1, step_2, step_3, step_4
end

function broutine!(st::AbstractVector, ::Val{:SWAP}, locs::Locations)
    @assert length(locs) == 2 "number of locations mismatch"
    step = _swap_step(locs)

    @fastmath @inbounds for b1 in 0:step[3]:size(st, 1)-step[3]
        for b2 in b1:step[1]:b1+step[4]-step[1]
            for b3 in b2:b2+step[2]-1
                unsafe_swap_kernel(st, b3, step[2], step[4])
            end
        end
    end
    return st
end

function broutine!(st::AbstractMatrix, ::Val{:SWAP}, locs::Locations)
    @assert length(locs) == 2 "number of locations mismatch"
    step = _swap_step(locs)

    @fastmath @inbounds for b1 in 0:step[3]:size(st, 2)-step[3]
        for b2 in b1:step[1]:b1+step[4]-step[1]
            for b3 in b2:b2+step[2]-1, b in axes(st, 1)
                unsafe_swap_kernel(st, b, b3, step[2], step[4])
            end
        end
    end
    return st
end

function broutine!(st::AbstractVector, ::Val{:SWAP}, locs::Locations, ctrl::CtrlLocations)
    @assert length(locs) == 2 "number of locations mismatch"
    step = _swap_step(locs)

    raw_ctrl = plain(ctrl)
    ctrl_mask = bmask(raw_ctrl)
    flag_mask = reduce(+, one(Int) << (raw_ctrl[i] - 1) for i in 1:length(ctrl) if ctrl.flags[i])

    @fastmath @inbounds for b1 in 0:step[3]:size(st, 1)-step[3]
        for b2 in b1:step[1]:b1+step[4]-step[1]
            for b3 in b2:b2+step[2]-1
                if ismatch(b3, ctrl_mask, flag_mask)
                    unsafe_swap_kernel(st, b3, step[2], step[4])
                end
            end
        end
    end
    return st
end

function broutine!(st::AbstractMatrix, ::Val{:SWAP}, locs::Locations, ctrl::CtrlLocations)
    @assert length(locs) == 2 "number of locations mismatch"
    step = _swap_step(locs)

    raw_ctrl = plain(ctrl)
    ctrl_mask = bmask(raw_ctrl)
    flag_mask = reduce(+, one(Int) << (raw_ctrl[i] - 1) for i in 1:length(ctrl) if ctrl.flags[i])

    @fastmath @inbounds for b1 in 0:step[3]:size(st, 2)-step[3]
        for b2 in b1:step[1]:b1+step[4]-step[1]
            for b3 in b2:b2+step[2]-1, b in axes(st, 1)
                if ismatch(b3, ctrl_mask, flag_mask)
                    unsafe_swap_kernel(st, b, b3, step[2], step[4])
                end
            end
        end
    end
    return st
end


"""
    AbstractContext

An imputation context records summary information about missing data for an imputation algorithm.
All `AbstractContext`s are callable with a function, which allows us to write code like:

```julia
context() do c
    # My imputation code using a clean context
end
```

This do-block will pass a fresh context to your code and apply the `on_complete` function on
the resulting data and context state. By default, `on_complete` will throw an
[ImputeError](@ref) if we have too many missing values.
"""
abstract type AbstractContext end

# We implement a version of copy for all contexts which reconstructs the context from the
# raw fields.
Base.copy(ctx::T) where {T <: AbstractContext} = T(fieldvalues(ctx)...)

"""
    ismissing(ctx::AbstractContext, x) -> Bool

Uses `ctx.is_missing` to determine if x is missing. If x is a `NamedTuple` or an `AbstractArray`
then `ismissing` will return true if `ctx.is_missing` returns true for any element.
The ctx.count is increased whenever whenever we return true and if `ctx.count / ctx.num`
exceeds our `ctx.limit` we throw an `ImputeError`

# Arguments
* `ctx::Context`: the contextual information about missing information.
* `x`: the value to check (may be an single values, abstract array or row)
"""
function Base.ismissing(ctx::AbstractContext, x)
    was_missing = if isa(x, NamedTuple)
        any(ctx.is_missing, Tuple(x))
    elseif isa(x, AbstractArray)
        any(ctx.is_missing, x)
    else
        ctx.is_missing(x)
    end

    missing_update!(ctx, was_missing)

    return was_missing
end

"""
    findfirst(ctx::AbstractContext, data::AbstractVector) -> Int

Returns the first non-missing index in `data`.

# Arguments
* `ctx::AbstractContext`: the context to pass into `ismissing`
* `data::AbstractVector`: the data array to search

# Returns
* `Int`: the first index in `data` that isn't missing
"""
function Base.findfirst(ctx::AbstractContext, data::AbstractVector)
    return findfirst(x -> !ismissing(ctx, x), data)
end

"""
    findlast(ctx::AbstractContext, data::AbstractVector) -> Int

Returns the last non-missing index in `data`.

# Arguments
* `ctx::AbstractContext`: the context to pass into `ismissing`
* `data::AbstractVector`: the data array to search

# Returns
* `Int`: the last index in `data` that isn't missing
"""
function Base.findlast(ctx::AbstractContext, data::AbstractVector)
    return findlast(x -> !ismissing(ctx, x), data)
end

"""
    findnext(ctx::AbstractContext, data::AbstractVector) -> Int

Returns the next non-missing index in `data`.

# Arguments
* `ctx::AbstractContext`: the context to pass into `ismissing`
* `data::AbstractVector`: the data array to search

# Returns
* `Int`: the next index in `data` that isn't missing
"""
function Base.findnext(ctx::AbstractContext, data::AbstractVector, idx::Int)
    return findnext(x -> !ismissing(ctx, x), data, idx)
end

@auto_hash_equals mutable struct Context <: AbstractContext
    num::Int
    count::Int
    limit::Float64
    is_missing::Function
    on_complete::Function
end

"""
    Context

Records base information about the missing data and assume all observations are equally
weighted.

# Keyword Arguments
* `n::Int`: number of observations
* `count::Int`: number of missing values found
* `limit::Float64`: portion of total values allowed to be imputed (should be between 0.0 and 1.0).
* `is_missing::Function`: must return a Bool indicating if the value counts as missing
* `on_complete::Function`: a function to run when imputation is complete
"""
function Context(;
    limit::Float64=0.1,
    is_missing::Function=ismissing,
    on_complete::Function=complete
)
    return Context(0, 0, limit, is_missing, on_complete)
end

function Base.empty(ctx::Context)
    _ctx = copy(ctx)
    _ctx.num = 0
    _ctx.count = 0

    return _ctx
end

function missing_update!(ctx::Context, was_missing)
    ctx.num += 1

    if was_missing
        ctx.count += 1
    end
end

function complete(ctx::Context, data)
    missing_ratio = ctx.count / ctx.num

    if missing_ratio > ctx.limit
        throw(ImputeError(
            "More than $(ctx.limit * 100)% of values were missing ($missing_ratio)."
        ))
    end

    return data
end

@auto_hash_equals mutable struct WeightedContext <: AbstractContext
    num::Int
    s::Float64
    limit::Float64
    is_missing::Function
    on_complete::Function
    wv::AbstractWeights
end

"""
    WeightedContext(wv; limit=1.0, is_missing=ismissing, on_complete=complete)

Records information about the missing data relative to a set of weights.
This context type can be useful if some missing observation are more important than others
(e.g., more recent observations in time series datasets)

# Arguments
* `wv::AbstractWeights`: a set of statistical weights to use when evaluating the importance
  of each observation. Will be accumulated during imputation.

# Keyword Arguments
* `num::Int`: number of observations
* `s::Float64`: sum of the weights of missing values
* `limit::Float64`: portion of total values allowed to be imputed (should be between 0.0 and 1.0).
* `is_missing::Function`: returns a Bool if the value counts as missing
* `on_complete::Function`: a function to run when imputation is complete
"""
function WeightedContext(
    wv::AbstractWeights;
    limit::Real=1.0,
    is_missing::Function=ismissing,
    on_complete::Function=complete
)
    return WeightedContext(0, 0.0, limit, is_missing, on_complete, wv)
end

function Base.empty(ctx::WeightedContext)
    _ctx = copy(ctx)
    _ctx.num = 0
    _ctx.s = 0.0

    return _ctx
end

function missing_update!(ctx::WeightedContext, was_missing)
    ctx.num += 1

    if was_missing
        ctx.s += ctx.wv[ctx.num]
    end
end

function complete(ctx::WeightedContext, data)
    missing_ratio = ctx.s / sum(ctx.wv)

    if missing_ratio > ctx.limit
        throw(ImputeError(
            "More than $(ctx.limit * 100)% of weighted values were missing ($missing_ratio)."
        ))
    end

    return data
end

#=
Define our callable methods for each context. Once we drop 1.0 we should be able to just
define this on the `AbstractContext`.
=#
for T in (Context, WeightedContext)
    @eval begin
        function (ctx::$T)(f::Function)
            _ctx = empty(ctx)
            return ctx.on_complete(_ctx, f(_ctx))
        end
    end
end


module Keldysh

# utility functions
export fermi

# Branch functions
export BranchEnum, forward_branch, backward_branch, imaginary_branch, BranchPoint, Branch

# Contour functions
export ContourEnum, full_contour, keldysh_contour, imaginary_contour, Contour, twist, heaviside, θ, get_branch, nbranches

# TimeGrid functions
export TimeGrid, TimeGridPoint, branch_bounds, integrate, step, realtimes, imagtimes

# TimeGF functions
export TimeGF, TimeGFTranspose, jump

# DOS functions
export dos_integrator, flat_dos, gaussian_dos, bethe_dos, chain_dos, square_dos

# generate_gf functions
export gf_1level, dos2gf

import Base.in, Base.size, Base.getindex, Base.setindex!, Base.IndexStyle, Base.step, Base.length, Base.similar, Base.adjoint, Base.read, Base.write

include("util.jl")
include("branch.jl")
include("contour.jl")
include("time_grid.jl")
include("time_gf.jl")
include("dos.jl")
include("generate_gf.jl")
include("hdf5.jl")
include("nessi.jl")

end # module


import Dates

"""
    NullCalendar()

A holiday calendar for which there is never a holiday. *sigh*
"""
struct NullCalendar <: Calendar end

is_holiday(::NullCalendar, ::Dates.Date) = false
holidays(::NullCalendar, from::Dates.Date, to::Dates.Date) = Base.repeat([false], (to - from).value + 1)

"""
    WeekendCalendar()

A calendar which will mark every Saturday and Sunday as a holiday
"""
struct WeekendCalendar <: Calendar end

function is_holiday(::WeekendCalendar, date::Dates.Date)
    signbit(5 - Dates.dayofweek(date))
end

function holidays(::WeekendCalendar, from::Dates.Date, to::Dates.Date)
    dow = Dates.dayofweek(from)
    days = (to - from).value + 1

    hols = Base.repeat([false], days)
    satstart = dow == 7 ? 7 : 7 - dow
    sunstart = 8 - dow
    for i in Iterators.flatten([satstart:7:days, sunstart:7:days])
        @inbounds hols[i] = true
    end
    hols
end

mutable struct CalendarCache
    bdays::Vector{Bool}
    bdayscounter::Vector{UInt32}
    dtmin::Dates.Date
    dtmax::Dates.Date
    initialised::Bool

    CalendarCache() = new([], [], Dates.Date(1900,1,1), Dates.Date(1900,1,1), false)
end

function needsupdate(cache::CalendarCache, d0::Dates.Date, d1::Dates.Date)
    !cache.initialised || d0 < cache.dtmin || d1 > cache.dtmax
end

"""
    CachedCalendar(cal::Calendar)

Creating a wrapping calendar that will cache the holidays lazily as retrieved
for a given year, rather than loading them in one go.
"""
struct CachedCalendar <: Calendar
    calendar::Calendar
    cache::CalendarCache
    period::UInt8
    CachedCalendar(cal::Calendar) = new(cal, CalendarCache(), 10)
end

Base.show(io::IO, mgr::CachedCalendar) = (print(io, "CachedCalendar("); show(io, mgr.calendar); print(io, ")"))

function updatecache!(cal::CachedCalendar, d0::Dates.Date, d1::Dates.Date)
    needsupdate(cal.cache, d0, d1) || return
    if cal.cache.initialised
        if d0 < cal.cache.dtmin
            days = (cal.cache.dtmin-d0).value
            bdays = holidays(cal.calendar, d0, cal.cache.dtmin-Dates.Day(1))
            cal.cache.bdays = vcat(bdays, cal.cache.bdays)
            counters = Base.sum(bdays) .+ cal.cache.bdayscounter

            lcounters = Vector{UInt32}(undef, days)
            lcounters[1] = UInt32(bdays[1])
            for i in 2:days
                @inbounds lcounters[i] = lcounters[i-1] + bdays[i]
            end
            cal.cache.bdayscounter = vcat(lcounters, counters)
            cal.cache.dtmin = d0
        end

        if d1 > cal.cache.dtmax
            days = (d1-cal.cache.dtmax).value
            bdays = holidays(cal.calendar, cal.cache.dtmax+Dates.Day(1), d1)
            cal.cache.bdays = vcat(cal.cache.bdays, bdays)
            rcounters = Vector{UInt32}(undef, days)
            rcounters[1] = cal.cache.bdayscounter[end] + UInt32(bdays[1])
            for i in 2:days
                @inbounds rcounters[i] = rcounters[i-1] + bdays[i]
            end
            cal.cache.bdayscounter = vcat(cal.cache.bdayscounter, rcounters)
            cal.cache.dtmax = d1
        end
    else
        days = (d1-d0).value + 1
        cal.cache.bdays = holidays(cal.calendar, d0, d1)
        cal.cache.bdayscounter = Vector{UInt32}(undef, days)
        cal.cache.bdayscounter[1] = UInt32(cal.cache.bdays[1])
        for i in 2:days
            @inbounds cal.cache.bdayscounter[i] = cal.cache.bdayscounter[i-1] + cal.cache.bdays[i]
        end
        cal.cache.dtmin = d0
        cal.cache.dtmax = d1
        cal.cache.initialised = true
    end
end

function is_holiday(cal::CachedCalendar, date::Dates.Date)
    if needsupdate(cal.cache, date, date)
        d0 = Dates.Date(cal.period*div(Dates.year(date), cal.period), 1, 1)
        d1 = Dates.Date(cal.period*(1+div(Dates.year(date), cal.period)), 12, 31)
        updatecache!(cal, d0, d1)
    end
    t0 = (date-cal.cache.dtmin).value + 1
    cal.cache.bdays[t0]
end

function holidaycount(cal::CachedCalendar, from::Dates.Date, to::Dates.Date)
    if needsupdate(cal.cache, from, to)
        d0 = Dates.Date(cal.period*div(Dates.year(from), cal.period), 1, 1)
        d1 = Dates.Date(cal.period*(1+div(Dates.year(to), cal.period)), 12, 31)
        updatecache!(cal, d0, d1)
    end

    t0 = (from-cal.cache.dtmin).value + 1
    t1 = (to-from).value

    Int(cal.cache.bdayscounter[t1+t0] - cal.cache.bdayscounter[t0] + cal.cache.bdays[t0])
end

function holidays(cal::CachedCalendar, from::Dates.Date, to::Dates.Date)
    if needsupdate(cal.cache, from, to)
        d0 = Dates.Date(cal.period*div(Dates.year(from), cal.period), 1, 1)
        d1 = Dates.Date(cal.period*(1+div(Dates.year(to), cal.period)), 12, 31)
        updatecache!(cal, d0, d1)
    end
    t0 = (from-cal.cache.dtmin).value + 1
    t1 = (to-from).value
    cal.cache.bdays[t0:t0+t1]
end

"""
    JointCalendar(calendars::Vector{Calendar}) <: Calendar

A grouping of calendars, for which it is a holiday if it's marked as a holiday
for any of the underlying calendars.

By default addition of calendars will generate a joint calendar for you.
"""
struct JointCalendar <: Calendar
    calendars::Vector{Calendar}
end

function is_holiday(cal::JointCalendar, date::Dates.Date)
    foldl(
        (acc, val) -> acc || is_holiday(val, date),
        cal.calendars;
        init = false,
    )
end
function holidays(cal::JointCalendar, from::Dates.Date, to::Dates.Date)
    hols = Base.repeat([false], (to-from).value+1)
    for subcal in cal.calendars
        hols = hols .| holidays(subcal, from, to)
    end
    hols
end

Base.:+(cal1::Calendar, cal2::Calendar) = JointCalendar([cal1, cal2])
Base.:+(cal1::Calendar, cal2::JointCalendar) =
    JointCalendar(vcat(cal1, cal2.calendars))
Base.:+(cal1::JointCalendar, cal2::Calendar) =
    JointCalendar(vcat(cal1.calendars, cal2))
Base.:+(cal1::JointCalendar, cal2::JointCalendar) =
    JointCalendar(vcat(cal1.calendars, cal2.calendars))

"""
    SimpleCalendarManager()
    SimpleCalendarManager(calendars::Dict{String, Calendar})

A basic calendar manager which just holds a reference to each underlying calendar, by name,
and will generate a joint calendar if multiple names are requested.

To set a calendar on this manager then use `setcalendar!`

```julia-repl
julia> mgr = SimpleCalendarManager()
julia> setcalendar!(mgr, "WEEKEND", WeekendCalendar())
WeekendCalendar()
```

If you want to set a cached wrapping of a calendar then use `setcachedcalendar!`

```julia-repl
julia> mgr = SimpleCalendarManager()
julia> setcachedcalendar!(mgr, "WEEKEND", WeekendCalendar())
CachedCalendar(WeekendCalendar())
```

### Examples
```julia-repl
julia> mgr = SimpleCalendarManager()
julia> setcalendar!(mgr, "WEEKEND", WeekendCalendar())
julia> is_holiday(calendar(mgr, ["WEEKEND"]), Date(2019,9,28))
true
```
"""
struct SimpleCalendarManager <: CalendarManager
    calendars::Dict{String,Calendar}
    SimpleCalendarManager(calendars) = new(calendars)
    SimpleCalendarManager() = new(Dict())
end
Base.show(io::IO, mgr::SimpleCalendarManager) = (print(io, "SimpleCalendarManager($(length(mgr.calendars)) calendars)"))

setcalendar!(mgr::SimpleCalendarManager, name::String, cal::Calendar) = (mgr.calendars[name] = cal;)
setcachedcalendar!(mgr::SimpleCalendarManager, name::String, cal::Calendar) = (mgr.calendars[name] = CachedCalendar(cal);)
setcachedcalendar!(mgr::SimpleCalendarManager, name::String, cal::CachedCalendar) = setcalendar!(mgr, name, cal)

function calendar(cal_mgr::SimpleCalendarManager, names::Vector)::Calendar
    if length(names) == 0
        NullCalendar()
    elseif length(names) == 1
        cal_mgr.calendars[names[1]]
    else
        foldl(
            (acc, val) -> acc + cal_mgr.calendars[val],
            names[2:end],
            cal_mgr.calendars[names[1]],
        )
    end
end

"""
    get_calendar_names(rdate::RDate)::Vector{String}

A helper method to get all of the calendar names that could potentially be requested by this rdate. This
mechanism can be used to mark the minimal set of calendars on which adjustments depend.
"""
get_calendar_names(::RDate) = Vector{String}()
get_calendar_names(rdate::Union{BizDays,CalendarAdj}) = rdate.calendar_names
get_calendar_names(rdate::Compound) = Base.foldl((val, acc) -> vcat(
        acc,
        get_calendar_names(val),
        rdate.parts,
        init = Vector{String}(),
    ))
get_calendar_names(rdate::Repeat) = get_calendar_names(rdate.part)


module Shapes

import Adapt

import Base: getindex, length, size
using Base: @pure, @propagate_inbounds, @_inline_meta

import Random
using Random: AbstractRNG

using Requires: @init, @require

import UnsafeArrays

import StaticArrays: Size, Length, similar_type, get
using StaticArrays: tuple_prod,
                    tuple_length,
                    tuple_minimum,
                    promote_tuple_eltype,
                    size_to_tuple,
                    SArray,
                    MArray,
                    StaticArray,
                    SOneTo

export AbstractShape,
       checkaxes,
       checksize,
       Shape,
       Size,
       Length,
       ScalarShape,
       VectorShape,
       MatrixShape,
       MultiShape,
       concrete_eltype,

       ShapedView,
       allocate

include("util.jl")
include("traits.jl")
include("ShapedView.jl")

getdims(shape::AbstractShape, dims) = getdims(Size(shape), dims)
getdims(s::Size, dims) = ((Tuple(s)..., dims...))
getdims(::Size{()}, dims) = dims

Base.Array{T}(::UndefInitializer, shape::AbstractShape, dims::Dims) where {T} =
    Array{T}(undef, getdims(shape, dims))
Base.Array(::UndefInitializer, shape::AbstractShape, dims::Dims) =
    Array{concrete_eltype(shape)}(undef, shape, dims)

Base.Array{T}(::UndefInitializer, shape::AbstractShape, dims::Integer...) where {T} =
    Array{T}(undef, shape, dims)
Base.Array(::UndefInitializer, shape::AbstractShape, dims::Integer...) =
    Array(undef, shape, dims)

Base.Matrix{T}(::UndefInitializer, shape::AbstractShape, dim::Integer) where {T} =
    Matrix{T}(undef, getdims(shape, (dim,)))
Base.Matrix(::UndefInitializer, shape::AbstractShape, dim::Integer) =
    Matrix{concrete_eltype(shape)}(undef, shape, dim)

Base.Vector{T}(::UndefInitializer, shape::AbstractShape) where {T} =
    Vector{T}(undef, length(shape))
Base.Vector(::UndefInitializer, shape::AbstractShape) =
    Vector{concrete_eltype(shape)}(undef, shape)

Base.zeros(T, shape::AbstractShape, dims::Dims) = zeros(T, getdims(shape, dims))
Base.zeros(shape::AbstractShape, dims::Dims) = zeros(concrete_eltype(shape), shape, dims)

Base.zeros(T, shape::AbstractShape, dims::Integer...) = zeros(T, shape, dims)
Base.zeros(shape::AbstractShape, dims::Integer...) = zeros(shape, dims)

allocate(shape::AbstractShape, dims::Dims) = Array(undef, shape, dims)
allocate(shape::AbstractShape, dims::Integer...) = allocate(shape, dims)

allocate(T::DataType, dims::Dims) = Array{T}(undef, dims)
allocate(T::DataType, dims::Integer...) = allocate(T, dims)

allocate(T::DataType, shape::AbstractShape, dims::Dims) = allocate(T, size(shape)..., dims...)
allocate(T::DataType, shape::AbstractShape, dims::Integer...) = allocate(T, shape, dims)

function Random.rand(rng::AbstractRNG, shape::AbstractShape, dims::Dims)
    _rand(rng, concrete_eltype(shape), getdims(shape, dims))
end
Random.rand(rng::AbstractRNG, shape::AbstractShape, dims::Integer...) = rand(rng, shape, dims)

Random.rand(shape::AbstractShape, dims::Dims) = _rand(concrete_eltype(shape), getdims(shape, dims))
Random.rand(shape::AbstractShape, dims::Integer...) = rand(shape, dims)

_rand(rng, T, ::Dims{0}) = rand(rng, T)
_rand(rng, T, dims::Dims) = rand(rng, T, dims)
_rand(T, ::Dims{0}) = rand(T)
_rand(T, dims::Dims) = rand(T, dims)

# StaticArrays support
similar_type(::Type{SH}) where {SH<:AbstractShape} =
    SArray{tuple_to_size(Tuple(Size(SH))),concrete_eltype(SH),ndims(SH),length(SH)}
similar_type(::SH) where {SH<:AbstractShape} = similar_type(SH)

similar_type(::Type{SA}, ::Type{SH}) where {SA<:StaticArray,SH<:AbstractShape} =
    similar_type(SA, concrete_eltype(SH), Size(SH))
similar_type(::SA, ::Type{SH}) where {SA<:StaticArray,SH<:AbstractShape} =
    similar_type(SA, concrete_eltype(SH), Size(SH))
similar_type(::Type{SA}, ::SH) where {SA<:StaticArray,SH<:AbstractShape} =
    similar_type(SA, concrete_eltype(SH), Size(SH))
similar_type(::SA, ::SH) where {SA<:StaticArray,SH<:AbstractShape} =
    similar_type(SA, concrete_eltype(SH), Size(SH))

@init @require ElasticArrays="fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4" begin
    import .ElasticArrays
    ElasticArrays.ElasticArray(::UndefInitializer, shape::AbstractShape, dims::Dims) = ElasticArrays.ElasticArray{eltype(shape)}(undef, getdims(shape, dims)...)
    ElasticArrays.ElasticArray(::UndefInitializer, shape::AbstractShape, dims::Integer...) = ElasticArrays.ElasticArray(undef, shape, dims)
end

end # module


#!/usr/bin/env julia

using Luxor

function get_png_files(folder)
    cd(folder)
    imagelist = filter(f -> !startswith(f, ".") && endswith(f, "png"), readdir(folder))
    return map(realpath, imagelist)
end

function addimagetile(imgfile, xcenter, ycenter, tilewidth, tileheight; cropping=true)
    gsave()
    box(xcenter, ycenter, tilewidth, tileheight, :clip)
    img = readpng(imgfile)
    w = img.width
    h = img.height
    if cropping == true
        # expand image to fill tile
        if w < h
            scalefactor = max(tilewidth, tileheight)/w
        else
            scalefactor = max(tilewidth, tileheight)/h
        end
    else
        # shrink image to fit inside tile
        if w > h
            scalefactor = min(tilewidth, tileheight)/w
        else
            scalefactor = min(tilewidth, tileheight)/h
        end
    end
    translate(xcenter, ycenter)
    scale(scalefactor, scalefactor)
    placeimage(img, -w/2, -h/2)
    clipreset()
    grestore()
end

imagelist = get_png_files(dirname(@__FILE__) * "/../docs/figures")

shuffle!(imagelist)

width, height = 1600, 1000
fname = "/tmp/tiled-images.png"
Drawing(width, height, fname)
origin()
background("grey50")
setopacity(0.5)
pagetiles = Tiler(width, height, 8, 8, margin=50)
for (pos, n) in pagetiles
  if n > length(imagelist) # run out of images
    break
  end
  addimagetile(imagelist[n], pos.x, pos.y, pagetiles.tilewidth, pagetiles.tileheight, cropping=true)
end
finish()
println("finished test: output in $(fname)")


"""
    abstract type Abstract_BC_QuadSolver_Result end

Quadratic solver result abstraction

"""
abstract type Abstract_BC_QuadSolver_Result end

"""
    converged(::Abstract_BC_QuadSolver_Result)

Return `true` if the solver converged
"""
converged(::Abstract_BC_QuadSolver_Result) = @assert(false,"To implement")

"""
    iteration_count(::Abstract_BC_QuadSolver_Result)

Return the number of consumed iteration
"""
iteration_count(::Abstract_BC_QuadSolver_Result) = @assert(false,"To implement")

"""
    objective_value(::Abstract_BC_QuadSolver_Result)

Returns objective value at the point [`solution`](@ref).
"""
objective_value(r::Abstract_BC_QuadSolver_Result) = @assert(false,"To implement")

"""
    multiplier_τ(::Abstract_BC_QuadSolver_Result)

Returns the multipliers stored in a compact form (see τ definition, TODO)
"""
multiplier_τ(r::Abstract_BC_QuadSolver_Result) = @assert(false,"To implement")

"""
    solution(::Abstract_BC_QuadSolver_Result)

Returns the founded solution
"""
solution(r::Abstract_BC_QuadSolver_Result) = @assert(false,"To implement")


using .Gadfly
using Colors
using StatsBase

function Gadfly.plot(
    zep::Zeppelin,
    rows;
    autoklms = true,
    klms = [],
    xmin = 0.0,
    xmax = missing,
    norm = :None,
    yscale = 1.05,
    ytransform = identity,
    style = NeXLSpectrum.NeXLSpectrumStyle,
    palette = NeXLCore.NeXLPalette,
)
    plot(
        Spectrum[filter(s->!ismissing(s), map(r -> spectrum(zep, r, false),rows))...],
        autoklms = autoklms,
        klms = klms,
        xmin = xmin,
        xmax = xmax,
        norm = norm,
        yscale = yscale,
        ytransform = ytransform,
        style = style,
        palette = palette,
    )
end

function Gadfly.plot(zeps::AbstractArray{Zeppelin}; offset=(0.0,0.0), point_size=1.5pt, coords=missing)
    names = NeXLSpectrum.name.(zeps)
    colors = map(i->NeXLPalette[1 + (i-1) % length(NeXLPalette)], eachindex(zeps))
    clsCx=collect(merge( [ StatsBase.countmap(repr.(z[:,:CLASS])) for z in zeps]...))
    sort!(clsCx, lt=(a,b)->a[2]<b[2], rev=true)
    shlen = length(Theme().point_shapes)
    if length(clsCx) >= shlen
        clsCx = clsCx[1:(shlen-1)]
        push!(clsCx,"Other"=>0)
    end
    shIdx=Dict(clsCx[i][1]=>i for i in eachindex(clsCx))
    xmin = minimum(z->minimum(z.data[:,:XABS]), zeps)
    xmax = maximum(z->maximum(z.data[:,:XABS]), zeps)
    ymin = minimum(z->minimum(z.data[:,:YABS]), zeps)
    ymax = maximum(z->maximum(z.data[:,:YABS]), zeps)
    dmax = ceil(max(xmax-xmin, ymax-ymin)/2)
    xc, yc = round((xmax+xmin)/2), round((ymax+ymin)/2)
    coords = Coord.cartesian(xmin=xc-dmax,xmax=xc+dmax,ymin=yc-dmax,ymax=yc+dmax)
    plot(
        Theme(point_size=1.5pt),
        map(enumerate(zeps)) do (i,z)
            layer(
                x=map(r->r.XABS+(i-1)*offset[1], eachrow(z.data)),
                y=map(r->r.YABS+(i-1)*offset[2], eachrow(z.data)),
                shape=map(r->get(shIdx, repr(r.CLASS), length(clsCx)), eachrow(z.data)),
                Geom.point,
                Theme(default_color=colors[i], point_size=point_size)
            )
        end...,
        Guide.xlabel("X (mm)"), Guide.ylabel("Y (mm)"),
        Guide.manual_color_key("Particle Data Set", names, colors),
        Guide.shapekey(title="Particle Class", labels=map(cc->cc[1],clsCx)),
        coords
    )
end

function Gadfly.plot(::Type{Histogram}, zeps::AbstractArray, sym::Symbol; bincount=40, limits=missing)
    names = NeXLSpectrum.name.(zeps)
    df = vcat(
        map(enumerate(zeps)) do (i, z)
            DataFrame(Dataset=fill(names[i],nrow(z.data)), v=z.data[:,sym])
        end...
    )
    if ismissing(limits)
        limits = ( min=floor(minimum(z->minimum(z.data[:,sym]), zeps)), max=ceil(maximum(z->maximum(z.data[:,sym]), zeps)) )
    end
    h = Geom.histogram(bincount=bincount, position=:identity, limits=limits)
    plot(df, x=:v, color=:Dataset, alpha=[0.5], h,
        Guide.xlabel(repr(sym)[2:end]), Guide.ylabel("Count")
    )
end


using RecursiveArrayTools: ArrayPartition
using ComponentArrays: ComponentVector
using ..RingPolymers: RingPolymers

get_positions(u::ComponentVector) = u.r
get_velocities(u::ComponentVector) = u.v
get_flat_positions(u::ComponentVector) = @view get_positions(u)[:]
get_flat_velocities(u::ComponentVector) = @view get_velocities(u)[:]

# These must be kept since the symplectic integrators (VelocityVerlet) in DiffEq use ArrayPartitions.
get_positions(u::ArrayPartition) = u.x[2]
get_velocities(u::ArrayPartition) = u.x[1]

function RingPolymers.transform_from_normal_modes!(p::RingPolymers.RingPolymerParameters, u)
    RingPolymers.transform_from_normal_modes!(p, DynamicsUtils.get_velocities(u))
    RingPolymers.transform_from_normal_modes!(p, DynamicsUtils.get_positions(u))
    return nothing
end

function RingPolymers.transform_to_normal_modes!(p::RingPolymers.RingPolymerParameters, u)
    RingPolymers.transform_to_normal_modes!(p, DynamicsUtils.get_velocities(u))
    RingPolymers.transform_to_normal_modes!(p, DynamicsUtils.get_positions(u))
    return nothing
end

function get_quantum_subsystem end

get_mapping_positions(u::ComponentVector) = u.qmap
get_mapping_momenta(u::ComponentVector) = u.pmap
get_mapping_positions(u::ComponentVector, i) = @view get_mapping_positions(u)[:,i]
get_mapping_momenta(u::ComponentVector, i) = @view get_mapping_momenta(u)[:,i]


using MultivariateMixtures
using MultivariateStats
using Distributions
using LinearAlgebra
using Test
using Random

@testset "MultivariateMixtures" begin
    include("utils.jl")
    include("gmm.jl")
    include("mfa.jl")
    # include("mppca.jl")
end


module Perf
using PhaseSpaceIO
using PhaseSpaceIO.Testing
using Test
using PhaseSpaceIO.Getters

@testset "Getters don't allocate" begin
    PS = [
        EGSParticle{Nothing},
        EGSParticle{Float32},
        IAEAParticle{0,0},
        IAEAParticle{2,1},
    ]
    @testset "P = $P" for P in PS
        ps = [arbitrary(P) for _ in 1:10^3]
        getters = [E,x,y,z, u,v,w]
        @testset "f = $f" for f in getters
             T = typeof(f(first(ps)))
             out = similar(ps, T)
             map!(f, out, ps)
             bytes = @allocated map!(f, out, ps)
             @test bytes == 0
        end
    end
end

@noinline function sum_E(ps)
    sum(p->p.E, ps)
end

function compute_sum_E_mmap(path)
    ps = PhspVector(path)
    sum_E(ps)
end
function compute_sum_E_iter(path)
    phsp_iterator(sum_E, path)
end

ncase = 10^6
for (P, ext) in [
        (EGSParticle{Float32}, ".egsphsp"),
        (IAEAParticle{0,0}, ".IAEAphsp"),
        ]
    println("*"^80)
    println("Benchmarking $P")
    ps = [arbitrary(P) for _ in 1:ncase]
    path = tempname() * ext
    path2 = tempname() * ext
    println("write $ncase particle")
    phsp_write(path2, ps[1:10])
    @time phsp_write(path, ps)

    println("compute_sum_E_mmap")
    compute_sum_E_mmap(path2)
    @time compute_sum_E_mmap(path)

    println("compute_sum_E_iter")
    compute_sum_E_iter(path2)
    @time compute_sum_E_iter(path)
end
println("*"^80)

end


abstract type PredictiveControlProblem end

struct PrimalNoCondensing <: PredictiveControlProblem end

struct PrimalFullCondensing <: PredictiveControlProblem end

struct DualFullCondensing <: PredictiveControlProblem end


#!/usr/bin/env julia6
using MyStats.condmean
using ReadGlobal: readfield, readfield!
using JLD2

function main(condname::AbstractString, inpname::AbstractVector{<:AbstractString}, outfilename::AbstractVector{<:AbstractString})
    @load condname ind bins
    field = readfield(inpname[1])
    for i in linearindices(inpname)
        readfield!(inpname[i],field)
        r,err = condmean(field,ind)
        writedlm(outfilename[i],hcat(bins,r,err))
    end
    return nothing
end

function main(args::Vector{<:AbstractString})
    condname = args[1]
    iseven(length(args)) && error()
    p = div(length(args) - 1, 2) + 1
    inpname = args[2:p]
    outfilename = args[p+1:end]
    main(condname,inpname,outfilename)
end

main(ARGS)


##############################################################################
##
## Iterate on terms
##
##############################################################################

eachterm(@nospecialize(x::AbstractTerm)) = (x,)
eachterm(@nospecialize(x::NTuple{N, AbstractTerm})) where {N} = x
TermOrTerms = Union{AbstractTerm, NTuple{N, AbstractTerm} where N}
hasintercept(@nospecialize(t::TermOrTerms)) =
    InterceptTerm{true}() ∈ terms(t) ||
    ConstantTerm(1) ∈ terms(t)
omitsintercept(@nospecialize(f::FormulaTerm)) = omitsintercept(f.rhs)
omitsintercept(@nospecialize(t::TermOrTerms)) =
    InterceptTerm{false}() ∈ terms(t) ||
    ConstantTerm(0) ∈ terms(t) ||
    ConstantTerm(-1) ∈ terms(t)
##############################################################################
##
## Parse IV
##
##############################################################################

function parse_iv(@nospecialize(f::FormulaTerm))
	for term in eachterm(f.rhs)
		if term isa FormulaTerm
            both = intersect(eachterm(term.lhs), eachterm(term.rhs))
            endos = setdiff(eachterm(term.lhs), both)
            exos = setdiff(eachterm(term.rhs), both)
            # otherwise empty collection. Later, there will be check
            isempty(endos) && throw("There are no endogeneous variables")
            length(exos) < length(endos) && throw("Model not identified. There must be at least as many instrumental variables as endogeneneous variables")
			formula_endo = FormulaTerm(ConstantTerm(0), tuple(ConstantTerm(0), endos...))
			formula_iv = FormulaTerm(ConstantTerm(0), tuple(ConstantTerm(0), exos...))
            formula_exo = FormulaTerm(f.lhs, tuple((term for term in eachterm(f.rhs) if !isa(term, FormulaTerm))..., both...))
            return formula_exo, formula_endo, formula_iv
		end
	end
	return f, nothing, nothing
end

##############################################################################
##
## Parse FixedEffect
##
##############################################################################
struct FixedEffectTerm <: AbstractTerm
    x::Symbol
end
StatsModels.termvars(t::FixedEffectTerm) = [t.x]
fe(x::Term) = FixedEffectTerm(Symbol(x))
fe(s::Symbol) = FixedEffectTerm(s)

has_fe(::FixedEffectTerm) = true
has_fe(::FunctionTerm{typeof(fe)}) = true
has_fe(t::InteractionTerm) = any(has_fe(x) for x in t.terms)
has_fe(::AbstractTerm) = false
has_fe(@nospecialize(t::FormulaTerm)) = any(has_fe(x) for x in eachterm(t.rhs))

fesymbol(t::FixedEffectTerm) = t.x
fesymbol(t::FunctionTerm{typeof(fe)}) = Symbol(t.args_parsed[1])

"""
    parse_fixedeffect(data, formula::FormulaTerm)
    parse_fixedeffect(data, ts::NTuple{N, AbstractTerm})

Construct any `FixedEffect` specified with a `FixedEffectTerm`.

# Returns
- `Vector{FixedEffect}`: a collection of all `FixedEffect`s constructed.
- `Vector{Symbol}`: names assigned to the fixed effect estimates (can be used as column names).
- `Vector{Symbol}`: names of original fixed effects.
- `FormulaTerm` or `NTuple{N, AbstractTerm}`: `formula` or `ts` without any term related to fixed effects (an intercept may be explicitly omitted if necessary).
"""
function parse_fixedeffect(data, @nospecialize(formula::FormulaTerm))
    fes = FixedEffect[]
    ids = Symbol[]
    fekeys = Symbol[]
     for term in eachterm(formula.rhs)
        result = _parse_fixedeffect(data, term)
        if result !== nothing
            push!(fes, result[1])
            push!(ids, result[2])
            append!(fekeys, result[3])
        end
    end
    if !isempty(fes)
        if any(fe.interaction isa UnitWeights for fe in fes)
            formula = FormulaTerm(formula.lhs, (InterceptTerm{false}(), (term for term in eachterm(formula.rhs) if !isa(term, Union{ConstantTerm,InterceptTerm}) && !has_fe(term))...))
        else
            formula = FormulaTerm(formula.lhs, Tuple(term for term in eachterm(formula.rhs) if !has_fe(term)))
        end
    end
    return fes, ids, unique(fekeys), formula
end

# Method for external packages
function parse_fixedeffect(data, @nospecialize(ts::NTuple{N, AbstractTerm})) where N
    fes = FixedEffect[]
    ids = Symbol[]
    fekeys = Symbol[]
    for term in eachterm(ts)
        result = _parse_fixedeffect(data, term)
        if result !== nothing
            push!(fes, result[1])
            push!(ids, result[2])
            append!(fekeys, result[3])
        end
    end
    if !isempty(fes)
        if any(fe.interaction isa UnitWeights for fe in fes)
            ts = (InterceptTerm{false}(), (term for term in eachterm(ts) if !isa(term, Union{ConstantTerm,InterceptTerm}) && !has_fe(term))...)
        else
            ts = Tuple(term for term in eachterm(ts) if !has_fe(term))
        end
    end
    return fes, ids, unique(fekeys), ts
end

# Construct FixedEffect from a generic term
function _parse_fixedeffect(data, @nospecialize(t::AbstractTerm))
    if has_fe(t)
        st = fesymbol(t)
        return FixedEffect(Tables.getcolumn(data, st)), Symbol(:fe_, st), [st]
    end
end

# Construct FixedEffect from an InteractionTerm
function _parse_fixedeffect(data, @nospecialize(t::InteractionTerm))
    fes = (x for x in t.terms if has_fe(x))
    interactions = (x for x in t.terms if !has_fe(x))
    if !isempty(fes)
        # x1&x2 from (x1&x2)*id
        fe_names = [fesymbol(x) for x in fes]
        v1 = _multiply(data, Symbol.(interactions))
        fe = FixedEffect((Tables.getcolumn(data, fe_name) for fe_name in fe_names)...; interaction = v1)
        interactions = string.(interactions)
        s = vcat(["fe_" * string(fe_name) for fe_name in fe_names], interactions)
        return fe, Symbol(reduce((x1, x2) -> x1*"&"*x2, s)), fe_names
    end
end

function _multiply(data, ss::AbstractVector)
    if isempty(ss)
        return uweights(size(data, 1))
    elseif length(ss) == 1
        # in case it has missing (for some reason *(missing) not defined) iv VERSION < 1.6
        # do NOT use ! since it would modify the vector
        return convert(AbstractVector{Float64}, replace(Tables.getcolumn(data, ss[1]), missing => 0))
    else
        return convert(AbstractVector{Float64}, replace!(.*((Tables.getcolumn(data, x) for x in ss)...), missing => 0))
    end
end


# Note: This test suite has many hard coded values in it.
# While that is normally a bad idea, we need to guarantee reproducible results
# when using the same seed.

global SEED = 0 # Do not change this
struct MyAlphabet <: Alphabet end

@testset "Random LongSequences" begin
function test_sampler(sampler, seed, elements, firstten, T)
    rng = StableRNG(seed)
    sampled1 = [rand(rng, sampler) for i in 1:1000]
    sampled2 = rand(StableRNG(seed), sampler, 1000)
    @test sampled1 == sampled2
    @test Set(sampled1) == Set(T[convert(T, i) for i in elements])
    @test eltype(sampler) == T
    @test eltype(firstten) == T
    @test rand(StableRNG(seed), sampler, 10) == firstten
end

function test_isseq(seq, alphabettype, len)
    @test seq isa LongSequence{alphabettype}
    @test length(seq) == len
end

@testset "SamplerUniform" begin
    # Cannot instantiate with empty collection
    @test_throws ArgumentError sampler = SamplerUniform{DNA}(DNA[])

    # DNA sampler with DNA array
    sampler = SamplerUniform{DNA}([DNA_A, DNA_C, DNA_G, DNA_W])
    firstten = DNA[DNA_C, DNA_W, DNA_G, DNA_C, DNA_C, DNA_C, DNA_C, DNA_A, DNA_A, DNA_A]
    test_sampler(sampler, SEED, sampler.elems, firstten, DNA)
    # Now RNA sampler from DNA array
    sampler = SamplerUniform{RNA}([DNA_A, DNA_C, DNA_G, DNA_W])
    firstten = RNA[RNA_C, RNA_W, RNA_G, RNA_C, RNA_C, RNA_C, RNA_C, RNA_A, RNA_A, RNA_A]
    test_sampler(sampler, SEED, sampler.elems, firstten, RNA)

    # Cannot make AA sampler from DNA
    @test_throws MethodError s = SamplerUniform{AminoAcid}([DNA_A, DNA_C, DNA_G, DNA_W])

    # Automatically infer eltype
    sampler1 = SamplerUniform{DNA}([DNA_A, DNA_C])
    sampler2 = SamplerUniform([DNA_A, DNA_C])
    @test typeof(sampler2) == SamplerUniform{DNA}
    @test eltype(sampler1) == eltype(sampler2)

    # Can also infer abstract eltype
    sampler3 = SamplerUniform([DNA_A, RNA_A])
    @test eltype(sampler3) == typejoin(DNA, RNA)

    sampler = SamplerUniform{RNA}([DNA_A, DNA_C, DNA_G, DNA_W])
    @test typeof(rand(sampler)) == RNA

end # SamplerUniform

@testset "SamplerWeighted" begin
    # Must have one less weight than elements
    @test_throws ArgumentError SamplerWeighted{DNA}([DNA_A], [0.5])
    @test_throws ArgumentError SamplerWeighted{DNA}([DNA_A, DNA_C], [0.5, 0.5])
    @test_throws ArgumentError SamplerWeighted{DNA}([DNA_A], [0.5, 0.5])

    # Weights cannot exceed one
    @test_throws ArgumentError SamplerWeighted{DNA}([DNA_A, DNA_C], [1.1])
    @test_throws ArgumentError SamplerWeighted{DNA}([DNA_A, DNA_C, DNA_G], [1.00001, 0.0])

    # Weights cannot be negative
    @test_throws ArgumentError SamplerWeighted{DNA}([DNA_A, DNA_C, DNA_G], [0.5, -0.001])
    @test_throws ArgumentError SamplerWeighted{DNA}([DNA_A, DNA_C, DNA_A], [1.1, -0.2])

    # Weights always sum to one after instantiation
    s = SamplerWeighted{DNA}([DNA_A, DNA_C, DNA_G, DNA_W], [0.03, 0.2, 0.7])
    @test sum(s.probs) == 1.0
    s = SamplerWeighted{DNA}([DNA_A], [])
    @test sum(s.probs) == 1.0
    s = SamplerWeighted{DNA}([DNA_A, DNA_C], [1.0])
    @test sum(s.probs) == 1.0
    s = SamplerWeighted{DNA}([DNA_A, DNA_C, DNA_G], [0.03, 0.20000001])
    @test sum(s.probs) == 1.0

    sampler = SamplerWeighted{DNA}([DNA_N, DNA_C, DNA_W, DNA_T], [0.1, 0.2, 0.3])
    firstten = DNA[DNA_C, DNA_N, DNA_T, DNA_T, DNA_T, DNA_N, DNA_W, DNA_C, DNA_W, DNA_N]
    test_sampler(sampler, 0, sampler.elems, firstten, DNA)

    sampler = SamplerWeighted{RNA}([DNA_N, DNA_C, DNA_W, DNA_T], [0.15, 0.5, 0.2])
    firstten = RNA[RNA_C, RNA_N, RNA_U, RNA_W, RNA_U, RNA_N, RNA_C, RNA_C, RNA_C, RNA_N]
    test_sampler(sampler, 0, sampler.elems, firstten, RNA)

    @test_throws MethodError s = SamplerWeighted{AminoAcid}([DNA_A, DNA_C], [0.1])

    # Casual test to see that it can use the global RNG automatically
    sampler = SamplerWeighted{RNA}([DNA_N, DNA_C, DNA_W, DNA_T], [0.15, 0.5, 0.2])
    @test typeof(rand(sampler)) == RNA

end # SamplerWeighted

@testset "randseq Sampler" begin
    # Cannot instantiate < 0-length seq
    @test_throws ArgumentError randseq(DNAAlphabet{2}(), SamplerUniform(dna"ACG"), -1)
    @test_throws ArgumentError randseq(AminoAcidAlphabet(), SamplerUniform(aa"VTW"), -1)
    @test_throws ArgumentError randseq(RNAAlphabet{4}(), SamplerWeighted(dna"ACG", ones(2)/3), -1)

    # CAN make empty sequence with mismatching alphabets
    # Or with matching alphabets
    @test randseq(DNAAlphabet{2}(), SamplerUniform(aa"AV"), 0) == LongSequence{DNAAlphabet{2}}()
    @test randseq(DNAAlphabet{2}(), SamplerUniform(rna"U"), 2) == LongSequence{DNAAlphabet{2}}(dna"TT")

    # Cannot make nonzero sequence with mismatching alphabets
    @test_throws MethodError randseq(RNAAlphabet{4}(), SamplerUniform(aa"AV", 1))
    @test_throws MethodError randseq(AminoAcidAlphabet(), SamplerUniform(dna"ACG", 10))

    # A few samplings
    sampler = SamplerUniform(aa"TVVWYAEDK")
    @test randseq(StableRNG(SEED), AminoAcidAlphabet(), sampler, 10) == aa"VEVTKTTADY"

    sampler = SamplerUniform(dna"TGAWYKN")
    @test randseq(StableRNG(SEED), DNAAlphabet{4}(), sampler, 10) == dna"GWNGGGKTTT"

    sampler = SamplerWeighted(rna"UGCMKYN", [0.1, 0.05, 0.2, 0.15, 0.15, 0.2])
    @test randseq(StableRNG(SEED), DNAAlphabet{4}(), sampler, 10) == dna"CTNYNTKCMT"

    # Casual tests to see that it can use the global RNG automatically
    sampler = SamplerUniform(aa"TVVWYAEDK")
    seq = randseq(AminoAcidAlphabet(), sampler, 25)
    test_isseq(seq, AminoAcidAlphabet, 25)

    sampler = SamplerWeighted(rna"UGCMKYN", [0.1, 0.05, 0.2, 0.15, 0.15, 0.2])
    seq = randseq(RNAAlphabet{4}(), sampler, 25)
    test_isseq(seq, RNAAlphabet{4}, 25)

    sampler = SamplerUniform(dna"AGC")
    seq = randseq(DNAAlphabet{2}(), sampler, 25)
    test_isseq(seq, DNAAlphabet{2}, 25)
end # randseq Sampler

@testset "randseq" begin
    sampler = SamplerUniform(aa"ACDEFGHIKLMNPQRSTVWY")
    automatic = randseq(StableRNG(SEED), AminoAcidAlphabet(), 1000)
    manual = randseq(StableRNG(SEED), AminoAcidAlphabet(), sampler, 1000)
    @test automatic == manual
    @test Set(automatic) == Set(aa"ACDEFGHIKLMNPQRSTVWY")

    @test randseq(StableRNG(SEED), DNAAlphabet{4}(), 20) == dna"CTCTTCGTATGCCGTACCGT"
    @test randseq(StableRNG(SEED), DNAAlphabet{2}(), 20) == dna"CATCCGCTCCTAGGCCATTT"

    @test randseq(StableRNG(SEED), RNAAlphabet{4}(), 20) == rna"CUCUUCGUAUGCCGUACCGU"
    @test randseq(StableRNG(SEED), RNAAlphabet{2}(), 20) == rna"CAUCCGCUCCUAGGCCAUUU"

    # Casual tests to see that it can use the global RNG automatically
    seq = randseq(DNAAlphabet{2}(), 100)
    test_isseq(seq, DNAAlphabet{2}, 100)

    seq = randseq(RNAAlphabet{4}(), 100)
    test_isseq(seq, RNAAlphabet{4}, 100)

    seq = randseq(AminoAcidAlphabet(), 100)
    test_isseq(seq, AminoAcidAlphabet, 100)
end # randseq

@testset "Simple constructors" begin
    manual = randseq(StableRNG(SEED), AminoAcidAlphabet(), 100)
    automatic = randaaseq(StableRNG(SEED), 100)
    @test automatic == manual

    manual = randseq(StableRNG(SEED), DNAAlphabet{4}(), 100)
    automatic = randdnaseq(StableRNG(SEED), 100)
    @test automatic == manual

    manual = randseq(StableRNG(SEED), RNAAlphabet{4}(), 100)
    automatic = randrnaseq(StableRNG(SEED), 100)
    @test automatic == manual

    # Casual tests to see that it can use the global RNG automatically
    seq = randdnaseq(10)
    test_isseq(seq, DNAAlphabet{4}, 10)

    seq = randrnaseq(10)
    test_isseq(seq, RNAAlphabet{4}, 10)

    seq = randaaseq(10)
    test_isseq(seq, AminoAcidAlphabet, 10)
end # simple constructors

@testset "Custom alphabet" begin
    Base.length(A::MyAlphabet) = 6
    BioSequences.symbols(A::MyAlphabet) = (DNA_A, DNA_C, DNA_G, DNA_T, RNA_U, DNA_N)
    BioSequences.BitsPerSymbol(A::MyAlphabet) = BioSequences.BitsPerSymbol{8}()
    BioSequences.encode(A::MyAlphabet, x::DNA) = reinterpret(UInt8, x)
    BioSequences.encode(A::MyAlphabet, x::RNA) = reinterpret(UInt8, x) | 0x10
    Base.eltype(A::MyAlphabet) = NucleicAcid

    function BioSequences.decode(A::MyAlphabet, x::UInt64)
        uint = UInt8(x)
        if uint & 0x10 == 0x10
            return reinterpret(RNA, uint & 0x0f)
        else
            return reinterpret(DNA, uint)
        end
    end

    seq = randseq(MyAlphabet(), 1000)
    @test typeof(seq) == LongSequence{MyAlphabet}
    @test length(seq) == 1000
    @test Set(seq) == Set(symbols(MyAlphabet()))
end # Custom Alphabet
end # Entire Random LongSequences testset


let
    gw = SimpleGridWorld(size=(2,2))

    pvec = fill(:left, length(states(gw)))

    solver = VectorSolver(pvec)

    p = solve(solver, gw)

    # test default
    @test sprint(showpolicy, gw, p) == " [1, 1] -> :left\n [2, 1] -> :left\n [1, 2] -> :left\n [2, 2] -> :left\n [-1, -1] -> :left"

    # test with small display
    iob = IOBuffer()
    io = IOContext(iob, :limit=>true, :displaysize=>(7, 7))
    showpolicy(io, gw, p, pre="@ ")
    @test String(take!(iob)) == "@ [1, 1] -> :left\n@ [2, 1] -> :left\n@ [1, 2] -> :left\n@ …"

    # test very long policy with small display
    struct M <: MDP{Int, Int}
        n::Int
    end
    iob = IOBuffer()
    io = IOContext(iob, :limit=>true, :displaysize=>(7, 7))
    POMDPs.states(m::MDP) = 1:m.n
    POMDPs.actions(m::MDP) = 1:3
    m = M(1_000_000_000)
    showpolicy(io, m, RandomPolicy(m))
    # Below, the actual values could be different because of the RandomPolicy, but length should be the same
    @test length(String(take!(iob))) == length(" 1 -> 2\n 2 -> 1\n 3 -> 3\n …")
end


"""
    GridData <: AbstractArray

Simulation data specific to a single grid.
"""
abstract type GridData{S,R,T,N} <: StaticArray{S,T,N} end

function (::Type{G})(d::GridData{S,R,T,N}) where {G<:GridData,S,R,T,N}
    args = source(d), dest(d), mask(d), proc(d), opt(d), boundary(d), padval(d),
        sourcestatus(d), deststatus(d)
    G{S,R,T,N,map(typeof, args)...}(args...)
end
function ConstructionBase.constructorof(::Type{T}) where T<:GridData{S,R} where {S,R}
    T.name.wrapper{S,R}
end

GridDataOrReps = Union{GridData, Vector{<:GridData}}

# Array interface
# Base.size(d::GridData{<:Tuple{S}}) where {S} = (Y, X)
# Base.axes(d::GridData) = map(Base.OneTo, size(d))
# Base.eltype(d::GridData{<:Any,<:Any,T}) where T = T
# Base.firstindex(d::GridData) = 1
# Base.lastindex(d::GridData{<:Tuple{S}}) where {S} = Y * X

# Getters
mask(d::GridData) = d.mask
radius(d::GridData{<:Any,R}) where R = R
proc(d::GridData) = d.proc
opt(d::GridData) = d.opt
boundary(d::GridData) = d.boundary
padval(d::GridData) = d.padval
source(d::GridData) = d.source
dest(d::GridData) = d.dest
sourcestatus(d::GridData) = d.sourcestatus
deststatus(d::GridData) = d.deststatus

gridsize(d::GridData) = size(d)
gridsize(A::AbstractArray) = size(A)
gridsize(nt::NamedTuple) = gridsize(first(nt))
gridsize(nt::NamedTuple{(),Tuple{}}) = 0, 0
gridview(d::GridData) = sourceview(d)
sourceview(d::GridData) = view(parent(source(d)), map(a -> a .+ radius(d), axes(d))...)
destview(d::GridData) = view(parent(dest(d)), map(a -> a .+ radius(d), axes(d))...)
source_array_or_view(d::GridData) = source(d) isa OffsetArray ? sourceview(d) : source(d)
dest_array_or_view(d::GridData) = dest(d) isa OffsetArray ? destview(d) : dest(d)

function Base.copy!(grid::GridData{<:Any,R}, A::AbstractArray) where R
    pad_axes = map(ax -> ax .+ R, axes(A))
    copyto!(parent(source(grid)), CartesianIndices(pad_axes), A, CartesianIndices(A))
    _updatestatus!(grid)
    return grid
end
function Base.copy!(A::AbstractArray, grid::GridData{<:Any,R}) where R
    pad_axes = map(ax -> ax .+ R, axes(A))
    copyto!(A, CartesianIndices(A), parent(source(grid)), CartesianIndices(pad_axes))
    return A
end
function Base.copy!(A::AbstractDimArray{T,N}, grid::GridData{<:Any,R}) where {T,N,R}
    copy!(parent(A), grid)
    return A
end
function Base.copy!(grid::GridData{<:Any,R}, A::AbstractDimArray{T,N}) where {R,T,N}
    copy!(grid, parent(A))
    return grid
end
function Base.copy!(dst::GridData{<:Any,RD}, src::GridData{<:Any,RS}) where {RD,RS}
    dst_axes = map(s -> RD:s + RD, gridsize(dst))
    src_axes = map(s -> RS:s + RS, gridsize(src))
    copyto!(parent(source(dst)), CartesianIndices(dst_axes),
            parent(source(src)), CartesianIndices(src_axes)
    )
    return dst
end

@propagate_inbounds Base.getindex(d::GridData{s}, I...) where s = getindex(source(d), I...)
@propagate_inbounds function Base.getindex(d::GridData{s}, i1::Int, I::Int...) where s
    getindex(source(d), i1, I...)
end

# _swapsource => ReadableGridData
# Swap source and dest arrays of a grid
_swapsource(d::Tuple) = map(_swapsource, d)
function _swapsource(grid::GridData)
    src = grid.source
    dst = grid.dest
    @set! grid.dest = src
    @set! grid.source = dst
    srcstatus = grid.sourcestatus
    dststatus = grid.deststatus
    @set! grid.deststatus = srcstatus
    return @set grid.sourcestatus = dststatus
end

# _updatestatus!
# Initialise the block status array.
# This tracks whether anything has to be done in an area of the main array.
_updatestatus!(grid::GridData) = _updatestatus!(opt(grid), grid)
function _updatestatus!(opt::SparseOpt, grid)
    blocksize = 2 * radius(grid)
    src = parent(source(grid))
    for I in CartesianIndices(src)
        # Mark the status block (by default a non-zero value)
        if _isactive(src[I], opt)
            bi = _indtoblock.(Tuple(I), blocksize)
            @inbounds sourcestatus(grid)[bi...] = true
            @inbounds deststatus(grid)[bi...] = true
        end
    end
    return nothing
end
_updatestatus!(opt, grid) = nothing

# _addpadding => OffsetArray{T,N}
# Find the maximum radius required by all rules
# Add padding around the original init array, offset into the negative
# So that the first real cell is still 1, 1
function _addpadding(init::AbstractArray{T,N}, r, padval) where {T,N}
    h, w = size(init)
    paddedsize = h + 4r, w + 2r
    paddedindices = -r + 1:h + 3r, -r + 1:w + r
    sourceparent = fill(convert(eltype(init), padval), paddedsize)
    source = OffsetArray(sourceparent, paddedindices...)
    # Copy the init array to the middle section of the source array
    for j in 1:size(init, 2), i in 1:size(init, 1)
        @inbounds source[i, j] = init[i, j]
    end
    return source
end

# _build_status => (Matrix{Bool}, Matrix{Bool})
# Build block-status arrays
# We add an additional block that is never used so we can
# index into it in the block loop without checking
function _build_status(opt::SparseOpt, source, r)
    hoodsize = 2r + 1
    blocksize = 2r
    nblocs = _indtoblock.(size(source), blocksize) .+ 1
    sourcestatus = zeros(Bool, nblocs)
    deststatus = zeros(Bool, nblocs)
    sourcestatus, deststatus
end
_build_status(opt::PerformanceOpt, init, r) = nothing, nothing

"""
    ReadableGridData <: GridData

    ReadableGridData(grid::GridData)
    ReadableGridData{S,R}(init::AbstractArray, mask, opt, boundary, padval)

Simulation data and storage passed to rules for each timestep.

# Type parameters

- `Y`: number of rows
- `X`: number of columns
- `R`: grid padding radius
"""
struct ReadableGridData{
    S<:Tuple,R,T,N,Sc,D,M,P<:Processor,Op<:PerformanceOpt,Bo,PV,SSt,DSt
} <: GridData{S,R,T,N}
    source::Sc
    dest::D
    mask::M
    proc::P
    opt::Op
    boundary::Bo
    padval::PV
    sourcestatus::SSt
    deststatus::DSt
end
function ReadableGridData{S,R}(
    source::Sc, dest::D, mask::M, proc::P, opt::Op, boundary::Bo,
    padval::PV, sourcestatus::SSt, deststatus::DSt
) where {S,R,Sc<:AbstractArray{T,N},D,M,P,Op,Bo,PV,SSt,DSt} where {T,N}
    ReadableGridData{S,R,T,N,Sc,D,M,P,Op,Bo,PV,SSt,DSt}(
        source, dest, mask, proc, opt, boundary, padval, sourcestatus, deststatus
    )
end
# Generate simulation data to match a ruleset and init array.
@inline function ReadableGridData{S,R}(
    init::AbstractArray, mask, proc, opt, boundary, padval
) where {S,R}
    # We add one extra row and column of status blocks so
    # we dont have to worry about special casing the last block
    if R > 0
        source = _addpadding(init, R, padval)
        dest = _addpadding(init, R, padval)
    else
        if opt isa SparseOpt
            opt = NoOpt()
        end
        source = deepcopy(init)
        dest = deepcopy(init)
    end
    sourcestatus, deststatus = _build_status(opt, source, R)

    grid = ReadableGridData{S,R}(
        source, dest, mask, proc, opt, boundary, padval, sourcestatus, deststatus
    )
    _updatestatus!(grid)
    return grid
end

Base.parent(d::ReadableGridData{S,<:Any,T,N}) where {S,T,N} = SizedArray{S,T,N}(sourceview(d))

"""
    WritableGridData <: GridData

    WritableGridData(grid::GridData)

Passed to rules as write grids, and can be written to directly as an array,
or preferably using `add!` etc. All writes handle updates to SparseOpt()
and writing to the correct source/dest array.

Reads are _always from the source array_, as rules must not be sequential between
cells. This means using e.g. `+=` is not supported, instead use `add!`.
"""
struct WritableGridData{
    S<:Tuple,R,T,N,Sc<:AbstractArray{T,N},D<:AbstractArray{T,N},
    M,P<:Processor,Op<:PerformanceOpt,Bo,PV,SSt,DSt
} <: GridData{S,R,T,N}
    source::Sc
    dest::D
    mask::M
    proc::P
    opt::Op
    boundary::Bo
    padval::PV
    sourcestatus::SSt
    deststatus::DSt
end
function WritableGridData{S,R}(
    source::Sc, dest::D, mask::M, proc::P, opt::Op,
    boundary::Bo, padval::PV, sourcestatus::SSt, deststatus::DSt
) where {S,R,Sc<:AbstractArray{T,N},D<:AbstractArray{T,N},M,P,Op,Bo,PV,SSt,DSt} where {T,N}
    WritableGridData{S,R,T,N,Sc,D,M,P,Op,Bo,PV,SSt,DSt}(
        source, dest, mask, proc, opt, boundary, padval, sourcestatus, deststatus
    )
end

Base.parent(d::WritableGridData{S}) where {S} = SizedArray{S}(destview(d))

### UNSAFE / LOCKS required

# Base.setindex!
# This is not safe for general use.
# It can be used where only identical transformations of a cell
# can happen from any other cell, such as setting all 1s to 2.
@propagate_inbounds function Base.setindex!(d::WritableGridData, x, I...)
    _setindex!(d, proc(d), x, I...)
end
@propagate_inbounds function Base.setindex!(d::WritableGridData, x, i1::Int, I::Int...)
    _setindex!(d, proc(d), x, i1, I...)
end

@propagate_inbounds function _setindex!(d::WritableGridData, opt::SingleCPU, x, I...)
    @boundscheck checkbounds(dest(d), I...)
    @inbounds _setdeststatus!(d, x, I...)
    @inbounds dest(d)[I...] = x
end
@propagate_inbounds function _setindex!(d::WritableGridData, proc::ThreadedCPU, x, I...)
    # Dest status is not threadsafe, even if the
    # setindex itself is safe. So we LOCK
    lock(proc)
    @inbounds _setdeststatus!(d, x, I...)
    unlock(proc)
    @inbounds dest(d)[I...] = x
end

# _setdeststatus!
# Sets the status of the destination block that the current index is in.
# It can't turn of block status as the block is larger than the cell
# But should be used inside a LOCK
_setdeststatus!(d::WritableGridData{<:Any,R}, x, I...) where R =
    _setdeststatus!(d, opt(d), x, I...)
function _setdeststatus!(d::WritableGridData{<:Any,R}, opt::SparseOpt, x, I...) where R
    blockindex = _indtoblock.(I .+ R, 2R)
    @inbounds deststatus(d)[blockindex...] |= !(opt.f(x))
    return nothing
end
_setdeststatus!(d::WritableGridData{<:Any,R}, opt::PerformanceOpt, x, I...) where R = nothing


function OnBoardAntennaPhaseModifier(arg0::Vector3D)
    return OnBoardAntennaPhaseModifier((Vector3D,), arg0)
end

function get_parameters_drivers(obj::OnBoardAntennaPhaseModifier)
    return jcall(obj, "getParametersDrivers", List, ())
end

function modify(obj::OnBoardAntennaPhaseModifier, arg0::EstimatedMeasurement)
    return jcall(obj, "modify", void, (EstimatedMeasurement,), arg0)
end


"""
    blockmodel(inblock, outblock, backbone)

From a `backbone` model, construct a model suitable for learning
a mapping from `inblock` to `outblock`.
"""
function blockmodel end

"""
    blockmodel(inblock::ImageTensor{N}, outblock::OneHotTensor{0}, backbone)
    blockmodel(inblock::ImageTensor{N}, outblock::OneHotTensorMulti{0}, backbone)

Construct a model for N-dimensional image classification. `backbone` should
be a convolutional feature extractor taking in batches of image tensors with
`inblock.nch` color channels.
"""
function blockmodel(
        inblock::ImageTensor{N},
        outblock::Union{OneHotTensor{0}, OneHotTensorMulti{0}},
        backbone) where N
    outsz = Flux.outputsize(backbone, (ntuple(_ -> 256, N)..., inblock.nchannels, 1))
    outch = outsz[end-1]
    head = Models.visionhead(outch, length(outblock.classes), p = 0.)
    return Chain(backbone, head)
end

"""
    blockmodel(inblock::ImageTensor{N}, outblock::OneHotTensor{N}, backbone; kwargs...)

Construct a model for N-dimensional image segmentation. `backbone` should
be a convolutional feature extractor taking in batches of image tensors with
`inblock.nch` color channels. Keyword arguments are passed to [`UNetDynamic`](#).
"""
function blockmodel(inblock::ImageTensor{N}, outblock::OneHotTensor{N}, backbone; kwargs...) where N
    return UNetDynamic(
        backbone,
        (ntuple(_ -> 256, N)..., inblock.nchannels, 1),
        length(outblock.classes);
        kwargs...)
end

"""
    blockmodel(inblock::ImageTensor{N}, outblock::Keypoints{N}, backbone)

Construct a model for image to keypoint regression. `backbone` should
be a convolutional feature extractor taking in batches of image tensors with
`inblock.nch` color channels.
"""
function blockmodel(inblock::ImageTensor{N}, outblock::KeypointTensor{N}, backbone) where N
    outsz = Flux.outputsize(backbone, (ntuple(_ -> 256, N)..., inblock.nchannels, 1))
    outch = outsz[end-1]
    head = Models.visionhead(outch, prod(outblock.sz)*N, p = 0.)
    return Chain(backbone, head)
end


function LeastSquaresAdapter(arg0::LeastSquaresProblem)
    return LeastSquaresAdapter((LeastSquaresProblem,), arg0)
end

function evaluate(obj::LeastSquaresAdapter, arg0::RealVector)
    return jcall(obj, "evaluate", LeastSquaresProblem_Evaluation, (RealVector,), arg0)
end

function get_convergence_checker(obj::LeastSquaresAdapter)
    return jcall(obj, "getConvergenceChecker", ConvergenceChecker, ())
end

function get_evaluation_counter(obj::LeastSquaresAdapter)
    return jcall(obj, "getEvaluationCounter", Incrementor, ())
end

function get_iteration_counter(obj::LeastSquaresAdapter)
    return jcall(obj, "getIterationCounter", Incrementor, ())
end

function get_observation_size(obj::LeastSquaresAdapter)
    return jcall(obj, "getObservationSize", jint, ())
end

function get_parameter_size(obj::LeastSquaresAdapter)
    return jcall(obj, "getParameterSize", jint, ())
end

function get_start(obj::LeastSquaresAdapter)
    return jcall(obj, "getStart", RealVector, ())
end


########################################################################
# Client
########################################################################\

using LinearAlgebra
using SparseArrays
using Random
using Printf

# Abstract class for client
abstract type AbstractClient end

# Client upload model to server
function update! end

# # Client download model to server
# function download end

include("fedProxClient.jl")
include("ScaffoldClient.jl")
include("fedDCDClient.jl")


using ACTRSimulators, Test, ACTRModels, DataFrames, Random
import ACTRSimulators: start!, press_key!
Random.seed!(8985)
include("task.jl")

function can_attend(actr)
    c1(actr) = !isempty(actr.visual_location.buffer)
    c2(actr) = !actr.visual.state.busy
    return c1, c2
end

function can_encode(actr)
    c1(actr) = !isempty(actr.visual.buffer)
    c2(actr) = !actr.imaginal.state.busy
    return c1, c2
end

function can_respond1(actr)
    c1(actr) = !isempty(actr.imaginal.buffer)
    c2(actr) = !actr.motor.state.busy
    return c1, c2
end

function can_respond2(actr)
    c1(actr) = !isempty(actr.imaginal.buffer)
    c2(actr) = !actr.motor.state.busy
    return c1, c2
end

function attend_action(actr, task)
    buffer = actr.visual_location.buffer
    chunk = deepcopy(buffer[1])
    clear_buffer!(actr.visual_location)
    attending!(actr, chunk)
    return nothing
end

function encode_action(actr, task)
    buffer = actr.visual.buffer
    chunk = deepcopy(buffer[1])
    clear_buffer!(actr.visual)
    encoding!(actr, chunk)
    return nothing
end

function motor_action1(actr, task)
    buffer = actr.imaginal.buffer
    chunk = deepcopy(buffer[1])
    clear_buffer!(actr.imaginal)
    key = "j"
    responding!(actr, task, key)
    return nothing
end

function motor_action2(actr, task)
    buffer = actr.imaginal.buffer
    chunk = deepcopy(buffer[1])
    clear_buffer!(actr.imaginal)
    key = "f"
    responding!(actr, task, key)
    return nothing
end

parms = (
    utility_noise = true,
)

scheduler = ACTRScheduler(;model_trace=false)
task = SimpleTask(;scheduler)
procedural = Procedural()
T = vo_to_chunk() |> typeof
visual_location = VisualLocation(buffer=T[])
visual = Visual(buffer=T[])
visicon = VisualObject[]
motor = Motor()
memory = [Chunk(;animal=:dog), Chunk(;animal=:cat)]
declarative = Declarative(;memory)
actr = ACTR(;
    scheduler,
    procedural,
    visual_location,
    visual,
    motor,
    declarative,
    visicon,
    parms...
)

rule1 = Rule(;conditions=can_attend, action=attend_action, actr, task, name="Attend")
push!(procedural.rules, rule1)

rule2 = Rule(;conditions=can_encode, action=encode_action, actr, task, name="Encode")
push!(procedural.rules, rule2)

rule3 = Rule(;conditions=can_respond1, action=motor_action1, actr, task, name="Respond j")
push!(procedural.rules, rule3)

rule4 = Rule(;conditions=can_respond2, action=motor_action2, actr, task, name="Respond f")
push!(procedural.rules, rule4)

map(_->run!(actr, task), 1:1000)

p_j = mean(task.data.key .== "j")

@test p_j ≈ .5 atol = .01