#
# Halley's method
#
# Reference:
# Y. Nakatsukasa, Z. Bai and F. Gygi, Optimizing Halley's iteration
# for computing the matrix polar decomposition, SIAM, J. Mat. Anal.
# Vol. 31, Num 5 (2010) pp. 2700-2720
#

type HalleyAlg{T} <: PolarAlg
    maxiter::Int
    verbose::Bool
    tol::T

    function HalleyAlg(;maxiter::Integer=100,
                       verbose::Bool=false,
                       tol::Real = cbrt(eps(T)))
        maxiter > 1 || error("maxiter must be greater than 1.")
        tol > 0 || error("tol must be positive.")

        new(maxiter,
            verbose,
            tol)
    end
end

function solve!{T}(alg::HalleyAlg,
                X::Matrix{T}, U::Matrix{T}, H::Matrix{T})
    common_iter!(HalleyUpdater(), X, U, H, alg.maxiter, alg.verbose, alg.tol)
end

immutable HalleyUpdater <: PolarUpdater end

function update_U!{T}(upd::HalleyUpdater, U::Matrix{T})
    UtU = Array(T, size(U))
    At_mul_B!(UtU, U, U)
    copy!(U, U * (3*I + UtU)* inv(I + 3*UtU))
end

#
# QR-based Dynamically Weighted Halley (QDWH) algorithm
#
# Reference: Optimizing Halley's iteration for computing the matrix
#            polar decomposition, Yuji Nakatsukasa, Zhaojun Bai and
#            Francois Gygi, SIAM, J. Mat. Anal. Vol. 31, Num 5 (2010)
#            pp. 2700-2720
#
# Limitations: 1. the QDWH method should support m > n matrix itself in the future.
#              2. the computing of alpha and L in solve! can be improved by using
#                 norm and condition number estimate, which are currently
#                 not available in Julia.
#
type QDWHAlg{T} <: PolarAlg
    maxiter::Int
    verbose::Bool
    piv::Bool       # whether to pivot
    tol::T

    function QDWHAlg(;maxiter::Integer=100,
                     verbose::Bool=false,
                     piv::Bool=true,
                     tol::Real=cbrt(eps(T)))
        maxiter > 1 || error("maxiter must be greater than 1.")
        tol > 0 || error("tol must be positive.")

        new(maxiter,
            verbose,
            piv,
            tol)
    end
end

function solve!{T}(alg::QDWHAlg,
                X::Matrix{T}, U::Matrix{T}, H::Matrix{T})
    # alpha is an estimate of the largest singular value of the
    # original matrix
    X_temp = Array(T, size(X))
    copy!(X_temp, X)

    n = size(X_temp, 1)
    alpha = vecnorm(X_temp)
    for i in length(X_temp)
        X_temp[i] /= alpha # form X0
    end

    # L is a lower bound for the smallest singular value of X0
    smin_est = norm(X_temp, 1)/cond(X_temp, 1)
    L  = smin_est/convert(T, sqrt(n))

    common_iter!(QDWHUpdater(alg.piv, L), X, U, H, alg.maxiter, alg.verbose, alg.tol)
end

type QDWHUpdater{T} <: PolarUpdater
    piv::Bool   # whether to pivot QR factorization
    L::T  # a lower bound for the smallest singluar value of each update matrix U
end

