"""
$(TYPEDEF)
Apex type for coordinate systems
"""
abstract type AbstractCoordinateSystem <: AbstractExtendableGridApexType end

"""
$(TYPEDSIGNATURES)

List possible coordinate systems. These describe the meaning of the grid coordinates.
"""
coordinatesystems()=AbstractTrees.print_tree(AbstractCoordinateSystem,5,indicate_truncation=false)

"""
$(TYPEDEF)
1D cartesion coordinate system (unknown `x`)
"""
abstract type Cartesian1D   <: AbstractCoordinateSystem end

"""
$(TYPEDEF)
2D cartesion coordinate system (unknowns `x,y`)
"""
abstract type Cartesian2D   <: AbstractCoordinateSystem end

"""
$(TYPEDEF)
2D cartesion coordinate system (unknowns `x,y,z`)
"""
abstract type Cartesian3D   <: AbstractCoordinateSystem end

"""
$(TYPEDEF)
3D cylindrical coordinate system (unknowns `r,ϕ,z`)
"""
abstract type Cylindrical3D <: AbstractCoordinateSystem end

"""
$(TYPEDEF)
2D cylindrical coordinate system (unknowns `r,z`)
"""
abstract type Cylindrical2D <: AbstractCoordinateSystem end

"""
$(TYPEDEF)
2D polar coordinate system (unknowns `r,ϕ`)
"""
abstract type Polar2D       <: AbstractCoordinateSystem end

"""
$(TYPEDEF)
1D polar coordinate system (unknown `r`)
"""
abstract type Polar1D       <: AbstractCoordinateSystem end

"""
$(TYPEDEF)
3D spheriacal coordinate system (unknowns `r,ϕ,θ`)
"""
abstract type Spherical3D   <: AbstractCoordinateSystem end

"""
$(TYPEDEF)
1D spheriacal coordinate system (unknown `r`)
"""
abstract type Spherical1D   <: AbstractCoordinateSystem end  #r (integral over ϕ,\theta)


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 4b149ad2-5bca-11eb-0104-0be8dd786177
begin
	import Pkg
	Pkg.activate(Base.current_project())
	using ReinforcementLearning
	using Flux
	using Statistics
	using Plots
end

# ╔═╡ 48977b2e-5bc9-11eb-2ad7-958ca5fd2ecb
md"""
# Chapter 8.6 Trajectory Sampling

The general function `run(policy, env, stop_condition, hook)` is very flexible and powerful. However, we are not restricted to use it only. In this notebook, we'll see how to use part of the components provided in `ReinforcementLearning.jl` to finish some specific experiments.

First, let's define the environment mentioned in **Chapter 8.6**:
"""

# ╔═╡ 6bf4f104-5bca-11eb-2b4a-69c4ba7ffadc
begin
	mutable struct TestEnv <: AbstractEnv
		transitions::Array{Int, 3}
		rewards::Array{Float64, 3}
		reward_table::Array{Float64, 2}
		terminate_prob::Float64
		# cache
		s_init::Int
		s::Int
		reward::Float64
		is_terminated::Bool
	end

	function TestEnv(;ns=1000, na=2, b=1, terminate_prob=0.1,init_state=1)
		transitions = rand(1:ns, b, na, ns)
		rewards = randn(b, na, ns)
		reward_table = randn(na, ns)
		TestEnv(
			transitions,
			rewards,
			reward_table,
			terminate_prob,
			init_state,
			init_state,
			0.,
			false
		)
	end

	function (env::TestEnv)(a::Int)
		t = rand() < 0.1
		bᵢ = rand(axes(env.transitions, 1))

		env.is_terminated = t
		if t
			env.reward = env.reward_table[a, env.s]
		else
			env.reward = env.rewards[bᵢ, a, env.s]
		end

		env.s = env.transitions[bᵢ, a, env.s]

	end

	RLBase.state_space(env::TestEnv) = Base.OneTo(1:size(env.rewards, 3))
	RLBase.action_space(env::TestEnv) = Base.OneTo(1:size(env.rewards, 2))

	function RLBase.reset!(env::TestEnv)
		env.s = env.s_init
		env.reward = 0.0
		env.is_terminated = false
	end

	RLBase.is_terminated(env::TestEnv) = env.is_terminated
	RLBase.state(env::TestEnv) = env.s
	RLBase.reward(env::TestEnv) = env.reward
end

# ╔═╡ 989f834a-5bca-11eb-363a-5f4afce32520
md"""
Note that this environment is not described very clearly on the book. Part of the information are inferred from the [lisp source code](http://incompleteideas.net/book/code/sampling2.lisp).

!!! info
	Actually the lisp code is also not perfect, I spent a whole afternoon to figure out the code logic. So good luck if you also want to understand it.

The definitions above are just like any other environment we've defined before in previous chapters. Now we'll add an extra function to make it work for our **planning** purpose.
"""

# ╔═╡ 97be6cb8-5bcb-11eb-3dc7-1b8105b206d7
"""
Return all the possible next states and corresponding reward.
"""
function successors(env::TestEnv, s, a)
    S = @view env.transitions[:, a, s]
    R = @view env.rewards[:, a, s]
    zip(R, S)
end

# ╔═╡ 0c52737a-5bcd-11eb-1546-21885703f303
γ = 0.9

# ╔═╡ aa061636-5bce-11eb-1381-bb00fb913f9b
n_sweep=10

# ╔═╡ a466db90-5bcc-11eb-3385-e5117be33410
"""
Here we are only interested in the performance of Q
with env starting at state `1`. Note here we're calculating
the discounted reward.
"""
function eval_Q(Q, env;n_eval=100)
	R = 0.
	for _ in 1:n_eval
		reset!(env)
		i = 0
		while !is_terminated(env)
			a = Q(state(env)) |> argmax  # greedy
			env(a)
			R += reward(env) * γ^i
			i += 1
		end
	end
	R/n_eval
end

# ╔═╡ 9bf4aea2-5bcd-11eb-230f-cf63a009a0f9
"""
Calculate the expected gain.
"""
function gain(Q,env,s,a)
    p = env.terminate_prob
    r = env.reward_table[a, s]
    p * r + (1-p) * mean(r̄ + γ * maximum(Q(s′)) for (r̄, s′) in  successors(env, s, a))
end

# ╔═╡ b92ba1e2-5bcb-11eb-16cf-f15929c2697b
function sweep(;b = 1, ns=1000)

    na = 2

    α=1.0
    p = 0.1

    env= TestEnv(;ns=ns, na=na, b=b, terminate_prob=p)
    Q = TabularQApproximator(;n_state=ns, n_action=na, opt=Descent(α))

    i = 1
    vals = [eval_Q(Q, env)]
    for _ in 1:n_sweep
        for s in 1:ns
            for a in 1:na
                G = gain(Q,env,s,a)
                update!(Q, (s,a) =>  Q(s, a) - G)
                if i % 100 == 0
                    push!(vals, eval_Q(Q, env))
                end
                i += 1
            end
        end
    end
    vals
end

# ╔═╡ 86538552-5bce-11eb-2cbf-ff6e175ef421
function on_policy(;b = 1, ns=1000)

    na = 2

    α=1.0
    p = 0.1

    env= TestEnv(;ns=ns, na=na, b=b, terminate_prob=p)
    Q = TabularQApproximator(;n_state=ns, n_action=na, opt=Descent(α))

    i = 1
    vals = [eval_Q(Q, env)]

	explorer = EpsilonGreedyExplorer(0.1)
	for i in 1:(n_sweep * ns * na)
		is_terminated(env) && reset!(env)
		s = state(env)
		a = Q(s) |> explorer
		env(a)
		G = gain(Q, env, s, a)
		update!(Q, (s,a) => Q(s,a) - G)
		if i % 100 == 0
			push!(vals, eval_Q(Q, env))
		end
	end

    vals
end

# ╔═╡ f29c9f18-5bcd-11eb-05b5-c72249e02975
begin
	fig_8_8 = plot(legend=:bottomright, xlabel="Computation time", ylabel="Value of start state under greedy policy")
	for b in [1, 3, 10]
		plot!(fig_8_8, mean(sweep(;b=b) for _ in 1:200), label="uniform b=$b")
		plot!(fig_8_8, mean(on_policy(;b=b) for _ in 1:200), label="on-policy b=$b")
	end
	fig_8_8
end

# ╔═╡ 8cba1ec8-5bcf-11eb-32c3-f38dee82139a
begin
	fig_8_8_2 = plot(legend=:bottomright, xlabel="Computation time", ylabel="Value of start state under greedy policy")
	plot!(fig_8_8_2, mean(sweep(;ns=10_000) for _ in 1:200), label="uniform")
	plot!(fig_8_8_2, mean(on_policy(;ns=10_000) for _ in 1:200), label="on-policy")
	fig_8_8_2
end

# ╔═╡ Cell order:
# ╟─48977b2e-5bc9-11eb-2ad7-958ca5fd2ecb
# ╠═4b149ad2-5bca-11eb-0104-0be8dd786177
# ╠═6bf4f104-5bca-11eb-2b4a-69c4ba7ffadc
# ╟─989f834a-5bca-11eb-363a-5f4afce32520
# ╠═97be6cb8-5bcb-11eb-3dc7-1b8105b206d7
# ╠═0c52737a-5bcd-11eb-1546-21885703f303
# ╠═aa061636-5bce-11eb-1381-bb00fb913f9b
# ╠═a466db90-5bcc-11eb-3385-e5117be33410
# ╠═9bf4aea2-5bcd-11eb-230f-cf63a009a0f9
# ╠═b92ba1e2-5bcb-11eb-16cf-f15929c2697b
# ╠═86538552-5bce-11eb-2cbf-ff6e175ef421
# ╠═f29c9f18-5bcd-11eb-05b5-c72249e02975
# ╠═8cba1ec8-5bcf-11eb-32c3-f38dee82139a


using Judycon, Test

#
# 1   2---3   4---5
# |   |   |   |   |
# 6---7   8   9   10
#

G1 = QuickFind(10);
G2 = QuickUnion(10, false, false);
G3 = QuickUnion(10, true, false);
G4 = QuickUnion(10, false, true);
G5 = QuickUnion(10, true, true);

function test(G)

    connect!(G, 1, 6)
    connect!(G, 6, 7)
    connect!(G, 7, 2)
    connect!(G, 2, 3)
    connect!(G, 3, 8)

    connect!(G, 9, 4)
    connect!(G, 4, 5)
    connect!(G, 5, 10)

    # `pts1` should form one set of connected components and `pts2` another, respectively.
    pts1 = [1, 2, 3, 6, 7, 8]
    pts2 = [4, 5, 9, 10]
    for p in pts1
        for q in pts1
            @test isconnected(G, p, q)
            @test isconnected(G, q, p)
        end
        for q in pts2
            @test !isconnected(G, p, q)
            @test !isconnected(G, q, p)
        end
    end
end

@testset "Test Judycon.jl" begin
    @testset "Test QuickFind" begin test(G1) end
    @testset "Test QuickUnion without weighting and path compression" begin test(G2) end
    @testset "Test QuickUnion with weighting and without path compression" begin test(G3) end
    @testset "Test QuickUnion without weighting and with path compression" begin test(G4) end
    @testset "Test QuickUnion with weighting and path compression" begin test(G5) end
end


using DelimitedFiles
using Plots

tensor = readdlm("C:path")
tensor = reshape(tensor, (420,420,31))
tensor = permutedims(tensor, [2,1,3])
tensor = tensor
heatmap(tensor[:,:,30])
contour(tensor[:,:,30], levels = 0.1:0.5)


using DelimitedFiles

filename = "./log-conv-CNS-1D-nosmooth-CFL.5"
file = readdlm(filename)

num_scheme = 2
Narr = [2;3]
Karr = [50;100;200;400;800]
block_size = 4
num_lines = size(file,1)
num_block = div(num_lines,block_size)
L1_err_dict = Dict()
for b in 1:num_block
    l = 1+block_size*(b-1)
    N = parse(Int,file[l,3][1])
    K = file[l,6]
    L1_err = file[l+1,4]
    @show N,K
    @show L1_err
    if !haskey(L1_err_dict,(N,K))
        L1_err_dict[(N,K)] = [L1_err]
    else
        push!(L1_err_dict[(N,K)],L1_err)
    end
end

num_row = length(Karr)-1
num_col = 2*length(Narr)
conv_table = zeros(num_row,num_col)
for i = 1:num_row
    for j = 1:num_col
        K = Karr[i]
        N = Narr[div(j-1,num_scheme)+1]
        s = mod1(j,num_scheme)
        conv_table[i,j] = log2(L1_err_dict[(N,K)][s]/L1_err_dict[(N,2*K)][s])
    end
end

err_table = zeros(length(Karr),2*length(Narr))
for i = 1:length(Karr)
    for j = 1:length(Narr)
        K = Karr[i]
        N = Narr[j]
        err_table[i,2*(j-1)+1] = L1_err_dict[(N,K)][1]
        err_table[i,2*(j-1)+2] = L1_err_dict[(N,K)][2]
    end
end


using SIAMFANLEquations
using SIAMFANLEquations.TestProblems
using SIAMFANLEquations.Examples
using Test
using LinearAlgebra

include("Chapter1/nsolsc_solution_test.jl")
include("Chapter1/ptcsolsc_test.jl")
include("Chapter2/basic2d_test.jl")
include("Chapter2/heq_lu_test.jl")
include("Chapter2/bvp_test.jl")
include("Chapter2/beam_test.jl")
include("Chapter2/pde_lin_test.jl")
include("Chapter2/nsolpde_test.jl")
include("Chapter2/knowsdt_test.jl")
include("Chapter3/gmres_test.jl")
include("Chapter3/mgs_test.jl")
include("Chapter3/bicgstab_test.jl")
include("Chapter3/Krylov_pde_test.jl")
include("Chapter3/ptcKrylovTest.jl")
include("Chapter3/ptcKrylovTestB.jl")
include("Chapter3/nk_test.jl")
include("Chapter3/nk_pde.jl")
include("Chapter3/nk_heq.jl")
include("Chapter4/reldiff.jl")
include("Chapter4/alex_test.jl")
include("Chapter4/ci_pde_aa.jl")
include("Chapter4/heq_aa.jl")
include("Chapter4/linear_aa.jl")

