module MeltingParam

# If you want to add a new method here, feel free to do so.
# Remember to also export the function name in GeoParams.jl (in addition to here)

using Parameters, LaTeXStrings, Unitful
using ..Units
using GeoParams:
    AbstractMaterialParam, PhaseDiagram_LookupTable, AbstractMaterialParamsStruct
import Base.show, GeoParams.param_info
using ..MaterialParameters: MaterialParamsInfo
using Setfield # allows modifying fields in immutable struct

abstract type AbstractMeltingParam{T} <: AbstractMaterialParam end

export  compute_meltfraction,
        compute_meltfraction!,    # calculation routines
        compute_dϕdT,             # derivative of melt fraction versus
        compute_dϕdT!,
        param_info,
        MeltingParam_Caricchi,
        MeltingParam_4thOrder,
        MeltingParam_5thOrder,
        MeltingParam_Quadratic,
        MeltingParam_Assimilation,
        SmoothMelting

include("../Utils.jl")
include("../Computations.jl")

# Caricchi  -------------------------------------------------------
"""
    MeltingParam_Caricchi()

Implements the T-dependent melting parameterisation used by Caricchi, Simpson et al. (as for example described in Simpson)
```math
    \\theta = {(a - (T + c)) \\over b}
```
```math
    \\phi_{melt} = {1.0 \\over (1.0 + e^\\theta)}
```

Note that T is in Kelvin. As default parameters we employ:
```math
b=23K,  a=800K , c=273.15K
```
Which gives a reasonable fit to experimental data of granodioritic composition (Piwinskii and Wyllie, 1968):

![MeltParam_Carrichi](./assets/img/MeltingParam_Caricchi.png)

References
====
- Simpson G. (2017) Practical finite element modelling in Earth Sciences Using MATLAB.
"""
@with_kw_noshow struct MeltingParam_Caricchi{T,U} <: AbstractMeltingParam{T}
    a::GeoUnit{T,U}     = 800.0K
    b::GeoUnit{T,U}     = 23.0K
    c::GeoUnit{T,U}     = 273.15K # shift from C to K
    apply_bounds::Bool  = true
end
MeltingParam_Caricchi(args...) = MeltingParam_Caricchi(convert.(GeoUnit, args)...)

function param_info(s::MeltingParam_Caricchi) # info about the struct
    return MaterialParamsInfo(;
        Equation=L"\phi = {1 \over {1 + \exp( {800-T[^oC] \over 23})}}"
    )
end

# Calculation routines
function (p::MeltingParam_Caricchi)(; T, kwargs...)
    if T isa Quantity
        @unpack_units a, b, c = p
    else
        @unpack_val a, b, c = p
    end
    θ = (a - (T - c)) / b
    return 1.0 / (1.0 + exp(θ))
end

function compute_meltfraction!(ϕ::AbstractArray, p::MeltingParam_Caricchi; T, kwargs...)
    for i in eachindex(T)
        @inbounds ϕ[i] = p(; T=T[i])
    end
    return nothing
end

function compute_dϕdT(p::MeltingParam_Caricchi; T, kwargs...)
    if T isa Quantity
        @unpack_units a, b, c = p
    else
        @unpack_val a, b, c = p
    end

    dϕdT = exp((a + c - T) / b) / (b * ((1.0 + exp((a + c - T) / b))^2))

    return dϕdT
end

function compute_dϕdT!(dϕdT::AbstractArray, p::MeltingParam_Caricchi; T, kwargs...)
    for i in eachindex(T)
        @inbounds dϕdT[i] = compute_dϕdT(p, T=T[i])
    end

    return nothing
end

# Print info
function show(io::IO, g::MeltingParam_Caricchi)
    return print(io, "Caricchi et al. melting parameterization")
end
#-------------------------------------------------------------------------

# MeltingParam_5thOrder  -------------------------------------------------------
"""
    MeltingParam_5thOrder(a,b,c,d,e,f,T_s,T_l)

Uses a 5th order polynomial to describe the melt fraction `phi` between solidus temperature `T_s` and liquidus temperature `T_l`
```math
    \\phi = a T^5 + b T^4 + c T^3 + d T^2 + e T + f  \\textrm{   for   } T_s ≤ T ≤ T_l
```
```math
    \\phi = 1  \\textrm{   if   } T>T_l
```
```math
    \\phi = 0  \\textrm{   if   } T<T_s
```
Temperature `T` is in Kelvin.

![MeltingParam_5thOrder](./assets/img/MeltingParam_5thorder.png)

The default values are for a composite liquid-line-of-descent:
- the upper part is for Andesite from: (Blatter, D. L. & Carmichael, I. S. (2001) Hydrous phase equilibria of a Mexican highsilica andesite: a candidate for a mantle origin? Geochim. Cosmochim. Acta 65, 4043–4065
- the lower part is extrapolated to the granitic minimum using the Marxer & Ulmer LLD for Andesite (Marxer, F. & Ulmer, P. (2019) Crystallisation and zircon saturation of calc-alkaline tonalite from the Adamello Batholith at upper crustal conditions: an experimental study. *Contributions Mineral. Petrol.* 174, 84)

"""
@with_kw_noshow struct MeltingParam_5thOrder{T,U,U1,U2,U3,U4,U5,U6} <:
                       AbstractMeltingParam{T}
    a::GeoUnit{T,U1} = 2.083291971482524e-12 / K^5
    b::GeoUnit{T,U2} = -1.239502833666574e-08 / K^4
    c::GeoUnit{T,U3} = 2.938887604687626e-05 / K^3
    d::GeoUnit{T,U4} = -0.034711533077108 / K^2
    e::GeoUnit{T,U5} = 20.425403874539178 / K
    f::GeoUnit{T,U6} = -4.790664658179178e+03 * NoUnits
    T_s::GeoUnit{T,U} = 963.15K
    T_l::GeoUnit{T,U} = 1388.2K
    apply_bounds::Bool  = true
end
MeltingParam_5thOrder(args...) = MeltingParam_5thOrder(convert.(GeoUnit, args)...)

function param_info(s::MeltingParam_5thOrder) # info about the struct
    return MaterialParamsInfo(; Equation=L"\phi = aT^5 + bT^4 + cT^3 + dT^2 + eT + f")
end

# Calculation routines
function (p::MeltingParam_5thOrder)(; T, kwargs...)
    if T isa Quantity
        @unpack_units a, b, c, d, e, f, T_s, T_l  = p
    else
        @unpack_val a, b, c, d, e, f, T_s, T_l = p
    end

    ϕ = a * T^5 + b * T^4 + c * T^3 + d * T^2 + e * T + f
    if p.apply_bounds
        if T < T_s
            ϕ = 0.0
        elseif T > T_l
            ϕ = 1.0
        end
    end

    return ϕ
end

function compute_meltfraction!(
    ϕ::AbstractArray, p::MeltingParam_5thOrder; T::AbstractArray, kwargs...
)
    for i in eachindex(T)
        @inbounds ϕ[i] = p(; T=T[i])
    end

    return nothing
end

function compute_dϕdT(p::MeltingParam_5thOrder; T::Real, kwargs...)
    if T isa Quantity
        @unpack_units a, b, c, d, e, T_s, T_l = p
    else
        @unpack_val a, b, c, d, e, T_s, T_l = p
    end

    dϕdT = 5 * a * T^4 + 4 * b * T^3 + 3 * c * T^2 + 2 * d * T + e
    if (T < T_s || T > T_l) && p.apply_bounds
        dϕdT = 0.0
    end

    return dϕdT
end

function compute_dϕdT!(
    dϕdT::AbstractArray, p::MeltingParam_5thOrder; T::AbstractArray, kwargs...
)

    for i in eachindex(T)
        @inbounds dϕdT[i] = compute_dϕdT(p, T=T[i])
    end

    return nothing
end

# Print info
function show(io::IO, g::MeltingParam_5thOrder)
    return print(
        io,
        "5th order polynomial melting curve: phi = $(NumValue(g.a)) T^5 + $(NumValue(g.b))T^4 + $(NumValue(g.c))T^3 + $(NumValue(g.d))T^2 + $(NumValue(g.e))T + $(NumValue(g.f))  $(Value(g.T_s)) ≤ T ≤ $(Value(g.T_l))",
    )
end
#-------------------------------------------------------------------------

# MeltingParam_4thOrder  -------------------------------------------------------
"""
    MeltingParam_4thOrder(b,c,d,e,f,T_s,T_l)

Uses a 4th order polynomial to describe the melt fraction `phi` between solidus temperature `T_s` and liquidus temperature `T_l`
```math
    \\phi = b T^4 + c T^3 + d T^2 + e T + f  \\textrm{   for   } T_s ≤ T ≤ T_l
```
```math
    \\phi = 1 \\textrm{   if   } T>T_l
```
```math
    \\phi = 0  \\textrm{   if   } T<T_s
```
Temperature `T` is in Kelvin.

![MeltingParam_4thOrder](./assets/img/MeltingParam_4thorder.png)

The default values are for Tonalite experiments from Marxer and Ulmer (2019):
- Marxer, F. & Ulmer, P. (2019) Crystallisation and zircon saturation of calc-alkaline tonalite from the Adamello Batholith at upper crustal conditions: an experimental study. *Contributions Mineral. Petrol.* 174, 84

"""
@with_kw_noshow struct MeltingParam_4thOrder{T,U,U2,U3,U4,U5,U6} <: AbstractMeltingParam{T}
    b::GeoUnit{T,U2} = -7.594512597174117e-10 / K^4
    c::GeoUnit{T,U3} = 3.469192091489447e-06 / K^3
    d::GeoUnit{T,U4} = -0.005923529809260 / K^2
    e::GeoUnit{T,U5} = 4.482855645604745 / K
    f::GeoUnit{T,U6} = -1.268730161921053e+03 * NoUnits
    T_s::GeoUnit{T,U} = 963.15K
    T_l::GeoUnit{T,U} = 1270.15K
    apply_bounds::Bool  = true
end
MeltingParam_4thOrder(args...) = MeltingParam_4thOrder(convert.(GeoUnit, args)...)

function param_info(s::MeltingParam_4thOrder) # info about the struct
    return MaterialParamsInfo(; Equation=L"\phi = bT^4 + cT^3 + dT^2 + eT + f")
end

# Calculation routines
function (p::MeltingParam_4thOrder)(; T::Real, kwargs...)
    if T isa Quantity
        @unpack_units b, c, d, e, f, T_s, T_l = p
    else
        @unpack_val  b, c, d, e, f, T_s, T_l = p
    end

    ϕ = b * T^4 + c * T^3 + d * T^2 + e * T + f
    if p.apply_bounds
        if T < T_s
            ϕ = 0.0
        elseif T > T_l
            ϕ = 1.0
        end
    end

    return ϕ
end

function compute_meltfraction!(
    ϕ::AbstractArray, p::MeltingParam_4thOrder; T::AbstractArray, kwargs...
)
    for i in eachindex(T)
        @inbounds ϕ[i] = p(; T=T[i])
    end

    return nothing
end

function compute_dϕdT(p::MeltingParam_4thOrder; T::Real, kwargs...)
    if T isa Quantity
        @unpack_units  b, c, d, e, T_s, T_l = p
    else
        @unpack_val b, c, d, e, T_s, T_l = p
    end

    dϕdT = 4 * b * T^3 + 3 * c * T^2 + 2 * d * T + e
    if (T < T_s || T > T_l) && p.apply_bounds
        dϕdT = 0.0
    end
    return dϕdT
end

function compute_dϕdT!(
    dϕdT::AbstractArray, p::MeltingParam_4thOrder; T::AbstractArray, kwargs...
)

    for i in eachindex(T)
        @inbounds dϕdT[i] = compute_dϕdT(p, T=T[i])
    end

    return nothing
end

# Print info
function show(io::IO, g::MeltingParam_4thOrder)
    return print(
        io,
        "4th order polynomial melting curve: phi = $(NumValue(g.b))T^4 + $(NumValue(g.c))T^3 + $(NumValue(g.d))T^2 + $(NumValue(g.e))T + $(NumValue(g.f))  $(Value(g.T_s)) ≤ T ≤ $(Value(g.T_l))",
    )
end
#-------------------------------------------------------------------------

# MeltingParam_Quadratic  -------------------------------------------------------
"""
    MeltingParam_Quadratic(T_s,T_l)

Quadratic melt fraction parameterisation where melt fraction ``\\phi`` depends only on solidus (``T_s``) and liquidus (``T_l``) temperature:
```math
    \\phi = 1.0 - \\left( {T_l - T} \\over {T_l - T_s} \\right)^2
```
```math
    \\phi = 1.0 \\textrm{ if } T>T_l
```
```math
    \\phi = 0.0 \\textrm{ if } T<T_s
```
Temperature `T` is in Kelvin.

![MeltingParam_Quadratic](./assets/img/MeltingParam_Quadratic.png)

This was used, among others, in Tierney et al. (2016) Geology

"""
@with_kw_noshow struct MeltingParam_Quadratic{T,U} <: AbstractMeltingParam{T}
    T_s::GeoUnit{T,U}   = 963.15K
    T_l::GeoUnit{T,U}   = 1273.15K
    apply_bounds::Bool  = true
end
MeltingParam_Quadratic(args...) = MeltingParam_Quadratic(convert.(GeoUnit, args)...)

function param_info(s::MeltingParam_Quadratic) # info about the struct
    return MaterialParamsInfo(; Equation=L"\phi = 1.0 - ((T_l - T)/(T_l - T_s))^2")
end

# Calculation routines
function (p::MeltingParam_Quadratic)(; T, kwargs...)
    if T isa Quantity
        @unpack_units T_s, T_l = p
    else
        @unpack_val T_s, T_l = p
    end

    ϕ = 1.0 - ((T_l - T) / (T_l - T_s))^2
    if p.apply_bounds
        if T > T_l
            ϕ = 1.0
        elseif T < T_s
            ϕ = 0.0
        end
    end
    return ϕ
end

function compute_meltfraction!(
    ϕ::AbstractArray, p::MeltingParam_Quadratic; T::AbstractArray, kwargs...
)

    for i in eachindex(T)
        @inbounds ϕ[i] = p(; T=T[i])
    end

    return nothing
