export inf, tropicalpolynomial

import Base: +, *, ==, ^, -

"""
    Tropical(val [, isinf=false])

Create a number in the Tropical semi ring where ``a ⨁ b = max(a,b)`` and ``a ⨀ b = a + b``.
"""
Tropical(val::T, isinf=false) where {T<:Real} = Tropical{T}(val, isinf)
inf(val::T) where {T<:Real} = Tropical(val, true)
inf(::Type{T}) where {T<:Real} = Tropical(zero(T), true)
inf(t::Tropical{T}) where {T<:Real} = Tropical(t.val, true)

Base.isinf(t::Tropical) = t.isinf

function (+)(a::Tropical{T}, b::Tropical{T}) where {T<:Real}
    if isinf(a)
        b
    elseif isinf(b)
        a
    else
        Tropical(max(a.val, b.val))
    end
end

# ⨁(a::Tropical{T}, b::Tropical{T}) where {T<:Real} = a + b
# ⨁(a::Tropical, b::Tropical) = +(promote(a,b)...)
# ⨁(a::T, b::T) where {T<:Real} = Tropical(max(a, b))
# ⨁(a::Real, b::Real) = ⨁(promote(a, b)...)

# this doesn't make sense, but we define it nonetheless to make MultivariatePolynomials
# happy
(-)(a::Tropical) = a

function (*)(a::Tropical{T}, b::Tropical{T}) where {T<:Real}
    if isinf(a)
        a
    elseif isinf(b)
        b
    else
        Tropical(a.val + b.val)
    end
end

# ⨀(a::Tropical{T}, b::Tropical{T}) where {T<:Real} = a * b
# ⨀(a::Tropical, b::Tropical) = *(promote(a,b)...)
# ⨀(a::T, b::T) where {T<:Real} = Tropical(a + b)
# ⨀(a::Real, b::Real) = ⨀(promote(a, b)...)

(^)(a::Tropical, p::Integer) = isinf(a) ? a : Tropical(p * a.val)

function Base.isequal(a::Tropical{T}, b::Tropical{T}) where T
    if isinf(a) || isinf(b)
        isinf(a) == isinf(b)
    else
        a.val == b.val
    end
end
(==)(a::Tropical{T}, b::Tropical{T}) where T = isequal(a,b)

Base.one(::Tropical{T}) where T = Tropical(zero(T))
Base.one(::Type{Tropical{T}}) where T = Tropical(zero(T))

Base.zero(::Tropical{T}) where {T} = Tropical(one(T), true)
Base.zero(::Type{Tropical{T}}) where {T} = Tropical(one(T), true)

Base.show(io::IO, t::Tropical{T}) where {T} = isinf(t) ? print(io, "-∞") : print(io, t.val)

function Base.promote_rule(::Type{Tropical{T}}, ::Type{Tropical{S}}) where {S,T}
    Tropical{promote_type(S,T)}
end
function Base.promote_rule(::Type{T}, ::Type{Tropical{S}}) where {S<:Real,T<:Real}
    Tropical{promote_type(S,T)}
end

Base.convert(::Type{Tropical{T}}, t::Tropical{T}) where {T<:Real} = t
function Base.convert(::Type{Tropical{T}}, t::Tropical{<:Real}) where {T<:Real}
    Tropical{T}(convert(T, t.val), t.isinf)
end
function Base.convert(::Type{Tropical{T}}, x::Real) where {T<:Real}
    if _isone(x)
        one(Tropical{T})
    elseif iszero(x)
        zero(Tropical{T})
    else
        Tropical(convert(T, x))
    end
end

# function Base.convert(::Type{T}, t::Tropical{<:Real}) where {T<:Number}
#     if _isone(t)
#         one(T)
#     elseif iszero(t)
#         zero(T)
#     else
#         convert(T, t.val)
#     end
# end

"""
    tropicalpolynomial(p::MP.AbstractPolynomial)

Interprets a normal polynomials a tropical one. This will especially convert all coefficents which are
1 (neutral element w.r.t. multiplication) to 0 (neutral element wr.t. tropical multiplication).
"""
function tropicalpolynomial(p::MP.AbstractPolynomial)
    mapcoefficients(c -> convert(Tropical{Float64}, c), p)
end

Base.float(t::Tropical) = isinf(t) ? -Inf : float(t.val)


# using Revise, Plots
using Test, ForwardDiff, Parameters, Setfield, LinearAlgebra
using BifurcationKit
const BK = BifurcationKit

norminf = x -> norm(x, Inf)
####################################################################################################
function COm(u, p)
	@unpack q1,q2,q3,q4,q5,q6,k = p
	x, y, s = u
	z = 1-x-y-s
	out = similar(u)
	out[1] = 2 * q1 * z^2 - 2 * q5 * x^2 - q3 * x * y
	out[2] = q2 * z - q6 * y - q3 * x * y
	out[3] = q4 * z - k * q4 * s
	out
end
dCOm = (z, p) -> ForwardDiff.jacobian(x -> COm(x, p), z)
jet = BK.get3Jet(COm, dCOm)

par_com = (q1 = 2.5, q2 = 2.0, q3 = 10., q4 = 0.0675, q5 = 1., q6 = 0.1, k = 0.4)
z0 = [0.07,0.2,05]

opts_br = ContinuationPar(pMin = 0.6, pMax = 2.5, ds = 0.002, dsmax = 0.01, nInversion = 4, detectBifurcation = 3, maxBisectionSteps = 25, nev = 2, maxSteps = 20000)
	br, = @time continuation(jet[1], jet[2], z0, par_com, (@lens _.q2), opts_br;
	printSolution = (x, p) -> (x = x[1], y = x[2]),
	plot = false, verbosity = 0, normC = norminf,
	bothside = true)

# plot(br, plotfold=false, markersize=4, legend=:topleft)
####################################################################################################
hp, = newton(jet[1:2]..., br, 1; options = NewtonPar( opts_br.newtonOptions; maxIter = 10),startWithEigen=true, d2F = jet[3])

hpnf = computeNormalForm(jet..., br, 4)
@test hpnf.nf.b |> real ≈ 1.070259e+01 rtol = 1e-3

hpnf = computeNormalForm(jet..., br, 1)
@test hpnf.nf.b |> real ≈ 4.332247e+00 rtol = 1e-2
####################################################################################################
# different tests for the Fold point
snpt = computeNormalForm(jet..., br, 2)
@set! opts_br.newtonOptions.verbose = false
@set! opts_br.newtonOptions.maxIter = 10

sn = newton(jet[1:2]..., br, 2; options = opts_br.newtonOptions, bdlinsolver = MatrixBLS())
# printstyled(color=:red, "--> guess for SN, p = ", br.specialpoint[2].param, ", psn = ", sn[1].p)
	# plot(br);scatter!([sn.x.p], [sn.x.u[1]])
@test sn[3] && sn[5] == 6

sn = newton(jet[1:2]..., br, 2; options = opts_br.newtonOptions, bdlinsolver = MatrixBLS(), d2F = jet[3])
@test sn[3] && sn[5] == 8

sn = newton(jet[1:2]..., br, 2; options = opts_br.newtonOptions, bdlinsolver = MatrixBLS(), d2F = jet[3], startWithEigen = true)
@test sn[3] && sn[5] == 8

sn_br, = continuation(jet[1:2]..., br, 2, (@lens _.k), ContinuationPar(opts_br, pMax = 1., pMin = 0., detectBifurcation = 1, maxSteps = 50, saveSolEveryStep = 1, detectEvent = 2), bdlinsolver = MatrixBLS(), d2F = jet[3], startWithEigen = true, updateMinAugEveryStep = 1, verbosity = 0, plot=false)
@test sn_br.specialpoint[1].type == :bt
@test ~isnothing(sn_br.eig)

# we test the jacobian and problem update
par_sn = set(br.params, br.lens, sn_br.sol[end].x.p)
par_sn = set(par_sn, sn_br.lens, sn_br.sol[end].p)
_J = dCOm(sn_br.sol[end].x.u, par_sn)
_eigvals, eigvec, = eigen(_J)
ind = argmin(abs.(_eigvals))
@test _eigvals[ind] ≈ 0 atol=1e-10
ζ = eigvec[:, ind]
@test sn_br.functional.b ./ norm(sn_br.functional.b) ≈ -ζ