function update_U!{T}(upd::QDWHUpdater, U::Matrix{T})
    piv = upd.piv
    L = upd.L
    m, n = size(U)
    B = Array(T, m+n, n)
    Q1 = Array(T, n, n)
    Q2 = Array(T, n, n)
    # Compute paramters L, a, b, c
    L2 = L^2
    dd = try
        (4 * (1 - L2)/L2^2)^(1/3)
        catch
        (complex(4 * (1 -L2)/L2^2, 0))^(1/3)
    end
    sqd = sqrt(1+dd)
    a = sqd + 0.5 * sqrt(8 - 4 * dd + 8 * (2 - L2)/(L2 * sqd))
    a = real(a)
    b = (a - 1)^2 / 4
    c = a + b - 1

    # update L
    upd.L = L * (a + b * L2)/(1 + c * L2)

    copy!(B, [sqrt(c)*U; eye(n)])
    if piv
        VERSION < v"0.4.0-dev+1827" ? F = qrfact(B, pivot = true) : F = qrfact(B, Val{true})
    else
        F = qrfact(B)
    end
    copy!(Q1, full(F[:Q])[1:m, :])
    copy!(Q2, full(F[:Q])[m+1:end, :])
    copy!(U, b / c * U + (a - b / c) / sqrt(c) * Q1 * Q2')

end


module AssemblersTests

using Test

using Test
using Gridap.Arrays
using Gridap.TensorValues
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.Fields
using Gridap.Algebra
using SparseArrays
using Gridap.FESpaces
using Gridap.CellData

domain =(0,1,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)

V = TestFESpace(
  model,
  ReferenceFE(lagrangian,Float64,1),
  dirichlet_tags=[1,2,3,4,6,5])

u(x) = x[1]+x[2]

U = TrialFESpace(u,V)

dv = get_cell_shapefuns(V)
du = get_cell_shapefuns_trial(U)

degree = 2
Ω = Triangulation(model)
dΩ = Measure(Ω,degree)

Γ = BoundaryTriangulation(model)
dΓ = Measure(Γ,degree)

a(u,v) = ∫(∇(u)⋅∇(v))*dΩ + ∫(u*v)*dΓ
ℓ(v) = ∫(v)*dΩ

mat_contribs = a(du,dv)
vec_contribs = ℓ(dv)

assem = SparseMatrixAssembler(U,V)

@test isa(U,TrialFESpace)
#@test_throws AssertionError assem = SparseMatrixAssembler(V,U)

data = collect_cell_matrix(U,V,mat_contribs)
A = assemble_matrix(assem,data)
@test size(A) == (num_free_dofs(V), num_free_dofs(U))

data = collect_cell_vector(V,vec_contribs)
b = assemble_vector(assem,data)
x = A\b
uh = FEFunction(U,x)

data = collect_cell_matrix_and_vector(U,V,mat_contribs,vec_contribs)
A,b = assemble_matrix_and_vector(assem,data)
x = A\b
uh = FEFunction(U,x)

uhd = zero(U)
data = collect_cell_matrix_and_vector(U,V,mat_contribs,vec_contribs,uhd)
A,b = assemble_matrix_and_vector(assem,data)
x = A\b
uh = FEFunction(U,x)

A = assemble_matrix(a,U,V)
b = assemble_vector(ℓ,V)
A,b = assemble_matrix_and_vector(a,ℓ,U,V)

A = assemble_matrix(a(du,dv),U,V)
b = assemble_vector(ℓ(dv),V)
A,b = assemble_matrix_and_vector(a(du,dv),ℓ(dv),U,V)

V = TestFESpace(
  model,
  ReferenceFE(lagrangian,Float64,1),
  vector_type=Vector{ComplexF64})
U = V

assem = SparseMatrixAssembler(U,V)
Ta = get_matrix_type(assem)
Tb = get_vector_type(assem)
@test eltype(Ta) == ComplexF64
@test eltype(Tb) == ComplexF64

# Now with an homogeneous linear form

a(u,v) = ∫(∇(u)⋅∇(v))*dΩ + ∫(u*v)*dΓ
ℓ(v) = 0

mat_contribs = a(du,dv)
vec_contribs = ℓ(dv)

data = collect_cell_matrix(U,V,mat_contribs)
A = assemble_matrix(assem,data)
@test size(A) == (num_free_dofs(V), num_free_dofs(U))

data = collect_cell_vector(V,vec_contribs)
b = assemble_vector(assem,data)
x = A\b
@test x ≈ b

data = collect_cell_matrix_and_vector(U,V,mat_contribs,vec_contribs)
A,b = assemble_matrix_and_vector(assem,data)
x = A\b
@test x ≈ b

uhd = zero(U)
data = collect_cell_matrix_and_vector(U,V,mat_contribs,vec_contribs,uhd)
A,b = assemble_matrix_and_vector(assem,data)
x = A\b
@test x ≈ b

end # module


"""
    (fx,ft)=fft3d(z)

Compute spatial (`fx`) and temporal (`ft`) Fourier spectra of field `z`.
"""
    function fft3d(z)

        nss = size(z); ns = nss[1]
        if (length(nss)>=3)
            nt=nss[3]
	else
	    nt=1
        end

        ns2 = div(ns,2)
        nt2 = div(nt,2)
        if ((nt>1)&&(mod(nt,2)!=0))
            z = z[:,:,1:(nt-1)]
            nt = nt-1
        end

        zf = abs.(fft(z)/(ns*ns*nt)).^2
        zf0 = zf
        zf[ns2+1,:,:] = zf0[ns2+1,:,:]/2
        zf[:,ns2+1,:] = zf0[:,ns2+1,:]/2
        if (nt>1)
          zf[:,:,nt2+1] = zf0[:,:,nt2+1]/2
        end

        fs = reshape(sum(zf, dims=3), ns, ns)
        ft = reshape(sum(sum(zf, dims=1), dims=2), nt, 1)

        fs = fftshift(fs)
        fs = fs/nt

        nn = zeros(ns, 1)
        fx = zeros(ns, 1)
        for j = -ns2:ns2-1
           for i = -ns2:ns2-1
              k2 = sqrt(i*i+j*j)
              ik2 = floor(Integer, k2+1.5)
              if (ik2>1)
                fx[ik2] = fx[ik2]+fs[i+ns2+1,j+ns2+1]
                nn[ik2] = nn[ik2]+1
              end
           end
        end
        if (nt>1)
            ft=ft[2:nt2+1]
        end
        fx = fx[2:ns2+1]./nn[2:ns2+1]
        return (fx, ft, fs)
    end


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ 6764c13e-681b-11eb-3527-77caf9d74f5b
using PlutoUI

# ╔═╡ 63a5111a-6817-11eb-2151-fb825d7879ba
using Test

# ╔═╡ d24d5ea0-6816-11eb-17d8-b5509f40beca
md"""
### Common Denominator

**Common denominators**

You will have a list of rationals in the form:
 [ ($numer_1$, $denom_1$) , ... ($numer_n$, $denom_n$) ]

where all numbers are positive integers.

You have to produce a result in the form
 ($n_1$, $d$) ... ($n_n$, $d$), in which $d$ is as small as possible and
         $\frac{n_1}{d} \equiv \frac{numer_1}{denom_1}$ ... $\frac{n_n}{d} \equiv \frac{numer_n}{denom_n}$.

Example:

`convert_fracs([(1, 2), (1, 3), (1, 4)]) should be [(6, 12), (4, 12), (3, 12)]`

Note:

- Due to the fact that first translations were written long ago - more than 4 years - these translations have only irreducible fractions.
- Newer translations have some reducible fractions. To be on the safe side it is better to do a bit more work by simplifying fractions even if they don't have to be.
"""

# ╔═╡ 69309f32-6817-11eb-0e36-2788da6513af
const VT{T} = Vector{Tuple{T, T}} where T

# ╔═╡ 7bdcddec-6817-11eb-09a5-53e299b8771f
begin

function convert_fracs(frac_lst::VT{T})::VT{T} where T <: Integer
    @assert size(frac_lst, 1) > 0

    ## 1 - simplify all fractions if possible
    frac_lst = map(t -> _simplify(t...), frac_lst)

    size(frac_lst, 1) == 1 && return frac_lst

    ## 2 - determine common denominator
    den_prod = find_common_den(frac_lst)

    ## 3 - build new fractions
	[(num * (den_prod / den), den_prod) for (num, den) in frac_lst]
end

function _simplify(num::T, den::T)::Tuple{T, T} where T <: Integer
	gcd = _gcd(num, den)
	(num ÷ gcd, den ÷ gcd)
end

"""
greatest common divisor

ex.
    a=24, b=36 => 12
    a=2, b=8 => 2
    a=3, b=2 => 1
    a=9, b=6 => 3
"""
function _gcd(a::T, b::T)::T where T <: Integer
	a, b =  a > b ? (a, b) : (b, a)
	while true
        r = a % b
		r == 0 && (return b)
		r == 1 && (return r)
		a, b = (b, r)
    end
end

"""
lowest common multiple

ex.
    a=4, b=3 => 12
    a=2, b=8 => 8
    a=9 (3*3), b=6 (2*3) => 18
"""
function _lcm(a::T, b::T)::T where T <: Integer
    (a * b) ÷ _gcd(a, b)
end

function find_common_den(frac_lst::VT{T})::T where T <: Integer
    den_ary = [den for (_, den) in frac_lst]
    com_den = _lcm(den_ary[1], den_ary[2])

    if size(frac_lst, 1) > 2
		com_den = foldl((com_den, cur_den) -> _lcm(com_den, cur_den),
			den_ary[3:end];
			init=com_den
		)
    end
    return com_den
end

end

# ╔═╡ 07856cdc-681b-11eb-28d3-3b7d32a3afd1
begin
	@test _gcd(3, 4) == 1
	@test _gcd(4, 3) == 1
	@test _gcd(24, 36) == 12
	@test _gcd(8, 2) == 2
	@test _gcd(3, 2) == 1
	@test _gcd(9, 6) == 3
end

# ╔═╡ 1e4f2bce-681b-11eb-00cb-1baae1e5d600
begin
	@test _lcm(3, 4) == 12
	@test _lcm(4, 3) == 12

 	@test _lcm(1, 6) == 6
 	@test _lcm(6, 1) == 6

	@test _lcm(6, 3) == 6
	@test _lcm(3, 5) == 15
	@test _lcm(5, 3) == 15

 	@test _lcm(2, 7) == 14
 	@test _lcm(7, 2) == 14

 	@test _lcm(2*3, 2*2) == 12
end

# ╔═╡ 4c877dc8-681b-11eb-2f8d-bb469d56f840
with_terminal() do
	println(convert_fracs([(1, 2), (1, 3), (1, 4)]))
end

# ╔═╡ 952b7954-681c-11eb-0d11-b7ba99db8b64
begin
	@test convert_fracs([(2, 4)]) == [(1, 2)]
	@test convert_fracs([(1, 2), (1, 3), (1, 4)]) == [(6, 12), (4, 12), (3, 12)]
	@test convert_fracs([(2, 4), (2, 6), (2, 8)]) == [(6, 12), (4, 12), (3, 12)]

	@test convert_fracs([(1, 2), (1, 3)]) == [(3, 6), (2, 6)]

	@test convert_fracs([(1, 3), (1, 2)]) == [(2, 6), (3, 6)]
end

# ╔═╡ Cell order:
# ╟─d24d5ea0-6816-11eb-17d8-b5509f40beca
# ╠═6764c13e-681b-11eb-3527-77caf9d74f5b
# ╠═63a5111a-6817-11eb-2151-fb825d7879ba
# ╠═69309f32-6817-11eb-0e36-2788da6513af
# ╠═7bdcddec-6817-11eb-09a5-53e299b8771f
# ╠═07856cdc-681b-11eb-28d3-3b7d32a3afd1
# ╠═1e4f2bce-681b-11eb-00cb-1baae1e5d600
# ╠═4c877dc8-681b-11eb-2f8d-bb469d56f840
# ╠═952b7954-681c-11eb-0d11-b7ba99db8b64


using Vectorize

## Use correct version of Base.Test
if VERSION >= v"0.5-"
    using Base.Test
else
    using BaseTestNext
    const Test = BaseTestNext
end

## Initialize testing globals
println("===== Testing Accelerate =====")
N = 1000
srand(13)

## LOGARITHM
for T in (Float32, Float64)
    @testset "Accelerate: Logarithmic::$T" begin
        X::Array{T} = exp(10*randn(N))
        @testset "Testing $f::$T" for f in [:log,:log2,:log10, :log1p]
            @eval fb = $f
            @eval fa = Vectorize.Accelerate.$f
            @test fa(X) ≈ fb(X)
        end
    end
end

## EXPONENIATION
for T in (Float32, Float64)
    @testset "Accelerate: Exponential::$T" begin
        @testset "Testing $f::$T" for f in [:exp,:exp2,:expm1]
            X::Array{T} = 10*randn(N)
            @eval fb = $f
            @eval fa = Vectorize.Accelerate.$f
            @test fa(X) ≈ fb(X)
        end
    end
end

## TRIGONOMETRIC
for T in (Float32, Float64)
    X::Array{T} = 10*randn(N)
    @testset "Accelerate: Trigonometric::$T" begin
        @testset "Testing $f::$T" for f in [:sin,:sinpi,:cos,:cospi,:tan,:atan] # tanpi not defined in Base
            @eval fb = $f
            @eval fa = Vectorize.Accelerate.$f
            @test fa(X) ≈ fb(X)
        end

        Y::Array{T} = 10*randn(N)
        @testset "Testing $f::$T" for f in [:atan2]
            @eval fb = $f
            @eval fa = Vectorize.Accelerate.$f
            @test fa(X,Y) ≈ fb(X,Y)
        end

        Z::Array{T} = 2*rand(N)-1
        @testset "Testing $f::$T" for f in [:asin,:acos]
            @eval fb = $f
            @eval fa = Vectorize.Accelerate.$f
            @test fa(Z) ≈ fb(Z)
        end
    end
end

## HYPERBOLIC
for T in (Float32, Float64)
    @testset "Accelerate: Hyperbolic::$T" begin
        X = 10*randn(N)
        @testset "Testing $f::$T" for f in [:sinh,:cosh,:tanh,:asinh]
            @eval fb = $f
            @eval fa = Vectorize.Accelerate.$f
            @test fa(X) ≈ fb(X)
        end

        Y = exp(10*randn(N))+1
        @testset "Testing $f::$T" for f in [:acosh]
            @eval fb = $f
            @eval fa = Vectorize.Accelerate.$f
            @test fa(Y) ≈ fb(Y)
        end

        Z = 2*rand(N)-1
        @testset "Testing $f::$T" for f in [:atanh]
            @eval fb = $f
            @eval fa = Vectorize.Accelerate.$f
            @test fa(Z) ≈ fb(Z)
        end
    end
end

## VECTOR OPERATIONS
for T in (Float32, Float64)
    @testset "Accelerate: Vector Operations::$T" begin
        X = 10*randn(N)
        Y = 10*randn(N)
        @testset "Testing $fa::$T" for (f, fa) in [(:+, :add), (:-, :sub), (:.*, :mul), (:./, :div)]
            @eval fb = $f
            @eval facc = Vectorize.Accelerate.$fa
            @test fb(X, Y) ≈ facc(X, Y)
        end
    end
end

println("===== Accelerate Tests Successful =====\n\n")


using LightGraphs, JuMP, Cbc, Plots
gr(size=(600,600))

function dollar_game(g, W)
  nv = length(vertices(g))

  model = Model(solver = CbcSolver())
  @variable(model, give[1:nv] >= 0, Int)
  @variable(model, take[1:nv] >= 0, Int)
  @objective(model, Min, sum(give[i] + take[i] for i = 1:nv))
  @expression(model, x[i=1:nv], W[i] +
              (take[i] - give[i]) * length(neighbors(g, i)) +
              sum(give[j] - take[j] for j = neighbors(g, i)))
  @constraint(model, x .>= 0)
  status = solve(model)

  println("status = $status")
  if status != :Optimal
    return zeros(nv), W
  end

  give = Int.(getvalue(give))
  take = Int.(getvalue(take))
  sol = Int.(getvalue(x))
  return give - take, sol
end

function plot_common!(p, g, W, game=[])
  nv = length(vertices(g))
  θ = 2π / nv
  t = (1:nv) * θ - 1.5θ
  r = 2.5
  x, y = r * sin.(t), r * cos.(t)

  for i = 1:nv
    for j = neighbors(g, i)
      plot!([x[i], x[j]], [y[i], y[j]], lw=2, c=:black)
    end
  end
  scatter!(x, y, ann=[(x[i],y[i],W[i]) for i = 1:nv], m=(30,:white,stroke(2,:black)))
  if length(game) > 0
    annotate!([(1.2x[i], 1.2y[i], game[i], :red) for i = 1:nv])
  end
  xlims!(-1.3r, 1.3r)
  ylims!(-1.3r, 1.3r)
  return p
end

function plot_common(g, W, game=[])
  p = plot(leg=false, axis=false, grid=false)
  return plot_common!(p, g, W, game)
end

function plot_graph(g, W, filename, game=[])
  p = plot_common(g, W, game)
  png(filename)
end

function plot_game(g, W, prefix, game)
  p = plot_common(g, W)
  nv = length(vertices(g))
  k = 0
  fname = @sprintf("%s-%03d", prefix, k)
  title!("Initial")
  png(p, fname)

  θ = 2π / nv
  t = (1:nv) * θ - 1.5θ
  r = 2.5
  x, y = r * sin.(t), r * cos.(t)

  for i = 1:nv
    for n = 1:abs(game[i])
      k += 1
      ng = neighbors(g, i)
      s = sign(game[i])
      W[i] -= s * length(ng)
      c = s > 0 ? RGB(0.1, 0.4, 0.9) : :red

      W[ng] .+= s

      p = plot(leg=false, axis=false, grid=false)
      plot_common!(p, g, W)
      for j = ng
        v = [x[j] - x[i]; y[j] - y[i]]
        σ = 0.4 / norm(v)
        X = [x[i] + v[1] * σ; x[j] - v[1] * σ]
        Y = [y[i] + v[2] * σ; y[j] - v[2] * σ]
        if s < 0
          X, Y = X[[2;1]], Y[[2;1]]
        end
        plot!(X, Y, c=c, lw=3, l=:arrow)
      end
      scatter!([x[i]], [y[i]], m=(30,c,stroke(3,:black)))
      scatter!([x[ng]], [y[ng]], m=(30,c,stroke(3,:black)))
      I = setdiff(1:nv, [i; ng])
      scatter!([x[I]], [y[I]], m=(30,:white,stroke(2,:black)))

      fname = @sprintf("%s-%03d", prefix, k)
      title!("Move $k")
      png(fname)
    end
  end
end

function example()
  g = Graph(5)
  W = [1; -2; 3; 2; -1]
  add_edge!(g, 1, 2)
  add_edge!(g, 1, 3)
  add_edge!(g, 2, 3)
  add_edge!(g, 2, 5)
  add_edge!(g, 3, 4)

  plot_graph(g, W, "example")
  game, sol = dollar_game(g, W)
  plot_graph(g, sol, "example-sol", game)
  plot_game(g, W, "example", game)
end

function example2()
  g = Graph(3)
  W = [0; 1; -1]
  add_edge!(g, 1, 2)
  add_edge!(g, 1, 3)
  add_edge!(g, 2, 3)

  game, sol = dollar_game(g, W)
end

function example3()
  g = Graph(3)
  W = [-2; -1; 4]
  add_edge!(g, 1, 2)
  add_edge!(g, 2, 3)

  game, sol = dollar_game(g, W)
end

example()
example2()
example3()


#= Ornstein-Uhlenbeck process
   without parameter estimation

The Ornstein-Uhlenbeck (OU) process is a continuous-time continuous-space
stochastic process governed by the stochastic differential equation (SDE)

          dY(t) = (Θ_1 - Θ_2 Y(t)) dt + Θ_3 dW(t)

where W(t) is the Wiener process (standard Brownian motion). This OU process is
stationary, Gaussian and Markovian. It is mean-reverting in the sense that it
drifts towards its long-term mean of Θ_1 / Θ_2.

For a step size Δt, the transition distribution can be obtained analytically as
a Normal with mean μ(Δt) and variance σ^2(Δt):

      Y(t + Δt) | Y(t) ~ Normal(μ(Δt), σ^2(Δt))
                 μ(Δt) = Θ_1 / Θ_2 + (Y(t) - Θ_1 / Θ_2) exp(- Θ_2 Δt)
               σ^2(Δt) = Θ_3^2 / (2 * Θ_2) * (1 - exp(- 2 Θ_2 Δt))

Given a sequence of noisy observations of the process Y(t_n) ≈ y_n for
1 <= n <= N, the task is to infer the posterior distribution over each Y(t_n).
This choice is for the sake of simplicity, the example can be modified to
infer posterior distributions over any Y(t).

This example has been inspired by the LibBi example at
http://libbi.org/packages/OrnsteinUhlenbeckBridge.html. =#

using Turing

#= The parameter values of Θ_1, Θ_2, Θ_3 follow [Aït-Sahalia, 1999],
   [Sun et al., 2013] and [Del Moral et al., 2015]. =#
const Θ_1 = 0.0187
const Θ_2 = 0.2610
const Θ_3 = 0.0224

#= [Del Moral, 2015] observe the OU process directly with no observation noise,
and their task is to simulate the diffusion bridges between the observed values.
Our examples uses observation noise with standard deviation σ_y > 0, as this is
required for inference to work at this stage. =#
const σ_y = 0.01

#= The observation times ts and the corresponding observed values.
   It is assumed that the observation times are sorted increasingly. =#
ts      = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
data_ys = [0.0, -0.009114765844660457, 0.029830170268001097,
0.021755657377028295, 0.013934983652697013, 0.01910503829939146,
0.049404191951179434, 0.03253467537333853, 0.07701445971172535,
0.04618793002053633, 0.04001971135990065]
N = length(ts)

# The transition distribution of the OU process
function trans(Yt, Δt)
  mean = Θ_1 / Θ_2 + (Yt - Θ_1 / Θ_2) * exp(- Θ_2 * Δt)
  variance = Θ_3 * Θ_3 / (2 * Θ_2) * (1 - exp(- 2 * Θ_2 * Δt))
  # Normal is initialized with mean and standard deviation, not variance!
  return Normal(mean, sqrt(variance))
end

# The generative model of the OU process
@model ou begin
  # use TArray to support array copying when particles are duplicated!
  Y = TArray(Float64, N)
  Y[1] = 0.0
  for n = 2:N
    Δt = ts[n] - ts[n-1]
    @assume(Y[n] ~ trans(Y[n-1], Δt))
    @observe(data_ys[n] ~ Normal(Y[n], σ_y))
  end
  @predict Y
end

#= This concludes the example. The remainder of this script verifies that
inference works properly by checking that (potentially weighted) samples from
the posterior over Y[N] (the last element of the sequence Y) approximate the
exact posterior over Y[N].

The exact posterior over Y[N] is obtained by recalling that our OU process
is a Gaussian process [Rasmussen and Williams, 2005] with constant mean function
μ(t) = Θ_1 / Θ_2 and stationary covariance function

          k(r) = Θ_3^2 / (2 * Θ_2) * exp(- Θ_2 * |r|)

which is also known as the exponential, Laplace, or Laplacian kernel. (See
Appendix B.2.1 in [Rasmussen and Williams, 2005] for a derivation.) Hence exact
posteriors can be obtained by performing standard GP regression, which involves
inverting a kernel matrix A := (K + σ_y^2 I). =#

k(x, y) = Θ_3^2 / (2 * Θ_2) * exp(- Θ_2 * abs(x - y))
K = [k(x1, x2)::Float64 for x1=ts, x2=ts]
noise_matrix = σ_y^2 * eye(N)
noise_matrix[1, 1] = 0.0 # observed 0.0 at time t=0 with no noise
A = K + noise_matrix
Ainv = inv(A)

# Compute true posterior over Y[N], the value of Y(t) at time t=ts[N]
x_star = ts[N]
k_star = [k(x, x_star)::Float64 for x=ts]
f_star_mean = transpose(k_star) * (A \ (data_ys - Θ_1 / Θ_2)) + Θ_1 / Θ_2
f_star_variance = k(x_star, x_star) - transpose(k_star) * (A \ k_star)
f_star_std = sqrt(f_star_variance)
f_star_posterior = Normal(f_star_mean[1], f_star_std[1])
println("True posterior of last state: ", f_star_posterior)

# KL-divergence between a Normal MLE-fitted to samples and true posterior
function ou_divergence(samples, weights)
  fitted = fit_mle(Normal, samples, weights)
  return Turing.kl(fitted, f_star_posterior)
end
function ou_divergence(samples :: Vector{Float64})
  weights = fill(Float64(1), length(samples))
  return ou_divergence(samples, weights)
end

#= Compute the KL-divergence between moment-matched Normal from samples and
exact posterior, as the number of samples increases =#
num_samples = [10, 30, 100, 300, 1000, 3000, 10000]
samples = sample(ou, SMC(maximum(num_samples)))
samples_lasty = [res[N]::Float64 for res=samples[:Y]]
kl_divergences = map(n -> ou_divergence(samples_lasty[1:n]), num_samples)

# Plot the results using Gadfly
using Gadfly
plot(x=num_samples, y=kl_divergences,
     Scale.x_log10, Scale.y_log10,
     Guide.title("KL-divergence between exact and moment matched posterior"),
     Guide.xlabel("number of samples"), Guide.ylabel("KL-divergence"))

#= References:
- Y. Aït-Sahalia. Transition densities for interest rate and other nonlinear
  diffusions. The Journal of Finance, 54(4):1361–1395, 1999. ISSN 1540-6261.
  doi: 10.1111/0022-1082.00149.
- L. Sun, C. Lee, and J. A. Hoeting. Penalized importance sampling for parameter
  estimation in stochastic differential equations. 2013.
  URL http://arxiv.org/abs/1305.4390.
- Del Moral, Pierre, and Lawrence M. Murray. "Sequential Monte Carlo with highly
  informative observations." SIAM/ASA Journal on Uncertainty Quantification 3,
  no. 1 (2015): 969-997. http://arxiv.org/pdf/1405.4081v2.pdf
- C. E. Rasmussen and C. K. I. Williams. Gaussian Processes for Machine Learning
  (Adaptive Computation and Machine Learning). The MIT Press, 2005.
  ISBN 978-0-262-18253-9. URL http://www.gaussianprocess.org/gpml/.
=#


vandercorput(num::Integer, base::Integer) = sum(d * Float64(base) ^ -ex for (ex,d) in enumerate(digits(num, base=base)))

function quasirandom(x, y;order=:none, squish=0.9)

	if order == :none
		grps = unique(x)
	else
		grps = order
	end

	nGroups = length(grps)
	width= squish/nGroups

	k = Array{UnivariateKDE}(undef, nGroups)
	maxDensity = zeros(nGroups)

	xTemp = zeros(length(x))
	for i in 1:nGroups
		xTemp[findall(elements -> elements == grps[i], x)] .= i
	end

	q = zeros(length(x))
	d = zeros(length(x))
	for i in 1:nGroups
		k[i] = kde(y[x .== grps[i]])
		d[x .== grps[i]] .= pdf(k[i], y[x .== grps[i]])
		maxDensity[i] = maximum(d[x .== grps[i]])
		q[x .== grps[i]] .= vandercorput.(1:length(q)
			, 2)[competerank(y[x .== grps[i]])]
	end

	x_jitter = xTemp .+ width./maxDensity[nGroups] .*  (q .- 0.5) .* 2 .* d

	return x_jitter .- 0.5
end

function quasirandomScatter(x, y; order=:none, squish=0.9, kwargs...)

	plt = scatter(x, y; markercolor=:white, markerstrokewidth=0,
		xlim=(0, length(unique(x))), margin=2mm)
	scatter!(quasirandom(x, y;order=order, squish), y; kwargs...)

	return plt
end


# by lazarusA # HIDE
# using GLMakie # HIDE
using GLMakie, NetCDF
using YAXArrays.Datasets: open_mfdataset
GLMakie.activate!() # HIDE
let
    dataset = open_mfdataset("./data/ETOPO1_halfdegree.nc")
    lon = dataset.lon.values
    lat = dataset.lat.values
    data = dataset.ETOPO1avg.data[:,:]
    function toCartesian(lon, lat; r = 1, cxyz = (0,0,0) )
        lat, lon = lat*π/180, lon*π/180
        x = cxyz[1] + r * cos(lat) * cos(lon)
        y = cxyz[2] + r * cos(lat) * sin(lon)
        z = cxyz[3] + r *sin(lat)
        x,y,z
    end
    function lonlat3D(lon, lat, data; cxyz = (0,0,0))
        xyzw = zeros(size(data)..., 3)
        for (i,lon) in enumerate(lon), (j,lat) in enumerate(lat)
            x, y, z = toCartesian(lon, lat; cxyz = cxyz)
            xyzw[i,j,1] = x
            xyzw[i,j,2] = y
            xyzw[i,j,3] = z
        end
        xyzw[:,:,1], xyzw[:,:,2], xyzw[:,:,3]
    end
    # this is needed in order to have a closed surface
    lonext  = cat(collect(lon), lon[1], dims = 1)
    dataext = begin
        tmpdata = zeros(size(lon)[1]+1,  size(lat)[1])
        tmpdata[1:size(lon)[1],:] = data
        tmpdata[size(lon)[1]+1,:] = data[1,:]
        tmpdata
    end
    #now let's plot that surface!
    xetopo, yetopo, zetopo = lonlat3D(lonext, lat, dataext)

    fig = Figure(resolution =(950,900),fontsize = 25,backgroundcolor = :black)
    ax = Axis3(fig, aspect = :data, azimuth = -0.2π, elevation = -0.05π)
    pltobj=surface!(ax,xetopo, yetopo, zetopo, color = dataext,
        colormap = :hot, backlight = 1.0f0, colorrange =(-6000,5000))
    cbar  =Colorbar(fig,pltobj,label = "ETOPO1 [m]",ticklabelsize = 20,
        tickalign =1, vertical = true, ticksize=15, labelcolor = :white,
        ticklabelcolor = :white, tickcolor = :grey90, height = Relative(0.5))
    hidespines!(ax)
    hidedecorations!(ax)
    fig[1,1] = ax
    fig[1,2] = cbar
    fig
    save(joinpath(@__DIR__, "output", "topographySphere.png"), fig, px_per_unit = 2.0) # HIDE
end

using Pkg # HIDE
Pkg.status(["GLMakie","NetCDF", "YAXArrays"]) # HIDE


# ******************************************************************************
# PACKAGES
# ******************************************************************************
using LinearAlgebra
using Random
using Distributions

# ******************************************************************************
# OPTIMIZATION FUNCTION DEFNITIONS
# ******************************************************************************
# A function for constructing a basis vector
basis(i,n) = [k==i ? 1.0 : 0.0 for k in 1:n]

# A function to perform a single step of the Hooke-Jeeves optimization method
function hooke_jeeves(f, x, y, α, n, γ = 0.5)
    improved = false
    x_best, y_best = x, y
    for i in 1:n
        for sign in (-1, 1)
            x′ = x + sign*α*basis(i,n)
            y′ = f(x′)
            if y′ < y_best
                x_best, y_best, improved = x′, y′, true
            end
        end
    end
    x, y = x_best, y_best
    if !improved
        α *= γ
    end
    return x, y, α
end

# A function to randomly sample a positive spanning set according to mesh
# adaptive direct search
function random_positive_spanning_set(α, n)
    δ = round(Int, 1/sqrt(α))
    lower_triangular_matrix = Matrix(Diagonal(δ*rand([1,-1], n)))
    for i in 1:(n-1)
        for j in 1:(i-1)
            lower_triangular_matrix[i,j] = rand(-δ + 1:(δ-1))
        end
    end
    D = lower_triangular_matrix[randperm(n),:][:, randperm(n)]
    D = hcat(D, -sum(D,dims=2))
    return [D[:,i] for i in 1:(n+1)]
end

# A function to perform a single step of mesh adaptive direct search
function mesh_adaptive_direct_search(f, x, y, α, n)
    improved = false
    for (i,d) in enumerate(random_positive_spanning_set(α, n))
        x′ = x + α*d; y′ = f(x′)
        if y′ < y
            x, y, improved = x′, y′, true
            x′ = x + 3*α*d; y′ = f(x′)
            if y′ < y
                x, y = x′, y′
            end
            break
        end
    end
    α = improved ? min(2*α, 1) : α/2
    return x, y, α
end

# A function to perform a single step of the cross entropy method
function cross_entropy(f, P, m=100, m_elite=10)
    samples = rand(P,m)
    order = sortperm([f(samples[:,i]) for i in 1:m])
    P = fit(typeof(P), samples[:, order[1:m_elite]])
    return P
end

# A penalty method mixture between a count and a quadratic penalty function
function penalty(c, x, ρ1, ρ2, δ)
    evals = c(x)
    max_arr = [max(evals[i] + δ, 0) for i = 1:length(evals)]
    greater_than_zero = [evals[i] > 0 for i = 1:length(evals)]
    max_arr = max_arr.^2
    return ρ1*sum(max_arr) + ρ2*sum(greater_than_zero)
end

# HJ parameters: ρ1=1.01, ρ2=0, γ=1.001, δ=0.5; 0.1378383301
# MADS parameters: ρ1=1.01, ρ2=0, γ=1.001, δ=0.5; 0.1395356843
function optimize_simple1(f, g, c, x, num_evals; ρ1=1.01, ρ2=0, γ=1.001, δ=0.5)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-8)
        #x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        ρ1 *= γ; ρ2 *= γ
    end
    return x