end

function compute_dϕdT(p::MeltingParam_Quadratic; T, kwargs...)
    if T isa Quantity
        @unpack_units T_s, T_l = p
    else
        @unpack_val T_s, T_l = p
    end

    dϕdT = (2T_l - 2T) / ((T_l - T_s)^2)
    if (T > T_l || T < T_s) && p.apply_bounds
        dϕdT = 0.0
    end
    return dϕdT
end

function compute_dϕdT!(
    dϕdT::AbstractArray, p::MeltingParam_Quadratic; T::AbstractArray, kwargs...
)

    for i in eachindex(T)
        @inbounds dϕdT[i] = compute_dϕdT(p, T=T[i])
    end

    return nothing
end

# Print info
function show(io::IO, g::MeltingParam_Quadratic)
    return print(
        io,
        "Quadratic melting curve:  ϕ = 1.0 - ((Tₗ-T)/(Tₗ-Tₛ))² with Tₛ=$(Value(g.T_s)), Tₗ=$(Value(g.T_l)) ",
    )
end
#-------------------------------------------------------------------------

# MeltingParam_Assimilation  -------------------------------------------------------
"""
    MeltingParam_Assimilation(T_s,T_l,a)

Melt fraction parameterisation that takes the assimilation of crustal host rocks into account, as used by Tierney et al. (2016) based upon a parameterisation of Spera and Bohrson (2001)

Here, the fraction of molten and assimilated host rocks ``\\phi`` depends on the solidus (``T_s``) and liquidus (``T_l``) temperatures of the rocks, as well as on a parameter ``a=0.005``
```math
    X = \\left( {T - T_s} \\over {T_l - T_s} \\right)
```
```math
    \\phi = a \\cdot \\left( \\exp^{2ln(100)X} - 1.0 \\right) \\textrm{ if } X ≤ 0.5
```
```math
    \\phi = 1- a \\cdot \\exp^{2ln(100)(1-X)}  \\textrm{ if } X > 0.5
```
```math
    \\phi = 1.0 \\textrm{ if } T>T_l
```
```math
    \\phi = 0.0 \\textrm{ if } T<T_s
```
Temperature `T` is in Kelvin.

![MeltingParam_Assimilation](./assets/img/MeltingParam_Assimilation.png)

This was used, among others, in Tierney et al. (2016), who employed as default parameters:
```math
   T_s=973.15, T_l=1173.15, a=0.005
```

References
==========
- Spera, F.J., and Bohrson, W.A., 2001, Energy-Constrained Open-System Magmatic Processes I: General Model and Energy-Constrained Assimilation and Fractional Crystallization (EC- AFC) Formulation: Journal of Petrology, v. 42, p. 999–1018.
- Tierney, C.R., Schmitt, A.K., Lovera, O.M., de Silva, S.L., 2016. Voluminous plutonism during volcanic quiescence revealed by thermochemical modeling of zircon. Geology 44, 683–686. https://doi.org/10.1130/G37968.1

"""
@with_kw_noshow struct MeltingParam_Assimilation{T,U,U1} <: AbstractMeltingParam{T}
    T_s::GeoUnit{T,U}   =   973.15K
    T_l::GeoUnit{T,U}   =   1173.15K
    a::GeoUnit{T,U1}    =   0.005NoUnits
    apply_bounds::Bool  =   true
end
MeltingParam_Assimilation(args...) = MeltingParam_Assimilation(convert.(GeoUnit,args)...)

function param_info(s::MeltingParam_Assimilation) # info about the struct
    return MaterialParamsInfo(Equation =  L"\phi = f(T_S,T_l,a) taking crustal assimilation into account.")
end

# Calculation routines
function (p::MeltingParam_Assimilation)(; T::Real, kwargs...)
    if T isa Quantity
        @unpack_units T_s,T_l,a   = p
    else
        @unpack_val T_s,T_l,a   = p
    end

    X = (T - T_s)/(T_l - T_s)

    if X <= 0.5
        ϕ   =    a * (exp(2*log(100)*X) - 1.0 )
    else
        ϕ   =    1.0 - a * exp(2*log(100)*(1-X) )
    end
    if p.apply_bounds
        if T>T_l
            ϕ = 1.0
        elseif T<T_s
            ϕ = 0.0
        end
    end
    return ϕ
end

function compute_meltfraction!(
    ϕ::AbstractArray, p::MeltingParam_Assimilation; T::AbstractArray, kwargs...
)

    for i in eachindex(T)
        @inbounds ϕ[i] = p(; T=T[i])
    end

    return nothing
end

function compute_dϕdT(p::MeltingParam_Assimilation; T::Real, kwargs...)
    if T isa Quantity
        @unpack_units T_s,T_l,a   = p
    else
        @unpack_val T_s,T_l,a   = p
    end

    X      =   (T - T_s)/(T_l - T_s)
    dϕdT   =   (9.210340371976184*a*exp((9.210340371976184*T - 9.210340371976184*T_s) / (T_l - T_s))) / (T_l - T_s)
    if X>0.5
        dϕdT   = (9.210340371976184*a*exp(9.210340371976184 + (9.210340371976184*T_s - 9.210340371976184*T) / (T_l - T_s))) / (T_l - T_s)
    end

    if (T>T_l || T<T_s) && p.apply_bounds
        dϕdT = 0.0
    end
    return dϕdT
end

function compute_dϕdT!(
    dϕdT::AbstractArray, p::MeltingParam_Assimilation; T::AbstractArray, kwargs...
)

    for i in eachindex(T)
        @inbounds dϕdT[i] = compute_dϕdT(p, T=T[i])
    end

    return nothing
end

# Print info
function show(io::IO, g::MeltingParam_Assimilation)
    return print(
        io,
        "Quadratic melting assimilation parameterisation after Spera & Bohrson (2001)",
    )
end

#-------------------------------------------------------------------------

# Smooth melting function ------------------------------------------------

"""
    SmoothMelting(; p=MeltingParam_4thOrder(), k_sol=0.2/K,  k_liq=0.2/K)

This smoothens the melting parameterisation ``p`` around the solidus ``T_{sol}`` and liquidus ``T_{liq}``
using a smoothened Heaviside step functions for the solidus:

```math
    H_{sol} =  {1.0 \\over { 1 + \\exp( -2 k_{sol} (T - T_{sol} - {2 \\over k_{sol}}) )  }}
```
and liquidus:
```math
    H_{liq} =  1.0 - {1.0 \\over { 1 + \\exp( -2 k_{liq} (T - T_{liq} + {2 \\over k_{liq}}) )  }}
```
The resulting melt fraction ``\\phi`` is computed from the original melt fraction ``\\phi_0`` (computed using one of the methods above) as:
```math
    \\phi =  \\phi_0 H_{sol} H_{liq} + 1.0 - H_{liq}
```
The width of the smoothening zones is controlled by ``k_{sol}, k_{liq}`` (larger values = sharper boundary).

This is important, as jumps in the derivative ``dϕ/dT`` can cause numerical instabilities in latent heat computations, which is prevented with this smoothening.

Example
====

Let's consider a 4th order parameterisation:
```julia
julia> using GeoParams, Plots
julia> p = MeltingParam_4thOrder();
julia> T= collect(650.0:1:1050.) .+ 273.15;
julia> T,phi,dϕdT =  PlotMeltFraction(p,T=T);
```

The same but with smoothening:
```julia
julia> p_s = SmoothMelting(p=MeltingParam_4thOrder(), k_liq=0.21/K);
4th order polynomial melting curve: phi = -7.594512597174117e-10T^4 + 3.469192091489447e-6T^3 + -0.00592352980926T^2 + 4.482855645604745T + -1268.730161921053  963.15 K ≤ T ≤ 1270.15 K with smooth Heaviside function smoothening using k_sol=0.1 K⁻¹·⁰, k_liq=0.11 K⁻¹·⁰
julia> T_s,phi_s,dϕdT_s =  PlotMeltFraction(p_s,T=T);
```

We can create plots of this with:
```julia
julia> plt1 = plot(T.-273.15, phi, ylabel="Melt Fraction ϕ", color=:red, label="original", xlabel="Temperature [C]")
julia> plt1 = plot(plt1, T.-273.15, phi_s,  color=:black, label="smoothened", legend=:bottomright)
julia> plt2 = plot(T.-273.15, dϕdT, ylabel="dϕ/dT", color=:red, label="original", xlabel="Temperature [C]")
julia> plt2 = plot(plt2, T.-273.15, dϕdT_s,  color=:black, label="smoothened", legend=:topright)
julia> plot!(plt1,plt2,   xlabel="Temperature [C]", layout=(2,1))
```
The derivative no longer has a jump now:

![MeltingParam_Smooth](./assets/img/MeltingParam_Smooth.png)

"""
struct SmoothMelting{P,T,U} <: AbstractMeltingParam{T}
    p::P
    k_sol::GeoUnit{T,U}
    k_liq::GeoUnit{T,U}
end

# Set default values:
function SmoothMelting(; p=MeltingParam_4thOrder(), k_sol=0.2/K,  k_liq=0.2/K)
    k_sol = convert(GeoUnit,k_sol)
    k_liq = convert(GeoUnit,k_liq)
    p = @set p.apply_bounds = false
    SmoothMelting(p, k_sol, k_liq)
end

SmoothMelting(p::AbstractMeltingParam) =  SmoothMelting(p=p)
#SmoothMelting(p::AbstractMeltingParam; k_liq) =  SmoothMelting(p=p, k_liq=convert.(GeoUnit,k_liq))

# Calculation routines
function (param::SmoothMelting)(; T, kwargs...)
    if T isa Quantity
        @unpack_units k_sol,k_liq  = param
    else
        @unpack_val k_sol,k_liq  = param
        k_sol = param.k_sol.val
    end

    ϕ = param.p(; T, kwargs...)     # Melt fraction computed in usual manner

    T_s = param.p.T_s
    H_s = 1.0/( 1.0 + exp(-2*k_sol*(T-T_s-(2/k_sol))));

    T_l = param.p.T_l
    H_l = 1.0 - 1.0/( 1.0 + exp(-2*k_liq*(T-T_l+(2/k_liq))));

    # Apply heaviside smoothening above liquidus & below solidus
    ϕ = ϕ*H_s*H_l + 1.0 - H_l

    return ϕ
end

function compute_meltfraction!(
    ϕ::AbstractArray, param::SmoothMelting;  T::AbstractArray, kwargs...
)
    for i in eachindex(T)
        @inbounds ϕ[i] = param(; T=T[i], kwargs=kwargs)
    end

    return nothing
end

function compute_dϕdT(param::SmoothMelting; T, kwargs...)

    if T isa Quantity
        @unpack_units k_sol,k_liq  = param
    else
        @unpack_val k_sol,k_liq  = param
    end

    # compute heaviside functions & derivatives of that vs. T
    T_s = param.p.T_s
    T_l = param.p.T_l
    H_s = 1.0/( 1.0 + exp(-2*k_sol*(T-T_s-(2/k_sol))));
    H_l = 1.0 - 1.0/( 1.0 + exp(-2*k_liq*(T-T_l+(2/k_liq))));

    dHs_dT = 2k_sol*( 1.0 / ((1.0 + exp(-2k_sol*(T + -2 / k_sol - T_s)))^2))*exp(-2k_sol*(T + -2 / k_sol - T_s))
    dHl_dT = 2k_liq*(-1.0 / ((1.0 + exp(-2k_liq*(T +  2 / k_liq - T_l)))^2))*exp(-2k_liq*(T +  2 / k_liq - T_l))

    # melt fraction & derivative
    dϕdT = compute_dϕdT(param.p, T=T)
    ϕ    = param.p(; T, kwargs...)

    # The derivative of the function
    # ϕ = ϕ(T)*H_s(T)*H_l(T)  + 1.0 - H_l
    # versus T is

    dϕdT_tot = dϕdT*H_s*H_l + ϕ*dHs_dT*H_l + ϕ*H_s*dHl_dT - dHl_dT

    return dϕdT_tot
end

function compute_dϕdT!(
    dϕdT::AbstractArray, p::SmoothMelting; T::AbstractArray, kwargs...
)

    for i in eachindex(T)
        @inbounds dϕdT[i] = compute_dϕdT(p, T=T[i])
    end

    return nothing
end

# Print info
function show(io::IO, g::SmoothMelting)
    param = show(io,g.p);
    return print(
        io,
        " with smooth Heaviside function smoothening using k_sol=$(Value(g.k_sol)), k_liq=$(Value(g.k_liq))",
    )
end
#-------------------------------------------------------------------------

"""
    compute_meltfraction(P,T, p::AbstractPhaseDiagramsStruct)

Computes melt fraction in case we use a phase diagram lookup table. The table should have the column `:meltFrac` specified.
"""
function compute_meltfraction(
    p::PhaseDiagram_LookupTable; P::_T, T::_T, kwargs...
) where {_T}
    return p.meltFrac.(T, P)
end

compute_meltfraction(p::PhaseDiagram_LookupTable, args) = compute_meltfraction(p; args...)
"""
    compute_meltfraction!(ϕ::AbstractArray{<:AbstractFloat}, P::AbstractArray{<:AbstractFloat},T:AbstractArray{<:AbstractFloat}, p::PhaseDiagram_LookupTable)

In-place computation of melt fraction in case we use a phase diagram lookup table. The table should have the column `:meltFrac` specified.
"""
function compute_meltfraction!(
    ϕ::AbstractArray{_T},
    p::PhaseDiagram_LookupTable;
    P::AbstractArray{_T},
    T::AbstractArray{_T},
    kwargs...,
) where {_T}
    ϕ[:] = p.meltFrac.(T, P)

    return nothing
end

compute_meltfraction!( ϕ::AbstractArray, p::PhaseDiagram_LookupTable, args) = compute_meltfraction!(p; args...)

"""
    compute_dϕdT(P,T, p::AbstractPhaseDiagramsStruct)

Computes derivative of melt fraction vs T in case we use a phase diagram lookup table. The table should have the column `:meltFrac` specified.
The derivative is computed by finite differencing.
"""
function compute_dϕdT(p::PhaseDiagram_LookupTable; P::_T, T::_T, kwargs...) where {_T}
    dT = (maximum(T) - minimum(T)) / 2.0 * 1e-6 + 1e-6   # 1e-6 of the average T value
    ϕ1 = p.meltFrac.(T .+ dT, P)
    ϕ0 = p.meltFrac.(T, P)
    dϕdT = (ϕ1 - ϕ0) / dT

    return dϕdT
