"""
    ShallowWater{T<:Real,Dim}

The shallow water equations with gravitational constant `g`
in `Dim` dimensions using `T` as scalar type.
"""
struct ShallowWater{T,Dim} <: AbstractBalanceLaw{1}
  g::T
end

function ShallowWater(g=1., Dim=1)
  ShallowWater{typeof(g), Dim}(g)
end

function Base.show(io::IO, model::ShallowWater{T,Dim}) where {T,Dim}
  print(io, "Shallow water equations with g=", model.g, " {T=", T, ", Dim=", Dim, "}")
end

"""
Conserved variables (h, hv) of the shallow water system in one space dimension.
"""
struct ShallowWaterVar1D{T} <: FieldVector{2,T}
  h ::T
  hv::T
end

function (::Type{ShallowWaterVar1D{T}})(val::Real) where {T}
  ShallowWaterVar1D{T}(val, val)
end

function similar_type(::ShallowWaterVar1D{T}) where {T}
  ShallowWaterVar1D{T}
end

@inline variables(model::ShallowWater{T,1}) where {T} = ShallowWaterVar1D{T}

@inline function primitive_variables(u::ShallowWaterVar1D)
  @unpack h, hv = u
  if h ≈ 0
    v = zero(h)
  else
    v = hv/h
  end
  h, v
end

@inline function primitive_variables(u::ShallowWaterVar1D, model::ShallowWater)
  primitive_variables(u)
end

@inline function entropy_variables(u::ShallowWaterVar1D, model::ShallowWater)
    entropy_variables(primitive_variables(u, model)..., model)
end

@inline function entropy_variables(h, v, model::ShallowWater)
    @unpack g = model
    SVector(g*h - v^2/2, v)
end

@inline function flux_potential(u::ShallowWaterVar1D, model::ShallowWater)
    flux_potential(primitive_variables(u, model)..., model)
end

@inline function flux_potential(h, v, model::ShallowWater)
    @unpack g = model
    g*h^2*v/2
end

@inline function flux(u::ShallowWaterVar1D{T}, model::ShallowWater{T,1}) where {T}
  h, v = primitive_variables(u)
  @unpack g = model

  SVector{2,T}(h*v, h*v^2 + g*h^2/2)
end

@inline function max_abs_speed(u::ShallowWaterVar1D, model::ShallowWater{T,1}) where {T}
  h, v = primitive_variables(u)
  @unpack g = model

  abs(v) + sqrt(g*h)
end

"""
    max_abs_speed(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)

Estimate the maximal speed in the solution of the Riemann problem with left and
right states `uₗ`, `uᵣ` for `model` via the two rarefaction approximation, see
Guermond and Popov (2016) Fast Estimation from above of the maximum wave speed in
the Riemann roblem for the Euler equations, Journal of Computational Physics 321,
pp. 908-926, and
Chen and Shu (2017) Entropy stable high order discontinuous Galerkin methods with
suitable quadrature rules for hyperbolic conservation laws, Journal of Computational
Physics 345, pp. 427-461.
"""
@inline function max_abs_speed(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)
    λ₋, λ₊ = min_max_speed(uₗ, uᵣ, model)

    max(abs(λ₋), abs(λ₊))
end

"""
    min_max_speed(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)

Estimate the maximal speed in the solution of the Riemann problem with left and
right states `uₗ`, `uᵣ` for `model` via the two rarefaction approximation, see
Guermond and Popov (2016) Fast Estimation from above of the maximum wave speed in
the Riemann roblem for the Euler equations, Journal of Computational Physics 321,
pp. 908-926, and
Chen and Shu (2017) Entropy stable high order discontinuous Galerkin methods with
suitable quadrature rules for hyperbolic conservation laws, Journal of Computational
Physics 345, pp. 427-461.
"""
@inline function min_max_speed(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)
    @unpack g = model
    hₗ, vₗ = primitive_variables(uₗ, model)
    aₗ = sqrt(g * hₗ)
    hᵣ, vᵣ = primitive_variables(uᵣ, model)
    aᵣ = sqrt(g * hᵣ)

    hₘ = (2*(aᵣ+aₗ) - (vᵣ-vₗ))^2 / 16g
    if hₘ <= hₗ
        # approximated by rarefaction wave
        λ₋ = vₗ - aₗ
    else
        # approximated by shock wave
        λ₋ = vₗ - hₘ * sqrt( g*(hₘ+hₗ) / (2*hₘ*hₗ) )
    end
    if hₘ <= hᵣ
        # approximated by rarefaction wave
        λ₊ = vᵣ + aᵣ
    else
        # approximated by shock wave
        λ₊ = vᵣ + hₘ * sqrt( g*(hₘ+hᵣ) / (2*hₘ*hᵣ) )
    end

    λ₋, λ₊
end

"""
    (::EnergyConservativeFlux1Param)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)

The one-parameter family of energy conservative numerical fluxes for the shallow
water equations, see
Ranocha (2017) Shallow water equations: Split-form, entropy stable, well-balanced,
and positivity-preserving numerical methods, GEM - International Journal on
Geomathematics 8(1), pp. 85-133.
"""
function (fvol::EnergyConservativeFlux1Param)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D,
                                              model::ShallowWater)
    h₋, v₋ = primitive_variables(uₗ, model)
    h₊, v₊ = primitive_variables(uᵣ, model)
    @unpack g = model
    @unpack a₁ = fvol

    fnum_h  = ( (3-a₁)*(h₋*v₋ + h₊*v₊) + (1+a₁)*(h₊*v₋ + h₋*v₊) ) / 8
    fnum_hv = g * ( (1+a₁)*(h₊^2 + h₋^2) + (2-2a₁)*(h₊*h₋) ) / 8 +
                ( (3-a₁)*(h₊*v₊^2 + h₋*v₋^2) + (1+a₁)*(h₊*v₋^2 + h₋*v₊^2) ) / 16 +
                (h₊*v₊*v₋ + h₋*v₊*v₋) / 4

    SVector(fnum_h, fnum_hv)
end

"""
    (::EnergyConservativeFlux2Param)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)

The two-parameter family of energy conservative numerical fluxes for the shallow
water equations, see
Ranocha (2017) Shallow water equations: Split-form, entropy stable, well-balanced,
and positivity-preserving numerical methods, GEM - International Journal on
Geomathematics 8(1), pp. 85-133.
"""
function (fvol::EnergyConservativeFlux2Param)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D,
                                              model::ShallowWater)
    h₋, v₋ = primitive_variables(uₗ, model)
    h₊, v₊ = primitive_variables(uᵣ, model)
    @unpack g = model
    @unpack a₁, a₂ = fvol

    fnum_h  = ( (3-a₁)*(h₋*v₋ + h₊*v₊) + (1+a₁)*(h₊*v₋ + h₋*v₊) ) / 8 +
                (a₁+3a₂-2) * (v₊^3 - v₊^2*v₋ - v₊*v₋^2 + v₋^3) / 16g
    fnum_hv = g * ( (1+a₁)*(h₊^2 + h₋^2) + (2-2a₁)*(h₊*h₋) ) / 8 -
                (2a₁+3a₂-5) * (h₊*v₊^2 + h₋*v₋^2) / 16 +
                (2a₁+3a₂-1) * (h₊*v₋^2 + h₋*v₊^2) / 16 +
                (h₊*v₊*v₋ + h₋*v₊*v₋) / 4 +
                (a₁+3a₂-2) * ((v₊^2)^2 - 2*v₊^2*v₋^2 + (v₋^2)^2)

    SVector(fnum_h, fnum_hv)
end

"""
    (diss::ScalarDissipation)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)

The scalar dissipation operator \$- \\frac{\\lambda}{2} R \\cdot R^T \\cdot  (w_r - w_l)\$
for the shallow water equations, where \$\\partial_w u = R \\cdot R^T\$ is evaluated
at the arithmetic mean values \$h = (h_- + h_+) / 2\$, \$v = (v_- + v_+) / 2\$.
"""
function (diss::ScalarDissipation)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)
    h₋, v₋ = primitive_variables(uₗ, model)
    h₊, v₊ = primitive_variables(uᵣ, model)
    @unpack g = model

    w₋ = entropy_variables(h₋, v₋, model)
    w₊ = entropy_variables(h₊, v₊, model)

    h = (h₋ + h₊) / 2
    v = (v₋ + v₊) / 2
    du_dw = SArray{Tuple{2,2}}(1, v, v, g*h+v^2) / g

    λ = diss.max_abs_speed(uₗ, uᵣ, model)

    -λ/2 * du_dw * (w₊ - w₋)
end

"""
    (diss::MatrixDissipation)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)

The scalar dissipation operator \$- \\frac{1}{2} R \\cdot |\\Lambda| \\cdot R^T \\cdot  (w_r - w_l)\$
for the shallow water equations, where \$f'(u) = R \\cdot \\Lambda \\cdot R^{-1}\$ is
evaluated at the arithmetic mean values \$h = (h_- + h_+) / 2\$, \$v = (v_- + v_+) / 2\$.
"""
function (diss::MatrixDissipation)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater)
    h₋, v₋ = primitive_variables(uₗ, model)
    h₊, v₊ = primitive_variables(uᵣ, model)
    @unpack g = model

    w₋ = entropy_variables(h₋, v₋, model)
    w₊ = entropy_variables(h₊, v₊, model)

    h = (h₋ + h₊) / 2
    v = (v₋ + v₊) / 2
    sqrt_gh = sqrt(g*h)
    R = SArray{Tuple{2,2}}(1, v-sqrt_gh, 1, v+sqrt_gh) / sqrt(2g)
    Λ = SArray{Tuple{2,2}}(abs(v-sqrt_gh), 0, 0, abs(v+sqrt_gh))

    -R * Λ * R' * (w₊ - w₋) / 2
end

"""
    (::SuliciuFlux)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater{T,1})

The Suliciu relaxation solver / numerical flux, see
Bouchut (2004) Nonlinear Stability of Finite Volume Methods for Hyperbolic
Conservation Laws and Well-Balanced Schemes for Sources.
"""
function (::SuliciuFlux)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater{T,1}) where T
    hl, vl = primitive_variables(uₗ, model)
    hr, vr = primitive_variables(uᵣ, model)
    @unpack g = model

    sqrt_ghl = sqrt(g*hl)
    sqrt_ghr = sqrt(g*hr)

    # compute speeds
    cl_hl = zero(T)
    cr_hr = zero(T)
    if hl <= hr && 0 < hr
      cl_hl = sqrt_ghl + 1.5*max(0, 0.5*g*(hr*hr-hl*hl)/(hr*sqrt_ghr) + vl - vr )
      cr_hr = sqrt_ghr + 1.5*max(0, 0.5*g*(hl*hl-hr*hr)/(cl_hl*hl) + vl - vr )
    elseif hr <= hl && 0 < hl
      cr_hr = sqrt_ghr + 1.5*max(0, 0.5*g*(hl*hl-hr*hr)/(hl*sqrt_ghl) + vl - vr )
      cl_hl = sqrt_ghl + 1.5*max(0, 0.5*g*(hr*hr-hl*hl)/(cr_hr*hr) + vl - vr )
    else
      cl_hl = cr_hr = zero(T)
    end

    # compute intermediate values
    cl = cl_hl*hl
    cr = cr_hr*hr
    vls = vrs = ( cl*vl+cr*vr+0.5*g*(hl*hl-hr*hr) ) / ( cl+cr )
    vls = vrs = ifelse(isnan(vls), zero(T), vls)
    πls = πrs = ( 0.5*g*(cr*hl*hl+cl*hr*hr)-cl*cr*(vr-vl) ) / (cl+cr)
    πls = πrs = ifelse(isnan(πls), zero(T), πls)
    hls = 1 / ( 1/hl + (cr*(vr-vl)+0.5*g*(hl*hl-hr*hr))/(cl*(cl+cr)) )
    hls = ifelse(isnan(hls), zero(T), hls)
    hrs = 1 / ( 1/hr + (cl*(vr-vl)+0.5*g*(hr*hr-hl*hl))/(cr*(cl+cr)) )
    hrs = ifelse(isnan(hrs), zero(T), hrs)

    # compute fluxes
    fnum_h  = zero(T)
    fnum_hv = zero(T)
    if 0 <= vl-cl_hl
      fnum_h  = hl*vl
      fnum_hv = hl*vl*vl + 0.5*g*hl*hl
    elseif 0 <= vls
      fnum_h  = hls*vls
      fnum_hv = hls*vls*vls + πls
    elseif 0 <= vr+cr_hr
      fnum_h  = hrs*vrs
      fnum_hv = hrs*vrs*vrs + πrs
    else
      fnum_h  = hr*vr
      fnum_hv = hr*vr*vr + 0.5*g*hr*hr
    end

    SVector(fnum_h, fnum_hv)
end

"""
    (::KineticFlux)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater{T,1})

The kinetic relaxation solver / numerical flux, see
Perthame and Simeoni (2001) A kinetic scheme for the Saint-Venant system with a
source term, Calcolo 38(4), pp. 201-231.
"""
function (::KineticFlux)(uₗ::ShallowWaterVar1D, uᵣ::ShallowWaterVar1D, model::ShallowWater{T,1}) where T
    hₗ, vₗ = primitive_variables(uₗ, model)
    hᵣ, vᵣ = primitive_variables(uᵣ, model)
    @unpack g = model

    # integrals for \xi \geq 0
    fnum_h, fnum_hv = kinetic_integrals_ξ_geq_0(hₗ, vₗ, g)
    # use symmetry for \xi \leq 0
    val_h, val_hv   = kinetic_integrals_ξ_geq_0(hᵣ, -vᵣ, g)
    fnum_h  -= val_h
    fnum_hv += val_hv

    SVector(fnum_h, fnum_hv)
end

# \int_{\xi \geq 0} (\xi, xi^2)  M_l(\xi) \dif \xi  for (fnum_h, fnum_hv)
function kinetic_integrals_ξ_geq_0(h, v, g)
    val_h  = zero(h)
    val_hv = zero(h)
    sqrt_2gh = sqrt(2g*h)

    if h > 0 && v >= sqrt_2gh
        val_h += h*v

        val_hv += h*v*v + g*h*h/2
    elseif h > 0 && sqrt_2gh+v > 0 && sqrt_2gh-v > 0
        sqrt_2ghmv2 = sqrt(2g*h-v^2)
        atan_v_sqrt_2ghmv2 = atan(v/sqrt_2ghmv2)

        val_h += h * v / 2
        val_h += h * sqrt_2ghmv2 * 2 / (3π)
        val_h += v^2 * sqrt_2ghmv2 / (6g*π)
        val_h += h*v * atan_v_sqrt_2ghmv2 / π

        val_hv += ( h*v*v + g*h*h/2 ) / 2
        val_hv += h*v*sqrt_2ghmv2 * 13 / (12π)
        val_hv += v^3*sqrt_2ghmv2 / (12g*π)
        val_hv += ( h*v^2 + g*h^2/2 ) * atan_v_sqrt_2ghmv2 / π
    end

    val_h, val_hv
end

################################################################################