@testset "Scalar Equations: Chapter 1" begin
    @test nsolsc_solution_test()
    @test ptcsolsc_test()
end
@testset "nsol and ptcsol: Chapter 2" begin
    @test basic2d_test()
    @test bvp_test(201)
    @test beam_test()
    @test heq_lu_test()
    @test pde_lin_test(31)
    @test nsolpde_test(31)
    @test knowsdt_test()
end
@testset "Newton-Krylov solvers: Chapter 3" begin
    @test nk_test()
    @test nk_pde()
    @test nk_heq()
    @test ptcKrylovTest()
    @test ptcKrylovTestB()
end
@testset "Krylov solvers: Chapter 3" begin
    @test gmres_test()
    @test mgs_test()
    @test bicgstab_test()
    @test gmres_test_pde(31)
    @test bicgstab_test_pde(31)
end
@testset "Anderson Acceleration: Chapter 4" begin
    @test ci_pde_aa()
    @test heq_aa()
    @test linear_aa()
    @test alex_test()
end


module TestTikZWiringDiagrams
using Test

using Catlab.Theories, Catlab.WiringDiagrams
using Catlab: Graphics
using Catlab.Graphics: TikZ

to_tikz(f; kw...) = Graphics.to_tikz(f; outer_ports_layout=:fixed, kw...)

function stmts(expr::TikZ.Expression, type::Type)
  children = hasfield(typeof(expr), :stmts) ? expr.stmts : []
  vcat(expr isa type ? [expr] : [],
       mapreduce(expr -> stmts(expr, type), vcat, children; init=[]))
end
stmts(doc::TikZ.Document, type::Type) = stmts(doc.picture, type)

A, B = Ob(FreeSymmetricMonoidalCategory, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, B, A)

pic = to_tikz(f)
@test length(stmts(pic, TikZ.Node)) == 2 # Includes outer box
@test length(stmts(pic, TikZ.Edge)) == 2

pic = to_tikz(id(A))
@test length(stmts(pic, TikZ.Node)) == 1
@test length(stmts(pic, TikZ.Edge)) == 1

pic = to_tikz(compose(f,g))
@test length(stmts(pic, TikZ.Node)) == 3
@test length(stmts(pic, TikZ.Edge)) == 3

pic = to_tikz(otimes(f,g))
@test length(stmts(pic, TikZ.Node)) == 3
@test length(stmts(pic, TikZ.Edge)) == 4

pic = to_tikz(braid(A,B))
@test length(stmts(pic, TikZ.Node)) == 1
@test length(stmts(pic, TikZ.Edge)) == 2

A, B = Ob(FreeBiproductCategory, :A, :B)

for expr in (mcopy(A), delete(A), mmerge(A), create(A))
  pic = to_tikz(expr)
  @test length(stmts(pic, TikZ.Node)) == 2
end

end


using DataStructures

function readPlayer(f, N)
	p = Queue{Int}()
	for i=1:N
		k = parse(Int, readline(f))
		enqueue!(p, k)
	end
	p
end

function parsegame(N)
	f = open("22.txt")
	readline(f)
	p1 = readPlayer(f, N)
	readline(f)
	readline(f)
	p2 = readPlayer(f, N)
	(p1, p2)
end

# returns solution of part 1
function x()
	N = 25
	(p1, p2) = parsegame(N)

	while length(p1) > 0 && length(p2) > 0
		a = dequeue!(p1)
		b = dequeue!(p2)
		if a < b
			enqueue!(p2, b)
			enqueue!(p2, a)
		else
			enqueue!(p1, a)
			enqueue!(p1, b)
		end
	end
	x = 0
	for i=(2*N):-1:1
		x += dequeue!(p1) * i
	end
	x
end

function game(p1, p2)
	hist = Set{Tuple{Vector{Int}, Vector{Int}}}()
	while length(p1) > 0 && length(p2) > 0
		#println(p1)
		#println(p2)
		if (collect(p1), collect(p2)) in hist
			#println("repeat of ", p1, p2)
			return (true, p1, p2)
		end
		push!(hist, (collect(p1), collect(p2)))
		a = dequeue!(p1)
		b = dequeue!(p2)
		if a > length(p1) || b > length(p2)
			if a < b
				enqueue!(p2, b)
				enqueue!(p2, a)
			else
				enqueue!(p1, a)
				enqueue!(p1, b)
			end
		else
			q1 = copy(collect(p1))
			q2 = copy(collect(p2))
			r1 = Queue{Int}()
			r2 = Queue{Int}()
			for i=1:a
				enqueue!(r1, q1[i])
			end
			for i=1:b
				enqueue!(r2, q2[i])
			end
			p1won = game(r1, r2)[1]
			if p1won
				enqueue!(p1, a)
				enqueue!(p1, b)
			else
				enqueue!(p2, b)
				enqueue!(p2, a)
			end
		end
	end
	(length(p2) == 0, p1, p2)
end

# returns solution of part 2 (slow solution: 5.997827 seconds)
function y()
	N = 25
	(p1, p2) = parsegame(N)
	who, q1, q2 = game(p1, p2)
	x = 0
	if who
		for i=(2*N):-1:1
			x += dequeue!(q1) * i
		end
	else
		for i=(2*N):-1:1
			x += dequeue!(q2) * i
		end
	end
	x
end


"""
    initialize_dual_potential!(::AbstractDivergence, a::DiscreteMeasure) -> Nothing

Apply an initialization for the dual potential, for use in
[`unbalanced_sinkhorn!`](@ref); falls back to zeroing out the dual potential.
Specialized implementations can improve performance, but should not affect
correctness.
"""
function initialize_dual_potential!(::AbstractDivergence, a::DiscreteMeasure)
    a.dual_potential .= 0
end

"""
    function unbalanced_sinkhorn!(
        D::AbstractDivergence,
        C,
        a::DiscreteMeasure,
        b::DiscreteMeasure,
        ϵ = 1e-1;
        tol = 1e-5,
        max_iters = 10^5,
        warn::Bool = true,
    ) -> NamedTuple

Implements algorithm 1 of [[SFVTP19](@ref)]. The `dual_potential` fields of `a`
and `b` are updated to hold the optimal dual potentials. The `density`,
`log_density`, and `set` fields are not modified. The parameters are

* `D`: the [`AbstractDivergence`](@ref) used for measuring the cost of
  creating or destroying mass
* `ϵ`: the regularization parameter
* `C`: either a function from `a.set` × `b.set` to real numbers; should satisfy
  `C(x,y) = C(y,x)` and `C(x,x)=0` when applicable, or a precomputed cost matrix, see [`cost_matrix`](@ref)
* `tol`: the convergence tolerance
* `max_iters`: the maximum number of iterations to perform.
* `warn`: whether or not to warn when the maximum number of iterations is reached.

Returns a NamedTuple of the number of iterations performed (`iters`), and the
maximum residual (`max_residual`), which is the maximum infinity norm difference
between consecutive iterates of the dual potentials, at the end of the process.
If `max_iters` is not reached, iteration stops when the `max_residual` falls
below `tol`.

"""
function unbalanced_sinkhorn!(
    D::AbstractDivergence,
    C::AbstractMatrix,
    a::DiscreteMeasure,
    b::DiscreteMeasure,
    ϵ = 1e-1;
    tol = 1e-5,
    max_iters = 10^5,
    warn::Bool = true,
)
    if D isa Balanced && warn && sum(a.density) ≉ sum(b.density)
        @warn "Should have `sum(a.density) ≈ sum(b.density)` for `D==Balanced()`"
    end

    initialize_dual_potential!(D, a)
    initialize_dual_potential!(D, b)

    f = a.dual_potential
    tmp_f = a.cache

    g = b.dual_potential
    tmp_g = b.cache

    max_residual = Inf
    iters = 0
    while iters < max_iters && max_residual > tol
        iters += 1
        max_residual = 0.0
        @inbounds for j in eachindex(g)
            for i in eachindex(a.log_density, f, tmp_f)
                tmp_f[i] = a.log_density[i] + (f[i] - C[i, j]) / ϵ
            end
            new_g = -ϵ * logsumexp!(tmp_f)
            new_g = -aprox(D, ϵ, -new_g)
            diff = abs(g[j] - new_g)
            if diff > max_residual
                max_residual = diff
            end
            g[j] = new_g
        end
        @inbounds for j in eachindex(f)
            for i in eachindex(b.log_density, g, tmp_g)
                tmp_g[i] = b.log_density[i] + (g[i] - C[j, i]) / ϵ
            end
            new_f = -ϵ * logsumexp!(tmp_g)
            new_f = -aprox(D, ϵ, -new_f)
            diff = abs(f[j] - new_f)
            if diff > max_residual
                max_residual = diff
            end
            f[j] = new_f
        end
    end

    if warn && iters == max_iters
        @warn "Maximum iterations ($max_iters) reached" max_residual
    end

    return (iters = iters, max_residual = max_residual)
end

"""
    function OT!(
        D::AbstractDivergence,
        C,
        a::DiscreteMeasure,
        b::DiscreteMeasure,
        ϵ = 1e-1;
        C = (x, y) -> norm(x - y),
        kwargs...,
    ) -> Number

Computes the optimal transport cost between `a` and `b`, using
[`unbalanced_sinkhorn!`](@ref); see that function for the meaning of the
parameters and the keyword arguments. Implements Equation (15) of
[[SFVTP19](@ref)].
"""
function OT!(
    D::AbstractDivergence,
    C::AbstractMatrix,
    a::DiscreteMeasure,
    b::DiscreteMeasure,
    ϵ = 1e-1;
    kwargs...,
)

    unbalanced_sinkhorn!(D, C, a, b, ϵ; kwargs...)

    f = a.dual_potential
    g = b.dual_potential

    T = promote_type(eltype(a), eltype(b))
    _nφstar = q -> -φstar(D, -q)

    t1 = fdot(_nφstar, a.density, f)
    t2 = fdot(_nφstar, b.density, g)
    t3 = zero(T)
    for i in eachindex(f), j in eachindex(g)
        t3 -=
            ϵ *
            a.density[i] *
            b.density[j] *
            (exp((f[i] + g[j] - C[i, j]) / ϵ) - one(T))
    end
    return t1 + t2 + t3
end

# Generic method
function _sinkhorn_divergence!(
    D::AbstractDivergence,
    C,
    a::DiscreteMeasure,
    b::DiscreteMeasure,
    ϵ;
    kwargs...,
)
    OT_aa = OT!(D, C, a, a, ϵ; kwargs...)
    OT_bb = OT!(D, C, b, b, ϵ; kwargs...)
    OT_ab = OT!(D, C, a, b, ϵ; kwargs...)

    OT_ab + (-OT_aa - OT_bb + ϵ * (sum(a.density) - sum(b.density))^2) / 2
end

"""
    sinkhorn_divergence!(
        D::AbstractDivergence,
        C,
        a::DiscreteMeasure,
        b::DiscreteMeasure,
        ϵ = 1e-1;
        kwargs...,
    ) -> Number

Computes the unbalanced sinkhorn divergence between `a` and `b` as defined in
Def. 6 of [[SFVTP19](@ref)], using [`unbalanced_sinkhorn!`](@ref); see that
function for the meaning of the parameters and the keyword arguments. Sets the
optimal `dual_potential`'s of `a` and `b`.
"""
sinkhorn_divergence!(
    D::AbstractDivergence,
    C,
    a::DiscreteMeasure,
    b::DiscreteMeasure,
    ϵ = 1e-1;
    kwargs...,
) = _sinkhorn_divergence!(D, C, a, b, ϵ; kwargs...)

sinkhorn_divergence!(
    D::AbstractDivergence,
    C::AbstractMatrix,
    a::DiscreteMeasure,
    b::DiscreteMeasure,
    ϵ = 1e-1;
    kwargs...,
) = throw(ArgumentError("Must pass a cost function `C`, not a cost matrix."))

"""
    function optimal_coupling!(
        D::AbstractDivergence,
        C,
        a::DiscreteMeasure,
        b::DiscreteMeasure,
        ϵ = 1e-1;
        dual_potentials_populated::Bool = false,
        kwargs...) -> Matrix

Computes the optimal coupling between `a` and `b` using the optimal dual
potentials, the regularization parameter `ϵ`, and the cost function `C`.

If `dual_potentials_populated = false`, [`unbalanced_sinkhorn!`](@ref) is
called to populate the dual potentials of `a` and `b`, using the divergence `D`.
If `dual_potentials_populated = true`, one of [`unbalanced_sinkhorn!`](@ref)
or [`OT!`](@ref) or [`sinkhorn_divergence!`](@ref) must be called first to
set the optimal dual potentials, with the same choice of `ϵ` and `C`. In this case,
`a` and `b` are not mutated.

This function implements Prop.
6 of [[SFVTP19](@ref)].
"""
function optimal_coupling!(
    D::AbstractDivergence,
    C::AbstractMatrix,
    a::DiscreteMeasure,
    b::DiscreteMeasure,
    ϵ = 1e-1;
    dual_potentials_populated::Bool = false,
    kwargs...,
)
    if !dual_potentials_populated
        unbalanced_sinkhorn!(D, C, a, b, ϵ; kwargs...)
    end

    f = a.dual_potential
    g = b.dual_potential
    return [
        exp((f[i] + g[j] - C[i, j]) / ϵ) * a.density[i] * b.density[j]
        for i in eachindex(a.density), j in eachindex(b.density)
    ]
end

# Provide default choice of norm as cost function
for fun in [:unbalanced_sinkhorn!, :OT!, :sinkhorn_divergence!, :optimal_coupling!]
    @eval begin
        $fun(
            D::AbstractDivergence,
            a::DiscreteMeasure,
            b::DiscreteMeasure,
            args...;
            kwargs...
        ) = $fun(D, (x,y)->norm(x-y), a, b, args...; kwargs...)
    end
end

# If a function is provided, precompute the cost matrix. sinkhorn_divergence is omitted since it requires three different cost matrices.
for fun in [:unbalanced_sinkhorn!, :OT!, :optimal_coupling!]
    @eval begin
        $fun(
            D::AbstractDivergence,
            C, # Not restricting to Function to allow callable structs etc.
            a::DiscreteMeasure,
            b::DiscreteMeasure,
            args...;
            kwargs...
        ) = $fun(D, cost_matrix(C, a, b), a, b, args...; kwargs...)
    end
end


module Dec9Part2

