using ThreeBodyDecay
using Parameters
using Test

@testset "Wigner angle permutations" begin

    mp = 0.938; mK = 0.49367; mpi  = 0.13957; mXic = 2.46867
    tbs = ThreeBodySystem(mp,mK,mpi; m0=mXic)
    σs = randomPoint(tbs.ms)

    @unpack m1,m2,m3,m0 = tbs.ms
    # (23) cosζ31_for1 = cosζ23_for1
    @test cosζk1_for1(2,[σs.σ1,σs.σ2,σs.σ3],[m1^2,m2^2,m3^2,m0^2]) ==
          cosζk1_for1(3,[σs.σ1,σs.σ3,σs.σ2],[m1^2,m3^2,m2^2,m0^2])
    # (31) cosζ21_for2 = cosζ32_for2
    @test cosζk2_for2(3,[σs.σ1,σs.σ2,σs.σ3],[m1^2,m2^2,m3^2,m0^2]) ==
          cosζk2_for2(1,[σs.σ3,σs.σ2,σs.σ1],[m3^2,m2^2,m1^2,m0^2])
    # (12) cosζ32_for3 = cosζ13_for3
    @test cosζk3_for3(1,[σs.σ1,σs.σ2,σs.σ3],[m1^2,m2^2,m3^2,m0^2]) ==
          cosζk3_for3(2,[σs.σ2,σs.σ1,σs.σ3],[m2^2,m1^2,m3^2,m0^2])

    # (123) cosζ31_for1 = cosζ21_for2
    @test cosζk2_for2(1,[σs.σ1,σs.σ2,σs.σ3],[m1^2,m2^2,m3^2,m0^2]) ==
          cosζk1_for1(3,[σs.σ2,σs.σ3,σs.σ1],[m2^2,m3^2,m1^2,m0^2])
    # (123)^2 cosζ31_for1 = cosζ23_for3
    @test cosζk3_for3(2,[σs.σ1,σs.σ2,σs.σ3],[m1^2,m2^2,m3^2,m0^2]) ==
          cosζk1_for1(3,[σs.σ3,σs[1],σs[2]],[m3^2,m1^2,m2^2,m0^2])
    #
end


@testset "Brunt-Vaisala" begin

    Γneutral = AvailablePotentialEnergyFramework.g/AvailablePotentialEnergyFramework.Dryair.cp
    Γstable = Γneutral - 10e-3u"K/m"
    Γunstable = Γneutral + 10e-3u"K/m"

    z = 1u"m" .* collect(0:50:10e3)
    tvprofile(Γ,z) = reshape(repeat(300 .- Γ*z,1,2),1,1,length(z),2)
    tvprofile(Γ :: Quantity, z :: Array{ <: Quantity}) = reshape(repeat(300u"K" .- Γ*z,1,2),1,1,length(z),2)

    function plot_N2(Γ,z)
        tv_profile = tvprofile(Γ,z)
        p1 = plot(profile1[1,1,:,1],z,title = "Tv profile, Γ = $Γ")
        p2 = plot(AvailablePotentialEnergyFramework.compute_N2_attempt(tv_profile,z),z, title = "N2")
        plot(p1,p2)
    end

    @test isapprox( 1u"1/s/s" .* zeros(length(z),2),compute_N2(tvprofile(Γneutral,z),z), atol=1e-10u"1/s/s")

    @test all(compute_N2(tvprofile(Γunstable,z),z) .< 0u"1/s/s" )

    @test all(compute_N2(tvprofile(Γstable,z),z) .> 0u"1/s/s" )

    @test isapprox(  zeros(length(z),2),compute_N2(tvprofile(ustrip(Γneutral),ustrip.(z)),ustrip.(z)), atol=1e-10)

    @test all(compute_N2(tvprofile(ustrip(Γunstable),ustrip.(z)),ustrip.(z)) .< 0 )

    @test all(compute_N2(tvprofile(ustrip(Γstable),ustrip.(z)),ustrip.(z)) .> 0 )

end

@testset "Thermodynamics" begin

    @test get_saturation_vapor_pressure(273.15u"K") == 6.112u"hPa"
    @test get_partial_vapor_pressure(0,1000u"hPa") == 0u"hPa"
    @test get_partial_vapor_pressure(1,1000u"hPa") == 1000u"hPa"/(18.016/28.966 + 1.0)
    @test get_mixing_ratio(0u"hPa",1000u"hPa") == 0
    @test get_mixing_ratio(get_partial_vapor_pressure(0.5,1000.0), 1000.0) == 0.5
    @test get_virtual_temperature(300u"K",0,0) == 300u"K"
    @test specific_humidity_to_mixing_ratio(mixing_ratio_to_specific_humidity(0.5)) ≈ 0.5
    @test mixing_ratio_to_specific_humidity(specific_humidity_to_mixing_ratio(0.5)) ≈ 0.5
    @test unit(get_specific_entropy(300u"K",0.2,1000u"hPa"))== u"J/K/kg"
    @test get_potential_temperature(300u"K",1000u"hPa",1000u"hPa") == 300u"K"
    @test get_potential_temperature(300u"K",1010u"hPa",1000u"hPa") < 300u"K"
    @test get_potential_temperature(300u"K",900u"hPa",1000u"hPa") > 300u"K"
    @test get_virtual_temperature(300u"K",0u"g/kg") == 300u"K"
    @test get_virtual_temperature(300u"K",0u"g/kg") == 300u"K"
    @test get_virtual_temperature(300u"K",10u"g/kg") > 300u"K"
    @test 1unit(surface_sensible_heat_flux_to_buoyancy(300u"K", 100u"W/m^2")) == 1unit(g)/u"s"*u"m"
    @test 1unit(surface_latent_heat_flux_to_buoyancy(300u"K", 100u"W/m^2")) == 1unit(g)/u"s"*u"m"
    @test 1unit(get_buoyancy(0.5u"K",10.5u"K")) == 1unit(g)
    @test 1unit(radiative_heating_rate_to_buoyancy(10u"K",5u"K/s")) == 1u"m/s^2/s"
    @test 1unit(get_density_ideal_gas(10.0u"Pa",5.0u"K")) == 1u"kg/m^3"
    # No units

    @test get_saturation_vapor_pressure(273.15) == 6.112
    @test get_partial_vapor_pressure(0,1000) == 0
    @test get_partial_vapor_pressure(1,1000) == 1000/(18.016/28.966 + 1.0)
    @test get_mixing_ratio(0,1000) == 0
    @test get_potential_temperature(300,1000,1000) == 300
    @test get_potential_temperature(300,1010,1000) < 300
    @test get_potential_temperature(300,900,1000) > 300
    @test get_virtual_temperature(300,0) == 300
    @test get_virtual_temperature(300,0) == 300
    @test get_virtual_temperature(300,10) > 300
    @test ustrip(surface_sensible_heat_flux_to_buoyancy(300u"K", 100u"W/m^2")) ==  surface_sensible_heat_flux_to_buoyancy(300, 100)
    @test ustrip(surface_latent_heat_flux_to_buoyancy(300u"K", 100u"W/m^2")) ==  surface_latent_heat_flux_to_buoyancy(300, 100)
    @test ustrip(get_buoyancy(0.5u"K",10.5u"K")) == get_buoyancy(0.5,10.5)
    @test ustrip(radiative_heating_rate_to_buoyancy(10u"K",5u"K/s")) == ustrip(radiative_heating_rate_to_buoyancy(10,5))
    @test ustrip(get_density_ideal_gas(10.0u"Pa",5.0u"K")) == get_density_ideal_gas(10.0,5.0)
    pres = Dataset(joinpath(@__DIR__,"testfiles/thermoprofile.nc")) do ds
        1u"hPa" .* variable(ds, "PRES")[:,:]
    end
    tabs = Dataset(joinpath(@__DIR__,"testfiles/thermoprofile.nc")) do ds
        1u"K" .* variable(ds, "TABS")[:,:]
    end
    qv = 1e-3u"kg/g" .* 1u"g/kg" .* Dataset(joinpath(@__DIR__,"testfiles/thermoprofile.nc")) do ds
        variable(ds, "QV")[:,:] #was originally in g/kg
    end
    @info size(pres) size(qv) size(tabs)
    r = specific_humidity_to_mixing_ratio.(qv)
    timeindex = 1200
    pparcel = pres[1,timeindex]
    tparcel = tabs[1,timeindex]
    rparcel = r[1,timeindex]

    #I will create a similar profile but with a perturbation to see what happens
    tabs_unstable = copy(tabs)
    tabs_unstable[2:40,:] .- 7.0u"K"
    tabs_unstable[41:end,:] .+ 7.0u"K"

    @test_broken unit(get_buoyancy_of_lifted_parcel(tparcel,rparcel,pparcel,tabs[:,timeindex],r[:,timeindex],pres[:,timeindex])[1]) == u"K"
    @test_broken get_buoyancy_of_lifted_parcel(tparcel,rparcel,pparcel,tabs_unstable[:,timeindex],r[:,timeindex],pres[:,timeindex])
end


import MathOptInterface
const MOI = MathOptInterface

import Linda:
    Column,
    Oracle, update!, optimize!, get_columns, get_dual_bound, get_objective_value
const House = DER.House

mutable struct HHOracle <: Oracle

    h::House
    model::MOI.ModelLike

    var2idx::Dict{Tuple{Symbol, Int, Symbol, Int}, MOI.VariableIndex}    # Variable indices in the model
    con2idx::Dict{Tuple{Symbol, Int, Symbol, Int}, MOI.ConstraintIndex}  # Constraint indices in the model

    # Shadow prices
    farkas::Bool
    π::Vector{Float64}
    σ::Float64
end

function update!(o::HHOracle, farkas::Bool, π::Vector{Float64}, σ::Float64)

    # Sanity checks
    T = o.h.num_timesteps
    length(π) == T || throw(DimensionMismatch("π has length $(length(π)) but should have length 2xT=$(2*T)."))

    # Update shadow prices
    o.farkas = farkas
    o.π .= π
    o.σ  = σ

    # Update objective coefficients
    pnet = [
        o.var2idx[(:house, o.h.index, :pnet, t)]
        for t in 1:T
    ]
    for t in 1:T
        MOI.modify(o.model,
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
            MOI.ScalarCoefficientChange(pnet[t], (!farkas) * o.h.price[t] - o.π[t])
        )
    end

    return nothing
end

function optimize!(o::HHOracle)
    MOI.optimize!(o.model)
    return nothing
end

function get_columns(o::HHOracle)

    T = o.h.num_timesteps

    # TODO: check solution status
    pnet = [
        MOI.get(o.model, MOI.VariablePrimal(),
            o.var2idx[(:house, o.h.index, :pnet, t)]
        )
        for t in 1:T
    ]

    # Compute real cost
    cost = dot(o.h.price, pnet)

    # Get reduced cost
    rc = MOI.get(o.model, MOI.ObjectiveValue()) - o.σ

    # Create column
    col = Column(
        cost,
        o.h.index,
        collect(1:2*T), pnet,
        true
    )

    return [(col, rc)]
end

function get_dual_bound(o::HHOracle)
    if o.farkas
        return -Inf
    else
        return MOI.get(o.model, MOI.ObjectiveValue()) - o.σ
    end
end

# Create an oracle from a house
function build_oracle(m::MOI.ModelLike, h::House)
    # Build pricing problem
    _, var2idx, con2idx = DR.DER.build_model!(h, m)

    # Setup oracle
    T = h.num_timesteps
    return HHOracle(h, m, var2idx, con2idx, false, zeros(T), 0.0)
end


using SemiseparableMatrices, BandedMatrices, LinearAlgebra, Test

include("test_semiseparable.jl")
include("test_almostbanded.jl")
include("test_invbanded.jl")


function simple_estimator(model::Potts, T::Real, Js::AbstractArray, _=nothing)
    nsites = numsites(model)
    nbonds = numbonds(model)
    invQ = 1.0/model.Q

    M = 0.0
    @inbounds for s in 1:nsites
        M += ifelse(model.spins[s]==1, 1.0-invQ, -invQ)
    end
    M /= nsites
    E = 0.0
    @inbounds for b in bonds(model)
        s1, s2 = source(b), target(b)
        E -= ifelse(model.spins[s1] == model.spins[s2], 1.0, 0.0) * Js[bondtype(b)]
    end
    E /= nsites

    res = Measurement()
    res["Magnetization"] = M
    res["|Magnetization|"] = abs(M)
    res["Magnetization^2"] = M^2
    res["Magnetization^4"] = M^4
    res["Energy"] = E
    res["Energy^2"] = E^2
    return res
end

@doc raw"""
    improved_estimator(model::Potts, T::Real, Js::AbstractArray, sw::SWInfo)

Returns the following observables as `Dict{String, Any}` using cluster information `sw`

# Observables
- `"Energy"`
    - energy density
- `"Energy^2"`
    - square of energy density
- `"Magnetization"`
    - magnetization density
- `"|Magnetization|"`
    - absolute value of magnetization density
- `"|Magnetization|^2"`
    - square of magnetization density
- `"|Magnetization|^4"`
    - quadruple of magnetization density
- `"Clustersize^2"`
    - ``\sum_c r_c^2 ``, where ``r_c`` is the size density of ``c``-th cluster
- `"Clustersize^4"`
    - ``\sum_c r_c^4 ``
- `"Clustersize^2 Clustersize^2"`
    - ``\sum_{c\ne c'} r_c^2 r_{c'}^2``
"""
function improved_estimator(model::Potts, T::Real, Js::AbstractArray, sw::SWInfo)
    nsites = numsites(model)
    Q = model.Q
    nc = numclusters(sw)
    invV = 1.0/nsites

    # magnetization
    I2 = (Q-1)/(Q*Q)
    I4 = (Q-1)*((Q-1)^3+1)/(Q^5)
    M = 0.0
    M2 = 0.0
    M4 = 0.0
    s2 = -1.0/Q
    s1 = 1.0+s2

    N4 = 0.0
    N2N2 = 0.0
    N2 = 0.0
    for (m,s) in zip(sw.clustersize, sw.clusterspin)
        M += (m*invV) * ifelse(s==1, s1, s2)
        m2 = (m*invV)^2
        N4 += m2*m2
        N2N2 += N2*m2
        N2 += m2
    end
    M4 = I4*N4 + 6*I2*I2*N2N2
    M2 = I2*N2

    # energy
    aJ = abs.(Js)
    mbeta = -1.0/T
    ns = sw.activated_bonds
    As = -aJ ./ expm1.(mbeta.*aJ)
    Ans = ns.*As
    E = 0.0
    E2 = 0.0
    for b in 1:numbondtypes(model)
        E2 += aJ[b]*Ans[b]
        E2 += Ans[b] * As[b]*(ns[b]-1)
        E2 += 2.0*Ans[b]*E
        E += Ans[b]
    end

    E *= -invV
    E2 *= invV*invV

    res = Measurement()
    res["Magnetization"] = M
    res["|Magnetization|"] = abs(M)
    res["Magnetization^2"] = M2
    res["Magnetization^4"] = M4
    res["Energy"] = E
    res["Energy^2"] = E2
    res["Clustersize^2"] = N2
    res["Clustersize^4"] = N4
    res["Clustersize^2 Clustersize^2"] = 2*N2N2

    return res
