text:
using Plots
using Random
Random.seed!(1)

include_ddp()

function visualize!(vis, model, q;
      r = 0.040, Δt = 0.1)

	default_background!(vis)

	torso = Cylinder(Point3f0(0.0, 0.0, 0.0), Point3f0(0.0, 0.0, model.l_torso),
		convert(Float32, 0.035))
	setobject!(vis["torso"], torso,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	thigh_1 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l_thigh1),
		convert(Float32, 0.035))
	setobject!(vis["thigh1"], thigh_1,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	calf_1 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l_calf1),
		convert(Float32, 0.035))
	setobject!(vis["calf1"], calf_1,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	foot_1 = Cylinder(Point3f0(0.0,0.0,0.0),
		Point3f0(0.0, 0.0, model.l_foot1 + model.d_foot1),
		convert(Float32, 0.035))
	setobject!(vis["foot1"], foot_1,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	thigh_2 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l_thigh2),
		convert(Float32, 0.035))
	setobject!(vis["thigh2"], thigh_2,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	calf_2 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l_calf2),
		convert(Float32, 0.035))
	setobject!(vis["calf2"], calf_2,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	foot_2 = Cylinder(Point3f0(0.0,0.0,0.0),
		Point3f0(0.0, 0.0, model.l_foot2 + model.d_foot2),
		convert(Float32, 0.035))
	setobject!(vis["foot2"], foot_2,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	setobject!(vis["heel1"], Sphere(Point3f0(0.0),
		convert(Float32, r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, 1.0)))
	setobject!(vis["heel2"], Sphere(Point3f0(0.0),
		convert(Float32, r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, 1.0)))
	setobject!(vis["toe1"], Sphere(Point3f0(0.0),
		convert(Float32, r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, 1.0)))
	setobject!(vis["toe2"], Sphere(Point3f0(0.0),
		convert(Float32, r)),
		MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0.0, 1.0)))
	setobject!(vis["knee1"], Sphere(Point3f0(0.0),
		convert(Float32, 0.035)),
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))
	setobject!(vis["knee2"], Sphere(Point3f0(0.0),
		convert(Float32, 0.035)),
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))
	setobject!(vis["hip"], Sphere(Point3f0(0.0),
		convert(Float32, 0.035)),
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))
	setobject!(vis["torso_top"], Sphere(Point3f0(0.0),
		convert(Float32, 0.035)),
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

	anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

	T = length(q)
	p_shift = [0.0; 0.0; r]
	for t = 1:T
		MeshCat.atframe(anim, t) do
			p = [q[t][1]; 0.0; q[t][2]] + p_shift

			k_torso = kinematics_1(model, q[t], body = :torso, mode = :ee)
			p_torso = [k_torso[1], 0.0, k_torso[2]] + p_shift

			k_thigh_1 = kinematics_1(model, q[t], body = :thigh_1, mode = :ee)
			p_thigh_1 = [k_thigh_1[1], 0.0, k_thigh_1[2]] + p_shift

			k_calf_1 = kinematics_2(model, q[t], body = :calf_1, mode = :ee)
			p_calf_1 = [k_calf_1[1], 0.0, k_calf_1[2]] + p_shift

			k_thigh_2 = kinematics_1(model, q[t], body = :thigh_2, mode = :ee)
			p_thigh_2 = [k_thigh_2[1], 0.0, k_thigh_2[2]] + p_shift

			k_calf_2 = kinematics_2(model, q[t], body = :calf_2, mode = :ee)
			p_calf_2 = [k_calf_2[1], 0.0, k_calf_2[2]] + p_shift

			k_toe_1 = kinematics_3(model, q[t], body = :foot_1, mode = :toe)
			p_toe_1 = [k_toe_1[1], 0.0, k_toe_1[2]] + p_shift

			k_heel_1 = kinematics_3(model, q[t], body = :foot_1, mode = :heel)
			p_heel_1 = [k_heel_1[1], 0.0, k_heel_1[2]] + p_shift

			k_toe_2 = kinematics_3(model, q[t], body = :foot_2, mode = :toe)
			p_toe_2 = [k_toe_2[1], 0.0, k_toe_2[2]] + p_shift

			k_heel_2 = kinematics_3(model, q[t], body = :foot_2, mode = :heel)
			p_heel_2 = [k_heel_2[1], 0.0, k_heel_2[2]] + p_shift

			settransform!(vis["thigh1"], cable_transform(p, p_thigh_1))
			settransform!(vis["calf1"], cable_transform(p_thigh_1, p_calf_1))
			settransform!(vis["foot1"], cable_transform(p_toe_1, p_heel_1))

			settransform!(vis["thigh2"], cable_transform(p, p_thigh_2))
			settransform!(vis["calf2"], cable_transform(p_thigh_2, p_calf_2))
			settransform!(vis["foot2"], cable_transform(p_toe_2, p_heel_2))

			settransform!(vis["torso"], cable_transform(p_torso,p))
			settransform!(vis["heel1"], Translation(p_heel_1))
			settransform!(vis["heel2"], Translation(p_heel_2))
			settransform!(vis["toe1"], Translation(p_toe_1))
			settransform!(vis["toe2"], Translation(p_toe_2))
			settransform!(vis["knee1"], Translation(p_thigh_1))
			settransform!(vis["knee2"], Translation(p_thigh_2))
			settransform!(vis["hip"], Translation(p))
			settransform!(vis["torso_top"], Translation(p_torso))
		end
	end

	settransform!(vis["/Cameras/default"],
	    compose(Translation(0.0, 0.0, -1.0), LinearMap(RotZ(-pi / 2.0))))

	MeshCat.setanimation!(vis, anim)
end

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)

contact_control_path = "/home/taylor/Research/ContactControl.jl/src"

using Parameters
# Utilities
include(joinpath(contact_control_path, "utils.jl"))

# Solver
include(joinpath(contact_control_path, "solver/cones.jl"))
include(joinpath(contact_control_path, "solver/interior_point.jl"))
include(joinpath(contact_control_path, "solver/lu.jl"))

# Environment
include(joinpath(contact_control_path, "simulator/environment.jl"))

# Dynamics
include(joinpath(contact_control_path, "dynamics/model.jl"))

# Simulator
include(joinpath(contact_control_path, "simulation/contact_methods.jl"))
include(joinpath(contact_control_path, "simulation/simulation.jl"))
include(joinpath(contact_control_path, "simulator/trajectory.jl"))

include(joinpath(contact_control_path, "dynamics/code_gen_dynamics.jl"))
include(joinpath(contact_control_path, "dynamics/fast_methods_dynamics.jl"))

# Models
include(joinpath(contact_control_path, "dynamics/quaternions.jl"))
include(joinpath(contact_control_path, "dynamics/mrp.jl"))
include(joinpath(contact_control_path, "dynamics/euler.jl"))

# include("dynamics/particle_2D/model.jl")
# include("dynamics/particle/model.jl")
# include(joinpath(contact_control_path, "dynamics/hopper_2D/model.jl"))
# include("dynamics/hopper_3D/model.jl")
# include("dynamics/hopper_3D_quaternion/model.jl")
# include("dynamics/quadruped/model.jl")
# include("dynamics/quadruped_simple/model.jl")
# include("dynamics/biped/model.jl")
include(joinpath(contact_control_path, "dynamics/flamingo/model.jl"))
# include("dynamics/pushbot/model.jl")
# include("dynamics/planarpush/model.jl")
# include("dynamics/planarpush_2D/model.jl")
# include("dynamics/rigidbody/model.jl")
# include("dynamics/box/model.jl")

# Simulation
include(joinpath(contact_control_path, "simulation/environments/flat.jl"))
# include("simulation/environments/piecewise.jl")
# include("simulation/environments/quadratic.jl")
# include("simulation/environments/slope.jl")
# include("simulation/environments/sinusoidal.jl")
# include("simulation/environments/stairs.jl")

include(joinpath(contact_control_path, "simulation/residual_approx.jl"))
include(joinpath(contact_control_path, "simulation/code_gen_simulation.jl"))

# Visuals
using MeshCatMechanisms
include(joinpath(contact_control_path, "dynamics/visuals.jl"))
include(joinpath(contact_control_path, "dynamics/visual_utils.jl"))

s = get_simulation("flamingo", "flat_2D_lc", "flat")

@load joinpath(pwd(), "examples/contact_implicit", "flamingo_stand_100hz.jld2") q u γ b ψ η μ h
@load joinpath(pwd(), "examples/contact_implicit", "flamingo_gait_100hz.jld2") qm um γm bm ψm ηm μm hm

nq = s.model.dim.q
m = s.model.dim.u

T = 25
# h = hm
# q0 = qm[1]
# q1 = qm[2]
# qT = qm[end]
q0 = q[1]
q1 = q[2]

qT1 = qm[1]
qT = qm[2]

strd_diff = kinematics(s.model, qT1)[3] - kinematics(s.model, q0)[3]

q0[1] += strd_diff
q1[1] += strd_diff

x1 = [q0; q1]
# xT = [q0; q1]#
xT = [qT1; qT]
u1 = u[1]
visualize!(vis, s.model, [q0, q1], Δt = h)

visualize!(vis, s.model, [qT1, qT], Δt = h)

q_ref = [q0, linear_interpolation(q1, qT1, T-1)..., qT]
visualize!(vis, s.model, q_ref, Δt = h)

struct Dynamics{T}
	s::Simulation
	ip_dyn::InteriorPoint
	ip_jac::InteriorPoint
	h::T
end

function gen_dynamics(s::Simulation, h;
		dyn_opts =  InteriorPointOptions{Float64}(
						r_tol = 1.0e-8,
						κ_tol = 1.0e-4,
						κ_init = 0.1,
						diff_sol = true),
		jac_opts =  InteriorPointOptions{Float64}(
						r_tol = 1.0e-8,
						κ_tol = 1.0e-4,
						κ_init = 0.1,
						diff_sol = true))

	z = zeros(num_var(s.model, s.env))
	θ = zeros(num_data(s.model))

	ip_dyn = interior_point(z, θ,
		idx_ineq = inequality_indices(s.model, s.env),
		r! = s.res.r!,
		rz! = s.res.rz!,
		rθ! = s.res.rθ!,
		rz = s.rz,
		rθ = s.rθ,
		opts = dyn_opts)

	ip_dyn.opts.diff_sol = false

	ip_jac = interior_point(z, θ,
		idx_ineq = inequality_indices(s.model, s.env),
		r! = s.res.r!,
		rz! = s.res.rz!,
		rθ! = s.res.rθ!,
		rz = s.rz,
		rθ = s.rθ,
		opts = jac_opts)

	ip_jac.opts.diff_sol = true

	Dynamics(s, ip_dyn, ip_jac, h)
end

d = gen_dynamics(s, h,
	dyn_opts = InteriorPointOptions{Float64}(κ_tol = 1.0e-4, κ_init = 0.1),
	jac_opts = InteriorPointOptions{Float64}(κ_tol = 1.0e-4, κ_init = 0.1))

function f!(d::Dynamics, q0, q1, u1, mode = :dynamics)
	s = d.s
	ip = (mode == :dynamics ? d.ip_dyn : d.ip_jac)
	h = d.h

	z_initialize!(ip.z, s.model, s.env, copy(q1))
	θ_initialize!(ip.θ, s.model, copy(q0), copy(q1), copy(u1), zeros(s.model.dim.w), s.model.μ_world, h)

	status = interior_point_solve!(ip)

	!status && (@warn "dynamics failure")
end

function f(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :dynamics)
	return copy(d.ip_dyn.z[1:d.s.model.dim.q])
end

f(d, q0, q1, zeros(m))

function fq0(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 1:d.s.model.dim.q])
end

fq0(d, q0, q1, zeros(m))

function fq1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, d.s.model.dim.q .+ (1:d.s.model.dim.q)])
end

fq1(d, q0, q1, zeros(m))

function fx1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 1:(2 * d.s.model.dim.q)])
end

fx1(d, q0, q1, zeros(m))

function fu1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 2 * d.s.model.dim.q .+ (1:d.s.model.dim.u)])
end

fu1(d, q0, q1, zeros(m))

struct FlamingoCI{I, T} <: Model{I, T}
    n::Int
    m::Int
    d::Int
	dynamics::Dynamics
end

model = FlamingoCI{Midpoint, FixedTime}(2 * s.model.dim.q, s.model.dim.u, 0, d)

function fd(model::FlamingoCI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]

	q2 = f(model.dynamics, q0, q1, u)

	return [q1; q2]
end

fd(model, x1, u1, zeros(0), h, 1)

function fdx(model::FlamingoCI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]
	dq2dx1 = fx1(model.dynamics, q0, q1, u)

	return [zeros(nq, nq) I; dq2dx1]
end

fdx(model, x1, u1, zeros(0), h, 1)

function fdu(model::FlamingoCI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]
	dq2du1 = fu1(model.dynamics, q0, q1, u)
	return [zeros(nq, model.m); dq2du1]
end

fdu(model, x1, u1, zeros(0), h, 1)

n = model.n
m = model.m

T = 36
# ū = [u[1] + 0.0 * randn(model.m) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
# @save joinpath(pwd(), "examples/contact_implicit", "flamingo_stand_100hz_v2.jld2") u_stand
@load joinpath(pwd(), "examples/contact_implicit", "flamingo_stand_100hz_v2.jld2") u_stand

ū = [u_stand + 0.00 * randn(model.m) for t = 1:T-1]

x̄ = rollout(model, x1, ū, w, h, T)

q̄ = state_to_configuration(x̄)
visualize!(vis, s.model, q̄, Δt = h)

# Objective
Q = [t < T ? 1.0 * Diagonal(ones(model.n)) : 1.0 * Diagonal(ones(model.n)) for t = 1:T]
q = [-2.0 * Q[t] * xT for t = 1:T]
R = [Diagonal(1.0e-1 * ones(model.m)) for t = 1:T-1]
r = [-2.0 * R[t] * u_stand for t = 1:T-1]

