export totalordermapcomponent, totalordermap

"""
$(TYPEDSIGNATURES)

Creates an `HermitemapComponent` of order `order` and dimension `Nx` for the basis `b`.
The features of the created maps are all the tensorial products of the basis elements up to the order `order`.
"""
function totalordermapcomponent(Nx::Int64, order::Int64; withconstant::Bool = false, b::String = "CstProHermiteBasis")
    @assert order >= 0 "Order should be positive"
    if b ∈ ["CstProHermiteBasis"; "CstPhyHermiteBasis"]
        B = MultiBasis(eval(Symbol(b))(order+2), Nx)
    elseif b ∈ ["CstLinProHermiteBasis"; "CstLinPhyHermiteBasis"]
        B = MultiBasis(eval(Symbol(b))(order+3), Nx)
    else
        error("Undefined basis")
    end

    idx = totalorder(order*ones(Int64, Nx))

    if withconstant == false
        idx = idx[2:end,:]
    end

    Nψ = size(idx, 1)

    f = ExpandedFunction(B, idx, zeros(Nψ))
    return HermiteMapComponent(IntegratedFunction(f))
end

"""
$(TYPEDSIGNATURES)

Creates an `Hermitemap` of order `order` for the basis `b`. The size of the map is extracted from the ensemble matrix `X`.
The features of the created maps are all the tensorial products of the basis elements up to the order `order`.
"""
function totalordermap(X::Array{Float64,2}, order::Int64; diag::Bool=true, factor::Float64 = 1.0, withconstant::Bool = false, b::String = "CstProHermite")

    L = LinearTransform(X; diag = diag, factor = factor)

    Nx = size(X,1)
    C = HermiteMapComponent[]
    @inbounds for i=1:Nx
        push!(C, totalordermapcomponent(i, order; withconstant = withconstant, b = b))
    end

    if b ∈ ["CstProHermiteBasis"; "CstPhyHermiteBasis"]
        m = order+2
    elseif b ∈ ["CstLinProHermiteBasis"; "CstLinPhyHermiteBasis"]
        m = order+3
    else
        error("Undefined basis")
    end

    return HermiteMap(m, Nx, L, C)
end


adj = spzeros(6,6)
adj[1,3] = 1
adj[1,5] = 1
adj[2,3] = 1
adj[2,4] = 1
adj[2,5] = 1
adj[2,6] = 1
labels = [1, 2, 1, 2, 1, 2]
counts = RE.count_neighbor_groups(adj, labels)
@test counts[1,1] == 2
@test counts[2,1] == 0
@test counts[1,2] == 2
@test counts[2,2] == 2
@test counts[1,3] == 0
@test counts[2,3] == 0
@test counts[1,4] == 0
@test counts[2,4] == 0

pairs = [(i,j) for i in [1,2] for j in [1,2]]
ns, ms = RE.compile_count_dicts(counts, labels, pairs)
@test ms[1,1] == 2
@test ms[1,2] == 0
@test ms[2,1] == 2
@test ms[2,2] == 2
@test ns[1] == 3
@test ns[2] == 3

pr = [0.1 0.2; 0.3 0.4]
ll = calculate_ll(ns, ms, pr)
@test_approx_eq_eps ll -7.6713 10e-4

i = 1
count_deltas = counts[:,i][:]
old_label = 1
new_label = 2
new_ns, new_ms = RE.adjust_counts(ns, ms, count_deltas, old_label, new_label)
@test new_ms[1,1] == 0
@test new_ms[1,2] == 0
@test new_ms[2,1] == 4
@test new_ms[2,2] == 2
@test new_ns[1] == 2
@test new_ns[2] == 4

new_ll = RE.calculate_ll(new_ns, new_ms, pr)
@test_approx_eq_eps new_ll -7.57119 10e-4

new_pr = zeros(2,2)
RE.maximize_parameters!(ns,ms,new_pr)
@test new_pr[1,1] == 2/9
@test new_pr[1,2] == 0
@test new_pr[2,1] == 2/9
@test new_pr[2,2] == 2/9

# test bipartite
labels = [1, 2, 3, 4, 3, 4]
counts = RE.count_neighbor_groups(adj, labels)
@test counts[1,1] == 0
@test counts[2,1] == 0
@test counts[3,1] == 2
@test counts[4,1] == 0
@test counts[1,2] == 0
@test counts[2,2] == 0
@test counts[3,2] == 2
@test counts[4,2] == 2
@test counts[1,3] == 0
@test counts[2,3] == 0
@test counts[3,3] == 0
@test counts[4,3] == 0
@test counts[1,4] == 0
@test counts[2,4] == 0
@test counts[3,4] == 0
@test counts[4,4] == 0
@test counts[1,5] == 0
@test counts[2,5] == 0
@test counts[3,5] == 0
@test counts[4,5] == 0
@test counts[1,6] == 0
@test counts[2,6] == 0
@test counts[3,6] == 0
@test counts[4,6] == 0
pairs = [(i,j) for i in [1,2] for j in [3,4]]
ns, ms = RE.compile_count_dicts(counts, labels, pairs)
@test ms[1,3] == 2
@test ms[1,4] == 0
@test ms[2,3] == 2
@test ms[2,4] == 2
@test ns[1] == 1
@test ns[2] == 1
@test ns[3] == 2
@test ns[4] == 2
@test RE.swap_group(1) == 2
@test RE.swap_group(2) == 1
@test RE.swap_group(3) == 4
@test RE.swap_group(4) == 3


module day24

dirs = Dict("e" => :e, "se" => :se, "sw" => :sw, "w" => :w, "ne" => :ne, "nw" => :nw)
offsets = Dict(:e => (2, 0), :se => (1, 1), :sw => (-1, 1), :w => (-2, 0), :ne => (1, -1), :nw => (-1, -1))

function parseline(line)
    result = []
    while length(line) > 0
        for dir in keys(dirs)
            if startswith(line, dir)
                push!(result, dirs[dir])
                line = line[length(dir)+1:end]
                break
            end
        end
    end
    result
end

function indextocoord(index)
    coord = CartesianIndex(0, 0)
    for i in index
        coord += CartesianIndex(offsets[i]...)
    end
    coord
end

function fliptile!(tiles, index)
    coord = indextocoord(index)
    if coord ∈ keys(tiles)
        tiles[coord] = !tiles[coord]
    else
        tiles[coord] = true
    end
end

input = parseline.(readlines("day24_input.txt"))
tiles = Dict()
map(i -> fliptile!(tiles, i), input)
println("Part 1: black = $(count(values(tiles)))")

# Part 2
adjacent(coord) = [coord + CartesianIndex(o) for o in values(offsets)]

function runstep!(tiles)
    candidates = Set(keys(tiles))
    for (k, v) in tiles
        if v
            union!(candidates, Set(adjacent(k)))
        end
    end
    oldtiles = copy(tiles)
    for c in candidates
        nblack = count([oldtiles[n] for n in adjacent(c) if n ∈ keys(oldtiles)])
        if c ∈ keys(oldtiles) && oldtiles[c]
            if nblack == 0 || nblack > 2 tiles[c] = false end
        else
            if nblack == 2 tiles[c] = true end
        end
    end
end

[runstep!(tiles) for i in 1:100]
println("Part 2: black = $(count(values(tiles)))")

end


using Test, PolynomialBases
using InteractiveUtils