end

default_estimator(model::Potts, update) = ifelse(update==SW_update!, improved_estimator, simple_estimator)


"""
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
-                                QUANTUM STATES                                -
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
"""

qubits(N::Int) = siteinds("Qubit", N)

qudits(N::Int; dim::Int = 3) = siteinds("Qudit",N; dim = dim)

#
# State-like gates, used to define product input states
#

# TODO: add an arbitrary state specified by angles

state(sn::String) = state(StateName(sn))
state(sn::String, dim::Int) = state(StateName(sn), dim)

function state(sn::String, i::Index)
  st = hastags(i, "Qubit") ? state(sn) : state(sn, dim(i))
  return ITensors.itensor(st, i)
end

# Pauli eingenstates
state(::StateName"X+") = [
  1 / sqrt(2)
  1 / sqrt(2)
]

state(::StateName"X-") = [
  1 / sqrt(2)
  -1 / sqrt(2)
]

state(::StateName"Y+") = [
  1 / sqrt(2)
  im / sqrt(2)
]

state(::StateName"Y-") = [
  1 / sqrt(2)
  -im / sqrt(2)
]

state(::StateName"Z+") = [
  1
  0
]

state(::StateName"Z-") = [
  0
  1
]

state(::StateName"0") = state("Z+")
state(::StateName"1") = state("Z-")

# SIC-POVMs

state(::StateName"T1") = state("Z+")
state(::StateName"T2") = [
  1/√3
  √2/√3
]
state(::StateName"T3") = [
  1/√3
  √2/√3 * exp(im*2π/3)
]
state(::StateName"T4") = [
  1/√3
  √2/√3 * exp(im*4π/3)
]

function state(::StateName{N}, dim::Int) where {N}
  n = parse(Int, String(N))
  st = zeros(Int64,dim)
  st[n + 1] = 1
  return st
end

"""
    productstate(N::Int)

    productstate(sites::Vector{<:Index})

Initialize qubits to an MPS wavefunction in the 0 state (`|ψ⟩ = |0⟩ ⊗ |0⟩ ⊗ …`).
"""
# TODO: add dimension to use qudit instead of qubit
function productstate(N::Int; dim::Int = 2, sitetype::String = "Qubit")
  dim > 2 && return productstate(siteinds("Qudit", N; dim = dim))
  return productstate(siteinds(sitetype,N))
end

productstate(sites::Vector{<:Index}) = productMPS(sites, "0")

"""
    productstate(M::Union{MPS,MPO,LPDO})

Initialize qubits on the Hilbert space of a reference state,
given as `MPS`, `MPO` or `LPDO`.
"""
productstate(M::Union{MPS,MPO,LPDO}) = productstate(hilbertspace(M))

"""
    productstate(N::Int, states::Vector{T})

    productstate(sites::Vector{<:Index}, states::Vector{T})

    productstate(M::Union{MPS,MPO,LPDO}, states::Vector{T})

Initialize the qubits to a given product state, where the state `T` can be specified either
with a Vector of states specified as Strings or bit values (0 and 1).
"""
function productstate(N::Int, states::Vector; dim::Int = 2, sitetype::String = "Qubit")
  dim > 2 && return productstate(siteinds("Qudit", N; dim = dim), states)
  return productstate(siteinds(sitetype, N), states)
end

function productstate(sites::Vector{<:Index}, states::Vector{<:Integer})
  return MPS(state.(string.(Int.(states)), sites))
end

function productstate(sites::Vector{<:Index}, states::Vector)
  return MPS(state.(states, sites))
end

function productstate(M::Union{MPS,MPO,LPDO}, states::Vector)
  return productstate(hilbertspace(M), states)
end

function productstate(sites::Vector{<:Index}, state::Union{String,Integer})
  return productstate(sites, fill(state, length(sites)))
end

function productstate(sites::Vector{<:Index}, states::Function)
  return productstate(sites, map(states, 1:length(sites)))
end

"""
    productoperator(N::Int)

    productoperator(sites::Vector{<:Index})

Initialize an MPO that is a product of identity operators.
"""
function productoperator(N::Int; dim::Int = 2, sitetype::String = "Qubit")
  dim > 2 && return productoperator(siteinds("Qudit", N; dim = dim))
  return productoperator(siteinds(sitetype, N))
end

productoperator(M::Union{MPS,MPO,LPDO}) = productoperator(hilbertspace(M))

productoperator(sites::Vector{<:Index}) = MPO([op("Id", s) for s in sites])


module FlightMechanicsUtils

using LinearAlgebra
using StaticArrays

export γ_AIR, R_AIR
export gD
export RAD2DEG, DEG2RAD, M2FT, FT2M, KEL2RANK, RANK2KEL, KG2LB, LB2KG, SLUG2KG, KG2SLUG
export PSF2PA, PA2PSF, SLUGFT32KGM3, KGM32SLUGFT3
include("constants.jl")

export atmosphere_isa
include("atmosphere.jl")

export rigid_body_velocity, rigid_body_acceleration
export pqr_2_ψθϕ_dot, pqr_2_quat_dot, ψθϕ_dot_2_pqr
export uvw_to_tasαβ, uvw_dot_to_tasαβ_dot, tasαβ_dot_to_uvw_dot
export rate_of_climb_constrain_no_wind
include("kinematics.jl")

export coordinated_turn_bank, steiner_inertia, translate_forces_moments
include("mechanics.jl")

export euler_angles, quaternions, rotation_matrix_zyx
export body2horizon, body2wind
export wind2body, wind2horizon
export horizon2body, horizon2wind
export horizon2ecef, ecef2horizon
include("rotations.jl")

export Ellipsoid
export Clarke1866, Clarke1880, International, Bessel, Everest, ModifiedEverest,
    AustralianNational, SouthAmerican1969, Airy, ModifiedAiry, Hough, Fischer1960SouthAsia,
    Fischer1960Mercury, Fischer1968, WGS60, WGS66, WGS72, WGS84
export llh2ecef, ecef2llh
include("ellipsoid.jl")

export qc2cas, qc2eas, qc2tas
export cas2eas, cas2tas
export eas2cas, eas2tas
export tas2cas, tas2eas
export compressible_qinf, incompressible_qinf
include("anemometry.jl")

end


module WishartTest

using Test
using ReactiveMP
using Distributions
using Random
using LinearAlgebra

@testset "Wishart" begin

    # Wishart comes from Distributions.jl and most of the things should be covered there
    # Here we test some extra ReactiveMP.jl specific functionality

    @testset "mean(:logdet)" begin
        @test mean(logdet, Wishart(3, [ 1.0 0.0; 0.0 1.0 ])) ≈ 0.845568670196936
        @test mean(logdet, Wishart(5, [1.4659658963311604 1.111775094889733 0.8741034114800605; 1.111775094889733 0.8746971141492232 0.6545661366809246; 0.8741034114800605 0.6545661366809246 0.5498917856395482])) ≈ -3.4633310802040693
    end

    @testset "vague" begin
        @test_throws MethodError vague(Wishart)

        d = vague(Wishart, 3)

        @test typeof(d) <: Wishart
        @test mean(d) == Matrix(Diagonal(3 * ReactiveMP.huge * ones(3)))
    end

    @testset "prod" begin
        v1 = [ 9.0 -3.4; -3.4 11.0 ]
        v2 = [ 10.2 -3.3; -3.3 5.0 ]
        v3 = [ 8.1 -2.7; -2.7 9.0 ]

        @test prod(ProdAnalytical(), Wishart(3, v1), Wishart(3, v2)) ≈ Wishart(3, [4.776325721474591 -1.6199382410125422; -1.6199382410125422 3.3487476649765537])
        @test prod(ProdAnalytical(), Wishart(4, v1), Wishart(4, v3)) ≈ Wishart(5, [4.261143738311623 -1.5064864332819319; -1.5064864332819319 4.949867121624725])
        @test prod(ProdAnalytical(), Wishart(5, v2), Wishart(4, v3)) ≈ Wishart(6, [4.51459128065395 -1.4750681198910067; -1.4750681198910067 3.129155313351499])
    end

    @testset "ndims" begin
        @test ndims(vague(Wishart, 3)) === 3
        @test ndims(vague(Wishart, 4)) === 4
        @test ndims(vague(Wishart, 5)) === 5
    end

end

end


@testset "Complex scaling" begin
    B = FEDVR(range(0,stop=20,length=71), 10)
    C = FEDVR(range(0,stop=20,length=71), 10, t₀=10.0, ϕ=π/3)
    @test C.t₀ ≥ 10.0
    @test_throws ArgumentError FEDVR(range(0,stop=20,length=71), 10, t₀=30.0, ϕ=π/3)
    @test complex_rotate(5, B) == 5
    @test complex_rotate(15, B) == 15
    @test complex_rotate(5, C) == 5
    @test complex_rotate(15, C) ≈ 10 + 5*exp(im*π/3)
end

@testset "Real locations" begin
    rₘₐₓ = 20
    R = FEDVR(range(0,stop=rₘₐₓ,length=11), 10, t₀=rₘₐₓ/2, ϕ=π/3)
    R′ = FEDVR(range(0,stop=rₘₐₓ,length=11), 10)
    @test norm(real_locs(R)-real_locs(R′)) == 0
end


# Copyright 2019 Tobias Frilling
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

@inline coolDown(method::Symbol, args...) = coolDown(Val(method), args...)

"""
    coolDown(::Val{:linear}, (start, stop), steps)

Return a linear cooldown function.

The range `start - stop` is devided into equidistant `steps`.
"""
function coolDown(::Val{:linear}, (start, stop), steps)
    return step -> start - ((start - stop) / steps) * (step - 1)
end

"""
    coolDown(::Val{:leadInOut}, (start, stop), steps; leadIn = 0.1, leadOut = 0.95)

Return a linear cooldown function with lead-in and lead-out.

The range `start - stop` is devided into equidistant `steps`. If the given
`step` is below the `leadIn` quantile, `start` is returned; if it's above the
`leadOut` quantile, `stop` is returned.
"""
function coolDown(::Val{:leadInOut}, (start, stop), steps; leadIn = 0.1, leadOut = 0.95)
    @assert leadIn < leadOut
    @assert leadOut - leadIn >= 0
    return step -> if (step/steps <= leadIn) start
                   elseif (step/steps >= leadOut) stop
                   else start - ((start - stop) / steps) * (step - 1)
                   end
end


g = DiGraph{Node{String}, Edge{Node{String}}}()

add_node!(g, "a")
add_node!(g, "b")
add_edge!(g, "a", "b")

@test length(nodes(g)) == 2
@test length(edges(g)) == 1
@test out_edges(g, "a") == in_edges(g, "b")

g2 = EvolvingGraph{Node{String}, String}()
add_edge!(g2, "a", "b", "t1")
add_edge!(g2, "b", "c", "t1")
add_edge!(g2, "c", "d", "t2")
add_edge!(g2, "a", "b", "t2")

g3 = aggregate_graph(g2)
display(g3)

A = adjacency_matrix(g3)
@test A[1,2] == 1
@test A[2,3] == 1
@test num_nodes(g3) == 4
@test num_edges(g3) == 3

# add a graph to an evolving graph
g = DiGraph{Node{String}, Edge{Node{String}}}()
add_node!(g, "1")
add_node!(g, "2")
add_edge!(g, "1", "2")
eg = EvolvingGraph{Node{String}, Int}()
add_graph!(eg, g, 1)
@test length(nodes(eg)) == 2
@test length(edges(eg)) == 1


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 457e4d2c-0225-11ec-1b21-7b47b5339342
using JuMP, Ipopt

# ╔═╡ 7f3a1eb1-63a7-4062-9dd6-e88e31e961da
model = Model(Ipopt.Optimizer)

# ╔═╡ af9691bd-e434-48e5-bb53-7b63b0d152a9
begin
  @variable(model, x)
  @variable(model, y)
end

# ╔═╡ c199a311-3874-4f2c-a6e9-270cbc40cedd
@NLobjective(model, Min, (1 - x)^2 + 100 * (y - x^2)^2)

# ╔═╡ 8b62d21e-9774-4c5c-9667-c6379e78733c
optimize!(model)

# ╔═╡ 5963c706-4143-443f-aaf0-d61001bfce04
value(x)

# ╔═╡ 2a365ed7-5e15-4b96-851f-2fa556941b51
value(y)

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
JuMP = "4076af6c-e467-56ae-b986-b466b2749572"

