### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# ╔═╡ 5b2ee40e-a2b8-11ea-0fef-c35fe6918860
md"""
# The tower of Hanoi Harry want to understand this

The tower of hanoi is a famous puzzle.

![setup of the tower of a hanoi](https://upload.wikimedia.org/wikipedia/commons/0/07/Tower_of_Hanoi.jpeg)

The game consists of three rods with disks stacked on top of them. The puzzle will start with all disks in a stack on one of the rods (like in the picture). The goal is to move all the discs to a single stack on  different rod.

To move the disks, you have to follow the following rules:

* You can move only one disk at a time.
* For each move, you have to take the upper disk from one of the stacks, and place it on top of another stack or empty rod.
* You cannot place a larger disk on top of a smaller disk.

This notebook will define a Julia implementation of the puzzle. It's up to you to write an algorithm that solves it.
"""

# ╔═╡ 95fbd0d2-a2b9-11ea-0682-fdf783251797
md"""
## Setting up the game pieces

What does a Julia implementation look like? We're not really interested in writing code that will manipulate physical disks. Our final goal is a function that will give us a _recipe_ to solve the tower of hanoi, which is just a list of moves to make. Because of that, we can use a lot of abstraction in our implementation, and keep the data structures as simple as possible.

To start, we have to define some representation of the disks and the stacks. The disks have one important property, which is that they are ordered. We can use integers to represent them.
"""

# ╔═╡ 620d6834-a2ba-11ea-150a-2132bb54e4b3
num_disks = 8

# ╔═╡ 35ada214-a32c-11ea-0da3-d5d494b28467
md"""(Side note: the number of disks is arbitrary. When testing your function, you may want to set it to 1 or 2 to start.)"""

# ╔═╡ 7243cc8e-a2ba-11ea-3f29-99356f0cdcf4
all_disks = 1:num_disks

# ╔═╡ 7e1ba2ac-a2ba-11ea-0134-2f61ed75be18
md"""
A single stack can be represented as an array with all the disks in it. We will list them from top to bottom.
"""

# ╔═╡ 43781a52-a339-11ea-3803-e56e2d08aa83
first_stack = collect(all_disks)

# ╔═╡ b648ab70-a2ba-11ea-2dcc-55b630e44325
md"""
Now we have to make three of those.
"""

# ╔═╡ 32f26f80-a2bb-11ea-0f2a-3fc631ada63d
starting_stacks = [first_stack, [], []]

# ╔═╡ e347f1de-a2bb-11ea-06e7-87cca6f2a240
md"""
## Defining the rules

Now that we have our "game board", we can implement the rules.

To start, we make two functions for states. A state of the game is just an array of stacks.

We will define a function that checks if a state is okay according to the rules. To be legal, all the stacks should be in the correct order, so no larger disks on top of smaller disks.

Another good thing to check: no disks should have appeared or disappeared since we started!
"""

# ╔═╡ 512fa6d2-a2bd-11ea-3dbe-b935b536967b
function islegal(stacks)
	order_correct = all(issorted, stacks)

	# check if we use the same disk set that we started with

	disks_in_state = sort([disk for stack in stacks for disk in stack])
	disks_complete = disks_in_state == all_disks

	order_correct && disks_complete
end

# ╔═╡ c56a5858-a2bd-11ea-1d96-77eaf5e74925
md"""
Another function for states: check if we are done! We can assume that we already checked if the state was legal. So we know that all the disks are there and they are ordered correctly.  To check if we are finished, we just need to check if the last stack contains all the disks.
"""

# ╔═╡ d5cc41e8-a2bd-11ea-39c7-b7df8de6ae3e
function iscomplete(stacks)
	last(stacks) == all_disks
end

# ╔═╡ 53374f0e-a2c0-11ea-0c91-97474780721e
md"""
Now the only rules left to implement are the rules for moving disks.

We could implement this as another check on states, but it's easier to just write a legal `move` function. Your solution will specify moves for the `move` function, so this will be the only way that the stacks are actually manipulated. That way, we are sure that nothing fishy is happening.

We will make our `move` function so that its input consists of a state of the game, and instructions for what to do. Its output will be the new state of the game.

So what should those instructions look like? It may seem intuitive to give a _disk_ that should be moved, but that's more than we need. After all, we are only allowed to take the top disk from one stack, and move it to the top of another. So we only have to say which _stacks_ we are moving between.

(Note that the `move` function is okay with moving a larger disk on top of a smaller disk. We already implemented that restriction in `islegal`.)
"""

# ╔═╡ e915394e-a2c0-11ea-0cd9-1df6fd3c7adf
function move(stacks, source::Int, target::Int)
	# check if the from stack if not empty
	if isempty(stacks[source])
		error("Error: attempted to move disk from empty stack")
	end

	new_stacks = deepcopy(stacks)

	disk = popfirst!(new_stacks[source]) # take disk
	pushfirst!(new_stacks[target], disk) # put on new stack

	return new_stacks
end

# ╔═╡ 87b2d164-a2c4-11ea-3095-916628943879
md"""
## Solving the problem

We have implemented the game pieces and the rules, so you can start working on your solution.

To do this, you can fill in the `solve(stacks)` function. This function should give a solution for the given `stacks`, by moving all the disks from stack 1 to stack 3.

As output, `solve` should give a recipe, that tells us what to do. This recipe should be an array of moves. Each moves is a `(source, target)` tuple, specifying from which stack to which stack you should move.

For example, it might look like this:
"""

# ╔═╡ 29b410cc-a329-11ea-202a-795b31ce5ad5
function wrong_solution(stacks)::Array{Tuple{Int,Int}}
	return [(1, 2), (2, 3), (2, 1), (1, 3)]
end

# ╔═╡ ea24e778-a32e-11ea-3f11-dbe9d36b1011
md"""
Now you can work on building an actual solution. Some tips:
* `solve(stacks)` can keep track of the board if you want, but it doesn't have to.
* The section below will actually run your moves, which is very useful for checking them.
* If you want, you can change `num_disks` to 1 or 2. That can be a good starting point.
"""

# ╔═╡ 010dbdbc-a2c5-11ea-34c3-837eae17416f
function solve(stacks)::Array{Tuple{Int,Int}}

	# what to do?

	return []
end

# ╔═╡ 3eb3c4c0-a2c5-11ea-0bcc-c9b52094f660
md"""
## Checking solutions

This is where we can check a solution. We start with a function that takes our recipe and runs it.
"""

# ╔═╡ 4709db36-a327-11ea-13a3-bbfb18da84ce
function run_solution(solver::Function, start = starting_stacks)
	moves = solver(deepcopy(start)) # apply the solver

	all_states = Array{Any,1}(undef, length(moves) + 1)
	all_states[1] = starting_stacks

	for (i, m) in enumerate(moves)
		try
			all_states[i + 1] = move(all_states[i], m[1], m[2])
		catch
			all_states[i + 1] = missing
		end
	end

	return all_states
end

# ╔═╡ 372824b4-a330-11ea-2f26-7b9a1ad018f1
md"""
You can use this function to see what your solution does.

If `run_solution` tries to make an impossible move, it will give `missing` from that point onwards. Look at what happens in the `wrong_solution` version and compare it to the moves in `wrong_solution`.
"""

# ╔═╡ d2227b40-a329-11ea-105c-b585d5fcf970
run_solution(wrong_solution)

# ╔═╡ 9173b174-a327-11ea-3a69-9f7525f2e7b4
run_solution(solve)

# ╔═╡ bb5088ec-a330-11ea-2c41-6b8b92724b3b
md"""
Now that we have way to run a recipe, we can check if its output is correct. We will check if all the intermediate states are legal and the final state is the finished puzzle.
"""

# ╔═╡ 10fb1c56-a2c5-11ea-2a06-0d8c36bfa138
function check_solution(solver::Function, start = starting_stacks)
	try
		# run the solution
		all_states = run_solution(solver, start)

		# check if each state is legal
		all_legal = all(islegal, all_states)

		# check if the final state is is the completed puzzle
		complete = (iscomplete ∘ last)(all_states)

		all_legal && complete
	catch
		# return false if we encountered an error
		return false
	end
end

# ╔═╡ 8ea7f944-a329-11ea-22cc-4dbd11ec0610
check_solution(solve)

# ╔═╡ e54add0a-a330-11ea-2eeb-1d42f552ba38
if check_solution(solve)
	if num_disks >= 8
		md"""
		Congratulations, your solution works!
		"""
	else
		md"""
		Your solution works for $(num_disks) disks. Change `num_disks` to see if it works for 8 or more.
		"""
	end
else
	md"""
	The `solve` function doesn't work yet. Keep working on it!
	"""
end

# ╔═╡ Cell order:
# ╟─5b2ee40e-a2b8-11ea-0fef-c35fe6918860
# ╟─95fbd0d2-a2b9-11ea-0682-fdf783251797
# ╠═620d6834-a2ba-11ea-150a-2132bb54e4b3
# ╟─35ada214-a32c-11ea-0da3-d5d494b28467
# ╠═7243cc8e-a2ba-11ea-3f29-99356f0cdcf4
# ╟─7e1ba2ac-a2ba-11ea-0134-2f61ed75be18
# ╠═43781a52-a339-11ea-3803-e56e2d08aa83
# ╟─b648ab70-a2ba-11ea-2dcc-55b630e44325
# ╠═32f26f80-a2bb-11ea-0f2a-3fc631ada63d
# ╟─e347f1de-a2bb-11ea-06e7-87cca6f2a240
# ╠═512fa6d2-a2bd-11ea-3dbe-b935b536967b
# ╟─c56a5858-a2bd-11ea-1d96-77eaf5e74925
# ╠═d5cc41e8-a2bd-11ea-39c7-b7df8de6ae3e
# ╟─53374f0e-a2c0-11ea-0c91-97474780721e
# ╠═e915394e-a2c0-11ea-0cd9-1df6fd3c7adf
# ╟─87b2d164-a2c4-11ea-3095-916628943879
# ╠═29b410cc-a329-11ea-202a-795b31ce5ad5
# ╟─ea24e778-a32e-11ea-3f11-dbe9d36b1011
# ╠═010dbdbc-a2c5-11ea-34c3-837eae17416f
# ╟─3eb3c4c0-a2c5-11ea-0bcc-c9b52094f660
# ╠═4709db36-a327-11ea-13a3-bbfb18da84ce
# ╟─372824b4-a330-11ea-2f26-7b9a1ad018f1
# ╠═d2227b40-a329-11ea-105c-b585d5fcf970
# ╠═9173b174-a327-11ea-3a69-9f7525f2e7b4
# ╟─bb5088ec-a330-11ea-2c41-6b8b92724b3b
# ╠═10fb1c56-a2c5-11ea-2a06-0d8c36bfa138
# ╠═8ea7f944-a329-11ea-22cc-4dbd11ec0610
# ╟─e54add0a-a330-11ea-2eeb-1d42f552ba38


#Copyright (c) 2015 Rex Computing and Isaac Yonemoto

#see LICENSE.txt

#this work was supported in part by DARPA Contract D15PC00135

#polynomial.jl
#solves an arbitrary 6-degree polynomial, using the arbitrary function solver

include("../unum.jl")
using Unums

const plength = 4
#const parray = [exp(100 * randn()) * (randbool() ? -1 : 1) for idx=1:plength]
const parray = [6, 5, -7, 1]
const tol = 0.0001

function pwr(x::Utype, n::Int)
  if n == 1
    x
  elseif iseven(n)
    sqr(pwr(x, n ÷ 2))
  else
    x * pwr(x, n - 1)
  end
end

function polynomial(v::Utype)
  T = isa(v, Ubound) ? typeof(v.lowbound) : typeof(v)
  sum = convert(T, parray[1])
  for (idx = 2:plength)
    sum += convert(T, parray[idx]) * pwr(v, idx - 1)
  end
  sum
end

a = solve(polynomial, tol, 0, 0)#, verbose = true)
map((x) -> println("result:", describe(x)), a)

println()
println()

print("solved polynomial: ")
for idx=plength:-1:2
  print(parray[idx], "x^", (idx - 1), " + ")
end
println(parray[1], " to relative tolerance ", tol)

println("solutions: ")
map((d) -> println(describe(d)), a)


using Choices
using Test

# from https://en.wikipedia.org/wiki/Sudoku
puzzle = [
  5 3 0 0 7 0 0 0 0;
  6 0 0 1 9 5 0 0 0;
  0 9 8 0 0 0 0 6 0;
  8 0 0 0 6 0 0 0 3;
  4 0 0 8 0 3 0 0 1;
  7 0 0 0 2 0 0 0 6;
  0 6 0 0 0 0 2 8 0;
  0 0 0 4 1 9 0 0 5;
  0 0 0 0 8 0 0 7 9;
]
solution = [
  5 3 4 6 7 8 9 1 2;
  6 7 2 1 9 5 3 4 8;
  1 9 8 3 4 2 5 6 7;
  8 5 9 7 6 1 4 2 3;
  4 2 6 8 5 3 7 9 1;
  7 1 3 9 2 4 8 5 6;
  9 6 1 5 3 7 2 8 4;
  2 8 7 4 1 9 6 3 5;
  3 4 5 2 8 6 1 7 9;
]

function block_is_consistent(block)
  have = fill(false, 9)
  for n ∈ block
    n == 0 && continue
    have[n] && return false
    have[n] = true
  end
  return true
end

"""
    cell_block(i, j)

Returns indices (a tuple of two ranges) for the 3×3 sudoku block containing the
cell at (i, j).
"""
cell_block(i, j) = map(x -> x .+ (1:3), trunc.(Int, ((i, j) .- 1) ./ 3) .* 3)

cell_is_consistent(p, i, j) =
  block_is_consistent(p[i,1:9]) &&
  block_is_consistent(p[1:9,j]) &&
  block_is_consistent(p[cell_block(i,j)...])

iscomplete(p) = !any(n == 0 for n ∈ p)

function solve(p)
  # start with the unsolved puzzle
  comp = choose([p])
  # for each "blank" (zero) cell
  for k ∈ keys(p)
    if (p[k] == 0)
      # add a piece of the computation that makes a guess for that cell and
      # checks it
      comp =
	comp |> before ->
	  choose(1:9) |> function(guess)
	    after = copy(before) # super inefficient, but whatever, it's a test
	    after[k] = guess
	    cell_is_consistent(after, Tuple(k)...) || fail()
	    return after
	  end
    end
  end
  s = first(comp)
  cut()
  return s
end

@test solve(puzzle) == solution


include("Potts.jl")
using IterTools:product

function pth_order_tensor(r)
    a =  Tuple([2 for i in 1:r])
    M =  zeros(a)
    for i in CartesianIndices(M)
        s = prod(2*[Tuple(i)...] .- 3)
        M[i] = s
    end
    return M
end

n =15
q = 2

Terms = Dict{}()

M =  pth_order_tensor(2)