_eigvals, eigvec, = eigen(_J')
ind = argmin(abs.(_eigvals))
ζstar = eigvec[:, ind]
@test sn_br.functional.a ≈ ζstar
####################################################################################################
# different tests for the Hopf point
hppt = computeNormalForm(jet..., br, 1)

@set! opts_br.newtonOptions.verbose = false

hp = newtonHopf(jet[1:2]..., br, 1; options = opts_br.newtonOptions, startWithEigen = true)
# printstyled(color=:red, "--> guess for HP, p = ", br.specialpoint[1].param, ", php = ", hp[1].p)
# plot(br);scatter!([hp[1].p[1]], [hp[1].u[1]])
@test hp[3] && hp[5] == 12

hp = newtonHopf(jet[1:2]..., br, 1; options = opts_br.newtonOptions, startWithEigen = false, bdlinsolver = MatrixBLS())
@test hp[3] && hp[5] == 12

hp = newtonHopf(jet[1:2]..., br, 1; options = opts_br.newtonOptions, startWithEigen = true, bdlinsolver = MatrixBLS(), verbose = true)
@test hp[3] && hp[5] == 8

# we check that we truly have a bifurcation point.
pb = hp[end]
ω = hp[1].p[2]
par_hp = set(br.params, br.lens, hp[1].p[1])
_J = dCOm(hp[1].u, par_hp)
_eigvals, eigvec, = eigen(_J)
ind = argmin(abs.(_eigvals .- Complex(0, ω)))
@test real(_eigvals[ind]) ≈ 0 atol=1e-9
@test abs(imag(_eigvals[ind])) ≈ abs(hp[1].p[2]) atol=1e-9
ζ = eigvec[:, ind]
# reminder: pb.b should be a null vector of (J+iω)
@test pb.b ≈ ζ atol = 1e-3

hp, = newton(jet[1:2]..., br, 1;
	options = NewtonPar( opts_br.newtonOptions; maxIter = 10),
	startWithEigen = true,
	bdlinsolver = MatrixBLS(),
	d2F = jet[3])
# printstyled(color=:red, "--> guess for HP, p = ", br.specialpoint[1].param, ", php = ", hp.p)
# plot(br);scatter!([hp.p[1]], [hp.u[1]])

hp, = newton(jet[1:2]..., br, 1; options = NewtonPar( opts_br.newtonOptions; maxIter = 10),startWithEigen=true, d2F = jet[3])

hp_br, = continuation(jet[1:2]..., br, 1, (@lens _.k), ContinuationPar(opts_br, ds = -0.001, pMax = 1., pMin = 0., detectBifurcation = 1, maxSteps = 50, saveSolEveryStep = 1, detectEvent = 2), bdlinsolver = MatrixBLS(), d2F = jet[3], d3F = jet[4], startWithEigen = true, updateMinAugEveryStep = 1, verbosity=0, plot=false)
@test hp_br.specialpoint[1].type == :gh
@test ~isnothing(hp_br.eig)


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ a4eba47e-3cde-11eb-2a1c-7f5102c9a9e1
begin
	using HDF5
	using Plots
	using PlutoUI
	using LaTeXStrings
	using Printf
	import Pda
end

# ╔═╡ b0c18c94-3cde-11eb-0c12-731d453f428b
md"""## Exponential RTP in free space"""

# ╔═╡ f2663b88-3cdf-11eb-2869-5f3d5ff56948
md"""The MSD of exponential RTPs with complete reorientation and initial condition $$P(\mathbf{x}, \theta,t) = \delta(\mathbf{x})/2\pi$$ is

$$\left<x^2\right> =\left<y^2\right>=  U^2_0\tau^2_R\left( \frac{t}{\tau_R} + \exp(-t/\tau_R) -1\right)$$

Ref: [Run-and-tumble bacteria slowly approaching the diffusive regime](https://doi.org/10.1103/PhysRevE.101.062607), Villa-Torrealba *et. al*, PRE, 2020
"""

# ╔═╡ 3635fca6-3ce0-11eb-0d4b-1310b1716b86
function msd_exponential(t::Real, U0::Real, tauR::Real)
	return U0^2*tauR^2*(t/tauR+exp(-t/tauR) -1)
end

# ╔═╡ 8b42ee14-3ce2-11eb-0702-5da8b9887bd9
function get_var(file::String)
	t = Pda.get_h5data(file, "x/t")[2:end]
	v = Pda.get_h5data(file, "x/v")[2:end]
	tauR = Pda.get_h5attr(file, "tauR")
	U0 = Pda.get_h5attr(file, "U0")
	ell = U0*tauR
	# non-dimensionalized by tauR and ell
	t ./= tauR
	v ./= ell^2
	return t,v
end

# ╔═╡ d0f3bf32-3cde-11eb-13c0-6d1d3bd0b070
begin
	file = "U1.000tauR1.200.h5"
	t,v = get_var(file)
	plot(t,v,
		framestyle=:box,
		label="Langevin",
		lw=2)
	plot!(t, msd_exponential.(t, 1, 1)/1^2,
		lw=2,
		linestyle=:dash,
		label="theory",
		legend=:topleft,
		xscale=:log10,
		yscale=:log10)
	xlabel!(L"t/\tau_R")
	ylabel!(L"\mathrm{var}(x)/\ell^2")
	title!("Exponential RTP in free space")
end

# ╔═╡ 48f91100-3ce2-11eb-1331-930d15c04e16
md"""## Pareto RTP in free space"""

# ╔═╡ 53054a04-3ce2-11eb-2869-973ecb35edeb
begin
	file1= "U1.000alpha1.200free.h5"
	t1,v1 = get_var(file1)
	plot(t1,v1,
		framestyle=:box,
		label="Langevin",
		lw=2)
	alpha = Pda.get_h5attr(file1, "alpha")
	# #theory is <x^2> \sim t^(3-\alpha)
	i = 3000
	plot!(t1[end-i:end], (t1[end-i:end]/t1[end]).^(3-alpha)*v1[end],
		lw=2,
		linestyle=:dash,
		label="theory",
		legend=:topleft,
		xscale=:log10,
		yscale=:log10)
	xlabel!(L"t/\tau_R")
	ylabel!(L"\mathrm{var}(x)/\ell^2")
	title!("Pareto RTP in free space")
end

# ╔═╡ Cell order:
# ╠═a4eba47e-3cde-11eb-2a1c-7f5102c9a9e1
# ╟─b0c18c94-3cde-11eb-0c12-731d453f428b
# ╟─f2663b88-3cdf-11eb-2869-5f3d5ff56948
# ╠═3635fca6-3ce0-11eb-0d4b-1310b1716b86
# ╠═8b42ee14-3ce2-11eb-0702-5da8b9887bd9
# ╠═d0f3bf32-3cde-11eb-13c0-6d1d3bd0b070
# ╟─48f91100-3ce2-11eb-1331-930d15c04e16
# ╠═53054a04-3ce2-11eb-2869-973ecb35edeb


# ## Пример №1
#
# Первый и последний пример, показывающий как использовать функцию.
#
# Работа алгоритма при разных параметрах ``\varepsilon`` и доли заполнения
# кадра (доли точек, которые прошел переходный слой) будет рассмотрена более
# подробно в секции численных экспериментов.

using NonLinearReactionAdvectionDiffusionWithFrontData
using Dierckx
using LinearAlgebra;
using Plots; gr();

# Входные параметры определим полностью, чтобы Вам было проще их быстро
# изменить, и быстро посмотреть на результат.
u_l(t) = -8;
u_r(t) =  4;
qf(x) = 4*sin(3 * π * x);       # Коэффициент линейного усиления,
ε = 0.03;                       # Малый параметр при старшей производной
a, b = 0, 1;                    # Область по X
t₀, T = 0, 0.42;                # Область по T
N, M = 200, 300;                # Кол-во разбиений по X, T
h = (b-a)/N;                    # шаг по X
τ = (T-t₀)/M;                   # шаг по T
Xₙ = [a  + n*h for n in 0:N];   # Сетка по Х
Tₘ = [t₀ + m*τ for m in 0:M];   # Сетка по Т
qₙ =     qf.(Xₙ);               # Сеточные значения коэффициента лин. усиления
ulₘ=    u_l.(Tₘ);               # Сеточные значения левого  ГУ
urₘ=    u_r.(Tₘ);               # Сеточные значения правого ГУ
u₀ = u_init.(Xₙ, ε=ε, x_tp = 0.1);               # Начальные условия
nothing #hide
#
u, XX, TP = solve(u₀, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ);
ϕl, ϕr, ϕ, f1_data, f2_data = generate_obs_data(u, Xₙ, N, Tₘ, M, qₙ, ulₘ, urₘ);
nothing # hide

q_guess = initial_guess(f1_data, Xₙ, N, Tₘ, M, ulₘ, urₘ, 0.005);

#md # !!! note
#md #   Работает, только с постоянными ГУ!


@testset "air_density" begin
  ad = @inferred air_density(25.0,100.0) # Tair, pressure
  # regression test
  @test ≈(ad, 1.168, atol =0.001)
end

@testset "pressure_from_elevation" begin
  pressure = @inferred pressure_from_elevation(500.0, 25.0) # elev, Tair
  # regression test
  @test ≈(pressure, 95.681, atol =0.001)
  #
  pressure = @inferred pressure_from_elevation(500.0, 25.0, 1.5) # elev, Tair, VPD
  # regression test
  @test ≈(pressure, 95.717, atol =0.001)
end

@testset "virtual_temp" begin
  Tair,pressure,VPD = 25.0,100.0,1.5
  vt = @inferred virtual_temp(Tair,pressure,VPD)
  # regression test
  @test ≈(vt, 26.9, atol =0.1)
end

@testset "kinematic_viscosity" begin
  Tair,pressure = 25.0,100.0
  vis = @inferred kinematic_viscosity(Tair,pressure)
  # regression test
  @test ≈(vis, 1.58e-5, atol =1e-7)
end

@testset "air_density" begin
  ad = @inferred air_density(25.0,100.0) # Tair, pressure
  # regression test
  @test ≈(ad, 1.168, atol =0.001)
end

@testset "dew_point" begin
  VPD = 1.5
  Tair = 25.0
  accuracy = 1e-2
  #@descend_code_warntype VPD_to_e(VPD,Tair)
  ea = @inferred VPD_to_e(VPD,Tair)
  Td = @inferred dew_point_from_e(ea,accuracy = accuracy)
  @test ≈(ea, Esat_from_Tair(Td), atol = accuracy)
  @test ≈(dew_point(Tair, VPD; accuracy = accuracy), Td, atol = accuracy)
  Td1 = @test_logs (:warn,r"set to 1")  @inferred dew_point(Tair, VPD; accuracy = 1.2)
  @test ≈(Td1, Td, atol = 1.0)
end

@testset "wetbulb temperature" begin
  Tair = 25.0
  VPD = 1.0
  pressure = 100.0
  accuracy = 1e-2
  #
  gamma  = @inferred psychrometric_constant(Tair,pressure)
  ea     = VPD_to_e(VPD,Tair)
  Tw = @inferred wetbulb_temp_from_e_Tair_gamma(ea, Tair, gamma; accuracy = accuracy)
  #@test ≈(ea, Esat_from_Tair(Tw) - gamma* (Tair - Tw), atol = accuracy)
  @test ≈(ea, Esat_from_Tair(Tw) - 0.93*gamma* (Tair - Tw), atol = accuracy)
  @test ≈(wetbulb_temp(Tair, pressure, VPD; accuracy = accuracy), Tw, atol = accuracy)
  Tw1 = @test_logs (:warn,r"set to 1") @inferred wetbulb_temp(Tair, pressure, VPD; accuracy = 1.2)
  @test ≈(Tw1, Tw, atol = 1.0)
end


module SpawnerRecruitModels

export
    AbstractSpawnerRecruitModel,
    BevertonHoltSpawnerRecruitModel,
    RickerSpawnerRecruitModel,
    LudwigWaltersSpawnerRecruitModel,
    CushingSpawnerRecruitModel,
    DerisoSchnuteSpawnerRecruitModel,
    ShepherdSpawnerRecruitModel,
    GammaSpawnerRecruitModel,
    recruit,
    max_recruits,
    max_spawnrecruits

abstract type AbstractSpawnerRecruitModel end

"""
    BevertonHoltSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel

Parameters:

- α: density-independent productivity parameter
- β: density-dependence parameter

Follows equation 3.6 (p. 88) of Quinn and Deriso. Asymptotic recruitment with
increasing spawner biomass.
"""
struct BevertonHoltSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T

    function BevertonHoltSpawnerRecruitModel(α::T, β::T) where T
        new{T}(α, β)
    end
end

function recruit(spawners, srmod::BevertonHoltSpawnerRecruitModel)
    srmod.α * spawners / (1 + srmod.β * spawners)
end

function max_recruits(srmod::BevertonHoltSpawnerRecruitModel)
    srmod.α / srmod.β
end

function max_spawnrecruits(srmod::BevertonHoltSpawnerRecruitModel)
    (Inf, max_recruits(srmod))
end

"""
    RickerSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- β: density-dependence parameter

Parameterized as in equation 3.8 of Quinn and Deriso. Dome-shaped recruitment
curve.
"""
struct RickerSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels
    α::T
    β::T

    function RickerSpawnerRecruitModel(α::T, β::T) where T
        new{T}(α, β)
    end
end

function recruit(spawners, srmod::RickerSpawnerRecruitModel)
    srmod.α * spawners * exp(-srmod.β * spawners)
end

function max_recruits(srmod::RickerSpawnerRecruitModel)
    srmod.α / (srmod.β * ℯ)
end

function max_spawnrecruits(srmod::RickerSpawnerRecruitModel)
    (1 / srmod.β, max_recruits(srmod))
end

"""
    LudwigWaltersSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- β: density-dependence parameter
- γ: density-dependence power parameter

Generalization of Ricker spawner recruit model. Parameterized as in equation 3.9
of Quinn and Deriso. Dome-shaped recruitment curve.
"""
struct LudwigWaltersSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T
    γ::T

    function LudwigWaltersSpawnerRecruitModel(α::T, β::T, γ::T) where T
        new{T}(α, β, γ)
    end
end

function recruit(spawners, srmod::LudwigWaltersSpawnerRecruitModel)
    srmod.α * spawners * exp(-srmod.β * spawners ^ srmod.γ)
end

function max_recruits(srmod::LudwigWaltersSpawnerRecruitModel)
    srmod.α * (srmod.β * srmod.γ) ^ -(1 / srmod.γ)
end

function max_spawnrecruits(srmod::LudwigWaltersSpawnerRecruitModel)
    (exp(-1 / srmod.γ), max_recruits(srmod))
end

"""
    CushingSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- γ: index of density dependence

Generalization of Ricker spawner recruit model. Parameterized as in equation 3.12
of Quinn and Deriso. Dome-shaped recruitment curve.
"""
struct CushingSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    γ::T

    function CushingSpawnerRecruitModel(α::T, γ::T) where T
        new{T}(α, γ)
    end
end

function recruit(spawners, srmod::CushingSpawnerRecruitModel)
    srmod.α * spawners ^ srmod.γ
end

max_recruits(srmod::CushingSpawnerRecruitModel) = Inf

max_spawnrecruits(srmod::CushingSpawnerRecruitModel) = (Inf, Inf)

"""
    DerisoSchnuteSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- β: optimality parameter
- γ: recruitment limitation or skewness parameter

Generalization of Beverton-Holt and Ricker spawner recruit models. Parameterized
as in equation 3.20 of Quinn and Deriso. May be asymptotic or dome-shaped
recruitment curve, depending on the value of γ.
"""
struct DerisoSchnuteSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T
    γ::T

    function DerisoSchnuteSpawnerRecruitModel(α::T, β::T, γ::T) where T
        new{T}(α, β, γ)
    end
end

function recruit(spawners, srmod::DerisoSchnuteSpawnerRecruitModel)
    srmod.α * spawners * (1 - srmod.β * srmod.γ * spawners) ^ (1 / srmod.γ)
end

function max_recruits(srmod::DerisoSchnuteSpawnerRecruitModel)
    (srmod.α / srmod.β) * (1 + srmod.γ) ^ -((1 + srmod.γ) / srmod.γ)
end

function max_spawnrecruits(srmod::DerisoSchnuteSpawnerRecruitModels)
    (1 / (srmod.β * (1 + srmod.γ)), max_recruits(srmod))
end

"""
    ShepherdSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModels

Parameters

- α: productivity parameter
- γ: index of density dependence

Generalization of Beverton-Holt, Ricker, and Cushing spawner recruit models.
Parameterized as in equation 3.21 of Quinn and Deriso. Unlimited recruitment for
γ < 1 as in the Cushing model, asymptotic recruitment (Beverton-Holt)
recruitment at γ = 1, and dome-shaped recruitment at γ > 1.
"""
struct ShepherdSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T
    γ::T

    function ShepherdSpawnerRecruitModel(α::T, β::T, γ::T) where T
        new{T}(α, β, γ)
    end
end

function recruit(spawners, srmod::ShepherdSpawnerRecruitModel)
    srmod.α * spawners / (1 + β * spawners ^ srmod.γ)
end

function max_recruits(srmod::ShepherdSpawnerRecruitModel)
    if srmod.γ > 1
        γm1 = srmod.γ - 1
        max_rec = srmod.α / (srmod.β * srmod.γ) * (srmod.β * γm1) ^ (γm1 / srmod.γ)
    elseif srmod.γ == 1
        max_rec = max_recruits(BevertonHoltSpawnerRecruitModel(srmod.α, srmod.β))
    else
        max_rec = Inf
    end
    max_rec
end

function max_spawnrecruits(srmod::ShepherdSpawnerRecruitModel)
    if srmod.γ > 1
        spawnmax = (srmod.β * (srmod.γ - 1)) ^ (-1 / srmod.γ)
    elseif srmod.γ ≤ 1
        spawnmax = Inf
    end
    (spawnmax, max_recruits(srmod))
end

"""
    GammaSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel

Parameters:

- α: productivity parameter
- β:
- γ:

Takes the form of an unnomalized gamma function. Contains the Ricker (γ = 1) and
Cushing (β = 0) as special cases. Approaches Beverton-Holt as γ → 0
"""
struct GammaSpawnerRecruitModel{T} <: AbstractSpawnerRecruitModel
    α::T
    β::T
    γ::T

    function GammaSpawnerRecruitModel(α::T, β::T, γ::T) where T
        γ > 0 || DomainError("γ must be positive for a valid spawner-recruit model")
        new{T}(α, β, γ)
    end
end

function recruit(spawners, srmod::GammaSpawnerRecruitModel)
    srmod.α * spawners ^ srmod.γ * exp(-srmod.β * spawners)
end

function max_recruits(srmod::GammaSpawnerRecruitModel)
    srmod.α * (srmod.γ / srmod.β) ^ srmod.γ * exp(-srmod.γ)
end

function max_spawnrecruits(srmod::GammaSpawnerRecruitModel)
    (srmod.γ / srmod.β, max_recruits(srmod))
end

end


julia> [i + j for i in 1:5, j in 1:5]
5×5 Array{Int64,2}:
 2  3  4  5   6
 3  4  5  6   7
 4  5  6  7   8
 5  6  7  8   9
 6  7  8  9  10

julia> [i + j for i in 1:5, j in 1:5, k in 1:2]
5×5×2 Array{Int64,3}:
[:, :, 1] =
 2  3  4  5   6
 3  4  5  6   7
 4  5  6  7   8
 5  6  7  8   9
 6  7  8  9  10

[:, :, 2] =
 2  3  4  5   6
 3  4  5  6   7
 4  5  6  7   8
 5  6  7  8   9
 6  7  8  9  10


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Plane(p, u, v)

A plane embedded in R³ passing through point `p`,
defined by non-parallel vectors `u` and `v`.

    Plane(p, n)

Alternatively specify point `p` and a given normal
vector `n` to the plane.
"""
struct Plane{T} <: Primitive{3,T}
  p::Point{3,T}
  u::Vec{3,T}
  v::Vec{3,T}
end

function Plane{T}(p::Point{3,T}, n::Vec{3,T}) where {T}
  # origin of coordinate system
  o = Vec{3,T}(zero(T), zero(T), zero(T))

  uv = Vec{3,T}[]
  for i in 1:3
    # subtract projection of Euclidean basis onto normal
    e = Vec{3,T}(ntuple(j -> j == i ? one(T) : zero(T), 3))
    v = e - (e⋅n)/(n⋅n) * n

    # check if it is a valid vector
    isapprox(v, o, atol=atol(T)) || push!(uv, v)

    # we need two vectors
    length(uv) == 2 && break
  end

  Plane{T}(p, uv[1], uv[2])
end

Plane(p::Point{3,T}, n::Vec{3,T}) where {T} = Plane{T}(p, n)

Plane(p::Tuple, u::Tuple, v::Tuple) = Plane(Point(p), Vec(u), Vec(v))

Plane(p::Tuple, n::Tuple) = Plane(Point(p), Vec(n))

paramdim(::Type{<:Plane}) = 2

isconvex(::Type{<:Plane}) = true

boundary(::Plane) = nothing

(p::Plane)(s, t) = p.p + s*p.u + t*p.v

"""
    origin(plane)

Return the origin of the `plane`.
"""
origin(p::Plane) = p.p

"""
    normal(plane)

Normal vector to the `plane`.
"""
normal(p::Plane) = p.u × p.v


using QuDirac

####################
# Ladder Operators #
####################
#
# `a` and `a_dag` are defined as the lowering and raising operators
# on individual states. `X` and `P` are defined as the position
# and momentum operators on individual states.
#
# See the "Constructing Operators with Functions" section
# of the QuDirac docs for more info on the use of the `@def_op`
# macro.
#
# Note that in practice, it is generally faster to use QuDirac's
# built-in `lower` and `raise` functions to apply ladder operations
# to states. The following simply serves to demonstrate the capabilities
# of the `@def_op` macro.

@def_op " a | n > = √n * | n - 1 >"
@def_op " < n | a = √(n + 1) * < n + 1 |"

a_dag = a'

# `X` and `P` are Hermitian, so:
@def_op " X | n > = √(1/2) * (a * | n > +  a_dag * | n >)"
@def_op " < n | X = √(1/2) * (< n | * a +  < n | * a_dag)"

@def_op " P | n > = im * √(1/2) * (a_dag * | n > - a * | n >)"
@def_op " < n | P = im * √(1/2) * (< n | * a_dag - < n | * a)"

#################################
# Hermite Polynomial Evaluation #
#################################
# This is a very naive implementation,
# likely prone to numerical error and slowness
# for high n. It should suffice for this example,
# however. Implements the expression found here:
# http://en.wikipedia.org/wiki/Hermite_polynomials#Recursion_relation
function hermite(n, x)
    if n == 0
        return 1
    elseif n == 1
        return 2 * x
    else
        return 2x*hermite(n-1, x) - 2*(n-1)*hermite(n-2, x)
    end
end

##################
# QHOInner Setup #
##################
# Custom inner product type
# the QHO problem
struct QHOInner <: AbstractInner end

qho_inner(x::Float64, k::BigInt) = e^((-x^2)/2) * hermite(k, x) * 1/√(2^k * factorial(k) * √π) # using natural units
qho_inner(x::Float64, k::Int) = qho_inner(x, BigInt(k))
qho_inner(n::Int, m::Int) = n==m ? 1.0 : 0.0

function QuDirac.inner_rule(::QHOInner, b, k)
    result = 1.0
    for (b0, k0) in zip(b, k)
        result *= qho_inner(b0, k0)
    end
    return result
end

QuDirac.inner_rettype(::QHOInner) = BigFloat

default_inner!(QHOInner())

# With the above, we've defined this behavior for basis states:
#
# < i::Int | * | j::Int > = ∫ ψᵢ'ψⱼ dx =  δᵢⱼ
#
# < x::Float64 | * | i::Int > = ψᵢ(x)

####################
# Make Some Plots! #
####################
# Here we're going to make some plots
# with Plotly. You'll have to have an
# account to do this, but it's free and
# totally awesome.

# Given an iterable of x and y points, generate a
# distribution for the state by taking the inner product
# Julia would correctly type-inferenced this as a BigFloat
# array, but we're going to convert it to Float64 to send it
# to Plotly
gen_z(kt::Ket, xpoints, ypoints) = Float64[d" < x, y | * kt " for x in xpoints, y in ypoints]

# Generate the distribution above, and package it for a Plotly surface plot
function gen_plot_data(kt::Ket{P,2}, xpoints, ypoints) where P
    return [
      [
        "z" => gen_z(kt, xpoints, ypoints),
        "x" => xpoints,
        "y" => ypoints,
        "type" => "surface"
      ]
    ]
end

# some default stuff
len = 50
max = pi
xpoints = range(-max, max, length=len)
ypoints = copy(xpoints)

# info("Loading the Plotly package, this could take a little while if it has to sign in...")

# using Plotly

# # Generate the distribution, sending it to Plotly.
# # Return the response URL, which you can then go to
# # to see and interact with your plot.
# #
# # To generate a plot of a wave function for a 2-factor Ket,
# # just call plot_wave2D(kt, xpoints, ypoints). This function
# # will build your plot and return the URL you should go to to
# # view the result. Here are some examples:
# #
# # Basis state:
# # julia> plot_wave2D(d\" | 1, 1 > \", xpoints, ypoints)
# #
# # Random superposition of the first 4 basis states:
# # julia> randkt = normalize!(sum(i -> rand() * ket(i), 0:3))^2
# #        plot_wave2D(randkt, xpoints, ypoints)
# #
# function plot_wave2D{P}(kt::Ket{P,2}, xpoints, ypoints)
#     response = Plotly.plot(gen_plot_data(kt, xpoints, ypoints))
#     return response["url"]
# end

# info("Finished loading Plotly. Make sure you're signed in before trying to plot!")


using Base: has_fast_linear_indexing, first_step_last_ascending
include("mydct.jl")

using MAT
using LinearAlgebra

matrix = [231 32 233 161 24 71 140 245
247 40 248 245 124 204 36 107
234 202 245 167 9 217 239 173
193 190 100 167 43 180 8 70
11 24 210 177 81 243 8 112
97 195 203 47 125 114 165 181
193 70 174 167 41 30 127 245
87 149 57 192 65 129 178 228]

MHmatrix = [
1.11e+03    4.40e+01    7.59e+01    -1.38e+02   3.50e+00    1.22e+02   1.95e+02    -1.01e+02
7.71e+01    1.14e+02    -2.18e+01   4.13e+01    8.77e+00    9.90e+01   1.38e+02    1.09e+01
4.48e+01    -6.27e+01   1.11e+02    -7.63e+01   1.24e+02    9.55e+01   -3.98e+01   5.85e+01
-6.99e+0    -4.02e+01   -2.34e+01   -7.67e+01   2.66e+01    -3.68e+01   6.61e+01    1.25e+02
-1.09e+0    -4.33e+01   -5.55e+01   8.17e+00    3.02e+01    -2.86e+01   2.44e+00    -9.41e+01
-5.38e+0    5.66e+01    1.73e+02    -3.54e+01   3.23e+01    3.34e+01   -5.81e+01   1.90e+01
7.88e+01    -6.45e+01   1.18e+02    -1.50e+01   -1.37e+02   -3.06e+01   -1.05e+02   3.98e+01
1.97e+01    -7.81e+01   9.72e-01    -7.23e+01   -2.15e+01   8.13e+01   6.37e+01    5.90e+00
]

cmatrix = [1118.7500000000002 44.02219262301787 75.91905032609426 -138.57241099707306 3.5000000000000284 122.07805518803062 195.04386762362958 -101.60490592795372; 77.1900790187901 114.86820590697009 -21.80144211456541 41.364135057036215 8.77720597607919 99.08296204231638 138.17151573211936 10.909279534077866; 44.83515365361005 -62.75244638273558 111.61411421683839 -76.37896579814382 124.42215968070717 95.59841936381827 -39.82879694884997 58.52376699413914; -69.98366474768355 -40.240894480279415 -23.497050834570445 -76.73205936982573 26.645774950452864 -36.83282895356625 66.18914845485334 125.42973061799121; -109.00000000000003 -43.34308566057378 -55.54369078865144 8.173470828327527 30.25000000000001 -28.660243733699012 2.4414982233614957 -94.14370254649862; -5.387835905170836 56.63450089125683 173.0215190411498 -35.42344938713219 32.38782492363563 33.45767275218619 -58.1167863718725 19.022561487011163; 78.84396931190855 -64.592409552383 118.67120305115003 -15.09048397585811 -137.31692787267204 -30.619666281198647 -105.11411421683843 39.813049706882936; 19.78824382848893 -78.18134089948808 0.9723118598351412 -72.34641800960696 -21.578163250360877 81.29990354778818 63.71037820527631 5.90618071066943]

v = [231 32 233 161 24 71 140 245]
c = [401.99 6.60002 109.167 -112.786 65.4074 121.831 116.656 28.8004]

function equalsMat(mat, emat)
    if round.(mat, digits=2)==round.(emat, digits=2)
        return true
    else
        return false
    end

end

function testDCT2(mat, cmat)
    if equalsMat(mDCT2n3(mat), cmat)
        return true
    else
        return false
    end
end

function testDCT2(mat)
    if equalsMat(mDCT2n3(mat), dct(mat))
        return true
    else
        return false
    end
end

function testiDCT2(cmat, mat)
    if equalsMat(miDCT2n3(cmat), mat)
        return true
    else
        return false
    end
end

function testiDCT2(cmat)
    if equalsMat(miDCT2n3(cmat), idct(cmat))
        return true
    else
        return false
    end
end

function testDCT(v, c)
    if equalsMat(mDCT(v), c)
        if equalsMat(v, miDCT(c))
            return true
        end
    else
        return false
    end
end

## VENGONO TESTATE TUTTE LE DCT

function testALL(mat, cmat)
    if (testDCT2(mat, cmat) && (testDCT2(mat))
        && testiDCT2(cmat, mat) && testiDCT2(cmat))
        return true
    else
        return false
    end
end

function testALL(mat)
    cmat = dct(mat)
    imat = idct(mat)
    return testALL(mat, cmat) && testALL(imat, mat)
end

if testALL(v,c)
    println("All vector test PASSED")
else
    println("All vector test NOT Passed")
end
if testALL(matrix, cmatrix)
    println("All matrix test PASSED")
else
    println("All matrix test NOT Passed")
end


using ArgParse
using Plots

function x_mother(c,omega,t)
    return cos(t)+c*cos(omega*t)
end

function y_mother(c,omega,t)
    return sin(t)+c*sin(omega*t)
end

function main()

    s = ArgParseSettings()
    @add_arg_table s begin
        "--output","-o"
        help = "output graph image file"
        arg_type = String
        default = "output.png"
        "-c"
        help = "coeff c"
        arg_type = Float64
        default = 1.0
        "-w"
        help = "coeff w"
        arg_type = Float64
        default = 1.0
    end

    parsed_args = parse_args(s);
    for (arg,val) in parsed_args
        println("   $arg => $val")
    end

    c = parsed_args["c"]
    w = parsed_args["w"]
    x(t) = x_mother(c,w,t)
    y(t) = y_mother(c,w,t)

    plot(x, y, 0:0.01:1000,aspect_ratio=1.0)

    savefig(parsed_args["output"])

end

main()


using BitSAD
using Makie, DataStructures
using Statistics: mean
using LinearAlgebra

circuit = @circuit IterativeSVD begin
    parameters : [
        rows::Int => 2,
        cols::Int => 2
    ]

    circuit : (dut::IterativeSVD)(A::Matrix{SBit}, v₀::Vector{SBit}) -> begin
        # Update right singular vector
        w = A * v₀
        wscaled = w .÷ sqrt(dut.rows)
        u = wscaled ./ norm(wscaled)

        # Update left singular vector
        z = permutedims(A) * u
        zscaled = z .÷ sqrt(dut.cols)
        σ = norm(zscaled)
        v = zscaled ./ σ

        return u, v, σ
    end
end

N = 10
T = 20000
m = 2
n = 2

# generate inputs
A = [2 .* rand(m, n) .- 1 for i in 1:N]
v₀ = [rand(n) for i in 1:N]
v₀ .= v₀ ./ norm.(v₀)
dut = [IterativeSVD(m, n) for i in 1:N]

# calculate scaling
α = 2 .* max.(norm.(A, Inf), norm.(A, 1))
A = A ./ α

# convert to bitstream
A = SBitstream.(A)
v₀ = SBitstream.(v₀)

# eval loop
BitSAD.clearops()
ϵ = zeros(T, N)
ubuffer = [CircularBuffer{Vector{Int}}(5000) for i in 1:N]
vbuffer = [CircularBuffer{Vector{Int}}(5000) for i in 1:N]
σbuffer = [CircularBuffer{Int}(5000) for i in 1:N]
Threads.@threads for trial in 1:N
    generate!.(A[trial], T)
    generate!.(v₀[trial], 1000)

    for t in 1:T
        # evaluate module
        output = dut[trial](pop!.(A[trial]), pop!.(v₀[trial]))
        (t >= 1000) && push!.(v₀[trial], decorrelate.(output[2]))

        # accumulate results in buffer
        u = estimate!(ubuffer[trial], output[1])
        v = estimate!(vbuffer[trial], output[2])
        σ = estimate!(σbuffer[trial], output[3])

        # record loss
        ϵ[t, trial] = norm(α[trial] * (float.(A[trial]) * v - u * σ * sqrt(n)))
    end

    println("Completed trial $trial")
end

u = @. estimate!(ubuffer)
v = @. estimate!(vbuffer)
σ = @. estimate!(σbuffer) * sqrt(n)
A = map(λ -> float.(λ), A)
f = svd(A[1])
println("u error: $(u[1] - f.U[1, :])")
println("v error: $(v[1] - f.V[1, :])")
println("σ error: $(σ[1] - f.S[1])")
println("  error: $(mean(norm.(α .* (A .* v .- u .* σ))))")
scene = lines(dropdims(mean(ϵ; dims = 2); dims = 2), color = :blue)
axis = scene[Axis]
axis[:names][:axisnames] = ("Iteration #", "Loss")
scene = title(scene, "Iterative SVD Loss Over $T Iterations", textsize = 15)

scene


"""
    FourierBasis([::Type,] num_inputs::Int, dorder::Int, iorder::Int [, full::Bool=false])

Creates a struct to generate fourier features up to a given order. Both coupled fourier features, e.g.,
``\\cos(π(3x₁ + 2x₂))`` and uncoupled features, ``[\\cos(πx₁), \\cos(π2x₁), …]`` can be generated with this basis function.
The dorder parameter controls order for the coupling features. The number of coupled features generated, ``(\\text{dorder}+1)^{\\text{num\\_inputs}}``,
 grows exponentially with dorder, so it is not reccomended for use with high deminsional vectors.
The iorder parameter controls the order of the independent features generated. The full parameter determines if both
``\\sin`` and ``\\cos`` features are generated, if false only cos features are generated.

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

# Examples
```jldoctest
julia> f = FourierBasis(2, 1, 2);

julia> x = [0.0, 0.5];

julia> feats = f(x)
6-element Array{Float64,1}:
  1.0
  6.123233995736766e-17
  1.0
  6.123233995736766e-17
  1.0
 -1.0

```
"""
struct FourierBasis{T,TF} <: Any where {T,TF}
    C::T

    function FourierBasis(::Type{T}, num_inputs::Int, dorder::Int, iorder::Int, both::Bool=false) where {T}
        C = make_Cmat(T, num_inputs, dorder, iorder)
        new{typeof(C),both}(C)
    end

    function FourierBasis(num_inputs::Int, dorder::Int, iorder::Int, both::Bool=false) where {T}
        return FourierBasis(Float64, num_inputs, dorder, iorder, both)
    end
end

"""
    FourierBasisBuffer(ϕ::FourierBasis)

Creates preallocated buffers for a fourier basis to use to avoid making allocations on every call of the basis.

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

# Examples
```jldoctest
julia> f = FourierBasis(2, 1, 2);

julia> buff = FourierBasisBuffer(f);

julia> x = [0.0, 0.5];

julia> feats = f(buff, x)
6-element Array{Float64,1}:
  1.0
  6.123233995736766e-17
  1.0
  6.123233995736766e-17
  1.0
 -1.0
"""
struct FourierBasisBuffer{T} <: Any
    y::T
    z::T

    function FourierBasisBuffer(ϕ::FourierBasis{TC,TB}) where {TC, TB}
        T = eltype(ϕ.C)
        n = size(ϕ.C, 2)
        y = zeros(T, n)
        if TB == true
            z = zeros(T, n*2)
        else
            z = zeros(T, n)
        end
        return new{typeof(y)}(y, z)
    end
end

function increment_counter!(counter::Vector{Int}, maxDigit::Int)
    for i in length(counter):-1:1
        counter[i] += 1
        if (counter[i] > maxDigit)
            counter[i] = 0
        else
            break
        end
    end
end

function make_Cmat(T::Type, num_inputs::Int, dorder::Int, iorder::Int)
    iTerms = iorder * num_inputs
    dTerms = (dorder+1)^num_inputs
    oTerms = min(iorder, dorder) * num_inputs
    num_feats = iTerms + dTerms - oTerms
    C = zeros(T, (num_inputs, num_feats))
    counter = zeros(Int, num_inputs)
    termCount::Int = 1

    while termCount <= dTerms
        for i in 1:num_inputs
            C[i, termCount] = counter[i]
        end
        increment_counter!(counter, dorder)
        termCount += 1
    end
    for i in 1:num_inputs
        for j in (dorder+1):(iorder)
            C[i, termCount] = j
            termCount += 1
        end
    end
    C .*= π

    return C
end

function (ϕ::FourierBasis{T,false})(x) where {T}
    y = zeros(size(ϕ.C, 2))
    mul!(y, ϕ.C', x)
    @. y = cos(y)
    return y
end

function (ϕ::FourierBasis{T,false})(buff::FourierBasisBuffer, x) where {T}
    y = buff.y
    mul!(y, ϕ.C', x)
    @. y = cos(y)
    return y
end

function (ϕ::FourierBasis{T,true})(x) where {T}
    n = size(ϕ.C, 2)
    y = zeros(n)
    z = zeros(n*2)
    mul!(y, ϕ.C', x)
    @. z[1:n] = cos(y)
    @. z[n+1:end] = sin(y)
    return z
end

function (ϕ::FourierBasis{T,true})(buff::FourierBasisBuffer, x) where {T}
    y, z = buff.y, buff.z
    n = length(y)
    mul!(y, ϕ.C', x)
    @. z[1:n] = cos(y)
    @. z[n+1:end] = sin(y)
    return z
end

"""
    length(ϕ::FourierBasis)

Returns the number of feautes produced by the Fourier basis.
"""
function length(ϕ::FourierBasis{T,false}) where {T}
    return size(ϕ.C, 2)
end

function length(ϕ::FourierBasis{T,true}) where {T}
    return size(ϕ.C, 2)*2
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _precess{T<:AbstractFloat}(ra::T, dec::T, equinox1::T, equinox2::T,
                                    FK4::Bool, radians::Bool)
    if radians
        ra_rad  = ra
        dec_rad = dec
    else
        # Convert to radians.
        ra_rad  = deg2rad(ra)
        dec_rad = deg2rad(dec)
    end
    a = cos(dec_rad)
    x = [a*cos(ra_rad), a*sin(ra_rad), sin(dec_rad)]
    r = premat(equinox1, equinox2, FK4=FK4)
    x2 = r*x
    ra_rad  = atan2(x2[2], x2[1])
    dec_rad = asin(x2[3])
    if radians
        return cirrange(ra_rad, 2.0*pi), dec_rad
    else
        return cirrange(rad2deg(ra_rad)), rad2deg(dec_rad)
    end
end

"""
    precess(ra, dec, equinox1, equinox2[, FK4=true, radians=true]) -> prec_ra, prec_dec

### Purpose ###

Precess coordinates from `equinox1` to `equinox2`.

### Explanation ###

The default `(ra, dec)` system is FK5 based on epoch J2000.0 but FK4 based on
B1950.0 is available via the `FK4` boolean keyword.

### Arguments ###

* `ra`: input right ascension, scalar or vector, in degrees, unless the
  `radians` keyword is set to `true`
* `dec`: input declination, scalar or vector, in degrees, unless the `radians`
  keyword is set to `true`
* `equinox1`: original equinox of coordinates, numeric scalar.
* `equinox2`: equinox of precessed coordinates.
* `FK4` (optional boolean keyword): if this keyword is set to `true`, the FK4
  (B1950.0) system precession angles are used to compute the precession matrix.
  When it is `false`, the default, use FK5 (J2000.0) precession angles.
* `radians` (optional boolean keyword): if this keyword is set to `true`, then
  the input and output right ascension and declination vectors are in radians
  rather than degrees.

### Output ###

The 2-tuple `(ra, dec)` of coordinates modified by precession.

### Example ###

The Pole Star has J2000.0 coordinates (2h, 31m, 46.3s, 89d 15' 50.6"); compute
its coordinates at J1985.0

``` julia
ra, dec = ten(2,31,46.3)*15, ten(89,15,50.6)
# => (37.94291666666666,89.26405555555556)
adstring(precess(ra, dec, 2000, 1985), precision=1)
# => " 02 16 22.73  +89 11 47.3"
```

Precess the B1950 coordinates of Eps Ind (RA = 21h 59m,33.053s, DEC = (-56d,
59', 33.053") to equinox B1975.

``` julia
ra, dec = ten(21, 59, 33.053)*15, ten(-56, 59, 33.053)
# => (329.88772083333333,-56.992514722222225)
adstring(precess(ra, dec, 1950, 1975, FK4=true), precision=1)
# => " 22 01 15.46  -56 52 18.7"
```

### Method ###

Algorithm from "Computational Spherical Astronomy" by Taff (1983),
p. 24. (FK4). FK5 constants from "Explanatory Supplement To The Astronomical
Almanac" 1992, page 104 Table 3.211.1
(https://archive.org/details/131123ExplanatorySupplementAstronomicalAlmanac).

### Notes ###

Accuracy of precession decreases for declination values near 90 degrees.
`precess` should not be used more than 2.5 centuries from 2000 on the FK5 system
(1950.0 on the FK4 system).  If you need better accuracy, use `bprecess` or
`jprecess` as needed.

Code of this function is based on IDL Astronomy User's Library.
"""
precess(ra::Real, dec::Real, equinox1::Real, equinox2::Real;
        FK4::Bool=false, radians::Bool=false) =
            _precess(promote(float(ra), float(dec),
                             float(equinox1), float(equinox2))...,
                     FK4, radians)

precess(radec::Tuple{Real, Real}, equinox1::Real, equinox2::Real;
        FK4::Bool=false,
        radians::Bool=false) =
            precess(radec[1], radec[2], equinox1, equinox2,
                    FK4=FK4, radians=radians)

function precess{R<:Real, D<:Real}(ra::AbstractArray{R},
                                   dec::AbstractArray{D}, equinox1::Real,
                                   equinox2::Real; FK4::Bool=false,
                                   radians::Bool=false)
    @assert length(ra) == length(dec) "ra and dec arrays should be of the same length"
    typera = typeof(float(one(R)))
    ra_out  = similar(ra,  typera)
    dec_out = similar(dec, typera)
    for i in eachindex(ra)
        ra_out[i], dec_out[i] = precess(ra[i], dec[i], equinox1, equinox2,
                                        FK4=FK4, radians=radians)
    end
    return ra_out, dec_out
end


srand(0)

facts("qraddcol") do

    context("β = 0") do
        m = 100
        A = randn(m,0)
        R = Array{Float64}(0,0)

        for i in 1:m
            a = randn(m)
            R = qraddcol(A, R, a)
            A = [A a]
            @fact  vecnorm( R'*R - A'*A ) --> less_than(1e-5)
        end
    end

    context("β > 0") do
        m = 100
        A = randn(m,0)
        R = Array{Float64}(0,0)
        β = 0.1
        for i in 1:m
            a = randn(m)
            B = [A a]
            R = qraddcol(A, R, a, β)
            @fact vecnorm( R'*R - B'*B - β^2*I ) --> less_than(1e-5)
        end
    end

end


import LinearAlgebra: norm

function get_errors(BIEsol, RCWAsol)
    x, y₁, y₂ = BIEsol.x, BIEsol.y₁, BIEsol.y₂
    err = [
        getproperty(BIEsol, sym) - getproperty(RCWAsol, sym)
        for sym in (:I₁, :O₁, :I₂, :O₂)
    ]
    (x=x, y₁=y₁, y₂=y₂, I₁=err[1], O₁=err[2], I₂=err[3], O₂=err[4])
end

function L2cubature(sol; dims)
    margL2 = [
        mapslices(norm, e; dims=dims) ./ sqrt(prod(size(e)[dims]))
        for e in (sol.I₁, sol.O₁, sol.I₂, sol.O₂)
    ]
    (I₁=margL2[1], O₁=margL2[2], I₂=margL2[3], O₂=margL2[4])
end

function errorL2(solerrors)
    err = L2cubature(solerrors; dims=1:2)
    (I₁=err.I₁[1], O₁=err.O₁[1], I₂=err.I₂[1], O₂=err.O₂[1])
end
function errorbydistance(solerrors)
    err = L2cubature(solerrors; dims=1)
    (y₁=solerrors.y₁, y₂=solerrors.y₂, I₁=vec(err.I₁), O₁=vec(err.O₁), I₂=vec(err.I₂), O₂=vec(err.O₂))
end


"""
$(TYPEDEF)

Backward Euler's integration scheme with given step-size.

### Fields

- `Δt` -- step-size

"""
struct BackwardEuler{N} <: AbstractSolver
    Δt::N
end

BackwardEuler(; Δt::N) where N = BackwardEuler(Δt)

function _init(alg::BackwardEuler, M, C, K)
    Δt = alg.Δt
    K̂ = K * Δt + C
    return K̂
end

function _solve(alg::BackwardEuler{N},
                ivp::InitialValueProblem{ST, XT},
                NSTEPS::Int; kwargs...) where {N, VT, ST, XT<:Tuple{VT, VT}}

    sys = system(ivp)
    (U₀, _) = initial_state(ivp)

    IMAX = NSTEPS + 1
    M, C, K, R = _unwrap(sys, IMAX)

    iszero(M) || throw(ArgumentError("this method assumes that the system is of first order"))

    K̂ = _init(alg, M, C, K)
    K̂⁻¹ = factorize(K̂)

    # initialize
    U = Vector{VT}(undef, IMAX)
    U[1] = U₀

    Δt = alg.Δt

    @inbounds for i in 1:NSTEPS
        R̂ᵢ₊₁ = R[i+1] * Δt + C*U[i]
        # solve
        U[i+1] = K̂⁻¹ \ R̂ᵢ₊₁
    end

    return _build_solution(alg, U, NSTEPS)
end

function _build_solution(alg::BackwardEuler{N}, U, NSTEPS) where {N}
    t = range(zero(N), step=alg.Δt, length=(NSTEPS+1))
    return Solution(alg, U, t)
end


################################################################################
#
#   fbs : Standard Forward Backward splitting
#
################################################################################

function fbs!{T<:FloatingPoint}(
    rmodel::SupervisedRiskModel,     # model definition
    θ::StridedArray{T},              # initial parameter guess
    x::StridedArray{T},              # inputs
    y::StridedArray{T};              # responses
    reg::Regularizer = ZeroReg(),    # regularizer
    λ::T             = one(T),       # step size
    cbinterval::Int  = 0,            # callback interval
    callback         = fbs_trace,    # callback function
    maxsteps::Int    = 100,          # max number of steps
    maxtime::Int     = 0         )   # max calculation time (0 = no limit)

    g    = similar(θ)
    θbar = similar(θ)
    starttime = time()
    for t in 1:maxsteps
        if (maxtime != 0) && (time() - starttime > maxtime)
            break
        end

        copy!(θbar, θ)
        v_risk, _ = value_and_addgrad!(rmodel, 0., g, 1., θbar, x, y)
        axpy!(-λ, g, θbar)
        prox!(reg, θ, θbar, λ)

        # callback
        if cbinterval > 0 && t % cbinterval == 0
            callback(t, θ, v_risk, value(reg, θ))
        end
    end

    θ
end

# Non mutating version
function fbs{T<:FloatingPoint}(
    rmodel::SupervisedRiskModel,     # model definition
    θ::StridedArray{T},              # initial parameter guess
    x::StridedArray{T},              # inputs
    y::StridedArray{T};              # responses
    vargs... )

    fbs!(rmodel, copy(θ), x, y; vargs...)
end


using ModelingToolkit, OrdinaryDiffEq, Test

@parameters t σ ρ β
@variables x(t) y(t) z(t) a(t) u(t) F(t)
@derivatives D'~t

test_equal(a, b) = @test isequal(simplify(a, polynorm=true), simplify(b, polynorm=true))

eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ a*y - β*z,
       0 ~ x - a]

lorenz1 = ODESystem(eqs,t,[x,y,z,a],[σ,ρ,β],name=:lorenz1)

lorenz1_aliased = alias_elimination(lorenz1)
@test length(equations(lorenz1_aliased)) == 3
@test length(states(lorenz1_aliased)) == 3

eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]

@test lorenz1_aliased == ODESystem(eqs,t,[x,y,z],[σ,ρ,β],observed=[a ~ x],name=:lorenz1)

# Multi-System Reduction

eqs1 = [D(x) ~ σ*(y-x) + F,
       D(y) ~ x*(ρ-z)-u,
       D(z) ~ x*y - β*z]

aliases = [u ~ x + y - z]

lorenz1 = ODESystem(eqs1,pins=[F],observed=aliases,name=:lorenz1)

eqs2 = [D(x) ~ F,
       D(y) ~ x*(ρ-z)-x,
       D(z) ~ x*y - β*z]

aliases2 = [u ~ x - y - z]

lorenz2 = ODESystem(eqs2,pins=[F],observed=aliases2,name=:lorenz2)

connections = [lorenz1.F ~ lorenz2.u,
               lorenz2.F ~ lorenz1.u]

connected = ODESystem([0 ~ a + lorenz1.x - lorenz2.y],t,[a],[],observed=connections,systems=[lorenz1,lorenz2])

# Reduced Unflattened System
#=

connections2 = [lorenz1.F ~ lorenz2.u,
                lorenz2.F ~ lorenz1.u,
                a ~ -lorenz1.x + lorenz2.y]
connected = ODESystem(Equation[],t,[],[],observed=connections2,systems=[lorenz1,lorenz2])
=#

# Reduced Flattened System

flattened_system = ModelingToolkit.flatten(connected)

aliased_flattened_system = alias_elimination(flattened_system)

@test isequal(states(aliased_flattened_system), [
        lorenz1.x
        lorenz1.y
        lorenz1.z
        lorenz2.x
        lorenz2.y
        lorenz2.z
       ])

@test setdiff(parameters(aliased_flattened_system), [
        lorenz1.σ
        lorenz1.ρ
        lorenz1.β
        lorenz1.F
        lorenz2.F
        lorenz2.ρ
        lorenz2.β
       ]) |> isempty

test_equal.(equations(aliased_flattened_system), [
       D(lorenz1.x) ~ lorenz1.σ*(lorenz1.y-lorenz1.x) + lorenz2.x - lorenz2.y - lorenz2.z,
       D(lorenz1.y) ~ lorenz1.x*(lorenz1.ρ-lorenz1.z)-(lorenz1.x + lorenz1.y - lorenz1.z),
       D(lorenz1.z) ~ lorenz1.x*lorenz1.y - lorenz1.β*lorenz1.z,
       D(lorenz2.x) ~ lorenz1.x + lorenz1.y - lorenz1.z,
       D(lorenz2.y) ~ lorenz2.x*(lorenz2.ρ-lorenz2.z)-lorenz2.x,
       D(lorenz2.z) ~ lorenz2.x*lorenz2.y - lorenz2.β*lorenz2.z])

test_equal.(observed(aliased_flattened_system), [
       lorenz1.F ~ lorenz2.x + -1 * (lorenz2.y + lorenz2.z),
       lorenz1.u ~ lorenz1.x + lorenz1.y + -1 * lorenz1.z,
       lorenz2.F ~ lorenz1.x + lorenz1.y + -1 * lorenz1.z,
       a ~ lorenz2.y + -1 * lorenz1.x,
       lorenz2.u ~ lorenz2.x + -1 * (lorenz2.y + lorenz2.z),
])

# issue #578

let
    @variables t x(t) y(t) z(t);
    @derivatives D'~t;
    eqs = [
        D(x) ~ x + y
        x ~ y
    ];
    sys = ODESystem(eqs, t, [x], []);
    asys = alias_elimination(ModelingToolkit.flatten(sys))

    test_equal.(asys.eqs, [D(x) ~ 2x])
    test_equal.(asys.observed, [y ~ x])
end


using LinearAlgebra

struct Simplex <: OptimizationMethod
  lambda::Float64
end

## SIMPLEX METHODS ##
function compute_simplex_size(vert::Vector{Vector{T}}) where {T}
  center = sum(vert)
  multiply_array_by_self!(center, 1 / length(vert))

  result = 0.0
  @inbounds @simd for i in vert
    tmp = i - center
    result += sqrt(dot(tmp, tmp))
  end
  return result / length(vert)
end

function minimize!(simplex::Simplex, p::Problem, end_criteria::EndCriteria)
  xtol = end_criteria.rootEpsilon
  max_stationary_state_iterations = end_criteria.maxStationaryStateIterations
  reset!(p)
  x = p.currentValue
  iter_num = 0

  # initialize the vertices of the simplex
  end_condition = false
  n = length(x)
  vertices = Vector{Vector{Float64}}(undef, n + 1)
  fill!(vertices, x)
  direction = zeros(n)
  @simd for i = 1:n
    @inbounds direction[i] = 1.0
    @inbounds vertices[i + 1] = update(p.constraint, vertices[i + 1], direction, simplex.lambda)
    # reset direction
    @inbounds direction[i] = 0.0
  end
  # initialize function values at the vertices of the simplex
  values = zeros(n + 1)
  # values = Vector{DD}(n + 1)
  @simd for i=1:n+1
    @inbounds values[i] = value!(p, vertices[i])
  end

  sum_array = zeros(n + 1)
  # sum_array = Vector{DD}(n + 1)
  # loop through looking for the minimum
  while !end_condition
    # sum_array = zeros(n)
    sum_array = sum(vertices)
    # determine the best (i_lowest) and worst (i_highest) and
    # 2nd worst (i_next_highest) vertices
    i_lowest = 1
    if values[1] < values[2]
      i_highest = 2
      i_next_highest = 1
    else
      i_highest = 1
      i_next_highest = 2
    end

    # we might be able to just do a sort here
    for i=2:n+1
      @inbounds if values[i] > values[i_highest]
        i_next_highest = i_highest
        i_highest = i
      else
        @inbounds if values[i] > values[i_next_highest] && i != i_highest
          i_next_highest = i
        end
      end

      @inbounds if values[i] < values[i_lowest]
        i_lowest = i
      end
    end

    simplex_size = compute_simplex_size(vertices)
    iter_num += 1
    if simplex_size < xtol || iter_num >= end_criteria.maxIterations
      ## TODO implement reason for exiting (EndResult:Type)
      @inbounds x = vertices[i_lowest]
      @inbounds low = values[i_lowest]
      p.functionValue = low
      p.currentValue = x

      return p
    end
    # if end criteria not met, continue
    factor = -1.0
    vTry = extrapolate!(p, i_highest, factor, values, sum_array, vertices)
    @inbounds if vTry <= values[i_lowest] && factor == -1.0
      factor = 2.0
      extrapolate!(p, i_highest, factor, values, sum_array, vertices)
    elseif abs(factor) > EPS_VAL
      @inbounds if vTry >= values[i_next_highest]
        @inbounds vSave = values[i_highest]
        factor = 0.5
        vTry = extrapolate!(p, i_highest, factor, values, sum_array, vertices)
        if vTry >= vSave && abs(factor) > EPS_VAL
          # println("vTry: $vTry")
          # println("vSave: $vSave")
          # println("i_highest: $i_highest")
          # println("i_next_highest: $i_next_highest")
          # println("ending: $iter_num")
          # error("FULL STOP")
          for i = 1:n + 1
            if i != i_lowest
              @inbounds vertices[i] = 0.5 * (vertices[i] + vertices[i_lowest])
              @inbounds values[i] = value!(p, vertices[i])
            end
          end
        end
      end
    end

    if abs(factor) <= EPS_VAL
      x = vertices[i_lowest]
      low = values[i_lowest]
      p.functionValue = low
      p.currentValue = x

      return p
    end
  end
end

centroid(p::Matrix) = reshape(mean(p, 2), size(p, 1))

function dominates(x::Vector, y::Vector)
  @simd for i in 1:length(x)
    @inbounds if x[i] <= y[i]
      return false
    end
  end
  return true
end

function dominates(x::Real, y::Vector)
  @simd for i in 1:length(y)
    @inbounds if x <= y[i]
      return false
    end
  end
  return true
end

nmobjective(y::Vector, m::Integer, n::Integer) = sqrt(var(y) * (m / n))

function minimize_2!(simplex::Simplex, prob::Problem, end_criteria::EndCriteria)
  xtol = ftol = end_criteria.rootEpsilon
  max_stationary_state_iterations = end_criteria.maxStationaryStateIterations
  reset!(prob)
  initial_x = prob.currentValue
  iter_num = 0
  m = length(initial_x)
  initial_step = ones(length(initial_x))
  a = 1.0
  g = 2.0
  b = 0.5

  n = m + 1
  p = repmat(initial_x, 1, n)
  for i = 1:m
    @inbounds p[i, i] += initial_step[i]
  end

  # Maintain a record of the value of f() at n points
  y = Array(Float64, n)
  for i = 1:n
    @inbounds y[i] = value!(prob, p[:, i])
  end

  f_x_previous, f_x = NaN, nmobjective(y, m, n)

  # Count iterations
  iteration = 0

  # Cache p_bar, y_bar, p_star, and p_star_star
  p_bar = Array(Float64, m)
  y_bar = Array(Float64, m)
  p_star = Array(Float64, m)
  p_star_star = Array(Float64, m)

  # Iterate until convergence or exhaustion
  f_converged = false

  while !f_converged && iteration < end_criteria.maxIterations
    # augment iter coutner
    iteration += 1

    # Find p_l and p_h, the min and max values of f() among p
    y_l, l = findmin(y)
    @inbounds p_l = p[:, l]
    y_h, h = findmax(y)
    @inbounds p_h = p[:, h]

    # Compute the centroid of the non-maximal points
    # also cache function values of all non-maximal points
    fill!(p_bar, 0.0)
    tmpindex = 0
    for i = 1:n
      if i != h
        tmpindex += 1
        @inbounds y_bar[tmpindex] = y[i]
        @inbounds p_bar[:] += p[:, i]
      end
    end

    # for j = 1:m
    #   @inbounds p_bar[j] /= m
    # end
    divide_array_by_self!(p_bar, m)

    # Compute a reflection
    for j = 1:m
      @inbounds p_star[j] = (1 + a) * p_bar[j] - a * p_h[j]
    end
    y_star = value!(prob, p_star)

    if y_star < y_l
      # Compute an expansion
      for j = 1:m
        @inbounds p_star_star[j] = g * p_star[j] + (1 - g) * p_bar[j]
      end
      y_star_star = value!(prob, p_star_star)

      if y_star_star < y_l
        @inbounds p_h[:] = p_star_star
        @inbounds p[:, h] = p_star_star
        @inbounds y[h] = y_star_star
      else
        p_h = p_star
        @inbounds p[:, h] = p_star
        @inbounds y[h] = y_star
      end
    else
      if dominates(y_star, y_bar)
        if y_star < y_h
          @inbounds p_h[:] = p_star
          @inbounds p[:, h] = p_h
          @inbounds y[h] = y_star
        end

        # Compute a contraction
        for j = 1:m
          @inbounds p_star_star[j] = b * p_h[j] + (1 - b) * p_bar[j]
        end
        y_star_star = value!(prob, p_star_star)

        if y_star_star > y_h
          for i = 1:n
            for j = 1:m
              @inbounds p[j, i] = (p[j, i] + p_l[j]) / 2.0
            end
            @inbounds y[i] = value!(prob, p[:, i])
          end
        else
          @inbounds p_h[:] = p_star_star
          @inbounds p[:, h] = p_h
          @inbounds y[h] = y_star_star
        end
      else
        @inbounds p_h[:] = p_star
        @inbounds p[:, h] = p_h
        @inbounds y[h] = y_star
      end
    end

    f_x_previous, f_x = f_x, nmobjective(y, m, n)

    if f_x <= ftol
      f_converged = true
    end
  end

  minimum = centroid(p)
  prob.functionValue = Float64(value!(prob, minimum))
  prob.currentValue = minimum

  return prob
end


using SortFilters
using Test
using Random
using Statistics

len = 100
Random.seed!(0)
T = Int
x = rand(T(-10):T(10), len)
window = 9

mx_ref = similar(x)
mn_ref = similar(x)
med_ref = similar(x, Float64)
p20_ref = similar(x, Float64)
p25_p75_ref = similar(x, NTuple{2,Float64})

for i = 1:len
    chunk = x[max(1, i-window+1) : i]
    mx_ref[i] = maximum(chunk)
    mn_ref[i] = minimum(chunk)
    med_ref[i] = median(chunk)
    p20_ref[i] = quantile(chunk, 0.20)
    p25_p75_ref[i] = quantile(chunk, (0.25, 0.75))
end

@testset "movsort" begin

mx = movsort(x, window, 1.0)
mx_ = movsort(x, window, 100)

mn = movsort(x, window, 0.0)
mn_ = movsort(x, window, 0)

med = movsort(x, window, 0.5)
med_ = movsort(x, window, 50)

p25_p75 = movsort(x, window, (0.25, 0.75))
p25_p75_ = movsort(x, window, (25, 75))

@test mx == mx_ == mx_ref
@test mn == mn_ == mn_ref
@test med[window:end] == med_[window:end] == med_ref[window:end]
@test p25_p75[window:end] == p25_p75_[window:end] == p25_p75_ref[window:end]

p20 = movsort(x, window, 0.20)
@test p20[window:end] == p20_ref[window:end] # should match nearest index quantiles

end

@testset "movsort!" begin

mx 		 = similar(mx_ref, Float64)
mx_ 	 = similar(mx_ref)
mn 		 = similar(mn_ref, Float64)
mn_ 	 = similar(mn_ref)
med 	 = similar(med_ref, Float64)
med_ 	 = similar(med_ref)
p25_p75  = similar(p25_p75_ref, NTuple{2,Float64})
p25_p75_ = similar(p25_p75_ref)
p20      = similar(p20_ref, Float64)

movsort!(mx, x, window, 1.0)
movsort!(mx_, x, window, 100)

movsort!(mn, x, window, 0.0)
movsort!(mn_, x, window, 0)

movsort!(med, x, window, 0.5)
movsort!(med_, x, window, 50)

movsort!(p25_p75, x, window, (0.25, 0.75))
movsort!(p25_p75_, x, window, (25, 75))

@test mx == mx_ == mx_ref
@test mn == mn_ == mn_ref
@test med[window:end] == med_[window:end] == med_ref[window:end]
@test p25_p75[window:end] == p25_p75_[window:end] == p25_p75_ref[window:end]

movsort!(p20, x, window, 0.20)
@test p20[window:end] == p20_ref[window:end] # should not match nearest index quantiles

end


function quadrotor_dynamics!(ẋ,X,u)
      #TODO change concatentations to make faster!
      # Quaternion representation
      # Modified from D. Mellinger, N. Michael, and V. Kumar,
      # "Trajectory generation and control for precise aggressive maneuvers with quadrotors",
      # In Proceedings of the 12th International Symposium on Experimental Robotics (ISER 2010), 2010.

      ## States: X ∈ R^13; q = [s;v]
      # x
      # y
      # z
      # q0
      # q1
      # q2
      # q3
      # xdot
      # ydot
      # zdot
      # omega1
      # omega2
      # omega3

      x = X[1:3]
      q = X[4:7]./norm(X[4:7]) #normalize quaternion
      v = X[8:10]
      omega = X[11:13]

      # Parameters
      m = .5 # mass
      J = Matrix(Diagonal([0.0023; 0.0023; 0.004])) # inertia matrix
      Jinv = Matrix(Diagonal(1.0./[0.0023; 0.0023; 0.004])) # inverted inertia matrix
      g = 9.81 # gravity
      L = 0.1750 # distance between motors

      w1 = u[1]
      w2 = u[2]
      w3 = u[3]
      w4 = u[4]

      kf = 1; # 6.11*10^-8;
      F1 = kf*w1;
      F2 = kf*w2;
      F3 = kf*w3;
      F4 = kf*w4;
      F = [0;0;F1+F2+F3+F4] #total rotor force in body frame

      km = 0.0245;
      M1 = km*w1;
      M2 = km*w2;
      M3 = km*w3;
      M4 = km*w4;
      tau = [L*(F2-F4);L*(F3-F1);(M1-M2+M3-M4)] #total rotor torque in body frame

      ẋ[1:3] = v # velocity in world frame
      ẋ[4:7] = 0.5*qmult(q,[0;omega]) #quaternion derivative
      ẋ[8:10] = [0;0;-g] + (1/m)*qrot(q,F) #acceleration in world frame
      ẋ[11:13] = Jinv*(tau - cross(omega,J*omega)) #Euler's equation: I*ω + ω x I*ω = constraint_decrease_ratio
end

function quadrotor_dynamics_mrp!(ẋ, X, u, params)
      #TODO change concatentations to make faster!
      # Modified Rodrigues Parameter (mrp) representation
      # Modified from D. Mellinger, N. Michael, and V. Kumar,
      # "Trajectory generation and control for precise aggressive maneuvers with quadrotors",
      # In Proceedings of the 12th International Symposium on Experimental Robotics (ISER 2010), 2010.
      # States: X ∈ R^12; r = modified Rodrigues parameters
      # x
      # y
      # z
      # r1
      # r2
      # r3
      # xdot
      # ydot
      # zdot
      # omega1
      # omega2
      # omega3

      x = X[1:3]
      r = X[4:6] # mrp
      v = X[7:9]
      omega = X[10:12]

      # Rotation
      R = mrptodcm(r)

      # Parameters
      m = .5 # mass
      J = Matrix(Diagonal([0.0023; 0.0023; 0.004])) # inertia matrix
      Jinv = Matrix(Diagonal(1.0./[0.0023; 0.0023; 0.004])) # inverted inertia matrix
      g = 9.81 # gravity
      L = 0.1750 # distance between motors

      w1 = u[1]
      w2 = u[2]
      w3 = u[3]
      w4 = u[4]

      kf = 1; # 6.11*10^-8;
      F1 = kf*w1;
      F2 = kf*w2;
      F3 = kf*w3;
      F4 = kf*w4;
      F = [0;0;F1+F2+F3+F4] #total rotor force in body frame

      km = 0.0245;
      M1 = km*w1;
      M2 = km*w2;
      M3 = km*w3;
      M4 = km*w4;
      tau = [L*(F2-F4);L*(F3-F1);(M1-M2+M3-M4)] #total rotor torque in body frame

      ẋ[1:3] = v # velocity in world frame
      ẋ[4:6] = 0.25*((1-r'*r)*ω - 2*cross(ω,r) + 2*(ω'*r)*r) #mrp derivative
      ẋ[7:9] = [0;0;-g]  + (1/m)*(R*F) #acceleration in world frame
      ẋ[11:13] = Jinv * (tau - cross(omega, J * omega)) #Euler's equation: I*ω + ω x I*ω = constraint_decrease_ratio
end

function quadrotor_dynamics(X,u)
      ẋ = zeros(13,1)
      quadrotor_dynamics!(ẋ,X,u)
      ẋ
end

function quadrotor_dynamics_mrp(X,u)
      ẋ = zeros(13,1)
      quadrotor_dynamics_mrp!(ẋ,X,u)
      ẋ
end

## Utilities
"""
@(SIGNATURES)
    Rotate a vector by a quaternion
"""
function qrot(q,r)
      r + 2*cross(q[2:4],cross(q[2:4],r) + q[1]*r)
end

"""
@(SIGNATURES)
    Multiplication of two quaternions (q = [s;v])
"""
function qmult(q1,q2)
      [q1[1]*q2[1] - q1[2:4]'*q2[2:4]; q1[1]*q2[2:4] + q2[1]*q1[2:4] + cross(q1[2:4],q2[2:4])]
end

"""
@(SIGNATURES)
  Transformation from modified Rodrigues parameters (mrp) to rotation matrix.
"""
function mrptodcm(mrp)
    # Converts a vector of Modified Rodrigues Parameters to a Rotation Matrix
    mrp2 = mrp' * mrp;
    S = [0       -mrp[3]   mrp[2];
         mrp[3]   0       -mrp[1];
        -mrp[2]   mrp[1]   0]
    I = diagm(0=>fill(1., 3))
    R = I + (8 * S * S + 4 * (1 - mrp2) * S) / ((1 + mrp2)^2)
    return R
end

# Model
n = 13
m = 4

model = Model(quadrotor_dynamics!,n,m)

# Unconstrained
Q = (1e-1)*Matrix(I,n,n)
Q[4,4] = 1.0; Q[5,5] = 1.0; Q[6,6] = 1.0; Q[7,7] = 1.0
R = (1.0)*Matrix(I,m,m)
# R = (1e-1)*Matrix(I,m,m)
Qf = (1000.0)*Matrix(I,n,n)
tf = 5.0
dt = 0.05

# -initial state
x0 = zeros(n)
x0[1:3] = [0.; 0.; 0.]
q0 = [1.;0.;0.;0.]
x0[4:7] = q0

# -final state
xf = copy(x0)
xf[1:3] = [0.;40.;0.] # xyz position
xf[4:7] = q0

obj_uncon = LQRObjective(Q, R, Qf, tf, x0, xf)

# Model + objective
quadrotor = [model, obj_uncon]

## Constrained

r_quad = 3.0
r_sphere = 3.0
spheres = ((0.,10.,0.,r_sphere),(0.,20.,0.,r_sphere),(0.,30.,0.,r_sphere))
n_spheres = 3

# -control limits
u_min = 0.0
u_max = 10.0

# 3 sphere obstacles
function cI_3obs_quad(c,x,u)
    for i = 1:n_spheres
        c[i] = TrajectoryOptimization.sphere_constraint(x,spheres[i][1],spheres[i][2],spheres[i][3],spheres[i][4]+r_quad)
    end
    c
end

# unit quaternion constraint
function unit_quaternion(c,x,u)
    c[1] = sqrt(x[4]^2 + x[5]^2 + x[6]^2 + x[7]^2) - 1.0
end

obj_uq = TrajectoryOptimization.ConstrainedObjective(obj_uncon,u_min=u_min,u_max=u_max,cE=unit_quaternion)
obj_3obs = TrajectoryOptimization.ConstrainedObjective(obj_uncon,u_min=u_min,u_max=u_max,cI=cI_3obs_quad)#,cE=unit_quaternion)

quadrotor_unit_quaternion = [model, obj_uq]
quadrotor_3obs = [model, obj_3obs, spheres]


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

export logforward!

"""
`logforward!(logalpha,Nt,a,A,y)`

In-place version of [`logforward`](@ref).
"""
function logforward!(alpha::AbstractArray{T},Nt::Integer,
                     a::AbstractVector{T},
                     A::AbstractMatrix{T},
                     y::AbstractArray{T}) where {T <: AbstractFloat}
  Ns = size(A,1)
  fill!(alpha,-T(Inf))
  for j = 1:Ns
    alpha[1,j] = a[j] + y[1,j]
  end

  for t = 2:Nt
    for j = 1:Ns
      for i = 1:Ns
        aij = A[i,j]
        if aij != -T(Inf)
          alpha[t,j] = logadd(alpha[t,j],alpha[t-1,i]+aij+y[t,j])
        end
      end
    end
  end

  logML = -T(Inf)
  for i = 1:Ns
    logML = logadd(logML,alpha[Nt,i])
  end
  return alpha, logML/Nt
end

# forward with constrained path
function logforward!(alpha::AbstractArray{T},
                     Nt::Integer,
                     t2IJ::Vector{Pair{Vector{D},Vector{D}}},
                     A::AbstractMatrix{T},
                     y::AbstractArray{T}) where {D<:Integer, T <: AbstractFloat}
  Ns = size(A,1)

  fill!(alpha,-T(Inf))
  idx0 = unique(t2IJ[1][1])
  for j in idx0
    alpha[1,j] = y[1,j] - log(length(idx0))
  end

  for t = 2:Nt
    I,J = t2IJ[t-1]
    for k in eachindex(I)
      i,j = I[k],J[k]
      alpha[t,j] = logadd(alpha[t,j],alpha[t-1,i]+A[i,j]+y[t,j])
    end
  end

  logML = -T(Inf)
  for i = 1:Ns
    logML = logadd(logML,alpha[Nt,i])
  end
  return alpha, logML/Nt
end

export logforward

"""
`logalpha, logML = logforward(Nt,a,A,y)`

Computes the scaled forward log-probabilities (``\\hat{\\alpha}``) of an Hidden Markov Model (HMM) of `Ns` states with log-transition matrix `A` (must be a `Ns` × `Ns` matrix), initial log-probabilities `a` (must be a `Ns`-long vector) and observation log-likelihoods `y` (must be of size `Nt2` × `Ns` with `Nt2` ≥ `Nt`).

Returns:
* `logalpha` a matrix of size `(Nt,Ns)` containing the forward log-probabilities
* `logML` the log likelihood of the observation normalized by the sequence length, i.e. ``\\log(P(\\mathbf{X}))/N_t``.

`logalpha, logML = logforward(Nt,t2tr,A,y)`

Returns the forward log-probabilities with a constrained path (see [`forward`](@ref)).
"""
function logforward(Nt::Integer,a,A::AbstractMatrix{T},y) where {T <: AbstractFloat}
  Ns = size(A,1)
  alpha = -T(Inf)*ones(T,Nt,Ns)
  logforward!(alpha,Nt,a,A,y)
end

# y is a tensor Nt2 × Ns × Nb
function logforward(Nts::Array{D},
                    as::AbstractMatrix{T},
                    As::AbstractArray{T,3},
                    y::AbstractArray{T,3}) where {D <: Integer,
                                          T <: AbstractFloat}
  Ns = size(As,1)
  Nb = size(y,3) # batch number
  @assert Nb == length(Nts) == size(as,2) == size(As,3)
  alphas = [-T(Inf)*ones(T,Nt,Ns) for Nt in Nts]
  logMLs = zeros(T,Nb)
  for i = 1:Nb
    _, logML = logforward!(alphas[i],Nts[i],view(as,:,i),view(As,:,:,i),view(y,:,:,i))
    logMLs[i] = logML
  end
  return alphas, logMLs
end

# y is a tensor Nt2 × Ns × Nb
function logforward(Nts::Array{D},
                    as::AbstractVector,
                    A::AbstractMatrix{T},
                    y::Array{T,3}) where {D <: Integer,
                                          T <: AbstractFloat}
  Ns = size(A,1)
  Nb = size(y,3) # batch number
  @assert Nb == length(Nts) == length(as)
  alphas = [-T(Inf)*ones(T,Nt,Ns) for Nt in Nts]
  logMLs = zeros(T,Nb)
  for i = 1:Nb
    _, logML = logforward!(alphas[i],Nts[i],as[i],A,view(y,:,:,i))
    logMLs[i] = logML
  end
  return alphas, logMLs
end


@doc raw"""
`DAE`: Differential Algebraic Equation

Defines a differential algebraic initial value problem
```math
\begin{aligned}
\dot{q} (t) &= v(t, q(t)) + u(t, q(t), \lambda(t)) , & q(t_{0}) &= q_{0} , \\
0 &= \phi (t, q(t), \lambda(t)) , & \lambda(t_{0}) &= \lambda_{0} ,
\end{aligned}
```
with vector field ``v``, projection ``u``, algebraic constraint ``\phi=0``,
initial conditions ``q_{0}`` and ``\lambda_{0}``, the dynamical variable ``q``
taking values in ``\mathbb{R}^{d}`` and the algebraic variable ``\lambda``
taking values in ``\mathbb{R}^{m}``.

### Parameters

* `DT <: Number`: data type
* `TT <: Real`: time step type
* `AT <: AbstractArray{DT}`: array type
* `vType <: Function`: type of `v`
* `uType <: Function`: type of `u`
* `ūType <: OptionalFunction`: type of `ū`
* `ϕType <: Function`: type of `ϕ`
* `ψType <: OptionalFunction`: type of `ψ`
* `v̄Type <: Function`: type of `v̄`
* `invType <: OptionalNamedTuple`: invariants type
* `parType <: OptionalNamedTuple`: parameters type
* `perType <: OptionalArray{AT}`: periodicity type

### Fields

* `d`: dimension of dynamical variable ``q`` and the vector field ``v``
* `m`: dimension of algebraic variable ``\lambda`` and the constraint ``\phi``
* `v`: function computing the vector field
* `u`: function computing the projection
* `ū`: function computing the secondary projection field ``\bar{u}`` (optional)
* `ϕ`: algebraic constraint
* `ψ`: secondary constraints (optional)
* `v̄`: function computing an initial guess for the velocity field ``v`` (defaults to `v`)
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q``
* `λ₀`: initial condition for algebraic variable ``\lambda``
* `μ₀`: initial condition for algebraic variable ``μ`` (optional)
* `invariants`: either a `NamedTuple` containing the equation's invariants or `nothing`
* `parameters`: either a `NamedTuple` containing the equation's parameters or `nothing`
* `periodicity`: determines the periodicity of the state vector `q` for cutting periodic solutions

The function `v`, providing the vector field, takes three arguments,
`v(t, q, v)`, the functions `u` and `ϕ`, providing the projection and the
algebraic constraint take four arguments, `u(t, q, λ, u)` and `ϕ(t, q, λ, ϕ)`,
where `t` is the current time, `q` and `λ` are the current solution vectors,
and `v`, `u` and `ϕ` are the vectors which hold the result of evaluating the
vector field ``v``, the projection ``u`` and the algebraic constraint ``\phi``
on `t`, `q` and `λ`.

### Constructors

```julia
DAE(v, u, ū, ϕ, ψ, v̄, t₀, q₀, λ₀, invariants, parameters, periodicity)

DAE(v, u, ϕ, t₀, q₀::StateVector, λ₀::StateVector, μ₀::StateVector=zero(λ₀); kwargs...)
DAE(v, u, ϕ, q₀::StateVector, λ₀::StateVector, μ₀::StateVector=zero(λ₀); kwargs...)
DAE(v, u, ϕ, t₀, q₀::State, λ₀::State, μ₀::State=zero(λ₀); kwargs...)
DAE(v, u, ϕ, q₀::State, λ₀::State, μ₀::State=zero(λ₀); kwargs...)

DAE(v, u, ū, ϕ, ψ, t₀, q₀::StateVector, λ₀::StateVector, μ₀::StateVector=zero(λ₀); kwargs...)
DAE(v, u, ū, ϕ, ψ, q₀::StateVector, λ₀::StateVector, μ₀::StateVector=zero(λ₀); kwargs...)
DAE(v, u, ū, ϕ, ψ, t₀, q₀::State, λ₀::State, μ₀::State=zero(λ₀); kwargs...)
DAE(v, u, ū, ϕ, ψ, q₀::State, λ₀::State, μ₀::State=zero(λ₀); kwargs...)
```

### Example

```julia
    function v(t, q, v)
        v[1] = q[1]
        v[2] = q[2]
    end

    function u(t, q, λ, u)
        u[1] = +λ[1]
        u[2] = -λ[1]
    end

    function ϕ(t, q, λ, ϕ)
        ϕ[1] = q[2] - q[1]
    end

    t₀ = 0.
    q₀ = [1., 1.]
    λ₀ = [0.]

    dae = DAE(v, u, ϕ, t₀, q₀, λ₀)

```

"""
struct DAE{vType <: Callable,
           uType <: Callable, ϕType <: Callable,
           ūType <: OptionalCallable, ψType <: OptionalCallable,
           v̄Type <: Callable,
           invType <: OptionalInvariants,
           parType <: OptionalParameters,
           perType <: OptionalPeriodicity} <: AbstractEquationDAE{invType,parType,perType,ψType}

    v::vType
    u::uType
    ϕ::ϕType
    ū::ūType
    ψ::ψType
    v̄::v̄Type

    invariants::invType
    parameters::parType
    periodicity::perType

    function DAE(v, u, ϕ, ū, ψ, v̄, invariants, parameters, periodicity)
        @assert !isempty(methods(v))
        @assert !isempty(methods(u))
        @assert !isempty(methods(ϕ))
        @assert !isempty(methods(ū)) || ū === nothing
        @assert !isempty(methods(ψ)) || ψ === nothing
        @assert !isempty(methods(v̄))

        new{typeof(v), typeof(u), typeof(ϕ), typeof(ū), typeof(ψ), typeof(v̄),
            typeof(invariants), typeof(parameters), typeof(periodicity)}(
                v, u, ϕ, ū, ψ, v̄, invariants, parameters, periodicity)
    end
end

DAE(v, u, ϕ, ū, ψ; v̄=v, invariants=NullInvariants(), parameters=NullParameters(), periodicity=NullPeriodicity()) = DAE(v, u, ϕ, ū, ψ, v̄, invariants, parameters, periodicity)
DAE(v, u, ϕ; kwargs...) = DAE(v, u, ϕ, nothing, nothing; kwargs...)

GeometricBase.invariants(equation::DAE) = equation.invariants
GeometricBase.parameters(equation::DAE) = equation.parameters
GeometricBase.periodicity(equation::DAE) = equation.periodicity

hasvectorfield(::DAE) = true

function check_initial_conditions(equ::DAE, ics::NamedTuple)
    haskey(ics, :q) || return false
    haskey(ics, :λ) || return false
    if hassecondary(equ)
        haskey(ics, :μ) || return false
        eltype(ics.λ) == eltype(ics.μ) || return false
        typeof(ics.λ) == typeof(ics.μ) || return false
        axes(ics.λ) == axes(ics.μ) || return false
    end
    return true
end

function check_methods(equ::DAE, tspan, ics::NamedTuple, params)
    applicable(equ.v, tspan[begin], ics.q, zero(ics.q), params) || return false
    applicable(equ.u, tspan[begin], ics.q, ics.λ, zero(ics.q), params) || return false
    applicable(equ.ϕ, tspan[begin], ics.q, zero(ics.λ), params) || return false
    return true
end

function datatype(equ::DAE, ics::NamedTuple)
    @assert check_initial_conditions(equ, ics)
    return eltype(ics.q)
end

function arrtype(equ::DAE, ics::NamedTuple)
    @assert check_initial_conditions(equ, ics)
    return typeof(ics.q)
end

_get_v(equ::DAE, params) = (t,q,v)   -> equ.v(t, q, v, params)
_get_u(equ::DAE, params) = (t,q,λ,u) -> equ.u(t, q, λ, u, params)
_get_ϕ(equ::DAE, params) = (t,q,ϕ)   -> equ.ϕ(t, q, ϕ, params)
_get_ū(equ::DAE, params) = (t,q,λ,u) -> equ.ū(t, q, λ, u, params)
_get_ψ(equ::DAE, params) = (t,q,v,ψ) -> equ.ψ(t, q, v, ϕ, params)
_get_v̄(equ::DAE, params) = (t,q,v)   -> equ.v̄(t, q, v, params)

function _functions(equ::DAE)
    if hassecondary(equ)
        (v = equ.v, u = equ.u, ϕ = equ.ϕ, ū = equ.ū, ψ = equ.ψ, v̄ = equ.v̄)
    else
        (v = equ.v, u = equ.u, ϕ = equ.ϕ, v̄ = equ.v̄)
    end
end

function _functions(equ::DAE, params::OptionalParameters)
    if hassecondary(equ)
        (v = _get_v(equ, params), u = _get_u(equ, params), ϕ = _get_ϕ(equ, params), ū = _get_ū(equ, params), ψ = _get_ψ(equ, params), v̄ = _get_v̄(equ, params))
    else
        (v = _get_v(equ, params), u = _get_u(equ, params), ϕ = _get_ϕ(equ, params), v̄ = _get_v̄(equ, params))
    end
end