end

compute_dϕdT(p::PhaseDiagram_LookupTable, args) = compute_dϕdT(p; args...)

"""
    compute_dϕdT!(dϕdT::AbstractArray{<:AbstractFloat}, P::AbstractArray{<:AbstractFloat},T:AbstractArray{<:AbstractFloat}, p::PhaseDiagram_LookupTable)

In-place computation of melt fraction in case we use a phase diagram lookup table. The table should have the column `:meltFrac` specified.
The derivative is computed by finite differencing.
"""
function compute_dϕdT!(
    dϕdT::AbstractArray{_T},
    p::PhaseDiagram_LookupTable;
    P::AbstractArray{_T},
    T::AbstractArray{_T},
    kwargs...,
) where {_T}
    dT = (maximum(T) - minimum(T)) / 2.0 * 1e-6 + 1e-6   # 1e-6 of the average T value
    ϕ1 = p.meltFrac.(T .+ dT, P)
    ϕ0 = p.meltFrac.(T, P)
    dϕdT = (ϕ1 - ϕ0) / dT

    return nothing
end

compute_dϕdT!( ϕ::AbstractArray, p::PhaseDiagram_LookupTable, args) = compute_dϕdT!(p; args...)

# fill methods programatically
for myType in (
    :MeltingParam_Caricchi,
    :MeltingParam_5thOrder,
    :MeltingParam_4thOrder,
    :MeltingParam_Quadratic,
    :MeltingParam_Assimilation,
    :SmoothMelting,
)
    @eval begin
        (p::$(myType))(args) = p(; args...)
        compute_meltfraction(p::$(myType), args) = p(; args...)
        function compute_meltfraction!(ϕ::AbstractArray, p::$(myType), args)
            return compute_meltfraction!(ϕ, p; args...)
        end
        compute_dϕdT(p::$(myType), args) = compute_dϕdT(p; args...)
        function compute_dϕdT!(dϕdT::AbstractArray, p::$(myType), args)
            return compute_dϕdT!(dϕdT, p; args...)
        end
    end
end

# Computational routines needed for computations with the MaterialParams structure
function compute_meltfraction(s::AbstractMaterialParamsStruct, args)
    if isempty(s.Melting) #in case there is a phase with no melting parametrization
        return zero(typeof(args).types[1])
    else
        return compute_meltfraction(s.Melting[1], args)
    end
end

function compute_dϕdT(s::AbstractMaterialParamsStruct, args)
    if isempty(s.Melting) #in case there is a phase with no melting parametrization
        return zero(typeof(args).types[1])
    else
        return compute_dϕdT(s.Melting[1], args)
    end
end

"""
    ϕ = compute_meltfraction(Phases::AbstractArray{<:Integer}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

Computation of melt fraction ϕ for the whole domain and all phases, in case an array with phase properties `MatParam` is provided, along with `P` and `T` arrays.
"""
compute_meltfraction(args...) = compute_param(compute_meltfraction, args...)

"""
    compute_meltfraction(ϕ::AbstractArray{<:AbstractFloat}, Phases::AbstractArray{<:Integer}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

In-place computation of melt fraction ϕ for the whole domain and all phases, in case an array with phase properties `MatParam` is provided, along with `P` and `T` arrays.
"""
compute_meltfraction!(args...) = compute_param!(compute_meltfraction, args...)

"""
    ϕ = compute_dϕdT(Phases::AbstractArray{<:Integer}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

Computates the derivative of melt fraction ϕ versus temperature `T` for the whole domain and all phases, in case an array with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This is employed in computing latent heat terms in an implicit manner, for example
"""
compute_dϕdT(args...) = compute_param(compute_dϕdT, args...)

"""
    compute_dϕdT!(ϕ::AbstractArray{<:AbstractFloat}, Phases::AbstractArray{<:Integer}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

Computates the derivative of melt fraction `ϕ` versus temperature `T`, ``{\\partial \\phi} \\over {\\partial T}`` for the whole domain and all phases, in case an array with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This is employed, for example, in computing latent heat terms in an implicit manner.
"""
compute_dϕdT!(args...) = compute_param!(compute_dϕdT, args...)

end


Base.@irrational Ωe 0.00007292115 big(0.00007292115)
Ωe, const Omegae = Ωe
gsw_f(latitude::T) where {T<:Real} = T(2) * Ωe * sind(latitude)


using Distributions
using SpinGlassNetworks

function ramp(t::T, τ::T, α::T, pi::T, pf::T) where T <: Real
    p = (pf + pi) + (pf - pi) * tanh(α * (2.0 * t / τ - 1.0))
    p / 2.0
end

@testset "Simple Coherent Ising Machine simulator for small Ising instances." begin
    L = 4
    ig = ising_graph("$(@__DIR__)/instances/basic/$(L)_001.txt")

    scale = 0.2
    noise = Normal(0.0, 0.1)

    x0 = zeros(L)
    sat = 1.0
    time = 1000.
    pi, pf, α = -15.0, 0.0, 2.0
    momentum = 0.4

    pump = [ramp(t, time, α, pi, pf) for t ∈ 0:time]

    opo = OpticalOscillators{Float64}(ig, scale, noise)
    dyn = OPODynamics{Float64}(x0, sat, pump, momentum)

    @testset  "OpticalOscillators and OPODynamics work properly." begin
        @test opo.scale ≈ scale
        @test opo.noise ≈ noise

        @test dyn.initial_state ≈ x0
        @test dyn.saturation ≈ sat
        @test dyn.pump ≈ pump
    end

    N = 500
    states = Vector{Vector{Int}}(undef, N)
    Threads.@threads for i ∈ 1:N
        states[i] = evolve_optical_oscillators(opo, dyn)
    end

    @test minimum(energy(states, ig)) ≈ brute_force(ig, :CPU, num_states=1).energies[1]
end


n = 2

mass = 1.0
gravity = 9.81
mass_matrix = Diagonal(mass * ones(n))
bias = [0.0; gravity * mass]
contact_jacobian = Diagonal(ones(n))
ϕ(q) = q[2]

h = 0.01

a = 500.0
softplus(x) = 1 / a * log(1.0 + exp(a * x))
softminus(x) = x - softplus(x)

plot(softplus)
plot(softminus)

function dynamics(x, u, t)
	q = x[1:n]
	v = x[n .+ (1:n)]

	k_impact = 500.0
	b_impact = 10.0
	# λ = [0.0; -k_impact * min(0.0, ϕ(q))]
	λ = [0.0; -k_impact * softminus(ϕ(q)) + max(0.0, - b_impact * v[2])]

	v⁺ = v + h * (mass_matrix \ (-bias + transpose(contact_jacobian) * λ))
	q⁺ = q + h * v⁺

	return [q⁺; v⁺]
end

q1 = [0.0; 1.0]
v1 = [0.0; 0.0]
x1 = [q1; v1]

x_hist = [x1]
T = 1000

for t = 1:T-1
	push!(x_hist, dynamics(x_hist[end], nothing, t))
	println(ϕ(x_hist[end][1:2]))
end

