module JackPolynomials

import DynamicPolynomials
export JackPolynomial
export ZonalPolynomial
export ZonalQPolynomial
export SchurPolynomial
export Zonal
export ZonalQ
export Schur
export Jack

function isPartition(lambda::Vector{<:Integer})
  return all(diff(lambda) .<= 0) && all(lambda .>= 0)
end

function dualPartition(lambda::Vector{T}) where {T<:Integer}
  out = T[]
  if !isempty(lambda)
    for i = 1:lambda[1]
      push!(out, sum(lambda .>= i))
    end
  end
  return out
end

function betaratio(
  kappa::Vector{I},
  mu::Vector{I},
  k::I,
  alpha::T,
) where {T<:Real,I<:Integer}
  muk = mu[k]
  t = k - alpha * muk
  u = map(i -> t + 1 - i + alpha * kappa[i], 1:k)
  v = map(i -> t - i + alpha * mu[i], 1:(k-1))
  muPrime = dualPartition(mu)
  w = map(i -> muPrime[i] - t - alpha * i, 1:(muk-1))
  prod1 = prod(u ./ (u .+ alpha .- 1))
  prod2 = prod((v .+ alpha) ./ v)
  prod3 = prod((w .+ alpha) ./ w)
  return alpha * prod1 * prod2 * prod3
end

function _N(lambda::Vector{I}, mu::Vector{I}) where {I<:Integer}
  prods = map(i -> prod(Iterators.drop(lambda .+ 1, i)), 1:length(lambda))
  return sum(mu .* prods)
end

# ------------------------------------------------------------------------------
#~~ Jack polynomial ~~~~##
# ------------------------------------------------------------------------------
"""
    Jack(x, lambda, alpha)

Evaluates a Jack polynomial.

# Arguments
- `x`: vector of real or complex numbers
- `lambda`: partition of an integer
- `alpha`: alpha parameter
"""
function Jack(
  x::Vector{C},
  lambda::Vector{I},
  alpha::T,
) where {T<:Real,I<:Integer,C<:Number}
  if !isPartition(lambda)
    error("`lambda` must be a partition of an integer")
  end
  if alpha <= 0
    error("`alpha` must be positive")
  end
  function jac(m::I, k::I, mu::Vector{I}, nu::Vector{I}, beta::Real)
    if isempty(nu) || nu[1] == 0 || m == 0
      return C(1)
    end
    if length(nu) > m && nu[m+1] > 0
      return C(0)
    end
    if m == 1
      return x[1]^(nu[1]) * prod(1 .+ alpha .* collect(1:(nu[1]-1)))
    end
    v = S[_N(lambda, nu), m]
    if k == 0 && !ismissing(v)
      return v
    end
    i = max(1, k)
    s = jac(m - 1, 0, nu, nu, T(1)) * beta * x[m]^(sum(mu) - sum(nu))
    while length(nu) >= i && nu[i] > 0
      if length(nu) == i || nu[i] > nu[i+1]
        nuPrime = copy(nu)
        nuPrime[i, 1] -= 1
        gamma = beta * betaratio(mu, nu, i, alpha)
        if nu[i] > 1
          s = s + jac(m, i, mu, nuPrime, gamma)
        else
          s =
            s +
            jac(m - 1, 0, nuPrime, nuPrime, T(1)) * gamma * x[m]^(sum(mu) - sum(nuPrime))
        end
      end
      i += 1
    end
    if k == 0
      S[_N(lambda, nu), m] = s
    end
    return s
  end # end jac --------------------------------------------------------------
  lx = length(x)
  S = Array{Union{Missing,C}}(missing, _N(lambda, lambda), lx)
  jac(lx, 0, lambda, lambda, T(1))
end

# ------------------------------------------------------------------------------
#~~ Symbolic Jack polynomial ~~~~##
# ------------------------------------------------------------------------------
function JackPolynomial0(
  m::I, lambda::Vector{I}, alpha::T
) where {T<:Real,I<:Integer}
  function jac(m::I, k::I, mu::Vector{I}, nu::Vector{I}, beta::T)
    if isempty(nu) || nu[1] == 0 || m == 0
      return T(1)
    end
    lnu = length(nu)
    if lnu > m && nu[m+1] > 0
      return T(0)
    end
    if m == 1
      return x[1]^(nu[1]) * prod(1 .+ alpha .* collect(1:(nu[1]-1)))
    end
    v = S[_N(lambda, nu), m]
    if k == 0 && !ismissing(v)
      return v
    end
    i = max(1, k)
    s = jac(m - 1, 0, nu, nu, T(1)) * beta * x[m]^(sum(mu) - sum(nu))
    while lnu >= i && nu[i] > 0
      if lnu == i || nu[i] > nu[i+1]
        nuPrime = copy(nu)
        nuPrime[i, 1] -= 1
        gamma = beta * betaratio(mu, nu, i, alpha)
        if nu[i] > 1
          s += jac(m, i, mu, nuPrime, gamma)
        else
          s += jac(m - 1, 0, nuPrime, nuPrime, T(1)) * gamma * x[m]^(sum(mu) - sum(nuPrime))
        end
      end
      i += 1
    end
    if k == 0
      S[_N(lambda, nu), m] = s
    end
    return s
  end
  DynamicPolynomials.@polyvar x[1:m]
  S = Array{Union{Missing,DynamicPolynomials.Polynomial{true,T}}}(
    missing,
    _N(lambda, lambda),
    m,
  )
  jac(m, 0, lambda, lambda, T(1))
end

"""
    JackPolynomial(m, lambda, alpha)

Symbolic Jack polynomial. The coefficients of the polynomial will have the
same type as `alpha`.

# Arguments
- `m`: integer, the number of variables
- `lambda`: partition of an integer
- `alpha`: alpha parameter
"""
function JackPolynomial(
  m::I, lambda::Vector{I}, alpha::T
) where {T<:Real,I<:Integer}
  if !isPartition(lambda)
    error("`lambda` must be a partition of an integer")
  end
  if alpha <= 0
    error("`alpha` must be positive")
  end
  jack = JackPolynomial0(m, lambda, alpha)
  if(typeof(jack) == T)
    DynamicPolynomials.@polyvar x[1:m]
    jack = sum(T(0) * x) + jack
  end
  return jack
end

# ------------------------------------------------------------------------------
#~~ Zonal polynomial ~~~~##
# ------------------------------------------------------------------------------
function _i(lambda::Vector{T}) where {T<:Integer}
  out = T[]
  for i in 1:length(lambda)
    out = vcat(out, repeat([i], lambda[i]))
  end
  return out
end

function _j(lambda::Vector{T}) where {T<:Integer}
  out = T[]
  mu = filter(x -> x > 0, lambda)
  for m in mu
    out = vcat(out, 1:m)
  end
  return out
end

function hookLengths(lambda::Vector{<:Integer}, alpha::Real)
  i = _i(lambda)
  j = _j(lambda)
  lambdap = dualPartition(lambda)
  upper = lambdap[j] .- i .+ alpha .* (lambda[i] .- j .+ 1)
  lower = lambdap[j] .- i .+ 1 .+ alpha .* (lambda[i] .- j)
  return vcat(upper, lower)
end

"""
    Zonal(x, lambda)

Evaluates a zonal polynomial.

# Arguments
- `x`: vector of real or complex numbers
- `lambda`: partition of an integer
"""
function Zonal(
  x::Vector{<:Union{R,Complex{R}}},
  lambda::Vector{<:Integer},
) where {R<:Real}
  jack = Jack(x, lambda, R(2))
  jlambda = prod(hookLengths(lambda, R(2)))
  n = sum(lambda)
  return jack * 2^n * factorial(n) / jlambda
end

# ------------------------------------------------------------------------------
#~~ Symbolic zonal polynomial ~~~~##
# ------------------------------------------------------------------------------
"""
    ZonalPolynomial(m, lambda[, type])

Symbolic zonal polynomial.

# Arguments
- `m`: integer, the number of variables
- `lambda`: partition of an integer
- `type`: the type of the coefficients of the polynomial; default `Rational`
"""
function ZonalPolynomial(
  m::I,
  lambda::Vector{I},
  type::Type = Rational
) where {I<:Integer}
  jack = JackPolynomial(m, lambda, type(2))
  jlambda = prod(hookLengths(lambda, type(2)))
  n = sum(lambda)
  return jack * 2^n * factorial(n) / jlambda
end

# ------------------------------------------------------------------------------
#~~ Quaternionic zonal polynomial ~~~~##
# ------------------------------------------------------------------------------
"""
    ZonalQ(x, lambda)

Evaluates a quaternionic zonal polynomial.

# Arguments
- `x`: vector of real or complex numbers
- `lambda`: partition of an integer
"""
function ZonalQ(
  x::Vector{<:Union{R,Complex{R}}},
  lambda::Vector{<:Integer},
) where {R<:Real}
  jack = Jack(x, lambda, R(0.5))
  jlambda = prod(hookLengths(lambda, R(0.5)))
  n = sum(lambda)
  return jack * factorial(n) / 2^n / jlambda
end

# ------------------------------------------------------------------------------
#~~ Symbolic quaternionic zonal polynomial ~~~~##
# ------------------------------------------------------------------------------
"""
    ZonalQPolynomial(m, lambda[, type])

Symbolic quaternionic zonal polynomial.

# Arguments
- `m`: integer, the number of variables
- `lambda`: partition of an integer
- `type`: the type of the coefficients of the polynomial; default `Rational`
"""
function ZonalQPolynomial(
  m::I,
  lambda::Vector{I},
  type::Type = Rational
) where {I<:Integer}
  jack = JackPolynomial(m, lambda, type(0.5))
  jlambda = prod(hookLengths(lambda, type(0.5)))
  n = sum(lambda)
  return jack * factorial(n) / 2^n / jlambda
end

# ------------------------------------------------------------------------------
#~~ Schur polynomial ~~~~##
# ------------------------------------------------------------------------------
"""
    Schur(x, lambda)

Evaluates a Schur polynomial.

# Arguments
- `x`: vector of real or complex numbers
- `lambda`: partition of an integer
"""
function Schur(x::Vector{T}, lambda::Vector{I}) where {T<:Number,I<:Integer}
  if !isPartition(lambda)
    error("`lambda` must be a partition of an integer")
  end
  function sch(m::Int64, k::Int64, nu::Vector{I})
    if isempty(nu) || nu[1] == 0 || m == 0
      return T(1)
    end
    if length(nu) > m && nu[m+1] > 0
      return T(0)
    end
    if m == 1
      return x[1]^nu[1]
    end
    v = S[_N(lambda, nu), m]
    if !ismissing(v)
      return v
    end
    s = sch(m - 1, 1, nu)
    lnu = length(nu)
    i = k
    while lnu >= i && nu[i] > 0
      if lnu == i || nu[i] > nu[i+1]
        nup = copy(nu)
        nup[i] -= 1
        if nu[i] > 1
          s += x[m] * sch(m, i, nup)
        else
          s += x[m] * sch(m - 1, 1, nup)
        end
      end
      i += 1
    end
    if k == 1
      S[_N(lambda, lambda), m] = s
    end
    return s
  end # end sch --------------------------------------------------------------
  lx = length(x)
  S = Array{Union{Missing,T}}(missing, _N(lambda, lambda), lx)
  sch(lx, 1, lambda)
end

# ------------------------------------------------------------------------------
#~~ Symbolic Schur polynomial ~~~~##
# ------------------------------------------------------------------------------
function SchurPolynomial0(
  m::I,
  lambda::Vector{I},
  T::Type
) where {I<:Integer}
  function sch(m::Int64, k::Int64, nu::Vector{I})
    if isempty(nu) || nu[1] == 0 || m == 0
      return T(1)
    end
    if length(nu) > m && nu[m+1] > 0
      return T(0)
    end
    if m == 1
      return x[1]^nu[1]
    end
    v = S[_N(lambda, nu), m]
    if !ismissing(v)
      return v
    end
    s = sch(m - 1, 1, nu)
    lnu = length(nu)
    i = k
    while lnu >= i && nu[i] > 0
      if lnu == i || nu[i] > nu[i+1]
        nup = copy(nu)
        nup[i] -= 1
        if nu[i] > 1
          s += x[m] * sch(m, i, nup)
        else
          s += x[m] * sch(m - 1, 1, nup)
        end
      end
      i += 1
    end
    if k == 1
      S[_N(lambda, lambda), m] = s
    end
    return s
  end # end sch --------------------------------------------------------------
  DynamicPolynomials.@polyvar x[1:m]
  S = Array{Union{Missing,DynamicPolynomials.Polynomial{true,T}}}(
    missing,
    _N(lambda, lambda),
    m
  )
  sch(m, 1, lambda)
end

"""
    SchurPolynomial(m, lambda[, type])

Symbolic Schur polynomial.

# Arguments
- `m`: integer, the number of variables
- `lambda`: partition of an integer
- `type`: the type of the coefficients of the polynomial; default `Int64`
"""
function SchurPolynomial(
  m::I,
  lambda::Vector{I},
  type::Type = Int64
) where {I<:Integer}
  if !isPartition(lambda)
    error("`lambda` must be a partition of an integer")
  end
  schur = SchurPolynomial0(m, lambda, type)
  if(typeof(schur) == type)
    DynamicPolynomials.@polyvar x[1:m]
    schur = sum(type(0) * x) + schur
  end
  return schur
end

end # module


## Compound objects

export Line
function Line(p₁::RealVector, p₂::RealVector, r::Real; infty=10^5)
    v = normalize(p₂-p₁)
    return Cylinder(p₁-infty*v,p₂+infty*v,r)
end

export HalfLine
function HalfLine(p₁::RealVector, p₂::RealVector, r::Real; infty=10^5)
    v = normalize(p₂-p₁)
    return Cylinder(p₁,p₂+infty*v,r)
end

export Arrow
function Arrow(end1::RealVector, end2::RealVector, r::Real)
    n = normalize(end2-end1)
    end3 = end2-6*r*n
    return csgUnion(Cylinder(end1,end3,r),Cone(end3,end2,2*r))
end

