using StaticOptim, Optim
using Test, Random

# All of these tests are from OptimTestProblems.jl

using StaticArrays
@testset "Unconstrained" begin
    sx = @SVector ones(2)
    sx = 3.2 * sx
    rosenbrock(x) =  (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
    res = soptimize(rosenbrock, sx)
    @test res.g_converged == true
    @test rosenbrock(res.minimizer) == res.minimum
    @test show(res) === nothing
    res = soptimize(rosenbrock, sx)
    @test res.g_converged == true
    res = soptimize(rosenbrock, sx, updating = true)
    @test res.g_converged == true
    res = soptimize(rosenbrock, sx, bto = StaticOptim.Order3())
    @test res.g_converged == true
    @test rosenbrock(res.minimizer) == res.minimum

    res = soptimize(rosenbrock, sx/2)
    @test res.g_converged == true
    @test rosenbrock(res.minimizer) == res.minimum

    function fletcher_powell(x::AbstractVector)
        function theta(x::AbstractVector)
            if x[1] > 0
                return atan(x[2] / x[1]) / (2.0 * pi)
            else
                return (pi + atan(x[2] / x[1])) / (2.0 * pi)
            end
        end

        return 100.0 * ((x[3] - 10.0 * theta(x))^2 +
        (sqrt(x[1]^2 + x[2]^2) - 1.0)^2) + x[3]^2
    end

    sx = @SVector ones(3)
    sx = 3.2 * sx
    res = soptimize(fletcher_powell, sx)
    @test res.g_converged == true
    @test fletcher_powell(res.minimizer) == res.minimum
    res = soptimize(fletcher_powell, sx/2)
    @test res.g_converged == true
    @test fletcher_powell(res.minimizer) == res.minimum

    function himmelblau(x::AbstractVector)
        return (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
    end
    sx = @SVector ones(2)
    sx = 3.2 * sx
    res = soptimize(himmelblau, sx)
    @test res.g_converged == true
    @test himmelblau(res.minimizer) == res.minimum
    res = soptimize(himmelblau, sx/2)
    @test res.g_converged == true
    @test himmelblau(res.minimizer) == res.minimum

    function powell(x::AbstractVector)
        return (x[1] + 10.0 * x[2])^2 + 5.0 * (x[3] - x[4])^2 +
        (x[2] - 2.0 * x[3])^4 + 10.0 * (x[1] - x[4])^4
    end
    sx = @SVector ones(4)
    sx = 3.2 * sx
    res = soptimize(powell, sx)
    @test res.g_converged == true
    @test powell(res.minimizer) == res.minimum
    res = soptimize(powell, sx/2)
    @test res.g_converged == true
    @test powell(res.minimizer) == res.minimum
end

# Constrained
@testset "Contrained" begin
    function U(h, w)
        h1, h2 = h[1], h[2]
        h1 >= 1 && return -Inf*one(h1)
        h2 >= 1 && return -Inf*one(h1)
        c = 2w*h1 + 1.5w*h2 + 0.1
        c <= 0 && return -Inf*one(h1)
        log(c) + log(1 - h1) + log(1 - h2)
    end

    l = SVector{2}(0., 0.4)
    sx = SVector{2}(0.7, 0.7)
    res = constrained_soptimize(x -> -U(x, 1), sx, lower = l)
    @test res.g_converged == true
    lbs = [SVector{2}(x, y) for x in 0:0.1:0.4, y in 0:0.1:0.4]
    ubs = [SVector{2}(x, y) for x in 0.6:.1:0.9, y in 0.6:0.1:0.9]
    ws = 0.7:0.1:1.01
    sx = SVector{2}(0.5, 0.5)
    constrained_soptimize(x -> -U(x, 1.), sx, lower = SVector{2}(0., 0.1))
    sols = [(w = w, l = l, res = constrained_soptimize(x -> -U(x, w), sx, lower = l)) for l in lbs, u in ubs, w in ws]
    @test all(sol.res.g_converged == true for sol in sols)

        ## Constrained, 4 people

        function U(h, w)
            h1, h2, h3, h4 = h[1], h[2], h[3], h[4]
            h1 >= 1 && return -Inf*one(h1)
            h2 >= 1 && return -Inf*one(h1)
            h3 >= 1 && return -Inf*one(h1)
            h4 >= 1 && return -Inf*one(h1)
            c = 2w*h1 + 1.5w*h2 + 1.2w*h3 + 0.5w*h4 + 0.1
            c <= 0 && return -Inf*one(h1)
            log(c) + log(1 - h1) + log(1 - h2) + log(1 - h3) + log(1 - h4)
        end

        l = SVector{4}(0., 0., 0., 0.)
        sx = SVector{4}(0.7, 0.7, 0.7, 0.7)
        res = constrained_soptimize(x -> -U(x, 1), sx, lower = l)
        @test res.g_converged == true
        lbs = [SVector{4}(x, y, z, t) for x in 0:0.1:0.4, y in 0:0.1:0.4, z in 0:0.1:0.4, t in 0:0.1:0.4]
        ubs = [SVector{4}(x, y, z, t) for x in 0.6:.1:0.9, y in 0.6:0.1:0.9, z in 0.6:0.1:0.9, t in 0.6:0.1:0.9]
        ws = 0.5:0.5:2.
        sx = SVector{4}(0.5, 0.5, 0.5, 0.5)
        constrained_soptimize(x -> -U(x, 1), sx, lower = SVector{4}(0.3, 0., 0.2, 0.2))
        #sols = [(w = w, l = l, res = constrained_soptimize(x -> -U(x, w), sx, lower = l, upper = u)) for l in lbs, u in ubs, w in ws];
        @test all(constrained_soptimize(x -> -U(x, w), sx, lower = l, upper = u).g_converged for l in lbs, u in ubs, w in ws)
        #sols = [(w = 1, l = l, res = constrained_soptimize(x -> -U(x, 1), sx, lower = l)) for l in lbs]
    end

    ### Univariate tests
    @testset "Univeriate" begin
        f(x) = x^2 + 2*x
        res = soptimize(f, 1.)
        @test res.g_converged == true
        @test f(res.minimizer) == res.minimum

        function U(x)
            x = max(0., x)
            log(x)
        end
        f(h) = -(U(0.2 + h) + U(1 - h))
        res = soptimize(f, 0.9)
        @test res.g_converged == true
        @test f(res.minimizer) == res.minimum

        ### Root finding
        up(c) = c <= 0 ? Inf*c : 1/c
        f(a) = up(2 - a) - .96up(2 + a)
        out = sroot(f, 0.5)
        @test out.fx < 1e-8
        @test f(out.x) == out.fx

        out = sroot(f, (-0.5, 0.5))
        @test  out.isroot == true
        @test f(out.x) == out.fx

        w = 3.
        beta = 0.96
        R = 1.01
        alpha_h = 1.5
        uc(c) = 1/c
        uh(h) = alpha_h/(1 - h)
        function eulerfun(x)
            a, h1, h2 = x[1], x[2], x[3]
            c1 = w*h1 - a
            c2 = w*h2 + R*a
            (h1 >= 1 || h2 >= 1) && return Inf*x
            c1 <= 0 && return Inf*x
            out1 = uc(c1) - beta*R*uc(c2)
            out2 = w*uc(c1) - uh(h1)
            out3 = w*uc(c2) - uh(h2)
            SVector(out1, out2, out3)
        end
        x = SVector(0., 0.5, 0.5)
        res = sroot(eulerfun, x)
        @test res.g_converged == true
    end
    @testset "Constrained vs Optim" begin
        # Test regular arrays (Nonlinear least squares)
        N_params = 15
        N_sim = 5000
        Random.seed!(1234)
        realparam = rand(N_params)
        data = rand(N_sim, N_params)
        data[:, 1] = ones(N_sim)
        y = data*realparam .+ data*exp.(realparam) .+ randn(N_sim)
        paramg = rand(N_params)
        function obj(param)
            yhat = data * param .+ data*exp.(param)
            sum((yy - yh)^2 for (yy, yh) in zip(y, yhat))
        end
        sres = soptimize(obj, paramg);
        @test sres.g_converged == true
        res = optimize(obj, paramg, method = Newton(), autodiff = :forward)
        @test all(x -> abs(x) < 1e-6, res.minimizer - sres.minimizer)
        l = rand(N_params)
        u = rand(N_params)
        u = [uu < ll ? Inf : uu for (uu, ll) in zip(u, l)]
        m = (u + l)/2
        cparamg = [isfinite(mm) ? mm : 2ll for (mm, ll) in zip(m, l)]
        df = TwiceDifferentiable(obj, cparamg, autodiff = :forward)
        dfc = TwiceDifferentiableConstraints(l, u)
        resoc = optimize(df, dfc, cparamg, IPNewton())
        sresoc = constrained_soptimize(obj, cparamg, lower = l, upper = u)

        @test all(x -> abs(x) < 1e-6, resoc.minimizer - sresoc.minimizer)
    end


"""Gradient noise generation (e.g. Perlin noises)
"""

function adjacent(p)
    fp1, fp2 = floor.(p)
    cp1, cp2 = ceil.(p)
    @SArray Int[fp1 cp1 fp1 cp1; fp2 fp2 cp2 cp2]
end

function lerp_weights(w)
    w1, w2 = w
    @SVector [1 - w1 - w2 + w1 * w2, w1 - w1 * w2, w2 - w1 * w2, w1 * w2]
end

function perlin!(A::AbstractArray, grid::AbstractArray)
    resolution, scale = size(A), size(grid)
    ranges = map.(Ref(x -> (1, x)), [resolution, scale])
    n = ndims(grid)
    neighbors = zeros(Int, n, 2^n)
    weights = zeros(2^n)
    # dot_prod_tmp = @SArray zeros(n)
    dot_prod_tmp = zeros(n)
    scaled_indices = zeros(n)
    # w = @SArray zeros(n)
    w = zeros(n)
    it_indices = collect(Base.Iterators.product(range.(Ref(1), resolution, step=1)...))

    # for indices in Base.Iterators.product(range.(Ref(1), resolution, step=1)...)
    #     scaled_indices .= remap.(indices, ranges...)
    #     @. w = fade(scaled_indices - floor(scaled_indices))
    #     lerp_weights!(weights, w)
    #     adjacent!(neighbors, scaled_indices)
    #     value = 0
    #     for i in 1:2^n
    #         neigh = neighbors[:, i]
    #         grad = getindex(grid, neigh...)
    #         @. dot_prod_tmp = (scaled_indices - neigh) * grad
    #         value += sum(dot_prod_tmp) * weights[i]
    #     end
    #     # distances = map(x -> scaled_indices .- x, neighbors)
    #     # dot_products = dot.(map(x -> getindex(grid, x...), neighbors), distances)
    #     # print(dot_products)
    #     # value = dot(dot_products, weights)
    #     setindex!(A, value, indices...)
    # end

    perlin_!(A, grid, weights, neighbors, scaled_indices, dot_prod_tmp, w, n, ranges, it_indices)

    A
end

function perlin_!(A, grid, weights, neighbors, scaled_indices, dot_prod_tmp, w, n, ranges, it_indices)
    for indices in it_indices
        scaled_indices .= remap.(indices, ranges[1], ranges[2])
        @. w = fade(scaled_indices - floor(scaled_indices))
        weights = lerp_weights(w)
        neighbors .= adjacent(scaled_indices)
        value = 0
        for i in 1:2^n
            @views neigh = neighbors[:, i]
            grad = getindex(grid, neigh[1], neigh[2])
            @. dot_prod_tmp = (scaled_indices - neigh) * grad
            value += sum(dot_prod_tmp) * weights[i]
        end
        # distances = map(x -> scaled_indices .- x, neighbors)
        # dot_products = dot.(map(x -> getindex(grid, x...), neighbors), distances)
        # print(dot_products)
        # value = dot(dot_products, weights)
        setindex!(A, value, indices...)
    end
end

# function perlin!(A::AbstractArray, grid::AbstractArray)
#     resolution, scale = size(A), size(grid)
#     ranges = map.(Ref(x -> (1, x)), [resolution, scale])
#     for indices in collect(Base.Iterators.product(range.(Ref(1), resolution, step=1)...))
#         scaled_indices = remap.(indices, ranges...)
#         subgrid_position = (scaled_indices .- floor.(scaled_indices))
#         w1, w2 = fade.(subgrid_position)
#         weights = (1 - w1 - w2 + w1 * w2, w1 - w1 * w2, w2 - w1 * w2, w1 * w2)
#         neighbors = adjacent(scaled_indices)
#         value = 0
#         for (neigh, weight) in zip(neighbors, weights)
#             grad = getindex(grid, neigh...)
#             value += dot(scaled_indices .- neigh, grad) * weight
#         end
#         # distances = map(x -> scaled_indices .- x, neighbors)
#         # dot_products = dot.(map(x -> getindex(grid, x...), neighbors), distances)
#         # print(dot_products)
#         # value = dot(dot_products, weights)
#         setindex!(A, value, indices...)
#     end

#     A
# end

function perlin(resolution, grid::AbstractArray)
    noise = zeros(resolution...)
    perlin!(noise, grid)
end

function create_gradients(scale)
    grid = zeros(scale...)
    grid = map(x -> normalize(rand(length(scale))), grid)
end

function perlin(resolution, scale)
    grid = create_gradients(scale)
    perlin(resolution, grid)
end

function perlin2d!(A::AbstractArray, grid::AbstractArray)
    resolution, scale = size(A), size(grid)
    for i in 1:resolution[1]
        for j in 1:resolution[2]
            x, y = remap.((i, j), map.(Ref(x -> (1, x)), [resolution, scale])...)
            x0 = Int(floor(x))
            x1 = Int(ceil(x))
            y0 = Int(floor(y))
            y1 = Int(ceil(y))
            sx = fade(x - x0)
            sy = fade(y - y0)

            n0 = dot_grid_gradient(grid, x0, y0, x, y)
            n1 = dot_grid_gradient(grid, x1, y0, x, y)
            ix0 = lerp([n0, n1], sx)

            n2 = dot_grid_gradient(grid, x0, y1, x, y)
            n3 = dot_grid_gradient(grid, x1, y1, x, y)
            ix1 = lerp([n2, n3], sx)
            # println((n0, n1, n2, n3))

            A[i, j] = lerp([ix0, ix1], sy)
        end
    end
    A
end

function dot_grid_gradient(grid, i, j, x, y)
    dx = x - i
    dy = y - j
    dot(grid[i, j], [dx, dy])
end
lerp(vals, w) = vals[1] * (1 - w) + vals[2] * w
fade(x) = 6x^5 - 15x^4 + 10x^3


using CUSOLVER
using CUDArt
using Base.Test

m = 15
n = 10
k = 1

@testset for elty in [Float32, Float64, Complex64, Complex128]
    @testset "potrf!" begin
        A    = rand(elty,n,n)
        A    = A*A' #posdef
        d_A  = CudaArray(A)
        d_A  = CUSOLVER.potrf!('U',d_A)
        h_A  = to_host(d_A)
        cA,_ = LAPACK.potrf!('U',A)
        @test h_A ≈ cA
        A    = rand(elty,m,n)
        d_A  = CudaArray(A)
        @test_throws DimensionMismatch CUSOLVER.potrf!('U',d_A)
        A    = zeros(elty,n,n)
        d_A  = CudaArray(A)
        @test_throws Base.LinAlg.SingularException CUSOLVER.potrf!('U',d_A)
    end
    @testset "potrs!" begin
        A     = rand(elty,n,n)
        A     = A*A' #posdef
        d_A   = CudaArray(A)
        d_A   = CUSOLVER.potrf!('U',d_A)
        h_A   = to_host(d_A)
        B     = rand(elty,n,n)
        d_B   = CudaArray(B)
        d_B   = CUSOLVER.potrs!('U',d_A,d_B)
        h_B   = to_host(d_B)
        @test h_B ≈ LAPACK.potrs!('U',h_A,B)
        A    = rand(elty,m,n)
        d_A  = CudaArray(A)
        @test_throws DimensionMismatch CUSOLVER.potrs!('U',d_A,d_B)
        A    = rand(elty,n,n)
        d_A  = CudaArray(A)
        B     = rand(elty,m,m)
        d_B   = CudaArray(B)
        @test_throws DimensionMismatch CUSOLVER.potrs!('U',d_A,d_B)
    end

    @testset "getrf!" begin
        A          = rand(elty,m,n)
        d_A        = CudaArray(A)
        d_A,d_ipiv = CUSOLVER.getrf!(d_A)
        h_A        = to_host(d_A)
        h_ipiv     = to_host(d_ipiv)
        alu        = Base.LinAlg.LU(h_A, convert(Vector{Int},h_ipiv), zero(Int))
        @test A ≈ full(alu)
        A    = zeros(elty,n,n)
        d_A  = CudaArray(A)
        @test_throws Base.LinAlg.SingularException CUSOLVER.getrf!(d_A)
    end

    @testset "getrs!" begin
        A          = rand(elty,n,n)
        d_A        = CudaArray(A)
        d_A,d_ipiv = CUSOLVER.getrf!(d_A)
        B          = rand(elty,n,n)
        d_B        = CudaArray(B)
        d_B        = CUSOLVER.getrs!('N',d_A,d_ipiv,d_B)
        h_B        = to_host(d_B)
        @test h_B  ≈ A\B
        A          = rand(elty,m,n)
        d_A        = CudaArray(A)
        @test_throws DimensionMismatch CUSOLVER.getrs!('N',d_A,d_ipiv,d_B)
        A          = rand(elty,n,n)
        d_A        = CudaArray(A)
        B          = rand(elty,m,n)
        d_B        = CudaArray(B)
        @test_throws DimensionMismatch CUSOLVER.getrs!('N',d_A,d_ipiv,d_B)
    end

    @testset "geqrf!" begin
        A         = rand(elty,m,n)
        d_A       = CudaArray(A)
        d_A,d_tau = CUSOLVER.geqrf!(d_A)
        h_A       = to_host(d_A)
        h_tau     = to_host(d_tau)
        qra       = Base.LinAlg.QR(h_A, h_tau)
        @test A ≈ full(qra)
    end

    @testset "ormqr!" begin
        A         = rand(elty,n,n)
        d_A       = CudaArray(A)
        d_A,d_tau = CUSOLVER.geqrf!(d_A)
        h_A       = to_host(d_A)
        h_tau     = to_host(d_tau)
        qra       = Base.LinAlg.QR(h_A, h_tau)
        B         = rand(elty,n,n)
        d_B       = CudaArray(B)
        d_B       = CUSOLVER.ormqr!('L','N',d_A,d_tau,d_B)
        h_B       = to_host(d_B)
        @test h_B ≈ qra[:Q]*B
    end

    @testset "sytrf!" begin
        A          = rand(elty,n,n)
        A          = A + A' #symmetric
        d_A        = CudaArray(A)
        d_A,d_ipiv = CUSOLVER.sytrf!('U',d_A)
        h_A        = to_host(d_A)
        h_ipiv     = to_host(d_ipiv)
        A, ipiv    = LAPACK.sytrf!('U',A)
        @test ipiv == h_ipiv
        @test A ≈ h_A
        A    = rand(elty,m,n)
        d_A  = CudaArray(A)
        @test_throws DimensionMismatch CUSOLVER.sytrf!('U',d_A)
        A    = zeros(elty,n,n)
        d_A  = CudaArray(A)
        @test_throws Base.LinAlg.SingularException CUSOLVER.sytrf!('U',d_A)
    end

    @testset "gebrd!" begin
        A                             = rand(elty,m,n)
        d_A                           = CudaArray(A)
        d_A, d_D, d_E, d_TAUQ, d_TAUP = CUSOLVER.gebrd!(d_A)
        h_A                           = to_host(d_A)
        h_D                           = to_host(d_D)
        h_E                           = to_host(d_E)
        h_TAUQ                        = to_host(d_TAUQ)
        h_TAUP                        = to_host(d_TAUP)
        A,d,e,q,p                     = LAPACK.gebrd!(A)
        #@test A ≈ h_A
        @test d ≈ h_D
        @test e ≈ h_E
        @test q ≈ h_TAUQ
        @test p ≈ h_TAUP
    end

    @testset "gesvd!" begin
        A              = rand(elty,m,n)
        d_A            = CudaArray(A)
        d_U, d_S, d_Vt = CUSOLVER.gesvd!('A','A',d_A)
        h_S            = to_host(d_S)
        h_U            = to_host(d_U)
        h_Vt           = to_host(d_Vt)
        svda           = svdfact(A,thin=false)
        @test h_U ≈ svda[:U]
        @test h_S ≈ svdvals(A)
        @test h_Vt ≈ svda[:Vt]
    end
end


# An implementation of DIOM for the solution of the square linear system Ax = b.
#
# This method is described in
#
# Y. Saad, Iterative methods for sparse linear systems.
# PWS Publishing Company, Boston, USA, 1996.
#
# Y. Saad, Practical use of some krylov subspace methods for solving indefinite and nonsymmetric linear systems.
# SIAM journal on scientific and statistical computing, 5(1), pp. 203--228, 1984.
#
# Alexis Montoison, <alexis.montoison@polymtl.ca>
# Montreal, September 2018.

export diom

"""Solve the consistent linear system Ax = b using direct incomplete orthogonalization method.

DIOM is similar to CG with partial reorthogonalization.

An advantage of DIOM is that nonsymmetric or symmetric indefinite or both nonsymmetric
and indefinite systems of linear equations can be handled by this single algorithm.

This implementation allows a right preconditioner M.
"""
function diom(A :: AbstractLinearOperator, b :: AbstractVector{T};
              M :: AbstractLinearOperator=opEye(size(A,1)),
              atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6,
              itmax :: Int=0, memory :: Int=20, verbose :: Bool=false) where T <: Number

  m, n = size(A)
  m == n || error("System must be square")
  length(b) == m || error("Inconsistent problem size")
  verbose && @printf("DIOM: system of size %d\n", n)

  # Initial solution x₀ and residual r₀.
  x = zeros(T, n)
  x_old = copy(x)
  # Compute β.
  rNorm = @knrm2(n, b) # rNorm = ‖r₀‖₂
  rNorm == 0.0 && return x, SimpleStats(true, false, [rNorm], [], "x = 0 is a zero-residual solution")

  iter = 0
  itmax == 0 && (itmax = 2*n)

  rNorms = [rNorm;]
  ε = atol + rtol * rNorm
  verbose && @printf("%5d  %7.1e\n", iter, rNorm)

  # Set up workspace.
  mem = min(memory, itmax) # Memory.
  V = [zeros(n) for i = 1 : mem] # Preconditioned Krylov vectors, orthogonal basis for {r₀, AM⁻¹r₀, (AM⁻¹)²r₀, ..., (AM⁻¹)ᵐ⁻¹r₀}.
  P = [zeros(n) for i = 1 : mem] # Directions for x : Pₘ = Vₘ(Uₘ)⁻¹.
  H = zeros(mem+2)  # Last column of the band hessenberg matrix Hₘ = LₘUₘ.
  # Each column has at most mem + 1 nonzero elements. hᵢ.ₘ is stored as H[m-i+2].
  # m-i+2 represents the indice of the diagonal where hᵢ.ₘ is located.
  # In addition of that, the last column of Uₘ is stored in H.
  L = zeros(mem) # Last mem Pivots of Lₘ.
  p = BitArray(undef, mem) # Last mem permutations.

  # Initial ξ₁ and V₁.
  ξ = rNorm
  @. V[1] = b / rNorm

  # Stopping criterion.
  solved = rNorm ≤ ε
  tired = iter ≥ itmax
  status = "unknown"

  while !(solved || tired)

    # Update iteration index.
    iter = iter + 1

    # Set position in circulars stacks.
    pos = mod(iter-1, mem) + 1 # Position corresponding to pₘ and vₘ in circular stacks P and V.
    next_pos = mod(iter, mem) + 1 # Position corresponding to vₘ₊₁ in the circular stack V.

    # Incomplete Arnoldi procedure.
    z = M * V[pos] # Forms pₘ
    w = A * z # Forms vₘ₊₁
    for i = max(1, iter-mem+1) : iter
      ipos = mod(i-1, mem) + 1 # Position corresponding to vᵢ in the circular stack V.
      diag = iter - i + 2
      H[diag] = @kdot(n, w, V[ipos]) # hᵢ.ₘ = < A * M⁻¹vₘ , vᵢ >
      @kaxpy!(n, -H[diag], V[ipos], w) # w ← w - hᵢ.ₘ * vᵢ
    end
    # Compute hₘ₊₁.ₘ and vₘ₊₁.
    H[1] = @knrm2(n, w) # hₘ₊₁.ₘ = ‖vₘ₊₁‖₂
    if H[1] ≠ 0.0 # hₘ₊₁.ₘ = 0 ⇒ "lucky breakdown"
      @. V[next_pos] = w / H[1] # vₘ₊₁ = w / hₘ₊₁.ₘ
    end
    # It's possible that uₘ₋ₘₑₘ.ₘ ≠ 0 when m ≥ mem + 1
    if iter ≥ mem + 2
      H[mem+2] = 0.0 # hₘ₋ₘₑₘ.ₘ = 0
    end

    # Update the LU factorization with partial pivoting of H.
    # Compute the last column of Uₘ.
    if iter ≥ 2
      for i = max(2,iter-mem+1) : iter
        lpos = mod(i-1, mem) + 1 # Position corresponding to lᵢ.ᵢ₋₁ in the circular stack L.
        diag = iter - i + 2
        next_diag = diag + 1
        if p[lpos]
          # The rows i-1 and i are permuted.
          H[diag], H[next_diag] = H[next_diag], H[diag]
        end
        # uᵢ.ₘ ← hᵢ.ₘ - lᵢ.ᵢ₋₁ * uᵢ₋₁.ₘ
        H[diag] = H[diag] - L[lpos] * H[next_diag]
      end
      # Compute ξₘ the last component of zₘ = β(Lₘ)⁻¹e₁.
      if !p[pos] # p[pos] ⇒ ξₘ = ξₘ₋₁
        # ξₘ = -lₘ.ₘ₋₁ * ξₘ₋₁
        ξ = - L[pos] * ξ
      end
    end

    # Compute the direction pₘ, the last column of Pₘ = Vₘ(Uₘ)⁻¹.
    for i = max(1,iter-mem) : iter-1
      ipos = mod(i-1, mem) + 1 # Position corresponding to pᵢ in the circular stack P.
      diag = iter - i + 2
      if ipos == pos
        # pₐᵤₓ ← -hₘ₋ₘₑₘ.ₘ * pₘ₋ₘₑₘ
        @kscal!(n, -H[diag], P[pos])
      else
        # pₐᵤₓ ← pₐᵤₓ - hᵢ.ₘ * pᵢ
        @kaxpy!(n, -H[diag], P[ipos], P[pos])
      end
    end
    # pₐᵤₓ ← pₐᵤₓ + M⁻¹vₘ
    @kaxpy!(n, 1.0, z, P[pos])

    # Determine if interchange between hₘ₊₁.ₘ and uₘ.ₘ is needed and compute next pivot lₘ₊₁.ₘ.
    if abs(H[2]) < H[1]
      p[next_pos] = true
      # pₘ = pₐᵤₓ / hₘ₊₁.ₘ
      @. P[pos] = P[pos] / H[1]
      # lₘ₊₁.ₘ = uₘ.ₘ / hₘ₊₁.ₘ
      L[next_pos] = H[2] / H[1]
    else
      p[next_pos] = false
      # pₘ = pₐᵤₓ / uₘ.ₘ
      @. P[pos] = P[pos] / H[2]
      # lₘ₊₁.ₘ = hₘ₊₁.ₘ / uₘ.ₘ
      L[next_pos] = H[1] / H[2]
    end

    # Compute solution xₘ.
    if p[pos]
      # xₘ = xₘ₋ₙ + ξₘ₋ₙ * pₘ
      # x_old = xₘ₋ₙ, with m-n is the last iteration without permutation at the next step
      @. x = x_old + ξ * P[pos]
    else
      # xₘ = xₘ₋₁ + ξₘ * pₘ
      @kaxpy!(n, ξ, P[pos], x)
    end

    # Update x_old and residual norm.
    if !p[next_pos]
      copyto!(x_old, x)
      # ‖ b - Axₘ ‖₂ = hₘ₊₁.ₘ * |ξₘ / uₘ.ₘ| without pivoting
      rNorm = H[1] * abs(ξ / H[2])
    else
      # ‖ b - Axₘ ‖₂ = hₘ₊₁.ₘ * |ξₘ / hₘ₊₁.ₘ| = |ξₘ| with pivoting
      rNorm = abs(ξ)
    end
    push!(rNorms, rNorm)

    # Update stopping criterion.
    solved = rNorm ≤ ε
    tired = iter ≥ itmax
    verbose && @printf("%5d  %7.1e\n", iter, rNorm)
  end
  verbose && @printf("\n")

  status = tired ? "maximum number of iterations exceeded" : "solution good enough given atol and rtol"
  stats = SimpleStats(solved, false, rNorms, T[], status)
  return (x, stats)
end


#module PowerFlowTests
include("../src/PowerFlowAnalysis.jl")

using Base.Test
using JuMP
using Ipopt
using PowerModels
using PowerFlowAnalysis

case6ww = PowerModels.parse_file("data/case6ww.m")
pm = ACPPowerModel(case6ww)
sol = run_ac_pf(case6ww, IpoptSolver())

net = Network(pm)
sys = PowerFlowSystem(pm, sol)
lf = LurieForm(sys)

known_pairs = [(1,4), (2,5), (3,6)] # lines in the orignal file
@testset "Network topology" for pair in known_pairs
  f, t = pair
  @test (f,t) in net.pairs
  pidx = net.pair_idx[(f,t)]
  @test net.from[pidx] == net.bus_idx[f]
  @test net.to[pidx]   == net.bus_idx[t]
end

@testset "Admittance matrix" begin
  Ymp = Dict(
      (5,5) => 5.293290152818541 - 14.137764902337823im,
      (3,6) => -1.923076923076923 + 9.615384615384615im,
      (1,3) => 0.0im
    ) # Several values extracted from MATPOWER
  for ((f,t), y) in Ymp
    @test y ≈ sys.Y[net.bus_idx[f], net.bus_idx[t]]
  end
end

p, q = calc_pq(sys.Y, sys.v⁰, sys.ϑ⁰)
@testset "Power flow solutions" begin
  @test p[sys.pq] ≈ sys.p⁰[sys.pq]
  @test p[sys.pv] ≈ sys.p⁰[sys.pv]
  @test q[sys.pq] ≈ sys.q⁰[sys.pq]
end

@testset "Lurie form" begin
  v¹ = (1 + 0.1*randn(sys.nb)).*sys.v⁰
  ϑ¹ = (1 + 0.1*randn(sys.nb)).*sys.ϑ⁰
  f¹ = calc_f(v¹, ϑ¹, lf)
  s¹ = lf.S*f¹
  p¹, q¹ = calc_pq(sys.Y, v¹, ϑ¹)

  @test p¹[sys.pv] ≈ s¹[1:sys.npv]
  @test p¹[sys.pq] ≈ s¹[(sys.npv + 1):(sys.npv + sys.npq)]
  @test q¹[sys.pq] ≈ s¹[(sys.npv + sys.npq + 1):(sys.npv + 2*sys.npq)]
end

@testset "Jacobian ∂f/∂x" begin
  dfdx = PowerFlowAnalysis.calc_dfdx(lf)
  dx = 1e-4*randn(size(lf.x⁰))
  v¹,ϑ¹ = PowerFlowAnalysis.calc_vϑ(lf.x⁰ + dx, lf)
  df = PowerFlowAnalysis.calc_f(v¹, ϑ¹, lf) - lf.f⁰
  @test isapprox(df, dfdx*dx; rtol = 1e-3)
end


module Distances

using LinearAlgebra
using Statistics, StatsBase

export euclidean, jensen_shannon

"""
    euclidean²(X)

Compute the pairwise squared distances between points `X` according to the Euclidean metric.
Assumes `X` is sized ``d \\times N`` where ``d`` and ``N`` denote dimensionality and cardinality respectively.
"""
function euclidean²(X)
    dotprod = X'*X
    vecnorm = vec(diag(dotprod))

    return (vecnorm' .+ vecnorm) .- 2*dotprod
end

"""
    euclidean(X)

Compute the pairwise distances between points `X` according to the Euclidean metric.
Assumes `X` is sized ``d \\times N`` where ``d`` and ``N`` denote dimensionality and cardinality respectively.
"""
euclidean(X) = .√euclidean²(X)

"""
    kullback_liebler(p, q)
Compute the pairwise distances between probability distributions `p` and `q` according to the Kullback-Liebler divergence.
Assumes `p` and `q` are normalized.
"""
kullback_liebler(p, q) = sum(x .* log(x ./ y) for (x,y) ∈ zip(p,q) if x > 0 && y > 0)

"""
    jensen_shannon(P)
Compute the pairwise distances between probability distributions `P` according to the Jensen-Shannon divergence.
Assumes `P` is sized ``d \\times N`` where ``d`` and ``N`` denote dimensionality and cardinality respectively.
"""
function jensen_shannon(P)
    D = zeros(size(P,2), size(P,2))
    Threads.@threads for i in 1:(size(P,2)-1)
        for j in (i+1):size(P,2)
            M = (P[:,i] + P[:,j]) / 2
            D[i,j] = (kullback_liebler(P[:,i],M) + kullback_liebler(P[:,j],M)) / 2
            D[j,i] = D[i,j]
        end
    end

    D = .√D
    return D / mean(D[:])
end

end


# First-order vector Stochastic Differential Equations

# NOTE: for scalar SDEs, checkstable can also be used
"""isstableode(A) -> Bool

Determine whether the matrix A corresponds to a stable ordinary differential equation
"""
isstableode(A) = all(real.(eigen(A).values).<0)

"""isdiagonalizable(A) -> Bool

Determine whether the matrix A is diagonalizable
"""
isdiagonalizable(A) = rank(eigen(A).vectors)==size(A,1)

abstract type TimeSeriesModel{T<:Number} end

abstract type AbstractSde{T<:Number} end

"""
Vector valued, first order stochastic differential equation with measurement model
$FIELDS
"""
struct VectorSde{T<:Number,U<:AbstractMatrix{T}} <: AbstractSde{T}
  A::U
  C::Matrix{T}
  Q::Matrix{T}
  R::Matrix{T}
  function VectorSde(A::U,C,Q,R) where {T,U<:AbstractMatrix{T}}
    p = LinearAlgebra.checksquare(A)
    q = size(C,1)
    size(C,2)==p || error("C must be q×",p)
    LinearAlgebra.checksquare(Q)==p || error("Q must be ",p,"×",p)
    LinearAlgebra.checksquare(R)==q || error("R must be ",q,"×",q)
    new{T,U}(A,C,Q,R)
  end
end

function Base.convert(::Type{AbstractSde{T}},x::VectorSde) where {T}
  return VectorSde(convert.(AbstractMatrix{T},[x.A,x.C,x.Q,x.R])...)
end

"State dimension"
statedim(m::VectorSde) = size(m.A,1)

"measurement dimension"
measdim(m::VectorSde) = size(m.R,1)

"Model order"
order(eq::VectorSde) = 1

isstable(eq::VectorSde) = isstableode(eq.A)

"""
    statcovz(m::VectorSde) -> R0

Stationary covariance matrix of the state variable `z`
"""
statcovz(m::VectorSde) = project_symmetric(lyap(m.A,m.Q))

"""
$SIGNATURES

Stationary covariance matrix of the observations `y`
"""
statcov(m::VectorSde) = m.C*statcovz(m)*m.C'+m.R

"""
Stationary state distribution
$SIGNATURES
"""
statstate(eq::AbstractSde) = State(statcovz(eq))

"""
    diag_vectorsde(m::VectorSde) -> VectorSde

Representation of the  vector SDE `m` with diagonalized `A` matrix.
Only if `scale` is `true, scale eigenvectors that the first component
is one.

NOTE: In general `diag_vectorsde` returns complex-valued models. Not all
functions work on complex-valued models, e.g. `rand`.
"""
function diag_vectorsde(m::VectorSde,scale=true)
  E = eigen(m.A)
  U = E.vectors
  if scale
    for j = 1:size(U,2)
      U[:,j] ./= U[1,j]
    end
  end
  A = Diagonal(E.values)
  C = m.C*U
  Q = U\m.Q/U'
  R = complex(m.R)
  VectorSde(A,C,Q,R)
end

function LinearAlgebra.lyap(A::Diagonal,C)
  a = A.diag
  p = length(a)
  ac = repeat(a,1,p)
  m = ac.+ac'
  -C./m
end

"Observation time and value"
struct Obs
  t::Float64
  y::Float64
end

## SDE process ————————————————————
"""
SDE process: Combines Stochastic Differential Equation with state
$FIELDS
"""
struct SdeProcess{T,U} <: TimeSeriesModel{T}
  "Stochastic Differential Equation"
  eq::U
  "State"
  state::State{T}
  function SdeProcess(eq::U,state::State{T}) where {T,U}
    if statedim(eq)≠dim(state) error("Model order mismatch between SDE and state") end
    new{T,U}(eq,state)
  end
end

SdeProcess(eq) = SdeProcess(eq,statstate(eq))

function Base.show(io::IO,proc::SdeProcess)
  print(io,proc.eq,"; state : ")
  if has_statstate(proc)
    print(io,"stationary")
  else
    print(io,proc.state)
  end
end

VectorSdeProcess = SdeProcess{T,U} where {T,U<:VectorSde}

function describe_model(proc::SdeProcess)
  describe_model(proc.eq)
  print("State: ")
  if has_statstate(proc)
    println("Stationary")
  else
    println(proc.state)
  end
end

order(proc::SdeProcess) = order(proc.eq)

Statistics.std(proc::SdeProcess) = std(proc.eq)

"Has stationary state?"
has_statstate(proc::SdeProcess;rtol=√eps()) = isstable(proc.eq) && isapprox(proc.state,statstate(proc.eq);rtol=rtol)

function Statistics.mean(procs::Vector{<:SdeProcess})
  eq = mean(getproperty.(procs,:eq))
  state = mean(getproperty.(procs,:state))
  return SdeProcess(eq,state)
end


using Test, Econometrics, CSV, RDatasets, LinearAlgebra

# Linear Models
@testset "Linear Models" begin
	@testset "Balanced Panel Data" begin
    # Balanced Panel Data
    data = dataset("Ecdat", "Crime")
    # Pooling
    model = fit(EconometricModel,
                @formula(CRMRTE ~ PrbConv + AvgSen + PrbPris),
                data)
    @test nullloglikelihood(model) ≈ 1633.305 rtol = 1e-3
    @test loglikelihood(model) ≈ 1643.304 rtol = 1e-3
    @test dof(model) == 5
    # This corresponds to BIC from R's lm, Stata uses rank(vcov(model))
    @test bic(model) ≈ -3254.379 rtol = 1e-3
    @test coef(model) ≈ [0.0186413, -0.0011647, 0.0002362, 0.0273394] rtol = 1e-5
    @test vcov(model) ≈ [  1.858177e-05 -1.620480e-07 -6.464872e-07 -2.860941e-05
                          -1.620480e-07  1.781365e-07 -1.715589e-09  1.286486e-07
                          -6.464872e-07 -1.715589e-09  7.193971e-08  8.182148e-09
                          -2.860941e-05  1.286486e-07  8.182148e-09  6.685387e-05
                        ] rtol = 1e-6
    @test confint(model) ≈ [ 0.0101761950  0.0271063941
                            -0.0019935581 -0.0003358998
                            -0.0002905263  0.0007628970
                             0.0112828681  0.0433959412
                           ] rtol = 1e-6
    # Between
    model = fit(EconometricModel,
                @formula(CRMRTE ~ PrbConv + AvgSen + PrbPris + between(County)),
                data)
    @test nullloglikelihood(model) ≈ 239.5638 rtol = 1e-3
    @test loglikelihood(model) ≈ 244.4479 rtol = 1e-3
    @test dof(model) == 5
    @test coef(model) ≈ [-0.00464339, -0.00354113, 0.00084805, 0.07302994] rtol = 1e-6
    @test vcov(model) ≈ [ 3.480883e-04 -5.985740e-06 -1.382250e-05 -5.104706e-04
                         -5.985740e-06  3.573593e-06 -1.212431e-08  8.538922e-06
                         -1.382250e-05 -1.212431e-08  1.356678e-06  3.953835e-06
                         -5.104706e-04  8.538922e-06  3.953835e-06  1.102622e-03
                        ] rtol = 1e-6
    # Using Stata's value, R's plm has confint incorrect
    @test confint(model) ≈ [-0.0417325 0.0324457
                            -0.0072991 0.0002169
                            -0.0014674 0.0031635
                             0.0070191 0.1390408
                           ] rtol = 1e-5
    # Within
    model = fit(EconometricModel,
                @formula(CRMRTE ~ PrbConv + AvgSen + PrbPris + absorb(County)),
                data)
    #
    @test_broken nullloglikelihood(model) ≈ 2277.488 rtol = 1e-3
    @test_broken loglikelihood(model) ≈ 1643.304 rtol = 1e-3
    xtreg_fe_i = fit(EconometricModel,
                     @formula(CRMRTE ~ PrbConv + AvgSen + PrbPris + absorb(County + Year)),
                     data)
    # Random Effects
    xtreg = fit(EconometricModel,
                @formula(CRMRTE ~ PrbConv + AvgSen + PrbPris + PID(County) + TID(Year)),
                data)
    # IV Pooling
    ivreg = fit(EconometricModel,
                @formula(CRMRTE ~ PrbConv + (AvgSen ~ PrbPris)),
                data)
    # IV Between
    ivreg_be = fit(EconometricModel,
                   @formula(CRMRTE ~ PrbConv + (AvgSen ~ PrbPris) + between(County)),
                   data)
    # IV Within
    ivreg_fe = fit(EconometricModel,
                   @formula(CRMRTE ~ PrbConv + (AvgSen ~ PrbPris) + absorb(County)),
                   data)
    ivreg_fe_i = fit(EconometricModel,
                     @formula(CRMRTE ~ PrbConv + (AvgSen ~ PrbPris) + absorb(County + Year)),
                     data)
    # IV Random
    ivreg = fit(EconometricModel,
                @formula(CRMRTE ~ PrbConv + (AvgSen ~ PrbPris) + PID(County) + TID(Year)),
                data)
  end
end
# Nominal Models
@testset "Nominal Models" begin
  data = joinpath(dirname(pathof(Econometrics)), "..", "data", "insure.csv") |>
	CSV.read |>
    (data -> data[[:insure, :age, :male, :nonwhite, :site]]) |>
    (data -> dropmissing!(data, disallowmissing = true)) |>
    (data -> categorical!(data, [:insure, :site]))
  model = fit(EconometricModel,
              @formula(insure ~ 1 + age + male + nonwhite + site),
              data,
              contrasts = Dict(:insure => DummyCoding(base = "Uninsure")))
  β, V = coef(model), vcov(model)
  @test β ≈ [1.286943, 0.0077961, -0.4518496, -0.2170589, 1.211563, 0.2078123,
             1.556656, -0.0039489, 0.1098438, 0.7577178, 1.324599, -0.3801756
             ] rtol = 1e-3
  @test V ≈ Hermitian(
  			[ 0.35084518 -0.00590456 -0.02519209 -0.01837512 -0.08995794 -0.08264268  0.29928937 -0.00510107 -0.02257270 -0.01585492 -0.07392512 -0.06808314
             -0.00590456  0.00013092 -0.00046225 -0.00028066  0.00041004  0.00040470 -0.00509454  0.00011354 -0.00039387 -0.00024122  0.00033847  0.00035236
             -0.02519209 -0.00046225  0.13504644  0.01565749  0.01009559  0.00933967 -0.02288847 -0.00039497  0.11365141  0.01369872  0.00945927  0.00802593
             -0.01837512 -0.00590456 -0.02519209  0.18116609  0.01886597 -0.02821348 -0.01638629 -0.00023179  0.01314666  0.15068111  0.01664999 -0.02345786
             -0.08995794 -0.00590456 -0.02519209 -0.01837512  0.22138224 0.064050660 -0.07438947  0.00035751  0.00845363  0.01594691  0.19895500  0.05189742
             -0.08264268  0.00040470  0.00933967 -0.02821348  0.06405066 0.134170290 -0.06796097  0.00034828  0.00826890 -0.02353016  0.05182727  0.11060487
              0.29928937 -0.00509454 -0.02288847 -0.01638629 -0.07438947 -0.06796097  0.35560782 -0.00601966 -0.02683161 -0.02105510 -0.08991231 -0.08165509
             -0.00510107  0.00011354 -0.00039497 -0.00023179  0.00035751  0.00034828 -0.00601966  0.00013455 -0.00047223 -0.00028680  0.00039990  0.0004179
             -0.02257270 -0.00039387  0.11365141  0.01314666  0.00845363  0.00826890 -0.02683161 -0.00047223  0.13336253  0.01625685  0.01101870  0.00922333
             -0.01585492 -0.00024122  0.01369872  0.15068111  0.01594691 -0.02353016 -0.02105510 -0.00028680  0.01625685  0.17604391  0.02099458 -0.03008391
             -0.07392512  0.00033847  0.00945927  0.01664999  0.19895500  0.05182727 -0.08991231  0.00039990  0.01101870  0.02099458  0.22070774  0.06273376
             -0.06808314  0.00035236  0.00802593 -0.02345786  0.05189742  0.11060487 -0.08165509  0.00041790  0.00922333 -0.03008391  0.06273376  0.13899387
            ]) rtol = 1e-3
  @test sqrt.(diag(V)) ≈ [0.5923219, 0.0114418, 0.3674867, 0.4256361, 0.4705127, 0.3662926,
  						  0.5963286, 0.0115994, 0.3651883, 0.4195759, 0.4697954, 0.3728188
						 ] rtol = 1e-4
end
# Ordinal Models
@testset "Ordinal Models" begin
  data = dataset("Ecdat", "Kakadu")[[:RecParks, :Sex, :Age, :Schooling]]
  data.RecParks = convert(Vector{Int}, data.RecParks)
  data.RecParks = levels!(categorical(data.RecParks, ordered = true), collect(1:5))
  model = fit(EconometricModel,
               @formula(RecParks ~ Age + Sex + Schooling),
               data,
               contrasts = Dict(:RecParks => DummyCoding(levels = collect(1:5)))
               )
  β, V = coef(model), vcov(model)
  @test β ≈ [0.009437926, -0.015143049, -0.103911316,
             -2.92391240, -1.549266200, -0.298963900, 0.6698249] rtol = 1e-4
  @test V ≈ [ 0.00000646 -0.0000007000  0.000013250 0.00031473	0.0003161358 0.000321276 0.0003270734
             -0.00000070  0.0071633800 -0.000184930 0.00294759	0.0029647442 0.002995366 0.0029904485
              0.00001325 -0.0001849300  0.000618730 0.00288365	0.0028548209 0.002791361 0.0027318928
              0.00031473  0.0029475900  0.002883650 0.03684071	0.0289239290 0.026986237 0.026498765
              0.00031614  0.0029647400  0.002854820 0.02892393	0.0293971651 0.027160433 0.0266006675
              0.000321276 0.0029953660  0.002791361 0.02698623	0.0271604331 0.027860143 0.0270590982
              0.0003270734 0.002990448  0.002731893 0.02649876	0.0266006675 0.027059098 0.0281002038
            ] rtol = 1e-3
  @test sqrt.(diag(V)) ≈ [0.002540789, 0.084636770, 0.024874381,
                          0.191939349, 0.171456015, 0.166913580, 0.167631154] rtol = 1e-4
  data = joinpath(dirname(pathof(Econometrics)), "..", "data", "auto.csv") |>
	CSV.read |>
    (data -> data[[:rep77, :foreign, :length, :mpg]]) |>
    (data -> dropmissing!(data, disallowmissing = true))
  data.rep77 = levels!(categorical(data.rep77; ordered = true),
                       ["Poor", "Fair", "Average", "Good", "Excellent"])
  model = fit(EconometricModel,
              @formula(rep77 ~ foreign + length + mpg),
              data)
  β, V = coef(model), vcov(model)
  @test β ≈ [2.89679875, 0.08282676, 0.23076532, 17.92728, 19.86486, 22.10311, 24.69193] atol = 1e-3
  @test V ≈ [0.62578335	0.0111044038 0.012321411 2.4421745	2.4836507	2.5750058	2.6957587
             0.0111044	0.0005186744 0.001222385 0.1243671	0.1258979	0.1284727	0.1319772
             0.01232141	0.0012223847 0.004974678 0.3328169	0.3368057	0.3430257	0.3549711
             2.44217451	0.1243670884 0.332816913 30.9544666	31.0279194	31.6005528	32.5099187
             2.48365067	0.1258979463 0.336805731 31.0279194	31.4396158	31.9986912	32.9170873
             2.57500582	0.1284726666 0.343025653 31.6005528	31.9986912	32.7044311	33.6189243
             2.69575871	0.1319772302 0.354971081 32.5099187	32.9170873	33.6189243	34.8533724
            ] rtol = 1e-2
  @test sqrt.(diag(V)) ≈ [0.79106469, 0.02277442, 0.07053140,
                          5.56367384, 5.60710405, 5.71877881, 5.90367448] rtol = 1e-2
end


using Pkg
cd(@__DIR__)
Pkg.activate("../"); Pkg.instantiate()

using Plots; gr()
using LaTeXStrings
using DataFrames
using CSV
using OrdinaryDiffEq
using LinearAlgebra
using DataInterpolations
using BenchmarkTools
using DiffEqFlux
using Optim
using Flux
using Zygote
using ModelingToolkit
using DataDrivenDiffEq

model_univ = FastChain(FastDense(6, 16, tanh),
                       FastDense(16, 2))
include("../optimization_results/neural_augmentation_results_u.jl")
include("../optimization_results/globaloptimizationresults.jl")

# Plot path
figpath = "../figs/";
# Linewidths and styles
LW1 = 2.5
LW2 = 1.5
LS1 = :solid
LS2 = :dot
LS3 = :dash
LC1 = :blue
LC2 = :red
LC3 = :green
LC4 = :black
LA1 = 1.0
LA2 = 0.7
LA3 = 0.3
;

df=CSV.read("../data/Lab-Helicopter_Experimental-data.csv")
ENV["COLUMNS"]=100 # change the number of IJulia columns from default 80 to 100
df[1:5,:]

u_th = df[:,1]
u_psi = df[:,2]
theta = df[:,3]
psi = df[:,4]
dt = df[1,5]
tm = 0:dt:(length(psi)-1)*dt
;

timeseriesdata = [theta';psi']
Nd = length(tm);

U_th = ConstantInterpolation(u_th,tm)
U_psi = ConstantInterpolation(u_psi,tm)
Theta = ConstantInterpolation(theta,tm)
Psi = ConstantInterpolation(psi,tm);
u_t = (th=U_th, psi=U_psi)
rad2deg2(t,u) = (t,u*180/pi)

#### Sparsify

# Create a Basis
@variables u[1:2]

# Array of potential basis functions
h = [cos.(u)...; sin.(u)...; u[1]^0; u[1]^2; u[2]^2; u[1]*u[2]]
basis_u = Basis(h, u)
println(basis_u)
# Create an optimizer for the SINDY problem
opt = SR3()
# Create the thresholds which should be used in the search process
λ = exp10.(-10:0.1:10)
# Target function to choose the results from; x = L0 of coefficients and L2-Error of the model
f_target(x, w) = iszero(x[1]) ? Inf : norm(w.*x, 2)

# Test on original data and without further knowledge
X = reduce(hcat,utimeseries)
DX = reduce(hcat,Ktimeseries)
Ψ = SInDy(X, DX, basis_u, λ, opt = opt, maxiter = 10000, f_target = f_target)
println(Ψ)
print_equations(Ψ)

###########
## Check the fit
###########

function helicopter_uode(dx,x,Ps,P0,u,p,t)
# P0 = (rc=rc,h=h,m=m,Jbc=Jbc,Js=Js,g=g,K_th_th=K_th_th,K_th_psi=K_th_psi,K_psi_th=K_psi_th,K_psi_psi=K_psi_psi,D_th=D_th,D_psi=D_psi);
    th = x[1]
    psi = x[2]
    w_th = x[3]
    w_psi = x[4]

    U = [u.th(t),u.psi(t)]
    K = Ψ.equations.f_(U,Ψ.parameters,t)

    #
    M11 = Ps[4]*P0.Jbc + Ps[3]*P0.m*((Ps[1]*P0.rc)^2+(Ps[2]*P0.h)^2)
    M22 = Ps[4]*P0.Jbc*cos(th)^2 + Ps[3]*P0.m*(Ps[1]*P0.rc*cos(th)-Ps[2]*P0.h*sin(th))^2 + Ps[5]*P0.Js
    #
    dLdth = -w_psi^2*(((Ps[4]*P0.Jbc)+(Ps[3]*P0.m)*((Ps[1]*P0.rc)^2-(Ps[2]*P0.h)^2))*cos(th)*sin(th)-(Ps[1]*P0.rc)*(Ps[2]*P0.h)*(cos(th)^2-sin(th)^2)) - (Ps[3]*P0.m)*(Ps[6]*P0.g)*((Ps[1]*P0.rc)*cos(th) - (Ps[2]*P0.h)*sin(th))
    #
    dp_psi_dth = -2*(((Ps[4]*P0.Jbc)+(Ps[3]*P0.m)*((Ps[1]*P0.rc)^2-(Ps[2]*P0.h)^2))*cos(th)*sin(th) + (Ps[3]*P0.m)*(Ps[1]*P0.rc)*(Ps[2]*P0.h)*(cos(th)^2-sin(th)^2))*w_psi
    #
    F_th = Ps[7]*P0.K_th_th*U[1] - Ps[8]*P0.K_th_psi*U[2] - P0.D_th*w_th + K[1]
    F_psi = Ps[9]*P0.K_psi_th*U[1] - Ps[10]*P0.K_psi_psi*U[2] - P0.D_psi*w_psi + K[2]

    #
    dx[1] = w_th
    dx[2] = w_psi
    dx[3] = (dLdth + F_th)/M11
    dx[4] = (-dp_psi_dth*w_th + F_psi)/M22
end

uode_f! = (dx,x,p,t) -> helicopter_uode(dx,x,Ps_best,P0,u_t,p,t)
prob = ODEProblem(uode_f!,x0_best,tspan)
sol_best = solve(prob,AutoTsit5(TRBDF2(autodiff=false)),reltol=1e-6)

model_univ_u = FastChain(FastDense(2, 16, tanh),
                       FastDense(16, 2))

function helicopter_uode(dx,x,Ps,P0,u,p,t)
# P0 = (rc=rc,h=h,m=m,Jbc=Jbc,Js=Js,g=g,K_th_th=K_th_th,K_th_psi=K_th_psi,K_psi_th=K_psi_th,K_psi_psi=K_psi_psi,D_th=D_th,D_psi=D_psi);
    th = x[1]
    psi = x[2]
    w_th = x[3]
    w_psi = x[4]

    U = [u.th(t),u.psi(t)]
    K = model_univ_u(U,p)
    #
    M11 = Ps[4]*P0.Jbc + Ps[3]*P0.m*((Ps[1]*P0.rc)^2+(Ps[2]*P0.h)^2)
    M22 = Ps[4]*P0.Jbc*cos(th)^2 + Ps[3]*P0.m*(Ps[1]*P0.rc*cos(th)-Ps[2]*P0.h*sin(th))^2 + Ps[5]*P0.Js
    #
    dLdth = -w_psi^2*(((Ps[4]*P0.Jbc)+(Ps[3]*P0.m)*((Ps[1]*P0.rc)^2-(Ps[2]*P0.h)^2))*cos(th)*sin(th)-(Ps[1]*P0.rc)*(Ps[2]*P0.h)*(cos(th)^2-sin(th)^2)) - (Ps[3]*P0.m)*(Ps[6]*P0.g)*((Ps[1]*P0.rc)*cos(th) - (Ps[2]*P0.h)*sin(th))
    #
    dp_psi_dth = -2*(((Ps[4]*P0.Jbc)+(Ps[3]*P0.m)*((Ps[1]*P0.rc)^2-(Ps[2]*P0.h)^2))*cos(th)*sin(th) + (Ps[3]*P0.m)*(Ps[1]*P0.rc)*(Ps[2]*P0.h)*(cos(th)^2-sin(th)^2))*w_psi
    #
    F_th = Ps[7]*P0.K_th_th*U[1] - Ps[8]*P0.K_th_psi*U[2] - P0.D_th*w_th + K[1]
    F_psi = Ps[9]*P0.K_psi_th*U[1] - Ps[10]*P0.K_psi_psi*U[2] - P0.D_psi*w_psi + K[2]

    #
    dx[1] = w_th
    dx[2] = w_psi
    dx[3] = (dLdth + F_th)/M11
    dx[4] = (-dp_psi_dth*w_th + F_psi)/M22
end

uode_f! = (dx,x,p,t) -> helicopter_uode(dx,x,Ps_best,P0,u_t,p,t)

tspan = (tm[1],tm[end])

#
prob = ODEProblem(uode_f!,x0_best,tspan,pnn)
#
sol_nn = solve(prob,AutoTsit5(TRBDF2(autodiff=false)),reltol=1e-6)

#
plot(sol_best,vars=(rad2deg2,0,1),lw=LW1,lc=LC1,label=L"$\theta$")
plot!(sol_nn,vars=(rad2deg2,0,1),lw=LW1,lc=LC3,label=L"$\theta$")
plot!(tm,theta*180/pi,lw=LW1,ls=LS2,lc=LC2,label=L"\theta^\mathrm{d}")
plot!(xlim = tspan,xlabel=L"time $t$ [s]", ylabel=L"$\theta$ [$^\circ$]")
plot!(title="Pitch angle: discovered (blue) vs. data (red) vs NN (green)",box=true)
figname="Helicopter_pitch-angle_model_fit_FNN_u_eq.svg"
savefig(figpath*figname)

plot(sol_best,vars=(rad2deg2,0,2),lw=LW1,lc=LC1,label=L"$\psi$")
plot!(sol_nn,vars=(rad2deg2,0,2),lw=LW1,lc=LC3,label=L"$\psi$")
plot!(tm,psi*180/pi,lw=LW1,ls=LS2,lc=LC2,label=L"\psi^\mathrm{d}")
plot!(xlim = tspan,xlabel=L"time $t$ [s]", ylabel=L"$\psi$ [$^\circ$]")
plot!(title="Yaw angle: discovered(blue) vs. data (red) vs NN (green)",box=true)
figname="Helicopter_yaw-angle_model_fit_FNN_u_eq.svg"
savefig(figpath*figname)

open("../optimization_results/equation_discovery_u.txt", "w") do f
    println(f,Ψ)
    println(f)
    print_equations(f,Ψ)
    println(f)
    print_equations(f,Ψ,show_parameter=true)
    println(f)
    write(f,"parameters = $(Ψ.parameters) \n")
end


using Distributions

# Unit-tests for the VAR estimation

lags = 3
series = 6
length = 1000

#
# # Function for simulating Vector autoregressive process with given coefficient matrix
#

# This function simulates Vector autoregressive process taking those values as input:
# - length: length of the resulting time series
# - typ: type of the process, whether it should or should not have a constant, etc.
# - coefficients: matrix of coefficients with dimensions (series, series*lags) with ordering in each line as x_lag1, y_lag1, x_lag2, y_lag2, etc...
# - other_c: matrix of coefficients corresponding to potential constants and trends with size (series, ...) where ... is 0, 1, or 2, according to type

function VARsim(length, typ, coefficients, other_c)
	# Should be enough
	burnout = 200

	a, b = size(coefficients)
	lags = convert(Int, max(a,b)/min(a,b))
	series = min(a,b)

	typ=="None" && (other=fill(0.0, length + burnout))
	typ=="Const" && (other=fill(1.0, length + burnout))
	typ=="Trend" && (other=[1:(length+burnout)])
	typ=="Const and trend" && (other=[fill(1.0, length + burnout) [1:(length+burnout)]])

	y = fill(0.0, length + burnout, series)
	eps = rand(Normal(0, 1/100), (length + burnout, series))

	for i=(lags+1):(size(y)[1])
			y[i,:] = mapreduce((x) -> y[i-x,:]*coefficients[(1:series) + (x-1)*series, :], +, 1:lags) + eps[i,:] + other[i,:]*other_c'
	end

	y = y[(burnout+1):end,:]

	return y
end

#
# # Function simulating random VAR of given specification
#
# This function is suited for performing a unit-test of the VAR system.
# Arguments:
# - lags: number of lags in the system
# - series: number of series in the system
# - length: length of the resulting time series
# - typ: whether the system should not have constant, trend or have any combination of those.

function VARdatarandsim(lags, series, length, typ)
	typ=="None" && (other_c=fill(0.0, series))
	typ=="Const" && (other_c=rand(Uniform(), series))
	typ=="Trend" && (other_c=rand(Uniform(), series))
	typ=="Const and trend" && (other_c=[rand(Uniform(), series) rand(Uniform(), series)])

	# Set them so that it does not blow up!
	coefficients = rand(Uniform(0, 1/(lags*series)), (series, lags*series))'
	y = VARsim(length, typ, coefficients, other_c)
	return (y, [other_c', coefficients])
end

apply(vcat,[coeficients[(1:series) + (x-1)*series, :] for x=[lags:-1:1]]) - varEstimate(y, lags, "None").C

mapreduce((x)->VARdifcoef(lags, series, length, burnout), +, 1:60)/60


"""

```
image(x,sides)
```

Move coordinates to minimum perioddic image.

"""
image(x,sides) = image_one.(x,sides)

# Translate a single coordinate according to periodic conditions
function image_one(x,side)
  xmod = mod(x,side)
  if xmod <= -side/2
    return xmod + side
  elseif xmod > side/2
    return xmod - side
  else
    return xmod
  end
end


# Load Julia packages (libraries)

using StanModels, CSV

df = convert(DataFrame,
  CSV.read(joinpath(@__DIR__, "..", "..", "data", "rugged.csv"), delim=';'));

dcc = filter(row -> !(ismissing(row[:rgdppc_2000])), df)
dcc[!, :log_gdp] = log.(dcc[!, :rgdppc_2000])
dcc[!, :cont_africa] = Array{Float64}(convert(Array{Int}, dcc[!, :cont_africa]))
dcc[!, :rugged] = convert(Array{Float64}, dcc[!, :rugged])

# Define the Stan language model

m8_1s = "
data{
    int N;
    vector[N] log_gdp;
    vector[N] cont_africa;
    vector[N] rugged;
    vector[N] rugged_cont_africa;
}
parameters{
    real a;
    real bR;
    real bA;
    real bAR;
    real sigma;
}
model{
    vector[N] mu = a + bR * rugged + bA * cont_africa + bAR * rugged_cont_africa;
    sigma ~ uniform( 0 , 10 );
    bAR ~ normal( 0 , 10 );
    bA ~ normal( 0 , 10 );
    bR ~ normal( 0 , 10 );
    a ~ normal( 0 , 100 );
    log_gdp ~ normal( mu , sigma );
}
";

# Define the Stanmodel and set the output format to :mcmcchains.

sm = SampleModel("m8.1s", m8_1s);

# Input data for cmdstan

m8_1_data = Dict(
  "N" => size(dcc, 1),
  "log_gdp" => dcc[!, :log_gdp],
  "cont_africa" => dcc[!, :cont_africa],
  "rugged" => dcc[!, :rugged],
  "rugged_cont_africa" => dcc[!, :rugged] .* dcc[!, :cont_africa] );

# Sample using cmdstan

(sample_file, log_file) = stan_sample(sm, data=m8_1_data);

# Result rethinking

rethinking = "
       mean   sd  5.5% 94.5% n_eff Rhat
a      9.22 0.14  9.00  9.46   282    1
bR    -0.21 0.08 -0.33 -0.08   275    1
bA    -1.94 0.24 -2.33 -1.59   268    1
bAR    0.40 0.14  0.18  0.62   271    1
sigma  0.96 0.05  0.87  1.04   339    1
"

# Describe the draws

if !(sample_file == nothing)
  chn = read_samples(sm)
  describe(chn)
end


using FIRITools
using Plots

const alt = FIRITools.ALTITUDE

# Compare mean and median profiles
day_mean = firi(chi=(0, 95))
day_med = FIRITools.quantile(0.5, chi=(0, 95))
day_30 = FIRITools.quantile(0.3, chi=(0, 95))
day_70 = FIRITools.quantile(0.7, chi=(0, 95))

plot(day_mean, alt/1000, label="mean", xscale=:log10, linewidth=2, ylims=(55, 110),
     legend=:topleft, title="Day", ylabel="altitude (km)", xlabel="Ne (m⁻³)")
plot!(day_med, alt/1000, label="median", xscale=:log10, linewidth=2)
plot!(day_30, alt/1000, label="30%", xscale=:log10)
plot!(day_70, alt/1000, label="70%", xscale=:log10)
savefig("day_mean.png")

night_mean = firi(chi=(100, 180))
night_med = FIRITools.quantile(0.5, chi=(100, 180))
night_30 = FIRITools.quantile(0.3, chi=(100, 180))
night_70 = FIRITools.quantile(0.7, chi=(100, 180))

plot(night_mean, alt/1000, label="mean", xscale=:log10, linewidth=2, ylims=(55, 110),
     legend=:topleft, title="Night", ylabel="altitude (km)", xlabel="Ne (m⁻³)")
plot!(night_med, alt/1000, label="median", xscale=:log10, linewidth=2)
plot!(night_30, alt/1000, label="30%", xscale=:log10)
plot!(night_70, alt/1000, label="70%", xscale=:log10)
savefig("night_mean.png")

# Compare mean and median around interpolation points
day_mean = firi(chi=(0, 95), lat=(30, 45))
day_med = FIRITools.quantile(0.5, chi=(0, 95), lat=(30, 45))
day_30 = FIRITools.quantile(0.3, chi=(0, 95), lat=(30, 45))
day_70 = FIRITools.quantile(0.7, chi=(0, 95), lat=(30, 45))

plot(day_mean, alt/1000, label="mean", xscale=:log10, linewidth=2, ylims=(55, 110),
     legend=:topleft, title="Day, limited latitude", ylabel="altitude (km)", xlabel="Ne (m⁻³)")
plot!(day_med, alt/1000, label="median", xscale=:log10, linewidth=2)
plot!(day_30, alt/1000, label="30%", xscale=:log10)
plot!(day_70, alt/1000, label="70%", xscale=:log10)
savefig("day_mean_latzoom.png")

day_mean = firi(chi=(30, 45))
day_med = FIRITools.quantile(0.5, chi=(30, 45))
day_30 = FIRITools.quantile(0.3, chi=(30, 45))
day_70 = FIRITools.quantile(0.7, chi=(30, 45))

plot(day_mean, alt/1000, label="mean", xscale=:log10, linewidth=2, ylims=(55, 110),
     legend=:topleft, title="Day, limited SZA", ylabel="altitude (km)", xlabel="Ne (m⁻³)")
plot!(day_med, alt/1000, label="median", xscale=:log10, linewidth=2)
plot!(day_30, alt/1000, label="30%", xscale=:log10)
plot!(day_70, alt/1000, label="70%", xscale=:log10)
savefig("day_mean_chizoom.png")

# Check to see if interpolation is acting as expected

real_lats = unique(FIRITools.HEADER[!,"Lat, deg"])
lats = 0:5:90
interp_lat = setdiff(lats, real_lats)
N = length(lats)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10, legend=:topleft,
         ylims=(55, 110), legendtitle="Latitude")
for la in real_lats
    prof = firi(45, la)
    plot!(p, prof, alt/1000, label=la, color=cmap[findfirst(isequal(la), lats)])
end
for la in interp_lat
    prof = firi(45, la)
    plot!(p, prof, alt/1000, label=nothing, color=cmap[findfirst(isequal(la), lats)],
          linestyle=:dash)
end
display(p)
savefig("lat.png")

real_sza = unique(FIRITools.HEADER[!,"Chi, deg"])
chis = 0:5:130
interp_sza = setdiff(chis, real_sza)
# interp_sza = [105, 110, 115, 120]
N = length(chis)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10,
         ylims=(55, 110), legend=:topleft, legendtitle="SZA")
for sza in real_sza
    prof = firi(sza, 45)
    plot!(p, prof, alt/1000, label=sza, color=cmap[findfirst(isequal(sza), chis)])
end

for sza in interp_sza
    prof = firi(sza, 45)
    plot!(p, prof, alt/1000, label=nothing, color=cmap[findfirst(isequal(sza), chis)],
          linestyle=:dash)
end
display(p)
savefig(p, "sza.png")

real_sza = [75, 80, 85, 90, 95, 100]
chis = 75:100
interp_sza = setdiff(chis, real_sza)
N = length(chis)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10,
         ylims=(55, 110), legend=:topleft, legendtitle="SZA")
for sza in real_sza
    prof = firi(sza, 45)
    plot!(p, prof, alt/1000, label=sza, color=cmap[findfirst(isequal(sza), chis)])
end

for sza in interp_sza
    prof = firi(sza, 45)
    plot!(p, prof, alt/1000, label=nothing, color=cmap[findfirst(isequal(sza), chis)],
          linestyle=:dash)
end
display(p)
savefig("sza_zoom.png")

###
# With extrapolation to ground

real_lats = unique(FIRITools.HEADER[!,"Lat, deg"])
lats = 0:5:90
interp_lat = setdiff(lats, real_lats)
N = length(lats)
cmap = palette(:rainbow, N, rev=true)
newalt = 0:110

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10, legend=:topleft,
         ylims=(0, 110), yticks=0:20:110, legendtitle="Latitude", xticks=exp10.([0, 3, 6, 9, 12]))
for la in real_lats
    prof = FIRITools.extrapolate(firi(45, la), newalt*1e3)
    plot!(p, prof, newalt, label=la, color=cmap[findfirst(isequal(la), lats)])
end
for la in interp_lat
    prof = FIRITools.extrapolate(firi(45, la), newalt*1e3)
    plot!(p, prof, newalt, label=nothing, color=cmap[findfirst(isequal(la), lats)],
          linestyle=:dash)
end
display(p)
savefig("lat_extrap.png")

real_sza = unique(FIRITools.HEADER[!,"Chi, deg"])
chis = 0:5:130
interp_sza = setdiff(chis, real_sza)
# interp_sza = [105, 110, 115, 120]
N = length(chis)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10,
         ylims=(0, 110), yticks=0:20:110, legend=:topleft, legendtitle="SZA", xticks=exp10.([0, 3, 6, 9, 12]))
