# bounds.jl

"""
   bound_remainder(::Type{TaylorModel1}f::Function, polf::Taylor1, polfI::Taylor1, x0::Interval, I::Interval)

Bound the absolute remainder of the polynomial approximation of `f` given
by the Taylor polynomial `polf` around `x0` on the interval `I`. It requires
the interval extension `polfI` of the polynomial that approximates `f` for
the whole interval `I`, in order to compute the Lagrange remainder.

If `polfI[end]` has a definite sign, then it is monotonic in the intervals
[I.lo, x0] and [x0.hi, I.hi], which is exploited; otherwise, it is used
to compute the Lagrange remainder. This corresponds to Prop 2.2.1 in Mioara
Joldes PhD thesis (pp 52).

"""
function bound_remainder(::Type{TaylorModel1}, f::Function, polf::Taylor1, polfI::Taylor1, x0, I::Interval)

    _order = get_order(polf) + 1
    fTIend = polfI[_order]
    if (sup(fTIend) < 0 || inf(fTIend) > 0)
        # Absolute remainder is monotonic
        a = interval(I.lo)
        b = interval(I.hi)
        Δlo = f(a) - polf(a-x0)
        # Δlo = f(a) - bound_taylor1(polf, a-x0)
        Δhi = f(b) - polf(b-x0)
        # Δhi = f(b) - bound_taylor1(polf, b-x0)
        Δx0 = f(x0) - polf[0]
        Δ = hull(hull(Δlo, Δx0), Δhi)
    else
        # Lagrange bound
        Δ = fTIend * (I-x0)^_order
    end
    return Δ
end

"""
   bound_remainder(::Type{RTaylorModel1}, f::Function, polf::Taylor1, polfI::Taylor1, x0::Interval, I::Interval)

Bound the relative remainder of the polynomial approximation of `f` given
by the Taylor polynomial `polf` around `x0` on the interval `I`. It requires
an the interval extension `polfI` of a polynomial that approximates `f` for
the whole interval `I`, in order to compute the Lagrange remainder.

If `polfI[end]` has a definite sign, then it is monotonic in the interval `I`,
which is exploited; otherwise, the last coefficients bounds the relative
remainder. This corresponds to Prop 2.3.7 in Mioara Joldes' PhD thesis (pp 67).

"""
function bound_remainder(::Type{RTaylorModel1}, f::Function, polf::Taylor1, polfI::Taylor1, x0, I::Interval)

    _order = get_order(polf) + 1
    fTIend = polfI[_order+1]
    a = interval(I.lo)
    b = interval(I.hi)
    if (sup(fTIend) < 0 || inf(fTIend) > 0) && isempty(a ∩ x0) && isempty(b ∩ x0)
        # Error is monotonic
        denom_lo = (a-x0)^_order
        Δlo = f(a) - polf(a-x0)
        # Δlo = f(a) - bound_taylor1(polf, a-x0)
        Δlo = Δlo / denom_lo
        denom_hi = (b-x0)^_order
        Δhi = f(b) - polf(b-x0)
        # Δhi = f(b) - bound_taylor1(polf, b-x0)
        Δhi = Δhi / denom_hi
        Δ = hull(Δlo, Δhi)
    else
        # Lagrange coefficient
        Δ = polfI[_order]
    end
    return Δ
end

"""
    bound_taylor1(fT::Taylor1, I::Interval)

Compute a *tight* polynomial bound for the Taylor polynomial `fT`
in the interval `I`.

Note: Algorithm 2.1.1 corresponds to `evaluate(fT, I)` or simply `fT(I).
This function uses the roots of the derivative of `ft` to obtain a
tighter bound.

"""
function bound_taylor1(fT::Taylor1, I::Interval)

    # Check if the fT is monotonous (the derivative has a given sign)
    fTd  = TaylorSeries.derivative(fT)
    range_deriv = fTd(I)
    (sup(range_deriv) ≤ 0 || inf(range_deriv) ≥ 0) && return bound_taylor1(fT, fTd, I)

    # Compute roots of the derivative using the second derivative
    # Fix some sort of relative tolerance for Newton root search
    fTd2 = TaylorSeries.derivative(fTd)
    rootsder = roots(x->fTd(x), x->fTd2(x), I, Newton, 1.0e-5*mag(I))

    # Bound the range of fT using the roots and end points
    num_roots = length(rootsder)
    num_roots == 0 && return fT(I)
    rangepoly = hull( fT(interval(I.lo)), fT(interval(I.hi)) )
    for ind in 1:num_roots
        rangepoly = hull(rangepoly, fT(rootsder[ind].interval))
    end

    return rangepoly
end

"""
    bound_taylor1(fT::Taylor1, fTd::Taylor1, I::Interval)

Compute a *tight* polynomial bound for the Taylor polynomial `fT`
in the interval `I`, considering whether its derivative `ftd` has
a definite sign.

"""
function bound_taylor1(fT::Taylor1{T}, fTd::Taylor1{T}, I::Interval{T}) where {T}
    #
    if inf(fTd(I)) ≥ 0
        return Interval(fT(I.lo), fT(I.hi))
    elseif sup(fTd(I)) ≤ 0
        return Interval(fT(I.hi), fT(I.lo))
    end
    return fT(I)
end
function bound_taylor1(fT::Taylor1{Interval{T}}, fTd::Taylor1{Interval{T}},
        I::Interval{S}) where {T,S}
    #
    if inf(fTd(I)) ≥ 0
        return hull(fT(I.lo), fT(I.hi))
    elseif sup(fTd(I)) ≤ 0
        return hull(fT(I.hi), fT(I.lo))
    end
    return fT(I)
end

"""
    bound_taylor1(fT::TaylorModel1, I=domain(fT)::Interval)

Compute a *tight* polynomial bound for the Taylor model `fT`
in the interval `I`, considering whether its derivative `ftd` has
a definite sign.

"""
bound_taylor1(fT::TaylorModel1, I=domain(fT)::Interval) = bound_taylor1(polynomial(fT), I)

"""
    linear_dominated_bounder(fT::TaylorModel1, ϵ=1e-3::Float64, max_iter=5::Int)

Compute a tighter polynomial bound for the Taylor model `fT` by the linear
dominated bounder algorithm. The linear dominated algorithm is applied until
the bound of `fT` gets tighter than `ϵ` or the number of steps reachs `max_iter`.
The returned bound corresponds to the improved polynomial bound with the remainder
of the `TaylorModel1` included.
"""
function linear_dominated_bounder(fT::TaylorModel1{T, S}; ϵ=1e-3, max_iter=5) where {T, S}
    d = 1.
    Dn = fT.dom
    Dm = fT.x0
    Pm = Taylor1(copy(fT.pol.coeffs))
    bound = interval(0.)
    n_iter = 0
    while d > ϵ && n_iter < max_iter
        x0 = mid(Dn - Dm)
        c = mid(Dn)
        update!(Pm, x0)
        linear = Taylor1(Pm.coeffs[1:2], Pm.order)
        non_linear = Pm - linear
        if T <: Interval
            Li = mid(linear[1])
        else
            Li = linear[1]
        end
        I1 = linear(Dn - c)
        Ih = non_linear(Dn - c)
        bound = I1.lo + Ih
        d = diam(bound)
        n_iter += 1
        if Li == 0
            break
        elseif Li > 0
            new_hi = min(Dn.lo + (d / abs(Li)), Dn.hi)
            Dm = Dn
            Dn = interval(Dn.lo, new_hi)
        else
            new_lo = max(Dn.hi - (d / abs(Li)), Dn.lo)
            Dm = Dn
            Dn = interval(new_lo, Dn.hi)
        end
    end
    hi = fT.pol(fT.dom - fT.x0).hi
    return interval(bound.lo, hi) + fT.rem
end

"""
    linear_dominated_bounder(fT::TaylorModelN, ϵ=1e-3::Float64, max_iter=5::Int)

Compute a tighter polynomial bound for the Taylor model `fT` by the linear
dominated bounder algorithm. The linear dominated algorithm is applied until
the bound of `fT` gets tighter than `ϵ` or the number of steps reachs `max_iter`.
The returned bound corresponds to the improved polynomial bound with the remainder
of the `TaylorModelN` included.
"""
function linear_dominated_bounder(fT::TaylorModelN{N, T, S}; ϵ=1e-5, max_iter=5) where {N, T, S}
    d = 1.
    Dn = fT.dom
    Dm = fT.x0
    Pm = TaylorN(deepcopy(fT.pol.coeffs))
    bound = interval(0.)
    n_iter = 0
    x0 = Array{S}(undef, N)
    linear_coeffs = Array{Float64}(undef, N)
    while d > ϵ && n_iter < max_iter
        x0 .= mid(Dn - Dm)
        c = mid(Dn)
        update!(Pm, x0)
        linear_part = Pm[1]
        if T <: Interval
            linear_coeffs .= mid.(linear_part.coeffs)
        else
            linear_coeffs .= linear_part.coeffs
        end
        linear = TaylorN(deepcopy(Pm.coeffs[1:2]), Pm.order)
        non_linear = Pm - linear
        centered_domain = Dn .- c
        I1 = linear(centered_domain)
        Ih = non_linear(centered_domain)
        bound = I1.lo + Ih
        d = diam(bound)
        n_iter += 1
        new_boxes = Interval[]
        for (idx, box) in enumerate(Dn)
            Li = linear_coeffs[idx]
            if Li == 0
                Dni = box
            elseif Li < 0
                lo = max(box.hi - (d / abs(Li)), box.lo)
                Dni = interval(lo, box.hi)
            else
                hi = min(box.lo + (d / abs(Li)), box.hi)
                Dni = interval(box.lo, hi)
            end
            push!(new_boxes, Dni)
        end
        Dm = Dn
        Dn = IntervalBox(new_boxes...)
    end
    hi = fT.pol(fT.dom - fT.x0).hi
    return interval(bound.lo, hi) + fT.rem
end

"""
    quadratic_fast_bounder(fT::TaylorModel1)

Compute a *tighter* polynomial bound by the quadratic fast bounder.
The returned bound corresponds to the "improved" polynomial bound
with the remainder of the `TaylorModel1` included. This "improved" bound
can be one of the following two:
    1) An improved bound: if the domain of `fT` has a local minimizer,
       then an improved bound is returned.
    2) Original TaylorModel bound: if the local minimizer criteria is not
       satisfied, then the original bound of `fT` is returned.

This algorithm is a slight modification to the Makino & Berz algorithm.
For this algorithm the linear part is bounded by solving a simple
set of linear equations (compared to the iterative procedure by Makino & Berz).
"""
function quadratic_fast_bounder(fT::TaylorModel1)
    P = fT.pol
    bound_tm = fT(fT.dom - fT.x0)
    if signbit(P[2])
        return bound_tm
    else
        x0 = -P[1] / (2 * P[2])
        x = Taylor1(fT.pol.order)
        Qx0 = (x - x0) * P[2] * (x - x0)
        bound_qfb = (P - Qx0)(fT.dom - fT.x0)
        hi = P(fT.dom - fT.x0).hi
        bound_qfb = interval(bound_qfb.lo, hi) + fT.rem
        bound = bound_qfb ∩ bound_tm
        return bound
    end
end

"""
    quadratic_fast_bounder(fT::TaylorModelN)

Compute a *tighter* polynomial bound by the quadratic fast bounder.
The returned bound corresponds to the "improved" polynomial bound
with the remainder of the `TaylorModelN` included. This "improved" bound
can be one of the following two:
    1) An improved bound: if the domain of `fT` has a local minimizer,
       then an improved bound is returned.
    2) Original TaylorModel bound: if the local minimizer criteria is not
       satisfied, then the original bound of `fT` is returned.

This algorithm is a slight modification to the Makino & Berz algorithm.
For this algorithm the linear part is bounded by solving a simple
set of linear equations (compared to the iterative procedure by Makino & Berz).
"""
function quadratic_fast_bounder(fT::TaylorModelN)
    P = fT.pol
    H = Matrix(TaylorSeries.hessian(P))
    bound_tm = fT(fT.dom - fT.x0)
    if isposdef(H)
        P1 = -P[1].coeffs
        xn = H \ P1
        x = set_variables("x", numvars=length(xn))
        Qxn = 0.5 * (x - xn)' * H * (x - xn)
        bound_qfb = (P - Qxn)(fT.dom - fT.x0)
        hi = P(fT.dom - fT.x0).hi
        bound_qfb = interval(bound_qfb.lo, hi) + fT.rem
        bound = bound_qfb ∩ bound_tm
        return bound
    else
        return bound_tm
    end
end


using Printf

"""
Iterator for the sifting algorithm.
"""
mutable struct SiftIterable{solT <: AbstractVector, vecT <: AbstractVector, numT <: Real}
    "Measurements/IMF"
    y::solT
    "Time series values"
    x::vecT
    "Positions of the local maxima"
    maxs::Vector{Int}
    "Positions of the local minima"
    mins::Vector{Int}
    "Indices of the zero crossings"
    crosses::Vector{Int}
    "Tolerance"
    tol::numT
    "Residuals used to check stopping criteria"
    residual::numT
    "Maximum number of iteration"
    maxiter::Int
    "Number of steps of a stable sift after which the sifting is aborted"
    stop_steps::Int
    "Number of steps on which the number of zero crossings was fix"
    fix_steps::Int
end

@inline converged(it::SiftIterable) = it.residual ≤ var(it.y) * it.tol

@inline start(it::SiftIterable) = 0

@inline done(it::SiftIterable, iteration::Int) = iteration ≥ it.maxiter || converged(it)

# kernel sifting algorithm
function iterate(it::SiftIterable, iteration::Int=start(it))
    # Check for termination first
    if done(it, iteration)
        return nothing
    end

    localmaxmin!(it.y, it.maxs, it.mins)
    maxlen = length(it.maxs)
    minlen = length(it.mins)

    it.fix_steps == it.stop_steps && return nothing

    zerocrossing!(it.y, it.crosses)

    abs(length(it.crosses) - maxlen - minlen) <= 1 && (it.fix_steps += 1)
    if maxlen < 4 || minlen < 4
        return nothing
    end
    smin = get_edgepoint(it.y, it.x, it.mins, first, isless)
    smax = get_edgepoint(it.y, it.x, it.maxs, first, !isless)
    emin = get_edgepoint(it.y, it.x, it.mins, last, isless)
    emax = get_edgepoint(it.y, it.x, it.maxs, last, !isless)

    # TODO: Format for clarity?
    maxTS = interpolate([first(it.x); it.x[it.maxs]; last(it.x)], [smax; it.y[it.maxs]; emax], it.x, DierckXInterp())
    minTS = interpolate([first(it.x); it.x[it.mins]; last(it.x)], [smin; it.y[it.mins]; emin], it.x, DierckXInterp())
    subs = 0.5*(maxTS + minTS)
    it.residual = sum(abs, subs)

    it.y .-= subs

    # Return the residual at item and iteration number as state
    it.residual, iteration + 1
end

# Utility functions

