using Isaac
using Isaac.TestSets
using Isaac.TestSets: precond

println("Testing nsolimod for index-1 saddles")
@testset "nsolidmod-index1" begin

#[1] Muller Test
xe = []
for init in (:near, :far)
   println("Testing nsolimod, Muller, init = $(init)")
   V = MullerPotential()
   x0, v0 = ic_dimer(V, init)
   E, dE = objective(V)
   x, ndE = nsolimod(dE, x0, 1; V0=v0, maxnumdE = 100, maxstep = 0.5)
   println("   num_dE = ", ndE)
   if init == :near
      xe = copy(x)
   else
      println("  |x - xe| = ", norm(x - xe, Inf))
      @test norm(x - xe, Inf) < 1e-8
   end
end

# [2] Vacancy Test
V = LJVacancy2D(bc = :clamped)
E, dE = objective(V)
x0, v0 = ic_dimer(V, :near)
P = precond(V, x0)
Iprep = (P, x) -> P
Pprep = (P, x) -> Isaac.TestSets.precond(V, x)
xe = []

for (init, precond, precon_prep) in
      [(:near, I, Iprep), (:near, P, Pprep), (:far, I, Iprep), (:far, P, Pprep)]

   x0, v0 = ic_dimer(V, init)
   pre = !(precond == I)

   println("Testing nsolimod, 2D LJ Vacancy, init = $init, precond = $pre")
   x, ndE = nsolimod(dE, x0, 1; V0=v0, P = precond, precon_prep=precon_prep)
   println("   num_dE = ", ndE)
   @test norm(dE(x), Inf) < 1e-5
   println("   |dE| = ", norm(dE(x), Inf))
   if xe == []
      # to make sure all 4 methods converge to the same saddle,
      # store the first one ...
      xe = x
      H2 = TestSets.hessian(V, x)
      σ = sort(eigvals(Symmetric(H2)))
      @test σ[1] * σ[2] < 0
      println("   min-eigs = ", σ[1:2])
   else
      # ... and then compare the others
      @test norm(x - xe, Inf) < 1e-6
      println("   |x - xe| = ", norm(x - xe, Inf))
   end
end

end  # @testset "nsolidstab-index1"


using CompScienceMeshes
using Test

for T in [Float32, Float64]
    faces = readmesh(joinpath(@__DIR__, "assets/rectangle32.in"), T=T)
    edges = skeleton(faces,1)
    verts = skeleton(faces,0)
    bnd = boundary(faces)

    int_edges = submesh(interior_tpredicate(faces), edges)
    # pred = interior_tpredicate(int_edges)

    bnd_verts = skeleton(bnd,0)
    function pred(node)
        !(node in cells(bnd_verts))
    end
    int_verts = submesh(pred, verts)

    @test numcells(int_verts) == 9
    @test numcells(edges) - numcells(int_edges) == 16
end


commands = readlines("day2/commands.txt")

function dive(commands)
    h, d = 0, 0
    for command in commands
        dir, x = split(command)
        x = parse(Int, x)
        if dir == "forward"
            h += x
        else
            d += dir == "down" ? x : -x
        end
    end
    h, d
end

h, d = dive(commands)
p1 = h * d

@show p1

function dive2(commands)
    h, d, a = 0, 0, 0
    for command in commands
        dir, x = split(command)
        x = parse(Int, x)
        if dir == "forward"
            h += x
            d += x * a
        else
            a += dir == "down" ? x : -x
        end
    end
    h, d, a
end

h, d, _ = dive2(commands)
p2 = h * d

@show p2


function fib_seeder{T<:Integer}(n::T)
    a = zeros(BigInt, n)
    a[1] = one(BigInt)
    return a
end

function fib{T<:Integer}(n::T, k::T)
    NFib(n, k, fib_seeder)
end


module DLM

export influence_matrix, influence_matrix!
export pressure_coefficients
export AbstractKernel, ParabolicKernel, QuarticKernel

# coefficients for 12-term approximation of kernel integrals
const a = (0.000319759140, -0.000055461471, 0.002726074362, 0.005749551566,
    0.031455895072, 0.106031126212, 0.406838011567, 0.798112357155,
    -0.417749229098, 0.077480713894, -0.012677284771, 0.001787032960)
const b = 0.009054814793

# kernel approximations
"""
    AbstractKernel()

Abstract type for approximation of the planar and nonplanar influence terms in
the numerator of the doublet lattice method kernel
"""
abstract type AbstractKernel end

"""
    ParabolicKernel()

Parabolic approximation for the planar and nonplanar influence terms in the
numerator of the doublet lattice method kernel.
"""
struct ParabolicKernel <: AbstractKernel end

"""
    QuarticKernel()

Quartic approximation for the planar and nonplanar influence terms in the
numerator of the doublet lattice method kernel.
"""
struct QuarticKernel <: AbstractKernel end

"""
    approximate_kernel_integrals(u, k)

Solves for the integrals I0 and J0 using a 12-term exponential approximation
given by Desmarais in "AN ACCURATE METHOD FOR EVALUATING THE KERNEL OF THE
INTEGRAL EQUATION RELATING LIFT TO DOWNWASH IN UNSTEADY POTENTIAL FLOW".

# Arguments
- u = (M*R - x0)/(β^2*x0)
- k = ω*r/U
"""
function approximate_kernel_integrals(u, k)

    TF = promote_type(typeof(u), typeof(k))

    I0 = zero(complex(TF))
    J0 = zero(complex(TF))
    @inbounds for j = 1:12
        p = b*2^j
        pk = complex(p, k)
        tmp = a[j]*exp(-pk*u)/pk
        I0 += tmp
        J0 += tmp*(u+1/pk)
    end

    return I0, J0
end

"""
    kernel_coefficients(x0, y0, z0, ω, U, M)

Solves for the kernel coefficients K1 and K2 and their steady state counterparts
K10 and K20.

See "A Doublet-Lattice Method for Calculating Lift Distributions
on Oscillating Surfaces in Subsonic Flows" by Edward Albano and William P Rodden.

# Arguments
- x0 = x-ξ
- y0 = y-η
- z0 = z-ζ
- ω = frequency
- U = Velocity
- M = Mach Number
"""
function kernel_coefficients(x0, y0, z0, ω, U, M)

    r = sqrt(y0^2 + z0^2)
    β = sqrt(1 - M^2)
    R = sqrt(x0^2+β^2*r^2)

    u = (M*R-x0)/(β^2*max(r, eps(typeof(r))))
    k = ω*r/U

    # use the symmetry properties of the integrand if u < 0
    if u < 0
        I0, J0 = approximate_kernel_integrals(0, k)
        I10 = 1 - 1im*k*I0
        I20 = 2 - 1im*k*I0 + k^2*J0
        u1pos = -u
        I0, J0 = approximate_kernel_integrals(u1pos, k)
        invR = 1/R
        invsqrt = 1/sqrt(1+u1pos^2)
        expk = exp(-1im*k*u1pos)
        I11 = (1-u1pos*invsqrt)*expk - 1im*k*I0
        I21 = ((2 + 1im*k*u1pos)*(1 - u1pos*invsqrt) -
            u1pos*invsqrt^3)*expk - 1im*k*I0 + J0*k^2
        I1 = complex(2*real(I10) - real(I11), imag(I11))
        I2 = complex(2*real(I20) - real(I21), imag(I21))
        expk = exp(-1im*k*u)
    else
        I0, J0 = approximate_kernel_integrals(u, k)
        invR = 1/R
        invsqrt = 1/sqrt(1+u^2)
        expk = exp(-1im*k*u)
        I1 = (1-u*invsqrt)*expk - 1im*k*I0
        I2 = ((2 + 1im*k*u)*(1 - u*invsqrt) -
            u*invsqrt^3)*expk - 1im*k*I0 + J0*k^2
    end

    K1 = I1 + M*r*invR*invsqrt*expk
    K2 = -I2 - 1im*k*invsqrt*expk*(M*r*invR)^2 -
       (M*r*invR)*((1 + u^2)*(β*r*invR)^2 +
       2 + M*r*u*invR)*expk*invsqrt^3

    K10 =  1 + x0/R
    K20 = -2 - x0/R*(2+β^2*r^2/R^2)

    return K1, K2, K10, K20
end

"""
    kernel_numerator(ω, U, M, x0, y0, z0, γr, γs)

Solves for the numerator terms of the kernel and subtracts out the steady-state
components of these terms.

See "Further Refinement of the Subsonic Doublet-Lattice Method" by William P. Rodden,
Paul F. Taylor and Samuel C. McIntosh Jr.
"""
function kernel_numerator(ω, U, M, x0, y0, z0, cγr, sγr, cγs, sγs)
    T1, T2 = cγr*cγs + sγr*sγs, (z0*cγr-y0*sγr)*(z0*cγs-y0*sγs)
    K1, K2, K10, K20 = kernel_coefficients(x0, y0, z0, ω, U, M)
    expn = exp(complex(0, -ω*x0/U))
    return (K1*expn-K10)*T1, (K2*expn-K20)*T2
end

"""
influence_coefficient(ω, U, M, (xr, yr, zr), (xi, yi, zi), (xo, yo, zo),
    chord, λ, γr, γs, kernel)

Solves for the influence coefficient of the sending panel on the receiving panel.

The planar and nonplanar influence terms in the numerator are approximated using
the approximation specified by `kernel`.

See "Further Refinement of the Subsonic Doublet-Lattice Method" by William P. Rodden,
Paul F. Taylor and Samuel C. McIntosh Jr.
"""
influence_coefficient

function influence_coefficient(ω, U, M, (xr, yr, zr),
    (xi, yi, zi), (xo, yo, zo), chord, cλ, cγr, sγr, cγs, sγs, ::ParabolicKernel)

    dx = xo-xi
    dy = yo-yi
    dz = zo-zi
    l = sqrt(dx^2 + dy^2 + dz^2)

    xm = (xo+xi)/2
    ym = (yo+yi)/2
    zm = (zo+zi)/2

    # Approximate numerator coefficients as parabolas
    x0i = xr - xi
    y0i = yr - yi
    z0i = zr - zi
    Ki1, Ki2 = kernel_numerator(ω, U, M, x0i, y0i, z0i, cγr, sγr, cγs, sγs)

    x0m = xr - xm
    y0m = yr - ym
    z0m = zr - zm

    Km1, Km2 = kernel_numerator(ω, U, M, x0m, y0m, z0m, cγr, sγr, cγs, sγs)

    x0o = xr - xo
    y0o = yr - yo
    z0o = zr - zo
    Ko1, Ko2 = kernel_numerator(ω, U, M, x0o, y0o, z0o, cγr, sγr, cγs, sγs)

    e = 1/2*l*cλ
    η = y0m*cγs + z0m*sγs
    ζ = -y0m*sγs + z0m*cγs

    if abs(ζ) == zero(ζ)
        F = 2*e/(η^2-e^2)
    else
        F =atan(2*e*abs(ζ),(η^2+ζ^2-e^2))/abs(ζ)
    end

    # steady term: from VLM
    d0 = steady_influence_coefficient(M, x0i, y0i, z0i, x0o, y0o, z0o, sγr, cγr)

    # planar term: ∫(A1*y^2 + B1*y + C1)/((η-y)^2+ζ^2)
    A1 = (Ki1 - 2*Km1 + Ko1)/(2*e^2)
    B1 = (Ko1 - Ki1)/(2*e)
    C1 = Km1

    d1 = ((η^2-ζ^2)*A1 + η*B1 + C1) * F +
        (1/2*B1 + η*A1) * log(((η-e)^2+ζ^2)/((η+e)^2 + ζ^2)) + 2*e*A1

    # nonplanar term: ∫(A1*y^2 + B1*y + C1)/((η-y)^2+ζ^2)^2
    if abs(ζ) == zero(ζ)
        D2 = zero(complex(typeof(d1)))
    else
        A2 = (Ki2 - 2*Km2 + Ko2)/(2*e^2)
        B2 = (Ko2 - Ki2)/(2*e)
        C2 = Km2

        α = (e/ζ)^2 * (1 - (η^2 + ζ^2 - e^2)/(2*e)*F)

        D2 = e/(η^2 + ζ^2 - e^2)*((
            (2*(η^2 + ζ^2 + e^2)*(e^2*A2 + C2) + 4*η*e^2*B2))/
            (((η + e)^2 + ζ^2)*((η - e)^2 + ζ^2)) -
             (α/e^2)*((η^2 + ζ^2)*A2 + η*B2 + C2))
    end

    return chord/(8*pi)*(d0 + d1 + D2)
end