plot(hcat(x_hist...)[1:2, :]', )


using Interpolations
using StaticArrays
using LinearAlgebra
using Statistics
using QuadGK
using Plots
# using Makie

function Rot2(θ)
    sθ,cθ = sincos(θ)
    @SMatrix [cθ -sθ; sθ cθ]
end

abstract type AbstractPath end

struct CarPath{T} <: AbstractPath
    s::Vector{T}
    ϕ::Vector{T}
    κ::Vector{T}
    X::Vector{T}
    Y::Vector{T}
end

function curvature(path::CarPath, s)
    k = searchsortedfirst(path.s, s)
    return path.κ[k]
end

struct CirclePath <: AbstractPath
    r::Float64
end
curvature(path::CirclePath, s) = 1/path.r

""" ArcPath """
struct ArcPath <: AbstractPath
    ang0::Float64
    radius::Float64
    ang::Float64
    xc::Float64
    yc::Float64
    function ArcPath(ang0, radius, ang)
        xc = -radius*sin(ang0)
        yc =  radius*cos(ang0)
        new(ang0, radius, ang, xc, yc)
    end
end

function position_change(arc::ArcPath)
    localToGlobal(arc, total_length(arc), 0)
end
total_length(arc::ArcPath) = arc.radius * arc.ang
curvature(arc::ArcPath, s) = 1/arc.radius
final_angle(arc::ArcPath) = arc.ang0 + arc.ang
function localToGlobal(arc::ArcPath, s, e)
    r = arc.radius
    θ = s ./ r
    radius = r .- e
    # radius = r
    ϕ = θ .- π/2 .+ arc.ang0
    x = radius .* cos.(ϕ) .+ arc.xc
    y = radius .* sin.(ϕ) .+ arc.yc
    return x,y
end

""" Straight Path """
struct StraightPath <: AbstractPath
    len::Float64
    ang::Float64
end

function position_change(path::StraightPath)
    x = path.len * cos(path.ang)
    y = path.len * sin(path.ang)
    return x,y
end
total_length(path::StraightPath) = path.len
final_angle(path::StraightPath) = path.ang
curvature(path::StraightPath, s) = 0.0
function localToGlobal(path::StraightPath, s, e)
    Rot2(path.ang)*@SVector [s, e]
end

"Build a straight path that continues from an arc"
function StraightPath(arc::ArcPath, len)
    StraightPath(len, final_angle(arc))
end

"Build an arc that continues from a line"
function ArcPath(line::StraightPath, radius, ang)
    ArcPath(final_angle(line), radius, ang)
end

""" Dubins Path """
struct DubinsPath{P} <: AbstractPath
    paths::P
    lengths::Vector{Float64}
    ds::Vector{Float64}
    dx::Vector{Float64}
    dy::Vector{Float64}
    curvature::Vector{Float64}
end

function DubinsPath(paths)
    num_paths = length(paths)
    lengths = cumsum([total_length(p) for p in paths])
    ds = circshift(lengths, 1)
    ds[1] = 0
    dx,dy = zeros(num_paths), zeros(num_paths)
    for i = 2:num_paths
        dx[i],dy[i] = position_change(paths[i-1])
    end
    dx = cumsum(dx)
    dy = cumsum(dy)
    κ = [curvature(path,0.0) for path in paths]
    DubinsPath(paths, lengths, ds, dx, dy, κ)
end

total_length(path::DubinsPath) = path.lengths[end]
final_angle(path::DubinsPath) = final_angle(path.paths[end])
function curvature(path::DubinsPath, s)
    i = match_segment(path, s)::Int
    return path.curvature[i]
end

function match_segment(path::DubinsPath, s::AbstractVector)
    [searchsortedfirst(path.ds, s_) for s_ in s]
end

function match_segment(path::DubinsPath, s::Real)
    j = searchsortedlast(path.ds, s)
    min(j, length(path.ds))
end

function localToGlobal(path::DubinsPath, s, e)
    x,y = zero(s), zero(s)
    for i in eachindex(s)
        j = match_segment(path, s[i])
        s_ = s[i] - path.ds[j]
        x[i],y[i] = localToGlobal(path.paths[j], s_, e[i])
        x[i] += path.dx[j]
        y[i] += path.dy[j]
    end
    return x,y
end

function Plots.plot(path::AbstractPath; kwargs...)
    s = range(0, total_length(path), length=101)
    e = zero(s)
    x,y = localToGlobal(path, s, e)
    plot(x,y; kwargs...)
end

function Plots.plot!(path::AbstractPath; kwargs...)
    s = range(0, total_length(path), length=101)
    e = zero(s)
    x,y = localToGlobal(path, s, e)
    plot!(x,y; kwargs...)
end

"""
Get a trajectory of s, ϕ, κ from a set of E,N pairs defining the path, where
    s is the distance along the path
    ϕ is the heading wrt the X-axis (E)
    κ is the curvature of the road, in 1/m
"""
function pathToLocal(X,Y)
    # Parameterize the curve from 0 to 1
    N = length(X)
    t = range(0,1,length=N)

    # Create cubic interpolations
    itpX = CubicSplineInterpolation(t,X)
    itpY = CubicSplineInterpolation(t,Y)

    # Compute the gradient vectors
    g = map(t) do i
        dx = Interpolations.gradient(itpX, i)[1]
        dy = Interpolations.gradient(itpY, i)[1]
        @SVector [dx,dy]
    end

    # Compute the hessian vectors
    h = map(t) do i
        ddx = Interpolations.hessian(itpX, i)[1]
        ddy = Interpolations.hessian(itpY, i)[1]
        @SVector [ddx,ddy]
    end

    # Calculate the heading
    ϕ = map(x->atan(x[2],x[1]), g)

    # Calculate the curvature
    κ = map(1:N) do k
        r′ = g[k]
        r″ = h[k]
        k = norm(r′ × r″) / norm(r′)^3
    end

    # Calulcate the path length: ∫ₐᵇ ‖r′(t)‖ dt
    x,w = gauss(N,0,1)
    @show any(isnan,x)
    r = map(x) do i
        dx = Interpolations.gradient(itpX, i)[1]
        dy = Interpolations.gradient(itpY, i)[1]
        sqrt(dx^2 + dy^2)
    end
    s = cumsum(r .* w)

    # Make the first path length zero
    s[1] = 0.0

    return s, ϕ, κ
end

"""
Attempt to recover the E,N coordanates from a set of local coordinates.
Doesn't work super well due to integration drift
"""
function localToGlobal(s,ϕ,κ)
    N = length(s)
    @assert length(ϕ) == N
    @assert length(κ) == N

    X = zeros(N)
    Y = zeros(N)
    for k = 2:N
        r = 1/mean(κ[k-1:k])
        ds = s[k] - s[k-1]
        θ = ds/r
        psi = mean(ϕ[k-1:k])
        psi = ϕ[k-1]
        xc = X[k-1] - r*sin(psi)
        yc = Y[k-1] + r*cos(psi)
        X[k] = xc + r*sin(psi + θ)
        Y[k] = yc - r*cos(psi + θ)
    end
    return X,Y
end

"""
Get the E,N coordinate given the distance along the path (s) and the longitudinal error (e)
    given a nominal path described by X,Y,S,Φ
Inputs s,e can be vectors or scalars
"""
function localToGlobal(path::CarPath, s, e)
    plot(path.X,path.Y, aspect_ratio=:equal)
    X,Y,S,Φ = path.X, path.Y, path.s, path.ϕ
    itpX = interpolate((S,),X, Gridded(Linear()))
    itpY = interpolate((S,),Y, Gridded(Linear()))
    itpΦ = interpolate((S,),Φ, Gridded(Linear()))
    nomX = itpX.(s)
    nomY = itpY.(s)
    nomΦ = itpΦ.(s)
    x = @. nomX - e*sin(nomΦ)
    y = @. nomY + e*cos(nomΦ)
    return x,y
end

function localToGlobal(path::CirclePath, s, e)
    r = path.r
    t = range(0,2pi,length=101)
    X,Y = r*cos.(t), r*sin.(t)
    plot(X,Y, aspect_ratio=:equal)
    θ = s ./ r
    radius = r .+ e
    x = radius .* cos.(θ)
    y = radius .* sin.(θ)
    return x,y
end


using ParamLevelSet
using jInv.Mesh
using SparseArrays
using LinearAlgebra

function heavi(x, type="global", eps=0.1, thr=0.0) # [h,d] =

    x .= x - thr;

    if type == "global"
        h = 0.5 * (1 + 2/pi*atan(pi* x /eps))
        d = 1.0 / (eps*((x .^ 2 * pi^2) / eps^2 + 1))

    elseif type == "compact"
        # compact
        # h = 0*x;
        # d = 0*x;
        # id = find((x < epsi) & (x > -epsi));
        # h(id) = 0.5*(1 + x(id)/epsi + 1/pi*sin(pi*x(id)/epsi));
        # h(x >= epsi) = 1;
        # h(x <=-epsi) = 0;
        # d(id) = 0.5*(1/epsi)*(1 + cos(pi*x(id)/epsi))
    end

    return h, d
end

function recon2d_slices_pals!(u::Array{T, 3}, A, p, niter::Int=50, nbasis::Int=4) where {T <: AbstractFloat}

    global Jt, J0, mesh, m, Hessian, grad, r

    mesh, m = init_pals(size(u), nbasis)
    sigmaH = getDefaultHeaviside();
    phat = similar(p)
    r = similar(p)
    At = A'
    stepsize = 0.01
    nslice = size(p,2)

    HW = size(u,1)*size(u,2)
    szpslice = size(p,1)*size(p,3)
    J = spzeros(prod(size(p)), nbasis*5) # A * J0

    for i=1:niter
        # u0, JBuilder = eval_f_grad(mesh, m; computeJacobian=1, sigma=sigmaH)
        u0, JBuilder = MeshFreeParamLevelSetModelFunc(mesh, m; computeJacobian = 1)

        J0 = getSparseMatrix(JBuilder)

        # construct J
        Threads.@threads for ss=1:nslice
            J[szpslice*(ss-1)+1:szpslice*ss, :] .= A * J0[HW*(ss-1)+1:HW*ss, :]

            u_slice = view(u, :, :, ss)
            Av = A*vec(u_slice)

            # phat_slice = reshape(Av, size(p,1), size(p,3))
            p_vec_ = reshape(Av, size(p,1), size(p,3))
            @views r[:,ss,:] = p_vec_ - p[:,ss,:]
        end

        Jt = J'
        Hessian = Jt*J
        Hinv = pinv(Matrix(Hessian))
        grad = Jt * vec(r)
        m .-= stepsize * Hinv * grad

        fres = sum(r .^ 2) / length(r)
        @show fres, m
        # println("residual: $fres, $m")
    end

    # u0, JBuilder = eval_f_grad(mesh,m;computeJacobian=1, sigma=sigmaH)
    u0, JBuilder = MeshFreeParamLevelSetModelFunc(mesh, m; computeJacobian = 1)
    u .= reshape(u0, size(u))
    return u, m
end


export lens
export micro_lens_array

"""
    lens(f, x, y; λ=550e-9, n=1)

Return complex transmission value at position `x` and `y`
for a lens.
The lens has a focal length `f` and the surrounding medium has
refractive index `n`. Wavelength is `λ`.
"""
function lens(f::Number, x::Number, y::Number; radius=Inf, λ=550e-9, n=1)
    κ = calc_κ(λ, n)
    r2 = @. x ^ 2 + y ^ 2

    out = @. circ(radius, x, y) * exp(-1im * π * κ / f * (x^2 + y^2))
    return out
end

"""
    lens(f, L, size=(512, 512); λ=550e-9, n=1)

Return complex lens transmission value for a field of `size`
and a length of `L`
The lens has a focal length `f` and the surrounding medium has
refractive index `n`. Wavelength is `λ`.

"""
function lens(f::Number, L::Number, size::Tuple=(512, 512); radius=Inf, λ=550e-9, n=1)
    # initialize output array with complex numbers
    out = zeros(ComplexF64, size)

    # iterate over values
    for (j, x) in enumerate(fftpos(L, size[2]))
        for (i, y) in enumerate(fftpos(L, size[1]))
            out[i, j] = circ(radius, x, y) * lens(f, x, y; radius=radius, λ=λ, n=n)
        end
    end

    return out
end

"""
    micro_lens_array(f, L, pitch, size=(512, 512); centered=true, λ=550e-9, n=1, aperture=true,
                     dtype=ComplexF32)

Return a micro lens array (MLA) with pitch size `pitch` in the field size of `L`.
All lenslets have the same focal length `f`.
As default `centered=true` means that the there is a micro lens centered around
the center of the array.
`aperture=true` indicates that there is a circular aperture cutting the field.
Otherwise the aperture are square shaped of the same size as the `pitch`.
"""
function micro_lens_array(f, L, pitch, size=(512, 512); centered=true, λ=550e-9, n=1, aperture=true, dtype=ComplexF32)
    # pitch size in indices
    mla_size = round.(Int, (size .* pitch) ./  L)
    # ÷ doesn't convert to Int tuple
    mla_size = convert(Tuple{Int, Int}, mla_size)

    # if centered, we embed a odd number of micro lenses
    # into a larger arrays
    # at the end, we do a center extract of this larger array
    # otherwise it is troublesome to calculate the offset and
    # cut off micro lenses
    if centered
        N_mla = ceil(Int, L / pitch)
        if mod(N_mla, 2) == 0
            N_mla = N_mla + 1
        end

        size_out = size
        size = N_mla .* mla_size
    end

    out = zeros(dtype, size)
    # put the micro lenses in the out field
    for j = 1:mla_size[2]:size[2]
        for i = 1:mla_size[1]:size[1]
            # min is required in the case when a full micro lens
            # doesn't fit anymore
            mla_i = min(mla_size[1], size[1] - (i - 1))
            mla_j = min(mla_size[2], size[2] - (j - 1))

            out_i = min(size[1], i - 1 + mla_size[1])
            out_j = min(size[2], j - 1 + mla_size[2])

            # add small microlens
            if aperture
                r = pitch / 2
            else
                r = Inf
            end
            out[i:out_i, j:out_j] = lens(f, pitch, mla_size, λ=λ, n=n, radius=r)[1:mla_i, 1:mla_j]
        end
    end

    if centered
        out = center_extract(out, size_out)
    end

    return out
end


"""
    $(SIGNATURES)
    quality function constructor for least square optimization
"""
function make_qfunc(model, x, y)
    function qfunc(p)
        mask = (y .!= 0)
        model_res = model(x, p)
        sum(((model_res[mask] .- y[mask]).^2) ./ y[mask])
    end
    qfunc
end

"""
    $(SIGNATURES)
    performs pre fit
    # Arguments
    - `chargedist => ChargeDist`: charge distribution on which pre fit is
                                  performed
"""
function pre_fit(chargedist::ChargeDist)
    x = chargedist.x
    y = chargedist.y
    qfunc_ped = make_qfunc(gauss, x, y)
    mxval_ped, mxidx_ped = findmax(y)
    p0_ped = [x[mxidx_ped], 0.01, mxval_ped]
    fit_ped = optimize(qfunc_ped, p0_ped, NewtonTrustRegion())
    popt_ped = Optim.minimizer(fit_ped)
    mask = (x .> (popt_ped[1] + 3 * popt_ped[2]))
    x_spe = x[mask]
    y_spe = y[mask]
    mxval_spe, mxidx_spe = findmax(y_spe)
    mask = (x_spe .< x_spe[mxidx_spe] * 2)
    x_spe = x_spe[mask]
    y_spe = y_spe[mask]
    qfunc_spe = make_qfunc(gauss, x_spe, y_spe)
    p0_spe = [x_spe[mxidx_spe], 0.1, mxval_spe]
    fit_spe = optimize(qfunc_spe, p0_spe, NewtonTrustRegion())
    popt_spe = Optim.minimizer(fit_spe)
    PreFit(popt_ped..., popt_spe...)
end

"""
    $(SIGNATURES)
    performs pmt response fit
    # Arguments
    - `charges => ChargeDist`: charge distribution on which fit is performed
    - `prefit_results => PreFitResults`: results from pre fit used as starting
                                         values for fit
"""
function pmtresp_fit(chargedist::ChargeDist, prefit::PreFit; mod=:default)
    p0 = [prefit.μₚ,
          prefit.σₚ,
          prefit.μₛ,
          prefit.σₛ,
          prefit.Aₛ / prefit.Aₚ,
          prefit.Aₚ + prefit.Aₛ]
    if mod == :uap
        push!(p0, prefit.μₛ / 5)
        push!(p0, prefit.σₛ / 5)
        push!(p0, prefit.Aₛ / 5)
        func = pmtresp_uap
        ResultStruct = PMTRespUapFit
    else
        func = pmtresp
        ResultStruct = PMTRespFit
    end
    qfunc = make_qfunc(func, chargedist.x, chargedist.y)
    fit = optimize(qfunc, p0, Newton())
    popt = Optim.minimizer(fit)

    ResultStruct(popt...)
end


struct FDPoisson2DProblem{TF, TBCT, TBCB, TBCL, TBCR}
    a # (a,b)×(c,d) define the domain
    b
    c
    d
    Δx # mesh spacing in x
    nx # number of points in the interior, excluding a and b
    Δy # mesh spacing in y
    ny # number of points in the interior, excluding c and d
    x # interior mesh
    y # interior mesh
    f::TF # value of f at the interior nodes
    A # the matrix
    rhs # right hand side of our system
    top_bc::TBCT
    bottom_bc::TBCB
    left_bc::TBCL
    right_bc::TBCR
end

struct DirichletBC
    u_bc
end

struct NeumannBC
    v_bc
end

struct PeriodicBC
end


include("euler/euler.jl")
using .Primes: prime_sieve
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(max_a::Int, max_b::Int)::Int
    prime = prime_sieve(max_a * max_b)
    prime_numbers = findall(prime[1:max_b])
    append!(prime_numbers, .-prime_numbers)
    max_sequence, result = 0, 0
    for a ∈ -max_a + (max_a % 2 == 0 ? 1 : 2):2:max_a - 1
        for b ∈ prime_numbers
            n = 0
            while true
                candidate = abs(n * (n + a) + b)
                if candidate ≠ 0 && !prime[candidate] break end
                n += 1
            end
            if n > max_sequence max_sequence, result = n, a * b end
        end
    end
    return result
end

compute(1_000, 1_000)

@benchmark compute(1_000, 1_000)


using BoardGames

mutable struct MCTSStrategy{G <: Game} <: Strategy
    nsteps::Int
    c::Float64
end

mutable struct TreeNode{Board}
    board::Board

    n::Float64
    w::Float64

    expanded::Bool
    children::Vector{TreeNode{Board}}
    parent::TreeNode{Board}

    function TreeNode(board::Board) where {Board}
        new{Board}(board, 0, 0, false, TreeNode{Board}[])
     end
end

function TreeNode(board::B, parent::TreeNode{B}) where {B}
    tn = TreeNode(board)
    tn.parent = parent
    return tn
end

struct MCTSTree{Board}
    root::TreeNode{Board}
end

MCTSTree(board::Board) where {Board} = MCTSTree(TreeNode(board))

function MCTSStrategy{G}(nsteps::Int) where {G <: Game}
    MCTSStrategy{G}(nsteps, √2)
end

function BoardGames.getmove(board, s::MCTSStrategy; verbose=false)
    tree = MCTSTree(board)
    player = playerturn(board)

    for i in 1:s.nsteps
        step(tree.root, player, s.c)
    end
    if verbose
        results = [(m, c.w, c.n) for (m,c) in zip(getmoves(board), tree.root.children)]
        @show results
    end
    move_idx = select(tree.root, player, 0)

    verbose && println("selected $move_idx")

    getmoves(board)[move_idx]
end

function expand(node::TreeNode)
    moves = getmoves(node.board)

    for move in moves
        push!(node.children, TreeNode(play(node.board, move), node))
    end
    node.expanded = true
    nothing
end

function select(node::TreeNode, player::Int, c=√2)
    params = zeros(Float64, size(node.children))

    for (i,child) in enumerate(node.children)
        if child.n == 0
            params[i] = Inf
        else
            if player == playerturn(node.board)
                params[i] =  child.w / child.n + c*sqrt(log(node.n)/child.n)
            else
                params[i] = -child.w / child.n + c*sqrt(log(node.n)/child.n)
            end
        end
    end
    return random_argmax(params)
end

function backpropagate(node::TreeNode, w::Float64)
    while(isdefined(node, :parent))
        node.w += w
        node.n += 1
        node = node.parent
    end
    node.w += w
    node.n += 1
    nothing
end

function step(node::TreeNode, player::Int, c::Float64)
    if !node.expanded
        expand(node)
    end
    if isempty(node.children)
        w = 0.
        if winner(node.board) == 0
            w = .5
        else
            w = Float64(winner(node.board) == player)
        end
        backpropagate(node, w)
    else
        i = select(node, player, c)
        step(node.children[i], player, c)
    end
    nothing
end

function random_argmax(v)
    @assert(length(v) > 0)
    best = v[1]
    best_idx = 1
    best_count = 1

    for i in 2:length(v)
        if v[i] > best
            best = v[i]
            best_idx = i
            best_count = 1
        elseif v[i] == best
            best_count += 1
            if rand() < 1/best_count
                best_idx = i
            end
        end
    end
    return best_idx
end

function Base.copy(s::MCTSStrategy{T}) where T
    MCTSStrategy{T}(s.nsteps, s.c)
end

function BoardGames.getvarsnames(s::MCTSStrategy)
    return ["iterations", "c"]
end

function BoardGames.getvalues(s::MCTSStrategy)
    (s.nsteps, s.c)
end

function BoardGames.setvalue!(s::MCTSStrategy, name::AbstractString, value)
    if name == "iterations"
        s.nsteps = parse(Int, value)
    elseif name == "c"
        s.c = parse(Float64, value)
    else
        error("invalid set for parameter $name for strategy $s")
    end
    nothing
end

BoardGames.name(::MCTSStrategy) = "MCTS"


using DelimitedFiles
using Test
using Base.Iterators

# Those corresponding verifiable data are obtained from orginal fortran subroutines at
# http://www.bosai.go.jp/study/application/dc3d/DC3Dhtml_E.html
@testset "Test Okada's displacement" begin
    @testset "Normal computations" begin
        u_truth = readdlm(joinpath(@__DIR__, "data/test_dc3d.dat"), ',', Float64)
        pos = product(-40.: -30., 20.: 30., 10.: 15.)
        fu = (x) -> dc3d(reverse(x)..., 2. /3, 50., 70., [-80., 120.], [-30., 25.], [200., -150., 100.])
        u_cal = map(fu, pos) |> vec
        cache = GeoGreensFunctions.dc3d_cache(Float64)
        for (i, x) in enumerate(pos)
            dc3d(reverse(x)..., 2. /3, 50., 70., [-80., 120.], [-30., 25.], [200., -150., 100.], cache)
            @test u_truth[i,:] ≈ u_cal[i] ≈ cache[1]
        end
    end

    @testset "Negative depth" begin
        u = dc3d(10., 20., 30., 2. /3, 50., 70., [-80., 120.], [-30., 25.], [200., -150., 100.])
        @test isapprox(zeros(Float64, 12), u, rtol=1e-8)
    end

    @testset "Singularity" begin
        u_singular = zeros(Float64, 12)

        u = dc3d(-80., 0., -50., 2. /3, 50., 70., [-80., 120.], [-30., 25.], [200., -150., 100.])
        @test isapprox(u_singular, u, atol=1e-8)

        u = dc3d(-70., 0., -10., 2. /3, 10., 90., [-80., 120.], [-20., 20.], [200., -150., 100.])
        @test isapprox(u_singular, u, atol=1e-8)

        u = dc3d(-80., 20., -50., 2. /3, 50., 70., [-80., 120.], [-30., 25.], [200., -150., 100.])
        u_truth = [
            -32.60576975, 59.39436703, 1.61816100, -0.43682428, 1.66503564, 0.27751175,
            0.64126234, -1.05836342, 0.03630938, -0.16974972, 0.61800758, 0.22842501]
        @test isapprox(u_truth, u, rtol=1e-8)

        u = dc3d(-10., 20., -50., 2. /3, 50., 90., [-80., 120.], [-30., 25.], [200., -150., 100.])
        u_truth = [
            -56.04511679, 52.38458483, 40.87786501, -0.10140811, -0.09099896, 0.02371607,
            1.60327214, -0.63369117, -1.13609685, 0.05783881, 0.88665943, 0.21946899,
        ]
        @test isapprox(u_truth, u, rtol=1e-8)

        u = dc3d(-10., 20., 0., 2. /3, 0., 0., [-80., 120.], [-30., 25.], [200., -150., 100.])
        # take a look: https://github.com/JuliaLang/julia/issues/23376
        @test isapprox(u_singular, u, atol=1e-8)

        u = dc3d(-100., 20., 0., 2. /3, 0., 0., [-80., -90.], [20., 25.], [200., -150., 100.])
        @test isapprox(u_singular, u, atol=1e-8)

        u = dc3d(-80., 20., 0., 2. /3, 0., 0., [-80., 100.], [25., 35.], [200., -150., 100.])
        @test isapprox(u_singular, u, atol=1e-8)
    end
end


names = ["Sprinkles", "PeanutButter", "Frosting", "Sugar"]
capacity = [5, -1, 0, -1]
durability = [-1, 3, -1, 0]
flavor = [0, 0, 4, 0]
texture = [0, 0, 0, 2]
calories = [5, 1, 6, 8]

maxscore = 0
for parts in partitions(100, 4)
	for perms in permutations(parts)
		if dot(perms, calories) != 500
			continue
		end
		score = max(0,dot(perms, capacity))*max(0,dot(perms, durability))*max(0,dot(perms, flavor))*max(0,dot(perms, texture))
		if score > maxscore
			maxscore = score
		end
	end
end
println(maxscore)


## make the surface struct
using HiQGA.transD_GP, Random
surface = transD_GP.Surface(imagegrid;useML=false, σ=noisegrid)
## stationary GP options
nmin, nmax= 2, 100
pnorm = 2.
K = transD_GP.GP.Mat32()
ftrain = surface.d[surface.select]
fbounds = permutedims([extrema(ftrain)...])
δ = 0.3
λ = [50_000,50_000.]
demean = false
sampledc = true
sdev_prop = 0.05*diff(fbounds, dims=2)[:]
sdev_pos = 0.05*abs.([diff([extrema(xall[1,:])...])[1], diff([extrema(xall[2,:])...])[1]])
updatenonstat = false
needλ²fromlog = false
## Initialize a stationary GP using these options
Random.seed!(12)
opt = transD_GP.OptionsStat(nmin = nmin,
                        nmax = nmax,
                        xbounds = [minimum(xall[1,:]) maximum(xall[1,:]);
                                   minimum(xall[2,:]) maximum(xall[2,:])],
                        fbounds = fbounds,
                        xall = xall,
                        λ = λ,
                        δ = δ,
                        demean = demean,
                        sampledc = sampledc,
                        sdev_prop = sdev_prop,
                        sdev_pos = sdev_pos,
                        pnorm = pnorm,
                        quasimultid = false,
                        K = K,
                        timesλ = 3,
                        needλ²fromlog = needλ²fromlog,
                        updatenonstat = updatenonstat
                        )


@testset "Test Function Spaces" begin

using jFEMTools
import Tensors: Vec, gradient, dot, norm, Tensor
using LinearAlgebra
jF = jFEMTools

# Load mesh
mesh = unitSquareMesh2(TriangleCell, (2,2));

# ### Define Function Spaces
dim = 2
finiteElement = ContinuousLagrange(:Triangle,1)
Wh = FEMFunctionSpace(mesh, finiteElement, 1)    #Scalar Space
Vh = FEMFunctionSpace(mesh, finiteElement, 2)    #Vector Space

# Basic Test
@test jF.getnlocaldofs(Vh) == 6
@test jF.getnlocaldofs(Wh) == 3
@inbounds for (cellcount, cell) in enumerate(jF.CellIterator(mesh))
    jF.reinit!(Wh, cell)
    jF.reinit!(Vh, cell)
end

# Wh Scalar Space testset
@test jF.getnquadpoints(Wh) == 3
@test jF.getnbasefunctions(Wh) == 3
@test jF.getdetJdV(Wh, 1) ≈ 1/24
@test jF.geometric_value(Wh, 1, 1) ≈ 2/3
@test jF.getdim(Wh) == 2
@test jF.shape_value(Wh,1,1) ≈ 2/3
@test jF.shape_gradient(Wh,1,1) ≈ Vec{2}((0.0,-2.0))
@test jF.shape_divergence(Wh,1,1) ≈ -2.0

# Vh Vector Space testset
@test jF.getnquadpoints(Vh) == 3
@test jF.getnbasefunctions(Vh) == 6
@test jF.getdetJdV(Vh, 1) ≈ 1/24
@test jF.geometric_value(Vh, 1, 1) ≈ 2/3
@test jF.getdim(Vh) == 2
@test jF.shape_value(Vh,1,1) ≈ [2/3, 0.0]
@test jF.shape_gradient(Vh,1,1) ≈ Tensor{2,2}([0.0 -2.0;0.0 0.0])
@test jF.shape_divergence(Vh,1,2) ≈ 2.0

end


using FractionalCalculus
using Test

@testset "Test Direct Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 0, 1.23, 1e-8, RLDirect())[1], 1.163931646862977; atol = 0.00001)
    @test isapprox(fracint(x->x^4, 0.25, 0, 1.23, 1e-8, RLDirect())[1], 1.64294134724409; atol = 0.00001)
    @test isapprox(fracint(x->x^4, 0.15, 0, 1.23, 1e-8, RLDirect())[1], 1.878978485456182; atol = 0.00001)
    @test isapprox(fracint(x->sin(x), 0.5, 0, 0.97, 1e-8, RLDirect())[1], 0.6442601778843754; atol = 0.00001)
    @test isapprox(fracint(x->sin(x), 0.25, 0, 0.97, 1e-8, RLDirect())[1], 0.7448298999485293; atol = 0.00001)
    @test isapprox(fracint(x->sin(x), 0.15, 0, 0.97, 1e-8, RLDirect())[1], 0.7802555196986999; atol = 0.00001)
end

@testset "Test Piecewise Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-4, RLPiecewise()), 1.163931646862977; atol = 1e-5)
    @test isapprox(fracint(x->x^4, 0.25, 1.23, 1e-4, RLPiecewise()), 1.64294134724409; atol = 1e-5)
    @test isapprox(fracint(sin, 0.5, 0.97, 1e-4, RLPiecewise()), 0.6442601778843754; atol = 1e-5)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-4, RLPiecewise())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
