using Donquicote, Test, BenchmarkTools
using Roots, ForwardDiff

f(x) = exp(x) - x^4
∂f(x) = exp(x) - 3*x^4
∂2f(x) = exp(x) - 2*x^3
D(f) = x -> ForwardDiff.derivative(f,float(x))

println("Bisection (built-in):")
@btime find_zero(f, (1,3), Bisection(); xatol = 1e-7)
println("FalsePosition (build-in):")
@btime find_zero(f, (1,3), FalsePosition(), xatol = 1e-7)
println("Newton (built-in):")
@btime find_zero((f, D(f)), 2, Roots.Newton(); xatol = 1e-7)
println("Newton (custom):")
@btime newton(f, D(f), 2)
println("Halley (built-in):")
@btime find_zero((f, D(f), D(D(f))), 2, Roots.Halley(); xatol = 1e-7)
println("Halley (custom):")
@btime haley(f, D(f), D(D(f)), 2)
println("Schroeder (built-in):")
@btime find_zero((f, D(f), D(D(f))), 2, Roots.Schroeder(); xatol = 1e-7)
println("Schroeder (custom):")
@btime schroeder(f, D(f), D(D(f)), 2)

##

V = ones(3)
V0 = zeros(V)


# ---
# title: 363. Max Sum of Rectangle No Larger Than K
# id: problem363
# author: Indigo
# date: 2021-05-29
# difficulty: Hard
# categories: Binary Search, Dynamic Programming, Queue
# link: <https://leetcode.com/problems/max-sum-of-rectangle-no-larger-than-k/description/>
# hidden: true
# ---
#
# Given a non-empty 2D matrix _matrix_ and an integer _k_ , find the max sum of
# a rectangle in the _matrix_ such that its sum is no larger than _k_.
#
# **Example:**
#
#
#
#     Input: matrix = [[1,0,1],[0,-2,3]], k = 2
#     Output: 2
#     Explanation:  Because the sum of rectangle [[0, 1], [-2, 3]] is 2,
#                  and 2 is the max number no larger than k (k = 2).
#
# **Note:**
#
#   1. The rectangle inside the matrix must have an area > 0.
#   2. What if the number of rows is much larger than the number of columns?
#
#
## @lc code=start
using LeetCode

