"""
    computeslip(mom, δ, force)

Calculates the force due to slippage on the substrate.

Since the model only alows for forces in `x` and `y` direction but not in `z` it is necessary to regularize the velocity at `` h \\rightarrow 0 ``.
The so called no-slip boundary condition can be in theory implemented using `δ` = 0.
However as soon as the film dewets and `height` becomes small, the solution diverges.
Choosing `δ` = 1 satisfies the weak slip condition, while choosing it much larger than 1 yields the strong slip solution.

# Math
One and two dimensional approaches are essentially equal.
Behind the force is the idea of a parabolic shape of the velocity along `z`.
With the boundary conditions

`` \\mathbf{u}|_{h=H} = \\mathbf{u}  ``

`` \\mathbf{u}|_{h=0} = 0 ``

The parameter `δ` allows to push the zero velocity into the substrate, such below te computable region.
Solving this equation with some more assumptions yields

`` \\mathbf{F}_{slip} = \\mu\\alpha_{\\delta}(h) \\mathbf{u} = \\frac{6\\mu h \\mathbf{u}}{2 h^2 + 6\\delta h + 3\\delta^2}. ``

# Example

# References
Here I shamelessly cite the paper from within our group
- [Lattice Boltzmann method for thin-liquid-film hydrodynamics](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.100.033313)

See also: [`Inputconstants`](@ref), [`calculatemoments`](@ref)

"""
function computeslip(mom::JuSwalbe.Macroquant{Vector{T}, Vector{T}}, forces::JuSwalbe.Forces{Vector{T}}, input::JuSwalbe.Inputconstants) where {T<:Number}
    # Measure the length and allocate a dummy array
    len = length(mom.height)
    slippage = zeros(T, len)
    num = zeros(T, len)
    denom = zeros(T, len)
    # Slippage can be calculated assuming a parabolic velocity profile in z
    num .= 6 .* mom.height .* mom.velocity
    denom .= 2 .* mom.height .+ 6 * mom.height * input.δ .+ 3 * input.δ^2
    slippage .= input.μ * (num ./ denom)
    # Write the result into the forcing array
    forces.slip = - slippage
    # Write it out as well
    return slippage
end

function computeslip(mom::JuSwalbe.Macroquant{Matrix{T}, JuSwalbe.Twovector{Matrix{T}}}, forces::JuSwalbe.Forces{JuSwalbe.Twovector{Matrix{T}}}, input::JuSwalbe.Inputconstants) where {T<:Number}
    # Measure the length and allocate a dummy array
    width, thick = size(mom.height)
    δ, μ = T(input.δ), T(input.μ)
    slippagex = zeros(T, (width, thick))
    slippagey = zeros(T, (width, thick))
    numx = zeros(T, (width, thick))
    numy = zeros(T, (width, thick))
    denom = zeros(T, (width, thick))
    # Slippage can be calculated assuming a parabolic velocity profile in z
    numx .= 6 .* mom.height .* mom.velocity.x
    numy .= 6 .* mom.height .* mom.velocity.y
    denom .= 2 .* mom.height .+ 6 * mom.height * δ .+ 3 * δ^2
    slippagex .= μ * ( numx ./ denom )
    slippagey .= μ * ( numy ./ denom )
    # Write the result into the forcing array
    slippage = JuSwalbe.Twovector(x= .- slippagex, y= .- slippagey)
    forces.slip = slippage
    # Write it out as well
    return slippage
end

"""
    computethermalcapillarywaves(mom, force, input)

Computes the forcing arising due to the inclusion of thermal capillary waves that undulate the free surface.

Adds a normal distributed term weighted with the thermal energy kbt to the forcing struct.
This addition makes it possible to simulate not only the thin film equation but also the stochastic thin film equation.
The output of the function should look like (of course dependent on input) this:
![Thermo capillary forcing distribution](https://github.com/Zitzeronion/JuSwalbe.jl/tree/thermalforce-feature/images/thermal_forcing.pdf)

# Math
The stochastic thin film equation (STF) much like the thin film equation is derived from the Landau Lifshitz Navier Stokes equation (LLNS).
Assuming the fluctuations are small and not leading order one can perform an integration of the stochastic stresses, which was done by Grün et al.,

`` \\frac{\\partial h}{\\partial t} = \\nabla\\cdot\\Big[\\frac{h^3}{3\\mu}\\nabla(\\Pi(h) - \\gamma \\Delta h) + \\int_0^h (h-y)\\mathcal{S}_{z||}(y)dy]. ``

Taking into account several assumption it can be shown that the above equation is equal to

`` \\frac{\\partial h}{\\partial t} = \\nabla\\cdot\\Big[\\frac{h^3}{3\\mu}\\nabla(\\Pi(h) - \\gamma \\Delta h) + \\sqrt{\\frac{2k_BT}{3\\mu}h^3}\\mathcal{N}(t)]. ``

Such the integral of `` \\mathcal{S} `` can be expressed with a single multipicative noise term `` \\mathcal{N} `` which is gaussian distributed with zero mean and a variance of one.

This force is adding exactly the term with `` \\mathcal{N} `` to the forcing struct, however it is a little modified due to slippage.

# Example
```jldoctest
julia> using JuSwalbe

julia> mom = simplemoment1d(1000)
JuSwalbe.Macroquant{Array{Float64,1},Array{Float64,1}}
  height: Array{Float64}((1000,)) [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0  …  1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
  velocity: Array{Float64}((1000,)) [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1  …  0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
  pressure: Array{Float64}((1000,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  energy: Array{Float64}((1000,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

julia> force = JuSwalbe.Forces(slip=fill(0.1, 1000), thermal=zeros(1000), h∇p=zeros(1000), bathymetry=zeros(1000))
JuSwalbe.Forces{Array{Float64,1}}
  slip: Array{Float64}((1000,)) [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1  …  0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
  h∇p: Array{Float64}((1000,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  bathymetry: Array{Float64}((1000,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
  thermal: Array{Float64}((1000,)) [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

julia> input = JuSwalbe.Inputconstants()
JuSwalbe.Inputconstants
  lx: Int64 512
  ly: Int64 512
  maxruntime: Int64 100000
  dumping: Int64 1000
  τ: Float64 1.0
  gravity: Float64 0.0
  γ: Float64 0.01
  δ: Float64 1.0
  μ: Float64 0.16666666666666666
  kbt: Float64 0.0

julia> thermal = computethermalcapillarywaves(mom, force, input); # No k_BT -> no force!

julia> input2 = JuSwalbe.Inputconstants(kbt = 0.001)
JuSwalbe.Inputconstants
  lx: Int64 512
  ly: Int64 512
  maxruntime: Int64 100000
  dumping: Int64 1000
  τ: Float64 1.0
  gravity: Float64 0.0
  γ: Float64 0.01
  δ: Float64 1.0
  μ: Float64 0.16666666666666666
  kbt: Float64 0.001

julia> thermal = computethermalcapillarywaves(mom, force, input2); # Now this force is not just zeros
```

# Code
This function is used to compute a force such one needs it for the collision operator, and if Luo forcing is used for the calculation of the moments
See also: [`collisionBGK`](@ref), [`calculatemoments`](@ref)

# References
The theoretical foundation for this was laid by G. Grün and H. Stone,
- [Thin-Film Flow Influenced by Thermal Noise](https://link.springer.com/article/10.1007/s10955-006-9028-8)
- [Spreading of Viscous Fluid Drops on a Solid Substrate Assisted by Thermal Fluctuations](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.95.244505)

Some more recent references are
- [Molecular simulation of thin liquid films: Thermal fluctuations and instability](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.100.023108)
- [Fully nonlinear dynamics of stochastic thin-film dewetting](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.92.061002)

Basically it reduces to drastically simplify the LLNS under various assumptions and show that the result is more or less equal to something simple.
"""
function computethermalcapillarywaves(mom::JuSwalbe.Macroquant{Vector{T},Vector{T}}, forces::JuSwalbe.Forces{Vector{T}}, input::Inputconstants) where {T<:Number}
  if input.kbt == 0.0
    return
  else
    len = length(mom.height)
    thermocap = zeros(T, len)
    slip = deepcopy(forces.slip)
    μ = T(input.μ)
    kbt = T(input.kbt)
    # Generate a Gaussian distribution with zero mean and variance of one
    gaussian = Normal()
    # Fill an array with random numbers distributed according to a Normal
    gaussianvec = rand(gaussian, len)
    # Compute the forces due to thermal capillary waves
    thermocap .= sqrt.(2 * kbt * μ * slip) .* gaussianvec
    forces.thermal = thermocap
    return thermocap
  end
end

function computethermalcapillarywaves(mom::JuSwalbe.Macroquant{Matrix{T},JuSwalbe.Twovector{Matrix{T}}}, forces::JuSwalbe.Forces{JuSwalbe.Twovector{Matrix{T}}}, input::Inputconstants) where {T<:Number}
  if input.kbt == 0.0
    return
  else
    width, thick = size(mom.height)
    thermocapx = zeros(T, (width, thick))
    thermocapy = zeros(T, (width, thick))
    slip = deepcopy(forces.slip)
    μ = T(input.μ)
    kbt = T(input.kbt)
    # Generate a Gaussian distribution with zero mean and a variance of one
    gaussian = Normal()
    # Fill an array with random numbers distributed according to a gaussian
    gaussianvecx = rand(gaussian, (width, thick))
    gaussianvecy = rand(gaussian, (width, thick))
    # Compute the forces due to thermal capillary waves
    thermocapx .= sqrt.(2 * kbt * μ * slip.x) .* gaussianvecx
    thermocapy .= sqrt.(2 * kbt * μ * slip.y) .* gaussianvecy

    forces.thermal = JuSwalbe.Twovector(x=thermocapx, y=thermocapy)
    return thermocapx, thermocapy
  end
end


x, y = generators(MPoly{Float64}, :x, :y)

@test (x+y)^2 == x^2 + 2x*y + y^2
@test x*(2x-y+1)^2 + 1 == 4x^3 - 4x^2*y + 4x^2 + y^2*x - 2x*y + x + 1
@test x + 1 == x + 1.0
@test x[1,0] == 1.0
@test x[0,1] == 0.0
@test eltype(x + im) == Complex{Float64}

@test  (x^(-1))[-1,0] == 1.0
@test_throws ArgumentError (x*y + x)^(-1)


# Import packages
using  OrdinaryDiffEq, DiffEqDevTools, Test, Random

# Define test problems
# Note that the time span in DiffEqProblemLibrary.ODEProblemLibrary is given by
# Float64 numbers

linear = (u,p,t) -> (p*u)
linear_analytic = (u0,p,t) -> u0*exp(p*t)
prob_ode_bigfloatlinear = ODEProblem(
                          ODEFunction(linear,analytic=linear_analytic),
                          big"0.5",(big"0.0",big"1.0"),big"1.01")

f_2dlinear = (du,u,p,t) -> (@. du = p*u)
f_2dlinear_analytic = (u0,p,t) -> @. u0*exp(p*t)
prob_ode_bigfloat2Dlinear = ODEProblem(
                    ODEFunction(f_2dlinear,analytic=f_2dlinear_analytic),
                  rand(BigFloat,(4,2)),(big"0.0",big"1.0"),big"1.01")

# Prepare tests
Random.seed!(100)
problem_array = [prob_ode_bigfloatlinear,prob_ode_bigfloat2Dlinear]
dts = 1 .//2 .^(8:-1:1)

testTol = 0.2

