@doc raw"""
Contains functions and types to work with the analytic solutions to the Dirac equation.

The bound state solutions to the Dirac equation are represented via the ``P_{n\kappa}(r)``
and ``Q_{n\kappa}(r)`` functions defined for each ``n`` and ``\kappa`` quantum number
combination and defined via the relation:

```math
\psi_{n\kappa m}(r,\theta,\varphi)
= \frac{1}{r} \begin{pmatrix}
  P_{n\kappa}(r) \chi_{\kappa m}(\theta,\varphi) \\
  i Q_{n\kappa}(r) \chi_{-\kappa m}(\theta, \varphi)
\end{pmatrix}
```

!!! note "Atomic units"

    All values are in **atomic units**, i.e. ``c = 1/α``, ``m = 1`` and ``\hbar = 1``.
"""
module DiracAnalytic
using SpecialFunctions: gamma
using GSL: sf_laguerre_n
using AtomicLevels: RelativisticOrbital

"""
The fine structure constants ``\\alpha``.

```@repl
α
1/α
```
"""
const α = 0.0072973525664

@doc raw"""
    energy(Z, n, κ) -> Real
    energy(Z, orbital::RelativisticOrbital) -> Real

Calculates the relativistic energy of a hydrogenic orbital in the ``V(r) = - Z / \alpha``
potential:

```math
E_{n\kappa}(Z) =
m c^2\left(1+\left[\frac{Z\alpha}{n-|k|+\sqrt{k^2-Z^2\alpha^2}}\right]^2\right)^{-1/2}
```

Also supports the following signatures

    energy(::Type{Complex}, ...) -> Complex

which allow the value to become complex. This happens after the critical nuclear charge
``Z > \abs{\kappa} / \alpha``.

```julia-repl
julia> energy(Complex, 200, ro"1s")
0.0 + 19962.685927945768im
```
"""
function energy end
energy(Z::Real, orb::RelativisticOrbital) = energy(Z, orb.n, orb.κ)
energy(T::Type{Complex}, Z, orb::RelativisticOrbital) = energy(T, Z, orb.n, orb.κ)
energy(Z::Real, n::Integer, κ::Integer) = _E(Z, n, κ, _Δ(Z, κ))
energy(::Type{Complex}, Z, n::Integer, κ::Integer) = _E(Z, n, κ, complex(_Δ(Z, κ)))
# Internal methods to handle the complex and non-complex versions of E
_E(Z, n, κ, Δ) = 1 / (α^2 * sqrt(1 + (Z*α / (n - abs(κ) + sqrt(Δ)))^2))
_Δ(Z, κ) = κ^2 - (Z*α)^2

# Radial solutions of the DiracEquation

abstract type DiracSolution end
function PQ end
Base.broadcastable(x::DiracSolution) = Ref(x)

"""
Implementation based on the expressions in Wikipedia:
<https://en.wikipedia.org/wiki/Hydrogen-like_atom#Solution_to_Dirac_equation>
"""
struct WikiSolution <: DiracSolution
    Z :: Float64
    n :: Int
    κ :: Int
end

function PQ(s::WikiSolution, r)
    E = energy(s.Z, s.n, s.κ)
    C, γ = _wiki_C(s.Z, s.n), _wiki_γ(s.Z, s.κ)
    A = _wiki_A(s.Z, s.n, s.κ)
    ρ = 2 * C * r

    if s.n == -s.κ
        A * (s.n + γ) * (ρ^γ) * exp(-ρ/2), A * s.Z * α * (ρ^γ) * exp(-ρ/2)
    else
        L1 = sf_laguerre_n(s.n - abs(s.κ) - 1, 2γ + 1, ρ)
        L2 = sf_laguerre_n(s.n - abs(s.κ), 2γ + 1, ρ)
        A * (ρ^γ) * exp(-ρ/2) * (
            s.Z * α * ρ * L1 +
            (γ - s.κ) * (γ/α^2 - s.κ*E) * L2 * α / C
        ),
        A * (ρ^γ) * exp(-ρ/2) * (
            (γ - s.κ) * ρ * L1 +
            s.Z * α * (γ/α^2 - s.κ*E) * L2 * α / C
        )
    end
end

function _wiki_A(Z, n, k)
    C, γ = _wiki_C(Z, n), _wiki_γ(Z, k)
    if n == -k
        sqrt(C / (2 * n * (n + γ) * γ * gamma(2γ)))
    else
        E = energy(Z, n, k)
        sqrt(
            0.5 *
            (C / (n - abs(k) + γ)) *
            (factorial(n - abs(k) - 1) / gamma(n - abs(k) + 2γ + 1)) *
            ((α^2 * E * k / γ)^2 + (α^2 * E * k / γ))
        ) / sqrt(2k * (k - γ))
    end
end

@doc raw"""
    _wiki_γ(Z, κ)

Calculates the ``\gamma = \sqrt{\kappa^2 - Z^2 \alpha^2}`` value.
"""
_wiki_γ(Z, κ) = sqrt(κ^2 - (Z*α)^2)

@doc raw"""
    _wiki_C(Z, n)

Calculates the radial decay constant ``C``:

```math
C = \frac{Z\alpha}{n}\frac{m c^2}{\hbar c} = \frac{Z}{n}
```
"""
_wiki_C(Z, n) = Z/n # μ = 1, ħ = 1 and α*c^2 / c = 1

end # module


"""
    bezier(line::LineString; resolution::Real=10000, sharpness::Real=0.85)

Take a LineString and returns a curved version
by applying a [Bezier spline](http://en.wikipedia.org/wiki/B%C3%A9zier_spline)
algorithm.
"""
function bezier(line::LineString; resolution::Real=10000., sharpness::Real=0.85)

    points::Vector{Point} = []

    for p in line.coordinates
        push!(points, Point(p))
    end

    spline = Spline(points=points, duration=resolution, sharpness=sharpness)

    coords = []
    for i in 1:10:spline.duration
        pos = Splines.pos(i)

        floor(i / 100) % 2 == 0 && push!(coords, [pos.coordinates[1], pos.coordinates[2]])

    end

    return LineString(coords)

end


cluster = [1 1 1 1; 2 2 3 3; 4 4 5 5]
X = Array(reshape(1:24, 2, 3, 4))

@testset "pool" begin
    @testset "GlobalPool" begin
        glb_cltr = [1 1 1 1; 1 1 1 1; 1 1 1 1]
        p = GlobalPool(:add, 3, 4)
        @test p(X) == sumpool(glb_cltr, X)
    end

    @testset "LocalPool" begin
        p = LocalPool(:add, cluster)
        @test p(X) == sumpool(cluster, X)
    end

    @testset "TopKPool" begin
        N = 10
        k, in_channel = 4, 7
        X = rand(in_channel, N)
        for T = [Bool, Float64]
            adj = rand(T, N, N)
            p = TopKPool(adj, k, in_channel)
            @test eltype(p.p) === Float32
            @test size(p.p) == (in_channel,)
            @test eltype(p.Ã) === T
            @test size(p.Ã) == (k, k)
            y = p(X)
            @test size(y) == (in_channel, k)
        end
    end

    for T = [UInt32, UInt64]
        @testset "$(T)" begin
            @testset "sumpool" begin
                y = [40, 44, 12, 14, 36, 38, 16, 18, 40, 42]
                @test sumpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:add, cluster, X) == reshape(y, 2, 5)
            end

            @testset "maxpool" begin
                y = [19, 20, 9, 10, 21, 22, 11, 12, 23, 24]
                @test maxpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:max, cluster, X) == reshape(y, 2, 5)
            end

            @testset "minpool" begin
                y = [1, 2, 3, 4, 15, 16, 5, 6, 17, 18]
                @test minpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:min, cluster, X) == reshape(y, 2, 5)
            end
        end
    end

    for T = [Int32, Int64]
        @testset "$(T)" begin
            @testset "sumpool" begin
                y = [40, 44, 12, 14, 36, 38, 16, 18, 40, 42]
                @test sumpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:add, cluster, X) == reshape(y, 2, 5)
            end

            @testset "subpool" begin
                y = [-40, -44, -12, -14, -36, -38, -16, -18, -40, -42]
                @test subpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:sub, cluster, X) == reshape(y, 2, 5)
            end

            @testset "maxpool" begin
                y = [19, 20, 9, 10, 21, 22, 11, 12, 23, 24]
                @test maxpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:max, cluster, X) == reshape(y, 2, 5)
            end

            @testset "minpool" begin
                y = [1, 2, 3, 4, 15, 16, 5, 6, 17, 18]
                @test minpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:min, cluster, X) == reshape(y, 2, 5)
            end
        end
    end

    for T = [Float16, Float32, Float64]
        @testset "$(T)" begin
            @testset "sumpool" begin
                y = [40, 44, 12, 14, 36, 38, 16, 18, 40, 42]
                @test sumpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:add, cluster, X) == reshape(y, 2, 5)
            end

            @testset "subpool" begin
                y = [-40, -44, -12, -14, -36, -38, -16, -18, -40, -42]
                @test subpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:sub, cluster, X) == reshape(y, 2, 5)
            end

            @testset "maxpool" begin
                y = [19, 20, 9, 10, 21, 22, 11, 12, 23, 24]
                @test maxpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:max, cluster, X) == reshape(y, 2, 5)
            end

            @testset "minpool" begin
                y = [1, 2, 3, 4, 15, 16, 5, 6, 17, 18]
                @test minpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:min, cluster, X) == reshape(y, 2, 5)
            end

            @testset "prodpool" begin
                y = [1729, 4480, 27, 40, 315, 352, 55, 72, 391, 432]
                @test prodpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:mul, cluster, X) == reshape(y, 2, 5)
            end

            @testset "divpool" begin
                y = 1 ./ [1729, 4480, 27, 40, 315, 352, 55, 72, 391, 432]
                @test divpool(cluster, X) ≈ reshape(y, 2, 5)
                @test pool(:div, cluster, X) ≈ reshape(y, 2, 5)
            end

            @testset "meanpool" begin
                y = [10., 11., 6., 7., 18., 19., 8., 9., 20., 21.]
                @test meanpool(cluster, X) == reshape(y, 2, 5)
                @test pool(:mean, cluster, X) == reshape(y, 2, 5)
            end
        end
    end
end


# using BoxCoxTrans
import YeoJohnsonTrans as YJT
using Tullio

function YJT.transform(X::AbstractArray, λs::Vector)
    Y = deepcopy(X)
    for j in size(X, 1)
        Y[j, :] = YJT.transform(X[j, :], λs[j])
    end
    return Y
end

"""
Yeo-Johnson transformation according to a matrix X
"""
function yeojohnson(X::AbstractArray)
    λs = Array{Float64, 1}(undef, size(X, 1))
    for (i, x) ∈ eachrow(X) |> enumerate
        λs[i], details = YJT.lambda(x; interval=(-10.0, 10.0))
    end
    return Y -> YJT.transform(Y, λs)
end
export yeojohnson


type MLEst
    rho
    transform

    MLEst() = new()
end

function MLEst(snr; transform = (x,_) -> x)
    est = MLEst()
    est.rho = 10^(0.1*snr);
    est.transform = transform
    est
end

function estimate(est::MLEst, y)
    z = est.transform(y, :notransp)
    (nAntennas,nCoherence,nBatches) = size(z)# input data in transformed domain

    cest = max.( 0, sum(abs2,z,2) - nCoherence/est.rho )
    cest = reshape(cest,nAntennas,nBatches)

    W = cest./(cest .+ nCoherence/est.rho)
    W = reshape(W,nAntennas,1,nBatches)

    hest = est.transform( W .* z, :transp)
end


export pairwise_marginal, pairwise_MI

using LinearAlgebra: diagind, diag
using CUDA

xlogx(x) =
    iszero(x) ? zero(x) : x * log(x)

xlogy(x, y) =
    iszero(x) && !isnan(y) ? zero(x) : x * log(y)

#############################
# Mutual Information on Binary Data
#############################

function pairwise_marginal(data::Union{Matrix{Bool}, CuMatrix{Bool}, BitMatrix};
        weights::Union{Vector, CuVector, Nothing}=nothing,
        pseudocount=0.0,
        Float=Float64)
    N = isnothing(weights) ? size(data, 1) : sum(weights)
    D = size(data, 2)
    base = N + pseudocount
    if data isa CuMatrix{Bool} && !isnothing(weights)
        weights = CuVector(weights)
    end

    not_data = Float.(.!data)
    data = Float.(data)
    pxy = similar(data, D, D, 4)

    if isnothing(weights)
        pxy[:, :, 4] = transpose(data) * data
        pxy[:, :, 3] = transpose(data) * not_data
        pxy[:, :, 2] = transpose(not_data) * data
    else
        pxy[:, :, 4] = transpose(data) * (data .* weights)
        pxy[:, :, 3] = transpose(data) * (not_data .* weights)
        pxy[:, :, 2] = transpose(not_data) * (data .* weights)
    end
    pxy[:, :, 1] = N .- pxy[:, :, 4] .- pxy[:, :, 3] .- pxy[:, :, 2]

    joint_count = pseudocount / 4
    pxy .+= joint_count

    # diagonal is the marginal of p(x)
    diag = diagind(pxy[:, :, 1])
    @view(pxy[:,:, 1])[diag] .+= joint_count
    @view(pxy[:,:, 2])[diag] .= zero(Float)
    @view(pxy[:,:, 3])[diag] .= zero(Float)
    @view(pxy[:,:, 4])[diag] .+= joint_count

    if data isa CuMatrix
        CUDA.unsafe_free!(data)
        CUDA.unsafe_free!(not_data)
    end
    pxy ./= base
end

function pairwise_MI(data::Union{Matrix{Bool}, CuMatrix{Bool}, BitMatrix};
        weights=nothing,
        pseudocount=0.0,
        Float=Float64)
    pxy = pairwise_marginal(data; weights, pseudocount, Float)
    p0 = diag(pxy[:, :, 1])
    p1 = diag(pxy[:, :, 4])
    D = size(data, 2)
    pxpy = Float.(similar(data, D, D, 4))
    pxpy[:,:,1] = p0 * p0'
    pxpy[:,:,2] = p0 * p1'
    pxpy[:,:,3] = p1 * p0'
    pxpy[:,:,4] = p1 * p1'
    pxy_log_pxy = @. xlogx(pxy)
    pxy_log_pxpy = @. xlogy(pxy, pxpy)
    MI = dropdims(sum(pxy_log_pxy - pxy_log_pxpy, dims=3), dims=3)
    if data isa CuMatrix{Bool}
        CUDA.unsafe_free!(pxy)
        CUDA.unsafe_free!(pxpy)
        CUDA.unsafe_free!(pxy_log_pxy)
        CUDA.unsafe_free!(pxy_log_pxpy)
    end
    MI
