@doc doc"""
    LPDistLoss{P} <: DistanceLoss

The P-th power absolute distance loss. It is Lipschitz continuous
iff `P == 1`, convex if and only if `P >= 1`, and strictly convex
iff `P > 1`.

```math
L(r) = |r|^P
```
"""
struct LPDistLoss{P} <: DistanceLoss end

LPDistLoss(p::Number) = LPDistLoss{p}()

value(loss::LPDistLoss{P}, difference::Number) where {P} = abs(difference)^P
function deriv(loss::LPDistLoss{P}, difference::T)::promote_type(typeof(P),T) where {P,T<:Number}
    if difference == 0
        zero(difference)
    else
        P * difference * abs(difference)^(P-convert(typeof(P), 2))
    end
end
function deriv2(loss::LPDistLoss{P}, difference::T)::promote_type(typeof(P),T) where {P,T<:Number}
    if difference == 0
        zero(difference)
    else
        (abs2(P)-P) * abs(difference)^P / abs2(difference)
    end
end

isminimizable(::LPDistLoss{P}) where {P} = true
issymmetric(::LPDistLoss{P}) where {P} = true
isdifferentiable(::LPDistLoss{P}) where {P} = P > 1
isdifferentiable(::LPDistLoss{P}, at) where {P} = P > 1 || at != 0
istwicedifferentiable(::LPDistLoss{P}) where {P} = P > 1
istwicedifferentiable(::LPDistLoss{P}, at) where {P} = P > 1 || at != 0
islipschitzcont(::LPDistLoss{P}) where {P} = P == 1
islocallylipschitzcont(::LPDistLoss{P}) where {P} = P >= 1
isconvex(::LPDistLoss{P}) where {P} = P >= 1
isstrictlyconvex(::LPDistLoss{P}) where {P} = P > 1
isstronglyconvex(::LPDistLoss{P}) where {P} = P >= 2

# ===========================================================

@doc doc"""
    L1DistLoss <: DistanceLoss

The absolute distance loss.
Special case of the [`LPDistLoss`](@ref) with `P=1`.
It is Lipschitz continuous and convex, but not strictly convex.

```math
L(r) = |r|
```

---
```
              Lossfunction                     Derivative
      ┌────────────┬────────────┐      ┌────────────┬────────────┐
    3 │\.                     ./│    1 │            ┌------------│
      │ '\.                 ./' │      │            |            │
      │   \.               ./   │      │            |            │
      │    '\.           ./'    │      │_           |           _│
    L │      \.         ./      │   L' │            |            │
      │       '\.     ./'       │      │            |            │
      │         \.   ./         │      │            |            │
    0 │          '\./'          │   -1 │------------┘            │
      └────────────┴────────────┘      └────────────┴────────────┘
      -3                        3      -3                        3
                 ŷ - y                            ŷ - y
```
"""
const L1DistLoss = LPDistLoss{1}

value(loss::L1DistLoss, difference::Number) = abs(difference)
deriv(loss::L1DistLoss, difference::T) where {T<:Number} = convert(T, sign(difference))
deriv2(loss::L1DistLoss, difference::T) where {T<:Number} = zero(T)

isdifferentiable(::L1DistLoss) = false
isdifferentiable(::L1DistLoss, at) = at != 0
istwicedifferentiable(::L1DistLoss) = false
istwicedifferentiable(::L1DistLoss, at) = at != 0
islipschitzcont(::L1DistLoss) = true
isconvex(::L1DistLoss) = true
isstrictlyconvex(::L1DistLoss) = false
isstronglyconvex(::L1DistLoss) = false

# ===========================================================

@doc doc"""
    L2DistLoss <: DistanceLoss

The least squares loss.
Special case of the [`LPDistLoss`](@ref) with `P=2`.
It is strictly convex.

```math
L(r) = |r|^2
```

---
```
              Lossfunction                     Derivative
      ┌────────────┬────────────┐      ┌────────────┬────────────┐
    9 │\                       /│    3 │                   .r/   │
      │".                     ."│      │                 .r'     │
      │ ".                   ." │      │              _./'       │
      │  ".                 ."  │      │_           .r/         _│
    L │   ".               ."   │   L' │         _:/'            │
      │    '\.           ./'    │      │       .r'               │
      │      \.         ./      │      │     .r'                 │
    0 │        "-.___.-"        │   -3 │  _/r'                   │
      └────────────┴────────────┘      └────────────┴────────────┘
      -3                        3      -2                        2
                 ŷ - y                            ŷ - y
```
"""
const L2DistLoss = LPDistLoss{2}

value(loss::L2DistLoss, difference::Number) = abs2(difference)
deriv(loss::L2DistLoss, difference::T) where {T<:Number} = convert(T,2) * difference
deriv2(loss::L2DistLoss, difference::T) where {T<:Number} = convert(T,2)

isdifferentiable(::L2DistLoss) = true
isdifferentiable(::L2DistLoss, at) = true
istwicedifferentiable(::L2DistLoss) = true
istwicedifferentiable(::L2DistLoss, at) = true
islipschitzcont(::L2DistLoss) = false
isconvex(::L2DistLoss) = true
isstrictlyconvex(::L2DistLoss) = true
isstronglyconvex(::L2DistLoss) = true

# ===========================================================

@doc doc"""
    PeriodicLoss <: DistanceLoss

Measures distance on a circle of specified circumference `c`.

```math
L(r) = 1 - \cos \left( \frac{2 r \pi}{c} \right)
```
"""
struct PeriodicLoss{T<:AbstractFloat} <: DistanceLoss
    k::T   # k = 2π/circumference
    function PeriodicLoss{T}(circ::T) where T
        circ > 0 || error("circumference should be strictly positive")
        new{T}(convert(T, 2π/circ))
    end
end
PeriodicLoss(circ::T=1.0) where {T<:AbstractFloat} = PeriodicLoss{T}(circ)
PeriodicLoss(circ) = PeriodicLoss{Float64}(Float64(circ))

value(loss::PeriodicLoss, difference::T) where {T<:Number} = 1 - cos(difference*loss.k)
deriv(loss::PeriodicLoss, difference::T) where {T<:Number} = loss.k * sin(difference*loss.k)
deriv2(loss::PeriodicLoss, difference::T) where {T<:Number} = abs2(loss.k) * cos(difference*loss.k)

isdifferentiable(::PeriodicLoss) = true
isdifferentiable(::PeriodicLoss, at) = true
istwicedifferentiable(::PeriodicLoss) = true
istwicedifferentiable(::PeriodicLoss, at) = true
islipschitzcont(::PeriodicLoss) = true
isconvex(::PeriodicLoss) = false
isstrictlyconvex(::PeriodicLoss) = false
isstronglyconvex(::PeriodicLoss) = false

# ===========================================================

@doc doc"""
    HuberLoss <: DistanceLoss

Loss function commonly used for robustness to outliers.
For large values of `d` it becomes close to the [`L1DistLoss`](@ref),
while for small values of `d` it resembles the [`L2DistLoss`](@ref).
It is Lipschitz continuous and convex, but not strictly convex.

```math
L(r) = \begin{cases} \frac{r^2}{2} & \quad \text{if } | r | \le \alpha \\ \alpha | r | - \frac{\alpha^3}{2} & \quad \text{otherwise}\\ \end{cases}
```

---
```
              Lossfunction (d=1)               Derivative
      ┌────────────┬────────────┐      ┌────────────┬────────────┐
    2 │                         │    1 │                .+-------│
      │                         │      │              ./'        │
      │\.                     ./│      │             ./          │
      │ '.                   .' │      │_           ./          _│
    L │   \.               ./   │   L' │           /'            │
      │     \.           ./     │      │          /'             │
      │      '.         .'      │      │        ./'              │
    0 │        '-.___.-'        │   -1 │-------+'                │
      └────────────┴────────────┘      └────────────┴────────────┘
      -2                        2      -2                        2
                 ŷ - y                            ŷ - y
```
"""
struct HuberLoss{T<:AbstractFloat} <: DistanceLoss
    d::T   # boundary between quadratic and linear loss
    function HuberLoss{T}(d::T) where T
        d > 0 || error("Huber crossover parameter must be strictly positive.")
        new{T}(d)
    end
end
HuberLoss(d::T=1.0) where {T<:AbstractFloat} = HuberLoss{T}(d)
HuberLoss(d) = HuberLoss{Float64}(Float64(d))

function value(loss::HuberLoss{T1}, difference::T2) where {T1,T2<:Number}
    T = promote_type(T1,T2)
    abs_diff = abs(difference)
    if abs_diff <= loss.d
        return convert(T,0.5)*abs2(difference)   # quadratic
    else
        return (loss.d*abs_diff) - convert(T,0.5)*abs2(loss.d)   # linear
    end
end
function deriv(loss::HuberLoss{T1}, difference::T2) where {T1,T2<:Number}
    T = promote_type(T1,T2)
    if abs(difference) <= loss.d
        return convert(T,difference)   # quadratic
    else
        return loss.d*convert(T,sign(difference))   # linear
    end
end
function deriv2(loss::HuberLoss{T1}, difference::T2) where {T1,T2<:Number}
    T = promote_type(T1,T2)
    abs(difference) <= loss.d ? one(T) : zero(T)
end

isdifferentiable(::HuberLoss) = true
isdifferentiable(l::HuberLoss, at) = true
istwicedifferentiable(::HuberLoss) = false
istwicedifferentiable(l::HuberLoss, at) = at != abs(l.d)
islipschitzcont(::HuberLoss) = true
isconvex(::HuberLoss) = true
isstrictlyconvex(::HuberLoss) = false
isstronglyconvex(::HuberLoss) = false
issymmetric(::HuberLoss) = true

# ===========================================================

@doc doc"""
    L1EpsilonInsLoss <: DistanceLoss

The ``ϵ``-insensitive loss. Typically used in linear support vector
regression. It ignores deviances smaller than ``ϵ``, but penalizes
larger deviances linarily.
It is Lipschitz continuous and convex, but not strictly convex.

```math
L(r) = \max \{ 0, | r | - \epsilon \}
```

---
```
              Lossfunction (ϵ=1)               Derivative
      ┌────────────┬────────────┐      ┌────────────┬────────────┐
    2 │\                       /│    1 │                  ┌------│
      │ \                     / │      │                  |      │
      │  \                   /  │      │                  |      │
      │   \                 /   │      │_      ___________!     _│
    L │    \               /    │   L' │      |                  │
      │     \             /     │      │      |                  │
      │      \           /      │      │      |                  │
    0 │       \_________/       │   -1 │------┘                  │
      └────────────┴────────────┘      └────────────┴────────────┘
      -3                        3      -2                        2
                 ŷ - y                            ŷ - y
```
"""
struct L1EpsilonInsLoss{T<:AbstractFloat} <: DistanceLoss
    ε::T

    function L1EpsilonInsLoss{T}(ɛ::T) where T
        ɛ > 0 || error("ɛ must be strictly positive")
        new{T}(ɛ)
    end
end
const EpsilonInsLoss = L1EpsilonInsLoss
@inline L1EpsilonInsLoss(ε::T) where {T<:AbstractFloat} = L1EpsilonInsLoss{T}(ε)
@inline L1EpsilonInsLoss(ε::Number) = L1EpsilonInsLoss{Float64}(Float64(ε))

function value(loss::L1EpsilonInsLoss{T1}, difference::T2) where {T1,T2<:Number}
    T = promote_type(T1,T2)
    max(zero(T), abs(difference) - loss.ε)
end
function deriv(loss::L1EpsilonInsLoss{T1}, difference::T2) where {T1,T2<:Number}
    T = promote_type(T1,T2)
    abs(difference) <= loss.ε ? zero(T) : convert(T,sign(difference))
end
deriv2(loss::L1EpsilonInsLoss{T1}, difference::T2) where {T1,T2<:Number} = zero(promote_type(T1,T2))

issymmetric(::L1EpsilonInsLoss) = true
isdifferentiable(::L1EpsilonInsLoss) = false
isdifferentiable(loss::L1EpsilonInsLoss, at) = abs(at) != loss.ε
istwicedifferentiable(::L1EpsilonInsLoss) = false
istwicedifferentiable(loss::L1EpsilonInsLoss, at) = abs(at) != loss.ε
islipschitzcont(::L1EpsilonInsLoss) = true
isconvex(::L1EpsilonInsLoss) = true
isstrictlyconvex(::L1EpsilonInsLoss) = false
isstronglyconvex(::L1EpsilonInsLoss) = false

# ===========================================================

@doc doc"""
    L2EpsilonInsLoss <: DistanceLoss

The quadratic ``ϵ``-insensitive loss.
Typically used in linear support vector regression.
It ignores deviances smaller than ``ϵ``, but penalizes
larger deviances quadratically. It is convex, but not strictly convex.

```math
L(r) = \max \{ 0, | r | - \epsilon \}^2
```

---
```
              Lossfunction (ϵ=0.5)             Derivative
      ┌────────────┬────────────┐      ┌────────────┬────────────┐
    8 │                         │    1 │                  /      │
      │:                       :│      │                 /       │
      │'.                     .'│      │                /        │
      │ \.                   ./ │      │_         _____/        _│
    L │  \.                 ./  │   L' │         /               │
      │   \.               ./   │      │        /                │
      │    '\.           ./'    │      │       /                 │
    0 │      '-._______.-'      │   -1 │      /                  │
      └────────────┴────────────┘      └────────────┴────────────┘
      -3                        3      -2                        2
                 ŷ - y                            ŷ - y
```
"""
struct L2EpsilonInsLoss{T<:AbstractFloat} <: DistanceLoss
    ε::T

    function L2EpsilonInsLoss{T}(ɛ::T) where T
        ɛ > 0 || error("ɛ must be strictly positive")
        new{T}(ɛ)
    end