for i in 1:n-1
    Terms[(i,i+1)] = (0.3 + rand())*M

end

Terms[(1,n)] = (0.3 + rand())*M

Terms[(1,3,5,7,9)] = 0.5*pth_order_tensor(5)

F = FactorGraph([2,5],n,q,Terms);

n_samples =10^6

s = raw_sampler_potts(F, n_samples, false)
samples=[2*a .- 3  for a in s]
sample_path ="saved_model/samples.csv"
writedlm( sample_path,hcat(samples...))


using Arrowhead, Base.Test

srand(0)

import Arrowhead: eigvals_FMM, synthesize_symarrow

@testset "Testing symmetric arrowhead methods" begin
    for n in (100, 300, 1000)

        A = SymArrow((1.0:n-1).*(2.0:n), n*randn(n-1), n*(n+1.0))

        @test issymmetric(A)

        B = SymArrow(ones(n-1), rand(n-1)./n^2, 1.0)

        @test issymmetric(B)

        R = chol(B)

        @test !issymmetric(R)

        @test norm(R'R-B)/norm(B) < n*eps()

        Ri = inv(R)

        @test !issymmetric(Ri)

        @test norm(Ri*R-I) < n*eps()

        SMA = Symmetric(Matrix(A))
        SMB = Symmetric(Matrix(B))

        λ = eigvals(A)
        Λ = eigvals(SMA)
        println("Testing eigenvalue error")
        println(norm(λ-Λ)/norm(Λ),"  ", 20*eps())
        @test norm(λ-Λ)/norm(Λ) < 20*eps()

        A1 = synthesize_symarrow(A, λ)
        Λ, Q = eig(A)

        println("Testing orthogonality")
        println(norm(Q'Q-I),"  ",20*eps())
        #@test norm(Q'Q-I) < 20*eps()

        println("Testing backward error")
        println(norm(A*Q-Q*Diagonal(Λ)),"  ",1.125*(norm(A.b-A1.b)+norm(A.c-A1.c)))
        @test norm(A*Q-Q*Diagonal(Λ)) < 1.125*(norm(A.b-A1.b)+norm(A.c-A1.c))

        λ = eigvals(A, B)
        Λ = eigvals(SMA, SMB)
        println("Testing generalized eigenvalue error")
        println(norm(λ-Λ)/norm(Λ),"  ", 20*eps())
        @test norm(λ-Λ)/norm(Λ) < 20*eps()
    end
end

@testset "Testing FMM acceleration" begin
    for n in (100, 300, 1000, 3000, 10000)
        println("n = ", n)
        A = SymArrow((1.0:n-1).*(2.0:n), n*ones(n-1), n*(n+1.0))
        if n ≤ 10_000
            @time λ = eigvals(A)
            @time λ_FMM = eigvals_FMM(A)

            println("Testing relative error")
            println(norm((λ-λ_FMM)./λ, Inf))
            @test norm((λ-λ_FMM)./λ, Inf) < 10*eps()
        else
            @time eigvals_FMM(A)
        end
    end
end


compute_R_star(rho_star, M_star) = cbrt(3.0 * M_star / (4.0 * π * rho_star))
compute_R_star_nom(G::Real) = 1.0
compute_R_star_nom(G) = 1.0u"Rsun"
compute_M_star(rho_star, R_star) = 4.0 * π * R_star^3 * rho_star / 3.0
compute_M_planet_nom(G::Real) = 0.0
compute_M_planet_nom(G) = 0.0u"Msun"
compute_M_tot(m1, m2) = m1 + m2
compute_M_tot(a, G, period) = 4.0 * π^2 * a^3 / (G * period^2)
compute_a(aR_star, R_star) = R_star * aR_star
compute_a(M_tot, period, G) = cbrt(G * M_tot * period^2 / (4.0 * π^2))
compute_a_X(a, m, M) = a * m / M
compute_period(M_tot, a, G) = 2.0 * π * sqrt(a^3 / (G * M_tot))
compute_rho_star(M_star, R_star) = 3.0 * M_star / (4.0 * π * R_star^3)
function compute_E0(ecc, cos_omega, sin_omega)
    y = sqrt(1.0 - ecc) * cos_omega
    x = sqrt(1.0 + ecc) * (1.0 + sin_omega)
    return 2.0 * atan(y, x)
end
compute_M0(ecc, E0) = E0 - ecc * sin(E0)
compute_M(t, t0, t_ref, n) = (t - t0 - t_ref) * n

# Spherical density
compute_rho(M, R) = 0.75 * M / (π * R^3)
compute_rho(M, R::Nothing) = nothing

# Semi-major axis / star radius ratio, assuming circular orbit
function compute_aor(duration, period, b; r=nothing)
    r = isnothing(r) ? 0.0 : r
    sin_ϕ, cos_ϕ = sincos(π * duration / period)
    return sqrt((1 + r)^2 - (b * cos_ϕ)^2) / sin_ϕ
end

# Impact radius
function compute_b(a_planet, R_star, duration, period, incl_factor_inv, ecc, sin_omega)
    c = sin(π * duration / (period * incl_factor_inv))
    c_sq = c^2
    ecc_sin_omega = ecc*sin_omega
    aor = a_planet / R_star
    num = aor^2 * c_sq - 1.0
    den = c_sq * ecc_sin_omega^2 + 2.0 * c_sq * ecc_sin_omega + c_sq - ecc^4 + 2.0 * ecc^2 - 1.0
    return sqrt(num / den) * (1.0 - ecc) * (1.0 + ecc)
end
compute_b(cos_incl, dcosi_db) = cos_incl / dcosi_db

# Inclination factor
compute_incl_factor_inv(ecc, sin_omega) = (1.0 - ecc) * (1.0 + ecc) / (1.0 + ecc * sin_omega)

# Jacobian for cos(i) -> b
compute_dcosi_db(a, R_star, incl_factor_inv) = R_star / (a * incl_factor_inv)

# Planet radius
compute_R_planet(R_star, r, R_planet) = R_planet
compute_R_planet(R_star, r, R_planet::Nothing) = iszero(r) ? zero(R_star) : R_star * r

# Transit times
compute_t0_tp(t0::Nothing, tp; M0, n) = (tp + M0 / n, tp)
compute_t0_tp(t0, tp::Nothing; M0, n) = (t0, t0 - M0 / n)
compute_t0_tp(t0::Nothing, tp::Nothing; kwargs...) = throw(ArgumentError("Please specify either `t0` or `tp`"))
compute_t0_tp(t0, tp; kwargs...) = throw(ArgumentError("Please only specify one of `t0` or `tp`"))


function weak_phs1(N,a,b)
	N1 = N+1   # variables x1, e1 and v1
	N2 = N   	#    variables x2, e2 and v2

	x1,w1,P = lglnodes(N1-1,a,b)    # discretization of x1 variables
	if N2 > 1
		x2,w2 = lgwt(N2,a,b)
	else
		x2,w2 = (a+b)/2, (b-a);
	end
	#M = massmatrix(ei,xi,we);
	#xquad, wquad = lgwt(Ne,a,b)\
	#M = massmatrix(ei,xi, xquad, wquad);
	M1 = massmatrix(x1,x1, w1);
	M2 = massmatrix([x2],[x2], [w2]);

	D = dermatrix(x1,x2,1)'*M2 ;
	p0 = map(i->leg_pol(a,x1,i), 1:length(x1))
	pL = map(i->leg_pol(b,x1,i), 1:length(x1))

    phJ = [zeros(N1,N1) D;
	   -D' zeros(N2, N2)];
    phB = [p0 -pL; zeros(N2,2)];
    phD = zeros(2,2);

	Q = blkdiag(inv(M1),inv(M2));
	#pflow = Collocation_points(x1, w1)
	#peffort = Collocation_points(x2, w2)

	p = Phs(phJ, phB, phD, Q)
	#p.Q = Q
	#p.disc_data = Discrete_domain(pflow, peffort)
	#func(x :: Array) = (0.5*x'*p.Q*x)[1]
	#p.Hamiltonian = func
	#p.GradHam = ForwardDiff.gradient(p.Hamiltonian)
	p
end

function weak_phs2(N,a,b)
	N1 = N   # variables x1, e1 and v1
	N2 = N+1   	#    variables x2, e2 and v2

	x2,w2,P = lglnodes(N2-1,a,b)    # discretization of x1 variables
	if N1 > 1
		x1,w1 = lgwt(N1,a,b)
	else
		x1,w1 = (a+b)/2, (b-a);
	end
	#M = massmatrix(ei,xi,we);
	#xquad, wquad = lgwt(Ne,a,b)\
	#M = massmatrix(ei,xi, xquad, wquad);
	M1 = massmatrix([x1],[x1], [w1]);
	M2 = massmatrix(x2,x2, w2);

	D = dermatrix(x2,x1,1)'*M1 ;
	p0 = map(i->leg_pol(a,x2,i), 1:length(x2))
	pL = map(i->leg_pol(b,x2,i), 1:length(x2))

    phJ = [zeros(N1,N1) -D';
	   D zeros(N2, N2)];
    phB = [zeros(N1,2); p0 -pL];
    phD = zeros(2,2);

		Q = blkdiag(inv(M1),inv(M2));

	#pflow = Collocation_points(x1, w1)
	#peffort = Collocation_points(x2, w2)

	p = Phs(phJ, phB, phD, Q)
	#p.Q = Q
	#p.disc_data = Discrete_domain(pflow, peffort)
	#func(x :: Array) = (0.5*x'*p.Q*x)[1]
	#p.Hamiltonian = func
	#p.GradHam = ForwardDiff.gradient(p.Hamiltonian)
	p
end

#coupled_gyrator(ph1 :: Phs, ports1, ph2 :: Phs, ports2, couple_matrix)
function finelem(Nelem,Nint,a,b)
	dx = (b-a)/Nelem;

	p = weak_phs1(Nint,0,dx)
	for i = 2:Nelem
		if (i%2 == 1)
			pn = weak_phs1(Nint,0,dx)
		else
			pn = weak_phs2(Nint,0,dx)
		end
		p = PortHamiltonian.coupled_gyrator(p,[2],pn,[1],[-1])
	end
	p
end


######### StanSample Bernoulli example  ###########

using StanSample

bernoulli_model = "
data {
  int<lower=1> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}
";

bernoulli_data = Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1])

# Keep tmpdir across multiple runs to prevent re-compilation
tmpdir = joinpath(@__DIR__, "tmp")

sm = SampleModel("bernoulli", bernoulli_model;
  method = StanSample.Sample(
    save_warmup=false,                           # Default
    thin=1,
    adapt = StanSample.Adapt(delta = 0.85)),
  tmpdir = tmpdir
);

rc = stan_sample(sm; data=bernoulli_data);

if success(rc)
  (samples, cnames) = read_samples(sm, output_format=:array)
end


quadratic_bowl_3D_lc = environment_3D(x -> transpose(x[1:2]) * x[1:2], cone = LinearizedCone)
quadratic_bowl_3D_nc = environment_3D(x -> transpose(x[1:2]) * x[1:2], cone = NonlinearCone)

circular_bowl_3D_nc = environment_3D(x -> -sqrt(2.5^2.0 - x[1]^2.0 - x[2]^2.0) + 2.5, cone = NonlinearCone)

# r = 2.0
# t = range(-r, stop = r, length = 1000)
# z = -sqrt.(r^2.0 .- t.^2.0 .- 0^2.0) .+ r
#
# plot(t, z, aspect_ratio = :equal)
# env = deepcopy(circular_bowl_3D_nc)
# env = deepcopy(flat_3D_nc)
# q = [-0.5; -0.5; env.surf([-0.5; -0.5])]
# n = [env.surf_grad(q[1:2]); -1.0]
# ns = n ./ sqrt(transpose(n) * n) #* model.r


using Flux, Zygote, StochasticDiffEq
using LinearAlgebra, Statistics
println("heat neural SDE")
using Test, NeuralPDE

using Random
Random.seed!(100)

println("high-dimensional heat equation")
d = 100 # number of dimensions
x0 = fill(8.0f0,d)
tspan = (0.0f0,2.0f0)
dt = 0.5
m = 50 # number of trajectories (batch size)

g(X) = sum(X.^2)
f(X,u,σᵀ∇u,p,t) = Float32(0.0)
μ_f(X,p,t) = zero(X) #Vector d x 1
σ_f(X,p,t) = Diagonal(ones(Float32,d)) #Matrix d x d
prob = TerminalPDEProblem(g, f, μ_f, σ_f, x0, tspan)

hls = 10 + d #hidden layer size
#sub-neural network approximating solutions at the desired point
u0 = Flux.Chain(Dense(d,hls,relu),
                Dense(hls,hls,relu),
                Dense(hls,1))
# sub-neural network approximating the spatial gradients at time point
σᵀ∇u = Flux.Chain(Dense(d+1,hls,relu),
                  Dense(hls,hls,relu),
                  Dense(hls,d))
opt = Flux.ADAM(0.005)  #optimizer
pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)

ans = solve(prob, pdealg, verbose=true, maxiters=500, trajectories=m,
                            alg=EM(), dt=dt, pabstol = 1f-6)

u_analytical(x,t) = sum(x.^2) .+ d*t
analytical_ans = u_analytical(x0, tspan[end])


# https://fenicsproject.org/docs/dolfin/latest/python/demos/biharmonic/demo_biharmonic.py.html

module demo_biharmonic

using FenicsPy

# Optimization options for the form compiler
parameters["form_compiler"]["cpp_optimize"] = true
parameters["form_compiler"]["optimize"] = true

# Make mesh ghosted for evaluation of DG terms
parameters["ghost_mode"] = "shared_facet"

# Create mesh and define function space
mesh = UnitSquareMesh(32, 32)
V = FunctionSpace(mesh, "CG", 2)

# Define Dirichlet boundary
@pydef mutable struct DirichletBoundary <: dolfin.SubDomain
    function inside(self, x, on_boundary)
        return on_boundary
    end
end

py"""
from math import sin, pi
from dolfin import UserExpression
class Source(UserExpression):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    def eval(self, values, x):
        values[0] = 4.0*pi**4*sin(pi*x[0])*sin(pi*x[1])
    def value_shape(self):
        return ()
"""

# Define boundary condition
u0 = Constant(0.0)
bc = DirichletBC(V, u0, SubDomain(DirichletBoundary()))

# Define trial and test functions
u = TrialFunction(V)
v = TestFunction(V)

# Define normal component, mesh size and right-hand side
h = CellDiameter(mesh)
h_avg = (h("+") + h("-"))/2.0
n = FacetNormal(mesh)
f = Expression(py"Source"(degree=2))

# Penalty parameter
alpha = Constant(8.0)

# Define bilinear form
a = inner(div(grad(u)), div(grad(v)))*dx -
   inner(avg(div(grad(u))), jump(grad(v), n))*dS -
   inner(jump(grad(u), n), avg(div(grad(v))))*dS +
   alpha/h_avg*inner(jump(grad(u),n), jump(grad(v),n))*dS

# Define linear form
L = f*v*dx

# Solve variational problem
u = FeFunction(V)
solve(a == L, u, bc)

# Save solution to file
file = File("biharmonic/biharmonic.pvd")
file << u

# Plot solution
#plot(u)

end # module demo_biharmonic


# p3.jl - band-limited interpolation

h = 1; xmax = 10; clf();
x = -xmax:h:xmax;                     # computational grid
xx = -xmax-h/20:h/10:xmax+h/20;       # plotting grid
v = zeros(length(x),3);
v[:,1] = @. float(x==0);
v[:,2] = @. float(abs(x) ≤ 3);
v[:,3] = @. max(0,1-abs(x)/3);
for plt = 1:3
    subplot(4,1,plt);
    plot(x,v[:,plt],".",markersize=6);
    p = 0;
    for i = 1:length(x)
        p = @. p + v[i,plt]*sin(pi*(xx-x[i])/h)/(pi*(xx-x[i])/h);
    end
    plot(xx,p,"-")
    axis([-xmax,xmax,-0.5,1.5]);
    xticks(1:0); yticks(0:1);
end


using OmegaGrav
using MatterPower
using Dierckx
using DelimitedFiles, Printf

# %% Cosmological parameters
Ωm = 0.315

# %% Specify a redshift
redshift = [0, 0.3, 0.5, 0.7, 1, 1.3, 1.5]
ired = 1
z = redshift[ired]

# %% Read in linear and non-linear power spectra computed by CLASS
filename = @sprintf("data/matterpower05_z%1d_pk.dat", ired)
d = readdlm(filename, comments = true)
pk_spl = Spline1D(d[:, 1], d[:, 2])
filename = @sprintf("data/matterpower05_z%1d_pk_nl.dat", ired)
d = readdlm(filename, comments = true)
pknl_spl = Spline1D(d[:, 1], d[:, 2])

# %% Compute `Ωgrav = Ωm * W / 2`
println("redshift = ", z)
# Linear matter power spectrum
Ωgrav_pklin = ograv_pk(x -> pk_spl(x), z, Ωm)
println("Ωgrav(linear pk) = ", Ωgrav_pklin)
# Non-linear matter power spectrum
Ωgrav_pknl = ograv_pk(x -> pknl_spl(x), z, Ωm)
println("Ωgrav(non-linear pk) = ", Ωgrav_pknl)


using Turing
using Test
using Random
using Distributions

Random.seed!(128)

alg1 = HMCDA(3000, 1000, 0.65, 0.015)
# alg2 = Gibbs(3000, HMCDA(1, 200, 0.65, 0.35, :m), HMC(1, 0.25, 3, :s))
alg3 = Gibbs(1500, PG(30, 10, :s), HMCDA(1, 500, 0.65, 0.005, :m))
# alg3 = Gibbs(2000, HMC(1, 0.25, 3, :m), PG(30, 3, :s))
# alg3 = PG(50, 2000)

@model gdemo(x) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt.(s))
  x[1] ~ Normal(m, sqrt.(s))
  x[2] ~ Normal(m, sqrt.(s))
  return s, m