function influence_coefficient(ω, U, M, (xr, yr, zr),
    (xi, yi, zi), (xo, yo, zo), chord, cλ, cγr, sγr, cγs, sγs, ::QuarticKernel)

    dx = xo-xi
    dy = yo-yi
    dz = zo-zi
    l = sqrt(dx^2 + dy^2 + dz^2)

    xm = (xi+xo)/2
    ym = (yi+yo)/2
    zm = (zi+zo)/2

    xmi = (xi+xm)/2
    ymi = (yi+ym)/2
    zmi = (zi+zm)/2

    xmo = (xm+xo)/2
    ymo = (ym+yo)/2
    zmo = (zm+zo)/2

    # Approximate numerator coefficients as quartic
    x0i = xr - xi
    y0i = yr - yi
    z0i = zr - zi
    Ki1, Ki2 = kernel_numerator(ω, U, M, x0i, y0i, z0i, cγr, sγr, cγs, sγs)

    x0mi = xr - xmi
    y0mi = yr - ymi
    z0mi = zr - zmi

    Kmi1, Kmi2 = kernel_numerator(ω, U, M, x0mi, y0mi, z0mi, cγr, sγr, cγs, sγs)

    x0m = xr - xm
    y0m = yr - ym
    z0m = zr - zm

    Km1, Km2 = kernel_numerator(ω, U, M, x0m, y0m, z0m, cγr, sγr, cγs, sγs)

    x0mo = xr - xmo
    y0mo = yr - ymo
    z0mo = zr - zmo

    Kmo1, Kmo2 = kernel_numerator(ω, U, M, x0mo, y0mo, z0mo, cγr, sγr, cγs, sγs)

    x0o = xr - xo
    y0o = yr - yo
    z0o = zr - zo
    Ko1, Ko2 = kernel_numerator(ω, U, M, x0o, y0o, z0o, cγr, sγr, cγs, sγs)

    e = 1/2*l*cλ
    η = y0m*cγs + z0m*sγs
    ζ = -y0m*sγs + z0m*cγs

    if abs(ζ) == zero(ζ)
        F = 2*e/(η^2-e^2)
    else
        F =atan(2*e*abs(ζ), (η^2+ζ^2-e^2))/abs(ζ)
    end

    # steady term: from VLM
    d0 = steady_influence_coefficient(M, x0i, y0i, z0i, x0o, y0o, z0o, sγr, cγr)

    # planar term: ∫(A1*y^4 + B1*y^3 + C1*y^2 + D1*y + E1)/((η-y)^2+ζ^2)
    A1 = -(1/(6*e^2))*(Ki1 - 16*Kmi1 + 30*Km1 - 16*Kmo1 + Ko1)
    B1 =  (1/(6*e  ))*(Ki1 -  8*Kmi1          +  8*Kmo1 - Ko1)
    C1 = Km1
    D1 = -(2/(3*e^3))*(Ki1 -  2*Kmi1          +  2*Kmo1 - Ko1)
    E1 =  (2/(3*e^4))*(Ki1 -  4*Kmi1 +  6*Km1 -  4*Kmo1 + Ko1)

    d1 = ((η^2-ζ^2)*A1 + η*B1 + C1 + η*(η^2 - 3*ζ^2)*D1 +
        (η^4 - 6*η^2*ζ^2 + ζ^4)*E1)*F + (η*A1 + 1/2*B1 +
        1/2*(3*η^2 - ζ^2)*D1 + 2*η*(η^2-ζ^2)*E1)*log(((η - e)^2 + ζ^2)/((η + e)^2 + ζ^2)) +
        2*e*(A1 + 2*η*D1 + (3*η^2 - ζ^2 + 1/3*e^2)*E1)

    # nonplanar term: ∫(A2*y^4 + B2*y^3 + C2*y^2 + D2*y + E2)/((η-y)^2+ζ^2)^2
    if abs(ζ) == zero(ζ)
        d2 = zero(complex(typeof(d1)))
    else
        A2 = -(1/(6*e^2))*(Ki2 - 16*Kmi2 + 30*Km2 - 16*Kmo2 + Ko2)
        B2 =  (1/(6*e  ))*(Ki2 -  8*Kmi2          +  8*Kmo2 - Ko2)
        C2 = Km2
        D2 = -(2/(3*e^3))*(Ki2 -  2*Kmi2          +  2*Kmo2 - Ko2)
        E2 =  (2/(3*e^4))*(Ki2 -  4*Kmi2 +  6*Km2 -  4*Kmo2 + Ko2)

        d2 = 1/(2*ζ^2)*(((η^2 + ζ^2)*A2 + η*B2 + C2 + η*(η^2 + 3*ζ^2)*D2 +
            (η^4 + 6*η^2*ζ^2 - 3*ζ^4)*E2)*F + 1/((η + e)^2 + ζ^2)*(((η^2 + ζ^2)*η +
            (η^2 - ζ^2)*e)*A2 + (η^2 + ζ^2 + η*e)*B2 + (η + e)*C2 +
            (η^4 - ζ^4 + (η^2 - 3*ζ^2)*η*e)*D2 + ((η^4 - 2*η^2*ζ^2 - 3*ζ^4)*η +
            (η^4 - 6*η^2*ζ^2 + ζ^4)*e)*E2) - 1/((η - e)^2 + ζ^2)*(((η^2 + ζ^2)*η -
            (η^2 - ζ^2)*e)*A2 + (η^2 + ζ^2 - η*e)*B2 + (η - e)*C2 +
            (η^4 - ζ^4 - (η^2 - 3*ζ^2)*η*e)*D2 + ((η^4 - 2*η^2*ζ^2 - 3*ζ^4)*η -
            (η^4 - 6*η^2*ζ^2 + ζ^4)*e)*E2) + (ζ^2*log(((η - e)^2 + ζ^2)/((η+e)^2 + ζ^2)))*D2 +
            4*ζ^2*(e + η*log(((η - e)^2 + ζ^2)/((η + e)^2 + ζ^2)))*E2)
    end

    return chord/(8*pi)*(d0 + d1 + d2)
end

"""
    steady_influence_coefficient(M, x0i, y0i, z0i, x0o, y0o, z0o, sγr, cγr)

Solves for the steady component of the influence of the sending panel on the
receiving panel using the vortex lattice method.
"""
function steady_influence_coefficient(M, x0i, y0i, z0i, x0o, y0o, z0o, sγr, cγr)

    # steady term (from VLM)
    ax = x0i/sqrt(1-M^2)
    ay = y0i
    az = z0i
    anrm = sqrt(ax^2 + ay^2 + az^2)
    ainv = 1/(anrm*(anrm - ax))

    bx = x0o/sqrt(1-M^2)
    by = y0o
    bz = z0o
    bnrm = sqrt(bx^2 + by^2 + bz^2)
    binv = 1/(bnrm*(bnrm - bx))

    vy =  az*ainv - bz*binv
    vz = -ay*ainv + by*binv

    ainv = 1/(anrm*bnrm*(anrm*bnrm + ax*bx + ay*by + az*bz))
    vy = vy + (az*bx - ax*bz)*(anrm + bnrm)*ainv
    vz = vz + (ax*by - ay*bx)*(anrm + bnrm)*ainv

    return -(sγr*vy - cγr*vz)
end

"""
    influence_matrix(ω, U, M, xyz::AbstractVector{<:AbstractArray{<:Number, 3}}, symmetric::AbstractVector; kernel=QuarticKernel())
    influence_matrix(ω, U, M, xyz::AbstractArray{<:Number, 3}, symmetric; kernel=QuarticKernel())
    influence_matrix(ω, U, M, xyzr::AbstractArray{<:Number, 3}, xyzs::AbstractArray{<:Number, 3}, symmetric; kernel=QuarticKernel())

Fills in aerodynamic influence coefficient matrix. Returns aerodynamic influence
coefficient matrix AIC with shape (nir, njr, nis, njs). For a non-allocating version,
see `influence_matrix!`.

Receiving and sending panel points must be aligned with the x-direction. The
freestream velocity must also be aligned in the x-direction.

# Arguments
- `ω`: oscillation frequency
- `U`: freestream velocity
- `M`: Mach number (for Prandtl-Glauert compressibility correction)
- `xyzr`: definition of receiving panel points with shape: (3, nir, njr)
- `xyzs`: definition of sending panel points with shape: (3, nis, njs)
- `symmetric`: flag indicating symmetry of each sending panel
- `kernel`: indicates which approximation of the kernel should be used
"""
influence_matrix

function influence_matrix(ω, U, M, xyz::AbstractVector{<:AbstractArray{<:Number, 3}},
        symmetric::AbstractVector; kernel=QuarticKernel())

    nx = sum(size.(xyz, 2))
    ny = sum(size.(xyz, 3))

    TF = promote_type(typeof(ω), typeof(U), typeof(M), eltype(eltype(xyz)))

    AIC = zeros(complex(TF), nx, ny, nx, ny)

    return influence_matrix!(AIC, ω, U, M, xyz, symmetric, kernel=kernel)
end

function influence_matrix(ω, U, M, xyzr::AbstractArray{<:Number, 3},
    xyzs::AbstractArray{<:Number, 3}, symmetric; kernel=QuarticKernel())

    nir = size(xyzr, 2)-1
    njr = size(xyzr, 3)-1
    nis = size(xyzs, 2)-1
    njs = size(xyzs, 3)-1

    TF = promote_type(typeof(ω), typeof(U), typeof(M), eltype(xyzr), eltype(xyzs))

    AIC = zeros(complex(TF), nir, njr, nis, njs)

    return influence_matrix!(AIC, ω, U, M, xyzr, xyzs, symmetric, kernel=kernel)
end

influence_matrix(ω, U, M, xyz::AbstractArray{<:Number, 3}, symmetric; kernel=QuarticKernel()) = influence_matrix(ω, U, M, xyz, xyz, symmetric, kernel=kernel)

"""
    influence_matrix!(AIC, ω, U, M, xyz::AbstractVector{<:AbstractArray{<:Number, 3}}, symmetric::AbstractVector; kernel=QuarticKernel())
    influence_matrix!(AIC, ω, U, M, xyz::AbstractArray{<:Number, 3}, symmetric; kernel=QuarticKernel())
    influence_matrix!(AIC, ω, U, M, xyzr::AbstractArray{<:Number, 3}, xyzs::AbstractArray{<:Number, 3}, symmetric; kernel=QuarticKernel())

Non-allocating version of [`influence_matrix`](@ref)
"""
influence_matrix!

function influence_matrix!(AIC, ω, U, M, xyz::AbstractVector{<:AbstractArray{<:Number, 3}}, symmetric::AbstractVector; kernel=QuarticKernel())

    nx = size(AIC, 1)
    ny = size(AIC, 2)
    nsurf = length(xyz)

    ir = 0
    jr = 0
    is = 0
    js = 0
    @inbounds for kr = 1:nsurf, ks = 1:nsurf
        nir = size(xyz[ir], 2)
        njr = size(xyz[ir], 3)
        nis = size(xyz[is], 2)
        njs = size(xyz[is], 3)

        ixr = ir+1:ir+nir
        iyr = jr+1:jr+njr
        ixs = is+1:is+nis
        iys = js+1:js+njs

        influence_matrix!(view(AIC, ixr, iyr, ixs, iys), ω, U, M, xyz[kr], xyz[ks], symmetric[ks], kernel=kernel)
    end

    return AIC
end

function influence_matrix!(AIC, ω, U, M, xyzr::AbstractArray{<:Number, 3}, xyzs::AbstractArray{<:Number, 3}, symmetric; kernel=QuarticKernel())
    nir = size(xyzr, 2)-1
    njr = size(xyzr, 3)-1
    nis = size(xyzs, 2)-1
    njs = size(xyzs, 3)-1

    @assert size(AIC) == (nir, njr, nis, njs)

    @inbounds for jr=1:njr, ir = 1:nir, js=1:njs, is=1:nis
        # receiving point at 3/4 chord
        xri = (1/4)*xyzr[1, ir, jr] + (3/4)*xyzr[1, ir+1, jr]
        yri = (1/4)*xyzr[2, ir, jr] + (3/4)*xyzr[2, ir+1, jr]
        zri = (1/4)*xyzr[3, ir, jr] + (3/4)*xyzr[3, ir+1, jr]
        xro = (1/4)*xyzr[1, ir, jr+1] + (3/4)*xyzr[1, ir+1, jr+1]
        yro = (1/4)*xyzr[2, ir, jr+1] + (3/4)*xyzr[2, ir+1, jr+1]
        zro = (1/4)*xyzr[3, ir, jr+1] + (3/4)*xyzr[3, ir+1, jr+1]
        xr = (xri+xro)/2
        yr = (yri+yro)/2
        zr = (zri+zro)/2
        # sending line at 1/4 chord
        xsi = (3/4)*xyzs[1, is, js] + (1/4)*xyzs[1, is+1, js]
        ysi = (3/4)*xyzs[2, is, js] + (1/4)*xyzs[2, is+1, js]
        zsi = (3/4)*xyzs[3, is, js] + (1/4)*xyzs[3, is+1, js]
        xso = (3/4)*xyzs[1, is, js+1] + (1/4)*xyzs[1, is+1, js+1]
        yso = (3/4)*xyzs[2, is, js+1] + (1/4)*xyzs[2, is+1, js+1]
        zso = (3/4)*xyzs[3, is, js+1] + (1/4)*xyzs[3, is+1, js+1]
        # average chord length of sending panel
        chord = (xyzs[1, is+1, js] - xyzs[1, is, js])/2 + (xyzs[1, is+1, js+1] - xyzs[1, is, js+1])/2
        # cos(sweep) of sending doublet line
        cλs = (yso-ysi)/(sqrt((xso-xsi)^2 + (yso-ysi)^2))
        # cos(dihedral) and sin(dihedral) of receiving point
        tmp = sqrt((yro-yri)^2+(zro-zri)^2)
        cγr = (yro-yri)/tmp
        sγr = (zro-zri)/tmp
        # cos(dihedral) and sin(dihedral) of sending line
        tmp = sqrt((yso-ysi)^2+(zso-zsi)^2)
        cγs = (yso-ysi)/tmp
        sγs = (zso-zsi)/tmp
        # compute influence of the sending panel on the receiving panel
        AIC[ir, jr, is, js] = influence_coefficient(ω, U, M, (xr, yr, zr),
            (xsi, ysi, zsi), (xso, yso, zso), chord, cλs, cγr, sγr, cγs, sγs, kernel)
        # add influence of reflected sending panel, if applicable
        if symmetric
            AIC[ir, jr, is, js] += influence_coefficient(ω, U, M, (xr, yr, zr),
                (xso, -yso, zso), (xsi, -ysi, zsi), chord, cλs, cγr, sγr, cγs, -sγs, kernel)
        end
    end

    return AIC
end

influence_matrix!(AIC, ω, U, M, xyz::AbstractArray{<:Number, 3}, symmetric; kernel=QuarticKernel()) = influence_matrix!(AIC, ω, U, M, xyz, xyz, symmetric, kernel=kernel)

"""
    pressure_coefficients(AIC, w)

Compute pressure coefficients `Cp` on each panel, given the downwash w.

# Arguments:
- `AIC`: aerodynamic influence coefficient matrix with shape (nir, njr, nis, njs)
- `w`: downwash on each panel, with shape (nir, njr)
"""
function pressure_coefficients(AIC, w)

    nir, njr, nis, njs = size(AIC)

    AIC_r = reshape(AIC, nir*njr, nis*njs)
    w_r = reshape(w, nir*njr)

    Cp = AIC_r\w_r

    return reshape(Cp, nir, njr)
end