end

# HJ parameters: ρ1=100, ρ2=1, γ=1.1, δ=0.75; 1.4993756041
# MADS parameters: ρ1=100, ρ2=1, γ=1.01, δ=1.0; 1.5604945284
function optimize_simple2(f, g, c, x, num_evals; ρ1=100, ρ2=1, γ=1.1, δ=0.75)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-8)
        #x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        ρ1 *= γ; ρ2 *= γ
    end
    return x
end

# HJ parameters: ρ1=1, ρ2=1, γ=1.01, δ=0.5; 0.0631003421
# MADS parameters: ρ1=1, ρ2=1, γ=1.01, δ=0.25; 0.0425930332
function optimize_simple3(f, g, c, x, num_evals; ρ1=1, ρ2=1, γ=1.01, δ=0.25)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-13)
        x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        #x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        ρ1 *= γ; ρ2 *= γ
    end
    return x
end

# ρ1=50, ρ2=1, γ=1.01, δ=0.15, m = 100; num_elite = 20, μ=x; Σ=10; -0.5935202993
function optimize_secret1(f, g, c, x, num_evals; ρ1=50, ρ2=1, γ=1.01, δ=0.15)
    μ = x; Σ = 10; P = MvNormal(μ, Σ)
    m = 100; num_elite = 20
    while count(f, g, c) < (num_evals - 2*m)
        P = cross_entropy(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), P, m, num_elite)
        ρ1 *= γ; ρ2 *= γ
    end
    return P.μ
end

# HJ parameters:  ρ1=1, ρ2=1, γ=1.001, δ=0.15; 0.657090133
# MADS parameters:  ρ1=1, ρ2=1, γ=1.001, δ=0.15; 0.3410080274
function optimize_secret2(f, g, c, x, num_evals; ρ1=1, ρ2=1, γ=1.001, δ=0.15)
    α, y, n, ϵ = 0.1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals - 308)
        x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        #x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        ρ1 *= γ; ρ2 *= γ
    end
    return x
end

function test_call()
    print("HELLO WORLD")
end

# *******************************************************************************
# MAIN OPTIMIZATION FUNCTION
# *******************************************************************************
function optimize(f, g, c, x0, n, prob)
    if prob == "simple1"
        x_best = optimize_simple1(f, g, c, x0, n)
    elseif prob == "simple2"
        x_best = optimize_simple2(f, g, c, x0, n)
    elseif prob == "simple3"
        x_best = optimize_simple3(f, g, c, x0, n)
    elseif prob == "secret1"
        x_best = optimize_secret1(f, g, c, x0, n)
    elseif prob == "secret2"
        x_best = optimize_secret2(f, g, c, x0, n)
    else
        print("Invalid Problem")
    end
end


"""
    InterpolateAndBin{L}(f::Function, left_bin_edges, intp::L, intp_grid,
        extrapolation_bc::Union{<:Real, Interpolations.BoundaryCondition})

Indicates that a dataset consisting of both indices and values should first
be interpolated to the `intp_grid` grid using the provided `intp` scheme
(e.g. `Linear()`). After interpolating, assign the interpolated values
to the bins defined by `left_bin_edges` and summarise the values falling
in each bin using the summary function `f` (e.g. `mean`).

## Example

```julia
using UncertainData, Interpolations, StatsBase

# Assume we have the following unevenly spaced data with some `NaN` values
T = 100
time = sample(1.0:T*5, T, ordered = true, replace = false)
y1 = rand(T)
time[rand(1:T, 10)] .= NaN
y1[rand(1:T, 10)] .= NaN

# We want to first intepolate the dataset linearly to a regular time grid
# with steps of `0.1` time units.
intp = Linear()
intp_grid = 0:0.1:1000
extrapolation_bc = Flat(OnGrid())

# Then, bin the dataset in time bins `50` time units wide, collect all
# values in each bin and summarise them using `f`.
f = mean
left_bin_edges = 0:50:1000

r = InterpolateAndBin(f, left_bin_edges, intp, intp_grid, extrapolation_bc)
```
"""
struct InterpolateAndBin{L}
    f::Function
    left_bin_edges
    intp::L # Linear
    intp_grid
    extrapolation_bc::Union{<:Real, Interpolations.BoundaryCondition}
end

export InterpolateAndBin


using Flux, Zygote, LinearAlgebra, Statistics
println("Heat han")
using Test, StochasticDiffEq

using NeuralPDE

using Random
Random.seed!(100)

opt = Flux.ADAM(0.005)  #optimizer
# high-dimensional heat equation
d = 30 # number of dimensions
x0 = fill(8.0f0,d)
tspan = (0.0f0,2.0f0)
dt = 0.5
time_steps = div(tspan[2]-tspan[1],dt)
m = 30 # number of trajectories (batch size)

g(X) = sum(X.^2)
f(X,u,σᵀ∇u,p,t) = 0.0
μ_f(X,p,t) = 0.0
σ_f(X,p,t) = 1.0
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,hls,relu),
                  Dense(hls,hls,relu),
                  Dense(hls,d)) for i in 1:time_steps]

alg = NNPDEHan(u0, σᵀ∇u, opt = opt)

ans = solve(prob, alg, verbose = true, abstol=1e-8, maxiters = 1000, dt=dt, trajectories=m)

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

println("high-dimensional heat equation")
# println("numerical = ", ans)
# println("analytical = " ,analytical_ans)
println("error_l2 = ", error_l2, "\n")
@test error_l2 < 1.0


module GaussQuadrature

using SpecialFunctions: gamma, lgamma

# October 2013 by Bill McLean, School of Maths and Stats,
# The University of New South Wales.
#
# Based on earlier Fortran codes
#
# gaussq.f original version 20 Jan 1975 from Stanford
# gaussq.f modified 21 Dec by Eric Grosse
# gaussquad.f95 Nov 2005 by Bill Mclean
#
# This module provides functions to compute the abscissae x[j] and
# weights w[j] for the classical Gauss quadrature rules, including
# the Radau and Lobatto variants.  Thus, the sum
#
#           n
#           ∑  w[j] f(x[j])
#          j=1
#
# approximates
#
#           hi
#          ∫  f(x) w(x) dx
#          lo
#
# where the weight function w(x) and interval lo < x < hi are as shown
# in the table below.
#
# Name                      Interval     Weight Function
#
# Legendre                 -1 < x < 1          1
# Chebyshev (first kind)   -1 < x < 1     1 / sqrt(1-x²)
# Chebyshev (second kind)  -1 < x < 1       sqrt(1-x²)
# Jacobi                   -1 < x < 1     (1-x)ᵅ (1+x)ᵝ
# Laguerre                  0 < x < ∞     xᵅ exp(-x)
# Hermite                  -∞ < x < ∞      exp(-x²)
#
# In addition to these classical rules, the module generates Gauss rules
# for logarithmic weights of the form
#
#    w(x) = x^ρ log(1/x)   for 0 < x < 1.
#
# For the Jacobi and Laguerre rules we require α > -1 and
# β > -1, so that the weight function is integrable.  Likewise, for
# log weight we require ρ > -1.
#
# Use the endpt argument to include one or both of the end points
# of the interval of integration as an abscissa in the quadrature
# rule, as follows.
#
# endpt = neither   Default      lo < x[j] < hi, j = 1:n.
# endpt = left      Left Radau   lo = x[1] < x[j] < hi, j = 2:n.
# endpt = right     Right Radau  lo < x[j] < x[n] = hi, j = 1:n-1.
# endpt = both      Lobatto      lo = x[1] < x[j] < x[n] = hi, j = 2:n-1.
#
# These labels make up an enumeration of type EndPt.
#
# The code uses the Golub and Welsch algorithm, in which the abscissae
# x[j] are the eigenvalues of a symmetric tridiagonal matrix whose
# entries depend on the coefficients in the 3-term recurrence relation
# for the othonormal polynomials generated by the weighted inner product.
#
# References:
#
#   1.  Golub, G. H., and Welsch, J. H., Calculation of Gaussian
#       quadrature rules, Mathematics of Computation 23 (April,
#       1969), pp. 221-230.
#   2.  Golub, G. H., Some modified matrix eigenvalue problems,
#       Siam Review 15 (april, 1973), pp. 318-334 (section 7).
#   3.  Stroud and Secrest, Gaussian Quadrature Formulas, Prentice-
#       Hall, Englewood Cliffs, N.J., 1966.