end

#############################
# Mutual Information on Categorical Data
#############################

function pairwise_marginal(data::Matrix;
                           weights::Union{Vector, Nothing} = nothing,
                           num_cats = maximum(data) - minimum(data) + 1,
                           pseudocount = zero(Float32),
                           Float = Float32)

    @assert minimum(data) >= 0 "Categorical data labels are assumed to be indexed starting 0"

    num_samples = size(data, 1)
    num_vars = size(data, 2)

    if weights === nothing
        weights = ones(Float32, num_samples)
    else
        pseudocount = pseudocount * sum(weights) / num_samples
    end

    num_vars = size(data, 2)
    pair_margs = Array{Float}(undef, num_vars, num_vars, num_cats, num_cats)
    Z = sum(weights) + pseudocount
    single_smooth = Float(pseudocount / num_cats)
    pair_smooth = Float(single_smooth / num_cats)

    for i = 1:num_vars
        # note: multithreading needs to be on inner loop for thread-safe copying across diagonal
        Threads.@threads for j = 1:num_vars
            if i<=j
                if i!=j
                    @inbounds pair_margs[i,j,:,:] .= pair_smooth
                else
                    @inbounds pair_margs[i,j,:,:] .= zero(eltype(pair_margs))
                    for l = 1:num_cats
                        @inbounds pair_margs[i,j,l,l] = single_smooth
                    end
                end
                @simd for k = 1:size(data,1) # @avx here gives incorrect results
                    @inbounds pair_margs[i,j,data[k,i]+1,data[k,j]+1] += weights[k]
                end
                @inbounds pair_margs[i,j,:,:] ./= Z
            else
                @inbounds pair_margs[i,j,:,:] .= transpose(@view pair_margs[j,i,:,:])
            end
            nothing
        end
    end

    return pair_margs
end

function balance_threads(num_items, num_examples, config; mine, maxe)
    block_threads = config.threads
    # make sure the number of example threads is a multiple of 32
    num_item_batches = cld(num_items, maxe)
    num_blocks = cld(num_item_batches * num_examples, block_threads)
    if num_blocks < config.blocks
        max_num_item_batch = cld(num_items, mine)
        max_num_blocks = cld(max_num_item_batch * num_examples, block_threads)
        num_blocks = min(config.blocks, max_num_blocks)
        num_item_batches = (num_blocks * block_threads) ÷ num_examples
    end
    item_work = cld(num_items, num_item_batches)
    @assert item_work*block_threads*num_blocks >= num_examples*num_items
    block_threads, num_blocks, num_examples, item_work
end

function pairwise_marginal_kernel(pair_margs, data, weights, num_ex_threads::Int32, vpair_work::Int32)
    threadid = ((blockIdx().x - one(Int32)) * blockDim().x) + threadIdx().x

    vpair_batch, ex_id = fldmod1(threadid, num_ex_threads)

    num_vars = size(pair_margs, 1)
    vpair_start = one(Int32) + (vpair_batch - one(Int32)) * vpair_work
    vpair_end = min(vpair_start + vpair_work - one(Int32), num_vars ^ Int32(2))

    @inbounds if ex_id <= size(data, 1)
        for vpair_id = vpair_start : vpair_end
            v1, v2 = fldmod1(vpair_id, num_vars)
            if v1 <= v2
                d1 = data[ex_id, v1]
                d2 = data[ex_id, v2]
                CUDA.@atomic pair_margs[v1,v2,d1+1,d2+1] += weights[ex_id]
            end
        end
    end
    nothing
end

function pairwise_marginal(data::CuMatrix;
                           weights::Union{CuVector, Vector, Nothing} = nothing,
                           num_cats = maximum(data) - minimum(data) + 1,
                           pseudocount = zero(Float32),
                           Float = Float32)

    @assert eltype(data) != Bool
    @assert minimum(data) >= 0 "Categorical data labels are assumed to be indexed starting 0"

    num_examples = size(data, 1)
    num_vars = size(data, 2)
    num_var_pairs = num_vars * num_vars

    if weights === nothing
        weights = ones(Float, num_examples)
    else
        pseudocount = pseudocount * sum(weights) / num_examples
    end

    num_vars = size(data,2)
    pair_margs = zeros(Float, num_vars, num_vars, num_cats, num_cats)
    Z = sum(weights) + pseudocount
    single_smooth = Float(pseudocount / num_cats)
    pair_smooth = Float(single_smooth / num_cats)

    # init pair_margs
    for i = 1 : num_vars
        for j = 1 : num_vars
            if i <= j
                if i != j
                    @inbounds pair_margs[i,j,:,:] .= pair_smooth
                else
                    @inbounds pair_margs[i,j,:,:] .= zero(eltype(pair_margs))
                    for l = 1 : num_cats
                        @inbounds pair_margs[i,j,l,l] = single_smooth
                    end
                end
            end
            nothing
        end
    end

    pair_margs = cu(pair_margs)
    weights = cu(weights)

    dummy_args = (pair_margs, data, weights, Int32(1), Int32(1))
    kernel = @cuda name="pairwise_marginal" launch=false pairwise_marginal_kernel(dummy_args...)
    config = launch_configuration(kernel.fun)

    threads, blocks, num_example_threads, vpair_work =
        balance_threads(num_var_pairs, num_examples, config; mine=2, maxe=32)

    args = (pair_margs, data, weights, Int32(num_example_threads), Int32(vpair_work))
    kernel(args...; threads, blocks)

    pair_margs ./= Z
    pair_margs = Array(pair_margs)

    for i = 1 : num_vars
        for j = 1 : num_vars
            if i > j
                @inbounds pair_margs[i,j,:,:] .= transpose(@view pair_margs[j,i,:,:])
            end
        end
    end

    pair_margs
end

