# This file is a part of Julia. License is MIT: http://julialang.org/license

# Floating point numbers - basic tests
@test 4.00000000000001 ≈ 4.0
@test 5.0 ≈ 4.999999999999993
@test 4.000000002 ≉ 4.00300002

# Other tolerance levels
@test 4.32 ≈ 4.3 rtol=0.1 atol=0.01
@test 1.001 ≈ 1.002 rtol=0.001 atol=0.0001
@test 4.5 ≉ 4.9 rtol=0.001 atol=0.001

# Complex numbers
@test 1.0 + 1.0im ≈ 1.0 + 1.00000000000001im
@test 0.9999999999999 + 1.0im ≈ 1.0 + 1.000000000000001im
@test 0.9999 + 1.0im ≈ 1.0 + 1.1im rtol=0.0001 atol=1.1

# Complex <-> reals
@test 1.0 + 0im ≈ 1.0000000000001
@test 0.9999999999999 ≈ 1.0 + 0im
@test 1.0+1im ≉ 1.000000000000001

# Comparing NaNs
@test 4.0 ≉ NaN
@test NaN ≉ 4.0
@test complex(2.3,NaN) ≉ complex(NaN,2.3)
@test NaN ≉ NaN
@test complex(NaN,NaN) ≉ complex(NaN,NaN)
@test complex(NaN,2.3) ≉ complex(NaN,2.3)
@test complex(2.3,NaN) ≉ complex(2.3,NaN)

# Comparing Infs
@test Inf ≈ Inf
@test Inf ≉ 1
@test Inf ≉ -Inf
@test complex(0.0,Inf) ≈ complex(0.0,Inf)
@test complex(0.0,Inf) ≉ complex(0.0,-Inf)

# Tests for integers and rationals
@test 4 ≈ 4
@test 4 ≈ 5 atol=1
@test 4 ≉ 6 atol=1
@test 4 ≉ 5

@test 1//2+3im ≈ 1//2+3im
@test 1//3 ≈ 0.33333333333333333
@test 1//3 ≈ 0.3333 rtol=0.0001 atol=0.0001
@test 1+1im ≈ 1im+1

# Notably missing from this test suite at the moment
# * Tests for other magnitudes of numbers - very small, very big, and combinations of small and big
# * Tests for various odd combinations of types, e.g. x::Integer ≈ y::Rational

# issue #12375:
@test 1e17 ≉ 1

# Tests for arrays:
@test [1,2,3] ≈ [1,2,3+1e-9]
@test [0,1] ≈ [1e-9, 1]
@test [0,Inf] ≈ [0,Inf]
@test [0,Inf] ≉ [0,-Inf]

# issue #19936
for elty in (Float16,Float32,Float64)
    nan  = elty(NaN)
    half = elty(0.5)
    @test nan ≉ nan
    @test nan ≈ nan nans=true
    @test [half, nan, half] ≉ [half, nan, half]
    @test [half, nan, half] ≈ [half, nan, half] nans=true
end


"""
     uniform_partition(n::Int, block_size::Int)

Compute a uniform block partition of the given size.

### Input

- `n`          -- number of dimensions of the partition
- `block_size` -- size of each block

### Output

A vector of ranges, `Vector{UnitRange{Int}}`, such that the size of each block
is the same, if possible.

### Examples

If the number of dimensions `n` is 2, we have two options: either two blocks
of size `1` or one block of size `2`:

```jldoctest partition
julia> LazySets.Approximations.uniform_partition(2, 1)
2-element Array{UnitRange{Int64},1}:
 1:1
 2:2

julia> LazySets.Approximations.uniform_partition(2, 2)
1-element Array{UnitRange{Int64},1}:
 1:2
```

If the block size argument is not compatible with (i.e. does not divide) `n`, the
output is filled with one block of the size needed to reach `n`:

```jldoctest partition
julia> LazySets.Approximations.uniform_partition(3, 1)
3-element Array{UnitRange{Int64},1}:
 1:1
 2:2
 3:3

julia> LazySets.Approximations.uniform_partition(3, 2)
2-element Array{UnitRange{Int64},1}:
 1:2
 3:3

julia> LazySets.Approximations.uniform_partition(10, 6)
2-element Array{UnitRange{Int64},1}:
 1:6
 7:10
```
"""
@inline function uniform_partition(n::Int, block_size::Int)
    m = div(n, block_size)
    r = n % block_size
    res = Vector{UnitRange{Int}}(undef, r > 0 ? m + 1 : m)
    k = 1
    @inbounds for i in 1:m
        l = k + block_size - 1
        res[i] = k:l
        k = l + 1
    end
    if r > 0
        res[m+1] = k:n
    end
    return res
end

"""
    decompose(S::LazySet{N},
              partition::AbstractVector{<:AbstractVector{Int}},
              block_options
             ) where {N}

Decompose a high-dimensional set into a Cartesian product of overapproximations
of the projections over the specified subspaces.

### Input

- `S`             -- set
- `partition`     -- vector of blocks (i.e., of vectors of integers) (see the
                     Notes below)
- `block_options` -- mapping from block indices in `partition` to a
                     corresponding overapproximation option; we only require
                     access via `[⋅]` (but see also the Notes below)

### Output

A `CartesianProductArray` containing the low-dimensional approximated
projections.

### Algorithm

For each block a specific `project` method is called, dispatching on the
corresponding overapproximation option.

### Notes

The argument `partition` requires some discussion.
Typically, the list of blocks should form a partition of the set
``\\{1, \\dots, n\\}`` represented as a list of consecutive blocks, where ``n``
is the ambient dimension of set `S`.

However, technically there is no problem if the blocks are not consecutive,
blocks are missing, blocks occur more than once, or blocks are overlapping.
This function will, however, stick to the order of blocks, so the resulting set
must be interpreted with care in such cases.
One use case is the need of a projection consisting of several blocks.

For convenience, the argument `block_options` can also be given as a single
option instead of a mapping, which is then interpreted as the option for all
blocks.

### Examples

This function supports different options: one can specify the target set, the
degree of accuracy, and template directions.
These options are exemplified below, where we use the following example.

```jldoctest decompose_examples
julia> using LazySets.Approximations: decompose

julia> S = Ball2(zeros(4), 1.);  # set to be decomposed (4D 2-norm unit ball)

julia> P2d = [1:2, 3:4];  # a partition with two blocks of size two

julia> P1d = [[1], [2], [3], [4]];  # a partition with four blocks of size one
```

#### Different set types

We can decompose using polygons in constraint representation:

```jldoctest decompose_examples
julia> all([ai isa HPolygon for ai in array(decompose(S, P2d, HPolygon))])
true
```

For decomposition into 1D subspaces, we can use `Interval`:

```jldoctest decompose_examples
julia> all([ai isa Interval for ai in array(decompose(S, P1d, Interval))])
true
```

However, if you need to specify different set types for different blocks, the
interface presented so far does not apply.
See the paragraph *Advanced input for different block approximations* below for
how to do that.

#### Refining the decomposition I: ``ε``-close approximation

The ``ε`` option can be used to refine a decomposition, i.e., obtain a more
accurate result.
We use the [Iterative refinement](@ref) algorithm from the `Approximations`
module.

To illustrate this, consider again the set `S` from above.
We decompose into two 2D polygons.
Using smaller ``ε`` implies a better precision, thus more constraints in each 2D
decomposition.
In the following example, we look at the number of constraints in the first
block.

```jldoctest decompose_examples
julia> d(ε, bi) = array(decompose(S, P2d, (HPolygon => ε)))[bi]
d (generic function with 1 method)

julia> [length(constraints_list(d(ε, 1))) for ε in [Inf, 0.1, 0.01]]
3-element Array{Int64,1}:
  4
  8
 32
```

#### Refining the decomposition II: template polyhedra

Another way to refine a decomposition is by using template polyhedra.
The idea is to specify a set of template directions and then to compute on each
block the polytopic overapproximation obtained by evaluating the support
function of the given input set over the template directions.

For example, octagonal 2D approximations of the set `S` are obtained with:

```jldoctest decompose_examples
julia> using LazySets.Approximations: OctDirections

julia> B = decompose(S, P2d, OctDirections);

julia> length(B.array) == 2 && all(dim(bi) == 2 for bi in B.array)
true
```

See [Template directions](@ref) for the available template directions.
Note that, in contrast to the polygonal ``ε``-close approximation from above,
this method can be applied to blocks of any size.

```jldoctest decompose_examples
julia> B = decompose(S, [1:4], OctDirections);

julia> length(B.array) == 1 && dim(B.array[1]) == 4
true
```

#### Advanced input for different block approximations

Instead of defining the approximation option uniformly for each block, we can
define different approximations for different blocks.
The third argument has to be a mapping from block index (in the partition) to
the corresponding approximation option.

For example:

```jldoctest decompose_examples
julia> res = array(decompose(S, P2d, Dict(1 => Hyperrectangle, 2 => 0.1)));

julia> typeof(res[1]), typeof(res[2])
(Hyperrectangle{Float64,Array{Float64,1},Array{Float64,1}}, HPolygon{Float64,Array{Float64,1}})
```
"""
function decompose(S::LazySet{N},
                   partition::AbstractVector{<:AbstractVector{Int}},
                   block_options
                  ) where {N}
    n = dim(S)
    result = Vector{LazySet{N}}(undef, length(partition))

    @inbounds for (i, block) in enumerate(partition)
        result[i] = project(S, block, block_options[i], n)
    end
    return CartesianProductArray(result)
end

# convenience method with uniform block options
function decompose(S::LazySet{N},
                   partition::AbstractVector{<:AbstractVector{Int}},
                   block_options::Union{Type{<:LazySet},
                                        Pair{<:UnionAll, <:Real},
                                        Real,
                                        Type{<:AbstractDirections},
                                        Nothing
                                       }
                  ) where {N}
    n = dim(S)
    result = Vector{LazySet{N}}(undef, length(partition))

    @inbounds for (i, block) in enumerate(partition)
        result[i] = project(S, block, block_options, n)
    end
    return CartesianProductArray(result)
end

# convenience method with uniform block size
function decompose(S::LazySet, block_options; block_size::Int=1)
    partition = uniform_partition(dim(S), block_size)
    return decompose(S, partition, block_options)
end

"""
    project(S::LazySet{N},
            block::AbstractVector{Int},
            [::Nothing=nothing],
            [n]::Int=dim(S)
           ) where {N}

Project a high-dimensional set to a given block by using a concrete linear map.

### Input

- `S`       -- set
- `block`   -- block structure - a vector with the dimensions of interest
- `nothing` -- (default: `nothing`) used for dispatch
- `n`       -- (optional, default: `dim(S)`) ambient dimension of the set `S`

### Output

A set representing the projection of the set `S` to block `block`.

### Algorithm

We apply the function `linear_map`.
"""
@inline function project(S::LazySet{N},
                         block::AbstractVector{Int},
                         ::Nothing=nothing,
                         n::Int=dim(S)
                        ) where {N}
    M = projection_matrix(block, n, N)
    return linear_map(M, S)
end

"""
    project(S::LazySet{N},
            block::AbstractVector{Int},
            set_type::Type{<:LinearMap},
            [n]::Int=dim(S)
           ) where {N}

Project a high-dimensional set to a given block by using a lazy linear map.

### Input

- `S`         -- set
- `block`     -- block structure - a vector with the dimensions of interest
- `LinearMap` -- used for dispatch
- `n`         -- (optional, default: `dim(S)`) ambient dimension of the set `S`

### Output

A lazy `LinearMap` representing the projection of the set `S` to block `block`.
"""
@inline function project(S::LazySet{N},
                         block::AbstractVector{Int},
                         set_type::Type{<:LinearMap},
                         n::Int=dim(S)
                        ) where {N}
    M = projection_matrix(block, n, N)
    return M * S
end

"""
    project(S::LazySet,
            block::AbstractVector{Int},
            set_type::Type{<:LazySet},
            [n]::Int=dim(S)
           )

Project a high-dimensional set to a given block and set type, possibly involving
an overapproximation.

### Input

- `S`        -- set
- `block`    -- block structure - a vector with the dimensions of interest
- `set_type` -- target set type
- `n`        -- (optional, default: `dim(S)`) ambient dimension of the set `S`

### Output

A set of type `set_type` representing an overapproximation of the projection of
`S`.

### Algorithm

1. Project the set `S` with `M⋅S`, where `M` is the identity matrix in the block
coordinates and zero otherwise.
2. Overapproximate the projected lazy set using `overapproximate` and
`set_type`.
"""
@inline function project(S::LazySet,
                         block::AbstractVector{Int},
                         set_type::Type{<:LazySet},
                         n::Int=dim(S)
                        )
    lm = project(S, block, LinearMap, n)
    return overapproximate(lm, set_type)
end

"""
    project(S::LazySet,
            block::AbstractVector{Int},
            set_type_and_precision::Pair{<:UnionAll, <:Real},
            [n]::Int=dim(S)
           )

Project a high-dimensional set to a given block and set type with a certified
error bound.

### Input

- `S`     -- set
- `block` -- block structure - a vector with the dimensions of interest
- `set_type_and_precision` -- pair `(T, ε)` of a target set type `T` and an
                              error bound `ε` for approximation
- `n`     -- (optional, default: `dim(S)`) ambient dimension of the set `S`

### Output

A set representing the epsilon-close approximation of the projection of `S`.

### Notes

Currently we only support `HPolygon` as set type, which implies that the set
must be two-dimensional.

### Algorithm

1. Project the set `S` with `M⋅S`, where `M` is the identity matrix in the block
coordinates and zero otherwise.
2. Overapproximate the projected lazy set with the given error bound `ε`.
"""
@inline function project(S::LazySet,
                         block::AbstractVector{Int},
                         set_type_and_precision::Pair{<:UnionAll, <:Real},
                         n::Int=dim(S)
                        )
    set_type = set_type_and_precision[1]
    ε = set_type_and_precision[2]
    @assert length(block) == 2 && set_type == HPolygon "currently only 2D " *
        "HPolygon decomposition is supported"

    lm = project(S, block, LinearMap, n)
    return overapproximate(lm, set_type, ε)
end