end

@testset "Test RLInt_Approx" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-4, RLIntApprox()), 1.163931646862977; atol = 1e-4)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-4, RLIntApprox())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
end

@testset "Test RLInt_Matrix" begin
    @test isapprox(fracint(x->x, 0.5, 1, 0.5, RLIntMatrix()), [0; 0.3535533905932738; 0.8838834764831844]; atol=1e-3)
end

@testset "Test RL_LinearInterp" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-4, RLLinearInterp()), 1.163931646862977; atol = 1e-4)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-4, RLLinearInterp())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
end

@testset "Test Simpson Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-7, RLIntSimpson()), 1.163931646862977; atol = 1e-2)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-7, RLIntSimpson())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-2)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-2)
    @test isapprox(result[3], 64.3280200256207; atol = 1)
end

@testset "Test Trapezoidal Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-7, RLIntTrapezoidal()), 1.163931646862977; atol = 1e-4)
#=
    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-7, RLInt_Trapezoidal())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
    =#
end

@testset "Test Rectangular Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-6, RLIntRectangular()), 1.163931646862977; atol = 1e-3)

    @test isapprox(fracint(x->x^4, 0.5, [1.23, 2, 3], 0.001, RLIntRectangular()), [1.1617493910169032; 10.363130001411157; 64.27898282619807]; atol=1e-2)
end

@testset "Test Cubic Spline Interpolation Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 0.001, RLIntCubicSplineInterp()), 1.163931646862977; atol = 1e-5)
    @test isapprox(fracint(sin, 0.5, 0.97, 0.01, RLIntCubicSplineInterp()), 0.6442601778843754; atol = 1e-6)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 0.01, RLIntCubicSplineInterp())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
end

@testset "Test fractional integral macros" begin
    @test isapprox(@fracint(x->x^4, 0.5, 1.23), 1.163931646862977; atol = 1e-4)
    @test isapprox(@semifracint(x->x^4, 1.23), 1.163931646862977; atol = 1e-4)
end

@testset "Test Hadamard integral" begin
    (a, b)=fracint(sin, 0.5, 1, 2, 1/3, 0.1, HadamardMat())

    @test isapprox(a, [ 1.0
    1.023136962522694
    1.077552576065
    1.1626221795557758
    1.2764357600635179
    1.414213562373095
    1.5668630279525204
    1.7202493081322234
    1.85605792647593
    1.954772501883529
    2.0]; atol=1e-3)
    @test isapprox(b, [Inf
    3.2545536740283265
    1.9422879792123904
    1.5034729488252023
    1.291625449283718
    1.1492602172225017
    1.0157975575699352
    0.8713525896641822
    0.725484956837982
    0.6074878993825767
    0.5500616274393197]; atol=1e-3)
end


struct AssetModel
    X0
    sigma
end

# stochastic process interface

function stateSize(self::AssetModel)   # dimension of X(t)
    return 1
end

function factors(self::AssetModel)   # dimension of W(t)
    return 1
end

function initialValues(self::AssetModel)
    return [ 0.0 ]
end

# calculate volatility taking into account exogenous hybrid adjuster
function volatility(self::AssetModel, t, Y)
    return self.sigma + Y[3]
end

function evolve(self::AssetModel, t0, Y0, dt, dW, Y1)
    sigma = volatility(self,t0,Y0)
    Y1[1] = Y0[1] + (Y0[2] - 0.5*sigma*sigma)*dt + sigma*dW[1]*sqrt(dt)
end

function asset(self::AssetModel, t, Y, alias)
    return self.X0 * exp(Y[1])
end


function gensys_uncertain_zlb(prob_vec::AbstractVector{S}, Th::AbstractMatrix{S}, Ch::AbstractVector{S},
                              Tzlbs::Vector{Matrix{S}}, Rzlbs::Vector{Matrix{S}}, Czlbs::Vector{Vector{S}},
                              Γ0_til::AbstractMatrix{S}, Γ1_til::AbstractMatrix{S}, Γ2_til::AbstractMatrix{S},
                              C_til::AbstractVector{S}, Ψ_til::AbstractMatrix{S}) where {S <: Real}

    Tbars = Vector{Matrix{S}}(undef, length(Tzlbs) + 1)
    Cbars = Vector{Vector{S}}(undef, length(Tzlbs) + 1)
    Tout  = Vector{Matrix{S}}(undef, length(Tzlbs) + 1)
    Rout  = Vector{Matrix{S}}(undef, length(Tzlbs) + 1)
    Cout  = Vector{Vector{S}}(undef, length(Tzlbs) + 1)

    # Calculate "uncertain" ZLB matrices and back out the transition equation
    for i in 1:length(Tzlbs)
        Tbars[i] = prob_vec[1] * Tzlbs[i] + prob_vec[2] * Th # it is assumed Tzlbs is a vector of the T_{t + 1}^{(zlb)} matrices
        Cbars[i] = prob_vec[1] * Czlbs[i] + prob_vec[2] * Ch

        Tout[i]  = (Γ2_til * Tbars[i] + Γ0_til) \ Γ1_til
        Rout[i]  = (Γ2_til * Tbars[i] + Γ0_til) \ Ψ_til
        Cout[i]  = (Γ2_til * Tbars[i] + Γ0_til) \ (C_til  - Γ2_til * Cbars[i])
    end

    # Add boundary condition
    Tout[end] = Tzlbs[end]
    Rout[end] = Rzlbs[end]
    Cout[end] = Czlbs[end]

    return Tout, Rout, Cout