@testset "Testing extrapolation methods" begin

    # Test AitkenNeville
    println("Testing AitkenNeville")
    @testset "Testing AitkenNeville" begin
        @testset "Testing sequential AitkenNeville" begin
            for prob in problem_array
            global dts

            #  Convergence test
            for j = 1:4
                sim = test_convergence(dts,prob,AitkenNeville(max_order = j,
                min_order = j, init_order = j, threading = false))
                @test sim.𝒪est[:final] ≈ j atol=testTol
            end

            # Regression test
            sol = solve(prob,AitkenNeville(max_order = 9, min_order = 1,
                init_order = 9, threading = false),reltol=1e-3)
            @test length(sol.u) < 15
            sol = solve(prob,AitkenNeville(max_order = 9, min_order = 1,
                init_order = 9, threading = false),reltol=1e-6)
            @test length(sol.u) < 18
            end
        end
    end # AitkenNeville

    # Define the subdividing sequences
    sequence_array =[:harmonic, :romberg, :bulirsch]

    println("Testing ImplicitEulerExtrapolation")
    @testset "Testing ImplicitEulerExtrapolation" begin
    for prob in problem_array, seq in sequence_array
        global dts

        newTol = 0.35
        #  Convergence test
        for j = 1:4
        alg = ImplicitEulerExtrapolation(min_order = j,
            init_order = j, max_order=j,
            sequence = seq, threading = false)
        sim = test_convergence(dts,prob,alg)
        @test sim.𝒪est[:final] ≈ alg.n_init+1.1 atol=newTol #Superconvergence
        end
        # Regression test
        sol = solve(prob,ImplicitEulerExtrapolation(max_order = 9, min_order = 1,
            init_order = 9,sequence = seq,threading = false),reltol=1e-3)
        @test length(sol.u) < 15
    end
    end

    println("Testing ImplicitEulerBarycentricExtrapolation")
    @testset "Testing ImplicitEulerBarycentricExtrapolation" begin
        for prob in problem_array, seq in sequence_array
            global dts

            newTol = 0.35
            #  Convergence test
            for j = 1:4
            alg = ImplicitEulerBarycentricExtrapolation(min_order = j,
                init_order = j, max_order=j,
                sequence = seq, threading = false)
            sim = test_convergence(dts,prob,alg)
            @test sim.𝒪est[:final] ≈ alg.n_init + 0.5 atol=newTol #Superconvergence
            end
            # Regression test
            sol = solve(prob,ImplicitEulerBarycentricExtrapolation(max_order = 9, min_order = 1,
                init_order = 9,sequence = seq, threading = false),reltol=1e-3)
            @test length(sol.u) < 15
        end
    end

    println("Testing ImplicitDeuflhardExtrapolation")
    @testset "Testing ImplicitDeuflhardExtrapolation" begin
    for prob in problem_array, seq in sequence_array
        global dts

        # Convergence test
        for j = 1:6
        alg = ImplicitDeuflhardExtrapolation(min_order = j,
            init_order = j, max_order=j,
            sequence = seq,threading=false)
        sim = test_convergence(dts,prob,alg)
        @test sim.𝒪est[:final] ≈ 2*(alg.n_init+1) atol=testTol
        end

        # Regression test
        alg = ImplicitDeuflhardExtrapolation(max_order=9, min_order=1,
        init_order=9, sequence=seq,threading=false)
        sol = solve(prob, alg, reltol=1e-3)
        @test length(sol.u) < 10
    end
    end

    println("Testing ImplicitHairerWannerExtrapolation")
    @testset "Testing ImplicitHairerWannerExtrapolation" begin
    for prob in problem_array, seq in sequence_array
        global dts

        # Convergence test
        for j = 1:6
        alg = ImplicitHairerWannerExtrapolation(min_order = j,
            init_order = j, max_order=j,
            sequence = seq, threading=false)
        sim = test_convergence(dts,prob,alg)
        @test sim.𝒪est[:final] ≈ 2*(alg.n_init+1) - 1 atol=testTol
        end

        alg = ImplicitHairerWannerExtrapolation(max_order=9, min_order=1,
        init_order=9, sequence=seq,threading=false)
        sol = solve(prob, alg, reltol=1e-3)
        @test length(sol.u) < 10
    end
    end

    # Test ExtrapolationMidpointDeuflhard

    println("Testing ExtrapolationMidpointDeuflhard")
    @testset "Testing ExtrapolationMidpointDeuflhard" begin
    @testset "Testing sequential ExtrapolationMidpointDeuflhard" begin
        for prob in problem_array, seq in sequence_array
        global dts

        # Convergence test
        for j = 1:6
            alg = ExtrapolationMidpointDeuflhard(min_order = j,
            init_order = j, max_order=j,
            sequence = seq, threading=false)
            sim = test_convergence(dts,prob,alg)
            @test sim.𝒪est[:final] ≈ 2*(alg.n_init+1) atol=testTol
        end

        # Regression test
        alg = ExtrapolationMidpointDeuflhard(max_order=9, min_order=1,
            init_order=9, sequence=seq, threading=false)
        sol = solve(prob, alg, reltol=1e-3)
        @test length(sol.u) < 10
        end
    end
    end # ExtrapolationMidpointDeuflhard

    # Test ExtrapolationMidpointHairerWanner
    println("Testing ExtrapolationMidpointHairerWanner")
    @testset "Testing ExtrapolationMidpointHairerWanner" begin
    @testset "Testing sequential ExtrapolationMidpointHairerWanner" begin
        for prob in problem_array, seq in sequence_array
        global dts

        # Convergence test
        for j = 1:6
            alg = ExtrapolationMidpointHairerWanner(min_order = j,
            init_order = j, max_order=j,
            sequence = seq, threading=false)
            sim = test_convergence(dts,prob,alg)
            @test sim.𝒪est[:final] ≈ 2*(alg.n_init+1) atol=testTol
        end

        # Regression test
        alg = ExtrapolationMidpointHairerWanner(max_order=9, min_order=2,
            init_order=9, sequence=seq, threading=false)
        sol = solve(prob, alg, reltol=1e-3)
        @test length(sol.u) < 10
        end
    end

    end # ExtrapolationMidpointHairerWanner

    println("Regression Test Float32 and Float64 Fallbacks")
    @testset "Regression Test Float32 and Float64 Fallbacks" begin
    prob_ode_2Dlinear = ODEProblem(
                        ODEFunction(f_2dlinear,analytic=f_2dlinear_analytic),
                        Float64.(prob_ode_bigfloat2Dlinear.u0),(0.0,1.0),1.01)
    s1 = solve(prob_ode_bigfloat2Dlinear,ExtrapolationMidpointDeuflhard())
    s2 = solve(prob_ode_2Dlinear,ExtrapolationMidpointDeuflhard())
    @test all(all(s1[i] - s2[i] .< 5e-14) for i in 1:length(s1))

    prob_ode_2Dlinear = ODEProblem(
                        ODEFunction(f_2dlinear,analytic=f_2dlinear_analytic),
                        Float32.(prob_ode_bigfloat2Dlinear.u0),(0.0f0,1.0f0),1.01f0)
    s1 = solve(prob_ode_bigfloat2Dlinear,ExtrapolationMidpointDeuflhard())
    s2 = solve(prob_ode_2Dlinear,ExtrapolationMidpointDeuflhard())
    @test all(all(s1[i] - s2[i] .< 5e-6) for i in 1:length(s1))
    end

end # Extrapolation methods


r=0.1

#theta(x::FloatOrVector,r::Float64)
#theta's vanishe in 0, are bounded by 1, increasing, concave
# positive for positive values, and negative for negative values.
@test MPCCSolver.theta(0.0, r) .== 0.0
@test MPCCSolver.theta(1.0+1e6, r) .<= 1
@test MPCCSolver.theta(-1e-6, r) .< 0.0
@test MPCCSolver.theta(1.0, r) .>= MPCCSolver.theta(0.5, r)

#dtheta(x::FloatOrVector,r::Float64)
#la dérivée des fonctions theta doit être positive
@test MPCCSolver.dtheta(1.0+1e6, r) >= 0.0
@test MPCCSolver.dtheta(-0.5, r)    >= 0.0

#ddtheta(x::FloatOrVector,r::Float64)
#la dérivée seconde des fonctions theta doit être négative
@test MPCCSolver.ddtheta(1.0+1e6, r) <= 0.0
@test MPCCSolver.ddtheta(-0.5, r)    <= 0.0

#invtheta(x::Vector,r::Float64)
@test abs(MPCCSolver.invtheta(MPCCSolver.theta(1.,r),r)-1.)      <= eps(Float64)*10.
@test abs(MPCCSolver.invtheta(MPCCSolver.theta(-1e-6,r),r)+1e-6) <= eps(Float64)*10.

#_poly_degree_2_positif
#Compute the smallest positive root of a second order polynomial
#or a negative/complex value if no positive real root.
#If no solution or infinite nb of solutions, return Inf.
@test abs(MPCCSolver._poly_degree_2_positif(0.5,-1.,0.25) - 0.2928932188134524) <= eps(Float64)
@test abs(MPCCSolver._poly_degree_2_positif(0.,1.,-0.5) - 0.5) <= eps(Float64)
@test !isfinite(MPCCSolver._poly_degree_2_positif(0.,0.,-0.))
@test !isfinite(MPCCSolver._poly_degree_2_positif(1.,1.,1.))

#alpha_theta_max(x::Float64,dx::Float64,y::Float64,dy::Float64,r::Float64,s::Float64,t::Float64)
#alpha_theta_max : résoud l'équation en alpha où psi(x,r,s,t)=s+t*theta(x,r)
r=1.0;s=1.0;t=0.0 #résoud l'éq.

x=1.0;y=2.0;dx=1.0;dy=-1.0;
#doit rendre (0,1)
@test isequal(MPCCSolver.alpha_theta_max(x,dx,y,dy,r,s,t), [0.0,1.0])

x=0.0;y=2.0;dx=1.0;dy=-2.0;
#solution : (1,0.5)

@test isequal(MPCCSolver.alpha_theta_max(x,dx,y,dy,r,s,t), [1.0,0.5])

r=1.0;s=1.0;t=1.0
#on tape la contrainte :
x=1.0;y=2.0;dx=1.0;dy=0.0;

@test minimum(MPCCSolver.alpha_theta_max(x,dx,y,dy,r,s,t))>=0.0
@test maximum(MPCCSolver.alpha_theta_max(x,dx,y,dy,r,s,t))==Inf

x=2.0;y=1.0;dx=-1.0;dy=2.0;

@test minimum(MPCCSolver.alpha_theta_max(x,dx,y,dy,r,s,t))>=0.0
@test maximum(MPCCSolver.alpha_theta_max(x,dx,y,dy,r,s,t))<Inf

#on est sous (s,s) et donc on ne tape pas la contrainte virtuelle
x=0.0;y=0.0;dx=0.0;dy=-2.0;

@test isequal(maximum(MPCCSolver.alpha_theta_max(x,dx,y,dy,r,s,t)),Inf)


module Utils

    using EllipsisNotation

    export dimensions, ∑

    # shorthand for sum operation
    ∑(fn, values) = sum(fn, values)

    """
    Gets the dimensions of an array of shape d x N x M ... ignoring
    the first dimension that here represents the 3 cartesian coordinates.
    E.g., d x N x M -> N x M
    """
    dimensions(X::AbstractArray) =  size(X)[2:end]

    """
    Get the number of nodes along each dimension (given d x N x M input array
    it returns (N-1)x(M-1)).
    """
    nnodes(X::AbstractArray) = ndims(X)==2 ? size(X, 2)-1 : dimensions(X) .- 1

    """
        x(X)

    Select the first coordinate from an array d x N ( x M)
    """
    x(X::AbstractArray) = X[1, ..]

    """
        y(X)

    Select the second coordinate from an array d x N ( x M)
    """
    y(X::AbstractArray) = X[2, ..]

    """
        z(X)

    Select the third coordinate from an array d x N ( x M)
    """
    z(X::AbstractArray) = X[3, ..]

    """
    Prepares standard parameters used in the creation of 1D splines
    """
    function get_curve_parameters(nodes::AbstractArray, δt)
        # parameter interval
        τ = 0:δt:1

        # dimesions
        N = 1
        η = size(nodes, 2) - 1
        d = size(nodes, 1)

        return τ, N, η, d
    end
end


"""
    linspace(start, stop, n)

Python linspace function
"""
linspace(start, stop, n::T) where {T<:Integer} =
    collect(range(start, stop = stop, length = n))

"""
    heaviside(x)

Heaviside step function
"""
heaviside(x::T) where {T<:Real} = ifelse(x >= 0, one(x), zero(x))

"""
    fortsign(x, y)

Fortran sign function
"""
fortsign(x, y) = abs(x) * sign(y)

"""
    mat_split(m)

Split matrix into row vectors
"""
function mat_split(m::AbstractMatrix{T}) where {T<:Number}

    if length(m[1, :]) == 2
        nx = eltype(m).([1.0 0.0])
        ny = eltype(m).([0.0 1.0])

        return nx * m, ny * m
    elseif length(m[1, :]) == 3
        nx = eltype(m).([1.0 0.0 0.0])
        ny = eltype(m).([0.0 1.0 0.0])
        nz = eltype(m).([0.0 0.0 1.0])

        return nx * m, ny * m, nz * m
    end

end

"""
    central_diff(y, x)
    central_diff(y, dx)

Central difference
"""
function central_diff(y::T, x::T) where {T<:AbstractVector{<:Number}}

    dy = zeros(eltype(y), axes(y))

    idx = eachindex(y) |> collect
    i0 = idx[1]
    i1 = idx[end]

    dy[i0] = (y[i0+1] - y[i0]) / (x[i0+1] - x[i0] + 1e-7)
    dy[i1] = (y[i1] - y[i1-1]) / (x[i1] - x[i1-1] + 1e-7)
    for i = i0+1:i1-1
        dy[i] = (y[i+1] - y[i-1]) / (x[i+1] - x[i-1] + 1e-7)
    end

    return dy

end

function central_diff(y::AbstractVector{T}, dx) where T
    x = ones(eltype(y), axes(y)) .* dx
    dy = central_diff(y, x)

    return dy
end

"""
    central_diff2(y, x)
    central_diff2(y, dx)

Central difference for second-order derivatives
"""
function central_diff2(y::T, x::T) where {T<:AbstractVector{<:Number}}
    dy = zeros(eltype(y), axes(y))

    i0 = eachindex(y) |> first
    i1 = eachindex(y) |> last

    for i = i0+1:i1-1
        dy[i] = (y[i+1] - 2.0 * y[i] + y[i-1]) / (x[i+1] - x[i-1])^2
    end

    return dy
end

function central_diff2(y::AbstractVector{T}, dx) where T
    x = ones(eltype(y), axes(y)) .* dx
    dy = central_diff2(y, x)

    return dy
end

"""
    central_diff!(dy, y, x)
    central_diff!(dy, y, dx)

Central difference
"""
function central_diff!(dy::AbstractVector{T1}, y::T2, x::T2) where {T1,T2<:AbstractVector{<:Number}}

    @assert axes(dy) == axes(y) == axes(x)

    idx = eachindex(y) |> collect
    i0 = idx[1]
    i1 = idx[end]

    dy[i0] = (y[i0+1] - y[i0]) / (x[i0+1] - x[i0] + 1e-7)
    dy[i1] = (y[i1] - y[i1-1]) / (x[i1] - x[i1-1] + 1e-7)
    for i = i0+1:i1-1
        dy[i] = (y[i+1] - y[i-1]) / (x[i+1] - x[i-1] + 1e-7)
    end

end

function central_diff!(dy::AbstractVector{T1}, y::AbstractVector{T2}, dx) where {T1,T2}
    x = ones(eltype(y), axes(y)) .* dx
    central_diff!(dy, y, x)
end

"""
    central_diff2!(dy, y, x)
    central_diff2!(dy, y, dx)

Central difference
"""
function central_diff2!(dy::AbstractVector{T1}, y::T2, x::T2) where {T1,T2<:AbstractVector{<:Number}}

    @assert axes(dy) == axes(y) == axes(x)

    i0 = eachindex(y) |> first
    i1 = eachindex(y) |> last

    dy[i0] = 0.0
    dy[i1] = 0.0
    for i = i0+1:i1-1
        dy[i] = (y[i+1] - 2.0 * y[i] + y[i-1]) / (x[i+1] - x[i-1])^2
    end

end

function central_diff2!(dy::AbstractVector{T1}, y::AbstractVector{T2}, dx) where {T1,T2}
    x = ones(eltype(y), axes(y)) .* dx
    central_diff2!(dy, y, x)
end

"""
    upwind_diff(y, x; stream)
    upwind_diff(y, dx; stream)

Upwind difference
"""
function upwind_diff(
    y::AbstractVector{T1},
    x::AbstractVector{T2};
    stream = :right::Symbol,
) where {T1,T2}

    dy = zeros(eltype(y), axes(y))

    idx = eachindex(y) |> collect
    i0 = idx[1]
    i1 = idx[end]

    if stream == :right
        dy[i0] = 0.0
        for i = i0+1:i1
            dy[i] = (y[i] - y[i-1]) / (x[i] - x[i-1] + 1e-7)
        end
    elseif stream == :left
        dy[i1] = 0.0
        for i = i0:i1-1
            dy[i] = (y[i+1] - y[i]) / (x[i+1] - x[i] + 1e-7)
        end
    else
        throw("streaming direction should be :left or :right")
    end

    return dy

