text:
spin_info.jl
spin_trotter.jl
spin_model.jl
spin_classical.jl
spin_quantum.jl
spin_classical_ising.jl
spin_classical_potts.jl
spin_classical_clock.jl
spin_quantum_ising.jl

functions from SpinInfo:
# spin_info.jl
	isZeroTemperature(spininfo)
	isClassical(spininfo)
	getModelname(spininfo)
	getStates(spininfo)
	isSymmetricFactorization(spininfo)
	getTemperature(spininfo)
	getExternalfield(spininfo)
	getEnvParameters(spininfo)
	setTemperature!(spininfo, temperature)
	setExternalfield!(spininfo, externalfield)
	setEnvParameters!(spininfo, temperature, externalfield)
	testSpinInfo()

functions from TrotterInfo:
# spin_trotter.jl
	TrotterInfo,
	getTrotterparameter,
	getTrotteriteration,
	getTrotterlayers,
	getBeta,
	getTemperature,
	setTrotterparameter!,
	setTrotterIteration!,
	testTrotterInfo

functions from SpinModel:
# spin_model.jl
	buildSpinSystem(spininfo)
	buildSpinSystem(spininfo, TrotterInfo)
	getMeasureOperator
	getFactorW
	isPottsModel
	# <functions for Hamiltonian>
	getHamiltonian(spinmodel)
	# <test function>
	testSpinModel()
	# <inherit from spininfo>
	#--- one input (spininfo)
	isZeroTemperature
	isClassical
	getModelname
	getStates
	isSymmetricFactorization
	getTemperature
	getExternalfield(spininfo)
	getEnvParameters(spininfo)
	#--- two input (spininfo, T)
	setTemperature!(spininfo, temperature)
	setExternalfield!(spininfo, externalfield)
	#--- three input (spininfo, T1, T2)
	setEnvParameters!(spininfo, temperature, externalfield)

QunatumSpinModel's field:
	#---- SpinInfo
	qunatumOrClassical
	numberOfState
	factorization
	modelname
	temperature -> zero or finite
	externalfield
	#----
	#---- TrotterInfo
	trotterparameter -> tau
	trotteriteration ->
		trotterlayers as a functions
	#----
	#= to think about
	trotterinfo & spininfo compatable
	temperature and trotteriteration -> initiate!
	trotteriteration and the whole iteration <- coming from simulator

	=#

	factorW
	factorWp

QunatumSpinModel's functions:
	getFactorWp
	# <inherit from trotterinfo>
	getTrotterparameter,
	getTrotteriteration,
	getTrotterlayers,
	getBeta,
	setTrotterparameter!,
	setTrotterIteration!,

	# <check info>
	checkInfo!(spininfo, trotterinfo)

QuantumIsingModel's functions:
	initialize!
	getFactorW
	getFactorWp
	getMeasureOperator
	getStates
	isSymmetricFactorization
	setTemperature!
	setEnvParameters!
	testQuantumIsingModel()
#----------------------------
#----------------------------
#----------------------------

ClassicalSpinModel's field:
	#---- SpinInfo
	qunatumOrClassical
	numberOfState
	factorization
	modelname
	temperature
	externalfield
	#----

	factorW

ClassicalSpinModel's functions:
	getBeta

	#----------------------------
	change the name of factorW and factorWp
	and according functions:
	spin_info.jl
	spin_model.jl
	spin_classical.jl
	spin_classical_ising.jl
	spin_classical_potts.jl
	spin_classical_clock.jl
	spin_trotter.jl
	spin_quantum.jl
	spin_quantum_ising.jl

	#=
		FUNCTIONS these are for spinmodel
		Hamiltonian::Array{T,2} # externalfield, numberOfState
		BoltzmannWeight::Array{T,2} # temperature, hamiltonian
		factorW::Array{T,2} # temperature, externalfield, model
			# hamiltonian, temperature
	=#


text:
##############################################################################
##
## Solution method: minimize by sparse least square optimization
##
##############################################################################

##############################################################################
##
## Factor Model (no regressors)
##
##############################################################################

function fit!(t::Union{Type{Val{:levenberg_marquardt}}, Type{Val{:dogleg}}},
                         fp::FactorModel,
                         fs::FactorSolution;
                         maxiter::Integer = 10_000,
                         tol::Real = 1e-9,
                         lambda::Number = 0.0)
    # initialize
    iter = 0
    converged = true
    fullrank = rank(fp)
    fp = FactorModel(deepcopy(fp.y), fp.sqrtw, fp.idrefs, fp.timerefs, 1)
    N = size(fs.idpool, 1)
    T = size(fs.timepool, 1)
    fg = FactorGradient(fp,
                        FactorSolution(Array(Float64, N), Array(Float64, T)),
                        FactorSolution(Array(Float64, N), Array(Float64, T))
                       )
    nls = LeastSquaresOptim.LeastSquaresProblem(
                view(fs, :, 1),
                similar(fp.y),
                (x,y) -> f!(fp, x, y),
                fg,
                g!)
    if t == Val{:levenberg_marquardt}
        optimizer = LeastSquaresOptim.LevenbergMarquardt()
    else
        optimizer = LeastSquaresOptim.Dogleg()
    end
    full = LeastSquaresOptim.LeastSquaresProblemAllocated(nls, optimizer, LeastSquaresOptim.LSMR())
    for r in 1:fullrank
        fsr = view(fs, :, r)
        full.x = fsr
        result = LeastSquaresOptim.optimize!(full,
            xtol = 1e-32, grtol = 1e-32, ftol = tol,  iterations = maxiter)
        iter += result.mul_calls
        converged = result.converged && converged
        if r < fullrank
            rescale!(fsr)
            subtract_factor!(fp, fsr)
        end
    end
    # rescale factors and loadings so that factors' * factors = Id
    return rescale(fs), iter, converged
end

##############################################################################
##
## Interactive FixedEffectModel
##
##############################################################################

function fit!{Rank}(t::Union{Type{Val{:levenberg_marquardt}}, Type{Val{:dogleg}}},
                         fp::InteractiveFixedEffectsModel{Rank},
                         fs::InteractiveFixedEffectsSolution{Rank};
                         maxiter::Integer = 10_000,
                         tol::Real = 1e-9,
                         lambda::Number = 0.0)
    scaleb = vec(sumabs2(fp.X, 1))
    fsT = InteractiveFixedEffectsSolutionT(fs.b, fs.idpool', fs.timepool')
    fg = InteractiveFixedEffectsGradientT(fp,
                similar(fsT),
                InteractiveFixedEffectsSolutionT(scaleb, similar(fsT.idpool), similar(fsT.timepool)))
    nls = LeastSquaresOptim.LeastSquaresProblem(fsT, similar(fp.y), (x,y) -> f!(fp, x, y), fg, g!)
    if t == Val{:levenberg_marquardt}
        optimizer = LeastSquaresOptim.LevenbergMarquardt()
    else
        optimizer = LeastSquaresOptim.Dogleg()
    end
    full = LeastSquaresOptim.LeastSquaresProblemAllocated(nls, optimizer, LeastSquaresOptim.LSMR())
    temp = similar(fp.y)
    result = LeastSquaresOptim.optimize!(full;
            xtol = 1e-32, grtol = 1e-32, ftol = tol,  iterations = maxiter)
    # rescale factors and loadings so that factors' * factors = Id
    fs = InteractiveFixedEffectsSolution(fsT.b, fsT.idpool', fsT.timepool')
    return rescale(fs), result.mul_calls, result.converged
end

##############################################################################
##
## Methods used in optimize! in LeastSquares
##
##############################################################################

function vecdot{T}(fs1::AbstractArray{T}, fs2::AbstractArray{T})
    out = zero(typeof(one(T) * one(T)))
    @inbounds @simd for i in eachindex(fs1)
        out += fs1[i] * fs2[i]
    end
    return out
end

function vecdot{T}(x::AbstractArray{T}, y::AbstractArray{T}, w::AbstractArray{T})
    out = zero(typeof(one(T) * one(T)))
    @inbounds @simd for i in eachindex(x)
        out += w[i] * x[i] * y[i]
    end
    return out
end

for t in (FactorSolution, InteractiveFixedEffectsSolution, InteractiveFixedEffectsSolutionT, HalfInteractiveFixedEffectsSolution)
    vars = [:(fill!(x.$field, α)) for field in fieldnames(t)]
    expr = Expr(:block, vars...)
    @eval begin
        function fill!(x::$t, α::Number)
             $expr
            return x
        end
    end

    vars = [:(scale!(x.$field, α)) for field in fieldnames(t)]
    expr = Expr(:block, vars...)
    @eval begin
        function scale!(x::$t, α::Number)
            $expr
            return x
        end
    end

    vars = [:(clamp!(x.$field, α, β)) for field in fieldnames(t)]
    expr = Expr(:block, vars...)
    @eval begin
        function clamp!(x::$t, α::Number, β::Number)
            $expr
            return x
        end
    end

    vars = [:(copy!(x1.$field, x2.$field)) for field in fieldnames(t)]
    expr = Expr(:block, vars...)
    @eval begin
        function copy!(x1::$t, x2::$t)
            $expr
            return x1
        end
    end

    vars = [:(axpy!(α, x1.$field, x2.$field)) for field in fieldnames(t)]
    expr = Expr(:block, vars...)
    @eval begin
        function axpy!(α::Number, x1::$t, x2::$t)
            $expr
            return x2
        end
    end

    vars = [:(map!(f, x1.$field, map(x -> x.$field, x2)...)) for field in fieldnames(t)]
    expr = Expr(:block, vars...)
    @eval begin
        function map!(f, x1::$t,  x2::$t...)
            $expr
            return x1
        end
    end
    vars = [:(vecdot(x1.$field, x2.$field)) for field in fieldnames(t)]
    expr = Expr(:call, :+, vars...)
    @eval begin
        function dot(x1::$t, x2::$t)
            $expr
        end
    end

    vars = [:(vecdot(x1.$field, x2.$field, w.$field)) for field in fieldnames(t)]
    expr = Expr(:call, :+, vars...)
    @eval begin
        function dot(x1::$t, x2::$t, w::$t)
            $expr
        end
    end

    vars = [:(sumabs2(x.$field)) for field in fieldnames(t)]
    expr = Expr(:call, :+, vars...)
    @eval begin
        function sumabs2(x::$t)
            $expr
        end
    end

    vars = [:(length(x.$field)) for field in fieldnames(t)]
    expr = Expr(:call, :+, vars...)
    @eval begin
        function length(x::$t)
            $expr
        end
    end

    vars = [:(similar(x.$field)) for field in fieldnames(t)]
    expr = Expr(:call, t, vars...)
    @eval begin
        function similar(x::$t)
            $expr
        end
    end

    vars = [:(maxabs(x.$field)) for field in fieldnames(t)]
    expr = Expr(:call, :max, vars...)
    @eval begin
        function maxabs(x::$t)
            $expr
        end
    end
end

norm(fs::AbstractFactorSolution) = sqrt(sumabs2(fs))
eltype(fg::AbstractFactorSolution) = Float64

function safe_scale!(x, β)
    if β != 1
        β == 0 ? fill!(x, zero(eltype(x))) : scale!(x, β)
    end
end

##############################################################################
##
## Factor Gradient
##
##############################################################################

abstract AbstractFactorGradient{T}

type FactorGradient{Rank, W, Rid, Rtime, Tid, Ttime, sTid, sTtime} <: AbstractFactorGradient{Rank}
    fp::FactorModel{Rank, W, Rid, Rtime}
    fs::FactorSolution{Rank, Tid, Ttime}
    scalefs::FactorSolution{Rank, sTid, sTtime}
end

function size(fg::FactorGradient, i::Integer)
    if i == 1
        length(fg.fp.y)
    elseif i == 2
      length(fg.fs.idpool) + length(fg.fs.timepool)
    end
end

type InteractiveFixedEffectsGradientT{Rank, W, Rid, Rtime, Tb, Tid, Ttime, sTb, sTid, sTtime} <: AbstractFactorGradient{Rank}
    fp::InteractiveFixedEffectsModel{Rank, W, Rid, Rtime}
    fs::InteractiveFixedEffectsSolutionT{Rank, Tb, Tid, Ttime}
    scalefs::InteractiveFixedEffectsSolutionT{Rank, sTb, sTid, sTtime}
end

function size(fg::InteractiveFixedEffectsGradientT, i::Integer)
    if i == 1
        length(fg.fp.y)
    elseif i == 2
      length(fg.fs.b) + length(fg.fs.idpool) + length(fg.fs.timepool)
    end
end

Base.rank{Rank}(f::AbstractFactorGradient{Rank}) = Rank
size(fg::AbstractFactorGradient) = (size(fg, 1), size(fg, 2))
eltype(fg::AbstractFactorGradient) = Float64
LeastSquaresOptim.colsumabs2!(fs::AbstractFactorSolution, fg::AbstractFactorGradient) = copy!(fs, fg.scalefs)

@generated function A_mul_B!{Rank}(α::Number, fg::AbstractFactorGradient{Rank}, fs::AbstractFactorSolution{Rank}, β::Number, y::AbstractVector{Float64})
    if Rank == 1
        ex = quote
            out += (fg.fs.idpool[idi] * fs.timepool[timei]
                + fg.fs.timepool[timei] * fs.idpool[idi])
        end
    else
        ex = quote
            @nexprs $Rank r -> begin
                out += (fg.fs.idpool[r, idi] * fs.timepool[r, timei]
                         + fg.fs.timepool[r, timei] * fs.idpool[r, idi])
            end
        end
    end
    quote
        mα = convert(Float64, -α)
        A_mul_B!_X(mα, fg.fp, fs, β, y)
        @fastmath @inbounds @simd for i in 1:length(y)
            timei = fg.fp.timerefs[i]
            idi = fg.fp.idrefs[i]
            out = 0.0
            $ex
            y[i] += mα * fg.fp.sqrtw[i] * out
        end
        return y
    end
end

function A_mul_B!_X(α::Number, fp::InteractiveFixedEffectsModel, fs::InteractiveFixedEffectsSolutionT, β::Number, y::AbstractVector{Float64})
    Base.BLAS.gemm!('N', 'N', α, fp.X, fs.b, convert(Float64, β), y)
end
function A_mul_B!_X(::Number, ::FactorModel, ::FactorSolution, β::Number, y::AbstractVector{Float64})
    safe_scale!(y, β)
end