for sza in real_sza
    prof = FIRITools.extrapolate(firi(sza, 45), newalt*1e3)
    plot!(p, prof, newalt, label=sza, color=cmap[findfirst(isequal(sza), chis)])
end

for sza in interp_sza
    prof = FIRITools.extrapolate(firi(sza, 45), newalt*1e3)
    plot!(p, prof, newalt, label=nothing, color=cmap[findfirst(isequal(sza), chis)],
          linestyle=:dash)
end
display(p)
savefig(p, "sza_extrap.png")

# Using median
real_sza = unique(FIRITools.HEADER[!,"Chi, deg"])
chis = 0:5:130
interp_sza = setdiff(chis, real_sza)
# interp_sza = [105, 110, 115, 120]
N = length(chis)
cmap = palette(:rainbow, N, rev=true)

p = plot(ylabel="altitude (km)", xlabel="Ne (m⁻³)", xscale=:log10,
         ylims=(0, 110), yticks=0:20:110, legend=:topleft, legendtitle="SZA", xticks=exp10.([0, 3, 6, 9, 12]))
for sza in real_sza
    prof = FIRITools.extrapolate(FIRITools.quantile(sza, 45, 0.5), newalt*1e3)
    plot!(p, prof, newalt, label=sza, color=cmap[findfirst(isequal(sza), chis)])