end

function upwind_diff(y::AbstractVector{T}, dx; stream = :right::Symbol) where T
    x = ones(eltype(y), axes(y)) .* dx
    dy = upwind_diff(y, x, stream = stream)

    return dy
end

"""
    upwind_diff!(dy, y, x; stream)
    upwind_diff!(dy, y, dx; stream)

Upwind difference
"""
function upwind_diff!(
    dy::AbstractVector{T1},
    y::AbstractVector{T2},
    x::AbstractVector{T3};
    stream = :right::Symbol,
) where {T1,T2,T3}

    @assert axes(dy) == axes(y) == axes(x)

    idx = eachindex(y) |> collect
    i0 = idx[1]
    i1 = idx[end]

    if stream == :right
        dy[i0] = 0.0
        for i = i0+1:i1
            dy[i] = (y[i] - y[i-1]) / (x[i] - x[i-1] + 1e-7)
        end
    elseif stream == :left
        dy[i1] = 0.0
        for i = i0:i1-1
            dy[i] = (y[i+1] - y[i]) / (x[i+1] - x[i] + 1e-7)
        end
    else
        throw("streaming direction should be :left or :right")
    end

    return dy

end

function upwind_diff!(
    dy::AbstractVector{T1},
    y::AbstractVector{T2},
    dx;
    stream = :right::Symbol,
) where {T1,T2}

    x = ones(eltype(y), axes(y)) .* dx
    upwind_diff!(dy, y, x, stream = stream)

end

"""
    unstruct_diff(u, x, nx; mode)
    unstruct_diff(u, x, nx, dim; mode)

Finite difference for pseudo-unstructured mesh
"""
function unstruct_diff(
    u::AbstractVector{T1},
    x::AbstractVector{T2},
    nx::Integer;
    mode = :central::Symbol,
) where {T1,T2}

    uu = reshape(u, (nx, :))
    xx = reshape(x, (nx, :))

    dux = similar(xx)
    for i = 1:nx
        if mode == :central
            dux[i, :] .= central_diff(uu[i, :], xx[i, :])
        elseif mode == :upwind
            dux[i, :] .= upwind_diff(uu[i, :], xx[i, :])
        else
            throw("difference mode should be central or upwind")
        end
    end

    return reshape(dux, (1, :))

end

function unstruct_diff(
    u::Function,
    x::AbstractVector{T},
    nx::Integer,
    dim::Integer;
    mode = :central::Symbol,
) where T

    uu = reshape(u.(x), (nx, :))
    xx = reshape(x, (nx, :))
    dux = zeros(eltype(x), axes(xx))

    if dim == 1
        for j in axes(xx, 2)
            if mode == :central
                dux[:, j] .= central_diff(uu[:, j], xx[:, j])
            elseif mode == :upwind
                dux[:, j] .= upwind_diff(uu[:, j], xx[:, j])
            end
        end
    elseif dim == 2
        for i in axes(xx, 1)
            if mode == :central
                dux[i, :] .= central_diff(uu[i, :], xx[i, :])
            elseif mode == :upwind
                dux[i, :] .= upwind_diff(uu[i, :], xx[i, :])
            end
        end
    end

    return dux[:]

end

"""
    lgwt(, a, b)

Gauss Legendre integral for fast spectral method

* @args: number of quadrature points N, integral range [a, b]
* @args: quadrature points x & weights w

"""
function lgwt(N::Integer, a::Real, b::Real)
    x = zeros(N)
    w = zeros(N)

    N1 = N
    N2 = N + 1

    y = zeros(N1)
    y0 = zeros(N1)
    Lp = zeros(N1)
    L = zeros(N1, N2)

    # initial guess
    for i = 1:N1
        y[i] =
            cos((2.0 * (i - 1.0) + 1.0) * 4.0 * atan(1.0) / (2.0 * (N - 1.0) + 2.0)) +
            0.27 / N1 *
            sin(4.0 * atan(1.0) * (-1.0 + i * 2.0 / (N1 - 1.0)) * (N - 1.0) / N2)
        y0[i] = 2.0
    end

    # compute the zeros of the N+1 legendre Polynomial
    # using the recursion relation and the Newton method
    while maximum(abs.(y .- y0)) > 0.0000000000001
        L[:, 1] .= 1.0
        L[:, 2] .= y
        for k = 2:N1
            @. L[:, k+1] = ((2.0 * k - 1.0) * y * L[:, k] - (k - 1) * L[:, k-1]) / k
        end
        @. Lp = N2 * (L[:, N1] - y * L[:, N2]) / (1.0 - y^2)
        @. y0 = y
        @. y = y0 - L[:, N2] / Lp
    end

    # linear map from [-1 1] to [a,b]
    @. x = (a * (1.0 - y) + b * (1.0 + y)) / 2.0
    @. w = N2^2 * (b - a) / ((1.0 - y^2) * Lp^2) / N1^2

    return x, w
end

"""
    extract_last(a, idx; mode)

Extract subarray except the last column
"""
function extract_last(a::AbstractArray{T}, idx::Integer; mode = :view::Symbol) where T
    if mode == :copy

        if ndims(a) == 2
            sw = a[:, idx]
        elseif ndims(a) == 3
            sw = a[:, :, idx]
        elseif ndims(a) == 4
            sw = a[:, :, :, idx]
        elseif ndims(a) == 5
            sw = a[:, :, :, :, idx]
        end

    elseif mode == :view

        if ndims(a) == 2
            sw = @view a[:, idx]
        elseif ndims(a) == 3
            sw = @view a[:, :, idx]
        elseif ndims(a) == 4
            sw = @view a[:, :, :, idx]
        elseif ndims(a) == 5
            sw = @view a[:, :, :, :, idx]
        end

    else

        throw("Error in extraction mode setup")

    end

    return sw
end


#=

  Euler #44 in Julia.

  Problem 44:
  """
  Pentagonal numbers are generated by the formula, P(n)=n(3n−1)/2.
  The first ten pentagonal numbers are:

  1, 5, 12, 22, 35, 51, 70, 92, 117, 145, ...

  It can be seen that P(4) + P(7) = 22 + 70 = 92 = P(8). However,
  their difference,  70 − 22 = 48, is not pentagonal.

  Find the pair of pentagonal numbers, P(j) and P(k), for which their sum
  and difference is pentagonal and D = |P(k) − P(j)| is minimised; what
  is the value of D?
  """

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

=#

include("Euler.jl")

# 0.04571280s
function euler44a()
    s = (1:2500).|>n->floor(Int,n*(3*n-1) / 2)
    t = Set(s)
    d = 10_000_000
    for j in s
        for k in s
            a = j+k
            b = abs(j-k)
            if j < k &&
                a < d &&
                a in t &&
                b < d &&
                b in t

                d = b
                break
            end
        end
    end

    return d
end

# 0.04486643s
function euler44b()
    n = 2500
    s = (1:2500).|>n->floor(Int,n*(3*n-1) / 2)
    t = Set(s)
    d = 10_000_000

    for j in 1:n
        for k = 1:n
            if s[j] >= s[k]
                continue;
            end
            a = s[j]+s[k]
            if a >= d
                continue
            end
            b = abs(s[j]-s[k])
            if b >= d
                continue
            end
            if a in t && b in t
                d = b
                break
            end
        end
    end

    return d
end

# run_euler(euler44a)
run_euler(euler44b)


# Tests related to validated integration

using TaylorModels
using LinearAlgebra: norm
using Test

# const _num_tests = 1000
# const α_mid = TaylorModels.α_mid

setformat(:full)

# This is based on an example from Buenger, Numer Algor (2018) 78:1001–1017,
# which shows the efectivity of shrink-wrapping
@testset "Testing `shrink_wrapping` 1" begin
    _order = 2
    set_variables("x", numvars=1, order=2*_order)
    x0 = IntervalBox(0..0, 1)
    dom = IntervalBox(-1..1, 1)
    for δ in 1/16:1/16:1
        Δ = @interval(-δ, δ)
        p = TaylorModelN(TaylorN(1, order=_order), Δ, x0, dom)
        # Usual TM case
        result = p * p
        rem_result = δ*(2+δ)*dom[1]
        range_result = dom[1]^2 + rem_result
        @test result(dom) == range_result
        @test remainder(result) == rem_result
        # Shrink-wrapping:
        q = [ p ]
        TaylorModels.shrink_wrapping!(q)
        @test p(dom) == q[1](dom)
        @test remainder(q[1]) ⊆ Δ
        @test remainder(q[1]) == x0[1]
        # Result using shrink-wrapped variables
        result_sw = q[1] * q[1]
        @test remainder(result_sw) ⊆ rem_result
        @test result_sw(dom) ⊆ range_result
    end
end

# This is based on the example of Sect 4 of M. Berz, K. Makino,
# in https://bt.pa.msu.edu/pub/papers/VIShrink06/VIShrink06.pdf
@testset "Testing `shrink_wrapping` 2" begin
    function two_state1!(v)
        x, y = v
        x_2 = x^2
        y_2 = y^2
        r_2 = x_2 + y_2
        s = sqrt(1 + r_2)
        v[1] = x * s
        v[2] = y * s
        return v
    end

    function two_state2!(v)
        x1, y1 = v
        x1_2 = x1^2
        y1_2 = y1^2
        r1_2 = x1_2 + y1_2
        r1 = sqrt( 1 + 4*r1_2 )
        s1 = 2 / (1+r1)
        ss = sqrt( s1 )
        v[1] = x1 * ss
        v[2] = y1 * ss
        return v
    end

    local B = IntervalBox( -1 .. 1, 2)
    local δ = 0.05
    local ib = δ * B
    local mib = IntervalBox( 0..0, 2)

    # Diverges using naive Interval arithmetic methods
    ib0 = [ib...]
    for iter = 1:210 # iter is twice the number of iterates
        two_state1!(ib0)
        two_state2!(ib0)
    end
    @test minimum(diam.(ib0)) > 1e6

    # Taylor model controls grow and behaves *essentially* as the identity
    # (due to the normalization to the symmetric box) after two iterates
    order = 10
    set_variables("x y", order=2*order)
    vm = [TaylorModelN(normalize_taylor(TaylorN(i, order=order), ib, true),
        0..0, mib, B) for i = 1:2]
    vm0 = deepcopy(vm)
    for iter = 1:1_000
        two_state1!(vm0)
        two_state2!(vm0)
    end
    @test maximum(remainder.(vm0)) < 1e-6 # 2.08e-7
    # Constant and linear terms are identic to the initial ones
    @test all(constant_term.(vm0) .== constant_term.(vm))
    @test all(linear_polynomial.(vm0) .== linear_polynomial.(vm))
    # Maximum difference of polynomials is very small
    @test norm(polynomial(vm0[1] - vm[1]), Inf) < 1e-19

    # Taylor model with shrink-wrapping after two iterates
    vm0 .= deepcopy.(vm)
    for iter = 1:1_000
        two_state1!(vm0)
        two_state2!(vm0)
        TaylorModels.shrink_wrapping!(vm0)
    end
    @test maximum(remainder.(vm0)) < 1e-8 #2.2e-13
    # The dominating difference is in the linear term
    @test norm(polynomial(vm0[1] - vm[1]), Inf) == norm(vm0[1][1] - vm[1][1], Inf)

    # Taylor model with shrink-wrapping after each iterate
    vm0 .= deepcopy.(vm)
    for iter = 1:1_000
        two_state1!(vm0)
        TaylorModels.shrink_wrapping!(vm0)
        two_state2!(vm0)
        TaylorModels.shrink_wrapping!(vm0)
    end
    @test maximum(remainder.(vm0)) < 2.2e-13
    # The dominating difference is in the linear term
    @test norm(polynomial(vm0[1] - vm[1]), Inf) == norm(vm0[1][1] - vm[1][1], Inf)

    # Test AssertionError
    vm = [TaylorModelN(TaylorN(i, order=order), 0..0, mib, ib) for i = 1:2]
    @test_throws AssertionError TaylorModels.shrink_wrapping!(vm)
end


immutable Skellam <: DiscreteUnivariateDistribution
    mu1::Float64
    mu2::Float64
    Skellam(m1::Real, m2::Real) = new(float64(m1), float64(m2))
end

isupperbounded(::Union(Skellam, Type{Skellam})) = false
islowerbounded(::Union(Skellam, Type{Skellam})) = false
isbounded(::Union(Skellam, Type{Skellam})) = false

minimum(::Union(Skellam, Type{Skellam})) = -Inf
maximum(::Union(Skellam, Type{Skellam})) = Inf

insupport(::Skellam, x::Real) = isinteger(x) && zero(x) <= x
insupport(::Type{Skellam}, x::Real) = isinteger(x) && zero(x) <= x

kurtosis(d::Skellam) = 1.0 / (d.mu1 + d.mu2)

mean(d::Skellam) = d.mu1 - d.mu2

function mgf(d::Skellam, t::Real)
    exp(-(d.mu1 + d.mu2) + d.mu1 * exp(t) + d.mu2 * exp(-t))
end

function cf(d::Skellam, t::Real)
    exp(-(d.mu1 + d.mu2) + d.mu1 * exp(im * t) + d.mu2 * exp(im * -t))
end

function pdf(d::Skellam, k::Real)
    isinteger(k) || return 0.0
    exp(-(d.mu1 + d.mu2)) * (d.mu1 / d.mu2)^(k / 2.0) *
        real(besseli(k, 2.0 * sqrt(d.mu1 * d.mu2)))
end

function logpdf(d::Skellam, k::Real)
    isinteger(k) || return 0.0
     -(d.mu1 + d.mu2) + (k / 2.0) * log(d.mu1) - log(d.mu2) +
           log(real(besseli(k, 2.0 * sqrt(d.mu1 * d.mu2))))
end

rand(d::Skellam) = rand(Poisson(d.mu1)) - rand(Poisson(d.mu2))

skewness(d::Skellam) = (d.mu1 - d.mu2) / (d.mu1 + d.mu2)^1.5

var(d::Skellam) = d.mu1 + d.mu2