function max_sum_submatrix(matrix::AbstractMatrix{Int}, k::Int)
    m, n = size(matrix)
    m > n && return max_sum_submatrix(matrix', k)
    pref_sum = OffsetArray(fill(0, m + 1, n + 1), -1, -1)

    for i in 1:m, j in 1:n
        pref_sum[i, j] = pref_sum[i - 1, j] + pref_sum[i, j - 1] + matrix[i, j] -
                         pref_sum[i - 1, j - 1]
    end
    res = typemin(Int)
    for tp in 1:m, bot in tp:m
        ss = SortedSet{Int}()
        push!(ss, 0)
        for r in 1:n
            rows_sum = pref_sum[bot, r] - pref_sum[tp - 1, r]
            tk = searchsortedfirst(ss, rows_sum - k)
            tk != pastendsemitoken(ss) && (res = max(res, rows_sum - deref((ss, tk))))
            push!(ss, rows_sum)
        end
    end
    return res
end
## @lc code=end


function calculate(f, g)
    local a = 0;
    for i = 2:length(f)
        a = a + f[i] * g[i];
    end
    return a;
end

function zeckendorf(j, f, g, y)
    if calculate(f, g) == y
        return (g, true);
    else
        for i = j:length(f)
            g[i] = 1;
            if zeckendorf(i + 2, f, g, y)[2]
                return (g, true);
            end
            g[i] = 0;
        end
    end
    return (g, false);
end

function printer(a)
    deleteat!(a, 1);
    if a[length(a)] == 0
        deleteat!(a, length(a));
    end
    println(join(a[end:-1:1], ""));
end

global x = parse(Int64, readline());
global f = Int64[1, 1];
while f[length(f)] < x
    resize!(f, length(f) + 1);
    f[length(f)] = f[length(f) - 1] + f[length(f) - 2];
end
g = zeros(Int64, length(f));
a = zeckendorf(2, f, g, x);
b = zeckendorf(3, f, g, x);
if a[2]
    printer(a[1]);
else
    printer(b[2]);
end


using ChaosTools, Test, Statistics

println("\nTesting transit time statistics...")
@testset "Transit time statistics" begin

@testset "Standard map (exact)" begin
    # INPUT
    ds = Systems.standardmap()
    T = 10000 # maximum time

    # period 3 of standard map
    u0 = SVector(0.8121, 1.6243)
    εs = sort!([4.0, 2.0, 0.01]; rev=true)
    # radius of 4.0 covers the first 2 points of the period 3 while 2.0 covers only
    # the first point. Therefore:
    # all first return times must be 1, and all second return times must be 2,
    # and 3rd must be same as second (because it doesn't matter how close are)
    exits, entries = transit_time_statistics(ds, u0, εs, T)
    transits, returns = transit_return(exits, entries)

    @test all(x -> length(x) > 5, exits)
    @test all(x -> length(x) > 5, entries)
    @test all(issorted, exits)
    @test all(issorted, entries)
    @test all(isequal(2), transits[1])
    @test all(isequal(1), transits[2])
    @test transits[2] == transits[3]
    @test all(isequal(1), returns[1])
    @test all(isequal(2), returns[2])
    @test returns[2] == returns[3]

    # quasiperiodic around period 3:
    u0 = SVector(0.877, 1.565)
    εs = sort!([4.0, 0.5, 0.1]; rev=true)
    exits, entries = transit_time_statistics(ds, u0, εs, T)
    transits, returns = transit_return(exits, entries)

    @test all(issorted, exits)
    @test all(issorted, entries)
    @test all(x -> length(x) > 5, exits)
    @test all(x -> length(x) > 5, entries)

    # For ε=4.0, nothing changes with the before
    @test all(isequal(1), returns[1])
    @test all(isequal(2), transits[1])

    # Similarly, 0.5 should be the same as before
    @test all(isequal(1), transits[2])
    @test all(isequal(2), returns[2])

    # But now, the third entry is different, because it has the size of the quasiperiodic
    # stability island torous
    @test returns[3] ≠ returns[2]
    @test transits[3] ≠ transits[2]
    @test any(>(3), returns[3])
    @test all(isequal(1), transits[3]) # still need only one step to exit
end

@testset "Towel map (boxes)" begin
    to = Systems.towel()
    tr = trajectory(to, 5000; Ttr = 10)
    u0 = tr[3000]

    # With these boxes, in the first 5 steps, the trajectory enters the y and z range
    # but not the x range. Therefore it should NOT enter the box. See figure!
    εs = [
        SVector(0.05, 0.05, 0.125),
        SVector(0.005, 0.005, 0.025),
    ]

    # Visual guidance
    # using PyPlot
    # tr0 = trajectory(to, 5, u0)
    # fig, axs = subplots(1,3)
    # comb = ((1, 2), (1, 3), (2, 3))
    # for i in 1:3
    #     j, k = comb[i]
    #     ax = axs[i]
    #     ax.scatter(tr[:, j], tr[:, k], s = 2, color = "C$(i-1)")
    #     ax.scatter([u0[j]], [u0[k]], s = 20, color = "k")
    #     ax.plot(tr0[:, j], tr0[:, k], color = "k")
    #     for l in 1:length(εs)
    #         rect = matplotlib.patches.Rectangle(
    #         u0[[j, k]] .- εs[l][[j, k]], 2εs[l][j], 2εs[l][k],
    #         alpha = 0.25, color = "k"
    #         )
    #         ax.add_artist(rect)
    #     end
    # end

    exits, entries = transit_time_statistics(to, u0, εs, 10000)
    transits, returns = transit_return(exits, entries)

    @test all(issorted, exits)
    @test all(issorted, entries)
    @test length(exits[1]) > length(exits[2])
    @test returns[1][1] > 5
    @test mean(returns[1]) < mean(returns[2])
end
#
# @testset "Continuous Roessler" begin
#
# # %%
# ro = Systems.roessler(ones(3))
# tr = trajectory(ro, 5000; Ttr = 100)
# u0 = trajectory(ro, 11; Ttr = 100)[end] # return center
# εs = sort!([1.0, 0.1, 0.01]; rev=true)
#
# # Visual guidance
# using PyPlot
# tr0 = trajectory(ro, 50, u0)
# fig, axs = subplots(1,3)
# comb = ((1, 2), (1, 3), (2, 3))
# for i in 1:3
#     j, k = comb[i]
#     ax = axs[i]
#     ax.plot(tr[:, j], tr[:, k], lw = 0.5, color = "C$(i-1)", alpha = 0.5)
#     ax.scatter([u0[j]], [u0[k]], s = 20, color = "k")
#     ax.plot(tr0[:, j], tr0[:, k], color = "k", lw = 1.0, ls = "--")
#     if eltype(εs[1]) <: Vector
#         for l in 1:length(εs)
#             rect = matplotlib.patches.Rectangle(
#             u0[[j, k]] .- εs[l][[j, k]], 2εs[l][j], 2εs[l][k],
#             alpha = 0.25, color = "k"
#             )
#             ax.add_artist(rect)
#         end
#     else
#         for l in 1:length(εs)
#             circ = matplotlib.patches.Circle(
#                 u0[[j, k]], εs[l]; alpha = 0.25, color = "k"
#             )
#             ax.add_artist(circ)
#         end
#     end
# end
#
# T = 100000.0 # maximum time
# exits, entries = transit_time_statistics(ro, u0, εs, T)
# transits, returns = transit_return(exits, entries)
#
# @test all(issorted, exits)
# @test all(issorted, entries)
# @test length(exits[1]) > length(exits[2])
# @test length(entries[1]) > length(entries[2])
#
# # We know the average period of the Roessler system. Therefore the return times
# # cannot be possibly larger than it (because u0 is in the xy plane)
# avg_period = 6.0
# @test all(r -> r > avg_period/2, returns[1])
#
#
# end

end


# Bicycle domain ported from rlpy https://github.com/rlpy/rlpy/blob/master/rlpy/Domains/Bicycle.py
# original paper Learning to Drive a Bicycle using Reinforcement Learning and Shaping, Jette Randlov, Preben Alstrom, 1998. https://www.semanticscholar.org/paper/Learning-to-Drive-a-Bicycle-Using-Reinforcement-and-Randl%C3%B8v-Alstr%C3%B8m/10bad197f1c1115005a56973b8326e5f7fc1031c

# Bicycle domain might not be correct.
struct BicycleParams{T} <:Any where {T<:Real}
	g::T    	# gravity 9.82
	v::T		# velocity of the bicycle (10 km/h paper value) default is 10. / 3.6 (10kph appox 6.2mph and 10kph is 2.77778 meters per second or 10/3)
	d_CM::T 	# vertical distance between the CM for the bicycle and for the cyclist 0.3m
	c::T    	# horizontal distance between the point where the front wheel touches the ground and the CM 0.66 m
	h::T		# height of the CM over the ground  0.94 cm
	M_c::T  	# mass bicycle  9kg ≈ 19.8lbs 15kg ≈ 33lbs (default is 15kg)
	M_d::T		# mass tire 1.0kg ≈ 2.2lbs 1.5kg ≈ 3.3lbs 3.0kg ≈ 4.4lbs (default is 1.7)
	M_p::T		# mass cyclist M_p = 45.4kg ≈ 100lbs 60.0kg ≈ 132.0lbs 90.7kg ≈ 200lbs (default is 60kg)
	M::T   		# mass of cyclist and bike M_p + M_c
	r::T	   	# radius of tire 0.34m
	dsigma::T  	# angular velcity of a tire (v/r)
	I::T    	# moment of inertia for bicycle and cyclist 13/3 * M_c * h^2 + M_p * (h + d_CM)^2
	I_dc::T 	# moment of inertia for tire M_d * r^2
	I_dv::T 	# moment of inertia for tire 3/2 * M_d * r^2
	I_dl::T 	# moment of inertia for tire 1/2 * M_d * r^2
	l::T    	# disance between the front and tire and the back tire at the point where they both touch the ground 1.11m

    function BicycleParams(T::Type=Float64)
		g = T(9.82)
		v = T(10. / 3.6)
		d_CM = T(0.3)
		c = T(0.66)
		h = T(0.94)
		M_c = T(15.)
		M_d = T(1.7)
		M_p = T(60.)
		M = M_c + M_p
		r = T(0.34)
		dsigma = T(v / r)
		I = T(13. / 3. * M_c * h^2 + M_p * (h+d_CM)^2)
		I_dc = T(M_d * r^2)
		I_dv = T(3. / 2.)
		I_dl = T(M_d / 2. * r^2)
		l = T(1.11)
		new{T}(g, v, d_CM, c, h, M_c, M_d, M_p, M, r, dsigma, I, I_dc, I_dv, I_dl, l)
	end

	function BicycleParams(::Type{T}, g, M_c, M_d, M_p, v) where {T}
		d_CM = T(0.3)
		c = T(0.66)
		h = T(0.94)
		M = M_c + M_p
		r = T(0.34)
		dsigma = T(v / r)
		I = T(13. / 3. * M_c * h^2 + M_p * (h+d_CM)^2)
		I_dc = T(M_d * r^2)
		I_dv = T(3. / 2. * M_d * r^2)
		I_dl = T(M_d / 2. * r^2)
		l = T(1.11)
		new{T}(g, v, d_CM, c, h, M_c, M_d, M_p, M, r, dsigma, I, I_dc, I_dv, I_dl, l)
	end
end

function random_bike_params()
	g = rand(Uniform(8., 10.))  		# gravity
	v = rand(Uniform(8., 15.)) / 3.6	# velocity of the bicycle mps
	M_c = rand(Uniform(9., 15.))  		# mass bicycle  9kg ≈ 19.8lbs 15kg ≈ 33lbs (default is 15kg)
	M_d = rand(Uniform(1., 3.))	    	# mass tire 1.0kg ≈ 2.2lbs 1.5kg ≈ 3.3lbs 3.0kg ≈ 4.4lbs (default is 1.7)
	M_p = rand(Uniform(45.4, 90.7))		# mass cyclist M_p = 45.4kg ≈ 100lbs 60.0kg ≈ 132.0lbs 90.7kg ≈ 200lbs (default is 60kg)

	params = BicycleParams(Float64, g, M_c, M_d, M_p, v)
	return params
end

function bicycle_balance(;randomize=false, maxT=30.0)
	if randomize
		params = random_bike_params()
	else
		params = BicycleParams()
	end
	m = create_bicycle_balance(params, maxT=maxT, dt=0.01)
end

function bicycle_goal(;randomize=false, maxT=30.0)
	if randomize
		params = random_bike_params()
	else
		params = BicycleParams()
	end
	m = create_bicycle_goal(params, maxT=maxT, dt=0.01)
end

function bicycle_spaces(maxT)
	X1 = zeros((5,2))
    X1[1,:] .= [-π * 12. / 180., π * 12 / 180]
	X1[2,:] .= [-π, π]
	X1[3,:] .= [-π * 80. / 180., π * 80. / 180.]
	X1[4,:] .= [-π, π]
	X1[5,:] .= [-π, π]
	X2 = zeros(2,2)
	X2[1,:] .= [-Inf, Inf]
	X2[2,:] .= [-Inf, Inf]

	S = ([0.0, maxT], X1, X2)
	A = 1:9
	return S, A
end

function create_bicycle_balance(params; maxT=60.0, dt=0.01)
	S, A = bicycle_spaces(maxT)
	X = S[2]

	p = (s,a)->bicycle_balance_step!(s, a, params, dt, maxT)
	d0 = bicycle_balance_initial

	meta = Dict{Symbol,Any}()
    # meta[:minreward] = -5.0
    # meta[:maxreward] = 5
    # meta[:minreturn] = -5 * ceil(maxT / (dt * 20))
    # meta[:maxreturn] =  5 * ceil(maxT / (dt * 20))
    meta[:stochastic] = true
    meta[:minhorizon] = 1  # not sure the the true minimum is. This seems like a good lower bound
    meta[:maxhorizon] = ceil(Int, maxT / dt)
    meta[:discounted] = false
	meta[:episodes] = 2000
	bplot = BicyclePlotData()
    render = (state,clearplot=false)->bicycleplot(bplot,state,clearplot)
    m = SequentialProblem(S,X,A,p,d0,meta,render)
	return m
end

function create_bicycle_goal(params; maxT=600.0, dt=0.01)
	S, A = bicycle_spaces(maxT)
	X = S[1]

	p = (s,a)->bicycle_goal_step!(s, a, params, dt, maxT)
	d0 = bicycle_goal_initial

	meta = Dict{Symbol,Any}()
    # meta[:minreward] = -5.0
    # meta[:maxreward] = 5
    # meta[:minreturn] = -5 * ceil(maxT / (dt * 20))
    # meta[:maxreturn] =  5 * ceil(maxT / (dt * 20))
    meta[:stochastic] = true
    meta[:minhorizon] = 1  # not sure the the true minimum is. This seems like a good lower bound
    meta[:maxhorizon] = ceil(Int, maxT / dt)
    meta[:discounted] = false
	meta[:episodes] = 3000
	bplot = BicyclePlotData()
    render = (state,clearplot=false)->bicycleplot(bplot,state,clearplot)
    m = SequentialProblem(S,X,A,p,d0,meta,render)
	return m
end

function bicycle_balance_initial()
	t = 0.0
	x = zeros(5)
	pos = zeros(2)

	return (t,x,pos), x
end

function bicycle_goal_initial()
	t = 0.0
	x = zeros(5)
	pos = zeros(2)
	pos[1] = -500.0  # 20 * randn()
	pos[2] = 0.0 # 20 * randn()
	return (t,x,pos), x
end

function bicycle_sim!(state, pos, constants::BicycleParams{T}, u::T, d::T, dt::T) where {T<:Real}
	omega, omegaDot, theta, thetaDot, psi = state
	g = constants.g
	v = constants.v
	d_CM = constants.d_CM
	c = constants.c
	h = constants.h
	M_c = constants.M_c
	M_d = constants.M_d
	M_p = constants.M_p
	M = constants.M
	r = constants.r
	dsigma = constants.dsigma
	I = constants.I
	I_dc = constants.I_dc
	I_dv = constants.I_dv
	I_dl = constants.I_dl
	l = constants.l

	w = 0.02 * (2.0 * rand() - 1.0)

	phi = omega + atan(d+w) / h
	invr_f = abs(sin(theta)) / l
	invr_b = abs(tan(theta)) / l
	if theta != 0.
		invr_CM = ((l-c)^2 + invr_b^(-2))^(-0.5)
	else
		invr_CM = 0.
	end

	nomega = omega + dt * omegaDot
	nomegaDot = omegaDot + dt * (M * h * g * sin(phi) - cos(phi) * (I_dc * dsigma * thetaDot + sign(theta) * v^2 * (M_d * r * (invr_f + invr_b) + M * h * invr_CM))) / I
	out = theta + dt * thetaDot

	if abs(out) < (π * 80. / 180.)
		ntheta = out
		nthetaDot = thetaDot + dt * (u - I_dv * dsigma * omegaDot) / I_dl
	else
		ntheta = sign(out) * (π * 80. / 180.)
		nthetaDot = 0.
	end
	# ntheta = mod(ntheta + π, 2. * π) - π

	npsi = psi + dt * sign(theta) * v * invr_b

	npsi = npsi % (2 * π)
	if npsi > π
		npsi -= 2 * π
	end

	# dposxf = v * dt * -sin(psi + theta + sign(psi + theta) * asin(v*dt*0.5*invr_f))
	# dposyf = v * dt * cos(psi + theta + sign(psi + theta) * asin(v*dt*0.5*invr_f))
	dposxb = v * dt * -sin(psi + sign(psi) * asin(v*dt*0.5*invr_b))
	dposyb = v * dt * cos(psi + sign(psi) * asin(v*dt*0.5*invr_b))

	state[1] = nomega
	state[2] = omegaDot = nomegaDot
	state[3] = ntheta
	state[4] = nthetaDot
	state[5] = npsi
	pos[1] += dposxb
	pos[2] += dposyb

	return npsi - psi

end

function compute_actions(action::Int)
	if action <= 0 || action > 10
        error("Action needs to be an integer in [1, 9]")
    end
    u = 0.
	d = 0.
	uaction = floor(Int, (action-1) / 3) + 1
	daction = action % 3

	if uaction == 1
        u = -2.
    elseif uaction == 2
    	u =  0.
	else
		u = 2.
    end

	if daction == 1
		d = -0.02
	elseif daction == 2
		d =  0.
	else
		d = 0.02
	end
	return u,d
end

function bicycle_balance_step!(state, action::Int, params, dt, maxT)
    u,d = compute_actions(action)
	t,x,pos = state
	psi_diff = bicycle_sim!(x, pos, params, u, d, dt)
	t += dt

	done = bike_fell(x) || t ≥ maxT

	reward = 1

	γ = 1.0
	if done
		γ = 0.0
	end

    return (t, x, pos), x, reward, γ
end

function bicycle_goal_step!(state, action::Int, params, dt, maxT)
    u,d = compute_actions(action)
	t,x,pos = state
	psi_diff = bicycle_sim!(x, pos, params, u, d, dt)
	t += dt

	if bike_fell(x)
		reward = -1.0
		γ = 0.0
	elseif bike_at_goal(pos)
		reward = 0.01
		γ = 0.0
	elseif t ≥ maxT
		reward = (4.0 - atan(pos...)^2) * 0.00004
		γ = 0.0
	else
		reward = (4.0 - atan(pos...)^2) * 0.00004
		γ = 1.0
	end

    return (t, x, pos), x, reward, γ
end

function bike_fell(x)
	omega = x[1]
	fallcond = abs(omega) > (π * 12. / 180.) # 12 degrees in rads
	return fallcond
end

function bike_at_goal(pos)
	d = √sum(pos.^2)
	return d ≤ 10  # within 10 meters
end

mutable struct BicyclePlotData <: Any
	ts::Vector{Float64}
	ωs::Vector{Float64}
	xs::Vector{Float64}
	ys::Vector{Float64}

	function BicyclePlotData()
		new(Vector{Float64}(),Vector{Float64}(),Vector{Float64}(),Vector{Float64}())
	end
end

@userplot BicyclePlot
@recipe function f(ap::BicyclePlot)
	data, state, clearplot = ap.args
	t, x, pos = state
	if clearplot
		empty!(data.ts)
		empty!(data.ωs)
		empty!(data.xs)
		empty!(data.ys)
	end

	push!(data.ts, t)
	push!(data.ωs, x[1])
	push!(data.xs, pos[1])
	push!(data.ys, pos[2])

	legend := false
	# xlims := (0., 1.)
	# ylims := (0., 1.)
	grid := true
	ticks := true
	layout := (2,1)
	foreground_color := :white
	# aspect_ratio := 1.

	@series begin
		seriestype := :path
		seriescolor := :blue
		subplot := 1
		# xlims := (0., 30.)
		fall = π * 12.0 / 180.0
		ylims := (-fall, fall)
		data.ts, data.ωs
	end

	@series begin
		seriestype := :path
		seriescolor := :black
		subplot := 2
		# xlims := (-5, 5)
		# ylims := (-3, 10)
		data.xs, data.ys
	end

	# goal
	@series begin
		seriestype := :shape
		linecolor := nothing
		seriescolor := :green
		aspect_ratio := 1.
		fillalpha := 0.4
		subplot := 2
		circle_shape(0.0, 0.0, 10)
	end

	# current pos
	@series begin
		seriestype := :shape
		linecolor := nothing
		seriescolor := :red
		aspect_ratio := 1.

		circle_shape(pos[1], pos[2], 1.0)
	end

end


const E = [
    0 0 1 1 0 0 0 0 0 0 0 0 0 0 1 1 0 0;
    0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0;
    0 1 0 0 0 1 0 0 0 0 0 0 1 0 0 0 1 0;
    0 1 0 0 1 0 1 0 0 0 0 1 0 1 0 0 1 0;
    0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0;
    1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1;
    0 1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1 0;
    0 1 0 0 0 0 0 1 0 0 1 0 0 0 0 0 1 0;
    0 1 0 0 0 0 0 0 1 1 0 0 0 0 0 0 1 0;
    0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0;
    0 0 1 1 0 0 0 0 0 0 0 0 0 0 1 1 0 0;
]

function frown!(E)
    E[7, 7] = E[7, 12] = E[9, 9] = E[9, 10] = 0
    E[9, 7] = E[9, 12] = E[7, 9] = E[7, 10] = 1
    E
end

frown(E) = frown!(copy(E))

rot270(E) = E'

# TODO This should probably be written out since rotr90 is unlike to be the solution a student would come up with
rot90(E) = rotr90(E)

stickerwall(E) = vcat(
    hcat(E, zeros(Int, size(E, 1)), ones(Int, size(E, 1)), zeros(Int, size(E, 1)), vcat(frown(E))),
    hcat(zeros(Int, size(E, 2) + 1)', [1], zeros(Int, size(E, 2) + 1)'),
    ones(Int, 2size(E, 2) + 3)',
    hcat(zeros(Int, size(E, 2) + 1)', [1], zeros(Int, size(E, 2) + 1)'),
    hcat(vcat(frown(E)), zeros(Int, size(E, 1)), ones(Int, size(E, 1)), zeros(Int, size(E, 1)), E)
)


using RecipesBase
using Printf

@recipe function plot(r::Roc{T}; convex_hull=false, traditional=false) where T
    ## defaults
    title --> "ROC"
    xguide --> "Pfa"
    yguide --> (traditional ? "Phit" : "Pmiss")
    legend --> false
    if convex_hull
        x = r.pfa[r.ch]
        y = r.pmiss[r.ch]
    else
        x = r.pfa
        y = r.pmiss
    end
    if traditional
        y = 1 .- y
    end
    (x, y)
end

@userplot DetPlot

@recipe function detplot(d::DetPlot; convex_hull=false)
    if (length(d.args) != 1 || ! (typeof(d.args[1]) <: Roc))
        error("Argument should be of type Roc, got $(length(d.args)) $(typeof(d.args[1]))", )
    end
    r, = d.args
    ##default
    title --> "DET plot"
    xguide --> "Pfa (%)"
    yguide --> "Pmiss (%)"
    legend --> false
    xlims --> (qnorm(0.001), qnorm(0.5))
    ylims --> (qnorm(0.001), qnorm(0.5))
    linewidth --> 2
    ## Not trivial in Plots.jl to control ticks and ticklabels
    ticks --> qnorm.([0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 40] ./ 100)
    formatter --> x -> begin
       p = 100 * pnorm(x)
       if p < 1
           Printf.@sprintf("%3.1f", p)
       else
           Printf.@sprintf("%2.0f", p)
       end
   end
    ## compulsory
    aspect_ratio := :equal

    if convex_hull
        x = qnorm.(r.pfa[r.ch])
        y = qnorm.(r.pmiss[r.ch])
    else
        x = qnorm.(r.pfa)
        y = qnorm.(r.pmiss)
    end
    (x, y)
end

@userplot ApePlot

@recipe function apeplot(ape::ApePlot; xmin = -7, xmax = 7)
    if !(length(ape.args) == 1 && typeof(ape.args[1]) <: Roc)
        error("Expecting one argument of type Roc, got $(typeof(ape.args[1]))")
    end
    r, = ape.args
    title --> "Applied Probability of Error"
    xguide --> "prior log odds"
    yguide --> "Bayes' error rate"
    label --> ["Bayes' Error" "Minimum" "Trivial"]
    linecolor --> [:red :green :black]
    lo = collect(xmin:0.01:xmax)
    be = ber(r, lo)
    mbe = minber(r, lo)
    defbe = 1 ./ ROCAnalysis.normfactor(lo)
    ylims --> (0, 1.1maximum(be))
    x = lo
    y = hcat(be, mbe, defbe)
    (x,y)
end

@userplot NbePlot

@recipe function nbeplot(nbe::NbePlot; xmin = -10, xmax = 5)
    if !(length(nbe.args) == 1 && typeof(nbe.args[1]) <: Roc)
        error("Expecting one argument of type Roc, got $(typeof(nbe.args[1]))")
    end
    r, = nbe.args
    title --> "Normalized Bayes' Error"
    xguide --> "prior log odds"
    yguide --> "normalized Bayes' error"
    label --> ["Bayes' Error" "Minimum BE" "False Alarms" "Misses" "Minimum FA" "Minimum miss"]
    legend --> :left
    linestyle --> [:solid :solid :dash :dot :dash :dot]
    linecolor --> [:red :green :red :red :green :green]
    lo = collect(xmin:0.01:xmax)
    norm = ROCAnalysis.normfactor(lo)
    bfa, bmiss = [be .* norm for be in ROCAnalysis.ber_famiss(r, lo)]
    nbe = bfa + bmiss
    mbfa, mbmiss = [mbe .* norm for mbe in ROCAnalysis.minber_famiss(r, lo)]
    mnbe = minber(r, lo) .* norm
    x = lo
    y = hcat(nbe, mnbe, bfa, bmiss, mbfa, mbmiss)
    (x, y)
end

@userplot LlrPlot

@recipe function llrplot(llr::LlrPlot)
    if !(length(llr.args) == 1 && typeof(llr.args[1]) <: Roc)
        error("Expecting one argument of type Roc, got $(typeof(nbe.args[1]))")
    end
    r, = llr.args
    title --> "LLR plot"
    xguide --> "score"
    yguide --> "log likelihood ratio"
    legend --> false
    llr_extrema = extrema(filter(x -> !isinf(x), r.llr))
    mi = max(llr_extrema[1], minimum(r.θ)) ## max of mins
    ma = min(llr_extrema[2], maximum(r.θ)) ## min of maxs
    xlims --> (mi, ma)
    ylims --> (mi, ma)
    (r.θ, r.llr)
end


using DelimitedFiles

function solve(input,n_generation)
    raw = readdlm(input,'\n')
    initial_state = parse(raw[1])
    spread = Dict{String,String}(map(parse, raw[2:end]))
    state = string("...",initial_state,"...")
    state_sum = Int64(0)
    zero_index=4
    sum_diffs = zeros(Int16,42)
    current_gen = 1
    for generation in 1:n_generation
        next_state = Array{String,1}()
        for pot in 3:length(state)-2
            note = state[pot-2:pot+2]
            if note in keys(spread)
                push!(next_state,spread[note])
            else
                push!(next_state,".")
            end
        end
        next_state = string("...",join(next_state),"...")
        zero_index+=1
        next_state_sum = sum((findall(x->x=='#',next_state) .-zero_index))
        sum_diff = next_state_sum-state_sum
        popfirst!(sum_diffs)
        push!(sum_diffs,sum_diff)
        state = next_state
        state_sum = next_state_sum
        current_gen += 1
        if length(unique(sum_diffs))==1
            break
        end
    end
    if length(unique(sum_diffs))==1
        current_gen -= 1
        state_sum += sum_diffs[1]*(n_generation-current_gen)
    end
    @show state_sum
end

function parse(s)
    m = match(r"(initial state: (?<initialstate>.+)|(?<left>.+) => (?<right>.+))", s)
    initial_state = m[:initialstate]
    if initial_state!= nothing
        return initial_state
    else
        return (m[:left],m[:right])
    end
end

solve("input_test.txt",20)
solve("input.txt",20)
solve("input.txt",50000000000)


using .PyPlot
export plotarea, plotarea_dem

"""
    plotit(x, y, dem)
    plotit(x, y, waterflows_output)

Plot DEM, uparea, flow-dir
"""
plotit(x, y, dem) = plotit(x, y, waterflows(dem))
function plotit(x, y, waterflows_output::Tuple)
    area, slen, dir, nout, nin, pits  = waterflows_output

    fig, axs = subplots(3,1)
    sca(axs[1])
    mycontour(x, y, dem)
    colorbar()
    sca(axs[2])
    heatmap(x, y, log10.(area))
    sca(axs[3])
    plotdir(x, y, dir)
end

function plotdir(x, y, dir)
    vecfield = dir2vec.(dir)
    vecfieldx = [v[1] for v in vecfield]
    vecfieldy = [v[2] for v in vecfield]
    quiver(repeat(x,1, length(y)), repeat(y,length(x),1), vecfieldx, vecfieldy)
end

pits2inds(pits) = ([p.I[1] for p in pits],
                   [p.I[2] for p in pits])
pits2vecs(x, y, pits) = (x[[p.I[1] for p in pits if p!=CartesianIndex(-1,-1)]],
                           y[[p.I[2] for p in pits if p!=CartesianIndex(-1,-1)]])

"""
    plotarea(x, y, area, pits; prefn=log10)

Plot uparea, or another variable.  If no pits should be plotted, use `[]`.
"""
function plotarea(x, y, area, pits; prefn=log10, cbar=true)
    px, py = pits2vecs(x, y, pits)
    fig, axs = subplots()
    heatmap(x, y, prefn.(area), cbar=cbar)
    scatter(px, py, 1, "r")
    fig
end

function plotarea_dem(x, y, dem, area, pits; levels=50, threshold=1/100, prefn=log10)
    fig, axs = subplots()
    mycontour(x, y, dem, levels=levels)

    threshold *= prod(size(area))
    area = copy(area)
    area[area.<threshold] .= NaN
    px, py = pits2vecs(x, y, pits)

    heatmap(x, y, prefn.(area))
    scatter(px, py, 1, "r")
end

function plotbnds(x,y,bnds)
    for b in bnds
        px, py = pits2vecs(x, y, b)
        scatter(px, py, 1, "g")
    end
end

function plotlakedepth(x, y, dem)
    area, slen, dir, nout, nin, pits  = waterflows(dem)
    demf = fill_dem(dem, pits, dir)
    heatmap(x, y, demf.-dem)
end

function heatmap(x, y, mat; cbar=true)
    dx2 = (x[2]-x[1])/2
    imshow(Array(mat'), origin="lower", extent=(x[1]-dx2,x[end]+dx2,y[1]-dx2,y[end]+dx2))
    cbar && colorbar()
end

function mycontour(x, y, mat; levels=50, cbar=true)
    contour(x, y, Array(mat'), levels, linewidths=1)
    cbar && colorbar()
end

function plot_catchments(x, y, c, lim=10)
    c = copy(c)
    tokeep = []
    for cc =1:maximum(c)
        inds = c.==cc
        if sum(inds)<lim
            c[inds] .= 0
        else
            push!(tokeep, inds)
        end
    end
    for (i,inds) in enumerate(tokeep)
        c[inds] .= i
    end
    heatmap(x,y,c)
end


# Tests for Revenue Additive DEA Models
@testset "RevenueAdditiveDEAModel" begin

    # Test using Book data
    X = [1; 1; 1; 1; 1; 1; 1; 1];
    Y = [7 7; 4 8; 8 4; 3 5; 3 3; 8 2; 6 4; 1.5 5];
    P = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1];

    # With :Ones
    dearevenueaddones = dearevenueadd(X, Y, P, :Ones)

    @test typeof(dearevenueaddones) == RevenueAdditiveDEAModel

    @test nobs(dearevenueaddones) == 8
    @test ninputs(dearevenueaddones) == 1
    @test noutputs(dearevenueaddones) == 2
    @test ismonetary(dearevenueaddones) == false

    @test efficiency(dearevenueaddones) == efficiency(dearevenueaddones, :Economic)
    @test efficiency(dearevenueaddones, :Economic)   ≈ [0.0; 2.0; 2.0; 6.0; 8.0; 4.0; 4.0; 7.5]
    @test efficiency(dearevenueaddones, :Technical)  ≈ [0.0; 0.0; 0.0; 6.0; 8.0; 2.0; 4.0; 7.5]
    @test efficiency(dearevenueaddones, :Allocative) ≈ [0.0; 2.0; 2.0; 0.0; 0.0; 2.0; 0.0; 0.0]
    @test normfactor(dearevenueaddones) ≈ [1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]

    # With :Ones and CRS (Same result in this example)
    dearevenueaddonescrs = dearevenueadd(X, Y, P, :Ones, rts = :CRS)

    @test efficiency(dearevenueaddonescrs, :Economic)   ≈ [0.0; 2.0; 2.0; 6.0; 8.0; 4.0; 4.0; 7.5]
    @test efficiency(dearevenueaddonescrs, :Technical)  ≈ [0.0; 0.0; 0.0; 6.0; 8.0; 2.0; 4.0; 7.5]
    @test efficiency(dearevenueaddonescrs, :Allocative) ≈ [0.0; 2.0; 2.0; 0.0; 0.0; 2.0; 0.0; 0.0]

    # Default is :Ones
    @test efficiency(dearevenueadd(X, Y, P)) == efficiency(dearevenueaddones)

    # With :MIP
    dearevenueaddmip = dearevenueadd(X, Y, P, :MIP)

    @test efficiency(dearevenueaddmip, :Economic)   ≈ [0.0; 0.5; 0.5; 2.0; 8/3; 2.0; 1.0; 5.0] atol = 1e-5
    @test efficiency(dearevenueaddmip, :Technical)  ≈ [0.0; 0.0; 0.0; 26/15; 8/3; 1.0; 11/12; 37/9] atol = 1e-5
    @test efficiency(dearevenueaddmip, :Allocative) ≈ [0.0; 0.5; 0.5; 4/15; 0.0; 1.0; 1/12; 8/9] atol = 1e-5
    @test normfactor(dearevenueaddmip) ≈ [7.0; 4.0; 4.0; 3.0; 3.0; 2.0; 4.0; 1.5]

    # With :MIP and monetary = true
    dearevenueaddmipmon = dearevenueadd(X, Y, P, :MIP, monetary = true)

    @test ismonetary(dearevenueaddmipmon) == true

    @test efficiency(dearevenueaddmipmon) ≈ efficiency(dearevenueaddmip) .* normfactor(dearevenueaddmip)
    @test efficiency(dearevenueaddmipmon, :Technical) ≈ efficiency(dearevenueaddmip, :Technical) .* normfactor(dearevenueaddmip)
    @test efficiency(dearevenueaddmipmon, :Allocative) ≈ efficiency(dearevenueaddmip, :Allocative) .* normfactor(dearevenueaddmip)

    # With custom weights
    dearevenueaddcustom = dearevenueadd(X, Y, P, rhoY = 1 ./ Y)

    @test efficiency(dearevenueaddcustom) ≈ efficiency(dearevenueaddmip)

    # Print
    show(IOBuffer(), dearevenueaddones)
    show(IOBuffer(), dearevenueadd(X, Y, P, :Ones, dispos = :Weak))

    # Test errors
    @test_throws DimensionMismatch dearevenueadd([1; 2 ; 3], [4 ; 5], [1; 1; 1] , :Ones) #  Different number of observations
    @test_throws DimensionMismatch dearevenueadd([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], :Ones) # Different number of observation in prices
    @test_throws DimensionMismatch dearevenueadd([1 1; 2 2; 3 3 ], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], :Ones) # Different number of output prices and Outputs
    @test_throws DimensionMismatch dearevenueadd([1; 2; 3] , [4; 5; 6], [1; 1; 1], rhoY = [1 1; 2 2; 3 3]) # Different size of output weights
    @test_throws ArgumentError dearevenueadd([1; 2; 3], [4; 5; 6], [1; 2; 3], :Ones, rts = :Error) # Invalid returns to scale
    @test_throws ArgumentError dearevenueadd([1; 2; 3], [1; 2; 3], [1; 1; 1], :Error) # Invalid inuts direction
    @test_throws ArgumentError dearevenueadd([1; 2; 3], [1; 2; 3], [1; 1; 1], :Ones, dispos = :Error) # Invalid inputs disposal
    @test_throws ArgumentError dearevenueadd([1; 2; 3], [1; 2; 3], [1; 1; 1], :Ones, rhoY= [1; 1; 1]) # Weights not allowed if model != :Custom

end


# Load Julia packages (libraries) needed.

using StatisticalRethinking
ProjDir = @__DIR__

# ### snippet 5.29

println()
df = DataFrame!(CSV.File(rel_path("..", "data", "milk.csv"), delim=';'));
df = filter(row -> !(row[:neocortex_perc] == "NA"), df);
df[!, :neocortex_perc] = parse.(Float64, df[:, :neocortex_perc])
df[!, :lmass] = log.(df[:, :mass])
#first(df, 5) |> display

# ### snippet 5.1

scale!(df, [:kcal_per_g, :neocortex_perc, :lmass])
println()

m_5_5 = "
data {
 int < lower = 1 > N; // Sample size
 vector[N] K; // Outcome
 vector[N] NC; // Predictor
}

parameters {
 real a; // Intercept
 real bN; // Slope (regression coefficients)
 real < lower = 0 > sigma;    // Error SD
}

model {
  vector[N] mu;               // mu is a vector
  a ~ normal(0, 0.2);           //Priors
  bN ~ normal(0, 0.5);
  sigma ~ exponential(1);
  mu = a + bN * NC;
  K ~ normal(mu , sigma);     // Likelihood
}
";

# Define the SampleModel and set the output format to :mcmcchains.

m5_5s = SampleModel("m5.5", m_5_5);

# Input data for cmdstan

m5_5_data = Dict("N" => size(df, 1), "NC" => df[!, :neocortex_perc_s],
    "K" => df[!, :kcal_per_g_s]);

# Sample using StanSample

rc = stan_sample(m5_5s, data=m5_5_data);

if success(rc)

  # Describe the draws

  dfa5 = read_samples(m5_5s; output_format=:dataframe)
  p = Particles(dfa5)
  quap(dfa5) |> display

end


# This definitely needs cleaning
using OrdinaryDiffEq, DiffEqDevTools, Test, Random
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_linear, prob_ode_2Dlinear
probArr = Vector{ODEProblem}(undef, 2)
probArr[1] = prob_ode_linear

probArr[2] = prob_ode_2Dlinear
Random.seed!(100)
## Convergence Testing
dts = 1 .//2 .^(8:-1:4)
dts1 = 1 .//2 .^(9:-1:5)
testTol = 0.2

for i = 1:2
  global dts
  prob = probArr[i]
  sim = test_convergence(dts,prob,Euler())
  @test abs(sim.𝒪est[:final]-1) < testTol
  sim2 = test_convergence(dts,prob,Heun())
  @test abs(sim2.𝒪est[:l∞]-2) < testTol
  sim2 = test_convergence(dts,prob,Ralston())
  @test abs(sim2.𝒪est[:l∞]-2) < testTol
  sim2 = test_convergence(dts,prob,Midpoint())
  @test abs(sim2.𝒪est[:l∞]-2) < testTol
  sim3 = test_convergence(dts,prob,RK4())
  @test abs(sim3.𝒪est[:l∞]-4) < testTol
  sim4 = test_convergence(dts,prob,BS3())
  @test abs(sim4.𝒪est[:l2]-3) < testTol
  alg = CarpenterKennedy2N54()
  @test abs(test_convergence(dts,prob,alg).𝒪est[:l∞]-OrdinaryDiffEq.alg_order(alg)) < testTol
  sim5 = test_convergence(dts, prob, AB3())
  @test abs(sim5.𝒪est[:l2]-3) < testTol
  sim6 = test_convergence(dts,prob,ABM32())
  @test abs(sim6.𝒪est[:l2]-3) < testTol
  sim7 = test_convergence(dts, prob, AB4())
  @test abs(sim7.𝒪est[:l2]-4) < testTol
  sim8 = test_convergence(dts1,prob,ABM43())  #using dts1 due to floating point error in convergence test
  @test abs(sim8.𝒪est[:l2]-4) < testTol
  sim9 = test_convergence(dts,prob,AB5())
  @test abs(sim9.𝒪est[:l2]-5) < testTol
  sim10 = test_convergence(dts,prob,ABM54())
  @test abs(sim10.𝒪est[:l2]-5) < testTol
  sim101 = test_convergence(dts,prob,VCAB3())
  @test abs(sim101.𝒪est[:l2]-3) < testTol
  sim102 = test_convergence(dts,prob,VCAB4())
  @test abs(sim102.𝒪est[:l2]-4) < testTol
  sim103 = test_convergence(dts,prob,VCAB5())
  @test abs(sim103.𝒪est[:l2]-5) < testTol
  sim104 = test_convergence(dts,prob,VCABM3())
  @test abs(sim104.𝒪est[:l2]-3) < testTol
  sim105 = test_convergence(dts,prob,VCABM4())
  @test abs(sim105.𝒪est[:l2]-4) < testTol
  sim106 = test_convergence(dts,prob,VCABM5())
  @test abs(sim106.𝒪est[:l2]-5) < testTol
  sim160 = test_convergence(dts,prob,Anas5(w=2))
  @test abs(sim160.𝒪est[:l2]-4) < 2*testTol

  println("Stiff Solvers")

  dts = 1 .//2 .^(9:-1:5)

  sim11 = test_convergence(dts,prob,ImplicitEuler(extrapolant = :linear))
  @test abs(sim11.𝒪est[:final]-1) < testTol

  sim112 = test_convergence(dts,prob,ImplicitEuler(nonlinsolve = NLFunctional()))
  @test abs(sim112.𝒪est[:final]-1) < testTol

  sim12 = test_convergence(dts,prob,
          GenericImplicitEuler(nlsolve=OrdinaryDiffEq.NLSOLVEJL_SETUP(autodiff=true)))
  @test abs(sim12.𝒪est[:final]-1) < testTol

  sim122 = test_convergence(dts,prob,
           GenericImplicitEuler(nlsolve=OrdinaryDiffEq.NLSOLVEJL_SETUP(autodiff=false)))

  sim13 = test_convergence(dts,prob,ImplicitMidpoint())
  @test abs(sim13.𝒪est[:final]-2) < testTol

  sim132 = test_convergence(dts,prob,ImplicitMidpoint(nonlinsolve = NLFunctional()))
  @test abs(sim132.𝒪est[:final]-2) < testTol

  sim13 = test_convergence(dts,prob,Trapezoid())
  @test abs(sim13.𝒪est[:final]-2) < testTol

  sim133 = test_convergence(dts,prob,Trapezoid(nonlinsolve = NLFunctional()))
  @test abs(sim133.𝒪est[:final]-2) < testTol

  sim14 = test_convergence(dts,prob,
          GenericTrapezoid(nlsolve=OrdinaryDiffEq.NLSOLVEJL_SETUP(autodiff=true)))
  @test abs(sim14.𝒪est[:final]-2) < testTol

  sim142 = test_convergence(dts,prob,
           GenericTrapezoid(nlsolve=OrdinaryDiffEq.NLSOLVEJL_SETUP(autodiff=false)))
  @test abs(sim142.𝒪est[:final]-2) < testTol

  sim14 = test_convergence(dts,prob,TRBDF2())
  @test abs(sim14.𝒪est[:final]-2) < testTol

  sim152 = test_convergence(dts,prob,TRBDF2(autodiff=false))
  @test abs(sim152.𝒪est[:final]-2) < testTol+0.1

  sim15 = test_convergence(dts,prob,SDIRK2())
  @test abs(sim15.𝒪est[:final]-2) < testTol

  sim152 = test_convergence(dts,prob,SSPSDIRK2())
  @test abs(sim152.𝒪est[:final]-2) < testTol

  sim16 = test_convergence(dts,prob,Kvaerno3())
  @test abs(sim16.𝒪est[:final]-3) < testTol

  sim162 = test_convergence(dts,prob,Kvaerno3(nonlinsolve = NLFunctional()))
  @test abs(sim162.𝒪est[:final]-3) < testTol

  sim17 = test_convergence(dts,prob,KenCarp3())
  @test abs(sim17.𝒪est[:final]-3) < testTol

  #####################################
  # BDF
  #####################################

  sim = test_convergence(dts,prob,ABDF2(extrapolant = :linear))
  @test abs(sim.𝒪est[:final]-2) < testTol
  @test abs(sim.𝒪est[:l2]-2) < testTol
  @test abs(sim.𝒪est[:l∞]-2) < testTol

  sim = test_convergence(dts,prob,ABDF2(nonlinsolve = NLFunctional()))
  @test abs(sim.𝒪est[:final]-2) < testTol
  @test abs(sim.𝒪est[:l2]-2) < testTol
  @test abs(sim.𝒪est[:l∞]-2) < testTol

  # QBDF
  sim = test_convergence(dts,prob,QBDF1())
  @test abs(sim.𝒪est[:final]-1) < testTol
  @test abs(sim.𝒪est[:l2]-1) < testTol
  @test abs(sim.𝒪est[:l∞]-1) < testTol

  sim = test_convergence(dts,prob,QBDF2())
  @test abs(sim.𝒪est[:final]-2) < testTol
  @test abs(sim.𝒪est[:l2]-2) < testTol
  @test abs(sim.𝒪est[:l∞]-2) < testTol

  # QNDF
  sim = test_convergence(dts,prob,QNDF1())
  @test abs(sim.𝒪est[:final]-1) < testTol
  @test abs(sim.𝒪est[:l2]-1) < testTol
  @test abs(sim.𝒪est[:l∞]-1) < testTol

  sim = test_convergence(dts,prob,QNDF2())
  @test abs(sim.𝒪est[:final]-2) < testTol
  @test abs(sim.𝒪est[:l2]-2) < testTol
  @test abs(sim.𝒪est[:l∞]-2) < testTol

  sim = test_convergence(dts,prob,QNDF2(nonlinsolve = NLFunctional()))
  @test abs(sim.𝒪est[:final]-2) < testTol
  @test abs(sim.𝒪est[:l2]-2) < testTol
  @test abs(sim.𝒪est[:l∞]-2) < testTol

  dts = 1 .//2 .^(7:-1:4)
  println("Higher Order")

  sim18 = test_convergence(dts,prob,Cash4())
  @test abs(sim18.𝒪est[:final]-4) < testTol

  sim19 = test_convergence(dts,prob,Hairer4())
  @test abs(sim19.𝒪est[:final]-4) < testTol

  sim110 = test_convergence(dts,prob,Hairer42())
  @test abs(sim110.𝒪est[:final]-4) < testTol

  sim111 = test_convergence(dts,prob,Kvaerno4())
  @test abs(sim111.𝒪est[:final]-4) < testTol

  sim112 = test_convergence(dts,prob,KenCarp4())
  @test abs(sim112.𝒪est[:final]-4) < testTol

  sim113 = test_convergence(dts,prob,Kvaerno5())
  @test abs(sim113.𝒪est[:final]-5) < testTol

  sim114 = test_convergence(dts,prob,KenCarp5())
  @test abs(sim114.𝒪est[:final]-5) < testTol

  sim115 = test_convergence(dts,prob,KenCarp5(nonlinsolve = NLFunctional()))
  @test_broken abs(sim115.𝒪est[:final]-5) < testTol
end


using jInv.InverseSolve
using Test
using jInv.Mesh
using jInv.Utils

# build regular mesh and Iact
domain = [0;5;0;5; 0;3]
n     = [12;13;10]
M     = getRegularMesh(domain,n)
idx   = ones(Int,(n[1],n[2],n[3]))
idx[:,:,1:5] .= 0
Iact  = sparse(I,M.nc,M.nc)
Iact  = Iact[:,vec(idx).==1]
mc    = randn(size(Iact,2))

regFuns = [
			(m,mref,M)->diffusionReg(m,mref,M,Iact=Iact),
			(m,mref,M)->wdiffusionReg(m,mref,M,Iact=Iact),
			smallnessReg,
			(m,mref,M)->wTVReg(m,mref,M,Iact=Iact),
			(m,mref,M)->anisoTVReg(m,mref,M,Iact=Iact)]
for k=1:length(regFuns)
	# checkDerivative of $(regFuns[k])

	function testFun(x,v=[])
		Sc,dS,d2S = regFuns[k](x,0.0.*x,M)
		if isempty(v)
			return Sc
		else
			return Sc,dot(dS,v)
		end
	end
	chkDer, = checkDerivative(testFun,mc,out=false)
	@test chkDer
end

regFuns = [wdiffusionRegNodal;wTVRegNodal]
idx   = ones(Int,(n[1]+1,n[2]+1,n[3]+1))
idx[:,:,1:5] .= 0
Iact  = sparse(I,prod(M.n.+1),prod(M.n.+1))
Iact  = Iact[:,vec(idx).==1]
mc    = randn(size(Iact,2))

for k=1:length(regFuns)
	# checkDerivative of $(regFuns[k])

	function testFun(x,v=[])
		Sc,dS,d2S = regFuns[k](x,0.0.*x,M,Iact=Iact)
		if isempty(v)
			return Sc
		else
			return Sc,dot(dS,v)
		end
	end
	chkDer, = checkDerivative(testFun,mc,out=false)
	@test chkDer
end

regFuns = [logBarrier;logBarrierSquared]
idx   = ones(Int,(n[1]+1,n[2]+1,n[3]+1))
idx[:,:,1:5] .= 0
Iact  = sparse(I,prod(M.n.+1),prod(M.n.+1))
Iact  = Iact[:,vec(idx).==1]
mc    = ones(size(Iact,2)) + 0.05*randn(size(Iact,2));

low   = ones(length(mc))*0.4;
high  = ones(length(mc))*1.6;
epsilon = ones(length(mc))*0.58;

for k=1:length(regFuns)
	# checkDerivative of $(regFuns[k])

	function testFun(x,v=[])
		Sc,dS,d2S = regFuns[k](x,0.0.*x,M,low,high,epsilon)
		if isempty(v)
			return Sc
		else
			return Sc,dot(dS,v)
		end
	end
	chkDer, = checkDerivative(testFun,mc,out=false)
	@test chkDer
end


#=
See also:
https://github.com/FluxML/Flux.jl/blob/master/src/onehot.jl
https://github.com/cossio/OneHot.jl
=#

"""
    onehot_encode(A, code)

Given an array `A` of `N` dimensions, returns a one-hot encoded `BitArray` of
`N + 1` dimensions where single entries of the first dimension are one.
"""
function onehot_encode(A::AbstractArray, code = sort(unique(A)))
    return reshape(A, 1, size(A)...) .== code
end

"""
    onehot_decode(X)

Given a onehot encoded array `X` of `N + 1` dimensions, returns the
equivalent categorical array of `N` dimensions.
"""
function onehot_decode(X::AbstractArray)
    idx = dropdims(argmax(X; dims=1); dims=1)
    return first.(Tuple.(idx))
end

"""
    categorical_sample_from_logits(logits)

Given a logits array `logits` of size `(q, *)` (where `q` is the number of classes),
returns an array `X` of size `(*)`, such that `X[i]` is a categorical random sample
from the distribution with logits `logits[:,i]`.
"""
function categorical_sample_from_logits(logits::AbstractArray)
    p = softmax(logits; dims=1)
    return categorical_sample(p)
end

"""
    categorical_sample_from_logits_gumbel(logits)

Like categorical_sample_from_logits, but using the Gumbel trick.
"""
function categorical_sample_from_logits_gumbel(logits::AbstractArray)
    z = logits .+ randgumbel.(eltype(logits))
    idx = dropdims(argmax(z; dims=1); dims=1)
    c = first.(Tuple.(idx))
    return c
end

"""
    categorical_sample(P)

Given a probability array `P` of size `(q, *)`, returns an array
`C` of size `(*)`, such that `C[i] ∈ 1:q` is a random sample from the
categorical distribution `P[:,i]`.
You must ensure that `P` defines a proper probability distribution.
"""
function categorical_sample(P::AbstractArray)
    idx = CartesianIndices(tail(size(P)))
    C = Array{Int}(undef, tail(size(P)))
    @inbounds for i in idx
        ps = @view P[:,i]
        C[i] = categorical_rand(ps)
    end
    return C
end

"""
    categorical_rand(ps)

Randomly draw `i` with probability `ps[i]`.
You must ensure that `ps` defines a proper probability distribution.
"""
function categorical_rand(ps::AbstractVector)
    i = 0
    u = rand(eltype(ps))
    for p in ps
        u -= p
        i += 1
        u ≤ 0 && break
    end
    return i
end

"""
    randgumbel(T = Float64)

Generates a random Gumbel variate.
"""
randgumbel(::Type{T} = Float64) where {T} = -log(randexp(T))


# ---
# title: 977. Squares of a Sorted Array
# id: problem977
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array, Two Pointers
# link: <https://leetcode.com/problems/squares-of-a-sorted-array/description/>
# hidden: true
# ---
#
# Given an integer array `nums` sorted in **non-decreasing** order, return _an
# array of **the squares of each number**  sorted in non-decreasing order_.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [-4,-1,0,3,10]
#     Output: [0,1,9,16,100]
#     Explanation: After squaring, the array becomes [16,1,0,9,100].
#     After sorting, it becomes [0,1,9,16,100].
#
#
# **Example 2:**
#
#
#
#     Input: nums = [-7,-3,2,3,11]
#     Output: [4,9,9,49,121]
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 104`
#   * `-104 <= nums[i] <= 104`
#   * `nums` is sorted in **non-decreasing** order.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#----------------------------------------------------------------------------------------------------
# Solve elastic net covariance estimation problem using subgradient method
#----------------------------------------------------------------------------------------------------
include("func.jl");
using Printf

function solve_enet(X0::Array{Float64,2},
                                Xtrue::Array{Float64,2},
                                stochErr::Float64,
                                maxIter::Int64,
                                stepSizes::Array{Float64,1};
                                method::String="subgradient",
                                verbose::Bool=true,
                                noiseType::String="gaussian"
                                )

    # Basic data
    (d,r) = size(X0);
    XTtrue = Xtrue';
    sqnrmXtrue = sum(abs2, Xtrue);

    #   Initialize
    X = copy(X0);
    a = zeros(d);
    XTa = zeros(r);
    b = 0;
    res = 0;
    G = zeros(d,r);
    η = 0;

    if method=="mirror"
        # Coefficients for the Bregman divergence polynomials p and Φ
        #       p(u) = a0 + a1 * u + a2 * u² + a3 * u³
        #       Φ(x) = c0 ||x||₂² + c1 ||x||₂³ + c2 ||x||₂⁴ + c3 ||x||₂⁵
        (a0, a1, a2, a3) = (1.0, 0.0, 1.0, 1.0);
        (c0, c1, c2, c3) = (a0*7/2,  a1*10/3,  a2*13/454, a3*16/5);

        V = zeros(d,r);
        λ = 0;
    end

    # for keeping track of progress
    err = NaN;
    err_hist =  fill(NaN, maxIter);  # to keep track of errors
    fun_val = NaN;
    fun_hist =  fill(NaN, maxIter);  # to keep track of function values (approximates)

    # draw the stochastic a, b
    (A,B) = get_ab(XTtrue, stochErr, maxIter, noiseType=noiseType)
    true_empirical_value = compute_empirical_function(Xtrue, A, B);

    #   Run subgradient method
    for k=1:maxIter
        # get stochastic a, b
        a = A[:,k];
        b = B[k];
        XTa = X'*a;

        # update subgradient - in place
        enet_subgrad!(G, XTa, a, b);

        if method=="subgradient"
            # update X - in place
            η = stepSizes[k];
            BLAS.axpy!(-η,G,X);
        elseif method=="mirror"
            # update V = ∇Φ(X) - η * G
            η = stepSizes[k];
            sqnrmX = sum(abs2, X);
            V = ( 2*c0 + 3*c1*sqrt(sqnrmX) + 4*c2*sqnrmX + 5*c3*sqnrmX^(3/2) ) * X  - η*G;
            nrmV = norm(V,2);

            # root finding problem to find λ = norm of X
            #           (2*c0) * λ + (3*c1) * λ^2 + (4*c2) * λ^3  + (5*c3) * λ^4 = nrmV
            # (note that direction of X is identical to direction of V)
            λ = get_root([-nrmV, 2*c0, 3*c1, 4*c2, 5*c3]);

            # update X - in place
            lmul!(0.0, X)  # reset to zero
            BLAS.axpy!(λ/nrmV,V,X)
        end

        # record error status and print to console
        err = sqnrmXtrue + sum(abs2, X) - 2 * sum(svdvals(XTtrue * X));
        normalized_err = err / sqnrmXtrue;
        err_hist[k] = normalized_err;

        fun_val = compute_enet_empirical_function(X, A, B);
        value_error = fun_val - true_empirical_value ;
        fun_hist[k] = value_error;

        if verbose
            @printf("iter %3d: emp val = %1.2e, error = %1.2e, stepsize = %1.2e\n",
                                k, value_error, normalized_err, η);
        end
    end

    return (err_hist, fun_hist)
end


# Author: Lukasz Bratos

include("diff.jl")

f(x) = abs(x)
g(x) = 1.0 / (1.0 + x^2)

for i in [5, 10, 15]
    Diff.rysujNnfx(f, -1.0, 1.0, i, "plot6a$i")
    Diff.rysujNnfx(g, -5.0, 5.0, i, "plot6b$i")
end


using RobustDMD, Random, Printf

# Generate DMD Synthetic Data
#------------------------------------------------------------------------------
# dimensions
m = 500;    # temporal dimension
n = 1000;    # spatial dimension
k = 3;      # number of modes
T = Float64;
sigman = 1e-3;
# generate data
X, t, at, Bt = simDMD(m, n, k, T; seed=123456);
Xdat = X + sigman*Random.randn!(copy(X))
# specify the loss
function lossFunc(r)
	return sum(abs2, r)
end
#
function lossGrad(r)
	conj!(r)
end
#
# Create object
#------------------------------------------------------------------------------

#
# Initial Guess
#------------------------------------------------------------------------------
# create initial guess using exact dmd
a0, B0 = dmdexactestimate(m, n, k, Xdat, t);

#
# Create Object
#------------------------------------------------------------------------------
params = DMDParams(k, Xdat, t, lossFunc, lossGrad, inner_directl2=true);
 copyto!(params.a, a0);
# copyto!(params.B, B0);
#Random.randn!(params.ar);
#
# Apply Solver
#------------------------------------------------------------------------------
res = solveDMD_withOptimBFGS(params);

@printf "%7.4e relative residual after BFGS\n" sqrt(dmdobjective(params)/lossFunc(X))

params2 = DMDParams(k,Xdat,t,lossFunc,lossGrad);
copyto!(params2.a,at)
@printf "%7.4e relative residual with true evals\n" sqrt(dmdobjective(params2)/lossFunc(X))
copyto!(params2.a,a0)
@printf "%7.4e relative residual with initial guess\n" sqrt(dmdobjective(params2)/lossFunc(X))

@printf "%7.4e abs err evals after BFGS\n" besterrperm(params.a,at)
@printf "%7.4e abs err evals initial guess\n" besterrperm(a0,at)


module NelderMead

include("Vertexes.jl")
include("Simplexes.jl")

"""
    optimise(f, initial_vertex_positions; kwargs...)

Find minimum of function, `f`, first creating a Simplex from vertices at
`initial_vertex_positions`, and options passed in via kwargs.
"""
function optimise(f::T, initial_vertex_positions::U; kwargs...
   ) where {T<:Function, W<:Number, V<:AbstractVector{W}, U<:AbstractVector{V}}
  return optimise!(Simplex(f, initial_vertex_positions), f; kwargs...)
end

"""
    optimise(f, initial_positions, initial_step; kwargs...)

Find minimum of function, `f`, first creating a Simplex using a starting
vertex position, `initial_position`, and other vertices `initial_step` away from
that point in all directions, and options passed in via kwargs.
"""
function optimise(f::T, initial_position::AbstractVector{U},
               initial_step::AbstractVector{V}; kwargs...
               ) where {T<:Function, U<:Number, V<:Number}
  return optimise!(Simplex(f, initial_position, initial_step), f; kwargs...)
end

"""
    optimise(f, lower, upper, gridsize; kwargs...)

Find minimum of function, `f`, first creating a `2 .* prod(gridsize)` Simplices
within a hyper-rectangle spanning from `lower` to `upper`, and
options passed in via kwargs.
"""
function optimise(f::F, lower::AbstractVector{T}, upper::AbstractVector{T},
    gridsize::AbstractVector{<:Integer}; kwargs...) where {F<:Function, T<:Number}

  all(gridsize .> 0) || throw(ArgumentError("gridsize, $gridsize, must .> 0"))
  if !(length(lower) == length(upper) == length(gridsize))
    throw(ArgumentError("The lengths of lower $lower, upper $upper, and
                        gridsize $gridsize must be the same"))
  end
  if !all(lower .< upper)
    throw(ArgumentError("lower $lower must .< upper $upper"))
  end

  dim = length(lower)
  index2position(i) = (i .- 1) ./ gridsize .* (upper .- lower) .+ lower
  index2values = Dict()
  totaltime = @elapsed for ii ∈ CartesianIndices(Tuple(gridsize .+ 1))
    index = collect(Tuple(ii))
    x = index2position(index)
    index2values[index] = f(x)
  end
  U = typeof(first(index2values)[2])
  simplices = Set{Simplex{T, U}}()
  function generatesimplices!(simplices, direction)
    for ii ∈ CartesianIndices(Tuple(((gridsize .+ 1) .* ones(Int, dim))))
      vertices = Vector{Vertex{T, U}}()
      index = collect(Tuple(ii))
      for i ∈ 1:dim + 1
        vertexindex = [index[j] + ((j == i) ? direction : 0) for j ∈ 1:dim]
        all(1 .<= vertexindex .<= gridsize .+ 1) || continue
        vertex = Vertex{T, U}(index2position(vertexindex),
                              index2values[vertexindex])
        push!(vertices, vertex)
      end
      length(vertices) != dim + 1 && continue
      push!(simplices, Simplex(vertices))
    end
  end

  totaltime += @elapsed generatesimplices!(simplices, 1)
  totaltime += @elapsed generatesimplices!(simplices, -1)
  @assert length(simplices) == 2 * prod(gridsize)
  totaltime += @elapsed solutions = map(s->optimise!(s, f; kwargs...), collect(simplices))
  return solutions
end

"""
    optimise!(s, f; kwargs...)

Find minimum of function, `f`, starting from Simplex, `s`, with options
passed in via kwargs.

# Keyword Arguments
-  stopval (default sqrt(eps())): stopping criterion when function evaluates
equal to or less than stopval
-  xtol_abs (default zeros(T)) .* ones(Bool, dimensionality(s)): stop if
the vertices of simplex get within this absolute tolerance
-  xtol_rel (default eps(T)) .* ones(Bool, dimensionality(s)): stop if
the vertices of simplex get within this relative tolerance
-  ftol_abs (default zero(real(U))): stop if function evaluations at the
vertices are close to one another by this absolute tolerance
-  ftol_rel (default 1000eps(real(U))): stop if function evaluations at the
vertices are close to one another by this relative tolerance
-  maxiters (default 1000): maximum number of iterations of the Nelder Mead
algorithm
-  timelimit (default Inf): stop if it takes longer than this in seconds
-  α (default 1): Reflection factor
-  β (default 0.5): Contraction factor
-  γ (default 2): Expansion factor
-  δ (default 0.5): Shrinkage factor
"""
function optimise!(s::Simplex{T,U}, f::F; kwargs...) where {F<:Function, T<:Real, U}
  kwargs = Dict(kwargs)
  stopval = get(kwargs, :stopval, sqrt(eps()))
  xtol_abs = get(kwargs, :xtol_abs, zeros(T)) .* ones(Bool, dimensionality(s))
  xtol_rel = get(kwargs, :xtol_rel, eps(T)) .* ones(Bool, dimensionality(s))
  ftol_abs = get(kwargs, :ftol_abs, zero(real(U)))
  ftol_rel = get(kwargs, :ftol_rel, 1000eps(real(U)))
  maxiters = get(kwargs, :maxiters, 1000)
  timelimit = get(kwargs, :timelimit, Inf)
  α = get(kwargs, :α, 1)
  β = get(kwargs, :β, 0.5)
  γ = get(kwargs, :γ, 2)
  δ = get(kwargs, :δ, 0.5)

  @assert α >= 0 "$α >= 0"
  @assert 0 <= β < 1 "0 <= $β < 1"
  @assert γ > 1 "$γ > 1"
  @assert γ > α "$γ > $α"

  reflect(this, other) = Vertex(newposition(this, α, other), f)
  expand(this, other) = Vertex(newposition(this, -γ, other), f)
  contract(this, other) = Vertex(newposition(this, -β, other), f)
  shrink(this, other) = Vertex(newposition(this, δ, other), f)

  function shrink!(s::Simplex)
    lengthbefore = length(s)
    best = bestvertex(s)
    newvertices = [shrink(best, v) for v ∈ s if !isequal(v, best)]
    remove!(s, findall(v->!isequal(v, best), s.vertices))
    push!(s, newvertices...)
    sort!(s, by=v->value(v))
    @assert length(s) == lengthbefore
    return nothing
  end

  iters, totaltime = 0, 0.0
  returncode = assessconvergence(s,
    xtol_abs, xtol_rel, ftol_abs, ftol_rel, stopval)
  history = deepcopy(s.vertices)
  while returncode == :CONTINUE && iters < maxiters && totaltime < timelimit
    totaltime += @elapsed begin
      iters += 1
      best = bestvertex(s)
      worst = worstvertex(s)
      secondworst = secondworstvertex(s)
      reflected = reflect(centroidposition(s), worst)

      if any(h->isequal(h, reflected), history)
        returncode = :ENDLESS_LOOP
        break
      end
      history .= circshift(history, 1)
      history[1] = reflected

      if best <= reflected < secondworst
        swapworst!(s, reflected)
      else
        centroid = Vertex(centroidposition(s), f)
        if reflected < best
          expanded = expand(centroid, reflected)
          expanded < reflected && swapworst!(s, expanded)
          expanded >= reflected && swapworst!(s, reflected)
        elseif secondworst <= reflected < worst
          contracted = contract(centroid, reflected)
          contracted <= reflected ? swapworst!(s, contracted) : shrink!(s)
        elseif reflected >= worst
          contracted = contract(centroid, worst)
          contracted < worst ? swapworst!(s, contracted) : shrink!(s)
        end
      end
      returncode = assessconvergence(s,
        xtol_abs, xtol_rel, ftol_abs, ftol_rel, stopval)
    end
  end

  iters == maxiters && (returncode = :MAXITERS_REACHED)
  best = bestvertex(s)
  return value(best), position(best), returncode, iters
end # optimise

end


using Base.Test
using ChaosQuantified

include("helpers.jl")
include("Rqa.jl")
include("entropy.jl")


## Julia program for pool-adjacent-violators algorithm
## author: weiya <szcfweiya@gmail.com>
## date: 2018-08-19

function pooladv(f::Array, w::Array)
   n = size(f, 1)
   lag = diff(f)
   if all(i -> i >= 0, lag)
        return(f)
   end
   while true
        for i = 1:(n-1)
            global idx
            idx = i+1
            lag[i] < 0 && break
        end
        newf = (w[idx]*f[idx] + w[idx-1]*f[idx-1])/(w[idx]+w[idx-1])
        f[idx] = newf
        f[idx-1] = newf
        lag = diff(f)
        if all(i -> i>=0, lag)
            return(f)
        end
   end
end

## example
g = pooladv([23, 27, 25, 28], ones(4))
for i in eachindex(g)
    println(g[i])
end


include("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/scripts_for_figures/common_functions.jl")
include("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/scripts_for_figures/FK_sampling_functions.jl")

using DataFrames, DataFramesMeta, RCall, Distributions, StatsBase
R"library(copula)"

function R_q(q, b, k_n, alpha)
    t1 = exp(- ((b/k_n)^alpha)*((1 + q^(1/alpha))^alpha - q))
    t2 =  ((q^(1/alpha))/(1+q^(1/alpha)))^(1-alpha)
    return t1*t2
end

function sample_xi_n(alpha, b, theta, k_n, n)
   u =  Uniform(0,1)
   g = Gamma(theta/alpha +n, 1/((b/k_n)^(alpha)))
   r_n = rand(u,1)[1]
   dz_n  = rand(g,1)[1]
   while r_n > R_q(dz_n, b, k_n,alpha)
       u =  Uniform(0,1)
       r_n = rand(u,1)[1]
       dz_n  = rand(g,1)[1]
       if r_n<= dz_n
        # println("true")
       end
         #return (b/k_n)*((dz_n)^(1/alpha))
   end
   return (b/k_n)*((dz_n)^(1/alpha))
end

function density_xi_n(x, alpha, b, theta, k_n, n)
  term1 = exp(-((b/k_n + x)^alpha))
  term2 = (b/k_n + x)^(alpha - 1)
  term3 = x^(theta + (n-1)*alpha)
  return term1*term2*term3
end

## test xi_n density

#x_t = collect(range(0,1.2*10, length=100))
#xi = Array{Float64}(undef, 10000)
#for i in 1:10000
#  xi[i] = sample_xi_n(0.75, 1, 0.75, 1,1)
#end

## TEST for xi_n

#histogram(xi,bins=250)
#x_t = collect(range(0,1.2*10^2, length=1000))
#plot!(x_t,2000*map(x ->density_xi_n(x,0.75, 1.0, 0.75, 1, 1),x_t))

#histogram(sim.^(1/0.25),bins=250, normalize=true)
#plot!(x_t,map(x ->density_xi_n(x,0.25, 1.0, 0.25, 1, 1),x_t)/20)
#plot!(x_t,2000*map(x ->density_xi_n(x,0.75, 1.0, 0.75, 1, 1),x_t))

function SB_NGG(alpha_, b, theta,N)
  u_n =Array{Float64}(undef, N-1)
  k_n = 1
  for i in (1:(N-1))
    if i == 1
      k_n  = 1
    else
      k_n =prod(vec(ones(1,i-1)).-u_n[1:(i-1)])
    end
      xi_n = sample_xi_n(alpha_, b, theta, k_n,i)
    #println(xi_n)
    R"zn = rgamma(1,1- $alpha_, $b/$k_n + $xi_n)"
    #g = Gamma(1- alpha_, 1/(b/k_n + xi_n))
    #z_n  = rand(g,1)[1]
    z_n  = @rget zn
    #h_n  = b/k_n + xi_n
    R" xn <- retstable($alpha_, 1, h =$b/$k_n + $xi_n,method = 'MH')"
    x_n = @rget xn
    u_n[i] = z_n/ (z_n + x_n)
 end

  p = Array{Float64}(undef, N)
  p[1] = u_n[1]
  for l in 2:(N-1)
    p[l] = prod(vec(ones(1,l - 1)).-u_n[1:(l - 1)])*u_n[l]
  end
  p[N]=prod(vec(ones(1,(N-1))).- u_n)
  return p
end

#n_tr=10
#p_k = SB_NGG(.25, 1,0.25,n_tr)
#plot(1:n_tr,p_k)

function Pkn_NGG_SB(n, β, σ, M; runs=10^4)
  theta = 1
  b = (σ*theta*β)^(1/σ)
  alpha_ = σ
    array_clust_num = Array{Int64}(undef, runs)
    for i in 1:runs
        weights_NGG_SB = SB_NGG(alpha_,b, theta, M)
        c = wsample(1:M, weights_NGG_SB, n, replace=true)
        n_c =  length(unique(c))
        array_clust_num[i] = n_c
    end
    return proportions(array_clust_num, n)
end

function Pkn_NGG_SB_1(n, β, σ, M; runs=10^4)
  theta = (σ*β)
  b = 1
  alpha_ = σ
    array_clust_num = Array{Int64}(undef, runs)
    for i in 1:runs
        weights_NGG_SB = SB_NGG(alpha_,b, theta, M)
        c = wsample(1:M, weights_NGG_SB, n, replace=true)
        n_c =  length(unique(c))
        array_clust_num[i] = n_c
    end
    return proportions(array_clust_num, n)
end


# Sample a gas fitness DGP and estimate the parameters for different lenghts of
# time series
using JLD2
using Hwloc

counts = Hwloc.histmap(Hwloc.topology_load())
ncores = counts[:Core]
if nworkers()<ncores
    addprocs(ncores - 1  - nworkers())
end
using DynNets
@everywhere using DynNets
#
N_est = 150
N = 20 ;  NTV =1
indTvPars = falses(N);indTvPars[1:NTV] = true
Tvals = [100,250,500,750]
N_T = length(Tvals)
est_times = SharedArray{Float64}(zeros(N_T,N_est))
est_w = SharedArray{Float64}(zeros(N_T,N_est,N))
est_A = SharedArray{Float64}(zeros(N_T,N_est,NTV))
est_B = SharedArray{Float64}(zeros(N_T,N_est,NTV))
est_conv_flag = SharedArray{Bool,2}(N_T,N_est)
est_flag = SharedArray{Bool,2}(N_T,N_est)

real_A = SharedArray{Float64}(zeros(N_T,N_est,NTV))
real_B = SharedArray{Float64}(zeros(N_T,N_est,NTV))
 @progress "T"   for ind = 1:N_T
    T = Tvals[ind]
    @sync @parallel    for n = 1:N_est
        sim_A = rand((0.5:0.00001:0.9),NTV)
        sim_B = rand((0.4:0.00001:0.999),NTV)
        real_A[ind,n,1:NTV] = sim_A
        real_B[ind,n,1:NTV] = sim_B
        Y_T,fit_T =  DynNets.sampleParDynNetsBin1(T,"Fit-Multi-gas",
                                                    indTvPar = [trues(NTV);falses(N-NTV)],
                                                    α = 5,β=1,
                                                    Agas =sim_A , Bgas = sim_B  )
        degsT = squeeze(sum(Y_T,2),2)
        model_data = DynNets.SdErgm1(degsT,[zeros(i)for i in [N,NTV,NTV]], indTvPars)
        #estimate
        est_times[ind,n] = @elapsed hat_pars ,hat_inds,conv_flag = DynNets.estimate(model_data)

        est_w[ind,n,:] = hat_pars[1]
        est_A[ind,n,:] = hat_pars[2]
        est_B[ind,n,:] = hat_pars[3]
        est_conv_flag[ind,n] = conv_flag
        est_flag[ind,n] = true
        #println([T,n])
    end

    # save
    save_path = "./data/estimatesTest/EstVarT.jld"
    @save(save_path)
    println([T,sum(est_times[ind,:])])
end


#Copyright (c) 2015 Rex Computing and Isaac Yonemoto

#see LICENSE.txt

#this work was supported in part by DARPA Contract D15PC00135

#hilbert.jl
#tests inverting a hilbert matrix.  All of the entries *should* be integers.

vtypes = [Float16, Float32, Float64, BigFloat]

#generate a hilbert matrix
function hilbert(T::Type, n)
  [1 / (convert(T, x + y - 1)) for x=1:n, y=1:n]
end

for T in vtypes
  for i in 1:10
    h = hilbert(T, i)

    invh = h^(-1)

    println("hilbert matrix $(i) upper corner in $(T) should be $(i ^ 2) but is: $(invh[1,1])")
  end
end


using QuBase
using Base.Test

include("multest.jl")
include("tensortest.jl")
include("labelbasistest.jl")
include("operatortest.jl")


using Distributions
using JSON

n = 15
p = 0.4
d = Binomial( n,p )

x = [ -3, -1, 0, 1, 3 ]

dmgf(t) = mgf(d, t )
y = map( dmgf, x )
println( y )

data = Dict([
	("n", n),
	("p", p),
	("data", x),
	("expected", y)
])

outfile = open("./test/fixtures/partial.json", "w")
JSON.json(data)

write( outfile, JSON.json(data) )


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:5] <= 1)

                     @variable(m, -3.1544947018849654 <= q <= 4.066926790342683)

                     add_NL_constraint(m, :(swish(-0.7544717811096784 + -0.6289249660430758*swish(-0.401848603561894 + 0.9879850116297502*swish(0.7471741672319268 + 0.1562161389212755*$(x[1]) + -0.06104000939937393*$(x[2]) + -0.24396039580062512*$(x[3]) + -0.35002310972719863*$(x[4]) + -0.47890036372503086*$(x[5])) + -0.6833595346541679*swish(0.07659783311726898 + 0.1400377608602934*$(x[1]) + -0.5920144548543678*$(x[2]) + -0.3715997206073691*$(x[3]) + -0.08423829417739581*$(x[4]) + -0.48744433011047184*$(x[5]))) + -0.656030629474813*swish(-0.9665555719654946 + 0.014645656858022438*swish(0.7471741672319268 + 0.1562161389212755*$(x[1]) + -0.06104000939937393*$(x[2]) + -0.24396039580062512*$(x[3]) + -0.35002310972719863*$(x[4]) + -0.47890036372503086*$(x[5])) + -0.5583055948877655*swish(0.07659783311726898 + 0.1400377608602934*$(x[1]) + -0.5920144548543678*$(x[2]) + -0.3715997206073691*$(x[3]) + -0.08423829417739581*$(x[4]) + -0.48744433011047184*$(x[5])))) + swish(-0.3137145052115069 + -0.6627216812930383*swish(-0.401848603561894 + 0.9879850116297502*swish(0.7471741672319268 + 0.1562161389212755*$(x[1]) + -0.06104000939937393*$(x[2]) + -0.24396039580062512*$(x[3]) + -0.35002310972719863*$(x[4]) + -0.47890036372503086*$(x[5])) + -0.6833595346541679*swish(0.07659783311726898 + 0.1400377608602934*$(x[1]) + -0.5920144548543678*$(x[2]) + -0.3715997206073691*$(x[3]) + -0.08423829417739581*$(x[4]) + -0.48744433011047184*$(x[5]))) + -0.39516873731093893*swish(-0.9665555719654946 + 0.014645656858022438*swish(0.7471741672319268 + 0.1562161389212755*$(x[1]) + -0.06104000939937393*$(x[2]) + -0.24396039580062512*$(x[3]) + -0.35002310972719863*$(x[4]) + -0.47890036372503086*$(x[5])) + -0.5583055948877655*swish(0.07659783311726898 + 0.1400377608602934*$(x[1]) + -0.5920144548543678*$(x[2]) + -0.3715997206073691*$(x[3]) + -0.08423829417739581*$(x[4]) + -0.48744433011047184*$(x[5])))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


##################################################
# Compute the NR orbit-averaged diffusion coefficients in (Jr,L)
##################################################

include("../sources/julia/Main.jl")

Jr = 0.2 # Radial action
L = 0.3 # Angular momentum
m = 0.1 # Mass of field star

println("--------------------")

println("Plummer potential")

println("q  = ", qCalc)
println("m  = ", m)
println("Jr = ", Jr)
println("L  = ", L)

println("--------------------")

println("Coulomb log = ", logCoulomb)

println("--------------------")

@time dJr, dL, dJrJr, dJrL, dLL = avr_action_coefficients(Jr,L,m)

println("D_Jr  =  ", dJr)
println("D_JrJr =  ", dJrJr)
println("D_L  =  ", dL)
println("D_LL =  ", dLL)
println("D_JrL =  ", dJrL)

println("--------------------")


# Demignot, S., D., D.,
# Effect of prosthetic sugar groups on the pharmacokinetics of glucose-oxidase
# https://pubmed.ncbi.nlm.nih.gov/2855567/
using Logging

using StructuralIdentifiability

logger = Logging.SimpleLogger(stdout, Logging.Info)
global_logger(logger)

ode = @ODEmodel(
    x0'(t) = a1 * (x1(t) - x0(t)) - (ka * n * x0(t)) / (kc * ka + kc * x2(t) + ka * x0(t)),
    x1'(t) = a2 * (x0(t) - x1(t)),
    x2'(t) = b1 * (x3(t) - x2(t)) - (kc * n * x2(t)) / (kc * ka + kc * x2(t) + ka * x0(t)),
    x3'(t) = b2 * (x2(t) - x3(t)),
    y1(t) = x0(t)
)

@time println(assess_global_identifiability(ode))


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -3.9360745052223884 <= q <= 4.0522696599704755)

                     add_NL_constraint(m, :((-0.02759545635089289 + 0.8396153655678034*$(x[1]) + -0.45360175821557247*$(x[2]) + 0.6295131713145015*$(x[3]))*(1 + erf((-0.02759545635089289 + 0.8396153655678034*$(x[1]) + -0.45360175821557247*$(x[2]) + 0.6295131713145015*$(x[3]))/sqrt(2)))/2 + (0.0856930337249362 + -0.9280577687232072*$(x[1]) + 0.795858881094551*$(x[2]) + -0.3475251376807962*$(x[3]))*(1 + erf((0.0856930337249362 + -0.9280577687232072*$(x[1]) + 0.795858881094551*$(x[2]) + -0.3475251376807962*$(x[3]))/sqrt(2)))/2 - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


"""
Intelligent driver model
"""

function IDM(;Vego::Float64, Vfront::Float64, Vref::Float64, Snet::Float64, T::Float64=0.1, Amax::Float64, Bdec::Float64, Smin::Float64=0.1, δ::Float64=3.0)
    """
        Vego:   Velocity of ego car;
        Vfront: Velocity of front car;
        Vref:   the velocity the vehicle would drive at in free traffic;
        Snet:   Distance between ego and front car;
        Smin:   a minimum desired net distance;
        T:      the minimum possible time to the vehicle in front, or Safe time headway;
        Amax:   the maximum vehicle acceleration of ego car;
        Bdec:   comfortable braking deceleration, must be positive;
        δ:      just a exponent, usually set to 4.
    """
    ΔV = Vego - Vfront
    Sstar = Smin + Vego*T + Vego*ΔV/(2*sqrt(Amax*Bdec))
    Aego = Amax*(1 - (Vego/Vref)^δ - (Sstar/Snet)^2)
end

#= # Test
@show IDM(Vego=4.0, Vfront=0.0, Vref=4.0, Snet=4.0, Smin=0.1, T=0.1, Amax=3.0, Bdec=1.0)
=#


# Test coherent state
o2 = osc(15)
v = coherentvec(o2, 1)
@test v ⋅ v ≈ 1
@test expectvec(v, o2('n')) ≈ 1 atol=1e-11
@test expectvec(v, o2('n')) ≈ weakvaluevec(v, v, o2('n'))
@test expectvec(v, o2('d')) ≈ weakvaluevec(v, groundvec(o2), o2('d')) atol=1e-11
@test weakvaluevec(v, groundvec(o2), o2('n')) ≈ 0

vp = projector(v)
@test vp ≈ coherent(o2, 1)

@test tr(vp * vp) ≈ 1
@test expect(vp, o2('n')) ≈ 1 atol=1e-11
@test expect(vp, o2('n')) ≈ weakvalue(vp, vp, o2('n'))
@test expect(vp, o2('d')) ≈ weakvalue(vp, ground(o2), o2('d')) atol=1e-11
@test weakvalue(vp, ground(o2), o2('n')) ≈ 0


#AbstractTensor
Base.:(*)(α::Number, A::AbstractTensor) = apply(A, x -> @. x *= α)
Base.:(*)(A::AbstractTensor, α::Number) = α   * A
Base.:(/)(A::AbstractTensor, α::Number) = 1/α * A
Base.conj!(A::AbstractTensor) = apply!(A, conj!)
Base.conj(A::AbstractTensor)  = apply(A, conj!)
function LA.pinv(A::AbstractTensor, inds = (1,2))
    u, s, vd = tensorsvd(A, inds)
    apply!(s, x -> x .= LA.pinv(x))
    @TO.tensor A2[1,2] := vd'[-1,1] * s[-1,-2] * u'[2,-2]
    return A2
end

#DTensor
apply!(A::DTensor, fun!) = (fun!(A.array); return A)
apply(A::DTensor, fun!) = apply!(deepcopy(A), fun!)
Base.adjoint(a::DTensor) = conj(a)

#= Operations on Tensors =#
function apply!(A::DASTensor, fun!)
    foreach(fun!, values(A))
    return A
end
apply(A::DASTensor, fun!) = apply!(deepcopy(A), fun!)
function Base.adjoint(A::DASTensor)
    B = apply(A, conj!)
    B.io = inv(B.io)
    setcharge!(B, -charge(A))
    return B
end


@testset "Context" begin
    let space = StateSpace([1,2], 2, [0.1, 1.0], 2)
        ctx = Context(0.5, space, [0,1])
        @test ctx.β == 0.5
        @test ctx.system_space === space
        @test ctx.system_state == [0,1]
        @test ctx.bath_space == space
        let a = e^-2.55, b = e^-1.1
            @test ctx.bath_state ≈ [2a/(2a + b), b/(2a + b)]
        end
    end
end

let rng = MersenneTwister(2018)
    @testset "Timestep" begin
        let ctx = Context(0.5, StateSpace([1,2], 4, [0.1, 1.0], 4), [1, 0, 0])
            re = relative_entropy(ctx.system_state, ctx.bath_state)
            for _ in 1:10000
                timestep(rng, ctx)
                r = relative_entropy(ctx.system_state, ctx.bath_state)
                @test isapprox(r, re; atol=1e-9) || r < re
                re = r
            end
            @test ctx.system_state ≈ ctx.bath_state
        end
        let ctx = Context(0.5, StateSpace(1:4, 4, [0.1, 1.0], 4), [1,0,0,0,0])
            re = relative_entropy(ctx.system_state, ctx.bath_state)
            for _ in 1:1000000
                timestep(rng, ctx)
                r = relative_entropy(ctx.system_state, ctx.bath_state)
                @test isapprox(r, re; atol=1e-9) || r < re
                re = r
            end
            @test ctx.system_state ≈ ctx.bath_state
        end
    end
end


struct SignedDistanceRefinery{F <: Function} <: AbstractRefinery
    signed_distance_func::F
    atol::Float64
    rtol::Float64
end

function needs_refinement(refinery::SignedDistanceRefinery, cell::Cell)
    minimum(cell.boundary.widths) > refinery.atol || return false

    for c in body_and_face_centers(cell.boundary)
        value_interp = evaluate(cell, c)
        value_true = refinery.signed_distance_func(c)
        if !isapprox(value_interp, value_true, rtol=refinery.rtol, atol=refinery.atol)
            return true
        end
    end
    return false
end

function refine_data(refinery::SignedDistanceRefinery, cell::Cell, indices)
    refine_data(refinery, child_boundary(cell, indices))
end

function refine_data(refinery::SignedDistanceRefinery, boundary::HyperRectangle)
    extrapolate(interpolate!(refinery.signed_distance_func.(vertices(boundary)),
                             BSpline(Linear()),
                             OnGrid()), Linear())
end


identity_fun(d::Circle)=Fun([d.center,0.,d.radius],Laurent(d))

Space(d::PeriodicDomain)=Fourier(d)
Space(d::Circle)=Laurent(d)

## Evaluation

Evaluation(d::PeriodicDomain,x::Number,n...) = Evaluation(Laurent(d),complex(x),n...)

## Definite Integral

DefiniteIntegral(d::PeriodicDomain) = DefiniteIntegral(Laurent(d))
DefiniteLineIntegral(d::PeriodicDomain) = DefiniteLineIntegral(Laurent(d))

## Toeplitz


using NeuralNetwork
import NeuralNetwork: accuracy

using RDatasets: dataset

iris = dataset("datasets", "iris")

X = Array(iris[:, 1:4])
y = zeros(Int, size(iris, 1))
for (i, label) in enumerate(["versicolor", "setosa", "virginica"])
    y[iris[:, 5] .== label] = i
end

indices = shuffle(1:size(iris, 1))

X = X[indices, :]'
y = y[indices]

N = Int(round(size(X, 2)*0.8))
X_train, X_test = X[:, 1:N], X[:, N+1:end]
y_train, y_test = y[1:N], y[N+1:end]

layer = LogisticRegression(4, 3, 5e-3, 1000)
fit!(layer, X_train, onehot(y_train))
Ŷ = predict(layer, X_test)

ŷ = maxindices(Ŷ)
println("Accuracy: $(accuracy(y_test, ŷ))")


import Dates: Date, year, month, day

cd(dirname(@__FILE__))

function jl_am_put_crr(S,K,r,c,σ,dt1,dt2,N)
    t = Int(dt2 - dt1)/365
    Δt = t/N
    u = exp(σ*√Δt)
    d = exp(-σ*√Δt)
    B = exp((r-c)*Δt)
    p = (B-d)/(u-d)
    q = (u-B)/(u-d)
    iR = exp(-Δt*r)

    # final time
    Z = map(0:N) do i
        x = K - S*exp((2*i-N)σ*√Δt)
        max(x,0)
    end

    for n = N-1:-1:0
        for i = 0:n
            x = K - S*exp((2*i-n)σ*√Δt)
            y = iR*(q*Z[i+1] + p*Z[i+2])
            Z[i+1] = max(x, y)
        end
    end
    return Z[1]
end

function ql_am_put_crr(S,K,r,c,σ,dt1::Date,dt2::Date, n)
    ccall((:ql_am_put_crr, :quantlib), Cdouble,
          (Cdouble, Cdouble, Cdouble, Cdouble, Cdouble,
           Cint, Cint, Cint,
           Cint, Cint, Cint,
           Cint),
          S, K, r, c, σ,
          year(dt1), month(dt1), day(dt1),
          year(dt2), month(dt2), day(dt2),
          n)
end

using BenchmarkTools
using Miletus

println("CRR Model")
println("=========")
println()

println("Miletus")
println(@benchmark value(CRRModel(Date(2010,01,01), Date(2011,01,01), 801, 36.0, 0.06, 0.0, 0.2),
                 AmericanPut(Date(2011,01,01), SingleStock(), 40.0)))

println("Miletus Currency")
println(@benchmark value(CRRModel(Date(2010,01,01), Date(2011,01,01), 801, 36.0USD, 0.06, 0.0, 0.2),
                 AmericanPut(Date(2011,01,01), SingleStock(), 40.0USD)))

if isfile("quantlib")
    println()
    println("QuantLib")
    println(@benchmark ql_am_put_crr(40.0, 36.0, 0.06, 0.0, 0.2,
                                     Date(2010,01,01), Date(2011,01,01), 801))
end

println()
println("Pure Julia")
println(@benchmark jl_am_put_crr(40.0, 36.0, 0.06, 0.0, 0.2,
           Date(2010,01,01), Date(2011,01,01), 801))
println()
println()

println("Monte Carlo Model")
println("=========")
println()

d1 = Dates.today()
d2 = d1 + Dates.Day(120)
m = GeomBMModel(d1, 100.00, 0.05, 0.0, 0.3)
o = EuropeanCall(d2, SingleStock(), 100.00)

println("Miletus")
println(@benchmark value($m, $o, $MonteCarloModel, 10_000))
println()

m = GeomBMModel(d1, 100.00USD, 0.05, 0.0, 0.3)
o = EuropeanCall(d2, SingleStock(), 100.00USD)

println("Miletus Currency")
println(@benchmark value($m, $o, $MonteCarloModel, 10_000))
println()

println("Least Squares Monte Carlo Model")
println("=========")
println()

d1 = Dates.today()
d2 = d1 + Dates.Day(120)
m = GeomBMModel(d1, 100.00, 0.05, 0.0, 0.3)
o = AmericanPut(d2, SingleStock(), 100.00)

println("Miletus")
println(@benchmark value($m, $o, $LeastSquaresMonteCarlo, 10_000, 3))
println()


using Plots
using Primes
using Memoize # https://github.com/hymd-research/Memoize.jl.git

@memoize function g(N::T) where {T<:Integer}
    if N == one(T) # meaning N==1
        return one(T)
    else
        fN = sum(e for (_, e) ∈ factor(N).pe)
        ifelse(iseven(fN), 1, -1) + g(N - 1)
    end
end

@time gs = [g(N) for N = 1:1000000]
@assert sum(gs .≥ 0) == 10
scatter(gs)


include("../interfaces.jl")

import PhysicalConstants.CODATA2018 as Constants

### Boundary conditions
Temperature = Model(
    T = 293.15u"K",
    port = HeatPort,
    equations = :[port.T=T]
)

# Base Classes
PartialResistance = Model(
    port_a = HeatPort,
    port_b = HeatPort,
    partialEquations  = :[
        0 = port_a.Q_flow + port_b.Q_flow
        port_a.Q_flow = (port_a.T - port_b.T)/R
    ]
)

### Resistance Models
Sphere = PartialResistance | Model(
    r_inner = 0.1u"m",
    r_outer = 1.0u"m",
    lambda = 1.0u"W/(m*K)",
    R = Var(start=1.0u"K/W"),
    π=π,
    equations = :[
        R = 1/(4*π*lambda)*(1/r_inner - 1/r_outer)
    ]
)

Convection = PartialResistance | Model(
    surfaceArea = 0.01u"m^2",
    alpha = 1000.0u"W/(m^2*K)",
    R = Var(start=1.0u"K/W"),
    equations = :[
        R = 1/(alpha*max(eps()u"m^2",surfaceArea))
    ]
)

Contact = PartialResistance | Model(
    surfaceArea = 0.01u"m^2",
    Rc_pp = 1.0u"m^2*K/W",
    R = Var(start=1.0u"K/W"),
    equations = :[
        R = Rc_pp/max(eps()u"m^2",surfaceArea)
    ]
)

Radiation = PartialResistance | Model(
    surfaceArea = 0.01u"m^2",
    epsilon = 1.0u"m/m", # no units... don't know how to represent that
    R = Var(start=1.0u"K/W"),
    useExact=true,
    sigma=Constants.StefanBoltzmannConstant, # how to use this within a model without specifying as a parameter/variable?
    equations = :[
        R = if useExact
                1/(surfaceArea*sigma*epsilon*(port_a.T^2+port_b.T^2)*(port_a.T + port_b.T))
            else
                1/(4*surfaceArea*sigma*epsilon*(0.5*(port_a.T + port_b.T))^3)
            end
    ]
)


function compute_average_speed_of_sound(settings, z, c_0, T_0)

    n_intermediate = 11
    dz = z ./ (n_intermediate-1)
    c_bar = c_0
    for k in 1:(n_intermediate-1)
        T_im = T_0 .+ settings.dT .- (z .- k .* dz) .* (-0.0065)
        c_im = sqrt.(1.4 * 287. .* T_im)
        c_bar = (k) / (k + 1) .* c_bar .+ c_im ./ (k + 1)
    end

    return c_bar
end

function compute_euler_transformation(n_vcr_a_1, n_vcr_a_2, n_vcr_a_3, theta_B, phi_B, psi_B)

    cth  = cos.(deg2rad.(theta_B))
    sth  = sin.(deg2rad.(theta_B))
    cphi = cos.(deg2rad.(phi_B))
    sphi = sin.(deg2rad.(phi_B))
    cpsi = cos.(deg2rad.(psi_B))
    spsi = sin.(deg2rad.(psi_B))
    n_vcr_s_1 = cth .* cpsi .* n_vcr_a_1 .+ cth .* spsi .* n_vcr_a_2 - sth .* n_vcr_a_3
    n_vcr_s_2 = (-spsi .* cphi .+ sphi .* sth .* cpsi) .* n_vcr_a_1 .+ ( cphi .* cpsi .+ sphi .* sth .* spsi) .* n_vcr_a_2 .+ sphi .* cth .* n_vcr_a_3
    n_vcr_s_3 = (spsi .* sphi .+ cphi .* sth .* cpsi) .* n_vcr_a_1 .+ ( -sphi .* cpsi .+ cphi .* sth .* spsi) .* n_vcr_a_2 .+ cphi .* cth .* n_vcr_a_3

    return n_vcr_s_1, n_vcr_s_2, n_vcr_s_3
end

function geometry(settings, x_obs::Array{Float64, 1}, x, y, z, alpha, gamma, t_s, c_0, T_0)

    # Size of inputs
    n_t = size(x)[1]

    # Initialize vectors
    T = eltype(x)
    r = zeros(T, n_t)
    beta = zeros(T, n_t)
    theta = zeros(T, n_t)
    phi = zeros(T, n_t)
    c_bar = zeros(T, n_t)
    t_o = zeros(T, n_t)

    # Compute body angles (psi_B, theta_B, phi_B): angle of body w.r.t. horizontal
    theta_B = alpha .+ gamma
    phi_B = zeros(T, n_t)
    psi_B = zeros(T, n_t)

    # Compute the relative observer-aircraft position vector i.e. difference between observer and ac coordinate
    r_1 =  x_obs[1] .- x
    r_2 =  x_obs[2] .- y
    r_3 = -x_obs[3] .+ z

    # Normalize the distance vector
    r = sqrt.(r_1 .^2 .+ r_2 .^2 .+ r_3 .^2)
    n_vcr_a_1 = r_1 ./ r
    n_vcr_a_2 = r_2 ./ r
    n_vcr_a_3 = r_3 ./ r

    # Define elevation angle
    beta = asind.(n_vcr_a_3)

    # Transformation direction cosines (Euler angles) to the source coordinate system (i.e. take position of the aircraft into account)
    n_vcr_s_1, n_vcr_s_2, n_vcr_s_3 = compute_euler_transformation(n_vcr_a_1, n_vcr_a_2, n_vcr_a_3, theta_B, phi_B, psi_B)

    # Define polar directivity angle
    theta = acosd.(n_vcr_s_1)

    # Define azimuthal directivity angle
    phi = -atand.(n_vcr_s_2, n_vcr_s_3)

    # Average speed of sound
    c_bar = compute_average_speed_of_sound(settings, z, c_0, T_0)

    # Define observed time
    t_o = t_s .+ r ./ c_bar

    return r, beta, theta, phi, c_bar, t_o
end


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

# %%
using Plots, LaTeXStrings
pyplot(fmt=:svg, fontfamily="sans-serif")

function pitick(start, stop, denom; mode=:text)
    a = Int(cld(start, π/denom))
    b = Int(fld(stop, π/denom))
    tick = range(a*π/denom, b*π/denom; step=π/denom)
    ticklabel = piticklabel.((a:b) .// denom, Val(mode))
    tick, ticklabel
end

function piticklabel(x::Rational, ::Val{:text})
    iszero(x) && return "0"
    S = x < 0 ? "-" : ""
    n, d = abs(numerator(x)), denominator(x)
    N = n == 1 ? "" : repr(n)
    d == 1 && return S * N * "π"
    S * N * "π/" * repr(d)
end

function piticklabel(x::Rational, ::Val{:latex})
    iszero(x) && return L"0"
    S = x < 0 ? "-" : ""
    n, d = abs(numerator(x)), denominator(x)
    N = n == 1 ? "" : repr(n)
    d == 1 && return L"%$S%$N\pi"
    L"%$S\frac{%$N\pi}{%$d}"
end

a, b = -2π, 2π
plot(sin, a, b; xtick=pitick(a, b, 4), label=L"y = \sin(x)", size=(720, 250))

# %%
plot(sin, a, b; xtick=pitick(a, b, 4; mode=:latex), label=L"y = \sin(x)", size=(720, 250),
    tickfontsize=10)

# %%
gr(fmt=:auto)

plot(sin, a, b; xtick=pitick(a, b, 4), label=L"y = \sin(x)", size=(720, 250),
    fontfamily="Computer Modern", legendfontsize=12)

# %%
plot(sin, a, b; xtick=pitick(a, b, 4; mode=:latex), label=L"y = \sin(x)", size=(720, 250),
    tickfontsize=10, fontfamily="Computer Modern", legendfontsize=12, bottom_margin=3Plots.mm)

# %%
function piticklabel(x::Rational, ::Val{:gr})
    iszero(x) && return L"0"
    S = x < 0 ? "-" : ""
    n, d = abs(numerator(x)), denominator(x)
    N = n == 1 ? "" : repr(n)
    d == 1 && return L"%$S%$N\pi"
    L"%$S\dfrac{%$N\pi}{%$d}"
end

plot(sin, a, b; xtick=pitick(a, b, 4; mode=:gr), label=L"y = \sin(x)", size=(720, 250),
    tickfontsize=10, fontfamily="Computer Modern", legendfontsize=12, bottom_margin=3Plots.mm)

# %%


struct IHTOptions
 epsTol::Float64
 maxIter::Int64
 checkEvery::Int64
end

IHTOptions(;epsTol::Float64=1e-2, maxIter::Int64=500, checkEvery::Int64=10) =
  IHTOptions(epsTol, maxIter, checkEvery)

  function covselLatentIHT_grad_S!(gS, Σ, S, L)
    T = S + L
    T = inv(T)
    gS .= Σ - T
    gS
  end

  function covselLatentIHT_grad_U!(gU, gS, Σ, S, L, U)
    covselLatentIHT_grad_S!(gS, Σ, S, L)
    scale!(gS, 2.)
    A_mul_B!(gU, gS, U)
  end

  function covselLatentIHT!(Σ::StridedMatrix, ηS, ηU, s, r;
    epsTol=1e-2, maxIter=100, checkEvery=5, callback=nothing)

    p = size(Σ, 1)

    gS = zeros(p, p)
    gU = zeros(p, r)

    S = zeros(p, p)
    L = zeros(p, p)

    # init
    tmp = full(inv(Symmetric(Σ)))
    HardThreshold!(S, tmp, s, false)

    tmp .= tmp - S
    U, d, V = svd(tmp)
    U = U[:,1:r] .* sqrt.(d[1:r])'
    L = U * U'

    fvals = []
    # gradient descent
    #
    fv = covSel_objective(Σ, S, L)
    push!(fvals, fv)
    for iter=1:maxIter
      # update S
      covselLatentIHT_grad_S!(gS, Σ, S, L)
      @. S -= ηS * gS
      HardThreshold!(S, S, s, false)
      max_gS = maximum( abs.(gS) )

      # update U and L
      covselLatentIHT_grad_U!(gU, gS, Σ, S, L, U)
      @. U -= ηU * gU
      A_mul_Bt!(L, U, U)
      max_gU = maximum( abs.(gU) )
      if callback != nothing
        callback(Σ, S, L)
      end

      done = max(max_gS, max_gU) < epsTol
      # check for convergence
      if mod(iter, checkEvery) == 0
        fv_new = covSel_objective(Σ, S, L)
        push!(fvals, fv_new)
        done = abs(fv_new - fv) <= epsTol
        fv = fv_new
      end
      done && break
    end

    (S, L, U, fvals, iter)
  end


using Test

include("lexicographic-perms.jl")

@testset "#all_perms base cases" begin
  @test all_perms([]) == [[]]

  @test all_perms([0]) == [[0]]

  @test all_perms([1, 2]) == [[1, 2], [2, 1]]
end

@testset "#prepend_all!" begin
  @test prepend_all!(3, [[]]) == [[3]]

  @test prepend_all!(3, [[1, 2], [2, 1]]) == [[3, 1, 2], [3, 2, 1]]

  @test prepend_all!(4, [[1, 2, 3], [1, 3, 2],
                         [2, 1, 3], [2, 3, 1],
                         [3, 1, 2], [3, 2, 1]]) == [[4, 1, 2, 3], [4, 1, 3, 2],
                                                    [4, 2, 1, 3], [4, 2, 3, 1],
                                                    [4, 3, 1, 2], [4, 3, 2, 1]]

end

@testset "#all_perms rec. cases" begin
  @test all_perms([true, false]) ==  [Bool[1, 0], Bool[0, 1]]

  @test all_perms([1, 2, 3]) == [[1, 2, 3], [1, 3, 2], [2, 1, 3],
                                 [2, 3, 1], [3, 1, 2], [3, 2, 1]]

  @test all_perms([1, 'a', "foo"]) == [Any[1, 'a', "foo"], Any[1, "foo", 'a'],
                                       Any['a', 1, "foo"], Any['a', "foo", 1],
                                       Any["foo", 1, 'a'], Any["foo", 'a', 1]]

  @test all_perms(Vector{Int}([1, 2, 3])) == Vector{Vector{Int}}([[1, 2, 3], [1, 3, 2], [2, 1, 3],
                                                                  [2, 3, 1], [3, 1, 2], [3, 2, 1]])

  @test all_perms(Vector{Float64}([1., 2., 3.])) == Vector{Vector{Float64}}([
                                                                             [1., 2., 3.], [1., 3., 2.], [2., 1., 3.],
                                                                             [2., 3., 1.], [3., 1., 2.], [3., 2., 1.]])

  @test all_perms([1, 2, 3, 4]) == Vector{Vector{Int}}([
                                    [1, 2, 3, 4], [1, 2, 4, 3],
                                    [1, 3, 2, 4], [1, 3, 4, 2],
                                    [1, 4, 2, 3], [1, 4, 3, 2],
                                    [2, 1, 3, 4], [2, 1, 4, 3],
                                    [2, 3, 1, 4], [2, 3, 4, 1],
                                    [2, 4, 1, 3], [2, 4, 3, 1],
                                    [3, 1, 2, 4], [3, 1, 4, 2],
                                    [3, 2, 1, 4], [3, 2, 4, 1],
                                    [3, 4, 1, 2], [3, 4, 2, 1],
                                    [4, 1, 2, 3], [4, 1, 3, 2],
                                    [4, 2, 1, 3], [4, 2, 3, 1],
                                    [4, 3, 1, 2], [4, 3, 2, 1],
  ])

  @test all_perms(['v', 'x', 'y', 'z']) == [['v', 'x', 'y', 'z'], ['v', 'x', 'z', 'y'],
                                            ['v', 'y', 'x', 'z'], ['v', 'y', 'z', 'x'],
                                            ['v', 'z', 'x', 'y'], ['v', 'z', 'y', 'x'],
                                            ['x', 'v', 'y', 'z'], ['x', 'v', 'z', 'y'],
                                            ['x', 'y', 'v', 'z'], ['x', 'y', 'z', 'v'],
                                            ['x', 'z', 'v', 'y'], ['x', 'z', 'y', 'v'],
                                            ['y', 'v', 'x', 'z'], ['y', 'v', 'z', 'x'],
                                            ['y', 'x', 'v', 'z'], ['y', 'x', 'z', 'v'],
                                            ['y', 'z', 'v', 'x'], ['y', 'z', 'x', 'v'],
                                            ['z', 'v', 'x', 'y'], ['z', 'v', 'y', 'x'],
                                            ['z', 'x', 'v', 'y'], ['z', 'x', 'y', 'v'],
                                            ['z', 'y', 'v', 'x'], ['z', 'y', 'x', 'v']
  ]

  @test all_perms(collect(0:4)) |> length == 120

  @test all_perms(collect(0:5)) |> length == 720

  @test all_perms(collect(0:6)) |> length == 5040
end

@time @testset "#all_perms of 10 digits => 10!" begin
  # 10! ≡ 3_628_800
  @test_skip all_perms(collect(0:9)) |> length == 3_628_800
end
#
# Test Summary:                  | Pass  Total
# #all_perms of 10 digits => 10! |    1      1
#   6.666260 seconds (36.33 M allocations: 3.653 GiB, 26.40% gc time, 5.47% compilation time)
#

@time @testset "#all_perms of 11 digits => 11!" begin
  # 11! ≡ 39_916_800
  @test_skip all_perms(collect(0:10)) |> length == 39_916_800
end
#
# Test Summary:                  | Pass  Total
# #all_perms of 11 digits => 11! |    1      1
# 169.726583 seconds (284.85 M allocations: 31.599 GiB, 46.52% gc time, 1.47% compilation time)
#

# way too long...
# @time @testset "#all_perms of 12 digits => 12!" begin
#   # 12! ≡ 479_001_600
#   @test_skip all_perms(collect(0:11)) |> length == 479_001_600
# end

@time @testset "millionnth permutation of 0..9" begin
  @test millionnth(collect(0:9)) == [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
end