end

res1 = sample(gdemo([1.5, 2.0]), alg1)

check_numerical(res1, [:s, :m], [49/24, 7/6])

# res2 = sample(gdemo([1.5, 2.0]), alg2)
#
# @test mean(res2[:s]) ≈ 49/24 atol=0.2
# @test mean(res2[:m]) ≈ 7/6 atol=0.2

res3 = sample(gdemo([1.5, 2.0]), alg3)

check_numerical(res3, [:s, :m], [49/24, 7/6])

# res1 = sample(gdemo([1.5, 2.0]), HMC(3000, 0.2, 4))
# println("HMC")
# println("E[s] = $(mean(res1[:s]))")
# println("E[m] = $(mean(res1[:m]))")


infile = "./src/Day14/input.txt"
input = readlines(infile)

function part1(input)
    sp = split.(input, " = ")
    instructions = first.(sp)
    data = last.(sp)
    ln = maximum(length.(data)) # 36
    mask = "X"^ln
    memory = Dict{Int,Int}()
    for i in 1:length(sp)
        if instructions[i] == "mask"
            mask = data[i]
        else
            pointer = parse(Int, match(r"\d+", instructions[i]).match)
            towrite = applymask(mask, parse(Int, data[i]))
            memory[pointer] = towrite
        end
    end
    return memory
end

function applymask(mask, data)
    # convert data to binary
    d = reverse(digits(data, base = 2, pad = 36))
    for (i,v) in enumerate(mask)
        if v != 'X'
            d[i] = parse(Int, v)
        end
    end
    # convert binary to int
    val = parse(Int, join(d,""), base = 2)
    return val
end

sp = split.(input, " = ")
instructions = first.(sp)
data = last.(sp)

function part2(instructions, data)
    mask = "X"^36
    memory = Dict{Int,Int}()
    for i in 1:length(instructions)
        if instructions[i] == "mask"
            mask = data[i]
        else
            pointer = parse(Int, match(r"\d+", instructions[i]).match)
            allpointers = applymemmask(mask, pointer)
            d = parse(Int, data[i])
            for p in allpointers
                memory[p] = d
            end
        end
    end
    return memory
end

function applymemmask(mask, pointer)
    allpointers = Int[]
    push!(allpointers, 0)
    l = length(mask)
    pbin = reverse(digits(pointer, base = 2, pad = l))
    for (i, v) in enumerate(mask)
        size = 2^(l - i)
        if v == 'X'
            # double up the entries
            added = allpointers .+ size
            push!(allpointers, added...)
        elseif v == '1'
            allpointers .+= size
        elseif v == '0'
            p = pbin[i]
            if p == 1
                allpointers .+= size
            else
                # nothing
            end
        end
    end
    return allpointers
end

applymemmask("000000000000000000000000000000X1001X", 42)

mem2 = part2(instructions, data)
sum(values(mem2))

using BenchmarkTools
@benchmark part2(instructions, data)


using Test, JLD2, UnicodePlots
using LinearAlgebra, ForwardDiff, Arpack
using SparseArrays, SparseDiffTools, Symbolics

# load mesh
meshfile = "$(@__DIR__)/data/squaremesh.jld2"
JLD2.@load(meshfile, p, t, ps, ts)
#p, t =  ps, ts

np, nt = size(p, 1), size(t, 1)
ε = 1e-4
x, dp = p, 2 * rand(size(p)...) .- 1
Ib = unique(FemAdjoint.btri(t)[:])
dp[Ib, :] .= 0.# fix Dirichlet points
xp, xm = p .+ ε * dp, p .- ε * dp

# finite differences
f = x -> FemAdjoint.costnormU(x, t)
fp, fm = f(xp), f(xm)
@time g = FemAdjoint.∇costnormU_direct(x, t, plotsolutions = true)
@time g = FemAdjoint.∇costnormU_direct(x, t, plotsolutions = false)
println("")
println(fp)
println(fm)
println(" ")
println((fp - fm) / (2 * ε))
println(dot(g, dp))
println("")
@testset "test eval direct gradient dot function for costnormU" begin
    @test abs(dot(g, dp) - (fp - fm) / (2 * ε)) < 1e-4
end


using ForwardDiff
using Distributions

# Real

x_real = randn(5)

dists = [Normal(0, 1)]

for dist in dists

    f(x::Vector) = sum(logpdf.(Ref(dist), x))

    ForwardDiff.gradient(f, x_real)

end

# Postive

x_positive = randn(5).^2

dists = [Gamma(2, 3)]

for dist in dists

    f(x::Vector) = sum(logpdf.(Ref(dist), x))

    g = x -> ForwardDiff.gradient(f, x)

end


## Creation

pseudo_ideal(::Type{R},m::RingInt) where R = R(m)
pseudo_ideal(::Type{R},mm::AbstractVector) where R = Ideal(Vector{R}(mm))
pseudo_ideal(::Type{R},m::Ideal{R}) where R = m

function Ideal(m::R) where R<:Ring
    if iszero(m)
        Ideal{R}(R[])
    elseif isunit(m)
        Ideal{R}([one(R)])
    else
        Ideal{R}([m])
    end
end
Ideal(m::T) where T<:Integer = Ideal(ZZ{T}(m))
Ideal(mm::RingInt...) = Ideal(collect(Base.promote_typeof(mm...), mm))
Ideal(mm::AbstractVector{<:RingInt}) = Ideal(gcd(mm))
Ideal(mm::AbstractVector{R}) where R<:MultivariatePolynomial = Ideal{R}(groebnerbase(mm))

# Arithmetic

import Base: in, issubset, ==, intersect, sum, +, *, ^, iszero, zero, one, rem

function in(a::R, id::Ideal{R}) where R<:MultivariatePolynomial
    a, r, d = pdivrem(a, id.base)
    isone(d) && iszero(r)
end

function issubset(id1::Ideal{R}, id2::Ideal{R}) where R<:Polynomial
    for b1 in id1.base
        in(b1, id2) || return false
    end
    true
end

==(id1::Ideal{R}, id2::Ideal{S}) where {R<:Ring,S<:Ring} = id1.base == id2.base
==(id::Ideal{R}, ::Type{R}) where R<:Ring = isone(id)
==(::Type{R}, id::Ideal{R}) where R<:Ring = isone(id)

function +(id1::Ideal{R}, id2::Ideal{R}) where R<:MultivariatePolynomial
    Ideal(vcat(id1.base, id2.base))
end
+(id::Ideal{R}, a::R) where R = Ideal([id.base; a])
+(a::R, id::Ideal{R}) where R = Ideal([a; id.base])

function *(id1::Ideal{R}, id2::Ideal{R}) where R<:MultivariatePolynomial
    C = [ b1 * b2 for b1 in id1.base for b2 in id2.base]
    Ideal(C)
end

*(id::Ideal{R}, a::R) where R = Ideal(id.base .* a)
*(a::R, id::Ideal{R}) where R = Ideal(a .* id.base)

function ^(id::Ideal{R}, p::Integer) where R<:MultivariatePolynomial
    if p == 1
        id
    elseif p == 0
        Ideal([one(R)])
    elseif p == 2
        id * id
    elseif iseven(p)
        (id^(p÷2))^2
    else
        (id^((p-1)÷2))^2 * id
    end
end

# remainder when dividing by an ideal
rem(a::R, id::Ideal{R}) where R<:Ring = divrem(a, id)[2]

function intersect(id1::Ideal{R}, id2::Ideal{R}) where R<:MultivariatePolynomial
    Ideal(_intersect(id1.base, id2.base))
end

function _intersect(v1::V, v2::V) where {R<:MultivariatePolynomial,V<:AbstractVector{R}}
    t = Symbol("intersect_extension")
    Rt = lextend(R, t)
    t, = generators(Rt)
    id3 = Ideal([Rt.(v1) .* t; Rt.(v2) .* (1 - t)])
    R.([x for x in id3.base if multideg(x)[1] == 0])
end

iszero(id::Ideal) = length(id.base) == 0
isone(id::Ideal) = length(id.base) == 1 && isone(id.base[1])
zero(::Type{I}) where {R,I<:Ideal{R}} = Ideal(R[])
one(::Type{I}) where {R,I<:Ideal{R}} = Ideal(R[1])
one(::R) where R<:Ideal = one(R)
zero(::R) where R<:Ideal = zero(R)


"""Compute the KL Divergence between the GP Prior and the variational distribution"""
GaussianKL(model::AbstractGP) = sum(broadcast(GaussianKL,model.f,get_Z(model)))

GaussianKL(gp::Abstract_GP,X::AbstractMatrix) = GaussianKL(gp.μ,gp.μ₀(X),gp.Σ,gp.K)

function GaussianKL(μ::AbstractVector{T},μ₀::AbstractVector,Σ::Matrix{T},K::PDMat{T,Matrix{T}}) where {T<:Real}
    0.5*(-logdet(Σ)+logdet(K)+tr(K\Σ)+invquad(K,μ-μ₀)-length(μ))
end

InverseGammaKL(α,β,αₚ,βₚ) = GammaKL(α,β,αₚ,βₚ)
"""
    KL(q(ω)||p(ω)), where q(ω) = Ga(α,β) and p(ω) = Ga(αₚ,βₚ)
"""
function GammaKL(α,β,αₚ,βₚ)
    sum((α-αₚ).*digamma(α) .- log.(gamma.(α)).+log.(gamma.(αₚ)) .+  αₚ.*(log.(β).-log.(βₚ)).+α.*(βₚ.-β)./β)
end

"""
    KL(q(ω)||p(ω)), where q(ω) = Po(ω|λ) and p(ω) = Po(ω|λ₀)
"""
function PoissonKL(λ::AbstractVector{T},λ₀::Real) where {T}
    λ₀*length(λ)-(one(T)+log(λ₀))*sum(λ)+sum(xlogx,λ)
end

"""
    KL(q(ω)||p(ω)), where q(ω) = Po(ω|λ) and p(ω) = Po(ω|λ₀) with ψ = E[log(λ₀)]
"""
function PoissonKL(λ::AbstractVector{<:Real},λ₀::AbstractVector{<:Real},ψ::AbstractVector{<:Real})
    sum(λ₀)-sum(λ)+sum(xlogx,λ)-dot(λ,ψ)
end

"""KL(q(ω)||p(ω)), where q(ω) = PG(b,c) and p(ω) = PG(b,0). θ = 𝑬[ω]"""
function PolyaGammaKL(b,c,θ)
    dot(b,logcosh.(0.5*c))-0.5*dot(abs2.(c),θ)
end

"""
    Entropy of GIG variables with parameters a,b and p and omitting the derivative d/dpK_p cf <https://en.wikipedia.org/wiki/Generalized_inverse_Gaussian_distribution#Entropy>
"""
function GIGEntropy(a,b,p)
    sqrtab = sqrt.(a.*b)
    return sum(0.5*log.(a./b))+sum(log.(2*besselk.(p,sqrtab)))+ sum(0.5*sqrtab./besselk.(p,sqrtab).*(besselk.(p+1,sqrtab)+besselk.(p-1,sqrtab)))
end


# Analyze the data using Lasso Path pacakge I wrote.
cd("C:\\Users\\victo\\source\\repos\\Silly_Julia_stuff\\Julia Packages\\LassoPath\\usage_examples\\crime_data")
using DelimitedFiles, CSV, DataFrames
using LassoPath
using BenchmarkTools

Df = CSV.read("process_data.csv", DataFrame)
Df = Df[!, Not([:state, :fold])]

# Extrac Numerics, predictors, predictants.
A, b = Matrix(Df[:, Not(end)]), Df[:, end]

PrintTitle("The Data is ready, we are ready for analyais now")

