text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

import Compat.ASCIIString

using ExprSearch, DerivationTrees, LinearDerivTrees
using RLESUtils, Confusion, TreeIterators, TreeUtils
using DecisionTrees

function logsystem()
    logsys = LogSystem()

    register_log!(logsys, "parameters", ["parameter", "value"], [ASCIIString, Any])
    register_log!(logsys, "computeinfo", ["parameter", "value"],
        [ASCIIString, ASCIIString])
    register_log!(logsys, "elapsed_cpu_s", ["decision_id", "CPUtime_s"],
        [Int64, Float64])
    register_log!(logsys, "split_result", ["decision_id", "expr", "fitness"],
        [Int64, ASCIIString, Float64])
    register_log!(logsys, "members", ["decision_id", "members"], [Int64, ASCIIString])
    register_log!(logsys, "classifier_metrics", ["truepos", "trueneg", "falsepos",
        "falseneg", "precision", "recall", "accuracy", "f1_score"],
        [Int64, Int64, Int64, Int64, Float64, Float64, Float64, Float64], "result",
        x->begin
            result = x[1]
            members = get_members(result.decision_tree)
            pred = map(x->classify(result.dt_params, result.decision_tree, x), members)
            pred = pred .== 1
            truth = get_truth(members, result.dt_params.userargs)
            truth = truth .== 1
            conf_mat = ConfusionMat(pred, truth)
            [conf_mat.truepos, conf_mat.trueneg, conf_mat.falsepos, conf_mat.falseneg,
                precision(conf_mat), recall(conf_mat), accuracy(conf_mat),
                f1_score(conf_mat)]
        end)
    register_log!(logsys, "interpretability_metrics", ["num_rules", "avg_rule_length",
        "num_nodes", "num_leaf", "avg_deriv_tree_num_nodes", "avg_deriv_tree_num_leafs"],
        [Int64, Float64, Int64, Int64, Float64, Float64], "result",
        x->begin
            result = x[1]
            #decision tree
            num_rules = DecisionTrees.count_nonleafs(result.decision_tree)
            num_leaf = DecisionTrees.count_leafs(result.decision_tree)
            sum_rule_length = traverse(x->begin
                    x.split_rule != nothing ? length(string(x.split_rule.expr)) : 0
                end,
                +, result.decision_tree.root)
            avg_rule_length = sum_rule_length / (num_rules + num_leaf)

            #each derivation tree
            num_derivnodes = Int64[]
            num_derivleafs = Int64[]
            for node in tree_iter(result.decision_tree.root)
                if node.split_rule != nothing
                    push!(num_derivnodes, length(get_derivtree(node.split_rule)))
                    push!(num_derivleafs, DerivationTrees.count_leafs(
                        get_derivtree(node.split_rule)))
                end
            end
            avg_deriv_num_nodes = mean(num_derivnodes)
            avg_deriv_num_leafs = mean(num_derivleafs)
            [num_rules, avg_rule_length, (num_rules + num_leaf), num_leaf, avg_deriv_num_nodes,
                avg_deriv_num_leafs]
        end)

    register_log!(logsys, "verbose1", ["msg"], [ASCIIString])
    register_log!(logsys, "split_result_print", ["msg"], [ASCIIString], "split_result_raw",
        x->begin
            decision_id, result = x
            fitness = result.fitness
            expr = string(get_expr(result))
            expr_short = take(expr, 50) |> join
            ["decision_id=$(decision_id), fitness=$(signif(fitness, 4)), expr=$(expr_short)"]
        end)

    logsys
end

max_stars_repo_name: sisl/ExprSearch

max_issues_repo_name: sisl/ExprSearch


text:
module TestSphere
    using Evolutionary
    using Base.Test

    function test_result(result::Vector, fitness::Float64, N::Int, tol::Float64)
        @test length(result) == N
        @test_approx_eq_eps fitness 0.0 tol
    end

    # Objective function
    sphere{T <: Vector}(x::T) = sum(x.*x)

    # Parameters
    N = 30
    P = 25
    terminate(strategy) = strategy[:σ] < 1e-10
    initial = ones(N)

    # Testing: (μ/μ_I, λ)-σ-Self-Adaptation-ES
    # with isotropic mutation operator y' := y + σ(N_1(0, 1), ..., N_N(0, 1))
    result, fitness, cnt = es(sphere, N;
        initStrategy = strategy(σ = 1.0, τ = 1/sqrt(2*N)),
        recombination = average, srecombination = averageSigma1,
        mutation = isotropic, smutation = isotropicSigma,
        termination = terminate, selection=:comma,
        μ = 3, λ = P, iterations = 1000)
    println("(3/3,$(P))-σ-SA-ES => F: $(fitness), C: $(cnt)")
    test_result(result, fitness, N, 1e-5)

    # Testing: GA
    result, fitness, cnt =
        ga( sphere, N;
            populationSize = 4P,
            mutationRate = 0.05,
            ɛ = 0.1,
            selection = sus,
            crossover = intermediate(0.25),
            mutation = domainrange(fill(0.5,N)),
            tol = 1e-5, tolIter = 15)
    println("GA(pop=$(P),xover=0.8,μ=0.1) => F: $(fitness), C: $(cnt)")
    test_result(result, fitness, N, 1e-2)

end

ext: jl


text:
"""

# Converts Euro into another currency.

    cconvert(x::Number, currency::String)

## Arguments

    - `x` : a Number
    - `currency` : A currency out of ["CZK", "PLN", "MXN", "TRY", "ISK", "USD", "RUB", "GBP", "JPY", "THB", "CNY", "KRW", "SEK", "BRL", "BGN", "HUF", "SGD", "CHF", "IDR", "NZD", "MYR", "ZAR", "CAD", "PHP", "HKD", "AUD", "INR", "HRK", "NOK", "DKK", "RON", "ILS"]

## Examples

```jldoctest
julia> cconvert(100, "USD")
julia> cconvert(10, "CAD")
julia> cconvert(1, "ASD")
julia>
julia> x = rand(100)
julia> cconvert.(x, "CAD")
```
"""
function cconvert(x::Number, symbol::String = "")

  url = "https://api.exchangeratesapi.io/"
  url *= "latest/"
  response = HTTP.get(url, cookies = true)
  body = JSON.parse(String(response.body))

  if symbol ∉ keys(body["rates"])
      error("Currency '$symbol' is not yet supported")
  end

  x *= body["rates"][symbol]
  return x
end


text:
mutable struct TimeStepping
    dt::Float64
    t_max::Float64
    t::Float64
    nstep::Int
    cfl_limit::Float64
    dt_max::Float64
    dt_io::Float64
end

function TimeStepping(namelist)
    dt = parse_namelist(namelist, "time_stepping", "dt_min"; default = 1.0)
    t_max = parse_namelist(namelist, "time_stepping", "t_max"; default = 7200.0)
    cfl_limit = parse_namelist(namelist, "time_stepping", "cfl_limit"; default = 0.5)
    dt_max = parse_namelist(namelist, "time_stepping", "dt_max"; default = 10.0)

    # set time
    t = 0.0
    dt_io = 0.0
    nstep = 0

    return TimeStepping(dt, t_max, t, nstep, cfl_limit, dt_max, dt_io)
end

function update(self)
    self.t += self.dt
    self.nstep += 1
    return
end


text:
module OffsetArrays

using Base: tail, @propagate_inbounds
@static if !isdefined(Base, :IdentityUnitRange)
    const IdentityUnitRange = Base.Slice
else
    using Base: IdentityUnitRange
end

export OffsetArray, OffsetMatrix, OffsetVector

include("axes.jl")
include("utils.jl")
include("origin.jl")

# Technically we know the length of CartesianIndices but we need to convert it first, so here we
# don't put it in OffsetAxisKnownLength.
const OffsetAxisKnownLength = Union{Integer, AbstractUnitRange}
const OffsetAxis = Union{OffsetAxisKnownLength, Colon}
const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}

## OffsetArray
"""
    OffsetArray(A, indices...)

Return an `AbstractArray` that shares element type and size with the first argument but
uses the supplied `indices` to infer its axes. If all the indices are `AbstractUnitRange`s then
these are directly used as the axis span along each dimension. Refer to the examples below for other
permissible types.

Alternatively it's possible to specify the coordinates of one corner of the array
and have the axes be computed automatically from the size of `A`.
This constructor makes it convenient to shift to
an arbitrary starting index along each axis, for example to a zero-based indexing scheme followed by
arrays in languages such as C and Python.
See [`Origin`](@ref) and the examples below for this usage.

# Example: offsets

There are two types of `indices`: integers and ranges-like types.

Integers are recognized as offsets, where `0` means no offsets are applied:

```jldoctest; setup=:(using OffsetArrays)
julia> A = OffsetArray(reshape(1:6, 2, 3), -1, -2)
2×3 OffsetArray(reshape(::UnitRange{$Int}, 2, 3), 0:1, -1:1) with eltype $Int with indices 0:1×-1:1:
 1  3  5
 2  4  6

julia> A[0, 1]
5
```

Examples of range-like types are: `UnitRange` (e.g, `-1:2`), `CartesianIndices`,
and `Colon()` (or concisely `:`). A `UnitRange` specifies the axis span along one particular dimension,
`CartesianIndices` specify the axis spans along multiple dimensions, and a `Colon` is a placeholder
that specifies that the `OffsetArray` shares its axis with its parent along that dimension.

```jldoctest; setup=:(using OffsetArrays)
julia> OffsetArray(reshape(1:6, 2, 3), 0:1, -1:1)
2×3 OffsetArray(reshape(::UnitRange{$Int}, 2, 3), 0:1, -1:1) with eltype $Int with indices 0:1×-1:1:
 1  3  5
 2  4  6

julia> OffsetArray(reshape(1:6, 2, 3), :, -1:1) # : as a placeholder to indicate that no offset is to be applied to the first dimension
2×3 OffsetArray(reshape(::UnitRange{$Int}, 2, 3), 1:2, -1:1) with eltype $Int with indices 1:2×-1:1:
 1  3  5
 2  4  6
```

Use `CartesianIndices` to specify the coordinates of two diagonally opposite corners:

```jldoctest; setup=:(using OffsetArrays)
julia> OffsetArray(reshape(1:6, 2, 3), CartesianIndex(0, -1):CartesianIndex(1, 1))
2×3 OffsetArray(reshape(::UnitRange{$Int}, 2, 3), 0:1, -1:1) with eltype $Int with indices 0:1×-1:1:
 1  3  5
 2  4  6
```

Integers and range-like types may not be combined in the same call:

```julia
julia> OffsetArray(reshape(1:6, 2, 3), 0, -1:1)
ERROR: [...]
```

# Example: origin

[`OffsetArrays.Origin`](@ref) may be used to specify the origin of the OffsetArray. The term origin here
refers to the corner with the lowest values of coordinates, such as the left edge for an `AbstractVector`,
the bottom left corner for an `AbstractMatrix` and so on. The coordinates of the origin sets the starting
index of the array along each dimension.

```jldoctest; setup=:(using OffsetArrays)
julia> a = [1 2; 3 4];

julia> OffsetArray(a, OffsetArrays.Origin(0, 1))
2×2 OffsetArray(::$(Array{Int,2}), 0:1, 1:2) with eltype $Int with indices 0:1×1:2:
 1  2
 3  4

julia> OffsetArray(a, OffsetArrays.Origin(0)) # set the origin to zero along each dimension
2×2 OffsetArray(::$(Array{Int, 2}), 0:1, 0:1) with eltype $Int with indices 0:1×0:1:
 1  2
 3  4
```

"""
struct OffsetArray{T,N,AA<:AbstractArray{T,N}} <: AbstractArray{T,N}
    parent::AA
    offsets::NTuple{N,Int}
    function OffsetArray{T, N, AA}(parent::AA, offsets::NTuple{N, Int}) where {T, N, AA<:AbstractArray{T,N}}
        # allocation of `map` on tuple is optimized away
        map(overflow_check, axes(parent), offsets)
        new{T, N, AA}(parent, offsets)
    end
end

function OffsetArray{T, N, AA}(parent::AA, offsets::NTuple{N, <:Integer}) where {T, N, AA<:AbstractArray{T,N}}
    OffsetArray{T, N, AA}(parent, map(x -> convert(Int, x)::Int, offsets))
end