end
L2EpsilonInsLoss(ε::T) where {T<:AbstractFloat} = L2EpsilonInsLoss{T}(ε)
L2EpsilonInsLoss(ε) = L2EpsilonInsLoss{Float64}(Float64(ε))

function value(loss::L2EpsilonInsLoss{T1}, difference::T2) where {T1,T2<:Number}
    T = promote_type(T1,T2)
    abs2(max(zero(T), abs(difference) - loss.ε))
end
function deriv(loss::L2EpsilonInsLoss{T1}, difference::T2) where {T1,T2<:Number}
    T = promote_type(T1,T2)
    absr = abs(difference)
    absr <= loss.ε ? zero(T) : convert(T,2)*sign(difference)*(absr - loss.ε)
end
function deriv2(loss::L2EpsilonInsLoss{T1}, difference::T2) where {T1,T2<:Number}
    T = promote_type(T1,T2)
    abs(difference) <= loss.ε ? zero(T) : convert(T,2)
end

issymmetric(::L2EpsilonInsLoss) = true
isdifferentiable(::L2EpsilonInsLoss) = true
isdifferentiable(::L2EpsilonInsLoss, at) = true
istwicedifferentiable(::L2EpsilonInsLoss) = false
istwicedifferentiable(loss::L2EpsilonInsLoss, at) = abs(at) != loss.ε
islipschitzcont(::L2EpsilonInsLoss) = false
isconvex(::L2EpsilonInsLoss) = true
isstrictlyconvex(::L2EpsilonInsLoss) = true
isstronglyconvex(::L2EpsilonInsLoss) = true

# ===========================================================

@doc doc"""
    LogitDistLoss <: DistanceLoss

The distance-based logistic loss for regression.
It is strictly convex and Lipschitz continuous.

```math
L(r) = - \ln \frac{4 e^r}{(1 + e^r)^2}
```

---
```
              Lossfunction                     Derivative
      ┌────────────┬────────────┐      ┌────────────┬────────────┐
    2 │                         │    1 │                   _--'''│
      │\                       /│      │                ./'      │
      │ \.                   ./ │      │              ./         │
      │  '.                 .'  │      │_           ./          _│
    L │   '.               .'   │   L' │           ./            │
      │     \.           ./     │      │         ./              │
      │      '.         .'      │      │       ./                │
    0 │        '-.___.-'        │   -1 │___.-''                  │
      └────────────┴────────────┘      └────────────┴────────────┘
      -3                        3      -4                        4
                 ŷ - y                            ŷ - y
```
"""
struct LogitDistLoss <: DistanceLoss end

function value(loss::LogitDistLoss, difference::Number)
    er = exp(difference)
    T = typeof(er)
    -log(convert(T,4)) - difference + 2log(one(T) + er)
end
function deriv(loss::LogitDistLoss, difference::T) where T<:Number
    tanh(difference / convert(T,2))
end
function deriv2(loss::LogitDistLoss, difference::Number)
    er = exp(difference)
    T = typeof(er)
    convert(T,2)*er / abs2(one(T) + er)
end

issymmetric(::LogitDistLoss) = true
isdifferentiable(::LogitDistLoss) = true
isdifferentiable(::LogitDistLoss, at) = true
istwicedifferentiable(::LogitDistLoss) = true
istwicedifferentiable(::LogitDistLoss, at) = true
islipschitzcont(::LogitDistLoss) = true
isconvex(::LogitDistLoss) = true
isstrictlyconvex(::LogitDistLoss) = true
isstronglyconvex(::LogitDistLoss) = false

# ===========================================================
@doc doc"""
    QuantileLoss <: DistanceLoss

The distance-based quantile loss, also known as pinball loss,
can be used to estimate conditional τ-quantiles.
It is Lipschitz continuous and convex, but not strictly convex.
Furthermore it is symmetric if and only if `τ = 1/2`.

```math
L(r) = \begin{cases} -\left( 1 - \tau  \right) r & \quad \text{if } r < 0 \\ \tau r & \quad \text{if } r \ge 0 \\ \end{cases}
```

---
```
              Lossfunction (τ=0.7)             Derivative
      ┌────────────┬────────────┐      ┌────────────┬────────────┐
    2 │'\                       │  0.3 │            ┌------------│
      │  \.                     │      │            |            │
      │   '\                    │      │_           |           _│
      │     \.                  │      │            |            │
    L │      '\              ._-│   L' │            |            │
      │        \.         ..-'  │      │            |            │
      │         '.     _r/'     │      │            |            │
    0 │           '_./'         │ -0.7 │------------┘            │
      └────────────┴────────────┘      └────────────┴────────────┘
      -3                        3      -3                        3
                 ŷ - y                            ŷ - y
```
"""
struct QuantileLoss{T <: AbstractFloat} <: DistanceLoss
    τ::T
end

function value(loss::QuantileLoss{T1}, diff::T2) where {T1, T2 <: Number}
    T = promote_type(T1, T2)
    diff * (convert(T,diff > 0) - loss.τ)
end
function deriv(loss::QuantileLoss{T1}, diff::T2) where {T1, T2 <: Number}
    T = promote_type(T1, T2)
    convert(T,diff > 0) - loss.τ
end
deriv2(::QuantileLoss{T1}, diff::T2) where {T1, T2 <: Number} = zero(promote_type(T1, T2))

issymmetric(loss::QuantileLoss) = loss.τ == 0.5
isdifferentiable(::QuantileLoss) = false
isdifferentiable(::QuantileLoss, at) = at != 0
istwicedifferentiable(::QuantileLoss) = false
istwicedifferentiable(::QuantileLoss, at) = at != 0
islipschitzcont(::QuantileLoss) = true
islipschitzcont_deriv(::QuantileLoss) = true
isconvex(::QuantileLoss) = true
isstrictlyconvex(::QuantileLoss) = false
isstronglyconvex(::QuantileLoss) = false

# ======================================================

@doc doc"""
    LogCoshLoss <: DistanceLoss

The log cosh loss is twice differentiable, strongly convex,
Lipschitz continous function.

```math
L(r) = log ( cosh ( x ))
```
---
```
           Lossfunction                     Derivative
      ┌────────────┬────────────┐      ┌────────────┬────────────┐
  2.5 │\                       /│    1 │                 .-------│
      │".                     ."│      │                |        │
      │ ".                   ." │      │               /         │
      │  ".                 ."  │      │_           . "         _│
    L │   ".               ."   │   L' │         /"              │
      │    '\.           ./'    │      │       ."                │
      │      \.         ./      │      │       |                 │
    0 │        "-. _ .-"        │   -1 │------"                  │
      └────────────┴────────────┘      └────────────┴────────────┘
      -3                        3      -3                        3
                 ŷ - y                            ŷ - y
```
"""

struct LogCoshLoss <: DistanceLoss end

_softplus(x::T) where T<:Number = x > zero(T) ? x + log1p(exp(-x)) : log1p(exp(x))
_log_cosh(x::T) where T<:Number = x + _softplus(-2x) - log(convert(T, 2))

function value(loss::LogCoshLoss, diff::T) where {T <: Number}
  _log_cosh(diff)
end

function deriv(loss::LogCoshLoss, diff::T) where {T <: Number}
  tanh.(diff)
end

function deriv2(::LogCoshLoss, diff::T) where {T <: Number}
  (sech.(diff))^2
end

issymmetric(loss::LogCoshLoss)           = true
isdifferentiable(::LogCoshLoss)          = true
isdifferentiable(::LogCoshLoss, at)      = true
istwicedifferentiable(::LogCoshLoss, at) = true
istwicedifferentiable(::LogCoshLoss)     = true
islipschitzcont(::LogCoshLoss)           = true
isconvex(::LogCoshLoss)                  = true
isstrictlyconvex(::LogCoshLoss)          = true
isstronglyconvex(::LogCoshLoss)          = true


@doc raw"""
```
binarize(Yen(), img)
```

Computes the binarization threshold value using Yen's maximum correlation criterion for
bilevel thresholding.

# Output

Returns the binarized image as an `Array{Gray{Bool},2}`.

# Details

This algorithm uses the concept of *entropic correlation* of a gray level histogram to produce a threshold
value.

Let ``f_1, f_2, \ldots, f_I`` be the frequencies in the various bins of the
histogram and ``I`` the number of bins. With ``N = \sum_{i=1}^{I}f_i``, let
``p_i = \frac{f_i}{N}`` (``i = 1, \ldots, I``) denote the probability
distribution of gray levels. From this distribution one derives two additional
distributions. The first defined for discrete values ``1`` to ``s`` and the
other, from ``s+1`` to ``I``. These distributions are

```math
A: \frac{p_1}{P_s}, \frac{p_2}{P_s}, \ldots, \frac{p_s}{P_s}
\quad \text{and} \quad
B: \frac{p_{s+1}}{1-P_s}, \ldots, \frac{p_n}{1-P_s}
\quad \text{where} \quad
P_s = \sum_{i=1}^{s}p_i.
```
The entropic correlations associated with each distribution are

```math
C(A) = -\ln \sum_{i=1}^{s} \left( \frac{p_i}{P_s} \right)^2 \quad \text{and} \quad C(B) = -\ln \sum_{i=s+1}^{I} \left( \frac{p_i}{1 - P_s} \right)^2.
```

Combining these two entropic correlation functions we have

```math
\psi(s) = -\ln \sum_{i=1}^{s} \left( \frac{p_i}{P_s} \right)^2 -\ln \sum_{i=s+1}^{I} \left( \frac{p_i}{1 - P_s} \right)^2.
```
Finding the discrete value ``s`` which maximises the function ``\psi(s)`` produces
the sought-after threshold value (i.e. the bin which determines the threshold).

# Arguments

The function argument is described in more detail below.

##  `img`

An `AbstractArray` representing an image. The image is automatically converted
to `Gray` in order to construct the requisite graylevel histogram.

# Example

Binarize the "cameraman" image in the `TestImages` package.

```julia
using TestImages, ImageBinarization

img = testimage("cameraman")
img_binary = binarize(Yen(), img)
```

# Reference

1. Yen JC, Chang FJ, Chang S (1995), “A New Criterion for Automatic Multilevel Thresholding”, IEEE Trans. on Image Processing 4 (3): 370-378, [doi:10.1109/83.366472](https://doi.org/10.1109/83.366472)
"""
function binarize(algorithm::Yen,  img::AbstractArray{T,2}) where T <: Colorant
  img₀₁ = zeros(Gray{Bool}, axes(img))
  edges, counts = build_histogram(img,  256)
  t = find_threshold(HistogramThresholding.Yen(), counts[1:end], edges)
  for i in CartesianIndices(img)
    img₀₁[i] = img[i] < t ? 0 : 1
  end
  img₀₁
end


@testset "CIC" begin

    nx = 64 # nombre de pts suivant x
    ny = 16   # nombre de pts suivant y

    alpha = 0.1
    kx = 0.5
    dimx = 2 * pi / kx
    dimy = 1

    mesh = TwoDGrid(dimx, nx, dimy, ny)

    dx = mesh.dx
    dy = mesh.dy

    rho = zeros(nx, ny)

    for i = 1:nx
        aux2 = alpha * cos(kx * mesh.x[i])
        for j = 1:ny
            rho[i, j] = aux2
        end
    end

    nbpart = 1000 * nx * ny

    p = ParticleGroup{2,2}(nbpart, charge = 1.0, mass = 1.0, n_weights = 1)

    sampler = LandauDamping(alpha, kx)

    ParticleInCell.sample!(p, mesh, sampler)

    kernel = CloudInCell()

    ρ = compute_rho(p, kernel, mesh)

    @test ρ ≈ rho atol=1e-2

end


using NURBS
using Base.Test

@testset "Order Increasing" begin

    @testset "From order 2 to 3" begin
        x = knot(4,2)
        b = [1. 2 3 4;1 2 2 1;1 1 1 1]
        @test typeof(raise23(b,x,4)) == Tuple{Array{Float64,2},Int64}
        @test raise23(b,x,4)[1] == [1 1.5 2 2.5 3 3.5 4; 1 1.5 2 2 2 1.5 1; 1 1 1 1 1 1 1]
        @test raise23(b,x,4)[2] == 7
        x = [0, 0., 1, 1, 2, 2]
        @test typeof(raise23(b,x,4)) == Tuple{Array{Float64,2},Int64}
        @test raise23(b,x,4)[1] == [1 1.5 2 3 3.5 4; 1 1.5 2 2 1.5 1; 1 1 1 1 1 1]
        @test raise23(b,x,4)[2] == 6
        x = knotu(4,2)
        @test typeof(raise23(b,x,4)) == Tuple{Array{Float64,2},Int64}
        @test raise23(b,x,4)[1] == [0.5 1 1.5 2 2.5 3 3.5 4 2;0.5 1 1.5 2 2 2 1.5 1 0.5;0.5 1 1 1 1 1 1 1 0.5]
        @test raise23(b,x,4)[2] == 9
    end

    @testset "From order 4 to 5" begin
        x = knot(4,4)
        b = [1. 2 3 4 ; 1 2 2 1 ; 1 1 1 1]
        @test typeof(raise45(b,x,4)) == Tuple{Array{Float64,2},Int64}
        @test raise45(b,x,4)[1] == [1 1.75 2.5 3.25 4; 1 1.75 2 1.75 1; 1 1 1 1 1]
        @test raise45(b,x,4)[2] == 5
        x = knot(5,4)
        b = [1. 2 3 4 5; 1 2 2 1 0; 1 1 1 1 1]
        @test typeof(raise45(b,x,5)) == Tuple{Array{Float64,2},Int64}
        @test raise45(b,x,5)[1] == [1 1.75 2.25 3 3.75 4.25 5; 1 1.75 2 1.875 1.25 0.75 0.0; 1 1 1 1 1 1 1]
        @test raise45(b,x,5)[2] == 7
        x = knot(6,4)
        b = [1. 2 3 4 5 6; 1 2 2 1 0 0; 1 1 1 1 1 1]
        @test typeof(raise45(b,x,6)) == Tuple{Array{Float64,2},Int64}
        @test isapprox(raise45(b,x,6)[1], [1 1.75 2.25 2.95833 3.5 4.04167 4.75 5.25 6; 1 1.75 2 1.91667 1.5 0.958333 0.25 0.0 0.0; 1 1 1 1 1 1 1 1 1], atol=5e-4)
        @test raise45(b,x,6)[2] == 9
        x = knot(7,4)
        b = [1. 2 3 4 5 6 7; 1 2 2 1 0 0 1; 1 1 1 1 1 1 1]
        @test typeof(raise45(b,x,7)) == Tuple{Array{Float64,2},Int64}
        @test isapprox(raise45(b,x,7)[1], [1 1.75 2.25 2.95833 3.5 4 4.5 5.04167 5.75 6.25 7; 1 1.75 2 1.91667 1.5 1 0.5 0.08333 0 0.25 1; 1 1 1 1 1 1 1 1 1 1 1]; atol=1e-5)
        @test raise45(b,x,7)[2] == 11
    end