function ShowResults(constructor::Function, fname::String)
    lassoIntance = constructor(A, b)
    lassoIntance.Tol = -1
    GetLassoPath(lassoIntance, 0.99)
    VisualizeLassoPath(lassoIntance, fname=fname)
    predictorIndices, λ, ws =
        CaptureImportantWeights(lassoIntance, 10, 1e-2)
    PrintTitle("Important Predictors are")
    predictors = names(Df[:, Not(end)])[predictorIndices]
    impacts = map(ws) do x
        if x >= 0
            return "+"
        else
            return "-"
        end
    end
    display([reshape(predictors, :) reshape(impacts, :)])
end

# @time ShowResults((x, y) -> LassoSCOP(x, y), "Lasso SCOP.png")
@time ShowResults((x, y) -> Lasso2NormProx(x, y), "lasso_proximal.png")


#=
    A simple conic problem using the second-order cone.

    The model is defined as
    min ‖u - u0‖₂
    st  p'u = q

    The conic formulation we will solve is given by
    min t
    st  p'u = q
        ‖u- u0‖₂ ≤ t  (second-order cone)

    Since second-order programming is a special case of SDP,
    it can be reformulated accordingly. This example illustrates
    the reformulation capabilities of JuMP/MathOptInterface.
    Indeed, second-order cones were not implemented in MySDPSolver,
    but MOI can handle them.
=#

using MySDPSolver, JuMP, LinearAlgebra

# Some data
u0 = ones(Float64, 10)
p = ones(Float64, 10)
q = 1.0

# Declare the model
model = Model(MySDPSolver.Optimizer)

# Add the variables and the objective
@variable(model, u[1:10])
@variable(model, t)
@objective(model, Min, t)

# Add the conic constraint, i.e. ‖u-u0‖₂ ≤ t
@constraint(model, c1, [t, (u - u0)...] in SecondOrderCone())  # see 'splatting' for the '...' syntax
@constraint(model, c2, u'*p == q)

# Call the solver
optimize!(model)

# Retrieve the solutions
println("t:")
display(value(t))
println("u: ")
display(value.(u))

# We can get the dual solutions as well
println("dual solutions:")
display(dual(c1))
display(dual(c2))

println("###############################################")


using LazySets

S = Singleton([0.0, 0.0])
B1 = BallInf([0.0, 0.0], 1.0)
B2 = BallInf([1.0, 0.0], 0.5)
S_proj = project(S, [1])
B1_proj = project(B1, [1])
B2_proj = project(B2, [1])

P1 = is_contained_in(B1)
@test P1(S) && P1(B1) && !P1(B2)

P2 = contains(S)
@test P2(S) && P2(B1) && !P2(B2)

# dim is only available for SetAtom types
A1 = Atom(x -> x > 1)
@test_throws ArgumentError dim(A1)

# project is only available for SetAtom (and their Negation) types
@test_throws ArgumentError project(A1, [1])

@test project(Negation(P1), [1]) == Negation(project(P1, [1]))

@test dim(P1) == dim(P2) == 2
P1_proj = project(P1, [1])
P2_proj = project(P2, [1])
@test P1_proj isa SetAtom && dim(P1_proj) == 1 && P1_proj.X == B1_proj
@test P1_proj(S_proj) && P1_proj(B1_proj) && !P1_proj(B2_proj)
@test P2_proj isa SetAtom && dim(P2_proj) == 1 && P2_proj.X == S_proj
@test P2_proj(S_proj) && P2_proj(B1_proj) && !P2_proj(B2_proj)

Pc = Conjunction([P1, P2])
@test Pc(S) && Pc(B1) && !Pc(B2)

Pd = Disjunction([P1, P2])
@test Pd(S) && Pd(B1) && !Pd(B2)

@test project(Pc, [1]) == Conjunction([P1_proj, P2_proj])
@test project(Pd, [1]) == Disjunction([P1_proj, P2_proj])

P1 = is_disjoint_from(S)
P2 = intersects(S)
@test !P1(S) && !P1(B1) && P1(B2)
@test P2(S) && P2(B1) && !P2(B2)


#
# TODO
#
FLOAT_MOD(s::State) = return

FLOAT_MUL(s::State) = if length(s.float) >= 2
  push!(s.float, pop!(s.float) * pop!(s.float))
end

FLOAT_PLUS(s::State) = if length(s.float) >= 2
  push!(s.float, pop!(s.float) + pop!(s.float))
end

FLOAT_SUB(s::State) = if length(s.float) >= 2
  push!(s.float, - pop!(s.float) + pop!(s.float))
end

FLOAT_DIV(s::State) = if length(s.float) >= 2
  divisor = peek(s.float)
  if divisor != zero(Float32)
    pop!(s.float)
    push!(s.float, pop!(s.float) / divisor)
  end
end

FLOAT_LT(s::State) = if length(s.float) >= 2
  push!(s.boolean, pop!(s.float) > pop!(s.float))
end

FLOAT_EQ(s::State) = if length(s.float) >= 2
  push!(s.boolean, pop!(s.float) == pop!(s.float))
end

FLOAT_GT(s::State) = if length(s.float) >= 2
  push!(s.boolean, pop!(s.float) < pop!(s.float))
end

FLOAT_COS(s::State) = if !isempty(s.float)
  push!(s.float, cos(pop!(s.float)))
end

FLOAT_DEFINE(s::State) = if !isempty(s.float) && !isempty(s.name)
  s.definitions[pop!(s.name)] = pop!(s.float)
end

FLOAT_DUP(s::State) = if !isempty(s.float)
  push!(s.float, peek(s.float))
end

FLOAT_FLUSH(s::State) = empty!(s.float)

FLOAT_FROM_BOOLEAN(s::State) = if !isempty(s.boolean)
  push!(s.float, pop!(s.boolean) ? one(Float32) : zero(Float32))
end

FLOAT_FROM_INTEGER(s::State) = if !isempty(s.integer)
  push!(s.float, convert(Float32, pop!(s.integer)))
end

FLOAT_MAX(s::State) = if length(s.float) >= 2
  push!(s.float, max(pop!(s.float), pop!(s.float)))
end

FLOAT_MIN(s::State) = if length(s.float) >= 2
  push!(s.float, min(pop!(s.float), pop!(s.float)))
end

FLOAT_POP(s::State) = if !isempty(s.float)
  pop!(s.float)
end

FLOAT_RAND(s::State) = push!(s.float, random_float(s))

FLOAT_ROT(s::State) = if length(s.float) > 2
  s.float[end], s.float[end-1], s.float[end-2] =
    s.float[end-1], s.float[end-2], s.float[end]
end

FLOAT_SHOVE(s::State) = if !isempty(s.integer) && !isempty(s.float)
  shove!(s.float, pop!(s.integer))
end

FLOAT_SIN(s::State) = if !isempty(s.float)
  push!(s.float, sin(pop!(s.float)))
end

FLOAT_STACK_DEPTH(s::State) = push!(s.integer, length(s.float))

FLOAT_SWAP(s::State) = if length(s.float) >= 2
  s.float[end], s.float[end-1] = s.float[end-1], s.float[end]
end

FLOAT_TAN(s::State) = if !isempty(s.float)
  push!(s.float, tan(pop!(s.float)))
end

FLOAT_YANK(s::State) = if !isempty(s.integer) && !isempty(s.float)
  yank!(s.float, pop!(s.integer))
end

FLOAT_YANKDUP(s::State) = if !isempty(s.integer) && !isempty(s.float)
  yankdup!(s.float, pop!(s.integer))
end

FLOAT_ERC(s::State) = return

Push.register("FLOAT.%",           FLOAT_MOD)
Push.register("FLOAT.*",           FLOAT_MUL)
Push.register("FLOAT.+",           FLOAT_PLUS)
Push.register("FLOAT.-",           FLOAT_SUB),
Push.register("FLOAT./",           FLOAT_DIV)
Push.register("FLOAT.<",           FLOAT_LT)
Push.register("FLOAT.=",           FLOAT_EQ)
Push.register("FLOAT.>",           FLOAT_GT)
Push.register("FLOAT.COS",         FLOAT_COS)
Push.register("FLOAT.DEFINE",      FLOAT_DEFINE)
Push.register("FLOAT.DUP",         FLOAT_DUP)
Push.register("FLOAT.FLUSH",       FLOAT_FLUSH)
Push.register("FLOAT.FROMBOOLEAN", FLOAT_FROM_BOOLEAN)
Push.register("FLOAT.FROMINTEGER", FLOAT_FROM_INTEGER)
Push.register("FLOAT.MAX",         FLOAT_MAX)
Push.register("FLOAT.MIN",         FLOAT_MIN)
Push.register("FLOAT.POP",         FLOAT_POP)
Push.register("FLOAT.RAND",        FLOAT_RAND)
Push.register("FLOAT.ROT",         FLOAT_ROT)
Push.register("FLOAT.SHOVE",       FLOAT_SHOVE)
Push.register("FLOAT.SIN",         FLOAT_SIN)
Push.register("FLOAT.STACKDEPTH",  FLOAT_STACK_DEPTH)
Push.register("FLOAT.SWAP",        FLOAT_SWAP)
Push.register("FLOAT.TAN",         FLOAT_TAN)
Push.register("FLOAT.YANK",        FLOAT_YANK)
Push.register("FLOAT.YANKDUP",     FLOAT_YANKDUP)
Push.register("FLOAT.ERC",         FLOAT_ERC)


const fp = (1e9, 64.0, 3.46e-8, 2.0e-3, 0.7, 0.02, 1.0)

"""
    DAMMfdata(n)

Generates a DataFrame of n fake data Tₛ, θ and Rₛ

# Examples
```julia-repl
julia> DAMMfdata(5)
5×3 DataFrame
 Row │ Tₛ       θ        Rₛ
     │ Float64  Float64  Float64
─────┼───────────────────────────
   1 │    10.8      0.3  2.04327
   2 │    31.5      0.1  7.8925
   3 │    38.7      0.7  1.6
   4 │    35.7      0.3  7.38025
   5 │    21.9      0.2  3.0012
"""
function DAMMfdata(n)
  Tₛ = rand(0:400, n)./10
  θ = rand(1:7, n)./10
  noise = rand(-20:20, n)./10
  Rₛ = DAMM(hcat(Tₛ, θ), fp) .+ noise
  Rₛ[Rₛ .< 0] .= 0
  df = DataFrame((Tₛ = Tₛ, θ = θ, Rₛ = Rₛ))
  return df
end


# # Spherical harmonic addition theorem
# This example confirms numerically that
# ```math
# f(z) = \frac{P_n(z\cdot y) - P_n(x\cdot y)}{z\cdot y - x\cdot y},
# ```
# is actually a degree-$(n-1)$ polynomial on $\mathbb{S}^2$, where $P_n$ is the degree-$n$
# Legendre polynomial, and $x,y,z \in \mathbb{S}^2$.
# To verify, we sample the function on a $N\times M$ equiangular grid
# defined by:
# ```math
# \begin{aligned}
# \theta_n & = (n+\tfrac{1}{2})\pi/N,\quad{\rm for}\quad 0\le n < N,\quad{\rm and}\\
# \varphi_m & = 2\pi m/M,\quad{\rm for}\quad 0\le m < M;
# \end{aligned}
# ```
# we convert the function samples to Fourier coefficients using
# `plan_sph_analysis`; and finally, we transform
# the Fourier coefficients to spherical harmonic coefficients using
# `plan_sph2fourier`.
#
# In the basis of spherical harmonics, it is plain to see the
# addition theorem in action, since $P_n(x\cdot y)$ should only consist of
# exact-degree-$n$ harmonics.
#
# For the storage pattern of the arrays, please consult the
# [documentation](https://MikaelSlevinsky.github.io/FastTransforms).

function threshold!(A::AbstractArray, ϵ)
    for i in eachindex(A)
        if abs(A[i]) < ϵ A[i] = 0 end
    end
    A
end

using FastTransforms, LinearAlgebra, Plots
const GENFIGS = joinpath(pkgdir(FastTransforms), "docs/src/generated")
!isdir(GENFIGS) && mkdir(GENFIGS)
plotlyjs()

# The colatitudinal grid (mod $\pi$):
N = 15
θ = (0.5:N-0.5)/N

# The longitudinal grid (mod $\pi$):
M = 2*N-1
φ = (0:M-1)*2/M

# Arbitrarily, we place $x$ at the North pole:
x = [0,0,1]

# Another vector is completely free:
y = normalize([.123,.456,.789])

# Thus $z \in \mathbb{S}^2$ is our variable vector, parameterized in spherical coordinates:
z = (θ,φ) -> [sinpi(θ)*cospi(φ), sinpi(θ)*sinpi(φ), cospi(θ)]

# On the tensor product grid, the Legendre polynomial $P_n(z\cdot y)$ is:
A = [(2k+1)/(k+1) for k in 0:N-1]
B = zeros(N)
C = [k/(k+1) for k in 0:N]
c = zeros(N); c[N] = 1
pts = vec([z(θ, φ)⋅y for θ in θ, φ in φ])
phi0 = ones(N*M)
F = reshape(FastTransforms.clenshaw!(c, A, B, C, pts, phi0, zeros(N*M)), N, M)

# We superpose a surface plot of $f$ on top of the grid:
X = [sinpi(θ)*cospi(φ) for θ in θ, φ in φ]
Y = [sinpi(θ)*sinpi(φ) for θ in θ, φ in φ]
Z = [cospi(θ) for θ in θ, φ in φ]
scatter3d(vec(X), vec(Y), vec(Z); markersize=1.25, markercolor=:violetred)
surface!(X, Y, Z; surfacecolor=F, legend=false, xlabel="x", ylabel="y", zlabel="f")
savefig(joinpath(GENFIGS, "sphere1.html"))
###```@raw html
###<object type="text/html" data="../sphere1.html" style="width:100%;height:400px;"></object>
###```

# We show the cut in the surface to help illustrate the definition of the grid.
# In particular, we do not sample the poles.
#
# We precompute a spherical harmonic--Fourier plan:
P = plan_sph2fourier(F)

# And an FFTW Fourier analysis plan on $\mathbb{S}^2$:
PA = plan_sph_analysis(F)

# Its spherical harmonic coefficients demonstrate that it is exact-degree-$n$:
V = PA*F
U = threshold!(P\V, 400*eps())

# The $L^2(\mathbb{S}^2)$ norm of the function is:
nrm1 = norm(U)

# Similarly, on the tensor product grid, our function samples are:
Pnxy = FastTransforms.clenshaw!(c, A, B, C, [x⋅y], [1.0], [0.0])[1]
F = [(F[n, m] - Pnxy)/(z(θ[n], φ[m])⋅y - x⋅y) for n in 1:N, m in 1:M]

# We superpose a surface plot of $f$ on top of the grid:
scatter3d(vec(X), vec(Y), vec(Z); markersize=1.25, markercolor=:violetred)
surface!(X, Y, Z; surfacecolor=F, legend=false, xlabel="x", ylabel="y", zlabel="f")
savefig(joinpath(GENFIGS, "sphere2.html"))
###```@raw html
###<object type="text/html" data="../sphere2.html" style="width:100%;height:400px;"></object>
###```

# Its spherical harmonic coefficients demonstrate that it is degree-$(n-1)$:
V = PA*F
U = threshold!(P\V, 400*eps())

# Finally, the Legendre polynomial $P_n(z\cdot x)$ is aligned with the grid:
pts = vec([z(θ, φ)⋅x for θ in θ, φ in φ])
F = reshape(FastTransforms.clenshaw!(c, A, B, C, pts, phi0, zeros(N*M)), N, M)

# We superpose a surface plot of $f$ on top of the grid:
scatter3d(vec(X), vec(Y), vec(Z); markersize=1.25, markercolor=:violetred)
surface!(X, Y, Z; surfacecolor=F, legend=false, xlabel="x", ylabel="y", zlabel="f")
savefig(joinpath(GENFIGS, "sphere3.html"))
###```@raw html
###<object type="text/html" data="../sphere3.html" style="width:100%;height:400px;"></object>
###```

# It only has one nonnegligible spherical harmonic coefficient.
# Can you spot it?
V = PA*F
U = threshold!(P\V, 400*eps())

# That nonnegligible coefficient should be
ret = eval("√(2π/($(N-1)+1/2))")

# which is approximately
eval(Meta.parse(ret))

# since the convention in this library is to orthonormalize.
nrm2 = norm(U)

# Note that the integrals of both functions $P_n(z\cdot y)$ and $P_n(z\cdot x)$ and their
# $L^2(\mathbb{S}^2)$ norms are the same because of rotational invariance. The integral of
# either is perhaps not interesting as it is mathematically zero, but the norms
# of either should be approximately the same.
nrm1 ≈ nrm2


using DifferentialEquations
using LinearAlgebra
using Plots
using ForwardDiff
using ProgressMeter
using DSP

N = 4
A = zeros(N,N)

xs = 1.0:100.0
ys = 1.0:100.0
x = zeros(length(xs),length(ys))

stencil = [0.0  1.0 0.0;
           1.0 -4.0 1.0;
           0.0  1.0 0.0]

# generate_u_prime(x, dx)

# function generate1DStencil!(A, N, Δx)
#     # Generate the centered difference matrix manually
#     for i in 1:N, j in 1:N

#     if abs(i-j)<=1
#         (A[i,j]+=1)
#     end
#     if i==j
#         (A[i,j]-=3)
#     end

#     end

#     A = A/(Δx^2)
# end

function double_partial(x)
    dx = zeros(size(x))
    for i in 2:length(dx)-1
      dx[i] = 1*x[i-1] - 2*x[i] + 1*x[i+1]
    end
    dx[1] = -(2*x[1] + -1*x[2])
    dx[end] = -(-1*x[end-1] + 2*x[end])
end

function double_partial_2D(x)
    padded_du = conv2(x, stencil)
    du = padded_du[2:end-1,2:end-1]
    du
end

f(u, p, t) = gen_forcing_2D(p)

function gen_forcing_2D(positions)
    out_vec = create_2D_gaussian_matrix(xs,ys,positions)
    out_vec
end

function heat_eq(u, p, t)
    du = double_partial_2D(u) + f(u,p,t)
    du
end

function gaussian(x, mu, sig)
    return exp(-(x - mu)^2.0 / (2 * sig^2.0))
end

function gaussian_2D(loc, mu, sig)
    x, y = loc
    x₀, y₀ = mu
    σ₁, σ₂ = sig
    return gaussian(x, x₀, σ₁) * gaussian(y, y₀, σ₂)
end

function gaussian_2D_multipos(loc, positions, sig)
    x, y = loc
    σ₁, σ₂ = sig
    val = 0
    for pos in positions
        x₀, y₀ = pos
        val += gaussian(x, x₀, σ₁) * gaussian(y, y₀, σ₂)
    end
    val
end

# Loss: Given a list of positions, what is the cost / loss of the
function loss(positions)
    simulation = predict(positions)
    # display(simulation)
    final_state = simulation[end]
    # display(size(final_state)[1])
    # display(final_state)
    return norm(final_state, 2)
end

function predict(positions)
    # Simply roll this forward in time and we get our solution
    tspan = (0.0, 1.0)
    u₀ = zeros(size(x))
    u_dual = convert.(eltype(positions[1]),u₀)
    prob = ODEProblem(heat_eq, u_dual, tspan, positions)
    sol = solve(prob)
    # display("Finished ODE Solve")
    sol
end

# function create_2D_gaussian_matrix_orig(xs,ys,positions)
#     out = zeros(size(xs)[1],size(ys)[1])
#     out = convert.(eltype(positions[1]),out)
#     x_idx = 1
#     y_idx = 1

#     for pos in positions
#         x_p, y_p = pos
#         for x in xs
#             for y in ys
#                 out[x_idx, y_idx] += gaussian_2D([x,y],[x_p,y_p],[5.0,5.0])
#                 y_idx += 1
#             end
#             y_idx = 1
#             x_idx += 1
#         end
#         x_idx = 1

#     end
#     out
# end

function create_2D_gaussian_matrix(xs,ys,positions)
    xsz = length(xs)
    ysz = length(ys)
    out = zeros(xsz,ysz)
    out = convert.(eltype(positions[1]),out)

    for i=1:xsz
        for j=1:ysz
            x, y = xs[i], ys[j]
            out[i, j] += gaussian_2D_multipos([x,y],positions,[5.0,5.0])
        end
    end

    out
end

function plot_gaussian()
    positions = [[20.0,30.0],[40.0,50.0]]
    xs = 1.0:100.0
    ys = 1.0:100.0
    gm = create_2D_gaussian_matrix(xs,ys,positions)
    gm_f(x, y) = gm[x, y]
    display(plot(1:length(xs), 1:length(ys), gm_f, st=:heatmap, color=:viridis))
end

function plot_heatmap(data)
    display(plot(1:size(data,1), 1:size(data,2), data, st=:heatmap, color=:viridis))
end

function plot_gaussian_2()
    positions = [[20.0,30.0],[40.0,50.0],[66.0,3.3]]
    xs = 1.0:0.5:100.0
    ys = 1.0:0.5:100.0
    gf = (x,y) -> gaussian_2D_multipos([x,y],positions,[5.0,5.0])
    display(plot(xs, ys, gf, st=:heatmap, color=:viridis))
end

function optimize_and_save_thermal(positions, η=0.1)
    solns = []
    append!(solns, [deepcopy(positions)])
    @showprogress for idx in 1:1000
        grads = ForwardDiff.gradient(s -> loss(s), positions) # Magic
        positions .-= η*grads
        # display(positions)
        # display(loss(positions))

        if loss(positions) < 0.1
            break
        end

        if idx in [500,1000]
            append!(solns, [deepcopy(positions)])
        end
    end

    solns
end


function derivative(f, x::Union{Number,AbstractArray{<:Number}},
                    integrator)
    local d
    alg = unwrap_alg(integrator, true)
    if get_current_alg_autodiff(integrator.alg, integrator.cache)
      d = ForwardDiff.derivative(f, x)
    else
      d = DiffEqDiffTools.finite_difference_gradient(f, x, alg.diff_type, eltype(x), Val{false})
    end
    d
end

function derivative!(df::AbstractArray{<:Number}, f, x::Union{Number,AbstractArray{<:Number}}, fx::AbstractArray{<:Number}, integrator::DEIntegrator)
    if alg_autodiff(integrator.alg)
        ForwardDiff.derivative!(df, f, fx, x)
    else
        RealOrComplex = eltype(integrator.u) <: Complex ? Val{:Complex} : Val{:Real}
        DiffEqDiffTools.finite_difference!(df, f, x, integrator.alg.diff_type, RealOrComplex, fx)
    end
    nothing
end

function jacobian(f, x,
                  integrator::DiffEqBase.DEIntegrator)
    local J
    alg = unwrap_alg(integrator, true)
    isarray = typeof(x) <: AbstractArray
    if get_current_alg_autodiff(integrator.alg, integrator.cache)
      if isarray
        J = ForwardDiff.jacobian(f,x)
      else
        J = ForwardDiff.derivative(f,x)
      end
    else
      if isarray
        J = DiffEqDiffTools.finite_difference_jacobian(f, x, alg.diff_type, eltype(x), Val{false})
      else
        J = DiffEqDiffTools.finite_difference_derivative(f, x, alg.diff_type, eltype(x))
      end
    end
    J
end

function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, fx::AbstractArray{<:Number}, integrator::DEIntegrator, jac_config)
    if alg_autodiff(integrator.alg)
      ForwardDiff.jacobian!(J, f, fx, x, jac_config)
    else
      DiffEqDiffTools.finite_difference_jacobian!(J, f, x, jac_config)
    end
    nothing