# """
#     downwash!(w, xyz, Vinf, α, β)
#
# Compute the downwash on each of the panels in `xyz` given the freestream
# velocity `Vinf`, angle of attack `α`, sideslip angle `β`, and panel twist `θ`.
# Panel twist has the shape (ni, nj). Non-allocating.
# """
# function downwash!(w, xyz, Vinf, α, β, θ)
#
#     ca, sa = cos(α), sin(α)
#     cb, sb = cos(β), sin(β)
#
#     V = (Vinf*ca*cb, -Vinf*sb, Vinf*sa*cb)
#
#     ni = size(xyz, 2)-1
#     nj = size(xyz, 3)-1
#
#     @assert size(w) == (ni, nj)
#
#     @inbounds for j = 1:nj, i = 1:ni
#         # compute cos(dihedral) and sin(dihedral)
#         yri = (1/4)*xyz[2, i, j] + (3/4)*xyz[2, i+1, j]
#         zri = (1/4)*xyz[3, i, j] + (3/4)*xyz[3, i+1, j]
#         yro = (1/4)*xyz[2, i, j+1] + (3/4)*xyz[2, i+1, j+1]
#         zro = (1/4)*xyz[3, i, j+1] + (3/4)*xyz[3, i+1, j+1]
#         tmp = sqrt((yro-yri)^2+(zro-zri)^2)
#         cd = (yro-yri)/tmp
#         sd = (zro-zri)/tmp
#         # compute cos(theta) and sin(theta)
#         ct = cos(θ[i,j])
#         st = sin(θ[i,j])
#         # compute downwash on each panel
#         w[i, j] = V[1]*st + V[2]*(-ct*sd) + V[3]*(ct*cd)
#     end
#     return w
# end

# """
#     downwash(xyz, Vinf, α, β)
#
# Compute the downwash on each of the panels in `xyz` given the freestream
# velocity `Vinf`, angle of attack `α`, sideslip angle `β`, and panel twist `θ`.
# Panel twist has the shape (ni, nj). For a non-allocating version see `downwash!`.
# """
# function downwash(xyz, Vinf, α, β, θ)
#
#     TF = promote_type(eltype(xyz), typeof(Vinf), typeof(α), typeof(β), eltype(θ))
#
#     w = zeros(TF, size(xyz,2)-1, size(xyz,3)-1)
#
#     return downwash!(w, xyz, Vinf, α, β, θ)
# end

end # module


export iauTpxev
"""
In the tangent plane projection, given celestial direction cosines
for a star and the tangent point, solve for the star's rectangular
coordinates in the tangent plane.

This function is part of the International Astronomical Union's
SOFA (Standards of Fundamental Astronomy) software collection.

Status:  support function.

Given:
   v         double[3]  direction cosines of star (Note 4)
   v0        double[3]  direction cosines of tangent point (Note 4)

Returned:
   *xi,*eta  double     tangent plane coordinates of star

Returned (function value):
            int        status: 0 = OK
                                 1 = star too far from axis
                                 2 = antistar on tangent plane
                                 3 = antistar too far from axis

Notes:

   1. The tangent plane projection is also called the "gnomonic
   projection" and the "central projection".

   2. The eta axis points due north in the adopted coordinate system.
   If the direction cosines represent observed (RA,Dec), the tangent
   plane coordinates (xi,eta) are conventionally called the
   "standard coordinates".  If the direction cosines are with
   respect to a right-handed triad, (xi,eta) are also right-handed.
   The units of (xi,eta) are, effectively, radians at the tangent
   point.

   3. The method used is to extend the star vector to the tangent
   plane and then rotate the triad so that (x,y) becomes (xi,eta).
   Writing (a,b) for the celestial spherical coordinates of the
   star, the sequence of rotations is (a+pi/2) around the z-axis
   followed by (pi/2-b) around the x-axis.

   4. If vector v0 is not of unit length, or if vector v is of zero
   length, the results will be wrong.

   5. If v0 points at a pole, the returned (xi,eta) will be based on
   the arbitrary assumption that the longitude coordinate of the
   tangent point is zero.

   6. This function is a member of the following set:

      spherical      vector         solve for

      iauTpxes    > iauTpxev <       xi,eta
      iauTpsts      iauTpstv          star
      iauTpors      iauTporv         origin

References:

   Calabretta M.R. & Greisen, E.W., 2002, "Representations of
   celestial coordinates in FITS", Astron.Astrophys. 395, 1077

   Green, R.M., "Spherical Astronomy", Cambridge University Press,
   1987, Chapter 13.

This revision:   2018 January 2

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# int iauTpxev(double v[3], double v0[3], double *xi, double *eta)

function iauTpxev(v::AbstractVector{<:Real}, v0::AbstractVector{<:Real})
   ref_xi = Ref{Float64}(0.0)
   ref_eta = Ref{Float64}(0.0)

   status = ccall((:iauTpxev, libsofa_c), Cint,
       (Ptr{Cdouble}, Ptr{Cdouble},
       Ref{Cdouble}, Ref{Cdouble}),
       convert(Array{Float64, 1}, v),
       convert(Array{Float64, 1}, v0),
       ref_xi, ref_eta)

   return status, ref_xi[], ref_eta[]
end


# Romeo LTI utils
#
# Various utility functions.

"""
    kernel_to_signel(k, Δt)

Given a kernel vector `k` compute the corresponding response signal.

This function is an unfortunate necesity of the discretization process, and the consequence
of the way the continuous-time convolutions are implemented using sampled signals. Input
particular, if the response of a system to a given excitation `u` is computed as `g ⋆ u`,
with `⋆` denoting discrete convolution, then the continuous-time impulse response of the
system is not `g`, but `g/Δt`.
"""
kernel_to_signal(k::Array{<:Number}, Δt::Real) = k / Δt

"""
    impulse_resp([T=typeof(1.0)], sys, N, Δt)

Given a system `sys` compute its impulse response vector of length `N`, uniformly
discretized with sample time `Δt`.
"""
impulse_resp(T::Type{<:Number}, sys::SisoLtiSystem, N::Integer, Δt::Real) = kernel_to_signal(kernel(T, sys, N, Δt), Δt)
impulse_resp(sys::SisoLtiSystem, N::Integer, Δt::Real) = impulse_resp(typeof(1.0), sys::SisoLtiSystem, N::Integer, Δt::Real)

"""
    invlaplace([T=typeof(1.0)], sys, N, Δt)

Given a Laplace domain operator (a system) `sys` compute its inverse Laplace transform
as a vector of length `N`, uniformly discretized with sample time `Δt`.
"""
invlaplace(T::Type{<:Number}, sys::SisoLtiSystem, N::Integer, Δt::Real) = impulse_resp(T, sys, N, Δt)
invlaplace(sys::SisoLtiSystem, N::Integer, Δt::Real) = impulse_resp(sys, N, Δt)

"""
    step_resp([T=typeof(1.0)], sys, N, Δt)

Given a system `sys` compute its step response vector of length `N`, uniformly
discretized with sample time `Δt`.
"""
step_resp(T::Type{<:Number}, sys::SisoLtiSystem, N::Integer, Δt::Real) = impulse_resp(T, sys / Diff(1), N, Δt)
step_resp(sys::SisoLtiSystem, N::Integer, Δt::Real) = step_resp(Float64, sys, N, Δt)

"""
    simulate(sys, input, Δt)

Simulate response of `sys` to input signal `input`, uniformly sampled with `Δt`.
"""
function simulate(sys::SisoLtiSystem, input::Array{T}, Δt::Real) where T <: Number
    convolve(kernel(T, sys, length(input), Δt), input)
end

export kernel_to_signal, impulse_resp, step_resp, simulate, invlaplace


using Xpress, JuMP, MathProgBase

include(Pkg.dir("JuMP","test","solvers.jl"))
lp_solvers = [XpressSolver()]
ip_solvers = [XpressSolver()]

include(Pkg.dir("JuMP","test","model.jl"))
include(Pkg.dir("JuMP","test","probmod.jl"))

FactCheck.exitstatus()


using LinearAlgebra

"""
    a, g, b, c = mapABCD(ABCD, form="CRFB")

Compute the coefficients for the specified structure.

See realizeNTF for a list of supported structures.

stuffABCD is the inverse function.
"""
function mapABCD(ABCD, form="CRFB")
    ABCD = copy(ABCD)

    order = size(ABCD, 1) - 1
    odd = rem(order, 2)
    even = 1- odd
    diagonal = 1:(order+2):order*(order+1)
    subdiag = diagonal[1:order] .+ 1
    supdiag = diagonal[2+odd:2:order] .- 1

    # Do the mapping.
    if form == "CRFB" || form == "CIFB" || form == "CRFBD"
        c = ABCD[subdiag]
        g = -ABCD[supdiag]
        if form == "CRFB"
            dly = (2+odd):2:order;  # row numbers of delaying integrators.
            ABCD[dly,:] = ABCD[dly,:] - diagm(c[dly.-1]) * ABCD[dly.-1,:]
        elseif form == "CRFBD"
            dly = (1+odd):2:order;  # row numbers of delaying integrators.
            ABCD[dly,:] = ABCD[dly,:] + diagm(g) * ABCD[dly.+1,:]
            if order > 2
                coupl = 2+even:2:order
                ABCD[coupl,:] = ABCD[coupl,:] - diagm(c[coupl.-1]) * ABCD[coupl.-1,:]
            end
        end
        a = -ABCD[1:order,order+2]
        b = ABCD[:,order+1]

    elseif form == "CRFF"
        c = [-ABCD[1,order+2]; ABCD[subdiag[1:end-1]]]
        g = -ABCD[supdiag]
        if even == 1
            multg = 1:2:order   # Rows to have g*(following row) added.
            ABCD[multg,:] = ABCD[multg,:] + diagm(g) * ABCD[multg.+1,:]
        end
        multc = 3:2:order # Rows to have c*(preceding row) subtracted.
        ABCD[multc,:] = ABCD[multc,:] - diagm(c[multc]) * ABCD[multc.-1,:]
        a = zeros(order)
        a[2:2:order] = ABCD[order+1,2:2:order]
        for i = 2:2:order     # Recover a coeff.
            ABCD[order+1,:] = ABCD[order+1,:] - a[i] * ABCD[i,:]
        end
        a[1:2:order] = ABCD[order+1,1:2:order]
        b = ABCD[:,order+1]

    elseif form == "CRFFD"
        # CRFFD has an extra order. Correct the common variables
        order = order - 1
        odd = rem(order, 2)
        even = 1 - odd
        diagonal = diagonal[1:order]
        subdiag = diagonal[1:order-1] .+ 1
        supdiag = diagonal[2+odd:2:order] .- 1

        g = -ABCD[supdiag]
        c = [-ABCD[1,order+3]; ABCD[subdiag]]
        a = zeros(order)
        for i = 1:2:order	# Recover the a coefficients
            a[i] = ABCD[order+1,i]
            ABCD[order+1,:] = ABCD[order+1,:] - a[i] * ABCD[i,:]
        end
        a[2:2:order] = ABCD[order+1,2:2:order]
        b = ABCD[1:order+1,order+2]
        for i = 2:2:order	# Recover the b coefficients
            b[i] = b[i] - c[i] * b[i-1]
            if odd == 1
                b[i] = b[i] + g[i÷2] * b[i+1]
            end
        end
        yscale = ABCD[order+2,order+1]
        a = a * yscale;
        b[end] = b[end] * yscale

    elseif form == "CIFF" || form == "Stratos"
        a = ABCD[order+1,1:order]
        c = [-ABCD[1,order+2]; ABCD[subdiag[1:end-1]]]
        g = -ABCD[supdiag]
        b = ABCD[:,order+1]

    elseif form == "DSFB" || form == "PFF"
        error("Form $form is not yet supported")
    else
        error("Unsupported form $form")
    end

    return a, g, b, c
end


using KernelEstimator
using Base.Test

# write your own tests here
include("testreg.jl")


module TIWT
export tidwt,
       tiwpt,
       tiwpd,
       itidwt,
       itiwpt,
       itiwpd

using Wavelets

using ..Utils

# ========== Single Step Translation-Invariant Wavelet Transform ==========
"""
    tidwt_step(v, d, h, g)

Performs one level of the translation-invariant discrete wavelet transform (TWDWT) on the
vector `v`, which is the `d`-th level scaling coefficients (Note the 0-th level scaling
coefficients is the raw signal). The vectors `h` and `g` are the detail and scaling filters.

# Arguments
- `v::AbstractVector{T} where T<:Number`: Vector of coefficients from a node at level `d`.
- `d::Integer`: Depth level of `v`.
- `h::Vector{S} where S<:Number`: High pass filter.
- `g::Vector{S} where S<:Number`: Low pass filter.

# Returns
- `w₁::Vector{T}`: Output from the low pass filter.
- `w₂::Vector{T}`: Output from the high pass filter.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
v = randn(8)
wt = wavelet(WT.haar)
g, h = WT.makereverseqmfpair(wt, true)

# One step of TIDWT
tidwt_step(v, 0, h, g)
```

**See also:** [`tidwt_step!`](@ref)
"""
function tidwt_step(v::AbstractVector{T}, d::Integer, h::Vector{S}, g::Vector{S}) where
                   {T<:Number, S<:Number}
    n = length(v)
    w₁ = zeros(T, n)
    w₂ = zeros(T, n)

    tidwt_step!(w₁, w₂, v, d, h, g)
    return w₁, w₂
end

"""
    tidwt_step!(w₁, w₂, v, d, h, g)

Same as `tidwt_step` but without array allocation.

# Arguments
- `w₁::AbstractVector{T} where T<:Number`: Vector allocation for output from low pass
  filter.
- `w₂::AbstractVector{T} where T<:Number`: Vector allocation for output from high pass
  filter.
- `v::AbstractVector{T} where T<:Number`: Vector of coefficients from a node at level `d`.
- `d::Integer`: Depth level of `v`.
- `h::Vector{S} where S<:Number`: High pass filter.
- `g::Vector{S} where S<:Number`: Low pass filter.

# Returns
- `w₁::Vector{T}`: Output from the low pass filter.
- `w₂::Vector{T}`: Output from the high pass filter.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
v = randn(8)
wt = wavelet(WT.haar)
g, h = WT.makereverseqmfpair(wt, true)
w₁ = zeros(8)
w₂ = zeros(8)

# One step of TIDWT
tidwt_step!(w₁, w₂, v, 0, h, g)
```

