text:
using Combinatorics, Simplicial

"""
The following highly specialized and technical function takes an array of integer vectors (faces), and an array of
births (a birth time for each face). This function implementation is extremely ugly due to speed optimization, alas!
-------------------------------------------
This function does _two_ different things:
 1) It removes all redundant faces
 2) For non-redundant faces it changes the births to the minimal possible birth in each equivalence class of faces.
------------------------------------------
This function can be called in the following two different ways:
RemoveDuplicates!(faces,births);
or, alternatively,
indx, p, births = RemoveDuplicates(faces,births);
----------------------------------------
Here,
p = p=sortperm(faces) ;
indx is the index of the redundant faces according to the new ordering of faces
births is the new vector of faces reordered according to p with modified births (and some deleted)
"""
function  RemoveDuplicates(faces2::Array{Array{TheIntegerType,1},1},  births2::Array{Int,1}) # ::Array{Int,1}
N=length(faces2)
if N != length(births2)
    error(" Mismatch of array lengths in the arguments of RemoveDuplicates!")
end
  p=sortperm(faces2); # lexicographically order faces
  faces=faces2[p]; # Do not mutate the arrays
  births=births2[p];
  # now that all the faces are sorted, we can go through the list
  # of those and since duplicated faces are adjacent, we can quickly compile the list of things to get rif of
  indx=Int[]
# now we loop all faces
i=2; last_i=1; last_face=faces[1];
while i<=N
              # here we scan for repeated faces until they are not repeated
              inner_loop_worked=false
              while isequal(faces[i],last_face)
                    inner_loop_worked=true
                    if i==N
                       break
                    else
                         i+=1
                    end
              end
              # Now we have reached the first non-repeated face, or just reached i==N
              if inner_loop_worked  # this means that during the last inner loop we had to skip some repeated faces
                                    # in particular, i points now to the first non-repeated face

                    if i<N # this is when the i-th element is distinct from the previous one
                       births[last_i]=minimum(births[last_i:i-1])
                       append!(indx,last_i+1:i-1) # add the list of indices to exclude
                    else # i.e. we ended up at  i==N
                          deletion_range_end= isequal(faces[N],last_face) ? N : N-1
                          births[last_i]=minimum(births[last_i:deletion_range_end])
                          append!(indx,last_i+1:deletion_range_end)
                          break # finally, since we reached the end we may break the loop here
                    end   # if i<N
                last_i=i-1
                else
                  last_i=i# since we know the i-th face is not duplicated we can start new checks here
                  i+=1
            end  #   if inner_loop_worked
              last_face=faces[last_i]
end # of outer while
deleteat!(births,indx);
return indx, p, births  # return all the information
end

 """
     Skeleton_a(FS::FiltrationOfSimplicialComplexes,dim::Int)::FiltrationOfSimplicialComplexes
     This Funcion takes a filtration of simplicial complexes and produces a filtration of their skeletons
     Usage:
     S=Skeleton_a(FS,max__dim);
 """
 function Skeleton_a(FS::FiltrationOfSimplicialComplexes,max__dim::Int)::FiltrationOfSimplicialComplexes

 if max__dim<=0; error("The maximal mimension needs to be positive"); end;
 if max__dim>MaximalHomologicalDimension;
   #error("This function is currently not designed to handle skeletons in dimension that is higher than $MaximalHomologicalDimension");
 end
 faces=sort.(collect.(FS.faces))     # here convert faces into an array form. This is done for improving speed
 is_below=FS.dimensions.<max__dim
 Index_of_Below_max__dim=findall(is_below)   # this array contains indices of faces of dimention< max__dim
 Index_of_high_dim=findall(.!is_below)       # this contains indices of faces of dimention>= max__dim
 N_small_faces= length(Index_of_Below_max__dim)  # the number of faces of dimention< max__dim
 There_are_large_facets= !isempty(Index_of_high_dim)

if There_are_large_facets
  # Now we create subsets of length  max__dim+1 for each Maximal face
  lengths_of_combinations= [ binomial(length(faces[i]) ,max__dim+1)  for i in Index_of_high_dim];
  N_large_faces=sum(lengths_of_combinations)

  # Now we assemble all these subsets in one array
  large_facets=Array{Array{TheIntegerType,1}}(undef,N_large_faces);
  large_facets_births= Array{Int}(undef,N_large_faces);
  counter=1
  for i=1:length(Index_of_high_dim)
    fac=faces[Index_of_high_dim[i]]
    b=FS.birth[Index_of_high_dim[i]]
    for f in combinations(fac, max__dim+1)
        large_facets[counter]=f
        large_facets_births[counter]=b
        counter+=1
    end
  end

  ###########################################
  # now we delete duplicates, while preserving correct births
  indx, p, large_facets_births = RemoveDuplicates(large_facets,large_facets_births);
  large_facets=large_facets[p];
  deleteat!(large_facets,indx);
  ##########################################
else # i.e. there are no large facets of size max__dim in this skeleton
    large_facets=Array{TheIntegerType,1}[]
    large_facets_births =Int[]
end

# construct the array of low-dim faces here
small_facets= FS.faces[Index_of_Below_max__dim]
small_facets_births= FS.birth[Index_of_Below_max__dim]
# colect the pieces together
AllFaces=vcat(small_facets,CodeWord.(large_facets));
AllBirths= vcat(small_facets_births,large_facets_births);
p=sortperm(AllBirths);
# Finally, we return the resulting filtration of simplicial complexes.
# Note that here we use "unsafe" version of the constructor, that does not check sanity of this
# filtration, since our construction is designed to produce a sane result
return FiltrationOfSimplicialComplexes(AllFaces[p], AllBirths[p],FS.vertices)
end
####################

"""
       Usage:
       TheSkeleton, f_vector=Skeleton_and_fvector(FS,max__dim);
       This Funcion takes a filtration of simplicial complexes and produces a filtration of their skeletons, as well as
       the f_vecors for each dimension in 1:max__dim (it does not outpur f_0, because it is useless in most use cases)
"""
       function Skeleton_and_fvector(FS::FiltrationOfSimplicialComplexes,max__dim::Int)::Tuple{FiltrationOfSimplicialComplexes,Vector{Vector{Int}}}
      # first, we compute the skeleton
        TheSkeleton=Skeleton_a(FS,max__dim);
       # Next, we create the  f_vector
       unique_births=unique(TheSkeleton.birth);
       N_births=length(unique_births);
       f_vector=Vector{Vector{Int}}(undef,max__dim);
       # first we compute f_vector[end]
       f_vector[max__dim]=zeros(Int,N_births);
       inx= findall(TheSkeleton.dimensions.==max__dim);
       N_facets=length(inx);
       f_vector[end]=[sum(TheSkeleton.birth[ inx].<=fill(unique_births[i], N_facets)) for i=1:N_births];
       # now we loop over the lower dimensions
       cofaces=sort.(collect.(TheSkeleton.faces[inx]));
       coface_births=TheSkeleton.birth[inx];
       for d=max__dim-1:-1:1
           ind_d_faces= findall(TheSkeleton.dimensions.==d);
           N_d_faces=length(ind_d_faces)
           # Now we append to d_dim_faces all the d-dim faces stemming from the cofaces
           N=length(ind_d_faces)+(d+2)*length(cofaces)
           faces=Array{Array{TheIntegerType,1}}(undef,N);
           births= Array{Int}(undef,N);
           faces[1:N_d_faces]=sort.(collect.(TheSkeleton.faces[ind_d_faces]));
           births[1:N_d_faces]=TheSkeleton.birth[ind_d_faces];
           counter=N_d_faces+1;
            # now we fill-out the rest of the arrays faces and births
            for i=1:length(cofaces)
              cf=cofaces[i]
              births[counter:counter+d+1].=coface_births[i]
              for j=1:d+2
                  c=copy(cf)
                  faces[counter]=deleteat!(c,j)
                  counter+=1
              end
            end
        # next we remove duplicates and adjust the births
        if !isempty(faces)
          indx,p,births = RemoveDuplicates(faces,births);
          faces=faces[p];
          deleteat!(faces,indx);
        end
        # fote that the "cleaned up" faces and births are ordered according to faces (and not births!)
        # next, we compute the d-th f-vector
        f_vector[d]=[sum(births.<=fill(unique_births[i], length(births))) for i=1:N_births];
        # Finally, if the loop is going to continue, then we need to update cofaces, and coface_births
          if d>1
          cofaces=faces
          coface_births=births
          end
        end # d=max__dim-1:-1:1
       return TheSkeleton, f_vector
       end
       ####################

function TestComplex(m::Int,n::Int,Threshhold=0.45)
A=rand(m,n);
K,r=Simplicial.DowkerComplex(A)
filt_index= findfirst(r.>Threshhold)-1
deleteat!(K.faces, filt_index:K.depth)
deleteat!(K.dimensions, filt_index:K.depth)
deleteat!(K.birth, filt_index:K.depth)
K.depth=maximum(K.birth);
return K
end

#######


text:
@testset "567.permutation-in-string.jl" begin
    @test check_inclusion("abc", "ccccbbbbaaaa") == false
    @test check_inclusion("ab", "eidbaooo") == true
    @test check_inclusion("ab", "eidboaoo") == false
    @test check_inclusion("hello", "ooolleoooleh") == false
    @test check_inclusion("adc", "dcda") == true
end

ext: jl


text:
# load

using Flux
using Flux: onehotbatch, argmax, crossentropy, throttle, @epochs
using BSON: @save, @load
using Base.Iterators: partition

m = Chain(
  Dense(28^2, 32, relu),
  Dense(32, 10),
  softmax) |> gpu

loss(x, y) = crossentropy(m(x), y)

accuracy(x, y) = mean(argmax(m(x)) .== argmax(y))

@load "fashion_mnist_mlp_model.bson" m
@load "fashion_mnist_mlp_weights.bson" weights
Flux.loadparams!(m, weights)

using MLDatasets # FashionMNIST
using ColorTypes: N0f8, Gray

const Img = Matrix{Gray{N0f8}}

function prepare_test()
    # load full test set
    test_x,  test_y  = FashionMNIST.testdata() # 10_000

    testrange = 1:1_000 # 1:10_000
    test_imgs = Img.([test_x[:,:,i] for i in testrange])
    tX = hcat(float.(reshape.(test_imgs, :))...) |> gpu
    tY = onehotbatch(test_y[testrange], 0:9) |> gpu
    tX, tY
end

tX, tY = prepare_test()
@show(accuracy(tX, tY))
# accuracy(tX, tY) = 0.839


text:
@testset "WT delta" begin
    N = 1024;
    scale = 100;

    X  = zeros(Float64, N); X[N÷2] = 1;
    Y  = map(t -> ContinuousWavelet.eval_analysis(CauchyWavelet(10), t/scale)/scale, LinRange(-N÷2+1,N÷2,N));

    wt = ContinuousWaveletTransform(CauchyWavelet(10), [scale]);
    wX = transform(wt, X);

    @test maximum(abs.(wX[:,1]-Y)) ≈ 0.0 atol=0.1
end


text:
struct AffineStrategy{nx, nu, TP<:SMatrix{nu, nx}, TA<:SVector{nu}}
    "The proportional gain."
    P::TP
    "The feed-forward term."
    α::TA
end

n_states(γ::AffineStrategy{nx}) where {nx} = nx
n_controls(γ::AffineStrategy{nx, nu}) where {nx, nu} = nu

Base.zero(γ::AffineStrategy) = zero(typeof(γ))
function Base.zero(::Type{<:AffineStrategy{nx, nu, TP, TA}}) where {nx, nu, TP, TA}
    return AffineStrategy(zero(TP), zero(TA))
end

"""
    $(TYPEDSIGNATURES)

Computes the control input for a given affine strategy γ, for a state reference
deviation of Δx and a control reference of ũ.
"""
control_input(γ::AffineStrategy, Δx::AbstractVector, ũ::AbstractVector) = ũ - γ.P*Δx - γ.α


text:
using ARDESPOT
using Random

mt = MersenneTwister(1)

rng = MemorizingRNG(mt)

rand(rng)
rand(rng, Float64)
rand(rng, [:a, :b, :c])
# rand(rng, Int) # XXX this doesn't work
randn(rng)
rand(rng, 1:1_000_000)
randn(rng, (3,3))


