### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 85e3f1ee-9158-11eb-2c00-672502866677
begin
using Pkg
	Pkg.add("JuMP")
	Pkg.add("Juniper")
	Pkg.add("Ipopt") # Juniper + Ipopt = solver for Quadratic optimization problems like this
using JuMP, Ipopt, Juniper
end

# ╔═╡ 2cfafa80-9157-11eb-35c2-51be3c25e814
html"""<style>
main {
    max-width: 1500px;
}"""

# ╔═╡ feb87f00-914f-11eb-2784-8f05d5919eba
md"""# Optimal Weights and Activation Values"""

# ╔═╡ 22f58180-9153-11eb-00c6-6586404966cc
md"""__Data:__  In Sanskrit we observe an alternation between several allomorphic roots for one and the same verb. Let's take _budh_ 'awake, know, $\ldots$' as a concrete example

+ bodh-ati (3sg. pres.)
+ bhot-sya-ti (3sg. fut.)
+ bhut (Noun, Nom.)
+ bhud-bhis (Noun, Instr.)
+ bud-dha (Past participle)

GSC offers a possible solution to this puzzle if we assume that the input for morphological derivations is a blend of different allomorphs with some partial activation values (see GSCNet).

But how do we find suitable values for these activations and how do we find the right constraint weights in a Gradient Harmonic Grammar?
"""

# ╔═╡ 13e8db00-9154-11eb-187a-0d4cd645dcaa
md"""The problem can be phrased as a mathematical optimization problem: We have a function:

$$H_{winner}(x) = c^Tx$$

expressing the Harmony (sum of violations) of the winner, that we want to maximize.
Our $x$s are the Harmonic constraints and the coefficients $c$ are the violations for the winning candidate.
This expression is subject to some constraints: for each other candidate it should hold that

$$H_{winner} > H_{candidate_i} $$ $$ \forall cand_i \in Tableaux$$

We also want the weights to be negative (these are penalties):

$$[x_i, x_2, x_3 ... x_n] \lt 0$$

and we know that the activation values are between 0 and 1: 0 meaning inactive and 1 meaning totally active (discrete).

$$1 \leq [\alpha .... \omega] \geq 0$$

"""

# ╔═╡ 65cb0890-9153-11eb-25b4-b19428afad2c
md""" ### Example: the past participle
__Inputs__:
root allomorphs:

+ (bhud)$^\alpha$
+ (budh)$^\beta$
+ (bhut)$^\gamma$

suffix allomorphs:

+ (-ta)$^\delta$
+ (-dha)$^\epsilon$

__Constraints__:
(all of them $\lt 0$)

+ Max : anti-deletion constraint
+ Dep : anti-epenthesis constraint
+ Voice : Avoid voice contrast ([+ voice][-voice] is bad)
+ License(lar) : Aspirated segments cannot occur before voiceless segments. Also avoid two aspirated next to each other.
+ Lazy : Reduce articulatory effort (avoid aspirated sounds)
+ *h]_W:Aspirated are not allowed in word final position
+ Id(lar) : Aspirated in the input == Aspirated in the output

"""

# ╔═╡ 1f066d00-9153-11eb-00ee-b760c15b6b67
md""" __Toy-example__

| $root_{(\alpha + \beta + \gamma)} + suff_{(\delta + \epsilon)}$ | Max   |   Dep | Voice | Lic | h]W | Lazy | Id(Lar) | Harmony |
|:---:|:---:|:---:|:-----:|:---:|:---:|:---:|:---:|:---:|
|bhud-ta | $\beta + \gamma + \epsilon$ | $(2-\alpha - \delta)$ | 1 | 0 | 0 | 0 | 0 | $M(\beta + \gamma + \epsilon) + D(2-\alpha - \delta) + Voice$ |
|:) bud-dha | $\alpha + \gamma + \delta$ | $(2-\beta - \epsilon)$ | 0 | 0 | 0 | 0 | 1 | $M(\alpha + \gamma + \epsilon) + D(2-\beta - \epsilon) + Id(lar)$ |
|budh-i-ta | $\alpha + \gamma + \epsilon$ | $(3-\beta - \delta)$ | 0 | 0 | 0 | 0 | 1 | $M(\alpha + \gamma + \epsilon) + D(3-\beta - \delta)$ |
| ... |... |... |... |... |... |... |... |... |

Notice that candidate 3 (*budh-i-ta*, built on the model of *kridh-i-ta*) has an extra epenthesis, hence the violation for dep (**3** - $\beta - \delta$)

Let's now solve this problem using the _JuMP_ library for mathematical optimization.
The expression we want to maximize is the Harmony of the winner (candidate 2):

$$max: M(\alpha + \gamma + \epsilon) + D(2-\beta - \epsilon) + Id(lar)$$

subject to:

$M(\alpha + \gamma + \epsilon) + D(2-\beta - \epsilon) + Id(lar) > M(\beta + \gamma + \epsilon) + D(2-\alpha - \delta) + Voice$
$M(\alpha + \gamma + \epsilon) + D(2-\beta - \epsilon) + Id(lar) > M(\alpha + \gamma + \epsilon) + D(3-\beta - \delta)$
$0 \lt [\alpha, \beta, \gamma, \delta, \epsilon] \lt 1$
$M, D, Id(lar), Voice, Lic, h]W < 0$

"""

# ╔═╡ 79976060-915a-11eb-325f-916ca3326950
# Initialize Model
begin
	optimizer = Juniper.Optimizer
	nl_solver = optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0)
	model = Model(optimizer_with_attributes(optimizer, "nl_solver" => nl_solver))
end

# ╔═╡ 75b962e0-915a-11eb-1b24-5d35a2b179e2
begin
	# Declare variables
	a = 1 # The difference between winner and losers should be greater than a
	b = -1 # Start value for the constraint weights
end

# ╔═╡ 4d9bf512-915b-11eb-2fee-e9a1ff31c9b7
begin
	unregister(model, :M)
	unregister(model, :D)
	unregister(model, :A)
	unregister(model, :L)
	unregister(model, :Li)
	unregister(model, :Coda)
	unregister(model, :Id)
	unregister(model, :α)
	unregister(model, :β)
	unregister(model, :γ)
	unregister(model, :δ)
	unregister(model, :ϵ)
end

# ╔═╡ 1bf6e310-9153-11eb-2fe5-5b4d277cff3a
begin
	@variable(model,  -10 <= Ma <= -1, start = b,integer = true) # MaxIO
	@variable(model,  -10 <= De <= -1, start = b,integer = true) # DepIO
	@variable(model,  -10 <= V <= -1, start = b,integer = true) # Assimilate == Voice
	@variable(model,  -10 <= Lic <= -1, start = b,integer = true) # License(lar)
	@variable(model, -10 <= hW <= -1, start = b,integer = true) # h]W
	@variable(model, -10 <= La <= -1, start = b,integer = true) #Lazy
	@variable(model, -10 <= IdL <= -3, start = 2 * b,integer = true) # Ident(lar)
	@variable(model, 0.1 <= α <= 0.9, start = 0.2) # bhud
	@variable(model, 0.1 <= β <= 0.9, start = 0.9) # budh
	@variable(model, 0.1 <= γ <= 0.9, start = 0.4) # bhut
	@variable(model, 0.1 <= δ <= 0.5, start = 0.5) # dha
	@variable(model, 0.5 <= ϵ <= 0.9, start = 0.7) # ta
end

# ╔═╡ 6d36f150-915f-11eb-252a-17639f320ed2
@objective(model, Max, Ma * (α + γ + δ) + De * (2 - β - ϵ) + IdL)

# ╔═╡ 900302e0-915b-11eb-3b97-855f0f06e4ba
# Objective function : Harmony of the winner
"""@objective(model, Max, Ma * (α + γ + δ) + De * (2 - β - ϵ) + IdL)"""

# ╔═╡ 8856e662-915b-11eb-1903-37e8e9ef1ed0
# Constraints
begin
	# Winner Harmony
	winner = @expression(model, Ma * (α + γ + δ) + De * (2 - β - ϵ) + IdL)

	bhudta = @expression(model, winner - (Ma * (β + γ + δ) + De * (2 - α - ϵ) + V))
	@constraint(model, bhudta >= a)

	budhta = @expression(model, winner - (Ma * (α + γ + δ) + De * (2 - β - ϵ) + V + Lic))
	@constraint(model, budhta >= a)

	budhita = @expression(model, winner - (Ma * (α + γ + δ) + De * (3 - β - ϵ)))
	@constraint(model, budhita >= a)
end

# ╔═╡ 6774d1e0-915c-11eb-1070-e98f93030697
"""$model"""

# ╔═╡ 7d023ed0-915c-11eb-25dc-e538e4c8f878
# Optimize / Solve
optimize!(model)

# ╔═╡ 7ae6d8e0-915c-11eb-39a3-bbfbfb998e5a
# Printing logic
# Objective value (Harmony of the winner after the optimization)
JuMP.objective_value(model)

# ╔═╡ 830bb260-915d-11eb-1afa-ffdeaa98ad5a
# Solver status
JuMP.primal_status(model)

# ╔═╡ 80785d50-915d-11eb-1441-2f013038b3e0
JuMP.termination_status(model)

# ╔═╡ 745246e0-915c-11eb-14f6-593066f40ef8
# MaxIO
value(Ma)

# ╔═╡ 6ecf24e0-915c-11eb-28cc-7d5e29ba29f1
#Dep
value(De)

# ╔═╡ b747204e-915d-11eb-1d2b-b77c4386f7f5
# Voice
value(V)

# ╔═╡ c1afebce-915d-11eb-3ba2-71fa0cc58aa1
# Coda
value(hW)

# ╔═╡ c20dea50-915d-11eb-2046-bb9095f0f080
value(α)

# ╔═╡ c22f06e0-915d-11eb-22da-3d03814961cd
value(β)

# ╔═╡ fcbe5450-915d-11eb-1080-c91bfcde29b5
value(γ)

# ╔═╡ 0187ea52-915e-11eb-2475-11a14ba0bc20
value(δ)

# ╔═╡ 062ce150-915e-11eb-1d7b-7558b1067064
value(ϵ)

# ╔═╡ ff36a752-915d-11eb-0835-33c93eaa8233
md"Of course this was only a toy-example. JuMP evaluated the expressions for just two candidates and some value were arbitrarly fixed but I think it gives the idea, how Mathematical Optimization (Quadratic Programming in this case) can be used to search suitable values for the constraints and the activation values"

# ╔═╡ Cell order:
# ╟─2cfafa80-9157-11eb-35c2-51be3c25e814
# ╟─feb87f00-914f-11eb-2784-8f05d5919eba
# ╟─22f58180-9153-11eb-00c6-6586404966cc
# ╟─13e8db00-9154-11eb-187a-0d4cd645dcaa
# ╟─65cb0890-9153-11eb-25b4-b19428afad2c
# ╟─1f066d00-9153-11eb-00ee-b760c15b6b67
# ╟─85e3f1ee-9158-11eb-2c00-672502866677
# ╠═79976060-915a-11eb-325f-916ca3326950
# ╠═75b962e0-915a-11eb-1b24-5d35a2b179e2
# ╟─4d9bf512-915b-11eb-2fee-e9a1ff31c9b7
# ╠═1bf6e310-9153-11eb-2fe5-5b4d277cff3a
# ╠═6d36f150-915f-11eb-252a-17639f320ed2
# ╠═900302e0-915b-11eb-3b97-855f0f06e4ba
# ╠═8856e662-915b-11eb-1903-37e8e9ef1ed0
# ╠═6774d1e0-915c-11eb-1070-e98f93030697
# ╠═7d023ed0-915c-11eb-25dc-e538e4c8f878
# ╠═7ae6d8e0-915c-11eb-39a3-bbfbfb998e5a
# ╠═830bb260-915d-11eb-1afa-ffdeaa98ad5a
# ╠═80785d50-915d-11eb-1441-2f013038b3e0
# ╠═745246e0-915c-11eb-14f6-593066f40ef8
# ╠═6ecf24e0-915c-11eb-28cc-7d5e29ba29f1
# ╠═b747204e-915d-11eb-1d2b-b77c4386f7f5
# ╠═c1afebce-915d-11eb-3ba2-71fa0cc58aa1
# ╠═c20dea50-915d-11eb-2046-bb9095f0f080
# ╠═c22f06e0-915d-11eb-22da-3d03814961cd
# ╠═fcbe5450-915d-11eb-1080-c91bfcde29b5
# ╠═0187ea52-915e-11eb-2475-11a14ba0bc20
# ╠═062ce150-915e-11eb-1d7b-7558b1067064
# ╟─ff36a752-915d-11eb-0835-33c93eaa8233


using DynamicPolynomials
using MultivariateMoments
using SemialgebraicSets
using ECOS

# -----------------------------------------------------------------------------
# Trivial tests
# -----------------------------------------------------------------------------

@testset "Guessing game --------- " begin
	@polyvar x y
	p  = (x - y)^2
	Sx = @set 1 - x^2 >= 0
	Sy = @set 1 - y^2 >= 0
	opt = lazy_relax(ECOS.Optimizer, SDDPBridge)

	expected_measure_x = AtomicMeasure(variables(Sx), [
		WeightedDiracMeasure([1], 0.5),
		WeightedDiracMeasure([-1], 0.5)])
	expected_measure_y = AtomicMeasure(variables(Sy), [
		WeightedDiracMeasure([0], 1)])
	expected_value = 1

	pvalue, pstrat_x, pstrat_y = solve_game(+p, Sx, Sy, opt, iteration=0x1)
	dvalue, dstrat_y, dstrat_x = solve_game(-p, Sy, Sx, opt, iteration=0x1)

	eps = 1e-3
	@test isapprox(expected_value, pvalue, atol=eps)
	@test isapprox(pvalue, -dvalue, atol=eps)
	@test isapprox(expected_measure_x, dstrat_x, atol=eps)
	@test isapprox(expected_measure_y, pstrat_y, atol=eps)
end


using JuMP, Cbc, DelimitedFiles

function distances_from_coords(node_coords::Array{Float32,2})
    """
    Recieves a 2D array of the X,Y coordinates for each of the n nodes.
    E.g.
    node_coords[node,1] : Latitude of the node

    Returns a n,n matrix of the euclidean distances between the nodes
    """
    node_number = size(node_coords,1)
    node_distances = zeros(node_number,node_number)
    for row in range(1,length=node_number)
        for column in range(1,length=node_number)
            #Applying Euclidean distance
            node_distances[row,column] = sqrt(
                Float32(node_coords[column,1]-node_coords[row,1])^2 +
                Float32(node_coords[column,2]-node_coords[row,2])^2 )
        end
    end

    return node_distances
end