**See also:** [`tidwt_step`](@ref)
"""
function tidwt_step!(w₁::AbstractVector{T},
                     w₂::AbstractVector{T},
                     v::AbstractVector{T},
                     d::Integer,
                     h::Vector{S},
                     g::Vector{S}) where {T<:Number, S<:Number}
    # Variable allocations to be used in Wavelets.Transforms.filtdown!
    n = length(v)                           # Signal length
    nc = Utils.nodelength(n, d+1)           # Node length at level d+1
    filtlen = length(h)                     # Filter length
    si = Vector{S}(undef, filtlen-1)        # Temp. filter vector

    # One step of translation-invariant transform
    for b in 0:(1<<d-1)
        # ----- Transform original block -----
        block = @view v[Utils.packet(d,b,n)]
        st = 2*b*nc + 1                     # Starting index in w₁, w₂
        # Decompose using low and high pass filter
        Transforms.filtdown!(g, si, w₁, st, nc, block, 1, 0, false)
        Transforms.filtdown!(h, si, w₂, st, nc, block, 1, -filtlen+1, true)
        # ----- Transform circularly shifted block -----
        block = circshift(block, 1)
        st = (2*b+1)*nc + 1                 # Starting index in w₁, w₂
        # Decompose using low and high pass filter
        Transforms.filtdown!(g, si, w₁, st, nc, block, 1, 0, false)
        Transforms.filtdown!(h, si, w₂, st, nc, block, 1, -filtlen+1, true)
    end
    return w₁, w₂
end

"""
    itidwt_step(w₁, w₂, d, g, h)

Perform one level of the inverse translation-invariant discrete wavelet transform (ITIDWT)
on the vector `w₁` and `w₂`, which are the `d+1`-th level scaling and detail coefficients
respectivaly (Note the 0-th level scaling coefficients correspond to the raw signal). The
vectors `h` and `g` are the high and low pass filters.

# Arguments
- `w₁::AbstractVector{T} where T<:Number`: `d+1`-level scaling coefficients.
- `w₂::AbstractVector{T} where T<:Number`: `d+1`-level detail coefficients.
- `d::Integer`: Depth level of output coefficients.
- `h::Vector{S} where S<:Number`: High pass filter.
- `g::Vector{S} where S<:Number`: Low pass filter.

# Returns
`v::Vector{T}`: Output of inverse translation-invariant wavelet transform.

# Examples
```julia
using WaveletsExt

# Setup
w₁, w₂ = randn(8), randn(8)
wt = wavelet(WT.haar)
g, h = WT.makereverseqmfpair(wt, false)

# One step of iTIDWT
tidwt_step(w₁, w₂, 0, h, g)
```

**See also:** [`itidwt_step!`](@ref)
"""
function itidwt_step(w₁::AbstractVector{T},
                     w₂::AbstractVector{T},
                     d::Integer,
                     h::Vector{S},
                     g::Vector{S}) where {T<:Number, S<:Number}
    v = similar(w₁)
    itidwt_step!(v, w₁, w₂, d, h, g)
    return v
end

"""
    itidwt_step!(v, w₁, w₂, d, h, g)

Same as `itidwt_step` but without array allocation.

# Arguments
- `v::AbstractVector{T} where T<:Number`: Output vector allocation.
- `w₁::AbstractVector{T} where T<:Number`: `d+1`-level scaling coefficients.
- `w₂::AbstractVector{T} where T<:Number`: `d+1`-level detail coefficients.
- `d::Integer`: Depth level of output coefficients.
- `h::Vector{S} where S<:Number`: High pass filter.
- `g::Vector{S} where S<:Number`: Low pass filter.

# Returns
`v::AbstractVector{T}`: Output from inverse translation-invariant wavelet transform of `w₁`
and `w₂`.

# Examples
```julia
using WaveletsExt

# Setup
w₁, w₂ = randn(8), randn(8)
v = similar(w₁)
wt = wavelet(WT.haar)
g, h = WT.makereverseqmfpair(wt, false)

# One step of iTIDWT
tidwt_step(v, w₁, w₂, 0, h, g)
```

**See also:** [`itidwt_step`](@ref)
"""
function itidwt_step!(v::AbstractVector{T},
                      w₁::AbstractVector{T},
                      w₂::AbstractVector{T},
                      d::Integer,
                      h::Vector{S},
                      g::Vector{S}) where {T<:Number, S<:Number}
    # Variable allocations to be used in Wavelets.Transforms.filtup!
    n = length(v)                           # Signal length
    np = Utils.nodelength(n, d)             # Node length at level d
    nc = Utils.nodelength(n, d+1)           # Node length at level d+1
    filtlen = length(h)                     # Filter length
    si = Vector{S}(undef, filtlen-1)        # Temp. filter vector
    v₁, v₂ = similar(v), similar(v)         # Temp. output storages

    # One step of translation-invariant transform
    for b in 0:(1<<d-1)
        sp = b*np + 1                       # Starting index in v
        # ----- Inverse transform original block -----
        sc = 2*b*nc + 1                     # Starting index in w₁, w₂
        # Reconstruct using low and high pass filter
        Transforms.filtup!(false, g, si, v₁, sp, np, w₁, sc, -filtlen+1, false)
        Transforms.filtup!(true, h, si, v₁, sp, np, w₂, sc, 0, true)
        # ----- Inverse transform circularly shifted block -----
        sc = (2*b+1)*nc + 1                 # Starting index in w₁, w₂
        # Reconstruct using low and high pass filter
        Transforms.filtup!(false, g, si, v₂, sp, np, w₁, sc, -filtlen+1, false)
        Transforms.filtup!(true, h, si, v₂, sp, np, w₂, sc, 0, true)
        # Adjust shift
        v₂[Utils.packet(d,b,n)] = circshift(v₂[Utils.packet(d,b,n)],-1)
    end
    # Compute result from v₁ and v₂
    for i in eachindex(v)
        @inbounds v[i] = (v₁[i] + v₂[i])/2
    end
    return v
end

# ========== Translation-Invariant DWT ==========
@doc raw"""
    tidwt(x, wt[, L])

Computes the translation-invariant discrete wavelet transform (TIDWT) for `L` levels.

# Arguments
- `x::AbstractVector{T} where T<:Number`: Original signal, preferably of size 2ᴷ where ``K
  \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`::Matrix{T}`: Output from TIDWT on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# TIDWT