function get_continuous_numbers(x::Array{Int64,1}, ref::Int64)::Union{Nothing,Array{Int64,1}}
    x = collect(filter(y -> y < ref, x))
    n_x = length(x)
    for i in 2:n_x
        for j in 1:n_x - (i + 1)
            if ref == sum(x[j:j + i])
                return x[j:j + i]
            end
        end
    end
    return nothing
end
function main(file_path::String)::Union{Int64,Nothing}
    numbers = parse.(Int64, readlines(file_path))
    ref = 144381670
    numbers = get_continuous_numbers(numbers, ref)
    return max(numbers...) + min(numbers...)
end
end


# written by Poom Chiarawongse <eight1911@gmail.com>

module util

    export gini, entropy, zero_one, q_bi_sort!, hypergeometric, check_input

    function assign(Y :: AbstractVector{T}, list :: AbstractVector{T}) where T
        dict = Dict{T, Int}()
        @simd for i in 1:length(list)
            @inbounds dict[list[i]] = i
        end

        _Y = Array{Int}(undef, length(Y))
        @simd for i in 1:length(Y)
            @inbounds _Y[i] = dict[Y[i]]
        end

        return list, _Y
    end

    function assign(Y :: AbstractVector{T}) where T
        set = Set{T}()
        for y in Y
            push!(set, y)
        end
        list = collect(set)
        return assign(Y, list)
    end

    @inline function zero_one(ns, n)
        return 1.0 - maximum(ns) / n
    end

    @inline function gini(ns, n)
        s = 0.0
        @simd for k in ns
            s += k * (n - k)
        end
        return s / (n * n)
    end

    # compute table of values i*log(i) for integers in 0 <= i <= maxvalue
    # where tables[i+1] = i * log(i)
    # (0*log(0) is set to 0 for convenience when computing entropy)
    function compute_entropy_terms(maxvalue)
        entropy_terms = zeros(Float64, maxvalue+1)
        for i in 1:maxvalue
            entropy_terms[i+1] = i * log(i)
        end
        return entropy_terms
    end

    # returns the entropy of ns/n, ns is an array of integers
    # and entropy_terms are precomputed entropy terms
    @inline function entropy(ns::AbstractVector{U}, n, entropy_terms) where {U <: Integer}
        s = 0.0
        for k in ns
            s += entropy_terms[k+1]
        end
        return log(n) - s / n
    end

    @inline function entropy(ns, n)
        s = 0.0
        @simd for k in ns
            if k > 0
                s += k * log(k)
            end
        end
        return log(n) - s / n
    end

    # adapted from the Julia Base.Sort Library
    @inline function partition!(v, w, pivot, region)
        i, j = 1, length(region)
        r_start = region.start - 1
        @inbounds while true
            while w[i] <= pivot; i += 1; end;
            while w[j]  > pivot; j -= 1; end;
            i >= j && break
            ri = r_start + i
            rj = r_start + j
            v[ri], v[rj] = v[rj], v[ri]
            w[i], w[j] = w[j], w[i]
            i += 1; j -= 1
        end
        return j
    end

    @inline function partitionv2!(v, w, pivot, region)
        i, j = 1, length(region)
        r_start = 0
        @inbounds while true
            while w[i] <= pivot; i += 1; end;
            while w[j]  > pivot; j -= 1; end;
            i >= j && break
            ri = r_start + i
            rj = r_start + j
            v[ri], v[rj] = v[rj], v[ri]
            w[i], w[j] = w[j], w[i]
            i += 1; j -= 1
        end
        return j
    end

    # adapted from the Julia Base.Sort Library
    function insert_sort!(v, w, lo, hi, offset)
        @inbounds for i = lo+1:hi
            j = i
            x = v[i]
            y = w[offset+i]
            while j > lo
                if x < v[j-1]
                    v[j] = v[j-1]
                    w[offset+j] = w[offset+j-1]
                    j -= 1
                    continue
                end
                break
            end
            v[j] = x
            w[offset+j] = y
        end
        return v
    end

    @inline function _selectpivot!(v, w, lo, hi, offset)
        @inbounds begin
            mi = (lo+hi)>>>1

            # sort the values in v[lo], v[mi], v[hi]

            if v[mi] < v[lo]
                v[mi], v[lo] = v[lo], v[mi]
                w[offset+mi], w[offset+lo] = w[offset+lo], w[offset+mi]
            end
            if v[hi] < v[mi]
                if v[hi] < v[lo]
                    v[lo], v[mi], v[hi] = v[hi], v[lo], v[mi]
                    w[offset+lo], w[offset+mi], w[offset+hi] = w[offset+hi], w[offset+lo], w[offset+mi]
                else
                    v[hi], v[mi] = v[mi], v[hi]
                    w[offset+hi], w[offset+mi] = w[offset+mi], w[offset+hi]
                end
            end

            # move v[mi] to v[lo] and use it as the pivot
            v[lo], v[mi] = v[mi], v[lo]
            w[offset+lo], w[offset+mi] = w[offset+mi], w[offset+lo]
            v_piv = v[lo]
            w_piv = w[offset+lo]
        end

        # return the pivot
        return v_piv, w_piv
    end

    # adapted from the Julia Base.Sort Library
    @inline function _bi_partition!(v, w, lo, hi, offset)
        pivot, w_piv = _selectpivot!(v, w, lo, hi, offset)
        # pivot == v[lo], v[hi] > pivot
        i, j = lo, hi
        @inbounds while true
            i += 1; j -= 1
            while v[i] < pivot; i += 1; end;
            while pivot < v[j]; j -= 1; end;
            i >= j && break
            v[i], v[j] = v[j], v[i]
            w[offset+i], w[offset+j] = w[offset+j], w[offset+i]
        end
        v[j], v[lo] = pivot, v[j]
        w[offset+j], w[offset+lo] = w_piv, w[offset+j]

        # v[j] == pivot
        # v[k] >= pivot for k > j
        # v[i] <= pivot for i < j
        return j
    end

    # adapted from the Julia Base.Sort Library
    # adapted from the Julia Base.Sort Library
    # this sorts v[lo:hi] and w[offset+lo, offset+hi]
    # simultaneously by the values in v[lo:hi]
    const SMALL_THRESHOLD  = 20
    function q_bi_sort!(v, w, lo, hi, offset)
        @inbounds while lo < hi
            hi-lo <= SMALL_THRESHOLD && return insert_sort!(v, w, lo, hi, offset)
            j = _bi_partition!(v, w, lo, hi, offset)
            if j-lo < hi-j
                # recurse on the smaller chunk
                # this is necessary to preserve O(log(n))
                # stack space in the worst case (rather than O(n))
                lo < (j-1) && q_bi_sort!(v, w, lo, j-1, offset)
                lo = j+1
            else
                j+1 < hi && q_bi_sort!(v, w, j+1, hi, offset)
                hi = j-1
            end
        end
        return v
    end

    # The code function below is a small port from numpy's library
    # library which is distributed under the 3-Clause BSD license.
    # The rest of DecisionTree.jl is released under the MIT license.

    # ported by Poom Chiarawongse <eight1911@gmail.com>

    # this is the code for efficient generation
    # of hypergeometric random variables ported from numpy.random
    function hypergeometric(good, bad, sample, rng)

        @inline function loggam(x)
            x0 = x
            n = 0
            if (x == 1.0 || x == 2.0)
                return 0.0
            elseif x <= 7.0
                n = Int(floor(7 - x))
                x0 = x + n
            end
            x2 = 1.0 / (x0*x0)
            xp = 6.2831853071795864769252867665590 # Tau
            gl0 = -1.39243221690590e+00
            gl0 = gl0 * x2 + 1.796443723688307e-01
            gl0 = gl0 * x2 - 2.955065359477124e-02
            gl0 = gl0 * x2 + 6.410256410256410e-03
            gl0 = gl0 * x2 - 1.917526917526918e-03
            gl0 = gl0 * x2 + 8.417508417508418e-04
            gl0 = gl0 * x2 - 5.952380952380952e-04
            gl0 = gl0 * x2 + 7.936507936507937e-04
            gl0 = gl0 * x2 - 2.777777777777778e-03
            gl0 = gl0 * x2 + 8.333333333333333e-02
            gl = gl0/x0 + 0.5*log(xp) + (x0-0.5)*log(x0) - x0
            if x <= 7.0
                @simd for k in 1:n
                    gl -= log(x0 - k)
                end
            end
            return gl
        end

        @inline function hypergeometric_hyp(good, bad, sample)
            d1 = bad + good - sample
            d2 = min(bad, good)

            Y = d2
            K = sample
            while Y > 0
                Y -= floor(UInt, rand(rng) + Y/(d1 + K))
                K -= 1
                if K == 0
                    break
                end
            end

            Z = d2 - Y
            return if good > bad
                sample - Z
            else
                Z
            end
        end

        @inline function hypergeometric_hrua(good, bad, sample)
            mingoodbad = min(good, bad)
            maxgoodbad = max(good, bad)
            popsize = good + bad
            m = min(sample, popsize - sample)
            d4 = mingoodbad / popsize
            d5 = 1.0 - d4
            d6 = m*d4 + 0.5
            d7 = sqrt((popsize - m) * sample * d4 * d5 / (popsize - 1) + 0.5)
            # d8 = 2*sqrt(2/e) * d7 + (3 - 2*sqrt(3/e))
            d8 = 1.7155277699214135*d7 + 0.8989161620588988
            d9 = floor(UInt, (m + 1) * (mingoodbad + 1) / (popsize + 2))
            d10 = (loggam(d9+1) + loggam(mingoodbad-d9+1) + loggam(m-d9+1) +
                   loggam(maxgoodbad-m+d9+1))
            d11 = min(m+1, mingoodbad+1, floor(UInt, d6+16*d7))

            while true
                X = rand(rng)
                Y = rand(rng)
                W = d6 + d8*(Y - 0.5)/X

                (W < 0.0 || W >= d11) && continue
                Z = floor(Int, W)
                T = d10 - (loggam(Z+1) + loggam(mingoodbad-Z+1) + loggam(m-Z+1) +
                           loggam(maxgoodbad-m+Z+1))
                (X*(4.0-X)-3.0) <= T && break
                (X*(X-T) >= 1) && continue
                (2.0*log(X) <= T) && break
            end

            if good > bad
                Z = m - Z
            end

            return if m < sample
                good - Z
            else
                Z
            end
        end

        return if sample > 10
            hypergeometric_hrua(good, bad, sample)
        else
            hypergeometric_hyp(good, bad, sample)
        end
    end

    function check_input(
            X                   :: AbstractMatrix{S},
            Y                   :: AbstractVector{T},
            W                   :: AbstractVector{U},
            max_features        :: Int,
            max_depth           :: Int,
            min_samples_leaf    :: Int,
            min_samples_split   :: Int,
            min_purity_increase :: Float64) where {S, T, U}
        n_samples, n_features = size(X)
        if length(Y) != n_samples
            throw("dimension mismatch between X and Y ($(size(X)) vs $(size(Y))")
        elseif length(W) != n_samples
            throw("dimension mismatch between X and W ($(size(X)) vs $(size(W))")
        elseif max_depth < -1
            throw("unexpected value for max_depth: $(max_depth) (expected:"
                * " max_depth >= 0, or max_depth = -1 for infinite depth)")
        elseif n_features < max_features
            throw("number of features $(n_features) is less than the number "
                * "of max features $(max_features)")
        elseif max_features < 0
            throw("number of features $(max_features) must be >= zero ")
        elseif min_samples_leaf < 1
            throw("min_samples_leaf must be a positive integer "
                * "(given $(min_samples_leaf))")
        elseif min_samples_split < 2
            throw("min_samples_split must be at least 2 "
                * "(given $(min_samples_split))")
        end
    end

    function check_input(
            X                   :: AbstractMatrix{S},
            Y                   :: AbstractVector{Float64},
            W                   :: AbstractVector{Int},
            indX                :: AbstractVector{Int},
            const_mtry          :: Bool,
            mpois               :: Int,
            max_depth           :: Int,
            min_samples_leaf    :: Int,
            min_samples_split   :: Int) where {S}
        n_samples, n_features = size(X)
        if length(Y) != n_samples
            throw("dimension mismatch between X and Y ($(size(X)) vs $(size(Y))")
        elseif length(W) != n_samples
            throw("dimension mismatch between X and T ($(size(X)) vs $(size(W))")
        elseif max_depth < -1
            throw("unexpected value for max_depth: $(max_depth) (expected:"
                * " max_depth >= 0, or max_depth = -1 for infinite depth)")
        elseif mpois < -1
            throw("unexpected value for max_depth: $(mpois) (expected:"
                * " mpois >= 0, or mpois = -1 for by default value)")
        elseif min_samples_leaf < 1
            throw("min_samples_leaf must be a positive integer "
                * "(given $(min_samples_leaf))")
        elseif min_samples_split < 2
            throw("min_samples_split must be at least 2 "
                * "(given $(min_samples_split))")
        end
    end

end


export Box

"""
Box set.
"""
struct Box{T <: Real} <: ConvexSet
    a::Vector{T}
    b::Vector{T}
    function Box(a::Vector{T}, b::Vector{T}) where T <: Real
        length(a) == length(b) || throw(ArgumentError("lengths of lower/upper bounds do not match"))
        new{T}(a, b)
    end
end

function project!(s::Box{T}, v::Vector{T}, y::Vector{T}) where T <: Real
    v .= clamp.(y, s.a, s.b)
end

function project(s::Box{T}, y::Vector{T}) where T <: Real
    project!(s, similar(y), y)
end


"""
`Quaternion` represents ``q_w + q_x \\bm{i} + q_y \\bm{j} + q_z \\bm{k}``.
The salar part and vector part can be accessed by `q.scalar` and `q.vector`, respectively.

# Examples
```jldoctest
julia> Quaternion(1,2,3,4)
1 + 2𝙞 + 3𝙟 + 4𝙠

julia> Quaternion(1)
1 + 0𝙞 + 0𝙟 + 0𝙠

julia> Quaternion(Vec(1,2,3))
0 + 1𝙞 + 2𝙟 + 3𝙠
```

See also [`quaternion`](@ref).

!!! note

    `Quaternion` is experimental and could change or disappear in future versions of Tensorial.
"""
struct Quaternion{T} <: Number
    data::NTuple{4, T}
    function Quaternion{T}(data::NTuple{4, Real}) where {T}
        new{T}(convert_ntuple(T, data))
    end