end

for sza in interp_sza
    prof = FIRITools.extrapolate(FIRITools.quantile(sza, 45, 0.5), newalt*1e3)
    plot!(p, prof, newalt, label=nothing, color=cmap[findfirst(isequal(sza), chis)],
          linestyle=:dash)
end
display(p)
savefig(p, "sza_median_extrap.png")


using Random
using Test
using FFTW
using CUDA

Random.seed!(0)

FT = Float64
ε = eps(FT)

# Test both even and odd sizes
for N in (4, 5)
    A_cpu = rand(FT, N)
    A_gpu = CuArray(A_cpu)

    B_cpu = FFTW.fft(A_cpu)
    B_gpu = CUDA.CUFFT.fft(A_gpu)

    @test isapprox.(B_cpu, Vector(B_gpu), atol=ε) |> all

    C_cpu = FFTW.ifft(B_cpu)
    C_gpu = CUDA.CUFFT.ifft(B_gpu)

    @test isapprox.(C_cpu, Vector(C_gpu), atol=ε) |> all

    @test isapprox.(A_cpu, C_cpu, atol=ε) |> all
    @test isapprox.(Vector(A_gpu), Vector(C_gpu), atol=ε) |> all
end


using DelimitedFiles
using StatsBase

input = joinpath(@__DIR__, "input")
fish = readdlm(input, ',', Int)