export neither, left, right, both
export legendre, legendre_coefs
export chebyshev, chebyshev_coefs
export jacobi, jacobi_coefs
export laguerre, laguerre_coefs
export hermite, hermite_coefs
export shifted_legendre_coefs
export logweight, logweight_coefs
export modified_moments, modified_chebyshev
export custom_gauss_rule, orthonormal_poly
export special_eigenproblem!

# Enumeration type used to specify which endpoints of the integration
# interval should be included amongst the quadrature points: neither,
# left, right or both.
struct EndPt
    label :: Char
end

const neither = EndPt('N')
const left    = EndPt('L')
const right   = EndPt('R')
const both    = EndPt('B')

#include("to_be_removed.jl")

# Maximum number of QL iterations used by steig!.
# You might need to increase this.
maxiterations = Dict(Float32 => 30, Float64 => 30, BigFloat => 40)

"""
    x, w = legendre(T, n, endpt=neither)

Returns points `x` and weights `w` for the `n`-point Gauss-Legendre rule
for the interval `-1 < x < 1` with weight function `w(x) = 1`.

Use `endpt=left`, `right` or `both` for the left Radau, right Radau or
Lobatto rules, respectively.
"""
function legendre(::Type{T}, n::Integer,
                  endpt::EndPt=neither) where {T<:AbstractFloat}
    @assert n ≥ 1
    a, b = legendre_coefs(T, n)
    return custom_gauss_rule(-one(T), one(T), a, b, endpt)
end

"""
Convenience function with type `T = Float64`:

    x, w = legendre(n, endpt=neither)
"""
legendre(n, endpt=neither) = legendre(Float64, n, endpt)

function legendre_coefs(::Type{T}, n::Integer) where {T<:AbstractFloat}
    a = zeros(T, n)
    b = zeros(T, n+1)
    b[1] = sqrt(convert(T, 2))
    for k = 2:n+1
        b[k] = (k-1) / sqrt(convert(T, (2k-1)*(2k-3)))
    end
    return a, b
end

"""
    x, w = chebyshev(T, n, kind=1, endpt=neither)

Returns points `x` and weights `w` for the `n`-point Gauss-Chebyshev rule
for the interval `-1 < x < 1` with weight function

    w(x) = 1 / sqrt(1-x²)   if kind=1
    w(x) = sqrt(1-x²)       if kind=2.

Use `endpt=left`, `right` or `both` for the left Radau, right Radau or
Lobatto rules, respectively.
"""
function chebyshev(::Type{T}, n::Integer, kind::Integer=1,
                   endpt::EndPt=neither) where {T<:AbstractFloat}
    @assert n ≥ 1
    a, b = chebyshev_coefs(T, n, kind)
    return custom_gauss_rule(-one(T), one(T), a, b, endpt)
end

"""
Convenience function with type `T = Float64`:

    x, w = chebyshev(n, kind=1, endpt=neither)

"""
chebyshev(n, kind=1, endpt=neither) = chebyshev(Float64, n, kind, endpt)