"""Calculate the R integral from Popel divided by x, (x = q*L0, L0 = zero-temp London depth).

Parameters
----------
a : float

b : float

x : float

Returns
-------
r : float
    The R integral from Popel, divided by x.

Note
----
See R. Pöpel (1989), doi: 10.1063/1.343622 for more details."""
function intR(a, b, x)

    z2 = a^2+b^2

    #for small x
    if x < 0.01*sqrt(z2)
        r = b/(3.0*z2) #This is really r/x

    #for large x
    elseif x > 100*sqrt(z2)
        r = (π*(1+(b^2-a^2)/x^2)/4-b/x)/x #This is really r/x

    #in between x
    else
        #calculate all the terms of r
        r = (1/x^2)*(-0.5*b*x+0.25*a*b*log((z2+x^2+2*a*x)/(z2+x^2-2*a*x))+
        0.25*(x^2+b^2-a^2)*atan2(2*b*x,(z2-x^2)))/x #This is really r/x
    end

    return r
end

"""Calculate the R integral from Popel divided by x, (x = q*L0, L0 = zero-temp London depth).

Parameters
----------
a : float

b : float

x : float

Returns
-------
s : float
    The S integral from Popel, divided by x.

Note
----
See R. Pöpel (1989), doi: 10.1063/1.343622 for more details."""
function intS(a, b, x)

    z2 = a^2+b^2

    #for small x
    if x < 0.01*sqrt(z2)
        s = a/(3.0*z2) #This is really s/x

    #for large x
    elseif x > 100*sqrt(z2)
        s = (a/x - a*b*π/(2*x^2))/x #This is really s/x

    #in between x
    else
        #calculate all the terms of s
        s = (1/x^2)*(0.5*(a*x)+
            0.125*(x^2+b^2-a^2)*log((z2+x^2+2*a*x)/(z2+x^2-2*a*x))-
            0.5*b*a*atan2(2*b*x,(z2-x^2)))/x #This is really s/x
    end

    return s
end

"""Similar to intS, but divided by a. This is specifically useful for when
a=0, as lim(a->0) intS/a should not blow up!"""
function intS_over_a(b, x)

    z2 = b^2

    #for small x
    if x < 0.01*sqrt(z2)
        s = 1.0/(3.0*z2) #This is really s/(x*a)

    #for large x
    elseif x > 100*sqrt(z2)
        s = (1.0/x - b*π/(2*x^2))/x #This is really s/(x*a)

    #in between x
    else
        #calculate all the terms of s
        s = (1/x^2)*(x-0.5*b*atan2(2*b*x,(z2-x^2)))/x #This is really s/(x*a)
    end

    return s
end