end


# for the DCHSBM representation
function linear_density(d_n::Tuple{AbstractDict{<:Integer, <:AbstractDict}, UnitRange{<:Integer}})
    d = d_n[1]
    edge_count = sum([sum(values(x)) for (k,x) in d])
    n = length(d_n[2])
    kmax = maximum(keys(d))
    possible_edges = sum([binomial(n,k) for k in 1:kmax])
    edge_count/possible_edges
end
function log_density(d_n::Tuple{AbstractDict{<:Integer, <:AbstractDict}, UnitRange{<:Integer}})
    size = hypergraphsize(d_n)
    size == 0 && return NaN
    n = length(d_n[2])
    log(size)/log(n)
end
function FBD.hypergraphsize(d_n::Tuple{AbstractDict{<:Integer, <:AbstractDict}, UnitRange{<:Integer}})
    sum([k*sum(values(x)) for (k,x) in d_n[1]])
end

# for the Kronecker representation
function linear_density(hedges::NTuple{3, <:AbstractVector{<:Integer}})
    edge_count = length(hedges[1])
    edge_count == 0 && return NaN
    n = maximum(maximum.(hedges))
    kmax = 3
    possible_edges = sum([binomial(n,k) for k in 1:kmax])
    edge_count/possible_edges
end
function log_density(hedges::NTuple{3, <:AbstractVector{<:Integer}})
    size = hypergraphsize(hedges)
    size == 0 && return NaN
    n = maximum(maximum.(hedges))
    log(size)/log(n)
end

# for the FBD representation
function linear_density(g::AbstractVector)
    edge_count = length(g)
    edge_count == 0 && return NaN
    n = maximum(maximum.(g))
    kmax = maximum(length.(g))
    possible_edges = sum([binomial(n,k) for k in 1:kmax])
    edge_count/possible_edges
end
function log_density(g::AbstractVector)
    size = hypergraphsize(g)
    size == 0 && return NaN
    n = maximum(maximum.(g))
    log(size)/log(n)
end


using COSMO
using Distributions
using Plots
using PortfolioOpt
using PortfolioOpt.TestUtils
plotly()

##################### Read Prices #####################
prices = get_test_data();
numD, numA = size(prices) # A: Assets    D: Days

##################### Calculating returns #####################
returns_series = percentchange(prices);

# market_history = VolumeMarketHistory(returns_series)

##################### backtest Parameters #####################
DEFAULT_SOLVER = optimizer_with_attributes(
    COSMO.Optimizer, "verbose" => false, "max_iter" => 900000
)

date_range = timestamp(returns_series)[100:end];

## Strategies parameters
# Look back days
k_back = 60
# results
backtest_results = Dict()

##################### backtest stochastic strategies #####################

backtest_results["markowitz_limit_R"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)
    Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

    # Parameters
    d = MvNormal(r̄, Σ)
    R = 0.0015 / market_budget(market)

    formulation = PortfolioFormulation(MIN_SENSE,
        ObjectiveTerm(SqrtVariance(d)),
        RiskConstraint(ExpectedReturn(d), GreaterThan(R)),
    )

    pointers = change_bids!(market, formulation, DEFAULT_SOLVER)
    return pointers
end

backtest_results["markowitz_limit_var"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)
    Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

    # Parameters
    d = MvNormal(r̄, Σ)
    V = 0.003 / market_budget(market)

    formulation = PortfolioFormulation(MAX_SENSE,
        ObjectiveTerm(ExpectedReturn(d)),
        RiskConstraint(SqrtVariance(d), LessThan(V)),
    )

    pointers = change_bids!(market, formulation, DEFAULT_SOLVER)
    return pointers
end

##################### backtest with RO strategies #####################