end

@inline Quaternion(data::NTuple{4, Any}) = Quaternion{promote_ntuple_eltype(data)}(data)
@inline (::Type{T})(data::Vararg{Any}) where {T <: Quaternion} = T(data)

# from scalar and vector
@inline Quaternion{T}(r::Real, v::Vec{3}) where {T} = @inbounds Quaternion{T}(r, v[1], v[2], v[3])
@inline Quaternion{T}(r::Real, v::Vec{2}) where {T} = @inbounds Quaternion{T}(r, v[1], v[2], zero(eltype(v)))
@inline Quaternion{T}(r::Real, v::Vec{1}) where {T} = @inbounds Quaternion{T}(r, v[1], zero(eltype(v)), zero(eltype(v)))
@inline Quaternion(r::Real, v::Vec) = Quaternion{promote_type(typeof(r), eltype(v))}(r, v)

# from vector
@inline Quaternion{T}(v::Vec{4}) where {T} = Quaternion{T}(Tuple(v))
for dim in 1:3
    @eval @inline Quaternion{T}(v::Vec{$dim}) where {T} = Quaternion{T}(zero(eltype(v)), v)
end
@inline Quaternion(v::Vec) = Quaternion{eltype(v)}(v)

# from scalar
@inline Quaternion{T}(r::Real) where {T} = (z = zero(r); Quaternion{T}(r, z, z, z))
@inline Quaternion(r::Real) = Quaternion{typeof(r)}(r)

@inline Vec(q::Quaternion) = Vec(Tuple(q))

Base.Tuple(q::Quaternion) = getfield(q, :data)

@inline function Base.getproperty(q::Quaternion, name::Symbol)
    name == :scalar && return @inbounds q[1]
    name == :vector && return @inbounds Vec(q[2], q[3], q[4])
    getfield(q, name)
end

Base.propertynames(q::Quaternion) = (:scalar, :vector, :data)

# conversion
Base.convert(::Type{Quaternion{T}}, x::Quaternion{T}) where {T} = x
Base.convert(::Type{Quaternion{T}}, x::Quaternion{U}) where {T, U} = Quaternion(map(T, Tuple(x)))
Base.convert(::Type{Quaternion{T}}, x::Real) where {T} = convert(Quaternion{T}, Quaternion(x))

# promotion
Base.promote_rule(::Type{Quaternion{T}}, ::Type{T}) where {T <: Real} = Quaternion{T}
Base.promote_rule(::Type{Quaternion{T}}, ::Type{U}) where {T <: Real, U <: Real} = Quaternion{promote_type(T, U)}
Base.promote_rule(::Type{Quaternion{T}}, ::Type{Quaternion{T}}) where {T <: Real} = Quaternion{T}
Base.promote_rule(::Type{Quaternion{T}}, ::Type{Quaternion{U}}) where {T <: Real, U <: Real} = Quaternion{promote_type(T, U)}

# used for `isapprox`
Base.real(q::Quaternion) = q.scalar
Base.isfinite(q::Quaternion) = prod(map(isfinite, Tuple(q)))

"""
    quaternion(θ, n::Vec; [normalize = true, degree = false])

Construct `Quaternion` from angle `θ` and axis `n` as

```math
q = \\cos\\frac{\\theta}{2} + \\bm{n} \\sin\\frac{\\theta}{2}
```

The constructed quaternion is normalized such as `norm(q) ≈ 1` by default.

# Examples
```jldoctest
julia> q = quaternion(π/4, Vec(0,0,1))
0.9238795325112867 + 0.0𝙞 + 0.0𝙟 + 0.3826834323650898𝙠

julia> x = rand(Vec{3})
3-element Vec{3, Float64}:
 0.5908446386657102
 0.7667970365022592
 0.5662374165061859

julia> (q * x / q).vector ≈ rotmatz(π/4) ⋅ x
true
```
"""
function quaternion(::Type{T}, θ::Real, x::Vec{3}; normalize::Bool = true, degree::Bool = false) where {T}
    if degree
        θ = deg2rad(θ)
    end
    ϕ = θ / 2
    if normalize
        n = LinearAlgebra.normalize(x) * sin(ϕ)
    else
        n = x * sin(ϕ)
    end
    @inbounds Quaternion{T}(cos(ϕ), n)
end
quaternion(T::Type, θ::Real, x::Vec{2}; normalize::Bool = true, degree::Bool = false) =
    @inbounds quaternion(T, θ, Vec(x[1], x[2], 0); normalize, degree)
quaternion(θ::Real, x::Vec; normalize::Bool = true, degree::Bool = false) =
    quaternion(promote_type(typeof(θ), eltype(x)), θ, x; normalize, degree)

Base.length(::Quaternion) = 4
Base.size(::Quaternion) = (4,)

@inline function Base.getindex(q::Quaternion, i::Int)
    @boundscheck 1 ≤ i ≤ 4 || throw(BoundsError(q, i))
    @inbounds Tuple(q)[i]
end

# quaternion vs quaternion
@inline Base.:-(q::Quaternion) = Quaternion(-Vec(q))
@inline Base.:+(q::Quaternion, p::Quaternion) = Quaternion(Vec(q) + Vec(p))
@inline Base.:-(q::Quaternion, p::Quaternion) = Quaternion(Vec(q) - Vec(p))
@inline Base.:/(q::Quaternion, p::Quaternion) = q * inv(p)
@inline function Base.:*(q::Quaternion, p::Quaternion)
    q₁, q₂, q₃, q₄ = Tuple(q)
    A = @Mat [ q₁ -q₂ -q₃ -q₄
               q₂  q₁ -q₄  q₃
               q₃  q₄  q₁ -q₂
               q₄ -q₃  q₂  q₁ ]
    Quaternion(A ⋅ Vec(p))
end

# quaternion vs number
@inline Base.:*(a::Number, q::Quaternion) = Quaternion(a * Vec(q))
@inline Base.:*(q::Quaternion, a::Number) = Quaternion(Vec(q) * a)
@inline Base.:/(q::Quaternion, a::Number) = Quaternion(Vec(q) / a)

# quaternion vs vector
@inline Base.:*(q::Quaternion, v::Vec) = q * Quaternion(v)
@inline Base.:*(v::Vec, q::Quaternion) = Quaternion(v) * q
@inline Base.:/(v::Vec, q::Quaternion) = v * inv(q)

function angleaxis(q::Quaternion)
    a = norm(q.vector)
    θ = 2atan(a, q.scalar)
    n = q.vector / a
    θ, n
end

"""
    rotate(x::Vec, q::Quaternion)

Rotate `x` by quaternion `q`.

# Examples
```jldoctest
julia> v = Vec(1.0, 0.0, 0.0)
3-element Vec{3, Float64}:
 1.0
 0.0
 0.0

julia> rotate(v, quaternion(π/4, Vec(0,0,1)))
3-element Vec{3, Float64}:
 0.7071067811865475
 0.7071067811865476
 0.0
```
"""
@inline rotate(v::Vec{3}, q::Quaternion) = (q * v / q).vector
@inline rotate(v::Vec{2}, q::Quaternion) = (v = (q * v / q).vector; @inbounds Vec(v[1], v[2]))

@inline Base.conj(q::Quaternion) = Quaternion(q.scalar, -q.vector)
@inline Base.abs2(q::Quaternion) = (v = Vec(q); dot(v, v))
@inline Base.abs(q::Quaternion) = sqrt(abs2(q))
@inline norm(q::Quaternion) = abs(q)
@inline inv(q::Quaternion) = conj(q) / abs2(q)

"""
    exp(::Quaternion)

Compute the exponential of quaternion as

```math
\\exp(q) = e^{q_w} \\left( \\cos\\| \\bm{v} \\| + \\frac{\\bm{v}}{\\| \\bm{v} \\|} \\sin\\| \\bm{v} \\| \\right)
```
"""
function Base.exp(q::Quaternion)
    v = q.vector
    v_norm = norm(v)
    if v_norm > 0
        n = v / v_norm
    else
        n = zero(v)
    end
    exp(q.scalar) * quaternion(2*v_norm, n; normalize = false)
end

"""
    log(::Quaternion)

Compute the logarithm of quaternion as

```math
\\ln(q) = \\ln\\| q \\| + \\frac{\\bm{v}}{\\| \\bm{v} \\|} \\arccos\\frac{q_w}{\\| q \\|}
```
"""
function Base.log(q::Quaternion)
    q_norm = norm(q)
    Quaternion(log(q_norm), normalize(q.vector) * acos(q.scalar/q_norm))
end

@inline normalize(q::Quaternion) = q / norm(q)

function rotmat_normalized(q::Quaternion)
    q₁, q₂, q₃, q₄ = Tuple(q)
    q₁² = q₁ * q₁
    q₂² = q₂ * q₂
    q₃² = q₃ * q₃
    q₄² = q₄ * q₄
    q₁q₂ = q₁ * q₂
    q₂q₃ = q₂ * q₃
    q₃q₄ = q₃ * q₄
    q₁q₃ = q₁ * q₃
    q₁q₄ = q₁ * q₄
    q₂q₄ = q₂ * q₄
    @Mat [q₁²+q₂²-q₃²-q₄² 2(q₂q₃-q₁q₄)    2(q₂q₄+q₁q₃)
          2(q₂q₃+q₁q₄)    q₁²-q₂²+q₃²-q₄² 2(q₃q₄-q₁q₂)
          2(q₂q₄-q₁q₃)    2(q₃q₄+q₁q₂)    q₁²-q₂²-q₃²+q₄²]
end

"""
    rotmat(::Quaternion)

Construct rotation matrix from quaternion.

# Examples
```jldoctest
julia> q = quaternion(π/4, Vec(0,0,1))
0.9238795325112867 + 0.0𝙞 + 0.0𝙟 + 0.3826834323650898𝙠

julia> rotmat(q)
3×3 Tensor{Tuple{3, 3}, Float64, 2, 9}:
 0.707107  -0.707107  0.0
 0.707107   0.707107  0.0
 0.0        0.0       1.0
```
"""
@inline rotmat(q::Quaternion) = rotmat_normalized(normalize(q))

function Base.show(io::IO, q::Quaternion)
    pm(x) = x < 0 ? " - $(-x)" : " + $x"
    print(io, q[1], pm(q[2]), "𝙞", pm(q[3]), "𝙟", pm(q[4]), "𝙠")
end


"""
	friction
"""
struct FrictionConstraints <: Constraints
    n
    ineq
    n_stage
end

function friction!(c, x⁺, u, model)

    v = x⁺[3]
    β = u[2:3]
    ψ = u[4]
    η = u[5:6]
    s = u[7]

    n = (model.mc + model.mp) * model.g

    c[1] = v + ψ - η[1]
    c[2] = -v + ψ - η[2]
    c[3] = model.μ * n - sum(β)
    c[4] = s - ψ * (model.μ * n - sum(β))
    c[5] = s - β' * η

    return nothing
end

function constraints!(c, Z, con::FrictionConstraints, model, idx, h, T)
    for t = 1:T-1
        x⁺ = view(Z, idx.x[t + 1])
        u = view(Z, idx.u[t])
        friction!(view(c, (t - 1) * con.n_stage .+ (1:con.n_stage)), x⁺, u, model)
    end

    nothing
end

function constraints_jacobian!(∇c, Z, con::FrictionConstraints, model, idx, h, T)
    n = model.n
    m = model.m

    c_tmp = zeros(con.n_stage)
    shift = 0

    for t = 1:T-1
        x⁺ = view(Z, idx.x[t + 1])
        u = view(Z, idx.u[t])

        cx!(c, z) = friction!(c, z, u, model)
        cu!(c, z) = friction!(c, x⁺, z, model)

        r_idx = (t - 1) * con.n_stage .+ (1:con.n_stage)

        c_idx = idx.x[t + 1]
        len = length(r_idx) * length(c_idx)
        ForwardDiff.jacobian!(reshape(view(∇c, shift .+ (1:len)), con.n_stage, n),
            cx!, c_tmp, x⁺)
        shift += len

        c_idx = idx.u[t]
        len = length(r_idx) * length(c_idx)
        ForwardDiff.jacobian!(reshape(view(∇c, shift .+ (1:len)), con.n_stage, m),
            cu!, c_tmp, u)
        shift += len
    end

    return nothing
end

function constraints_sparsity(con::FrictionConstraints, model, idx, T;
    shift_row = 0, shift_col = 0)
    row = []
    col = []

    for t = 1:T-1
        r_idx = shift_row + (t - 1) * con.n_stage .+ (1:con.n_stage)

        c_idx = shift_col .+ idx.x[t + 1]
        row_col!(row, col, r_idx, c_idx)

        c_idx = shift_col .+ idx.u[t]
        row_col!(row, col, r_idx, c_idx)
    end

    return collect(zip(row, col))
end


### A Pluto.jl notebook ###
# v0.17.0

using Markdown
using InteractiveUtils

# ╔═╡ dc9d1ef7-23e2-4bc3-ab47-e9ca80af5811
begin
	using PrettyTables
	using PlutoUI
end

# ╔═╡ bd0378cb-ca70-4803-b5fc-7bd08ad1b9f3
md"""
#### Practice Problem: Multiple Chemical Reactions in a Single CSTR

Component $A$ is converted to $B$ (desired product) and $C$ (undesired product) in single well-mixed continuous stirred tank reactor (CSTR) with a single input and output and the following reaction chemistry:

$$\begin{eqnarray}
\hat{r}_{1}: A &\rightleftharpoons& B\\
\hat{r}_{2}: A &\rightleftharpoons& C\\
\end{eqnarray}$$

The substrate $A$ (component 1) is supplied to the reactor at 10 mol/hr, where the input stream has a volumetric flow rate of $\dot{F}_{1}$ = 2 L/hr. The fractional conversion of $A$ was measured and found to be equal to 70%. Lastly, the extent of the desired reaction was estimated to $\dot{\epsilon}_{1}$ = 1.5 mol/hr.

__Assumptions__:
* The CSTR is well-mixed, operating at steady-state, and has a constant volume V = 14L
* The CSTR has a single input stream and a single output stream
* Neither of the chemical reactions occur in the streams
* There is no $B$ or $C$ in the input stream

__Compute__:
* The chemical reaction rate(s) $\hat{r}_{1}$ and $\hat{r}_{2}$
* The mol flow rates of all components in the CSTR exit stream
* If $\hat{r}_{1}$ and $\hat{r}_{2}$ are irreversible, and first-order in $A$, compute the forward rate constants $k_{1}$ and $k_{2}$ for reaction 1 and 2, respectively.

"""

