# grids.jl : introduces the Grid type for easy use of multigrid algorithms

# Ax=b
struct Grid{matT<:SparseMatrixCSC, vecT<:AbstractVector, opeT, tupT<:NTuple}
    A::matT # matrix A
    x::vecT # solution vector
    b::vecT # rhs vector
    R::opeT # restriction matrix
    P::opeT # prolongation matrix
    sz::tupT # PDE grid size
end

ndims(grid::Grid) = length(size(grid))

size(grid::Grid) = grid.sz

size(grid::Grid, n::Integer) = grid.sz[n]

show(io::IO, grid::Grid) = print(io, string(join(grid.sz," x ")," grid"))

# coarsen the problem A into a sequence of coarser grids
"""
    coarsen(A, sz, R_op, P_op, ngrids)

Coarsen the matrix `A` using `R_op` as restriction operator and `P_op` as interpolation operator. The matrix `A` is the discrete version of a PDE defined on an `sz`-point mesh, and `ngrids` is the number of levels.  A Galerkin approach is used to compose the coarse matrices, unless a function `f` is provided for direct discretization.
"""
function coarsen(A::SparseMatrixCSC, sz::NTuple, R_op::TransferKind, P_op::TransferKind, ngrids::Int)
    grids = Vector{Grid{typeof(A),Vector{eltype(A)},typeof(A),typeof(sz)}}(undef, ngrids)
    grids[1] = Grid(A,zero_x(A),zero_x(A),R(R_op,sz...),spzeros(0,0),sz)
    for i in 2:ngrids
        sz_c = grids[i-1].sz.>>1
        R_mat = R(R_op,sz_c...)
        P_mat = P(P_op,sz_c...)
        A_c = grids[i-1].R*grids[i-1].A*P_mat
        grids[i] = Grid(A_c,zero_x(A_c),zero_x(A_c),R_mat,P_mat,sz_c)
    end
    grids
end

function coarsen(f::Function, sz::NTuple, R_op::TransferKind, P_op::TransferKind, ngrids::Int)
	A0 = f(sz...)
    grids = Vector{Grid{typeof(A0),Vector{eltype(A0)},typeof(A0),typeof(sz)}}(undef, ngrids)
    grids[1] = Grid(A0,zero_x(A0),zero_x(A0),R(R_op,sz...),spzeros(0,0),sz)
    for i in 2:ngrids
        sz_c = grids[i-1].sz.>>1
        R_mat = R(R_op, sz_c...)
        P_mat = P(P_op, sz_c...)
        A_c = f(sz_c...)
        grids[i] = Grid(A_c, zero_x(A_c), zero_x(A_c), R_mat, P_mat, sz_c)
    end
    grids
end

# automatically determine number of grids
# TODO: because of the way the grid transfer operators are implemented
# we can only do prolongation for grids with an even number of points...
factor_twos(n::Int) = sum(factor(Vector, n) .== 2)

zero_x(A::SparseMatrixCSC) = zeros(eltype(A),size(A,1))

# compute the residu on this level
residu(grid::Grid) = grid.b - grid.A*grid.x

# compute the h-norm of the residu on this level
h_norm(v::AbstractVector, sz) = 1.0/prod(sz) * norm(v)
norm_of_residu(grid::Grid) = h_norm(residu(grid), grid.sz)

# apply smoother
smooth!(grid::Grid, ν::Int, smoother::Smoother) = smooth!(grid.x, grid.A, grid.b, ν, smoother)

# TODO this should be improved...
function smooth!(grid::Grid, ν::Int, smoother::RedBlackGaussSeidel)
    for color in 1:2
        x_copy = copy(grid.x)
        smooth!(x_copy, grid.A, grid.b, ν, GaussSeidel())
        for i in color:2:size(grid.A, 1)
            grid.x[i] = x_copy[i]
        end
    end
end


"""
    Sphere{T,N} <: ParametricSurface{T,N}

Spherical surface centered at the origin.

```julia
Sphere(radius::T = 1.0; interface::NullOrFresnel{T} = nullinterface(T))
```
"""
struct Sphere{T,N} <: ParametricSurface{T,N}
    radius::T
    interface::NullOrFresnel{T}

    function Sphere(radius::T; interface::NullOrFresnel{T} = NullInterface(T)) where {T<:Real}
        @assert !isnan(radius)
        @assert radius > zero(T)
        return new{T,3}(radius, interface)
    end
end
export Sphere

interface(a::Sphere{T,N}) where {T<:Real,N} = a.interface
radius(a::Sphere{T}) where {T<:Real} = a.radius

uvrange(::Type{Sphere{T,N}}) where {T<:Real,N} = ((-T(π), T(π)), (zero(T), T(π)))

onsurface(sph::Sphere{T,3}, x::T, y::T, z::T) where {T<:Real} = samepoint(x^2 + y^2 + z^2, radius(sph)^2)

inside(sph::Sphere{T,3}, x::T, y::T, z::T) where {T<:Real} = x^2 + y^2 + z^2 - radius(sph)^2 < zero(T)

point(sph::Sphere{T,3}, ϕ::T, θ::T) where {T<:Real} = SVector{3,T}(radius(sph) * sin(θ) * cos(-ϕ), radius(sph) * sin(θ) * sin(-ϕ), radius(sph) * cos(θ))

normal(::Sphere{T,3}, ϕ::T, θ::T) where {T<:Real} = SVector{3,T}(sin(θ) * cos(-ϕ), sin(θ) * sin(-ϕ), cos(θ))

partials(sph::Sphere{T,3}, ϕ::T, θ::T) where {T<:Real} = SVector{3,T}(radius(sph) * sin(θ) * -sin(-ϕ), radius(sph) * sin(θ) * cos(-ϕ), 0.0), SVector{3,T}(radius(sph) * cos(θ) * cos(-ϕ), radius(sph) * cos(θ) * sin(-ϕ), radius(sph) * -sin(θ))
function uv(::Sphere{T,3}, x::T, y::T, z::T) where {T<:Real}
    # avoid divide by zero for ForwardDiff
    ϕ = -NaNsafeatan(y, x)
    if x == zero(T) && y == zero(T)
        θ = zero(T)
    else
        θ = NaNsafeatan(sqrt(x^2 + y^2), z)
    end
    return SVector{2,T}(ϕ, θ)
end

# Assumes the ray has been transformed into the canonical sphere coordinate frame which has the vertical axis passing through (0,0,0) and aligned with the z axis.
function surfaceintersection(sph::Sphere{T,N}, r::AbstractRay{T,N}) where {T<:Real,N}
    ox, oy, oz = origin(r)
    dx, dy, dz = direction(r)
    rad = radius(sph)

    a = dx^2 + dy^2 + dz^2
    b = 2 * (ox * dx + oy * dy + oz * dz)
    c = (ox^2 + oy^2 + oz^2) - rad^2

    temp = quadraticroots(a, b, c)

    if temp === nothing
        return EmptyInterval(T) # no intersection with sphere and ray not contained entirely in sphere
    end

    t1, t2 = temp

    if isapprox(t1, t2, rtol = 1e-12, atol = 2 * eps(T))
        return EmptyInterval(T) # single root which indicates a tangent ray sphere intersection
    end

    if t1 > zero(T)
        pt1 = point(r, t1)
    else
        pt1 = nothing
    end

    if t2 > zero(T)
        pt2 = point(r, t2)
    else
        pt2 = nothing
    end

    let int1 = nothing, int2 = nothing
        if pt1 !== nothing
            θ, ρ = uv(sph, pt1)
            int1 = Intersection(t1, pt1, pt1, θ, ρ, interface(sph))
        end

        if pt2 !== nothing
            θ, ρ = uv(sph, pt2)
            int2 = Intersection(t2, pt2, pt2, θ, ρ, interface(sph))
        end

        if int1 !== nothing && int2 !== nothing
            if t1 <= t2
                return Interval(int1, int2)
            else
                return Interval(int2, int1)
            end
        elseif int1 !== nothing
            return rayorigininterval(int1)
        elseif int2 !== nothing
            return rayorigininterval(int2)
        else
            return EmptyInterval(T)
        end
    end
end

BoundingBox(sph::Sphere{T,3}) where {T<:Real} = BoundingBox(-radius(sph), radius(sph), -radius(sph), radius(sph), -radius(sph), radius(sph))


module GaussianProcesses

using Optim, PDMats, ElasticPDMats, ElasticArrays, Distances, FastGaussQuadrature, RecipesBase, Distributions
using StaticArrays
using StatsFuns, SpecialFunctions

using LinearAlgebra, Printf, Random, Statistics
import Statistics: mean

# Functions that should be available to package
# users should be explicitly exported here
export GPBase, GP, GPE, GPMC, predict_f, predict_y, Kernel, Likelihood, CompositeKernel, SumKernel, ProdKernel, Masked, FixedKernel, fix, Noise, Const, SE, SEIso, SEArd, Periodic, Poly, RQ, RQIso, RQArd, Lin, LinIso, LinArd, Matern, Mat12Iso, Mat12Ard, Mat32Iso, Mat32Ard, Mat52Iso, Mat52Ard, #kernel functions
    MeanZero, MeanConst, MeanLin, MeanPoly, SumMean, ProdMean, #mean functions
    GaussLik, BernLik, ExpLik, StuTLik, PoisLik, BinLik,       #likelihood functions
    mcmc, optimize!,                                           #inference functions
    set_priors!,set_params!, update_target!

const MatF64 = AbstractMatrix{Float64}
const VecF64 = AbstractVector{Float64}

const φ = normpdf
const Φ = normcdf
const invΦ = norminvcdf

# all package code should be included here
include("means/means.jl")
include("kernels/kernels.jl")
include("likelihoods/likelihoods.jl")
include("common.jl")
include("utils.jl")
include("chol_utils.jl")
include("GP.jl")
include("GPE.jl")
include("GPEelastic.jl")
include("GPMC.jl")
include("mcmc.jl")
include("optimize.jl")
include("plot.jl")

# ScikitLearnBase, which is a skeleton package.
include("ScikitLearn.jl")

end # module


module IntegerMathUtils
export iroot, ispower, rootrem, find_exponent, is_probably_prime

iroot(x::Integer, n::Integer) = iroot(x, Cint(n))

# TODO: Add more efficient implimentation
iroot(x::T, n::Cint) where {T<:Integer} = T(iroot(big(x), Cint(n)))

function iroot(x::BigInt, n::Cint)
    n <= 0 && throw(DomainError(n, "Exponent must be > 0"))
    x <= 0 && iseven(x) && throw(DomainError(n, "This is a math no-no"))
    ans = BigInt()
    ccall((:__gmpz_root, :libgmp), Cint, (Ref{BigInt}, Ref{BigInt}, Cint), ans, x, n)
    ans
end

# TODO: Add more efficient implimentation for smaller types
function rootrem(x::T, n::Integer) where {T<:Integer}
    x = big(x)
    n = Cint(n)
    n <= 0 && throw(DomainError(n, "Exponent must be > 0"))
    x <= 0 && iseven(x) && throw(DomainError(n, "This is a math no-no"))
    root = BigInt()
    rem  = BigInt()
    ccall((:__gmpz_rootrem, :libgmp), Nothing,(Ref{BigInt}, Ref{BigInt}, Ref{BigInt}, Int), root, rem, x, n)
    return (root, T(rem))
end

# TODO: Add more efficient implimentation for smaller types
ispower(x::Integer) = ispower(big(x))

function ispower(x::BigInt)
    return 0 != ccall((:__gmpz_perfect_power_p, :libgmp), Cint, (Ref{BigInt},), x)
end

# TODO: Add more efficient implimentation for smaller types
function find_exponent(x::Integer)
    x <= 1 && return 1
    for exponent in ndigits(x, base=2):-1:2
        rootrem(x, exponent)[2] == 0 && return exponent
    end
    1
end

function is_probably_prime(x::Integer; reps=25)
    return ccall((:__gmpz_probab_prime_p, :libgmp), Cint, (Ref{BigInt}, Cint), x, reps) != 0
end

end