end

function build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
  if !has_jac(f)
    if alg_autodiff(alg)
      jac_config = ForwardDiff.JacobianConfig(uf,du1,uprev,ForwardDiff.Chunk{determine_chunksize(u,alg)}())
    else
      if alg.diff_type != Val{:complex}
        jac_config = DiffEqDiffTools.JacobianCache(tmp,du1,du2,alg.diff_type)
      else
        jac_config = DiffEqDiffTools.JacobianCache(Complex{eltype(tmp)}.(tmp),Complex{eltype(du1)}.(du1),nothing,alg.diff_type)
      end
    end
  else
    jac_config = nothing
  end
  jac_config
end


#!/usr/bin/env julia

export pe4

function pe4(st::Integer=100,ed::Integer=999)
    ispalindrome(n)=digits(n)==reverse(digits(n))
    maximum(i*j for i in st:ed, j in st:ed if ispalindrome(i*j))
end

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe4()|>print
end


# m5.5s.jl

using Pkg, DrWatson

using MonteCarloMeasurements
using StanSample
using StatisticalRethinking

# ### snippet 5.29

println()
df = CSV.read(sr_datadir("milk.csv"), DataFrame; delim=';');
df = filter(row -> !(row[:neocortex_perc] == "NA"), df);
df[!, :neocortex_perc] = parse.(Float64, df[:, :neocortex_perc])
df[!, :lmass] = log.(df[:, :mass])
scale!(df, [:kcal_per_g, :neocortex_perc, :lmass])

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

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

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

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

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

# Input data for cmdstan

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

# Sample using StanSample

rc5_5s = stan_sample(m5_5s, data=m5_5_data);

if success(rc5_5s)

    # Describe the draws

    part5_5s = read_samples(m5_5s, :particles)
    part5_5s |> display

    retinking = "
            mean   sd  5.5% 94.5%
    a     0.04 0.15 -0.21  0.29
    bN    0.13 0.22 -0.22  0.49
    sigma 1.00 0.16  0.74  1.26
    "

end


# dynamics tests
tf = 10.0
t = LinRange(0.,tf,Nt)
sim2 = Sim(sim;γ = 0.0,tf=tf,t=t)

#make initial state
ϕi = sol[end]
@pack! sim2 = ϕi

sol2,err = testsim(sim2)
@test err == false

psitest = xspace(sol2[end],sim2)
npeak = abs2.(psitest[256])
@test isapprox(g*npeak,μ,rtol=1e-2)


""" [intended for internal use only]
Enforce small numbers to be zero
Helper function used for testing."""
function set_small_to_zero(a::Number)
	# Note: this tolerates up to approximately 1e-8
   	return(a + 1 ≈ 1 ? 0 : a)
end