struct ShallowWaterRiemannSolution{T,T1} <: AbstractRiemannSolution
  prob::RiemannProblem{ShallowWater{T,1},ShallowWaterVar1D{T},T1}
  uₘ::ShallowWaterVar1D{T} # middle state
  σ₁⁻::T # slow speed of first family
  σ₁⁺::T # fast speed of first family
  σ₂⁻::T # slow speed of second family
  σ₂⁺::T # fast speed of second family

  function ShallowWaterRiemannSolution{T,T1}(prob::RiemannProblem{ShallowWater{T,1},ShallowWaterVar1D{T},T1}) where {T,T1}
    uₘ, σ₁⁻, σ₁⁺, σ₂⁻, σ₂⁺ = compute_state_and_speeds(prob.uₗ, prob.uᵣ, prob.model)
    new(prob, uₘ, σ₁⁻, σ₁⁺, σ₂⁻, σ₂⁺)
  end
end

function solve(prob::RiemannProblem{ShallowWater{T,1},ShallowWaterVar1D{T},T1}) where {T,T1}
  ShallowWaterRiemannSolution{T,T1}(prob)
end

function minmax_speeds(sol::ShallowWaterRiemannSolution)
  sol.σ₁⁻, sol.σ₂⁺
end

function (sol::ShallowWaterRiemannSolution)(ξ::Real)
  @unpack uₘ, σ₁⁻, σ₁⁺, σ₂⁻, σ₂⁺ = sol
  @unpack uₗ, uᵣ = sol.prob
  @unpack g = sol.prob.model

  if ξ <= σ₁⁻
    uₗ
  elseif ξ <= σ₁⁺
    rarefaction_wave_1(ξ, uₗ, g)
  elseif ξ <= σ₂⁻
    uₘ
  elseif ξ <= σ₂⁺
    rarefaction_wave_2(ξ, uᵣ, g)
  else # σ₂⁺ < ξ
    uᵣ
  end
end

function rarefaction_wave_1(ξ, uₗ::ShallowWaterVar1D{T}, g) where {T}
  hₗ, vₗ = primitive_variables(uₗ)
  h  = (vₗ + 2*sqrt(g*hₗ) -  ξ)^2 / 9g
  hv = (vₗ + 2*sqrt(g*hₗ) + 2ξ)*h / 3
  ShallowWaterVar1D{T}(h, hv)
end

function rarefaction_wave_2(ξ, uₗ::ShallowWaterVar1D{T}, g) where {T}
  hₗ, vₗ = primitive_variables(uₗ)
  h  = (vₗ - 2*sqrt(g*hₗ) -  ξ)^2 / 9g
  hv = (vₗ - 2*sqrt(g*hₗ) + 2ξ)*h / 3
  ShallowWaterVar1D{T}(h, hv)
end

function (sol::ShallowWaterRiemannSolution)(t::Real, x::Real)
  @unpack x₀, t₀ = sol.prob

  sol((x-x₀)/(t-t₀))
end

function compute_state_and_speeds(uₗ::ShallowWaterVar1D{T}, uᵣ::ShallowWaterVar1D{T}, model::ShallowWater) where {T}
  @unpack g = model
  hₗ, vₗ = primitive_variables(uₗ)
  hᵣ, vᵣ = primitive_variables(uᵣ)

  if v_S1_SWE(hᵣ,hₗ,vₗ,g) <= vᵣ && vᵣ <= v_R2_SWE(hᵣ,hₗ,vₗ,g)
    # Case I: wave1 = shock, wave2 = rarefaction
    if hₗ ≈ 0
      hₘ = hₗ
      vₘ = v_R2_SWE(hₘ,hᵣ,vᵣ,g)
      uₘ = ShallowWaterVar1D{T}(hₘ, hₘ*vₘ)
      σ₁⁻ = σ₁⁺ = vₘ - sqrt(g*hₗ)
      σ₂⁻ = vₘ + sqrt(g*hₘ)
      σ₂⁺ = vᵣ + sqrt(g*hᵣ)
    else
      hₘ = Roots.fzero(h -> vᵣ-v_R2_SWE(hᵣ,h,v_S1_SWE(h,hₗ,vₗ,g),g), hₗ, hᵣ)
      # this is better than vₘ = v_S1_SWE(hₘ,hₗ,vₗ) for hₗ≈0
      vₘ = v_R2_SWE(hₘ,hᵣ,vᵣ,g)
      uₘ = ShallowWaterVar1D{T}(hₘ, hₘ*vₘ)
      # this is better than σ₁⁻ = σ₁⁺ = vₗ - sqrt(hₘ+hₘ^2/hₗ) / sqrt(2) for hₗ≈0
      σ₁⁻ = σ₁⁺ = vₘ - sqrt(hₗ+hₗ^2/hₘ) * sqrt(g/2)
      σ₂⁻ = vₘ + sqrt(g*hₘ)
      σ₂⁺ = vᵣ + sqrt(g*hᵣ)
    end
  elseif v_R1_SWE(hᵣ,hₗ,vₗ,g) <= vᵣ && v_R2_SWE(hᵣ,hₗ,vₗ,g) <= vᵣ
    # Case II: wave1 = rarefaction, wave2 = rarefaction
    if 2*(sqrt(g*hᵣ)+sqrt(g*hₗ)) >= vᵣ-vₗ
      # no vacuum
      hₘ = (2*(sqrt(g*hᵣ)+sqrt(g*hₗ)) -(vᵣ-vₗ))^2 / 16g
      vₘ = v_R1_SWE(hₘ,hₗ,vₗ,g)
      uₘ = ShallowWaterVar1D{T}(hₘ,hₘ*vₘ)
      σ₁⁻ = vₗ - sqrt(g*hₗ)
      σ₁⁺ = vₘ - sqrt(g*hₘ)
      σ₂⁻ = vₘ + sqrt(g*hₘ)
      σ₂⁺ = vᵣ + sqrt(g*hᵣ)
    else
      # vacuum
      uₘ = ShallowWaterVar1D{T}(0,0)
      σ₁⁻ = vₗ - sqrt(g*hₗ)
      σ₁⁺ = vₗ + 2*sqrt(g*hₗ)
      σ₂⁻ = vᵣ - 2*sqrt(g*hᵣ)
      σ₂⁺ = vᵣ + sqrt(g*hᵣ)
    end
  elseif v_S2_SWE(hᵣ,hₗ,vₗ,g) <= vᵣ && vᵣ <= v_R1_SWE(hᵣ,hₗ,vₗ,g)
    # Case III: wave1 = rarefaction, wave2 = shock
    if hᵣ ≈ 0
      hₘ = hᵣ
      vₘ = v_R1_SWE(hₘ,hₗ,vₗ,g)
      uₘ = ShallowWaterVar1D{T}(hₘ,hₘ*vₘ)
      σ₁⁻ = vₗ - sqrt(g*hₗ)
      σ₁⁺ = vₘ - sqrt(g*hₘ)
      σ₂⁻ = σ₂⁺ = vₘ + sqrt(g*hᵣ)
    else
      hₘ = Roots.fzero(h->vᵣ-v_S2_SWE(hᵣ,h,v_R1_SWE(h,hₗ,vₗ,g),g), hᵣ, hₗ)
      vₘ = v_R1_SWE(hₘ,hₗ,vₗ,g)
      uₘ = ShallowWaterVar1D{T}(hₘ,hₘ*vₘ)
      σ₁⁻ = vₗ - sqrt(g*hₗ)
      σ₁⁺ = vₘ - sqrt(g*hₘ)
      σ₂⁻ = σ₂⁺ = vₘ + sqrt(hᵣ+hᵣ^2/hₘ) * sqrt(g/2)
    end
  else #vᵣ <= v_S1_SWE(hᵣ,hₗ,vₗ) && vᵣ <= v_S2_SWE(hᵣ,hₗ,vₗ)
    # Case IV: wave1 = shock, wave2 = shock
    # TODO: Chose some interval?
    hₘ = Roots.fzero(h->vᵣ-v_S2_SWE(hᵣ,h,v_S1_SWE(h,hₗ,vₗ,g),g), max(hₗ,hᵣ))
    vₘ = v_S1_SWE(hₘ,hₗ,vₗ,g)
    uₘ = ShallowWaterVar1D{T}(hₘ,hₘ*vₘ)
    σ₁⁻ = σ₁⁺ = vₗ - hₘ * sqrt(1/hₘ+1/hₗ) * sqrt(g/2)
    σ₂⁻ = σ₂⁺ = vₘ + hᵣ * sqrt(1/hᵣ+1/hₘ) * sqrt(g/2)
  end

  uₘ, σ₁⁻, σ₁⁺, σ₂⁻, σ₂⁺
end

# wave curves parameterised by the height h
v_R1_SWE(h, hₗ, vₗ, g) = vₗ - 2*(sqrt(g*h) - sqrt(g*hₗ))
v_R2_SWE(h, hₗ, vₗ, g) = vₗ + 2*(sqrt(g*h) - sqrt(g*hₗ))
v_S1_SWE(h, hₗ, vₗ, g) = vₗ - (h-hₗ) * sqrt( (1/h+1/hₗ)*g/2 )
v_S2_SWE(h, hₗ, vₗ, g) = vₗ + (h-hₗ) * sqrt( (1/h+1/hₗ)*g/2 )

@recipe function f(ξu::Tuple{Ξ,Vector{ShallowWaterVar1D{T}}}) where {Ξ,T}
  ξ, u = ξu

  h  = mappedarray(u->u.h, u)
  hv = mappedarray(u->u.hv, u)

  size --> (1000, 400)
  layout --> (1,2)
  legend --> false

  @series begin
    subplot := 1
    ylabel --> L"h"
    label  --> L"h"
    ξ, h
  end

  @series begin
    subplot := 2
    ylabel --> L"hv"
    label  --> L"hv"
    ξ, hv
  end
end


#=

  Survo puzzle in Julia + ConstraintSolver.jl

  http://en.wikipedia.org/wiki/Survo_Puzzle
  """
  Survo puzzle is a kind of logic puzzle presented (in April 2006) and studied
  by Seppo Mustonen. The name of the puzzle is associated to Mustonen's
  Survo system which is a general environment for statistical computing and
  related areas.

  In a Survo puzzle the task is to fill an m * n table by integers 1,2,...,m*n so
  that each of these numbers appears only once and their row and column sums are
  equal to integers given on the bottom and the right side of the table.
  Often some of the integers are given readily in the table in order to
  guarantee uniqueness of the solution and/or for making the task easier.
  """

  See also
  http://www.survo.fi/english/index.html
  http://www.survo.fi/puzzles/index.html

  References:
  Mustonen, S. (2006b). "On certain cross sum puzzles"
  http://www.survo.fi/papers/puzzles.pdf
  Mustonen, S. (2007b). "Enumeration of uniquely solvable open Survo puzzles."
  http://www.survo.fi/papers/enum_survo_puzzles.pdf
  Kimmo Vehkalahti: "Some comments on magic squares and Survo puzzles"
  http://www.helsinki.fi/~kvehkala/Kimmo_Vehkalahti_Windsor.pdf
  R code: http://koti.mbnet.fi/tuimala/tiedostot/survo.R

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

=#
using ConstraintSolver, JuMP
using Cbc
const CS = ConstraintSolver
include("constraints_utils.jl")
include("survo_puzzle_problems.jl")