function Blocks³(p₁::RealVector, p₂::RealVector, p₃::RealVector, thickness::Float64)
    center = Circumcenter(p₁,p₂,p₃)
    e₃ = NormalVector(p₁,p₂,p₃)
    R = norm(p₁-center)
    e₁ = normalize(p₁-center)
    e₂ = cross(e₃,e₁)
    box = Box([R+thickness,R+thickness,thickness],-[R+thickness,0,thickness])
    box1 = AffineTransform(box,hcat(e₁,e₂,e₃),center)
    e₁ = -normalize(p₃-center)
    e₂ = cross(e₃,e₁)
    box2 = AffineTransform(box,hcat(e₁,e₂,e₃),center)
    if det(hcat(p₁-center,p₃-center,e₃)) > 0
        return csgIntersection(box1,box2)
    else
        return csgMerge(box1,box2)
    end
end

export Arc
function Arc(p₁::RealVector, p₂::RealVector, p₃::RealVector, r::Float64; ε=10^(-4))
    α₁,α₂,α₃ = Angles(p₁,p₂,p₃)
    if (α₁>π-ε)
        return csgUnion(HalfLine(p₁,p₁+(p₂-p₃),r),HalfLine(p₃,p₃-(p₂-p₃),r))
    elseif (α₂>π-ε)
        return Cylinder(p₁,p₃,r)
    elseif (α₃>π-ε)
        return csgUnion(HalfLine(p₁,p₁+(p₁-p₂),r),HalfLine(p₃,p₃-(p₁-p₂),r))
    else
        return csgBound(csgIntersection(Torus(p₁,p₂,p₃,r),Blocks³(p₁,p₂,p₃,2r)),Torus(p₁,p₂,p₃,r))
    end
end

export RoundedBox
function RoundedBox(vertex1::RealVector, vertex2::RealVector, radius::Real)
    if radius == 0
        return Box(vertex1, vertex2)
    elseif radius < 0
        error("radius must be non-negative")
    end
    x=vertex1[1], vertex2[1]
    y=vertex1[2], vertex2[2]
    z=vertex1[3], vertex2[3]
    x_min, x_max=minimum(x), maximum(x)
    y_min, y_max=minimum(y), maximum(y)
    z_min, z_max=minimum(z), maximum(z)

    v_min=[x_min, y_min, z_min] .+ radius
    v_max=[x_max, y_max, z_max] .- radius

    x_r=v_min[1], v_max[1]
    y_r=v_min[2], v_max[2]
    z_r=v_min[3], v_max[3]

    spheres=[Sphere([x_r[i], y_r[j], z_r[k]], radius) for i in 1:2, j in 1:2, k in 1:2]

    cylinders_x=[Cylinder([x_r[1], y_r[j], z_r[k]], [x_r[2], y_r[j], z_r[k]], radius) for j in 1:2, k in 1:2]
    cylinders_y=[Cylinder([x_r[i], y_r[1], z_r[k]], [x_r[i], y_r[2], z_r[k]], radius) for i in 1:2, k in 1:2]
    cylinders_z=[Cylinder([x_r[i], y_r[j], z_r[1]], [x_r[i], y_r[j], z_r[2]], radius) for i in 1:2, j in 1:2]

    box_x=Box([x_min, y_r[1], z_r[1]], [x_max, y_r[2], z_r[2]])
    box_y=Box([x_r[1], y_min, z_r[1]], [x_r[2], y_max, z_r[2]])
    box_z=Box([x_r[1], y_r[1], z_min], [x_r[2], y_r[2], z_max])

    return csgMerge(spheres..., cylinders_x..., cylinders_y..., cylinders_z..., box_x, box_y, box_z)
end


import Distributions
const Dists = Distributions
using Distributions: mean
import MeasureTheory
const MT = MeasureTheory
using HypercubeTransform
using Test

@testset "Scalar" begin
    a = Dists.Normal()
    b = Dists.Uniform()
    c = Dists.Gamma()

    ca = ascube(a)
    cb = ascube(b)
    cc = ascube(c)

    @test transform(ca, 0.5) == 0.0
    @test transform(cb, 0.5) == 0.5
    @test transform(cc, 0.5) == Dists.quantile(c, 0.5)

    @inferred Float64 transform(ca, 0.5)
end

@testset "ArrayHC" begin
    d = Dists.Product([Dists.Normal(), Dists.Uniform(), Dists.Normal()])
    m = MT.For(d.v) do p
            return p
    end
    cd = ascube(d)
    cm = ascube(m)
    @test transform(cd, [0.5,0.5,0.5]) == transform(cm, [0.5,0.5,0.5])
    cc = ascube.(d.v)
    @test transform.(cc, [0.5,0.5,0.5]) == transform(cd, [0.5,0.5,0.5])
    @inferred Vector{Float64} transform(cm, [0.5, 0.5, 0.5])
end

@testset "Dirichlet" begin
    d = Dists.Dirichlet([1.0,1.0,1.0])
    a = Dists.Normal()
    dc = ascube(d)
    @inferred transform(dc, [0.5, 0.5, 0.5])
    p = rand(3, 10_000_00)
    x = transform.(Ref(dc), eachcol(p))
    @test isapprox(mean(x), mean(d), atol=1e-3)

    dad = ascube((a, d))
    @inferred transform(dad, [0.5,0.5,0.5,0.5])
    pad = rand(4, 10_000_00)
    xad = transform.(Ref(dad), eachcol(pad))
    mN = mean(first.(xad))
    mD = mean(last.(xad))
    @test isapprox(mN, mean(a), atol=1e-3)
    @test isapprox(mD, mean(d), atol=1e-3)

end


using Primes
using Base.Iterators: drop, product

proper_divisors(n) = begin
    if n == 1
        return [1]
    end
    factors_of_n = factor(n)
    nfactors = prod(x -> x[2] + 1, factors_of_n)
    res = Vector{Int}(undef, nfactors)
    res .= 1
    i = 1
    for (factor, power) in factors_of_n
        if i == 1
            res[1] = 1
            res[2:power+1] .= factor .^ (1:power)
            i = power + 1
        else
            next_few = i * power
            next_few_vals = [
                existing * more
                for (existing, more) in product(res[1:i], factor .^ (1:power))
            ]
            res[i+1:i+next_few] = vec(next_few_vals)
            i += next_few
        end
    end
    res[1:end-1]
end


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

function cell_interface_error(quad, cellmap, facescale, distance_to_interface)
    err = 0.0
    for (idx, (p, w)) in enumerate(quad)
        x = cellmap(p)
        err += (distance_to_interface(x)[1])^2 * facescale[idx] * w
    end
    return err
end

function interface_error(levelset, levelsetcoeffs, cutmesh, numqp, distance_to_interface)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    cellids = findall(CutCell.cell_sign(cutmesh) .== 0)

    err = 0.0
    for cellid in cellids
        quad = interfacequads[+1, cellid]
        if length(quad) > 0
            normals = CutCell.interface_normals(interfacequads, cellid)
            cellmap = CutCell.cell_map(cutmesh, cellid)
            facescale = CutCell.scale_area(cellmap, normals)

            err += cell_interface_error(quad, cellmap, facescale, distance_to_interface)
        end
    end
    return sqrt(err)
end

function shrinking_circle_distance(x, xc, initialradius, speed, time)
    r = initialradius - speed * time
    return circle_distance_function(x, xc, r)
end

function interface_error_over_timesteps(
    levelset,
    mesh,
    speed,
    dt,
    nsteps,
    numqp,
    xc,
    initialradius,
)
    err = zeros(nsteps + 1)

    levelsetcoeffs = CutCell.levelset_coefficients(
        x -> shrinking_circle_distance(x, xc, initialradius, speed, 0.0),
        mesh,
    )
    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)

    err[1] = interface_error(
        levelset,
        levelsetcoeffs,
        cutmesh,
        numqp,
        x -> shrinking_circle_distance(x, xc, initialradius, speed, 0.0),
    )/(2pi*initialradius)

    levelsetspeed = speed * ones(length(levelsetcoeffs))

    for t = 1:nsteps
        paddedmesh = CutCell.BoundaryPaddedMesh(cutmesh, 1)
        refseedpoints, spatialseedpoints, seedcellids =
            CutCell.seed_zero_levelset(2, levelset, levelsetcoeffs, cutmesh)
        paddedlevelset = CutCell.BoundaryPaddedLevelset(
            paddedmesh,
            refseedpoints,
            spatialseedpoints,
            seedcellids,
            levelset,
            levelsetcoeffs,
            cutmesh,
            1e-10,
        )
        levelsetcoeffs =
            CutCell.step_first_order_levelset(paddedlevelset, levelsetspeed, dt)

        cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
        err[t+1] = interface_error(
            levelset,
            levelsetcoeffs,
            cutmesh,
            numqp,
            x -> shrinking_circle_distance(x,xc,initialradius,speed,t*dt),
        )/(2pi*(initialradius - speed*t*dt))
    end
    return err
end

function final_interface_error(
    levelset,
    mesh,
    speed,
    dt,
    nsteps,
    numqp,
    xc,
    initialradius,
)

    levelsetcoeffs = CutCell.levelset_coefficients(
        x -> shrinking_circle_distance(x, xc, initialradius, speed, 0.0),
        mesh,
    )
    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    levelsetspeed = speed * ones(length(levelsetcoeffs))

    for t = 1:nsteps
        paddedmesh = CutCell.BoundaryPaddedMesh(cutmesh, 1)
        refseedpoints, spatialseedpoints, seedcellids =
            CutCell.seed_zero_levelset(2, levelset, levelsetcoeffs, cutmesh)
        paddedlevelset = CutCell.BoundaryPaddedLevelset(
            paddedmesh,
            refseedpoints,
            spatialseedpoints,
            seedcellids,
            levelset,
            levelsetcoeffs,
            cutmesh,
            1e-10,
        )
        levelsetcoeffs =
            CutCell.step_first_order_levelset(paddedlevelset, levelsetspeed, dt)

        cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    end
    err = interface_error(
        levelset,
        levelsetcoeffs,
        cutmesh,
        numqp,
        x -> shrinking_circle_distance(x,xc,initialradius,speed,nsteps*dt),
    )/(2pi*(initialradius - speed*nsteps*dt))
    return err
end

function grid_size(mesh)
    w = CutCell.widths(mesh)
    nn = CutCell.nodes_per_mesh_side(mesh)
    return w ./ (nn .- 1)
end

function time_step_size(speed, mesh; CFL = 0.5)
    dx = minimum(grid_size(mesh))
    return CFL * dx / abs(speed)
end

function numsteps(stoptime,dt)
    r = stoptime/dt
    @assert isinteger(r)
    return round(Int,r)
end

function error_for_numelmts(nelmts)
    x0 = [0.0, 0.0]
    L, W = 1.0, 1.0
    numghostlayers = 1
    polyorder = 2
    numqp = 3

    xc = [0.5, 0.5]
    initialradius = 0.50
    speed = 1.0
    stoptime = 0.4

    basis = TensorProductBasis(2, polyorder)
    levelset = InterpolatingPolynomial(1, basis)

    mesh = CutCell.Mesh(x0, [L, W], [nelmts, nelmts], basis)
    dt = time_step_size(speed, mesh)
    nsteps = numsteps(stoptime,dt)

    return final_interface_error(
        levelset,
        mesh,
        speed,
        dt,
        nsteps,
        numqp,
        xc,
        initialradius,
    )
end

nelmts = [5,10,20,40,80]

# err80 = error_for_numelmts(80)
err = [error_for_numelmts(ne) for ne in nelmts]

dx = 1.0 ./ nelmts
rate = diff(log.(err)) ./ diff(log.(dx))

using CSV,DataFrames
df = DataFrame([dx,err],["Element Size", "Error"])
CSV.write("examples/levelset-propagate/circle-convergence.csv",df)

#
# using Plots
# fig = plot(legend=:topleft)
# plot!(range(0,stop=0.4,length=length(err[1])),err[1],label="5x5")
# plot!(range(0,stop=0.4,length=length(err[2])),err[2],label="10x10")
# plot!(range(0,stop=0.4,length=length(err[3])),err[3],label="20x20")
# plot!(range(0,stop=0.4,length=length(err[4])),err[4],label="40x40")


include("pascaltriangle.jl") # define pascaltriangle(m,n)
# include("pascalmax.jl") # degine pascalmax(m,n)

function b2in(vecmb::SparseVector{Int64})

    # this function return the index of many body state at the index of Fock state

    # ignore the last element since it is the coefficient
    Msize = length(vecmb)-1
    Np = sum(vecmb[1:Msize])

    # Lm = mbsize+1
    # Ln = Np + 1
    indfk = 1
    matp = zeros(Int,Msize+1,Np+1);
    pascaltriangle!(Msize,Np,matp) # the size is Msize+1 times Np+1
    # matp = pascaltriangle(Msize,Np) # the size is Msize+1 times Np+1
    # note the indices are m+1 and n+1 for N^m_n

    for jj = 1:Msize
        for ii = 0:vecmb[jj]

            if ii == vecmb[jj]
               continue
            end

            sumparticles = sum(vecmb[1:jj-1]) # numnber of particles taken account so far
            indfk = indfk + matp[Msize-jj+1,Np-ii-sumparticles+1] # (n+m-1)!/n!(m-1)!
            # note the indices are m+1 and n+1 for N^m_n
            # indfk = indfk + pascalmax(Msize-jj,Np-ii-sumparticles) # pascalmax(m,n) = (n+m-1)!/n!(m-1)!

        end
    end

    # unlike Ponomarev's way, I count the index of Fock state in the backwards
    # so that "1" means |Np,0,0,...>.
    indfk = matp[Msize+1,Np+1] + 1 - indfk # the indices are m+1 and n+1 for N^m_n

    return indfk

end


function gen_carrier!(
    carrier::StructArray{<:Complex{T}},
    frequency,
    sampling_frequency,
    start_phase,
    amplitude::T
) where T
    c_re = carrier.re; c_im = carrier.im
    phase_step = upreferred(frequency / sampling_frequency)
    @avx for i in 1:length(carrier)
        c_im_temp, c_re_temp =
            sincos(T(2π) * ((i - 1) * T(phase_step) + T(start_phase)))
        c_im[i] = c_im_temp * amplitude
        c_re[i] = c_re_temp * amplitude
    end
    carrier