"""Squared mahalanobis distance for covariance estimator S and center µ"""
function mahalanobis(arrX::AbstractVector{Float64},
					 µ::AbstractVector{Float64}, S::AbstractArray{Float64,2})
    return((arrX - µ)'*inv(S)*(arrX - µ))
end

# Allows the computation to be mapped on rows of a DataFrame
mahalanobis(x::DataFrameRow, µ::AbstractVector{Float64}, S::AbstractArray{Float64,2}) =
    mahalanobis(convert(Vector, x), µ, S)

"""Squared hellinger distance for covariance estimators S and centers µ"""
function hellinger(µ1::AbstractVector{Float64}, S1::AbstractArray{Float64,2},
				   µ2::AbstractVector{Float64}, S2::AbstractArray{Float64,2})
    S = (S1 + S2)/2
    # Singularity sometimes lead to negative determinant values
    # (computational artifacts) and values need to be set to zero
    dS1 = set_small_to_zero(det(S1))
    dS2 = set_small_to_zero(det(S2))
    dS = set_small_to_zero(det(S))
    H = 1 - (dS1^(1/4))*(dS2^(1/4))/(dS^(1/2))*
    	exp((-1/8)*(µ1-µ2)'*inv(S)*(µ1-µ2))
    return(H)
end


module StateSpaceStudy

    using LinearAlgebra, Random, Distributions, PyPlot

    include("filteringrecursion.jl")

    include("smoothingrecursion.jl")

    include("simulatedata.jl")

    include("Example_4_3.jl")

    include("Example_4_2.jl")

    export simulatedata, filteringrecursion, smoothingrecursion

end


using RAFF
using DelimitedFiles
using PyPlot

#
# This example shows how solve a problem given by data in a
# file. Also, we show how to use the `model_list` utility structure to
# retrieve pre-defined models in RAFF.
#

# Retrieve the number of parameters, the model function and the model
# function string (not used) for the "Cubic" model.
n, model, modelstr = RAFF.model_list["cubic"]

# Read data from file
open("C1.txt") do fp

    global data = readdlm(fp)

end

# Set starting point
initguess = [0.0, 0.0, 0.0, 0.0]

# Set the number of multistart iterations.
maxms = 10

# Call RAFF. In this case, the model is not a multivariate one. The
# last column of the file indicates which observation is the outlier.
rsol = raff(model, data[:, 1:end - 1], n; MAXMS=maxms, initguess=initguess)

println(rsol)

# Now we plot the solution of the problem

x = data[:, 1]
y = data[:, 2]
c = data[:, 3]

# Plot the obtained solution
modl1 = (x) -> model(x, rsol.solution)

t = minimum(x):0.01:maximum(x)
PyPlot.plot(t, modl1.(t), "r", label="RAFF")

# Plot the 'true' solutios, i. e., the one use to generate the
# problem.
θSol = [2.0, 0, -4.0, -10]
modl2 = (x) -> model(x, θSol)

PyPlot.plot(t, modl2.(t), "b--", label="True Solution")

PyPlot.legend(loc=4)

# Color the obtained outliers
c[rsol.outliers] .= 5.0

# Plot all the points
PyPlot.scatter(x, y, c=c, marker="o", s=50.0, linewidths=0.2,
               cmap=PyPlot.cm."Paired", alpha=0.6)

PyPlot.show()


"""
    Plexification{T <: Number} <: Number

An abstract plexification.
"""
abstract type Plexification{T <: Number} <: Number end

"""
    unreal(::T) where T <: Real

The unreal part of a real number is an empty array.
"""
function unreal(::T) where T <: Real
    T[]
end

"""
    array(z::Plexification) -> Array

Return an array with all the `Real` components of `z`.
"""
function array(z::Plexification)
    vcat(real(z), unreal(z))
end

"""
    iszerodivisor(a::Real) -> Bool

Return `true` if a is zero. Equivalent to `iszero(a)`.
"""
function iszerodivisor(a::Real)
    iszero(a)
end

"""
    iszero(z::Plexification) -> Bool

Return `true` if all `Real` components of `z` are zero.
"""
function iszero(z::Plexification)
    iszero(array(z))
end

"""
    isreal(z::Plexification) -> Bool

Return `true` if the unreal part of `z` is zero.
"""
function isreal(z::Plexification)
    iszero(unreal(z))
end

"""
    (+)(z::Plexification)

The identity operation on Cayley-Dickson constructs.
"""
function (+)(z::Plexification)
    z
end

"""
    commutator(x::Plexification, y::Plexification)

Measure the failure of commutativity of the multiplication operation:
```julia
    (x * y) - (y * x)
```
If `commutator(x, y)` is zero, then the multiplication operation is commutative.
"""
function commutator(x::Plexification, y::Plexification)
    (x * y) - (y * x)
end

"""
    associator(x::Plexification, y::Plexification, z::Plexification)

Measure the failure of associativity of the multiplication operation:
```julia
    ((x * y) * z) - (x * (y * z))
```
If `associator(x, y, z)` is zero, then the multiplication operation is associative.
"""
function associator(x::Plexification, y::Plexification, z::Plexification)
    ((x * y) * z) - (x * (y * z))
end

"""
    jacobiator(x::Plexification, y::Plexification, z::Plexification)

Measure the failure of the multiplication operation to satisfy the Jacobi identity:
```julia
    ((x * y) * z) + ((z * x) * y) + ((y * z) * x)
```
If `jacobiator(x, y, z)` is zero, then the multiplication operation satisfies the Jacobi
identity.
"""
function jacobiator(x::Plexification, y::Plexification, z::Plexification)
    ((x * y) * z) + ((z * x) * y) + ((y * z) * x)
end

"""
    alternatorL(x::Plexification, y::Plexification)

Measure the failure of the multiplication operation to be left-alternative:
```julia
    ((x * x) * y) - (x * (x * y))
```
If `alternatorL(x, y)` is zero, then the multiplication operation is left-alternative.
"""
function alternatorL(x::Plexification, y::Plexification)
    associator(x, x, y)
end

"""
    alternatorR(x::Plexification, y::Plexification)

Measure the failure of the multiplication operation to be right-alternative:
```julia
    ((x * y) * y) - (x * (y * y))
```
If `alternatorR(x, y)` is zero, then the multiplication operation is right-alternative.
"""
function alternatorR(x::Plexification, y::Plexification)
    associator(x, y, y)
end

"""
    flexator(x::Plexification, y::Plexification)

Measure the failure of the multiplication operation to be flexible:
```julia
    ((x * y) * x) - (x * (y * x))
```
If `flexator(x, y)` is zero, then the multiplication operation is flexible.
"""
function flexator(x::Plexification, y::Plexification)
    associator(x, y, x)
end

"""
    (/)(x::Plexification, y::Plexification)

Naive right quotient of two Cayley-Dickson constructs.
Equivalent to multiplying `x` on the right by the inverse of `y`.
"""
function (/)(x::Plexification, y::Plexification)
    x * inv(y)
end

"""
    (\\)(y::Plexification, x::Plexification)

Naive left quotient of two Cayley-Dickson constructs.
Equivalent to multiplying `x` on the left by the inverse of `y`.
"""
function (\)(y::Plexification, x::Plexification)
    inv(y) * x
end

"""
    (/)(a::Real, z::Plexification)

Equivalent to multiplying the inverse of `z` by `a`.
"""
function (/)(a::Real, z::Plexification)
    a * inv(z)
end

"""
    (\\)(z::Plexification, a::Real)

Equivalent to multiplying the inverse of `z` by `a`.
"""
function (\)(z::Plexification, a::Real)
    inv(z) * a
end


using BenchmarkTools
using Measurements
using SpecialFunctions
using QuadGK

const SUITE = BenchmarkGroup()

const a = 12.3 ± 4.5
const b = 67.8 ± 9.0
const vector = [1 ± 0.1 for _ in 1:10000]

int1(x) = exp(x / a)

SUITE["base"] = BenchmarkGroup()
SUITE["uncertain"] = BenchmarkGroup()
SUITE["vector"] = BenchmarkGroup()
SUITE["specfun"] = BenchmarkGroup()
SUITE["quadgk"] = BenchmarkGroup()

SUITE["base"]["create"] = @benchmarkable 4.7 ± 0.3
SUITE["base"]["sum"] = @benchmarkable $(a) + $(b)
SUITE["base"]["multiplication"] = @benchmarkable $(a) * $(b)
SUITE["base"]["sqrt"] = @benchmarkable sqrt($(b))
SUITE["base"]["sin"] = @benchmarkable sqrt($(a))

SUITE["uncertain"]["zeta"] = @benchmarkable @uncertain zeta($(2 ± 0.13))
SUITE["uncertain"]["log"] = @benchmarkable @uncertain log($(9.4 ± 1.3), $(58.8 ± 3.7))
SUITE["uncertain"]["atan"] = @benchmarkable @uncertain atan($(10), $(13.5 ± 0.8))

SUITE["vector"]["sqrt"] = @benchmarkable sqrt.($(vector))
SUITE["vector"]["sin"] = @benchmarkable sin.($(vector))
SUITE["vector"]["gamma"] = @benchmarkable gamma.($(vector))
SUITE["vector"]["sin²+cos²"] = @benchmarkable cos.($(vector)) .^ 2 .+ sin.($(vector)) .^ 2
SUITE["vector"]["sum - all different"] = @benchmarkable sum($(vector))
SUITE["vector"]["mean - all different"] = @benchmarkable mean($(vector))
SUITE["vector"]["prod - all different"] = @benchmarkable prod($(vector))
SUITE["vector"]["sum - all equal"] = @benchmarkable sum(v) setup = (x = 1 ± 0.1; v = [x for _ in eachindex(vector)])
SUITE["vector"]["mean - all equal"] = @benchmarkable mean(v) setup = (x = 1 ± 0.1; v = [x for _ in eachindex(vector)])
SUITE["vector"]["prod - all equal"] = @benchmarkable prod(v) setup = (x = 1 ± 0.1; v = [x for _ in eachindex(vector)])

SUITE["specfun"]["gamma"] = @benchmarkable gamma($(a))

SUITE["quadgk"]["int1"] = @benchmarkable quadgk(int1, 1, 7)
SUITE["quadgk"]["cos"] = @benchmarkable quadgk(cos, $(a), $(b))
SUITE["quadgk"]["sin"] = @benchmarkable quadgk(sin, $(-a), $(a))


#=
  Set covering in Julia ConstraintSolver.jl

  Placing of firestations, from Winston "Operations Research", page 486

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK
const CS = ConstraintSolver
include("constraints_utils.jl")

function set_covering()

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS, # <-
                                                            # "traverse_strategy"=>:DBFS,

                                                            "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            # "branch_split"=>:InHalf, # <-

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>16,

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                        ))
    min_distance = 15 # minimum distance
    # distances between the cities
    distance = [[0,10,20,30,30,20],
               [10,0,25,35,20,10],
               [20,25,0,15,30,20],
               [30,35,15,0,15,25],
               [30,20,30,15,0,14],
               [20,10,20,25,14,0]]
    num_cities = length(distance)

    # where to place the fire stations: 1 if placed in this city.
    @variable(model, x[1:num_cities], Bin)
    @variable(model, 0 <= z <= num_cities, Int)

    # calculate the number of covered fire stations
    for j in 1:num_cities
        @constraint(model,sum([x[i] for i in 1:num_cities if distance[i][j] <= min_distance]) >= 1)
    end

    # objective: minimize the number of fire stations
    @constraint(model, z == sum(x))

    @objective(model, Min, z)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                xx = convert.(Integer,JuMP.value.(x; result=sol))
                zz = convert.(Integer,JuMP.value.(z; result=sol))
                println("x:$xx z:$zz")

            end
        end
    else
        println("status:$status")
    end

    return status
end

@time set_covering()


# INTERVAL TYPES

"""
Abstract interval type
"""
abstract type AbstractInterval{T<:AbstractFloat} end

const PersistenceDiagram{T} = AbstractVector{<:AbstractInterval{T}}

"""
    birth(i::AbstractInterval)

Return a birth value of the interval `i`.
"""
birth(i::AbstractInterval) = error("`birth` is not implemented ")

"""
    death(i::AbstractInterval)

Return a death value of the interval `i`.
"""
death(i::AbstractInterval) = error("`death` is not implemented ")

# auxilary methods
show(io::IO, i::AbstractInterval) = print(io, "[$(birth(i)),$(death(i)))")
birthx(i::AbstractInterval) = birth(i) - death(i)
deathx(i::AbstractInterval) = birth(i) + death(i)
pair(i::AbstractInterval) = birth(i) => death(i)
isempty(i::AbstractInterval) = birth(i) == death(i)
eltype(i::AbstractInterval{T}) where {T<:AbstractFloat} = T
in(e::T, i::AbstractInterval{T}) where {T<:AbstractFloat} = birth(i) <= e < death(i)
(==)(i1::AbstractInterval, i2::AbstractInterval) =
    birth(i1) == birth(i2) && death(i1) == death(i2)
isless(i1::AbstractInterval, i2::AbstractInterval) =
    birth(i1) < birth(i2) ? true : ( birth(i1) == birth(i2) ? death(i1) < death(i2) : false )

"""
    Interval{T<:AbstractFloat} <: AbstractInterval{T}

Simple implementation of the `AbstractInterval` type.
"""
struct Interval{T<:AbstractFloat} <: AbstractInterval{T}
    b::T
    d::T
end
Interval(p::Pair{T,T}) where {T<:AbstractFloat} = Interval(p.first, p.second)
Interval() = Interval(0.0=>Inf)

birth(i::Interval) = i.b
death(i::Interval) = i.d
diag(i::Interval) = let c = deathx(i)/2.0; Interval(c, c) end

"""
    AnnotatedInterval{T<:AbstractFloat, C<:AbstractChain} <: AbstractInterval{T}

Interval type annotated with a generator chain.
"""
mutable struct AnnotatedInterval{T<:AbstractFloat, C<:AbstractChain} <: AbstractInterval{T}
    b::T
    d::T
    g::C
end
AnnotatedInterval(b::T, d::T) where {T<:AbstractFloat} = AnnotatedInterval(b, d, EmptyChain())

birth(i::AnnotatedInterval) = i.b
death(i::AnnotatedInterval) = i.d
generator(i::AnnotatedInterval) = i.g
show(io::IO, i::AnnotatedInterval) = print(io, "[$(i.b),$(i.d)) => $(i.g)")

"""
    diagram(d, pts)

Construct persistence diagram of dimension `d` from point pairs `pts`.
"""
diagram(pts::Pair...) = [Interval(p) for p in pts]

#
# I/O
#

function write(io::IO, pdd::Dict{Int, PD}) where {T<:Real, PD<:PersistenceDiagram{T}}
    for (d, pd) in pdd
        write(io, "$d\n")
        for i in pd
            write(io, "$(birth(i)) $(death(i))\n")
        end
    end
end

function read(io::IO, ::Type{I}) where {T<:Real, I<:AbstractInterval{T}}
    pd = Dict{Int, Vector{I}}()
    d = -1
    while !eof(io)
        l = readline(io)
        vals = split(l, ' ')
        if length(vals) == 1
            d = parse(Int, vals[1])
            pd[d] = I[]
        else
            vs = parse.(T, vals)
            push!(pd[d], Interval(vs[1],vs[2]))
        end
    end
    return pd
end


module Layers

export Linear, Sequential, forward, backward, parms, ∇parms

include("./tensor_dt.jl")
using ..TensorDT: Tensor

include("./abstract_layers.jl")
import ..AbstractLayers: AbstractLayer, AL, forward, backward, parms, ∇parms, idims, odims

include("./initializations.jl")
using ..Initializations: init_xavier, INIT_FNs

## ======================================================================
## Linear (or Dense) Layer
## ======================================================================
struct Linear <: AbstractLayer
  idim::Integer   ## input dim
  odim::Integer   ## output dim
  ## parameters
  w::Tensor
  b::Tensor
  #
  _type::Symbol
  store::Dict # {Any. Any}

  ## Assume type is DT == Float64
  function Linear(idim::Integer, odim::Integer;
      DT=Float64, init_fn=init_xavier, seed=42, kwargs...)
    ##
    @assert 1 ≤ idim
    @assert 1 ≤ odim
    @assert init_fn ∈ INIT_FNs
    ##
    w = init_fn((odim, idim); seed, DT, kwargs...)
    b = zeros(DT, (odim, 1))
    new(idim, odim, w, b, :dense, Dict())
  end
end

# gen_key(self::Linear, prefix::Symbol) = Symbol(prefix, self._id[1])

function forward(self::Linear, input::Tensor)::Tensor
  self.store[:input] = input  ## Storing the input for backward pass
  self.w * input + self.b
end