end


module TestFunctions

struct Rosenbrock
    f
    ∇f
    argmin
    Rosenbrock(; a = 1, b = 5) =
    new(x->(a - x[1])^2 + b * (x[2] - x[1]^2)^2,
        x->[
            -2 * (a - x[1]) - 4 * b * x[1] * (x[2] - x[1]^2),
            2 * b * (x[2] - x[1]^2)
        ],
        [a, a^2])
end

end # module


###############################################################################
#
# Linear regression of Y~X, test slope
#
###############################################################################
"""
    test_slope(xs::Array,
               ys::Array;
               slope::Number = 0,
               intercept::Bool = true
    )

Make linear regression and return the p value of whether the regression slope
    differs from the given slope, given
- `xs` Array of x, can be NaN
- `ys` Array of y, can be NaN
- `slope` Slope to test
- `intercept` Optional: if true use intercept in the fitting
"""
function test_slope(
            xs::Array,
            ys::Array;
            slope::Number = 0,
            intercept::Bool = true
)
    # filter out NaN from the lists
    new_x = Float64[];
    new_y = Float64[];
    for (x,y) in zip(xs, ys)
        if ~isnan(x) && ~isnan(y)
            push!(new_x, x);
            push!(new_y, y - slope*x);
        end
    end

    # put new_x and new_y into a DataFrame
    df = DataFrame(X=new_x, Y=new_y);

    # run the fitting
    lr = linear_df_xy(df, intercept);

    # calculate the p value
    if intercept
        _slope_p = coeftable(lr).cols[4][2];
    else
        _slope_p = coeftable(lr).cols[4][1];
    end

    return _slope_p
end


using Plots

@fastmath function SIR_fs(N,bet,gamm)
    final_size = zeros(N+1,1);
    final_size[2] = 1;
    for Z2 in 0:N
        @inbounds for Z1 in Z2+1:N-1
            p1 = 1 / ( 1 + gamm/(bet*(N-Z1)));
            final_size[Z1+2] = final_size[Z1+2] + final_size[Z1+1]*p1;
            final_size[Z1+1] = final_size[Z1+1]*(1-p1);
        end
    end
    return final_size;
end

N = 20;
bet = 2/(N-1);
gamm = 1.0;

final_size = SIR_fs(N,bet,gamm);

@time final_size = SIR_fs(N,bet,gamm);

bar(0:N,final_size)


"""
    fourier_pade(u, degree_num, degree_den, num_output=length(u))

Compute the Fourier-Padé reconstruction of `u` with degrees
`(degree_num, degree_den)` and evaluate it at `num_output`
equispaced points, cf.
Driscoll and Fornberg (2001) A Padé-based algorithm for overcoming the Gibbs phenomenon,
doi: 10.1023/A:1016648530648.
"""
function fourier_pade(u, degree_num, degree_den, num_output=length(u))
  N = degree_num + degree_den
  modes = length(u) ÷ 2 + 1
  if modes <= N
    throw(ArgumentError("Cannot perform a Fourier-Padé reconstruction with degrees ($degree_num, $degree_den) given $(length(u)) real coefficients corresponding to $modes complex modes."))
  end

  uhat = rfft(u)
  uhat ./= length(u)

  # denominator
  col_den = uhat[degree_num+2:N+1]
  row_den = zeros(eltype(col_den), degree_den+1)
  row_den[1:min(degree_num+2,degree_den+1)] = uhat[degree_num+2:-1:max(1, degree_num+2-degree_den)]
  Z = nullspace(Matrix(Toeplitz(col_den, row_den)))
  den_p = zeros(eltype(uhat), num_output)
  den_p[1:degree_den+1] = Z[:, end]
  den_p ./= den_p[findfirst(!iszero, den_p)]

  # numerator
  col_num = uhat[1:degree_num+1]
  col_num[1] /= 2
  row_num = zeros(eltype(col_num), degree_den+1)
  row_num[1] = col_num[1]
  A = Toeplitz(col_num, row_num)
  num_p = zeros(eltype(uhat), num_output)
  mul!(view(num_p, 1:degree_num+1), A, view(den_p, 1:degree_den+1))

  # evaluate Fourier-Padé approximation
  bfft_plan = plan_bfft(num_p)
  2 .* real.((bfft_plan * num_p) ./ (bfft_plan * den_p))
end

function reference(::typeof(fourier_pade))
"""
@article{driscoll2001pade,
  title={A {P}ad{\'e}-based algorithm for overcoming the {G}ibbs phenomenon},
  author={Driscoll, Tobin A and Fornberg, Bengt},
  journal={Numerical Algorithms},
  volume={26},
  number={1},
  pages={77--92},
  year={2001},
  publisher={Springer},
  doi={10.1023/A:1016648530648}
}
"""
end


"""
    Student's t

* Parametrization
parametrized in \\nu

* Score

* Fisher Information

* `time_varying_params` map.

* Default link
"""
TDist

function score!(score_til::Matrix{T}, y::T, ::Type{TDist}, param::Matrix{T}, t::Int) where T
    score_til[t, 1] = 0.5 * (((y^2)*(param[t, 1] + 1)/(param[t, 1] * y^2 + param[t, 1]^2)) -
                      1/param[t, 1] -
                      log((y^2)/(param[t, 1]) + 1) -
                      digamma(param[t, 1] / 2) + digamma((param[t, 1] + 1) / 2))
    return
end

function fisher_information!(aux::AuxiliaryLinAlg{T}, ::Type{TDist}, param::Matrix{T}, t::Int) where T
    aux.fisher[1, 1] = 0.5 * (0.5 * trigamma(0.5 * param[t, 1]) - 0.5 * trigamma( 0.5 * (param[t, 1] + 1.0) ) -
                             (param[t, 1] + 5.0) / (param[t, 1] * (param[t, 1] + 3.0) * (param[t, 1] + 1.0)))
    return
end

function log_likelihood(::Type{TDist}, y::Vector{T}, param::Matrix{T}, n::Int) where T
    loglik = zero(T)
    for t in 1:n
        loglik -= 0.5 * log(param[t, 1]) + logbeta(0.5, param[t, 1]/2) + ((param[t, 1] + 1)/2) *
                  log(1 + (y[t]^2)/param[t, 1])
    end
    return -loglik
end

# Links
function link!(param_tilde::Matrix{T}, ::Type{TDist}, param::Matrix{T}, t::Int) where T
    param_tilde[t, 1] = link(LogLink, param[t, 1], zero(T))
    return
end
function unlink!(param::Matrix{T}, ::Type{TDist}, param_tilde::Matrix{T}, t::Int) where T
    param[t, 1] = unlink(LogLink, param_tilde[t, 1], zero(T))
    return
end
function jacobian_link!(aux::AuxiliaryLinAlg{T}, ::Type{TDist}, param::Matrix{T}, t::Int) where T
    aux.jac[1] = jacobian_link(LogLink, param[t, 1], zero(T))
    return
end

# utils
function update_dist(::Type{TDist}, param::Matrix{T}, t::Int) where T
    return TDist(param[t, 1])
end

function params_sdm(d::TDist)
    return Distributions.params(d)
end

function num_params(::Type{TDist})
    return 1
end


include("../../src/MutualInformation/mutual-information.jl")
include("reader_MI_tests.jl")

function get_layer_bins_vector(vector; neuron_nbins = 5, nbins = 5, lower_lim=-1, upper_lim=1)
    columns = separate_columns(vector)
    hists = create_bin_histogram.(columns, [neuron_nbins], [lower_lim], [upper_lim])
    bounds = (get_weights_length.(hists)...,)
    indices = [getindex.(get_bin.(hists, v),1) for v in vector]
    indices = [correct_upper_bounds(ind, bounds) for ind in indices]
    combination_values = get_bin_from_indices.(indices, [bounds])
    max_value = foldr(*, bounds)
    create_bin_vector(combination_values, nbins, 1, max_value)
end

function calculate_joint_probability_MI(X, Y, lim_X, lim_Y, nbins)
    number_neurons_X = length(X[1])
    bins_X = nbins^(number_neurons_X)
    if number_neurons_X > 1
        linearized_X = get_layer_bins_vector(X, neuron_nbins=bins_X, nbins=nbins, lower_lim=lim_X[1], upper_lim=lim_X[2])
        limits_X = (minimum(linearized_X), maximum(linearized_X))
    else
        linearized_X = getindex.(X, 1)
        limits_X = lim_X
    end

    number_neurons_Y = length(Y[1])
    bins_Y = nbins^(number_neurons_Y)
    if number_neurons_Y > 1
        linearized_Y = get_layer_bins_vector(Y, neuron_nbins=nbins, nbins=bins_Y, lower_lim=lim_Y[1], upper_lim=lim_Y[2])
        limits_Y = (minimum(linearized_Y), maximum(linearized_Y))
    else
        linearized_Y = getindex.(Y, 1)
        limits_Y = lim_Y
    end

    calculate_joint_probability(linearized_X, linearized_Y, limits_X, limits_Y, (bins_X+1, round(Int, (bins_Y))))
end

function get_layer_probabilities(vector; neuron_nbins = 5, nbins = 5, lower_lim=-1, upper_lim=1)
    columns = separate_columns(vector)
    hists = create_bin_histogram.(columns, [neuron_nbins], [lower_lim], [upper_lim])
    probabilities = map(h -> h.weights/sum(h.weights), hists)
    probability = fit(Histogram, (colunms...), nbins = neuron_nbins)
end

#mutual_information(calculate_joint_probability()...)

# dt2 = get_numerical_dataset_test2()
#p2 = calculate_joint_probability_MI(dt2[1], dt2[2], (-2, 1.5), (-1, 1), 3)
#MI = mutual_information(p2...)

#=
Cálculo da informação mútua (caso 1):
-> Valor teórico calculado: I(X;T) = 0.047056227
-> Valor obtido pela aproximação: I(X;T) = 0.04700065437122643 (3 bins)
-> Valor obtido pela aproximação: I(X;T) = 0.04510842519666598 (2 bins)

Cálculo da informação mútua utilizando os valores teóricos (caso 2):
-> Valor teórico calculado: I(X;T) = 0.1685846799340601
-> Valor calculado pela aproximação: I(X;T) = 0.16843457465541672 (9 bins)
-> Valor calculado pela aproximação: I(X;T) = 0.05633855652784133 (4 bins)
=#


"""
    @polyeval(z, p, N)

Evaluate the polynomial ``\\sum_k c[k] z^{k-1}`` for the coefficients `c[1]`, `c[2]`, ..., `c[N]`;
that is, the coefficients are given in ascending order by power of `z`.  This macro expands
to efficient inline code that uses either Horner's method.

```jldoctest

julia> p = randn(5)
5-element Array{Float64,1}:
  0.106455
  0.0351716
 -0.204764
 -0.87642
  0.260407

julia> @polyeval(0.5, p, 5)
-0.020427265391245605

```
"""
macro polyeval(x, p, N)
    if N==1
        return :($(esc(p))[1])
    elseif N==2
        return :($(esc(p))[1] + $(esc(x))*$(esc(p))[2])
    else
        nm1 = N-1
        ex = :($(esc(p))[$nm1] + $(esc(x))*$(esc(p))[$N])
        for i in N-2:-1:1
            ex = :($(esc(p))[$i] + $(esc(x))*$ex)
        end
    end
    return ex
end

"""
    ConvergenceError([msg, val, niter, err])

Exception that is to be thrown when an iterative procedure fails
to converge.

"""
mutable struct ConvergenceError <: Exception
    msg::String
    val::Float64
    niter::Int
    err::Float64
    ConvergenceError() = new("Calculations failed to converge", NaN, -1, NaN)
    ConvergenceError(msg) = new(msg, NaN, -1, NaN)
    ConvergenceError(msg, val) = new(msg, val, -1, NaN)
    ConvergenceError(msg, val, niter) = new(msg, val, niter, NaN)
    ConvergenceError(msg, val, niter, err) = new(msg, val, niter, err)
end