"""
    OffsetVector(v, index)

Type alias and convenience constructor for one-dimensional [`OffsetArray`](@ref)s.
"""
const OffsetVector{T,AA<:AbstractVector{T}} = OffsetArray{T,1,AA}

"""
    OffsetMatrix(A, index1, index2)

Type alias and convenience constructor for two-dimensional [`OffsetArray`](@ref)s.
"""
const OffsetMatrix{T,AA<:AbstractMatrix{T}} = OffsetArray{T,2,AA}

function overflow_check(r, offset::T) where T
    # This gives some performance boost https://github.com/JuliaLang/julia/issues/33273
    throw_upper_overflow_error() = throw(OverflowError("Boundary overflow detected: offset should be <= $(typemax(T) - last(r)) for offsets of type $T, received $offset"))
    throw_lower_overflow_error() = throw(OverflowError("Boundary overflow detected: offset should be >= $(typemin(T) - first(r)) for offsets of type $T, received $offset"))

    if offset > 0 && last(r) > typemax(T) - offset
        throw_upper_overflow_error()
    elseif offset < 0 && first(r) < typemin(T) - offset
        throw_lower_overflow_error()
    end
end

# Tuples of integers are treated as offsets
# Empty Tuples are handled here
@inline function OffsetArray(A::AbstractArray, offsets::Tuple{Vararg{Integer}})
    _checkindices(A, offsets, "offsets")
    OffsetArray{eltype(A), ndims(A), typeof(A)}(A, offsets)
end

# These methods are necessary to disallow incompatible dimensions for
# the OffsetVector and the OffsetMatrix constructors
for (FT, ND) in ((:OffsetVector, :1), (:OffsetMatrix, :2))
    @eval @inline function $FT(A::AbstractArray{<:Any,$ND}, offsets::Tuple{Vararg{Integer}})
        _checkindices(A, offsets, "offsets")
        OffsetArray{eltype(A), $ND, typeof(A)}(A, offsets)
    end
    FTstr = string(FT)
    @eval @inline function $FT(A::AbstractArray, offsets::Tuple{Vararg{Integer}})
        throw(ArgumentError($FTstr*" requires a "*string($ND)*"D array"))
    end
end

## OffsetArray constructors
for FT in (:OffsetArray, :OffsetVector, :OffsetMatrix)
    # Nested OffsetArrays may strip off the wrapper and collate the offsets
    @eval @inline function $FT(A::OffsetArray, offsets::Tuple{Vararg{Integer}})
        _checkindices(A, offsets, "offsets")
        $FT(parent(A), map(+, A.offsets, offsets))
    end

    # In general, indices get converted to AbstractUnitRanges.
    # CartesianIndices{N} get converted to N ranges
    @eval @inline function $FT(A::AbstractArray, inds::Tuple{Any,Vararg{Any}})
        $FT(A, _toAbstractUnitRanges(to_indices(A, axes(A), inds)))
    end

    # convert ranges to offsets
    @eval @inline function $FT(A::AbstractArray, inds::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}})
        _checkindices(A, inds, "indices")
        # Performance gain by wrapping the error in a function: see https://github.com/JuliaLang/julia/issues/37558
        throw_dimerr(lA, lI) = throw(DimensionMismatch("supplied axes do not agree with the size of the array (got size $lA for the array and $lI for the indices"))
        lA = size(A)
        lI = map(length, inds)
        lA == lI || throw_dimerr(lA, lI)
        $FT(A, map(_offset, axes(A), inds))
    end

    @eval @inline $FT(A::AbstractArray, inds::Vararg) = $FT(A, inds)

    @eval @inline $FT(A::AbstractArray, origin::Origin) = $FT(A, origin(A))
end

# array initialization
@inline function OffsetArray{T,N}(init::ArrayInitializer, inds::Tuple{Vararg{OffsetAxisKnownLength}}) where {T,N}
    _checkindices(N, inds, "indices")
    AA = Array{T,N}(init, map(_indexlength, inds))
    OffsetArray{T, N, typeof(AA)}(AA, map(_indexoffset, inds))
end
@inline function OffsetArray{T, N}(init::ArrayInitializer, inds::Tuple) where {T, N}
    OffsetArray{T, N}(init, _toAbstractUnitRanges(inds))
end
@inline OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg) where {T,N} = OffsetArray{T,N}(init, inds)

@inline OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N} = OffsetArray{T,N}(init, inds)
@inline function OffsetArray{T}(init::ArrayInitializer, inds::Tuple) where {T}
    OffsetArray{T}(init, _toAbstractUnitRanges(inds))
end
@inline OffsetArray{T}(init::ArrayInitializer, inds::Vararg) where {T} = OffsetArray{T}(init, inds)

Base.IndexStyle(::Type{OA}) where {OA<:OffsetArray} = IndexStyle(parenttype(OA))
parenttype(::Type{OffsetArray{T,N,AA}}) where {T,N,AA} = AA
parenttype(A::OffsetArray) = parenttype(typeof(A))

Base.parent(A::OffsetArray) = A.parent

# TODO: Ideally we would delegate to the parent's broadcasting implementation, but that
#       is currently broken in sufficiently many implementation, namely RecursiveArrayTools, DistributedArrays
#       and StaticArrays, that it will take concentrated effort to get this working across the ecosystem.
#       The goal would be to have `OffsetArray(CuArray) .+ 1 == OffsetArray{CuArray}`.
# Base.Broadcast.BroadcastStyle(::Type{<:OffsetArray{<:Any, <:Any, AA}}) where AA = Base.Broadcast.BroadcastStyle(AA)

Base.eachindex(::IndexCartesian, A::OffsetArray) = CartesianIndices(axes(A))
Base.eachindex(::IndexLinear, A::OffsetVector)   = axes(A, 1)

@inline Base.size(A::OffsetArray) = size(parent(A))
@inline Base.size(A::OffsetArray, d) = size(parent(A), d)

@inline Base.axes(A::OffsetArray) = map(IdOffsetRange, axes(parent(A)), A.offsets)
@inline Base.axes(A::OffsetArray, d) = d <= ndims(A) ? IdOffsetRange(axes(parent(A), d), A.offsets[d]) : IdOffsetRange(axes(parent(A), d))
@inline Base.axes1(A::OffsetArray{T,0}) where {T} = IdOffsetRange(axes(parent(A), 1))  # we only need to specialize this one

Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T =
    similar(parent(A), T, dims)
function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T
    B = similar(A, T, map(_indexlength, inds))
    return OffsetArray(B, map(_offset, axes(B), inds))
end

# reshape accepts a single colon
Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
    AR = reshape(A, map(_indexlength, inds))
    return OffsetArray(AR, map(_offset, axes(AR), inds))
end

# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
    OffsetArray(reshape(parent(A), map(_indexlength, inds)), map(_indexoffset, inds))
# And for non-offset axes, we can just return a reshape of the parent directly
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
Base.reshape(A::OffsetArray, inds::Dims) = reshape(parent(A), inds)
Base.reshape(A::OffsetArray, ::Colon) = reshape(parent(A), Colon())
Base.reshape(A::OffsetVector, ::Colon) = A
Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(parent(A), inds)
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = reshape(parent(A), inds)