function tsp_main(args)
    """
    Use Cbc to calculate Travelling Salesperson solution for a given graph.
    Args: Filename of the parsed graph file.
        The file should follow this standard
        http://www.math.uwaterloo.ca/tsp/world/countries.html#DJ

    """

    node_coords = DelimitedFiles.readdlm(args[1],' ',Float32,'\n')

    """
    nxn Matrix of floats representing the distances between the nodes
    """
    node_distances = distances_from_coords(node_coords)

    node_number = size(node_distances)[1]

    tsp_model = Model(with_optimizer(Cbc.Optimizer , threads=4))

    #Limits execution time
    set_time_limit_sec(tsp_model , 7200)

    #Decision variable - binary nxn Matrix representing the edges taken between nodes
    @variable(tsp_model, edges_taken[1:node_number, 1:node_number], Bin)
	@variable(tsp_model,  0 <= u[1:node_number]  <= (node_number-1),Int)

    #Min Sum ( node_distances[i][j] * edges_taken[i][j] for all nodes)
    @objective(tsp_model,Min, sum( map(*,node_distances,edges_taken) ) )

    fix(u[1], 0,force=true) # change the 1 for the position of initial node

	for i in 1:node_number
		@constraint(tsp_model,sum(edges_taken[i,j] for j in 1:node_number if i != j ) == 1)
	end
	# Restrições para garantir a chegada em cada vértice
	for j in 1:node_number
		@constraint(tsp_model,sum(edges_taken[i,j] for i in 1:node_number if i != j) == 1)
	end
	for i in 1:node_number
		for j in 2:node_number
			if i != j
				@constraint(tsp_model,u[i] - u[j] + node_number * edges_taken[i,j] <= node_number - 1)
			end
		end
	end
	for i in 2:node_number
		@constraint(tsp_model, 0 <= u[i] <= node_number - 1)
	end

    optimize!(tsp_model)

	#Código para facilitar a impressão das cidades percorridas
	auxU = value.(u)
	ordem = zeros(UInt64 , (node_number,1))
	for i in 1:node_number
		ordem[Int64(round(auxU[i]+1))] = i
	end

	print("O valor do caminho otimo é: ")
	println(objective_value(tsp_model))
	print("A ordem dos pontos percorridos: ")
	for i in 1:node_number
		print(ordem[i])
		print(" ")
	end
end

tsp_main(ARGS)


using Combinatorics
include("WrapDist.jl")

function findMinWrapDistUnique(f_true,f_hat)
# Assign each of the ground truth frequencies to closest
# (squared error) canditate of estimated frequencies.
# Metric is the wrap-aroud distance on the torus $\mathds{T}=[0,1]$.
# Note: each frequency in f_hat can only be assigned to the closest frequency
# in f_true and can only be used to explain multiple components
# (see findMinWrapDist).
# Note: Complexity of O(n!) where n=length(f_true)!
#
# Inputs:
# f_true: ground truth frequencies
# f_hat: estimated frequencies
#
# Output:
# dists: wrap-around distances of the frequencies in f_true and f_hat

f_true = f_true[:]
f_hat = f_hat[:]
K = length(f_true)

perms = collect(permutations(f_hat))
dists = Array{Float64,1}(undef,factorial(K))
for ii = 1:factorial(K)
    dists[ii] = norm(WrapDist.(f_true,perms[ii]))
end
(mindist, idx) = findmin(dists)

return WrapDist.(f_true,perms[idx])
end


function solve()
    n, s, k = [parse(Int, x) for x in split(readline())]
    s = n-s
    d = gcd(n, s, k)
    n ÷= d
    s ÷= d
    k ÷= d
    if gcd(k, n) ≠ 1
        -1
    else
        i = invmod(k, n)
        mod(i*s, n)
    end
end

function main()
    t = parse(Int, readline())
    for _=1:t
        println(solve())
    end
end

main()


@testset "likelihood.jl" begin
    println("             > likelihood.jl")
    @testset "consistency" begin
        using BinomialSynapses: likelihood_indices

        m_out = 1024
        m_in  = 1024
        Ns = CuArray(rand(1:128, m_out))
        ps = CUDA.rand(m_out)
        qs = CUDA.rand(m_out)
        σs = CUDA.rand(m_out)
        τs = CUDA.rand(m_out)
        model = BinomialModel(Ns, ps, qs, σs, τs);

        ns = CuArray(rand(1:128, m_out, m_in))
        ks = CUDA.zeros(Int, m_out, m_in);

        ks .= ns .% 64

        u1      = likelihood(ks, model, 0.3f0)
        u2, idx = likelihood_indices(ks, model, 0.3f0)

        @test u1 ≈ u2
        @test maximum(idx) <= m_in
        @test minimum(idx) >= 1

        idx = Array(idx)
        for r in eachrow(idx)
            @test issorted(r)
        end

        if RUN_BENCHMARKS
            println("")
            println("Benchmarking function likelihood!: should take about 300μs")
            display(@benchmark CUDA.@sync likelihood($ks, $model, 0.3f0))
            println("")
            println("")
            println("Benchmarking function likelihood_resample!: should take about 4ms")
            state = BinomialState(ns, ks)
            obs   = BinomialObservation(0.3f0, 0.1f0)
            display(@benchmark CUDA.@sync likelihood_resample!($state, $model, $obs))
            println("")
        end
    end
    @testset "correctness of values" begin
        observation = 3.0
        M_out = 6
        M_in = 4
        Ns = 10 .* CUDA.ones(Int, M_out)
        ps = CUDA.rand(M_out)
        qs = cu([1f0, 1f0, 3f0, 3f0, 1f0, 3f0])
        σs = Float32(0.1) .* CUDA.ones(M_out)
        τs = CUDA.rand(M_out)
        model = BinomialModel(Ns, ps, qs, σs, τs);

        ks = cu([3 3 3 3; 1 1 1 1; 1 1 1 1; 3 3 3 3; 3 1 1 1; 1 3 3 3])

        u, idx = likelihood_indices(ks, model, observation)
        u = Array(u)
        idx = Array(idx)

        # Combinations of q and k that correspond to the observation should have
        # a high likelihood
        @test u[1] > 3.9894
        @test u[3] > 3.9894

        # Combinations of q and k that do not match the observation should have
        # a low likelihood
        @test u[2] < 1.5f-10
        @test u[4] < 1.5f-10

        # Resampling should pick the particles that match q and the observation
        @test all(idx[5:6,:].==1)
    end
end


function get_boundary_filt(basin)

    mx = maximum(basin)
    mn = minimum(basin)
    # Kernel
    w = centered([1 1 1; 1 -8 1 ; 1 1 1])
    #w = centered([0 1 0; 1 -4 1 ; 0 1 0])

    # replicate for boundary conditions
    res = imfilter(basin,w,"replicate")

    # segmentation: take a low threshold in order to not loose some structure.
    # This detection can be improved.
    #res = abs.(res) .> (mx-mn)
    res = abs.(res) .> 0.
    return res
end

function get_list(basin,xg ,yg)

    num_att = length(unique(basin))

    # v_list empty container that will receive the different coordinates corresponding to each basins,
    # Each element of this v_set is an array of coordinates of different sizes.
    v_list = Vector{SVector{2,Float64}}()
    v_set = Vector{typeof(v_list)}()

    # original boundary
    bnd = get_boundary_filt(basin)
    I1=findall(bnd .== 1);
    for ind in I1; push!(v_list , [xg[ind[1]]; yg[ind[2]]]); end
    #@show v_list
    push!(v_set,v_list)

    # Merging !!
    basin_i_j = zeros(Int8,size(basin))
    for k in unique(basin)
        v_list = Vector{Vector{Float64}}()
        fill!(basin_i_j, zero(0)) # erase matrix
        I = findall(basin .== k)
        [basin_i_j[x] = 1 for x in I];
        # compute boundary
        bnd = get_boundary_filt(basin_i_j)
        I1=findall(bnd .== 1);
        # Get coordinates lists from matrix
        #push!(v_list , hcat([[xg[ind[1]]; yg[ind[2]]] for ind in I1 ]...))
        for ind in I1; push!(v_list , [xg[ind[1]]; yg[ind[2]]]); end
        push!(v_set,v_list)
    end

    return v_set
end

"""
    detect_wada_merge_method(xg,yg, basin)
The algorithm gives the maximum and minimum Haussdorff distances between merged basins. These two distances can help to decide if the basin has the Wada property.

[A. Daza, A. Wagemakers and M. A. F. Sanjuán, Ascertaining when a basin is Wada: the merging method, Sci. Rep., 8, 9954 (2018)]

## Arguments
* `basin` : the matrix containing the information of the basin.
* `xg`, `yg` : 1-dim range vector that defines the grid of the initial conditions to test.

## Example
```
max_dist,min_dist = detect_wada_merge_method(basin,xg,yg)
# grid resolution
epsilon = xg[2]-xg[1]
# if dmax is large then this is not Wada
@show dmax = max_dist/epsilon
@show dmin = min_dist/epsilon
```

"""
function detect_wada_merge_method(xg,yg,bsn::BasinInfo)
    ind  = findall(iseven.(bsn.basin) .== true)
    basin_test = deepcopy(bsn.basin)
    for k in ind; basin_test[k] =basin_test[k]+1; end
    return detect_wada_merge_method(xg,yg,basin_test)
end

function detect_wada_merge_method(xg,yg,basin)

    num_att = length(unique(basin))

    v_list=get_list(basin, xg, yg)

    # compute distances using combbinations of 2 elements from a collection
    ind = combinations(1:num_att,2)

    min_dist = Inf;
    max_dist = 0
    for (k, patron) in enumerate(ind)

       v01 = v_list[patron[1]]
       v02 = v_list[patron[2]]

       hd = haussdorff_dist(v01,v02)

       if hd  > max_dist
           max_dist = hd
       end

       if hd  < min_dist
           min_dist = hd
       end

  end

   return max_dist, min_dist
end

function haussdorff_dist(v01,v02)
       kdtree1 = KDTree(v01)
       idxs12, dists12 = knn(kdtree1, v02, 1, true)
       kdtree2 = KDTree(v02)
       idxs21, dists21 = knn(kdtree2, v01, 1, true)
       max12 = maximum(dists12);
       max21 = maximum(dists21);
       hd =max(max12[1],max21[1])
       return hd
end

mutable struct ds_info{I}
    bsn_nfo::BasinInfo # basin info for BA routine
    integ::I               # integrator
end

function reset_point_data!(ds_nfo)
    reset_bsn_nfo!(ds_nfo.bsn_nfo)
end

"""
    detect_wada_grid_method(integ, bsn_nfo::BasinInfo; max_iter=10)
The algorithm test for Wada basin in a dynamical system. It uses the dynamical system to look if all the atractors are represented in the boundary.

[A. Daza, A. Wagemakers, M. A. F. Sanjuán and J. A. Yorke, Testing for Basins of Wada, Sci. Rep., 5, 16579 (2015)]

## Arguments
* `integ` : the matrix containing the information of the basin.
* `bsn_nfo` : structure that holds the information of the basin as well as the map function. This structure is set when the basin is first computed with `basin_stroboscopic_map` or `basin_poincare_map`.

## Keyword arguments
* `max_iter` : set the maximum depth of subdivisions to look for an atractor. The number of points doubles at each step.

"""
function detect_wada_grid_method(integ, bsn_nfo::BasinInfo; max_iter=10)

   ds_nfo = ds_info(bsn_nfo, integ)
   num_att = bsn_nfo.Na

   if findfirst(x->x==-1, bsn_nfo.basin) != nothing
       @error "The basin contains escapes or undefined attractors, cannot test for Wada"
       return nothing
   end

   # helper function to obtain coordinates
   index_to_coord(p) = [bsn_nfo.xg[p[1]], bsn_nfo.yg[p[2]]]

   # We remove the atractors for this computation
   tmp_bsn = deepcopy(bsn_nfo.basin)
   ind  = findall(iseven.(tmp_bsn) .== true)
   [tmp_bsn[k] = tmp_bsn[k]+1 for k in ind ]

   # obtain points in the boundary
   bnd = get_boundary_filt(tmp_bsn)
   p1_ind = findall(bnd .> 0)

   # initialize empty array of indices and collection of empty sets of colors
   clr_mat = [Set{Int16}() for i=1:length(p1_ind)]
   p2_ind = typeof(p1_ind)(undef, length(p1_ind))
   W = zeros(num_att, max_iter)

   # Initialize matrices (step 1)
   for (k,p1) in enumerate(p1_ind)
       p2, nbgs = get_neighbor_and_colors(tmp_bsn, [p1[1], p1[2]])
       if length(nbgs) > 1
           # keep track of different colors and neighbor point
           push!(clr_mat[k],nbgs...)
           p2_ind[k] = p2
       else
           println("Not in the boundary: weird...")
       end
       W[length(clr_mat[k]),1] += 1
   end

   # Do the iteration!
   for n = 2:max_iter
       for k in 1:length(p1_ind)
           pc1 = index_to_coord(p1_ind[k])
           pc2 = index_to_coord(p2_ind[k])
           # update number of colors
           clr_mat[k]=divide_and_test_W(ds_nfo, pc1, pc2, n, clr_mat[k], num_att)
           # update W matrix
           W[length(clr_mat[k]),n] += 1
       end
       @show W

       # Stopping criterion: if W[Na] in % increases less than ε  then stop or if we have more than 95% boxes in Na
       if  (abs(W[num_att,n] - W[num_att,n-1]) <0.01*W[num_att,n] || W[num_att,n] == 0.0 || W[num_att,n]/sum(W[:,n])>0.95) &&  n > num_att # make at least num_att iterations for stats
           return W[:,n]./sum(W[:,1])
       end

   end

   return W[:,end]./sum(W[:,1])
end

function get_neighbor_and_colors(basin, p)
    radius = 1
    n=p[1]
    m=p[2]
    v=Int16[]
    p2=CartesianIndex(-1,-1)
    # check neihbors and collect basin colors
    for k=n-radius:n+radius, l=m-radius:m+radius
        try
            push!(v,basin[k,l])
            if k != n || l != m
                if basin[n,m] != basin[k,l]
                    p2=CartesianIndex(k,l)
                end
            end
        catch
            #println("error in push!")
        end
    end
    return p2,unique(v)
end

function divide_and_test_W(ds_nfo, p1, p2, nstep, clrs, Na)

    if length(clrs)  == Na
        return clrs
    end

    lerp(t) = (1-t)*p1 + t*p2 # linear interpolation function

    # generates coordinates for initial conditions to test with linear interpolation
    npts = 2^(nstep)
    vecs = Set([lerp(ratio) for ratio in range(0,1, step=1/npts)])  # create a set with the points interspeded
    npts = 2^(nstep-1)
    vecs_before = Set([lerp(ratio) for ratio in range(0,1, step=1/npts)])  # remove points already computed in the previous step
    pnt_to_test = setdiff(vecs,vecs_before) # get the points we have to test

    # get colors and update color set for this box!
    for pnt in pnt_to_test
        clr = get_color_point!(ds_nfo.bsn_nfo, ds_nfo.integ, pnt)
        push!(clrs, clr)
        if length(clrs)  == Na
            break
        end
    end

    return clrs

end


# Import the functions for which we define methods to be called when using PeriodicNumbers
import Base.show
import Base: +, -, *, /, abs, abs2
using LinearAlgebra