text:
struct MLM{T,S}
    Ω::T
    Φ::S
    fdspec::Matrix{Float64}
    D::Matrix{Float64}
    function MLM(S, Ω, Φ)
        length(S) == length(Ω) || throw(ArgumentError("Ω and S are not the same size"))
        Dtr = zeros(length(Φ), length(Ω)) # will transpose at the end
        G = zeros(ComplexF64,1,3)
        for (i, (ω,s)) ∈ enumerate(zip(Ω,S))
            for (j,ϕ) ∈ enumerate(Φ)
                allocateG!(G,ω,ϕ)
                Dtr[j, i] = 1.0/real((G) * inv(s) * (G'))[1]
            end
            k = length(Φ)/(sum(Dtr[:, i]) * 2pi)
            Dtr[:, i] .*= k
        end
        Dout = permutedims(Dtr)
        fdspec = zeros(length(Ω), length(Φ))
        for j ∈ 1:length(Φ), (i,s) ∈ enumerate(S)
            fdspec[i, j] = Dout[i, j] * real(s[1, 1])
        end
        new{typeof(Ω), typeof(Φ)}(Ω, Φ, fdspec, Dout)
    end
end

function allocateG!(G,ω,ϕ)
    G[1] = 1.0
    G[2] = -1im * cos(ϕ)
    G[3] = -1im * sin(ϕ)
    return nothing
end


text:
#! format: off

using SIIPExamples

using PowerSystems
using PowerSimulations
const PSI = PowerSimulations
using PowerSystemCaseBuilder

using Dates
using DataFrames

using HiGHS # solver
solver = optimizer_with_attributes(HiGHS.Optimizer, "mip_rel_gap" => 0.05)
odir = mktempdir(".", cleanup = true) #tmpdir for build steps

c_sys5_hy_wk = build_system(SIIPExampleSystems, "5_bus_hydro_wk_sys")
c_sys5_hy_uc = build_system(SIIPExampleSystems, "5_bus_hydro_uc_sys")
c_sys5_hy_ed = build_system(SIIPExampleSystems, "5_bus_hydro_ed_sys")

c_sys5_hy_wk_targets = build_system(SIIPExampleSystems, "5_bus_hydro_wk_sys_with_targets")
c_sys5_hy_uc_targets = build_system(SIIPExampleSystems, "5_bus_hydro_uc_sys_with_targets")
c_sys5_hy_ed_targets = build_system(SIIPExampleSystems, "5_bus_hydro_ed_sys_with_targets")

PSI.JuMP._wrap_in_math_mode(str) = "\$\$ $(replace(str, "__"=>"")) \$\$"

print_tree(HydroGen)

print_tree(PSI.AbstractHydroFormulation)

template = ProblemTemplate()
set_device_model!(template, HydroEnergyReservoir, HydroDispatchRunOfRiver)
set_device_model!(template, HydroDispatch, FixedOutput)
set_device_model!(template, PowerLoad, StaticPowerLoad)

prob = DecisionModel(template, c_sys5_hy_uc, horizon = 2)
build!(prob, output_dir = odir)

PSI.get_jump_model(prob)

template = ProblemTemplate()
set_device_model!(template, HydroEnergyReservoir, HydroDispatchReservoirBudget)
set_device_model!(template, PowerLoad, StaticPowerLoad)

prob = DecisionModel(template, c_sys5_hy_uc, horizon = 2)
build!(prob, output_dir = odir)

PSI.get_jump_model(prob)

template = ProblemTemplate()
set_device_model!(template, HydroEnergyReservoir, HydroDispatchReservoirStorage)
set_device_model!(template, PowerLoad, StaticPowerLoad)

prob = DecisionModel(template, c_sys5_hy_uc_targets, horizon = 24)
build!(prob, output_dir = odir)

PSI.get_jump_model(prob)

template_md = ProblemTemplate()
set_device_model!(template_md, ThermalStandard, ThermalDispatchNoMin)
set_device_model!(template_md, PowerLoad, StaticPowerLoad)
set_device_model!(template_md, HydroEnergyReservoir, HydroDispatchReservoirStorage)

template_da = ProblemTemplate()
set_device_model!(template_da, ThermalStandard, ThermalBasicUnitCommitment)
set_device_model!(template_da, PowerLoad, StaticPowerLoad)
set_device_model!(template_da, HydroEnergyReservoir, HydroDispatchReservoirStorage)

problems = SimulationModels(
    decision_models = [
        DecisionModel(
            template_md,
            c_sys5_hy_wk_targets,
            name = "MD",
            optimizer = solver,
            system_to_file = false,
            initialize_model = false,
        ),
        DecisionModel(
            template_da,
            c_sys5_hy_uc_targets,
            name = "DA",
            optimizer = solver,
            system_to_file = false,
            initialize_model = false,
        ),
    ],
)

sequence = SimulationSequence(
    models = problems,
    feedforwards = Dict(
        "DA" => [
            EnergyLimitFeedforward(
                component_type = HydroEnergyReservoir,
                source = ActivePowerVariable,
                affected_values = [ActivePowerVariable],
                number_of_periods = get_forecast_horizon(c_sys5_hy_uc_targets),
            ),
        ],
    ),
    ini_cond_chronology = InterProblemChronology(),
);

sim = Simulation(
    name = "hydro",
    steps = 1,
    models = problems,
    sequence = sequence,
    simulation_folder = odir,
)

build!(sim)

PSI.get_jump_model(sim.models.decision_models[1])

PSI.get_jump_model(sim.models.decision_models[2])

transform_single_time_series!(c_sys5_hy_wk, 2, Hour(24)) # TODO fix PSI to enable longer intervals of stage 1

template_ed = ProblemTemplate()
set_device_model!(template_ed, ThermalStandard, ThermalDispatchNoMin)
set_device_model!(template_ed, PowerLoad, StaticPowerLoad)
set_device_model!(template_ed, HydroEnergyReservoir, HydroDispatchReservoirStorage)

problems = SimulationModels(
    decision_models = [
        DecisionModel(
            template_md,
            c_sys5_hy_wk_targets,
            name = "MD",
            optimizer = solver,
            system_to_file = false,
            initialize_model = false,
        ),
        DecisionModel(
            template_da,
            c_sys5_hy_uc_targets,
            name = "DA",
            optimizer = solver,
            system_to_file = false,
            initialize_model = false,
        ),
        DecisionModel(
            template_ed,
            c_sys5_hy_ed_targets,
            name = "ED",
            optimizer = solver,
            system_to_file = false,
            initialize_model = false,
        ),
    ],
)

sequence = SimulationSequence(
    models = problems,
    feedforwards = Dict(
        "DA" => [
            EnergyLimitFeedforward(
                component_type = HydroEnergyReservoir,
                source = ActivePowerVariable,
                affected_values = [ActivePowerVariable],
                number_of_periods = get_forecast_horizon(c_sys5_hy_uc_targets),
            ),
        ],
        "ED" => [
            EnergyLimitFeedforward(
                component_type = HydroEnergyReservoir,
                source = ActivePowerVariable,
                affected_values = [ActivePowerVariable],
                number_of_periods = get_forecast_horizon(c_sys5_hy_ed_targets),
            ),
        ],
    ),
    ini_cond_chronology = InterProblemChronology(),
);

sim = Simulation(
    name = "hydro",
    steps = 1,
    models = problems,
    sequence = sequence,
    simulation_folder = odir,
)

build!(sim)

PSI.get_jump_model(sim.models.decision_models[1])

PSI.get_jump_model(sim.models.decision_models[2])

PSI.get_jump_model(sim.models.decision_models[3])

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl

max_stars_repo_name: NREL-SIIP/Examples

max_issues_repo_name: NREL-SIIP/Examples

max_forks_repo_name: NREL-SIIP/Examples


text:
# Test type relations

using Distributions

@assert UnivariateDistribution <: Distribution
@assert MultivariateDistribution <: Distribution
@assert MatrixDistribution <: Distribution

@assert DiscreteDistribution <: Distribution
@assert ContinuousDistribution <: Distribution

@assert DiscreteUnivariateDistribution <: DiscreteDistribution
@assert DiscreteUnivariateDistribution <: UnivariateDistribution
@assert ContinuousUnivariateDistribution <: ContinuousDistribution
@assert ContinuousUnivariateDistribution <: UnivariateDistribution
@assert DiscreteMultivariateDistribution <: DiscreteDistribution
@assert DiscreteMultivariateDistribution <: MultivariateDistribution
@assert ContinuousMultivariateDistribution <: ContinuousDistribution
@assert ContinuousMultivariateDistribution <: MultivariateDistribution
@assert DiscreteMatrixDistribution <: DiscreteDistribution
@assert DiscreteMatrixDistribution <: MatrixDistribution
@assert ContinuousMatrixDistribution <: ContinuousDistribution
@assert ContinuousMatrixDistribution <: MatrixDistribution


text:
#primary structure for portfolio math
#some auxillery structures to improve missing value efficiency
struct Portfolio
  U::Universe
  wₜ::Vector{Float64}
  R::Vector{Float64}
  pμ::Base.RefValue{Float64}

  #the below vector contains all valid weights, and does not include invalid weights
  #thus the shape of wvec is a ragged array, with each array of length equal to the
  #number of valid values
  wmat::Matrix{Float64} #stores the weight vectors
end

#update the portfolio object
#innermost loop and likely a performance bottleneck, so highly optimized
function Portfolio!(P, wₜ)::Portfolio

  P.wₜ .= wₜ
  P.wmat .= P.U.missingreturns .* wₜ
  P.wmat ./= sum(P.wmat, dims=1)

  #build the returns (Note: Threads.@threads provides signficant performance gains)
  P.R .= 0.0
  for c ∈ 1:P.U.N
    for r ∈ P.U.missingindices[c]
      P.R[c] += P.wmat[r,c] * P.U.returns[r,c]
    end
  end

  #update the pointer
  P.pμ[] = mean(P.R)
  #print(P.R)
  return P
end

#construct a brand new portfolio object
function Portfolio(U, wₜ)::Portfolio
  #wvecsize::Vector{Int} = (v->length(v)).(U.missingmap)
  #wvec::Vector{Vector{Float64}} = (n->Vector{Float64}(undef, n)).(wvecsize)
  wmat::Matrix{Float64} = Matrix{Float64}(undef, U.M, U.N)

  P::Portfolio = Portfolio( #pre-allocate if needed, reduces code reuse
    U,
    Vector{Float64}(undef, U.M),
    Vector{Float64}(undef, U.N),
    Ref(-1.0),
    #Matrix{Float64}(undef, U.M, U.N))
    wmat)

  return Portfolio!(P, wₜ)
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
"""
    dist(blob, c)

Measure the distance between a blob and a coordinate `c` using the Mahalanobis distance induced the the blobs measurement covariance
"""
function dist(blob::Blob, c)
    k = blob.kf
    d = Mahalanobis(inv(k.C*covariance(k)*k.C'))
    v1 = SVector(c.I)
    v2 = SVector(k.x[1], k.x[2])
    d(v1,v2)
    # sqrt((c[1]-blob.kf.x[1])^2 + (c[2]-blob.kf.x[2])^2)
end

function dist(sample, c)
    sqrt((c[1]-sample[1])^2 + (c[2]-sample[2])^2)
end

too_far(c::AbstractCorrespondence,blob,coord) = dist(blob, coord) > dist_th(c)

"""
    measurement = assign(c::AbstractCorrespondence, blobs, coordinates)

Assign measurements to blobs using the `AbstractCorrespondence`. Returns `m::Measurement`

#Arguments:
- `blobs`: A vector of blobs
- `coordinates`: A vector of CartesianIndex
"""
assign

function assign(c::HungarianCorrespondence, blobs, coordinates)
    isempty(blobs) && (return Measurement(coordinates, Int[]))
    isempty(coordinates) && (return Measurement(coordinates, zeros(Int, length(blobs))))
    p = c.p
    DM = [dist(b,c)^p for b in blobs, c in coordinates]
    # DM[DM .> bt.dist_th] .= 100000
    assi = hungarian((DM))[1]
    measurement = Measurement(coordinates, assi)
end

function assign(c::NearestNeighborCorrespondence, blobs, coordinates)
    isempty(blobs) && (return Measurement(coordinates, Int[]))
    isempty(coordinates) && (return Measurement(coordinates, zeros(Int, length(blobs))))
    assi = zeros(Int, length(blobs))
    # kdtree = KDTree(Float32.(Matrix(location.(blobs))'))
    kdtree = KDTree(Float32.(Matrix(coordinates)'))
    for (bi,blob) in enumerate(blobs)
        I,d = knn(kdtree, to_static(location(blob)), 1)
        if d[] > dist_th(c)
            continue
        end
        assi[bi] = I[]
    end
    measurement = Measurement(coordinates, assi)
end

function assign(c::MCCorrespondence, blobs, coordinates)
    isempty(blobs) && (return Measurement(coordinates, Int[]))
    isempty(coordinates) && (return Measurement(coordinates, zeros(Int, length(blobs))))
    measurements = map(1:c.num_samples) do i
        samples = sample_new_locations(blobs)
        measurement = assign(c.inner, samples, coordinates)
    end
end

function sample_new_locations(blobs)
    map(blobs) do blob
        d = MvNormal(blob.kf.x[1:2], covariance(blob.kf)[1:2,1:2])
        sample = rand(d)
    end
end


text:
# Newey-West covariance estimator
# from https://github.com/mcreel/Econometrics/blob/508aee681ca42ff1f361fd48cd64de6565ece221/src/NP/NeweyWest.jl
# under MIT licence https://github.com/mcreel/Econometrics/blob/508aee681ca42ff1f361fd48cd64de6565ece221/LICENSE
# and adapted

"""
    function newey_west(Z,nlags=0)

Returns the Newey-West estimator of the asymptotic variance matrix
INPUTS: Z, a nxk matrix with rows the vector zt'
        nlags, the number of lags
OUTPUTS: omegahat, the Newey-West estimator of the covariance matrix
"""
function newey_west(Z,nlags=0)
    n,k = size(Z)
    # de-mean the variables
    Z = Z .- mean(Z,dims=1)
    omegahat = Z'*Z/n       # sample variance
    # automatic lags?
    if nlags == 0
        nlags = max(1, round(Int, n^0.25))
    end
    # sample autocovariances
    for i = 1:nlags
       Zlag = @view(Z[1:n-i,:])
       ZZ = @view(Z[i+1:n,:])
       gamma = (ZZ'*Zlag)/n
       weight = 1.0 - (i/(nlags+1.0))
       omegahat += weight*(gamma + gamma')
    end
    return omegahat
end


text:
module DiffPrecTest

using Statistics, LinearAlgebra, SparseArrays
using ProximalBase, CoordinateDescent, CovSel
using StatsBase, Distributions

export
  DiffPrecResultBoot,
  DiffPrecResultNormal,
  SimulationResult,
  AsymmetricOracleBoot,
  SymmetricOracleBoot,
  AsymmetricOracleNormal,
  SymmetricOracleNormal,
  SymmetricNormal,
  AsymmetricNormal,
  SeparateNormal,
  DecorrelatedScoreNormal,
  DecorrelatedScoreOracleNormal,
  ReducedNormal,
  ReducedOracleNormal,
  estimate,
  compute,
  variance,
  computeSimulationResult,

  # first and second stage functions
  diffEstimation, invAsymHessianEstimation, invQSymHessian, invHessianReduced,

  # support estimation
  CDInverseKroneckerLoss, CDInverseReducedSymKroneckerLoss,
  ANTSupport,
  BootStdSupport,
  BootMaxSupport,
  DTraceValidationSupport,
  supportEstimate,

  # bootstrap
  BootstrapEstimates,
  bootstrap,
  simulCI, simulCIstudentized,
  testGlobalNull

####################################
#
#   different solvers
#
####################################

abstract type DiffPrecMethod end
struct AsymmetricOracleBoot <: DiffPrecMethod end
struct SymmetricOracleBoot <: DiffPrecMethod end
struct AsymmetricOracleNormal <: DiffPrecMethod end
struct SymmetricOracleNormal <: DiffPrecMethod end
struct AsymmetricNormal <: DiffPrecMethod end
struct SymmetricNormal <: DiffPrecMethod end
struct SeparateNormal <: DiffPrecMethod end

struct OneStep <: DiffPrecMethod end             # not implemented yet
struct OneStepRefit <: DiffPrecMethod end        # not implemented yet
struct DecorrelatedScoreNormal <: DiffPrecMethod end
struct DecorrelatedScoreOracleNormal <: DiffPrecMethod end

struct ReducedNormal <: DiffPrecMethod end
struct ReducedOracleNormal <: DiffPrecMethod end

####################################
#
#   different outputs
#
####################################

struct DiffPrecResultBoot
  p::Float64
  boot_p::Vector{Float64}
end

struct DiffPrecResultNormal
  p::Float64
  std::Float64
end

struct SimulationResult
  bias::Float64
  coverage::Float64
  lenCoverage::Float64
  p_mean::Float64
  p_var::Float64
end

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

include("skron.jl")
include("ustats.jl")
include("support.jl")
include("variance.jl")
include("diffEstimation.jl")
include("invHessianEstimation.jl")
include("util.jl")
include("bootstrap.jl")
include("reducedDiffEstim.jl")

##############################
#
#   our procedure
#
##############################

function estimate(::SymmetricNormal, X, Y, row, col;
    Sx::Union{Symmetric, Nothing}=nothing,
    Sy::Union{Symmetric, Nothing}=nothing,
    Δ::Union{SymmetricSparseIterate, Nothing}=nothing,
    suppΔ::Union{BitArray{2}, Nothing}=nothing,
    ω::Union{SparseIterate, Nothing}=nothing,
    suppω::Union{BitArray{2}, Nothing}=nothing
    )

  nx, p = size(X)
  ny    = size(Y, 1)

  if Sx === nothing
      Sx = Symmetric( X'X / nx )
  end
  if Sy === nothing
      Sy = Symmetric( Y'Y / ny )
  end

  # first stage
  λ = 1.01 * quantile( Normal(), 1. - 0.1 / (p*(p+1)) )
  if Δ === nothing && suppΔ === nothing
    Δ = diffEstimation(Sx, Sy, X, Y, λ)
    suppΔ = getSupport(Δ)
  end
  if suppΔ === nothing
    suppΔ = getSupport(Δ)
  end

  # second stage
  if ω === nothing && suppω === nothing
    j = sub2indLowerTriangular(p, row, col)
    ω = invQSymHessian(Sx, Sy, j, X, Y, λ)
    suppω = getSupport(ω, p)
  end
  if suppω === nothing
    suppω = getSupport(ω, p)
  end

  # refit stage
  indS = getLinearSupport(row, col, suppΔ, suppω)
  estimate(SymmetricOracleNormal(), Sx, Sy, X, Y, row, col, indS), Δ, suppΔ, ω, suppω, indS
end

####################################
#
#   competing procedures
#
####################################

function estimate(::AsymmetricNormal, X, Y, row, col;
    Sx::Union{Symmetric, Nothing}=nothing,
    Sy::Union{Symmetric, Nothing}=nothing,
    Δ::Union{SymmetricSparseIterate, Nothing}=nothing,
    suppΔ::Union{BitArray{2}, Nothing}=nothing,
    ω::Union{SparseIterate, Nothing}=nothing,
    suppω::Union{BitArray{2}, Nothing}=nothing
    )

  nx, p = size(X)
  ny    = size(Y, 1)

  if Sx === nothing
      Sx = Symmetric( X'X / nx )
  end
  if Sy === nothing
      Sy = Symmetric( Y'Y / ny)
  end

  # first stage
  λ = 1.01 * quantile( Normal(), 1. - 0.1 / (p*(p+1)) )
  if Δ === nothing && suppΔ === nothing
      Δ = diffEstimation(Sx, Sy, X, Y, λ)
      suppΔ = getSupport(Δ)
  end
  if suppΔ === nothing
      suppΔ = getSupport(Δ)
  end

  # second stage
  if ω === nothing && suppω === nothing
      ω = invAsymHessianEstimation(Sx, Sy, row, col, X, Y, λ)
      suppω = getSupport(ω, p)
  end
  if suppω === nothing
      suppω = getSupport(ω, p)
  end

  # refit stage
  indS = getLinearSupport(row, col, suppΔ, suppω)
  estimate(AsymmetricOracleNormal(), Sx, Sy, X, Y, row, col, indS), Δ, suppΔ, ω, suppω, indS
end

function estimate(
    ::SeparateNormal,
    X,
    Y,
    ri,
    ci;
    Sx::Union{Symmetric, Nothing}=nothing,
    Sy::Union{Symmetric, Nothing}=nothing
    )

  nx, p = size(X)
  ny    = size(Y, 1)

  if Sx === nothing
      Sx = Symmetric( X'X / nx )
  end
  if Sy === nothing
      Sy = Symmetric( Y'Y / ny)
  end

  if ri == ci
      lasso_x_r = lasso(X[:, 1:p .!= ri], X[:, ri], 2. * sqrt(Sx[ri, ri] * log(p) / nx))
      lasso_y_r = lasso(Y[:, 1:p .!= ri], Y[:, ri], 2. * sqrt(Sy[ri, ri] * log(p) / ny))

      T1 = 1. / lasso_x_r.σ^2.
      T2 = 1. / lasso_y_r.σ^2.

      Δab = T1 - T2

      v1 = 2. / (nx * lasso_x_r.σ^4.)
      v2 = 2. / (ny * lasso_y_r.σ^4.)

      return DiffPrecResultNormal(Δab, sqrt(v1 + v2))
  else
      if ri > ci
          ri, ci = ci, ri
      end

      lasso_x_r = lasso(X[:, 1:p .!= ri], X[:, ri], 2. * sqrt(Sx[ri, ri] * log(p) / nx))
      lasso_x_c = lasso(X[:, 1:p .!= ci], X[:, ci], 2. * sqrt(Sx[ci, ci] * log(p) / nx))
      lasso_y_r = lasso(Y[:, 1:p .!= ri], Y[:, ri], 2. * sqrt(Sy[ri, ri] * log(p) / ny))
      lasso_y_c = lasso(Y[:, 1:p .!= ci], Y[:, ci], 2. * sqrt(Sy[ci, ci] * log(p) / ny))

      T1 = dot(lasso_x_r.residuals, lasso_x_c.residuals) / nx + lasso_x_r.σ^2. * lasso_x_c.x[ri] + lasso_x_c.σ^2. * lasso_x_r.x[ci-1]
      T2 = dot(lasso_y_r.residuals, lasso_y_c.residuals) / ny + lasso_y_r.σ^2. * lasso_y_c.x[ri] + lasso_y_c.σ^2. * lasso_y_r.x[ci-1]

      T1 = T1 / ( lasso_x_r.σ^2. * lasso_x_c.σ^2.)
      T2 = T2 / ( lasso_y_r.σ^2. * lasso_y_c.σ^2.)

      Δab = T2 - T1

      v1 = (1. + lasso_x_c.x[ri]^2. * lasso_x_r.σ^2. / lasso_x_c.σ^2.) / (nx * lasso_x_r.σ^2. * lasso_x_c.σ^2.)
      v2 = (1. + lasso_y_c.x[ri]^2. * lasso_y_r.σ^2. / lasso_y_c.σ^2.) / (ny * lasso_y_r.σ^2. * lasso_y_c.σ^2.)

      return DiffPrecResultNormal(Δab, sqrt(v1 + v2))
  end
end

####################################
#
#   oracle procedures
#
####################################

# indS --- coordinates for the nonzero element of the true Δ (LinearIndices)
function estimate(::AsymmetricOracleBoot,
    Sx::Symmetric,
    Sy::Symmetric,
    X,
    Y,
    row,
    col,
    indS::Vector{Int64};
    bootSamples::Int64=1000
    )

  nx, p = size(X)
  ny    = size(Y, 1)

  makeFirst!(indS, p, row, col)

  A = kron_sub(Sy, Sx, indS, indS)
  Δ = A \ (Sy[indS] - Sx[indS])
  Δab = row == col ? Δ[1] : ( Δ[1] + Δ[2] ) / 2.

  boot_est = zeros(Float64, bootSamples)
  for b=1:bootSamples
     bX_ind = sample(1:nx, nx)
     bY_ind = sample(1:ny, ny)
     bSx = cov(X[bX_ind, :])
     bSy = cov(Y[bY_ind, :])
     kron_sub!(A, bSy, bSx, indS, indS)
     Δb = A \ (bSy[indS] - bSx[indS])
     boot_est[b] = row == col ? Δb[1] : ( Δb[1] + Δb[2] ) / 2.
  end

  DiffPrecResultBoot(Δab, boot_est)
end

# indS --- coordinates for the nonzero element of the true Δ (LinearIndices)
function estimate(::SymmetricOracleBoot,
    Sx::Symmetric,
    Sy::Symmetric,
    X,
    Y,
    row,
    col,
    indS::Vector{Int64};
    bootSamples::Int64=1000
    )

  nx, p = size(X)
  ny    = size(Y, 1)

  makeFirst!(indS, p, row, col)

  C = skron_sub(Sx, Sy, indS, indS)
  Δab = (C \ (Sy[indS] - Sx[indS]))[1]

  boot_est = zeros(Float64, bootSamples)
  for b=1:bootSamples
     bX_ind = sample(1:nx, nx)
     bY_ind = sample(1:ny, ny)
     bSx = cov(X[bX_ind, :])
     bSy = cov(Y[bY_ind, :])
     skron_sub!(C, bSx, bSy, indS, indS)
     boot_est[b] = (C \ (bSy[indS] - bSx[indS]))[1]
  end

  DiffPrecResultBoot(Δab, boot_est)
end

function estimate(
    ::AsymmetricOracleNormal,
    Sx::Symmetric,
    Sy::Symmetric,
    X,
    Y,
    row,
    col,
    indS::Vector{Int64}
    )

  nx, p = size(X)
  ny    = size(Y, 1)

  makeFirst!(indS, p, row, col)

  A = kron_sub(Sy, Sx, indS, indS)
  Δ = A \ (Sy[indS] - Sx[indS])

  tmp = zeros(length(indS))
  tmp[1] = 1.
  ω = A \ tmp

  v = variance(AsymmetricOracleNormal(), Sx, Sy, X, Y, ω, indS, Δ, indS)
  DiffPrecResultNormal(Δ[1], sqrt(v))
end

function estimate(
    ::SymmetricOracleNormal,
    Sx::Symmetric,
    Sy::Symmetric,
    X,
    Y,
    row,
    col,
    indS::Vector{Int64}
    )

  nx, p = size(X)
  ny    = size(Y, 1)

  makeFirst!(indS, p, row, col)
  C = skron_sub(Sx, Sy, indS, indS)

  tmp = zeros(length(indS))
  tmp[1] = 1.
  ω = C \ tmp
  Δ = C \ (Sy[indS] - Sx[indS])
  v = variance(SymmetricOracleNormal(), Sx, Sy, X, Y, ω, indS, Δ, indS)

  DiffPrecResultNormal(Δ[1], sqrt(v))
end

# indΔ does not contain coordinates we are doing inference about
function estimate(
    ::DecorrelatedScoreOracleNormal,
    Sx::Symmetric,
    Sy::Symmetric,
    X,
    Y,
    row,
    col,
    ω::Vector{Float64},
    indω::Vector{Int64},
    Δn::Vector{Float64},
    indΔn::Vector{Int64}
    )

  nx, p = size(X)
  ny    = size(Y, 1)

  v = 0.
  for ci=1:length(indΔ)
      for ri=1:length(indω)
          @inbounds v += _getElemSKron(Sx, Sy, indω[ri], indΔn[ci]) * ω[ri] * Δn[ci]
      end
  end
  for ri=1:length(indω)
      @inbounds v -= ω[ri] * (Sy[indω[ri]] - Sx[indω[ri]])
  end
  a = 0.
  if row == col
      ind = (col-1)*p + row
      for ri=1:length(indω)
          @inbounds a += ω[ri] * _getElemSKron(Sx, Sy, indω[ri], ind)
      end
  else
      ind1 = (col-1)*p + row
      ind2 = (row-1)*p + col
      for ri=1:length(indω)
          @inbounds a += ω[ri] * (_getElemSKron(Sx, Sy, indω[ri], ind1) + _getElemSKron(Sx, Sy, indω[ri], ind2))
      end
  end
  Δab = -v / a

  v = variance(SymmetricOracleNormal(), Sx, Sy, X, Y, ω, indω, Δ, indΔ)

  DiffPrecResultNormal(Δab, sqrt(v))
end

end


text:
##############################################################################
##
## Define a type that iterates through q-grams of a string
##
##############################################################################

immutable QGramIterator{S <: AbstractStringorGraphemeIterator, T <: Integer}
	s::S # grapheme
	l::Int # length of string
	q::T # length of q-grams
end

function Base.start(qgram::QGramIterator)
	(1, qgram.l < qgram.q ? endof(qgram.s) + 1 : chr2ind(qgram.s, qgram.q))
end
function Base.next(qgram::QGramIterator, state)
	istart, iend = state
	element = SubString(qgram.s, istart, iend)
	nextstate = nextind(qgram.s, istart), nextind(qgram.s, iend)
	element, nextstate
end
function Base.done(qgram::QGramIterator, state)
	istart, idend = state
	done(qgram.s, idend)
end
Base.eltype{S <: AbstractString, T}(qgram::QGramIterator{S, T}) = SubString{typeof(qgram.s)}
Base.eltype{S <: GraphemeIterator, T}(qgram::QGramIterator{S, T}) = SubString{typeof(qgram.s.s)}
Base.length(qgram::QGramIterator) = max(qgram.l - qgram.q + 1, 0)
function Base.collect(qgram::QGramIterator)
	x = Array(eltype(qgram), length(qgram))
	i = 0
	for q in qgram
		i += 1
		@inbounds x[i] = q
	end
	x
end
Base.sort(qgram::QGramIterator) = sort!(collect(qgram))

##############################################################################
##
## For each element in union{v1, v2}, this iterator output numbers of times it appears in v1 and the number of times it appears in v2
## v1 and v2 must be sorted vectors
##
##############################################################################

type CountInterator{T1 <: AbstractVector, T2 <: AbstractVector}
	v1::T1
	v2::T2
end
Base.start(s::CountInterator) = (1, 1)

function Base.next(s::CountInterator, state)
	state1, state2 = state
	iter1 = done(s.v2, state2)
	iter2 = done(s.v1, state1)
	if iter1
		@inbounds x1 = s.v1[state1]
	elseif iter2
		@inbounds x2 = s.v2[state2]
	else
		@inbounds x1 = s.v1[state1]
		@inbounds x2 = s.v2[state2]
		iter1 = x1 <= x2
		iter2 = x2 <= x1
	end
	nextstate1 = iter1 ? searchsortedlast(s.v1, x1, state1, length(s.v1), Base.Forward) + 1 : state1
	nextstate2 = iter2 ? searchsortedlast(s.v2, x2, state2, length(s.v2), Base.Forward) + 1 : state2
	((nextstate1 - state1, nextstate2 - state2), (nextstate1, nextstate2))
end

function Base.done(s::CountInterator, state)
	state1, state2 = state
	done(s.v2, state2) && done(s.v1, state1)
end

##############################################################################
##
## Distance on strings is computed by set distance on qgram sets
##
##############################################################################
abstract AbstractQGram <: SemiMetric

function evaluate(dist::AbstractQGram, s1::AbstractStringorGraphemeIterator, s2::AbstractStringorGraphemeIterator, len1::Integer, len2::Integer)
	sort1 = sort(QGramIterator(s1, len1, dist.q))
	sort2 = sort(QGramIterator(s2, len2, dist.q))
	evaluate(dist, CountInterator(sort1, sort2))
end

##############################################################################
##
## q-gram
## Define v(s) a vector on the space of q-uple which contains number of times it appears in s
## For instance v("leila")["il"] =1
## q-gram is ∑ |v(s1, p) - v(s2, p)|
##
##############################################################################

immutable QGram{T <: Integer} <: AbstractQGram
	q::T
end
QGram() = QGram(2)

function evaluate(dist::QGram, countiterator)
	n = 0
	for (n1, n2) in countiterator
		n += abs(n1 - n2)
	end
	n
end

##############################################################################
##
## cosine
##
## 1 - v(s1, p).v(s2, p)  / ||v(s1, p)|| * ||v(s2, p)||
##############################################################################

immutable Cosine{T <: Integer} <: AbstractQGram
	q::T
end
Cosine() = Cosine(2)

function evaluate(dist::Cosine, countiterator)
	norm1, norm2, prodnorm = 0, 0, 0
	for (n1, n2) in countiterator
		norm1 += n1^2
		norm2 += n2^2
		prodnorm += n1 * n2
	end
	1.0 - prodnorm / (sqrt(norm1) * sqrt(norm2))
end

##############################################################################
##
## Jaccard
##
## Denote Q(s, q) the set of tuple of length q in s
## 1 - |intersect(Q(s1, q), Q(s2, q))| / |union(Q(s1, q), Q(s2, q))|
##
##############################################################################

immutable Jaccard{T <: Integer} <: AbstractQGram
	q::T
end
Jaccard() = Jaccard(2)

function evaluate(dist::Jaccard, countiterator)
	ndistinct1, ndistinct2, nintersect = 0, 0, 0
	for (n1, n2) in countiterator
		ndistinct1 += n1 > 0
		ndistinct2 += n2 > 0
		nintersect += (n1 > 0) & (n2 > 0)
	end
	1.0 - nintersect / (ndistinct1 + ndistinct2 - nintersect)
end

##############################################################################
##
## SorensenDice
##
## 1 - 2 * |intersect(Q(s1, q), Q(s2, q))| / (|Q(s1, q)| + |Q(s2, q))|)
##############################################################################

immutable SorensenDice{T <: Integer} <: AbstractQGram
	q::T
end
SorensenDice() = SorensenDice(2)

function evaluate(dist::SorensenDice, countiterator)
	ndistinct1, ndistinct2, nintersect = 0, 0, 0
	for (n1, n2) in countiterator
		ndistinct1 += n1 > 0
		ndistinct2 += n2 > 0
		nintersect += (n1 > 0) & (n2 > 0)
	end
	1.0 - 2.0 * nintersect / (ndistinct1 + ndistinct2)
end

##############################################################################
##
## overlap
##
## 1 -  |intersect(Q(s1, q), Q(s2, q))| / min(|Q(s1, q)|, |Q(s2, q)))
##############################################################################

immutable Overlap{T <: Integer} <: AbstractQGram
	q::T
end
Overlap() = Overlap(2)

function evaluate(dist::Overlap, countiterator)
	ndistinct1, ndistinct2, nintersect = 0, 0, 0
	for (n1, n2) in countiterator
		ndistinct1 += n1 > 0
		ndistinct2 += n2 > 0
		nintersect += (n1 > 0) & (n2 > 0)
	end
	1.0 - nintersect / min(ndistinct1, ndistinct2)
end


text:
using COMCall
using Base.Test

CoInitialize()
iecls = COMCall.CLSIDFromProgID("InternetExplorer.Application")
ie = COM[iecls]
@test typeof(COMCall.QueryInterface(ie, COMCall.BaseIID.IDispatch)) ==  COMCall._IDispatch

# Goal:
#@test ie = COM["InternetExplorer.Application"]
#@test ie[:Navigate]("www.julialang.org")


text:
# This file define interfaces with the filters defined in the filters folder
abstract type KalmanFilter end

const HALF_LOG_2_PI = 0.5 * log(2 * pi)

# Default loglikelihood function for optimization
function optim_loglike(
    model::StateSpaceModel, filter::KalmanFilter, unconstrained_hyperparameters::Vector{Fl}
) where Fl
    reset_filter!(filter)
    update_model_hyperparameters!(model, unconstrained_hyperparameters)
    update_filter_hyperparameters!(filter, model)
    # A way to stabilize the objective function is to take the mean of the
    # log like only for the optimizer
    return optim_kalman_filter(model.system, filter) / size(model.system.y, 1)
end

function update_model_hyperparameters!(
    model::StateSpaceModel, unconstrained_hyperparameters::Vector{Fl}
) where Fl
    register_unconstrained_values!(model, unconstrained_hyperparameters)
    constrain_hyperparameters!(model)
    fill_model_system!(model)
    return model
end

function update_filter_hyperparameters!(filter::KalmanFilter, model::StateSpaceModel)
    fill_model_filter!(filter, model)
    return filter
end

function fill_model_filter!(::KalmanFilter, ::StateSpaceModel)
    return nothing
end

function loglike(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    return optim_loglike(model, filter, get_free_unconstrained_values(model)) * size(model.system.y, 1)
end

"""
    FilterOutput{Fl<:Real}

Structure with the results of the Kalman filter:
* `v`: innovations
* `F`: variance of innovations
* `a`: predictive state
* `att`: filtered state
* `P`: variance of predictive state
* `Ptt`: variance of filtered state
* `Pinf`: diffuse part of the covariance
"""
struct FilterOutput{Fl<:AbstractFloat}
    v::Vector{Vector{Fl}}
    F::Vector{Matrix{Fl}}
    a::Vector{Vector{Fl}}
    att::Vector{Vector{Fl}}
    P::Vector{Matrix{Fl}}
    Ptt::Vector{Matrix{Fl}}
    Pinf::Vector{Matrix{Fl}}

    function FilterOutput(model::StateSpaceModel)
        Fl = typeof_model_elements(model)
        n = size(model.system.y, 1)

        v = Vector{Vector{Fl}}(undef, n)
        F = Vector{Matrix{Fl}}(undef, n)
        a = Vector{Vector{Fl}}(undef, n + 1)
        att = Vector{Vector{Fl}}(undef, n)
        P = Vector{Matrix{Fl}}(undef, n + 1)
        Ptt = Vector{Matrix{Fl}}(undef, n)
        Pinf = Vector{Matrix{Fl}}(undef, n)
        return new{Fl}(v, F, a, att, P, Ptt, Pinf)
    end
end

function assert_possible_to_filter(model::StateSpaceModel)
    if !has_hyperparameter(model) || isfitted(model)
        return true
    end
    error("Model has not been estimated yet, please use `fit!`.")
end

"""
    get_innovations

Returns the innovations `v` obtained with the Kalman filter.
"""
function get_innovations end

get_innovations(filter::FilterOutput) = permutedims(cat(filter.v...; dims=2))

function get_innovations(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    return get_innovations(kalman_filter(model; filter=filter))
end

"""
    get_innovations_variance

Returns the variance `F` of innovations obtained with the Kalman filter.
"""
function get_innovations_variance end

get_innovations_variance(filter_output::FilterOutput) = cat(filter_output.F...; dims=3)

function get_innovations_variance(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    return get_innovations_variance(kalman_filter(model; filter=filter))
end

"""
    get_filtered_state

Returns the filtered state `att` obtained with the Kalman filter.
"""
function get_filtered_state end

get_filtered_state(filter::FilterOutput) = permutedims(cat(filter.att...; dims=2))

function get_filtered_state(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    return get_filtered_state(kalman_filter(model; filter=filter))
end

"""
    get_filtered_state_variance

Returns the variance `Ptt` of the filtered state obtained with the Kalman filter.
"""
function get_filtered_state_variance end

get_filtered_state_variance(filter::FilterOutput) = cat(filter.Ptt...; dims=3)

function get_filtered_state_variance(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    return get_filtered_state_variance(kalman_filter(model; filter=filter))
end

"""
    get_predictive_state

Returns the predictive state `a` obtained with the Kalman filter.
"""
function get_predictive_state end

get_predictive_state(filter::FilterOutput) = permutedims(cat(filter.a...; dims=2))

function get_predictive_state(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    return get_predictive_state(kalman_filter(model; filter=filter))
end

"""
    get_predictive_state_variance

Returns the variance `P` of the predictive state obtained with the Kalman filter.
"""
function get_predictive_state_variance end

get_predictive_state_variance(filter::FilterOutput) = cat(filter.P...; dims=3)

function get_predictive_state_variance(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    return get_predictive_state_variance(kalman_filter(model; filter=filter))
end

"""
    kalman_filter(model::StateSpaceModel; filter::KalmanFilter=default_filter(model)) -> FilterOutput

Filter the data using the desired `filter` and the estimated hyperparameters in `model`.
"""
function kalman_filter(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    filter_output = FilterOutput(model)
    reset_filter!(filter)
    if has_hyperparameter(model)
        free_unconstrained_values = get_free_unconstrained_values(model)
        update_model_hyperparameters!(model, free_unconstrained_values)
        update_filter_hyperparameters!(filter, model)
    end
    return kalman_filter!(filter_output, model.system, filter)
end

"""
    SmootherOutput{Fl<:Real}

Structure with the results of the smoother:
* `alpha`: smoothed state
* `V`: variance of smoothed state
"""
struct SmootherOutput{Fl<:Real}
    alpha::Vector{Vector{Fl}}
    V::Vector{Matrix{Fl}}

    function SmootherOutput(model::StateSpaceModel)
        Fl = typeof_model_elements(model)
        n = size(model.system.y, 1)

        alpha = Vector{Vector{Fl}}(undef, n)
        V = Vector{Matrix{Fl}}(undef, n)
        return new{Fl}(alpha, V)
    end
end

"""
    kalman_smoother(model::StateSpaceModel; filter::KalmanFilter=default_filter(model)) -> FilterOutput

Apply smoother to data filtered by `filter` and the estimated hyperparameters in `model`.
"""
function kalman_smoother(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    filter_output = FilterOutput(model)
    kalman_filter!(filter_output, model.system, filter)
    smoother_output = SmootherOutput(model)
    return kalman_smoother!(smoother_output, model.system, filter_output)
end

"""
    get_smoothed_state

Returns the smoothed state `alpha` obtained with the smoother.
"""
function get_smoothed_state end

get_smoothed_state(smoother::SmootherOutput) = permutedims(cat(smoother.alpha...; dims=2))

function get_smoothed_state(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    return get_smoothed_state(kalman_smoother(model; filter=filter))
end

"""
    get_smoothed_state_variance

Returns the variance `V` of the smoothed state obtained with the smoother.
"""
function get_smoothed_state_variance end

get_smoothed_state_variance(smoother::SmootherOutput) = cat(smoother.V...; dims=3)

function get_smoothed_state_variance(model::StateSpaceModel; filter::KalmanFilter=default_filter(model))
    assert_possible_to_filter(model)
    return get_smoothed_state_variance(kalman_smoother(model; filter=filter))
end


text:
using BenchmarkTools
using LoopVectorization
using LinearAlgebra
using Plots
using UnPack

using EntropyStableEuler.Fluxes2D
import EntropyStableEuler.Fluxes1D: wavespeed_1D # for LF flux
import EntropyStableEuler: γ
using FluxDiffUtils
using StartUpDG
using StartUpDG.ExplicitTimestepUtils

include("HybridizedSBPUtils.jl")
using .HybridizedSBPUtils

N = 3
K1D = 8
T = .25
CFL = .5

# init ref element and mesh
rd = init_reference_tri(N)
VX,VY,EToV = uniform_tri_mesh(K1D)
md = init_DG_mesh(VX,VY,EToV,rd)

# Make domain periodic
LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
build_periodic_boundary_maps!(md,rd,LX,LY)

#####
##### Define initial coefficients and time-stepping
#####

# "Define the initial conditions by interpolation"
@unpack x,y = md
rho = @. 2 + .5*exp(-100*(x^2+y^2))
u = @. 0*x
v = @. 0*x
p = @. rho^γ

Q = primitive_to_conservative(rho,u,v,p)
Qrhskew,Qshskew,Vh,Ph,VhP = build_hSBP_ops(rd)
SBP_ops = (Matrix(Qrhskew'),Matrix(Qshskew'),Vh,Ph,VhP)

# interpolate geofacs to both vol/surf nodes, pack back into md
@unpack Vq,Vf = rd
@unpack rxJ, sxJ, ryJ, syJ = md
rxJ, sxJ, ryJ, syJ = (x->[Vq;Vf]*x).((rxJ, sxJ, ryJ, syJ)) # interp to hybridized points
@pack! md = rxJ, sxJ, ryJ, syJ

# e = 1
# Qh = (x->Vh*x).(Q)
# Qh_log = (Qh...,log.(first(Qh)),log.(last(Qh)))
# rhse = zero.(getindex.(Qh,:,e))
# hadamard_sum!(rhse,(Qrhskew,Qshskew),euler_fluxes,getindex.(Qh_log,:,e))
# rhse_old = dense_hadamard_sum(getindex.(Qh,:,e),(Qrhskew,Qshskew),euler_fluxes)
# @show sum(norm.(rhse .- rhse_old))
#
# function time_had()
#     hadamard_sum!(rhse,(Qrhskew,Qshskew),euler_fluxes,getindex.(Qh_log,:,e))
# end

# # old version
# function dense_hadamard_sum(Qhe,ops,flux_fun)
#
#     (QxTr,QyTr) = ops
#
#     # precompute logs for logmean
#     (rho,u,v,beta) = Qhe
#     Qlog = (log.(rho), log.(beta))
#
#     n = size(QxTr,1)
#     nfields = length(Qhe)
#
#     QF = zero.(Qhe) #ntuple(x->zeros(n),nfields)
#     QFi = zeros(nfields)
#     for i = 1:n
#         Qi = getindex.(Qhe,i)
#         Qlogi = getindex.(Qlog,i)
#
#         fill!(QFi,0)
#         for j = 1:n
#             Qj = getindex.(Qhe,j)
#             Qlogj = getindex.(Qlog,j)
#
#             Fx,Fy = flux_fun(Qi...,Qlogi...,Qj...,Qlogj...)
#             @. QFi += QxTr[j,i]*Fx + QyTr[j,i]*Fy
#         end
#
#         for field in eachindex(Qhe)
#             QF[field][i] = QFi[field]
#         end
#     end
#
#     return QF
# end

# function time_had2()
#     dense_hadamard_sum(getindex.(Qh,:,e),(Qrhskew,Qshskew),euler_fluxes)
# end
# @btime time_had()
# @btime time_had2()

# convert to rho,u,v,beta vars
function cons_to_prim_withlogs(rho,rhou,rhov,E)
    rho,u,v,beta = conservative_to_primitive_beta(rho,rhou,rhov,E)
    return (rho,u,v,beta),(log.(rho),log.(beta))
end

mxm_accum!(X,x,e) = X[:,e] .+= 2*Ph*x

function rhs(Q, SBP_ops, rd::RefElemData, md::MeshData)

    QrTr,QsTr,Vh,Ph,VhP = SBP_ops
    @unpack LIFT,Vf,Vq = rd
    @unpack rxJ,sxJ,ryJ,syJ,J,K = md
    @unpack nxJ,nyJ,sJ = md
    @unpack mapP,mapB = md

    Nh,Nq = size(VhP)

    # entropy var projection
    VU = (x->VhP*x).(v_ufun((x->Vq*x).(Q)...))
    Uh = u_vfun(VU...)
    Qh,logQh = cons_to_prim_withlogs(Uh...)
    Qh_log = (Qh...,logQh...)

    # compute face values
    QM = (x->x[Nq+1:Nh,:]).(Qh_log)
    QP = (x->x[mapP]).(QM)

    # simple lax friedrichs dissipation
    Uf =  (x->x[Nq+1:Nh,:]).(Uh)
    (rhoM,rhouM,rhovM,EM) = Uf
    rhoUM_n = @. (rhouM*nxJ + rhovM*nyJ)/sJ
    lam = abs.(wavespeed_1D(rhoM,rhoUM_n,EM))
    LFc = .25*max.(lam,lam[mapP]).*sJ

    fSx,fSy = euler_fluxes(QM...,QP...)
    normal_flux(fx,fy,u) = fx.*nxJ + fy.*nyJ - LFc.*(u[mapP]-u)
    flux = map(normal_flux,fSx,fSy,Uf)
    rhsQ = (x->LIFT*x).(flux)

    # compute volume contributions using flux differencing
    rhse = zero.(getindex.(Qh,:,1))
    for e = 1:K
        QxTr = rxJ[1,e]*QrTr + sxJ[1,e]*QsTr
        QyTr = ryJ[1,e]*QrTr + syJ[1,e]*QsTr
        # hadamard_sum!(rhse,(QxTr,QyTr),euler_fluxes,getindex.(Qh_log,:,e))

        # will compute sum(Ax.*Fx + Ay.*Fy,dims=2)
        # euler_fluxes(QL_log_e...,QR_log_e...)
        hadamard_sum!(rhse,(QxTr,QyTr),euler_fluxes,getindex.(Qh_log,:,e);
                      skip_index=(i,j)->(i>Nq)&(j>Nq))
        # dense_hadamard_sum(getindex.(Qh,:,e),(QxTr,QyTr),euler_fluxes)
        mxm_accum!.(rhsQ,rhse,e)
    end

    rhstest = zero(eltype(first(rhsQ)))
    compute_rhstest = false
    if compute_rhstest
        for fld in eachindex(rhsQ)
            VUq = VU[fld][1:Nq,:]
            rhstest += sum(md.wJq.*VUq.*(Vq*rhsQ[fld]))
        end
    end

    return map(x -> -x./J,rhsQ),rhstest
end

# rhsQ,rhstest = rhs(Q,SBP_ops,rd,md)
# error("d")

# "Time integration coefficients"
rk4a,rk4b,rk4c = ck45()
CN = (N+1)*(N+2)/2  # estimated trace constant
dt = CFL * 2 / (CN*K1D)
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

# "Perform time-stepping"
resQ = zero.(Q)
for i = 1:Nsteps
    for INTRK = 1:5
        rhsQ,rhstest = rhs(Q,SBP_ops,rd,md)
        bcopy!.(resQ, @. rk4a[INTRK]*resQ + dt*rhsQ)
        bcopy!.(Q,@. Q + rk4b[INTRK]*resQ)
    end

    if i%10==0 || i==Nsteps
        println("Time step $i out of $Nsteps")
    end
end

# plotting nodes
gr(aspect_ratio=1,legend=false)
scatter((x->rd.Vp*x).((x,y,Q[1])),zcolor=rd.Vp*Q[1],msw=0,ms=2,cam=(0,90))

#
#
# # dopri initialization
# rka,rkE,rkc = dp56()
# PIparams = init_PI_controller(5)
# Qtmp = similar.(Q)
# rhsQrk = ntuple(x->zero.(Q),length(rkE))
# prevErrEst = nothing
# rhsQ = rhs_avx(Q,rd,md)
# bcopy!.(rhsQrk[1],rhsQ) # initialize DOPRI rhs (FSAL property)
#
# t = 0.0
# i = 0
# dthist = [dt]
#
# while t < T
#     for INTRK = 2:7
#         k = zero.(Qtmp)
#         for s = 1:INTRK-1
#             bcopy!.(k, @. k + rka[INTRK,s]*rhsQrk[s])
#         end
#         bcopy!.(Qtmp, @. Q + dt*k)
#         rhsQ = rhs_avx(Qtmp,rd,md)
#         bcopy!.(rhsQrk[INTRK],rhsQ)
#     end
#
#     global t,dt,i,prevErrEst
#     accept_step, dtnew, prevErrEst
#         = compute_adaptive_dt(Q,rhsQrk,dt,rkE,PIparams,prevErrEst)
#     if accept_step
#         t += dt
#         bcopy!.(Q, Qtmp)
#         bcopy!.(rhsQrk[1], rhsQrk[7]) # use FSAL property
#         push!(dthist,dt) # store dt history
#     end
#     dt = min(T-t,dtnew)
#     i = i + 1  # number of total steps attempted
#
#     if i%10==0
#         println("i = $i, t = $t, dt = $dtnew, errEst = $prevErrEst")
#     end
# end
#
# # plotting nodes
# gr(aspect_ratio=1,legend=false)
# scatter((x->rd.Vp*x).((x,y,Q[1])),zcolor=rd.Vp*Q[1], msw=0,ms=2,cam=(0,90))

max_stars_repo_name: yiminllin/ESDG-PosLimit

max_issues_repo_name: yiminllin/ESDG-PosLimit

max_forks_repo_name: yiminllin/ESDG-PosLimit


text:
module FluidQueues
using DiscretisedFluidQueues
import LinearAlgebra
using JSON

CMEParams = JSON.parsefile((@__DIR__)*"/iltparams.json")

include("inverse_laplace.jl")
include("SFM_operators.jl")

end


text:
module GeometricFlux

using DelimitedFiles
using SparseArrays
using Statistics: mean
using LinearAlgebra: Adjoint, norm, Transpose
using Random
using Reexport

using CUDA, CUDA.CUSPARSE
using ChainRulesCore
using ChainRulesCore: @non_differentiable
using FillArrays: Fill
using Flux
using Flux: glorot_uniform, leakyrelu, GRUCell, @functor
@reexport using GraphSignals
using Graphs
using NNlib, NNlibCUDA
using Zygote

import Word2Vec: word2vec, wordvectors, get_vector

const ConcreteFeaturedGraph = Union{FeaturedGraph,FeaturedSubgraph}

export
    # layers/graphlayers
    AbstractGraphLayer,

    # layers/gn
    GraphNet,

    # layers/msgpass
    MessagePassing,

    # layers/conv
    GCNConv,
    ChebConv,
    GraphConv,
    GATConv,
    GatedGraphConv,
    EdgeConv,
    GINConv,
    CGConv,

    # layer/pool
    GlobalPool,
    LocalPool,
    TopKPool,
    topk_index,

    # models
    GAE,
    VGAE,
    InnerProductDecoder,
    VariationalEncoder,

    # layer/selector
    bypass_graph,

    # utils
    generate_cluster,

    #node2vec
    node2vec

include("datasets.jl")

include("utils.jl")

include("layers/graphlayers.jl")
include("layers/gn.jl")
include("layers/msgpass.jl")

include("layers/conv.jl")
include("layers/pool.jl")
include("models.jl")
include("layers/misc.jl")

include("sampling.jl")
include("embedding/node2vec.jl")

using .Datasets

end


text:
using Distributions
using StatsBase

"""
ABC rejection for parameter inference

Algorithm:
	1. Sample parameters from prior distribtion
	2. simulate data with parameters accodring to given model
	3. if distance(data,simulated_data) < epsilon then accept the parameters
	4. repeat until reached the set number of repeats

Input:

	- alg: an internal type of the algorithm
	- model: an internal type of the model

Output:

	- internal type of ABC_population

"""
function abc_rejection(alg::ABC_Algorithm,
						model::ABC_Model,
						path::ASCIIString)

	counter = 0 #general counter for the number of times the simulation runs
	num_accepted = 0 #counter for the number of accepted parameters
	data = Array{Array{Union{Array{Array{Float64,1},1},Array{Random_Walk,1}},1},1}(alg.particles) #output array of the simulated data from each particle
	posterior = Array{Array{Float64,1},1}(alg.particles) #output array, for each parameter we have a posterior distribution
	weights = Array{Float64,1}(alg.particles)
	distances = Array{Float64,1}(alg.particles)

	while num_accepted < alg.particles

		counter += 1

		#sample a new parameter accodring to the prior
		params = sample_params_from_prior(model.priors)

		#simulate a dataset with the sampled parameters and check if the parameter were accepted
		(accepted_sim,accepted_dis,sim_data,sim_TMs) = simulate_and_compare_data(alg,model,params,1)

		if sim_data==false
			#params need to be rejected
			continue
		end

		is_accepted = false
		#for the random walk model with gradient, each spatio-temporal cluster needs to have a distance < epsilon
		if model.integration_mode=="random_walk_gradient"
			"""
			if accepted_sim==model.n_clusters
				#accept params only is all clusters were below epsilon
				is_accepted=true
			end
			"""
			if sum(accepted_dis)<(model.n_clusters*alg.epsilon[1]) && accepted_sim>0
				#accept params when the sum of the distances was smaller than the sum of epsilons
				#and at least one simulation was accepted
				is_accepted=true

				#we want to plot the TMs 10% of the times
				if num_accepted%(alg.particles/10)==0 && num_accepted!=0
					plot_TMs(sim_TMs,path,num_accepted,1,num_groups=round(Int,sqrt(model.n_clusters)))
				end
			end

		else
			if accepted_sim>0
				#at least one simulated dataset was accepted
				is_accepted=true
			end
		end

		if is_accepted
			#add to output vector
			num_accepted += 1
			data[num_accepted] = sim_data
			posterior[num_accepted] = params
			#the weights in this case is the number of accepted simulations
			weights[num_accepted] = accepted_sim
			distances[num_accepted] = maximum(accepted_dis)#accepted_dis[1]
			#TODO: change that all distances are saved

			println("accepted particles: ", num_accepted)
			#println(" particles: ", params)
			#println("accepted distances: ", accepted_dis)

		end

	end

	#normalize the weights
	weights = weights ./ sum(weights)

	return [ABC_population(1,counter,alg.epsilon[1],data,posterior,weights,distances)]

end

"""
ABC SMC for parameter inference

Algorithm:
	1. create populations according to alg specifications
	2. for the first population : run ABC rejection algorithm with the first epsilon
	3. for any other population:
		3.1. sample particle from the previous population
		3.2 perturb particle according to a pertrubation kernel
		3.3. simulate a data set using the perturbed particle and compare with original data set with the corresponding epsilon
		3.4 if particle is accepted, calculate its weight
	4. return all populations

Input:

	- alg: an internal type of the algorithm
	- model: an internal type of the model
	- pop_num : number of populations when using adaptive epsilons

Output:

	- internal type of ABC_population
"""
function abc_smc(alg::ABC_Algorithm,
					model::ABC_Model,
					path::ASCIIString;
					pop_num::Int64=5)

	#get epsilon values for each population
	epsilon_arr = alg.epsilon

	num_pop = length(alg.epsilon)

	#adaptive epsilon
	if length(alg.epsilon)==1
		num_pop = pop_num
	end

	#save results from each population
	results = Array{ABC_population,1}(num_pop)

	#run simulation for each population separately

	for pop in 1:num_pop

		println("population number: ", pop)
		tic()

		if pop==1

			#the parameters will be sampled from the prior only
			#simply run ABC rejection once
			current_pop = abc_rejection(alg,model,path)[1]

		else
			prev_pop = results[pop-1]
			current_pop = create_population(alg,model,pop,prev_pop.particles,prev_pop.weights,prev_pop.distances,path)

		end

		results[pop] = current_pop
		#save the results of the current population
		save_pop(pop,current_pop,model,path)
		toc()

	end

	return results

end

"""
ABC SMC for model selection
"""
function abc_smc_model_selection()
#TODO
end

########################################################## HELPER FUNCTIONS ######################################################################

"""
sample parameters from priors

Input:
	- priors (from the model) : an array with prior distribution (or constant Float64) for each parameter

Output:
	- params : array of Float64 with one sample per parameter
"""
function sample_params_from_prior(priors::Array{Union{Float64,Distributions.Uniform,Distributions.DiscreteUniform,Distributions.Normal,Distributions.LogNormal},1})

	params = zeros(length(priors))

	for i in 1:length(priors)

		if typeof(priors[i]) == Float64
			#constant prior
			params[i] = priors[i]

		else
			#prior is taken from the corresponding distribution
			params[i] = rand(priors[i])

		end

	end

	return params

end

"""
run simulation in order to generate simulated data set with the given parameters

Input:
	- alg : the internal type of the algorithm
	- model : the internal type of the model
	- params : the parameters to generate the dataset with

Output:

	- An array containing all generated datasets (one for deterministic, multiple for stochastic)
		* each elemnt in this array is an array of array, where each array includes data values for all species at a specific time time point

"""
function run_simulation(alg::ABC_Algorithm,model::ABC_Model,params::Array{Float64,1})

	#generate a number of data-sets (as set in the alg fields)
	if model.integration_mode=="Gillespie"
		integration_repeats = model.integration_repeats
	elseif model.integration_mode=="random_walk_gradient"
		integration_repeats= model.n_clusters
	else
		integration_repeats=1
	end

	data = Array{Union{Array{Array{Float64,1},1},Array{Random_Walk,1}},1}(integration_repeats)

	if model.integration_mode=="ODE"

		#there is only one repeat for the integration
		res = abc_ode(model.integration_function,model.initial_conditions,alg.time_points,params)
		data[1] = res

	elseif model.integration_mode=="Gillespie"

		for i in 1:integration_repeats

			#println("Gillespie run ", i)

			#data_i = abc_gillespie(model.integration_function,stoch_matrix::Array{Int64,2},model.initial_conditions,time_end::Float64,params)
			#println("params : ", params)
			data_i = Gillespie_abc(params,model.initial_conditions,model.integration_matrix,model.integration_hazards,alg.time_points)
			data[i] = data_i
		end

	elseif model.integration_mode=="random_walk"

		walks = Array{Random_Walk,1}(model.n_walks)

		for i in 1:model.n_walks
			#simulating random walks (BM)
			walks[i] = simulate_random_walk(model.n_steps,model.start_point,params[1],params[2],params[3],model.bias_angle)
		end

		data[1]=walks

	elseif model.integration_mode=="random_walk_gradient"

		for j in 1:model.n_clusters

			walks = Array{Random_Walk,1}(model.n_walks)

			for i in 1:model.n_walks
				#simulating random walks (BM)
				current_walk = simulate_random_walk_attractant_diffusion_gradient(model.n_steps[j][i],model.start_point[j][i],model.cell_radius,model.bias_angle,params,start_time=model.start_time[j][i])
				if current_walk==false
					#could not run the simulation with the given parameters -> parameters should be discarded
					return false
				else
					walks[i] = current_walk
				end
			end

			data[j]=walks

		end

	end

	return data

end

"""
create the next population for the ABC SMC method
"""
function create_population(alg::ABC_Algorithm,
							model::ABC_Model,
							pop_num::Int64,
							prev_particles::Array{Array{Float64,1},1},
							prev_weights::Array{Float64,1},
							prev_dis::Array{Float64,1},
							path::ASCIIString)

	particle = 0
	accepted_particles = 0
	data = Array{Array{Union{Array{Array{Float64,1},1},Array{Random_Walk,1}},1},1}(alg.particles) #output array of the simulated data from each particle
	particles = Array{Array{Float64,1},1}(alg.particles) #array of particles : each item is an array of parameters
	weights = Array{Float64,1}(alg.particles)
	distances = Array{Float64,1}(alg.particles)

	#check if epsilon needs to be updated according to the previous population
	if length(alg.epsilon)<pop_num
		#in this case the epsilon for populations >1 is taken to be the 10% quartile of the distances from the previous population
		alg.epsilon = push!(alg.epsilon,percentile(prev_dis,10))
		println("next epsilon ", alg.epsilon[pop_num])
	end

	while accepted_particles < alg.particles

		particle +=1

		#sample a new particle

		sampled_particles = sample_and_perturb(alg,model,prev_particles,prev_weights)

		#simulate data set with the sampled perturbed parameters
		(num_accepted,accepted_dis,sim_data,sim_TMs) = simulate_and_compare_data(alg,model,sampled_particles,pop_num)

		if sim_data==false
			#params need to be rejected
			continue
		end

		is_accepted = false
		#for the random walk model with gradient, each spatio-temporal cluster needs to have a distance < epsilon
		if model.integration_mode=="random_walk_gradient"
			"""
			if num_accepted==model.n_clusters
				#all clusters were accepted
				is_accepted=true
			end
			"""
			if sum(accepted_dis)<(model.n_clusters*alg.epsilon[pop_num]) && num_accepted>0
				#accept params when the sum of the distances was smaller than the sum of epsilons
				is_accepted=true

				#we want to plot the TMs 10% of the times
				if accepted_particles%(alg.particles/10)==0 && accepted_particles!=0
					plot_TMs(sim_TMs,path,accepted_particles,pop_num,num_groups=round(Int,sqrt(model.n_clusters)))
				end
			end

		else
			if num_accepted>0
				#at least one simulated dataset was accepted
				is_accepted=true
			end
		end

		if is_accepted #at least one simulated dataset was accepted
			#add to output vector
			accepted_particles += 1
			data[accepted_particles] = sim_data
			particles[accepted_particles] = sampled_particles

			#calculate weight for the accepted particle
			p_weight = calculate_particle_weight(sampled_particles,model.priors,num_accepted,prev_weights,prev_particles,alg.kernels)
			weights[accepted_particles] = p_weight
			distances[accepted_particles] = maximum(accepted_dis)#accepted_dis[1]

			println("accepted particles: ", accepted_particles)
			#println(" particles: ", sampled_particles)

		end

	end

	#normalise the weights
	weights = weights ./ sum(weights)

	#update the kernels according to the current population
	update_kernels(particles,alg,model)

	return ABC_population(pop_num,particle,alg.epsilon[pop_num],data,particles,weights,distances)

end

"""
simulate a data set with the given parameters and comapre to the original dataset

Input:
	- alg : internal representation of the algorithm
	- model : internal representation of the model
	- params : the parameters for generating the simulated data set
"""
function simulate_and_compare_data(alg::ABC_Algorithm,
									model::ABC_Model,
									params::Array{Float64,1},
									pop_num::Int64)

	#simulate a dataset with the sampled parameters
	sim_data = run_simulation(alg,model,params)

	if sim_data==false
		#parameters needs to be discarded
		return (0,0,sim_data,0)
	end

	#save the distances for all simulation data
	distances = Array{Float64,1}(length(sim_data))

	#save the transition matrix of the simulated data
	simulated_TMs = Array{Array{Float64,2},1}(length(sim_data))

	#go over all data-sets (one for determinstic, multiple for stochastic)
	for i in 1:length(sim_data)

		sim_data_i = sim_data[i]

		#check the distance between the original dataset and the current simulated dataset
		distance = 0

		if alg.distance_func == "euclidian"
			distance = euclidian_distance(alg.data,sim_data_i)
		elseif alg.distance_func == "hellinger"
			#check the distance between each spatio-temporal cluster separately
			sim_data_TM = transition_matrix(sim_data_i)
			simulated_TMs[i] = sim_data_TM
			distance = hellinger_distance(transition_matrix(alg.data[i]),sim_data_TM)
		else
			error("simulate_and_compare_data: the requested distance function: " ,alg.distance_func, " is not implemented")
			return
		end

		distances[i] = distance

	end

	#go over all the distances and count how many where accepted
	num_accepted_distances = 0
	#accepted_dis = 0
	accepted_dis = Array{Float64,1}()

	for dis in distances

		if dis<=alg.epsilon[pop_num]
			num_accepted_distances +=1
			#accepted_dis=dis
			push!(accepted_dis,dis)
		end

	end

	return (num_accepted_distances,accepted_dis,sim_data,simulated_TMs)
end

function calculate_particle_weight(particle::Array{Float64,1},
									priors::Array{Union{Float64,Distributions.Uniform,Distributions.DiscreteUniform,Distributions.Normal,Distributions.LogNormal},1},
									num_accepted::Int64,
									prev_weights::Array{Float64,1},
									prev_particles::Array{Array{Float64,1},1},
									kernels::Array{Distributions.Uniform,1})

	#get the prior of the particle - multiple all parameters priors
	p_prior = 1.0

	for i in 1:length(particle)

		param_prior = 1.0

		if typeof(priors[i]) == Float64
			#constant prior
				if particle[i]==priors[i]
					param_prior = 1.0
				else
					param_prior = 0.0
				end
		else
			param_prior = pdf(priors[i],particle[i])
		end

		p_prior = p_prior*param_prior
	end

	numerator = p_prior * num_accepted

	#go over all previous weights and multiple by the pertrubation kernel
	denominator = 0.0

	for j in 1:length(prev_particles)

		#go over each parameter and get the pdf of the kernel - in order to get the pdf of the particle
		kernel_prob = 1.0
		for k in 1:length(particle)

			#the parameters of the kernel distribution are now accodring to the previous parameter
			(k_min,k_max) = Distributions.params(kernels[k])
			current_kernel = Uniform(prev_particles[j][k] + k_min , prev_particles[j][k] + k_max)

			kernel_prob = kernel_prob * pdf(current_kernel,particle[k])

			"""kernel_prob = kernel_prob * pdf(kernels[k],particle[k])"""
		end

		denominator += (prev_weights[j] * kernel_prob)
	end

	return (numerator/denominator)

end

function sample_and_perturb(alg::ABC_Algorithm,
							model::ABC_Model,
							prev_particles::Array{Array{Float64,1},1},
							prev_weights::Array{Float64,1})

	zero_prior = true
	params = Array{Float64,1}()

	while zero_prior
		#as long as the prior for the perturb particle is zero, keeping sampling and perturbing

		"""
		#sample parameters from previous population

		rand_u = rand(Uniform(0,1))
		cum_prob = 0

		#sum over all particles weights until reached the random number
		particle_index=0

		for p in 1:alg.particles
			#go over each particle
			#each particle is a vector of parameters
			#the entire vector is associated with a weight
			particle_index=p

			cum_prob = cum_prob + prev_weights[p]

			if (cum_prob>rand_u)
				break
			end
		end

		particle_index = p

		#particle_index is the chosen particle from the previous population

		sampled_from_prev = prev_particles[particle_index]
		"""

		sampled_from_prev = StatsBase.sample(prev_particles,WeightVec(prev_weights))

		#perturb the sampled parameters (accodring to kernel)
		#the kernel is part of the alg and is a specified distribution
		#go over each parameter and perturb it (change in place of array)

		#prior probability of the particle
		particle_prior_prob = 1.0
		params = deepcopy(sampled_from_prev)

		#go over each parameter of the particle
		for i in 1:model.num_params

			params[i] += rand(alg.kernels[i])

			if typeof(model.priors[i]) == Distributions.DiscreteUniform
				#need to round the value of the perturbed parameter
				params[i] = round(Int,params[i])
			end

			#get the prior probability of the perturbed parameter
			if typeof(model.priors[i]) == Float64
				#round the value
				params[i] = round(params[i])
				#constant prior
				if params[i]==model.priors[i]
					param_prior_prob = 1.0
				else
					param_prior_prob = 0.0
				end
			else
				param_prior_prob = pdf(model.priors[i],params[i])
			end

			particle_prior_prob = particle_prior_prob * param_prior_prob

		end

		if particle_prior_prob>0
			zero_prior=false
		end
	end

	return params

end

"""
function to change (in place) perturbation kernels (a field in the algorim type) according to the previous population
"""
function update_kernels(particles::Array{Array{Float64,1},1},alg::ABC_Algorithm,model::ABC_Model)

	#go over each parameter in the population and update the kernel accordingly
	for i in 1:model.num_params
		#get all the current parameter values
		p_vals = [particles[j][i] for j in 1:length(particles)]
		scale = maximum(p_vals) - minimum(p_vals)
		alg.kernels[i] = Uniform(-scale/2.0,scale/2.0)
	end

end


text:
##################################################################
# function to simulate one iteration of the experiment. For a    #
# given number of beetles, function determines if and when each  #
# beetle dies.                                                   #
##################################################################

# total_beetle: positive integer, refers to total number of beetles in experiment
# α_i, β_i, γ_i, p_i: positive reals, parameter values
# z_0: positive real, initial dose of insecticide

# returns length 14 vector where first 13 values correspond
# to the number of beetles that died on that day and final
# value corresponds to the number of beetles that survived.
function beetle_sim(total_beetle, α_i, β_i, γ_i, p_i, z_0)

    results = zeros(14)

    S_prev = 1

    for i = 1:total_beetle
        # number of releases occuring up to or on 13th day
        V_count = rand(Poisson(α_i * 13))

        # if no releases, count beetle as survived
        if V_count == 0
            results[14] = results[14] + 1
            continue
        end

        # draw inter-arrival times for releases
        T_inter = rand(Exponential(1/α_i), V_count)

        # set amount of insecticide
        z_i = z_0

        # death status of beetle
        death = false

        # time in loop begins at 0
        current_t = 0

        # calculate probability of death for each inter-release
        # period. Then remove amount of toxin released at end of
        # period and move on to next period
        for j = 1:(V_count + 1)

            # last iteration, for time between last release and
            # end of 13th day
            if j == V_count + 1

                # prob of death in inter-release interval
                death_p = z_i * γ_i + p_i * z_i *
                        (2*current_t*(13 - T_inter[j-1]) +
                        (13 - T_inter[j-1])^2) / 2

                if death_p > rand(Uniform(0,1))
                    # if the beetle dies, set death true and
                    # break out of for loop
                    death = true
                    # add to death counter of final interval
                    results[13] = results[13] + 1
                end

                break
            end

            # update time for next loop
            current_t = sum(T_inter[1:j])

            # make sure we haven't gone above 13 days
            if current_t > 13
                break
            end

            S_cur = exp(-(γ_i * current_t * z_i + z_i * p_i * current_t^2 / 2))

            death_p = S_prev - S_cur

            # set S_prev for next iteration
            S_prev = S_cur

            # use Uniform to generate random deaths
            if death_p > rand(Uniform(0,1))
                # if the beetle dies, set death true and
                # break out of for loop
                death = true
                # add to death counter
                index = convert(Int,round(current_t, RoundUp))

                results[index] = results[index] + 1
                break
            end

            # remove insecticide released at end of this period
            z_i = z_i - rand(Exponential(1 / β_i))

            # stop loop if negative insecticide
            if z_i <= 0
                break
            end

        end

        if death == false
            results[14] = results[14] + 1
        end

    end

    return results

end

function beetle_sim_ext(total_beetle, α_i, β_i, γ_i, p_i, z_0)

    results = zeros(14)

    for i = 1:total_beetle
        # number of releases occuring up to or on 13th day
        V_count = rand(Poisson(α_i * 13))
        if V_count == 0
            results[14] = results[14] + 1
            continue
        end

        # draw inter-arrival times for first release
        T_inter = rand(Exponential(1/α_i))

        # prob of death in inter-release interval
        death_p = p_i * z_0 *
                        (T_inter^2) / 2

        if death_p > rand(Uniform(0,1))
            # if the beetle dies, set death true and
            # break out of for loop
            death = true
            # add to death counter
            index = convert(Int,round(T_inter, RoundUp))
            if index > 14
                results[14] = results[14] + 1
                continue
            end

            results[index] = results[index] + 1
        else
            results[14] = results[14] + 1
        end

    end

    return results

end

# ABC simulation function for the beetle experiment. Takes as inputs
# original data from experiment, number of iterations, and initial
# dose of insecticide. For each iteration it records the parameter
# values drawn from the priors and the distance of the simulated
# data from the observed data. Returns a dataframe with those values
# for each iteration.

function abc_sim(data, z_0, n, beetle_n, α_prior, β_prior, γ_prior, p_prior)
    results = DataFrame( α = Float64[], β = Float64[], γ = Float64[],
                p = Float64[], dist = Float64[])

    # these lines put data in same format given in Diggle 1984
    sum_1 = sum(data[6:7])
    sum_2 = sum(data[8:13])
    data = vcat(data[1:5], sum_1, sum_2, data[14])
    for i = 1:n

        # draw instances of parameters from prior distributions
        α_i = rand(α_prior)
        β_i = rand(β_prior) + 5
        γ_i = rand(γ_prior)
        p_i = rand(p_prior)

        # run the experiment simulation
        sim_i = beetle_sim(beetle_n, α_i, β_i, γ_i, p_i,z_0)

        # convert data into form from Diggle 1984
        sum_i1 = sum(sim_i[6:7])
        sum_i2 = sum(sim_i[8:13])
        sim_i  = vcat(sim_i[1:5], sum_i1, sum_i2, sim_i[14])

        # calculate distance b/t simulated and observed data
        dist = norm(sim_i - data, 1)

        # record results
        push!(results, [α_i, β_i, γ_i, p_i, dist])

    end
    return results
end

max_stars_repo_name: aidan-cg/undergrad-thesis-ABC

max_issues_repo_name: aidan-cg/undergrad-thesis-ABC

max_forks_repo_name: aidan-cg/undergrad-thesis-ABC


text:
# Spinnaker.jl: wrapper for FLIR/Point Grey Spinnaker SDK
# Copyright (C) 2019 Samuel Powell

# CameraImage.jl: an AbstractArray interface to Image object data
import Base: size, getindex, setindex!, IndexStyle
export CameraImage

# Raw format map
# Fefine how pixel formats relate to native integers
const raw_fmtmap = Dict(PixelFormat_Mono8  => UInt8,
                        PixelFormat_Mono10 => UInt16,
                        PixelFormat_Mono12 => UInt16,
                        PixelFormat_Mono14 => UInt16,
                        PixelFormat_Mono16 => UInt16)

# Normalised format map
# Define how pixel formats relate to normalised fixed point types
const nrm_fmtmap =  Dict(PixelFormat_Mono8  => N0f8,
                         PixelFormat_Mono10 => N6f10,
                         PixelFormat_Mono12 => N4f12,
                         PixelFormat_Mono14 => N2f14,
                         PixelFormat_Mono16 => N0f16)

struct CameraImage{T,N} <: AbstractArray{T,N}
    data::Array{T,N}
    id::Int64
    timestamp::Int64
    exposure::Float64
end

size(a::CameraImage) = size(a.data)
getindex(a::CameraImage, i::Int) = a.data[i]
setindex!(a::CameraImage, v, i::Int) = (a.data[i] = v)
IndexStyle(::Type{<:CameraImage}) = IndexLinear()

id(image::CameraImage) = image.id
timestamp(image::CameraImage) = image.timestamp
exposure(image::CameraImage) = image.exposure

"""
    CameraImage(::SpinImage, ::Type{T}, normalize=false)

    Convert SpinImage to CameraImage, with type conversion and optional normalisation.
    See `getimage` for details of options.
"""
function CameraImage(spinim::SpinImage, ::Type{T}; normalize=false) where T
    width, height = size(spinim)
    imdat = Array{T,2}(undef, width, height)
    _copyimage!(spinim.handle, width, height, imdat, normalize)
    return CameraImage(imdat, id(spinim), timestamp(spinim), exposure(spinim))
end

function _copyimage!(himage_ref::spinImage,
                     width::Integer,
                     height::Integer,
                     image::AbstractArray{T,2},
                     normalize::Bool) where T

    @assert prod(size(image)) == width*height
    hpixfmt = Ref(spinPixelFormatEnums(0))
    spinImageGetPixelFormat(himage_ref, hpixfmt)

    # Map the pixel format to a native integer format. For un-normalized output this is
    # just an unsigned integer of the correct size. For a normalized output a number type
    # from FixedPointNumbers is used to maintain normalisation.
    Ti = UInt8
    if(normalize)
      try
        Ti = nrm_fmtmap[hpixfmt[]]
      catch e
        spinImageRelease(himage_ref)
        throw(e)
      end
    else
      try
        Ti = raw_fmtmap[hpixfmt[]]
      catch e
        spinImageRelease(himage_ref)
        throw(e)
      end
    end

    # Make sure this is a good idea
    sz = Ref(Csize_t(0))
    spinImageGetBufferSize(himage_ref, sz)
    @assert (sizeof(Ti)*width*height) <= sz[]

    # Wrap the image data in an array of the correct pointer type
    rawptr = Ref(Ptr{Cvoid}(0))
    spinImageGetData(himage_ref, rawptr)
    data = unsafe_wrap(Array,  Ptr{Ti}(rawptr[]), (width, height));

    # Convert and copy data from buffer
    copyto!(image, T.(data))
    return image

  end


text:
using Plots, StatsBase, Statistics, LaTeXStrings, JLD

function ColoringTheNode(i, j, Network, Num)
    check = false
    if Network[i, j] == 0
        Network[i, j] = Num
        check = true
        Num += 1
    end
    return Network, Num, check
end

function FindNeighbors(i, j, Network, dim)
    Numbereds = 0
    colorList = []
    for neighbor in ([0,1], [1,0], [0,-1], [-1,0])
        if j + neighbor[2] >= 1 && j + neighbor[2] <= dim && Network[([i, j] + neighbor)...] > 0
            push!(colorList,Network[([i, j] + neighbor)...])
            Numbereds += 1
        end
    end
    push!(colorList, Network[i, j])
    return colorList, Numbereds #colorList contain the node and its neighbors
end

function JoinColors(Network, dim, colorList)
    mincolor = min(colorList...)
    for i in 1:dim
        for j in 1:dim
            if Network[i, j] in colorList
                Network[i, j] = mincolor
            end
        end
    end
    return Network
end

function InitialNetwork(dim, P)
    Network = sample([-1,0], Weights([1-P, P]),(dim,dim))
    return Network
end

function PercolationCheck(dim, P)
    Num = 2
    maxnum = 100000
    Network = InitialNetwork(dim, P)
    Network[1,:] = ones(dim)
    Network[end,:] = fill(maxnum,dim)
    for i in 2:dim-1
        for j in 1:dim
            Network, Num, check = ColoringTheNode(i, j, Network, Num)
            if check == true
                colorList, Numbereds = FindNeighbors(i, j, Network, dim)
                if Numbereds == 0
                    continue
                elseif Numbereds == 1
                    Network[i, j] = colorList[1]
                else
                    Network = JoinColors(Network, dim, colorList)
                end
            end
        end
    end
    if 1 in Network[end,:]
        return 1
    else
        return 0
    end
end

PercolationCheck(40,0.55)


text:
using SimpleGraphs
using SimpleGraphAlgorithms

include("solver.jl")

function cell_dist(a::Cell, b::Cell)
    return sqrt((a.row - b.row)^2 + (a.col - b.col)^2)
end

function adj_matrix(puzzle)
    d, n = size(puzzle)

    # fill cells array
    cells = Array{Cell}(undef, d, d)
    fill_cells!(cells, puzzle)
    adj = zeros(Int, n * n, n * n)

    for i = 1:n*n
        for j = 1:n*n
            if (
                (cells[i].row == cells[j].row) ||
                (cells[i].col == cells[j].col) ||
                (cells[i].group == cells[j].group) ||
                cell_dist(cells[i], cells[j]) < 2
            )
                adj[i, j] = 1
            end
        end
    end

    return adj, cells
end

function graph_solve(puzzle)
    adj, cells = adj_matrix(puzzle)
    g = SimpleGraph(adj)
    m = max_indep_set(g)

    for i in m
        cells[i].is_star = true
    end

    return StarBattleResult(cells, 0)
end


text:
using FileIO
using CSV
using DataFrames
using Plots
using Measures      # for better positioning of figures
using PDFmerger

# if you have python installed, you can use pyplot for better formatting
pyplot()
# gr()

include("processing_sim.jl")

function plot_all_forecasts(
    file_date,
    states,
    dates,
    local_case_dict;
    zoom = false,
    confidence_level = "both",
    truncation_days = 7,
)
    """
    Plots all the forecasts and saves the result as a pdf.
    """

    forecast_origin = string(Dates.Date(file_date) - Dates.Day(truncation_days))

    # read in the sim results and kept TP's
    sim_all_states = CSV.read(
        "results/UoA_forecast_output/" * forecast_origin * "/UoA_samples_" * forecast_origin * ".csv",
        DataFrame,
    )

    # read in the case data
    case_dates = collect(dates)

    # initialise the plot
    fig = plot(
        layout = (4, 2),
        dpi = 200,
        size = (700, 1000),
        margin = 3mm,
        rightmargin = 4mm,
        framestyle = :box,
        legend = false
    )

    for (sp, state) in enumerate(states)

        # get relevant columns of the dataframe and merge to matrix
        D = Matrix(sim_all_states[sim_all_states.state .== state, 3:end - 1])
        # if anything in D is missing, it means the sims didn't work so 0 out those entries
        D[ismissing.(D)] .= 0
        D = Matrix{Int}(D)

        onset_dates = sim_all_states[sim_all_states.state .== state, "onset date"]
        # onset_dates = Date.(onset_dates, dateformat"m yyyy")
        df_D_summary = summarise_forecast_for_plotting(D)

        local_cases = local_case_dict[state][dates .>= onset_dates[1]]
        # boolean for correct case dates
        case_dates_ind = [d ∈ onset_dates ? true : false for d in case_dates]

        onset_dates_lims = Dates.value.([onset_dates[1], onset_dates[end]])

        if confidence_level == "50" || confidence_level == "both"
            plot!(
                fig,
                title = state,
                subplot = sp,
                onset_dates,
                df_D_summary[!, "median"],
                linecolor = 1,
                linewidth = 1,
                color = 1,
                fillalpha = 0.2,
                label = false,
            )
            plot!(
                fig,
                subplot = sp,
                onset_dates,
                [df_D_summary[!, "median"] df_D_summary[!, "median"]],
                fillrange = [df_D_summary[!, "lower"] df_D_summary[!, "upper"]],
                fillalpha = 0.4,
                c = 1,
                label = ["50%" nothing]
            )

        end

        if confidence_level == "95" || confidence_level == "both"
            plot!(
                fig,
                title = state,
                subplot = sp,
                onset_dates,
                df_D_summary[!, "median"],
                linecolor = 1,
                linewidth = 1,
                color = 1,
                fillalpha = 0.2,
                label = false,
            )

            plot!(
                fig,
                subplot = sp,
                onset_dates,
                [df_D_summary[!, "median"] df_D_summary[!, "median"]],
                fillrange = [df_D_summary[!, "bottom"] df_D_summary[!, "top"]],
                fillalpha = 0.2,
                c = 1,
                label = ["95%" nothing]
            )
        end

        plot!(
            fig,
            subplot = sp,
            case_dates[case_dates_ind],
            local_cases,
            lc = "black",
            label = false,
        )

        # plot the last date used for conditioning the simulations
        vline!(
            fig,
            subplot = sp,
            [Dates.Date(file_date) - Dates.Day(truncation_days)],
            lc = :gray,
            ls = :dash,
            label = false,
        )

        # adjust limits to make plots a little nicer to look at
        xlims!(fig, subplot = sp, onset_dates_lims...)
        if zoom
            ylims!(fig, subplot = sp, 0, 1.25 * maximum(local_cases))
        end

    end

    dir_name = joinpath("figs", "case_forecasts", forecast_origin)
    if !ispath(dir_name)
        mkpath(dir_name)
    end

    if zoom
        savefig(
            fig,
            dir_name *
            "/UoA_forecast_" *
            forecast_origin *
            "_" *
            "zoomed_" *
            confidence_level *
            "_intervals.pdf"
        )
    else
        savefig(
            fig,
            dir_name *
            "/UoA_forecast_" *
            forecast_origin *
            "_" *
            confidence_level *
            "_intervals.pdf"
        )
    end

    return nothing

end

function plot_all_forecast_intervals(
    file_date,
    states,
    dates,
    local_case_dict;
    truncation_days = 7,
)
    """
    Simple wrapper function to plot the forecasts with various zoom and confidence levels.
    """

    plot_all_forecasts(
        file_date,
        states,
        dates,
        local_case_dict,
        confidence_level = "both",
    )
    plot_all_forecasts(
        file_date,
        states,
        dates,
        local_case_dict,
        zoom = true,
        confidence_level = "both",
    )
    plot_all_forecasts(
        file_date,
        states,
        dates,
        local_case_dict,
        confidence_level = "50",
    )
    plot_all_forecasts(
        file_date,
        states,
        dates,
        local_case_dict,
        confidence_level = "95",
    )

    forecast_origin = string(Dates.Date(file_date) - Dates.Day(truncation_days))
    dir_name = "figs/case_forecasts/" * forecast_origin * "/"
    file_name_tmp = "UoA_forecast_"

    pdf_filenames = [
        dir_name * file_name_tmp * forecast_origin * "_zoomed_both_intervals.pdf",
        dir_name * file_name_tmp * forecast_origin * "_both_intervals.pdf",
        dir_name * file_name_tmp * forecast_origin * "_50_intervals.pdf",
    ]
    # merge the pdfs and delete the files
    merge_pdfs(
        pdf_filenames,
        dir_name * file_name_tmp * forecast_origin * ".pdf",
        cleanup=true,
    )

    return nothing

end


text:
using Random
import NonUniformRandomVariateGeneration: sampleBeta, samplePoisson

struct BPθ
   λ::Float64
   a::Float64
   b::Float64
end

const BPσ = 5

function makeBPModel(θ::BPθ, ys::Vector{Float64} = Vector{Float64}(undef, 0))
  n::Int64 = length(ys)
  invσ²over2::Float64 = 0.5/BPσ^2
  logncG::Float64 = -0.5 * log(2 * π * BPσ^2)
  @inline function G(p::Int64, px::Int64Particle)
    @inbounds v::Float64 = px.x - ys[p]
    lG::Float64 = logncG - v * invσ²over2 * v
    @inbounds return exp(lG)
  end
  @inline function μ!(p::Int64Particle, rng::RNG) where RNG <: AbstractRNG
    S::Float64 = sampleBeta(θ.a, θ.b, rng)
    p.x = samplePoisson(θ.λ*S, rng)
  end
  @inline function simulate(p::Int64Particle, rng::RNG) where RNG <: AbstractRNG
    return p.x + BPσ*randn(rng)
  end
  return Model(μ!, G, n, simulate, Int64Particle, Float64)
end