function advance(fish, days)
    for _ in 1:days
        next = Dict()
        for (age, n) in fish
            if age == 0
                haskey(next, 6) ? next[6] += n : next[6] = n
                next[8] = n
            else
                haskey(next, age - 1) ? next[age - 1] += n : next[age - 1] = n
            end
        end
        fish = next
    end
    return fish
end

fish = countmap(vec(fish))

p1 = sum(map(fishes -> last(fishes), collect(advance(fish, 80))))
p2 = sum(map(fishes -> last(fishes), collect(advance(fish, 256))))

println("-----------------------------------------------------------------------")
println("lanternfish -- part one :: $p1")
println("lanternfish -- part two :: $p2")
println("-----------------------------------------------------------------------")

@assert(p1 == 383160)
@assert(p2 == 1721148811504)


# This includes parts that were formerly a part of Julia. License is MIT: http://julialang.org/license
__precompile__()
module Primes

if VERSION >= v"0.5.0-dev+4340"
    if isdefined(Base,:isprime)
        import Base: isprime, primes, primesmask, factor
    else
        export isprime, primes, primesmask, factor
    end

    using Base: BitSigned
    using Base.Checked.checked_neg
else
    typealias BitSigned Union{Int128,Int16,Int32,Int64,Int8}
    function checked_neg(x::Integer)
        y = -x
        (y < 0) == (x < 0) && throw(OverflowError())
        return y
    end