backtest_results["normal_limit_soyster"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)
    Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

    # Parameters
    d = MvNormal(r̄, Σ)
    R = -0.05 / market_budget(market)

    formulation = PortfolioFormulation(MAX_SENSE,
        ObjectiveTerm(ExpectedReturn(d)),
        RiskConstraint(ExpectedReturn(BudgetSet(d, maximum(abs.(returns); dims=1)'[:] .- r̄, numA * 1.0), WorstCase), GreaterThan(R)),
    )

    pointers = change_bids!(market, formulation, DEFAULT_SOLVER)
    return pointers
end

backtest_results["normal_limit_bertsimas_30"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)
    Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

    # Parameters
    d = MvNormal(r̄, Σ)
    R = -0.05 / market_budget(market)

    formulation = PortfolioFormulation(MAX_SENSE,
        ObjectiveTerm(ExpectedReturn(d)),
        RiskConstraint(ExpectedReturn(BudgetSet(d, maximum(abs.(returns); dims=1)'[:] .- r̄, numA * 0.3), WorstCase), GreaterThan(R)),
    )

    pointers = change_bids!(market, formulation, DEFAULT_SOLVER)
    return pointers
end

backtest_results["normal_limit_bertsimas_60"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)
    Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

    # Parameters
    d = MvNormal(r̄, Σ)
    R = -0.05 / market_budget(market)

    formulation = PortfolioFormulation(MAX_SENSE,
        ObjectiveTerm(ExpectedReturn(d)),
        RiskConstraint(ExpectedReturn(BudgetSet(d, maximum(abs.(returns); dims=1)'[:] .- r̄, numA * 0.6), WorstCase), GreaterThan(R)),
    )

    pointers = change_bids!(market, formulation, DEFAULT_SOLVER)
    return pointers
end

backtest_results["normal_limit_bental"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)
    Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

    # Parameters
    d = MvNormal(r̄, Σ)
    R = -0.05 / market_budget(market)

    formulation = PortfolioFormulation(MAX_SENSE,
        ObjectiveTerm(ExpectedReturn(d)),
        RiskConstraint(ExpectedReturn(EllipticalSet(d, 0.1), WorstCase), GreaterThan(R)),
    )

    pointers = change_bids!(market, formulation, DEFAULT_SOLVER)
    return pointers
end

##################### backtest with Data-Driven RO strategies #####################

backtest_results["mixed_signals_limit_betina"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)

    # Parameters
    j_robust = 25
    R = -0.015 / market_budget(market)
    num_train = 20
    kst = 14
    klt = 30
    kmom = 20
    Q = 2

    # Predict return
    r̄ = zeros(numA)
    for i in 1:numA
        r̄[i] = mixed_signals_predict_return(
            returns[:, i], num_train, kst, klt, kmom, Q
        )
    end

    formulation = PortfolioFormulation(MAX_SENSE,
        ObjectiveTerm(ExpectedReturn(Dirac(r̄))),
        RiskConstraint(ConditionalExpectedReturn{WorstCase}(Inf, DeterministicSamples(returns'[:,:]), j_robust), GreaterThan(R)),
    )

    pointers = change_bids!(market, formulation, DEFAULT_SOLVER)
    return pointers
end

##################### backtest with Moment DRO strategies #####################

backtest_results["delage_utility"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)
    Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

    # Parameters
    d = MvNormal(r̄, Σ)
    R = -0.015 / market_budget(market)

    formulation = PortfolioFormulation(MAX_SENSE,
        ObjectiveTerm(ExpectedUtility(MomentUncertainty(d, 0.05, 1.3),
            PieceWiseUtility(
                [1.0], [0.0]
            ),
            WorstCase
        ))
    )

    pointers = change_bids!(market, formulation, DEFAULT_SOLVER)
    return pointers
end

##################### extra strategies #####################

backtest_results["max_sharpe"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)
    returns = values(past_returns)
    Σ, r̄ = mean_variance(returns[(end - k_back):end, :])

    change_bids!(market, market_budget(market) * max_sharpe(Σ, r̄, risk_free_rate(market)))

    return nothing
end

backtest_results["equal_weights"], _ = sequential_backtest_market(
    VolumeMarketHistory(returns_series), date_range,
) do market, past_returns, ext
    # Prep
    numD, numA = size(past_returns)

    change_bids!(market, market_budget(market) * equal_weights(numA))

    return nothing
end

##################### plot results  #####################
using Plots
using Plots.PlotMeasures

plt = plot(;title="Culmulative Wealth",
    xlabel="Time",
    ylabel="Wealth",
    legend=:outertopright,
    left_margin=10mm
);
for (strategy_name, recorders) in backtest_results
    plot!(plt,
        axes(get_records(recorders[:wealth]), 1), get_records(recorders[:wealth]).data;
        label=strategy_name,
    )
end
plt

plt2 = scatter(;
    title="Mean Vs Std",
    xlabel="σ",
    ylabel="E[r]",
    legend=:outertopright,
    left_margin=10mm
);
for (strategy_name, recorders) in backtest_results
    scatter!(plt2,
    [std(get_records(recorders[:returns]).data)], [mean(get_records(recorders[:returns]).data)];
        label=strategy_name,
    )
end
plt2


module Jennrich
using Manifolds, LinearAlgebra

export jennrich

"""
	jennrich(A, r :: Int)

Compute the rank r canonical polyadic decomposition of A using a pencil-based algorithm
"""
jennrich(A, r :: Int) = jennrich(A, ones(Int, r))

"""
	jennrich(A, L :: AbstractVector{Int})

Compute the (Lᵣ,Lᵣ,1)-BTD of A using a pencil-based algorithm
"""
function jennrich(𝔗 :: AbstractArray{<:Real, 3}, L :: AbstractVector{Int})
	n₁, n₂, n₃ = size(𝔗)
	@assert n₁ ≥ sum(L) && n₂ ≥ sum(L)

	𝔗Tuck = TuckerPoint(𝔗, (sum(L), sum(L), 2))
	𝔊     = 𝔗Tuck.hosvd.core

	evd   = eigen(𝔊[:,:,1] * 𝔊[:,:,2]^-1)
	makeReal!(evd)
	A :: Matrix{eltype(𝔗)} = 𝔗Tuck.hosvd.U[1] * real(evd.vectors)

	B⨀C   = (A \ reshape(𝔗, n₁, :))'
	B     = Matrix{eltype(A)}(undef, n₂, sum(L))
	C     = Matrix{eltype(A)}(undef, n₃, length(L))
	blkInd = cumsum(vcat(1, L))
	for i ∈ 1:length(L)
		# Bᵢ⨀Cᵢ is the concatenation of rank 1 matrices Bᵢ[:,j] ⊗ cᵢ
		# Compute the rank 1 SVD of these matrices (they all have cᵢ as a right singular vector)
		from, to = blkInd[i], blkInd[i+1] - 1
		Bᵢ⨀Cᵢ = B⨀C[:,from:to]
		rk1mtx = reshape(Bᵢ⨀Cᵢ * randn(L[i]), (n₂, n₃))
		C[:,i] = leftSingularVector(rk1mtx', maxiter=2) # right singular vector
		for j = 1:L[i]
			rk1mtx = reshape(Bᵢ⨀Cᵢ[:,j], (n₂, n₃))
			B[:,from-1+j] = rk1mtx * C[:,i]  # left singular vector * singular value
		end
	end
	A, B, C
end

# Compute the dominant left singular vector of A
# This only works well if the second singular value of A is significantly less than the first
function leftSingularVector(A; maxiter=1)
	σv = randn(size(A,2))
	u = normalize(A * σv)
	for i = 1:maxiter - 1
		mul!(σv, A', u)
		normalize!(mul!(u, A, σv))
	end
	u
end

function makeReal!(evd::Union{GeneralizedEigen, Eigen})
	λ, V = evd
	i = firstindex(λ)
	while i ≤ lastindex(λ)
		if imag(λ[i]) ≠ 0
			λ[i+1] = conj(λ[i+1])
			V[:,i] = real.(V[:,i])
			V[:,i+1] = imag.(V[:,i+1])
			i += 1
		end
		i += 1
	end
	evd
end

end # module


using LinearAlgebra

"""
This file contains a list of cost functions that are ready to be used in optimisations
I'm going to try and keep some standard inputs that make sense
"""

"""
Target-gate fidelity
"""
function C1(KT, KN, D)
    1 - abs2(tr(KT' * KN)/D)
end

"""
Target-state infidelity
"""
function C2(ψT, ψN)
    1 - abs2(ψT' * ψN)
end

"""
Control amplitudes
"""
function C3(u)
    sum(abs2.(u))
end

"""
Control variations
"""
function C4(u)
    (K, J) = size(u)
    sum(abs2.(diff(u, dims = 2)))
end

"""
Occupation of forbidden states
"""
function C5(ψF, ψj)
    # since ψj is the intermediate states we can assume its an array
    sum(abs2(ψF' * ψ) for ψ in ψj)
end

"""
Evolution time (target gate)
"""
function C6(KT, KJ, N, D)
    1 - 1/N * sum(abs2(tr(KT' * Kj)/D) for Kj in KJ)
end

"""
Evolution time (target state)
"""
function C7(ψT, ψJ, N)
    1 - 1/N * sum(abs2(ψT, ψj) for ψj in ψJ)
end


using Serialization

function julia_func(z,frame)
    k = frame*0.002
    c = complex(-0.5, 0.66-k)
    z = z^2 + c
    return z
end

function gen_julia(julia_set,frame,im_width,im_height)

    max_it = 300
    max_z2 = 100
    min_x, max_x = -1, 1
    min_y, max_y = -1, 1

    for ix in 1:im_width
        for iy in 1:im_height
            real_part = ((max_x-min_x)/im_width)*ix + min_x
            im_part = ((max_y-min_y)/im_height)*iy + min_y
            z = complex(real_part, im_part)
            it = 0
            while(abs2(z) < max_z2 && it < max_it)
                z = julia_func(z,frame)
                it += 1
            ratio = it/max_it
            julia_set[ix,iy,frame] = ratio
            end
        end
    end
    return julia_set
end

function main()
    frames = 50
    im_width, im_height = 1000, 1000
    julia_set = zeros(im_width,im_height,frames)
    for frame in 1:frames
        gen_julia(julia_set,frame,im_width,im_height)
    end

    serialize("./data/julia_data",julia_set)
end

main()


# function that returns the joint tree for an n-rotor vtol vehicle
# include("simulate.jl")
function gen_nRotor(n::Integer)
    # Inertial and mass properties from:
    # Hossain, M. Raju, and Nicholas Krouglicof. "Multi-body dynamics modeling & control of quadrotor helicopter using bond graph." Proceedings of the International Conference on Bond Graph Modeling and Simulation. 2016.

    mb = 0.30155; mp = 0.011721;

    Ib = zeros(3,3);
    Ib[1,1] = 6.4e-4; Ib[2,2] = 6.4e-4; Ib[3,3] = 1.22e-3;

    Ip = zeros(3,3);
    Ip[1,1] = 6.4061e-7; Ip[2,2] = 6.2732e-6; Ip[3,3] = 3.008e-5;

    InFrame = InertialFrameAsRB()
    QuadFrame = RigidBody(mb,Ib,2)
    x0Quad = [zeros(3);[1;zeros(3)];zeros(3);zeros(4)];
    initialiseRigidBody!(QuadFrame,x0Quad);

    Props = Vector{RigidBody}(undef,n);
    j = Vector{Joint}(undef, n+1);
    j[1] = Joint(InFrame,QuadFrame,zeros(3),zeros(3));

    th = 2*pi/n;
    l = 1.0;
    axis = [0.0 0.0 1.0][:];
    for i=1:n
        Props[i] = RigidBody(mp, Ip, 2+i);
        ang =(i-1)*th; # angle
        r0Prop_i = [l*cos(ang); l*sin(ang); 0.0];
        β0Prop_i = abs.([cos(ang/2); sin(ang/2).*[0.0;0.0;1.0]]);
        x0Prop_i = [r0Prop_i; β0Prop_i; zeros(7)];
        initialiseRigidBody!(Props[i],x0Prop_i);

        j[i+1] = Joint(QuadFrame, Props[i], r0Prop_i, zeros(3), type = "Revolute", axis = axis);
    end
    return j
end
## simulate quadcopter
# clearconsole();
# qRotor_j = gen_nRotor(4);
#
# tEnd = 0.1;
# tSpan = 0.01;
# g =([0.0,0.0,-9.81]); # Gravity Vector.
# tSim, solFinal = simulate(tEnd, tSpan, qRotor_j, g=g);
#
# solQuad = solFinal[1];
# solProp1 = solFinal[2];
# solProp2 = solFinal[3];
# solProp3 = solFinal[4];
# solProp4 = solFinal[5];


# Decision Tree Module

module Tree
export DecisionTree, create_decision_tree, predict

"""
Data structure for Decision Tree\\
One of `children` and `target` will be `nothing`
"""
mutable struct DecisionTree
    col_id::Integer
    children::Union{AbstractDict{Function,DecisionTree},Nothing}
    target::Union{AbstractFloat,Nothing}
end

"""
C4.5 implementation for Decision Tree\\
`X_data` has 2 dims, `Y_data` has 1 dims\\
Returns a computed `DecisionTree` structure
"""
function create_decision_tree(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number; max_depth::Integer=10)::DecisionTree
    @assert ndims(X_data) == ndims(Y_data) + 1 == 2
    @assert size(X_data)[1] == size(Y_data)[1]
    @assert max_depth >= 1

    """
    Output the most frequent value in the vector
    """
    function majority(Y_vec::Array)::Number
        unique_vals = Dict{Number, Integer}()
        for Y_val in Y_vec
            if !haskey(unique_vals, Y_val)
                unique_vals[Y_val] = 1
            else
                unique_vals[Y_val] += 1
            end
        end
        result = sort(collect(unique_vals), by=m->m[2])
        return result[end][1]
    end

    """
    Check if vector has same value
    """
    function is_same_class(Y_vec::Array)::Bool
        val = Y_vec[1]
        for m in Y_vec[2:end]
            if m != val
                return false
            end
        end
        return true
    end

    """
    Get an array of unique values from a vector
    """
    function get_unique_array(X_vec::Array)::Array
        X_max = maximum(X_vec)
        X_min = minimum(X_vec)
        round_digits = -Integer(trunc(log10(X_max-X_min)-2))
        result = []
        for m in X_vec
            if !(trunc(m, digits=round_digits) in result)
                push!(result, m)
            end
        end
        return result
    end

    """
    The entropy function
    """
    function info(Y_data::Array)::AbstractFloat
        portion = Dict{Number,Integer}()
        for m in Y_data
            if !haskey(portion, m)
                portion[m] = 1
            else
                portion[m] += 1
            end
        end
        T = size(Y_data)[1]
        result = 0.0
        for m in keys(portion)
            p = portion[m] / T
            result += -p * log2(p)
        end
        return result
    end

    """
    Detect types for each column of `X_data`\\
    `true` if continuous\\
    `false` if categorical
    """
    function detect_types(X_data::Array)::Array
        result = Array{Bool}(undef, size(X_data)[2])
        for i in 1:size(X_data)[2]
            X_unique = get_unique_array(X_data[:, i])
            if length(X_unique) > 10
                result[i] = true
            else
                result[i] = false
            end
        end
        return result
    end

    function create_decision_tree_recursive(X_data::Array, Y_data::Array, X_type::Array, max_depth::Integer)::DecisionTree
        if is_same_class(Y_data)
            return DecisionTree(-1, nothing, Y_data[1])
        end
        if max_depth <= 0
            return DecisionTree(-1, nothing, majority(Y_data))
        end
        max_index = -1
        max_ratio = 0.0
        max_fn = nothing
        for i in 1:size(X_data)[2]
            X_vec_unique = get_unique_array(X_data[:, i])
            if X_type[i]
                # continous data, then choose binary threshold
                fn_max = nothing
                ratio_max = 0
                for j in 1:(length(X_vec_unique)-1)
                    # choose average as threshold
                    threshold = (X_vec_unique[j] + X_vec_unique[j+1]) / 2.0
                    fn = (m) -> m <= threshold
                    # compute information gain
                    Ii = info(Y_data)
                    Y_1 = Y_data[fn.(X_data[:, i])]
                    Y_2 = Y_data[(!fn).(X_data[:, i])]
                    I1 = info(Y_1)
                    I2 = info(Y_2)
                    gain = Ii - (length(Y_1)/length(Y_data)*I1 + length(Y_2)/length(Y_data)*I2)
                    splitInfo = -(length(Y_1)/length(Y_data))*log2(length(Y_1)/length(Y_data)) - (length(Y_2)/length(Y_data))*log2(length(Y_2)/length(Y_data))
                    ratio = gain / splitInfo
                    if fn_max === nothing
                        fn_max = fn
                        ratio_max = ratio
                    else
                        if ratio > ratio_max
                            fn_max = fn
                            ratio_max = ratio
                        end
                    end
                end
                if ratio_max > max_ratio
                    max_index = i
                    max_fn = fn_max
                    max_ratio = ratio_max
                end
            else
                # skip visited categorical feature (homogeneous)
                if length(X_vec_unique) <= 1
                    continue
                end
                # categorical data, then compute by splitting on each category
                Ii = info(Y_data)
                gain = Ii
                splitInfo = 0.0
                for val in X_vec_unique
                    fn = (m) -> m == val
                    Y0 = Y_data[fn.(X_data[:, i])]
                    I0 = info(Y0)
                    gain -= (length(Y0)/length(Y_data))*I0
                    splitInfo -= (length(Y0)/length(Y_data))*log2(length(Y0)/length(Y_data))
                end
                ratio = gain / splitInfo
                if ratio > max_ratio
                    max_index = i
                    max_fn = nothing
                    max_ratio = ratio
                end
            end
        end
        if max_index < 0
            return DecisionTree(-1, nothing, majority(Y_data))
        end
        children = Dict{Function,DecisionTree}()
        if max_fn === nothing
            # categorical
            X_vec_unique = get_unique_array(X_data[:, max_index])
            for val in X_vec_unique
                fn = (m) -> m == val
                identify = fn.(X_data[:, max_index])
                X_part = X_data[identify, :]
                Y_part = Y_data[identify]
                children[fn] = create_decision_tree_recursive(X_part, Y_part, X_type, max_depth-1)
            end
        else
            # continuous
            identify_1 = max_fn.(X_data[:, max_index])
            X_part_1 = X_data[identify_1, :]
            Y_part_1 = Y_data[identify_1]
            children[max_fn] = create_decision_tree_recursive(X_part_1, Y_part_1, X_type, max_depth-1)
            identify_2 = (!max_fn).(X_data[:, max_index])
            X_part_2 = X_data[identify_2, :]
            Y_part_2 = Y_data[identify_2]
            children[(!max_fn)] = create_decision_tree_recursive(X_part_2, Y_part_2, X_type, max_depth-1)
        end
        return DecisionTree(max_index, children, nothing)
    end

    return create_decision_tree_recursive(X_data, Y_data, detect_types(X_data), max_depth)
end

"""
Prediction function for Decision Tree\\
`X_data` has ndims of 1 or 2\\
Returns an array of predictions\\
Print log when error occurs (undecided children tree)
"""
function predict(X_data::Array{T} where T<:Number, tree::DecisionTree)::Array
    if ndims(X_data) == 1
        X_data = reshape(X_data, (1, size(X_data)[1]))
    end
    @assert ndims(X_data) == 2
    prediction = []
    for i in 1:size(X_data)[1]
        X_vec = X_data[i, :]
        Y_pred = nothing
        tree_copy = tree
        while true
            if tree_copy.target !== nothing
                Y_pred = tree_copy.target
                break
            end
            @assert tree_copy.children !== nothing
            updated = false
            for key_fn in keys(tree_copy.children)
                if key_fn(X_vec[tree_copy.col_id])
                    tree_copy = tree_copy.children[key_fn]
                    updated = true
                    break
                end
            end
            if !updated
                println("Error occured for $X_vec")
                break
            end
        end
        push!(prediction, Y_pred)
    end
    return prediction
end

end


type MSECriterion{A <: AbstractArray, B <: AbstractArray} <: Criterion{A, B}
    output :: B
    gradInput :: A
end

MSECriterion(T :: Type, batch_size, input_size) =
    MSECriterion(Array(T, (batch_size,)), Array(T, (input_size, batch_size)))

function updateOutput!{A,B}(self :: MSECriterion{A,B}, input :: A, target :: A)
    local input_size = size(self.gradInput)[1]
    local batch_size = size(self.gradInput)[2]
    local sum = zero(eltype(A))
    for b = 1:batch_size
        sum = 0
        for i = 1:input_size
            local r = input[i,b] - target[i,b]
            sum += r * r
        end
        sum /= input_size
        self.output[b] = sum
    end
    return self.output
end

function updateGradInput!{A,B}(self :: MSECriterion{A,B},
                               input :: A,
                               target :: A)
    local input_size = size(self.gradInput)[1]
    local batch_size = size(self.gradInput)[2]
    local sum = zero(eltype(A))
    local d = eltype(A)(2) / input_size
    for b = 1:batch_size
        sum = 0
        for i = 1:input_size
            self.gradInput[i,b] = (input[i,b] - target[i,b]) * d
        end
    end
    return self.gradInput
end

export MSECriterion


# Original implementations by @nantonel
# https://github.com/maxmouchet/HMMBase.jl/pull/6

function viterbilog!(
    T1::AbstractMatrix,
    T2::AbstractMatrix,
    z::AbstractVector,
    a::AbstractVector,
    A::AbstractMatrix,
    LL::AbstractMatrix,
)
    T, K = size(LL)
    (T == 0) && return

    fill!(T1, 0.0)
    fill!(T2, 0)

    al = log.(a)
    Al = log.(A)

    for i in OneTo(K)
        T1[1, i] = al[i] + LL[1, i]
    end

    @inbounds for t = 2:T
        for j in OneTo(K)
            amax = 0
            vmax = -Inf

            for i in OneTo(K)
                v = T1[t-1, i] + Al[i, j]
                if v > vmax
                    amax = i
                    vmax = v
                end
            end

            T1[t, j] = vmax + LL[t, j]
            T2[t, j] = amax
        end
    end

    z[T] = argmax(T1[T, :])
    for t = T-1:-1:1
        z[t] = T2[t+1, z[t+1]]
    end
end

"""
    viterbi(a, A, LL) -> Vector

Find the most likely hidden state sequence, see [Viterbi algorithm](https://en.wikipedia.org/wiki/Viterbi_algorithm).
"""
function viterbi(a::AbstractVector, A::AbstractMatrix, LL::AbstractMatrix; logl = nothing)
    ## < v1.1 compatibility
    (logl !== nothing) && deprecate_kwargs("logl")
    (logl == false) && (LL = log.(LL))
    ## --------------------
    T1 = Matrix{Float64}(undef, size(LL))
    T2 = Matrix{Int}(undef, size(LL))
    z = Vector{Int}(undef, size(LL, 1))
    viterbilog!(T1, T2, z, a, A, LL)
    z
end

"""
    viterbi(hmm, observations; robust) -> Vector

**Example**
```julia
using Distributions, HMMBase
hmm = HMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
y = rand(hmm, 1000)
zv = viterbi(hmm, y)
```
"""
function viterbi(hmm::AbstractHMM, observations; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations; robust = robust)
    viterbi(hmm.a, hmm.A, LL)
end


using QuantumCumulants
using OrdinaryDiffEq
using ModelingToolkit
using Test

@testset "correlation" begin

hf = FockSpace(:cavity)
ha = NLevelSpace(:atom,(:g,:e))
h = hf⊗ha

a = Destroy(h,:a)
σ = Transition(h,:σ,:g,:e)

# Single-atom laser
@cnumbers Δ g κ γ ν

H = Δ*a'*a + g*(a'*σ + σ'*a)
J = [a,σ,σ']
he_laser = meanfield([a'*a,σ'*σ,a*σ'],H,J;rates=[κ,γ,ν])

he_avg = cumulant_expansion(he_laser,2)
he_comp = complete(he_avg)

ps = (Δ, g, γ, κ, ν)
sys = ODESystem(he_comp)

# Numerical solution
# p0 = [0.0,0.5,1.0,0.1,0.9]
p0 = ps .=> [1, 1.5, 0.25, 1, 4]
u0 = states(sys) .=> zeros(ComplexF64,length(he_comp))
tmax = 10.0

prob = ODEProblem(sys,u0,(0.0,tmax),p0)
sol = solve(prob,RK4())
n = getindex.(sol.u,1)
pe = getindex.(sol.u,2)

@test all(iszero.(imag.(n)))
@test all(iszero.(imag.(pe)))
@test all(real.(n) .>= 0.0)
@test all(1.0 .>= real.(pe) .>= 0.0)

# Correlation function
c_steady = CorrelationFunction(a', a, he_comp; steady_state=true)
csys = ODESystem(c_steady)

u0_c = correlation_u0(c_steady, sol.u[end])
p0_c = correlation_p0(c_steady, sol.u[end], p0)
# p0_c = (p0..., (c_steady.de0.lhs .=> sol.u[end])...)
τ = range(0.0, 10tmax; length=1001)

prob_c = ODEProblem{true}(csys,u0_c,(0.0,τ[end]),[p0_c...])
sol_c = solve(prob_c,RK4(),saveat=τ,save_idxs=1)

# Spectrum via FFT of g(τ)
using QuantumOptics.timecorrelations: correlation2spectrum
ω, S1 = correlation2spectrum(sol_c.t, sol_c.u)
# using PyPlot; pygui(true)
# plot(ω, S1 ./ maximum(S1), label="FFT of g(τ)")

# Spectrum via Laplace transform
S = Spectrum(c_steady, ps)
usteady = sol.u[end]

S2 = S(ω,usteady,getindex.(p0, 2))
# plot(ω, S2 ./ maximum(S2), label="Laplace transform")
# xlim(-2pi,2pi)

S1_ = S1 ./ maximum(S1)
S1_ .-= minimum(S1_)
S_check = abs.(S2 ./ maximum(S2) .- S1_)
@test maximum(S_check) < 0.02

# Phase invariant case
has_phase(x) = !iszero(phase(x))
import SymbolicUtils
phase(avg::Average) = phase(avg.arguments[1])
phase(op::Destroy) = -1
phase(op::Create) = 1
function phase(t::Transition)
    lvls = QuantumCumulants.levels(t.hilbert, t.aon)
    i = findfirst(isequal(t.i), lvls)
    j = findfirst(isequal(t.j), lvls)
    if i < j
        -1
    elseif i > j
        1
    else
        0
    end
end
phase(op::QuantumCumulants.QMul) = sum(phase(arg) for arg in op.args_nc)
c_nophase = CorrelationFunction(a', a, he_avg; steady_state=true, filter_func=!has_phase)

S_nophase = Spectrum(c_nophase, ps)
S3 = S_nophase(ω,usteady,getindex.(p0, 2))
# plot(ω, S3 ./ maximum(S3), label="Laplace transform (phase invariant)")
# legend()

# Test Mollow triplet
h = NLevelSpace(:atom, (:g,:e))
@cnumbers Δ Ω γ
@qnumbers σ::Transition(h)
H = Δ*σ(:e,:e) + Ω*(σ(:g,:e) + σ(:e,:g))
J = [σ(:g,:e)]
eqs = meanfield([σ(:e,:g),σ(:e,:e)], H, J; rates=[γ])

ps = (Δ,Ω,γ)
p0 = (20.0,5.0,1.0)
u0 = zeros(ComplexF64, 2)
sys = ODESystem(eqs)
prob = ODEProblem(sys,u0,(0.0,20.0),ps .=> p0)
sol = solve(prob,RK4())

@test sol.retcode == :Success

c = CorrelationFunction(σ(:e,:g), σ(:g,:e), eqs; steady_state=true)
csys = ODESystem(c)
cu0 = correlation_u0(c, sol.u[end])
@test length(cu0) == 3
cp0 = correlation_p0(c, sol.u[end], ps .=> p0)
@test length(cp0) == 5

cprob = ODEProblem(csys,cu0,(0.0,20.0),cp0)
csol = solve(cprob, RK4())

@test csol.retcode == :Success

# Mollow when not in steady state
c = CorrelationFunction(σ(:e,:g), σ(:g,:e), eqs; steady_state=false)
csys = ODESystem(c)
cu0 = correlation_u0(c, sol.u[end])
@test length(cu0) == 3
cp0 = correlation_p0(c, sol.u[end], ps .=> p0)
@test length(cp0) == 4

cprob = ODEProblem(csys,cu0,(0.0,20.0),cp0)
csol_ns = solve(cprob, RK4())

@test all(csol_ns.u .≈ csol.u)

# When not in steady state -- cavity that decays
h = FockSpace(:fock)
a = Destroy(h,:a)
@cnumbers ωc κ
H = ωc*a'*a
he = meanfield(a'*a,H,[a];rates=[κ])
ps = (ωc,κ)
sys = ODESystem(he)
n0 = 20.0
u0 = [n0]
p0 = (1,1)
prob = ODEProblem(sys,u0,(0.0,10.0),p0)
sol = solve(prob,RK4())

c = CorrelationFunction(a', a, he)
csys = ODESystem(c)
idx = 5
u0_c = correlation_u0(c, sol.u[idx])
prob_c = ODEProblem(csys,u0_c,(0.0,10.0),p0)
sol_c = solve(prob_c,RK4(),save_idxs=1)
# plot(sol_c.t,real.(sol_c.u), label="Re(g) -- numeric")
# plot(sol_c.t,imag.(sol_c.u), label="Im(g) -- numeric")

gfunc(τ) = @. sol.u[idx] * exp((im*p0[1]-0.5p0[2])*τ)
# plot(sol_c.t, real.(gfunc(sol_c.t)), ls="dashed", label="Re(g) -- analytic")
# plot(sol_c.t, imag.(gfunc(sol_c.t)), ls="dashed", label="Re(g) -- analyitc")
# legend()

@test isapprox(sol_c.u, gfunc(sol_c.t), rtol=1e-4)

end # testset


import FFTW
import Luna: RK45
import Test: @test

function testinit()
    trange = 32
    ωrange = 160
    samples = 2^(ceil(Int, log2(trange*ωrange/2π)))
    δt = trange/samples
    δω = 2π/trange
    n = collect(range(0, length=samples))
    t = @. (n-samples/2)*δt
    ω = @. (n-samples/2)*δω
    N = 5
    z0 = π/2
    zmax = z0*2

    At = @. N*sech(t)
    Aω = FFTW.fftshift(FFTW.fft(At))

    At = zero(Aω)
    nl = zero(Aω)
    nlω = zero(Aω)

    FT = FFTW.plan_fft!(Aω)
    IFT = FFTW.plan_ifft!(Aω)
    f! = let FT=FT, IFT=IFT, ω=ω
        function f!(out, Aω, z)
                out .= Aω
                IFT*out
                out .= abs2.(out).*out
                FT*out
                @. out = im*out - im/2*ω^2*Aω
                return out
        end
    end

    Lin = @. -im/2*ω^2

    linout = similar(Aω)
    Linfunc = let ω=ω
        function Linfunc(out, z)
            @. out = -im/2*ω^2
        end
    end

    fnl! = let FT=FT, IFT=IFT
        function fnl!(out, Aω, z)
                out .= Aω
                IFT*out
                out .= abs2.(out).*out
                FT*out
                @. out = im*out
                return out
        end
    end

    return t, ω, zmax, Aω, f!, Lin, fnl!, Linfunc
end

function test_precon(plot=false)
    t, ω, zmax, Aω, f!, Lin, fnl!, Linfunc = testinit()
    z = 0
    dz = 1e-3

    saveN = 501

    zarr, Aarr = RK45.solve_precon(fnl!, Lin, Aω, z, dz, zmax, saveN, rtol=1e-6, status_period=2)

    if plot
        Atarr = FFTW.ifft(FFTW.ifftshift(Aarr, 1), 1)
        energy = dropdims(sum(abs2.(Atarr), dims=1), dims=1)
        pygui(true)
        plt.figure()
        plt.pcolormesh(t, zarr, abs2.(transpose(Atarr)))
        plt.colorbar()
        plt.figure()
        plt.pcolormesh(ω, zarr, abs2.(transpose(Aarr)))
        plt.figure()
        plt.plot(zarr, 1 .- energy/energy[1])
        plt.figure()
        plt.plot(t, abs2.(Atarr[:, 1]))
        plt.plot(t, abs2.(Atarr[:, end]))
    end

    return zarr, Aarr
end

function test_noprecon(plot=false)
    t, ω, zmax, Aω, f!, Lin, fnl!, Linfunc = testinit()
    z = 0
    dz = 1e-3

    saveN = 201

    zarr, Aarr = RK45.solve(f!, Aω, z, dz, zmax, saveN)

    if plot
        Atarr = FFTW.ifft(FFTW.ifftshift(Aarr, 1), 1)
        energy = dropdims(sum(abs2.(Atarr), dims=1), dims=1)
        pygui(true)
        plt.figure()
        plt.pcolormesh(t, zarr, abs2.(transpose(Atarr)))
        plt.colorbar()
        plt.figure()
        plt.pcolormesh(ω, zarr, abs2.(transpose(Aarr)))
        plt.figure()
        plt.plot(zarr, 1 .- energy/energy[1])
        plt.figure()
        plt.plot(t, abs2.(Atarr[:, 1]))
        plt.plot(t, abs2.(Atarr[:, end]))
    end

    return zarr, Aarr
end

t, ω, zmax, Aω, f!, Lin, fnl!, Linfunc = testinit()
z = 0
dz = 1e-3

zarr, Aarr = RK45.solve(f!, copy(Aω), z, dz, zmax, rtol=1e-8, output=true, outputN=501)
zarrp, Aarrp = RK45.solve_precon(fnl!, Lin, copy(Aω), z, dz, zmax,
                                 rtol=1e-8, output=true, outputN=501)
zarrpf, Aarrpf = RK45.solve_precon(fnl!, Linfunc, copy(Aω), z, dz, zmax,
                                   rtol=1e-8, output=true, outputN=501)
# Is the initial spectrum restored after 2 soliton periods? (without preconditioner)
@test isapprox(abs2.(Aarr[:, 1]), abs2.(Aarr[:, end]), rtol=1e-4)
# Is the initial spectrum restored after 2 soliton periods? (with preconditioner)
@test isapprox(abs2.(Aarrp[:, 1]), abs2.(Aarrp[:, end]), rtol=1e-3)
# Is the initial spectrum restored after 2 soliton periods?
# (with preconditioner and z-dependent linear part)
@test isapprox(abs2.(Aarrpf[:, 1]), abs2.(Aarrpf[:, end]), rtol=1e-3)
# Is there a difference if the linear part is a function (but constant)?
@test all(abs2.(Aarrp) .== abs2.(Aarrpf))


function SeisSmooth2(A;L1=7,L2=7,Nrepeat=1)
    a1 = isodd(L1) ? 0 : 1
    a2 = isodd(L2) ? 0 : 1
    kernel = ones(eltype(A),L1,L2)/(L1*L2);
    for irepeat = 1 : Nrepeat
	A = conv2(A,kernel)
	A = A[floor(Int,L1/2) + 1:end-floor(Int,L1/2) + a1,floor(Int,L2/2) + 1:end-floor(Int,L2/2) + a2]
	end
    return A
end

function SeisSmooth2(A,h::Array{Header,1};L1=7,L2=7)
    A = SeisSmooth2(A,L1=L1,L2=L2,Nrepeat=Nrepeat)
    return A,h
end

function SeisSmooth2(in::String,out::String;key=["iaz";"iang"],L1=7,L2=7,Nrepeat=1)
    d, h, ext = SeisRead(in)
    nang = size(d,3)
    for iang = 1 : nang
	A = d[:,1,iang,1,:]
	A = SeisSmooth2(A,L1=L1,L2=L2,Nrepeat=Nrepeat)
	d[:,1,iang,1,:] = A
    end
    SeisWrite(out,d,h,ext)

    #@compat parameters = Dict(:L1=>L1,:L2=>L2)
    #SeisProcess(in,out,[SeisSmooth2],[parameters];group="gather",key=key)
end


"""
    numericShift!(hh, qq, iq, qrows, qcols, neq, condn)

Compute the numeric shiftrights and store them in q.
"""
function numericShift!(hh::Array{Float64,2}, qq::Array{Float64,2}, iq::Int64, qRows::Int64, qCols::Int64, neq::Int64, condn::Float64)

    # total number of shifts
    nnumeric = 0

    # functions to seperate hh
    left = 1:qCols
    right = (qCols + 1):(qCols + neq)

    # preform QR factorization on right side of hh
    F = Compat.qr(hh[:, right], Val(true))

    # filter R only keeping rows that are zero
    Q, R = F
    zerorows = abs.(diag(R))
    zerorows = findall(x->(float(x) <= condn), zerorows)

    while (length(zerorows) != 0) && (iq <= qRows)
        # update hh with matrix multiplication of Q and hh
        hh = *(Q', hh)

        # need total number of zero rows
        nz = size(zerorows, 1)

        # update qq to keep track of rows that are zero
        qq[(iq + 1):(iq + nz), :] = hh[zerorows, left]

        # update hh by shifting right
        hh[zerorows, :] = shiftRight!(hh[zerorows, :], neq)

        # increment our variables by number of shifts
        iq = iq + nz
        nnumeric = nnumeric + nz

        # redo QR factorization and filter R as before
        F = Compat.qr(hh[:, right], Val(true))
	Q, R = F
        zerorows = abs.(diag(R))
        zerorows = findall(x->(float(x) <= condn), zerorows)

    end # while

    return(hh, qq, iq, nnumeric)

end # numericShift