end


## cauchy

function cauchycircleS(cfs::AbstractVector,z::Number,s::Bool)
    ret=zero(Complex{Float64})

    if s
        zm = one(Complex{Float64})

        #odd coefficients are pos
        @simd for k=1:2:length(cfs)
            @inbounds ret += cfs[k]*zm
            zm *= z
        end
    else
        z=1./z
        zm = z

        #even coefficients are neg
        @simd for k=2:2:length(cfs)
            @inbounds ret -= cfs[k]*zm
            zm *= z
        end
    end

    ret
end

function stieltjes{DD<:Circle}(sp::Laurent{DD},f::AbstractVector,z,s::Bool)
    d=domain(sp)
    if !d.orientation
        return -stieltjes(reverseorientation(Fun(f,sp)),z,!s)
    end
    @assert in(z,d)
    -2π*im*cauchycircleS(f,mappoint(d,Circle(),z),s)
end
function stieltjes{DD<:Circle}(sp::Laurent{DD},f::AbstractVector,z::Number)
    d=domain(sp)
    if !d.orientation
        return -stieltjes(reverseorientation(Fun(f,sp)),z)
    end

    z=mappoint(d,Circle(),z)
    -2π*im*cauchycircleS(f,z,abs(z) < 1)
end

stieltjes{DD<:Circle}(sp::Laurent{DD},f,z::Vector)=[stieltjes(sp,f,zk) for zk in z]
stieltjes{DD<:Circle}(sp::Laurent{DD},f,z::Matrix)=reshape(stieltjes(sp,f,vec(z)),size(z,1),size(z,2))

stieltjes{DD<:Circle}(sp::Fourier{DD},f,z,s...)=stieltjes(Laurent(domain(sp)),coefficients(f,sp,Laurent(domain(sp))),z,s...)

# we implement cauchy ±1 as canonical
hilbert{DD<:Circle}(sp::Laurent{DD},f,z)=(stieltjes(sp,f,z,true)+stieltjes(sp,f,z,false))/(-2π)

## stieltjesintegral and logkernel

function stieltjesintegral{DD<:Circle}(sp::Laurent{DD},f,z::Number,s...)
    d=domain(sp)
    @assert d==Circle()  #TODO: radius
    ζ=Fun(d)
    r=stieltjes(integrate(f-f[2]/ζ),z,s...)
    abs(z)<1?r:r+2π*im*f[2]*log(z)
end

stieltjesintegral{DD<:Circle}(sp::Fourier{DD},f,z::Number,s...)=stieltjesintegral(Fun(Fun(f,sp),Laurent),z,s...)

function logkernel{DD<:Circle}(sp::Fourier{DD},g,z::Number)
    d=domain(sp)
    c,r=d.center,d.radius
    z=z-c
    if abs(z) ≤r
        ret=2r*log(r)*g[1]
        for j=2:2:length(g)
            k=div(j,2)
            ret+=-g[j]*sin(k*angle(z))*abs(z)^k/(k*r^(k-1))
        end
        for j=3:2:length(g)
            k=div(j,2)
            ret+=-g[j]*cos(k*angle(z))*abs(z)^k/(k*r^(k-1))
        end
        ret
    else
        ret=2r*logabs(z)*g[1]
        for j=2:2:length(g)
            k=div(j,2)
            ret+=-g[j]*sin(k*angle(z))*r^(k+1)/(k*abs(z)^k)
        end
        for j=3:2:length(g)
            k=div(j,2)
            ret+=-g[j]*cos(k*angle(z))*r^(k+1)/(k*abs(z)^k)
        end
        ret
    end
end
logkernel{DD<:Circle}(sp::Fourier{DD},g,z::Vector) = promote_type(eltype(g),eltype(z))[logkernel(sp,g,zk) for zk in z]
logkernel{DD<:Circle}(sp::Fourier{DD},g,z::Matrix) = reshape(promote_type(eltype(g),eltype(z))[logkernel(sp,g,zk) for zk in z],size(z))

logkernel{DD<:Circle}(sp::Laurent{DD},g,z)=logkernel(Fun(Fun(g,sp),Fourier),z)


using Test

using BEAST
using CompScienceMeshes

o, x, y = point(0,0,0), point(1,0,0), point(0,1,0)
G = Mesh([o,x,y],[index(1,2,3)])

@test numvertices(G) == 3
@test numcells(G) == 1

K = MWDoubleLayerTDIO(1.0, 1.0, 0)

X = raviartthomas(G, skeleton(G,1))
Y = buffachristiansen(G, boundary(G))
@test numfunctions(X) == 3
@test numfunctions(Y) == 3

Δt = 20.0
Nt = 5
δ = timebasisdelta(Δt, Nt)
T = timebasiscxd0(Δt, Nt)
T2 = timebasisshiftedlagrange(Δt, Nt, 2)
V = X ⊗ T
W = Y ⊗ T

fr1, store1 = BEAST.allocatestorage(K, W, V,
    Val{:densestorage}, BEAST.LongDelays{:ignore})

BEAST.assemble!(K, W, V, store1)
Z = fr1()
@test all(==(0), Z[:,:,1])

W = X⊗δ
V = Y⊗T2

K = TDMaxwell3D.doublelayer(speedoflight=1.0, numdiffs=1)
fr2, store2 = BEAST.allocatestorage(K, W, V, Val{:densestorage}, BEAST.LongDelays{:ignore})
BEAST.assemble!(K, W, V, store2)
Z2 = fr2()
@test all(==(0), Z2[:,:,1])

γ = geometry(Y)
ch1 = chart(G, first(cells(G)))
qps = quadpoints(ch1, 3)

x = qps[1][1]
w = qps[1][2]

test_els = BEAST.elements(G)
trial_els = BEAST.elements(γ)
time_els = [(0.0, 20.0)]

qs = BEAST.defaultquadstrat(K,X,Y)
qdt = BEAST.quaddata(K, refspace(X), refspace(Y), refspace(T2), test_els, trial_els, nothing, qs)
qrl = BEAST.quadrule(K, refspace(X), refspace(Y), refspace(T2),
    1, test_els[1], 1, trial_els[1], 1, time_els[1], qdt, qs)
z = zeros(Float64, 3, 3, 10)
BEAST.innerintegrals!(z, K, x, refspace(X), refspace(Y), refspace(T2), test_els[1], trial_els[1], (0.0, 20.0), qrl, w)
@test_broken !any(isnan.(z))

verts = trial_els[1].vertices
import WiltonInts84
iG, ivG, igG = WiltonInts84.wiltonints(verts[1], verts[2], verts[3], cartesian(x), 0.0, 20.0, Val{2}, qrl.workspace)
@test_broken !any(isnan(iG))

ctr = WiltonInts84.contour!(verts[1],verts[2],verts[3],cartesian(x),0.0,20.0,qrl.workspace)
iG, ivG, igG = WiltonInts84.wiltonints(ctr,cartesian(x),Val{2})
@test_broken !any(isnan(iG))

import WiltonInts84
a = -0.23570226039551564
b = 1.5700924586837734e-16
p = -6.371760312437358e-16
h = 0.0
using StaticArrays
m = SVector(-0.7071067811865474, 0.7071067811865478, -0.0)
iG, vG = WiltonInts84.segintsg(a, b, p, h, m, Val{2})
@test_broken !any(isnan.(iG))


mutable struct NilakanthaSeries
  p::Float64
  a::Float64
  s::Float64
  iterate::Function
  function NilakanthaSeries(_p, _a, _s)
    this = new(_p, _a, _s)
    this.iterate = function()
      s = this.s
      a = this.a
      this.p += s*(4.0/(a*(a*(a+3.0)+2.0)))
      this.a += 2.0
      this.s = -s
    end
    this
  end
end

approximator =  NilakanthaSeries(3.0, 2.0, 1.0)

for _ = 0:150000
  approximator.iterate()
end

println("pi = ",approximator.p)


"""
    Chi(ν)

The *Chi distribution* `ν` degrees of freedom has probability density function

```math
f(x; \\nu) = \\frac{1}{\\Gamma(\\nu/2)} 2^{1 - \\nu/2} x^{\\nu-1} e^{-x^2/2}, \\quad x > 0
```

It is the distribution of the square-root of a [`Chisq`](@ref) variate.

```julia
Chi(ν)       # Chi distribution with ν degrees of freedom

params(d)    # Get the parameters, i.e. (ν,)
dof(d)       # Get the degrees of freedom, i.e. ν
```

External links

* [Chi distribution on Wikipedia](http://en.wikipedia.org/wiki/Chi_distribution)

"""
struct Chi{T<:Real} <: ContinuousUnivariateDistribution
    ν::T
    Chi{T}(ν::T) where {T} = new{T}(ν)
end

function Chi(ν::Real; check_args::Bool=true)
    @check_args Chi (ν, ν > zero(ν))
    return Chi{typeof(ν)}(ν)
end

Chi(ν::Integer; check_args::Bool=true) = Chi(float(ν); check_args=check_args)

@distr_support Chi 0.0 Inf

### Conversions
convert(::Type{Chi{T}}, ν::Real) where {T<:Real} = Chi(T(ν))
convert(::Type{Chi{T}}, d::Chi{S}) where {T <: Real, S <: Real} = Chi(T(d.ν), check_args=false)

#### Parameters

dof(d::Chi) = d.ν
params(d::Chi) = (d.ν,)
@inline partype(d::Chi{T}) where {T<:Real} = T

#### Statistics