# ╔═╡ 6b07d7ae-8b5d-4c41-8184-9329cdaa94a6
begin

	# Setup problem parameters -
	F_dot = 2.0 	# units: L/hr
	f1 = 0.70 		# units: dimensionless
	ϵ₁_dot = 1.5 	# units: mol/hr
	n11_dot = 10.0 	# units: mol/hr
	V = 14.0 		# units: L

	# setup stoichiometric matrix for "easy" computation -
	S = [

		# ϵ₁_dot ϵ₂_dot
		-1.0 -1.0 	; # 1 A
		1.0 0.0 	; # 2 B
		0.0 1.0 	; # 3 C
	];

	# setup input flow rate -
	n_dot_in = [
		n11_dot ; 	# 1 A
		0.0 	; 	# 2 B
		0.0 	; 	# 3 C
	];

	# show -
	nothing
end

# ╔═╡ 7e8aaae7-555f-4b70-bae7-4adc1d1c9923
md"""
#### Compute: the chemical reaction rate(s) $\hat{r}_{1}$ and $\hat{r}_{2}$

To compute the reaction rates $\hat{r}_{1}$ and $\hat{r}_{2}$, we first must estimate the open extent for the second reaction $\dot{\epsilon}_{2}$
by relating the $\dot{\epsilon}_{2}$ to the fractional conversion of A. To start, we know that for component $i$ the species mole balance is given by:

$$\dot{n}_{i1} - \dot{n}_{i2} + \sum_{j=1}^{\mathcal{R}}\sigma_{ij}\dot{\epsilon}_{j} = 0$$

where $\epsilon_{j} = \hat{r}_{j}V$. Rearranging the species mole balance for $\dot{n}_{i2}$:

$$\dot{n}_{i2} = \dot{n}_{i1} + \sum_{j=1}^{\mathcal{R}}\sigma_{ij}\dot{\epsilon}_{j}$$

However, the fractional conversion is given by:

$$\dot{n}_{i2} = \dot{n}_{i1} - f_{i}\dot{n}_{i1}$$

Thus, we can eliminate $\dot{n}_{i2}$ from the species mol balance ans solve for the fractional conversion $f_{i}$:

$$f_{i} = -\frac{1}{\dot{n_{i,1}}}\left(\sum_{j=1}^{\mathcal{R}}\sigma_{ij}\dot{\epsilon}_{j}\right)\qquad{i=1}$$

For the currenr problem, we know the fractional concersion of A (i=1), and we have two reactions where $\dot{\epsilon}_{2}$ is known, thus
we can solve for $\dot{\epsilon}_{2}$:

$$\dot{\epsilon}_{2} = \dot{n}_{11}f_{1} - \dot{\epsilon}_{1}$$

Once we have $\dot{\epsilon}_{1}$ and $\dot{\epsilon}_{2}$, we can compute the rates: $\hat{r}_{j}=\dot{\epsilon}_{j}V^{-1}$.

"""

# ╔═╡ fddd2b82-8eba-4517-bdfd-0953bd10c259
begin

	# compute -
	ϵ₂_dot = n11_dot*f1  - ϵ₁_dot;

	# display -
	with_terminal() do
		# compute ϵ₂_dot:
		println("ϵ₂_dot = $(ϵ₂_dot) mol/hr")
	end
end

# ╔═╡ 86964693-9968-4524-a2bc-8c3a8f202cc1
begin

	# compute the rates -
	r1_hat = (1/V)*ϵ₁_dot
	r2_hat = (1/V)*ϵ₂_dot

	# display -
	with_terminal() do
		println("r1_hat = $(r1_hat) mol/L-hr and r2_hat = $(r2_hat) mol/L-hr")
	end
end

# ╔═╡ 34ddef24-462a-4a1d-8a61-9c760f278883
md"""

#### Compute: the mole flow rates of all components in the CSTR exit stream

We know have both $\dot{\epsilon}_{1}$ and $\dot{\epsilon}_{2}$, so we can use the steady-state species mole balance for component $i=1,2,3$

$$\dot{n}_{i1} - \dot{n}_{i2} + \sum_{j=1}^{\mathcal{R}}\sigma_{ij}\dot{\epsilon}_{j} = 0\qquad{i=1,2,3}$$

to compute the composition of the exit stream.

"""

# ╔═╡ fff53850-a627-4349-bef2-c5aaf8828c3d
begin

	# setup extent vector -
	ϵ_vector = [ϵ₁_dot ; ϵ₂_dot];

	# compute n_dot coming out of reactor -
	n_dot_out = n_dot_in + S*ϵ_vector

	# setup data array =
	state_table_array = Array{Any,2}(undef, 3,3)

	# A -
	state_table_array[1,1] = "A"
	state_table_array[1,2] = n_dot_in[1]
	state_table_array[1,3] = n_dot_out[1]

	# B -
	state_table_array[2,1] = "B"
	state_table_array[2,2] = n_dot_in[2]
	state_table_array[2,3] = n_dot_out[2]

	# C -
	state_table_array[3,1] = "C"
	state_table_array[3,2] = n_dot_in[3]
	state_table_array[3,3] = n_dot_out[3]

	# setup a pretty table =
	# header row -
	state_table_header_row = (["Species","n_dot_in","n_dot_out"],["","mol/hr","mol/hr"]);

	with_terminal() do
		# write the table -
		pretty_table(state_table_array; header=state_table_header_row)
	end
end

# ╔═╡ 8b5a451f-5d31-49f5-91fb-3e24a00b87ac
md"""
#### Compute the forward rate constants $k^{+}_{1}$ and $k^{+}_{2}$

To compute the rate constants, we can solve the rate expression $\hat{r}_{j} = k^{+}_{j}\left[A\right]$ once we have the concentration of A.
To compute concentration, we know that in stream $s$

$$\left[C_{i,s}\right] = \dot{n}_{i,s}\left(\dot{F}_{s}\right)^{-1}$$

where $\left[C_{i,s}\right]$ denotes the concentration of component $i$ in stream $s$, and $\dot{F}_{s}$ denotes the volumetric flow rate of stream $s$. Because the reactor is well-mixed, and is constant volume, the concentration of $A$ (denoted in index notation as $C_{12}$) is given by:

$$C_{12} = \dot{n}_{12}\left(\dot{F}\right)^{-1}$$

This, the rate constant(s) are given by:

$$k_{j}^{+} = \left(\frac{\dot{F}}{\dot{n}_{12}}\right)\hat{r}_{j}\qquad{j=1,2}$$

"""

# ╔═╡ 86b16793-3647-4bd6-870e-3ca2be3035ae
begin

	# concetration -
	A_concentration = (n_dot_out[1]/F_dot)

	# compute -
	k1f = r1_hat*(1/A_concentration)
	k2f = r2_hat*(1/A_concentration)

	# display -
	with_terminal() do
		println("k1f = $(k1f) 1/hr and k2f = $(k2f) 1/hr")
	end
end

# ╔═╡ 4bab103c-3d4b-11ec-1ceb-9566ca26ce6b
html"""
<style>
main {
    max-width: 1200px;
    width: 85%;
    margin: auto;
    font-family: "Roboto, monospace";
}

a {
    color: blue;
    text-decoration: none;
}

.H1 {
    padding: 0px 30px;
}
</style>"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"

[compat]
PlutoUI = "~0.7.19"
PrettyTables = "~1.2.3"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "0bc60e3006ad95b4bb7497698dd7c6d649b9bc06"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.1"

[[ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[DataAPI]]
git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.9.0"

[[DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

[[IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "ae4bbcadb2906ccc085cf52ac286dc1377dceccc"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.1.2"

[[Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "e071adf21e165ea0d904b595544a8e514c8bb42c"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.19"

[[PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "d940010be611ee9d67064fe559edbb305f8cc0eb"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.2.3"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "TableTraits", "Test"]
git-tree-sha1 = "fed34d0e71b91734bf0a7e10eb1bb05296ddbcd0"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.6.0"

[[Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─bd0378cb-ca70-4803-b5fc-7bd08ad1b9f3
# ╠═6b07d7ae-8b5d-4c41-8184-9329cdaa94a6
# ╟─7e8aaae7-555f-4b70-bae7-4adc1d1c9923
# ╠═fddd2b82-8eba-4517-bdfd-0953bd10c259
# ╠═86964693-9968-4524-a2bc-8c3a8f202cc1
# ╟─34ddef24-462a-4a1d-8a61-9c760f278883
# ╟─fff53850-a627-4349-bef2-c5aaf8828c3d
# ╟─8b5a451f-5d31-49f5-91fb-3e24a00b87ac
# ╠═86b16793-3647-4bd6-870e-3ca2be3035ae
# ╠═dc9d1ef7-23e2-4bc3-ab47-e9ca80af5811
# ╟─4bab103c-3d4b-11ec-1ceb-9566ca26ce6b
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


md"""
# Элементы функционального программирования

## Функции - как объекты 1-го класса
"""
x = 0

s = sin  # <: Function

md"""
Пусть, например, мы хотим перемещать робота в заданном направлении до тех пор, пока не наступит некоторое событие,
выраженное соответствующим условием.

Это реализуется следующей функцией:
"""

function movements!(condition::Function, r, side)
    n=0
    while !condition(r)
        move!(r,side)
        n+=1
    end
    return n
end

#Теперь эту функцию можно вызывать, например, так:

movements!(isborder, r, side)

md"""
## Анонимные функции
Анонимные функции представляют собой функциональные значения (<: Function) без имени
 (подобно тому, как бывают значения типа Int без имени)

Например:

condition = robot -> !isborder(robot)

Здесь переменной condition присвоено некоторое функциональное значение (значение анонимной функции).

Другие римеры:

f = x -> 2x

f(2)
4

(x -> 2x)(3)
6

f2= (x,y) -> x^2+y^2

f2(2,3)
13

g= _ -> 1

g(1)
1

g(2)
1

g(3)
1

g()
ERROR

g(1,2)
ERROR

g2 = (_,_) -> 1

g2(1,2)
1

Анонимные функции часто также называю ЛЯМБДА- функциями
(это название пришло из соответствующего раздела математической логики, из так называемого лямбда-исчислниия)

В языке Python, например, определение фанкции

x -> 2x

выглядело бы так:

lambda x: 2*x

При этом, однако, в отличие от языка Python, в языке Julia нет тех сильных ограничений, которые есть в языке Python на вид выражения,
определяющее анонимную функцию.

## Функции высшего порядка

Функция называется  функцией высшего порядка, если она имеет хотя бы один аргумент функционального типа
или если она возвращает значение функционального типа.

Пример встроенной функции высшего порядка - map

julia> map(x->x^2, [1,2,3,4])
4-element Vector{Int64}:
  1
  4
  9
 16

julia> map(x->x^2, (1,2,3,4))
(1, 4, 9, 16)

julia> map(sin, (1,2,3,4))
(0.8414709848078965, 0.9092974268256817, 0.1411200080598672, -0.7568024953079282)

## do-синтаксис

julia> map([1,2,3,4,5]) do x
    x^2
end
5-element Vector{Int64}:
1
4
9
16
25

julia> map([1,2,3,4,5]) do _
    0
end
5-element Vector{Int64}:
0
0
0
0
0

do-синтаксис применим не только к встоеннной функции map, но и к любой другой, у котрой первый аргумент
имеет функциональный тип.

Этот синтаксис просто позволяет определить фактическое значение этого параметра в этой компактной конструкци,
без того, чтобы определять эту функцию отдельно стандартным образом.

## Пример функции высшего порядка, возвращающей функциональное значение. Понятие замыкания.

julia> function g(a)
           return x -> a*x # возвращается функциональное значение (функция)
       end
g (generic function with 1 method)

julia> h = g(2)
#33 (generic function with 1 method)

julia> h(2)
4

julia> h(3)
6

julia> h = g(3)
#33 (generic function with 1 method)

julia> h(3)
9

## Примеры практического использования функций высшего порядка.

Ранее были у на  такие функции

    movements!(robot, side) # перемещает робота "до упора" в заданном направлении и возвращает число сделанных "шагов"

    movements!(robot, side, num_steps) # делает num_steps шагов в заданном напралении

Их можно было бы дополнить ещё двумя методами - двумя функциями высшего порядка:

    movements!(condition::Function, robot, side)

и, соответственно,

    movements!(condition::Function, robot, side, num_steps)

где в в обоих случаях через функциональный параметр condition в функцию передаётся условие,
только и только при выполнении которого робот должен перемещаться.
При этом во втором случае через параметр num_steps передается еще наибольшее число шагов, которое робот может
сделать в заданном направлении (даже если условие, определяемое condition все время выполняется).

Функции высшего порядка movements! (новые методы функции movements!, в терминологии языка julia) могут быть
реализованы так, что функция condition будет иметь один аргумент (robot), или, в другом варианте реализации,
вообще не будут иметь аргументов.

### Варианты реализации функции movements!(condition::Function, robot, side)

1-вариант
"""
"""
movements!(condition::Function, robot, side)

-- возвращает число шагов, сделанных "до тех пор, пока выпоняется условие condition"

-- condition - функция одного аргумента (типа робот), возвращающая логическое значение
"""
function movements!(condition::Function, r, side)

    while condition(r) == true
        move!(r,side)
    end
end

# Примеры использования этой функции:

r=Ribot() # тип робота мог бы быть и другим

movements!(r->!isborder(r,Ost), r, Ost) # - это вызов функции

movements!(r, Ost) do robot
    robot->!isborder(robot, Ost)
end # - это вызов функции с использованием do-синтаксиса

movements!(ismarker, r, Ost)

movements!(robot -> ismarker(robot), r, Ost)

md"""
2-ой вариант реализации той же функции movements!(condition::Function, r, side)
(наверное, более предпочтительный)
"""

"""
movements!(condition::Function, robot, side)

-- возвращает число шагов, сделанных "до тех пор, пока выпоняется условие condition"

