"""
```
ct_kalman_filter(y, T, R, C, Q, Z, D, E, tspan, s_0 = Vector(), P_0 = Matrix();
    outputs = [:loglh, :pred, :filt], Nt0 = 0)
tspan - Number of time units betwen each observation
"""
function ct_kalman_filter(y::Matrix{S},
    T::Matrix{S}, R::Matrix{S}, C::Vector{S},
    Q::Matrix{S}, Z::Matrix{S}, D::Vector{S}, E::Matrix{S}, tspan::Float64,
    s_0::Vector{S} = Vector{S}(undef, 0), P_0::Matrix{S} = Matrix{S}(undef, 0, 0);
    outputs::Vector{Symbol} = [:loglh, :pred, :filt],
    Nt0::Int = 0, method = Tsit5(),
    reltol::Float64 = 1e-8, abstol::Float64 = 1e-8) where {S<:AbstractFloat}

    # Determine outputs
    return_loglh = :loglh in outputs
    return_pred  = :pred  in outputs
    return_filt  = :filt  in outputs

    # Dimensions
    Ns = size(T, 1) # number of states
    Nt = size(y, 2) # number of periods of data

    # Initialize inputs and outputs
    k = StateSpaceRoutines.KalmanFilter(T, R, C, Q, Z, D, E, s_0, P_0)

    mynan = convert(S, NaN)
    loglh  = return_loglh ? fill(mynan, Nt)         : Vector{S}(0)
    s_pred = return_pred  ? fill(mynan, Ns, Nt)     : Matrix{S}(0, 0)
    P_pred = return_pred  ? fill(mynan, Ns, Ns, Nt) : Array{S, 3}(0, 0, 0)
    s_filt = return_filt  ? fill(mynan, Ns, Nt)     : Matrix{S}(0, 0)
    P_filt = return_filt  ? fill(mynan, Ns, Ns, Nt) : Array{S, 3}(0, 0, 0)

    # Populate initial states
    s_0 = k.s_t
    P_0 = k.P_t

    # Loop through periods t
    for t = 1:Nt
        # Forecast
        forecast!(k, tspan; method = method, reltol = reltol, abstol = abstol)
        if return_pred
            s_pred[:,    t] = k.s_t
            P_pred[:, :, t] = k.P_t
        end

        # Update and compute log-likelihood
        StateSpaceRoutines.update!(k, y[:, t]; return_loglh = return_loglh)
        if return_filt
            s_filt[:,    t] = k.s_t
            P_filt[:, :, t] = k.P_t
        end
        if return_loglh
            loglh[t]        = k.loglh_t
        end

        # Update s_0 and P_0 if Nt0 > 0
        if t == Nt0
            s_0 = k.s_t
            P_0 = k.P_t
        end
    end

    # Populate final states
    s_T = k.s_t
    P_T = k.P_t

    # Remove presample periods
    loglh, s_pred, P_pred, s_filt, P_filt =
        StateSpaceRoutines.remove_presample!(Nt0, loglh, s_pred, P_pred, s_filt, P_filt; outputs = outputs)

    return loglh, s_pred, P_pred, s_filt, P_filt, s_0, P_0, s_T, P_T
end

"""
```
ct_kalman_filter(y, T, R, C, Q, Z, D, E, tspan, s_0 = Vector(), P_0 = Matrix();
    outputs = [:loglh, :pred, :filt], Nt0 = 0)
tspan - Number of time units betwen each observation
n_subinterval - Number of subintervals to use ODE integration with
"""
function ct_kalman_filter(y::Matrix{S},
                          T::Matrix{S}, R::Matrix{S}, C::Vector{S}, Q::Matrix{S},
                          Z::Matrix{S}, D::Vector{S}, E::Matrix{S}, tspan::Float64, n_subinterval::Int64,
                          s_0::Vector{S} = Vector{S}(0), P_0::Matrix{S} = Matrix{S}(0, 0);
                          outputs::Vector{Symbol} = [:loglh, :pred, :filt],
                          Nt0::Int = 0, method = Tsit5(),
                          reltol::Float64 = 1e-8, abstol::Float64 = 1e-8) where {S<:AbstractFloat}

    # Determine outputs
    return_loglh = :loglh in outputs
    return_pred  = :pred  in outputs
    return_filt  = :filt  in outputs

    # Dimensions
    Ns = size(T, 1) # number of states
    Nt = size(y, 2) # number of periods of data

    # Initialize inputs and outputs
    k = StateSpaceRoutines.KalmanFilter(T, R, C, Q, Z, D, E, s_0, P_0)

    mynan = convert(S, NaN)
    loglh  = return_loglh ? fill(mynan, Nt)         : Vector{S}(0)
    s_pred = return_pred  ? fill(mynan, Ns, Nt)     : Matrix{S}(0, 0)
    P_pred = return_pred  ? fill(mynan, Ns, Ns, Nt) : Array{S, 3}(0, 0, 0)
    s_filt = return_filt  ? fill(mynan, Ns, Nt)     : Matrix{S}(0, 0)
    P_filt = return_filt  ? fill(mynan, Ns, Ns, Nt) : Array{S, 3}(0, 0, 0)

    # Populate initial states
    s_0 = k.s_t
    P_0 = k.P_t

    # Loop through periods t
    for t = 1:Nt
        # Forecast
        forecast!(k, tspan, n_subinterval; method = method, reltol = reltol, abstol = abstol)
        if return_pred
            s_pred[:,    t] = k.s_t
            P_pred[:, :, t] = k.P_t
        end

        # Update and compute log-likelihood
        StateSpaceRoutines.update!(k, y[:, t]; return_loglh = return_loglh)
        if return_filt
            s_filt[:,    t] = k.s_t
            P_filt[:, :, t] = k.P_t
        end
        if return_loglh
            loglh[t]        = k.loglh_t
        end

        # Update s_0 and P_0 if Nt0 > 0
        if t == Nt0
            s_0 = k.s_t
            P_0 = k.P_t
        end
    end

    # Populate final states
    s_T = k.s_t
    P_T = k.P_t

    # Remove presample periods
    loglh, s_pred, P_pred, s_filt, P_filt =
        remove_presample!(Nt0, loglh, s_pred, P_pred, s_filt, P_filt; outputs = outputs)

    return loglh, s_pred, P_pred, s_filt, P_filt, s_0, P_0, s_T, P_T
end

"""
Forecast functions
"""
function forecast!(k::StateSpaceRoutines.KalmanFilter, tspan::Float64; method = Tsit5(),
                   reltol::Float64 = 1e-8, abstol::Float64 = 1e-8) where {S<:AbstractFloat}
    T, R, Q = k.T, k.R, k.Q
    s0, P0 = k.s_t, k.P_t
    f(u,p,t) = T*u
    g(u,p,t) = T*u + u*T' + R*Q*R'
    tspan_int = (0., tspan)
    s_prob = ODEProblem(f, s0, tspan_int)
    P_prob = ODEProblem(g, P0, tspan_int)
    s_sol = DifferentialEquations.solve(s_prob, method, reltol = reltol, abstol = abstol)
    P_sol = DifferentialEquations.solve(P_prob, method, reltol = reltol, abstol = abstol)

    k.s_t = s_sol.u[end]
    k.P_t = P_sol.u[end]
end

# INCOMPLETE
# Not sure how to write down the ODE system correctly since I don't have the same
# kind of recursion that I can do otherwise. Moreover, I effectively will get the
# same predicted states regardless, so I'm not sure if this will improve the
# measurement equation or accuracy.
function forecast!(k::StateSpaceRoutines.KalmanFilter, tspan::Float64, n_subinterval::Int64; method = Tsit5(),
                   reltol::Float64 = 1e-8, abstol::Float64 = 1e-8) where {S<:AbstractFloat}
    T, R, Q = k.T, k.R, k.Q
    s_filt, P0 = k.s_t, k.P_t
    indiv_dim = size(T, 1)
    s0 = s_filt[1:indiv_dim]
    f(u,p,t) = T*u
    g(u,p,t) = T*u + u*T' + R*Q*R'
    tspan_int = (0., span)
    s_prob = ODEProblem(f, s0, tspan_int)
    P_prob = ODEProblem(g, P0, tspan_int)
    s_sol = solvect(s_prob, method, reltol = reltol, abstol = abstol)
    P_sol = solvect(P_prob, method, reltol = reltol, abstol = abstol)
    k.s_t = s_sol.u
    k.P_t = P_sol.u
end


# Text provided under a Creative Commons Attribution license, CC-BY.  All code is made available under the FSF-approved BSD-3 license.  (c) Lorena A. Barba, Gilbert F. Forsyth 2017. Thanks to NSF for support via CAREER award #1149784.
# [@LorenaABarba](https://twitter.com/LorenaABarba)
# Translation to Julia by [@miguelraz_](https://twitter.com/miguelraz_)

# ## 12 steps to Navier–Stokes
# =====
# ***

# This will be a milestone! We now get to Step 8: Burgers' equation. We can learn so much more from this equation. It plays a very important role in fluid mechanics, because it contains the full convective nonlinearity of the flow equations, and at the same time there are many known analytical solutions.

# ## Step 8: Burgers' Equation in 2D
# ----
# ***

# Remember, Burgers' equation can generate discontinuous solutions from an initial condition that is smooth, i.e., can develop "shocks." We want to see this in two dimensions now!

# Here is our coupled set of PDEs:

# $$
# \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} + v \frac{\partial u}{\partial y} = \nu \; \left(\frac{\partial ^2 u}{\partial x^2} + \frac{\partial ^2 u}{\partial y^2}\right)$$

# $$
# \frac{\partial v}{\partial t} + u \frac{\partial v}{\partial x} + v \frac{\partial v}{\partial y} = \nu \; \left(\frac{\partial ^2 v}{\partial x^2} + \frac{\partial ^2 v}{\partial y^2}\right)$$

# We know how to discretize each term: we've already done it before!

# $$
# \begin{split}
# & \frac{u_{i,j}^{n+1} - u_{i,j}^n}{\Delta t} + u_{i,j}^n \frac{u_{i,j}^n-u_{i-1,j}^n}{\Delta x} + v_{i,j}^n \frac{u_{i,j}^n - u_{i,j-1}^n}{\Delta y} = \\
# & \qquad \nu \left( \frac{u_{i+1,j}^n - 2u_{i,j}^n+u_{i-1,j}^n}{\Delta x^2} + \frac{u_{i,j+1}^n - 2u_{i,j}^n + u_{i,j-1}^n}{\Delta y^2} \right)
# \end{split}
# $$

# $$
# \begin{split}
# & \frac{v_{i,j}^{n+1} - v_{i,j}^n}{\Delta t} + u_{i,j}^n \frac{v_{i,j}^n-v_{i-1,j}^n}{\Delta x} + v_{i,j}^n \frac{v_{i,j}^n - v_{i,j-1}^n}{\Delta y} = \\
# & \qquad \nu \left( \frac{v_{i+1,j}^n - 2v_{i,j}^n+v_{i-1,j}^n}{\Delta x^2} + \frac{v_{i,j+1}^n - 2v_{i,j}^n + v_{i,j-1}^n}{\Delta y^2} \right)
# \end{split}
# $$

# And now, we will rearrange each of these equations for the only unknown: the two components $u,v$ of the solution at the next time step:

# $$
# \begin{split}
# u_{i,j}^{n+1} = & u_{i,j}^n - \frac{\Delta t}{\Delta x} u_{i,j}^n (u_{i,j}^n - u_{i-1,j}^n)  - \frac{\Delta t}{\Delta y} v_{i,j}^n (u_{i,j}^n - u_{i,j-1}^n) \\
# &+ \frac{\nu \Delta t}{\Delta x^2}(u_{i+1,j}^n-2u_{i,j}^n+u_{i-1,j}^n) + \frac{\nu \Delta t}{\Delta y^2} (u_{i,j+1}^n - 2u_{i,j}^n + u_{i,j-1}^n)
# \end{split}
# $$

# $$
# \begin{split}
# v_{i,j}^{n+1} = & v_{i,j}^n - \frac{\Delta t}{\Delta x} u_{i,j}^n (v_{i,j}^n - v_{i-1,j}^n) - \frac{\Delta t}{\Delta y} v_{i,j}^n (v_{i,j}^n - v_{i,j-1}^n) \\
# &+ \frac{\nu \Delta t}{\Delta x^2}(v_{i+1,j}^n-2v_{i,j}^n+v_{i-1,j}^n) + \frac{\nu \Delta t}{\Delta y^2} (v_{i,j+1}^n - 2v_{i,j}^n + v_{i,j-1}^n)
# \end{split}
# $$

using Plots

# ## Variable declarations
#####
nₓ = 41
ny = 41
nₜ = 120
c = 1
Δx = 2 / (nₓ - 1)
Δy = 2 / (ny - 1)
σ = 0.0009
ν = 0.01
Δt = σ * Δx * Δy / ν

rangex = range(start = 0, stop = 2, length = nₓ)
rangey = range(start = 0, stop = 2, length = ny)
indices = CartesianIndices((rangex, rangey))

uₙ = ones(ny, nₓ) ##create a 1xn vector of 1's
u  = ones(ny, nₓ)
v  = ones(ny, nₓ)
vₙ  = ones(ny, nₓ)
comb  = ones(ny, nₓ)

# ## Assign initial conditions
##set hat function I.C. : u(.5<=x<=1 && .5<=y<=1 ) is 2
u[indices] .= 2.0
##set hat function I.C. : v(.5<=x<=1 && .5<=y<=1 ) is 2
v[indices] .= 2.0

#####

### (plot ICs)
# TODO PLOT
fig = pyplot.figure(figsize=(11, 7), dpi=100)
ax = fig.gca(projection='3d')
X, Y = numpy.meshgrid(x, y)
ax.plot_surface(X, Y, u[:], cmap=cm.viridis, rstride=1, cstride=1)
ax.plot_surface(X, Y, v[:], cmap=cm.viridis, rstride=1, cstride=1)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$');

# TODO fix variable names, make as function
for n in 1:nₜ
    un = u.copy()
    vn = v.copy()

    u[2:end, 2:end] = @views (un[2:end, 2:end] -
                     dt / dx * un[2:end, 2:end] *
                     (un[2:end, 2:end] - un[2:end, 1:end-2]) -
                     dt / dy * vn[2:end, 2:end] *
                     (un[2:end, 2:end] - un[1:end-2, 2:end]) +
                     nu * dt / dx^2 *
                     (un[2:end,3:end] - 2 * un[2:end, 2:end] + un[2:end, 1:end-2]) +
                     nu * dt / dy^2 *
                     (un[2:, 2:end] - 2 * un[2:end, 2:end] + un[1:end-2, 2:end]))

    v[2:end, 2:end] = (vn[2:end, 2:end] -
                     dt / dx * un[2:end, 2:end] *
                     (vn[2:end, 2:end] - vn[2:end, 1:end-2]) -
                     dt / dy * vn[2:end, 2:end] *
                    (vn[2:end, 2:end] - vn[1:end-2, 2:end]) +
                     nu * dt / dx^2 *
                     (vn[2:end, 3:end] - 2 * vn[2:end, 2:end] + vn[2:end, 1:end-2]) +
                     nu * dt / dy^2 *
                     (vn[2:, 2:end] - 2 * vn[2:end, 2:end] + vn[1:end-2, 2:end]))

    u[begin, :] = 1
    u[end, :] = 1
    u[:, begin] = 1
    u[:, end] = 1

    v[begin, :]   = 1
    v[end, :] = 1
    v[:, end] = 1
    v[:, begin] = 1