#######################################################################################################################
#######################################################################################################################
### Definition of 'PeriodicNumber' with its respective methods for basic operations
#######################################################################################################################
#######################################################################################################################
"""
    PeriodicNumber{T}(value::T, L::T)

Define a number whose value is always contained between 0 and L; more precisely ∈ [0,L).

When it is initialized with value>L or value<0 the value is transformed to be its L-modulus.

"""
struct PeriodicNumber{T<:Real} <:Number
    value::T # Value to be converted if not ∈ [0,L)
    L::T # upper bound
    function PeriodicNumber{T}(value, L) where {T<:Real}
        zero_T = convert(T, 0)
        if L<=zero_T
            error("The period should be a positive number.")
        end
        if (value<zero_T)||(value>=L)
            value = mod(value, L)
            if value==L
                value=zero_T
            end
        end
    new(value, L)
    end
end

# Simple constructor for periodic number
"Constructor for 'PeriodicNumber', with L=1 by default."
PeriodicNumber(x::T, L::T=1.0) where {T<:Real} = PeriodicNumber{T}(x, L)

#### Format the output of periodic numbers
function Base.show(io::IO, PN::PeriodicNumber)
    L = PN.L
    if L==1.0
        print(io, "PN(", PN.value, ")" )
    else
        print(io, "PN(", PN.value, " L=", L, ")" )
    end
end
PeriodicNumber(1.2)

#################
### Definition of basic operations
#################

# Define sum of periodic values among themselves or with real numbers. The result is always a PeriodicNumber type
function Base.:+(PN1::PeriodicNumber, PN2::PeriodicNumber)::PeriodicNumber
    x1=PN1.value; x2=PN2.value
    L1=PN1.L; L2=PN2.L
    if (L1 !=L2)
        error("Comparing PeriodicNumbers of different periodicity")
    end

    return PeriodicNumber(x1+x2, L1)
end

function Base.:+(PN1::PeriodicNumber, x2::Real)::PeriodicNumber
    x1=PN1.value;  L=PN1.L
    return PeriodicNumber(x1+x2, L)
end

Base.:+(x1::Real, PN2::PeriodicNumber)::PeriodicNumber = PN2+x1

# Define product periodic numbers among themselves or with real numbers. The result is always a PeriodicNumber type
function Base.:*(PN1::PeriodicNumber, PN2::PeriodicNumber)::PeriodicNumber
    x1=PN1.value; x2=PN2.value
    L1=PN1.L; L2=PN2.L
    if (L1 !=L2)
        error("Comparing quantities of different periodicity")
    end
    return PeriodicNumber(x1*x2, L1)
end

function Base.:*(PN1::PeriodicNumber, x2::Real)::PeriodicNumber
    x1=PN1.value; L=PN1.L
    return PeriodicNumber(x1*x2, L)
end

Base.:*(x1::Real, PN2::PeriodicNumber)::PeriodicNumber = PN2*x1

#= Define difference of periodic values among themselves or with real numbers. The result is always a PeriodicNumber type, and hence positive.
The difference between periodic values is somewhat tricky, because we need to consider that two particles far apart in the real space might be very close in the periodic one.
For instance, if x1=0 and x2=L, the real distance is L, but in the periodic space it's 0
# NB: The following function is essentially computing a distance in the periodic space, since the value is always positive and
does not depend on the order, i.e. x1-x2 = x2-x1.
=#
"""
    -(PN1::PeriodicNumber, PN2::PeriodicNumber)

Return the difference of two periodic numbers, also as a 'PeriodicNumber' type.

Given that the output is a periodic number, and hence its value is positive, this 'difference
function' is actually commutative.
Moreover, the output's value can never be larger than L/2.
"""
function Base.:-(PN1::PeriodicNumber, PN2::PeriodicNumber)::PeriodicNumber
    x1=PN1.value; x2=PN2.value
    L1=PN1.L; L2=PN2.L
    if (L1 !=L2)
        error("Comparing quantities of different periodicity")
    end
    arg=x1-x2
    arg=min(PeriodicNumber(arg, L1).value, PeriodicNumber(-arg, L1).value )
    return PeriodicNumber(arg, L1)
end

#=
NB: This is also an UNsigned difference.
Thus, it is made to work only with PeriodicNumber's whose lower limit is equal to 0
=#
function Base.:-(PN1::PeriodicNumber, x2::Real)::PeriodicNumber
    x1=PN1.value; L=PN1.L
    dif=abs(x1-x2)
    while dif>0.5*L
        dif -= L
    end
    return PeriodicNumber(dif, L)
end

#NB: Note also that since the difference of a periodic number and any real number is UNsigned, the '-' operator is commutative
Base.:-(x2::Real, PN1::PeriodicNumber)::PeriodicNumber = PN1-x2

#Simple first calls for compilation purposes
PeriodicNumber(0.1, 3.0)+ PeriodicNumber(2.5, 3.0); PeriodicNumber(2.5, 3.0) - PeriodicNumber(0.1, 3.0)
PeriodicNumber(0.1, 3.0) +2.0; PeriodicNumber(0.1, 3.0) - 2.0
rand()*PeriodicNumber(2.5, 3.0); PeriodicNumber(2.5, 3.0)*rand()

# Define other useful methods for computing absolute values, norm, etc.
value(PN::PeriodicNumber{<:Real}) = PN.value
abs(PN::PeriodicNumber{<:Real}) = Base.abs(value(PN))
abs2(PN::PeriodicNumber{<:Real}) = Base.abs2(value(PN))

# Define a 'distance' function to work with numbers with periodic BC
"""
    periodic_distance(PN1::PeriodicNumber, PN2::PeriodicNumber)

Compute the *un*-signed distance between a pair of periodic numbers.

It is implemented simply as the 'value' field of '-' method for periodic numbers.
Since such method always returns a non-negative value and is commutative, this function
is a well defined distance.
"""
periodic_distance(P1::PeriodicNumber, P2::PeriodicNumber)::Float64=(P1-P2).value


"""
    Parameterizations

Useful parameterizations that map various constrained parameters to
the real line or ``R^n``.  The naming convention follows:

```julia
p = XXX_param(x, args...)
x = XXX_value(p, args...)
log_jacobian = XXX_logjac(x, p, args...)
```

where `log_jacobian` gives ``\\log \\left| \\frac{\\mathrm{d}
x}{\\mathrm{d} p} \\right|``, which is the approprate factor to multiply
a density in `x` to produce a density in `p`.  In other words,
`log_jacobian` is the factor that should be introduced into the prior
to properly account for the reparameterization.

Scalar parameterizations also have vector-based versions.
"""
module Parameterizations

"""
    bounded_param(x, low, high)

Variables that are bounded between `low` and `high`.
"""
bounded_param(x, low, high) = log(x-low) - log(high-x)

function bounded_value(p, low, high)
    if p > 0
        ep = exp(-p)
        (high + ep*low)/(one(p) + ep)
    else
        ep = exp(p)
        (high*ep + low)/(one(p) + ep)
    end
end

bounded_value(x::Vector, low, high) = [bounded_value(x[i], low, high) for i in eachindex(x)]

bounded_logjac(x, p, low, high) = log((high-x)*(x-low)/(high-low))
function bounded_logjac(x::Vector, p::Vector, low, high)
    lj = 0.0
    for i in eachindex(x)
        lj += bounded_logjac(x[i], p[i], low, high)
    end
    lj
end

"""
    positive_param(x)

Variable that is constrained to be positive.
"""
function positive_param(x)
    log(x)
end
positive_param(x::Vector) = [positive_param(x[i]) for i in eachindex(x)]

function positive_value(p)
    exp(p)
end
positive_value(p::Vector) = [positive_value(p[i]) for i in eachindex(p)]

function positive_logjac(x, p)
    p
end
function positive_logjac(x::Vector, p::Vector)
    lj = zero(x[1])
    for i in eachindex(x)
        lj += positive_logjac(x[i], p[i])
    end
    lj
end

"""
    increasing_params(x)

Variables that are constrained to be increasing: `x[1] < x[2] < ...`.
"""
function increasing_params(x)
    n = length(x)
    for i in 2:n
        @assert(x[i] > x[i-1], "values must be increasing")
    end

    p = zeros(x)
    p[1] = x[1]
    for i in 2:n
        p[i] = log(x[i]-x[i-1])
    end

    p
end

function increasing_values(p)
    n = length(p)

    x = zeros(p)
    x[1] = p[1]
    for i in 2:n
        x[i] = x[i-1] + exp(p[i])
    end

    x
end

function increasing_logjac(x, p)
    n = length(p)

    lj = 0.0
    for i in 2:n
        lj += p[i]
    end

    lj
end

"""
    increasing_positive_params(x)

Variables that are both positive and increasing.

"""
function increasing_positive_params(x)
    @assert all(x[2:end] .> x[1:end-1])

    n = size(x, 1)
    p = zeros(n)

    p[1] = log(x[1])
    for i in 2:n
        p[i] = log(x[i]-x[i-1])
    end

    p
end

function increasing_positive_values(p)
    n = size(p, 1)
    x = zeros(n)

    x[1] = exp(p[1])
    for i in 2:n
        x[i] = x[i-1] + exp(p[i])
    end

    x
end

function increasing_positive_logjac(x, p)
    sum(p)
end

"""
    simplex_params(x)

Variables that are constrained to be positive and sum to 1.  The
resulting parameterization is reduced in dimension by 1.
"""
function simplex_params(x)
    @assert(abs(sum(x) - 1) < 1e-8, "values must sum to one")

    n = length(x)

    p = zeros(n-1)
    remaining = one(p[1])
    for i in 1:n-1
        p[i] = bounded_param(x[i], zero(x[i]), remaining)
        remaining = remaining - x[i]
    end

    p
end

function simplex_values(p)
    n = length(p)

    x = zeros(n+1)
    remaining = one(p[1])
    for i in 1:n
        x[i] = bounded_value(p[i], zero(p[i]), remaining)
        remaining = remaining - x[i]
    end

    x[n+1] = remaining

    x
end

function simplex_logjac(x, p)
    lj = 0.0
    remaining = 1.0

    for i in 1:length(p)
        lj += bounded_logjac(x[i], p[i], zero(x[i]), remaining)
        remaining = remaining - x[i]
    end

    lj
end

"""
    unit_disk_param(z)

Two-dimensional variable `z` that is constrained to live in the unit
disk: ``|z| < 1``.
"""
function unit_disk_param(z)
    x = z[1]
    y = z[2]

    r2 = x*x + y*y

    a = 1.0 - r2

    Float64[x/a, y/a]
end

function unit_disk_value(p)
    u = p[1]
    v = p[2]

    R2 = u*u + v*v
    R = sqrt(R2)

    r = 0.5*(sqrt(4.0*R2 + 1.0) - 1.0)/R

    x = u*r/R
    y = v*r/R

    Float64[x, y]
end

function unit_disk_logjac(z, p)
    x = z[1]
    y = z[2]

    r2 = x*x + y*y

    3.0*log1p(-r2) - log1p(r2)
end

end


using Printf
using Statistics
using LaTeXStrings

"""
    UncertainValue

Represents a floating point numerical value and an associated uncertainty (1σ).
"""
struct UncertainValue <: AbstractFloat
    value::Float64
    sigma::Float64

    function UncertainValue(val::Real, sigma::Real)
        @assert sigma >= 0.0
        return new(convert(Float64, val), convert(Float64, sigma))
    end
end

Base.convert(::Type{UncertainValue}, n::Real) = UncertainValue(convert(Float64, n), 0.0)
Base.convert(::Type{UncertainValue}, n::UncertainValue) = n

Base.zero(::Type{UncertainValue}) = UncertainValue(0.0, 0.0)
Base.one(::Type{UncertainValue}) = UncertainValue(1.0, 0.0)
Base.add_sum(x::UncertainValue, y::UncertainValue)::Real = x.value + y.value
Base.add_sum(x::AbstractFloat, y::UncertainValue)::Real = x + y.value
Base.hash(uv::UncertainValue, h::UInt) = Base.hash(uv.value, Base.hash(uv.sigma, h))

"""
    Base.round(uva::UncertainValue; defdigits=4)

Provides a (somewhat) intelligent formatter for `UncertainValue` objects that uses the
uncertainty to suggest the number of digits to display.
"""
function Base.round(uva::UncertainValue; defdigits = 4)
    r =
        (σ(uva) >= 1.0e-10 * value(uva)) && (abs(value(uva)) > 1.0e-100) ?
        convert(Int, ceil(max(1, log10(max(σ(uva), abs(value(uva))) / σ(uva))))) + 2 :
        defdigits
    return uv(
        round(value(uva), RoundNearestTiesUp, sigdigits = r),
        round(σ(uva), RoundNearestTiesUp, sigdigits = 2),
    )
end

"""
    poisson(val::Int)

Creates an UncertainValue from an integer which is assumed to have σ = sqrt(val).
"""
poisson(val::Int) = UncertainValue(convert(Float64, val), sqrt(convert(Float64, val)))

"""
    uv(val::Real, σ::Real)

Create an UncertainValue from a real value and 1σ uncertainty.
"""
uv(val::Real, σ::Real=0.0) = UncertainValue(val, σ)
uv(uuvv::UncertainValue) = uuvv

"""
    max(uv1::UncertainValue, uv2::UncertainValue)

Determines the maximum of two values by comparing first the value and second the
uncertainty. (Equal value defer to the larger uncertainty being considered the 'max'.)

If `value(uv1)==value(uv2)` then it is undefined which is larger but to ensure a constant
ordering I define the one with a larger uncertainty to be larger since it is more probable
that it could be larger.
"""
Base.max(uv1::UncertainValue, uv2::UncertainValue) =
    uv1.value ≠ uv2.value ? (uv1.value > uv2.value ? uv1 : uv2) :
    (uv1.sigma > uv2.sigma ? uv1 : uv2)

"""
    min(uv1::UncertainValue, uv2::UncertainValue)

Determines the maximum of two values by comparing first the value and second the
uncertainty. (Equal value defer to the larger uncertainty being considered the 'min'.)

If `value(uv1)==value(uv2)` then it is undefined which is smaller but to ensure a constant
ordering I define the one with a larger uncertainty to be smaller since it is more probable
that it could be smaller.
"""
Base.min(uv1::UncertainValue, uv2::UncertainValue) =
    uv1.value ≠ uv2.value ? (uv1.value > uv2.value ? uv2 : uv1) :
    (uv1.sigma > uv2.sigma ? uv1 : uv2)

Base.minimum(uvs::AbstractVector{UncertainValue}) = reduce(min, uvs)

Base.minimum(xs::Tuple{UncertainValue,Vararg{UncertainValue}}) = reduce(min, xs)

Base.maximum(uvs::AbstractVector{UncertainValue}) = reduce(max, uvs)

Base.maximum(xs::Tuple{UncertainValue,Vararg{UncertainValue}}) = reduce(max, xs)

Base.sum(xs::UncertainValue...) =
    UncertainValue(sum(value(x) for x in xs), sqrt(sum(variance(x) for x in xs)))

Base.sum(uvs::AbstractVector{UncertainValue}) =
    UncertainValue(sum(value(x) for x in uvs), sqrt(sum(variance(x) for x in uvs)))