-- condition - функция без аргументов, возвращающая логическое значение
"""
function movements!(condition::Function, robot, side)
    #  число шагов, сделанных "до тех пор, пока выпоняется условие condition"
    while condition() == true
        move!(robot, side)
    end
end

# Пример использования так определенной функции movements!:

r=Ribot() # тип робота мог бы быть и другим

movements!(() -> !isborder(r, Ost), r, Ost) # - это вызов функции

movements!(r, Ost) do
    robot->!isborder(r, Ost)
end # - это вызов функции с использованием do-синтаксиса

md"""
### Реализация функции movements!(condition::Function, robot, side)

Определим эту функцию (обобщенную), так, чтобы она перемещала робота в заданном направлении не более чем на
num_steps шагов, пока выпоняется условие condition()==true". При этом она будет возвращать число сделанных роботом
"шагов" до своей остановки.
"""

"""
movements!(condition::Function, robot, side, num_steps)

-- делает не более чем num_steps шагов в заданном направлении, пока выпоняется условие condition()==true",
и возвращает число сделанных шагов

-- condition - функция без аргументов, возвращающая логическое значение
"""
function movements!(condition::Function, r, side, num_steps)
    n = 0
    for _ in 1:num_steps
        if condition() == true
            move!(r, side)
            n += 1
        else
            return n
        end
    end
    return n
end

md"""
Такая функция была бы удобной, например, для решения задачи, в которой требуется
на неограниченном поле нати маркер (была у нас такая задача).

Так же её можно было бы использовать для реализации функции try_move!(robot, side::HorizonSide).

Рассмотрим ещё примеры некоторых других полезных (обобщенных) функций высшего порядка, которыми стоило бы
поплнить нашу библиотеку функций.

1. Функция

    shuttle(condtion::Function, robot, side)

перемещающая робота взад-вперед (челноком, as shuttle) с увелисением амплитуды перемещений
до тех пор, пока не будет выполено заданное условие,
где side - начальное направление,

2. Функция, перемещающая робота змейкой

    snake(condition::Function, robot, side::NTuple{HorizonSide})

до тех пор, пока не будет выполено заданное условие condition()==true,
где side - пара направлений, определяющих траекторию змейки.

3. Функция, перемещающая робота по спирали

    spiral(condition::Function, robot, side)

до тех пор, пока не будет выполено заданное условие condition()==true,
где side - начальное направление.
"""


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export wfst2tr
"""
  `wfst2tr(fst,Nt; get_label=get_ilabel)`

Returns an array `time2tr` of length `Nt-1`.
The `t`-th element of `time2tr` is a dictionary mapping the transition index between input label at time `t` to the input label at time `t+1`.

```julia
julia> using FiniteStateTransducers

julia> H = WFST(["a1","a2","a3"],["a"]);

julia> add_arc!(H,1,2,"a1","a");

julia> add_arc!(H,2,3,"a2","<eps>");

julia> add_arc!(H,2,2,"a2","<eps>");

julia> add_arc!(H,3,4,"a3","<eps>");

julia> add_arc!(H,3,2,"a3","a");

julia> initial!(H,1);

julia> final!(H,4)
WFST #states: 4, #arcs: 5, #isym: 3, #osym: 1
|1/0.0f0|
a1:a/0.0f0 → (2)
(2)
a2:ϵ/0.0f0 → (3)
a2:ϵ/0.0f0 → (2)
(3)
a3:ϵ/0.0f0 → (4)
a3:a/0.0f0 → (2)
((4/0.0f0))

julia> Nt = 10;

julia> time2tr = wfst2tr(H,Nt)
9-element Array{Dict{Int64,Array{Int64,1}},1}:
 Dict(1 => [2])
 Dict(2 => [2, 3])
 Dict(2 => [2, 3],3 => [2])
 Dict(2 => [2, 3],3 => [2])
 Dict(2 => [2, 3],3 => [2])
 Dict(2 => [2, 3],3 => [2])
 Dict(2 => [2, 3],3 => [2])
 Dict(2 => [2],3 => [2])
 Dict(2 => [3])

```

"""
function wfst2tr(fst::WFST{W,A,D},Nt; get_label=get_ilabel) where {W,A,D}
  isym=get_isym(fst)
  Ns = length(isym)
  X = matrix2wfst(isym,ones(W,Ns,Nt); W = W)

  Z = ∘(X,fst; connect=true)

  i = get_initial(Z)
  current_states = Set{D}([i])
  next_states = Set{D}()
  time2transitions = [Dict{D,Vector{D}}() for t in 1:Nt-1]
  t = 1
  while !isempty(current_states)
    dict = Dict{D,Set{D}}()
    for c in current_states
      for c_arc in Z[c]
        src = D(get_label(c_arc))
        n = get_nextstate(c_arc)
        push!(next_states,n)
        dest = Set{D}([D(get_label(n_arc)) for n_arc in Z[n]])
        if !isempty(dest)
          if haskey(dict,src)
            union!(dict[src],dest)
          else
            dict[src] = dest
          end
        end
      end
    end
    if t < Nt
      time2transitions[t] = Dict(k=>[dict[k]...] for k in keys(dict))
    else
      if !isempty(dict)
        error("something is wrong with the graph")
      end
    end
    t += 1
    current_states, next_states = next_states, Set{D}()
  end
  return time2transitions
end

"""
  `wfst2tr(fst; get_label=get_ilabel, convert_weight=identity)`

Extract transition matrix and initial probabilities from a WFST representing an HMM model.

