"""
    LinearOettliPrager <: AbstractDirectSolver

Type for the OettliPrager solver of the interval linear system ``Ax=b``. The solver first
converts the system of interval equalities into a system of real inequalities using
Oettli-Präger theorem [[OET64]](@ref) and then finds the feasible set by solving a LP
problem in each orthant using `LazySets.jl`.

### Notes
- You need to import `LazySets.jl` to use this functionality.
- An object of type `LinearOettliPrager` is a function with methods

        (op::LinearOettliPrager)(A::AbstractMatrix{T},
                                 b::AbstractVector{T}) where {T<:Interval}

    #### Input
    - `A`   -- N×N interval matrix
    - `b`   -- interval vector of length N

### Examples

```julia-repl
julia> A = [2..4 -2..1;-1..2 2..4]
2×2 Matrix{Interval{Float64}}:
  [2, 4]  [-2, 1]
 [-1, 2]   [2, 4]

julia> b = [-2..2, -2..2]
2-element Vector{Interval{Float64}}:
 [-2, 2]
 [-2, 2]

julia> polytopes = solve(A, b, LinearOettliPrager());

julia> typeof(polytopes)
Vector{HPolytope{Float64, SparseArrays.SparseVector{Float64, Int64}}}
```
"""
struct LinearOettliPrager <: AbstractDirectSolver end

"""
    NonLinearOettliPrager <: AbstractIterativeSolver

Type for the OettliPrager solver of the interval linear system ``Ax=b``. The solver first
converts the system of interval equalities into a system of real inequalities using
Oettli-Präger theorem [[OET64]](@ref) and then finds the feasible set using
the forward-backward contractor method [[JAU14]](@ref) implemented in
`IntervalConstraintProgramming.jl`.

### Fields
- `tol` -- tolerance for the paving, default 0.01.

### Notes
- You need to import `IntervalConstraintProgramming.jl` to use this functionality.
- An object of type `NonLinearOettliPrager` is a function with methods

        (op::NonLinearOettliPrager)(A::AbstractMatrix{T},
                                    b::AbstractVector{T},
                                    [X]::AbstractVector{T}=enclose(A, b)) where {T<:Interval}

        (op::NonLinearOettliPrager)(A::AbstractMatrix{T},
                                    b::AbstractVector{T},
                                    X::IntervalBox) where {T<:Interval}

    #### Input
    - `A`   -- N×N interval matrix
    - `b`   -- interval vector of length N
    - `X`   -- (optional) initial enclosure for the solution of ``Ax = b``. If not given,
               it is automatically computed using [`enclose`](@ref enclose)

### Examples

```julia-repl
julia> A = [2..4 -2..1;-1..2 2..4]
2×2 Matrix{Interval{Float64}}:
  [2, 4]  [-2, 1]
 [-1, 2]   [2, 4]

julia> b = [-2..2, -2..2]
2-element Vector{Interval{Float64}}:
 [-2, 2]
 [-2, 2]

julia> solve(A, b, NonLinearOettliPrager(0.1))
Paving:
- tolerance ϵ = 0.1
- inner approx. of length 1195
- boundary approx. of length 823
```
"""
struct NonLinearOettliPrager <: AbstractIterativeSolver
    tol::Float64
end
NonLinearOettliPrager() = NonLinearOettliPrager(0.01)


# using Cbc # slow, not recommended
# solver = CbcSolver(logLevel=0, integerTolerance=1e-9, primalTolerance=1e-9, ratioGap=1e-8)

using Gurobi
solver = GurobiSolver(OutputFlag=0)

# using CPLEX
# solver = CplexSolver(CPX_PARAM_SCRIND=0)

using JuMP
using PiecewiseLinearOpt
using Base.Test
using HDF5

methods_1D = (:CC,:MC,:Logarithmic,:LogarithmicIB,:ZigZag,:ZigZagInteger,:SOS2,:GeneralizedCelaya,:SymmetricCelaya,:Incremental,:DisaggLogarithmic)
methods_2D = (:CC,:MC,:Logarithmic,:LogarithmicIB,:ZigZag,:ZigZagInteger,:SOS2,:GeneralizedCelaya,:SymmetricCelaya,:DisaggLogarithmic)
# methods_2D = (:OptimalIB,) # very slow on all solvers
patterns_2D = (:Upper,:Lower,:UnionJack,:K1,:Random) # not :BestFit because requires function values at midpoints, :OptimalTriangleSelection and :Stencil not supported currently

# tests on network flow model with piecewise-linear objective
# instance data loaded from .h5 files
instance_data_1D = joinpath(dirname(@__FILE__),"1D-pwl-instances.h5")
instance_data_2D = joinpath(dirname(@__FILE__),"2D-pwl-instances.h5")

println("\nunivariate tests")
@testset "instance $instance" for instance in ["10104_1_concave_1"]
    demand = h5read(instance_data_1D, string(instance,"/demand"))
    supply = h5read(instance_data_1D, string(instance,"/supply"))
    ndem = size(demand, 1)
    nsup = size(supply, 1)

    rawd = h5read(instance_data_1D, string(instance,"/d"))
    d = rawd[1,:]
    rawfd = h5read(instance_data_1D, string(instance,"/fd"))
    fd = rawfd[1,:]

    objval1 = NaN
    @testset "1D: $method" for method in methods_1D
        model = Model(solver=solver)
        @variable(model, x[1:nsup,1:ndem] ≥ 0)
        @constraint(model, [j in 1:ndem], sum(x[i,j] for i in 1:nsup) == demand[j])
        @constraint(model, [i in 1:nsup], sum(x[i,j] for j in 1:ndem) == supply[i])
        @objective(model, Min, sum(piecewiselinear(model, x[i,j], d, fd, method=method) for i in 1:nsup, j in 1:ndem))

        @test solve(model) == :Optimal
        if isnan(objval1)
            objval1 = getobjectivevalue(model)
        else
            @test getobjectivevalue(model) ≈ objval1 rtol=1e-4
        end
    end
end

println("\nbivariate tests")
# @testset "numpieces $numpieces, variety $variety, objective $objective" for numpieces in [4,8], variety in 1:5, objective in 1:20
@testset "numpieces $numpieces, variety $variety, objective $objective" for numpieces in [4], variety in 1:1, objective in 1:1
    instance = string(numpieces,"_",variety,"_",objective)

    demand = h5read(instance_data_2D, string(instance,"/demand"))
    supply = h5read(instance_data_2D, string(instance,"/supply"))
    ndem = size(demand, 1)
    nsup = size(supply, 1)

    rawd = h5read(instance_data_2D, string(instance,"/d"))
    d = rawd[1,:]
    rawfd = h5read(instance_data_2D, string(instance,"/fd"))
    fˣ = reshape(rawfd[1,:], length(d), length(d))
    ˣtoⁱ = Dict(d[p] => p for p in 1:length(d))
    f = (pˣ,pʸ) -> fˣ[ˣtoⁱ[pˣ],ˣtoⁱ[pʸ]]

    objval1 = NaN
    @testset "2D: $method, $pattern" for method in methods_2D, pattern in patterns_2D
        model = Model(solver=solver)
        @variable(model, x[1:nsup,1:ndem] ≥ 0)
        @constraint(model, [j in 1:ndem], sum(x[i,j] for i in 1:nsup) == demand[j])
        @constraint(model, [i in 1:nsup], sum(x[i,j] for j in 1:ndem) == supply[i])
        @variable(model, y[1:nsup,1:ndem] ≥ 0)
        @constraint(model, [j in 1:ndem], sum(y[i,j] for i in 1:nsup) == demand[j])
        @constraint(model, [i in 1:nsup], sum(y[i,j] for j in 1:ndem) == supply[i])

        @objective(model, Min, sum(piecewiselinear(model, x[i,j], y[i,j], BivariatePWLFunction(d, d, f, pattern=pattern), method=method, subsolver=solver) for i in 1:nsup, j in 1:ndem))

        @test solve(model) == :Optimal
        if isnan(objval1)
            objval1 = getobjectivevalue(model)
        else
            @test getobjectivevalue(model) ≈ objval1 rtol=1e-4
        end
    end
end