"""
    calcz(vm, b0, c0, [EPS, [MAXITER, [relax]]])

Calculates the compressibility factor using a virial equation:

`z = 1 + b₀/z + c₀/vₘ²`

The algorithm initially uses as an initial guess only the b₀ virial
coefficient. Then it uses a Newton-Raphson algorithm to compute the
compressibility factor

 * `b0` b₀ parameter of the virial equation
 * `c0` c₀ parameter of the virial equation
 * `EPS` Convergence criterium
 * `MAXITER` Maximim number of iterations
 * `relax` Sub-relaxation parameter if convergence is difficult.
"""
function calcz(b0, c0, EPS=1e-8, MAXITER=200, relax=1.0)

    z = (1.0 + sqrt(1.0 + 4*b0)) / 2.0
    dz = 0.0

    for i = 1:MAXITER
        f = -c0 + z*(-b0 + z*(-1.0 + z))
        df = -b0 + z*(-2.0 + 3.0*z)
        dz = -f / df
        if abs(dz) < EPS
            return z + dz
        end
        z = z + relax*dz
    end
    throw(ConvergenceError("Compressibility factor failed to converge properly!",
                           z, MAXITER, abs(dz)))

end


using RedBlackGraph
using LinearAlgebra
using Test

@testset "Avus Linear Algebra Tests" begin
    r = red_one(1)
    o = one(r)

    # test dot product of myself and father
    u = [r, 2, 3, 0, 0]
    v = [2, r, 0, 0, 0]
    @test u ⋅ v == 2

    # test dot product of myself and mother
    u = [r, 2, 3, 0, 0]
    v = [3, 0, o, 0, 0]
    @test u ⋅ v == 3

    # sample from docs
    u = [2, 0, 0, 0, r]
    v = [3, 0, r, 0, 0]
    @test u ⋅ v == 5

    # test dot product of my daughter and my father
    u = [2, 0, 0, 0, o]
    v = [2, r, 0, 0, 0]
    @test u ⋅ v == 4

    A = [r 2 3 0 0;
         0 r 0 2 0;
         0 0 1 0 0;
         0 0 0 r 0;
         2 0 0 0 1]
    expected_result =  [r 2 3 4 0;
                        0 r 0 2 0;
                        0 0 1 0 0;
                        0 0 0 r 0;
                        2 4 5 0 1]
    result = A * A
    @test result == expected_result

    expected_result =  [r 2 3 4 0;
                        0 r 0 2 0;
                        0 0 1 0 0;
                        0 0 0 r 0;
                        2 4 5 8 1]
    result = A * result
    @test result == expected_result

    A =  [r 2 3 4 0;
          0 r 0 2 0;
          0 0 1 0 0;
          0 0 0 r 0;
          2 4 5 0 1]
    u = [2 0 0 0 0]
    v = [0, 3, 0, 0, 0]

    result = u * A
    u_lambda = [2 4 5 8 0]
    @test result == u_lambda

    result = A * v
    v_lambda = [5; 3; 0; 0; 9]
    @test result == v_lambda

    A = [r 2 3 0 0;
         0 r 0 2 0;
         0 0 1 0 0;
         0 0 0 r 0;
         2 0 0 0 1]
    @test A == A * I

end


# ---
# title: 1515. Best Position for a Service Centre
# id: problem1515
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Geometry
# link: <https://leetcode.com/problems/best-position-for-a-service-centre/description/>
# hidden: true
# ---
#
# A delivery company wants to build a new service centre in a new city. The
# company knows the positions of all the customers in this city on a 2D-Map and
# wants to build the new centre in a position such that **the sum of the
# euclidean distances to all customers is minimum**.
#
# Given an array `positions` where `positions[i] = [xi, yi]` is the position of
# the `ith` customer on the map, return _the minimum sum of the euclidean
# distances_ to all customers.
#
# In other words, you need to choose the position of the service centre
# `[xcentre, ycentre]` such that the following formula is minimized:
#
# ![](https://assets.leetcode.com/uploads/2020/06/25/q4_edited.jpg)
#
# Answers within `10^-5` of the actual value will be accepted.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/06/25/q4_e1.jpg)
#
#
#
#     Input: positions = [[0,1],[1,0],[1,2],[2,1]]
#     Output: 4.00000
#     Explanation: As shown, you can see that choosing [xcentre, ycentre] = [1, 1] will make the distance to each customer = 1, the sum of all distances is 4 which is the minimum possible we can achieve.
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2020/06/25/q4_e3.jpg)
#
#
#
#     Input: positions = [[1,1],[3,3]]
#     Output: 2.82843
#     Explanation: The minimum possible sum of distances = sqrt(2) + sqrt(2) = 2.82843
#
#
# **Example 3:**
#
#
#
#     Input: positions = [[1,1]]
#     Output: 0.00000
#
#
# **Example 4:**
#
#
#
#     Input: positions = [[1,1],[0,0],[2,0]]
#     Output: 2.73205
#     Explanation: At the first glance, you may think that locating the centre at [1, 0] will achieve the minimum sum, but locating it at [1, 0] will make the sum of distances = 3.
#     Try to locate the centre at [1.0, 0.5773502711] you will see that the sum of distances is 2.73205.
#     Be careful with the precision!
#
#
# **Example 5:**
#
#
#
#     Input: positions = [[0,1],[3,2],[4,5],[7,6],[8,9],[11,1],[2,12]]
#     Output: 32.94036
#     Explanation: You can use [4.3460852395, 4.9813795505] as the position of the centre.
#
#
#
#
# **Constraints:**
#
#   * `1 <= positions.length <= 50`
#   * `positions[i].length == 2`
#   * `0 <= positions[i][0], positions[i][1] <= 100`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using EconPDEs

struct WangWangYangModel
    μ::Float64
    σ::Float64
    r::Float64
    ρ::Float64
    γ::Float64
    ψ::Float64
    wmax::Float64
end

function WangWangYangModel(;μ = 0.015, σ = 0.1, r = 0.035, ρ = 0.04, γ = 3, ψ = 1.1, wmax = 5000.0)
    WangWangYangModel(μ, σ, r, ρ, γ, ψ, wmax)
end

function initialize_stategrid(m::WangWangYangModel; n = 500)
    OrderedDict(:w => range(0.0, stop = sqrt(m.wmax), length = n).^2)
end

function initialize_y(m::WangWangYangModel, stategrid)
    OrderedDict(:p => 1 .+ stategrid[:w])
end

function (m::WangWangYangModel)(state::NamedTuple, y::NamedTuple)
    μ = m.μ ;  σ = m.σ ;  r = m.r ;  ρ = m.ρ ;  γ = m.γ ;  ψ = m.ψ ; wmax = m.wmax
    w = state.w
    p, pw, pww = y.p, y.pw, y.pww
    c = (r + ψ * (ρ - r)) * p * pw^(-ψ)
    μw = (r - μ + σ^2) * w + 1 - c
    # One only needs a ghost node if μw <= 0 (since w^2p_ww = 0). In this case, we obtain a formula for pw so that c <= 1
    if w ≈ 0.0 && μw <= 0.0
       pw = ((r + ψ * (ρ - r)) * p)^(1 / ψ)
       c = (r + ψ * (ρ - r)) * p * pw^(-ψ)
       μw = (r - μ + σ^2) * w + 1 - c
    end
    # At the top, I use the solution of the unconstrainted, i.e. pw = 1 (I could also do reflecting boundary but less elegant)
    pt = (((r + ψ * (ρ - r)) * pw^(1 - ψ) - ψ * ρ) / (ψ - 1) + μ - γ * σ^2 / 2) * p + ((r - μ + γ * σ^2) * w + 1) * pw + σ^2 * w^2 / 2  * (pww - γ * pw^2 / p)
    μw = 1 + (r - μ + σ^2) * w - c
    return (pt,), (μw,), (w = w, p = p, pw = pw, pww = pww, μw = μw, c = c)
end

m = WangWangYangModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
y, result, distance = pdesolve(m, stategrid, y0; bc = OrderedDict(:pw => (1.0, 1.0)))


export jacobian,  StabilityType

@with_kw struct jacobian
    ode
    jac_func
    u
    p
    t
end

"""
    StabilityType

Store the information about stability.
"""
@with_kw struct StabilityType
    stable ::Bool
    unstable ::Bool
    saddle  ::Bool
    damping ::Bool
end

"""
Construct the jacobian struct type

Argument
--------
- `ode` {ODE function}: with the form `f(du,u,p,t)`
- `u` {Array}: initial values
- `p` {Array}: parameters

Return
------
- `Jacobian` {struct}

Warning
-------
NameTuple definition of `ode` function is unrecommended. Due to the incompatibility with `ModelingToolkit.modelingtoolkitize`
"""
function jacobian(ode, u, p; t=nothing)
    prob = ODEProblem(ode, u, t, p)
    de =  ModelingToolkit.modelingtoolkitize(prob)

    jac_func = ModelingToolkit.generate_jacobian(de, expression=Val{false})[2]
    return jacobian(ode=ode, jac_func=jac_func, u=u, p=p, t=t)
end

"""Use DEmeta for jacobian generation"""
function jacobian(de::T; t= nothing) where T<:DEmeta
    return jacobian(de.func, de.u0, de.p; t=t)
end

"""
Calculate the jacobian of given initial valuables and parameters.
"""
function (self::jacobian)(u,p;t=nothing)
    l = length(u)
    j = zeros(l,l)
    self.jac_func(j, u, p, t)
    return j
end

"""
Get jacobian from state variable with default paramter set
"""
function (self::jacobian)(u;t=nothing)
    return self(u, self.p; t=t)
end

"""
Get jacobian from default state.
"""
function (self::jacobian)()
    return self(self.u, self.p; t=self.t)
end

function StabilityType(jac_matrix)
    es = eigvals(jac_matrix)

    return StabilityType(
    stable = is_stable(es),
    unstable = is_unstable(es),
    saddle = is_saddle(es),
    damping = is_damping(es)
    )
end

"""
Stability identification
"""
function is_stable(es) :: Bool
    method(e) = real(e) >= 0. ? true : false
    return judge_jac(es, method, true)
end

"""
Unstability identification
"""
function is_unstable(es) :: Bool
    method(e) = real(e) >= 0. ? true : false
    return judge_jac(es, method, false)
end

"""
Damping identification
"""
function is_damping(es) :: Bool
   method(e) = imag(e) != 0. ? true : false
   return judge_jac(es, method, false)
end

"""
Saddle point identification
"""
function is_saddle(es) :: Bool
    real_es = real.(es)
    if (maximum(real_es) > 0. ) & (minimum(real_es) < 0.)
        saddle = true
    else
        saddle = false
    end
    return saddle
end

"""
General method for judging jacobian
"""
function judge_jac(es, eigval_method, default_bool)
    feature = default_bool
    for e in es
        if eigval_method(e)
            feature = !feature
            break
        end
    end
    return feature
end


# Estimate polynomial linear regression model with a half-T prior.

using DynamicHMCModels

ProjDir = @__DIR__
cd(ProjDir)

# Import the dataset.

data = CSV.read(joinpath("..", "..", "data", "Howell1.csv"), DataFrame)

# Use only adults and standardize

df = filter(row -> row[:age] >= 18, data);
df[!, :weight] = convert(Vector{Float64}, df[:, :weight]);
df[!, :weight_s] = (df[:, :weight] .- mean(df[:, :weight])) / std(df[:, :weight]);
df[!, :weight_s2] = df[:, :weight_s] .^ 2;

# Define a structure to hold the data: observables, covariates,
# and the degrees of freedom for the prior.

"""
Linear regression model ``y ∼ Xβ + ϵ``, where ``ϵ ∼ N(0, σ²)`` IID.
Flat prior for `β`, half-T for `σ`.
"""
Base.@kwdef mutable struct LinearRegressionModel{Ty <: AbstractVector, Tx <: AbstractMatrix,
Tv <: Real}
    "Observations."
    y::Ty
    "Covariates"
    x::Tx
    "Degrees of freedom for prior."
    v::Tv
end

# Write a function to return a properly dimensioned transformation.

function make_transformation(model::LinearRegressionModel)
  as((β = as(Array, size(model.x, 2)), σ = asℝ₊))
end

N = size(df, 1)
x = hcat(ones(N), hcat(df[:, :weight_s], df[:, :weight_s2]));
model = LinearRegressionModel(;y = df[:, :height], x=x, v=1.0)

# Pack parameters *as a single argument*.

function (model::LinearRegressionModel)(θ)
    @unpack y, x, v = model   # extract data
    @unpack β, σ = θ                 # extract parameters
    loglikelihood(Normal(0, σ), y .- x*β) + logpdf(TDist(v), σ)
end

# Evaluate at model function at some initial valuues

println()
model((β = [1.0, 2.0, 3.0], σ = 1.0)) |> display
println()

# Wrap the problem with a transformation, then use Flux for the gradient.

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);

# Tune and sample.

results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

p = as_particles(posterior)
display(p)

stan_result = "
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
           Mean         SD       Naive SE       MCSE      ESS
    a 154.609019750 0.36158389 0.0057171433 0.0071845548 1000
   b1   5.838431778 0.27920926 0.0044146860 0.0048693502 1000
   b2  -0.009985954 0.22897191 0.0036203637 0.0047224478 1000
sigma   5.110136300 0.19096315 0.0030193925 0.0030728192 1000

Quantiles:
          2.5%        25.0%        50.0%       75.0%        97.5%
    a 153.92392500 154.3567500 154.60700000 154.8502500 155.32100000
   b1   5.27846200   5.6493250   5.83991000   6.0276275   6.39728200
   b2  -0.45954687  -0.1668285  -0.01382935   0.1423620   0.43600905
sigma   4.76114350   4.9816850   5.10326000   5.2300450   5.51500975
";

# end of m4.5d.jl


using NumericalMethodsforEngineers
using Test

@sym begin
  CollocationWeightedResidualMethod(xi_, yi_, doprint_) := Module([Y, R],
    begin
      Y(x_) := LagrangePolynomial(xi, yi)
      If(doprint, Println("\nY(x) = ", Y(x), "\n"))
      #
      # Can be formulated as ytile(x) = F(x) + C1(a) * Ψ(x)
      #
      # F(x_) := 2*x^2 - x
      # C(a_) := -4a
      # Ψ(x_) := x^2 - x
      #
      #Y(x_) := F(x) + C(a)*Ψ(x)
      #
      R(x_) = Simplify(D(Y(x), x, 2) - 3*x - 4*Y(x))
      R(x_) = Simplify(R(x) ./ (a => -C1/4))
      If(doprint, Println("R(x) = ", R(x), "\n"))
      Return(R(x))
    end
  )