end

# Primes generating functions
#     https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes
#     https://en.wikipedia.org/wiki/Wheel_factorization
#     http://primesieve.org
#     Jonathan Sorenson, "An analysis of two prime number sieves", Computer Science Technical Report Vol. 1028, 1991
const wheel         = [4,  2,  4,  2,  4,  6,  2,  6]
const wheel_primes  = [7, 11, 13, 17, 19, 23, 29, 31]
const wheel_indices = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 7, 7]

@inline function wheel_index(n)
    d, r = divrem(n - 1, 30)
    return 8d + wheel_indices[r + 2]
end
@inline function wheel_prime(n)
    d, r = (n - 1) >>> 3, (n - 1) & 7
    return 30d + wheel_primes[r + 1]
end

function _primesmask(limit::Int)
    limit < 7 && throw(ArgumentError("The condition limit ≥ 7 must be met."))
    n = wheel_index(limit)
    m = wheel_prime(n)
    sieve = ones(Bool, n)
    @inbounds for i = 1:wheel_index(isqrt(limit))
        if sieve[i]
            p = wheel_prime(i)
            q = p^2
            j = (i - 1) & 7 + 1
            while q ≤ m
                sieve[wheel_index(q)] = false
                q += wheel[j] * p
                j = j & 7 + 1
            end
        end
    end
    return sieve
