# ##############################################################################
#
# State space representation of Max-Plus linear systems.
# Creation of Max-Plus dynamical linear systems in implicit state form:
#    X(n) = D.X(n) + A.X(n-1) + B.U(n),
#    Y(n) = C.X(n)
#
# ##############################################################################

# Utility function: return an unified size for a matrix or a vector
size2(A::SparseMatrixCSC) = (size(A, 1), size(A, 2))
size2(A::Array) = (size(A, 1), size(A, 2))
size2(A::Vector) = (size(A, 1), size(A, 2))

# ==============================================================================
# Implicit dynamic linear Max-Plus system.

"""
    MPSysLin(A::AbsMatMP, B::AbsMatMP, C::AbsMatMP [, D::AbsMatMP, [x0::AbsMatMP]])

Structure for state space representation of Max-Plus linear systems.
Creation of max-plus dynamical linear systems in implicit state form:
```
    X(n) = D ⨂ X(n) ⨁ A ⨂ X(n-1) ⨁ B ⨂ U(n),
    Y(n) = C ⨂ X(n)
```

`MPSysLin` is the equivalent to ScicosLab function `mpsyslin`. It accepts dense
or sparse Max-Plus arrays. `D` and `x0` can be omited: they will be filled
with Max-Plus zeros `ε`.

# Arguments
- `A` shall be squared.
- `B` shall has its row dimension in accordance with dimensions of `A`.
- `C` shall has its column dimension in accordance with dimensions of `A`.
- `D` shall has its dimension in accordance with dimensions of `A`.
- `x0` shall has its dimensions in accordance with dimensions of `A`.

Sizes are checked and an error is thrown during the compilation.

Elements of the system can access directly ie. `S.A`, `S.B`, `S.C`, `S.D`,
`S.x0`.

# Examples
```julia-repl
julia> S = MPSysLin(MP([1.0 2 3;  4 5 6;  7 8 9]), # Max-Plus matrix A
                    MP([    0.0;      0;      0]), # Max-Plus column vector B
                    MP([    0.0       0       0]), # Max-Plus row vector C
                    mpeye(3, 3),                   # Max-Plus matrix D
                    full(zeros(MP, 3, 1)))         # Max-Plus column vector x0

Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 3×3 Max-Plus dense matrix:
  0   .   .
  .   0   .
  .   .   0

A = 3×3 Max-Plus dense matrix:
  1   2   3
  4   5   6
  7   8   9

B = 3-element Max-Plus vector:
  0
  0
  0

C = 1×3 Max-Plus dense matrix:
  0   0   0

x0 = 3-element Max-Plus vector:
  .
  .
  .

julia> typeof(S)
MPSysLin

julia> S.A
3×3 Max-Plus dense matrix:
  1   2   3
  4   5   6
  7   8   9
```
"""
struct MPSysLin
    A::AbsMatMP
    B::AbsMatMP
    C::AbsMatMP
    D::AbsMatMP
    x0::AbsMatMP

    # Constructor with Max-Plus matrices: A, B, C, D and x0
    function MPSysLin(A::AbsMatMP, B::AbsMatMP, C::AbsMatMP, D::AbsMatMP, x0::AbsMatMP)
        # show(stdout, (size2(A),size2(B),size2(C),size2(D),size2(x0)))

        (ma,na) = size2(A)
        (ma != na) &&
            error("Matrix A shall be squared")

        (mb,nb) = size2(B)
        (mb != na) &&
            error("The row dimension of B: $(mb) is not in accordance with dimensions of A: $(na)")

        (mc,nc) = size2(C)
        (nc != na) &&
            error("The column dimension of C: $(nc) is not in accordance with dimensions of A: $(na)")

        (mx0,nx0) = size2(x0)
        (mx0 != na) &&
            error("Dimensions of x0: $(mx0) do not in accordance with dimensions of A: $(na)")

        (md,nd) = size2(D)
        ((md != na) || (nd != na)) &&
            error("The column dimension of D: $(md)×$(nd) is not in accordance with dimensions of A: $(na)")

        new(A, B, C, D, x0)
    end

    # Constructor with Max-Plus matrices: A, B, C, D (and implicit x0 set as zeros)
    function MPSysLin(A::AbsMatMP, B::AbsMatMP, C::AbsMatMP, D::AbsMatMP)
        na = size(A,2)
        MPSysLin(A, B, C, D, spzeros(MP, na, 1))
    end

    # Constructor with Max-Plus matrices: A, B, C (and implicit D and x0 set as zeros)
    function MPSysLin(A::AbsMatMP, B::AbsMatMP, C::AbsMatMP)
        na = size(A,2)
        MPSysLin(A, B, C, spzeros(MP, na, na), spzeros(MP, na, 1))
    end
end # MPSysLin

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

"""
    Base.:(==)(x::MPSysLin, y::MPSysLin)

Test whether two Max-Plus linear system are equal.
```
"""
function Base.:(==)(x::MPSysLin, y::MPSysLin)
    (x.A == y.A) && (x.B == y.B) && (x.C == y.C) && (x.D == y.D) && (x.x0 == y.x0)
end

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

"""
    mpshow(io::IO, S::MPSysLin)

Base function for displaying a Max-Plus linear system.

# Examples
```julia-repl
julia>
```
"""
function mpshow(io::IO, S::MPSysLin)
    (@printf io "Implicit dynamic linear Max-Plus system:\n")
    (@printf io "  x(n) = D*x(n) + A*x(n-1) + B*u(n)\n  y(n) = C*x(n)\n  x(0) = x0\n\nwith:")
    (@printf io "\nD = "); mpshow(io, S.D)
    (@printf io "\nA = "); mpshow(io, S.A)
    (@printf io "\nB = "); mpshow(io, S.B)
    (@printf io "\nC = "); mpshow(io, S.C)
    (@printf io "\nx0 = "); mpshow(io, S.x0)
end

"""
    show(io::IO, S::MPSysLin)

Display a Max-Plus linear system.

# Examples
```julia-repl
julia>
```
"""
Base.show(io::IO, S::MPSysLin) = mpshow(io, S)

"""
    show(io::IO, ::MIME"text/plain", S::MPSysLin)

Display a Max-Plus linear system.

# Examples
```julia-repl
julia>
```
"""
Base.show(io::IO, ::MIME"text/plain", S::MPSysLin) = mpshow(io, S)

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

"""
    LaTeX(S::MPSysLin)

Return the LaTeX code as string from the given Max-Plus linear system.

# Examples
```julia-repl
julia>
```
"""
function LaTeX(S::MPSysLin)
    "\\left\\{\\begin{array}{lcl}\nx_n & = & " *
    LaTeX(S.D) * " x_n \\oplus " *
    LaTeX(S.A) * " x_{n-1} \\oplus " *
    LaTeX(S.B) * " u_n\\\\ y_n & = & " *
    LaTeX(S.C) * " x_n\\\\ x_0 & = & " *
    LaTeX(S.x0) * "\\end{array}\\right."
end

"""
    LaTeX(io::IO, S::MPSysLin)

Display the LaTeX code from the given Max-Plus linear system.

# Examples
```julia-repl
julia>
```
"""
function LaTeX(io::IO, S::MPSysLin)
    (@printf io "\\left\\{\\begin{array}{lcl}\n")
    (@printf io "x_n & = & ")
    LaTeX(io, S.D)
    (@printf io " x_n \\oplus ")
    LaTeX(io, S.A)
    (@printf io " x_{n-1} \\oplus ")
    LaTeX(io, S.B)
    (@printf io " u_n\\\\")
    (@printf io "y_n & = & ")
    LaTeX(io, S.C)
    (@printf io " x_n\\\\")
    (@printf io "x_0 & = & ")
    LaTeX(io, S.x0)
    (@printf io "\\end{array}\\right.")
end

Base.show(io::IO, ::MIME"text/latex", S::MPSysLin) = LaTeX(io, S)

# ==============================================================================
# Parallel composition.
# From ScicosLab file: %mpls_a_mpls.sci