# regression test
for basis_type in subtypes(PolynomialBases.NodalBasis{PolynomialBases.Line})
    println(devnull, "  ", basis_type(1))
    for p in 0:20, T in (Float32, Float64)
        basis = try
            basis_type(p, T)
        catch m
            isa(m, DomainError) && continue
            throw(m)
        end
        xmin, xmax = 0.25, 0.9
        x = similar(basis.nodes)
        ξ = similar(x)
        map_from_canonical!(x, basis.nodes, xmin, xmax, basis)
        map_to_canonical!(ξ, x, xmin, xmax, basis)
        @test ξ ≈ basis.nodes atol=10*eps(T)
    end
end


using LinearManifoldCluster
import Random
using Statistics
using MultivariateStats
using LinearAlgebra

import LaTeXStrings: latexstring
using Plots
upscale = 2
Plots.scalefontsizes(upscale)
default(size=(600*upscale,400*upscale))
pgfplots()
# gr()
# clibrary(:colorcet)
# theme(:wong2)

const SEED = 921369837
Random.seed!(SEED)

R(θ) = [cos(θ) -sin(θ); sin(θ) cos(θ)]

function elliptical_bound(S::AbstractMatrix{<:Real}, χ::Real)
    F = eigen(Symmetric(S))
    λ = F.values
    FI = sortperm(λ, rev=true)
    EV = F.vectors[:,FI[1]]
    ϕ = atan(EV[2], EV[1])
    if ϕ < 0
        ϕ += 2π
    end

    θ = range(0.0, 2π, step=0.1)
    a, b = χ*sqrt(λ[FI[1]]), χ*sqrt(λ[FI[2]])
    x, y  = a.*cos.(θ), b.*sin.(θ)
    ellipse = R(ϕ) * hcat(x, y)'
    return ellipse
end

# generate
n = 200          # number of points in generated manifold
N = 2            # space dimensionality
M = [1]         # number of manifolds
τ = .2 => .8   # translation vector bounds
E = [[.12]]      # bound on an extend from manifold for all manifolds)
Φ = Vector{Float64}[[.9]]   # bounds on manifold points (per manifold wrt dims)
X = generate(n,N,M,τ,Φ,E)[1][1]

# rotate
RX = R(π/5)*X
RX .+=  [0.5; -0.13]

# perform PCA
M = fit(PCA, RX)
P = projection(M)
O = mean(M)

e = 0.1
L1 = [O+P[:,1] O-P[:,1]]
L2 = [O+P[:,2]*e O-P[:,2]*e]

plot(RX[1,:], RX[2,:], seriestype=:scatter, aspect_ratio=:equal,
     xlims=(0.0,1.0), ylims=(0.2,0.85), markersize=2,
     legend=:bottomright, lab=latexstring("LM cluster, \$C_{\\Lambda,\\theta}\$"), palette=:viridis)
plot!(L1[1,:], L1[2,:], seriestype=:line, linewidth=2, lab=latexstring("linear manifold, \$\\Lambda\$"))
plot!(L2[1,:], L2[2,:], seriestype=:line, linewidth=2, linestyle=:dash, lab="orthoganal subspace")

# boundary
C = cov(RX.-O, dims=2)

for χ in [2.0, 2.5, 3.0]
    B = elliptical_bound(C, χ).+O
    plot!(B[1,:], B[2,:], seriestype=:path, linewidth=1,
          lab=latexstring("geom. LM cluster, \$C_{\\Lambda,\\chi=$χ}\$"))
end

Plots.pdf("linear-manifold-cluster-1d")


using StaticArrays, Test
using OrdinaryDiffEq

u0 = fill(zero(MVector{2,Float64}), 2)
u0[1] = ones(MVector{2,Float64}) .+ 1
f = (du,u,p,t) -> du .= u
ode = ODEProblem(f, u0, (0.,1.))
sol = solve(ode, Euler(), dt=1e-2)
sol = solve(ode, Tsit5())

u0 = fill(zero(SVector{2,Float64}), 2) .+ 1
u0[1] = ones(SVector{2,Float64}) .+ 1
ode = ODEProblem(f, u0, (0.,1.))
sol = solve(ode, Euler(), dt=1e-2)
sol = solve(ode, Tsit5())

sol = solve(ode, SSPRK22(), dt=1e-2)

u0 = ones(MVector{2,Float64})
ode = ODEProblem(f, u0, (0.,1.))
sol = solve(ode, Euler(), dt=1e-2)
sol = solve(ode, Tsit5(), dt=1e-2)

u0 = ones(SVector{2,Float64})
f = (u,p,t) -> u
ode = ODEProblem(f, u0, (0.,1.))
sol = solve(ode, Euler(), dt=1e-2)
integrator = init(ode, ImplicitEuler())
@test OrdinaryDiffEq.get_W(integrator.cache.nlsolver) isa StaticArrays.LU
sol = solve(ode, ImplicitEuler())
sol = solve(ode, ImplicitEuler(nlsolve=NLAnderson()))
sol = solve(ode, Tsit5(), dt=1e-2)

#https://github.com/JuliaDiffEq/DifferentialEquations.jl/issues/373
function lorenz_static(u,p,t)
 dx = 10.0*(u[2]-u[1])
 dy = u[1]*(28.0-u[3]) - u[2]
 dz = u[1]*u[2] - (8/3)*u[3]
 @SVector [dx,dy,dz]
end

u0 = @SVector [1.0,0.0,0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz_static,u0,tspan)
solve(prob,dt=0.1,Rosenbrock23(autodiff=false))


module PoincareInvariant2ndTest

    using PoincareInvariants
    using GeometricIntegrators
    using GeometricIntegrators.Utils
    using SymPy

    const nx = 100
    const ny = 100
    const Δt = 10.
    const B₀ = 1.
    const r₀ = 0.5
    const z₀ = 0.0
    const z₁ = 0.1
    const u₀ = 5E-1
    const u₁ = 5E-2
    const kz = 2
    const ku = 4

    function B(t, q)
        B₀ * (1 + q[1]^2 + q[2]^2)
    end

    function gcϑ(t, q, ϑ)
        ϑ[1] =-q[2] * (1 + B(t,q)) / 4
        ϑ[2] =+q[1] * (1 + B(t,q)) / 4
        ϑ[3] = q[4]
        ϑ[4] = 0
        nothing
    end

    function gcω(t, q, Ω)
        Ω[1,1] = 0
        Ω[1,2] =-B(t,q)
        Ω[1,3] = 0
        Ω[1,4] = 0

        Ω[2,1] =+B(t,q)
        Ω[2,2] = 0
        Ω[2,3] = 0
        Ω[2,4] = 0

        Ω[3,1] = 0
        Ω[3,2] = 0
        Ω[3,3] = 0
        Ω[3,4] =+q[4]

        Ω[4,1] = 0
        Ω[4,2] = 0
        Ω[4,3] =-q[4]
        Ω[4,4] = 0

        nothing
    end

    function gc_surface_q(s,t)
        x  = r₀*(s-0.5)
        y  = r₀*(t-0.5)
        z  = z₀ + z₁ * cos(2π*kz*s) * cos(2π*kz*t)
        u  = u₀ + u₁ * sin(2π*ku*s) * sin(2π*ku*t)

        [x, y, z, u]
    end

    function gc_surface_p(s,t)
        q = gc_surface_q(s,t)
        p = zero(q)
        gcϑ(zero(eltype(q)), q, p)
        p
    end

    function gc_dummy(t, a, b, c)
        nothing
    end

    function gc_dummy_pode(q₀, p₀)
        PODE(gc_dummy, gc_dummy, q₀, p₀)
    end

    function gc_dummy_iode(q₀)
        p₀ = zero(q₀)

        if ndims(q₀) == 1
            gcϑ(zero(eltype(q₀)), q₀, p₀)
        else
            tq = zeros(eltype(q₀), size(q₀,1))
            tp = zeros(eltype(p₀), size(p₀,1))

            for i in 1:size(q₀,2)
                simd_copy_xy_first!(tq, q₀, i)
                gcϑ(zero(eltype(q₀)), tq, tp)
                simd_copy_yx_first!(tp, p₀, i)
            end
        end

        IODE(gc_dummy, gc_dummy, gc_dummy, gc_dummy, q₀, p₀)
    end

    function compute_canonical_invariant_numerical()
        pinv = PoincareInvariant2ndCanonical(gc_dummy_pode, gc_surface_q, gc_surface_p, Δt, 4, nx, ny, 0)
        sol  = Solution(pinv.equ, Δt, 0)

        I = evaluate_poincare_invariant(pinv, sol)

        return I[1]
    end

    function compute_invariant_numerical()
        pinv = PoincareInvariant2nd(gc_dummy_iode, gc_surface_q, gcω, Δt, 4, nx, ny, 0)
        sol  = Solution(pinv.equ, Δt, 0)

        I, J, K, L = evaluate_poincare_invariant(pinv, sol)

        return I[1], J[1]
    end

    function compute_invariant_analytical()
        s, t = Sym("s, t")
        x, y, z, u = gc_surface_q(s,t)

        S = simplify(B₀ * (1 + x^2 + y^2) * ( diff(x, s) * diff(y, t) + diff(x, t) * diff(y, s) )
                   + u * ( diff(z, s) * diff(u, t) + diff(z, t) * diff(u, s) ) )
        I = N(SymPy.integrate(S, (s, 0, 1), (t, 0, 1)), 20)
    end

    export compute_canonical_invariant_numerical, compute_invariant_numerical, compute_invariant_analytical