Types = (Float32, Float64, BigFloat)
f(x, y, z) = 0.5x^2 - 0.1y^3 + z
f_wraparound2(x, y, z) = 5*sin(x) + 2*cos(y) + z
f_wraparound3(x, y, z) = 5*sin(x) + 2*cos(y) - 3*cos(z)
for T in Types
    grid = range(zero(T), stop=2π*one(T), length=40)
    f_uw = f.(grid, grid', reshape(grid, 1, 1, :))
    f_wr = f_uw .% (2π)
    uw_test = unwrap(f_wr)
    difference = first(f_uw) - first(uw_test)
    @test isapprox(uw_test .+ difference, f_uw, atol=1e-8)
    # test in-place version
    unwrap!(f_wr)
    difference = first(f_uw) - first(f_wr)
    @test isapprox(f_wr .+ difference, f_uw, atol=1e-8)

    f_uw = f_wraparound2.(grid, grid', reshape(grid, 1, 1, :))
    f_wr = f_uw .% (2π)
    uw_test = unwrap(f_wr, (true, true, false))
    difference = first(f_uw) - first(uw_test)
    @test isapprox(uw_test .+ difference, f_uw, atol=1e-8)
    # test in-place version
    unwrap!(f_wr, (true, true, false))
    difference = first(f_uw) - first(f_wr)
    @test isapprox(f_wr .+ difference, f_uw, atol=1e-8)

    f_uw = f_wraparound3.(grid, grid', reshape(grid, 1, 1, :))
    f_wr = f_uw .% (2π)
    uw_test = unwrap(f_wr, (true, true, true))
    difference = first(f_uw) - first(uw_test)
    @test isapprox(uw_test .+ difference, f_uw, atol=1e-8)
    # test in-place version
    unwrap!(f_wr, (true, true, true))
    difference = first(f_uw) - first(f_wr)
    @test isapprox(f_wr .+ difference, f_uw, atol=1e-8)
end


using StaticArrays
using Combinatorics: permutations, multiset_permutations
using Random: randperm

export lambdamatrix, lambdaperms, periodicorbits, find_period

"""
    lambdamatrix(λ, inds::Vector{Int}, sings) -> Λk
Return the matrix ``\\mathbf{\\Lambda}_k`` used to create a new
dynamical system with some unstable fixed points turned to stable in the function
[`periodicorbits`](@ref).

## Arguments

1. `λ<:Real` : the multiplier of the ``C_k`` matrix, with `0<λ<1`.
2. `inds::Vector{Int}` :
   The `i`th entry of this vector gives the *row* of the nonzero element of the `i`th
   column of ``C_k``.
3. `sings::Vector{<:Real}` : The element of the `i`th column of ``C_k`` is +1
   if `signs[i] > 0` and -1 otherwise (`sings` can also be `Bool` vector).

Calling `lambdamatrix(λ, D::Int)`
creates a random ``\\mathbf{\\Lambda}_k`` by randomly generating
an `inds` and a `signs` from all possible combinations. The *collections*
of all these combinations can be obtained from the function [`lambdaperms`](@ref).

## Description
Each element
of `inds` *must be unique* such that the resulting matrix is orthogonal
and represents the group of special reflections and permutations.

Deciding the appropriate values for `λ, inds, sings` is not trivial. However, in
ref. [2] there is a lot of information that can help with that decision. Also,
by appropriately choosing various values for `λ`, one can sort periodic
orbits from e.g. least unstable to most unstable, see [3] for details.

## References

[2] : D. Pingel *et al.*, Phys. Rev. E **62**, pp 2119 (2000)

[3] : F. K. Diakonos *et al.*, Phys. Rev. Lett. **81**, pp 4349 (1998)
"""
function lambdamatrix(λ::Real, inds::AbstractVector{<:Integer},
    sings::AbstractVector{<:Real})
    # this function seems to be super inefficient

    D = length(inds)
    D != length(sings) && throw(ArgumentError("inds and sings must have equal size."))
    0 < λ < 1 || throw(ArgumentError("λ must be in (0,1)"))
    unique(inds)!=inds && throw(ArgumentError("All elements of inds must be unique."))
    # This has to be improved to not create intermediate arrays!!!
    a = zeros(typeof(λ), (D,D))
    for i in 1:D
        a[(i-1)*D + inds[i]] = λ*(sings[i] > 0 ? +1 : -1)
    end
    return SMatrix{D,D}(a)
end

function lambdamatrix(λ::T, D::Integer) where {T<:Real}
    positions = randperm(D)
    signs = rand(Bool, D)
    lambdamatrix(λ, positions, signs)
end

"""
    lambdaperms(D) -> indperms, singperms
Return two collections that each contain all possible combinations of indices (total of
``D!``) and signs (total of ``2^D``) for dimension `D` (see [`lambdamatrix`](@ref)).
"""
function lambdaperms(D::Integer)
    indperms = collect(permutations([1:D;], D))
    p = trues(D)
    singperms = [p[:]] #need p[:] because p is mutated afterwards
    for i = 1:D
        p[i] = false; append!(singperms, multiset_permutations(p, D))
    end
    return indperms, singperms
end

"""
    periodicorbits(ds::DiscreteDynamicalSystem,
                   o, ics [, λs, indss, singss]; kwargs...) -> FP
Find fixed points `FP` of order `o` for the map `ds`
using the algorithm due to Schmelcher & Diakonos [1].
`ics` is a collection of initial conditions (container of vectors) to be evolved.

## Optional Arguments
The optional arguments `λs, indss, singss` *must be containers* of appropriate
values, besides `λs` which can also be a number. The elements of those containers
are passed to: [`lambdamatrix(λ, inds, sings)`](@ref), which creates the appropriate
``\\mathbf{\\Lambda}_k`` matrix. If these arguments are not given,
a random permutation will be chosen for them, with `λ=0.001`.

## Keyword Arguments
* `maxiters::Int = 100000` : Maximum amount of iterations an i.c. will be iterated
   before claiming it has not converged.
* `disttol = 1e-10` : Distance tolerance. If the 2-norm of a previous state with
   the next one is `≤ disttol` then it has converged to a fixed point.
* `inftol = 10.0` : If a state reaches `norm(state) ≥ inftol` it is assumed that
   it has escaped to infinity (and is thus abandoned).
* `roundtol::Int = 4` : The found fixed points are rounded
   to `roundtol` digits before pushed into the list of returned fixed points `FP`,
   *if* they are not already contained in `FP`.
   This is done so that `FP` doesn't contain duplicate fixed points (notice
   that this has nothing to do with `disttol`). Turn this to `typemax(Int)`
   to get the full precision of the algorithm.

## Description
The algorithm used can detect periodic orbits
by turning fixed points of the original
map `ds` to stable ones, through the transformation
```math
\\mathbf{x}_{n+1} = \\mathbf{x}_n +
\\mathbf{\\Lambda}_k\\left(f^{(o)}(\\mathbf{x}_n) - \\mathbf{x}_n\\right)
```
with ``f`` = `eom`. The index ``k`` counts the various
possible ``\\mathbf{\\Lambda}_k``.

## Performance Notes
*All* initial conditions are
evolved for *all* ``\\mathbf{\\Lambda}_k`` which can very quickly lead to
long computation times.

## References

[1] : P. Schmelcher & F. K. Diakonos, Phys. Rev. Lett. **78**, pp 4733 (1997)
"""
function periodicorbits(ds::DDS{false, S},
                        o::Int,
                        ics,
                        λs,
                        indss,
                        singss;
                        maxiters::Int = 100000,
                        disttol::Real = 1e-10,
                        inftol::Real = 10.0,
                        roundtol::Int = 4) where {S}

    FP = S[]
    integ = integrator(ds)
    for λ in λs
        for inds in indss
            for sings in singss
                Λ = lambdamatrix(λ, inds, sings)
                _periodicorbits!(FP, integ, o, ics, Λ, maxiters, disttol, inftol, roundtol)
            end
        end
    end
    return Dataset(FP)
end
function periodicorbits(ds::DDS{false, S},
                        o::Int,
                        ics;
                        maxiters::Int = 100000,
                        disttol::Real = 1e-10,
                        inftol::Real = 10.0,
                        roundtol::Int = 4) where {S}

    integ = integrator(ds)
    FP = S[]
    Λ = lambdamatrix(0.001, dimension(ds))
    _periodicorbits!(FP, integ, o, ics, Λ, maxiters, disttol, inftol, roundtol)
    return Dataset(FP)
end
periodicorbits(::DDS{true}, args...; kwargs...) = error(
"Currently `periodicorbits` is implemented only for out-of-place maps. Sorry!")

function _periodicorbits!(
    FP, integ, o, ics, Λ, maxiter, disttol, inftol, roundtol)

    for st in ics
        reinit!(integ, st)
        prevst = st
        for i in 1:maxiter
            prevst, st = Sk(integ, prevst, o, Λ)
            norm(st) > inftol && break

            if norm(prevst - st) < disttol
                unist = round.(st, digits = roundtol)
                unist ∉ FP && push!(FP, unist)
                break
            end
            prevst = st
        end
    end
end

function Sk(integ::MDI{false}, prevst, o, Λ)
    integ.u = prevst
    step!(integ, o)
    return prevst, prevst + Λ*(integ.u - prevst)
end


include("common.jl")

function f(x, y)
    0.0001*(10/(1+x^2) + x * y + 10*y^2)
end

Δt = 0.01
mmesh = Mesh(30, 30, 1/30)
bdnode = bcnode(mmesh)
bdnode = [bdnode; bdnode .+ mmesh.ndof]
# nu = constant(eval_f_on_gauss_pts(f, mmesh))

nu = Variable(ones(get_ngauss(mmesh))) * 0.0001

nodes = fem_nodes(mmesh)
u = @. sin(2π * nodes[:,1])
v = @. sin(2π * nodes[:,2])
u0 = [u;v]
u0[bdnode] .= 0.0
us = solve_burgers(u0, 10, nu)

U = matread("fenics/fwd2.mat")["U"]
loss = sum((us - U)^2)

_loss = Float64[]
cb = (vs, iter, loss)->begin
    global _loss
    push!(_loss, loss)
    printstyled("[#iter $iter] loss=$loss\n", color=:green)
    if mod(iter, 10)==1
        @info iter
        close("all")
        visualize_scalar_on_gauss_points(vs[1], mmesh)
        matwrite("fenics/bwd3-$iter.mat", Dict("iter"=>iter,"loss"=>_loss, "nu"=>vs[1]))
        savefig("fenics/bwd3_pixel$iter.png")
    end
end

# @info run(sess, loss)
sess = Session(); init(sess)
BFGS!(sess, loss, vars = [nu], callback = cb)


using JuMP, Gurobi, Combinatorics, Cbc
import XLSX

#### read data ####
xf = XLSX.readxlsx("../data/data1.xlsx")
l = xf["General!B1"]     # work day duration
N = xf["General!B3"] + 1  # number of clients + 1
M = xf["General!B2"]  # number of providers
# + 1 is for including the headquarter (HS14)

# Clients
s = xf["Clients!A2:A$N"]   # client service start times
d = vcat([0], xf["Clients!B2:B$N"])   # durations for clients, set d[1] = 0 for HS14

# Providers
w = xf["Providers!B2:B$(M+1)"]  # start time of providers
f = xf["Providers!A2:A$(M+1)"]  # hiring cost of providers

# make separate s vector for each provider, inserting into the first element the provider's work start time
s_test = ones(M, N)
for k in 1:M
    global s_test[k,:] = vcat([w[k]], s)
end
s = s_test

# Distances
t = xf["Distances"][:]  # 2-d array of travel times
#### ######

L = 10000 # large number

E = collect(powerset(collect(2:N), 2, N-1)) # for subtour 'E'limination constraints

########

model = Model();
set_optimizer(model, Gurobi.Optimizer);
#(13)
@variable(model, x[1:N, 1:N, 1:M], Bin)

######################## routing constraints ######################

# (2) each client served once by one provider
@constraint(model, served[j in 2:N], sum(x[i,j,k] for i in 1:N for k in 1:M) == 1)

# (3) route continuity
@constraint(model, continuity[p in 2:N, k in 1:M], sum(x[i,p,k] for i in 1:N) == sum(x[p,j,k] for j in 1:N) )

# (4) not all providers need to leave 14HS (and start working)
@constraint(model, leave_depot[k in 1:M], sum(x[1,j,k] for j in 2:N) <= 1)

# (5) not all providers need to return to 14HS (i.e. they didn't leave 14HS in the first place)
@constraint(model, return_depot[k in 1:M], sum(x[i,1,k] for i in 2:N) <= 1)

# (6) no self looping
@constraint(model, diag[k in 1:M], sum(x[i,i,k] for i in 1:N) == 0)

# (7) at least 1 provider must leave 14HS and start working to serve clients
@constraint(model, sum(x[1,j,k] for j in 2:N for k in 1:M) >= 1)

# (8) all providers must start their work at 14HS, if they do work
@constraint(model, large[k in 1:M], sum(x[i,j,k] for i in 2:N for j in 2:N) <= L*sum(x[1,j,k] for j in 2:N))

#(9) subtour elimination for each provider
for k in 1:M
    for e in 1:length(E)
        @constraint(model, sum(x[i,j,k] for i in E[e] for j in E[e]) <= length(E[e])-1)
    end
end

######################## end routing constraints ######################

######################## time constraints ######################

# (10) initial client service times
@constraint(model, start[j in 2:N, k in 1:M], x[1,j,k]*(s[k,1] + t[1,j] + d[1]) <= s[k,j])

# (11) in between start and end of work, service times
@constraint(model, bet[i in 2:N, j in 2:N, k in 1:M], x[i,j,k]*(s[k,i] + t[i,j] + d[i]) <= s[k,j])

# (12) come back to 14HS by hour l.
@constraint(model, endofwork[i in 2:N, k in 1:M], x[i,1,k]*(s[k,i] + t[i,1] + d[i]) <= l)

######################## end time constraints ######################

# (1) minimise hiring cost to serve all clients
@objective(model, Min, sum(f[k]*x[1,j,k] for j in 1:N for k in 1:M))

optimize!(model)

(value.(x[:,:,1]))

for k in 1:M
    if sum(value.(x[1,:,k])) == 1
        println("vehicle $k provides service.")
        # come back hr
        for i in 2:N
            if value.(x[i,1,k]) == 1
                println("vehicle $k comes back to 14HS at $(t[i,1] + s[k,i] + d[i])hr. \n")
            end
        end
    end
end

path = [1 2 1]
#path = [1 7 1]
#path = [1 6 1]
#path = [1 3 4 5 1]

worked_hr = 0
#delivered  = 0
for i in 2:length(path)
    global worked_hr += t[path[i-1], path[i]] + d[path[i]]
    #global delivered += d[path[i]]
end
worked_hr


"""
There are extremely high values in the data due to fires, gas flare etc. You may find some values even greater than the aggregate radiance of large cities. Such pixels also have high standard deviation. These pixels may not be of importantance from the point of view of measureming prosperity. The ```outlier_mask``` function generates a mask of pixels with standard deviation less that the 99.9th percentile. Essentially, this function can be used to removed top 1 percent of pixels by standard deviation. A mask can be provided to the function, so that it calculates the percentile based on the lit pixel of the mask. For example, if the datacube is a box around India and the mask is the polygon mask of India, the outlier_mask function will calculate the 99th percentile of the standard deviation of the pixels inside India's boundary.

```julia
outlier_mask(datacube, mask)
```
"""
function outlier_mask(datacube, mask)
    stds = zeros(size(datacube)[1], size(datacube)[2])
    @showprogress for i in 1:size(datacube)[1]
        for j in 1:size(datacube)[2]
            if mask[i, j]==0
                continue
            end
            if count(i->(ismissing(i)),datacube[i, j, :])/length(datacube[i, j, :]) > 0.50  # Don't do anything if there are too many missings
                continue
            end
            stds[i, j] = std(detrend_ts(filter(x -> !ismissing(x), datacube[i, j, :])))
        end
    end
    function std_mask(std, threshold)
        if std < threshold
            return 1
        else
            return 0
        end
    end
    threshold   = quantile(mask_vals(stds, mask), 0.999)
    outlierMask = std_mask.(stds, threshold)
    outlierMask = outlierMask.*mask
    outlierMask = convert(Array{Int8, 2}, outlierMask)
    return outlierMask
end

"""
The time series of a pixel may show a few outliers, but as a whole the pixel may be of importantance in measuring economic activity. The outlier_ts function uses replaces the outlier observations with interpolated values. This is done using the tsclean function the forecast package of R.

```julia
sample_timeseries = datacube[1, 2, :] # The time series of pixel [1, 2]
outlier_ts(sample_timeseries)
```
"""

function outlier_ts(timeseries, window_size = 5, n_sigmas = 3)
    # Credit: https://gist.github.com/erykml/d15525855f2ef455bd7969240f6f4073#file-hampel_filter_forloop-py
    missings = findall(ismissing, timeseries)
    new_series = Array{Union{Float64, Missing}}(filter!(!ismissing, copy(timeseries)))
    n = length(new_series)
    k = 1.4826 # scale factor for Gaussian distribution
    indices = []
    for i in (window_size):(n - window_size)
        x0 = median(new_series[(i - window_size+1):(i + window_size)]) # Use this for the time series outlier package
        # x0 = missing
        S0 = k * median(abs.(new_series[(i - window_size +1):(i + window_size)] .- x0))
        if (abs(new_series[i] - x0) > n_sigmas * S0)
            # new_series[i] = missing # Can use x0
            push!(indices, i)
        end
    end
    new_series[indices] .= missing
    for i in missings #putting back the missing values
        new_series = insert!(new_series, i, missing)
    end
    return new_series
    # Think about how this would behave in the presence of missing.
end


using PolyhedralRelaxations, JuMP, Cbc, Test

"""
    example_trig(; verbose = true)

``trig'' is an MINLPLib test instance with trigonometric functions
(see http://www.minlplib.org/trig.html). The best known primal feasible
objective value for this problem is -3.76250036. This example illustrates
how to obtain lower bounds to the optimal objective value using the MILP and
LP relaxations obtained using the PolyhedralRelaxations package.

trig is a non-linear non-convex minimization problem with one decision variable
and one inequality constraint.

Variables: -2 <= x <= 5; start = 1
Constraint: 5 * sin(x) - x <= 0.0
Objective: Minimize sin(11 * x) - cos(13 * x) + sin(17 * x) + cos(19 * x)

We reformulate this problem using additional variables as
Variables:
    -2 <= x <= 5; start = 1
    -1 <= y[1:5] <= 1
Constraints:
    5 * y[1] - x <= 0.0
    y[1] = sin(x)
    y[2] = sin(11 * x)
    y[3] = cos(13 * x)
    y[4] = sin(17 * x)
    y[5] = cos(19 * x)
Objective: Minimize y[2] + y[3] - y[4] - y[5]

"""

function example_trig(; verbose = true)
    best_known_objective = -3.76250036
    cbc_optimizer = JuMP.optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    error_tolerances = [NaN64, 1e-1, 1e-2]
    # base partition holds the inflection points and the end-points for the trigonometric functions
    base_partition = Dict{Int,Vector{Float64}}()
    base_partition[1] = [-2.0, 0.0, π, 5.0]
    base_partition[2] = [-2.0, collect(-7*π/11:π/11:17*π/11)..., 5.0]
    base_partition[3] = [-2.0, collect(-15*π/26:π/13:41*π/26)..., 5.0]
    base_partition[4] = [-2.0, collect(-10*π/17:π/17:27*π/17)..., 5.0]
    base_partition[5] = [-2.0, collect(-23*π/38:π/19:59*π/38)..., 5.0]
    # functions
    functions = Dict{Int,Function}()
    functions[1] = x -> sin(x)
    functions[2] = x -> sin(11 * x)
    functions[3] = x -> cos(13 * x)
    functions[4] = x -> sin(17 * x)
    functions[5] = x -> cos(19 * x)
    for i = 1:length(error_tolerances)
        err_tol = error_tolerances[i]
        milp = Model(cbc_optimizer)
        @variable(milp, -2.0 <= x <= 5.0)
        @variable(milp, -1.0 <= y[1:5] <= 1.0)

        # construct MILP relaxations
        for j = 1:5
            f = functions[j]
            p = deepcopy(base_partition[j])
            construct_univariate_relaxation!(
                milp,
                f,
                x,
                y[j],
                p,
                true,
                error_tolerance = err_tol,
            )
        end
        @constraint(milp, 5 * y[1] - x <= 0.0)
        @objective(milp, Min, y[2] + y[3] - y[4] - y[5])
        optimize!(milp)
        relaxation_objective = objective_value(milp)
        @test relaxation_objective ≤ best_known_objective
        relative_gap =
            abs(best_known_objective - relaxation_objective) / abs(relaxation_objective)
        if verbose
            println(
                "Optimal solution: $best_known_objective; MILP for error tolerance $err_tol: $relaxation_objective; relative gap: $relative_gap",
            )
        end
    end

    for i = 1:length(error_tolerances)
        err_tol = error_tolerances[i]
        # construct LP relaxations
        lp = Model(cbc_optimizer)
        @variable(lp, -2.0 <= x <= 5.0)
        @variable(lp, -1.0 <= y[1:5] <= 1.0)
        for j = 1:5
            f = functions[j]
            p = deepcopy(base_partition[j])
            construct_univariate_relaxation!(
                lp,
                f,
                x,
                y[j],
                p,
                false,
                error_tolerance = err_tol,
            )
        end
        @constraint(lp, 5 * y[1] - x <= 0.0)
        @objective(lp, Min, y[2] + y[3] - y[4] - y[5])
        optimize!(lp)
        relaxation_objective = objective_value(lp)
        @test relaxation_objective ≤ best_known_objective
        relative_gap =
            abs(best_known_objective - relaxation_objective) / abs(relaxation_objective)
        if verbose
            println(
                "Optimal solution: $best_known_objective; LP for error tolerance $err_tol: $relaxation_objective; relative gap: $relative_gap",
            )
        end
    end

end

example_trig()


#=
This example classifies CIFAR 10(https://www.cs.toronto.edu/~kriz/cifar.html)
dataset using a convolutional neural network with batch normalization

For details of applying batch normalization, see:
"Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift"
Sergey Ioffe  & Christian Szegedy
https://arxiv.org/abs/1502.03167

The trained architecture has the form of:
conv -> bn -> relu -> pool ->
conv -> bn -> relu -> pool ->
conv -> bn -> relu ->
fc -> bn  -> relu -> output

=#

using Knet, Statistics, MLDatasets

function loaddata()
    @info("Loading CIFAR 10...")
    xtrn, ytrn = CIFAR10.traindata()
    xtst, ytst = CIFAR10.testdata()
    xtrn = convert(Knet.array_type[], xtrn)
    xtst = convert(Knet.array_type[], xtst)
    #= Subtract mean of each feature
    where each channel is considered as
    a single feature following the CNN
    convention=#
    mn = mean(xtrn, dims=(1,2,4))
    xtrn = xtrn .- mn
    xtst = xtst .- mn
    @info("Loaded CIFAR 10")
    return (xtrn, ytrn), (xtst, ytst)
end

##Model definition

#=
Initialization is from
He et al., 2015,
Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification
https://arxiv.org/abs/1502.01852
=#
kaiming(et, h, w, i, o) = et(sqrt(2 / (w * h * o))) .* randn(et, h, w, i, o)

function init_model(;et=Float32)
    # Use bnparams() to initialize gammas and betas
    w = Any[
        kaiming(et, 3, 3, 3, 16),    bnparams(et, 16),
        kaiming(et, 3, 3, 16, 32),   bnparams(et, 32),
        kaiming(et, 3, 3, 32, 64),   bnparams(et, 64),
        xavier(et, 100, 8 * 8 * 64), bnparams(et, 100),
        xavier(et, 10, 100),         zeros(et, 10, 1)
    ]
    # Initialize a moments object for each batchnorm
    m = Any[bnmoments() for i = 1:4]
    w = map(Knet.array_type[], w)
    return w, m
end

function conv_layer(w, m, x; maxpool=true)
    o = conv4(w[1], x; padding=1)
    o = batchnorm(o, m, w[2])
    o = relu.(o)
    if maxpool; o=pool(o); end
    return o
end

function lin_layer(w, m, x)
    o = w[1] * x
    o = batchnorm(o, m, w[2])
    return relu.(o)
end

function predict(w, m, x)
    o = conv_layer(w[1:2] , m[1], x)
    o = conv_layer(w[3:4] , m[2], o)
    o = conv_layer(w[5:6] , m[3], o; maxpool=false)
    o = lin_layer( w[7:8] , m[4], mat(o))
    return w[9] * o .+ w[10]
end

function loss(w, m, x, classes)
    ypred = predict(w,m, x)
    return nll(ypred, classes)
end

lossgrad = grad(loss)

# Training
function epoch!(w, m, o, xtrn, ytrn;  mbatch=64)
    data = minibatch(xtrn, ytrn, mbatch;
                   shuffle=true,
                   xtype=Knet.array_type[])
    for (x, y) in data
        g = lossgrad(w, m, x, y)
        update!(w, g, o)
    end
end

# Accuracy computation
function acc(w, m, xtst, ytst; mbatch=64)
    data = minibatch(xtst, ytst, mbatch;
                     partial=true,
                     xtype=Knet.array_type[])
    accuracy(x->predict(w,m,x), data)
end

# TODO: add command line options
function train(;optim=Momentum, epochs=5,
               lr=0.1, oparams...)
    w, m = init_model()
    o = map(_->Momentum(;lr=lr, oparams...), w)
    (xtrn, ytrn), (xtst, ytst) = loaddata()
    for epoch = 1:epochs
        println("epoch: ", epoch)
        @time epoch!(w, m, o, xtrn, ytrn)
        println("train accuracy: ", acc(w, m, xtrn, ytrn))
        println("test accuracy: ", acc(w, m, xtst, ytst))
    end
end

endswith(string(PROGRAM_FILE), "cnn_batchnorm.jl") && train()

nothing


#  Solves a linear elliptic PDE in 1D
  using FEMfunctions

  using LinearAlgebra
  using SparseArrays

  nElements = 40

  include("onedElliptic.jl")

  order = 1
  error = onedElliptic(nElements,order)
  @test error<1e-10

  order = 2
  error = onedElliptic(nElements,order)
  @test error<1e-10

  order = 3
  error = onedElliptic(nElements,order)
  @test error<1e-10


import Obj
import Obj.Objective

function linreg_g()
    println("\nTesting gradient of Obj.linreg(X, y)...")
    X, y, w = _obj_randXyw() # init
    _obj_g(Obj.linreg(X, y), w)
end

function logreg_g()
    println("\nTesting gradient of Obj.logreg(X, y)...")
    X, y, w = _obj_randXyw() # init
    _obj_g(Obj.logreg(X, y), w)
end

_obj_randXyw(n=10, m=100) = randn(m,n), randn(m), randn(n)

function _obj_g(obj::Objective, w::Array{Float64,1})
    g = Obj.g(obj, w)
    g_b = sum([ Obj.g(obj, w, [i]) for i = 1:obj.dim ]) ./ obj.dim

    # assert
    success = (g ≈ g_b) # approx equal?
    if (success)
        println("SUCCEEDED: Complete grad at once matches complete grad using g_b")
    else
        println("FAILED! 2norm-distance between grad and stochastic grad is $(vecnorm(g-g_b, 2))")
    end
    return nothing
end


"""
Demonstration of solution of SYK equation by Matsubara frequency/imaginary time
alternation method using the discrete Lehmann representation

 The SYK equation is the nonlinear Dyson equation corresponding to self-energy

     Σ(τ) = J^2 * G(τ) * G(τ) * G(β-τ).

 We solve the Dyson equation self-consistently by a weighted
 fixed point iteration, with weight w assigned to the new iterate
 and weight 1-w assigned to the previous iterate. The self-energy
 is evaluated in the imaginary time domain, and each linear Dyson
 equation, corresponding to fixed self-energy, is solved in the
 Matsubara frequency domain, where it is diagonal.

 To solve the equation with a desired chemical potential μ,
 we pick a number nmu>1, and solve a sequence intermediate
 problems to obtain a good initial guess. First we solve the
 equation with chemical potential μ_0 = 0, then use this
 solution as an initial guess for the fixed point iteration with
 μ_1 = μ/nmu, then use this the solution as an initial guess
 for the fixed point iteration with μ = 2*μ/nmu, and so on,
 until we reach μ_{nmu} = μ.
"""

using Lehmann

diff(a, b) = maximum(abs.(a - b)) # return the maximum deviation between a and b

conformal_tau(τ, β) = -π^(1 / 4) / sqrt(2β) * 1 / sqrt(sin(π * τ / β))

function syk_sigma_dlr(d, G_x, J = 1.0)

    tau_k = d.τ # DLR imaginary time nodes
    tau_k_rev = d.β .- tau_k # Reversed imaginary time nodes

    G_k = dlr2tau(d, G_x) # G at tau_k
    G_k_rev = dlr2tau(d, G_x, tau_k_rev) # G at beta - tau_k

    # for i in 1:length(G_k)
    #     println("$(d.τ[i])    $(real(G_k[i]))     $(imag(G_k[i]))")
    # end

    # for i = 1:length(tau_k)
    #     G_k[i] = conformal_tau(tau_k[i], d.β)
    #     G_k_rev[i] = conformal_tau(tau_k_rev[i], d.β)
    # end

    Sigma_k = J .^ 2 .* G_k .^ 2 .* G_k_rev # SYK self-energy in imaginary time
    Sigma_x = tau2dlr(d, Sigma_k) # DLR coeffs of self-energy

    # println("sigma diff: ", diff(Sigma_k, dlr2tau(d, Sigma_x)))
    # for i in 1:length(Sigma_x)
    #     println("$(d.τ[i])    $(real(Sigma_x[i]))     $(imag(Sigma_x[i]))")
    # end
    # exit(0)
    return Sigma_x
end

function solve_syk_with_fixpoint_iter(d, mu, tol = d.rtol, mix = 0.3, maxiter = 100)

    Sigma_q = zeros(length(d)) # Initial guess
    G_q = zeros(ComplexF64, length(d))
    for iter in 1:maxiter

        G_q .= -1 ./ (d.ωn * 1im .- mu .+ Sigma_q) # Solve Dyson
        G_x = matfreq2dlr(d, G_q) # Get DLR coeffs
        Sigma_x_new = syk_sigma_dlr(d, G_x)
        Sigma_q_new = dlr2matfreq(d, Sigma_x_new)

        # println(diff(Sigma_q_new, Sigma_q))
        if maximum(abs.(Sigma_q_new .- Sigma_q)) < tol
            break
        end
        Sigma_q = mix * Sigma_q_new + (1 - mix) * Sigma_q # Linear mixing
    end
    return G_q
end

d = DLRGrid(Euv = 5.0, β = 1000.0, isFermi = true, rtol = 1e-14) # Initialize DLR object
G_q = solve_syk_with_fixpoint_iter(d, 0.0)

G_q = matfreq2tau(d, G_q)
println(G_q)
for i in 1:length(G_q)
    println("$(d.τ[i])    $(real(G_q[i]))     $(imag(G_q[i]))")
end


"""
    unsafe_min(x, y)
    unsafe_max(x, y)

of (x, y) obtain the one closer to -Inf [Inf], respectively
-  where x == y, obtain either

Only safe for use where neither x nor y is NaN
""" unsafe_min, unsafe_max

"""
    unsafe_minmag(x, y)
    unsafe_maxmag(x, y)

Only safe for use where neither x nor y is NaN
""" unsafe_minmag, unsafe_maxmag

unsafe_min(x::Float64, y::Float64) = signbit(x - y) ? x : y
unsafe_max(x::Float64, y::Float64) = signbit(y - x) ? x : y

unsafe_minmax(x::Float64, y::Float64) = ifelse(x < y, (x, y), (y, x))
unsafe_maxmin(x::Float64, y::Float64) = ifelse(y < x, (x, y), (y, x))

function fast_min(x::Float64, y::Float64)
    x_minus_y = x - y
    signbit(x_minus_y) ? x : ifelse(!isnan(x_minus_y), y, NaN64)
end

function fast_max(x::Float64, y::Float64)
    x_minus_y = x - y
    signbit(x_minus_y) ? y : ifelse(!isnan(x_minus_y), x, NaN64)
end

unsafe_minmag(x::Float64, y::Float64) = ifelse(abs(x) < abs(y), x, y)
unsafe_maxmag(x::Float64, y::Float64) = ifelse(abs(y) < abs(x), x, y)

function fast_minmag(x::Float64, y::Float64)
    x_minus_y = abs(x) - abs(y)
    signbit(x_minus_y) ? x : ifelse(!isnan(x_minus_y), y, NaN64)
end

function fast_maxmag(x::Float64, y::Float64)
    x_minus_y = abs(x) - abs(y)
    signbit(x_minus_y) ? y : ifelse(!isnan(x_minus_y), x, NaN64)
end

function unsafe_minmaxmag(x::Float64, y::Float64)
    ifelse(abs(x) < abs(y), (x, y), (y, x))
end

function unsafe_maxminmag(x::Float64, y::Float64)
    ifelse(abs(y) < abs(x), (x, y), (y, x))
end


using Test
using QuantLib
using Dates

# build a fixed rate bond to test cash flows
settlement_date = Date(2008, 9, 18)
set_eval_date!(settings, settlement_date - Dates.Day(3))
bbCurveRate = 0.055
bbDayCount = QuantLib.Time.ActualActualBond()
bbIR = InterestRate(bbCurveRate, bbDayCount, CompoundedCompounding(), QuantLib.Time.Semiannual())
ts = FlatForwardTermStructure(settlement_date, bbIR.rate, bbIR.dc, bbIR.comp, bbIR.freq)
settlement_days = 3
face_amount = 100.0
fixed_schedule = QuantLib.Time.Schedule(Date(2007, 5, 15), Date(2017, 5, 15),
                       QuantLib.Time.TenorPeriod(QuantLib.Time.Semiannual()), QuantLib.Time.Unadjusted(),
                       QuantLib.Time.Unadjusted(), QuantLib.Time.DateGenerationBackwards(), false,
                       QuantLib.Time.USGovernmentBondCalendar())

pe = DiscountingBondEngine(ts)
fixedrate_bond = FixedRateBond(settlement_days, face_amount, fixed_schedule, 0.045,
                        QuantLib.Time.ISMAActualActual(), QuantLib.Time.ModifiedFollowing(), 100.0,
                        Date(2007, 5, 15), fixed_schedule.cal, pe)

# generate a floating rate bond
floating_cal = QuantLib.Time.TargetCalendar()
settlement_days = 3
cap_vol = ConstantOptionVolatility(3, floating_cal, QuantLib.Time.ModifiedFollowing(), 0.0, QuantLib.Time.Actual365())

fb_issue_date = Date(2005, 10, 21)
fb_bond_engine = DiscountingBondEngine(ts)
fb_dc = QuantLib.Time.Actual360()
conv = QuantLib.Time.ModifiedFollowing()

fb_schedule = QuantLib.Time.Schedule(Date(2005, 10, 21), Date(2010, 10, 21), QuantLib.Time.TenorPeriod(QuantLib.Time.Quarterly()),
                                        QuantLib.Time.Unadjusted(), QuantLib.Time.Unadjusted(), QuantLib.Time.DateGenerationBackwards(), false,
                                        QuantLib.Time.USNYSECalendar())
fixing_days = 2
in_arrears = true
gearings = ones(length(fb_schedule.dates) - 1)
spreads = fill(0.001, length(fb_schedule.dates) - 1)
libor_3m = usd_libor_index(QuantLib.Time.TenorPeriod(Dates.Month(3)), ts)
floating_bond = FloatingRateBond(settlement_days, face_amount, fb_schedule, libor_3m, fb_dc, conv, fixing_days, fb_issue_date, fb_bond_engine, in_arrears, 100.0, gearings, spreads, cap_vol=cap_vol)

# test
@test accrual_start_date(fixedrate_bond.cashflows.coupons[1]) == Date(2007, 5, 15)
@test accrual_end_date(fixedrate_bond.cashflows.coupons[1]) == Date(2007, 11, 15)
@test ref_period_start(fixedrate_bond.cashflows.coupons[1]) == Date(2007, 5, 15)
@test ref_period_end(fixedrate_bond.cashflows.coupons[1]) == Date(2007, 11, 15)
@test QuantLib.accrual_period(fixedrate_bond.cashflows.coupons[1]) == 0.5
@test QuantLib.get_latest_coupon(fixedrate_bond.cashflows) == fixedrate_bond.cashflows.coupons[end]

# test NPV calculations
@test npv(fixedrate_bond.cashflows, ts, ts.referenceDate, ts.referenceDate) == 94.67245111530346
@test npv(fixedrate_bond.cashflows, bbIR, true, ts.referenceDate, ts.referenceDate) == 94.71618240316593
@test QuantLib.npvbps(fixedrate_bond.cashflows, ts, ts.referenceDate, ts.referenceDate) == (94.67245111530346,0.07152580522996536)

# duration and yield
@test duration(ModifiedDuration(), fixedrate_bond.cashflows, bbIR, bbIR.dc, true, ts.referenceDate, ts.referenceDate) == 6.899980120600781
@test QuantLib.yield(fixedrate_bond.cashflows, 94.67245111530346, bbIR.dc, bbIR.comp, bbIR.freq, true, ts.referenceDate, ts.referenceDate, 1.0e-10, 100, 0.05) == 0.05506693320721388

# Other calculations
@test accrued_amount(fixedrate_bond.cashflows, ts.referenceDate) == 1.5407608695652275
@test previous_cashflow_date(fixedrate_bond.cashflows, ts.referenceDate) == Date(2008, 5, 15)
@test accrual_days(fixedrate_bond.cashflows, bbIR.dc, ts.referenceDate) == 126
@test next_cashflow(fixedrate_bond.cashflows, ts.referenceDate) == 3
@test next_coupon_rate(fixedrate_bond.cashflows, ts.referenceDate) == 0.045
@test has_occurred(fixedrate_bond.cashflows.coupons[end-1], ts.referenceDate) == false
@test maturity_date(fixedrate_bond.cashflows) == Date(2017, 5, 15)
@test amount(fixedrate_bond.cashflows.coupons[1]) == 2.2499999999999964
@test QuantLib.calc_rate(fixedrate_bond.cashflows.coupons[1]) == 0.045

# Floating coupon calculations
@test QuantLib.calc_rate(floating_bond.cashflows[12]) == 0.05443944335694899
@test amount(floating_bond.cashflows[12]) == 1.3912302191220296
@test QuantLib.index_fixing(floating_bond.cashflows[12]) == 0.05343944335694899
@test accrued_amount(floating_bond.cashflows[12], ts.referenceDate) == 0.8922019883499973
@test accrued_amount(floating_bond.cashflows[1], ts.referenceDate) == 0.0


@inline function _line_interp_weights(x, x1, x2)
  lx = x2 - x1
  dx = (x - x1)/lx
  return (1 - dx, dx)
end

@inline function _tri_interp_weights(x, x1, x2, y, y1, y2)
  lx = (x2 - x1)
  dx = (x - x1)/lx

  ly = (y2 - y1)
  dy = (y - y1)/ly

  return (1 - dx, -dy + dx, dy)
end

@inline function _square_interp_weights(x, x1, x2, y, y1, y2)
  lx = (x2 - x1)
  dx = (x - x1)/lx

  ly = (y2 - y1)
  dy = (y - y1)/ly

  w = ((1 - dx) * (1 - dy),
       (1 - dx) * dy,
       dx * (1 - dy),
       dx * dy)

  return w
end

@inline function _line_interp(x, x1, x2, f1, f2)
  w = _line_interp_weights(x, x1, x2)
  return w[1] * f1 + w[2] * f2
end

@inline function _line_interp!(f, x, x1, x2, f1, f2)
  w = _line_interp_weights(x, x1, x2)
  return f .= w[1] .* f1 .+ w[2] .* f2
end

@inline function _tri_interp(x, x1, x2, y, y1, y2, v1, v2, v3)
  w = _tri_interp_weights(x, x1, x2, y, y1, y2)
  return w[1] * v1 + w[2] * v2 + w[3] * v3
end

@inline function _tri_interp!(f, x, x1, x2, y, y1, y2, v1, v2, v3)
  w = _tri_interp_weights(x, x1, x2, y, y1, y2)
  return f .= w[1] .* v1 .+ w[2] .* v2 .+ w[3] .* v3
end

@inline function _square_interp(x, x1, x2, y, y1, y2, f11, f12, f21, f22)
  w = _square_interp_weights(x, x1, x2, y, y1, y2)
  return w[1] * f11 + w[2] * f12 + w[3] * f21 + w[4] * f22
end

@inline function _square_interp!(f, x, x1, x2, y, y1, y2, f11, f12, f21, f22)
  w = _square_interp_weights(x, x1, x2, y, y1, y2)
  return f .= w[1] .* f11 .+ w[2] .* f12 .+ w[3] .* f21 .+ w[4] .* f22
end


using Test
using SDPT3

# Example page 4 of SeDuMi_Guide_105R5.pdf
@testset "Linear Programming example" begin
    blk = ["l"  4]
    c = [ 1.0
         -1.0
          0.0
          0.0]
    A = [10.0 -7.0 -1.0 0.0
          1.0  0.5  0.0 1.0]
    b = [5.0
         3.0]
    obj, X, y, Z, info, runhist = sdpt3(blk, [Matrix(A')], [c], b)
    tol = 1e-5
    @test obj[1] ≈ -1/8 atol=tol rtol=tol
    @test obj[2] ≈ -1/8 atol=tol rtol=tol
    @test X[1] ≈ [47/24, 25//12, 0, 0] atol=tol rtol=tol
    @test y ≈ [1/8, -1/4] atol=tol rtol=tol
    @test Z[1] ≈ [0, 0, 1/8, 1/4] atol=tol rtol=tol
    @test info["termcode"] == 0.0
end


## Ported Joel CFD's example to Julia with extensions: https://github.com/joel-ho/Field-inversion-example
# Currently uses Zygote for autodiff
using LinearAlgebra
using Zygote
using StaticArrays
using SparseArrays
using Base.Iterators

# Newton iteration
prod_size(tensor) = (prod ∘ size)(tensor)

newton_step(u, R, ∂R∂u) = reshape(∂R∂u \ -R[:], size(u))

## Finite differencing
central_2nd(Tm, T, Tp) = Tm - 2T + Tp
central_2nd_diff(Tm, T, Tp, h) = central_2nd(Tm, T, Tp) / h^2
central_1st_diff(Tm, Tp, h) = (Tm - Tp) / 2h

compute_q(x, y) = 5e5 * (1 - (x - 0.5)^2) * (1 - (y - 0.5)^2)

function get_T(T_grid, i, j, T_boundary)
    if i < 1 || j < 1 || i > length(T_grid[:,1])|| j > length(T_grid[1,:])
        T_boundary
    else
        T_grid[i,j]
    end
end

stencil_1(T, i, j, T_boundary) =
    [ get_T(T, i, j, T_boundary)   ;
      get_T(T, i-1, j, T_boundary) ;
      get_T(T, i+1, j, T_boundary) ;
      get_T(T, i, j-1, T_boundary) ;
      get_T(T, i, j+1, T_boundary) ]

# Difference operators specific to stencil_1
Δ²x(T, h) = central_2nd_diff(T[2], T[1], T[3], h)
Δ²y(T, k) = central_2nd_diff(T[4], T[1], T[5], k)
Δx(T, h) = central_1st_diff(T[3], T[2], h)
Δy(T, k) = central_1st_diff(T[5], T[4], k)

compute_k(β₁, β₂, β₃, β₄,  T) = β₁/T + β₂ + β₃ * T + β₄ * T^2
# + β₃ * T + β₄ * T^2
∇compute_k(β₁, β₂, β₃, β₄, T) = gradient(compute_k, β₁, β₂, β₃, β₄, T)

function residual_cell(T, β, q, dx, dy)
    k = compute_k(β..., first(T))
    ∂k∂T = (last ∘ ∇compute_k)(β..., first(T))

    k * (Δ²x(T, dx) + Δ²y(T, dy)) + ∂k∂T * (Δx(T, dx)^2 + Δy(T, dy)^2 ) + q
end

residual(T, β, q, dx, dy, T_boundary) =
        [ let Ts = stencil_1(T, T_ind.I..., T_boundary);
            residual_cell(Ts, β, q[T_ind], dx, dy) end
            for T_ind in CartesianIndices(T) ]

# Derivative stuff
function solve_direct(x, u, ∂R∂x, ∂R∂u)
    ∂R∂u_sq = reshape(∂R∂u, (length(u[:]), length(u[:])))
    reshape(reduce(hcat, ∂R∂u_sq \ -(∂R∂x)[:,:,i][:] for i in eachindex(x)), (size(u)..., length(x)))
end

function solve_adjoint(u, ∂R∂u, dfdu)
    reshape(∂R∂u, (length(u[:]), length(u[:])))' \ -(dfdu)'[:]
end

total_derivative_direct(∂f∂x, ψ, ∂f∂u) = ∂f∂x + [ sum(∂f∂u * ψ[n]) for n in eachindex(∂f∂x) ]

total_derivative_adjoint(∂f∂x, φ, ∂R∂x) = ∂f∂x + [ sum(permutedims(φ) * reshape(∂R∂x, (length(R[:]), length(∂f∂x)))[:,n]) for n in eachindex(∂f∂x) ]

∇residual_cell(Ts, βs, q, dx, dy) = gradient(residual_cell, Ts, βs, q, dx, dy)

function ∇residual_cell_T_update!(∂R∂T, grad, k, l)
    a, b, c, d = size(∂R∂T)
    ∂R∂T[k,l,k,l]            = grad[1]
    if k > 1 ∂R∂T[k,l,k-1,l] = grad[2] end
    if k < a ∂R∂T[k,l,k+1,l] = grad[3] end
    if l > 1 ∂R∂T[k,l,k,l-1] = grad[4] end
    if l < b ∂R∂T[k,l,k,l+1] = grad[5] end
end

∇residual(T, β, q, dx, dy, T_boundary) = gradient(residual, T, β, q, dx, dy, T_boundary)

∇residual_β(T, β, q₀, dx, dy, T_boundary) =
        [ let Ts = stencil_1(T, T_ind.I..., T_boundary);
            ∇residual_cell(Ts, β, q₀[T_ind], dx, dy)[2][i] end for
            T_ind in CartesianIndices(T), i in eachindex(β) ]

function compute_residual_and_grad(T, β, q, dx, dy)
    # Residual
    R = residual(T, β, q, dx, dy, T_boundary)

    # Jacobian
    T_size = size(T)
    ∂R∂T = zeros(T_size..., T_size...)
    for T_ind in CartesianIndices(T)
        Ts = stencil_1(T, T_ind.I..., T_boundary)
        grad = ∇residual_cell(Ts, β, q₀[T_ind], dx, dy)[1]
        ∇residual_cell_T_update!(∂R∂T, grad, T_ind.I...)
    end

    R, sparse(reshape(∂R∂T, size(R) .* reverse(T_size)))
end

## Grid setup
dx = 0.1
dy = 0.2
xs, ys = -1:dx:1, -1:dy:1
grid = [ (x, y) for x in xs, y in ys ]
T_boundary = 200.
βs = [13221., 27., 5., 1.]
ω = 1.0

T₀ = T_boundary * ones(size(grid))
q₀ = map(xs -> compute_q(xs...), grid);

n = length(xs)
T_truth = deepcopy(T₀)

num_iters = 30
ε = zeros(num_iters);

## NEWTON ITERATIONS
#=================================================================#

for i in 1:num_iters
    R, ∂R∂T = compute_residual_and_grad(T_truth, βs, q₀, dx, dy)
    ΔT = newton_step(T_truth, R, ∂R∂T)
    ε[i] = maximum(abs.(ΔT))
    # display(R)
    println("Newton step error L² norm: $(ε[i])")
    T_truth .+= ω * ΔT
end

##
k_truth = compute_k.(βs..., T_truth)

##
using Plots
gr()

##
plot(ε)

##
surf1 = surface(xs, ys, q₀)

##
surf2 = surface(xs, ys, T_truth)

##
surf3 = surface(xs, ys, k_truth)

## ADJOINT ITERATIONS
#=================================================================#

function cost(T, T_truth, indices)
    norm(T[index] - T_truth[index] for index in indices)
end

∇cost(T, T_truth, indices) = gradient(cost, T, T_truth, indices)

# indices = rand(1:length(T_truth[:]), (Int ∘ floor)(length(T_truth[:]) / 2))
indices = eachindex(T_truth)

##
β =  [13000., 10., 10., 5.]
γs = [100., 5., 5., 5.]

num_adjoint = 100
num_iters   = 30

T₀   = T_boundary * ones(size(grid))
q₀   = map(xs -> compute_q(xs...), grid);

T    = copy(T₀)
ε    = zeros(num_adjoint, num_iters)
εa   = zeros(num_adjoint)
β_l  = [β'; zeros(num_adjoint - 1, length(β)) ];
R    = zeros(size(T))
∂R∂T = zeros(size(R)..., size(T)...);

##
for i in 1:num_adjoint

    println("\nAdjoint Iteration — $i")
    display(β)

    for j in 1:num_iters
        R, ∂R∂T = compute_residual_and_grad(T, β, q₀, dx, dy)
        ΔT      = newton_step(T, R, ∂R∂T)
        ε[i,j]  = maximum(abs.(ΔT))
        # println("Newton step error: $(ε[i,j])")
        T     .+= ω * ΔT
    end

    εa[i] = cost(T, T_truth, indices)
    println("\nCost function value: $(εa[i])")

    # println("Computing df/dT")
    dfdT = ∇cost(T, T_truth, indices)[1]

    # println("Computing ∂f/∂x")
    ∂f∂x = zeros(length(β))

    # println("Computing ∂R/∂x")
    ∂R∂x = ∇residual_β(T, β, q₀, dx, dy, T_boundary)

    # display(∂R∂x)

    # println("Solving adjoint problem: (∂R/∂T)ᵗ φ = -(df/dT)ᵗ")
    # φ = solve_adjoint(T, ∂R∂T, dfdT)
    # dfdβ = total_derivative_adjoint(∂f∂x, φ, ∂R∂x)

    dTdx = solve_direct(β, T, ∂R∂x, ∂R∂T)
    dfdβ = total_derivative_direct(∂f∂x, dTdx, dfdT)

    # display(dfdβ)

    # println("Updating design variables β")
    @. β -= γs * dfdβ
    β_l[i,:] = β'
end

##
display(β)

##
k = compute_k.(β..., T)

##
using Plots
plotlyjs()

##
plot(εa, xlabel = "Iterations", ylabel = "Grid Error")

##
surface(xs, xs, T)
surface!(xs, xs, T_truth)

##
surface(xs, xs, k)
surface!(xs, xs, k_truth)

##
# residual_discrete_k(T, k, q, h) = first(k) * (Δ²x(T, h) + Δ²y(T, h)) + Δx(k, h) * Δx(T, h) + Δy(k, h) * Δy(T, h) + q
# ∇residual_discrete_k(Ts, βs, q, h) = gradient(residual_discrete_k, Ts, βs, q, h)

# function update_k_grid_ghost_points!(k_grid)
#     k_grid[1,:] = 2 * k_grid[2,:] .- k_grid[3,:]
#     k_grid[end,:] = 2 * k_grid[end-1,:] .- k_grid[end-2,:]
#     k_grid[:,1] = 2 * k_grid[:,2] .- k_grid[:,3]
#     k_grid[:,end] = 2 * k_grid[:,end-1] .- k_grid[:,end-2]

#     k_grid
# end


using CalcDist
using Test

# tolerance level
TOL8 = 1.0e-8
TOL6 = 1.0e-6
TOL4 = 1.0e-4
TOL = TOL8

@testset "CalcDist.jl" begin
   @testset "Binomial" begin
      # binompdf
      @test_throws DomainError binompdf(10, -0.1, 4)
      @test_throws DomainError binompdf(-5, 0.2, 3)
         # single
      @test only(binompdf(10, 0.5, 5)) == 0.24609375
      @test only(binompdf(20, 0.95, 8)) ≈ 2.0403e-11 atol=TOL
         # vector
      @test binompdf(5, 0.6, [3,4,5]) ≈ [0.3456, 0.2592, 0.07776] atol=TOL
         # missing x
      @test binompdf(5, 0.5) == [0.03125, 0.15625, 0.3125, 0.3125, 0.15625, 0.03125]

      # binomcdf
      @test_throws DomainError binomcdf(10, -0.1, 4)
      @test_throws DomainError binomcdf(-5, 0.2, 3)
         # single
      @test only(binomcdf(5, 0.6, 3)) ≈ 0.66304 atol=TOL
         # vector
      @test binomcdf(5, 0.6, [3,4,5]) ≈ [0.66304, 0.92224, 1.0] atol=TOL
         # missing x
      @test binomcdf(4, 0.25) ≈ [0.31640625, 0.73828125, 0.94921875, 0.99609375, 1.0] atol=TOL
   end

   @testset "Poisson" begin
      # poissonpdf
      @test_throws DomainError poissonpdf(-2, 5)
         # single
      @test only(poissonpdf(6, 10)) ≈ 0.0413030934 atol=TOL
         # vector
      @test poissonpdf(6, [8,9,10]) ≈ [0.10325773, 0.06883849, 0.04130309] atol=TOL

      # poissoncdf
      @test_throws DomainError poissoncdf(-4, 7)
         # single
      @test only(poissoncdf(0.1875, 2)) ≈ 0.99904486 atol=TOL
         # vector
      @test poissoncdf(0.126, [0,1,2,3]) ≈ [0.88161485, 0.99269832, 0.99969658, 0.99999050] atol=TOL
   end

   @testset "Geometric" begin
      # geometpdf
      @test_throws DomainError geometpdf(-0.5, 6)
      @test_throws DomainError geometpdf(1.2, 6)
         # single
      @test only(geometpdf(0.4, 6)) ≈ 0.031104 atol=TOL
         # vector
      @test geometpdf(0.25, [2,3,4]) == [0.1875, 0.140625, 0.10546875]

      # geometcdf
      @test_throws DomainError geometcdf(-0.4, [4,5,6])
      @test_throws DomainError geometcdf(1.4, [4,5,6])
         # single
      @test only(geometcdf(0.5, 2)) == 0.75
         # vector
      @test geometcdf(0.5, [1,2,3]) == [0.5, 0.75, 0.875]
   end

   @testset "Counting" begin
      # nCr
      @test_throws DomainError nCr(-9, 4)
      @test_throws DomainError nCr(9, [4,-5,6])
         # single
      @test only(nCr(9,4)) == 126
         # vector
      @test nCr(9, [5,4,3]) == [126, 126, 84]
      @test nCr([8,7,6], 3) == [56, 35, 20]

      # nPr
      @test_throws DomainError nCr(-8, 4)
      @test_throws DomainError nCr(8, [2,-4,5])
         # single
      @test only(nPr(7,3)) == 210
         # vector
      @test nPr(7, [1,2,3]) == [7, 42, 210]
      @test nPr([5,6,7], 2) == [20, 30, 42]
   end
end


# Arithmetic Operators

x = 1;
y = 2;

+x;
-x;
x + y;
x - y;
x * y;
x / y; # Floating point division
x ÷ y; # Integer division
x \ y; # Inverse division
y ^ 2;
x % y;

# Bitwise Operators

b = Int8(10);
println(bitstring(b));

~b; # bitwise NOT
b & b; # bitwise and
b | b; # bitwise or
b ⊻ b; # bitwise xor
b >>> 2; # logical right shift
b >> 2; # bitwise right shift
b << 2; # bitwise left shift

# Logical Operators

a = true;
b = false;

a && b;
a || b;
!a;

# Assignment Operators

x = 1;
x += 1;
x -= 1;
x *= 1;
x /= 1;

# Vectorized/dot Operators
# - tells compiler to do operator to all elements in array

A = [1, 2, 3, 4];
println(A);
A = A .^ 2;
println(A);
A .^= 2;
println(A);

# Relational Operators

x = 2;
y = 3;

x == y;
x < y;
x > y;
x <= y;
x >= y;
x != y;


@testset "$(rpad("Tableau coefficients",80))" begin

    # test Gauss-Legendre Runge-Kutta coefficients

    glrk2_tab1 = getCoefficientsGLRK(2)
    glrk2_tab2 = getCoefficientsGLRK(2, high_precision=false)
    @test glrk2_tab1.a ≈ glrk2_tab2.a atol=4eps()
    @test glrk2_tab1.b == glrk2_tab2.b
    @test glrk2_tab1.c ≈ glrk2_tab2.c atol=4eps()

    glrk3_tab1 = getCoefficientsGLRK(3)
    glrk3_tab2 = getCoefficientsGLRK(3, high_precision=false)
    @test glrk3_tab1.a ≈ glrk3_tab2.a atol=4eps()
    @test glrk3_tab1.b == glrk3_tab2.b
    @test glrk3_tab1.c ≈ glrk3_tab2.c atol=4eps()

    @test getCoefficientsGLRK(1) == getCoefficientsGLRK1()
    @test getCoefficientsGLRK(2) == getCoefficientsGLRK2()
    @test getCoefficientsGLRK(3) == getCoefficientsGLRK3()
    @test getCoefficientsGLRK(4) == getCoefficientsGLRK4()
    @test getCoefficientsGLRK(5) == getCoefficientsGLRK5()

    # test Gauss-Lobatto Runge-Kutta coefficients

    @test getCoefficientsLobIII(2) == getCoefficientsLobIII2()
    @test getCoefficientsLobIII(3) == getCoefficientsLobIII3()
    @test getCoefficientsLobIII(4) == getCoefficientsLobIII4()

    @test getCoefficientsLobIIIA(2) == getCoefficientsLobIIIA2()
    @test getCoefficientsLobIIIA(3) == getCoefficientsLobIIIA3()
    @test getCoefficientsLobIIIA(4) == getCoefficientsLobIIIA4()

    @test getCoefficientsLobIIIB(2) == getCoefficientsLobIIIB2()
    @test getCoefficientsLobIIIB(3) == getCoefficientsLobIIIB3()
    @test getCoefficientsLobIIIB(4) == getCoefficientsLobIIIB4()

    @test getCoefficientsLobIIIC(2) == getCoefficientsLobIIIC2()
    @test getCoefficientsLobIIIC(3) == getCoefficientsLobIIIC3()
    @test getCoefficientsLobIIIC(4) == getCoefficientsLobIIIC4()

    @test getCoefficientsLobIIID(2) == getCoefficientsLobIIID2()
    @test getCoefficientsLobIIID(3) == getCoefficientsLobIIID3()
    @test getCoefficientsLobIIID(4) == getCoefficientsLobIIID4()

    @test getCoefficientsLobIIIE(2) == getCoefficientsLobIIIE2()
    @test getCoefficientsLobIIIE(3) == getCoefficientsLobIIIE3()
    @test getCoefficientsLobIIIE(4) == getCoefficientsLobIIIE4()

    # test PGLRK coefficients

    pglrk2 = getCoefficientsPGLRK(2)
    @test typeof(pglrk2) <: CoefficientsPGLRK

end


using BayesianNonparametrics
using Test
using Distributions

# write your own tests here
include("distributionTests.jl")
include("dpmTests.jl")


using Printf
using LinearAlgebra

include("../common/PWGrid_v01.jl")
include("../common/wrappers_fft.jl")
include("../common/gen_lattice.jl")

include("structure_factor.jl")
include("gen_dr.jl")
include("gen_rho.jl")
include("write_xsf.jl")

function test_main()
    Ns = [64,64,64]
    a = 16.0
    LatVecs = gen_lattice_sc(a)
    pwgrid = PWGrid(Ns,LatVecs)

    Npoints = prod(Ns)
    Ω = pwgrid.Ω
    G2 = pwgrid.G2

    # Atomic positions and nuclear charge
    #Xpos = reshape( [5.0, 5.0, 5.0,
    #                 6.5, 5.0, 5.0], (3,2) )
    Xpos = reshape( [0.0, 0.0, 0.0], (3,1) )

    # Calculate structure factor
    @printf("Npoints = %8d\n", Npoints)
    Sf = structure_factor( Xpos, pwgrid.G )

    center = 0.5*a*ones(3)
    println(center)
    dr = gen_dr( pwgrid.r, center )
    rho = gen_rho( Ns, dr, 0.25, Sf )

    sumSf = sum(Sf)
    @printf("sum(Sf) = (%15.10e,%15.10e)\n", real(sumSf), imag(sumSf))
    @printf("integRho = %18.10f\n", sum(rho)*Ω/Npoints)

    write_xsf( "rho.xsf", LatVecs, Xpos )
    write_xsf_3d_crystal( "rho.xsf", Ns, LatVecs, rho )
end

@time test_main()


### contains functions to compute probabilities in subsets
### which we used to call partitions
### the functions are in the process of being rewritten for
### multisets

"""
	matrix_phi(k, U, occupancy_vector, n)
"""
function matrix_phi(k, U, occupancy_vector, n)

	"""matrix P(phi_k) in the notes

	n = number of photons
	occupancy_vector is a Int64 vector of {0,1}^m with one if the mode belongs to the partition, zero otherwise

	remark : opposite sign convention from the notes called notes_singlemode_output
	by default for the r = (1,1, ..., 1) input """

	if k > n
		throw(ArgumentError("k > n"))
	end

	check_at_most_one_particle_per_mode(occupancy_vector)

	m = size(U,1)
	mat = Matrix{eltype(U)}(I, size(U))

	for i in 1:m
		if occupancy_vector[i] == 1
			mat[i,i] = exp(-2im * pi * k / (n+1))
		end
	end

	mat

end

"""
	proba_partition_partial(;U, S, occupancy_vector, input_state, checks=true)

Return a ``n+1`` sized array giving the probability to find `[0,1,...]`, photons
inside the bins given by `occupancy_vector` at the output of `U`.

!!! note
	- We take `U` of dimension ``m`` while `M` is the scattering matrix, as in
	  Tichy, ``M_ij = U_{d_i}, _{d_j}``.
	- Given ``n`` photons, generally in the first modes, the distinguishability
	  matrix is defined as in Tichy, ``S_{ij} = <phi_{d_i}|phi_{d_j}>``.
	  This is not a problem as it does not depend on the output partition but be aware of it.
"""
function proba_partition_partial(; U, S, occupancy_vector, input_state, checks = true)

	"""returns a n+1 sized array giving the probability of have [zero, one, ...]
	photons inside the bins given by occupancy_vector for the interferometer
	given by the matrix U (perfectly indistinguishable photons)

	like proba_partition but with partial distinguishability defined through the S
	matrix, with conventions as in Tichy (note that we take the S matrix to be n*n
	while the interferometer has m modes)

	NOTE :
	in the following, we take U to be m*m
	while M is the scattering matrix, as in Tichy, M_ij = U_{d_i}, _{d_j}
	and have n photons, generally in the first n modes
	the distinguishability matrix is defined as in tichy, to be n*n,
	S_{ij} = <phi_{d_i}|phi_{d_j}>
	this is not a problem as it does not depend on the output partition but be aware of it
	"""

	m = size(U,1)
	n = sum(input_state)

	if n == 0
		throw(ArgumentError("number of photons ill defined, zero cannot be computed but trivial"))
	end

	if size(S,1) != n
		throw(ArgumentError("S matrix does not have the size required for this number of photons"))
	end

	if checks
		if !is_unitary(U)
			throw(ArgumentError("U not unitary"))
		elseif !is_a_gram_matrix(S)
			throw(ArgumentError("S not gram"))
		else
			println("Note : checking unitarity of U, that S is gram, slows down significantly probabilities computations")
		end
	end

	function scattering_matrix_amplitudes(U, input_state, occupancy_vector, k)

	    """U tilde in the handwritten notes, corresponds to U^dagger Lambda(eta) U"""

	    scattering_matrix(U' * matrix_phi(k, U, occupancy_vector, sum(input_state)) * U, input_state, input_state)

	end

	proba_fourier = 1/vector_factorial(input_state) .* [ryser(copy(transpose(S)) .* scattering_matrix_amplitudes(U, input_state, occupancy_vector, k)) for k in 0:n]

	if length(proba_fourier) == 0
		throw(ArgumentError("cannot compute the idft of an empty array"))
	end

	p_recovered = [1/(n+1) * sum([exp(2im * pi * j * k / (n+1)) * proba_fourier[k+1] for k in 0:n]) for j in 0:n] # direct inverse fourier transform

	if any([!isapprox(imag(p_recovered[i]),  0., atol = 1e-7, rtol = 1e-5) for i in 1:length(p_recovered)])
		println("WARNING : proba[$j] has a significant imaginary part of $(imag(tmp))")
	elseif !(isapprox(sum(real.(p_recovered)), 1., rtol = 1e-5))
		println("WARNING : probabilites do not sum to one ($(sum(real.(tmp)))), maybe be erreneous")
	end

	real.(p_recovered)

end

"""
	proba_partition_bosonic(;U, occupancy_vector, input_state=ones(Int, size(U,1)), checks=true)

Indistinguishable version of [`proba_partition_partial`](@ref).
"""
function proba_partition_bosonic(;U, occupancy_vector, input_state = ones(Int, size(U,1)), checks = true)

	"""indistinguishable version of proba_partition_partial"""

	n = sum(input_state)

	proba_partition_partial(U = U , S = ones(n,n), occupancy_vector = occupancy_vector, input_state = input_state, checks = checks)

end

function proba_partition_distinguishable(;U, occupancy_vector, input_state = ones(Int, size(U,1)), checks = true)

	"""indistinguishable version of proba_partition_partial"""

	n = sum(input_state)

	proba_partition_partial(U = U , S = Matrix{eltype(U)}(I,n,n), occupancy_vector = occupancy_vector, input_state = input_state, checks = checks)

end

"""
	partition_probability_distribution_distinguishable_rand_walk(part, U)

Generate a vector giving the probability to have ``k`` photons in the partition
`part` at the output of the interferomter `U`.
"""
function partition_probability_distribution_distinguishable_rand_walk(part, U)

	"""generates a vector giving the probability to have k photons in
	the partition part for the interferometer U by the random walk method
	discussed in a 22/02/21 email
	part is written like (1,2,4) if it is the output modes 1, 2 and 4"""

	function proba_photon_in_partition(i, part, U)
		"""returns the probability that the photon i ends up in the set of outputs part, for distinguishable particles only"""

		sum([abs(U[i, j])^2 for j in part]) # note the inversion line column
	end

	function walk(probability_vector, walk_number, part, U)

		new_probability_vector = similar(probability_vector)
		n = size(U)[1]
		proba_this_walk = proba_photon_in_partition(walk_number, part, U)

		for i in 1 : n+1
			new_probability_vector[i] = (1-proba_this_walk) * probability_vector[i] + proba_this_walk * probability_vector[i != 1 ? i-1 : n+1]
		end

		new_probability_vector

	end

	n = size(U)[1]
	probability_vector = zeros(n+1)
	probability_vector[1] = 1

	for walk_number in 1 : n
		probability_vector = walk(probability_vector, walk_number, part, U)
	end

	if !(isapprox(sum(probability_vector), 1., rtol = 1e-8))
		println("WARNING : probabilites do not sum to one ($(sum(probability_vector)))")
	end

	probability_vector

end


using LinearAlgebra, Statistics
using Distributions, QuasiMonteCarlo, Random

n = 10000
d = 10

# synthetic d-dimensional least squares problem
wgen = rand(d);
w0 = zeros(d);

# Quasi-random distributions
qrngs = Dict([("IID Uniform"     , UniformSample()),
              ("Sobol"           , SobolSample()),])

qrng_names = collect(keys(qrngs))

# data distribution
# x ~ N(0,I_d)
# y ~ N(x'wopt, 1)
# goal is to learn w, which is to minimize E[0.5*(x'w - y)^2]

# online mode:  draw samples from this integral at every SGD iteration
# offline mode: draw n samples iid uniformly from this integral
#               to form a training set,
#               then draw one sample (possibly following a low-discrepancy sequence)
#               from the training set at every SGD iteration

function gen_xy(q::Vector{Float64})
    xi = quantile(Normal(), q[1:d])
    yi = quantile(Normal(xi'wgen, 1), q[end])
    return (xi, yi)
end

function gen_syn_data(num_examples::Int64, qrng_name::String, dim::Int64, seed::Int64)
    Random.seed!(seed)

    lb = zeros(dim)
    ub = ones(dim)

    distrib = qrngs[qrng_name]
    qmc_samples = QuasiMonteCarlo.sample(num_examples, lb, ub, distrib)

    syn_data = []
    for i in 1:num_examples
        q = mod.(qmc_samples[:,i], 1)
        push!(syn_data, gen_xy(q))

    end
    return syn_data
end

function sgd(T::Int64, w0::Vector{Float64}, wopt::Vector{Float64},
	samples::Vector{Any}, ordering::Vector{Int64})
	subopt = zeros(T)
	wk = w0
	# diminishing step size optimized for IID Uniform
	alpha = 1/(d + 2 + d/norm(wgen)^2);
	for t in 1:T
		(xi, yi) = samples[ordering[t]]
		alpha = 1/(1/alpha + (1 - alpha*(d+2))/(1 - alpha))
		wk -= alpha * (xi'wk - yi) * xi
		subopt[t] = norm(wk - wopt).^2
	end
	return subopt
end


using Test
using MechGlueModelingToolkit
import MechGlueModelingToolkit.Unitfu
using Unitfu:s, @u_str
using ModelingToolkit
using ModelingToolkit: Symbolic, Variable, Differential, @syms, @variables

t = Variable{typeof(1.0s)}(:t)()
x = Variable{u"kg"}(:x)(t)
y = Variable{u"kg"}(:y)(t)
D = Differential(t)

eq1 = x ~ y*t
eq2 = x*10u"s" ~ y*t

@test ModelingToolkit.instantiate(t) == 1u"s"
@test ModelingToolkit.instantiate(x) == 1u"kg"
@test ModelingToolkit.instantiate(y) == 1u"kg"

@test !ModelingToolkit.validate(eq1)
@test ModelingToolkit.validate(eq2)

eqs = [
        D(x) ~ y/t
        D(y) ~ (x*y)/(t*10u"kg")
]

sys = ODESystem(eqs,t,[x,y],[])
@test ModelingToolkit.validate(sys)

@testset "MechGlueModelingToolkit.jl" begin
    # Write your tests here.
end


export Isosceles, Trapezoid  # constructor-like methods (a.k.a factory methods)

#= Special polygons =#
# The constructors for the special polygons are defined as "factory methods" that construct
# vertices and pass them to Polygon().

# To-dos: parallegram, rhombus, ...

# Isosceles triangle
function Isosceles(base::NTuple{2,SVector{2,<:Real}},  # (end point 1, end point 2): two end points of base
                   h::Real)  # height drawn normal to base; direction is π/2 from base[2]-base[1]
    m = (base[1] + base[2]) / 2  # midpoint of base
    b̂ = normalize(base[2] - base[1])  # unit direction of base
    ĥ = @SVector [-b̂[2], b̂[1]]  # unit direction of height
    p = m + h.*ĥ  # apex

    v = [base[1] base[2] p]  # vertices

    return Polygon(v)
end

Isosceles(base::NTuple{2,AbstractVector{<:Real}}, h::Real) = Isosceles(SVector{2}.(base), h)

function Isosceles(c::SVector{2,<:Real},  # midpoint of base
                   b::Real,  # base length
                   h::Real,  # height
                   θ::Real=0.0)  # height direction measured from +y-direction (= base direction measured from +x-direction )
    b2 = 0.5b
    b̂ = @SVector [cos(θ), sin(θ)]
    base = (c - b2.*b̂, c + b2.*b̂)

    return Isosceles(base, h)
end

Isosceles(c::AbstractVector{<:Real}, b::Real, h::Real, θ::Real=0.0) = Isosceles(SVector{2}(c), b, h, θ)

# Trapezoid
function Trapezoid(base::NTuple{2,SVector{2,<:Real}},  # (end point 1, end point 2): two end points of base
                   h::Real,  # height drawn normal to base; direction is π/2 from base[2]-base[1]
                   ϕ::NTuple{2,Real})  # (base angle 1, base angle 2)
    b̂ = normalize(base[2] - base[1])  # unit direction of base
    ĥ = @SVector [-b̂[2], b̂[1]]  # unit direction of height

    t₁ = base[1] + (h*cot(ϕ[1])) .* b̂ + h.*ĥ
    t₂ = base[2] - (h*cot(ϕ[2])) .* b̂ + h.*ĥ

    v = [base[1] base[2] t₂ t₁]

    return Polygon(v)
end

Trapezoid(base::NTuple{2,AbstractVector{<:Real}}, h::Real, ϕ::NTuple{2,Real}) =
    Trapezoid(SVector{2}.(base), h, ϕ)

function Trapezoid(c::SVector{2,<:Real},  # midpoint of base
                   b::Real,  # base length
                   h::Real,  # height
                   ϕ::NTuple{2,Real},  # (base angle 1, base angle 2)
                   θ::Real=0.0)  # height direction measured from +y-direction (= base direction measured from +x-direction)
    b2 = 0.5b
    b̂ = @SVector [cos(θ), sin(θ)]
    base = (c - b2.*b̂, c + b2.*b̂)

    return Trapezoid(base, h, ϕ)
end

Trapezoid(c::AbstractVector{<:Real}, b::Real, h::Real, ϕ::NTuple{2,Real}, θ::Real=0.0) =
    Trapezoid(SVector{2}(c), b, h, ϕ, θ)

# Isosceles trapezoid
Trapezoid(base::NTuple{2,AbstractVector{<:Real}}, h::Real, ϕ::Real) = Trapezoid(base, h, (ϕ,ϕ))
Trapezoid(c::AbstractVector{<:Real}, b::Real, h::Real, ϕ::Real, θ::Real=0.0) = Trapezoid(c, b, h, (ϕ,ϕ), θ)


# Generate Small example, d=800, n=1601, 2n=3202
example='L'
K,M=GenerateExamples(example)

# Create DampedSystem structure
DS=DampedSystem(M,K,[1.0 0;0 1])

# Compute changes of basis
Ω,Φ,Ξ=ChangeOfBasis(DS)

# Eigenvalue decomposition of a CSymDPR1 matrix.
# l is position of damper, ρ is viscosity
T=ComplexF64
l=50
ρ=50.0
y=Φ[l,:]
A=CSymDPR1(Ξ,y,map(T,ρ))
println("Time of single eigenvalue decomposition of a CSymDPR1 matrix.")
println("Second timing is relevant\n")
@time Λ,S=eigen(A)
@time Λ,S=eigen(A)

# Computation of trace(X) for a single damper in position l with viscosity ρ, and
# the number of damped frequencies s
s=27

function f(ρ::Number)
    A=CSymDPR1(Ξ,Φ[l,:],ComplexF64(ρ))
    Λ,S=eigen(A)
    Tr=traceX(S,s)
    return real(Tr)
end
println("\nTime of trace computation for a single damper.")
println("Second timing is relevant.\n")
@time Tr=f(50.0)
@time Tr=f(50.0)

#  MULTIPLE DAMPERS WITH DIFFERENT VISCOSITIES
# Computation of trace(X) for three dampers in positions l with viscosities ρ, and
# the number of damped frequencies s

# Positions of dampers and viscosities
l=[50,950,220]
ρ=[100.0,200.0,300.0]

# Auxiliary arrays
n₁=800
n=2*n₁+1
U=Matrix{T}(undef,2n,length(l))
Uvec=Matrix{T}(undef,2n,length(l))

function f(ρ::Vector)
    Uvec[:,1]=Φ[l[1],:]
    U[:,1]=Φ[l[1],:]
    A=CSymDPR1(Ξ,Φ[l[1],:],ComplexF64(ρ[1]))
    Λ,S=eigen(A)
    y=similar(Λ)
    # Loop
    for j=2:length(l)
        U[:,j]=Φ[l[j],:]
        if j==2
           U[:,j]-=Φ[l[j+1]+n₁,:]
        end
        lmult!(y,S,U[:,j])
        Uvec[:,j]=y
        A=CSymDPR1(Λ,y,ComplexF64(ρ[j]))
        Λ,S₁=eigen(A)
        # Multiplication of linked Cauchy-like matrices
        S=S*S₁
    end
    Tr=traceX(S,s)
    println("ρ = ", ρ," trace = ", Tr)
    return real(Tr)
end

println("\nTime of trace computation for three dampers.")
println("Second timing is relevant.\n")
@time Tr=f(ρ)
@time Tr=f(ρ)

# Optimization. Given starting ρ, we compute the optmial viscosities
ρ=[100.0,100,100]
println("\nTime for optimization of viscosities for three dampers.\n")
@time result=optimize(f,ρ,method=ConjugateGradient(eta=0.01),
    x_tol=1e-2,f_tol=1e-6)

println(result)
println("\nOptimal viscosities: ",result.minimizer)
println("\nMinimal trace: ", result.minimum)


@doc """
Takes the singles sided FFT. returns (frequencies, complex 1-sided FT)
"""-> function FFT_1s(TimeDomain::Vector; fs = 1, Window = "Hanning", PlotOn = false)
    ## This function takes in a time domain vector, optional sampling rate "fs," and optional Hanning window
    L = length(TimeDomain)
    if Window == "Hanning"
        CorFac = 2
        Window = WindowHanning(L)
    elseif (Window == "none")|(Window == "None")
        CorFac = 1
        Window = ones(L)
    else
        error("Window not recognized")

    end
    WindowedData = TimeDomain .* Window
    if L % 2 == 1
        EndIndex = Int(round(L / 2))
    else
        EndIndex = Int(round(L / 2))
    end
    FTData_1s = fft(WindowedData)[1:EndIndex]
    FTData_1s = FTData_1s ./ L #Corrects for vector length
    freqs = fftfreq(L, fs)[1:EndIndex]

    if PlotOn
        plot(fftfreq(L), abs.(fft(WindowedData)) ./ L)
        plot(freqs, abs.(FTData_1s), "ro")
    end

    return freqs, FTData_1s
end

@doc """
Creates a Hanning window of length N
"""-> function WindowHanning(N)
    return 0.5 .- 0.5 .* cos.(2 .* pi .* collect(0:(N-1)) ./ (N - 1))
end

@doc """
This is a documentation test2
"""-> function known_aliasing(
    Data::Vector,
    fs,
    f_unwrap::Union{Float64,Vector};
    FreqBuf::Int = 1,
    oversamp::Int = 1,
    PlotOn=false
)

    freqs, ftData = FFT_1s(Data, fs = fs)
    L = length(Data)
    L_1s = length(ftData)
    T = 1 / fs
    T_vec = 0:T:((length(Data)-1)*T)
    df = fs / L

    MaxFreq = maximum(f_unwrap)
    ScaleFac = Int(ceil(MaxFreq / (fs / 2))) * oversamp
    new_fs = ScaleFac * fs
    new_L = ScaleFac * L
    new_L_1s = Int(round(new_L / 2))

    newFreqs = fftfreq(ScaleFac * L, new_fs)[1:new_L_1s]
    NewSpectra = Complex.(zeros(length(newFreqs)))
    NewSpectra[1:L_1s] = ftData[:]

    MovedIndex = [0] #freqs which have been moved to a new location due to aliasing already

for i = 1:length(f_unwrap)
    if f_unwrap[i] > fs / 2
        AliasedFreq = calcAlias(fs, f_unwrap[i]) #the freq. the frequency has aliased to
        AliasedIndex = Int(round(AliasedFreq / df)) #the index of the aliased frequency
        if AliasedIndex == 0
            AliasedIndex = 1
        end
        NewIndex = Int(round(f_unwrap[i] / df)) #The new (corrected) index

        if findfirst(x -> x == NewIndex, MovedIndex) == nothing
            push!(MovedIndex, NewIndex)

            IndexDistFromEdge = [
                AliasedIndex,
                length(AliasedFreq) - AliasedIndex,
                NewIndex,
                new_L_1s - NewIndex,
            ]
            if minimum(abs.(IndexDistFromEdge)) < FreqBuf
                FreqBuf = minimum(abs.(IndexDistFromEdge))
            end

            StartInd_Alias = AliasedIndex - FreqBuf
            StopInd_Alias = AliasedIndex + FreqBuf
            StartInd_New = NewIndex - FreqBuf
            StopInd_New = NewIndex + FreqBuf
            if StopInd_Alias>length(ftData)
                FreqBuf = FreqBuf-(StopInd_Alias-length(ftData))
            end
            if StartInd_Alias<0
                FreqBuf = 0
            end
            if StopInd_New>length(NewSpectra)
                FreqBuf = FreqBuf-(StopInd_New-length(NewSpectra))
            end
            if StartInd_New<0
                FreqBuf = 0
            end

            StartInd_Alias = AliasedIndex - FreqBuf
            StopInd_Alias = AliasedIndex + FreqBuf
            StartInd_New = NewIndex - FreqBuf
            StopInd_New = NewIndex + FreqBuf

            println(length(ftData))
            println(StartInd_Alias)
            println(StopInd_Alias)
            println(StartInd_New)
            println(StopInd_New)

            NewSpectra[StartInd_New:StopInd_New] =
                ftData[StartInd_Alias:StopInd_Alias]
            NewSpectra[StartInd_Alias:StopInd_Alias] =
                zeros(length(StartInd_Alias:StopInd_Alias))

            println("Actual freq: $(f_unwrap[i]). Aliased to: $(AliasedFreq)")
        end
    end
end

FullSpectraNew = FFT_1s_to2s(NewSpectra)
NewTimeDomain = ifft(FullSpectraNew .* length(FullSpectraNew))
T_vec_new = 0:T/ScaleFac:((length(NewTimeDomain)-1)*T/ScaleFac)
if PlotOn
    figure(1)
    subplot(211)
    plot(freqs, abs.(ftData), "g",label="Orig data")
    stem(newFreqs, abs.(NewSpectra),label="New data")
    legend()

    subplot(212)

    # figure(2)
    plot(T_vec, Data, "r",label="Orig data")
    plot(T_vec_new, NewTimeDomain, "g",label="New data")
    legend()
end
    return FullSpectraNew, NewTimeDomain, T_vec_new

end

@doc """
Calculates the frequency some other freq. will alias to given a sampling rate, fs
calcAlias(fs, f)
fs is sampling rate
f is the original freq
"""-> function calcAlias(fs, f)
    #fs is the sampling rate
    #df is the df in the DFT
    #f is the actual waveform frequency

    Nyqu = fs / 2

    Wraps = Int(floor(f / Nyqu)) #numer of full wraparounds
    WrapDist = f % Nyqu # how far past the last wrap it will be

    if (Wraps % 2) == 0#if it is even, the wrap will be "backwards"
        return AliasedFreq = WrapDist
    else
        return AliasedFreq = Nyqu - WrapDist
    end
end

function FFT_1s_to2s(ftData)
    return cat(ftData[:], conj.(reverse(ftData[2:end]))[:], dims = 1)
end


#===========================================================================================
	Sub-Module		Advection Equation
============================================================================================#
export initialGaussian, rhsAdvection

# The state vector describing how we evolve the advection equation in time
export AdvectionState
struct AdvectionState{T}
    time::T
    u::FinDiff2018.PLCFun{T, T}
end

# State vectors can be scaled and added
function Base. *(a::T, s::AdvectionState{T})::AdvectionState{T} where {T}
    AdvectionState{T}(s.time, a * s.u)
end

function Base. +(s1::AdvectionState{T},
                 s2::AdvectionState{T})::AdvectionState{T} where {T}
    @assert abs(s1.time - s2.time) <= 100*eps(T)
    AdvectionState{T}(s1.time, s1.u + s2.u)
end

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

# A Gaussian, centred at x=1/2, with a width of 1/10
#export guassian
function gaussian(x::T)::T where {T}
    exp(- ((x - 0.5) * 10)^2)
end

# Define initial conditions: A Gaussian at t=0
#export initialGaussian
function initialGaussian(nlines::Int)::AdvectionState{Float64}
    t = 0
    u = samplePLC(Float64, Float64, nlines, gaussian)
    AdvectionState{Float64}(t, u)
end

# The RHS (Right Hand Side) of the advection equation, including boundary conditions
#export rhsAdvection
function rhsAdvection(s::AdvectionState{T})::AdvectionState{T} where {T}
    t = s.time
    u = s.u
    nlines = length(u.points) - 1
    # Calculate spatial derivative
    ux = derivRight(u)
    # Define time derivative as the spatial derivative almost everywhere,
    # except at the right boundary where we set the time derivative to zero
    ut = PLCFun{T, T}([[ux.points[i] for i in 1:nlines]; 0])
    AdvectionState{T}(t, ut)
end


using FrameworkPacote
using NLPModels
using Test

@testset "Básicos" begin
  nlp = ADNLPModel(
    x -> (x[1] - 1)^2 + (x[2] - 2)^2 / 4,
    zeros(2)
  )
  output = gradiente(nlp)
  @test isapprox(output.solution, [1.0; 2.0], rtol=1e-4)
  @test output.objective < 1e-4
  @test output.dual_feas < 1e-4
  @test output.status == :first_order

  output = gradiente(nlp, max_iter = 0)
  @test output.status == :max_iter
  output = gradiente(nlp, max_time = 0)
  @test output.status == :max_time
end


function flux(u, q, w, v, x, t, mu, eta)
    f = mu[1]*q + u[1]*mu[2];
    return f;
end
function source(u, q, w, v, x, t, mu, eta)
    x1 = x[1];
    nu = mu[1];
    c = mu[2];
    c2 = c^2;
    s = c2*nu*(sin(c*x1) - cos(c*x1));
    return s;
end
function ubou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    ub = 0.0;
    return ub;
end
function fbou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    f = flux(u, q, w, v, x, t, mu, eta);
    fb = f[1]*n[1] + tau[1]*(u[1]-uhat[1]);
    return fb;
end
function initu(x, mu, eta)
    u0 = 0.0;
    return u0;
end


using DiffEqFlux, OrdinaryDiffEq, Test

function f(u, p, t)
    p .* u
end
rc=62
ps = repeat([-0.001], rc)
tspan = (7.0, 84.0)
u0 = 3.4 .+ ones(rc)
t = collect(range(minimum(tspan), stop=maximum(tspan), length=157))
prob = ODEProblem(f, u0, tspan, ps)
data = Array(solve(prob, Tsit5(), saveat = t, abstol=1e-12, reltol=1e-12))
ptest = ones(rc)

u′,u = collocate_data(data,t,SigmoidKernel())
@test sum(abs2,u - data) < 1e-8

function loss(p)
    cost = zero(first(p))
    for i in 1:size(u′,2)
      _du = f(@view(u[:,i]),p,t[i])
      u′i = @view u′[:,i]
      cost += sum(abs2,u′i .- _du)
    end
    sqrt(cost)
end
@test loss(ptest) ≈ 418.3400017500223


module helper
#Solves for steady state distribution of workforce across age stages
function solveStages(var, fixedParam)
    NC_f    = var[1];
    YC_f    = var[2];
    PL_f    = var[3];
    D_f     = var[4];

    NC_m    = var[5];
    YC_m    = var[6];
    PL_m    = var[7];
    D_m     = var[8];

    # 4 equations for women
    eq1f = NC_f + YC_f + PL_f + D_f - 0.5;
    eq2f = fixedParam.chi*(NC_f+YC_f) - (fixedParam.etaf+fixedParam.gamma)*PL_f; # flows into and out of PL
    eq3f = fixedParam.etaf*PL_f - (fixedParam.chi+fixedParam.gamma)*YC_f; # flows into and out of YC
    eq4f = fixedParam.gamma*(NC_f+YC_f+PL_f) - fixedParam.phi*D_f;

    # 4 equations for men
    eq1m = NC_m + YC_m + PL_m + D_m - 0.5;
    eq2m = fixedParam.chi*(NC_m+YC_m) - (fixedParam.etam+fixedParam.gamma)*PL_m; # flows into and out of PL
    eq3m = fixedParam.etam*PL_m - (fixedParam.chi+fixedParam.gamma)*YC_m; # flows into and out of YC
    eq4m = fixedParam.gamma*(NC_m+YC_m+PL_m) - fixedParam.phi*D_m;

    return [eq1f, eq2f, eq3f, eq4f, eq1m, eq2m, eq3m, eq4m];
end
export solveStages
end


#
# This example assumes below packages are available in your environment
#

using StatsModelComparisons, StanSample
using StatsFuns, CSV, Random
using ParetoSmooth
#using StatisticalRethinking

ProjDir = @__DIR__

df = CSV.read(joinpath(ProjDir, "cars_logprob.csv"), DataFrame)

log_lik = Matrix(df)';
n_sam, n_obs = size(log_lik)
lppd = reshape(logsumexp(log_lik .- log(n_sam); dims=1), n_obs);

pwaic = [var(log_lik[:, i]) for i in 1:n_obs];
@show -2(sum(lppd) - sum(pwaic))
println()

@show waic(log_lik)
println()

#Random.seed!(123)
loo, loos, pk = psisloo(log_lik);
@show -2loo

#Random.seed!(123)
ll = reshape(Matrix(df), 50, 1000, 4);
psis_ll = psis(ll);

lwp = deepcopy(ll);
lwp += psis_ll.weights;
lwpt = Matrix(reshape(lwp, 50, 4000)');
loos = reshape(logsumexp(lwpt; dims=1), size(lwpt, 2));

@show loo = sum(loos)
@show 2loo

if isdefined(Main, :StatisticalRethinking)
    pk_plot(pk)
    savefig(joinpath(ProjDir, "pk.png"))

    pk_plot(psis_ll.pareto_k)
    savefig(joinpath(ProjDir, "pareto_k.png"))
end


"""
Convert a graph game to a simple game based on Myerson (1977). Namely, assume v(S) > 0 iff S is connected subgraph.
"""
function graphtoMyerson(G::GraphGame)::SimpleGame
    function vᴸ(S::Players)
        subgraph, vmap = induced_subgraph(G.g, S)
        comps = [
            [vmap[i] for i in comp]
            for comp in connected_components(subgraph)]

        return ∑(G.v(T) for T in comps)
    end

    return SimpleGame(G.N, vᴸ)
end


export
    HMM,
    LinearGaussian,
    NonLinearGaussian,
    GaussianHMM,
    generate

abstract type AbstractHMM end
abstract type GaussianHMM end

struct HMM <: AbstractHMM
    transmean::Function
    transloglik::Function # log transition function:  f(xk|xkm1)
    obsmean::Function
    obsloglik::Function   # log observation function: g(y|xk)
    dimx::Int
    dimy::Int
end

struct LinearGaussian <: GaussianHMM
    #=
        x <- Ax + chol(Q)' * randn
        y <- Bx + chol(R)' * randn
    =#
    A::Matrix{Float}
    B::Matrix{Float}
    Q::Union{Float,Matrix{Float}}
    R::Union{Float,Matrix{Float}}
    # implicit
    dimx::Int
    dimy::Int
    cholQ::Union{Float,UpperTriangular{Float, Matrix{Float}}}
    cholR::Union{Float,UpperTriangular{Float, Matrix{Float}}}
    transmean::Function
    obsmean::Function
    function LinearGaussian(A,B,Q,R)
        dimx=size(A,1)
        dimy=size(B,1)
        @assert dimx==size(A,2)==size(Q,1)==size(Q,2) "dimensions don't match"
        @assert dimy==size(R,1)==size(R,2) "dimensions don't match"
        @assert issymmetric(Q) && issymmetric(R) "cov mat must be symmetric"
        @assert isposdef(Q) && isposdef(R) "cov mat must be pos def"
        new(A,B,Q,R,dimx,dimy,chol(Q),chol(R),
            (k,xkm1) -> A*xkm1,
            (k,xk)   -> B*xk )
    end
end

struct NonLinearGaussian <: GaussianHMM
    #=
        x <- f(x) + chol(Q)' * randn
        y <- g(x) + chol(R)' * randn
    =#
    transmean::Function
    obsmean::Function
    Q::Union{Float,Matrix{Float}}
    R::Union{Float,Matrix{Float}}
    # implicit
    dimx::Int
    dimy::Int
    cholQ::Union{Float,UpperTriangular{Float, Matrix{Float}}}
    cholR::Union{Float,UpperTriangular{Float, Matrix{Float}}}
    function NonLinearGaussian(f, g, Q, R)
        dimx = size(Q, 1)
        dimy = size(R, 1)
        @assert issymmetric(Q) && issymmetric(R) "cov mat must be symmetric"
        @assert isposdef(Q) && isposdef(R) "cov mat must be pos def"
        new((k,xkm1)-> f(k,xkm1),
            (k,xk)  -> g(k,xk),
            Q,R,dimx,dimy,chol(Q),chol(R)  )
    end
end

function HMM(g::GaussianHMM)
    transloglik = (k,xkm1,xk) -> -norm(g.cholQ'\(xk - g.transmean(k,xkm1)))^2/2
    obsloglik   = (k,yk,  xk) -> -norm(g.cholR'\(yk - g.obsmean(k,xk)))^2/2
    HMM(g.transmean, transloglik, g.obsmean, obsloglik, g.dimx, g.dimy)
end

### Generation of observations

"""
    generate(lg, x0, K)

Generate observations following a given (non)linear Gaussian dynamic for `K`
time steps.
"""
function generate(g::GaussianHMM, x0::Vector{Float}, K::Int
                    )::Tuple{Matrix{Float},Matrix{Float}}
    @assert length(x0)==g.dimx "dimensions don't match"
    # allocate states/observations
    states, observations = zeros(g.dimx, K), zeros(g.dimy, K)
    # assign first state
    states[:,1] = x0
    # pre-generate noise
    noisex = g.cholQ' * randn(g.dimx,K)
    noisey = g.cholR' * randn(g.dimy,K)
    # use noise in iterative linear system
    for k = 1:(K-1)
        observations[:,k] = g.obsmean(k, states[:,k]) + noisey[:,k]
        states[:,k+1]     = g.transmean(k+1, states[:,k]) + noisex[:,k+1]
    end
    # last observation
    observations[:,K] = g.obsmean(K, states[:,K]) + noisey[:,K]
    # package and return
    return (states, observations)
end


include("euler/euler.jl")
using .Numbers: is_s_number

function compute(n::Int)::Int
    result = 0
    for i ∈ 2:n if is_s_number(i, i * i) result += i * i end end
    return result
end


using ArgCheck
using LinearAlgebra
using DocStringExtensions
using .Scattering
import QuadGK

"""
Contains parameters for running K-Matrix calculations.

Takes as input the mesh size, equivalent to the number of
iteration steps. The time required to solve KMatrix(N) goes as 𝒪(N²).
"""
struct KMatrix <: Method
    N::Int  # Number of iteration steps
end

"""
Compute the phase shift of S-waves by computing the reactance matrix.

$(SIGNATURES)
where `k₀` is the momentum for which to compute, `m` is the mass, and `P` the potential.

# Details

The goal is to obtain the S-wave phase shift of the potential P for the momentum `k`.
This is done by solving the matrix equation
             `K = A⁻¹V`
where K is the reactance matrix. V is the matrix of the potential evaluated at the mesh points,
and A describes the Lippmann-Schwinger equation for K. The phase shift is then related to
the last diagonal element of K as
             `K(k₀, k₀) = 1/mk₀ ⋅ tan(δₗ₌₀(k₀))`
"""
function (method::KMatrix)(k₀, m, P::Potential)
    A, V = createA(k₀, m, method.N, P)
    # Create the reactance matrix K using the matrix equation AK = V
    K = inv(A)*V
    @assert size(A) == size(V) == size(K) == (method.N+1, method.N+1)

    # Diagonal elements of K are related to the phase shifts δₗ
    # as K(k₀, k₀) = -tan(δₗ)/mk₀
    Kk0 = K[end, end]
    δᵢ = atan(-Kk0*m*k₀) / pi * 180
    #energy = k₀^2 / m * 197
    δᵢ
end

"""
Construct the A matrix involved in K-Matrix calculations.

$(SIGNATURES)

where `k₀` is the momentum for which to compute, `m` is the mass, `N` the
mesh size and `P` the potential.

# Details

The goal is the set up the matrices V and A of the matrix equation

    AK = V

In the first phase, the mesh points are constructed by Gaussian quadrature,
and V is constructed as the potential evaluated at P(kᵢ, kⱼ) ∀ i,j in the mesh.
The set up of A is more complicated, and is delegated to another function.
"""
function createA(k₀, m, N::Int, P::Potential)
    k, ω = QuadGK.gauss(N)
    transform!(k, ω)
    @assert k₀ ∉ k "k₀ can not be in the mesh points k"
    push!(k, k₀)
    V = [P(ki, kj) for ki in k, kj in k]
    A = createA(V, k, ω, m)
    return A, V
end

"""
Construct the A matrix involved in K-Matrix calculations.

$(SIGNATURES)

where `V` matrix of the potential evaluated at the mesh points, `k` is the momenta,
`ω` are the quadrature weights, and `m` is the mass.
# Details

The goal is the set up the matrices V and A of the matrix equation

    K = VA.

Here A is constructed as described in
https://manybodyphysics.github.io/NuclearForces/doc/pub/scatteringtheory/html/scatteringtheory.html
"""
function createA(V, k, ω, m)
    @argcheck length(k) - 1 == length(ω)
    @argcheck size(V, 1) == size(V, 2) == length(k)
    @argcheck m > 0

    N = length(k)-1
    k₀ = k[end]
    A = diagm(0 => repeat([1.0,], N+1))  # δᵢⱼ
    uⱼ = 0.0
    @inbounds for j in 1:N+1
        uⱼ = 0.0
        if j ≠ N+1
            uⱼ = 2/π * ω[j]*k[j]^2/((k₀^2 - k[j]^2)/m)
        else
            for n in 1:N
                uⱼ += ω[n]*k₀^2/((k₀^2 - k[n]^2)/m)
            end
            uⱼ *= -2/π
        end

        for i in 1:N+1
            A[i, j] -= V[i, j] * uⱼ
        end
    end
    return A
end

"""
Map the Gaussian quadrature points from [0, 1] to (-∞, ∞)

$(SIGNATURES)

"""
function transform!(xs, weights)
    # Times 197 for MeV
    @. weights = π/4 * weights / cos(π/4*(1.0 + xs))^2
    @. xs      = tan(π/4*(1.0 + xs))
end


# disease parameters
using Distributions

#R0 = 2.14
#beta = R0 / 3.56 - computed using N = 100
beta = R0 / 3.83 # computed using N = 500 - see matlab calibration scripts
p_asymp =  0.33 # fraction of asymptomatic infections
#p_asymp =  1.0 # fraction of asymptomatic infections
#latent period
global latent_period = dt

# incubation period
inc_mu = 1.62
inc_sig = 0.418
inc_dist = LogNormal(inc_mu, inc_sig)
# generate a random sample with rand(rng, inc_dist)

# post-incubation (recovery) period
rec_min = 5.0
rec_max = 10.0
rec_dist = Uniform(rec_min, rec_max)

# infectiousness (to capture secondary case dispersion)
b_dispersion = 0.1
b_scale = beta / b_dispersion
b_dist = Gamma(b_dispersion, b_scale)

Vmax = 7.0
inc_plat_fac = 0.1
rec_plat_fac = 0.0

# function that determines agent's contagiousness (beta)
function compute_kinc(t_inc::Float64)
    t_plat_inc = inc_plat_fac * t_inc
    k_inc = log(Vmax)/(t_inc - t_plat_inc)
    return k_inc
end

function compute_krec(t_rec::Float64)
    t_plat_rec = rec_plat_fac * t_rec
    k_rec = log(1.0/Vmax) / (t_rec - t_plat_rec)
    return k_rec
end

# test sensitivity
c_lims = [1.0, 5.11];
b1_lims = [0.8, 2.31]
b2_lims = [1.26, 3.47]
b3_lims = [1.05, 1.14]


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCellDG
include("../useful_routines.jl")
include("circular_bc_transformation_elasticity_solver.jl")

function compute_L2_stress_error(polyorder, nelmts, width)
    K1, K2 = 247.0, 192.0    # Pa
    mu1, mu2 = 126.0, 87.0   # Pa
    lambda1 = lame_lambda(K1, mu1)
    lambda2 = lame_lambda(K2, mu2)
    stiffness = CutCellDG.HookeStiffness(lambda1, mu1, lambda2, mu2)

    theta0 = -0.067
    stiffness = CutCellDG.HookeStiffness(lambda1, mu1, lambda2, mu2)
    transfstress =
        CutCellDG.plane_strain_transformation_stress(lambda1, mu1, theta0)

    meshwidth = [width, width]
    numqp = required_quadrature_order(polyorder) + 2
    penaltyfactor = 1e2

    dx = width / nelmts
    penalty = penaltyfactor / dx * 0.5 * (lambda1 + mu1 + lambda2 + mu2)

    elasticitybasis = LagrangeTensorProductBasis(2, polyorder)
    levelsetbasis = HermiteTensorProductBasis(2)
    quad = tensor_product_quadrature(2, 4)
    dim, nf = size(interpolation_points(levelsetbasis))
    refpoints = interpolation_points(elasticitybasis)

    interfacecenter = [0.5, 0.5]
    interfaceradius = minimum(meshwidth) / 3.0
    outerradius = 2.0
    analyticalsolution = AnalyticalSolution(
        interfaceradius,
        outerradius,
        interfacecenter,
        lambda1,
        mu1,
        lambda2,
        mu2,
        theta0,
    )

    mesh, cellquads, facequads, interfacequads = construct_mesh_and_quadratures(
        meshwidth,
        nelmts,
        elasticitybasis,
        interfacecenter,
        interfaceradius,
        numqp,
    )
    nodaldisplacement = nodal_displacement(
        mesh,
        elasticitybasis,
        cellquads,
        facequads,
        interfacequads,
        stiffness,
        theta0,
        analyticalsolution,
        penalty,
    )

    stresserror = stress_L2_error(
        nodaldisplacement,
        elasticitybasis,
        cellquads,
        stiffness,
        transfstress,
        theta0,
        mesh,
        x -> core_stress(analyticalsolution),
        x -> shell_stress(analyticalsolution, x),
    )

    den = integral_norm_on_mesh(
        x -> exact_stress(analyticalsolution, x),
        cellquads,
        mesh,
        4,
    )

    normalizedstresserr = stresserror ./ den

    return normalizedstresserr
end

function test_quadratic_elements()
    width = 1.0
    polyorder = 2
    powers = [3, 4, 5]
    nelmts = [2^p + 1 for p in powers]

    err = [compute_L2_stress_error(polyorder, ne, width) for ne in nelmts]

    stresserror = Array(transpose(hcat(err...)))
    dx = width ./ nelmts
    stressrates = mapslices(x -> convergence_rate(dx, x), stresserror, dims = 1)

    @test all(stressrates .> 1.8)
end

function test_cubic_elements()
    width = 1.0
    polyorder = 3
    powers = [3, 4, 5]
    nelmts = [2^p + 1 for p in powers]

    err = [compute_L2_stress_error(polyorder, ne, width) for ne in nelmts]

    dx = width ./ nelmts
    stresserror = Array(transpose(hcat(err...)))

    stressrates = mapslices(x -> convergence_rate(dx, x), stresserror, dims = 1)
    @test all(stressrates .> 2.7)
end

test_quadratic_elements()
test_cubic_elements()


export id, X, Y, Z, H, S, Sdg, T, Tdg, U1, U2, U3, CU, CNOT, CCX, SWAP

"""
    id(n=1)

Identity operator for n qubits system.
"""
function id(n::Int=1)
    return identityoperator(SpinBasis(1//2)^n)
end
const X = sigmax
const Y = sigmay
const Z = sigmaz
const H = hadamard
const S = sgate
const Sdg = sdggate
const T = tgate
const Tdg = tdggate
const U1 = phaseshift
"""
    U2(ϕ::Real, λ::Real)

Equivalent to `singleunitary(π/2, ϕ, λ)`.

```math
U(\\phi, \\lambda) = \\frac{1}{\\sqrt{2}} \\begin{bmatrix}
       1 & -e^{i \\lambda} \\\\
       e^{i \\phi} & e^{i (\\lambda + \\phi)}
\\end{bmatrix}
```
"""
U2(ϕ::Real, λ::Real) = singleunitary(π/2, ϕ, λ)
const U3 = singleunitary
const CU = controlsgate
const CNOT = cnot
const CCX = toffoli
const SWAP = swap


type NormalArm <: BanditArm
  mu::Float64
  sigma::Float64
end

function draw(arm::NormalArm)
  rand(Normal(mu, sigma))
end