end

function wrm(x1,y1, doprint=true)
    setsymata(:xin, List(x1...))
    setsymata(:yin, List(y1...))
    setsymata(:doprint, doprint ? true : false)
    symparseeval("Compile(Evaluate(CollocationWeightedResidualMethod(xin,yin,doprint)))")
end

println("\nExample 7.13: y'' = 3x + 4y, y(0)=0, y(1)=1")
println("Residual for Weighted Residual Method using 1 point Lagragian Polynomial\n")

#
# A consequence of the Julia function based approach is that Y(x) and R(x) are local
# to the Symata function and are not available for examination, e.g. @sym R(x)
#

lp_eq = @sym ToString(LagrangePolynomial([0.0,0.5,1.0], [0.0,a,1.0]))
println(lp_eq)
@test lp_eq == "x*(-1.0 + 2.0*x + (-4.0)*a*(-1.0 + x))"
wrm_eq = @sym ToString(CollocationWeightedResidualMethod([0.0,0.5,1.0], [0.0,a,1.0], false))
println(wrm_eq)
@test wrm_eq == "4.0 + 2.0*C1 - 3x - 4x*(-1.0 + 2.0*x + C1*(-1.0 + x))"

#
# Another consideration is that below yt is yt(C1::Any, x::Any)
#

yt = wrm([0.0,0.5,1.0], [0.0,:a,1.0], true)

println("( Example 7.13 gives: R = -4x^2*(2 + C1) + x*(1 + 4C1) + 2*(2 + C1) )", "\n")

"
```
The function ytilde(x, C1) contains the 1-point Weighted Residual Method
polynomial with parameter C1.
```
Arguments to ytilde(x, C1):

x::Float64          : x value\n
C1::Float64         : Weighted Residual Method constant.

C1 can be estimated by e.g. collocation, subdomain, least squares or Galerkin.
See examples 7.15 through 7.18 and Fig.7.14.jl
"
ytilde(x, C1) = yt(C1, x)
println("ytilde(1.0, 0.0) = $(ytilde(1.0, 0.0))")
@test ytilde(1.0, 0.0) == -3.0

println("ytilde(0.0, 1.0) = $(ytilde(0.0, 1.0))")
@test ytilde(0.0, 1.0) == 6.0
#
# For now I will continue to use the E.7.13.jl approach.
#


include("misc.jl")

function rbfBasis(X1, X2, sigma)
  return exp.(-distancesSquared(X1,X2)/(2*sigma^2))
end

function leastSquaresRBFL2(X, y, l, sigma)
	# Make RBF basis
	n = size(X,1)
	Z = rbfBasis(X, X, sigma)

	# Find regression weights minimizing squared error with L2 regularization
	w = (Z'*Z + eye(size(Z,2))*l)\(Z'*y)

	display(size(w))

	# Make linear prediction function
	predict(Xtilde) = rbfBasis(Xtilde, X, sigma)*w

	# Return model
	return LinearModel(predict,w)
end


using NLsolve, StaticArrays

function f!(F, x)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
end

function j!(J, x)
    J[1, 1] = x[2]^3-7
    J[1, 2] = 3*x[2]^2*(x[1]+3)
    u = exp(x[1])*cos(x[2]*exp(x[1])-1)
    J[2, 1] = x[2]*u
    J[2, 2] = u
end

nlsolve(f!, j!, @MArray [ 0.1; 1.2], )


@testset "Utils" begin
    @testset "sec2rad" begin
        # Radians <-> arcseconds
        @test sec2rad(3600) == deg2rad(1)
        @test rad2sec(1) == rad2deg(1) * 3600
    end
    @testset "normalize_angle" begin
        @testset for a = -15.0:0.1:15.0, b = -15.0:0.2:15.0
            c = normalize_angle(a, b)
            @test (b - π) <= c
            @test c <= (b + π)
            two_k = round((a - c) / π)
            @test c ≈ a - two_k * π atol=1.0e-14
        end
    end
end


using LinearAlgebra: I, nullspace
using UniqueVectors: UniqueVector
using StaticArrays: MMatrix

"""
    parametrize_subspace(M::AbstractMatrix{<:Real}, v::AbstractVector{<:Real})

`M` is a `K-n × K` full rank matrix and `v` is a `K-n` length vector. The equation `Mx = v`
describes an `n`-dimensional affine subspace. Compute `a` and `b` so that this space is
parametrized as `x = a + b * x[1:n]`. The columns of `b` are a basis for the null space of
`M` and `a` is a solution of `Mx=v`.
"""
function parametrize_subspace(M::AbstractMatrix{<:Real}, v::AbstractVector{<:Real})
    n = size(M, 2) - size(M, 1)
    @assert n >= 0
    b = nullspace(M)
    N = inv([transpose(b); M]) # this is square because M is full rank
    a = N[:, (n+1):end] * v
    return a, b
end

"""
    intersect_linear_quadratic(A::AbstractMatrix{<:Real},
    B::AbstractVector{<:Real}, C::Real, a::AbstractVector{<:Real},
    b::AbstractVector{<:Real})

Intersect the line `x = a + b * t` with the quadratic `xᵀAx + Bᵀx + C = 0` and return the
two intersection points.
"""
function intersect_linear_quadratic(A::AbstractMatrix{<:Real},
    B::AbstractVector{<:Real}, C::Real, a::AbstractVector{<:Real},
    b::AbstractVector{<:Real})
    quadratic = transpose(b) * A * b
    linear = transpose(a) * A * b + transpose(b) * A * a + transpose(B) * b
    constant = transpose(a) * A * a + transpose(B) * a + C
    t1, t2 = solve_quadratic(quadratic, linear, constant)
    return a + b * t1, a + b * t2
end

"""
    solve_quadratic(a::Real, b::Real, c::Real)

Find the two solutions of `ax² + bx + c = 0`.
"""
function solve_quadratic(a::Real, b::Real, c::Real)
    t1 = -b/(2 * a)
    t2 = sqrt(b^2 - 4 * a * c)/(2 * a)
    return t1 + t2, t1 - t2
end

"""
    sign_of_permutation(a1::AbstractVector{T}, a2::AbstractVector{T}) where T

Given two arrays which are permutations of each other, find the sign of the permutation.
"""
function sign_of_permutation(a1::AbstractVector{T}, a2::AbstractVector{T}) where T
    n = length(a1)
    @assert length(a2) == n
    inds = [findfirst(isequal(x), a2) for x in a1]
    m = MMatrix{n,n,Int}(zeros(n,n)) # det is faster for small matrices using StaticArrays
    for i in 1:n
        m[i, inds[i]] = 1
    end
    return sign(det(m))
end

"""
    first_setdiff_index(l1, l2)

Given two collections `l1` and `l2` find the index of the first element of `l1` that is
not in `l2`. If no such element exists return `0`.
"""
function first_setdiff_index(l1, l2)
    for (i, e) in enumerate(l1)
        if !(e in l2)
            return i
        end
    end
    return 0
end

"""
    first_setdiff(l1, l2)

Given two collections `l1` and `l2` find the first element of `l1` that is not in `l2`.
If no such element exists return `nothing`.
"""
function first_setdiff(l1, l2)
    i = first_setdiff_index(l1, l2)
    return i > 0 ? l1[i] : nothing
end

"""
    triangulated_lattice(v1::Vector{<:Real}, v2::Vector{<:Real}, n::Int, m::Int)

Subdivide a parallelogram with sides `v1` and `v2` into `n*m` similar parallelograms.
Divide each in two pieces to find a triangulation. Return the points and a
TriangulatedComplex.
"""
function triangulated_lattice(v1::Vector{<:Real}, v2::Vector{<:Real}, n::Int, m::Int)
    @assert length(v1) == length(v2) == 2
    inds = UniqueVector([(i, j) for i in 1:(n+1) for j in 1:(m+1)])
    points = [Point(v1 * (i-1)/n + v2 * (j-1)/m) for (i, j) in inds]
    simplices = Simplex{2, 3}[]
    for i in 1:n
        for j in 1:m
            p11 = points[findfirst(isequal((i,j)), inds)]
            p12 = points[findfirst(isequal((i,j+1)), inds)]
            p21 = points[findfirst(isequal((i+1,j)), inds)]
            p22 = points[findfirst(isequal((i+1,j+1)), inds)]
            push!(simplices, Simplex([p11, p12, p21]))
            push!(simplices, Simplex([p12, p22, p21]))
        end
    end
    return points, TriangulatedComplex(simplices)
end


# calculate probability by bayes theorem

module ProbsBayes
    using DataFrames, CSV, FreqTables

    function main()
        data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_600.txt")
        df_org = CSV.read(data_path, DataFrame,
                      header=["date", "time", "ir", "lidar"],
                      delim=' ')

        hours = [Int64(floor(e/10000)) for e in df_org.time]

        # add hour array to df
        df_h = DataFrame(hour=hours)
        df_new = hcat(df_org, df_h)

        # P(t)
        freqs = freqtable(df_new, :lidar, :hour)
        probs = freqs / length(df_org.lidar)
        keys = names(probs, 2)
        p_t = [sum(probs[begin:end, Name(key)]) for key in keys]
        println("P(t) = $(p_t)")

        # P(z)
        probs_trans = transpose(probs)
        keys_trans = names(probs_trans, 2)
        p_z = [sum(probs_trans[begin:end, Name(key)]) for key in keys_trans]
        println("P(z) = $(p_z)")

        # P(t|z)
        for (i, key) in enumerate(keys_trans)
            probs_trans[begin:end, Name(key)] /= p_z[i]
        end
        cond_t_z = probs_trans

        p_z_630 = p_z[findfirst(keys_trans .== 630)]
        p_t_13 = p_t[findfirst(keys .== 13)]
        p_t_13_z_630 = cond_t_z[Name(13), Name(630)]

        bayes = p_t_13_z_630 * p_z_630 /p_t_13

        cond_z_t = probs / p_t[findfirst(keys .== 13)]
        answer = cond_z_t[Name(630), Name(13)]

        println("P(z=630) = $(p_z_630)")
        println("P(t=13) = $(p_t_13)")
        println("P(t=13|z=630) = $(p_t_13_z_630)")
        println("Bayes P(z=630|t=13) = $(bayes)")
        println("Answer P(z=630|t=13) = $(answer)")
    end
end


const Prime = UInt64

struct Primes
    rangei :: Int64
    primebits :: BitArray{1}
    function Primes(n :: Int64)
        if n < 3
          if n < 2 return new(-1, falses(0)) # no elements
          else return new((0, trues(0))) end # n = 2: meaning is 1 element of 2
        end
        limi :: Int = (n - 1) ÷ 2 # calculate the required array size
        isprimes :: BitArray = trues(limi)
        @inbounds(
        for i in 1:limi
            p = i + i + 1
            start = (p * p - 1) >>> 1 # shift/divide if LLVM doesn't optimize
            if start > limi
                return new(limi, isprimes)
            end
            if isprimes[i]
                for j in start:p:limi
                  isprimes[j] = false
                end
            end
        end)
    end
end

Base.eltype(::Type{Primes}) = Prime

function Base.length(P::Primes)::Int64
    if P.rangei < 0 return 0 end
    return 1 + count(P.primebits)
end

function Base.iterate(P::Primes, state::Int = 0)::
                                        Union{Tuple{Prime, Int}, Nothing}
    lmt = P.rangei
    if state > lmt return nothing end
    if state <= 0 return (UInt64(2), 1) end
    let
        prmbts = P.primebits
        i = state
        @inbounds(
        while i <= lmt && !prmbts[i] i += 1 end)
        if i > lmt return nothing end
        return (i + i + 1, i + 1)
    end
end


import Base.length
import Base.isempty
import Base.getindex
import Base.setindex!
import Base.firstindex
import Base.lastindex
import Base.sum

# Use a binary tree to describe a categorical distribution
# This allows for O(1) read, O(log n) write and O(log n) generation of a random index
# Optimized for categorical distributions with a large number of possible outcomes

# Should this be a subtype of AbstractVector?
struct CategoricalTree{T<:Real}
    ns::Vector{<:Integer} # number of nodes at each layer of the tree (`ns[end]` should be 1, i.e. the root of the tree
    as::Vector{Vector{T}} # store each layer into a separate array. Each parent node contains the sum of its children
end

function CategoricalTree(a::Vector{T}) where T<:Real
    n = length(a)
    ns::Vector{Int64} = [n]
    while n > 1
        if isodd(n)
            n = Int((n+1)/2)
            push!(ns,n)
        else
            n = Int(n/2)
            push!(ns,n)
        end
    end
    as = [zeros(T,k) for k in ns]
    as[1] = a
    for i in 1:(length(ns)-1)
        for j in 2:2:ns[i]
            as[i+1][Int(j/2)] = as[i][j-1] + as[i][j]
        end
        if isodd(ns[i])
            as[i+1][ns[i+1]] = as[i][ns[i]]
        end
    end
    return CategoricalTree{T}(ns,as)
end

isempty(ct::CategoricalTree{<:Real}) = isempty(ct.as[1])
length(ct::CategoricalTree{<:Real}) = length(ct.as[1])

firstindex(::CategoricalTree) = 1
lastindex(ct::CategoricalTree) = length(ct)

getindex(ct::CategoricalTree{<:Real}, i) = ct.as[1][i]