end

# TODO PLOT
fig = pyplot.figure(figsize=(11, 7), dpi=100)
ax = fig.gca(projection='3d')
X, Y = numpy.meshgrid(x, y)
ax.plot_surface(X, Y, u, cmap=cm.viridis, rstride=1, cstride=1)
ax.plot_surface(X, Y, v, cmap=cm.viridis, rstride=1, cstride=1)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$');

## Learn More

# The video lesson that walks you through the details for Steps 5 to 8 is **Video Lesson 6** on You Tube:

# TODO YOUTUBE
# src="https://www.youtube.com/embed/tUg_dE3NXoY"


ProjDir = @__DIR__

include("./gen6-6-8.jl")

fig = plot_models([m6_6s, m6_7s, m6_8s], :waic)
plot(fig)


export batch_normalize,
    batch_normalize!,
    rotmat,
    rand_hermitian,
    rand_unitary,
    sprand_hermitian,
    sprand_unitary,
    general_controlled_gates,
    general_c1_gates,
    linop2dense,
    # kron
    hilbertkron,
    batched_kron!,
    batched_kron,
    kron!,
    # norms
    trnorm,
    nucnorm,
    # fidelity
    pure_state_fidelity,
    density_fidelity,
    purification_fidelity,
    # matrix tools
    autostatic,
    rot_mat

using LuxurySparse, LinearAlgebra, BitBasis, SparseArrays
import LinearAlgebra: svdvals

"""
nucnorm(m)

Computes the nuclear norm of a matrix `m`.
"""
function nucnorm(m::AbstractMatrix)
    norm(svdvals(m),1)
end

"""
trnorm(m)

Computes the trace norm of a matrix `m`.
"""
trnorm(m::AbstractMatrix) = nucnorm(m)

"""
    batch_normalize!(matrix)

normalize a batch of vector.
"""
function batch_normalize!(s::AbstractMatrix, p::Real=2)
    B = size(s, 2)
    for i = 1:B
        normalize!(view(s, :, i), p)
    end
    s
end

"""
    batch_normalize

normalize a batch of vector.
"""
function batch_normalize(s::AbstractMatrix, p::Real=2)
    ts = copy(s)
    batch_normalize!(ts, p)
end

"""
    hilbertkron(num_bit::Int, gates::Vector{AbstractMatrix}, locs::Vector{Int}) -> AbstractMatrix

Return general kronecher product form of gates in Hilbert space of `num_bit` qubits.

* `gates` are a list of matrices.
* `start_locs` should have the same length as `gates`, specifing the gates starting positions.
"""
function hilbertkron(num_bit::Int, ops::Vector{<:AbstractMatrix}, start_locs::Vector{Int})
    sizes = [log2dim1(op) for op in ops]
    start_locs = num_bit .- start_locs .- sizes .+ 2

    order = sortperm(start_locs)
    sorted_ops = ops[order]
    sorted_start_locs = start_locs[order]
    num_ids = vcat(sorted_start_locs[1]-1, diff(push!(sorted_start_locs, num_bit+1)) .- sizes[order])

    _wrap_identity(sorted_ops, num_ids)
end

# kron, and wrap matrices with identities.
function _wrap_identity(data_list::Vector{T}, num_bit_list::Vector{Int}) where T<:AbstractMatrix
    length(num_bit_list) == length(data_list) + 1 || throw(ArgumentError())
    ⊗ = kron
    reduce(zip(data_list, num_bit_list[2:end]); init=IMatrix(1 << num_bit_list[1])) do x, y
        x ⊗ y[1] ⊗ IMatrix(1<<y[2])
    end
end

batched_kron(a, b, c, xs...) = Base.afoldl(batched_kron, (batched_kron)((batched_kron)(a,b),c), xs...)

function batched_kron(A::AbstractArray{T, 3}, B::AbstractArray{S, 3}) where {T, S}
    @assert size(A, 3) == size(B, 3) "batch size mismatch"
    C = Array{Base.promote_op(*,T,S), 3}(undef, size(A, 1) * size(B, 1), size(A, 2) * size(B, 2), size(A, 3))
    return batched_kron!(C, A, B)
end

function batched_kron!(C::Array{T, 3}, A::AbstractArray{T1, 3}, B::AbstractArray{T2, 3}) where {T, T1, T2}
    @assert !Base.has_offset_axes(A, B)
    m, n = size(A)
    p, q = size(B)
    @inbounds for k in 1:size(C, 3)
        for s in 1:n, r in 1:m, w in 1:q, v in 1:p
            C[p * (r - 1) + v, q * (s - 1) + w, k] = A[r, s, k] * B[v, w, k]
        end
    end
    return C
end

# NOTE: JuliaLang/julia/pull/31069 includes this function
function kron!(C::AbstractMatrix{T}, A::AbstractMatrix{T1}, B::AbstractMatrix{T2}) where {T,T1,T2}
    @assert !Base.has_offset_axes(A, B)
    m = 1
    @inbounds for j = 1:size(A, 2), l = 1:size(B, 2), i = 1:size(A, 1)
        aij = A[i,j]
        for k = 1:size(B, 1)
            C[m] = aij * B[k,l]
            m += 1
        end
    end
    return C
end

"""
    general_controlled_gates(num_bit::Int, projectors::Vector{Tp}, cbits::Vector{Int}, gates::Vector{AbstractMatrix}, locs::Vector{Int}) -> AbstractMatrix

Return general multi-controlled gates in hilbert space of `num_bit` qubits,

* `projectors` are often chosen as `P0` and `P1` for inverse-Control and Control at specific position.
* `cbits` should have the same length as `projectors`, specifing the controling positions.
* `gates` are a list of controlled single qubit gates.
* `locs` should have the same length as `gates`, specifing the gates positions.
"""
function general_controlled_gates(
    n::Int,
    projectors::Vector{<:AbstractMatrix},
    cbits::Vector{Int},
    gates::Vector{<:AbstractMatrix},
    locs::Vector{Int}
)
    IMatrix(1<<n) - hilbertkron(n, projectors, cbits) +
        hilbertkron(n, vcat(projectors, gates), vcat(cbits, locs))
end

"""
    general_c1_gates(num_bit::Int, projector::AbstractMatrix, cbit::Int, gates::Vector{AbstractMatrix}, locs::Vector{Int}) -> AbstractMatrix

general (low performance) construction method for control gate on different lines.
"""
general_c1_gates(num_bit::Int, projector::Tp, cbit::Int, gates::Vector{Tg}, locs::Vector{Int}) where {Tg<:AbstractMatrix, Tp<:AbstractMatrix} =
    hilbertkron(num_bit, [IMatrix(2) - projector], [cbit]) + hilbertkron(num_bit, vcat([projector], gates), vcat([cbit], locs))

"""
    rotmat(M::AbstractMatrix, θ::Real)

Returns rotated `M`: ``exp(-\\frac{imθ}{2} M)``.
"""
rotmat(M::AbstractMatrix, θ::Real) = exp(-im * θ/2 * M)

"""
    linop2dense([T=ComplexF64], linear_map!::Function, n::Int) -> Matrix

Returns the dense matrix representation given linear map function.
"""
linop2dense(linear_map!::Function, n::Int) = linop2dense(ComplexF64, linear_map!, n)
linop2dense(::Type{T}, linear_map!::Function, n::Int) where T = linear_map!(Matrix{T}(I, 1<<n, 1<<n))

################### Fidelity ###################

"""
    density_fidelity(ρ1, ρ2)

General fidelity (including mixed states) between two density matrix for qubits.

# Definition

```math
F(ρ, σ)^2 = tr(ρσ) + 2 \\sqrt{det(ρ)det(σ)}
```
"""
function density_fidelity(ρ1::AbstractMatrix, ρ2::AbstractMatrix)
    return sqrt( tr(ρ1 * ρ2) + 2 * sqrt(det(ρ1) * det(ρ2)) )
end