[compat]
Ipopt = "~0.7.0"
JuMP = "~0.21.9"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[ASL_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "370cafc70604b2522f2c7cf9915ebcd17b4cd38b"
uuid = "ae81ac8f-d209-56e5-92de-9978fef736f9"
version = "0.1.2+0"

[[ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[BenchmarkTools]]
deps = ["JSON", "Logging", "Printf", "Statistics", "UUIDs"]
git-tree-sha1 = "aa3aba5ed8f882ed01b71e09ca2ba0f77f44a99e"
uuid = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
version = "1.1.3"

[[BinaryProvider]]
deps = ["Libdl", "Logging", "SHA"]
git-tree-sha1 = "ecdec412a9abc8db54c0efc5548c64dfce072058"
uuid = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
version = "0.5.10"

[[Bzip2_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "19a35467a82e236ff51bc17a3a44b69ef35185a2"
uuid = "6e34b625-4abd-537c-b88f-471c36dfa7a0"
version = "1.0.8+0"

[[Calculus]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "f641eb0a4f00c343bbc32346e1217b86f3ce9dad"
uuid = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
version = "0.5.1"

[[ChainRulesCore]]
deps = ["Compat", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "bdc0937269321858ab2a4f288486cb258b9a0af7"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.3.0"

[[CodecBzip2]]
deps = ["Bzip2_jll", "Libdl", "TranscodingStreams"]
git-tree-sha1 = "2e62a725210ce3c3c2e1a3080190e7ca491f18d7"
uuid = "523fee87-0ab8-5b00-afb7-3ecf72e48cfd"
version = "0.7.2"

[[CodecZlib]]
deps = ["TranscodingStreams", "Zlib_jll"]
git-tree-sha1 = "ded953804d019afa9a3f98981d99b33e3db7b6da"
uuid = "944b1d66-785c-5afd-91f1-9de20f533193"
version = "0.7.0"

[[CommonSubexpressions]]
deps = ["MacroTools", "Test"]
git-tree-sha1 = "7b8a93dba8af7e3b42fecabf646260105ac373f7"
uuid = "bbf7d656-a473-5ed7-a52c-81e309532950"
version = "0.3.0"

[[Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "727e463cfebd0c7b999bbf3e9e7e16f254b94193"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.34.0"

[[CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "7d9d316f04214f7efdbb6398d545446e246eff02"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.10"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[DiffResults]]
deps = ["StaticArrays"]
git-tree-sha1 = "c18e98cba888c6c25d1c3b048e4b3380ca956805"
uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
version = "1.0.3"

[[DiffRules]]
deps = ["NaNMath", "Random", "SpecialFunctions"]
git-tree-sha1 = "3ed8fa7178a10d1cd0f1ca524f249ba6937490c0"
uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
version = "1.3.0"

[[Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "a32185f5428d3986f47c2ab78b1f216d5e6cc96f"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.5"

[[Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[ForwardDiff]]
deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "NaNMath", "Printf", "Random", "SpecialFunctions", "StaticArrays"]
git-tree-sha1 = "b5e930ac60b613ef3406da6d4f42c35d8dc51419"
uuid = "f6369f11-7733-5829-9624-2563aa707210"
version = "0.10.19"

[[HTTP]]
deps = ["Base64", "Dates", "IniFile", "Logging", "MbedTLS", "NetworkOptions", "Sockets", "URIs"]
git-tree-sha1 = "44e3b40da000eab4ccb1aecdc4801c040026aeb5"
uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3"
version = "0.9.13"

[[IniFile]]
deps = ["Test"]
git-tree-sha1 = "098e4d2c533924c921f9f9847274f2ad89e018b8"
uuid = "83e8ac13-25f8-5344-8a64-a9f2b223428f"
version = "0.5.0"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[Ipopt]]
deps = ["BinaryProvider", "Ipopt_jll", "Libdl", "LinearAlgebra", "MathOptInterface", "MathProgBase"]
git-tree-sha1 = "380786b4929b8d18d76e909c6b2eca355b7c3bd6"
uuid = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
version = "0.7.0"

[[Ipopt_jll]]
deps = ["ASL_jll", "Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "MUMPS_seq_jll", "OpenBLAS32_jll", "Pkg"]
git-tree-sha1 = "82124f27743f2802c23fcb05febc517d0b15d86e"
uuid = "9cc047cb-c261-5740-88fc-0cf96f7bdcc7"
version = "3.13.4+2"

[[IrrationalConstants]]
git-tree-sha1 = "f76424439413893a832026ca355fe273e93bce94"
uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
version = "0.1.0"

[[JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "642a199af8b68253517b80bd3bfd17eb4e84df6e"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.3.0"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[JSONSchema]]
deps = ["HTTP", "JSON", "ZipFile"]
git-tree-sha1 = "b84ab8139afde82c7c65ba2b792fe12e01dd7307"
uuid = "7d188eb4-7ad8-530c-ae41-71a32a6d4692"
version = "0.3.3"

[[JuMP]]
deps = ["Calculus", "DataStructures", "ForwardDiff", "JSON", "LinearAlgebra", "MathOptInterface", "MutableArithmetics", "NaNMath", "Printf", "Random", "SparseArrays", "SpecialFunctions", "Statistics"]
git-tree-sha1 = "4f0a771949bbe24bf70c89e8032c107ebe03f6ba"
uuid = "4076af6c-e467-56ae-b986-b466b2749572"
version = "0.21.9"

[[LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[LogExpFunctions]]
deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "3d682c07e6dd250ed082f883dc88aee7996bf2cc"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.0"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[METIS_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "2dc1a9fc87e57e32b1fc186db78811157b30c118"
uuid = "d00139f3-1899-568f-a2f0-47f597d42d70"
version = "5.1.0+5"

[[MUMPS_seq_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "METIS_jll", "OpenBLAS32_jll", "Pkg"]
git-tree-sha1 = "1a11a84b2af5feb5a62a820574804056cdc59c39"
uuid = "d7ed1dd3-d0ae-5e8e-bfb4-87a502085b8d"
version = "5.2.1+4"

[[MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "0fb723cd8c45858c22169b2e42269e53271a6df7"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.7"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[MathOptInterface]]
deps = ["BenchmarkTools", "CodecBzip2", "CodecZlib", "JSON", "JSONSchema", "LinearAlgebra", "MutableArithmetics", "OrderedCollections", "SparseArrays", "Test", "Unicode"]
git-tree-sha1 = "575644e3c05b258250bb599e57cf73bbf1062901"
uuid = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
version = "0.9.22"

[[MathProgBase]]
deps = ["LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "9abbe463a1e9fc507f12a69e7f29346c2cdc472c"
uuid = "fdba3010-5040-5b88-9595-932c9decdf73"
version = "0.7.8"

[[MbedTLS]]
deps = ["Dates", "MbedTLS_jll", "Random", "Sockets"]
git-tree-sha1 = "1c38e51c3d08ef2278062ebceade0e46cefc96fe"
uuid = "739be429-bea8-5141-9913-cc70e7f3736d"
version = "1.0.3"

[[MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[MutableArithmetics]]
deps = ["LinearAlgebra", "SparseArrays", "Test"]
git-tree-sha1 = "3927848ccebcc165952dc0d9ac9aa274a87bfe01"
uuid = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
version = "0.2.20"

[[NaNMath]]
git-tree-sha1 = "bfe47e760d60b82b66b61d2d44128b62e3a369fb"
uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
version = "0.3.5"

[[NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[OpenBLAS32_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "ba4a8f683303c9082e84afba96f25af3c7fb2436"
uuid = "656ef2d0-ae68-5445-9ca0-591084a874a2"
version = "0.3.12+1"

[[OpenSpecFun_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.5+0"

[[OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "438d35d2d95ae2c5e8780b330592b6de8494e779"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.0.3"

[[Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[Preferences]]
deps = ["TOML"]
git-tree-sha1 = "00cfd92944ca9c760982747e9a1d0d5d86ab1e5a"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.2"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[SpecialFunctions]]
deps = ["ChainRulesCore", "LogExpFunctions", "OpenSpecFun_jll"]
git-tree-sha1 = "a322a9493e49c5f3a10b50df3aedaf1cdb3244b7"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "1.6.1"

[[StaticArrays]]
deps = ["LinearAlgebra", "Random", "Statistics"]
git-tree-sha1 = "3240808c6d463ac46f1c1cd7638375cd22abbccb"
uuid = "90137ffa-7385-5640-81b9-e52037218182"
version = "1.2.12"

[[Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[TranscodingStreams]]
deps = ["Random", "Test"]
git-tree-sha1 = "216b95ea110b5972db65aa90f88d8d89dcb8851c"
uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
version = "0.9.6"

[[URIs]]
git-tree-sha1 = "97bbe755a53fe859669cd907f2d96aee8d2c1355"
uuid = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
version = "1.3.0"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[ZipFile]]
deps = ["Libdl", "Printf", "Zlib_jll"]
git-tree-sha1 = "c3a5637e27e914a7a445b8d0ad063d701931e9f7"
uuid = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
version = "0.9.3"

[[Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╠═457e4d2c-0225-11ec-1b21-7b47b5339342
# ╠═7f3a1eb1-63a7-4062-9dd6-e88e31e961da
# ╠═af9691bd-e434-48e5-bb53-7b63b0d152a9
# ╠═c199a311-3874-4f2c-a6e9-270cbc40cedd
# ╠═8b62d21e-9774-4c5c-9667-c6379e78733c
# ╠═5963c706-4143-443f-aaf0-d61001bfce04
# ╠═2a365ed7-5e15-4b96-851f-2fa556941b51
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


#=
Copyright 2021 BlackRock, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=#

function run_test(f::PiecewiseQuadratic, h::PiecewiseQuadratic)
    return (simplify(envelope(f)) ≈ h) && (simplify(envelope(reverse(f))) ≈ reverse(h))
end

function plot_test(f::PiecewiseQuadratic, h::PiecewiseQuadratic)
    clf()
    X, Y = get_plot(f)
    plot(X, Y)
    X, Y = get_plot(simplify(envelope(f)))
    plot(X, Y)
    X, Y = get_plot(h)
    plot(X, Y; linestyle=:dash)
    return legend(["f", "env(f)", "f**"])
end

@testset "Midpoint to endpoint" begin
    f = PiecewiseQuadratic([BoundedQuadratic(-Inf, 1.0, 1.0, 2.0, 2.0),
                            BoundedQuadratic(1.0, 1.0, 1.0, 2.0, 1.0),
                            BoundedQuadratic(1.0, Inf, 1.0, 2.0, 2.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(-Inf, 0.0, 1.0, 2.0, 2.0),
                            BoundedQuadratic(0.0, 1.0, 0.0, 2.0, 2.0),
                            BoundedQuadratic(1.0, 2.0, 0.0, 6.0, -2.0),
                            BoundedQuadratic(2.0, Inf, 1.0, 2.0, 2.0)])
    @test run_test(f, h)
    @test run_test(reverse(f), reverse(h))
end

@testset "Convex functions" begin
    f = PiecewiseQuadratic([BoundedQuadratic(-Inf, 1.0, 0.0, -1.0, 1.0),
                            BoundedQuadratic(1.0, Inf, 0.0, 1.0, -1.0)])
    @test run_test(f, f)
    @test run_test(reverse(f), reverse(f))

    f = PiecewiseQuadratic([BoundedQuadratic(1.0, -2.0, 1.0)])
    @test run_test(f, f)
    @test run_test(reverse(f), reverse(f))
end

@testset "Nonconvex piecewise affine" begin
    f = PiecewiseQuadratic([BoundedQuadratic(-Inf, 1.0, 0.0, 1.0, 2.0),
                            BoundedQuadratic(1.0, 1.0, 0.0, 2.0, 0.0),
                            BoundedQuadratic(1.0, Inf, 0.0, 3.0, 0.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(-Inf, 1.0, 0.0, 1.0, 1.0),
                            BoundedQuadratic(1.0, Inf, 0.0, 3.0, -1.0)])
    @test run_test(f, h)
    @test run_test(reverse(f), reverse(h))
end

@testset "Midpoint to midpoint" begin
    f = PiecewiseQuadratic([BoundedQuadratic(-Inf, 0.0, 1.0, 1.0, 1.0),
                            BoundedQuadratic(0.0, Inf, 1.0, -1.0, 1.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(-Inf, -0.5, 1.0, 1.0, 1.0),
                            BoundedQuadratic(-0.5, 0.5, 0.0, 0.0, 0.75),
                            BoundedQuadratic(0.5, Inf, 1.0, -1.0, 1.0)])
    @test run_test(f, h)
end

@testset "Upper to lower" begin
    f = PiecewiseQuadratic([BoundedQuadratic(-Inf, 0.0, 1.0, -1.0, 1.0),
                            BoundedQuadratic(0.0, Inf, 1.0, 1.0, 1.0)])
    @test run_test(f, f)

    f = PiecewiseQuadratic([BoundedQuadratic(-Inf, 0.0, 0.0, -1.0, 1.0),
                            BoundedQuadratic(0.0, Inf, 0.0, 1.0, 1.0)])
    @test run_test(f, f)
end

@testset "Irrelevant point" begin
    f = PiecewiseQuadratic([BoundedQuadratic(0.0, 0.0, 0.0, 0.0, 1.0),
                            BoundedQuadratic(0.0, Inf, 0.0, 1.0, 0.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(0.0, Inf, 0.0, 1.0, 0.0)])
    @test run_test(f, h)
end

@testset "Relevant point" begin
    f = PiecewiseQuadratic([BoundedQuadratic(0.0, 0.0, 0.0, 0.0, -1.0),
                            BoundedQuadratic(0.0, Inf, 0.0, 1.0, 0.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(0.0, Inf, 0.0, 1.0, -1.0)])
    @test run_test(f, h)

    f = PiecewiseQuadratic([BoundedQuadratic(0.0, 0.0, 0.0, 0.0, -1.0),
                            BoundedQuadratic(0.0, 1.0, 0.0, 1.0, 0.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(0.0, 1.0, 0.0, 2.0, -1.0)])
    @test run_test(f, h)
end

@testset "Lower to upper" begin
    f = PiecewiseQuadratic([BoundedQuadratic(-1.0, 0.0, 0.0, 1.0, 1.0),
                            BoundedQuadratic(0.0, 1.0, 0.0, -1.0, 1.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(-1.0, 1.0, 0.0, 0.0, 0.0)])
    @test run_test(f, h)
end

@testset "Upper to upper" begin
    f = PiecewiseQuadratic([BoundedQuadratic(-1.0, 0.0, 0.0, -1.0, 0.0),
                            BoundedQuadratic(0.0, 1.0, 0.0, -1.0, 1.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(-1.0, 0.0, 0.0, -1.0, 0.0),
                            BoundedQuadratic(0.0, 1.0, 0.0, 0.0, 0.0)])
    @test run_test(f, h)
end

@testset "Two points" begin
    f = PiecewiseQuadratic([BoundedQuadratic(0.0, 0.0, 0.0, 0.0, 1.0),
                            BoundedQuadratic(0.0, 0.0, 0.0, 0.0, 0.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(0.0, 0.0, 0.0, 0.0, 0.0)])
    @test run_test(f, h)

    f = PiecewiseQuadratic([BoundedQuadratic(0.0, 0.0, 0.0, 0.0, 0.0),
                            BoundedQuadratic(0.0, 0.0, 0.0, 0.0, 1.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(0.0, 0.0, 0.0, 0.0, 0.0)])
    @test run_test(f, h)
end

@testset "Test case 1" begin
    f = PiecewiseQuadratic([BoundedQuadratic(0.0, 1.0, 0.0, 0.0, 0.0),
                            BoundedQuadratic(1.0, 2.0, 0.0, 1.0, -1.0),
                            BoundedQuadratic(2.0, Inf, 1.0, -4.0, 5.0)])
    z1 = 0.8284271247461898
    z2 = 2.414213562373095
    h = PiecewiseQuadratic([BoundedQuadratic(0.0, 1.0, 0.0, 0.0, 0.0),
                            BoundedQuadratic(1.0, z2, 0.0, z1, -z1),
                            BoundedQuadratic(z2, Inf, 1.0, -4.0, 5.0)])

    @test run_test(f, h)
end

@testset "Test case 2" begin
    f = PiecewiseQuadratic([BoundedQuadratic(-2.0, -1.0, 0.0, 1.0, -1.0),
                            BoundedQuadratic(-1.0, 0.0, 0.0, 2.0, 0.0),
                            BoundedQuadratic(0.0, Inf, 0.0, 0.0, 0.0)])
    h = PiecewiseQuadratic([BoundedQuadratic(-2.0, Inf, 0.0, 0.0, -3.0)])

    @test run_test(f, h)
end

@testset "Test case 3" begin
    f = PiecewiseQuadratic([BoundedQuadratic(0.00000000000000000000, 0.02157170544841656434,
                                             11.64694718222335900748,
                                             -0.04528496679275312298,
                                             0.00031129231510998337),
                            BoundedQuadratic(0.02157170544841656434, 0.02157170544841656434,
                                             11.64694718222335900748,
                                             -0.04528496679275312298,
                                             0.00031129231510998337),
                            BoundedQuadratic(0.02157170544841656434, Inf,
                                             11.64694718222335900748,
                                             -0.04528497091670485170,
                                             0.00031129240407065537)])
    h = deepcopy(f)
    @test run_test(simplify(f), simplify(h))
end


function calcsort!(A,values::UnitRange{Int})
    store = zeros(Int, size(values))
    for i in 1:length(A)
        for j in 1:length(values)
            if A[i]==values[j]
                store[j]+=1
            end
        end
    end
    k=1
    for i in 1:length(values)
        while store[i]>0
            A[k]=values[i]
            k+=1
            store[i]-=1
        end
    end
    return A
end

function calcsort!(A,values::Vector{Int})
    store = zeros(Int, size(values))
    for i in 1:length(A)
        for j in 1:length(values)
            if A[i]==values[j]
                store[j]+=1
            end
        end
    end
    k=1
    for i in 1:length(values)
        while store[i]>0
            A[k]=values[i]
            k+=1
            store[i]-=1
        end
    end
    return A
end


# Showcasing the available analysis tools for the Bayes (Junction) tree.
# using Revise

using IncrementalInference
using DistributedFactorGraphs # For `isSolvable` function.
using Combinatorics # For creating the variable ordering `permutations`.
using SuiteSparse.CHOLMOD: SuiteSparse_long # For CCOLAMD constraints.
using Gadfly # For histogram and scatter plots.
Gadfly.set_default_plot_size(35cm,25cm)

latex_fonts = Theme(major_label_font="CMU Serif", major_label_font_size=16pt,
                    minor_label_font="CMU Serif", minor_label_font_size=14pt,
                    key_title_font="CMU Serif", key_title_font_size=12pt,
                    key_label_font="CMU Serif", key_label_font_size=10pt)
Gadfly.push_theme(latex_fonts)

# Get tree for each variable ordering in a factor graph.
fg = generateGraph_Kaess(graphinit=false)
all_trees = getAllTrees(deepcopy(fg))

# scores stores: (tree key ID, nnz, cost fxn 1, cost fxn 2).
unsorted_scores = Vector{Tuple{Int, Float64, Float64, Float64}}()
for key in keys(all_trees)
    e = all_trees[key] # (Bayes tree, var order, nnz
    tree = e[1] # Get the Bayes tree.
    cost1 = getTreeCost_01(tree)
    cost2 = getTreeCost_02(tree)
    push!(unsorted_scores, (key, e[3], cost1, cost2))
end

# Sort them to make sure the keys are in order.
scores = sort(unsorted_scores)

# Separate scores into vectors for plotting.
all_nnzs = (x->(x[2])).(scores)
costs_01 = (x->(x[3])).(scores)
costs_02 = (x->(x[4])).(scores)

min_ids_02 = findall(x->x == minimum(costs_02), costs_02)
max_ids_02 = findall(x->x == maximum(costs_02), costs_02)

min_ids_nnz = findall(x->x == minimum(all_nnzs), all_nnzs)
max_ids_nnz = findall(x->x == maximum(all_nnzs), all_nnzs)

# Find the intersection between best on both rubrics (lower left quadrant).
best_ids = findall(x->x in min_ids_02, min_ids_nnz)
# Find good factorizations but bad trees (upper left quadrant).
bad_trees_good_mats_ids = findall(x->x in max_ids_02, min_ids_nnz)
# Find good trees with bad matrix factorizations (lower right quadrant).
good_trees_bad_mats_ids = min_ids_02[findall(x->x == maximum(all_nnzs[min_ids_02]), all_nnzs[min_ids_02])]

# Get AMDs variable ordering.
amd_ordering = getEliminationOrder(fg)
amd_tree = buildTreeReset!(deepcopy(fg), amd_ordering)
amd_tree_nnz = nnzTree(amd_tree)
amd_tree_cost02 = getTreeCost_02(amd_tree)

# Get CCOLAMD variable ordering. First bring in CCOLAMD.
include(normpath(Base.find_package("IncrementalInference"), "..", "ccolamd.jl"))
A, varsym, fctsym = getBiadjacencyMatrix(fg)
colamd_ordering = varsym[Ccolamd.ccolamd(A)]
colamd_tree = buildTreeReset!(deepcopy(fg), colamd_ordering)
colamd_tree_nnz = nnzTree(colamd_tree)
colamd_tree_cost02 = getTreeCost_02(colamd_tree)

# Now add the iSAM2 constraint.
cons = zeros(SuiteSparse_long, length(A.colptr) - 1)
cons[findall(x->x == :x3, varsym)[1]] = 1 # NOTE(tonioteran) hardcoded for Kaess' example.
ccolamd_ordering = varsym[Ccolamd.ccolamd(A, cons)]
ccolamd_tree = buildTreeReset!(deepcopy(fg), ccolamd_ordering)
ccolamd_tree_nnz = nnzTree(ccolamd_tree)
ccolamd_tree_cost02 = getTreeCost_02(ccolamd_tree)

# Plot data points and underlying histogram.
bincnt = 20
layers = []
push!(layers, Gadfly.layer(x=[amd_tree_nnz],
                           y=[amd_tree_cost02],
                           Theme(default_color=colorant"green")))
push!(layers, Gadfly.layer(x=[colamd_tree_nnz],
                           y=[colamd_tree_cost02],
                           Theme(default_color=colorant"blue")))
push!(layers, Gadfly.layer(x=[ccolamd_tree_nnz],
                           y=[ccolamd_tree_cost02],
                           Theme(default_color=colorant"red")))
push!(layers, Gadfly.layer(x=all_nnzs,
                           y=costs_02,
                           Geom.hexbin(xbincount=bincnt, ybincount=bincnt)))
pl = Gadfly.plot(layers...,
            Guide.xlabel("Number of non zeros [int]"),
            Guide.ylabel("Tree cost [cfxn2]"),
            Guide.manual_color_key("",
             ["AMD", "COLAMD", "iSAM2"],
             ["green", "blue", "red"]))

img = SVG("vo_cost_canon_kaess.svg", 6inch, 6inch)
Gadfly.draw(img, pl)


using Test

function run(nums, target_turn=2020)
    last_said_on_turn = Dict()

    turn = 1
    # initialize
    for n in nums[1:length(nums) - 1]
        last_said_on_turn[n] = turn
        turn += 1
        # println("Turn = ", turn, " Val = ", n)
    end

    prev = last(nums)
    while turn < target_turn
        # compute next number to say
        # If that was the first time the number has been spoken, the current player
        # says 0. (i.e. default to 0)
        next = 0
        if haskey(last_said_on_turn, prev)
            # Otherwise, the number had been spoken before; the current player
            # announces how many turns apart the number is from when it was previously
            # spoken.
            next = turn - last_said_on_turn[prev]
        end

        # speak!
        last_said_on_turn[prev] = turn
        # println("Turn = ", turn, " Val = ", prev)

        prev = next
        turn += 1
    end

    println("Turn = ", turn, " Val = ", prev)
    return prev
end

println("\n[Part 1] Examples:")
@test run([0,3,6], 4) == 0
@test run([0,3,6], 5) == 3
@test run([0,3,6], 6) == 3
@test run([0,3,6], 7) == 1
@test run([0,3,6], 8) == 0
@test run([0,3,6], 9) == 4
@test run([0,3,6], 10) == 0
@test run([0,3,6]) == 436

println("\n[Part 1] Result:")
run([19,0,5,1,10,13])

println("\n[Part 2] Examples:")
@test run([0,3,6], 30000000) == 175594.
# Given 1,3,2, the 30000000th number spoken is 2578.
# Given 2,1,3, the 30000000th number spoken is 3544142.
# Given 1,2,3, the 30000000th number spoken is 261214.
# Given 2,3,1, the 30000000th number spoken is 6895259.
# Given 3,2,1, the 30000000th number spoken is 18.
# Given 3,1,2, the 30000000th number spoken is 362.

println("\n[Part 2] Result:")
run([19,0,5,1,10,13], 30000000)


#=
Challenge #256 [Easy] Oblique and De-Oblique

The oblique function slices a matrix (2d array) into diagonals.
The de-oblique function takes diagonals of a matrix, and reassembles the original
rectangular one.

input for oblique

 0  1  2  3  4  5
 6  7  8  9 10 11
12 13 14 15 16 17
18 19 20 21 22 23
24 25 26 27 28 29
30 31 32 33 34 35
(and the output to de-oblique)

output for oblique

0
1 6
2 7 12
3 8 13 18
4 9 14 19 24
5 10 15 20 25 30
11 16 21 26 31
17 22 27 32
23 28 33
29 34
35
(and the input to de-oblique)

bonus deambiguated de-oblique matrices

There's only one de-oblique solution for a square matrix, but when the result is
not square, another input is needed to indicate whether the output should be tall
or wide or provide specific dimentsions of output:

rectangular oblique data input

0
1 6
2 7 12
3 8 13
4 9 14
5 10 15
11 16
17

output for (wide) deoblique (3 6, INPUT) or deoblique (WIDE, INPUT)

 0  1  2  3  4  5
 6  7  8  9 10 11
12 13 14 15 16 17

output for (tall) deoblique (6 3, INPUT) or deoblique (TALL, INPUT)

 0  1  2
 6  7  3
12  8  4
13  9  5
14 10 11
15 16 17

Note

The main use of these functions in computer science is to operate on the diagonals
of a matrix, and then revert it back to a rectangular form. Usually the rectangular
dimensions are known.

(source: https://redd.it/48a4pu)
=#

#   exit at EOF
read_matrix( data::DataType=Any, input::IO=STDIN, spacing::Char=' ' ) = readdlm(
input, spacing, data, header=false, skipblanks=true, ignore_invalid_chars=true,
quotes=false, comments=false )

function de_oblique( grid::Matrix{ Int } )
end

function oblique( grid::Matrix{ Int } )
    local len::Integer = size( grid, 1 )
    local line::Array{ Int, 1 } = Array{ Int, 1 }( )
    local tmp::Integer = 0
    local new_grid::Array{ Array{ Int, 1 }, 1 } = Array{ Array{ Int, 1 }, 1 }( )

    for i in 1:len
        tmp = i
        for j in 1: i
            push!( line, grid[ j , tmp ] )
            tmp -= 1
        end
        push!( new_grid, line )
        line = Array{ Int, 1 }( )
    end

    for i in 2:len
        tmp = i
        for j in len:-1:i
            push!( line, grid[ tmp, j ] )
            tmp += 1
        end
        push!( new_grid, line )
        line = Array{ Int, 1 }( )
    end

    return new_grid
end

function main( )
    local grid::Matrix{ Int } = read_matrix( Int )
    local new_grid::Array{ Array{ Int, 1 }, 1 } = Array{ Array{ Int, 1 }, 1 }( )
    local len::Integer = wid::Integer = 0
    ( len, wid ) = size( grid )

    if len != wid
        new_grid = de_oblique( grid )
    else
        new_grid = oblique( grid )
    end

    for i in 1:size( new_grid, 1 )
        println( new_grid[ i ] )
    end
end

main( )


# Code written entirely by Eric Alcaide
# 8-puzzle example for SearchMethods.jl

using SearchMethods

# The goal is to order the 8-puzzle from an unordered state.
# 0 represents a blank space.
#  |-----------------|
#  |  0  |  1  |  2  |
#  |-----------------|
#  |  3  |  4  |  5  |
#  |-----------------|
#  |  6  |  7  |  8  |
#  |-----------------|

# Formalize the initial game state as an array
state = [3,1,2,4,0,5,6,7,8]

# Generate new nodes from a parent node.
function expand_node(node, tree)
	neighbors = []
	# Locate 0 -> determine moves
	loc = findin(node, 0)[1]
	# Check if moves are possible (up, down, left, right) + expand
	# expand up node
	if loc > 3
		up = copy(node)
		up[loc], up[loc-3] = up[loc-3], up[loc]
		push!(neighbors, up)
	end
	# expand down node
	if loc < 7
		down = copy(node)
		down[loc], down[loc+3] = down[loc+3], down[loc]
		push!(neighbors, down)
	end
	# expand left node
	if loc % 3 != 1
		left = copy(node)
		left[loc], left[loc-1] = left[loc-1], left[loc]
		push!(neighbors, left)
	end
	# expand right node
	if loc % 3 != 0
		right = copy(node)
		right[loc], right[loc+1] = right[loc+1], right[loc]
		push!(neighbors, right)
	end

	return neighbors
end

# Check if the puzzle is ordered
function goal_check(node)
	if node == [0,1,2,3,4,5,6,7,8] return true
	else return false
	end
end

# TESTS FOR UNINFORMED SEARCH ALGORITHMS
println("UNINFORMED METHODS")
# Sample test for the Iterative Deepening Search algorithm.
println("IDS Example")
@time println(SearchMethods.ids(state, false, goal_check, expand_node, 2))
# Sample test for the Breadth Frirst Search algorithm.
println("BFS Example")
@time println(SearchMethods.bfs(state, false, goal_check, expand_node))
# Sample test for the Depth Limited Search algorithm.
println("DLS Example")
@time println(SearchMethods.dls(state, false, goal_check, expand_node, 3))


module DemoLoopVectorization

export comparemethods
using LoopVectorization
using BenchmarkTools

#Entry point when called as an executable
function julia_main()::Cint
    try
        comparemethods()
        return (0)
    catch e
        @error "$e"
        return (1)
    end
end

"""
    threemeans(x::Array{Float64,2})
Returns a tuple of column means of x, column means of positive part of x,
column means of negative part of x.
"""
function threemeans(x::Array{Float64,2})
    nr, nc = size(x)
    epes = Vector{Float64}(undef, nc)
    enes = Vector{Float64}(undef, nc)
    ees = Vector{Float64}(undef, nc)
    @avx for j = 1:nc
        epe = 0.0
        ene = 0.0
        ee = 0.0
        for i = 1:nr
            ee += x[i, j]
            epe += max(0.0, x[i, j])
            ene += min(0.0, x[i, j])
        end
        epes[j] = epe / nr
        enes[j] = ene / nr
        ees[j] = ee / nr
    end
    epes, enes, ees
end

function threemeans_noavx(x::Array{Float64,2})
    nr, nc = size(x)
    epes = Vector{Float64}(undef, nc)
    enes = Vector{Float64}(undef, nc)
    ees = Vector{Float64}(undef, nc)
    for j = 1:nc
        epe = ene = ee = 0.0
        for i = 1:nr
            ee += x[i, j]
            if x[i, j] > 0
                epe += x[i, j]
            else
                ene += x[i, j]
            end
        end
        epes[j] = epe / nr
        enes[j] = ene / nr
        ees[j] = ee / nr
    end
    epes, enes, ees
end

function threemeans_naive(x::Array{Float64,2})
    epe = vec(mean(max.(x, 0.0), dims = 1))
    ene = vec(mean(min.(x, 0.0), dims = 1))
    ee = vec(mean(x, dims = 1))
    epe, ene, ee
end

#= 13 Dec 2021
julia> comparemethods()
threemeans  475.800 μs (3 allocations: 23.81 KiB)
threemeans_noavx  1.537 ms (3 allocations: 23.81 KiB)
threemeans_naive  6.343 ms (31 allocations: 15.28 MiB)
res1[1] ≈ res2[1] && (res1[2] ≈ res2[2] && res1[3] ≈ res2[3]) = true
res1[1] ≈ res3[1] && (res1[2] ≈ res3[2] && res1[3] ≈ res3[3]) = true
=#
function comparemethods()
    x = rand(1000, 1000)
    print("threemeans")
    res1 = @btime(threemeans($x))
    print("threemeans_noavx")
    res2 = @btime(threemeans_noavx($x))
    print("threemeans_naive")
    res3 = @btime(threemeans_naive($x))

    @show res1[1] ≈ res2[1] && res1[2] ≈ res2[2] && res1[3] ≈ res2[3]
    @show res1[1] ≈ res3[1] && res1[2] ≈ res3[2] && res1[3] ≈ res3[3]
    nothing
end

end #module


# Run all unit tests
using LinearLeastSquares
using Base.Test

include("helpers.jl")
include("add_subtract.jl")
include("multiply_divide.jl")
include("stack.jl")
include("diag.jl")
include("repmat.jl")
include("index.jl")
include("reshape.jl")
include("transpose.jl")
include("sum_squares.jl")
include("compound.jl")
include("constraints.jl")
include("problems.jl")


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Julia 1.4.0
#     language: julia
#     name: julia-1.4
# ---

# # 情報幾何学のお勉強のぉと

# ## 導入
#
# $p$,$q$ を正の実数とする. $y=px^2$, $x=qy^2$ という $xy$ - 平面上の曲線を考える．この二つの曲線は原点ともう一つの点 $P=(x_0,y_0)$ と交わる．簡単な計算によってそれは $p$, $q$ を用いると $P = (x_0,y_0) = (p^{-2/3}q^{-1/3}, p^{-1/3}q^{-2/3})$ で与えられることがわかる．これをプログラミング言語 Julia を用いて確認する．

using Plots

using Interact

function vis_crosspoint(p::Float64,q::Float64)
    x = range(0,3,length=30) |> collect
    y = range(0,3,length=30) |> collect
    plot!(x,p*x.^2,color=:blue)
    plot!(q*y.^2,y,color=:orange)
    scatter!([p^(-2/3)*q^(-1/3)],[p^(-1/3)*q^(-2/3)])
end

# ## 試しに描画

plot(xlim=(0.,3.), ylim=(0.,3.), legend=:false)
vis_crosspoint(1.,2.4)

# ## Interact.jl で $p$ と $q$ を変動させた様子を描画

# +
p_slider = slider(0.001:0.1:3,label="p")
q_slider = slider(0.001:0.1:3,label="q")

plot(xlim=(0.,3.), ylim=(0.,3.), legend=:false)
cross_pt = map(
    (p,q)->vis_crosspoint(p,q),
    map(observe,[p_slider,q_slider])...,
)

[p_slider,q_slider,cross_pt] .|> display
# -
# ## $q$ を固定して $p$ だけ動かす

# +
plot(xlim=(0.,3.), ylim=(0.,3.), legend=:false)
anim = @animate for p in range(0.1,1,length=10)
    p=vis_crosspoint(p,1.4)
    plot!(p)
end

gif(anim, "/tmp/anim_fps15.gif", fps = 3)


#module DataSet

# using AbstractFFTs
using FFTW
using Statistics

export Data1D,ind2pos,pos2ind,val,ind,cut
export FourierTransform,PhaseCorrect,BaseLineCorrect
export integral,derivative,set!
export plot,plot!
export length,shift,hard_shift
export interp
export medianBaseline

ENV["GKS_ENCODING"]="utf-8"
import Plots

import Base.iterate

#import NMR.SimplePlot

# Data structures for storing digitised data sets
# Definitions:
#
# - Index: the horizontal index of a data point
# - Position: the numerical position (integer) of the data point in
#           the internal storage array

mutable struct Data1D{Tdata,Tindex}
   dat::Array{Tdata,1}
   istart::Tindex
   istop::Tindex
end

Base.iterate(d::NMR.Data1D,state=1) =
    state > length(d) ? nothing :
      ((d.istart+state*(d.istop-d.istart)/(length(d)-1),d.dat[state]),state+1)

function ind2pos(d::Data1D,ind)
    return round(Int64,(ind-d.istart)*(length(d.dat)-1)/(d.istop-d.istart))+1
end

function pos2ind(d::Data1D,pos::Integer)
   return ind(d)[pos]
end

function ind(d::Data1D)
	return LinRange(d.istart,d.istop,length(d.dat))
end

function val(d::Data1D)
    return d.dat
end

function val(d::Data1D,ind)
    i1=(ind-d.istart)*length(d.dat)/(d.istop-d.istart)
    p1=floor(Int64,i1)+1
    ic=i1-p1+1
    return ((1-ic)*d.dat[p1]+ic*d.dat[p1+1])
end

function cut(d::Data1D,i1,i2)
    p1=ind2pos(d,i1)
    p2=ind2pos(d,i2)
    return Data1D(d.dat[p1:p2],pos2ind(d,p1),pos2ind(d,p2))
end

resolution(d::Data1D) = (d.istop-d.istart)/length(d.dat)

"""
`FourierTransform(d::Data1D)` performs an FFT assuming that the
input data set `d` is a free induction decay, and returns
a `DataSet` object with the resulting complex data.
"""
function FourierTransform(fid::Data1D;
        CTR=0.0,
        LB=0.0,
        SI=length(fid.dat),
        PPM=1,FFTSHIFT=true)

    apod=[exp(-t*LB) for t=ind(fid)];
    s=zeros(Complex{Float64},SI);
    s[1:length(apod)]=apod.*fid.dat;
    s[1]*=0.5;
    SW=length(fid.dat)/(fid.istop-fid.istart);
    spect=fft(s);
    if(FFTSHIFT)
        spect=fftshift(spect)
    end;
    return Data1D(spect,CTR-SW/2/PPM,CTR+SW/2/PPM)
end

"""
`PhaseCorrect(spect::Data1D;Ph0=0.0,Ph1=0.0,Pivot=0.0)`
returns a phase corrected `DataSet`, using the values indicated.
The parameters are given in rad, k, and rad/k, where k are the units
of the spectral domain (typically, Hz or ppm).
"""
function PhaseCorrect(spect::Data1D;Ph0=0.0,Ph1=0.0,Pivot=0.0)
    out=deepcopy(spect);
    out.dat = spect.dat .* [exp(im*(Ph0+(k-Pivot)*Ph1)) for k=ind(spect)]
    return out;
end

polyfit(x::Vector, y::Vector, deg::Int) = pinv(collect(v ^ p for v in x, p in 0:deg)) * y

function horner(x::Vector, c::Vector)
	r=ones(length(x))*c[end]
    for k=(length(c)-1):-1:1
        r.*=x
        r.+=c[k]
    end
    return r
end

@doc raw"""
## Automatic Baseline Correction
`BaseLineCorrect(spect::Data1D;regions=128,kfactor=5,wdw=32)` corrects the
base line of `spect` using the algorithm
by S. Golotvin, A. Williams, J. Magn. Reson. 146 (2000) 122-125.

### Optional Parameters:
- `regions`: the number of regions the spectrum is divided into to estimate
   the noise variance. This should be chosen large enough so that at least one
   of these regions is free from signals, and shows a flat baseline. The
   variance ``\\sigma^2_{min}`` in this region is then used to distinguish signals and noise.

- `wdw`: the window size around each data point that is used to decide whether the
  data point is part of a signal, or should be counted as base line.

- `kfactor`: a point will be counted as part of the baseline if and only if the local
  variance in its window is less than `kfactor`*``\\sigma^2_{min}``.

- `output`: defaults to `:spectrum`. `output=:points` will return an array of points,
  which indicate baseline positions (locations free from peaks) as determined
  by the algorithm. This is useful in situations where a large number of similar
  spectra need to be processed.

### Notes
- if the input data is complex, the imaginary part is ignored,
  and the result is always real.

- the first and last point of `spect` are always used as part of the baseline.
  The first and last points of the output are therefore guaranteed to be zero. This
  helps to keep the interpolation numerically stable.

## Baseline Correction With Known Locations

`BaseLineCorrect(spect::Data1D,pts::Array{Float64,1})`
uses the points in `pts` to determine the baseline. This should be a list
of positions in the spectrum that are known not to contain signal; it can be obtained
by using the option `output=:points` in the first method of `BaseLineCorrect()`.

"""
function BaseLineCorrect(spect::Data1D;regions=128,kfactor=5,lw=0.0,wdw=32,order=5,verbose=false,output=:spectrum)

    out=Data1D(real.(spect.dat),spect.istart,spect.istop);
    n=length(out.dat)
    lreg=floor(Int64,n/regions)

    # compute wdw from lw parameter if set
    if lw>0.0
        wdw = round(Int32,4*lw/resolution(spect))
    end

    # find smallest standard deviation in regions
    mindev=minimum(abs.([std(out.dat[k:k+lreg]) for k=1:lreg:(n-lreg)]))

    # find base line points
    select=[abs(std(out.dat[(k-wdw):(k+wdw)])) < kfactor*mindev for k=(wdw+1):(n-wdw)]

    pts=((wdw+1):(n-wdw))[select];
    verbose && println("Fitting $(length(pts)) Points.")
    push!(pts,1);
    push!(pts,n);

    # polynomial fitting
    coeffs=polyfit(pts/n,out.dat[pts],order)
    verbose && println("Coefficients: $coeffs")

    res=sum(abs.(out.dat[pts]-horner(pts/n,coeffs)))
    verbose && println("Residual: $res")

    out.dat .-= horner(collect(1:n)/n,coeffs)

    if output==:points
        return([pos2ind(spect,p) for p in pts])
    else
        return (out)
    end
end

function BaseLineCorrect(spect::Data1D,ind::Array{Float64,1};order=5)
    out=Data1D(real.(spect.dat),spect.istart,spect.istop);
    n=length(out.dat)
    pts=[ind2pos(out,k) for k in ind]
    coeffs=polyfit(pts/n,out.dat[pts],order)
    out.dat .-= horner(collect(1:n)/n,coeffs)
    return(out)
end

wrap(n,l)=[mod(k,l) for k in n]

function extrema(X::Array{Float64,1})
    Y=Array{Float64,1}()
    for k=2:(length(X)-1)
        if (X[k]>X[k-1] && X[k]>X[k+1]) || X[k]<X[k-1] && X[k]<X[k+1]
            push!(Y,X[k])
        end
    end
    return Y
end

@doc raw"""
    function medianBaseline(s::NMR.Data1D;wdw=length(s)>>6)

Compute baseline for the real part of `s` by the algorithm of M. S. Friedrichs,
*Journal of Biomolecular NMR*,  **5** (1995) 147  153.
The window is automatically calculated as ``\Delta\omega/2^6``, unless
overridden by giving a value to the optional key `wdw`.

"""
function medianBaseline(s::NMR.Data1D{Tdata,Tindex};wdw=length(s)>>6) where { Tdata<:Real, Tindex}
    r=s.dat
    L=length(r)
    b=zeros(L)
    c=zeros(L)
    for k=1:L
        b[k]=median(extrema(r[ wrap(k.+(-wdw:wdw),1:L)]))
    end
    gauss=exp.(-((wdw:wdw)./wdw).^2)
    gauss=gauss/sum(gauss)
    for k=1:length(r)
        c[k]=sum(gauss.*b[wrap(k.+(-wdw:wdw),1:L)])/(2*wdw+1)
    end
    return NMR.Data1D(c,s.istart,s.istop)
end

function medianBaseline(s;opts...)
    return medianBaseline(real(s))+im*medianBaseline(imag(s);opts...)
end

"""
`integral(spect::Data1D)` computes the numerical integral of `spect` by
the trapezoid rule.
"""
function integral(spect::Data1D)
    s=sum(spect.dat[2:(end-1)])+0.5*spect.dat[1]+0.5*spect.dat[end]
    return s*(spect.istop-spect.istart)/length(spect.dat)
end

"""
`derivative(spect::Data1D)` computes the derivative of `spect` and returns
it as another Data1D object
"""
function derivative(spect::Data1D)
    s=spect.dat
    inc=(spect.istop-spect.istart)/length(spect.dat)
#    d[1:(end-1)]=(spect.dat[2:end]-spect.dat[1:(end-1)] )/inc
#    d[end]=d[end-1]
    d = 1.0/12*(8*[s[2:end];0]-8*[0;s[1:(end-1)]] + [s[3:end];0;0] - [0;0;s[1:(end-2)]] )/inc
    return Data1D(d,spect.istart,spect.istop)
end

function Plot(s::Data1D;opts...)
    return Plot(ind(s),real(val(s)),style=Dict(["stroke-width"=>"1"]),Reverse=[true,false];opts...)
end

function set!(spect::Data1D,start,stop,value=0)
    rge=ind2pos(spect,start):ind2pos(spect,stop)
    spect.dat[rge]=value
end

import Base.*, Base./, Base.+,Base.-,Base.abs

function *(d::Data1D,n::Number)
    c = n * d.dat;
    return Data1D(c,d.istart,d.istop)
end

*(n::Number,d::Data1D)=d*n

function *(d1::Data1D,d2::Data1D)
	ind(d1) == ind(d2) || error("Data1D objects incompatible for multiplication")
	return Data1D(d1.dat.*d2.dat,d1.istart,d1.istop)
end

function /(d::Data1D,n::Number)
    return Data1D(d.dat./n,d.istart,d.istop)
end

function /(d1::Data1D,d2::Data1D)
	ind(d1) == ind(d2) || error("Data1D objects incompatible for division")
	return Data1D(d1.dat./d2.dat,d1.istart,d1.istop)
end

function +(d::Data1D,n::Number)
    return Data1D(n.+d.dat,d.istart,d.istop)
end

+(n::Number,d::Data1D)=d+n

function +(d1::Data1D,d2::Data1D)
	ind(d1) == ind(d2) || error("Data1D objects incompatible for addition")
	return Data1D(d1.dat.+d2.dat,d1.istart,d1.istop)
end

function -(d::Data1D,n::Number)
    return Data1D(-n.+d.dat,d.istart,d.istop)
end

function -(d1::Data1D,d2::Data1D)
	ind(d1) == ind(d2) || error("Data1D objects incompatible for subtraction")
	return Data1D(d1.dat.-d2.dat,d1.istart,d1.istop)
end

function abs(d::Data1D)
    c=Data1D(abs.(d.dat),d.istart,d.istop)
    return c
end

function plot(d::Data1D;opts...)
	return Plots.plot(real.(ind(d)),real.(val(d));opts...)
end

import Plots.RecipesBase.plot!

function plot!(d::Data1D;opts...)
	return Plots.plot!(real.(ind(d)),real.(val(d));opts...)
end

import Base.length

function length(d::Data1D)
	return length(d.dat)
end

function shift(d::Data1D,δ::Number)
	c=deepcopy(d);
	c.istart += δ
	c.istop += δ
	return c
end

"""
`function hard_shift(d::Data1D,δ::Number)`

shifts the data in `d` by `δ`. The data is re-interpolated as needed. The
resulting `Data1D` object is guaranteed to have the same index as `d`, such that
they are compatible for the purposes of addition etc. Data shifted out of the
window is lost; data points shifted into view are set to zero.
"""
function hard_shift(d::Data1D,δ::Number)
    # compute integer offset and fractional contribution from left
    inc=(d.istop-d.istart)/(length(d.dat)-1)
    n = convert(Integer,div(δ,inc))
    f = mod(δ,inc)/inc
    ndata=zero(d.dat)
    nstart=d.istart
    nstop=d.istop

    if n<0
        ndata[1:(end+n-1)]=(1-f)*d.dat[(1-n):(end-1)]+f*d.dat[(2-n):end]
    else
        ndata[n+1:end-1]=(1-f)*d.dat[1:end-n-1]+f*d.dat[2:end-n]
    end

    return Data1D(ndata,nstart,nstop)
end

import Base.real, Base.imag, Base.abs, Base.angle

function Base.real(s::NMR.Data1D)
    return(NMR.Data1D(real.(s.dat),s.istart,s.istop))
end

# function Base.abs(s::NMR.Data1D)
#     return(NMR.Data1D(abs.(s.dat),s.istart,s.istop))
# end

function Base.imag(s::NMR.Data1D)
    return(NMR.Data1D(imag.(s.dat),s.istart,s.istop))
end

function Base.angle(s::NMR.Data1D)
    return(NMR.Data1D(angle.(s.dat),s.istart,s.istop))
end

#end


module WignerSymbols
export δ, Δ, clebschgordan, wigner3j, wigner6j, racahV, racahW, HalfInteger

using HalfIntegers
using RationalRoots
using LRUCache
const RRBig = RationalRoot{BigInt}
import RationalRoots: _convert

include("growinglist.jl")
include("bigint.jl") # additional GMP BigInt functionality not wrapped in Base.GMP.MPZ
include("primefactorization.jl")
convert(BigInt, primefactorial(401)) # trigger compilation and generate some fixed data

const Key3j = Tuple{UInt,UInt,UInt,Int,Int}
const Key6j = NTuple{6,UInt}

const Wigner3j = LRU{Key3j,Tuple{Rational{BigInt},Rational{BigInt}}}(; maxsize = 10^6)
const Wigner6j = LRU{Key6j,Tuple{Rational{BigInt},Rational{BigInt}}}(; maxsize = 10^6)

function set_buffer3j_size!(; maxsize)
    resize!(Wigner3j; maxsize = maxsize)
end
function set_buffer6j_size!(; maxsize)
    resize!(Wigner6j; maxsize = maxsize)
end

# check integerness and correctness of (j,m) angular momentum
ϵ(j, m) = (abs(m) <= j && ishalfinteger(j) && isinteger(j-m) && isinteger(j+m))

# check triangle condition
"""
    δ(j₁, j₂, j₃) -> ::Bool

Checks the triangle conditions `j₃ <= j₁ + j₂`, `j₁ <= j₂ + j₃` and `j₂ <= j₃ + j₁`.
"""
δ(j₁, j₂, j₃) = (j₃ <= j₁ + j₂) && (j₁ <= j₂ + j₃) && (j₂ <= j₃ + j₁) && isinteger(j₁+j₂+j₃)

# triangle coefficient
"""
    Δ(T::Type{<:Real} = RationalRoot{BigInt}, j₁, j₂, j₃) -> ::T

Computes the triangle coefficient
`Δ(j₁, j₂, j₃) = √((j₁+j₂-j₃)!*(j₁-j₂+j₃)!*(j₂+j₃-j₁)! / (j₁+j₂+j₃+1)!)`
as a type `T` real number.

Returns `zero(T)` if the triangle condition `δ(j₁, j₂, j₃)` is not satisfied, but
throws a `DomainError` if the `jᵢ`s are not (half)integer
"""
Δ(j₁, j₂, j₃) = Δ(RRBig, j₁, j₂, j₃)
function Δ(T::Type{<:Real}, j₁, j₂, j₃)
    for jᵢ in (j₁, j₂, j₃)
        (ishalfinteger(jᵢ) && jᵢ >= zero(jᵢ)) || throw(DomainError("invalid jᵢ", jᵢ))
    end
    if !δ(j₁, j₂, j₃)
        return zero(T)
    end
    n, d = Δ²(j₁, j₂, j₃)
    r = Base.unsafe_rational(n, d)
    return convert(T, signedroot(RationalRoot{BigInt}, r))
end

"""
    wigner3j(T::Type{<:Real} = RationalRoot{BigInt}, j₁, j₂, j₃, m₁, m₂, m₃ = -m₂-m₁) -> ::T

Compute the value of the Wigner-3j symbol
    ⎛ j₁  j₂  j₃ ⎞
    ⎝ m₁  m₂  m₃ ⎠
as a type `T` real number. By default, `T = RationalRoot{BigInt}` and `m₃ = -m₁-m₂`.

Returns `zero(T)` if the triangle condition `δ(j₁, j₂, j₃)` is not satisfied, but
throws a `DomainError` if the `jᵢ`s and `mᵢ`s are not (half)integer or `abs(mᵢ) > jᵢ`.
"""
wigner3j(j₁, j₂, j₃, m₁, m₂, m₃ = -m₁-m₂) = wigner3j(RRBig, j₁, j₂, j₃, m₁, m₂, m₃)
function wigner3j(T::Type{<:Real}, j₁, j₂, j₃, m₁, m₂, m₃ = -m₁-m₂)
    # check angular momenta
    for (jᵢ,mᵢ) in ((j₁, m₁), (j₂, m₂), (j₃, m₃))
        ϵ(jᵢ, mᵢ) || throw(DomainError((jᵢ, mᵢ), "invalid combination (jᵢ, mᵢ)"))
    end
    return _wigner3j(T, HalfInteger.((j₁, j₂, j₃, m₁, m₂, m₃))...)
end

function _wigner3j(T::Type{<:Real}, j₁::HalfInteger, j₂::HalfInteger, j₃::HalfInteger,
                                    m₁::HalfInteger, m₂::HalfInteger, m₃::HalfInteger)
    # check triangle condition and m₁+m₂+m₃ == 0
    if !δ(j₁, j₂, j₃) || !iszero(m₁+m₂+m₃)
        return zero(T)
    end

    # we reorder such that j₁ >= j₂ >= j₃ and m₁ >= 0 or m₁ == 0 && m₂ >= 0
    j₁, j₂, j₃, m₁, m₂, m₃, sgn = reorder3j(j₁, j₂, j₃, m₁, m₂, m₃)
    # TODO: do we also want to use Regge symmetries?
    α₁ = convert(Int, j₂ - m₁ - j₃ ) # can be negative
    α₂ = convert(Int, j₁ + m₂ - j₃ ) # can be negative
    β₁ = convert(UInt, j₁ + j₂ - j₃ )
    β₂ = convert(UInt, j₁ - m₁ )
    β₃ = convert(UInt, j₂ + m₂ )

    # extra sign in definition: α₁ - α₂ = j₁ + m₂ - j₂ + m₁ = j₁ - j₂ + m₃
    sgn = isodd(α₁ - α₂) ? -sgn : sgn

    # dictionary lookup or compute
    if haskey(Wigner3j, (β₁, β₂, β₃, α₁, α₂))
        r, s = Wigner3j[(β₁, β₂, β₃, α₁, α₂)]
    else
        s1n, s1d = Δ²(j₁, j₂, j₃)
        s2n = prod(map(primefactorial, (β₂, β₁ - α₁, β₁ - α₂, β₃, β₃ - α₁, β₂ - α₂)))

        snum, rnum = splitsquare(s1n*s2n)
        sden, rden = splitsquare(s1d)
        snum, sden = divgcd!(snum, sden)
        rnum, rden = divgcd!(rnum, rden)
        s = Base.unsafe_rational(convert(BigInt, snum), convert(BigInt, sden))
        r = Base.unsafe_rational(convert(BigInt, rnum), convert(BigInt, rden))
        s *= compute3jseries(β₁, β₂, β₃, α₁, α₂)
        Wigner3j[(β₁, β₂, β₃, α₁, α₂)] = (r,s)
    end
    return _convert(T, sgn*s)*convert(T, signedroot(r))
end

"""
    clebschgordan(T::Type{<:Real} = RationalRoot{BigInt}, j₁, m₁, j₂, m₂, j₃, m₃ = m₁+m₂) -> ::T

Compute the value of the Clebsch-Gordan coefficient <j₁, m₁; j₂, m₂ | j₃, m₃ >
as a type `T` real number. By default, `T = RationalRoot{BigInt}` and `m₃ = m₁+m₂`.

Returns `zero(T)` if the triangle condition `δ(j₁, j₂, j₃)` is not satisfied, but
throws a `DomainError` if the `jᵢ`s and `mᵢ`s are not (half)integer or `abs(mᵢ) > jᵢ`.
"""
clebschgordan(j₁, m₁, j₂, m₂, j₃, m₃ = m₁+m₂) =
    clebschgordan(RRBig, j₁, m₁, j₂, m₂, j₃, m₃)
function clebschgordan(T::Type{<:Real}, j₁, m₁, j₂, m₂, j₃, m₃ = m₁+m₂)
    s = wigner3j(T, j₁, j₂, j₃, m₁, m₂, -m₃)
    iszero(s) && return s
    s *= convert(T, signedroot(RRBig, j₃+j₃+one(j₃)))
    return isodd(convert(Int,j₁ - j₂ + m₃)) ? -s : s
end

"""
    racahV(T::Type{<:Real} = RationalRoot{BigInt}, j₁, j₂, j₃, m₁, m₂, m₃ = -m₁-m₂) -> ::T

Compute the value of Racah's V-symbol
V(j₁, j₂, j₃; m₁, m₂, m₃) = (-1)^(-j₁+j₂+j₃) * ⎛ j₁  j₂  j₃ ⎞
                                               ⎝ m₁  m₂  m₃ ⎠
as a type `T` real number. By default, `T = RationalRoot{BigInt}` and `m₃ = -m₁-m₂`.

Returns `zero(T)` if the triangle condition `δ(j₁, j₂, j₃)` is not satisfied, but
throws a `DomainError` if the `jᵢ`s and `mᵢ`s are not (half)integer or `abs(mᵢ) > jᵢ`.
"""
racahV(j₁, j₂, j₃, m₁, m₂, m₃ = -m₁-m₂) = racahV(RRBig, j₁, j₂, j₃, m₁, m₂, m₃)
function racahV(T::Type{<:Real}, j₁, j₂, j₃, m₁, m₂, m₃ = -m₁-m₂)
    s = wigner3j(T, j₁, j₂, j₃, m₁, m₂, m₃)
    return isodd(convert(Int, -j₁ + j₂ + j₃)) ? -s : s
end

"""
    wigner6j(T::Type{<:Real} = RationalRoot{BigInt}, j₁, j₂, j₃, j₄, j₅, j₆) -> ::T

Compute the value of the Wigner-6j symbol
    _⎧ j₁  j₂  j₃ ⎫_
     ⎩ j₄  j₅  j₆ ⎭
as a type `T` real number. By default, `T = RationalRoot{BigInt}`.

Returns `zero(T)` if any of triangle conditions `δ(j₁, j₂, j₃)`, `δ(j₁, j₆, j₅)`,
`δ(j₂, j₄, j₆)`, `δ(j₃, j₄, j₅)` are not satisfied, but throws a `DomainError` if
the `jᵢ`s are not integer or halfinteger.
"""
wigner6j(j₁, j₂, j₃, j₄, j₅, j₆) = wigner6j(RRBig, j₁, j₂, j₃, j₄, j₅, j₆)
function wigner6j(T::Type{<:Real}, j₁, j₂, j₃, j₄, j₅, j₆)
    # check validity of `jᵢ`s
    for jᵢ in (j₁, j₂, j₃, j₄, j₅, j₆)
        (ishalfinteger(jᵢ) && jᵢ >= zero(jᵢ)) || throw(DomainError("invalid jᵢ", jᵢ))
    end
    return _wigner6j(T, HalfInteger.((j₁, j₂, j₃, j₄, j₅, j₆))...)
end

function _wigner6j(T::Type{<:Real}, j₁::HalfInteger, j₂::HalfInteger, j₃::HalfInteger,
                                    j₄::HalfInteger, j₅::HalfInteger, j₆::HalfInteger)
    α̂₁ = (j₁, j₂, j₃)
    α̂₂ = (j₁, j₆, j₅)
    α̂₃ = (j₂, j₄, j₆)
    α̂₄ = (j₃, j₄, j₅)

    # check triangle conditions
    if !(δ(α̂₁...) && δ(α̂₂...) && δ(α̂₃...) && δ(α̂₄...))
        return zero(T)
    end
    # reduce
    α₁ = convert(UInt, +(α̂₁...))
    α₂ = convert(UInt, +(α̂₂...))
    α₃ = convert(UInt, +(α̂₃...))
    α₄ = convert(UInt, +(α̂₄...))
    β₁ = convert(UInt, j₁+j₂+j₄+j₅)
    β₂ = convert(UInt, j₁+j₃+j₄+j₆)
    β₃ = convert(UInt, j₂+j₃+j₅+j₆)

    # we should have αᵢ < βⱼ, ∀ i, j and ∑ᵢ αᵢ = ∑ⱼ βⱼ
    # now order them as β₁ >= β₂ >= β₃ >= α₁ >= α₂ >= α₃ >= α₄
    (β₁, β₂, β₃, α₁, α₂, α₃, α₄) = reorder6j(β₁, β₂, β₃, α₁, α₂, α₃, α₄)

    # dictionary lookup or compute
    if haskey(Wigner6j, (β₁, β₂, β₃, α₁, α₂, α₃))
        r, s = Wigner6j[(β₁, β₂, β₃, α₁, α₂, α₃)]
    else
        # order irrelevant: product remains the same under action of reorder6j
        n₁, d₁ = Δ²(α̂₁...)
        n₂, d₂ = Δ²(α̂₂...)
        n₃, d₃ = Δ²(α̂₃...)
        n₄, d₄ = Δ²(α̂₄...)

        snum, rnum = splitsquare(n₁ * n₂ * n₃ * n₄)
        sden, rden = splitsquare(d₁ * d₂ * d₃ * d₄)
        snum, sden = divgcd!(snum, sden)
        rnum, rden = divgcd!(rnum, rden)
        s = Base.unsafe_rational(convert(BigInt, snum), convert(BigInt, sden))
        r = Base.unsafe_rational(convert(BigInt, rnum), convert(BigInt, rden))
        s *= compute6jseries(β₁, β₂, β₃, α₁, α₂, α₃, α₄)

        Wigner6j[(β₁, β₂, β₃, α₁, α₂, α₃)] = (r, s)
    end
    return _convert(T, s) * convert(T, signedroot(r))
end

"""
    racahW(T::Type{<:Real} = RationalRoot{BigInt}, j₁, j₂, J, j₃, J₁₂, J₂₃) -> ::T

Compute the value of Racah's W coefficient
`W(j₁, j₂, J, j₃; J₁₂, J₂₃) = <(j₁,(j₂j₃)J₂₃)J | ((j₁j₂)J₁₂,j₃)J> / sqrt((2J₁₂+1)*(2J₁₃+1))`
as a type `T` real number. By default, `T = RationalRoot{BigInt}`.

Returns `zero(T)` if any of triangle conditions `δ(j₁, j₂, J₁₂)`, `δ(j₂, j₃, J₂₃)`,
`δ(j₁, J₂₃, J)`, `δ(J₁₂, j₃, J)` are not satisfied, but throws a `DomainError` if
the `jᵢ`s and `J`s are not integer or halfinteger.
"""
racahW(j₁, j₂, J, j₃, J₁₂, J₂₃) = racahW(RRBig, j₁, j₂, J, j₃, J₁₂, J₂₃)
function racahW(T::Type{<:Real}, j₁, j₂, J, j₃, J₁₂, J₂₃)
    s = wigner6j(T, j₁, j₂, J₁₂, j₃, J, J₂₃)
    if !iszero(s) && isodd(convert(Int, j₁ + j₂ + j₃ + J))
        return -s
    else
        return s
    end
end

# COMPUTATIONAL ROUTINES
#------------------------
# squared triangle coefficient
function Δ²(j₁, j₂, j₃)
    # also checks the triangle conditions by converting to unsigned integer:
    n1 = copy(primefactorial( convert(UInt, + j₁ + j₂ - j₃) ))
    n2 = primefactorial( convert(UInt, + j₁ - j₂ + j₃) )
    n3 = primefactorial( convert(UInt, - j₁ + j₂ + j₃) )
    num = mul!(mul!(n1, n2), n3)
    den = copy(primefactorial( convert(UInt, j₁ + j₂ + j₃ + 1) ))
    # result
    return divgcd!(num, den)
end

# reorder parameters determining the 3j symbol to canonical order:
# j₁ >= j₂ >= j₃ and m₁ >= 0 or m₁ == 0 && m₂ >= 0
function reorder3j(j₁, j₂, j₃, m₁, m₂, m₃, sign = one(Int8))
    if j₁ < j₂
        return reorder3j(j₂, j₁, j₃, m₂, m₁, m₃, -sign)
    elseif j₂ < j₃
        return reorder3j(j₁, j₃, j₂, m₁, m₃, m₂, -sign)
    elseif m₁ < zero(m₁)
        return reorder3j(j₁, j₂, j₃, -m₁, -m₂, -m₃, -sign)
    elseif iszero(m₁) && m₂ < zero(m₂)
        return reorder3j(j₁, j₂, j₃, -m₁, -m₂, -m₃, -sign)
    else
        # sign doesn't matter if total J=j₁ + j₂ + j₃ is even
        if iseven(convert(UInt,j₁ + j₂ + j₃))
            sign = one(sign)
        end
        return (j₁, j₂, j₃, m₁, m₂, m₃, sign)
    end
end

# reorder parameters determining the 6j symbol to canonical order:
# β₁ >= β₂ >= β₃ >= α₁ >= α₂ >= α₃ >= α₄
function reorder6j(β₁, β₂, β₃, α₁, α₂, α₃, α₄)
    if β₁ < β₂
        return reorder6j(β₂, β₁, β₃, α₁, α₂, α₃, α₄)
    elseif β₂ < β₃
        return reorder6j(β₁, β₃, β₂, α₁, α₂, α₃, α₄)
    elseif α₁ < α₂
        return reorder6j(β₁, β₂, β₃, α₂, α₁, α₃, α₄)
    elseif α₂ < α₃
        return reorder6j(β₁, β₂, β₃, α₁, α₃, α₂, α₄)
    elseif α₃ < α₄
        return reorder6j(β₁, β₂, β₃, α₁, α₂, α₄, α₃)
    else
        return (β₁, β₂, β₃, α₁, α₂, α₃, α₄)
    end
end

# compute the sum appearing in the 3j symbol
function compute3jseries(β₁, β₂, β₃, α₁, α₂)
    krange = max(α₁, α₂, zero(α₁)):min(β₁, β₂, β₃)
    T = PrimeFactorization{eltype(eltype(factorialtable))}

    nums = Vector{T}(undef, length(krange))
    dens = Vector{T}(undef, length(krange))
    for (i, k) in enumerate(krange)
        num = iseven(k) ? one(T) : -one(T)
        den = copy(primefactorial(k))
        den = mul!(mul!(den, primefactorial(k-α₁)), primefactorial(k-α₂))
        den = mul!(mul!(mul!(den, primefactorial(β₁-k)),
                                    primefactorial(β₂-k)),
                                        primefactorial(β₃-k))
        nums[i], dens[i] = num, den
    end
    den = commondenominator!(nums, dens)
    totalnum = sumlist!(nums)
    totalden = convert(BigInt, den)
    for n = 1:length(den.powers)
        p = bigprime(n)
        while den.powers[n] > 0
            q, r = divrem(totalnum, p)
            if iszero(r)
                totalnum = q
                den.powers[n] -= 1
            else
                break
            end
        end
    end
    totalden = convert(BigInt, den)
    return Base.unsafe_rational(totalnum, totalden)
end

# compute the sum appearing in the 6j symbol
function compute6jseries(β₁, β₂, β₃, α₁, α₂, α₃, α₄)
    krange = max(α₁, α₂, α₃, α₄):min(β₁, β₂, β₃)
    T = PrimeFactorization{eltype(eltype(factorialtable))}

    nums = Vector{T}(undef, length(krange))
    dens = Vector{T}(undef, length(krange))
    for (i, k) in enumerate(krange)
        num = iseven(k) ? copy(primefactorial(k+1)) : neg!(copy(primefactorial(k+1)))
        den = copy(primefactorial(k-α₁))
        den = mul!(mul!(mul!(den, primefactorial(k-α₂)),
                                    primefactorial(k-α₃)),
                                        primefactorial(k-α₄))
        den = mul!(mul!(mul!(den, primefactorial(β₁-k)),
                                    primefactorial(β₂-k)),
                                        primefactorial(β₃-k))
        nums[i], dens[i] = divgcd!(num, den)
    end
    den = commondenominator!(nums, dens)
    totalnum = sumlist!(nums)
    for n = 1:length(den.powers)
        p = bigprime(n)
        while den.powers[n] > 0
            q, r = divrem(totalnum, p)
            if iszero(r)
                totalnum = q
                den.powers[n] -= 1
            else
                break
            end
        end
    end
    totalden = convert(BigInt, den)
    return Base.unsafe_rational(totalnum, totalden)
end

function _precompile_()
    @assert precompile(wigner3j, (Type{Float64}, Int, Int, Int, Int, Int, Int))
    @assert precompile(wigner6j, (Type{Float64}, Int, Int, Int, Int, Int, Int))
    @assert precompile(wigner3j, (Type{BigFloat}, Rational{Int}, Rational{Int}, Rational{Int}, Rational{Int}, Rational{Int}, Rational{Int}))
    @assert precompile(wigner6j, (Type{BigFloat}, Rational{Int}, Rational{Int}, Rational{Int}, Rational{Int}, Rational{Int}, Rational{Int}))
    @assert precompile(wigner3j, (Type{RRBig}, HalfInt, HalfInt, HalfInt, HalfInt, HalfInt, HalfInt))
    @assert precompile(wigner6j, (Type{RRBig}, HalfInt, HalfInt, HalfInt, HalfInt, HalfInt, HalfInt))
end
_precompile_()

end # module


### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ 120da254-1c28-11ec-1938-a55c8cde5727
using Pkg; Pkg.develop(path=".."); Pkg.activate("..")

# ╔═╡ c99f9fc5-0ab1-444a-80f6-b5241edf2f8f
begin
	using QuantumStateBase
	using Plots
	using LinearAlgebra
end

# ╔═╡ db69bedc-2803-4268-9f56-34e4a88a324e
state = SqueezedThermalState(Complex{Float64}, ξ(0.5, π/4), 0.5, )

# ╔═╡ cd19d00d-bf2b-479d-8c88-493311d97303
d = gaussian_state_sampler(Float64, state, 4096)

# ╔═╡ b9ea40f4-1ed1-4862-8858-be07f2c04633
scatter(d[1, :], d[2, :])

# ╔═╡ Cell order:
# ╟─120da254-1c28-11ec-1938-a55c8cde5727
# ╠═c99f9fc5-0ab1-444a-80f6-b5241edf2f8f
# ╠═db69bedc-2803-4268-9f56-34e4a88a324e
# ╠═cd19d00d-bf2b-479d-8c88-493311d97303
# ╠═b9ea40f4-1ed1-4862-8858-be07f2c04633


using TrajectoryOptimization
using HDF5
include("N_plots.jl")

# Set up the problem
model, obj0 = Dynamics.cartpole_analytical
n,m = model.n, model.m

x0 = zeros(n)
xf = [0.5;pi;0;0]
tf = 2.0
obj = LQRObjective(obj0.cost.Q, obj0.cost.R, obj0.cost.Qf, tf, x0, xf)
u_bnd = 20
x_bnd = [0.6,Inf,Inf,Inf]

obj_c = ConstrainedObjective(obj,u_min=-u_bnd, u_max=u_bnd)
obj_min = update_objective(obj_c,tf=:min)
dt = 0.1

# Params
N = 51
method = :hermite_simpson

# Initial Trajectory
U0 = ones(1,N)
X0 = line_trajectory(obj.x0,obj.xf,N)

X0_rollout = copy(X0)
solver = Solver(model,obj_c,N=N)
rollout!(X0_rollout,U0,solver)

# Dircol functions
# eval_f, eval_g, eval_grad_f, eval_jac_g = gen_usrfun_ipopt(solver,method)

function comparison_plot(results,sol_ipopt; kwargs...)
    n = size(results.X,1)
    if n == 2 # pendulum
        state_labels = ["pos" "vel"]
    else
        state_labels = ["pos" "angle"]
    end
    time_dircol = range(0,stop=obj.tf,length=size(sol_ipopt.X,2))
    time_ilqr = range(0,stop=obj.tf,length=size(results.X,1))

    colors = [:blue :red]

    p_u = plot(time_ilqr[1:end-1],to_array(results.U)',width=2,label="iLQR",color=colors)
    plot!(time_dircol,sol_ipopt.U',width=2, label="IPOPT",color=colors,style=:dashdot,ylabel="control")

    p_x = plot(time_ilqr,to_array(results.X)[1:2,:]',width=2,label=state_labels,color=colors)
    plot!(time_dircol,sol_ipopt.X[1:2,:]',width=2, label="",color=colors,style=:dashdot,ylabel="state")

    p = plot(p_x,p_u,layout=(2,1),xlabel="time (s)"; kwargs...)

    return p
end

function convergence_plot(stat_i,stat_d)
    plot(stat_i["cost"],width=2,label="iLQR")
    plot!(stat_d["cost"], ylim=[0,10], ylabel="Cost",xlabel="iterations",width=2,label="dircol")
end

colors_X = [:red :blue :orange :green]

#####################################
#          UNCONSTRAINED            #
#####################################

# Solver Options
opts = SolverOptions()
opts.cost_tolerance = 1e-6
opts.outer_loop_update_type = :default
opts.constraint_decrease_ratio = 0.75

# iLQR
solver = Solver(model,obj,N=N,opts=opts,integration=:rk3)
solver.opts.verbose = false
res_i, stat_i = solve(solver,U0)
plot(to_array(res_i.X)')
plot(to_array(res_i.U)')
max_violation(res_i)
cost(solver, res_i)
stat_i["runtime"]
stat_i["iterations"]
var_i = DircolVars(to_array(res_i.X),to_array(res_i.U))

# DIRCOL
res_d, stat_d = solve_dircol(solver, X0_rollout, U0; method=method)
plot(res_d.X')
plot(res_d.U')
stat_d["c_max"][end]
stat_d["cost"][end]
stat_d["runtime"]
stat_d["iterations"]

comparison_plot(res_i,res_d)
convergence_plot(stat_i,stat_d)

time_per_iter = stat_i["runtime"]/stat_i["iterations"]
time_per_iter = stat_d["runtime"]/stat_d["iterations"]

#####################################
#           CONSTRAINED             #
#####################################

# Solver Options
opts = SolverOptions()
opts.cost_tolerance = 1e-8
opts.cost_tolerance_intermediate = 1e-2
opts.constraint_tolerance = 1e-8
opts.outer_loop_update_type = :feedback
opts.constraint_decrease_ratio = .25
opts.square_root = true
opts.penalty_initial = 1
opts.penalty_scaling = 50

# iLQR
obj_c2 = update_objective(obj_c,tf=2)
solver = Solver(model,obj_c2,N=N,opts=opts,integration=:rk3)
res_i, stat_i = solve(solver,U0)
stat_i["iterations"]
comparison_plot(res_i,res_d)
plot(to_array(res_i.X)')
plot(to_array(res_i.U)')
stat_i["c_max"][end]
stat_i["runtime"]
constraint_plot(solver,U0)

# DIRCOL
res_d, stat_d = solve_dircol(solver, X0_rollout, U0; method=:hermite_simpson)
plot(res_d.X')
plot(res_d.U')
stat_d["c_max"][end]
stat_d["cost"][end]
stat_d["runtime"]
stat_d["iterations"]

convergence_plot(stat_i,stat_d)

cost(solver,res_i)
cost(solver,res_d)

time_per_iter = stat_i["runtime"]/stat_i["iterations"]
time_per_iter = stat_d["runtime"]/stat_d["iterations"]

#####################################
#        INFEASIBLE START           #
#####################################

# Solver Options
opts = SolverOptions()
opts.cost_tolerance = 1e-6
opts.cost_tolerance_intermediate = 5e-1
opts.constraint_tolerance = 1e-4
opts.outer_loop_update_type = :feedback
opts.constraint_decrease_ratio
opts.resolve_feasible = false
opts.square_root = true  # this makes a big difference!
opts.R_infeasible = 1
opts.penalty_initial = 1
opts.penalty_scaling = 10
opts.use_nesterov = false
opts.al_type = :default

# iLQR
solver = Solver(model,obj_c,N=N,opts=opts,integration=:rk3)
res_i, stat_i = solve(solver,X0,U0)
plot(to_array(res_i.X)')
plot(to_array(res_i.U)')
stat_i["c_max"][end]
stat_i["cost"][end]
stat_i["runtime"]
stat_i["iterations"]
var_i = DircolVars(res_i)

# DIRCOL
res_d, stat_d = solve_dircol(solver, X0, U0; method=:hermite_simpson)
plot(res_d.X')
plot(res_d.U')
stat_d["c_max"][end]
stat_d["cost"][end]
stat_d["runtime"]
stat_d["iterations"]
stat_d["cost"]

comparison_plot(res_i,res_d)
convergence_plot(stat_i,stat_d)

cost(solver,res_i)
cost(solver,res_d)

time_per_iter = stat_i["runtime"]/stat_i["iterations"]
time_per_iter = stat_d["runtime"]/stat_d["iterations"]

#####################################
#          MINIMUM TIME             #
#####################################

X0_warm = to_array(res_i.X)
U0_warm = to_array(res_i.U)

# Solver Options
opts = SolverOptions()
opts.max_dt = 0.7
opts.verbose = false
opts.cost_tolerance = 1e-6
opts.cost_tolerance_intermediate = 1e-5
opts.constraint_tolerance = 0.01
opts.R_minimum_time = 100
opts.penalty_initial_minimum_time_equality = 100.
opts.penalty_scaling_minimum_time_equality = 10

opts.outer_loop_update_type = :default
opts.constraint_decrease_ratio
opts.resolve_feasible = false
opts.square_root = true  # this makes a big difference!
opts.penalty_initial = 1
opts.penalty_scaling = 10
opts.use_nesterov = false
opts.al_type = :default
opts.live_plotting = true
opts.resolve_feasible = false
opts.R_infeasible = 1
opts.minimum_time_tf_estimate = 1.5

# iLQR
solver_mintime = Solver(model,obj_min,N=N,opts=opts,integration=:rk3)
solver_mintime.opts.verbose = true
res_i, stat_i = solve(solver_mintime,X0_warm,U0_warm)
plot(to_array(res_i.X)')
plot(to_array(res_i.U)')
max_violation(res_i)
_cost(solver, res_i)
stat_i["runtime"]
stat_i["iterations"]
total_time(solver, res_i)
plot(stat_i["c_max"],yscale=:log10)

# DIRCOL
solver.opts.verbose = false
res_d, stat_d = solve_dircol(solver, X0_rollout, U0; method=:hermite_simpson)
stat_d["info"]
plot(res_d.X')
plot(res_d.U')
stat_d["c_max"][end]
stat_d["cost"][end]
stat_d["runtime"]
stat_d["iterations"]
stat_d["cost"]
total_time(solver, res_d)

comparison_plot(res_i,res_d)
convergence_plot(stat_i,stat_d)

eval_f(var_i.Z)
eval_f(res_d.Z)

stat_i["cost"][30]
stat_d["cost"][30]

time_per_iter = stat_i["runtime"]/stat_i["iterations"]
time_per_iter = stat_d["runtime"]/stat_d["iterations"]

dt_i = [res_i.U[k][end]^2 for k = 1:N]
dt_d = res_d.U[end,:]

plot(dt_i)
plot!(dt_d)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#
#                                                                              #
#                       KNOT POINT COMPARISONS                                 #
#                                                                              #
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~#

Ns = [21,41,51,81,101,201,401,501,801,1001]
obj.tf ./ (Ns.-1)
dt_truth = 1e-3

#####################################
#          UNCONSTRAINED            #
#####################################
group = "cartpole/unconstrained"
run_step_size_comparison(model, obj, U0, group, Ns, integrations=[:midpoint,:rk3,:rk3_foh,:rk4],dt_truth=1e-3,opts=opts)
plot_stat("iterations",group)

#####################################
#            CONSTRAINED            #
#####################################
Ns = [21,41,51,81,101,201]
group = "cartpole/constrained"
plot_stat("error",group)

#####################################
#            INFEASIBLE             #
#####################################
solver_truth, res_truth,  = run_dircol_truth(model, obj_c, 1e-3, X0, U0, "cartpole/infeasible")
time_truth = get_time(solver_truth)
plot(res_truth.U')

err_mid, eterm_mid, stats_mid = run_Ns(model, obj_c, Ns, :midpoint, infeasible=true)
err_rk3, eterm_rk3, stats_rk3 = run_Ns(model, obj_c, Ns, :rk3, infeasible=true)
err_foh, eterm_foh, stats_foh = run_Ns(model, obj_c, Ns, :rk3_foh, infeasible=true)
err_rk4, eterm_rk4, stats_rk4 = run_Ns(model, obj_c, Ns, :rk4, infeasible=true)

save_data("cartpole/infeasible")


using FeedbackParticleFilters, Distributions, Random, PDMats

println("Testing LinearDiffusionStateModel.jl:")

@testset "LinearDiffusionStateModel.jl" begin

    A      = [0. 1.; -0.7 -0.3]
    B      = zeros(2,1)
    B[2,1] = sqrt(2.)
    init   = MvNormal(LinearAlgebra.Diagonal(FillArrays.Fill(1., 2)))

    print("  inner constructor")
    model = LinearDiffusionStateModel(A, B, init)
    B2 = zeros(Int, 2, 1)
    B2[2,1] = 2
    print("-")
    @test_throws MethodError LinearDiffusionStateModel(A, B2, init) # wrong element type of B2
    B3 = zeros(3, 1)
    print("-")
    @test_throws DimensionMismatch LinearDiffusionStateModel(A, B3, init) # wrong size of B3
    println("DONE")

    print("  method find_stationary_variance")
    Sigma = find_stationary_variance(A, B*B')
    print("-")
    @test Sigma == [4.746796851882334 -5.543824348282954e-5; -5.543824348282954e-5 3.3192089048776765]
    println("DONE")

    print("  method drift")
    print("-")
    @test drift(model) == A
    println("DONE")

    print("  method diffusion")
    print("-")
    @test diffusion(model) == B
    println("DONE")

    print("  method initial_condition")
    print("-")
    @test initial_condition(model) == model.init
    println("DONE")

    print("  method initialize")
    print("-")
    x0 = initialize(model)
    @test length(x0) == 2
    println("DONE")

    print("  method state_dim")
    print("-")
    @test state_dim(model) == 2
    println("DONE")

    print("  method noise_dim")
    print("-")
    @test noise_dim(model) == 1
    println("DONE")

    print("  method drift_function")
    print("-")
    @test drift_function(model)([1., 3.]) == [3.0, -1.5999999999999999]
    println("DONE")

    print("  method diffusion_function")
    print("-")
    @test diffusion_function(model)([1., 3.]) == hcat([0.0, 1.4142135623730951]...)'
    println("DONE")

    print("  calling instance of DiffusionStateModel")
    x0 = ones(2)
    print("-")
    @test model(x0, 0.01) != x0
    xx0 = hcat(ones(2), zeros(2))
    print("-")
    @test model(xx0, 0.01) != xx0
    println("DONE")

    print("  method propagate!")
    x1 = copy(x0)
    propagate!(x1, model, 0.01)
    @test x1 != x0
    print("-")
    xx1 = copy(xx0)
    propagate!(xx1, model, 0.01)
    print("-")
    @test xx1 != xx0
    println("DONE")

end; #LinearDiffusionStateModel.jl


using NMF: nndsvd, solve!, CoordinateDescent

"""
    NMF(;ncomps=nothing)
    NMF(ncomps)

Use non-negative matrix factorization (NMF, NNMF) to form a non-negative, low-rank, and orthonormal basis of the input. The implementation of the underlying NMF is provided by [NMF.jl](https://github.com/JuliaStats/NMF.jl). The implementation uses a non-negative SVD for initialization and a coordinate-descent solver to fit.

If `ncomps` is `nothing`, it will be set to the number of frames in the reference cube when processed.

!!! warning "Non-negativity constraint"
    NMF is not designed to fit negative values. This algorithm will warn you (but will not error) if a target or reference cube contains negative values. The output may seem reasonable, but it is not well-defined with respect to the NMF algorithm. To overcome this, rescaling the data by its minimum before processing is necessary
    ```julia
    target = cube .- minimum(cube)
    S = reconstruct(NMF(), target, angles)
    ```
    When doing full-frame reduction (e.g. `NMF()(cube, angles)`) *this is handled automatically*, so this constraint only applies to the lower-level API and methods which rely on those, such as [`GreeDS`](@ref). **In general, if you see warnings, heed them.**

# References
1. [Ren et al. 2018](http://adsabs.harvard.edu/abs/2018ApJ...852..104R) Non-negative Matrix Factorization: Robust Extraction of Extended Structures
"""
struct NMF <: ADIAlgorithm
    ncomps
end
NMF(; ncomps=nothing) = NMF(ncomps)

function subtract(alg::NMF, cube; kwargs...)
    target = normalize_nmf_input(cube)
    if :ref in keys(kwargs)
        ref_ = normalize_nmf_input(kwargs[:ref])
        S = reconstruct(alg, target; kwargs..., ref=ref_)
    else
        S = reconstruct(alg, target; kwargs...)
    end
    return target .- S
end

function normalize_nmf_input(data)
    if any(<(0), data)
        return data .- minimum(data)
    end
    return data
end

function fit(alg::NMF, data::AbstractMatrix{T}; ref=data, kwargs...) where T
    if any(<(0), data) || any(<(0), ref)
        @warn "Negative values encountered in input. Make sure to rescale your data"
    end
    if alg.ncomps === nothing
        k = size(data, 1)
    else
        k = min(alg.ncomps, size(data, 1))
    end

    X = collect(ref)
    W, H = nndsvd(X, k)
    solve!(CoordinateDescent{T}(), X, W, H)
    if ref !== data
        Y = collect(data)
        W = Y * H'
        solve!(CoordinateDescent{T}(update_H=false), Y, W, H)
    end
    return LinearDesign(H, W)
end