end

function _primesmask(lo::Int, hi::Int)
    7 ≤ lo ≤ hi || throw(ArgumentError("The condition 7 ≤ lo ≤ hi must be met."))
    lo == 7 && return _primesmask(hi)
    wlo, whi = wheel_index(lo - 1), wheel_index(hi)
    m = wheel_prime(whi)
    sieve = ones(Bool, whi - wlo)
    hi < 49 && return sieve
    small_sieve = _primesmask(isqrt(hi))
    @inbounds for i = 1:length(small_sieve)  # don't use eachindex here
        if small_sieve[i]
            p = wheel_prime(i)
            j = wheel_index(2 * div(lo - p - 1, 2p) + 1)
            q = p * wheel_prime(j + 1)
            j = j & 7 + 1
            while q ≤ m
                sieve[wheel_index(q) - wlo] = false
                q += wheel[j] * p
                j = j & 7 + 1
            end
        end
    end
    return sieve
end

"""
    primesmask([lo,] hi)

Returns a prime sieve, as a `BitArray`, of the positive integers (from `lo`, if specified)
up to `hi`. Useful when working with either primes or composite numbers.
"""
function primesmask(lo::Int, hi::Int)
    0 < lo ≤ hi || throw(ArgumentError("The condition 0 < lo ≤ hi must be met."))
    sieve = falses(hi - lo + 1)
    lo ≤ 2 ≤ hi && (sieve[3 - lo] = true)
    lo ≤ 3 ≤ hi && (sieve[4 - lo] = true)
    lo ≤ 5 ≤ hi && (sieve[6 - lo] = true)
    hi < 7 && return sieve
    wheel_sieve = _primesmask(max(7, lo), hi)
    lsi = lo - 1
    lwi = wheel_index(lsi)
    @inbounds for i = 1:length(wheel_sieve)   # don't use eachindex here
        sieve[wheel_prime(i + lwi) - lsi] = wheel_sieve[i]
    end
    return sieve