function sift_iterator(y, x;
                      reltol::Real = 0.1,
                      maxiter::Int = 6,
                      stop_steps::Int = 4)
    maxs, mins = Int[], Int[]
    crosses = Int[]
    residual = sum(abs, y)
    tolerance = 0.1
    fix_steps = 0
    # Return the iterable
    return SiftIterable(y, x, maxs, mins, crosses,
        tolerance, residual, maxiter, stop_steps, fix_steps)
end

"""
    sift(y, x=1:length(y); kwargs...) -> imf

Same as [`sift!`](@ref), but allocates a solution vector `imf` initialized with zeros.
"""
sift(y, x=1:length(y); kwargs...) = sift!(zeros(eltype(y), size(y)), y, x; kwargs...)

"""
    sift!(imf, y, x; kwargs...) -> imf

# Arguments

- `imf`: intrinsic mode function, will be updated in-place;
- `y`: data measurements;
- `x`: time stamp series.

## Keywords

- `reltol::Real = 0.1`: relative tolerance for the stopping condition;
- `maxiter::Int = 10`: maximum number of iterations;
- `verbose::Bool = false`: print method information;
- `log::Bool = false`: keep track of the residual in each iteration.

# Output

**if `log` is `false`**

- `imf`: intrinsic mode function.
"""
function sift!(imf, y, x;
             reltol::Real = 0.1,
             maxiter::Int = 6,
             stop_steps::Int = 4,
             log::Bool = false,
             verbose::Bool = false,
             kwargs...)

    imf .= y
    # Actually perform sifting
    iterable = sift_iterator(imf, x; reltol, maxiter, stop_steps)

    for (iteration, item) = enumerate(iterable)
        verbose && @printf("%3d\t%1.2e\n", iteration, iterable.residual)
    end

    verbose && println()

    iterable.y
end


"""
### mmread(filename, infoonly::Bool=false, retcoord::Bool=false)
Read the contents of the Matrix Market file 'filename' into a matrix,
which will be either sparse or dense, depending on the Matrix Market format
indicated by 'coordinate' (coordinate sparse storage), or 'array' (dense
array storage).
If infoonly is true (default: false), only information on the size and
structure is returned from reading the header. The actual data for the
matrix elements are not parsed.
If retcoord is true (default: false), the rows, column and value vectors
are returned, if it is a sparse matrix, along with the header information.
"""
function mmread(mmstream::TranscodingStream, infoonly::Bool=false, retcoord::Bool=false)
    # Read first line
    firstline = chomp(readline(mmstream))
    tokens = split(firstline)
    if length(tokens) != 5
        throw(ParseError(string("Not enough words on first line: ", firstline)))
    end
    if tokens[1] != "%%MatrixMarket"
        throw(ParseError(string("Expected start of header `%%MatrixMarket`, got `$(tokens[1])`")))
    end
    (head1, rep, field, symm) = map(lowercase, tokens[2:5])
    if head1 != "matrix"
        throw(ParseError("Unknown MatrixMarket data type: $head1 (only \"matrix\" is supported)"))
    end

    eltype = field == "real" ? Float64 :
             field == "complex" ? ComplexF64 :
             field == "integer" ? Int64 :
             field == "pattern" ? Bool :
             throw(ParseError("Unsupported field $field (only real and complex are supported)"))

    symlabel = symm == "general" ? identity :
               symm == "symmetric" ? symmetric! :
               symm == "hermitian" ? hermitian! :
               symm == "skew-symmetric" ? skewsymmetric! :
               throw(ParseError("Unknown matrix symmetry: $symm (only general, symmetric, skew-symmetric and hermitian are supported)"))

    # Skip all comments and empty lines
    ll   = readline(mmstream)
    while length(chomp(ll))==0 || (length(ll) > 0 && ll[1] == '%')
        ll = readline(mmstream)
    end
    # Read matrix dimensions (and number of entries) from first non-comment line
    dd = map(_parseint, split(ll))
    if length(dd) < (rep == "coordinate" ? 3 : 2)
        throw(ParseError(string("Could not read in matrix dimensions from line: ", ll)))
    end
    rows = dd[1]
    cols = dd[2]
    entries = (rep == "coordinate") ? dd[3] : (rows * cols)
    infoonly && return (rows, cols, entries, rep, field, symm)

    rep == "coordinate" ||
        return symlabel(reshape([parse(Float64, readline(mmstream)) for i in 1:entries],
                                (rows,cols)))

    rr = Vector{Int}(undef, entries)
    cc = Vector{Int}(undef, entries)
    xx = Vector{eltype}(undef, entries)
    for i in 1:entries
        line = readline(mmstream)
        splits = find_splits(line, eltype == ComplexF64 ? 3 : (eltype == Bool ? 1 : 2))
        rr[i] = _parseint(line[1:splits[1]])
        cc[i] = _parseint(eltype == Bool
                          ? line[splits[1]:end]
                          : line[splits[1]:splits[2]])
        if eltype == ComplexF64
            real = parse(Float64, line[splits[2]:splits[3]])
            imag = parse(Float64, line[splits[3]:length(line)])
            xx[i] = ComplexF64(real, imag)
        elseif eltype == Bool
            xx[i] = true
        else
            xx[i] = parse(eltype, line[splits[2]:length(line)])
        end
    end
    (retcoord
     ? (rr, cc, xx, rows, cols, entries, rep, field, symm)
     : symlabel(sparse(rr, cc, xx, rows, cols)))
end


using Test

using MarketData
using TimeSeries

using MarketTechnicals

@testset "Utilities" begin

@testset "typical price" begin
    ts = collect(Date(2017, 5, 1):Day(1):Date(2017, 5, 5))
    ta = typical(TimeArray(ts, reshape(1:15, (5, 3)), [:High, :Low, :Close]))

    @test timestamp(ta)[1]   == Date(2017, 5, 1)
    @test timestamp(ta)[end] == Date(2017, 5, 5)
    @test values(ta)         == [6, 7, 8, 9, 10]
end

@testset "relu" begin
    @test MarketTechnicals.relu(10) == 10
    @test MarketTechnicals.relu(-5) == 0
    @test MarketTechnicals.relu(0)  == 0

    date = collect(Date(2017, 5, 1):Day(1):Date(2017, 5, 3))
    ta = TimeArray(date, [10, -5, 0], [:magic], "meta")
    ta = MarketTechnicals.relu(ta)
    @test timestamp(ta)[1] == Date(2017, 5, 1)
    @test values(ta)       == [10, 0, 0]
    @test colnames(ta)     == [:magic]
    @test meta(ta)         == "meta"
end

@testset "wilder_smooth" begin
    date = collect(Date(2017, 5, 1):Day(1):Date(2017, 5, 5))
    ta = TimeArray(date, [0, 10, 5, 6, 7], [:magic], "meta")
    ta = MarketTechnicals.wilder_smooth(ta, 3)
    @test timestamp(ta)[1] == Date(2017, 5, 3)
    @test values(ta)[1]    == 15.
    @test values(ta)[2]    == 16.
    @test colnames(ta)     == [:magic]
    @test meta(ta)         == "meta"

    ta = TimeArray(date, [0, 10, 5, 6, 7], [:magic], "meta")
    ta = MarketTechnicals.wilder_smooth(ta, 3, padding=true)
    @test timestamp(ta)[1]     == Date(2017, 5, 1)
    @test isnan(values(ta)[1])
    @test isnan(values(ta)[2])
    @test values(ta)[3]        == 15.
    @test values(ta)[4]        == 16.
    @test colnames(ta)         == [:magic]
    @test meta(ta)             == "meta"

    ta = TimeArray(date, [0, 4, 5, 6, 7], [:magic], "meta")
    ta = MarketTechnicals.wilder_smooth(ta, 3, dx=true)
    @test timestamp(ta)[1] == Date(2017, 5, 3)
    @test values(ta)[1]    == 3.
    @test values(ta)[2]    == 4.
    @test colnames(ta)     == [:magic]
    @test meta(ta)         == "meta"
end

end  # @testset "Utilities"


module TestTranspose
using ParallelAccelerator

@acc transpose_t(A) = A'

function test()
    A = [1. 2. 3.; 4. 5. 6.]
    B = transpose_t(A)
    return B
end

@acc function test2()
    A = [1 2 3; 4 5 6]
    B = transpose(A)
    return B
end

end

using Base.Test
println("testing transpose...")
@test_approx_eq TestTranspose.test() [1. 4.; 2. 5.; 3. 6.]
@test TestTranspose.test2()==[1 4; 2 5; 3 6]
println("Done testing transpose.")


"""
    DiscreteNonParametric(xs, ps)

A *Discrete nonparametric distribution* explicitly defines an arbitrary
probability mass function in terms of a list of real support values and their
corresponding probabilities

```julia
d = DiscreteNonParametric(xs, ps)

params(d)  # Get the parameters, i.e. (xs, ps)
support(d) # Get a sorted AbstractVector describing the support (xs) of the distribution
probs(d)   # Get a Vector of the probabilities (ps) associated with the support
```

External links

* [Probability mass function on Wikipedia](http://en.wikipedia.org/wiki/Probability_mass_function)
"""
struct DiscreteNonParametric{T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:AbstractVector{P}} <: DiscreteUnivariateDistribution
    support::Ts
    p::Ps

    function DiscreteNonParametric{T,P,Ts,Ps}(xs::Ts, ps::Ps; check_args::Bool=true) where {
            T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:AbstractVector{P}}
        check_args || return new{T,P,Ts,Ps}(xs, ps)
        @check_args(
            DiscreteNonParametric,
            (length(xs) == length(ps), "length of support and probability vector must be equal"),
            (ps, isprobvec(ps), "vector is not a probability vector"),
            (xs, allunique(xs), "support must contain only unique elements"),
        )
        sort_order = sortperm(xs)
        new{T,P,Ts,Ps}(xs[sort_order], ps[sort_order])
    end
end

DiscreteNonParametric(vs::AbstractVector{T}, ps::AbstractVector{P}; check_args::Bool=true) where {
        T<:Real,P<:Real} =
    DiscreteNonParametric{T,P,typeof(vs),typeof(ps)}(vs, ps; check_args=check_args)

Base.eltype(::Type{<:DiscreteNonParametric{T}}) where T = T

# Conversion
convert(::Type{DiscreteNonParametric{T,P,Ts,Ps}}, d::DiscreteNonParametric) where {T,P,Ts,Ps} =
    DiscreteNonParametric{T,P,Ts,Ps}(convert(Ts, support(d)), convert(Ps, probs(d)), check_args=false)

# Accessors
params(d::DiscreteNonParametric) = (d.support, d.p)

"""
    support(d::DiscreteNonParametric)

Get a sorted AbstractVector defining the support of `d`.
"""
support(d::DiscreteNonParametric) = d.support

"""
    probs(d::DiscreteNonParametric)

Get the vector of probabilities associated with the support of `d`.
"""
probs(d::DiscreteNonParametric)  = d.p

==(c1::D, c2::D) where D<:DiscreteNonParametric =
    (support(c1) == support(c2) || all(support(c1) .== support(c2))) &&
    (probs(c1) == probs(c2) || all(probs(c1) .== probs(c2)))

Base.isapprox(c1::D, c2::D) where D<:DiscreteNonParametric =
    (support(c1) ≈ support(c2) || all(support(c1) .≈ support(c2))) &&
    (probs(c1) ≈ probs(c2) || all(probs(c1) .≈ probs(c2)))

# Sampling