function backward(self::Linear, ∇p::Tensor)::Tensor
  ## as each bᵢ is added to output oᵢ, ∇b is the same as output ∇
  self.store[:∇b] = ∇p
  self.store[:∇w] = ∇p * self.store[:input]'
  r = sum.(self.w' * ∇p)
  @assert(size(self.w) == size(self.store[:∇w]),
          "w and ∇w should have same shape: $(size(self.w)) == $(size(self.store[:∇w]))")
  @assert(size(self.b) == size(self.store[:∇b]),
          "b and ∇b should have same shape: $(size(self.b)) == $(size(self.store[:∇b]))")
  r
end

parms(self::Linear) = [self.w, self.b]
∇parms(self::Linear) = [self.store[:∇w], self.store[:∇b]]

## Extra
idims(self::Linear) = self.idim
odims(self::Linear) = self.odim

## ======================================================================
## Neural Network as a sequence of Layers
## ======================================================================

const CHECK_DIMS_EXCLUSION = [:activation, :regularization] #:embedding]

struct Sequential <: AbstractLayer
  layers::Vector{AbstractLayer}
  _type::Symbol

  function Sequential(layers::Vector{AbstractLayer})
    @assert length(layers) > 0
    check_dim_layer(layers)
    #
    new(layers, :sequential)
  end
end

function forward(self::Sequential, input::Tensor)::Tensor
  for (ix, l) ∈ enumerate(self.layers)
    input = forward(l, input)
  end
  input
end

function backward(self::Sequential, ∇p::Tensor)::Tensor
  for (ix, l) ∈ enumerate(reverse(self.layers))
    ∇p = backward(l, ∇p)
  end
  ∇p
end

parms(self::Sequential) = [p for l ∈ self.layers for p ∈ parms(l)]
∇parms(self::Sequential) = [∇p for l ∈ self.layers for ∇p ∈ ∇parms(l)]

idims(self::Sequential) = nothing
odims(self::Sequential) = nothing

##
## Internal helpers
##

function check_dim_layer(layers::Vector{AbstractLayer})
  """
  output of prev. layer == input of curr layer
  """
  @assert length(layers) > 0
  pl = layers[1]._type ∈ CHECK_DIMS_EXCLUSION ? layers[1] : nothing

  for cl ∈ layers[2:end]
    cl._type ∈ CHECK_DIMS_EXCLUSION && continue
    isnothing(pl) && (pl = cl; continue)

    odims(pl) ≠ idims(cl) &&
      throw(ArgumentError("incompatible shape between pl.odim; $(odims(pl)) vs cl.idim: $(idims(cl))"))
    pl = cl
  end
end

end ## Module


using Plots
using PyCall
using DifferentialEquations
using Statistics
using Base.Iterators

function rossler(du,u,p,t)
    x, y, z = u
    a, b, c = p
    du[1] = -y -z
    du[2] = x + a*y
    du[3] = b + z*(x-c)
end

p = [0.1,0.1,18]
u0 = [5., 5., 1.]
tspan = (0.0,1000.0)

prob = ODEProblem(rossler, u0, tspan, p)
tmp = solve(prob, abstol=1e-9)
prob = ODEProblem(rossler, tmp.u[end], (0.0,5000.0), p)
sol = solve(prob, abstol=1e-9)

pyplot()
pygui(false)
plot(sol,vars=(1,2,3), lw=0.1, legend=false)
savefig("rossler_c13.pdf")

function local_maxima(x::Vector{T}) where T<:Real
    y = []
    for i in 2:length(x)-1
        if x[i-1] < x[i] && x[i] > x[i+1]
            push!(y, x[i])
        end
    end
    return y
end

z = map(v->v[3], sol.u)
zmax = local_maxima(z)

pygui(false)
scatter(
    [zmax[i] for i in 1:(length(zmax)-1)],
    [zmax[i] for i in 2:length(zmax)],
    markersize=1,
    markerstrokealpha=0.0,
    markercolor = :black,
    legend=false,
    xlabel = "z[k]",
    ylabel = "z[k+1]"
    )

include("FractalDimensions.jl")
d, N, ε = box_counting_dimension(sol.u, 8)
loglog_regression(N,ε)

plot(-log.(ε), log.(N), xlabel="-ln ε", ylabel="ln N", marker=:d, legend=false)

C, r = correlation_dimension(sol.u, 15)
loglog_regression(C,r)

plot(log.(r), log.(C), xlabel="ln r", ylabel="ln C", marker=:d, legend=false)
savefig("cordim_rossler.pdf")

using DynamicalSystems
using ChaosTools
# almost feels like cheating
r = Systems.roessler(a=0.1,b=0.1,c=18)
lyapunov(r,1000.0)
λ = lyapunovs(r,1000.0)
kaplanyorke_dim(λ) # Lyapunov dimension

crange = LinRange(1,18,1000)
n = 5000
Ttr = 2000
# bifurcation diagram of the Poincaré map of the z coordinate on the plane y=0
diagram = produce_orbitdiagram(r, (2,0.0), 3, 3, crange; n=n, Ttr=Ttr)

L = length(crange)
x = []
y = []
for j in 1:L, z in diagram[j]
    push!(x,crange[j])
    push!(y,z)
end

plotlyjs()
scatter(
    x,
    y,
    markersize=0.5,
    markerstrokealpha=0.0,
    markercolor = :black,
    legend=false,
    grid = false,
    xlabel = "c",
    ylabel = "z"
    )

savefig("rossler_bifurcation.pdf")


# sample Binomial r.v. using inverse transform
@inline function sampleInverseCDFBinom(n::Int64, p::Float64, rng::RNG) where
  RNG <: AbstractRNG
  q::Float64 = 1-p
  s::Float64 = p/q
  k::Int64 = 0
  f::Float64 = q^n
  u::Float64 = rand(rng)
  while u > f
    u -= f
    k += 1
    f *= (n + 1 - k) / k * s
  end
  return k
end

let
  btrd_fc_table = [
    0.08106146679532726,
    0.04134069595540929,
    0.02767792568499834,
    0.02079067210376509,
    0.01664469118982119,
    0.01387612882307075,
    0.01189670994589177,
    0.01041126526197209,
    0.009255462182712733,
    0.008330563433362871]
  global @inline function btrd_fc(k::Int64)
    if k < 10
      return btrd_fc_table[k+1]
    else
      tmp::Float64 = 1/(k+1)
      return (1/12 - (1/360 - tmp * tmp / 1260) * tmp * tmp) * tmp
    end
  end
end

## Generate a Binomial r.v. using the algorithm of
## W. Hörmann. The Generation of Binomial Random Variates, Journal of
## Statistical Computation and Simulation 46, 1993.
## intended for use when n*p >= 10 and p <= 0.5
@inline function btrd(n::Int64, p::Float64, rng::RNG) where RNG <: AbstractRNG
  # 0. setup
  m::Int64 = floor(Int64, (n + 1) * p)
  r::Float64 = p / (1 - p)
  nr::Float64 = (n + 1) * r
  npq::Float64 = n * p * (1 - p)
  sqrt_npq::Float64 = sqrt(npq)
  b::Float64 = 1.15 + 2.53 * sqrt_npq
  a::Float64 = -0.0873 + 0.0248 * b + 0.01 * p
  c::Float64 = n * p + 0.5
  α::Float64 = (2.83 + 5.1 / b) * sqrt_npq
  vr::Float64 = 0.92 - 4.2 / b
  ur_vr::Float64 = 0.86 * vr
  u::Float64 = 0.0
  v::Float64 = 0.0
  us::Float64 = 0.0
  k::Int64 = 0
  km::Int64 = 0
  f::Float64 = 0.0
  i::Int64 = 0
  while true
    # 1.
    v = rand(rng)
    if v <= ur_vr
      u = v / vr - 0.43
      return floor(Int64, (2 * a / (0.5 - abs(u)) + b) * u + c)
    end
    # 2.
    if v >= vr
      u = rand(rng) - 0.5
    else
      u = v / vr - 0.93
      u = sign(u) * 0.5 - u
      v = vr * rand(rng)
    end
    # 3.0
    us = 0.5 - abs(u)
    k = floor(Int64, (2 * a / us + b) * u + c)
    if k < 0 || k > n
      continue
    end
    v = v * α / (a / (us * us) + b)
    km = abs(k - m)
    # [ if km > 15.0 goto 3.2 ]
    if km <= 15.0
      # 3.1 [ recursive evaluation of f(k) ]
      f = 1.0
      i = min(m, k)
      if m < k
        while i != k
          i += 1
          f *= (nr / i - r)
        end
      elseif m > k
        while i != m
          i += 1
          v *= (nr / i - r)
        end
      end
      if v <= f
        return k
      else
        continue
      end
    else
      # 3.2
      v = log(v)
      ρ = (km / npq) * (((km / 3 + 0.625) * km + 1/6) / npq + 0.5)
      t = - km * km / (2 * npq)
      if v < t - ρ
        return k
      end
      if v > t + ρ
        continue
      end
      # 3.3
      nm = n - m + 1
      h = (m + 0.5) * log((m + 1) / (r * nm)) + btrd_fc(m) + btrd_fc(n - m)
      # 3.4
      nk = n - k + 1
      if v <= h + (n+1)*log(nm/nk)+(k+0.5)*log(nk*r/(k+1))-btrd_fc(k)-btrd_fc(n-k)
        return k
      end
    end
  end
end

# choose which algorithm to use for given (n,p)
@inline function sampleBinomial(n::Int64, p::Float64,
  rng::RNG = GLOBAL_RNG) where RNG <: AbstractRNG
  if n * p >= 10 && n * (1.0 - p) >= 10
    if p <= 0.5
      return btrd(n, p, rng)
    else
      return n - btrd(n, 1.0 - p, rng)
    end
  else
    if p <= 0.5
      return sampleInverseCDFBinom(n, p, rng)
    else
      return n - sampleInverseCDFBinom(n, 1.0 - p, rng)
    end
  end
end


#= -------------------------------------------------------------------------
# @file testing_transformations.jl
#
# @date 08/07/19 08:34:43
# @author Martin Noblia
# @email mnoblia@disroot.org
#
# @brief
#
# @detail
#
#  Licence:
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License

---------------------------------------------------------------------------=#
# testing
using Robotics, LinearAlgebra
# NOTE(elsuizo:2019-08-07): no usar == cuando se comparan Floats
@testset "transformations" begin
   @testset "rotx" begin
      R1 = rotx(deg2rad(30))
      R2 = rotx(deg2rad(30 + 360))
      @test R1 ≈ R2
      @test det(R1) ≈ det(R2) ≈ 1.0
   end
   @testset "roty" begin
      R3 = roty(deg2rad(30))
      R4 = roty(deg2rad(30 + 360))
      @test R3 ≈ R4
      @test det(R3) ≈ det(R4) ≈ 1.0
   end
   @testset "rotz" begin
      R5 = rotz(deg2rad(30))
      R6 = rotz(deg2rad(30 + 360))
      @test R5 ≈ R6
      @test det(R5) ≈ det(R6) ≈ 1.0
   end
end


using SEC, Test
import SEC.Emulate: mean, cov, GaussianProcess, condition

@testset "GaussianProcess unit test" begin
    μ(x) = 0.0
    k(x,y) = exp(-(x-y)^2/2)
    gp = GaussianProcess(mean = μ, covariance = k)
    @test μ === mean(gp)
    @test k === cov(gp)
end

@testset "1D GPR mean and cov test" begin
    μ(x) = 0.0
    k(x,y) = exp(-(x-y)^2/2)
    gp = GaussianProcess(mean = μ, covariance = k)
    N = 10
    X = [i for i in 1:N]
    μX = mean(gp, X)
    σX = cov(gp, X)
    @test all(μX .== zeros(Float64, N))
    @test σX[1,1] ≈ 1.0
    @test σX[1, N] ≈ k(1, N)
end

@testset "1D GPR conditional mean and cov test" begin
    μ(x) = 0.0
    k(x,y) = exp(-(x-y)^2/2)
    gp = GaussianProcess(mean = μ, covariance = k)
    N = 10
    X = [i for i in 1:N]
    Y = [2*i for i in 1:N]
    data = (; X, Y)
    cgp = condition(gp, data)
    cμX =  mean(cgp, X)
    tolerance = 1e-12
    @test all(abs.(cμX .- Y) .≤ tolerance)
    covXX =  [cgp.covariance(x,x) for x in X]
    @test all(abs.(covXX) .≤ tolerance)
end


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 4afee820-f909-11ea-06bc-d51a333b2baa
using Distributions, Random, Statistics, FreqTables

# ╔═╡ 3c23623a-f90a-11ea-2b40-05b297cb37dd
using Optim, GLM, LinearAlgebra

# ╔═╡ d989e292-f90a-11ea-3e73-516abb783304
using HTTP, CSV, DataFrames

# ╔═╡ 08f4baf6-f90c-11ea-0ead-4330f04896db
include("lgwt.jl")

# ╔═╡ 502cd58e-f911-11ea-134b-b720cf9f5592
md"""
# Loading packages and Data
"""

# ╔═╡ 78789bb2-f921-11ea-3f1d-2d6d33d9fb8b
md"""
# Question $1$
Estimate a multinomial logit (with alternative-specific covariates $Z$) on the following data set, which is a panel form of the same data as Problem Set 3. You should be able to simply re-use your code from Problem Set 3. However, I would ask that you use automatic differentiation to speed up your estimation, and to obtain the standard errors of your estimates.
"""

# ╔═╡ c40137b0-f921-11ea-0655-45cdc040cd54
begin
	url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS4-mixture/nlsw88t.csv"
	df = CSV.read(HTTP.get(url).body)
	X = [df.age df.white df.collgrad]
	Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
         df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
	y = df.occ_code
end

# ╔═╡ d2596706-f921-11ea-23fd-ef3d19f4e322
md"""
**Note:** this took my machine about 30 minutes to estimate using random starting values. You might consider using the estimated values from Question 1 of PS3.

The choice set is identical to that of Problem Set 3 and represents possible occupations and is structured  as follows.

1. Professional/Technical
2. Managers/Administrators
3. Sales
4. Clerical/Unskilled
5. Craftsmen
6. Operatives
7. Transport
8. Other
"""

# ╔═╡ c58ae742-f922-11ea-0a1a-63f0c6a48759
begin
	function mlogit_with_Z(theta, X, Z, y)

	        alpha = theta[1:end-1]
	        gamma = theta[end]
	        K = size(X,2)
	        J = length(unique(y))
	        N = length(y)
	        bigY = zeros(N,J)
	        for j=1:J
	            bigY[:,j] = y.==j
	        end
	        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

	        T = promote_type(eltype(X),eltype(theta))
	        num   = zeros(T,N,J)
	        dem   = zeros(T,N)
	        for j=1:J
	            num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma)
	            dem .+= num[:,j]
	        end

	        P = num./repeat(dem,1,J)

	        loglike = -sum( bigY.*log.(P) )

	        return loglike
	    end
	    startvals = [2*rand(7*size(X,2)).-1; .1]
	    td = TwiceDifferentiable(theta -> mlogit_with_Z(theta, X, Z, y), startvals; autodiff = :forward)
	    # run the optimizer
	    theta_hat_optim_ad = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
	    theta_hat_mle_ad = theta_hat_optim_ad.minimizer
	    # evaluate the Hessian at the estimates
	    H  = Optim.hessian!(td, theta_hat_mle_ad)
	    theta_hat_mle_ad_se = sqrt.(diag(inv(H)))

end

# ╔═╡ dd2547f4-f921-11ea-1ba1-fd07f1bed50a
md"""
# Question $2$

Does the estimated coefficient $\hat{\gamma}$ make more sense now than in Problem Set 3?
"""

# ╔═╡ b3340562-f925-11ea-1983-a1eca1e95821
γ = (theta_hat_mle_ad[end]/100)