mean(d::Chi) = (h = d.ν/2; sqrt2 * exp(loggamma(h + 1//2) - loggamma(h)))

var(d::Chi) = d.ν - mean(d)^2
_chi_skewness(μ::Real, σ::Real) = (σ2 = σ^2; σ3 = σ2 * σ; (μ / σ3) * (1 - 2σ2))

function skewness(d::Chi)
    μ = mean(d)
    σ = sqrt(d.ν - μ^2)
    _chi_skewness(μ, σ)
end

function kurtosis(d::Chi)
    μ = mean(d)
    σ = sqrt(d.ν - μ^2)
    γ = _chi_skewness(μ, σ)
    (2/σ^2) * (1 - μ * σ * γ - σ^2)
end

entropy(d::Chi{T}) where {T<:Real} = (ν = d.ν;
    loggamma(ν/2) - T(logtwo)/2 - ((ν - 1)/2) * digamma(ν/2) + ν/2)

function mode(d::Chi; check_args::Bool=true)
    ν = d.ν
    @check_args(
        Chi,
        (ν, ν >= 1, "Chi distribution has no mode when ν < 1"),
    )
    sqrt(ν - 1)
end

#### Evaluation

function logpdf(d::Chi, x::Real)
    ν, _x = promote(d.ν, x)
    xsq = _x^2
    val = (xlogy(ν - 1, xsq / 2) - xsq + logtwo) / 2 - loggamma(ν / 2)
    return x < zero(x) ? oftype(val, -Inf) : val
end

gradlogpdf(d::Chi{T}, x::Real) where {T<:Real} = x >= 0 ? (d.ν - 1) / x - x : zero(T)

for f in (:cdf, :ccdf, :logcdf, :logccdf)
    @eval $f(d::Chi, x::Real) = $f(Chisq(d.ν; check_args=false), max(x, 0)^2)
end

for f in (:quantile, :cquantile, :invlogcdf, :invlogccdf)
    @eval $f(d::Chi, p::Real) = sqrt($f(Chisq(d.ν; check_args=false), p))
end

#### Sampling

rand(rng::AbstractRNG, d::Chi) =
    (ν = d.ν; sqrt(rand(rng, Gamma(ν / 2.0, 2.0one(ν)))))

struct ChiSampler{S <: Sampleable{Univariate,Continuous}} <:
    Sampleable{Univariate,Continuous}
    s::S
end

rand(rng::AbstractRNG, s::ChiSampler) = sqrt(rand(rng, s.s))

sampler(d::Chi) = ChiSampler(sampler(Chisq(d.ν)))


"""Specification for a "time-local" pulse parametrization.

The parametrization is given as a collection of three functions:

* ``ϵ(u(t))``
* ``u(ϵ(t))``
* ``∂ϵ/∂u`` as a function of ``u(t)``.
"""
struct PulseParametrization
    epsilon_of_u :: Function
    u_of_epsilon :: Function
    de_du_derivative :: Function
end

"""Parametrization where ``ϵ(t) ≡ u(t)``."""
NoParametrization() = PulseParametrization(u->u, ϵ->ϵ, u->1.0)

"""Parametrization ϵ(t) = u²(t), enforcing pulse values ≥ 0."""
SquareParametrization() = PulseParametrization(
    u -> begin # epsilon_of_u
        ϵ = u^2
    end,
    ϵ -> begin # u_of_epsilon
        a = max(ϵ, 0.0)
        u = √a
    end,
    u -> begin # de_du_derivative
        ∂ϵ╱∂u = 2u
    end
)

"""Parametrization with a tanh function that enforces ϵ_min < ϵ(t) < ϵ_max.
"""
function TanhParametrization(ϵ_min, ϵ_max)

    Δ = ϵ_max - ϵ_min
    Σ = ϵ_max + ϵ_min
    εₚ = eps(1.0)  # 2⋅10⁻¹⁶ (machine precision)
    @assert ϵ_max > ϵ_min

    PulseParametrization(
        u -> begin # epsilon_of_u
            ϵ = tanh(u) * Δ / 2 + Σ / 2
        end,
        ϵ -> begin # u_of_epsilon
            a = clamp(2ϵ/Δ -  Σ/Δ, -1+εₚ, 1-εₚ)
            u = atanh(a)  # -18.4 < u < 18.4
        end,
        u -> begin # de_du_derivative
            ∂ϵ╱∂u = (Δ / 2) * sech(u)^2
        end
    )

end

"""Parametrization with a tanh² function that enforces ``0 ≤ ϵ(t) < ϵ_max``.
"""
function TanhSqParametrization(ϵ_max)

    εₚ = eps(1.0)  # 2⋅10⁻¹⁶ (machine precision)
    @assert ϵ_max > 0

    PulseParametrization(
        u -> begin # epsilon_of_u
            ϵ = ϵ_max * tanh(u)^2
        end,
        ϵ -> begin # u_of_epsilon
            a = clamp(ϵ / ϵ_max, 0, 1-εₚ)
            u = atanh(√a)
        end,
        u -> begin # de_du_derivative
            ∂ϵ╱∂u = 2ϵ_max * tanh(u) * sech(u)^2
        end
    )

end

"""Parametrization with a Logistic function that enforces ϵ_min < ϵ(t) < ϵ_max.
"""
function LogisticParametrization(ϵ_min, ϵ_max; k=1.0)

    Δ = ϵ_max - ϵ_min
    ε₀ = eps(0.0)  # 5⋅10⁻³²⁴
    @assert ϵ_max > ϵ_min

    PulseParametrization(
        u -> begin # epsilon_of_u
            ϵ = Δ / (1 + exp(-k * u)) + ϵ_min
        end,
        ϵ -> begin # u_of_epsilon
            ϵ′ = ϵ - ϵ_min
            a = max(ϵ′/(Δ - ϵ′), ε₀)
            u = log(a) / k
        end,
        u -> begin # de_du_derivative
            e⁻ᵏᵘ = exp(-k * u)
            ∂ϵ╱∂u = Δ * k * e⁻ᵏᵘ / (1 + e⁻ᵏᵘ)^2
        end
    )

end

"""Parametrization with a Logistic-Square function that enforces 0 ≤ ϵ(t) < ϵ_max.
"""
function LogisticSqParametrization(ϵ_max; k=1.0)

    ε₀ = eps(0.0)  # 5⋅10⁻³²⁴
    @assert ϵ_max > 0

    PulseParametrization(
        u -> begin # epsilon_of_u
            ϵ = ϵ_max * (2 / (1 + exp(-k * u)) - 1)^2
        end,
        ϵ -> begin # u_of_epsilon
            ρ = clamp(ϵ / ϵ_max, 0.0, 1.0)
            a = clamp((2 / (√ρ + 1)) - 1, ε₀, 1.0)
            u = -log(a) / k
        end,
        u -> begin # de_du_derivative
            eᵏᵘ = exp(k * u)
            ∂ϵ╱∂u = 4k * ϵ_max * eᵏᵘ * (eᵏᵘ  - 1) / (eᵏᵘ + 1)^3
        end
    )

end


using Cosmology

default_cosmology() = cosmology(h=0.70, OmegaM=0.3)   #S11


using JFVM
using PyPlot
include("CRT.jl")

# define the domain and constants
Lx = 1.0 # [m]
Nx = 100  # number of grids
ϕ  = 0.4 # porosity
u  = 1.0/(3600*24) # m/s, eq. 1 m/day
D  = [1e-9, 1e-9, 1e-9] # m^2/s diffusivity
K_eq = 10.0
k_f = 100.0
k_b = k_f/K_eq
nu_A = 1.0 # stoichiometric coefficient of component A
nu_B = 1.0 # stoichiometric coefficient of component B
nu_C = 1.0 # stoichiometric coefficient of component C
a = 1.0 #order of reaction
b = 1.0 #order of reaction
c = 1.0 #order of reaction
N_comp = 3 # number of components

# create the mesh
m = createMesh1D(Nx, Lx)  # mesh
# m = createMeshRadial2D(Nx, Nx, Lx, Lx)

# boundary conditions
C_bc = [1.0, 0.5, 0.0] # boundary concentration
C_0 = zeros(N_comp) # initial concentration

dt = Lx/(u/ϕ)/Nx
T = Nx/2*dt
tol = 1e-10 # convergence tolerance

#Calculating reference "true solution"
C_old = CRT.CRT_GIA_eps(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T, tol)
#C_old = CRT.CRT_SNIA(m, Nx, ϕ, u, D, K_eq, nu_A, nu_B, nu_C, C_bc, C_0, N_comp, dt, T)
# Visualizing
fig, ax = subplots()
x = [C_old[1].domain.facecenters.x[1]; C_old[1].domain.cellcenters.x; C_old[1].domain.facecenters.x[end]]
  phi = [0.5*(C_old[1].value[1]+C_old[1].value[2]); C_old[1].value[2:end-1]; 0.5*(C_old[1].value[end-1]+C_old[1].value[end])]
  plot(x, phi, label=L"$C_{A},\ \nu_A=$""$nu_A"L",\ a=""$a")

  x = [C_old[2].domain.facecenters.x[1]; C_old[2].domain.cellcenters.x; C_old[2].domain.facecenters.x[end]]
    phi = [0.5*(C_old[2].value[1]+C_old[2].value[2]); C_old[2].value[2:end-1]; 0.5*(C_old[2].value[end-1]+C_old[2].value[end])]
    plot(x, phi, label=L"$C_{B},\ \nu_B=$""$nu_B"L",\ b=""$b")

    x = [C_old[3].domain.facecenters.x[1]; C_old[3].domain.cellcenters.x; C_old[3].domain.facecenters.x[end]]
      phi = [0.5*(C_old[3].value[1]+C_old[3].value[2]); C_old[3].value[2:end-1]; 0.5*(C_old[3].value[end-1]+C_old[3].value[end])]
      plot(x, phi, label=L"$C_{C},\ \nu_C=$""$nu_C"L",\ c=""$c")
ax[:legend]()
grid("on")
xlabel(L"$L_x\ [m]$")
ylabel(L"$C_{\alpha}\ [mol/m^3]$")
#ax[:set_title](L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(\mathcal{D}_{\alpha}\phi\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (GIA)")
ax[:set_title](L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (GIA)")
ax[:legend]()
savefig("C:/Users/Frederik/Dropbox/Backup/DTU/9. semester/Coupled Reactive Transport Modelling/Code/concentration_GIA_eps.png")
#savefig("C:/Users/Frederik/Dropbox/Backup/DTU/9. semester/Coupled Reactive Transport Modelling/Code/concentration_SNIA.png")

#CRT_SIA title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (SIA)"

#CRT_SNIA title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (SNIA)"

#CRT_GIA_eps title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}\epsilon=0,\quad K_{eq}=$""$K_eq"" (GIA)"

#CRT_GIA_kin_eq title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}(K_{eq}C_A^aC_B^b-C_C^c)=0,\quad K_{eq}=$""$K_eq"" (GIA)"

#CRT_GIA_kin_fb title:
#L"$\phi\frac{\partial C_{\alpha}}{\partial t}+\nabla\cdot(uC_{\alpha})-\nabla\cdot(D_{\alpha}\nabla C_{\alpha})+\nu_{\alpha}(k_{f}C_A^aC_B^b-k_{b}C_C^c)=0,\quad K_{eq}=$""$K_eq"" (GIA)"


"""
        electre(decisionMat, weights, fns)

Apply ELECTRE (ELimination Et Choice Translating REality) method
for a given matrix and weights.

# Arguments:
 - `decisionMat::DataFrame`: n × m matrix of objective values for n candidate (or strategy) and m criteria
 - `weights::Array{Float64, 1}`: m-vector of weights that sum up to 1.0. If the sum of weights is not 1.0, it is automatically normalized.
 - `fns::Array{Function, 1}`: m-vector of function that are either minimize or maximize.

# Description
electre() applies the ELECTRE method to rank n strategies subject to m criteria which are supposed to be either maximized or minimized.
C and D values are used to determine the best strategy. If the strategy with the highest C value
is same as the strategy with the lowest D value than the solution is unique. Otherwise, two strategies
are reported as the solution.

# Output
- `::ElectreResult`: ElectreResult object that holds multiple outputs including scores and best index.

# Examples
```julia-repl
julia> w =  [0.110, 0.035, 0.379, 0.384, 0.002, 0.002, 0.010, 0.077];
julia> Amat = [
      100 92 10 2 80 70 95 80 ;
      80  70 8  4 100 80 80 90 ;
      90 85 5 0 75 95 70 70 ;
      70 88 20 18 60 90 95 85
    ];

julia> dmat = makeDecisionMatrix(Amat);
julia> fns = makeminmax([maximum for i in 1:8]);
julia> result = electre(dmat, w, fns)

julia> dmat
4×8 DataFrame
 Row │ Crt1     Crt2     Crt3     Crt4     Crt5     Crt6     Crt7     Crt8
     │ Float64  Float64  Float64  Float64  Float64  Float64  Float64  Float64
─────┼────────────────────────────────────────────────────────────────────────
   1 │   100.0     92.0     10.0      2.0     80.0     70.0     95.0     80.0
   2 │    80.0     70.0      8.0      4.0    100.0     80.0     80.0     90.0
   3 │    90.0     85.0      5.0      0.0     75.0     95.0     70.0     70.0
   4 │    70.0     88.0     20.0     18.0     60.0     90.0     95.0     85.0

julia> result.bestIndex
(4,)

julia> result.C
4-element Array{Float64,1}:
  0.3693693693693696
  0.01501501501501501
 -2.473473473473473
  2.0890890890890894

julia> result.D
4-element Array{Float64,1}:
  0.1914244325928971
 -0.19039293350192432
  2.884307608766315
 -2.885339107857288
```

# References
Celikbilek Yakup, Cok Kriterli Karar Verme Yontemleri, Aciklamali ve Karsilastirmali
Saglik Bilimleri Uygulamalari ile. Editor: Muhlis Ozdemir, Nobel Kitabevi, Ankara, 2018
"""
function electre(decisionMat::DataFrame, weights::Array{Float64,1}, fns::Array{Function,1})::ElectreResult

    w = unitize(weights)

    nalternatives, ncriteria = size(decisionMat)

    normalizedMat = normalize(decisionMat)
    weightednormalizedMat = w * normalizedMat

    fitnessTable = []
    nonfitnessTable = []

    for i in 1:nalternatives
        for j in 1:nalternatives
            if i != j
                betterlist = []
                worselist = []
                for h in 1:ncriteria
                    if fns[h] == maximum
                        if weightednormalizedMat[i, h] >= weightednormalizedMat[j, h]
                            push!(betterlist, h)
                        else
                            push!(worselist, h)
                        end
                    else
                        if weightednormalizedMat[i, h] <= weightednormalizedMat[j, h]
                push!(betterlist, h)
                        else
                            push!(worselist, h)
                        end
                    end
                end
                fitnesstableelement = Dict(:i => i, :j => j, :set => betterlist)
                push!(fitnessTable, fitnesstableelement)

                nonfitnesstableelement = Dict(:i => i, :j => j, :set => worselist)
                push!(nonfitnessTable, nonfitnesstableelement)
            end
        end
    end

    fitnessmatrix = zeros(Float64, nalternatives, nalternatives)
    nonfitnessmatrix = zeros(Float64, nalternatives, nalternatives)

    for elements in fitnessTable
        i = elements[:i]
        j = elements[:j]
        elemlist = elements[:set]

        CC = sum(w[elemlist])
        fitnessmatrix[i, j] = CC
    end

    nonfitnessmatrix = zeros(Float64, nalternatives, nalternatives)
    for elements in nonfitnessTable
        i = elements[:i]
        j = elements[:j]
        elemlist = elements[:set]

        r_ik       = weightednormalizedMat[i, elemlist]
        r_jk       = weightednormalizedMat[j, elemlist]
        r_ik_full  = weightednormalizedMat[i, :]
        r_jk_full  = weightednormalizedMat[j, :]

        if length(r_ik) > 0 && length(r_jk) > 0
            nom = maximum(abs.(r_ik - r_jk))
            dom = maximum(abs.(r_ik_full - r_jk_full))
            nonfitnessmatrix[i, j] = nom / dom
        else
            nonfitnessmatrix[i, j] = Inf64
        end  # end if

    end # end for

    C = zeros(Float64, nalternatives)
    D = zeros(Float64, nalternatives)
    for i in 1:nalternatives
        C[i] = sum(fitnessmatrix[i, :]) - sum(fitnessmatrix[:, i])
        D[i] = sum(nonfitnessmatrix[i, :]) - sum(nonfitnessmatrix[:, i])
    end

    best_C_index = sortperm(C) |> last
    best_D_index = sortperm(D) |> first
    best = nothing

    if best_C_index == best_D_index
        best = (best_C_index,)
    else
        best = (best_C_index, best_D_index)
    end

    result = ElectreResult(
        decisionMat,
        w,
        weightednormalizedMat,
        fitnessTable,
        nonfitnessTable,
        fitnessmatrix,
        nonfitnessmatrix,
        C,
        D,
        best
    )

    return result
end

"""
        electre(setting)

Apply ELECTRE (ELimination Et Choice Translating REality) method
for a given matrix and weights.

# Arguments:
 - `setting::MCDMSetting`: MCDMSetting object.

# Description
electre() applies the ELECTRE method to rank n strategies subject to m criteria which are supposed to be either maximized or minimized.
C and D values are used to determine the best strategy. If the strategy with the highest C value
is same as the strategy with the lowest D value than the solution is unique. Otherwise, two strategies
are reported as the solution.

# Output
- `::ElectreResult`: ElectreResult object that holds multiple outputs including scores and best index.
"""
function electre(setting::MCDMSetting)::ElectreResult
    electre(
        setting.df,
        setting.weights,
        setting.fns
    )
end


using Equate
using Statistics

@testset "Mean" begin
    res = Mean(sgtabX, sgtabY)
    popXonY = ExpandTable(res.table.lYx, ACTmath.xcount)
    @test round(mean(popXonY); digits = 4) == 18.9798
end

@testset "Linear" begin
    res = Linear(sgtabX, sgtabY)
    popXonY = ExpandTable(res.table.lYx, ACTmath.xcount)
    @test round(mean(popXonY); digits = 4) == 18.9798
    @test round(std(popXonY; corrected = false); digits = 4) == 8.9394 # 8.9393 in KB p.50
end

@testset "Equipercentile (lower)" begin
    res = Equipercentile(sgtabX, sgtabY)
    popXonY = ExpandTable(res.table.eYx, ACTmath.xcount)
    @test round(mean(popXonY); digits = 4) == 18.9799
    @test round(std(popXonY; corrected = false); digits = 4) == 8.9352
end

@testset "Percentile Rank Function" begin
    P̂ = round2.(Equate.PRF.(0:1:40, Ref(sgtabX)); digits = 2)
    @test sum(KBp48 .== P̂) == 41
end


using NLsolve

λ_p    = 0.15             # net markup
α      = 1-0.60/(1+λ_p)
δ    = 0.025
π_ss = 1.02^(1/4)
γ    = log(1.02)/4
disc = 0.99

h    = 0.5
g_x  = 0.18

s    = 0.995
θ    = 1/4
f_s1 = 3/4
f_h1 = 1/4
A2   = 0.7
v    = 0.7
t_x  = 0.2
τ_x  = 0
#ψ_h1 = θ * f_h1 / 1.1
#ψ_h2 = θ * (1 - f_h1) / 1.1
ψ_h1 = 0
ψ_h2 = 0
H1   = 1
H2   = 1

L1   = ((1-θ) * f_s1 + θ * f_h1) * H1
L2   = ((1-θ) * (1 - f_s1) + θ * (1 - f_h1)) * H2

ρ    = exp(γ) / disc - (1 - δ)

w1   = (α ^ α * (1 - α) ^ (1 - α) / (1 + λ_p) / (ρ ^ α)) ^ (1 / (1 - α))
w2   = A2 * w1

klr1 = w1 / ρ * α / (1 - α)
klr2 = A2 * klr1
k    = klr1 * L1

F1   = λ_p / (1 + λ_p) * L1 * klr1 ^ α
y1   = L1 * klr1 ^ α - F1
y    = y1 / v
y2   = (1 - v) * y
F2   = λ_p * y2 / A2
x    = y

is   = (1 - (1 - δ) * exp(-γ)) * exp(γ) * k / (1 - θ)
g    = g_x * x
c    = y - g - (1 - θ) * is

τ    = τ_x * x
τ_h1 = τ * ψ_h1 / θ / (1 - f_h1)
τ_h2 = τ * ψ_h2 / θ / (1 - f_h1)
τ_s  = τ - θ * (f_h1 * τ_h1 + (1 - f_h1) * τ_h2)

t    = t_x * x
t_h1 = t * ψ_h1 / θ / f_h1
t_h2 = t * ψ_h2 / θ / (1 - f_h1)
t_s  = (t - θ * ( f_h1 * t_h1 + (1 - f_h1) * t_h2)) / (1 - θ)

c_h10 = c / θ / f_h1
c_h20 = c / θ / (1 - f_h1)
c_s0  = c / (1 - θ)
λ_h10 = exp(γ) / (exp(γ) * c_h10 - h * c)
λ_h20 = exp(γ) / (exp(γ) * c_h20 - h * c)
λ_s0  = exp(γ) / (exp(γ) * c_s0 - h * c)
bR0   = (g - t + τ) / (1 - (1.025 / 4) / exp(γ) / π_ss)
R0    = (1.025 / 4)

function distSS(x)
global w1, w2, H1, H2, th1, th2, τ_h1, τ_h2, c, g, t, τ
global s, θ, f_h1, f_s1, γ, π_ss, h, disc

c_h1 = x[1]
c_h2 = x[2]
c_s   = x[3]
λ_h1 = x[4]
λ_h2 = x[5]
λ_s  = x[6]
bR   = x[7]
R    = x[8]

r = zeros(8)

r[1] = w1 * H1 - t_h1 + τ_h1 + R * bR * f_s1 * (1 - s) / θ / f_h1 / exp(γ) / π_ss - c_h1
r[2] = w2 * H2 - t_h2 + τ_h2 + R * bR * (1-f_s1) * (1 - s) / θ / (1 - f_h1) / exp(γ) / π_ss - c_h2
r[3] = c / (1 - θ) - θ * (f_h1 * c_h1 + (1 - f_h1) * c_h2) / (1 - θ) - c_s

r[4] = exp(γ) / (exp(γ) - h) / c_h1 - λ_h1
r[5] = exp(γ) / (exp(γ) - h) / c_h2 - λ_h2
r[6] = exp(γ) / (exp(γ) - h) / c_s  - λ_s

r[7] = (g - t + τ) / (1 - R / exp(γ) / π_ss) - bR

r[8] = exp(γ) * π_ss / disc * λ_s / (s * λ_s + (1 - s) * (f_s1 * λ_h1 + (1 - f_s1) * λ_h2)) - R

return r

end

init = [c_h10 c_h20 c_s0 λ_h10 λ_h20 λ_s0 bR0 R0]
ss = nlsolve(distSS, init)

solution = ss.zero

println(solution)


using OASA
using Base.Test

# Synthesize random data
X = hcat(ones(100, 1), rand(100, 3))
truth_theta = [.5 1.2 3.3 5.0]
Y = truth_theta * X'

# Run some basic tests
function TestA()
    yh = infer(myModel, X)
    @test size(yh) == size(Y)
    @test assess(myModel, Y, yh) > 0.0
    @test size(assessgrad(myModel, Y, yh, X)) == (1, size(X,2))
end

# Ensure that all of the objective functions are operational
@testset "Objective Functions" begin
    @testset "L1 Regression" begin
        myModel = Model(L1_Obj(),
                    Linear(),
                    GradientDescent(),
                    SingleLayer(zeros(1, size(X,2))),
                    Dict())
        TestA()
    end
    @testset "L2 Regression" begin
        myModel.obj = L2_Obj()
        TestA()
    end
    @testset "Logistic Regression" begin
        myModel.obj = Logit_Obj()
        TestA()
    end
end


# The package contains tools to solve an N-body problem in Newtonian gravity,
# and plot and animate the orbits and the consverved quantities.
# This package is based on the DifferentialEquations package tutorial:
# https://github.com/JuliaDiffEq/DiffEqTutorials.jl/blob/master/PhysicalModels/Outer-Solar-System.ipynb

module PlanetarySystem

using DifferentialEquations
using RecursiveArrayTools
using Plots

const G = 2.95912208286e-4 # Newton's constant, units: [L]= AU, [M]= M_sun, [t]= Earth days
const ∑ = sum # symbol for sum function
const N = 5 #Number of bodies (central star included)

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

# Physical model:

# Gravitational newtonian potential for the NBsolution() solver.
# p is a parameter, t the time, the coordinates x, y, z are arrays running over the bodies,
# M is the array of the bodies' masses
potential(p, t, x, y, z, M) = -G*∑(i->∑(j->(M[i]*M[j])/sqrt((x[i]-x[j])^2 + (y[i]-y[j])^2 + (z[i]-z[j])^2), 1:i-1), 2:N)

# Kinetic energy
# the velocity components  vx, vy, vz are 2D arrays running over the bodies and the time steps,
# M is the array of the bodies' masses
kinetic(vx, vy, vz, M) = 1/2 * ( ∑(M.*vx.*vx,dims=1) + ∑(M.*vy.*vy,dims=1) + ∑(M.*vz.*vz,dims=1) )
# potential, to compute the Hamiltonian
# the coordinates x, y, z are 2D arrays running over the bodies and the time steps,
# M is the array of the bodies' masses
pot(x, y, z, M)= - G*∑(i->∑(j->(M[i]*M[j])./sqrt.((x[i,:]-x[j,:]).^2+(y[i,:]-y[j,:]).^2+(z[i,:]-z[j,:]).^2), 1:i-1), 2:N)

# Hamiltonian (total energy) - variables as above
H(x, y, z, vx, vy, vz, M) = kinetic(vx, vy, vz, M) + pot(x, y, z, M)'

# Angular momentum components  - variables as above
Lx(x, y, z, vx, vy, vz, M) =  ∑(i-> M[i] .*  (vy[i,:] .* x[i,:] - vx[i,:] .* y[i,:]), 1:N)'
Ly(x, y, z, vx, vy, vz, M) =  ∑(i-> M[i] .*  (-vz[i,:] .* x[i,:] + vx[i,:] .* z[i,:]), 1:N)'
Lz(x, y, z, vx, vy, vz, M) =  ∑(i-> M[i] .*  (vz[i,:] .* y[i,:] - vy[i,:] .* z[i,:]), 1:N)'

export H, Lx, Ly, Lz

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

# Main
export SolveAndPlot
function SolveAndPlot(M::Array{Float64,1}, vel, pos, tspan::Tuple{Float64,Float64}, planets::Array{String,1}, plotfilename::String, animfilename::String, fifilename::String)

sol=NBsolution(M, vel, pos, tspan)
myplot(sol,plotfilename,planets)
animation(sol,animfilename,planets)
plot_first_integrals(sol, M, fifilename, planets)

end

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

# Solves the N-body problem with a sympletic method, Yoshida6.

# M is the array of the bodies' masses,
# vel contains the initial velocity vectors of all bodies (see example in Readme)
# pos contains the initial position of all bodies (see example in Readme)
# tspan is the time interval

export NBsolution
function NBsolution(M::Array{Float64,1}, vel, pos, tspan::Tuple{Float64,Float64})
    nprob = NBodyProblem(potential, M, vel, pos, tspan)
    sol = solve(nprob, Yoshida6(), dt=10) # can change the method or the time step dt
end

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

# Plots the orbits in coordinate space and saves to file
# sol is the output of NBsolution()
# filename is the name of the file where the plot is saved
# planets is the array of the planets' names

export myplot
function myplot(sol, filename::String, planets::Array{String,1})
    plot(size=(500, 400),
         title  = "Planetary orbits",
         legend=:bottomleft,
         xlim=(-30,30), ylim=(-30,30), zlim=(-30,30)
         ,markerstrokewidth=0
         #,bg=:black
         #,xlabel = "x", ylabel = "y", zlabel = "z"
         )

    for i in 1:N
    scatter!(sol,
              vars=(3*N+i,4*N+i,5*N+i),
              label=planets[i],
              markersize=1,
              markerstrokewidth=0)
    end

    savefig(filename)
end

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

# Animates the planets and saves to file
# see above

export animation
function animation(sol, filename::String, planets::Array{String,1})

    nplots=500
    element(i)=round(Int, length(sol.t)/nplots-1)*i

    anim = @animate for i=1:nplots #
        scatter((sol[3*N+1,element(i)],sol[4*N+1,element(i)],sol[5*N+1,element(i)]),
                #xlabel = "x", ylabel = "y", zlabel = "z",
                xlim=(-30,30), ylim=(-30,30), zlim=(-30,30),
                label=planets[1],
                markerstrokewidth=0,
                markercolor=:yellow,
                marker=10,
                title  = "Planetary system",
                size=(500, 400))
        if(N>1)
            for j in 2:(N)
        scatter!((sol[3*N+j,element(i)],sol[4*N+j,element(i)],sol[5*N+j,element(i)]),
                 marker = 5,
                 markerstrokewidth=0,
                 label=planets[j])
            end
        end
    end

    gif(anim, filename, fps = 17) #change fps to adjust the speed of the animation

end

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

# Plots the fractional variation of conserved quantities (energy and angular momentum)
# see above

export plot_first_integrals
function plot_first_integrals(sol, M::Array{Float64,1}, filename::String, planets::Array{String,1})

    vx=sol[1:N,:]
    vy=sol[N+1:2*N,:]
    vz=sol[2*N+1:3*N,:]
    x= sol[3*N+1:4*N,:]
    y= sol[4*N+1:5*N,:]
    z= sol[5*N+1:6*N,:]

    plot(sol.t, (H(x, y, z, vx, vy, vz, M)[1] .- H(x, y, z, vx, vy, vz, M)')./H(x, y, z, vx, vy, vz, M)[1],
         lab="Energy variation",
         title="Conserved Quantities",
         xlabel="t",
         size=(500, 400))
    plot!(sol.t, (Lx(x, y, z, vx, vy, vz, M)[1].-Lx(x, y, z, vx, vy, vz, M)')./Lx(x, y, z, vx, vy, vz, M)[1], lab="Angular momentum variation, x")
    plot!(sol.t, (Ly(x, y, z, vx, vy, vz, M)[1].-Ly(x, y, z, vx, vy, vz, M)')./Ly(x, y, z, vx, vy, vz, M)[1], lab="Angular momentum variation, y")
    plot!(sol.t, (Lz(x, y, z, vx, vy, vz, M)[1].-Lz(x, y, z, vx, vy, vz, M)')./Lz(x, y, z, vx, vy, vz, M)[1], lab="Angular momentum variation, z")

    savefig(filename)

end

end # module


import Base: +, -, *, <
using HCubature, ForwardDiff

# Function composition
(+)(a::Function, b::Function) = (x, y) -> (@inline a)(x, y) + (@inline b)(x, y)
(*)(a::Function, b::Function) = (x, y) -> (@inline a)(x, y) * (@inline b)(x, y)

(+)(a::Function, b::Number) = (x, y) -> (@inline a)(x, y) + b
(+)(a::Number, b::Function) = (x, y) -> (@inline b)(x, y) + a
(-)(a::Function, b::Number) = (x, y) -> (@inline a)(x, y) - b
(-)(a::Number, b::Function) = (x, y) -> (@inline b)(x, y) - a

# Cartesian position comparison
(<)(a::Tuple, b::Tuple) = all(a .< b)

# Derivatives
∂x(f) = (x, y) -> ForwardDiff.derivative(x -> f(x, y), x)
∂y(f) = (x, y) -> ForwardDiff.derivative(y -> f(x, y), y)

# Integrals
∫, ∫∫ = hquadrature, hcubature
∑ = sum

# Notice valid only for linear functions
∫dlX(f, a, b, y, slope) = ∫(x -> f(x, y(x)), a, b)[1] * sqrt(1 + slope^2) # helper
∫dlY(f, a, b, x, slope) = ∫(y -> f(x(y), y), a, b)[1] * sqrt(1 + slope^2) # helper

∫dl(f, e::EdgeX) = ∫dlX(f, e.range..., e.y, e.a)
∫dl(f, e::EdgeY) = ∫dlY(f, e.range..., e.x, e.a)
∫dl(f, edges) = sum(∫dl.(f, edges))

initdiv = trunc(Int, ceil(max(1/dx,1/dy, 1)))
∫∫dS(f, A, B) = hcubature(v -> f(v...), A, B, initdiv=initdiv)[1]


using ControlSystems;
using DifferentialEquations;
using PyPlot;
using Polynomials;

A=[0 1;1 0];
B=[0;1];
C=[0 0];
D=0;

sys=ss(A,B,C,D,0);

u(x,t)=[-2.11803 -2.11803 ]*x;
x0=[1;1];
tspan=0:0.01:20;
y, t, x, uout = lsim(sys,u,tspan,x0=x0);

plot(t,x[:,1],label="x1")
plot(t,x[:,2],label="x2")
xlabel("time (t)")
legend()
grid(true)


module Day17

struct TargetArea
    x::UnitRange{Int}
    y::UnitRange{Int}
end

function parse_target(input::String)
    input = split(input, ": ")[2]
    xpart, ypart = split(input, ", ")
    xpart = split(split(xpart, "=")[2], "..")
    ypart = split(split(ypart, "=")[2], "..")

    x = parse(Int, xpart[1]):parse(Int, xpart[2])
    y = parse(Int, ypart[1]):parse(Int, ypart[2])

    TargetArea(x, y)
end

function part1(target::TargetArea)
    S = abs(target.y[1])
    Int((S - 1) * S / 2)
end

function match_x_in_steps(range::UnitRange{Int}, n)
    # for 2 steps,
    # l <= v + v-1 <= r
    # for 3 steps,
    # l <= v + v-1 + v-2 <= r --> l+3 <= 3v <= r+3 --> l/3+1 <= v <= r/3+1
    l = range[1]
    r = range[end]
    triangle = n * (n-1) / 2
    lower = (l+triangle)/n |> ceil |> Int
    upper = (r + triangle)/n |> floor |> Int
    return lower:upper
end

function match_y_in_steps(range::UnitRange{Int}, n)
    # the problem is symmetric, so only need to consider negative vy.
    # for 1 step,
    # t <= v <= b
    # for 2 steps,
    # t <= v + v+1 <= b
    # for 3 steps,
    # t <= v + v+1 + v+2 <= b
    # ...
    # t-tri <= nv <= b-tri
    t = range[end] |> abs
    b = range[1] |> abs
    triangle = n * (n-1) / 2
    lower = (t - triangle) / n |> ceil |> Int
    upper = (b - triangle) / n |> floor |> Int
    return -(lower:upper)
end

function match(target::TargetArea)
    vxs = Dict()
    stalled_out_vxs = Int[]
    n = 1
    while n*(n+1)/2 <= target.x[end]
        matches = match_x_in_steps(target.x, n)
        push!(vxs, n => Array(matches))
        # For all of the x velocities which arrive in vx steps,
        # they are "stalled out".
        if n in matches
            push!(stalled_out_vxs, n)
        end
        n += 1
    end

    vys = Dict()
    n = 1
    while n*(n+1)/2 <= abs(target.y[1])
        push!(vys, n => Array(match_y_in_steps(target.y, n)))
        n += 1
    end
    i = 0
    k = 0
    for (n, vy) in vys
        vy = vy[vy .< 0]
        parabola_steps = 2abs.(vy) .- 1
        original_vy = -1 .- vy
        for (pn, vy) in zip(parabola_steps, original_vy)
            if (pn+n) ∈ keys(vys)
                push!(vys[pn+n], vy)
            else
                push!(vys, (pn+n) => [vy])
            end
        end
        i += 1
    end

    velocities = Set()
    # Need to find a solution for each vy
    for (n, vy) in vys
        if n ∈ keys(vxs)
            n_step_velocities = tuple.(vxs[n], Array(vys[n])') |> collect
            union!(velocities, vec(n_step_velocities))
        else
            stalled_out_matches = stalled_out_vxs[stalled_out_vxs .<= n]
            vs = tuple.(stalled_out_vxs, Array(vys[n])') |> collect
            union!(velocities, vs)
        end
    end
    velocities
end

function run()
    test = "target area: x=20..30, y=-10..-5"
    testarea = parse_target(test)

    @assert part1(testarea) == 45
    @assert length(match(testarea)) == 112

    input = "target area: x=155..182, y=-117..-67"
    target = parse_target(input)
    display(part1(target))
    display(length(match(target)))
end

function test_part2()
    expected = """23,-10  25,-9   27,-5   29,-6   22,-6   21,-7   9,0     27,-7   24,-5
25,-7   26,-6   25,-5   6,8     11,-2   20,-5   29,-10  6,3     28,-7
8,0     30,-6   29,-8   20,-10  6,7     6,4     6,1     14,-4   21,-6
26,-10  7,-1    7,7     8,-1    21,-9   6,2     20,-7   30,-10  14,-3
20,-8   13,-2   7,3     28,-8   29,-9   15,-3   22,-5   26,-8   25,-8
25,-6   15,-4   9,-2    15,-2   12,-2   28,-9   12,-3   24,-6   23,-7
25,-10  7,8     11,-3   26,-7   7,1     23,-9   6,0     22,-10  27,-6
8,1     22,-8   13,-4   7,6     28,-6   11,-4   12,-4   26,-9   7,4
24,-10  23,-8   30,-8   7,0     9,-1    10,-1   26,-5   22,-9   6,5
7,5     23,-6   28,-10  10,-2   11,-1   20,-9   14,-2   29,-7   13,-3
23,-5   24,-8   27,-9   30,-7   28,-5   21,-10  7,9     6,6     21,-5
27,-10  7,2     30,-9   21,-8   22,-7   24,-9   20,-6   6,9     29,-5
8,-2    27,-8   30,-5   24,-7"""
    expected = split(expected) |> Set

    test = "target area: x=20..30, y=-10..-5"
    testarea = parse_target(test)
    velocities = match(testarea)
    actual = [join(v, ",") for v in velocities] |> Set

    setdiff(actual, expected) |> display
    setdiff(expected, actual) |> display
end

end


using Test
using Distributions
import Distributions: normpdf, normcdf, normlogpdf, normlogcdf

@testset "SkewNormal" begin
    @test_throws ArgumentError SkewNormal(0.0, 0.0, 0.0)
    d1 = SkewNormal(1, 2, 3)
    d2 = SkewNormal(1.0f0, 2, 3)
    @test partype(d1) == Float64
    @test partype(d2) == Float32
    @test params(d1) == (1.0, 2.0, 3.0)
    # Azzalini sn: sprintf("%.17f",dsn(3.3, xi=1, omega=2, alpha=3))
    @test pdf(d1, 3.3) ≈ 0.20587854616839998
    @test minimum(d1) ≈ -Inf
    @test maximum(d1) ≈  Inf
    @test logpdf(d1, 3.3) ≈ log(pdf(d1, 3.3))
    ## cdf and quantile: when we get Owen's T
    #@test cdf(d, 4.5) ≈ 1.0 #when we get Owen's T
    @test mean(d1) ≈ 2.513879513212096
    @test std(d1) ≈ 1.306969326142243
    #
    d0 = SkewNormal(0.0, 1.0, 0.0)
    @test SkewNormal() == d0
    d3 = SkewNormal(0.5, 2.2, 0.0)
    d4 = Normal(0.5, 2.2)
    #
    @test pdf(d3, 3.3) == Distributions.pdf(d4, 3.3)
    @test pdf.(d3, 1:3) == Distributions.pdf.(d4, 1:3)
    a = mean(d3), var(d3), std(d3)
    b = Distributions.mean(d4), Distributions.var(d4), Distributions.std(d4)
    @test a == b
    @test skewness(d3) == Distributions.skewness(d4)
    @test kurtosis(d3) == Distributions.kurtosis(d4)
    @test mgf(d3, 2.25) == Distributions.mgf(d4, 2.25)
    @test cf(d3, 2.25) == Distributions.cf(d4, 2.25)
end

# R code using Azzalini's package sn
#library("sn")
#sprintf("%.17f", dsn(3.3, xi=1, omega=2, alpha=3))
#f1 <- makeSECdistr(dp=c(1,2,3), family="SN", name="First-SN")
#sprintf("%.15f", mean(f1))
#sprintf("%.15f", sd(f1))


using Plots
using DataStructures

function inbounds(canvas_size, loc)
    if minimum(Tuple(loc)) < 1
        return false
    elseif loc[2] > canvas_size[2]
        return false
    elseif loc[1] > canvas_size[1]
        return false
    else
        return true
    end
end

function create_circle!(canvas, loc::CartesianIndex, radius, threshold,
                        old_val)
    for i = 1:size(canvas)[1]
        for j = 1:size(canvas)[2]
            if ((radius-threshold)^2
                < (i-loc[1])^2+(j-loc[2])^2
                < (radius+threshold)^2)
                canvas[i,j] = old_val
            end
        end
    end
end

function color!(canvas, loc::CartesianIndex, old_val, new_val)
    # bounds check
    if !inbounds(size(canvas), loc)
        return
    end

    if (old_val == new_val)
        return
    elseif (canvas[loc] != old_val)
        return
    else
        canvas[loc] = new_val
    end
end

function find_neighbors(canvas, loc::CartesianIndex, old_val, new_val;
                       direction = "all")

    neighbors =  []
    if direction == "all"
        possible_neighbors = [loc + CartesianIndex(0, 1),
                              loc + CartesianIndex(1, 0),
                              loc + CartesianIndex(0, -1),
                              loc + CartesianIndex(-1, 0)]
    elseif direction == "north"
        temp_loc = loc + CartesianIndex(1, 0)
        possible_neighbors = [temp_loc]
    elseif direction == "south"
        temp_loc = loc + CartesianIndex(-1, 0)
        possible_neighbors = [temp_loc]
    elseif direction == "east"
        temp_loc = loc + CartesianIndex(0, 1)
        possible_neighbors = CartesianIndex[]
        while canvas[temp_loc] == old_val
            push!(possible_neighbors, temp_loc)
            temp_loc = temp_loc + CartesianIndex(0, 1)
        end
    elseif direction == "west"
        temp_loc = loc + CartesianIndex(0, -1)
        possible_neighbors = CartesianIndex[]
        while canvas[temp_loc] == old_val
            push!(possible_neighbors, temp_loc)
            temp_loc = temp_loc + CartesianIndex(0, -1)
        end
    else
        error("direction ", direction, " not found!")
    end
    for possible_neighbor in possible_neighbors
        if inbounds(size(canvas), possible_neighbor) &&
           canvas[possible_neighbor] == old_val
            push!(neighbors, possible_neighbor)
        end
    end

    return neighbors
end

function stack_fill!(canvas, loc::CartesianIndex, old_val, new_val)
    if new_val == old_val
        return
    end

    s = Stack{CartesianIndex}()
    #s = CartesianIndex[]
    push!(s, loc)

    while length(s) > 0
        temp_loc = pop!(s)
        color!(canvas, temp_loc, old_val, new_val)
        possible_neighbors = find_neighbors(canvas, temp_loc, old_val, new_val)
        for neighbor in possible_neighbors
            push!(s,neighbor)
        end

    end
end

function queue_fill!(canvas, loc::CartesianIndex, old_val, new_val)
    if new_val == old_val
        return
    end

    q = Queue{CartesianIndex}()
    enqueue!(q, loc)

    while length(q) > 0
        temp_loc = dequeue!(q)
        color!(canvas, temp_loc, old_val, new_val)
        possible_neighbors = find_neighbors(canvas, temp_loc, old_val, new_val)
        for neighbor in possible_neighbors
            enqueue!(q,neighbor)
        end

    end
end

function recursive_fill!(canvas, loc::CartesianIndex, old_val, new_val)
    if (old_val == new_val)
        return
    elseif (canvas[loc] != old_val)
        return
    else
        canvas[loc] = new_val
    end

    possible_neighbors = find_neighbors(canvas, loc, old_val, new_val)
    for possible_neighbor in possible_neighbors
        recursive_fill!(canvas, possible_neighbor, old_val, new_val)
    end
end

# TODO: find array of *all* west and east elements in find_neighbor
function loop_fill!(canvas, loc::CartesianIndex, old_val, new_val)

    if (old_val == new_val)
        return
    elseif (canvas[loc] != old_val)
        return
    end

    q = Queue{CartesianIndex}()
    enqueue!(q, loc)

    for element in q
        west_elements = find_neighbors(canvas, element, old_val, new_val;
                       direction = "west")
        east_elements = find_neighbors(canvas, element, old_val, new_val;
                       direction = "east")
        for temp_loc in vcat(west_elements, east_elements)
            color!(canvas, temp_loc, old_val, new_val)

            north_element = find_neighbors(canvas, temp_loc, old_val,
                            new_val; direction = "north")
            south_element = find_neighbors(canvas, temp_loc, old_val,
                            new_val; direction = "south")

            if length(north_element) > 0
                enqueue!(q,north_element[1])
            end
            if length(south_element) > 0
                enqueue!(q,south_element[1])
            end
        end

        color!(canvas, element, old_val, new_val)
    end
end

function init_canvas()
    canvas = zeros(100,100)
    loc = CartesianIndex(50,50)
    create_circle!(canvas, loc, 7, 2, 0.75)
    return canvas
end

function main()
    canvas = init_canvas()
    loc2 = CartesianIndex(50,50)
    #recursive_fill!(canvas, loc2, 0.0, 0.5)
    #queue_fill!(canvas, loc2, 0.0, 0.5)
    #stack_fill!(canvas, loc2, 0.0, 0.5)
    loop_fill!(canvas, loc2, 0.0, 0.5)

    heatmap(canvas, color=:coolwarm)
end

main()


cd(@__DIR__)
using DocStringExtensions
using Statistics
using StatsBase
using LinearAlgebra
using DataFrames
using FastTransforms
using Distributions
using Random
# Base.:|>(args::Tuple, f) = f(args...)

C(r, a = 2.0) = (1 + r^2)^(-a / 2)

periodic_dist(x, L) = x ≤ L / 2 ? x : L - x

"""
    generate_corr_matrix(dims::Tuple, C::Function; T::Type = Float64)

Calculate the function `C` for each distance `r` and returns an array where each C(r) is stored.
The distnace `r` is measured from upper left corner and implies periodic boundary conditions.

# Arguments:
- `dims::Tuple`: dimensions of the desired lattice
- `C::Function`: desired correlation function `C(r)`
- `T::Type = Float64`: type of the retured array
"""
function generate_corr_matrix(dims::Tuple, C::Function; T::Type = Float64)
	Λ = Array{T}(undef, dims)
	fill_periodic_distances!(Λ, C)
	return Λ
end

"""
    fill_periodic_distances!(Λ::AbstractArray, C::Function)

Helper function for the `generate_corr_matrix` function.
The type of `Λ` is known on call and therefore can this inner function works faster.

# Arguments:
- `Λ::AbstractArray`: array to save the `C(r)` to
- `C::Function`: desired correlation function `C(r)`
"""
function fill_periodic_distances!(Λ::AbstractArray, C::Function)
	for i ∈ CartesianIndices(Λ)
		r = norm(periodic_dist.(Tuple(i) .- 1, size(Λ)))
		Λ[i] = C(r)
	end
end

"""
    generate_continuous_disorder(dims::Tuple, S::AbstractArray)

Produces an array with dimensions `dims` with correlated disorder where the correlation follows the function `C(r)`

# Arguments:
- `dims::Tuple`: dimensions of the desired lattice
- `C::Function`: desired correlation function `C(r)`
"""
function generate_continuous_disorder(dims::Tuple, S::AbstractArray)
	φq = sqrt.(max.(S, 0)) .* rand.(Normal(0, √(2 * prod(dims)))) # 2V = 2prod(dims) is used because the random numbers are produced only in real space
	φ = real.(ifft(φq))
	return φ
end

"""
    generate_S(dims::Tuple, C::Function)

Produces a spectral density array with dimensions `dims` of the array with `C(r)` entries

# Arguments:
- `dims::Tuple`: dimensions of the desired lattice
- `C::Function`: desired correlation function `C(r)`
"""
generate_S(dims::Tuple, C::Function) = real.(fft(generate_corr_matrix(dims, C)))

"""
    transform_to_discrete_disorder(Λ::AbstractArray, p::Number; T = Int8(1), F = Int8(0))

Produces a truncated array similar to `Λ` but only with `T` and `F` values.
The concentration of `T` is given through `p`.

# Arguments:
- `Λ::AbstractArray`: expect an array with continuous correlated numbers
- `p::Number`: concentration of `T` values
- `T = Int8(1)`: values with concentration `p` to fill into the returned array
- `F = Int8(0)`: values with concentration `1 - p` to fill into the returned array
"""
transform_to_discrete_disorder(Λ::AbstractArray, p::Number; T = Int8(1), F = Int8(0)) = [i ≤ quantile(Normal(0, 1), p) ? T : F for i ∈ Λ]

######################### TEST #########################
Random.seed!(1)
L1 = 1000
L2 = 2000
a = Inf
p = 0.8

S = generate_S((L1, L2), r -> C(r, a))
Λc = generate_continuous_disorder((L1, L2), S)
Λd = transform_to_discrete_disorder(Λc, p)
@show mean(Λc)
@show std(Λc)
@show check_p = counts(Λd, 1:1)[1] / prod(size(Λd))

using PyPlot
clf()
imshow(Λc; interpolation = "none")
savefig("$(L1)x$(L2)_a$(a)_p$(p)_continuous_disorder.pdf")

clf()
imshow(Λd, interpolation = "none")
savefig("$(L1)x$(L2)_a$(a)_p$(p)_discrete_disorder.pdf")

# ######################### TEST SPEED #########################
# using BenchmarkTools

# @btime generate_continuous_disorder((L1, L2), r -> C(r, a))
# @btime transform_to_discrete_disorder(Λc, p)

A = -1.38

Int(round(abs(A))) # reads: perform Int of round of abs of A
round(abs(A)) |> Int
abs(A) |> round |> Int
A |> abs |> round |> Int # reads: take A, do abs, round it, make it an Int

### dot version

A = [0.1, -1.3, 2.5]

Int.(round.(abs.(A)))
round.(abs.(A)) .|> Int
abs.(A) .|> round .|> Int
A .|> abs .|> round .|> Int

### absurd

(1 + 3)^(1 / 2)
(1 + 3) |> x -> ^(x, 1 / 2)

############### in-place operations and |>

A = rand(10000)

B = round.(abs.(A))
A .= round.(abs.(A)) # in-place
A .= A .|> abs .|> round # in-place


import RobotDynamics: jacobian!

""""
Specifies whether the constraint is an equality or inequality constraint.
Valid subtypes are `Equality`, `Inequality` ⟺ `NegativeOrthant`, and `SecondOrderCone`.

The sense of a constraint can be queried using `sense(::AbstractConstraint)`

If `sense(con) <: Conic` (i.e. not `Equality`), then the following operations are supported:
* `Base.in(::Conic, x::StaticVector)`. i.e. `x ∈ cone`
* `projection(::Conic, x::StaticVector)`
* `∇projection(::Conic, J, x::StaticVector)`
* `∇²projection(::Conic, J, x::StaticVector, b::StaticVector)`
"""
abstract type ConstraintSense end
abstract type Conic <: ConstraintSense end

"""
Equality constraints of the form ``g(x) = 0`.
Type singleton, so it is created with `Equality()`.
"""
struct Equality <: ConstraintSense end
"""
Inequality constraints of the form ``h(x) \\leq 0``.
Type singleton, so it is created with `Inequality()`.
Equivalent to `NegativeOrthant`.
"""
struct NegativeOrthant <: Conic end
const Inequality = NegativeOrthant

struct PositiveOrthant <: Conic end

"""
The second-order cone is defined as
``\\|x\\| \\leq t``
where ``x`` and ``t`` are both part of the cone.
TrajectoryOptimization assumes the scalar part ``t`` is
the last element in the vector.
"""
struct SecondOrderCone <: Conic end

dualcone(::NegativeOrthant) = NegativeOrthant()
dualcone(::PositiveOrthant) = PositiveOrthant()
dualcone(::SecondOrderCone) = SecondOrderCone()

projection(::NegativeOrthant, x) = min.(0, x)
projection(::PositiveOrthant, x) = max.(0, x)

function projection(::SecondOrderCone, x::StaticVector)
    # assumes x is stacked [v; s] such that ||v||₂ ≤ s
    s = x[end]
    v = pop(x)
    a = norm(v)
    if a <= -s          # below the cone
        return zero(x)
    elseif a <= s       # in the cone
        return x
    elseif a >= abs(s)  # outside the cone
        return 0.5 * (1 + s/a) * push(v, a)
    else
        throw(ErrorException("Invalid second-order cone projection"))
    end
end

function ∇projection!(::SecondOrderCone, J, x::StaticVector{n}) where n
    s = x[end]
    v = pop(x)
    a = norm(v)
    if a <= -s                               # below cone
        J .*= 0
    elseif a <= s                            # in cone
        J .*= 0
        for i = 1:n
            J[i,i] = 1.0
        end
    elseif a >= abs(s)                       # outside cone
        # scalar
        b = 0.5 * (1 + s/a)
        dbdv = -0.5*s/a^3 * v
        dbds = 0.5 / a

        # dvdv = dbdv * v' + b * oneunit(SMatrix{n-1,n-1,T})
        for i = 1:n-1, j = 1:n-1
            J[i,j] = dbdv[i] * v[j]
            if i == j
                J[i,j] += b
            end
        end

        # dvds
        J[1:n-1,n] .= dbds * v

        # ds
        dsdv = dbdv * a + b * v / a
        dsds = dbds * a
        ds = push(dsdv, dsds)
        J[n,:] .= ds
    else
        throw(ErrorException("Invalid second-order cone projection"))
    end
    return J
end

function Base.in(x, ::SecondOrderCone)
    s = x[end]
    v = pop(x)
    a = norm(v)
    return a <= s
end

function cone_status(::SecondOrderCone, x)
    s = x[end]
    v = pop(x)
    a = norm(v)
    if a <= -s
        return :below
    elseif a <= s
        return :in
    elseif a > abs(s)
        return :outside
    else
        return :invalid
    end
end

function ∇²projection!(::SecondOrderCone, hess, x::StaticVector, b::StaticVector)
    n = length(x)
    s = x[end]
    v = pop(x)
    bs = b[end]
    bv = pop(b)
    a =  norm(v)

    if a <= -s
        return hess .= 0
    elseif a <= s
        return hess .= 0
    elseif a > abs(s)
        dvdv = -s/norm(v)^2/norm(v)*(I - (v*v')/(v'v))*bv*v' +
            s/norm(v)*((v*(v'bv))/(v'v)^2 * 2v' - (I*(v'bv) + v*bv')/(v'v)) +
            bs/norm(v)*(I - (v*v')/(v'v))
        dvds = 1/norm(v)*(I - (v*v')/(v'v))*bv;
        dsdv = bv'/norm(v) - v'bv/norm(v)^3*v'
        dsds = 0
        hess[1:n-1,1:n-1] .= dvdv*0.5
        hess[1:n-1,n] .= dvds*0.5
        hess[n:n,1:n-1] .= dsdv*0.5
        hess[n,n] = 0
        return hess
        # return 0.5*[dvdv dvds; dsdv dsds]
    else
        throw(ErrorException("Invalid second-order cone projection"))
    end
end

function ∇projection!(::NegativeOrthant, J, x::StaticVector{n}) where n
    for i = 1:n
        J[i,i] = x <= 0 ? 1 : 0
    end
end

function ∇²projection!(::NegativeOrthant, hess, x::StaticVector, b::StaticVector)
    hess .= 0
end

Base.in(::NegativeOrthant, x::StaticVector) = all(x->x<=0, x)

"""
    AbstractConstraint

Abstract vector-valued constraint for a trajectory optimization problem.
May be either inequality or equality (specified by `sense(::AbstractConstraint)::ConstraintSense`),
and be function of single or adjacent knotpoints.

Interface:
Any constraint type must implement the following interface:
```julia
n = state_dim(::MyCon)
m = control_dim(::MyCon)
p = Base.length(::MyCon)
sense(::MyCon)::ConstraintSense
c = evaluate(::MyCon, args...)
jacobian!(∇c, ::MyCon, args...)
```

All constraints are categorized into the following type tree:
```text
                        AbstractConstraint
                        ↙                ↘
           StageConstraint               CoupledConstraint
            ↙        ↘                       ↙           ↘
StageConstraint ControlConstraint CoupledStateConstraint CoupledControlConstraint
```

The state and control dimensions (where applicable) can be queried using
`state_dim(::AbstractConstraint)` and `control_dim(::AbstractConstraint)`.
The dimensions of a constraint can be verified using [`check_dims`](@ref).
The width of the constraint Jacobian is given by [`get_inds`](@ref) or [`widths`](@ref).

The number of constraint values associated with the constraint (length of the constraint vector)
is given with `length(::AbstractConstraint)`.

# Evaluation methods
Refer to the doc strings for the following methods for more information on the required
signatures.
* [`evaluate`](@ref)
* [`jacobian!`](@ref)
* [`∇jacobian!`](@ref)
"""
abstract type AbstractConstraint end

"Only a function of states and controls at a single knotpoint"
abstract type StageConstraint <: AbstractConstraint end
"Only a function of states at a single knotpoint"
abstract type StateConstraint <: StageConstraint end
"Only a function of controls at a single knotpoint"
abstract type ControlConstraint <: StageConstraint end
"Only a function of states and controls at two adjacent knotpoints"
abstract type CoupledConstraint <: AbstractConstraint end
"Only a function of states at adjacent knotpoints"
abstract type CoupledStateConstraint <: CoupledConstraint end
"Only a function of controls at adjacent knotpoints"
abstract type CoupledControlConstraint <: CoupledConstraint end

const StateConstraints =
    Union{StageConstraint,StateConstraint,CoupledConstraint,CoupledStateConstraint}
const ControlConstraints =
    Union{StageConstraint,ControlConstraint,CoupledConstraint,CoupledControlConstraint}

"Get constraint sense (Inequality vs Equality)"
sense(::C) where {C<:AbstractConstraint} = throw(NotImplemented(:sense, Symbol(C)))

"Dimension of the state vector"
RobotDynamics.state_dim(::C) where {C<:StateConstraint} =
    throw(NotImplemented(:state_dim, Symbol(C)))

"Dimension of the control vector"
RobotDynamics.control_dim(::C) where {C<:ControlConstraint} =
    throw(NotImplemented(:control_dim, Symbol(C)))

"Return the constraint value"
evaluate(::C) where {C<:AbstractConstraint} = throw(NotImplemented(:evaluate, Symbol(C)))

"Length of constraint vector"
Base.length(::C) where {C<:AbstractConstraint} = throw(NotImplemented(:length, Symbol(C)))

# widths(con::StageConstraint, n=state_dim(con), m=control_dim(con)) = (n+m,)
# widths(con::StateConstraint, n=state_dim(con), m=0) = (n,)
# widths(con::ControlConstraint, n=0, m=control_dim(con)) = (m,)
# widths(con::CoupledConstraint, n=state_dim(con), m=control_dim(con)) = (n+m, n+m)
# widths(con::CoupledStateConstraint, n=state_dim(con), m=0) = (n,n)
# widths(con::CoupledControlConstraint, n=0, m=control_dim(con)) = (m,m)

"Upper bound of the constraint, as a vector, which is 0 for all constraints
(except bound constraints)"
@inline upper_bound(con::AbstractConstraint) =
    upper_bound(sense(con)) * @SVector ones(length(con))
@inline upper_bound(::Inequality) = 0.0
@inline upper_bound(::Equality) = 0.0

"Upper bound of the constraint, as a vector, which is 0 equality and -Inf for inequality
(except bound constraints)"
@inline lower_bound(con::AbstractConstraint) =
    lower_bound(sense(con)) * @SVector ones(length(con))
@inline lower_bound(::Inequality) = -Inf
@inline lower_bound(::Equality) = 0.0

"""
    primal_bounds!(zL, zU, con::AbstractConstraint)

Set the lower `zL` and upper `zU` bounds on the primal variables imposed by the constraint
`con`. Return whether or not the vectors `zL` or `zU` could be modified by `con`
(i.e. if the constraint `con` is a bound constraint).
"""
primal_bounds!(zL, zU, con::AbstractConstraint) = false

"Is the constraint a bound constraint or not"
@inline is_bound(con::AbstractConstraint) = false

"Check whether the constraint is consistent with the specified state and control dimensions"
@inline check_dims(con::StateConstraint, n, m) = state_dim(con) == n
@inline check_dims(con::ControlConstraint, n, m) = control_dim(con) == m
@inline check_dims(con::AbstractConstraint, n, m) =
    state_dim(con) == n && control_dim(con) == m

get_dims(con::Union{StateConstraint,CoupledStateConstraint}, nm::Int) =
    state_dim(con), nm - state_dim(con)
get_dims(con::Union{ControlConstraint,CoupledControlConstraint}, nm::Int) =
    nm - control_dim(con), control_dim(con)
get_dims(con::AbstractConstraint, nm::Int) = state_dim(con), control_dim(con)

con_label(::AbstractConstraint, i::Int) = "index $i"

"""
    get_inds(con::AbstractConstraint)

Get the indices of the joint state-control vector that are used to calculate the constraint.
If the constraint depends on more than one time step, the indices start from the beginning
of the first one.
"""
get_inds(con::StateConstraint, n, m) = (1:n,)
get_inds(con::ControlConstraint, n, m) = (n .+ (1:m),)
get_inds(con::StageConstraint, n, m) = (1:n+m,)
get_inds(con::CoupledConstraint, n, m) = (1:n+m, n+m+1:2n+2m)

"""
    widths(::AbstractConstraint)
    widths(::AbstractConstraint, n, m)

Return a tuple of the widths of the Jacobians for a constraint. If `n` and `m` are not passed
in, they are assumed to be consistent with those returned by `state_dim` and `control_dim`.
"""
@inline widths(con::AbstractConstraint, n, m) = length.(get_inds(con, n, m))
@inline widths(con::StageConstraint) = (state_dim(con) + control_dim(con),)
@inline widths(con::StateConstraint) = (state_dim(con),)
@inline widths(con::ControlConstraint) = (control_dim(con),)
@inline widths(con::CoupledConstraint) =
    (state_dim(con) + control_dim(con), state_dim(con) + control_dim(con))
@inline widths(con::CoupledStateConstraint) = (state_dim(con), state_dim(con))
@inline widths(con::CoupledControlConstraint) = (control_dim(con), control_dim(con))

############################################################################################
# 								EVALUATION METHODS 										   #
############################################################################################
"""
    evaluate!(vals, con::AbstractConstraint, Z, [inds])

Evaluate constraints for entire trajectory. This is the most general method used to evaluate
constraints along the trajectory `Z`, and should be the one used in other functions.
The `inds` argument determines at which knot points the constraint is evaluated.

If `con` is a `StageConstraint`, this will call `evaluate(con, z)` by default, or
`evaluate(con, z1, z2)` if `con` is a `CoupledConstraint`.
"""
function evaluate!(
    vals::Vector{<:AbstractVector},
    con::StageConstraint,
    Z::AbstractTrajectory,
    inds = 1:length(Z),
)
    for (i, k) in enumerate(inds)
        vals[i] .= evaluate(con, Z[k])
    end
end

function evaluate!(
    vals::Vector{<:AbstractVector},
    con::CoupledConstraint,
    Z::AbstractTrajectory,
    inds = 1:length(Z)-1,
)
    for (i, k) in enumerate(inds)
        vals[i] .= evaluate(con, Z[k], Z[k+1])
    end
end

"""
    jacobian!(∇c, con::AbstractConstraint, Z, [inds, is_const])

Evaluate constraints for entire trajectory. This is the most general method used to evaluate
constraints along the trajectory `Z`, and should be the one used in other functions.
The `inds` argument determines at which knot points the constraint is evaluated.
The optional `is_const` argument is a `BitArray` of the same size as `∇c`, and captures
the output of `jacobian!`, which should return a Boolean specifying if the Jacobian is
constant or not.

The values are stored in `∇c`, which should be a matrix of matrices. If `con` is a
`StageConstraint`, `size(∇c,2) = 1`, and `size(∇c,2) = 2` if `con` is a `CoupledConstraint`.

If `con` is a `StageConstraint`, this will call `jacobian!(∇c, con, z)` by default, or
`jacobian!(∇c, con, z1, z2, i)` if `con` is a `CoupledConstraint`.
"""
function jacobian!(
    ∇c::VecOrMat{<:AbstractMatrix},
    con::StageConstraint,
    Z::AbstractTrajectory,
    inds = 1:length(Z),
    is_const = BitArray(undef, size(∇c))
)
    for (i, k) in enumerate(inds)
        is_const[i] = jacobian!(∇c[i], con, Z[k])
    end
end

function jacobian!(
    ∇c::VecOrMat{<:AbstractMatrix},
    con::CoupledConstraint,
    Z::AbstractTrajectory,
    inds = 1:size(∇c, 1),
    is_const = BitArray(undef, size(∇c))
)
    for (i, k) in enumerate(inds)
        is_const[i,1] = jacobian!(∇c[i, 1], con, Z[k], Z[k+1], 1)
        is_const[i,2] = jacobian!(∇c[i, 2], con, Z[k], Z[k+1], 2)
    end
end

"""
    ∇jacobian!(G, con::AbstractConstraint, Z, λ, inds, is_const, init)
    ∇jacobian!(G, con::AbstractConstraint, Z::AbstractKnotPoint, λ::AbstractVector)

Evaluate the second-order expansion of the constraint `con` along the trajectory `Z`
after multiplying by the lagrange multiplier `λ`.
The optional `is_const` argument is a `BitArray` of the same size as `∇c`, and captures
the output of `jacobian!`, which should return a Boolean specifying if the Jacobian is
constant or not. The `init` flag will force re-calculation of constant Jacobians when true.

The method for each constraint should calculate the Jacobian of the vector-Jacobian product,
    and therefore should be of size n × n if the input dimension is n.

Importantly, this method should ADD and not overwrite the contents of `G`, since this term
is dependent upon all the constraints acting at that time step.
"""
function ∇jacobian!(
    G::VecOrMat{<:AbstractMatrix},
    con::StageConstraint,
    Z::AbstractTrajectory,
    λ::Vector{<:AbstractVector},
    inds = 1:length(Z),
    is_const = ones(Bool, length(inds)),
    init::Bool = false,
)
    for (i, k) in enumerate(inds)
        if init || !is_const[i]
            is_const[i] = ∇jacobian!(G[i], con, Z[k], λ[i])
        end
    end
end

function ∇jacobian!(
    G::VecOrMat{<:AbstractMatrix},
    con::CoupledConstraint,
    Z::AbstractTrajectory,
    λ::Vector{<:AbstractVector},
    inds = 1:length(Z),
    is_const = ones(Bool, length(inds)),
    init::Bool = false,
)
    for (i, k) in enumerate(inds)
        if init || !is_const[i]
            is_const[i] = ∇jacobian!(G[i, 1], con, Z[k], Z[k+1], λ[i], 1)
            is_const[i] = ∇jacobian!(G[i, 2], con, Z[k], Z[k+1], λ[i], 2)
        end
    end
end

# Default methods for converting KnotPoints to states and controls for StageConstraints
"""
    evaluate(con::AbstractConstraint, z)       # stage constraint
    evaluate(con::AbstractConstraint, z1, z2)  # coupled constraint

Evaluate the constraint `con` at knot point `z`. By default, this method will attempt to call

    evaluate(con, x)

if `con` is a `StateConstraint`,

    evaluate(con, u)

if `con` is a `ControlConstraint`, or

    evaluate(con, x, u)

if `con` is a `StageConstraint`. If `con` is a `CoupledConstraint` the constraint should
define

    evaluate(con, z1, z2)

"""
@inline evaluate(con::StateConstraint, z::AbstractKnotPoint) = evaluate(con, state(z))
@inline evaluate(con::ControlConstraint, z::AbstractKnotPoint) = evaluate(con, control(z))
@inline evaluate(con::StageConstraint, z::AbstractKnotPoint) =
    evaluate(con, state(z), control(z))

"""
    jacobian!(∇c, con::AbstractConstraint, z, i=1)       # stage constraint
    jacobian!(∇c, con::AbstractConstraint, z1, z2, i=1)  # coupled constraint

Evaluate the constraint `con` at knot point `z`. By default, this method will attempt to call

    jacobian!(∇c, con, x)

if `con` is a `StateConstraint`,

    jacobian!(∇c, con, u)

if `con` is a `ControlConstraint`, or

    jacobian!(∇c, con, x, u)

if `con` is a `StageConstraint`. If `con` is a `CoupledConstraint` the constraint should
define

    jacobian!(∇c, con, z, i)

where `i` determines which Jacobian should be evaluated. E.g. if `i = 1`, the Jacobian
with respect to the first knot point's stage and controls is calculated.

# Automatic Differentiation
If `con` is a `StateConstraint` or `ControlConstraint` then this method is automatically
defined using ForwardDiff.
"""
jacobian!(∇c, con::StateConstraint, z::AbstractKnotPoint, i = 1) =
    jacobian!(∇c, con, state(z))
jacobian!(∇c, con::ControlConstraint, z::AbstractKnotPoint, i = 1) =
    jacobian!(∇c, con, control(z))
jacobian!(∇c, con::StageConstraint, z::AbstractKnotPoint, i = 1) =
    jacobian!(∇c, con, state(z), control(z))

# ForwardDiff jacobians that are of only state or control
function jacobian!(∇c, con::StageConstraint, x::StaticVector)
    eval_c(x) = evaluate(con, x)
    ∇c .= ForwardDiff.jacobian(eval_c, x)
    return false
end

@inline ∇jacobian!(G, con::StateConstraint, z::AbstractKnotPoint, λ, i = 1) =
    ∇jacobian!(G, con, state(z), λ)
@inline ∇jacobian!(G, con::ControlConstraint, z::AbstractKnotPoint, λ, i = 1) =
    ∇jacobian!(G, con, control(z), λ)
@inline ∇jacobian!(G, con::StageConstraint, z::AbstractKnotPoint, λ, i = 1) =
    ∇jacobian!(G, con, state(z), control(z), λ)

function ∇jacobian!(G, con::StageConstraint, x::StaticVector, λ)
    eval_c(x) = evaluate(con, x)'λ
    G_ = ForwardDiff.hessian(eval_c, x)
    G .+= G_
    return false
end

function ∇jacobian!(
    G,
    con::StageConstraint,
    x::StaticVector{n},
    u::StaticVector{m},
    λ,
) where {n,m}
    ix = SVector{n}(1:n)
    iu = SVector{m}(n .+ (1:m))
    eval_c(z) = evaluate(con, z[ix], z[iu])'λ
    G .+= ForwardDiff.hessian(eval_c, [x; u])
    return false
end

function gen_jacobian(con::AbstractConstraint, i = 1)
    ws = widths(con)
    p = length(con)
    C1 = SizedMatrix{p,ws[i]}(zeros(p, ws[i]))
end

function gen_views(∇c::AbstractMatrix, con::StateConstraint, n = state_dim(con), m = 0)
    view(∇c, :, 1:n), view(∇c, :, n:n-1)
end

function gen_views(∇c::AbstractMatrix, con::ControlConstraint, n = 0, m = control_dim(con))
    view(∇c, :, 1:0), view(∇c, :, 1:m)
end

function gen_views(
    ∇c::AbstractMatrix,
    con::AbstractConstraint,
    n = state_dim(con),
    m = control_dim(con),
)
    if size(∇c, 2) < n + m
        view(∇c, :, 1:n), view(∇c, :, n:n-1)
    else
        view(∇c, :, 1:n), view(∇c, :, n .+ (1:m))
    end
end


using ExBroadcast
using Test
using OffsetArrays, CUDA, StructArrays, Interpolations, Adapt
@testset "@tab.jl" begin
    a = randn(1000)
    b, c = @tab sincos.(a)
    @test b == sin.(a)
    @test c == cos.(a)
    b′ = @view b[1:999]
    @test_throws ArgumentError b′, c .= sincos.(a)
    b″ = OffsetArray(b,-1)
    @test_throws ArgumentError b″, c .= sincos.(a)
end

@testset "@mtb for OffsetVector" begin
    a = randn(1000)
    aᵒ = OffsetArray(a,-1)
    @mtb b = parent(aᵒ .+ 1)
    @test (a .+ 1) == b
end

@testset "@mtb for small array" begin
    a = [1]
    @mtb 2 a .+= 1
    @test a[1] == 2
end

@testset "@mtb for BitArray" begin
    a = randn(4096 * 6 - 2048)
    b = a .> 0
    @test @mtb (a .> 0) == b
end

@testset "CUDA supports" begin
    a = CUDA.randn(1000)
    b, c = @tab sincos.(a)
    @test b ≈ sin.(a)
    @test c ≈ cos.(a)
    d = StructArray{ComplexF32}((b, c))
    @test abs.(d) isa CuVector
    e = OffsetArray(d, -1) .+ 1
    @test e isa StructArray
    @test eachindex(e) == 0:999
end

@testset "Lazy test" begin
    a = ((i,k...) for i in 1:10, k in Iterators.product(1:10,1:2))
    @test identity.(a) == identity.(Lazy(a))
    a = (randn() for i in 1:100)
    b = Lazy(a) .+ [1 2] .- [1 2]
    @test b[:,1] != b[:,2]
    a = (randn() for i in 1:100, j in tuple())
    @test identity.(a) == identity.(Lazy(a))
end

@testset "CUDA Interp test" begin
    a = randn(Float32,101,101)
    itp = CubicSplineInterpolation((-50f0:50f0,-50f0:50f0), a)
    cuitp = adapt(CuArray, itp)
    res = cuitp.(-10:10,-10:10)
    @test res isa CuArray
    @test collect(res) ≈ itp.(-10:10,-10:10)

    itp = itp.itp
    a,b,c = preweight(itp,-1:0.01f0:1,-1:0.01f0:1)
    @test a.(b,c) == itp.(-1:0.01f0:1,-1:0.01f0:1)
end