@generated function Ac_mul_B!{Rank}(α::Number, fg::AbstractFactorGradient{Rank}, y::AbstractVector{Float64}, β::Number, fs::AbstractFactorSolution{Rank})
    if Rank == 1
        ex = quote
           fs.idpool[idi] += sqrtwi * fg.fs.timepool[timei]
           fs.timepool[timei] += sqrtwi * fg.fs.idpool[idi]
        end
    else
        ex = quote
            @nexprs $Rank r -> begin
                fs.idpool[r, idi] += sqrtwi * fg.fs.timepool[r, timei]
                fs.timepool[r, timei] += sqrtwi * fg.fs.idpool[r, idi]
            end
        end
    end
    quote
        mα = convert(Float64, -α)
        Ac_mul_B!_X(mα, fg.fp, y, β, fs)
        @inbounds @simd for i in 1:length(y)
            sqrtwi = mα * y[i] * fg.fp.sqrtw[i]
            idi = fg.fp.idrefs[i]
            timei = fg.fp.timerefs[i]
            $ex
        end
        return fs
    end
end

function Ac_mul_B!_X(α::Number, fp::InteractiveFixedEffectsModel, y::AbstractVector{Float64}, β::Number, fs::InteractiveFixedEffectsSolutionT)
    safe_scale!(fs, β)
    @inbounds @simd for k in 1:length(fs.b)
        out = zero(Float64)
        for i in 1:length(y)
           out += y[i] * fp.X[i, k]
        end
        fs.b[k] += α * out
    end
end

function Ac_mul_B!_X(::Number, ::FactorModel, ::AbstractVector{Float64}, β::Number, fs::FactorSolution) safe_scale!(fs, β)
end

##############################################################################
##
## f! for LeastSquaresOptim
##
##############################################################################

@generated function f!{Rank}(fp::AbstractFactorModel{Rank},
    fs::AbstractFactorSolution{Rank}, out::Vector{Float64})
    if Rank == 1
        ex = quote
            out[i] -= sqrtwi * fs.idpool[idi] * fs.timepool[timei]
        end
    else
        ex = quote
           @nexprs $Rank r -> begin
               out[i] -= sqrtwi * fs.idpool[r, idi] * fs.timepool[r, timei]
           end
        end
    end
    quote
        copy!(out, fp.y)
        A_mul_B!_X(-1.0, fp, fs, 1.0, out)
        @fastmath @inbounds @simd for i in 1:length(out)
            sqrtwi = fp.sqrtw[i]
            idi = fp.idrefs[i]
            timei = fp.timerefs[i]
            $ex
        end
        return out
    end
end

##############################################################################
##
## g! for LeastSquaresOptim
##
##############################################################################

@generated function g!{Rank}(fs::AbstractFactorSolution{Rank}, fg::AbstractFactorGradient{Rank})
    if Rank == 1
        ex = quote
            fg.scalefs.idpool[idi] += abs2(sqrtwi * fg.fs.timepool[timei])
            fg.scalefs.timepool[timei] += abs2(sqrtwi * fg.fs.idpool[idi])
        end
    else
        ex = quote
            @nexprs $Rank r -> begin
                fg.scalefs.idpool[r, idi] += abs2(sqrtwi * fg.fs.timepool[r, timei])
                fg.scalefs.timepool[r, timei] += abs2(sqrtwi * fg.fs.idpool[r, idi])
            end
        end
    end
    quote
        copy!(fg.fs, fs)
        fill!(fg.scalefs.idpool, zero(Float64))
        fill!(fg.scalefs.timepool, zero(Float64))
        @inbounds @simd for i in 1:length(fg.fp.y)
            sqrtwi = fg.fp.sqrtw[i]
            idi = fg.fp.idrefs[i]
            timei = fg.fp.timerefs[i]
            $ex
        end
    end
end

ext: jl


text:
using Flux, DiffEqFlux, DifferentialEquations, Plots

u0 = Float32[2.; 0.]
datasize = 30
tspan  = (0.0f0,1.5f0)
function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
t = range(tspan[1],tspan[2],length=datasize)
prob = ODEProblem(trueODEfunc,u0,tspan)
ode_data = Array(solve(prob,Tsit5(),saveat=t))
dudt = Chain(x -> x.^3,
             Dense(2,50,tanh),
             Dense(50,2))
ps = Flux.params(dudt)
n_ode = x->neural_ode(dudt,x,tspan,Tsit5(),saveat=t,reltol=1e-7,abstol=1e-9)
pred = n_ode(u0) # Get the prediction using the correct initial condition
scatter(t,ode_data[1,:],label="data")
scatter!(t,Flux.data(pred[1,:]),label="prediction")

function predict_n_ode()
  n_ode(u0)
end
loss_n_ode() = sum(abs2,ode_data .- predict_n_ode())
data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function ()
  display(loss_n_ode())
  cur_pred = Flux.data(predict_n_ode())
  pl = scatter(t,ode_data[1,:],label="data")
  scatter!(pl,t,cur_pred[1,:],label="prediction")
  display(plot(pl))
end

cb()
Flux.train!(loss_n_ode, ps, data, opt, cb = cb)


text:
# Model
include_model("hopper")
model_ft = free_time_model(model)

# Horizon
T = 101

# Time step
tf = 1.0
h = tf / (T - 1)

# Bounds
_uu = Inf * ones(model_ft.m)
_uu[model_ft.idx_u] = model_ft.uU
_uu[end] = 2.0 * h
_ul = zeros(model_ft.m)
_ul[model_ft.idx_u] = model_ft.uL
_ul[end] = 0.5 * h
ul, uu = control_bounds(model_ft, T, _ul, _uu)

# Initial and final states
z_h = 0.25
q1 = [0.0, 0.5 + z_h, 0.0, 0.25]

xl, xu = state_bounds(model_ft, T,
		[model_ft.qL; model_ft.qL],
		[model_ft.qU; model_ft.qU],
        x1 = [q1; Inf * ones(model.nq)])

# Objective
include_objective("velocity")
obj_tracking = quadratic_time_tracking_objective(
    [Diagonal(zeros(model_ft.n)) for t = 1:T],
    [Diagonal([1.0, 1.0, zeros(model_ft.m - model_ft.nu - 1)..., 0.0]) for t = 1:T-1],
    [zeros(model_ft.n) for t = 1:T],
    [zeros(model_ft.m) for t = 1:T],
    1.0)
obj_contact_penalty = PenaltyObjective(1.0e5, model_ft.m - 1)
obj_velocity = velocity_objective(
    [Diagonal(1.0 * ones(model_ft.nq)) for t = 1:T-1],
    model_ft.nq,
    h = h,
    idx_angle = collect([3]))
obj = MultiObjective([obj_tracking, obj_contact_penalty, obj_velocity])

# Constraints
include_constraints(["free_time", "contact", "loop"])
con_free_time = free_time_constraints(T)
con_contact = contact_constraints(model_ft, T)
con_loop = loop_constraints(model, (1:model.n), 1, T)
con = multiple_constraints([con_free_time, con_contact, con_loop])

# Problem
prob = trajectory_optimization_problem(model_ft,
               obj,
               T,
               xl = xl,
               xu = xu,
               ul = ul,
               uu = uu,
               con = con)

# Trajectory initialization
x0 = [[q1; q1] for t = 1:T] # linear interpolation on state
u0 = [[1.0e-3 * rand(model_ft.m-1); h] for t = 1:T-1] # random controls

# Pack trajectories into vector
z0 = pack(x0, u0, prob)

#NOTE: may need to run examples multiple times to get good trajectories
# Solve nominal problem

optimize = true
if optimize
	# include_snopt()
	@time z̄ , info = solve(prob, copy(z0),
		nlp = :ipopt,
		tol = 1.0e-5, c_tol = 1.0e-5, mapl = 5,
		time_limit = 60)
	@show check_slack(z̄, prob)
	x̄, ū = unpack(z̄, prob)
	tf, t, h̄ = get_time(ū)

	# projection
	Q = [Diagonal(ones(model_ft.n)) for t = 1:T]
	R = [Diagonal(0.1 * ones(model_ft.m)) for t = 1:T-1]
	x_proj, u_proj = lqr_projection(model_ft, x̄, ū, h̄[1], Q, R)

	@show tf
	@show h̄[1]
	@save joinpath(pwd(), "examples/trajectories/hopper_vertical_gait.jld2") x̄ ū h̄ x_proj u_proj
else
	@load joinpath(pwd(), "examples/trajectories/hopper_vertical_gait.jld2") x̄ ū h̄ x_proj u_proj
end

using Plots
plot(hcat(ū...)[1:2, :]',
    linetype = :steppost,
    label = "",
    color = :red,
    width = 2.0)

plot!(hcat(u_proj...)[1:2, :]',
    linetype = :steppost,
    label = "", color = :black)

plot(hcat(state_to_configuration(x̄)...)',
    color = :red,
    width = 2.0,
	label = "")

plot!(hcat(state_to_configuration(x_proj)...)',
    color = :black,
	label = "")

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model_ft, state_to_configuration(x_proj), Δt = h̄[1])

# Reference trajectory
x_track = deepcopy(x_proj)
u_track = deepcopy(u_proj)
for i = 1:5
	x_track = [x_track..., x_track[2:T]...]
	u_track = [u_track..., u_track[1:T-1]...]
end

T_track = length(x_track)

# x_shift = 0.25
# x_track_shift = deepcopy(x_track)
# for t = 1:T_track
#     x_track_shift[t] = x_track_shift[t] + [x_shift; 0.0; 0.0; 0.0; x_shift; 0.0; 0.0; 0.0]
# end
# x_track_shift
plot(hcat(state_to_configuration(x_track)...)',
    color = :black,
	label = "")

# plot(hcat(state_to_configuration(x_track_shift)...)',
#     color = :black,
# 	label = "")

plot(hcat(u_track...)',
	color = :black,
	label = "")

K, P = tvlqr(model_ft, x_track, u_track, h̄[1],
	[Diagonal([1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]) for t = 1:T_track],
	[Diagonal(1.0 * ones(model_ft.m)) for t = 1:T_track - 1])

# K_shift, P_shift = tvlqr(model_ft, x_track_shift, u_track, h̄[1],
# 	[Diagonal(1.0 * ones(model_ft.n)) for t = 1:T_track],
# 	[Diagonal(1.0 * ones(model_ft.m)) for t = 1:T_track - 1])

K_vec = [vec(K[t]) for t = 1:T_track-1]
P_vec = [vec(P[t]) for t = 1:T_track-1]