function Base.similar(::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where {T<:AbstractArray}
    P = T(undef, map(_indexlength, shape))
    OffsetArray(P, map(_offset, axes(P), shape))
end

Base.fill(v, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
    fill!(similar(Array{typeof(v)}, inds), v)
Base.zeros(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =
    fill!(similar(Array{T}, inds), zero(T))
Base.ones(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =
    fill!(similar(Array{T}, inds), one(T))
Base.trues(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
    fill!(similar(BitArray, inds), true)
Base.falses(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
    fill!(similar(BitArray, inds), false)

## Indexing

# Note this gets the index of the parent *array*, not the index of the parent *range*
# Here's how one can think about this:
#   Δi = i - first(r)
#   i′ = first(r.parent) + Δi
# and one obtains the result below.
parentindex(r::IdOffsetRange, i) = i - r.offset

@inline function Base.getindex(A::OffsetArray{T,N}, I::Vararg{Int,N}) where {T,N}
    @boundscheck checkbounds(A, I...)
    J = map(parentindex, axes(A), I)
    @inbounds parent(A)[J...]
end

@inline function Base.getindex(A::OffsetVector, i::Int)
    @boundscheck checkbounds(A, i)
    @inbounds parent(A)[parentindex(Base.axes1(A), i)]
end
@propagate_inbounds Base.getindex(A::OffsetArray, i::Int)  = parent(A)[i]

@inline function Base.setindex!(A::OffsetArray{T,N}, val, I::Vararg{Int,N}) where {T,N}
    @boundscheck checkbounds(A, I...)
    J = map(parentindex, axes(A), I)
    @inbounds parent(A)[J...] = val
    A
end

@inline function Base.setindex!(A::OffsetVector, val, i::Int)
    @boundscheck checkbounds(A, i)
    @inbounds parent(A)[parentindex(Base.axes1(A), i)] = val
    A
end
@propagate_inbounds function Base.setindex!(A::OffsetArray, val, i::Int)
    parent(A)[i] = val
    A
end

Base.strides(A::OffsetArray) = strides(parent(A))
Base.elsize(::Type{OffsetArray{T,N,A}}) where {T,N,A} = Base.elsize(A)
@inline Base.unsafe_convert(::Type{Ptr{T}}, A::OffsetArray{T}) where {T} = Base.unsafe_convert(Ptr{T}, parent(A))

# For fast broadcasting: ref https://discourse.julialang.org/t/why-is-there-a-performance-hit-on-broadcasting-with-offsetarrays/32194
Base.dataids(A::OffsetArray) = Base.dataids(parent(A))
Broadcast.broadcast_unalias(dest::OffsetArray, src::OffsetArray) = parent(dest) === parent(src) ? src : Broadcast.unalias(dest, src)

### Special handling for AbstractRange

const OffsetRange{T} = OffsetArray{T,1,<:AbstractRange{T}}
const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer

Base.step(a::OffsetRange) = step(parent(a))

@propagate_inbounds Base.getindex(a::OffsetRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
@propagate_inbounds function Base.getindex(a::OffsetRange, r::IdOffsetRange)
    OffsetArray(a.parent[r.parent .+ (r.offset - a.offsets[1])], r.offset)
end
@propagate_inbounds Base.getindex(r::OffsetRange, s::IIUR) =
    OffsetArray(r[s.indices], s)
@propagate_inbounds Base.getindex(a::OffsetRange, r::AbstractRange) = a.parent[r .- a.offsets[1]]
@propagate_inbounds Base.getindex(a::AbstractRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)

for OR in [:IIUR, :IdOffsetRange]
    for R in [:StepRange, :StepRangeLen, :LinRange, :UnitRange]
        @eval @propagate_inbounds Base.getindex(r::$R, s::$OR) = OffsetArray(r[_unwrap(s)], axes(s))
    end

    # this method is needed for ambiguity resolution
    @eval @propagate_inbounds Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T =
    OffsetArray(r[_unwrap(s)], axes(s))
end

#= Integer UnitRanges may return an appropriate AbstractUnitRange{<:Integer}, as the result may be used in indexing, and
indexing is faster with ranges =#
@propagate_inbounds Base.getindex(r::UnitRange{<:Integer}, s::IdOffsetRange) = IdOffsetRange(r[_unwrap(s)] .- s.offset, s.offset)
@propagate_inbounds Base.getindex(r::UnitRange{<:Integer}, s::IIUR) = IdentityUnitRange(r[_unwrap(s)])

function Base.show(io::IO, r::OffsetRange)
    show(io, r.parent)
    o = r.offsets[1]
    print(io, " with indices ", o+1:o+length(r))
end
Base.show(io::IO, ::MIME"text/plain", r::OffsetRange) = show(io, r)

### Some mutating functions defined only for OffsetVector ###

Base.resize!(A::OffsetVector, nl::Integer) = (resize!(A.parent, nl); A)
Base.push!(A::OffsetVector, x...) = (push!(A.parent, x...); A)
Base.pop!(A::OffsetVector) = pop!(A.parent)
Base.append!(A::OffsetVector, items) = (append!(A.parent, items); A)
Base.empty!(A::OffsetVector) = (empty!(A.parent); A)

# These functions keep the summary compact
function Base.inds2string(inds::Tuple{Vararg{Union{IdOffsetRange, IdentityUnitRange{<:IdOffsetRange}}}})
    Base.inds2string(map(UnitRange, inds))
end
Base.showindices(io::IO, ind1::IdOffsetRange, inds::IdOffsetRange...) = Base.showindices(io, map(UnitRange, (ind1, inds...))...)

function Base.showarg(io::IO, @nospecialize(a::OffsetArray), toplevel)
    print(io, "OffsetArray(")
    Base.showarg(io, parent(a), false)
    Base.showindices(io, axes(a)...)
    print(io, ')')
    if toplevel
        print(io, " with eltype ", eltype(a))
    end
end

function Base.replace_in_print_matrix(A::OffsetArray{<:Any,2}, i::Integer, j::Integer, s::AbstractString)
    J = map(parentindex, axes(A), (i,j))
    Base.replace_in_print_matrix(parent(A), J..., s)
end
function Base.replace_in_print_matrix(A::OffsetArray{<:Any,1}, i::Integer, j::Integer, s::AbstractString)
    ip = parentindex(axes(A,1), i)
    Base.replace_in_print_matrix(parent(A), ip, j, s)
end

"""
    no_offset_view(A)

Return an `AbstractArray` that shares structure and underlying data with the argument,
but uses 1-based indexing. May just return the argument when applicable.
Not exported.

The default implementation uses `OffsetArrays`, but other types should use something more
specific to remove a level of indirection when applicable.

```jldoctest; setup=:(using OffsetArrays)
julia> A = [1 3 5; 2 4 6];

julia> O = OffsetArray(A, 0:1, -1:1)
2×3 OffsetArray(::$(Matrix{Int}), 0:1, -1:1) with eltype $Int with indices 0:1×-1:1:
 1  3  5
 2  4  6

julia> OffsetArrays.no_offset_view(O)[1,1] = -9
-9

julia> A
2×3 $(Matrix{Int}):
 -9  3  5
  2  4  6
```
"""
function no_offset_view(A::AbstractArray)
    if Base.has_offset_axes(A)
        OffsetArray(A, Origin(1))
    else
        A
    end
end

no_offset_view(A::OffsetArray) = no_offset_view(parent(A))
no_offset_view(a::Array) = a

####
# work around for segfault in searchsorted*
#  https://github.com/JuliaLang/julia/issues/33977
####

function _safe_searchsorted(v::OffsetArray, x, ilo::T, ihi::T, o::Base.Ordering) where T<:Integer
    u = T(1)
    lo = ilo - u
    hi = ihi + u
    @inbounds while lo < hi - u
        m = (lo + hi) ÷ 2
        if Base.lt(o, v[m], x)
            lo = m
        elseif Base.lt(o, x, v[m])
            hi = m
        else
            a = searchsortedfirst(v, x, max(lo,ilo), m, o)
            b = searchsortedlast(v, x, m, min(hi,ihi), o)
            return a : b
        end
    end
    return (lo + 1) : (hi - 1)
end
function _safe_searchsortedfirst(v::OffsetArray, x, lo::T, hi::T, o::Base.Ordering) where T<:Integer
    u = T(1)
    lo = lo - u
    hi = hi + u
    @inbounds while lo < hi - u
        m = (lo + hi) ÷ 2
        if Base.lt(o, v[m], x)
            lo = m
        else
            hi = m
        end
    end
    return hi
end
function _safe_searchsortedlast(v::OffsetArray, x, lo::T, hi::T, o::Base.Ordering) where T<:Integer
    u = T(1)
    lo = lo - u
    hi = hi + u
    @inbounds while lo < hi - u
        m = (lo + hi) ÷ 2
        if Base.lt(o, x, v[m])
            hi = m
        else
            lo = m
        end
    end
    return lo
end

if VERSION ≤ v"1.2"
    # ambiguity warnings in earlier versions
    Base.searchsorted(v::OffsetArray, x, ilo::Int, ihi::Int, o::Base.Ordering) =
        _safe_searchsorted(v, x, ilo, ihi, o)
    Base.searchsortedfirst(v::OffsetArray, x, lo::Int, hi::Int, o::Base.Ordering) =
        _safe_searchsortedfirst(v, x, lo, hi, o)
    Base.searchsortedlast(v::OffsetArray, x, lo::Int, hi::Int, o::Base.Ordering) =
        _safe_searchsortedlast(v, x, lo, hi, o)
end

Base.searchsorted(v::OffsetArray, x, ilo::T, ihi::T, o::Base.Ordering) where T<:Integer =
    _safe_searchsorted(v, x, ilo, ihi, o)
Base.searchsortedfirst(v::OffsetArray, x, lo::T, hi::T, o::Base.Ordering) where T<:Integer =
    _safe_searchsortedfirst(v, x, lo, hi, o)
Base.searchsortedlast(v::OffsetArray, x, lo::T, hi::T, o::Base.Ordering) where T<:Integer =
    _safe_searchsortedlast(v, x, lo, hi, o)

if VERSION < v"1.1.0-DEV.783"
    Base.copyfirst!(dest::OffsetArray, src::OffsetArray) = (maximum!(parent(dest), parent(src)); return dest)
end

##
# Adapt allows for automatic conversion of CPU OffsetArrays to GPU OffsetArrays
##
import Adapt
Adapt.adapt_structure(to, x::OffsetArray) = OffsetArray(Adapt.adapt(to, parent(x)), x.offsets)

if Base.VERSION >= v"1.4.2"
    include("precompile.jl")
    _precompile_()
end

end # module


text:
export sample_blockaded, distance, distance_matrix, Box, BoxPBC, NoisyChain, NoisyChainPBC

## TODO: abstract Lattice geometries and put chains as subtypes
##       lattices need to know at which site to generate a new point
## TODO: orthogonalize PBC types?

abstract type Geometry end

"""
    generate_point(geometry)

Generate a new (random) point in the geometry.
"""
function generate_point end

# default
_euclidean_point(lengths) = rand(length(lengths)) .* lengths

"""
    distance(geometry, p1, p2)

Compute the distance of two points p1 and p2 within the given geometry.
"""
function distance end

## two default implementations

function _euclidean(p1, p2)
    sqrt(sum(@. (p1 - p2)^2 ))
end

function _euclidean_pbc(p1, p2, lengths)
    diffs = _euclidean_pbc.(p1, p2, lengths)
    sqrt(sum(diffs.^2))
end

function _euclidean_pbc(x1::Number, x2::Number, length::Number)
    diff = abs(mod(x1, length) - mod(x2, length))
    diff > length/2 ? length-diff : diff
end

"""
    distance_matrix(geometry, points)

Compute the distance pairwise between all points.
"""
function distance_matrix(geometry::Geometry, points::Vector)
    N = length(points)
    res = zeros(Float64, (N, N))
    for i in 1:N
        for j in i+1:N
            res[i,j] = distance(geometry, points[i], points[j])
        end
    end
    Symmetric(res)
end

distance_matrix(geometry::Geometry, points::AbstractMatrix) = distance_matrix(geometry, collect(eachcol(points)))

"""
    sample_blockaded(geometry, N; blockade_radius, max_iter)

Generate N points with distance blockade_radius (default 1) within the given geometry.
"""
function sample_blockaded(geometry, N; blockade_radius=1, max_iter=1000)
    res = [generate_point(geometry)]
    sizehint!(res, N)
    while length(res) < N
        push!(res, _find_new_point(geometry, res, blockade_radius, max_iter))
    end
    res
end

function _find_new_point(geometry, points, blockade_radius, max_iter)
    for _ in 1:max_iter
        p = generate_point(geometry)
        if all(map(p2 -> distance(geometry, p, p2) > blockade_radius, points))
            return p
        end
    end
    error("Iteration limit reached! Could not find a $(length(points)+1)th point in $geometry geometry with $max_iter iterations.")
end

"""
    Box{T}

Geometry representing a simple box with given lengths.

# Fields
- `lengths::Vector{T}`: stores the lengths of the individual dimensions
"""
struct Box{T} <: Geometry
    lengths::Vector{T}
end

generate_point(box::Box) = _euclidean_point(box.lengths)
distance(::Box, p1, p2) = _euclidean(p1, p2)

"""
    BoxPBC{T}

Geometry representing a box with given lengths respecting periodic boundary conditions.

# Fields
- `lengths::Vector{T}`: stores the lengths of the individual dimensions
"""
struct BoxPBC{T} <: Geometry
    lengths::Vector{T}
end

distance(box::BoxPBC, p1, p2) = _euclidean_pbc(p1, p2, box.lengths)
generate_point(box::BoxPBC) = _euclidean_point(box.lengths)

"""
    NoisyChain

Represents a 1D lattice with given spacing and length.
Each position is subjected to noise drawn from a box distribution on [-σ, σ].

# Fields
- `L::Int64`: number of sites in the chain
- `spacing::Float64`: (mean) distance between two sites
- `σ::Float64`: width of box distribution for position noise
"""
struct NoisyChain <: Geometry
    L::Int64
    spacing::Float64
    σ::Float64
end

distance(::NoisyChain, p1, p2) = _euclidean(p1, p2)

"""
    NoisyChainPBC

Represents a 1D periodic lattice with given spacing and length.
Each position is subjected to noise drawn from a box distribution on [-σ, σ].

# Fields
- `L::Int64`: number of sites in the chain
- `spacing::Float64`: (mean) distance between two sites
- `σ::Float64`: width of box distribution for position noise
"""
struct NoisyChainPBC <: Geometry
    L::Int64
    spacing::Float64
    σ::Float64
end

distance(chain::NoisyChainPBC, p1, p2) = _euclidean_pbc(p1, p2, chain.L * chain.spacing)

## chain types need to now at which site to generate the points

generate_point(chain::Union{NoisyChain, NoisyChainPBC}, site=1) = chain.σ * (2*rand()-1) + site*chain.spacing

function _find_new_point(chain::Union{NoisyChain, NoisyChainPBC}, points, blockade_radius, max_iter)
    site = length(points) + 1
    for _ in 1:max_iter
        p = generate_point(chain, site)
        if all(map(p2 -> distance(chain, p, p2) > blockade_radius, points))
            return p
        end
    end
    error("Iteration limit reached! Could not find a $(length(points)+1)th point in $chain geometry with $max_iter iterations.")
end


text:
module QueryStrategies

abstract type QueryStrategy end

abstract type DataBasedQs <: QueryStrategy end
abstract type ModelBasedQs <: QueryStrategy end
abstract type HybridQs <: QueryStrategy end

using MLKernels
using NearestNeighbors
using Statistics
using LinearAlgebra
using InteractiveUtils
using PyCall
using SVDD

const gaussian_kde = PyNULL()

function __init__()
    copy!(gaussian_kde, pyimport_conda("scipy.stats", "scipy")[:gaussian_kde])
end

include("qs_utils.jl")

include("TestQs.jl")
include("RandomQs.jl")
include("RandomOutlierQs.jl")
include("MinimumMarginQs.jl")
include("ExpectedMinimumMarginQs.jl")
include("ExpectedMaximumEntropyQs.jl")
include("MinimumLossQs.jl")
include("HighConfidenceQs.jl")
include("DecisionBoundaryQs.jl")
include("NeighborhoodBasedQs.jl")
include("BoundaryNeighborCombinationQs.jl")

export
    QueryStrategy,
    DataBasedQs,
    ModelBasedQs, HybridQs,

    # data-based query strategies
    TestQs, RandomQs, MinimumMarginQs, ExpectedMinimumMarginQs, ExpectedMaximumEntropyQs,
    MinimumLossQs,
    # model-based query strategies
    RandomOutlierQs, HighConfidenceQs, DecisionBoundaryQs,
    # hybrid query strategies
    NeighborhoodBasedQs, BoundaryNeighborCombinationQs,

    qs_score,
    initialize_qs,
    filter_array,
    multi_kde, KDEException, MissingLabelTypeException

end


text:
module GenParticleFilters

using Gen, Distributions, Parameters
using Gen: ParticleFilterState

export ParticleFilterState, ParticleFilterSubState, ParticleFilterView

include("view.jl")
include("utils.jl")
include("initialize.jl")
include("translate.jl")
include("update.jl")
include("resample.jl")
include("rejuvenate.jl")

end


text:
# -*- coding: utf-8 -*-
# # Benchmarking `MeshArrays` performance
#
# - Started in 2018/09 (`gaelforget`)
# - Reformated in 2019/08 (`gaelforget`)
# - Added to repo in 2020/01 (`gaelforget`)

# **2021/02/26 using julia 1.5**
#
# _note the speedup in `GridLoad` & `LatitudeCircles` & `demo3` benchmarks_
#
# ```
# gradient          274.532 μs (231 allocations: 3.69 MiB)
# smooth             24.581 ms (10845 allocations: 238.25 MiB)
#
# GridLoad               80.133 ms (8285 allocations: 663.26 MiB)
# read_bin              151.670 μs (71 allocations: 2.01 MiB)
# LatitudeCircles       647.875 ms (254111 allocations: 3.17 GiB)
# ThroughFlow            43.826 ms (751490 allocations: 29.59 MiB)
# demo3                 688.477 ms (1005446 allocations: 3.20 GiB)
# ```
#
# **2019/08/08 using julia 1.1**
#
# ```
# gradient          299.661 μs (318 allocations: 3.69 MiB)
# smooth             34.006 ms (15572 allocations: 238.32 MiB)
#
# GridLoad          337.105 ms (1960 allocations: 530.35 MiB)
# read_bin          191.141 μs (61 allocations: 1.61 MiB)
# LatitudeCircles     2.904 s (1511608 allocations: 4.42 GiB)
# ThroughFlow loop   65.804 ms (752743 allocations: 24.46 MiB)
# demo3               2.581 s (2264189 allocations: 4.44 GiB)
# ```

using BenchmarkTools
using MeshArrays
p=dirname(pathof(MeshArrays))
include(joinpath(p,"../examples/Demos.jl"))

cd("/Users/gforget/mywork/data/")

# ## Benchmark 1
#
# This uses `demo2 / smooth` as done in `2018/09` to test exchanges and array operations.

γ,Γ=GridOfOnes("CubeSphere",6,100);
(Rini,Rend,DXCsm,DYCsm)=demo2(Γ);
@btime (dFLDdx, dFLDdy)=gradient(Rini,Γ);
@btime Rend=smooth(Rini,DXCsm,DYCsm,Γ);

# **2019/08/08**
# ```
# 299.661 μs (318 allocations: 3.69 MiB)
#  34.006 ms (15572 allocations: 238.32 MiB)
# ```

# ## Benchmark 2
#
# This uses `demo3` to test indexing type operations.

γ=GridSpec("LatLonCap","GRID_LLC90/")
Γ=GridLoad(γ)
TrspX=γ.read(γ.path*"TrspX.bin",MeshArray(γ,Float32))
TrspY=γ.read(γ.path*"TrspY.bin",MeshArray(γ,Float32))
TauX=γ.read(γ.path*"TauX.bin",MeshArray(γ,Float32))
TauY=γ.read(γ.path*"TauY.bin",MeshArray(γ,Float32))
SSH=γ.read(γ.path*"SSH.bin",MeshArray(γ,Float32))
(UV, LC, Tr)=demo3(TrspX,TrspY,Γ);

@btime Γ=GridLoad(γ)
@btime γ.read(γ.path*"TrspX.bin",MeshArray(γ,Float32))
@btime LC=LatitudeCircles(-89.0:89.0,Γ)
@btime for i=1:length(LC); ThroughFlow(UV,LC[i],Γ); end
@btime (UV, LC, Tr)=demo3(TrspX,TrspY,Γ);

# **2019/08/08**
# ```
# 257.936 ms (1960 allocations: 530.35 MiB)
# 177.847 μs (61 allocations: 1.61 MiB)
#   2.297 s  (1511608 allocations: 4.42 GiB)
#  60.967 ms (752743 allocations: 24.46 MiB)
#   2.305 s  (2264189 allocations: 4.44 GiB)
# ```


text:
import Base: size, axpy!
import Base.LinAlg.scale!
import Base.BLAS: gemm!

@compat abstract type AbstractGLRM end

export AbstractGLRM, GLRM, getindex, size, scale_regularizer!

@compat const ObsArray = @compat(Union{Array{Array{Int,1},1}, Array{UnitRange{Int},1}})

### GLRM TYPE
type GLRM<:AbstractGLRM
    A                            # The data table
    losses::Array{Loss,1}        # array of loss functions
    rx::Array{Regularizer,1}               # Array of regularizers to be applied to each column of X
    ry::Array{Regularizer,1}     # Array of regularizers to be applied to each column of Y
    k::Int                       # Desired rank
    observed_features::ObsArray  # for each example, an array telling which features were observed
    observed_examples::ObsArray  # for each feature, an array telling in which examples the feature was observed
    X::AbstractArray{Float64,2}  # Representation of data in low-rank space. A ≈ X'Y
    Y::AbstractArray{Float64,2}  # Representation of features in low-rank space. A ≈ X'Y
end

# Initialize with nothing; could be useful for copying
# GLRM{L<:Loss, R<:Regularizer}() = GLRM([],Loss[],ZeroReg(),Regularizer[],0,UnitRange{Int}[],UnitRange{Int}[],Array(Float64,(0,0)),Array(Float64,(0,0)))

# usage notes:
# * providing argument `obs` overwrites arguments `observed_features` and `observed_examples`
# * offset and scale are *false* by default to avoid unexpected behavior
# * convenience methods for calling are defined in utilities/conveniencemethods.jl
function GLRM(A, losses::Array, rx::Array, ry::Array, k::Int;
# the following tighter definition fails when you form an array of a tighter subtype than the abstract type, eg Array{QuadLoss,1}
# function GLRM(A::AbstractArray, losses::Array{Loss,1}, rx::Array{Regularizer,1}, ry::Array{Regularizer,1}, k::Int;
              X = randn(k,size(A,1)), Y = randn(k,embedding_dim(losses)),
              obs = nothing,                                    # [(i₁,j₁), (i₂,j₂), ... (iₒ,jₒ)]
              observed_features = fill(1:size(A,2), size(A,1)), # [1:n, 1:n, ... 1:n] m times
              observed_examples = fill(1:size(A,1), size(A,2)), # [1:m, 1:m, ... 1:m] n times
              offset = false, scale = false,
              checknan = true, sparse_na = true)
    # Check dimensions of the arguments
    m,n = size(A)
    if length(losses)!=n error("There must be as many losses as there are columns in the data matrix") end
    if length(rx)!=m error("There must be either one X regularizer or as many X regularizers as there are rows in the data matrix") end
    if length(ry)!=n error("There must be either one Y regularizer or as many Y regularizers as there are columns in the data matrix") end
    if size(X)!=(k,m) error("X must be of size (k,m) where m is the number of rows in the data matrix. This is the transpose of the standard notation used in the paper, but it makes for better memory management. \nsize(X) = $(size(X)), size(A) = $(size(A)), k = $k") end
    if size(Y)!=(k,sum(map(embedding_dim, losses))) error("Y must be of size (k,d) where d is the sum of the embedding dimensions of all the losses. \n(1 for real-valued losses, and the number of categories for categorical losses).") end

    # Determine observed entries of data
    if obs==nothing && sparse_na && isa(A,SparseMatrixCSC)
        I,J = findn(A) # observed indices (vectors)
        obs = [(I[a],J[a]) for a = 1:length(I)] # observed indices (list of tuples)
    end
    if obs==nothing # if no specified array of tuples, use what was explicitly passed in or the defaults (all)
        # println("no obs given, using observed_features and observed_examples")
        glrm = GLRM(A,losses,rx,ry,k, observed_features, observed_examples, X,Y)
    else # otherwise unpack the tuple list into arrays
        # println("unpacking obs into array")
        glrm = GLRM(A,losses,rx,ry,k, sort_observations(obs,size(A)...)..., X,Y)
    end

    # check to make sure X is properly oriented
    if size(glrm.X) != (k, size(A,1))
        # println("transposing X")
        glrm.X = glrm.X'
    end
    # check none of the observations are NaN
    if checknan
        for i=1:size(A,1)
            for j=glrm.observed_features[i]
                if isnan(A[i,j])
                    error("Observed value in entry ($i, $j) is NaN.")
                end
            end
        end
    end

    if scale # scale losses (and regularizers) so they all have equal variance
        equilibrate_variance!(glrm)
    end
    if offset # don't penalize the offset of the columns
        add_offset!(glrm)
    end
    return glrm
end

parameter_estimate(glrm::GLRM) = (glrm.X, glrm.Y)

function scale_regularizer!(glrm::GLRM, newscale::Number)
    scale!(glrm.rx, newscale)
    scale!(glrm.ry, newscale)
    return glrm
end


text:
using DrWatson
quickactivate(@__DIR__)
@show projectdir()
##
using CSV
using DataFrames
# using DataFramesMeta
using Dates
using StatsBase
using CovidSurvey
using StatsPlots
using LaTeXStrings

using KernelFunctions
using AbstractGPs
using Stheno

# Load standard packages from the Julia ecosystem
using Optim # Standard optimisation algorithms.
using ParameterHandling # Helper functionality for dealing with model parameters.
using Zygote # Algorithmic Differentiation
using Turing
using ReverseDiff
using Memoization
setadbackend(:reversediff)
Turing.setrdcache(true)

using ParameterHandling: value, flatten
plotlyjs()
#---------------------------------------------------------------------------
# utils
struct ELBO{X,Y,U}
    x::X
    y::Y
    u::U
end

function (g::ELBO)(θ)
    ps = unpack(θ)
    f  = build_gp(ps)
    fx = f(g.x, ps.v)
    fu = f(g.u, 1e-6)
    return -elbo(fx, g.y, fu) #-logpdf(f(X1, ps.v), y_friends)
end

struct NLML{X,Y}
    x::X
    y::Y
end

function (g::NLML)(θ)
    ps = unpack(θ)
    f  = build_gp(ps)
    fx = f(g.x, ps.v)
    return -logpdf(fx, g.y) #-logpdf(f(X1, ps.v), y_friends)
end

#----------------------------------------------------------------------------
# survey data
holidays = load( projectdir("data", "holidays.csv") ) |> DataFrame
Rt_SSI = National.readsurvey("/home/and/data/covidsurvey/contacts.csv")
Rt_SSI = Rt_SSI[:,[:Rt,:uncertainty_lower,:uncertainty_upper, :date]]
# survey = National.readsurvey("/home/and/data/covidsurvey/contacts_0425.csv")
# survey = National.readsurveys_raw()
survey = load( "/home/and/data/covidsurvey/smoothed_contacts.csv" ) |> DataFrame
survey[!,:weekend] = Float64.( dayofweek.(survey.date) .> 5 )
#Rt = load(projectdir("out", "Rt.csv"))|>DataFrame
Rt = load( projectdir("reports/2021-04-27/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2020-12-16_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
t = survey.date
holidays = holidays[holidays.date .∈ Ref(t),:]
#----------------------------------------------------------------------------
# plot survey data
plot(t, Rt.mq, ribbon=(Rt.mq - Rt.llq, Rt.uuq - Rt.mq), linewidth=0, label="95% quantiles")
plot!(t, Rt.mq, ribbon=(Rt.mq - Rt.lq, Rt.uq - Rt.mq), linewidth=2, label="50% quantiles", hover=mq)
plot!(survey.date, survey.colleagues, ticks=:native, label="colleagues")
plot!(survey.date, survey.family, ticks=:native, label="family")
plot!(survey.date, survey.strangers, ticks=:native, label="strangers")
plot!(survey.date, survey.friends, ticks=:native, label="friends")
#--------------------------------------------------------------------
# time series decomposition
survey_col = :friends_2m_p95
d = survey[:, survey_col]
y = zscore(d)

params = (
    s1 = positive(1.),
    s2 = positive(1.),
    s3 = positive(1.),
    l1 = positive(20.),
    l2 = positive(1.),
    # l3 = positive(50.),
    v  = positive(0.1),
)

function build_gp(ps::NamedTuple)
    return @gppp let
        f₁ = GP( ps.s1 * Matern52Kernel() ∘ ScaleTransform(1/ps.l1) )
        f₂ = GP( ps.s2 * Matern12Kernel() ∘ ScaleTransform(1/ps.l2) )
        f₃ = GP( ps.s3 * PeriodicKernel() ∘ ScaleTransform(1/7)  )
        f₀ = f₁ + f₂ + f₃
    end
end

x_ = 1. : length(survey.date)
x = GPPPInput(:f₀, x_)

flat_initial_params, unflatten = flatten(params)
unpack(θ::Vector) = value(unflatten(θ))
##
# obj = ELBO(x, y, u)
obj = NLML(x,y)
obj(flat_initial_params)

opt = Optim.optimize(obj, flat_initial_params, NelderMead(), Optim.Options(show_trace=true))
# opt = Optim.optimize(obj, opt.minimizer, LBFGS(), Optim.Options(show_trace=true), autodiff=:forward)
flat_final_params = opt.minimizer
final_params = value(unflatten(flat_final_params))
##
ps = final_params
f  = build_gp(final_params)
fx = f(x, final_params.v)
fp = posterior(fx, y);
#--------------------------------------------------------------------
x1 = GPPPInput(:f₁, x_);
x2 = GPPPInput(:f₂, x_);
x3 = GPPPInput(:f₃, x_);
##
pgfplotsx()
# Plots.scalefontsizes(.5)
p = let
    p1 = let
        p = plot(t, y, c=:black, label="data");
        plot!(p, t, fp(x); color=:red, ls=:dash, label="fit")
        plot!(p, t, fp(x1); color=:red, label="slow trend")
        ih = [findfirst(==(h),t) for h in holidays.date]
        th = t[ih]
        yh = y[ih]
        scatter!(th, yh, label="holidays", mα=0, msa=1, hover=holidays.name)
        # annotate!([(Dates.value(th[i]) , yh[i], Plots.text(holidays.name[i], :top, 8) ) for (i,n) in enumerate(holidays.name)])
    end

    p2 = let
        p = plot()
        plot!(p, t, fp(x2); color=:orange, label="fast trend");
        hline!([0], label=nothing, c=:black, lw=2)
    end

    # p3 = let
    #     p = plot()
    #     plot!(p, t, fp(x3); color=:green, label="quasi-periodic trend");
    #     hline!([0], label=nothing, c=:black, lw=2)
    # end

    p4 = let
        df = DataFrame(
            :date => t,
            :weekly_effect => mean.( marginals( fp(x3) ) ),
            :day_of_week => dayofweek.(t)
        )
        gdf = groupby(df, :day_of_week, sort=true)
        cdf = combine(gdf, :weekly_effect => mean)

        p = plot()
        plot!(p, dayname.(cdf.day_of_week), cdf.weekly_effect_mean; color=:green, label="weekly effect", lw=2);
        hline!([0], label=nothing, c=:black, lw=2)
    end

    p5 = let
        ϵ = y-mean.( marginals( fp(x) ) )
        ih = [findfirst(==(h),t) for h in holidays.date]
        ϵh = ϵ[ih]
        p = plot()
        plot!(p, t, ϵ, label="residuals", c=:orange)
        scatter!(p, holidays.date, ϵh, label="holidays", mα=0, msa=1, hover=holidays.name)
        hline!([0], label=nothing, c=:black, lw=2)
    end
    plot(p1, p2, p4, p5, layout=(4,1), shapex=true, size=(1500, 1000))
end
savefig( p, projectdir("figures/", "survey_$(survey_col)_decomposition.png") )
run(`firefox $( projectdir("figures/", "survey_$(survey_col)_decomposition.png") )`, wait=false)
#---------------------------------------------------------------------
p = let
    m = marginals( fp(x1) )
    mf = mean.(m)
    sf = std.(m)
    sr = std(Rt.mq)
    smf = std(mf)

    mf .-= mf[1]
    mf .*= sr / smf
    mf .+= Rt.mq[1]
    sf .*= sr / smf

    p = plot(Rt.date, Rt.mq, ribbon=(Rt.mq - Rt.llq, Rt.uuq - Rt.mq), linewidth=0, label="R_t 95% quantiles", size=(1200, 800))
    plot!(Rt.date, Rt.mq, ribbon=(Rt.mq - Rt.lq, Rt.uq - Rt.mq), linewidth=2, label="R_t 50% quantiles", hover=Rt.mq)
    plot!(t, mf, ribbon=(sf,sf), label="contacts to $(survey_col) (smoothed trend)")
    #plot!(t, survey.Rt)
end
savefig( p, projectdir("figures/", "survey_$(survey_col)_vs_Rt.png") )
run(`firefox $( projectdir("figures/", "survey_$(survey_col)_vs_Rt.png") )`, wait=false)
#-----------------------------------------------------------------------------
smoothed[!, survey_col] = mean.( marginals( fp(x1) ) )
# smoothed2 = rename( smoothed, :family_2m_p95=>:family, :colleagues_2m_p95=>:colleagues, :strangers_2m_p95=>:strangers, :friends_2m_p95=>:friends)
# save( "/home/and/data/covidsurvey/smoothed_contacts_0425_2m.csv", smoothed)
# smoothed = DataFrame(:date => t)
##
plotlyjs()
smoothed = load( "/home/and/data/covidsurvey/smoothed_contacts.csv" ) |> DataFrame
contact = [:family, :colleagues, :friends, :strangers]
##
plotlyjs()
p = let
    p = let
            plot(legend=:outertopright, ticks=:native)
            plot!(Rt.date, Rt.mq, ribbon=(Rt.mq - Rt.lq, Rt.uq - Rt.mq), linewidth=2, label="R_t 50% quantiles")
            plot!(Rt_SSI.date, Rt_SSI.Rt, ribbon=(Rt_SSI.Rt - Rt_SSI.uncertainty_lower, Rt_SSI.uncertainty_upper - Rt_SSI.Rt), linewidth=2, label="R_t (SSI)", hover=Rt_SSI.Rt)
    end;

    ps = [p]
    for i in 1:4
        p = plot(legend=:outertopright)
        plot!(smoothed.date, smoothed[:, contact[i]], label="$(contact[i]) (smoothed)")
        y = zscore(survey[:, contact[i]])
        plot!(survey.date, y, label="$(contact[i]) (raw)")
        push!(ps, p)
    end;
    plot(ps..., layout=(5,1), size=(800, 1000), sharex=true, link=:x)
end
##
# p = let
pgfplotsx()
# rt = load( projectdir("reports/2021-04-27/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2020-12-16_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
# rt = load( projectdir("reports/2021-04-28/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2021-02-07_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
# rt = load(projectdir("out", "Rt.csv"))|>DataFrame
# rt = load(projectdir("data", "Rt_SSI.csv"))|>DataFrame
# rt = load(projectdir("reports/test_single_observation_AlternativeDelays/Rt_cases=2020-06-01_chains=10_hospit=true_model=cases_observ=2021-02-06_sero=true_steps=2000_warmup=1000.csv")) |> DataFrame
rt = load(projectdir("reports/test_single_observation_AlternativeDelays/Rt_cases=2020-06-01_chains=10_hospit=true_model=hospit_observ=2021-02-06_sero=true_steps=2000_warmup=1000.csv")) |> DataFrame

function plot_Rt_vs_contacts(c)
    p = plot(title="Rt (hospit) vs. $c")
    #plot!(Rt.date, Rt.mq, ribbon=(Rt.mq - Rt.lq, Rt.uq - Rt.mq), linewidth=2, label="R_t 50% quantiles")
    plot!(rt.date, rt.mq
        , ribbon=(rt.mq - rt.llq
        , rt.uuq - rt.mq)
        , lw=2
        , label=L"R_t"
        , hover=rt.mq
        , c=:black
        , ylabel = "reproduction rate"
        , legend=:topleft)
    y = zscore(survey[:, c])
    plot!(twinx()
        , survey.date .+ Day(0), y
        , label="Daily contacts"
        , lc=:darkorange
        , lw=4
        , ylabel="10-day lagged contacts (standardized)"
        , legend=:topright
        , xticks = :none)
    return p
end

p = let
    ps = []
    for c in [:family, :colleagues, :friends, :strangers]
        p = plot_Rt_vs_contacts(c)
        push!(ps, p)
    end
    plot(ps..., layout=(2,2), size=(1200, 1000))
end
##
savefig(p, projectdir("figures/" ,"Rt-hospit_vs_smoothed-contacts.png"))
##
function plot_Rt_vs_contacts(i)
    p = plot(title="$(contact[i])")
    #plot!(Rt.date, Rt.mq, ribbon=(Rt.mq - Rt.lq, Rt.uq - Rt.mq), linewidth=2, label="R_t 50% quantiles")
    plot!(Rt_SSI.date, Rt_SSI.Rt
        , ribbon=(Rt_SSI.Rt - Rt_SSI.uncertainty_lower
        , Rt_SSI.uncertainty_upper - Rt_SSI.Rt)
        , lw=2
        , label="R_t (SSI)"
        , hover=Rt_SSI.Rt
        , c=:black
        , ylabel = "reproduction rate"
        , legend=:topleft)
    # y = zscore(survey[:, contact[i]])
    y = smoothed[:, contact[i]]
    plot!(twinx()
        , smoothed.date .+ Day(10), y
        , label="smoothed daily contacts"
        , lc=:darkblue
        , ylabel="10-day lagged contacts (standardized)"
        , legend=:topright
        , xticks = :none)
    return p
end

p = let
    ps = []
    for i in 1:4
        p = plot_Rt_vs_contacts(i)
        push!(ps, p)
    end
    plot(ps..., layout=(2,2), size=(1200, 1000))
end

# savefig(p, projectdir("figures/" ,"decomposed_contacts_smoothed.png"))
##

ps = [p]
for i in 1:4
    p = plot(legend=:outertopright)
    plot!(smoothed.date, smoothed[:, contact[i]], label="$(contact[i]) (smoothed)")
    y = zscore(survey[:, contact[i]])
    plot!(survey.date, y, label="$(contact[i]) (raw)")
    push!(ps, p)
end;
plot(ps..., layout=(5,1), size=(800, 1000), sharex=true, link=:x)
# end
## =============================================================================
#                 fit survey data to Rt
# ==============================================================================
function getx(survey; suffix="", norm=true)
    xs = []
    cols = [
        Symbol("family"*suffix),
        Symbol("colleagues"*suffix),
        Symbol("friends"*suffix),
        Symbol("strangers"*suffix)
    ]
    for col in cols
        x = survey[:,col] .|> Float64
        norm && National.surveynorm!(x)
        push!(xs, x)
    end
    return hcat(xs...)
end

function make_consistent(df1, df2; startdate=nothing, enddate=nothing)
    df1 = deepcopy(df1)
    df2 = deepcopy(df2)
    ts = isnothing(startdate) ? max(df1.date[1], df2.date[1]) : Date(startdate)
    te = isnothing(enddate) ? min(df1.date[end], df2.date[end]) : Date(enddate)
    df1 = df1[ findfirst(==(ts), df1.date) : findfirst(==(te), df1.date), :]
    df2 = df2[ findfirst(==(ts), df2.date) : findfirst(==(te), df2.date), :]
    @assert all(df1.date .== df2.date)
    return df1, df2
end

function getcov(survey, rt; suffix="")
    x = covariates(survey; suffix) .|> Float64
    t = survey.date
    y = rt.mq .|> Float64
    # lq = rt.llq
    # uq = rt.uuq
    # u = (y .- lq) / 1.65 #1.95
    # y = Rt.mq #Rt_SSI.Rt
    # lq = Rt.llq
    return (; t, x, y)
end

function gety_SSI(rt)
    lq = rt.uncertainty_lower
    uq = rt.uncertainty_upper
    y = rt.Rt
    u = max.(abs.(y .- lq), abs.(uq .- y)) / 1.65 #1.95
    t = rt.date
    return (; t, y, u, lq, uq)
end

function gety(rt)
    lq = rt.llq
    uq = rt.uuq
    y = rt.mq
    u = max.(abs.(y .- lq), abs.(uq .- y)) / 1.95
    t = rt.date
    return (; t, y, u, lq, uq)
end
## plot Rt vs Rt_SSI
p = let
    # rt = load( projectdir("reports/2021-04-28/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2021-02-07_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
    rt2 = load( projectdir("reports/constant_iar/Rt_cases=2020-06-01_chains=10_hospit=true_model=const-iar_observ=2021-02-06_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
    rt = load( projectdir("reports/const_iar_cases_only/Rt_cases=2020-06-01_chains=5_hospit=true_model=const-iar_observ=2021-02-06_sero=true_steps=5000_warmup=2000.csv")) |> DataFrame
    rt_SSI, rt = make_consistent(Rt_SSI, rt)
    plot(legend=:outertopright)
    plot!(rt.date, rt.mq, ribbon=(rt.mq - rt.lq, rt.uq - rt.mq), linewidth=2, label="R_t")
    plot!(rt_SSI.date, rt_SSI.Rt, ribbon=(rt_SSI.Rt - rt_SSI.uncertainty_lower, rt_SSI.uncertainty_upper - rt_SSI.Rt), linewidth=2, label="R_t (SSI)", hover=Rt_SSI.Rt)
    plot!(rt2.date, rt2.mq, ribbon=(rt2.mq - rt2.lq, rt2.uq - rt2.mq), linewidth=2, label="R2_t")
end
##
# Plots.scalefontsizes(1.3)
# pgfplotsx()
plotlyjs()
p = let
    rt = load( projectdir("data/Rt.csv") ) |> DataFrame
    y1 = rt[:,"Rt-observ-2021-02-07"]
    y2 = rt[:,"Rt-SSI-shifted"]
    t  = rt.date
    p = plot(legend=:bottomright, ylabel="R_t", xticks=:native)
    plot!(t, y1, lw=2, label="without predictors")
    plot!(t, y2, lw=2, label="SSI")
    # xlims!((t[1],t[end]))
end
savefig(projectdir("figures/compare_Rt_vs_Rt-SSI.html"))
##
survey_fname = "/home/and/data/covidsurvey/contacts_0425.csv"

suffix=""
d = National.readsurvey(survey_fname)
d = National.readsurveys()
fname = normpath( homedir(), "data/covidsurvey/smoothed_contacts_0425_2m.csv" )
fname = normpath( homedir(), "data/covidsurvey/smoothed_contacts.csv" )
# fname = projectdir("data", "Rt_SSI.csv")
# fname = projectdir("reports/test_single_observation_AlternativeDelays/Rt_cases=2020-06-01_chains=10_hospit=true_model=cases_observ=2021-02-06_sero=true_steps=2000_warmup=1000.csv")
d = CSV.File(fname)|>DataFrame
d.date += Day(12)

rt = load( projectdir("reports/2021-04-27/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2020-12-16_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
rt = load( projectdir("reports/2021-04-28/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2021-02-07_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
rt = load(projectdir("out", "Rt.csv"))|>DataFrame
rt = deepcopy(Rt_SSI)
rt = load(projectdir("data", "Rt_SSI.csv"))|>DataFrame
rt = load(
    projectdir("reports/test_single_observation_AlternativeDelays/Rt_cases=2020-06-01_chains=10_hospit=true_model=cases_observ=2021-02-06_sero=true_steps=2000_warmup=1000.csv")
) |> DataFrame
d, rt = make_consistent(d, rt, enddate="2020-12-16")
# x = d.mq
# National.surveynorm!(x)
# x = x[:,:]
x = getx(d; suffix)
t, y, u, lq, uq = gety(rt)
# t, y, u, lq, uq = gety_SSI(rt)

# survey = National.readsurveys()
# survey = National.readsurveys_raw()
# survey = National.readsurvey(survey_fname)

# masks_start = findfirst(==(Date("2020-10-29")), t)
# @model function linear_regression_masks(y, x, u, i, ::Type{TV} = Vector{Float64}) where TV
#     n, m = size(x)
#     α_with_mask  ~ filldist(Beta(1, 50), m)
#     α_without_mask  ~ filldist(Beta(1, 50), m)
#     σ2 ~ InverseGamma(2.,.1)
#     ŷ  = TV(undef, n)
#     ŷ[1:i]  .= log(y[1]) .+ x[1:i,:] * α_without_mask
#     ŷ[i+1:end] .= log(y[1]) .+ x[i+1:end,:] * α_with_mask
#     uy ~ MvNormal(y, u)
#     uy ~ MvNormal(exp.(ŷ), σ2)
#     return exp.(ŷ)
# end

@model function linear_regression(y, x, u; link=exp, invlink=log)
    n, m = size(x)
    # α ~ filldist(Beta(1, 5), m)
    α ~ filldist(InverseGamma(2.,1.), m)
    σ2 ~ InverseGamma(2.,.1)
    ŷ = invlink(y[1]) .+ x * α
    # uy ~ MvNormal(y, u)
    # uy ~ MvNormal(link.(ŷ), σ2)
    y ~ MvNormal(link.(ŷ), σ2)
    return link.(ŷ)
end

# m = linear_regression_masks(y, x, u, i; link=KLogistic(3.), invlink=KLogit(3.))
m = linear_regression(y, x, u; link=KLogistic(4.), invlink=KLogit(4.))
c = sample(m, NUTS(200, 0.6), 1200)
c = c[200+1:end,:,:]
ŷs = generated_quantities(m,c)
ŷs = hcat(ŷs...)
##

##
# pgfplotsx()
plotlyjs()
# plot(c)
p = let
    ŷ = mean(ŷs, dims=2)
    e = sum(abs2, y-ŷ) / length(y)
    p = plot(legend=:outertop)
    plot!(t, y
        , ribbon=(y - lq, uq - y)
        # , ribbon=(y - Rt.llq, Rt.uuq - y)
        , lw=2
        , label="R_t (SSI)"
        , c=:black
        , ylabel = "reproduction rate"
        # , legend=:topright
        , α=.4)
    CovidSurvey.plot_confidence_timeseries!(p, t, ŷs, label = "fit (MSE = $e)")
    # vline!([Date("2020-10-29")], label="masks introduced")
    p
end
# plot(c)
# savefig(p, projectdir("figures/", "linearregression_smoothed_survey_on_Rt.png") )
# savefig(p, projectdir("figures/", "linearregression_smoothed_survey_on_Rt.html") )
#---------------------------------------------------------------------------
# make new prior for Rt
# rt = deepcopy(Rt_SSI)
rt = load( projectdir("reports/2021-04-28/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2021-02-07_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
d = National.readsurveys()
d, rt = make_consistent(d, rt)

# t, y, u, lq, uq = gety_SSI(rt)
t, y, u, lq, uq = gety(rt)

# prior = [Normal(μ, σ) for (μ, σ) in zip(y,u)]
save(projectdir("reports/test_contacts/RtPrior.csv"), DataFrame((; t, y, u)))

# prior = [fit(Normal{Float64}, ŷ) for ŷ in eachrow(ŷs)]
# plot(mean.(prior), ribbon=std.(prior))
# plot!(y, ribbon=(y - lq, uq - y), linewidth=2, label="R_t")
# plot!(mean(ŷs, dims=2), ribbon=std(ŷs, dims=2))
#---------------------------------------------------------------------------
# plot posterior effect size distribution vs. prior
using Plots.PlotMeasures
pgfplotsx()
p = let
    # [:family, :colleagues, :friends, :strangers]
    xs = c["α[1]"]|>vec
    p1 = histogram(xs, normalize=true, title="contacts to family", label=nothing, fc=:gray, lw=.1, xlabel="effect size", bottom_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α[2]"]|>vec
    p2 = histogram(xs, normalize=true, title="contacts to colleagues", label=nothing, fc=:gray, lw=.1, xlabel="effect size", bottom_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α[3]"]|>vec
    p3 = histogram(xs, normalize=true, title="contacts to friends", label=nothing, fc=:gray, lw=.1, xlabel="effect size", top_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α[4]"]|>vec
    p4 = histogram(xs, normalize=true, title="contacts to strangers", label=nothing, fc=:gray, lw=.1, xlabel="effect size", top_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    plot(p1,p2,p3,p4, layout=(2,2), legend=:topright, )
end
#----------------------------------------------------------------------------
using Plots.PlotMeasures
pgfplotsx()
p1 = let
    # [:family, :colleagues, :friends, :strangers]
    xs = c["α_without_mask[1]"]|>vec
    p1 = histogram(xs, normalize=true, title="contacts to family", label=nothing, fc=:gray, lw=.1, xlabel="effect size", bottom_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α_without_mask[2]"]|>vec
    p2 = histogram(xs, normalize=true, title="contacts to colleagues", label=nothing, fc=:gray, lw=.1, xlabel="effect size", bottom_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α_without_mask[3]"]|>vec
    p3 = histogram(xs, normalize=true, title="contacts to friends", label=nothing, fc=:gray, lw=.1, xlabel="effect size", top_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α_without_mask[4]"]|>vec
    p4 = histogram(xs, normalize=true, title="contacts to strangers", label=nothing, fc=:gray, lw=.1, xlabel="effect size", top_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    plot(p1,p2,p3,p4, layout=(2,2), legend=:topright, )
end

p2 = let
    # [:family, :colleagues, :friends, :strangers]
    xs = c["α_with_mask[1]"]|>vec
    p1 = histogram(xs, normalize=true, title="contacts to family", label=nothing, fc=:gray, lw=.1, xlabel="effect size", bottom_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α_with_mask[2]"]|>vec
    p2 = histogram(xs, normalize=true, title="contacts to colleagues", label=nothing, fc=:gray, lw=.1, xlabel="effect size", bottom_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α_with_mask[3]"]|>vec
    p3 = histogram(xs, normalize=true, title="contacts to friends", label=nothing, fc=:gray, lw=.1, xlabel="effect size", top_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    xs = c["α_with_mask[4]"]|>vec
    p4 = histogram(xs, normalize=true, title="contacts to strangers", label=nothing, fc=:gray, lw=.1, xlabel="effect size", top_margin=3mm, right_margin=0mm);
    # plot!(Beta(1,1), xlims=(0,0.1), lw=2, label="prior", legend=:topright);
    vline!([mean(xs)], lw=2, label="mean = $(round(mean(xs), sigdigits=2))", c=:black, ls=:dash)
    plot(p1,p2,p3,p4, layout=(2,2), legend=:topright, )
end

plot(p1, p2, layout=(2,1), size=(1000, 1500))
# savefig(p, projectdir("figures/", "linearregression_smoothed_survey_on_Rt_effect_analysis.png") )
# savefig(p, projectdir("figures/", "linearregression_smoothed_survey_on_Rt_effect_analysis.html") )
##
survey_fname = "/home/and/data/covidsurvey/contacts_0425.csv"
survey = National.readsurvey(survey_fname)
#--------------------------------------------------------------------------
# variable day
@model function linear_regression(y, x, u; link=exp, invlink=log)
    n, m = size(x)
    # α ~ filldist(Beta(1, 5), m)
    α ~ filldist(InverseGamma(2.,1.), m)
    σ2 ~ InverseGamma(2.,.1)
    ŷ = invlink(y[1]) .+ x * α
    # uy ~ MvNormal(y, u)
    # uy ~ MvNormal(link.(ŷ), σ2)
    y ~ MvNormal(link.(ŷ), σ2)
    return link.(ŷ)
end

fname = normpath( homedir(), "data/covidsurvey/smoothed_contacts.csv" )
# fname = projectdir("data", "Rt_SSI.csv")
# fname = projectdir("reports/test_single_observation_AlternativeDelays/Rt_cases=2020-06-01_chains=10_hospit=true_model=cases_observ=2021-02-06_sero=true_steps=2000_warmup=1000.csv")
d = CSV.File(fname)|>DataFrame

function getdata(rt, d, lag)
    d  = deepcopy(d)
    rt = deepcopy(rt)
    d.date += Day(lag)
    d, rt = make_consistent(d, rt, enddate="2021-01-15")
    x = getx(d)
    t, y, u, lq, uq = gety(rt)
    # t, y, u, lq, uq = gety_SSI(rt)
    return t, x, y, u
end

# rt = load( projectdir("reports/2021-04-28/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2021-02-07_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
# rt = load(
#     projectdir("reports/test_single_observation_AlternativeDelays/Rt_cases=2020-06-01_chains=10_hospit=true_model=cases_observ=2021-02-06_sero=true_steps=2000_warmup=1000.csv")
# ) |> DataFrame
rt = load(
    projectdir("data/Rt_SSI.csv")
) |> DataFrame
function delay_analysis(rt, d; lags=0:1:25)
    es = []
    ŷs = []
    ts = []
    ps = []
    @progress for l in lags
        t, x, y, u = getdata(rt, d, l)
        m = linear_regression(y, x, u; link=KLogistic(3.), invlink=KLogit(3.))
        c = sample(m, NUTS(200, 0.6), 1200)
        c = c[200+1:end,:,:]
        gq = generated_quantities(m,c)
        ŷ = mean(gq)
        e = sum(abs2, y-ŷ)/length(y)
        @show e
        p = plot(t, ŷ, title="$l: $e");
        plot!(rt.date, y);
        push!(ps, p)
        push!(ts, t)
        push!(es, e)
        push!(ŷs, ŷ)
    end
    (; ps, ts, es, ŷs)
end
lags = 0:1:20
res = delay_analysis(rt, d; lags)
# res_SSI = deepcopy(res)
##
pgfplotsx()
p = let
    p1 = plot(lags, res.es, xlabel="lag [days]", ylabel="mean squared error", label="mse after shift \\& fit", title="find time lag that minimizes error") #best lag = 11
    i = 12#argmin(res.es)
    vline!([i], label="best fit")
    p2 = plot(rt.date, rt.mq, label=L"R_t")
    plot!(res.ts[i], res.ŷs[i], label="best fit", ylabel=L"R_t")
    plot(p1, p2, layout=(2,1), size=(800, 500))
end
savefig(projectdir("figures/Rt_vs_predictors_delayed.png"))
## cross correlation analysis
let
    lags = 0:100
    rt = load( projectdir("reports/2021-04-28/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2021-02-07_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
    df1, df2 = make_consistent(smoothed, rt)
    y1 = df2.mq
    p1 = plot()
    p2 = plot()
    for n in [:family, :friends, :strangers, :colleagues]
        y2 = df1[:,n]
        cc = crosscor(y1, y2, lags)
        err = [sum(abs2, y1[1:end-t+1] - y2[t:end])/length(y2[t:end]) for t in 1:101]
        plot!(p1, lags, cc, label=String(n))
        plot!(p2, lags, err, label=String(n))
    end
    plot(p1, p2, layout=(2,1))
end
##
let
    lags = 0:10
    # rt = load( projectdir("reports/2021-04-28/Rt_cases=2020-06-01_chains=10_hospit=true_model=v2_observ=2021-02-07_seed=10000_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
    rt = load( projectdir("reports/test_single_observation_AlternativeDelays/Rt_cases=2020-06-01_chains=10_hospit=true_model=cases_observ=2021-02-06_sero=true_steps=2000_warmup=1000.csv") ) |> DataFrame
    rt = rt[rt.date .<= Date("2020-12-16"),:]
    Rt_SSI = National.readsurvey("/home/and/data/covidsurvey/contacts.csv")
    df1, df2 = make_consistent(Rt_SSI, rt)
    y1 = df1.Rt
    y2 = df2.mq
    cc = crosscor(y1, y2, lags)
    err = [sum(abs2, y1[1:end-t+1] - y2[t:end])/length(y2[t:end]) for t in lags.+1]
    p1 = plot(lags, cc)
    p2 = plot(lags, err)
    plot(p1, p2, layout=(2,1))
end

##
plotlyjs()
pgfplotsx()
scalefontsizes(.8)
p = let
    p = let
            plot(legend=:outertopright)
            plot!(t, y, ribbon=(y - lq, uq - y), linewidth=2, title=L"R_t", label=nothing)
            # plot!(Rt_SSI.date, Rt_SSI.Rt, ribbon=(Rt_SSI.Rt - Rt_SSI.uncertainty_lower, Rt_SSI.uncertainty_upper - Rt_SSI.Rt), linewidth=2, label="R_t (SSI)", hover=Rt_SSI.Rt)
    end

    ps = [p]
    for (i,n) in enumerate(["family", "colleagues", "friends", "strangers"])
        p = plot(legend=:outertopright)
        plot!(t, x[:, i], title="$n", label=nothing)
        push!(ps, p)
    end;
    plot(ps..., layout=(5,1), size=(800, 1000), sharex=true)
end
savefig("/home/and/dev/CovidSurvey/figures/zero_contacts.png")

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
export slideshow, slide, fragment

@api slideshow => (SlideShow <: Selection) begin
    doc("A slideshow.")
    curry(tiles::TileList, doc="An array of slides.")
    kwarg(selected::Int=1, doc="Index of the current slide.")
end

render_slide(x, state) =
    packitems(center, packacross(center, vbox([x]))) |>
        (t -> render(t, state) & @d(:attributes => @d(:fit => :fit)))

@api slide => (Slide <: Tile) begin
    doc("A slide. Content of the slide is centered vertically and horizontally.")
    curry(tile::Tile, doc="Content of the slide.")
    kwarg(
        transitions::AbstractString="slide-from-right cross-fade-all",
        doc=md"""A space-separated list of transitions to use for the current
                slide. Valid transitions are `"cross-fade"`,
                `"hero-transition"`,`"list-cascade"`, `"scale-up"`,
                `"slide-down"`, `"slide-from-bottom"`,`"slide-from-right"`,
                `"slide-up"` and `"tile-cascade"`."""
        )
end

@api fragment => (Fragment <: Tile) begin
    typedarg(index::Int=0)
    curry(tile::Tile)
end

render(f::Fragment, state) =
    Elem("slide-fragment", render(f.tile, state),
        index=f.index)

render(s::Slide, state) =
    Elem("slide-tile", render(s.tile, state)) & @d(
        "inTransitions" => s.transitions)

wrapslide(s::Slide) = s
wrapslide(s) = slide(s)

render(x::SlideShow, state) =
    Elem("slide-show",
        map(x -> render(wrapslide(x), state), x.tiles.tiles),
        selected=x.selected-1
    )


text:
include("Conserved.jl")
include("HydroSnapshot.jl")

mutable struct HDSim
    time::Float64 # Time
    cycle::Int64 # Cycle number
    state::HydroSnapshot # Primitive variables and grid
    extensives::Array{Conserved, 1} # Conserved quantities
    eos # Equation of state
    tsf # Time step function
    vm # Vertex motion
    fc # Flux calculator
    pg # Physical geometry
    st # Source term
    cu # Cell updater
end

function primitive2extensive(p::PrimitiveVars,
                             volume::Float64,
                             eos)::Conserved
    mass = p.density*volume
    momentum = mass*p.velocity
    kinetic_specific_energy = 0.5*p.velocity^2
    thermal_specific_energy =
     calcSpecificThermalEnergy(eos,
                        p.density,
                        p.pressure)
    specific_energy =
     kinetic_specific_energy+
     thermal_specific_energy
    energy = specific_energy*mass
    return Conserved(mass,
                     momentum,
                     energy)
end

function initConserved(state::HydroSnapshot, pg, eos)
    cumulative_volume = [calcVolume(pg, r) for r in state.grid]
    volumes = diff(cumulative_volume)
    conserved_list = [primitive2extensive(p,V,eos)
                      for (p,V) in zip(state.cells, volumes)]
    return conserved_list
end

function initHDSim(init_cond::HydroSnapshot,
                   eos,
                   tsf,
                   vm,
                   fc,
                   pg,
                   st,
                   cu)
    conserved_list = initConserved(init_cond, pg, eos)
    return HDSim(0,
                 0,
                 init_cond,
                 conserved_list,
                 eos,
                 tsf,
                 vm,
                 fc,
                 pg,
                 st,
                 cu)
end

function updateConserved!(conserved_list::Array{Conserved,1},
                          fluxes::Array{Conserved,1},
                          grid::Array{Float64,1},
                          pg,
                          dt::Float64)
    areas = [calcArea(pg,r) for r in grid]
    currents = dt.*areas.*fluxes
    conserved_list .-= diff(currents)
end

function updateSourceContribution!(st,
                                  extensives::Array{Conserved,1},
                                  state::HydroSnapshot,
                                  t::Float64,
                                  dt::Float64)
    extensives .+= dt*calcSource(st,state,t,dt)
end

#function updatePositions!(grid::Array{Float64, 1},
#                          velocities::Array{Float64, 1},
#                          dt::Float64)
#    grid .+= dt*velocities
#end

function timeAdvance(hdsim::HDSim)
    cached = Dict("sound_speeds"=>[calcSoundSpeed(hdsim.eos,
                                                  c.density,
                                                  c.pressure)
                                   for c in hdsim.state.cells],
                   "energies"=>[calcSpecificThermalEnergy(hdsim.eos,
                                                          c.density,
                                                          c.pressure)
                                for c in hdsim.state.cells])
    dt = calcTimeStep(hdsim.tsf, hdsim.state, hdsim.eos, cached)
    edge_velocities = calcEdgeVelocities(hdsim.vm, hdsim.state)
    fluxes = calcFluxes(hdsim.fc,
                        hdsim.state,
                        edge_velocities,
                        hdsim.eos,
                        cached)
    updateConserved!(hdsim.extensives,
                     fluxes,
                     hdsim.state.grid,
                     hdsim.pg,
                     dt)
    updateSourceContribution!(hdsim.st,
                              hdsim.extensives,
                              hdsim.state,
                              hdsim.time,
                              dt)
    #updatePositions!(hdsim.state.grid,
    #                 edge_velocities,
    #                 dt)
    #hdsim.state.grid .+= dt*edge_velocities
    for n in 1:length(edge_velocities)
        hdsim.state.grid[n] += dt*edge_velocities[n]
    end

    updateCells!(hdsim.cu,
                 hdsim.state,
                 hdsim.extensives,
                 hdsim.pg,
                 hdsim.eos)

    hdsim.time += dt
    hdsim.cycle += 1
end

max_stars_repo_name: bolverk/Sirion

max_issues_repo_name: bolverk/Sirion

max_forks_repo_name: bolverk/Sirion


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    ZScore()

Applies the z-score transform (a.k.a. normal score) to all columns of the table.
The z-score transform of the column `x`, with mean `μ` and standard deviation `σ`,
is defined by `(x .- μ) ./ σ`.
"""
struct ZScore <: Colwise end

assertions(::Type{ZScore}) = [assert_continuous]

isrevertible(::Type{ZScore}) = true

function colcache(::ZScore, x)
  μ = mean(x)
  σ = std(x, mean=μ)
  (μ=μ, σ=σ)
end

colapply(::ZScore, x, c)  = @. (x - c.μ) / c.σ

colrevert(::ZScore, y, c) = @. c.σ * y + c.μ


text:
#=
    All rights reserved. Ali Mert Ceylan 2020
=#

using Test
using AutoDiff

@testset "AutoDiff" begin

@time include("arthmtc.jl")
@time include("promtn.jl")
@time include("brdcstng.jl")
# @time include("specfncs.jl")

end


text:
"""
    all_mode_configurations(n, n_subset; only_photon_number_conserving=false)
    all_mode_configurations(input_state::Input, part::Partition; only_photon_number_conserving=false)
    all_mode_configurations(input_state::Input, sub::Subset; only_photon_number_conserving=false)

Generate all possible photon counts of `n` photons in a partition/subset
of `n_subset` subsets.

!!! note
    - Does not take into account photon number conservation by default
    - This is the photon counting in partitions and not events outputs but it
      can be used likewise
"""
function all_mode_configurations(n,n_subset; only_photon_number_conserving = false)

        array = []
        for i in 1:(n+1)^(n_subset)

                this_vector = digits(i-1, base = n+1, pad = n_subset)

                if only_photon_number_conserving
                        if sum(this_vector) == n
                                push!(array,this_vector)
                        end
                else
                        push!(array,this_vector)
                end

        end
        array

end

all_mode_configurations(input_state::Input,part::Partition; only_photon_number_conserving = false) = all_mode_configurations(input_state.n,part.n_subset; only_photon_number_conserving = only_photon_number_conserving)

all_mode_configurations(input_state::Input,sub::Subset; only_photon_number_conserving = false) = all_mode_configurations(input_state.n,1; only_photon_number_conserving = only_photon_number_conserving)

"""
    remove_trivial_partitions!(part_list)

In a list of partitions sizes, ex. `[[2,0],[1,1],[0,2]]`, keeps only
the elements with non trivial subset size, in this ex. only `[1,1]`.
"""
function remove_trivial_partitions!(part_list)

    filter!(x -> !any(x .== 0), part_list)

end

"""
    ranked_partition_list(part_list)

Remove partitions such as `[1,2]` when `[2,1]` is already counted as only the
size of the partition counts; only keeps vectors of decreasing count.
"""
function ranked_partition_list(part_list)

    output_partitions = []

    for part in part_list

        keep_this_part = true
        for i in 1:length(part)-1
            if part[i] < part[i+1]
                keep_this_part = false
                break
            end
        end

        if keep_this_part
            push!(output_partitions, part)
        end
    end

    output_partitions
end

"""
    photon_number_conserving_events(physical_indexes, n; partition_spans_all_modes=false)

Return only the events conserving photon number `n`.

!!! note
    - If `partition_spans_all_modes`=`false`, gives all events with less than `n` or `n` photons
    - If `partition_spans_all_modes` = `true` only exact photon number conserving physical_indexes
"""
function photon_number_conserving_events(physical_indexes, n; partition_spans_all_modes = false)

        results = []
        for index in physical_indexes
                if partition_spans_all_modes == false
                        if sum(index) <= n
                                push!(results, index)
                        end
                else
                        if sum(index) == n
                                push!(results, index)
                        end
                end
        end
        results

end

"""
    photon_number_non_conserving_events(physical_indexes, n; partition_spans_all_modes=false)

Return the elements not conserving the number of photons.
"""
function photon_number_non_conserving_events(physical_indexes,n ; partition_spans_all_modes = false)

        setdiff(physical_indexes, photon_number_conserving_events(physical_indexes, n, ; partition_spans_all_modes = partition_spans_all_modes))

end

"""
    check_photon_conservation(physical_indexes,  pdf, n; atol=ATOL, partition_spans_all_modes=false)

Check if probabilities corresponding to non photon number conserving events are zero.
"""
function check_photon_conservation(physical_indexes,  pdf, n; atol = ATOL, partition_spans_all_modes = false)

        events_to_check = photon_number_non_conserving_events(physical_indexes,n; partition_spans_all_modes = partition_spans_all_modes)

        for (i, index) in enumerate(physical_indexes)
                if index in events_to_check
                        @argcheck isapprox(clean_proba(pdf[i]),0, atol=atol)# "forbidden event has non zero probability"
                end
        end

end

"""
    compute_probabilities_partition(physical_interferometer::Interferometer, part::Partition, input_state::Input)

Compute the probability to find a certain photon counts in a partition `part` of
the output modes for the given interferometer.

Return `(counts = physical_indexes, probabilities = pdf) corresponding to the
occupation numbers in the partition and the associated probability.
"""
function compute_probabilities_partition(physical_interferometer::Interferometer, part::Partition, input_state::Input)

        @argcheck at_most_one_photon_per_bin(input_state.r) "more than one input per mode is not implemented"

        occupies_all_modes(part) ? (@warn "inefficient if no loss: partition occupies all modes thus extra calculations made that are unnecessary") : nothing

        n = input_state.n
        m = input_state.m
        mode_occupation_list = fill_arrangement(input_state)
        S = input_state.G.S

        fourier_indexes = all_mode_configurations(n,part.n_subset, only_photon_number_conserving = false)
        probas_fourier = Array{ComplexF64}(undef, length(fourier_indexes))
        virtual_interferometer_matrix = similar(physical_interferometer.U)

        for (index_fourier_array, fourier_index) in enumerate(fourier_indexes)

                # for each fourier index, we recompute the virtual interferometer
                virtual_interferometer_matrix  = physical_interferometer.U
                diag = [one(eltype(virtual_interferometer_matrix)) for i in 1:m]

                for (i,fourier_element) in enumerate(fourier_index)

                        this_phase = exp(2*pi*1im/(n+1) * fourier_element)

                        for j in 1:length(diag)

                                if part.subsets[i].subset[j] == 1
                                        diag[j] *= this_phase

                                end

                        end

                end

                virtual_interferometer_matrix *= Diagonal(diag)
                virtual_interferometer_matrix *= physical_interferometer.U'

                # beware, only the modes corresponding to the
                # virtual_interferometer_matrix[input_config,input_config]
                # must be taken into account !
                probas_fourier[index_fourier_array] = permanent(virtual_interferometer_matrix[mode_occupation_list,mode_occupation_list] .* S)
        end

        physical_indexes = copy(fourier_indexes)

        probas_physical(physical_index) = 1/(n+1)^(part.n_subset) * sum(probas_fourier[i] * exp(-2pi*1im/(n+1) * dot(physical_index, fourier_index)) for (i,fourier_index) in enumerate(fourier_indexes))

        pdf = [probas_physical(physical_index) for physical_index in physical_indexes]

        pdf = clean_pdf(pdf)

        check_photon_conservation(physical_indexes, pdf, n; partition_spans_all_modes = occupies_all_modes(part))

        (physical_indexes, pdf)
end

"""
    compute_probability_partition_occupancy(physical_interferometer::Interferometer, part_occupancy::PartitionOccupancy, input_state::Input)

Compute the probability to find a partition occupancy.

!!! note
    Inefficient to use multiple times for the same physical setting, rather use
    compute_probabilities_partition.
"""
function compute_probability_partition_occupancy(physical_interferometer::Interferometer, part_occupancy::PartitionOccupancy, input_state::Input)

        (physical_indexes, pdf) = compute_probabilities_partition(physical_interferometer, part_occupancy.partition, input_state::Input)

        for (i,counts) in enumerate(physical_indexes)
                if counts == part_occupancy.counts.state
                        return pdf[i]
                end
        end
        nothing

end

function print_pdfs(physical_indexes, pdf, n; physical_events_only = false, partition_spans_all_modes = false)

        indexes_to_print = physical_events_only ? photon_number_conserving_events(physical_indexes, n; partition_spans_all_modes = partition_spans_all_modes) : physical_indexes

        println("---------------")
        println("Partition results : ")
        for (i, index) in enumerate(physical_indexes)
                if index in indexes_to_print
                        println("index = $index, p = $(pdf[i])")
                end

        end
        println("---------------")
end

"""
    compute_probability!(ev::Event{TIn,TOut}) where {TIn<:InputType, TOut<:PartitionCount}
    compute_probability!(ev::Event{TIn,TOut}) where {TIn<:InputType, TOut<:PartitionCountsAll}

Given a defined [`Event`](@ref), computes/updates its probability or set of probabilities
(for instance if looking at partition outputs, with `MultipleCounts` begin filled).

This function is defined separately as it is most often the most time consuming
step of calculations and one may which to separate the evaluation of probabilities
from preliminary definitions.
"""
function compute_probability!(ev::Event{TIn,TOut}) where {TIn<:InputType, TOut<:PartitionCount}

        check_probability_empty(ev)

        ev.proba_params.precision = eps()
        ev.proba_params.failure_probability = 0

        ev.proba_params.probability = compute_probability_partition_occupancy(ev.interferometer, ev.output_measurement.part_occupancy, ev.input_state)

end

function compute_probability!(ev::Event{TIn,TOut}) where {TIn<:InputType, TOut<:PartitionCountsAll}

        check_probability_empty(ev)

        ev.proba_params.precision = eps()
        ev.proba_params.failure_probability = 0

        i = ev.input_state
        part = ev.output_measurement.part

        (part_occ,  pdf) = compute_probabilities_partition(ev.interferometer, ev.output_measurement.part, i)

        # clean up to keep photon number conserving events (possibly lossy events in the partition occupies all modes)

        part_occ_physical = []
        pdf_physical = []

        n = i.n

        for (i,occ) in enumerate(part_occ)
            if sum(occ) <= n
                push!(part_occ_physical, occ)
                push!(pdf_physical, pdf[i])
            end
        end

        mc = MultipleCounts([PartitionOccupancy(ModeOccupation(occ),n,part) for occ in part_occ_physical], pdf_physical)

        ev.proba_params.probability = EventProbability(mc).probability

end

#
# check_probability_empty(ev)
#
# interf = ev.interferometer
# part = ev.output_measurement.part_occupancy
# input_state = ev.input_state
#
# ev.proba_params.precision = eps()
# ev.proba_params.failure_probability = 0
#
# ev.proba_params.probability = nothing ####################33compute_probabilities_partition(ev.interferometer, ev.part::Partition, input_state::Input)


text:
# ____________________________________________________________________
# SpecLineLorentz
#
mutable struct SpecLineLorentz <: AbstractComponent
    norm::Parameter
    center::Parameter
    fwhm::Parameter
    voff::Parameter
    norm_integrated::Bool

    function SpecLineLorentz(center::Number)
        out = new(Parameter(1),
                  Parameter(center),
                  Parameter(3000),
                  Parameter(0),
                  true)
        @assert center > 0
        out.norm.low = 0
        out.center.low = 0
        out.fwhm.low = 0
        out.voff.low = 0
        out.center.fixed = true
        return out
    end
end

function prepare!(comp::SpecLineLorentz, domain::Domain{1})
    return fill(NaN, length(domain))
end

function evaluate!(buffer, comp::SpecLineLorentz, x::Domain{1},
                   norm, center, fwhm, voff)
    x0 = center - (voff / 3.e5) * center
    hwhm = fwhm / 3.e5 * center / 2  # Note: this is in the same units as `center`
    map!(x -> begin
         (abs.(x) .> 20)  &&  (return 0.)
         ret = norm ./ (1 .+ x.^2.)
         comp.norm_integrated  &&  (ret /= (pi * hwhm))
         return ret
         end,
         buffer, (x .- x0) ./ hwhm)
end

#=
    x = Domain(500:1:1500.)
    comp = QSFit.SpecLineLorentz(1000.)
    comp.fwhm.val = 3e4
    ceval = GFit.CompEval(comp, x)
    evaluate!(ceval)
    @gp x[:] ceval.buffer ./ maximum(ceval.buffer) "w l"
=#


text:
"""
    UniformContexts(limits::Matrix{Float64})

Construct a callable object to generate uniform contexts.
"""
struct UniformContexts <: AbstractContextSampler
    limits::Matrix{Float64}

    function UniformContexts(limits::Matrix{Float64})
        check_limits(limits)
        return new(limits)
    end
end

function (sampler::UniformContexts)(n::Int64)
    d = size(sampler.limits, 1)
    X = zeros(d, n)
    for i in 1:d
        X[i, :] = rand(Uniform(sampler.limits[i, 1], sampler.limits[i, 2]), n)
    end
    return X
end


text:
## cholesky.jl : Gaussian random field generator based on the Cholesky decomposition of the covariance matrix

"""
    Cholesky <: GaussianRandomFieldGenerator

A [`GaussiandRandomFieldGenerator`](@ref) based on a Cholesky factorization of the covariance matrix.

# Examples
```jldoctest
julia> m = Matern(0.1,1.0)
Matérn (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> c = CovarianceFunction(2,m)
2d Matérn covariance function (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> pts1 = 0:0.02:1; pts2 = 0:0.02:1
0.0:0.02:1.0

julia> grf = GaussianRandomField(c,Cholesky(),pts1,pts2)
Gaussian random field with 2d Matérn covariance function (λ=0.1, ν=1.0, σ=1.0, p=2.0) on a 51x51 structured grid, using a Cholesky decomposition

julia> plot(grf)
[...]

```
Note that the Cholesky factorization requires the covariance matrix to be `Symmetric` and positive definite. If the covariance matrix is not `Symmetric`, a warning will be thrown but the method will try to continue with an approximate symmetric matrix. If the covariance matrix is not positive definite, an error will be thrown. Try using the `Spectral` method in that case.

See also: [`Spectral`](@ref), [`KarhunenLoeve`](@ref), [`CirculantEmbedding`](@ref)
"""
struct Cholesky <: GaussianRandomFieldGenerator end

# compose a GaussianRandomField using Cholesky factorization
function _GaussianRandomField(mean, cov::CovarianceFunction, method::Cholesky, pts...)
    C = apply(cov,pts,pts)

    # error checking
    issymmetric(C) || @warn "to use a Cholesky factorization, the covariance matrix must be symmetric/hermitian"

    # compute Cholesky factorization
    isposdef!(Symmetric(C)) || throw(ArgumentError("to use a Cholesky factorization, the covariance matrix must be positive definite"))
    L = LowerTriangular(C')

    GaussianRandomField{Cholesky,typeof(cov),typeof(pts),typeof(mean),typeof(L)}(mean,cov,pts,L)
end

# returns the required dimension of the random points
randdim(grf::GaussianRandomField{Cholesky}) = size(grf.data, 1)

# sample from the GaussianRandomField using Cholesky factorization
function _sample(grf::GaussianRandomField{Cholesky}, xi)
	grf.mean + std(grf.cov) * reshape(grf.data * xi, size(grf.mean))
end

Base.show(io::IO, ::Cholesky) = print(io, "Cholesky decomposition")