function setindex!(ct::CategoricalTree{<:Real}, x, i)
    ct.as[1][i] = x
    j = i
    for k in 2:length(ct.ns)
        j = isodd(j) ? Int((j+1)/2) : Int(j/2)
        if isodd(ct.ns[k-1]) && j == ct.ns[k]
                ct.as[k][j] = ct.as[k-1][2*j-1]
        else
            ct.as[k][j] = ct.as[k-1][2*j-1] + ct.as[k-1][2*j]
        end
    end
end

# since we already computed the sum, we might as well make use of it
sum(ct::CategoricalTree{<:Real}) = ct.as[end][1]

"""
    rand_categorical(ct::CategoricalTree{<:Real}, a0)

Pick one index according to a categorical distribution described by `ct`.
"""
function rand_categorical(ct::CategoricalTree{<:Real}, a0 = sum(ct))
    r = rand()*sum(ct)
    i = length(ct.ns) - 1
    j = 1
    while i > 1
        if r <= ct.as[i][j]
            j = 2*j - 1
        else
            r -= ct.as[i][j]
            j = 2*j + 1
        end
        i -= 1
    end
    if r <= ct.as[i][j]
        return j
    else
        return j+1
    end
end


using PyPlot

A = [1.0 2.0
     2.0 1.0]

evals, evecs = eig(A)
evecs =(evecs[:,1], evecs[:,2])

fig, ax = subplots()
for spine in ["left", "bottom"]
    ax[:spines][spine][:set_position]("zero")
end

for spine in ["right", "top"]
    ax[:spines][spine][:set_color]("none")
end
ax[:grid](alpha=0.4)

xmin, xmax = -3, 3
ymin, ymax = -3, 3
ax[:set_xlim](xmin, xmax)
ax[:set_ylim](ymin, ymax)

for v in evecs
    # Plot each eigenvector
    ax[:annotate](" ", xy=v, xytext=[0, 0],
                  arrowprops={"facecolor"=>"blue",
                              "shrink"=>0,
                              "alpha"=>0.6,
                              "width"=>0.5})

    # Plot the image of each eigenvector
    v = A * v
    ax[:annotate](" ", xy=v, xytext=[0, 0],
                  arrowprops={"facecolor"=>"red",
                              "shrink"=>0,
                              "alpha"=>0.6,
                              "width"=>0.5})
end

x = linspace(xmin, xmax, 3)
for v in evecs
    a = v[2] / v[1]
    ax[:plot](x, a .* x, "b-", lw=0.4)
end


@testset "Boltzmann" begin
    @test RTFOOL.boltzmann(1.0, 1.0) == e^-1.00
    @test RTFOOL.boltzmann(1.0, 2.0) == e^-2.00
    @test RTFOOL.boltzmann(0.5, 1.5) == e^-0.75

    @test RTFOOL.boltzmann(1.0, [1.0]) == [1.0]
    @test RTFOOL.boltzmann(1.0, [2.0]) == [1.0]
    @test RTFOOL.boltzmann(1.0, [1.0, 1.0]) == [0.5, 0.5]
    @test RTFOOL.boltzmann(1.0, [1.0, 2.0]) == [e/(1+e), 1/(1+e)]
    @test RTFOOL.boltzmann(0.5, [0.5, 1.5]) ≈ [e/(e + √e), √e/(e + √e)]
    @test RTFOOL.boltzmann(1/3, [1, 2, 3]) ==
        [e^-(1/3), e^-(2/3), e^-1] / (e^-(1/3) + e^-(2/3) + e^-1)

    @test RTFOOL.boltzmann(1.0, [1.0], [2]) == [1.0]
    @test RTFOOL.boltzmann(1.0, [2.0], [3]) == [1.0]
    @test RTFOOL.boltzmann(2.0, [1.0, 1.0], [1, 2]) == [1.0/3, 2.0/3]
    @test RTFOOL.boltzmann(1.0, [1.0, 2.0], [2, 1]) == [2e/(1+2e), 1/(1+2e)]
    @test RTFOOL.boltzmann(0.5, [0.5, 1.5], [2, 3]) ≈ [2e/(2e + 3*√e), 3*√e/(2e + 3*√e)]
    @test RTFOOL.boltzmann(1/3, [1, 2, 3], [1, 2, 3]) ==
        [e^-(1/3), 2e^-(2/3), 3e^-1] / (e^-(1/3) + 2e^-(2/3) + 3e^-1)

    let space = StateSpace{1}([(1,)], [1.0], BigInt[2])
        @test RTFOOL.boltzmann(1.0, space) == [1.0]
    end
    let space = StateSpace{1}([(1,)], [2.0], BigInt[3])
        @test RTFOOL.boltzmann(1.0, [2.0], [3]) == [1.0]
    end
    let space = StateSpace{2}([(1,1), (2,2)], [1.0, 1.0], BigInt[1,2])
        @test RTFOOL.boltzmann(2.0, space) == [1.0/3, 2.0/3]
    end
    let space = StateSpace{2}([(1,1),(2,2)], [1.0, 2.0], BigInt[2,1])
        @test RTFOOL.boltzmann(1.0, space) == [2e/(1+2e), 1/(1+2e)]
    end
    let space = StateSpace{2}([(1,1),(2,2)], [0.5, 1.5], BigInt[2,3])
        @test RTFOOL.boltzmann(0.5, space) ≈ [2e/(2e + 3*√e), 3*√e/(2e + 3*√e)]
    end
    let space = StateSpace{3}([(1,1,1), (1,1,2), (1,2,2)], [1,2,3], BigInt[1,2,3])
        @test RTFOOL.boltzmann(1/3, space) ≈
            [e^-(1/3), 2e^-(2/3), 3e^-1] / (e^-(1/3) + 2e^-(2/3) + 3e^-1)
    end
end


#=

Purpose: Use the Schur method to find the bounded solution of the
         discrete Lyapunov equation:

         X = B + A*X*A'.

Inputs:  A (n*n)
         B (n*n) = Symmetric positive semidefinite

Outputs: X (n*n) = Symmetric positive semidefinite solution.

Based on Octave code written by A.S. Hodel.

=#