"""
    Statistics.mean(uvs::AbstractVector{UncertainValue}, minvar=0.0e-10)

The variance weighted mean of a collection of UncertainValue items.
(see https://en.wikipedia.org/wiki/Weighted_arithmetic_mean#Variance_weights).
Be careful! Adding a single value with value!=0 and σ=0.0 will cause NaN.
"""
function Statistics.mean(uvs::AbstractVector{UncertainValue})
    return UncertainValue(
        sum(value(x) / variance(x) for x in uvs) / sum(1.0 / variance(x) for x in uvs),
        sqrt(1.0 / sum(1.0 / variance(x) for x in uvs)),
    )
end
function Statistics.mean(uvs::UncertainValue...)
    return UncertainValue(
        sum(value(x) / variance(x) for x in uvs) / sum(1.0 / variance(x) for x in uvs),
        sqrt(1.0 / sum(1.0 / variance(x) for x in uvs)),
    )
end

"""
    Statistics.std(uvs::UncertainValue...)

The standard deviation of the value portion of the collection of UncertainValue items.
"""
Statistics.std(uvs::UncertainValue...) = std(value(x) for x in uvs)
Statistics.std(uvs::AbstractVector{UncertainValue}) = std(value(x) for x in uvs)

"""
    variance(uv::UncertainValue)

Returns σ².
"""
variance(uv::UncertainValue) = uv.sigma^2

"""
    variance(f::Real)

Returns 0.
"""
variance(f::Real) = zero(typeof(f))

Base.isless(uv1::UncertainValue, uv2::UncertainValue) =
    isequal(uv1.value, uv2.value) ? !isless(uv1.sigma, uv2.sigma) :
    isless(uv1.value, uv2.value)

"""
    σ(uv::UncertainValue)

Returns the 1-σ uncertainty)
"""
σ(uv::UncertainValue) = uv.sigma

"""
    σ(r::Real)

Returns 0.
"""
σ(::Real) = zero(Float64)
σ(m::Missing) = m

"""
    uncertainty(uv::UncertainValue, k::Real=1.0)

Returns the k-σ uncertainty (defaults to k=1.0)
"""
uncertainty(uv::UncertainValue, k::Float64 = 1.0) = k * uv.sigma

"""
    uncertainty(f::Real, k::Real=1.0)

Returns 0.0.
"""
uncertainty(::Real, ::Real) = zero(Float64)

"""
    fractional(uv::UncertainValue)

Computes the fractional uncertainty.
"""
fractional(uv::UncertainValue)::Float64 =
    abs(uv.value) > 1.0e-15 * abs(uv.sigma) ? abs(uv.sigma / uv.value) : Inf64

"""
    fractional(f::Real)

Returns 0
"""
fractional(f::Real) = zero(typeof(f))

"""
    value(uv::UncertainValue)

Returns the value portion. (uv.value)
"""
value(uv::UncertainValue) = uv.value

"""
    value(f::Real)

Returns f
"""
value(f::Real) = Float64(f)
value(m::Missing) = m

"""
    pearson(uv1::UncertainValue, uv2::UncertainValue, covar::Real)

Computes the Pearson correlation coefficient given the covariance between two UncertainValue.
"""
function pearson(uv1::UncertainValue, uv2::UncertainValue, covar::Real)
    res = covar / (σ(uv1) * σ(uv2))
    @assert (res >= -1.0) & (res <= 1.0) "The Pearson correlation coefficient must be on [-1.0, 1.0] - $res"
    return res
end

"""
    covariance(uv1::UncertainValue, uv2::UncertainValue, correlation::Real)

Computes the covariance given the correlation coefficient between two UncertainValue.
"""
function covariance(uv1::UncertainValue, uv2::UncertainValue, correlation::Real)
    @assert (correlation >= -1.0) & (correlation <= 1.0) "The Pearson correlation coefficient must be on [-1.0, 1.0] - $correlation"
    return correlation * σ(uv1) * σ(uv2)
end

"""
    divide(n::UncertainValue, d::UncertainValue, cc::AbstractFloat)

Computes `n/d` where `n` and `d` are UncertainValue and `cc` is the correlation coefficient
defined as `cc = covar(n,d)/( σ(n), σ(d) )`
"""
function divide(a::UncertainValue, b::UncertainValue, cc::Float64=0.0)
    @assert (cc >= -1.0) & (cc <= 1.0) "The Pearson correlation coefficient must be on [-1.0, 1.0] - $cc"
    f, sab = a.value / b.value, covariance(a, b, cc)
    return UncertainValue(
        f, abs(1.0/b.value)*sqrt(a.sigma^2 + ((b.sigma*a.value) / b.value)^2 - 2.0 * sab * a.value / b.value)
    )
end

"""
    multiply(a::UncertainValue, b::UncertainValue, cc::AbstractFloat)

Computes `a*b` where `a` and `b` are UncertainValue and `cc` is the correlation coefficient
defined as `cc = covar(a,b)/( σ(a), σ(b) )`
"""
function multiply(a::UncertainValue, b::UncertainValue, cc::Float64=0.0)
    @assert (cc >= -1.0) & (cc <= 1.0) "The Pearson correlation coefficient must be on [-1.0, 1.0] - $cc"
    f, sab = a.value * b.value, covariance(a, b, cc)
    x, y = a.sigma * b.value, b.sigma*a.value
    return UncertainValue(f, sqrt(x^2 + y^2 + 2.0 * cc * x * y))
end

"""
    add(ka::Real, a::UncertainValue, kb::Real, b::UncertainValue, cc::AbstractFloat)

Computes `ka*a + kb*b` where `a` and `b` are UncertainValue and `cc` is the correlation
coefficient defined as `cc = covar(a,b)/( σ(a), σ(b) )`
"""
add(ka::Real, a::UncertainValue, kb::Real, b::UncertainValue, cc::AbstractFloat) =
    UncertainValue(
        convert(Float64, ka) * a.value + convert(Float64, kb) * b.value,
        sqrt((ka * a.sigma)^2 + (kb * b.sigma)^2 + 2.0 * ka * kb * covariance(a, b, cc)),
    )

Base.:*(a::Real, B::UncertainValue) =
    UncertainValue(convert(Float64, a) * B.value, abs(convert(Float64, a) * B.sigma))
Base.:*(A::UncertainValue, b::Real) =
    UncertainValue(convert(Float64, b) * A.value, abs(convert(Float64, b) * A.sigma))

Base.:-(A::UncertainValue) = UncertainValue(-A.value, A.sigma)
Base.:+(A::UncertainValue) = UncertainValue(A.value, A.sigma)

Base.abs(A::UncertainValue) = A.value >= 0.0 ? A : UncertainValue(abs(A.value), A.sigma)

Base.:/(a::Real, B::UncertainValue) = UncertainValue(
    convert(Float64, a) / B.value,
    abs((convert(Float64, a) * B.sigma) / (B.value * B.value)),
)
Base.:/(A::UncertainValue, b::Real) =
    UncertainValue(A.value / convert(Float64, b), abs(A.sigma / convert(Float64, b)))

Base.inv(B::UncertainValue) = UncertainValue(1.0 / B.value, abs(B.sigma / (B.value^2)))

Base.:^(A::UncertainValue, b::Real) =
    UncertainValue(A.value^b, abs(A.value^(b - one(b)) * b * A.sigma))

Base.log(A::UncertainValue) = UncertainValue(log(A.value), abs(A.sigma / A.value))

Base.exp(A::UncertainValue) = UncertainValue(exp(A.value), exp(A.value) * A.sigma)

Base.sin(A::UncertainValue) = UncertainValue(sin(A.value), abs(cos(A.value) * A.sigma))

Base.cos(A::UncertainValue) = UncertainValue(cos(A.value), abs(sin(A.value) * A.sigma))

Base.tan(A::UncertainValue) = UncertainValue(tan(A.value), sec(A.value)^2 * A.sigma)

Base.sqrt(A::UncertainValue) =
    UncertainValue(sqrt(A.value), abs(A.value^(-0.5) * 0.5 * A.sigma))

Base.isequal(A::UncertainValue, B::UncertainValue)::Bool =
    isequal(A.value, B.value) && isequal(A.sigma, B.sigma)

Base.isapprox(
    A::UncertainValue,
    B::UncertainValue;
    atol::Float64 = 0.0,
    rtol::Float64 = atol == 0.0 ? 0.0 : √eps(Float64),
)::Bool =
    isapprox(A.value, B.value, rtol = rtol, atol = atol) &&
    isapprox(A.sigma, B.sigma, rtol = rtol, atol = atol)

"""
    equivalent(A::UncertainValue, B::UncertainValue, k=1.0)

A and B are equivalent if |A-B| <= k σ(A-B)
"""
equivalent(A::UncertainValue, B::UncertainValue, k = 1.0) =
    abs(A.value - B.value) <= k * sqrt(A.sigma^2 + B.sigma^2)

function Base.parse(::Type{UncertainValue}, str::AbstractString)::UncertainValue
    sp = split(str, r"(?:\+\-|\-\+|±)")
    val = parse(Float64, sp[1])
    sigma = length(sp) == 2 ? parse(Float64, sp[2]) : 0.0
    UncertainValue(val, sigma)
end

function _showstr(uv::UncertainValue, pm = "±")
    smin = max(uv.sigma, 1.0e-6 * abs(uv.value), 1.0e-10)
    ls, lr = map(v -> floor(Int, log10(abs(v))), ( smin, max(1.0e-10, uv.value / smin)))
    return if (ls < -4) || (ls > 6)
        if lr < 1
            @sprintf("%0.1e %s %0.1e", uv.value, pm, uv.sigma)
        elseif lr < 2
            @sprintf("%0.2e %s %0.1e", uv.value, pm, uv.sigma)
        elseif lr < 3
            @sprintf("%0.3e %s %0.1e", uv.value, pm, uv.sigma)
        elseif lr < 4
            @sprintf("%0.4e %s %0.1e", uv.value, pm, uv.sigma)
        elseif lr < 5
            @sprintf("%0.5e %s %0.1e", uv.value, pm, uv.sigma)
        else
            @sprintf("%0.6e %s %0.1e", uv.value, pm, uv.sigma)
        end
    else # lv < 0 && lv > -4
        if lr > 6
            @sprintf("%e %s %0.1e", uv.value, pm, uv.sigma)
        elseif ls >= 0
            @sprintf("%0.0f %s %0.0f", uv.value, pm, uv.sigma)
        elseif ls == -1
            @sprintf("%0.2f %s %0.2f", uv.value, pm, uv.sigma)
        elseif ls == -2
            @sprintf("%0.3f %s %0.3f", uv.value, pm, uv.sigma)
        else# if ls==-3
            @sprintf("%0.4f %s %0.4f", uv.value, pm, uv.sigma)
        end
    end
end

Base.show(io::IO, uv::UncertainValue) = print(io, _showstr(uv))

"""
    LaTeXStrings.latexstring(uv::UncertainValue; fmt=nothing, mode=:normal[|:siunitx])::LaTeXString

Converts an `UncertainValue` to a `LaTeXString` in a reasonable manner.
`mode=:siunitx" requires `\\usepackage{siunitx}` which defines `\\num{}`.
`fmt` is a C-style format string like "%0.2f" or nothing for an "intelligent" default.
"""
function LaTeXStrings.latexstring(
    uv::UncertainValue;
    fmt = nothing,
    mode = :normal,
)::LaTeXString
    pre, post = (mode == :siunitx ? ("\\num{", "}") : ("", ""))
    if isnothing(fmt)
        return latexstring(pre * _showstr(uv, raw"\pm") * post)
    else
        fmtrfunc = generate_formatter(fmt)
        return latexstring(
            pre * fmtrfunc(value(uv)) * " \\pm " * fmtrfunc(uncertainty(uv)) * post,
        )
    end
end


# ------------------------------------------------------------------------------------------
# # Multiple Dispatch
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# In this notebook we'll explore **multiple dispatch**, which is a key feature of Julia.
#
# Multiple dispatch makes software *generic* and *fast*!
#
# #### Starting with the familiar
#
# To understand multiple dispatch in Julia, let's start with what we've already seen.
#
# We can declare functions in Julia without giving Julia any information about the types of
# the input arguments that function will receive:
# ------------------------------------------------------------------------------------------

square(x) = x^2

square(10)

square("Hello ")

square([1,2,3])

# ------------------------------------------------------------------------------------------
# #### Specifying the types of our input arguments
#
# However, we also have the *option* to tell Julia explicitly what types our input arguments
# are allowed to have.
#
# For example, let's write a function `f` that only takes `Number`s as inputs.
# ------------------------------------------------------------------------------------------

f(a::Integer, b::Integer) = "a and b are both integers"

f(3, 4)

f(1.2, 3.4)

# ------------------------------------------------------------------------------------------
# But we can define that method!
# ------------------------------------------------------------------------------------------

f(a::Float64, b::Float64) = "a and b are both Float64s"

f(1.2, 3.4)

# ------------------------------------------------------------------------------------------
# ### Basic dispatch
# ------------------------------------------------------------------------------------------

f(a, b) = "fallback"
f(a::Number, b::Number) = "a and b are both numbers"
f(a::Number, b) = "a is a number"
f(a, b::Number) = "b is a number"
f(a::Integer, b::Integer) = "a and b are both integers"

methods(f)

f(1.5, 2)

f(1, "bar")

f(1, 2)

f("foo", [1,2])

f(1, 2, 3)

# ------------------------------------------------------------------------------------------
# ### Ambiguities
# ------------------------------------------------------------------------------------------

g(a::Int, b::Number) = 1
g(a::Number, b::Int) = 2

g(1, 2.5)

g(1.5, 2)

g(1, 2)

g(x::Int, y::Int) = 3

g(1, 2)

# ------------------------------------------------------------------------------------------
# ### "Diagonal" dispatch
# ------------------------------------------------------------------------------------------

f(a::T, b::T) where {T<:Number} = "a and b are both $(T)s"

methods(f)

f(big(1.5), big(2.5))

f(big(1), big(2)) # <== integer rule is more specific

f(a::T, b::T) where {T<:Integer} = "both are $T integers"

methods(f)

f(big(1), big(2))

f("foo", "bar") # <== still doesn't apply to non-numbers

# ------------------------------------------------------------------------------------------
# ### Varargs methods
# ------------------------------------------------------------------------------------------

f(args::Number...) = "$(length(args))-ary heterogeneous call"
f(args::T...) where {T<:Number} = "$(length(args))-ary homogeneous call"

f(1)

f(1, 2, 3)

f(1, 1.5, 2)

f()

f(1, 2) # <== previous 2-arg method is more specific

f("foo") # <== still doesn't apply to non-numbers

# "splat" (more below)
f([1, 2, 3]...)

# ------------------------------------------------------------------------------------------
# ### Optional Arguments
# ------------------------------------------------------------------------------------------

h(x, y = 0) = 2x + 3y

methods(h)

# ------------------------------------------------------------------------------------------
# Shorthand for this:
# ```
# h(x, y) = 2x + 3y
# h(x) = h(x, 0)
# ```
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Keyword Arguments
# ------------------------------------------------------------------------------------------