plot(hcat(K_vec...)', label = "")
plot(hcat(P_vec...)', label = "")

include(joinpath(pwd(), "src/contact_simulator/simulator_variable_time_step.jl"))
model_sim = model

T_sim = 1 * T_track
tf_track = h̄[1] * (T_track - 1)
t_sim = range(0, stop = tf_track, length = T_sim)
t_ctrl = range(0, stop = tf_track, length = T_track)
h_sim = tf_track / (T_sim - 1)
x_track_stack = hcat(x_track...)

x_sim = [copy(x_proj[1])]
u_sim = []
T_horizon = T_sim - T
x_shift = 0.0
x_shift_vec = [x_shift; 0.0; 0.0; 0.0; x_shift; 0.0; 0.0; 0.0]

for tt = 1:T_horizon-1
	t = t_sim[tt]
	i = searchsortedlast(t_ctrl, t)
	println("t: $t")
	println("	i: $i")
	# ii = searchsortedlast(t_sim, t)

	# x_cubic = zeros(model_sim.n)
	# for j = 1:model_sim.n
	# 	interp_cubic = CubicSplineInterpolation(t_ctrl, x_track_stack[j, :])
	# 	x_cubic[j] = interp_cubic(t)
	# end

	# push!(u_sim, u_track[i][1:end-1] - K[i] * (x_sim[end] - x_track[i]))

	w0 = (tt == 101 ? 1000.0 * [1.0; 0.0; 0.0; 0.0] : 1.0e-5 * randn(model.nq))

	_q_sim = [x_sim[end][1:nq], x_sim[end][nq .+ (1:nq)]]
	_v_sim = [(x_sim[end][nq .+ (1:nq)] - x_sim[end][1:nq]) / h̄[1]]

	d = 3
	_h = h̄[1] / convert(Float64, d)
	for p = 1:d

		# x_cubic = zeros(model_sim.n)
		# for j = 1:model_sim.n
		# 	interp_cubic = CubicSplineInterpolation(t_ctrl, x_track_stack[j, :])
		# 	x_cubic[j] = interp_cubic(t + p * _h)
		# end

		ii = max(1, searchsortedlast(t_sim, t - h̄[1] +  p * _h))
		_x = [x_sim[ii][(1:nq)]; _q_sim[end]]

		_q = step_contact(model,
			_v_sim[end], _q_sim[end-1], _q_sim[end],
			(u_track[i][1:end-1] - K[i] * (_x - (x_track[i] + 0.0 * x_shift_vec)))[1:2],
			# u_sim[end][1:2],
			p == 1 ? 0.0 * w0 : 0.0 * randn(model.nq),
			_h)
		_v = (_q - _q_sim[end]) / _h

		push!(_q_sim, _q)
		push!(_v_sim, _v)
	end

	push!(x_sim, [x_sim[end][nq .+ (1:nq)]; _q_sim[end]])
end

plot(hcat(state_to_configuration(x_track[1:1:T_horizon])...)',
    labels = "", legend = :bottomleft,
    width = 2.0, color = ["red" "green" "blue" "orange"], linestyle = :dash)

plot!(hcat(state_to_configuration(x_sim[1:1:T_horizon])...)',
    labels = "", legend = :bottom,
    width = 1.0, color = ["red" "green" "blue" "orange"])

# plot(hcat(u_track...)[1:2, :]',
# 	width = 2.0,
# 	linetype = :steppost,
# 	linestyle = :dash)
#
# plot!(hcat(u_sim...)[1:2, :]',
# 	width = 1.0,
# 	linetype = :steppost)
#
# plot(hcat(u_track...)[3:5, 100:200]',
# 	width = 2.0,
# 	color = ["red" "green" "blue"],
# 	linetype = :steppost,
# 	linestyle = :dash)
#
# plot!(hcat(u_sim...)[3:5, 100:200]',
# 	width = 1.0,
# 	color = ["red" "green" "blue"],
# 	linetype = :steppost)

vis = Visualizer()
render(vis)
visualize!(vis, model_ft, state_to_configuration(x_sim), Δt = h̄[1])

# # Simulator
# function dynamics(model::Hopper, v1, q1, q2, q3, u, λ, b, w, h)
# 	(M_func(model, q1) * v1
# 	- M_func(model, q2) * (SVector{4}(q3) - SVector{4}(q2)) / h
# 	+ h * (transpose(B_func(model, q3)) * SVector{2}(u)
# 	+ transpose(N_func(model, q3)) * SVector{1}(λ)
# 	+ transpose(P_func(model, q3)) * SVector{2}(b)
# 	- G_func(model, q2)))
# end
#
# function maximum_dissipation(model::Hopper, q2, q3, ψ, η, h)
# 	ψ_stack = ψ[1] * ones(model.nb)
# 	return P_func(model, q3) * (q3 - q2) / h + ψ_stack - η
# end
#
# function friction_cone(model::Hopper, λ, b)
# 	return @SVector [model.μ * λ[1] - sum(b)]
# end
#
# mutable struct MOISimulator <: MOI.AbstractNLPEvaluator
#     num_var::Int                 # number of decision variables
#     num_con::Int                 # number of constraints
#     primal_bounds
#     constraint_bounds
#
# 	model
# 	slack_penalty
# 	v1
# 	q1
# 	q2
# 	u
# 	w
# 	h
# end
#
# function simulator_problem(model, v1, q1, q2, u, w, h; slack_penalty = 1.0e5)
# 	num_var = model.nq + model.nc + model.nb + model.nc + model.nb + model.ns
# 	num_con = model.nq + model.nc + model.nc + model.nb + 3
#
# 	zl = zeros(num_var)
# 	zl[1:model.nq] .= -Inf
# 	zu = Inf * ones(num_var)
#
# 	cl = zeros(num_con)
# 	cu = zeros(num_con)
# 	cu[model.nq .+ (1:model.nc)] .= Inf
# 	cu[model.nq + model.nc .+ (1:model.nc)] .= Inf
# 	cu[model.nq + model.nc + model.nc + model.nb + 1] = Inf
# 	cu[model.nq + model.nc + model.nc + model.nb + 2] = Inf
# 	cu[model.nq + model.nc + model.nc + model.nb + 3] = Inf
#
# 	MOISimulator(num_var,
# 		num_con,
# 		(zl, zu),
# 		(cl, cu),
# 		model,
# 		slack_penalty,
# 		v1,
# 		q1,
# 		q2,
# 		u,
# 		w,
# 		h)
# end
# prob_sim = simulator_problem(model, ones(nq), ones(nq), ones(nq), zeros(nu), ones(nq), h)
# z0 = rand(prob_sim.num_var)
#
# function MOI.eval_objective(prob::MOISimulator, x)
#     return prob.slack_penalty * x[prob.num_var]
# end
#
# MOI.eval_objective(prob_sim, z0)
#
# function MOI.eval_objective_gradient(prob::MOISimulator, grad_f, x)
#     grad_f .= 0.0
# 	grad_f[prob.num_var] = prob.slack_penalty
# 	return nothing
# end
#
# ∇obj = zeros(prob_sim.num_var)
# MOI.eval_objective_gradient(prob_sim, ∇obj, z0)
#
# function MOI.eval_constraint(prob::MOISimulator, g, x)
# 	model = prob.model
# 	q3 = view(x, 1:model.nq)
# 	λ = view(x, model.nq .+ (1:model.nc))
# 	b = view(x, model.nq + model.nc .+ (1:model.nb))
# 	ψ = view(x, model.nq + model.nc + model.nb .+ (1:model.nc))
# 	η = view(x, model.nq + model.nc + model.nb + model.nc .+ (1:model.nb))
# 	s = view(x, model.nq + model.nc + model.nb + model.nc + model.nb .+ (1:model.ns))
# 	# λ = 0
# 	# b = 0
#     g[1:model.nq] = dynamics(model,
# 		prob.v1, prob.q1, prob.q2, q3, prob.u, λ, b, prob.w, prob.h)
# 	g[model.nq .+ (1:model.nc)] = ϕ_func(model, q3)
# 	g[model.nq + model.nc .+ (1:model.nc)] = friction_cone(model, λ, b)
# 	g[model.nq + model.nc + model.nc .+ (1:model.nb)] = maximum_dissipation(model,
# 		prob.q2, q3, ψ, η, prob.h)
# 	g[model.nq + model.nc + model.nc + model.nb + 1] = s[1] - (λ' * ϕ_func(model, q3))[1]
# 	g[model.nq + model.nc + model.nc + model.nb + 2] = s[1] - (ψ' * friction_cone(model, λ, b))[1]
# 	g[model.nq + model.nc + model.nc + model.nb + 3] = s[1] - (η' * b)[1]
#
#     return nothing
# end
# prob_sim.h
# c0 = zeros(prob_sim.num_con)
# MOI.eval_constraint(prob_sim, c0, z0)
#
# function MOI.eval_constraint_jacobian(prob::MOISimulator, jac, x)
#     con!(g, z) = MOI.eval_constraint(prob, g, z)
# 	ForwardDiff.jacobian!(reshape(jac, prob.num_var, prob.num_con), con!, zeros(prob.num_con), x)
#     return nothing
# end
#
# ∇c0 = vec(zeros(prob_sim.num_con, prob_sim.num_var))
# MOI.eval_constraint_jacobian(prob_sim, ∇c0, z0)
#
# function sparsity_jacobian(prob::MOISimulator)
#     row = []
#     col = []
#
#     r = (1:prob.num_con)
#     c = (1:prob.num_var)
#
#     row_col!(row, col, r, c)
#
#     return collect(zip(row, col))
# end
#
# sparsity_jacobian(prob_sim)
#
# @time solve(prob_sim, z0)
#
# function step_contact(model, v1, q1, q2, u, w, h)
#     prob = simulator_problem(model, v1, q1, q2, u, w, h)
#     z0 = [copy(q2); 1.0e-5 * ones(model.nc + model.nb + model.nc + model.nb + model.ns)]
#     @time z , info = solve(prob, copy(z0), tol = 1.0e-8, c_tol = 1.0e-8, mapl = 0)
#
#     @assert z[end] < 1.0e-8
#
#     return z[1:model.nq]
# end
#
q_proj = state_to_configuration(x_proj)
q_sim = [q_proj[1], q_proj[2]]
v_sim = [(q_proj[2] - q_proj[1]) / h̄[1]]

for t = 1:T-1
	# x rate
	_q_sim = [q_sim[end-1], q_sim[end]]
	_v_sim = [v_sim[end]]

	d = 10
	_h = h̄[1] / convert(Float64, d)
	for i = 1:d
		_q = step_contact(model,
			_v_sim[end], _q_sim[end-1], _q_sim[end], u_proj[t][model.idx_u], zeros(model.d), _h)
		_v = (_q - _q_sim[end]) / _h

		push!(_q_sim, _q)
		push!(_v_sim, _v)
	end

	push!(q_sim, _q_sim[end])
	push!(v_sim, _v_sim[end])
end

plot(hcat(q_proj...)',
    color = :red,
    width = 2.0,
	label = "")
plot!(hcat(q_sim...)',
    color = :black,
    width = 1.0,
	label = "")

plot(hcat(v_sim...)',
    color = :black,
    width = 1.0,
	label = "")


text:
include("twostep_utils.jl")
using LinearAlgebra, PyPlot, MultivariateStats, BSON

### load model ###
using BSON: @load
@load "models/twostep.bson" m

##### lets do some evaluation ####

function common(action, state)
    if (action == 1 && Bool(state[2]))
        return 1
    elseif (action == 2 && Bool(state[3]))
        return 1
    else
        return 0
    end
end

test_as = zeros(T, 0)
test_rews = zeros(T, 0)

types = []
stays = []
for i = 1:1000
    L, ys, rews, as, ps, hs, states = run_episode(m, hidden = true)
    global test_as = [test_as as]
    global test_rews = [test_as as]

    for t = 11:2:(T-2)
        a = as[t] #action on THIS trial
        state = states[:, t+1] #NEXT state
        comm = common(a, state) #was the transition common?
        rew = rews[t+1] #NEXT reward (after the NEXT state)
        type = 2*(1-rew) + 1 - 1*comm #0 (rew common), 1 (rew uncommon), 2 (unrew common), 3 (unrew uncommon)
        stay = (a == as[t+2]) #was my next action from state 1 the same?
        global types = [types; type]
        global stays = [stays; stay]
    end

end

stayprobs = [mean(stays[types .== i]) for i = 0:3]
println(stayprobs)

f = figure()
bar([-2; -1; 1; 2], stayprobs, color = ["b"; "r"; "b"; "r"])
bar([0], [0], color = "r")
xlabel("trial type", fontsize = 12)
ylabel("stay probability", fontsize = 12)
ylim(0.5, 1)
legend(["common", "uncommon"])

ax = f.get_axes()[1]
ax.spines["top"].set_visible(false)
ax.spines["right"].set_visible(false)
xticks([-1.5, 1.5], ["rewarded", "unrewarded"])
savefig("figs/twostep_history_dependence.png", bbox_inches = "tight")
close()

max_stars_repo_name: KrisJensen/metaRL

max_issues_repo_name: KrisJensen/metaRL

max_forks_repo_name: KrisJensen/metaRL


text:
# Set random seed:
srand(1234)

#uniform.jl

"""
Create a NxM uniformly distributed matrix using the method rand()*(b-a)-a where a,b are the Unif Parameters.

Arguments
---------
a
    Lower Bound
b
    Upper Bound
N
    Row Dimension
M
    Column Dimension

"""
function uniform(a,b,N,M)
    #Returns a NxM uniformly distributed matrix
    #for random number in range [a,b] we note that it can be generated by: rand() * (b-a) - a
    #where rand() prints real numbers:
    rand(N,M)*(b-a)+a
end

#sigmoid.jl

"""
Return the sigmoid of x:

Arguments
---------
x
    input

"""

#Later on we will add a hyperparemeter to this function:

function sigmoid(x)
    return 1/(1+exp(-x))
    end

#sigmoid_deriv.jl

"""
Return the derivative of the sigmoid of x:

Arguments
---------
x
    input

"""

function sigmoid_deriv(x)
    return 1/(1+exp(-x))*(1-1/(1+exp(-x)))
end

#Train_Neural_Net.jl

"""
Train the XOR neural net for a given number of epochs:

Arguments
---------
epochs
    number of times the network will foward and back propagate
Layer_1
    weighting for hidden layer
Layer_2
    weighting for output layer
learning_rate
    Learning Rate for Gradient Descent
"""

function Train_Neural_Net_Loop(epochs,Layer_1,Layer_2,learning_rate)

    #Initialise XOR truth values:
    X=[ 0 0 ; 0 1 ; 1 0 ; 1 1] #When selected X[i,:] we have column vectors so need to transpose for matmult
    Y=[ 0;1;1;0]

    MSE=0

    for i=1:epochs #For the weight matrix wij is the weight for input i going to Neuron j:

        s_1=map(sigmoid,X*Layer_1) #Applies Sigmoid to X*Layer_1 the weighting matrix
        s_2=map(sigmoid,s_1*Layer_2) #Applies Sigmoid to a2 * second weighting matrix
        direct_error=s_2-Y
        MSE=0.5*sum(direct_error.*direct_error) #element wise squares direct_error then summs up and divides by length

    #Begin updating:
    delta_outer=-1.0*direct_error.*map(sigmoid_deriv,s_1*Layer_2) #this is the delta (which is the raw loss element-wise
                                                        #multiplied by z(3) see explanation in readme

    delta_inner=delta_outer*transpose(Layer_2).*map(sigmoid_deriv,X*Layer_1)
    # map(sigmoid_deriv,X*Layer_1) is f'(z2)

    #Update the weights:

    Layer_2 +=learning_rate*transpose(s_1)*delta_outer #Outer layer has been updated.
    Layer_1 += learning_rate*transpose(X)*delta_inner

    print(string(MSE,"\r"))

    end
    return Layer_1, Layer_2, MSE
    end

#Train_Neural_Net.jl

"""
Train the XOR neural net for a given mean square error it needs to reach:

Arguments
---------
MSE_Min
    Error the network trains till
Layer_1
    weighting for hidden layer
Layer_2
    weighting for output layer
learning_rate
    Learning Rate for Gradient Descent
"""

function Train_Neural_Net_MSE(MSE_Min,Layer_1,Layer_2,learning_rate)

    #Initialise XOR truth values:
    X=[ 0 0 ; 0 1 ; 1 0 ; 1 1] #When selected X[i,:] we have column vectors so need to transpose for matmult
    Y=[ 0;1;1;0]

    MSE=MSE_Min+1

    while MSE>=MSE_Min #For the weight matrix wij is the weight for input i going to Neuron j:

        s_1=map(sigmoid,X*Layer_1) #Applies Sigmoid to X*Layer_1 the weighting matrix
        s_2=map(sigmoid,s_1*Layer_2) #Applies Sigmoid to a2 * second weighting matrix
        direct_error=s_2-Y
        MSE=0.5*sum(direct_error.*direct_error) #element wise squares direct_error then summs up and divides by length

#This is not actually MSE it is the square error
    #Begin updating:
    delta_outer=-1.0*direct_error.*map(sigmoid_deriv,s_1*Layer_2) #this is the delta (which is the raw loss element-wise
                                                        #multiplied by z(3) see explanation in readme

    delta_inner=delta_outer*transpose(Layer_2).*map(sigmoid_deriv,X*Layer_1)
    # map(sigmoid_deriv,X*Layer_1) is f'(z2)

    #Update the weights:

    Layer_2 +=learning_rate*transpose(s_1)*delta_outer #Outer layer has been updated.
    Layer_1 += learning_rate*transpose(X)*delta_inner

    print(string(MSE,"\r"))

    end
    return Layer_1, Layer_2, MSE
    end

#XOR_Net.jl

"""
Output the values of a given XOR_Net

Arguments
---------
M
    Input for the XOR
w1
    Layer 1
w2
    Layer 2

"""

function XOR_Net(M,w1,w2)
    #Where M is a matrix that
    s_1=map(sigmoid,M*w1) #Applies Sigmoid to X*Layer_1 the weighting matrix
    s_2=map(sigmoid,s_1*w2)
    return s_2

end

#----------------------------- We can add tanh functions and hyperparameters-----------#

# #TEST RUN:---

# #Initialise Layers of Neuron Weights:

# Layer_1=uniform(0,1,2,2) #(In video above W1)
# Layer_2=uniform(0,1,2,1) #Column vector for the outer layer (in video above W2)

# # #Initialise Learning Rate:

# # learning_rate1=0.1
# # learning_rate2=0.01
# # # learning_rate3=0.4

# # w1, w2, MSE1 =Train_Neural_Net_Loop(100000,Layer_1,Layer_2,learning_rate1)
# x1, x2, MSE2 =Train_Neural_Net_Loop(1000000,Layer_1,Layer_2,learning_rate2)
# # z1, z2, MSE3=Train_Neural_Net_Loop(100000,Layer_1,Layer_2,learning_rate3)

# print(string(MSE1,""),string(MSE2,""),string(MSE3,""))

# # #================================================================================================#
# # #If kernal has hyper-parameter: set it here and uncomment the curry function named ker:
# # # function c_ker(hyp,ker)
# # #     return ((x,y) -> ker(x,y,hyp))
# # # end
# # # new_kernal = c_ker(0.1,hyper_exp_square_ker)


text:
using AutomotiveDrivingModels
using AutoRisk
using HDF5
using NGSIM

# extraction settings and constants
timestep_delta = 10 # timesteps between feature extractions
record_length = 10 # number of frames for record to track in the past
offset = 500 # from ends of the trajectories
maxframes = nothing # nothing for no max

output_filename = "ngsim.h5"
output_filepath = joinpath("../../data/trajectories/", output_filename)

models = Dict{Int, DriverModel}() # dummy, no behavior available
println("output filepath: $(output_filepath)")

# feature extractor (note the lack of behavioral features)
# use these feature extractors for lidar imputation:
subexts = [
        CoreFeatureExtractor(),
        TemporalFeatureExtractor(),
        WellBehavedFeatureExtractor(),
        CarLidarFeatureExtractor(16, carlidar_max_range = 50.)
    ]

ext = MultiFeatureExtractor(subexts)
n_features = length(ext)
features = Dict{Int, Dict{Int, Array{Float64}}}()

tic()
n_traj = 6
# extract
for traj_idx in 1:n_traj

    # setup
    trajdata = load_trajdata(traj_idx)
    roadway = get_corresponding_roadway(traj_idx)
    max_n_objects = maximum(n_objects_in_frame(trajdata, i) for i in 1 : nframes(trajdata))
    scene = Scene(max_n_objects)
    rec = SceneRecord(record_length, 0.1, max_n_objects)
    features[traj_idx] = Dict{Int, Array{Float64}}()
    ctr = 0

    for frame in offset : (nframes(trajdata) - offset)
        ctr += 1
        if maxframes != nothing && ctr >= maxframes
            break
        end

        print("\rtraj: $(traj_idx) / $(n_traj)\tframe $(frame) / $(nframes(trajdata) - offset)")

        # update the rec
        AutomotiveDrivingModels.update!(rec, get!(scene, trajdata, frame))

        # every timestep_delta step, extract features
        if frame % timestep_delta == 0

            for (vidx, veh) in enumerate(scene)
                # extract features
                veh_features = pull_features!(ext, rec, roadway, vidx, models)

                # add entry to features if vehicle not yet encountered
                if !in(veh.id, keys(features[traj_idx]))
                    features[traj_idx][veh.id] = zeros(n_features, 0)
                end

                # stack onto existing features
                features[traj_idx][veh.id] = cat(2, features[traj_idx][veh.id],
                    reshape(veh_features, (n_features, 1)))
            end
        end
    end
end
toc()

# compute max length across samples
maxlen = 0
for (traj_idx, feature_dict) in features
    for (veh_id, veh_features) in feature_dict
        maxlen = max(maxlen, size(veh_features, 2))
    end
end
println("max length across samples: $(maxlen)")

# write trajectory features
h5file = h5open(output_filepath, "w")
for (traj_idx, feature_dict) in features

    feature_array = zeros(n_features, maxlen, length(feature_dict))
    for (idx, (veh_id, veh_features)) in enumerate(feature_dict)
        feature_array[:, 1:size(veh_features, 2), idx] = reshape(veh_features, (n_features, size(veh_features, 2), 1))
    end
    h5file["$(traj_idx)"] = feature_array

end

# write feature names
attrs(h5file)["feature_names"] = feature_names(ext)
close(h5file)


text:
using Test

using ScoreAD: score_AD, score_AD_log, reject_nonfinite

using Distributions: Normal, logpdf
using Statistics: mean
import ForwardDiff: Dual, value, derivative

"""
    BernoulliData(N, s)

`N` draws from a Bernoulli(θ) distribution, `s` of which are `1`, `N-s` are `0`.
"""
struct BernoulliData
    "Total number of observations."
    N::Int
    "Observations =1."
    s::Int
end

(data::BernoulliData)(θ) = data.s/data.N*score_AD(θ)

@testset "bernoulli" begin
    N = 105
    s = 61
    data = BernoulliData(N, s)
    θ = 0.5
    for θ in range(0.1; stop = 0.9, length = 100)
        @test data(θ) ≈ s/N
        @test derivative(data, θ) ≈ s/N * 1/θ
    end
end

"""
    NormalData(x)

A vector of draws from a Normal(0, σ) distribution.
"""
struct NormalData{T}
    x::Vector{T}
end

function (data::NormalData)(σ)
    d = Normal(0.0, σ)
    mean(x * score_AD_log(logpdf(d, x)) for x in data.x)
end

@testset "normal" begin
    σ = 1.0
    x = rand(Normal(0, σ), 100)
    data = NormalData(x)
    for σ in range(0.1; stop = 5, length = 100)
        @test data(σ) ≈ mean(x)
        @test derivative(data, σ) ≈ mean(x * (x^2/(σ^3) - 1/σ) for x in x)
    end
end

@testset "reject nonfinite" begin
    @test (@inferred reject_nonfinite(-1.0)) ≡ -1.0
    @test (@inferred reject_nonfinite(-Inf)) ≡ -Inf
    @test (@inferred reject_nonfinite(Inf)) ≡ -Inf
    @test (@inferred reject_nonfinite(NaN)) ≡ -Inf

    @test (@inferred reject_nonfinite(Dual(-1.0, (2.0, 3.0)))) ==
        Dual(-1.0, (2.0, 3.0))

    function test_rejected_Dual(ℓ)
        r = @inferred reject_nonfinite(ℓ)
        @test value(r) ≡ oftype(value(ℓ), -Inf)
        @test typeof(r) ≡ typeof(ℓ)
    end

    test_rejected_Dual(Dual(-Inf, (2.0, 3.0)))
    test_rejected_Dual(Dual(Inf, (2.0, 3.0)))
    test_rejected_Dual(Dual(NaN, (2.0, 3.0)))
    test_rejected_Dual(Dual(1.0, (NaN, 3.0)))
    test_rejected_Dual(Dual(1.0, (2.0, -Inf)))
end


text:
using IntervalArithmetic
using Test

setformat(:standard, decorations=true, sigfigs=6)

setprecision(128)

@testset "Parse string to Interval" begin
    @test parse(Interval{Float64}, "1") == Interval(1, 1)
    @test parse(Interval{Float64}, "[1, 2]") == Interval(1, 2)
    @test parse(Interval{Float64}, "[-0x1.3p-1, 2/3]") == @interval(-0x1.3p-1, 2/3)

    @test parse(Interval{BigFloat}, "1") == @biginterval(1)
    @test parse(Interval{BigFloat}, "[-0x1.3p-1, 2/3]") == @biginterval(-0x1.3p-1, 2/3)

    @test parse(Interval{Float64}, "[Empty]") == emptyinterval(Float64)
    @test parse(Interval{BigFloat}, "[Empty]") == emptyinterval(BigFloat)

    @test parse(Interval{Float64}, "3 ±  4") == (-1)..7
    @test parse(Interval{Float64}, "0.2 ± 0.1") == 0.1..0.3
    @test parse(Interval{BigFloat}, "0.2 ± 0.1") == big"0.2" ± big"0.1"
end

@testset "Parse string to DecoratedInterval" begin
    @test parse(DecoratedInterval{Float64}, "[3]") == DecoratedInterval(3)

    @test parse(DecoratedInterval{Float64}, "[3, 4]") == DecoratedInterval(3, 4)

    @test parse(DecoratedInterval{Float64}, "[3, 4]_dac") == DecoratedInterval(3, 4, dac)
end


text:
struct TypeRef{T} end

Base.adjoint(x::Ref) = x[]
Base.adjoint(::Type{T}) where {T} = TypeRef{T}()
Base.getindex(::TypeRef{T}) where {T} = Ref{T}()
Base.getindex(::TypeRef{T}, x::S) where {T, S} = Ref{T}(x)

const SymString = Union{AbstractString, Symbol}

macro map(expr...)
    T, ex =
        if length(expr) == 1
            nothing, first(expr)
        else
            first(expr), last(expr)
    end

    args = map(ex.args) do arg
        if Meta.isexpr(arg, :(=))
            key, value = arg.args
            return :( $(QuoteNode(key)) => $value )
        elseif Meta.isexpr(arg, :(:=))
            key, value = arg.args
            return :( $key => $value )
        else
            error("Invalid syntax for @map, expected expression of type `foo = bar` or `foo := bar`")
        end
    end

    result =
        if isnothing(T)
            Expr(:call, :Dict, args...)
        else
            Expr(:call, Expr(:curly, :Dict, T.args...), args...)
        end

    return esc(result)
end

const GL_ENUM_TYPE = @map Symbol, GLenum {
    GLboolean = GL_BOOL,
    GLchar    = GL_BYTE,
    GLdouble  = GL_DOUBLE,
    GLfloat   = GL_FLOAT,
    GLint     = GL_INT,
    GLshort   = GL_SHORT,
    GLubyte   = GL_UNSIGNED_BYTE,
    GLuint    = GL_UNSIGNED_INT,
    GLushort  = GL_UNSIGNED_SHORT,
}

macro genumof(type)
    :( GL_ENUM_TYPE[ $(QuoteNode(type)) ] )
end

genumof(::Type{Cint})    = GL_INT
genumof(::Type{Cchar})   = GL_BYTE
genumof(::Type{Cfloat})  = GL_FLOAT
genumof(::Type{Cshort})  = GL_SHORT
genumof(::Type{Int64})   = GL_DOUBLE
genumof(::Type{Cuint})   = GL_UNSIGNED_INT
genumof(::Type{Cuchar})  = GL_UNSIGNED_BYTE
genumof(::Type{Cushort}) = GL_UNSIGNED_SHORT

gbool(x::GLenum) = x == GL_TRUE
gbool(x::Bool) = x ? GL_TRUE : GL_FALSE

"""
        @g_str(name) -> GLenum

Returns a OpenGL constant.

# Examples

```julia
g"bool" == GL_BOOL
g"Texture 2D" == GL_TEXTURE_2D
g"Clamp.to.edge" == GL_CLAMP_TO_EDGE
g"texture-wrap-t" == GL_TEXTURE_WRAP_T
```
"""
macro g_str(ex)
    name = uppercase(replace(ex, r"(\s|\.|\-)+" => "_"))
    esc( Symbol("GL_$name") )
end


text:
module Error

export LispError

type LispError <: Exception
    msg::AbstractString
end

end #module

max_stars_repo_name: scottcrossen/JuliaDevelopment

max_issues_repo_name: scottcrossen/JuliaDevelopment

max_forks_repo_name: scottcrossen/JuliaDevelopment


text:
# read_generated_quantities

"""

Read generated_quantities output files created by StanSample.jl.

$(SIGNATURES)

### Required arguments
```julia
* `model`                    : SampleModel
```
"""
function read_generated_quantities(model::SampleModel, chains=[1];
    output_format=:namedtuple,
    kwargs...
  )

read_csv_files(model::SampleModel, output_format::Symbol;
  include_internals=false,
  chains=1:model.n_chains[1],
  start=1,
  kwargs...)

  read_csv_files(model::SampleModel, :namedtuple;
  include_internals=false,
  chains=[1],
  start=1,
  kwargs...)

end   # end of read_generated_quantities


text:
using ZXCalculus, Graphs

function gen_cir()
    cir = ZXDiagram(5)
    push_gate!(cir, Val{:X}(), 5, 1//1)
    push_gate!(cir, Val{:H}(), 5)
    push_gate!(cir, Val{:Z}(), 5)
    push_gate!(cir, Val{:CNOT}(), 5, 4)
    push_gate!(cir, Val{:Z}(), 5, 7//4)
    push_gate!(cir, Val{:CNOT}(), 5, 1)
    push_gate!(cir, Val{:Z}(), 5, 1//4)
    push_gate!(cir, Val{:CNOT}(), 5, 4)
    push_gate!(cir, Val{:Z}(), 4, 1//4)
    push_gate!(cir, Val{:Z}(), 5, 7//4)
    push_gate!(cir, Val{:CNOT}(), 5, 1)
    push_gate!(cir, Val{:CNOT}(), 4, 1)
    push_gate!(cir, Val{:Z}(), 5, 1//4)
    push_gate!(cir, Val{:Z}(), 1, 1//4)
    push_gate!(cir, Val{:Z}(), 4, 7//4)
    push_gate!(cir, Val{:CNOT}(), 4, 1)
    push_gate!(cir, Val{:CNOT}(), 5, 4)
    push_gate!(cir, Val{:Z}(), 5, 7//4)
    push_gate!(cir, Val{:CNOT}(), 5, 3)
    push_gate!(cir, Val{:Z}(), 5, 1//4)
    push_gate!(cir, Val{:CNOT}(), 5, 4)
    push_gate!(cir, Val{:Z}(), 4, 1//4)
    push_gate!(cir, Val{:Z}(), 5, 7//4)
    push_gate!(cir, Val{:CNOT}(), 5, 3)
    push_gate!(cir, Val{:CNOT}(), 4, 3)
    push_gate!(cir, Val{:Z}(), 5, 1//4)
    push_gate!(cir, Val{:Z}(), 3, 1//4)
    push_gate!(cir, Val{:Z}(), 4, 7//4)
    push_gate!(cir, Val{:H}(), 5)
    push_gate!(cir, Val{:Z}(), 5)
    push_gate!(cir, Val{:CNOT}(), 4, 3)
    push_gate!(cir, Val{:CNOT}(), 5, 4)
    push_gate!(cir, Val{:H}(), 5)
    push_gate!(cir, Val{:Z}(), 5)
    push_gate!(cir, Val{:CNOT}(), 5, 3)
    push_gate!(cir, Val{:Z}(), 5, 7//4)
    push_gate!(cir, Val{:CNOT}(), 5, 2)
    push_gate!(cir, Val{:Z}(), 5, 1//4)
    push_gate!(cir, Val{:CNOT}(), 5, 3)
    push_gate!(cir, Val{:Z}(), 3, 1//4)
    push_gate!(cir, Val{:Z}(), 5, 7//4)
    push_gate!(cir, Val{:CNOT}(), 5, 2)
    push_gate!(cir, Val{:CNOT}(), 3, 2)
    push_gate!(cir, Val{:Z}(), 5, 1//4)
    push_gate!(cir, Val{:H}(), 5)
    push_gate!(cir, Val{:Z}(), 2, 1//4)
    push_gate!(cir, Val{:Z}(), 3, 7//4)
    push_gate!(cir, Val{:Z}(), 5)
    push_gate!(cir, Val{:CNOT}(), 3, 2)
    push_gate!(cir, Val{:CNOT}(), 5, 3)
    push_gate!(cir, Val{:H}(), 5)
    push_gate!(cir, Val{:Z}(), 5)
    push_gate!(cir, Val{:CNOT}(), 5, 2)
    push_gate!(cir, Val{:Z}(), 5, 7//4)
    push_gate!(cir, Val{:CNOT}(), 5, 1)
    push_gate!(cir, Val{:Z}(), 5, 1//4)
    push_gate!(cir, Val{:CNOT}(), 5, 2)
    push_gate!(cir, Val{:Z}(), 2, 1//4)
    push_gate!(cir, Val{:Z}(), 5, 7//4)
    push_gate!(cir, Val{:CNOT}(), 5, 1)
    push_gate!(cir, Val{:CNOT}(), 2, 1)
    push_gate!(cir, Val{:Z}(), 5, 1//4)
    push_gate!(cir, Val{:Z}(), 1, 1//4)
    push_gate!(cir, Val{:Z}(), 2, 7//4)
    push_gate!(cir, Val{:H}(), 5)
    push_gate!(cir, Val{:Z}(), 5)
    push_gate!(cir, Val{:CNOT}(), 2, 1)
    push_gate!(cir, Val{:CNOT}(), 5, 2)
    push_gate!(cir, Val{:CNOT}(), 5, 1)
    return cir
end

cir = gen_cir()
@test tcount(cir) == 28
cir2 = phase_teleportation(cir)
@test tcount(cir2) == 8


text:
module test_skip_testsets

module TestTools

export @testset

import Test: @testset
using Test: Test, AbstractTestSet, DefaultTestSet, Error, Random, testset_forloop, get_testset_depth, get_testset, push_testset, pop_testset, finish, record, _check_testset
if VERSION >= v"1.3.0-DEV.565"
    default_rng = Random.default_rng
else
    GLOBAL_RNG = Random.GLOBAL_RNG
end

skip_testsets = []

# code from https://github.com/JuliaLang/julia/blob/master/stdlib/Test/src/Test.jl#L1065

macro testset(args...)
    isempty(args) && error("No arguments to @testset")

    tests = args[end]

    # Determine if a single block or for-loop style
    if !isa(tests,Expr) || (tests.head !== :for && tests.head !== :block)
        error("Expected begin/end block or for loop as argument to @testset")
    end

    if tests.head === :for
        ex = testset_forloop(args, tests, __source__)
    else
        if VERSION >= v"1.8.0-DEV.809"
            testset_beginend_call = Test.testset_beginend_call
        else
            testset_beginend_call = Test.testset_beginend
        end
        ex = testset_beginend_call(args, tests, __source__)
    end
    desc = first(args)
    !(desc in skip_testsets) && return ex
end

end # module TestTools

using Test
using .TestTools
push!(TestTools.skip_testsets, "testset2")

@testset "testset1" begin
@test true
end

@testset "testset2" begin
@test true
end

@testset "testset3" begin
@test true
end

end # module test_skip_testsets


text:
module PartialOrder
export subDAGof, ordered_subDAGof, copyDAG, raw_filter!

type Node{T}
    value::T
    gt::Set{Node{T}}

    Node(value) = new(value, Set{Node{T}}())
    Node(value, gt) = new(value, gt)
end
Node{T}(value::T) = Node{T}(value)

copyDAG{T}(top::Node{T}) = copyDAG(Dict{Node{T},Node{T}}(), top)
function copyDAG{T}(subs::Dict{Node{T},Node{T}}, node::Node{T})
    if haskey(subs, node); return subs[node] end

    gt = Set{Node{T}}([copyDAG(subs, child) for child in node.gt])
    subs[node] = Node{T}(node.value, gt)
end

raw_filter!{T}(n::Node{T}, keep::Set{T}) = raw_filter!(Set{Node{T}}(), n, keep)
function raw_filter!{T}(seen::Set{Node{T}}, node::Node{T}, keep::Set{T})
    if node in seen; return end
    push!(seen, node)
    node.gt = Set{Node{T}}(filter(node->(node.value in keep), node.gt))
    for child in node.gt;  raw_filter!(child, keep)  end
end

subDAGof{T}(node::Node{T}) = (sub = Set{Node{T}}(); addsubDAG!(sub, node); sub)
function addsubDAG!{T}(seen::Set{Node{T}}, node::Node{T})
    if node in seen; return; end
    push!(seen, node)
    for below in node.gt; addsubDAG!(seen, below); end
end

function ordered_subDAGof{T}(node::Node{T})
    seen, order = Set{Node{T}}(), Node{T}[]
    addsubDAG!(seen, order, node)
    order
end
function addsubDAG!{T}(seen::Set{Node{T}},order::Vector{Node{T}},node::Node{T})
    if node in seen; return; end
    push!(seen, node); push!(order, node)
    for below in node.gt; addsubDAG!(seen, order, below); end
end

insert!{T}(at::Node{T}, node::Node{T}) = insert!(Dict{Node{T},Bool}(), at, node)
function insert!{T}(seen::Dict{Node{T},Bool}, at::Node{T}, node::Node{T})
    if haskey(seen, at); return seen[at] end
    if node.value >= at.value
        if at.value >= node.value
            at.value = node.value  # at == node
            return seen[at] = true
        end
        # node > at
        push!(node.gt, at)
        at_above_node = false
    else
        at_above_node = any(Bool[insert!(seen, below, node) for below in at.gt])
    end
    if !at_above_node
        if at.value >= node.value
            setdiff!(at.gt, intersect(at.gt, node.gt))
            push!(at.gt, node)
            at_above_node = true
        end
    end
    seen[at] = at_above_node
end

end # module


text:
"""
Surprise-REINFORCE
------------------

REINFORCE with surprise-modulated learning rate

See REINFORCE.jl for some more implementational info
See learner/REINFORCEmodulatedlearner.jl for the core implementation.

Note 1
------
With particle filtering we need to run multiple times with different seeds.
As final LL we report the MEAN across seeds.
For signals/results we keep track of the ones of all seeds
and in the end report the one with the MAXIMUM LL.

Note 2
------
get_all_puzzles(), get_all_episodes_in_puzzle(), get_data_by_puzzle_episode()
are in rl_data_toolsforalgos.jl
"""
function getLL_REINFORCE_modulated(data::Array{Int,2},
                                    alpha::Float64,
                                    surpriseprobability::Float64,
                                    stochasticity::Float64,
                                    gamma::Float64,
                                    temp::Float64,
                                    nparticles::Union{Int, Float64};
                                    particlefilterseeds = [1,2,3],
                                    do_debug_print = false)
    nparticles = converttoInt(nparticles)
    data_stats, nr_states, nr_actions, nr_samples = get_data_stats_and_dimensions(data)
    if data_stats.nr_subjects != 1
        error("Invalid argument: getLL_REINFORCE_modulated supports only one subject")
    end
    if nr_actions!=2; error("REINFORCE supports two actions only"); end
    all_puzzles = get_all_puzzles(data)

    action_selection_probs_seeds = Array{Array{Float64, 1}, 1}(undef, length(particlefilterseeds))
    [action_selection_probs_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    LL_seeds = zeros(Float64, length(particlefilterseeds))

    theta_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    Tsas_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    #signals:
    Surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [Surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    γ_surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [γ_surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    SPE_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [SPE_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    PolicyParametersDiffChosenVsUnchosen_record_seed =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyParametersDiffChosenVsUnchosen_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    PolicyDiffChosenVsUnchosen_record_seed = Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyDiffChosenVsUnchosen_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    finaltheta_record = 0
    for i_seed in 1:length(particlefilterseeds)
        action_counter = 0
        Testimate = []
        learner = []
        for i_puzzle in all_puzzles
            # Puzzle = independent block of the experiment.
            # The code was made for general purposes, in this experiment we
            # have only ONE puzzle.
            Testimate = TParticleFilterJump(nr_states = nr_states,
                                    nr_actions = nr_actions,
                                    nparticles = nparticles,
                                    surpriseprobability = surpriseprobability,
                                    stochasticity = stochasticity,
                                    seed = particlefilterseeds[i_seed])

            learner = REINFORCEModulated(nr_states = nr_states,
                            nr_actions = nr_actions,
                            α = alpha,
                            γ = gamma,
                            do_debug_print = do_debug_print)

            all_epi_in_puzzle = get_all_episodes_in_puzzle(data, i_puzzle)
            for i_episode in all_epi_in_puzzle
                episode_data = get_data_by_puzzle_episode(data, i_puzzle, i_episode)
                path_length = size(episode_data)[1]  # = T
                if do_debug_print
                    println("i_puzzle:$i_puzzle | i_episode:$i_episode | path_length:$path_length")
                end
                reward = 0. # we have a single reward at goal state. If the episode was
                # not completed, this code still correctly works (updates are 0)
                for t in 1:path_length
                    action_counter += 1
                    state = episode_data[t, Int(c_state)]
                    action = episode_data[t, Int(c_action)]
                    reward = episode_data[t, Int(c_reward)]
                    next_state = episode_data[t, Int(c_nextState)]
                    done = reward>0 ? true : false

                    p_actions = getSoftmaxProbs([learner.θ[state,1],learner.θ[state,2]], temp, [0., 0.])

                    LL_seeds[i_seed] += log(p_actions[action])
                    action_selection_probs_seeds[i_seed][action_counter] = copy(p_actions[action])

                    # Particle Filtering model learner
                    updatet!(Testimate, state, action, next_state, done)
                    γ_surprise_seeds[i_seed][action_counter] = copy(Testimate.γ_surprise[1])
                    Surprise_seeds[i_seed][action_counter] = copy(Testimate.Sgm[1])
                    Tsas_seeds[i_seed]= deepcopy(Testimate.Ps1a0s0)
                    SPE_seeds[i_seed][action_counter] = copy(Testimate.SPE[1])

                    PolicyParametersDiffChosenVsUnchosen_record_seed[i_seed][action_counter] = learner.θ[state, action] - learner.θ[state, mod(action,2)+1]
                    PolicyDiffChosenVsUnchosen_record_seed[i_seed][action_counter] = p_actions[action] - p_actions[mod(action,2)+1]

                    # REINFORCE modulated
                    update!(learner, state, action, reward, p_actions[action], (1. - Testimate.γ_surprise[1]))

                    if do_debug_print
                        println("t:$t | s:$state | a:$action | R:$reward | s':$next_state | p(s,a):$(round(p_action,digits=2))")
                    end
                end # path_length (in single episode)
          end # episode
      end # puzzle
      theta_seeds[i_seed] = copy(learner.θ)
  end # seed

  LL = get_meanLL_acrossseeds(LL_seeds)

  action_selection_probs_cat = hcat(action_selection_probs_seeds...)
  action_selection_probs = mean(action_selection_probs_cat, dims=2)[:]
  trial_by_trial_p_sa = (action_selection_probs)

  # For signals: return the ones corresponding to the Max LL across seeds
  LL_max, seed_idx = findmax(LL_seeds)
  trial_by_trial_p_sa_max = (action_selection_probs_seeds[seed_idx])

  SPE_record = copy(SPE_seeds[seed_idx])

  algo = "REINFORCE_modulated"
  param_v = [alpha, surpriseprobability, stochasticity, gamma, temp, nparticles]
  param_str = "alpha=$(round(alpha,digits=4)),
              surpriseprobability=$(round(surpriseprobability,digits=4)),
              stochasticity=$(round(stochasticity,digits=4)),
              gamma=$(round(gamma,digits=4)),
              temp=$(round(temp,digits=4)),
              nparticles=$nparticles"

  results_dict = Dict{String,Any}("final_Theta" => theta_seeds[seed_idx],
                                "final_Tsas" => Tsas_seeds[seed_idx],
                                "Surprise" => Surprise_seeds[seed_idx],
                                "gamma_surprise" => γ_surprise_seeds[seed_idx],
                                "nparticles" => nparticles,
                                "particlefilterseeds" => particlefilterseeds,
                                "trial_by_trial_p_sa_max" => trial_by_trial_p_sa_max,
                                "PolicyParametersDiffChosenVsUnchosen" => PolicyParametersDiffChosenVsUnchosen_record_seed[seed_idx],
                                "PolicyDiffChosenVsUnchosen" => PolicyDiffChosenVsUnchosen_record_seed[seed_idx],)

  signals = Step_by_step_signals(
    data_stats.all_subject_ids[1],
    [], # no RPE
    SPE_record, # no SPE
    trial_by_trial_p_sa,
    LL,
    data_stats,
    algo,
    param_v,
    param_str,
    99, #biased_action_idx, dummy value
    results_dict)
  return LL, signals

end


text:
struct ZIPoissonExpert{T<:Real} <: ZIDiscreteExpert
    p::T
    λ::T
    ZIPoissonExpert{T}(p::T, λ::T) where {T<:Real} = new{T}(p, λ)
end

function ZIPoissonExpert(p::T, λ::T; check_args=true) where {T <: Real}
    check_args && @check_args(ZIPoissonExpert, 0 <= p <= 1 && λ >= zero(λ))
    return ZIPoissonExpert{T}(p, λ)
end

## Outer constructors
ZIPoissonExpert(p::Real, λ::Real) = ZIPoissonExpert(promote(p, λ)...)
ZIPoissonExpert(p::Integer, λ::Integer) = ZIPoissonExpert(float(p), float(λ))
ZIPoissonExpert() = ZIPoissonExpert(0.5, 1.0)

## Conversion
function convert(::Type{ZIPoissonExpert{T}}, p::S, λ::S) where {T <: Real, S <: Real}
    ZIPoissonExpert(T(p), T(λ))
end
function convert(::Type{ZIPoissonExpert{T}}, d::ZIPoissonExpert{S}) where {T <: Real, S <: Real}
    ZIPoissonExpert(T(d.p), T(d.λ), check_args=false)
end
copy(d::ZIPoissonExpert) = ZIPoissonExpert(d.p, d.λ, check_args=false)

## Loglikelihood of Expoert
logpdf(d::ZIPoissonExpert, x...) = isinf(x...) ? 0.0 : Distributions.logpdf.(Distributions.Poisson(d.λ), x...)
pdf(d::ZIPoissonExpert, x...) = isinf(x...) ? -Inf : Distributions.pdf.(Distributions.Poisson(d.λ), x...)
logcdf(d::ZIPoissonExpert, x...) = isinf(x...) ? 0.0 : Distributions.logcdf.(Distributions.Poisson(d.λ), x...)
cdf(d::ZIPoissonExpert, x...) = isinf(x...) ? 1.0 : Distributions.cdf.(Distributions.Poisson(d.λ), x...)

## expert_ll, etc
expert_ll_exact(d::ZIPoissonExpert, x::Real) = (x == 0.) ? log.(p_zero(d) + (1-p_zero(d))*exp.(LRMoE.logpdf(LRMoE.PoissonExpert(d.λ), x))) : log.(1-p_zero(d)) + LRMoE.logpdf(LRMoE.PoissonExpert(d.λ), x)
function expert_ll(d::ZIPoissonExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_ll_pos = LRMoE.expert_ll(LRMoE.PoissonExpert(d.λ), tl, yl, yu, tu)
    # Deal with zero inflation
    p0 = p_zero(d)
    expert_ll = (yl == 0.) ? log.(p0 + (1-p0)*exp.(expert_ll_pos)) : log.(0.0 + (1-p0)*exp.(expert_ll_pos))
    return expert_ll
end
function expert_tn(d::ZIPoissonExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_tn_pos = LRMoE.expert_tn(LRMoE.PoissonExpert(d.λ), tl, yl, yu, tu)
    # Deal with zero inflation
    p0 = p_zero(d)
    expert_tn = (tl == 0.) ? log.(p0 + (1-p0)*exp.(expert_tn_pos)) : log.(0.0 + (1-p0)*exp.(expert_tn_pos))
    return expert_tn
end
function expert_tn_bar(d::ZIPoissonExpert, tl::Real, yl::Real, yu::Real, tu::Real)
    expert_tn_bar_pos = LRMoE.expert_tn_bar(LRMoE.PoissonExpert(d.λ), tl, yl, yu, tu)
    # Deal with zero inflation
    p0 = p_zero(d)
    expert_tn_bar = (tl > 0.) ? log.(p0 + (1-p0)*exp.(expert_tn_bar_pos)) : log.(0.0 + (1-p0)*exp.(expert_tn_bar_pos))
    return expert_tn_bar
end

exposurize_expert(d::ZIPoissonExpert; exposure = 1) = ZIPoissonExpert(d.p, d.λ*exposure)

## Parameters
params(d::ZIPoissonExpert) = (d.p, d.λ)
p_zero(d::ZIPoissonExpert) = d.p
function params_init(y, d::ZIPoissonExpert)
    μ, σ2 = mean(y), var(y)
    λp = (σ2-μ) / μ
    tmp = λp / μ
    p_init = tmp / (1+tmp)
    λ_init = λp / p_init
    try
        ZIPoissonExpert(p_init, λ_init)
    catch;
        ZIPoissonExpert()
    end
end

## Simululation
sim_expert(d::ZIPoissonExpert) = (1 .- Distributions.rand(Distributions.Bernoulli(d.p), 1)[1]) .* Distributions.rand(Distributions.Poisson(d.λ), 1)[1]

## penalty
penalty_init(d::ZIPoissonExpert) = [2.0 1.0]
no_penalty_init(d::ZIPoissonExpert) = [1.0 Inf]
penalize(d::ZIPoissonExpert, p) = (p[1]-1)*log(d.λ) - d.λ/p[2]

## statistics
mean(d::ZIPoissonExpert) = (1-d.p)*mean(Distributions.Poisson(d.λ))
var(d::ZIPoissonExpert) = (1-d.p)*var(Distributions.Poisson(d.λ)) + d.p*(1-d.p)*(mean(Distributions.Poisson(d.λ)))^2
quantile(d::ZIPoissonExpert, p) = p <= d.p ? 0.0 : quantile(Distributions.Poisson(d.λ), p-d.p)

## EM: M-Step
function EM_M_expert(d::ZIPoissonExpert,
                     tl, yl, yu, tu,
                     exposure,
                     z_e_obs, z_e_lat, k_e;
                     penalty = true, pen_pararms_jk = [2.0 1.0])

    # Old parameters
    λ_old = d.λ
    p_old = p_zero(d)

    if p_old > 0.999999
        return d
    end

    expert_ll_pos = fill(0.0, length(yl))
    expert_tn_bar_pos = fill(0.0, length(yl))

    tmp_exp = PoissonExpert(d.λ)
    for i in 1:length(yl)
        d_expo = exposurize_expert(tmp_exp, exposure = exposure[i])
        expert_ll_pos[i] = expert_ll.(d_expo, tl[i], yl[i], yu[i], tu[i])
        expert_tn_bar_pos[i] = expert_tn_bar.(d_expo, tl[i], yl[i], yu[i], tu[i])
    end

    # Update zero probability
    z_zero_e_obs = z_e_obs .* EM_E_z_zero_obs(yl, p_old, expert_ll_pos)
    z_pos_e_obs = z_e_obs .- z_zero_e_obs
    z_zero_e_lat = z_e_lat .* EM_E_z_zero_lat(tl, p_old, expert_tn_bar_pos)
    z_pos_e_lat = z_e_lat .- z_zero_e_lat
    p_new = EM_M_zero(z_zero_e_obs, z_pos_e_obs, z_zero_e_lat, z_pos_e_lat, k_e)

    # Update parameters: call its positive part
    tmp_update = EM_M_expert(tmp_exp,
                            tl, yl, yu, tu,
                            exposure,
                            z_pos_e_obs, z_pos_e_lat, k_e,
                            penalty = penalty, pen_pararms_jk = pen_pararms_jk)

    return ZIPoissonExpert(p_new, tmp_update.λ)
end

## EM: M-Step, exact observations
function EM_M_expert_exact(d::ZIPoissonExpert,
                     ye, exposure,
                     # expert_ll_pos,
                     z_e_obs;
                     penalty = true, pen_pararms_jk = [Inf 1.0 Inf])

    # Old parameters
    λ_old = d.λ
    p_old = p_zero(d)

    if p_old > 0.999999
        return d
    end

    # Update zero probability
    expert_ll_pos = fill(NaN, length(exposure))
    for i in 1:length(exposure)
        expert_ll_pos[i] = expert_ll_exact(exposurize_expert(LRMoE.PoissonExpert(λ_old), exposure = exposure[i]), ye[i])
    end

    z_zero_e_obs = z_e_obs .* EM_E_z_zero_obs(ye, p_old, expert_ll_pos)
    z_pos_e_obs = z_e_obs .- z_zero_e_obs
    p_new = EM_M_zero(z_zero_e_obs, z_pos_e_obs, 0.0, 0.0, 0.0)
        # EM_M_zero(z_zero_e_obs, z_pos_e_obs, z_zero_e_lat, z_pos_e_lat, k_e)

    # Update parameters: call its positive part
    tmp_exp = PoissonExpert(d.λ)
    tmp_update = EM_M_expert_exact(tmp_exp,
                            ye, exposure,
                            # expert_ll_pos,
                            z_pos_e_obs;
                            penalty = penalty, pen_pararms_jk = pen_pararms_jk)

    return ZIPoissonExpert(p_new, tmp_update.λ)
end

max_stars_repo_name: sparktseung/JuliaLRMoE

max_issues_repo_name: sparktseung/JuliaLRMoE

max_forks_repo_name: sparktseung/JuliaLRMoE


text:
using GPUArrays, Test
using GPUArrays.TestSuite

using Pkg

@testset "JLArray" begin
    GPUArrays.test(JLArray)
end

function test_package(package, branch=nothing)
    mktempdir() do devdir
        withenv("JULIA_PKG_DEVDIR" => devdir) do
            # try to install from the same branch of GPUArrays
            try
                if branch === nothing
                    branch = chomp(read(`git -C $(@__DIR__) rev-parse --abbrev-ref HEAD`, String))
                    branch == "HEAD" && error("in detached HEAD state")
                end
                Pkg.add(PackageSpec(name=package, rev=String(branch)))
                @info "Installed $package from $branch branch"
            catch ex
                @warn "Could not install $package from same branch as GPUArrays, trying master branch" exception=ex
                Pkg.add(PackageSpec(name=package, rev="master"))
            end

            Pkg.test(package)
        end
    end
end

if haskey(ENV, "GITLAB_CI")
    branch = ENV["CI_COMMIT_REF_NAME"]
    @testset "CuArray" begin
        test_package("CuArrays", branch)
    end
end


text:
immutable modwt <: WaveletTransform

  original::Array{Float64}
  L::Int
  series::Int
  level::Int
  filter::waveletFilter
  boundary::ASCIIString
  W::Array{Float64, 3}
  V::Array{Float64, 3}
  aligned::Bool

  function modwt(X::Array{Float64}, filter::ASCIIString, nLevels::Int, boundary::ASCIIString)
    filter = eval(Expr(:call, symbol(string(filter,"Filter")), 1, true))

    if isa(X, Vector)
      (N,) = size(X)
      nSeries = 1
    else
      (N, nSeries) = size(X)
    end

    (W, V) = modwtDo(X, filter, nLevels, boundary, nSeries, N)

    new(X, N, nSeries, nLevels, filter, boundary, W, V, false)

  end

end

function modwtForward(V::Vector{Float64}, filter::waveletFilter, j::Int)
  (N,) = size(V)
  Wj = fill(NaN, N)
  Vj = fill(NaN, N)

  for t = 0:(N-1)
    k = t
    Wjt = filter.h[1]*V[k+1]
    Vjt = filter.g[1]*V[k+1]
    for n = 1:(filter.L-1)
      k = k - 2^(j-1)
      if (k < 0)
        k = k + ceil(-k/N)*N
      end
      Wjt = Wjt + filter.h[n+1]*V[k+1]
      Vjt = Vjt + filter.g[n+1]*V[k+1]
    end
    Wj[t+1] = Wjt
    Vj[t+1] = Vjt
  end

  return (Wj, Vj)
end

function modwtBackward(W::Vector{Float64}, V::Vector{Float64}, filter::waveletFilter, j::Int)
  (N,) = size(V)
  (L,) = size(filter.h)
  Vj = fill(NaN, N)

  for t = 0:(N-1)
    k = t
    Vjt = filter.h[1]*W[k+1] + filter.g[1]*V[k+1]
    for n = 1:(L-1)
      k = k + 2^(j-1)
      if k >= N
        k = k - floor(k/N)*N
      end
      Vjt = Vjt + filter.h[n+1]*W[k+1] + filter.g[n+1]*V[k+1]
    end
    Vj[t+1] = Vjt
  end

  return Vj
end

function modwtDo(X::Array{Float64}, filter::waveletFilter, nLevels::Int, boundary::ASCIIString, nSeries::Int, N::Int)
  # reflect X for reflection method
  if (boundary == "reflection")
    X = extendSeries(X, boundary, "double")
    N = 2*N
  end

  # initialize variables for pyramid algorithm
  nBoundary = fill(0.0, nLevels)
  WCoefs = fill(0.0, N, nLevels, nSeries)
  VCoefs = fill(0.0, N, nLevels, nSeries)

  # implement the pyramid algorithm
  for i=1:nSeries
    Vj = X[:,i]
    for j=1:nLevels
      (WCoefs[:,j,i], VCoefs[:,j,i]) = modwtForward(Vj,filter,j)
      Vj = convert(Vector, VCoefs[:,j,i])
      Lj = (2^j-1)*(filter.L-1)+1
      nBoundary[j] = min(Lj,N)
    end
  end

  return (WCoefs, VCoefs)
end

function imodwtDo(wt::modwt)
  output = fill(0.0, wt.L, wt.series)

  for i=1:wt.series
    Vj = convert(Vector, wt.V[:,end,i])
    for j=wt.level:-1:1
      Vj = modwtBackward(convert(Vector, wt.W[:,j,i]), Vj, wt.filter, j)
      Vj = convert(Vector, Vj)
    end
    output[:,i] = Vj
  end

  return output

end


text:
if !(@isdefined BUNDLES_PATH)
    include("./bundlepaths.jl")
end

const CONFIG_BEGIN_MARKER = "###JULIA-WEBIO-CONFIG-BEGIN"
const CONFIG_END_MARKER = "###JULIA-WEBIO-CONFIG-END"

_tryrun(args...; kwargs...) = try run(args...; kwargs...) catch nothing end

"""
    install_notebook_config()

Install necessary configuration for the `jlstaticserve` notebook (server)
extension. This function only configures the notebook extension, not the browser
nbextension or Jupyter Lab extension.
* Adds the path to `./deps` to Python's `sys.path` so that we can load the
    `jlstaticserve.py` extension. This is done in `jupyter_notebook_config.py`
    because there's no way to add to `sys.path` from the JSON config file.
* Adds `jlstaticserve` to the list of extensions loaded in the notebook server.
    This is done in `jupyter_notebook_config.json` because that file has higher
    precedence when both the `.py` and `.json` files exist (IPyWidgets, for
    example, writes to the JSON file, so if we only wrote to the `.py` file,
    that directive would take precedence and the `jlstaticserve` extension would
    **not** be loaded).
"""
function install_jupyter_serverextension()
    config_dir = jupyter_config_dir()
    mkpath(config_dir)
    config_file_py = joinpath(config_dir, "jupyter_notebook_config.py")
    config_py = isfile(config_file_py) ? read(config_file_py, String) : ""

    # Remove previous config
    config_regex = Regex(
        "\n?" * CONFIG_BEGIN_MARKER * ".*" * CONFIG_END_MARKER * "\n?",
        "s",
    )
    config_py = replace(config_py, config_regex => "")

    # We do a repr to make sure that all the necessary characters are escaped
    # and the whole path is wrapped in quotes.
    deps_dir = dirname(@__FILE__)
    deps_dir_esc = repr(deps_dir)
    error_msg = strip("""
        Directory $deps_dir could not be found; WebIO will not work as expected.
        """)
    config_py *= """

    $CONFIG_BEGIN_MARKER
    # Add the path to WebIO/deps so that we can load the `jlstaticserve` extension.
    import os, sys, warnings
    webio_deps_dir = $deps_dir_esc
    if os.path.isfile(os.path.join(webio_deps_dir, "jlstaticserve.py")):
        sys.path.append(webio_deps_dir)
    else:
        warning_msg = (
            'Directory %s could not be found; WebIO.jl will not work as expected. '
            + 'Make sure WebIO.jl is installed correctly (try running '
            + 'Pkg.add("WebIO") and Pkg.build("WebIO") from the Julia console to '
            + 'make sure it is).'
        ) % webio_deps_dir
        warnings.warn(warning_msg)
    $CONFIG_END_MARKER
    """

    config_file_json = joinpath(config_dir, "jupyter_notebook_config.json")
    config_json = isfile(config_file_json) ? read(config_file_json, String) : "{}"
    config_json = try
        JSON.parse(config_json)
    catch exc
        @error "Unable to parse Jupyter notebook config file ($config_file_json). Please fix it and rebuild WebIO." exception=exc
        rethrow()
    end

    # Magic to safely access nested JSON objects and set them to empty objects
    # if they don't exist yet.
    app_config = get!(config_json, "NotebookApp", Dict())
    extensions_config = get!(app_config, "nbserver_extensions", Dict())
    extensions_config["jlstaticserve"] = true

    # Defer writing until both files until the end to avoid inconsistent state.
    write(config_file_py, config_py)
    write(config_file_json, json(config_json, 4))
    return nothing
end

function find_condajl_jupyter_cmd(;
    check_nodejs::Bool=false,
)::Cmd
    conda_root = joinpath(first(Base.DEPOT_PATH), "conda", "3")
    conda_bin_dir = joinpath(
        conda_root,
        @static(Sys.iswindows() ? "Scripts" : "bin"),
    )

    jupyter = joinpath(conda_bin_dir, exe("jupyter"))
    if !isfile(jupyter)
        error(
            "Could not find the Conda.jl `jupyter` executable " *
            "(is it installed?)."
        )
    end

    # We need to add Conda's bin directory to the PATH so that jupyterlab can
    # detect nodejs.
    conda_env = copy(ENV)
    conda_env["PATH"] = string(conda_bin_dir, ":", conda_env["PATH"])
    cmd = Cmd(`$jupyter`, env=conda_env)

    if check_nodejs
        node_exe = joinpath(conda_bin_dir, exe("node"))
        if !Sys.isfile(node_exe)
            install_node = Base.prompt(
                "NodeJS is not installed in your Conda environment " *
                "but is neccessary to install the JupyterLab extension. " *
                "Install it? [Y/n]"
            )
            if isyes(install_node)
                @eval Main using IJulia
                @eval Main.IJulia Conda.add("nodejs")
            end
        end
    end

    return cmd
end

function find_path_jupyter_cmd()::Cmd
    jupyter = Sys.which("jupyter")
    if jupyter === nothing
        error("Could not find `jupyter` executable in PATH.")
    end
    return `$jupyter`
end

"""
    find_jupyter_cmd([; conda=false])

Find the most likely candidate for the `jupyter` executable.
This will locate `jupyter` by searching the `PATH` environment variable and,
if not found, tries to return Conda.jl's jupyter.
If both of these approaches fail, an error is thrown.
"""
function find_jupyter_cmd(;
    # DEPRECATED: use `condajl` keyword argument
    force_conda_jupyter::Union{Nothing, Bool}=nothing,
    condajl::Union{Nothing, Bool}=force_conda_jupyter,
    check_nodejs::Bool=false,
)::Cmd
    # Try to find the "system" Jupyter (unless we explicitly want to force
    # using IJulia's Jupyter installation).
    if condajl === nothing || condajl === false
        try
            return find_path_jupyter_cmd()
        catch
            # If conda is false, we explicitly don't want to look for jupyter
            # in Conda.jl's installation, so we fail now.
            # Otherwise, conda is nothing and we can continue searching.
            if condajl === false
                rethrow()
            end
        end
    end

    try
        return find_condajl_jupyter_cmd(check_nodejs=check_nodejs)
    catch
        if condajl === nothing
            error("Could not find `jupyter` executable in PATH or Conda.jl.")
        end
        rethrow()
    end
end

"""
    install_jupyter_labextension([jupyter]; condajl=false)

Install the Jupyter Lab extension for WebIO using the specified `jupyter`
executable.
The executable defaults to the first one found in the `PATH` or Jupyter
installed via IJulia/Conda.jl.
To force using Conda.jl's jupyter, specify the `force_conda_jupyter=true`
keyword; this might be necessary if you launch Jupyter via IJulia in the Julia
REPL.

The [IJulia provider documentation](https://juliagizmos.github.io/WebIO.jl/latest/providers/ijulia/)
provides some more information (and caveats) about the relationship between
Jupyter Lab and WebIO.
"""
function install_jupyter_labextension(
        jupyter::Union{Cmd, Nothing}=nothing;
        # Deprecated: use `condajl` keyword argument.
        force_conda_jupyter::Union{Bool, Nothing}=nothing,
        condajl::Union{Bool, Nothing}=force_conda_jupyter,
        dev::Bool=isdev()
)
    if jupyter === nothing
        jupyter = find_jupyter_cmd(; condajl=condajl, check_nodejs=true)
        @info(
            "Using default Jupyter executable at $jupyter; to use a different "
            * "executable, see the documentation by running "
            * "`?WebIO.install_jupyter_labextension`."
        )
    end
    install_jupyter_serverextension()

    _tryrun(`$jupyter labextension unlink --no-build @webio/webio`)
    _tryrun(`$jupyter labextension uninstall --no-build @webio/jupyter-lab-provider`)

    if dev
        @info "Installing Jupyter labextension in dev mode..."
        core_path = joinpath(PACKAGES_PATH, "webio")
        lab_provider_path = joinpath(PACKAGES_PATH, "jupyter-lab-provider")
        run(`$jupyter labextension link --no-build $core_path`)
        run(`$jupyter labextension install --no-build $lab_provider_path`)
    else
        run(`$jupyter labextension install --no-build @webio/jupyter-lab-provider@$(WEBIO_VERSION)`)
    end
    run(`$jupyter lab build`)
end

"""
    install_jupyter_nbextension([jupyter])

Install the Jupyter Notebook extension (nbextension) for WebIO.
This copies the nbextension code to the appropriate place and writes the
appropriate configuration files.
"""
function install_jupyter_nbextension(
        jupyter::Cmd=find_jupyter_cmd();
        nbextension_flags::Cmd=`--user`,
)
    install_jupyter_serverextension()

    # Copy the nbextension files.
    install_cmd = `$jupyter nbextension install $nbextension_flags $JUPYTER_NBEXTENSION_PATH`
    @info "Installing Jupyter WebIO extension..." cmd=install_cmd
    run(install_cmd)

    enable_cmd = `$jupyter nbextension enable $nbextension_flags $JUPYTER_NBEXTENSION_NAME`
    @info "Enabling Jupyter WebIO extension..." cmd=enable_cmd
    run(enable_cmd)

    return nothing
end

### BEGIN BORROWED CODE ###
# This code is adapted from IJulia.jl (MIT license).
# https://github.com/JuliaLang/IJulia.jl/blob/68748c2b88b4b394e2802e911eb154b7008ebdd2/deps/kspec.jl
@static if Sys.iswindows()
    function appdata() # return %APPDATA%
        path = zeros(UInt16, 300)
        CSIDL_APPDATA = 0x001a
        result = ccall((:SHGetFolderPathW,:shell32), stdcall, Cint,
            (Ptr{Cvoid},Cint,Ptr{Cvoid},Cint,Ptr{UInt16}),C_NULL,CSIDL_APPDATA,C_NULL,0,path)
        return result == 0 ? transcode(String, resize!(path, findfirst(iszero, path)-1)) : get(ENV, "APPDATA", "")
    end
    function default_jupyter_data_dir()
        APPDATA = appdata()
        return !isempty(APPDATA) ? joinpath(APPDATA, "jupyter") : joinpath(get(ENV, "JUPYTER_CONFIG_DIR", joinpath(homedir(), ".jupyter")), "data")
    end
elseif Sys.isapple()
    default_jupyter_data_dir() = joinpath(homedir(), "Library/Jupyter")
else
    function default_jupyter_data_dir()
        xdg_data_home = get(ENV, "XDG_DATA_HOME", "")
        data_home = !isempty(xdg_data_home) ? xdg_data_home : joinpath(homedir(), ".local", "share")
        joinpath(data_home, "jupyter")
    end
end

function jupyter_data_dir()
    env_data_dir = get(ENV, "JUPYTER_DATA_DIR", "")
    !isempty(env_data_dir) ? env_data_dir : default_jupyter_data_dir()
end
### END BORROWED CODE ###

function jupyter_config_dir()
    env_config_dir = get(ENV, "JUPYTER_CONFIG_DIR", "")
    !isempty(env_config_dir) ? env_config_dir : joinpath(homedir(), ".jupyter")
end
jupyter_nbextensions_dir() = joinpath(jupyter_data_dir(), "nbextensions")
jupyter_nbconfig_dir() = joinpath(jupyter_config_dir(), "nbconfig")

isyes(s) = isempty(s) || lowercase(strip(s)) in ("y", "yes")

@static if Sys.iswindows()
    exe(x::String) = endswith(x, "exe") ? x : string(x, ".exe")
else
    exe(x::String) = x
end


text:
#!/usr/bin/env julia
using ArgParse

s = ArgParseSettings()

@add_arg_table s begin
    "--N"
    "--a"
    "--dphi"
    "--o"
end

parameters = parse_args(s)

# Output
const odir = parameters["o"]
@assert isdir(odir)

# System geometry
using QuantumOptics, CollectiveSpins
const N = parse(Int, parameters["N"])
const a = parse(Float64, parameters["a"])
const dphi = parse(Float64, parameters["dphi"])

keyparameters = Dict(
        "a"=>parameters["a"],
        "N"=>parameters["N"],
    )

omega_eff = NaN
gamma_eff = NaN

tic()
omega_eff, gamma_eff = CollectiveSpins.effective_interaction_rotated.chain_orthogonal(a, N, dphi)
t = toc()

name = QuantumOptics.io.dict2filename(keyparameters)
f = open(joinpath(odir, name), "w")
write(f, string(omega_eff))
write(f, ";")
write(f, string(gamma_eff))
write(f, "\n#calctime=")
write(f, string(t))
close(f)


text:
using Statistics
using Plots
using Flux
using OceanParameterizations
using Oceananigans.Grids
using BSON
using OrdinaryDiffEq, DiffEqSensitivity
using WindMixing
using JLD2
using FileIO

PATH = joinpath(pwd(), "extracted_training_output")
# PATH = "D:\\University Matters\\Massachusetts Institute of Technology\\CLiMA Project\\OceanParameterizations.jl\\training_output"

DATA_NAME = "parameter_optimisation_18sim_windcooling_windheating_5params_BFGS_T0.8_grad"

DATA_PATH = joinpath(PATH, "$(DATA_NAME)_extracted.jld2")
ispath(DATA_PATH)

FILE_PATH = joinpath(pwd(), "Output", DATA_NAME)

if !ispath(FILE_PATH)
    mkdir(FILE_PATH)
end

file = jldopen(DATA_PATH, "r")

train_files = file["training_info/train_files"]
train_parameters = file["training_info/parameters"]
loss_scalings = file["training_info/loss_scalings"]
mpp_parameters = file["parameters"]

𝒟train = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)
close(file)

ν₀, ν₋, ΔRi, Riᶜ, Pr = mpp_parameters

# ν₀ = 1f-4
# ν₋ = 1f-1
# ΔRi = 1f-1
# Riᶜ = 0.25f0
# Pr = 1f0

# ν₀ = train_parameters["ν₀"]
# ν₋ = train_parameters["ν₋"]
# ΔRi = train_parameters["ΔRi"]
# Riᶜ = train_parameters["Riᶜ"]
# Pr = train_parameters["Pr"]

N_inputs = 96
hidden_units = 400
N_outputs = 31

weights, re = Flux.destructure(Chain(Dense(N_inputs, hidden_units, relu), Dense(hidden_units, N_outputs)))

uw_NN = re(zeros(Float32, length(weights)))
vw_NN = re(zeros(Float32, length(weights)))
wT_NN = re(zeros(Float32, length(weights)))

to_run = [
    "wind_-5e-4_cooling_3e-8_new",
    "wind_-5e-4_cooling_1e-8_new",
    "wind_-2e-4_cooling_3e-8_new",
    "wind_-2e-4_cooling_1e-8_new",
    "wind_-5e-4_heating_-3e-8_new",
    "wind_-2e-4_heating_-1e-8_new",
    "wind_-2e-4_heating_-3e-8_new",
    "wind_-5e-4_heating_-1e-8_new",

    "wind_-3.5e-4_cooling_2e-8_new",
    "wind_-3.5e-4_heating_-2e-8_new",

    "wind_-5e-4_cooling_2e-8_new",
    "wind_-3.5e-4_cooling_3e-8_new",
    "wind_-3.5e-4_cooling_1e-8_new",
    "wind_-2e-4_cooling_2e-8_new",
    "wind_-3.5e-4_heating_-3e-8_new",
    "wind_-3.5e-4_heating_-1e-8_new",
    "wind_-2e-4_heating_-2e-8_new",
    "wind_-5e-4_heating_-2e-8_new",
  ]

for test_file in to_run
    @info "running $test_file"
    test_files = [test_file]
    𝒟test = WindMixing.data(test_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)
    trange = 1:1:1153
    plot_data = NDE_profile(uw_NN, vw_NN, wT_NN, test_file, 𝒟test, 𝒟train, trange,
                            modified_pacanowski_philander=true,
                            ν₀=ν₀, ν₋=ν₋, ΔRi=ΔRi, Riᶜ=Riᶜ, Pr=Pr,
                            convective_adjustment=false,
                            smooth_NN=false, smooth_Ri=false,
                            zero_weights=true,
                            loss_scalings=loss_scalings)

    animation_type = "Pre-Training"
    n_trainings = length(train_files)
    training_types = "Modified Pacanowski-Philander"
    VIDEO_NAME = "$(test_file)_mpp"
    animate_profiles_fluxes_comparison(plot_data, plot_data, plot_data, joinpath(FILE_PATH, VIDEO_NAME), fps=30,
                                                    animation_type=animation_type, n_trainings=n_trainings, training_types=training_types)
end


text:
using Pkg;Pkg.activate(".")
using Test, Plots, BenchmarkTools, JLD2
using VortexDistributions
gr(transpose=true)

## load test data
loadpath="/Users/abradley/Dropbox/Julia/Vortices - simple results"
f2 = joinpath(loadpath,"nv20-t_200.jld2")

## load good data
@load f2 𝛹200 x y
psi200 = Torus(𝛹200,x,y)
vfound = findvortices(psi200,periodic=true)
p1=heatmap(x,y,angle.(psi200.ψ))
vdata = vortex_array(vfound)
p1

## make a plot function
function plot_vortices!(p,vdata)
    for j in 1:size(vdata)[1]
        vx,vy = vdata[j,1],vdata[j,2]
        scatter!(p1,[vx],[vy],marker=vortex_marker(vdata[j,3]),label=false)
    end
return p1
end

## plot Vortices
plot_vortices!(p1,vdata)

##


text:
using CodecZlib

b64 = b"""eNrtVE1rwzAMvedXaKdeRn7ENrb21rHCzmrs1m49K9gOJv9+cko/HBcGg0LHcpOfnq2np0QL
    2FuKgBbICDAoeoiKwEc0hqIUgLAxfV0tQJCdhQM7qh68kheswKeBt5ROYetTemYMCC3rii//
    WMS3WkhXVyuFAaLT261JuBWwu4iDbvYp1tYzHVS68VEIObwFgaDB0KizuFs38aSdqKv3TgcJ
    uPYdn2B1opwIpeKE53qPftxRd88Y6uoVbdPzWxznrQ3ZUi3DudQ/bcELbevqM32iCIrj3IIh
    W6plOJf6L6xaajZjzqW/qAsKIvITBGs9Nm3glboZzkVP5l6Y+0bHLnedD0CttIyrpEU5Kv7N
    Mz3XkPBc/TSN3yxGiqMiipHRekycK0ZwMhM8jerGC9zuZaoTho3kMKSfJjLaF8v8wLzmXMqM
    zJvGew/jnZPzclA08yAkikegDTTUMfzwDXBcwoE="""
println(String(transcode(ZlibDecompressor(), base64decode(b64))))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
#Julia Visualizations in Jupyter

#Pkg.add("Gadfly")
using Gadfly

max_stars_repo_name: PacktPublishing/Learning-Jupyter

max_issues_repo_name: PacktPublishing/Learning-Jupyter

max_forks_repo_name: PacktPublishing/Learning-Jupyter


text:
using DataValues
using Test
using Dates
using InteractiveUtils

@testset "Core" begin

@testset "Missing integration" begin

@test DataValue(missing) == NA
@test DataValue{Int}(missing) == DataValue{Int}()
@test convert(Missing, NA) === missing
@test false == (DataValue(missing) != NA)

end

@testset "conversion" begin

@test DataValue{Float64}(2) == DataValue(2.)
@test convert(Union{Missing,Int}, DataValue(2)) == 2
@test convert(Union{Missing,Int}, DataValue{Int}()) === missing
@test convert(Union{Missing,Int}, NA) === missing

@test convert(DataValue{Int}, DataValue(3)) == DataValue(3)

end

@testset "isna" begin

@test DataValues.isna(NA) == true
@test DataValues.isna(DataValue()) == true
@test DataValues.isna(DataValue{Int}()) == true
@test DataValues.isna(DataValue(3)) == false

end

@testset "Show" begin

io = IOBuffer()
show(IOContext(io, :compact => true), DataValue(3))
@test String(take!(io)) == "3"

io = IOBuffer()
show(IOContext(io, :compact => true), DataValue{Int}())
@test String(take!(io)) == "#NA"

end

@testset "Unpack" begin

@test get(DataValue("test"), "default") == "test"
@test get(DataValue{String}(), "default") == "default"

@test DataValues.unsafe_get(DataValue(3)) == 3

end

@testset "Hasvalue" begin

@test hasvalue(DataValue(3)) == true
@test hasvalue(DataValue{Int}()) == false

end

@testset "zero" begin

@test zero(DataValue{Int}) == DataValue(0)
@test zero(DataValue(3)) == DataValue(0)

@test zero(DataValue{Float64}) == DataValue(0.)
@test zero(DataValue(3.)) == DataValue(0.)

for T in (subtypes(Dates.DatePeriod)..., subtypes(Dates.TimePeriod)...)
    @test zero(DataValue{T}()) == T(0)
end

@test zero(DataValue(Year(3))) == zero(Year)
@test zero(DataValue{Year}) == zero(Year)

end

@testset "Comparisons" begin

@test (DataValue(3) == NA) == false
@test (DataValue{Int}() == NA) == true

@test (NA == DataValue(3)) == false
@test (NA == DataValue{Int}()) == true

@test (DataValue(3) != NA) == true
@test (DataValue{Int}() != NA) == false

@test (NA != DataValue(3)) == true
@test (NA != DataValue{Int}()) == false

@test (NA == NA) == true

@test isless(DataValue{Int}(), DataValue{Int}()) == false
@test isless(DataValue{Int}(), DataValue{Int}(3)) == false
@test isless(DataValue{Int}(3), DataValue{Int}()) == true
@test isless(DataValue{Int}(3), DataValue{Int}(5)) == true

@test isless(3, DataValue{Int}()) == true
@test isless(3, DataValue{Int}(2)) == false

@test isless(DataValue{Int}(), 3) == false
@test isless(DataValue{Int}(2), 3) == true

@test isless(NA, NA) == false
@test isless(3, NA) == true
@test isless(NA, 3) == false

end

@testset "Promotion" begin

@test promote(DataValue(3), 5) == (DataValue(3), DataValue(5))
@test promote(DataValue(3), 5.) == (DataValue(3.), DataValue(5.))

end

# 3VL

@test DataValue(true) & DataValue(true) == DataValue(true)
@test DataValue(true) & DataValue(false) == DataValue(false)
@test DataValue(true) & DataValue{Bool}() == DataValue{Bool}()
@test DataValue(false) & DataValue(true) == DataValue(false)
@test DataValue(false) & DataValue(false) == DataValue(false)
@test DataValue(false) & DataValue{Bool}() == DataValue(false)
@test DataValue{Bool}() & DataValue(true) == DataValue{Bool}()
@test DataValue{Bool}() & DataValue(false) == DataValue(false)
@test DataValue{Bool}() & DataValue{Bool}() == DataValue{Bool}()

@test true & DataValue(true) == DataValue(true)
@test true & DataValue(false) == DataValue(false)
@test true & DataValue{Bool}() == DataValue{Bool}()
@test false & DataValue(true) == DataValue(false)
@test false & DataValue(false) == DataValue(false)
@test false & DataValue{Bool}() == DataValue(false)

@test DataValue(true) & true == DataValue(true)
@test DataValue(true) & false == DataValue(false)
@test DataValue(false) & true == DataValue(false)
@test DataValue(false) & false == DataValue(false)
@test DataValue{Bool}() & true == DataValue{Bool}()
@test DataValue{Bool}() & false == DataValue(false)

@test DataValue(true) | DataValue(true) == DataValue(true)
@test DataValue(true) | DataValue(false) == DataValue(true)
@test DataValue(true) | DataValue{Bool}() == DataValue(true)
@test DataValue(false) | DataValue(true) == DataValue(true)
@test DataValue(false) | DataValue(false) == DataValue(false)
@test DataValue(false) | DataValue{Bool}() == DataValue{Bool}()
@test DataValue{Bool}() | DataValue(true) == DataValue(true)
@test DataValue{Bool}() | DataValue(false) == DataValue{Bool}()
@test DataValue{Bool}() | DataValue{Bool}() == DataValue{Bool}()

@test true | DataValue(true) == DataValue(true)
@test true | DataValue(false) == DataValue(true)
@test true | DataValue{Bool}() == DataValue(true)
@test false | DataValue(true) == DataValue(true)
@test false | DataValue(false) == DataValue(false)
@test false | DataValue{Bool}() == DataValue{Bool}()

@test DataValue(true) | true == DataValue(true)
@test DataValue(true) | false == DataValue(true)
@test DataValue(false) | true == DataValue(true)
@test DataValue(false) | false == DataValue(false)
@test DataValue{Bool}() | true == DataValue(true)
@test DataValue{Bool}() | false == DataValue{Bool}()

@test !DataValue(true) == DataValue(false)
@test !DataValue(false) == DataValue(true)
@test !DataValue{Bool}() == DataValue{Bool}()

# NA comparisons
# @test (DataValue(5)==NA) == false
# @test (DataValue{Int}()==NA) == true
# @test (NA==DataValue(5)) == false
# @test (NA==DataValue{Int}()) == true

# @test (DataValue(5)!=NA) == true
# @test (DataValue{Int}()!=NA) == false
# @test (NA!=DataValue(5)) == true
# @test (NA!=DataValue{Int}()) == false

:+, :-, :!, :~
@test +DataValue(1) == DataValue(+1)
@test +DataValue{Int}() == DataValue{Int}()
@test -DataValue(1) == DataValue(-1)
@test -DataValue{Int}() == DataValue{Int}()
@test ~DataValue(1) == DataValue(~1)
@test ~DataValue{Int}() == DataValue{Int}()

# TODO add ^, / back
for op in (:+, :-, :*, :%, :&, :|, :<<, :>>)
    @eval begin
        @test $op(DataValue(3), DataValue(5)) == DataValue($op(3, 5))
        @test $op(DataValue{Int}(), DataValue(5)) == DataValue{Int}()
        @test $op(DataValue(3), DataValue{Int}()) == DataValue{Int}()
        @test $op(DataValue{Int}(), DataValue{Int}()) == DataValue{Int}()

        @test $op(DataValue{Int}(3), 5) == DataValue($op(3, 5))
        @test $op(3, DataValue{Int}(5)) == DataValue($op(3, 5))
        @test $op(DataValue{Int}(), 5) == DataValue{Int}()
        @test $op(3, DataValue{Int}()) == DataValue{Int}()
    end
end

@test DataValue(Int16(4)) / DataValue(Int32(2)) == DataValue(2.)
@test DataValue{Int16}() / DataValue(Int32(2)) == DataValue{Float64}()
@test DataValue(Int16(4)) / DataValue{Int32}() == DataValue{Float64}()
@test DataValue{Int16}() / DataValue{Int32}() == DataValue{Float64}()
@test Int16(4) / DataValue(Int32(2)) == DataValue(2.)
@test Int16(4) / DataValue{Int32}() == DataValue{Float64}()
@test DataValue{Int16}(4) / Int32(2) == DataValue{Float64}(2.)
@test DataValue{Int16}() / Int32(2) == DataValue{Float64}()

@test DataValue(3)^2 == DataValue(9)
@test DataValue{Int}()^2 == DataValue{Int}()

@test DataValue(3) == DataValue(3)
@test !(DataValue(3) == DataValue(4))
@test !(DataValue{Int}() == DataValue(3))
@test !(DataValue{Float64}() == DataValue(3))
@test !(DataValue(3) == DataValue{Int}())
@test !(DataValue(3) == DataValue{Float64}())
@test DataValue{Int}() == DataValue{Int}()
@test DataValue{Int}() == DataValue{Float64}()

@test DataValue(3) == 3
@test 3 == DataValue(3)
@test !(DataValue(3) == 4)
@test !(4 == DataValue(3))
@test !(DataValue{Int}() == 3)
@test !(3 == DataValue{Int}())

@test !(DataValue(3) != DataValue(3))
@test DataValue(3) != DataValue(4)
@test DataValue{Int}() != DataValue(3)
@test DataValue{Float64}() != DataValue(3)
@test DataValue(3) != DataValue{Int}()
@test DataValue(3) != DataValue{Float64}()
@test !(DataValue{Int}() != DataValue{Int}())
@test !(DataValue{Int}() != DataValue{Float64}())

@test !(DataValue(3) != 3)
@test !(3 != DataValue(3))
@test DataValue(3) != 4
@test 4 != DataValue(3)
@test DataValue{Int}() != 3
@test 3 != DataValue{Int}()

@test DataValue(4) > DataValue(3)
@test !(DataValue(3) > DataValue(4))
@test !(DataValue(4) > DataValue{Int}())
@test !(DataValue{Int}() > DataValue(3))
@test !(DataValue{Int}() > DataValue{Int}())

@test DataValue(4) > 3
@test !(DataValue(3) > 4)
@test !(DataValue{Int}() > 3)

@test 4 > DataValue(3)
@test !(3 > DataValue(4))
@test !(4 > DataValue{Int}())

@test lowercase(DataValue("TEST"))==DataValue("test")
@test lowercase(DataValue{String}())==DataValue{String}()

@test DataValue("TEST")[2:end]==DataValue("EST")
@test DataValue{String}()[2:end]==DataValue{String}()

@test length(DataValue("TEST"))==DataValue(4)
@test length(DataValue{String}())==DataValue{Int}()

@test DataValue{Int} == DataValue{Int}
@test DataValue(43) == DataValue(43)

io = IOBuffer()

show(io, DataValue(enum_val_a))
@test String(take!(io)) == "DataValue{TestEnum}(enum_val_a)"

@testset "string" begin
    @test rstrip("asdf  ") == rstrip(DataValue("asdf"))
    @test rstrip(DataValue{String}()) == NA

    @test rstrip("asdf  ", ' ') == rstrip(DataValue("asdf"), ' ')
    @test rstrip(DataValue{String}(), ' ') == NA

    @test rstrip(isspace, "asdf  ") == rstrip(isspace, DataValue("asdf"))
    @test rstrip(isspace, DataValue{String}()) == NA

    @test lstrip(" asdf") == lstrip(DataValue("asdf"))
    @test lstrip(DataValue{String}()) == NA

    @test lstrip(" asdf", ' ') == lstrip(DataValue("asdf"), ' ')
    @test lstrip(DataValue{String}(), ' ') == NA

    @test lstrip(isspace, " asdf") == lstrip(isspace, DataValue("asdf"))
    @test lstrip(isspace, DataValue{String}()) == NA
end

using DataValueInterfaces

@testset "DataValueInterfaces" begin

    @test DataValueInterfaces.nondatavaluetype(DataValue{Union{}}) == Missing
    @test DataValueInterfaces.nondatavaluetype(DataValue{Int}) == Union{Int, Missing}
    @test DataValueInterfaces.nondatavaluetype(Union{}) == Union{}
    @test DataValueInterfaces.nondatavaluetype(Int) == Int

    @test DataValueInterfaces.datavaluetype(Union{}) == Union{}
    @test DataValueInterfaces.datavaluetype(Int) == Int
    @test DataValueInterfaces.datavaluetype(Missing) == DataValue{Union{}}
    @test DataValueInterfaces.datavaluetype(Union{Int, Missing}) == DataValue{Int}

end

end


text:
function calc_B_C(reactor, R_0, B_0)
  tmp_epsilon_B = reactor.epsilon + reactor.b / R_0

  B_0 / ( 1 - tmp_epsilon_B )
end