@testset "Hessian of Lagrangian" begin

    # horizon
    T = 3

    # acrobot
    nx = 4
    nu = 1
    nw = 0
    w_dim = [nw for t = 1:T]

    function acrobot(x, u, w)
        # dimensions
        n = 4
        m = 1
        d = 0

        # link 1
        mass1 = 1.0
        inertia1 = 0.33
        length1 = 1.0
        lengthcom1 = 0.5

        # link 2
        mass2 = 1.0
        inertia2 = 0.33
        length2 = 1.0
        lengthcom2 = 0.5

        gravity = 9.81
        friction1 = 0.1
        friction2 = 0.1

        # mass matrix
        function M(x, w)
            a = (inertia1 + inertia2 + mass2 * length1 * length1
                + 2.0 * mass2 * length1 * lengthcom2 * cos(x[2]))

            b = inertia2 + mass2 * length1 * lengthcom2 * cos(x[2])

            c = inertia2

        return [a b; b c]
        end

        # dynamics bias
        function τ(x, w)
            a = (-1.0 * mass1 * gravity * lengthcom1 * sin(x[1])
                - mass2 * gravity * (length1 * sin(x[1])
                + lengthcom2 * sin(x[1] + x[2])))

            b = -1.0 * mass2 * gravity * lengthcom2 * sin(x[1] + x[2])

            return [a; b]
        end

        function C(x, w)
            a = -2.0 * mass2 * length1 * lengthcom2 * sin(x[2]) * x[4]
            b = -1.0 * mass2 * length1 * lengthcom2 * sin(x[2]) * x[4]
            c = mass2 * length1 * lengthcom2 * sin(x[2]) * x[3]
            d = 0.0

            return [a b; c d]
        end

        # input Jacobian
        function B(x, w)
            [0.0; 1.0]
        end

        # dynamics
        q = view(x, 1:2)
        v = view(x, 3:4)

        qdd = M(q, w) \ (-1.0 * C(x, w) * v
                + τ(q, w) + B(q, w) * u[1] - [friction1; friction2] .* v)

        return [x[3]; x[4]; qdd[1]; qdd[2]]
    end

    function midpoint_implicit(y, x, u, w)
        h = 0.05 # timestep
        y - (x + h * acrobot(0.5 * (x + y), u, w))
    end

    dt = Dynamics(midpoint_implicit, nx, nx, nu, nw=nw, eval_hess=true)
    dyn = [dt for t = 1:T-1]
    model = DynamicsModel(dyn, w_dim=w_dim)

    # initial state
    x1 = [0.0; 0.0; 0.0; 0.0]

    # goal state
    xT = [0.0; π; 0.0; 0.0]

    # objective
    ot = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4]) + 0.1 * dot(u, u)
    oT = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4])
    objt = Cost(ot, nx, nu, nw, [t for t = 1:T-1], eval_hess=true)
    objT = Cost(oT, nx, 0, nw, [T], eval_hess=true)
    obj = [objt, objT]

    # constraints
    x_init = Bound(nx, nu, [1], xl=x1, xu=x1)
    x_goal = Bound(nx, 0, [T], xl=xT, xu=xT)
    ct = (x, u, w) -> [-5.0 * ones(nu) - cos.(u) .* sum(x.^2); cos.(x) .* tan.(u) - 5.0 * ones(nx)]
    cT = (x, u, w) -> sin.(x.^3.0)
    cont = StageConstraint(ct, nx, nu, nw, [t for t = 1:T-1], :inequality, eval_hess=true)
    conT = StageConstraint(cT, nx, 0, nw, [T], :equality, eval_hess=true)
    cons = ConstraintSet([x_init, x_goal], [cont, conT])

    # problem
    trajopt = TrajectoryOptimizationProblem(obj, model, cons)
    s = Solver(trajopt, eval_hess=true)

    # Lagrangian
    function lagrangian(z)
        x1 = z[1:nx]
        u1 = z[nx .+ (1:nu)]
        x2 = z[nx + nu .+ (1:nx)]
        u2 = z[nx + nu + nx .+ (1:nu)]
        x3 = z[nx + nu + nx + nu .+ (1:nx)]
        λ1_dyn = z[nx + nu + nx + nu + nx .+ (1:nx)]
        λ2_dyn = z[nx + nu + nx + nu + nx + nx .+ (1:nx)]

        λ1_stage = z[nx + nu + nx + nu + nx + nx + nx .+ (1:(nu + nx))]
        λ2_stage = z[nx + nu + nx + nu + nx + nx + nx + nu + nx .+ (1:(nu + nx))]
        λ3_stage = z[nx + nu + nx + nu + nx + nx + nx + nu + nx + nu + nx .+ (1:nx)]

        L = 0.0
        L += ot(x1, u1, zeros(nw))
        L += ot(x2, u2, zeros(nw))
        L += oT(x3, zeros(0), zeros(nw))
        L += dot(λ1_dyn, midpoint_implicit(x2, x1, u1, zeros(nw)))
        L += dot(λ2_dyn, midpoint_implicit(x3, x2, u2, zeros(nw)))
        L += dot(λ1_stage, ct(x1, u1, zeros(nw)))
        L += dot(λ2_stage, ct(x2, u2, zeros(nw)))
        L += dot(λ3_stage, cT(x3, zeros(0), zeros(nw)))
        return L
    end

    nz = nx + nu + nx + nu + nx + nx + nx + nu + nx + nu + nx + nx
    np = nx + nu + nx + nu + nx
    nd = nx + nx + nu + nx + nu + nx + nx
    @variables z[1:nz]
    L = lagrangian(z)
    Lxx = Symbolics.hessian(L, z[1:np])
    Lxx_sp = Symbolics.sparsehessian(L, z[1:np])
    spar = [findnz(Lxx_sp)[1:2]...]
    Lxx_func = eval(Symbolics.build_function(Lxx, z)[1])
    Lxx_sp_func = eval(Symbolics.build_function(Lxx_sp.nzval, z)[1])

    z0 = rand(nz)
    nh = length(s.p.sp_hess_lag)
    h0 = zeros(nh)

    σ = 1.0
    fill!(h0, 0.0)
    DirectTrajectoryOptimization.trajectory!(s.p.trajopt.x, s.p.trajopt.u, z0[1:np],
        s.p.trajopt.model.idx.x, s.p.trajopt.model.idx.u)
    DirectTrajectoryOptimization.duals!(s.p.trajopt.λ_dyn, s.p.trajopt.λ_stage, z0[np .+ (1:nd)], s.p.idx.dyn_con, s.p.idx.stage_con)
    DirectTrajectoryOptimization.eval_obj_hess!(h0, s.p.idx.obj_hess, s.p.trajopt.obj, s.p.trajopt.x, s.p.trajopt.u, s.p.trajopt.w, σ)
    DirectTrajectoryOptimization.eval_hess_lag!(h0, s.p.idx.dyn_hess, s.p.trajopt.model.dyn, s.p.trajopt.x, s.p.trajopt.u, s.p.trajopt.w, s.p.trajopt.λ_dyn)
    DirectTrajectoryOptimization.eval_hess_lag!(h0, s.p.idx.stage_hess, s.p.trajopt.con.stage, s.p.trajopt.x, s.p.trajopt.u, s.p.trajopt.w, s.p.trajopt.λ_stage)

    sp_obj_hess = DirectTrajectoryOptimization.sparsity_hessian(obj, model.dim.x, model.dim.u)
    sp_dyn_hess = DirectTrajectoryOptimization.sparsity_hessian(dyn, model.dim.x, model.dim.u)
    sp_con_hess = DirectTrajectoryOptimization.sparsity_hessian(cons.stage, model.dim.x, model.dim.u)
    sp_hess = collect([sp_obj_hess..., sp_dyn_hess..., sp_con_hess...])
    sp_key = sort(unique(sp_hess))

    idx_obj_hess = DirectTrajectoryOptimization.hessian_indices(obj, sp_key, model.dim.x, model.dim.u)
    idx_dyn_hess = DirectTrajectoryOptimization.hessian_indices(dyn, sp_key, model.dim.x, model.dim.u)
    idx_con_hess = DirectTrajectoryOptimization.hessian_indices(cons.stage, sp_key, model.dim.x, model.dim.u)

    # indices
    @test sp_key[vcat(idx_obj_hess...)] == sp_obj_hess
    @test sp_key[vcat(idx_dyn_hess...)] == sp_dyn_hess
    @test sp_key[vcat(idx_con_hess...)] == sp_con_hess

    # Hessian
    h0_full = zeros(np, np)
    for (i, h) in enumerate(h0)
        h0_full[sp_key[i][1], sp_key[i][2]] = h
    end
    @test norm(h0_full - Lxx_func(z0)) < 1.0e-8
    @test norm(norm(h0 - Lxx_sp_func(z0))) < 1.0e-8

    # h0 = zeros(nh)
    # MOI.eval_hessian_lagrangian(s.p, h0, z0[1:np], σ, z0[np .+ (1:nd)])
    # @test norm(norm(h0 - Lxx_sp_func(z0))) < 1.0e-8

    # a = z0[1:np]
    # b = z0[np .+ (1:nd)]
    # info = @benchmark MOI.eval_hessian_lagrangian($s.p, $h0, $a, $σ, $b)
end


### wiener-cameraman.jl
#
# Copyright (C) 2016 Mosè Giordano.
#
# Maintainer: Mosè Giordano <mose AT gnu DOT org>
# Keywords: deconvolution, signal processing
#
# This file is a part of Deconvolution.jl.
#
# License is MIT "Expat".
#
### Commentary:
#
# This file produces an example image used in Deconvolution.jl documentation.
#
### Code:

using Images, TestImages, Deconvolution, Random, FFTW

Random.seed!(42)