function dlyap{T<:AbstractFloat}(a::Array{T,2},b::Array{T,2})

  n = size(a,1)
  x = zeros(n,n)
  j = n

  (s,u) = schur(a)
  b = u'b*u

  while j > 0
    j1 = j
    if j == 1
      block = 1
    elseif s[j,j-1] > 2*eps(T)*norm(s)
      block = 2
      j -= 1
    else
      block = 1
    end
    lhs = kron(s[j:j1,j:j1],s) - I # I = eye(block*n)
    rhs = vec(b[:,j:j1])
    if j1 < n
      rhs2 = s*(x[:,(j1+1):n]*s[j:j1,(j1+1):n]')
      rhs += vec(rhs2)
    end
    w = -lhs\rhs
    x[:,j] = w[1:n]
    if block == 2
      x[:,j1] = w[(n+1):block*n]
    end
    j -= 1
  end

  x = u*x*u'

  return x

end


using Plots
using Optim

# Calculate dose-response
# flux reaction:
function reaction(conc,Km,Vmax)
    10 .^Vmax.*conc./(10 .^Km .+ conc)
end

Km = 0
Vmax = 1
conc = collect(0:.01:2)

flux = reaction(conc, Km, Vmax)
plot(conc,flux)

# simulate data with 7% std
conc_sim = [.1,.5,1,2]
flux_sim = reaction(conc_sim,  Km, Vmax)
flux_sim = randn(length(flux_sim))*0.07.*flux_sim+flux_sim
std_sim = flux_sim * .07

scatter!(conc_sim,flux_sim,yerror=std_sim,grid=false,linealpha = 0)

# negative log-likelihood for simulated data
function L(conc,flux,std,p)
    X = reaction(conc,p[1],p[2])
    LL = -sum(-length(X)*log(2π)/2 .- length(X).*log.(abs.(std)) .- (sum((flux .- X).^2) ./ (2std.^2)))
end

# maximum likelihood fit
init = [1.0,1.0]        # start parameters Km,Vmax
optimum = optimize(p -> L(conc_sim,flux_sim,std_sim,p),init,Newton())

### We will continue with the profile likelihood on July, 10th


include("../algoritmos-busqueda/generar-vecino.jl")
include("../learner/euclidean-1-NN.jl")
include("../algoritmos-busqueda/busqueda-local.jl")
include("../utils/funcion-objetivo.jl")

function HeuristicaEnfriamiento(F_w, F_w_nuevo, T)
    return rand() < exp((F_w - F_w_nuevo)/ (T*1.380649*10^(-23)))
end

"""
    Expresón para pasar avanzar en la construcción de la T
"""
macro SiguienteT(ß, T)
    quote
        $(esc(T))  / (1 + $(esc(ß))*$(esc(T)))
    end
end

"""
    Expresión para calcular la beta
"""
macro calculaß(T, T_f, M)
    quote
        ($(esc(T)) - $(esc(T_f)))/($(esc(M))*$(esc(T)) *$(esc(T_f)))
    end
end

"""
function EnfriamientoSimulado(
    evaluaciones_maximas_funcion_objetivo::Int,
    numero_atributos::Int,
    umbral_atributo::Real,
    F,  # fitness
    probabilidad_aceptar_peor = 0.3,
    temperatura_final = 10^(-3),
    maximo_vecinos,
    maximo_vecinos_aceptados
    )::Vector{<:Real}
Devuelve mejor vector de peso encontrando de acorde
al algoritmo de enfriamiento simulado
"""
function EnfriamientoSimulado(
    evaluaciones_maximas_funcion_objetivo::Int,
    numero_atributos::Int,
    umbral_atributo::Real,
    F,  # fitness
    probabilidad_aceptar_peor,
    temperatura_final,
    maximo_vecinos,
    maximo_vecinos_aceptados
    )::Vector{<:Real}
    µ = 0.3
    M = 15000/maximo_vecinos

    # inicializaciones
    w = rand(numero_atributos) # vector de tamaño atributo uniformement inicializado
    # Ponemos a cero atributos menores que umbral atributo
    w = map(x-> (x<umbral_atributo) ? 0 : x , w)
    w_mejor = w
    F_w = F(w)
    F_w_mejor = F_w

    vecinos_aceptados = 1
    evaluaciones = 0
    primera_iteracion = true

    T = µ * F_w / -log(probabilidad_aceptar_peor)
    ß = (T - temperatura_final) / (M * T * temperatura_final)
    while ((evaluaciones < evaluaciones_maximas_funcion_objetivo
        && vecinos_aceptados > 0
        &&
        T <= temperatura_final
    )
        ||
        primera_iteracion
    )
        #reseteamos valores
        vecinos_aceptados = 0
        vecinos_generados = 0
        primera_iteracion = false

        while( vecinos_generados < maximo_vecinos &&
            vecinos_aceptados < maximo_vecinos_aceptados)
            w_n = GenNeighbourhood(w, 0, 1)
            F_w_vecino = F(w_n)
            vecinos_generados += 1
            ß = @calculaß(T, temperatura_final, M)

            if(F_w < F_w_vecino || HeuristicaEnfriamiento(F_w, F_w_vecino, T))
                vecinos_aceptados += 1
                w = w_n
                F_w = F_w_vecino
                if F_w_vecino > F_w_mejor
                    w_mejor = w_n
                    F_w_mejor = F_w_vecino
                end
            end
        end
        evaluaciones += 1
        T = @SiguienteT(ß, T)
    end
    return w_mejor
end

"""
    BL_LearnerOneNN(data::Matrix{<:Real}, labels::Vector{<:Real})
Devuelve un clasificado 1-NN refinado por unos pesos en la distancia euclídea.
Ese vector de pesos se ha obtenido a partir de un algoritmo de búsqueda local
"""
function EnfriamientoSimulado_LearnerOneNN(data::Matrix{<:Real}, labels)
    # Creamos función objetivo
    a = 0.5
    umbral_tasa_reduccion = 0.1
    F = CrearFuncionObjetivo(data, labels, a, umbral_tasa_reduccion )

    numero_atributos = size(data)[2] # los atributos son las columnas
    # Criterios de acorde a los requisitos
    probabilidad_aceptar_peor = 0.3
    temperatura_final = 10^(-3)
    numero_maximo_vecinos_sin_mejora = 20 * numero_atributos
    evaluaciones_maximas_funcion_objetivo = 15000
    maximo_vecinos = 10*numero_atributos
    maximo_vecinos_aceptados = 0.1*maximo_vecinos
    # Buscamos pesos de acorde a parámetros calculados
    w = EnfriamientoSimulado(
        evaluaciones_maximas_funcion_objetivo,
        numero_atributos,
        umbral_tasa_reduccion,
        F,  # fitness
        probabilidad_aceptar_peor,
        temperatura_final,
        maximo_vecinos,
        maximo_vecinos_aceptados
    )

    return WeightedLearnerEuclideanOneNN(w, data, labels), F(w), w
end


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export LogWeight

"""
`LogWeight(x)`

| Set                                 |     ``\\oplus``      |  ``\\otimes``  | ``\\bar{0}`` | ``\\bar{1}`` |
|:-----------------------------------:|:--------------------:|:--------------:|:------------:|:------------:|
|``\\mathbb{R}\\cup\\{\\pm\\infty\\}``|``\\log(e^{x}+e^{y})``|     ``+``      |``-\\infty``  |   ``0``      |
"""
struct LogWeight{T <: AbstractFloat} <: Semiring
  x::T
end

LogWeight(x::Number) = LogWeight(float(x))

zero(::Type{LogWeight{T}}) where T = LogWeight{T}(T(-Inf))
one(::Type{LogWeight{T}}) where T = LogWeight{T}(zero(T))

*(a::LogWeight{T}, b::LogWeight{T}) where {T <: AbstractFloat} = LogWeight{T}(a.x + b.x)
+(a::LogWeight{T}, b::LogWeight{T}) where {T <: AbstractFloat} = LogWeight{T}(logadd(a.x,b.x))
/(a::LogWeight{T}, b::LogWeight{T}) where {T <: AbstractFloat} = LogWeight{T}(a.x - b.x)
reverse(a::LogWeight) = a
function logadd(y::T, x::T) where {T <: AbstractFloat}
  if isinf(x) return y end
  if isinf(y) return x end
  if x < y
    diff = x-y
    x = y
  else
    diff = y-x
  end
  return x + log1p(exp(diff))
end

# parsing
parse(::Type{S},str) where {T, S <: LogWeight{T}} = S(parse(T,str))

#properties
iscommulative(::Type{W}) where {W <: LogWeight} = true
isleft(::Type{W}) where {W <: LogWeight} = true
isright(::Type{W}) where {W <: LogWeight}= true
isweaklydivisible(::Type{W}) where {W <: LogWeight}= true
iscomplete(::Type{W}) where {W <: LogWeight}= true


"""
    cheby1(n, x)

Compute the Chebyshev polynomial of the first kind of order n at point x.

This family of Chebyshev polynomials are a special case of the Jacobi polynomials
with α, β = -1/2.
"""
function cheby1(n, x)
    return Jacobi.jacobi(n, -0.5, -0.5, x)
end

"""
    cheby2(n, x)

Compute the Chebyshev polynomial of the second kind of order n at point x.

This family of Chebyshev polynomials are a special case of the Jacobi polynomials
with α, β = 1/2.
"""
function cheby2(n, x)
    return Jacobi.jacobi(n, 0.5, 0.5, x)
end


#=
gauss2.jl
=#

#using ImagePhantoms #: Object, Object2d

export Gauss2
export phantom, radon, spectrum, fwhm2spread

"""
    Gauss2 <: AbstractShape2
"""
struct Gauss2 <: AbstractShape2 end

# constructors

"""
    Gauss2(cx, cy, wx, wy=wx, ϕ=0, value::Number=1)
    Gauss2(center::NTuple{2,RealU}, width::NTuple{2,RealU}=(1,1), ϕ::RealU=0, v=1)
    Gauss2([6-vector])
    Gauss2(w, v=1) (isotropic of width `w`)
Construct `Gauss2` object from parameters;
here `width` = FWHM (full-width at half-maximum).

In 1D, the formula is `g(x) = exp(-π ((x - cx) / sx)^2)`
where `sx = fwhm2spread(w) = w * sqrt(π / log(16))`,
which, for `cx=0`,  has 1D FT `G(ν) = sx^2 exp(π (sx νx)^2)`.
"""
function Gauss2(
    cx::RealU,
    cy::RealU,
    wx::RealU,
    wy::RealU = wx,
    ϕ::RealU = 0,
    value::Number = 1,
)
    (cx, cy, wx, wy) = promote(cx, cy, wx, wy)
    Object(Gauss2(), (cx,cy), (wx,wy), ϕ, value)
end

function Gauss2(
    center::NTuple{2,RealU},
    width::NTuple{2,RealU} = (1,1) .* oneunit(center[1]),
    ϕ::RealU = 0,
    value::Number = 1,
)
    Gauss2(center..., width..., ϕ, value)
end

function Gauss2(v::AbstractVector{<:Number})
    length(v) == 6 || throw(ArgumentError("$v wrong length"))
    Gauss2(v...)
end

Gauss2(w::RealU, v::Number = 1) = Gauss2((zero(w),zero(w)), (w,w), 0, v)

# helper

"""
    s = fwhm2spread(w)
Convert FWHM `w` to equivalent Gaussian spread `s` for ``\\exp(-π (x/s)^2)``.
`exp(-π (fwhm/2/s)^2) = 1/2` means
`fwhm/2/s) = sqrt(log(2)/π)`
`2s/fwhm = sqrt(π/log(2))`
`2s = fwhm * sqrt(π/log(2))`
`s = fwhm * sqrt(π / log(16))`
"""
@inline fwhm2spread(w) = w * sqrt(π / log(16))

# methods

"""
    phantom(ob::Object2d{Gauss2})
Returns function of `(x,y)` for making image.
"""
function phantom(ob::Object2d{Gauss2})
    return (x,y) -> ob.value * # trick due to fwhm to spread scaling:
        exp(-π * sum(abs2.(coords(ob, x, y))) / fwhm2spread(1)^2)
end

"""
    radon_gauss2(r, ϕ, cx, cy, wx, wy, θ)
Radon transform at `(r,ϕ)` of 2D Gaussian.
"""
function radon_gauss2(r, ϕ, cx, cy, wx, wy, θ)
    (sx, sy) = fwhm2spread.((wx, wy))
    (sinϕ, cosϕ) = sincos(ϕ)
    r -= cx * cosϕ + cy * sinϕ # Radon translation property
    (sinϕ, cosϕ) = sincos(ϕ - θ) # Radon rotation property
    s = sqrt(abs2(sx * cosϕ) + abs2(sy * sinϕ)) # by Fourier-slice Thm.
    return sx * sy / s * exp(-π * abs2(r / s))
end

"""
    radon(ob::Object2d{Gauss2})
Returns function of `(r,ϕ)` for making a sinogram.
"""
radon(ob::Object2d{Gauss2}) = (r,ϕ) -> ob.value *
    radon_gauss2(r, ϕ, ob.center..., ob.width..., ob.angle[1])

function spectrum_gauss2(fx, fy, cx, cy, wx, wy, θ)
    (sx, sy) = fwhm2spread.((wx, wy))
    (kx, ky) = rotate2d(fx, fy, θ) # rotate first, then translate
    return sx * exp(-2im*π*fx*cx) *
           sy * exp(-2im*π*fy*cy) * exp(-π * (abs2(sx*kx) + abs2(sy*ky)))
end

"""
    spectrum(ob::Object2d{Gauss2})
Returns function of ``(f_x,f_y)`` for the spectrum (2D Fourier transform).
"""
spectrum(ob::Object2d{Gauss2}) = (fx,fy) -> ob.value *
    spectrum_gauss2(fx, fy, ob.center..., ob.width..., ob.angle[1])


export cwt

"""
    cwt(signal, Λ)

Calculate the continuous wavelet transform of a signal given filter bank Λ.

# Arguments
 - `signal::Vector{Float64}`: Time series.
 - `Λ::Array{Float64, 2}`: Fourier transformed filters arranged by [frequency, scale].
"""
function cwt(signal::Vector{Float64}, Λ::Array{Float64,2})
    signal_ft = fft(signal, 1)
    W = ifft(signal_ft.*Λ, 1)
    return W
end

"""
    cwt(signal, Λ, ω0, δt, δj)

Convenience method for cwt that constructs the filter bank Λ given a signal.

# Arguments
 - `signal::Vector{Float64}`: Time series.
 - `F::FilterBank1dBlock`: Fourier transformed filters and metadata
"""
function cwt(signal::Vector{Float64}, F::FilterBank1dBlock)
    W = cwt(signal, F.Λ)
    return W
end


@doc raw"""
    rating_constraint!(canonical::Canonical,
                            rating_data::Vector{Tuple{String, Float64}},
                            cons_name::Symbol,
                            var_names::Tuple{Symbol, Symbol})

Constructs constraint from rating data and related variable tuple.

# Constraint

``` var1[r[1], t]^2 + var2[r[1], t]^2 <= r[2]^2 ```

where r in rating data and t in time steps.

# LaTeX

`` x_1^2 + x_2^2 \leq r^2 ``

# Arguments
* canonical::Canonical : the canonical model built in PowerSimulations
* rating_data::Vector{Tuple{String, Float64}} : rating data name (1) and value (2)
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol} : the names of the variables
- : var_names[1] : var1
- : var_names[2] : var2
"""
function rating_constraint!(canonical::Canonical,
                            rating_data::Vector{Tuple{String, Float64}},
                            cons_name::Symbol,
                            var_names::Tuple{Symbol, Symbol})
    time_steps = model_time_steps(canonical)
    var1 = get_variable(canonical, var_names[1])
    var2 = get_variable(canonical, var_names[2])
    add_cons_container!(canonical, cons_name, (r[1] for r in rating_data), time_steps)
    constraint = get_constraint(canonical, cons_name)

    for r in rating_data
        for t in time_steps
          constraint[r[1], t] = JuMP.@constraint(canonical.JuMPmodel, var1[r[1], t]^2 + var2[r[1], t]^2 <= r[2]^2)
        end
    end

    return

end


# Hessian Eigenmaps (HLLE)
# ---------------------------
# Hessian eigenmaps: Locally linear embedding techniques for high-dimensional data,
# D. Donoho and C. Grimes, Proc Natl Acad Sci U S A. 2003 May 13; 100(10): 5591–5596

import Combinatorics: combinations

#### HLLE type
struct HLLE{T <: AbstractFloat} <: SpectralResult
    k::Int
    λ::AbstractVector{T}
    proj::Projection{T}

    HLLE{T}(k::Int, λ::AbstractVector{T}, proj::Projection{T}) where T = new(k, λ, proj)
end

## properties
outdim(M::HLLE) = size(M.proj, 1)
projection(M::HLLE) = M.proj

eigvals(M::HLLE) = M.λ
neighbors(M::HLLE) = M.k

## show & dump
function show(io::IO, M::HLLE)
    print(io, "Hessian Eigenmaps(outdim = $(outdim(M)), neighbors = $(neighbors(M)))")
end

function dump(io::IO, M::HLLE)
    show(io, M)
    println(io, "eigenvalues: ")
    Base.showarray(io, transpose(M.λ), header=false, repr=false)
    println(io)
    println(io, "projection:")
    Base.showarray(io, M.proj, header=false, repr=false)
end

## interface functions
function transform(::Type{HLLE}, X::DenseMatrix{T}; d::Int=2, k::Int=12) where T<:AbstractFloat
    n = size(X, 2)

    # Identify neighbors
    D, I = find_nn(X, k)

    # Obtain tangent coordinates and develop Hessian estimator
    hs = round(Int, d*(d+1)/2)
    W = spzeros(hs*n,n)
    for i=1:n
        # re-center points in neighborhood
        μ = mean(X[:,I[:,i]], dims=2)
        N = X[:,I[:,i]] .- μ
        # calculate tangent coordinates
        #tc = svd(transpose(N)).U[:,1:d]
        tc = svd(N).V[:,1:d]

        # Develop Hessian estimator
        Yi = [ones(k) tc zeros(k,hs)]
        for ii=1:d
            Yi[:,d+ii+1] = tc[:,ii].^2
        end
        yi = 2(1+d)
        for (ii,jj) in combinations(1:d,2)
            Yi[:, yi] = tc[:, ii] .* tc[:, jj]
            yi += 1
        end
        F = qr(Yi)
        H = transpose(Matrix(F.Q)[:,d+2:end])
        W[(i-1)*hs .+ (1:hs),I[:,i]] = H
    end

    # decomposition
    λ, V = decompose(transpose(W)*W, d)
    return HLLE{T}(k, λ, transpose(V) .* sqrt(n))
end


###########################################################
# Load (old) functions
include("bils_chang_kim_2011/steady_state_equilibrium.jl")

###########################################################
# Import necessary packages
using BasisMatrices
using Optim
using Distributions
using Parameters
using LinearAlgebra
using Distributions
using Optim
using QuantEcon
using SpecialFunctions
using SparseArrays
using Plots
using LaTeXStrings
using ThreadTools

###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
# Model parameters
@with_kw mutable struct ModelParameters
    A::Float64              = 1.0                           # output productivity factor
    r::Float64              = 0.06                          # annual interest rate
    r_mon::Float64          = (1.0 + r)^(1.0/12.0) - 1.0    # monthly interest rate
    worker_surplus::Float64 = 0.5                           # worker's share of surplus
    β::Float64              = 0.9948                        # discount factor
    β_mon::Float64          = 1.0/(1.0 + r_mon)             # monthly depreciation rate
    γ::Float64              = 1.0                           # curvature of utility function
    λ::Float64              = 0.02                          # exogenous separation rate
    b::Float64              = 0.4                           # UI benefits
    B::Float64              = 0.15                          # utility from leisure
    κ::Float64              = 0.522                         # vacancy posting cost with log utility
    θ::Float64              = 1.0                           # labor market tightness, v/u, normalized to 1
    α::Float64              = 0.5                           # power term in the matching function
    η::Float64              = 0.3133                        # scale parameter in the matching function
    p_θ::Float64            = η*θ^α                         # unemployed worker matching rate
    q_θ::Float64            = η*θ^(α-1)                     # vacancy matching probability
    amin::Float64           = -6.0                          # minimum asset
    amax::Float64           = 120.0                         # maximum asset
    agrid::Vector{Float64}  = LinRange(amin, amax, 20)      # asset grid
    ρ_x::Float64            = 0.97                          # persistence of idiosyncratic productivity ln(x)
    σ_x::Float64            = 0.13                          # std. dev. of innovation to ln(x)
    π_x                     = TauchenApprox(9, ρ_x, σ_x).distribution           # probability weights on ln(x)
    xgrid                   = exp.(TauchenApprox(9, ρ_x, σ_x).state_values)     # match productivity grid
    spline_A                = SplineParams(agrid,0,1)
    spline_X                = SplineParams(xgrid,0,1)
    basis_A                 = Basis(spline_A)
    basis_AX                = Basis(spline_A,spline_X)
    nodes_A                 = nodes(basis_A)[1]
    nodes_AX                = nodes(basis_AX)[1]
end;

###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
# utility function
function utility(para::ModelParameters, c::Real)

    γ = para.γ # Store risk-aversion parameter

    # Define the utility mapping
    # condition on the input `c`
    if c > 0
        if γ == 1
          return log(c) # if σ = 1, CRRA utility is log
        else
          return (c^(1-γ)-1)/(1-γ) # if σ != 1, regular specification
        end
    else
        return -Inf # infinite disutility if c∈(-∞,0]
    end
end;

function utility(para::ModelParameters, cvec)

  u(c) = utility(para,c)

  return u.(cvec)
end;

###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
# Initialize interpolated worker value functions
function InitializeValues(para::ModelParameters)

    # Retrieve asset and productivity shock grids
    @unpack basis_A, basis_AX, nodes_A, nodes_AX = para

    # Create functions to map all (a,x) and (a) to 1
    # for both W(a,x) and U(a)
    g1d = x -> x[1]^0
    gvals = [g1d(nodes_AX[i,:]) for i in 1:size(nodes_A,1)]
    f2d = x -> x[1]^0 * x[2]^0
    fvals = [f2d(nodes_AX[i,:]) for i in 1:size(nodes_AX,1)]

    # Interpolate W(a,x) and U(a)
    W_init = Interpoland(basis_AX, fvals);
    U_init = Interpoland(basis_A, gvals)

    # Return bases for A and A×X
    # as well as interpolated initial W(a,x) and U(a)
    return W_init, U_init
end;

###########################################################
# Initialize interpolated worker value functions
W_init, U_init = InitializeValues(ModelParameters());

###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
# Initialize interpolated worker wage schedule
function InitializeWage(para::ModelParameters)

    # Retrieve asset and productivity shock grids
    # as well as the market tightness paramter θ
    @unpack θ, basis_AX, nodes_AX = para

    # Map all (a,x) to θ*0.9 = 0.9
    f2d = x -> x[1]^0 * x[2]^0 * θ * 0.9
    fvals = [f2d(nodes_AX[i,:]) for i in 1:size(nodes_AX,1)]

    # Interpolate wage schedule over A×X
    wage_init = Interpoland(basis_AX, fvals);

    # Return basis for A×X
    # and interpolated wage schedule
    return wage_init
end;

###########################################################
# Initialize interpolated wage schedule
wage_init = InitializeWage(ModelParameters());

###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
# Interpolated invariant productivity shock process density function
function ShockDensity(para::ModelParameters, xval)

    @unpack ρ_x, σ_x = para

    distr = LogNormal(0, σ_x/sqrt(1-ρ_x^2))
    density = pdf(distr, xval)

    return density
end

###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
# Single iteration of the Bellman map w interpolated objects
function WorkerBellmanMap(para::ModelParameters, wage, W_old, U_old)

    # Import parameters
    @unpack xgrid, π_x, β, b, B, p_θ, r, basis_A, basis_AX, nodes_A, nodes_AX = para

    # Shorthand for utility
    u(c) = utility(para,c)

    # Create empty W and U objects
    W_new           = similar(nodes_AX[:,1])
    U_new           = similar(nodes_A)
    emp_policy      = similar(W_new)
    unemp_policy    = similar(U_new)

    # Initialize index for U(a)
    k=0

    # Fill out W_new and U_new matrices
     Threads.@threads for i in 1:length(nodes_AX[:,1])
    #for i in 1:length(nodes_AX[:,1])

        # Create (a,x) vector
        ax = vec(nodes_AX[i,:])

        # Create Bellman mapping under employment
        c_emp(a′) = (1+r)*ax[1] + wage(ax) - a′
        obj_emp     = tmap(a′ -> (1)*(u(c_emp(a′)) + β * sum( π_x[i]*max.(W_old(vec([a′ xgrid[i]])), U_old(xgrid[i])) for i in 1:length(xgrid))), nodes_A)
        # Solve for maximum W(a,x)
        # and maximizing a′
        W_new[i], optimal_a′ = findmax(obj_emp)
        emp_policy[i] = nodes_A[optimal_a′]

        # If x equals E[x], then
        # start solving for max. U(a)
        # and maximizing a′
        if ax[2] ≈ 1.0
            # Update index for U(a)
            k = k + 1
            # Create Bellman mapping under unemployment
            c_unemp(a′) = (1+r)*ax[1] + b - a′
            obj_unemp   = tmap( a′ -> u(c_unemp(a′)) + B + β*(1-p_θ)*U_old(ax[1]) .+ β*(p_θ)*W_old(ax), nodes_A)
            # Solve for maximum W(a,x)
            # and maximizing a′
            U_new[k], optimal_a′ = findmax(obj_unemp)
            unemp_policy[k] = nodes_A[optimal_a′]
        end
    end

    # Interpolate W_new & U_new
    W_new = Interpoland(basis_AX, W_new)
    U_new = Interpoland(basis_A, U_new)

    return W_new, U_new, emp_policy, unemp_policy
end;

###########################################################
# Iterate Bellman map
para = ModelParameters()
para.agrid = LinRange(para.amin, para.amax, 18)
W_new, U_new, emp_policy, unemp_policy = WorkerBellmanMap(para, wage_init, W_init, U_init)

###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
# Value function iteration
function SolveWorkerBellman(para::ModelParameters, wage, W0, U0, ϵ=1e-5)

    @unpack nodes_A = para

    W_old = W0
    U_old = U0

    diff = 1.
    iter = 0

    while diff > ϵ
        iter = iter + 1
        W_new, U_new = WorkerBellmanMap(para, wage, W_old, U_old)
        diff = norm(U_new.(nodes_A)-U_old.(nodes_A))
        W_old = W_new
        U_old = U_new
        println("Iteration: $(iter); Norm: $(diff)")
    end

    W, U, emp_policy, unemp_policy = WorkerBellmanMap(para, wage, W_old, U_old)

    return W, U, emp_policy, unemp_policy
end;

###########################################################
# Iterate Bellman map
W_init, U_init = InitializeValues(ModelParameters());
wage_init = InitializeWage(ModelParameters());
para = ModelParameters()
para.agrid = LinRange(para.amin, para.amax, 100)
para.B = 0.0
W_new, U_new, emp_policy, unemp_policy = SolveWorkerBellman(para,
                                                            wage_init,
                                                            W_init,
                                                            U_init);

plot(a -> U_new(a), LinRange(-6.0,120.0,5000))

U_new(0)
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
###########################################################
#

using Profile

@profile InitializeValues(ModelParameters())
Profile.print()


using SymPy
if VERSION >= v"0.5.0-dev+7720"
    using Base.Test
else
    using BaseTestNext
    const Test = BaseTestNext
end

@testset "Matrix" begin
    ## matrices
    (x,) = @syms x
    A = [x 1; 1 x]
    b = [x, 2]

    ## These fail for older installations of SymPy
    @test simplify(det(A)) == x^2 - 1

    ## we use inverse for A[:inv]()
    inv(A) # aliased to use inverse
    @test simplify((A[:inv]() - inv(A))[1,1]) == 0

    adjoint(A)
    dual(A)
    cholesky(A)
    ## other functions, could wrap
    b = subs(A, x, 2)
    QRdecomposition(b)

    # is_lower, is_square, is_symmetric much slower than julia only counterparts. May deprecate, but for now they are here
    @test is_lower(A) == istril(A)
    @test is_square(A) == true
    @compat @test is_symmetric(A) == issymmetric(A)
    @test Set(eigvals(A)) == Set([x-1, x+1])

    A = [1 0 0; 0 1 0; 0 0 x]
    evecs = eigvecs(A)
    @test evecs[1] == [1, 0, 0]

    ##
    M = Sym[1 2 3; 3 6 2; 2 0 1]
    q,r = QRdecomposition(M)
    @test (q * r - M)[1,1] == 0
    if VERSION >= v"0.5"
        L = [Sym[2,3,5], Sym[3,6,2], Sym[8,3,6]]
        out = GramSchmidt(L)
    end
    A = Sym[4 3; 6 3]
    L, U, _ = LUdecomposition(A)
    @test L == Sym[1 0; 3//2 1]

end


using LinearAlgebra

function Cn(v, n::Integer)
    θ = 2*π / n
    return rotation_matrix(v, θ)
end

"""
    Molecules.rotation_matrix(v, θ::AbstractFloat)

Returns a rotation matrix by θ about a rotation axis v
"""
function rotation_matrix(V, θ::AbstractFloat)
    cosθ = cos(θ)
    sinθ = sin(θ)
    v = normalize(V)
    a = [1,2,3]
    O = zeros(eltype(v), (3,3))
    O .+= 1 - cosθ
    for i = 1:3, j = 1:3
        if i == j
            O[i,i] *= v[i]^2
            O[i,i] += cosθ
        else
            O[i,j] *= v[i]*v[j]
            b = [i,j]
            C = [i for i in a if i ∉ b][1]
            if i < j
                O[i,j] += (-1)^(i+j) * v[C]*sinθ
            else
                O[i,j] += (-1)^(i+j-1) * v[C]*sinθ
            end
        end
    end
    return O
end

function rotation_matrixd(v, θ::AbstractFloat)
    r = deg2rad(θ)
    return rotation_matrix(v, r)
end

"""
    Molecules.reflection_matrix(v::Vector)

Returns a reflection matrix through the plane with normal vector v
"""
function reflection_matrix(V)
    O = zeros(eltype(V), (3,3))
    v = normalize(V)
    for i = 1:3, j = i:3
        if i == j
            O[i,i] = 1 - 2*v[i]^2
        else
            O[i,j] = -2 * v[i] * v[j]
            O[j,i] = O[i,j]
        end
    end
    return O
end

function σ(v)
    return reflection_matrix(v)
end

"""
    Molecules.Sn(v::Vector, n::int)

Returns Sn improper rotation about vector v
"""
function Sn(v, n)
    return Cn(v, n) * σ(v)
end

"""
    Molecules.inversion_matrix()

Returns a diagonal matrix with -1 along the diagonal
"""
function inversion_matrix()
    a = zeros((3,3))
    for i = 1:3
        a[i,i] = -1
    end
    return a
end

function i()
    return inversion_matrix()
end


      function percentile(perc::Float64,p::Array{Float64,2})

      sig::Float64=0.0
      sig0::Float64=0.0
      sig1::Float64=0.0
      dsig::Float64=0.0
      pmin::Float64=0.0
      delta::Float64=0.0
      nvec::Int64= size(p,1)*size(p,2)
      nsig::Int64= Int64(floor(nvec/2))
      itnwt::Int64=0

      phi::Float64=0.0
      phi0::Float64=0.0
      phi1::Float64=0.0
      dphi::Float64=0.0
      TOL::Float64 = 0.000001
      #
      pmin=p[1,1]
      delta=pmin
      for i2=1:size(p,2)
          for i1=1:size(p,1)
              delta=max(delta,p[i1,i2])
              pmin=min(pmin,p[i1,i2])
          end
      end
      delta=delta-pmin
      if delta <= 0.0
          println("Null interval ")
          return nothing
      end

      sig0=0.0
      phi0=0.0
      for i2=1:size(p,2)
          for i1=1:size(p,1)
              if (p[i1,i2]-pmin) >= sig0
                 phi0 += 1.0
              end
          end
      end
      phi0 = perc - phi0/Float64(nvec)
      sig1=delta
      phi1=0.0
      for i2=1:size(p,2)
          for i1=1:size(p,1)
              if (p[i1,i2]-pmin) >= sig1
                 phi1+=1.0
              end
          end
      end

      phi1=perc - phi1/Float64(nvec)
      dsig=-phi0*(sig1-sig0)/(phi1-phi0)
      sig = sig0+dsig

      for itnwt=1:99
      phi=0.0
      for i2=1:size(p,2)
          for i1=1:size(p,1)
              if (p[i1,i2]-pmin) >= sig
                 phi=phi+1.0
              end
          end
      end
      phi = perc - phi/Float64(nvec)

      if phi*phi0 < 0.0
         sig1 = sig
         phi1 = phi
      else
         sig0 = sig
         phi0 = phi
      end
      dsig=-phi0*(sig1-sig0)/(phi1-phi0)
      sig = sig0+dsig

      if  abs(phi) < TOL || abs(dsig) < 1.e-16*delta
          println("iter = ",itnwt)
          break
      end
      end

      sig=sig+pmin
      println("sig = ",sig, " phi = ",phi, " dsig = ",dsig)

      return nothing
      end # function


include("getdigit_n.jl")
include("myprintf.jl")
using Printf
using .GetDigit_n
using .MyPrintf

const DIGIT = 100
const MAXN = 500

function viete(n)
    numerator = sqrt(BigFloat(2))
    result = BigFloat(1)

    for i = 1:n
        result *= (numerator / BigFloat(2))
        numerator = sqrt(BigFloat(2) + numerator)
    end

    return BigFloat(2) / result
end

function main()
    n = GetDigit_n.getn(viete, MAXN, DIGIT, true)
    precision = GetDigit_n.getdigit(viete, n, DIGIT, true)
    precision = GetDigit_n.getdigit_2(viete, n, DIGIT, precision)

    setprecision(precision)
    MyPrintf.myprintf("π（計算値） = ", DIGIT, viete(n))
    MyPrintf.myprintf("π（厳密値） = ", DIGIT, BigFloat(π))
    @printf("n = %d, 計算%d桁で、%d桁求まりました\n", n, precision, DIGIT)
end

main()


function u_test(x, y, σ)
    exp( - ( (x-π)^2 + (y-π)^2 ) / (2σ^2) )
end

function testTrapezoidalQuadrature()
    M = 64
    N = 64

    σ = 0.5

    grid = Grid2d(M,N)
    u    = get_field(grid)

    for i in 1:M
        for j in 1:N
            u[i,j]  = u_test(grid.x[i], grid.y[j], σ)
        end
    end

    uint = trapezoidal_quadrature(u, u, grid)

    @test uint - π*σ^2 ≈ zero(eltype(u)) atol=1E-14
end

testTrapezoidalQuadrature()