k(x, y = 0; opt::Bool = false) = opt ? 2x+y : x+2y

k(2)

k(2, 3)

k(2, opt=true)

k(2, 3, opt=true)

foo(x, y; req::Bool) = req ? 2x+y : x+2y

foo(2, 3)

methods(k)

k(2, opt=true)

# ------------------------------------------------------------------------------------------
# ### Keyword arguments: slurp and splat
# ------------------------------------------------------------------------------------------

function allkw(; kw...)
    @show keys(kw)
end

allkw(a=1,b=2)

# ------------------------------------------------------------------------------------------
# Just like iterators can be splatted as positional arguments, dict-like collections and
# named tuples can be splatted as keyword arguments.
# ------------------------------------------------------------------------------------------

function rect(;width=1,height=1,fill="#")
    for i in 1:height
        println(fill^width)
    end
end

params = (width=8,height=3,fill='A')

rect(; params...)

# ------------------------------------------------------------------------------------------
# ### Exercises
#
# #### Exercise 1
#
# Write a function that repeats a string an integer number of times which takes the
# arguments in either order.
#
# #### Exercise 2a
#
# Write a function `F` that returns the tuple `(x, y, k)` where:
# - `x` is the first positional argument and is mandatory
# - `y` is the second positional argument and is optional
# - `k` is an optional keyword argument
#
# The optional arguments should have the following defaults:
# - `y` defaults to `2x`
# - `k` defaults to `2y`
#
# #### Exercise 2b
#
# Write a function `G` just like `F` but with differet defaults:
# - `k` defaults to `2x`
# - `y` defaults to `2k`
# ------------------------------------------------------------------------------------------


# test without resolution
A = ones(Int, 8,8)
edge_count = count_edges(A)
@test length(edge_count) == 0

A = [[ones(Int, 4,4) 2*ones(Int, 4,4)]; [3*ones(Int, 4,4) 4*ones(Int, 4,4)]]
edge_count = count_edges(A)
@test edge_count[(1,2)] == 4
@test edge_count[(1,3)] == 4
@test edge_count[(2,4)] == 4
@test edge_count[(3,4)] == 4

A = ones(Int, 8,8)
A[2:7,2:7] = 2*ones(Int, 6,6)
A[3:6,3:6] = 3*ones(Int, 4,4)
A[4:5,4:5] = 4*ones(Int, 2,2)
edge_count = count_edges(A)
@test edge_count[(1,2)] == 24
@test edge_count[(2,3)] == 16
@test edge_count[(3,4)] == 8

A = ones(Int, 0,0)
edge_count = count_edges(A)
@test length(edge_count) == 0

A = ones(Int, 4,4,4)
A[2:3,2:3,2:3] = 2*ones(Int, 2,2,2)
edge_count = count_edges(A)
@test edge_count[(1,2)] == 24

A = ones(Int, 4,4,4,4)
A[2:3,2:3,2:3,2:3] = 2*ones(Int, 2,2,2,2)
edge_count = count_edges(A)
@test edge_count[(1,2)] == 64

# test with resolution
A = [[ones(Int, 4,4) 2*ones(Int, 4,4)]; [3*ones(Int, 4,4) 4*ones(Int, 4,4)]]
edge_count = count_edges(A, [1,2])
@test edge_count[(1,2)] == 4*2
@test edge_count[(1,3)] == 4*1
@test edge_count[(2,4)] == 4*1
@test edge_count[(3,4)] == 4*2

A = ones(Int, 4,4,4)
A[2:3,2:3,2:3] = 2*ones(Int, 2,2,2)
edge_count = count_edges(A, [1,2,2])
@test edge_count[(1,2)] == 40

A = ones(Int, 4,4,4)
A[2:3,2:3,2:3] = 2*ones(Int, 2,2,2)
@test_throws AssertionError count_edges(A, [2])


# ### No axis

#md # ![FILE_NAME.png](images/FILE_NAME.png)

plot(sin, 0, 2π, axis=false)


# # DUMMY MODELS FOR TESTING

# ## PARTICLE TRACKER (without methods lifted)

# Consider an object that tracks a particle in one dimension, moving,
# in discrete time, at a speed proportional to the distance away from
# some moving target. The particle is initially at rest at the origin,
# which is where the target also begins.

# Calling `train!` on the object moves it along for the specified
# number of time steps, while calling `ingest!` updates the target
# position.

mutable struct Particle
    position::Float64
    target::Float64
    velocity::Float64
    η::Float64 # learning rate
    training_losses::Vector{Float64}
    Particle(η) = new(0.0, 0.0, 0.0, η, Float64[])
end

Particle(; η=0.1) = Particle(η)

# native train!
function train!(model::Particle)
    model.velocity = model.η*(model.target - model.position)
    model.position = model.position + model.velocity
end

loss(model::Particle) =  abs(model.target - model.position)

function train!(model, n)
    training_losses = map(1:n) do _
        train!(model)
        loss(model)
    end
    model.training_losses = training_losses
    return nothing
end

training_losses(model::Particle) = model.training_losses

function ingest!(model::Particle, target)
    model.target = target
    return nothing
end

# temporary testing area

# model = Particle()
# ingest!(model, 1)
# train!(model, 1)
# @assert loss(model) ≈ 0.9
# ingest!(model, -0.9)
# train!(model, 1)
# @assert loss(model) ≈ 0.9

# model = Particle()
# ingest!(model, 1)
# train!(model, 2)
# @assert training_losses(model) ≈ [0.9, 0.81]


"""
Performs an exhaustive search for a function using a discretized search space.

The total number of function calls will be d^n, where d is the number of
dimensions, and n is the length of the discretized search space.

    Args:
      fn (Function): Objective function
      x_range (Array): Range of x values to explore
      dimensions (Int64): Dimensionality of search space
      spacing (Float64, optional): Discretization
      verbose (Bool, optional): Print results

    Returns:
      Tuple: (x, fx, f_calls) where x is the minimizer, fx is the minimum, and
      f_calls is the number of function calls made.
"""
function exhaustive_search(fn::Function, x_range::Array, dimensions::Int64;
                           spacing=.1, verbose=true)
    f_calls = 1

    i = 1
    k = [1 for _ in 1:dimensions]
    min_x, max_x = x_range
    x_values = collect(min_x:spacing:max_x)
    n_x_values = length(x_values)
    x = [x_values[begin] for _ in 1:dimensions]
    x_star = copy(x)
    f_star = fn(x_star)
    n = length(x)

    while i <= n
        k[i] += 1

        if k[i] > n_x_values
            x[i] = min_x
            k[i] = 1
            i += 1
            continue
        end

        x[i] = x_values[k[i]]

        fnew = fn(x)
        f_calls += 1

        if verbose && (f_calls % 1000000 == 0)
            println("Iterations: $(f_calls)\nx*: $(x_star)\nf*: $(f_star)\n")
        end

        if fnew < f_star
            f_star = fnew
            x_star = copy(x)
        end

        i = 1
    end

    if verbose
        println("Iterations: $(f_calls)\nx*: $(x_star)\nf*: $(f_star)\n")
    end

    return x_star, f_star, f_calls
end


#src # Adding a new equation: scalar conservation laws

# If you want to use Trixi for your own research, you might be interested in
# a new physics model that's not already included in Trixi.jl. In this tutorial,
# we will implement the cubic conservation law
# ```math
# \partial_t u(t,x) + \partial_x u(t,x)^3 = 0
# ```
# in a periodic domain in one space dimension. In Trixi.jl, such a mathematical model
# is encoded as a subtype of [`Trixi.AbstractEquations`](@ref).

# ## Basic setup
using Trixi