# ╔═╡ 1b8742de-f90b-11ea-3f90-991ccc07fcae
md"""
# Question 3

Now we will estimate the mixed logit version of the model in Question 1, where $\gamma$ is distributed $N\left(\tilde{\gamma},\sigma^2_{\gamma}\right)$. Following the notes, the  formula for the choice probabilities will be

While this looks daunting, we can slightly modify the objective function from Question 1.

The first step is to recognize that we will need to \textit{approximate} the integral in \eqref{eq:intlike}. There are many ways of doing this, but we will use something called Gauss-Legendre quadrature.\footnote{Another popular method of approximating the integral is by simulation.} We can rewrite the integral in \eqref{eq:intlike} as a (weighted) discrete summation:

$
\begin{align}
\label{eq:quadlike}
\ell\left(X,Z;\beta,\mu_{\gamma},\sigma_{\gamma}\right)&=\sum_{i=1}^N\sum_{t=1}^T \log\left\{\sum_{k}\omega_{r}\prod_{j}\left[\frac{\exp\left(X_{it}\left(\beta_{j}-\beta_{J}\right)+\xi_r\left(Z_{itj}-Z_{itJ}\right)\right)}{\sum_k \exp\left(X_{it}\left(\beta_{k}-\beta_{J}\right)+\xi_r\left(Z_{itk}-Z_{itJ}\right)\right)}\right]^{d_{itj}}f\left(\xi_r\right)\right\}
\end{align}
$

where $\omega_r$ are the quadrature weights and $\xi_r$ are the quadrature points. $\tilde{\gamma}$ and $\sigma_{\gamma}$ are parameters of the distribution function $f\left(\cdot\right)$.

1. Before we dive in, let's learn how to use quadrature. In the folder where this problem set is posted on GitHub, there is a file called \texttt{lgwt.jl}. This is a function that returns the $\omega$'s and $\xi$'s for a given choice of $K$ (number of quadrature points) and bounds of integration $[a,b]$.

Let's practice doing quadrature using the density of the Normal distribution.
"""

# ╔═╡ d0ee42d2-f90e-11ea-2409-dd20c2aba3fe
d = Normal(0,1); # μ = 0, σ = 1

# ╔═╡ dea66a7e-f90e-11ea-3e27-f9f08da3cdf1
md"""
Once we have the distribution defined, we can do things like evaluate its density or probability, take draws from it, etc.

We want to verify that $\int \phi(x)dx$ is equal to 1 (i.e. integrating over the density of the support equals 1). We also want to verify that $\int x\phi(x)dx$ is equal to $\mu$ (which for the distribution above is 0).

When using quadrature, we should try to pick a number of points and bounds that will minimize computational resources, but still give us a good approximation to the integral. For Normal distributions, $\pm 4\sigma$ will get us there.

get quadrature node and weights for 7 grid points
"""

# ╔═╡ 2909be18-f90c-11ea-3974-d347c718d61f
nodes, weights = lgwt(10, -5,5); # ±5σ

# ╔═╡ 0edfe3f2-f90f-11ea-0388-550a5584c468
md"""
now compute the integral over the density and verify it's 1
"""

# ╔═╡ e73ce9e6-f90c-11ea-3924-1b474a3182b3
sum(weights .* pdf.(d, nodes))

# ╔═╡ 22f79f3a-f90f-11ea-2447-c77328685e30
md"""
now compute the expectation and verify it's $0$
"""

# ╔═╡ e74a2fde-f90c-11ea-2d85-736fb0aa9d1d
sum(weights.* nodes .* pdf.(d, nodes))

# ╔═╡ f5d6438a-f90c-11ea-01a9-8f7d4f1434c8
md"""
To get some more practice, I'd like you to use quadrature to compute the following integrals:

$\int_{-5\sigma}^{5\sigma}x^{2}f\left(x\right)dx$ where $f\left(\cdot\right)$ is the pdf of a $N\left(0,2\right)$ distribution using 10 quadrature points
"""

# ╔═╡ 79061b6a-f910-11ea-2340-478981473b5f
sum(weights .* nodes.^2 .* pdf.(Normal(0,2), nodes))

# ╔═╡ d568097e-f90c-11ea-2d78-b7d7ab4ba704
md"""
## Monte Carlo integration

An alternative to quadrature is Monte Carlo integration. Under this approach, we approximate the integral of $f$ by averaging over a function of many random numbers. Formally, we have that

$\int_a^b f\left(x\right)dx \approx \left(b-a\right)\frac{1}{D}\sum_{i=1}^D f\left(X_{i}\right)$

where $D$ is the number of random draws and where each $X_i$ is drawn from a $U[a,b]$ interval

- With $D=1,000,000$, use the formula abobe to approximate $\int_{-5\sigma}^{5\sigma}x^{2}f\left(x\right)dx$ where $f\left(\cdot\right)$ is the pdf of a $N\left(0,2\right)$ and verify that it is (very) close to 4
-  Do the same for $\int_{-5\sigma}^{5\sigma}xf\left(x\right)dx$ where $f\left(\cdot\right)$ and verify that it is very close to 0
- Do the same for $\int_{-5\sigma}^{5\sigma}f\left(x\right)dx$ and verify that it is very close to 1
-  Comment on how well the simulated integral approximates the true value when $D=1,000$ compared to when $D=1,000,000$.
"""

# ╔═╡ 0463c252-f911-11ea-2ca2-09693e84fd6f
σ = 2

# ╔═╡ d1259856-f90c-11ea-1a7b-491d175ceba9
function montecarlo(a, b, D)
	xrand = rand(Uniform(a,b), D)
	σ = ((b-a)/D) * sum(xrand.^2 .* pdf(Normal(0,2), xrand))
	μ = ((b-a)/D) * sum(xrand .* pdf(Normal(0,2), xrand))
	FDP = ((b-a)/D) * sum(pdf(Normal(0,2), xrand))

	return σ, μ, FDP
end

# ╔═╡ b4e8ce6c-f914-11ea-31ba-cf912ec061c4
montecarlo(-5σ, 5σ, 10^6)

# ╔═╡ Cell order:
# ╟─502cd58e-f911-11ea-134b-b720cf9f5592
# ╠═4afee820-f909-11ea-06bc-d51a333b2baa
# ╠═3c23623a-f90a-11ea-2b40-05b297cb37dd
# ╠═d989e292-f90a-11ea-3e73-516abb783304
# ╠═08f4baf6-f90c-11ea-0ead-4330f04896db
# ╟─78789bb2-f921-11ea-3f1d-2d6d33d9fb8b
# ╠═c40137b0-f921-11ea-0655-45cdc040cd54
# ╟─d2596706-f921-11ea-23fd-ef3d19f4e322
# ╠═c58ae742-f922-11ea-0a1a-63f0c6a48759
# ╟─dd2547f4-f921-11ea-1ba1-fd07f1bed50a
# ╠═b3340562-f925-11ea-1983-a1eca1e95821
# ╟─1b8742de-f90b-11ea-3f90-991ccc07fcae
# ╠═d0ee42d2-f90e-11ea-2409-dd20c2aba3fe
# ╟─dea66a7e-f90e-11ea-3e27-f9f08da3cdf1
# ╠═2909be18-f90c-11ea-3974-d347c718d61f
# ╟─0edfe3f2-f90f-11ea-0388-550a5584c468
# ╠═e73ce9e6-f90c-11ea-3924-1b474a3182b3
# ╟─22f79f3a-f90f-11ea-2447-c77328685e30
# ╠═e74a2fde-f90c-11ea-2d85-736fb0aa9d1d
# ╟─f5d6438a-f90c-11ea-01a9-8f7d4f1434c8
# ╠═79061b6a-f910-11ea-2340-478981473b5f
# ╟─d568097e-f90c-11ea-2d78-b7d7ab4ba704
# ╠═0463c252-f911-11ea-2ca2-09693e84fd6f
# ╠═d1259856-f90c-11ea-1a7b-491d175ceba9
# ╠═b4e8ce6c-f914-11ea-31ba-cf912ec061c4


import Pkg; Pkg.add(Pkg.PackageSpec(url="https://github.com/JuliaComputing/JuliaAcademyData.jl"))
using JuliaAcademyData; activate("Deep learning with Flux")

#

#
# ## Neural networks
#
# Now that we know what neurons are, we are ready for the final step: the neural network!. A neural network is literally made out of a network of neurons that are connected together.
#
# So far, we have just looked at single neurons, that only have a single output.
# What if we want multiple outputs?
#
#
# ### Multiple output models
#
# What if we wanted to distinguish between apples, bananas, *and* grapes? We could use *vectors* of `0` or `1` values to symbolize each output.
#
#
#
# The idea of using vectors is that different directions in the space of outputs encode information about different types of inputs.

#-

# Now we extend our previous model to give multiple outputs by repeating it with different weights. For the first element of the array we'd use:
#
# $$\sigma(x;w^{(1)},b^{(1)}) := \frac{1}{1 + \exp(-w^{(1)} \cdot x + b^{(1)})};$$
#
# then for the second we'd use
#
# $$\sigma(x;w^{(2)},b^{(2)}) := \frac{1}{1 + \exp(-w^{(2)} \cdot x + b^{(2)})};$$
#
# and if you wanted $n$ outputs, you'd have for each one
#
# $$\sigma(x;w^{(i)},b^{(i)}) := \frac{1}{1 + \exp(-w^{(i)} \cdot x + b^{(i)})}.$$

#-

# Notice that these equations are all the same, except for the parameters, so we can write this model more succinctly, as follows. Let's write $b$ in an array:
#
# $$b=\left[\begin{array}{c}
# b_{1}\\
# b_{2}\\
# \vdots\\
# b_{n}
# \end{array}\right]$$
#
# and put our array of weights as a matrix:
#
# $$ \mathsf{W}=\left[\begin{array}{c}
# \\
# \\
# \\
# \\
# \end{array}\begin{array}{cccc}
# w_{1}^{(1)} & w_{2}^{(1)} & \ldots & w_{n}^{(1)}\\
# w_{1}^{(2)} & w_{2}^{(2)} & \ldots & w_{n}^{(2)}\\
# \vdots & \vdots &  & \vdots\\
# w_{1}^{(n)} & w_{2}^{(n)} & \ldots & w_{n}^{(n)}
# \end{array}\right]
# $$
#
# We can write this all in one line as:
#
# $$\sigma(x;w,b)= \left[\begin{array}{c}
# \sigma^{(1)}\\
# \sigma^{(2)}\\
# \vdots\\
# \sigma^{(n)}
# \end{array}\right] = \frac{1}{1 + \exp(-\mathsf{W} x + b)}$$
#
# $\mathsf{W} x$ is the operation called "matrix multiplication"
#
# [Show small matrix multiplication]

W = [10 1;
     20 2;
     30 3]
x = [3;
     2]
W*x

# It takes each column of weights and does the dot product against $x$ (remember, that's how $\sigma^{(i)}$ was defined) and spits out a vector from doing that with each column. The result is a vector, which makes this version of the function give a vector of outputs which we can use to encode larger set of choices.
#
# Matrix multiplication is also interesting since **GPUs (Graphics Processing Units, i.e. graphics cards) are basically just matrix multiplication machines**, which means that by writing the equation this way, the result can be calculated really fast.

#-

# This "multiple input and multiple output" version of the sigmoid function is known as a *layer of neurons*.
#
# Previously we worked with a single neuron, which we visualized as
#
#
#
# where we have two pieces of data (green) coming into a single neuron (pink) that returned a single output. We could use this single output to do binary classification - to identify an image of a fruit as `1`, meaning banana or as `0`, meaning not a banana (or an apple).
#
# To do non-binary classification, we can use a layer of neurons, which we can visualize as
#
#
#
# We now have stacked a bunch of neurons on top of each other to hopefully work together and train to output results of more complicated features.
#
# We still have two input pieces of data, but now have several neurons, each of which produces an output for a given binary classification:
# * neuron 1: "is it an apple?"
# * neuron 2: "is it a banana?"
# * neuron 3: "is it a grape?"

#-

# # Multiple outputs with Flux.jl

#-

# First step: load the data.

using CSV, DataFrames, Flux, Plots
## Load apple data in CSV.read for each file
apples1 = DataFrame(CSV.File(datapath("data/Apple_Golden_1.dat"), delim='\t', allowmissing=:none, normalizenames=true))
apples2 = DataFrame(CSV.File(datapath("data/Apple_Golden_2.dat"), delim='\t', allowmissing=:none, normalizenames=true))
apples3 = DataFrame(CSV.File(datapath("data/Apple_Golden_3.dat"), delim='\t', allowmissing=:none, normalizenames=true))
## And then concatenate them all together
apples = vcat(apples1, apples2, apples3)
bananas = DataFrame(CSV.File(datapath("data/Banana.dat"), delim='\t', allowmissing=:none, normalizenames=true))
grapes1 = DataFrame(CSV.File(datapath("data/Grape_White.dat"), delim='\t', allowmissing=:none, normalizenames=true))
grapes2 = DataFrame(CSV.File(datapath("data/Grape_White_2.dat"), delim='\t', allowmissing=:none, normalizenames=true))
grapes = vcat(grapes1, grapes2)

#-

## Extract out the features and construct the corresponding labels
x_apples  = [ [apples[i, :red], apples[i, :blue]] for i in 1:size(apples, 1) ]
x_bananas  = [ [bananas[i, :red], bananas[i, :blue]] for i in 1:size(bananas, 1) ]
x_grapes = [ [grapes[i, :red], grapes[i, :blue]] for i in 1:size(grapes, 1) ]
xs = vcat(x_apples, x_bananas, x_grapes)
ys = vcat(fill([1,0,0], size(x_apples)),
          fill([0,1,0], size(x_bananas)),
          fill([0,0,1], size(x_grapes)))

# ### One-hot vectors

#-

# Recall:
#
#

#-

# `Flux.jl` provides an efficient representation for one-hot vectors, using advanced features of Julia so that it does not actually store these vectors, which would be a waste of memory; instead `Flux` just records in which position the non-zero element is. To us, however, it looks like all the information is being stored:

using Flux: onehot

onehot(2, 1:3)

#-

ys = vcat(fill(onehot(1, 1:3), size(x_apples)),
          fill(onehot(2, 1:3), size(x_bananas)),
          fill(onehot(3, 1:3), size(x_grapes)))

# ## The core algorithm from the previous lecture

## model = Dense(2, 1, σ)
## L(x,y) = Flux.mse(model(x), y)
## opt = SGD(params(model))
## Flux.train!(L, zip(xs, ys), opt)

# ### Visualization

using Plots
plot()

contour!(0:0.01:1, 0:0.01:1, (x,y)->model([x,y]).data[1], levels=[0.5, 0.51], color = cgrad([:blue, :blue]))
contour!(0:0.01:1, 0:0.01:1, (x,y)->model([x,y]).data[2], levels=[0.5,0.51], color = cgrad([:green, :green]))
contour!(0:0.01:1, 0:0.01:1, (x,y)->model([x,y]).data[3], levels=[0.5,0.51], color = cgrad([:red, :red]))

scatter!(first.(x_apples), last.(x_apples), m=:cross, label="apples", color = :blue)
scatter!(first.(x_bananas), last.(x_bananas), m=:circle, label="bananas", color = :green)
scatter!(first.(x_grapes), last.(x_grapes), m=:square, label="grapes", color = :red)

Courses/Deep learning with Flux/1500.Intro-to-neural-networks.jl