end

using .PoincareInvariant2ndTest

### compute analytical invariant ###

Ia = compute_invariant_analytical()

### compute and check canonical invariant ###

Jn = compute_canonical_invariant_numerical()

@assert Jn ≈ Ia atol=2eps()

### compute and check noncanonical invariant ###

In, Jn = compute_invariant_numerical()

@assert In ≈ Ia atol=2eps()
@assert Jn ≈ Ia atol=2eps()


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
methods(factorial)

# %%
factorial(20)

# %%
factorial(21)

# %%
@which factorial(21)

# %%
Base._fact_table64

# %%
Base._fact_table64[1] = 666

# %%
factorial(1)

# %%
factorial(0.5)

# %%
using SpecialFunctions

# %%
factorial(0.5)

# %%
gamma(1.5)

# %%
@which factorial(0.5)

# %%
methods(factorial)

# %%
factorial(20)

# %%
@which factorial(20)

# %%


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# %%
using DynamicalSystems

# %%
f(u, p, t) = SVector(u[2], u[1]+u[2])
u0 = [0, 1]
prob = DiscreteDynamicalSystem(f, u0, nothing)

# %% tags=[]
tr = trajectory(prob, 20, Ttr=10)

# %% tags=[]
tr = trajectory(prob, 20, Ttr=0)

# %%
dump(tr)

# %%


"""
    truncresponse(response, L, noisepower; warn=true)
Truncate the given array to the given length, printing a warning if the
truncated tail has more average power than the given noise floor.

`noisepower` can be a scalar or a row vector with one entry per channel

Warning can be disabled with the kwarg `warn=false`
"""
function truncresponse(response::AbstractArray, L, noisepower=eps(); warn=true)
    RL = size(response, 1)
    if RL > L
        truncenergy = sum(response[L+1:end, :].^2; dims=1)
        headroom = 1.05 # 5% over noise
        atol = sqrt(eps()) # absolute tolerance (for when noise floor = 0.0)
        if warn && any(truncenergy .> (atol + noisepower * headroom) * (RL-L))
            @warn string("nonsilent samples past end of analysis window. Check ",
                         "your test signal is long enough for the response")
            # Base.show_backtrace(STDERR, backtrace())
        end
        response = timeslice(response, 1:L)
    end

    response
end

# internal functions to do the actual truncation while keeping the
# dimensionality
timeslice(response::AbstractVector, I) = response[I]
timeslice(response::AbstractMatrix, I) = response[I, :]

# used to help parse arguments
striphz(f) = f
striphz(f::Frequency{T}) where T = T(f/Hz)
stripsec(t) = t
stripsec(t::Time{T}) where T = T(t/s)

energy(v) = sum(x->x^2, v)
power(v) = energy(v)/length(v)

"""
Normalize a frequency with a given samplerate to cycles/sample.
"""
freqnorm(f, sr) = uconvert(NoUnits, f/sr)

"""
Normalize a duration with a given samplerate to integer samples (rounded)
"""
durnorm(dur, sr) = round(Int, uconvert(NoUnits, dur*sr))

"""
    snr(mixture, signal)

Returns the signal-to-noise power ratio assuming that the mixture contains a
scaled version of the signal.
"""
function snr(mixture, signal)
    mixcoef = dot(mixture, signal) / dot(signal, signal)
    proj = mixcoef * signal

    energy(proj)/energy(mixture-proj)
end


using SPH
using Test, Plots

default(;color = :thermal, aspect_ratio = 1, label = "", xaxis=("x", (-0.1, 1.1)), yaxis = ("y", (-0.1, 1.1)))

function M6_hardcoded(q)
    if q < 0.0
        res = 0.0
    elseif q < 1.0
        res = (3 - q)^5 - 6*(2 - q)^5 + 15*(1 - q)^5
    elseif q < 2.0
        res = (3 - q)^5 - 6*(2 - q)^5
    elseif q < 3.0
        res = (3 - q)^5
    else
        res = 0.0
    end
    return res
end

@testset "SPH.jl" begin
   xs = rand(1000) * 6 .- 1
   fs_hardcoded = M6_hardcoded.(xs) .* SPH.σ6[1]
   fs_procedural = M6[1].(xs)
   @test isapprox.(fs_hardcoded, fs_procedural, atol=1e-12) |> all
   #@show xs[a], fs_hardcoded[a], fs_procedural[a]
end

@testset "Density calculations" begin

    particles = SPH.particles_in_box(1)
    @test SPH.compute_distances(particles) ≈ [0.0]

    particles = [SPH.Particle(0.0, 0.0), SPH.Particle(0.0, 1.0)]
    @test SPH.compute_distances(particles) ≈ [0.0 1.0; 1.0 0.0]

    push!(particles, SPH.Particle(1.0, 1.0))
    @test SPH.compute_distances(particles) ≈ [0.0 1.0 √2; 1.0 0.0 1.0; √2 1.0 0.0]

end

N = 100
particles = SPH.particles_in_box(N)

NEdge = floor(Int, sqrt(N))
xs = LinRange(0, 1, NEdge)
ys = LinRange(0, 1, NEdge)

#left = SPH.Particle.(zeros(NEdge), ys)
#bottom = SPH.Particle.(xs, zeros(NEdge))
#top = SPH.Particle.(xs, ones(NEdge))
#right = SPH.Particle.(ones(NEdge), ys)

#append!(particles, left)
#append!(particles, right)
#append!(particles, top)
#append!(particles, bottom)

timestep = 0.01
anim = @animate for i in 1:500
    global particles = SPH.update(particles, dt=timestep);
    if i % 2 == 0
        SPH.plot_particles(particles, clims = (0, 120))
    end