function pairwise_MI(data::Union{Matrix,CuMatrix};
                     weights::Union{Vector, Nothing} = nothing,
                     num_cats = maximum(data) - minimum(data) + 1,
                     pseudocount = zero(Float32),
                     Float = Float32)
    num_samples = size(data, 1)
    num_vars = size(data, 2)

    if weights === nothing
        weights = ones(Int32, num_samples)
    else
        pseudocount = pseudocount * sum(weights) / num_samples
    end

    # `joint_cont[i, j, k, w]' is the total weight of samples whose i- and j-th variable are k and w, respectively
    joint_cont = pairwise_marginal(data; weights, num_cats, pseudocount)

    # `marginal_cont[i, j]' is the total weight of sample whose i-th variable is j
    marginal_cont = zeros(Float, num_vars, num_cats)
    for i = 1:num_vars
        for j = 1:num_cats
            @inbounds marginal_cont[i,j] = joint_cont[i,i,j,j]
        end
    end

    # Compute mutual information
    MI = zeros(Float, num_vars, num_vars)
    for var1_idx = 1 : num_vars
        for var2_idx = var1_idx : num_vars
            @inbounds MI[var1_idx, var2_idx] = sum(
                xlogx.(joint_cont[var1_idx, var2_idx, :, :]) .-
                xlogy.(joint_cont[var1_idx, var2_idx, :, :], marginal_cont[var1_idx, :] .* marginal_cont[var2_idx, :]'))
        end
    end

    for var1_idx = 2 : num_vars
        for var2_idx = 1 : var1_idx - 1
            @inbounds MI[var1_idx, var2_idx] = MI[var2_idx, var1_idx]
        end
    end

    MI
end


using Surrogates

#1D
x = [1.0,2.0,3.0]
y = [3.0,5.0,7.0]
lb = 0.0
ub = 5.0
my_wend = Wendland(x,y,lb,ub)
add_point!(my_wend,0.5,4.0)
val = my_wend(0.5)

#ND
lb = [0.0,0.0]
ub = [4.0,4.0]
x = sample(5,lb,ub,SobolSample())
f = x -> x[1]+x[2]
y = f.(x)
my_wend_ND = Wendland(x,y,lb,ub)
est = my_wend_ND((1.0,2.0))
add_point!(my_wend_ND,(3.0,3.5),4.0)
add_point!(my_wend_ND,[(9.0,10.0),(12.0,13.0)],[10.0,11.0])
#todo


function nll(m::NystromModel)
	if !same_hyp(m)
		recompute_kernel_eig!(m)
	end
	refresh_cache!(m)
	sigma2 = exp(2m.lik.log_sigma)
	neig = length(m.nystrom.values)
	LambdaInv = diagm(neig./m.nystrom.values)
	Z = (m.PhiTPhi + sigma2 * LambdaInv)
	yTQy = (m.yTy - dot(m.PhiTy, Z \ m.PhiTy)) / sigma2
	logdetZ = logdet(Z)
	n = size(m.X, 1)
	scale = n - length(m.nystrom.values)
	logdetQ = scale*log(sigma2) + logdetZ + sum(log(m.nystrom.values))
	yTQy / 2 + logdetQ / 2 + scale*log(2pi)
end

function nll(m::FullModel)
	refresh_cache!(m)
	n = size(m.X, 1)
	dot(m.y, m.CinvY) / 2 + logdet(m.C) / 2 + n*log(2pi)
end

function nll(m::FITCModel)
	sigma2 = exp(2m.lik.log_sigma)
	Kmn = kernel(m.k, m.u, m.X)
	Km = kernel(m.k, m.u)
	Knn = kernel(m.k, m.X, m.X)
	Lambda = diagm(diag(Knn - Kmn' * (Km \ Kmn)))
	C = Kmn' * (Km \ Kmn) + Lambda + sigma2 * eye(Lambda)
	n = size(m.y)[1]
	(dot(m.y,(C \ m.y))  .+ logdet(C) .+ n*log(2pi))/2
end

# function nll(m::LaplacianModel)
# 	refresh_cache!(m)
# 	n = size(m.X, 1)
# 	sigma2 = exp(2m.lik.log_sigma)
# 	scale = n - length(m.lambdas)
# 	Ssqrtlambda = Float64[S(m.k, sqrt(lambda)) for lambda in m.lambdas]
# 	LambdaInv = diagm(1./Ssqrtlambda)
# 	Z = (m.PhiTPhi + sigma2 * LambdaInv)
# 	logdetZ = logdet(Z)
# 	logdetQ = scale*log(sigma2) + logdetZ + sum(log(Ssqrtlambda))
# 	yTQy = (m.yTy - dot(m.PhiTy, Z \ m.PhiTy)) / sigma2
# 	yTQy / 2 + logdetQ / 2 + scale*log(2pi)
# end

include("negloglik_laplacian.jl")


include("graph.jl")
include("opinion.jl")

function main()
    (W, U, n) = read_edgelist("./data/raw/cloister.edgelist")
    m = 2
    normalize_edgelist!(W)
    normalize_edgelist!(U)
    y = generate_opinions(n, m)
    α, β, γ = generate_factors(n)
    M = Matrix(β .* W - γ .* U)
    Q = get_fundamental_matrix(β, W, γ, U)
    z∞ = get_equlibrium(α, Q, y)
    return y, Q, z∞
end

main()


using BarycentricInterpolation
using LinearAlgebra
using Test

# These tests arbitrarily use 5 and 6 degree polynomials (odd and even are used to
# ensure coverage)

equi₅ = Equispaced{5}()
equi₆ = Equispaced{6}()
cheb1₅ = Chebyshev1{5}()
cheb1₆ = Chebyshev1{6}()
cheb2₅ = Chebyshev2{5}()
cheb2₆ = Chebyshev2{6}()
leg₅ = Legendre{5}()
leg₆ = Legendre{6}()

x5 = [0.814723686393179,1.7205156234688,1.8475024397623,2.76087829590132,3.39323754212673,3.49077794712614]
x6 = [0.278498218867048,0.825379738072032,1.78288657350633,2.74777510870561,2.90538819038315,3.87598097214377,4.83314792038672]
arb₅ = ArbitraryPolynomial(x5)
arb₆ = ArbitraryPolynomial(x6)

@testset "Degree" begin
    @test degree(equi₅) == 5
    @test degree(equi₆) == 6
end

@testset "Nodes" begin
    @test nodes(equi₅) ≈ [-1,-0.6,-0.2,0.2,0.6,1]
    @test nodes(equi₆) ≈ [-1,-0.666666666666667,-0.333333333333333,0,0.333333333333333,0.666666666666667,1]

    @test nodes(cheb1₅) ≈ -[0.965925826289068,0.707106781186548,0.258819045102521,-0.258819045102521,-0.707106781186547,-0.965925826289068]
    @test nodes(cheb1₆) ≈ -[0.974927912181824,0.78183148246803,0.433883739117558,6.12323399573677e-17,-0.433883739117558,-0.781831482468029,-0.974927912181824]

    @test nodes(cheb2₅) ≈ -[1,0.809016994374947,0.309016994374947,-0.309016994374947,-0.809016994374947,-1]
    @test nodes(cheb2₆) ≈ -[1,0.866025403784439,0.5,6.12323399573677e-17,-0.5,-0.866025403784439,-1]

    @test nodes(leg₅) ≈ [-0.932469514203152,-0.661209386466264,-0.238619186083197,0.238619186083197,0.661209386466264,0.932469514203152]
    @test nodes(leg₆) ≈ [-0.949107912342758,-0.741531185599394,-0.405845151377397,0,0.405845151377397,0.741531185599394,0.949107912342758]

    @test nodes(arb₅) ≈ x5
    @test nodes(arb₆) ≈ x6

    @test nodes(Equispaced{5, Float64}) ≈ range(-1, stop=1, length=6)
    @test collect(nodes(Equispaced{5, Float32})) isa Vector{Float32}
end

@testset "Weights" begin
    @test weights(equi₅) ≈ [-1,5,-10,10,-5,1]
    @test weights(equi₆) ≈ [1,-6,15,-20,15,-6,1]

    @test weights(cheb1₅) ≈ -[0.258819045102521,-0.707106781186547,0.965925826289068,-0.965925826289068,0.707106781186548,-0.258819045102521]
    @test weights(cheb1₆) ≈ -[0.222520933956314,-0.623489801858733,0.900968867902419,-1,0.900968867902419,-0.623489801858734,0.222520933956314]

    @test weights(cheb2₅) ≈ -[0.5,-1,1,-1,1,-0.5]
    @test weights(cheb2₆) ≈ -[0.5,-1,1,-1,1,-1,0.5]

    leg₅w_arb = weights(ArbitraryPolynomial(nodes(leg₅)))
    leg₅w = weights(leg₅)
    leg₅w .*= leg₅w_arb[3] / leg₅w[3]
    @test leg₅w_arb ≈ leg₅w
    leg₆w_arb = weights(ArbitraryPolynomial(nodes(leg₆)))
    leg₆w = weights(leg₆)
    leg₆w .*= leg₆w_arb[3] / leg₆w[3]
    @test leg₆w_arb ≈ leg₆w

    @test weights(arb₅) ≈ [-0.079601734468871,2.82206348049722,-3.28654127480991,1.17155138701322,-2.4317744194857,1.80430256125404]
    @test weights(arb₆) ≈ [0.0114361913589559,-0.0390633084690722,0.100394094241955,-0.588770483980979,0.552871013562898,-0.0415349717143145,0.00466746500055649]

    @test collect(weights(Chebyshev2{5, Float32})) isa Vector{Float32}
end

xnew = [-0.514624351277159,0.285656117611641,0.427542456238857,0.849303738865132] # randomly generated

@testset "Interpolation matrix" begin
    @test interpolation_matrix(cheb1₅, xnew) ≈ [-0.0289845478236897 0.0959627734159381 -0.207065941756692 0.626076869404405 0.609098232472457 -0.0950873857124182;-0.0105623302174427 0.0465782216908936 0.999201970730136 -0.0492504664603377 0.0197735254445248 -0.00574092118777367;-0.0706583814898062 0.371759729030326 0.841448464571141 -0.206847346420822 0.0915972568975815 -0.0272997225884207;0.362858671520288 0.813049119299893 -0.267459002051222 0.142520710886204 -0.0742818875343108 0.0233123878791484][:, end:-1:1]
    @test interpolation_matrix(cheb1₆, xnew) ≈ [0.0127857533810382 -0.0411607067288632 0.0812980146502686 -0.166310956920905 0.955056229650331 0.199706640020739 -0.0413749740526084;0.041385687119839 -0.161088285741265 0.779202698741526 0.448772538954965 -0.160518366877509 0.0748749109083769 -0.0226291831059332;0.00285521790063521 -0.0123604217994083 0.997916710684666 0.0164279128299457 -0.00734604100263827 0.00362101552128418 -0.00111439413448407;0.184921983997812 0.964706061561875 -0.226419209876787 0.122921499942711 -0.0733011112746609 0.0399052720352278 -0.0127344963861772][:, end:-1:1]
    @test interpolation_matrix(cheb1₅, nodes(cheb1₅)) ≈ I
    @test interpolation_matrix(cheb1₆, nodes(cheb1₆)) ≈ I
end

@testset "Interpolate" begin
    c1 = Chebyshev1{10}()
    @test interpolate(c1, sin.(nodes(c1)), xnew) ≈ sin.(xnew)
    c2 = Chebyshev2{20}(0, 10)
    x = 10*rand(5)
    @test interpolate(c2, sin.(nodes(c2)), x) ≈ sin.(x)
end

@testset "Differentiation matrix" begin
    @test differentiation_matrix(cheb1₅) ≈ -[7.20976852010751 -10.5558337350587 5.27791686752937 -3.04720672422855 1.63299316185545 -0.517638090205042;1.4142135623731 0.707106781186547 -3.04720672422855 1.4142135623731 -0.707106781186548 0.218779599482357;-0.378937381963012 1.63299316185545 0.138700708242031 -1.93185165257814 0.757874763926024 -0.218779599482357;0.218779599482357 -0.757874763926024 1.93185165257814 -0.138700708242031 -1.63299316185545 0.378937381963012;-0.218779599482357 0.707106781186547 -1.4142135623731 3.04720672422855 -0.707106781186545 -1.4142135623731;0.517638090205041 -1.63299316185545 3.04720672422854 -5.27791686752936 10.5558337350587 -7.2097685201075]
    @test differentiation_matrix(cheb1₆) ≈ -[9.84466088119818 -14.5105621059791 7.48352452737972 -4.60952974192497 2.87399478545474 -1.59494677776481 0.512858431636276;1.84827792218219 1.0055981139743 -4.15304279314467 2.05143372654511 -1.18863516903897 0.639524003844967 -0.203155804362921;-0.4564869487803 1.98886610088761 0.267253750914115 -2.55809601537987 1.15238243548124 -0.569229914492255 0.17531059136945;0.22824347439015 -0.797473388882404 2.07652139657234 -1.02695629777827e-15 -2.07652139657234 0.797473388882405 -0.22824347439015;-0.175310591369451 0.569229914492254 -1.15238243548124 2.55809601537987 -0.26725375091411 -1.98886610088761 0.456486948780299;0.203155804362921 -0.639524003844966 1.18863516903897 -2.05143372654511 4.15304279314467 -1.00559811397431 -1.84827792218218;-0.512858431636278 1.59494677776481 -2.87399478545474 4.60952974192498 -7.48352452737972 14.5105621059791 -9.84466088119817]
    c1 = Chebyshev2{20}(0, 10)
    x = 10*rand(5)
    D = differentiation_matrix(c1)
    @test interpolate(c1, D*sin.(nodes(c1)), x) ≈ cos.(x)
end


abstract Constraint
abstract ConvexConstraint <: Constraint
immutable BoxConstraint <: ConvexConstraint
    lower
    upper
end

immutable BallConstraint <: ConvexConstraint
    center # center
    radius # radius
    p # l^p norm
end

immutable EqualityConstraint <: Constraint
    c
    Jc!
    cJc!
end

function EqualityConstraint(c, Jc!)
    function cJc!(x::Array, storage::Array)
        Jc!(x, storage)
        return c(x)
    end
    return EqualityConstraint(c, Jc!, cJc!)
end

type AugmentedLagrangian{Tl<:Union{Real, Vector}}
   F::DifferentiableFunction
   C::EqualityConstraint
   lambda::Tl
   mu::Real
   Dc::Matrix
end

typealias AL AugmentedLagrangian
function AugmentedLagrangian(F::DifferentiableFunction,
                 C::EqualityConstraint,
                 x0::AbstractVector;
                 lambda=:auto, mu=10.0)
   # TODO: maybe not ideal to evaluate F, C here, maybe better to set
   # dimensions when C is first called??? But see #282
   C0 = C.c(x0)
   dim_x = length(x0)
   dim_c = length(C0)
   if lambda == :auto
      lambda = - mu * C0
   end
   return AugmentedLagrangian(F, C, lambda, mu, zeros(dim_c, dim_x))
end


M_example = M(model,u₀,u₀_type,tspan,[],[10],15,Rodas5)

rates_example = ks_γϵ
sol = solve_prob(M_example,ks_γϵ;tᵢ=[])

ŷ = sum(sol,dims=1)
mean_t = mean_lifetime_linear(sol)

L = ŷ / mean_t

pyplot()
# plot(sol)
plot(sol.t,ŷ',linewidth=5,label="P")
plot!(sol.t,L',linewidth=5,label="pdf")
xlabel!("t")
ylabel!("distribution")


@testset "extrema_series" begin
    @testset "test problems" begin
        demo_problems = demo_problem.(eachindex(demo_problems_definitions), true)
        for (f, (a, b), min) in demo_problems
            min1, _, y1 = ArbExtras.extrema_series(f, a, b)
            _, min2, y2 = .-ArbExtras.extrema_series(x -> -f(x), a, b)
            min3, y3 = ArbExtras.minimum_series(f, a, b)
            min4, y4 = .-ArbExtras.maximum_series(x -> -f(x), a, b)

            @test isequal(min1, min2)
            @test isequal(min2, min3)
            @test isequal(min3, min4)

            @test isequal(y1, y2)
            @test isequal(y3, y3)
            @test isequal(y3, y4)

            @test Arblib.overlaps(min1, min)

            @test Arblib.overlaps(y1, f(Arb(a + b) / 2))
        end
    end

    @testset "special cases" begin
        # Thin interval
        @test ArbExtras.extrema_series(x -> -2x, Arf(2), Arf(2)) == (-4, -4, -4)
        @test ArbExtras.minimum_series(x -> -2x, Arf(2), Arf(2)) == (-4, -4)
        @test ArbExtras.maximum_series(x -> -2x, Arf(2), Arf(2)) == (-4, -4)
        @test ArbExtras.extrema_series(x -> -2x, Arf(2), Arf(2), abs_value = true) ==
              (4, 4, 4)
        @test ArbExtras.minimum_series(x -> -2x, Arf(2), Arf(2), abs_value = true) == (4, 4)
        @test ArbExtras.maximum_series(x -> -2x, Arf(2), Arf(2), abs_value = true) == (4, 4)

        # Non-ordered endpoints
        @test_throws ArgumentError ArbExtras.extrema_series(identity, Arf(2), Arf(1))
        @test_throws ArgumentError ArbExtras.minimum_series(identity, Arf(2), Arf(1))
        @test_throws ArgumentError ArbExtras.maximum_series(identity, Arf(2), Arf(1))

        # Infinite endpoints
        @test_throws ArgumentError ArbExtras.extrema_series(identity, Arf(2), Arf(Inf))
        @test_throws ArgumentError ArbExtras.minimum_series(identity, Arf(2), Arf(Inf))
        @test_throws ArgumentError ArbExtras.maximum_series(identity, Arf(2), Arf(Inf))

        # Monotonic on interval
        @test isequal(ArbExtras.extrema_series(identity, Arf(-2), Arf(4)), (-2, 4, NaN))
        @test isequal(ArbExtras.minimum_series(identity, Arf(-2), Arf(4)), (-2, NaN))
        @test isequal(ArbExtras.maximum_series(identity, Arf(-2), Arf(4)), (4, NaN))

        @test isequal(
            ArbExtras.extrema_series(identity, Arf(-2), Arf(4), abs_value = true),
            (0, 4, NaN),
        )
        @test isequal(
            ArbExtras.extrema_series(identity, Arf(-3), Arf(-1), abs_value = true),
            (1, 3, NaN),
        )
        @test isequal(
            ArbExtras.minimum_series(identity, Arf(-2), Arf(4), abs_value = true),
            (0, NaN),
        )
        @test isequal(
            ArbExtras.minimum_series(identity, Arf(-3), Arf(-1), abs_value = true),
            (1, NaN),
        )
        @test isequal(
            ArbExtras.maximum_series(identity, Arf(-6), Arf(4), abs_value = true),
            (6, NaN),
        )

        # Check that it gives a correct result even if the
        # computations with the endpoints cannot be done exactly. I
        # have not been able to come up with an example for which we
        # get a wrong result without adding the endpoints explicitly.
        # If we are working with such large numbers the enclosure we
        # get is typically very large and includes the true result
        # anyway.

        @test all(
            Arblib.overlaps.(
                ArbExtras.extrema_series(x -> x^2, Arf(1), Arf(1e300))[1:2],
                (Arb(-1), Arb(1e300)^2),
            ),
        )
        @test all(
            Arblib.overlaps.(
                ArbExtras.extrema_series(x -> x^2, Arf(-1e300), Arf(-1))[1:2],
                (Arb(-1), Arb(1e300)^2),
            ),
        )
        @test Arblib.overlaps(
            ArbExtras.minimum_series(x -> x^2, Arf(1), Arf(1e300))[1],
            Arb(-1),
        )
        @test Arblib.overlaps(
            ArbExtras.minimum_series(x -> x^2, Arf(-1e300), Arf(-1))[1],
            Arb(-1),
        )
        @test Arblib.overlaps(
            ArbExtras.maximum_series(x -> -x^2, Arf(1), Arf(1e300))[1],
            Arb(1),
        )
        @test Arblib.overlaps(
            ArbExtras.maximum_series(x -> -x^2, Arf(-1e300), Arf(-1))[1],
            Arb(1),
        )

        # Non-finite restterm
        @test all(isnan, ArbExtras.extrema_series(x -> inv(x - sin(x)), Arf(0.1), Arf(1)))
        @test all(isnan, ArbExtras.minimum_series(x -> inv(x - sin(x)), Arf(0.1), Arf(1)))
        @test all(isnan, ArbExtras.maximum_series(x -> inv(x - sin(x)), Arf(0.1), Arf(1)))
    end
end


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module RationalTrees

using Nemo, NumberTheory, Counts

export ModuleRationalTrees
export EuclidTree, CalkinWilfTree, SchinzelSierpinskiEncoding

"""

Rational trees as understood here are binary trees enumerating the positive or
nonnegative rational numbers. Examples are the Euclid tree, the Kepler tree and the
Stern-Brocot tree (a.k.a. Farey tree). They are closely related to binary partitions
and to Stern's diatomic sequence or Dijkstra's fusc function.

Malter, Schleicher, Zagier, [New looks at old number theory](https://pdfs.semanticscholar.org/6d28/dcef911dd91f47e6ca4bd2c564c1f3099a05.pdf), Amer. Math. Monthly, 120(3), 2013, pp. 243-264.

* EuclidTree, CalkinWilfTree, SchinzelSierpinskiEncoding
"""
const ModuleRationalTrees = ""

# http://oeis.org/wiki/User:Peter_Luschny/SchinzelSierpinskiConjectureAndCalkinWilfTree
# A294442 A294446

"""

```
julia> for n ∈ 1:4 Println(EuclidTree(n)) end
[1//1]
[1//2, 2//1]
[1//3, 3//2, 2//3, 3//1]
[1//4, 4//3, 3//5, 5//2, 2//5, 5//3, 3//4, 4//1]
```
"""
function EuclidTree(n)

    function DijkstraFusc(m)
        a, b, k = 1, 0, m
        while k > 0
            k % 2 == 1 ? b += a : a += b
            k = k >> 1
        end
        b
    end

    DF = [DijkstraFusc(k) for k ∈ 2^(n-1):2^n]
    [fmpq(DF[j], DF[j+1]) for j ∈ 1:2^(n-1)]
end

"""

Alias for the (much better named) EuclidTree. See Malter, Schleicher, Zagier, New looks at old number theory, Amer. Math. Monthly, 120(3), 2013, pp. 243-264.
"""
CalkinWilfTree(n) = EuclidTree(n)

"""

Return the Schinzel-Sierpinski encoding of the positive rational number r.

```
julia> for n ∈ 1:4 println([SchinzelSierpinski(l) for l ∈ EuclidTree(n)]) end
[1//1]
[2//5, 5//2]
[3//11, 5//3, 3//5, 11//3]
[2//11, 3//2, 11//19, 19//7, 7//19, 19//11, 2//3, 11//2]
```
"""
function SchinzelSierpinskiEncoding(l, searchLimit = 500000)

    a, b = numerator(l), denominator(l)
    sgn = a < b ? -1 : 1
    p, q = 1, 2

    while q < searchLimit
        r = a * (q + 1)
        r % b == 0 && (p = div(r, b) - 1)
        isPrime(p) && return fmpq(p, q)
        q = NextPrime(q)
    end
    warn("Search limit reached for ", l)
    return 0
end

#START-TEST-########################################################

using Test, SeqUtils, PrimesIterator

function test()
    @testset "EuclidTree" begin
        S = [numerator(sum(r for r ∈ EuclidTree(n))) for n ∈ 1:9]
        L = fmpz[1, 5, 11, 23, 47, 95, 191, 383, 767] # A052940
        @test all(S[1:9] .== L[1:9])
    end
end

function demo()
    println("\nEuclidTree")
    for n ∈ 1:5
        Println(EuclidTree(n))
    end
    println()

    println("\nSchinzelSierpinskiEncoding of the Euclid tree")
    for n ∈ 1:5
        Println([SchinzelSierpinskiEncoding(l) for l ∈ EuclidTree(n)])
        # println(sum([SchinzelSierpinskiEncoding(l) for l ∈ EuclidTree(n)]))
    end
    println()

    println("\nSchinzelSierpinskiEncoding of primes")
    Println([SchinzelSierpinskiEncoding(fmpq(p, 1)) for p ∈ Primes(100)])
    println()
end

"""

EuclidTree(100) ::
    0.000005 seconds (4 allocations: 224 bytes)
[SchinzelSierpinskiEncoding(l) for l ∈ EuclidTree(100)] ::
    0.000007 seconds (6 allocations: 320 bytes)
"""
function perf()
    GC.gc()
    @time EuclidTree(100)
    GC.gc()
    @time [SchinzelSierpinskiEncoding(l) for l ∈ EuclidTree(100)]
end

function main()
    test()
    demo()
    perf()
end

main()

end # module

#=
# A062251 := n -> SchinzelSierpinskiPrimes(n,"numer"):
# A062251 = 2, 5, 11, 11, 19, 17, 41, 23, 53, 29, 43, 47, 103, 41,
# 59, 47, 67, 53, 113, 59, 83, 131, 137, 71, 149, 103, 107, 83,
# A060324 := n -> SchinzelSierpinskiPrimes(n,"denom"):
# A060324 = 2, 2, 3, 2, 3, 2, 5, 2, 5, 2, 3, 3, 7, 2, 3, 2,
# 3, 2, 5, 2, 3, 5, 5, 2, 5, 3, 3, 2, 5, 2, 13, 3
=#


__precompile__(false)
module SymbolicControlSystems
using LinearAlgebra
using ControlSystems, SymPy, Latexify
import Symbolics
using InteractiveUtils

export sp,
    sym2num,
    latextf,
    Sym,
    @vars,
    simplify,
    tustin,
    doubleeuler,
    s,
    z,
    print_c_array,
    show_construction

# export symbolics2sym

const sp = SymPy.PyCall.PyNULL()
const s = SymPy.Sym("s")
const z = SymPy.Sym("z")

function __init__()
    # global const s = Sym("s")
    # global const z = Sym("z")
    copy!(sp, SymPy.sympy)
end

function SymPy.Sym(sys::StateSpace{<:Any,Sym})
    A, B, C, D = ControlSystems.ssdata(sys)
    if isdiscrete(sys)
        (C*inv(z * I(size(A, 1)) - A)*B+D)[1]
    else
        (C*inv(s * I(size(A, 1)) - A)*B+D)[1]
    end
end

function SymPy.Sym(sys::TransferFunction)
    if isdiscrete(sys)
        ssys = sp.Poly(numvec(sys)[], z) / sp.Poly(denvec(sys)[], z)
    else
        ssys = sp.Poly(numvec(sys)[], s) / sp.Poly(denvec(sys)[], s)
    end
end

"""
    ControlSystems.tf(sys::Sym, h = nothing)

Convert a symbolic, rational expression into a transfer function. `h` denotes the sample time if the system is discrete.
"""
function ControlSystems.tf(sys::Sym, h = nothing)
    n, d = sp.fraction(sys)
    d = sp.Poly(d)
    # d = d.monic() # Don't do this here
    n = n isa Number ? n : sp.Poly(n)
    if h === nothing
        tf(expand_coeffs(n, s), expand_coeffs(d, s))
    else
        tf(expand_coeffs(n, z), expand_coeffs(d, z), h)
    end
end

function expand_coeffs(n, var; numeric = false)
    n = sp.Poly(n, var)
    deg = n.degree() |> Float64 |> Int
    c = n.all_coeffs() # This fails if the coeffs are symbolic
    numeric && (c = Float64.(c))
    [c; zeros(eltype(c), deg - length(c) + 1)]
end
expand_coeffs(n::Real, args...; numeric = false) = n

function ControlSystems.minreal(sys::TransferFunction{<:Any,<:ControlSystems.SisoTf{<:Sym}})
    Sym(sys) |> simplify |> tf
end

function ControlSystems.tf(sys::StateSpace{<:Any,Sym})
    n, p = simplify.(sp.Poly.(simplify.(sp.fraction(simplify(Sym(sys)))), s))
    tf(simplify(n / p))
end

function Symbolics.Num(sys::StateSpace{<:Any,Symbolics.Num})
    A,B,C,D = ssdata(sys)
    if isdiscrete(sys)
        sz = Symbolics.@variables z
    else
        sz = Symbolics.@variables s
    end
    sz = only(sz)
    ex = (C*inv(sz * I(size(A, 1)) - A)*B + D)
    # if length(ex) == 1
    #     return ex[1] # don't do this as this complicates later function building
    # else
    #     return ex
    # end
    # ex = Symbolics.expand(ex)
    # ds = Symbolics.degree(ex, sz)
    # Symbolics.get_variables(ex)
    # Symbolics.polynormalize(ex)
end

function ControlSystems.tf(sys::StateSpace{<:Any,Symbolics.Num})
    n, p = simplify.(sp.Poly.(simplify.(sp.fraction(simplify(Sym(sys)))), s))
    tf(simplify(n / p))
end

function ControlSystems.minreal(sys::StateSpace{<:Any,Symbolics.Num})
    # sys |> Symbolics.Num .|> Symbolics.symbolics_to_sympy .|> sp.simplify
    nsys = Symbolics.Num(sys)
    nsys = Symbolics.simplify.(nsys)
    nsys = Symbolics.simplify_fractions.(nsys)
end

sym2num(P::TransferFunction, args...) = sym2num(Sym(P), args...)

"""
    sym2num(G, h::Real, pairs::Pair...)

Replace symbols by numbers. `h` indicates the sample rate if the system is discrete.
"""
function sym2num(P::Sym, h::Real, pairs::Pair...)
    P.has(s) && error(
        "Found `z` in symbolic expression, provide sample time as second argument to `sym2num`",
    )
    for (sym, val) in pairs
        P = subs(P, (sym, val))
    end
    n, d = sp.Poly.(sp.fraction(P), z)
    nn = float.(expand_coeffs(n, z))
    nd = float.(expand_coeffs(d, z))
    if nd[1] != 1 # make monic
        nn, nd = monic(nn, nd)
    end
    G = tf(nn, nd isa Number ? [nd] : nd, h)
    G
end

"""
    sym2num(G, pairs::Pair...)

Replace symbols by numbers.
"""
function sym2num(P::Sym, pairs::Pair...)
    # for i = 1:2
    P.has(z) && error(
        "Found `z` in symbolic expression, provide sample time as second argument to `sym2num`",
    )
    P.has(s) || error("Found no `s` in symbolic expression, provide")
    for (sym, val) in pairs
        P = subs(P, (sym, val))
    end
    n, d = sp.Poly.(sp.fraction(P), s)
    nn = float.(expand_coeffs(n, s))
    nd = float.(expand_coeffs(d, s))
    if nd[1] != 1 # make monic
        nn, nd = monic(nn, nd)
    end
    G = tf(nn, nd isa Number ? [nd] : nd)
    G
end

"""
    tustin(G::LTISystem, h)

Discretize `G` using the Tustin (bi-linear) method for sample time `h`.
"""
function tustin(G::LTISystem, h)
    iscontinuous(G) || throw(DomainError("Cannot discretize a discrete-domain system."))
    tustin(Sym(G), h)
end

function tustin(Gs::Sym, h)
    hh = Sym("hh")
    Gt = Gs.subs(s, hh * (z - 1) / (z + 1)) |> simplify
    @info "Perforing pole-zero cancellation (this might take a while)"
    @time Gt = sp.cancel(Gt) # Simplify takes forever with floating point values
    @time Gt = sp.collect(Gt, z) # Simplify takes forever with floating point values
    Gt = Gt.subs(hh, (2 / h))
    Gtf = tf(Gt, h)
    n, d = numvec(Gtf)[], denvec(Gtf)[]
    d[1] == 1 && (return Gtf)
    n, d = monic(n, d)
    tf(n, d, h)
end

"""
    doubleeuler(sys::AbstractStateSpace{<:Continuous}, Ts)

Discretize `sys` with a second-order approximation to the exponential map (Forward Euler is a first-order approximation). This is useful if the symbolic expressions for the true ZoH-discretization becomes too complicated. A second-order approximation is in many cases indistinguishable from the true ZoH-discretization.

# Example
```julia
w = 2pi .* exp10.(LinRange(-1, log10(25), 400))
sys = ssrand(1,1,4)
bodeplot(sys, w, lab="cont")
bodeplot!(c2d(sys, 0.001, :fwdeuler), w, label="fwdeuler")
bodeplot!(doubleeuler(sys, 0.001), w, label="double euler")
```
"""
function doubleeuler(sys::AbstractStateSpace{<:Continuous}, Ts)
    A, B, C, D = ssdata(sys)
    T = promote_type(eltype.((A, B, C, D))...)
    ny, nu = size(sys)
    nx = sys.nx
    Ad, Bd, Cd, Dd = (I + Ts * A + Ts^2 / 2 * A^2), Ts * B, C, D
    ss(Ad, Bd, Cd, Dd, Ts)
end

"""
    ccode(G; simplify = identity, cse = true)

Return a string with C-code for filtering a signal `u` through `G`.

# Arguments:
- `G`: A linear system
- `simplify`: A function for symbolic simplification. You may try `Sympy.simplify`, but for large systems, this will take a long time to compute.
- `cse`: Perform common subexpression elimination. This generally improvems the performance of the generated code.
"""
function ccode(G::TransferFunction; simplify = identity, cse = true)
    P = Sym(G)
    P.has(z) || error("Did not find `z` in symbolic expression")
    P.has(s) && error("Found `s` in symbolic expression, provide expression in `z`")
    n, d = numvec(G)[], denvec(G)[]
    if d[1] != 1
        @info "Denominator polynomial not monic, dividing by the leading coefficient."
        n = simplify.(n ./ d[1])
        d = simplify.(d ./ d[1])
    end
    @info "Calling free_symbols"
    @info vars = P.free_symbols
    vars.remove(z)
    vars = collect(vars)
    vars = sort(vars, by = string)
    var_str = ""
    for var in vars
        var_str *= ", double $(var)"
    end

    nu, ny = length(n), length(d)
    code = """
#include <stdio.h>\n
#include <math.h>\n
double transfer_function(double ui$(var_str)) {
    static double u[$(nu)] = {0};
    static double y[$(ny)] = {0};
    int i;
    for (i=$(nu-1); i > 0; --i) {
        u[i] = u[i-1];
    }
    u[0] = ui;
    for (i=$(ny-1); i > 0; --i) {
        y[i] = y[i-1];
    }
    y[0] = 0;
"""
    if cse
        @info "Finding common subexpressions"
        subex, final = sp.cse([n; d])
        if length(final) == 1
            final = final[]
        end
        n = final[1:length(n)]
        d = final[length(n)+1:end]
        for se in subex
            code *= "    double $(se[1]) = $(sp.ccode(se[2]));\n"
        end
    end
    for (i, n) in enumerate(n)
        @info "Writing numerator term $i/$(length(n))"
        code *= "    y[0] += ($(sp.ccode(n)))*u[$(i-1)];\n"
    end
    for (i, n) in enumerate(d[2:end])
        @info "Writing denominator term $i/$(length(d)-1)"
        code *= "    y[0] += ($(sp.ccode(-n)))*y[$(i)];\n"
    end
    code *= "    return y[0];\n}"
    println(code)
    try
        clipboard(code)
    catch
    end
    code
end

function ccode(sys::StateSpace{<:Discrete}; cse = true, function_name = "transfer_function")
    sys.nu == 1 || throw(ArgumentError("Multiple input not yet supported"))
    nx = sys.nx
    ny = sys.ny
    u = Sym("u")
    x = [Sym("x[$(i-1)]") for i = 1:nx]
    # @show P
    if ControlSystems.numeric_type(sys) <: SymPy.Sym
        P = Sym(sys)
        vars = P.free_symbols
        vars.remove(z)
        vars = collect(vars)
        vars = sort(vars, by = string)
        var_str = ""
        for var in vars
            var_str *= ", double $(var)"
        end
    else
        var_str = ""
    end
    x1 = zeros(Sym, nx) # workaround for strange bug with undefinied referece appearing in Pluto only
    y = zeros(Sym, ny)
    @show x1 = mul!(x1, sys.A, x) + sys.B * u
    @show y = mul!(y, sys.C, x) + sys.D * u
    # @show y = sp.collect.(y, x)

    code = """
#include <stdio.h>\n
#include <math.h>\n
void $(function_name)(double *y, double u$(var_str)) {
    static double x[$(nx)] = {0};  // Current state
    double xp[$(nx)] = {0};        // Next state
    int i;
"""
    if cse
        @info "Finding common subexpressions"
        subex, final = sp.cse([x1; y])
        x1 = final[][1:length(x1)]
        y = final[][length(x1)+1:end]
        for se in subex
            code *= "    double $(se[1]) = $(sp.ccode(se[2]));\n"
        end
    end
    code *= "\n    // Advance the state xp = Ax + Bu\n"
    for (i, n) in enumerate(x1)
        code *= "    xp[$(i-1)] = ($(sp.ccode(n)));\n"
    end
    code *= """

        // Accumulate the output y = C*x + D*u
    """

    for (i, n) in enumerate(y)
        code *= "    y[$(i-1)] = ($(sp.ccode(n)));\n"
    end
    code *= """

        // Make the predicted state the current state
        for (i=0; i < $(nx); ++i) {
            x[i] = xp[i];
        }
    """
    code *= "\n}"

    println(code)
    try
        clipboard(code)
    catch
    end
    code
end

function structured_text(code)

    code = replace(code, "pow" => "EXPT")
    code = replace(code, "static " => "")

    code = replace(code, r"double (\w+) =" => s"\1 : LREAL :=")
    code = replace(code, r"double (\w+)" => s"\1 : LREAL")
    code = replace(code, r"(.{1,20}?) \+= " => s"\1 := \1 + ")
    code = replace(code, " = " => " := ")
    println(code)
    try
        clipboard(code)
    catch
    end
    code
end

latextf(x::TransferFunction, args...) = latextf(Sym(x), args...)

Latexify.Latexify(x::TransferFunction) = latextf(x, args...)

function monic(n, d)
    n = n ./ d[1]
    d = d ./ d[1]
    if n isa Number
        n = [n]
    end
    n, d
end

"""
    latextf(x, mon = true)

Return a latex string representing `x`. If `mon`, then the denominator polynomial will be made monic (leading coefficient = 1).
"""
function latextf(x::Sym, mon = true)
    var = x.has(z) ? z : s
    n, d = sp.fraction(x)
    if mon
        n, d = monic(n, d)
    end
    n, d = sp.collect.((n, d), var)
    str = "\$\\dfrac{$((n))}{$((d))}\$"
    str = replace(str, '*' => "")
    try
        clipboard(code)
    catch
    end
    str
end

show_construction(sys::LTISystem; kwargs...) = show_construction(stdout, sys; kwargs...)
function show_construction(io::IO, sys::LTISystem; letb = true)
    # sys = StateSpace(sys)
    letb && println(io, "sys = let")
    prestr = letb ? "    " : ""
    println(io, prestr*"tempA = ", sys.A)
    println(io, prestr*"tempB = ", sys.B)
    println(io, prestr*"tempC = ", sys.C)
    println(io, prestr*"tempD = ", sys.D)
    letb || print(io, "sys = ")
    if isdiscrete(sys)
        println(io, prestr*"ss(tempA, tempB, tempC, tempD, $(sys.Ts))")
    else
        println(io, prestr*"ss(tempA, tempB, tempC, tempD)")
    end
    letb && println(io, "end")
    nothing
end

"""
sys2vec = @(sys) [
        size(sys.A,1);
        size(sys.B,2);
        size(sys.C,1);
        sys.A(:);
        sys.B(:);
        sys.C(:);
        sys.D(:)
    ]
"""
function vec2sys(v, ts = nothing)
    nx = Int(v[1])
    nu = Int(v[2])
    ny = Int(v[3])
    ai = (1:nx^2) .+ 3
    bi = (1:nx*nu) .+ ai[end]
    ci = (1:nx*ny) .+ bi[end]
    di = (1:nu*ny) .+ ci[end]
    A = reshape(v[ai], nx, nx)
    B = reshape(v[bi], nx, nu)
    C = reshape(v[ci], ny, nx)
    D = reshape(v[di], ny, nu)
    ts === nothing ? ss(A, B, C, D) : ss(A, B, C, D, ts)
end

print_c_array(a::AbstractArray, args...; kwargs...) =
    print_c_array(stdout, a, args...; kwargs...)
function print_c_array(io, a::AbstractVector, name = "vec"; cse = false, s = "", struct_name=nothing)
    l = length(a)
    a = float.(a)
    if cse
        a = write_array_cse(io, a, name, s)
    end
    if struct_name === nothing
        s == "" && println(io, "double $name[$l];")
        struct_name = ""
    else
        struct_name = struct_name*"->"
    end
    for i = 1:l
        println(io, "$(struct_name)" * s * "$name[$(i-1)] = $(a[i]);")
    end
    println(io)
end
function print_c_array(io, a::AbstractMatrix, name = "mat"; cse = false, s = "", struct_name=nothing)
    r, c = size(a)
    a = float.(a)
    if cse
        a = write_array_cse(io, a, name, s)
    end
    if struct_name === nothing
        s == "" && println(io, "double $name[$r][$c];")
        struct_name = ""
    else
        struct_name = struct_name*"->"
    end
    for i = 1:r, j = 1:c
        println(io, "$(struct_name)" * s * "$name[$(i-1)][$(j-1)] = $(a[i,j]);")
    end
    println(io)
end

function print_c_array(io, a::AbstractArray{<:Any,3}, name = "array"; cse = false, s = "", struct_name=nothing)
    r, c, d = size(a)
    a = float.(a)
    if cse
        a = write_array_cse(io, a, name, s)
    end
    if struct_name === nothing
        s == "" && println(io, "double $name[$r][$c][$d];")
        struct_name = ""
    else
        struct_name = struct_name*"->"
    end
    for i = 1:r, j = 1:c, k = 1:d
        println(io, "$(struct_name)" * s * "$name[$(i-1)][$(j-1)][$(k-1)] = $(a[i,j,k]);")
    end
    println(io)
end

# Interpolation
function print_c_array(
    io,
    a::Vector{<:AbstractArray},
    t::AbstractVector,
    name          = "mat";
    cse           = false,
    s             = "",
    print_vector  = true,
    print_logic   = true,
    struct_name::Union{Nothing, String} = nothing,
    struct_type   = nothing,
    ivecname = name * "_interp_vect",
)
    length(a) == length(t) || throw(ArgumentError("length mismatch between a and t"))
    print_vector && print_c_array(io, t, ivecname; cse, s, struct_name)
    # for i in eachindex(t)
    #     iname = name*"_$(i-1)"
    #     print_c_array(io, a[i], iname; cse, s)
    # end
    a = cat(a..., dims = 3)

    if struct_name !== nothing
        struct_type isa String || throw(ArgumentError("Expected a struct type since you provided struct name"))
        print_c_array(io, a, name * "_interp"; cse, s, struct_name)
        println(io, "void interpolate_$(name)($(struct_type) *$(struct_name), double t) {")
        ivecname = struct_name*"->"*ivecname
    else
        print_c_array(io, a, name * "_interp"; cse, s)
        println(
        io,
        "void interpolate_$(name)(double *$(name), double *$(name)_interp, double *$(ivecname), double t) {"
        )
    end

    println(
        io,
        """
    int k = 0;
    for (k = 0; k < $(length(t)-2); ++k) { // Loop to find correct interval
        if (t < $(ivecname)[k+1]) {
            break;
        }
    } // t is now between indices k and k+1 modolu edge cases below
    double t0 = $(ivecname)[k];
    double t1 = $(ivecname)[k+1];
    double l = t1 - t0;      // Length of interval between points
    double alpha = (t-t0)/l; // Between 0 and 1, proportion of k+1 point
    if (t < $(ivecname)[0]) { // edge cases
        alpha = 0;
    }else if (t > $(ivecname)[$(length(t)-1)]) {
        alpha = 1;
    }
$(interpolator_string(a, name, struct_name, struct_type))}
""",
    )
end

function interpolator_string(a::AbstractArray{<:Any,3}, name, struct_name, struct_type)
    if struct_name === nothing
        """
            for (int i = 0; i < $(size(a, 1)); ++i) {     // $name has dimensions $((size(a)[1:2]))
                for (int j = 0; j < $(size(a, 2)); ++j) { // $(name)_interp has dimensions $((size(a)))
                    $(name)[i][j] = (1-alpha)*$(name)_interp[i][j][k] + alpha*$(name)_interp[i][j][k+1];
                }
            }
        """
    else
        name = struct_name*"->"*name
        """
            for (int i = 0; i < $(size(a, 1)); ++i) {     // $name has dimensions $((size(a)[1:2]))
                for (int j = 0; j < $(size(a, 2)); ++j) { // $(name)_interp has dimensions $((size(a)))
                    $(name)[i][j] = (1-alpha)*$(name)_interp[i][j][k] + alpha*$(name)_interp[i][j][k+1];
                }
            }
        """
    end
end

# StateSpace
function print_c_array(io, sys::AbstractStateSpace, args...; s = "", en = "", struct_name=nothing)
    print_c_array(io, sys.A, "A" * en; s, struct_name)
    print_c_array(io, sys.B, "B" * en; s, struct_name)
    print_c_array(io, sys.C, "C" * en; s, struct_name)
    print_c_array(io, sys.D, "D" * en; s, struct_name)
end

# Interpolated StateSpace
function print_c_array(
    io,
    sys::Vector{<:AbstractStateSpace},
    t::AbstractVector,
    name = "sys";
    cse = false,
    s = "",
    en = "",
    struct_name::Union{Nothing, String} = nothing,
    struct_type   = nothing,
)
    length(sys) == length(t) || throw(ArgumentError("length mismatch between a and t"))
    ivecname = name * en * "_interp_vect"
    print_c_array(io, t, ivecname; cse, s, struct_name)
    for p in (:A, :B, :C, :D)
        A = getproperty.(sys, p)
        print_c_array(io, A, t, name * "_" * string(p) * en; cse, s, print_vector = false, struct_name, struct_type, ivecname)
    end
end

function write_array_cse(io, a, name = "x", s = "")
    subex, final = sp.cse(a, symbols = [SymPy.Sym(name * string(i)) for i = 1:200])
    new_a = final[]
    for se in subex
        println(io, "double $(se[1]) = $(sp.ccode(se[2]));")
    end
    new_a
end

end


using Revise
using Soss

m = @model x begin
    α ~ Cauchy()
    β ~ Normal()
    σ ~ HalfNormal()
    y ~ For(x) do xj
        Normal(α + β * xj, σ)
    end
end;

truth = rand(m(x=randn(10)));
using Plots
xx = range(extrema(truth.x)...,length=100)
plot(xx, truth.α .+ truth.β .* xx, legend=false)
scatter!(truth.x,truth.y)

post = dynamicHMC(m(X=truth.X), (y=truth.y,));

particles(post)

pred = predictive(m,:β)

function fit(m::Model, verbosity::Integer, X, y)
    fitresult = dynamicHMC(m(X=X), (y=y,))
    cache = nothing
    report = nothing
    (fitresult, cache, report)
end

struct MLJpredictor <: Distributions.Sampleable{Univariate, Continuous}
    pred :: Soss.Model
    Xrow :: Matrix
    βs   :: Vector{NamedTuple}
end

function Base.rand(p::MLJpredictor)
    args = merge(rand(p.βs), (X=p.Xrow,))
    rand(p.pred(args)).y
end

function predict(m, fitresult, Xnew)
    pred = predictive(m, setdiff(variables(m),[:X,:y])...)
    map(eachrow(Xnew)) do x
        X = reshape(x, 1, :)
        MLJpredictor(pred, X, fitresult)
    end
end

fitresult = fit(m, 0, truth.X, truth.y)[1];
p = predict(m,fitresult, rand(4,3));

map(rand, p)

# mutable struct MLJmodel <: MLJ.Probabilistic
#     sampler #function that can do MCMC
#     m :: Soss.Model
# end

# MLJBase.fit(model::SomeSupervisedModel, verbosity::Integer, X, y) -> fitresult, cache, report
# MLJBase.predict(model::SomeSupervisedModel, fitresult, Xnew) -> yhat

# function fit(mljModel, verbosity, X, y)
#     # This line is wrong, we'll need to grab the right NamedTuple keys
#     (samples, report) = mljModel.sampler(mljModel.m(X=X), (y=y,))
#     cache = ...
#     (samples, cache, report)
# end

# function Base.rand(p::MLJpredictor)
#     β = rand(p.βs)

#     # This line is wrong, we'll need to grab the right NamedTuple keys
#     rand(p.pred(X=p.Xrow, β=β))
# end

# function predict(mljModel,βs,Xnew)
#     pred = predictive(mljModel.m, mljModel.X)
#     [MLJpredictor(pred, Xrow, βs) for Xrow in eachrow(Xnew)]
# end


# Author: Vincent Arel-Bundock
# Contact: varel@umich.edu
# License: BSD-3
# Original code from the python statsmodels project
# https://github.com/statsmodels/statsmodels

function qreg_coef(y, X::Matrix, q::Real, method::IRLS)
    n, p = size(X)
    xstar = copy(X)
    diff = Inf

    beta0 = zeros(p)
    beta  = zeros(p)
    xtx   = Array{Float64}(undef, p, p)
    xty   = Array{Float64}(undef, p)
    xbeta = Array{Float64}(undef, n)
    resid = Array{Float64}(undef, n)

    for itr in 1:method.maxiter
        if diff > method.tol
            copyto!(beta0, beta)

            mul!(xtx, xstar', X)
            mul!(xty, xstar', y)
            beta .= xtx \ xty
            mul!(xbeta, X, beta)

            @. resid = y - xbeta

            for i in 1:n
                if abs(resid[i]) < method.threshold
                    @inbounds resid[i] = sign(resid[i]) * method.threshold
                end
                if resid[i] < 0
                    @inbounds resid[i] = abs(q * resid[i])
                else
                    @inbounds resid[i] = abs((1 - q) * resid[i])
                end
            end

            xstar .= X ./ resid

            diff = norm(beta0 - beta, Inf)
        end
    end
    return beta
end


@testset "Spin Sampling" begin
    function test_sampling_spin(
        sampling_type::Symbol, pg::ParticleGroup{D,V}, df::CosSumGaussianSpin
    ) where {D,V}
        mean = zeros(2)
        sigma = zeros(2)

        n_particles = pg.n_particles

        sampling = ParticleSampler{D,V}(sampling_type, n_particles)

        sample!(pg, sampling, df, mesh)

        for i_part in 1:n_particles
            xi = get_x(pg, i_part)
            vi = get_v(pg, i_part)
            mean[1] += xi[1]
            mean[2] += vi[1]
        end

        mean = mean / n_particles

        for i_part in 1:n_particles
            xi = get_x(pg, i_part)
            vi = get_v(pg, i_part)
            sigma[1] += (xi[1] - mean[1])^2
            sigma[2] += (vi[1] - mean[2])^2
        end

        sigma = sigma / (n_particles - 1)

        return mean, sigma
    end

    σ, μ = 0.02, 0.0
    kx, α = 1.004355, 0.001
    n_particles = 100000
    xmin = 0.0
    xmax = 2π / kx
    Lx = xmax - xmin
    nx = 512

    mesh = OneDGrid(xmax, xmin, nx)

    pg = ParticleGroup{1,1}(n_particles; n_spin=3)

    params = (k=[[kx]], α=[α], σ=[[σ]], μ=[[μ]])

    df = CosSumGaussianSpin(params...)

    mean, sigma = test_sampling_spin(:sobol, pg, df)

    @test true
end


using DifferentialEquations
using Plots
n_iter=50
barrier=4
HH=range(.125001,stop=.499,length=n_iter)
exit_time=zeros(length(HH))
dE=zeros(length(HH))
max_time=exp(8)
tspan = (0.0,max_time)
pert=1e-6

ODE1(z,w)=conj(  im * w.*( 1 ./(z.^2-w.^2)+1 ./(1+w.^2) ))
ODE2(z,w)=conj(  im * z.*( 1 ./(w.^2-z.^2)+1 ./(1+z.^2) ))
condition(u,t,integrator)=abs(u[1])>barrier || abs(u[2])>barrier || abs(u[3])>barrier||abs(u[4])>barrier
affect!(integrator) = terminate!(integrator)
cb = DiscreteCallback(condition,affect!)
function Eq_of_M(du,u,p,t)
    du[1]=real(ODE1(u[1]+im*u[2],u[3]+im*u[4]))
    du[2]=imag(ODE1(u[1]+im*u[2],u[3]+im*u[4]))
    du[3]=real(ODE2(u[1]+im*u[2],u[3]+im*u[4]))
    du[4]=imag(ODE2(u[1]+im*u[2],u[3]+im*u[4]))
    return
end
# refernce
# z=X+iP W=Q+iY . Plane Q=0, P=0
# z=u[1]+i u[2] ; w=u[3]+i u[4]

for j=1:length(HH)
    H=HH[j]
    h=1/(2*H)
    Y_0=1/sqrt(h+1)
    u0= [0;0;pert;Y_0;0;0]
    prob = ODEProblem(Eq_of_M,u0,tspan)
    t,A=solve(prob, RK4(),reltol=1e-8,abstol=1e-10,maxiters=1e15,callback=cb)
    exit_time[j]=t.t[end]
    Z=A[1,:]+im *A[2,:];
    W=A[3,:]+im *A[4,:];
    N=length(t.t)
    Energy=zeros(N)
    for k=1:N
        Energy[k]=abs(1/(1+Z[k]^2)-1/(1+W[k]^2))
    end
    dE[j]=Energy[end]-Energy[1]

end

p=plot(HH,log.(exit_time),seriestype=:scatter,marker=:circle,msw = 0,ms=2,seriescolor = :black,legend=false)
savefig(p,"ExitLogLog1000b.pdf")
 plot(-HH,-log.(abs.(dE)))


using DataFrames
using CSV
using Plots
#%%
data_v = open("pendulum.dat")
df_v=CSV.read(data_v, header=false, delim=" ")
t_v=df_v[1]
theta_v=df_v[2]
omega_v=df_v[3]

data_rk = open("pendulum_rk.dat")
df_rk=CSV.read(data_rk, header=false, delim=" ")
t_rk=df_rk[1]
theta_rk=df_rk[2]
omega_rk=df_rk[3]
#%%
plot(df_v[1], df_v[2], label="theta", dpi=600)
plot(df_rk[1], df_rk[2], label="theta", dpi=600)
# xaxis!(:log10)
# yaxis!(:log10)
#%%
E_v = (omega_v.^2)/2 + (ones(length(omega_v))-cos.(theta_v))
E_rk = (omega_rk.^2)/2 + (ones(length(omega_rk))-cos.(theta_rk))

plot(t_v, E_v, label="Verlet", dpi=600)
plot!(t_rk, E_rk, label="RK2", dpi=600)
plot!(legend=:topleft)
# ylims!(0,0.4)
ylabel!("E/(m L^2)")
xlabel!("step")
savefig("pendulum.png")


"""
SamplerState for the stochastic gradient Nosé-Poincaré Hamiltonian Dynamcics algorithm following

Roychowdhury, A., Kulis, B., & Parthasarathy, S. (2016). Robust Monte Carlo Sampling using Riemannian Nos’{e}-Poincar’{e} Hamiltonian Dynamics. In ICML (pp. 2673–2681).

Fields:

    - x: the current state
    - p: current momentum
    - xi: thermostat variable
    - t: iterations counter
    - stepsize: s.stepsize(s.t) gives the stepsize to be used in iteration s.t
    - mu: Constant in the thermostat equation
    - D: friction parameters
"""
type SGNHTState <: SamplerState
    x::Array{Float64,1}
    p::Array{Float64,1}
    xi::Float64

    t::Int
    stepsize::Function
    mu::Float64
    D::Float64
    function SGNHTState(x,stepsize::Function;mu=1.,D=1.0)
        p = randn(length(x))
        xi = D
        new(x,p,xi,0,stepsize,mu,D)
    end
    function SGNHTState(x,stepsize::Float64;mu=1.,D=1.0)
        p = randn(length(x))
        xi = D
        eps = stepsize
        stepsize = niters::Int -> eps
        new(x,p,xi,0,stepsize,mu,D)
    end

end

"""
sample! performs one SGLD update and returns the updated SGLDState.
"""
function sample!(s::SGNHTState,grad::Function)
    nparams = length(s.x)
    eps = s.stepsize(s.t)
    Base.LinAlg.axpy!(.5*eps,grad(s.x),s.p)
    s.t += 1
    Base.LinAlg.axpy!(0.5*eps,s.p,s.x)
    s.xi += .5*eps/s.mu * (sum(s.p.*s.p) - nparams)
    if abs(eps*s.xi) > 1e-6
    s.p = exp(-eps*s.xi)*s.p + sqrt(.5*(1-exp(-2.0*eps*s.xi))/s.xi * s.D) .* randn(nparams)
    else
    Base.LinAlg.axpy!(sqrt(eps*s.D), randn(nparams),s.p)
    end
    s.xi += .5*eps/s.mu * (sum(s.p.*s.p) - nparams)
    Base.LinAlg.axpy!(.5*eps,s.p,s.x)
    Base.LinAlg.axpy!(.5*eps,grad(s.x),s.p)
    s
end


using ManifoldsBase: LinearAlgebra
include("../utils.jl")

@testset "Tucker" begin
    n⃗ = (4, 5, 6)
    r⃗ = (2, 3, 4)
    r⃗_small = (2, 3, 3)
    M = Tucker(n⃗, r⃗)
    M_small = Tucker(r⃗, (2, 3, 3))

    types = [Float64]
    TEST_FLOAT32 && push!(types, Float32)

    for T in types
        @testset "Type $T" begin
            U = ntuple(d -> Matrix{T}(LinearAlgebra.I, n⃗[d], r⃗[d]), 3)
            C₁ = reshape(
                T[
                    0.0 0.0 1.0 0.0
                    1.0 1/√2 0.0 0.0
                    1.0 -1/√2 0.0 0.0
                    0.0 0.0 0.0 0.0
                    0.0 0.0 0.0 1.0
                    0.0 0.0 0.0 0.0
                ],
                (2, 3, 4),
            )
            C₂ = reshape(
                T[
                    0.0 0.0 -1/√2 1/√2
                    -1/√2 -0.5 0.0 0.0
                    1.0 -1/√2 0.0 0.0
                    0.0 0.0 0.0 0.0
                    0.0 0.0 1/√2 1/√2
                    1/√2 0.5 0.0 0.0
                ],
                (2, 3, 4),
            )
            C₃ = reshape(
                T[
                    1/√2 0.5 1/√2 0.0
                    1/√2 0.5 -1/√2 0.0
                    1/√2 -0.5 0.0 0.0
                    -1/√2 0.5 0.0 0.0
                    0.0 0.0 0.0 1/√2
                    0.0 0.0 0.0 -1/√2
                ],
                (2, 3, 4),
            )

            p = TuckerPoint(C₁, U...)
            p_small = TuckerPoint(C₂, r⃗_small)
            p_ambient = embed(M, p)

            U⊥ = ntuple(d -> nullspace(U[d]') * Matrix{Bool}(I, n⃗[d] - r⃗[d], r⃗[d]), 3)
            v = TuckerTVector(reshape(collect(1.0:prod(r⃗)), r⃗), U⊥)
            U⊥ = ntuple(d -> nullspace(U[d]') * Matrix{Bool}(I, n⃗[d] - r⃗[d], r⃗[d]), 3)
            w = TuckerTVector(reshape((1.0:prod(r⃗)) .^ 2, r⃗), U⊥)

            norm_v = norm(M, p, v)
            @test norm(M, p, 2 * v - v * 2) ≤ √eps(T) * norm_v
            @test norm(M, p, (v / 2) - (2 \ v)) ≤ eps(T) * norm_v
            @test norm(M, p, (v / 2) - 0.5v) ≤ eps(T) * norm_v
            @test +v == v

            @test embed(M, p) ≈ reshape(kron(reverse(U)...) * vec(C₁), n⃗)
            @test embed(M, TuckerPoint(p_ambient, r⃗)) ≈ p_ambient
            @test inner(M, p, v, w) ≈ dot(embed(M, p, v), embed(M, p, w))
            @test inner(M, p, v, v) ≈ norm(M, p, v)^2
            w_ambient = embed(M, p, w)
            @test inner(M, p, v, w_ambient) ≈ inner(M, p, v, w) ≈ inner(M, p, w_ambient, v)

            @test p ≈ TuckerPoint(p_ambient, r⃗)
            @test v == TuckerTVector(v.Ċ, v.U̇)

            @test manifold_dimension(M) ==
                  prod(r⃗) + sum(ntuple(d -> r⃗[d] * (n⃗[d] - r⃗[d]), length(r⃗)))

            @test is_point(M_small, p_small)
            @test is_point(M, p)
            @test is_point(M, embed(M, p))
            @test !is_point(M, zeros(n⃗))
            @test !is_point(M, TuckerPoint(zeros(T, r⃗), U...))
            @test !is_point(M, TuckerPoint(C₁, map(rk -> Matrix{T}(I, rk, rk), r⃗)...))
            @test !is_point(M, p_small)
            @test !is_point(M, embed(M_small, p_small))
            @test !is_point(M_small, p)
            q = allocate(p)
            copyto!(q, p)
            q.hosvd.core .= 2 * p.hosvd.core
            @test !is_point(M, q)

            u = allocate(v)
            u .= v
            @test u == v
            u = allocate(v)
            copyto!(u, v)
            @test u == v

            # broadcasting
            @test axes(v) === ()
            u = copy(v)
            # test that the copy is equal to the original, but represented by
            # a new array
            @test u.U̇ !== v.U̇
            @test u.Ċ == v.Ċ
            x = u .+ v .* 2
            @test x isa TuckerTVector
            @test x == v + u * 2
            x .= 2 .* u .+ v
            @test x == 2 * u + v

            @test is_vector(M, p, v)
            @test !is_vector(M, p_small, v)
            @test !is_vector(
                M,
                p,
                TuckerTVector(ones(T, r⃗), map(u⊥ -> u⊥[:, 1:(end - 1)], U⊥)),
            )

            ℬ = get_basis(M, p, DefaultOrthonormalBasis())
            J = convert(Matrix, ℬ)
            @test J'J ≈ LinearAlgebra.I
            vectors = get_vectors(M, p, DefaultOrthonormalBasis())
            vectors_mtx = hcat(map(ξ -> vec(embed(M, p, ξ)), vectors)...)
            @test J ≈ vectors_mtx

            shiftprint(x) = replace(sprint(show, "text/plain", x), "\n" => "\n ")

            @test sprint(show, "text/plain", M) == "Tucker((4, 5, 6), (2, 3, 4), ℝ)"
            @test sprint(show, "text/plain", p) == """
            $(summary(p))
            U factor 1:
             $(shiftprint(p.hosvd.U[1]))\n
            U factor 2:
             $(shiftprint(p.hosvd.U[2]))\n
            U factor 3:
             $(shiftprint(p.hosvd.U[3]))\n
            Core:
             $(shiftprint(p.hosvd.core))"""

            @test sprint(show, "text/plain", v) == """
            $(summary(v))
            U̇ factor 1:
             $(shiftprint(v.U̇[1]))\n
            U̇ factor 2:
             $(shiftprint(v.U̇[2]))\n
            U̇ factor 3:
             $(shiftprint(v.U̇[3]))\n
            Ċ factor:
             $(shiftprint(v.Ċ))"""

            @test sprint(show, "text/plain", ℬ) == """
            $(summary(ℬ)) ≅ $(shiftprint(J))
            """

            pts = [
                p,
                TuckerPoint(C₂, U[1][[4, 2, 1, 3], :], U[2:end]...),
                TuckerPoint(C₃, U[1][[4, 3, 1, 2], :], U[2:end]...),
            ]
            test_manifold(
                M,
                pts;
                is_mutating=false, # avoid allocations of the wrong type
                basis_types_to_from=(DefaultOrthonormalBasis(),),
                basis_types_vecs=(DefaultOrthonormalBasis(),),
                test_exp_log=false,
                default_inverse_retraction_method=ProjectionInverseRetraction(),
                test_injectivity_radius=false,
                default_retraction_method=PolarRetraction(),
                test_is_tangent=false,
                test_project_tangent=false,
                test_default_vector_transport=false,
                test_vector_spaces=false,
                test_vee_hat=false,
                test_tangent_vector_broadcasting=true,
                test_representation_size=false,
                projection_atol_multiplier=15,
                retraction_methods=[PolarRetraction()],
                inverse_retraction_methods=[ProjectionInverseRetraction()],
                mid_point12=nothing,
            )
        end
    end
end


"""
We define a `MultiRandomWalk` struct herein. This is a discrete-time, discrete-space
stochastic process. A `MultiRandomWalk` is a `RandomWalk` with multiple agents walking on
the same domain.
"""

"""
    MultiRandomWalk{S,T,U<:Sampleable} <: AbstractMarkovChain

##### Fields
- `domain::S` : the domain within which walkers move (e.g. an integer lattice or a graph).
- `initialstate::T<:AbstractVector` : the state at which the random walk begins (at
time-step 1) -- contains each walker's state.
- `stepsampler::U<:MultiSampler` : sampler that allows i.i.d. steps to be drawn.
"""
struct MultiRandomWalk{S,T<:AbstractVector,U<:MultiSampler} <: AbstractMarkovChain
    domain::S
    initialstate::T
    stepsampler::U
end

"""
    Base.show(io::IO, mrw::MultiRandomWalk) -> nothing

Show a `MultiRandomWalk` in a human-friendly manner.
"""
function Base.show(io::IO, mrw::MultiRandomWalk)
    println(io, typeof(mrw))
    println(io, " Domain:")
    println(io, "  ", mrw.domain)
    println(io, " Number of walkers:")
    println(io, "  ", length(mrw.initialstate))
    println(io, " Initial state:")
    println(io, "  ", mrw.initialstate)
    println(io, "")
    println(io, " MultiSampler that allows i.i.d. steps to be drawn:")
    println(io, "  ", mrw.stepsampler)
    return nothing
end

"""
    Base.eltype(mrw::MultiRandomWalk) -> DataType

Return the type associated with `mrw`'s domain.
"""
function Base.eltype(mrw::MultiRandomWalk)
    return typeof(mrw.initialstate[1])
end

"""
    MultiRandomWalk(nwalkers::Integer, latticedim::Integer = 1,
                    initialstate::AbstractVector{<:AbstractVector{<:Integer}} = (
                                                  repeat(zeros(Int, latticedim), nwalkers)),
                    stepsampler::MultiSampler = MultiSampler(
                            repeat([EVecSampler(latticedim)], nwalkers))) -> MultiRandomWalk

Construct and return a `MultiRandomWalk` on the integer lattice with dimension `latticedim`.

##### Complexity
- Time complexity: `O(1)`.
"""
function MultiRandomWalk(nwalkers::Integer, latticedim::Integer = 1,
                         initialstate::AbstractVector = (
                                              [zeros(Int, latticedim) for _ in 1:nwalkers]),
                         stepsampler::MultiSampler = MultiSampler(
                                               repeat([EVecSampler(latticedim)], nwalkers)))
    # check arguments
    latticedim < 1 && throw(ArgumentError("latticedim must be positive."))
    nwalkers < 1 && throw(ArgumentError("nwalkers must be positive."))
    for i in 1:length(initialstate)
        length(initialstate[i]) != latticedim &&
            throw(DimensionMismatch("initialstate's elements' lengths should equal
                                    latticedim"))
    end

    if latticedim == 1
        return MultiRandomWalk(Int, initialstate, stepsampler)
    else
        return MultiRandomWalk(repeat([Int], latticedim), initialstate, stepsampler)
    end
end

"""
    MultiRandomWalk(graph::AbstractGraph, nwalkers::Integer,
                    initialvertex::AbstractVector = repeat([rand(1:nv(graph))], nwalkers);
                    checkargs::Bool = true) -> MultiRandomWalk

Construct and return a `MultiRandomWalk` on a graph.

##### Complexity
- Time complexity: `O(n)` when `checkargs == true`; otherwise, `O(1)`.
"""
function MultiRandomWalk(graph::AbstractGraph, nwalkers::Integer,
                        initialvertex::AbstractVector=repeat([rand(1:nv(graph))], nwalkers);
                        checkargs::Bool = true)

    if checkargs && !(initialvertex in vertices(graph))
        throw(ArgumentError("initialvertex must be a vertex in graph)."))
    else
        return RandomWalk(graph, initialvertex, NeighborSampler(graph, initialvertex))
    end
end

"""
    nwalkers(mrw::MultiRandomWalk) -> Int

Return the number of walkers in `mrw`.

##### Complexity
- Time complexity: `O(1)`.
"""
function nwalkers(mrw::MultiRandomWalk)
    return length(initialstate)
end


# Descriptor of a cartesian grid

"""
    struct CartesianDescriptor{D,T,F<:Function}
      origin::Point{D,T}
      sizes::NTuple{D,T}
      partition::NTuple{D,Int}
      map::F
    end

Struct that stores the data defining a Cartesian grid.
"""
struct CartesianDescriptor{D,T,F<:Function} <: GridapType
  origin::Point{D,T}
  sizes::NTuple{D,T}
  partition::NTuple{D,Int}
  map::F
  isperiodic::NTuple{D,Bool}
  @doc """
      CartesianDescriptor(
        origin::Point{D},
        sizes::NTuple{D},
        partition;
        map::Function=identity,
        isperiodic::NTuple{D,Bool}=tfill(false,Val{D})) where D

  `partition` is a 1D indexable collection of arbitrary type.
  """
  function CartesianDescriptor(
    origin::Point{D},
    sizes::NTuple{D},
    partition;
    map::Function=identity,
    isperiodic::NTuple{D,Bool}=tfill(false,Val{D}())) where D

    for i in 1:D
      if isperiodic[i]
        @assert (partition[i] > 2) "A minimum of 3 elements is required in any
          periodic direction"
      end
    end

    T = eltype(sizes)
    F = typeof(map)
    new{D,T,F}(origin,sizes,Tuple(partition),map,isperiodic)
  end
end

"""
    CartesianDescriptor(
      domain,
      partition;
      map::Function=identity,
      isperiodic::NTuple{D,Bool}=tfill(false,Val{D}))

`domain` and `partition` are 1D indexable collections of arbitrary type.
"""
function CartesianDescriptor(
  domain,
  partition;
  map::Function=identity,
  isperiodic::NTuple=tfill(false,Val{length(partition)}()))

  D = length(partition)
  limits = [(domain[2*d-1],domain[2*d]) for d in 1:D]
  sizes = Tuple([(limits[d][2]-limits[d][1])/partition[d] for d in 1:D])
  origin = Point([ limits[d][1] for d in 1:D]...)
  CartesianDescriptor(origin,sizes,partition;map=map,isperiodic=isperiodic)
end

"""
    CartesianDescriptor(
      pmin::Point{D},
      pmax::Point{D},
      partition;
      map::Function=identity,
      isperiodic::NTuple{D,Bool}=tfill(false,Val{D})) where D

`partition` is a 1D indexable collection of arbitrary type.
"""
function CartesianDescriptor(
  pmin::Point{D},
  pmax::Point{D},
  partition;
  map::Function=identity,
  isperiodic::NTuple{D,Bool}=tfill(false,Val{D}())) where D

  T = eltype(pmin)
  domain = zeros(T,2*D)
  for d in 1:D
    domain[2*(d-1)+1] = pmin[d]
    domain[2*(d-1)+2] = pmax[d]
  end
  CartesianDescriptor(domain,partition;map=map,isperiodic=isperiodic)
end

# Deprecated signatures for backwards compatibility

@deprecate CartesianDescriptor(origin::Point{D}, sizes::NTuple{D}, partition, map::Function) where D CartesianDescriptor(origin, sizes, partition; map=map)

@deprecate CartesianDescriptor(domain,partition,map::Function) CartesianDescriptor(domain,partition;map=map)

@deprecate CartesianDescriptor(pmin::Point{D}, pmax::Point{D}, partition, map::Function) where D CartesianDescriptor(pmin, pmax, partition; map=map)

# Coordinates

struct CartesianCoordinates{D,T,F} <: AbstractArray{Point{D,T},D}
  data::CartesianDescriptor{D,T,F}
end

Base.size(a::CartesianCoordinates) = a.data.partition .+ 1

Base.IndexStyle(::Type{<:CartesianCoordinates}) = IndexCartesian()

function Base.getindex(a::CartesianCoordinates{D,T}, I::Vararg{Int,D}) where {D,T}
  p = zero(mutable(Point{D,T}))
  x0 = a.data.origin
  dx = a.data.sizes
  @inbounds for d in 1:D
    p[d] =  x0[d] + (I[d]-1)*dx[d]
  end
  a.data.map(Point(p))
end

# Cell nodes

struct CartesianCellNodes{D} <: AbstractArray{Vector{Int},D}
  partition::NTuple{D,Int}
  function CartesianCellNodes(partition)
    D = length(partition)
    new{D}(partition)
  end
end

Base.size(a::CartesianCellNodes) = a.partition

Base.IndexStyle(::Type{<:CartesianCellNodes}) = IndexCartesian()

function array_cache(a::CartesianCellNodes{D}) where D
  zeros(Int,2^D)
end

@inline function getindex!(cache,a::CartesianCellNodes,i::Integer)
  cis = CartesianIndices(size(a))
  ci = cis[i]
  getindex!(cache,a,Tuple(ci)...)
end

@inline function getindex!(v,a::CartesianCellNodes{1},i::Integer)
  v[1] = i
  v[2] = i+1
  v
end

@inline function getindex!(v,a::CartesianCellNodes{D},i::Integer...) where D
  nodes = LinearIndices(size(a).+1)
  lnodes = CartesianIndices(tfill(2,Val{D}()))
  j = i .- 1
  for (i,lnode) in enumerate(lnodes)
    k = Tuple(lnode) .+ j
    v[i] = nodes[ k... ]
  end
  v
end

function Base.getindex(a::CartesianCellNodes,i::Integer...)
  cache = array_cache(a)
  getindex!(cache,a,i...)
end

# Cartesian Grid

"""
    struct CartesianGrid{D,T,F} <: Grid{D,D}
      # private fields
    end
"""
struct CartesianGrid{D,T,F} <: Grid{D,D}
  node_coords::CartesianCoordinates{D,T,F}
  cell_nodes::CartesianCellNodes{D}
  cell_type::Fill{Int8,1,Tuple{Base.OneTo{Int}}}
  @doc """
      CartesianGrid(desc::CartesianDescriptor)
  """
  function CartesianGrid(desc::CartesianDescriptor{D,T,F}) where {D,T,F}
    node_coords = CartesianCoordinates(desc)
    cell_nodes = CartesianCellNodes(desc.partition)
    cell_type = Fill(Int8(1),length(cell_nodes))
    new{D,T,F}(node_coords,cell_nodes,cell_type)
  end
end

OrientationStyle(::Type{<:CartesianGrid}) = Val{true}()

"""
    get_cartesian_descriptor(grid::CartesianGrid)

Get the descriptor of the Cartesian grid
"""
function get_cartesian_descriptor(a::CartesianGrid)
  a.node_coords.data
end

get_node_coordinates(g::CartesianGrid) = g.node_coords

get_cell_type(g::CartesianGrid) = g.cell_type

get_cell_nodes(g::CartesianGrid) = g.cell_nodes

function get_reffes(g::CartesianGrid{D}) where D
  p = Polytope(tfill(HEX_AXIS,Val{D}()))
  order = 1
  reffe = LagrangianRefFE(Float64,p,order)
  [reffe,]
end

get_reffes(g::CartesianGrid{1}) = [SEG2,]

get_reffes(g::CartesianGrid{2}) = [QUAD4,]

get_reffes(g::CartesianGrid{3}) = [HEX8,]

"""
    CartesianGrid(args...;kwargs...)

Same args needed to construct a `CartesianDescriptor`
"""
function CartesianGrid(args...;kwargs...)
  desc = CartesianDescriptor(args...;kwargs...)
  CartesianGrid(desc)
end

# Cell map

struct CartesianMap{D,T,L} <: AbstractArray{AffineMap{D,T,L},D}
  data::CartesianDescriptor{D,T,typeof(identity)}
  function CartesianMap(des::CartesianDescriptor{D,T}) where {D,T}
    L = D*D
    new{D,T,L}(des)
  end
end

Base.size(a::CartesianMap) = a.data.partition

Base.IndexStyle(::Type{<:CartesianMap}) = IndexCartesian()

function Base.getindex(a::CartesianMap{D,T},I::Vararg{Int,D}) where {D,T}
  p = zero(mutable(Point{D,T}))
  x0 = a.data.origin
  dx = a.data.sizes
  @inbounds for d in 1:D
    p[d] =  x0[d] + (I[d]-1)*dx[d]
  end
  origin =  Point(p)
  jacobian = diagonal_tensor(VectorValue(dx))
  AffineMap(jacobian,origin)
end

function field_array_gradient(a::CartesianMap)
  dx = a.data.sizes
  jacobian = diagonal_tensor(VectorValue(dx))
  j = AffineMapGrad(jacobian)
  Fill(j,length(a))
end

function field_array_gradient(a::Reindexed{T,N,A}) where {T,N,A<:CartesianMap}
  g = field_array_gradient(a.i_to_v)
  reindex(g,a.j_to_i)
end

function get_cell_map(grid::CartesianGrid{D,T,typeof(identity)} where {D,T})
  CartesianMap(grid.node_coords.data)
end

# Cartesian grid topology with periodic BC

function _cartesian_grid_topology_with_periodic_bcs(grid::UnstructuredGrid,
  isperiodic::NTuple,
  partition)

  cell_to_vertices, vertex_to_node =
    _generate_cell_to_vertices_from_grid(grid, isperiodic, partition)
  _generate_grid_topology_from_grid(grid,cell_to_vertices,vertex_to_node)
end

function _generate_cell_to_vertices_from_grid(grid::UnstructuredGrid,
  isperiodic::NTuple, partition)

  if is_first_order(grid)
    nodes = get_cell_nodes(grid)
    cell_to_vertices = copy(nodes)

    nnodes = num_nodes(grid)
    num_nodes_x_dir = [partition[i]+1 for i in 1:length(partition)]
    point_to_isperiodic, slave_point_to_point, slave_point_to_master_point =
      _generate_slave_to_master_point(num_nodes_x_dir,isperiodic, nnodes)

    vertex_to_point = findall( .! point_to_isperiodic)
    point_to_vertex = fill(-1,length(point_to_isperiodic))
    point_to_vertex[vertex_to_point] = 1:length(vertex_to_point)
    point_to_vertex[slave_point_to_point] = point_to_vertex[slave_point_to_master_point]

    cell_to_vertices = Table(LocalToGlobalArray(nodes,point_to_vertex))

    vertex_to_node = vertex_to_point
    node_to_vertex = point_to_vertex
  else
    @notimplemented
  end
  (cell_to_vertices,vertex_to_node, node_to_vertex)
end

function _generate_slave_to_master_point(num_nodes_x_dir::Vector{Int},
  isperiodic::NTuple, num_nodes::Int)

  point_to_isperiodic = fill(false,num_nodes)

  slave_ijk_bounds = Array{Any,1}(undef,length(isperiodic))
  master_ijk_bounds = Array{Any,1}(undef,length(isperiodic))

  linear_indices = LinearIndices(Tuple(num_nodes_x_dir))
  periodic_dirs = findall(x->x==true, isperiodic)
  for periodic_dir in periodic_dirs
    for i in 1:length(isperiodic)
      if i == periodic_dir
        slave_ijk_bounds[i] = num_nodes_x_dir[i]
      else
        slave_ijk_bounds[i] = 1:num_nodes_x_dir[i]
      end
    end
    slave_ijk_set = CartesianIndices(Tuple(slave_ijk_bounds))
    point_to_isperiodic[linear_indices[slave_ijk_set]] .= true
  end

  slave_point_to_point = findall( point_to_isperiodic)
  slave_point_to_master_point = Array{Int,1}(undef,length(slave_point_to_point))

  cartesian_indices = CartesianIndices(Tuple(num_nodes_x_dir))
  for (i,point) in enumerate(slave_point_to_point)
    ijk = collect(cartesian_indices[point].I)
    for i in periodic_dirs
      if ijk[i] == num_nodes_x_dir[i]
        ijk[i] = 1
      end
    end

    master_point_ijk = CartesianIndex(Tuple(ijk))
    slave_point_to_master_point[i] = linear_indices[master_point_ijk]
  end

  point_to_isperiodic, slave_point_to_point, slave_point_to_master_point
end


# Path that the walk takes
function path(
    u,
    w::Walker2D,
    timesteps::Int64,
    rate::Float64,
    width::Int=1,
    length::Int=1)
# W is a type parameter
ws = [position(w)]

for i in 1:timesteps
    w = u(w, step(rate),width,length)
    push!(ws, position(w))
end
return ws
end

# Create local functions
function period_path(t,r,l,w)
    path(
        updateperiod,
        Walker2D(0,0),
        Int(round(t,digits=0)),
        r,w,l)
end

function period_path_df(t,r,l,w)
    DataFrame(
        reduce(
            vcat,
            [[i[1] i[2]] for i in period_path(t,r,l,w)]),
        [:x, :y])
end

# Taking more than one path
function paths(
Nwalks::Int,
u,
w::Walker2D,
timesteps::Int64,
rate::Float64,
width::Int=1,
length::Int=1)
displace = []

for k in 1:Nwalks
    t = path(
                u,w,timesteps,rate,width,length)
    tup = endpoint(t)
    data = WalkersData(
        Nwalks,
        timesteps,
        rate,
        width,
        length,
        tup.x,
        tup.y)

    push!(displace,data)
end
return displace
end

# Last point in a path
function endpoint(t)
    return Walker2D(t[end][1],t[end][2])
end

function paths_df(
    paramRange::Int,
    largestTime::Int,
    Nwalks::Int,
    updateRule,
    w::Walker2D,
    width::Int=1,
    length::Int=1
    )
timesteps = range(1,largestTime,length=paramRange)
rate = range(0.001, 0.25, length=paramRange)
dis_vec = []
for i in 1:paramRange,j in 1:paramRange
Walks = paths(
    Nwalks,
    updateRule,
    w,
    Int(round(timesteps[i],digits=0)),
    rate[j],
    width,
    length)
push!(dis_vec,Walks)
end
df = WalkersData2DataFram(dis_vec)

df.taxcabdisplacement = abs.(df.x) .+ abs.(df.y)
df.taxcabdisplacement² = abs.(df.x).^2 .+ abs.(df.y).^2
df.sqrttaxcabdisplacement² = sqrt.(df.taxcabdisplacement²)

return df
end


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

using Markdown
using InteractiveUtils

# ╔═╡ cb402178-8d6c-11eb-0162-0929b93a7cdb
sqrt(2)

# ╔═╡ 1b81c7a4-8d6d-11eb-1fff-050c4510b871
f = sqrt∘sum

# ╔═╡ 3652cce0-8d6d-11eb-204f-6d4f9f3e16ff
f([1,2])

# ╔═╡ 3ce8e698-8d6d-11eb-1666-5d62cbf21d1b
sum([1,2,3]) |> abs |> sqrt

# ╔═╡ 4cae5ade-8d6d-11eb-0169-7154cca2b874
([1,2,3])<| sum

# ╔═╡ 8b124a94-8d6d-11eb-0bb1-9f777e1d98f6
map(x -> x +1 ,[1,2,3]) == [2,3,4]

# ╔═╡ 24c947fa-8d6e-11eb-1b1b-dfdea765d7e6
foreach(println,[1,2,3])

# ╔═╡ 571228d8-8d6e-11eb-381f-4bb319df187f
foldl(=>,1:4) == (1 => 2 => 3 => 4)

# ╔═╡ 9943f8f0-8d6e-11eb-1bd0-4f90bdfb4543
foldl(=>,1:4)

# ╔═╡ 8e198494-8d6f-11eb-012d-034e88313e4e
begin
	a = 1
	b = 2
end

# ╔═╡ faa90d6c-8d6f-11eb-06b1-530a08ed4846
⊗(a,b) = kron(a,b)

# ╔═╡ 168e6272-8d70-11eb-2f8b-4dabdadaa454
⊗([1,2],[2,3])

# ╔═╡ 28272208-8d70-11eb-12ad-8fc5c7f16c98
foldl(⊗,[[1,2],[1,2],[1,2]])

# ╔═╡ 7c054e56-8d70-11eb-043d-cfebcc74b872
g(x) = (x==1||x==2) ? "x=1" : "x=2"

# ╔═╡ 62a9c638-8d72-11eb-253a-cb0f833dc233
g(2)

# ╔═╡ 9b18de58-8d77-11eb-335b-4b1775144c05

# ╔═╡ ae4013f2-8d77-11eb-13f6-3fd35a98cf09

# ╔═╡ Cell order:
# ╠═cb402178-8d6c-11eb-0162-0929b93a7cdb
# ╠═1b81c7a4-8d6d-11eb-1fff-050c4510b871
# ╠═3652cce0-8d6d-11eb-204f-6d4f9f3e16ff
# ╠═3ce8e698-8d6d-11eb-1666-5d62cbf21d1b
# ╠═4cae5ade-8d6d-11eb-0169-7154cca2b874
# ╠═8b124a94-8d6d-11eb-0bb1-9f777e1d98f6
# ╠═24c947fa-8d6e-11eb-1b1b-dfdea765d7e6
# ╠═571228d8-8d6e-11eb-381f-4bb319df187f
# ╠═9943f8f0-8d6e-11eb-1bd0-4f90bdfb4543
# ╠═8e198494-8d6f-11eb-012d-034e88313e4e
# ╠═faa90d6c-8d6f-11eb-06b1-530a08ed4846
# ╠═168e6272-8d70-11eb-2f8b-4dabdadaa454
# ╠═28272208-8d70-11eb-12ad-8fc5c7f16c98
# ╠═7c054e56-8d70-11eb-043d-cfebcc74b872
# ╠═62a9c638-8d72-11eb-253a-cb0f833dc233
# ╠═9b18de58-8d77-11eb-335b-4b1775144c05
# ╠═ae4013f2-8d77-11eb-13f6-3fd35a98cf09


module Tutorial

# 1 Getting started

using Modia
@usingModiaPlot

# Define model
SimpleModel = Model(
    T = 0.4,
    x = Var(init=0.2),
    equations = :[T * der(x) + x = 1],
)

# Transform to ODE form
simpleModel = @instantiateModel(SimpleModel)

# Simulate with a default integrator of DifferentialEquations
simulate!(simpleModel, stopTime = 1.2)

# Simulate with a specific integrator (Tsit5) and use a unit for stopTime
simulate!(simpleModel, Tsit5(), stopTime = 1.2u"s")

# Produce a line plot with GLMakie
plot(simpleModel, ("x", "der(x)"), figure=1)

# 2.1 - Equation oriented models

using Modia

LowPassFilter = Model(
    T = 0.2,
    u = input,
    y = output | Var(:x),
    x = Var(init=0),
    equations = :[T * der(x) + x = u],
)

# 2.2 Merging models

HighPassFilter = LowPassFilter | Model( y = :(-x + u) )

setLogMerge(true)

LowAndHighPassFilter = LowPassFilter | Model(
        y = nothing,
        low = output | Var(:x),
        high = output | Var(:(-x + u)),
    )

setLogMerge(false)

@showModel LowAndHighPassFilter

# 2.3 Functions and tables

TestLowAndHighPassFilter = LowAndHighPassFilter | Model(
        u = :(sin( (time+1u"s")*u"1/s/s" * time)*u"V"),
        x = Var(init=0.2u"V")
    )

using Interpolations

table = CubicSplineInterpolation(0:0.5:2.0, [0.0, 0.7, 2.0, 1.8, 1.2])
TestLowAndHighPassFilter2 = TestLowAndHighPassFilter | Map(u = :(table(time*u"1/s")*u"V"))

function ref(time)
    y1 = sin(time)
    y2 = cos(time)
    return (y1,y2)
end

using Modia
TestMultiReturningFunction1 = Model(
    equations = :[
        (y1,y2) = ref(time)
        y3 = y1+y2
    ]
)
testMultiReturningFunction1 = @instantiateModel(TestMultiReturningFunction1, unitless=true)
simulate!(testMultiReturningFunction1, stopTime=10.0)
plot(testMultiReturningFunction1, ["y1", "y2", "y3"], figure=2)

# 2.4 Hierarchical modeling

TwoFilters = (
    high = HighPassFilter,
    low = LowPassFilter,
)

BandPassFilter = (
    u = input,
    y = output,
    high = HighPassFilter | Map(T=0.5, x=Var(init=0.1u"V")),
    low = LowPassFilter | Map(x=Var(init=0.2u"V")),
    equations = :[
        high.u = u,
        low.u = high.y,
        y = low.y]
)

TestBandPassFilter = BandPassFilter | Map(
        u = :(sin( (time+1u"s")*u"1/s/s" * time)*u"V")
    )

#bandPassFilter = @instantiateModel(TestBandPassFilter, logStateSelection=true)
#simulate!(bandPassFilter, Tsit5(), stopTime = 50u"s")
#plot(bandPassFilter, ["u", "y"], figure=3)

StateSpace = Model(
    A = fill(0.0, 0, 0),
    B = fill(0.0, 0, 0),
    C = fill(0.0, 0, 0),
    D = fill(0.0, 0, 0),
    u = input,
    y = output,
    x = Var(init = zeros(0)),
    equations = :[
        der(x) = A*x + B*u
             y = C*x + D*u
    ]
)

SecondOrder = Model(
    w = 20.0,
    D =  0.1,
    k =  2.0,
    sys = StateSpace | Map(A = :([  0        1;
                                 -w^2  -2*D*w]),
                           B = :([0; w^2;;]),    # Julia 1.7: Trailing ";" defines a column matrix
                           C = :([k 0]),
                           D = :(zeros(1,1)),
                           x = Var(init = zeros(2)) ),
    equations = :[sys.u = [1.0]]
)
secondOrder = @instantiateModel(SecondOrder)
simulate!(secondOrder, stopTime=2.0)
plot(secondOrder, ("sys.u", "sys.x", "sys.y"), figure=3)

using ModiaLang.StaticArrays
TestArray1 = Model(
    v = Var(init=SVector{3,Float64}(1.0, 2.0, 3.0)),
    equations = :[der(v) = -v]
)
testArray1 = @instantiateModel(TestArray1, logCode=true)
simulate!(testArray1, stopTime=2.0)
plot(testArray1, "v", figure=4)

TestArray2 = Model(
    v = Var(init=[1.0, 2.0, 3.0]),
    equations = :[der(v) = -v]
)
testArray2 = @instantiateModel(TestArray2)
simulate!(testArray2, stopTime=2.0, merge=Map(v = [4.0, 3.0, 2.0, 1.0]))
plot(testArray2, "v", figure=5)

TwoInertiasAndIdealGearTooManyInits = Model(
    J1    = 50.0,
    J2    = 100.0,
    ratio = 2.0,
    f     = 3.0, # Hz

    phi1 = Var(init = 0.0), # Absolute angle of inertia1
    w1   = Var(init = 0.0), # Absolute angular velocity of inertia1
    phi2 = Var(init = 0.0), # Absolute angle of inertia2
    w2   = Var(init = 0.0), # Absolute angular velocity of inertia2

    equations = :[
        tau = 2.0*sin(2*3.14*f*time/u"s")

        # inertia1
        w1 = der(phi1)
        J1*der(w1) = tau - tau1

        # ideal gear
        phi1 = ratio*phi2
        ratio*tau1 = tau2

        # inertia2
        w2 = der(phi2)
        J2*der(w2) = tau2
    ]
)

drive1 = @instantiateModel(TwoInertiasAndIdealGearTooManyInits, logStateSelection=true)

end


using ImageDraw
using Test

@testset "Polygon" begin
    vert=CartesianIndex{2}[]
    push!(vert, CartesianIndex(1,1))
    push!(vert, CartesianIndex(1,3))
    push!(vert, CartesianIndex(1,5))
    push!(vert, CartesianIndex(3,3))
    push!(vert, CartesianIndex(3,1))

    img = @inferred draw(zeros(Gray{Bool},5,5), Polygon(vert))
    @test all(x->x==true, img[1,:])==true
    @test all(x->x==true, img[1:3,1])==true
    @test all(x->x==true, img[3,1:3])==true
    @test all(x->x==false, img[2,2:3])==true
    @test img[2,4]==true

    img=zeros(Gray{Bool},5,5)
    @inferred draw!(img, Polygon(vert))
    @test all(x->x==true, img[1,:])==true
    @test all(x->x==true, img[1:3,1])==true
    @test all(x->x==true, img[3,1:3])==true
    @test all(x->x==false, img[2,2:3])==true
    @test img[2,4]==true

    img=zeros(RGB{N0f8},5,5)
    @inferred draw!(img, Polygon(vert), RGB(1,0,0))
    @test all(x->x==RGB{N0f8}(1,0,0), img[1,:])==true
    @test all(x->x==RGB{N0f8}(1,0,0), img[1:3,1])==true
    @test all(x->x==RGB{N0f8}(1,0,0), img[3,1:3])==true
    @test all(x->x==RGB{N0f8}(0,0,0), img[2,2:3])==true
    @test img[2,4]==RGB{N0f8}(1,0,0)

    poly_tuples = [(1,1),(3,1),(5,1),(3,3),(1,3)]
    img = @inferred draw(zeros(Gray{Bool},5,5), Polygon(poly_tuples))
    @test all(x->x==true, img[1,:])==true
    @test all(x->x==true, img[1:3,1])==true
    @test all(x->x==true, img[3,1:3])==true
    @test all(x->x==false, img[2,2:3])==true
    @test img[2,4]==true
end

@testset "Path" begin
    vert=CartesianIndex{2}[]
    push!(vert, CartesianIndex(1,1))
    push!(vert, CartesianIndex(1,3))
    push!(vert, CartesianIndex(1,5))
    push!(vert, CartesianIndex(3,3))
    push!(vert, CartesianIndex(3,1))

    img=zeros(RGB{N0f8},5,5)
    @inferred draw!(img, Path(vert))
    @test all(x->x==RGB{N0f8}(1,1,1), img[1,:])==true
    @test all(x->x==RGB{N0f8}(1,1,1), img[3,1:3])==true
    @test all(x->x==RGB{N0f8}(0,0,0), img[2,1:3])==true
    @test img[2,4]==RGB{N0f8}(1,1,1)

    poly_tuples = [(1,1),(3,1),(5,1),(3,3),(1,3)]

    img=zeros(RGB{N0f8},5,5)
    @inferred draw!(img, Path(poly_tuples))
    @test all(x->x==RGB{N0f8}(1,1,1), img[1,:])==true
    @test all(x->x==RGB{N0f8}(1,1,1), img[3,1:3])==true
    @test all(x->x==RGB{N0f8}(0,0,0), img[2,1:3])==true
    @test img[2,4]==RGB{N0f8}(1,1,1)

    # No longer invalid, the program will draw what it can on the image

    # img = zeros(Gray{N0f8},5,5)
    # invalid_points = [(6,1), (4,4), (1,7), (7,6)]
    # @test_throws ErrorException draw!(img, Path(invalid_points))
    #
    # invalid_points = [(4,4), (1,7), (7,6)]
    # @test_throws ErrorException draw!(img, Path(invalid_points))
end

@testset "RegularPolygon" begin
    img = zeros(Gray, 10, 10)
    expected = copy(img)
    expected[3, 2:6] .= Gray(1)
    expected[7, 2:6] .= Gray(1)
    expected[3:7, 2] .= Gray(1)
    expected[3:7, 6] .= Gray(1)
    @test all(expected .== @inferred draw(img, RegularPolygon(CartesianIndex(5,4), 4, 4, π/4))) == true
end