using Optim, HTTP, GLM, LinearAlgebra, Random, Statistics, DataFrames, CSV, FreqTables
function q4()
    # Question 1
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS3-gev/nlsw88w.csv"
    df = CSV.read(HTTP.get(url).body, DataFrame)
    X = [df.age df.white df.collgrad]
    Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
    df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
    y = df.occupation

    function mlogit(Beta, X, Z,y)

        K = size(X,2)+1
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigBeta = [reshape(Beta,K,J-1) zeros(K)]

        bigZ = zeros(N,J)
        for j=1:J
            bigZ[:,j] = Z[:,j]-Z[:,J]
        end

        num = zeros(N,J)
        dem = zeros(N)
        for j=1:J
            XZ=cat(X,bigZ[:,j],dims=2)
            num[:,j] = exp.(XZ*bigBeta[:,j])
            dem .+= num[:,j]
        end

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

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

        return loglike
    end

    mlogit_hat_optim = optimize(b-> mlogit(b,X,Z, y), rand(7*(size(X,2)+1)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    println(mlogit_hat_optim.minimizer)

    # Question 2
    println

    # Question 3
    function nlogit(βwc,βbc,λwc,λbc, X, Z,y)

        K = size(X,2)+1
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)

        bigYwc = zeros(N,3)
        for j=1:3
            bigYwc[:,j] = y.==j
        end
        bigβwc = [reshape(βwc,K,3) zeros(K)]

        bigZwc = zeros(N,3)
        for j=1:3
            bigZwc[:,j] = Z[:,j]-Z[:,J]
        end

        numwc1 = zeros(N,3)
        demwc1 = zeros(N)
        numwc = zeros(N,3)
        for j=1:3
            XZwc=cat(X,bigZwc[:,j],dims=2)
            numwc1[:,j] = exp.((XZwc*bigβwc[:,j])./λwc)
            demwc1 .+= numwc1[:,j]
            demwc=(repeat(demwc1,1,3))^λwc
            numwc2=(repeat(demwc1,1,3))^(λwc-1)
            numwc[:,j]=numwc1[:,j].*numwc2
        end

        bigYbc = zeros(N,4)
        for j=1:4
            bigYbc[:,j] = y.==j+3
        end
        bigβbc = [reshape(βbc,K,4) zeros(K)]

        bigZbc = zeros(N,4)
        for j=1:4
            bigZbc[:,j] = Z[:,j+3]-Z[:,J]
        end

        numbc1 = zeros(N,4)
        dembc1 = zeros(N)
        numbc = zeros(N,4)
        for j=1:4
            XZbc=cat(X,bigZbc[:,j],dims=2)
            numbc1[:,j] = exp.((XZbc*bigβbc[:,j])./λbc)
            dembc1 .+= numbc1[:,j]
            dembc=(repeat(dembc1,1,4))^λbc
            numbc2=(repeat(dembc1,1,4))^(λbc-1)
            numbc[:,j]=numbc1[:,j].*numbc2
        end

        num=cat(numwc,numbc,dims=2)
        dem=1+demwc+dembc
        P = num./dem

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

        return loglike
    end

    nlogit_hat_optim = optimize((βwc,βbc,λwc,λbc)->nlogit(βwc,βbc,λwc,λbc, X, Z,y), rand(7*(size(X,2)+1)+2), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    println(nlogit_hat_optim.minimizer)

    end

    # Question 4
    q4()


@testset "Signature" begin
    @test Signature("+++") == Signature(3, 0)
    @test Signature("++-") == Signature(2, 1)
    @test Signature("+-𝟎") == Signature(1, 1, 1)

    @test !is_degenerate(ℝ⁴)
    @test !is_degenerate(spacetime)
    @test is_degenerate(sig_111)

    @test metric(SIGNATURE, Val(1), Val(1)) == 1
    @test metric(SIGNATURE, Val(1), Val(2)) == 0
    @test metric(SIGNATURE, Val(3), Val(3)) == -1
end


# ============================================================
# Theories
# ============================================================

export prim_conserve,
       conserve_prim,
       mixture_prim_conserve,
       mixture_conserve_prim,
       em_coefficients,
       advection_flux,
       burgers_flux,
       euler_flux,
       euler_jacobi
export gauss_moments,
       mixture_gauss_moments,
       moments_conserve,
       diatomic_moments_conserve,
       mixture_moments_conserve,
       pdf_slope,
       mixture_pdf_slope,
       moments_conserve_slope,
       mixture_moments_conserve_slope,
       discrete_moments,
       pressure,
       stress,
       heat_flux,
       maxwellian,
       maxwellian!,
       mixture_maxwellian,
       mixture_maxwellian!,
       shakhov,
       shakhov!,
       rykov!,
       reduce_distribution,
       full_distribution,
       ref_vhs_vis,
       vhs_collision_time,
       rykov_zr,
       aap_hs_collision_time,
       aap_hs_prim,
       aap_hs_diffeq!,
       shift_pdf!,
       hs_boltz_kn,
       kernel_mode,
       boltzmann_fft,
       boltzmann_fft!,
       boltzmann_ode!,
       boltzmann_nuode!,
       bgk_ode!,
       chapman_enskog
export heat_capacity_ratio,
       sound_speed
export sample_maxwell,
       next_collision_time,
       boundary_time

include("theory_continuum.jl")
include("theory_plasma.jl")
include("theory_maxwellian.jl")
include("theory_shakhov.jl")
include("theory_diatomic.jl")
include("theory_atom.jl")
include("theory_pdf.jl")
include("theory_fsm.jl")
include("theory_moments_pure.jl")
include("theory_moments_mixture.jl")
include("theory_thermo.jl")
include("theory_particle.jl")


function fft(input::Vector{Int}, n::Int)::Vector{Int}
    for i in 1:n
        input = fft(input)
    end
    return input
end

function fft(input::Vector{Int})::Vector{Int}
    output = Array{Int}(undef, length(input))
    base_pattern = [0, 1, 0, -1]
    for phase in eachindex(input)
        pattern = collect(Iterators.drop(Iterators.take(Iterators.cycle(repeat(base_pattern, inner=phase)), length(input)+1),1))
        output[phase] = digits(abs(sum(pattern.*input)))[1]
    end
    return output
end

# input = reverse(digits(12345678))
# println(fft(input, 4))
# input = reverse(digits(80871224585914546619083218645595))
# println(fft(input, 100)[1:8])
# input = reverse(digits(19617804207202209144916044189917))
# println(fft(input, 100)[1:8])
# input = reverse(digits(69317163492948606335995924319873))
# println(fft(input, 100)[1:8])
input = reverse(digits(59762677844514231707968927042008409969419694517232887554478298452757853493727797530143429199414189647594938168529426160403829916419900898120019486915163598950118118387983556244981478390010010743564233546700525501778401179747909200321695327752871489034092824259484940223162317182126390835917821347817200000199661513570119976546997597685636902767647085231978254788567397815205001371476581059051537484714721063777591837624734411735276719972310946108792993253386343825492200871313132544437143522345753202438698221420628103468831032567529341170616724533679890049900700498413379538865945324121019550366835772552195421407346881595591791012185841146868209045))
println("Part 1: ", join(fft(input, 100)[1:8]))


"""
    CentralReconstruction{K}

A simple, central, polynomial reconstruction (linear procedure) without adaptive
techniques using polynomials of degee `K-1`.
"""
struct CentralReconstruction{K} <: AbstractReconstruction
    function CentralReconstruction{K}() where {K}
        @argcheck (K>0 && isodd(K)) ArgumentError("`K` must be odd (K = $K).")
        new{K}()
    end
end

function CentralReconstruction(::Val{K}=Val{3}()) where {K}
    CentralReconstruction{K}()
end

function CentralReconstruction(K::Integer)
    CentralReconstruction{K}()
end

@inline stencil_width(::CentralReconstruction{K}) where {K} = K
@inline stencil_width_val(::CentralReconstruction{K}) where {K} = Val{K}()

@inline order(::CentralReconstruction{K}) where {K} = K

function (recons::CentralReconstruction{3})(edge_u, cell, u_m1, u_0, u_p1, balance_law, meshx)
    edge_u[1,cell] = ( 5*u_0 - u_p1 + 2*u_m1 ) / 6
    edge_u[2,cell] = ( 5*u_0 + 2*u_p1 - u_m1 ) / 6

    nothing
end

function (recons::CentralReconstruction{5})(edge_u, cell, u_m2, u_m1, u_0, u_p1, u_p2, balance_law, meshx)
    edge_u[1,cell] = ( 47*u_0 - 13*u_p1 + 2*u_p2 + 27*u_m1 - 3*u_m2 ) / 60
    edge_u[2,cell] = ( 47*u_0 + 27*u_p1 - 3*u_p2 - 13*u_m1 + 2*u_m2 ) / 60

    nothing
end

function (recons::CentralReconstruction{7})(edge_u, cell, u_m3, u_m2, u_m1, u_0, u_p1, u_p2, u_p3, balance_law, meshx)
    edge_u[1,cell] = ( 319*u_0 - 101*u_p1 + 25*u_p2 - 3*u_p3 + 214*u_m1 - 38*u_m2 + 4*u_m3 ) / 420
    edge_u[2,cell] = ( 319*u_0 + 214*u_p1 - 38*u_p2 + 4*u_p3 - 101*u_m1 + 25*u_m2 - 3*u_m3 ) / 420

    nothing
end

function interpolate!(uval, ξ, u_m1, u_0, u_p1, balance_law, mesh, recons::CentralReconstruction{3})
    for i in eachindex(uval)
        uval[i] = @evalpoly(ξ[i], 5*u_0/6 - u_p1/6 + u_m1/3, u_0 - u_m1, -u_0 + u_p1/2 + u_m1/2)
    end

    nothing
end

function interpolate!(uval, ξ, u_m2, u_m1, u_0, u_p1, u_p2, balance_law, mesh, recons::CentralReconstruction{5})
    for i in eachindex(uval)
        uval[i] = @evalpoly(ξ[i], 47*u_0/60 - 13*u_p1/60 + u_p2/30 + 9*u_m1/20 - u_m2/20, 5*u_0/4 - u_p1/12 - 5*u_m1/4 + u_m2/12, -u_0 + 3*u_p1/4 - u_p2/8 + u_m1/4 + u_m2/8, -u_0/2 + u_p1/6 + u_m1/2 - u_m2/6, u_0/4 - u_p1/6 + u_p2/24 - u_m1/6 + u_m2/24)
    end

    nothing
end

function interpolate!(uval, ξ, u_m3, u_m2, u_m1, u_0, u_p1, u_p2, u_p3, balance_law, mesh, recons::CentralReconstruction{7})
    for i in eachindex(uval)
        uval[i] = @evalpoly(ξ[i], 319*u_0/420 - 101*u_p1/420 + 5*u_p2/84 - u_p3/140 + 107*u_m1/210 - 19*u_m2/210 + u_m3/105, 49*u_0/36 - 5*u_p1/36 + u_p2/90 - 49*u_m1/36 + 5*u_m2/36 - u_m3/90, -23*u_0/24 + 7*u_p1/8 - 19*u_p2/80 + 7*u_p3/240 + u_m1/16 + 21*u_m2/80 - u_m3/30, -7*u_0/9 + 11*u_p1/36 - u_p2/36 + 7*u_m1/9 - 11*u_m2/36 + u_m3/36, 23*u_0/72 - 13*u_p1/48 + 5*u_p2/48 - u_p3/72 - u_m1/6 + u_m2/48 + u_m3/144, u_0/12 - u_p1/24 + u_p2/120 - u_m1/12 + u_m2/24 - u_m3/120, -u_0/36 + u_p1/48 - u_p2/120 + u_p3/720 + u_m1/48 - u_m2/120 + u_m3/720)
    end

    nothing
end


using Polynomials, Roots
export get_polynomial
export get_polynomial_coefficients
export compute_fwd_theta
export compute_bnd_rel_bwd_err
export compute_bwd_theta

"""
    p = get_polynomial(graph::Compgraph)

Return the polynomial underlying the computational `graph`. The `graph` is
assumed to involve only linear combinations and multiplications.

`p` is polynomial of type `Polynomial` from the package
[`Polynomials.jl`](https://juliamath.github.io/Polynomials.jl/stable/).

See also [`get_polynomial_coefficients`](@ref).
"""
function get_polynomial(graph::Compgraph)
    graph = Compgraph(Any, graph) # Allow for symbolic values
    p = eval_graph(graph, Polynomial("x"))
    return p
end

"""
    c = get_polynomial_coefficients(graph::Compgraph)

Return the coefficients of the polynomial underlying the computational `graph`.
The `graph` is assumed to involve only linear combinations and multiplications.

The coefficients `c` is a vector containing the coefficients as expressed in
the monomial basis and sorted from the leading coefficient to the constant term.

See also [`get_polynomial`](@ref).
"""
function get_polynomial_coefficients(graph::Compgraph)
    return get_polynomial(graph).coeffs
end

"""
    (e_fwd,theta)=compute_fwd_theta(
        graph::Compgraph{T},
        f;
        coefftype = T,
        tolerance = eps(coefftype) / 2,
        theta_init = big"0.1",
    )

Return relative forward error and corresponding theta for approximation to `f`.

The first output ``e_{fwd}`` is the function that computes the relative forward
error``e_{fwd}(z)=|f(z) - p(z)| / |z|``, where ``p`` is the polynomial
underlying the computational graph `graph`. This is meaningful only if ``p``
approximates ``f`` in some sense.

The second output ``θ`` is the largest positive real number such that
``e_{fwd}(θ) ≤ tolerance``, where ``tolerance`` is by default the unit roundoff
of the data type `coefftype`, which in turn defaults to the type of the
coefficients of `graph`.

The value of ``θ`` is approximated by using the built-in function `fzero` with
starting value set to `theta_init`. By default, this root-finding procedure
uses high-precision arithmetic.
"""
function compute_fwd_theta(
    graph::Compgraph{T},
    f;
    coefftype = T,
    tolerance = eps(coefftype) / 2,
    theta_init = big"0.1",
) where {T}
    # Obtain coefficients of polynomial.
    coeff = get_polynomial_coefficients(graph)

    # Convert coefficients to required type.
    coeff = convert.(coefftype, coeff)
    if isreal(coeff)
        coeff = real(coeff)
    end

    # Find point where the relative forward error equals tolerance.
    p = Polynomial(abs.(coeff), :x)
    e_fwd(z) = abs.(f.(z) - p.(z)) ./ abs.(z)
    g(z) = e_fwd(z) - tolerance
    theta_fwd = fzero(g, theta_init, verbose = true)
    return e_fwd, convert(coefftype, theta_fwd)
end

"""
    e_bwd=compute_bnd_rel_bwd_err(
        f,
        graph::Compgraph{T};
        coefftype = T,
        numterms = 100,
        numdigits = 100,
    )

Compute a bound on the relative backward error of the function `f`.

The returned function ``e_{bwd}`` is a bound on the relative backward error of
the polynomial underlying the computational graph `graph`, seen as an
approximant to `f`. For the underlying polynomial ``p``, the
function returns a bound on `|δ|` where `δ` is such that ``f(z+δ) = p(z)``.

The first argument `f` is a symbol. Currently supported values include:

  - `:exp` The bound is computed by means of the identity δ = log((exp(-z)
    p(z)-1)+1). The code computes the series expansion of the right-hand side of the
    equation, truncates it to the first `numterms` coefficients, bounds each
    coefficient of the ensuing polynomial with its absolute value. The computation
    uses `numdigits` digits of precision.
"""
function compute_bnd_rel_bwd_err(
    f,
    graph::Compgraph{T};
    coefftype = T,
    numterms = 100,
    numdigits = 100,
) where {T}
    if f != :exp
        error("Currently :exp is the only supported function.")
    end
    # Set precision to numdigits decimal digits.
    setprecision(Integer(ceil(log2(10.0^numdigits))))

    # Obtain coefficients of polynomial p.
    coeff = get_polynomial_coefficients(graph)

    # Convert coefficients of p to required type.
    coeff = convert.(coefftype, coeff)
    if isreal(coeff)
        coeff = real(coeff)
    end

    # Series expansion of polynomial approximant p.
    papproximant = Polynomial(coeff)

    # (Truncated) series expansions of exp(-z) and log(1+z).
    expminusz =
        (-1) .^ (0:1:numterms) ./
        factorial.(collect(big(0.0):big(1.0):numterms))
    pexpminusz = Polynomial(expminusz)
    logzplusone = [0; (-1) .^ (0:1:numterms-1) ./ (1:1:numterms)]
    plogzplusone = Polynomial(logzplusone)

    # From exp(z+δ) ≈ p(z), approximate δ ≈ log((exp(-z) p(z)-1)+1).
    # The variables are as follows:
    #    * pexpzpz: the coefficients of exp(-z) p(z)
    #    * presult: the coefficients of log((exp(-z) p(z)-1)+1)
    pexpzpz = Polynomial((pexpminusz*papproximant).coeffs[1:numterms])
    presult = Polynomial((plogzplusone(pexpzpz - 1)).coeffs[1:numterms])

    # Compute bound on |δ|.
    bnd_bwd_err = Polynomial(abs.(presult.coeffs))
    e_bwd(z) = abs.(bnd_bwd_err(z)) ./ abs.(z)

    return e_bwd
end

"""
    (e_bwd,theta)=compute_bwd_theta(;
        bnd_rel_err = compute_bnd_rel_bwd_err(:exp, graph),
        tolerance = eps(coefftype) / 2,
        theta_init = big"0.2",
        use_log = false,
    )

    (e_bwd,theta)=compute_bwd_theta(
        graph::Compgraph{T};
        coefftype = T,
        numterms = 100,
        numdigits = 100,
        tolerance = eps(coefftype) / 2,
        theta_init = big"0.2",
        use_log = false,
    )

Compute bound on relative backward error with corresponding theta.

The first output ``e_{bwd}`` is a function that returns a bound on the relative
backward error of the polynomial underlying the computational graph `graph`,
seen as an approximant to the exponential. For the underlying polynomial ``p``,
the function returns a bound on `|δ|` where `δ` is such that exp(z+δ) = p(z).
This function is computed using `compute_bnd_rel_bwd_err(:graph)`.

The second output ``θ`` is the largest positive real number such that
``e_{bwd}(θ) ≤ tolerance``, where ``tolerance`` is by default the unit roundoff
of the data type `coefftype`, which in turn defaults to the type of the
coefficients of `graph`.

The value of ``theta`` is estimated by approximately solving the equation
``e_{bwd}(z) = tolerance``, using the built-in function `fzero` with starting
value set to `theta_init`. By default this root-finding procedure uses
high-precision arithmetic.

If the kwarg `use_log` is set to `true`, then the value of ``theta`` is computed
by approximating a solution to the equation ``log(e_{bwd}(z)) = log(tolerance)``
instead.

The alternative form accepts a function that returns a bound on the relative
backward error. By default, the function is constructed with
`compute_bnd_rel_bwd_err(:exp, ...)` and the default values for the kwargs in
the first form.
"""
function compute_bwd_theta(;
    bnd_rel_err = compute_bnd_rel_bwd_err(:exp, graph),
    tolerance = eps(coefftype) / 2,
    theta_init = big"0.2",
    use_log = false,
) where {T}
    # Find point where bound on relative backward error equals tolerance.
    e_bwd = bnd_rel_err
    h(z) = use_log ? log(e_bwd(z) / tolerance) : e_bwd(z) - tolerance
    theta_bwd = fzero(h, theta_init)

    return e_bwd, theta_bwd
end

function compute_bwd_theta(
    graph::Compgraph{T};
    coefftype = T,
    numterms = 100,
    numdigits = 100,
    tolerance = eps(coefftype) / 2,
    theta_init = big"0.2",
    use_log = false,
) where {T}
    return compute_bwd_theta(
        bnd_rel_err = compute_bnd_rel_bwd_err(
            :exp,
            graph;
            coefftype = coefftype,
            numterms = numterms,
            numdigits = numdigits,
        ),
        tolerance = tolerance,
        theta_init = theta_init,
        use_log = use_log,
    )
end


using BenchmarkTools, Test, CUDA

a = CUDA.zeros(1024)

function kernel(a)
    i = threadIdx().x
    a[i] += 1
    return
end

@cuda threads=length(a) kernel(a)

##

N = 2^20
x_d = CUDA.fill(1.0f0, N)  # a vector stored on the GPU filled with 1.0 (Float32)
y_d = CUDA.fill(2.0f0, N)  # a vector stored on the GPU filled with 2.0
y_d .+= x_d

function add_broadcast!(y, x)
    CUDA.@sync y .+= x
    return
end

##
@btime add_broadcast!($y_d, $x_d)

##

function gpu_add1!(y, x)
    for i = 1:length(y)
        @inbounds y[i] += x[i]
    end
    return nothing
end

fill!(y_d, 2)
@cuda gpu_add1!(y_d, x_d)
@test all(Array(y_d) .== 3.0f0)

##

function bench_gpu1!(y, x)
    CUDA.@sync begin
        @cuda gpu_add1!(y, x)
    end
end
@btime bench_gpu1!($y_d, $x_d)

##

const nx = 1024  # do 1024 x 1024 2D FFT
xc = CuArray{ComplexF64}(CUDA.randn(Float64, nx, nx))
p = plan_fft!( xc )

##
@btime CUDA.@sync(p * x) setup=(
    x=CuArray{ComplexF64}(CUDA.randn(Float64, nx, nx)));

##
for device in CUDA.devices()
    @show capability(device)
end

##

using AbstractFFTs
using CUDA.CUFFT
##
b = CUDA.rand(ComplexF32,64,64,64)
# pa = plan_fft( a )
@btime fft(b);


using DataDrivenDiffEq
using ModelingToolkit
using LinearAlgebra
@info "Loading DiffEqSensitivity"
using DiffEqSensitivity
@info "Loading Optim"
using Optim
@info "Loading DiffEqFlux"
using DiffEqFlux
@info "Loading Flux"
using Flux
@info "Loading OrdinaryDiffEq"
using OrdinaryDiffEq
using Test
using JLD2
@info "Finished loading packages"

const losses = []

@info "Started Lotka Volterra UODE Testset"
@info "Generate data"

function lotka_volterra(du, u, p, t)
    α, β, γ, δ = p
    du[1] = α*u[1] - β*u[2]*u[1]
    du[2] = γ*u[1]*u[2]  - δ*u[2]
end

tspan = (0.0f0,3.0f0)
u0 = Float32[0.44249296,4.6280594]
p_ = Float32[1.3, 0.9, 0.8, 1.8]
prob = ODEProblem(lotka_volterra, u0,tspan, p_)
solution = solve(prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 0.1)

# Ideal data
tsdata = Array(solution)
# Add noise to the data
noisy_data = tsdata + Float32(1e-5)*randn(eltype(tsdata), size(tsdata))

@info "Setup neural network and auxillary functions"
ann = FastChain(FastDense(2, 32, tanh),FastDense(32, 64, tanh),FastDense(64, 32, tanh), FastDense(32, 2))
p = initial_params(ann)

function dudt_(u, p,t)
    x, y = u
    z = ann(u,p)
    [p_[1]*x + z[1],
    -p_[4]*y + z[2]]
end

prob_nn = ODEProblem(dudt_,u0, tspan, p)
s = solve(prob_nn, Tsit5(), u0 = u0, p = p, saveat = solution.t)

function predict(θ)
    Array(solve(prob_nn, Vern7(), u0 = u0, p = θ, saveat = solution.t,
                         abstol=1e-6, reltol=1e-6))
end

# No regularisation right now
function loss(θ)
    pred = predict(θ)
    sum(abs2, noisy_data .- pred), pred # + 1e-5*sum(sum.(abs, params(ann)))
end

callback(θ,l,pred) = begin
    push!(losses, l)
    if length(losses)%10 == 0
        println("Loss after $(length(losses)) iterations $(losses[end])")
    end
    false
end

if !isfile(joinpath(dirname(@__FILE__), "partial_lotka_volterra.jld2"))
    @info "Initial loss $(loss(p)[1])"
    @info "Train neural network until converged"
    res1 = DiffEqFlux.sciml_train(loss, p, ADAM(0.01), cb=callback, maxiters = 100)
    @info "Finished initial training with loss $(losses[end])"
    res2 = DiffEqFlux.sciml_train(loss, res1.minimizer, BFGS(initial_stepnorm=0.01), cb=callback, maxiters = 10000)
    @info "Finished extended training with loss $(losses[end])"
    p_trained = res2.minimizer
    @save "$(joinpath(dirname(@__FILE__), "partial_lotka_volterra.jld2"))" p_trained

    # Check for the convergence of loss
    # Necessary for result
    @test losses[end] <= 1e-2
else
    @info "Loading pretrained parameters"
    @load "$(joinpath(dirname(@__FILE__), "partial_lotka_volterra.jld2"))" p_trained
end

# Plot the data and the approximation
NNsolution = predict(p_trained)

@test norm(NNsolution .- solution, 2) < 0.1

Y = ann(noisy_data, p_trained)
X = noisy_data
# Create a Basis
@variables u[1:2]
# Lots of polynomials
polys = Any[1]
for i ∈ 1:5
    push!(polys, u[1]^i)
    push!(polys, u[2]^i)
    for j ∈ i:5
        if i != j
            push!(polys, (u[1]^i)*(u[2]^j))
            push!(polys, u[2]^i*u[1]^i)
        end
    end
end

# And some other stuff
h = [cos.(u)...; sin.(u)...; polys...]
basis = Basis(h, u)

# Create an optimizer for the SINDy problem
opt = SR3()
# Create the thresholds which should be used in the search process
λ = exp10.(-10:0.05:-0.5)
# Target function to choose the results from; x = L0 of coefficients and L2-Error of the model
g(x) = x[1] < 1 ? Inf : norm(x, 2)
@info "Start SINDy regression with unknown threshold"
# Test on uode derivative data
Ψ = SINDy(X[:, 2:end], Y[:, 2:end], basis, λ, opt, g = g, maxiter = 10000, normalize = true, denoise = true) # Succeed
print_equations(Ψ)
p̂ = parameters(Ψ)
@info "Build initial guess system"
# The parameters are a bit off, so we reiterate another SINDy term to get closer to the ground truth
# Create function
unknown_sys = ODESystem(Ψ)
unknown_eq = ODEFunction(unknown_sys)
# Just the equations
b = Basis((u, p, t)->unknown_eq(u, ones(size(p̂)), t), u)
# Test on uode derivative data
@info "Refine the guess"
Ψ = SINDy(X[:, 2:end], Y[:, 2:end],b, SR3(0.1), maxiter = 1000) # Succeed
p̂ = parameters(Ψ)

@info "Checking equations"
found_basis = map(x->simplify(x.rhs), equations(Ψ.equations))
pps = parameters(Ψ.equations)

expected_eqs = Vector{Any}()
for _p in pps
    push!(expected_eqs, _p*u[1]*u[2])
end
@test all(isequal.(found_basis, expected_eqs))
@test isapprox(abs.(p̂), p_[2:3], atol = 9e-2)

# The parameters are a bit off, so we reiterate another SINDy term to get closer to the ground truth
# Create function
unknown_sys = ODESystem(Ψ)
unknown_eq = ODEFunction(unknown_sys)

# Build a ODE for the estimated system
function approx(du, u, p, t)
    # Add SINDy Term
    α, δ, β, γ = p
    z = unknown_eq(u, [β; γ], t)
    du[1] = α*u[1] + z[1]
    du[2] = -δ*u[2] + z[2]
end
@info "Simulate system"
# Create the approximated problem and solution
ps = [p_[[1,4]]; p̂]
a_prob = ODEProblem(approx, u0, tspan, ps)
a_solution = solve(a_prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = solution.t)
@test norm(a_solution-solution, 2) < 5e-1
@test norm(a_solution-solution, Inf) < 1.5e-1

@info "Finished"


using SchumakerSpline
x = Array{Union{Missing,Float64}}(collect(0.0:0.01:2.0))
y = Array{Union{Missing,Float64}}(sqrt.(x))
y[5] = missing
x[2] = missing
s1 = Schumaker{Float64}(x, y)

# This should plot something over [0,2]
plt = plot(s1)
# This should plot something over [0,1] but only two points because it is input as an array and this is the exact array of plotpoints.
plt = plot(s1, [0.0,1.0])
# Now it is a tuple and so there will be grid_len intermediate points put it (200 by default)
plt = plot(s1, (0.0,1.0))
# And we can plot the derivatives too
plt = plot(s1, (0.0,1.0); derivs = true)
# This should add a second spline just below with same derivaitves
glt = plot(s1 -10.0, [0.5,1.0]; plot_options = (label = "shifted",), deriv_plot_options = (label = "shifted deriv1",), deriv2_plot_options = (label = "shifted deriv2",), plt = plt)
# This should only plot the first spline.
qlt = plot(s1 + 0.2, (0.0,1.0); plot_options = (label = "shifted",))

ss = Array{Schumaker,1}([s1, s1+0.2, s1*0.8+0.2])
plt  = plot(ss, [0.0,1.0])

plt = plot(ss)
plt2 = plot(0-2.0*s1 + 0.1, [0.0,1.0]; derivs = false, plot_options = (label = "new one",), plt = plt)

# I can't really test anything programatically here. Just have to eyeball it.
# So just giving a true so it all runs happily.
true


# Test linear regression

using Regression
using Base.Test

### auxiliary

function _affmul(A::Matrix, x::VecOrMat, bias::Float64)
	d = size(x,1) - 1
	if ndims(x) == 1
		A * x[1:d] + 2.0 * x[d+1]
	else
		A * x[1:d,:] .+ 2.0 * x[d+1,:]
	end
end

### linear least square

A = randn(8, 3)
y = randn(8)
Y = randn(8, 5)

function verify_llsq(A::Matrix, y::VecOrMat, m::Symbol)
	x = llsq(A, y; method=m)
	u = A * x
	@test_approx_eq_eps 0.0 dot(vec(y - u), vec(u)) 1.0e-10

	At = A'
	x2 = llsq(At, y; method=m, trans=true)
	@test_approx_eq_eps x x2 1.0e-12

	x = llsq(A, y; method=m, bias=2.0)
	u = _affmul(A, x, 2.0)
	@test_approx_eq_eps 0.0 dot(vec(y - u), vec(u)) 1.0e-10

	x2 = llsq(At, y; method=m, bias=2.0, trans=true)
	@test_approx_eq_eps x x2 1.0e-12
end

println("    on llsq")
for m in [:qrlq, :orth, :svd]
	verify_llsq(A, y, m)
	verify_llsq(A, Y, m)
end

### Ridge regression

function verify_ridgereg(A::Matrix, y::VecOrMat, r::Float64)
	d = size(A, 2)
	At = A'

	x = ridgereg(A, y, r)
	u = A * x
	z = A' * (u - y) + r * x
	@test_approx_eq_eps zeros(size(x)) z 1.0e-10

	x2 = ridgereg(At, y, r; trans=true)
	@test_approx_eq_eps x x2 1.0e-12

	x = ridgereg(A, y, r; bias=2.0)
	u = _affmul(A, x, 2.0)
	_x = ndims(x) == 1 ? x[1:d] : x[1:d,:]
	z = A' * (u - y) + r * _x
	@test_approx_eq_eps zeros(size(z)) z 1.0e-10

	x2 = ridgereg(At, y, r; bias=2.0, trans=true)
	@test_approx_eq_eps x x2 1.0e-10
end

println("    on ridgereg")
verify_ridgereg(A, y, 0.6)
verify_ridgereg(A, Y, 0.6)


module VND
using StaticArrays

import Base: +, *, /, ==, <, >, norm, dot
export VecN, norm, dist

type VecN
    elem::SVector;
end

sizeof(a::VecN) = length(a.elem)
sOK(a::VecN, b::VecN) = (sizeof(a) == sizeof(b)) ? true :
                         throw(BoundsError("Vector of different lengths"));

(+)(a::VecN, b::VecN) = [a.elem[i] + b.elem[i] for i in 1:sizeof(a) if sOK(a,b)]
(*)(x::Real, a::VecN) = [a.elem[i]*x for i in 1:sizeof(a)]
(*)(a::VecN, x::Real) = x*a
(/)(a::VecN, x::Real) = [a.elem[i]/x for i in 1:sizeof(a)]

(==)(a::VecN, b::VecN) = any([(a.elem[i] == b.elem[i]) for i in 1:sizeof(a) if sOK(a,b)])

dot(a::VecN, b::VecN) = sum([a.elem[i]*b.elem[i] for i in 1:sizeof(a) if sOK(a,b)])
norm(a::VecN) = sqrt(dot(a,a));
dist(a::VecN, b::VecN) = sum(map(x -> x*x,[a.elem[i]-b.elem[i] for i in 1:sizeof(a) if sOK(a,b)]))

end

#
# ----------
#  Examples
# ----------
#

using VND, StaticArrays

a = @SVector [1.1,2.2,3.3];
b = @SVector [2.1,3.2,4.3];

aa = VecN(a);
bb = VecN(b);

aa2 = 2.0*aa
bb3 = bb/3.0

cc = aa + bb
aa == bb

aanorm = norm(aa)
abdist = dist(aa,bb)

#
# Π can be calculated in higher dimensions
# For example the unit 4-ball has a volume of (Π * Π)/2.0
#
# see: https://en.wikipedia.org/wiki/Volume_of_an_n-ball
#


# diffs.jl

import MIRT
using LinearMapsAA: LinearMapAA
using Random: seed!
using BenchmarkTools: @btime

# old 2D versions for comparison with the new N-D versions

"""
    d = diff2d_forw_old(X)

2D finite differences along both dimensions, for anisotropic TV regularization.
Performs the same operations as
d = ``[(I_N \\otimes D_M); (D_N \\otimes I_M)] X[:]``
where ``D_N`` denotes the ``N-1 \\times N`` 1D finite difference matrix
and ``\\otimes`` denotes the Kronecker product,
but does it efficiently
without using `spdiagm` (or any `SparseArrays` function).

in
- `X` `M × N` array (typically a 2D image).
It cannot be a Vector! (But it can be a `M×1` or `1×N` 2D array.)

out
- `d` vector of length `N*(M-1) + (N-1)*M`
"""
function diff2d_forw_old(x::AbstractMatrix{<:Number})
    return [diff(x,dims=1)[:]; diff(x,dims=2)[:]]
end

"""
    z = diff2d_adj_old(d, M, N; out2d=false)

Adjoint of 2D finite differences along both dimensions.
Performs the same operations as
``z = [(I_N \\otimes D_M); (D_N \\otimes I_M)]' * d``
where `D_N` denotes the `N-1 × N` 1D finite difference matrix
and `\\otimes` denotes the Kronecker product, but does it efficiently
without using `spdiagm` (or any `SparseArrays` function).

in
- `d` vector of length `N*(M-1) + (N-1)*M`
- `M,N` desired output size

option
- `out2d` if true then return `M × N` array, else `M*N` vector

out
- `z` `M*N` vector or `M × N` array (typically a 2D image)

"""
function diff2d_adj_old(d::AbstractVector{<:Number}, M::Int, N::Int ; out2d=false)

    length(d) != N*(M-1) + (N-1)*M && throw("length(d)")

    dx = d[1:(N*(M-1))]
    dx = reshape(dx, M-1, N)
    dy = d[1+(N*(M-1)):end]
    dy = reshape(dy, M, N-1)
    zx = [-transpose(dx[1,:]);
        (@views dx[1:end-1,:] - dx[2:end,:]);
        transpose(dx[end,:])]
    zy = [-dy[:,1] (@views dy[:,1:end-1] - dy[:,2:end]) dy[:,end]]
    z = zx + zy # M by N

    return out2d ? z : z[:] # array or vector?
end

"""
    T = diff2d_map_old(M::Int, N::Int)
"""
function diff2d_map_old(M::Int, N::Int)
    return LinearMapAA(
        x -> diff2d_forw_old(reshape(x,M,N)),
        d -> diff2d_adj_old(d, M, N),
        (N*(M-1)+M*(N-1), N*M), (name="diff2_map",),
    )
end

# timing runs
function diff_map_time( ; M::Int=2^7, N::Int=2^7+1)
    seed!(0)
    x = randn(M * N)
    T2d = diff2d_map_old(M, N)
    Tnd = MIRT.diff_map((M, N))

    # check consistency
    d2d = T2d * x
    dnd = Tnd * x
    @assert d2d == dnd

    d = d2d
    y2d = T2d' * d
    ynd = Tnd' * d
    @assert y2d == ynd[:]

    # time diff_forw
    @btime $T2d * $x
    @btime $Tnd * $x

    # time diff_adj
    @btime $T2d' * $d
    @btime $Tnd' * $d
    nothing
end

#= ir71 with 1.4.2
  98.381 μs (16 allocations: 762.73 KiB)
  58.344 μs (18 allocations: 508.77 KiB)
  159.237 μs (39 allocations: 1.00 MiB)
  102.423 μs (115 allocations: 384.05 KiB)
=#

diff_map_time()


@testset "dijkstra" begin

    @testset "small_network" begin
        file = "data/networks/distance/weighted_elist_small.txt"
        graph = read_graph(file, "distance")
        distances, pred = dijkstra(graph, 1)
        @test distances[5] == 36.0
        @test distances[10] == 19.0
        dist_5, path_5 = dijkstra(graph, 1, 5)
        @test dist_5 == 36.0
        @test path_5 == [1, 12, 5]
        dist_10, path_10 = dijkstra(graph, 1, 10)
        @test dist_10 == 19.0
        @test path_10 == [1, 10]
    end

    @testset "large_network" begin
        file = "data/networks/distance/weighted_elist_large.txt"
        graph = read_graph(file, "distance")
        distances, pred = dijkstra(graph, 1)
        @test distances[130] == 49.0
        @test distances[180] == 25.0
        @test distances[200] == 46.0
        dist_130, path_130 = dijkstra(graph, 1, 130)
        @test dist_130 == 49.0
        @test path_130 == [1, 96, 30, 130]
        dist_200, path_200 = dijkstra(graph, 1, 200)
        @test dist_200 == 46.0
        @test path_200 == [1, 96, 140, 14, 200]
    end

    @testset "vlarge_network" begin
        file = "data/networks/distance/weighted_elist_vlarge.txt"
        graph = read_graph(file, "distance")
        distances, pred = dijkstra(graph, 1)
        @test distances[250] == 19.0
        @test distances[1000] == 20.0
        @test distances[1500] == 16.0
        dist_250, path_250 = dijkstra(graph, 1, 250)
        @test dist_250 == 19.0
        @test path_250 == [1, 293, 680, 139, 1439, 1109, 250]
        dist_1500, path_1500 = dijkstra(graph, 1, 1500)
        @test dist_1500 == 16.0
        @test path_1500 == [1, 293, 680, 139, 1371, 980, 1500]
    end

end


# Utilities for demo.jl
module DemoTools

using LinearAlgebra
using TensorKit
using JLD2
using MERAKit

"""
    block_op(op, num_sites)

Take an operator `op` that defines the local term of a global operator, and block sites
together to form a new operator for which each site corresponds to `num_sites` sites of the
original operator, and which sums up to the same global operator. `num_sites` should be a
power of 2.
"""
function block_op(op::MERAKit.SquareTensorMap{2}, num_sites)
    while num_sites > 1
        VL = space(op, 1)
        VR = space(op, 2)
        eyeL = id(VL)
        eyeR = id(VR)
        opcross = eyeR ⊗ op ⊗ eyeL
        opleft = op ⊗ eyeL ⊗ eyeR
        opright = eyeL ⊗ eyeR ⊗ op
        op_new_unfused = opcross + 0.5*(opleft + opright)
        fusionspace = domain(op)
        fusetop = isomorphism(fuse(fusionspace), fusionspace)
        fusebottom = isomorphism(fusionspace, fuse(fusionspace))
        op = (fusetop ⊗ fusetop) * op_new_unfused * (fusebottom ⊗ fusebottom)
        num_sites /= 2
    end
    if num_sites != 1
        msg = "`num_sites` needs to be a power of 2"
        throw(ArgumentError(msg))
    end
    return op
end

function block_op(op::MERAKit.SquareTensorMap{1}, num_sites)
    while num_sites > 1
        V = space(op, 1)
        eye = id(V)
        opleft = op ⊗ eye
        opright = eye ⊗ op
        op_new_unfused = opleft + opright
        fusionspace = domain(op_new_unfused)
        fusetop = isomorphism(fuse(fusionspace), fusionspace)
        fusebottom = isomorphism(fusionspace, fuse(fusionspace))
        op = fusetop * op_new_unfused * fusebottom
        num_sites /= 2
    end
    if num_sites != 1
        msg = "`num_sites` needs to be a power of 2"
        throw(ArgumentError(msg))
    end
    return op
end

function goldenchain_hamiltonian(JI, Jτ; block_size = 1)
    V = GradedSpace[FibonacciAnyon](:I => 0, :τ => 1)
    H = TensorMap(zeros, Float64, V*V ← V*V)
    H.data[FibonacciAnyon(:I)] .= JI
    H.data[FibonacciAnyon(:τ)] .= Jτ
    H = block_op(H, block_size)
    H = H / block_size
    return H
end

"""
    XXZ_hamiltonian(J_xy, J_z; symmetry=:none, block_size=1)

Return the local Hamiltonian term for the XXZ model: J_xy*XX + J_xy*YY + J_z*ZZ.

`symmetry` should be `:none` or `:group`, and determmines whether the Hamiltonian should be
an explicitly U(1) symmetric `TensorMap` or a dense one. `block_size` determines how many
sites to block together, and should be a power of 2. If sites are blocked, then the
Hamiltonian is divided by a constant to make the energy per site be the same as for the
original.
"""
function XXZ_hamiltonian(J_xy, J_z; symmetry=:none, block_size=1)
    if symmetry == :U1 || symmetry == :group
        V = ℂ[U₁](-1=>1, 1=>1)
        Z = TensorMap(zeros, Float64, V ← V)
        Z.data[U₁(1)] .= 1.0
        Z.data[U₁(-1)] .= -1.0
        ZZ = Z ⊗ Z
        XXplusYY = TensorMap(zeros, Float64, V ⊗ V ← V ⊗ V)
        XXplusYY.data[U₁(0)] .= [0.0 2.0; 2.0 0.0]
    elseif symmetry == :none
        V = ℂ^2
        Z = TensorMap(zeros, Float64, V ← V)
        Z.data .= [1.0 0.0; 0.0 -1.0]
        ZZ = Z ⊗ Z
        XXplusYY = TensorMap(zeros, Float64, V ⊗ V ← V ⊗ V)
        XXplusYY.data[2,3] = 2.0
        XXplusYY.data[3,2] = 2.0
    else
        error("Unknown symmetry: $symmetry")
    end
    H = J_xy*XXplusYY + J_z*ZZ
    H = block_op(H, block_size)
    H = H / block_size
    return H
end

"""
    ising_hamiltonian(J=-1.0, h=1.0; symmetry=:none, block_size=1)

Return the local Hamiltonian term for the Ising model: J*XX + h*Z

`symmetry` should be `:none`, `:group`, or `:anyons` and determmines whether the Hamiltonian
should be an explicitly Z2 symmetric or anyonic `TensorMap`, or a dense one. `block_size`
determines how many sites to block together, and should be a power of 2. If sites are
blocked, then the Hamiltonian is divided by a constant to make the energy per site be the
same as for the original.
"""
function ising_hamiltonian(J=-1.0, h=1.0; symmetry=:none, block_size=1)
    if symmetry == :Z2
        V = ℂ[ℤ₂](0=>1, 1=>1)
        # Pauli Z
        Z = TensorMap(zeros, Float64, V ← V)
        Z.data[ℤ₂(0)] .= 1.0
        Z.data[ℤ₂(1)] .= -1.0
        eye = id(V)
        ZI = Z ⊗ eye
        IZ = eye ⊗ Z
        # Pauli XX
        XX = TensorMap(zeros, Float64, V ⊗ V ← V ⊗ V)
        XX.data[ℤ₂(0)] .= [0.0 1.0; 1.0 0.0]
        XX.data[ℤ₂(1)] .= [0.0 1.0; 1.0 0.0]
        H = J*XX + h/2 * (ZI+IZ)
    elseif symmetry == :anyons
        V = RepresentationSpace{IsingAnyon}(:I => 0, :ψ => 0, :σ => 1)
        H = TensorMap(zeros, Float64, V ⊗ V ← V ⊗ V)
        if J == 1.0 && h == 1.0
            H.data[IsingAnyon(:I)] .= 2.0
            H.data[IsingAnyon(:ψ)] .= -2.0
        elseif J == -1.0 && h == 1.0
            H.data[IsingAnyon(:I)] .= -2.0
            H.data[IsingAnyon(:ψ)] .= 2.0
        else
            msg = "Can't create an anyon symmetric Ising Hamiltonian that isn't critical."
            throw(ArgumentError(msg))
        end
    elseif symmetry == :none
        V = ℂ^2
        # Pauli matrices
        X = TensorMap(zeros, Float64, V ← V)
        Z = TensorMap(zeros, Float64, V ← V)
        eye = id(V)
        X.data .= [0.0 1.0; 1.0 0.0]
        Z.data .= [1.0 0.0; 0.0 -1.0]
        XX = X ⊗ X
        ZI = Z ⊗ eye
        IZ = eye ⊗ Z
        H = J*XX + h/2 * (ZI+IZ)
    else
        error("Unknown symmetry: $symmetry")
    end
    H = block_op(H, block_size)
    H = H / block_size
    return H
end

"""
    magop(; symmetry = :none, block_size = 1)

Return the magnetization operator for the Ising model, blocked over `block_size` sites.
"""
function magop(; symmetry = :none, block_size = 1)
    if symmetry == :none
        V = ℂ^2
    elseif symmetry == :group || symmetry == :Z2
        V = Z2Space(0 => 1, 1 => 1)
    else
        msg = "Unknown symmetry: $symmetry"
        raise(ArgumentError(msg))
    end
    X = TensorMap(zeros, Float64, V ← V)
    if symmetry == :none
        X.data .= [0.0 1.0; 1.0 0.0]
    end
    magop = block_op(X, block_size)
    return magop
end

"""
    store_mera(path, m)

Store the MERA `m` on disk at `path`.

The storage format is a JLD2 file of `pseudoserialize(m)`.  We use pseudoserialization to
both work around limitations of JLD(2) and improve compatibility between different versions
of MERAKit.jl

See also: [`load_mera`](@ref)
"""
function store_mera(path, m)
    deser = pseudoserialize(m)
    @save path deser
end

"""
    load_mera(path)

Load the MERA at stored at `path`.

The storage format is a JLD2 file of `pseudoserialize(m)`.  We use pseudoserialization to
both work around limitations of JLD(2) and improve compatibility between different versions
of MERAKit.jl

See also: [`store_mera`](@ref)
"""
function load_mera(path)
    @load path deser
    m = depseudoserialize(deser...)
    return m
end

end  # module


using Printf

function readNums(filePath)
    lines = filter(l -> length(l) > 1, split(open(f->read(f, String), filePath), "\n"))
    nums = map(ln -> [parse(Int8, char) for char in split(ln, "")], lines)
    return hcat(nums...)'
end

function solution1(nums)
    mask = trues(size(nums))
    # left to right
    for column = 1:(size(nums, 2)-1)
        mask[:,column] = mask[:,column] .& (nums[:,column] .< nums[:,column+1])
    end
    # right to left
    for column = 2:size(nums, 2)
        mask[:,column] = mask[:,column] .& (nums[:,column] .< nums[:,column-1])
    end
    # top to bottom
    for row = 1:(size(nums, 1)-1)
        mask[row,:] = mask[row,:] .& (nums[row,:] .< nums[row+1,:])
    end
    # bottom to top
    for row = 2:size(nums, 1)
        mask[row,:] = mask[row,:] .& (nums[row,:] .< nums[row-1,:])
    end
    return sum(nums[mask] .+ 1)
end

function solution2(nums)
    borders = nums .== 9
    function bfs(p_row, p_col)
        # out of bounds
        if (p_row < 1 || p_row > size(borders, 1) || p_col < 1 || p_col > size(borders, 2))
            return 0
        end
        # already visited or border
        if borders[p_row, p_col]
            return 0
        end
        borders[p_row, p_col] = 1
        return (1
                + bfs(p_row + 1, p_col    )
                + bfs(p_row - 1, p_col    )
                + bfs(p_row    , p_col + 1)
                + bfs(p_row    , p_col - 1))
    end
    basin_counts = zeros(Int64, 0)
    for col = 1:(size(borders, 2))
        for row = 1:(size(borders, 1))
            res = bfs(row, col)
            if res != 0
                append!(basin_counts, res)
            end
        end
    end
    largest3 = sort(basin_counts)[end-2:end]
    return largest3[1] * largest3[2] * largest3[3]
end

if length(ARGS) < 1
	println("Provide path of a file you'd like to solve")
	exit(1)
end

nums = readNums(ARGS[1])
println("Day 9:");
Printf.@printf "Solution 1: %d\n" solution1(nums)
Printf.@printf "Solution 2: %d\n" solution2(nums)


"""
    time-varying LQR
"""
function tvlqr(A, B, Q, R)
    T = length(Q)

    P = [zero(A[1]) for t = 1:T]
    K = [zero(B[1]') for t = 1:T-1]
    P[T] = Q[T]

    for t = T-1:-1:1
        K[t] = (R[t] + B[t]' * P[t+1] *  B[t]) \ (B[t]' * P[t+1] * A[t])
        P[t] = (Q[t] + K[t]' * R[t] * K[t]
                + (A[t] - B[t] * K[t])' * P[t+1] * (A[t] - B[t] * K[t]))
    end

    return K, P
end

function tvlqr(model::Model{<: Integration, FixedTime}, x, u, h, Q, R)
    A, B = jacobians(model, x, u, h)
    K, P = tvlqr(A, B, Q, R)
    return K, P
end

function tvlqr(model::Model{<: Integration, FreeTime}, x, u, h, Q, R)
    A, B = jacobians(model, x, u, h)
    K, P = tvlqr(A, B, Q, [r[1:end-1, 1:end-1] for r in R])
    return K, P
end

"""
    jacobians along trajectory
"""
function _jacobians(model, x, u, h)
    A = []
    B = []

    w = zeros(model.d)

    T = length(x)

    for t = 1:T-1
        xt = x[t]
        ut = u[t]
        xt⁺ = x[t+1]

        fx(z) = fd(model, xt⁺, z, ut, w, h, t)
        fu(z) = fd(model, xt⁺, xt, z, w, h, t)
        fx⁺(z) = fd(model, z, xt, ut, w, h, t)

        A⁺ = ForwardDiff.jacobian(fx⁺, xt⁺)
        push!(A, -1.0 * A⁺ \ ForwardDiff.jacobian(fx, xt))
        push!(B, -1.0 * A⁺ \ ForwardDiff.jacobian(fu, ut))
    end

    return A, B
end

jacobians(model::Model{<: Integration, FixedTime}, x, u, h) = _jacobians(model, x, u, h)

function jacobians(model::Model{<: Integration, FreeTime}, x, u, h)
    A, B = _jacobians(model, x, u, h)
    T = length(x)
    return A, [B[t][:, 1:end-1] for t = 1:T-1]
end

"""
    projection
        use time-varying LQR to get dynamically feasible solution
"""
function lqr_projection(model::Model{<: Integration, FixedTime}, x̄, ū, h̄, Q, R)
    K, P = tvlqr(model, x̄, ū, h̄, Q, R)

    x_proj = [copy(x̄[1])]
    u_proj = []

    T = length(x̄)

    for t = 1:T-1
        push!(u_proj, ū[t] - K[t] * (x_proj[end] - x̄[t]))
        push!(x_proj, propagate_dynamics(model, x_proj[end], u_proj[end],
            zeros(model.d), h̄[1], t, tol_r = 1.0e-12, tol_d = 1.0e-12))
    end

    return x_proj, u_proj
end

function lqr_projection(model::Model{<: Integration, FreeTime}, x̄, ū, h̄, Q, R)
    K, P = tvlqr(model, x̄, ū, h̄, Q, R)

    x_proj = [copy(x̄[1])]
    u_proj = []

    T = length(x̄)

    for t = 1:T-1
        push!(u_proj, [ū[t][1:end-1] - K[t] * (x_proj[end] - x̄[t]); h̄[1]])
        push!(x_proj, propagate_dynamics(model, x_proj[end], u_proj[end],
            zeros(model.d), h̄[1], t, tol_r = 1.0e-12, tol_d = 1.0e-12))
    end

    return x_proj, u_proj
end


"""
is_GLVlocallystable(A, r)

Computes if the pair (A,r) has a feasible interior equilibrium that is stable.

# Arguments
- `A::Array{Real,2}`: community matrix
- `r::Array{Real, 1}`: growth vector

# Outputs
- `stability::Bool`: true = stable, false = not stable
"""
function is_GLVlocallystable(A::Array{<:Real,2}, r::Array{<:Real,1})::Bool
    N = size(A, 1)

    # compute interior equilibrium for the local community
    xss = -pinv(A)*r # r[1,:] instead of r, necessary when r is an Array{Float64, 2}

    # Check if equilibrium is stable.
    # First, check if xss < 0 (i.e., it unfeasible)
    if (minimum(xss) < 0)
              coexist = false
    else # it is feasible
              # check for stability
              M = diagm(xss)*A #community matrix
              λmax = maximum(real(eigvals(M)))

              coexist = λmax < 0 ? true : false

    end
    return coexist
end

"""
is_GLVpermanent(A, r; <keyword arguments>)

Computes permanence of the pair (A,r) using Jansen's criterion of mutual invasibility.
is_permanente(A,r;) -> Boolean

# Arguments
- `A::Array{<:Real,2}`: community matrix
- `r::Array{<:Real, 1}`: growth vector
- `regularization::Real=0`:
- `z_tolerance::Real=-1e-60`:
- `iterations::Integer=1e4`:

# Outputs
- `permanence::Bool`: true = permanent, false = not permanent
"""
function is_GLVpermanent(A::Array{<:Real, 2}, r::Array{<:Real, 1}; regularization::Real = 0, z_tolerance::Real = -1e-60, iterations::Int = Int(1e4))::Bool
    N = size(A, 1)

    # Regularize A by adding a small negative diagonal term
    regularization > 0 ? Atemp = A .+ Diagonal(Distributions.rand(Uniform(-1.0*regularization, 0), N)) : Atemp = copy(A)

    if N == 1 # just scalar equation
        coexistence = Atemp[1] < 0 && r[1] > 0 ?  true : false
    else # N > 1

        # check if it has an interior feasible equilibria
        if rank(Atemp) .== rank([Atemp r]) # if solution exists

            xeqint = -pinv(Atemp)*r

            if minimum(xeqint) > 0 # if interior is feasible

                equilibria = []
                @inbounds for localcommunityID = 1: 2^N-1

                #@inbounds @showprogress for localcommunityID = 1: 2^N-1

                    # calculate present species for that local communtiy
                    z = digits(localcommunityID, base = 2, pad = N) # integer to binary
                    indx = findall( z .== 1)

                    # build local (Alocal, rlocal) with only the present species
                    Alocal = Atemp[indx,indx]
                    rlocal = r[indx]

                    # Check if there exists a solution to Alocal*x + rlocal
                    if rank(Alocal) .== rank([Alocal rlocal]) # solution exists
                        #@show localcommunityID
                        xss = -pinv(Alocal)*rlocal

                        if (minimum(xss) > 0) # if solution is feasible
                            #@show localcommunityID
                            xtemp = zeros(N)
                            xtemp[indx] = xss
                            push!(equilibria, xtemp)
                            #equilibria = hcat(equilibria, convert(Array{Float64,2}, xtemp))
                        else
                        end

                    else # solution does not exists

                    end

                end
                Xss = hcat(equilibria...)

                if size(Xss, 2) >= 1 # at least one boundary equilibrium
                    # Test for Jansen's criterion of permanence
                    h = Variable(N)
                    z = z_tolerance

                    M = Atemp*Xss
                    v = ones(1, size(Xss,2))
                    #problem = minimize(1, [h'*M + (h'*r)*v + z*v >= 0, h >= 1e-60, sum(h) == 1])
                    problem = minimize(1, [h'*M + (h'*r)*v + z*v >= 0, h >= 1e-60, sum(h) == 1])
                    # Solve the problem by calling solve!
                    #Convex.solve!(problem, SCSSolver(max_iters = iterations, verbose=false))
                    # Convex.solve!(problem, () -> SCS.Optimizer(verbose = false, max_iters = iterations), verbose=false)
                    solve!(problem, () -> Optimizer(verbose = false, max_iters = iterations), verbose=false)

                    # Check if problem.status = "OPTIMAL::TerminationStatusCode = 1"
                    Integer(problem.status) .== 1 ? coexistence = true : coexistence = false
                else # no boundary equilibrium, true by default
                    coexistence = true
                end

            else # interior is not feasible
                coexistence = false
            end

        else # if there is no interior equilibrium
            coexistence = false
        end
    end

    return coexistence
end

"""

assembly_hypergraph_GLV(A,r; <keyword arguments>)
Computes the assembly hypergraph for (A,r).

# Arguments
- `A::Array{<:Real,2}`: community matrix
- `r::Array{<:Real, 1}`: growth vector
- `method::String="permanence"`: will return a randomized growthvector using one of the following methods
    - `"permanence"`: uses the is_GLVpermantent(@ref) to determine interspecies coexistance
    - `"localstability"`: uses the is_GLVlocallystable(@ref) to determine interspecies coexistance
- `regularization::Real=0`:

# Outputs
- `hypergraph::Array{Array{Int64,1},1}`: array of hyper edges: species in edge => species coexist
"""
function assembly_hypergraph_GLV(A::Array{<:Real,2}, r::Array{<:Real,1}; method::String = "permanence", regularization::Real = 0)::Array{Array{Int64,1},1}
    N = size(A, 1)

    hyperdges = Array{Int64, 1}[]
    #hyperdges = []

    if method .== "permanence"
        print("\nCalculating coexistence hypergraph (method: permanence) \n")
        @inbounds for community_id = 1 : 2^N-1
        #@inbounds @showprogress 0.0001 for community_id = 1 : 2^N-1
            z = digits(community_id, base = 2, pad = N) # integer to binary
            indx = findall( z .== 1)

            if is_GLVpermanent(A[indx, indx], r[indx]; regularization = regularization) .== true
                push!(hyperdges, indx)
            end
        end

    elseif method .== "localstability"
        print("\nCalculating coexistence hypergraph (method: local stability) \n")
        @inbounds @showprogress 0.0001 for community_id = 1 : 2^N-1
            z = digits(community_id, base = 2, pad = N) # integer to binary
            indx = findall( z .== 1)

            if is_GLVlocallystable(A[indx, indx], r[indx]) .== true
                push!(hyperdges, indx)
            end
        end
    end
    return hyperdges

end

"""
save_hypergraph_dat(file, H)

Saves the hypergraph H into file as a list of hyperedges

# Arguments
- `file::String`: full file name (including the path where you want to save it)
- `H::Array{Array{Int64,1},1}`: hypergraph

See also: [`assembly_hypergraph_GLV`](@ref)
"""
function save_hypergraph_dat(file::String, H::Array{Array{Int64,1},1})::Nothing
    out = [filter(x -> !isspace(x), chop(repr(h), head = 1, tail =1))*"\n" for h in H] |> join
    path = dirname(file);
    if startswith(path, '~')
        path = replace(path, '~' => homedir())
        file = replace(file, '~' => homedir())
    end
    mkpath(path)
    io = open(file, "w")
    write(io, out)
    close(io)
end

"""
random_growthvector(N, μ, σ; seed=nothing)

Generates a random growth vector (the "r" vector in the generalized
[Lotka-Voltera](https://en.wikipedia.org/wiki/Generalized_Lotka%E2%80%93Volterra_equation)  equation)

# Arguments
- `N::Real`: length of returned growth vector
- `μ::Real`: mean of LogNormal distribution used to generate each value
- `σ::Real`: standard deviation of LogNormal distribution used to generate each value
- `seed::Union{Nothing, <:Int}=nothing`: if specified, this seed will be used in the random number generator, allowing reproducibility

# Output
- `growth_vector::Array{<:Real,1}`
See also: [`randomize_growthvector`](@ref)
"""
function random_growthvector(N::Int, μ::Real, σ::Real; seed::Union{<:Signed, Nothing}=nothing)::Array{<:Real,1}
    N = Int(N);
    if seed != nothing; seed!(seed); end
    return Distributions.rand(LogNormal(μ, σ), N)

end

"""
randomize_growthvector(r; <keyword arguments>)

# Arguments
- `r::Array{<:Real,1}`: growth vector (this can be generated randomly by [`random_growthvector`](@ref))
- `method::String="preserve_norm"`: will return a randomized growthvector using one of the following methods
    - `"preserve_norm"`: generated using a normal distribution for each entry, and then scaled to have the same norm as growth vector input (r)
    - `"shuffle"`: randomly permute growth vector input (r)
    - `"sample"`: randomly sample (with replacement) entries of growth vector input (r)
    - `"preserve_sign_shuffle"`: same as "shuffle" but the signs are not modified
    - `"preserve_sign_sample"`: same as "sample" but the signs are not modified
- `seed::Union{Nothing, <:Int}=nothing`: if specified, this seed will be used in the random number generator, allowing reproducibility

# Output
- `growth_vector::Array{<:Real,1}`

See also: [`random_growthvector`](@ref)
"""
function randomize_growthvector(r::Array{<:Real,1}; method::String="preserve_norm", seed::Union{Nothing, <:Signed}=nothing)::Array{<:Real,1}
    if seed != nothing; seed!(seed); end # set seed

    if  method == "preserve_norm"               # new random growth vector but same norm as r
        v = Distributions.rand(Normal(), length(r))
        return norm(r,2).*v/norm(v,2)
    elseif method == "preserve_sign_sample"     # unifmormly sampled from values of r, preserve signs
        values = abs.(r)
        return sign.(r).*sample(values, length(r))
    elseif method == "preserve_sign_shuffle"    # randomly permuted from values of r, preserve signs
        values = abs.(r)
        return sign.(r).*shuffle(values)
    elseif method == "shuffle"                  # randomly permuted from values of r
        return shuffle(r)
    elseif method == "sample"                   # unifmormly sampled from values of r
        return sample(r, length(r))
    end

    @warn "Method Not Recognized: returning original growth vector"
    return r
end

"""
random_communitymatrix(N, σ, p)

Generates a random community matrix (the "A" matrix in the generalized
[Lotka-Voltera](https://en.wikipedia.org/wiki/Generalized_Lotka%E2%80%93Volterra_equation) equation).
entries according to a Bernoulli distribution with success rate parameter, p (the
entries that are not "populated" are set to 0)

# Arguments
- `N::Real`: dimension of returned square matrix (N x N)
- `σ::Real`: standard deviation of normal distribution used to generate each entry (μ = 0)
- `p::Real`: success rate of Bernoulli distribution used to populate the returned matrix
- `seed::Union{Nothing, <:Int}=nothing`: if specified, this seed will be used in the random number generator, allowing reproducibility

# Output
- `community_matrix::Array{<:Real,2}`

See also: [`randomize_communitymatrix`](@ref)
"""
function random_communitymatrix(N::Number, σ::Real, p::Real; seed::Union{Nothing, <:Signed}=nothing)::Array{<:Real,2}
    N = Int(N);
    if seed != nothing; seed!(seed); end # set seed
    W = Distributions.rand(Normal(0, σ), N, N)
    Z = Distributions.rand(Bernoulli(p), N, N)
    temp = W.*Z
    temp[diagind(temp)] .= -1

    return temp
end

"""
randomize_communitymatrix(A; <keyword arguments>)

# Arguments
- `r::Array{<:Real,2}`: growth vector (this can be generated randomly by [`random_growthvector`](@ref))
- `method::String="shuffle"`: will return a randomized community matrix using one of the following methods
    - `"shuffle"`: shuffles all of the entries except for the ones on the diagonals
    - `"preserve_sign_shuffle"`: same as "shuffle" but the signs are not modified
- `seed::Union{Nothing, <:Int}=nothing`: if specified, this seed will be used in the random number generator, allowing reproducibility

# Output
- `community_matrix::Array{<:Real,2}`

See also: [`random_communitymatrix`](@ref)

"""
function randomize_communitymatrix(A::Array{<:Real,2}; method::String="shuffle", seed::Union{Nothing, <:Signed}=nothing)::Array{<:Real,2}
    if seed != nothing; seed!(seed); end # set seed
    Atemp = deepcopy(A)
    if method == "shuffle"
        # list tuples (i,j) excepting  if i == j
        tuples_orig = [(i,j) for i = 1:size(A, 1) for j = setdiff(1:size(A, 2), i)]
        # shuffel tuples
        tuples_shuff = shuffle(tuples_orig)
        # randomize matrix using shuffled tuples
        Atemp[CartesianIndex.(tuples_orig)] = Atemp[CartesianIndex.(tuples_shuff)]
        return Atemp
    elseif method == "sample"
        # list tuples (i,j) excepting  if i == j
        tuples_orig = [(i,j) for i = 1:size(A, 1) for j = setdiff(1:size(A, 2), i)]
        # shuffel tuples
        tuples_shuff = sample(tuples_orig, length(tuples_orig))
        # randomize matrix using shuffled tuples
        Atemp[CartesianIndex.(tuples_orig)] = Atemp[CartesianIndex.(tuples_shuff)]
        return Atemp
    elseif method == "preserve_sign_shuffle"
        diagonal = shuffle(abs.(diag(Atemp))) # shuffle diags
        Atemp2 = deepcopy(Atemp)
        Atemp2[diagind(Atemp2)] .= 0
        non_diag_non_zero = shuffle(abs.(Atemp2[ Atemp2 .!= 0])) # strength of non-diagonal and non-zero entries
        d_count = 1;
        n_count = 1;
        Ar = zeros(size(Atemp))
        for i = 1: size(Atemp, 1), j = 1: size(Atemp, 2)
            if i == j
                Ar[i,j] = diagonal[d_count]*sign(Atemp[i,j])
                d_count+=1;
            else
                Ar[i,j] = non_diag_non_zero[n_count]*sign(Atemp[i,j])
                n_count+=1*(Atemp[i,j]!=0)
            end
        end

        return Ar
    elseif method == "preserve_sign_sample"
        diagonal = abs.(diag(Atemp)) # stength diagonal entries
        Atemp2 = deepcopy(Atemp)
        Atemp2[diagind(Atemp2)] .= 0
        non_diag_non_zero = abs.(Atemp2[ Atemp2 .!= 0]) # strength of non-diagonal and non-zero entries

        Ar = zeros(size(Atemp))
        for i = 1: size(Atemp, 1), j = 1: size(Atemp, 2)
            if i == j
                Ar[i,j] = sample(diagonal)*sign(Atemp[i,j])
            else
                Ar[i,j] = sample(non_diag_non_zero)*sign(Atemp[i,j])
            end
        end

        return Ar
    end
    @warn "Method Not Recognized: returning original growth community matrix"
    return A
end


using BenchmarkTools
using FinancialMonteCarlo

S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01;
D = 90.0;

Nsim = 10000;
Nstep = 30;
sigma = 0.2;
mc = MonteCarloConfiguration(Nsim, Nstep);
toll = 1e-3;

rfCurve = ZeroRate(r);

FwdData = Forward(T)
EUData = EuropeanOption(T, K)
AMData = AmericanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model = BlackScholesProcess(sigma, Underlying(S0, d));

@btime FwdPrice = pricer(Model, rfCurve, mc, FwdData);
@btime EuPrice = pricer(Model, rfCurve, mc, EUData);
@btime AmPrice = pricer(Model, rfCurve, mc, AMData);
@btime BarrierPrice = pricer(Model, rfCurve, mc, BarrierData);
@btime AsianPrice1 = pricer(Model, rfCurve, mc, AsianFloatingStrikeData);
@btime AsianPrice1 = pricer(Model, rfCurve, mc, AsianFixedStrikeData);

optionDatas = [FwdData, EUData, AMData, BarrierData, AsianFloatingStrikeData, AsianFixedStrikeData]

@btime (FwdPrice, EuPrice, AMPrice, BarrierPrice, AsianPrice1, AsianPrice2) = pricer(Model, rfCurve, mc, optionDatas)


function Kprec( pw::PWGrid, psi::Array{ComplexF64,2}, ik::Int )

    Ngwx  = size(psi)[1]
    Ncols = size(psi)[2]

    Ngwk = pw.gkvec.Ngw[ik]
    idx_gkk = pw.gkvec.idx_gk[ik]
    G = pw.gvec.G[:,idx_gkk]

    Kpsi  = zeros( ComplexF64, Ngwk, Ncols )

    for ic = 1:Ncols
        for ig = 1:Ngwk
            Gk = G[:,ig] + pw.gkvec.kpts[:,ik]
            Kpsi[ig,ic] = psi[ig,ic] / ( 1.0 + Gk[1]^2 + Gk[2]^2 + Gk[3]^2 )
        end
    end
    return Kpsi
end