end
gif(anim, "test.gif", fps = 25)


import HMCUtilities:
    VariableConstraint,
    constrain,
    free

struct TestAffineConstraint{N,TM,TMinv,Tb} <: HMCUtilities.VariableConstraint{N,N}
    M::TM
    Minv::TMinv
    b::Tb

    function TestAffineConstraint(M, b)
        n = length(b)
        @assert size(M) == (n, n)
        Minv = inv(M)
        return new{n,typeof(M),typeof(Minv),typeof(b)}(M, Minv, b)
    end
end

HMCUtilities.constrain(c::TestAffineConstraint, y) = c.Minv * (y - c.b)
HMCUtilities.free(c::TestAffineConstraint, x) = c.M * x + c.b

struct TestSphericalConstraint <: HMCUtilities.VariableConstraint{3,2} end

function HMCUtilities.constrain(c::TestSphericalConstraint, y)
    sinθ, cosθ = sincos(y[1])
    sinϕ, cosϕ = sincos(y[2])
    return [sinθ * cosϕ, sinθ * sinϕ, cosθ]
end

function HMCUtilities.free(c::TestSphericalConstraint, x)
    ϕ = atan(x[2], x[1])
    θ = acos(x[3])
    return [θ, ϕ]
end

struct TestSquareConstraint <: HMCUtilities.UnivariateConstraint end

HMCUtilities.constrain(c::TestSquareConstraint, y::Real) = y^2

HMCUtilities.free(c::TestSquareConstraint, x::Real) = sqrt(x)