function rand(rng::AbstractRNG, d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    n = length(p)
    draw = rand(rng, float(eltype(p)))
    cp = p[1]
    i = 1
    while cp <= draw && i < n
        @inbounds cp += p[i +=1]
    end
    return x[i]
end

sampler(d::DiscreteNonParametric) =
    DiscreteNonParametricSampler(support(d), probs(d))

# Override the method in testutils.jl since it assumes
# an evenly-spaced integer support
get_evalsamples(d::DiscreteNonParametric, ::Float64) = support(d)

# Evaluation

pdf(d::DiscreteNonParametric) = copy(probs(d))

function pdf(d::DiscreteNonParametric, x::Real)
    s = support(d)
    idx = searchsortedfirst(s, x)
    ps = probs(d)
    if idx <= length(ps) && s[idx] == x
        return ps[idx]
    else
        return zero(eltype(ps))
    end
end
logpdf(d::DiscreteNonParametric, x::Real) = log(pdf(d, x))

function cdf(d::DiscreteNonParametric, x::Real)
    ps = probs(d)
    P = float(eltype(ps))

    # trivial cases
    x < minimum(d) && return zero(P)
    x >= maximum(d) && return one(P)
    isnan(x) && return P(NaN)

    n = length(ps)
    stop_idx = searchsortedlast(support(d), x)
    s = zero(P)
    if stop_idx < div(n, 2)
        @inbounds for i in 1:stop_idx
            s += ps[i]
        end
    else
        @inbounds for i in (stop_idx + 1):n
            s += ps[i]
        end
        s = 1 - s
    end

    return s
end

function ccdf(d::DiscreteNonParametric, x::Real)
    ps = probs(d)
    P = float(eltype(ps))

    # trivial cases
    x < minimum(d) && return one(P)
    x >= maximum(d) && return zero(P)
    isnan(x) && return P(NaN)

    n = length(ps)
    stop_idx = searchsortedlast(support(d), x)
    s = zero(P)
    if stop_idx < div(n, 2)
        @inbounds for i in 1:stop_idx
            s += ps[i]
        end
        s = 1 - s
    else
        @inbounds for i in (stop_idx + 1):n
            s += ps[i]
        end
    end

    return s
end

function quantile(d::DiscreteNonParametric, q::Real)
    0 <= q <= 1 || throw(DomainError())
    x = support(d)
    p = probs(d)
    k = length(x)
    i = 1
    cp = p[1]
    while cp < q && i < k #Note: is i < k necessary?
        i += 1
        @inbounds cp += p[i]
    end
    x[i]
end

minimum(d::DiscreteNonParametric) = first(support(d))
maximum(d::DiscreteNonParametric) = last(support(d))
insupport(d::DiscreteNonParametric, x::Real) =
    length(searchsorted(support(d), x)) > 0

mean(d::DiscreteNonParametric) = dot(probs(d), support(d))

function var(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return var(x, Weights(p, one(eltype(p))); corrected=false)
end

function skewness(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return skewness(x, Weights(p, one(eltype(p))))
end

function kurtosis(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return kurtosis(x, Weights(p, one(eltype(p))))
end

entropy(d::DiscreteNonParametric) = entropy(probs(d))
entropy(d::DiscreteNonParametric, b::Real) = entropy(probs(d), b)

function mode(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return mode(x, Weights(p, one(eltype(p))))
end
function modes(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return modes(x, Weights(p, one(eltype(p))))
end

function mgf(d::DiscreteNonParametric, t::Real)
    x, p = params(d)
    s = zero(Float64)
    for i in 1:length(x)
        s += p[i] * exp(t*x[i])
    end
    s
end

function cf(d::DiscreteNonParametric, t::Real)
    x, p = params(d)
    s = zero(Complex{Float64})
    for i in 1:length(x)
       s += p[i] * cis(t*x[i])
    end
    s
end

# Sufficient statistics

struct DiscreteNonParametricStats{T<:Real,W<:Real,Ts<:AbstractVector{T},
                                  Ws<:AbstractVector{W}} <: SufficientStats
    support::Ts
    freq::Ws
end

function suffstats(::Type{<:DiscreteNonParametric}, x::AbstractArray{T}) where {T<:Real}

    N = length(x)
    N == 0 && return DiscreteNonParametricStats(T[], Float64[])

    n = 1
    vs = Vector{T}(undef,N)
    ps = zeros(Float64, N)
    x = sort(vec(x))

    vs[1] = x[1]
    ps[1] += 1.

    xprev = x[1]
    @inbounds for i = 2:N
        xi = x[i]
        if xi != xprev
            n += 1
            vs[n] = xi
        end
        ps[n] += 1.
        xprev = xi
    end

    resize!(vs, n)
    resize!(ps, n)
    DiscreteNonParametricStats(vs, ps)

end

function suffstats(::Type{<:DiscreteNonParametric}, x::AbstractArray{T},
                   w::AbstractArray{W}; check_args::Bool=true) where {T<:Real,W<:Real}

    @check_args DiscreteNonParametric (length(x) == length(w))

    N = length(x)
    N == 0 && return DiscreteNonParametricStats(T[], W[])

    n = 1
    vs = Vector{T}(undef, N)
    ps = zeros(W, N)

    xorder = sortperm(vec(x))
    x = vec(x)[xorder]
    w = vec(w)[xorder]

    vs[1] = x[1]
    ps[1] += w[1]

    xprev = x[1]
    @inbounds for i = 2:N
        xi = x[i]
        wi = w[i]
        if xi != xprev
            n += 1
            vs[n] = xi
        end
        ps[n] += wi
        xprev = xi
    end

    resize!(vs, n)
    resize!(ps, n)
    DiscreteNonParametricStats(vs, ps)

end

# # Model fitting

fit_mle(::Type{<:DiscreteNonParametric},
        ss::DiscreteNonParametricStats{T,W,Ts,Ws}) where {T,W,Ts,Ws} =
    DiscreteNonParametric{T,W,Ts,Ws}(ss.support, normalize!(copy(ss.freq), 1), check_args=false)


"""
    latexarray{T}(arr::AbstractArray{T, 2})
Create a LaTeX array environment using [`latexraw`](@ref).

# Examples
```julia
arr = [1 2; 3 4]
latexarray(arr)
```
```math
"\\begin{equation}\n\\left[\n\\begin{array}{cc}\n1 & 2\\\\ \n3 & 4\\\\ \n\\end{array}\n\\right]\n\\end{equation}\n"
```
"""
function latexarray(arr::AbstractArray; adjustment::Symbol=:c, transpose=false, double_linebreak=false,
    starred=false, kwargs...)
    transpose && (arr = permutedims(arr))
    rows = first(size(arr))
    columns = length(size(arr)) > 1 ? size(arr)[2] : 1

    eol = double_linebreak ? " \\\\\\\\\n" : " \\\\\n"

    str = "\\begin{equation$(starred ? "*" : "")}\n"
    str *= "\\left[\n"
    str *= "\\begin{array}{" * "$(adjustment)"^columns * "}\n"

    arr = latexraw(arr; kwargs...)
    for i=1:rows, j=1:columns
        str *= arr[i,j]
        j==columns ? (str *= eol) : (str *= " & ")
    end

    str *= "\\end{array}\n"
    str *= "\\right]\n"
    str *= "\\end{equation$(starred ? "*" : "")}\n"
    latexstr = LaTeXString(str)
    COPY_TO_CLIPBOARD && clipboard(latexstr)
    return latexstr
end

latexarray(args::AbstractArray...; kwargs...) = latexarray(hcat(args...); kwargs...)
latexarray(arg::AbstractDict; kwargs...) = latexarray(collect(keys(arg)), collect(values(arg)); kwargs...)
latexarray(arg::Tuple...; kwargs...) = latexarray([collect(i) for i in arg]...; kwargs...)

function latexarray(arg::Tuple; kwargs...)
    if first(arg) isa Tuple || first(arg) isa AbstractArray
        return latexarray([collect(i) for i in arg]...; kwargs...)
    end
    return latexarray(collect(arg); kwargs...)
end


using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

@inline function exp_ps_m11(A)
    T=promote_type(eltype(A),Float64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_ps_m11!(A_copy)
end

@inline function exp_ps_m11!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=9
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    memslots5 = similar(A,T)
    memslots6 = similar(A,T)
    memslots7 = similar(A,T)
    memslots8 = similar(A,T)
    memslots9 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    value_one=ValueOne()
    # Computation order: A2 A3 A4 A5 A6 A7 P5 C4 P4 C3 P3 C2 P2 C1 P1 C0 P0
    # Computing A2 with operation: mult
    mul!(memslots2,memslots1,memslots1)
    # Computing A3 with operation: mult
    mul!(memslots3,memslots1,memslots2)
    # Computing A4 with operation: mult
    mul!(memslots4,memslots1,memslots3)
    # Computing A5 with operation: mult
    mul!(memslots5,memslots1,memslots4)
    # Computing A6 with operation: mult
    mul!(memslots6,memslots1,memslots5)
    # Computing A7 with operation: mult
    mul!(memslots7,memslots1,memslots6)
    # Computing P5 = x*I+x*A+x*A2+x*A3+x*A4+x*A5+x*A6+x*A7
    coeff1=9.67759295863189e-41
    coeff2=2.6882202662866363e-42
    coeff3=7.265460179153071e-44
    coeff4=1.911963205040282e-45
    coeff5=4.902469756513544e-47
    coeff6=1.2256174391283858e-48
    coeff7=2.9893108271424046e-50
    coeff8=7.117406731291439e-52
    memslots8 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots4 .+ coeff6.*memslots5 .+ coeff7.*memslots6 .+ coeff8.*memslots7
    mul!(memslots8,true,I*coeff1,true,true)
    # Computing C4 with operation: mult
    mul!(memslots9,memslots8,memslots7)
    # Deallocating P5 in slot 8
    # Computing P4 = x*I+x*A+x*A2+x*A3+x*A4+x*A5+x*A6+x*C4
    coeff1=3.279889237069838e-30
    coeff2=1.1309962886447716e-31
    coeff3=3.7699876288159054e-33
    coeff4=1.216125041553518e-34
    coeff5=3.8003907548547434e-36
    coeff6=1.151633562077195e-37
    coeff7=3.387157535521162e-39
    coeff8=1.0
    # Smart lincomb recycle C4
    memslots9 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots4 .+ coeff6.*memslots5 .+ coeff7.*memslots6 .+ coeff8.*memslots9
    mul!(memslots9,true,I*coeff1,true,true)
    # Computing C3 with operation: mult
    mul!(memslots8,memslots9,memslots7)
    # Deallocating P4 in slot 9
    # Computing P3 = x*I+x*A+x*A2+x*A3+x*A4+x*A5+x*A6+x*C3
    coeff1=1.9572941063391263e-20
    coeff2=8.896791392450574e-22
    coeff3=3.868170170630684e-23
    coeff4=1.6117375710961184e-24
    coeff5=6.446950284384474e-26
    coeff6=2.4795962632247976e-27
    coeff7=9.183689863795546e-29
    coeff8=1.0
    # Smart lincomb recycle C3
    memslots8 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots4 .+ coeff6.*memslots5 .+ coeff7.*memslots6 .+ coeff8.*memslots8
    mul!(memslots8,true,I*coeff1,true,true)
    # Computing C2 with operation: mult
    mul!(memslots9,memslots8,memslots7)
    # Deallocating P3 in slot 8
    # Computing P2 = x*I+x*A+x*A2+x*A3+x*A4+x*A5+x*A6+x*C2
    coeff1=1.1470745597729725e-11
    coeff2=7.647163731819816e-13
    coeff3=4.779477332387385e-14
    coeff4=2.8114572543455206e-15
    coeff5=1.5619206968586225e-16
    coeff6=8.22063524662433e-18
    coeff7=4.110317623312165e-19
    coeff8=1.0
    # Smart lincomb recycle C2
    memslots9 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots4 .+ coeff6.*memslots5 .+ coeff7.*memslots6 .+ coeff8.*memslots9
    mul!(memslots9,true,I*coeff1,true,true)
    # Computing C1 with operation: mult
    mul!(memslots8,memslots9,memslots7)
    # Deallocating P2 in slot 9
    # Computing P1 = x*I+x*A+x*A2+x*A3+x*A4+x*A5+x*A6+x*C1
    coeff1=0.0001984126984126984
    coeff2=2.48015873015873e-5
    coeff3=2.7557319223985893e-6
    coeff4=2.755731922398589e-7
    coeff5=2.505210838544172e-8
    coeff6=2.08767569878681e-9
    coeff7=1.6059043836821613e-10
    coeff8=1.0
    # Smart lincomb recycle C1
    memslots8 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots4 .+ coeff6.*memslots5 .+ coeff7.*memslots6 .+ coeff8.*memslots8
    mul!(memslots8,true,I*coeff1,true,true)
    # Computing C0 with operation: mult
    mul!(memslots9,memslots8,memslots7)
    # Deallocating P1 in slot 8
    # Deallocating A7 in slot 7
    # Computing P0 = x*I+x*A+x*A2+x*A3+x*A4+x*A5+x*A6+x*C0
    coeff1=1.0
    coeff2=1.0
    coeff3=0.5
    coeff4=0.16666666666666666
    coeff5=0.041666666666666664
    coeff6=0.008333333333333333
    coeff7=0.001388888888888889
    coeff8=1.0
    # Smart lincomb recycle A
    memslots1 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots4 .+ coeff6.*memslots5 .+ coeff7.*memslots6 .+ coeff8.*memslots9
    mul!(memslots1,true,I*coeff1,true,true)
    # Deallocating A2 in slot 2
    # Deallocating A3 in slot 3
    # Deallocating A4 in slot 4
    # Deallocating A5 in slot 5
    # Deallocating A6 in slot 6
    # Deallocating C0 in slot 9
    return memslots1 # Returning P0
end


using ReachabilityBenchmarks, MathematicalSystems, Reachability, Plots

include(@relpath "building_model.jl")
include(@relpath "building_specifications.jl")

S = building_model()
X0, options = building_specification()
options = Dict(options)

# initial value problem
problem = InitialValueProblem(S, X0)

# partition with 1D blocks
partition = [i:i for i in 1:48]

# reachability algorithms
algorithm_reach = BFFPSV18(:δ => 1e-3, :vars => [25], :partition => partition)
algorithm_check = BFFPSV18(:δ => 2e-3, :vars => [25], :partition => partition)

# check property
options[:mode] = "check"
solution = solve(problem, options; op=algorithm_check)
@assert solution.satisfied

# compute flowpipe
options[:mode] = "reach"
solution = solve(problem, options; op=algorithm_reach)

# project flowpipe to time and x₂₅
solution.options[:plot_vars] = [0, 25]
solution_proj = project(solution)

# plot projection
plot(solution_proj)
savefig("building")


module DTU_BrayamValqui_SP2021
using JuMP
using Statistics
using StatsBase
using CSV
using DataFrames
using Distributed
using SharedArrays
import MathOptInterface: AbstractOptimizer

#Data OPF
#include("Data_RTS3_OPF.jl")
#include("Data_RTS24_OPF.jl")
#Data Investment Planning-OPF
#include("Data_RTS3.jl")
include("Data_RTS24.jl")
#Functions
include("OPF.jl")
include("OPF_dual.jl")
include("solver_OPF.jl")
include("solver_OPF_dual.jl")
include("Investment_OPF_stage1.jl")
include("Investment_OPF_stage2.jl")
include("Investment_OPF_original.jl")
include("Solver_Benders.jl")
include("Solver_CentralPlanner.jl")
include("loadinputs.jl")

export OPF,solver_OPF,OPF_dual,solver_OPF_dual,Investment_OPF_stage1,Investment_OPF_stage2,Solver_Benders,loadinputs,Solver_CentralPlanner,Investment_OPF_original

end
#cd("C:\\Users\\braya\\.julia\\dev\\DTU_BrayamValqui_SP2021\\src")


### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# ╔═╡ c14d5f69-8ed3-4ad7-b094-8de94394ad8e
using Revise

# ╔═╡ 3bbf5ec7-c772-4a1a-9eab-c0f26a99e94f
using PolaronMobility

# ╔═╡ f088c85f-4d12-4f4c-8e3e-318652b68a43
using Plots

# ╔═╡ 2b5f21a3-1130-47cd-9938-f49a60048e27
using DataFrames

# ╔═╡ b38c2c58-689b-48f9-b5b1-1499651e1928
using CSV

# ╔═╡ e9685260-9ece-4c3b-b110-096a7c39857b
using QuadGK

# ╔═╡ e495c578-ac0e-46c3-8f2f-aa0541cf3574
using Optim

# ╔═╡ 7f124d90-0430-11ec-129e-bf0ec7dbda06
# ESA Results

# ╔═╡ bf2bb3d0-6f5d-412a-8f9e-b10cb30560bd
begin
	ħ = 1.05457162825e-34
	eV = 1.602176487e-19
	me = 9.10938188e-31
	kB =  1.3806504e-23
	ϵ_0 = 8.854E-12
	amu = 1.660_539_066_60e-27
	Ha = 4.35974820e-18
	Bohr = 5.29177249e-11
	c = 2.99792458e8
end

# ╔═╡ 802b2d3b-d1d8-40c7-9f34-407780385457
α_range = 1:12

# ╔═╡ 30fe032c-f911-4f7e-aa24-1fb473ecf9e0
β_range = vcat([i for i in 0.1:0.1:8.0], [100.0])

# ╔═╡ 71355321-bb21-42d5-9966-60a9438b762c
Ω_range = vcat([1e-200], [i for i in 0.01:0.1:20.01])

# ╔═╡ a1161327-a8e8-44a0-8af7-b14d437ae14c
ω = 1

# ╔═╡ f40b9810-6a0e-4627-a290-80a446d101c8
m = 1

# ╔═╡ 95538cd0-3588-423a-ab1a-95915ee60828
begin
	# Equation 31: The <|X(t) - X(s)|^{-1}> * exp(-|t-w|) effective action.
	A_integrand(x, v, w) = (w^2 * x + (v^2 - w^2) / v * (abs(1 - exp(-v * x))))^(-0.5) * exp(-x)

	A(v, w, α) = π^(-0.5) * α * v * QuadGK.quadgk(x -> A_integrand(x, v, w), BigFloat(0.0), Inf)[1]

	# Equation 33: Lowest Free energy E = -B - A where B = -3/(4v)*(v-w)^2.
	free_energy(v, w, α) = (3 / (4 * v)) * (v - w)^2 - A(v, w, α)
end

# ╔═╡ 5eaec90c-f922-496e-a388-f40a49f38b26
begin
	# Equation 62d in Hellwarth.
	Y(x, v, β) = 1 / (1 - exp(-v * β)) * (1 + exp(-v * β) - exp(-v * x) - exp(v * (x -     β)))

	# Integrand of Equation 62c in Hellwarth.
	A_integrand(x, v, w, β) =  (exp(β - x) + exp(x)) / sqrt(1e-10 + abs(w^2 * x * (1 - x / β)     + Y(x, v, β) * (v^2 - w^2) / v))

	# Equation 62c in Hellwarth.
	A(v, w, α, β) = α * v / (sqrt(π) * (exp(BigFloat(β)) - 1)) * QuadGK.quadgk(x ->       A_integrand(x, v, w, β), BigFloat(0), BigFloat(β / 2))[1]

	# Equation 62b in Hellwarth. Equation 20 in Osaka.
	B(v, w, β) = 3 / β * (log(v / w) - 1 / 2 * log(2 * π * BigFloat(β)) - log(sinh(v *     BigFloat(β) / 2) / sinh(w * BigFloat(β) / 2)))

	# Equation 62e in Hellwarth. Equation 17 in Osaka.
	C(v, w, β) = 3 / 4 * (v^2 - w^2) / v * (coth(v * BigFloat(β) / 2) - 2 / (v * β))

	# Equation 62a in Hellwarth. In paragraph below Equation 22 in Osaka; has extra 1/	  β due to different definition of A, B & C.
	function free_energy(v, w, α, β)
		setprecision(BigFloat, 32)
		a = A(v, w, α, β)
		b = B(v, w, β)
		c = C(v, w, β)
		-(a + b + c)
	end
end

# ╔═╡ 5568b92e-5a20-4771-b440-4de309b6173c
function variation(α, β; v = 0.0, w = 0.0)

    initial = [v, w]

    # Limits of the optimisation.
    lower = [1, 1]
    upper = [200, 200]

    # Osaka Free Energy function to minimise.
    f(x) = free_energy(x[1], x[2], α, β)

    # Use Optim to optimise the free energy function w.r.t v and w.
    solution = Optim.optimize(
        Optim.OnceDifferentiable(f, initial; autodiff = :forward),
        lower,
        upper,
        initial,
        Fminbox(BFGS()),
    )

    # Get v and w values that minimise the free energy.
    v, w = Optim.minimizer(solution)

	println("v: $v, w: $w")

    # Return variational parameters that minimise the free energy.
    return v, w
end

# ╔═╡ 6fc1e642-77ef-49bb-895d-a4006b563f5e
function feynamn_mass(α, v, w, m, ω)
	j(τ) = 1 + (v * ω / w^2) * (1 - w^2 / v^2) * (1 - exp(-v * τ / ω)) / τ
	integrand(τ) = exp(-τ) * τ^(1/2) * j(τ)^(-3/2)
	return m * (1 + α * (v / w)^3 / (3 * π^(1/2)) * quadgk(τ -> integrand(τ), 0.0, Inf)[1])
end

# ╔═╡ ba42efa0-8f26-4ff1-a0ac-a56c626c7525
begin
	v = Array{Float64}(undef, length(β_range), length(α_range))
	w = Array{Float64}(undef, length(β_range), length(α_range))
	FE = Array{Float64}(undef, length(β_range), length(α_range))
	M = Array{Float64}(undef, length(β_range), length(α_range))
	for α in 1:length(α_range), β in 1:length(β_range)
		println("α: $(α_range[α]), β: $(β_range[β])")
		if β == 1
			v[β, α], w[β, α] = variation(α_range[α], β_range[β]; v = 4.0, w = 2.0)
		else
			v[β, α], w[β, α] = variation(α_range[α], β_range[β]; v = v[β-1, α], w = w[β-1, α])
		end
		FE[β, α] = free_energy(v[β, α], w[β, α], α_range[α], β_range[β])
		M[β, α] = feynamn_mass(α_range[α], v[β, α], w[β, α], m, ω)
		CSV.write("v_data.csv", DataFrames.DataFrame([[0.0, [i for i in α_range]...]'; [β_range v]], :auto))
		CSV.write("w_data.csv", DataFrames.DataFrame([[0.0, [i for i in α_range]...]'; [β_range w]], :auto))
		CSV.write("F_data.csv", DataFrames.DataFrame([[0.0, [i for i in α_range]...]'; [β_range FE]], :auto))
		CSV.write("M_data.csv", DataFrames.DataFrame([[0.0, [i for i in α_range]...]'; [β_range M]], :auto))
	end
end

# ╔═╡ c6a84269-a382-4fbd-a102-8a949fe531ec
v

# ╔═╡ 0e9a3b4a-4626-4efb-8925-f493f0ab5f34
w

# ╔═╡ 441cd179-c28f-4e23-ae23-e6a65d7edb46
FE

# ╔═╡ 3b93da85-80a2-466e-b28e-dcf22b4ff0fe
M

# ╔═╡ 2405eb4e-b81c-46fb-94a3-50456b4d7512
begin
	# Integrand of (31) in Feynman I (Feynman 1955, Physical Review, "Slow 	electrons...")
	fF(τ, v, w) = (abs(w^2 * τ + (v^2 - w^2) / v * (1 - exp(- v * τ))))^-0.5 * exp(-τ)

	# (31) in Feynman I
	AF(v,w,α) = π^(-0.5) * α * v * quadgk(τ -> fF(τ, v, w), 0, Inf)[1]

	# (33) in Feynman I
	F(v,w,α) = (3 / (4 * v)) * (v - w)^2 - AF(v, w, α)
end

# ╔═╡ ae352c06-69fe-41fa-ad4c-bafb23235012
function complex_conductivity(Ω, β, α, v, w)

	println("α = $α, Ω = $Ω")

    # FHIP1962, page 1011, eqn (47c).
    R = (v^2 - w^2) / (w^2 * v)

    # FHIP1962, page 1009, eqn (35c).
    D(x) = w^2 / v^2 * (R * (1 - cos(v * x)) * coth(β * v / 2) + x^2 / β - 1im * (R * sin(v * x) + x))

    # FHIP1962, page 1009, eqn (36).
    S(x) = 2 * α / (3 * √π) * (exp(1im * x) + 2 * cos(x) / (exp(β) - 1)) / (D(x))^(3 / 2)

    # FHIP1962, page 1009, eqn (35a).
    integrand(x) = (1 - exp(1im * Ω * x)) * imag(S(x)) / Ω
    1 / (-1im * Ω + 1im * QuadGK.quadgk(x -> integrand(x), 0.0, Inf)[1])
end

# ╔═╡ 6d015de5-3c44-400c-ba60-80d645a2789c
begin
	σ = Array{ComplexF64}(undef, length(Ω_range), length(β_range))
	for α in 1:length(α_range)
		for β in 1:length(β_range), Ω in 1:length(Ω_range)
			println("α: $(α_range[α]), β: $(β_range[β]), Ω: $(Ω_range[Ω])")
			σ[Ω, β] = complex_conductivity(Ω_range[Ω], β_range[β], α_range[α], v[β, α], w[β, α])
		end
		CSV.write("conductivity_data_$α.csv", DataFrames.DataFrame([[0.0, [i for i in β_range]...]'; [Ω_range σ]], :auto))
	end
end

# ╔═╡ 8f052fb6-65bd-4c5f-b931-af616548fb79
σ

# ╔═╡ 38f4dfbc-9a01-494c-8674-0d1b1cc324d3
contour_real = Plots.contourf(Ω_range, β_range, log.(abs.(real.(σ)))', xlabel = "Ω (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12)

# ╔═╡ 22c03bea-a0d5-4291-a568-f29c21d46443
begin
	plot(β_range[1:end-1], v[1:end-1])
	plot!(β_range[1:end-1], w[1:end-1])
end

# ╔═╡ 35d0a3c8-025c-4e63-a65d-0010baf5b61c
plot(β_range[1:end-1], M[1:end-1])

# ╔═╡ Cell order:
# ╠═7f124d90-0430-11ec-129e-bf0ec7dbda06
# ╠═c14d5f69-8ed3-4ad7-b094-8de94394ad8e
# ╠═3bbf5ec7-c772-4a1a-9eab-c0f26a99e94f
# ╠═f088c85f-4d12-4f4c-8e3e-318652b68a43
# ╠═2b5f21a3-1130-47cd-9938-f49a60048e27
# ╠═b38c2c58-689b-48f9-b5b1-1499651e1928
# ╠═e9685260-9ece-4c3b-b110-096a7c39857b
# ╠═e495c578-ac0e-46c3-8f2f-aa0541cf3574
# ╠═bf2bb3d0-6f5d-412a-8f9e-b10cb30560bd
# ╠═802b2d3b-d1d8-40c7-9f34-407780385457
# ╠═30fe032c-f911-4f7e-aa24-1fb473ecf9e0
# ╠═71355321-bb21-42d5-9966-60a9438b762c
# ╠═a1161327-a8e8-44a0-8af7-b14d437ae14c
# ╠═f40b9810-6a0e-4627-a290-80a446d101c8
# ╠═95538cd0-3588-423a-ab1a-95915ee60828
# ╠═5568b92e-5a20-4771-b440-4de309b6173c
# ╠═5eaec90c-f922-496e-a388-f40a49f38b26
# ╠═6fc1e642-77ef-49bb-895d-a4006b563f5e
# ╠═ba42efa0-8f26-4ff1-a0ac-a56c626c7525
# ╠═c6a84269-a382-4fbd-a102-8a949fe531ec
# ╠═0e9a3b4a-4626-4efb-8925-f493f0ab5f34
# ╠═441cd179-c28f-4e23-ae23-e6a65d7edb46
# ╠═3b93da85-80a2-466e-b28e-dcf22b4ff0fe
# ╠═2405eb4e-b81c-46fb-94a3-50456b4d7512
# ╠═ae352c06-69fe-41fa-ad4c-bafb23235012
# ╠═6d015de5-3c44-400c-ba60-80d645a2789c
# ╠═8f052fb6-65bd-4c5f-b931-af616548fb79
# ╠═38f4dfbc-9a01-494c-8674-0d1b1cc324d3
# ╠═22c03bea-a0d5-4291-a568-f29c21d46443
# ╠═35d0a3c8-025c-4e63-a65d-0010baf5b61c


using HybridSystems
using Polyhedra
using SwitchOnSafety
using Test

include("solvers.jl")

include("jsr.jl")
include("invariant.jl")

include("../examples/run_examples.jl")


using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

@inline function denman_beavers_opt(A)
    return denman_beavers_opt!(copy(A))
end

@inline function denman_beavers_opt!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    max_memslots=7
    memslots=Vector{Matrix{T}}(undef,max_memslots)
    n=size(A,1)
    for j=1:max_memslots
        memslots[j]=Matrix{T}(undef,n,n)
    end
    # The first slots are precomputed nodes [:A]
    memslots[1]=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    value_one=ValueOne()
    # Computation order: Bb2 Ba2 B2 Ba3 Bb3 B3 Ba4 Bb4 B4 Ba5 Bb5 B5 T2k7
    # Computing Bb2 = x*I+x*A
    coeff1=1.0
    coeff2=-0.011141376535527443
    memslots[2] .= coeff2.*memslots[1]
    mul!(memslots[2],true,I*coeff1,true,true)
    # Computing Ba2 = x*I+x*A
    coeff1=1.0
    coeff2=0.7610413081657074
    memslots[3] .= coeff2.*memslots[1]
    mul!(memslots[3],true,I*coeff1,true,true)
    # Computing B2 with operation: ldiv
    memslots[4]=memslots[3]\memslots[2]
    # Deallocating Ba2 in slot 3
    # Deallocating Bb2 in slot 2
    # Computing Ba3 = x*I+x*A+x*B2
    coeff1=1.0
    coeff2=-0.07127247848266706
    coeff3=0.9541397430184381
    memslots[2] .= coeff2.*memslots[1] .+ coeff3.*memslots[4]
    mul!(memslots[2],true,I*coeff1,true,true)
    # Computing Bb3 = x*I+x*A+x*B2
    coeff1=1.0
    coeff2=-0.02575361014723645
    coeff3=-0.004282806676350282
    memslots[3] .= coeff2.*memslots[1] .+ coeff3.*memslots[4]
    mul!(memslots[3],true,I*coeff1,true,true)
    # Computing B3 with operation: ldiv
    memslots[5]=memslots[2]\memslots[3]
    # Deallocating Ba3 in slot 2
    # Deallocating Bb3 in slot 3
    # Computing Ba4 = x*I+x*A+x*B2+x*B3
    coeff1=1.0
    coeff2=0.5152503668289159
    coeff3=-0.10150239642416546
    coeff4=0.016950835467818528
    memslots[2] .= coeff2.*memslots[1] .+ coeff3.*memslots[4] .+ coeff4.*memslots[5]
    mul!(memslots[2],true,I*coeff1,true,true)
    # Computing Bb4 = x*I+x*A+x*B2+x*B3
    coeff1=1.0
    coeff2=0.03845623901890617
    coeff3=0.026744368550352622
    coeff4=-0.029043605652931636
    memslots[3] .= coeff2.*memslots[1] .+ coeff3.*memslots[4] .+ coeff4.*memslots[5]
    mul!(memslots[3],true,I*coeff1,true,true)
    # Computing B4 with operation: ldiv
    memslots[6]=memslots[2]\memslots[3]
    # Deallocating Ba4 in slot 2
    # Deallocating Bb4 in slot 3
    # Computing Ba5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=0.49421282325538585
    coeff2=0.05943417585870563
    coeff3=0.5017024742690791
    coeff4=0.5276435956563207
    coeff5=1.0
    memslots[2] .= coeff2.*memslots[1] .+ coeff3.*memslots[4] .+ coeff4.*memslots[5] .+ coeff5.*memslots[6]
    mul!(memslots[2],true,I*coeff1,true,true)
    # Computing Bb5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=1.0
    coeff2=0.12762852235945304
    coeff3=-0.0015101375211018255
    coeff4=-0.14486471812413823
    coeff5=0.01583234295751802
    memslots[3] .= coeff2.*memslots[1] .+ coeff3.*memslots[4] .+ coeff4.*memslots[5] .+ coeff5.*memslots[6]
    mul!(memslots[3],true,I*coeff1,true,true)
    # Computing B5 with operation: ldiv
    memslots[7]=memslots[2]\memslots[3]
    # Deallocating Ba5 in slot 2
    # Deallocating Bb5 in slot 3
    # Computing T2k7 = x*I+x*A+x*B2+x*B3+x*B4+x*B5
    coeff1=0.28457285753903816
    coeff2=0.11495283295141033
    coeff3=-0.0344930890633602
    coeff4=0.34606871085737384
    coeff5=-0.006442438999085179
    coeff6=1.4670072325444539
    # Smart lincomb recycle A
    memslots[1] .= coeff2.*memslots[1] .+ coeff3.*memslots[4] .+ coeff4.*memslots[5] .+ coeff5.*memslots[6] .+ coeff6.*memslots[7]
    mul!(memslots[1],true,I*coeff1,true,true)
    # Deallocating B2 in slot 4
    # Deallocating B3 in slot 5
    # Deallocating B4 in slot 6
    # Deallocating B5 in slot 7
    return memslots[1] # Returning T2k7
end


module TestDataMigration
using Test

using Catlab, Catlab.Theories, Catlab.Graphs, Catlab.CategoricalAlgebra
using Catlab.Graphs.BasicGraphs: TheoryGraph, TheoryWeightedGraph
using Catlab.Graphs.BipartiteGraphs: TheoryUndirectedBipartiteGraph

# Pullback data migration
#########################

@present TheoryDDS(FreeSchema) begin
  X::Ob
  Φ::Hom(X,X)
end

@abstract_acset_type AbstractDDS
@acset_type DDS(TheoryDDS, index=[:Φ]) <: AbstractDDS

h = Graph(3)
add_parts!(h, :E, 3, src = [1,2,3], tgt = [2,3,1])

# Identity data migration.
@test h == Graph(h, Dict(:V => :V, :E => :E),
                    Dict(:src => :src, :tgt => :tgt))

# Migrate DDS → Graph.
dds = DDS()
add_parts!(dds, :X, 3, Φ=[2,3,1])
X = TheoryDDS[:X]
@test h == Graph(dds, Dict(:V => :X, :E => :X),
                 Dict(:src => id(X), :tgt => :Φ))

h2 = copy(h)
migrate!(h2, dds, Dict(:V => :X, :E => :X),
                  Dict(:src => id(X), :tgt => :Φ))
@test h2 == ob(coproduct(h, h))

# Migrate DDS → DDS by advancing four steps.
@test dds == DDS(dds, Dict(:X => :X),
                      Dict(:Φ => [:Φ, :Φ, :Φ, :Φ]))

@present TheoryLabeledDDS <: TheoryDDS begin
  Label::AttrType
  label::Attr(X, Label)
end
@acset_type LabeledDDS(TheoryLabeledDDS, index=[:Φ, :label])

S, ϕ, Label, label = generators(TheoryLabeledDDS)
V, E, s, t, Weight, weight = generators(TheoryWeightedGraph)

ldds = LabeledDDS{Int}()
add_parts!(ldds, :X, 4, Φ=[2,3,4,1], label=[100, 101, 102, 103])

wg = WeightedGraph{Int}(4)
add_parts!(wg, :E, 4, src=[1,2,3,4], tgt=[2,3,4,1], weight=[101, 102, 103, 100])

@test wg == WeightedGraph{Int}(ldds,
  Dict(:V => :X, :E => :X),
  Dict(:src => id(S), :tgt => :Φ, :weight => [:Φ, :label]))

@test Presentation(Graph(1)) == TheoryGraph
@test Presentation(ldds) == TheoryLabeledDDS

F = FinFunctor(
  Dict(V => S, E => S),
  Dict(s => id(S), t => ϕ, weight => compose(ϕ, label)),
  TheoryWeightedGraph, TheoryLabeledDDS
)
ΔF = Delta(F, LabeledDDS{Int}, WeightedGraph{Int})
@test wg == ΔF(ldds)

idF = FinFunctor(
  Dict(X => X, Label => Label),
  Dict(ϕ => ϕ, label => label),
  TheoryLabeledDDS, TheoryLabeledDDS
)
@test ldds == LabeledDDS{Int}(ldds, idF)

# Left pushforward data migration
#################################

Σ = Sigma

V1B, V2B, EB = generators(TheoryUndirectedBipartiteGraph, :Ob)
srcB, tgtB = generators(TheoryUndirectedBipartiteGraph, :Hom)

VG, EG = generators(TheoryGraph, :Ob)
srcG, tgtG = generators(TheoryGraph, :Hom)

F = FinFunctor(
  Dict(V1B => VG, V2B => VG, EB => EG),
  Dict(srcB => srcG, tgtB => tgtG),
  TheoryUndirectedBipartiteGraph, TheoryGraph
)

idF = FinFunctor(
  Dict(VG => VG, EG => EG),
  Dict(srcG => srcG, tgtG => tgtG),
  TheoryGraph, TheoryGraph
)

ΣF = Σ(F, UndirectedBipartiteGraph, Graph)
X = UndirectedBipartiteGraph()

Y = ΣF(X)
@test nparts(Y, :V) == 0
@test nparts(Y, :E) == 0

X = @acset UndirectedBipartiteGraph begin
  V₁ = 4
  V₂ = 3
  E = 4
  src = [1,2,2,3]
  tgt = [1,1,2,3]
end

Y = ΣF(X)
@test nparts(Y, :V) == 7
@test nparts(Y, :E) == 4
@test length(Y[:src] ∩ Y[:tgt]) == 0

@test Σ(idF, Graph, Graph)(Y) == Y

@present ThSpan(FreeSchema) begin
  (L1, L2, A)::Ob
  l1::Hom(A, L1)
  l2::Hom(A, L2)
end
@acset_type Span(ThSpan, index=[:l1, :l2])

X = @acset Span begin
  L1 = 3
  L2 = 4
  A = 3
  l1 = [1,1,2]
  l2 = [1,2,3]
end

@present ThInitial(FreeSchema) begin
  I::Ob
end
@acset_type Initial(ThInitial)

L1, L2, A = generators(ThSpan, :Ob)
l1, l2 = generators(ThSpan, :Hom)
I = generators(ThInitial)[1]

bang = FinFunctor(
  Dict(L1 => I, L2 => I, A => I),
  Dict(l1 => id(I), l2 => id(I)),
  ThSpan, ThInitial
)

Σbang = Σ(bang, Span, Initial)
Y = Σbang(X)

@test nparts(Y, :I) == 4

vertex = FinFunctor(Dict(I => VG), Dict(), ThInitial, TheoryGraph)
edge = FinFunctor(Dict(I => EG), Dict(), ThInitial, TheoryGraph)

Z = Σ(vertex, Initial, Graph)(Y)
@test nparts(Z, :V) == 4
@test nparts(Z, :E) == 0

Z = Σ(edge, Initial, Graph)(Y)
@test nparts(Z, :V) == 8
@test nparts(Z, :E) == 4
@test Z[:src] ∪ Z[:tgt] == 1:8

end


using DataFrames
using Effects
using GLM
using LinearAlgebra
using StableRNGs
using StandardizedPredictors
using Statistics
using StatsModels
using Test

b0, b1, b2, b1_2 = beta = [0.0, 1.0, 1.0, -1.0]

@testset "simple" begin
    x = collect(-10:10)
    dat = (; :x => x, :y => x .* b1 .+ b0 + randn(StableRNG(42), length(x)))
    model = lm(@formula(y ~ x), dat)

    design = Dict(:x => 1:20)
    eff = effects(design, model)
    # test effect
    @test eff.y ≈ eff.x .* last(coef(model)) .+ first(coef(model))
    # test error
    pred = [1 first(eff.x)]
    # if we drop support for Julia < 1.4, this first can become only
    @test first(eff.err) ≈ first(sqrt(pred * vcov(model) * pred'))
    # test CI
    @test eff.lower ≈ eff.y - eff.err
    @test eff.upper ≈ eff.y + eff.err

    @testset "contrasts" begin
        # need to check this works when the response has also
        # had a name change
        contrasts = Dict(:x => Center(10), :y => ZScore())
        model_centered = lm(@formula(y ~ x), dat; contrasts=contrasts)

        eff_centered = effects(design, model_centered)

        μ = mean(dat.y)
        σ = std(dat.y)
        # different contrasts shouldn't affect the predictions/effects
        @test eff_centered[!, "y(centered: 0.21 scaled: 6.27)"] ≈ zscore(eff.y, µ, σ)
        @test eff_centered.lower ≈ zscore(eff.lower, µ, σ)
        @test eff_centered.upper ≈ zscore(eff.upper, µ, σ)
    end
end

@testset "multiple" begin
    x = collect(-10:10)
    z = @. sqrt(x + 100)
    dat = DataFrame(; x=x, z=z, y=randn(StableRNG(42), length(x)))

    @testset "additive" begin
        form = @formula(y ~ 1 + x + z)
        dat[!, :y] += modelmatrix(form, dat) * beta[1:(end - 1)]
        model = lm(form, dat)

        design = Dict(:x => 1:20)
        maximin(v) = mean(extrema(v))
        for typical in (maximin, mean)
            eff = effects(design, model; typical=typical)
            # test effect
            intercept = ones(length(design[:x]))
            z_typical = typical(z) * intercept
            bs = coef(model)
            @test eff.y ≈ @. bs[1] + eff.x * bs[2] + bs[3] * z_typical
            # test error
            pred = [intercept eff.x z_typical]
            @test eff.err ≈ sqrt.(diag(pred * vcov(model) * pred'))
            # test CI
            @test eff.lower ≈ eff.y - eff.err
            @test eff.upper ≈ eff.y + eff.err
        end
    end

    @testset "multiplicative" begin
        form = @formula(y ~ 1 + x * z)
        dat[!, :y] += modelmatrix(form, dat) * beta
        model = lm(form, dat)

        design = Dict(:x => 1:20)
        maximin(v) = mean(extrema(v))
        for typical in (maximin, mean)
            eff = effects(design, model; typical=typical)
            # test effect
            intercept = ones(length(design[:x]))
            z_typical = typical(z) * intercept
            zx_typical = z_typical .* design[:x]
            bs = coef(model)
            @test eff.y ≈ @. bs[1] + eff.x * bs[2] + bs[3] * z_typical + bs[4] * zx_typical
            # test error
            pred = [intercept eff.x z_typical zx_typical]
            @test eff.err ≈ sqrt.(diag(pred * vcov(model) * pred'))
            # test CI
            @test eff.lower ≈ eff.y - eff.err
            @test eff.upper ≈ eff.y + eff.err
        end
    end
end


function determineK(L₁₁,L₂₂,Lₘ)
    K = Lₘ/√(L₁₁*L₂₂)
end

function determineCancelledCenter(L₁₁,L₂₂)
    Lₘ = L₁₁
    Lₐ = Lₘ
    Lᵦ = L₂₂-Lₘ
    Lc = L₁₁-Lₘ
    K = Lₘ/√(L₁₁*L₂₂)
    return K,Lₐ,Lᵦ,Lc
end


using DEBmicroTrait, Test

n_polymers = 5
n_enzymes  = 6
P          = rand(n_polymers)
E          = rand(n_enzymes)
#
Χ_0        = [10,10,10,10,10]
ρ_P        = [1.35, 1.5, 1.4, 0.4, 1.0]
R_E        = [3.29e-9, 4e-9, 4e-9, 4.4e-9, 5e-9,3.9e-9]
V_E        = [5, 0.24, 4,0.19, 3,4]
MW_E       = [48e3, 65e3, 59.7e3, 70e3, 50e3, 52e3]
D_E        = DEBmicroTrait.aqueous_diffusivity(MW_E)
#
R_P        = 20e-3*ones(n_polymers)
#
α_kin_P    = DEBmicroTrait.binding_sites(ρ_P, R_E)
K_EP_0     = DEBmicroTrait.polymer_affinity(R_E, R_P, V_E, D_E)

#######################################################################
#                     test paper values
# enzyme characteristics
R_E        = [3.29e-9]
V_E        = [5.0]
MW_E       = [48e3]
#D_E        = DEBmicroTrait.aqueous_diffusivity(MW_E)
D_E        = [1e-10]
# Polymer characteristics
ρ_P        = [1.35]
R_P        = [1].*1e-3
#
K_EP_0     = DEBmicroTrait.polymer_affinity(R_E, R_P, V_E, D_E)
@test K_EP_0 ≈ [0.244] atol=0.001
#
k_max_RP   = DEBmicroTrait.max_hydrolyis_rate(R_E, R_P, V_E)
#
R_P        = [1, 10, 20, 30, 40, 50].*1e-3
K_EP_0     = DEBmicroTrait.polymer_affinity(R_E, R_P, V_E, D_E)

V_E        = [42.3, 9.2, 0.41, 34.1, 31.5, 113.5, 128.6, 0.18, 112.2, 0.092, 4.6, 32896, 165, 92.8, 152.13, 69, 0.02, 54, 7590, 16.0, 18.6,
              300.3, 0.39, 0.39]

k_max_RP   = DEBmicroTrait.max_hydrolyis_rate(R_E, R_P, V_E)

kmaxs      = zeros(24,50)
R_P        = collect(range(1,50, length=50))

for i in 1:24
    for j in 1:50
        kmaxs[i,j] = DEBmicroTrait.max_hydrolyis_rate(R_E, [R_P[j]], [V_E[i]])[1]
    end
end

using JLD
save("/Users/glmarschmann/.julia/dev/DEBmicroTrait/files/cellulase_depolymerization.jld", "kmaxs", kmaxs)

using Statistics

median(V_E)


@testset "'MESS' ............................... " begin

  ss = SampleSpace(5:20, 100)

  params = MESS(
      ss, .2, .4, alpha = .05, beta = .2, pess = .4
  )

  @test null(params) == .2

  @test mtoer(params) == .05

  @test samplespace(params) == ss

  @test maxsamplesize(params) == 100

  @test minconditionalpower(params) == 0.0

  @test typeof(params) == MESS{typeof(1.0)}

  params = MESS(
      ss, .2, .4, minconditionalpower = .7
  )
  @test minconditionalpower(params) == 0.7

  design = optimaldesign(10, params, solver)

  @test score(design) < 26.5

end


using Optim, StatsFuns

function quap(loglikf, x0, lower, upper)
  res = optimize(loglikf, lower, upper, x0, Fminbox(GradientDescent()))
  (Optim.minimizer(res), res)
end


## Matrix helpers
#===========================================================================#

function wake_panels(panels, bound, num)
    lastx, lasty = (p2 ∘ last)(panels)
    bounds = range(lastx, bound, length = num)
    @. WakePanel2D(SVector(bounds[1:end-1], lasty), SVector(bounds[2:end], lasty))
end

ψp(str1, str2, panel_i :: AbstractPanel2D, point) = panel_scalar(vortex_stream_plus, str1, panel_i, first(point), last(point)) + panel_scalar(vortex_stream_minus, str2, panel_i, first(point), last(point))

ψm(str1, str2, panel_i :: AbstractPanel2D, point) = panel_scalar(vortex_stream_plus, str1, panel_i, first(point), last(point)) - panel_scalar(vortex_stream_minus, str2, panel_i, first(point), last(point))

trailing_edge(str1, str2, panel_i :: AbstractPanel2D, point) = panel_scalar(source_stream, str1, panel_i, first(point), last(point)) + panel_scalar(vortex_stream_plus, str2, panel_i, first(point), last(point))

source(str, panel_i, point) = panel_scalar(source_stream, str, panel_i, first(point), last(point))

trailing_edge_source(str, panel_i, point) = panel_scalar(source_stream_trailing_edge, str, panel_i, first(point), last(point))

## Matrix construction
#===========================================================================#

# Need to check singularity conditions
ψp_matrix(panels :: Vector{<: AbstractPanel2D}, points) = [ ψp(1., 1., panel_j, point_i) for (point_i, panel_j) in product(points, panels) ]

ψm_matrix(panels :: Vector{<: AbstractPanel2D}, points) = [ ψm(1., 1., panel_j, point_i) for (point_i, panel_j) in product(points, panels) ]

kutta_condition(points) = [ 1 zeros(length(points) - 2)' 1 ]

streamline(f_pan, l_pan) = (p2(l_pan) - p1(l_pan) + p1(f_pan) - p2(f_pan)) / 2

trailing_edge_influence(s, t, te_panel, pt) = ifelse(p1(te_panel) == pt, 0.0, trailing_edge_source(1., te_panel, pt) * norm(t × s) + ψp(1., 1., te_panel, pt) * (abs ∘ dot)(t, s))

function trailing_edge(panels :: Vector{<: AbstractPanel2D}, pts)
    te_panel  = trailing_edge_panel(panels)
    s         = p2(te_panel) - p1(te_panel)
    t         = streamline(panels[1], panels[end])
    te_vector = trailing_edge_influence.(Ref(s), Ref(t), Ref(te_panel), pts)

    te_panel, te_vector
end

function influence_matrix(panels :: Vector{<: AbstractPanel2D})
    rev_panels  = reverse_panel.(panels)[end:-1:1]
    pts         = panel_points(rev_panels)
    num_pts     = length(pts)

    ψM      = [ ψm_matrix(rev_panels, pts) zeros(num_pts) ]
    ψP      = [ zeros(num_pts) ψp_matrix(rev_panels, pts) ]
    ψM, ψP

    ψ_matrix = ψM .+ ψP

    # Trailing edge setup
    te_panel, te_vector = trailing_edge(rev_panels, pts)
    te_matrix           = [ te_vector zeros(num_pts, num_pts - 2) -te_vector ]
    r_te                = panel_length(te_panel)

    inf_matrix          = ψ_matrix + te_matrix

    ψ0_vector           = fill(-1, length(pts))

    [     inf_matrix        ψ0_vector  ;
      kutta_condition(pts)      0      ]
end

boundary(colpoints, uni) = - [ [ stream(uni, x, y) for (x, y) in colpoints ]; 0 ]

# function aerodynamic_coefficients(panels :: Vector{<: AbstractPanel2D}, γs, u)
#       speed = norm(u)

#       vels  = sum(γs)
#       Δrs   = midpair_map(distance, panels)
#       cps   = @. pressure_coefficient(speed, vels)
#       cls   = @. lift_coefficient(cps, Δrs / 2, panel_angle(panels))

#       cps, cls
# end

# source_matrix(panels_1 :: Vector{<: AbstractPanel2D}, panels_2 :: Vector{<: AbstractPanel2D}) = [ source_influence(panel_j, panel_i) for panel_i ∈ panels_1, panel_j ∈ panels_2 ]

# source_strengths(panels :: Vector{<: AbstractPanel2D}, u) = dot.(Ref(u), panel_normal.(panels))

# boundary_vector(panels :: Vector{<: AbstractPanel2D}, u) = [ - source_matrix(panels, panels) * source_strengths(panels, u); 0 ]

# function boundary_vector(panels :: Vector{<: AbstractPanel2D}, wakes :: Vector{<: AbstractPanel2D}, u)
#       source_panels = [ panels; wakes ]
#       [ - source_matrix(panels, source_panels) * source_strengths(source_panels, u); 0 ]
# end


function _chol!(A::StridedMatrix{<:BlasFloat}, ::Type{UpperTriangular})
    C, info = LAPACK.potrf!('U', A)
    return UpperTriangular(C), info
end
function _chol!(A::StridedMatrix{<:BlasFloat}, ::Type{LowerTriangular})
    C, info = LAPACK.potrf!('L', A)
    return LowerTriangular(C), info
end


#Common functions and data structures
include("common.jl")
include("orthogonalize.jl")

# Hessenberg problem solver
include("hessenberg.jl")

#Stationary solvers
include("stationary.jl")

#Conjugate gradients
include("cg.jl")

#BiCGStab(l)
include("bicgstabl.jl")

#MINRES
include("minres.jl")

#GMRES
include("gmres.jl")

#IDRS
include("idrs.jl")

# QMR
include("qmr.jl")

#Chebyshev
include("chebyshev.jl")

#Simple Eigensolvers
include("simple_eigensolvers.jl")
include("lobpcg.jl")

#Golub-Kahan-Lanczos singular values computation
include("svdl.jl")

include("lsqr.jl")
include("lsmr.jl")

#History data structure
include("history.jl")

# Test deprecations
include("deprecations.jl")


#import recombination
include("recombination.jl")

T=300           # Temperature; Kelvin

kb=8.6173324E-5 # in units of eV
β=1/(kb*T)      # Thermodynamic Beta; units

#N=1000
#pot=0.13*randn(N) # Normal distribution
#pot=0.13*sin(range(0,N)) # Sinusoidal variations

N,pot = recombination.starrynight_read_potential(ARGS[1])

function potential_statistics(N,pot)
    println("Boltzmann statistics...")
    recombination.calc_recombination(N, pot, En -> exp.(-β*En) )
    println("Fermi-Dirac statistics...")
    recombination.calc_recombination(N, pot, En -> 1./(exp.(β*En) + 1.0) )
    println("Fermi-Dirac, chem pot = 0.8 eV...")
    recombination.calc_recombination(N, pot, En -> 1./(exp.(β*(En+0.8)) + 1.0) )
    recombination.calc_mobility(N,pot)
end

potential_statistics(N,pot)

# Apply knowledge of polaron Gaussian wavefunction by Gaussian blurring the underlying potential

using Images
n=50# Terrible hard-wired hack
pot=reshape(pot,n,n,n)
#print(pot)

## Make up a pot to test Gaussian filter
#n=10
#pot=zeros(n,n,n)
#pot[n/2,n/2,n/2]=1.0

const l=49 # size of Gaussian Kernel; make this larger to avoid edge effects
# defaults to: l = 4*ceil(Int,σ)+1) ; which is why it had a discretisation error
# c.f.
# https://github.com/JuliaImages/ImageFiltering.jl/blob/f9195cc14f24dc5bcb3d45adb1608a7d7d103018/src/kernelfactors.jl#L267

#print(pot)
for i=0:0.05:10
    println(STDERR,"Gaussian filter with Sigma=$i")
    R=recombination.calc_recombination(N,
                                       imfilter(pot,reflect(Kernel.gaussian([i,i,i], [l,l,l])), "circular"),
                                       En -> 1./(exp.(β*(En+0.8)) + 1.0) )
    # circular gives periodic boundaries - actually not sure of this, may have broken.
    # reflect() makes it operation a convolution
    # (i,i,i) define the standard deviation of the Gaussian
    # (l,l,l) are the size of the Kernel (stencil); set larger than default to avoid discretisation error
    println("$i $R") # STDOUT, for later printing.
end


function foo()
    lines = readlines(open("data/6.txt", "r"))
    total_distance = 10000
    points = zeros(Int, length(lines), 2)
    maxx = 0
    for (i, line) in enumerate(lines)
        point = parse.(Int, split(line, ','))
        if any(point .> maxx)
            maxx = maximum(point)
        end
        points[i, :] = point
    end

    stop = maxx + 10
    start = -10
    area = 0
    for i in start:stop
        for j in start:stop
            dist = abs.(points[:, 1] .- i) .+ abs.(points[:, 2] .- j)
            area += sum(dist) < total_distance
        end
    end
    return area
end

println(foo())


@testset "Turing" begin
    @testset "GP regression" begin
        k = SqExponentialKernel()
        y = randn(3)
        X = randn(3, 1)
        x = [rand(1) for _ in 1:3]

        Turing.@model function GPRegression(y, X)
            # Priors.
            α ~ LogNormal(0.0, 0.1)
            ρ ~ LogNormal(0.0, 1.0)
            σ² ~ LogNormal(0.0, 1.0)

            # Realized covariance function
            kernel = α * (SqExponentialKernel() ∘ ScaleTransform(1 / ρ))
            f = GP(kernel)

            # Sampling Distribution.
            return y ~ f(X, σ²)
        end
        # Test for matrices
        m = GPRegression(y, RowVecs(X))
        @test length(Turing.sample(m, Turing.HMC(0.5, 1), 5)) == 5

        # Test for vectors of vector
        m = GPRegression(y, x)
        @test length(Turing.sample(m, Turing.HMC(0.5, 1), 5)) == 5
    end
    @testset "latent GP regression" begin
        X = randn(3, 1)
        x = [rand(1) for _ in 1:3]
        y = rand.(Poisson.(exp.(randn(3))))

        Turing.@model function latent_gp_regression(y, X)
            f = GP(Matern32Kernel())
            u ~ f(X)
            λ = exp.(u)
            return y .~ Poisson.(λ)
        end
        m = latent_gp_regression(y, RowVecs(X))
        @test length(Turing.sample(m, Turing.NUTS(), 5)) == 5

        # Test for vectors of vector
        m = latent_gp_regression(y, x)
        @test length(Turing.sample(m, Turing.NUTS(), 5)) == 5
    end
end


#-------------------------------------------------------------------------------------------------
#--
#-- Ususal libraries
#--
using CSV, Dates
using DataFrames, DataFramesMeta
using Plots, PyPlot
using DifferentialEquations

# The dynamics of the epidemy is a function that mutates its argument with a precise signature
# Don't pay too much attetion to the print debugs/

function epiDynamics!(dP, P, params, t)
    c = 0
    S = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    E = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    I = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    J = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    H = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    C = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    R = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    F = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    D = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    K = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1
    L = P[c*nAgeGroup + 1:c*nAgeGroup + nAgeGroup]; c += 1

    (BED, ICU) = P[nAgeGroup*COMPARTMENTS_N:end]
    BED = BED[1]
    ICU = ICU[1]

    r₀, tₗ, tᵢ, tₕ, tᵤ, tᵣ,
        γₑ, γᵢ, γⱼ, γₖ, γᵣ,
        δₖ, δₗ, δᵤ ,
        mitigation,
        BED_max, ICU_max, Population = params

    ####################################
    # Arrows reflecting epidemiology - Check signs (just in case)
    # EI means from  E to I
    # EI > 0  means flow from E to I (E goes down, I goes up)
    EI = ones(nAgeGroup) .* E / tₗ;  EI = max.(EI, 0.0001);     IE = -EI
    IJ = mₐ              .* I / tᵢ;  IJ = max.(IJ, 0.0001);     JI = -IJ
    JK = cₐ              .* J / tₕ;  JK = max.(JK, 0.0001);     KJ = -JK
    HL = cₐ              .* H / tₕ;  HL = max.(HL, 0.0001);     LH = -HL

    # Asymptomatic recovery arrows
    IR = (1 .- mₐ)       .* I / tᵢ;  IR = max.(IR, 0.0001);     RI = -IR
    JR = (1 .- cₐ)       .* J / tₕ;  JR = max.(JR, 0.0001);     RJ = -JR
    HR = (1 .- cₐ)       .* H / tₕ;  HR = max.(HR, 0.0001);     RH = -HR
    KR = (1 .- δₖ .* fₐ) .* K / tᵤ;  KR = max.(KR, 0.0001);     RK = -KR
    LR = (1 .- δₗ .* fₐ) .* L / tᵤ;  LR = max.(LR, 0.0001);     RL = -LR
    CR = (1 .- δᵤ .* fₐ) .* C / tᵤ;  CR = max.(CR, 0.0001);     RC = -CR

    # Full recovery arrows
    RF = ones(nAgeGroup) .* R / tᵣ;  RF = max.(RF, 0.0001);     FR = -RF

    # Deaths
    KD =       δₖ .* fₐ  .* K / tᵤ;  KD = max.(KD, 0.0001);     DK = -KD
    LD =       δₗ .* fₐ  .* L / tᵤ;  LD = max.(LD, 0.0001);     DL = -LD
    CD =       δᵤ .* fₐ  .* C / tᵤ;  CD = max.(CD, 0.0001);     DC = -CD

    ####################################
    # Bed transfers

    ####### Step 1:
    # Decrease in ICU usage after 14 days (recall that CD and CR are vectors over the age groups)
    dICU = -(sum(CD) + sum(CR))
    dICU = ensurePositive(dICU, ICU)
    if PRINT_DEBUG @show dICU end

    # ICU beds available
    ICU_free = max(0.0, ICU_max - (ICU + dICU))
    if PRINT_DEBUG
        @show ICU_max
        @show ICU
        @show dICU
        @show ICU_free
    end

    # Move as many patients as possible from $L$ to $C$ in proportion of each group
    ICU_transfer = min(sum(L), ICU_free)
    LC = ICU_transfer / sum(L) .* L;                            CL = -LC
    if PRINT_DEBUG
        @show ICU_transfer
        @show sum(LC)
    end

    # Overall change in ICU bed becomes
    dICU = dICU + ICU_transfer
    dICU = ensurePositive(dICU, ICU)

    # And some normal beds are freed
    dBED = -ICU_transfer
    dBED = ensurePositive(dBED, BED)
    # println(dBED); println(BED); println(floor(BED_max)); println(ICU_transfer)

    ####### Step 2:
    # Beds available
    BED_free = max(0.0, BED_max - (BED + dBED))

    # Move as many patients as possible from $K$ to $L$ in proportion of each group
    BED_transfer = min(sum(K), BED_free)
    KL = BED_transfer / sum(K) .* K;                            LK = -KL

    # Overall change in normal bed becomes
    dBED = dBED + BED_transfer
    dBED = ensurePositive(dBED, BED)

    ####### Step 3:
    # Beds available
    BED_free = max(0.0, BED_max - (BED + dBED))

    # Move as many patients as possible from $J$ to $H$ in proportion of each group
    BED_transfer = min(sum(J), BED_free)
    JH = BED_transfer / sum(J) .* J;                            HJ = -JH

    # Overall change in ICU bed becomes
    dBED = dBED + BED_transfer
    dBED = ensurePositive(dBED, BED)

    ####################################
    # Sum of all flows + Check never negative compartment

    # Susceptible
    # Calculation of β
    mitigationRatio = getCurrentRatio(t; start = BASE_DAYS, schedule = mitigation)
    β = (r₀ * mitigationRatio) .* zₐ

    dS = - (sum(S)  / Population) .* β .* (γₑ.*E + γᵢ.*I + γⱼ.*J + γₖ.*K + γᵣ.*R)
    dS = min.(-0.0001, dS)
    dS = ensurePositive(dS, S)

    # Exposed
    dE = -dS + IE
    dE = ensurePositive(dE, E)

    # Infected.
    dI = EI + JI + RI
    dI = ensurePositive(dI, I)
    # Infected no hospital
    dJ = IJ + HJ + KJ + RJ
    dJ = ensurePositive(dJ, J)

    # Infected in hospital
    dH = JH + LH + RH
    dH = ensurePositive(dH, H)

    # Critical no hospital
    dK = JK + LK + DK + RK
    dK = ensurePositive(dK, K)

    # Critical in hospital
    dL = KL + HL + CL + DL + RL
    dL = ensurePositive(dL, L)

    # Critical in ICU
    dC = LC + DC + RC
    dC = ensurePositive(dC, C)

    # Asymptomatic recovery (can only increase)
    dR = IR + JR + HR + KR + LR + CR + FR
    dR = ensurePositive(dR, R)

    # Full recovery (can only increase)
    dF = RF
    dF = max.(dF, 0.0001)

    # Dead (can only increase)
    dD = KD + LD + CD
    dD = max.(dD, 0.0001)

    # Vector change of population and update in place
    result = vcat(dS, dE, dI, dJ, dH, dC, dR, dF, dD, dK, dL, [dBED], [dICU])

    for i = 1:length(result)
        dP[i] = result[i]
    end

    if PRINT_DEBUG
        pf = map(sum, [dS, dE, dI, dJ, dH, dK, dL, dC, dD, dR, dF, [dBED], [BED], [BED_max], [dICU], [ICU], [ICU_max],
                       dS + dE + dI + dJ + dH + dK + dL + dC + dD + dR + dF])
        @show pf
        #readline()
    end
end

function calculateSolution(country, diseaseparams, countryparams;
                           finalDate::Union{Nothing, Date} = nothing)

    # Deconstruct the parameters
    r₀, tₗ, tᵢ, tₕ, tᵤ, tᵣ,
        γₑ, γᵢ, γⱼ, γₖ, γᵣ,
        δₖ, δₗ, δᵤ = diseaseparams

    modelStart, infectedM, infectiousM, mv0, mv1, mv2, mv3, mv4, mv5, mv6, mv7, mv8, mv9 = countryparams

    mitigation = [(0, mv0), (7, mv1), (14, mv2), (21, mv3), (35, mv4),
                  (49, mv5), (63, mv6), (77, mv7), (91, mv8), (105, mv9)]

    # First date should the date of the last death reported
    startRecordDate = first(countryData[country][:cases].time)

    # The model always starts at 0 by definition.
    startModelDay = 0

    # Final date should the date of the last death reported unless given
    # The date is shifted from actual time to model time
    if finalDate == nothing
        endRecordDate = last(countryData[country][:cases].time)
    else
        endRecordDate = finalDate
    end
    endModelDay =  timeReal2Model(endRecordDate, modelStart)

    tSpan = (startModelDay, endModelDay)

    ## Country-specific constants
    # Those are constants which cannot be changed to improve the model.

    BED_max = countryData[country][:hospital_capacity]
    ICU_max = countryData[country][:ICU_capacity]

    Age_Pyramid = Array{Float64}(countryData[country][:age_distribution])
    Age_Pyramid_frac = Age_Pyramid ./ sum(Age_Pyramid)
    Population = sum(Age_Pyramid)

    #TotalConfirmedAtStart = @where(countryData[country][:cases], :time .== startDate)[!, :cases][1]
    #ConfirmedAtStart = TotalConfirmedAtStart .* Age_Pyramid_frac

    # TotalDeathsAtStart = @where(countryData[country][:cases], :time .== startDate)[!, :deaths][1]
    TotalDeathsAtStart = DEATH_AT_MODEL_START
    DeathsAtStart = TotalDeathsAtStart .* Age_Pyramid_frac

    ## Parameter vector
    # Those are parameters which can be changed to improve the model
    TotalInfected = infectedM * TotalDeathsAtStart
    InfectedAtStart = TotalInfected .* Age_Pyramid_frac

    TotalInfectious = infectiousM .* TotalDeathsAtStart
    InfectiousAtStart = TotalInfectious .* Age_Pyramid_frac

    model_params = [r₀,
                   tₗ, tᵢ, tₕ, tᵤ, tᵣ,
                   γₑ, γᵢ, γⱼ, γₖ, γᵣ,
                   δₖ, δₗ, δᵤ,
                   mitigation,
                   BED_max,
                   ICU_max,
                   Population]

    ## Compartment vector
    # Note that values are initialised at 1 to avoid division by zero
    S0 = Age_Pyramid .- InfectedAtStart .- InfectiousAtStart .- DeathsAtStart
    E0 = InfectedAtStart
    I0 = InfectiousAtStart
    J0 = 0.0001 .* ones(Float64, nAgeGroup)
    H0 = 0.0001 .* ones(Float64, nAgeGroup)
    C0 = 0.0001 .* ones(Float64, nAgeGroup)
    R0 = 0.0001 .* ones(Float64, nAgeGroup)
    F0 = 0.0001 .* ones(Float64, nAgeGroup)
    D0 = DeathsAtStart
    K0 = 0.0001 .* ones(Float64, nAgeGroup)
    L0 = 0.0001 .* ones(Float64, nAgeGroup)

    # Everybody confirmed is in hospital. Assume 1 ICU bed to stay away from zeros.
    BED = [TotalInfectious]
    ICU = [1.0]

    P0 = vcat(S0, E0, I0, J0, H0, C0, R0, F0, D0, K0, L0, BED, ICU)

    # Differential equation solver
    model = ODEProblem(epiDynamics!, P0, tSpan, model_params)

    # Note: progress steps might be too quick to see!
    sol = solve(model, Tsit5(); progress = false)

    return sol
end

# Calculate the forecast total individual in a given compartment.
# The calculation is performed at each time steps of the solutions generated by the model
function calculateTotalCompartment(sol, comp::String)
    # The solutions in 'soll are presented as a vector of vectors:
    #  - it is a vector of size the number of timesteps
    #  - each element of the vector is a vector of all the variables
    nSteps = length(sol.t);
    nVars  = length(sol.u[1]);

    # Creates an Array (variables x time steps)
    solMat = reduce(hcat, sol.u)

    # Select the rows of Dx and sums to have total deaths at each time period
    total = sum(solMat[compIndex(comp):compIndex(comp) + nAgeGroup - 1, :]; dims = 1)

    return total[:]
end

# Calculate the forecast total in a given compartment summed across al ages.
# The calculation is performed at each time steps
# of the solutions generated by the model
function getSummedCompartment(sol, C::String)
    # The solutions in 'sol' are presented as a vector of vectors:
    #  - it is a vector of size the number of timesteps
    #  - each element of the vector is a vector of all the variables
    # Creates an Array (variables x time steps)
    solMat = reduce(hcat, sol.u)

    # Select the rows of Dx and sums to have total deaths at each time period
    return sum(solMat[compIndex(C):compIndex(C) + nAgeGroup - 1, :]; dims=1)
end

# Calculate the forecast total deaths. The calculation is performed at each time steps
# of the solutions generated by the model
function getVariableForecast(sol, V::String)
    # Creates an Array (variables x time steps)
    solMat = reduce(hcat, sol.u)

    # Select the rows of Dx and sums to have total deaths at each time period
    return solMat[varIndex(V), :]
end

# Calculate the total deaths forecast by the model on the dates for which there is an actual
# record.
# calculateTotalDeaths calculates for each time step on model time
# forecastOnActualDates calculates on the time the actual dates fall on when converted
#   to model time.
function forecastCompartmentOnActualDates(sol, C::String, country::String)
    # The model always starts at time 0 to deals with the mitigation ratio (which starts at 0)
    # The 'modelStart' parameter deals with shifting the model compared to the actual record to go
    # from 'model time' to 'real time'

    # Generate the forecast deaths on the dates of the model (in model time)
    c = getSummedCompartment(sol, C)

    # Get the start and final dates of the deaths record (in real time)
    startRecordDate = first(countryData[country][:cases].time)
    finalRecordDate = last(countryData[country][:cases].time)

    startRecordDay = date2days(startRecordDate)
    finalRecordDay = date2days(finalRecordDate)

    # Translates the dates into model time
    startModelDay = timeReal2Model(startRecordDate, country)
    finalModelDay = timeReal2Model(finalRecordDate, country)

    # Make a linear approximation of the forecast to match the actual days

    # -- How many steps to forecast?
    l = length(countryData[country][:cases].time)

    return [ linearInterpolation(t, sol.t, c) for t in range(startModelDay, stop = finalModelDay, length = l)]
end

function forecastVariableOnActualDates(sol, C::String, country::String)
    # The model always starts at time 0 to deals with the mitigation ratio (which starts at 0)
    # The 'modelStart' parameter deals with shifting the model compared to the actual record to go
    # from 'model time' to 'real time'

    # Generate the forecast deaths on the dates of the model (in model time)
    v = getVariableForecast(sol, C)

    # Get the start and final dates of the deaths record (in real time)
    startRecordDate = first(countryData[country][:cases].time)
    finalRecordDate = last(countryData[country][:cases].time)

    startRecordDay = date2days(startRecordDate)
    finalRecordDay = date2days(finalRecordDate)

    # Translates the dates into model time
    startModelDay = timeReal2Model(startRecordDate, country)
    finalModelDay = timeReal2Model(finalRecordDate, country)

    # Make a linear approximation of the forecast to match the actual days

    # -- How many steps to forecast?
    l = length(countryData[country][:cases].time)

    return [ linearInterpolation(t, sol.t, v) for t in range(startModelDay, stop = finalModelDay, length = l)]
end

# Calculate the total deaths forecast by the model on the dates for which there is an actual
# record.
# calculateTotalDeaths calculates for each time step on model time
# forecastOnActualDates calculates on the time the actual dates fall on when converted
#   to model time.
function forecastOnActualDates(sol, country::String)
    # The model always starts at time 0 to deals with the mitigation ratio (which starts at 0) and
    # assumes 5 deaths at that date.
    # The 'modelStart' parameter deals with shifting the model compared to the actual record to go
    # from 'model time' to 'real time'

    # Generate the forecast deaths on the dates of the model (in model time)
    forecastDeaths = calculateTotalDeaths(sol)

    # Get the start and final dates of the deaths record (in real time)
    startRecordDate = first(countryData[country][:cases].time)
    finalRecordDate = last(countryData[country][:cases].time)

    startRecordDay = date2days(startRecordDate)
    finalRecordDay = date2days(finalRecordDate)

    # Translates the dates into model time
    startModelDay = timeReal2Model(startRecordDate, country)
    finalModelDay = timeReal2Model(finalRecordDate, country)

    # Make a linear approximation of the forecast to match the actual days

    # -- How many steps to forecast?
    l = length(countryData[country][:cases].time)

    forecast = [ linearInterpolation(t, sol.t, forecastDeaths)
                     for t in range(startModelDay, stop = finalModelDay, length = l)]

    # Select the actual deaths record on those dates
    relevantCases = countryData[country][:cases][:, :deaths]
    relevantCases = convert(Array, relevantCases)

    return [[ t for t in startRecordDay:finalRecordDay],
            relevantCases,
            forecast]
end


#julia
using Makie
scene = Scene()
N = 32
function xy_data(x,y,i, N)
    x = ((x/N)-0.5f0)*i
    y = ((y/N)-0.5f0)*i
    r = sqrt(x*x + y*y)
    res = Float32(sin(r)/r)
    isnan(res) ? 1f0 : res
end

surf_func(i) = [Float32(xy_data(x, y, i, 32)) + 0.5 for x=1:32, y=1:32]

z = surf_func(20)
range = linspace(0, 3, N)
surf = surface(range, range, z, colormap = :Spectral)
wf = wireframe(range, range, surf[:z] .+ 1.0,
    linewidth = 2f0, color = lift_node(x-> x[5], surf[:colormap])
)
axis(linspace(0, 3, 4), linspace(0, 3, 4), linspace(0, 3, 4))
center!(scene)

wf[:linewidth] = 1
surf[:colormap] = :YlGnBu
io = VideoStream(scene)
for i in linspace(0, 60, 200)
    surf[:z] = surf_func(i)
    recordframe!(io)
end
io


export AbBasis, LinearUnit, Relu, Requ

abstract type AbBasis{F<:AbstractFloat} end
abstract type AbUniDirBasis{F}<:AbBasis{F} end
calc{F<:AbUniDirBasis}(::Type{F}, X, u) = calc(F, X*u)

################################################################# Linear
struct LinearUnit{F} <: AbUniDirBasis{F}
  u::Vector{F}
end
calc{F}(::Type{LinearUnit{F}}, x) = x
derv{F}(::Type{LinearUnit{F}}, x) = ones(x)
calc{F}(b::LinearUnit{F}, X) = X*b.u

################################################################# Signed
abstract type AbSignedUniDirBasis{F}<:AbUniDirBasis{F} end
calc(b::B, X) where B<:AbSignedUniDirBasis = b.sgn * calc(B, X*b.u)
derv(b::B, X) where B<:AbSignedUniDirBasis = b.sgn * derv.(B, X*b.u)

################################################################# Relu
mutable struct Relu{F} <: AbSignedUniDirBasis{F}
  u::Vector{F}
  sgn::F
end
Relu{F}(u::Vector{F}) = Relu(u, one(F))
calc{F}(::Type{Relu{F}}, x) = max.(x, 0.)
derv{F}(::Type{Relu{F}}, x) = (x .> 0.) * 1.

################################################################# ReQU
struct Requ{F} <: AbSignedUniDirBasis{F}
  u::Vector{F}
  sgn::F
end
Requ{F}(u::Vector{F}) = Requ(u, one(F))
calc{F}(::Type{Requ{F}}, x) = max.(x, 0) .^ 2
derv{F}(::Type{Requ{F}}, x) = 2. * (x.>0) .* x

################################################################# Huber
struct Huber{t, F} <: AbSignedUniDirBasis{F}
    u::Vector{F}
    sgn::F
end
Huber{t}(u::Vector{F}, s::F) where {t, F} = Huber{t, F}(u, s)
Huber{t}(u::Vector{F}) where {t, F} = Huber{t, F}(u, one(F))
Huber(u::Vector{F}, s::F) where {F} = Huber{1., F}(u, s)
Huber(u::Vector{F}) where {F} = Huber{1., F}(u, one(F))
calc(::Type{Huber{t, F}}, x) where {F, t} = (x.>0) .* ((x*t .< 1) .* (t*x.*x/2) + (x*t .≥ 1) .* (x-.5/t))
derv(::Type{Huber{t, F}}, x) where {F, t} = (x.>0) .* ((x*t .< 1) .* (t*x)      + (x*t .≥ 1)            )

################################################################# Softplus
struct Softplus{t, F} <: AbSignedUniDirBasis{F}
    u::Vector{F}
    sgn::F
end
Softplus{t}(u::Vector{F}, s::F) where {t, F} = Softplus{t, F}(u, s)
Softplus{t}(u::Vector{F}) where {t, F} = Softplus{t, F}(u, one(F))
Softplus(u::Vector{F}, s::F) where {F} = Softplus{1., F}(u, s)
Softplus(u::Vector{F}) where {F} = Softplus{1., F}(u, one(F))
calc(::Type{Softplus{t, F}}, x) where {F, t} = log1plusexp(t*x)/t
derv(::Type{Softplus{t, F}}, x) where {F, t} = sigmoid(t*x)
###############################################################################


"""
   otherm_upp(pk, z, Ωm, h0[, Ωcb=Ωm]; massbias = 1.266, Mmin = 1e11, Mmax = 5e15, t10MF = false, xmax = 6.0)

Comoving density parameter of thermal energy of ionized baryons in collapsed structures using Planck's universal pressure profile.

*Reference*: Equation (1,6-8) of Chiang, Makiya, Komatsu & Ménard, arXiv:2007.01679

# Arguments
- `pk`(k): a function which returns a matter power spectrum with the argument k being the comoving wavenumber.
   - `pk` is in units of (Mpc/h)^3 and the argument k is in units of h/Mpc.
- `z::Real`: redshift.
- `Ωm::Real`: present-day total matter density parameter.
- `h0::Real`: present-day Hubble constant in units of 100 km/s/Mpc.

# Optional Arguments
- `Ωcb::Real=Ωm`: present-day baryon + cold dark matter density parameter.

# Optional keyword arguments
- `massbias::Real=1.266`: mass bias parameter for Planck's universal pressure profile, `B=Mtrue/Mplanck`.
- `Mmin::Real=1e11`: minimum mass for integration, ``∫_{Mmin}^{Mmax} dn/dM ∫dV Pe``.
- `Mmax::Real=5e15`: maximum mass for integration, ``∫_{Mmin}^{Mmax} dn/dM ∫dV Pe``.
- `t10MF::Bool=false`: if `true`, use `tinker10MF` for the halo multiplicity function. If `false` (the default), use `tinker08MF`.
- `xmax::Real=6.0`: pressure profile is integrated up to ``r = xmax * r500``.
"""
function otherm_upp(
   pk,
   z::Real,
   Ωm::Real,
   h0::Real,
   Ωcb = Ωm;
   massbias = 1.266,
   Mmin = 1e11,
   Mmax = 5e15,
   t10MF = false,
   xmax = 6.0,
)
   ρc = 2.775e11 # in units of h^2 M⊙/Mpc^3
   ρceVcm3 = 1.05375e4 # in units of h^2 eV/cm^3
   Δc = 500 # 500 times the critical density of the Universe
   E2 = Ωm * (1 + z)^3 + 1 - Ωm  # flat Universe
   Δm = Δc * E2 / Ωm / (1 + z)^3 # Δm for mass function
   nmass = 100
   func = zeros(nmass)
   lnMh = range(log(Mmin), length = nmass, log(Mmax))
   for i = 1:nmass
      func[i] = dndlnMh(pk, z, Δm, Ωcb, lnMh[i], t10MF = t10MF)
   end
   spl = Spline1D(lnMh, func)
   # %% Pressure profile integral
   # Planck's universal pressure profile
   # Reference: Planck Collaboration, A&A, 550, A131 (2013)
   c, γ, α, β, P0 = [1.81, 0.31, 1.33, 4.13, 6.41]
   upp(x) =
      x^2 * P0 * (0.7 / h0)^1.5 * (c * x)^(-γ) * (1 + (c * x)^α)^((γ - β) / α)
   uppint, err = quadde(upp, 0, xmax)
   # %% Compute ρth = \int dlnM dn/dlnM \int dV Pe, in units of h^2 eV/cm^3
   function dρdlnMh(lnMh)
      αp = 0.12
      E2 = Ωm * (1 + z)^3 + 1 - Ωm # for a flat Universe
      RΔh = cbrt(exp(lnMh) * 3 / 4π / (ρc * E2 * Δc)) # in h^-1 Mpc
      Pe =
         1.65 * # eV/cm^3
         (1 / 0.7)^2 * # h^2 cancels h^2 in the critical density in the denominator
         E2^(4 / 3) *
         (exp(lnMh) / massbias / 3e14 / 0.7)^(2 / 3 + αp)
      dρdlnMh = spl(lnMh) * Pe * 4π * RΔh^3 / massbias
   end
   ρe, err = hquadrature(dρdlnMh, log(Mmin), log(Mmax))
   redshift = z
   if t10MF
      fnu(x) = tinker10MF(x, redshift, Δm)
      factor, err = hquadrature(fnu, -50, 5)
      ρe /= factor
   end
   Y = 0.24 # helium abundance
   Ωtherm = (8 - 5Y) / (4 - 2Y) * ρe * uppint / ρceVcm3
end

"""
   otherm_ks(pk, z, Ωm, fb[, Ωcb=Ωm]; Mmin = 1e11, Mmax = 5e15, virial = false, t10MF = false, xmax = 3.0)

Comoving density parameter of thermal energy of ionized baryons in collapsed structures using Komatsu-Seljak pressure profile.

*Reference*: Equation (1,6-8) of Chiang, Makiya, Komatsu & Ménard, arXiv:2007.01679

# Arguments
- `pk`(k): a function which returns a matter power spectrum with the argument k being the comoving wavenumber.
    - **pk times k^3 must be dimensionless**. For example, if k is in units of h/Mpc, `pk` must be in units of Mpc^3/h^3.
- `z::Real`: redshift.
- `Ωm::Real`: present-day total matter density parameter.
- `fb::Real`: mean baryon fraction.

# Optional Arguments
- `Ωcb::Real=Ωm`: present-day baryon + cold dark matter density parameter.

# Optional keyword arguments
- `Mmin::Real=1e11`: minimum mass for integration, ``∫_{Mmin}^{Mmax} dn/dM ∫dV Pe``.
- `Mmax::Real=5e15`: maximum mass for integration, ``∫_{Mmin}^{Mmax} dn/dM ∫dV Pe``.
- `virial::Bool=false`: if `true`, use the virial overdensity `Δvir`. If `false` (the default), use `Δm=200`.
- `t10MF::Bool=false`: if `true`, use `tinker10MF` for the halo multiplicity function. If `false` (the default), use `tinker08MF`.
- `xmax::Real=3.0`: pressure profile is integrated up to ``r = xmax * c * rs = xmax * r200m (if virial = false) or rvir (if virial = true)``.
"""
function otherm_ks(
   pk,
   z::Real,
   Ωm::Real,
   fb::Real,
   Ωcb = Ωm;
   Mmin = 1e11,
   Mmax = 5e15,
   virial = false,
   t10MF = false,
   xmax = 3.0,
)
   ρc = 2.775e11 # in units of h^2 M⊙/Mpc^3
   GN = 3 / 8π / 2998^2 / ρc # = 4.786e-20 Mpc/M⊙ is Newton's constant
   if virial # Determine Δm used by the mass function
      E2 = Ωm * (1 + z)^3 + 1 - Ωm # flat Universe
      Ωmz = Ωm * (1 + z)^3 / E2
      Δvir = 18 * π^2 + 82 * (Ωmz - 1) - 39 * (Ωmz - 1)^2
      Δm = Δvir / Ωmz
   else
      Δm = 200
   end
   nmass = 100
   func = zeros(nmass)
   lnMh = range(log(Mmin), length = nmass, log(Mmax))
   for i = 1:nmass
      func[i] = dndlnMh(pk, z, Δm, Ωcb, lnMh[i], t10MF = t10MF)
   end
   spl = Spline1D(lnMh, func)
   ## Pressure profile integral
   # Komatsu-Seljak Pressure Profile
   # Reference: Komatsu & Seljak, MNRAS, 327, 1353 (2001)
   # Fitting formulae for γ(c) and η(c) are valid for 1<c<25.
   # Reference: Equations (17,18) of Komatsu & Seljak, MNRAS, 336, 1256 (2002)
   γ(c) = 1.137 + 8.94e-2 * log(c / 5) - 3.68e-3 * (c - 5)
   η0(c) = 2.235 + 0.202 * (c - 5) - 1.16e-3 * (c - 5)^2
   B(c) = 3 / η0(c) * (γ(c) - 1) / γ(c) / (log(1 + c) / c - 1 / (1 + c))
   ygas(x, c) = (1 - B(c) * (1 - log(1 + x) / x))^(1 / (γ(c) - 1))
   ksint = zeros(25)
   for c = 1:25
      ks(x) = x^2 * ygas(x, c)^γ(c)
      ksint[c], err = quadde(ks, 0, xmax * c) # integrated out to r = xmax*r200m or xmax*rvir
   end
   c = 1:25
   ksints = Spline1D(c, ksint)
   # %% Compute ρth = \int dlnM dn/dlnM \int dV Pe, in units of h^2 eV/cm^3
   ## Gas thermal energy density: Komatsu-Seljak profile
   function dρdlnMh(lnMh)
      if virial # For Δ = Δvir(z)
         RΔh = cbrt(exp(lnMh) * 3 / 4π / (ρc * E2 * Δvir)) # in h^-1 Mpc
         A0, B0, C0 = 7.85, -0.081, -0.71 # Avir, Bvir, Cvir in Table 1 of Duffy et al.
      else # For Δ = Δm = 200
         RΔh = cbrt(exp(lnMh) * 3 / 4π / (ρc * Ωcb * Δm)) / (1 + z) # in h^-1 Mpc
         A0, B0, C0 = 10.14, -0.081, -1.01 # Amean, Bmean, Cmean in Table 1 of Duffy et al.
      end
      c = A0 * (exp(lnMh) / 2e12)^B0 * (1 + z)^C0
      mc = log(1 + c) - c / (1 + c)
      ρgnorm = fb * (c * (1 + c)^2 * mc * ygas(c, c))^-1 # = 4πρgas(0)rs^3/Mgas
      dρdlnMh =
         spl(lnMh) * GN * exp(2 * lnMh) / RΔh * η0(c) / 3 * ρgnorm * ksints(c) # in units of h^2 M⊙/Mpc^3
   end
   res, err = hquadrature(dρdlnMh, log(Mmin), log(Mmax))
   redshift = z
   if t10MF
      fnu(x) = tinker10MF(x, redshift, Δm)
      factor, err = hquadrature(fnu, -50, 5)
      res /= factor
   end
   Ωtherm = res / ρc
end

"""
   by(pk, z, Ωm, [, Ωcb=Ωm]; αp = 0.12, Mmin = 1e11, Mmax = 5e15, t10MF = false)

Compton-y-weighted halo bias.

*Reference*: Equation (B1) of Chiang, Makiya, Ménard & Komatsu, arXiv:2006.14650

# Arguments
- `pk`(k): a function which returns a matter power spectrum with the argument k being the comoving wavenumber.
   - `pk` is in units of (Mpc/h)^3 and the argument k is in units of h/Mpc.
- `z::Real`: redshift.
- `Ωm::Real`: present-day total matter density parameter.

# Optional Arguments
- `Ωcb::Real=Ωm`: present-day baryon + cold dark matter density parameter.

# Optional keyword arguments
- `αp::Real=0.12`: the empirical correction factor for the pressure-mass relation (Equation (7) of Arnaud et al., Astronomy & Astrophysics, 517, A92 (2010))
- `Mmin::Real=1e11`: minimum mass for integration, ``∫_{Mmin}^{Mmax} dn/dM ∫dV Pe``.
- `Mmax::Real=5e15`: maximum mass for integration, ``∫_{Mmin}^{Mmax} dn/dM ∫dV Pe``.
- `t10MF::Bool=false`: if `true`, use `tinker10MF` for the halo multiplicity function. If `false` (the default), use `tinker08MF`.
"""
function by(
   pk,
   z::Real,
   Ωm::Real,
   Ωcb = Ωm;
   αp = 0.12,
   Mmin = 1e11,
   Mmax = 5e15,
   t10MF = false,
)
   Δc = 500 # 500 times the critical density of the Universe
   E2 = Ωm * (1 + z)^3 + 1 - Ωm  # flat Universe
   Δm = Δc * E2 / Ωm / (1 + z)^3 # Δm for mass function
   nmass = 100
   func1, func2 = zeros(nmass), zeros(nmass)
   lnMh = range(log(Mmin), length = nmass, log(Mmax))
   for i = 1:nmass
      func1[i] = dndlnMh(pk, z, Δm, Ωcb, lnMh[i], t10MF = t10MF)
      func2[i] = bdndlnMh(pk, z, Δm, Ωcb, lnMh[i], t10MF = t10MF)
   end
   spl1, spl2 = Spline1D(lnMh, func1), Spline1D(lnMh, func2)
   dbylnMh1(lnMh) = spl1(lnMh) * exp(lnMh * (5 / 3 + αp))
   dbylnMh2(lnMh) = spl2(lnMh) * exp(lnMh * (5 / 3 + αp))
   by1, err = hquadrature(dbylnMh1, log(Mmin), log(Mmax))
   by2, err = hquadrature(dbylnMh2, log(Mmin), log(Mmax))
   by = by2 / by1
end


using BlockArrays
using ControlSystems
using LinearAlgebra
using PredictiveControl
using Test

# Setup the problem using a test case from Richter's paper
n = 4
m = 2
N = 5

A = [ 0.7 -0.1  0.0  0.0;
      0.2 -0.5  0.1  0.0;
      0.0  0.1  0.1  0.0;
      0.5  0.0  0.5  0.5 ]

B = [ 0.0  0.1;
      0.1  1.0;
      0.1  0.0;
      0.0  0.0 ]

C = Matrix{Float64}( I, 4, 4 )

D = fill( 0.0, (4, 2) )

sys = StateSpace( A, B, C, D, 0.01 )

Gₛ = [ 1.0  0.0;
      -1.0  0.0;
       0.0  1.0;
       0.0 -1.0 ]
gₛ = [ 1.0;
      -1.0;
       1.0;
      -1.0 ]

clqr = ConstrainedTimeInvariantLQR( sys, N, Matrix( 1.0*I(n) ), Matrix( 1.0*I(m) ), :Q )

H = PredictiveControl.hessian( clqr, PrimalFullCondensing )
G = kron( I(N), Gₛ )
g = reshape( kron( ones( Float64, N, 1 ), gₛ ), (:) )

@test isapprox( FGM.coldstartdelta( G, g, H = Matrix( H ) ), 19.026675, atol = 1e-4 )
@test isapprox( FGM.coldstartdelta( G, g, L = eigmax( Matrix( H ) ) ), 19.026675, atol = 1e-4 )


using LocalApproximationPolicyEvaluation
using Test
using POMDPModels
using POMDPs
using GridInterpolations
using LocalFunctionApproximation

# Step 1 - Setup the gridworld problem
g_size = (9,9)
g = SimpleGridWorld(size = g_size, rewards = Dict(GWPos(g_size...) => 1, GWPos(1,1) => 0), tprob = 1., discount=1)
action_probability(g::SimpleGridWorld, s, a) = 0.25
POMDPs.convert_s(::Type{AbstractArray}, s::GWPos, g::SimpleGridWorld) = SVector{4, Float64}(1., s[1], s[2], s[1]*s[2])
POMDPs.gen(mdp::SimpleGridWorld, s, a, rng = Random.GLOBAL_RNG) = (sp =rand(transition(g, s, a )), r=reward(g, s, a))
POMDPs.initialstate(g::SimpleGridWorld) = initialstate_distribution(g)

# Step 2 - Solve the problem semi-exactly using local approximation
grid = RectangleGrid([1:g_size[1] ...], [1:g_size[2]...])
interp = LocalGIFunctionApproximator(grid)
solver = LocalPolicyEvalSolver(interp, action_probability, is_mdp_generative = true, n_generative_samples = 1, max_iterations = 2000, belres = 1e-6)
policy = solve(solver, g)
values = [value(policy, s) for s in states(g)]
@test value(policy, GWPos(9,9)) == 1.0
@test value(policy, GWPos(8,9)) > .25
@test value(policy, GWPos(1,1)) == 0