end
primesmask{T<:Integer}(lo::T, hi::T) = lo ≤ hi ≤ typemax(Int) ? primesmask(Int(lo), Int(hi)) :
    throw(ArgumentError("Both endpoints of the interval to sieve must be ≤ $(typemax(Int)), got $lo and $hi."))

primesmask(limit::Int) = primesmask(1, limit)
primesmask(n::Integer) = n ≤ typemax(Int) ? primesmask(Int(n)) :
    throw(ArgumentError("Requested number of primes must be ≤ $(typemax(Int)), got $n."))

"""
    primes([lo,] hi)

Returns a collection of the prime numbers (from `lo`, if specified) up to `hi`.
"""
function primes(lo::Int, hi::Int)
    lo ≤ hi || throw(ArgumentError("The condition lo ≤ hi must be met."))
    list = Int[]
    lo ≤ 2 ≤ hi && push!(list, 2)
    lo ≤ 3 ≤ hi && push!(list, 3)
    lo ≤ 5 ≤ hi && push!(list, 5)
    hi < 7 && return list
    lo = max(2, lo)
    sizehint!(list, 5 + floor(Int, hi / (log(hi) - 1.12) - lo / (log(lo) - 1.12 * (lo > 7))) ) # http://projecteuclid.org/euclid.rmjm/1181070157
    sieve = _primesmask(max(7, lo), hi)
    lwi = wheel_index(lo - 1)
    @inbounds for i = 1:length(sieve)   # don't use eachindex here
        sieve[i] && push!(list, wheel_prime(i + lwi))
    end
    return list
end
primes(n::Int) = primes(1, n)

const PRIMES = primes(2^16)

"""
    isprime(x::Integer) -> Bool

Returns `true` if `x` is prime, and `false` otherwise.

```jldoctest
julia> isprime(3)
true
```
"""
function isprime(n::Integer)
    # Small precomputed primes + Miller-Rabin for primality testing:
    #     https://en.wikipedia.org/wiki/Miller–Rabin_primality_test
    #     https://github.com/JuliaLang/julia/issues/11594
    for m in (2, 3, 5, 7, 11, 13, 17, 19, 23)
        n % m == 0 && return n == m
    end
    n < 841 && return n > 1
    s = trailing_zeros(n - 1)
    d = (n - 1) >>> s
    for a in witnesses(n)::Tuple{Vararg{Int}}
        x = powermod(a, d, n)
        x == 1 && continue
        t = s
        while x != n - 1
            (t -= 1) ≤ 0 && return false
            x = oftype(n, widemul(x, x) % n)
            x == 1 && return false
        end
    end
    return true
end

"""
    isprime(x::BigInt, [reps = 25]) -> Bool

Probabilistic primality test. Returns `true` if `x` is prime with high probability (pseudoprime);
and `false` if `x` is composite (not prime). The false positive rate is about `0.25^reps`.
`reps = 25` is considered safe for cryptographic applications (Knuth, Seminumerical Algorithms).

```jldoctest
julia> isprime(big(3))
true
```
"""
isprime(x::BigInt, reps=25) = ccall((:__gmpz_probab_prime_p,:libgmp), Cint, (Ptr{BigInt}, Cint), &x, reps) > 0

# Miller-Rabin witness choices based on:
#     http://mathoverflow.net/questions/101922/smallest-collection-of-bases-for-prime-testing-of-64-bit-numbers
#     http://primes.utm.edu/prove/merged.html
#     http://miller-rabin.appspot.com
#     https://github.com/JuliaLang/julia/issues/11594
#     Forišek and Jančina, "Fast Primality Testing for Integers That Fit into a Machine Word", 2015
#         (in particular, see function FJ32_256, from which the hash and bases were taken)
const bases = UInt16[
    15591,  2018,   166,  7429,  8064, 16045, 10503,  4399,  1949,  1295,  2776,  3620,
      560,  3128,  5212,  2657,  2300,  2021,  4652,  1471,  9336,  4018,  2398, 20462,
    10277,  8028,  2213,  6219,   620,  3763,  4852,  5012,  3185,  1333,  6227,  5298,
     1074,  2391,  5113,  7061,   803,  1269,  3875,   422,   751,   580,  4729, 10239,
      746,  2951,   556,  2206,  3778,   481,  1522,  3476,   481,  2487,  3266,  5633,
      488,  3373,  6441,  3344,    17, 15105,  1490,  4154,  2036,  1882,  1813,   467,
     3307, 14042,  6371,   658,  1005,   903,   737,  1887,  7447,  1888,  2848,  1784,
     7559,  3400,   951, 13969,  4304,   177,    41, 19875,  3110, 13221,  8726,   571,
     7043,  6943,  1199,   352,  6435,   165,  1169,  3315,   978,   233,  3003,  2562,
     2994, 10587, 10030,  2377,  1902,  5354,  4447,  1555,   263, 27027,  2283,   305,
      669,  1912,   601,  6186,   429,  1930, 14873,  1784,  1661,   524,  3577,   236,
     2360,  6146,  2850, 55637,  1753,  4178,  8466,   222,  2579,  2743,  2031,  2226,
     2276,   374,  2132,   813, 23788,  1610,  4422,  5159,  1725,  3597,  3366, 14336,
      579,   165,  1375, 10018, 12616,  9816,  1371,   536,  1867, 10864,   857,  2206,
     5788,   434,  8085, 17618,   727,  3639,  1595,  4944,  2129,  2029,  8195,  8344,
     6232,  9183,  8126,  1870,  3296,  7455,  8947, 25017,   541, 19115,   368,   566,
     5674,   411,   522,  1027,  8215,  2050,  6544, 10049,   614,   774,  2333,  3007,
    35201,  4706,  1152,  1785,  1028,  1540,  3743,   493,  4474,  2521, 26845,  8354,
      864, 18915,  5465,  2447,    42,  4511,  1660,   166,  1249,  6259,  2553,   304,
      272,  7286,    73,  6554,   899,  2816,  5197, 13330,  7054,  2818,  3199,   811,
      922,   350,  7514,  4452,  3449,  2663,  4708,   418,  1621,  1171,  3471,    88,
    11345,   412,  1559,   194
]

function _witnesses(n::UInt64)
    i = ((n >> 16) $ n) * 0x45d9f3b
    i = ((i >> 16) $ i) * 0x45d9f3b
    i = ((i >> 16) $ i) & 255 + 1
    @inbounds return (Int(bases[i]),)
end
witnesses(n::Integer) =
        n < 4294967296      ? _witnesses(UInt64(n)) :
        n < 2152302898747   ? (2, 3, 5, 7, 11) :
        n < 3474749660383   ? (2, 3, 5, 7, 11, 13) :
                              (2, 325, 9375, 28178, 450775, 9780504, 1795265022)

isprime(n::UInt128) =
    n ≤ typemax(UInt64) ? isprime(UInt64(n)) : isprime(BigInt(n))
isprime(n::Int128) = n < 2 ? false :
    n ≤ typemax(Int64)  ? isprime(Int64(n))  : isprime(BigInt(n))

# Trial division of small (< 2^16) precomputed primes +
# Pollard rho's algorithm with Richard P. Brent optimizations
#     https://en.wikipedia.org/wiki/Trial_division
#     https://en.wikipedia.org/wiki/Pollard%27s_rho_algorithm
#     http://maths-people.anu.edu.au/~brent/pub/pub051.html
#
function factor!{T<:Integer,K<:Integer}(n::T, h::Associative{K,Int})
    # check for special cases
    if n < 0
        h[-1] = 1
        if isa(n, BitSigned) && n == typemin(T)
            h[2] = 8 * sizeof(T) - 1
            return h
        else
            return factor!(checked_neg(n), h)
        end
    elseif n == 1
        return h
    elseif n == 0 || isprime(n)
        h[n] = 1
        return h
    end

    local p::T
    for p in PRIMES
        if n % p == 0
            h[p] = get(h, p, 0) + 1
            n = div(n, p)
            while n % p == 0
                h[p] = get(h, p, 0) + 1
                n = div(n, p)
            end
            n == 1 && return h
            isprime(n) && (h[n] = 1; return h)
        end
    end
    T <: BigInt || widemul(n - 1, n - 1) ≤ typemax(n) ? pollardfactors!(n, h) : pollardfactors!(widen(n), h)
end

"""
    factor(n::Integer) -> Dict

Compute the prime factorization of an integer `n`. Returns a dictionary. The
keys of the dictionary correspond to the factors, and hence are of the same type as `n`.
The value associated with each key indicates the number of times the factor appears in the
factorization.

```jldoctest
julia> factor(100) # == 2^2 * 5^2
Dict{Int64,Int64} with 2 entries:
  2 => 2
  5 => 2
```

For convenience, a negative number `n` is factored as `-1*(-n)` (i.e. `-1` is considered
to be a factor), and `0` is factored as `0^1`:

```jldoctest
julia> factor(-9) # == -1 * 3^2
Dict{Int64,Int64} with 2 entries:
  -1 => 1
   3 => 2

julia> factor(0)
Dict{Int64,Int64} with 1 entries:
  0 => 1
```
"""
factor{T<:Integer}(n::T) = factor!(n, Dict{T,Int}())

"""
    factor(ContainerType, n::Integer) -> ContainerType

Return the factorization of `n` stored in a `ContainerType`, which must be a
subtype of `Associative` or `AbstractArray`, a `Set`, or an `IntSet`.

```jldoctest
julia> factor(DataStructures.SortedDict, 100)
DataStructures.SortedDict{Int64,Int64,Base.Order.ForwardOrdering} with 2 entries:
  2 => 2
  5 => 2
```

When `ContainerType <: AbstractArray`, this returns the list
of all prime factors of `n` with multiplicities, in sorted order.

```jldoctest
julia> factor(Vector, 100)
4-element Array{Int64,1}:
 2
 2
 5
 5

julia> prod(factor(Vector, 100)) == 100
true
```

When `ContainerType == Set`, this returns the distinct prime
factors as a set.

```jldoctest
julia> factor(Set, 100)
Set([2,5])
```
"""
factor{T<:Integer, D<:Associative}(::Type{D}, n::T) = factor!(n, D(Dict{T,Int}()))
factor{T<:Integer, A<:AbstractArray}(::Type{A}, n::T) = A(factor(Vector{T}, n))
factor{T<:Integer}(::Type{Vector{T}}, n::T) =
    sort!(mapreduce(collect, vcat, Vector{T}(), [repeated(k, v) for (k, v) in factor(n)]))
factor{T<:Integer, S<:Union{Set,IntSet}}(::Type{S}, n::T) = S(keys(factor(n)))
factor{T<:Any}(::Type{T}, n) = throw(MethodError(factor, (T, n)))