"""
    project(S::LazySet,
            block::AbstractVector{Int},
            ε::Real,
            [n]::Int=dim(S)
           )

Project a high-dimensional set to a given block and set type with a certified
error bound.

### Input

- `S`     -- set
- `block` -- block structure - a vector with the dimensions of interest
- `ε`     -- error bound for approximation
- `n`     -- (optional, default: `dim(S)`) ambient dimension of the set `S`

### Output

A set representing the epsilon-close approximation of the projection of `S`.

### Algorithm

1. Project the set `S` with `M⋅S`, where `M` is the identity matrix in the block
coordinates and zero otherwise.
2. Overapproximate the projected lazy set with the given error bound `ε`.
The target set type is chosen automatically.
"""
@inline function project(S::LazySet,
                         block::AbstractVector{Int},
                         ε::Real,
                         n::Int=dim(S)
                        )
    # currently we only support HPolygon
    if length(block) == 2
        set_type = HPolygon
    else
        throw(ArgumentError("ε-close approximation is only supported for 2D " *
                            "blocks"))
    end
    return project(S, block, set_type => ε, n)
end

"""
    project(S::LazySet,
            block::AbstractVector{Int},
            directions::Type{<:AbstractDirections},
            [n]::Int
           )

Project a high-dimensional set to a given block using template directions.

### Input

- `S`          -- set
- `block`      -- block structure - a vector with the dimensions of interest
- `directions` -- template directions
- `n`          -- (optional, default: `dim(S)`) ambient dimension of the set `S`

### Output

The template direction approximation of the projection of `S`.
"""
@inline function project(S::LazySet,
                         block::AbstractVector{Int},
                         directions::Type{<:AbstractDirections},
                         n::Int=dim(S)
                        )
    lm = project(S, block, LinearMap, n)
    return overapproximate(lm, directions(length(block)))
end

"""
    project(H::HalfSpace, block::AbstractVector{Int})

Concrete projection of a half-space.

### Input

- `H`        -- set
- `block`    -- block structure, a vector with the dimensions of interest

### Output

A set representing the projection of the half-space `H` on the dimensions
specified by `block`.

### Algorithm

If the unconstrained dimensions of `H` are a subset of the `block` variables,
the projection is applied to the normal direction of `H`.
Otherwise, the projection results in the universal set.

The latter can be seen as follows.
Without loss of generality consider a projection onto a single and constrained
dimension ``xₖ`` (projections in multiple dimensions can be modeled as repeated
one-dimensional projections).
We can write the projection as an existentially quantified linear constraint:

```math
    ∃xₖ: a₁x₁ + … + aₖxₖ + … + aₙxₙ ≤ b
```

Since ``aₖ ≠ 0``, there is always a value for ``xₖ`` that satisfies the
constraint for any valuation of the other variables.

### Examples

Consider the half-space ``x + y + 0⋅z ≤ 1``, whose ambient dimension is `3`.
The (trivial) projection in the three dimensions is achieved letting the block
of variables to be `[1, 2, 3]`:

```jldoctest project_halfspace
julia> H = HalfSpace([1.0, 1.0, 0.0], 1.0)
HalfSpace{Float64,Array{Float64,1}}([1.0, 1.0, 0.0], 1.0)

julia> project(H, [1, 2, 3])
HalfSpace{Float64,Array{Float64,1}}([1.0, 1.0, 0.0], 1.0)
```

Projecting along dimensions `1` and `2` only:

```jldoctest project_halfspace
julia> project(H, [1, 2])
HalfSpace{Float64,Array{Float64,1}}([1.0, 1.0], 1.0)
```

In general, use the call syntax `project(H, constrained_dimensions(H))` to return
the half-space projected on the dimensions where it is constrained only:

```jldoctest project_halfspace
julia> project(H, constrained_dimensions(H))
HalfSpace{Float64,Array{Float64,1}}([1.0, 1.0], 1.0)
```

If a constrained dimension is projected, we get the universal set of the
dimension corresponding to the projection.

```jldoctest project_halfspace
julia> project(H, [1, 3])
Universe{Float64}(2)

julia> project(H, [1])
Universe{Float64}(1)
```
"""
function project(H::HalfSpace, block::AbstractVector{Int})
    if constrained_dimensions(H) ⊆ block
        return HalfSpace(H.a[block], H.b)
    else
        return Universe(length(block))
    end
end

"""
    project(P::AbstractPolyhedron{N}, block::AbstractVector{Int}) where {N}

Concrete projection of a polyhedral set.

### Input

- `P`        -- set
- `block`    -- block structure, a vector with the dimensions of interest

### Output

A polyhedron representing the projection of `P` on the dimensions specified by
`block`.
If `P` was bounded, the result is an `HPolytope`; otherwise the result is an
`HPolyhedron`.

### Algorithm

- If the unconstrained dimensions of `P` are a subset of the `block` variables,
  each half-sace `c` of `P` is transformed to `HalfSpace(c.a[block], c.b)`.
- In the general case, we compute the concrete linear map of the projection
  matrix associated to the given block structure.

### Examples

Consider the four-dimensional cross-polytope (unit ball in the 1-norm):

```jldoctest project_polyhedron
julia> P = Ball1(zeros(4), 1.0);
```

All dimensions are constrained, and computing the (trivial) projection on the whole
space behaves as expected:

```jldoctest project_polyhedron
julia> constrained_dimensions(P)
4-element Array{Int64,1}:
 1
 2
 3
 4

julia> P_1234 = project(P, [1, 2, 3, 4]);

julia> P_1234 == convert(HPolytope, P)
true
```
Each constraint of the cross polytope is constrained in all dimensions.

Now let's take a ball in the infinity norm and remove some constraints:

```jldoctest project_polyhedron
julia> B = BallInf(zeros(4), 1.0);

julia> c = constraints_list(B)[1:2]
2-element Array{HalfSpace{Float64,LazySets.Arrays.SingleEntryVector{Float64}},1}:
 HalfSpace{Float64,LazySets.Arrays.SingleEntryVector{Float64}}([1.0, 0.0, 0.0, 0.0], 1.0)
 HalfSpace{Float64,LazySets.Arrays.SingleEntryVector{Float64}}([0.0, 1.0, 0.0, 0.0], 1.0)

julia> P = HPolyhedron(c);

julia> constrained_dimensions(P)
2-element Array{Int64,1}:
 1
 2
```

Finally we take the concrete projection onto variables `1` and `2`:

```jldoctest project_polyhedron
julia> project(P, [1, 2]) |> constraints_list
2-element Array{HalfSpace{Float64,Array{Float64,1}},1}:
 HalfSpace{Float64,Array{Float64,1}}([1.0, 0.0], 1.0)
 HalfSpace{Float64,Array{Float64,1}}([0.0, 1.0], 1.0)
```
"""
function project(P::AbstractPolyhedron{N}, block::AbstractVector{Int}) where {N}
    if constrained_dimensions(P) ⊆ block
        clist = [HalfSpace(c.a[block], c.b) for c in constraints_list(P)]
    else
        n = dim(P)
        M = projection_matrix(block, n, N)
        lm = linear_map(M, P)
        clist = constraints_list(lm)
    end
    T = isbounded(P) ? HPolytope : HPolyhedron
    return T(clist)
end

function project(Z::Zonotope{N}, block::AbstractVector{Int}) where {N}
    n = dim(Z)
    M = projection_matrix(block, n, N)
    lm = remove_zero_generators(linear_map(M, Z))
    return lm
end

function project(H::AbstractHyperrectangle, block::AbstractVector{Int})
    πc = center(H)[block]
    πr = radius_hyperrectangle(H)[block]
    return Hyperrectangle(πc, πr, check_bounds=false)
end

function project(V::Union{<:VPolygon{N}, <:VPolytope{N}},
                 block::AbstractVector{Int}) where {N}
    n = dim(V)
    M = projection_matrix(block, n, N)
    πvertices = broadcast(v -> M * v, vertices_list(V))

    m = size(M, 1)
    if m == 1
        # convex_hull in 1d returns the minimum and maximum points, in that order
        aux = convex_hull(πvertices)
        a = first(aux[1])
        b = length(aux) == 1 ? a : first(aux[2])
        return Interval(a, b)
    elseif m == 2
        return VPolygon(πvertices; apply_convex_hull=true)
    else
        return VPolytope(πvertices)
    end
end


function corpFlightBookings(bookings::Array{Array{Int64,1},1}, n::Int64)
    result::Array{Int64,1} = zeros(Int64,20000)
    for book in bookings
        for ind in book[1]:book[2]
            result[ind] += book[3]
        end
    end
    result[1:n]
end


module LinearSVM

include("DualCD.jl")
include("Pegasos.jl")
include("ADMMDriver.jl")
#include("ADMMWorker.jl")
include("Utilities.jl")

# DualCD.jl:
export dualcd
export pegasos

# ADMMDriver.jl:
export distributeData
export admm

# Utilities.jl:
export loadDataset
export testResult

end # module


using CompositeGrids, Test, StaticArrays, LinearAlgebra, Printf, Random, Statistics, JLD2, FileIO
# import Test: @test, @testset

if isempty(ARGS)
    include("chebyshev.jl")
    include("grid.jl")
    include("interpolate.jl")
    include("io.jl")
else
    include(ARGS[1])
end


import Base.copy

#*********************************#
#       COST FUNCTION CLASS       #
#*********************************#

abstract type CostFunction end
CostTrajectory = Vector{C} where C <: CostFunction

function cost(c::CostTrajectory, X::VectorTrajectory{T}, U::VectorTrajectory{T},dt::Vector{T})::T where T
    N = length(X)
    J = 0.0
    for k = 1:N-1
        J += stage_cost(c[k],X[k],U[k],dt[k])
    end
    J += stage_cost(c[N],X[N])
    return J
end

"$(TYPEDEF) Expansion of cost function"
struct Expansion{T<:AbstractFloat,Q<:AbstractMatrix,R<:AbstractMatrix}
    x::Vector{T}
    u::Vector{T}
    xx::Q
    uu::R
    ux::Matrix{T}
end

function Expansion{T,Q,R}(n::Int, m::Int) where {T,Q,R}
    x = zeros(T,n)
    u = zeros(T,m)
    xx = Q(zeros(T,n,n))
    uu = R(zeros(T,m,m))
    ux = zeros(T,m,n)
    Expansion(x, u, xx, uu, ux)
end
Expansion{T}(n::Int, m::Int) where T = Expansion{T,Matrix{T},Matrix{T}}(n,m)

import Base./, Base.*
function *(e::Expansion, a::Real)
    e.x .*= a
    e.u .*= a
    e.xx .*= a
    e.uu .*= a
    e.ux .*= a
    return nothing
end
*(a::Real, e::Expansion) = e*a

function /(e::Expansion,a::Real)
    e.x ./= a
    e.u ./= a
    e.xx ./= a
    e.uu ./= a
    e.ux ./= a
    return nothing
end

function copy(e::Expansion)
    Expansion(copy(e.x),copy(e.u),copy(e.xx),copy(e.uu),copy(e.ux))
end

function reset!(e::Expansion)
    !isempty(e.x) ? e.x .= zero(e.x) : nothing
    !isempty(e.u) ? e.u .= zero(e.u) : nothing
    !isempty(e.xx) ? e.xx .= zero(e.xx) : nothing
    !isempty(e.uu) ? e.uu .= zero(e.uu) : nothing
    !isempty(e.ux) ? e.ux .= zero(e.ux) : nothing
    return nothing
end

ExpansionTrajectory{T} = Vector{Expansion{T}} where T <: AbstractFloat

function reset!(et::ExpansionTrajectory)
    for e in et
        reset!(e)
    end
end

#######################################################
#              COST FUNCTION INTERFACE                #
#######################################################

#TODO this documentation breaks functionality by create ambiguous methods...

# "$(SIGNATURES) Evaluate the cost at state `x` and control `u`"
# stage_cost(cost::CostFunction, x, u) = 0.0
# "$(SIGNATURES) Evaluate the cost at the terminal state `xN`"
# stage_cost(cost::CostFunction, xN) = 0.0
# "$(SIGNATURES) Evaluate the second order expansion at state `x` and control `u`"
# cost_expansion!(Q::Expansion, cost::CostFunction, x, u) = nothing
# "$(SIGNATURES) Evaluate the second order expansion at the terminal state `xN`"
# cost_expansion!(Q::Expansion, cost::CostFunction, xN) = nothing

"""
$(TYPEDEF)
Cost function of the form
    1/2xₙᵀ Qf xₙ + qfᵀxₙ +  ∫ ( 1/2xᵀQx + 1/2uᵀRu + xᵀHu + qᵀx  rᵀu ) dt from 0 to tf
R must be positive definite, Q and Qf must be positive semidefinite

Constructor use any of the following constructors:
```julia
QuadraticCost(Q, R, H, q, r, c)
QuadraticCost(Q, R; H, q, r, c)
QuadraticCost(Q, q, c)
```
Any optional or omitted values will be set to zero(s).
"""
mutable struct QuadraticCost{T} <: CostFunction
    Q::AbstractMatrix{T}                 # Quadratic stage cost for states (n,n)
    R::AbstractMatrix{T}                 # Quadratic stage cost for controls (m,m)
    H::AbstractMatrix{T}                 # Quadratic Cross-coupling for state and controls (n,m)
    q::AbstractVector{T}                 # Linear term on states (n,)
    r::AbstractVector{T}                 # Linear term on controls (m,)
    c::T                                 # constant term
    function QuadraticCost(Q::AbstractMatrix{T}, R::AbstractMatrix{T}, H::AbstractMatrix{T},
            q::AbstractVector{T}, r::AbstractVector{T}, c::T) where T
        if !isposdef(R)
            @warn "R is not positive definite"
        end
        if !ispossemidef(Q)
            err = ArgumentError("Q must be positive semi-definite")
            throw(err)
        end

        new{T}(Q,R,H,q,r,c)
    end
end

get_sizes(cost::QuadraticCost) = (size(cost.Q,1),size(cost.R,1))

function QuadraticCost(Q,R; H=zeros(size(R,1), size(Q,1)), q=zeros(size(Q,1)),
        r=zeros(size(R,1)), c=0.0)
    QuadraticCost(Q,R,H,q,r,c)
end

function QuadraticCost(Q,q,c)
    QuadraticCost(Q,zeros(0,0),zeros(0,size(Q,1)),q,zeros(0),c)
end