obj = StageCosts([QuadraticCost(Q[t], q[t],
	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = obj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

# Constraints
ul = -1.0 * ones(model.m)
uu = 1.0 * ones(model.m)
p = [t < T ? 2 * m : n for t = 1:T]
info_t = Dict(:ul => ul, :uu => uu, :inequality => (1:2 * m))
info_T = Dict(:xT => xT)
con_set = [StageConstraint(p[t], t < T ? info_t : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	p = cons.con[t].p

	if t < T
		ul = cons.con[t].info[:ul]
		uu = cons.con[t].info[:uu]
		c .= [ul - u; u - uu]
	else
		c .= x - cons.con[T].info[:xT]
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T,
	analytical_dynamics_derivatives = true)

# Solve
@time constrained_ddp_solve!(prob,
	max_iter = 1000, max_al_iter = 3,
	ρ_init = 1.0, ρ_scale = 10.0,
	con_tol = 1.0e-3)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

q̄ = state_to_configuration(x̄)

vis = Visualizer()
render(vis)
visualize!(vis, s.model, q̄, Δt = h)

plot(hcat(ū...)', linetype = :steppost)
function visualize!(vis, model, q;
		Δt = 0.1, scenario = :vertical)

    r_foot = 0.05
    r_leg = 0.5 * r_foot

	default_background!(vis)

    setobject!(vis["body"], Sphere(Point3f0(0),
        convert(Float32, 0.1)),
        MeshPhongMaterial(color = RGBA(0, 1, 0, 1.0)))

    setobject!(vis["foot"], Sphere(Point3f0(0),
        convert(Float32, r_foot)),
        MeshPhongMaterial(color = RGBA(1.0, 165.0 / 255.0, 0, 1.0)))

    n_leg = 100
    for i = 1:n_leg
        setobject!(vis["leg$i"], Sphere(Point3f0(0),
            convert(Float32, r_leg)),
            MeshPhongMaterial(color = RGBA(0, 0, 0, 1.0)))
    end

    p_leg = [zeros(3) for i = 1:n_leg]
    anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

    for t = 1:length(q)
        p_body = [q[t][1], 0.0, q[t][2]]
        p_foot = [kinematics(model, q[t])[1], 0.0, kinematics(model, q[t])[2]]

        q_tmp = Array(copy(q[t]))
        r_range = range(0, stop = q[t][4], length = n_leg)
        for i = 1:n_leg
            q_tmp[4] = r_range[i]
            p_leg[i] = [kinematics(model, q_tmp)[1], 0.0, kinematics(model, q_tmp)[2]]
        end
        q_tmp[4] = q[t][4]
        p_foot = [kinematics(model, q_tmp)[1], 0.0, kinematics(model, q_tmp)[2]]

        z_shift = [0.0; 0.0; r_foot]

        MeshCat.atframe(anim, t) do
            settransform!(vis["body"], Translation(p_body + z_shift))
            settransform!(vis["foot"], Translation(p_foot + z_shift))

            for i = 1:n_leg
                settransform!(vis["leg$i"], Translation(p_leg[i] + z_shift))
            end
        end
    end

	if scenario == :vertical
		settransform!(vis["/Cameras/default"],
			compose(Translation(0.0, 0.5, -1.0),LinearMap(RotZ(-pi / 2.0))))
	end

    MeshCat.setanimation!(vis, anim)
end

# # Model
# include_model("double_integrator")
#
# function f(model::DoubleIntegratorContinuous, x, u, w)
#     [x[2]; (1.0 + w[1]) * u[1]]
# end
#
# function fd(model::DoubleIntegratorContinuous{Midpoint, FixedTime}, x, u, w, h, t)
# 	x + h * f(model, x + 0.5 * h * f(model, x, u, w), u, w)
# end
#
# model = DoubleIntegratorContinuous{Midpoint, FixedTime}(2, 1, 1)
# n = model.n
# m = model.m
#
# # Time
# T = 11
# h = 0.1
#
# # Initial conditions, controls, disturbances
# x1 = [1.0; 0.0]
# x_ref = [[0.0; 0.0] for t = 1:T]
# xT = [0.0; 0.0]
# ū = [1.0 * randn(model.m) for t = 1:T-1]
# u_ref = [zeros(model.m) for t = 1:T-1]
# w = [zeros(model.d) for t = 1:T-1]
#
# # Rollout
# x̄ = rollout(model, x1, ū, w, h, T)
#
# # Objective
# Q = [(t < T ? h : 1.0) * (t < T ?
# 	 Diagonal([1.0; 1.0])
# 		: Diagonal([1.0; 1.0])) for t = 1:T]
# q = [-2.0 * Q[t] * x_ref[t] for t = 1:T]
# R = h * [Diagonal(1.0 * ones(model.m)) for t = 1:T-1]
# r = [zeros(model.m) for t = 1:T-1]
#
# obj = StageCosts([QuadraticCost(Q[t], q[t],
# 	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)
#
# function g(obj::StageCosts, x, u, t)
# 	T = obj.T
#     if t < T
# 		Q = obj.cost[t].Q
# 		q = obj.cost[t].q
# 	    R = obj.cost[t].R
# 		r = obj.cost[t].r
#         return x' * Q * x + q' * x + u' * R * u + r' * u
#     elseif t == T
# 		Q = obj.cost[T].Q
# 		q = obj.cost[T].q
#         return x' * Q * x + q' * x
#     else
#         return 0.0
#     end
# end
#
# # Constraints
# ul = [-5.0]
# uu = [5.0]
# p = [t < T ? 2 * m : n for t = 1:T]
# info_t = Dict(:ul => ul, :uu => uu, :inequality => (1:2 * m))
# info_T = Dict(:xT => xT)
# con_set = [StageConstraint(p[t], t < T ? info_t : info_T) for t = 1:T]
#
# function c!(c, cons::StageConstraints, x, u, t)
# 	T = cons.T
# 	p = cons.con[t].p
#
# 	if t < T
# 		ul = cons.con[t].info[:ul]
# 		uu = cons.con[t].info[:uu]
# 		c .= [ul - u; u - uu]
# 	else
# 		c .= x - cons.con[T].info[:xT]
# 	end
# end
#
# prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T)
#
# # Solve
# @time constrained_ddp_solve!(prob,
# 	max_iter = 1000, max_al_iter = 10,
# 	ρ_init = 1.0, ρ_scale = 10.0,
# 	con_tol = 1.0e-5)
#
# x, u = current_trajectory(prob)
# x̄, ū = nominal_trajectory(prob)
#
# # Visualize
# using Plots
# # plot(hcat([[0.0; 0.0] for t = 1:T]...)',
# #     width = 2.0, color = :black, label = "")
# plt = plot(hcat(x...)',
# 	width = 2.0,
# 	color = [:cyan :orange],
# 	label = ["x" "ẋ"],
# 	xlabel = "time step",
# 	ylabel = "state")
#
# savefig(plt,
# 	joinpath("/home/taylor/Research/parameter_optimization_manuscript/figures/di_base_state.png"))
#
# plt = plot(hcat(u..., u[end])',
# 	width = 2.0,
# 	color = :magenta,
# 	linetype = :steppost,
# 	xlabel = "time step",
# 	ylabel = "control",
# 	label = "")
#
# savefig(plt,
# 	joinpath("/home/taylor/Research/parameter_optimization_manuscript/figures/di_base_control.png"))
#
# # Simulate policy
# include(joinpath(@__DIR__, "simulate.jl"))
#
# # Model
# model_sim = DoubleIntegratorContinuous{RK3, FixedTime}(model.n, model.m, model.d)
# x1_sim = copy(x1)
# T_sim = 10 * T
#
# # Time
# tf = h * (T - 1)
# t = range(0, stop = tf, length = T)
# t_sim = range(0, stop = tf, length = T_sim)
# dt_sim = tf / (T_sim - 1)
#
# # Policy
# K = [K for K in prob.p_data.K]
# plot(vcat(K...))
# K = [prob.p_data.K[t] for t = 1:T-1]
# # K, _ = tvlqr(model, x̄, ū, h, Q, R)
# # # K = [-k for k in K]
# # K = [-K[1] for t = 1:T-1]
# # plot(vcat(K...))
#
# # Simulate
# N_sim = 1
# x_sim = []
# u_sim = []
# J_sim = []
# Random.seed!(1)
# for k = 1:N_sim
# 	wi_sim = 0.0 * min(0.1, max(-0.1, 1.0e-1 * randn(1)[1]))
# 	w_sim = [wi_sim for t = 1:T-1]
# 	println("sim: $k - w = $(wi_sim[1])")
#
# 	x_ddp, u_ddp, J_ddp, Jx_ddp, Ju_ddp = simulate_linear_feedback(
# 		model_sim,
# 		K,
# 	    x̄, ū,
# 		x_ref, u_ref,
# 		Q, R,
# 		T_sim, h,
# 		x1_sim,
# 		w_sim,
# 		ul = ul,
# 		uu = uu)
#
# 	push!(x_sim, x_ddp)
# 	push!(u_sim, u_ddp)
# 	push!(J_sim, J_ddp)
# end
#
# # Visualize
# idx = (1:2)
# plt = plot(t, hcat(x̄...)[idx, :]',
# 	width = 2.0, color = :black, label = "",
# 	xlabel = "time (s)", ylabel = "state",
# 	title = "double integrator (J_avg = $(round(mean(J_sim), digits = 3)), N_sim = $N_sim)")
#
# for xs in x_sim
# 	plt = plot!(t_sim, hcat(xs...)[idx, :]',
# 	    width = 1.0, color = :magenta, label = "")
# end
# display(plt)
#
# plt = plot(t, hcat(ū..., ū[end])',
# 	width = 2.0, color = :black, label = "",
# 	xlabel = "time (s)", ylabel = "control",
# 	linetype = :steppost,
# 	title = "double integrator (J_avg = $(round(mean(J_sim), digits = 3)), N_sim = $N_sim)")
#
# for us in u_sim
# 	plt = plot!(t_sim, hcat(us..., us[end])',
# 		width = 1.0, color = :magenta, label = "",
# 		linetype = :steppost)
# end
# display(plt)
# # u_sim
# # plot(vcat(K...))

max_issues_repo_name: thowell/DirectMotionPlanning


text:
function set_plane_strain_oop_stress(σᵢⱼ,r,D ; abstol=1e-12, maxiter=100, σoop_guess=nothing)
  # If the initial out of plane is different from σ₃₃ insert a guess
  !isnothing(σoop_guess) && (σᵢⱼ = insert_σoop(σᵢⱼ,σoop_guess))
  σoop = σᵢⱼ[3,3]
  # get damage constants:
  A1,B1 = compute_A1B1(r,D)

  for i in 1:maxiter

    dϵ_oop, ϵ_oop = Tensors.gradient(σᵢⱼ -> compute_ϵ_oop(r,A1,B1,σᵢⱼ), σᵢⱼ, :all)
    #println(dϵ_oop)
    dϵ_oop = dϵ_oop[3,3]
    # Newton update
    σoop = σoop - ϵ_oop/dϵ_oop
    σᵢⱼ = insert_σoop(σᵢⱼ,σoop)
    # exit condition
    (abs(ϵ_oop) <= abstol) && break
  end
  return σᵢⱼ
end

function set_plane_strain_oop_stress_rate(σᵢⱼ,σ̇ᵢⱼ,r,D,Δt ; abstol=1e-16, maxiter=100, damaged_allowed=true)
  # If the initial out of plane is different from σ₃₃ insert a guess
  #!isnothing(σoop_guess) && (σᵢⱼ = insert_σoop(σᵢⱼ,σoop_guess)) # Removed because of type inference issues
  σoop = σᵢⱼ[3,3]
  σoop_rate = σ̇ᵢⱼ[3,3]
  σᵢⱼnext = σᵢⱼ + σ̇ᵢⱼ*Δt
  # get damage constants:
  A1,B1 = compute_A1B1(r,D)
  # preallocate differentiation results
  result = DiffResults.GradientResult(σᵢⱼnext) # careful, result stores values in Matrix not Tensor
  for i in 1:maxiter
    # Joint computation of value and grad of ϵ_oop
    ForwardDiff.gradient!(result, σᵢⱼnext -> compute_ϵ̇ij(r,D,σᵢⱼ,σᵢⱼnext,Δt ; damaged_allowed)[1][3,3], σᵢⱼnext)
    dϵ̇_oop = DiffResults.gradient(result)
    ϵ̇_oop  = DiffResults.value(result)
    #println(dϵ̇_oop)
    dϵ̇_oop_dσ̇_oop = dϵ̇_oop[3,3]
    # Newton update
    σoop_rate = σoop_rate - ϵ̇_oop/dϵ̇_oop_dσ̇_oop
    σᵢⱼnext = insert_σoop(σᵢⱼnext,σoop_rate)
    # exit condition
    (abs(ϵ̇_oop) <= abstol) && break
  end
  σ̇ᵢⱼ = (σᵢⱼnext - σᵢⱼ)/Δt
  return σ̇ᵢⱼ
end

function set_plane_strain_oop_stress_rate_2(r,σᵢⱼ,σ̇ᵢⱼ,D,Ḋ ; abstol=1e-16, maxiter=100, damaged_allowed=true)
  # If the initial out of plane is different from σ₃₃ insert a guess
  #!isnothing(σoop_guess) && (σᵢⱼ = insert_σoop(σᵢⱼ,σoop_guess)) # Removed because of type inference issues
  σoop = σᵢⱼ[3,3]
  σoop_rate = σ̇ᵢⱼ[3,3]
  # get damage constants:
  A1,B1 = compute_A1B1(r,D)
  # use a static array version of σ̇ᵢⱼ for AD
  #σ̇ᵢⱼ_sa = SA[σ̇ᵢⱼ[1,1] σ̇ᵢⱼ[1,2] σ̇ᵢⱼ[1,3] ; σ̇ᵢⱼ[2,1] σ̇ᵢⱼ[2,2] σ̇ᵢⱼ[2,3] ; σ̇ᵢⱼ[3,1] σ̇ᵢⱼ[3,2] σ̇ᵢⱼ[3,3]]
  for i in 1:maxiter
    # Joint computation of value and grad of ϵ_oop
    dϵ̇_oop = ForwardDiff.gradient(σ̇ᵢⱼ -> compute_ϵ̇ij_2(r,D,Ḋ,σᵢⱼ,σ̇ᵢⱼ; damaged_allowed)[3,3], σ̇ᵢⱼ)
    ϵ̇_oop  = compute_ϵ̇ij_2(r,D,Ḋ,σᵢⱼ,σ̇ᵢⱼ; damaged_allowed)[3,3]
    #println(dϵ̇_oop)
    dϵ̇_oop_dσ̇_oop = dϵ̇_oop[3,3]
    # Newton update
    σoop_rate = σoop_rate - ϵ̇_oop/dϵ̇_oop_dσ̇_oop
    σ̇ᵢⱼ = insert_σoop(σ̇ᵢⱼ,σoop_rate)
    #σ̇ᵢⱼ_sa = setindex(σ̇ᵢⱼ_sa,σoop_rate,9)
    # exit condition
    (abs(ϵ̇_oop) <= abstol) && break
  end
  return σ̇ᵢⱼ
end

# function get_strain_from_stress(r,D,σᵢⱼ)
#   ϵᵢⱼ = compute_ϵij(r,D,σᵢⱼ)

# end

function get_damage_onset(r::Rheology,σ₃,D)
  func_to_minimize = S -> abs(zero_at_damage_onset(r,first(S),σ₃,D)) #optimize function needs an array of variables
  Sc = optimize(func_to_minimize, [3.0],method=LBFGS(),autodiff = :forward).minimizer[1]
  return Sc
end
get_damage_onset(r::Rheology,σ₃) = get_damage_onset(r,σ₃,r.D₀)

function zero_at_damage_onset(r,S,σ₃,D)
  σᵢⱼ = build_principal_stress_tensor(r,S,σ₃,D)
  σ = 1/3*tr(σᵢⱼ)
  sᵢⱼ = dev(σᵢⱼ)
  τ = get_τ(sᵢⱼ)
  A, B = compute_AB(r,D)
  return abs(τ/σ) - A/B
end

function get_KI_minimizer_S(r,D,σ₃)
  func_to_minimize = S -> abs(zero_at_damage_onset(r,S,σ₃,D)) #optimize function needs an array of variables
  Sc = optimize(func_to_minimize, 1, 1e4,rel_tol=1e-15).minimizer[1]
  return Sc
end

#tested
function get_KI_minimizer_D_S(r,σ₃)
  func_to_minimize = D -> -get_KI_minimizer_S(r,D,σ₃)
  Dc = optimize(func_to_minimize, r.D₀, 0.999, rel_tol=1e-15).minimizer[1]
  Sc = get_KI_minimizer_S(r,Dc,σ₃)
  return Dc, Sc
end

function get_KI_minimizer_D(r,S,σ₃)
  func_to_minimize = D -> KI_from_external_load(r,S,σ₃,D) #optimize function needs an array of variables
  Dc = optimize(func_to_minimize, r.D₀, 0.999, rel_tol=1e-10).minimizer[1]
  return Dc
end

function get_KI_minimizer_D_on_S_range(r,Smin,Smax,σ₃ ; len=1000)
  S_vec = range(Smin,Smax ; length=len)
  D_vec = range(r.D₀,0.999 ; length=len)
  σᵢⱼ_mat = [build_principal_stress_tensor(r,S,σ₃,D) for S in S_vec, D in D_vec]
  D_mat = repeat(D_vec',outer=(length(S_vec),1))
  KI_mat = compute_KI.(Ref(r),σᵢⱼ_mat,D_mat)
  Dc = 0.0
  S = 0.0
  for iS in reverse(axes(S_vec,1))
    valmin, idmin = findmin(KI_mat[iS,:])
    if valmin > 0
      continue
    else
      S = S_vec[iS]
      Dc = get_KI_minimizer_D(r,S,σ₃)
      break
    end
  end
  return Dc, S
end

function KI_from_external_load(r,S,σ₃,D)
  σij = build_principal_stress_tensor(r,S,σ₃,D)
  KI = compute_KI(r,σij,D)
  return KI
end

function build_principal_stress_tensor(r,S,σ₃ ; abstol=1e-15)
  σᵢⱼ_guess = SymmetricTensor{2,3}(SA[S*σ₃ 0 0 ; 0 σ₃ 0 ; 0 0 r.ν*(S+1)*σ₃])
  return set_plane_strain_oop_stress(σᵢⱼ_guess,r,r.D₀ ; abstol)
end

function build_principal_stress_tensor(r,S,σ₃,D ; abstol=1e-15)
  σᵢⱼ_guess = SymmetricTensor{2,3}(SA[S*σ₃ 0 0 ; 0 σ₃ 0 ; 0 0 r.ν*(S+1)*σ₃])
  return set_plane_strain_oop_stress(σᵢⱼ_guess,r,D ; abstol)
end

function build_stress_tensor(r,σxx,σyy,σxy,D ; abstol=1e-15)
  σᵢⱼ_guess = SymmetricTensor{2,3}(SA[σxx σxy 0 ; σxy σyy 0 ; 0 0 r.ν*(σxx+σyy)])
  return set_plane_strain_oop_stress(σᵢⱼ_guess,r,D ; abstol)
end

function band_coords(σᵢⱼ,θ)
  Q = Tensor{2,2}([sind(θ) cosd(θ) ; -cosd(θ) sind(θ)])
  σ2D = SymmetricTensor{2,2,eltype(σᵢⱼ)}(SA[σᵢⱼ[1,1] σᵢⱼ[1,2] ; σᵢⱼ[1,2] σᵢⱼ[2,2]])
  σ2D_rotated = Q ⋅ σ2D ⋅ Q'
  σᵢⱼ_rotated = SymmetricTensor{2,3,eltype(σᵢⱼ)}((σ2D_rotated[1,1],σ2D_rotated[2,1],σᵢⱼ[3,1],σ2D_rotated[2,2],σᵢⱼ[3,2],σᵢⱼ[3,3]))
  #σᵢⱼ_rotated = Tensor{2,3}([σ2D_rotated zeros(eltype(σᵢⱼ),2) ; [zero(eltype(σᵢⱼ)) zero(eltype(σᵢⱼ)) σᵢⱼ[3,3]]])
  return σᵢⱼ_rotated
end

function principal_coords(σᵢⱼ,θ)
  Q = Tensor{2,2}([sind(θ) -cosd(θ) ; cosd(θ) sind(θ)])
  σ2D = SymmetricTensor{2,2,eltype(σᵢⱼ)}(SA[σᵢⱼ[1,1] σᵢⱼ[1,2] ; σᵢⱼ[1,2] σᵢⱼ[2,2]])
  σ2D_rotated = Q ⋅ σ2D ⋅ Q'
  σᵢⱼ_rotated = SymmetricTensor{2,3,eltype(σᵢⱼ)}((σ2D_rotated[1,1],σ2D_rotated[2,1],σᵢⱼ[3,1],σ2D_rotated[2,2],σᵢⱼ[3,2],σᵢⱼ[3,3]))
  #σᵢⱼ_rotated = Tensor{2,3}([σ2D_rotated zeros(eltype(σᵢⱼ),2) ; [zero(eltype(σᵢⱼ)) zero(eltype(σᵢⱼ)) σᵢⱼ[3,3]]])
  return σᵢⱼ_rotated
end

function get_stress_deviation_from_far_field(σᵢⱼ ; offdiagtol=1e-5, compression_axis=:x, shear_mode=:pure, shear_strain_sign=:positive)
  σᵢⱼ2D = SymmetricTensor{2,2}(SA[σᵢⱼ[1,1] σᵢⱼ[1,2] ; σᵢⱼ[2,1] σᵢⱼ[2,2]])
  if shear_mode==:simple
    if (σᵢⱼ2D[1,1]==σᵢⱼ2D[2,2]) && (σᵢⱼ2D[1,2]==0)
      (shear_strain_sign==:positive) && (return 45.0)
      (shear_strain_sign==:negative) && (return -45.0)
    end
  end

  σᵢⱼ = filter_offdiagonal(σᵢⱼ ; tol=offdiagtol)
  F = eigen(σᵢⱼ[1:2,1:2])
  #any(isnan.([F.values F.vectors])) && (F = eigen(Array(σᵢⱼ)))
  ind = findall(==(minimum(F.values)),F.values)
  σ₁_xy_direction = length(ind) == 1 ? Vec(F.vectors[:,ind]...) : Vec(1.0,0.0)
  σ₁_far_field_direction = (compression_axis==:x) ? Vec(1.0,0.0) : Vec(0.0,1.0)
  if all(σ₁_xy_direction .≈ σ₁_far_field_direction)
    angle = 0.0
  else
    angle = atand(σ₁_xy_direction[2],σ₁_xy_direction[1]) # wrt x direction
    (compression_axis==:y) && (angle = 90 + angle)
  end
  #(angle == 90) && (angle = zero(angle))
  (angle > 90) && (angle -= 180)
  return angle
end

get_stress_deviation_from_y(σᵢⱼ ; shear_mode=:pure, shear_strain_sign = :positive) = get_stress_deviation_from_far_field(σᵢⱼ ; offdiagtol=1e-5, compression_axis=:y, shear_mode, shear_strain_sign)

include("solve.jl")
include("single_point_analysis.jl")
include("two_points_analysis.jl")
#############################
## Solve related functions ##
#############################

# function adaptative_time_integration(r::Rheology,p::Params,σᵢⱼ_i,ϵᵢⱼ_i,D_i,ϵ̇11,Δt,tspan)
#   #unpack
#   time_maxiter = p.solver.time_maxiter

#   # initializations
#   t_vec = Float64[tspan[1]]
#   tsim = tspan[1]
#   σᵢⱼ_vec = SymmetricTensor{2,3}[σᵢⱼ_i]
#   ϵᵢⱼ_vec = SymmetricTensor{2,3}[ϵᵢⱼ_i]
#   D_vec = Float64[D_i]
#   σᵢⱼnext = σᵢⱼ_i
#   ϵᵢⱼnext = ϵᵢⱼ_i
#   Dnext = D_i
#   Δt_next = Δt
#   last_tsim_printed = 0.0
#   last_tsim_saved = 0.0
#   i = 1 # iter counter
#   while tsim < tspan[2]
#     print_flag, last_tsim_printed = get_print_flag(p,i,tsim,last_tsim_printed)
#     print_flag && print_time_iteration(i,tsim)

#     σᵢⱼnext, ϵᵢⱼnext, Dnext, Δt_used, Δt_next = adaptative_Δt_solve(r,p,σᵢⱼnext,ϵᵢⱼnext,Dnext,ϵ̇11,Δt_next)

#     tsim += Δt_used

#     save_flag, last_tsim_saved = get_save_flag(p,i,tsim,last_tsim_saved)
#     if save_flag
#       push!(σᵢⱼ_vec,σᵢⱼnext)
#       push!(ϵᵢⱼ_vec,ϵᵢⱼnext)
#       push!(D_vec,Dnext)
#       push!(t_vec,tsim)
#     end

#     i += 1
#     if !isnothing(time_maxiter)
#       (length(t_vec)==time_maxiter+1) && break
#     end
#   end
#   return t_vec, σᵢⱼ_vec, ϵᵢⱼ_vec, D_vec
# end

# function adaptative_Δt_solve(r,p,σᵢⱼ_i,ϵᵢⱼ_i,D_i,ϵ̇11,Δt)
#   e₀ = p.solver.e₀

#   σᵢⱼnext1, ϵᵢⱼnext1, Dnext1, u1 = solve(r,p,σᵢⱼ_i,ϵᵢⱼ_i,D_i,ϵ̇11,Δt)
#   σᵢⱼmid, ϵᵢⱼmid, Dmid, umid = solve(r,p,σᵢⱼ_i,ϵᵢⱼ_i,D_i,ϵ̇11,Δt/2)
#   σᵢⱼnext2, ϵᵢⱼnext2, Dnext2, u2 = solve(r,p,σᵢⱼmid,ϵᵢⱼmid,Dmid,ϵ̇11,Δt/2)

#   # compute errors for each unknowns physical quantity and ponderate
#   eD = Dnext2-Dnext1
#   eσ = norm((u2-u1)[1:2])
#   eϵ = norm((u2-u1)[3])
#   if e₀ isa Real
#     e₀ref = e₀
#     e_normalized = (eD, eσ/r.G, eϵ)
#     ok_flag = all(e_normalized.<e₀)
#     e = maximum(e_normalized)
#   elseif e₀ isa NamedTuple
#     e₀ref = e₀.D
#     ok_flag = (eD<e₀.D) && (eσ<(e₀.σ)) && (eϵ<e₀.ϵ)
#     e_normalized = (eD, eσ*(e₀ref/e₀.σ), eϵ*(e₀ref/e₀.ϵ))
#     e,ind = findmax(e_normalized)
#     ok_flag || @debug("maximum error comes from indice $(ind) of (D,σ,ϵ)")
#   end

#   if ok_flag
#     # increse timestep
#     Δt_next = min(Δt*abs(e₀ref/e),Δt*2)
#     # keep best solution
#     return σᵢⱼnext2, ϵᵢⱼnext2, Dnext2, Δt, Δt_next
#   else
#     # recursively run with decreased timestep
#     adaptative_Δt_solve(r,p,σᵢⱼ_i,ϵᵢⱼ_i,D_i,ϵ̇11, Δt*abs(e₀ref/e))
#     #initialy Δt*abs(e₀ref/e)^2 but without the square seems to generaly require less iterations.
#   end

# end

# function adaptative_time_integration_2_points(r::Rheology,p::Params,S_i,σ₃,D_i,ϵ̇11,ϵ̇ⁱξη,Δt,θ,tspan ; damage_growth_out=true, bifurcate_on=:KI)
#   #unpack
#   time_maxiter = p.solver.time_maxiter

#   # fill first values
#   σᵢⱼ_i_principal = build_principal_stress_tensor(r,S_i,σ₃,D_i ; abstol=1e-15) # takes care of the plane strain constraint by solving non linear out of plane strain wrt σ₃₃ using Newton algorithm
#   ϵᵢⱼ_i_principal = compute_ϵij(r,D_i,σᵢⱼ_i_principal)
#   σᵢⱼ_i = band_coords(σᵢⱼ_i_principal,θ)
#   ϵᵢⱼ_i = band_coords(ϵᵢⱼ_i_principal,θ)

#   # initialize containers
#   t_vec = Float64[tspan[1]]
#   S_vec = Float64[S_i]
#   σⁱᵢⱼ_vec = SymmetricTensor{2,3}[σᵢⱼ_i]
#   σᵒᵢⱼ_vec = SymmetricTensor{2,3}[σᵢⱼ_i]
#   ϵⁱᵢⱼ_vec = SymmetricTensor{2,3}[ϵᵢⱼ_i]
#   Dⁱ_vec = Float64[D_i]
#   Dᵒ_vec = Float64[D_i]

#   # initialize values
#   tsim::Float64 = tspan[1]
#   Snext::Float64 = S_i
#   σⁱᵢⱼnext::typeof(σᵢⱼ_i) = σᵢⱼ_i
#   σᵒᵢⱼnext::typeof(σᵢⱼ_i) = σᵢⱼ_i
#   ϵⁱᵢⱼnext::typeof(ϵᵢⱼ_i) = ϵᵢⱼ_i
#   Dⁱnext::Float64 = D_i
#   Dᵒnext::Float64 = D_i
#   Δt_next::Float64 = Δt
#   Δt_used::Float64 = Δt

#   last_tsim_printed = 0.0
#   last_tsim_saved = 0.0
#   bifurcation_flag = false
#   i = 1 # iter counter
#   while tsim < tspan[2]
#     print_flag, last_tsim_printed = set_print_flag!(p,i,tsim,last_tsim_printed)
#     print_flag && print_time_iteration(i,tsim)

#     # solving procedure depends on the sign of S derivative
#     if bifurcation_flag
#       Snext, σⁱᵢⱼnext, σᵒᵢⱼnext, ϵⁱᵢⱼnext, Dⁱnext, Dᵒnext, Δt_used, Δt_next = adaptative_solve_2_points(r,p,Snext,σ₃,σⁱᵢⱼnext,σᵒᵢⱼnext,ϵⁱᵢⱼnext,Dᵒnext,Dⁱnext,ϵ̇ⁱξη,θ,Δt_next ; damage_growth_out)
#     else
#       σⁱᵢⱼnext_test, ϵⁱᵢⱼnext_test, Dⁱnext_test, Δt_used_test, Δt_next_test = adaptative_solve_1_point(r,p,σⁱᵢⱼnext,ϵⁱᵢⱼnext,Dⁱnext,ϵ̇11,θ,Δt)
#       σᵒᵢⱼnext_test = σⁱᵢⱼnext_test
#       Dᵒnext_test = Dⁱnext_test
#       σᵒᵢⱼnext_principal_test = principal_coords(σᵒᵢⱼnext_test,θ)
#       Snext_test = σᵒᵢⱼnext_principal_test[1,1]/σᵒᵢⱼnext_principal_test[2,2]

#       # compute bifurcation criterion
#       (bifurcate_on==:KI) && ( bifurcation_criterion = (compute_KI(r,σⁱᵢⱼnext_test,Dⁱnext_test)>0) )
#       (bifurcate_on==:S)  && ( bifurcation_criterion = ((Snext_test - Snext)<=0) )

#       if bifurcation_criterion
#         bifurcation_flag = true
#         Snext, σⁱᵢⱼnext, σᵒᵢⱼnext, ϵⁱᵢⱼnext, Dⁱnext, Dᵒnext, Δt_used, Δt_next = adaptative_solve_2_points(r,p,Snext,σ₃,σⁱᵢⱼnext,σᵒᵢⱼnext,ϵⁱᵢⱼnext,Dᵒnext,Dⁱnext,ϵ̇ⁱξη,θ,Δt_next ; damage_growth_out)
#       else
#         Snext, σⁱᵢⱼnext, σᵒᵢⱼnext, ϵⁱᵢⱼnext, Dⁱnext, Dᵒnext, Δt_used, Δt_next = Snext_test, σⁱᵢⱼnext_test, σᵒᵢⱼnext_test, ϵⁱᵢⱼnext_test, Dⁱnext_test, Dᵒnext_test, Δt_used_test, Δt_next_test
#       end

#     end

#     tsim += Δt_used

#     save_flag, last_tsim_saved = get_save_flag(p,i,tsim,last_tsim_saved)
#     if save_flag
#       push!(S_vec,Snext)
#       push!(σⁱᵢⱼ_vec,σⁱᵢⱼnext)
#       push!(σᵒᵢⱼ_vec,σᵒᵢⱼnext)
#       push!(ϵⁱᵢⱼ_vec,ϵⁱᵢⱼnext)
#       push!(Dⁱ_vec,Dⁱnext)
#       push!(Dᵒ_vec,Dᵒnext)
#       push!(t_vec,tsim)
#     end

#     i += 1

#     # break loop under conditions
#     if !isnothing(time_maxiter)
#       (length(t_vec)==time_maxiter+1) && break
#     end
#     (Dⁱnext > 0.999) && break
#     (Dᵒnext > 0.999) && break
#   end
#   return t_vec, S_vec, σⁱᵢⱼ_vec, σᵒᵢⱼ_vec, ϵⁱᵢⱼ_vec, Dⁱ_vec, Dᵒ_vec
# end

# find S_i and D_i automaticaly so that S_i is maximized with constraint KI(S_i,D_i)<=0
# and solve for in and out from the beginning of the simulation with fixed D_out = D_i
# function adaptative_time_integration_2_points(r::Rheology,p::Params,σ₃,ϵ̇ⁱξη,Δt,θ,tspan ; damage_growth_out=false)
#   #unpack
#   time_maxiter = p.solver.time_maxiter

#   D_i, S_i =  get_KI_mininizer_D_on_S_range(r,1,50,σ₃ ; len=1000)
#   # fill first values
#   σᵢⱼ_i_principal = build_principal_stress_tensor(r,S_i,σ₃,D_i ; abstol=1e-15) # takes care of the plane strain constraint by solving non linear out of plane strain wrt σ₃₃ using Newton algorithm
#   ϵᵢⱼ_i_principal = compute_ϵij(r,D_i,σᵢⱼ_i_principal)
#   σᵢⱼ_i = band_coords(σᵢⱼ_i_principal,θ)
#   ϵᵢⱼ_i = band_coords(ϵᵢⱼ_i_principal,θ)

#   # initialize containers
#   t_vec = Float64[tspan[1]]
#   S_vec = Float64[S_i]
#   σⁱᵢⱼ_vec = SymmetricTensor{2,3}[σᵢⱼ_i]
#   σᵒᵢⱼ_vec = SymmetricTensor{2,3}[σᵢⱼ_i]
#   ϵⁱᵢⱼ_vec = SymmetricTensor{2,3}[ϵᵢⱼ_i]
#   Dⁱ_vec = Float64[D_i]
#   Dᵒ_vec = Float64[D_i]

#   # initialize values
#   tsim::Float64 = tspan[1]
#   Snext::Float64 = S_i
#   σⁱᵢⱼnext::typeof(σᵢⱼ_i) = σᵢⱼ_i
#   σᵒᵢⱼnext::typeof(σᵢⱼ_i) = σᵢⱼ_i
#   ϵⁱᵢⱼnext::typeof(ϵᵢⱼ_i) = ϵᵢⱼ_i
#   Dⁱnext::Float64 = D_i
#   Dᵒnext::Float64 = D_i
#   Δt_next::Float64 = Δt
#   Δt_used::Float64 = Δt

#   last_tsim_printed = 0.0
#   last_tsim_saved = 0.0
#   i = 1 # iter counter
#   while tsim < tspan[2]
#     set_print_flag!(p,i,tsim,last_tsim_printed)
#     p.flags.print_flag && print_time_iteration(i,tsim)

#     # solving procedure depends on the sign of S derivative
#     Snext, σⁱᵢⱼnext, σᵒᵢⱼnext, ϵⁱᵢⱼnext, Dⁱnext, Dᵒnext, Δt_used, Δt_next = adaptative_solve_2_points(r,p,Snext,σ₃,σⁱᵢⱼnext,σᵒᵢⱼnext,ϵⁱᵢⱼnext,Dᵒnext,Dⁱnext,ϵ̇ⁱξη,θ,Δt_next ; damage_growth_out)

#     tsim += Δt_used # update simulation time

#     # save
#     set_save_flag!(p,i,tsim)
#     if p.flags.save_flag
#       push!(S_vec,Snext)
#       push!(σⁱᵢⱼ_vec,σⁱᵢⱼnext)
#       push!(σᵒᵢⱼ_vec,σᵒᵢⱼnext)
#       push!(ϵⁱᵢⱼ_vec,ϵⁱᵢⱼnext)
#       push!(Dⁱ_vec,Dⁱnext)
#       push!(Dᵒ_vec,Dᵒnext)
#       push!(t_vec,tsim)
#     end

#     # increment iter counter
#     i += 1

#     # break loop under conditions
#     if !isnothing(time_maxiter)
#       (length(t_vec)==time_maxiter+1) && break
#     end
#     (Dⁱnext > 0.999) && break
#   end
#   return t_vec, S_vec, σⁱᵢⱼ_vec, σᵒᵢⱼ_vec, ϵⁱᵢⱼ_vec, Dⁱ_vec, Dᵒ_vec
# end

ext: jl


text:
# This file includes an example script to plot the basis functions of PSK modulator.

using DigitalCommunications
using Plots; plotlyjs()

# Simulation parameters
k = 2
M = 2^k
nsymbols = 2
nbits = k * nsymbols
ebno = collect(0 : 10)          # Snr per bit
esno = ebno .+ 10 * log10(k)    # Snr ber symbol

# Communcation system components
gen = BitGenerator(nbits)
coding = GrayCoding(M)
modulator = Modulator(PSK(M))
channel = WaveformAWGNChannel(esno[1], basis(modulator), modulator.tsample, modulator.pulse.duration)

# Simulate the system
@run gen.bits |> coding |> modulator
tx = gen.bits |> coding |> modulator
rx = tx |> channel

# Plots
sps = Int(modulator.pulse.duration / modulator.tsample)
t = collect(Iterators.partition(0 : modulator.tsample : modulator.pulse.duration * nsymbols, sps))
plt = plot()
for (ti, txi, rxi) in zip(t, tx, rx)
    plot!(ti, txi)
    plot!(ti, rxi)
end
display(plt)


text:
#------------------------------
# some plotting functions for simulated data
#------------------------------

struct simdata
    xs
    x_baseline
    tvals
    group1
    group2
end

function plot_truesolution(group, data::simdata, sol_group1, sol_group2; showdata=true)
    if group == 1
        sol = sol_group1
        groupinds = data.group1
        legendposition = :topright
    else
        sol = sol_group2
        groupinds = data.group2
        legendposition = :topleft
    end
    curplot = plot(sol.t, sol',
                label = [L"\mathrm{true~solution~}u_1" L"\mathrm{true~solution~}u_2"],
                legend = legendposition,
                legendfontsize = 12,
                line=(3, ["#ff7f0e" "#1f77b4"])
                )
    if !showdata
        return curplot
    else
        for ind in 1:length(data.xs[groupinds])
            for var in 1:size(data.xs[groupinds][1],1)
                color = "#ffbb78"
                if var > 5
                    color = "#aec7e8"
                end
                Plots.scatter!(cat(0,data.tvals[groupinds][ind], dims=1), data.xs[groupinds][ind][var,:], label="", marker=(:c,6,color))
            end
        end
    end
    return curplot
end

function eval_z_trajectories(xs, x_params,tvals, group1, sol_group1, sol_group2, m, dt; swapcolorcoding::Bool=false) # look at trajectories of first 9 individuals during training
    plotarray=[]
    for ind in 2:5#9
        curgroup = ind ∈ group1 ? 1 : 2
        colors_truesol = swapcolorcoding ? ["#ff7f0e" "#1f77b4"] : ["#1f77b4" "#ff7f0e"]
        curxs = xs[ind]
        curmu, cursi = m.encodedμ(m.encoder(curxs)), m.encodedlogσ(m.encoder(curxs))
        curz = latentz.(curmu, cursi)
        learnedparams = m.paramNN(x_params[ind])
        if length(m.paramNN.layers[end].α) == 2
            curparams = Float32[learnedparams[1], m.ODEprob.p[2], m.ODEprob.p[3], learnedparams[2]]
        else
            curparams = learnedparams
        end
        curts = tvals[ind] .* (1 ./dt) .+ 1
        origt1s = repeat([tvals[ind]], length(curxs[:,1]))
        origxst0 = curxs[:,1]
        origxsotherts = curxs[:,2:end]
        #curprob = ODEProblem(linear_2d_system,curmu[:,1],tspan,curparams)
        cursol = solve(m.ODEprob, Tsit5(), u0 = curmu[:,1], p=curparams, saveat = dt)
        #cursol = solve(curprob, Tsit5(), saveat = dt)
        curplot = curgroup == 1 ? plot(sol_group1.t, sol_group1'; legend=false, line=([:dot :dot], 3, colors_truesol)) : plot(sol_group2.t, sol_group2'; legend=false, line=([:dot :dot], 3, colors_truesol))
        plot!(cursol.t, cursol'; line=(2, ["#1f77b4" "#ff7f0e"]))
        Plots.scatter!(cat(0,tvals[ind], dims=1), curmu[1,:], marker = (:c, 4, "#1f77b4"))
        Plots.scatter!(cat(0,tvals[ind], dims=1), curmu[2,:], marker = (:c, 4, "#ff7f0e"))
        Plots.scatter!(zeros(10,1), origxst0; marker=(:c, 3, "#bab0ac"), alpha=0.5)
        for tp in 1:length(tvals[ind])
            Plots.scatter!(repeat([tvals[ind][tp]], length(curxs[:,1])), origxsotherts[:,tp]; marker=(:c, 3, "#bab0ac"), alpha=0.5)
        end
        push!(plotarray, curplot)
    end
    myplot = plot(plotarray[:]..., layout=(2,2))#layout=(3,3))
    display(myplot)
end

function plot_individual_solutions(ind, xs, x_baseline,tvals, group1, sol_group1, sol_group2, m, dt; swapcolorcoding::Bool=false, showlegend::Bool=true) # look at trajectories of one individual after training
    curgroup = ind ∈ group1 ? 1 : 2
    colors_truesol = swapcolorcoding ? ["#ff7f0e" "#1f77b4"] : ["#1f77b4" "#ff7f0e"]
    if curgroup == 1
        sol = sol_group1
        legend = showlegend ? :topright : false
        truesollabels = ""
        smoothlabels = ""
        label1 = ""
        label2 = ""
    else
        sol = sol_group2
        legend = showlegend ? :topleft : false
        truesollabels = [L"\mathrm{true~}u_1" L"\mathrm{true~}u_2"]
        smoothlabels = [L"\mathrm{smooth~}\mu_1" L"\mathrm{smooth~}\mu_2"]
        label1 = L"\mu_1 \mathrm{~from~encoder}"
        label2 = L"\mu_2 \mathrm{~from~encoder}"
    end
    curxs = xs[ind]
    curmu, cursi = m.encodedμ(m.encoder(curxs)), m.encodedlogσ(m.encoder(curxs))
    learnedparams = m.paramNN(x_baseline[ind])
    if length(m.paramNN.layers[end].α) == 2
        curparams = Float32[learnedparams[1], m.ODEprob.p[2], m.ODEprob.p[3], learnedparams[2]]
    else
        curparams = learnedparams
    end
    cursol = solve(m.ODEprob, Tsit5(), u0 = curmu[:,1], p=curparams, saveat = dt)
    curplot = plot(sol.t, sol'; labels=truesollabels, legend=legend, line=([:dot :dot], 3, colors_truesol))
    plot!(cursol.t, cursol'; labels=smoothlabels, line=(2, ["#1f77b4" "#ff7f0e"]))
    Plots.scatter!(cat(0,tvals[ind], dims=1), curmu[1,:], label = label1, marker = (:c, 6, "#1f77b4"))
    Plots.scatter!(cat(0,tvals[ind], dims=1), curmu[2,:], label = label2, marker = (:c, 6, "#ff7f0e"))
    plot!(xlab="time", ylab="value of latent representation")
    #display(curplot)

    return curplot

end

function allindsplot(group, data::simdata, m, sol_group1, sol_group2; swapcolorcoding::Bool=false, showlegend::Bool=true)
    # get data
    xs, x_params, tvals = data.xs, data.x_baseline, data.tvals
    # set parameters
    if group == 1
        sol = sol_group1
        groupinds = group1
        legendposition = showlegend ? :topright : false
        ylims = (-0.3,4.2)
    else
        sol = sol_group2
        groupinds = group2
        legendposition = showlegend ? :topleft : false
        ylims = (-0.5,12)
    end

    # plot true solution
    colors_truesol = swapcolorcoding ? ["#ff7f0e" "#1f77b4"] : ["#1f77b4" "#ff7f0e"]
    allindsplot1 = plot(sol.t, sol';
                    #ylims = ylims,
                    labels=[L"\mathrm{true~}u_1" L"\mathrm{true~}u_2"],
                    legend=legendposition,
                    line=([:dot :dot], 4, colors_truesol)#["#e6550d" "#3182bd"])
    )
    # get data from group currently looked at
    groupxs = xs[groupinds]
    groupx_params = x_params[groupinds]
    grouptvals = tvals[groupinds]
    # plot ODE solutions (= smooth latent μs as function of t)
    for ind in 1:length(xs[groupinds])
        curxs = groupxs[ind]
        curmu, cursi = m.encodedμ(m.encoder(curxs)), m.encodedlogσ(m.encoder(curxs))
        learnedparams = m.paramNN(groupx_params[ind])
        if length(m.paramNN.layers[end].α) == 2
            curparams = Float32[learnedparams[1], m.ODEprob.p[2], m.ODEprob.p[3], learnedparams[2]]
        else
            curparams = learnedparams
        end
        cursol = solve(m.ODEprob, Tsit5(), u0 = curmu[:,1], p=curparams, saveat = dt)
        if ind == 2
            labels = [L"\mathrm{smooth~}\mu_1" L"\mathrm{smooth~}\mu_2"]
            label1 = L"\mu_1 \mathrm{~from~encoder}"
            label2 = L"\mu_2 \mathrm{~from~encoder}"
        else
            labels = ""
            label1 = ""
            label2 = ""
        end
        plot!(cursol.t, cursol'; label=labels, line=(0.5, ["#3182bd" "#e6550d"])) #["#6baed6" "#fd8d3c"]
    end
    # plot latent μs obtained directly from the encoder before the ODE solving step
    for ind in 1:length(xs[groupinds])
        curxs = groupxs[ind]
        curmu, cursi = m.encodedμ(m.encoder(curxs)), m.encodedlogσ(m.encoder(curxs))
        if ind == 2
            labels = [L"\mathrm{smooth~}\mu_1" L"\mathrm{smooth~}\mu_2"]
            label1 = L"\mu_1 \mathrm{~from~encoder}"
            label2 = L"\mu_2 \mathrm{~from~encoder}"
        else
            labels = ""
            label1 = ""
            label2 = ""
        end
        Plots.scatter!(cat(0,grouptvals[ind], dims=1),curmu[1,:]; label=label1, marker=(:c,4,"#9ecae1"))
        Plots.scatter!(cat(0,grouptvals[ind], dims=1),curmu[2,:]; label=label2, marker=(:c,4,"#fdae6b"))
    end
    plot!(xlab="time", ylab="value of latent representation")

    return allindsplot1

end

function plot_batch_solution(ind, xs, x_baseline, tvals, group1, sol_group1, sol_group2, m, dt; swapcolorcoding::Bool=false, showlegend::Bool=true)
    distmat = getdistmat_odesols_mean(m, xs, x_baseline; centralise=true);
    minibatches, batch_weights = findminibatches_distmat(distmat, batchsize, kernel);
    batch_xs = collect(xs[minibatches[i]] for i in 1:length(xs));
    batch_x_baseline = collect(x_baseline[minibatches[i]] for i in 1:length(x_baseline));
    batch_tvals = collect(tvals[minibatches[i]] for i in 1:length(tvals));

    curgroup = ind ∈ group1 ? 1 : 2
    colors_truesol = swapcolorcoding ? ["#fd8d3c" "#6baed6"] : ["#6baed6" "#fd8d3c"]
    curxs = batch_xs[ind][1] # reference individual
    curmu, cursi = m.encodedμ(m.encoder(curxs)), m.encodedlogσ(m.encoder(curxs))
    curparams = m.paramNN(batch_x_baseline[ind][1])
    cursol = solve(m.ODEprob, Tsit5(), u0 = curmu[:,1], p = curparams, saveat = dt)
    if curgroup == 1
        curplot = plot(sol_group1.t, sol_group1';
                        labels=[L"\mathrm{true~}u_1" L"\mathrm{true~}u_2"],
                        legend=showlegend ? :topright : false,
                        line=([:dot, :dot], 3, colors_truesol)#orange, blue
                        #size=(1000,500)
                        )
    else
        curplot = plot(sol_group2.t, sol_group2';
                        labels=[L"\mathrm{true~}u_1" L"\mathrm{true~}u_2"],
                        legend=showlegend ? :topleft : false,
                        line=([:dot, :dot], 3, colors_truesol),
                        #size=(1000,500)
                        )
    end
    for i in 2:batchsize
        bcurxs = batch_xs[ind][i]
        bcurmu, bcursi = m.encodedμ(m.encoder(bcurxs)), m.encodedlogσ(m.encoder(bcurxs))
        bcurparams = m.paramNN(batch_x_baseline[ind][i])
        bcursol = solve(m.ODEprob, Tsit5(), u0 = bcurmu[:,1], p = bcurparams, saveat = dt)
        if i==batchsize
            labels = [L"\mathrm{smooth~}\mu_1\mathrm{~batch~individual}" L"\mathrm{smooth~}\mu_2\mathrm{~batch~individual}"]
            label1 = L"\mu_1\mathrm{~from~encoder,~batch~individual}"
            label2 = L"\mu_2\mathrm{~from~encoder,~batch~individual}"
        else
            labels, label1, label2="", "", ""
        end
        plot!(bcursol.t, bcursol'; labels=labels, line=(0.5, ["#9ecae1" "#fdae6b"])) # blue, orange
        scatter!(cat(0,batch_tvals[ind][i], dims=1), bcurmu[1,:]; label=label1, marker=(:c,4,"#9ecae1"))
        scatter!(cat(0,batch_tvals[ind][i], dims=1), bcurmu[2,:]; label=label2, marker=(:c,4,"#fdae6b"))
    end
    plot!(cursol.t, cursol';
        labels=[L"\mathrm{smooth~}\mu_1\mathrm{~reference~individual}" L"\mathrm{smooth~}\mu_2\mathrm{~reference~individual}"],
        line=(2.5, ["#1f77b4" "#e6550d"])
        )
    scatter!(cat(0,batch_tvals[ind][1], dims=1),curmu[1,:];
        label = L"\mu_1\mathrm{~from~encoder,~reference~individual}", marker=(:c,7,"#1f77b4"))
    scatter!(cat(0,batch_tvals[ind][1], dims=1),curmu[2,:];
        label = L"\mu_2\mathrm{~from~encoder,~reference~individual}", marker=(:c,7,"#e6550d"))
    plot!(xlab="time", ylab="value of latent representation")
    display(curplot)
    return curplot
end


text:
# 2 or more args graph operators
function Base.intersect(g1::GNNGraph, g2::GNNGraph)
    @assert g1.num_nodes == g2.num_nodes
    @assert graph_type_symbol(g1) == graph_type_symbol(g2)
    graph_type = graph_type_symbol(g1)
    num_nodes = g1.num_nodes

    idx1, _ = edge_encoding(edge_index(g1)..., num_nodes)
    idx2, _ = edge_encoding(edge_index(g2)..., num_nodes)
    idx = intersect(idx1, idx2)
    s, t = edge_decoding(idx, num_nodes)
    return GNNGraph(s, t; num_nodes, graph_type)
end


text:
struct Conditional{S <: OneOrMoreSymbols, F}
  name::S
  stepper::F
end

struct Gibbs{M<:Model, T<:Tuple} <: InferenceAlgorithm
  model::M
  algs::T
end
Gibbs(m, algs...) = Gibbs(m, algs)

function update(gibbs::Gibbs, state::S; kwargs...) where S
  # TODO: Allow user to
  # - time variable updates.
  # - fix parameters.
  for s in gibbs.algs
    if s.name isa NTuple
      newvalues = s.stepper(gibbs.model, state)
      state = setproperties!!(state, newvalues)
    else
      newvalue = s.stepper(gibbs.model, state)
      state = setproperty!!(state, s.name, newvalue)
    end
  end
  return state
end


text:
__precompile__()
module MakePkgUpdatePrecompileInTheBackground

export start_cycling, stop_cycling, report_progress, supplant_repl_backend, progress_bar, recompile_everything

# We use these subsystems.  A lot.
using Base.REPL, Base.Terminals, Base.LineEdit

# Way better than the TV show.
const CSI = Terminals.CSI

# Global variables help keep the code quality high
btask_report = Channel(10000)
prompt_state = nothing

# The last time a newline caused a new prompt to be written out, signifying that
# we can reset our memory of how many lines we've written out
last_enter = 0
last_draw_enter = 0
last_draw_lines = 0

# This gets locked and unlocked to control ownership of the screen so that we
# don't have two tasks overwriting eachother
running_mutex = Channel()

include("util.jl")
include("pkgupdate.jl")
include("btask.jl")
include("backend.jl")

function __init__()
    # Create timer to continually draw background task progress updates
    btask_draw_timer = Timer(timer -> btask_draw(), 0, .05)

    # Try to auto-supplant the default REPL backend
    supplant_repl_backend(attempt=1)

    # Try to auto-supplant Pkg.update()
    supplant_pkgupdate()
end

Base.precompile(recompile_everything, ())

end # module


text:
module SVG

export Inline, Class, Svg, Polyline, Circle, Line, write, bounds, bounds_viewbox, scaled, points2circles

import Base.==

dp(n; digits=2) = round(n; digits)

abstract type SvgObject end

abstract type Style end

"""
	Inline
Inline style object
# Properties
- `attribs` Dict of CSS attributes
# Constructor
	Inline(;fill="none", stroke="black", stroke_width=1)
All keyword arguments
"""
struct Inline <: Style
	attribs
	function Inline(pairs=[];fill="none", stroke="black", stroke_width=1)
		attribs = Dict()
		attribs["fill"] = fill
		attribs["stroke"] = stroke
		attribs["stroke-width"] = stroke_width
		for (k,v) in pairs
			attribs[k] = v
		end
		new(attribs)
	end
end

"""
	Class
CSS class name Object
# Constructor
	Class(class::String)
"""
struct Class <: Style
	class::String
end

"""
	Polyline
Polyline object
# Properties
- `xs` Vector of X coordinates, one per vertex
- `ys` Vector of Y coordinates, one per vertex
# Constructors
Quite the choice.
- Construct with separate vectors or vectors of 2d points.
- Supply an optional transformation function `fx`, `fy`, `fxy`
- Or empty or with pre-allocated vectors of `xs`/`ys`
- Supply `Style` using kwarg
```
Polyline(; style::Style=Inline())
Polyline(npoints::Int; style::Style=Inline())
Polyline(xs::Vector, ys::Vector; style::Style=Inline())
Polyline(xs::Vector, ys::Vector, fx::Function, fy::Function; style::Style=Inline())
Polyline(xys::Vector; style::Style=Inline())
Polyline(xys::Vector, fxy::Function; style::Style=Inline())
Polyline(xys::Vector, fx::Function, fy::Function; style::Style=Inline())
```
"""
struct Polyline <: SvgObject
	xs
	ys
	style::Style
	Polyline(; style::Style=Inline()) = new(Vector{Float64}(), Vector{Float64}(), style)
	Polyline(npoints::Int; style::Style=Inline()) = new(Vector{Float64}(undef, npoints), Vector{Float64}(undef, npoints), style)
	Polyline(xs::Vector, ys::Vector; style::Style=Inline()) = new(xs, ys, style)
	Polyline(xs::Vector, ys::Vector, fx::Function, fy::Function; style::Style=Inline()) = Polyline(map(fx, xs), map(fy, ys); style)
	Polyline(xys::Vector; style::Style=Inline()) = Polyline(map(xy->xy[1], xys), map(xy->xy[2], xys); style)
	Polyline(xys::Vector, fxy::Function; style::Style=Inline()) = Polyline(map(xy->fxy(xy), xys); style)
	Polyline(xys::Vector, fx::Function, fy::Function; style::Style=Inline()) = Polyline(map(xy->fx(xy[1]), xys), map(xy->fy(xy[2]), xys); style)
end

==(p1::Polyline, p2::Polyline) = p1.xs == p2.xs && p1.ys == p2.ys && p1.style.attribs == p2.style.attribs

scaled(p::Polyline, fx, fy) = Polyline(p.xs, p.ys, fx, fy; style=p.style)

"""
	Line
# Properties
- `x1`, `y1` `x2`, `y2` - Endpoints of the line
- `style` - Style object for this line
# Constructor
	Line(x1, y1, x2, y2; style::Style=Inline())
"""
struct Line <: SvgObject
	x1
	y1
	x2
	y2
	style::Style
	Line(x1, y1, x2, y2; style::Style=Inline()) = new(x1, y1, x2, y2, style)
end

scaled(l::Line, fx, fy) = Line(fx(l.x1), fy(l.y1), fx(l.x2), fy(l.y2); style=l.style)

"""
	Circle
Circle object
# Properties
- `x`, `y`, `r` - position and radius
- `style` = Style object
Radius is invariant under the module supplied scaling operations
"""
struct Circle <: SvgObject
	x
	y
	r
	style::Style
	Circle(x, y, r; style::Style=Inline()) = new(x, y, r, style)
end

scaled(c::Circle, fx, fy) = Circle(fx(c.x), fy(c.y), c.r; style=c.style)

"""
	Svg
Svg object
# Properties
- `objects` - Vector of SVG objects
# Constructors
Call either empty or with Vector of existing objects
- `Svg()`
- `Svg(objs::SvgObject)`
"""
struct Svg
	objects::Vector{SvgObject}
	Svg() = new(Vector{SvgObject}())
	Svg(objs) = new(objs)
end
"""
	scaled(s::Svg, w, h; digits, flip_y)
	scaled(s::Svg, fx, fy)
	scaled(s::SvgObject, fx, fy)
Scale either the whole Svg or individual Objects
- `s` either the whole Svg or individual objects
- `fx`, `fy` - functions to scale the positional co-ordinates
- `w`, `h` - scale the entire Svg up to the smallest of `w` / `h` and maintain aspect ratio
- `digits::Int` - round to this many digits
- `flip_y::Bool` - flip vertically
"""
scaled(s::Svg, fx::Function, fy::Function) = Svg(map(o->scaled(o, fx, fy), s.objects))
function scaled(svg, width::Real, height::Real; flip_y=false, digits=2)
    xmin, ymin, xmax, ymax = bounds(svg)
    xmx = xmax - xmin
    ymx = ymax - ymin
    scale = min(width, height) / min(xmx, ymx)
    fx = x -> round(scale * (x - xmin); digits)
    fy = y -> round(flip_y ? height - scale * (y - ymin) : scale * (y - ymin); digits)
	scaled(svg, fx, fy)
end

"""
	bounds(svg::Svg)
	bounds(svg::SvgObject)
Return the bounding box of the entire `Svg` or an individual `SvgObject`
# Returns
NamedTuple(xmin, ymin, xmax, ymax)
"""
function bounds(svg::Svg)
	xmin, ymin, xmax, ymax = Inf, Inf, -Inf, -Inf
	for b in map(bounds, svg.objects)
		if b.xmin < xmin
			xmin = b.xmin
		end
		if b.ymin < ymin
			ymin = b.ymin
		end
		if b.xmax > xmax
			xmax = b.xmax
		end
		if b.ymax > ymax
			ymax = b.ymax
		end
	end
	(;xmin, ymin, xmax, ymax)
end

"""
	bounds_viewbox(svg::Svg)
	bounds_viewbox(bounds::NamedTuple)
Create a viewbox string using either the Svg or a set of bounds
"""
bounds_viewbox(svg::Svg) = bounds_viewbox(bounds(svg))
bounds_viewbox(bounds::NamedTuple) = "$(bounds.xmin) $(bounds.ymin) $(bounds.xmax) $(bounds.ymax)"

"""
	write(fn::String, svg::Svg, width, height; viewbox="", inhtml=false, digits=2, objwrite_fn=write_objs)
	write(io::IO, svg::Svg, width, height; viewbox="", inhtml=false, digits=2, objwrite_fn=write_objs)
write the Svg to a file or IO object.
# Arguments
- `fn` - write to filename
- `io` - write to IO Object
- `svg` - the Svg to write
- `width`, `height` - the width / height of the SVG object to supply
- `viewbox` - specify a particular viewbox: "xmin ymin xmax ymax"
- `inhtml` - Boolean flag whether to wrap the SVG in an HTML document for viewing in a browser.
- `digits` - round all SVG Object co-ordinates to this many digits
- `objwrite_fn` - a callback so the caller can modify the objects as they are written
"""
Base.write(fn::String, svg::Svg, width, height; viewbox="", inhtml=false, stylesheet="", digits=2, objwrite_fn=write_objs) = open(fn, "w+") do io write(io, svg, width, height; viewbox, inhtml, stylesheet, digits, objwrite_fn) end

write_objs(io::IO, svg::Svg) = foreach(o->write(io, o), svg.objects)

function Base.write(io::IO, svg::Svg, width, height; viewbox="", inhtml=false, stylesheet="", digits=2, objwrite_fn=write_objs)
	open_html(io; inhtml)
	open_svg(io, width, height; viewbox, stylesheet)
	objwrite_fn(io, svg)
	close_svg(io)
	close_html(io; inhtml)
end

function open_svg(io, width, height; viewbox="", stylesheet="")
	if viewbox != ""
		viewbox = " viewBox=\"$viewbox\""
	end
	println(io, """<svg width="$width" height="$height"$viewbox xmlns="http://www.w3.org/2000/svg">""")
	if stylesheet != ""
		println(io, "")
	end
end
close_svg(io) = println(io, "</svg>")
open_html(io; inhtml) = if inhtml println(io, "
") end
close_html(io; inhtml) = if inhtml println(io, "
") end

function Base.write(io::IO, s::Inline)
	print(io, "style=\"")
	for (k,v) in s.attribs
		print(io, "$k:$v; ")
	end
	print(io, "\"")
end
Base.write(io::IO, s::Class) = print(io, "class=\"$(s.class)\"")

function Base.write(io::IO, p::Polyline)
	print(io, "<polyline points=\"")
	foreach(i->print(io, p.xs[i], ",", p.ys[i], " "), 1:length(p.xs))
	print(io, "\" ")
	write(io, p.style)
	println(io, " />")
end

bounds(p::Polyline) = (xmin=minimum(p.xs), ymin=minimum(p.ys), xmax=maximum(p.xs), ymax=maximum(p.ys))

function Base.write(io::IO, L::Line)
	print(io, "<line x1=\"", L.x1, "\" y1=\"", L.y1, "\" x2=\"", L.x2, "\" y2=\"", L.y2, "\" ")
	write(io, L.style)
	println(io, " />")
end

bounds(l::Line) = (xmin=min(l.x1,l.x2), ymin=min(l.y1,l.y2), xmax=max(l.x1,l.x2), ymax=max(l.y1,l.y2))

function Base.write(io::IO, c::Circle)
	print(io, "<circle cx=\"", c.x, "\" cy=\"", c.y, "\" r=\"", c.r, "\" ")
	write(io, c.style)
	println(io, " />")
end

"""
	points2circles(xs, ys, r::Real)
	points2circles(xys, r::Real)
	points2circles(xs, ys, rs::Vector{Real})
	points2circles(xys, rs::Vector{Real})
Helper function to create a collection of points of a given radius to a collection of Circle objects
# Arguments
- `xs`, `ys` Vectors or x/y corrdinates
- `xys` Vector of 2D points
- `r` Radius
- `rs` Vector of Radii
"""
points2circles(xs, ys, r::Real=1) = map(i->Circle(xs[i], ys[i], r), 1:length(xs))
points2circles(xs, ys, rs::Vector{Real}) = map(i->Circle(xs[i], ys[i], r[i]), 1:length(xs))
points2circles(xys, r::Real=1) = map(i->Circle(xys[i][1], xys[i][2], r), 1:length(xys))
points2circles(xys, rs::Vector{Real}) = map(i->Circle(xys[i][1], xys[i][2], r[i]), 1:length(xys))

bounds(c::Circle) = (xmin=c.x-c.r, ymin=c.y-c.r, xmax=c.x+c.r, ymax=c.y+c.r)

###
end


text:
@testset "Vocabulary" begin
    multi(x, n) = collect(Iterators.repeated(x, n))

    v = Vocabulary(["a", "b", "c", "d", "e"], "?")

    before_enc = ["a", "a", "c" ,"d", "e","x", "sad"]
    before_unk_enc = ["a", "a", "c" ,"d", "e","?", "?"]
    after_enc = [2,2,4,5,6,1,1]
    multi_after_enc = hcat(multi(after_enc, 5)...)

    @test length(v) == 6

    @test v("a") == 2
    @test v[2] == "a"
    @test v(before_enc) == after_enc
    @test v[after_enc] == before_unk_enc
    @test v(multi(before_enc, 5)) == multi_after_enc
    @test v[multi_after_enc] == multi(before_unk_enc, 5)
    @test v[multi(after_enc, 5)] == multi(before_unk_enc, 5)

    @test encode(v, "a") == 2
    @test decode(v, 2) == "a"
    @test encode(v, before_enc) == after_enc
    @test decode(v, after_enc) == before_unk_enc
    @test encode(v, multi(before_enc, 5)) == multi_after_enc
    @test decode(v, multi_after_enc) == multi(before_unk_enc, 5)
    @test decode(v, multi(after_enc, 5)) == multi(before_unk_enc, 5)
end


text:
using Test
using OpenTelemetryAPI
using OpenTelemetrySDK
using OpenTelemetryProto
using OpenTelemetryExporterOtlpProtoGrpc

exporter = InMemoryExporter()
# exporter = OtlpProtoGrpcTraceExporter()

p = TracerProvider(span_processor = CompositSpanProcessor(SimpleSpanProcessor(exporter)))

tracer = Tracer(provider = p)

@testset "OtlpProtoGrpc" begin
    with_span(Span("foo", tracer)) do
        with_span(Span("bar", tracer)) do
            with_span(Span("baz", tracer)) do
                println("Hello world!")
            end
        end
    end

    req = convert(
        OpenTelemetryProto.OpentelemetryClients.ExportTraceServiceRequest,
        exporter.pool,
    )
end


text:
struct Risk
	λ::Float64
	α::Float64
	offset::Union{Dict{Leaf,Float64},Nothing}
	bound::Union{Float64,Nothing}
	penalty::Float64
end

"""
	RiskNeutral()

Create a risk-neutral risk measure.
"""
function RiskNeutral()
	Risk(1.0,1.0)
end

"""
	Risk(λ::Float64,
         α::Float64;
         offset::Union{Dict{Leaf,Float64},Nothing}=nothing,
         bound::Union{Float64,Nothing}=nothing,
         penalty::Float64=100000.0)

Define the CVaR risk measure to be applied to the accumulated profits at the leaf nodes.

### Required Arguments
`λ` is weighting applied for the risk measure (max sum of weightings should be 1.0),
if sum of weightings is less than 1.0, expected value will make up the rest.

`α` is the probability in the tail of the distribution

### Optional Arguments
`offset` applies a negative offset to each leaf node. This can be used to reorder
the outcomes prior to applying the risk measure.

`bound` if used, this will create a constraint on CVaR(α) with this as the upper bound.

`penalty` if a constraint on CVaR is applied, then the marginal cost of violating the
constraint is `penalty`.
"""
function Risk(λ::Float64, α::Float64; offset::Union{Dict{Leaf,Float64},Nothing}=nothing, bound::Union{Float64,Nothing}=nothing, penalty::Float64=100000.0)
	if α<=0.0 || α>1.0
		error("α must be >0 and <=1")
	elseif λ<0.0 || λ>1.0
		error("λ must be >=0 and <=1")
	end
	Risk(λ,α,offset,bound,penalty)
end

"""
	Risk(λ::Float64,
         α::Float64;
         offset::Union{Dict{Leaf,Float64},Nothing}=nothing,
         bound::Union{Float64,Nothing}=nothing,
         penalty::Float64=100000.0)

Define the CVaR risk constraint to be applied to the accumulated profits at the leaf nodes.

### Required Arguments
`α` is the probability in the tail of the distribution

### Optional Arguments
`offset` applies a negative offset to each leaf node. This can be used to reorder
the outcomes prior to applying the risk measure.

`bound` if used, this will create a constraint on CVaR(α) with this as the upper bound.

`penalty` if a constraint on CVaR is applied, then the marginal cost of violating the
constraint is `penalty`.
"""
function Risk(α::Float64; offset::Union{Dict{Leaf,Float64},Nothing}=nothing, bound::Union{Float64,Nothing}=nothing, penalty::Float64=100000.0)
	if α<=0.0 || α>1.0
		error("α must be >0 and <=1")
	end
	if bound==nothing
		error("Risk constraint cannot be created without setting a bound")
	end
	Risk(0.0,α,offset,bound,penalty)
end


text:
using Pkg
 Pkg.activate(".")

cd("RvSpectML")
Pkg.activate("examples")
verbose = true
  make_plots = true
  if verbose && !isdefined(Main,:RvSpectML)  println("# Loading RvSpecML")    end
  using RvSpectML
  if verbose && make_plots && !isdefined(Main,:RvSpectMLPlots)  println("# Loading RvSpecMLPlots")    end
  if make_plots
     using RvSpectMLPlots
  end
  if verbose   println("# Loading other packages")    end
  using Statistics

all_spectra = include(joinpath(pkgdir(EchelleInstruments),"examples/read_expres_data_101501.jl"))
#all_spectra = include(joinpath(pkgdir(EchelleInstruments),"examples/read_neid_solar_data_20190918.jl"))
if typeof(get_inst(all_spectra)) <: AnyEXPRES   continuum_normalize_spectra!(all_spectra)   end

order_list_timeseries = extract_orders(all_spectra,pipeline_plan, orders_to_use=12:85, recalc=true )

function calc_depth_and_expected_rv_precission(spectrum::ST, pixels::AR, order::Integer) where { ST<:AbstractSpectra2D, AR<:AbstractRange{Int64}, AA1<:AbstractArray{Int64,1} }
  logλ = log.(view(spectrum.λ,pixels,order))
  flux = view(spectrum.flux,pixels,order)
  var = view(spectrum.var,pixels,order)
  (f_mean, f_deriv) = RvSpectML.TemporalGPInterpolation.predict_mean_and_deriv(logλ, flux, logλ;sigmasq_obs=var, use_logx=false, use_logy=false, smooth_factor=2)
  depth = 1.0 .- minimum(f_mean)/maximum(f_mean)
  exp_sigma_rv = RvSpectMLBase.speed_of_light_mps / sqrt( sum(f_deriv.^2 ./ var) )
  return (depth=depth, exp_σ_rv=exp_sigma_rv)
end

#order_info = get_order_info(all_spectra, orders_to_use=1:85)

#if !make_plots     make_no_plots!(pipeline_plan)    end
line_list_df = prepare_line_list(linelist_for_ccf_filename, all_spectra, pipeline_plan,  orders_to_use=38:76, v_center_to_avoid_tellurics=ccf_mid_velocity, Δv_to_avoid_tellurics = RvSpectMLBase.max_bc, recalc=true)
#line_list_orig = deepcopy(line_list_df)

#line_list_merged = RvSpectML.assign_lines_to_orders(line_list_df, order_info)
#line_list_reweighted = RvSpectML.calc_snr_weights_for_lines!(deepcopy(line_list_merged), all_spectra)
#result = calc_weights_for_lines(copy(line_list_merged), all_spectra )

# range_no_mask_change=RvSpectMLBase.max_bc, step=v_step,  max=RvSpectMLBase.max_bc

(ccfs, v_grid) = ccf_total(order_list_timeseries, line_list_df, pipeline_plan,  mask_type=:tophat, mask_scale_factor=7.0,
  v_step=400, ccf_mid_velocity=ccf_mid_velocity, recalc=true, calc_ccf_var=false)
rvs_ccf = calc_rvs_from_ccf_total(ccfs, pipeline_plan, v_grid=v_grid, times = order_list_timeseries.times, recalc=true).-ccf_mid_velocity

#alg_fit_rv = EchelleCCFs.MeasureRvFromCCFQuadratic()
#alg_fit_rv = EchelleCCFs.MeasureRvFromCCFTemplate(v_grid=v_grid2, template=vec(mean(ccfs2,dims=2)), frac_of_width_to_fit=1.0, measure_width_at_frac_depth=0.5)
alg_fit_rv = EchelleCCFs.MeasureRvFromCCFGaussian(frac_of_width_to_fit=1.0, measure_width_at_frac_depth=0.5, init_guess_ccf_σ=ccf_mid_velocity)
((ccfs2, ccf_vars2), v_grid2) = ccf_total(order_list_timeseries, line_list_df, pipeline_plan,  mask_type=:tophat, mask_scale_factor=7.0,
      ccf_mid_velocity=ccf_mid_velocity, recalc=true, calc_ccf_var=true)
rvs_ccf2 = calc_rvs_from_ccf_total(ccfs2, ccf_vars2, pipeline_plan, v_grid=v_grid2, times = order_list_timeseries.times, alg_fit_rv=alg_fit_rv, recalc=true)

using Plots
plot(v_grid,ccfs./maximum(ccfs,dims=1),color=:blue,label=:none)
 plot!(v_grid2,ccfs2./maximum(ccfs2,dims=1),markersize=1.2, color=:red,label=:none)

x_plt = order_list_timeseries.times # mod.(order_list_timeseries.times,365.25)
scatter(x_plt,rvs_ccf.-mean(rvs_ccf), color=:blue); scatter!(x_plt,rvs_ccf2.-mean(rvs_ccf2),color=:red)

(ccfs2, v_grid2) = ccf_total(order_list_timeseries, line_list_reweighted, pipeline_plan,  mask_type=:tophat, mask_scale_factor=2.0,
  ccf_mid_velocity=ccf_mid_velocity, recalc=true, calc_ccf_var=false)

((ccfs2, ccf_vars2), v_grid2) = ccf_total(order_list_timeseries, line_list_df, pipeline_plan,  mask_type=:halfcos, mask_scale_factor=11.0,
    ccf_mid_velocity=ccf_mid_velocity, recalc=true, calc_ccf_var=true)

rvs_ccf2 = calc_rvs_from_ccf_total(ccfs2, ccf_vars2, pipeline_plan, v_grid=v_grid2, times = order_list_timeseries.times, alg_fit_rv=alg_fit_rv, recalc=true)

#scatter(line_list_reweighted[:depth]./line_list_reweighted[:weight_snr],line_list_reweighted[:weight], markersize=1.2)
scatter(line_list_reweighted[!,:depth],line_list_reweighted[!,:weight], markersize=1.2)
scatter(1.0 ./line_list_reweighted[!,:weight_snr],line_list_reweighted[!,:weight], markersize=1.2)
scatter(line_list_reweighted[!,:depth].*line_list_reweighted[!,:weight_snr],line_list_reweighted[!,:weight], markersize=1.2)
scatter(line_list_reweighted[!,:weight_snr] ./line_list_reweighted[!,:exp_sigma_rv].^2,line_list_reweighted[!,:weight], markersize=1.2)

#snr_of_spectra = vec(mean(result[:,:depth],dims=2))
#plot((result./snr_of_spectra)', label=:none)

plot(v_grid2,ccfs2./maximum(ccfs2,dims=2).-mean(ccfs2./maximum(ccfs2,dims=2),dims=2),markersize=1.2, color=:blue,label=:none)

ref_obs_idx = RvSpectMLBase.choose_obs_idx_for_init_guess(df_files_use,first(all_spectra).inst)

line_width_05 = RvSpectMLBase.calc_line_width(v_grid,view(ccfs,:,ref_obs_idx),frac_depth=0.05)

line_list_df = prepare_line_list(linelist_for_ccf_filename, all_spectra, pipeline_plan,  v_center_to_avoid_tellurics=ccf_mid_velocity, Δv_to_avoid_tellurics = line_width_05)

((ccfs, ccf_vars), v_grid) = ccf_total(order_list_timeseries, line_list_df, pipeline_plan,  mask_type=:tophat, mask_scale_factor=2.0, ccf_mid_velocity=ccf_mid_velocity, recalc=true,
      calc_ccf_var=true)
# Code working towards propagating uncertainties through CCF
#((ccfs2, ccf_vars2), v_grid2) = ccf_total(order_list_timeseries, line_list_df, pipeline_plan,  mask_type=:gaussian, mask_scale_factor=10.0, ccf_mid_velocity=ccf_mid_velocity, recalc=true,  calc_ccf_var=true)

rvs_ccf = calc_rvs_from_ccf_total(ccfs, pipeline_plan, v_grid=v_grid, times = order_list_timeseries.times, recalc=true)

if make_plot(pipeline_plan, :ccf_total)
   using Plots
   plt = RvSpectMLPlots.make_plot_ccf_vs_time(ccfs,v_grid = v_grid, output_path=output_dir, save_fig = save_plot(pipeline_plan,:ccf_total) )
   #plt = RvSpectMLPlots.make_plot_ccf_vs_time(ccfs2,v_grid = v_grid2, output_path=output_dir, save_fig = save_plot(pipeline_plan,:ccf_total) )
   display(plt)
end

if make_plot(pipeline_plan, :ccf_total)
   plt = RvSpectMLPlots.make_heatmap_ccf_vs_time(ccfs,v_grid = v_grid, times=order_list_timeseries.times, output_path=output_dir, save_fig = save_plot(pipeline_plan,:ccf_total) )
   #plt = RvSpectMLPlots.make_heatmap_ccf_vs_time(ccfs2,v_grid = v_grid2, times=order_list_timeseries.times, output_path=output_dir, save_fig = save_plot(pipeline_plan,:ccf_total) )
   display(plt)
end

(order_ccfs, v_grid_order_ccfs_novars) = ccf_orders(order_list_timeseries, line_list_df, pipeline_plan, orders_to_use=1:85, ccf_mid_velocity=ccf_mid_velocity, recalc=true)

plot(order_ccfs_novars[:,:,1],label=:none)
scatter!(order_ccfs_novars[:,5,1],markersize=1.2,label=:none)

mask_shape = EchelleCCFs.TopHatCCFMask(order_list_timeseries.inst) #, scale_factor=8)
#line_list = EchelleCCFs.BasicLineList(line_list_df.lambda, line_list_df.weight)
line_list = EchelleCCFs.BasicLineList2D(line_list_df)
v_step = 400
ccf_plan = EchelleCCFs.BasicCCFPlan(mask_shape = mask_shape, line_list=line_list, midpoint=ccf_mid_velocity, range_no_mask_change=RvSpectMLBase.max_bc, step=v_step,  max=RvSpectMLBase.max_bc  )
v_grid_orders_ccfs = EchelleCCFs.calc_ccf_v_grid(ccf_plan)
#(order_ccfs, order_ccf_vars )
 order_ccfs = EchelleCCFs.calc_order_ccf_chunklist_timeseries(order_list_timeseries, ccf_plan)

ccf_plan2 = EchelleCCFs.BasicCCFPlan(mask_shape = mask_shape, line_list=EchelleCCFs.BasicLineList2D(line_list_df), midpoint=ccf_mid_velocity, range_no_mask_change=RvSpectMLBase.max_bc, step=v_step,  max=RvSpectMLBase.max_bc )
(order_ccfs2, order_ccf_vars2 ) = EchelleCCFs.calc_order_ccf_and_var_chunklist_timeseries(order_list_timeseries, ccf_plan2)

plt = plot()
 obsid = 9
 order_ccfs = order_ccfs_novars
 for chid in 1:size(order_ccfs,2)
   plot!(plt,v_grid_orders_ccfs,order_ccfs[:,chid,1]./maximum(order_ccfs[:,chid,1]),label=:none)
   #scatter!(plt,v_grid_orders_ccfs,order_ccfs2[:,chid,1]./maximum(order_ccfs2[:,chid,1]),markersize=1.2,label=:none)
   #plot!(plt,v_grid_orders_ccfs,sqrt.(order_ccf_vars[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]),label=:none)

 end
 display(plt)

chid = 84; scatter!(plt,v_grid_orders_ccfs,order_ccfs[:,chid,1]./maximum(order_ccfs[:,chid,1]),markersize=1.3, label=:none)

plt = plot()
  chid = 52
  mean_for_order = vec(mean(order_ccfs[:,chid,:]./maximum(order_ccfs[:,chid,:],dims=1),dims=2))
  for obsid in 1:size(order_ccfs,3)
    plot!(plt,v_grid_orders_ccfs,(order_ccfs[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]).-mean_for_order .+ 0.003*obsid,label=:none)
    #if obsid ==1
  #    scatter!(plt,v_grid_orders_ccfs,(order_ccfs[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]),yerr=sqrt.(order_ccf_vars[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]),label=:none)
  #  end
  end
  display(plt)
  size(mean_for_order)

ccf_comp  = reshape(sum(order_ccfs,dims=2),size(order_ccfs,1),size(order_ccfs,3))

ccfs .-ccf_comp

scatter(map(chid->minimum(order_ccfs[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]),1:85))

plt = plot()
    chid = 50
    for obsid in 1:size(order_ccfs,3)
      plot!(plt,v_grid_orders_ccfs,(order_ccfs[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]) .-
                     mean(order_ccfs[:,chid,:]./maximum(order_ccfs[:,chid,:]),dims=2) ,label=:none)
      #if obsid ==1
      #  scatter!(plt,v_grid_orders_ccfs,(order_ccfs[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]),yerr=sqrt.(order_ccf_vars[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]),label=:none)
      #end
    end
    display(plt)
obsid = 6; scatter!(plt,v_grid_orders_ccfs,(order_ccfs[:,chid,obsid])./maximum(order_ccfs[:,chid,obsid]) .-
               mean(order_ccfs[:,chid,:]./maximum(order_ccfs[:,chid,:]),dims=2) , markersize=1.2, label=string(obsid))

if make_plot(pipeline_plan, :ccf_orders)
   plt = make_heatmap_ccf_vs_order(order_ccfs, v_grid=v_grid_orders_ccfs, order_labels=map(c->order_list_timeseries.chunk_list[1].data[c].λ.indices[2], 1:size(order_ccfs,2)),
                      output_path=output_dir * target_subdir, save_fig=save_plot(pipeline_plan,:ccf_orders)   )
   display(plt)
end

if make_plot(pipeline_plan, :ccf_orders)
   plt = make_heatmap_ccf_vs_order(order_ccfs2, v_grid=v_grid_orders_ccfs, order_labels=map(c->order_list_timeseries.chunk_list[1].data[c].λ.indices[2], 1:size(order_ccfs,2)),
                      output_path=output_dir * target_subdir, save_fig=save_plot(pipeline_plan,:ccf_orders)   )
   display(plt)
end

order_ccfs.-order_ccfs2

if need_to(pipeline_plan,:scalpels)
   rvs_scalpels = map(n->Scalpels.clean_rvs_scalpels(rvs_ccf, ccfs, num_basis=n), 1:5)
   println("RMS RVs cleaned by Scalpels: ",std.(rvs_scalpels) )
   dont_need_to!(pipeline_plan,:scalpels)
end

make_plot!(pipeline_plan,:scalpels)
if make_plot(pipeline_plan, :scalpels)
   @assert !need_to(pipeline_plan, :rvs_ccf_total)
   @assert !need_to(pipeline_plan, :ccf_total)
   plt = make_plots_scalpels(rvs_ccf, ccfs, max_num_basis=2, v_grid=v_grid, times=order_list_timeseries.times, output_path=output_dir, save_fig = save_plot(pipeline_plan, :scalpels) )
   display(plt)
end

if true || need_to(pipeline_plan,:scalpels)
   rvs_scalpels2 = map(n->Scalpels.clean_rvs_scalpels(rvs_ccf, ccfs2, num_basis=n), 1:5)
   #rvs_scalpels2 = map(n->Scalpels.clean_rvs_scalpels(rvs_ccf, reshape(order_ccfs[:,order_to_use,:],155*25,44), num_basis=n), 1:5)
   println("RMS RVs cleaned by Scalpels: ",std.(rvs_scalpels2) )
   dont_need_to!(pipeline_plan,:scalpels)
end

make_plot!(pipeline_plan,:scalpels)
if true || make_plot(pipeline_plan, :scalpels)
   @assert !need_to(pipeline_plan, :rvs_ccf_total)
   @assert !need_to(pipeline_plan, :ccf_total)
   plt = make_plots_scalpels(rvs_ccf2, ccfs2, max_num_basis=2, v_grid=v_grid2, times=order_list_timeseries.times, output_path=output_dir, save_fig = save_plot(pipeline_plan, :scalpels) )
   display(plt)
end


text:
using LsqFit

@testset "fits" begin
    pol1(x,a,b) = a + b * x
    pol2(x, a, b, c) = a + b*x + c*x^2
    pol3(x, a, b, c, d) = a + b*x + c*x^2 + d*x^3

    x=collect(LinRange(0., 10., 100))
    y = pol1.(x,(10.0,), (3.0,),)
    fr = ATools.fit_pol1(x,y)
    @test fr.fitpar[1]≈10.0 && fr.fitpar[2] ≈ 3.0

    y = pol2.(x,(10.0,), (0.5,), (1.,),)
    fr = ATools.fit_pol2(x,y)
    @test fr.fitpar[1]≈10.0 && fr.fitpar[2] ≈ 0.5 && fr.fitpar[3] ≈ 1.0

    y = pol3.(x,(10.0,), (1.0,), (0.7,),(0.5,))
    fr = ATools.fit_pol3(x,y)
    @test fr.fitpar[1]≈10.0 && fr.fitpar[2] ≈ 1.0 && fr.fitpar[3] ≈ 0.7 && fr.fitpar[4] ≈ 0.5

    hg1,pg1 = ATools.hist1d(qs, "gaussian μ=100, σ=10",  50, 50.0, 150.0, norm=true)
    fg = ATools.fit_gauss(hg1)
    @test fg.mu[1] > 99.5 && fg.mu[1] < 100.5
    @test fg.std[1] > 9.5 && fg.std[1] < 10.5

    fr, p1 = ATools.fit_profile(xs, qs, "x", "y", "pol1")
    fpars = fr.fitpar
    fstds = fr.fitstd
    par =[100.0, 0.0]
    @test all([fpars[i] - 5* fstds[i] < par[i] for i in 1:2])

    ## Test cfit with errors using pol3 values above.
    err_sigma = 0.04
    y += rand(Normal(0, err_sigma), length(y))
    p0 = [10.0, 1.0, 0.7, 0.5]
    lb = fill( 0.0, length(p0))
    ub = fill(20.0, length(p0))
    pol3_fit = @. pol(x, p) = p[1] + p[2] * x + p[3] * x^2 + p[4] * x^3
    fit_result = ATools.cfit(pol3_fit, x, y, fill(err_sigma, length(y)), p0, lb, ub)
    @test fit_result.converged
    @test all(isapprox.(coef(fit_result), p0; atol=err_sigma))
end

max_stars_repo_name: jjgomezcadenas/ATools

max_issues_repo_name: jjgomezcadenas/ATools

max_forks_repo_name: jjgomezcadenas/ATools


text:
# test simple methods
#---------------------
# test tensorcopy
A=randn((3,5,4,6))
p=randperm(4)
C1=permutedims(A,p)
C2=tensorcopy(A,1:4,p)
@test vecnorm(C1-C2)<eps()*sqrt(length(C1))*vecnorm(C1+C2)
@test_throws TensorOperations.IndexError tensorcopy(A,1:3,1:4)
@test_throws TensorOperations.IndexError tensorcopy(A,[1,2,2,4],1:4)

# test tensoradd
B=randn((5,6,3,4))
p=[3,1,4,2]
C1=tensoradd(A,p,B,1:4)
C2=A+permutedims(B,p)
@test vecnorm(C1-C2)<eps()*sqrt(length(C1))*vecnorm(C1+C2)
@test_throws DimensionMismatch tensoradd(A,1:4,B,1:4)

# test tensortrace
A=randn(50,100,100)
C1=tensortrace(A,[:a,:b,:b])
C2=zeros(50)
for i=1:50
    for j=1:100
        C2[i]+=A[i,j,j]
    end
end
@test vecnorm(C1-C2)<eps()*sqrt(length(C1))*vecnorm(C1+C2)
A=randn(3,20,5,3,20,4,5)
C1=tensortrace(A,[:a,:b,:c,:d,:b,:e,:c],[:e,:a,:d])
C2=zeros(4,3,3)
for i1=1:4, i2=1:3, i3=1:3
    for j1=1:20,j2=1:5
        C2[i1,i2,i3]+=A[i2,j1,j2,i3,j1,i1,j2]
    end
end
@test vecnorm(C1-C2)<eps()*sqrt(length(C1))*vecnorm(C1+C2)

# test tensorcontract
A=randn(3,20,5,3,4)
B=randn(5,6,20,3)
C1=tensorcontract(A,[:a,:b,:c,:d,:e],B,[:c,:f,:b,:g],[:a,:g,:e,:d,:f];method=:BLAS)
C2=tensorcontract(A,[:a,:b,:c,:d,:e],B,[:c,:f,:b,:g],[:a,:g,:e,:d,:f];method=:native)
C3=zeros(3,3,4,3,6)
for a=1:3, b=1:20, c=1:5, d=1:3, e=1:4, f=1:6, g=1:3
    C3[a,g,e,d,f] += A[a,b,c,d,e]*B[c,f,b,g]
end
@test vecnorm(C1-C3)<eps()*sqrt(length(C1))*vecnorm(C1+C3)
@test vecnorm(C2-C3)<eps()*sqrt(length(C1))*vecnorm(C2+C3)
@test_throws TensorOperations.IndexError tensorcontract(A,[:a,:b,:c,:d],B,[:c,:f,:b,:g])
@test_throws TensorOperations.IndexError tensorcontract(A,[:a,:b,:c,:a,:e],B,[:c,:f,:b,:g])

# test tensorproduct
A=randn(5,5,5,5)
B=rand(Complex128,(5,5,5,5))
C1=reshape(tensorproduct(A,[1,2,3,4],B,[5,6,7,8],[1,2,5,6,3,4,7,8]),(5*5*5*5,5*5*5*5))
C2=kron(reshape(B,(25,25)),reshape(A,(25,25)))
@test vecnorm(C1-C2)<eps()*sqrt(length(C1))*vecnorm(C1+C2)
@test_throws TensorOperations.IndexError tensorproduct(A,[:a,:b,:c,:d],B,[:d,:e,:f,:g])
@test_throws TensorOperations.IndexError tensorproduct(A,[:a,:b,:c,:d],B,[:e,:f,:g,:h],[:a,:b,:c,:d,:e,:f,:g,:i])

# test index notation
#---------------------
# Da=10
# Db=15
# Dc=4
# Dd=8
# De=6
# Df=7
# Dg=3
# Dh=2
# A=rand(Complex128,(Da,Dc,Df,Da,De,Db,Db,Dg))
# B=rand(Complex128,(Dc,Dh,Dg,De,Dd))
# C=rand(Complex128,(Dd,Dh,Df))
# D=rand(Complex128,(Dd,Df,Dh))
# D[l"d,f,h"]=A[l"a,c,f,a,e,b,b,g"]*B[l"c,h,g,e,d"]+0.5*C[l"d,h,f"]
# @test_approx_eq(vecnorm(D),sqrt(abs(scalar(D[l"d,f,h"]*conj(D[l"d,f,h"])))))
# @test_throws IndexError D[l"a,a,a"]
# @test_throws IndexError D[l"a,b,c,d"]
# @test_throws IndexError D[l"a,b"]

# test in-place methods
#-----------------------
# test different versions of in-place methods,
# with changing element type and with nontrivial strides

# tensorcopy!
Abig=randn((30,30,30,30))
A=sub(Abig,1+3*(0:9),2+2*(0:6),5+4*(0:6),4+3*(0:8))
p=[3,1,4,2]
Cbig=zeros(Complex128,(50,50,50,50))
C=sub(Cbig,13+(0:6),11+4*(0:9),15+4*(0:8),4+3*(0:6))
Acopy=tensorcopy(A,1:4,1:4)
Ccopy=tensorcopy(C,1:4,1:4)
TensorOperations.tensorcopy!(A,1:4,C,p)
TensorOperations.tensorcopy!(Acopy,1:4,Ccopy,p)
@test vecnorm(C-Ccopy)<eps()*sqrt(length(C))*vecnorm(C+Ccopy)
@test_throws TensorOperations.IndexError TensorOperations.tensorcopy!(A,1:3,C,p)
@test_throws DimensionMismatch TensorOperations.tensorcopy!(A,p,C,p)
@test_throws TensorOperations.IndexError TensorOperations.tensorcopy!(A,1:4,C,[1,1,2,3])

# tensoradd!
Cbig=zeros(Complex128,(50,50,50,50))
C=sub(Cbig,13+(0:6),11+4*(0:9),15+4*(0:8),4+3*(0:6))
Acopy=tensorcopy(A,1:4,p)
Ccopy=tensorcopy(C,1:4,1:4)
alpha=randn()
beta=randn()
TensorOperations.tensoradd!(alpha,A,1:4,beta,C,p)
Ccopy=beta*Ccopy+alpha*Acopy
@test vecnorm(C-Ccopy)<eps()*sqrt(length(C))*vecnorm(C+Ccopy)
@test_throws TensorOperations.IndexError TensorOperations.tensoradd!(1.2,A,1:3,0.5,C,p)
@test_throws DimensionMismatch TensorOperations.tensoradd!(1.2,A,p,0.5,C,p)
@test_throws TensorOperations.IndexError TensorOperations.tensoradd!(1.2,A,1:4,0.5,C,[1,1,2,3])

# tensortrace!
Abig=rand((30,30,30,30))
A=sub(Abig,1+3*(0:8),2+2*(0:14),5+4*(0:6),7+2*(0:8))
Bbig=rand(Complex128,(50,50))
B=sub(Bbig,13+(0:14),3+5*(0:6))
Acopy=tensorcopy(A,1:4)
Bcopy=tensorcopy(B,1:2)
alpha=randn()
beta=randn()
TensorOperations.tensortrace!(alpha,A,[:a,:b,:c,:a],beta,B,[:b,:c])
Bcopy=beta*Bcopy
for i=1+(0:8)
    Bcopy+=alpha*slice(A,i,:,:,i)
end
@test vecnorm(B-Bcopy)<eps()*vecnorm(B+Bcopy)*sqrt(length(B))
@test_throws TensorOperations.IndexError TensorOperations.tensortrace!(alpha,A,[:a,:b,:c],beta,B,[:b,:c])
@test_throws DimensionMismatch TensorOperations.tensortrace!(alpha,A,[:a,:b,:c,:a],beta,B,[:c,:b])
@test_throws TensorOperations.IndexError TensorOperations.tensortrace!(alpha,A,[:a,:b,:a,:a],beta,B,[:c,:b])
@test_throws DimensionMismatch TensorOperations.tensortrace!(alpha,A,[:a,:b,:a,:c],beta,B,[:c,:b])

# tensorcontract!
Abig=rand((30,30,30,30))
A=sub(Abig,1+3*(0:8),2+2*(0:14),5+4*(0:6),7+2*(0:8))
Bbig=rand(Complex128,(50,50,50))
B=sub(Bbig,3+5*(0:6),7+2*(0:7),13+(0:14))
Cbig=rand(Complex64,(40,40,40))
C=sub(Cbig,3+2*(0:8),13+(0:8),7+3*(0:7))
Acopy=tensorcopy(A,1:4)
Bcopy=tensorcopy(B,1:3)
Ccopy=tensorcopy(C,1:3)
alpha=randn()
beta=randn()
Ccopy=beta*Ccopy
for d=1+(0:8),a=1+(0:8),e=1+(0:7)
    for b=1+(0:14),c=1+(0:6)
        Ccopy[d,a,e]+=alpha*A[a,b,c,d]*conj(B[c,e,b])
    end
end
TensorOperations.tensorcontract!(alpha,A,[:a,:b,:c,:d],'N',B,[:c,:e,:b],'C',beta,C,[:d,:a,:e];method=:BLAS)
@test vecnorm(C-Ccopy)<eps(Float32)*vecnorm(C+Ccopy)*sqrt(length(C))
Cbig=rand(Complex64,(40,40,40))
C=sub(Cbig,3+2*(0:8),13+(0:8),7+3*(0:7))
Ccopy=tensorcopy(C,1:3)
Ccopy=beta*Ccopy
for d=1+(0:8),a=1+(0:8),e=1+(0:7)
    for b=1+(0:14),c=1+(0:6)
        Ccopy[d,a,e]+=alpha*A[a,b,c,d]*conj(B[c,e,b])
    end
end
TensorOperations.tensorcontract!(alpha,A,[:a,:b,:c,:d],'N',B,[:c,:e,:b],'C',beta,C,[:d,:a,:e];method=:native)
@test vecnorm(C-Ccopy)<eps(Float32)*vecnorm(C+Ccopy)*sqrt(length(C))
@test_throws TensorOperations.IndexError TensorOperations.tensorcontract!(alpha,A,[:a,:b,:c,:a],'N',B,[:c,:e,:b],'N',beta,C,[:d,:a,:e])
@test_throws TensorOperations.IndexError TensorOperations.tensorcontract!(alpha,A,[:a,:b,:c,:d],'N',B,[:c,:b],'N',beta,C,[:d,:a,:e])
@test_throws TensorOperations.IndexError TensorOperations.tensorcontract!(alpha,A,[:a,:b,:c,:d],'N',B,[:c,:e,:b],'N',beta,C,[:d,:e])
@test_throws DimensionMismatch TensorOperations.tensorcontract!(alpha,A,[:a,:b,:c,:d],'N',B,[:c,:e,:b],'N',beta,C,[:d,:e,:a])


text:
function main(window)
    push!(window.assets, "tex")
    push!(window.assets, "widgets")

    texᵗ = Signal("f(x) = \\int_{-\\infty}^\\infty
               \\hat f(\\xi)\\,e^{2 \\pi i \\xi x}
               \\,d\\xi")

    modeᵗ = Signal(false)

    map(texᵗ, modeᵗ) do t, mode
        vbox(h1("LaTeX"),
             hbox("LaTeX support is via ", hskip(1em), tex("\\KaTeX")),
             textinput() >>> texᵗ,
             hbox("Show as a block", hskip(1em), checkbox(mode) >>> modeᵗ),
             vskip(1em),
             tex(t, block=mode)) |> Escher.pad(1em) |> maxwidth(30em)
    end
end


text:
module SDC
export SDC, make_SDC, make_spectral_int_matrix, gauss_lobatto_nodes

using FastGaussQuadrature: gausslobatto
using InvertedIndices: Not
using Polynomials: fromroots, integrate
#import BenchmarkTools

#"""
#Solves the equation:
#    w + alpha*fI(t,w) = b
#For w given α, t, and b
#Do I need a best guess for w, as well?
#"""
#function implicit_solve(α, t, b)
#end

"""
Get n Gauss-Lobatto quadrature nodes, on interval [0-1].
"""
function gauss_lobatto_nodes(n::Int)
    nodes, weights =  gausslobatto(n)
    return 0.5*(nodes .+ 1.0)
end

"""
Get the spectral integration matrix
"""
function make_spectral_int_matrix(n::Int)
    nodes = gauss_lobatto_nodes(n) # Perhaps add option to change quadrature in the future
    S = zeros(n,n)

    for j = 1:n
        Lt = fromroots(nodes[Not(j)])
        Li = (1.0/Lt(nodes[j]))*Lt
        for i = 1:n
            S[i,j] = integrate(Li,0.0,nodes[i])
        end
    end
    return S
end

"""
Notes for improvement:

The main usage for this function will probably be something like:
for i in 1:number_of_timesteps
    SDC(...)

So I don't want to waste resources recomputing S everytime, or setting up all
the equations.

Perhaps I should have a create_SDC function? This would be provided all
parameters except the initial guess, and create an SDC function which only
needs the initial values as arguments.

From a UX perspective, I like that they will have an SDC function whose only
arguments are the initial data. That would be very easy to use. But I don't
like that they have to call a function to create a function. That seems like it
might be confusing.

Actually, I think it would be good to also have Δt as an argument to the
user-SDC function. I can imagine a user might want to use variable timestep
size (but they probably don't need to vary the number of collocation nodes,
implicit/explicit f, etc).
"""
function sdc(n, fE, fI, t, Δt, u, implicit_solver)
    nodes = gauss_lobatto_nodes(n)
    n_cor = 2*n-3
    S = make_spectral_int_matrix(n) # nodes will be computed twice; should make more efficient
    subtimes = (Δt*nodes).+ t
    subvals = Vector{typeof(u)}(undef, n)

    subvals[1] = u
    # Prediction
    for m in 1:n-1
        subvals[m+1] = begin
            u_m = subvals[m]
            t_m = subtimes[m]
            t_mp1 = subtimes[m+1]
            Δt_m = t_mp1 - t_m
            RHS = u_m + Δt_m*fE(t_m,u_m)
            implicit_solver(-Δt_m, t_m, RHS)
        end
    end

    f(t, u) = fE(t, u) + fI(t, u)
    # Correction
    for k in 1:n_cor
        # B15 - Not needed for gauss lobatto nodes, just leaves subvals as they were
        # But B15 will be needed if start point is not a node
        f_subvals = [f(subtimes[i], subvals[i]) for i in 1:n]
        u_m_k = subvals[1]
        for m in 1:n-1
            #B16
            u_m_k_next = subvals[m+1] # Should figure out better notation for this
            subvals[m+1] = begin
                u_m_kp1 = subvals[m]
                t_m = subtimes[m]
                t_mp1 = subtimes[m+1]
                Δt_m = t_mp1 - t_m
                u_mp1_k = subvals[m+1]
                ΔfE = fE(t_m, u_m_kp1) - fE(t_m, u_m_k)
                RHS = subvals[m] + Δt_m*(ΔfE - f(t_mp1, u_mp1_k))
                RHS += sum(q -> (S[m+1,q]-S[m,q])*f_subvals[q], 1:size(S)[2])
                implicit_solver(-Δt_m, fI, RHS)
            end
            u_m_k = u_m_k_next
        end
    end
    # Will be different if endpoint is a node, need to use B6
    return subvals[end]
end

"""
This function returns a function: an SDC solver which needs only the initial
conditions as the arguments. Moreover, the quadrature and spectral integration
matrix are built-in to the function, not calculated each run.
"""
function make_SDC(n, fE, fI, implicit_solver)
    nodes = gauss_lobatto_nodes(n)
    n_cor = 2*n-3
    S = make_spectral_int_matrix(n) # nodes will be computed twice; should make more efficient

    return function small_arg_list_SDC(u, t, Δt)
        subtimes = (Δt*nodes).+ t
        subvals = Vector{typeof(u)}(undef, n)
        subvals[1] = u
        # Prediction
        for m in 1:n-1
            subvals[m+1] = begin
                u_m = subvals[m]
                t_m = subtimes[m]
                t_mp1 = subtimes[m+1]
                Δt_m = t_mp1 - t_m
                RHS = u_m + Δt_m*fE(t_m,u_m)
                implicit_solver(-Δt_m, fI, RHS)
            end
        end

        f(t, u) = fE(t, u) + fI(t, u)
        # Correction
        for k in 1:n_cor
            # B15 - Not needed for gauss lobatto nodes, just leaves subvals as they were
            # But B15 will be needed if start point is not a node
            f_subvals = [f(subtimes[i], subvals[i]) for i in 1:n]
            u_m_k = subvals[1]
            for m in 1:n-1
                #B16
                u_m_k_next = subvals[m+1] # Should figure out better notation for this
                subvals[m+1] = begin
                    u_m_kp1 = subvals[m]
                    t_m = subtimes[m]
                    t_mp1 = subtimes[m+1]
                    Δt_m = t_mp1 - t_m
                    u_mp1_k = subvals[m+1]
                    ΔfE = fE(t_m, u_m_kp1) - fE(t_m, u_m_k)
                    RHS = subvals[m] + Δt_m*(ΔfE - f(t_mp1, u_mp1_k))
                    RHS += sum(q -> (S[m+1,q]-S[m,q])*f_subvals[q], 1:size(S)[2])
                    implicit_solver(-Δt_m, fI, RHS)
                end
                u_m_k = u_m_k_next
            end
        end
        # Will be different if endpoint is a node, need to use B6
        return subvals[end]
    end
end

#function graph(u_np1_saves, T::Timing;
#              graph_name::String="SDC_evolution.png",
#              display_plot::Bool=false)
#
#    # Get exact value for comparison purposes
#    u_exact = u(T.t_n + T.Δt)
#    # Compute errors
#    u_np1_saves_errors = Vector{Tuple{Int64, Float64}}(undef, length(u_np1_saves))
#    for (i, (j, u_np1)) in enumerate(u_np1_saves)
#        #u_np1_saves_errors[i] = (j, abs(u_np1-u_exact))
#        u_np1_saves_errors[i] = (j, u_np1 - u_exact)
#    end
#
#    # Plotting
#    plot = Plots.plot(
#        getindex.(u_np1_saves_errors, 1), # No of Corrections
#        getindex.(u_np1_saves_errors, 2), # Error
#        label="Approximation of u(t_n+Δt)"
#    )
#    p = length(T.Δtimes)+1
#    Plots.plot!(
#        plot,
#        title="Error Over 1 Timestep Vs No of Corrections \nλ=$λ, Δt=$(T.Δt), t_n=$(T.t_n), p=$p (Gauss-Lobatto)",
#        xlabel="# of Corrections (0 ⟹ Initial Prediction)",
#        ylabel="Error",
#        xlim=(0, 2*p),
#        #yaxis=:log,
#    )
#    # Make a line on y=0, for easier comprehension
#    Plots.hline!(plot, [0], label="")
#
#    Plots.savefig(plot, graph_name)
#
#    if display_plot
#        display(plot)
#    end
#    return nothing
#end

#"""
#Make sure collocation nodes have proper scale and order.
#"""
#function validate_collocation_nodes(c_nodes)
#    for node in c_nodes
#        if node < 0 || node > 1
#            throw(DomainError(node, "Collocation nodes must be ∈ [0,1]."))
#        end
#    end
#    for i in 1:length(c_nodes)-1
#        if c_nodes[i+1] <= c_nodes[i]
#            throw(DomainError(node, "Collocation nodes must be obey c_m+1 > c_m ∀ m = 1, ... , p."))
#        end
#    end
#end

end