function survo_puzzle(p,print_solution=true,all_solutions=false,timeout=Inf)
    row_sums = p[:row_sums]
    col_sums = p[:col_sums]
    println("row_sums: $row_sums col_sums:$col_sums")
    grid = resize_matrix(p[:problem])
    rows,cols = size(grid)  # 9

    m = Model(optimizer_with_attributes(CS.Optimizer, "all_solutions"=>all_solutions,"logging"=>[],
                                        # "traverse_strategy"=>:BFS, # Seems to be quite fast!
                                        # "traverse_strategy"=>:DFS,
                                        # "traverse_strategy"=>:DBFS,

                                        # "branch_split"=>:Smallest, # Seems to be the best
                                        # "branch_split"=>:Biggest,
                                        # "branch_split"=>:InHalf,

                                        # "simplify"=>false,
                                        # "simplify"=>true,
                                        "time_limit"=>timeout
                                        ))

    # define the variables
    @variable(m, 1 <= x[1:rows,1:cols] <= rows*cols, Int)

    # hints
    for r=1:rows, c=1:cols
        if grid[r,c] != 0
            @constraint(m, x[r,c] == grid[r,c])
        end
    end

    @constraint(m, vec(x[:,:]) in CS.AllDifferentSet())

    # Row and column constraints
    for r = 1:rows
        @constraint(m, row_sums[r] == sum(x[r,:]))
    end
    for c = 1:cols
        @constraint(m, col_sums[c] == sum(x[:,c]))
    end

    # Solve the problem
    optimize!(m)

    status = JuMP.termination_status(m)
    println("status:$status")
    if status == MOI.OPTIMAL
        x = convert.(Integer,JuMP.value.(x))
        if print_solution
            # display(@show x)
            print_grid(x)
        end
        num_sols = MOI.get(m, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        return x
    end
end

print_solution=true
all_solutions=true
timeout=10

for p in sort(collect(keys(survo_puzzle_problems)))
    println("\nproblem $p")
    @time survo_puzzle(survo_puzzle_problems[p],print_solutions,all_solutions,timeout)

    println("\n")
end


using DiffEqFlux, OrdinaryDiffEq, Test

u0 = Float32[2.; 0.]
datasize = 30
tspan = (0.0f0,1.5f0)

function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
t = range(tspan[1],tspan[2],length=datasize)
prob = ODEProblem(trueODEfunc,u0,tspan)
ode_data = Array(solve(prob,Tsit5(),saveat=t))

model = Chain(x -> x.^3,
             Dense(2,50,tanh),
             Dense(50,2))
neuralde = NeuralODE(model,tspan,Rodas5(),saveat=t,reltol=1e-7,abstol=1e-9)

function predict_n_ode()
  neuralde(u0)
end
loss_n_ode() = sum(abs2,ode_data .- predict_n_ode())

data = Iterators.repeated((), 10)
opt = ADAM(0.1)
cb = function () #callback function to observe training
  display(loss_n_ode())
end

# Display the ODE with the initial parameter values.
cb()

neuralde = NeuralODE(model,tspan,Rodas5(),saveat=t,reltol=1e-7,abstol=1e-9)
ps = Flux.params(neuralde)
loss1 = loss_n_ode()
Flux.train!(loss_n_ode, ps, data, opt, cb = cb)
loss2 = loss_n_ode()
@test loss2 < loss1

neuralde = NeuralODE(model,tspan,KenCarp4(),saveat=t,reltol=1e-7,abstol=1e-9)
ps = Flux.params(neuralde)
loss1 = loss_n_ode()
Flux.train!(loss_n_ode, ps, data, opt, cb = cb)
loss2 = loss_n_ode()
@test loss2 < loss1

neuralde = NeuralODE(model,tspan,RadauIIA5(),saveat=t,reltol=1e-7,abstol=1e-9)
ps = Flux.params(neuralde)
loss1 = loss_n_ode()
Flux.train!(loss_n_ode, ps, data, opt, cb = cb)
loss2 = loss_n_ode()
@test loss2 < loss1


abstract type MCDMResult end
abstract type SCDMResult end
abstract type MCDMMethod end

"""
    struct MCDMSetting
        df::DataFrame
        weights::Array{Float64, 1}
        fns::Array{Function, 1}
    end

    Immutable data structure for a MCDM setting.

# Arguments
- `df::DataFrame`: The decision matrix in type of DataFrame.
- `weights::Array{Float64,1}`: Array of weights for each criterion.
- `fns::Array{Function, 1}`: Array of functions. The elements are either minimum or maximum.

# Description
Many methods including Topsis, Electre, Waspas, etc., use a decision matrix, weights, and directions
of optimizations in types of DataFrame, Vector, and Vector, respectively. The type MCDMSetting simply
holds these information to pass them into methods easly. Once a MCDMSetting object is created, the problem
can be passed into several methods like topsis(setting), electre(setting), waspas(setting), etc.

# Examples

```julia-repl
julia> df = DataFrame();
julia> df[:, :x] = Float64[9, 8, 7];
julia> df[:, :y] = Float64[7, 7, 8];
julia> df[:, :z] = Float64[6, 9, 6];
julia> df[:, :q] = Float64[7, 6, 6];

julia> w = Float64[4, 2, 6, 8];

julia> fns = makeminmax([maximum, maximum, maximum, maximum]);

julia> setting = MCDMSetting(df, w, fns)

julia> result = topsis(setting);
julia> # Same result can be obtained using
julia> result2 = mcdm(setting, TopsisMethod())
```

"""
struct MCDMSetting
    df::DataFrame
    weights::Array{Float64, 1}
    fns::Array{Function, 1}
end

struct TopsisResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    normalizedDecisionMatrix::DataFrame
    normalizedWeightedDecisionMatrix::DataFrame
    bestIndex::Int64
    scores::Array{Float64,1}
end

struct VikorResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    weightedDecisionMatrix::DataFrame
    bestIndex::Int64
    scores::Array{Float64,1}
end

struct ElectreResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    weightedDecisionMatrix::DataFrame
    fitnessTable::Array{Dict,1}
    nonfitnessTable::Array{Dict,1}
    fitnessMatrix::Array{Float64,2}
    nonfitnessMatrix::Array{Float64,2}
    C::Array{Float64,1}
    D::Array{Float64,1}
    bestIndex::Tuple
end

struct MooraResult <: MCDMResult
    mooraType::Symbol
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    weightedDecisionMatrix::DataFrame
    referenceMatrix::Union{DataFrame, Nothing}
    scores::Array{Float64,1}
    bestIndex::Int64
end

struct DematelResult <: MCDMResult
    comparisonMatrix::DataFrame
    NDMatrix::Array{Float64,2}
    relationShipMatrix::Array{Float64,2}
    c::Array{Float64,1}
    r::Array{Float64,1}
    c_plus_r::Array{Float64,1}
    c_minus_r::Array{Float64,1}
    threshold::Float64
    influenceMatrix::Array{Float64,2}
    weights::Array{Float64,1}
end

struct AHPConsistencyResult <: MCDMResult
    comparisonMatrix::DataFrame
    normalizedComparisonMatrix::DataFrame
    consistencyVector::Array{Float64,1}
    priority::Array{Float64,1}
    pc::Array{Float64,1}
    lambda_max::Float64
    CI::Float64
    RI::Float64
    CR::Float64
    isConsistent::Bool
end

struct AHPResult <: MCDMResult
    comparisonMatrixList::Array{DataFrame,1}
    criteriaComparisonMatrix::DataFrame
    criteriaConsistency::AHPConsistencyResult
    decisionMatrix::DataFrame
    scores::Array{Float64,1}
    weights::Array{Float64,1}
    bestIndex::Int64
end

struct NDSResult <: MCDMResult
    ranks::Array{Int64,1}
    bestIndex::Int64
end

struct LaplaceResult <: SCDMResult
    expected_values::Array{Float64,1}
    bestIndex::Int64
end

struct MaximinResult <: SCDMResult
    rowmins::Array{Float64,1}
    bestIndex::Int64
end

struct MaximaxResult <: SCDMResult
    rowmaxs::Array{Float64,1}
    bestIndex::Int64
end

struct MinimaxResult <: SCDMResult
    rowmaxs::Array{Float64,1}
    bestIndex::Int64
end

struct MiniminResult <: SCDMResult
    rowmins::Array{Float64,1}
    bestIndex::Int64
end

struct SavageResult <: SCDMResult
    regretMatrix::DataFrame
    scores::Array{Float64,1}
    bestIndex::Int64
end

struct HurwiczResult <: SCDMResult
    scores::Array{Float64,1}
    bestIndex::Int64
end

struct MLEResult <: SCDMResult
    scores::Array{Float64,1}
    bestIndex::Int64
end

struct ExpectedRegretResult <: SCDMResult
    scores::Array{Float64,1}
    bestIndex::Int64
end

struct GameResult <: MCDMResult
    row_player_probabilities::Array{Float64,1}
    value::Float64
end

struct DataEnvelopResult <: MCDMResult
    efficiencies::Array{Float64,1}
    references::DataFrame
    orderedcases::Array{Symbol,1}
end

struct GreyResult <: MCDMResult
    referenceRow::Array{Float64,1}
    normalizedMat::Matrix
    absoluteValueMat::Matrix
    greyTable::Matrix
    scores::Array{Float64,1}
    ordering::Array{Int64,1}
    bestIndex::Int64
end

struct SawResult <: MCDMResult
    decisionMatrix::DataFrame
    normalizedDecisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct ARASResult <: MCDMResult
    referenceRow::Array{Float64,1}
    extendMat::Array{Float64,2}
    normalizedMat::Array{Float64,2}
    optimalitydegrees::Array{Float64,1}
    scores::Array{Float64,1}
    orderings::Array{Int64,1}
    bestIndex::Int64
end

struct WPMResult <: MCDMResult
    decisionMatrix::DataFrame
    normalizedDecisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct WASPASResult <: MCDMResult
    decisionMatrix::DataFrame
    normalizedDecisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct EDASResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct MARCOSResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct MABACResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct MAIRCAResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct COPRASResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct PrometheeResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct CoCoSoResult <: MCDMResult
    decisionMatrix::DataFrame
    weights::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct CRITICResult <: MCDMResult
    decisionMatrix::DataFrame
    w::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct EntropyResult <: MCDMResult
    decisionMatrix::DataFrame
    w::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct CODASResult <: MCDMResult
    decisionMatrix::DataFrame
    w::Array{Float64,1}
    scores::Array{Float64,1}
    ranking::Array{Int64,1}
    bestIndex::Int64
end

struct TopsisMethod <: MCDMMethod end

struct ElectreMethod <: MCDMMethod
end

struct ArasMethod <: MCDMMethod
end

struct CocosoMethod <: MCDMMethod
    lambda::Float64
end

CocosoMethod()::CocosoMethod = CocosoMethod(0.5)

struct CodasMethod <: MCDMMethod
    tau::Float64
end

CodasMethod()::CodasMethod = CodasMethod(0.02)

struct CoprasMethod <: MCDMMethod
end

struct CriticMethod <: MCDMMethod
end

struct EdasMethod <: MCDMMethod
end

struct GreyMethod <: MCDMMethod
    zeta::Float64
end

GreyMethod() :: GreyMethod = GreyMethod(0.5)

struct MabacMethod <: MCDMMethod
end

struct MaircaMethod <: MCDMMethod
end

struct MooraMethod <: MCDMMethod
    method::Symbol
end

MooraMethod() :: MooraMethod = MooraMethod(:reference)

struct PrometheeMethod <: MCDMMethod
    pref::Array{Function, 1}
    qs::Array{Union{Nothing, Float64}, 1}
    ps::Array{Union{Nothing, Float64}, 1}
end

PrometheeMethod(pref::Array{Function, 1}, qs::Array{Float64,1}, ps::Array{Float64, 1}) :: PrometheeMethod = PrometheeMethod(pref, qs, ps)

struct SawMethod <: MCDMMethod
end

struct VikorMethod <: MCDMMethod
    v::Float64
end

VikorMethod()::VikorMethod = VikorMethod(0.5)

struct WaspasMethod <: MCDMMethod
    lambda::Float64
end

WaspasMethod() :: WaspasMethod = WaspasMethod(0.5)

struct WPMMethod <: MCDMMethod
end

struct MarcosMethod <: MCDMMethod
end


## Turing Complete Interpreter
### A Very Tiny Turing Complete Programming Language defined with denotational semantics

# semantica dalle dispense degano
using Metatheory

import Base.ImmutableDict
Mem = Dict{Symbol, Union{Bool, Int}}

read_mem = @theory begin
	(v::Symbol, σ::Mem) |> if v == :skip σ else σ[v] end
end

@testset "Reading Memory" begin
	ex = :((x), $(Mem(:x => 2)))
	@test true == areequal(read_mem, ex, 2)
end

arithm_rules = @theory begin
	(a + b, σ::Mem) 	=> (a, σ) + (b, σ)
	(a * b, σ::Mem) 	=> (a, σ) * (b, σ)
	(a - b, σ::Mem) 	=> (a, σ) - (b, σ)
	(n::Int, σ::Mem) 	=> n
	(a::Int + b::Int) 	|> a + b
	(a::Int * b::Int) 	|> a * b
	(a::Int - b::Int) 	|> a - b
end

@testset "Arithmetic" begin
	@test areequal(read_mem ∪ arithm_rules, :((2 + 3), $(Mem())), 5)
end

# don't need to access memory
bool_rules = @theory begin
	(a < b, σ::Mem) => (a, σ) < (b, σ)
	(a ∨ b, σ::Mem) => (a, σ) ∨ (b, σ)
	(a ∧ b, σ::Mem) => (a, σ) ∧ (b, σ)
	(¬(a), σ::Mem) => ¬((a, σ))

	(a::Bool, σ::Mem) 	|> a
	(¬a::Bool) 			|> !a
	(a::Bool ∨ b::Bool) 	|> (a || b)
	(a::Bool ∧ b::Bool) 	|> (a && b)
	(a::Int < b::Int) 	|> (a < b)
end

t = read_mem ∪ arithm_rules ∪ bool_rules

@testset "Booleans" begin
	@test areequal(t, :((false ∨ false), $(Mem())), false)

	exx = :((false ∨ false) ∨ ¬(false ∨ false), $(Mem(:x => 2)))
	g = EGraph(exx)
	saturate!(g, t)
	ex = extract!(g, astsize)
	@test ex == true
	params=SaturationParams(timeout=12)
	@test areequal(t, exx, true; params=params)

	@test areequal(t, :((2 < 3) ∧ (3 < 4), $(Mem(:x => 2))), true)
	@test areequal(t, :((2 < x) ∨ ¬(3 < 4), $(Mem(:x => 2))), false)
	@test  areequal(t, :((2 < x) ∨ ¬(3 < 4), $(Mem(:x => 4))), true)
end

if_rules = @theory begin
	(if guard; t end) => (if guard; t else :skip end)
	(if guard; t else f end, σ::Mem) => (if (guard, σ); t else f end, σ)
	(if true; t else f end, σ::Mem) => (t, σ)
	(if false; t else f end, σ::Mem) => (f, σ)
end

if_language = read_mem ∪ arithm_rules ∪ bool_rules ∪ if_rules

@testset "If Semantics" begin
	@test areequal(if_language, 2, :(if true x else 0 end, 	$(Mem(:x => 2))))
	@test areequal(if_language, 0, :(if false x else 0 end, 	$(Mem(:x => 2))))
	@test areequal(if_language, 2, :(if ¬(false) x else 0 end, $(Mem(:x => 2))))
	params=SaturationParams(timeout=10)
	@test areequal(if_language, 0, :(if ¬(2 < x) x else 0 end, $(Mem(:x => 3))); params=params)
end

while_rules = @theory begin
	(:skip, σ::Mem) => σ
	((c1; c2), σ::Mem) => ((c1, σ); c2)
	(c1::Int; c2) => c2
	(c1::Bool; c2) => c2
	(σ::Mem; c2) => (c2, σ)
	(while guard body end, σ::Mem) =>
		(if guard; (body; while guard body end) else :skip end, σ)
end

write_mem = @theory begin
	(sym::Symbol = val, σ::Mem) => (sym = (val, σ), σ)
	(sym::Symbol = val::Int, σ::Mem) |> merge(σ, Dict(sym => val))
end

while_language = if_language ∪ write_mem ∪ while_rules;

@testset "While Semantics" begin
	exx = :((x = 3), $(Mem(:x => 2)))
	g = EGraph(exx)
	saturate!(g, while_language)
	ex = extract!(g, astsize)

	@test areequal(while_language, Mem(:x => 3), exx)

	exx = :((x = 4; x = x + 1), $(Mem(:x => 3)))
	g = EGraph(exx)
	saturate!(g, while_language)
	ex = extract!(g, astsize)

	params=SaturationParams(timeout=10)
	@test areequal(while_language, Mem(:x => 5), exx; params=params)

	# FIXME bug!
	params=SaturationParams(timeout=14)
	exx = :((if x < 10 x = x + 1 else skip end), $(Mem(:x => 3)))
	# g = EGraph(exx)
	# params=SaturationParams(timeout=12)
	# saturate!(g, while_language, params)
	# println(extract!(g, astsize))
	@test areequal(while_language, Mem(:x => 4), exx; params=params)
	# exit(0)

	exx = :(
		(while x < 10
			x = x + 1
		end; x), $(Mem(:x => 3)))
	g = EGraph(exx)
	params=SaturationParams(timeout=100, scheduler=Schedulers.SimpleScheduler)
	saturate!(g, while_language, params)
	@test 10 == extract!(g, astsize)
end


"""
See [1, "Online Implementation of IRL: A Hybrid Optimal Adaptive Controller"].

# Refs
[1] “Reinforcement Learning and Feedback Control: Using Natural Decision Methods to Design Optimal Adaptive Controllers,” IEEE Control Syst., vol. 32, no. 6, pp. 76–105, Dec. 2012, doi: 10.1109/MCS.2012.2214134.

# Notes
- T: Data stack period
- N: The maximum length of stacked data
"""
mutable struct LinearIRL <: AbstractIRL
    Q::AbstractMatrix
    R_inv::AbstractMatrix
    B::AbstractMatrix
    T::Real
    N::Int
    i::Int
    i_init::Int
    terminated::Bool
    function LinearIRL(Q::AbstractMatrix, R::AbstractMatrix, B::AbstractMatrix;
            T=0.04,
            N=nothing,
            i_init=0,  # iteration number
        )
        n1, n2 = size(Q)
        @assert n1 == n2
        N_min = Int(n1*(n1 + 1)/2)
        if N == nothing
            N = N_min
        else
            @assert N >= N_min
        end
        @assert T > 0 && N > 0
        n1, n2 = size(Q)
        @assert n1 == n2
        N_min = Int(n1*(n1 + 1)/2)
        if N == nothing
            N = N_min
        else
            @assert N >= N_min
        end
        @assert T > 0 && N > 0
        R_inv = inv(R)
        i = i_init
        terminated = false
        new(Q, R_inv, B, T, N, i, i_init, terminated)
    end
end

function reset!(irl::LinearIRL)
    irl.i = irl.i_init
end

"""
Value iteration [1, Eq. 99]; updated in least-square sense
# Notes
w: critic parameter (vectorised)
- ϕs_prev: the vector of bases (evaluated)
- V̂: the vector of approximate values (evaluated)
"""
function value_iteration!(irl::LinearIRL, buffer::DataBuffer, w;
        sc::AbstractStopCondition=DistanceStopCondition(),
    )
    @unpack i, N = irl
    @unpack data_array = buffer
    data_filtered = filter(x -> x.i == i, data_array)  # data from the current policy
    if length(data_filtered) >= N + 1
        data_sorted = sort(data_filtered, by=x -> x.t)  # sort by time index
        ϕs_prev = data_sorted[end-N:end-1] |> Map(datum -> datum.ϕ) |> collect
        # V̂s = data_sorted[end-(N-1):end] |> Map(datum -> datum.V̂) |> collect
        xs = data_sorted[end-(N-1):end] |> Map(datum -> datum.x) |> collect
        ∫rs = data_sorted[end-(N-1):end] |> Map(datum -> datum.∫r) |> collect
        P = convert_to_matrix(w)
        V̂s_with_prev_P = xs |> Map(x -> value(irl, P, x)) |> collect
        V̂s = ∫rs .+ V̂s_with_prev_P
        # update the critic parameter
        w_new = ( hcat(V̂s...) * pinv(hcat(ϕs_prev...)) )'[:]  # to reduce the computation time; [:] for vectorisation
        if !irl.terminated
            if is_terminated(sc, w, w_new)
                irl.terminated = true
            else
                w .= w_new
                update_index!(irl)
            end
        end
    end
end

"""
Policy iteration [1, Eq. 98]; updated in least-square sense
# Notes
w: critic parameter (vectorised)
- Δϕs: the vector of (ϕ - ϕ_prev) (evaluated)
- ∫rs: the vector of integral running cost by numerical integration (evaluated)
"""
function policy_iteration!(irl::LinearIRL, buffer::DataBuffer, w;
        sc::AbstractStopCondition=DistanceStopCondition(),
    )
    @unpack i, N = irl
    @unpack data_array = buffer
    data_filtered = filter(x -> x.i == i, data_array)  # data from the current policy
    if length(data_filtered) >= N + 1
        data_sorted = sort(data_filtered, by=x -> x.t)  # sort by time index
        ϕs_prev_and_present = data_sorted[end-N:end] |> Map(datum -> datum.ϕ) |> collect
        Δϕs = diff(ϕs_prev_and_present)
        ∫rs = data_sorted[end-(N-1):end] |> Map(datum -> datum.∫r) |> collect
        # update the critic parameter
        w_new = ( hcat(∫rs...) * pinv(hcat(-Δϕs...)) )'[:]  # to reduce the computation time; [:] for vectorisation
        if !irl.terminated
            if is_terminated(sc, w, w_new)
                irl.terminated = true
            else
                w .= w_new
                update_index!(irl)
            end
        end
    end
end

function update_index!(irl::LinearIRL)
    irl.i += 1
end

"""
Policy improvement [1, Eq. 96].
"""
function _optimal_input(R_inv, B, P, x)
    -0.5 * R_inv * B' * (2 * P * x)
end

function optimal_input(irl::LinearIRL, x, w::AbstractVector)
    @unpack R_inv, B = irl
    P = convert_to_matrix(w)
    _optimal_input(R_inv, B, P, x)
end

function optimal_input(irl::LinearIRL, x, P::AbstractMatrix)
    @unpack R_inv, B = irl
    _optimal_input(R_inv, B, P, x)
end

function value(irl::LinearIRL, P::AbstractMatrix, x)
    x' * P * x
end

function value(irl::LinearIRL, w::AbstractVector, x)
    P = convert_to_matrix(w)
    value(irl, P, x)
end

"""
w: critic parameter (vectorised)
"""
function Base.push!(buffer::DataBuffer, irl::LinearIRL, cost::AbstractCost;
        t, x, u, w,
    )
    # P = convert_to_matrix(w)
    @unpack data_array = buffer
    @unpack i = irl
    data_sorted = sort(data_array, by = x -> x.t)  # sorted by t
    # V̂_with_prev_P = value(irl, P, x)
    # prev data
    if length(data_sorted) != 0
        t_prev = data_sorted[end].t
        x_prev = data_sorted[end].x
        u_prev = data_sorted[end].u
        # numerical integration
        Δt = t - t_prev
        r = cost(x, u)
        r_prev = cost(x_prev, u_prev)
        ∫r = 0.5 * (r + r_prev) * Δt # trapezoidal
        # V̂ = ∫r + V̂_with_prev_P
    else
        # V̂ = missing
        ∫r = missing
    end
    ϕ = convert_quadratic_to_linear_basis(x)  # x'Px = w'ϕ(x)
    datum = (;
             t=t,
             x=x,
             u=u,
             w=w,  # logging
             ϕ=ϕ,
             ∫r=∫r,
             # V̂=V̂,
             i=i,  # iteration number
            )
    push!(buffer.data_array, datum)
end


"""
    Fill <: Imputor

Fills in the missing data with a specific value.

# Fields
* `value::Any`: A scalar missing value or a function that returns the a scalar if
    passed the data with missing data removed (e.g, `mean`)
"""
struct Fill{T} <: Imputor
    value::T
end

"""
    Fill() -> Fill

By default `Fill()` will use the mean of the existing values as the fill value.
"""
Fill() = Fill(mean)

"""
    impute!(imp::Fill, ctx::Context, data::AbstractVector)

Computes the fill value if `imp.value` is a `Function` (i.e., `imp.value(drop(copy(data)))`)
and replaces all missing values in the `data` with that value.
"""
function impute!(imp::Fill, ctx::Context, data::AbstractVector)
    fill_val = if isa(imp.value, Function)
        imp.value(Iterators.drop(copy(data)))
    else
        imp.value
    end

    for i in 1:length(data)
        if ismissing(ctx, data[i])
            data[i] = fill_val
        end
    end

    return data
end


using Test
using DataWrangler

@testset "d" begin

    x = [1,2,3,4,5]
    dx = d(x,center=true)
    mdx = ismissing.(dx)
    @test mdx == [0,0,0,0,1]
    @test dx[.!Bool.(mdx)] == [1,1,1,1]

    dx = d(x,2,center=true)
    mdx = ismissing.(dx)
    @test mdx == [1,0,0,0,1]
    @test dx[.!Bool.(mdx)] == [0,0,0]

    dx = d(x,1,2)
    mdx = ismissing.(dx)
    @test mdx == [0,0,0]
    @test dx[.!Bool.(mdx)] == [2,2,2]

    x = reshape(collect(1:20),10,2)
    dx = d(x,2,2)
    @test dx == repeat([0],6,2)

    # Fractional Difference
    x = rand(100)
    x ≈ d(d(x,.5),-.5)
    x ≈ d(d(x,2.5),-2.5)
    d(x,1.) ≈ [x[1]; d(x)]

    n = 1000
    wx = 2*pi
    w = Array(LinRange(0,wx,n))
    x = sin.(w)
    x10 = sin.(w.+pi/2)
    x20 = sin.(w.+pi)
    scale = wx
    dx = 1.0; isapprox(x10[10:end], d(x,dx;scale)[10:end]; rtol= .01)
    dx = 2.0; isapprox(x20[10:end], d(x,dx;scale)[10:end]; rtol= .01)

end


using LinearAlgebra, SparseArrays, DiffEqOperators, Random, Test
################################################################################
# Test 2d extension
################################################################################

Random.seed!(7373)

#Create Array
n = 8
m = 15
A = rand(n,m)

#Create atomic BC
q1 = RobinBC((1.0, 2.0, 3.0), (0.0, -1.0, 2.0), 0.1, 4.0)
q2 = PeriodicBC(Float64)

BCx = vcat(fill(q1, div(m,2)), fill(q2, m-div(m,2)))  #The size of BCx has to be all size components *except* for x
BCy = vcat(fill(q1, div(n,2)), fill(q2, n-div(n,2)))

Qx = MultiDimBC{1}(BCx)
Qy = MultiDimBC{2}(BCy)

Ax = Qx*A
Ay = Qy*A

@test size(Ax)[1] == size(A)[1]+2
@test size(Ay)[2] == size(A)[2]+2

for j in 1:m
    @test Ax[:, j]  == Array(BCx[j]*A[:, j])
end
for i in 1:n
    @test Ay[i,:] == Array(BCy[i]*A[i,:])
end

################################################################################
# Test 3d extension
################################################################################

#Create Array
n = 8
m = 11
o = 12
A = rand(n,m, o)

#Create atomic BC
q1 = RobinBC((1.0, 2.0, 3.0), (0.0, -1.0, 2.0), 0.1, 4.0)
q2 = PeriodicBC(Float64)

BCx = vcat(fill(q1, (div(m,2), o)), fill(q2, (m-div(m,2), o)))  #The size of BCx has to be all size components *except* for x
BCy = vcat(fill(q1, (div(n,2), o)), fill(q2, (n-div(n,2), o)))
BCz = fill(Dirichlet0BC(Float64), (n,m))

Qx = MultiDimBC{1}(BCx)
Qy = MultiDimBC{2}(BCy)
Qz = MultiDimBC{3}(Dirichlet0BC(Float64), size(A)) #Test the other constructor

Ax = Qx*A
Ay = Qy*A
Az = Qz*A

Q = compose(Qx,Qy,Qz)
QL, Qb = Array(Q, size(A))
QLs, Qbs = sparse(Q, size(A))

A_conc = QL*reshape(A, prod(size(A))) .+ Qb
A_conc_sp = QLs*reshape(A,prod(size(A))) .+ Qbs

#test BC concretization
A_arr = Array(Q*A)
@test reshape(A_arr, prod(size(A_arr))) ≈ A_conc_sp ≈ A_conc

@test size(Ax)[1] == size(A)[1]+2
@test size(Ay)[2] == size(A)[2]+2
@test size(Az)[3] == size(A)[3]+2
for j in 1:m, k in 1:o
    @test Ax[:, j, k] == Array(BCx[j, k]*A[:, j, k])
end
for i in 1:n, k in 1:o
    @test Ay[i, :, k] == Array(BCy[i, k]*A[i, :, k])
end
for i in 1:n, j in 1:m
    @test Az[i, j, :] == Array(BCz[i, j]*A[i, j, :])
end

#test compositions to higher dimension
for N in 2:6
    sizes = rand(4:7, N)
    local A = rand(sizes...)

    Q1_N = RobinBC(Tuple(rand(3)), Tuple(rand(3)), fill(0.1, N), 4.0, size(A))

    local Q = compose(Q1_N...)

    A1_N = Q1_N.*fill(A, N)

    local A_arr = Array(Q*A)
    Q_l, Q_b = sparse(Q, size(A))

    @test A_arr ≈ Array(compose(A1_N...))
    @test A_arr ≈ reshape(Q_l*reshape(A, length(A)) .+ Q_b, size(A_arr)) #Test concretization
end


using Triangulate
using WriteVTK
using PolynomialBasis
using ImplicitDomainQuadrature
using Revise
using CutCell
include("../../../test/useful_routines.jl")

function update_product_stress!(
    qpstress,
    basis,
    stiffness,
    transfstress,
    theta0,
    celldisp,
    points,
    jac,
    vectosymmconverter,
)
    dim = CutCell.dimension(basis)
    lambda, mu = CutCell.lame_coefficients(stiffness, +1)
    nump = size(points)[2]

    for qpidx = 1:nump
        p = points[:, qpidx]
        grad = CutCell.transform_gradient(gradient(basis, p), jac)
        NK = sum([CutCell.make_row_matrix(vectosymmconverter[k], grad[:, k]) for k = 1:dim])
        symmdispgrad = NK * celldisp

        inplanestress = (stiffness[+1] * symmdispgrad) - transfstress
        s33 = lambda * (symmdispgrad[1] + symmdispgrad[2]) - (lambda + 2mu / 3) * theta0

        numericalstress = vcat(inplanestress, s33)

        append!(qpstress, numericalstress)
    end
end

function update_parent_stress!(
    qpstress,
    basis,
    stiffness,
    celldisp,
    points,
    jac,
    vectosymmconverter,
)

    dim = CutCell.dimension(basis)
    lambda, mu = CutCell.lame_coefficients(stiffness, -1)
    nump = size(points)[2]
    for qpidx = 1:nump
        p = points[:, qpidx]
        grad = CutCell.transform_gradient(gradient(basis, p), jac)
        NK = sum([CutCell.make_row_matrix(vectosymmconverter[k], grad[:, k]) for k = 1:dim])
        symmdispgrad = NK * celldisp

        inplanestress = stiffness[-1] * symmdispgrad
        s33 = lambda * (symmdispgrad[1] + symmdispgrad[2])

        numericalstress = vcat(inplanestress, s33)

        append!(qpstress, numericalstress)
    end
end

function compute_stress_at_quadrature_points(
    nodaldisplacement,
    basis,
    stiffness,
    transfstress,
    theta0,
    cellquads,
    cutmesh,
)

    dim = CutCell.dimension(basis)
    ncells = CutCell.number_of_cells(cutmesh)
    qpstress = zeros(0)
    jac = CutCell.jacobian(cutmesh)
    vectosymmconverter = CutCell.vector_to_symmetric_matrix_converter()

    for cellid = 1:ncells
        s = CutCell.cell_sign(cutmesh, cellid)
        @assert s == -1 || s == 0 || s == +1
        if s == +1 || s == 0
            nodeids = CutCell.nodal_connectivity(cutmesh, +1, cellid)
            celldofs = CutCell.element_dofs(nodeids, dim)
            celldisp = nodaldisplacement[celldofs]
            points = cellquads[+1, cellid].points
            update_product_stress!(
                qpstress,
                basis,
                stiffness,
                transfstress,
                theta0,
                celldisp,
                points,
                jac,
                vectosymmconverter,
            )
        end
        if s == -1 || s == 0
            nodeids = CutCell.nodal_connectivity(cutmesh, -1, cellid)
            celldofs = CutCell.element_dofs(nodeids, dim)
            celldisp = nodaldisplacement[celldofs]
            points = cellquads[-1, cellid].points
            update_parent_stress!(
                qpstress,
                basis,
                stiffness,
                celldisp,
                points,
                jac,
                vectosymmconverter,
            )
        end
    end
    return reshape(qpstress, 4, :)
end

function compute_quadrature_points(cellquads, cutmesh)
    ncells = CutCell.number_of_cells(cutmesh)
    coords = zeros(2, 0)
    for cellid = 1:ncells
        s = CutCell.cell_sign(cutmesh, cellid)
        cellmap = CutCell.cell_map(cutmesh, cellid)
        @assert s == -1 || s == 0 || s == +1
        if s == +1 || s == 0
            coords = hcat(coords, cellmap(cellquads[+1, cellid].points))
        end
        if s == -1 || s == 0
            coords = hcat(coords, cellmap(cellquads[-1, cellid].points))
        end
    end
    return coords
end

function onboundary(x, L, W)
    return x[2] ≈ 0.0 || x[1] ≈ L || x[2] ≈ W || x[1] ≈ 0.0
end

function bulk_modulus(l, m)
    return l + 2m / 3
end

function lame_lambda(k, m)
    return k - 2m / 3
end

function analytical_coefficient_matrix(inradius, outradius, ls, ms, lc, mc)
    a = zeros(3, 3)
    a[1, 1] = inradius
    a[1, 2] = -inradius
    a[1, 3] = -1.0 / inradius
    a[2, 1] = 2 * (lc + mc)
    a[2, 2] = -2 * (ls + ms)
    a[2, 3] = 2ms / inradius^2
    a[3, 2] = 2(ls + ms)
    a[3, 3] = -2ms / outradius^2
    return a
end

function analytical_coefficient_rhs(ls, ms, theta0)
    r = zeros(3)
    Ks = bulk_modulus(ls, ms)
    r[2] = -Ks * theta0
    r[3] = Ks * theta0
    return r
end

struct AnalyticalSolution
    inradius::Any
    outradius::Any
    center::Any
    A1c::Any
    A1s::Any
    A2s::Any
    ls::Any
    ms::Any
    lc::Any
    mc::Any
    theta0::Any
    function AnalyticalSolution(inradius, outradius, center, ls, ms, lc, mc, theta0)
        a = analytical_coefficient_matrix(inradius, outradius, ls, ms, lc, mc)
        r = analytical_coefficient_rhs(ls, ms, theta0)
        coeffs = a \ r
        new(
            inradius,
            outradius,
            center,
            coeffs[1],
            coeffs[2],
            coeffs[3],
            ls,
            ms,
            lc,
            mc,
            theta0,
        )
    end
end

function radial_displacement(A::AnalyticalSolution, r)
    if r <= A.inradius
        return A.A1c * r
    else
        return A.A1s * r + A.A2s / r
    end
end

function (A::AnalyticalSolution)(x)
    relpos = x - A.center
    r = sqrt(relpos' * relpos)
    ur = radial_displacement(A, r)
    if ur ≈ 0.0
        [0.0, 0.0]
    else
        costheta = (x[1] - A.center[1]) / r
        sintheta = (x[2] - A.center[2]) / r
        u1 = ur * costheta
        u2 = ur * sintheta
        return [u1, u2]
    end
end

function solve_cylindrical_bc_elasticity(
    width,
    center,
    inradius,
    outradius,
    stiffness,
    transfstress,
    theta0,
    nelmts,
    polyorder,
    numqp,
    penaltyfactor,
)

    lambda1, mu1 = CutCell.lame_coefficients(stiffness, +1)
    lambda2, mu2 = CutCell.lame_coefficients(stiffness, -1)

    analyticalsolution =
        AnalyticalSolution(inradius, outradius, center, lambda1, mu1, lambda2, mu2, theta0)

    dx = width / nelmts
    meanmoduli = 0.5 * (lambda1 + lambda2 + mu1 + mu2)
    penalty = penaltyfactor / dx * meanmoduli

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCell.Mesh([0.0, 0.0], [width, width], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs = CutCell.levelset_coefficients(
        x -> -circle_distance_function(x, center, inradius),
        mesh,
    )

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition =
        CutCell.incoherent_interface_condition(basis, interfacequads, stiffness, cutmesh, penalty)

    displacementbc = CutCell.DisplacementCondition(
        analyticalsolution,
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onboundary(x, width, width),
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh)
    CutCell.assemble_bulk_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        cellquads,
        cutmesh,
    )
    CutCell.assemble_incoherent_interface_transformation_rhs!(
        sysrhs,
        transfstress,
        basis,
        interfacequads,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(sysmatrix, sysrhs, displacementbc, cutmesh)
    CutCell.assemble_penalty_displacement_transformation_rhs!(
        sysrhs,
        transfstress,
        basis,
        facequads,
        cutmesh,
        x -> onboundary(x, width, width),
    )

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    sol = matrix \ rhs

    return sol, basis, cellquads, cutmesh
end

function solve_zero_traction_bc_elasticity(
    width,
    center,
    inradius,
    outradius,
    stiffness,
    transfstress,
    theta0,
    nelmts,
    polyorder,
    numqp,
    penaltyfactor,
)

    lambda1, mu1 = CutCell.lame_coefficients(stiffness, +1)
    lambda2, mu2 = CutCell.lame_coefficients(stiffness, -1)

    dx = width / nelmts
    meanmoduli = 0.5 * (lambda1 + lambda2 + mu1 + mu2)
    penalty = penaltyfactor / dx * meanmoduli

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCell.Mesh([0.0, 0.0], [width, width], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs = CutCell.levelset_coefficients(
        x -> -circle_distance_function(x, center, inradius),
        mesh,
    )

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition =
        CutCell.incoherent_interface_condition(basis, interfacequads, stiffness, cutmesh, penalty)

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh)
    CutCell.assemble_bulk_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        cellquads,
        cutmesh,
    )
    CutCell.assemble_incoherent_interface_transformation_rhs!(
        sysrhs,
        transfstress,
        basis,
        interfacequads,
        cutmesh,
    )

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    topleftnodeid = CutCell.nodes_per_mesh_side(mesh)[2]
    CutCell.apply_dirichlet_bc!(matrix, rhs, [1, topleftnodeid], 1, 0.0, 2)
    CutCell.apply_dirichlet_bc!(matrix, rhs, [1], 2, 0.0, 2)

    sol = matrix \ rhs

    return sol, basis, cellquads, cutmesh
end

K1, K2 = 247.0, 192.0
mu1, mu2 = 126.0, 87.0
lambda1 = lame_lambda(K1, mu1)
lambda2 = lame_lambda(K2, mu2)

theta0 = -0.067
stiffness = CutCell.HookeStiffness(lambda1, mu1, lambda2, mu2)

width = 1.0
penaltyfactor = 1e2

nelmts = 11
polyorder = 2
numqp = required_quadrature_order(polyorder) + 4
center = [width / 2, width / 2]
inradius = width / 4
outradius = width

transfstress = CutCell.plane_strain_transformation_stress(lambda1, mu1, theta0)

# sol, basis, cellquads, cutmesh = solve_cylindrical_bc_elasticity(
#     width,
#     center,
#     inradius,
#     outradius,
#     stiffness,
#     transfstress,
#     theta0,
#     nelmts,
#     polyorder,
#     numqp,
#     penaltyfactor,
# )

sol, basis, cellquads, cutmesh = solve_zero_traction_bc_elasticity(
    width,
    center,
    inradius,
    outradius,
    stiffness,
    transfstress,
    theta0,
    nelmts,
    polyorder,
    numqp,
    penaltyfactor,
)

qpstress = compute_stress_at_quadrature_points(
    sol,
    basis,
    stiffness,
    transfstress,
    theta0,
    cellquads,
    cutmesh,
)

numqp = size(qpstress)[2]
pressure = -[(qpstress[1, i] + qpstress[2, i] + qpstress[4, i])/3 for i = 1:numqp]

quadcoords = compute_quadrature_points(cellquads, cutmesh)

triin = Triangulate.TriangulateIO()
triin.pointlist = quadcoords
(triout, vorout) = triangulate("", triin)
connectivity = triout.trianglelist
cells = [
    MeshCell(VTKCellTypes.VTK_TRIANGLE, connectivity[:, i]) for i = 1:size(connectivity)[2]
]
vtkfile = vtk_grid(
    "examples/transformation-strain/cylindrical-particle/zero-traction-pressure",
    quadcoords[1, :],
    quadcoords[2, :],
    cells,
)
vtkfile["pressure"] = pressure
vtkfile["s11"] = qpstress[1,:]
vtkfile["s22"] = qpstress[2,:]
vtkfile["s12"] = qpstress[3,:]
vtkfile["s33"] = qpstress[4,:]
outfiles = vtk_save(vtkfile)


module SymbolicUtils

export @syms, term, @fun, showraw

# Sym, Term and other types
include("types.jl")

# Methods on symbolic objects
using SpecialFunctions, NaNMath
export cond
include("methods.jl")

# LinkedList, simplification utilities
include("utils.jl")

export Rewriters

# A library for composing together expr -> expr functions
include("rewriters.jl")

using .Rewriters

using Combinatorics: permutations, combinations
export @rule, @acrule, @arule, RuleSet

# Rule type and @rule macro
include("rule.jl")

# Matching a Rule
include("matchers.jl")

# Convert to an efficient multi-variate polynomial representation
import AbstractAlgebra.Generic: MPoly, PolynomialRing, ZZ, exponent_vector
using AbstractAlgebra: ismonomial, symbols
using DataStructures
include("abstractalgebra.jl")

# Term ordering
include("ordering.jl")

# Default rules for expression simplification
include("simplify_rules.jl")

# API = simplify + substitute
export simplify, substitute
include("api.jl")

end # module


"""
```
eval(slm::SLM{T}, x::AbstractVector{T}, evalmode::Int = 0) where {T <: Real}
```

evaluates the spline prescribed by `slm` at the values `x`. The optional input
`evalmode` specifies the evaluation of the spline according to

0  -> spline's value at each point in `x`
1  -> first derivative of the spline
2  -> second derivative of the spline
3  -> third derivative of the spline

Note that the inverse of points which fall above the maximum
or below the minimum value of the function will be returned as a NaN.

Currently, no constant extrapolation is supported.
"""
function eval(slm::SLM{T}, x::AbstractVector{T}, evalmode::Int = 0) where {T <: Real}

    ## Set up
    nₓ     = length(x)
    knots  = get_knots(slm)
    nk     = length(knots)
    dknots = diff(knots)
    coef   = get_coef(slm)

    # Extrapolation
    if slm[:extrap] != :none
        error("Extrapolation method $(slm[:extrap]) is not supported.")
    end

    ## Evaluate spline

    # Sort points into bins. Inverse case is handled later -> just zeros for now.
    x_bins = evalmode >= 0 ? bin_sort(x, knots) : zeros(Int, nₓ)

    if get_type(slm) == :cubic
        y = if evalmode == 0
            eval_cubic(x, knots, dknots, coef, x_bins)
        elseif evalmode == 1
            eval_cubic_derivative_1(x, knots, dknots, coef, x_bins)
        elseif evalmode == 2
            eval_cubic_derivative_2(x, knots, dknots, coef, x_bins)
        elseif evalmode == 3
            eval_cubic_derivative_3(x, knots, dknots, coef, x_bins)
        elseif evalmode == -1
            error("The function inverse has not been implemented yet.")
        else
            error("The input `evalmode` cannot be $(evalmode). It must be one of [0, 1, 2, 3].")
        end
    else
        error("Cannot use the spline type $(get_type(slm))")
    end

    return y
end

function eval_cubic(x::AbstractVector{T}, knots::AbstractVector{T}, dknots::AbstractVector{T},
                    coef::AbstractMatrix{T}, x_bins::AbstractVector{Int}) where {T <: Real}

    # Set up
    t  = (x - knots[x_bins]) ./ dknots[x_bins]
    t² = t .^ 2
    t³ = t .^ 3
    s² = (1. .- t) .^ 2
    s³ = (1. .- t) .^ 3

    # No extrapolation yet so directly return
    x_bins_up = x_bins .+ 1

    return (coef[x_bins, 2] .* (s² - s³) +
            coef[x_bins_up, 2] .* (t³ - t²)) .* dknots[x_bins] +
            coef[x_bins, 1] .* (3 .* s² - 2. .* s³) +
            coef[x_bins_up, 1] .* (3. .* t² - 2. .* t³)

end

function eval_cubic_derivative_1(x::AbstractVector{T}, knots::AbstractVector{T}, dknots::AbstractVector{T},
                                 coef::AbstractMatrix{T}, x_bins::AbstractVector{Int}) where {T <: Real}

    # Set up
    t  = (x - knots[x_bins]) ./ dknots[x_bins]
    t² = t .^ 2
    s  = 1. .- t
    s² = (1. .- t) .^ 2

    # No extrapolation yet so directly return
    x_bins_up = x_bins .+ 1

    return coef[x_bins, 2] .* (3. .* s² - 2. .* s) +
        coef[x_bins_up, 2] .* (3. .* t² - 2. .* t) +
        6 .* (coef[x_bins, 1] .* (-s + s²) +
              coef[x_bins_up, 1] .* (t - t²)) ./ dknots[x_bins]
end

function eval_cubic_derivative_2(x::AbstractVector{T}, knots::AbstractVector{T}, dknots::AbstractVector{T},
                                 coef::AbstractMatrix{T}, x_bins::AbstractVector{Int}) where {T <: Real}

    # Set up
    t  = (x - knots[x_bins]) ./ dknots[x_bins]
    s  = 1. .- t

    # No extrapolation yet so directly return
    x_bins_up = x_bins .+ 1

    return (coef[x_bins, 2] .* (2. .- 6. .* s) +
            coef[x_bins_up, 2] .* (6 .* t .- 2.)) ./ dknots[x_bins] +
            (coef[x_bins, 1] .* (6. .- 12. .* s) +
             coef[x_bins_up, 1] .* (6. .- 12. .* t)) ./ (dknots[x_bins] .^ 2)

end

function eval_cubic_derivative_3(x::AbstractVector{T}, knots::AbstractVector{T}, dknots::AbstractVector{T},
                                 coef::AbstractMatrix{T}, x_bins::AbstractVector{Int}) where {T <: Real}

    # No extrapolation yet so directly return
    x_bins_up = x_bins .+ 1

    return 6. .* (coef[x_bins, 2] .+ coef[x_bins_up, 2]) ./ (dknots[x_bins] .^ 2) +
        12. .* (coef[x_bins, 1] .- coef[x_bins_up, 1]) ./ (dknots[x_bins] .^ 3)
end


"""
    j_fluence_DA_inf_TD(t, β::Array{Float64,1}, ρ::Float64, nmed::Float64)

Compute the jacobian of fluence_DA_inf_TD to be used in Levenberg-Marquardt.

# Arguments
- `t`: the time vector (ns).
- `β::Array{Float64,1}`: the optical properties [μa, μs'] (cm⁻¹)
- `ρ::Float64`: the source detector separation (cm⁻¹)
- `nmed::Float64`: the sample medium's index of refraction

# Examples
julia> j_fluence_DA_inf_TD(1:1:5, [0.1,10.0], 1.0, 1.0)
"""
function j_fluence_DA_inf_TD(t, β::Array{Float64,1}, ρ::Float64, nmed::Float64 = 1.0)
    μa::Float64 = β[1]
    μsp::Float64 = β[2]
    D::Float64 = 1/3μsp
    ν::Float64 = 29.9792345/nmed

    jϕ = Array{Float64}(undef, length(t), 2)

    Threads.@threads for n in eachindex(t)

        aux = exp(-(ρ^2/(4*D*ν*t[n])) - μa*ν*t[n])

        jϕ[n, 1] = -3^(3/2)*t[n]*ν^2*μsp^(3/2)*aux/(8*(π*t[n]*ν)^(3/2))
        jϕ[n, 2] = 3^(5/2)*sqrt(μsp)*(2*t[n]*ν - ρ^2*μsp)/(32*(t[n]*π*ν)^(3/2)*t[n])*aux

    end

    return jϕ
end

t = 0.01:0.01:3

ydata =  fluence_DA_inf_TD(t, [0.01,1.], 1.0) # + 0.00001*randn(length(t))

function curvetest(ydata, t)

    fit = curve_fit((t, β) -> fluence_DA_inf_TD(t, β, 1.0), t, ydata, [0.1, 60.])
end

function curvetestj(ydata, t)
    fit = curve_fit((t, β) -> fluence_DA_inf_TD(t, β, 1.0), (t, β) -> j_fluence_DA_inf_TD(t, β, 1.0), t, ydata, [0.1, 60.])
end

function curvetestff(ydata, t)
    fit = curve_fit((t, β) -> fluence_DA_inf_TD(t, β, 1.0), t, ydata, [0.1, 60.]; autodiff=:finiteforward)
end

function j_refl_DA_slab_TD(t, β::Array{Float64,1}, ρ::Float64, ndet::Float64, nmed::Float64)
    n::Float64 = nmed/ndet
    μa::Float64 = β[1]
    μsp::Float64 = β[2]
  	ν::Float64 = 29.9792345/nmed
    Rt = Array{Float64}(undef, length(t))

	if n > 1.0
		  A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		  5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
	elseif n < 1.0
		  A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
	else
		  A = 1.0
	end

    jϕ = Array{Float64}(undef, length(t), 2)

    Threads.@threads for n in eachindex(t)

        jϕ[n, 1] = (3*μsp)^(3/2)*exp(-t[n]*ν*μa - 3*ρ*μsp/(4*t[n]*ν))
        jϕ[n, 1] *= (-(4*A + 3)*exp(-(4*A + 3)^2/(12*t[n]*ν*μsp))/(3*μsp) - exp(-3/(4*t[n]*ν*μsp))/μsp)
        jϕ[n, 1] /= 16*π^(3/2)*t[n]^(3/2)*sqrt(ν)

        jϕ[n, 2] = ((27*ρ*μsp^2 - 18*t[n]*ν*μsp - 27)*exp((4*A + 3)^2/(12*t[n]*ν*μsp)))
        jϕ[n, 2] += ((36*A + 27)*ρ*μsp^2 + (-24*A - 18)*t[n]*ν*μsp - 64*A^3 - 144*A^2 - 108*A - 27)*exp(3/(4*t[n]*ν*μsp))
        jϕ[n, 2] *= -exp(-3*ρ*μsp/(4*t[n]*ν) - (4*A + 3)^2/(12*ν*t[n]*μsp) - 3/(4*t[n]*ν*μsp) - t[n]*ν*μa)
        jϕ[n, 2] /= 64*sqrt(3)*(π*μsp)^(3/2)*t[n]^(7/2)*ν^(5/2)
    end
    return jϕ

end

t = 0.01:0.01:3

ydata =  refl_DA_slab_TD(t, [0.01,1.], 1.0, 1.0, 1.0) # + 0.00001*randn(length(t))

function curvetest(ydata, t)

    fit = curve_fit((t, β) -> refl_DA_slab_TD(t, β, 1.0, 1.0, 1.0), t, ydata, [0.1, 20.])
end

#this is much slower.... jacobian takes >3x longer to calculate
function curvetestj(ydata, t)
    fit = curve_fit((t, β) -> refl_DA_slab_TD(t, β, 1.0, 1.0, 1.0), (t, β) -> j_refl_DA_slab_TD(t, β, 1.0, 1.0, 1.0), t, ydata, [0.1, 20.])
end

function curvetestff(ydata, t)
    fit = curve_fit((t, β) -> refl_DA_slab_TD(t, β, 1.0, 1.0, 1.0), t, ydata, [0.1, 20.]; autodiff=:finiteforward)
end


using Test

using SymbolicWedderburn
using SymbolicWedderburn.FiniteFields
using LinearAlgebra

using PermutationGroups
using Cyclotomics

include("gf.jl")
include("eigenspacedecomposition.jl")
include("ccmatrix.jl")
include("dixon.jl")


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

using Combinatorics
import JuLIP: evaluate

struct PureBasis{TB}
   basis::TB
   Ibas::Vector{Int}
end

Base.length(b::PureBasis) = length(b.Ibas)

(b::PureBasis)(args...) = evaluate(b, args...)

PureBasis(ship::SHIPBasis) = PureBasis(ship, collect(1:length(ship)))

function PureBasis(ship::SHIPBasis, N::Integer)
   @assert length(ship.zlist) == 1
   Ibas = ACE.Utils.findall_basis_N(ship, N)
   return PureBasis(ship, Ibas)
end

function evaluate(b::PureBasis, Rs::AbstractVector)
   N = length(Rs)
   ship = b.basis
   @assert N <= 5
   @assert length(ship.zlist) == 1
   Zs = zeros(Int16, N)
   tmp = alloc_temp(ship, Rs)

   # one-particle basis
   phi = precompute_phi(tmp, Rs, Zs, ship)

   # permutation symmetrised tensor products
   Phi = precompute_prod_phi(tmp, phi, ship)

   # get B via the same rotation-symmetrisation as in the density trick case
   _my_mul!(tmp.Bc[1], ship.A2B[1], Phi)
   return real.(tmp.Bc[1])[b.Ibas]
end

function precompute_phi(tmp, Rs, Zs, ship)
   alist = ship.alists[1]
   phi = zeros(Complex{Float64}, length(alist), length(Rs))
   for (iR, (R, Z)) in enumerate(zip(Rs, Zs))
      # evaluate the r-basis and the R̂-basis for the current neighbour at R
      evaluate!(tmp.J, tmp.tmpJ, ship.J, norm(R))
      evaluate!(tmp.Y, tmp.tmpY, ship.SH, R)
      # write the products into phi
      iz = z2i(ship, Z)
      for i = alist.firstz[iz]:(alist.firstz[iz+1]-1)
         zklm = alist[i]
         phi[i, iR] = tmp.J[zklm.k+1] * tmp.Y[index_y(zklm.l, zklm.m)]
      end
   end
   return phi
end

function precompute_prod_phi(tmp, phi, ship)
   aalist = ship.aalists[1]
   Phi = fill!(tmp.AA[1], 0)
   for i = 1:length(aalist)
      Phi[i] = _VN(phi, (@view aalist.i2Aidx[i, 1:aalist.len[i]]))
   end
   return Phi
end

function _VN(phi, phiidx)
   nR = size(phi, 2)     # number of neighbours
   N  = length(phiidx)   # order of the basis function
   Phi = zero(ComplexF64)
   for A in combinations(1:nR, N), σA in permutations(A)
      Phi += prod( phi[phiidx[α], σα]  for (α, σα) in enumerate(σA) )
   end
   return Phi / factorial(N)
end


# find hyperparameter values using gird search

include(pwd()*"/dlmodel.jl")

p = [0]

lambda_1 = collect(0:0.02:0.5)
lambda_2 = collect(0:0.02:0.5)
lambda_3 = [0]

model_fit_regularization = zeros(length(p),length(lambda_1),length(lambda_2),length(lambda_3),3+3+1+1)
model_fit_train = zeros(length(p),length(lambda_1),length(lambda_2),length(lambda_3),2000)
model_fit_val = zeros(length(p),length(lambda_1),length(lambda_2),length(lambda_3),2000)

# find hyperparameters using gridsearch
@time for i in 1:length(p)
    for j in 1:length(lambda_1)
        @printf "Percentage done %.2f\n" (j-1)/length(lambda_1)
        for k in 1:length(lambda_2)
            for l in 1:length(lambda_3)

                Random.seed!(1234) # set random numbers

                w = Any[ xavier(n_hidden_1,n_input), zeros(n_hidden_1,1),
                         #xavier(n_hidden_2,n_hidden_1), zeros(n_hidden_2,1),
                         xavier(n_out,n_hidden_1), zeros(n_out,1)]

                optim = optimizers(w, Adam)

                dropout_percentage = p[i]
                lambda = [lambda_1[j], lambda_2[k]]
                loss_train_vec, loss_val_vec = train(2000, dropout_percentage, lambda, w, optim)

                model_fit_train[i,j,k,l,:] = loss_train_vec
                model_fit_val[i,j,k,l,:] = loss_val_vec

                # calc predictions
                y_pred_train_proc = predict(w, x_train, 0)
                y_pred_val_proc = predict(w, x_val, 0)

                y_pred_train = class_nll(y_pred_train_proc,[1,2])
                y_pred_val = class_nll(y_pred_val_proc,[1,2])

                # loss
                model_fit_regularization[i,j,k,l,end-1] = loss(w,x_train, y_train, 0, lambda)/length(y_train)
                model_fit_regularization[i,j,k,l,end] = loss(w,x_val, y_val, 0, lambda)/length(y_val)

                # classification results
                model_fit_regularization[i,j,k,l,1:3] = classification_results(y_pred_train, y_train, [1,2])
                model_fit_regularization[i,j,k,l,4:6] = classification_results(y_pred_val, y_val, [1,2])

            end
        end
    end
end

# results without regularization
training_loss_no_reg = model_fit_regularization[1,1,1,1,end-1]
test_loss_no_reg = model_fit_regularization[1,1,1,1,end]
training_acc_no_reg = model_fit_regularization[1,1,1,1,1]
test_acc_no_reg = model_fit_regularization[1,1,1,1,4]

# find best hyperparameter settings
(~, idx_best_training_loss) = findmin(model_fit_regularization[:,:,:,:,end-1])
(~, idx_best_val_loss) = findmin(model_fit_regularization[:,:,:,:,end])

(~, idx_best_training_tp) = findmax(model_fit_regularization[:,:,:,:,2])
(~, idx_best_val_tp) = findmax(model_fit_regularization[:,:,:,:,5])

(~, idx_best_training_tn) = findmax(model_fit_regularization[:,:,:,:,3])
(~, idx_best_val_tn) = findmax(model_fit_regularization[:,:,:,:,6])

(~, idx_best_training_acc) = findmax(model_fit_regularization[:,:,:,:,1])
(~, idx_best_val_acc) = findmax(model_fit_regularization[:,:,:,:,4])

# 5,1,1,5
# 3,1,8,1

@printf "Best dropout percentage %.2f\n" p[1]
@printf "Best lambda_1 %.10f\n" lambda_1[1]
<<<<<<< HEAD
@printf "Best lambda_2 %.10f\n" lambda_2[4]
@printf "Best lambda_3 %.10f\n" lambda_3[1]
=======
@printf "Best lambda_2 %.10f\n" lambda_2[1]
@printf "Best lambda_3 %.10f\n" lambda_3[11]
>>>>>>> 2layernetwork


mutable struct LIF{VT<:Real, IT<:Integer} <: AbstractCell
    # required fields
    voltage::VT
    current::VT

    # model specific fields
    lastt::IT
    τm::VT
    vreset::VT
    R::VT
end

mutable struct ADEXP{VT<:Real, IT<:Integer} <: AbstractCell

    # {VFT=Vector{Float32},VBT=Vector{Bool}}  <: AbstractCell
    # required fields
    voltage::VT
    current::VT

    # model specific fields
    lastt::IT
    τm::VT
    vreset::VT
    R::VT
    #    param::ADEXParameter = ADEXParameter()
    #end
    #@with_kw mutable struct AD
    param::ADEXParameter = ADEXParameter(a,
                                        b,
                                        cm,
                                        v_rest,
                                        tau_m,
                                        tau_w,
                                        v_thresh,
                                        delta_T,
                                        v_reset,
                                        spike_height)
    N::SNNInt = 1
    cnt::SNNInt = 2
    v::Vector{SNNFloat} = fill(param.v_rest, N)
    w::Vector{SNNFloat} = zeros(N)
    fire::Vector{Bool} = zeros(Bool, N)
    I::Vector{SNNFloat} = zeros(N)
    spike_raster::Vector{SNNInt} = zeros(N)
    records::Dict = Dict()
end

@snn_kw struct ADEXParameter{FT=Float32}
    a::FT = 4.0
    b::FT = 0.0805
    cm::FT = 0.281
    v_rest::FT = -70.6
    tau_m::FT = 9.3667
    tau_w::FT = 144.0
    v_thresh::FT = -50.4
    delta_T::FT = 2.0
    v_reset::FT = -70.6
    spike_height::FT = 30
end

"""
LIF(τm, vreset, vth, R = 1.0)
Create a ADEXP neuron with zero initial voltage and empty current queue.
"""
LIF(τm::Real, vreset::Real, R::Real = 1.0) = LIF{Float32, Int}(vreset, 0, 0, τm, vreset, R)

isactive(neuron::LIF, t::Integer; dt::Real = 1.0) = (neuron.current > 0)
getvoltage(neuron::LIF) = neuron.voltage

"""
    excite!(neuron::LIF, current)
    excite!(neurons::AbstractArray{<:LIF}, current)
Excite a `neuron` with external `current`.
"""
excite!(neuron::LIF, current) = (neuron.current += current)
excite!(neurons::T, current) where T<:AbstractArray{<:LIF} = (neurons.current .+= current)

"""
    spike!(neuron::LIF, t::Integer; dt::Real = 1.0)
    spike!(neurons::AbstractArray{<:LIF}, spikes; dt::Real = 1.0)
Record a output spike from the threshold function with the `neuron` body.
Sets `neuron.lastt`.
"""
spike!(neuron::LIF, t::Integer; dt::Real = 1.0) = (t > 0) && (neuron.voltage = neuron.vreset)
function spike!(neurons::T, spikes; dt::Real = 1.0) where T<:AbstractArray{<:LIF}
    map!((x, y, s) -> (s > 0) ? y : x, neurons.voltage, neurons.voltage, neurons.vreset, spikes)
end

"""
    lif(t::Real, I, V; R, tau)
    lif!(t::AbstractArray{<:Real}, I::AbstractArray{<:Real}, V::AbstractArray{<:Real}; R::AbstractArray{<:Real}, tau::AbstractArray{<:Real})
    lif!(t::CuVector{<:Real}, I::CuVector{<:Real}, V::CuVector{<:Real}; R::CuVector{<:Real}, tau::CuVector{<:Real})
Evaluate a leaky integrate-and-fire neuron.
Use `CuVector` instead of `Vector` to evaluate on GPU.
# Fields
- `t`: time since last evaluation in seconds
- `I`: external current
- `V`: current membrane potential
- `vrest`: resting membrane potential
- `R`: resistance constant
- `tau`: time constant
"""
function lif(t::Real, I, V; vrest, R, tau)
    V = V * exp(-t / tau) + vrest
    V += I * (R / tau)

    return V
end
function lif!(t::AbstractArray{<:Real}, I::AbstractArray{<:Real}, V::AbstractArray{<:Real}; vrest::AbstractArray{<:Real}, R::AbstractArray{<:Real}, tau::AbstractArray{<:Real})
    # account for leakage
    @. V = V * exp(-t / tau) + vrest

    # apply update step
    @. V += I * (R / tau)

    return V
end
function lif!(t::CuVector{<:Real}, I::CuVector{<:Real}, V::CuVector{<:Real}; vrest::CuVector{<:Real}, R::CuVector{<:Real}, tau::CuVector{<:Real})
    # account for leakage
    @. V = V * exp(-t / tau) + vrest

    # apply update step
    @. V += I * (R / tau)

    return V
end

"""
    evaluate!(neuron::LIF, t::Integer; dt::Real = 1.0)
    (neuron::LIF)(t::Integer; dt::Real = 1.0)
    evaluate!(neurons::AbstractArray{<:LIF}, t::Integer; dt::Real = 1.0)
Evaluate the neuron model at time `t`. Return the resulting membrane potential.
"""
function evaluate!(neuron::LIF, t::Integer; dt::Real = 1.0)
    neuron.voltage = lif((t - neuron.lastt) * dt, neuron.current, neuron.voltage; vrest = 0, R = neuron.R, tau = neuron.τm)
    neuron.lastt = t
    neuron.current = 0

    return neuron.voltage
end
(neuron::LIF)(t::Integer; dt::Real = 1.0) = evaluate!(neuron, t; dt = dt)
function evaluate!(neurons::T, t::Integer; dt::Real = 1.0) where T<:AbstractArray{<:LIF}
    lif!((t .- neurons.lastt) .* dt, neurons.current, neurons.voltage;
         vrest = zero(neurons.voltage), R = neurons.R, tau = neurons.τm)
    neurons.lastt .= t
    neurons.current .= 0

    return neurons.voltage
end

"""
    reset!(neuron::LIF)
    reset!(neurons::AbstractArray{<:LIF})
Reset `neuron` by setting the membrane potential to `neuron.vreset`.
"""
function reset!(neuron::LIF)
    neuron.lastt = 0
    neuron.voltage = neuron.vreset
    neuron.current = 0
end
function reset!(neurons::T) where T<:AbstractArray{<:LIF}
    neurons.lastt .= 0
    neurons.voltage .= neurons.vreset
    neurons.current .= 0
end

Base.show(io::IO, ::MIME"text/plain", neuron::LIF) =
    print(io, """LIF:
                     voltage: $(neuron.voltage)
                     current: $(neuron.current)
                     τm:      $(neuron.τm)
                     vreset:  $(neuron.vreset)
                     R:       $(neuron.R)
    """)
Base.show(io::IO, neuron::LIF) =
    print(io, "LIF(voltage = $(neuron.voltage), current = $(neuron.current))")


using SparseArrays
using LinearAlgebraicRepresentation
LAR = LinearAlgebraicRepresentation
using ViewerGL
GL = ViewerGL

V1, (VV,EV,FV,CV) = LAR.simplex(3, true)
tetra1 = V1, EV, FV, CV
V2 = [
  0.0 0.5 0.1 0.1
  0.0 0.1 0.5 0.1
  0.0 0.1 0.1 0.5
]
tetra2 = V2, EV, FV, CV
twotetra = LAR.Struct([ tetra, tetra2 ])
V,EV,FV,CV = LAR.struct2lar(twotetra)
GL.VIEW([ GL.GLGrid(V,FV, GL.Point4d(1,1,1,0.2)) ]);
m = Lar.Model(V);
Lar.addModelCells!(m, 1, convert(LAR.ChainOp, LAR.coboundary_0(EV::LAR.Cells)))
Lar.addModelCells!(m, 2, LAR.coboundary_1(V, FV::LAR.Cells, EV::LAR.Cells))
Lar.viewModel(m,1)

sp_idx = LarA.spaceIndex(m, 2)

de_models = [
	LarA.face_decomposition(m, face_idx, sp_idx[face_idx])
	for face_idx = 1 : size(m, 2, 1)
]

mfrag = Lar.mergeMultipleModels(de_models)

triangulated_faces = LAR.triangulate(convert(Lar.Points, mfrag.G'), mfrag.T[1:2])
FVs = convert(Array{Lar.Cells}, triangulated_faces);
GL.VIEW( GL.GLExplode(mfrag.G,FVs,expl,expl,expl,99,1) );

mArranged = Lar.deepcopy(mfrag)
Lar.addModelCells!(mArranged, 3, LAR.Arrangement.minimal_3cycles(
		convert(Lar.Points, mfrag.G'), mfrag.T[1], mfrag.T[2])
)

V,CVs,FVs,EVs = LAR.pols2tria(mArranged.G, mArranged.T[1], mArranged.T[2], mArranged.T[3])

GL.VIEW(GL.GLExplode(V,FVs,1.5,1.5,1.5,99,1));
GL.VIEW(GL.GLExplode(V,EVs,1.5,1.5,1.5,99,1));
meshes = GL.GLExplode(V,CVs[1:end],8,4,6,99,0.5);
push!( meshes, GL.GLFrame)
GL.VIEW( meshes );

#=
cop_EV = LAR.coboundary_0(EV::LAR.Cells);
cop_EW = convert(LAR.ChainOp, cop_EV);
cop_FE = LAR.coboundary_1(V, FV::LAR.Cells, EV::LAR.Cells);
W = convert(LAR.Points, V');

V, copEV, copFE, copCF = LAR.Arrangement.spatial_arrangement( W::LAR.Points, cop_EW::LAR.ChainOp, cop_FE::LAR.ChainOp)

EV = LAR.cop2lar(copEV)
FE = [findnz(copFE[k,:])[1] for k=1:size(copFE,1)]
FV = [collect(Set(cat(EV[e] for e in FE[f]))) for f=1:length(FE)]
FV = convert(LAR.Cells, FV)
W = convert(LAR.Points, V')
WW = [[k] for k=1:size(W,2)]

GL.VIEW(GL.numbering(.5)((W,[WW,EV]) ));

triangulated_faces = LAR.triangulate(V, [copEV, copFE])
FVs = convert(Array{LAR.Cells}, triangulated_faces)
V = convert(LAR.Points, V')
GL.VIEW(GL.GLExplode(V,FVs,1.5,1.5,1.5,99));

EVs = LAR.FV2EVs(copEV, copFE) # polygonal face fragments
GL.VIEW(GL.GLExplode(V,EVs,1.5,1.5,1.5,99,1));


module BinaryDecSSA

    """
    Need to more properly document and add ability to plot/output m intstead of N₊.
    """

    export BinDecMod, Ensemble, EnsProb, CentralMoments

    using Distributions, Parameters, StatsBase, LinearAlgebra, Colors

    """
    Define the binary decision model struct for the SSA.
    """
    @with_kw struct BinDecMod
        N::Int64 = 100
        pars::Vector{Float64}; # can be [ϵ,μ], [ϵ1,ϵ2,μ] or [ϵ1,ϵ2,μ1,μ2]
        fixedIC::Bool = true
        p₀::Float64 = 0.5 # initial prob of decided to be on rh food source if fixedIC = false.
    end

    """
    Define the struct for an ensemble of sims.
    """
    @with_kw struct Ensemble
        BD::BinDecMod
        e_size::Int64 = 1 # size of ensemble
        τₘ::Real = 1E3 # max system time
        Δτ::Real = 1.0 # time spacing of storage in trajectory
        sims = SSA(e_size, τₘ, Δτ, BD)
        @assert e_size >= 1
        @assert τₘ > 0
        @assert Δτ > 0 && Δτ < τₘ
    end

    """
    Propensities.
    """
    function props(n::Int64, BD::BinDecMod)
        @unpack pars, N = BD;
        props = zeros(2) # propensity vector.
        if length(pars)==2
            ϵ = pars[1];
            μ = pars[2];
            props[1] = (N-n)*ϵ+μ*n*(N-n)/(N-1);
            props[2] = n*ϵ+μ*n*(N-n)/(N-1);
            return props
        elseif length(pars)==3
            ϵ1 = pars[1];
            ϵ2 = pars[2];
            μ = pars[3];
            props[1] = (N-n)*ϵ1+μ*n*(N-n)/(N-1);
            props[2] = n*ϵ2+μ*n*(N-n)/(N-1);
            return props
        elseif length(pars)==4
            ϵ1 = pars[1];
            ϵ2 = pars[2];
            μ1 = pars[3];
            μ2 = pars[4];
            props[1] = (N-n)*ϵ1+μ1*n*(N-n)/(N-1);
            props[2] = n*ϵ2+μ2*n*(N-n)/(N-1);
            return props
        else
            error("length of pars must be 2, 3 or 4.")
        end
    end

    """
    Mean-field fully connected SSA: species is n. Does much faster mf SSA sims.
    """
    function SSA(ens_its::Int64, τₘ::Real, Δτ::Real, BD::BinDecMod)
        @unpack p₀, fixedIC, N = BD;
        # create the storage array for the ensembles
        ens_store = Array{Array{Real}}(undef, ens_its);
        # assign the initial condition - either fixed or binomially drawn.
        if fixedIC == true
            n = ceil(Int64, N*p₀)
        else
            n = rand(Binomial(N,p₀),1)[1]
        end
        # define the sim times for storage of n
        times = convert(Array{Float64,1},LinRange(τₘ,0.0,floor(Int,τₘ/Δτ)+1))
        # loop over for each sim in the ensemble
        for it in 1:ens_its
            # n at t=0 is either fixed or binomially drawn.
            fixedIC == true ? n = ceil(Int64, N*p₀) : n = rand(Binomial(N,p₀),1)[1]
            # initialise the state vector and sys time
            sim_times = copy(times)
            store = zeros(length(times));
            τ = 0.0;
            m = 1; # counter for updating storage
            while τ < τₘ
                ps = props(n, BD); # get propensities
                f = sum(ps);
                r1, r2 = rand(2); # collect 2 ind rand nums ∈ [0,1]
                u = (1/f)*log(1/r1); # waiting time until next reaction
                k = findfirst(x -> x>=r2*f,cumsum(ps)); # find fired reaction
                # update state vector storage
                while τ+u >= sim_times[length(sim_times)]
                    store[m] = copy(n)
                    pop!(sim_times)
                    m += 1
                    if length(sim_times) == 0
                        break
                    end
                end
                # update the state variable
                k == 1 ? n += 1 : n -= 1;
                # update the time.
                τ += u;
            end
            ens_store[it] = store ./ N; # return the value of m not N₊
        end
        return (reverse(times),ens_store)
    end

    """
    Get the central moments from the ensemble
    """
    function CentralMoments(ens::Ensemble,n::Int64)
        ns = convert(Matrix{Float64},transpose(cat(ens.sims[2]...,dims=2)))
        store = Array{Float64,1}(undef, length(ns[1,:]))
        if n == 1
            for t in 1:length(ns[1,:])
                store[t] = mean(ns[:,t])
            end
        else
            for t in 1:length(ns[1,:])
                store[t] = moment(ns[:,t],n)
            end
        end
        # need to sort out times
        @unpack Δτ = ens;
        times = LinRange(0:Δτ:length(ns[1,:])-1);
        return (times, store)
    end

    """
    Get prob distribution at all times for an ensemble
    """
    function EnsProb(ens::Ensemble, T::Int64)
        @unpack BD = ens;
        @unpack N = BD;
        ns = convert(Matrix{Float64},transpose(cat(ens.sims[2]...,dims=2)))
        mod_bins = LinRange(0.0-(1/N),1.0+(1/N),N+2);
        mid_pts = LinRange(0.0,1.0,N+1);
        bin_vals = normalize(fit(Histogram, ns[:,T], mod_bins), mode=:probability).weights;
        return (mid_pts, bin_vals.*N)
    end

end # module end


module OceanGrids

using Unitful
using Unitful: °
using Interpolations
using LinearAlgebra, SparseArrays
using NearestNeighbors
using Inpaintings

abstract type OceanGrid end

"""
    OceanRectilinearGrid

Ocean rectilinear grid.
Here I use the slightly abuse the term rectilinear,
because the grid supposedly represents layers around the (curved) globe.
Here "rectilinear" means that the grid is defined entirely by just its
`lat`, `lon`, `depth`, `δlat`, `δlon`, and `δdepth` vectors.
For example, the Ocean Circulation Inverse Model (OCIM) version 1
grid is rectilinear in that sense.
"""
struct OceanRectilinearGrid <: OceanGrid
    lat          # °
    lon          # °
    depth        # m
    δlat         # °
    δlon         # °
    δdepth       # m
    lat_3D       # °
    lon_3D       # °
    depth_3D     # m
    δy           # m
    δx_3D        # m
    δy_3D        # m
    δz_3D        # m
    volume_3D    # m³
    depth_top    # m
    depth_top_3D # m
    A_2D         # m²
    wet3D::BitArray
    nlon
    nlat
    ndepth
    nboxes
end

"""
    OceanCurvilinearGrid

Ocean curvilinear grid.
The grid is curvilinear if it can be accessed using Cartesian Indices.
I.e., it maps to a rectilinear grid.
This type of grid typically cannot be inferred from only
`lat`, `lon`, `depth`, `δlat`, `δlon`, and `δdepth` vectors.
Instead, the full 3-dimensional arrays
`lat_3D`, `lon_3D`, `D`epth_3D`, `δx_3D`, `δy_3D`, and `δz_3D`
are required.
For exaample, the displaced-pole grid of the POP model is curcvilinear.
"""
struct OceanCurvilinearGrid <: OceanGrid
    lat_3D       # °
    lon_3D       # °
    depth_3D     # m
    δx_3D        # m
    δy_3D        # m
    δz_3D        # m
    volume_3D    # m³
    depth_top    # m
    depth_top_3D # m
    A_2D         # m²
    wet3D::BitArray
    nlon
    nlat
    ndepth
    nboxes
end

const TU = AbstractArray{<:Quantity}
const T1 = AbstractArray{<:Number}

"""
    OceanGrid(elon::T, elat::U, edepth::V; R=6371.0u"km") where {T<:TU, U<:TU, V<:TU}

Returns an `OceanRectilinearGrid` with boxes whose edges are defined by the
`elon`, `elat`, and `edepth` vectors.
The globe radius can be changed with the keyword `R` (default value 6371 km)
The 3D array of wet boxes, `wet3D` is true everywhere by default.
"""
function OceanGrid(elon::TU, elat::TU, edepth::TU; R=6371.0u"km")
    unitful_data, unitless_data = generate_grid_data_no_wet3D(elon, elat, edepth, R)
    #
    nlat, nlon, ndepth = unitless_data[[2,1,3]]
    wet3D = trues(nlat, nlon, ndepth)
    return OceanRectilinearGrid(unitful_data..., wet3D, unitless_data...)
end

"""
    OceanGrid(elon::T, elat::U, edepth::V, wet3D::BitArray; R=6371.0u"km")

Same as `OceanGrid(elon, elat, edepth; R)` but with specified `wet3D`.
"""
function OceanGrid(elon::TU, elat::TU, edepth::TU, wet3D::BitArray; R=6371.0u"km")
    unitful_data, unitless_data = generate_grid_data_no_wet3D(elon, elat, edepth, R)
    return OceanRectilinearGrid(unitful_data..., wet3D, unitless_data...)
end

function generate_grid_data_no_wet3D(elon::TU, elat::TU, edepth::TU, R)
    R = R |> u"m"
    edepth = edepth .|> u"m"
    elat = elat .|> °
    edon = elon .|> °
    # δ objects
    nlon, nlat, ndepth = length(elon) - 1, length(elat) - 1, length(edepth) - 1
    nboxes = nlon * nlat * ndepth
    lon = edges_to_centers(elon)
    lat = edges_to_centers(elat)
    depth = edges_to_centers(edepth)
    δlon, δlat, δdepth = diff(elon), diff(elat), diff(edepth)
    # depth objects
    δz_3D = repeat(reshape(δdepth, (1,1,ndepth)), outer=(nlat,nlon,1))
    depth_3D = repeat(reshape(depth, (1,1,ndepth)), outer=(nlat,nlon,1))
    depth_top = cumsum(δdepth) - δdepth
    depth_top_3D = repeat(reshape(depth_top, (1,1,ndepth)), outer=(nlat,nlon,1))
    # lat objects
    lat_3D = repeat(reshape(lat, (nlat,1,1)), outer=(1,nlon,ndepth))
    δy = R * δlat ./ 360°
    δy_3D = repeat(reshape(δy, (nlat,1,1)), outer=(1,nlon,ndepth))
    # lon objects
    lon_3D = repeat(reshape(lon, (1,nlon,1)), outer=(nlat,1,ndepth))
    # For δx_3D, first calculate the area and then infer the equivalent distance in meters
    A_2D = R^2 * abs.(sin.(elat[1:end-1]) - sin.(elat[2:end])) * ustrip.(δlon .|> u"rad")'
    A_3D = repeat(A_2D, outer=(1, 1, ndepth))
    δx_3D = A_3D ./ δy_3D
    # volume
    volume_3D = δx_3D .* δy_3D .* δz_3D
    unitful_data = (lat, lon, depth, δlat, δlon, δdepth, lat_3D, lon_3D, depth_3D,
                    δy, δx_3D, δy_3D, δz_3D, volume_3D, depth_top, depth_top_3D, A_2D)
    unitless_data = (nlon, nlat, ndepth, nboxes)
    return unitful_data, unitless_data
end

# Convert units in case not provided
function OceanGrid(elon::TU, elat::TU, edepth::T1; R=6371.0u"km")
    @warn "`edepth` provided without units — assuming meters"
    return OceanGrid(elon, elat, edepth * u"m"; R=R)
end
function OceanGrid(elon::T1, elat::TU, edepth::TU; R=6371.0u"km")
    @warn "`elon` provided without units — assuming degrees"
    return OceanGrid(elon * °, elat, edepth; R=R)
end
function OceanGrid(elon::TU, elat::T1, edepth::TU; R=6371.0u"km")
    @warn "`elat` provided without units — assuming degrees"
    return OceanGrid(elon, elat * °, edepth; R=R)
end
function OceanGrid(elon::T1, elat::T1, edepth::TU; R=6371.0u"km")
    @warn "`elat` and `elon` provided without units — assuming degrees"
    return OceanGrid(elon * °, elat * °, edepth; R=R)
end
function OceanGrid(elon::T1, elat::T1, edepth::T1; R=6371.0u"km")
    @warn "`elat`, `elon`, and `edepth` provided without units — assuming degrees and meters"
    return OceanGrid(elon * °, elat * °, edepth * u"m"; R=R)
end
function OceanGrid(elon::TU, elat::T1, edepth::T1; R=6371.0u"km")
    @warn "`elat` and `edepth` provided without units — assuming and meters"
    return OceanGrid(elon, elat * °, edepth * u"m"; R=R)
end
function OceanGrid(elon::T1, elat::TU, edepth::T1; R=6371.0u"km")
    @warn "`elon` and `edepth` provided without units — assuming degrees and meters"
    return OceanGrid(elon * °, elat, edepth * u"m"; R=R)
end

edges_to_centers(x::Vector) = 0.5 * (x[1:end-1] + x[2:end])
edges_to_centers(x::AbstractRange) = x[1:end-1] .+ 0.5step(x)

"""
    OceanGrid(nlat::Int, nlon::Int, ndepth::Int)

Returns a regularly spaced `OceanRectilinearGrid` with size `nlat`, `nlon`, and `ndepth`.
"""
function OceanGrid(nlat::Int, nlon::Int, ndepth::Int)
    elat = range(-90,90,length=nlat+1) * °
    elon = range(0,360,length=nlon+1) * °
    edepth = range(0,3682,length=ndepth+1) * u"m"
    return OceanGrid(elon, elat, edepth)
end

function Base.show(io::IO, g::OceanGrid)
    print("OceanGrid of size $(g.nlat)×$(g.nlon)×$(g.ndepth) (lat×lon×depth)")
end

"""
    OceanGridBox

Ocean grid box.
Each grid can be looped over, where the `OceanGridBox` are the elements of the grid.
This useful to investigate specific boxes in the grid.
"""
struct OceanGridBox
    I
    lat          # °
    lon          # °
    depth        # m
    δlat         # °
    δlon         # °
    δdepth       # m
    δx           # m
    δy           # m
    δz           # m
    volume       # m³
    depth_top    # m
    A            # m²
    iswet::Bool
end

"""
    box(g::OceanGrid, i, j, k)

Accesses the individual box of `g::OceanGrid` at index `(i,j,k)`.
Each grid can be looped over, where the `OceanGridBox` are the elements of the grid.
This useful to investigate specific boxes in the grid.
"""
function box(g::OceanGrid, i, j, k)
    return OceanGridBox(
                        CartesianIndex(i,j,k),
                        g.lat[i],
                        g.lon[j],
                        g.depth[k],
                        g.δlat[i],
                        g.δlon[j],
                        g.δdepth[k],
                        g.δx_3D[i,j,k],
                        g.δy_3D[i,j,k],
                        g.δz_3D[i,j,k],
                        g.volume_3D[i,j,k],
                        g.depth_top[k],
                        g.A_2D[i,j],
                        g.wet3D[i,j,k]
                       )
end

"""
    box(g::OceanGrid, I)

Accesses the individual box of `g::OceanGrid` at index `I`.
Each grid can be looped over, where the `OceanGridBox` are the elements of the grid.
This useful to investigate specific boxes in the grid.
"""
function box(g::OceanGrid, I)
    i,j,k = CartesianIndices((g.nlat,g.nlon,g.ndepth))[I].I
    return box(g::OceanGrid, i, j, k)
end

Base.getindex(g::OceanGrid, I) = box(g,I)

"""
    Base.size(g::OceanGrid)

Size of the grid.
"""
Base.size(g::OceanGrid) = g.nlat, g.nlon, g.ndepth

"""
    Base.length(g::OceanGrid)

Size of the grid.
"""
Base.length(g::OceanGrid) = g.nlat * g.nlon * g.ndepth

function Base.show(io::IO, b::OceanGridBox)
    println("$(wet_print(b.iswet)) OceanGridBox at $(b.I):")
    println("  location: $(round(b.lat,digits=1))N, $(round(b.lon,digits=1))E")
    println("  depth: $(round(b.depth,digits=1))")
    println("  size: $(round(b.δx |> u"km",digits=1)) × $(round(b.δy |> u"km",digits=1)) × $(round(b.δz,digits=1)) (δx × δy × δz)")
end

wet_print(iswet) = iswet ? "Wet" : "Dry"

Base.round(q::Quantity; digits=0) = round(q |> ustrip, digits=digits) * unit(q)

area(b::OceanGridBox) = b.A |> u"km^2"
volume(b::OceanGridBox) = b.volume

Base.iterate(g::OceanGrid) = box(g::OceanGrid, 1), 1
function Base.iterate(g::OceanGrid, i)
    if i == g.nboxes
        return nothing
    else
        return box(g::OceanGrid, i+1), i+1
    end
end

Base.isequal(g₁::T, g₂::T) where {T<:OceanGrid} = prod([getfield(g₁,f) == getfield(g₂,f) for f in fieldnames(T)])
Base.:(==)(g₁::T, g₂::T) where {T<:OceanGrid} = isequal(g₁, g₂)

export OceanGrid, OceanCurvilinearGrid, OceanRectilinearGrid, box, OceanGridBox

include("functions.jl")

end # module


export RCWAplot

"""
    plot(::DeltaRCWASolution{1})

Visualize the RCWA solution in terms of the incident and scattered fields

Keyword arguments:
- combine{Bool} :: whether or not to plot the total fields or their components
- method{Symbol} :: choose between :series and :fft (default, fallback)
- part{Function} :: A function that gives a real from a complex (default: real)
"""
RCWAplot

@recipe function f(sol::DeltaRCWASolution{1}; combine=false, part=real, method=:fft)
    # Nz = sol.modes.dims[1][1]
    # dz = sol.modes.dims[1][2] / Nz
    # z⃗ = range(0, step=dz, length=Nz)
    pw, stack, I₁, I₂, O₁, O₂ = sol.pw, sol.stack, sol.I₁, sol.I₂, sol.O₁, sol.O₂
    x⃗ = pw.x⃗[1]
    z⃗₂ = x⃗
    z⃗₁ = -reverse(z⃗₂)
    kz₁, kz₂ = _get_kz.(Ref(pw), stack.media[[1, end]])
    if method == :series
        AO₂ = zeros(length.((x⃗, z⃗₂)))
        AI₁ = zeros(length.((x⃗, z⃗₁)))
        AO₁ = zeros(length.((x⃗, z⃗₁)))
        AI₂ = zeros(length.((x⃗, z⃗₂)))
        k⃗ = pw.k⃗[1]
        for j in eachindex(x⃗)
            for i in eachindex(z⃗₂)
                AO₂[j, i] = part.(sum(exp.((z⃗₂[i] * im) .*  kz₂ .+ (x⃗[j] * im) .* k⃗) .* O₂))
                AI₂[j, i] = part.(sum(exp.((z⃗₂[i] * im) .* -kz₂ .+ (x⃗[j] * im) .* k⃗) .* I₂))
                AI₁[j, i] = part.(sum(exp.((z⃗₁[i] * im) .*  kz₁ .+ (x⃗[j] * im) .* k⃗) .* I₁))
                AO₁[j, i] = part.(sum(exp.((z⃗₁[i] * im) .* -kz₁ .+ (x⃗[j] * im) .* k⃗) .* O₁))
            end
        end
    else
        # using bfft here because it matches the summation
        # but in reality ifft (resp. bfft) are off by a factor of
        # sqrt(length(x⃗)) (resp. 1/sqrt(length(x⃗))) when it comes to preserving
        # the norm/energy of the frequency domain representation.
        # Note that fft doesn't preserve the norm (increases by sqrt(length(x⃗)))
        # Ultimately this is just a choice of convention that can be modified.
            AO₂ = part.(bfft(exp.( kz₂ * transpose(im * z⃗₂)) .* O₂, 1))
            AI₂ = part.(bfft(exp.(-kz₂ * transpose(im * z⃗₂)) .* I₂, 1))
            AI₁ = part.(bfft(exp.( kz₁ * transpose(im * z⃗₁)) .* I₁, 1))
            AO₁ = part.(bfft(exp.(-kz₁ * transpose(im * z⃗₁)) .* O₁, 1))
    end
    if combine
       @series begin
            seriestype := :heatmap
            title := "port 1 | port 2"
            xguide := "z"
            yguide := "x"
            vcat(z⃗₁, z⃗₂), x⃗, cat(AO₁ .+ AI₁, AO₂ .+ AI₂; dims=2)
        end
    else
        layout := 4
        # plot_title := String(part) * " part of $(typeof(sol.pol)) field"
        # subplots are in row-major order
        @series begin
            seriestype := :heatmap
            subplot := 1
            title := "Scattered, port 1"
            xaxis := nothing
            yguide := "x"
            z⃗₁, x⃗, AO₁
        end
        @series begin
            seriestype := :heatmap
            subplot := 2
            title := "Incident, port 2"
            xaxis := nothing
            yaxis := nothing
            z⃗₂, x⃗, AI₂
        end
        @series begin
            seriestype := :heatmap
            subplot := 3
            title := "Incident, port 1"
            xguide := "z"
            yguide := "x"
            z⃗₁, x⃗, AI₁
        end
        @series begin
            seriestype := :heatmap
            subplot := 4
            title := "Scattered, port 2"
            yaxis := nothing
            xguide := "z"
            z⃗₂, x⃗, AO₂
        end
    end
end


using Revise
using Random
using RDatasets
using ForwardDiff
using ForwardDiff: Dual
partiali(f, x, i, args...) = ForwardDiff.partials(f([Dual{}(x[j], 1.0*(i==j)) for j in eachindex(x)], args...))[]

using ZigZagBoomerang
g(y, μ, σ) = exp(-(y-μ)^2/(2σ^2))/sqrt(2*π*σ^2) # Gaussian
lo(x) = exp(x)/(1+exp(x))

galaxies = dataset("mass", "galaxies")
Yobs = galaxies.x1/1000
const N = 82
@assert N == length(Yobs)
Random.seed!(1)

const K = 3
const d = 3K - 1
const σ = 1.0
const R = 10 # subsampling

function probs(x)
    if K == 2
        (lo(x[end]), 1-lo(x[end]))
    elseif K == 3
        a, b = exp(x[end]), exp(x[end-1])
        c = 1.0 + a + b
        (a/c, b/c, 1 - a/c - b/c)
    else
        a = exp.(x[end-K+2:end])
        sa = sum(a)
        [a/(1 + sa); 1 - sa/(1 + sa)]
    end
end
lprior(x) = sum(log(g(x[k], 1.0, 10.0)) for k in 1:d)
f(y, x, ps=probs(x)) = sum(ps[k]*g(y, x[k], σ*x[K+k]) for k in 1:K)
ϕ(x, Y) = -lprior(x) - N/R*sum(log(f(y, x)) for y in rand(Y, R))
∇ϕ(x, i, Y) = partiali(ϕ, x, i, Y)

using LinearAlgebra
using SparseArrays
using Statistics

t0 = 0.0
x0 = collect(1:K)*10.0
x0 = [x0; ones(d-K)]
θ0 = rand([-1.0,1.0], d)

Γ = sparse(I, d, d)

c = [20.0 for i in 1:d]

μ = zeros(d)

Z = LocalZigZag(Γ, μ)
T = 2000.0

@time trace, (tT, xT, θT), (acc, num), c = pdmp(∇ϕ, t0, x0, θ0, T, c, Z, Yobs; adapt=true)
@show acc, acc/num, mean(c)

tr = collect(trace)
ts = first.(tr)
xs = last.(tr)

ms = [getindex.(xs, k) for k in 1:d]
ps_ = [probs(x) for x in xs]
ps = [getindex.(ps_, k) for k in 1:K]

using Makie
using Bridge
col = Bridge.viridis(1:K)
p1 = scatter(T*ones(K), xT[1:K])
for k in 1:K
    band!(p1, ts, ms[k] + 1.96ms[k+K], ms[k] - 1.96ms[k+K], color=(col[k],0.3) )
    scatter!(p1, ts, ms[k], markersize=1.8, color=col[k] )
    lines!(p1, ts, ms[k],  color=col[k] )

end

p2 = lines(ts, ps[1], color=col[1])
for k in 2:K
    lines!(p2, ts, ps[k], color=col[k])
end

m = median.(ms)
r = range(0, 40, length=200)
p3 = lines(r, [f(y, m) for y in r])
linesegments!(p3, [repeat(Yobs, inner=2) repeat([0,0.05], outer=N)])

p = hbox(title(p3, "est. density and obs."), title(p2, "Trace p[k]"), title(p1, "Trace μ[k] ± σ[k]"))
save(joinpath("figures", "galaxy.png"), title(p, "Galaxy dataset"))