@testset "defaults" begin
    @testset "one-to-one" begin
        M = [1.0 2.0;
             3.0 4.0]
        b = [5.0, 6.0]

        c = TestAffineConstraint(M, b)
        xs = [zeros(2), ones(2), [1.0, 2.0]]
        ys = [b, [8.0, 13.0], [10.0, 17.0]]
        logπxs = [-dot(x, x) / 2 for x in xs]
        ∇x_logπxs = [-x for x in xs]
        Minv = inv(M)
        logπys = logπxs .- log(2)
        ∇y_logπys = [Minv'x for x in ∇x_logπxs]

        for (x, y, logπx, ∇x_logπx, y_exp, logπy_exp, ∇y_logπy_exp) in zip(xs, ys, logπxs, ∇x_logπxs, ys, logπys, ∇y_logπys)
            test_constraint(
                c,
                x,
                y,
                logπx,
                ∇x_logπx,
                y_exp,
                logπy_exp,
                ∇y_logπy_exp,
                test_type_stability = false
            )
        end
    end

    @testset "non one-to-one" begin
        c = TestSphericalConstraint()
        xs = [[sqrt(3 / 2), sqrt(3 / 2), 1] ./ 2]
        ys = [π .* [1 / 3, 1 / 4]]
        κμ = 10.0 * [1.0, 0.0, 0.0]
        logπxs = [dot(x, κμ) for x in xs]  # vMF
        ∇x_logπxs = [κμ for x in xs]
        logπys = [(5 * sqrt(6) - log(4 / 3)) / 2]
        ∇y_logπys = [[5 / sqrt(2) + 1 / sqrt(3), -5 * sqrt(3 / 2)]]

        for (x, y, logπx, ∇x_logπx, y_exp, logπy_exp, ∇y_logπy_exp) in zip(xs, ys, logπxs, ∇x_logπxs, ys, logπys, ∇y_logπys)
            test_constraint(
                c,
                x,
                y,
                logπx,
                ∇x_logπx,
                y_exp,
                logπy_exp,
                ∇y_logπy_exp,
                test_type_stability = false
            )
        end
    end

    @testset "univariate" begin
        c = TestSquareConstraint()
        x = rand()
        y = -sqrt(x)
        y_exp = abs(y)
        logπx = x^2 / 2
        ∇x_logπx = x
        logπy_exp = logπx + log(2 * abs(y))
        ∇y_logπy_exp = 2 * y^3 + inv(y)

        @testset "scalar" begin
            vtypes = [Float64, Float32]
            test_constraint(
                c,
                x,
                y,
                logπx,
                ∇x_logπx,
                y_exp,
                logπy_exp,
                ∇y_logπy_exp;
                cvtypes=vtypes,
                fvtypes=vtypes,
                test_type_stability = false
            )
        end

        @testset "vector" begin
            vtypes = [Vector{Float64}, Vector{Float32}]
            test_constraint(
                c,
                [x],
                [y],
                logπx,
                [∇x_logπx],
                [y_exp],
                logπy_exp,
                [∇y_logπy_exp];
                cvtypes=vtypes,
                fvtypes=vtypes,
                test_type_stability = false
            )
        end
    end
end

@testset "IdentityConstraint" begin
    @testset "vector n = $n" for n in 1:3
        for m in 1:3
            c = HMCUtilities.IdentityConstraint(n)
            x = randn(n)
            y = copy(x)
            logπx = -dot(x, x) / 2
            ∇x_logπx = -x
            logπy_exp = copy(logπx)
            ∇y_logπy_exp = copy(∇x_logπx)
            test_constraint(c, x, y, logπx, ∇x_logπx, y, logπy_exp, ∇y_logπy_exp)
        end
    end

    @testset "scalar" begin
        c = HMCUtilities.IdentityConstraint(1)
        vtypes = [Float64, Float32]
        x = randn()
        y = copy(x)
        logπx = -dot(x, x) / 2
        ∇x_logπx = -x
        logπy_exp = copy(logπx)
        ∇y_logπy_exp = copy(∇x_logπx)
        test_constraint(
            c,
            x,
            y,
            logπx,
            ∇x_logπx,
            y,
            logπy_exp,
            ∇y_logπy_exp;
            cvtypes=vtypes,
            fvtypes=vtypes
        )
    end
end

@testset "LowerBoundedConstraint" begin
    vtypes = [Float64, Float32]

    @testset "lb=-5, y=0" begin
        c = HMCUtilities.LowerBoundedConstraint(-5.0)
        (y, x, logπx, ∇x_logπx, logπy_exp, ∇y_logπy_exp) = (0, -4, -8, 4, -8, 5)
        test_constraint(
            c,
            x,
            y,
            logπx,
            ∇x_logπx,
            y,
            logπy_exp,
            ∇y_logπy_exp;
            cvtypes=vtypes,
            fvtypes=vtypes
        )

        @testset "clamp" begin
            @test constrain(c, -Inf) ≈ -5 + eps()
        end
    end

    @testset "lb=5, y=-5" begin
        c = HMCUtilities.LowerBoundedConstraint(5.0)
        (y, x, logπx, ∇x_logπx, logπy_exp, ∇y_logπy_exp) = (
            -5, 5.00673794700, -12.5337124350, -5.00673794700, -17.5337124350, 0.966264865075
        )
        test_constraint(
            c,
            x,
            y,
            logπx,
            ∇x_logπx,
            y,
            logπy_exp,
            ∇y_logπy_exp;
            cvtypes=vtypes,
            fvtypes=vtypes
        )

        @testset "clamp" begin
            @test constrain(c, -Inf) ≈ 5 + eps()
        end
    end
end

@testset "UpperBoundedConstraint" begin
    vtypes = [Float64, Float32]

    @testset "ub=-5, y=0" begin
        c = HMCUtilities.UpperBoundedConstraint(-5.0)
        (y, x, logπx, ∇x_logπx, logπy_exp, ∇y_logπy_exp) = (0, -6, -18, 6, -18, -5)
        test_constraint(
            c,
            x,
            y,
            logπx,
            ∇x_logπx,
            y,
            logπy_exp,
            ∇y_logπy_exp;
            cvtypes=vtypes,
            fvtypes=vtypes
        )

        @testset "clamp" begin
            xbound = -5 - eps()
            @test constrain(c, -Inf) ≈ xbound
            @test constrain_with_pushlogpdf(c, -Inf)[1] ≈ xbound
            @test constrain_with_pushlogpdf_grad(c, -Inf)[1] ≈ xbound
        end
    end

    @testset "ub=5, y=-5" begin
        c = HMCUtilities.UpperBoundedConstraint(5.0)
        (y, x, logπx, ∇x_logπx, logπy_exp, ∇y_logπy_exp) = (
            -5, 4.99326205300, -12.4663329650, -4.99326205300, -17.4663329650, 1.03364433507
        )
        test_constraint(
            c,
            x,
            y,
            logπx,
            ∇x_logπx,
            y,
            logπy_exp,
            ∇y_logπy_exp;
            cvtypes=vtypes,
            fvtypes=vtypes
        )

        @testset "clamp" begin
            xbound = 5 - eps()
            @test constrain(c, -Inf) ≈ xbound
            @test constrain_with_pushlogpdf(c, -Inf)[1] ≈ xbound
            @test constrain_with_pushlogpdf_grad(c, -Inf)[1] ≈ xbound
        end
    end
end

@testset "BoundedConstraint" begin
    vtypes = [Float64]

    @testset "lb=0, ub=1" begin
        c = HMCUtilities.BoundedConstraint(0.0, 1.0)
        @testset "y=-20" begin
            (y, x, logπx, ∇x_logπx, logπy_exp, ∇y_logπy_exp) = (
                -20, 2.0611536181902035814e-9, 0.0, 0.0, -20, 1.0
            )
            test_constraint(
                c,
                x,
                y,
                logπx,
                ∇x_logπx,
                y,
                logπy_exp,
                ∇y_logπy_exp;
                cvtypes=vtypes,
                fvtypes=vtypes
            )
        end

        @testset "y=20" begin
            (y, x, logπx, ∇x_logπx, logπy_exp, ∇y_logπy_exp) = (
                20, 0.99999999793884638181, -0.5, -1, -20.5, -1
            )
            test_constraint(
                c,
                x,
                y,
                logπx,
                ∇x_logπx,
                y,
                logπy_exp,
                ∇y_logπy_exp;
                cvtypes=vtypes,
                fvtypes=vtypes
            )
        end

        @testset "clamp" begin
            xbound = 0 + eps()
            @test isapprox(constrain(c, -Inf), xbound, atol=1e-12)
            @test isapprox(constrain_with_pushlogpdf(c, -Inf)[1], xbound, atol=1e-12)
            @test isapprox(constrain_with_pushlogpdf_grad(c, -Inf)[1], xbound, atol=1e-12)

            xbound = 1 - eps()
            @test isapprox(constrain(c, Inf), xbound, atol=1e-12)
            @test isapprox(constrain_with_pushlogpdf(c, Inf)[1], xbound, atol=1e-12)
            @test isapprox(constrain_with_pushlogpdf_grad(c, Inf)[1], xbound, atol=1e-12)
        end
    end

    @testset "lb=-5, ub=5" begin
        c = HMCUtilities.BoundedConstraint(-5.0, 5.0)
        @testset "y=1" begin
            (y, x, logπx, ∇x_logπx, logπy_exp, ∇y_logπy_exp) = (
                1, 2.310585786, -2.669403338, -2.310585786, -1.993341620, -5.005004541
            )
            test_constraint(
                c,
                x,
                y,
                logπx,
                ∇x_logπx,
                y,
                logπy_exp,
                ∇y_logπy_exp;
                cvtypes=vtypes,
                fvtypes=vtypes
            )

        end

        @testset "clamp" begin
            xbound = -5 + eps()
            @test constrain(c, -Inf) ≈ xbound
            @test constrain_with_pushlogpdf(c, -Inf)[1] ≈ xbound
            @test constrain_with_pushlogpdf_grad(c, -Inf)[1] ≈ xbound

            xbound = 5 - eps()
            @test constrain(c, Inf) ≈ xbound
            @test constrain_with_pushlogpdf(c, Inf)[1] ≈ xbound
            @test constrain_with_pushlogpdf_grad(c, Inf)[1] ≈ xbound
        end
    end
end

@testset "TransformConstraint" begin
    (lb, ub) = (-rand(), rand())
    @test_throws AssertionError HMCUtilities.TransformConstraint(ub, lb)

    c = HMCUtilities.TransformConstraint(lb, Inf)
    @test c === HMCUtilities.LowerBoundedConstraint(lb)

    c = HMCUtilities.TransformConstraint(-Inf, ub)
    @test c === HMCUtilities.UpperBoundedConstraint(ub)

    c = HMCUtilities.TransformConstraint(lb, ub)
    @test c === HMCUtilities.BoundedConstraint(lb, ub)

    c = HMCUtilities.TransformConstraint()
    @test c === HMCUtilities.IdentityConstraint(1)
    c = HMCUtilities.TransformConstraint(-Inf, Inf)
    @test c === HMCUtilities.IdentityConstraint(1)
end

@testset "UnitVectorConstraint" begin
    vtypes = [Vector{Float64}, Vector{Float32}]

    @testset "n=$n" for n in 2:4
        c = HMCUtilities.UnitVectorConstraint(n)
        κμ = rand() * normalize(randn(n))

        y = randn(n)
        x = normalize(y)
        y_exp = copy(x)
        logπx = dot(x, κμ)  # vMF
        ∇x_logπx = κμ
        logπy_exp = logπx - dot(y, y) / 2
        ∇y_logπy_exp = (I - x * x') * ∇x_logπx / norm(y) - y

        test_constraint(
            c,
            x,
            y,
            logπx,
            ∇x_logπx,
            y_exp,
            logπy_exp,
            ∇y_logπy_exp;
            cvtypes=vtypes,
            fvtypes=vtypes,
            test_logdetJ_consistency = false
        )

        @testset "clamp" begin
            xbound = normalize(y)
            @test constrain(c, y) ≈ xbound
            @test constrain_with_pushlogpdf(c, y)[1] ≈ xbound
            @test constrain_with_pushlogpdf_grad(c, y)[1] ≈ xbound
        end
    end
end

@testset "UnitVectorScaledConstraint" begin
    vtypes = [Vector{Float64}, Vector{Float32}]

    @testset "n=$n, r=$r" for n in 2:4, r in (1.0, 3.0, 1.0)
        c = HMCUtilities.UnitVectorScaledConstraint(n, r)
        κμ = rand() * normalize(randn(n))

        y = randn(n) .* r
        x = normalize(y)
        y_exp = copy(x) .* r
        logπx = dot(x, κμ)  # vMF
        ∇x_logπx = κμ
        logπy_exp = logπx - dot(y, y) / r^2 / 2
        ∇y_logπy_exp = (I - x * x') * ∇x_logπx / norm(y) - y ./ r^2

        test_constraint(
            c,
            x,
            y,
            logπx,
            ∇x_logπx,
            y_exp,
            logπy_exp,
            ∇y_logπy_exp;
            cvtypes=vtypes,
            fvtypes=vtypes,
            test_logdetJ_consistency = false
        )

        @testset "clamp" begin
            xbound = normalize(y)
            @test constrain(c, y) ≈ xbound
            @test constrain_with_pushlogpdf(c, y)[1] ≈ xbound
            @test constrain_with_pushlogpdf_grad(c, y)[1] ≈ xbound
        end
    end
end

@testset "UnitSimplexConstraint" begin
    vtypes = [Vector{Float64}, Vector{Float32}]

    ys = (
        [-log(2)],
        [-log(2), log(2)],
        [-log(2), 0.0, log(2)]
    )

    xs = (
        [1/3, 2/3],
        [1/5, 8/15, 4/15],
        [1/7, 2/7, 8/21, 4/21]
    )

    Js = (
        [2/9, -2/9],
        [4/25 0.0; -8/75 8/45; -4/75 -8/45],
        [6/49 0.0 0.0;
        -2/49 4/21 0.0;
        -8/147 -8/63 8/63;
        -4/147 -4/63 -8/63]
    )

    ∇y_logdetJs = (
        [1/3],
        [2/5, -1/3],
        [3/7, 0.0, -1/3]
    )

    @testset "n=$n" for n in 2:4
        vtypes = [Vector{Float64}, Vector{Float32}]

        c = HMCUtilities.UnitSimplexConstraint(n)
        alpha = rand(n)

        x = xs[n - 1]
        y = ys[n - 1]
        J = Js[n - 1]
        ∇y_logdetJ = ∇y_logdetJs[n - 1]

        logπx, back = Zygote.pullback(x -> sum((alpha .- 1) .* log.(x)), x)
        ∇x_logπx = back(1.0)[1]

        logdetJ = _logabsdet(J'J)[1] / 2

        logπy_exp = logπx + logdetJ
        ∇y_logπy_exp = J' * ∇x_logπx .+ ∇y_logdetJ

        test_constraint(
            c,
            x,
            y,
            logπx,
            ∇x_logπx,
            y,
            logπy_exp,
            ∇y_logπy_exp;
            cvtypes=vtypes,
            fvtypes=vtypes
        )

        @testset "clamp" begin
            y = [-Inf, randn(n-2)...]
            xbound = constrain(c, y)
            @test isapprox(xbound[1], eps(); atol=1e-16)
            xbound = constrain_with_pushlogpdf(c, y)[1]
            @test isapprox(xbound[1], eps(); atol=1e-16)
            xbound = constrain_with_pushlogpdf_grad(c, y)[1]
            @test isapprox(xbound[1], eps(); atol=1e-16)

            y = [Inf, randn(n-2)...]
            xbound = constrain(c, y)
            @test xbound ≈ [1 - eps(), eps() * ones(n-1)...]
            @test xbound[1] < 1
            @test all(xbound[2:end] .> 0)
            xbound = constrain_with_pushlogpdf(c, y)[1]
            @test xbound ≈ [1 - eps(), eps() * ones(n-1)...]
            @test xbound[1] < 1
            @test all(xbound[2:end] .> 0)
            xbound = constrain_with_pushlogpdf_grad(c, y)[1]
            @test xbound ≈ [1 - eps(), eps() * ones(n-1)...]
            @test xbound[1] < 1
            @test all(xbound[2:end] .> 0)
        end
    end
end

@testset "JointConstraint" begin
    vtypes = [Vector{Float64}, Vector{Float32}]

    cs = [HMCUtilities.IdentityConstraint(3),
          HMCUtilities.LowerBoundedConstraint(0.0),
          HMCUtilities.UpperBoundedConstraint(1.0),
          HMCUtilities.BoundedConstraint(0.0, 1.0),
          HMCUtilities.UnitVectorConstraint(4),
          HMCUtilities.UnitSimplexConstraint(3)]

    y = Float64[]
    x = Float64[]
    y_exp = Float64[]
    logπx = randn()
    ∇x_logπx = Float64[]
    ∇y_logπy_exp = Float64[]
    logdetJ_exp = sum(cs) do (c)
        yi = randn(HMCUtilities.free_dimension(c))
        xi = constrain(c, yi)
        yi_exp = free(c, xi)
        ∇x_logπxi = randn(length(xi))
        _, pushlogpdf_grad = constrain_with_pushlogpdf_grad(c, yi)
        logdetJi, ∇y_logπyi = pushlogpdf_grad(0.0, ∇x_logπxi)
        push!(y, yi...)
        push!(x, xi...)
        push!(y_exp, yi_exp...)
        push!(∇x_logπx, ∇x_logπxi...)
        push!(∇y_logπy_exp, ∇y_logπyi...)
        return logdetJi
    end
    logπy_exp = logπx + logdetJ_exp

    jc = HMCUtilities.JointConstraint(cs...)

    test_constraint(
        jc,
        x,
        y,
        logπx,
        ∇x_logπx,
        y_exp,
        logπy_exp,
        ∇y_logπy_exp;
        cvtypes=vtypes,
        fvtypes=vtypes,
        test_logdetJ_consistency = false
    )

    @testset "IdentityConstraint merged" begin
        cs = [HMCUtilities.IdentityConstraint(3),
              HMCUtilities.IdentityConstraint(2),
              HMCUtilities.UpperBoundedConstraint(1.0),
              HMCUtilities.IdentityConstraint(2),
              HMCUtilities.IdentityConstraint(1)]
        jc = HMCUtilities.JointConstraint(cs...)

        cs2 = [HMCUtilities.IdentityConstraint(5),
               HMCUtilities.UpperBoundedConstraint(1.0),
               HMCUtilities.IdentityConstraint(3)]
        jc2 = HMCUtilities.JointConstraint(cs...)

        @test jc == jc2
    end
end


"""
    Thirty360Excel()

**US (NASD) 30/360** day count convention, as computed via Microsoft Excel `YEARFRAC` with the basis option of `0`.

This differs from [`Thirty360`](@ref) when:
*  if the start date is the last day of February, then
  -  ``d_1`` is 30, and
  - if the end date is also the last day of February ``d_2`` is also 30.

# Reference
- [Microsoft Excel `YEARFRAC` function](https://support.office.com/en-us/article/yearfrac-function-3844141e-c76d-4143-82b6-208454ddc6a8)
- [David A. Wheeler (2008) "YEARFRAC Incompatibilities between Excel 2007 and OOXML (OXML), and the Definitions Actually Used by Excel 2007"](https://dwheeler.com/yearfrac/excel-ooxml-yearfrac.pdf)
"""
struct Thirty360Excel <: DayCount
end
function frac(startdate::Date, enddate::Date, dc::Thirty360Excel)
    if startdate > enddate
        return -frac(enddate, startdate, dc)
    end

    y1 = year(startdate)
    y2 = year(enddate)
    m1 = month(startdate)
    m2 = month(enddate)
    d1 = day(startdate)
    d2 = day(enddate)
    if d1 >= 30
        d1 = 30
        if d2 >= 30
            d2 = 30
        end
    elseif m1 == 2 && startdate == lastdayofmonth(startdate)
        d1 = 30
        if m2 == 2 && enddate == lastdayofmonth(enddate)
            d2 = 30
        end
    end
    return thirty360(y2-y1,m2-m1,d2-d1)
end

"""
    ActualActualExcel()

**Actual/Actual** day count convention, as computed via Microsoft Excel `YEARFRAC` with the basis option of `1`.

The year fraction is
```math
\\frac{\\text{# of days}}{D}
```
where:
* if start date and end date are in the same calendar year, then ``D`` is the number of days in that calendar year,
* otherwise if the end date is less than or equal to one year after the start date, then
  - ``D = 366`` if the interval includes February 29 (including both the start and end dates), or
  - ``D = 365`` otherwise.
* otherwise it is the average length of the years included in the interval (this does not depend on where the start or end dates fall within those years).

# Reference
- [Microsoft Excel `YEARFRAC` function](https://support.office.com/en-us/article/yearfrac-function-3844141e-c76d-4143-82b6-208454ddc6a8)
- [David A. Wheeler (2008) "YEARFRAC Incompatibilities between Excel 2007 and OOXML (OXML), and the Definitions Actually Used by Excel 2007"](https://dwheeler.com/yearfrac/excel-ooxml-yearfrac.pdf)
"""
struct ActualActualExcel <: DayCount
end
function frac(startdate::Date, enddate::Date, dc::ActualActualExcel)
    if startdate > enddate
        return -frac(enddate, startdate, dc)
    end
    y1 = year(startdate)
    y2 = year(enddate)
    if y1 == y2
        return Dates.value(enddate-startdate) / daysinyear(y1)
    elseif startdate + Year(1) >= enddate
        if (isleapyear(y1) && startdate <= Date(y1,2,29)) ||
            (isleapyear(y2) && Date(y2,2,29) <= enddate)
            return Dates.value(enddate-startdate) / 366
        else
            return Dates.value(enddate-startdate) / 365
        end
    else
        yrange = y1:y2
        return Dates.value(enddate-startdate) / (sum(daysinyear, yrange) / length(yrange))
    end
end


# Function to create matrix from array of arrays. This version is faster than vcat
function fast_mat(S)
    X = Array{Float32, 2}(undef, length(S[1]), length(S))
    for i=1:length(S)
        X[:,i] = S[i]
    end
    X
end

# This function samples episodes from the policy using importance sampling to weight them
function sample_episodes(mdp, policy::Policy, Neps::Int64; max_steps::Int64 = 1000,
                         a_and_p = action_and_probability)
    S, G, avg_ep_return = [], Float32[], 0.
    for i=1:Neps
        s, steps = initialstate(mdp), 0
        Si, Ri, ρi = [], [], []
        while !isterminal(mdp, s)
            push!(Si, convert_s(Array{Float64,1}, s, mdp))
            a, prob = a_and_p(policy, s)
            push!(ρi, policy.action_probability(mdp, s, a) / prob)
            s, r = gen(mdp, s, a)
            push!(Ri, r)
            steps += 1
            steps >= max_steps && break
        end
        steps >= max_steps && println("Episode timeout at ", max_steps, " steps")
        weighted_returns = reverse(cumsum(reverse(Ri))) .* reverse(cumprod(reverse(ρi)))
        push!(S, Si...)
        push!(G, weighted_returns...)
        avg_ep_return += sum(Ri)
    end
    S, G, avg_ep_return = fast_mat(S), G', avg_ep_return/Neps
    @assert size(S,2) == size(G,2)
    @assert size(G,1) == 1
    S, G, avg_ep_return
end


cd("c:\\users\\mgjo\\Box Sync\\Documents\\Research projects\\juliaprogram\\constlab")

using Tensors,PGFPlotsX,LaTeXStrings
using ForwardDiff
using DelimetedFiles

include("materials.jl");

model=2

if model==1
#model=1: linear isotropic elasticity
  nstate=1 #cannot set to 0
  npara  =2
  Emod=1000.0; nu=0.3
  mpara=[Emod,nu]
  tol=1e-3
elseif model==2  #von mises, nonlinear mixed hardening

  nstate=10
  npara  =7
  Emod=200000.0; nu=0.3; Sy=100; Hiso=Emod/2; kappainf=30; Hkin=Emod/2; alphainf=50;
  mpara=[Emod,nu,Sy,Hiso,kappainf,Hkin,alphainf]
  tol=1e-3

elseif model==3  #maxwell viscoelasticity

  nstate=1 #cannot set to 0
  npara  =4
  Emod=200000.0; nu=0.3; kappa=1*10^5; mu=1*10^5
  mpara=[Emod,nu,kappa,mu]
  tol=1e-3

end

stateold=zeros(nstate)

#Define load control
#strainctrlv=[1,0,0,0,0,0,0,0,0]
#strain control, uniaxial stress
unkn=[2,3,4,5,6,7,8,9];
notunkn=[1];
#pure strain control
#unkn=[];
#notunkn=[1,2,3,4,5,6,7,8,9];

nunkn=size(unkn,1);
nnotunkn=size(notunkn,1);
stress_unbalance=zeros(nunkn);

tol_unb=1e-3;

#Define time stepping
nsteps=300;
tend=10.; dt=tend/(nsteps-1);
timehist=0.:dt:tend;

#Strain stepping
ϵ_max =0.005; dϵ=ϵ_max/(nsteps-1)
ϵ_hist=0.:dϵ:ϵ_max;

#Stress stepping
σ_max =500.; dσ=(σ_max/(nsteps-1))
σ_hist=0.:dσ:σ_max;
σ_hist=σ_hist.*0.

aa=zeros(nsteps,13)
for i=1:nsteps
   aa[i,:]=[timehist[i],ϵ_hist[i],0.,0.,0.,0.,0.,σ_hist[i],0.,0.,0.,0.,0.]
end
writedlm("loadhist.txt",aa)

ϵ_out = zeros(nsteps) #Float64[]
σ_out = zeros(nsteps) #Float64[]
σ_out2 = zeros(nsteps) #Float64[]

#
ϵ_out,σ_out,σ_out2=time_hist(nsteps,tol_unb,model,tol,σ_hist,ϵ_hist,nstate,stateold,dt,npara,mpara,unkn,notunkn)

println("stress=",σ_out)

#μ = Emod

#=@pgf Axis({thick, no_marks, legend_pos={north_west}, xmajorgrids, ymajorgrids, ylabel=L"\bullet / E", xlabel=L"\epsilon_{11}",width="12cm", height="8cm"},
    Plot(Coordinates(ϵ_hist, σ_out./Emod)),
    LegendEntry(L"\sigma_{11}"),
    Plot(Coordinates(ϵ_out, σ_out2./Emod)),
    LegendEntry(L"\sigma_{22}")
)
=#

@pgf Axis({thick, no_marks, legend_pos={north_west}, xmajorgrids, ymajorgrids, ylabel=L"{\rm [MPa]}", xlabel=L"\epsilon_{11}",width="12cm", height="8cm"},
    Plot(Coordinates(ϵ_hist, σ_out)),
    LegendEntry(L"\sigma_{11}"),
    Plot(Coordinates(ϵ_out, σ_out2)),
    LegendEntry(L"\sigma_{22}")
)


numbers = [0,14,1,3,7,9]
# numbers = [1,3,2]
# numbers = [0,3,6]
# numbers = [2,1,3]

function get_next_number!(spoken, last_number, current_idx)
    if !(last_number in keys(spoken))
        spoken[last_number] = current_idx - 1
        return 0
    else
        result = current_idx - spoken[last_number] - 1
        spoken[last_number] = current_idx - 1
        return result
    end
end

spoken_history = Dict(n[2] => n[1] for n in enumerate(numbers[1:end-1]))
last_spoken = numbers[end]
for i in length(numbers)+1:30000000
    global spoken_history
    global last_spoken = get_next_number!(spoken_history, last_spoken, i)
end

println(last_spoken)


#
# Test a sparse rosenbrock, i.e. we have many factors but only a few are
# actually active in determining the response from the function.
#
function rosenbrock(x)
  n = length(x)
  return sum( 100*( x[2:n] - x[1:(n-1)].^2 ).^2 + ( x[1:(n-1)] - 1 ).^2 )
end
# FIXME this should be a more efficient implementation, but switching to it would reset all benchmarks
#rosenbrock(x) = sum(i -> 100*abs2(x[i+1] - x[i]^2) + abs2(x[i] - 1), Base.OneTo(length(x)-1))

function sphere(x)
  sum(x.^2)
end
# FIXME this should be a more efficient implementation, but switching to it would reset all benchmarks
#sphere(x) = sum(abs2, x)

mutable struct SparseShiftedProblem
  active_factors::Array{Int, 1}
  shift::Array{Float64, 2}
  func::Function

  function SparseShiftedProblem(n = 1000, k = 20, baseProblem = rosenbrock)
    afs = shuffle(collect(1:n))[1:k]
    new(afs, randn(k,1), baseProblem)
  end
end

function calc(sr::SparseShiftedProblem, x)
  sr.func(x[sr.active_factors] - sr.shift)
end

function sample_points_and_calc_response(n, s, func)
  xs = randn(n, s);
  y = zeros(s)
  [(y[i] = func(xs[:,i])) for i in 1:s]
  return xs, y
end

# Correlation selection is based on correlating each covariate with the response
# and keeping the covariates with the highest correlations with the response.
function correlation_selection(xs, y, maxSelected)
  correlations = y' * xs'
  find_top_largest_correlations = sort(correlations[:], rev=true)[1:maxSelected]
  index_of_largest_correlations = findall(in(find_top_largest_correlations), correlations)
  index_of_largest_correlations
end

# Absolute Correlation selection is based on correlating each covariate with the response
# and keeping the covariates with the highest absolute correlations with the response.
function abs_correlation_selection(xs, y, maxSelected)
  correlations = abs(y' * xs')
  find_top_largest_correlations = sort(correlations[:], rev=true)[1:maxSelected]
  index_of_largest_correlations = findall(in(find_top_largest_correlations), correlations)
  index_of_largest_correlations
end

# Stepwise filtering abs correlation selection. This is based on the idea that if
# there are complex interactive effects among the covariates the covariates that
# are part of an interaction should also have some individual effect. This is similar
# to the assumption of "effect heredity". Thus we do a stepwise filtering of the possible
# covariates to include by first eliminating half of them, then redoing the filtering
# on the interactions between the first selection, and so on.

# Random selection.
function random_selection(xs, y, maxSelected)
  n = size(xs, 1)
  shuffle(collect(1:n))[1:maxSelected]
end

# Call everything once to ensure they have been compiled.
a = rosenbrock([10.0, 4.5])
b = sphere([10.0, 4.5])
p = SparseShiftedProblem(100, 5, rosenbrock)
xs, y = sample_points_and_calc_response(100, 10, (x) -> calc(p, x))
s = random_selection(xs, y, 5)
s2 = correlation_selection(xs, y, 5)

# A problem is characterized mainly by:
#   1. How complex the interrelations are among the active covariates, basically determined by the objective func itself
#   2. Sparseness degree, i.e. how many of all available covariates that actually affects the response (0.0-1.0)
#   3. Dimensionality, the total number of covariates (100-10000)
#
# A selection scheme then saves an incomplete "trail" of evaluations of the objective function
# and uses them to predict which covariates are important. The size of this set is also
# important, as well as how it evolves over time.
#
# So we want to study how well a covariate subset selection (CSS) scheme predicts
# which covariates are important as we vary the four main parameters above.
# Dependent variables:
#   O1. Percentage of missed active covariates over NR repetitions
#   O2. Variation in missed covariates over NR repetitions
#   O3. Average time for a prediction
# Independent variables:
Problems = [rosenbrock, sphere]
Dimensionality = [1000, 10000, 100000]
ActualSparsenessDegree = [0.001, 0.01, 0.10, 0.25]
ArchiveSizeDegree = [0.1, 0.50, 1.00]
GuessedSparsenessDegree = [0.001, 0.01, 0.05, 0.10]

NumRepeats = 10
CovariateSelectionSchemes = [abs_correlation_selection, correlation_selection]
ncss = length(CovariateSelectionSchemes)

csvfile = open("comparing_covariate_selection_schemes.csv", "w")

using HypothesisTests

# Return the name of the best method of two methods given a Wilcoxon test.
# By default we are minimizing so better means smaller average, set minimizing
# to false if maximizing.
function best_of_if_better(x, y, names = ["1", "2"], minimizing = true, pvalue_treshold = 0.05)
  if pvalue(MannWhitneyUTest(x, y)) < pvalue_treshold
    op = minimizing ? < : >
    op(mean(x), mean(y)) ? names[1] : names[2]
  else
    "no_diff"
  end
end

println(csvfile, "CovariateSelectionScheme,Problem,Dimensionality,ActualSparseness,GuessedSparseness,ArchiveSize,MeanTime,StdMissed,MeanMissed")
for problem in Problems
  for dim in Dimensionality
    for asd in ActualSparsenessDegree
      for sizedegree in ArchiveSizeDegree
        for gsd in GuessedSparsenessDegree

          k = int(asd * dim)

          num_selected = int(gsd * dim)
          percent_missed = zeros(NumRepeats, ncss)
          times = zeros(NumRepeats, ncss)
          size = int(sizedegree * dim)

          for ir in 1:NumRepeats
            ssp = SparseShiftedProblem(dim, k, problem)
            xs, y = sample_points_and_calc_response(dim, size, (x) -> calc(ssp, x))

            for cssi in 1:ncss
              css = CovariateSelectionSchemes[cssi]
              tic()
              selected_factors = css(xs, y, num_selected)
              times[ir, cssi] = toq()
              missed = setdiff(Set(ssp.active_factors...), Set(selected_factors...))
              percent_missed[ir, cssi] = length(missed)/k*100.0
            end
          end

          for cssi in 1:ncss
            css = CovariateSelectionSchemes[cssi]
            println(csvfile, css, ",", problem, ",", dim, ",", asd, ",", gsd, ",", size, ",",
              mean(times[:,cssi]), ",",
              std(percent_missed[:,cssi]), ",", mean(percent_missed[:,cssi]))
            flush(stdout)
          end
          ms = mean(percent_missed, 1) # ms = mean(percent_missed, 1)
          best_method = best_of_if_better(percent_missed[:,1], percent_missed[:,2], CovariateSelectionSchemes)
          println(join([problem, dim, asd, gsd, size, minimum(ms), ms[2]-ms[1], best_method], ","))
        end
      end
    end
  end
end

# The conclusions from all these runs are:
#  For both problems:
#    1. Absolute correlation selection is almost always better than correlation selection (as expected)
#  For Sphere:
#    1. Only if the actual sparseness degree is very low (i.e. with < 30 actual active covariates) are the methods very good
#    2. And even so one must guess at least 10-50 times more active covariates
#  For Rosenbrock:
#    1. 1000: Only with an intermediate actual sparseness (0.01) are the methods very good
#    2. 10000: Only with very low actual sparseness (0.001) are the methods very good
#    3. 10000: Increasing the number of samples has quite a large effect i.e. lowers the missed percentage
#
# What is not clear is how good the method needs to be to actually give an advantage since they
# are applied in rounds. Maybe a small improvement can have quite a large effect in speeding up convergence.
#
# OTOH, we already have information about the correlation of different covariates in the diagonal
# of the covar matrix so maybe these additional selection algs are not needed...