function chebyshev_coefs(::Type{T}, n::Integer,
                         kind::Integer) where {T<:AbstractFloat}
    half = convert(T, 1//2)
    a = zeros(T, n)
    b = fill(half, n+1)
    if kind == 1
        b[1] = sqrt(convert(T, pi))
	if n >= 2
            b[2] = sqrt(half)
        end
    elseif kind == 2
        b[1] = sqrt(half * convert(T, pi))
    else
        error("Unsupported value for kind")
    end
    return a, b
end

"""
    x, w = jacobi(n, α, β, endpt=neither)

Returns points `x` and weights `w` for the `n`-point Gauss-Jacobi rule
for the interval `-1 < x < 1` with weight function

    w(x) = (1-x)ᵅ (1+x)ᵝ.

Use `endpt=left`, `right` or `both` for the left Radau, right Radau or
Lobatto rules, respectively.
"""
function jacobi(n::Integer, α::T, β::T,
                endpt::EndPt=neither) where {T<:AbstractFloat}
    @assert n ≥ 1
    @assert α > -1.0 && β > -1.0
    a, b = jacobi_coefs(n, α, β)
    custom_gauss_rule(-one(T), one(T), a, b, endpt)
end

function jacobi_coefs(n::Integer, α::T, β::T) where {T<:AbstractFloat}
    a = zeros(T, n)
    b = zeros(T, n+1)
    ab = α + β
    abi = ab + 2
    b[1] = 2^((ab+1)/2) * exp(
             ( lgamma(α+1) + lgamma(β+1) - lgamma(abi) )/2 )
    a[1] = ( β - α ) / abi
    b[2] = sqrt( 4*(α+1)*(β+1) / ( (ab+3)*abi*abi ) )
    a2b2 = β*β - α*α
    for i = 2:n
        abi = ab + 2i
        a[i]   = a2b2 / ( (abi-2)*abi )
        b[i+1] = sqrt( 4i*(α+i)*(β+i)*(ab+i) /
                     ( (abi*abi-1)*abi*abi ) )
    end
    return a, b
end

"""
    x, w = laguerre(n, α, endpt=neither)

Returns points `x` and weights `w` for the `n`-point Gauss-Laguerre rule
for the interval `0 < x < ∞` with weight function

    w(x) = xᵅ exp(-x),   α > -1.

Use `endpt=left` for the left Radau rule.
"""
function laguerre(n::Integer, α::T,
                  endpt::EndPt=neither) where {T<:AbstractFloat}
    @assert n ≥ 1
    @assert α > -1.0
    @assert endpt in [neither, left]
    a, b = laguerre_coefs(n, α)
    custom_gauss_rule(zero(T), convert(T, Inf), a, b, endpt)
end

function laguerre_coefs(n::Integer, α::T) where {T<:AbstractFloat}
    a = zeros(T, n)
    b = zeros(T, n+1)
    b[1] = sqrt(gamma(α+1))
    for i = 1:n
        a[i] = 2i - 1 + α
        b[i+1] = sqrt( i*(α+i) )
    end
    return a, b
end

"""
    x, w = hermite(T, n)

Returns points `x` and weights `w` for the `n`-point Gauss-Laguerre rule
for the interval `-∞ < x < ∞` with weight function

    w(x) = exp(-x²).
"""
function hermite(::Type{T}, n::Integer) where {T<:AbstractFloat}
    @assert n ≥ 1
    a, b = hermite_coefs(T, n)
    custom_gauss_rule(convert(T, -Inf), convert(T, Inf), a, b, neither)
end

"""
Convenience function with type `T = Float64`:

    x, w = hermite(n)
"""
hermite(n) = hermite(Float64, n)

function hermite_coefs(::Type{T}, n::Integer) where {T<:AbstractFloat}
    a = zeros(T, n)
    b = zeros(T, n+1)
    b[1] = sqrt(sqrt(convert(T, pi)))
    for i = 1:n
        iT = convert(T, i)
        b[i+1] = sqrt(iT/2)
    end
    return a, b
end

"""
    x, w = logweight(T, n, r=0, endpt=neither)

Returns points `x` and weights `w` for the n-point Gauss rule on
the interval `0 < x < 1` with weight function

    w(x) = xʳ log(1/x),    r ≥ 0.
"""
function logweight(::Type{T}, n::Integer, r::Integer,
                   endpt::EndPt=neither) where {T<:AbstractFloat}
    @assert n ≥ 1
    @assert r ≥ 0
    α, β = logweight_coefs(T, n, r)
    custom_gauss_rule(zero(T), one(T), α, β, endpt)
end

"""
Convenience function with type `T = Float64`:

    x, w = logweight(n, r=0, endpt=neither)
"""
logweight(n, r=0, endpt=neither) = logweight(Float64, n, r, endpt)

"""
More general method works when `w(x) = x^ρ log(1/x)` for real `ρ > -1`:

    x, w = logweight(n, ρ, endpt=neither)
"""
function logweight(n::Integer, ρ::T,
                   endpt::EndPt=neither) where {T<:AbstractFloat}
    @assert n ≥ 1
    @assert ρ > -1
    α, β = logweight_coefs(n, ρ)
    custom_gauss_rule(zero(T), one(T), α, β, endpt)
end

function logweight_coefs(::Type{T}, n::Integer,
                         r::Integer) where {T<:AbstractFloat}
    a, b = shifted_legendre_coefs(T, 2n)
    ν = modified_moments(T, n, r)
    α, β, σ = modified_chebyshev(a, b, ν)
    return α, β
end

function logweight_coefs(n::Integer, ρ::T) where {T<:AbstractFloat}
    a, b = shifted_legendre_coefs(T, 2n)
    ν = modified_moments(n, ρ)
    α, β, σ = modified_chebyshev(a, b, ν)
    return α, β
end

function modified_moments(::Type{T}, n::Integer,
                          r::Integer) where {T<:AbstractFloat}
    @assert r ≥ 0
    @assert n ≥ 0
    m = min(2n, r+1)
    rrp1 = 1 / convert(T, r+1) # reciprocal r+1
    B = one(T)
    S = rrp1
    ν = zeros(T, 2n)
    ν[1] = rrp1 * S
    for l = 2:m
        j = l - 1
	rrpj = 1 / convert(T, r + 1 + j)
        rmj = convert(T, r + 1 - j)
	rrmj = 1 / rmj
        B *= rmj * rrpj
        S +=  rrpj - rrmj
        tlm1 = convert(T, 2l-1)
        ν[l] = rrp1 * B * S * sqrt(tlm1)
    end
    if 2n > r+1
        p = one(T)
        for j = 1:r
            p *= j / convert(T, 2(2j+1))
        end
        ν[r+2] = - sqrt(convert(T, 2r+3)) *p / ( 2(r+1) )
        for l = r+2:2n-1
            Tl = convert(T, l)
            ν[l+1] = - ν[l] * ( (Tl-r-1)/(Tl+r+1) ) * sqrt((2Tl+1)/(2Tl-1) )
        end
    end
    return ν
end

function modified_moments(n::Integer, ρ::T) where {T<:AbstractFloat}
    @assert ρ > -1
    @assert n ≥ 0
    r = (ρ<0) ? 0 : round(Integer, ρ)
    m = min(2n, r+1)
    S = 1 / ( 1 + ρ )
    B = one(T)
    ν = zeros(T, 2n)
    ν[1] = S / (1+ρ)
    for l = 2:m
        j = l - 1
        B *= (ρ+1-j) / (ρ+1+j)
	S += 1 / (ρ+1+j) - 1 / (ρ+1-j)
        ν[l] = B * S * sqrt(convert(T,2l-1))
    end
    if 2n > r+1
        l = r + 2
	j = l - 1
	X = ( (ρ-r) / (ρ+r+2) - 1 ) / (ρ+r+2)
        tlm1 = convert(T, 2l-1)
	ν[l] = ( B / (1+ρ) ) * ( X + ( (ρ-r)/(ρ+r+2) ) * S ) * sqrt(tlm1)
        for l = r+3:2n
	    j = l - 1
            B *= (ρ+1-j) / (ρ+1+j)
            S += 1 / (ρ+1+j) - 1 / (ρ+1-j)
        tlm1 = convert(T, 2l-1)
	    ν[l] = ( B / (1+ρ) ) * ( X + ( (ρ-r)/(ρ+r+2) ) * S ) * sqrt(tlm1)
        end
    end
    return ν
end

function shifted_legendre_coefs(T, n)
    a = zeros(T, n)
    b = zeros(T, n+1)
    fill!(a, 1/convert(T,2))
    b[1] = one(T)
    for k = 2:n+1
        b[k] = (k-1) / ( 2 * sqrt( convert(T, (2k-1)*(2k-3)) ) )
    end
    return a, b
end

"""
    x, w = custom_gauss_rule(lo, hi, a, b, endpt, maxits=maxiterations[T])

Generates the points `x` and weights `w` for a Gauss rule with weight
function `w(x)` on the interval `lo < x < hi`.

The arrays `a` and `b` hold the coefficients (as given, for instance, by
`legendre_coeff`) in the three-term recurrence relation for the monic
orthogonal polynomials `p(0,x)`, `p(1,x)`, `p(2,x)`, ... , that is,

    p(k, x) = (x-a[k]) p(k-1, x) - b[k]² p(k-2, x),    k ≥ 1,

where `p(0, x) = 1` and, by convention, `p(-1, x) = 0` with

              hi
    b[1]^2 = ∫  w(x) dx.
             lo

Thus, `p(k, x) = xᵏ + lower degree terms` and

     hi
    ∫  p(j, x) p(k, x) w(x) dx = 0 if j ≠ k.
    lo
"""
function custom_gauss_rule(lo::T, hi::T,
             a::Array{T,1}, b::Array{T,1}, endpt::EndPt,
             maxits::Integer=maxiterations[T]) where {T<:AbstractFloat}
    n = length(a)
    @assert length(b) == n+1
    if endpt == left
        if n == 1
            a[1] = lo
        else
            a[n] = solve(n, lo, a, b) * b[n]^2 + lo
        end
    elseif endpt == right
        if n == 1
            a[1] = hi
        else
            a[n] = solve(n, hi, a, b) * b[n]^2 + hi
        end
    elseif endpt == both
        if n == 1
            error("Must have at least two points for both ends.")
        end
        g = solve(n, lo, a, b)
        t1 = ( hi - lo ) / ( g - solve(n, hi, a, b) )
        b[n] = sqrt(t1)
        a[n] = lo + g * t1
    end
    w = zero(a)
    special_eigenproblem!(a, b, w, maxits)
    for i = 1:n
        w[i] = (b[1] * w[i])^2
    end
    idx = sortperm(a)
    # Ensure end point values are exact.
    if endpt in (left, both)
        a[idx[1]] = lo
    elseif endpt in (right, both)
        a[idx[n]] = hi
    end
    return a[idx], w[idx]
end

"""
    α, β, σ = modified_chebyshev(a, b, ν)

Implements the modified Chebyshev algorithm described in `doc/notes.tex`
and used in `logweight_coefs`.
"""
function modified_chebyshev(a::Vector{T}, b::Vector{T},
                            ν::Vector{T}) where {T<:AbstractFloat}
    m = length(ν)
    @assert m % 2 == 0 && m >= 2
    n = div(m, 2)
    @assert length(a) >= max(1, 2n-1)
    @assert length(b) >= max(1, 2n-1)
    α = zeros(T, n)
    β = zeros(T, n+1)
    β[1] = sqrt(ν[1])
    σ = zeros(T, 2n, n)
    for l = 1:2n
        σ[l,1] = ν[l] / β[1]
    end
    α[1] = a[1] + b[2] * ( σ[2,1] / σ[1,1] )
    if n >= 2
        # handle k=1 separately since we have omitted σ[l,0]=0.
        t = ( 1 + (b[3]/b[2]) * (σ[3,1]/σ[1,1])
                + ( (a[2]-α[1])/b[2] ) * (σ[2,1]/σ[1,1]) )
        if t < 0
            error("modified Chebyshev algorithm failed at k = 1")
        end
        β[2] = b[2] * sqrt(t)
        for l = 1:2n-2
            σ[l+1,2] = ( (b[l+2]/β[2])*σ[l+2,1]
                       + ((a[l+1]-α[1])/β[2]) * σ[l+1,1]
  	               + (b[l+1]/β[2]) * σ[l,1] )
        end
        α[2] = a[2] + b[3] * (σ[3,2]/σ[2,2]) - β[2] * (σ[2,1]/σ[2,2])
        # general k
        for k = 2:n-1
            t = ( 1 + (b[k+2]/b[k+1]) * (σ[k+2,k]/σ[k,k])
                    + ( (a[k+1]-α[k])/b[k+1] ) * (σ[k+1,k]/σ[k,k])
                    - (β[k]/b[k+1]) * (σ[k+1,k-1]/σ[k,k]) )
            β[k+1] = b[k+1] * sqrt(t)
            if t < 0
                error("modified Chebyshev algorithm failed at k = $k")
            end
            for l = k:2n-k-1
                σ[l+1,k+1] = ( (b[l+2]/β[k+1]) * σ[l+2,k]
                             + ( (a[l+1]-α[k])/β[k+1] ) * σ[l+1,k]
                             + (b[l+1]/β[k+1]) * σ[l,k]
                             - (β[k]/β[k+1]) * σ[l+1,k-1] )
            end
	    α[k+1] = ( a[k+1] + b[k+2] * (σ[k+2,k+1]/σ[k+1,k+1])
                              - β[k+1] * (σ[k+1,k]/σ[k+1,k+1]) )
        end
    end
    return α, β, σ
end

function solve(n::Integer, shift::T, a::Array{T,1},
               b::Array{T,1}) where {T<:AbstractFloat}
    #
    # Perform elimination to find the nth component s = delta[n]
    # of the solution to the nxn linear system
    #
    #     ( J_n - shift I_n ) delta = e_n,
    #
    # where J_n is the symmetric tridiagonal matrix with diagonal
    # entries a[i] and off-diagonal entries b[i], and e_n is the nth
    # standard basis vector.
    #
    t = a[1] - shift
    for i = 2:n-1
        t = a[i] - shift - b[i]^2 / t
    end
    return one(t) / t
end

function special_eigenproblem!(d::Array{T,1}, e::Array{T,1}, z::Array{T,1},
                               maxits::Integer) where {T<:AbstractFloat}
    #
    # Finds the eigenvalues and first components of the normalised
    # eigenvectors of a symmetric tridiagonal matrix by the implicit
    # QL method.
    #
    # d[i]   On entry, holds the ith diagonal entry of the matrix.
    #        On exit, holds the ith eigenvalue.
    #
    # e[i]   On entry, holds the [i,i-1] entry of the matrix for
    #        i = 2, 3, ..., n.  (The value of e[1] is not used.)
    #        On exit, e is overwritten.
    #
    # z[i]   On exit, holds the first component of the ith normalised
    #        eigenvector associated with d[i].
    #
    # maxits The maximum number of QL iterations.
    #
    # Martin and Wilkinson, Numer. Math. 12: 377-383 (1968).
    # Dubrulle, Numer. Math. 15: 450 (1970).
    # Handbook for Automatic Computation, Vol ii, Linear Algebra,
    #        pp. 241-248, 1971.
    #
    # This is a modified version of the Eispack routine imtql2.
    #
    n = length(z)
    z[1] = one(T)
    z[2:n] .= zero(T)
    e[n+1] = zero(T)

    if n == 1 # Nothing to do for a 1x1 matrix.
        return
    end
    for l = 1:n
        for j = 1:maxits
            # Look for small off-diagonal elements.
            m = n
            for i = l:n-1
                if abs(e[i+1]) <= eps(T) * ( abs(d[i]) + abs(d[i+1]) )
                    m = i
                    break
                end
            end
            p = d[l]
            if m == l
                continue
            end
            if j == maxits
                msg = string("No convergence after ", j, " iterations",
                             " (try increasing maxits)")
                error(msg)
            end
            # Form shift
            g = ( d[l+1] - p ) / ( 2 * e[l+1] )
            r = hypot(g, one(T))
            g = d[m] - p + e[l+1] / ( g + copysign(r, g) )
            s = one(T)
            c = one(T)
            p = zero(T)
            for i = m-1:-1:l
                f = s * e[i+1]
                b = c * e[i+1]
                if abs(f) <  abs(g)
                    s = f / g
                    r = hypot(s, one(T))
                    e[i+2] = g * r
                    c = one(T) / r
                    s *= c
                else
                    c = g / f
                    r = hypot(c, one(T))
                    e[i+2] = f * r
                    s = one(T) / r
                    c *= s
                end
                g = d[i+1] - p
                r = ( d[i] - g ) * s + 2 * c * b
                p = s * r
                d[i+1] = g + p
                g = c * r - b
                # Form first component of vector.
                f = z[i+1]
                z[i+1] = s * z[i] + c * f
                z[i]   = c * z[i] - s * f
            end # loop over i
            d[l] -= p
            e[l+1] = g
            e[m+1] = zero(T)
        end # loop over j
    end # loop over l
end

function orthonormal_poly(x::AbstractVector{T}, a::AbstractVector{T},
                          b::AbstractVector{T}) where {T<:AbstractFloat}
    # p[i,j] = value at x[i] of orthonormal polynomial of degree j-1.
    m = length(x)
    n = length(a)
    p = zeros(T, m, n+1)
    rb1 = one(T) / b[1]
    rb2 = one(T) / b[2]
    for i = 1:m
        p[i,1] = rb1
        p[i,2] = rb2 * ( x[i] - a[1] ) * p[i,1]
    end
    for j = 2:n
       rb = one(T) / b[j+1]
       for i = 1:m
           p[i,j+1] = rb * ( (x[i]-a[j]) * p[i,j]
                                - b[j] * p[i,j-1] )
       end
    end
    return p
end

end


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ f7e8e8ce-7dfc-11eb-0677-cb442926e519
using Distributions

# ╔═╡ e9b0055a-7dfc-11eb-2270-e30e833cee00
using StatsPlots

# ╔═╡ ef30e082-7dfd-11eb-3116-b590d75e8735
begin
	k = @bind K html"<input type='range' min='1.0' max='20.0' step='1.0' value='10'>"
	P = @bind p html"<input type='range' min='0.0' max='1.0' step='0.01' value='0.5'>"
	md"""**distribution parameters:**
	$$K$$: $k
	$$p$$: $P
	"""
end

# ╔═╡ 7f8de914-7dfd-11eb-100b-a33c8cfbd43d
x = [0:1:K;]

# ╔═╡ b887bd70-7e03-11eb-2ed9-4127c977a6d7
md"""
 $$K$$ = $K

 $$p$$ = $p
"""

# ╔═╡ c7a6a7aa-7dfe-11eb-1110-f313180ab599
begin
	m = K * p
	v = m * (1 - p)
	md"""
	Mean: $$K * p$$ = $m

	Variance: $$K * p * (1-p)$$ = $v
	"""
end

# ╔═╡ 821bd6d6-7dfe-11eb-1e6a-c5e52652461b
B = Binomial(K, p)

# ╔═╡ 7d111506-7dff-11eb-1b4d-f5b37e333a55
plot(x, pdf.(B, x),
	t = "scatter", label = "Binomial($K, $p)",
	ylims = (0, maximum(pdf.(B, x)) + 0.1))

# ╔═╡ Cell order:
# ╠═7f8de914-7dfd-11eb-100b-a33c8cfbd43d
# ╟─ef30e082-7dfd-11eb-3116-b590d75e8735
# ╟─b887bd70-7e03-11eb-2ed9-4127c977a6d7
# ╟─c7a6a7aa-7dfe-11eb-1110-f313180ab599
# ╠═821bd6d6-7dfe-11eb-1e6a-c5e52652461b
# ╟─7d111506-7dff-11eb-1b4d-f5b37e333a55
# ╟─f7e8e8ce-7dfc-11eb-0677-cb442926e519
# ╟─e9b0055a-7dfc-11eb-2270-e30e833cee00


if !isdefined(Main, :CustomStencil)
    include("../module/CustomStencil.jl")
    using Main.CustomStencil
end
include("misc.jl")
coefs = Float64.([1 1 1])

thumbstack_stencil = @def_stencil_expression(
            @sum(i,-2,2,(
                @sum(j,-2,2, c[max(abs(i),abs(j))]*D[x+i,y+j,z])
            )) + @sum(i,1,2, c[i]*(D[x,y,z+i]+D[x,y,z-i]))
)

st_inst = NewStencilInstance(thumbstack_stencil, 2, coefs)

radius = 2
nx = 5
ny = 5
nz = 5
dx = 1<<(nx)
dy = 1<<(ny)
dz = 1<<(nz)

t_steps = 5
## Actual time steps
data = CreateData(dx,dy,dz,radius)
dout = ApplyStencil(st_inst, data, t_steps)
#heatmap(dout1[:,:,Int(round(dz/2))])
o = Array(dout)

function cpu_thumbstack_stencil(data, radius,coefs)
    out = zeros(Float32,size(data))
    for i = (radius+1):(size(data,1)-radius)
        for j = (radius+1):(size(data,2)-radius)
            for k = (radius+1):(size(data,3)-radius)
                c = 0#coefs[1]*data[i,j,k]
                for rx = -radius:radius, ry = -radius:radius
                    cr = max(abs.([rx,ry])...)
                    c += coefs[cr+1]*(data[i+rx,j+ry,k])
                end
                for rz = 1:radius
                    c += coefs[rz+1]*(data[i,j,k+rz] + data[i,j,k-rz])
                end
                out[i,j,k] = c
            end
        end
    end
    return out
end

cout = cpu_thumbstack_stencil(data, radius, coefs)
for i = 2:t_steps
    global cout
    cout = cpu_thumbstack_stencil(cout, radius, coefs)
end

@assert(sum(abs.(cout .- o))./(reduce(*, size(data))) < 0.1)
mean_err = sum(abs.(cout .- o))./(reduce(*, size(data)))
# using Plots
# shift = 6
# heatmap(cout[:,:,Int(round(dz/2))+shift])
# heatmap(o[:,:,Int(round(dz/2))+shift])
#
# heatmap(o[:,:,Int(round(dz/2))+shift] .- cout[:,:,Int(round(dz/2))+shift])
#
#
# diff =  abs.(cout .- o)#./(reduce(*, size(data)))
# findmax(diff)
#


module LocalSearch

export fill_in, check_value, sudoku_cost, sudoku_greedydesc, flip, fliprow, makeflip, makefliprow, search

"""
    fill_in(sudoku::Matrix)

Fills a sudoku with non-repeated random numbers per row,
so each row has all numbers from 1 to 9.
"""
function fill_in(sudoku::Matrix)
    grid = deepcopy(sudoku)
    for i = 1:9
      row = grid[i,:]
      missing_entries = Set(findall(row .== 0))
      for n = 1:9
          if !(n in row)
                  ms = rand(missing_entries)
                  grid[i, ms] = n
                  delete!(missing_entries, ms)
              end
      end
    end
    grid
  end

"""
    check_value(sudoku, val_i::Int, val_j::Int, value =nothing)

Gives the number of repetitions in row, column and subgrid for a number in position [val_i, val_j] in the sudoku
it is used to check if a particular number meets the sudoku constraints.
It takes as inputs a filled sudoku grid, two coordinates of the evaluated position (row, column) and an optional argument.
If the optional argument is given then it will return the number of constrains violations when this value is placed in the input position.
"""
function check_value(sudoku::Matrix, val_i::Int, val_j::Int, value =nothing)

    if value === nothing
        val = sudoku[val_i, val_j]
    else
        val = value
    end

    constr = 0.0
    #for a particular position [val_i][val_j] how many repetions are in the row.
    for j in 1:length(sudoku[val_i,:])
        if sudoku[val_i, val_j] != 0 && sudoku[val_i, j] == val && j != val_j
            constr += 1
        end
    end
    #for a particular position [val_i][val_j] how many repetions are in the column.
    for i in 1:length(sudoku[val_j,:])
        if sudoku[val_i, val_j] != 0 && sudoku[i, val_j] == val && i != val_i
            constr += 1
        end
    end
    #for a particular position [val_i][val_j] how many repetions are in the subgrid.
    sqrt_n = isqrt(length(sudoku))
    x_i = val_i/3
    x_j = val_j/3
    #create subgrid indices
    if x_i <=1
        bloq_i = 1
        vars_i = [1, 2, 3]
    elseif 1 < x_i <= 2
        bloq_i = 2
        vars_i = [4, 5, 6]
    else
        bloq_i = 3
        vars_i = [7, 8, 9]
    end

    if x_j <=1
        bloq_j = 1
        vars_j = [1, 2, 3]
    elseif 1 < x_j <= 2
        bloq_j = 2
        vars_j = [4, 5, 6]
    else
        bloq_j = 3
        vars_j = [7, 8, 9]
    end
    #looping over each element in the subgrid
    for i in vars_i
        for j in vars_j
            if (i, j) != (val_i, val_j) && sudoku[i, j] != 0 && sudoku[i, j] == val
                constr += 1
            end
        end
    end
    return Int(constr)
end

"""
    sudoku_cost(sudoku::Matrix)

Returns the total number of constraints violations in the sudoku,
Ex. if a number is found two times in a row the cost will be 4 (2 for each number).
"""
function sudoku_cost(sudoku::Matrix)
    total_cost = 0
    for i in 1:9
        for j in 1:9
            total_cost += check_value(sudoku, i, j)
        end
    end
    return Int(total_cost)
end

""""
    sudoku_greedydesc(sudoku::Matrix, empty::Matrix, max_iter::Int)

Selects a random position in the sudoku and changes its number with a different random number,
if the number of constraints violations increase, tries another random, otherwise,
assigns the new number to that position. Then repeats the same process 'max_iter' times
"""
function sudoku_greedydesc(sudoku::Matrix, empty::Matrix, max_iter::Int)
    board = deepcopy(sudoku)
    i = 0
    while i < max_iter
        global total_cost
        #calculate number of conflicts
        total_cost = sudoku_cost(board)
        #check if there are no repeated numbers return the solution
        if total_cost == 0
            return board, 0
        end

        #Select random position to change
        rand_i = rand(1:9)
        rand_j = rand(1:9)

        #check if that position is in fixed variables, if true then randomly select another position
        if empty[rand_i, rand_j] != 0
            continue
        end
        #check which number in [0:9] decrease the number of conflicts and assign it to the current position
        for new_val in Set(1:9)
            if check_value(board, rand_i, rand_j) >= check_value(board, rand_i, rand_j, new_val)
                board[rand_i, rand_j] = new_val
            end
        end
        i += 1
    end

    return board, Int(total_cost)
end

"""
    flip(sudoku::Matrix, empty::Matrix)

Randomly selects two positions in the sudoku and makes a swap.
Returns the new sudoku and the indices of changed positions
"""
function flip(sudoku::Matrix, empty::Matrix)
    (i1, j1, i2, j2) = (rand(1:9), rand(1:9), rand(1:9), rand(1:9))
    res = deepcopy(sudoku)
    while empty[i1, j1] != 0 || empty[i2, j2] != 0
        (i1, j1, i2, j2) = (rand(1:9), rand(1:9), rand(1:9), rand(1:9))
    end
    (res[i1,j1], res[i2,j2]) = (res[i2,j2], res[i1,j1])
    return res, (i1,j1), (i2,j2)
end

"""
    fliprow(sudoku::Matrix, empty::Matrix)

Randomly selects two positions in the same row in the sudoku and makes a swap.
Returns the new sudoku and the indices of changed positions
"""
function fliprow(sudoku::Matrix, empty::Matrix)
    (i1, j1, j2) = (rand(1:9), rand(1:9), rand(1:9))
    res = deepcopy(sudoku)
    while empty[i1, j1] != 0 || empty[i1, j2] != 0
        (i1, j1, j2) = (rand(1:9), rand(1:9), rand(1:9))
    end
    (res[i1,j1], res[i1,j2]) = (res[i1,j2], res[i1,j1])
    return res, (i1,j1), (i1,j2)
end

""""
    makeflip(sudoku::Matrix, empty::Matrix, max_iter::Int)

Evaluates if the swap make by 'flip' function increased/decreased the total cost of the sudoku, if increase,
doesnt make the swap and selects another swap, otherwise makes the swap.
Repeat the process 'max_iter' times
"""
function makeflip(sudoku::Matrix, empty::Matrix, max_iter::Int)
    resp = deepcopy(sudoku)
    i = 0
    while i < max_iter
        resp2, pos1, pos2 = flip(resp, empty)
        if sudoku_cost(resp2) <= sudoku_cost(resp)
            (resp[pos1[1], pos1[2]], resp[pos2[1],pos2[2]]) = (resp[pos2[1], pos2[2]], resp[pos1[1],pos1[2]])
        else
            (resp2[pos1[1], pos1[2]], resp2[pos2[1], pos2[2]]) = (resp2[pos2[1], pos2[2]], resp2[pos1[1], pos1[2]])
        end
        i +=1
    end
    return resp
end

""""
    makefliprow(sudoku::Matrix, empty::Matrix, max_iter::Int)

Evaluates if the swap made by 'flip_row' function increased/decreased the total cost of the sudoku, if increase,
doesnt make the swap and selects another swap, otherwise makes the swap.
Repeat the process 'max_iter' times
"""
function makefliprow(sudoku::Matrix, empty::Matrix, max_iter::Int)
    resp = deepcopy(sudoku)
    i = 0
    while i < max_iter
        resp2, pos1, pos2 = fliprow(resp, empty)
        if sudoku_cost(resp2) <= sudoku_cost(resp)
            (resp[pos1[1], pos1[2]], resp[pos2[1],pos2[2]]) = (resp[pos2[1], pos2[2]], resp[pos1[1],pos1[2]])
        else
            (resp2[pos1[1], pos1[2]], resp2[pos2[1], pos2[2]]) = (resp2[pos2[1], pos2[2]], resp2[pos1[1], pos1[2]])
        end
        i +=1
    end
    return resp
end

""""
    search(sudoku::Matrix, max_repl::Int, max_flips::Int)

Takes an empty Sudoku, and search for the solution that minimizes the number of constraint violations
To find the solution the search must be done at least 1000 times. This algorithm is not efficient.
"""
function search(sudoku::Matrix, max_repl::Int, max_flips::Int)
    @assert length(sudoku[1,:]) == length(sudoku[:,1]) == 9 "Sudoku must be 9x9"
    grid = fill_in(sudoku)
    sol, cost = sudoku_greedydesc(grid, sudoku, max_repl)
    if cost == 0
        return sol, cost
    else res = makeflip(sol, sudoku, max_flips)
    end
    return res
end

include("print_sudoku.jl")
include("examples.jl")
end #module


string_to_numbers(T::DataType,str::AbstractString) = parse.(T,split(str," ",keep=false))

function readneutral(inputfile::AbstractString)
    info("Reading $inputfile file")
    open(inputfile) do input

        NN = parse(Int,readline(input))
        #nodes = Vector{Vec2D{Float64}}(NN) #Pre-allocation a vector of Vec2D{Float64} with NN elemetns
        nodes = Vec2DArray{Float64}(NN) #Pre-allocation a vector of Vec2D{Float64} with NN elemetns

        for i=1:NN
            nodes[i] = Vec2D(string_to_numbers(Float64,readline(input))...)
        end

        NC = parse(Int,readline(input))

        cells = Vector{TriangleCell}(NC)
        bc_cell = Vector{UInt}(NC)

        for i=1:NC
            a = string_to_numbers(UInt,readline(input))
            bc_cell[i] = a[1]
            cells[i] = TriangleCell((a[2],a[3],a[4]))
        end

        NBF = parse(Int,readline(input))

        bf2n =  Vector{Face2Node{2}}(NBF)
        bc_face = Vector{UInt}(NBF)

        for i=1:NBF
            a = string_to_numbers(UInt,readline(input))
            bc_face[i] = a[1]
            bf2n[i] = Face2Node{2}((a[2],a[3]))
        end
        info("Number of nodes: $NN\nNumber of Cells: $NC\nNumber of Boundary Faces: $NBF\n\n")
        return nodes, cells, bc_cell, bf2n, bc_face
    end
end


#############################################################################
# JuMP
# An algebraic modelling language for Julia
# See https://github.com/jump-dev/JuMP.jl
#############################################################################
# speed2.jl
#
# Runs some JuMP benchmarks to test for speed-related regressions.
# Test time to build model in memory.
#############################################################################

include("../solvers.jl")

using JuMP
using Random

function pMedian(solver, numFacility::Int, numCustomer::Int, numLocation::Int)
    Random.seed!(10)
    customerLocations = [rand(1:numLocation) for a in 1:numCustomer]

    buildTime = @elasped begin
        m = Model(solver = solver)

        # Facility locations
        @variable(m, 0 <= s[1:numLocation] <= 1)

        # Aux. variable: x_a,i = 1 iff the closest facility to a is at i
        @variable(m, 0 <= x[1:numLocation, 1:numCustomer] <= 1)

        # Objective: min distance
        @objective(
            m,
            Max,
            sum(
                abs(customerLocations[a] - i) * x[i, a] for
                a in 1:numCustomer, i in 1:numLocation
            )
        )

        # Constraints
        for a in 1:numCustomer
            # Subject to linking x with s
            for i in 1:numLocation
                @constraint(m, x[i, a] - s[i] <= 0)
            end
            # Subject to one of x must be 1
            @constraint(m, sum(x[i, a] for i in 1:numLocation) == 1)
        end

        # Subject to must allocate all facilities
        @constraint(m, sum(s[i] for i in 1:numLocation) == numFacility)
    end

    writeTime = @elapsed JuMP.build(m)

    return buildTime, writeTime
end

function cont5(solver, n)
    m = n
    n1 = n - 1
    m1 = m - 1
    dx = 1 / n
    T = 1.58
    dt = T / m
    h2 = dx^2
    a = 0.001
    yt = [0.5 * (1 - (j * dx)^2) for j in 0:n]

    buildTime = @elapsed begin
        mod = Model(solver = solver)
        @variable(mod, 0 <= y[0:m, 0:n] <= 1)
        @variable(mod, -1 <= u[1:m] <= 1)
        @objective(
            mod,
            Min,
            0.25 *
            dx *
            (
                (y[m, 0] - yt[1])^2 +
                2 * sum((y[m, j] - yt[j+1])^2 for j in 1:n1) +
                (y[m, n] - yt[n+1])^2
            ) + 0.25 * a * dt * (2 * sum(u[i]^2 for i in 1:m1) + u[m]^2)
        )

        # PDE
        for i in 0:m1
            for j in 1:n1
                @constraint(
                    mod,
                    h2 * (y[i+1, j] - y[i, j]) ==
                    0.5 *
                    dt *
                    (
                        y[i, j-1] - 2 * y[i, j] + y[i, j+1] + y[i+1, j-1] -
                        2 * y[i+1, j] + y[i+1, j+1]
                    )
                )
            end
        end

        # IC
        for j in 0:n
            @constraint(mod, y[0, j] == 0)
        end

        # BC
        for i in 1:m
            @constraint(mod, y[i, 2] - 4 * y[i, 1] + 3 * y[i, 0] == 0)
            @constraint(
                mod,
                y[i, n-2] - 4 * y[i, n1] + 3 * y[i, n] ==
                (2 * dx) * (u[i] - y[i, n])
            )
        end
    end

    writeTime = @elapsed JuMP.build(mod)

    return buildTime, writeTime
end

function RunTests()
    # Pmedian
    if !isempty(lp_solvers)
        pmedian_build = Float64[]
        pmedian_write = Float64[]
        for runs in 1:9
            bt, wt = pMedian(first(lp_solvers), 100, 100, 5000)
            push!(pmedian_build, bt)
            push!(pmedian_write, wt)
        end
        sort!(pmedian_build)
        sort!(pmedian_write)
        print(
            "PMEDIAN BUILD MIN=",
            minimum(pmedian_build),
            "  MED=",
            pmedian_build[5],
            "\n",
        )
        print(
            "PMEDIAN INTRN MIN=",
            minimum(pmedian_write),
            "  MED=",
            pmedian_write[5],
            "\n",
        )
    else
        @warn("PMEDIAN NOT RUN!")
    end

    # Cont5
    if !isempty(quad_solvers)
        cont5_build = Float64[]
        cont5_write = Float64[]
        for runs in 1:9
            bt, wt = cont5(first(quad_solvers), 500)
            push!(cont5_build, bt)
            push!(cont5_write, wt)
        end
        sort!(cont5_build)
        sort!(cont5_write)
        print(
            "CONT5 BUILD   MIN=",
            minimum(cont5_build),
            "  MED=",
            cont5_build[5],
            "\n",
        )
        print(
            "CONT5 INTRN   MIN=",
            minimum(cont5_write),
            "  MED=",
            cont5_write[5],
            "\n",
        )
    else
        @warn("CONT5 NOT RUN!")
    end
end

RunTests()


using Test
using CryptoGroups
import CryptoGroups: F2PB, F2GNB, FP, Reducer, Weierstrass, AffinePoint, double, Reducer, BinaryCurve, oncurve, specialize

eq = Weierstrass{1, 1}

x1 = FP{23}(3)
y1 = FP{23}(10)

p1 = AffinePoint{eq}(x1, y1)

x2 = FP{23}(9)
y2 = FP{23}(7)

p2 = AffinePoint{eq}(x2, y2)

@test p1 + p2 == AffinePoint{eq, FP{23}}(17, 20)
@test double(p1) == AffinePoint{eq, FP{23}}(7, 12)
@test oncurve(p1) == true
@test oncurve(p2) == true

f = reverse(bin"10011")
R = Reducer(f)

α = F2PB{R}(bin"0100")

a = α^4
b = one(α)

#eq = BEQ(a, b)
eq = specialize(BinaryCurve, a, b)

p1 = AffinePoint{eq}(α^6, α^8)
p2 = AffinePoint{eq}(α^3, α^13)

@test p1 + p2 == AffinePoint{eq}(one(α), α^13)
@test double(p1) == AffinePoint{eq}(α^10, α^8)

@test oncurve(p1) == true
@test oncurve(p2) == true

α = F2GNB{4, 1}(bin"1100")

#eq = BEQ(zero(α), α^3)

eq = specialize(BinaryCurve, zero(α), α^3)

#error("here")

G = AffinePoint{eq}(α^3, α^5)

@test double(G) == AffinePoint{eq}(α^4, α^3)
@test double(G) + G == AffinePoint{eq}(α^13, α^2)

@test oncurve(G) == true
@test oncurve(G*3) == true

### Experimenting with 𝔽₂ field in Weierstrass equation. The doubling operation though does not work and thus is type constrained.

α = F2GNB{4, 1}(bin"1100")

#eq = WeierstrassEQ(one(α), one(α))
eq = specialize(Weierstrass, one(α), one(α))

p1 = AffinePoint{eq}(α^5, α^10)
p2 = AffinePoint{eq}(α^14, α^5)

@test oncurve(p1) == true
@test oncurve(p2) == true

@test oncurve(p1 + p2) == true

# Testing point multiplication
eq = Weierstrass{1, 1}

p = AffinePoint{eq, FP{23}}(3, 10)

@test oncurve(3*p) == true
@test double(p) == 2p
@test double(p) + p == 3p

@test double(p) + p + double(p) == 5p
@test double(double(p)) + p == 5p

############################### Real curve examples #####################

##################### P_192 ####
let

    p = 6277101735386680763835789423207666416083908700390324961279
    n = 6277101735386680763835789423176059013767194773182842284081

    a = p - 3
    b = parse(BigInt, "64210519 e59c80e7 0fa7e9ab 72243049 feb8deec c146b9b1", base=16)

    Gx = parse(BigInt, "188da80e b03090f6 7cbf20eb 43a18800 f4ff0afd 82ff1012", base=16)
    Gy = parse(BigInt, "07192b95 ffc8da78 631011ed 6b24cdd5 73f977a1 1e794811", base=16)

    #point = AffinePoint{WeierstrassEQ(a, b), (FP/p)}(Gx, Gy)
    point = AffinePoint{specialize(Weierstrass, a, b), specialize(FP, p)}(Gx, Gy)

    @test oncurve(point) == true
    @test point * n == zero(point)

end

######################  K_163 #################

import CryptoGroups: hex2bits, <|

let

    a = 1
    b = 1

    ########### Polynomial basis

    #R = Reducer([163, 7, 6, 3, 0])
    #F = F2PB{R}

    F = specialize(F2PB, [163, 7, 6, 3, 0])

    EQ = specialize(BinaryCurve, F <| a, F <| b)

    Gx = "2 fe13c053 7bbc11ac aa07d793 de4e6d5e 5c94eee8"
    Gy = "2 89070fb0 5d38ff58 321f2e80 0536d538 ccdaa3d9"

    point = AffinePoint{EQ, F}(Gx, Gy)
    @test oncurve(point)

    ######### Normal Basis

    T = 4

    Gx = "0 5679b353 caa46825 fea2d371 3ba450da 0c2a4541"
    Gy = "2 35b7c671 00506899 06bac3d9 dec76a83 5591edb2"

    F = F2GNB{163, 4}
    eq = specialize(BinaryCurve, F <| a, F <| b)

    point = AffinePoint{eq, F}(Gx, Gy)

    @test oncurve(point)
    @test oncurve(point*3)

end


@testset "Special Functions" begin
    @testset "Gamma Function" begin
        @test gamma(Arb(2)) ≈ gamma(2)
        @test gamma(Arb(3)) ≈ gamma(3)
        @test gamma(Acb(2 + 2im)) ≈ gamma(2 + 2im)
        @test gamma(Acb(3 + 3im)) ≈ gamma(3 + 3im)
        @test gamma(ArbSeries([2, 1]))[0] ≈ gamma(2)
        @test gamma(AcbSeries([2 + 2im, 1]))[0] ≈ gamma(2 + 2im)

        @test loggamma(Arb(2)) ≈ loggamma(2)
        @test loggamma(Arb(3)) ≈ loggamma(3)
        @test loggamma(Acb(2 + 2im)) ≈ loggamma(2 + 2im)
        @test loggamma(Acb(3 + 3im)) ≈ loggamma(3 + 3im)
        @test loggamma(ArbSeries([2, 1]))[0] ≈ loggamma(2)
        @test loggamma(AcbSeries([2 + 2im, 1]))[0] ≈ loggamma(2 + 2im)

        @test digamma(Arb(2)) ≈ digamma(2)
        @test digamma(Arb(3)) ≈ digamma(3)
        @test digamma(Acb(2 + 2im)) ≈ digamma(2 + 2im)
        @test digamma(Acb(3 + 3im)) ≈ digamma(3 + 3im)
        @test digamma(ArbSeries([2, 1]))[0] ≈ digamma(2)
        @test digamma(AcbSeries([2 + 2im, 1]))[0] ≈ digamma(2 + 2im)

        @test trigamma(Acb(2 + 2im)) ≈ trigamma(2 + 2im)
        @test trigamma(Acb(3 + 3im)) ≈ trigamma(3 + 3im)

        @test polygamma(Acb(2), Acb(3 + 3im)) ≈ polygamma(2, 3 + 3im)
        @test polygamma(Acb(3), Acb(4 + 4im)) ≈ polygamma(3, 4 + 4im)

        @test gamma(Arb(2), Arb(3)) ≈ gamma(2, 3)
        @test gamma(Arb(3), Arb(4)) ≈ gamma(3, 4)
        @test gamma(Acb(2 + 2im), Acb(3 + 3im)) ≈ gamma(2 + 2im, 3 + 3im)
        @test gamma(Acb(3 + 3im), Acb(4 + 4im)) ≈ gamma(3 + 3im, 4 + 4im)
        @test gamma(Arb(2), ArbSeries([3, 1]))[0] ≈ gamma(2, 3)
        @test gamma(Acb(2 + 2im), AcbSeries([3 + 3im, 1]))[0] ≈ gamma(2 + 2im, 3 + 3im)

        @test all(gamma_inc(Arb(2), Arb(3)) .≈ gamma_inc(2, 3))
        @test all(gamma_inc(Arb(3), Arb(4)) .≈ gamma_inc(3, 4))
        @test all(gamma_inc(Acb(2), Acb(3)) .≈ gamma_inc(2, 3))
        @test all(gamma_inc(Acb(3), Acb(4)) .≈ gamma_inc(3, 4))
        @test all(getindex.(gamma_inc(Arb(2), ArbSeries([3, 1])), 0) .≈ gamma_inc(2, 3))
        @test all(getindex.(gamma_inc(Acb(2), AcbSeries([3, 1])), 0) .≈ gamma_inc(2, 3))

        @test all(beta_inc(Arb(2), Arb(3), Arb(1 // 2)) .≈ beta_inc(2.0, 3.0, 0.5))
        @test all(beta_inc(Arb(3), Arb(4), Arb(1 // 4)) .≈ beta_inc(3.0, 4.0, 0.25))
        @test all(beta_inc(Acb(2), Acb(3), Acb(1 // 2)) .≈ beta_inc(2.0, 3.0, 0.5))
        @test all(beta_inc(Acb(3), Acb(4), Acb(1 // 4)) .≈ beta_inc(3.0, 4.0, 0.25))
        @test all(
            getindex.(beta_inc(Arb(2), Arb(3), ArbSeries([1 // 2, 1])), 0) .≈
            beta_inc(2.0, 3.0, 0.5),
        )
        @test all(
            getindex.(beta_inc(Acb(2), Acb(3), AcbSeries([1 // 2, 1])), 0) .≈
            beta_inc(2.0, 3.0, 0.5),
        )
    end

    @testset "Trigonometric Integrals" begin
        @test expint(Arb(2), Arb(3)) ≈ expint(2, 3)
        @test expint(Arb(3), Arb(4)) ≈ expint(3, 4)
        @test expint(Acb(2 + 2im), Acb(3 + 3im)) ≈ expint(2 + 2im, 3 + 3im)
        @test expint(Acb(3 + 3im), Acb(4 + 4im)) ≈ expint(3 + 3im, 4 + 4im)
        @test expint(Arb(2), ArbSeries([3, 1]))[0] ≈ expint(2, 3)
        @test expint(Acb(2 + 2im), AcbSeries([3 + 3im, 1]))[0] ≈ expint(2 + 2im, 3 + 3im)

        @test expint(Arb(2)) ≈ expint(2)
        @test expint(Arb(3)) ≈ expint(3)
        @test expint(Acb(2 + 2im)) ≈ expint(2 + 2im)
        @test expint(Acb(3 + 3im)) ≈ expint(3 + 3im)
        @test expint(ArbSeries([2, 1]))[0] ≈ expint(2)
        @test expint(AcbSeries([2 + 2im, 1]))[0] ≈ expint(2 + 2im)

        @test expinti(Arb(2)) ≈ expinti(2)
        @test expinti(Arb(3)) ≈ expinti(3)
        @test expinti(Acb(2)) ≈ expinti(2)
        @test expinti(Acb(3)) ≈ expinti(3)
        @test expinti(ArbSeries([2, 1]))[0] ≈ expinti(2)
        @test expinti(AcbSeries([2, 1]))[0] ≈ expinti(2)

        @test sinint(Arb(2)) ≈ sinint(2)
        @test sinint(Arb(3)) ≈ sinint(3)
        @test sinint(Acb(2)) ≈ sinint(2)
        @test sinint(Acb(3)) ≈ sinint(3)
        @test sinint(ArbSeries([2, 1]))[0] ≈ sinint(2)
        @test sinint(AcbSeries([2, 1]))[0] ≈ sinint(2)

        @test cosint(Arb(2)) ≈ cosint(2)
        @test cosint(Arb(3)) ≈ cosint(3)
        @test cosint(Acb(2)) ≈ cosint(2)
        @test cosint(Acb(3)) ≈ cosint(3)
        @test cosint(ArbSeries([2, 1]))[0] ≈ cosint(2)
        @test cosint(AcbSeries([2, 1]))[0] ≈ cosint(2)
    end

    @testset "Error Functions" begin
        @test erf(Arb(2)) ≈ erf(2)
        @test erf(Arb(3)) ≈ erf(3)
        @test erf(Acb(2 + 2im)) ≈ erf(2 + 2im)
        @test erf(Acb(3 + 3im)) ≈ erf(3 + 3im)
        @test erf(ArbSeries([2, 1]))[0] ≈ erf(2)
        @test erf(AcbSeries([2 + 2im, 1]))[0] ≈ erf(2 + 2im)

        @test erfc(Arb(2)) ≈ erfc(2)
        @test erfc(Arb(3)) ≈ erfc(3)
        @test erfc(Acb(2 + 2im)) ≈ erfc(2 + 2im)
        @test erfc(Acb(3 + 3im)) ≈ erfc(3 + 3im)
        @test erfc(ArbSeries([2, 1]))[0] ≈ erfc(2)
        @test erfc(AcbSeries([2 + 2im, 1]))[0] ≈ erfc(2 + 2im)

        @test erfi(Arb(2)) ≈ erfi(2)
        @test erfi(Arb(3)) ≈ erfi(3)
        @test erfi(Acb(2 + 2im)) ≈ erfi(2 + 2im)
        @test erfi(Acb(3 + 3im)) ≈ erfi(3 + 3im)
        @test erfi(ArbSeries([2, 1]))[0] ≈ erfi(2)
        @test erfi(AcbSeries([2 + 2im, 1]))[0] ≈ erfi(2 + 2im)
    end

    @testset "Airy Functions" begin
        @test airyai(Arb(2)) ≈ airyai(2)
        @test airyai(Arb(3)) ≈ airyai(3)
        @test airyai(Acb(2 + 2im)) ≈ airyai(2 + 2im)
        @test airyai(Acb(3 + 3im)) ≈ airyai(3 + 3im)
        @test airyai(ArbSeries([2, 1]))[0] ≈ airyai(2)
        @test airyai(AcbSeries([2 + 2im, 1]))[0] ≈ airyai(2 + 2im)

        @test airyaiprime(Arb(2)) ≈ airyaiprime(2)
        @test airyaiprime(Arb(3)) ≈ airyaiprime(3)
        @test airyaiprime(Acb(2 + 2im)) ≈ airyaiprime(2 + 2im)
        @test airyaiprime(Acb(3 + 3im)) ≈ airyaiprime(3 + 3im)
        @test airyaiprime(ArbSeries([2, 1]))[0] ≈ airyaiprime(2)
        @test airyaiprime(AcbSeries([2 + 2im, 1]))[0] ≈ airyaiprime(2 + 2im)

        @test airybi(Arb(2)) ≈ airybi(2)
        @test airybi(Arb(3)) ≈ airybi(3)
        @test airybi(Acb(2 + 2im)) ≈ airybi(2 + 2im)
        @test airybi(Acb(3 + 3im)) ≈ airybi(3 + 3im)
        @test airybi(ArbSeries([2, 1]))[0] ≈ airybi(2)
        @test airybi(AcbSeries([2 + 2im, 1]))[0] ≈ airybi(2 + 2im)

        @test airybiprime(Arb(2)) ≈ airybiprime(2)
        @test airybiprime(Arb(3)) ≈ airybiprime(3)
        @test airybiprime(Acb(2 + 2im)) ≈ airybiprime(2 + 2im)
        @test airybiprime(Acb(3 + 3im)) ≈ airybiprime(3 + 3im)
        @test airybiprime(ArbSeries([2, 1]))[0] ≈ airybiprime(2)
        @test airybiprime(AcbSeries([2 + 2im, 1]))[0] ≈ airybiprime(2 + 2im)
    end

    @testset "Bessel Functions" begin
        @test besselj(Arb(2), Arb(3)) ≈ besselj(2, 3)
        @test besselj(Arb(3), Arb(4)) ≈ besselj(3, 4)
        @test besselj(Acb(2), Acb(3 + 3im)) ≈ besselj(2, 3 + 3im)
        @test besselj(Acb(3), Acb(4 + 4im)) ≈ besselj(3, 4 + 4im)

        @test besselj0(Arb(2)) ≈ besselj0(2)
        @test besselj0(Arb(3)) ≈ besselj0(3)
        @test besselj0(Acb(2 + 2im)) ≈ besselj0(2 + 2im)
        @test besselj0(Acb(3 + 3im)) ≈ besselj0(3 + 3im)

        @test besselj1(Arb(2)) ≈ besselj1(2)
        @test besselj1(Arb(3)) ≈ besselj1(3)
        @test besselj1(Acb(2 + 2im)) ≈ besselj1(2 + 2im)
        @test besselj1(Acb(3 + 3im)) ≈ besselj1(3 + 3im)

        @test sphericalbesselj(Arb(2), Arb(3)) ≈ sphericalbesselj(2, 3)
        @test sphericalbesselj(Arb(3), Arb(4)) ≈ sphericalbesselj(3, 4)

        @test bessely(Arb(2), Arb(3)) ≈ bessely(2, 3)
        @test bessely(Arb(3), Arb(4)) ≈ bessely(3, 4)
        @test bessely(Acb(2), Acb(3 + 3im)) ≈ bessely(2, 3 + 3im)
        @test bessely(Acb(3), Acb(4 + 4im)) ≈ bessely(3, 4 + 4im)

        @test bessely0(Arb(2)) ≈ bessely0(2)
        @test bessely0(Arb(3)) ≈ bessely0(3)
        @test bessely0(Acb(2 + 2im)) ≈ bessely0(2 + 2im)
        @test bessely0(Acb(3 + 3im)) ≈ bessely0(3 + 3im)

        @test bessely1(Arb(2)) ≈ bessely1(2)
        @test bessely1(Arb(3)) ≈ bessely1(3)
        @test bessely1(Acb(2 + 2im)) ≈ bessely1(2 + 2im)
        @test bessely1(Acb(3 + 3im)) ≈ bessely1(3 + 3im)

        @test besselh(Arb(2), Arb(3)) ≈ besselh(2, 3)
        @test besselh(Arb(3), Arb(4)) ≈ besselh(3, 4)
        @test besselh(Acb(2), Acb(3)) ≈ besselh(2, 3)
        @test besselh(Acb(3), Acb(4)) ≈ besselh(3, 4)
        @test besselh(Arb(2), 1, Arb(3)) ≈ besselh(2, 1, 3)
        @test besselh(Arb(3), 1, Arb(4)) ≈ besselh(3, 1, 4)
        @test besselh(Acb(2), 1, Acb(3)) ≈ besselh(2, 1, 3)
        @test besselh(Acb(3), 1, Acb(4)) ≈ besselh(3, 1, 4)
        @test besselh(Arb(2), 2, Arb(3)) ≈ besselh(2, 2, 3)
        @test besselh(Arb(3), 2, Arb(4)) ≈ besselh(3, 2, 4)
        @test besselh(Acb(2), 2, Acb(3)) ≈ besselh(2, 2, 3)
        @test besselh(Acb(3), 2, Acb(4)) ≈ besselh(3, 2, 4)
        @test_throws SpecialFunctions.AmosException(1) besselh(Arb(2), 3, Arb(3))
        @test_throws SpecialFunctions.AmosException(1) besselh(Acb(2), 3, Acb(3))

        @test besseli(Arb(2), Arb(3)) ≈ besseli(2, 3)
        @test besseli(Arb(3), Arb(4)) ≈ besseli(3, 4)
        @test besseli(Acb(2), Acb(3 + 3im)) ≈ besseli(2, 3 + 3im)
        @test besseli(Acb(3), Acb(4 + 4im)) ≈ besseli(3, 4 + 4im)

        @test besselk(Arb(2), Arb(3)) ≈ besselk(2, 3)
        @test besselk(Arb(3), Arb(4)) ≈ besselk(3, 4)
        @test besselk(Acb(2), Acb(3 + 3im)) ≈ besselk(2, 3 + 3im)
        @test besselk(Acb(3), Acb(4 + 4im)) ≈ besselk(3, 4 + 4im)

        @test besselkx(Arb(2), Arb(3)) ≈ besselkx(2, 3)
        @test besselkx(Arb(3), Arb(4)) ≈ besselkx(3, 4)
        @test besselkx(Acb(2), Acb(3 + 3im)) ≈ besselkx(2, 3 + 3im)
        @test besselkx(Acb(3), Acb(4 + 4im)) ≈ besselkx(3, 4 + 4im)
    end

    @testset "Elliptic Integrals" begin
        @test ellipk(Acb(-1)) ≈ ellipk(-1)
        @test ellipk(Acb(-2)) ≈ ellipk(-2)
        @test ellipk(AcbSeries([-1, 1]))[0] ≈ ellipk(-1)

        @test ellipe(Acb(-1)) ≈ ellipe(-1)
        @test ellipe(Acb(-2)) ≈ ellipe(-2)
    end

    @testset "Zeta Functions" begin
        @test eta(Acb(2 + 2im)) ≈ eta(2 + 2im)
        @test eta(Acb(3 + 3im)) ≈ eta(3 + 3im)

        @test zeta(Arb(2)) ≈ zeta(2)
        @test zeta(Arb(3)) ≈ zeta(3)
        @test zeta(Acb(2 + 2im)) ≈ zeta(2 + 2im)
        @test zeta(Acb(3 + 3im)) ≈ zeta(3 + 3im)
        @test zeta(ArbSeries([2, 1]))[0] ≈ zeta(2)
        @test zeta(AcbSeries([2 + 2im, 1]))[0] ≈ zeta(2 + 2im)
        @test zeta(Arb(2), Arb(3)) ≈ zeta(2, 3)
        @test zeta(Arb(3), Arb(4)) ≈ zeta(3, 4)
        @test zeta(Acb(2 + 2im), Acb(3 + 3im)) ≈ zeta(2 + 2im, 3 + 3im)
        @test zeta(Acb(3 + 3im), Acb(4 + 4im)) ≈ zeta(3 + 3im, 4 + 4im)
        @test zeta(ArbSeries([2, 1]), Arb(3))[0] ≈ zeta(2, 3)
        @test zeta(AcbSeries([2 + 2im, 1]), Acb(3 + 3im))[0] ≈ zeta(2 + 2im, 3 + 3im)
    end
end


function pcor(x::VecOrMat, y::VecOrMat, z::Vector)
    n = length(z)
    ρxz = cor(x,z)
    ρzy = cor(z,y)
    ρxy = cor(x,y)
    @inbounds for j in 1:size(y, 2)
        ρzy_j = ρzy[1, j]
        @simd for i in 1:size(x, 2)
            ρxz_i = ρxz[i, 1]
            ρxy[i,j] = (ρxy[i,j] - ρxz_i * ρzy_j)/sqrt((1.0 - ρxz_i * ρxz_i))/sqrt(1.0 - ρzy_j * ρzy_j)
        end
    end
    ρxy
end


using Primes

function squarefree_sum(rows)
    numbers = Set(binomial(n, k) for n in 1:rows for k in 0:Int(floor(n/2)))

    limit = Int(floor(sqrt(maximum(numbers))))
    squares = primes(limit).^2

    println(sum(square_check(num, squares) for num in numbers))
end

function square_check(num, squares)
    for square in squares
        if num % square == 0
            return 0
        end

        if square > num
            return num
        end
    end

    return num
end

squarefree_sum(50)


## Make sure to activate the VectorAutoregressions environment before running # hide
using CSV, DataFrames
using VectorAutoregressions

exdir = pkgdir(VectorAutoregressions) * "/examples";
datadf = DataFrame(CSV.File(exdir * "/kilian_data.csv"));
first(datadf, 6)

# ## Get Lag Selection Criteria
ls = LagSelection(datadf, 8)

# ## Estimate the VAR
v = VAR(datadf, ls.selection["AIC"])

# ## Check Stability of the VAR
vstable = StabilityCheck(v)


using LinearAlgebra: cross, norm
import AstroTime: seconds

export kepler, period

function period(a, μ)
    2π * sqrt(abs(a)^3/μ) * seconds
end

function kepler(μ, r₀, v₀, Δt, numiter=50, rtol=sqrt(eps()))
    if abs(Δt) < rtol
        return r₀, v₀
    end

    rm = norm(r₀)
    rdv = r₀ ⋅ v₀
    vdv = v₀ ⋅ v₀
    α = -vdv / μ + 2 /rm

    # Elliptic orbit
    if α > rtol
        # For α == 1 the first guess will be too close to converge
        if α ≈ 1
            χ₁ = √μ * Δt * α * 0.97
        else
            χ₁ = √μ * Δt * α
        end
    # Parabolic orbit
    elseif abs(α) < rtol
        h = cross(r₀, v₀)
        hm = norm(h)
        p = hm^2 / μ
        s = 0.5 * (π/2 - atan(3 * sqrt(μ / p^3) * Δt))
        w = atan(tan(s)^(1/3))
        χ₁ = √p * (2 * cot(2 * w))
        # TODO: Check if this is necessary.
        #= α = 0.0 =#
    # Hyperbolic orbit
    else
        a = 1/α
        χ₁ = (sign(Δt) * sqrt(-a)
            * log(-2 * μ * Δt
            / (a * (rdv + sign(Δt) * sqrt(-μ * α) * (1 - rm * alpha)))))
    end

    counter = 0
    converged = false
    χ = 0.0
    c₂ = 0.0
    c₃ = 0.0
    r = 0.0
    ψ = 0.0
    while counter < numiter
        counter += 1
        χ = χ₁
        χ² = χ^2
        ψ = χ² * α
        c₂ = c2(ψ)
        c₃ = c3(ψ)
        r = χ²*c₂ + rdv / √μ * χ * (1 - ψ * c₃) + rm * (1 - ψ * c₂)
        δt = χ^3 * c₃ + rdv / √μ * χ² * c₂ + rm * χ * (1 - ψ * c₃)
        χ₁ = χ + (Δt * √μ - δt) / r
        if abs(χ - χ₁) < rtol
            converged = true
            break
        end
    end

    if !converged
        error("Not converged.")
    end

    f = 1 - χ^2 / rm * c₂
    g = Δt - χ^3 / √μ * c₃
    fdot = √μ / (r * rm) * χ * (ψ * c₃ - 1)
    gdot = 1 - χ^2 / r * c₂
    @assert f * gdot - fdot * g ≈ 1

    return f * r₀ + g * v₀, fdot * r₀ + gdot * v₀
end


using ModelingToolkit, SparseArrays, BenchmarkTools,
    LinearAlgebra

# Generate RAM matrices of factor models
function gen_model(nfact, nitem)
    nfact = Int64(nfact)
    nitem = Int64(nitem)

    ## Model definition
    nobs = Int64(nfact*nitem)
    nnod = nfact+nobs
    @ModelingToolkit.variables λ[1:nobs], σ[1:nobs]

    #F
    Ind = collect(1:nobs)
    Jnd = collect(1:nobs)
    V = fill(1,nobs)
    F = sparse(Ind, Jnd, V, nobs, nnod)

    #S
    Ind = collect(1:nnod)
    Jnd = collect(1:nnod)
    V = [σ; fill(1.0, nfact)]
    S = sparse(Ind, Jnd, V, nnod, nnod)

    #A
    Ind = collect(1:nobs)
    Jnd = vcat([fill(nobs+i, nitem) for i in 1:nfact]...)
    V = λ
    A = sparse(Ind, Jnd, V, nnod, nnod)
    return(A, S, F)
end

function neumann_series(mat::SparseMatrixCSC)
    inverse = I + mat
    next_term = mat^2

    while nnz(next_term) != 0
        inverse += next_term
        next_term *= mat
    end

    return inverse
end

function implied_cov(A, S, F)
    invia = neumann_series(A)

    imp_cov_sym = F*invia*S*permutedims(invia)*permutedims(F)

    imp_cov_sym = Array(imp_cov_sym)
    imp_cov_sym = ModelingToolkit.simplify.(imp_cov_sym)
    return imp_cov_sym
end

function symbolic_ULS(Sigma, S)
    diff = S - Sigma
    F = sum(x -> ^(x, 2), diff)
    #F = diff*permutedims(diff)
    #F = F[diagind(F)]
    F = simplify.(F)
    return diff
end

# small factor models
A, S, F = gen_model(1, 4)

Sigma = implied_cov(A, S, F)

nobs = size(Sigma, 1)

@ModelingToolkit.variables s[1:nobs, 1:nobs]

F = symbolic_ULS(Sigma, s)


using Z3

# Optimize model
# Min x + y
# s.t x > 3
#     y > 5

# Model
m = Z3Model()
@assert isa(m,Z3Model)

# Variables
x = variable(m)
y = variable(m,:Bin)

# Constraints
constraint(m, x == 5)
implies(m, x > 0, y)

# Objective
objective(m, :Min, x*10 + y*100)

#status = solve(m)


include("Graph.jl")
include("NodeCentrality.jl")

using SparseArrays
using LinearAlgebra

function spA(G)
    Is = zeros(Int, 2*G.m)
    Js = zeros(Int, 2*G.m)
    Vs = zeros(2*G.m)
    for (ID, u, v, w) in G.E
        Is[ID] = u
        Js[ID] = v
        Vs[ID] = w
        Is[G.m + ID] = v
        Js[G.m + ID] = u
        Vs[G.m + ID] = w
    end
    return sparse(Is, Js, Vs, G.n, G.n)
end

function EdgeKatz(G; fac1 = 0.85)
    KC = KatzCentrality(G, fac=fac1)
    EK = zeros(G.m)
    for (ID, u, v, w) in G.E
        EK[ID] = (1.0 + KC[u]) * (1.0 + KC[v])
    end
    return EK
end

function getLp(G)
    L = zeros(G.n, G.n)
    for (ID, u, v, w) in G.E
        L[u, v] -= w
        L[v, u] -= w
        L[u, u] += w
        L[v, v] += w
    end
    L .+= (1.0 / G.n)
    Lp = inv(L)
    Lp .-= (1.0 / G.n)
    return Lp
end

function BiharmonicDistanceRelatedCentrality(G)
    Lp = getLp(G)
    C = zeros(G.m)
    for (ID, u, v, w) in G.E
        C[ID] = norm(Lp[:, u] - Lp[:, v])^2
    end
    return C
end

function SpanningEdgeCentrality(G)
    Lp = getLp(G)
    C = zeros(G.m)
    for (ID, u, v, w) in G.E
        C[ID] = Lp[u, u] + Lp[v, v] - 2*Lp[u, v]
    end
    return C
end

function EdgeBetweenness(G)
    g = getAdjacentList(G)
    idx = Dict{Tuple{Int32, Int32}, Int32}()
    for (ID, u, v, w) in G.E
        idx[(u, v)] = ID
        idx[(v, u)] = ID
    end
    Cb = zeros(G.m)
    p = Array{Array{Int32, 1}, 1}(undef, G.n)
    d = zeros(Int32, G.n)
    S = zeros(Int32, G.n+10)
    sigma = zeros(G.n)
    d = zeros(Int32, G.n)
    Q = zeros(Int32, G.n+10)
    delta = zeros(G.n)

    for s = 1 : G.n
        foreach(i -> p[i] = [], 1 : G.n)
        top = 0
        sigma .= 0
        sigma[s] = 1.0
        d .= -1
        d[s] = 0
        front = 1
        rear = 1
        Q[1] = s

        while front <= rear
            v = Q[front]
            front += 1
            top += 1
            S[top] = v
            for w in g[v]
                if d[w] < 0
                    rear += 1
                    Q[rear] = w
                    d[w] = d[v] + 1
                end
                if d[w] == (d[v] + 1)
                    sigma[w] += sigma[v]
                    push!(p[w], v)
                end
            end
        end

        delta .= 0

        while top > 0
            w = S[top]
            top -= 1
            for v in p[w]
                delta[v] += ((sigma[v] / sigma[w]) * (1 + delta[w]))
                Cb[idx[(v, w)]] += ((sigma[v] / sigma[w]) * (1 + delta[w]))
            end
        end

    end

    return Cb
end

function ApproxEdgeBetweenness(G; eps = 0.5)
    g = getAdjacentList(G)
    idx = Dict{Tuple{Int32, Int32}, Int32}()
    for (ID, u, v, w) in G.E
        idx[(u, v)] = ID
        idx[(v, u)] = ID
    end
    pivots = selectPivots(G, eps)
    Cb = zeros(G.m)
    p = Array{Array{Int32, 1}, 1}(undef, G.n)
    d = zeros(Int32, G.n)
    S = zeros(Int32, G.n+10)
    sigma = zeros(G.n)
    d = zeros(Int32, G.n)
    Q = zeros(Int32, G.n+10)
    delta = zeros(G.n)

    for s in pivots
        foreach(i -> p[i] = [], 1 : G.n)
        top = 0
        sigma .= 0
        sigma[s] = 1.0
        d .= -1
        d[s] = 0
        front = 1
        rear = 1
        Q[1] = s

        while front <= rear
            v = Q[front]
            front += 1
            top += 1
            S[top] = v
            for w in g[v]
                if d[w] < 0
                    rear += 1
                    Q[rear] = w
                    d[w] = d[v] + 1
                end
                if d[w] == (d[v] + 1)
                    sigma[w] += sigma[v]
                    push!(p[w], v)
                end
            end
        end

        delta .= 0

        while top > 0
            w = S[top]
            top -= 1
            for v in p[w]
                delta[v] += ((sigma[v] / sigma[w]) * (1 + delta[w]))
                Cb[idx[(v, w)]] += ((sigma[v] / sigma[w]) * (1 + delta[w]))
            end
        end

    end

    return Cb
end

function getW(G)
    IpL = zeros(G.n, G.n)
    for (ID, u, v, w) in G.E
        IpL[u, v] -= w
        IpL[v, u] -= w
        IpL[u, u] += w
        IpL[v, v] += w
    end
    foreach(i -> IpL[i, i] += 1, 1 : G.n)
    return inv(IpL)
end

function ForestCentrality(G)
    W = getW(G)
    C = zeros(G.m)
    for (ID, u, v, w) in G.E
        C[ID] = (W[u, u] + W[v, v] - 2*W[u, v]) / W[u, v]
    end
    return C
end


# SPDX-License-Identifier: X11
# 2020-09-23
# : (Colon) (300pt)

lawofcosines(a, b, θ) = sqrt(a * a + b * b - 2a * b * cos(θ))
angleofhour(h,m) = π * (60h + m) / (6 * 60)
angleofminute(m) = 2π * m / 60

A, B, H, M = parse.(Int, split(readline()))

theta = abs(angleofminute(M) - angleofhour(H,M))
println(lawofcosines(A, B, theta))