# Open the test image
img = channelview(testimage("cameraman"))
# Create the blurring kernel in frequency domain
x = hcat(ntuple(x -> collect((1:512) .- 257), 512)...)
k = 0.001
blurring_ft = @. exp(-k*(x ^ 2 + x ^ 2)^(5//6))
# Create additive noise
noise = rand(Float64, size(img))
# Fourier transform of the blurred image, with additive noise
blurred_img_ft = fftshift(blurring_ft) .* fft(img) .+ fft(noise)
# Get the blurred image from its Fourier transform
blurred_img = real(ifft(blurred_img_ft))
# Get the blurring kernel in the space domain
blurring = ifft(fftshift(blurring_ft))
# Polish the image with Deconvolution deconvolution
polished = wiener(blurred_img, img, noise, blurring)

# Wiener deconvolution works also when you don't have the real image and noise,
# that is the most common and useful case.  This happens because the Wiener
# filter only cares about the power spectrum of the signal and the noise, so you
# don't need to have the exact signal and noise but something with a similar
# power spectrum.
img2 = channelview(testimage("livingroom")) # Load another image
noise2 = rand(Float64, size(img)) # Create another additive noise
# Polish the image with Deconvolution deconvolution
polished2 = wiener(blurred_img, img2, noise2, blurring)
# Wiener also works using a real number instead of a noise array
polished3 = 0.5 .* wiener(blurred_img, img2, 1000, blurring)
polished4 = 0.5 .* wiener(blurred_img, img2, 10000, blurring)

save("original.jpg", Images.clamp01nan.(img))
save("blurred.jpg", Images.clamp01nan.(blurred_img))
save("polished.jpg", Images.clamp01nan.(polished))
save("polished2.jpg", Images.clamp01nan.(polished2))
save("polished3.jpg", Images.clamp01nan.(polished3))
save("polished4.jpg", Images.clamp01nan.(polished4))


#################################################################
### This test was lifted from the vingette for the R gmm package.
#################################################################
using Distributions
using Linearalgebra
srand(42)
x = rand(Normal(4, 2), 1000)

function h_norm(θ)
    m1 = [θ[1]] .- x
    m2 = [θ[2]].^2 .- (x .- [θ[1]]).^2
    m3 = x.^3 .- [θ[1]].*([θ[1]].^2 .+ 3*[θ[2]].^2)
    return [m1 m2 m3]
end

step_1 = GMMEstimator(h_norm, [1.0, 1.0], initialW = Matrix(I, 3, 3), mgr = OneStepGMM())
@time estimate!(step_1);

ch(θ) = [1 -1]*θ
hlb = [0.]
hub = [0.]

cstep_1 = GMMEstimator(h_norm, [1.0, 1.0], initialW = Matrix(I, 3, 3), mgr = OneStepGMM(), constraints = Constrained(ch, hlb, hub, 1))

step_2_hac = GMMEstimator(h_norm, coef(step_1), initialW = MomentBasedEstimators.optimal_W(step_1, QuadraticSpectralKernel(0.91469)));

estimate!(step_2_hac);

step_2_iid = GMMEstimator(h_norm, coef(step_1), initialW = optimal_W(step_1, HC0()));
estimate!(step_2_iid);

gmm_qs_mgr = GMMEstimator(h_norm, [1.,1.], initialW = Matrix(I, 3, 3), mgr = TwoStepGMM(QuadraticSpectralKernel(0.91469)));
estimate!(gmm_qs_mgr);

gmm_iid_mgr = GMMEstimator(h_norm, [1.,1.], initialW = Matrix(I, 3, 3), mgr = TwoStepGMM(HC0()));
estimate!(gmm_iid_mgr);

### MDE
kl_iid = MDEstimator(h_norm, [1.0, 1.0])
estimate!(kl_iid)

kl_smt = MDEstimator(h_norm, [1.0, 1.0], kernel = MomentBasedEstimators.TruncatedSmoother(2))
estimate!(kl_smt)

el_iid = MDEstimator(h_norm, [1.0, 1.0], div = Divergences.ReverseKullbackLeibler())
cu_iid = MDEstimator(h_norm, [1.0, 1.0], div = Divergences.ChiSquared())


"""
Contains functions that take in time-invariant network structure parameters,
such as the coupling matrix J and gain parameter g, and run simulations of
network activity. Includes functions that numerically calculate network
properties (such as Lyapunov spectrum) in parallel with the simulation.
"""

using LinearAlgebra

#-------------------------------------------------------------------------------
"""
	sim_activity(J::Matrix{T}, g::T, h0::Vector{T}; dt=0.1, t_bound=1000.0,
		s=0.0, savestep=1, minimum_norm=0) where T <: AbstractFloat

Integrates a tanh rate network with coupling matrix `J` and strength scaling `g`
from initial condition `h0`.

When `minimum_norm > 0` the integration will be cut off when the network state
falls too close to the origin to prevent floating point underflow.
"""
function sim_activity(J::Matrix{T}, g::T, h0::Vector{T}; dt=0.1, t_bound=1000.0,
		s=0.0, savestep=1, minimum_norm=0) where T <: AbstractFloat

	N = size(J)[1]
	# Check network size and type consistency
	@assert N == size(J)[2] == length(h0)
	@assert T == typeof(dt) == typeof(t_bound) == typeof(s)

	maxstep = round(Int, t_bound/dt) # Discretized simulation length

	# Allocate memory for output
	out = zeros((N, Int(maxstep/savestep)))

	h = copy(h0) # Set initial condition
	r = Vector{T}(undef, N)

	J = g .* J + s * I

	step = 0
	while step < maxstep
		step += 1
		map!(tanh, r, h)
		BLAS.gemv!('N', dt, J, r, 1 - dt, h)
		if step % savestep == 0 # Save every specified time interval
			savept = step ÷ savestep
			out[:,savept] .= h
		end
		if norm(h) < minimum_norm
			D = -I + J .* ((sech.(h)).^2)'
			if all(real(eigvals(D)) .< 0)
				break
			end
		end
	end

	lastindex = step ÷ savestep
	return out[:,1:lastindex]
end

"""
	sim_lyapunov(J::Matrix{T}, g::T, h0::Vector{T}, dh::Vector{T},
		Q0::Matrix{T}; s=0.0, dt=0.1, t_bound=1000.0, t_ONS=1.0, wait=0.0,
		return_max_LE=false, printstuff=false) where T <: AbstractFloat

Integrates a tanh network with a given weight matrix and gain parameter from a
specified initial condition, and returns a Lyapunov spectrum calculated along
the trajectory.

Optionally returns the maximum Lyapunov exponent calculated using a direct
numerical simulation of a perturbed trajectory.
"""
function sim_lyapunov(J::Matrix{T}, g::T, h0::Vector{T}, dh::Vector{T},
		Q0::Matrix{T}; s=0.0, dt=0.1, t_bound=1000.0, t_ONS=1.0, wait=0.0,
		return_max_LE=false, printstuff=false) where T <: AbstractFloat

	N = size(J)[1]
	# Check dimensional consistency
	@assert size(J) == (length(h0), length(dh)) == size(Q0) == (N,N)
	@assert T == typeof(s) == typeof(dt) == typeof(t_bound) == typeof(t_ONS) == typeof(wait)

	# Allocate memory for output
	spectrum = zeros(T, N)

	maxstep = Int(t_bound / dt) # When to stop the simulation
	startstep = Int(wait / dt) # When to start tracking Lypunov exponents
	tONSstep = Int(t_ONS / dt) # Interval between re-orthonormalization
	try
		Int(t_bound/t_ONS)
		Int(wait/t_ONS)
	catch e
		if isa(e, InexactError)
			msg = "t_ONS must divide t_bound and wait"
			throw(ArgumentError(msg))
		end
	end

	# Allocate vector for activity and firing rate of current state
	h = copy(h0)
	r = Vector{T}(undef, N)

	# Allocate memory for Jacobian calculation and orthonormal basis
	sech2(x::T)::T = (sech(x))^2
	sech2h = Matrix{T}(undef, 1, N)
	D = Matrix{T}(undef, N, N)
	Q = copy(Q0)
	dQ = Matrix{T}(undef, N, N)

	if return_max_LE
		max_LE = 0.0
		p_factor = T(1e-6)  # Order of the perturbation handled by the sim
		dh = copy(dh) ./ norm(dh)
		h2 = h .+ p_factor .* dh
		r2 = Vector{T}(undef, N)
	end

	# The effective coupling matrix J
	J = g .* J + s * I

	step = 0
	while step < maxstep
		step += 1

		#Evolve the trajectory and the orthonormal basis in the tangent space
		map!(tanh, r, h)
		BLAS.gemv!('N', dt, J, r, 1 - dt, h)
		map!(sech2, sech2h, h)
		copyto!(D, J)
		D .*= sech2h
		D[diagind(D)] .-= 1
		BLAS.gemm!('N', 'N', dt, D, Q, T(0), dQ)
		Q .+= dQ

		#Evolve the perturbed trajectory and calculate how much it has grown
		if return_max_LE
			map!(tanh, r2, h2)
			BLAS.gemv!('N', dt, J, r2, 1 - dt, h2)
			dh .= (h2 .- h) ./ p_factor
			scale = norm(dh)

			# If the perturbation scale has grown or shrunk too much, then
			# reset it to 1 and track the cumulative sum of log scales
			if scale > 100 || scale < 0.01 || step == startstep
				h2 .= h .+ dh ./ scale .* p_factor
				if step > startstep
					max_LE += log(scale)
				end
			end
		end

		# At each re-orthonormalization step, the log diagonal entries of R are
		# saved in a cumulative sum
		if step % tONSstep == 0
			printstuff && println("Time step: ", step)
			copyto!(dQ, Q)
			F = qr!(dQ)
			diags = view(F.factors, diagind(F.factors))
			if step > startstep  # forcing positive R-diagonal elements
				spectrum .+= log.(abs.(diags))
			end
			fill!(Q, 0)
			map!(sign, view(Q, diagind(Q)), diags)
			lmul!(F.Q, Q)
		end
	end

	# Update the maximum Lyapunov exponent after the simulation has ended so
	# that the full simulation time is used for the calculation
	if return_max_LE
		max_LE += log(norm((h2 .- h) ./ p_factor))
		return (spectrum ./ (t_bound - wait), max_LE / (t_bound - wait))
	else
		return spectrum ./ (t_bound - wait)
	end
end

"""
	sim_direct_max_LE(J::Matrix{Float64}, g::Float64, h0::Vector{Float64},
		dh::Vector{Float64}; dt=0.1, t_bound=1000.0, wait=0.0)

Calculates the maximum Lyapunov exponent of a tanh rate network given by
coupling matrix J, dependent on a given initial condition and initial
perturbation, using a direct method that tracks the growth of the perturbed
trajectory relative to the unperturbed one.
"""
function sim_direct_max_LE(J::Matrix{T}, g::T, h0::Vector{T}, dh::Vector{T};
		s=0.0, dt=0.1, t_bound=1000.0, wait=0.0, printstuff=false) where T <: AbstractFloat

	N = size(J)[1]
	# Check dimensional consistency
	@assert size(J) == (length(h0), length(dh)) == (N,N)
	p_factor = 10^-10  # Order of the perturbation handled by the sim

	# Allocate memory for output
	max_LE = 0.0

	maxstep = Int(t_bound/dt) # When to stop the simulation
	startstep = Int(wait/dt) # When to start tracking Lypunov exponents

	# Set the initial condition and perturbed IC
	h = copy(h0)
	dh = copy(dh) ./ norm(dh)
	h2 = h .+ p_factor .* dh

	r = Vector{Float64}(undef, N)
	r2 = Vector{Float64}(undef, N)

	step = 0
	while step < maxstep
		step += 1

		# Evolve the trajectory and its perturbation
		map!(tanh, r, h)
		map!(tanh, r2, h2)
		BLAS.gemv!('N', dt*g, J, r, 1.0-dt, h)
		BLAS.gemv!('N', dt*g, J, r2, 1.0-dt, h2)
		copyto!(dh, h2)
		dh .-= h
		dh ./= p_factor
		scale = norm(dh)

		# If the perturbation scale has grown or shrunk too much, then
		# reset it to 1 and track the cumulative sum of log scales
		if scale > 10 || scale < 0.1 || step == startstep
			printstuff && println("Time step: ", step)
			copyto!(h2, h)
			dh ./= (scale / p_factor)
			h2 .+= dh
			if step > startstep
				max_LE += log(scale)
			end
		end
	end

	# Update the maximum Lyapunov exponent after the simulation has ended so
	# that the full simulation time is used for the calculation
	max_LE += log(norm((h2 .- h) ./ p_factor))
	return max_LE / (t_bound - wait)
end

"""
    sim_lyapunov(J::Matrix{Float64}, g::Float64, A::Matrix{Float64},
        onb::Matrix{Float64}; dt=0.1, t_bound=1000.0, t_ONS=1.0, wait=0.0,
		printstuff=false)

Calculates the Lyapunov spectrum of a tanh rate network with coupling matrix
`J` and gain parameter `g` by integrating an orthonormal system according to
the Jacobians in the tangent space along a trajectory given by activity
matrix `A`.
"""
function sim_lyapunov(J::Matrix{Float64}, g::Float64, A::Matrix{Float64},
		onb::Matrix{Float64}; dt=0.1, t_bound=1000.0, t_ONS=1.0, wait=0.0,
		printstuff=false)

	N = size(J)[1]
	# Check dimensional consistency
	@assert(size(J)==(size(A,1),N)==size(onb)==(N,N))

	# Ranges that index the diagonal entries of useful matrix sizes
	diagonals = diagind(J)

	# Allocate memory for output
	spectrum = zeros(N)

	maxstep = Int(t_bound/dt) # When to stop the simulation
	startstep = Int(wait/dt) # When to start tracking Lypunov exponents
	tONSstep = Int(t_ONS/dt) # Interval between re-orthonormalization
	@assert maxstep==size(A,2) "Simulation length must match time parameters"
	try
		Int(t_bound/t_ONS)
		Int(wait/t_ONS)
	catch e
		if isa(e, InexactError)
			msg = "t_ONS must divide t_bound and wait"
			throw(ArgumentError(msg))
		end
	end

	# Allocate memory for Jacobian calculation and orthonormal basis
	sech2(x::Float64)::Float64 = (sech(x))^2
	sech2h = Matrix{Float64}(undef, 1, N)
	D = Matrix{Float64}(undef, N, N)
	Q = copy(onb)
	dQ = Matrix{Float64}(undef, N, N)

	J = g .* J

	step = 0
	while step < maxstep
		step += 1

		#Evolve the trajectory and the orthonormal basis in the tangent space
		map!(sech2, sech2h, view(A, :, step))
		copyto!(D, J)
		D .*= sech2h
		D[diagonals] .-= 1.0
		BLAS.gemm!('N', 'N', dt, D, Q, 0.0, dQ)
		Q .+= dQ

		# At each re-orthonormalization step, the log diagonal entries of R are
		# saved in a cumulative sum
		if step % tONSstep == 0
			printstuff && println("Time step: ", step)
			copyto!(dQ, Q)
			F = qr!(dQ)
			diags = view(F.factors, diagonals)
			if step > startstep  # forcing positive R-diagonal elements
				spectrum .+= log.(abs.(diags))
			end
			fill!(Q, 0.0)
			map!(sign, view(Q, diagonals), diags)
			lmul!(F.Q, Q)
		end
	end

	# Update the maximum Lyapunov exponent after the simulation has ended so
	# that the full simulation time is used for the calculation
	return spectrum ./ (t_bound - wait)
end

"""
	sim_lyapunov_halfspec(J::Matrix{Float64}, g::Float64,
		A::Matrix{Float64}, onb::Matrix{Float64}; dt=0.1, t_bound=1000.0,
		t_ONS=1.0, wait=0.0, printstuff=false)

Calculates the Lyapunov spectrum of a tanh rate network with coupling matrix
`J` and gain parameter `g` by integrating an orthonormal system according to
the Jacobians in the tangent space along a trajectory given by activity
matrix `A`. Saves time by only calculating half of the Lyapunov spectrum
"""
function sim_lyapunov_halfspec(J::Matrix{Float64}, g::Float64,
		A::Matrix{Float64}, onb::Matrix{Float64}; dt=0.1, t_bound=1000.0,
		t_ONS=1.0, wait=0.0, printstuff=false)

	N = size(J)[1]
	@assert N%2==0 "Calculating half the spectrum only works for even N"
	Nh = Int(N/2)
	# Check dimensional consistency
	@assert(size(J)==(size(A,1),N)==size(onb)==(N,N))

	# Ranges that index the diagonal entries of useful matrix sizes
	diagonals = diagind(J)
	diags_1halfrng = range(1, step=N+1, length=Nh)
	diags_2halfrng = range(1 + Nh * (N + 1), step=N+1, length=Nh)

	# Allocate memory for output
	spectrum = zeros(Nh)

	maxstep = Int(t_bound/dt) # When to stop the simulation
	startstep = Int(wait/dt) # When to start tracking Lypunov exponents
	tONSstep = Int(t_ONS/dt) # Interval between re-orthonormalization
	@assert maxstep==size(A,2) "Simulation length must match time parameters"
	try
		Int(t_bound/t_ONS)
		Int(wait/t_ONS)
	catch e
		if isa(e, InexactError)
			msg = "t_ONS must divide t_bound and wait"
			throw(ArgumentError(msg))
		end
	end

	# Allocate memory for Jacobian calculation and orthonormal basis
	sech2(x::Float64)::Float64 = (sech(x))^2
	sech2h = Matrix{Float64}(undef, 1, N)
	D = Matrix{Float64}(undef, N, N)
	Q = copy(onb)
	dQ = Matrix{Float64}(undef, N, N)
	Qh = Matrix{Float64}(undef, N, Nh)

	J = g .* J

	step = 0
	while step < maxstep
		step += 1

		#Evolve the orthonormal basis in the tangent space
		map!(sech2, sech2h, view(A, :, step))
		copyto!(D, J)
		D .*= sech2h
		D[diagonals] .-= 1.0
		BLAS.gemm!('N', 'N', dt, D, Q, 0.0, dQ)
		Q .+= dQ

		# At each re-orthonormalization step, the log diagonal entries of R are
		# saved in a cumulative sum
		if step % tONSstep == 0
			printstuff && println("Time step: ", step)
			copyto!(Qh, view(Q, :, 1:Nh))
			F = qr!(Qh)
			diags = view(F.factors, diags_1halfrng)
			if step > startstep  # forcing positive R-diagonal elements
				spectrum .+= log.(abs.(diags))
			end
			fill!(Q, 0.0)
			map!(sign, view(Q, diags_1halfrng), diags)
			Q[diags_2halfrng] .= 1.0
			lmul!(F.Q, Q)
		end
	end

	# Update the maximum Lyapunov exponent after the simulation has ended so
	# that the full simulation time is used for the calculation
	return spectrum ./ (t_bound - wait)
end


using Base.Threads
using TensorOperations
import Base.+
import Base.*
import LinearAlgebra.dot

"""
Addition of two ttvector
"""
function +(x::ttvector{T},y::ttvector{T}) where T<:Number
    @assert(x.ttv_dims == y.ttv_dims, "Dimensions mismatch!")
    d = length(x.ttv_dims)
    ttv_vec = Array{Array{T,3},1}(undef,d)
    rks = x.ttv_rks + y.ttv_rks
    rks[1] = 1
    rks[d+1] = 1
    #initialize ttv_vec
    @threads for k in 1:d
        ttv_vec[k] = zeros(T,x.ttv_dims[k],rks[k],rks[k+1])
    end
   #first core
    ttv_vec[1][:,:,1:x.ttv_rks[2]] = x.ttv_vec[1]
    ttv_vec[1][:,:,(x.ttv_rks[2]+1):rks[2]] = y.ttv_vec[1]
    #2nd to end-1 cores
    @threads for k in 2:(d-1)
        ttv_vec[k][:,1:x.ttv_rks[k],1:x.ttv_rks[k+1]] = x.ttv_vec[k]
        ttv_vec[k][:,(x.ttv_rks[k]+1):rks[k],(x.ttv_rks[k+1]+1):rks[k+1]] = y.ttv_vec[k]
    end
    #last core
    ttv_vec[d][:,1:x.ttv_rks[d],1] = x.ttv_vec[d]
    ttv_vec[d][:,(x.ttv_rks[d]+1):rks[d],1] = y.ttv_vec[d]
    return ttvector{T}(ttv_vec,x.ttv_dims,rks,zeros(d))
end

"""
Addition of two ttoperators
"""
function +(x::ttoperator{T},y::ttoperator{T}) where T<:Number
    @assert(x.tto_dims == y.tto_dims, DimensionMismatch)
    d = length(x.tto_dims)
    tto_vec = Array{Array{T,4},1}(undef,d)
    rks = x.tto_rks + y.tto_rks
    rks[1] = 1
    rks[d+1] = 1
    #initialize tto_vec
    @threads for k in 1:d
        tto_vec[k] = zeros(T,x.tto_dims[k],x.tto_dims[k],rks[k],rks[k+1])
    end
    #first core
    tto_vec[1][:,:,:,1:x.tto_rks[1+1]] = x.tto_vec[1]
    tto_vec[1][:,:,:,(x.tto_rks[2]+1):rks[2]] = y.tto_vec[1]
    #2nd to end-1 cores
    @threads for k in 2:(d-1)
        tto_vec[k][:,:,1:x.tto_rks[k],1:x.tto_rks[k+1]] = x.tto_vec[k]
        tto_vec[k][:,:,(x.tto_rks[k]+1):rks[k],(x.tto_rks[k+1]+1):rks[k+1]] = y.tto_vec[k]
    end
    #last core
    tto_vec[d][:,:,1:x.tto_rks[d],1] = x.tto_vec[d]
    tto_vec[d][:,:,(x.tto_rks[d]+1):rks[d],1] = y.tto_vec[d]
    return ttoperator{T}(tto_vec,x.tto_dims,rks,zeros(d))
end

#matrix vector multiplication in TT format
function *(A::ttoperator{T},v::ttvector{T}) where T<:Number
    @assert(A.tto_dims==v.ttv_dims, DimensionMismatch)
    d = length(A.tto_dims)
    Y = Array{Array{T,3},1}(undef, d)
    A_rks = A.tto_rks #R_0, ..., R_d
    v_rks = v.ttv_rks #r_0, ..., r_d
    @threads for k in 1:d
		M = zeros(A.tto_dims[k], A_rks[k],v_rks[k], A_rks[k+1],v_rks[k+1])
		@tensor M[a,b,c,d,e] = A.tto_vec[k][a,z,b,d]*v.ttv_vec[k][z,c,e]
        Y[k] = reshape(M, A.tto_dims[k], A_rks[k]*v_rks[k], A_rks[k+1]*v_rks[k+1])
    end
    return ttvector{T}(Y,A.tto_dims,A.tto_rks.*v.ttv_rks,zeros(Integer,d))
end

#matrix matrix multiplication in TT format
function *(A::ttoperator{T},B::ttoperator{T}) where T<:Number
    @assert(A.tto_dims==B.tto_dims, DimensionMismatch)
    d = length(A.tto_dims)
    Y = Array{Array{T,4},1}(undef, d)
    A_rks = A.tto_rks #R_0, ..., R_d
    B_rks = B.tto_rks #r_0, ..., r_d
    @threads for k in 1:d
		M = zeros(A.tto_dims[k],A.tto_dims[k],A_rks[k],B_rks[k],A_rks[k+1],B_rks[k+1])
		@tensor M[a,b,c,d,e,f] = A.tto_vec[k][a,z,c,e]*B.tto_vec[k][z,b,d,f]
        Y[k] = reshape(M, A.tto_dims[k], A.tto_dims[k], A_rks[k]*B_rks[k], A_rks[k+1]*B_rks[k+1])
    end
    return ttoperator{T}(Y,A.tto_dims,A.tto_rks.*B.tto_rks,zeros(Integer,d))
end

#dot returns the dot product of two ttvector
function dot(A::ttvector{T},B::ttvector{T}) where T<:Number
    @assert A.ttv_dims==B.ttv_dims "TT dimensions are not compatible"
    d = length(A.ttv_dims)::Int
    A_rks = A.ttv_rks
    B_rks = B.ttv_rks
	out = zeros(T,maximum(A_rks),maximum(B_rks))
    out[1,1] = convert(T,1.0)
    @inbounds for k in 1:d
        M = @view(out[1:A_rks[k+1],1:B_rks[k+1]])
		@tensoropt((α,β), M[a,b] = A.ttv_vec[k][z,α,a]*B.ttv_vec[k][z,β,b]*out[1:A_rks[k],1:B_rks[k]][α,β]) #size R^A_{k} × R^B_{k}
    end
    return out[1,1]::T
end

"""
`dot_par(x_tt,y_tt)' returns the dot product of `x_tt` and `y_tt` in a parallelized algorithm
"""
function dot_par(A::ttvector{T},B::ttvector{T}) where T<:Number
    @assert A.ttv_dims==B.ttv_dims "TT dimensions are not compatible"
    d = length(A.ttv_dims)
    Y = Array{Array{T,2},1}(undef,d)
    A_rks = A.ttv_rks
    B_rks = B.ttv_rks
	C = zeros(T,maximum(A_rks.*B_rks))
    @threads for k in 1:d
		M = zeros(T,A_rks[k],B_rks[k],A_rks[k+1],B_rks[k+1])
		@tensor M[a,b,c,d] = A.ttv_vec[k][z,a,c]*B.ttv_vec[k][z,b,d] #size R^A_{k-1} ×  R^B_{k-1} × R^A_{k} × R^B_{k}
		Y[k] = reshape(M, A_rks[k]*B_rks[k], A_rks[k+1]*B_rks[k+1])
    end
    @inbounds C[1:length(Y[d])] = Y[d][:]
    for k in d-1:-1:1
        @inbounds C[1:size(Y[k],1)] = Y[k]*C[1:size(Y[k],2)]
    end
    return C[1]::T
end

function *(a::S,A::ttvector) where S<:Number
    i = findfirst(isequal(0),A.ttv_ot)
    T = typejoin(typeof(a),eltype(A))
    X = copy(A.ttv_vec)
    X[i] = a*X[i]
    return ttvector{T}(X,A.ttv_dims,A.ttv_rks,A.ttv_ot)
end

function *(a::S,A::ttoperator) where S<:Number
    i = findfirst(isequal(0),A.tto_ot)
    T = typejoin(typeof(a),eltype(A))
    X = copy(A.tto_vec)
    X[i] = a*X[i]
    return ttoperator{T}(X,A.tto_dims,A.tto_rks,A.tto_ot)
end


using LinearAlgebra
using SparseArrays

function expmv!(
        f::AbstractVecOrMat, t::StepRangeLen, A, b::AbstractVecOrMat,
        M = nothing, norm = LinearAlgebra.norm, opnorm = LinearAlgebra.opnorm;
        precision = :double, shift = DEFAULT_SHIFT, check_positive = false
    )

    n = size(A, 2)
    T = eltype(b)
    Tr = real(T)
    q = t.len - 1

    t0 = Tr(t.ref)
    tmax = Tr(t.ref + (t.len - 1.) * t.step)

    if shift == true && !hasmethod(tr, Tuple{typeof(A)})
        shift = false
        @warn "Shift set to false as $(typeof(A)) doesn't support tr"
    end

    force_estm = !hasmethod(opnorm, Tuple{typeof(A), typeof(1)})
    Anorm_dt = (tmax - t0) * normAm(A, 1; check_positive = check_positive);

    if (precision == :single || precision == :half && Anorm_dt > 85.496) || ( precision == :double && Anorm_dt > 63.152)
       force_estm = true;
    end

    if M == nothing
        (M, alpha, unA) = select_taylor_degree(A, size(b,2); precision=precision, shift=shift, force_estm=force_estm, check_positive=check_positive)
    end

    tol =
      if precision == :double
          Tr(2.0)^(-53)
      elseif precision == :single
          Tr(2.0)^(-24)
      elseif precision == :half
          Tr(2.0)^(-10)
      end

    X = zeros(T, n, q+1);
    (m_max, p) = size(M);
    U = diagm(0 => 1:m_max);

    _, s = degree_selector(tmax - t0, M, U, p)
    h = (tmax - t0)/q;

    expmv!(view(X,:,1), t0, A, b, M, norm, opnorm;
        precision=precision, shift=shift, check_positive=check_positive);

    mu = zero(Tr)
    if shift
        mu = tr(A)/n
        A = diagshift!(A, mu) #A = A - mu * I
    end

    d = max(1, Integer(floor(q/s)))
    j = Integer(floor(q/d))
    r = q - d * j
    m_opt, _ = degree_selector(d, M, U, p)
    dr = d

    z = X[:,1]
    K = zeros(T, n, m_opt+1)
    temp = similar(b)

    for i = 1:j+1
        if i > j
            dr = r
        end
        @views K[:,1] .= z
        m = 0
        for k = 1:dr
            f = copyto!(f, z)
            c1 = norm(z, Inf)

            p = 1
            for p = 1:m_opt # outer?
                if p > m
                    temp = mul!(temp, A, view(K,:,p)) # A*K[:,p]
                    @views K[:,p+1] .= (h/p) .* temp
                end

                @views temp .= (Tr(k)^p) .* K[:,p+1]
                f .+= temp
                c2 = norm(temp, Inf)
                if c1 + c2 <= tol * norm(f, Inf)
                    break
                end
                c1 = c2
            end

            m = max(m,p)
            if shift
                @views X[:, k + (i-1)*d + 1] .= exp(k*h*mu) .* f
            else
                @views X[:, k + (i-1)*d + 1] .= f
            end
        end

        if i <= j
            @views z .= X[:, i*d + 1]
        end
    end

    return X
end

function expmv(
        t::StepRangeLen, A, b::AbstractVecOrMat,
        M = nothing, norm = LinearAlgebra.norm, opnorm = LinearAlgebra.opnorm;
        precision = :double, shift = DEFAULT_SHIFT, check_positive = false
    )

    return expmv!(
        similar(b), t, A, b, M, norm, opnorm;
        precision = precision, shift = shift, check_positive = check_positive
    )
end


using PyPlot
using Triangulate

N = Float64[ 0  -2   0  2  2  1 -1 -2
             0  -2  -2 -2  0  2  2  0 ]

T = [ 2  3  3  1  1  1
      3  5  4  5  6  7
      8  8  5  6  7  8 ]

E = [ 2 3 4 5 6 7 8
      3 4 5 6 7 8 2 ]

tri = Triangulation(N, T, E, 7)

figure(1)
draw_triangles(tri)
enumerate_triangles(tri)
enumerate_vertices(tri)
#mark_first_local_vertex(tri, 0.3)
axis("off")
axis("equal")
savefig("bad_triangulation.pdf")


module Day18

export get_inputs, get_solution1, get_solution2

## Input getting
function get_inputs()
    test_input1 = test_input2 = readlines(IOBuffer(
        """
        1 + 2 * 3 + 4 * 5 + 6
        1 + (2 * 3) + (4 * (5 + 6))
        2 * 3 + (4 * 5)
        5 + (8 * 3 + 9 + 3 * 4 * 3)
        5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))
        ((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2
        """))
    test_output1 = 71 + 51 + 26 + 437 + 12240 + 13632
    test_output2 = 231 + 51 + 46 + 1445 + 669060 + 23340
    data = readlines(joinpath(@__DIR__, "input.txt"))
    return (; test_input1, test_input2, test_output1, test_output2, data)
end

## Solution functions
# Define new addition operator from same precedence pool as *
const ⊠ = +

# And new multiplication operator from same precence pool as +
const ⊞ = *

# Replace the desired operation rules and parse as an Expr object to be evaled
parse_expr(str, rules...) = Meta.parse(reduce(replace, rules, init=str))

# Create expressions and eval them
get_solution(data, rules...) = sum(eval.(parse_expr.(data, rules...)))

# Part 1
get_solution1(data) = get_solution(data, '+'=>'⊠')

# Part 2
get_solution2(data) = get_solution(data, '+'=>'⊠', '*'=>'⊞')

end


using QuasiNewtonMethods, StrideArrays, Aqua
using Test

struct Rosenbrock end
function QuasiNewtonMethods.logdensity(::Rosenbrock, θ)
    s = zero(eltype(θ))
    N = length(θ) >> 1
    @inbounds @simd for i ∈ 1:N
        s -= 100(θ[i+N] - θ[i]^2)^2 + (1 - θ[i])^2
    end
    if isodd(length(θ))
        @inbounds δ = 1 - θ[end]
        -muladd(δ, δ, -s)
    else
        s
    end
end

function QuasiNewtonMethods.∂logdensity!(∇, ::Rosenbrock, θ)
    s = zero(eltype(θ))
    N = length(θ) >> 1
    @inbounds @simd for i ∈ 1:N
        s -= 100(θ[i+N] - θ[i]^2)^2 + (1 - θ[i])^2
        ∇[i] = 400(θ[i+N] - θ[i]^2)*θ[i] + 2(1 - θ[i])
        ∇[i+N] = 200(θ[i]^2 - θ[i+N])
    end
    if isodd(length(θ))
        @inbounds δ = 1 - θ[end]
        s = -muladd(δ, δ, -s)
        @inbounds ∇[end] = 2δ
    end
    s
end

@testset "QuasiNewtonMethods.jl" begin
    Aqua.test_all(QuasiNewtonMethods)
    for n ∈ 2:24
        @show n
        state = QuasiNewtonMethods.BFGSState{n}(undef);
        x = @StrideArray randn(StaticInt(n));
        # 2nd order
        @test abs(optimize!(state, Rosenbrock(), x)) < 2eps()
        @show QuasiNewtonMethods.optimum(state) .- 1;
        @test all(x -> x ≈ 1, QuasiNewtonMethods.optimum(state))
        @test maximum(abs, QuasiNewtonMethods.gradient(state)) < 1e-8
        # 3rd order
        @test abs(optimize!(state, Rosenbrock(), x, QuasiNewtonMethods.BackTracking{3}())) < eps()
        @show QuasiNewtonMethods.optimum(state) .- 1;
        @test all(x -> x ≈ 1, QuasiNewtonMethods.optimum(state))
        @test maximum(abs, QuasiNewtonMethods.gradient(state)) < 1e-8
    end
end


#
# spg2.jl --
#
# Implements Spectral Projected Gradient Method (Version 2: "continuous
# projected gradient direction") to find the local minimizers of a given
# function with convex constraints, described in:
#
# [1] E. G. Birgin, J. M. Martinez, and M. Raydan, "Nonmonotone spectral
#     projected gradient methods on convex sets", SIAM Journal on Optimization
#     10, pp. 1196-1211 (2000).
#
# [2] E. G. Birgin, J. M. Martinez, and M. Raydan, "SPG: software for
#     convex-constrained optimization", ACM Transactions on Mathematical
#     Software (TOMS) 27, pp. 340-349 (2001).
#
# ----------------------------------------------------------------------------
#
# This file is part of OptimPack.jl which is licensed under the MIT
# "Expat" License:
#
# Copyright (C) 2014: Éric Thiébaut.
#
# ----------------------------------------------------------------------------

type SpgResult{T,N}
    x::DenseArray{T,N}
    xbest::DenseArray{T,N}
    f::Float64
    fbest::Float64
    pgtwon::Float64
    pginfn::Float64
    fcnt::Int
    pcnt::Int
    iter::Int
    status::Symbol
    SpgResult(x::DenseArray{T,N}, xbest::DenseArray{T,N}) = new(x, xbest, 0.0, 0.0, 0.0, 0.0, 0, 0, 0, :WORK_IN_PROGRESS)
end

function spg2{T,N}(fg!::Function, prj!::Function, x0::DenseArray{T,N}, m::Integer;
                   maxit::Integer=-1, maxfc::Integer=-1,
                   eps1::Real=1e-6, eps2::Real=1e-6,
                   eta::Real=1.0, ftol::Real=1e-4,
                   lmin::Real=1e-30, lmax::Real=1e+30,
                   amin::Real=0.1, amax::Real=0.9,
                   printer=nothing, verb::Bool=false)

    # Allocate workspace.
    x0 = copy(x0)
    dims = size(x0)
    space = DenseVectorSpace(T, dims)
    d = Array(T, dims)
    g = Array(T, dims)
    x = Array(T, dims)
    g = Array(T, dims)
    g0 = Array(T, dims)
    xbest = Array(T, dims)
    _x = wrap(space, x)
    _g = wrap(space, g)
    _d = wrap(space, d)
    _x0 = wrap(space, x0)
    _g0 = wrap(space, g0)
    _xbest = wrap(space, xbest)

    # Initialization.
    ws = SpgResult(x, xbest)
    local sty, sts, f0, f
    if m > 1
        lastfv = Array(T, m)
        fill!(lastfv, inf(T))
    else
        lastfv = nothing
    end

    # Project initial guess.
    prj!(x, x0)
    ws.pcnt += 1

    # Evaluate function and gradient.
    f = fg!(x, g)
    ws.fcnt += 1

    # Initialize best solution and best function value.
    ws.fbest = f
    copy!(_xbest, _x)

    # Main loop.
    while true

        # Compute continuous projected gradient (and its norms).
        axpby!(_d, 1.0, _x, eta, _g)
        prj!(_d, _d)
        axpby!(_d, 1.0/eta, _x, -1.0/eta, _d)
        ws.pcnt += 1
        ws.pgtwon = norm2(_d)
        ws.pginfn = norminf(_d)

        # Print iteration information
        if printer <: Function
            printer(ws)
        end
        if verb
            @printf("ITER = %-5d  EVAL = %-5d  PROJ = %-5d  F(%s) =%24.17e  ||PG||oo = %17.10e\n",
                    ws.iter, ws.fcnt, ws.pcnt, (f <= ws.fbest ? "+" : "-"), f, ws.pginfn)
        end

        # Test stopping criteria.
        if ws.pginfn <= eps1
            # Gradient infinite-norm stopping criterion satisfied, stop.
            ws.status = :CONVERGENCE_WITH_INFNORM
            return ws
        end
        if ws.pgtwon <= eps2
            # Gradient 2-norm stopping criterion satisfied, stop.
            ws.status = :CONVERGENCE_WITH_TWONORM
            return ws
        end
        if maxit >= 0 && ws.iter >= maxit
            # Maximum number of iterations exceeded, stop.
            ws.status = :TOO_MANY_ITERATIONS
            return ws
        end
        if maxfc >= 0 && ws.fcnt >= maxfc
            # Maximum number of function evaluations exceeded, stop.
            ws.status = :TOO_MANY_EVALUATIONS
            return ws
        end

        # Store function value for the nonmonotone line search and
        # find maximum function value since m last calls.
        if m > 1
            lastfv[(ws.iter%m) + 1] = f
            fmax = maximum(lastfv)
        else
            fmax = f
        end

        # Compute spectral steplength.
        if iter == 0
            # Initial steplength.
            lambda = min(lmax, max(lmin, 1.0/ws.pginfn))
        else
            _s = _x0 # alias
            _y = _g0 # alias
            axpby!(_s, 1.0, _x, -1.0, _x0)
            axpby!(_y, 1.0, _g, -1.0, _g0)
            sty = dot(_s, _y)
            if sty > 0.0
                # Safeguarded Barzilai & Borwein spectral steplength.
                sts = dot(_s, _s)
                lambda = min(lmax, max(lmin, sts/sty))
            else
                lambda = lmax
            end
        end

        # Save current point.
        copy!(_x0, _x)
        copy!(_g0, _g)
        f0 = f

        # Compute the spectral projected gradient direction and <G,D>
        axpby!(_x, 1.0, _x, -lambda, _g)
        prj(_x, _x)
        ws.pcnt += 1
        axpby!(_d, 1.0, _x, -1.0, _x0)
        delta = dot(_g0, _d)

        # Nonmonotone line search.
        stp = 1.0 # Step length for first trial.
        while true
            # Evaluate function and gradient at trial point.
            f = fg!(x, g)
            ws.fcnt += 1

            # Compare the new function value against the best function
            # value and, if smaller, update the best function value and the
            # corresponding best point.
            if f < ws.fbest
                ws.fbest = f
                copy!(_xbest, _x)
            end

            # Test stopping criteria.
            if f <= fmax + stp*ftol*delta
                # Nonmonotone Armijo-like stopping criterion satisfied, stop.
                break
            end
            if maxfc >= 0 && ws.fcnt >= maxfc
                # Maximum number of function evaluations exceeded, stop.
                ws.status = :TOO_MANY_EVALUATIONS
                return ws
            end

            # Safeguarded quadratic interpolation.
            q = -delta*(stp*stp)
            r = (f - f0 - stp*delta)*2.0
            if r > 0.0 && amin*r <= q <= amax*stp*r
                stp = q/r
            else
                stp /= 2.0
            end

            # Compute trial point.
            axpby!(_x, 1.0, _x0, stp, _d)
        end

        if ws.status != :WORK_IN_PROGRESS
            # The number of function evaluations was exceeded inside the line
            # search.
            return ws
        end

        # Proceed with next iteration.
        ws.iter += 1

    end
end


include("bkfiltertest.jl")

x1,p1 = augment(s,1e-10*eye(2))
@test x1 == [1.0,0.0,0.0,0.0]
@test p1 == [4*eye(2,2) zeros(2,2);
             zeros(2,2) 1e-10*eye(2)]

function fa(x::Vector)
    x1 = zeros(2)
    x1[1] = x[1] + 0.01*x[2] + x[3]
    x1[2] = x[2] - 0.01*x[1] + x[4]
    x1
end

fam = AugmentedUnscentedModel(fa,1e-10*eye(2))

function ha(x)
    x1 = zeros(1)
    x1[1] = x[1] + x[3]
end

zam = AugmentedUnscentedObservationModel(ha,[0.1]')

kaf = AugmentedUnscentedKalmanFilter(s,fam,zam)

kaf1 = predict(kaf)
kaf2 = update(kaf1,y)

@test_approx_eq kaf2.x.x kf2.x.x
@test_approx_eq kaf2.x.p kf2.x.p


#= basic support for sparse matrices, more Magma style than Julia
  a sparse matrix is a collection (Array) of sparse rows (SRow)
  A SRow is two Arrays, one (pos) containing the columns, values
  contains the values.
  A[i,j] = A.rows[i].values[A.rows[i].pos[j]]
  to be formal

  The upper_triangular stuff performs very elementary transformations
  until the matrix becomes dense. At this point, the dense bit is extraced and
  converted to an fmpz_mat which is then hnf'ed.

  Missing:
   full HNF, Howell, modular and not
   better elemination strategy

  TODO: sort out the various upper_triangular implementations.
  One (+one with trafo) is probably enough
=#

import Base.push!, Base.max, Nemo.nbits, Base.Array,
       Base.hcat,
       Base.vcat, Base.max, Base.min

export upper_triangular, vcat!, show, sub, SMat, SRow, random_SMatSLP,
       fmpz_mat, rows, cols, copy, push!, mul, mul!, toNemo, sparse,
       valence_mc, swap_rows!, elementary_divisors,
       rand_row, hcat, hcat!, vcat, vcat!, mod!, mod_sym!

include("Sparse/Matrix.jl")
include("Sparse/Row.jl")
include("Sparse/Module.jl")
include("Sparse/Trafo.jl")
include("Sparse/HNF.jl")
include("Sparse/Solve.jl")
include("Sparse/UpperTriangular.jl")


struct WeekDayEffect1{Tα} <: ContinuousMultivariateDistribution
	α::Tα
end
WeekDayEffect1() = WeekDayEffect1(ones(Int64, 7))

Distributions.rand(rng::AbstractRNG, d::WeekDayEffect1{Tα}) where {Tα} = begin
	weekdayeffect_simplex = rand( Dirichlet(d.α) )
	weekdayeffect_simplex .* length(d.α)
end

Distributions.logpdf(d::WeekDayEffect1{Tα}, x::AbstractVector{T}) where {Tα,T} =
    logpdf( Dirichlet( d.α ), x ./ length(d.α) )

Bijectors.bijector(d::WeekDayEffect1) =
	Bijectors.bijector(Dirichlet(d.α)) ∘ Bijectors.Scale(1/7, dim=Val(1))

#Bijectors.bijector(Dirichlet( d.α ))
@model function test(y, d, ::Type{TV} = Vector{Float64}) where TV
	dwe ~ WeekDayEffect1()
	e ~ Normal(4, 0.1)
	ŷ = TV(undef, length(y))
	for i in eachindex(ŷ)
		ŷ = e + dwe[d[i]]
	end
	ŷ ~ MvNormal(y, 0.1)
	return ŷ
end

m = 10
d = repeat(1:7, outer=m)
n = length(d)
y = repeat(1:7, outer=m) + randn(n) * 0.1

m = test(y, d)
plot!(m())
c = sample(m, NUTS(1000, 0.8), 2000; progress=true)

gq = generated_quantities(m, c[1001:end])|>vec
ŷ = mean(gq)
plot(y)
plot!(ŷ)

plot(group(c, "dwe"))

logpdf(Dirichlet(ones(7)), ones(7) / 7.)
logpdf(WeekDayEffect1(ones(7)), ones(7))

bs = Bijectors.Scale(1/7, dim=Val(1))
bs isa Bijectors.Bijector
x = 1.
y = bs(x)
inv(bs)(y) ≈ x

bd = bijector(Dirichlet(ones(7)))
bd isa Bijectors.Bijector
x = ones(7)/7
y = bd(x)
inv(bd)(y) ≈ x

b = bd ∘ bs
x = ones(7)
y = b(x)
inv(b)(y) ≈ x


module GridapTypesTests

using Test
using Gridap.Helpers

import Gridap.Helpers: operate

struct Foo{A} <: GridapType
  bar::A
end

object = Foo(nothing)

@test sprint(show,"text/plain",object) == "$(nameof(typeof(object)))()"

function operate(op,a::Foo)
  Foo(op(a.bar))
end

function operate(op,a::Foo,b::Foo)
  Foo(op(a.bar,b.bar))
end

function operate(op,a::Number,b::Foo)
  operate(op,Foo(a),b)
end

function operate(op,a::Foo,b::Number)
  operate(op,a,Foo(b))
end

a = Foo(1)

b = -a
@test b.bar == -1

c = a + b

@test c.bar == 0

c = a - 2
@test c.bar == -1

c = 2 - a
@test c.bar == 1

using Gridap.TensorValues

c = inner(2,a)
@test c.bar == 2

c = a / b
@test c.bar == -1

end # module


using ITensors

let
  N = 100
  cutoff = 1E-8
  tau = 0.1
  ttotal = 5.0

  # Make an array of 'site' indices
  s = siteinds("S=1/2", N; conserve_qns=true)

  # Make gates (1,2),(2,3),(3,4),...
  gates = ITensor[]
  for j in 1:(N - 1)
    s1 = s[j]
    s2 = s[j + 1]
    hj =
      op("Sz", s1) * op("Sz", s2) +
      1 / 2 * op("S+", s1) * op("S-", s2) +
      1 / 2 * op("S-", s1) * op("S+", s2)
    Gj = exp(-1.0im * tau / 2 * hj)
    push!(gates, Gj)
  end
  # Include gates in reverse order too
  # (N,N-1),(N-1,N-2),...
  append!(gates, reverse(gates))

  # Initialize psi to be a product state (alternating up and down)
  psi = productMPS(s, n -> isodd(n) ? "Up" : "Dn")

  c = div(N, 2) # center site

  # Compute and print <Sz> at each time step
  # then apply the gates to go to the next time
  for t in 0.0:tau:ttotal
    Sz = expect(psi, "Sz"; sites=c)
    println("$t $Sz")

    t ≈ ttotal && break

    psi = apply(gates, psi; cutoff)
    normalize!(psi)
  end

  return nothing
end


"""
    simple_normalized_residuals(data::Dict, se_sol::Dict, R::Matrix)
It normalizes the residuals only based on their standard deviation, no sensitivity matrix involved.
R^2 replaces Ω, where Ω = S ⋅ R
Avoids all the matrix calculations but is a "simplified" method
"""
function simple_normalized_residuals(data::Dict, se_sol::Dict, R::Matrix)
    normalized_residuals(data, se_sol, R.^2)
end
"""
Adds the normalized residuals to the solution dictionary and returns the largest normalized residual,
its index (i.e., the measurement it refers to), and whether it exceeds the given threshold `t` or not.
"""
function normalized_residuals(data::Dict, se_sol::Dict, Ω::Matrix; t::Float64=3.0)
    lnr = ("0", 0)
    count = 1
    for (m, meas) in data["meas"]
        if !haskey(se_sol["solution"][string(meas["cmp"])], string(meas["cmp_id"]))
            r = [0.0, 0.0, 0.0]
            se_sol["solution"]["meas"][m]["r"] = [0.0, 0.0, 0.0]
        else
            h_x = se_sol["solution"][string(meas["cmp"])][string(meas["cmp_id"])][string(meas["var"])]
            z = _DST.mean.(meas["dst"])
            r = abs.(h_x-z)
            se_sol["solution"]["meas"][m]["r"] = r
        end
        se_sol["solution"]["meas"][m]["nr"] = [r[i]/sqrt(abs(Ω[count+i-1, count+i-1])) for i in 1:length(meas["dst"])]
        for i in 1:length(meas["dst"])
            if  se_sol["solution"]["meas"][m]["nr"][i] > last(lnr) lnr = (m,  se_sol["solution"]["meas"][m]["nr"][i]) end
        end
        count+=length(meas["dst"])
    end
    return lnr, last(lnr) > t
end
"""
Returns the Measurement Jacobian H
"""
function build_H_matrix(functions::Vector, state::Array)::Matrix{Float64}
    H = Matrix{Float64}(undef, length(functions), length(state))
    for row in 1:length(functions)
        H[row,:] = ForwardDiff.gradient(functions[row], state)
    end
    return H
end
"""
Returns the Gain Matrix G
"""
build_G_matrix(H::Matrix, R::Matrix)::Matrix{Float64} = transpose(H)*inv(R)*H
"""
Returns the Measurement Error Covariance Matrix R
"""
function build_R_matrix(data::Dict)::Matrix{Float64}
    meas_row_order = [m for (m, _) in data["meas"]]
    R_entries = vcat([_DST.std.(data["meas"][mid]["dst"])[1:length(data["meas"][mid]["dst"])] for mid in meas_row_order]...)
    return LinearAlgebra.diagm(R_entries.^2)
end
"""
Returns the Residual Covariance Matrix Ω, where:
Ω = R - H*G^(-1)*H^T
"""
build_omega_matrix(R::Matrix{Float64}, H::Matrix{Float64}, G::Matrix{Float64}) = R - H*inv(G)*transpose(H)
"""
Returns the Residual Covariance Matrix Ω, where:
Ωᵢᵢ = Rᵢᵢ⋅Sᵢᵢ
"""
build_omega_matrix(S::Matrix{Float64}, R::Matrix{Float64}) = diagm((sqrt.(abs.(diag(S).*diag(R)))))
"""
Returns the sensitivity matrix S starting from the hat matrix K
"""
build_S_matrix(K::Matrix{Float64}) = Matrix{Int64}(I, size(K)) - K
"""
Returns the hat matrix K
"""
build_K_matrix(H::Matrix{Float64}, G::Matrix{Float64}, R::Matrix{Float64}) = H*inv(G)*transpose(H)*inv(R)


using Plots, Statistics, StatsBase, LaTeXStrings, JLD, ProgressBars
#mutable struct -> instances of the composite types can be modified.
mutable struct MDSystem{FT<:AbstractFloat}
    L::FT
    step::FT
    Eᵤ::FT
    Eₖ::FT
    T::FT
    P::FT
    N::Integer
    F::Matrix{FT}
    r::Matrix{FT}
    v::Matrix{FT}
    function MDSystem(N::Integer, T₀::FT, step::FT, L::FT) where {FT<:AbstractFloat}
        r = rand(FT, 2, N) * L
        v = rand(FT, 2, N) .- 0.5
        v *= sqrt((T₀ ÷ mean(v .^ 2)))
        F = Matrix{FT}(undef, 2, N)
        return new{FT}(L, step, FT(NaN), FT(NaN), FT(NaN), FT(NaN), N, F, r, v)
    end
end

function Initialize(N::Integer,L::FT, T₀::FT, step::FT) where {FT<:AbstractFloat}
    md = MDSystem(N,  T₀, step, L)
    alignleft(md)
    for i in 1:md.N
        Force_cal(md, i)
    end
    VelVerlet(md)
    Boundary_condition(md)
    Potential_cal(md)
    Kinetic_cal(md)
    Temprature_cal(md)
    Pressure_cal(md)
    return md
end

function Sim_process(md :: MDSystem)
    VelVerlet(md)
    Boundary_condition(md)
    Potential_cal(md)
    Kinetic_cal(md)
    Temprature_cal(md)
    Pressure_cal(md)
end

function alignleft(md::MDSystem)
    xcount = ceil(Integer, (2 * md.N)^(1 / 2) / 2)
    ycount = 2 * xcount

    xs = collect(range(0.1 * md.L / 2, 0.9 * md.L / 2, length = xcount))
    md.r[1, :] .= repeat(xs, outer = ycount)[1:md.N]

    ys = collect(range(0.1 * md.L, 0.9 * md.L, length = ycount))
    md.r[2, :] .= repeat(ys, inner = xcount, outer = 1)[1:md.N]
end

function Force_cal(md :: MDSystem, i :: Integer)
    md.F[:, i] *= 0
    for j in setdiff(1:md.N, i)
        rp = md.r[:, j]
        Δx = md.r[1, i] - rp[1]
        Δy = md.r[2, i] - rp[2]
        if Δx < -md.L / 2
            rp[1] -= md.L
        elseif Δx > md.L / 2
            rp[1] += md.L
        end
        if Δy < -md.L / 2
            rp[2] -= md.L
        elseif Δy > md.L / 2
            rp[2] += md.L
        end
        Δr = sqrt(sum((md.r[:, i] - rp) .^ 2))
        md.F[:, i] += 48 * (((Δr)^-14) - 0.5 * ((Δr)^-8)) * (md.r[:, i] - rp)
    end
end

function VelVerlet(md :: MDSystem)
    for i in 1:md.N
        a₁ = md.F[:, i]
        md.r[:, i] += md.step * md.v[:, i] + 1/2 * md.step  * md.step * a₁
        Force_cal(md, i)
        a₂ = md.F[:, i]
        md.v[:, i] += md.step * (a₁ + a₂) / 2
    end
end

function Boundary_condition(md :: MDSystem)
    md.r = (md.r .+ md.L) .% md.L
end

function Potential_cal(md :: MDSystem)
    ΣΔr = 0.0
    for i in 1:(md.N - 1)
        for j in (i+1):md.N
            rp = md.r[:, j]
            Δx = md.r[1, i] - rp[1]
            Δy = md.r[2, i] - rp[2]
            if Δx < -md.L / 2
                rp[1] -= md.L
            elseif Δx > md.L / 2
                rp[1] += md.L
            end
            if Δy < -md.L / 2
                rp[2] -= md.L
            elseif Δy > md.L / 2
                rp[2] += md.L
            end
            Δr = sqrt(sum((md.r[:, i] - rp) .^ 2))
            ΣΔr += ((Δr) ^ (-12)) - ((Δr) ^ (-6))
        end
    end
    md.Eᵤ = 4 * ΣΔr
end

function Kinetic_cal(md :: MDSystem)
    md.Eₖ = sum(md.v .^2) / 2 #kinetic Energy per molecule
end

function Temprature_cal(md :: MDSystem)
    md.T = md.Eₖ / md.N
end

function Pressure_cal(md :: MDSystem)
    Σrterm = 0.0
    Σvᵢ² = sum(md.v .^2)
    for i in 1:(md.N-1)
        for j in (i+1):md.N
            rp = md.r[:, j]
            Δx = md.r[1, i] - rp[1]
            Δy = md.r[2, i] - rp[2]
            if Δx < -md.L / 2
                rp[1] -= md.L
            elseif Δx > md.L / 2
                rp[1] += md.L
            end
            if Δy < -md.L / 2
                rp[2] -= md.L
            elseif Δy > md.L / 2
                rp[2] += md.L
            end
            Δr = sqrt(sum(md.r[:, i] - rp) .^ 2)
            Σrterm += ((Δr) ^ (-12)) - (1/2) * ((Δr) ^ (-6))
        end
    end
    md.P = (Σvᵢ² + 48 * Σrterm)/(2 * md.L * md.L)
end

N = 100; T₀ = 1.0; h = 0.001; L = 50.0; step_num = 20000
md = Initialize(N, L, T₀ , h)

Poses = zeros(step_num, 2, N)
Velocities = zeros(step_num, 2, N)
Potentials = zeros(step_num)
Kinetics = zeros(step_num)
Temps = zeros(step_num)
pressures = zeros(step_num)

for i in ProgressBar(1:step_num)
    Sim_process(md)
    Poses[i, :, :] = md.r
    Velocities[i, :, :] = md.v
    Potentials[i] = md.Eᵤ
    Kinetics[i] = md.Eₖ
    Temps[i] = md.T
    pressures[i] = md.P
end
Poses
save("../../MD/SingleData.jld",
    "Poses", Poses, "Velocities", Velocities, "Potentials", Potentials,
    "Kinetics", Kinetics, "Temps", Temps, "Pressures", pressures)

plot(collect(0:h:(h*step_num))[1:20000], [count(<=(md.L / 2), Poses[i, 1, :]) / md.N for i ∈ 1:step_num][1:20000],
    ylabel=L"N_l_e_f_t/N", xlabel=L"t/\tau", dpi=400, title=L"Molecules'\ Densities\ at\ the\ left\ half", legend=false)
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\MD\\Figs\\LeftHalfDensity.png")

plot(collect(0:h:(h*step_num))[1:2000],Kinetics[1:2000], label = L"E_k")
plot!(collect(0:h:(h*step_num))[1:2000],Potentials[1:2000], label = L"E_U")
plot!(collect(0:h:(h*step_num))[1:2000],Potentials[1:2000]+Kinetics[1:2000], label = L"E_t_o_t_a_l"
    , title = L"Enery\ of\ the\ System",dpi = 400)
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\MD\\Figs\\Energies.png")

plot(collect(0:h:(h*step_num))[1:20000],120*Temps[1:20000], label = L"Temprature", title=L"Temprature\ of\ the\ System")
plot!([0.0, 20.0], 120 * [mean(Temps[10000:end]), mean(Temps[10000:end])],
    label = L"\langle T_{eq} \rangle = %$(round(120 * mean(Temps[10000:end]), digits = 3))")
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\MD\\Figs\\Temprature.png")


using ITensors
using ITensorGPU

import ITensors: space, state, op

space(::SiteType"Qubit") = 2
state(::SiteType"Qubit", ::StateName"0") = 1
state(::SiteType"Qubit", ::StateName"1") = 2

op_matrix(s::String) = op_matrix(OpName(s))

op_matrix(::OpName"Id") = [
  1 0
  0 1
]

op_matrix(::OpName"I") = op_matrix("Id")

op_matrix(::OpName"X") = [
  0 1
  1 0
]

function op_matrix(on::OpName, s::Index...; kwargs...)
  rs = reverse(s)
  return itensor(op_matrix(on; kwargs...), prime.(rs)..., dag.(rs)...)
end

op_matrix(gn::String, s::Index...; kwargs...) = op_matrix(OpName(gn), s...; kwargs...)

function op_matrix(gn::String, s::Vector{<:Index}, ns::Int...; kwargs...)
  return op_matrix(OpName(gn), s[[ns...]]...; kwargs...)
end

op(gn::OpName, ::SiteType"Qubit", s::Index...; kwargs...) = op_matrix(gn, s...; kwargs...)

N = 10
s = siteinds("Qubit", N)
X = cu.(ops(s, [("X", n) for n in 1:N]))

initstate = fill("0", N)

ψ0 = productCuMPS(s, initstate)

gates = [X[n] for n in 1:2:N]
ψ = apply(gates, ψ0)


import NbodyGradient: init_nbody

@testset "init_nbody" begin

    elements = "elements.txt"
    H = [3,1,1]
    t0 = 7257.93115525 - 7300.0
    init = ElementsIC(t0, H, elements)
    init.elements[:,3] .-= 7300.0
    x, v, jac_init = init_nbody(init)
    elements_big = big.(init.elements)
    t0big = big(t0)
    init_big = ElementsIC(t0big, H, elements_big)
    xbig, vbig, jac_init_big = init_nbody(init_big)

    # elements = readdlm("elements.txt",',')

    # elements[2:3,1] /= 100.0
    # n_body = 4
    n_body = 3

    # jac_init     = zeros(Float64,7*n_body,7*n_body)
    # jac_init_big = zeros(BigFloat,7*n_body,7*n_body)
    jac_init_num = zeros(BigFloat, 7 * n_body, 7 * n_body)
    # x,v = init_nbody(elements,t0,n_body,jac_init)
    # elements_big = big.(elements)
    # t0big = big(t0)
    # xbig,vbig = init_nbody(elements_big,t0big,n_body,jac_init_big)
    elements0 = copy(init.elements)
    # dq = big.([1e-10,1e-5,1e-6,1e-6,1e-6,1e-5,1e-5])
    # dq = big.([1e-10,1e-8,1e-8,1e-8,1e-8,1e-8,1e-8])
    dq = big.([1e-12,1e-10,1e-10,1e-10,1e-10,1e-10,1e-10])

    # dq = big.([1e-15,1e-15,1e-15,1e-15,1e-15,1e-15,1e-15])
    # t0big = big(t0)
    # Now, compute derivatives numerically:
    for j = 1:n_body
        for k = 1:7
            elementsbig = big.(elements0)
            dq0 = dq[k]; if j == 1 && k == 1 ; dq0 = big(1e-15); end
            elementsbig[j,k] += dq0
            initp = ElementsIC(t0big, H, elementsbig)
            xp, vp, _ = init_nbody(initp)
            elementsbig[j,k] -= 2dq0
            initm = ElementsIC(t0big, H, elementsbig)
            xm, vm, _ = init_nbody(initm)
            for l in 1:n_body, p in 1:3
                i1 = (l - 1) * 7 + p
                if k == 1
                    j1 = j * 7
                else
                    j1 = (j - 1) * 7 + k - 1
                end
                jac_init_num[i1,  j1] = (xp[p,l] - xm[p,l]) / dq0 * .5
                jac1 = jac_init[i1,j1]; jac2 = jac_init_num[i1,j1]
                # if abs(jac1-jac2) > 1e-4*abs(jac1+jac2) && abs(jac1+jac2) > 1e-14
                #  println(l," ",p," ",j," ",k," ",jac_init_num[i1,j1]," ",jac_init[i1,j1]," ",jac_init_num[i1,j1]/jac_init[i1,j1])
                # end
                jac_init_num[i1 + 3,j1] = (vp[p,l] - vm[p,l]) / dq0 * .5
                jac1 = jac_init[i1 + 3,j1]; jac2 = jac_init_num[i1 + 3,j1]
                # if abs(jac1-jac2) > 1e-4*abs(jac1+jac2) && abs(jac1+jac2) > 1e-14
                # println(l," ",p+3," ",j," ",k," ",jac_init_num[i1+3,j1]," ",jac_init[i1+3,j1]," ",jac_init_num[i1+3,j1]/jac_init[i1+3,j1])
                # end
            end
        end
        jac_init_num[j * 7,j * 7] = 1.0
    end

    # println("Maximum jac_init-jac_init_num: ",maximum(abs.(jac_init-jac_init_num)))
    # println("Maximum jac_init-jac_init_num: ",maximum(abs.(asinh.(jac_init)-asinh.(jac_init_num))))
    # println("Maximum jac_init-jac_init_big: ",maximum(abs.(asinh.(jac_init)-asinh.(jac_init_big))))
    # println("Maximum jac_init/jac_init_big-1: ",maximum(abs.(jac_init[jac_init_big .!= 0.0]./jac_init_big[jac_init_big .!= 0.0].-1)))
    # println("Maximum jac_init_big-jac_init_num: ",maximum(abs.(asinh.(jac_init_num)-asinh.(jac_init_big))))
    # println(convert(Array{Float64,2},jac_init_big-jac_init_num))
    # @test isapprox(jac_init_num,jac_init)
    @test isapprox(jac_init_num, jac_init;norm=maxabs)
    @test isapprox(jac_init, convert(Array{Float64,2}, jac_init_big);norm=maxabs)
    # end
    for i in 1:21, j in 1:21
        if jac_init_big[i,j] != 0.0
            dij = abs(jac_init[i,j] / jac_init_big[i,j] - 1)
            if dij > 1e-12
                jac_max = dij
                println(i, " ", j, " ", convert(Float64, jac_max), " ", jac_init[i,j], " ", convert(Float64, jac_init_big[i,j]), " ", convert(Float64, jac_init[i,j] - jac_init_big[i,j]))
            end
        end
    end
end


using ProgressMeter

function main()
    out_dir="result"
    PARAMDAT="PARAMDAT.json"
    xmax,ymax,nodes,vecAx,vecAy = read_allgrid()
    nt,dt,init_rho,init_u,init_v,init_p,specific_heat_ratio,bdcon,every_outnum,out_file_front,out_ext,restart_file,restartnum,time_integ,init_small,norm_ok,nt_lusgs = input_para(PARAMDAT)

    Qbase,cellxmax,cellymax,restartnum = set_initQbase(xmax,ymax,restart_file,init_rho,init_u,init_v,init_p,specific_heat_ratio,out_file_front,out_ext,out_dir,restartnum)

    # init Delta_Qcon_hat
    Delta_Qcon_hat = zeros(cellxmax,cellymax,4)

    # main loop
    println("threads num ")
    println(Threads.nthreads())

    prog = Progress(nt,1)
    for k in 1:nt
        next!(prog)

        volume = set_volume(nodes,cellxmax,cellymax)

        Qbase = set_boundary(Qbase,cellxmax,cellymax,vecAx,vecAy,bdcon)
        Qcon = base_to_conservative(Qbase,cellxmax,cellymax,specific_heat_ratio)

        Qcon_hat = setup_Qcon_hat(Qcon,cellxmax,cellymax,volume)

        cell_Ahat_plas,cell_Ahat_minus,cell_Bhat_plas,cell_Bhat_minus = cal_jacobi(Qbase,Qcon,cellxmax,cellymax,specific_heat_ratio,vecAx,vecAy)
        cell_E_hat_plas,cell_E_hat_minus,cell_F_hat_plas,cell_F_hat_minus = setup_cell_flux_hat(Qcon,cellxmax,cellymax,cell_Ahat_plas,cell_Ahat_minus,cell_Bhat_plas,cell_Bhat_minus)

        E_hat,F_hat = FVS(cell_E_hat_plas,cell_E_hat_minus,cell_F_hat_plas,cell_F_hat_minus,cellxmax,cellymax)

        RHS = setup_RHS(cellxmax,cellymax,E_hat,F_hat,Qcon_hat)

        evalnum = k+restartnum
        if time_integ == "1"
            # exlicit scheme
            Qcon_hat = time_integration_explicit(dt,Qcon_hat,RHS,cellxmax,cellymax)
        elseif time_integ == "2"
            # implicit scheme
            A_adv_hat_p, A_adv_hat_m, B_adv_hat_p, B_adv_hat_m, A_beta_sig, B_beta_sig = one_wave(Qbase,Qcon,cellxmax,cellymax,vecAx,vecAy,specific_heat_ratio,volume)
            Lx,Ly,Ux,Uy,D = set_LDU(dt,Qcon_hat,A_adv_hat_p, A_adv_hat_m, B_adv_hat_p, B_adv_hat_m, A_beta_sig, B_beta_sig,RHS,cellxmax,cellymax,volume)

            Delta_Qcon_hat_temp = zeros(cellxmax,cellymax,4)

            for lusgs_t in 1:nt_lusgs
                nt_temp = 100
                norm2 = 0.0

                for i in 1:nt_temp
                    #println(Delta_Qcon_hat[2,2,:])
                    Delta_Qcon_hat = lusgs(dt,RHS,cellxmax,cellymax,volume,Lx,Ly,Ux,Uy,D,Delta_Qcon_hat)
                    #println(Delta_Qcon_hat[2,2,:])

                    res = set_res(Delta_Qcon_hat,Delta_Qcon_hat_temp,cellxmax,cellymax)
                    norm2 = check_converge(res,RHS,cellxmax,cellymax,init_small)

                    Delta_Qcon_hat_temp = deepcopy(Delta_Qcon_hat)

                    #throw(UndefVarError(:x))

                    #println(norm2)
                    #println(Delta_Qcon_hat[2,2,:])
                    #println(dt*RHS[2,2,:])
                    println("\n ---------------------------------- \n")
                end

                if norm2[1] < norm_ok && norm2[4] < norm_ok
                    println("\n ---------------------------------- \n")
                    println("nt : "*string(round(evalnum))*"     lusgs_t : "*string(lusgs_t))
                    println("density res:"*string(norm2[1]) * "  energy res:"*string(norm2[4]))
                    println("\n ---------------------------------- \n")
                    break
                end

            end
            #Delta_Qcon_hat = lusgs(dt,RHS,cellxmax,cellymax,volume,Lx,Ly,Ux,Uy,D,Delta_Qcon_hat)
            #Delta_Qcon_hat = lusgs(dt,RHS,cellxmax,cellymax,volume,Lx,Ly,Ux,Uy,D,Delta_Qcon_hat)

            # lusgs = update
            for i in 2:cellxmax-1
                for j in 2:cellymax-1
                    for l in 1:4
                        Qcon_hat[i,j,l] = Qcon_hat[i,j,l] + Delta_Qcon_hat[i,j,l]
                    end
                end
            end

            #println(dt*RHS[2,2,:])

            # 内部反復
            #=
            nt_inner = 10
            Qcon_hat_n = copy(Qcon_hat)
            Qcon_hat_m1 = copy(Qcon_hat)
            Delta_Qcon_hat_m = copy(Delta_Qcon_hat)

            for inner_t in 1:nt_inner
                Delta_Qcon_hat_temp = zeros(cellxmax,cellymax,4)

                for lusgs_t in 1:nt_lusgs
                    Qcon_m1 = Qhat_to_Q(Qcon_hat_m1,cellxmax,cellymax,volume)
                    Qbase_m1 = conservative_to_base(Qcon_m1,cellxmax,cellymax,specific_heat_ratio)

                    cell_Ahat_plas,cell_Ahat_minus,cell_Bhat_plas,cell_Bhat_minus = cal_jacobi(Qbase_m1,Qcon_m1,cellxmax,cellymax,specific_heat_ratio,vecAx,vecAy)
                    cell_E_hat_plas,cell_E_hat_minus,cell_F_hat_plas,cell_F_hat_minus = setup_cell_flux_hat(Qcon_m1,cellxmax,cellymax,cell_Ahat_plas,cell_Ahat_minus,cell_Bhat_plas,cell_Bhat_minus)
                    E_hat,F_hat = FVS(cell_E_hat_plas,cell_E_hat_minus,cell_F_hat_plas,cell_F_hat_minus,cellxmax,cellymax)

                    # lusgsと違うので注意
                    RHS = set_RHS_inner(cellxmax,cellymax,dt,Qcon_hat_n,Qcon_hat_m1,E_hat,F_hat)
                    Lx,Ly,Ux,Uy,Dinv = set_LDU_inner(dt,cell_Ahat_plas,cell_Ahat_minus,cell_Bhat_plas,cell_Bhat_minus,cellxmax,cellymax,volume)

                    Delta_Qcon_hat_m = inner_iteration(dt,RHS,cellxmax,cellymax,volume,Lx,Ly,Ux,Uy,Dinv,Delta_Qcon_hat_m)

                    res = set_res(Delta_Qcon_hat,Delta_Qcon_hat_temp,cellxmax,cellymax)
                    norm2 = check_converge(res,RHS,cellxmax,cellymax,init_small)
                    Delta_Qcon_hat_temp = copy(Delta_Qcon_hat_m)

                    if norm2[1] < norm_ok && norm2[4] < norm_ok
                        println("\n ---------------------------------- \n")
                        println("nt : "*string(round(inner_t)))
                        println("density res:"*string(norm2[1]) * "  energy res:"*string(norm2[4]))
                        println("\n ---------------------------------- \n")
                        break
                    end
                end

                # inner update
                for i in 2:cellxmax-1
                    for j in 2:cellymax-1
                        for l in 1:4
                            Qcon_hat_m1[i,j,l] = Qcon_hat_m1[i,j,l] + Delta_Qcon_hat_m[i,j,l]
                        end
                    end
                end
            end

            Qcon_hat = copy(Qcon_hat_m1)
            =#

            #println(Qcon_hat[:,:,1])
        else
            println(time_integ)
            println("pressure check time_integ")
            throw(UndefVarError(:x))
        end

        Qcon = Qhat_to_Q(Qcon_hat,cellxmax,cellymax,volume)
        Qbase = conservative_to_base(Qcon,cellxmax,cellymax,specific_heat_ratio)

        if round(evalnum) % every_outnum == 0
            println("nt_______________________________"*string(round(evalnum)))
            output_result(evalnum,Qbase,cellxmax,cellymax,specific_heat_ratio,out_file_front,out_ext,out_dir)
        end

        if Qcon[2,2,4] < 0
            println("pressure are minus !!")
            break
        end
        if Qcon[2,2,1] == NaN
            println(" NaN !!")
            break
        end

    end
end

# -- main --
main()


export Portfolio

function Portfolio(cbmo::Float64,mu::Vector,
                   Q::Any,u::Vector{Float64},
                   kappa::Int64,z0::Vector{Float64}
                   ;relax::Function=NL,
                   t::Float64=0.0,
		   n::Int64=5,
		   tol::Float64=1e-6)

	c1=z->sum(z[1:n]);lc1=1.0;uc1=1.0;
	c2=z->sum(z[n+1:2*n]);lc2=n-kappa;uc2=Inf;
	c3=z->z[1:n].*z[n+1:2*n];lc3=-Inf*ones(n);uc3=zeros(n);
	#c3=z->sum(z[1:n].*z[n+1:2*n]);lc3=-Inf;uc3=zeros(1);
	c3=z->relax(z[1:n],z[n+1:2*n],t=t);lc3=-Inf*ones(n);uc3=zeros(n);
	cons=z->vcat(c1(z),c2(z),c3(z));
	fobj=z->cbmo*sqrt(z[1:n]'*Q*z[1:n])-z[1:n]'*mu;
	lvar=zeros(2*n);uvar=vcat(u,ones(n));
	mp=ADNLPModel(fobj, z0, lvar=lvar, uvar=uvar, c=cons, lcon=vcat(lc1,lc2,lc3), ucon=vcat(uc1,uc2,uc3))

 return mp
end

function Portfolio(cbmo::Float64,mu::Vector,
                   Q::Any,u::Vector{Float64},
                   kappa::Int64,z0::Vector{Float64},
		   tol::Float64=1e-6)
	n=size(Q,1)

	c1=z->sum(z[1:n]);lc1=1.0;uc1=1.0;
	c2=z->sum(z[n+1:2*n]);lc2=n-kappa;uc2=Inf;
	c3=z->z[1:n].*z[n+1:2*n];lc3=-Inf*ones(n);uc3=zeros(n);
	cons=z->vcat(c1(z),c2(z),c3(z));
	fobj=z->cbmo*sqrt(z[1:n]'*Q*z[1:n])-z[1:n]'*mu;
	lvar=zeros(2*n);uvar=vcat(u,ones(n));
	mp=ADNLPModel(fobj, z0, lvar=lvar, uvar=uvar, c=cons, lcon=vcat(lc1,lc2,lc3), ucon=vcat(uc1,uc2,uc3))

	#ex1=JuMP.Model()
	#JuMP.@variable(ex1,x[1:2],start=1.0)
	#JuMP.@NLobjective(ex1,Min,x[1]-x[2])
	#JuMP.@constraint(ex1,1-x[2]>=0)
	#ex1=MathProgNLPModel(ex1)

	#cG=z->z[1:n]
	#G=ADNLPModel(x->(), z0, c=cG, lcon=zeros(n))
	#cH=z->z[n+1:2*n]
	#H=ADNLPModel(x->(), z0, c=cH, lcon=zeros(n))

	G=JuMP.Model()
	JuMP.@variable(G,x[1:2*n],start=1.0)
	JuMP.@constraint(G,x[1:n].>=0)
	JuMP.@NLobjective(G,Min,0.0)
	G=MathProgNLPModel(G)
	H=JuMP.Model()
	JuMP.@variable(H,x[1:2*n],start=1.0)
	JuMP.@constraint(H,x[n+1:2*n].>=0)
	JuMP.@NLobjective(H,Min,0.0)
	H=MathProgNLPModel(H)

 return mp,G,H
end


struct Lengths{T}
    a::T # height of pivot
    b::T # length of long arm
    c::T # length of short arm
    d::T # lenght of weight arm
    e::T # length of sling
    u::T # radius of weight
    z::T # radius of projectile
end

struct Masses{T}
    w::T # mass of weight
    p::T # mass of projectile
    a::T # mass of arm
end

struct Angles{T} # in radians
    aq::T # angle between pivot stand and arm
    wq::T # angle between arm and weight
    sq::T # angle between arm and sling
end

struct AnglularVelocities{T}
    aw::T # for aq
    ww::T # for wq
    sw::T # for sq
end

struct Inertias{T}
    iw::T # inertia of weight
    ia::T # inertia of arm
end

struct Vec
    x
    y
end

struct Constants{T}
    w::T # wind speed
    ρ::T # Density of Air
    Cd::T # Drag Co-efficient of Projectile
    Grav::T
    r::T # release angle of projectile [input]
end

mutable struct Solution
    WeightCG
    WeightArm
    ArmSling
    Projectile
    SlingEnd
    ArmCG
    Time
    ReleaseVelocity
    ReleasePositon
end

mutable struct TrebuchetState
    l::Lengths
    m::Masses
    a::Angles # [stage 1 & 2]
    aw::AnglularVelocities # [stage 1 & 2]
    c::Constants
    i::Inertias
    stage::Union{Val{:Ground},Val{:Hang},Val{:Released},Val{:End}}
    rate
    p # projectile point [stage 3]
    v # projectile speed [stage 3]
    sol::Solution
    Tn # tension on string
    terminated
    function TrebuchetState(l::Lengths, m::Masses, c::Constants, rate)
        θ = asin(l.a/l.b)
        sq = π - θ
        aq = π/2 + θ
        wq = -aq
        ai = m.a*((l.b + l.c)^2)/12
        wi = m.w*(l.u)^2/2 # inertia for disk
        a = Angles(aq, wq, sq)
        aw = AnglularVelocities(0.0, 0.0, 0.0)
        i = Inertias(wi, ai)
        new(l, m, a, aw, c, i, Val{:Ground}(), rate, -1, -1, Solution(), 0, false)
    end
end


# Newton and quasi Newton solvers

## LOGISTIC + 0/L2 and Huber ==============

"""
$SIGNATURES

Fit a GLR using Newton's method.

## Complexity

Assuming `n` dominates `p`, O(κnp²), dominated by the construction of the
Hessian at each step with κ the number of Newton steps.
"""
function _fit(glr::GLR{<:Union{LogisticLoss,RobustLoss},<:L2R},
              solver::Newton, X, y)
    p     = size(X, 2) + Int(glr.fit_intercept)
    θ₀    = zeros(p)
    _fgh! = fgh!(glr, X, y)
    opt   = Optim.only_fgh!(_fgh!)
    res   = Optim.optimize(opt, θ₀, Optim.Newton())
    return Optim.minimizer(res)
end

"""
$SIGNATURES

Fit a GLR using Newton's method combined with an iterative solver  (conjugate
gradient) to solve the Newton steps (∇²f)⁻¹∇f.

## Complexity

Assuming `n` dominates `p`, O(κ₁κ₂np), dominated by the application of the
Hessian at each step where κ₁ is the number of Newton steps and κ₂ is the
average number of CG steps per Newton step (which is at most p).
"""
function _fit(glr::GLR{<:Union{LogisticLoss,RobustLoss},<:L2R},
              solver::NewtonCG, X, y)
    p    = size(X, 2) + Int(glr.fit_intercept)
    θ₀   = zeros(p)
    _f   = objective(glr, X, y)
    _fg! = (g, θ) -> fgh!(glr, X, y)(0.0, g, nothing, θ) # Optim.jl/issues/738
    _Hv! = Hv!(glr, X, y)
    opt  = Optim.TwiceDifferentiableHV(_f, _fg!, _Hv!, θ₀)
    res  = Optim.optimize(opt, θ₀, Optim.KrylovTrustRegion())
    return Optim.minimizer(res)
end

"""
$SIGNATURES

Fit a GLR using LBFGS.

## Complexity

Assuming `n` dominates `p`, O(κnp), dominated by the computation of the
gradient at each step with κ the number of LBFGS steps.
"""
function _fit(glr::GLR{<:Union{LogisticLoss,RobustLoss},<:L2R},
              solver::LBFGS, X, y)
    p    = size(X, 2) + Int(glr.fit_intercept)
    θ₀   = zeros(p)
    _fg! = (f, g, θ) -> fgh!(glr, X, y)(f, g, nothing, θ)
    opt  = Optim.only_fg!(_fg!)
    res  = Optim.optimize(opt, θ₀, Optim.LBFGS())
    return Optim.minimizer(res)
end

## MULTINOMIAL + 0/L2 ==============

"""
$SIGNATURES

Fit a multiclass GLR using Newton's method with an iterative solver (conjugate
gradient).

## Complexity

Assuming `n` dominates `p`, O(κ₁κ₂npc), where `c` is the number of classes. The
computations are dominated by the application of the Hessian at each step with
κ₁ the number of Newton steps and κ₂ the average number of CG steps per Newton
step.
"""
function _fit(glr::GLR{MultinomialLoss,<:L2R}, solver::NewtonCG, X, y)
    p    = size(X, 2) + Int(glr.fit_intercept)
    c    = maximum(y)
    θ₀   = zeros(p * c)
    _f   = objective(glr, X, y; c=c)
    _fg! = (g, θ) -> fg!(glr, X, y)(0.0, g, θ) # XXX: Optim.jl/issues/738
    _Hv! = Hv!(glr, X, y)
    opt  = Optim.TwiceDifferentiableHV(_f, _fg!, _Hv!, θ₀)
    res  = Optim.optimize(opt, θ₀, Optim.KrylovTrustRegion())
    return Optim.minimizer(res)
end

"""
$SIGNATURES

Fit a multiclass GLR using LBFGS.

## Complexity

Assuming `n` dominates `p`, O(κnpc), with `c` the number of classes, dominated
by the computation of the gradient at each step with κ the number of LBFGS
steps.
"""
function _fit(glr::GLR{MultinomialLoss,<:L2R}, solver::LBFGS, X, y)
    p    = size(X, 2) + Int(glr.fit_intercept)
    c    = maximum(y)
    θ₀   = zeros(p * c)
    _fg! = fg!(glr, X, y)
    opt  = Optim.only_fg!(_fg!)
    res  = Optim.optimize(opt, θ₀, Optim.LBFGS())
    return Optim.minimizer(res)
end


## Taylor coeffcients of different functions
taylor_coeff(::typeof(exp), i, x) = exp(x) / factorial(i)

function taylor_coeff(::typeof(sin), i, x)
    j = i % 4

    (j == 0) && (return sin(x) / factorial(i))
    (j == 1) && (return cos(x) / factorial(i))
    (j == 2) && (return -sin(x) / factorial(i))
    (j == 3) && (return -cos(x) / factorial(i))
end

taylor_coeff(::typeof(inv), i, x) = (-1)^i / (x^(i+1))

taylor_coeffs(f, n, x0) = taylor_coeff.(f, 0:n, x0)

doc"""
Make a Taylor1Model for a given function over a given domain.
"""
function Taylor1Model(f, n, x0, I)

    a = taylor_coeffs(f, n, x0)
    p = Taylor1(a)

    Γ = taylor_coeff(f, n+1, I)

    if sup(Γ) ≤ 0 || inf(Γ) ≥ 0

        lo = Interval(inf(I))
        hi = Interval(sup(I))

        Δlo = f(lo) - bound(p, x0, lo)
        Δhi = f(hi) - bound(p, x0, hi)
        Δx0 = f(x0) - bound(p, x0, x0)
        Δ = Interval(min(inf(Δlo), inf(Δx0), inf(Δhi)), max(sup(Δlo), sup(Δx0), sup(Δhi)))

    else
        V = (I - x0)^(n+1)
        Δ = V * Γ
    end

    return Taylor1Model(n, x0, I, p, Δ)
end

doc"""
Evaluate a polynomial of a polynomial-like object (Taylor1Model, Taylor1) `f`.
`a` are the coefficients of the polynomial, expressed as a `Taylor1`.
"""
function evaluate_polynomial(a::Taylor1, f)

    M = zero(f)

    # Horner method:
    for i in degree(f):-1:0
        M = (M * f) + a[i]
    end

    return M
end

doc"""
Calculate the Taylor1Model of `(g∘f)` given a function `g` and a Taylor1Model `f`.
"""
function TMcomposition(g, f::Taylor1Model)
    n = f.n

    # centre ("x0") for g is f(x0), which is first coefficient of expansion of f:
    f_of_x0 = f[0]  # constant_term(f.p)

    # calculate interval for g: image of f
    Bf = bound(f)  # Bf + Δf in Joldes thesis; basically f(I)

    Mg = Taylor1Model(g, n, f_of_x0, Bf)

    b, Δg = Mg.p, Mg.Δ

    f_non_constant_part = [0; f.p[1:end]]  # f - f(x0)

    M1 = Taylor1Model(f, f_non_constant_part, f.Δ)

    M = evaluate_polynomial(b, M1)

    return Taylor1Model(M, M.Δ + Δg)   # same polynomial, different interval
end

# allow exp(t) etc. for a Taylor1Model t:
for g in (:exp, :sin, :inv)
    @eval $g(f::Taylor1Model) = TMcomposition($g, f)
end

function /(f::Taylor1Model, g::Taylor1Model)
    @assert data(f) == data(g)

    M = f * inv(g)
end


# minimal use of the package, user manages all the plotting

using Pseudospectra
using Plots

A = randn(100,100)
# we want a true Schur decomposition, so force complexity
Tschur,U,eigA  = schur(A .+ 0im)
ax = [-12,12,-12,12]
npts = 40
opts = Dict{Symbol,Any}(:real_matrix => true)
Z,x,y,levels,err,Tproj,eigAproj = psa_compute(Tschur,npts,
                                              ax,eigA,opts)
p = scatter(real(eigA),imag(eigA),color="black",label="")
if isempty(levels)
    contour!(p,x,y,log10.(Z),linewidth=2)
else
    contour!(p,x,y,log10.(Z),levels=levels,linewidth=2)
end


store = [KPCSupply("beef", 38//10, 36),
         KPCSupply("pork", 54//10, 43),
         KPCSupply("ham", 36//10, 90),
         KPCSupply("greaves", 24//10, 45),
         KPCSupply("flitch", 4//1, 30),
         KPCSupply("brawn", 25//10, 56),
         KPCSupply("welt", 37//10, 67),
         KPCSupply("salami", 3//1, 95),
         KPCSupply("sausage", 59//10, 98)]

sack = solve(store, 15)
println("The store contains:\n - ", join(store, "\n - "))
println("\nThe thief should take::\n - ", join(sack, "\n - "))
@printf("\nTotal value in the sack: %.2f €\n", sum(getfield.(sack, :value)))


using ITensors.NDTensors,
      Test
using LinearAlgebra

@testset "random_orthog" begin
  n,m = 10,4
  O1 = random_orthog(n,m)
  @test norm(transpose(O1)*O1 - Diagonal(fill(1.,m))) < 1E-14

  n,m = 4,10
  O2 = random_orthog(n,m)
  @test norm(O2*transpose(O2) - Diagonal(fill(1.,n))) < 1E-14
end


function find_surface(r0, dr, psi, psi0)

    r = r0[:]
    sgn = psi(r0...) - psi0
    while true
        rnew += dr
        if sgn*(psi(rnew...) - psi0) < 0
            return 0.5*(rnew + r)
        end
        r[:] = rnew
    end
end

function rootfind_bisection(f::Function, val::Float64, bracket1::Vector{Float64}, bracket2::Vector{Float64})

    a = bracket1[:]
    b = bracket2[:]
    fa = f(a) - val
    fb = f(b) - val
    c = (a + b)/2
    # c = (a*fb - b*fa) / (fb - fa)
    fc = f(c) - val

    while abs(fc) > 1e-8
        if fa*fc < 0
            b = c
        elseif fc*fb < 0
            a = c
        else
            @error "root finding algorithm has failed."
            break
        end
        fa = f(a) - val
        fb = f(b) - val
        c = (a + b)/2
        # c = (a*fb - b*fa) / (fb - fa)
        fc = f(c) - val
    end

    return c
end

function rk4_step(r::Vector{Float64}, f::Function, ds::Float64)

    k1 = f(r)
    k2 = f(r + ds*k1/2)
    k3 = f(r + ds*k2/2)
    k4 = f(r + ds*k3)

    return r + ds*(k1 + 2*k2 + 2*k3 + k4) / 6
end

function conservative_step(r::Vector{Float64}, b::Function, psi::Function, ds::Float64)

    # compute predictor step using rk4
    rp = rk4_step(r, b, ds)

    # compute correction using root finding algorithm
    en = zeros(length(rp))
    en[1:2] = [0 -1; 1 0] * b(rp)[1:2] / norm(b(rp)[1:2])
    bracket1 = rp + ds*en
    bracket2 = rp - ds*en
    rc = rootfind_bisection(psi, psi(r), bracket1, bracket2)

    # compute with correct displacement
    return rc, norm(rc - r)
end

function trace_fieldline(stop_condition::Function, r0::Vector{Float64}, b::Function, psi::Function, ds::Float64)

    r = r0[:]
    deltaS = 0.0

    prog = ProgressUnknown("Tracing fieldline... ", spinner=true)
    points = Vector{Float64}[r0]
    while !stop_condition(r)
        rnew, dsnew = conservative_step(r, b, psi, ds)
        r[:] = rnew
        deltaS += dsnew
        append!(points, [rnew])
        next!(prog)
    end

    finish!(prog)
    return hcat(points...), deltaS
end

function trace_fieldline(N::Int64, r0::Vector{Float64}, b::Function, psi::Function, ds::Float64)

    r = r0[:]
    deltaS = 0.0

    points = zeros(length(r0), N)
    points[:,1] = r[:]
    for s=2:N
        rnew, dsnew = conservative_step(r, b, psi, ds)
        r[:] = rnew
        deltaS += dsnew
        points[:,s] = r
    end

    return points, deltaS
end


#=
utilities for complex polynomial interpolation

TODO make interface for complex and real interpolation consistent, use dispatch
=#

"""
$(SIGNATURES)

Compute interpolation data for a complex weighted sum-of-squares conic constraint
on a domain.
"""
function interpolate(
    R::Type{Complex{T}},
    halfdeg::Int,
    n::Int,
    gs::Vector,
    g_halfdegs::Vector{Int};
    sample_factor::Int = 10,
    use_qr::Bool = false,
    ) where {T <: Real}
    # generate interpolation
    # TODO use more numerically-stable basis for columns, and evaluate in a more
    # numerically stable way by multiplying the columns
    L = binomial(n + halfdeg, n)
    U = L^2
    L_basis = [a for t in 0:halfdeg for a in Combinatorics.multiexponents(n, t)]
    mon_pow(z, ex) = prod(z[i]^ex[i] for i in eachindex(ex))
    V_basis = [z -> mon_pow(z, L_basis[k]) * mon_pow(conj(z), L_basis[l]) for
        l in eachindex(L_basis) for k in eachindex(L_basis)]
    @assert length(V_basis) == U

    # sample from domain (inefficient for general domains, only samples from
    # unit box and checks feasibility)
    num_samples = sample_factor * U
    samples = Vector{Vector{Complex{T}}}(undef, num_samples)
    k = 0
    randbox() = 2 * rand(T) - 1
    while k < num_samples
        z = [Complex(randbox(), randbox()) for i in 1:n]
        if all(g -> g(z) > zero(T), gs)
            k += 1
            samples[k] = z
        end
    end

    # select subset of points to maximize abs(det(V)) in heuristic QR-based
    # procedure (analogous to real case)
    V = [b(z) for z in samples, b in V_basis]
    VF = qr(Matrix(transpose(V)), ColumnNorm())
    keep = VF.p[1:U]
    points = samples[keep]
    V = V[keep, :]

    # setup P matrices
    P0 = V[:, 1:L]
    if use_qr
        P0 = Matrix(qr(P0).Q)
    end
    Ps = [P0]
    for i in eachindex(gs)
        gi = gs[i].(points)
        @views P0i = P0[:, 1:binomial(n + halfdeg - g_halfdegs[i], n)]
        Pi = Diagonal(sqrt.(gi)) * P0i
        if use_qr
            Pi = Matrix(qr(Pi).Q)
        end
        push!(Ps, Pi)
    end

    return (points, Ps)
end


export Pade21

"""
    Pade21()

This uses a Padé-approximation of type (2,1) for prediction.
"""
struct Pade21 <: AbstractPredictor end
struct Pade21Cache{T,AV<:AbstractVector{T}} <: AbstractPredictorCache
    x²::Vector{T}
    x³::Vector{T}
    x_h::AV

    u::Vector{T}
    u₁::Vector{T}
    u₂::Vector{T}
end

function cache(::Pade21, H, x, ẋ, t)
    x², x³, x_h = copy(ẋ), copy(ẋ), copy(x)
    u = evaluate(H, x, t)
    u₁, u₂ = copy(u), copy(u)
    Pade21Cache(x², x³, x_h, u, u₁, u₂)
end

@inline function g!(u, H, x, ẋ, t, h, x_h)
    @inbounds for i in eachindex(x)
        x_h[i] = x[i] + h * ẋ[i]
    end
    evaluate!(u, H, x_h, t + h)
end

@inline function g!(u, H, x, ẋ, ẍ, t, h, x_h)
    h² = 0.5 * h * h
    @inbounds for i in eachindex(x)
        x_h[i] = x[i] + h * ẋ[i] + h² * ẍ[i]
    end
    evaluate!(u, H, x_h, t + h)
end

function update!(cache::Pade21Cache, H, x, ẋ, t, Jac::JacobianMonitor, eval_err::Float64)
    # unpack stuff to make the rest easier to read
    @unpack u, u₁, u₂, x_h, x², x³ = cache
    ψ = 10eval_err
    h₂ = nthroot(ψ, 4)
    g!(u₁, H, x, ẋ, t, h₂, x_h)
    g!(u₂, H, x, ẋ, t, -h₂, x_h)

    h² = h₂^2
    @inbounds for i in eachindex(u₁)
        u[i] = -(u₁[i] + u₂[i]) / h²
    end
    LA.ldiv!(x², Jac, u)

    h₃ = nthroot(ψ, 5)
    g!(u₁, H, x, ẋ, x², t, h₃, x_h)
    g!(u₂, H, x, ẋ, x², t, -h₃, x_h)

    h³ = h₃^3
    @inbounds for i in eachindex(u₁)
        u[i] = -3 * (u₁[i] - u₂[i]) / h³
    end
    LA.ldiv!(x³, Jac, u)

    cache
end

function predict!(x̂, cache::Pade21Cache, H::HomotopyWithCache, x, t, Δt, ẋ, J)
    @unpack x², x³ = cache
    @inbounds for i in eachindex(x)
        δ = 1 - Δt * x³[i] / (3 * x²[i])
        if isnan(δ) || iszero(δ)
            x̂[i] = x[i] + Δt * ẋ[i] + 0.5 * Δt^2 * x²[i]
        else
            x̂[i] = x[i] + Δt * ẋ[i] + (0.5 * Δt^2 * x²[i]) / δ
        end
    end
    nothing
end

order(::Pade21Cache) = 4
@inline highest_derivative(cache::Pade21Cache) = (cache.x³, 3)
second_derivative(cache::Pade21Cache) = cache.x²
third_derivative(cache::Pade21Cache) = cache.x³


# anonymous functions:
(x, y) -> x^3 - y + x * y
f = (x, y) -> x^3 - y + x * y
println(f(3, 2) ) # 31

function (x)
    x + 2
end

(x) -> x + 2
x -> x + 2  # lambda syntax
() -> println("hello, Julia")
(x, y, z) -> 3x + 2y - z


## the learning algorithm: PLS2 - multiple targets
function trainer(model::PLS2Model{T},
                X::AbstractArray{T}, Y::Matrix{T}) where T<:AbstractFloat
    W,b,P,Q  = model.W,model.b,model.P,model.Q
    model.ntargetcols = size(Y,2)
    nfactors = model.nfactors

    for i = 1:nfactors
        b[:,i]    = Y[:,1] #u: arbitrary col. Thus, I set to the first.
        Rold = b[:,i]
        local R::Vector{T}
        while true

            W[:,i]  = X'b[:,i]
            W[:,i] /= norm(W[:,i])#sqrt.(W[:,i]'*W[:,i])
            R       = X*W[:,i]

            Rold      = R
            Q[:,i]    = Y'R
            Q[:,i]   /= norm(Q[:,i])
            b[:,i]    = Y*Q[:,i]
            if  all(abs.(R - Rold) .<= 1.0e-3)
                break
            end

        end
        Rn     = R'/(R'R) # change to use function...
        P[:,i] = Rn*X

        X      = X - R * P[:,i]'
        c      = Rn'b[:,i:i]'./(R'R)
        Yp     = c*R*Q[:,i]'
        Y      = Y - Yp
        if  all(abs.(Y - Yp) .<= 1.0e-3)
            print("PLS2 converged. No need learning with more than $(i) factors")
            model.nfactors = i
            return model
        end

    end

    return model

end

function predictor(model::PLS2Model{T},
                   X::AbstractArray{T}) where T<:AbstractFloat

    W,Q,P    = model.W,model.Q,model.P
    nfactors   = model.nfactors
    Y          = zeros(T,size(X,1),model.ntargetcols)
    #println("nfactors: ",nfactors)
    for i = 1:nfactors
        R      = X*W[:,i]
        X      = X - R * P[:,i]'
        Y      = Y + R * Q[:,i]'
    end

    return Y
end


include("euler/Julia/grid.jl")
include("euler/Julia/math.jl")
using .Grid: adjacent_product
using .Math: fast_max

function compute(n::Int, grid::String)::Int
    grid = [parse.(Int, row) for row in split.(split(grid, "\n"), " ")]
    width, height = length(grid[1]), length(grid)
    result = 0
    for y ∈ 1:height
        for x ∈ 1:width
            if x + n - 1 ≤ width result = fast_max(result, adjacent_product(grid, n, (x, y), (1, 0))) end
            if y + n - 1 ≤ height result = fast_max(result, adjacent_product(grid, n, (x, y), (0, 1)))
                if x + n - 1 ≤ width result = fast_max(result, adjacent_product(grid, n, (x, y), (1, 1))) end
                if x - n ≥ 0 result = fast_max(result, adjacent_product(grid, n, (x, y), (-1, 1))) end
            end
        end
    end
    return result
end