"""
$(SIGNATURES)
Cost function of the form
``(x-x_f)^T Q (x_x_f) + u^T R u``
R must be positive definite, Q must be positive semidefinite
"""
function LQRCost(Q::AbstractArray, R::AbstractArray, xf::AbstractVector)
    H = zeros(size(R,1),size(Q,1))
    q = -Q*xf
    r = zeros(size(R,1))
    c = 0.5*xf'*Q*xf
    return QuadraticCost(Q, R, H, q, r, c)
end

"""
$(SIGNATURES)
Cost function of the form
``(x-x_f)^T Q (x_x_f)``
Q must be positive semidefinite
"""
function LQRCostTerminal(Qf::AbstractArray,xf::AbstractVector)
    qf = -Qf*xf
    cf = 0.5*xf'*Qf*xf
    return QuadraticCost(Qf,zeros(0,0),zeros(0,size(Qf,1)),qf,zeros(0),cf)
end

function stage_cost(cost::QuadraticCost, x::AbstractVector, u::AbstractVector)
    0.5*x'cost.Q*x + 0.5*u'*cost.R*u + cost.q'x + cost.r'u + cost.c + u'*cost.H*x
end

function stage_cost(cost::QuadraticCost, x::AbstractVector, u::AbstractVector, dt)
    (0.5*x'cost.Q*x + 0.5*u'*cost.R*u + cost.q'x + cost.r'u + cost.c + u'*cost.H*x)*dt
end

function stage_cost(cost::QuadraticCost, xN::AbstractVector{T}) where T
    0.5*xN'cost.Q*xN + cost.q'*xN + cost.c
end

function cost_expansion!(Q::Expansion{T}, cost::QuadraticCost,
        x::AbstractVector{T}, u::AbstractVector{T}, dt::T) where T
    Q.x .= cost.Q*x + cost.q + cost.H'*u
    Q.u .= cost.R*u + cost.r + cost.H*x
    Q.xx .= cost.Q
    Q.uu .= cost.R
    Q.ux .= cost.H
    Q*dt
    return nothing
end

function cost_expansion!(S::Expansion{T}, cost::QuadraticCost, xN::AbstractVector{T}) where T
    S.xx .= cost.Q
    S.x .= cost.Q*xN + cost.q
    return nothing
end

function gradient!(grad, cost::QuadraticCost,
        x::AbstractVector, u::AbstractVector,dt)
    grad.x .= (cost.Q*x + cost.q + cost.H'*u)
    grad.u .= (cost.R*u + cost.r + cost.H*x)

    grad .*= dt
    return nothing
end

function gradient!(grad, cost::QuadraticCost, xN::AbstractVector)
    grad .= cost.Q*xN + cost.q
    return nothing
end

function hessian!(hess, cost::QuadraticCost,
        x::AbstractVector, u::AbstractVector, dt)
    hess.xx .= cost.Q
    hess.uu .= cost.R
    hess.ux .= cost.H
    hess.xu .= cost.H'

    hess .*= dt
    return nothing
end

function hessian!(hess, cost::QuadraticCost, xN::AbstractVector)
    hess .= cost.Q
    return nothing
end

function copy(cost::QuadraticCost)
    return QuadraticCost(copy(cost.Q), copy(cost.R), copy(cost.H), copy(cost.q), copy(cost.r), copy(cost.c))
end

"""
$(TYPEDEF)
Cost function of the form
    ℓf(xₙ) + ∫ ℓ(x,u) dt from 0 to tf
"""
struct GenericCost <: CostFunction
    ℓ::Function             # Stage cost
    ℓf::Function            # Terminal cost
    expansion::Function     # 2nd order Taylor Series Expansion of the form,  Q,R,H,q,r = expansion(x,u)
    n::Int                  #                                                     Qf,qf = expansion(xN)
    m::Int
end

"""
$(SIGNATURES)
Create a Generic Cost, specifying the gradient and hessian of the cost function analytically

# Arguments
* hess: multiple-dispatch function of the form,
    Q,R,H = hess(x,u) with sizes (n,n), (m,m), (m,n)
    Qf = hess(xN) with size (n,n)
* grad: multiple-dispatch function of the form,
    q,r = grad(x,u) with sizes (n,), (m,)
    qf = grad(x,u) with size (n,)

"""
function GenericCost(ℓ::Function, ℓf::Function, grad::Function, hess::Function, n::Int, m::Int)
    function expansion(x::Vector{T},u::Vector{T}) where T
        Q,R,H = hess(x,u)
        q,r = grad(x,u)
        return Q,R,H,q,r
    end
    expansion(xN) = hess(xN), grad(xN)
    GenericCost(ℓ,ℓf, expansion, n,m)
end

"""
$(SIGNATURES)
Create a Generic Cost. Gradient and Hessian information will be determined using ForwardDiff

# Arguments
* ℓ: stage cost function of the form J = ℓ(x,u)
* ℓf: terminal cost function of the form J = ℓ(xN)
"""
function GenericCost(ℓ::Function, ℓf::Function, n::Int, m::Int)
    linds = LinearIndices(zeros(n+m,n+m))
    xinds = 1:n
    uinds = n .+(1:m)
    inds = (x=xinds, u=uinds, xx=linds[xinds,xinds], uu=linds[uinds,uinds], ux=linds[uinds,xinds])
    expansion = auto_expansion_function(ℓ,ℓf,n,m)

    GenericCost(ℓ,ℓf, expansion, n,m)
end

function auto_expansion_function(ℓ::Function,ℓf::Function,n::Int,m::Int)
    z = zeros(n+m)
    hess = zeros(n+m,n+m)
    grad = zeros(n+m)
    qf,Qf = zeros(n), zeros(n,n)

    linds = LinearIndices(hess)
    xinds = 1:n
    uinds = n .+(1:m)
    inds = (x=xinds, u=uinds, xx=linds[xinds,xinds], uu=linds[uinds,uinds], ux=linds[uinds,xinds])
    function ℓ_aug(z::Vector{T}) where T
        x = view(z,xinds)
        u = view(z,uinds)
        ℓ(x,u)
    end
    function expansion(x::Vector{T},u::Vector{T}) where T
        z[inds.x] = x
        z[inds.u] = u
        ForwardDiff.hessian!(hess, ℓ_aug, z)
        Q = view(hess,inds.xx)
        R = view(hess,inds.uu)
        H = view(hess,inds.ux)

        ForwardDiff.gradient!(grad, ℓ_aug, z)
        q = view(grad,inds.x)
        r = view(grad,inds.u)
        return Q,R,H,q,r
    end
    function expansion(xN::Vector{T}) where T
        ForwardDiff.gradient!(qf,ℓf,xN)
        ForwardDiff.hessian!(Qf,ℓf,xN)
        return Qf, qf
    end
end

stage_cost(cost::GenericCost, x::Vector{T}, u::Vector{T}) where T = cost.ℓ(x,u)
stage_cost(cost::GenericCost, xN::Vector{T}) where T = cost.ℓf(xN)

function cost_expansion!(Q::Expansion{T}, cost::GenericCost, x::Vector{T},
        u::Vector{T}) where T

    e = cost.expansion(x,u)

    Q.x .= e[4]
    Q.u .= e[5]
    Q.xx .= e[1]
    Q.uu .= e[2]
    Q.ux .= e[3]
    return nothing
end

function cost_expansion!(S::Expansion{T}, cost::GenericCost, xN::Vector{T}) where T
    Qf, qf = cost.expansion(xN)
    S.xx .= Qf
    S.x .= qf
    return nothing
end

copy(cost::GenericCost) = GenericCost(cost.ℓ,cost.ℓ,cost.n,cost.m)


using ADSeismic
using ADCME
using PyPlot
using PyCall
using DelimitedFiles
using Statistics
use_gpu()

function mkpath_by_force(path)
  if !isdir(path)
    mkpath(path)
  end
end

output_dir = "data/acoustic"
figure_dir = "figure/acoustic/layer-model/"
result_dir = "result/acoustic/layer-model/"

mkpath_by_force(output_dir)
mkpath_by_force(figure_dir)
mkpath_by_force(result_dir)

################### Generate synthetic data #####################
reset_default_graph()

ap_sim = load_acoustic_model("models/layer-model-anomaly.mat"; IT_DISPLAY=0)
src = load_acoustic_source("models/layer-model-anomaly.mat")
rcv_sim = load_acoustic_receiver("models/layer-model-anomaly.mat")

Rs_ = compute_forward_GPU(ap_sim, src, rcv_sim)

sess = Session(); init(sess)
Rs = run(sess, Rs_)

for i = 1:length(src)
    writedlm(joinpath(output_dir, "layermodel-r$i.txt"), Rs[i])
end

u = run(sess, ap_sim(src[div(length(src),2)+1]).u)
visualize_file(u, ap_sim.param, dirs=figure_dir)

################### Inversion using Automatic Differentiation #####################
reset_default_graph()

ap_sim = load_acoustic_model("models/layer-model-smooth.mat"; inv_vp=true, IT_DISPLAY=0)
src = load_acoustic_source("models/layer-model-smooth.mat")
rcv_sim = load_acoustic_receiver("models/layer-model-smooth.mat")

Rs = Array{Array{Float64,2}}(undef, length(src))
for i = 1:length(src)
    Rs[i] = readdlm(joinpath(output_dir, "layermodel-r$i.txt"))
end

losses, gs = compute_loss_and_grads_GPU(ap_sim, src, rcv_sim, Rs, get_collection()[1])
g = sum(gs)
loss = sum(losses)
sess = Session(); init(sess)

vp = (run(sess, ap_sim.vp.contents)).^0.5
scale = mean(vp)
vmax = maximum(vp)
vmin = minimum(vp)

function callback(x, iter, loss)
  if iter%10==0
    X = Array(reshape(x, ap_sim.param.NY+2, ap_sim.param.NX+2))
    close("all")
    pcolormesh([0:ap_sim.param.NX+1;]*ap_sim.param.DELTAX,[0:ap_sim.param.NY+1;]*ap_sim.param.DELTAY,  X*scale, vmax=vmax, vmin=vmin)
    axis("scaled")
    colorbar(shrink=0.4)
    xlabel("x (km)")
    ylabel("z (km)")
    gca().invert_yaxis()
    title("Iteration = $iter")
    savefig(joinpath(figure_dir, "inv_$(lpad(iter,5,"0")).png"))
    writedlm(joinpath(result_dir, "inv_$(lpad(iter,5,"0")).txt"), X*scale)
    open(joinpath(result_dir, "loss.txt"), "a") do io
      writedlm(io, loss)
    end
  else
    return
  end
end

# callback(run(sess, get_collection()[1])', 0, 0)
LBFGS!(sess, loss, g, get_collection()[1]; callback=callback)


using ImageFiltering, Base.Test

@testset "extrema_filter" begin
    function groundtruth(f, A, window::Tuple)
        Aex = copy(A)
        hshift = map(x->x>>1+1, window)
        for Ishift in CartesianRange(window)
            for I in CartesianRange(size(A))
                Aex[I] = f(Aex[I], A[map(d->clamp(I[d]+Ishift[d]-hshift[d], 1, size(A,d)), 1:ndims(A))...])
            end
        end
        Aex
    end
    groundtruth(f, A, window) = groundtruth(f, A, (window...,))
    # 1d case
    A = [0.1,0.3,0.2,0.3,0.4]
    mm = extrema_filter(A, 1)
    @test mm == [(a,a) for a in A]
    mm = extrema_filter(A, 2)
    @test mm == [(0.1,0.1),(0.1,0.3),(0.2,0.3),(0.2,0.3),(0.3,0.4)]
    mm = extrema_filter(A, 3)
    @test mm == [(0.1,0.3),(0.1,0.3),(0.2,0.3),(0.2,0.4),(0.3,0.4)]
    mm = extrema_filter(A, 4)
    @test mm == [(0.1,0.3),(0.1,0.3),(0.1,0.3),(0.2,0.4),(0.2,0.4)]
    mm = extrema_filter(A, 5)
    @test mm == [(0.1,0.3),(0.1,0.3),(0.1,0.4),(0.2,0.4),(0.2,0.4)]
    mm = extrema_filter(A, 6)
    @test mm == [(0.1,0.3),(0.1,0.3),(0.1,0.4),(0.1,0.4),(0.2,0.4)]
    mm = extrema_filter(A, 7)
    @test mm == [(0.1,0.3),(0.1,0.4),(0.1,0.4),(0.1,0.4),(0.2,0.4)]
    A = [0.1,0.3,0.5,0.4,0.2]
    mm = extrema_filter(A, 1)
    @test mm == [(a,a) for a in A]
    mm = extrema_filter(A, 2)
    @test mm == [(0.1,0.1),(0.1,0.3),(0.3,0.5),(0.4,0.5),(0.2,0.4)]
    mm = extrema_filter(A, 3)
    @test mm == [(0.1,0.3),(0.1,0.5),(0.3,0.5),(0.2,0.5),(0.2,0.4)]
    mm = extrema_filter(A, 4)
    @test mm == [(0.1,0.3),(0.1,0.5),(0.1,0.5),(0.2,0.5),(0.2,0.5)]
    mm = extrema_filter(A, 5)
    @test mm == [(0.1,0.5),(0.1,0.5),(0.1,0.5),(0.2,0.5),(0.2,0.5)]
    mm = extrema_filter(A, 6)
    @test mm == [(0.1,0.5),(0.1,0.5),(0.1,0.5),(0.1,0.5),(0.2,0.5)]
    mm = extrema_filter(A, 7)
    @test mm == [(0.1,0.5),(0.1,0.5),(0.1,0.5),(0.1,0.5),(0.2,0.5)]
    # 2d case
    A = rand(5,5)/10
    A[2,2] = 0.8
    A[4,4] = 0.6
    for w in ((2,2), (2,3), (3,2), (3,3), (2,5))
        mm = extrema_filter(A, w)
        maxval = map(last, mm)
        Amax = groundtruth(max, A, w)
        @test maxval == Amax
        minval = map(first, mm)
        Amin = groundtruth(min, A, w)
        @test minval == Amin
    end
    # 3d case
    A = rand(5,5,5)/10
    A[2,2,2] = 0.7
    A[4,4,2] = 0.4
    A[2,2,4] = 0.5
    for w in ((2,2,2), (2,3,2), (3,2,2), (2,2,3), (3,3,3), (2,5,3))
        mm = extrema_filter(A, w)
        maxval = map(last, mm)
        Amax = groundtruth(max, A, w)
        @test maxval == Amax
        minval = map(first, mm)
        Amin = groundtruth(min, A, w)
        @test minval == Amin
    end
end

nothing


# A quadratic irrational : (√a + b)/c
type QuadraticIrrational
  a_square::Integer
  b::Integer
  c::Integer
end

qi_sqrt(a::Integer) = QuadraticIrrational(a, 0, 1)

eval(q::QuadraticIrrational) = (sqrt(q.a_square) + q.b) / q.c

function -(q::QuadraticIrrational, i::Integer)
  return QuadraticIrrational(q.a_square, q.b - q.c * i, q.c)
end

function +(q::QuadraticIrrational, i::Integer)
  return QuadraticIrrational(q.a_square, q.b + q.c * i, q.c)
end

function /(i::Integer, q::QuadraticIrrational)
  # (i*c)/(√a + b) = ((i*c)*(√a + b))/(a - b^2)
  den = q.a_square - q.b^2
  num_f = i * q.c
  common_div = gcd(num_f, den)
  den /= common_div
  num_f /= common_div
  return QuadraticIrrational(num_f^2 * q.a_square, -num_f * q.b, den)
end

function show(io::IO, q::QuadraticIrrational)
  print(io, "(√$(q.a_square) + $(q.b)) / $(q.c)")
end


# This file solve the flap reponses of blades
function sflapre(ψ,λ_α,θ_cp,θ_lat,θ_lon) #staticflapre
  # 其中θ应当输入平均值，例如mean(θ_cp)
  γ_ = (ρ*8.2*(0.06)*R^4)/Iβ
  μ = μ_air
  F0 = 1+μ^2/2
  F_B1 = μ
  F_λ = 1
  A0 = 2*μ
  A_B1 = 1+3*μ^2/4
  A_λ = μ
  A_a1 = 1-μ^2/4
  B_β0 = μ
  B_A1 = 1+μ^2/4
  B_b1 = 1+μ^2/4
  γ_B = γ_/2*B_β0/B_b1
  θ0_ = θcp[1]
  θ_lat_ = θ_lat[1]
  θ_lon_ = θ_lon[1]
  Matrix_β = [γ_/2*F0 0 γ_/2*F_B1 γ_/2*F_λ;
              A0/A_a1 0 A_B1/A_a1 A_λ/A_a1;
              γ_B*F0 -B_A1/B_b1  γ_B*F_B1 γ_B*F_λ]*[θ0_,θ_lat_,θ_lon_,λ_α]
  β0 = Matrix_β[1]
  β_lon = Matrix_β[2]
  β_lat = Matrix_β[3]
  β = β0-β_lon*cos(ψ)-β_lat*sin(ψ)
  dβ = Ω*(β_lon*sin(ψ)-β_lat*cos(ψ))
  ddβ = Ω*(β_lon*cos(ψ)+β_lat*cos(ψ))
  return β,dβ,ddβ,β_lon,β_lat,β0
end

function dflapre(betanow,betaxnow,Mbeta_aero)	#动态挥舞响应
  # 输入：当前挥舞角、挥舞角速度
  # 输出：步进挥舞角、挥舞角速度
	# 常规参数
	Mcen = ecut*R*Ω^2*m_*1/2*(R-ecut*R)^2
	MG = m_*g*1/2*(R-ecut*R)^2
  nite = 10 # 迭代步数
  h = dψ/nite # 龙格库塔求解间隔
  ite = 1
  betanext = 0.0
  betaxnext = 0.0

	# 挥舞方程
	funcbetax(xxbeta) = 1/(Iβ*Ω^2)*(Mbeta_aero-MG*cos(xxbeta)-
											Mcen*sin(xxbeta))-
											cos(xxbeta)*sin(xxbeta)

	# 龙格库塔求解
  while ite<=nite
  	k1 = betaxnow
  	kx1 = funcbetax(betanow)
  	k2 = betaxnow+h/2*k1
  	kx2 = funcbetax(betanow+h/2*kx1)
  	k3 = betaxnow+h/2*k2
  	kx3 = funcbetax(betanow+h/2*kx2)
  	k4 = betaxnow+h*k3
  	kx4 = funcbetax(betanow+h*kx3)
  	betanext = betanow + h/6*(k1+2*k2+2*k3+k4)
  	betaxnext = betaxnow + h/6*(kx1+2*kx2+2*kx3+kx4)
    betanow = betanext
    betaxnow = betaxnext
    ite = ite+1
  end

	return betanext,betaxnext
end


#!/usr/bin/env julia

#=
 Licensing Information: You are free to use or extend these projects for
 education or reserach purposes provided that (1) you retain this notice
 and (2) you provide clear attribution to UC Berkeley, including a link
 to http://barc-project.com

 Attibution Information: The barc project ROS code-base was developed
 at UC Berkeley in the Model Predictive Control (MPC) lab by Jon Gonzales
 (jon.gonzales@berkeley.edu). The cloud services integation with ROS was developed
 by Kiet Lam  (kiet.lam@berkeley.edu). The web-server app Dator was
 based on an open source project by Bruce Wootton
=#

using RobotOS
@rosimport barc.msg: ECU, Encoder, Ultrasound, Z_KinBkMdl
@rosimport data_service.msg: TimeData
@rosimport geometry_msgs.msg: Vector3
@rosimport ackermann_msgs.msg: AckermannDrive, AckermannDriveStamped
@rosimport nav_msgs.msg: Odometry
rostypegen()
using barc.msg
using data_service.msg
using geometry_msgs.msg
using ackermann_msgs.msg
using nav_msgs.msg
using JuMP
using Ipopt

# define model parameters
L_a     = 0.125         # distance from CoG to front axel
L_b     = 0.125         # distance from CoG to rear axel
dt      = 0.1           # time step of system
a_max   = 1             # maximum acceleration

# preview horizon
N       = 5

# define targets [generic values]
x_ref   = 10
y_ref   = 0

# define decision variables
# states: position (x,y), yaw angle, and velocity
# inputs: acceleration, steering angle
println("Creating kinematic bicycle model ....")
mdl     = Model(solver = IpoptSolver(print_level=3))
@defVar( mdl, x[1:(N+1)] )
@defVar( mdl, y[1:(N+1)] )
@defVar( mdl, psi[1:(N+1)] )
@defVar( mdl, v[1:(N+1)] )
@defVar( mdl, a[1:N] )
@defVar( mdl, d_f[1:N] )

# define objective function
@setNLObjective(mdl, Min, (x[N+1] - x_ref)^2 + (y[N+1] - y_ref)^2 )

# define constraints
# define system dynamics
# Reference: R.Rajamani, Vehicle Dynamics and Control, set. Mechanical Engineering Series,
#               Spring, 2011, page 26
@defNLParam(mdl, x0     == 0); @addNLConstraint(mdl, x[1]     == x0);
@defNLParam(mdl, y0     == 0); @addNLConstraint(mdl, y[1]     == y0);
@defNLParam(mdl, psi0   == 0); @addNLConstraint(mdl, psi[1]   == psi0 );
@defNLParam(mdl, v0     == 0); @addNLConstraint(mdl, v[1]     == v0);
@defNLExpr(mdl, bta[i = 1:N], atan( L_a / (L_a + L_b) * tan(d_f[i]) ) )
for i in 1:N
    @addNLConstraint(mdl, x[i+1]    == x[i]      + dt*(v[i]*cos( psi[i] + bta[i] ))  )
    @addNLConstraint(mdl, y[i+1]    == y[i]      + dt*(v[i]*sin( psi[i] + bta[i] ))  )
    @addNLConstraint(mdl, psi[i+1]  == psi[i]    + dt*(v[i]/L_b * sin(bta[i]))  )
    @addNLConstraint(mdl, v[i+1]    == v[i]      + dt*(a[i])  )
    @addConstraint(mdl, 0 <=  a[i] <= a_max  )
end

# status update
println("initial solve ...")
solve(mdl)
println("finished initial solve!")

function SE_callback(msg::Odometry)
    # update mpc initial condition
    setValue(x0,    msg.pose.pose.position.x)
    setValue(y0,    msg.pose.pose.position.y)
    setValue(psi0,  msg.pose.pose.orientation)
    setValue(v0,    msg.twist.linear.z)
end

function main()
    # initiate node, set up publisher / subscriber topics
    init_node("mpc")
    pub = Publisher("ecu", AckermannDriveStamped, queue_size=10)
    sub  = Subscriber("state_estimate", Odometry, SE_callback, queue_size=10)
    loop_rate = Rate(10)

    while ! is_shutdown()
        # run mpc, publish command
        solve(mdl)

        # get optimal solutions
        a_opt   = getValue(a[1])
        d_f_opt = getValue(d_f[1])
        cmd = AckermannDriveStamped()
        cmd.drive.steering_angle = d_f_opt
        cmd.drive.speed = a_opt
        # publish commands
        publish(pub, cmd)
        rossleep(loop_rate)
    end
end

if ! isinteractive()
    main()
end


using DSP, DelimitedFiles, LsqFit

function DT_model(t, β::Array{Float64,1}, ρ::Float64, ndet::Float64, nmed::Float64)
	#optimzed semi-infinite model

    n::Float64 = nmed/ndet
    μa::Float64 = β[1]
    μsp::Float64 = β[2]
    D::Float64 = 1/3μsp
	ν::Float64 = 29.9792345/nmed

    Rt = Array{Float64}(undef, length(t))

	if n > 1.0
		A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		 5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
	elseif n < 1.0
		A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
	else
		A = 1.0
	end

	zs::Float64 = 1/μsp
	ze::Float64 = 2A*D

	z3m::Float64 = - zs
    z4m::Float64 = 2ze +zs

	Threads.@threads for n in eachindex(t)

   		Rt[n] = -exp(-(ρ^2/(4D*ν*t[n])) - μa*ν*t[n])
   		Rt[n] = Rt[n]/(2*(4π*D*ν)^(3/2)*t[n]^(5/2))

   		Rt[n] = Rt[n]*(z3m*exp(-(z3m^2/(4D*ν*t[n]))) - z4m*exp(-(z4m^2/(4D*ν*t[n]))))

	end

   replace!(Rt, NaN => 0)
   return Rt./maximum(Rt)

end

function get_fit_window(counts, risefactor, tailfactor)

    maxval = maximum(counts)
    ind1::Int = findfirst(x -> x > maxval*risefactor, counts)
    ind2::Int =  findlast(x -> x >maxval*tailfactor, counts)

    return ind1, ind2
end

function pads(u::Array{Float64,1})
    outsize = 2 .* length(u) .- 1
    nffts = nextfastfft(outsize)
    out = zeros(eltype(u), nffts)
end

function get_fftplan(u::Array{Float64,1}, out)
    upad = zeros(eltype(u), length(out))
    copyto!(upad,u)
    pl = plan_rfft(upad)
    upadRfft = pl * upad
    return pl, upadRfft
end

function conv_DT(t, β::Array{Float64,1}, data::input_data, out, pl,upadRfft;
    vpad = zeros(eltype(data.DTOF), length(out)),
    RtDT = Array{Float64}(undef, length(data.IRF)),
    vpadRfft = zeros(eltype(upadRfft), length(upadRfft))
    )

        RtDT = DT_model(data.t, β, data.ρ, data.nmed, data.ndet)
        copyto!(vpad, RtDT)

        vpadRfft = pl * vpad
        vpadRfft .*= upadRfft
        out = irfft(vpadRfft, length(out))

        out = out./maximum(out)

	    tidx = findfirst(x -> x == t[1], data.t)

	    out = out[tidx:tidx+length(t)-1]
	return log.(out)
end

function getfit(input_data, model_params)
    ## get FFT padding and plans
    out::Array{eltype(input_data.IRF)} = pads(input_data.IRF)
    pl, upadRfft = get_fftplan(input_data.IRF, out)

    #get fit windows
    ind1, ind2 = get_fit_window(input_data.DTOF, model_params.risefactor, model_params.tailfactor)

    curve_fit((t, β) -> conv_DT(t, β, input_data, out, pl, upadRfft), input_data.t[ind1:ind2], log.(input_data.DTOF[ind1:ind2]),
    model_params.initparams, lower=model_params.lb, upper=model_params.ub)

end

#= 12/3/2020
@benchmark getfit(data,modelp)
BenchmarkTools.Trial:
  memory estimate:  30.68 MiB
  allocs estimate:  9561
  --------------
  minimum time:     15.015 ms (0.00% GC)
  median time:      16.626 ms (5.95% GC)
  mean time:        18.319 ms (3.91% GC)
  maximum time:     64.462 ms (1.59% GC)
  --------------
  samples:          273
  evals/sample:     1
=#


#---------------------------------------------------------------------#
#This function computes the Grid and Elements in 1D.
#Written by F.X. Giraldo on April 19, 2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
function create_grid(Np,Npoin_cg,Npoin_dg,Ne,x,DFloat)

    coord_cg=zeros(DFloat,Npoin_cg)
    coord_dg=zeros(DFloat,Npoin_dg)
    intma_cg=zeros(Int64,Np,Ne)
    intma_dg=zeros(Int64,Np,Ne)
    periodicity_cg=zeros(Int64,Npoin_cg)
    periodicity_dg=zeros(Int64,Npoin_dg)

    #Set some constants
    xmin=-1
    xmax=+1
    dx=(xmax-xmin)/Ne

    #Generate Grid Points
    I=1
    coord_cg[1]=xmin
    for e=1:Ne
        x0=xmin + (e-1)*dx
        intma_cg[1,e]=I
        for i=2:Np
            I+=1
            coord_cg[I]=( x[i]+1 )*dx/2 + x0
            intma_cg[i,e]=I
        end
    end

    #Generate Periodicity_CG array
    for I=1:Npoin_cg
        periodicity_cg[I]=I
    end
    periodicity_cg[Npoin_cg]=1

    #Generate INTMA_DG
    I=0
    for e=1:Ne, i=1:Np
        I+=1
        intma_dg[i,e]=I
    end

    #Generate COORD_DG
    for e=1:Ne, i=1:Np
        Icg=intma_cg[i,e]
        Idg=intma_dg[i,e]
        coord_dg[Idg]=coord_cg[Icg]
    end

    #Generate Periodicity_DG array
    for I=1:Npoin_dg
        periodicity_dg[I]=I
    end

    return (coord_cg, coord_dg, intma_cg, intma_dg, periodicity_cg, periodicity_dg)
end


"""
    gbrand(typeorrange, nrows, ncols, density; kwargs...)::GBMatrix
    gbrand(rng::AbstractRNG, typeorrange, nrows, ncols, density; kwargs...)::GBMatrix

Construct a random `GBMatrix`, analogous to `sprand` from SparseArrays

# Arguments
- `rng::AbstractRNG`: Random number generator for both values and indices.
- `typeorrange`: Either a type such as `Float64`, or a range such as `1:10`.
Any input which supports `eltype(typeorrange)`.
- `nrows::Integer`, `ncols::Integer`: Dimensions of the result.
- `density::AbstractFloat`: The approximate density of result.

# Keywords
- `symmetric::Bool`: The result matrix is symmetric, Aᵀ = A.
- `pattern::Bool`: The result matrix consists solely of `one(eltype(typeorrange))`.
- `skewsymmetric::Bool`: The result matrix is skew-symmetric, Aᵀ = -A.
- `hermitian::Bool`: The result matrix is hermitian, aᵢⱼ = āⱼᵢ.
- `nodiagonal::Bool`: The result matrix has no values on the diagonal.

# Returns
- `GBMatrix`
"""
function gbrand(
    rng::AbstractRNG, typeorrange, nrows::Integer, ncols::Integer, density::AbstractFloat;
    symmetric=false, pattern=false, skewsymmetric=false, hermitian=false, nodiagonal=false
)
    type = eltype(typeorrange)
    A = GBMatrix{type}(nrows, ncols)
    (type <: Unsigned || type == Bool) && (skewsymmetric = false)
    if nrows != ncols
        symmetric = false
        skewsymmetric = false
        hermitian = false
    end
    if pattern || symmetric
        skewsymmetric = false
        hermitian = false
    end
    if skewsymmetric
        hermitian = false
        nodiagonal = false
    end
    if !(type <: Complex)
        hermitian = false
    end

    for _ ∈ 1:round(Int64, nrows * ncols * density)
        i = rand(rng, 1:nrows)
        j = rand(rng, 1:ncols)
        nodiagonal && (i == j) && continue
        if pattern
            x = one(type)
        else
            x = rand(rng, typeorrange)
        end
        A[i, j] = x
        symmetric && (A[j, i] = x)
        skewsymmetric && (A[j, i] = -x)
        hermitian && (A[j, i] = conj(x))
    end

    return A
end

function gbrand(
    typeorrange, nrows::Integer, ncols::Integer, density::AbstractFloat;
    symmetric=false, pattern=false, skewsymmetric=false, hermitian=false, nodiagonal=false
)
    return gbrand(
        default_rng(), typeorrange, nrows, ncols, density;
        symmetric, pattern, skewsymmetric, hermitian, nodiagonal
    )
end


using TopOpt
import Makie
using TopOpt.TopOptProblems.Visualization: visualize

using FromFile
@from "new_problems.jl" import NewPointLoadCantilever

using TimerOutputs

# function run_topopt()
    println("Start running.")
    # https://github.com/KristofferC/TimerOutputs.jl
    to = TimerOutput()
    reset_timer!(to)
    Nonconvex.show_residuals[] = true

    # Define the problem
    E = 1.0 # Young’s modulus
    v = 0.3 # Poisson’s ratio
    f = 1.0; # downward force

    # Parameter settings
    V = 0.3 # volume fraction
    # xmin = 0.001 # minimum density
    xmin = 1e-6 # minimum density
    rmin = 2.0; # density filter radius

    nels = (20, 10, 2)
    sizes = (1.0, 1.0, 1.0)
    @timeit to "problem def" problem = NewPointLoadCantilever(Val{:Linear}, nels, sizes, E, v, f);

    # Define a finite element solver
    @timeit to "penalty def" penalty = TopOpt.PowerPenalty(3.0)
    @timeit to "solver def" solver = FEASolver(Displacement, Direct, problem, xmin = xmin,
        penalty = penalty);

    # Define compliance objective
    @timeit to "objective def" begin
        # Define compliance objective
        comp = Compliance(problem, solver)
        filter = DensityFilter(solver, rmin = rmin)
        obj = x -> comp(filter(x))
    end

    # Define volume constraint
    @timeit to "constraint def" begin
        volfrac = TopOpt.Volume(problem, solver)
        constr = x -> volfrac(filter(x)) - V
    end

    # Define subproblem optimizer
    # TODO MMA02 converge to weird results under this criteria
    mma_options = options = Nonconvex.MMAOptions(
        maxiter = 400,
        tol = Nonconvex.Tolerance(x = 1e-3, f = 1e-6),
        )
    ipopt_options = Nonconvex.IpoptOptions(
        # maxiter = 100,
        # tol = Nonconvex.Tolerance(x = 1e-3, f = 1e-6),
    )

    convcriteria = Nonconvex.GenericCriteria()
    # convcriteria = Nonconvex.KKTCriteria()

    x0 = fill(V, length(solver.vars))
    @timeit to "optimizer def" optimizer = Optimizer(obj, constr, x0,
        Nonconvex.MMA87(),
        options = mma_options,
        # Nonconvex.IpoptAlg(),
        # options = ipopt_options,
        convcriteria = convcriteria
        );
    # Nonconvex.MMA02(),

    # Define SIMP optimizer
    @timeit to "simp def" simp = SIMP(optimizer, solver, penalty.p);

    # Solve
    # initial solution, critical to set it to volfrac! (blame non-convexity :)
    @timeit to "simp run" result = simp(x0);

    # Print the timings in the default way
    show(to)
    println()

    @show result.convstate
    @show result.objval
    try
        @show optimizer.workspace.iter
    catch
        # IpoptWorkspace has no field iter
    end

    # # Visualize the result using Makie.jl
    fig = visualize(problem; topology=result.topology,
        default_exagg_scale=0.07, scale_range=10.0, vector_linewidth=3, vector_arrowsize=0.5)
    Makie.display(fig)

    # return problem, result
# end

# problem, result = run_topopt();


using PyPlot
using Peacock
using CUDA
GPU = 1

function epf(x,y)
    # equation of a circle with radius 0.2a
    if x^2+y^2 <= 0.2^2
        # dielectric inside the circle
        return 8.9
    else
        # air outside the circle
        return 1
    end
end

# Permeability is unity everywhere
function muf(x,y)
    return 1
end

a1 = [1, 0]  # first lattice vector
a2 = [0, 1]  # second lattice vector
d1 = 0.01  # resolution along first lattice vector
d2 = 0.01  # resolution along second lattice vector
geometry = Geometry(epf, muf, a1, a2, d1, d2)
fourier_space_cutoff = 7
solver = Solver(geometry, fourier_space_cutoff)
G = BrillouinZoneCoordinate(  0,   0, "Γ")
X = BrillouinZoneCoordinate(1/2,   0, "X")
M = BrillouinZoneCoordinate(1/2, 1/2, "M")
ks = [G,X,M,G]

figure(figsize=(4,3))
@time plot_band_diagram(solver, ks, TE, GPU, color="red",
            bands=1:4, dk=0.1, frequency_scale=1/2pi)
@time plot_band_diagram(solver, ks, TM, GPU, color="blue",
            bands=1:4, dk=0.1, frequency_scale=1/2pi)
ylim(0,0.8)


using DataDrivenDiffEq
using ModelingToolkit
using OrdinaryDiffEq

using LinearAlgebra

# Create a test problem
function lorenz(u,p,t)
    x, y, z = u

    ẋ = 10.0*(y - x)
    ẏ = x*(28.0-z) - y
    ż = x*y - (8/3)*z
    return [ẋ, ẏ, ż]
end

u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
dt = 0.1
prob = ODEProblem(lorenz,u0,tspan)
sol = solve(prob, Tsit5(), saveat = dt, progress = true)

## Start the automatic discovery
ddprob = ContinuousDataDrivenProblem(sol)

@variables t x(t) y(t) z(t)
u = [x;y;z]
basis = Basis(polynomial_basis(u, 5), u, iv = t)
opt = STLSQ(exp10.(-5:0.1:-1))
ddsol = solve(ddprob, basis, opt, normalize = true)
print(ddsol, Val{true})


# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

using UnitCommitment, Test, LinearAlgebra

@testset "Sensitivity" begin
    @testset "Susceptance matrix" begin
        instance = UnitCommitment.read_benchmark("test/case14")
        actual = UnitCommitment.susceptance_matrix(instance.lines)
        @test size(actual) == (20, 20)
        expected = Diagonal([29.5,  7.83,  8.82, 9.9,  10.04,
                             10.2,  41.45, 8.35, 3.14, 6.93,
                             8.77,  6.82,  13.4, 9.91, 15.87,
                             20.65, 6.46,  9.09, 8.73, 5.02])
        @test round.(actual, digits=2) == expected
    end

    @testset "Reduced incidence matrix" begin
        instance = UnitCommitment.read_benchmark("test/case14")
        actual = UnitCommitment.reduced_incidence_matrix(lines=instance.lines,
                                                         buses=instance.buses)
        @test size(actual) == (20, 13)
        @test actual[1,   1]  == -1.0
        @test actual[3,   1]  ==  1.0
        @test actual[4,   1]  ==  1.0
        @test actual[5,   1]  ==  1.0
        @test actual[3,   2]  == -1.0
        @test actual[6,   2]  ==  1.0
        @test actual[4,   3]  == -1.0
        @test actual[6,   3]  == -1.0
        @test actual[7,   3]  ==  1.0
        @test actual[8,   3]  ==  1.0
        @test actual[9,   3]  ==  1.0
        @test actual[2,   4]  == -1.0
        @test actual[5,   4]  == -1.0
        @test actual[7,   4]  == -1.0
        @test actual[10,  4]  ==  1.0
        @test actual[10,  5]  == -1.0
        @test actual[11,  5]  ==  1.0
        @test actual[12,  5]  ==  1.0
        @test actual[13,  5]  ==  1.0
        @test actual[8,   6]  == -1.0
        @test actual[14,  6]  ==  1.0
        @test actual[15,  6]  ==  1.0
        @test actual[14,  7]  == -1.0
        @test actual[9,   8]  == -1.0
        @test actual[15,  8]  == -1.0
        @test actual[16,  8]  ==  1.0
        @test actual[17,  8]  ==  1.0
        @test actual[16,  9]  == -1.0
        @test actual[18,  9]  ==  1.0
        @test actual[11, 10]  == -1.0
        @test actual[18, 10]  == -1.0
        @test actual[12, 11]  == -1.0
        @test actual[19, 11]  ==  1.0
        @test actual[13, 12]  == -1.0
        @test actual[19, 12]  == -1.0
        @test actual[20, 12]  ==  1.0
        @test actual[17, 13]  == -1.0
        @test actual[20, 13]  == -1.0
    end

    @testset "Injection Shift Factors (ISF)" begin
        instance = UnitCommitment.read_benchmark("test/case14")
        actual = UnitCommitment.injection_shift_factors(lines=instance.lines,
                                                        buses=instance.buses)
        @test size(actual) == (20, 13)
        @test round.(actual, digits=2) == [
            -0.84 -0.75 -0.67 -0.61 -0.63 -0.66 -0.66 -0.65 -0.65 -0.64 -0.63 -0.63 -0.64;
            -0.16 -0.25 -0.33 -0.39 -0.37 -0.34 -0.34 -0.35 -0.35 -0.36 -0.37 -0.37 -0.36;
             0.03 -0.53 -0.15 -0.1  -0.12 -0.14 -0.14 -0.14 -0.13 -0.13 -0.12 -0.12 -0.13;
             0.06 -0.14 -0.32 -0.22 -0.25 -0.3  -0.3  -0.29 -0.28 -0.27 -0.25 -0.26 -0.27;
             0.08 -0.07 -0.2  -0.29 -0.26 -0.22 -0.22 -0.22 -0.23 -0.25 -0.26 -0.26 -0.24;
             0.03  0.47 -0.15 -0.1  -0.12 -0.14 -0.14 -0.14 -0.13 -0.13 -0.12 -0.12 -0.13;
             0.08  0.31  0.5  -0.3  -0.03  0.36  0.36  0.28  0.23  0.1  -0.0   0.02  0.17;
             0.0   0.01  0.02 -0.01 -0.22 -0.63 -0.63 -0.45 -0.41 -0.32 -0.24 -0.25 -0.36;
             0.0   0.01  0.01 -0.01 -0.12 -0.17 -0.17 -0.26 -0.24 -0.18 -0.14 -0.14 -0.21;
            -0.0  -0.02 -0.03  0.02 -0.66 -0.2  -0.2  -0.29 -0.36 -0.5  -0.63 -0.61 -0.43;
            -0.0  -0.01 -0.02  0.01  0.21 -0.12 -0.12 -0.17 -0.28 -0.53  0.18  0.15 -0.03;
            -0.0  -0.0  -0.0   0.0   0.03 -0.02 -0.02 -0.03 -0.02  0.01 -0.52 -0.17 -0.09;
            -0.0  -0.01 -0.01  0.01  0.11 -0.06 -0.06 -0.09 -0.05  0.02 -0.28 -0.59 -0.31;
            -0.0  -0.0  -0.0  -0.0  -0.0  -0.0  -1.0  -0.0  -0.0  -0.0  -0.0  -0.0   0.0 ;
             0.0   0.01  0.02 -0.01 -0.22  0.37  0.37 -0.45 -0.41 -0.32 -0.24 -0.25 -0.36;
             0.0   0.01  0.02 -0.01 -0.21  0.12  0.12  0.17 -0.72 -0.47 -0.18 -0.15  0.03;
             0.0   0.01  0.01 -0.01 -0.14  0.08  0.08  0.12  0.07 -0.03 -0.2  -0.24 -0.6 ;
             0.0   0.01  0.02 -0.01 -0.21  0.12  0.12  0.17  0.28 -0.47 -0.18 -0.15  0.03;
            -0.0  -0.0  -0.0   0.0   0.03 -0.02 -0.02 -0.03 -0.02  0.01  0.48 -0.17 -0.09;
            -0.0  -0.01 -0.01  0.01  0.14 -0.08 -0.08 -0.12 -0.07  0.03  0.2   0.24 -0.4  ]
    end

    @testset "Line Outage Distribution Factors (LODF)" begin
        instance = UnitCommitment.read_benchmark("test/case14")
        isf_before = UnitCommitment.injection_shift_factors(lines=instance.lines,
                                                            buses=instance.buses)
        lodf = UnitCommitment.line_outage_factors(lines=instance.lines,
                                                  buses=instance.buses,
                                                  isf=isf_before)
        for contingency in instance.contingencies
            for lc in contingency.lines
                prev_susceptance = lc.susceptance
                lc.susceptance = 0.0
                isf_after =  UnitCommitment.injection_shift_factors(lines=instance.lines,
                                                                    buses=instance.buses)
                lc.susceptance = prev_susceptance
                for lm in instance.lines
                    expected = isf_after[lm.offset, :]
                    actual = isf_before[lm.offset, :] +
                                 lodf[lm.offset, lc.offset] * isf_before[lc.offset, :]
                    @test norm(expected - actual) < 1e-6
                end
            end
        end
    end
end


using ModelingToolkit, OrdinaryDiffEq, CardiovascularModels

@variables t
D = Differential(t)

@named vessel1 = ElasticVessel(Ees = 5e6, Vd = 0.0)
@named resistor = Resistor(R = 10e6)
@named vessel2 = ElasticVessel(Ees = 1e7, Vd = 0.0)

eqs_con = [
connect(vessel1.out, resistor.in),
connect(resistor.out, vessel2.in),
vessel1.in.Q ~ 0.,
vessel2.out.Q ~ 0.
]

volume_start = [
    vessel1.V => 1e-4,
    vessel2.V => 2e-4

]

time_span = (0.0, 5.0)

@named connect_sys = ODESystem(eqs_con, t; systems = [vessel1, vessel2, resistor])
problem = ODEProblem(structural_simplify(connect_sys), volume_start, time_span, [])

sol = solve(problem, Tsit5(), dtmax = 0.01, reltol = 1e-6)

# Plots.plot(sol)


module Digitization
using ParallelKMeans
using Statistics

export digitization_k, Model

mutable struct Model
   centers::Array{Float64,}
   hashmap::Dict{Char,Int64}
end

function digitization_k(array, k::Int, nthreads::Int)
    features = collect(Matrix(array[:, 1:2])');
    result = kmeans(features, k; n_threads=nthreads);
    labels = result.assignments;
    symbols = map(x -> Char(x+33), labels);
    hashmap = Dict(zip(symbols, labels));
    centers = zeros((0, 2));
    features = features';
    for i in sort(unique(labels))
        condition = (labels .== i);
        group = features[condition, :];
        center = mean!([1. 1.], group);
        centers = vcat(centers, center);
    end
    model = Model(centers, hashmap);
    return symbols, model;
end

end


#__precompile__()

module ExperimentalDesign

using Random
using Logging
using Primes
using DataFrames
using Distributions
using StatsModels
using DocStringExtensions
using LinearAlgebra

import Random: rand

# Types

export    AbstractDesign,   AbstractScreeningDesign,    AbstractFactorialDesign,
    AbstractOptimalDesign,  PlackettBurman, FullFactorial,  FractionalFactorial,
    OptimalDesign, RandomDesign

# Methods

export fullfactorial,  explicit_fullfactorial, plackettburman, isplackettburman,
    paley,     optimize_design,     expanded_design,     generate_random_design,
    rand, random_design, d_criterion, next_offset_divisible_prime

include("design.jl")
include("plackett_burman.jl")
include("factorial.jl")
include("random.jl")
include("d_optimal/variance_predictions.jl")

end


# Note: poor performance

# this may not be working properly

cd(@__DIR__)
using AdvancedPS, Random, StatsBase, Distributions

N = 30
# Parameters for each cluster, we assume that each cluster is Gaussian distributed in the example.
μs = [-3.5, 0.0]
# Construct the data points.
data = mapreduce(c -> rand(MvNormal([μs[c], μs[c]], 1.), N), hcat, 1:2)

priors = (
    idx=(Categorical([.5,.5]),N),
    μs = (Normal(0,1),2),
)

bounds = ((1,2),(0.0,Inf))

function loglike(idx, μs, data)
    LL = 0.0
    for (i,id) in enumerate(idx)
        LL += logpdf(MvNormal([μs[id], μs[id]], 1.0), data[:,i])
    end
    return LL
end

loglike(θ) = loglike(θ..., data)

model = DEModel(priors=priors, model=loglike)

de = DE(bounds=bounds, burnin=1000, priors=priors, progress=true, Np=15)
n_iter = 2000
@elapsed chains = sample(model, de, n_iter)
println(chains)


import StatsFuns: log2π

@node MvNormalWeightedMeanPrecision Stochastic [ out, (ξ, aliases = [ xi, weightedmean ]), (Λ, aliases = [ invcov, precision ]) ]

conjugate_type(::Type{ <: MvNormalWeightedMeanPrecision }, ::Type{ Val{ :Λ } })   = Wishart

@average_energy MvNormalWeightedMeanPrecision (q_out::Any, q_ξ::Any, q_Λ::Any) = begin
    (m_mean, v_mean) = mean_cov(q_ξ)
    (m_out, v_out)   = mean_cov(q_out)
    0.5 * (ndims(q_out) * log2π + logdet(cholinv(mean(q_Λ))) + tr(mean(q_Λ)*(v_out + v_mean + (m_out - m_mean)*(m_out - m_mean)')))
end

@average_energy MvNormalWeightedMeanPrecision (q_out_ξ::Any, q_Λ::Any) = begin
    (m, V) = mean_cov(q_out_ξ)
    d = Int64(ndims(q_out_ξ)/2)
    @views 0.5*(d*log2π + logdet(cholinv(mean(q_Λ))) + tr(mean(q_Λ)*( V[1:d,1:d] - V[1:d,d+1:end] - V[d+1:end,1:d] + V[d+1:end,d+1:end] + (m[1:d] - m[d+1:end])*(m[1:d] - m[d+1:end])' ) ))
end


function main()
    # Part 1: compute "FFT" on input
    input = open(f -> parse.(Int64, collect(readline(f))), "day16_input.txt")
    base(i, n) = [0, 1, 0, -1][(i ÷ n) % 4 + 1]
    fft(input, n) = abs(sum([d * base(i, n) for (i, d) in enumerate(input)])) % 10
    result = foldl((x, y) -> [fft(x, n) for n in 1:length(input)], 1:3, init=input)
    println("Part 1: After 100 iterations: $(join(result[1:8], ""))")

    # Part 2: repeat input 10000 times and then use first 7 digits as offset to find 8 digit message
    # Note that in each phase, a digit only relies on digits in its position or later from prev phase
    big_input = repeat(input, 10000)
    N = length(big_input)
    offset = parse(Int64, join(input[1:7]))
    if offset < N÷2
        println("Oh dear, my clever accumulator hack won't work")
        return
    end
    current = copy(big_input)
    for phase in 1:100
        # Last half of output only relies on the [00001111]-style base sequence, i.e. upper triangular type
        # We can just compute backwards for the digits we need
        result = zeros(Int64, N)
        for d in N:-1:(N÷2 + 1)
            result[d] = (d == N) ? current[d] : ((current[d] + result[d+1]) % 10)
        end
        current = result
    end
    println("Part 2: After 100 iterations, message: $(join(current[offset+1:offset+8]))")
end
main()


#%% cell
#=
Now that we have a model, we will construct one more helper function that takes
the control variables $(S, F, C)$ and shocks $(\delta, R, \eta_G, \eta_a,
\eta_L, \eta_R)$ and applies equilibrium conditions to produce $(\pi, \delta',
Y, L, Y_n, R')$. We will use this function in both the solution and simulation
routines below.
=#

#%% cell
function Base.step(m::Model, S, F, C, δ0, R0, ηG, ηa, ηL, ηR)
    @unpack Θ, ϵ, gbar, ϑ, γ, β, μ, ϕ_π, ϕ_y, πstar = m.p
    πs = m.s.π

    # Compute pie(t) from condition (35) in MM (2015)
    π0 = @. ((1-(1-Θ)*(S/F)^(1-ϵ))/Θ)^(1/(ϵ-1))

    # Compute delta(t) from condition (36) in MM (2015)
    δ1 = @. ((1-Θ)*((1-Θ*π0^(ϵ-1))/(1-Θ))^(ϵ/(ϵ-1))+Θ*π0^ϵ/δ0)^(-1.0)

    # Compute Y(t) from condition (38) in MM (2015)
    Y0 = @. C/(1-gbar/exp(ηG))

    # Compute L(t) from condition (37) in MM (2015)
    L0 = @. Y0/exp(ηa)/δ1

    #  Compute Yn(t) from condition (31) in MM (2015)
    Yn0 = @. (exp(ηa)^(1+ϑ)*(1-gbar/exp(ηG))^(-γ)/exp(ηL))^(1/(ϑ+γ))

    # Compute R(t) from conditions (27), (39) in MM (2015) -- Taylor rule
    R1 = @. πstar / β * (R0*β/πstar)^μ*((π0/πstar)^ϕ_π * (Y0/Yn0)^ϕ_y)^(1-μ)*exp(ηR)

    π0, δ1, Y0, L0, Yn0, R1
end
#%% cell
#=
### Solution routine
=#

#%% cell
# construct an initial guess for the solution
function initial_coefs(m::Model, degree)
    npol = size(m.g.X0_G[degree], 2)
    coefs = fill(1e-5, npol, 3)
    coefs[1, :] = [m.s.S, m.s.F, m.s.C^(-m.p.γ)]
    coefs
end

function solve(m::Model; verbose::Bool=false)
    # simplify notation
    n, n_nodes = size(m.g.ηR, 1), length(m.g.ω_nodes)

    ## allocate memory
    # euler equations
    e = zeros(n, 3)

    # previous iteration S, F, C
    S0_old_G = ones(n)
    F0_old_G = ones(n)
    C0_old_G = ones(n)

    # current iteration S, F, C
    S0_new_G = ones(n)
    F0_new_G = ones(n)
    C0_new_G = ones(n)

    # future S, F, C
    S1 = zeros(n, n_nodes)
    F1 = zeros(n, n_nodes)
    C1 = zeros(n, n_nodes)
    π1 = zeros(n, n_nodes)

    local coefs, start_time

    degs = m.p.deg > 1 ? [1, m.p.deg] : [m.p.deg]

    for deg in degs
        # set up matrices for this degree
        X0_G = m.g.X0_G[deg]

        # future basis matrix for S, F, C
        X1 = Array{Float64}(undef, n, n_complete(8, deg))

        # initialize coefs
        if deg <= 1
            coefs = initial_coefs(m, deg)
        else
            # for higher order, start with degree 2 coefficients
            coefs = X0_G\e
            # old_coefs = copy(coefs)
            # coefs = initial_coefs(m, deg)
            # coefs[1:size(old_coefs, 1), :] = old_coefs
        end

        err = 1.0
        it = 0

        start_time = time()
        # solve at this degree of complete polynomial
        while err > m.p.tol
            it += 1;
            # Current choices (at t)
            # ------------------------------
            S0 = X0_G*coefs[:, 1]                # Compute S(t) using coefs
            F0 = X0_G*coefs[:, 2]                # Compute F(t) using coefs
            C0 = X0_G*coefs[:, 3]
            C0 .^= (-1/m.p.γ)                    # Compute C(t) using coefs

            π0, δ1, Y0, L0, Yn0, R1 = step(m, S0, F0, C0, m.g.δ, m.g.R, m.g.ηG,
                                           m.g.ηa, m.g.ηL, m.g.ηR)

            if m.p.zlb R1 .= max.(R1, 1.0) end

            for u in 1:n_nodes
                # Form complete polynomial of degree "Degree" (at t+1) on future state
                complete_polynomial!(
                    X1,
                    hcat(log.(R1), log.(δ1), m.g.ηR1[:, u], m.g.ηa1[:, u],
                         m.g.ηL1[:, u], m.g.ηu1[:, u], m.g.ηB1[:, u], m.g.ηG1[:, u]),
                    deg
                )

                S1[:, u] = X1*coefs[:, 1]                # Compute S(t+1)
                F1[:, u] = X1*coefs[:, 2]                # Compute F(t+1)
                C1[:, u] = (X1*coefs[:, 3]).^(-1/m.p.γ)  # Compute C(t+1)
            end

            # Compute next-period π using condition
            # (35) in MM (2015)
            @. π1 = ((1-(1-m.p.Θ).*(S1./F1).^(1-m.p.ϵ))/m.p.Θ).^(1/(m.p.ϵ-1))

            # Evaluate conditional expectations in the Euler equations
            #---------------------------------------------------------
            e[:, 1] = @.(exp(m.g.ηu)*exp(m.g.ηL)*L0^m.p.ϑ*Y0/exp(m.g.ηa) + (m.p.β*m.p.Θ*π1^m.p.ϵ*S1))*m.g.ω_nodes
            e[:, 2] = @.(exp(m.g.ηu)*C0^(-m.p.γ)*Y0 + (m.p.β*m.p.Θ*π1^(m.p.ϵ-1)*F1))*m.g.ω_nodes
            e[:, 3] = @.(m.p.β*exp(m.g.ηB)/exp(m.g.ηu)*R1*((exp(m.g.ηu1).*C1^(-m.p.γ)/π1)))*m.g.ω_nodes

            # Variables of the current iteration
            #-----------------------------------
            copy!(S0_new_G, S0)
            copy!(F0_new_G, F0)
            copy!(C0_new_G, C0)

            # Compute and update the coefficients of the decision functions
            # -------------------------------------------------------------
            coefs_hat = X0_G\e   # Compute the new coefficients of the decision
                                 # functions using a backslash operator

            # Update the coefficients using damping
            coefs = m.p.damp*coefs_hat + (1-m.p.damp)*coefs

            # Evaluate the percentage (unit-free) difference between the values
            # on the grid from the previous and current iterations
            # -----------------------------------------------------------------
            # The convergence criterion is adjusted to the damping parameters
            err = mean(abs, 1.0.-S0_new_G./S0_old_G) +
                  mean(abs, 1.0.-F0_new_G./F0_old_G) +
                  mean(abs, 1.0.-C0_new_G./C0_old_G)

            if (it % 20 == 0) && verbose
                @printf "On iteration %d err is %6.7e\n" it err
            end

            # Store the obtained values for S(t), F(t), C(t) on the grid to
            # be used on the subsequent iteration in Section 10.2.6
            #-----------------------------------------------------------------------
            copy!(S0_old_G, S0_new_G)
            copy!(F0_old_G, F0_new_G)
            copy!(C0_old_G, C0_new_G)

        end
    end

    coefs, time() - start_time
end

#%% cell
#=
### Simulation
=#

#%% cell
struct Simulation
    # shocks
    ηR::Vector{Float64}
    ηa::Vector{Float64}
    ηL::Vector{Float64}
    ηu::Vector{Float64}
    ηB::Vector{Float64}
    ηG::Vector{Float64}

    # variables
    δ::Vector{Float64}
    R::Vector{Float64}
    S::Vector{Float64}
    F::Vector{Float64}
    C::Vector{Float64}
    π::Vector{Float64}
    Y::Vector{Float64}
    L::Vector{Float64}
    Yn::Vector{Float64}
    w::Vector{Float64}
end

function Simulation(m::Model, coefs::Matrix)

    # 11. Simualating a time-series solution
    #---------------------------------------
    _path = joinpath(@__DIR__, "..", "epsi_test_NK.mat")
    rands = (matread(_path)["epsi_test_NK"])::Matrix{Float64}
    capT = size(rands, 1)

    # Initialize the values of 6 exogenous shocks
    #--------------------------------------------
    ηR = zeros(capT)
    ηa = zeros(capT)
    ηL = zeros(capT)
    ηu = zeros(capT)
    ηB = zeros(capT)
    ηG = zeros(capT)

    # Generate the series for shocks
    #-------------------------------
    @inbounds for t in 1:capT-1
        ηR[t+1] = m.p.ρηR*ηR[t] + m.p.σηR*rands[t, 1]
        ηa[t+1] = m.p.ρηa*ηa[t] + m.p.σηa*rands[t, 2]
        ηL[t+1] = m.p.ρηL*ηL[t] + m.p.σηL*rands[t, 3]
        ηu[t+1] = m.p.ρηu*ηu[t] + m.p.σηu*rands[t, 4]
        ηB[t+1] = m.p.ρηB*ηB[t] + m.p.σηB*rands[t, 5]
        ηG[t+1] = m.p.ρηG*ηG[t] + m.p.σηG*rands[t, 6]
    end

    δ  = ones(capT+1) # Time series of delta(t)
    R  = ones(capT+1) # Time series of R(t)
    S  = Array{Float64}(undef, capT)   # Time series of S(t)
    F  = Array{Float64}(undef, capT)   # Time series of F(t)
    C  = Array{Float64}(undef, capT)   # Time series of C(t)
    π  = Array{Float64}(undef, capT)   # Time series of π(t)
    Y  = Array{Float64}(undef, capT)   # Time series of Y(t)
    L  = Array{Float64}(undef, capT)   # Time series of L(t)
    Yn = Array{Float64}(undef, capT)   # Time series of Yn(t)
    w  = Array{Float64}(undef, capT)

    pol_bases = Array{Float64}(undef, 1, size(coefs, 1))
    @inbounds for t in 1:capT
        # Construct the matrix of explanatory variables "pol_bases" on the series
        # of state variables; columns of "pol_bases" are given by the basis
        # functions of the polynomial of degree 2
        complete_polynomial!(
            pol_bases,
            hcat(log(R[t]), log(δ[t]), ηR[t], ηa[t], ηL[t], ηu[t], ηB[t], ηG[t]),
            m.p.deg
        )
        S[t], F[t], MU = pol_bases*coefs
        C[t] = (MU).^(-1/m.p.γ)

        π[t], δ[t+1], Y[t], L[t], Yn[t], R[t+1] = step(m, S[t], F[t], C[t], δ[t],
                                                       R[t], ηG[t], ηa[t], ηL[t], ηR[t])

        # Compute real wage
        w[t] = exp(ηL[t])*(L[t]^m.p.ϑ)*(C[t]^m.p.γ)

        # If ZLB is imposed, set R(t)=1 if ZLB binds
        if m.p.zlb; R[t+1] = max(R[t+1],1.0); end
    end

    Simulation(ηR, ηa, ηL, ηu, ηB, ηG, δ, R, S, F, C, π, Y, L, Yn, w)
end

#%% cell
#=
### Accuracy
=#

#%% cell
struct Residuals
    resids::Matrix{Float64}
end

function Residuals(m::Model, coefs::Matrix, sim::Simulation; burn::Int=200)
    capT = length(sim.w)
    resids = zeros(9, capT)

    # Integration method for evaluating accuracy
    # ------------------------------------------
    # Monomial integration rule with 2N^2+1 nodes
    ϵ_nodes, ω_nodes = qnwmonomial2(vcov(m.p))
    n_nodes = length(ω_nodes)

    # Allocate for arrays needed in the loop
    basis_mat = Array{Float64}(undef, n_nodes, 8)
    X1 = Array{Float64}(undef, n_nodes, size(coefs, 1))

    ηR1 = Array{Float64}(undef, n_nodes)
    ηa1 = Array{Float64}(undef, n_nodes)
    ηL1 = Array{Float64}(undef, n_nodes)
    ηu1 = Array{Float64}(undef, n_nodes)
    ηB1 = Array{Float64}(undef, n_nodes)
    ηG1 = Array{Float64}(undef, n_nodes)

    for t in 1:capT                 # For each given point,
        # Take the corresponding value for shocks at t
        #---------------------------------------------
        ηR0 = sim.ηR[t]  # ηR(t)
        ηa0 = sim.ηa[t]  # ηa(t)
        ηL0 = sim.ηL[t]  # ηL(t)
        ηu0 = sim.ηu[t]  # ηu(t)
        ηB0 = sim.ηB[t]  # ηB(t)
        ηG0 = sim.ηG[t]  # ηG(t)

        # Exctract time t values for all other variables (and t+1 for R, δ)
        #------------------------------------------------------------------
        R0  = sim.R[t]   # R(t-1)
        δ0  = sim.δ[t]   # δ(t-1)
        R1  = sim.R[t+1] # R(t)
        δ1  = sim.δ[t+1] # δ(t)

        L0  = sim.L[t]   # L(t)
        Y0  = sim.Y[t]   # Y(t)
        Yn0 = sim.Yn[t]  # Yn(t)
        π0  = sim.π[t]   # π(t)
        S0  = sim.S[t]   # S(t)
        F0  = sim.F[t]   # F(t)
        C0  = sim.C[t]   # C(t)

        # Fill basis matrix with R1, δ1 and shocks
        #-----------------------------------------
        # Note that we do not premultiply by standard deviations as ϵ_nodes
        # already include them. All these variables are vectors of length n_nodes
        copy!(ηR1, ηR0*m.p.ρηR .+ ϵ_nodes[:, 1])
        copy!(ηa1, ηa0*m.p.ρηa .+ ϵ_nodes[:, 2])
        copy!(ηL1, ηL0*m.p.ρηL .+ ϵ_nodes[:, 3])
        copy!(ηu1, ηu0*m.p.ρηu .+ ϵ_nodes[:, 4])
        copy!(ηB1, ηB0*m.p.ρηB .+ ϵ_nodes[:, 5])
        copy!(ηG1, ηG0*m.p.ρηG .+ ϵ_nodes[:, 6])

        basis_mat[:, 1] .= log.(R1)
        basis_mat[:, 2] .= log.(δ1)
        basis_mat[:, 3] = ηR1
        basis_mat[:, 4] = ηa1
        basis_mat[:, 5] = ηL1
        basis_mat[:, 6] = ηu1
        basis_mat[:, 7] = ηB1
        basis_mat[:, 8] = ηG1

        # Future choices at t+1
        #----------------------
        # Form a complete polynomial of degree "Degree" (at t+1) on future state
        # variables; n_nodes-by-npol
        complete_polynomial!(X1, basis_mat, m.p.deg)

        # Compute S(t+1), F(t+1) and C(t+1) in all nodes using coefs
        S1 = X1*coefs[:, 1]
        F1 = X1*coefs[:, 2]
        C1 = (X1*coefs[:, 3]).^(-1/m.p.γ)

        # Compute π(t+1) using condition (35) in MM (2015)
        π1 = @. ((1-(1-m.p.Θ)*(S1./F1).^(1-m.p.ϵ))/m.p.Θ).^(1/(m.p.ϵ-1))

        # Compute residuals for each of the 9 equilibrium conditions
        #-----------------------------------------------------------
        resids[1, t] = 1-dot(ω_nodes,
            @. (exp(ηu0)*exp(ηL0)*L0^m.p.ϑ*Y0/exp(ηa0) + m.p.β*m.p.Θ*π1.^m.p.ϵ.*S1)/S0
        )
        resids[2, t] = 1 - dot(ω_nodes,
            @. (exp(ηu0)*C0^(-m.p.γ)*Y0 + m.p.β*m.p.Θ*π1.^(m.p.ϵ-1).*F1)/F0
        )
        resids[3, t] = 1.0 -dot(ω_nodes,
            @. (m.p.β*exp(ηB0)/exp(ηu0)*R1*exp.(ηu1).*C1.^(-m.p.γ)./π1)/C0^(-m.p.γ)
        )
        resids[4, t] = 1-((1-m.p.Θ*π0^(m.p.ϵ-1))/(1-m.p.Θ))^(1/(1-m.p.ϵ))*F0/S0
        resids[5, t] = 1-((1-m.p.Θ)*((1-m.p.Θ*π0^(m.p.ϵ-1))/(1-m.p.Θ))^(m.p.ϵ/(m.p.ϵ-1)) + m.p.Θ*π0^m.p.ϵ/δ0)^(-1)/δ1
        resids[6, t] = 1-exp(ηa0)*L0*δ1/Y0
        resids[7, t] = 1-(1-m.p.gbar/exp(ηG0))*Y0/C0
        resids[8, t] = 1-(exp(ηa0)^(1+m.p.ϑ)*(1-m.p.gbar/exp(ηG0))^(-m.p.γ)/exp(ηL0))^(1/(m.p.ϑ+m.p.γ))/Yn0
        resids[9, t] = 1-m.s.π/m.p.β*(R0*m.p.β/m.s.π)^m.p.μ*((π0/m.s.π)^m.p.ϕ_π * (Y0/Yn0)^m.p.ϕ_y)^(1-m.p.μ)*exp(ηR0)/R1   # Taylor rule

        # If the ZLB is imposed and R>1, the residuals in the Taylor rule (the
        # 9th equation) are zero
        if m.p.zlb && R1 <= 1; resids[9, t] = 0.0; end

    end
    # discard the first burn observations
    Residuals(resids[:, burn+1:end])
end

Statistics.mean(r::Residuals) = log10(mean(abs, r.resids))
Base.max(r::Residuals) = log10(maximum(abs, r.resids))
max_E(r::Residuals) = log10.(maximum(abs, r.resids, dims=2))[:]


# P(n) = (1/2)*n*(sin(2*pi/n))

P(n) = pi

function c(n = BigFloat)
    if(n==2)
        return 0
    else
        return BigFloat(sqrt((1/2)*(1+c(n-1))))
    end
end

function s(n = BigFloat)
    if(n==2)
        return 1
    else
        return BigFloat(BigFloat(sqrt((1/2)*(1-c(n-1)))))    end
end

function s2(n = BigFloat)
    if (n==2)
        return 1
    else
        return (BigFloat(s2(2*n)))/(2*(BigFloat(c(n))))
    end
end

A(n = BigFloat) =  BigFloat((2^(n-1))*s(n))

println("ZADANIE 5")

setprecision(128) do
    for i in 2:128
        print(i)
        print("  ")
        print(P(2^i))
        print("  ")
        print(A(i))
        print("  ")
        println(abs(P(2^i) - A(i)))
    end
end

II rok/I semestr/ANM/Cwiczenia 02/Zadanie5.jl


using JuAFEM
using JuAFEM:Vec
using Tensors
using SparseArrays
using LinearAlgebra
include("../src/FEMRefinements.jl")
using .FEMRefinements

const ∇ = Tensors.gradient
const Δ = Tensors.hessian;

v1 = Vec((2.0, 0.0))
v2 = Vec((0.0, 2.0))
grid = generate_grid2D(40, 40, v1, v2)

dim = 2
ip = Lagrange{dim,RefCube,1}()
qr = QuadratureRule{dim,RefCube}(2)
qr_face = QuadratureRule{dim - 1,RefCube}(2)
cellvalues = CellScalarValues(qr, ip);
facevalues = FaceScalarValues(qr_face, ip);

dh = DofHandler(grid)
push!(dh, :u, 1)
close!(dh)

function u_ana(x::Vec{2,T}) where {T}
    xs = (Vec{2}((0.5,  1.5)),
          Vec{2}((1.5, 1.5)),
          Vec{2}((0.5,  0.5)))
    σ = 1 / 8
    s = zero(eltype(x))
    for i in 1:3
        s += exp(- norm(x - xs[i])^2 / σ^2)
    end
    return max(1e-15 * one(T), s) # Denormals, be gone
end;

dbcs = ConstraintHandler(dh)
dbc = Dirichlet(:u, union(getfaceset(grid, "top"), getfaceset(grid, "right")), (x, t) -> u_ana(x))
add!(dbcs, dbc)
close!(dbcs)
update!(dbcs, 0.0)

K = create_sparsity_pattern(dh);

function doassemble(cellvalues::CellScalarValues{dim}, facevalues::FaceScalarValues{dim},
                         K::SparseMatrixCSC, dh::DofHandler) where {dim}
    b = 1.0
    f = zeros(ndofs(dh))
    assembler = start_assemble(K, f)

    n_basefuncs = getnbasefunctions(cellvalues)
    global_dofs = zeros(Int, ndofs_per_cell(dh))

    fe = zeros(n_basefuncs) # Local force vector
    Ke = zeros(n_basefuncs, n_basefuncs) # Local stiffness mastrix

    @inbounds for (cellcount, cell) in enumerate(CellIterator(dh))
        fill!(Ke, 0)
        fill!(fe, 0)
        coords = getcoordinates(cell)

        reinit!(cellvalues, cell)
        for q_point in 1:getnquadpoints(cellvalues)
            dΩ = getdetJdV(cellvalues, q_point)
            coords_qp = spatial_coordinate(cellvalues, q_point, coords)
            f_true = -LinearAlgebra.tr(hessian(u_ana, coords_qp)) + u_ana(coords_qp)
            for i in 1:n_basefuncs
                δu = shape_value(cellvalues, q_point, i)
                ∇δu = shape_gradient(cellvalues, q_point, i)
                fe[i] += (δu * f_true) * dΩ
                for j in 1:n_basefuncs
                    u = shape_value(cellvalues, q_point, j)
                    ∇u = shape_gradient(cellvalues, q_point, j)
                    Ke[i, j] += (∇δu ⋅ ∇u + δu * u) * dΩ
                end
            end
        end
        for face in 1:nfaces(cell)
            if (cellcount, face) ∈ getfaceset(grid, "left") ||
                    (cellcount, face) ∈ getfaceset(grid, "bottom")
                reinit!(facevalues, cell, face)
                for q_point in 1:getnquadpoints(facevalues)
                    coords_qp = spatial_coordinate(facevalues, q_point, coords)
                    n = getnormal(facevalues, q_point)
                    g = gradient(u_ana, coords_qp) ⋅ n
                    dΓ = getdetJdV(facevalues, q_point)
                    for i in 1:n_basefuncs
                        δu = shape_value(facevalues, q_point, i)
                        fe[i] += -(δu * g) * dΓ
                        for j in 1:n_basefuncs
                            ∇u = shape_gradient(cellvalues, q_point, j)
                            Ke[i, j] += (δu * ∇u ⋅ n) * dΓ
                        end
                    end
                end
            end
        end

        celldofs!(global_dofs, cell)
        assemble!(assembler, global_dofs, fe, Ke)
    end
    return K, f
end;

K, f = doassemble(cellvalues, facevalues, K, dh);
apply!(K, f, dbcs)
u = Symmetric(K) \ f;

AdaptiveMesh.plot_solution3D(dh, u; plotgrid=true)


using ThreeBodyDecay

let
    ev = LinRange(2.0, 10.0, 100)
    cal = iRhoQTB.(ev.^2, 3.0, 1.0, 0.2, 1.0)
    # plot(ev, [real.(cal) imag(cal)], lab="")
end

# iRhoQTB(5.0^2, 3, 1, 0.2, 2)


#
# function cholesky_U_deravitive(A::Symmetric{T,MMatrix{N,N,T,N2}}, Ȧ::Symmetric{T,MMatrix{N,N,T,N2}}) where {N,T,N2}
#     # Consider the factorization of a PD matrix A: A = U' * U
#     # The derivative of this factorization is Ȧ = U̇' * U + U' * U̇
#     # This function finds U̇ given Ȧ using the algorithm from: https://homepages.inf.ed.ac.uk/imurray2/pub/16choldiff/choldiff.pdf
#
#     cf = cholesky(A)
#     U = cf.U
#     U⁻¹ = inv(U)
#     term_in = U⁻¹' * Ȧ * U⁻¹
#     @inbounds begin
#         for k = 1:(N+1):N2
#             term_in[k] *= 0.5
#         end
#     end
#     U̇ = UpperTriangular(term_in) * U
#     return U, U⁻¹, U̇
# end


h = 100
v = 20
g = 9.81

println((h + v^2/2g)^2 * pi * v^2/g)


# The Hausdorff measure generalizes the notion of volume or area to a manifold.
# For manifolds embedded in real space, it can be intuitively understoof as the addition
# of an infinitesimal dimension and the limit of the volume in the embedded space as the
# infinitesimal dimension goes to 0.
# A more useful description comes from the area formula. Given an m-
# dimensional manifold M and a Lipschitz map f: ℝᵐ → M ⊆ ℝⁿ that embeds M in ℝⁿ, then the Hausdorff
# measure dμ(x) on M is is the pushforward of the Lebesgue measure through f, given as
# dμ(x) = det(G(x))^(-1/2) dλ(f⁻¹(x)), where G(x) is the matrix representation of the metric tensor
# in M, given by G(x) = J(x)ᵀ J(x), where J is the Jacobian of f.
# This is the un-normalized Hausdorff measure, which integrates to the area/volume of M.

"""
    Hausdorff{M,A} <: PrimitiveMeasure

The un-normalized Hausdorff measure on a manifold.

The Hausdorff measure generalizes the notion of area or volume to a manifold that is embedded
in a metric space. That is, the mass of the measure over some region of the manifold is the
area/volume of that region in the embedded space.

# Constructors

    Hausdorff(M::AbstractManifold)

Constructs the Hausdorff measure for the manifold `M` using the default embedding of the
manifold.
"""
struct Hausdorff{M,A} <: PrimitiveMeasure
    manifold::M
    atlas::A  # optional, when the Hausdorff measure is constructed from a Lebesgue measure
end
Hausdorff(M) = Hausdorff(M, nothing)

function Base.show(io::IO, mime::MIME"text/plain", μ::Hausdorff)
    print(io, "Hausdorff")
    return show(io, mime, (μ.manifold, μ.atlas))
end

Manifolds.base_manifold(μ::Hausdorff) = μ.manifold

MeasureTheory.logdensity(::Hausdorff, x) = zero(eltype(x))

function Base.rand(rng::AbstractRNG, T::Type, μ::Normalized{<:Hausdorff})
    p = default_point(μ, T)
    return Random.rand!(rng, p, μ)
end

# Stiefel

# In general, given matrix z ∈ 𝔽^(n × k) with IID std normal elements, p=z(z'z)^(-1/2) is drawn
# from the Hausdorff measure on St(n,k,𝔽).
# This implementation uses the unique QR decomposition z=QR where R[i,i] is real and positive
# for all i, so that Q ~ Hausdorff(St(n,k,𝔽)).
# See Theorem 2.3.19 of Gupta AK, Nagar DK. Matrix variate distributions. CRC Press; 2018
function Random.rand!(
    rng::AbstractRNG, p::AbstractMatrix, ::Normalized{<:Hausdorff{<:Stiefel}}
)
    randn!(rng, p)
    Q, _ = qr_unique!(p)
    return Q
end

# Chikuse, 2003 Eq. 1.4.8
function logmass(::Hausdorff{Stiefel{n,k,ℝ}}) where {n,k}
    halfn = n//2
    return k * logtwo + (k * halfn) * logπ - logmvgamma(k, halfn)
end
# vol(St(n,k,𝔽)) = vol(𝔽𝕊ⁿ⁻¹) * vol(St(n-1,k-1,𝔽)), Chikuse, 2003 §1.4.4
# then, vol(St(n,k,𝔽)) = 2ᵏ π^(k(2n-k+1)d/4) / ∏ᵢ₌₁ᵏΓ(d(n-k+i)/2)),
# where d = real_dimension(𝔽)
function logmass(::Hausdorff{Stiefel{n,k,ℂ}}) where {n,k}
    a = n - k
    lfact = logfactorial(a)
    lm = k * logtwo + (k * (2n - k + 1)//2) * logπ - lfact
    for _ in 1:(k - 1)
        a += 1
        lfact += log(a)
        lm -= lfact
    end
    return lm
end
function logmass(::Hausdorff{Stiefel{n,k,ℍ}}) where {n,k}
    a = 2(n - k) + 1
    lfact = logfactorial(a)
    lm = k * logtwo + (k * (a + k)) * logπ - lfact
    for _ in 1:(k - 1)
        a += 2
        lfact += log(a) + log(a - 1)
        lm -= lfact
    end
    return lm
end

# Grassmann

function Random.rand!(
    rng::AbstractRNG, p::AbstractMatrix, ::Normalized{<:Hausdorff{Grassmann{n,k,𝔽}}}
) where {n,k,𝔽}
    return rand!(rng, p, normalize(Hausdorff(Stiefel(n, k, 𝔽))))
end

# because Gr(n,k,𝔽) = St(n,k,𝔽) / U(k,𝔽),
# then vol(Gr(n,k,𝔽)) = vol(St(n,k,𝔽)) / vol(St(k,k,𝔽))
#                     = π^(k((n-k)d/2)) ∏ᵢ₌₁ᵏ Γ(id/2) / Γ((n-k)d/2 + id/2),
# where d = real_dimension(𝔽)
function logmass(::Hausdorff{Grassmann{n,k,ℝ}}) where {n,k}
    return k * (n - k)//2 * logπ + logmvgamma(k, k//2) - logmvgamma(k, n//2)
end
function logmass(::Hausdorff{Grassmann{n,k,ℂ}}) where {n,k}
    a = n - k
    lfact1 = logfactorial(a)
    lfact2 = log(1)
    lm = k * a * logπ + lfact2 - lfact1
    for i in 1:(k - 1)
        a += 1
        lfact1 += log(a)
        lfact2 += log(i)
        lm += lfact2 - lfact1
    end
    return lm
end
function logmass(::Hausdorff{Grassmann{n,k,ℍ}}) where {n,k}
    a = 2(n - k) + 1
    b = 1
    lfact1 = logfactorial(a)
    lfact2 = log(b)
    lm = 2 * k * (n - k) * logπ + lfact2 - lfact1
    for _ in 1:(k - 1)
        a += 2
        b += 2
        lfact1 += log(a) + log(a - 1)
        lfact2 += log(b) + log(b - 1)
        lm += lfact2 - lfact1
    end
    return lm
end

# Circle

default_point(::Normalized{<:Hausdorff{Circle{ℝ}}}, T::Type) = zero(float(real(T)))
default_point(::Normalized{<:Hausdorff{Circle{ℂ}}}, T::Type) = zero(complex(float(real(T))))

function Random.rand!(rng::AbstractRNG, p::Real, ::Normalized{<:Hausdorff{Circle{ℝ}}})
    return rand(rng, typeof(p)) * twoπ - π
end
function Random.rand!(rng::AbstractRNG, p::Complex, ::Normalized{<:Hausdorff{Circle{ℂ}}})
    return sign(randn(rng, typeof(p)))
end

logmass(::Hausdorff{<:Circle}) = log2π

# Sphere

function Random.rand!(
    rng::AbstractRNG, p::AbstractArray, μ::Normalized{<:Hausdorff{<:AbstractSphere}}
)
    return normalize!(randn!(rng, p))
end

function logmass(μ::Hausdorff{<:AbstractSphere{𝔽}}) where {𝔽}
    m = prod(representation_size(base_manifold(μ)))
    ν = real_dimension(𝔽) * m//2
    return logtwo + ν * logπ - loggamma(ν)
end

# ProjectiveSpace

function Random.rand!(
    rng::AbstractRNG, p::AbstractArray, μ::Normalized{<:Hausdorff{<:AbstractProjectiveSpace}}
)
    return normalize!(randn!(rng, p))
end

# because 𝔽ℙⁿ = 𝔽𝕊ⁿ / 𝔽𝕊⁰, then vol(𝔽ℙⁿ) = vol(𝔽𝕊ⁿ) / vol(𝔽𝕊⁰)
function logmass(μ::Hausdorff{<:AbstractProjectiveSpace{ℝ}})
    m = prod(representation_size(base_manifold(μ)))
    ν = m//2
    return ν * logπ - loggamma(ν)
end
function logmass(μ::Hausdorff{<:AbstractProjectiveSpace{ℂ}})
    m = prod(representation_size(base_manifold(μ)))
    n = m - 1
    return n * logπ - logfactorial(n)
end
function logmass(μ::Hausdorff{<:AbstractProjectiveSpace{ℍ}})
    m = prod(representation_size(base_manifold(μ)))
    n = m - 1
    return 2n * logπ - loggamma(2(n + 1))
end

# Rotations/SpecialOrthogonal

function Random.rand!(
    rng::AbstractRNG,
    p::AbstractMatrix,
    ::Normalized{<:Hausdorff{<:Union{Rotations,SpecialOrthogonal}}},
)
    # draw p ~ O(n)
    rand!(rng, p, Hausdorff(Stiefel(n, n)))
    # fix to p ~ SO(n)
    if det(p) < 0
        @views p[:, [1, 2]] = p[:, [2, 1]]
    end
    return p
end

# since O(n) = St(n,n) is essentially two copies of SO(n), then vol(SO(n)) = vol(St(n,n)) / 2
function logmass(::Hausdorff{<:Union{Rotations{n},SpecialOrthogonal{n}}}) where {n}
    return logmass(Hausdorff(Stiefel(n, n))) - logtwo
end


using KernelFunctions, Test, Distances
#TODO Test metric weights for ARD, test equivalency for different constructors,
# test type conversion
l = 2.0
vl = [l,l]
s = ScaleTransform(l)

## SqExponentialKernel
@testset "SqExponentialKernel" begin
    @test KernelFunctions.metric(SqExponentialKernel(l)) == SqEuclidean()
    @test isequal(transform(SqExponentialKernel(l)),s)
    @test KernelFunctions.transform(SqExponentialKernel(vl)) == ARDTransform(vl)
    @test isequal(KernelFunctions.transform(SqExponentialKernel(s)),s)
end

## MaternKernel
@testset "MaternKernel" begin
    @test KernelFunctions.metric(MaternKernel(l)) == Euclidean()
    @test KernelFunctions.metric(Matern32Kernel(l)) == Euclidean()
    @test KernelFunctions.metric(Matern52Kernel(l)) == Euclidean()
    @test isequal(KernelFunctions.transform(MaternKernel(l)),s)
    @test isequal(KernelFunctions.transform(Matern32Kernel(l)),s)
    @test isequal(KernelFunctions.transform(Matern52Kernel(l)),s)
    @test KernelFunctions.transform(MaternKernel(vl)) == ARDTransform(vl)
    @test KernelFunctions.transform(Matern32Kernel(vl)) == ARDTransform(vl)
    @test KernelFunctions.transform(Matern52Kernel(vl)) == ARDTransform(vl)
    @test KernelFunctions.transform(MaternKernel(s)) == s
    @test KernelFunctions.transform(Matern32Kernel(s)) == s
    @test KernelFunctions.transform(Matern52Kernel(s)) == s
end