"""
    pure_state_fidelity(v1::Vector, v2::Vector)

fidelity for pure states.
"""
pure_state_fidelity(v1::Vector, v2::Vector) = abs(v1'*v2)

"""
    purification_fidelity(m1::Matrix, m2::Matrix)

Fidelity for mixed states via purification.

Reference:
    http://iopscience.iop.org/article/10.1088/1367-2630/aa6a4b/meta
"""
function purification_fidelity(m1::Matrix, m2::Matrix)
    O = m1'*m2
    return tr(sqrt(O*O'))
end

"""
    rand_unitary([T=ComplexF64], N::Int) -> Matrix

Create a random unitary matrix.
"""
rand_unitary(N::Int) = rand_unitary(ComplexF64, N)
rand_unitary(::Type{T}, N::Int) where T = qr(randn(T, N, N)).Q |> Matrix

"""
    sprand_unitary([T=ComplexF64], N::Int, density) -> SparseMatrixCSC

Create a random sparse unitary matrix.
"""
sprand_unitary(N::Int, density::Real) = sprand_unitary(ComplexF64, N, density)
sprand_unitary(::Type{T}, N::Int, density::Real) where T = SparseMatrixCSC(qr(sprandn(T, N, N, density)).Q)

"""
    rand_hermitian([T=ComplexF64], N::Int) -> Matrix

Create a random hermitian matrix.
"""
rand_hermitian(N::Int) = rand_hermitian(ComplexF64, N)

function rand_hermitian(::Type{T}, N::Int) where T
    A = randn(T, N, N)
    A + A'
end

"""
    sprand_hermitian([T=ComplexF64], N, density)

Create a sparse random hermitian matrix.
"""
sprand_hermitian(N::Int, density) = sprand_hermitian(ComplexF64, N, density)

function sprand_hermitian(::Type{T}, N::Int, density::Real) where T
    A = sprandn(T, N, N, density)
    return A + A'
end

@static if VERSION < v"1.1.0"
    function SparseArrays.sprandn(::Type{T}, n::Int, m::Int, density::Real) where T <: Real
        return T.(sprandn(n, m, density))
    end

    function SparseArrays.sprandn(::Type{Complex{T}}, n::Int, m::Int, density::Real) where T <: Real
        return T.(sprandn(n, m, density)) + im * T.(sprandn(n, m, density))
    end
end

"""
    autostatic(A[; threshold=8])

Staticize dynamic array `A` by a `threshold`.
"""
autostatic(A::AbstractVecOrMat; threshold::Int=8) = length(A) > (1 << threshold) ? A : staticize(A)

# General definition
function rot_mat(::Type{T}, gen::AbstractMatrix, theta::Real) where {N, T}
    I = IMatrix{size(gen, 1), T}()
    m = I * cos(theta / 2) - im * sin(theta / 2) * gen
    if eltype(m) != T
        m2 = similar(m, T)
        copyto!(m2, m)
        return m2
    else
        return m
    end
end


"""
Hequation.jl

This file contains the function/Jacobian evaluations for
the Chandrasekhar H-equation examples and everything you should
need to run them.

If you only want to run the examples, you should not have to look
at the code.
"""
### And now for the functions ...
"""
function heqJ!(FP,F,x,pdata)

The is the Jacobian evaluation playing by nsol rules. The
precomputed data is a big deal for this one.
"""
function heqJ!(FP::Array{T,2}, F, x, pdata) where {T<:Real}
    pseed = pdata.pseed
    mu = pdata.mu
    n = length(x)
    #
    # Look at the formula in the notebook and you'll see what I did here.
    #
    pmu = pdata.pmu
    Gfix = pdata.gtmp
    @views Gfix .= x - F
    @views Gfix .= -(Gfix .* Gfix .* pmu)
    @views @inbounds for jfp = 1:n
        FP[:, jfp] .= Gfix[:, 1] .* pseed[jfp:jfp+n-1]
        FP[jfp, jfp] = 1.0 + FP[jfp, jfp]
    end
end

"""
heqf!(F,x,pdata)

The function evaluation as per nsold rules.

The precomputed data is a big deal for this example. In particular,
the output pdata.FFB from plan_fft! goes to the fixed point map
computation. Things get very slow if you do not use plan_fft or plan_fft!
"""
function heqf!(F, x, pdata)
    HeqFix!(F, x, pdata)
    #
    # naked BLAS call to fix the allocation blues
    #
    # Using any variation of F.=x-F really hurts
    #
    axpby!(1.0, x, -1.0, F)
end

"""
function HeqFix!(Gfix,x,pdata)
The fixed point map. Gfix goes directly into the function and
Jacobian evaluations for the nonlinear equations formulation.

The precomputed data is a big deal for this example. In particular,
the output pdata.FFA from plan_fft goes to the fixed point map
computation. Things get very slow if you do not use plan_fft.
"""
function HeqFix!(Gfix, x, pdata)
    n = length(x)
    Gfix .= x
    heq_hankel!(Gfix, pdata)
    Gfix .*= pdata.pmu
    Gfix .= 1.0 ./ (1.0 .- Gfix)
end

"""
heqinit(x0::Array{T,1}, c) where T :< Real

Initialize H-equation precomputed data.
"""
function heqinit(x0::Array{T,1}, c) where {T<:Real}
    n = length(x0)
    cval = ones(1)
    cval[1] = c
    vsize = (n,)
    bsize = (2 * n,)
    ssize = (2 * n - 1,)
    FFA = plan_fft(ones(bsize))
    mu = collect(0.5:1:n-0.5)
    pmu = mu * c
    mu = mu / n
    hseed = zeros(ssize)
    for is = 1:2*n-1
        hseed[is] = 1.0 / is
    end
    hseed = (0.5 / n) * hseed
    pseed = hseed
    gtmp = zeros(vsize)
    rstore = zeros(bsize)
    zstore = zeros(bsize) * (1.0 + im)
    hankel = zeros(bsize) * (1.0 + im)
    FFB = plan_fft!(zstore)
    bigseed = zeros(bsize)
    @views bigseed .= [hseed[n:2*n-1]; 0; hseed[1:n-1]]
    @views hankel .= conj(FFA * bigseed)
    return (
        cval = cval,
        mu = mu,
        hseed = hseed,
        pseed = pseed,
        gtmp = gtmp,
        pmu = pmu,
        rstore = rstore,
        zstore = zstore,
        hankel = hankel,
        FFB = FFB,
    )
end

"""
setc!(pdata, cin)

If you are varying c in a computation, this function
lets you set it.
"""
function setc!(pdata, cin)
    c = pdata.cval[1]
    cfix = cin / c
    pdata.pmu .*= cfix
    pdata.cval[1] = cin
end

"""
heq_hankel!(b,pdata)
Multiply an nxn Hankel matrix with seed in R^(2N-1) by a vector b
FFA is what you get with plan_fft before you start computing
"""
function heq_hankel!(b, pdata)
    reverse!(b)
    heq_toeplitz!(b, pdata)
end

"""
heq_toeplitz!(b,pdata)
Multiply an nxn Toeplitz matrix with seed in R^(2n-1) by a vector b
"""
function heq_toeplitz!(b, pdata)
    n = length(b)
    y = pdata.rstore
    y .*= 0.0
    @views y[1:n] = b
    heq_cprod!(y, pdata)
    @views b .= y[1:n]
end

"""
heq_cprod!(b,pdata)
Circulant matrix-vector product with FFT
compute u = C b

Using in-place FFT
"""

function heq_cprod!(b, pdata)
    xb = pdata.zstore
    xb .*= 0.0
    xb .+= b
    pdata.FFB \ xb
    hankel = pdata.hankel
    xb .*= hankel
    pdata.FFB * xb
    b .= real.(xb)
end

"""
Alternative formulation for CI. Tuned to match the paper

author="P. B. Bosma and W. A. DeRooij",
title="Efficient Methods to Calculate Chandrasekhar's H-functions",
journal="Astron. Astrophys.",
volume=126,
year=1983,
pages=283

"""
function heqbos!(F, x, pdata)
    c = pdata
    n = length(x)
    mu = 0.5:1:n-0.5
    mu = mu / n
    h = 1.0 / n
    cval = sqrt(1.0 - c)
    A = zeros(n, n)
    for j = 1:n
        for i = 1:n
            A[i, j] = mu[j] / (mu[i] + mu[j])
        end
    end
    A = (c / 2) * h * A
    F .= (A * x)
    for ig = 1:n
        F[ig] = 1.0 / (cval + F[ig])
    end
    F .= x - F
end


export
    DiscreteCPT

"""
    function DiscreteCPT(range::Vector{O}, paramdict::Dict{I, Vector{Float64}}) where {I <: Tuple, O}

Constructs an sfunc that represents a Discrete Conditional Probability table.

`DiscreteCPT`s are implemented as a `Table` with a `Cat`.

The `range` parameter defines all the possible outputs of the `DiscreteCPT`.  The `paramdict`
parameter defines the input(s) and the actual CPT.  For example,

```
    range = [1, 2]
    paramdict = Dict((1,1) => [0.3, 0.7], (1,2) => [0.6, 0.4], (2,1) =>[0.4, 0.6],
    (2,2) => [0.7, 0.3], (3,1) => [0.5, 0.5], (3,2) => [0.8, 0.2])
```

can create a `DiscreteCPT` which has two(2) inputs (the length of the key) and, given each input
as defined by the key, selects either `1` or `2` (the range) with the given probability.  So, if
the input is `(2,1)`, `1` is selected with probability `0.4` and `2` is selected with probability
`0.6`.

See also: [`Table`](@ref), [`Cat`](@ref)
"""
function DiscreteCPT(range::Vector{O}, paramdict::Dict{I, Vector{Float64}}) where {I <: Tuple, O}
    NumInputs = length(collect(keys(paramdict))[1])
    sfmaker(probs) = Cat(range, probs)
    return Table(Tuple{}, O, NumInputs, paramdict, sfmaker)
end

#=
# This form is provided for backward compatibility
# The ranges of the parents and child are all integer ranges from 1 to the number of values
function DiscreteCPT(params :: Array{Vector{Float64}, N}) where N
    rangesizes = size(params)
    parranges = [collect(1:r) for r in rangesizes]
    parcombos = cartesian_product(parranges)
    childrange = collect(1:length(params[1]))
    I = NTuple{N, Int}
    paramdict = Dict{I, Vector{Float64}}()
    for (i, combo) in enumerate(parcombos)
        paramdict[tuple(combo...)] = params[i]
    end
    return DiscreteCPT(childrange, paramdict)
end
=#


# Create an admittance matrix. The implementation is a modification of
# MATPOWER's makeYbus. We attach the original MATPOWER's license in makeYbus.m:
#
# MATPOWER
# Copyright (c) 1996-2016, Power Systems Engineering Research Center (PSERC)
# by Ray Zimmerman, PSERC Cornell
#
# Covered by the 3-clause BSD License.
#
# This function returns the following:
#
#  Ybus : nb  x nb admittance
#  Yf_br: nbr x nb from-bus admittance of non-transformer branches
#  Yt_br: nbr x nb to-bus admittance of non-transformer branches
#  Yf_tr: ntr x nb from-bus admittance of transformer branches
#  Yt_tr: ntr x nb to-bus admittance of transformer branches
#
# where nb is the number of buses, nbr is the number of non-transformer
# branches, and ntr is the number of transformer branches.
function makeYbus(data, bus_to_indexes)
    baseMVA = data["baseMVA"]
    bus     = data["bus"]
    branch  = data["branch"]
    F_BUS, T_BUS, BR_R, BR_X, BR_B, RATE_A, RATE_B, RATE_C, TAP, SHIFT, BR_STATUS,
    ANGMIN, ANGMAX, PF, QF, PT, QT, MU_SF, MU_ST, MU_ANGMIN, MU_ANGMAX = IndexSet.idx_branch()
    BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
    LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()

    # constants
    nb = size(bus, 1)          # number of buses
    nl = size(branch, 1)       # number of lines

    # define named indices into bus, branch matrices

    # for each branch, compute the elements of the branch admittance matrix where
    #
    #      | If |   | Yff  Yft |   | Vf |
    #      |    | = |          | * |    |
    #      | It |   | Ytf  Ytt |   | Vt |
    #
    stat = branch[:, BR_STATUS]                     # ones at in-service branches
    Ys = stat ./ (branch[:, BR_R] + 1im * branch[:, BR_X])  # series admittance
    Bc = stat .* branch[:, BR_B]                           # line charging susceptance
    tap = ones(nl)                               # default tap ratio = 1
    i = findall(branch[:, TAP] .!= 0)              # indices of non-zero tap ratios
    tap[i] = branch[i, TAP]                         # assign non-zero tap ratios
    tap = tap .* exp.(1im*pi/180 * branch[:, SHIFT]) # add phase shifters
    Ytt = Ys + 1im*Bc/2
    Yff = Ytt ./ (tap .* conj(tap))
    Yft = - Ys ./ conj(tap)
    Ytf = - Ys ./ tap

    # compute shunt admittance
    # if Psh is the real power consumed by the shunt at V = 1.0 p.u.
    # and Qsh is the reactive power injected by the shunt at V = 1.0 p.u.
    # then Psh - j Qsh = V * conj(Ysh * V) = conj(Ysh) = Gs - j Bs,
    # i.e. Ysh = Psh + j Qsh, so ...
    Ysh = (bus[:, GS] + 1im * bus[:, BS]) / baseMVA[1] # vector of shunt admittances

    # build connection matrices
    f = [bus_to_indexes[e] for e in branch[:, F_BUS]] # list of "from" buses
    t = [bus_to_indexes[e] for e in branch[:, T_BUS]] # list of "to" buses

    Cf = sparse(1:nl, f, ones(nl), nl, nb)       # connection matrix for line & from buses
    Ct = sparse(1:nl, t, ones(nl), nl, nb)       # connection matrix for line & to buses

    # build Yf and Yt such that Yf * V is the vector of complex branch currents injected
    # at each branch's "from" bus, and Yt is the same for the "to" bus end
    i = [1:nl; 1:nl]
    Yf = sparse(i, [f; t], [Yff; Yft], nl, nb)
    Yt = sparse(i, [f; t], [Ytf; Ytt], nl, nb)

    # build Ybus
    Ybus = Cf' * Yf + Ct' * Yt + sparse(1:nb, 1:nb, Ysh, nb, nb)

    return Ybus

end

"""
    bustypeindex(bus, gen)

Returns vectors indexing buses by type: ref, pv, pq.

"""
function bustypeindex(bus, gen, bus_to_indexes)
    # retrieve indeces
    BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
    LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()

    GEN_BUS, PG, QG, QMAX, QMIN, VG, MBASE, GEN_STATUS, PMAX, PMIN, PC1, PC2, QC1MIN,
    QC2MIN, QC2MAX, RAMP_AGC, RAMP_10, RAMP_30, RAMP_Q, APF, MU_PMAX, MU_PMIN, MU_QMAX,
    MU_QMIN = IndexSet.idx_gen()

    # form vector that lists the number of generators per bus.
    # If a PV bus has 0 generators (e.g. due to contingency)
    # then that bus turns to a PQ bus.

    # Design note: this might be computed once and then modified for each contingency.
    gencon = zeros(Int8, size(bus, 1))

    for i in 1:size(gen, 1)
        if gen[i, GEN_STATUS] == 1
            id_bus = bus_to_indexes[gen[i, GEN_BUS]]
            gencon[id_bus] += 1
        end
    end

    bustype = copy(bus[:, BUS_TYPE])

    for i in 1:size(bus, 1)
        if (bustype[i] == PV_BUS_TYPE) && (gencon[i] == 0)
            bustype[i] = PQ_BUS_TYPE
        elseif (bustype[i] == PQ_BUS_TYPE) && (gencon[i] > 0)
            bustype[i] = PV_BUS_TYPE
        end
    end

    # form vectors
    ref = findall(x -> x==REF_BUS_TYPE, bustype)
    pv = findall(x -> x==PV_BUS_TYPE, bustype)
    pq = findall(x -> x==PQ_BUS_TYPE, bustype)

    return ref, pv, pq, bustype
end

"""
    assembleSbus(gen, load, SBASE, nbus)

Assembles vector of constant power injections (generator - load). Since
we do not have voltage-dependent loads, this vector only needs to be
assembled once at the beginning of the power flow routine.

"""
function assembleSbus(gen, bus, baseMVA, bus_to_indexes)

    ngen = size(gen, 1)
    nbus = size(bus, 1)
    sbus = zeros(Complex{Float64}, nbus)
    sload = zeros(Complex{Float64}, nbus)

    # retrieve indeces
    BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
    LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()

    GEN_BUS, PG, QG, QMAX, QMIN, VG, MBASE, GEN_STATUS, PMAX, PMIN, PC1, PC2, QC1MIN,
    QC2MIN, QC2MAX, RAMP_AGC, RAMP_10, RAMP_30, RAMP_Q, APF, MU_PMAX, MU_PMIN, MU_QMAX,
    MU_QMIN = IndexSet.idx_gen()

    for i in 1:ngen
        if gen[i, GEN_STATUS] == 1
            id_bus = bus_to_indexes[gen[i, GEN_BUS]]
            sbus[id_bus] += (gen[i, PG] + 1im*gen[i, QG])/baseMVA
        end
    end

    for i in 1:nbus
        id_bus = bus_to_indexes[bus[i, BUS_I]]
        load = (bus[i, PD] + 1im*bus[i, QD])/baseMVA
        sbus[id_bus] -= load
        sload[id_bus] = load
    end

    return sbus, sload
end


module statmech_tm_solver

# Write your package code here.
__precompile__(true)

using LinearAlgebra
using Zygote
using Optim
using Random
using KrylovKit
using TensorKit
using TensorOperations
using JLD2
using Tullio
using LoopVectorization

using TensorKitAD
using ChainRules
using ChainRulesCore
import ChainRulesCore: rrule, frule
import Base: +, -, *, iterate, length, getindex, similar
import TensorKit: leftorth, rightorth
import LinearAlgebra: normalize, dot, ldiv!

# imps.jl
export  act,
        transf_mat,
        transf_mat_T,
        ovlp,
        nonherm_cost_func,
        lambda_gamma,
        left_canonical,
        iTEBD_truncate,
        ln_fidelity,
        free_energy,
        mps_add,
        left_canonical_QR,
        right_canonical_QR,
        entanglement_spectrum,
        entanglement_entropy,
        expand,
        tangent_map_tn,
        tangent_map,
        calculate_ALR,
        biorth_right_canonical,
        biorth_left_canonical

# utils.jl
export  toarray,
        get_chi,
        get_d,
        arr_to_TensorMap,
        convert_to_tensormap,
        quicksave,
        quickload,
        logsumexp

# mpo_zoo.jl
export  mpo_triangular_AF_ising,
        mpo_triangular_AF_ising_alternative,
        mpo_triangular_AF_ising_adapter,
        mpo_kink_processor,
        mpo_square_ising

# bi_direction.jl
export  bimps,
        A_canonical,
        B_canonical

# cmpo.jl
export  cmpo,
        get_phy,
        get_vir

# cmps.jl
export  cmps,
        get_matrices,
        convert_to_cmps,
        right_canonical,
        K_mat,
        log_ovlp,
        convergence_measure,
        finite_env,
        kinetic,
        particle_density,
        pairing,
        point_interaction,
        optimize_conv_meas,
        compress,
        truncation_check,
        gauge_fixing_proj,
        tangent_map,
        preconditioner,
        precond_grad,
        lieb_liniger_h_tangent_map

# cmpo_zoo.jl
export  cmpo_ising,
        cmpo_xy,
        cmpo_xz,
        energy_quantum_ising,
        cmpo_xxz,
        energy_quantum_xxz,
        cmpo_ising_realtime,
        energy_lieb_liniger,
        energy_pairing_boson

# bi_direction_q.jl
export  qbimps

# TensorKitAD_supp.jl
export  convert_to_array,
        elem_mult

include("imps.jl")
include("utils.jl")
include("mpo_zoo.jl")
include("bi_direction.jl")
include("cmpo.jl")
include("cmps.jl")
include("cmpo_zoo.jl")
include("bi_direction_q.jl")
include("TensorKitAD_supp.jl")

end


include("twoplayergames.jl")

function threatpoint(x::Array{Float64, 2}})

    threat_point = minimum(x, 2)[:]

    return threat_point
end

function Q_operator(tpg::TwoPlayerGame, a1::Int, a2::Int, W::Matrix{Float64}, u::Vector{Float64})

    δ = tpg.δ
    # Flow from action profile and gain from deviation
    _g = g(tpg, a1, a2)
    _h = h(tpg, a1, a2)

    # Find the lines to intersect set with
    w1 = ((1-δ)/δ)*_h[1] + u[1]
    w2 = ((1-δ)/δ)*_h[2] + u[2]

    Q_aWu = intersect_AS(W, w1, w2)

    return Q_aWu
end

function C_operator(tpg::TwoPlayerGame, a1::Int, a2::Int, W::Matrix{Float64}, u::Vector{Float64})

    δ = tpg.δ

    # Flow from action profile and gain from deviation
    _g = g(tpg, a1, a2)
    _h = h(tpg, a1, a2)

    if all(δ*(_g - u) .>= (1-δ)*_h)
        # Only one relevant point if the above holds
        out_ep = [ga]
        out_nep = 1
    else
        # Allocate space for points and give it a size hint
        out_ep = []
        out_nep = 0
        sizehint!(out_ep, 3)

        # Compute elements in Q(a, W, u)
        Q_aWu = Q_operator(tpg, a1, a2, W, u)

        # Check each extreme point of Q(a, W, u) for binding IC
        for ep in Q_aWu
            # Checking binding IC for both agents
            bIC_1 = isapprox((1-δ)*_h[1] - δ*(g[1] - u[1]), 0.0) ? true : false
            bIC_2 = isapprox((1-δ)*_h[2] - δ*(g[2] - u[2]), 0.0) ? true : false

            # Keep points for which either IC is binding
            if bIC_1 || bIC_2
                push!(out_ep, ep)
                out_nep += 1
            end
        end
    end

    return out_ep, out_nep
end

function R_operator(tpg::TwoPlayerGame, W::Matrix{Float64}, u::Vector{Float64})
    # Create space for the extreme points of R(W, u) -- Abreu Sannikov
    # prove that it will be less than 3|A| so we can give a size hint.
    R_Wu = []
    sizehint!(R_Wu, 3*na1*na2)

    for a1=1:na1
        for a2=1:na2
            # Compute the flow payoff
            flow_today = (1-δ)*g(tpg, a1, a2)

            # Get set of possible continuation values
            C_aWu, nC_aWu = C_operator(tpg, a1, a2, W, u)
            for cind=1:nC_aWu
                val_today = flow_today + δ*C_aWu[cind]
                push!(R_Wu, val_today)
            end
        end
    end

    # Take convex hull
    R_Wu = convexhull(R_Wu)

    return R_Wu
end


# Define and test Monoid instance for Int
# ---------------------------------------

TypeClasses.combine(a::Int, b::Int) = a + b

@test reduce_monoid([1,2,3,1]) == 7
@test foldl_monoid([1,2,3,4]) == 10
@test foldr_monoid([1,2,3,100], init=3000) == 3106

# Test default Semigroup instance for String
# ------------------------------------------

@test reduce_monoid(["hi"," ","du"], init="!") == "!hi du"

# Test `neutral` singleton value
# ------------------------------

@test combine(neutral, :anything) == :anything
@test combine("anything", neutral) == "anything"
@test combine(neutral, neutral) == neutral


using Plots
pyplot()
using LaTeXStrings
using StatPlots     # needed for box plot support

n = 500
x = randn(n)  # N(0, 1)
x = exp.(x)   # Map x to lognormal
y = randn(n) + 2.0  # N(2, 1)
z = randn(n) + 4.0  # N(4, 1)
data = vcat(x, y, z)
l = [LaTeXString("\$X\$") LaTeXString("\$Y\$")  LaTeXString("\$Z\$") ]
xlabels = reshape(repmat(l, n), n*3, 1)

boxplot(xlabels, data, label="", ylims=(-2, 14))


#########################
### global constants ####
#########################

#' Constants Used in the bigleaf Package
#'
#' This function defines the following constants:
#'
#' - cp           Specific heat of air for constant pressure (J K-1 kg-1)
#' - Rgas         Universal gas constant (J mol-1 K-1)
#' - Rv           Gas constant of water vapor (J kg-1 K-1) (Stull 1988 p.641)
#' - Rd           Gas constant of dry air (J kg-1 K-1) (Foken p. 245)
#' - Md           Molar mass of dry air (kg mol-1)
#' - Mw           Molar mass of water vapor (kg mol-1)
#' - eps          Ratio of the molecular weight of water vapor to dry air (=Mw/Md)
#' - g            Gravitational acceleration (m s-2)
#' - solar_constant Solar constant (W m-2)
#' - pressure0    Reference atmospheric pressure at sea level (Pa)
#' - Tair0        Reference air temperature (K)
#' - k            von Karman constant
#' - Cmol         Molar mass of carbon (kg mol-1)
#' - Omol         Molar mass of oxygen (kg mol-1)
#' - H2Omol       Molar mass of water (kg mol-1)
#' - sigma        Stefan-Boltzmann constant (W m-2 K-4)
#' - Pr           Prandtl number
#' - Sc_CO2       Schmidt number for CO2
#' - Kelvin       Conversion degree Celsius to Kelvin
#' - DwDc         Ratio of the molecular diffusivities for water vapor and CO2
#' - days2seconds Seconds per day
#' - kPa2Pa       Conversion kilopascal (kPa) to pascal (Pa)
#' - Pa2kPa       Conversion pascal (Pa) to kilopascal (kPa)
#' - umol2mol     Conversion micromole (umol) to mole (mol)
#' - mol2umol     Conversion mole (mol) to micromole (umol)
#' - kg2g         Conversion kilogram (kg) to gram (g)
#' - g2kg         Conversion gram (g) to kilogram (kg)
#' - kJ2J         Conversion kilojoule (kJ) to joule (J)
#' - J2kJ         Conversion joule (J) to kilojoule (kJ)
#' - se_median    Conversion standard error (SE) of the mean to SE of the median
#' - frac2percent Conversion between fraction and percent
#'
#' # Details
 This function is passed as an argument to every function that uses one
#'          or more constants. Individual constants passed to a function can be
#'          easily altered. E_g. the following command will change the value of
#'          the von Karman constant from 0.41 to 0.4:
#'
#'          `bigleaf_constants(k=0.4)`
#'
#'          the value of a constant can be returned by calling:
#'
#'          `bigleaf_constants$*name_of_constant*`
#'
#'          To permanently change the constants contained within this function (which
#'          makes sense for some of them, e.g. for the von Karman constant),
#'          the command `\link[utils]{fixInNamespace`} can be used. E_g.
#'
#'          `fixInNamespace(bigleaf_constants,ns="bigleaf")`
#'
#'          Note that this has to be repeated every time the package is newly installed/loaded.
#'
"""
"""
function bigleaf_constants(
  ## Physical constants
  cp         = 1004.834,        # specific heat of air for constant pressure (J K-1 kg-1)
  Rgas       = 8.31451,         # universal gas constant (J mol-1 K-1)
  Rv         = 461.5,           # gas constant of water vapor (J kg-1 K-1) (Stull 1988 p.641)
  Rd         = 287.0586,        # gas constant of dry air (J kg-1 K-1) (Foken 2008 p. 245)
  Md         = 0.0289645,       # molar mass of dry air (kg mol-1)
  Mw         = 0.0180153,       # molar mass of water vapor (kg mol-1)
  eps        = 0.622,           # ratio of the molecular weight of water vapor to dry air (=Mw/Md)
  g          = 9.81,            # gravitational acceleration (m s-2)
  solar_constant = 1366.1,      # solar constant, i.e. solar radation at earth distance from the sun (W m-2)
  pressure0  = 101325,          # reference atmospheric pressure at sea level (Pa)
  Tair0      = 273.15,          # reference air temperature (K)
  k          = 0.41,            # von Karman constant
  Cmol       = 0.012011,        # molar mass of carbon (kg mol-1)
  Omol       = 0.0159994,       # molar mass of oxygen (kg mol-1)
  H2Omol     = 0.01801528,      # molar mass of water (kg mol-1)
  sigma      = 5.670367e-08,    # Stefan-Boltzmann constant (W m-2 K-4)
  Pr         = 0.71,            # Prandtl number
  Sc_CO2     = 1.07,            # Schmidt number for CO2 (Hicks et al. 1987)

  ## Conversion constants
  Kelvin       = 273.15,         # conversion degree Celsius to Kelvin
  DwDc         = 1.6,            # Ratio of the molecular diffusivities for water vapor and CO2
  days2seconds = 86400,          # seconds per day
  kPa2Pa       = 1000,           # conversion kilopascal (kPa) to pascal (Pa)
  Pa2kPa       = 0.001,          # conversion pascal (Pa) to kilopascal (kPa)
  umol2mol     = 1e-06,          # conversion micromole (umol) to mole (mol)
  mol2umol     = 1e06,           # conversion mole (mol) to micromole (umol)
  kg2g         = 1000,           # conversion kilogram (kg) to gram (g)
  g2kg         = 0.001,          # conversion gram (g) to kilogram (kg)
  kJ2J         = 1000,           # conversion kilojoule (kJ) to joule (J)
  J2kJ         = 0.001,          # conversion joule (J) to kilojoule (kJ)
  se_median    = 1.253,          # conversion standard error (SE) of the mean to SE of the median (http://influentialpoints_com/Training/standard_error_of_median_htm)
  frac2percent = 100             # conversion between fraction and percent
)

  list(
    cp = cp, Rgas = Rgas, Rv = Rv, Rd = Rd, Md = Md, Mw = Mw, eps = eps, g = g,
    solar_constant = solar_constant,
    pressure0 = pressure0, Tair0 = Tair0, k = k, Cmol = Cmol, Omol = Omol,
    H2Omol = H2Omol,
    sigma = sigma, Pr = Pr, Sc_CO2 = Sc_CO2, Kelvin = Kelvin, DwDc = DwDc,
    days2seconds = days2seconds, kPa2Pa = kPa2Pa, Pa2kPa = Pa2kPa, umol2mol = umol2mol,
    mol2umol = mol2umol, kg2g = kg2g, g2kg = g2kg, kJ2J = kJ2J, J2kJ = J2kJ,
    se_median = se_median, frac2percent = frac2percent
  )

end


# generate the typical outcome plots found in
# https://arxiv.org/abs/2201.05967

using DyadicKDE
include("./plot_helpers.jl")

# specify parameters
n_data = 100
n_evals = 100
degeneracies = ["total", "partial", "none"]
kernel_name = "epanechnikov_order_2"
evals = collect(range(-2.0, stop=2.0, length=n_evals))
sdp_solver = "cosmo"
n_resample = 10000
significance_level = 0.05
results = []
ps = Dict(
    "total" => [0.5, 0.0, 0.5],
    "partial" => [0.25, 0.0, 0.75],
    "none" => [0.2, 0.2, 0.6],
)

println("Running experiments")
for degen in degeneracies

    println("Degeneracy: $degen")
    data = make_data(n_data, ps[degen])
    h_ROT = estimate_ROT_bandwidth(data, "epanechnikov_order_2")

    est = DyadicKernelDensityEstimator(
        kernel_name, h_ROT, significance_level,
        n_resample, sdp_solver, evals, data,
        Dict("degen" => degen, "p" => ps[degen]))

    fit(est)
    push!(results, est)

end

println("Making plots")
linewidth = 1.0
y_lim = [-0.01, 0.45]
e = 0.22
pci_marker = PyPlot.matplotlib.path.Path([[-1,0], [-1,e], [-1,-e], [-1,0], [1,0], [1,e], [1,-e], [1,0]])
handle_f = PyPlot.matplotlib.lines.Line2D([0], [0], color="k", lw=linewidth, label="\$f_W(w)\$")
handle_fhat = PyPlot.matplotlib.lines.Line2D([0], [0], color="k", lw=linewidth, linestyle=(0, (1,1)), label="\$\\widehat f_W(w)\$")
handle_ucb = PyPlot.matplotlib.patches.Patch(facecolor="lightgray", edgecolor="lightgray", label="UCB")
handle_pci = PyPlot.matplotlib.lines.Line2D([0], [0], color="black", lw=0, markersize=17, marker=pci_marker, label="PCI")
handles = [handle_f, handle_fhat, handle_ucb, handle_pci]

for est in results

    degen = est.meta["degen"]
    p = est.meta["p"]
    f = generate_f(est.evals, p)

    # plot f
    fig, ax = plt.subplots(figsize=(4,4))
    ax.plot(est.evals, f, color = "black", linewidth=linewidth)

    # plot fhat
    ax.plot(est.evals, est.fhat,
        color = "black", linewidth=linewidth, linestyle=(0, (1,1)))

    # plot ucb
    ax.fill_between(est.evals, est.ucb[1,:], est.ucb[2,:],
        color="lightgray", linewidth=0.0)

    # plot pci
    plot_confidence_intervals(ax, est.evals, est.pci[1,:],
        est.pci[2,:], 10, "black", linewidth, "-")

    # save plot
    PyPlot.xlabel("Evaluation point, \$w\$")
    plt.ylim(y_lim)
    x_min = minimum(est.evals)
    x_max = maximum(est.evals)
    plt.xlim((x_min, x_max))
    plt.yticks(range(0.0, stop=0.4, step=0.1))
    legend(handles=handles, loc="upper left")
    plt.ylabel("Density", labelpad=4.0)
    plt.tight_layout()
    PyPlot.savefig("outcome_plot_$degen.pdf")
    close("all")

end


function coldist(Z,i,j,N)
    s = 0
    @inbounds @simd for k in 1:N
        s += Z[k,i] != Z[k,j]
    end
    s
end

function compute_theta(Z::Matrix{Ti}) where Ti <: Integer
    N,M = size(Z)
    meanfracid = @distributed (+) for i = 1:M-1
        s = 0.0
        for j = i+1:M
            s += 1.0 - coldist(Z,i,j,N) / N
        end
        s
    end
    meanfracid /= 0.5 * M * (M-1)
    theta = min(0.5, 0.38 * 0.32 / meanfracid)
    return theta
end

"""
    compute_weights(Z::Matrix{Ti}, theta::Real) where Ti <: Integer
Compute the normalized counts of the number of sequences at hamming distance ≤ `theta` from any given sequence
in `Z`.
"""
function compute_weights(Z::Matrix{Ti}, theta::Real) where Ti <: Integer
    N, M = size(Z)
    thresh = floor(N * theta)
    W = zeros(M) |> SharedArray
    @sync @distributed for i in 1:M
        @simd for j in 1:M
            dij = coldist(Z,i,j,N)
            val = (dij < thresh)
            W[i] += val
        end
    end

    @sync @distributed for i in 1:M
        W[i] = 1.0 / W[i]
    end
    Meff = sum(W)
    println("M = $M N = $N Meff = $Meff theta = $theta")
    return sdata(W)./Meff , Meff
end

"""
    compute_weights(Z::Matrix{Ti}, theta::Symbol) where Ti<:Integer
Compute the normalized counts of the number of sequences at hamming distance ≤ of a precomputed optimal threshold.
See `Fast and accurate multivariate Gaussian modeling of protein families: Predicting residue contacts and
protein-interaction partners` [https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0092721]
and in particular the Supplementary Information at section `2 Reweighting Scheme` for details.
"""
function compute_weights(Z::Matrix{Ti}, theta::Symbol) where Ti<:Integer
    theta !== :auto && error("only theta=:auto is supported")
    thr = compute_theta(Z)
    return compute_weights(Z,thr)
end


"""
Perform a union between two HyperRectangles.
"""
function union(h1::HyperRectangle{N, T}, h2::HyperRectangle{N, T}) where {T,N}
    m = min.(minimum(h1), minimum(h2))
    mm = max.(maximum(h1), maximum(h2))
    HyperRectangle{N, T}(m, mm-m)
end

"""
Perform a difference between two HyperRectangles.
"""
diff(h1::HyperRectangle, h2::HyperRectangle) = h1

"""
Perform a intersection between two HyperRectangles.
"""
function intersect(h1::HyperRectangle{N, T}, h2::HyperRectangle{N, T}) where {T,N}
    m = max.(minimum(h1), minimum(h2))
    mm = min.(maximum(h1), maximum(h2))
    HyperRectangle{N, T}(m, mm-m)
end

function intersect(a::SimpleRectangle, b::SimpleRectangle)
    min_n = max.(minimum(a), minimum(b))
    max_n = min.(maximum(a), maximum(b))
    w = max_n - min_n
    SimpleRectangle(min_n[1], min_n[2], w[1], w[2])
end


using LinearAlgebra
using Plots
using Distributions
using Random

function VanillaSGD(gradcalc::Function, nn_predict::Function, maxIter, nHidden,nParams, xtrain, ytrain, xvalid, yvalid,a=.0001, B=1)
    print("Vanilla SGD Running...")
    n = size(xtrain,1)
    W= randn(nParams,1)
    flist=[]
    valid=[]
    wbest=zeros(nParams)
    vallow=Inf
    for t in 1:maxIter

    	i = Base.rand(1:n,B)
    	f,g = gradcalc(W, xtrain[i,:], ytrain[i], nHidden)
    	W = W - a*g

        if (mod(t-1,round(maxIter/100)) == 0)
            push!(flist,f)
            ComputeValid(nn_predict, W, xvalid, yvalid, valid)
            if valid[end]<vallow
                wbest=W
                vallow=valid[end]
            end
        end
    end
    display(plot(1:length(valid), valid, xlabel="Number of Iterations", ylabel="Validation Error", legend=false))
    return wbest, vallow
end

function SGDBabysitter(gradcalc::Function, nn_predict::Function, maxIter, nHidden,nParams, xtrain, ytrain, xvalid, yvalid, maxbatch=500)

    n = size(xtrain,1)
    valid=[]
    flist=[]
    alist=[]
    Blist=[]

    a,B,W=BS_initialize(gradcalc, nn_predict, maxIter, nHidden,
                            nParams, xtrain, ytrain, xvalid, yvalid)
    push!(Blist,B)
    push!(alist,a)
    wbest=zeros(nParams)
    vallow=Inf
    triedDec=false
    for t in 1:maxIter

        idxshuffle = shuffle(1:n)
        i = rand(idxshuffle, B)
    	# i = rand(1:n,B)

    	f,g = gradcalc(W, xtrain[i,:], ytrain[i], nHidden)
    	W = W - a*g

        #Every few iterations, compute validation and append f,g,validation
        if (mod(t-1,round(maxIter/100)) == 0)
              push!(flist,f)
              ComputeValid(nn_predict, W, xvalid, yvalid, valid)
              if valid[end]<vallow
                  wbest=W
                  vallow = valid[end]
              end
        end
        if (mod(t-1,round(maxIter/50)) == 0)
            a,B,W,triedDec=BS_Select(validation_array=valid,f_array=flist, gradcalc=gradcalc,xtrain=xtrain,ytrain=ytrain,W=W,nHidden=nHidden, alist=alist, Blist=Blist,maxbatch=maxbatch, wbest=wbest,  maxIter=maxIter,triedDec=triedDec)

          end
    end

    display(plot(1:length(valid), valid, xlabel="Number of Iterations", ylabel="Validation Error", legend=false))

    return wbest, vallow
end

#INITIALIZATION FUNCTION--------------------------------------------------------
function BS_initialize(gradcalc::Function, nn_predict::Function, maxIter, nHidden,nParams, xtrain, ytrain, xvalid, yvalid)
switch=1
n = size(xtrain,1)
B=5
a=rand(Uniform(.095,.09999999))
j=1
W_init = randn(nParams,1)
W = zeros(nParams,1)
while switch==1 && j<10
    W=W_init
    flist=[]
    validlist=[]
    for t in 1:10
        i = rand(1:n,B)
        f,g = gradcalc(W, xtrain[i,:], ytrain[i], nHidden)
        W = W - a*g
        push!(flist,f)
        ComputeValid(nn_predict, W, xvalid, yvalid, validlist)
    end
    if validlist[end]<(4/5)*validlist[1] && flist[end]<(4/5)*flist[1]
        switch=0
    else
        a=rand(Uniform(.5,.75))*a
        print("Initializing... ")
        j+=1
    end
end
return a,B,W
end

#Selection Function-------------------------------------------------------------

ctr = 0
maxIter = 10000
decLimit = round(maxIter/(100*20))
threshold = round(maxIter/(100*10))
function BS_Select(;validation_array,f_array, gradcalc::Function,xtrain,ytrain,W,nHidden, alist, Blist, maxbatch, wbest,  maxIter,triedDec)
    global ctr, decLimit, threshold
    #Get most recent step and batch sizes
    a=alist[end]
    B=Blist[end]
    n,d=size(xtrain)

    #Get Gradient Angle
    i=rand(1:n,B)
    j=rand(1:n,B)
    _,gi=gradcalc(W, xtrain[i,:],ytrain[i],nHidden)
    _,gj=gradcalc(W,xtrain[j,:],ytrain[j], nHidden)
    costheta=dot(gi,gj)/(norm(gi)*norm(gj))

    tol = 1e-5
    if costheta > 1
        costheta = 1 - tol
    elseif costheta < -1
        costheta = -1 + tol
    end
    theta = acos(costheta)

    #Case 1: alpha decrease if validation error does not

    if length(validation_array) >= 4 && ((validation_array[end-3] < validation_array[end]||validation_array[end-2] < validation_array[end]) && (validation_array[end-1]<validation_array[end])) || length(validation_array) >= 4 && ((mean(validation_array[end-3:end-1]) <= 1.05*validation_array[end]) && (mean(validation_array[end-3:end-1]) >= 0.95*validation_array[end]))
        if triedDec==false
            a=(3/4)*a
            triedDec = true
        elseif triedDec==true
            #go back to best w
            W = wbest
            a = (4/5)*a
            # a = (4/3)*a
            triedDec = false
        end

    #---- Alterative to case 1 ----#
    #Don't delete, might use later
    #Tries increasing alpha

    # if length(validation_array) >= 4 && (((validation_array[end-3] < validation_array[end]||validation_array[end-2] < validation_array[end]) && (validation_array[end-1]<validation_array[end])) || ((mean(validation_array[end-3:end-1]) <= 1.05*validation_array[end]) && (mean(validation_array[end-3:end-1]) >= 0.95*validation_array[end])))
    #     if ctr <= decLimit
    #         a = (3/4)*a
    #         print("case 2 part 1")
    #         print(" Ctr is :",ctr, " ")
    #         ctr += 1
    #     elseif ctr > decLimit && ctr <= threshold
    #         a = (10)*a
    #         # ctr += 1
    #         ctr = 0
    #         decLimit += 1
    #         print("case 2 part 2")
    #     elseif ctr > threshold
    #         W = wbest
    #         ctr = 0
    #         print("case 2 part 3")
    #     end

    #Case 2: alpha decreases based on validation error
    # elseif length(validation_array) >= 4 && ((validation_array[end-3] < validation_array[end]||validation_array[end-2] < validation_array[end]) && (validation_array[end-1]<validation_array[end]))
    elseif length(validation_array) >= 4 && validation_array[end-2]>validation_array[end]
        a_adjust = (1 - abs(validation_array[end-2] - validation_array[end])/validation_array[end-2])
        a = a*a_adjust
        print("case 1")
    end

    #case 3 batchsize as a function of theta
    if exp(-.5/(pi-theta))==NaN
        B=maxbatch
    else
        B=convert(Int64,round(maxbatch*(1/(1+round(exp(-.5/(pi-theta+.00001)),digits=5)))))
    end

    # if B > maxbatch
    #     B = maxbatch
    # elseif B < 1
    #     B = 1
    # else
    #     if round(maxbatch*cos(theta/2)) == NaN
    #         B = maxbatch
    #     else
    #         B = convert(Int64, round(maxbatch*cos(theta/2)))
    #     end
    # end

    if a > 1
        a = 1
    end
    push!(alist,a)
    push!(Blist,B)
    print("Step Size: ", a," ","Batch Size: ", B, " ")
    return a,B,W,triedDec
end

#Compute Validation-------------------------------------------------------------
function ComputeValid(predict_func::Function, W, xvalid, yvalid, valid)
    yhat = predict_func(W,[ones(size(xvalid,1)) xvalid],nHidden)
    push!(valid,sum((yhat-yvalid).^2)/length(yvalid))
 end


using ITensors
using OptimKit
using Zygote
using Random: seed!

include("circuit.jl")

function gate(::Val{:Ry}; θ)
  return [
    cos(θ / 2) -sin(θ / 2)
    sin(θ / 2) cos(θ / 2)
  ]
end

function gate(::Val{:CX})
  return [
    1 0 0 0
    0 1 0 0
    0 0 0 1
    0 0 1 0
  ]
end

function Rylayer(N, θ⃗)
  return [("Ry", (n,), (θ=θ⃗[n],)) for n in 1:N]
end

function CXlayer(N)
  return [("CX", (n, n + 1)) for n in 1:2:(N - 1)]
end

# The variational circuit we want to optimize
function variational_circuit(θ⃗)
  N = length(θ⃗)
  return vcat(Rylayer(N, θ⃗), CXlayer(N), Rylayer(N, θ⃗), CXlayer(N))
end

N = 4
θ⃗ = 2π .* rand(N)
gates = variational_circuit(θ⃗)

s = siteinds("Qubit", N)
ψₘₚₛ = MPS(s, "0")
ψ = prod(ψₘₚₛ)
U = buildcircuit(gates, s)
# Create the target state
Uψ = apply(U, ψ)

@show inner_circuit(Uψ, U, ψ);

function loss(θ⃗)
  gates = variational_circuit(θ⃗)
  U = buildcircuit(gates, s)
  return -abs(inner_circuit(Uψ, U, ψ))^2
end

seed!(1234)
θ⃗₀ = randn!(copy(θ⃗))
@show θ⃗
@show loss(θ⃗)
@show loss'(θ⃗)
@show θ⃗₀
@show loss(θ⃗₀)
@show loss'(θ⃗₀)

fg(x) = (loss(x), convert(Vector, loss'(x)))
θ⃗ₒₚₜ, fₒₚₜ, gₒₚₜ, numfg, normgradhistory = optimize(fg, θ⃗₀, GradientDescent())
@show θ⃗ₒₚₜ
@show fₒₚₜ
@show gₒₚₜ
@show numfg
println("normgradhistory = ")
display(normgradhistory)

nothing


using ..MaterialParameters: MaterialParamsInfo
import GeoParams.param_info

export  DiffusionCreep,
        SetDiffusionCreep

const AxialCompression, SimpleShear, Invariant = 1,2,3

#=----Diffusion Creep---
Defines diffusion creep law parameters

n is the power-law exponent
r is the water-fugacity exponent
p is the negative defined grain size exponent (value is set negative in the calculation of σ and ε)
A is a   material specific rheological parameter
E is the activation energy
V is the activation volume
R is the universal gas constant
Apparatus defines the appartus type that shall be recreated (Axial Compression, Simple Shear, Invariant)
=#

@with_kw_noshow mutable struct DiffusionCreep{T,N,U1,U2,U3,U4,U5} <: AbstractCreepLaw{T}
    Name::NTuple{N,Char}        = ""
    n::GeoUnit{T,U1}            = 1.0NoUnits         # power-law exponent
    r::GeoUnit{T,U1}            = 0.0NoUnits         # exponent of water-fugacity
    p::GeoUnit{T,U1}            = -3.0NoUnits         # grain size exponent
    A::GeoUnit{T,U2}            = 1.5MPa^(-n-r)*s^(-1)*m^(-p)    # material specific rheological parameter
    E::GeoUnit{T,U3}            = 500kJ/mol          # activation energy
    V::GeoUnit{T,U4}            = 6e-6m^3/mol        # activation volume
    R::GeoUnit{T,U5}            = 8.314J/mol/K       # universal gas constant
    Apparatus::Int32            = AxialCompression # type of experimental apparatus, either AxialCompression, SimpleShear or Invariant
end

DiffusionCreep(args...) = DiffusionCreep(NTuple{length(args[1]), Char}(collect.(args[1])), convert.(GeoUnit,args[2:end-1])..., args[end])

function param_info(s::DiffusionCreep)
    name = String(collect(s.Name))
    eq = L"\tau_{ij} = 2 \eta  \dot{\varepsilon}_{ij}"
    if name == ""
        return MaterialParamsInfo(Equation=eq)
    end
    inf = DiffusionCreep_info[name][2]
    return MaterialParamsInfo(Equation=eq, Comment=inf.Comment, BibTex_Reference=inf.BibTex_Reference)
end

# Calculation routines for linear viscous rheologies
# All inputs must be non-dimensionalized (or converted to consitent units) GeoUnits
function computeCreepLaw_EpsII(TauII, a::DislocationCreep, c::CreepLawVariables)
    @unpack_val n,r,p,A,E,V,R = a
    @unpack_val P,T,f,d       = c

    FT, FE = CorrectionFactor(a);

    return A*(TauII*FT)^n*f^r*d^p*exp(-(E + P*V)/(R*T))/FE;
end

# EpsII .= A.*(TauII.*FT).^n.*f.^r.*d.^p.*exp.(-(E.+P.*V)./(R.*T))./FE; Once we have a
# All inputs must be non-dimensionalized (or converted to consistent units) GeoUnits
function computeCreepLaw_TauII(EpsII, a::DiffusionCreep, c::CreepLawVariables)
    @unpack_val n,r,A,E,V,R = a
    @unpack_val P,T,f,d     = c

    FT, FE = CorrectionFactor(a);

    return A^(-1/n)*(EpsII*FE)^(1/n)*f^(-r/n)*d^(-p/n)*exp((E + P*V)/(n * R*T))/FT;
end

# Print info
function show(io::IO, g::DiffusionCreep)
    print(io, "DiffusionCreep: Name = $(String(collect(g.Name))), n=$(g.n.val), r=$(g.r.val), p=$(g.p.val), A=$(g.A.val), E=$(g.E.val), V=$(g.V.val), Apparatus=$(g.Apparatus)" )
end

# This computes correction factors to go from experimental data to tensor format
# A nice discussion
function CorrectionFactor(a::DiffusionCreep{_T}) where {_T}

    FT = one(_T)
    FE = one(_T)
    if a.Apparatus == AxialCompression
        FT = sqrt(one(_T)*3)               # relation between differential stress recorded by apparatus and TauII
        FE = one(_T)*2/sqrt(one(_T)*3)     # relation between gamma recorded by apparatus and EpsII
    elseif a.Apparatus == SimpleShear
        FT = one(_T)*2                     # it is assumed that the flow law parameters were derived as a function of differential stress, not the shear stress. Must be modidified if it is not the case
        FE = one(_T)*2
    end
    return FT,FE
end

# Add pre-defined creep laws
"""
    SetDiffusionCreep["Name of Diffusion Creep"]
This is a dictionary with pre-defined creep laws
"""
SetDiffusionCreep(name::String) = DiffusionCreep_info[name][1]

# predefined diffusion creep laws are to be added in the dictionary as it is done for dislocation creep laws (see 'DislocationCreep.jl')!

DiffusionCreep_info = Dict([])


export Map_Quadratic
export eval_map

# ****************************************************************

@doc raw"""
```julia
Map_Quadratic(X_hat_A => X_A, X_hat_B => X_B, X_hat_C => X_C)

Map_Quadratic(X_hat_A, X_hat_B, X_hat_C)
```

Same semantic as [`Map_Affine`](@ref)
"""
struct Map_Quadratic{T} <: Abstract_Map
    _hat_to_A::Pair{T,T}
    _hat_to_B::Pair{T,T}
    _hat_to_C::Pair{T,T}
end

function Map_Quadratic(x_hat_A::T,x_hat_B::T,x_hat_C::T) where {T}
    Map_Quadratic(x_hat_A=>one(T),
                  x_hat_B=>one(T),
                  x_hat_C=>one(T))
end

parameter_size(m::Map_Quadratic) = 3

function eval_map(m::Map_Quadratic{T},
                  X_hat::AbstractVector{T_HAT},
                  θ::AbstractVector{T_θ}) where {T_HAT,T,T_θ}

    @assert length(θ) == parameter_size(m)

    X_T = promote_type(T_HAT,T,T_θ)

    θ_A = θ[1]
    θ_B = θ[2]
    θ_C = θ[3]

    X_hat_A, X_A = m._hat_to_A
    X_hat_B, X_B = m._hat_to_B
    X_hat_C, X_C = m._hat_to_C

    basis_A(x) = (X_hat_B - x)*(X_hat_C - x)/((X_hat_B - X_hat_A)*(X_hat_C - X_hat_A))
    basis_B(x) = (x - X_hat_A)*(X_hat_C - x)/((X_hat_B - X_hat_A)*(X_hat_C - X_hat_B))
    basis_C(x) = (x - X_hat_A)*(x - X_hat_B)/((X_hat_C - X_hat_A)*(X_hat_C - X_hat_B))

    transform(x) = basis_A(x)*X_A*θ_A + basis_B(x)*X_B*θ_B + basis_C(x)*X_C*θ_C

    X=similar(X_hat,X_T)
    @. X = transform(X_hat)

    X
end


using Test, Random
import LuxurySparse: PermMatrix, pmrand
using SparseArrays: sprand
using LinearAlgebra

Random.seed!(2)
p1 = PermMatrix([1, 4, 2, 3], [0.1, 0.2, 0.4im, 0.5])
p2 = PermMatrix([2, 1, 4, 3], [0.1, 0.2, 0.4, 0.5])
#p3 = PermMatrix([4,1,2,3],[0.5, 0.4im, 0.3, 0.2])
p3 = pmrand(4)
sp = sprand(4, 4, 0.3)
v = [0.5, 0.3im, 0.2, 1.0]

@testset "basic" begin
    @test p1 == copy(p1)
    @test eltype(p1) == ComplexF64
    @test eltype(p2) == Float64
    @test eltype(p3) == Float64
    @test size(p1) == (4, 4)
    @test size(p3) == (4, 4)
    @test size(p1, 1) == size(p1, 2) == 4
    @test Matrix(p1) == [0.1 0 0 0; 0 0 0 0.2; 0 0.4im 0 0; 0 0 0.5 0]
    p0 = similar(p1)
    @test p0.perm == p1.perm
    @test p0.perm !== p1.perm
    @test p0.vals !== p1.vals
end

@testset "linalg" begin
    @test inv(p1) * p1 ≈ Matrix(I, 4, 4)
    @test p1 * transpose(p1) == diagm(0 => p1.vals .^ 2)
    #@test p1*adjoint(p1) == diagm(0=>abs.(p1.vals).^2)
    #@test all(isapprox.(adjoint(p3), transpose(conj(Matrix(p3)))))
    @test p1 * p1' == diagm(0 => abs.(p1.vals) .^ 2)
    @test all(isapprox.(p3', transpose(conj(Matrix(p3)))))
end

@testset "mul" begin
    @test p3 * p2 == SparseMatrixCSC(p3) * p2 == Matrix(p3) * p2

    # Multiply vector
    @test p3 * v == Matrix(p3) * v
    @test v' * p3 == v' * Matrix(p3)
    @test p3 * collect(1:4) == p3.perm .* p3.vals

    # Diagonal matrices
    Dv = Diagonal(v)
    @test p3 * Dv == Matrix(p3) * Dv
    @test Dv * p3 == Dv * Matrix(p3)
end

@testset "elementary" begin
    @test all(isapprox.(conj(p1), conj(Matrix(p1))))
    @test all(isapprox.(real(p1), real(Matrix(p1))))
    @test all(isapprox.(imag(p1), imag(Matrix(p1))))
end

@testset "basicmath" begin
    @test p1 * 2 == Matrix(p1) * 2
    @test p1 / 2 == Matrix(p1) / 2
end

@testset "memorysafe" begin
    @test p1 == PermMatrix([1, 4, 2, 3], [0.1, 0.2, 0.4im, 0.5])
    @test p2 == PermMatrix([2, 1, 4, 3], [0.1, 0.2, 0.4, 0.5])
    @test v == [0.5, 0.3im, 0.2, 1.0]
end

@testset "sparse" begin
    Random.seed!(2)
    pm = pmrand(10)
    out = zeros(10, 10)
    @test LuxurySparse.nnz(pm) == 10
    @test LuxurySparse.nonzeros(pm) == pm.vals
    @test LuxurySparse.dropzeros!(pm) == pm
end

@testset "identity sparse" begin
    p1 = Diagonal(randn(10))
    @test LuxurySparse.nnz(p1) == 10
    @test LuxurySparse.nonzeros(p1) == p1.diag
    @test LuxurySparse.dropzeros!(p1) == p1
end

@testset "setindex" begin
    pm = PermMatrix([3, 2, 4, 1], [0.0, 0.0, 0.0, 0.0])
    pm[3, 4] = 1.0
    @test_throws BoundsError pm[3, 1] = 1.0
    @test pm[3, 4] == 1.0
end


using HJBSolver, FactCheck

type TestProblem2D{T<:Real}
    gamma::Vector{T}
    c::Vector{T}
    model::HJBTwoDim{T}
    truevaluefun::Function
    truecontrolfun::Function
end

function createmodel(;T::Float64=1.)
    γ = [5e-2, 0.1]
    c = [-1., -0.9]

    xmin = zeros(2); xmax = [0.5,0.5]

    truesol(t,x) = (t-T)*(x[1]+x[2])^2+vecdot(c,x)
    truepol(t,x) = (1+c+2*(t-T)*(x[1]+x[2]-γ.^2))./(2-2*(t-T)*γ.^2)
    hatf(t,x) = (x[1]+x[2])^2+sum((1-c-2*(t-T)*(x[1]+x[2])).^2./(4-4*(t-T)*γ.^2))

    amin = truepol(0., xmax) - 1e-1
    amax = truepol(0., xmin) + 1e-1

    b(t,x,a) = a-1.
    σ(t,x,a) = (γ.*b(t,x,a)).^2
    f(t,x,a) = -vecdot(a,b(t,x,a)) - hatf(t,x)
    g(x) = truesol(T,x)

    model = HJBTwoDim(b, σ, f, g, T, amin, amax, xmin, xmax, truesol)
    return TestProblem2D(γ, c, model, truesol, truepol)
end

function calculateerror(prob::TestProblem2D, K::Vector{Int}, v,pol)
    model = prob.model

    x1 = linspace(model.xmin[1], model.xmax[1], K[1])
    x2 = linspace(model.xmin[2], model.xmax[2], K[2])
    x = (collect(x1), collect(x2))

    i,j = ceil(Int, K/2)
    idxi = K[2]*(i-1) + j
    xij = [x1[i], x2[j]]
    w = prob.truevaluefun(0., xij)
    α = prob.truecontrolfun(0., xij)

    return abs(w-v[idxi,1])/abs(w), abs(α[1]-pol[1][idxi,1])/abs(α[1]),
    abs(α[2]-pol[2][idxi,1])/abs(α[2])
end

facts("2D, Policy iteration") do
    K = [26, 26]; N = 10
    Δt = 5e-2
    prob = createmodel(T=Δt*N)
    model = prob.model

    @time v, pol = solve(model, K, N)
    errv, erra1, erra2 = calculateerror(prob, K, v, pol)

    @fact errv --> roughly(0., 5e-3)
    @fact erra1 --> roughly(0., 5e-2)
    @fact erra2 --> roughly(0., 5e-2)
end

facts("2D, Policy timestepping") do
    K = [21, 21]; N = 5; M = (25,25)
    Δt = 1e-1
    prob = createmodel(T=Δt*N)
    model = prob.model

    avals1 = linspace(model.amin[1], model.amax[1], M[1])
    avals2 = linspace(model.amin[2], model.amax[2], M[2])

    @time v, pol = solve(model, K, N, (avals1,avals2))

    errv, erra1, erra2 = calculateerror(prob, K, v, pol)
    @fact errv --> roughly(0., 1e-2)
    @fact erra1 --> roughly(0., 0.3)
    @fact erra2 --> roughly(0., 0.3)
end

#==

function calctruesols(t, x, prob, flat=false)
    K = [length(xi) for xi in x]
    v = zeros(reverse(K)...)
    pol1 = zeros(v)
    pol2 = zeros(v)

    for i = 1:K[1], j = 1:K[2]
        xij = [x[1][i],x[2][j]]
        v[j,i] = prob.truevaluefun(t, xij)
        pol1[j,i], pol2[j,i] = prob.truecontrolfun(t, xij)
    end
    if flat == true
        v = v[:]
        pol1 = pol1[:]
        pol2 = pol2[:]
    end
    return v, (pol1, pol2)
end

K = [51, 51]; N = 50
Δt = 2e-2
prob = createmodel(T=Δt*N)
model = prob.model
x1 = linspace(model.xmin[1], model.xmax[1], K[1])
x2 = linspace(model.xmin[2], model.xmax[2], K[2])
x = (collect(x1), collect(x2))
v, pol = solve(prob.model, K, N)

tv = zeros(v)
tpol1 = zeros(pol[1])
tpol2 = zeros(pol[2])
for i = 1:N
tv[:,i], tmppol = calctruesols((i-1)*Δt, x, prob, true)
tpol1[:,i] = tmppol[1]
tpol2[:,i] = tmppol[2]
end
tv[:,end], tmppol = calctruesols(model.T,x,prob, true)
==#


using ModelingToolkit
using NonlinearSolve
using Test

@variables x y z u
@parameters σ ρ β

eqs = [0 ~ σ * (y - x),
    0 ~ x * (ρ - z) - y,
    0 ~ x * y - β * z]

par = [
    σ => 1,
    ρ => 0.1 + σ,
    β => ρ * 1.1,
]
u0 = [
    x => u,
    y => σ, # default u0 from default p
    z => u - 0.1,
]
ns = NonlinearSystem(eqs, [x, y, z], [σ, ρ, β], name = :ns, defaults = [par; u0])
ns.y = u * 1.1
resolved = ModelingToolkit.varmap_to_vars(Dict(), parameters(ns),
                                          defaults = ModelingToolkit.defaults(ns))
@test resolved == [1, 0.1 + 1, (0.1 + 1) * 1.1]

prob = NonlinearProblem(ns, [u => 1.0], Pair[])
@test prob.u0 == [1.0, 1.1, 0.9]
@show sol = solve(prob, NewtonRaphson())

@variables a
@parameters b
top = NonlinearSystem([0 ~ -a + ns.x + b], [a], [b], systems = [ns], name = :top)
top.b = ns.σ * 0.5
top.ns.x = u * 0.5

res = ModelingToolkit.varmap_to_vars(Dict(), parameters(top),
                                     defaults = ModelingToolkit.defaults(top))
@test res == [0.5, 1, 0.1 + 1, (0.1 + 1) * 1.1]

prob = NonlinearProblem(top, [states(ns, u) => 1.0, a => 1.0], [])
@test prob.u0 == [1.0, 0.5, 1.1, 0.9]
@show sol = solve(prob, NewtonRaphson())

# test NullParameters+defaults
prob = NonlinearProblem(top, [states(ns, u) => 1.0, a => 1.0])
@test prob.u0 == [1.0, 0.5, 1.1, 0.9]
@show sol = solve(prob, NewtonRaphson())

# test initial conditions and parameters at the problem level
pars = @parameters(begin
                       x0
                       t
                   end)
vars = @variables(begin x(t) end)
der = Differential(t)
eqs = [der(x) ~ x]
@named sys = ODESystem(eqs, t, vars, [x0])
pars = [
    x0 => 10.0,
]
initialValues = [
    x => x0,
]
tspan = (0.0, 1.0)
problem = ODEProblem(sys, initialValues, tspan, pars)
@test problem.u0 isa Vector{Float64}


"""
    μϵF_induced_by_ϵjs(pts, ϵF, ϵjs)

## Returns
The measures ``\\{\\mu(b_f)\\}`` of the bins ``b_f \\in B_{\\epsilon_{F}}``
(``B_{\\epsilon_{F}}`` is a partition given by the binning scheme `ϵF`)
induced by the measures of the bins ``b_j \\in B_{\\epsilon_{j}}``
(``B_{\\epsilon_{j}}`` is a partition given by the binning scheme `ϵⱼ`)
for several different ``B_{\\epsilon_{j}}`` constructed from the
binning schemes `ϵⱼs`

This is essentially averaging the invariant measures obtained at different
resolutions, expressing the average at the final resolution given by
the binning scheme `ϵF`.

## Explanation

Consider a cloud of `points` and two separate binning schemes, `ϵF` and `ϵⱼ`.
The binning schemes dictate how to create a rectangular partition covering the
points.

Create the partition ``B_{\\epsilon_{F}}`` from `ϵF` and a separate partition
``B_{\\epsilon_{j}}`` from `ϵⱼ`. Then, compute an invariant measure over each
partition separately.

The boxes in ``B_{\\epsilon_{F}}`` are have sizes different from
the boxes in ``B_{\\epsilon_{j}}``. Thus, the invariant measure associated
with each of the partitions convey something about the dynamics at different
scales. We may want to know how the dynamics exposed by
``B_{\\epsilon_{j}}`` is expressed in terms of another coarse-graining
``B_{\\epsilon_{F}}``.

This function computes the measure of the bins ``b_f \\in B_{\\epsilon_{F}}``
induced by the bins ``b_j \\in B_{\\epsilon_{j}}``. That is, for each ``b_f``,
find the set of bins
``\\{b_k \\in B_{\\epsilon_{j}} : b_k \\cap b_f \\neq \\emptyset \\}``.

Extending this to multiple partitions (many different resolutions),
we can decide on a final rectangular grid (resolution) and compute the measure
over that grid as an average over multiple separate partitions (finer, coarser, or both,
depending on the specific application).

## Arguments
- **`points`**: The points representing observations from the system from
    which to compute an invariant measure.
- **`ϵF`**: Edge lengths along each coordinate axis in the partition ``B_{\\epsilon_{F}}``.
- **`ϵⱼ`**: Edge lengths along each coordinate axis in the partition ``B_{\\epsilon_{j}}``.
"""
function μϵF_induced_by_ϵjs(points, ϵF, ϵjs)
    n_partitions = length(ϵjs)

    μ_across_binsizes = Vector{Vector{Float64}}(undef, n_partitions)

    for i=1:n_partitions
        μ_across_binsizes[i] = μϵF_induced_by_ϵj(points, ϵF, ϵjs[i])
    end

    # Gather in a matrix and transpose. Now the i-th column
    # and a-th row give the measure indiced for the i-th
    # state of the target partition induced by the a-th
    # bin size ϵₐ.
    μ = transpose(hcat(μ_across_binsizes...,))
end

export μϵF_induced_by_ϵjs

"""
    average_measure(points, ϵF, ϵjs)

Express the invariant measure of a rectangular box covering of `points`, at
the resolution given by the binning scheme `ϵⱼ` as the invariant measure at
a separate resolution given by the binning scheme `ϵF`.

See the documentation for [`μϵF_induced_by_ϵjs`](@ref) for more details.

## Arguments
- **`points`**: The points representing observations from the system from
    which to compute an invariant measure.
- **`ϵF`**: Edge lengths along each coordinate axis in the partition ``B_{\\epsilon_{F}}``.
- **`ϵjs`**: Edge lengths along each coordinate axis in the (multiple) partitions ``B_{\\epsilon_{j}}``.
"""
function average_measure(points, ϵF, ϵjs)
    n_partitions = length(ϵjs)

    measures = μϵF_induced_by_ϵjs(points, ϵF, ϵjs)

    # Compute the average measure for the final partition
    n_states_μF = size(measures, 2)
    avg_measure = zeros(Float64, n_states_μF)
    for i in 1:n_states_μF
        avg_measure[i] = (1/n_partitions) * sum(measures[:, i])
    end

    return avg_measure
end

export average_measure


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

@userplot ExcessPlot

@recipe function f(p::ExcessPlot)
  # get user input
  xs = p.args[1]

  ks = 2:length(xs)
  ξs = meanexcess(xs, ks)

  seriestype --> :path
  xguide --> "number of maxima"
  yguide --> "extreme value index"

  ks, ξs
end


export readSrcRcvLocationFile,generateSrcRcvProjOperators,readDataFileToDataMat,writeSrcRcvLocFile,writeDataFile,limitDataToOffset

function cs2loc(cs_loc::Int64, n::Array{Int64,1})
loc = zeros(Int64,length(n))
if length(n)==2
	@inbounds loc[1] = mod(cs_loc-1,n[1])+1;
	@inbounds loc[2] = div(cs_loc-1,n[1])+1;
elseif length(n)==3
	@inbounds loc[1] = mod(cs_loc-1,n[1]) + 1;
	@inbounds loc[2] = div(mod(cs_loc-1,n[1]*n[2]),n[1]) + 1;
	@inbounds loc[3] = div(cs_loc-1,n[1]*n[2])+1;
end
return loc;
end

function loc2cs(loc1::Int64,loc2::Int64,n::Array{Int64,1})
@inbounds cs = loc1 + (loc2-1)*n[1];
return cs;
end

function loc2cs3D(loc1::Int64,loc2::Int64,loc3::Int64,n::Array{Int64,1})
@inbounds cs = loc1 + (loc2-1)*n[1] + (loc3-1)*n[1]*n[2];
return cs;
end

function readSrcRcvLocationFile(filename,M::RegularMesh)
# read files for sources and receivers location in UTM and translate them into indices locations.
# file contains a table: idx xUTM yUTM zUTM
A = readdlm(filename);
domainBoundaryUTM = M.domain;
n = M.n;
# domainBoundaryUTM: [x_start, x_end, y_start, y_end, z_start, z_end];
# n: number of cells. (n-1: number of cells)

for k=1:M.dim
	A[:,k+1] = A[:,k+1] .- domainBoundaryUTM[2*k-1];
	A[:,k+1] = A[:,k+1] ./ (domainBoundaryUTM[2*k] - domainBoundaryUTM[2*k-1]);
	A[:,k+1] = round.(A[:,k+1] .* (n[k])) .+ 1.0;
end
A = round.(Int,A);
return A;
end

function generateSrcRcvProjOperators(n_nodes::Vector,SrcRcvNodeLoc::Array{Int64})
# gets a sources/receivers location table (in indices) and generates operators.
nSrcRcv = size(SrcRcvNodeLoc,1);
I = zeros(Int64,nSrcRcv);
J = zeros(Int64,nSrcRcv);
V = zeros(Float64,nSrcRcv);
for k=1:nSrcRcv
	J[k] = k;
	if length(n_nodes)==2
		I[k] = loc2cs(SrcRcvNodeLoc[k,2],SrcRcvNodeLoc[k,3],n_nodes);
	else
		I[k] = loc2cs3D(SrcRcvNodeLoc[k,2],SrcRcvNodeLoc[k,3],SrcRcvNodeLoc[k,4],n_nodes)
	end
	V[k] = 1.0;
end
P = sparse(I,J,V,prod(n_nodes),nSrcRcv);
end

function readDataFileToDataMat(filename,SrcNodeLoc::Array{Int64},RcvNodeLoc::Array{Int64})
# reads a data file into a data matrix D for the code to work with
Data = readdlm(filename);
# Data is a table of #src #rcv tau(src,rcv) Wd(src,rcv) and the rest is zero.
# CellLocMap is a table of #src i j k
numSrc = size(SrcNodeLoc,1);
numRcv = size(RcvNodeLoc,1);

if size(Data,2) == 4
	D = zeros(Float64,numRcv,numSrc);
	Wd = zeros(Float64,numRcv,numSrc);
	SrcIdxs = SrcNodeLoc[:,1];
	RcvIdxs = RcvNodeLoc[:,1];
	for k=1:size(Data,1)
		srcIdx = convert(Int64,round(Data[k,1]));
		rcvIdx = convert(Int64,round(Data[k,2]));
		isrc = binarySearch(SrcIdxs,srcIdx);
		ircv = binarySearch(RcvIdxs,rcvIdx);
		D[ircv,isrc] = Data[k,3] ;#+ 1im*Data[k,4];
		Wd[ircv,isrc] = Data[k,4] ;#+ 1im*Data[k,5];
	end
elseif size(Data,2) == 6
	D = zeros(ComplexF64,numRcv,numSrc);
	Wd = zeros(ComplexF64,numRcv,numSrc);
	SrcIdxs = SrcNodeLoc[:,1];
	RcvIdxs = RcvNodeLoc[:,1];
	for k=1:size(Data,1)
		srcIdx = convert(Int64,round(Data[k,1]));
		rcvIdx = convert(Int64,round(Data[k,2]));
		isrc = binarySearch(SrcIdxs,srcIdx);
		ircv = binarySearch(RcvIdxs,rcvIdx);
		D[ircv,isrc] = Data[k,3] + 1im*Data[k,5];
		Wd[ircv,isrc] = Data[k,4] + 1im*Data[k,6];
	end

end
return D,Wd;
end

function binarySearch(arr, value)
    low = 1
    high = length(arr)
    while low <= high
        mid = round(Int,(low+high)/2)
        if arr[mid] > value
          high = mid-1
        elseif arr[mid] < value
          low = mid+1
        else
          return mid
        end
    end
    return -1
end

function writeSrcRcvLocFile(filename,Msh::RegularMesh,pad::Int64,jump::Int64)
# writes an equally distanced sources/receivers file.
n_nodes = Msh.n.+1;
Q = zeros(Bool, tuple(n_nodes...));

if Msh.dim==2
	Q[pad+1:jump:end-pad,1] .= true;
else
	Q[pad+1:jump:end-pad,pad+1:jump:end-pad,1] .= true;
end
numSrcRcv = round(Int,sum(Q));
Q = findall(Q[:]);
SrcRcvTable = zeros(numSrcRcv,1+Msh.dim);
SrcRcvIDs = randperm(numSrcRcv);
SrcRcvIDs = sort(SrcRcvIDs[1:numSrcRcv]);
SrcRcvTable[:,1] .= SrcRcvIDs;
startLocations = Msh.domain[1:2:end];
for k=1:numSrcRcv
	SrcRcvTable[k,2:end] .= (cs2loc(Q[k],n_nodes).-1).*Msh.h .+ startLocations;
end
writedlm(filename,SrcRcvTable);
end

function writeDataFile(filename,D::Union{Array{Float64,2},Array{ComplexF64,2}},Wd::Union{Array{Float64,2},Array{ComplexF64,2}},SrcNodeLoc::Array{Int64},RcvNodeLoc::Array{Int64})
srcIDs = SrcNodeLoc[:,1];
rcvIDs = RcvNodeLoc[:,1];

if eltype(D)==Float64
	println("Writing real data file");
	Data = zeros(Float32,count(!iszero, Wd),4);
	k = 1;
	for j = 1:size(D,2)
		for i = 1:size(D,1)
			if Wd[i,j] != 0.0
				Data[k,1] = srcIDs[j];
				Data[k,2] = rcvIDs[i];
				Data[k,3] = D[i,j];
				Data[k,4] = Wd[i,j];
				k+=1;
			end
		end
	end
else
	println("Writing ``complex'' data file");
	Data = zeros(Float32,count(!iszero, Wd),6);
	k = 1;
	for j = 1:size(D,2)
		for i = 1:size(D,1)
			if Wd[i,j] != 0.0
				Data[k,1] = srcIDs[j];
				Data[k,2] = rcvIDs[i];
				Data[k,3] = real(D[i,j]);
				Data[k,4] = real(Wd[i,j]);
				Data[k,5] = imag(D[i,j]);
				Data[k,6] = imag(Wd[i,j]);
				k+=1;
			end
		end
	end
end
writedlm(filename,Data);
return;
end

function limitDataToOffset(Wd::Union{Array{Float64,2},Array{ComplexF64,2}},SrcNodeLoc::Array{Int64},RcvNodeLoc::Array{Int64},offsetInNodes::Int64)
for j = 1:size(Wd,2)
	srcLoc = SrcNodeLoc[j,2:end];
	for i = 1:size(Wd,1)
		rcvLoc = RcvNodeLoc[i,2:end];
		dist = norm(srcLoc .- rcvLoc);
		if dist > offsetInNodes || dist < 2
			Wd[i,j] = 0.0;
		end
	end
end
return Wd;
end


# Implementation of an elementary cellular automata
# according to https://mathworld.wolfram.com/ElementaryCellularAutomaton.html
# uses random initialization or single cell initialization and uses wolfram codes to specify the rule
# synchronous update of the 1D lattice

using Agents, Random
using CairoMakie
using InteractiveDynamics
using CSV

"""
The automaton living in a 1D space
"""
mutable struct Cell <: AbstractAgent
    id::Int
    pos::Dims{1}
    status::Int # either 0 or 1, where 1 is 'alive'
end

"""
Returns an array with the rule for the next status configurations, 1-indexed.

Thus, the first position corresponds to the rule for the cell neighborhood status (0,0,0)
and the last to (1,1,1)
"""
function rule_from_code(wolfram_code)
    return digits(wolfram_code, base=2, pad=8)
end

"""
Takes the status of a neighborhood and returns the corresponding
index in the model rule (0,0,0) -> 1 (0,1,0) -> 3
"""
function configuration_index(cell_statuses)
    # takes the tuples and forms a string (0,1,1) -> "011"
    binary_code = string(cell_statuses...)
    # turns the string binary code into the integer
    index = parse(Int, binary_code, base=2) + 1
    return index
end

"""
Given a cell checks its neighbors and
decides the next status of the cell based on the model rule
"""
function next_status(cell, model)
    neighbors = collect(nearby_agents(cell, model))
    cell_statuses = (c.status for c in [neighbors[1], cell, neighbors[2]])
    index = configuration_index(cell_statuses)
    return model.rule[index]
end

"""
Initializes the ABM

# Arguments
- `n_cells::Int`: total of cell automaton
- `wolfram_code::Int`: wolfram code for the rule
- `random_start::Bool`: random start or single cell at center
- `seed::Int`: random seed
"""
function build_model(; n_cells = 101, wolfram_code=30, random_start=true, seed = 30)
    space = GridSpace((n_cells,); metric=:chebyshev)

    properties = Dict(:rule => rule_from_code(wolfram_code),)
    model = ABM(
        Cell,
        space; properties,
        rng = MersenneTwister(seed))

	if random_start
		for x in 1:n_cells
			cell = Cell(nextid(model), (x,), rand([0,1]))
			add_agent_pos!(cell, model)
		end
	else
		for x in 1:n_cells
			initial_state = x == div(n_cells, 2) + 1 ? 1 : 0
			cell = Cell(nextid(model), (x,), initial_state)
        	add_agent_pos!(cell, model)
		end
	end

    return model
end

"""
Dummy update of the cells
"""
function cell_step!(_, _)

end

"""
Performs a synchronous update of all cells
"""
function ca_step!(model)
    new_statuses = fill(0, nagents(model))
    for agent in allagents(model)
        new_statuses[agent.id] = next_status(agent, model)
    end

    for id in allids(model)
        model[id].status = new_statuses[id]
    end
end

# Initialize model
model = build_model(n_cells=101, wolfram_code=30, random_start=false)

# Runs the model and collects data
data, _ = run!(model, cell_step!, ca_step!, 100; adata=[:status]);

# The data contains the step, id/position and status of the cell (1/0)
data

CSV.write("data/ca_data.csv", data);

# Lets plot the time evolution in the y axis
function plot_automata(data)
    fig, ax = heatmap(data.id, data.step, data.status, colormap=:Blues_3)
    ax.yreversed = true
    return fig
end

plot_automata(data)


using Test

using OpticsPolynomials

include("jacobi.jl")
include("cheby.jl")
include("legendre.jl")
include("zernike.jl")


# simple discrete hmm with testable output
A = eye(2)
B_matrix = [0.0  1.0  0.0 ;
            1.0  0.0  0.0 ]

model = HMM(A,B_matrix)
model.p = vec([1 0])

# state sequence s, observation sequence o
n_obs = 10
s,o = rand(model,n_obs)
@test(length(s) == length(o) == n_obs)
@test(all(s .== 1))
@test(all(o .== 2))

# check forward-backward algorithm on perfect sequence
alpha,beta,log_p_obs = forward_backward(model,o; scaling=false)
@test(log_p_obs == 0.0)
@test(all(alpha[:,1] .== 1))
@test(all(alpha[:,2] .== 0))
@test(all(beta[end,:] .== 1))
@test(all(beta[:,1] .== 1))

# test forward-backward with scaling
alpha,beta,log_p_obs = forward_backward(model,o)
@test(log_p_obs == 0.0)
@test(all(alpha[:,1] .== 1))
@test(all(alpha[:,2] .== 0))
@test(all(beta[end,:] .== 1))
@test(all(beta[:,1] .== 1))

# check viterbi algorithm on perfect sequence
@test(all(viterbi(model,o) .== s))

# check forward-backward algorithm on impossible sequence
alpha,beta,log_p_obs = forward_backward(model,o+1; scaling=false)
@test(log_p_obs == -Inf)
@test(all(alpha .== 0))
@test(all(beta[1:end-1,:] .== 0))

######
# Compare output to Michael Hamilton's "dishonest casino" example
# see Hamilton's HMM implementation in python:
# http://www.cs.colostate.edu/~hamiltom/code.html#python-hidden-markov-model

A = [ 0.99 0.01 ;
      0.01 0.99 ]
B = [ 1/6   1/6   1/6   1/6   1/6   1/6 ;
      1/10  1/10  1/10  1/10  1/10  1/2 ]
p = [0.5, 0.5]
model = HMM(A,B,p)

@test(model.n == 2)
for i = 1:model.n
  @test(typeof(model.B[i]) == Categorical)
  @test(model.B[i].K == 6)
end
@test(length(model.p) == model.n)

o = [1, 2, 1, 6, 6] # observation sequence

# check forward/backward (without scaling)
true_alpha = [ 8.3333e-02   1.3833e-02   2.2909e-03   3.7885e-04   6.295e-05  ;
               5.0000e-02   5.0333e-03   5.1213e-04   2.6496e-04   1.3305e-04 ]

true_beta = [ 8.9630e-04   5.2841e-03   3.0533e-02   1.7000e-01   1.0000e+00 ;
              2.4262e-03   2.4418e-02   2.4613e-01   4.9667e-01   1.0000e+00 ]

alpha,beta,log_p_obs = forward_backward(model, o; scaling=false)

@test(all(round(alpha,6) .== round(true_alpha',6)))
@test(all(round(beta,5) .== round(true_beta',5)))

# check forward/backward (with scaling)
alpha_scaled = [ 0.625   0.7332  0.8173  0.5884  0.3212 ;
                 0.375   0.2668  0.1827  0.4116  0.6788 ]

beta_scaled = [  4.5729   3.5946   2.9391   2.4312  3.2847 ;
                 12.3785  16.6108  23.6922  7.1028  3.2847 ]

alpha,beta,log_p_obs = forward_backward(model,o; scaling=true)

@test(all(round(alpha,4) .== round(alpha_scaled',4)))
@test(all(round(beta,4) .== round(beta_scaled',4)))

# check viterbi
true_vit = [2,2,2,2,2]
@test(all(viterbi(model,o) .== true_vit))
@test(all(viterbi(model,[1,2,3,4,5]) .== ones(Int,5)))

# another check for viterbi
model.A = [ 0.7 0.3 ;
          0.3 0.7 ]
o = [ 1, 2, 1, 6, 6, 3, 3, 2, 1 ]
true_vit = [ 1, 1, 1, 2, 2, 1, 1, 1, 1 ]
@test(all(viterbi(model,o) .== true_vit))

# check Baum-Welch algorithm
o = [ 1, 2, 1, 6, 6, 3, 3, 2, 1 ]
expected_result_A = [ 0.75  0.25 ;
                      0.25  0.75 ]
expected_result_B = [ 0.6  0.4  0.0  0.0  0.0  0.0 ;
                      0.0  0.0  0.5  0.0  0.0  0.5 ]
expected_result_p = [1.0, 0.0]

model1 = deepcopy(model) # check without scaling alpha/beta
fit!(model1, o; max_iter=60, scaling=false)

@test(all(round(model1.A,4) .== expected_result_A))
@test(all(round(model1.B[1].p,4) .== expected_result_B[1,:]'))
@test(all(round(model1.B[2].p,4) .== expected_result_B[2,:]'))
@test(all(round(model1.p,4) .== expected_result_p))

model2 = deepcopy(model) # check with scaling alpha/beta
fit!(model2, o; max_iter=60, scaling=true)
@test(all(round(model2.A,4) .== expected_result_A))
@test(all(round(model2.B[1].p,4) .== expected_result_B[1,:]'))
@test(all(round(model2.B[2].p,4) .== expected_result_B[2,:]'))
@test(all(round(model2.p,4) .== expected_result_p))


Rectangular() = Uniform(-0.5, +0.5)

function _string(kernel::Uniform)
    @unpack a, b = kernel
    "Rectangular kernel (U[$a,$b])"
end

const LocationScaleDists = Union{Cosine,Epanechnikov,SymTriangularDist,Triweight}

_string(kernel::SymTriangularDist) = "Triangular kernel"
_string(kernel::LocationScaleDists) = "$(Base.typename(typeof(kernel))) kernel"

const SupportedKernels = Union{LocationScaleDists,Uniform}

function setbandwidth(kernel::LocationScaleDists, h::Number)
    typeof(kernel)(zero(kernel.μ), h * kernel.σ)
end

function setbandwidth(kernel::Uniform, h::Number)
    @unpack a, b = kernel
    if a + b != 0
        error("Use symmetric Uniform/Rectangular kernel")
    end
    Uniform(a * h, b * h)
end

weights(::Uniform, ZsR::RunningVariable) = uweights(length(ZsR))
weights(D::Distribution, ZsR::RunningVariable) = pdf.(D, ZsR)

struct EquivalentKernel{K,ET,T}
    kernel::K #only order 1 for now
    E::ET
    kernel_mult::T
end

function _expectation(ub, n)
    rawNodes, rawWeights = gausslegendre(n)
    _nodes = map(x -> (0.5 * ub) * x + ub / 2, rawNodes)
    _weights = map(x -> x * 1 / 2 * ub, rawWeights) # (result of doing 1/(b-a) * (b-a)/2)
    f -> dot(f.(_nodes), _weights)
end

function EquivalentKernel(kernel::SupportedKernels)
    @unpack ub = support(kernel)
    E = _expectation(ub, 10)
    cov_x = [x->1 x->x; x->x x->x^2]
    kernel_mult = [1 0] * inv(map(f -> E(u -> f(u) * pdf(kernel, u)), cov_x))
    EquivalentKernel(kernel, E, kernel_mult)
end

Distributions.support(eq::EquivalentKernel) = eq.kernel

function Distributions.pdf(eq::EquivalentKernel, u)
    @unpack kernel_mult, kernel = eq
    dot(kernel_mult, [1; u]) * pdf(kernel, u)
end

function kernel_moment(kernel, ::Val{j}) where {j}
    @unpack ub = support(kernel)
    quadgk(x -> x^j * pdf(kernel, x), 0, ub)[1]
end

kernel_moment(::SupportedKernels, ::Val{0}) = 0.5
kernel_moment(::SupportedKernels, ::Val{1}) = 0
kernel_moment(::SupportedKernels, ::Val{3}) = 0

function kernel_moment(eq::EquivalentKernel, ::Val{j}) where {j}
    @unpack E = eq
    E(x -> x^j * pdf(eq, x))
end

function squared_kernel_moment(kernel, ::Val{j}) where {j}
    @unpack ub = support(kernel)
    quadgk(x -> x^j * pdf(kernel, x)^2, 0, ub)[1]
end

squared_kernel_moment(::SupportedKernels, ::Val{1}) = 0
squared_kernel_moment(::SupportedKernels, ::Val{3}) = 0

function squared_kernel_moment(eq::EquivalentKernel, ::Val{j}) where {j}
    @unpack E = eq
    E(x -> x^j * pdf(eq, x)^2)
end


# In covariance adaptation we want to use cholesky factorization instead
# of a full eigendecomposition. Can we speed this up with sparse matrices in
# Julia?

# Objective function to be maximized.
function my_fun1(x)
  exp(-(x-2).^2) + 0.8 * exp(-(x+2).^2)
end

Ns = [10, 100, 1000, 2000]
num_reps = 10
Ds = [0.01, 0.10, 0.50]

function non_sparse_sqrtm(C, randsamples)
  sqrtC = sqrtm(C)
  s = sqrtC * randsamples
  C + (s * s')
end

function sparse_cholfact(C, randsamples)
  sqrtC = sparse(cholfact(C))
  s = sqrtC' * randsamples
  C = C + (s * s')
end

for(n in Ns)
  # Non-sparse, sqrtm
  C = eye(n)
  if n <= 2000
    tic()
    for(rep in 1:num_reps)
      C = non_sparse_sqrtm(C, randn(n, n))
    end
    t = toq()
    tpdn = t/(num_reps*n)
    println("Non-sparse, sqrtm, n = $(n), time = $(t/num_reps), time = $(tpdn)")
  end

  for(d in Ds)
    # Sparse, cholfact, randn
    C = speye(n)
    tic()
    for(rep in 1:num_reps)
      C = sparse_cholfact(C, sprandn(n, n, d))
    end
    t = toq()
    tpds = t/(num_reps*n)
    println("Sparse, cholfact, d = $(d), n = $(n), time = $(t/num_reps), time = $(tpds), $(tpdn/tpds) times faster")
  end
end


@testset "arithmetic.jl" begin

    # Sampling based validation

    Nsamps = 10^5

    a = interval(0,2) ∪ interval(3,4)
    b = interval(-10,-9) ∪ interval(20, 25)

    aSamps = zeros(Nsamps); bSamps = zeros(Nsamps)

    # Sample inside set
    for i = 1: Nsamps
        if rand() < 0.5
            aSamps[i] = rand() * 2
        else
            aSamps[i] = rand() + 3
        end

        if rand() < 0.5
            bSamps[i] = (rand() * 5) + 20
        else
            bSamps[i] = rand() - 10
        end
    end

    @test all(aSamps .∈ a)
    @test all(bSamps .∈ b)

    BivOps = [+, -, *, / , min, max, ^] # Error with /
    UniOps = [-, sin, cos, tan, exp, log, sqrt]

    for op in BivOps

        cInt = op(a, b)
        cSamps = op.(aSamps, bSamps)
        @test all( cSamps .∈ cInt)

    end

    for op in UniOps

        cInt = op(a)
        cSamps = op.(aSamps)
        @test all( cSamps .∈ cInt)

    end

    # With 0 ∈ b

    a = interval(0,2) ∪ interval(3,4)
    b = interval(-5,5) ∪ interval(20, 25)

    for i = 1: Nsamps
        if rand() < 0.5
            aSamps[i] = rand() * 2
        else
            aSamps[i] = rand() + 3
        end

        if rand() < 0.5
            bSamps[i] = (rand() * 5) + 20
        else
            bSamps[i] = rand() * 10  - 5
        end
    end

    @test all(aSamps .∈ a)
    @test all(bSamps .∈ b)

    BivOps = [+, -, *, / , min, max] # Error with /

    for op in BivOps

        cInt = op(a, b)
        cSamps = op.(aSamps, bSamps)
        @test all( cSamps .∈ cInt)

    end

end


using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("colorscheme.jl"))

using InteractiveDynamics
using DynamicalSystems
using GLMakie

using OrdinaryDiffEq: Vern9
ds = Systems.henonheiles()

u0s = (
    [0.0, -0.25, 0.42, 0.0], # chaotic
    [0.0, 0.1, 0.5, 0.0], # quasiperiodic
    [0.0, 0.30266571044921875, 0.4205654433900762, 0.0], # periodic
)

diffeq = (alg = Vern9(), dtmax = 0.05)
idxs = (1, 2, 4)

fig, obs = interactive_evolution(
    ds, u0s; idxs, tail = 2000, diffeq, colors = COLORS[1:3],
)

ax = content(fig[1,1])
ax.xlabel = "q₁"
ax.ylabel = "q₂"
ax.zlabel = "v₂"
# main.figure[AbstractPlotting.Axis][:names, :axisnames] = ("q₁", "q₂", "v₂")

# %%
record_interaction(fig, string(@__FILE__)[1:end-2]*"mp4"; total_time = 15)


struct MAParams <: AbstractParams
    q::Vector{<:Integer}
end

MAParams(q::Integer) = MAParams(collect(1:q))

function fit(params::MAParams, z::AbstractVector{T}; n::Integer=20) where T
    q = params.q
    Q = maximum(q)
    if q == collect(1:Q)
        μ = mean(z)
        z = z .- μ
        θ, σ2 = innovations(z, Q, n=n)
        ma = MAModel{Q}(μ, θ, σ2)
    else
        ma = fit(ARMAParams(false, eltype(q)[], q), z, n=n)
    end
    return ma
end

const MAModel{q,T} = ARMAXModel{0,q,0,T}
MAModel{q}(μ::T,θ::AbstractVector{T},σ2::T) where {q,T} = ARMAXModel{0,q,0,T}(μ, T[], θ, T[], σ2)

isstationary(model::MAModel) = true

isinvertible(model::MAModel) = all(norm.(roots(Polynomial([1;.-model.θ]))) .> 1)

function innovations!(θ::AbstractMatrix{T}, σ2::AbstractVector{T}, γ::AbstractVector{T}, q::Integer, n::Integer) where {T}
    θ[n,n] = γ[1 + n] / σ2[1]
    for k in 2:n
        tmp = zero(T)
        for j in 1:k
            tmp += θ[k - 1,(k - 1) - (j - 1) + 1] * θ[n,n - (j - 1)] * σ2[1 + (j - 1)]
        end
        θ[n,n - (k - 1)] =  n - (k - 1) > q ? zero(T) : (-one(T) / σ2[1 + (k - 1)]) * (γ[1 + n - (k - 1)] - tmp)
    end
    σ2[n + 1] = γ[1]
    for j in 1:n
        σ2[n + 1] -= θ[n,n - (j - 1)]^2 * σ2[1 + (j - 1)]
    end
end

function innovations(z::AbstractVector, q::Integer; n::Integer=20)
    n >= q || throw(ArgumentError("number of iterations should be higher than the number of coefficients q"))
    γ = map(i -> autocovariance(z, i), 0:n)
    T = typeof(zero(eltype(z)) / 1)
    σ2 = Vector{T}(undef, n + 1)
    σ2[1] = γ[1]
    θ = zeros(T, n, n)
    for k in 1:n
        innovations!(θ, σ2, γ, q, k)
    end
    return θ[n,1:q], σ2[n + 1]
end


module MatrixDecomposition

using BenchmarkTools
using StructuredOptimization
using Images, ImageView

function set_up()
    # dataset from: http://pione.dinf.usherbrooke.ca/static/dataset/clutter/IndianTraffic3.zip
	Frames = 900
	dir = "street"
	img = load("$dir/in000000.jpg")
	n,m = size(img,1),size(img,2)
    frames = collect(1:135:Frames)
	N = length(frames)

	Y = zeros(Float64,n,m,N)
	for f in eachindex(frames)
		a = @sprintf("%6.6i",frames[f])
		img = load("$dir/in$(a).jpg")
		Y[:,:,f] .= convert(Array{Float64},Gray.(img))
		img = []
	end

	Y = reshape(Y,n*m,N)

	F = Variable(n*m,N)
	B = Variable(n*m,N)

    println("frames: $N, number of variables: $(2*N*m*n)")

	R, lambda = 1, 3e-2
	return B, F, Y, R, lambda, n, m, N
end

function run_demo()
	slv = PANOC(tol = 1e-4, verbose = 2)
	setup = set_up()
	solve_problem!(slv,setup...)
	return setup
end

function solve_problem!(slv, B, F, Y, R, lambda, n, m, N)
	it, = @minimize ls(B+F-Y)+lambda*norm(F,1) st rank(B) <= R with slv
	return it
end

function benchmark(;verb = 0, samples = 5, seconds = 100, tol = 1e-4)

	suite = BenchmarkGroup()

	solvers = [
             #  "ZeroFPR",
               "PANOC",
               "PG"
              ]
	slv_opt = [
              # "(verbose = $verb, tol = $tol)",
               "(verbose = $verb, tol = $tol)",
               "(verbose = $verb, tol = $tol)"
              ]

	its = Dict([(sol,0.) for sol in solvers])
	for i in eachindex(solvers)

		setup = set_up()
		solver = eval(parse(solvers[i]*slv_opt[i]))

		suite[solvers[i]] =
		@benchmarkable(it = solve_problem!(solver, setup...),
                       setup = (
                                it = 0;
                                setup = deepcopy($setup);
                                solver = deepcopy($solver) ),
                       teardown = (
                                   $its[$solvers[$i]] = it;
                                  ),
                       evals = 1, samples = samples, seconds = seconds)
	end

	results = run(suite, verbose = (verb != 0))
	println("MatrixDecomposition its")
	println(its)
	return results
end

function show_results(B, F, Y, R, lambda, n, m, N)
	F = F.x
	B = B.x

	F[F .!=0] = F[F.!=0]+reshape(B,size(F))[F.!=0]
	F[F.== 0] = 1. #put white in null pixels
	F[F.>1] .= 1; F[F.<0.] .= 0.
	B[B.>1] .= 1; B[B.<0.] .= 0.

	#convert back to images
	Y = reshape(Y,n,m,N)
	B = reshape(B,n,m,N)
	F = reshape(F,n,m,N)
	#
	####videos
	#imshow(Y)
	#imshow(F)
	#imshow(B)
	#
	##images
	idx = round.(Int,linspace(1,N,4))
	imshow(hcat([Y[:,:,i] for i in idx]...))
	imshow(hcat([F[:,:,i] for i in idx]...))
	imshow(hcat([B[:,:,i] for i in idx]...))

    save_stuff = false
    if save_stuff == true
        for (i,idxi) in enumerate(idx)
            save("Y$i.jpg", Y[:,:,idxi])
            save("F$i.jpg", F[:,:,idxi])
            save("B$i.jpg", B[:,:,idxi])
        end
    end

end

end