xw = tidwt(x, wt)
```

**See also:** [`itidwt`](@ref)
"""
function tidwt(x::AbstractVector{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert L ≤ maxtransformlevels(x) ||
            throw(ArgumentError("Too many transform levels (length(x) < 2ᴸ)"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    n = length(x)
    w = Array{T,2}(undef, (n,L))
    v = copy(x)

    # TIDWT for L levels
    @inbounds begin
        for d in 0:(L-1)
            v, w[:,L-d] = tidwt_step(v, d, h, g)
        end
    end
    return [v w]
end

"""
    itidwt(xw, wt)

Computes the inverse translation-invariant discrete wavelet transform (iTIDWT) on `xw`.

# Arguments
- `xw::AbstractArray{T,2} where T<:Number`: TIDWT-transformed array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.

# Returns
`::Vector{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# TIDWT
xw = tidwt(x, wt)

# iTIDWT
x̂ = itidwt(xw, wt)
```

**See also:** [`tidwt`](@ref)
"""
function itidwt(xw::AbstractArray{T,2}, wt::OrthoFilter) where T<:Number
    # Setup
    _, k = size(xw)
    g, h = WT.makereverseqmfpair(wt, false)
    x = xw[:,1]             # Output
    L = k-1                 # Number of decomposition levels

    # iTIDWT
    for i in 1:(k-1)
        @inbounds x = itidwt_step(x, xw[:,i+1], L-i, h, g)
    end
    return x
end

# ========== Translation-Invariant WPT ==========
@doc raw"""
    tiwpt(x, wt[, L])

Computes `L` levels of translation-invariant wavelet packet transform (TIWPT) on `x`.

# Arguments
- `x::AbstractVector{T} where T<:Number`: Original signal, preferably of size 2ᴷ where ``K
  \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`::Matrix{T}`: Output from TIWPT on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# TIWPT
xw = tiwpt(x, wt)
```

**See also:** [`itiwpt`](@ref), [`tidwt`](@ref)
"""
function tiwpt(x::AbstractVector{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert L ≤ maxtransformlevels(x) ||
            throw(ArgumentError("Too many transform levels (length(x) < 2ᴸ)"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    n = length(x)
    xw = Matrix{T}(undef, (n, 1<<L))
    xw[:,1] = x

    # TIWPT for L levels
    for d in 0:(L-1)
        nblock = 1<<d
        for b in 0:(nblock-1)
            # Parent index
            np = (1<<L)÷nblock
            i = b*np + 1
            # Children indices
            nc = np÷2
            j₁ = 2*b*nc + 1
            j₂ = (2*b+1)*nc + 1
            # Overwrite output of TIDWT directly onto xw
            xw[:,j₁], xw[:,j₂] = tidwt_step(xw[:,i], d, h, g)
        end
    end
    return xw
end

"""
    itiwpt(xw, wt)

Computes the inverse translation-invariant wavelet packet transform (iTIWPT) on `xw`.

# Arguments
- `xw::AbstractArray{T,2} where T<:Number`: TIDWT-transformed array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.

# Returns
`::Vector{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# TIWPT
xw = tiwpt(x, wt)

# iTIWPT
x̂ = itiwpt(xw, wt)
```

**See also:** [`tiwpt`](@ref), [`itidwt`](@ref)
"""
function itiwpt(xw::AbstractArray{T,2}, wt::OrthoFilter) where T<:Number
    # Sanity check
    n, m = size(xw)
    @assert isdyadic(m) || throw(ArgumentError("Number of columns of xw is not dyadic."))
    @assert m ≤ maxtransformlevels(n) ||
            throw(ArgumentError("Number of nodes in `xw` is more than possible number of nodes at any depth for signal of length `n`"))

    # Setup
    g, h = WT.makereverseqmfpair(wt, false)
    temp = copy(xw)                         # Temp. array to store intermediate outputs
    L = ndyadicscales(m)                    # Number of decompositions from xw

    # iTIWPT
    for d in (L-1):-1:0
        nblock = 1<<d
        for b in 0:(nblock-1)
            # Parent index
            np = (1<<L)÷nblock
            i = b*np + 1
            # Children indices
            nc = np÷2
            j₁ = 2*b*nc + 1
            j₂ = (2*b+1)*nc + 1
            # Overwrite output of iTIWPT directly onto temp
            temp[:,i] = itidwt_step(temp[:,j₁], temp[:,j₂], d, h, g)
        end
    end
    return temp[:,1]
end

# ========== Translation-Invariant WPD ==========
@doc raw"""
    tiwpd(x, wt[, L])

Computes `L` levels of translation-invariant wavelet packet decomposition (TIWPD) on `x`.

# Arguments
- `x::AbstractVector{T} where T<:Number`: Original signal, preferably of size 2ᴷ where ``K
  \in \mathbb{N}``.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(x)`) Number of levels of decomposition.

# Returns
`::Matrix{T}`: Output from TIWPD on `x`.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# TIWPD
xw = tiwpd(x, wt)
```

**See also:** [`itiwpd`](@ref), [`tiwpt`](@ref)
"""
function tiwpd(x::AbstractVector{T},
               wt::OrthoFilter,
               L::Integer = maxtransformlevels(x)) where T<:Number
    # Sanity check
    @assert L ≤ maxtransformlevels(x) ||
            throw(ArgumentError("Too many transform levels (length(x) < 2ᴸ)"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))

    # Setup
    g, h = WT.makereverseqmfpair(wt, true)
    n = length(x)                       # Size of signal
    n₀ = 1<<(L+1)-1                     # Total number of nodes
    n₁ = n₀ - (1<<L)                    # Total number of nodes excluding leaf nodes
    xw = Matrix{T}(undef, (n, n₀))      # Output allocation
    xw[:,1] = x

    # TIWPD
    for i in 1:n₁
        d = floor(Integer, log2(i))     # Current depth
        xw[:,left(i)], xw[:,right(i)] = tidwt_step(xw[:,i], d, h, g)
    end
    return xw
end

"""
    itiwpd(xw, wt[, L])
    itiwpd(xw, wt, tree)

Computes the inverse translation-invariant wavelet packet decomposition (iTIWPD) on `xw`.

# Arguments
- `xw::AbstractArray{T,2} where T<:Number`: TIDWT-transformed array.
- `wt::OrthoFilter`: Orthogonal wavelet filter.
- `L::Integer`: (Default: `maxtransformlevels(size(xw,1))`) Depth of inverse decomposition.
- `tree::BitVector`: Representation of binary tree used as basis for the signal `x`.

# Returns
`::Vector{T}`: Inverse transformed signal.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
x = generatesignals(:heavysine)
wt = wavelet(WT.haar)

# TIWPD
xw = tiwpd(x, wt)

# iTIWPD
x̂ = itiwpd(xw, wt)
```

**See also:** [`tiwpd`](@ref), [`itiwpt`](@ref)
"""
function itiwpd(xw::AbstractArray{T,2},
                wt::OrthoFilter,
                L::Integer = maxtransformlevels(size(xw,1))) where T<:Number
    # Sanity check
    @assert L ≤ maxtransformlevels(x) ||
            throw(ArgumentError("Too many transform levels (length(x) < 2ᴸ)"))
    @assert L ≥ 1 || throw(ArgumentError("L must be ≥ 1"))
    return itiwpd(xw, wt, maketree(size(xw,1), L, :full))
end

function itiwpd(xw::AbstractArray{T,2}, wt::OrthoFilter, tree::BitVector) where
                T<:Number
    # Sanity check
    @assert isvalidtree(xw[:,1], tree)

    # Setup
    g, h = WT.makereverseqmfpair(wt, false)
    tmp = copy(xw)

    # iTIWPD
    for (i, haschild) in Iterators.reverse(enumerate(tree))
        # Current node has child => Compute one step of inverse transform
        if haschild
            # Current depth
            d = floor(Integer, log2(i))
            # Parent node
            v = @view tmp[:,i]
            # Children nodes
            w₁ = @view tmp[:,left(i)]
            w₂ = @view tmp[:,right(i)]
            # Inverse transform
            itidwt_step!(v, w₁, w₂, d, h, g)
        end
    end
    return tmp[:,1]
end

end


using PooksoftOptionsKit
using PyPlot
using Dates
using Distributions

# simulation logic -
function simulate(pathToSimulationFile::String)

    # initialize -
    lattice_model = nothing
    asset_set = nothing
    asset_price_array = nothing
    data_array_at_exp = nothing

    # load the lattice model -
    result = build_simulation_lattice_data_structure(pathToSimulationFile)
    if (typeof(result.value) == PSError)
        return result.value
    else
        lattice_model = result.value
    end

    # load the asset set -
    result = build_simulation_asset_set(pathToSimulationFile)
    if (typeof(result.value) == PSError)
        return result.value
    else
        asset_set = result.value
    end

    # load asset price array -
    result = build_simulation_price_array(pathToSimulationFile)
    if (typeof(result.value) == PSError)
        return result.value
    else
        asset_price_array = result.value
    end

    # main loop -
    number_of_time_steps = length(asset_price_array)
    data_array = zeros(number_of_time_steps,2)
    for (index,asset_price) in enumerate(asset_price_array)

        # build the pricing tree -
        tree = build_ternary_price_tree(asset_set, lattice_model, asset_price)

        # compute -
        local result = option_contract_price(tree, lattice_model)

        # grab -
        data_array[index,1] = asset_price
        data_array[index,2] = result.value
    end

    # run simulation at expiration -
    # compute -
    result = compute_option_profit_and_loss_at_expiration(asset_set, asset_price_array)
    if (typeof(result.value) == PSError)
        return result.value
    else
        data_array_at_exp = result.value
    end

    # return -
    return (data_array, data_array_at_exp)
end

# setup -
path_to_design_file = "$(pwd())/tests/config/IronCondor.json"
(data_array, data_array_at_exp) = simulate(path_to_design_file)

# make a plot from the simulation -


# (1) パッケージをインストールする（初回のみ）
## ] + addで

# (2) ライブラリの読み出し
## dplyr的に操作したいのでDataFramesMetaを使う
using DataFrames
using DataFramesMeta
using HypothesisTests
using RCall
@rlibrary readr

# (3) データの読み込み
## CSV.read()だとネット上のファイルは読み込めないのでRCallを使う。遅い。
email_data =
  rcopy(R"readr::read_csv('http://www.minethatdata.com/Kevin_Hillstrom_MineThatData_E-MailAnalytics_DataMiningChallenge_2008.03.20.csv')")

# (4) データの準備
## 女性向けメールが配信されたデータを削除したデータを作成
@linq male_df =
  email_data |> where(:segment .!= "Womens E-Mail") |> # 女性向けメールが配信されたデータを削除
  transform(treatment = ifelse.(:segment .== "Mens E-Mail", 1, 0)) #介入を表すtreatment変数を追加

# (5) 集計による比較
## byでgroupbyとtransformを同時にできる
@linq summary_by_segment =
  male_df |> by(
    :treatment,
    conversion_rate = mean(:conversion), # グループごとのconversionの平均
    spend_mean = mean(:spend), # グループごとのspendの平均
    count = length(:conversion),
  ) # グループごとのデータ数

# (6) t検定を行う
## (a)男性向けメールが配信されたグループの購買データを得る
@linq mens_mail = male_df |> where(:treatment .== 1)

## (b)メールが配信されなかったグループの購買データを得る
@linq no_mail = male_df |> where(:treatment .== 0)

## (a)(b)の平均の差に対して有意差検定を実行する
## HypothesisTestsの等分散t検定を使う
rct_ttest = EqualVarianceTTest(mens_mail[:, :spend], no_mail[:, :spend])

# (7) セレクションバイアスのあるデータの作成
## seedを固定する
## 書籍とseedを合わせるためRCallを使う
R"set.seed(1)"

## 条件に反応するサンプルの量を半分にする
obs_rate_c = 0.5
obs_rate_t = 0.5

## バイアスのあるデータの作成
@linq biased_data =
  male_df |> transform(
    obs_rate_c = ifelse.(
      (:history .> 300) .| (:recency .< 6) .| (:channel .== "Multichannel"),
      obs_rate_c,
      1,
    ),
    obs_rate_t = ifelse.(
      (:history .> 300) .| (:recency .< 6) .| (:channel .== "Multichannel"),
      1,
      obs_rate_t,
    ),
    random_number = rcopy(R"runif($(nrow(male_df)))"), # 乱数を合わせるためRに作らせる
  ) |> where(
    ((:treatment .== 0) .& (:random_number .< :obs_rate_c)) .|
    ((:treatment .== 1) .& (:random_number .< :obs_rate_t)),
  )

# (8) セレクションバイアスのあるデータで平均を比較
## group_byとsummairseを使って集計(Biased)
## byだけでできる
@linq summary_by_segment_biased =
  biased_data |> by(
    :treatment,
    conversion_rate = mean(:conversion),
    spend_mean = mean(:spend),
    count = length(:conversion),
  )

# (9) Rの関数であるt.testを使ってt検定を行う(Biased)
## (a)男性向けメールが配信されたグループの購買データを得る
@linq mens_mail_biased = biased_data |> where(:treatment .== 1)

## (b)メールが配信されなかったグループの購買データを得る
@linq no_mail_biased = biased_data |> where(:treatment .== 0)

## (a)(b)の平均の差に対して有意差検定を実行
rct_ttest_biased =
  EqualVarianceTTest(mens_mail_biased[:, :spend], no_mail_biased[:, :spend])


export densitymatrix, ket

"""
    densitymatrix(net, prob, norm=true)

Returns the Density matrix encoded by the neural network `net`, and normalizes
it if `norm==true`.
"""
function densitymatrix(net::NeuralNetwork, hilb::AbstractSuperOpBasis, norm=true)
    if !(net isa CachedNet)
        net = cached(net)
    end

    v = state(hilb, net)

    p_hilb = physical(hilb)
    ρ = zeros(out_type(net), spacedimension(p_hilb), spacedimension(p_hilb))
    #if v isa DiagonalStateWrapper
    #    v = v.parent
    #end
    for i=1:spacedimension(p_hilb)
        set_index!(row(v), p_hilb, i)
        for j=1:spacedimension(p_hilb)
            set_index!(col(v), p_hilb, j)
            ρ[i,j] = exp(net(v))
        end
    end
    if norm
        ρ ./= tr(ρ)
    end
    ρ
end

densitymatrix(net::NeuralNetwork, hilb::AbstractHilbert, norm=true) =
    densitymatrix(net, SuperOpSpace(hilb), norm)

Base.Matrix(net::NeuralNetwork, hilb, norm=true) = densitymatrix(net, hilb, norm)
Base.Vector(net::MatrixNet, hilb, norm=true) = vec(densitymatrix(net, hilb, norm))

"""
    ket(net, prob, norm=true)

Returns the state (ket) encoded by the neural network `net`, and normalizes
it if `norm==true`.
"""
function ket(net, hilb::AbstractHilbert, norm=true)
    if !(net isa CachedNet)
        net = cached(net)
    end

    psi = zeros(out_type(net), spacedimension(hilb))

    v = state(hilb, net)
    for i=1:spacedimension(hilb)
        set_index!(v, hilb, i)
        psi[i] = exp(net(v))
    end
    if norm
        normalize!(psi)
    end
    return psi
end

ket(net::MatrixNet, hilb::AbstractSuperOpBasis, norm=true) =
    ket(net, hilb, norm)

ket(net::MatrixNet, hilb::AbstractHilbert, norm=true) =
    ket(net, SuperOpSpace(hilb), norm)

Base.Vector(net::NeuralNetwork, hilb, norm=true) = ket(net, hilb, norm)

QuantumOpticsBase.DenseOperator(net::NeuralNetwork, hilb::AbstractHilbert, norm=true) =
    DenseOperator(convert(CompositeBasis, hilb), Matrix(net, hilb, norm))

QuantumOpticsBase.Ket(hilb::AbstractHilbert, net::KetNeuralNetwork, norm=true) =
    Ket(convert(CompositeBasis, hilb), ket(net, hilb, norm))

QuantumOpticsBase.expect(Ô::AbsLinearOperator, state) =
    expect(SparseOperator(Ô), state)
##


# This file was generated, do not modify it. # hide
using Gnuplot, Colors, ColorSchemes, Random
Random.seed!(123)
function custom_palette(colormap=:viridis)
    cmap = get(colorschemes[colormap], LinRange(0,1,256))
    ctmp = "0 '#$(hex(cmap[1]))',"
    for i in 2:256; ctmp = ctmp*"$(i-1) '#$(hex(cmap[i]))'," end;
    "set palette defined("*ctmp[1:end-1]*")"
end
# Archimedes spiral
a = 1.5
b = -2.4
t = LinRange(0,5*π,500)
x = (a .+ b*t) .* cos.(t)
y = (a .+ b*t) .* sin.(t)
@gp " " :-  # this is just to reset session
@gp :- "set multiplot layout 3,3; set key off;
    unset ytics; unset xtics; unset border" :-
colormaps = [:magma, :viridis, :plasma, :inferno, :berlin,
    :leonardo, :devon, :spring, :ice]
for i in 1:9
    @gp :- i title = "$(colormaps[i])" "set size square" :-
    @gp(:-, x, y, t, "w l lw 3 dt 1 lc palette",
        custom_palette(colormaps[i]),"set cbtics out nomirror", :-)
end
@gp
save(term="pngcairo size 900,800", output="plt_ex9.png") # hide


type ConfusionMatrix
    classes::Vector
    matrix::Matrix{Int}
    accuracy::Float64
    kappa::Float64
end

function show(io::IO, cm::ConfusionMatrix)
    print(io, "Classes:  ")
    show(io, cm.classes)
    print(io, "\nMatrix:   ")
    display(cm.matrix)
    print(io, "\nAccuracy: ")
    show(io, cm.accuracy)
    print(io, "\nKappa:    ")
    show(io, cm.kappa)
end

function _hist_add!{T}(counts::Dict{T,Int}, labels::Vector{T}, region::Range1{Int})
    for i in region
        lbl = labels[i]
        counts[lbl] = get(counts, lbl, 0) + 1
    end
    return counts
end

function _hist_sub!{T}(counts::Dict{T,Int}, labels::Vector{T}, region::Range1{Int})
    for i in region
        lbl = labels[i]
        counts[lbl] -= 1
    end
    return counts
end

function _hist_shift!{T}(counts_from::Dict{T,Int}, counts_to::Dict{T,Int}, labels::Vector{T}, region::Range1{Int})
    for i in region
        lbl = labels[i]
        counts_from[lbl] -= 1
        counts_to[lbl] = get(counts_to, lbl, 0) + 1
    end
    return nothing
end

_hist{T}(labels::Vector{T}, region::Range1{Int} = 1:endof(labels)) =
    _hist_add!(Dict{T,Int}(), labels, region)

function _set_entropy{T}(counts::Dict{T,Int}, N::Int)
    entropy = 0.0
    for v in values(counts)
        if v > 0
            entropy += v * log(v)
        end
    end
    entropy /= -N
    entropy += log(N)
    return entropy
end

_set_entropy(labels::Vector) = _set_entropy(_hist(labels), length(labels))

function _info_gain{T}(N1::Int, counts1::Dict{T,Int}, N2::Int, counts2::Dict{T,Int})
    N = N1 + N2
    H = - N1/N * _set_entropy(counts1, N1) - N2/N * _set_entropy(counts2, N2)
    return H
end

function _neg_z1_loss{T<:Real}(labels::Vector, weights::Vector{T})
    missmatches = labels .!= majority_vote(labels)
    loss = sum(weights[missmatches])
    return -loss
end

function _weighted_error{T<:Real}(actual::Vector, predicted::Vector, weights::Vector{T})
    mismatches = actual .!= predicted
    err = sum(weights[mismatches]) / sum(weights)
    return err
end

function _mse_loss{T<:Float64, U<:Real}(labels::Vector{T}, features::Vector{U}, thresh)
    s_l = s_r = s2_l = s2_r = 0.0
    nl = nr = 0
    @inbounds for i in 1:length(labels)
      l = labels[i]
      if features[i] < thresh
        s_l += l
        s2_l += l*l
        nl += 1
      else
        s_r += l
        s2_r += l*l
        nr += 1
      end
    end
    loss = s2_l - s_l^2/nl + s2_r - s_r^2/nr
    return -loss
end

function majority_vote(labels::Vector)
    if length(labels) == 0
        return nothing
    end
    counts = _hist(labels)
    top_vote = labels[1]
    top_count = -1
    for (k,v) in counts
        if v > top_count
            top_vote = k
            top_count = v
        end
    end
    return top_vote
end

### Classification ###

function confusion_matrix(actual::Vector, predicted::Vector)
    @assert length(actual) == length(predicted)
    N = length(actual)
    _actual = zeros(Int,N)
    _predicted = zeros(Int,N)
    classes = sort(unique([actual; predicted]))
    N = length(classes)
    for i in 1:N
        _actual[actual .== classes[i]] = i
        _predicted[predicted .== classes[i]] = i
    end
    CM = zeros(Int,N,N)
    for i in zip(_actual, _predicted)
        CM[i[1],i[2]] += 1
    end
    accuracy = trace(CM) / sum(CM)
    prob_chance = (sum(CM,1) * sum(CM,2))[1] / sum(CM)^2
    kappa = (accuracy - prob_chance) / (1.0 - prob_chance)
    return ConfusionMatrix(classes, CM, accuracy, kappa)
end

function _nfoldCV(classifier::Symbol, labels, features, args...)
    nfolds = args[end]
    if nfolds < 2
        return nothing
    end
    if classifier == :tree
        pruning_purity = args[1]
    elseif classifier == :forest
        nsubfeatures = args[1]
        ntrees = args[2]
        partialsampling = args[3]
    elseif classifier == :stumps
        niterations = args[1]
    end
    N = length(labels)
    ntest = _int(floor(N / nfolds))
    inds = randperm(N)
    accuracy = zeros(nfolds)
    for i in 1:nfolds
        test_inds = falses(N)
        test_inds[(i - 1) * ntest + 1 : i * ntest] = true
        train_inds = !test_inds
        test_features = features[inds[test_inds],:]
        test_labels = labels[inds[test_inds]]
        train_features = features[inds[train_inds],:]
        train_labels = labels[inds[train_inds]]
        if classifier == :tree
            model = build_tree(train_labels, train_features, 0)
            if pruning_purity < 1.0
                model = prune_tree(model, pruning_purity)
            end
            predictions = apply_tree(model, test_features)
        elseif classifier == :forest
            model = build_forest(train_labels, train_features, nsubfeatures, ntrees, partialsampling)
            predictions = apply_forest(model, test_features)
        elseif classifier == :stumps
            model, coeffs = build_adaboost_stumps(train_labels, train_features, niterations)
            predictions = apply_adaboost_stumps(model, coeffs, test_features)
        end
        cm = confusion_matrix(test_labels, predictions)
        accuracy[i] = cm.accuracy
        println("\nFold ", i)
        println(cm)
    end
    println("\nMean Accuracy: ", mean(accuracy))
    return accuracy
end

nfoldCV_tree(labels::Vector, features::Matrix, pruning_purity::Real, nfolds::Integer)                                           = _nfoldCV(:tree, labels, features, pruning_purity, nfolds)
nfoldCV_forest(labels::Vector, features::Matrix, nsubfeatures::Integer, ntrees::Integer, nfolds::Integer, partialsampling=0.7)  = _nfoldCV(:forest, labels, features, nsubfeatures, ntrees, partialsampling, nfolds)
nfoldCV_stumps(labels::Vector, features::Matrix, niterations::Integer, nfolds::Integer)                                         = _nfoldCV(:stumps, labels, features, niterations, nfolds)

### Regression ###

function mean_squared_error(actual, predicted)
    @assert length(actual) == length(predicted)
    return mean((actual - predicted).^2)
end

function R2(actual, predicted)
    @assert length(actual) == length(predicted)
    ss_residual = sum((actual - predicted).^2)
    ss_total = sum((actual .- mean(actual)).^2)
    return 1.0 - ss_residual/ss_total
end

function _nfoldCV{T<:Float64, U<:Real}(regressor::Symbol, labels::Vector{T}, features::Matrix{U}, args...)
    nfolds = args[end]
    if nfolds < 2
        return nothing
    end
    if regressor == :tree
        maxlabels = args[1]
    elseif regressor == :forest
        nsubfeatures = args[1]
        ntrees = args[2]
        maxlabels = args[3]
        partialsampling = args[4]
    end
    N = length(labels)
    ntest = _int(floor(N / nfolds))
    inds = randperm(N)
    R2s = zeros(nfolds)
    for i in 1:nfolds
        test_inds = falses(N)
        test_inds[(i - 1) * ntest + 1 : i * ntest] = true
        train_inds = !test_inds
        test_features = features[inds[test_inds],:]
        test_labels = labels[inds[test_inds]]
        train_features = features[inds[train_inds],:]
        train_labels = labels[inds[train_inds]]
        if regressor == :tree
            model = build_tree(train_labels, train_features, maxlabels, 0)
            predictions = apply_tree(model, test_features)
        elseif regressor == :forest
            model = build_forest(train_labels, train_features, nsubfeatures, ntrees, maxlabels, partialsampling)
            predictions = apply_forest(model, test_features)
        end
        err = mean_squared_error(test_labels, predictions)
        corr = cor(test_labels, predictions)
        r2 = R2(test_labels, predictions)
        R2s[i] = r2
        println("\nFold ", i)
        println("Mean Squared Error:     ", err)
        println("Correlation Coeff:      ", corr)
        println("Coeff of Determination: ", r2)
    end
    println("\nMean Coeff of Determination: ", mean(R2s))
    return R2s
end

nfoldCV_tree{T<:Float64, U<:Real}(labels::Vector{T}, features::Matrix{U}, nfolds::Integer, maxlabels::Integer=5)      = _nfoldCV(:tree, labels, features, maxlabels, nfolds)
nfoldCV_forest{T<:Float64, U<:Real}(labels::Vector{T}, features::Matrix{U}, nsubfeatures::Integer, ntrees::Integer, nfolds::Integer, maxlabels::Integer=5, partialsampling=0.7)  = _nfoldCV(:forest, labels, features, nsubfeatures, ntrees, maxlabels, partialsampling, nfolds)


using OptimizationDynamics
const iLQR = OptimizationDynamics.IterativeLQR
using LinearAlgebra
using Random

# ## visualization
vis = Visualizer()
render(vis)

# ## mode
MODE = :friction
MODE = :frictionless

# ## state-space model
h = 0.05
T = 51

if MODE == :friction
	im_dyn = ImplicitDynamics(cartpole_friction, h, eval(r_cartpole_friction_func), eval(rz_cartpole_friction_func), eval(rθ_cartpole_friction_func);
		r_tol=1.0e-8, κ_eval_tol=1.0e-4, κ_grad_tol=1.0e-3, no_impact=true)
        cartpole_friction.friction .= [0.35; 0.35]
        ## cartpole_friction.friction .= [0.25; 0.25]
        ## cartpole_friction.friction .= [0.1; 0.1]
        ## cartpole_friction.friction .= [0.01; 0.01]
else
	im_dyn = ImplicitDynamics(cartpole_frictionless, h, eval(r_cartpole_frictionless_func), eval(rz_cartpole_frictionless_func), eval(rθ_cartpole_frictionless_func);
    	        r_tol=1.0e-8, κ_eval_tol=1.0, κ_grad_tol=1.0, no_impact=true, no_friction=true)
end

nx = 2 * cartpole_friction.nq
nu = cartpole_friction.nu

# ## iLQR model
ilqr_dyn = iLQR.Dynamics((d, x, u, w) -> f(d, im_dyn, x, u, w),
					(dx, x, u, w) -> fx(dx, im_dyn, x, u, w),
					(du, x, u, w) -> fu(du, im_dyn, x, u, w),
					nx, nx, nu)
model = [ilqr_dyn for t = 1:T-1];

# ## initial conditions
q0 = [0.0; 0.0]
q1 = [0.0; 0.0]
qT = [0.0; π]
q_ref = [0.0; π]

x1 = [q1; q1]
xT = [qT; qT]

# ## objective
function objt(x, u, w)
	J = 0.0
	J += transpose(u) * u
	return J
end

function objT(x, u, w)
        J = 0.0
        J += transpose(x - xT) * (x - xT)
        return J
end

ct = iLQR.Cost(objt, nx, nu)
cT = iLQR.Cost(objT, nx, 0)
obj = [[ct for t = 1:T-1]..., cT];

# ## constraints
function terminal_con(x, u, w)
    [
    x - xT; # goal
    ]
end

cont = iLQR.Constraint()
conT = iLQR.Constraint(terminal_con, nx, 0)
cons = [[cont for t = 1:T-1]..., conT];

# ## rollout
ū = [(t == 1 ? -1.5 : 0.0) * ones(nu) for t = 1:T-1] # set value to -1.0 when friction coefficient = 0.25
x̄ = iLQR.rollout(model, x1, ū)
q̄ = state_to_configuration(x̄)
visualize!(vis, cartpole_friction, q̄, Δt=h);

# ## solver
solver = iLQR.solver(model, obj, cons,
    opts=iLQR.Options(linesearch = :armijo,
    α_min=1.0e-5,
    obj_tol=1.0e-5,
    grad_tol=1.0e-3,
    max_iter=100,
    max_al_iter=20,
    con_tol=0.005,
    ρ_init=1.0,
    ρ_scale=10.0,
    verbose=false))
iLQR.initialize_controls!(solver, ū)
iLQR.initialize_states!(solver, x̄);

# ## solve
iLQR.reset!(solver.s_data)
@time iLQR.solve!(solver);

@show iLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w)
@show solver.s_data.iter[1]
@show norm(terminal_con(solver.m_data.x[T], zeros(0), zeros(0)), Inf)
@show solver.s_data.obj[1] # augmented Lagrangian cost

# ## solution
x_sol, u_sol = iLQR.get_trajectory(solver)
q_sol = state_to_configuration(x_sol)
visualize!(vis, cartpole_friction, q_sol, Δt=h);

# ## benchmark
solver.options.verbose = false
@benchmark iLQR.solve!($solver, x̄, ū) setup=(x̄=deepcopy(x̄), ū=deepcopy(ū));


using Revise
using ADCME
using AdFem
using SparseArrays
using PyPlot

function plot_disp(u, m, n, h)
    close("all")
    u = reshape(u, m+1, n+1)
    x = collect((0:m) * h)
    y = collect((0:n) * h)
    pcolormesh(x, y, u')
    colorbar()
    axis("scaled")
end

function plot_disp_line(u, m, n, h)
   close("all")
    u = reshape(u, m+1, n+1)
   plot(u[:, n÷2], label="middle")
   plot(u[:, n÷3], label="upper")
   plot(u[:, n÷3*2], label="lower")
   legend()
end


export FixedPatch

"""
    FixedPatch()

The `FixedPatch` is similar to the [`OrthogonalPatch`](@ref) but it doesn't change during
the tracking. Instead it only updates at the start of the tracking.
"""
struct FixedPatch <: AbstractAffinePatch end

struct FixedPatchState{T,N} <: AbstractAffinePatchState{N}
    v̄::PVector{T,N} # this is already conjugated
end

function state(::FixedPatch, x::PVector)
    v = copy(x)
    conj!(v.data)
    FixedPatchState(v)
end

function init!(state::FixedPatchState, x::AbstractVector)
    @boundscheck length(x) == length(state.v̄)
    LinearAlgebra.normalize!(x)
    @inbounds for i in eachindex(state.v̄)
        state.v̄[i] = conj(x[i])
    end
    state
end

onpatch!(x::AbstractVector, state::FixedPatchState) = onpatch!(x, state.v̄)
evaluate!(u, state::FixedPatchState, x::PVector) = evaluate_patch!(u, state.v̄, x)
jacobian!(U, state::FixedPatchState, x::PVector) = jacobian_patch!(U, state.v̄, x)


function testNLEQ2()
    refSol = Dict{Int64,Vector}();

    refSol[2] = [0.21132486540517284, 0.78867513459482719];

    refSol[3] = [0.14644660863174364, 0.5, 0.85355339136825636];

    refSol[4] = [0.10267276384825176, 0.40620376278203324, 0.5937962372179667,
                    0.89732723615174825];

    refSol[5] = [0.083751256499240576, 0.31272929522379189, 0.5, 0.68727070477620811,
                    0.91624874350075947];

    refSol[6] = [0.066876587754941222, 0.28874077956671557, 0.36668218955839399,
                    0.63331781044160607, 0.71125922043328438, 0.93312341224505879];

    refSol[7] = [0.058069149620871653, 0.23517161236154782, 0.33804409473596703, 0.5,
                    0.66195590526403303, 0.76482838763845207, 0.94193085037912838];

    refSol[8] = [0.054730284645850648, 0.20965065342027825, 0.26705029739773067,
                    0.50002862809115289, 0.50003386025132168, 0.7329693255337042,
                    0.79033329229948646, 0.94516665683719758];

    refSol[9] = [0.044205346135779117, 0.19949067230774983, 0.23561910847320633,
                    0.41604690789257159, 0.5, 0.5839530921074283, 0.76438089152679367,
                    0.80050932769225014, 0.95579465386422091];

    testResult = false;
    for dim = 2:9
        n1 = dim + 1

        # Initialize the options
        opt = OptionsNLEQ(OPT_JACGEN            => 1,
                          OPT_JACFCN            => chebyQuadJac,
                          OPT_MSTOR             => 0,
                          OPT_NOROWSCAL         => 0,
                          OPT_NITMAX            => 200,
                          OPT_RTOL              => 1e-5)

        x0    = collect(1:dim)./n1
        xScal = zeros(x0)

        (x0, _, retCode) = nleq2(chebyQuad,x0,xScal,opt);

        relNormDiff = norm(x0-refSol[dim],Inf)/norm(refSol[dim],Inf);

        if dim == 2
            testResult = relNormDiff < eps(Float64);
        else
            testResult &= relNormDiff < eps(Float64);
        end
    end
    return testResult
end

testNLEQ2()


#=
Plots consumption, income and debt for the simple infinite horizon LQ
permanent income model with Gaussian iid income.

@author : Spencer Lyon

@date: 07/09/2014
=#

using Plots
pyplot()

const r = 0.05
const beta = 1.0 / (1.0 + r)
const T = 60
const sigma = 0.15
const mu = 1.0

function time_path2()
    w = randn(T+1)
    w[1] =  0.0
    b = Array{Float64}(T+1)
    for t=2:T+1
        b[t] = sum(w[1:t])
    end
    b .*= -sigma
    c = mu + (1.0 - beta) .* (sigma .* w .- b)
    return w, b, c
end

# == Figure showing a typical realization == #
function single_realization()
    w, b, c = time_path2()
    p = plot(0:T, mu + sigma .* w, color=:green, label="non-financial income")
    plot!(c, color=:black, label="consumption")
    plot!(b, color=:blue, label="debt")
    plot!(xlabel="Time", linewidth=2, alpha=0.7, xlims=(0, T))

    return p
end

# == Figure showing multiple consumption paths == #
function consumption_paths(n=250)  # n is number of paths
    time_paths = []

    for i=1:n
        push!(time_paths, time_path2()[3])
    end

    p = plot(time_paths, linewidth=0.8, alpha=0.7, legend=:none)
    plot!(xlabel="Time", ylabel="Consumption", xlims=(0, T))
    return p
end


module aoc_09
    using Test
    export read_input, part1, part2
    export test_part1, test_part2, test
    export input, demo

    const input_r = r""

    function read_input(filename = "input.txt")
        #lines = readlines(filename)
        # for line in eachline(filename)
        #for line in lines
        #end
        [parse_line(line) for line in eachline(filename)]
        #lines
    end

    function parse_line(line)
        parse(Int,line)
    end

    const input = read_input("input.txt")
    const demo = read_input("demo.txt")

    function part1(input = input, preamble = 25)
        for i = preamble + 1 : length(input)
            prev = input[i-preamble:i-1]
            valid = any(in.(input[i] .- prev, (prev,)))
            if !valid
                return input[i]
            end
        end
        0
    end

    function test_part1()
        @testset "Part 1" begin
            @test part1( demo , 5) == 127
            @test part1() == 2089807806
        end
    end

    function part2(input = input, preamble = 25)
        target = part1(input, preamble)
        cs = cumsum(input)
        # min_cs = cs[1]
        # gt_cs = filter(>=(target+min_cs), cs)
        # lt_cs = filter(<=(gt_cs[end] - target), cs)
        cs_range = ()
        for gt in cs
            valid = gt .- cs .== target
            if any( valid )
                cs_range = ( cs[ findfirst(valid) ], gt )
                break
            end
        end
        ind = ( (x,y)->findfirst(isequal( x ), cs) + y ).(cs_range,(1,0))
        sum( extrema( input[ind[1]:ind[2]] ) )
    end

    function test_part2()
        @testset "Part 2" begin
            @test part2( demo , 5) == 62
            @test part2() == 245848639
        end
    end

    function test()
        test_part1()
        test_part2()
    end
end


include("utils\\categorical_tree.jl")
include("utils\\randutils.jl")

a = collect(1:5)
ct = CategoricalTree{Int64}(a)
ct[1] = 0
ct
ct[2] = 0
ct
ct[3] = 0
ct
ct[4] = 0
ct
ct[5] = 0
ct
sum(ct)
length(ct)
rand_categorical(ct)
a
N = 1000
counts = zeros(Int, length(ct))
for i in 1:N
    counts[rand_categorical(ct)] += 1
end

using Plots
bar(1:length(ct), counts/N, label="Empirical probability", legend=:topleft)
scatter!(1:length(ct), a/sum(ct), label="theoretical probability")

using Distributions
χ = sum(((counts-N*a/sum(ct)).^2)./(N*a/sum(ct)))
quantile(Chisq(length(a)-1), 0.95)

using Statistics

n = 10000
m = 5000
step = 100
function compare_times(n,m,step)
    N = 2:step:n
    t_a = zeros(length(N))
    t_ct = zeros(length(N))
    sd_a = zeros(length(N))
    sd_ct = zeros(length(N))
    ta = zeros(m)
    tct = zeros(m)
    for i in 1:length(N)
        a = collect(1:N[i])
        ct = CategoricalTree{Int64}(a)
        for k in 1:m
            ta[k] = @elapsed rand_categorical(a)
            tct[k] = @elapsed rand_categorical(ct)
        end
        t_a[i] = mean(ta)
        t_ct[i] = mean(tct)
        sd_a[i] = std(ta)
        sd_ct[i] = std(tct)
    end
    return N, t_a, t_ct, sd_a, sd_ct
end

N, ta, tct, sda, sdct = compare_times(n,m,step)

using Plots
plot(N, log2.(ta), label="Array based algorithm",
    legend=:topleft,
    xlabel="N",
    ylabel="log2 t")
plot!(N, log2.(tct), label="Tree based algorithm")

a = spzeros(5)
a[2] = 2
a[3] = 3
a[5] = 5
a
rand_categorical(a)
collect(zip(findnz(a)))

N = 10000
counts = zeros(Int, length(a))
for i in 1:N
    counts[rand_categorical(a)] += 1
end

using Plots
bar(1:length(a), counts/N, label="Empirical probability", legend=:topleft)
scatter!(1:length(a), a/sum(a), label="Theoretical probability")

using Distributions
χ = sum(((counts-N*a/sum(a)).^2)./(N*a/sum(a)))
quantile(Chisq(length(a)-1), 0.95)


import Random.rand!

"""
    _rescale!(mat)

Changes the matrix `mat` so that it is between `0` and `1`.
"""
function _rescale!(mat)
    mat .-= NaNMath.minimum(mat)
    mat ./= NaNMath.maximum(mat)
end

"""
    mask!(array::AbstractArray{<:AbstractFloat}, maskarray::AbstractArray{<:AbstractBool})

Modifies `array` so that the positions at which `maskarray` is `false` are
replaced by `NaN`.
"""
function mask!(array::AbstractArray{<:Float64}, maskarray::AbstractArray{<:Bool})
    (size(array) == size(maskarray)) || throw(DimensionMismatch("The dimensions of array, $(size(array)), and maskarray, $(size(maskarray)), must match. "))
    array[.!maskarray] .= NaN
    array
end

"""
    rand(alg, dims::Tuple{Vararg{Int64,2}}; mask=nothing) where {T <: Integer}

Creates a landscape of size `dims` (a tuple of two integers) following the model
defined by `alg`. The `mask` argument accepts a matrix of boolean values, and is
passed to `mask!` if it is not `nothing`.
"""
function Base.rand(alg::T, dims::Tuple{Int64,Int64}; mask=nothing) where {T <: NeutralLandscapeMaker}
    ret = Matrix{Float64}(undef, dims...)
    rand!(ret, alg; mask=mask)
end
Base.rand(alg::T, dims::Integer...; mask=nothing) where {T <: NeutralLandscapeMaker} = rand(alg, dims; mask = mask)

"""
    rand!(mat, alg) where {IT <: Integer}

Fill the matrix `mat` with a landscape created following the model defined by
`alg`. The `mask` argument accepts a matrix of boolean values, and is passed to
`mask!` if it is not `nothing`.
"""
function rand!(mat::AbstractArray{<:AbstractFloat,2} where N, alg::T; mask=nothing) where {T <: NeutralLandscapeMaker}
    _landscape!(mat, alg)
    isnothing(mask) || mask!(mat, mask)
    _rescale!(mat)
end


order(crs::AbstractAlgebraicCoarsening) = crs.order

function _add_to_neighborhood!(N, k, v, maxsize)
    push!(N, (k, v)) # TODO Smarter implementation
    if length(N) > maxsize
        sort!(N; by=x->x[1], rev=true)
        pop!(N)
    end
    return nothing
end

function coarse_neighborhoods(W, C, invseed, order)
    nc = length(C)
    nf = length(invseed) - length(C)
    # TODO Other strategies for sparsifying neighborhood?
    N = [Vector{Tuple{Float64, Int}}() for i in 1:nf]
    rows = rowvals(W)
    vals = nonzeros(W)

    for (j, c) in enumerate(C)
        for idx in nzrange(W, c)
            row = rows[idx]
            val = vals[idx]
            if val != 0 && invseed[row] > nc
                idxf = invseed[row] - nc
                _add_to_neighborhood!(N[idxf], val, j, order)
            end
        end
    end
    return N
end

function _fillcoarseop_coarse!(Pi, Pj, Pv, C, idx)
    for (i, c) in enumerate(C)
        idx += 1
        Pi[idx] = c
        Pj[idx] = i
        Pv[idx] = 1.0
    end
    return idx
end

function _fillcoarseop_fine!(Pi, Pj, Pv, F, N, idx)
    for (i, f) in enumerate(F)
        total = sum(x -> x[1], N[i])
        for (w, c) in N[i]
            idx += 1
            Pi[idx] = f
            Pj[idx] = c
            Pv[idx] = w / total
        end
    end
    return idx
end

function _fillcoarseop!(Pi, Pj, Pv, C, F, N)
    idx = 0
    idx = _fillcoarseop_coarse!(Pi, Pj, Pv, C, idx)
    idx = _fillcoarseop_fine!(Pi, Pj, Pv, F, N, idx)
    return nothing
end

function formcoarseop(crs::AbstractAlgebraicCoarsening, C, F, N)
    nc = length(C)
    nf = length(F)
    k = sum(length, N) + nc
    Pi = Vector{Int}(undef, k)
    Pj = Vector{Int}(undef, k)
    Pv = Vector{Float64}(undef, k)
    _fillcoarseop!(Pi, Pj, Pv, C, F, N)
    return sparse(Pi, Pj, Pv, nc + nf, nc)
end

function fix_adjacency!(A)
    for i in 1:size(A, 1)
        A[i, i] = 0
    end
    return A
end

function fix_laplacian!(L)
    rows = rowvals(L)
    vals = nonzeros(L)
    for c in axes(L, 2)
        acc = 0
        for idx in nzrange(L, c)
            if rows[idx] != c
                acc -= vals[idx]
            end
        end
        L[c, c] = acc
    end
    return L
end


include("basic_binary_tree.jl")
"""
array-based binary search tree
left tree values < root value < right tree values
"""
abstract type AbstractBinarySearchTree_arr <: AbstractBinaryTree_arr end
function Base.findfirst(
    value::T,
    tree::Tree,
    place::Int = tree.root,
) where {T,Tree<:AbstractBinarySearchTree_arr}
    if tree.val[place] == value
        return place
    elseif tree.val[place] > value
        if tree.lch[place] == 0
            return 0
        end
        return findfirst(value, tree, tree.lch[place])
    else
        if tree.rch[place] == 0
            return 0
        end
        return findfirst(value, tree, tree.rch[place])
    end
end


module FlowNodeLayersPermutationLayerTest

using Test
using ReactiveMP
using LinearAlgebra

using ReactiveMP: getP, getmat, getdim, forward, forward!, backward, backward!, jacobian, inv_jacobian
using ReactiveMP: det_jacobian, absdet_jacobian, logdet_jacobian, logdet_jacobian, logabsdet_jacobian

@testset "Permutation Layer" begin

    @testset "Constructor" begin

        # check for placeholder creation
        layer = PermutationLayer()
        @test typeof(layer)   == ReactiveMP.PermutationLayerPlaceholder

        # check for layer with available P matrix
        P = PermutationMatrix(5)
        layer = PermutationLayer(5, P)
        @test layer.P == P
        @test typeof(layer)   <: PermutationLayer
        @test typeof(layer)   <: ReactiveMP.AbstractLayer
        @test typeof(layer.P) <: PermutationMatrix

        # check for layer with unknown P matrix
        layer = PermutationLayer(4)
        @test size(layer.P)   == (4, 4)
        @test typeof(layer)   <: PermutationLayer
        @test typeof(layer)   <: ReactiveMP.AbstractLayer
        @test typeof(layer.P) <: PermutationMatrix

    end

    @testset "Get" begin

        # check for layer with available P matrix
        P = PermutationMatrix(5)
        layer = PermutationLayer(5, P)
        @test layer.P           == getP(layer)
        @test layer.P           == getmat(layer)
        @test getP(layer)       == P
        @test layer.dim         == getdim(layer)
        @test getdim(layer)     == 5
        @test typeof(layer.P)   <: PermutationMatrix

        # check for layer with unknown P matrix
        layer = PermutationLayer(4)
        @test layer.P           == getP(layer)
        @test layer.P           == getmat(layer)
        @test size(layer.P)     == (4, 4)
        @test getdim(layer)     == layer.dim
        @test getdim(layer)     == 4
        @test typeof(layer.P)   <: PermutationMatrix

    end

    @testset "Prepare-Compile" begin

        layer = ReactiveMP._prepare(3, ReactiveMP.PermutationLayerPlaceholder())
        @test typeof(layer)     <: Tuple
        @test typeof(layer[1])  <: PermutationLayer
        @test layer[1].dim      == 3
        @test size(layer[1].P)  == (3,3)

        P = PermutationMatrix(4)
        layer = PermutationLayer(4, P)
        @test compile(layer)  == layer
        @test_throws ArgumentError compile(layer, 1)

        @test nr_params(layer) == 0

    end

    @testset "Base" begin

        # check for layer with available P matrix
        P = PermutationMatrix(5)
        layer = PermutationLayer(5, P)
        @test eltype(layer) <: Integer

        # check for layer with unknown P matrix
        layer = PermutationLayer(4)
        @test eltype(layer) <: Integer

    end

    @testset "Forward-Backward" begin

        # check forward function
        P = PermutationMatrix(3)
        layer = PermutationLayer(3, P)
        @test forward(layer, [5.0, 1.5, 8.5])  == P*[5.0, 1.5, 8.5]
        @test forward(layer, [4.0, 2.5, -9.0]) == P*[4.0, 2.5, -9.0]
        @test forward.(layer, [[5.0, 1.5, 8.5], [4.0, 2.5, -9.0]]) == [P*[5.0, 1.5, 8.5], P*[4.0, 2.5, -9.0]]

        # check forward! function
        P = PermutationMatrix(3)
        layer = PermutationLayer(3, P)
        output = zeros(3)
        forward!(output, layer, [5.0, 1.5, 8.5])
        @test output == P*[5.0, 1.5, 8.5]
        forward!(output, layer, [4.0, 2.5, -9.0])
        @test output == P*[4.0, 2.5, -9.0]

        # check bakward function
        P = PermutationMatrix(3)
        layer = PermutationLayer(3, P)
        @test backward(layer, [5.0, 1.5, 8.5])  == P'*[5.0, 1.5, 8.5]
        @test backward(layer, [4.0, 2.5, -9.0]) == P'*[4.0, 2.5, -9.0]
        @test backward.(layer, [[5.0, 1.5, 8.5], [4.0, 2.5, -9.0]]) == [P'*[5.0, 1.5, 8.5], P'*[4.0, 2.5, -9.0]]

        # check backward! function
        P = PermutationMatrix(3)
        layer = PermutationLayer(3, P)
        input = zeros(3)
        backward!(input, layer, [5.0, 1.5, 8.5])
        @test input == P'*[5.0, 1.5, 8.5]
        backward!(input, layer, [4.0, 2.5, -9.0])
        @test input == P'*[4.0, 2.5, -9.0]

    end

    @testset "Jacobian" begin

        # check jacobian function
        P = PermutationMatrix(3)
        layer = PermutationLayer(3, P)
        @test jacobian(layer, randn(3)) == P
        @test jacobian(layer, randn(3)) == P
        @test jacobian.(layer, [randn(3), randn(3)]) == [P, P]

        # check invjacobian function
        P = PermutationMatrix(3)
        layer = PermutationLayer(3, P)
        @test inv_jacobian(layer, randn(3)) == P'
        @test inv_jacobian(layer, randn(3)) == P'
        @test inv_jacobian.(layer, [randn(3), randn(3)]) == [P', P']

    end

    @testset "Utility Jacobian" begin

        # check utility functions jacobian
        P = PermutationMatrix(3)
        layer = PermutationLayer(3, P)
        @test det_jacobian(layer, randn(3)) == det(P)
        @test det_jacobian(layer) == det(P)
        @test absdet_jacobian(layer, randn(3)) == 1.0
        @test absdet_jacobian(layer) == 1.0
        @test logdet_jacobian(layer, randn(3)) == 0.0
        @test logdet_jacobian(layer) == 0.0
        @test logabsdet_jacobian(layer, randn(3)) == 0.0
        @test logabsdet_jacobian(layer) == 0.0

    end

end

end


module Day04

import Base.show
import ..data_dir # from parent module

export part1, part2

# Matrix conversion is hard
as_matrix(vectors::Vector{Vector{Int64}}) = Matrix(transpose(hcat(vectors...)))

# Define BingoGame and BingoBoard structs --- a BingoGame is a BingoBoard with state

const bingo_board_size = 5 # number of rows and columns

struct BingoBoard
    definition::Matrix{Int64}
end

function BingoBoard(text_definition::Vector{String})
    parse_row(row) = map(x -> parse(Int, x), split(row))
    as_vectors = map(x -> parse_row(x), text_definition)
    BingoBoard(as_matrix(as_vectors))
end

mutable struct BingoGame
    board::BingoBoard
    state::Matrix{Bool}
end

BingoGame(board::BingoBoard) = BingoGame(
    board,
    zeros(Int64, size(board.definition)...)
)

BingoGame(text_definition::Vector{String}) = BingoGame(BingoBoard(text_definition))

function show(io::IO, game::BingoGame)
    defn = game.board.definition
    rows, cols = size(defn)
    to_print = Vector{String}()
    for row = 1:rows
        row_to_print = Vector{String}()
        for col = 1:cols
            number_to_print = lpad(defn[row, col], 2, " ")
            if game.state[row, col]
                number_to_print = "\033[;1m" * number_to_print * "\033[0m"
            end
            push!(row_to_print, number_to_print)
        end
        push!(to_print, join(row_to_print, " "))
    end
    println(io, join(to_print, "\n"))
end

function show(io::IO, board::BingoBoard)
    show(io, BingoGame(board))
end

# Functions for playing a BingoGame

function update!(game::BingoGame, number::Int)
    match_coords = findfirst(x -> x == number, game.board.definition)
    if !isnothing(match_coords)
        game.state[match_coords] = true
    end
end

function has_won(state::Matrix{Bool})
    winning = false
    for i = 1:bingo_board_size
        if all(state[i, :]) || all(state[:, i])
            winning = true
            break
        end
    end
    winning
end

has_won(game::BingoGame) = has_won(game.state)

function score(game::BingoGame, number)
    negated_state = ones(Bool, size(game.state)...) - game.state
    unmarked_numbers_only = negated_state .* game.board.definition # Hadamard product
    sum_unmarked_numbers = sum(unmarked_numbers_only)
    number * sum_unmarked_numbers
end

# Parse the input
const data_file = joinpath(data_dir, "day04")
const input = readlines(data_file) # read into an array, each line is an element

function extract_numbers_and_boards(input)
    numbers = map(x -> parse(Int, x), split(input[1], ","))
    input_without_numbers = filter(x -> x != "", input)[2:end]
    board_inputs = [
        input_without_numbers[i:i+4]
        for i in range(1, Int(length(input_without_numbers)), step = bingo_board_size)
    ]
    bingo_boards = [BingoBoard(x) for x in board_inputs]
    numbers, bingo_boards
end

"""
Iterate through all of the numbers until a winning game is identified, then stop.
If multiple games win in a single round, we take the first in the array.
"""
function part1(input = input; silent = false)
    numbers, bingo_boards = extract_numbers_and_boards(input)
    games = [BingoGame(board) for board in bingo_boards]
    winning_game = nothing
    number_index = 0
    while isnothing(winning_game)
        number_index += 1
        number = numbers[number_index]
        if !silent
            println("Turn " * string(number_index) * " with number " * string(number))
        end
        for game in games
            update!(game, number)
            # Don't check for a win before turn 5
            if number_index >= 5 && has_won(game)
                winning_game = game
                break
            end
        end
    end
    winning_score = score(winning_game, numbers[number_index])
    if !silent
        println()
        println("Winner with score " * string(winning_score))
        println(winning_game)
        println()
    end
    winning_score
end

"""
A little different to part 1. We need to find the _last_ game to win. The approach here
is to update every game, and then identify the first game in the array to win.
Then we remove all winning games and continue until all of the games are won. This way
we're constantly updating _winning_game_ with the most recent victory, and the last
iteration of the `while` loop will evaluate the last games to win.
"""
function part2(input = input; silent = false)
    numbers, bingo_boards = extract_numbers_and_boards(input)
    games = [BingoGame(board) for board in bingo_boards]
    winning_game = nothing
    number_index = 0
    while length(games) > 0
        number_index += 1
        number = numbers[number_index]
        if !silent
            println("Turn " * string(number_index) * " with number " * string(number))
        end
        for game in games
            update!(game, number)
        end
        for game in games
            # Don't check for a win before turn 5
            if number_index >= 5 && has_won(game)
                winning_game = game
                break
            end
        end
        games = games[(!has_won).(games)]
    end
    winning_score = score(winning_game, numbers[number_index])
    if !silent
        println()
        println("Winner with score " * string(winning_score))
        println(winning_game)
        println()
    end
    winning_score
end

end # module


"""
```
label_components(OneComponent3D(), binary_volume, local_window)
```

Computes the connected components of a binary volume using 26-connectivity.

# Output

Returns a three-dimensional `Int` array that represents the label of every voxel.

# Arguments

## `binary_volume::AbstractArray`

The binary volume made up of values 0 (black) and 1 (white).

## `local_window::NTuple{3, Int}`

The size of the local window for each iteration.

# Example

Computes the labels for a randomly generated 20 x 20 x 20 array using 5 x 7 x 7 local window.

```julia
using ImageComponentAnalysis

binary_volume = Int.(rand(Bool, 20, 20, 20))

labels = label_components(OneComponent3D(), binary_volume, (5, 7, 7))
```

# Reference

[1] Q. Hu, G. Qian, and W.L. Nowinski, “Fast connected-component labelling in three-dimensional binary images based on iterative recursion”, Computer Vision and Image Understanding, vol. 99, no. 3, pp. 414–434, Sep. 2005. [doi:10.1016/j.cviu.2005.04.001](https://doi.org/10.1016/j.cviu.2005.04.001)
"""
function label_components(algorithm::OneComponent3D, binary_volume::AbstractArray, local_window::NTuple{3, Int})

	@inbounds begin
		height, width, length = size(binary_volume)
		queue = Queue{NTuple{3, Int}}()

		n₁, n₂, n₃ = map((b, l) -> ceil(Int, (max(1, min(b, l)-1)/2)), size(binary_volume), local_window)
		labels = padarray(Int, binary_volume, Fill(0, (n₁ + 1, n₂ + 1, n₃ + 1)))
		labelindex = 3

		for x = 1:height, y = 1:width, z = 1:length
			if (labels[x, y, z] == 1)
				enqueue!(queue, (x, y, z))
				while !isempty(queue)
					vᵢx, vᵢy, vᵢz = dequeue!(queue)
					label_subvolume!(algorithm, vᵢx, vᵢy, vᵢz, n₁, n₂, n₃, queue, labels, labelindex)
				end
				labelindex += 1
			end
		end

		labels = view(labels, 1:height, 1:width, 1:length)
		replace!(s -> s >= 3 ? s-2 : 0, labels)
	end
end

# Label 26-connected voxels in the subvolume with the labelindex.
function label_subvolume!(algorithm, vᵢx, vᵢy, vᵢz, n₁, n₂, n₃, queue, labels, labelindex)

	window_size = 0
	center = (vᵢx, vᵢy, vᵢz)

	# Set center voxel of subvolume as labelindex.
	labels[vᵢx, vᵢy, vᵢz] = labelindex

	# Label 3 x 3 x 3 neighbourhood of center voxel.
	for (x, y, z) = get_neighbours(algorithm, 1, center, (n₁, n₂, n₃))
		if (labels[x, y, z] == 1)
			labels[x, y, z] = -10
		end
	end

	# Grow window size to local window.
	while window_size != max(n₁, n₂, n₃) + 1
		for (vx, vy, vz) = get_neighbours(algorithm, window_size, center, (n₁, n₂, n₃))

			# Label object voxels with distance of window_size.
			if (labels[vx, vy, vz] == -10)
				if (abs(vx - vᵢx) == n₁ || abs(vy - vᵢy) == n₂ || abs(vz - vᵢz) == n₃)
					enqueue!(queue, (vx, vy, vz))
				end

				# Label 3 x 3 x 3 neighbourhood of current object voxel.
				for (x, y, z) = get_neighbours(algorithm, 1, (vx, vy, vz), (n₁, n₂, n₃), center)
					if (labels[x, y, z] == 1)
						labels[x, y, z] = -10
						if (abs(x - vᵢx) == n₁ || abs(y - vᵢy) == n₂ || abs(z - vᵢz) == n₃)
							enqueue!(queue, (x, y, z))
						end
					end
				end

				labels[vx, vy, vz] = labelindex
			end
		end
		window_size += 1
	end

	# Replace and queue all object voxels with label -10.
	for (x, y, z) = Iterators.product(vᵢx-n₁:vᵢx+n₁, vᵢy-n₂:vᵢy+n₂, vᵢz-n₃:vᵢz+n₃)
		if labels[x, y, z] == -10
			enqueue!(queue, (x, y, z))
			labels[x, y, z] = labelindex
		end
	end
end

# Get the neighbourhood of the center voxel.
function get_neighbours(algorithm::OneComponent3D, distance, center, n, truecenter = center)

	neighbour_ranges = map((c, n, t) -> max(t-n,c-distance):min(t+n,c+distance), center, n, truecenter)

	function select_neighbours((x, y, z))
		max(abs(x-center[1]), abs(y-center[2]), abs(z-center[3])) == distance
	end

	Iterators.filter(select_neighbours, Iterators.product(neighbour_ranges...))
end


function equilibrate!(M, aff, opt)
    max_iters=opt.equilibration_iters
    lb = opt.equilibration_lb
    ub = opt.equilibration_ub

    @timeit "eq init" begin
        α = (aff.n / (aff.m + aff.p)) ^ .25
        β = ((aff.m + aff.p) / aff.n ) ^ .25
        α2, β2 = α^2, β^2
        γ = .1

        u, v           = zeros(aff.m + aff.p), zeros(aff.n)
        u_, v_         = zeros(aff.m + aff.p), zeros(aff.n)
        u_grad, v_grad = zeros(aff.m + aff.p), zeros(aff.n)
        row_norms, col_norms = zeros(aff.m + aff.p), zeros(aff.n)
        E = LinearAlgebra.Diagonal(u)
        D = LinearAlgebra.Diagonal(v)
        M_ = copy(M)

        rows_M_ = SparseArrays.rowvals(M_)
    end

    for iter in 1:max_iters
        @timeit "update diags" begin
            E.diag .= exp.(u)
            D.diag .= exp.(v)
        end
        @timeit "M_" begin
            LinearAlgebra.mul!(M_, M, D)
            LinearAlgebra.mul!(M_, E, M_)
        end

        step_size = 2. / (γ * (iter + 1.))

        # u gradient step
        @timeit "norms" begin
            fill!(row_norms, 0.0)
            fill!(col_norms, 0.0)
            for col in 1:aff.n
                for line in SparseArrays.nzrange(M_, col)
                    row_norms[rows_M_[line]] += abs2(M_[rows_M_[line], col])
                    col_norms[col] += abs2(M_[rows_M_[line], col])
                end
            end
        end
        @timeit "grad" begin
            @. u_grad = row_norms - α2 + γ * u
            @. v_grad = col_norms - β2 + γ * v
        end
        @timeit "proj " begin
            # u
            @. u -= step_size * u_grad
            box_project!(u, lb, ub)
            # v
            @. v -= step_size * v_grad
            sum_v = sum(v)
            v .= sum_v / aff.n
            box_project!(v, 0., ub)
        end
        # Update averages.
        @timeit "update" begin
            @. u_ = 2 * u / (iter + 2) + iter * u_ / (iter + 2)
            @. v_ = 2 * v / (iter + 2) + iter * v_ / (iter + 2)
        end
    end

    @timeit "update diags" begin
        E.diag .= exp.(u_)
        D.diag .= exp.(v_)
    end

    return E, D
end

function box_project!(y, lb, ub)
    y .= min.(ub, max.(y, lb))
    nothing
end


"""
This module defines types for Pauli Matrix
"""
module Pauli

abstract type PauliMatrix{T <: Complex} <: AbstractMatrix{T} end

struct X{T} <: PauliMatrix{T} end
struct Y{T} <: PauliMatrix{T} end
struct Z{T} <: PauliMatrix{T} end

Base.Matrix(x::PauliMatrix{T}) where T = Matrix{T}(x)
Base.Matrix{T}(::X) where T = T[0 1;1 0]
Base.Matrix{T}(::Y) where T = T[0 -im;im 0]
Base.Matrix{T}(::Z) where T = T[1 0;0 -1]

using SparseArrays

SparseArrays.SparseMatrixCSC(x::PauliMatrix{T}) where T = SparseMatrixCSC{T}(x)
SparseArrays.SparseMatrixCSC{T}(x::X) where T = sparse(Matrix{T}(x))
SparseArrays.SparseMatrixCSC{T}(x::Y) where T = sparse(Matrix{T}(x))
SparseArrays.SparseMatrixCSC{T}(x::Z) where T = sparse(Matrix{T}(x))

SparseArrays.nonzeros(::X{T}) where T = T[1, 1]
SparseArrays.nonzeros(::Y{T}) where T = T[im, -im]
SparseArrays.nonzeros(::Z{T}) where T = T[1, -1]

function Base.checkbounds(::Type{Bool}, x::PauliMatrix, ind1, ind2)
    1 <= ind1 <= 2 || throw(BoundsError(x, (ind1, ind2)))
    1 <= ind2 <= 2 || throw(BoundsError(x, (ind1, ind2)))
    true
end

function Base.getindex(x::X{T}, inds...) where T
    @boundscheck checkbounds(Bool, x, inds...)

    inds == (1, 2) || inds == (2, 1) ? one(T) : zero(T)
end

function Base.getindex(x::Y{T}, inds...) where T
    @boundscheck checkbounds(Bool, x, inds...)

    inds == (1, 2) || inds == (2, 1) ? T(-im) : zero(T)
end

function Base.getindex(x::Z{T}, inds...) where T
    @boundscheck checkbounds(Bool, x, inds...)

    inds == (1, 1) ? -one(T) :
    inds == (2, 2) ? one(T)  :
    zero(T)
end

Base.size(::PauliMatrix) = (2, 2)

Base.summary(io::IO, ::X{T}) where T = print(io, "Pauli X{$T}")
Base.summary(io::IO, ::Y{T}) where T = print(io, "Pauli Y{$T}")
Base.summary(io::IO, ::Z{T}) where T = print(io, "Pauli Z{$T}")

end