```julia
julia> add_arc!(H,1,2,"a1","a");

julia> add_arc!(H,2,2,"a1","<eps>",-log(0.3));

julia> add_arc!(H,2,3,"a2","<eps>",-log(0.7));

julia> add_arc!(H,3,3,"a2","<eps>",-log(0.3));

julia> add_arc!(H,3,4,"a3","<eps>",-log(0.7));

julia> add_arc!(H,4,4,"a3","<eps>",-log(0.3));

julia> add_arc!(H,4,2,"a1","a",-log(0.7));

julia> initial!(H,1);

julia> final!(H,4)
WFST #states: 4, #arcs: 12, #isym: 3, #osym: 1
|1/0.0f0|
a1:a/0.0f0 → (2)
a1:a/0.0f0 → (2)
(2)
a2:ϵ/0.0f0 → (3)
a2:ϵ/0.0f0 → (2)
a1:ϵ/1.2039728f0 → (2)
a2:ϵ/0.35667494f0 → (3)
(3)
a3:ϵ/0.0f0 → (4)
a3:a/0.0f0 → (2)
a2:ϵ/1.2039728f0 → (3)
a3:ϵ/0.35667494f0 → (4)
((4/0.0f0))
a3:ϵ/1.2039728f0 → (4)
a1:a/0.35667494f0 → (2)

julia> a,A = wfst2tr(H; convert_weight = w -> exp(-get(w)))
(Float32[1.0, 0.0, 0.0], Float32[0.3 0.7 0.0; 0.0 0.3 0.7; 0.7 0.0 0.3])
```
"""
function wfst2tr(fst::WFST{W,A,D};
                 get_label=get_ilabel, convert_weight=identity) where {W,A,D}
  isym=get_isym(fst)
  Ns = length(isym)
  W2 = typeof(convert_weight(one(W)))
  a = zeros(W2,Ns)
  H = zeros(W2,Ns,Ns)
  init = get_initial(fst)
  state2outtr = Dict(i => (get_ilabel.(s),get_weight.(s)) for (i,s) in enumerate(fst))
  for (p,s,n,d,e,a) in FiniteStateTransducers.DFS(fst,init)
    if d
      intr = get_ilabel(a)
      outtr, w = state2outtr[n]
      for i in eachindex(outtr)
        H[intr,outtr[i]] = convert_weight(w[i])
      end
    end
  end
  for arc in fst[init]
    a[get_label(arc)] = convert_weight(get_weight(arc))
  end
  return a, H
end


# CENG0038 Droplet Distribution research project

dir = "./src"
push!(LOAD_PATH,dir)

module researchProject

import trachea
import bifurcations

# straightTube(T, ro, Q, mu, ro_d, d_d, d, phi, L)
T = 309.5 # Temperature of Air
ro = 1.12027 # Density of Air (http://www.anaesthesia.med.usyd.edu.au/resources/lectures/humidity_clt/humidity.html#:~:text=While%20nose%20breathing%20at%20rest,%25%20to%2070%25%20relative%20humidity.)
Q = 0.001 # Based on 60L/min flowrate of ventilators
mu = 1.895e-5 # Dynamic viscosity of air (https://www.engineersedge.com/physics/viscosity_of_air_dynamic_and_kinematic_14483.htm#:~:text=The%20viscosity%20of%20air%20depends,10%2D5%20Pa%C2%B7s%20.&text=At%2025%20%C2%B0C%2C%20the,the%20kinematic%20viscosity%2015.7%20cSt.)
ro_d = 1000 # Droplet density of water
d = 0.0189 # Diameter of trachea (https://www.google.com/url?sa=i&url=https%3A%2F%2Fwww.slideshare.net%2Fjinojustinj%2Ftracheal-pathologies&psig=AOvVaw38sp20tGz3PGVizR_R8F_N&ust=1605869019853000&source=images&cd=vfe&ved=0CAIQjRxqFwoTCICtjbi2ju0CFQAAAAAdAAAAABAT)
phi = (pi/2) # Inclination angle in radians (phi) SI units: rad
# IBA = 1.040216 # Interbronchial angle, average for all ages, 59.6 degrees SI Units: radians
# phi = pi-(pi/2)-(IBA/2) # Inclination angle in radians (phi) SI units: rad
L = 0.118 # Avg length of trachea (https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5900092/#:~:text=On%20average%2C%20the%20length%20of,to%20be%20shorter%20in%20females.&text=Structure%20of%20the%20trachea.)

## Dataset 1
# Source: "The Mechanism of Breath Formation"
dataset1 = [5.4444E-07, 5.8472E-07, 6.2656E-07, 6.6727E-07, 7.2656E-07, 7.6727E-07, 8.3492E-07, 9.0000E-07, 9.6627E-07, 1.0557E-06, 1.1621E-06, 1.2616E-06, 1.3686E-06, 1.4746E-06, 1.5741E-06, 1.6736E-06, 1.7828E-06, 1.8863E-06, 1.9858E-06, 2.1567E-06, 2.3432E-06, 2.5136E-06, 2.6939E-06, 2.8643E-06, 3.0672E-06, 3.3113E-06, 3.5555E-06, 3.8230E-06]

post_trachea = []

for d_d = dataset1
    penetration = trachea.straightTube(T, ro, Q, mu, ro_d, d_d, d, phi, L)
    push!(post_trachea, penetration)
end

print(post_trachea)

# bifurcation(ro_d, d_d, d_0, d_1, mu, Q, alpha)
ro_d = 1000 # Droplet density of water
d_0 = 0.018 # Diameter of trachea (https://www.google.com/url?sa=i&url=https%3A%2F%2Fwww.slideshare.net%2Fjinojustinj%2Ftracheal-pathologies&psig=AOvVaw38sp20tGz3PGVizR_R8F_N&ust=1605869019853000&source=images&cd=vfe&ved=0CAIQjRxqFwoTCICtjbi2ju0CFQAAAAAdAAAAABAT)
d_1 = 0.0122 # Diameter of daughter (https://www.google.com/url?sa=i&url=https%3A%2F%2Fwww.slideshare.net%2Fjinojustinj%2Ftracheal-pathologies&psig=AOvVaw38sp20tGz3PGVizR_R8F_N&ust=1605869019853000&source=images&cd=vfe&ved=0CAIQjRxqFwoTCICtjbi2ju0CFQAAAAAdAAAAABAT)
mu = 1.895e-5 # Dynamic viscosity of air (https://www.engineersedge.com/physics/viscosity_of_air_dynamic_and_kinematic_14483.htm#:~:text=The%20viscosity%20of%20air%20depends,10%2D5%20Pa%C2%B7s%20.&text=At%2025%20%C2%B0C%2C%20the,the%20kinematic%20viscosity%2015.7%20cSt.)
Q = 0.001 # Based on 60L/min flowrate of ventilators
alpha = 1.06116 # SCA, all ages. Haskin Goodman et. al. "Normal Tracheal Bifurcation Angle"

post_bifurc0 = []

for d_d = dataset1
    bifurc_penetration = bifurcations.bifurcation(ro_d, d_d, d_0, d_1, mu, Q, alpha)
    push!(post_bifurc0, bifurc_penetration)
end

# straightTube(T, ro, Q, mu, ro_d, d_d, d, phi, L)
T = 309.5 # Temperature of Air
ro = 1.12027 # Density of Air (http://www.anaesthesia.med.usyd.edu.au/resources/lectures/humidity_clt/humidity.html#:~:text=While%20nose%20breathing%20at%20rest,%25%20to%2070%25%20relative%20humidity.)
Q = 0.001 # Based on 60L/min flowrate of ventilators
mu = 1.895e-5 # Dynamic viscosity of air (https://www.engineersedge.com/physics/viscosity_of_air_dynamic_and_kinematic_14483.htm#:~:text=The%20viscosity%20of%20air%20depends,10%2D5%20Pa%C2%B7s%20.&text=At%2025%20%C2%B0C%2C%20the,the%20kinematic%20viscosity%2015.7%20cSt.)
ro_d = 1000 # Droplet density of water
d = 0.0189 # Diameter of trachea (https://www.google.com/url?sa=i&url=https%3A%2F%2Fwww.slideshare.net%2Fjinojustinj%2Ftracheal-pathologies&psig=AOvVaw38sp20tGz3PGVizR_R8F_N&ust=1605869019853000&source=images&cd=vfe&ved=0CAIQjRxqFwoTCICtjbi2ju0CFQAAAAAdAAAAABAT)
# phi = (pi/2) # Inclination angle in radians (phi) SI units: rad
IBA = 1.040216 # Interbronchial angle, average for all ages, 59.6 degrees SI Units: radians Haskin Goodman et. al. "Normal Tracheal Bifurcation Angle"
phi = pi-(pi/2)-(IBA/2) # Inclination angle in radians (phi) SI units: rad
L = 0.118 # Avg length of trachea (https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5900092/#:~:text=On%20average%2C%20the%20length%20of,to%20be%20shorter%20in%20females.&text=Structure%20of%20the%20trachea.)

post_inclinedTube0 = []

for d_d = dataset1
    penetration = trachea.straightTube(T, ro, Q, mu, ro_d, d_d, d, phi, L)
    push!(post_inclinedTube0, penetration)
end

# print(post_trachea)
# print("\n")
# print(post_inclinedTube0)
# print("\n")
# print(size(post_trachea))
# print("\n")
# print("\n")
# print(post_bifurc0)
# print("\n")
# print(size(post_bifurc0))

# Deposition Efficiency after first bifurcation
n_i = post_trachea .* post_bifurc0

# print(n_i)

# Distribution for dataset1
d_i1 = [0.056718217, 0.060143555, 0.068012484, 0.071244953, 0.073582454, 0.076460071, 0.078809187, 0.077779268, 0.072560295, 0.065129111, 0.056718217, 0.045163553, 0.037531843, 0.031410246, 0.026885611, 0.02266377, 0.018430315, 0.01499732, 0.01121713, 0.00960475, 0.007514077, 0.005690533, 0.003839032, 0.002862148, 0.002179373, 0.001434928, 0.000799432, 0.000618126]
# print(size(d_i1))

# Penetration * Distribution for data set 1 (one bifurcation)
final_distribution = n_i .* d_i1

print(final_distribution)

# L = 0.1
# print(L)

# Write your package code here.

end


# uncomment if first time using LightXML
# Pkg.add("LightXML");
using LightXML;

xWidth = 19900.0
yWidth = 13500.0
latRoot = 46.9877911
lonRoot = 7.131336
pathString = "M"

function mappedLongitude(lon)

    if (lon < 0)
        (((180 + (180 - abs(lon))) / 360) * xWidth)
    else
        ((lon / 360) * xWidth)
    end
end

function mappedLatitude(lat)
    with_rounding(Float64,RoundDown) do
        ((180 - (lat + 90)) / 180) * yWidth
    end
end

gpxDoc = parse_file("/Users/rene/Development/gps-test/github/gpx/1k.gpx")

xroot = root(gpxDoc)
trk = get_elements_by_tagname(xroot, "trk")[1]
trkseg = get_elements_by_tagname(trk, "trkseg")[1]
trkpt = get_elements_by_tagname(trkseg, "trkpt")
nodeCounter = 1
lastMappedLongitude = 1.0
lastMappedLatitude = 1.0
for c in child_nodes(trkseg)
    if is_elementnode(c)
        #thisNode = XMLElement(c)
        ad = attributes_dict(XMLElement(c))
        if (nodeCounter == 1)
            lastMappedLongitude = mappedLongitude(parse(Float32,(ad["lon"])))
            lastMappedLatitude = mappedLatitude(parse(Float32,(ad["lat"])))
            mappedLonRoot = mappedLongitude(lonRoot)
            mappedLatRoot = mappedLatitude(latRoot)
            pathString = string(pathString,round(lastMappedLongitude - mappedLonRoot,5) , "," , round(lastMappedLatitude - mappedLatRoot,5))
        else
            currentMappedLongitude = mappedLongitude(parse(Float32,(ad["lon"])))
            currentMappedLatitude = mappedLatitude(parse(Float32,(ad["lat"])))
            pathString = string(pathString,"l",round((currentMappedLongitude - lastMappedLongitude),5) , "," , round((currentMappedLatitude - lastMappedLatitude),5))
            lastMappedLongitude = currentMappedLongitude
            lastMappedLatitude = currentMappedLatitude
        end
        nodeCounter = nodeCounter + 1
    end
end

println(pathString)


using GeometricIntegrators
using ChargedParticleDynamics.GuidingCenter4d.SymmetricSurface
using GeometricExamples

const Δt     = 2.
const ntime  = 5000
const nx     = 100
const ny     = 100
const nsave  = 10
const nplot  = 4
const run_id = "poincare_2nd_symmetric_vprk_dt2"

include("../guiding_center_4d_settings_vprk.jl")

tableau_list = get_tableau_list_vprk_projection()

pinv = guiding_center_4d_iode_poincare_invariant_2nd(Δt, nx, ny, ntime, nsave)

include("guiding_center_4d_symmetric_poincare_invariant_2nd.jl")


@memoize function ξ(Pk,ℓ,n; q="best") # Eq 18 in SVM + Eq 3 in 2Fast.
    N = 16*1024
    kmino = exp(-25)
    kmaxo = exp(25)
    r0o = 1/kmaxo
    rmini = exp(-11)#exp(-14.555)
    rmaxi = exp(11)#exp(11)

    if q == "auto"
        if ℓ == 2 && n == -2
            q = 0
        else
            q = max(.5, 1+n)
        end
    end

    if q == "best"
        q = best_q(Pk,ℓ,n, kmino, kmaxo)
        if ℓ == 0
            q += .4
        end
        #if ℓ == 1
        #    if n == 0
        #        q -= 1
        #    end
        #end
    end

    if q == "show"
        q = show_best_q(Pk,ℓ,n, kmino, kmaxo)
        if ℓ == 0
            println("Xi l0 correction")
            q += .4
        end
        #if ℓ == 1
        #    if n == 0
        #        println("Xi l1n0 correction")
        #        q -= 1
        #    end
        #end
    end

    if q == "henry"
        q = henry_q(Pk,ℓ,n, kmino, kmaxo, r=kmaxo)#r=rmaxi)
    end

    r, xi = xicalc(Pk, ℓ, -n, N=N, kmin=kmino, kmax=kmaxo, r0=r0o; q=q)
    return r, xi .* (r .^ (-n))
end


using KitBase, KitBase.FastGaussQuadrature
using ProgressMeter, Plots

begin
    # case
    matter = "photon"
    case = "linesource"
    space = "1d1f1v"
    nSpecies = 1
    flux = "kfvs"
    collision = "bgk"
    interpOrder = 2
    limiter = "vanleer"
    boundary = "fix"
    cfl = 0.5
    maxTime = 0.2

    # physical space
    x0 = 0
    x1 = 1
    nx = 100
    pMeshType = "uniform"
    nxg = 0

    # velocity space
    vMeshType = "rectangle"
    umin = -1
    umax = 1
    nu = 28
    nug = 0

    # gas
    knudsen = 0.01
    mach = 0.0
    prandtl = 1
    inK = 0
    omega = 0.5
    alphaRef = 1.0
    omegaRef = 0.5
end

set = Setup(
    matter,
    case,
    space,
    flux,
    collision,
    nSpecies,
    interpOrder,
    limiter,
    boundary,
    cfl,
    maxTime,
)

pSpace = PSpace1D(x0, x1, nx, nxg)

points, weights = gausslegendre(nu)
vSpace = VSpace1D(points[1], points[end], nu, points, ones(nu).*(points[end]-points[1])/(nu-1), weights)

μᵣ = ref_vhs_vis(knudsen, alphaRef, omegaRef)
gas = Gas(knudsen, mach, prandtl, inK, 3.0, omega, alphaRef, omegaRef, μᵣ)

fL = 0.5 * ones(nu)
wL = moments_conserve(fL, points, weights)
primL = conserve_prim(wL, 3.0)
fR = 0.0001 * ones(nu)
wR = moments_conserve(fR, points, weights)
primR = conserve_prim(wR, 3.0)
ib = IB1F(wL, primL, fL, primL, wR, primR, fR, primR)

ks = SolverSet(set, pSpace, vSpace, gas, ib, @__DIR__)

ctr = Array{ControlVolume1D1F}(undef, ks.pSpace.nx)
face = Array{Interface1D1F}(undef, ks.pSpace.nx + 1)
for i in eachindex(ctr)
    ctr[i] = ControlVolume1D1F(
        ks.pSpace.x[i],
        ks.pSpace.dx[i],
        ks.ib.wR,
        ks.ib.primR,
        ks.ib.fR,
    )
end

for i = 1:ks.pSpace.nx+1
    face[i] = Interface1D1F(ks.ib.wL, ks.ib.fL)
end

function fboundary(
    ff::T2,
    f::T4,
    u::T5,
    dt,
    rot = 1,
) where {
    T2<:AbstractArray{<:AbstractFloat,1},
    T4<:AbstractArray{<:AbstractFloat,1},
    T5<:AbstractArray{<:AbstractFloat,1},
}

    δ = heaviside.(u .* rot)

    fWall = 0.5 .* δ .+ f .* (1.0 .- δ)

    @. ff = u * fWall * dt

    return nothing

end

function step(
    ffL::T2,
    w::T3,
    f::T4,
    ffR::T2,
    u::T5,
    weights::T5,
    τ,
    dx,
    dt,
) where {
    T2<:AbstractArray{<:AbstractFloat,1},
    T3<:AbstractArray{<:AbstractFloat,1},
    T4<:AbstractArray{<:AbstractFloat,1},
    T5<:AbstractArray{<:AbstractFloat,1},
}
    M = sum(weights .* f)
    for i in eachindex(u)
        f[i] += (ffL[i] - ffR[i]) / dx + (M - f[i]) / τ * dt
    end
    w[1] = sum(weights .* f)
end

dt = cfl * pSpace.dx[1]
nt = maxTime / dt |> floor |> Int

anim = @animate for iter = 1:150#nt
    #reconstruct!(ks, ctr)

    fboundary(
        face[1].ff,
        ctr[1].f,
        vSpace.u,
        dt,
    )
    @inbounds for i = 2:nx
        flux_kfvs!(
            face[i].ff,
            ctr[i-1].f,
            ctr[i].f,
            vSpace.u,
            dt,
            ctr[i-1].sf,
            ctr[i].sf,
        )
    end

    @inbounds for i = 1:nx-1
        step(
            face[i].ff,
            ctr[i].w,
            ctr[i].f,
            face[i+1].ff,
            vSpace.u,
            vSpace.weights,
            1.0,
            pSpace.dx[i],
            dt,
        )
    end

    pltx = ks.pSpace.x[1:ks.pSpace.nx]
    plty = zeros(ks.pSpace.nx, 6)
    for i in eachindex(pltx)
        plty[i, 1] = ctr[i].w[1]
    end
    plot(pltx, plty[:, 1], label = "Density", lw = 2, xlabel = "x")
end

gif(anim, "propagation.gif")

begin
    pltx = ks.pSpace.x[1:ks.pSpace.nx]
    plty = zeros(ks.pSpace.nx, 6)
    for i in eachindex(pltx)
        plty[i, 1] = ctr[i].w[1]
    end
    plot(pltx, plty[:, 1], label = "Density", lw = 2, xlabel = "x")
end


using PooksoftOptionsKit
using PyPlot

# call -
result_buyer = compute_call_option_profit_loss_at_expiration(:buy, 20.0, 1.20, 0.0, 40.0)
result_seller = compute_call_option_profit_loss_at_expiration(:sell, 20.0, 1.20, 0.0, 40.0)

profit_loss_buyer = result_buyer.value
profit_loss_seller = result_seller.value

# plot -
plot(profit_loss_buyer[:,1], profit_loss_buyer[:,3], "b")       # buyer - payoff
#plot(profit_loss_buyer[:,1], profit_loss_buyer[:,2], "b--")     # buyer - P/L
#plot(profit_loss_seller[:,1], profit_loss_seller[:,2], "r")       # buyer - payoff
#plot(profit_loss_seller[:,1], profit_loss_seller[:,2], "r--")     # buyer - P/L


function flux_Gxp_kernel(nx, ny, idx, shared, op::Function, thread_idx, block_dim)

    u1 = shared[thread_idx + block_dim * 4]
    u2 = shared[thread_idx + block_dim * 5]
    rho = shared[thread_idx + block_dim * 6]
    pr = shared[thread_idx + block_dim * 7]

    ut = u1*ny - u2*nx
    un = u1*nx + u2*ny

    beta = 0.5*rho/pr
    S1 = ut*CUDAnative.sqrt(beta)
    B1 = 0.5*CUDAnative.exp(-S1*S1)/CUDAnative.sqrt(pi*beta)
    A1pos = 0.5*(1 + CUDAnative.erf(S1))

    pr_by_rho = pr/rho
    u_sqr = ut*ut + un*un

    shared[thread_idx] = op((rho*(ut*A1pos + B1)), shared[thread_idx])
    temp1 = pr_by_rho + ut*ut
    temp2 = temp1*A1pos + ut*B1
    shared[thread_idx + block_dim] = op((rho*temp2), shared[thread_idx + block_dim])
    temp1 = ut*un*A1pos + un*B1
    shared[thread_idx + block_dim * 2] = op((rho*temp1), shared[thread_idx + block_dim * 2])

    temp1 = (7*pr_by_rho) + u_sqr
    temp2 = 0.5*ut*temp1*A1pos
    temp1 = (6*pr_by_rho) + u_sqr
    shared[thread_idx + block_dim * 3] = op((rho*(temp2 + 0.5*temp1*B1)), shared[thread_idx + block_dim * 3])

    return nothing
end

function flux_Gxn_kernel(nx, ny, idx, shared, op::Function, thread_idx, block_dim)

    u1 = shared[thread_idx + block_dim * 4]
    u2 = shared[thread_idx + block_dim * 5]
    rho = shared[thread_idx + block_dim * 6]
    pr = shared[thread_idx + block_dim * 7]

    ut = u1*ny - u2*nx
    un = u1*nx + u2*ny

    beta = 0.5*rho/pr
    S1 = ut*CUDAnative.sqrt(beta)
    B1 = 0.5*CUDAnative.exp(-S1*S1)/CUDAnative.sqrt(pi*beta)
    A1neg = 0.5*(1 - CUDAnative.erf(S1))

    pr_by_rho = pr/rho
    u_sqr = ut*ut + un*un
    shared[thread_idx] = op((rho*(ut*A1neg - B1)), shared[thread_idx])

    temp1 = pr_by_rho + ut*ut
    temp2 = temp1*A1neg - ut*B1
    shared[thread_idx + block_dim] = op((rho*temp2), shared[thread_idx + block_dim])
    temp1 = ut*un*A1neg - un*B1
    shared[thread_idx + block_dim * 2] = op((rho*temp1), shared[thread_idx + block_dim * 2])
    temp1 = (7*pr_by_rho) + u_sqr
    temp2 = 0.5*ut*temp1*A1neg
    temp1 = (6*pr_by_rho) + u_sqr
    shared[thread_idx + block_dim * 3] = op((rho*(temp2 - 0.5*temp1*B1)), shared[thread_idx + block_dim * 3])
    return nothing
end

function flux_Gyp_kernel(nx, ny, idx, shared, op::Function, thread_idx, block_dim)
    u1 = shared[thread_idx + block_dim * 4]
    u2 = shared[thread_idx + block_dim * 5]
    rho = shared[thread_idx + block_dim * 6]
    pr = shared[thread_idx + block_dim * 7]

    ut = u1*ny - u2*nx
    un = u1*nx + u2*ny

    beta = 0.5*rho/pr
    S2 = un*CUDAnative.sqrt(beta)
    B2 = 0.5*CUDAnative.exp(-S2*S2)/CUDAnative.sqrt(pi*beta)
    A2pos = 0.5*(1 + CUDAnative.erf(S2))

    pr_by_rho = pr/rho
    u_sqr = ut*ut + un*un

    shared[thread_idx] = op((rho*(un*A2pos + B2)), shared[thread_idx])
    temp1 = pr_by_rho + un*un
    temp2 = temp1*A2pos + un*B2
    temp1 = ut*un*A2pos + ut*B2
    #TODO - Verify this is correct
    shared[thread_idx + block_dim] = op((rho*temp1), shared[thread_idx + block_dim])
    shared[thread_idx + block_dim * 2] = op((rho*temp2), shared[thread_idx + block_dim * 2])

    temp1 = (7*pr_by_rho) + u_sqr
    temp2 = 0.5*un*temp1*A2pos
    temp1 = (6*pr_by_rho) + u_sqr
    shared[thread_idx + block_dim * 3] = op((rho*(temp2 + 0.5*temp1*B2)), shared[thread_idx + block_dim * 3])
    return nothing
end

function flux_Gyn_kernel(nx, ny, idx, shared, op::Function, thread_idx, block_dim)
    u1 = shared[thread_idx + block_dim * 4]
    u2 = shared[thread_idx + block_dim * 5]
    rho = shared[thread_idx + block_dim * 6]
    pr = shared[thread_idx + block_dim * 7]

    ut = u1*ny - u2*nx
    un = u1*nx + u2*ny

    beta = 0.5*rho/pr
    S2 = un*CUDAnative.sqrt(beta)
    B2 = 0.5*CUDAnative.exp(-S2*S2)/CUDAnative.sqrt(pi*beta)
    A2neg = 0.5*(1 - CUDAnative.erf(S2))

    pr_by_rho = pr/rho
    u_sqr = ut*ut + un*un

    shared[thread_idx] = op((rho*(un*A2neg - B2)), shared[thread_idx])
    temp1 = pr_by_rho + un*un
    temp2 = temp1*A2neg - un*B2
    temp1 = ut*un*A2neg - ut*B2
    #TODO - Verify this is correct
    shared[thread_idx + block_dim] = op((rho*temp1), shared[thread_idx + block_dim])
    shared[thread_idx + block_dim * 2] = op((rho*temp2), shared[thread_idx + block_dim * 2])

    temp1 = (7*pr_by_rho) + u_sqr
    temp2 = 0.5*un*temp1*A2neg
    temp1 = (6*pr_by_rho) + u_sqr
    shared[thread_idx + block_dim * 3] = op((rho*(temp2 - 0.5*temp1*B2)), shared[thread_idx + block_dim * 3])

    return nothing
end


__precompile__()

module NearestNeighbors

using Distances
import Distances: Metric, result_type, eval_reduce, eval_end, eval_op, eval_start, evaluate

using StaticArrays

import Base.show
import Compat.view

export BruteTree, KDTree, BallTree, DataFreeTree
export knn, inrange # TODOs? , allpairs, distmat, npairs
export injectdata

export Euclidean,
       Cityblock,
       Minkowski,
       Chebyshev,
       Hamming

# Change this to enable debugging
const DEBUG = false

abstract NNTree{V <: AbstractVector, P <: Metric}

typealias MinkowskiMetric Union{Euclidean, Chebyshev, Cityblock, Minkowski}

function check_input{V1, V2 <: AbstractVector}(::NNTree{V1}, ::Vector{V2})
    if length(V1) != length(V2)
        throw(ArgumentError(
            "dimension of input points:$(length(V2)) and tree data:$(length(V1)) must agree"))
    end
end

function check_input{V1, V2 <: Number}(::NNTree{V1}, point::Vector{V2})
    if length(V1) != length(point)
        throw(ArgumentError(
            "dimension of input points:$(length(point)) and tree data:$(length(V1)) must agree"))
    end
end

get_T{T <: AbstractFloat}(::Type{T}) = T
get_T{T}(::T) = Float64

include("debugging.jl")
include("evaluation.jl")
include("tree_data.jl")
include("datafreetree.jl")
include("knn.jl")
include("inrange.jl")
include("hyperspheres.jl")
include("hyperrectangles.jl")
include("utilities.jl")
include("brute_tree.jl")
include("kd_tree.jl")
include("ball_tree.jl")
include("tree_ops.jl")

end # module


###################################################
# Quick guide for the Geodesy package
###################################################

#
# Step 1 - Grab some points
#

# Starting with points in a coordinate system defined by an SRID (https://en.wikipedia.org/wiki/SRID)
raw_data = [
        573105.43200000003    6086900.3839999996    277.42700000000002    "EPSG"    32755;
        572734.61300000001    6087631.6480000000    273.80000000000001    "EPSG"    32755;
        572688.82100000000    6086535.8650000002    267.334               "EPSG"    32755;
        573667.63899999997    6086370.3410000000    287.68799999999999    "EPSG"    32755
]

# build an srid type from the info in the raw data
auth = symbol(raw_data[1, end-1])
code = raw_data[1, end]
srid = SRID{auth, code}
typealias Utm55S CRS{srid}  # use a type alias if it suits

# create a vector of the points in this SRID
# Transformations for the CRS type are handled by Proj4 package so anything supported by Proj4 is fine, BUT
# its up to the user to correctly interpret the x,y,z fields of the CRS type as they could mean anything (lon lat ... / east north ... / x y ... etc)

Xin = convert(Vector{Utm55S}, raw_data[:, 1:3]; row=true) # = convert(Vector{CRS{srid}}, raw_data[:, 1:3]; row=true)

#
# Step 2 - Transform the points into an longitude / latitude / altitude format with the WGS84 ellipsoid (i.e. GPS coordinates)
#
# N.B.  The below step requires Geodesy to recognise the Proj4 projection for type LLA{WGS84}.  If you're using a point type
# that geodesy doesn't know the Proj4 projection for, overload Geodesy.get_projection() to return the appropriate Proj4 projection e.g.:
#    Geodesy.get_projection(::Type{LLA{WGS84}}) = Proj4.Projection("+proj=longlat +datum=WGS84 +no_defs")
#

Xwgs84 = geotransform(LLA{WGS84}, Xin)

#
# Step 3 - transform them into a local (East, North, Up) coordindate frame
#

# use the center of the points as the center of the local frame
bbox = Bounds(Xwgs84)
lla_ref = center(bbox)

# option a: embed the local frame's origin in the point type.  This saves passing around lla_ref
Xlocal_a = geotransform(ENU{lla_ref}, Xwgs84)
println(typeof(Xlocal_a))

# option b: don't embed the local frame's origin in the point type.  Use this approach if you use loads of reference points
Xlocal_b = geotransform(ENU, Xwgs84, lla_ref)
println(typeof(Xlocal_b))

#
# Step 4 - do something with the points.  For this guide we'll just add some new points in between the old ones
#

for pair in combinations(1:length(Xlocal_a), 2)
    push!(Xlocal_a, (Xlocal_a[pair[1]] + Xlocal_a[pair[2]]) / 2.0)
    push!(Xlocal_b, (Xlocal_b[pair[1]] + Xlocal_b[pair[2]]) / 2.0)
end

#
# Step 5 - convert back to a world coordinate frame (ECEF this time to be different)
#

Xecef_a = geotransform(ECEF, Xlocal_a)           # the reference point is in the type of Xlocal_a
Xecef_b = geotransform(ECEF, Xlocal_b, lla_ref)  # need to supply the reference point

#
# Step 6 - Lastly go back to the original format
#

Xout = geotransform(CRS{srid},  Xecef_a)

# N.B. the input srid was actually a utm projection (zone 55) for the WGS84 datum.  We can use this package to select the SRID that best matches another point
(zone, band) = Geodesy.utm_zone(lla_ref) # I think this is the same for all reference ellipses...
srid_out = Geodesy.utm_srid(lla_ref)     # Geodesy can only do this for the WGS84 datum at present

###############################################
# ECEF Type Warning
###############################################

# whether ECEF is truly an ECEF point depends on the datum
# e.g.:
lla_osgb36 = LLA{Geodesy.OSGB36}(0, 0, 0)          # OSGB36 is a good match to the Earth in the UK but not elsewhere
ecef_fake = geotransform(ECEF, lla_osgb36)         # = 6.377563396e6, 0.0, 0.0
# isn't a true ECEF point because the OSGB36 ellipsoid isn't geocentric.

# If in doubt, datum conversions can be using the CRS type to get Proj4 to do it, e.g.
ecef_srid = geotransform(SRID(ECEF{WGS84}), lla_osgb36)     # = 6.377879171552554e6,-99.12039106890559, 534.423089412207
ecef_wgs84 = convert(ECEF{WGS84}, ecef_srid)                # convert to a type native to this package (type conversion not a transformation)

# or equivilently
srid = SRID(lla_odgb36)
lla_srid = convert(CRS{srid}, lla_osgb36)                   # type conversion not a transformation
ecef_wgs84_2 = geotransform(ECEF{WGS84}, lla_srid)          # = 6.377879171552554e6,-99.12039106890559, 534.423089412207


# Changes from original file:
#   pyimport tinyarray
#   range(A) -> range(0,length=A)
#   A.append(...) -> append!(A,...)

tinyarray = pyimport("tinyarray")
    sigma_0 = tinyarray.array([[1, 0], [0, 1]])
    sigma_x = tinyarray.array([[0, 1], [1, 0]])
    sigma_y = tinyarray.array([[0, -1im], [1im, 0]])
    sigma_z = tinyarray.array([[1, 0], [0, -1]])

function make_system(t=1.0, alpha=0.5, e_z=0.08, W=10, L=30)
    # Start with an empty tight-binding system and a single square lattice.
    # `a` is the lattice constant (by default set to 1 for simplicity).
    lat = kwant.lattice.square()

    syst = kwant.Builder()

    #### Define the scattering region. ####
    syst[[lat(x, y) for x in range(0,length=L) for y in range(0,length=W)]] =
        4 * t * sigma_0 + e_z * sigma_z # note [...] for comprehension and 1:L/W for range(L/W) also no \ for line break
    # hoppings in x-direction
    syst[kwant.builder.HoppingKind((1, 0), lat, lat)] =
        -t * sigma_0 + 1im * alpha * sigma_y / 2 ## note 1im not 1j
    # hoppings in y-directions
    syst[kwant.builder.HoppingKind((0, 1), lat, lat)] =
        -t * sigma_0 - 1im * alpha * sigma_x / 2

    #### Define the left lead. ####
    lead = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))

    lead[[lat(0, j) for j in range(0,length=W)]] = 4 * t * sigma_0 + e_z * sigma_z
    # hoppings in x-direction
    lead[kwant.builder.HoppingKind((1, 0), lat, lat)] =
        -t * sigma_0 + 1im * alpha * sigma_y / 2
    # hoppings in y-directions
    lead[kwant.builder.HoppingKind((0, 1), lat, lat)] =
        -t * sigma_0 - 1im * alpha * sigma_x / 2

    #### Attach the leads and return the finalized system. ####
    syst.attach_lead(lead)
    syst.attach_lead(lead.reversed())

    return syst
end

function plot_conductance(syst, energies)
    # Compute conductance
    data = []
    for energy in energies
        smatrix = kwant.smatrix(syst, energy)
        append!(data,smatrix.transmission(1, 0))
    end

    pyplot.figure()
    pyplot.plot(energies, data)
    pyplot.xlabel("energy [t]")
    pyplot.ylabel("conductance [e^2/h]")
    pyplot.gcf() ## note gcf not show
end


using HierarchicalMatrices
using Test, LinearAlgebra, Random, InteractiveUtils

for r in map(BLOCKRANK, subtypes(AbstractFloat))
    @test iseven(r)
end

Random.seed!(0)

@hierarchical UpperTriangularHierarchicalMatrix Matrix (UpperTriangular, LowRankMatrix) Diagonal

for T in (Float32, Float64)
    A = rand(T, 10, 5)
    x = rand(T, 40)

    y = zero(x)
    HierarchicalMatrices.mul!(y, A, x, 1, 1)
    @test norm(y[1:10] - A*x[1:5]) ≤ eps(T)*norm(A*x[1:5])

    fill!(y, 0.0)
    HierarchicalMatrices.mul!(y, A, x, 5, 5, 2, 2)
    @test norm(y[5:2:23] - A*x[5:2:14]) ≤ eps(T)*norm(A*x[5:2:14])

    fill!(y, 0.0)
    HierarchicalMatrices.mul!(y, transpose(A), x, 1, 5, 2, 1)
    @test norm(y[1:5] - A'x[5:2:23]) ≤ eps(T)*norm(A'x[5:2:23])

    fill!(y, 0.0)
    HierarchicalMatrices.mul!(y, transpose(A), x, 6, 3, 1, 3)
    @test norm(y[6:3:18] - A'*x[3:12]) ≤ eps(T)*norm(A'*x[3:12])

    A = map(big, A)
    x = map(big, x)

    y = zero(x)
    HierarchicalMatrices.mul!(y, A, x, 1, 1)
    @test y[1:10] == A*x[1:5]

    fill!(y, 0.0)
    HierarchicalMatrices.mul!(y, A, x, 5, 5, 2, 2)
    @test y[5:2:23] == A*x[5:2:14]

    fill!(y, 0.0)
    HierarchicalMatrices.mul!(y, transpose(A), x, 1, 5, 2, 1)
    @test y[1:5] == A'x[5:2:23]

    fill!(y, 0.0)
    HierarchicalMatrices.mul!(y, transpose(A), x, 6, 3, 1, 3)
    @test y[6:3:18] == A'*x[3:12]

    UTHM = UpperTriangularHierarchicalMatrix(T, 2, 2)

    @test typeof(UTHM.UpperTriangularblocks) == Matrix{UpperTriangular{T,LowRankMatrix{T}}}
end

include("../examples/Kernel.jl")