struct CubicEquation <: Trixi.AbstractEquations{1 #= number of spatial dimensions =#,
                                                1 #= number of primary variables, i.e. scalar =#};
end

# We create `CubicEquation` as an empty `struct` since we do not use any parameters
# for this equation. Other models could bundle arbitrary parameters, e.g., the
# ideal gas constant for the compressible Euler equations.

# Next, we define the physical flux `f(u) = u^3` using the calling structure
# used in Trixi.jl.

Trixi.flux(u, orientation, equation::CubicEquation) = u.^3
Trixi.varnames(_, ::CubicEquation) = ("scalar",)

# In Trixi.jl, the conserved variables `u` are usually passed as `SVector`s of variables
# at a single physical location. Hence, we must use `u.^3` instead of the scalar
# operation `u^3`.

# That's already enough to run a simple simulation with a standard DGSEM discretization
# using the non-dissipative central flux at interfaces.

using OrdinaryDiffEq

## Create a simulation setup
equation = CubicEquation()

initial_condition_sine(x, t, equation::CubicEquation) = SVector(sinpi(x[1]))

mesh = TreeMesh(-1.0, 1.0, # min/max coordinates
                initial_refinement_level=4,
                n_cells_max=10^4)

solver = DGSEM(3 #= polynomial degree =#, flux_central)

semi = SemidiscretizationHyperbolic(mesh, equation, initial_condition_sine, solver)

# We wrap the return value of the `initial_condition_sine` inside an `SVector` since that's the approach
# used in Trixi.jl also for systems of equations. We need to index the spatial coordinate `x[1]`,
# since it is an `SVector` with one component. In multiple space dimensions, all spatial coordinates
# are passed together.

# Next, we create an `ODEProblem` from the SciML/DifferentialEquations ecosystem.
# We can solve this ODE numerically using any time integration method,
# e.g., `SSPRK43` from [OrdinaryDiffEq.jl](https://github.com/SciML/OrdinaryDiffEq.jl).
# Before, we set up a [callback](@ref callbacks-id) to summarize the simulation setup.

## Create ODE problem with given time span
tspan = (0.0, 0.09)
ode = semidiscretize(semi, tspan)

summary_callback = SummaryCallback()
callbacks = CallbackSet(summary_callback)

## OrdinaryDiffEq's `solve` method evolves the solution in time and executes the passed callbacks
sol = solve(ode, SSPRK43(),
            save_everystep=false, callback=callbacks);

# That's it, you ran your first simulation using your new equation with Trixi! Now, we can plot
# the solution at the final time using Plots.jl.
using Plots
plot(sol)

# You can already see that discontinuities will develop and oscillations start to
# occur around steep parts of the wave. That's expected from our central discretization.
# To avoid these issues, we need to use dissipative numerical fluxes (approximate
# Riemann solvers) at interfaces.

# ## Advanced setup

# Thus, we add a Godunov's flux for our cubic equation. That is easy for this equation
# since the wave speed `f'(u) = 3u^2` is always non-negative.

@inline Trixi.flux_godunov(u_ll, u_rr, orientation, equation::CubicEquation) = flux(u_ll, orientation, equation)

# Let's run the example again but with a dissipative numerical flux at interfaces.
# `remake` will recreate the semidiscretization we used before and only change
# selected parameters, in this case the `solver`.

## A new setup with dissipation
semi = remake(semi, solver=DGSEM(3, flux_godunov))
ode = semidiscretize(semi, tspan)
sol = solve(ode, SSPRK43(), save_everystep=false)
plot!(sol)

# You can see that there are fewer oscillations, in particular around steep edges.
# Now let's increase the final time (and also the spatial resolution).

## A larger final time: Nonclassical shocks develop (you can even increase the refinement to 12)
semi = remake(semi, mesh=TreeMesh(-1.0, 1.0, initial_refinement_level=8, n_cells_max=10^5))
ode = semidiscretize(semi, (0.0, 0.5) #= tspan =#)
sol = solve(ode, SSPRK43(), save_everystep=false)
plot(sol)

# You can observe that nonclassical shocks develop and are stable under grid refinement,
# e.g. for `initial_refinement_level=12`. In this case, these nonclassical shocks
# can be avoided by using an entropy-dissipative semidiscretization. Thus, we need
# to define an entropy-conservative numerical flux

@inline function Trixi.flux_ec(u_ll, u_rr, orientation, equation::CubicEquation)
  return SVector(0.25 * (u_ll[1]^3 + u_ll[1]^2 * u_rr[1] + u_ll[1] * u_rr[1]^2 + u_rr[1]^3))
end

# and use a [`VolumeIntegralFluxDifferencing`](@ref) instead of the standard
# [`VolumeIntegralWeakForm`](@ref) in the DGSEM.

## Let's use a provably entropy-dissipative semidiscretization
semi = remake(semi, solver=DGSEM(3, flux_godunov, VolumeIntegralFluxDifferencing(flux_ec)))
ode = semidiscretize(semi, (0.0, 0.5))
sol = solve(ode, SSPRK43(), save_everystep=false);
plot(sol)

# Possible next steps could be
# - to define `Trixi.max_abs_speeds(u, equations::CubicEquation) = 3 * u[1]^2`
#   to use CFL-based time step control via a [`StepsizeCallback`](@ref)
# - to define quantities of interest like `Trixi.entropy(u, equations::CubicEquation) = u[1]^2`
#   and integrate them in a simulation using the [`AnalysisCallback`](@ref)
# - to experiment with shock-capturing volume integrals [`VolumeIntegralShockCapturingHG`](@ref)
#   and adaptive mesh refinement [`AMRCallback`](@ref)

# ## Summary of the code

# To sum up, here is the complete code that we used (without the callbacks since these create a
# lot of unnecessary output in the doctests of this tutorial).
# In addition, we create the `struct` inside the new module `CubicConservationLaw`. That
# ensures that we can re-create `struct`s defined therein without having to restart Julia.

## Define new physics
module CubicConservationLaw

using Trixi

struct CubicEquation <: Trixi.AbstractEquations{1 #= number of spatial dimensions =#,
                                                1 #= number of primary variables, i.e. scalar =#}
end

@inline Trixi.flux(u, orientation, equation::CubicEquation) = u.^3
Trixi.varnames(_, ::CubicEquation) = ("scalar",)

@inline Trixi.flux_godunov(u_ll, u_rr, orientation, equation::CubicEquation) = flux(u_ll, orientation, equation)
@inline function Trixi.flux_ec(u_ll, u_rr, orientation, equation::CubicEquation)
  return SVector(0.25 * (u_ll[1]^3 + u_ll[1]^2 * u_rr[1] + u_ll[1] * u_rr[1]^2 + u_rr[1]^3))
end

end # module

## Create a simulation setup
import .CubicConservationLaw
using Trixi
using OrdinaryDiffEq
using Plots

equation = CubicConservationLaw.CubicEquation()

initial_condition_sine(x, t, equation::CubicConservationLaw.CubicEquation) = SVector(sinpi(x[1]))

mesh = TreeMesh(-1.0, 1.0, # min/max coordinates
                initial_refinement_level=4,
                n_cells_max=10^4)

solver = DGSEM(3 #= polynomial degree =#, flux_central)

semi = SemidiscretizationHyperbolic(mesh, equation, initial_condition_sine, solver)

## Create ODE problem with given time span
tspan = (0.0, 0.1)
ode = semidiscretize(semi, tspan)

## OrdinaryDiffEq's `solve` method evolves the solution in time and executes the passed callbacks
sol = solve(ode, SSPRK43(), save_everystep=false)
plot(sol)

## A new setup with dissipation
semi = remake(semi, solver=DGSEM(3, flux_godunov))
ode = semidiscretize(semi, tspan)
sol = solve(ode, SSPRK43(), save_everystep=false)
plot!(sol)

## A larger final time: Nonclassical shocks develop (you can even increase the refinement to 12)
semi = remake(semi, mesh=TreeMesh(-1.0, 1.0, initial_refinement_level=8, n_cells_max=10^5))
ode = semidiscretize(semi, (0.0, 0.5))
sol = solve(ode, SSPRK43(), save_everystep=false)
plot(sol)

## Let's use a provably entropy-dissipative semidiscretization
semi = remake(semi, solver=DGSEM(3, flux_godunov, VolumeIntegralFluxDifferencing(flux_ec)))
ode = semidiscretize(semi, (0.0, 0.5))
sol = solve(ode, SSPRK43(), save_everystep=false)
plot(sol)


struct Kron{T,N,I} <: AbstractArray{T,N}
    arrays::I
    function Kron{T,N,I}(A::I) where {T,N,I}
        isempty(A) && throw(ArgumentError("Cannot take kronecker product of empty vectors"))
        new{T,N,I}(A)
    end
end

_kron_dims() = 0
_kron_dims(A, B...) = max(ndims(A), _kron_dims(B...))

Kron{T,N}(A::AbstractArray...) where {T,N} = Kron{T,N,typeof(A)}(A)
Kron{T}(A::AbstractArray...) where {T} = Kron{T,_kron_dims(A...)}(A...)
Kron(A...) = Kron{promote_type(eltype.(A)...)}(A...)

size(K::Kron, j::Int) = prod(size.(K.arrays, j))
size(a::Kron{<:Any,1}) = (size(a,1),)
size(a::Kron{<:Any,2}) = (size(a,1), size(a,2))
size(a::Kron{<:Any,N}) where {N} = (@_inline_meta; ntuple(M -> size(a, M), Val(N)))

getindex(K::Kron{T,1,<:Tuple{<:AbstractVector}}, k::Int) where T =
    first(K.arrays)[k]

function getindex(K::Kron{T,1,<:Tuple{<:AbstractVector,<:AbstractVector}}, k::Int) where T
    A,B = K.arrays
    K,κ = divrem(k-1, length(B))
    A[K+1]*B[κ+1]
end

getindex(K::Kron{T,2,<:Tuple{<:AbstractMatrix}}, k::Int, j::Int) where T =
    first(K.arrays)[k,j]

function getindex(K::Kron{T,2,<:Tuple{<:AbstractArray,<:AbstractArray}}, k::Int, j::Int) where T
    A,B = K.arrays
    K,κ = divrem(k-1, size(B,1))
    J,ξ = divrem(j-1, size(B,2))
    A[K+1,J+1]*B[κ+1,ξ+1]
end

## Adapted from julia/stdlib/LinearAlgebra/src/dense.jl kron definition
function _kron2!(R, K)
    size(R) == size(K) || throw(DimensionMismatch("Matrices have wrong dimensions"))
    a,b = K.arrays
    @assert !has_offset_axes(a, b)
    m = 1
    @inbounds for j = 1:size(a,2), l = 1:size(b,2), i = 1:size(a,1)
        aij = a[i,j]
        for k = 1:size(b,1)
            R[m] = aij*b[k,l]
            m += 1
        end
    end
    R
end

copyto!(R::AbstractMatrix, K::Kron{<:Any,2,<:Tuple{<:AbstractMatrix,<:AbstractMatrix}}) =
    _kron2!(R, K)
copyto!(R::AbstractVector, K::Kron{<:Any,1,<:Tuple{<:AbstractVector,<:AbstractVector}}) =
    _kron2!(R, K)

struct Diff{T, N, Arr} <: AbstractArray{T, N}
    v::Arr
    dims::Int
end

Diff(v::AbstractVector{T}) where T = Diff{T,1,typeof(v)}(v, 1)

function Diff(A::AbstractMatrix{T}; dims::Integer) where T
    dims == 1 || dims == 2 || throw(ArgumentError("dimension must be 1 or 2, got $dims"))
    Diff{T,2,typeof(A)}(A, dims)
end

IndexStyle(::Type{<:Diff{<:Any,1}}) = IndexLinear()
IndexStyle(::Type{<:Diff{<:Any,2}}) = IndexCartesian()

size(D::Diff{<:Any,1}) = (length(D.v)-1,)
function size(D::Diff{<:Any,2})
    if D.dims == 1
        (size(D.v,1)-1,size(D.v,2))
    else #dims == 2
        (size(D.v,1),size(D.v,2)-1)
    end
end

getindex(D::Diff{<:Any, 1}, k::Integer) = D.v[k+1] - D.v[k]
function getindex(D::Diff, k::Integer, j::Integer)
    if D.dims == 1
        D.v[k+1,j] - D.v[k,j]
    else # dims == 2
        D.v[k,j+1] - D.v[k,j]
    end
end

struct Cumsum{T, N, Arr} <: AbstractArray{T, N}
    v::Arr
    dims::Int
end

Cumsum(v::AbstractVector{T}) where T = Cumsum{T,1,typeof(v)}(v, 1)

function Cumsum(A::AbstractMatrix{T}; dims::Integer) where T
    dims == 1 || dims == 2 || throw(ArgumentError("dimension must be 1 or 2, got $dims"))
    Cumsum{T,2,typeof(A)}(A, dims)
end

IndexStyle(::Type{<:Cumsum{<:Any,1}}) = IndexLinear()
IndexStyle(::Type{<:Cumsum{<:Any,2}}) = IndexCartesian()

size(Q::Cumsum) = size(Q.v)

getindex(Q::Cumsum{<:Any, 1}, k::Integer) = k == 1 ? Q.v[1] : Q.v[k] + Q[k-1]
function getindex(Q::Cumsum, k::Integer, j::Integer)
    if Q.dims == 1
        k == 1 ? Q.v[1,j] : Q.v[k,j] + Q[k-1,j]
    else # dims == 2
        j == 1 ? Q.v[k,1] : Q.v[k,j] + Q[k,j-1]
    end
end


### A Pluto.jl notebook ###
# v0.18.4

using Markdown
using InteractiveUtils

# ╔═╡ 066c5b84-c19e-11ec-3f27-115d87575e0d
struct Dual{T}
	val::T
	der::T
end

# ╔═╡ fd3fcf64-b334-4ba7-862f-cdc02b221d6e
a = Dual(3.0,5.0)

# ╔═╡ 393d4637-21a3-4ce3-b720-bca1439c3233
typeof(a.val)

# ╔═╡ ca1251ce-b63f-4075-94c8-c0d184655bdc
Base.:+(f::Dual, g::Dual) = Dual(f.val + g.val, f.der + g.der)

# ╔═╡ 5d59360d-fd53-4996-8400-7d218e4eb29f
Base.:+(f::Dual, α::Number) = Dual(f.val + α, f.der)

# ╔═╡ 32d7e717-6a90-4975-a20a-80c7cabcb85c
Base.:+(α::Number, f::Dual) = f + α

# ╔═╡ 8cdaede1-9825-4404-8898-295c4b0f6e94
Base.:-(f::Dual, g::Dual) = Dual(f.val - g.val, f.der - g.der)

# ╔═╡ 4acc7af2-7eba-4d68-a106-245fb310b9b9
Base.:*(f::Dual, g::Dual) = Dual(f.val * g.val, f.der*g.val + g.der*f.val)

# ╔═╡ 0b9d1885-e113-42f3-b808-90064338f9e2
Base.:^(f::Dual, n::Integer) = Base.power_by_squaring(f,n)

# ╔═╡ 0336e41f-aabd-40a9-af47-dfeed25daa68
begin
	import Base:exp
	exp(f::Dual) = Dual( exp(f.val), exp(f.val)*f.der )
end

# ╔═╡ 9f4518e5-b444-4bd9-8ecf-23450a4d9125
b = Dual(2.0,7.0)

# ╔═╡ e63e723a-e023-4a35-bdaa-be6a3bf4b628
a + b

# ╔═╡ eb39443c-1b71-4137-b777-086700d44a51
 a - b

# ╔═╡ 6b25fab1-aee7-497d-8055-444a4c03aedb
a * b

# ╔═╡ 85588519-64e0-41ba-872a-1942c5a6925d
a^3

# ╔═╡ d1227abf-b2e8-43f8-8efb-aa353455d5ba
a*(a+b)

# ╔═╡ 0609e93e-f7b4-462c-82ba-d674871dc4d5
exp(a)

# ╔═╡ Cell order:
# ╠═066c5b84-c19e-11ec-3f27-115d87575e0d
# ╠═fd3fcf64-b334-4ba7-862f-cdc02b221d6e
# ╠═393d4637-21a3-4ce3-b720-bca1439c3233
# ╠═ca1251ce-b63f-4075-94c8-c0d184655bdc
# ╠═5d59360d-fd53-4996-8400-7d218e4eb29f
# ╠═32d7e717-6a90-4975-a20a-80c7cabcb85c
# ╠═8cdaede1-9825-4404-8898-295c4b0f6e94
# ╠═4acc7af2-7eba-4d68-a106-245fb310b9b9
# ╠═0b9d1885-e113-42f3-b808-90064338f9e2
# ╠═0336e41f-aabd-40a9-af47-dfeed25daa68
# ╠═9f4518e5-b444-4bd9-8ecf-23450a4d9125
# ╠═e63e723a-e023-4a35-bdaa-be6a3bf4b628
# ╠═eb39443c-1b71-4137-b777-086700d44a51
# ╠═6b25fab1-aee7-497d-8055-444a4c03aedb
# ╠═85588519-64e0-41ba-872a-1942c5a6925d
# ╠═d1227abf-b2e8-43f8-8efb-aa353455d5ba
# ╠═0609e93e-f7b4-462c-82ba-d674871dc4d5


module SampleJuliaProj

# Write your package code here.

using Combinatorics
export find_subset
greet() = "Hello World"
greet_random() = "Hello World Random"
# Write your package code here.
export func

"""
    func(x)

Returns double the number `x` plus `1`.
"""
func(x) = 2x + 1

"""
	find_subset(n::Int64,p,q,RV::AbstractString="X")
```@docs
Given `i` and `j` compute `Κ ⊆ \U1d4a9 \\{i,j}`; i.e., all non exclusive subsets. `i,j` can also be empty (i.e. []), in which case the non-empty superset gets listed. An optional prefix can be added (default is `X`).
#### Examples
```

Given `i` and `j` compute `Κ ⊆ \U1d4a9 \\{i,j}`; i.e., all non exclusive subsets. `i,j` can also be empty (i.e. []), in which case the non-empty superset gets listed. An optional prefix can be added (default is `X`).
#### Examples
```julia-repl
julia>find_subset(4,1,3,"X")
["" "X1" "X1,X2" "X1,X2,X5" "X1,X5" "X2" "X2,X5" "X5"]

julia>find_subset(4,1,3,"")
["" "1" "1,2" "1,2,5" "1,5" "2" "2,5" "5"]

julia> find_subset(5,4,3,"🍒")
["" "🍒3" "🍒3,🍒4" "🍒4"]

julia> find_subset(5,[],[],"🍓")
["🍓1" "🍓1,🍓2" "🍓1,🍓2,🍓3" "🍓1,🍓3" "🍓2" "🍓2,🍓3" "🍓3"]
```
"""
function find_subset(n::Int64,p,q,RV::AbstractString="X")
	s="$(RV)" .* string.(1:n)
	s1="$(RV)" .* string(p)
	s2="$(RV)" .* string(q)
	PowerSet = collect(powerset(s,1,n)) # Skip empty set
	setXk=[]
	for kk = 1:Int64(exp2(n)-1)
		subSet1 = vcat(PowerSet[kk,:][1])
		subSet2 = setdiff(subSet1,[s1,s2])
		setXk=[setXk; join(subSet2,",")]
	end

	return sort(unique(setXk))
end

end


# julia1_4.4.jl Perceptron Part 2
# using UCI Machine learning laboratory data
# you need to know your training/test data
# it generally doesn't come in a ready-to-use form
# This sonar data is 60 magnitudes in a time sequence
#  followed by a letter M=mine and R =rock
# read in training data
# Create Sonar Data Set
using Printf, Random, LinearAlgebra
numsonar=208   # total number of sonar sweeps
trainln=138    # number of training samples
testln=numsonar-trainln  # number of test samples
burstln=60     # sonar echo burst length
#
traindata=Array{Float64}(undef,trainln,burstln)
trainlabel=Array{Int32}(undef,trainln)
testdata=Array{Float64}(undef,testln,burstln)
testlabel=Array{Int32}(undef,testln)
#
function create_sonar_data_arrays()
    global traindata, trainlabel, testdata, testlabel
    sdata=Array{Float64}(undef,numsonar,burstln)
    stype=Array{Int32}(undef,numsonar)
    sonardata=zeros(Float64,burstln)
    f1=open("sonar-data.csv","r")
    for i=1:numsonar
        sonarlist=split(readline(f1),",")
        stype[i]=Char(sonarlist[burstln+1][1])
        sdata[i,:]=parse.(Float64,sonarlist[1:burstln])
    end
    # randonize and partition training from test data
    seq=randperm(numsonar) # create random index sequence
    for i=1:numsonar
        sonardata[:]=sdata[seq[i],:]
        sonartype=(stype[seq[i]]==Int32('M')) ? 1 : 0
        if i <= trainln
            traindata[i,:]=sonardata
            trainlabel[i]=sonartype
        else
            testdata[i-trainln,:]=sonardata
            testlabel[i-trainln]=sonartype
        end
    end
    close(f1)
end
create_sonar_data_arrays()
#
###
#
# Training module
inputsets=traindata
labels=trainlabel
# parameters
epocs=2000
global learning_rate=.00045 #learning_rate can also be a function that
        # decreases with increasing epochs eg. (1.0/(1.0+log10(epoch#)))
global bias=0.0
global weights=zeros(Float64,burstln)  #could also be randn(Float64,burstln)

# prediction function

predict(sinput,weights,bias) = sinput'*weights + bias > 0.0 ? 1 : 0

# training
function train(inputsets,weights,bias)
    sinput=zeros(Float64,burstln) # will be sonar input sample
    save_weights=zeros(Float64,burstln)
    save_bias=0.0
    last_serror=70
    for i=1:epocs                 # loop through epochs
        score=0
        error=0.0
        serror=0.0
        seq=randperm(trainln)
        lr=learning_rate   # *(1.0/(1.0+log10(i)))
        for j=1:trainln
            sinput[:] = inputsets[seq[j],:]
            p = predict(sinput,weights,bias)
            if p != labels[seq[j]]
                error=labels[seq[j]]-p # either 1 of -1 (could also be the
                               # actual amount that crosses zero in predict)
                serror+=error^2        # number of misses (or actual ∑miss^2 of 0)
                weights .+= lr * error .* sinput
                bias += lr * error
            else
                score+=1
            end
            #println(seq[j]," ",p," ",labels[seq[j]]," ",error," ",sinput[1]," ",weights[1]," ",bias)
        end
        if serror<last_serror
            @printf("epoch=%4d, perscore=%5.1f, serror=%5f1\n",i,score*100.0/trainln,serror)
            last_serror=serror
            save_weights=weights
            save_bias=bias
        end
    end
    return (save_weights,save_bias)
end
(weights,bias)=train(inputsets,weights,bias)

#
###
#
# test the new estimated weights and bias
inputsets=testdata
labels=testlabel

function test(inputsets,weights,bias)
    sinput=zeros(Float64,burstln) # will be sonar input sample
    error=0
    serror=0
    score=0
    for j=1:testln
        sinput[:] = inputsets[j,:]
        p = predict(sinput,weights,bias)
        if p != labels[j]
            error=labels[j]-p
            serror+=error^2
        else
            score+=1
        end
        #println(seq[j]," ",p," ",labels[seq[j]]," ",error," ",sinput[1]," ",weights[1]," ",bias)
    end
    @printf("perscore=%5.2f, serror=%3d\n",score*100.0/testln,serror)
end
test(inputsets,weights,bias)


@testset "Models of elliptic curves" begin
  @testset "Weierstraß model computation" begin
    E = EllipticCurve([1,2,3,4,5])
    EE, f, g = @inferred short_weierstrass_model(E)
    @test is_short_weierstrass_model(EE)
    @test a_invars(EE) == (0, 0, 0, fmpq(61, 16), fmpq(127, 32))
    P = E([1, 2])
    @test P == g(f(P))

    R = GF(5)
    E = EllipticCurve(map(R, [1, 2, 3, 4, 5]))
    EE, f, g = @inferred short_weierstrass_model(E)
    @test is_short_weierstrass_model(EE)
    @test a_invars(EE) == (0, 0, 0, R(1), R(1))
    P = rand(EE)
    @test P == f(g(P))
    # @inferred will break the tests
  end
  @testset "Simplified model computation" begin
    #j!=0
    R = GF(2, 5)
    o = gen(R)
    E = EllipticCurve(map(R, [o^3, 1+o, 1, 1, o^2+o+1]))
    EE, f, g = @inferred simplified_model(E)
    @test is_simplified_model(EE)
    @test a_invars(EE) == (1, o^4 + o^3 + o + 1, 0, 0, o^2 + 1)
    P = rand(EE)
    @test P == f(g(P))

    #j=0
    o = gen(R)
    E = EllipticCurve(map(R, [0, o^3 + o, o^3 + o^2 + 1, o^3 + o^2 + o, o + 1]))
    EE, f, g = @inferred simplified_model(E)
    @test is_simplified_model(EE)
    @test a_invars(EE) == (0, 0, o^3 + o^2 + 1, 0, 0)
    P = rand(EE)
    @test P == f(g(P))

    #j!=0
    R = GF(3, 5)
    o = gen(R)
    E = EllipticCurve(map(R, [o^3, 1+o, 1, 1, o^2+o+1]))
    EE, f, g = @inferred simplified_model(E)
    @test is_simplified_model(EE)
    @test a_invars(EE) == (0, o^2 + 1, 0, 0, o^4 + 2*o^3 + o + 2)
    P = rand(EE)
    @test P == f(g(P))

    #j=0
    E = EllipticCurve(map(R, [2*o^4 + 2*o^3 + o^2 + o + 1, 2*o^4 + 2*o^3 + o + 1, 0, o^4 + 2*o^3 + o^2, o^3 + o + 2]))
    EE, f, g = @inferred simplified_model(E)
    @test is_simplified_model(EE)
    @test a_invars(EE) == (0, 0, 0, o^4 + 2*o^3 + o^2, o^3 + o + 2)
    P = rand(EE)
    @test P == f(g(P))

  end

end


using Base.Iterators
using Zygote

"""
    xlogx(x::Number)

Compute `x * log(x)`, returning zero if `x` is zero. (Copied from StatsFuns.jl package.)
"""
function xlogx(x::Number)
    result = x * log(x)
    ifelse(iszero(x), zero(result), result)
end

"""
Compute the von Neumann entropy of a density matrix `ρ`.
"""
neumann_entropy(ρ::AbstractMatrix) = -sum(xlogx.(real(eigvals(Matrix(ρ)))))

"""
    binary_entropy(p::Number)

Compute the binary entropy of `p` (natural logarithm).
"""
binary_entropy(p::Number) = -(xlogx(p) + xlogx(1 - p))

"""
Compute `-log(Z)`, with `Z` the partition function.
"""
thermal_logZ(βH::AbstractMatrix) = -log(tr(exp(-Matrix(βH))))

"""
Trace distance to target density matrix `σ`.
"""
trdistance(θlist, σ, cgc) = 0.5 * opnorm(real.(matrix(apply(latent_density(θlist), cgc))) - σ, 1)

"""
Use representation of density matrix in terms of Pauli matrices; cos(θ) is Bloch vector coefficient of Pauli-Z for each qubit
"""
latent_density(θlist) = DensityMatrix(kron([[1, 0, 0, cos(θ)] for θ in θlist]...), length(θlist))
# """
# Use representation of density matrix in terms of Pauli matrices; cos(θ) is Bloch vector coefficient of Pauli-Z for each qubit
# """
# function latent_density2(θlist)
#     ρ = kron([[θ1, 0, 0, θ2] for (θ1,θ2) in partition(θlist,2)]...)
#     return DensityMatrix(ρ ./ ρ[1], length(θlist)÷2)
# end

# Zygote.@adjoint function partition(θlist, step::Int)
#     return partition(θlist, step), θ->(collect(Base.Iterators.flatten(θ)), 0)
# end

function ngradient(f, xs::AbstractArray...)
    grads = zero.(xs)
    for (x, Δ) in zip(xs, grads), i in 1:length(x)
        δ = sqrt(eps())
        tmp = x[i]
        x[i] = tmp - δ/2
        y1 = f(xs...)
        x[i] = tmp + δ/2
        y2 = f(xs...)
        x[i] = tmp
        Δ[i] = (y2-y1)/δ
        if eltype(x) <: Complex
            # derivative with respect to imaginary part
            x[i] = tmp - im*δ/2
            y1 = f(xs...)
            x[i] = tmp + im*δ/2
            y2 = f(xs...)
            x[i] = tmp
            Δ[i] += im*(y2-y1)/δ
        end
    end
    return grads
end


### type definitions ###

"""
A celestial object, such as a planet, star, or moon.
"""
abstract type CelestialObject end

"""
A Keplerian orbit, defined, by semimajor axis (`.a`), eccentricity (`.e`), inclination(`.i`),
right ascension of the ascending node (`.Ω`), argument of the periapsis (`.ω`),
mean anomaly at epoch(`.Mo`), and an epoch (`.epoch`), as well as the celestial body around
which the orbit takes place (.`primary`).
"""
struct Orbit{P}
    a::Float64
    e::Float64
    i::Float64
    Ω::Float64
    ω::Float64
    Mo::Float64
    epoch::Float64
    primary::P
    period::Float64
    basis::RotZXZ{Float64}
end

Base.hash(orb::Orbit, h::UInt) = hash_struct(orb, h, (:a, :e, :i, :Ω, :ω, :Mo, :epoch, :primary))

### helper methods ###
"""
    Δt = period(a, μ)

Compute the period of an orbit with semimajor axis `a` around a body with
gravity parameter `μ`.
"""
period(a, μ) = 2π*√(abs(a)^3 / μ)

## angle-based methods ##
"""
    ϕ = flight_path_angle(θ, e)
    ϕ = flight_path_angle(θ, orbit)

Compute the flight path angle of an orbit with eccentricity `e` at true anomaly `θ`.
"""
flight_path_angle(θ, e) = atan(e*sin(θ)/(1+e*cos(θ)))
flight_path_angle(θ, orb::Orbit) = flight_path_angle(θ, orb.e)

"""
    r = orbital_distance(θ, a, e)
    r = orbital_distance(θ, orbit)

Compute the orbital distance of an orbit with semimajor axis `a` and eccentricity `e`
at true anomaly `θ`, via the [conic-section orbit equation](https://en.wikipedia.org/wiki/Kepler_orbit#Development_of_the_laws).
"""
orbital_distance(θ, a, e) = a*(1-e^2)/(1+e*cos(θ))
orbital_distance(θ, orb::Orbit) = orbital_distance(θ, orb.a, orb.e)

"""
    θ = orbital_angle(r, a, e)
    θ = orbital_angle(r, orbit)

Compute the true anomaly of an orbit with semimajor axis `a` and eccentricity `e`
at orbital distance `r`, via the [conic-section orbit equation](https://en.wikipedia.org/wiki/Kepler_orbit#Development_of_the_laws).
"""
orbital_angle(r, a, e) = wrap_acos((a*(1-e^2)/r - 1)/e)
orbital_angle(r, orb::Orbit) = orbital_angle(r, orb.a, orb.e)

"""
    ψ = orbital_direction(θ, e)
    ψ = orbital_direction(r, a, e)

Compute the direction of velocity of an orbit with semimajor axis `a` at true anomaly `θ`.

Alternatively to true anomaly, orbital distance `r` and eccentricity `e` can be provided.
"""
orbital_direction(θ, e) = θ + π/2 - flight_path_angle(θ, e)
orbital_direction(r, a, e) = orbital_direction(orbital_angle(r,a,e), e)

"""
    r̄ = orbital_position(θ, a, e, basis)
    r̄ = orbital_position(θ, a, e, i, Ω, ω)
    r̄ = orbital_position(θ, orbit)

Compute the position vector at true anomaly `θ` of an orbit with with semimajor axis `a`,
eccentricity `e`, and a rotation `basis` defining the orbit's perifocal plane.

The rotation can be computed from inclinaion `i`, RAAN `Ω`, and the argument of the periapsis `ω`.
"""
function orbital_position(θ, a, e, basis::Rotation{3,Float64})
    r = orbital_distance(θ, a, e)
    ō = r * @SVector([cos(θ),sin(θ),0])
    return perifocal_to_inertial_bases(ō, basis)
end
orbital_position(θ, a, e, i, Ω, ω) = orbital_position(θ, a, e, basis_rotation(Ω, i, ω))
orbital_position(θ, orb::Orbit) = orbital_position(θ, orb.a, orb.e, orb.basis)

"""
    v̄ = orbital_velocity(θ, a, e, μ, basis)
    v̄ = orbital_velocity(θ, a, e, μ, i, Ω, ω)
    v̄ = orbital_velocity(θ, orbit)

Compute the velocity vector at true anomaly `θ` of an orbit with with semimajor axis `a`,
eccentricity `e`, gravity paremter `μ`, and a rotation `basis` defining the orbit's perifocal plane.

The rotation can be computed from inclinaion `i`, RAAN `Ω`, and the argument of the periapsis `ω`.
"""
function orbital_velocity(θ, a, e, μ, basis::Rotation{3,Float64})
    r = orbital_distance(θ, a, e)
    ϕ = orbital_direction(θ, e)
    v = √(μ*(2/r - 1/a))
    dōdt = v * @SVector([cos(ϕ),sin(ϕ),0])
    return perifocal_to_inertial_bases(dōdt, basis)
end
orbital_velocity(θ, a, e, μ, i, Ω, ω) = orbital_velocity(θ, a, e, μ, basis_rotation(Ω, i, ω))
orbital_velocity(θ, orb::Orbit) = orbital_velocity(θ, orb.a, orb.e, orb.primary.μ, orb.basis)

"""
    r̄, v̄ = state_vector(θ, a, e, μ, basis)
    r̄, v̄ = state_vector(θ, a, e, μ, i, Ω, ω)
    r̄, v̄ = state_vector(θ, orbit)

Compute the state vector at true anomaly `θ` of an orbit with with semimajor axis `a`,
eccentricity `e`, gravity paremter `μ`, and a rotation `basis` defining the orbit's perifocal plane.

The rotation can be computed from inclinaion `i`, RAAN `Ω`, and the argument of the periapsis `ω`.
"""
state_vector(θ, a, e, μ, basis) = (orbital_position(θ, a, e, basis), orbital_velocity(θ, a, e, μ, basis))
state_vector(θ, a, e, μ, i, Ω, ω) = state_vector(θ, a, e, μ, basis_rotation(Ω, i, ω))
state_vector(θ, orb::Orbit) = state_vector(θ, orb.a, orb.e, orb.primary.μ, orb.basis)

## time-based methods ##

"""
    r̄ = time_orbital_position(t, a, e, i, Ω, ω, Mo, epoch, μ)
    r̄ = time_orbital_position(t, orbit)

Compute the position vector at time `t` from a full set of orbital paremeters, as well as a gravity parameter `μ`.
"""
time_orbital_position(t, a, e, i, Ω, ω, Mo, epoch, μ
    ) = orbital_position(time_to_true(t, e, period(a,μ), Mo, epoch), a, e, i, Ω, ω)
time_orbital_position(t, orb::Orbit
    ) = orbital_position(time_to_true(t, orb.e, orb.period, orb.Mo, orb.epoch), orb.a, orb.e, orb.basis)

"""
    r̄ = time_orbital_velocity(t, a, e, i, Ω, ω, Mo, epoch, μ)
    r̄ = time_orbital_velocity(t, orbit)

Compute the velocity vector at time `t` from a full set of orbital paremeters, as well as a gravity parameter `μ`.
"""
time_orbital_velocity(t, a, e, i, Ω, ω, Mo, epoch, μ
    ) = orbital_velocity(time_to_true(t, e, period(a,μ), Mo, epoch), a, e, μ, i, Ω, ω)
time_orbital_velocity(t, orb::Orbit
    ) = orbital_velocity(time_to_true(t, orb.e, orb.period, orb.Mo, orb.epoch), orb.a, orb.e, orb.primary.μ, orb.basis)

"""
    r̄ = time_orbital_velocity(t, a, e, i, Ω, ω, Mo, epoch, μ)
    r̄ = time_orbital_velocity(t, orbit)

Compute the state vector at time `t` from a full set of orbital paremeters, as well as a gravity parameter `μ`.
"""
time_state_vector(t, a, e, i, Ω, ω, Mo, epoch, μ
    ) = state_vector(time_to_true(t, e, period(a,μ), Mo, epoch), a, e, μ, i, Ω, ω)
time_state_vector(t, orb::Orbit
    ) = state_vector(time_to_true(t, orb.e, orb.period, orb.Mo, orb.epoch), orb.a, orb.e, orb.primary.μ, orb.basis)

### Alternate Constructors ###

Orbit(a,body::CelestialObject
    ) = Orbit(a,0.,0.,0.,0.,0.,0.,body,period(a,body.μ),basis_rotation(0.,0.,0.))
Orbit(a,e,i,Ω,ω,Mo,body::CelestialObject
    ) = Orbit(a,e,i,Ω,ω,Mo,0.,body,period(a,body.μ),basis_rotation(Ω, i, ω))
Orbit(a,e,i,Ω,ω,Mo,epoch,body::CelestialObject
    ) = Orbit(a,e,i,Ω,ω,Mo,epoch,body,period(a,body.μ),basis_rotation(Ω, i, ω))

"""
    orb = Orbit(t, r̄, v̄, prim, epoch)

Compute an orbit from a time `t`, state vector (`r̄, v̄`), primary body `prim`, and epoch.
"""
function Orbit(t, r̄::AbstractVector{<:Real}, v̄::AbstractVector{<:Real}, prim::CelestialObject, epoch=t)
    r = norm(r̄)
    v = norm(v̄)
    μ = prim.μ
    if r == 0
        throw(ArgumentError("Invalid position"))
    end

    a = 1 / (2/r - (v^2)/μ)
    h̄ = cross(r̄,v̄)
    i = acos(h̄[3]/norm(h̄))

    ē = cross(v̄,h̄)/μ - r̄/r
    e = norm(ē)
    if isapprox(e, 0, atol=1e-15)       # if circular
        e = 0.
        ê = @SVector([1.,0.,0.])        # set periapsis at reference x
    else
        ê = ē/e
    end

    n̄ = cross(@SVector([0.,0.,1.]), h̄)
    n = norm(n̄)
    if isapprox(n, 0, atol=1e-15)       # if uninclined
        n = 0.
        n̂ = ê                           # set Ω at periapsis
    else
        n̂ = n̄/n
    end

    Ω = wrap_angle(copysign(1,n̂[2]) * wrap_acos(n̂[1]))
    ω = wrap_angle(copysign(1,ē[3]) * wrap_acos(dot(n̂,ê)))

    θ = angle_in_plane(r̄, basis_rotation(Ω, i, ω))
    M = true_to_mean(θ, e)
    Δt = t - epoch
    if e < 1
        Mo = wrap_angle(M - time_to_mean(Δt, period(a,μ)))
    else
        Mo = M - time_to_mean(Δt, period(a,μ))
    end

    return Orbit(a,e,i,Ω,ω,Mo,epoch,prim)
end

### comparison and sorting methods ###

function Base.isless(orb1::Orbit, orb2::Orbit)
    fnames = fieldnames(Orbit)
    for fn in fnames
        if orb1.fn != orb2.fn
            return isless(orb1.fn, orb2.fn)
        end
    end
    return false
end

# TO DO: account for possibility of negative angles (i.e. inclination)
function Base.isapprox(orb1::Orbit, orb2::Orbit; atol=1e-6, rtol=1e-6)
    a = isapprox(orb1.a, orb2.a; atol=atol, rtol=rtol)
    e = isapprox(orb1.e, orb2.e; atol=atol, rtol=rtol)
    b = isapprox(orb1.basis, orb2.basis;  atol=atol, rtol=rtol)
    pos = isapprox(time_orbital_position(0,orb1), time_orbital_position(0,orb2); atol=atol, rtol=rtol)
    return a && e && b && pos
end

### descriptive display ###

Base.show(io::IO, orb::Orbit) = println(io,
    "Orbit above $(orb.primary.name):\n",
    "  semi-major axis:\t\t$(orb.a) m\n",
    "  eccentricity:\t\t\t$(orb.e)\n",
    "  inclination:\t\t\t$(rad2deg(orb.i))°\n",
    "  RA of ascending node:\t\t$(rad2deg(orb.Ω))°\n",
    "  argument of periapsis:\t$(rad2deg(orb.ω))°\n",
    "  mean anomaly at epoch:\t$(rad2deg(orb.Mo))°\n",
    "  epoch:\t\t\t$(orb.epoch) s"
)


module Simulation

using CSV, DataFrames, Distributions, IntervalCensoredEventTimes,
Ipopt, NLopt, Random, SparseArrays, Statistics, Test

Random.seed!(123)
# number of observations
n = 1_000
# 2 covariates in base model: one binary and one continuous
Z  = [rand(Bernoulli(0.5), n) randn(n)]
βz = [1.0, 0.8]
# auto-correlated covariates for variable selection
p  = 100
ρ  = 0.1
Σx = [ρ^(abs(i - j)) for i in 1:p, j in 1:p]
X  = Matrix(transpose(rand!(MvNormal(Σx), Matrix{Float64}(undef, p, n))))
# first 5 coefficients are signal
βx = zeros(p)
βx[1:5] .= 0.25
# systematic component
η  = Z * βz + X * βx
expη = exp.(-η)
# actual event times
T = rand.(Exponential.(expη))
# examination time points
s = quantile(T, 0.0:0.1:1.0)
s = [s[1]/2; s; s[end]+1]
display(s); println()
# censoring
C = rand([:left_censored, :right_censored, :interval_censored], n)
L = Vector{Float64}(undef, n)
R = Vector{Float64}(undef, n)
for i in 1:n
    if C[i] == :left_censored
        L[i] = 0
        R[i] = s[findfirst(sj -> sj ≥ T[i], s)]
    elseif C[i] == :right_censored
        L[i] = s[findlast(sj -> sj < T[i], s)]
        R[i] = Inf
    elseif C[i] == :interval_censored
        L[i] = s[findlast(sj -> sj < T[i], s)]
        R[i] = s[findfirst(sj -> sj ≥ T[i], s)]
    end
end
@info "Simulated interval-censored data:"
show(DataFrame(T=T, C=C, L=L, R=R)); println()
@info "Base model:"
icm = IntervalCensoredModel(Z, L, R)
solver = Ipopt.IpoptSolver(print_level=3)
@time fit!(icm, solver, init=initialize_uniform!(icm))
display(icm); println()
@info "Sparse model by IHT:"
icsm = IntervalCensoredSparseModel(icm, X, 5)
@time iht!(icsm, verbose=true, debiasing=true, printfreq=1)
βx = sparsevec(icsm.βx)
display(DataFrame(j=βx.nzind, βj=βx.nzval)); println()
end


function priors_onegroup_alpha_delta_variant_cities(θ)
	@unpack β₀,β_school,β_home,β_other,β_work,ϵ,χ,p_test,E₀,inc_R_αvar,time_scale_αvar,mid_point_αvar,inc_R_δvar,time_scale_δvar,mid_point_δvar,init_scale= θ #,S₁,S₂,S₃,S₄,S₅,S₆

    LP = 0.
    LP += logpdf(Beta(50,50),ϵ)
    LP += logpdf(Gamma(3,4.5/3),χ)
    LP += logpdf(Gamma(3,5/3),p_test) ##units relative to 1% detection chance

	LP += logpdf(Gamma(3,1000/3),E₀)

    LP += logpdf(Gamma(10,1.5/10),β₀)
    LP += logpdf(Gamma(10,1.5/10),β_other)
    LP += logpdf(Gamma(10,1.5/10),β_work)
    LP += logpdf(Gamma(10,1.5/10),β_school)
    LP += logpdf(Gamma(10,1.5/10),β_home)

	LP += logpdf(Gamma(15,0.4/15),inc_R_αvar)
	LP += logpdf(Gamma(15,0.15/15),time_scale_αvar)
	LP += logpdf(Gamma(5,0.5/5),mid_point_αvar) #In units of 30 days after Feb 1st
	LP += logpdf(Gamma(15,0.6/15),inc_R_δvar)
	LP += logpdf(Gamma(15,0.2/15),time_scale_δvar)
	LP += logpdf(Gamma(5,0.5/5),mid_point_δvar)
	LP += logpdf(Gamma(5,1.5/5),init_scale)

    return LP
end

function priors_onegroup_alpha_delta_variant_outside_cities(θ)
	@unpack β₀,β_school,β_home,β_other,β_work,ϵ,χ,p_test,E₀,inc_R_αvar,time_scale_αvar,mid_point_αvar,inc_R_δvar,time_scale_δvar,mid_point_δvar,init_scale = θ
	LP = 0.
    LP += logpdf(Beta(50,50),ϵ)
    LP += logpdf(Gamma(3,4.5/3),χ)
    LP += logpdf(Gamma(3,2/3),p_test) #units relative to 1% detection chance

	LP += logpdf(Gamma(3,1000/3),E₀)

    LP += logpdf(Gamma(10,1.5/10),β₀)
    LP += logpdf(Gamma(10,1.5/10),β_other)
    LP += logpdf(Gamma(10,1.5/10),β_work)
    LP += logpdf(Gamma(10,1.5/10),β_school)
    LP += logpdf(Gamma(10,1.5/10),β_home)

	LP += logpdf(Gamma(15,0.4/15),inc_R_αvar)
	LP += logpdf(Gamma(15,0.15/15),time_scale_αvar)
	LP += logpdf(Gamma(5,0.5/5),mid_point_αvar) #In units of 30 days after Feb 1st
	LP += logpdf(Gamma(15,0.6/15),inc_R_δvar)
	LP += logpdf(Gamma(15,0.2/15),time_scale_δvar)
	LP += logpdf(Gamma(5,0.5/5),mid_point_δvar)
	LP += logpdf(Gamma(5,1.5/5),init_scale)

    return LP
end

function priors_onegroup_alpha_delta_variant_noncities_fitted(θ)
    @unpack β₀, β_school, β_home, β_other, β_work, ϵ, χ, p_test, E₀, inc_R_αvar, time_scale_αvar, mid_point_αvar, inc_R_δvar, time_scale_δvar, mid_point_δvar, init_scale = θ

    LP = 0.0
    LP += logpdf(Beta(10, 10), ϵ)

    LP += logpdf(fitted_priors.d_χ, χ)  #
    LP += logpdf(fitted_priors.d_ptest, p_test) # units relative to 1% detection chance

    LP += logpdf(Gamma(3, 1000 / 3), E₀)

    LP += logpdf(Gamma(5, 1.5 / 5), β₀)
    LP += logpdf(Gamma(5, 1.5 / 5), β_other)
    LP += logpdf(Gamma(5, 1.5 / 5), β_work)
    LP += logpdf(Gamma(5, 1.5 / 5), β_school)
    LP += logpdf(Gamma(5, 1.5 / 5), β_home)

    LP += logpdf(Gamma(15, 0.4 / 15), inc_R_αvar)
    LP += logpdf(Gamma(15, 0.15 / 15), time_scale_αvar)
    LP += logpdf(Gamma(5, 0.5 / 5), mid_point_αvar) #In units of 30 days after Feb 1st
    LP += logpdf(Gamma(15, 0.6 / 15), inc_R_δvar)
    LP += logpdf(Gamma(15, 0.2 / 15), time_scale_δvar)
    LP += logpdf(Gamma(5, 0.5 / 5), mid_point_δvar)

	LP += logpdf(Gamma(3, 1.2 / 3), init_scale)

    return LP
end


using MittagLefflerFunctions
using PyPlot

α = 1.85
β = 1.0
x = range(0, 10, length=201)
E1 = MLQuad3(α, β, 6)
E2 = MLQuad3(α, β, 8)

y = mlf.(α, β, -x, 10)
y1 = E1.(x)
y2 = E2.(x)

figure(1)
#plot(x, y1-y, x, y2-y)
semilogy(x, abs.(y1-y), x, abs.(y2-y))
grid(true)


"""
Constructs LTI system with quadratic reward and uniformly sampled initial conditions.
"""
function LTIQuadraticEnv(
    A::AbstractVecOrMat,
    B::AbstractVecOrMat,
    C::AbstractVecOrMat,
    D::AbstractVecOrMat,
    Q::AbstractVecOrMat,
    R::AbstractVecOrMat,
    s0_lb::Union{Real,Vector{<:Real}},
    s0_ub::Union{Real,Vector{<:Real}},
    tspan::Tuple{<:Real,<:Real},
    dt::Real;
    #= Keyword arguments =#
    o_ub::Union{Nothing,Real,Vector{<:Real}}=nothing, # upper bound for observation space
    o_lb::Union{Nothing,Real,Vector{<:Real}}=nothing, # lower bound for observation space
    a_ub::Union{Nothing,Real,Vector{<:Real}}=nothing, # upper bound for action space
    a_lb::Union{Nothing,Real,Vector{<:Real}}=nothing, # lower bound for action space
    solver::DiffEqBase.AbstractODEAlgorithm=Tsit5(),
    reltol::Real=1e-8,
    abstol::Real=1e-8, # TODO: add more kwargs for integrator
    T=Float32
    )

    n_states, n_actions, _ = state_space_validation(A, B, C, D, Continuous())
    ic_sampler = UniformSampler(s0_lb, s0_ub)

    # Sample random IC for dimension check and to define ODEProblem
    s0 = ic_sampler()
    length(s0) == n_states || throw(ArgumentError("Length $(length(s0)) of s0 doesn't match state dimension $(n_states)"))

    function ode(s, a, t)
        ṡ = A * s + B * a
    end

    problem = ODEProblem(ode, s0, tspan)
    reward_fn = QuadraticReward(Q, R)
    observation_fn = LinearObservation(C, D)

    return DiffEqEnv(problem, reward_fn, n_actions, dt, s0_lb, s0_ub;
        observation_fn=observation_fn,
        o_ub=o_ub, o_lb=o_lb, a_ub=a_ub, a_lb=a_lb,
        solver=solver, reltol=reltol, abstol=abstol,
        T=T
    )
end

"""
Constructs LTI system with quadratic reward and constant initial conditions.
"""
function LTIQuadraticEnv(
    A::AbstractVecOrMat,
    B::AbstractVecOrMat,
    C::AbstractVecOrMat,
    D::AbstractVecOrMat,
    Q::AbstractVecOrMat,
    R::AbstractVecOrMat,
    s0::Union{Real,Vector{<:Real}},
    tspan::Tuple{<:Real,<:Real},
    dt::Real;
    #= Keyword arguments =#
    o_ub::Union{Nothing,Real,Vector{<:Real}}=nothing, # upper bound for observation space
    o_lb::Union{Nothing,Real,Vector{<:Real}}=nothing, # lower bound for observation space
    a_ub::Union{Nothing,Real,Vector{<:Real}}=nothing, # upper bound for action space
    a_lb::Union{Nothing,Real,Vector{<:Real}}=nothing, # lower bound for action space
    solver::DiffEqBase.AbstractODEAlgorithm=Tsit5(),
    reltol::Real=1e-8,
    abstol::Real=1e-8, # TODO: add more kwargs for integrator
    T=Float32
    )

    s0_lb = s0
    s0_ub = s0

    return LTIQuadraticEnv(A, B, C, D, Q, R,
        s0_lb, s0_ub, tspan, dt;
        o_ub=o_ub, o_lb=o_lb, a_ub=a_ub, a_lb=a_lb,
        solver=solver, reltol=reltol, abstol=abstol, T=T)
end


using StructuralCausalModels
using GraphRecipes, Plots

ProjDir = @__DIR__
cd(ProjDir) #do

d_string = "dag {A -> {E Z}; B -> {D Z}; Z -> {D E}; E -> D}"

dag = DAG("conf_triangles", d_string);
show(dag)

graphplot(dag.e, names=names(dag.e, 1), curvature_scalar=0.1, nodesize=0.2,
  method=:spring, fontsize=10, arrow=1.0, nodeshape=:circle)

gui()


# This file includes an example file for 1D integrate.

using FractalTools
using Makie

# Generate data
xi, dx, xf = 0., 0.01, 10.
x = collect(xi : dx : xf)
y = sin.(x)
pts = collect.(zip(x, y))
res = integrate(pts, Interp1D(0.0001))