"""
    Base.:(+)(y::MPSysLin, x::MPSysLin)

Parallel composition of two Max-Plus linear systems.

# Examples
```julia-repl
julia> S1 = MPSysLin(MP([1.0 2 3;  4 5 6;  7 8 9]),
                     MP([    1.0;      2;      3]),
                     MP([    4.0       5       6]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S2 = MPSysLin(MP([10.0 11 12;  13 14 15;  16 17 18]),
                     MP([    0.0;      0;      0]),
                     MP([    0.0       0       0]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S1 + S2
Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 6×6 Max-Plus dense matrix:
  0   .   .   .   .   .
  .   0   .   .   .   .
  .   .   0   .   .   .
  .   .   .   0   .   .
  .   .   .   .   0   .
  .   .   .   .   .   0

A = 6×6 Max-Plus dense matrix:
  1   2   3    .    .    .
  4   5   6    .    .    .
  7   8   9    .    .    .
  .   .   .   10   11   12
  .   .   .   13   14   15
  .   .   .   16   17   18

B = 6-element Max-Plus vector:
  1
  2
  3
  0
  0
  0

C = 1×6 Max-Plus dense matrix:
  4   5   6   0   0   0

x0 = 6-element Max-Plus vector:
  .
  .
  .
  .
  .
  .
```
"""
function Base.:(+)(x::MPSysLin, y::MPSysLin)
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    MPSysLin([x.A spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.A],
             [x.B; y.B],
             [x.C y.C],
             [x.D spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# Series composition.
# From ScicosLab file: %mpls_m_mpls.sci

"""
    Base.:(*)(y::MPSysLin, x::MPSysLin)

Series composition of two Max-Plus linear systems.

# Examples
```julia-repl
julia> S1 = MPSysLin(MP([1.0 2 3;  4 5 6;  7 8 9]),
                     MP([    1.0;      2;      3]),
                     MP([    4.0       5       6]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S2 = MPSysLin(MP([10.0 11 12;  13 14 15;  16 17 18]),
                     MP([    0.0;      0;      0]),
                     MP([    0.0       0       0]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S1 * S2
TODO
```
"""
function Base.:(*)(y::MPSysLin, x::MPSysLin)
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nb1 = size(x.B, 2)
    nc2 = size(y.C, 1)
    MPSysLin([x.A spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.A],
             [x.B; spzeros(MP, n2, nb1)],
             [spzeros(MP, nc2, n1) y.C],
             [x.D spzeros(MP, n1, n2); y.B * x.C y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# Diagonal composition.
# From ScicosLab file: %mpls_g_mpls.sci

"""
    Base.:(|)(y::MPSysLin, x::MPSysLin)

Diagonal composition of two Max-Plus linear systems.

# Examples
```julia-repl
julia> S1 = MPSysLin(MP([1.0 2 3;  4 5 6;  7 8 9]),
                     MP([    1.0;      2;      3]),
                     MP([    4.0       5       6]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S2 = MPSysLin(MP([10.0 11 12;  13 14 15;  16 17 18]),
                     MP([    0.0;      0;      0]),
                     MP([    0.0       0       0]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S1 | S2

Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 6×6 Matrix{MP{Float64}}:
  0   .   .   .   .   .
  .   0   .   .   .   .
  .   .   0   .   .   .
  .   .   .   0   .   .
  .   .   .   .   0   .
  .   .   .   .   .   0

A = 6×6 Matrix{MP{Float64}}:
  1   2   3    .    .    .
  4   5   6    .    .    .
  7   8   9    .    .    .
  .   .   .   10   11   12
  .   .   .   13   14   15
  .   .   .   16   17   18

B = 6×2 Matrix{MP{Float64}}:
  1   .
  2   .
  3   .
  .   0
  .   0
  .   0

C = 2×6 Matrix{MP{Float64}}:
  4   5   6   .   .   .
  .   .   .   0   0   0

x0 = 6-element Vector{MP{Float64}}:
  .
  .
  .
  .
  .
  .
```
"""
function Base.:(|)(x::MPSysLin, y::MPSysLin)
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nb1 = size(x.B, 2)
    nb2 = size(y.B, 2)
    nc1 = size(x.C, 1)
    nc2 = size(y.C, 1)
    MPSysLin([x.A spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.A],
             [x.B spzeros(MP, n1, nb2); spzeros(MP, n2, nb1) y.B],
             [x.C spzeros(MP, nc1, n2); spzeros(MP, nc2, n1) y.C],
             [x.D spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# computes [S1;S2] that is : inputs in common, concatenation of outputs.
# From ScicosLab file: %mpls_f_mpls.sci

"""
    Base.:vcat(x::MPSysLin, y::MPSysLin)

Composition of two Max-Plus linear systems: inputs in common, concatenation of
outputs.

# Examples
```julia-repl
julia> S1 = MPSysLin(MP([1.0 2 3;  4 5 6;  7 8 9]),
                     MP([    1.0;      2;      3]),
                     MP([    4.0       5       6]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S2 = MPSysLin(MP([10.0 11 12;  13 14 15;  16 17 18]),
                     MP([    0.0;      0;      0]),
                     MP([    0.0       0       0]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> [S1 ; S2]
Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 6×6 Matrix{MP{Float64}}:
  0   .   .   .   .   .
  .   0   .   .   .   .
  .   .   0   .   .   .
  .   .   .   0   .   .
  .   .   .   .   0   .
  .   .   .   .   .   0

A = 6×6 Matrix{MP{Float64}}:
  1   2   3    .    .    .
  4   5   6    .    .    .
  7   8   9    .    .    .
  .   .   .   10   11   12
  .   .   .   13   14   15
  .   .   .   16   17   18

B = 6-element Vector{MP{Float64}}:
  1
  2
  3
  0
  0
  0

C = 2×6 Matrix{MP{Float64}}:
  4   5   6   .   .   .
  .   .   .   0   0   0

x0 = 6-element Vector{MP{Float64}}:
  .
  .
  .
  .
  .
  .
```
"""
function Base.:vcat(x::MPSysLin, y::MPSysLin)
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nc1 = size(x.C, 1)
    nc2 = size(y.C, 1)
    MPSysLin([x.A spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.A],
             [x.B; y.B],
             [x.C spzeros(MP, nc1, n2); spzeros(MP, nc2, n1) y.C],
             [x.D spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# computes [S1,S2] that is concatenation of inputs, addition of outputs
# From ScicosLab file: %mpls_c_mpls.sci

"""
    Base.:hcat(x::MPSysLin, y::MPSysLin)

Composition of two Max-Plus linear systems: concatenation of inputs, addition of
outputs.

# Examples
```julia-repl
julia> S1 = MPSysLin(MP([1.0 2 3;  4 5 6;  7 8 9]),
                     MP([    1.0;      2;      3]),
                     MP([    4.0       5       6]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S2 = MPSysLin(MP([10.0 11 12;  13 14 15;  16 17 18]),
                     MP([    0.0;      0;      0]),
                     MP([    0.0       0       0]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> [S1 S2]

Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 6×6 Matrix{MP{Float64}}:
  0   .   .   .   .   .
  .   0   .   .   .   .
  .   .   0   .   .   .
  .   .   .   0   .   .
  .   .   .   .   0   .
  .   .   .   .   .   0

A = 6×6 Matrix{MP{Float64}}:
  1   2   3    .    .    .
  4   5   6    .    .    .
  7   8   9    .    .    .
  .   .   .   10   11   12
  .   .   .   13   14   15
  .   .   .   16   17   18

B = 6×2 Matrix{MP{Float64}}:
  1   .
  2   .
  3   .
  .   0
  .   0
  .   0

C = 1×6 Matrix{MP{Float64}}:
  4   5   6   0   0   0

x0 = 6-element Vector{MP{Float64}}:
  .
  .
  .
  .
  .
  .
```
"""
function Base.:hcat(x::MPSysLin, y::MPSysLin)
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nb1 = size(x.B, 2)
    nb2 = size(y.B, 2)
    MPSysLin([x.A spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.A],
             [x.B spzeros(MP, n1, nb2); spzeros(MP, n2, nb1) y.B],
             [x.C y.C],
             [x.D spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# Feedback composition: computes star(S1*S2)*S1 in state-space form.
# From ScicosLab file: %mpls_v_mpls.sci

"""
    Base.:(/)(y::MPSysLin, x::MPSysLin)

Feedback composition: computes `mpstar(S1*S2)*S1` in state-space form.

# Examples
```julia-repl
julia> S1 = MPSysLin(MP([1.0 2 3;  4 5 6;  7 8 9]),
                     MP([    1.0;      2;      3]),
                     MP([    4.0       5       6]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S2 = MPSysLin(MP([10.0 11 12;  13 14 15;  16 17 18]),
                     MP([    0.0;      0;      0]),
                     MP([    0.0       0       0]),
                     mpeye(3, 3),
                     full(zeros(MP, 3, 1)));

julia> S1 / S2

Implicit dynamic linear Max-Plus system:
Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 6×6 Matrix{MP{Float64}}:
  0   .   .   1   1   1
  .   0   .   2   2   2
  .   .   0   3   3   3
  4   5   6   0   .   .
  4   5   6   .   0   .
  4   5   6   .   .   0

A = 6×6 Matrix{MP{Float64}}:
  1   2   3    .    .    .
  4   5   6    .    .    .
  7   8   9    .    .    .
  .   .   .   10   11   12
  .   .   .   13   14   15
  .   .   .   16   17   18

B = 6-element Vector{MP{Float64}}:
  1
  2
  3
  .
  .
  .

C = 1×6 Matrix{MP{Float64}}:
  4   5   6   .   .   .

x0 = 6-element Vector{MP{Float64}}:
  .
  .
  .
  .
  .
  .
```
"""
function Base.:(/)(x::MPSysLin, y::MPSysLin)
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nb1 = size(x.B, 2)
    nc1 = size(y.C, 1)
    MPSysLin([x.A spzeros(MP, n1, n2); spzeros(MP, n2, n1) y.A],
             [x.B; spzeros(MP, n2, nb1)],
             [x.C spzeros(MP, nc1, n2)],
             [x.D x.B * y.C; y.B * x.C y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# From ScicosLab file: mplssize.sci
function Base.:size(S::MPSysLin)
    [size(S.B, 1), size(S.B, 2), size(S.C, 1)]
end

# ==============================================================================
# From ScicosLab file: explicit.sci

"""
    mpexplicit(S::MPSysLin)

Convert to an explicit linear system.

# Examples
```julia-repl
julia> S = MPSysLin(MP([1.0 2; 3 4]),
                    MP([0.0; 0]),
                    MP([0.0 0]),
                    mpeye(Float64, 2,2))

Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 2×2 Matrix{MP{Float64}}:
  0   .
  .   0

A = 2×2 Matrix{MP{Float64}}:
  1   2
  3   4

B = 2-element Vector{MP{Float64}}:
  0
  0

C = 1×2 Matrix{MP{Float64}}:
  0   0

x0 = 2-element Vector{MP{Float64}}:
  .
  .

julia> mpexplicit(S)

Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 2×2 Matrix{MP{Float64}}:
  0   .
  .   0

A = 2×2 Matrix{MP{Float64}}:
  1   2
  3   4

B = 2-element Vector{MP{Float64}}:
  0
  0

C = 1×2 Matrix{MP{Float64}}:
  0   0

x0 = 2-element Vector{MP{Float64}}:
  .
  .
```
"""
function mpexplicit(S::MPSysLin)
    ds = mpstar(S.D)
    A = ds * S.A
    B = ds * S.B
    AC = [A; S.C]
    MPSysLin(AC, B, S.C,
             spzeros(MP, size(AC,1), size(AC,2)), S.x0)
end

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

"""
    mpsimul(S::MPSysLin, u::AbsMatMP, history::Bool)

Compute states X of an autonomous linear max-plus system:
```
    x(n+1) = A ⨂ x(n)`  for n = 0 .. k
```

Where:
- `A` is a system matrix
- `x0` is a initial state vector,
- `k` is the number of iterations
- when history is set to true save all computed states, else return the last one.

# Arguments
- u: inputs to inject in the system
- history: if true then return a vector holding all results, else return the
  last result.

# Examples
```julia-repl
julia> S = MPSysLin(MP([1.0 2; 3 4]),
                    MP([0.0; 0]),
                    MP([0.0 0]),
                    mpeye(Float64, 2,2))

Implicit dynamic linear Max-Plus system:
  x(n) = D*x(n) + A*x(n-1) + B*u(n)
  y(n) = C*x(n)
  x(0) = x0

with:
D = 2×2 Matrix{MP{Float64}}:
  0   .
  .   0

A = 2×2 Matrix{MP{Float64}}:
  1   2
  3   4

B = 2-element Vector{MP{Float64}}:
  0
  0

C = 1×2 Matrix{MP{Float64}}:
  0   0

x0 = 2-element Vector{MP{Float64}}:
  .
  .

julia> mpsimul(S, MP(1:10))
1×10 Array{MP{Float64},2}:
 1  5  9  13  17  21  25  29  33  37

julia> mpsimul(S1, MP(1:10), history=false)
1×1 Array{MP{Float64},2}:
 37
```
"""
function mpsimul(S::MPSysLin, u::AbsMatMP, history::Bool)
    x = S.x0
    k = size(u, 1)
    if history
        Y = mpones(k, size(S.C, 1))
        for i = 1:k
            x = S.A * x + S.B * u[i,:]
            Y[i,:] = S.C * x
        end
        Y
    else
        for i = 1:k
            x = S.A * x + S.B * u[i,:]
        end
        Y = S.C * x
        Y[1,:]
    end
end

function mpsimul(S::MPSysLin, u::VecMP, history::Bool)
    mpsimul(S, map(x -> MP(x.λ), reshape(u, length(u), 1)), history)
end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 175c73f6-f2e5-11ea-3458-eddf1465c5c4
using Pkg, DrWatson

# ╔═╡ 175d464e-f2e5-11ea-318a-dd4ad9573a95
using StatisticalRethinking

# ╔═╡ 0b40c1ca-f2e4-11ea-1658-0d38b13854c0
md"## Clip-03-02-05t.jl"

# ╔═╡ 175ca6a0-f2e5-11ea-066c-49e37107aea2
@quickactivate "StatisticalRethinkingTuring"

# ╔═╡ 176c899a-f2e5-11ea-3f17-59b92f0a2483
md"### snippet 3.2"

# ╔═╡ 176d23ca-f2e5-11ea-3a5c-93353273f6eb
begin
	p_grid = range(0, stop=1, length=1000)
	prior = ones(length(p_grid))
	likelihood = pdf.(Binomial.(9, p_grid), 6)
	posterior = likelihood .* prior
	posterior = posterior / sum(posterior)
end

# ╔═╡ 177885bc-f2e5-11ea-0d11-f9ae48c08d8f
md"### snippet 3.3"

# ╔═╡ 17798f98-f2e5-11ea-299e-f3424848225c
md"##### Draw 10000 samples from this posterior distribution."

# ╔═╡ 1784a202-f2e5-11ea-051b-09f8abb69875
begin
	N = 10000
	samples = sample(p_grid, Weights(posterior), N)
end;

# ╔═╡ 178bd8ba-f2e5-11ea-21dd-61c16076a938
md"##### Create an MCMCChains.Chains object."

# ╔═╡ 178d0d36-f2e5-11ea-2593-bf4c46edb720
chn = MCMCChains.Chains(reshape(samples, N, 1, 1), [:p]);

# ╔═╡ df229f44-7dd0-11eb-2ebe-9b12eff28aa1
typeof(chn)

# ╔═╡ 179c95ce-f2e5-11ea-2358-714343a64afa
md"##### Describe the chain."

# ╔═╡ 17a5374c-f2e5-11ea-1bfe-3fb1cce48d48
Text(sprint(show, "text/plain", summarize(chn)))

# ╔═╡ 17a70c16-f2e5-11ea-2e3f-2dc8006c1ee7
md"##### Plot the chain."

# ╔═╡ 17aece92-f2e5-11ea-311a-1b4591b1172a
p1 = plot(chn; seriestype=:traceplot)

# ╔═╡ 17bdf552-f2e5-11ea-28e5-1f5f22b43146
md"### snippet 3.4"

# ╔═╡ 17c6af64-f2e5-11ea-0b28-c55bf548b3ea
p2 = plot(chn; seriestype=:density)

# ╔═╡ 17cf1094-f2e5-11ea-1a80-858534e81348
md"### snippet 3.5"

# ╔═╡ 17d80aa0-f2e5-11ea-1c8f-01770152e142
md"##### Compare with analytical (conjugate) solution."

# ╔═╡ 17e08cca-f2e5-11ea-31a6-8931f98059ba
begin
	w = 6
	n = 9
	x = 0:0.01:1
	plot( x, pdf.(Beta( w+1 , n-w+1 ) , x ), lab="Conjugate solution")
	density!(samples, lab="Sample density")
end

# ╔═╡ 17f4a430-f2e5-11ea-048a-73ae33e192ed
begin
	density(samples, lab="Sample2 density")
	vline!(hpdi(samples), lab="hpdi samples2")
	vline!(quantile(samples, [0.25, 0.75]), lab="quantiles [0.25, 0.75]")
end

# ╔═╡ 17fecdb6-f2e5-11ea-2c5b-7bc56e8b1360
md"## End of clip-03-02-05t.jl"

# ╔═╡ Cell order:
# ╟─0b40c1ca-f2e4-11ea-1658-0d38b13854c0
# ╠═175c73f6-f2e5-11ea-3458-eddf1465c5c4
# ╠═175ca6a0-f2e5-11ea-066c-49e37107aea2
# ╠═175d464e-f2e5-11ea-318a-dd4ad9573a95
# ╟─176c899a-f2e5-11ea-3f17-59b92f0a2483
# ╠═176d23ca-f2e5-11ea-3a5c-93353273f6eb
# ╟─177885bc-f2e5-11ea-0d11-f9ae48c08d8f
# ╟─17798f98-f2e5-11ea-299e-f3424848225c
# ╠═1784a202-f2e5-11ea-051b-09f8abb69875
# ╟─178bd8ba-f2e5-11ea-21dd-61c16076a938
# ╠═178d0d36-f2e5-11ea-2593-bf4c46edb720
# ╠═df229f44-7dd0-11eb-2ebe-9b12eff28aa1
# ╟─179c95ce-f2e5-11ea-2358-714343a64afa
# ╠═17a5374c-f2e5-11ea-1bfe-3fb1cce48d48
# ╟─17a70c16-f2e5-11ea-2e3f-2dc8006c1ee7
# ╠═17aece92-f2e5-11ea-311a-1b4591b1172a
# ╟─17bdf552-f2e5-11ea-28e5-1f5f22b43146
# ╠═17c6af64-f2e5-11ea-0b28-c55bf548b3ea
# ╟─17cf1094-f2e5-11ea-1a80-858534e81348
# ╟─17d80aa0-f2e5-11ea-1c8f-01770152e142
# ╠═17e08cca-f2e5-11ea-31a6-8931f98059ba
# ╠═17f4a430-f2e5-11ea-048a-73ae33e192ed
# ╟─17fecdb6-f2e5-11ea-2c5b-7bc56e8b1360


module PCHIPInterpolation
#
# Simple PCHIP implementation following Fritsch and Carlson, SIAM J. NUMER. ANAL. 17 (1980) 238-246.
# Derivative calculated in a fashion similar to SciPy's PchipInterpolate
#

export Interpolator, integrate

using ArgCheck: @argcheck

struct Interpolator{Xs,Ys,Ds}
    xs::Xs
    ys::Ys
    ds::Ds

    function Interpolator(xs::AbstractVector, ys::AbstractVector)
        @argcheck length(xs) ≥ 2
        @argcheck length(xs) == length(ys) DimensionMismatch
        foldl(xs) do a,b
            @argcheck a < b "xs must be stricly increasing"
            return b
        end

        ds = _initial_ds_scipy(xs, ys)
        new{typeof(xs),typeof(ys),typeof(ds)}(xs, ys, ds)
    end
end

ϕ(t) = 3t^2 - 2t^3
ψ(t) = t^3 - t^2

function _value_with_index(pchip::Interpolator, x::Number, i)
    x1, x2 = pchip.xs[i:i+1]
    @assert x1 ≤ x ≤ x2
    y1, y2 = pchip.ys[i:i+1]
    d1, d2 = pchip.ds[i:i+1]
    h = x2 - x1

    return (y1 * ϕ((x2-x)/h)
           + y2 * ϕ((x-x1)/h)
           - d1*h * ψ((x2-x)/h)
           + d2*h * ψ((x-x1)/h))
end

(pchip::Interpolator)(x::Number) = _value_with_index(pchip, x, _pchip_index(pchip, x))

function _integrate_segment(pchip::Interpolator, i, x1=nothing, x2=nothing)
    if isnothing(x1)
        x1 = pchip.xs[i]
        y1 = pchip.ys[i]
    else
        y1 = _value_with_index(pchip, x1, i)
    end

    if isnothing(x2)
        x2 = pchip.xs[i+1]
        y2 = pchip.ys[i+1]
    else
        y2 = _value_with_index(pchip, x2, i)
    end

    return (x2 - x1)/6 * (y1 + 4*_value_with_index(pchip, (x1 + x2)/2, i) + y2)  # Simpson's rule
end

function integrate(pchip::Interpolator, a::Number, b::Number)
    if b < a
        return -integrate(pchip, b, a)
    end

    i = _pchip_index(pchip, a)
    j = _pchip_index(pchip, b)

    if i == j
        return _integrate_segment(pchip, i, a, b)
    end

    integral = _integrate_segment(pchip, i, a, nothing) + _integrate_segment(pchip, j, nothing, b)

    for k in range(i+1, stop=j-1)
        integral += _integrate_segment(pchip, k)
    end

    return integral
end

function _pchip_index(pchip :: Interpolator, x)
    @argcheck (pchip.xs[1] <= x <= pchip.xs[end]) DomainError
    N = length(pchip.xs)
    if N < 200  # Approximate performance cross-over on my old intel i7-3517U
        i = _pchip_index_linear_search(pchip, x)
    else
        i = _pchip_index_bisectional_search(pchip, x)
    end
    if i == N
        # Treat right endpoint as part of rightmost interval
        i = N-1
    end
    i
end

function _pchip_index_linear_search(pchip :: Interpolator, x)
    i = 1
    N = length(pchip.xs)
    while i < N  &&  x >= pchip.xs[i+1]
        i = i + 1
    end
    i
end

function _pchip_index_bisectional_search(pchip :: Interpolator, x)
    N = length(pchip.xs)
    imin, imax = 1, N

    i = imin + div(imax - imin + 1, 2)
    while imin < imax
        if x < pchip.xs[i]
            imax = i - 1
        elseif x >= pchip.xs[i+1]
            imin = i + 1
        else
            break
        end
        i = imin + div(imax - imin + 1, 2)
    end
    i
end

"Similar to how SciPy's PCHIP does it"
function _initial_ds_scipy(xs::AbstractVector, ys::AbstractVector)
    h(i) = xs[i+1]-xs[i]
    Δ(i) = (ys[i+1]-ys[i]) / h(i)

    N = length(xs)
    ds = similar(ys./xs)
    if N == 2
        ds[:] .= Δ(1)
    else
        Δl = Δ(1)
        hl = h(1)
        for i ∈ 2:N-1
            Δr = Δ(i)
            hr = h(i)
            if sign(Δl) != sign(Δr) || Δl ≈ 0.0 || Δr ≈ 0.0
                ds[i] = 0.0
            else
                wl = 2hl + hr
                wr = hl + 2hr
                ds[i] = (wl + wr) / (wl/Δl + wr/Δr)
            end
            Δl = Δr
            hl = hr
        end
        ds[1] = _edge_derivative(h(1), h(2), Δ(1), Δ(2))
        ds[N] = _edge_derivative(h(N-1), h(N-2), Δ(N-1), Δ(N-2))
    end
    ds
end

function _edge_derivative(h1, h2, Δ1, Δ2)
    d = ((2h1 + h2)*Δ1 - h2*Δ2) / (h1 + h2)
    if sign(d) != sign(Δ1)
        d = 0.0
    elseif sign(Δ1) != sign(Δ2)  &&  abs(d) > abs(3Δ1)
        d = 3Δ1
    end
    d
end

end  # module


module LinearModels

using RobotDynamics

"""
    DoubleIntegrator(D=1)

Create a LinearModel of a `D`-dimensional double integrator, such that the state dimension
is `2*D` (`D` positions and `D` velocities). The model is continuous and non-affine.

This is, in practice, identical to the `DoubleIntegrator` type in the base `RobotZoo`;
however the other is a child of `AbstractModel` whereas this one is explicitly a
`LinearModel`. Both are included for reference and performance comparisons.
"""
function DoubleIntegrator(D=1)
    A = zeros(2D,2D)
    B = zeros(2D,D)
    for i = 1:D
        A[i, D+i] = 1
        B[D+i, i] = 1
    end
    return LinearModel(A, B)
end

end


# some random benchmarks; for j=50000 you need more than 8GB of memory
using WignerSymbols

@time wigner3j(BigFloat, 15, 30, 40, 2, 2, -4)
@time wigner3j(BigFloat, 200, 200, 200, -10, 60, -50)
@time wigner3j(BigFloat, 50000, 50000, 50000, 1000, -6000, 5000)

j = 8
@time wigner6j(BigFloat, j, j, j, j, j, j)
j = 200
@time wigner6j(BigFloat, j, j, j, j, j, j)
j = 600
@time wigner6j(BigFloat, j, j, j, j, j, j)
j = 10000
@time wigner6j(BigFloat, j, j, j, j, j, j)
j = 50000
@time wigner6j(BigFloat, j, j, j, j, j, j)

function compute3jmax(jmax)
    for j₁ = 0:1//2:jmax
        for j₂ = 0:1//2:j₁
            for j₃ = abs(j₁-j₂):j₂
                if isinteger(j₁)
                    m₁ = 0
                    M₂ = min(j₂, j₃-m₁)
                    for m₂ = M₂:(-1):0
                        wigner3j(Float64, j₁,j₂,j₃,0, m₂)
                    end
                    for m₁ = 1:j₁
                        M₂ = min(j₂, j₃-m₁)
                        for m₂ = -M₂:M₂
                            wigner3j(Float64, j₁, j₂, j₃, m₁, m₂)
                        end
                    end
                else
                    for m₁ = 1//2:j₁
                        M₂ = min(j₂, j₃-m₁)
                        for m₂ = -M₂:M₂
                            wigner3j(Float64, j₁, j₂, j₃, m₁, m₂)
                        end
                    end
                end
            end
        end
    end
    return nothing
end

@time compute3jmax(10)

function computeinteger3jmax(jmax)
    for j₁ = 0:jmax
        for j₂ = 0:j₁
            for j₃ = abs(j₁-j₂):j₂
                m₁ = 0
                M₂ = min(j₂, j₃-m₁)
                for m₂ = M₂:(-1):0
                    wigner3j(Float64, j₁,j₂,j₃,0, m₂)
                end
                for m₁ = 1:j₁
                    M₂ = min(j₂, j₃-m₁)
                    for m₂ = -M₂:M₂
                        wigner3j(Float64, j₁, j₂, j₃, m₁, m₂)
                    end
                end
            end
        end
    end
    return nothing
end

function compute6jmax(jmax)
    for j₁ = 0:1//2:jmax
        for j₂ = 0:1//2:j₁
            for j₃ = abs(j₁-j₂):j₂
                for j₄ = 0:1//2:jmax
                    for j₅ = abs(j₃-j₄):min(jmax,j₃+j₄)
                        for j₆ = min(abs(j₂-j₄),abs(j₁-j₅)):min(j₂+j₄,j₁+j₅,jmax)
                            wigner6j(Float64, j₁, j₂, j₃, j₄, j₅, j₆)
                        end
                    end
                end
            end
        end
    end
    return nothing
end
@time compute6jmax(10)

function computeinteger6jmax(jmax)
    for j₁ = 0:jmax
        for j₂ = 0:j₁
            for j₃ = abs(j₁-j₂):j₂
                for j₄ = 0:jmax
                    for j₅ = abs(j₃-j₄):min(jmax,j₃+j₄)
                        for j₆ = min(abs(j₂-j₄),abs(j₁-j₅)):min(j₂+j₄,j₁+j₅,jmax)
                            wigner6j(Float64, j₁, j₂, j₃, j₄, j₅, j₆)
                        end
                    end
                end
            end
        end
    end
    return nothing
end


@testset "Polyhedral" begin
    @testset "Start Solutions Iterator" begin
        f = equations(cyclic(5))
        iter = HC.PolyhedralStartSolutionsIterator(f)

        mv = 0
        for (cell, X) in iter
            mv += cell.volume
            @test size(X,2) == cell.volume
            @test all(1:size(X,2)) do j
                x = @view X[:,j]
                all(1:length(cell.indices)) do i
                    a, b = cell.indices[i]
                    v_i = iter.start_coefficients[i][a] * prod(x.^iter.support[i][:,a]) +
                          iter.start_coefficients[i][b] * prod(x.^iter.support[i][:,b])
                    abs(v_i) < 1e-12
                end
            end
        end
        @test mv == 70

        Random.seed!(130793) # this seed should result in the use of 128 bit hnf
        f = equations(PolynomialTestSystems.tritangents())
        iter = HC.PolyhedralStartSolutionsIterator(f)
        @test sum(cell_X -> first(cell_X).volume, iter) == 12636
    end

    @testset "Polyhedral solve" begin
        f = equations(cyclic(5))
        result = solve(f; start_system=:polyhedral)
        @test nfinite(result) == 70

        @polyvar x y
        result1 = solve([(x-3),(y-2)], start_system=:polyhedral, system_scaling=nothing)
        @test [3,2] ≈ solutions(result1)[1] atol=1e-8
        result1 = solve([(x-3),(y-2)], start_system=:polyhedral, system_scaling=:equations_and_variables)
        @test [3,2] ≈ solutions(result1)[1] atol=1e-8
    end

    @testset "All affine solutions" begin
        @polyvar x y
        f = [2y + 3y^2 - x*y^3, x + 4*x^2 - 2*x^3*y]
        res = solve(f, start_system=:polyhedral, only_torus=true)
        @test nsolutions(res) == 3
        @test ntracked(res) == 3

        res = solve(f, start_system=:polyhedral, affine_tracking=false, only_torus=false)
        @test nsolutions(res) == 6
        @test ntracked(res) == 8

        res = solve(f, start_system=:polyhedral, affine_tracking=true, only_torus=false)
        @test nsolutions(res) == 6
        @test ntracked(res) == 8
    end
end


#___  ____       ___ _____   _   _       _ _
#|  \/  | |     |_  /  __ \ | | | |     (_| |
#| .  . | |       | | /  \/ | | | |_ __  _| |_ ___
#| |\/| | |       | | |     | | | | '_ \| | __/ _ \
#| |  | | |___/\__/ | \__/\ | |_| | | | | | || (_) |
#_______\_____\____/ \____/  _____|_|___|_____\_____ _____ _____
#| ___ \        (_)         | | \ \ / / / __  |  _  / __  |  _  |
#| |_/ _ __ ___  _  ___  ___| |_ \ V /  `' / /| |/' `' / /| |/' |
#|  __| '__/ _ \| |/ _ \/ __| __|/   \    / / |  /| | / / |  /| |
#| |  | | | (_) | |  __| (__| |_/ /^\ \ ./ /__\ |_/ ./ /__\ |_/ /
#\_|  |_|  \___/| |\___|\___|\__\/   \/ \_____/\___/\_____/\___/
#              _/ |
#             |__/
#
# This code is part of the proposal of the team "MLJC UniTo" - University of Turin
# for "ProjectX 2020" Climate Change for AI.
# The code is licensed under MIT 3.0
# Please read readme or comments for credits and further information.

# Compiler: Julia 1.5.3

# Short description of this file: Poisson's 2D Equations with GPU

using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Plots
using CUDA

@parameters x y θ
@variables u(..)
@derivatives Dxx''~x
@derivatives Dyy''~y

# 2D PDE
eq  = Dxx(u(x,y,θ)) + Dyy(u(x,y,θ)) ~ sin(2π*x)*cos(2π*y)

# Boundary conditions
bcs = [u(0,y,θ) ~ 0.f0, u(1,y,θ) ~ -sin(pi*1)*sin(pi*y),
       u(x,0,θ) ~ 0.f0, u(x,1,θ) ~ -sin(pi*x)*sin(pi*1)]
# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0),
           y ∈ IntervalDomain(0.0,1.0)]

# Neural network
dim = 2 # number of dimensions
chain = Chain(Dense(dim,16,Flux.σ),Dense(16,16,Flux.σ),Dense(16,1))

# Discretization
dx = 0.05
discretization = PhysicsInformedNN(dx,
                                   chain,
                                   strategy = GridTraining())

pde_system = PDESystem(eq,bcs,domains,[x,y],[u])
prob = discretize(pde_system,discretization)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

res = GalacticOptim.solve(prob, BFGS(); cb = cb, maxiters=1000)
phi = discretization.phi

xs,ys = [domain.domain.lower:dx/10:domain.domain.upper for domain in domains]
analytic_sol_func(x,y) = -sin(2π*x)*cos(2π*y)/(8π^2)

u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
diff_u = abs.(u_predict .- u_real)

p1 = Plots.plot(xs, ys, u_real, linetype=:contourf,title = "analytic");
p2 = Plots.plot(xs, ys, u_predict, linetype=:contourf,title = "predict");
p3 = Plots.plot(xs, ys, diff_u,linetype=:contourf,title = "error");
plot(p1,p2,p3)

p1 = Plots.plot(xs, ys, u_real, linetype=:surface,title = "analytic");
p2 = Plots.plot(xs, ys, u_predict, linetype=:surface,title = "predict");
p3 = Plots.plot(xs, ys, diff_u,linetype=:surface,title = "error");
plot(p1,p2,p3)
savefig("graph.pdf")


using .C2BasisDefinition: C1, C1Norm

function opnormbound(N::Type{<:C1}, M, B)
    n = size(M, 2)

    est = 0

    for i in 1:n
        v = zeros(n)
        v[i] = 1
        z = C1Norm(B, v)
        est+= C1Norm(B, M*v)/z
    end
    return est.hi
end

function norms_of_powers_basis(B::Basis, m::Integer, Q::DiscretizedOperator, f::AbstractArray;
    normv0::Real=-1., #used as "missing" value
    normQ::Real=-1.,
    normE::Real=-1.,
    normEF::Real=-1.,
    normIEF::Real=-1.,
    normN::Real=-1.)

    @assert eltype(f) <: Interval
    T = typeof(zero(eltype(Q.L)).hi) # gets "Float64" from Q.L
    n = size(Q.L, 1)
    M = mid.(Q.L)
    R = radius.(Q.L)
    #δ = opnormbound(N, R)
    #γz = gamma(T, max_nonzeros_per_row(Q.L))
    #γn = gamma(T, n+3) # not n+2 like in the paper, because we wish to allow for f to be the result of rounding
    #ϵ = zero(T)
    midf = mid.(f)

    norm(v) = C2BasisDefinition.C1Norm(B, v)

    norms = zeros(m)

    S = zeros((n, m))

    for v in AverageZero(B)
        norm_0 = norm(v)
        v/= norm_0.lo
        for i in 1:m
            w = M*v
            v = w - Q.e * (midf*w)[1]
            S[:, i]+= abs.(v)
        end
    end

    return [norm(S[:, i]) for i in 1:m]

end


using Distributions, Turing, Random, Parameters
ProjDir = @__DIR__
cd(ProjDir)
include("LNR.jl")

@model model(data, x) = begin
    μ = Array{Real,1}(undef,Nr)
    β0 ~ Normal(0, 1)
    β1 ~ Normal(1, 1)
    μ = β0 .+ β1.*x
    s ~ Uniform(0, pi/2)
    σ = tan(s)
    for i in 1:length(data)
        data[i] ~ LNR(μ, σ, 0.0)
    end
end

Random.seed!(3423)
Nreps = 30
error_count = fill(0.0, Nreps)
Nr = 20
Nobs = 5

for i in 1:Nreps
    x = rand(Normal(0, 1), Nr)
    β0 = rand(Normal(0, .5))
    β1 = rand(Normal(1, .5))
    μ = β0 .+ β1.*x
    σ = rand(Uniform(.5, 2))
    dist = LNR(μ=μ, σ=σ, ϕ=0.0)
    data = rand(dist, Nobs)
    chain = sample(model(data, x), NUTS(1000, .8), 2000, discard_adapt = false, progress = false)
    error_count[i] = sum(get(chain,:numerical_error)[1])
end

describe(error_count)


module Moments

export make_mon_expo,
       get_ℝ_block_diag,
       make_xxᵀ_tens_yyᵀ,
       get_ℂ_block_diag,
       make_xx̄ᵀ_tens_yȳᵀ,
       var_kron,
       var_kron_C,
       eᵢ,
       get_ℜℑααᶥββᶥᴿ,
       get_xx̄yȳMM_blocks,
       get_γγᶥζζᶥ_δδᶥηηᶥ

## Utils
"""The standard basis vector eᵢ in dimension n"""
eᵢ(n::Int,i::Int) = [Int(j==i) for j in 1:n]

""" input: A,B (arrays of integer tupples) output: A ⊗ B """
function var_kron(A,B)
   C_temp = [ a + b for  b in B, a ∈ A]
   rcs =  [(i,j) for i in 1:size(A)[1] for j in 1:size(B)[1]]
   return [C_temp[ij[2],kl[2],ij[1],kl[1]] for ij in rcs , kl in rcs ]
end

function var_kron_C(A,B)
   C_temp = [ [a] + b for  b in B, a ∈ A]
   rcs =  [(i,j) for i in 1:size(A)[1] for j in 1:size(B)[1]]
   return [C_temp[ij[2],kl[2],ij[1],kl[1]] for ij in rcs , kl in rcs ]
end

## Moments
"""[x]≦ₜ or [x]₌ₜ"""
function make_mon_expo(n::Int,t::Int; isle::Bool = true)
    @assert typeof(t) == Int64
    t == 0 ? (return [eᵢ(n,0)]) : 0
    tmp = make_mon_expo(n,t-1;isle=isle)
    M_vec = reshape([m + eᵢ(n,i) for i ∈ 1:n, m ∈ tmp],:,1)
    return unique(isle ? vcat(tmp,M_vec) : M_vec)
end

"""[x]≦ₜ[x]ᵀ≦ₜ or [x]₌ₜ[x]ᵀ₌ₜ"""
function make_mon_expo(n::Int,t::Tuple{Int,Int}; isle::Bool = true)
    M_vec1      = make_mon_expo(n,t[1]; isle=isle)
    M_vec2      = make_mon_expo(n,t[2]; isle=isle)
    M_mat_flat  = [ mi+ mj for mi in M_vec1 for mj in M_vec2]
    M_mat       = reshape(M_mat_flat, (length(M_vec1), length(M_vec2)) )
    return M_mat
end

## Complex moments
"""[x,̄x]≦ₜ or [x,̄x]ᵀ₌ₜ"""
make_mon_expo(d::Tuple{Int,Int},t::Int; isle::Bool = true) = make_mon_expo(sum(2 .* d),t; isle=isle)

"""[x,̄x]≦ₜ[x,̄x]*≦ₜ or [x,̄x]₌ₜ[x,̄x]*₌ₜ
but is actually [xᵣₑ,xᵢₘ]≦ₜ[xᵣₑ,xᵢₘ]ᵀ≦ₜ or [xᵣₑ,xᵢₘ]₌ₜ[xᵣₑ,xᵢₘ]ᵀ₌ₜ
"""
function make_mon_expo(d::Tuple{Int,Int},t::Tuple{Int,Int}; isle::Bool = true)
    M_vec1      = make_mon_expo(d,t[1]; isle=isle)
    M_vec2      = make_mon_expo(d,t[1]; isle=isle)
    M_mat      = [ mi+ mj for mi in M_vec1, mj in M_vec2]
    # M_mat       = reshape(M_mat_flat, (length(M_vec1), length(M_vec2)) )
    return M_mat
end
## Real
"""xxᵀ⊗yyᵀ"""
function make_xxᵀ_tens_yyᵀ(d)
    n = sum(d)
    pre_expo = make_mon_expo(n,1;isle = false)
    x, y = pre_expo[1:d[1]], pre_expo[d[1]+1:n]
    return var_kron(x .+ reshape(x,1,d[1]), y .+ reshape(y,1,d[2]))
end

"""Returns diagonal block of the momoment matrix with partition:
    xᵃ⁺ᶜyᵇ⁺ᵈ where |a+c|,|b+d| ∈ 2N """
function get_ℝ_block_diag(d::Tuple{Int,Int},t::Tuple{Int,Int})
    mom_mat = Moments.make_mon_expo(sum(d),t)
    # return Dict("Default" => mom_mat)
    mom_vec = mom_mat[1,:]
    halfsum(arr) = [sum(arr[1:d[1]]),sum(arr[(d[1]+1):end])]
    tf(arr)      = map(x -> (isodd(x[1]),isodd(x[2])), halfsum.(arr))
    tf2(arr)     = map(p -> p[1], findall(arr))

    isevev(pair) = pair == (false,false)
    isodod(pair) = pair == (true,true)
    isevod(pair) = pair == (false,true)
    isodev(pair) = pair == (true,false)

    ee = tf2(isevev.(tf(mom_vec)))
    oo = tf2(isodod.(tf(mom_vec)))
    eo = tf2(isevod.(tf(mom_vec)))
    oe = tf2(isodev.(tf(mom_vec)))

    ℝ_block_diag = Dict("ee" => mom_mat[ee,ee],
                        "oo" => mom_mat[oo,oo],
                        "eo" => mom_mat[eo,eo],
                        "oe" => mom_mat[oe,oe])

    for key in ["ee","oo","eo","oe"]
         isempty(ℝ_block_diag[key]) ? delete!(ℝ_block_diag,key) : continue
    end
    return ℝ_block_diag
end

## Complex

""" ℜe(xx*⊗yy*), ℑm(xx*⊗yy*)
xx*⊗yy* = (xᵣₑ + i xᵢₘ)(xᵣₑ - i xᵢₘ)ᵀ ⊗ (yᵣₑ + i yᵢₘ)(yᵣₑ - i yᵢₘ)ᵀ
= (xᵣₑxᵣₑᵀ + xᵢₘxᵢₘᵀ + i (xᵢₘxᵣₑᵀ - xᵣₑxᵢₘᵀ)) ⊗ (yᵣₑyᵣₑᵀ + yᵢₘyᵢₘᵀ + i (yᵢₘyᵣₑᵀ - yᵣₑyᵢₘᵀ))
⟹ ℜe(xx*⊗yy*) = (xᵣₑxᵣₑᵀ + xᵢₘxᵢₘᵀ ) ⊗ (yᵣₑyᵣₑᵀ + yᵢₘyᵢₘᵀ ) - (xᵢₘxᵣₑᵀ - xᵣₑxᵢₘᵀ) ⊗ (yᵢₘyᵣₑᵀ - yᵣₑyᵢₘᵀ)
⟹ ℑm(xx*⊗yy*) = (xᵣₑxᵣₑᵀ + xᵢₘxᵢₘᵀ) ⊗ (yᵢₘyᵣₑᵀ - yᵣₑyᵢₘᵀ) + (xᵢₘxᵣₑᵀ - xᵣₑxᵢₘᵀ)⊗(yᵣₑyᵣₑᵀ + yᵢₘyᵢₘᵀ) """
U(vec,dₐ) = vec .+ reshape(vec,1,dₐ)
W(vec1,vec2,dₐ) = vec1 .+ reshape(vec2,1,dₐ)
function make_xx̄ᵀ_tens_yȳᵀ(d)
    d₁,d₂ = d
    n = sum(2 .* d)
    pre_expo = make_mon_expo(n,1;isle =false)
    xᵣₑ, xᵢₘ = pre_expo[1:d₁], pre_expo[d₁+1:2*d₁]
    yᵣₑ, yᵢₘ = pre_expo[2*d₁+1:2*d₁+d₂], pre_expo[2*d₁+d₂+1:2*d₁+2*d₂]

    xᵣₑxᵣₑᵀ, xᵢₘxᵢₘᵀ, yᵣₑyᵣₑᵀ, yᵢₘyᵢₘᵀ = U(xᵣₑ,d₁), U(xᵢₘ,d₁), U(yᵣₑ,d₂), U(yᵢₘ,d₂)
    xᵢₘxᵣₑᵀ, xᵣₑxᵢₘᵀ, yᵢₘyᵣₑᵀ, yᵣₑyᵢₘᵀ = W(xᵢₘ,xᵣₑ,d₁), W(xᵣₑ,xᵢₘ,d₁), W(yᵢₘ,yᵣₑ,d₂), W(yᵣₑ,yᵢₘ,d₂)

    Real_dict =  Dict(("+ 1")  => var_kron(xᵣₑxᵣₑᵀ,yᵣₑyᵣₑᵀ),
                      ("+ 2")  => var_kron(xᵣₑxᵣₑᵀ,yᵢₘyᵢₘᵀ),
                      ("+ 3")  => var_kron(xᵢₘxᵢₘᵀ,yᵣₑyᵣₑᵀ),
                      ("+ 4")  => var_kron(xᵢₘxᵢₘᵀ,yᵢₘyᵢₘᵀ),
                      ("- 5")  => var_kron(xᵢₘxᵣₑᵀ,yᵢₘyᵣₑᵀ),
                      ("+ 6")  => var_kron(xᵢₘxᵣₑᵀ,yᵣₑyᵢₘᵀ),
                      ("+ 7")  => var_kron(xᵣₑxᵢₘᵀ,yᵢₘyᵣₑᵀ),
                      ("- 8")  => var_kron(xᵣₑxᵢₘᵀ,yᵣₑyᵢₘᵀ))

    Imag_dict =  Dict(("+ 1")  => var_kron(xᵣₑxᵣₑᵀ,yᵢₘyᵣₑᵀ),
                      ("- 2") => var_kron(xᵣₑxᵣₑᵀ,yᵣₑyᵢₘᵀ),
                      ("+ 3")  => var_kron(xᵢₘxᵢₘᵀ,yᵢₘyᵣₑᵀ),
                      ("- 4") => var_kron(xᵢₘxᵢₘᵀ,yᵣₑyᵢₘᵀ),
                      ("+ 5")  => var_kron(xᵢₘxᵣₑᵀ,yᵣₑyᵣₑᵀ),
                      ("+ 6")  => var_kron(xᵢₘxᵣₑᵀ,yᵢₘyᵢₘᵀ),
                      ("- 7") => var_kron(xᵣₑxᵢₘᵀ,yᵣₑyᵣₑᵀ),
                      ("- 8") => var_kron(xᵣₑxᵢₘᵀ,yᵢₘyᵢₘᵀ))

    return Dict("real" => Real_dict,
                "imag" => Imag_dict)
end

## Block diagonalization
# The complex blocks
"""Splits a vector of exponents into the register components ααᶥββᶥ →  (α,αᶥ,β,βᶥ)"""
split_expo(ααᶥββᶥ,d) =     (ααᶥββᶥ[1:d[1]],
                            ααᶥββᶥ[d[1]+1:2*d[1]],
                            ααᶥββᶥ[1+2*d[1]:2*d[1]+d[2]],
                            ααᶥββᶥ[1+2*d[1]+d[2]:end])

"""Gets the principal block of the moment matrix after block diagonalization"""
function get_xx̄yȳMM_blocks(d,t)
    Iᵗ = Moments.make_mon_expo(d,t[1])
    Iᵗ_temp = map(x ->  sum.(split_expo(x,d)),Iᵗ)
    r_list  = map(v -> v[1]-v[2],Iᵗ_temp) ; s_list  = map(v -> v[3]-v[4],Iᵗ_temp)

    Iᵗ_d = Dict()
    for r ∈ -t[1]:t[1], s ∈ -t[1]:t[1]
        T = Iᵗ[(r_list .== r) .& (s_list .== s) ]
        isempty(T) ? continue : Iᵗ_d[r,s] = T
    end
    return Dict(zip(keys(Iᵗ_d),[Iᵗ_d[k] .+ reshape(Iᵗ_d[k],1,:) for k in keys(Iᵗ_d)]))
end

## The real analogue
"""Returns all  (γ,γᶥ,ζ,ζᶥ),(δ,δᶥ,η,ηᶥ) ∈ (ℕᵈ)⁴ s.t.
(γ,γᶥ,ζ,ζᶥ)+(δ,δᶥ,η,ηᶥ)=(α,αᶥ,β,βᶥ)  ∀  ααᶥββᶥ ∈ Moment matrixₜ
"""
function get_γγᶥζζᶥ_δδᶥηηᶥ(d,t)
    M = Moments.make_mon_expo(d,t)
    M_vec = M[:,1]
    γγᶥζζᶥ_δδᶥηηᶥ_dict = Dict()
    for k in Moments.make_mon_expo(d,2*t[1])# Define variables in the moment matrix.
        fd = hcat(map(x -> [x[1],x[2]],findall(M .== [k]))...)
        γγᶥζζᶥ_δδᶥηηᶥ_dict[k] = [M_vec[fd[1,:]], M_vec[fd[2,:]]]
    end
    return γγᶥζζᶥ_δδᶥηηᶥ_dict
end

"""The coefficients function"""
pf(x) = prod(factorial.(x))
function get_coef(part1,part2,d)
    γ,γᶥ,ζ,ζᶥ = split_expo(part1,d) ; δ,δᶥ,η,ηᶥ = split_expo(part2,d)
    a = ((-1.0)^sum(δᶥ+ηᶥ))*((1.0im)^sum(δ+δᶥ+η+ηᶥ))*prod(pf.([[γ,γᶥ,ζ,ζᶥ]+[δ,δᶥ,η,ηᶥ]...]))
    return a/prod(pf.([γ,γᶥ,ζ,ζᶥ,δ,δᶥ,η,ηᶥ]))
end
get_coef(γ_δ_pair,d) = map((x,y)->get_coef(x,y,d),γ_δ_pair...)

""" get_expo(γγᶥδδᶥ,ζζᶥηηᶥ,d) = [γ+γᶥ,δ+δᶥ,η+ηᶥ,ζ+ζᶥ] """
function get_expo(γ_,δ_,d)
    γ,γᶥ,ζ,ζᶥ = split_expo(γ_,d) ; δ,δᶥ,η,ηᶥ = split_expo(δ_,d)
    return vcat(γ+γᶥ,δ+δᶥ,ζ+ζᶥ,η+ηᶥ)
end
get_expo(γ_δ_,d) = map((x,y)->get_expo(x,y,d),γ_δ_...)

"""f"""
function get_ℜℑααᶥββᶥᴿ(d,ααᶥββᶥ::Array{Int64,1},γdict)
    γ_δ_ = γdict[ααᶥββᶥ]
    return   get_expo(γ_δ_,d),get_coef(γ_δ_,d)
end

function get_ℜℑααᶥββᶥᴿ(d,B,γdict)
    MMex = map(x->get_ℜℑααᶥββᶥᴿ(d,x,γdict)[1],B)
    MMCoef = map(x->get_ℜℑααᶥββᶥᴿ(d,x,γdict)[2],B)

    real_mask = isreal.(MMCoef)

    RMMex = copy(MMex) ; IMMex = copy(MMex)
    RMMex[.!real_mask] .= [[]]
    IMMex[real_mask] .= [[]]

    RMMCoef = copy(MMCoef) ; IMMCoef = copy(MMCoef)
    RMMCoef[.!real_mask] .= [[0.0]] ;RMMCoef = real.(RMMCoef)
    IMMCoef[real_mask] .= [[0.0]] ; IMMCoef = imag.(IMMCoef)

    MMexᴿ = hcat(vcat(RMMex,IMMex),vcat(IMMex,RMMex))
    MMCoefᴿ = hcat(vcat(RMMCoef,IMMCoef),vcat((-1.0)*IMMCoef,RMMCoef))
    return MMexᴿ,MMCoefᴿ
end

""""""
function get_ℂ_block_diag(d,t;noBlock = false)
    if noBlock
        MMexᴿ = Dict("Default" => map(x->[x],make_mon_expo(d,t)))
        MMCoefᴿ = Dict("Default" => fill([1.0],size(MMexᴿ["Default"])...))
        return MMexᴿ,MMCoefᴿ
    end
    xx̄yȳMM_blocks = get_xx̄yȳMM_blocks(d,t)
    γγᶥζζᶥ_δδᶥηηᶥ = get_γγᶥζζᶥ_δδᶥηηᶥ(d,t)

    MMexᴿ= Dict() ; MMCoefᴿ = Dict()
    for B in keys(xx̄yȳMM_blocks)
        MMexᴿ[B], MMCoefᴿ[B] = get_ℜℑααᶥββᶥᴿ(d,xx̄yȳMM_blocks[B],γγᶥζζᶥ_δδᶥηηᶥ)
    end
    return MMexᴿ,MMCoefᴿ
end

end


# This can probably be cleaned up and simplifed with a little thought
@inline wavelength{T<:Unitful.Frequency}(s::T) = c0/s
@inline wavelength{T<:Unitful.Length}(s::T) = s
@inline wavelength(attr::OpticalAttr) = wavelength.(attr.source)
@inline wavelength(laser::AbstractLaser) = wavelength(laser.frequency)
@inline frequency{T<:Unitful.Frequency}(s::T) = s
@inline frequency{T<:Unitful.Length}(s::T) = c0/s
@inline frequency(attr::OpticalAttr) = frequency.(attr.source)
@inline frequency(laser::AbstractLaser) = frequency(laser.frequency)
@inline frequency(model::Model) = frequency(model.laser)
@inline frequency(model::ToyModel) = frequency(model.ω0)

function get_attr(attr::OpticalAttr) attr.property end

function get_label(attr::OpticalAttr) attr.label end
"""
OpticalAttr is callable. Giving a source source will return the
value of the OpticalAttr interpolated at that point
"""
(attr::OpticalAttr)(source) = attr.fit_func(frequency(source))
"""
Alias for der
"""
function der(attr::OpticalAttr, query; order=1)
    der(attr.fit_func, query; order=order)
end

function fit(attr::OpticalAttr, poly_order=12)
    ω = frequency(attr)
    μ = mean(ω)
    σ = std(ω)
    attr_unit = oneunit(first(get_attr(attr)))
    p = attr_unit*polyfit((ω-μ)/σ, get_attr(attr)./attr_unit, poly_order)
    attr.fit_func = ScaledFit(μ, σ, p)
    return attr
end


using Flux
using LinearAlgebra: I # Identity matrix
using Distributions
using Random
using Plots

using Minijyro

Random.seed!(42)
theme(:ggplot2)

# Convenience functions
function get_num_params(nn)
    # For each parameter vector count the number of elements and collect the total
    # sum.
    return mapreduce(p -> prod(size(p)), +, Flux.params(nn))
end

function set_params!(nn, weights)
    new_params = Flux.params([])
    i = 1
    for p in Flux.params(nn)
        num_elements = prod(size(p))
        new_p = weights[i:i+num_elements-1]
        push!(new_params, reshape(new_p, size(p)))
        i += num_elements
    end
    Flux.loadparams!(nn, new_params)
end

@jyro function model(xs, neural_net, sigma)
    num_params = get_num_params(neural_net)
    w ~ MvNormal(num_params, 1.0)

    set_params!(neural_net, w)
    y ~ MvNormal(neural_net(xs), sigma*I)
end

D = 1 # Dimensionality of data
N = 100 # Number of data points
sigma = 0.1 # Noise level
dnn = Chain(Dense(D, 5, tanh), Dense(5, 1), x -> vec(x))

# Create some data
xs = collect(range(-3, stop=3, length=N))
xs = reshape(xs, (D,N))
ys = cos.(xs) + sigma * randn(D,N)
ys = reshape(ys, (N,))

display(scatter(xs[1,:], ys, grid=:none))

# Run NUTS sampler
cond_model = condition(model, Dict(:y => ys))
trace!(cond_model)
samples, stats = nuts(cond_model, (xs, dnn, sigma), 2000; autodiff=:reverse, progress=true)

# Plot posterior samples
p = plot(xlims=(-10,10))
for posterior_sample in samples[end-500:end]
    wide_xs = range(-10, stop=10, length=400)
    wide_xs = reshape(wide_xs, (1,400))
    set_params!(dnn, posterior_sample)
    preds = dnn(wide_xs)
    plot!(p, wide_xs[1,:], preds, color=:red, alpha=0.3)
end
display(scatter!(p, xs[1,:], ys, color=:blue, grid=:none, legend=false))


# Type and constructors
# Default constructor is for the line through two points
"""
	(type) Line{T<:AnyComplex} in the complex plane

Each `Line` type is parameterized according to the common type of its complex input arguments.
"""
struct Line{T<:AnyComplex} <: AbstractClosedCurve
	base::T
	direction::T
	Line{T}(a,b) where T = new(a,sign(b-a))
end
"""
	Line(a,b)

Construct the line passing through the two given points.

	Line(a,direction=z)

Construct the line through the given point and parallel to the complex sign of the given `direction` value. """
function Line(a::Number,b::Number)
	a,b = promote(complex(float(a)),complex(float(b)))
	Line{typeof(a)}(a,b)
end

# Construct by giving a keyword
Line(z::Number;direction) = Line(z,z+direction)
#Line(z::Number;angle) = Line(z,z+exp(complex(0,angle)))

# Complex type converters
for ctype in [:Spherical,:Polar,:Complex]
	@eval begin
		function $ctype(L::Line{T}) where T<:AnyComplex
			Line($ctype(L.base),$ctype(L.base+L.direction))
		end
	end
end

# Required methods
arclength(::Line) = Inf
point(L::Line,t::Real) = L.base + (2t-1)/(t-t^2)*L.direction
(C::Line)(t::Real) = point(C,t)

ispositive(C::Line) = true  # this seems arbitrary...

"""
	arg(L::Line,z)

Find a parameter argument `t` such that `L(t)==z` is true. For an infinite `z`, return zero (but note that `L(1)` is also infinity).

This gives undefined results if `z` is not actually on the line.
"""
function arg(L::Line,z::Number)
	isinf(z) && return float(0)
	del = z-L.base
	if abs(real(del)) > abs(imag(del))
		del = real(del)
		α = real(L.direction)
	else
		del = imag(del)
		α = imag(L.direction)
	end
	for t in realroots(del,2α-del,-α)
		0-eps(del) ≤ t ≤ 1+eps(del) && return t
	end
	return []
end
unittangent(L::Line,t::Real=0) = L.direction
function tangent(L::Line{T},t::Real) where T <: AnyComplex
	T( (1/t^2+1/(1-t)^2)*L.direction )
end

+(L::Line,z::Number) = Line(L.base+z,direction=L.direction)
-(L::Line) = Line(-L.base,direction=-L.direction)
*(L::Line,z::Number) = Line(L.base*z,direction=L.direction*sign(z))

"""
	inv(L)
Invert a line `L` through the origin. In general the inverse is a `Circle` through the inverse of any three points on the line.
"""
function inv(L::Line)
	w = 1 ./ point(L,[0.25,0.5,0.75])
	Circle(w...)
end

# Other methods
isfinite(::Line) = false
slope(L::Line) = imag(L.direction)/real(L.direction)
angle(L::Line) = angle(L.direction)
conj(L::Line) = Line(conj(L.base),direction=conj(L.direction))
reverse(L::Line) = Line(L.base,direction=-L.direction)

"""
	isapprox(L1::Line,L2::Line; tol=<default>)
	L1 ≈ L2

Determine if `L1` and `L2` represent the same line, irrespective of the type or values of its parameters. Identity is determined by agreement within `tol`, which is interpreted as the weaker of absolute and relative differences.
"""
function isapprox(L1::Line,L2::Line;tol=DEFAULT[:tol])
	dz = L1.base - L2.base
	w1,w2 = L1.direction,L2.direction
	return isapprox(real(w1)*imag(w2),imag(w1)*real(w2),rtol=tol,atol=tol) &&
		isapprox(real(w1)*imag(dz),imag(w1)*real(dz),rtol=tol,atol=tol)
end

"""
	isleft(z,L::Line)

Determine whether the number `z` lies "to the left" of line `L`. This means that the angle it makes with `tangent(L)` is in the interval (0,π).

Note that `isleft` and `isright` are *not* logical opposites; a point on the curve should give `false` in both cases.
"""
isleft(z::Number,L::Line) = π > angle((z-L.base)/L.direction) > 0

"""
	isright(z,L::Line)

Determine whether the number `z` lies "to the right" of line `L`. This means that the angle it makes with `tangent(L)` is in the interval (-π,0).

Note that `isleft` and `isright` are *not* logical opposites; a point on the curve should give `false` in both cases.
"""
isright(z::Number,L::Line) = -π < angle((z-L.base)/L.direction) < 0

winding(L::Line,z::Number) = isleft(z,L) ? 1 : 0

"""
	dist(z,L::Line)

Compute the distance from number `z` to the line `L`.
"""
dist(z::Number,L::Line) = imag( (z-L.base)/L.direction )
"""
	closest(z,L::Line)

Find the point on line `L` that lies closest to `z`.
"""
function closest(z::Number,L::Line)
	s = L.direction
	L.base + real((z-L.base)/s)*s
end
"""
	reflect(z,L::Line)

Reflect the value `z` across the line `L`. (For reflection of a line through a point, use translation and negation.)
"""
function reflect(z::Number,L::Line)
	ζ = z - L.base
	L.base + L.direction*conj(ζ/L.direction)
end

function show(io::IO,L::Line)
	print(IOContext(io,:compact=>true),"Line(...",L(0.5),"...",L((sqrt(5)-1)/2),"...)")
end
function show(io::IO,::MIME"text/plain",L::Line{T}) where {T}
	print(io,"Line{$T} in the complex plane:\n   through (",L.base,") parallel to (",L.direction,")")
end

# In the plane (but not on the sphere), two points are enough to draw a line, and we want to avoid infinity.
plotdata(L::Line{T}) where T<:Union{Complex,Polar} = point(L,[0.1,0.9])


using Base: Char
# one-body matrix, nearest neighbour, PBC
function kinetic_matrix_generator(NsX::Int64, NsY::Int64, t::Float64)

    KMatrix = zeros(NsX * NsY, NsX * NsY)

    for i = 1 : NsX
        for j = 1 : NsY
            # indices of nearest neighbours (nn) of (i, j)
            nn_up = mod(i - 2, NsX) + 1
            nn_dn = mod(i, NsX) + 1
            nn_lf = mod(j - 2, NsY) + 1
            nn_rg = mod(j, NsY) + 1

            # parse indices into the kinetic matrix
            KMatrix[(i - 1) * NsY + j, (nn_up - 1) * NsY + j] = -t
            KMatrix[(i - 1) * NsY + j, (nn_dn - 1) * NsY + j] = -t
            KMatrix[(i - 1) * NsY + j, (i - 1) * NsY + nn_lf] = -t
            KMatrix[(i - 1) * NsY + j, (i - 1) * NsY + nn_rg] = -t
        end
    end

    return KMatrix

end

# auxiliary field matrix generator
@inline function auxfield_matrix_generator(σ::Array{Int64,1}, γ::Float64, Δτ::Float64, U::Float64)

    AFMatrixUp = diagm(0 => exp.(2 * γ * σ .- Δτ * U / 2))
    AFMatrixDn = diagm(0 => exp.(-2 * γ * σ .- Δτ * U / 2))
    return AFMatrixUp, AFMatrixDn

end

function singlestep_matrix(σ::Array{Int64,1}, lattice::Lattice, qmc::QMC)

    AFMatrixUp, AFMatrixDn = auxfield_matrix_generator(σ, qmc.γ, qmc.Δτ, lattice.U)
    Bl_Up = qmc.KMatrixExp * AFMatrixUp * qmc.KMatrixExp
    Bl_Dn = qmc.KMatrixExp * AFMatrixDn * qmc.KMatrixExp

    return Bl_Up, Bl_Dn

end

# Calculate Q', D', T' for Q'D'T' = QDT * B or B * QDT
# Q => orthogonal matrix
# D => diagonal matrix with descending entries
# arguments Q, D, T would be overwritten
function QRCP_update!(
    Q::Array{Float64,2}, D::Array{Float64,2}, T::Array{Float64,2},
    B::Array{Float64,2}, direction::Char
    )

    # Q'D'T' = B * QDT
    if direction == 'L'
        BQD = B * Q * D
        # column-pivoted QR (QRCP) decomposition
        QRCP_BQD = LinearAlgebra.qr(BQD, Val(true))
        Q .= QRCP_BQD.Q * I
        D .= diagm(0 => diag(QRCP_BQD.R))
        T .= (inv(D) * QRCP_BQD.R) * (QRCP_BQD.P' * T)

        return Q, D, T

    # Q'D'T' = QDT * B
    elseif direction == 'R'
        DTB = D * T * B
        # column-pivoted QR (QRCP) decomposition
        QRCP_DTB = LinearAlgebra.qr(DTB, Val(true))
        Q .= Q * QRCP_DTB.Q
        D .= diagm(0 => diag(QRCP_DTB.R))
        T .= (inv(D) * QRCP_DTB.R) * QRCP_DTB.P'

        return Q, D, T

    else
        ArgumentError("direction can only be 'L' or 'R'")
    end

end

# full propagation using column-pivoted QR (QRCP) decomposition
function matrix_product_stablized(
    σlist::Array{Int64,2}, lattice::Lattice, qmc::QMC
    )

    MatProdUp_Q = Matrix{Float64}(I, lattice.Ns, lattice.Ns)
    MatProdUp_D = Matrix{Float64}(I, lattice.Ns, lattice.Ns)
    MatProdUp_T = Matrix{Float64}(I, lattice.Ns, lattice.Ns)

    MatProdDn_Q = Matrix{Float64}(I, lattice.Ns, lattice.Ns)
    MatProdDn_D = Matrix{Float64}(I, lattice.Ns, lattice.Ns)
    MatProdDn_T = Matrix{Float64}(I, lattice.Ns, lattice.Ns)

    for i = 1 : div(qmc.L, qmc.cluster_size)

        MatProdUp = Matrix{Float64}(I, lattice.Ns, lattice.Ns)
        MatProdDn = Matrix{Float64}(I, lattice.Ns, lattice.Ns)

        for j = 1 : qmc.cluster_size
            σ = σlist[:, (i - 1) * qmc.cluster_size + j]
            #AFMatrixUp, AFMatrixDn = auxfield_matrix_generator(σ, qmc.γ, qmc.Δτ, lattice.U)
            #MatProdUp = Symmetric(qmc.KMatrixExp * AFMatrixUp * qmc.KMatrixExp) * MatProdUp
            #MatProdDn = Symmetric(qmc.KMatrixExp * AFMatrixDn * qmc.KMatrixExp) * MatProdDn
            Bj_Up, Bj_Dn = singlestep_matrix(σ, lattice, qmc)
            MatProdUp = Bj_Up * MatProdUp
            MatProdDn = Bj_Dn * MatProdDn
        end

        QRCP_update!(MatProdUp_Q, MatProdUp_D, MatProdUp_T,
                    MatProdUp, 'L')
        QRCP_update!(MatProdDn_Q, MatProdDn_D, MatProdDn_T,
                    MatProdDn, 'L')

    end

    return MatProdUp_Q, MatProdUp_D, MatProdUp_T, MatProdDn_Q, MatProdDn_D, MatProdDn_T

end

# generate overlap "tensor" / column major
function overlap_coefficient(Ns::Int64,
    UUp::Array{T1,2}, UDn::Array{T2,2}
) where {T1<:FloatType, T2<:FloatType}

    UUp_Inv, UDn_Inv = inv(UUp), inv(UDn)
    UUp_Array = complex(zeros(Ns, Ns, Ns))
    UDn_Array = complex(zeros(Ns, Ns, Ns))

    for i = 1 : Ns
        for j = 1 : Ns
            for k = 1 : Ns
                UUp_Array[k ,i, j] = UUp_Inv[k, i] * UUp[j, k]
                UDn_Array[k, i, j] = UDn_Inv[k, i] * UDn[j, k]
            end
        end
    end

    return UUp_Array, UDn_Array

end


import FEM.stiffness
import FEM.GaussPoint2
import FEM.Point2

facts("FEM.AbstractMaterial") do

context("FEM.Material.LinearIsotropic") do

    mat = FEM.linearisotropicmod().LinearIsotropic(200e9, 0.3)
    gp = GaussPoint2(Point2(0.0, 0.0), 0.0)
    context("FEM.Material.LinearIsotropic.stiffness") do

        # Stiffness
        #########################################################################################
        stiff = stiffness(mat, gp)
        # Below calculated using Calfem with command hooke(4, 200e9, 0.3)
        stiff_calfem = [
        [269230769230.769    115384615384.615    115384615384.615                  0];
        [115384615384.615    269230769230.769    115384615384.615                  0];
        [115384615384.615    115384615384.615    269230769230.769                  0];
        [               0                   0                   0   76923076923.0769]]

        @fact norm(stiff - stiff_calfem) / norm(stiff) => roughly(0.0, atol=10.0^(-14))
    end

end # context

end # facts


@testset "Horn" begin
    p, q = synthetic_gradient(SynthSphere(), radius = 50)
    Z = convert_gradient(Horn(), p, q)
    @test Z[38,38] ≈ -0.0027751467456787315
    @test Z[1,1] ≈ 0.0
    @test Z[76,76] ≈ 47.05959259743248
end


export k_guessing_game, ambiguous_guessing_game

export maximum_likelihood_facet, anti_guessing_facet, ambiguous_guessing_facet, k_guessing_facet, non_negativity_facet
export coarse_grained_input_ambiguous_guessing_facet

"""
The ``k``-guessing game is a general family of Bell inequalities for signaling polytopes.
For any integer ``k\\in[0,Y]`` a ``k``-guessing game Bell inequality
``(\\mathbf{G}_{\\text{K}},\\gamma_{\\text{K}})``
bounds the signaling polytope ``\\mathcal{C}_d^{\\binom{Y}{k}\\to Y }``.
The columns of matrix ``\\mathbf{G}_{\\text{K}}\\in \\mathbf{R}^{Y \\times \\binom{Y}{k}}``
consist of all ``\\binom{Y}{k}`` ways to arrange ``k`` unit elements and ``(Y-k)``
null elements.
Hence each input ``x\\in[X]`` corresponds to a unique set of ``k`` correct
answers out of ``Y`` possible answers.
The upper bound is  ``\\gamma_{\\text{K}} = \\binom{Y}{k} - \\binom{Y-d}{k}``.

A ``k``-guessing game Bell inequality is constructed with the `k_guessing_game`
method,

```julia
k_guessing_game(Y :: Int64, d :: Int64, k :: Int64) :: BellGame
```

Parameters:
* `Y` - The number of outputs
* `d` - The signaling dimension
* `k` - The number of unit elements in each column. Must satisfy `Y ≥ k ≥ 0`.

For example,

```jldoctest k_guessing_game
using SignalingDimension

G_K = k_guessing_game(6,2,3)

# output

6×20 BellScenario.BellGame:
 1  1  1  1  1  1  1  1  1  1  0  0  0  0  0  0  0  0  0  0
 1  1  1  1  0  0  0  0  0  0  1  1  1  1  1  1  0  0  0  0
 1  0  0  0  1  1  1  0  0  0  1  1  1  0  0  0  1  1  1  0
 0  1  0  0  1  0  0  1  1  0  1  0  0  1  1  0  1  1  0  1
 0  0  1  0  0  1  0  1  0  1  0  1  0  1  0  1  1  0  1  1
 0  0  0  1  0  0  1  0  1  1  0  0  1  0  1  1  0  1  1  1
```

The upper bound is then

```jldoctest k_guessing_game
G_K.β

# output

16
```

The `k_guessing_game` can alternatively take a `BellScenario.LocalSignaling(X, Y, d)` scenario as input.

```julia
k_guessing_game( scenario :: LocalSignaling, k :: Int64 ) :: BellGame
```

A `DomainError` is thrown if `scenario.X !=  binomial(scenario.Y, k)`.
"""
function k_guessing_game(scenario :: LocalSignaling, k :: Int64) :: BellGame
    if scenario.X != binomial(scenario.Y, k)
        throw(DomainError(scenario.X, "Number of inputs `X` must be `binomial(Y, k)`."))
    end

    k_guessing_game(scenario.Y, scenario.d, k)
end
function k_guessing_game(Y :: Int64, d :: Int64, k :: Int64) :: BellGame
    if !(Y ≥ k ≥ 0)
        throw(DomainError(k, "Input k must satisfy `Y  ≥ k ≥ 0`"))
    end

    game = hcat(map( combo -> begin
        a = zeros(Int64, Y)
        a[combo] .= 1
        return a
    end, combinations(1:Y, k))...)

    max_score = binomial(Y,k) - binomial(Y-d,k) #sum(map(i -> binomial(Y-i, k-1), 1:d))

    BellGame(game, max_score)
end

"""
Ambiguous guessing games are a family of Bell inequalities for signaling polytopes.
An ambiguous guessing game, denoted by ``(\\mathbf{G}_{\\text{Q}},\\gamma_{\\text{Q}})``, bounds the
signaling polytope ``\\mathcal{C}_d^{X \\to Y}``.
The ``Y \\times X`` matrix ``\\mathbf{G}_{\\text{Q}}`` is described as having two types of rows:
* *Guessing Rows*: one column contains a non-zero element of value ``(X - d + 1)``;
* *Ambiguous Rows*: each column contains a ``1``.

For any integer ``k \\in [0,Y]`` An ambiguous guessing game is defined to have ``k`` guessing rows
and ``(Y-k)`` ammbiguous rows.
The upper bound on the  ambiguous guessing game Bell inequality is then ``\\gamma_{\\text{Q}} = d(X-d+1)``.

The `ambiguous_guessing_game` method constructs an ambiguous guessing game Bell inequality with `k`
guessing rows.

```julia
ambiguous_guessing_game(X::Int64, Y::Int64, d::Int64, k::Int64) :: BellGame
```

Parameters:
* `X` - The number of inputs
* `Y` - The number of outputs
* `d` - The signaling dimension
* `k` - Th number of guessing rows. Must satisfy `Y ≥ k ≥ 0`.

For example, the following constructed ambiguous guessing game has four guessing
rows and  3 ambiguous rows.

```jldoctest ambiguous_guessing_game
using SignalingDimension

(X, Y, d) = (6, 7, 3)

k = 4        # num guessing rows

G_Q = ambiguous_guessing_game(X, Y, d, k)

# output

7×6 BellScenario.BellGame:
 4  0  0  0  0  0
 0  4  0  0  0  0
 0  0  4  0  0  0
 0  0  0  4  0  0
 1  1  1  1  1  1
 1  1  1  1  1  1
 1  1  1  1  1  1
```

```jldoctest ambiguous_guessing_game
G_Q.β

# output

12
```

Alternatively, the `ambiguous_guessing_game` can accept a `BellScenario.LocalSignaling(X, Y, d)` scenario.

```julia
ambiguous_guessing_game(scenario :: LocalSignaling, k :: Int64) :: BellGame
```
"""
function ambiguous_guessing_game(scenario :: LocalSignaling, k :: Int64) :: BellGame
    ambiguous_guessing_game(scenario.X,scenario.Y,scenario.d,k)
end
function ambiguous_guessing_game(X::Int64, Y::Int64, d::Int64, k::Int64) :: BellGame
    if !(Y ≥ k ≥ 0)
        throw(DomainError(k, "num guessing rows `k` should satisfy `Y ≥ k ≥ 0`."))
    elseif !(min(X,Y) ≥ d ≥ 1)
        throw(DomainError(d, "`scenario.d` should satisfy `min(X, Y) ≥ d ≥ 1`."))
    end

    ambiguous_scalar = (X - d + 1)

    γ = d * ambiguous_scalar
    G = zeros(Int64, Y, X)

    for y in 1:k
        if y ≤ X
            G[y,y] = ambiguous_scalar
        else
            G[y,X] = ambiguous_scalar
        end
    end

    ambiguous_row = ones(Int64, X)
    for y in (k+1):Y
        G[y,:] =  ambiguous_row
    end

    BellGame(G, γ)
end

"""
The maximum likelihood facet ``(\\mathbf{G}_{\\text{ML}},d)`` tightly bounds the
``\\mathcal{C}_d^{N \\to N}``  signaling polytope for ``N > d > 1``.
The matrix ``\\mathbf{G}_{\\text{ML}}`` is a ``(k=1)``-guessing game (see [`k_guessing_game`](@ref)),
hence, ``\\mathbf{G}_{\\text{ML}}=\\mathbb{I}_{N}`` the ``N\\times N`` identity matrix.
The upper bound is  ``d \\geq \\langle \\mathbf{G}_{\\text{ML}}, \\mathbf{P}\\rangle``
for any channel ``\\mathbf{P}\\in\\mathcal{C}_d^{N \\to N}``.

```julia
maximum_likelihood_facet( N :: Int64, d :: Int64 ) :: BellGame
```

Construct the maximum likelihood facet for the ``\\mathcal{C}_d^{N \\to N}``
signaling polytope. Note that `N` specifies the number of inputs and outputs.
For example,

```jldoctest maximum_likelihood_facet
using SignalingDimension

N = 4    # number of inputs and outputs
d = 2    # bit signaling

G_ML = maximum_likelihood_facet(N, d)

# output

4×4 BellScenario.BellGame:
 1  0  0  0
 0  1  0  0
 0  0  1  0
 0  0  0  1
```

```jldoctest maximum_likelihood_facet
G_ML.β   # the upper bound

# output

2
```

A `DomainError` is thrown if the following requirements aren't satisfied:
* `N > 2
* `N > d > 1`
"""
function maximum_likelihood_facet(N :: Int64, d :: Int64) :: BellGame
    if N ≤ 2
        throw(DomainError(N, "Input N must satisfy `N > 2`"))
    elseif !(N > d > 1)
        throw(DomainError(d, "Input d must satisfy `N > d > 1`"))
    end

    BellGame(Matrix{Int64}(I,N,N), d)
end

"""
    anti_guessing_facet( N :: Int64, d :: Int64, ε :: Int64 ) :: BellGame

Constructs the anti-guessing facet of the ``\\mathcal{C}_d^{N \\to N}`` signaling polytope.
Input `ε` sets the size of the ``(k=N-1)``-guessing game block while the remaining
unit elements are arranged along the diagonal.
The upper bound of the anti-guessing game Bell inequality is ``\\gamma = \\varepsilon + d - 2``.
For example,

```jldoctest anti_guessing_game
using SignalingDimension

N = 6    # number of inputs and outputs
d = 2    # dit signaling
ε = 4    # anti-guessing block size

G_A = anti_guessing_facet(N,d,ε)

# output

6×6 BellScenario.BellGame:
 0  1  1  1  0  0
 1  0  1  1  0  0
 1  1  0  1  0  0
 1  1  1  0  0  0
 0  0  0  0  1  0
 0  0  0  0  0  1
```
```jldoctest anti_guessing_game
G_A.β    # upper bound

# output

4
```

Note in the above example that

A `DomainError` is thrown if the following requirements aren't satisfied:
* `N ≥ 4`
* `(N - 2) ≥ d ≥ 2`
* `(N - d + 1) ≥ ε ≥ 3`
"""
function anti_guessing_facet(N :: Int64, d :: Int64, ε :: Int64) :: BellGame
    if !(N ≥ 4)
        throw(DomainError(N, "Input N must satisfy `N ≥ 4`"))
    elseif !(N-2 ≥ d ≥ 2)
        throw(DomainError(d, "Input d must satisfy `(N - 2) ≥ d ≥ 2`"))
    elseif !(N-d+1 ≥ ε ≥ 3)
        throw(DomainError(ε, "Input ε must satisfy `(N - d + 1) ≥ ε ≥ 3`"))
    end

    err_game = ones(Int64, ε, ε) - I
    succ_game = Matrix{Int64}(I, (N-ε),(N-ε))

    BellGame(cat(
        cat(err_game, zeros(Int64, ε, (N-ε)), dims=2),
        cat(zeros(Int64, (N-ε), ε), succ_game, dims=2),
        dims = 1
    ), ε + d - 2)
end

"""
    ambiguous_guessing_facet( Y :: Int64, d :: Int64 ) :: BellGame

Constructs the ambiguous guessing facet for the ``\\mathcal{C}_d^{(Y-1)\\to Y} polytope.
This is a special case of the [`ambiguous_guessing_game`](@ref) where `Y = X - 1`.
For example,

```jldoctest ambiguous_guessing_facet
using SignalingDimension

Y = 5
d = 3

G_Q = ambiguous_guessing_facet(Y, d)

# output

5×4 BellScenario.BellGame:
 2  0  0  0
 0  2  0  0
 0  0  2  0
 0  0  0  2
 1  1  1  1
```
```jldoctest ambiguous_guessing_facet
G_Q.β

# output

6
```

A `DomainError` is thrown if the inputs don't satisfy the following requirements:
* `Y ≥ 4`
* `(Y - 2) ≥ d ≥ 2`
"""
function ambiguous_guessing_facet(Y :: Int64, d :: Int64) :: BellGame
    if !(Y ≥ 4)
        throw(DomainError(Y, "Input Y must satisfy `Y ≥ 4`"))
    elseif !(Y-2 ≥ d ≥ 2)
        throw(DomainError(d, "Input d must satisfy `(Y - 2) ≥ d ≥ 2`"))
    end

    succ_game = (Y-d)*Matrix{Int64}(I, Y-1, Y-1)

    BellGame( cat(succ_game, ones(Int64, 1, Y-1), dims = 1), d*(Y-d))
end

"""
    k_guessing_facet( Y :: Int64, d :: Int64, k :: Int64 ) :: BellGame

Constructs the ``k``-guessing facet for the ``\\mathcal{C}_d^{\\binom{Y}{k}\\to Y}``
signaling polytope. A ``k``-guessing game is tight when ``Y = d + k``.
Note that `k` is the number of unit elements in each column.

A `DomainError` is satisfied if the following requirements aren't satisfied:
* `Y == k + d`
* `Y ≥ 3`
"""
function k_guessing_facet(Y :: Int64, d :: Int64, k :: Int64) :: BellGame
    if Y != k + d
        throw(DomainError(d, "Input `d` must satisfy `Y == k + d`."))
    elseif !(Y ≥ 3)
        throw(DomainError(Y, "Input N must satisfy `Y ≥ 3`"))
    end

    k_guessing_game(Y, d, k)
end

"""
A non-negativity facet reflects the fact that ``P(y|x) \\geq 0``. These facets
bound all signaling polytopes.

    non_negativity_facet( X :: Int64, Y :: Int64 ) :: BellGame

Constructs the non-negativity game for a channel with `X` inputs and `Y` outputs.
For example

```jldoctest non_negativity_facet
using SignalingDimension

G = non_negativity_facet(3, 4)

# output

4×3 BellScenario.BellGame:
 1  0  0
 1  0  0
 1  0  0
 0  0  0
```
```jldoctest non_negativity_facet
G.β

# output

1
```

A `DomainError` is thrown if `X` or `Y` is not greater than 1.
"""
function non_negativity_facet(X :: Int64, Y :: Int64) :: BellGame
    if !(Y > 1)
        throw(DomainError(Y, "Inputs must satisfy `Y > 1`"))
    elseif !(X > 1)
        throw(DomainError(X, "Inputs must satisfy `X > 1`"))
    end

    m = zeros(Int64, Y, X)
    m[1:end-1,1] .= 1

    BellGame(m, 1)
end

"""
    coarse_grained_input_ambiguous_guessing_facet(
        Y :: Int64,
        d :: Int64
    ) :: BellGame

Constructs a canonical form for a input coarse-grained ambiguous game.
For example

```jldoctest
using SignalingDimension

Y = 4
d = 2

G = coarse_grained_input_ambiguous_guessing_facet(Y, d)

# output

4×4 BellScenario.BellGame:
 2  0  0  0
 0  2  0  0
 0  0  1  1
 1  1  1  0
```

A `DomainError` is thrown if the inputs don't satisfy the following requirements:
* `Y ≥ 4`
* `(Y - 2) ≥ d ≥ 2`
"""
function coarse_grained_input_ambiguous_guessing_facet(Y :: Int64, d :: Int64) :: BellGame
    G_Q = ambiguous_guessing_facet(Y, d)

    game = zeros(Int64,Y,Y)
    game[1:Y,1:end-1] = G_Q
    game[Y-1,Y] = game[Y-1,Y-1] - 1
    game[Y-1,Y-1] = 1

    BellGame(game, G_Q.β)
end


using DataFrames, Distributions, Gadfly

set_default_plot_size(6inch, 3inch)

beta = Beta(2,2)
Dbeta = DataFrame(x=rand(beta, 10^4))
layer1 = layer(x->pdf(beta, x), 0, 1, Theme(default_color="black"))
p1 = plot(Dbeta, layer1,
    layer(x=:x, Geom.histogram(bincount=20, density=true, limits=(min=0,))),
)
p2 = plot(Dbeta, layer1,
    layer(x=:x, Geom.histogram(bincount=20, density=true, limits=(min=0, max=1)))
)

hstack(p1, p2)


using Primes

function sieve_01(n::Int64)::Vector{Int64}
    [ x for x in 2:n if all( x % i != 0 for i in 2:x^0.5 + 1) || x == 2 ]
end

function sieve_02(n::Int64)::Vector{Int64}
    L::Vector{Int64} = []
    for i::Int64 in 2:n
        if isprime(i)
            push!(L,i)
        end
    end
    return L
end

function sieve_03(n::Int64)::Vector{Int64}
    [ x for x in 2:n if isprime(x) ]::Vector{Int64}
end


module argmaxUniProb
  export argmaxProb
  using Statistics

  """
      argmaxProb( r::Matrix{Float64} )

      compute maximum value per row, if several cols hold the maximum value, randomly choose among them.
      #Arguments
        - `r::Matrix{Float64}`: risk matrix each column represents one hypothesis.
      #Return
        - `Vector{Float64}`:    indices of maximal values.
  """
  function argmaxProb( r::Matrix{Float64} )
    index = Vector{Int64}( undef, size( r, 1 ) );
    for i in 1:1:size( index, 1 )
      index[i] = argmaxProb( r[i,:] );
    end
    return index;
  end

  """
      argmaxProb!( r::Matrix{Float64}, v::Vector{Int64} )

      compute maximum value per row, if several cols have maximum value randomly choose among them.
      #Arguments
        - `r::Matrix{Float64}`: risk matrix each column represents one hypothesis.
        - `v::Vector{Int64} `:  resulting vector.
      #Return
        - `nothing`:            nothing.
  """
  function argmaxProb!( r::Matrix{Float64}, v::AbstractVector )
    @assert size( v,1 ) == size( r, 2 );
    for i in 1:1:size( v, 1 )
     v[i] = argmaxProb( r[i,:] );
    end
  end

  """
      argmaxProb( values::Vector{Float64} )

      compute minimum, if several indices have maximum value randomly choose among them.
      #Arguments
        - `values::Vector{Float64}`: values to compute argmax on.
      #Return
        - `Vector{Float64}`:         indices of maximal values.
  """
  function argmaxProb( values::AbstractVector )
    maxPositions = collect( 1:1:size( values, 1 ) )[ values .== maximum( values ) ];
    if size( maxPositions, 1 ) > 1
      res = maxPositions[ rand( 1:size( maxPositions, 1 ) ) ];
    else
      res = maxPositions[1];
    end
    return res;
  end
end


"""
This code is based on the original work on the Scilab DoE package by the
following individuals
    Copyright (C) 2012 - 2013 - Michael Baudin
    Copyright (C) 2012 - Maria Christopoulou
    Copyright (C) 2010 - 2011 - INRIA - Michael Baudin
    Copyright (C) 2009 - Yann Collette
    Copyright (C) 2009 - CEA - Jean-Marc Martinez
and
    Copyright (C) 2014, Abraham D. Lee
who converted the package to Python.
"""

"""
$(TYPEDSIGNATURES)

Constructs a central composite design.

```jldoctest
julia> ccdesign(3)
22×3 Matrix{Float64}:
 -1.0      -1.0      -1.0
  1.0      -1.0      -1.0
 -1.0       1.0      -1.0
  1.0       1.0      -1.0
 -1.0      -1.0       1.0
  1.0      -1.0       1.0
 -1.0       1.0       1.0
  1.0       1.0       1.0
  0.0       0.0       0.0
  0.0       0.0       0.0
  ⋮
  1.82574   0.0       0.0
  0.0      -1.82574   0.0
  0.0       1.82574   0.0
  0.0       0.0      -1.82574
  0.0       0.0       1.82574
  0.0       0.0       0.0
  0.0       0.0       0.0
  0.0       0.0       0.0
  0.0       0.0       0.0

```
"""
function ccdesign(n::Int, center::Array{Int}=[4, 4], alpha::Symbol=:orthogonal, face::Symbol=:circumscribed)
    @assert n>1
    @assert alpha in [:orthogonal, :rotatable]
    @assert face in [:circumscribed, :inscribed, :faced]
    @assert length(center) == 2

    H2, a = star(n, alpha, center)
    if face == :inscribed
        H1 = explicit_fullfactorial(fullfactorial(Tuple(fill([-1, 1], n))))
        H1 = H1/a
        H2, a = star(n, :faced, [1,1])
    elseif face == :faced
        H2, a = star(n, :faced, [1,1])
        H1 = explicit_fullfactorial(fullfactorial(Tuple(fill([-1, 1], n))))
    elseif face == :circumscribed
        H1 = explicit_fullfactorial(fullfactorial(Tuple(fill([-1, 1], n))))
    else
        print("Invalid value of face")
    end

    C1 = zeros(center[1], n)
    C2 = zeros(center[2], n)

    H1 = vcat(H1, C1)
    H2 = vcat(H2, C2)
    H = vcat(H1, H2)
end

function star(n, alpha, center)
    if alpha  == :faced
        a = 1
    elseif alpha == :orthogonal
        nc = 2.0^n
        nco = center[1]
        na = 2.0*n
        nao = center[2]
        a = (n*(1 + nao/na)/(1+nco/nc))^0.5
    elseif alpha == :rotatable
        nc = 2.0^n
        a = nc^0.25
    else
        print("Invalid value of alpha")
    end

    H = zeros(2*n, n)
    for i in 1:n
         H[2*i-1:2*i, i] = [-1, 1]
    end
    H *= a
    return H, a
end


using LinearAlgebraicRepresentation, ViewerGL, SparseArrays
Lar = LinearAlgebraicRepresentation; GL = ViewerGL

n,m,p = 1,1,1
V,(VV,EV,FV,CV) = Lar.cuboidGrid([n,m,p],true)
cube = V,FV,EV

threecubes = Lar.Struct([ cube,
    Lar.t(.3,.4,.25), Lar.r(pi/5,0,0), Lar.r(0,0,pi/12), cube,
    Lar.t(-.2,.4,-.2), Lar.r(0,pi/5,0), Lar.r(0,pi/12,0), cube ])

V,FV,EV = Lar.struct2lar(threecubes)
GL.VIEW([ GL.GLGrid(V,FV,GL.COLORS[1],.5), GL.GLFrame ])

cop_EV = convert(Lar.ChainOp, Lar.coboundary_0(EV::Lar.Cells));
cop_FE = Lar.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);
W = convert(Lar.Points, V');

V, copEV, copFE, copCF = Lar.space_arrangement( W, cop_EV, cop_FE)

#triangulated_faces = Lar.triangulate(V, [copEV, copFE])
W = convert(Lar.Points, V')
V,CVs,FVs,EVs = Lar.pols2tria(W, copEV, copFE, copCF)

GL.VIEW(GL.GLExplode(V,FVs,1.5,1.5,1.5,99,1));
GL.VIEW(GL.GLExplode(V,EVs,1.5,1.5,1.5,99,1));
meshes = GL.GLExplode(V,CVs[1:end],8,4,6,99,0.5);
GL.VIEW( push!( meshes, GL.GLFrame) );


#=
HackerRank problem - 10 Days of Statistics: Day 1: Quartiles
https://www.hackerrank.com/challenges/s10-quartiles/problem

Calculate quartiles from user inputs.

Date: 2020-06-26

Issues:
  - The problem does not specify the screen output formatting,
e.g. 4 or 4.0. Possible answers are generally rational number, but the tests
seem to assume integer formatting output such as '4'. The test rejects
outputs in decimal numbers.
=#

using Printf
n = parse(Int, readline(stdin))
arr = sort(parse.(Int, split(readline(stdin))))

"""
  median(arr)

Calculate a median value in the way given in the problem.

# Arguments
- `arr`: An array of numbers
"""
function median(arr)
    n = size(arr)[1]
    n % 2 == 0 ? (arr[n ÷ 2] + arr[n ÷ 2 + 1]) / 2.0 : arr[n ÷ 2 + 1]
end

println("$(@sprintf("%d", median(arr[1:(n ÷ 2)])))")
println("$(@sprintf("%d", (median(arr))))")
println("$(@sprintf("%d", median(arr[(n % 2 == 0 ? n ÷ 2 + 1 : n ÷ 2 + 2):end])))")


## These are the figures that appear in the paper
# "Spectra of Jacobi operators via connection coefficients" by Marcus Webb and Sheehan Olver (2018)

using ApproxFun, SpectralMeasures, Plots; Plots.gr(linewidth=3,legend=false)
using ComplexPhasePortrait, PyPlot, FileIO; pygui(false) # PyPlot is good for Wegert Plots

x = linspace(-3, 3, 300)
  Z = x .+ flipdim(x,1)'*im

mkdir("Figs")

# Example for nondecreasingness of resolvents in intervals of continuity
MyResolvent = x -> real(1 + 2*sqrt(complex(x-1.)).*sqrt(complex(x+1.))-2*x + x./(9-x.^2))
  Plots.plot([-1,1],[0,0],color=:black)
  plot!([-3,-3,-3],[-3,0,3],line=:dash,color=:black,marker=:circle,markersize=2,linewidth=1)
  plot!([3,3,3],[-3,0,3],line=:dash,color=:black,marker=:circle,markersize=2,linewidth=1)
  plot!([-1,-1],[-3,3],line=:dash,color=:black,linewidth=1)
  plot!([1,1],[-3,3],line=:dash,color=:black,linewidth=1)
  plot!(-8:0.001:8,[MyResolvent(-8:0.001:-1);NaN*(-1.001:0.001:0.999);MyResolvent(1:0.001:8)],xlabel="\\lambda",ylabel="f\\(\\lambda\\)",xlims=(-7.5,7.5),ylims=(-2,2))
  Plots.savefig("Figs/secular.pdf")

# Pertrubation of the [1,1] element
for a in [0, 0.15, 0.35, 0.5, 0.75, 1.]
  Plots.plot(spectralmeasure([a],[.5]),xlims=(-2,2),ylims=(0,2),title="\\alpha=$(a)")
  Plots.savefig("Figs/spectralmeasurealpha=$(round(Int,100*a)).pdf")

  R = principalresolvent([a],[.5])
  save("Figs/principalresolventalpha=$(round(Int,100*a)).png", portrait(R(Z),PTstepmod))
  imshow(imread("Figs/principalresolventalpha=$(round(Int,100*a)).png"),extent=[-3,3,-3,3])
  PyPlot.plot([-1.,1.],[0.,0.],linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/principalresolventalpha=$(round(Int,100*a)).png",transparent=true,bbox_inches="tight")
  clf()

  r = discresolvent([a],[.5])
  save("Figs/discresolventalpha=$(round(Int,100*a)).png", portrait(r(Z),PTstepmod))
  imshow(imread("Figs/discresolventalpha=$(round(Int,100*a)).png"),extent=[-3,3,-3,3])
  PyPlot.plot(cos(linspace(0,2pi,100)),sin(linspace(0,2pi,100)),linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/discresolventalpha=$(round(Int,100*a)).png",transparent=true,bbox_inches="tight")
  clf()
end

# Perturbation of the [1,2] element
for b in [.5, 0.707, .85, 1, 1.2, 1.5]
  Plots.plot(spectralmeasure([0.],[b]/sqrt(2)),xlims=(-2,2),ylims=(0,2),title="\\beta=$(b)")
  Plots.savefig("Figs/spectralmeasurebeta=$(round(Int,1000*b)).pdf")

  R = principalresolvent([0.],[b]/sqrt(2))
  save("Figs/principalresolventbeta=$(round(Int,1000*b)).png", portrait(R(Z),PTstepmod))
  imshow(imread("Figs/principalresolventbeta=$(round(Int,1000*b)).png"),extent=[-3,3,-3,3])
  PyPlot.plot([-1.,1.],[0.,0.],linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/principalresolventbeta=$(round(Int,1000*b)).png",transparent=true,bbox_inches="tight")
  clf()

  r = discresolvent([0.],[b]/sqrt(2))
  save("Figs/discresolventbeta=$(round(Int,1000*b)).png", portrait(r(Z),PTstepmod))
  imshow(imread("Figs/discresolventbeta=$(round(Int,1000*b)).png"),extent=[-3,3,-3,3])
  PyPlot.plot(cos(linspace(0,2pi,100)),sin(linspace(0,2pi,100)),linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/discresolventbeta=$(round(Int,1000*b)).png",transparent=true,bbox_inches="tight")
  clf()
end

# 3 by 3 Perturbation
a = [3/4,-1/4,1/2]; b=[1,3/4]
  μ = spectralmeasure(a,b)
  R = principalresolvent(a,b)
  r = discresolvent(a,b)

  Plots.plot(μ,xlims=(-3,3),ylims=(0,2))
  Plots.savefig("Figs/spectralmeasure3by3pert.pdf")

  save("Figs/principalresolvent3by3pert.png", portrait(R(Z),PTstepmod))
  imshow(imread("Figs/principalresolvent3by3pert.png"),extent=[-3,3,-3,3])
  PyPlot.plot([-1.,1.],[0.,0.],linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/principalresolvent3by3pert.png",transparent=true,bbox_inches="tight")
  clf()

  save("Figs/discresolvent3by3pert.png", portrait(r(Z),PTstepmod))
  imshow(imread("Figs/discresolvent3by3pert.png"),extent=[-3,3,-3,3])
  PyPlot.plot(cos(linspace(0,2pi,100)),sin(linspace(0,2pi,100)),linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/discresolvent3by3pert.png",transparent=true,bbox_inches="tight")
  clf()

# Legendre
for n in [1,2,3,10,30,100]
  bLeg = (1:n-1)./sqrt(4*(1:n-1).^2-1)
  Plots.plot(spectralmeasure(zeros(n),bLeg),title="Legendre approximation, n = $(n)",xlims=(-1.5,1.5),ylims=(0,1),yticks=[0,.25,.5,.75,1])
  Plots.savefig("Figs/spectralmeasureLegn=$(n).pdf")

  R = principalresolvent([0.],bLeg)
  save("Figs/principalresolventLegn=$(n).png", portrait(R(Z),PTstepmod))
  imshow(imread("Figs/principalresolventLegn=$(n).png"),extent=[-3,3,-3,3])
  PyPlot.plot([-1.,1.],[0.,0.],linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/principalresolventLegn=$(n).png",transparent=true,bbox_inches="tight")
  clf()

  r = discresolvent([0.],bLeg)
  save("Figs/discresolventLegn=$(n).png", portrait(r(Z),PTstepmod))
  imshow(imread("Figs/discresolventLegn=$(n).png"),extent=[-3,3,-3,3])
  PyPlot.plot(cos(linspace(0,2pi,100)),sin(linspace(0,2pi,100)),linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/discresolventLegn=$(n).png",transparent=true,bbox_inches="tight")
  clf()
end

# Ultraspherical polynomials
γ = 0.6
  for n in [1,2,3,10,30,100]
  bUlt = .5*sqrt(((1:n).*(2γ+(0:n-1)))./((γ+(0:n-1)).*(γ+(1:n))))
  Plots.plot(spectralmeasure(zeros(n),bUlt),title="Ultraspherical\\($(γ)\\), n = $(n)",xlims=(-1.5,1.5),ylims=(0,1),yticks=[0,.25,.5,.75,1])
  Plots.savefig("Figs/spectralmeasureUltra0pt6n=$(n).pdf")

  R = principalresolvent([0.],bUlt)
  save("Figs/principalresolventUltra0pt6n=$(n).png", portrait(R(Z),PTstepmod))
  imshow(imread("Figs/principalresolventUltra0pt6n=$(n).png"),extent=[-3,3,-3,3])
  PyPlot.plot([-1.,1.],[0.,0.],linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/principalresolventUltra0pt6n=$(n).png",transparent=true,bbox_inches="tight")
  clf()

  r = discresolvent([0.],bUlt)
  save("Figs/discresolventUltra0pt6n=$(n).png", portrait(r(Z),PTstepmod))
  imshow(imread("Figs/discresolventUltra0pt6n=$(n).png"),extent=[-3,3,-3,3])
  PyPlot.plot(cos(linspace(0,2pi,100)),sin(linspace(0,2pi,100)),linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/discresolventUltra0pt6n=$(n).png",transparent=true,bbox_inches="tight")
  clf()
end

# Jacobi polynomials
α = 0.4; β = 1.9
  for n in [1,2,3,10,30,100]
  aJac = (β.^2-α.^2)./((2.*(0:n)+α+β).*(2.*(1:n+1)+α+β))
  bJac = 2*sqrt(((1:n).*(α+(1:n)).*(β+(1:n)).*(α+β+(1:n)))./((2.*(1:n)+α+β-1).*((2.*(1:n)+α+β).^2).*(2.*(1:n)+α+β+1)))
  Plots.plot(spectralmeasure(aJac,bJac),title="Jacobi\\($(α),$(β)\\), n = $(n)",xlims=(-1.5,1.5),ylims=(0,1),yticks=[0,.25,.5,.75,1])
  Plots.savefig("Figs/spectralmeasureJacobin=$(n).pdf")

  R = principalresolvent(aJac,bJac)
  save("Figs/principalresolventJacobin=$(n).png", portrait(R(Z),PTstepmod))
  imshow(imread("Figs/principalresolventJacobin=$(n).png"),extent=[-3,3,-3,3])
  PyPlot.plot([-1.,1.],[0.,0.],linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/principalresolventJacobin=$(n).png",transparent=true,bbox_inches="tight")
  clf()

  r = discresolvent(aJac,bJac)
  save("Figs/discresolventJacobin=$(n).png", portrait(r(Z),PTstepmod))
  imshow(imread("Figs/discresolventJacobin=$(n).png"),extent=[-3,3,-3,3])
  PyPlot.plot(cos(linspace(0,2pi,100)),sin(linspace(0,2pi,100)),linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/discresolventJacobin=$(n).png",transparent=true,bbox_inches="tight")
  clf()
end

# Random
mkdir("Figs")
srand(200)
longaRand = 3*(2rand(100)-1)./(1:100).^2
for n in [1,2,3,10,50,100]
  aRand = longaRand[1:n]
  Plots.plot(spectralmeasure(aRand,Float64[]),title="Random, n = $(n)",xlims=(-1.5,1.5),ylims=(0,1),yticks=[0,.25,.5,.75,1])
  Plots.savefig("Figs/spectralmeasurerandomn=$(n).pdf")

  R = principalresolvent(aRand,Float64[])
  save("Figs/principalresolventrandomn=$(n).png", portrait(R(Z),PTstepmod))
  imshow(imread("Figs/principalresolventrandomn=$(n).png"),extent=[-3,3,-3,3])
  PyPlot.plot([-1.,1.],[0.,0.],linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/principalresolventrandomn=$(n).png",transparent=true,bbox_inches="tight")
  clf()

  r = discresolvent(aRand,Float64[])
  save("Figs/discresolventrandomn=$(n).png", portrait(r(Z),PTstepmod))
  imshow(imread("Figs/discresolventrandomn=$(n).png"),extent=[-3,3,-3,3])
  PyPlot.plot(cos(linspace(0,2pi,100)),sin(linspace(0,2pi,100)),linewidth=2,color=:black)
  PyPlot.grid("on")
  PyPlot.savefig("Figs/discresolventrandomn=$(n).png",transparent=true,bbox_inches="tight")
  clf()
end

a = zeros(100); a[1] = 1
  Plots.plot(spectralmeasure(a,Float64[]),title="\\alpha=[1,0,...]",xlims=(-1,1.5),ylims=(0,4.2),yticks=[0,1,2,3,4])
  Plots.savefig("Figs/spectralmeasuredelta1.pdf")

a = zeros(100); a[2] = 1
  Plots.plot(spectralmeasure(a,Float64[]),title="\\alpha=[0,1,0,...]",xlims=(-1,1.5),ylims=(0,4.2),yticks=[0,1,2,3,4])
    Plots.savefig("Figs/spectralmeasuredelta2.pdf")

a = zeros(100); a[5] = 1
  Plots.plot(spectralmeasure(a,Float64[]),title="\\alpha=[0,0,0,0,1,0,...]",xlims=(-1,1.5),ylims=(0,4.2),yticks=[0,1,2,3,4])
  Plots.savefig("Figs/spectralmeasuredelta5.pdf")


struct RollingWindow
    WindowSize::Int
    observations::Int
    maxiter::Int
    skip::Int
end

"""
    RollingWindow(samples::Int,windowsize::Int)

Creates a RollingWindow iterator from a number of `samples` and a static `windowsize`. The iterator can be used
in for loops to iteratively return indices of a dynamic rolling window.
"""
RollingWindow(samples::Int,windowsize::Int) = RollingWindow(windowsize, samples,samples - windowsize + 1, 1)

"""
    RollingWindow(samples::Int,windowsize::Int,skip::Int)

Creates a RollingWindow iterator from a number of `samples` and a static `windowsize` where every iteration `skip` steps are skipped.
The iterator can be used in for loops to iteratively return indices of a dynamic rolling window.
"""
RollingWindow(samples::Int,windowsize::Int,skip::Int) = RollingWindow(windowsize, samples,samples - windowsize + 1, skip)

function Base.iterate( iter::RollingWindow, state = 1 )
    window = 0
    if state <= iter.maxiter
        window = collect( state : ( state + iter.WindowSize - 1) )
    else
        return nothing
    end
    return ( window, state + iter.skip  )
end

mutable struct ewma
    lambda::Float64
    center::Float64
    lastval::Float64
    rv::RunningVar
end

"""
    EWMA(Initial::Float64, Lambda::Float64) = ewma(Lambda, Initial, Initial, RunningVar(Initial))

Constructs an exponentially weighted moving average object from an initial scalar property value `Initial` and
the decay parameter `Lambda`. This defaults the center value to be the initial value.
"""
EWMA(Initial::Float64, Lambda::Float64) = ewma(Lambda, Initial, Initial, RunningVar(Initial))

"""
    EWMA(Initial::Float64, Lambda::Float64) = ewma(Lambda, Initial, Initial, RunningVar(Initial))

Constructs an exponentially weighted moving average object from an vector of scalar property values `Initial` and
the decay parameter `Lambda`. This computes the running statistcs neccesary for creating the EWMA model using the
interval provided and updates the center value to the mean of the provided values.
"""
function EWMA(Initial::Array, Lambda::Float64)
    burnin = EWMA(Initial[1], Lambda)#Call constructor above
    len = length(Initial)
    for sample in 1:len
        burnin(Initial[sample])
    end
    burnin.center = burnin.rv.m.mu
    return burnin
end

"""
    EWMA(P::ewma)(New; train = true)

Provides an EWMA score for a `New` scalar value. If ```train == true``` the model is updated to include this new value.
"""
function (P::ewma)(New; train = true)
    P.lastval = ( P.lambda * New ) + ( 1 - P.lambda ) * P.lastval
    if train == true
        Update!(P.rv, New)
    end
    return P.lastval
end

"""
    ChangeCenter(P::ewma, new::Float64)

This is a convenience function to update the center of a `P` EWMA model, to a `new` scalar value.
"""
function ChangeCenter(P::ewma, new::Float64)
    P.center .= new
end

"""
    Variance(P::ewma)

This function returns the EWMA control variance.
"""
Variance(P::ewma) = (P.lambda / (2.0 - P.lambda) ) * Variance(P.rv)

"""
    Limits(P::ewma; k = 3.0)

This function returns the upper and lower control limits with a `k` span of variance for an EWMA object `P`.
"""
Limits(P::ewma; k = 3.0) = (P.center + (k * sqrt( Variance( P ) ) ), P.center - (k * sqrt( Variance( P ) ) )  )

mutable struct NaiveForecaster
    lastval::Float64
end
"""
    NaiveForecaster( univariate::Array )

Makes a forecaster model that simply predicts the last value it has seen for all future values.
"""
NaiveForecast(univariate::Array) = NaiveForecaster( last( univariate ) )

"""
    update!( model::NaiveForecaster, newdata::Float64 )

Makes a forecast model that simply predicts the last value it has seen for all future values.
"""
update!( model::NaiveForecaster, newdata::Float64 ) = model.lastval .= newdata

"""
    update( model::NaiveForecaster, newdata::Float64 )

returns an updated forecast model that simply predicts the last value it has seen for all future values.
"""
update(model::NaiveForecaster, newdata::Float64) = return NaiveForecaster( newdata )

"""
    (nf::NaiveForecaster)()

Predicts with a naive forecaster model.
"""
(nf::NaiveForecaster)() = nf.lastval

mutable struct SimpleAverage
    runmean::RunningMean
end

"""
    SimpleAverage( univariate )
"""
SimpleAverage(initial::Float64) = SimpleAverage( RunningMean( initial) )

"""
    SimpleAverage( univariate::Array )
"""
function SimpleAverage(univariate::Array)
    rm = RunningMean( univariate[1])
    for i in 2:length(univariate)
        Update!( rm, univariate[ i ] )
    end
    return SimpleAverage( rm )
end

"""
    Update( model::naiveforecaster, newdata::Float64 )
"""
function Update(model::SimpleAverage, newdata::Float64)
    return Update!( model.runmean, newdata )
end

"""
    (sa::SimpleAverage)()

Predicts with a naive forecaster model.
"""
(sa::SimpleAverage)() = sa.runmean.mu

struct EchoStateNetwork
    Winput::Array
    W::Array
    Woutput::Array
    States::Array
    LastState::Array
    L2::Float64
    alpha::Float64
    bias
end

"""
    EchoStateNetwork(X, Y, Reservoir = 1000;
                        L2 = 1e-8, alpha = 0.25, SpectralRadius = 1.00, Sparsity = 0.99, Noise = -1.00,
                        bias = true, burnin = 0)

    Creates an echo state network model/object.
"""
function EchoStateNetwork(X, Y, Reservoir = 1000;
                        L2 = 1e-8, alpha = 0.25, SpectralRadius = 1.00, Sparsity = 0.99, Noise = -1.00,
                        bias = true, burnin = 0)
    X = forceMatrix( X ) ; Y = forceMatrix( Y )#Ensure inputs are Array2's
    (Obs, Vars) = size( X )
    #Allocate memory for the collected states matrix
    States = zeros(bias + Vars + Reservoir, Obs - burnin - 1)
    Predictions = zeros( size( Y ) )
    #Initialize input weight matrix and recurrent weight matrix
    Winput = rand( Reservoir, bias + Vars) .* -2 .+ 1
    #Generate Resevoir matrix with random uniform values from -1 to +1
    W = rand( Reservoir, Reservoir ) .* -2 .+ 1
    #Create a 'mask' from the Bernouli distribution (binomial distribution size = 1)
    Dropout = rbinomial(1.0 - Sparsity, Reservoir, Reservoir )
    #Now we need to introduce sparsity in our reservoir connections based on the mask
    W = W .* Dropout
    #Normalize and set spectral radius of reservoir weight matrix
    W = W .* ( SpectralRadius / abs( LinearAlgebra.eigen( W ).values[ 1 ] ) )
    #Run the reservoir with the data and collect X
    state = zeros(Reservoir, 1)
    #Go through each time step and compute the model
    for t in 1 : (Obs - 1)
        u = bias ? vcat( 1, X[t,:] ) : X[t,:]
        state = ((1.0 - alpha) * state) .+ (alpha * tanh.( Winput * u .+ W * state ))
        #If we are past the burnin number, start saving the model
        if t > burnin
            States[ :, t - burnin ] = vcat(u, state)
        end
    end
    if Noise > 0.0
        VS = size(S)[2]
        for obs in 1:Obs
            mini = reduce( min, States[obs,:] )
            maxi = reduce( max, States[obs,:] )
            States[obs,:] .+= States[obs,:] .* ( rand( 1, VS ) .* Noise .- (Noise/2.0) )
        end
    end
    #Train the output layer from the random reservoir via least squares
    Woutput = Y[ (burnin + 2) : end, :]' * States' * Base.inv( States * States' .+ L2 * LinearAlgebra.Diagonal(ones(bias + Vars + Reservoir) ) )
    return EchoStateNetwork( Winput, W, Woutput, States, state, L2, alpha, bias )
end

"""
    TuneRidge(Y, model::EchoStateNetwork, L2, burnin = 0 )

    This function modifies an ESN model's output weights without reconstructing the dynamic resevoir.
    Basically this is a time saver!
"""
function TuneRidge(Y, model::EchoStateNetwork, L2, burnin = 0 )
    model.L2 .= L2
    model.Woutput .= Y[ (burnin + 2) : end]' * model.States' * Base.inv( model.States * model.States' .+ L2 * diag(ones(bias + Vars + Reservoir) ) )
end

"""
    PredictFn(model::EchoStateNetwork, X; UseLastState = true, burnin = 0)

    This function uses X data to predict Y values from a stored ESN model.
    Burnin can be used here, but it simply rejects the first `N` samples from storage/allocation.
    UseLastState will employ the last known state from the training set as the first state estimate (reccomended)
"""
function PredictFn(model::EchoStateNetwork, X; UseLastState = true, burnin = 0)
    (Obs,Vars) = size(X)
    #Run the trained ESN using last state as starting point
    state = UseLastState ? model.LastState : state = zeros( size( model.Reservoir ) )
    Prediction = zeros( size(model.Woutput)[1], Obs - burnin)

    u = X[1,:]
    for t in 1 : Obs
        u = model.bias ? vcat( 1, X[t,:]) : X[t,:]
        state = ((1.0 - model.alpha ) * state) .+ (model.alpha * tanh.( model.Winput * u .+ model.W * state ) )
        if t > burnin
            Prediction[ : , t - burnin ] = model.Woutput * vcat(u, state)
        end
        if t < Obs
            u = X[t + 1,:] #Load next sample for prediction
        end
    end#End for

    return Prediction
end


using PointSampler
using Test

#%% Unit Line Design (1D) with 10 points, including zero
P1 = monte_carlo_th( 10, 1 )

@test [0.0] ∈ P1
@test length(P1) == 10
@test max( P1... )[end] <= 1.0
@test min( P1... )[end] >= 0.0

#%% Unit Square Design (2D) with 20 points, including zero vector
P2 = monte_carlo_th( 20, 2 )

@test zeros(2) ∈ P2
@test length(P2) == 20
@test all( maximum( hcat(P2...), dims = 2 ) .<= 1.0 )
@test all( minimum( hcat(P2...), dims = 2 ) .>= 0.0 )

#%% Designs in 3D using seeds

seed_set = [ rand(3) for i = 1 : 10 ];
seed_base = deepcopy(seed_set);

# %% Less points requested than provided as seeds
P3_1 = monte_carlo_th( 5, 3; seeds = seed_set )
@test P3_1 ⊆ seed_set
@test all(seed_set .== seed_base)

#%% test whether seeds are contained in returned set
P3_2 = monte_carlo_th( 15, 3; seeds = seed_set )
#@test seed_set ⊆ P3_2
@test all( P3_2[1:10] .≈ seed_set )
@test all( seed_set .== seed_base )

#%% Weed out seeds that are violating box constraints
lb = .45 .* ones(3);
ub = .55 .* ones(3);

P3_3 = monte_carlo_th( 20, lb, ub; seeds = seed_set, clean_seeds = true )

@test all(seed_set .== seed_base)
@test all( maximum( hcat(P3_3...), dims = 2 ) .<= ub )
@test all( minimum( hcat(P3_3...), dims = 2 ) .>= lb )

# %%Keep seeds that are violating box constraints but sample only in box
P3_4 = monte_carlo_th( 20, lb, ub; seeds = seed_set, clean_seeds = false )

#@test seed_set ⊆ P3_4
@test all( P3_4[1:10] .≈ seed_set )
P3_4 = P3_4[11 : end ]
@test all( all( p .<= ub ) for p ∈ P3_4 )
@test all( all( p .>= lb ) for p ∈ P3_4 )


@testset "$TEST $G" begin

g5 = DG(4)
add_edge!(g5,1,2); add_edge!(g5,2,3); add_edge!(g5,1,3); add_edge!(g5,3,4)

y = closeness_centrality(g5; normalize=false)
z = closeness_centrality(g5)

@test y == [0.75, 0.6666666666666666, 1.0, 0.0]
@test z == [0.75, 0.4444444444444444, 0.3333333333333333, 0.0]

g = G(5)
add_edge!(g,1,2)
z = closeness_centrality(g)
@test z[1] == z[2] == 0.25
@test z[3] == z[4] == z[5] == 0.0

end # testset