function pollardfactors!{T<:Integer,K<:Integer}(n::T, h::Associative{K,Int})
    while true
        local c::T = rand(1:(n - 1)), G::T = 1, r::K = 1, y::T = rand(0:(n - 1)), m::K = 1900
        local ys::T, q::T = 1, x::T
        while c == n - 2
            c = rand(1:(n - 1))
        end
        while G == 1
            x = y
            for i in 1:r
                y = y^2 % n
                y = (y + c) % n
            end
            local k::K = 0
            G = 1
            while k < r && G == 1
                for i in 1:(m > (r - k) ? (r - k) : m)
                    ys = y
                    y = y^2 % n
                    y = (y + c) % n
                    q = (q * (x > y ? x - y : y - x)) % n
                end
                G = gcd(q, n)
                k += m
            end
            r *= 2
        end
        G == n && (G = 1)
        while G == 1
            ys = ys^2 % n
            ys = (ys + c) % n
            G = gcd(x > ys ? x - ys : ys - x, n)
        end
        if G != n
            isprime(G) ? h[G] = get(h, G, 0) + 1 : pollardfactors!(G, h)
            G2 = div(n,G)
            isprime(G2) ? h[G2] = get(h, G2, 0) + 1 : pollardfactors!(G2, h)
            return h
        end
    end
end

end # module


include("../src/TrueSkill.jl")
using .TrueSkill
global const ttt = TrueSkill
using CSV; using Dates

println("Code 1")
mu = 0.0; sigma = 6.0; beta = 1.0; gamma = 0.03; draw = 0.0

println("Code 2")
a1 = ttt.Player(ttt.Gaussian(mu, sigma), beta, gamma); a2 = ttt.Player()
a3 = ttt.Player(); a4 = ttt.Player()

println("Code 3")
team_a = [ a1, a2 ]
team_b = [ a3, a4 ]
teams = [team_a, team_b]
g = ttt.Game(teams)

@time ttt.posteriors(ttt.Game(teams))

println("Code 4")
lhs = g.likelihoods
ev = g.evidence
ev = round(ev, digits=3)
println(ev)

println("Code 5")
pos = ttt.posteriors(g)
print(pos[1][1])
print(lhs[1][1] * a1.prior)

println("Code 6")
ta = [a1]
tb = [a2, a3]
tc = [a4]
teams_3 = [ta, tb, tc]
result = [1., 0., 0.]
g = ttt.Game(teams_3, result, p_draw=0.25)

@time ttt.posteriors(ttt.Game(teams_3, result, p_draw=0.25))

println("Code 7")
c1 = [["a"],["b"]]
c2 = [["b"],["c"]]
c3 = [["c"],["a"]]
composition = [c1, c2, c3]
h = ttt.History(composition, gamma = 0.0)

@time ttt.History(composition)

println("Code 8")
lc = ttt.learning_curves(h)
print(lc["a"])
print(lc["b"])

println("Code 9")
ttt.convergence(h)
lc = ttt.learning_curves(h)
print(lc["a"])
print(lc["b"])

@time ttt.convergence(h, iterations=1)

println("Code 10")
using Random; Random.seed!(999); N = 1000
function skill(experience, middle, maximum, slope)
    return maximum/(1+exp(slope*(-experience+middle)))
end
target = skill.(1:N, 500, 2, 0.0075)
opponents = Random.randn.(1000)*0.5 .+ target

println("Code 11")
composition = [[["a"], [string(i)]] for i in 1:N]
results = [ r ? [1.,0.] : [0.,1.] for r in (Random.randn(N).+target .> Random.randn(N).+opponents) ]
times = [i for i in 1:N]
priors = Dict{String,ttt.Player}()
for i in 1:N  priors[string(i)] = ttt.Player(ttt.Gaussian(opponents[i], 0.2))  end
@time h = ttt.History(composition, results, times, priors, gamma=0.015)
ttt.convergence(h)
@time ttt.convergence(h, iterations=1)

mu = [tp[2].mu for tp in ttt.learning_curves(h)["a"]]

println("Code 12")
data = CSV.read("input/history.csv")

times = Dates.value.(data[:,"time_start"] .- Date("1900-1-1"))
composition = [ r.double == "t" ? [[r.w1_id,r.w2_id],[r.l1_id,r.l2_id]] :
[[r.w1_id],[r.l1_id]] for r in eachrow(data) ]

h = ttt.History(composition=composition, times = times, sigma = 1.6, gamma = 0.036)
ttt.convergence(h,epsilon=0.01, iterations=10)

println("Code 13")
players = Set(vcat((composition...)...))
priors = Dict([(p, ttt.Player(ttt.Gaussian(0., 1.6), 1.0, 0.036) ) for p in players])

composition_ground = [ r.double == "t" ? [[r.w1_id, r.w1_id*r.ground, r.w2_id, r.w2_id*r.ground],[r.l1_id, r.l1_id*r.ground, r.l2_id, r.l2_id*r.ground]] : [[r.w1_id, r.w1_id*r.ground],[r.l1_id, r.l1_id*r.ground]] for r in eachrow(data) ]

h_ground = ttt.History(composition=composition_ground, times = times, sigma = 1.0, gamma = 0.01, beta = 0.0, priors = priors)
ttt.convergence(h_ground,epsilon=0.01, iterations=10)

println("Code 14")
N1 = ttt.Gaussian(mu = 1.0, sigma = 1.0); N2 = ttt.Gaussian(1.0, 2.0)

println("Code 15")
p1 = ttt.performance(a1)
p2 = ttt.performance(a2)
p3 = ttt.performance(a3)
p4 = ttt.performance(a4)

println("Code 16")
ta = p1 + p2; tb = p3 + p4

println("Code 17")
d = ta - tb

println("Code 18")
e = 1.0 - ttt.cdf(d, 0.0)

println("Code 19")
na = length(team_a)
nb = length(team_b)
sd = sqrt(na + nb)*beta
p_draw = 0.25
margin = ttt.compute_margin(p_draw, sd)

println("Code 20")
g = ttt.Game(teams, p_draw = 0.25)
post = ttt.posteriors(g)

println("Code 21")
d_approx = ttt.approx(d, margin, false)

println("Code 22")
approx_lh_d = d_approx / d

println("Code 23")
mu = a1.prior.mu
sigma2 = a1.prior.sigma^2
phi = d.mu
v2 = d.sigma^2
phi_div = approx_lh_d.mu
v2_div = approx_lh_d.sigma^2
prior = a1.prior
posterior = post[1][1]
println( prior * ttt.Gaussian(mu-phi+phi_div, sqrt(v2 + v2_div - sigma2)) )

println(posterior)


function labeldims(n)
    if n == 1
        ("\$t\$", "\$x\$")
    elseif n == 2
        ("\$x\$", "\$y\$")
    else
        ("\$x\$", "\$y\$", "\$z\$")
    end
end

function stationarycharacteristics(P::Process; n = length(getX0(P)), labels=labeldims(n), xlabel=labels[1], ylabel=labels[2], zlabel=(length(labels) == 3 ? labels[3] : nothing), size=(400, 350))
    p1 = marginaltrajectory(P; xlabel, ylabel, zlabel, size, mainalpha=0.35)
    display(p1)
    AMI = selfmutualinfo(P)
    p2 = plot(AMI.dims[1].val, Array((AMI isa AbstractVector ? AMI : AMI[:, end:-1:1]));  legend=true,   foreground_color_legend = nothing,
    background_color_legend = nothing, legendtitle=nothing, size,
    #xticks=[0, optimize_ticks(extrema(dims(AMI, 1).val)...; k_min = 3, k_max = 5)[1]...],
    yticks=optimize_ticks(extrema(AMI)...; k_min = 3, k_max = 6)[1],
    xlim=(0, Inf), xlabel="\$\\tau\$", ylabel="AMI")
    display(p2)
    return nothing
end
export stationarycharacteristics

function plotlyapunovresponse(P::Process, p, prange; size=(400, 350), N=1000, k=length(getX0(P)), kwargs...)
    λs = lyapunovresponse(P, p, prange, N, k)
    plot(prange, λs; size, xlabel="\$\\theta\$", ylabel="\$\\lambda\$", kwargs...)
end
export plotlyapunovresponse


#Reference:http://nbviewer.jupyter.org/gist/genkuroki/b46b0a094ee232302d41780defaffd6e

import PyPlot # Juliaでmatplotlib.pyplotを使うためのパッケージ
plt = PyPlot # as plt の代わり

using PyCall # JuliaからPythonの任意のライブラリを使うためのパッケージ
animation = pyimport("matplotlib.animation") # as animation の代わり

# このノートブックでは imagemagick でGIF動画を作成し,
# ノートブック内にbase64エンコードして読み込んで表示する.
function displayfile(mimetype, filename)
    open(filename) do f
        base64text = base64encode(f)
        display("text/html", """<img src="data:$mimetype;base64,$base64text">""")
    end
end

fig,ax=plt.subplots()
ax[:set_xlim](-4.5,10.5)
ax[:set_ylim](-4.5,4.5)

b=0.8
x=10.0
y=3.0
times=200
srand(12345) # np.random.seed(12345) の代わり. 上のセルと同じ結果になるように挿入した.

function update(data) # defをfunctionに置き換えて, コロンを消した.
    u1 = rand()
    u2 = rand()
    global x, y
    x_old=x
    x = sqrt(-2*log(u1))*cos(2*pi*u2) + b*y
    ax[:plot]([x_old, x], [y, y], lw=1)
    y_old = y
    y = sqrt(-2*log(u1))*sin(2*pi*u2) + b*x
    ax[:plot]([x, x], [y_old, y], lw=1)
    ax[:set_title]("$data")
end # end を追加した.
ani=animation[:FuncAnimation](fig,update,frames=times,interval=100)
ani[:save]("output.gif", writer="imagemagick") # GIFファイルに出力

plt.clf() # ←必ずしも必要ない. 単に余計な表示を消すためのコード

displayfile("image/gif", "output.gif")


"""
    jackknife_error(g::Function, x; [binsize=10])

Computes the jackknife standard deviation of `g(<x>)` by binning
x and performing leave-one-out analysis.
"""
function jackknife_error(g::Function, x::AbstractVector{T}; binsize::Int=10) where T<:Real
    rem(length(x), binsize) == 0 || throw(ArgumentError("Bin size not compatible with sample size"))

    xblock = map(mean, Base.Iterators.partition(x, binsize))
    gis = Jackknife.leaveoneout(g, xblock)
    return sqrt(Jackknife.var(g, xblock; gis=gis))
end

function jackknife_error(g::Function, x::AbstractVector{T}; binsize::Int=10) where T<:Complex
    sqrt(jackknife_error(g, real(x), binsize=binsize)^2 + jackknife_error(g, imag(x), binsize=binsize)^2)
end

function jackknife_error(g::Function, x::AbstractArray{T}; binsize::Int=10) where T<:Number
    ndimsx = ndims(x)
    mapslices(y->binning_error(y; binsize=binsize), x, ndimsx)[(Colon() for _ in 1:ndimsx-1)...,1]
end

# Data: arrays
function jackknife_error(g::Function, x::Union{AbstractArray{<:Number},AbstractVector{T} where T<:AbstractArray{<:Number}}; binsize::Int=10)
    error("Jackknife error calculation for array like data not yet supported.")
end

"""
**Jackknife** errors for functions of uncertain data, i.e. g(<x>)
"""
module Jackknife

# https://github.com/ararslan/Jackknife.jl was useful resource

"""
    leaveoneout(g::Function, x)
Estimate the `g(x)` for each subsample of `x`, systematically
omitting each index one at a time. The result is a vector with the resulting `g(x_[])`.
"""
function leaveoneout(g::Function, x::AbstractVector{T}) where T<:Real
    length(x) > 1 || throw(ArgumentError("The sample must have size > 1"))
    return [g(circshift(x, -i)[2:end]) for i in 0:length(x)-1]
end
# function leaveoneout(g::Function, x::AbstractVector{T}) where T<:Real
#     length(x) > 1 || throw(ArgumentError("The sample must have size > 1"))
#     inds = eachindex(x)
#     return [g(mean(x[filter(j -> j != i, inds)])) for i in inds]
# end

"""
    var(g::Function, x)
Compute the jackknife estimate of the variance of `g(<x>)`, where `g` is given as a
function that computes a point estimate when passed a real-valued vector `x`.
"""
function var(g::Function, x::AbstractVector{T}; gis::AbstractVector{T}=leaveoneout(g, x)) where T<:Real
    n = length(x)
    return Base.var(gis) * (n - 1)^2 / n # Eq. (3.35) in QMC Methods book
end

"""
    bias(g::Function, x)
Compute the jackknife estimate of the bias of `g`, which is given as a
function that computes a point estimate when passed a real-valued vector `x`.
"""
function bias(g::Function, x::AbstractVector{T}; gis::AbstractVector{T}=leaveoneout(g, x)) where T<:Real
    return (length(x) - 1) * (mean(gis) - g(x)) # Basically Eq. (3.33)
end

"""
    estimate(g::Function, x)
Compute the bias-corrected jackknife estimate of the parameter `g(<x>)`.
"""
function estimate(g::Function, x::AbstractVector{T}; gis::AbstractVector{T}=leaveoneout(g, x)) where T<:Real
    n = length(x)
    return n * g(x) - (n - 1) * mean(gis) # Eq. (3.34) in QMC Methods book
end

end # module

export Jackknife