module Int

include("Computation.jl")
using .Computation

include("Format.jl")
using .Format

export int, vint

# Determines upper bound alpha with Interval Method, optionally outputs proof
# Authors: Antara Hebbar and Jason Liu
function int(m::Int64, s::Int64; output::Int64 = 2)
    output > 0 && printHeader(center("INTERVAL METHOD"))

    if m <= 0 || s <= 0
        if output > 0
            printf("Interval Method does not apply", line = true)
            printEnd()
        end
        return 1 // 1
    elseif m < s
        alpha = m // s * int(s, m, output = 0)
        if output > 1
            printfT(
                "Duality Theorem",
                "muffins($m,$s) = $m/$s · muffins($s,$m)",
                "Bounding muffins($s,$m) instead",
            )
            int(s, m, output = output)
        elseif output > 0
            printfT(
                "Interval Method",
                "Upper bound α of muffins($m,$s) is $(formatFrac(alpha))",
            )
            printEnd()
        end
        return alpha
    elseif m % s == 0
        if output > 0
            printfT("Interval Method", "Since $m % $s = 0, muffins($m,$s) = 1")
            printEnd()
        end
        return 1 // 1
    end

    (V, W, sV, sW) = sv(m, s)
    numV = (V)sV
    numW = (W)sW

    # Derive alpha using Int Algorithm
    if numW > numV
        f = Int64(floor(numV / sW))                                     # Upper bound of minimum # of large W-shs a student can have
        g = Int64(floor((numW - numV) / sW))                            # Lower bound of minimum # of small W-shs a student can have
        alpha = min(
            ((W - f)W - (W - f + 1)m // s + f) // ((W - f - 1)W + 2f),  # Value for alpha derived by solving f(1-α) + (W-f)(1-y) = m/s
            ((g - 1)W + (W - 2g + 1)m // s) // (W^2 - g),               # Value for alpha derived by solving gy + (W-g)(1-x) = m/s
        )
        alpha = alpha < 1 / 3 ? 1 // 3 : alpha
        if vint(m, s, alpha, output = output)
            return alpha
        end
    elseif numW < numV
        f = Int64(floor((numV - numW) / sV))                            # Upper bound of minimum # of large V-shs a student can have
        g = Int64(floor(numW / sV))                                     # Lower bound of minimum # of small V-shs a student can have
        alpha = min(
            ((V - f)W + (2f - V - 1)m // s) // ((V - f)W + (f - 1)V),   # Value for alpha derived by solving f(1-x) + (V-f)(1-y) = m/s
            ((V - g + 1)m // s + g - V) // ((V - g - 1)V + 2g),         # Value for alpha derived by solving gα + (V-g)(1-x) = m/s
        )
        alpha = alpha < 1 / 3 ? 1 // 3 : alpha
        if vint(m, s, alpha, output = output)
            return alpha
        end
    end

    if output > 0
        printf("Interval Method inconclusive", line = true)
        printEnd()
    end
    1 // 1
end

# Helper function for int -- verifies whether int(m, s, alpha) is conclusive
function vint(m::Int64, s::Int64, alpha::Rational{Int64}; output::Int64 = 2)
    if !(m > s > 0) || m % s == 0
        output > 0 && printf("VInt does not apply", line = true)
        false
    elseif alpha < 1 / 3
        output > 0 && printfT("No piece size < 1/3", "For m ≥ s, α must be ≥ 1/3")
        false
    elseif alpha > 1
        output > 0 && printfT("No piece size > 1", "α must be ≤ 1")
        false
    else
        (V, W, sV, sW) = sv(m, s)
        numV = (V)sV
        numW = (W)sW

        # Initialize Interval Method proof
        if output > 1
            # Define and format variables for proof
            alphaF = formatFrac(alpha)
            size = formatFrac(m // s)
            a = formatFrac(m // (V + 1)s)
            b = max(1 - m // (W - 1)s, 0)
            b1 = formatFrac(1 - b)
            b = formatFrac(b)

            # Establish assumptions and premises
            printHeader("CLAIM")
            printfT(
                "Claim",
                "For all ($m, $s) procedures where $s students each receive $size muffins, the smallest piece size is bounded above by α = $alphaF",
                "Proven by contradicting the assumption",
            )

            printHeader("ASSUMPTION")
            printfT("Assumption", "Assume muffins($m,$s) > α ≥ 1/3")
            printfT(
                "Property of Buddies",
                "Since muffins($m,$s) > 1/3, each muffin must be cut into 2 shs, totaling $(2m) shs",
                "",
                "Since each muffin is cut into 2 shs that are buddies, each sh of size x must imply the existence of a sh of size 1-x",
            )

            # Begin casework
            printHeader("CASEWORK")

            printfT(
                "Case 1",
                "Alice has ≥ $(V+1) shs",
                "One of them is ≤ ($size × 1/$(V+1)) = $a",
                "",
                "Contradicts assumption if α ≥ $a",
            )

            if W - 1 <= 1 && m / s > 1
                printfT(
                    "Case 2",
                    "Bob cannot have ≤ $(W-1) shs since $size > 1, so this case is impossible",
                )
            else
                printfT(
                    "Case 2",
                    "Bob has ≤ $(W-1) shs",
                    "One of them is ≥ ($size × 1/$(W-1)) = $b1",
                    "Its buddy is ≤ (1 - $b1) = $b",
                    "",
                    "Contradicts assumption if α ≥ $b",
                )
            end
        end

        # Check if V-Conjecture applies
        if m // s * 1 // (V + 1) > alpha || 1 - m // s * 1 // (W - 1) > alpha
            output > 0 && printf("V-Conjecture does not apply, VInt failed", line = true)
            return false
        end

        ((_, x), (y, _)) = findend(m, s, alpha, V)
        xOriginal = m // s - alpha * (V - 1)
        yOriginal = m // s - (1 - alpha) * (V - 2)

        # Check if FindEnd works
        if x != xOriginal && y != yOriginal
            output > 0 && printf("FindEnd inconclusive, VInt failed", line = true)
            return false
        end

        # Define and format variables for proof
        cd = lcm(s, denominator(x), denominator(y))
        alphaF = formatFrac(alpha, cd)
        alpha1 = formatFrac(1 - alpha, cd)
        size = formatFrac(m // s)
        a = formatFrac(m // s - x, cd)
        b = formatFrac(m // s - y, cd)
        xF = formatFrac(x, cd)
        x1 = formatFrac(1 - x, cd)
        yF = formatFrac(y, cd)
        y1 = formatFrac(1 - y, cd)
        diff = abs(numV - numW)

        # Continue proof
        if output > 1
            # Continue casework
            printfT(
                "Note",
                "The remaining cases deal with everyone having either $W or $V shs, so:",
                "",
                "$(W)·s_$W + $(V)·s_$V = $(2m)  (total shs)",
                "s_$W + s_$V = $s  (total students)",
                "where s_N = # of students with N shs",
                "",
                "The solution to the system is:",
                "s_$W = $sW, s_$V = $sV",
                "So there are $numW $W-shs and $numV $V-shs",
            )

            if x == xOriginal
                printfT(
                    "Case 3",
                    "Alice has a $V-sh ≥ $xF",
                    "Her other $(V-1) $V-shs sum to ≤ ($size - $xF) = $a",
                    "One of them is ≤ ($a × 1/$(V-1)) = $alphaF",
                    "",
                    "Contradicts assumption if α ≥ $alphaF",
                )
            else
                printfT("Case 3", "FindEnd did not produce a conclusive bound for $V-shs")
            end

            if y == yOriginal
                printfT(
                    "Case 4",
                    "Bob has a $W-sh ≤ $yF",
                    "His other $(W-1) $W-shs sum to ≥ ($size - $yF) = $b",
                    "One of them is ≥ ($b × 1/$(W-1)) = $alpha1",
                    "Its buddy is ≤ (1 - $alpha1) = $alphaF",
                    "",
                    "Contradicts assumption if α ≥ $alphaF",
                )
            else
                printfT("Case 4", "FindEnd did not produce a conclusive bound for $W-shs")
            end

            printHeader("CASE 5: INTERVAL ANALYSIS")
            println()
            printf("The following intervals capture the negation of the previous cases:")
            if alpha < x <= y < 1 - alpha
                println(
                    "\n",
                    interval(
                        ["(", alphaF],
                        [")[", xF],
                        ["](", yF],
                        [")", alpha1],
                        labels = ["$numV $V-shs", "0", "$numW $W-shs"],
                    ),
                )
                printLine()
            else
                if xF != alphaF
                    println(
                        "\n",
                        interval(
                            ["(", alphaF],
                            [")[", xF],
                            [")", alpha1],
                            labels = ["$numV $V-shs", "0 $V-shs"],
                        ),
                    )
                else
                    println("\n", "No negation interval for $V-shs")
                end

                if yF != alpha1
                    println(
                        "\n",
                        interval(
                            ["(", alphaF],
                            ["](", yF],
                            [")", alpha1],
                            labels = ["0 $W-shs", "$numW $W-shs"],
                        ),
                    )
                else
                    println("\n", "No negation interval for $W-shs")
                end

                printfT(
                    "Case 5",
                    "The Interval Method is inconclusive on these intervals, VInt failed",
                )
            end
        end

        # Fail if interval inconclusive
        if x > y || x == alpha || y == 1 - alpha
            output == 1 &&
                printf("Could not generate disjoint intervals, VInt failed", line = true)
            return false
        end

        # Define and format variables for proof
        (vMin, vMax, sMin, sMax) = (V, W, sV, sW)
        (f, g) = (Int64(floor(numV / sW)), Int64(floor(diff / sW)))
        (lim1, lim2) = ((vMin)sMin, diff)
        (i, j, k, l) = (xF, yF, y1, x1)
        (rngMin, rngMax, rngS, rngL) = ((alphaF, xF), (x1, alpha1), (yF, y1), (x1, alpha1))
        if numV > numW
            (vMin, vMax, sMin, sMax) = (W, V, sW, sV)
            (f, g) = (Int64(floor(diff / sV)), Int64(floor(numW / sV)))
            (lim1, lim2) = (diff, (vMin)sMin)
            (i, j, k, l) = (y1, x1, xF, yF)
            (rngMin, rngMax, rngS, rngL) =
                ((yF, alpha1), (alphaF, y1), (alphaF, y1), (x1, xF))
        end
        numMin = (vMin)sMin

        # Conclude interval case
        if output > 1
            printfT(
                "Property of Buddies",
                "Because 0 shs lie in [$xF,$yF], 0 shs must also lie in [$y1,$x1]",
                "",
                "Similary, the $numMin $vMin-shs in ($(rngMin[1]),$(rngMin[2])) implies $numMin $vMax-shs in ($(rngMax[1]),$(rngMax[2]))",
            )
            println()
            printf("The following intervals capture the previous statements:")
            println(
                "\n",
                interval(
                    ["(", alphaF],
                    [")[", i],
                    ["](", j],
                    [")[", k],
                    ["](", l],
                    [")", alpha1],
                    labels = [
                        "$numMin $V-shs",
                        "0",
                        "$diff $vMax-shs",
                        "0",
                        "$numMin $W-shs",
                    ],
                ),
            )
            printLine()

            printfT(
                "Note",
                "Let the $vMax-shs that lie in ($(rngS[1]),$(rngS[2])) be small $vMax-shs",
                "Let the $vMax-shs that lie in ($(rngL[1]),$(rngL[2])) be large $vMax-shs",
            )
        end

        upperB = (vMax - f) * toFrac(rngS[2]) + (f) * toFrac(rngL[2])
        lowerB = (vMax - g) * toFrac(rngL[1]) + (g) * toFrac(rngS[1])
        u = formatFrac(upperB, cd)
        l = formatFrac(lowerB, cd)

        if upperB <= m // s
            if output > 1
                printfT(
                    "Case 5.1",
                    "Alice has ≤ $f large $vMax-shs",
                    "She must also have ≥ $(vMax-f) small $vMax-shs, so her maximum amount of muffins is:",
                    "< ($(vMax-f) × $(rngS[2])) + ($f × $(rngL[2])) = $u ≤ $size",
                    "",
                    "Since Alice's muffin amount can not reach $size, this case is impossible",
                )
                printfT(
                    "Case 5.2",
                    "All $sMax $vMax-sh students have at least $(f+1) large $vMax-shs",
                    "",
                    "This is impossible because $(f+1) × $sMax = $((f+1)sMax) ≥ $lim1",
                )
            end
        elseif lowerB >= m // s
            if output > 1
                printfT(
                    "Case 5.1",
                    "Bob has ≤ $g small $vMax-shs",
                    "He must also have ≥ $(vMax-g) large $vMax-shs, so his minimum amount of muffins is:",
                    "> ($g × $(rngS[1])) + ($(vMax-g) × $(rngL[1])) = $l ≥ $size",
                    "",
                    "Since Bob's muffin amount can not reach $size, this case is impossible",
                )
                printfT(
                    "Case 5.2",
                    "All $sMax $vMax-sh students have at least $(g+1) small $vMax-shs",
                    "",
                    "This is impossible because $(g+1) × $sMax = $((g+1)sMax) ≥ $lim2",
                )
            end
        else
            if output > 1
                printfT(
                    "Bound # of large $vMax-shs",
                    "All $sMax $vMax-sh students have at least $(f+1) large $vMax-shs",
                    "",
                    "This is impossible because $(f+1) × $sMax = $((f+1)sMax) ≥ $lim1",
                )
                printfT(
                    "Negation of prev. bound",
                    "∃ student Alice with ≤ $f large $vMax-shs",
                    "She must also have ≥ $(vMax-f) small $vMax-shs, so her maximum amount of muffins is:",
                    "< ($(vMax-f) × $(rngS[2])) + ($f × $(rngL[2])) = $u",
                    "",
                    "Since $size lies below the interval, this bound is inconclusive",
                )

                printfT(
                    "Bound # of small $vMax-shs",
                    "All $sMax $vMax-sh students have at least $(g+1) small $vMax-shs",
                    "",
                    "This is impossible because $(g+1) × $sMax = $((g+1)sMax) ≥ $lim2",
                )
                printfT(
                    "Negation of prev. bound",
                    "∃ student Bob with ≤ $g small $vMax-shs",
                    "He must also have ≥ $(vMax-g) large $vMax-shs, so his minimum amount of muffins is:",
                    "> ($g × $(rngS[1])) + ($(vMax-g) × $(rngL[1])) = $l",
                    "",
                    "Since $size lies above the interval, this bound is inconclusive",
                )

                printf(
                    "Bounding # of large and small $vMax-shs inconclusive, restarting with individual case analysis",
                    line = true,
                )
            end

            for k = 0:vMax
                upperB = (vMax - k) * toFrac(rngL[2]) + (k) * toFrac(rngS[2])
                lowerB = (vMax - k) * toFrac(rngL[1]) + (k) * toFrac(rngS[1])
                u = formatFrac(upperB, cd)
                l = formatFrac(lowerB, cd)
                if upperB <= m // s || lowerB >= m // s
                    if output > 1
                        printfT(
                            "Case 5.$(k+1)",
                            "Alice has $k small $vMax-shs and $(vMax-k) large $vMax-shs",
                            "Her possible muffin amount lies in:",
                            "",
                            "($k × $(rngS[1]) + $(vMax-k) × $(rngL[1]),",
                            "$k × $(rngS[2]) + $(vMax-k) × $(rngL[2]))",
                            "= ($l, $u)",
                            "",
                            "Since $size lies outside this interval, this case is impossible",
                        )
                    end
                else
                    # Fail if interval analysis inconclusive
                    if output > 1
                        printfT(
                            "Case 5.$(k+1)",
                            "Bob has $k small $vMax-shs and $(vMax-k) large $vMax-shs",
                            "His possible muffin amount lies in:",
                            "",
                            "($k × $(rngS[1]) + $(vMax-k) × $(rngL[1]),",
                            "$k × $(rngS[2]) + $(vMax-k) × $(rngL[2]))",
                            "= ($l, $u)",
                            "",
                            "Since $size lies inside this interval, this case is inconclusive",
                        )
                    elseif output > 0
                        printf("Interval analysis inconclusive, VInt failed", line = true)
                    end
                    return false
                end
            end
        end

        # Conclude proof
        alpha = formatFrac(alpha)
        if output > 1
            # Conclude with alpha's value
            printHeader("CONCLUSION")
            printfT(
                "Conclusion",
                "All possible cases contradict the assumption iff. α ≥ $alpha",
                "muffins($m,$s) ≤ α, ∀ α ≥ $alpha",
                "",
                "muffins($m,$s) ≤ $alpha",
            )
            printEnd()
        elseif output > 0
            printfT("Interval Method", "Upper bound α of muffins($m,$s) is $alpha")
            printEnd()
        end

        true
    end
end

end


using QuasiArrays, ArrayLayouts, LinearAlgebra, Test
import QuasiArrays: apply

@testset "Multiplication" begin
    @testset "Diag * Inclusion" begin
        A = QuasiDiagonal(Inclusion(0:0.1:1))
        @test rowsupport(A,0.1) == colsupport(A,0.1) == 0.1
        b = Inclusion(0:0.1:1)
        Ab = A*b
        @test Ab[0.1] ≈ 0.1^2
        @test_throws DimensionMismatch A*Inclusion(1:2)

        @test A*A isa QuasiDiagonal
        @test (A*A)[0.1,0.1] ≈ A[0.1,0.1]^2
    end
    @testset "Quasi * Quasi" begin
        A = QuasiArray(rand(3,3),(0:0.5:1,0:0.5:1))
        @test *(A) == A
        @test A*A isa QuasiArray
        @test QuasiArray(A*A) == QuasiArray(A.parent*A.parent,A.axes)
    end
    @testset "Quasi * Array" begin
        A = QuasiArray(rand(3,3),(0:0.5:1,Base.OneTo(3)))
        B = rand(3,3)
        @test A*B isa QuasiArray
        @test QuasiArray(A*B) ≈ QuasiArray(A.parent*B,A.axes)
        @test B*A' isa QuasiArray
        @test (A*B)[0.5,1] ≈ (Array(A)*B)[2,1]
    end

    @testset "Triple" begin
        A = QuasiArray(rand(3,3),(0:0.5:1,0:0.5:1))
        @test A isa QuasiArray
        @test Array(A*A) ≈ Array(A)^2
        @test Array(A*A*A) ≈ Array(A)^3
    end

    @testset "absvec" begin
        a = QuasiArray(rand(3),(0:0.5:1,))
        c = QuasiArray(rand(3) .+ im .* randn(3),(0:0.5:1,))
        @test a'a ≈ transpose(a)a ≈ dot(a,a)
        @test c'c ≈ dot(c,c)
        @test transpose(c)c ≈ transpose(c.parent)*c.parent
    end
end


include("nnlib.jl")
include("fields.jl")
using Test, Random
using Distributions, StatsBase, LinearAlgebra

#Base.broadcastable(x::LinearField) = Ref(x)

@testset "field" begin
    lf = LinearField(0.5)
    out, x = 0.0, 2.0
    @instr get_field(lf, out, x)
    @test out == 1.0
    @test check_inv(PVar, (Dup(0.5),))

    x = 2.0
    y = 0.5
    θ = 0.5
    @instr update_field(LinearField(θ), y, x; dt=0.1)
    @test y === 0.6
    @i function gf(θ, y, x)
        LinearField(θ)
        get_field(θ, y, x)
    end

    @i function uf(θ, y, x; dt)
        LinearField(θ)
        update_field(θ, y, x; dt=dt)
    end
    @test check_inv(update_field, (LinearField(θ), y, x), kwargs=(dt=0.1,))
    @test check_grad(gf, (θ, Loss(y), x); verbose=true)
    @test check_grad(uf, (θ, Loss(y), x); kwargs=(dt=0.1,), verbose=true)

    x = PVar(2.0)
    y = PVar(0.5)
    θ = 0.5
    @instr update_field(LinearField(θ), y, x; dt=0.1)
    @test value(y) === 0.6
    @test y.logp ≈ -0.05
end

@i function get_field(ps::typeof(⊕(sin)), field_out, y)
    field_out += sin(y)
end

function f(F, y0, N, dt)
    y = y0
    for i=1:N
       y += F(y)*dt
    end
    return y
end

@testset "leapfrog" begin
    x = 0.8
    truth = f(sin, 0.8, 10000, 0.0001)
    field = ⊕(sin)
    @instr Dup(x)
    @instr leapfrog(field, x; Nt=10000, dt=0.0001)
    @test isapprox(x, truth; atol=1e-3)

    # integrate sin(θ), get -log(det(∂cos(θ)/∂θ))
    lf = LinearField(0.5)
    x = PVar(0.8)
    truth = value(x)
    for i=1:1000
        truth += 0.001*truth*lf.θ
    end
    @instr Dup(x)
    @test check_inv(leapfrog, (lf, x); kwargs=(Nt=1000, dt=0.001))
    @instr leapfrog(lf, x; Nt=1000, dt=0.001)
    @test isapprox(value(x), truth; atol=1e-3)
    @test isapprox(x.x.logp, -0.5, atol=1e03) # TODO: manual check
end

Base.adjoint(gv::GVar) = GVar(gv.x', gv.g')
Base.conj(gv::GVar) = GVar(conj(gv.x), conj(gv.g))
@testset "ode loss" begin
    Random.seed!(3)
    lf = LinearField(0.5)
    xs = randn(10)
    μ, σ = 0.0, 1.0
    loss_out = 0.0
    @instr Dup.(xs)
    @test check_inv(ode_loss, (μ, σ, lf, xs, loss_out); kwargs=Dict(:Nt=>10, :dt=>0.1), verbose=true)
    @test check_grad(ode_loss, (μ, σ, lf, xs, Loss(loss_out)); kwargs=Dict(:Nt=>100, :dt=>0.01), verbose=true)
    @instr ode_loss(μ, σ, lf, xs, loss_out; Nt=100, dt=0.01)
    @show loss_out
end

@testset "normal log pdf" begin
    out = 0.0
    x = 1.0
    μ = 0.3
    σ = 1.5
    l1 = normal_logpdf(out, x, μ, σ)
    distri = Normal(μ, σ)
    l2 = logpdf(distri, x)
    @test l1[1] ≈ l2
    @test check_inv(normal_logpdf, (out, x, μ, σ))
    x2 = [0.2, 0.9]
    μ2 = [0.0, 0.2]
    σ = 1.0
    distri = MultivariateNormal(μ2, [σ 0; 0 σ])
    @test check_inv(normal_logpdf2d, (out, x2, μ2, σ))
    @test logpdf(distri, x2) ≈ normal_logpdf2d(0.0, x2, μ2, σ)[1]
end

#=
@testset "neural ode" begin
    Random.seed!(2)
    θ = 0.5
    nsample = 100
    Nt = 50
    dt = 0.01

    source_distri = Normal(1, 0.5)
    xs0 = rand(source_distri, nsample)
    logp0 = max.(logpdf.(source_distri, xs0), -1000)

    # solving the ode, and obtain the probability change
    v_xs = copy(xs0)
    v_ks = copy(xs0)
    v_logp = copy(logp0)
    v_θ = θ
    field_out = 0.0
    tape = ode_with_logp!(v_xs, v_ks, v_logp, field, field_out, v_θ, Nt, dt)
    @test check_inv(tape; verbose=true, atol=1e-3)
end
=#


include(joinpath(@__DIR__, "branch_variable.jl"))
include(joinpath(@__DIR__, "branch_special.jl"))

_branchvar_lower_solution(m,i) = m._lower_solution[m._branch_to_sol_map[i]]
function _fractional_integer_branch!(t::ExtensionType, m::GlobalOptimizer{N,T,S}) where {N,T<:AbstractFloat,S}
    if !m._current_node.cont
        integer_count = 0
        for i in m.branch_variables
            if _is_unfixed_integer(m, i)
                isinteger(_lower_solution(BranchVar, m, i))
                integer_count += 1
            end
        end
    end
    return false
end

#=
5.3.1 of Belotti, Pietro, et al. "Branching and bounds tighteningtechniques for
non-convex MINLP." Optimization Methods & Software 24.4-5 (2009): 597-634.
=#
function _select_branch_point!(m::GlobalOptimizer{N,T,S}, i) where {N,T<:Real,S}
    l = _lower_bound(BranchVar, m, i)
    u = _upper_bound(BranchVar, m, i)
    s = _lower_solution(BranchVar, m, i)
    α = _branch_cvx_α(m)
    b = _branch_offset_β(m)*(u - l)
    return max(l + b, min(u - b, α*s + (one(T) - α)*_mid(BranchVar, m, i)))
end

function _continuous_branch!(t::ExtensionType, m::GlobalOptimizer{N,T,S}) where {N,T<:AbstractFloat,S}

    bi = _select_branch_variable!(m)
    bp = _select_branch_point!(m)

    # rounds into branch points, which in turn prevents the
    # solution at the branch point from being discarded
    lower_bound = max(n.lower_bound, m._lower_objective_value)
    upper_bound = min(n.upper_bound, m._upper_objective_value)
    new_depth = n.depth + 1

    m._maximum_node_id += 1
    push!(m._stack, NodeBB(_upper_bound(:branch, m),
                           ntuple(i -> (i == bi ? nextfloat(bp) : _upper_bound(BranchVar, m, i)), N),
                           lower_bound, upper_bound, new_depth, m._maximum_node_id,
                           bi, BD_NEG))

    m._maximum_node_id += 1
    push!(m._stack, NodeBB(ntuple(i -> (i == bi ? prevfloat(bp) : _lower_bound(BranchVar, m, i)), N),
                           _upper_bound(:branch, m),
                           lower_bound, upper_bound, new_depth, m._maximum_node_id,
                           bi, BD_POS))

    m._node_repetitions = 1
    m._node_count += 2
    return true
end

"""
$(SIGNATURES)

Creates two nodes from `current_node` using information available the `x`
and stores them to the stack. By default, relative width bisection is perfomed
at a point `branch_pnt` which is a convex combination
(parameter: `branch_cvx_factor`) of the solution to the relaxation and
the midpoint of the node. If this solution lies within `branch_offset/width` of
a bound then the branch point is moved to a distance of `branch_offset/width`
from the bound.
"""
function branch_node!(t::ExtensionType, m::GlobalOptimizer)
    was_branched = _conic_branch!(t, m)
    !was_branched && (was_branched &= _fractional_integer_branch!(t, m))
    !was_branched && (was_branched &= _continuous_branch!(t, m))
    return nothing
end


# INVESTIGATE THE FOLLOWING:
# * Orthogonalize the conjugate vector and the residual vector together, periodically against previous vectors, does it
#   reduce the number of iterations?

mutable struct ConjGradModified{T <: Number}
    A::Function              # Linear opeartor
    b::AbstractArray{T}      # RHS vector
    tensor_size::Tuple       # The size of the tensor the linear operator is acting on.

    x::Vector{T}        # current guess, started with initial guess.
    r::Vector{T}        # previous computed residual
    rnew::Vector{T}     # Current residual
    d::Vector{T}        # Conjugate Direction
    Ad::Vector{T}       # for reducing garbage collector time.

    itr::UInt64              # Iteration count.
    Q::Vector{Vector{T}}     # list of normalized residuals vectors.
    Ds::Vector{Vector{T}}    # ist of normalized Ad vector during the iterations.

    function ConjGradModified(
        A::Function,
        b::AbstractArray,
        x0::Union{AbstractArray,Nothing}=nothing
    )
        x0 = x0===nothing ? b .+ 0.1 : x0
        r = reshape(b - A(x0), :)
        T = typeof(r[1])
        this = new{T}()
        this.r = r
        this.A = A
        this.b = b
        this.tensor_size = size(b)
        this.x = x0
        this.rnew = similar(this.r)
        this.d = this.r
        this.Ad = similar(this.r)
        ComputeVec!(this, this.r, this.Ad)
        this.x = x0
        this.itr = 0

        this.Q = Vector{Vector{T}}()
        this.Ds = Vector{Vector{T}}()
        push!(this.Q, this.r/norm(this.r))
        push!(this.Ds, this.d/norm(this.d))
        return this
    end

    function ConjGradModified(
            A::AbstractMatrix,
            b::AbstractVector,
            x0::Union{AbstractArray,Nothing}=nothing
        )
        return ConjGradModified((x)-> A*x, b, x0)
    end

end

"""
    Apply the linear transformation on the tensor and return the
    residual vector after transformation.
"""
function ComputeResidualVec(this::ConjGradModified, x::AbstractArray)
    resVec = this.b - this.A(reshape(x, this.tensor_size))
    return reshape(resVec, length(resVec))
end

"""
    Get the norm of the current residual of Modified Conjugate
    Gradient.
"""
function GetResidualNorm(this::ConjGradModified)
    return norm(this.r)
end

"""
    Compute the vector using the linear transformation.
"""
function ComputeVec(this::ConjGradModified, x::AbstractArray)
    vec = this.A(reshape(x, this.tensor_size))
    return reshape(vec, length(vec))
end

"""
    Pre-allocates a vector and get the values to the given
    pre-allocated vector.
"""
function ComputeVec!(this::ConjGradModified, x::AbstractArray, vec::AbstractArray)
    vec .= reshape(this.A(reshape(x, this.tensor_size)), :)
    return nothing
end

"""
    Force the orthogonalization happen between the residual vectors against
    all previous residual vectors, and the conjugate vector as well.
"""
function ForceOrthogonalize(this::ConjGradModified)
    for Idx in 1: length(this.Q)
        for Jdx in 1: Idx - 1
            # this.Q[Idx] .-= dot(this.Q[Jdx], this.Q[Idx])*this.Q[Jdx]
            this.Ds[Idx] .-= dot(this.Ds[Jdx], this.Ds[Idx])*this.Ds[Jdx]
        end
    end
    for  Aq in this.Q
        this.d .-= dot(Aq, this.d)*Aq
    end

    return
end

"""
    Performs One step of conjugate Gradient.
"""
function (this::ConjGradModified)()
    r = this.r
    if norm(r) == 0
        return 0 # The problem is solved already.
    end
    d = this.d
    Ad = this.Ad
    a = dot(r, r)/dot(d, Ad)

    if a < 0
        error("CG got a non-definite matrix")
    end

    this.x += a*d
    this.rnew .= r - a*Ad                    # update rnew
    rnewNorm = norm(this.rnew)

    b = dot(this.rnew, this.rnew)/dot(r, r)
    # @assert abs(dot(rnew + β*d, Ad)) < 1e-8 "Not conjugate"
    this.d = this.rnew + b*d
    this.r = copy(this.rnew)                      # Override
    this.itr += 1
    ComputeVec!(this, this.d, Ad)
    return convert(Float64, rnewNorm)
end

"""
    Performs j iterations of conjugate gradient, collect the
    2norm of the residual and return all of them in a vector.
"""
function (this::ConjGradModified)(j::Int64)
    results = Vector()
    for _ in 1: j
        push!(results, this())
    end
    return results
end

### Using stuff.
using Test, LinearAlgebra
@testset begin

    function Test1()
        n = 10
        A = rand(n, n)
        A = A'*A
        display(A)
        b = rand(n)
        cg = ConjGradModified(A, b)
        for I in 1:n + 1
            println("Iter: $I; "*"$(cg())")
        end
        return norm(b - A*cg.x) < 1e-8

    end
    function Test2()
        n = 5
        A = rand(n, n)
        A = A'*A
        display(A)
        b = rand(n)
        cg = ConjGradModified(A, b)
        for I in 1:n + 1
            println("Iter: $I; "*"$(cg())")
            ForceOrthogonalize(cg)
        end
        return norm(b - A*cg.x) < 1e-8
    end

    @test Test1()
    @test Test2()
end


using Geodesics
using Test

@testset "Forward" begin
    @test_throws ArgumentError Geodesics.forward(0, 2π, 0, 0, 1, 0)
    @test_throws ArgumentError Geodesics.forward(0, 2π, 0, 0, 0, 0)
    @test_throws ArgumentError Geodesics.inverse(0, 0, 1, 1, 1, 1)

    # Compare to Haversine computation on the unit sphere
    let lon0 = 0, lat0 = 0, azimuth = 1, distance = 1, a = 1, f = 0
        lon1, lat1, baz = Geodesics.forward(lon0, lat0, azimuth, distance, a, f)
        @test lon1 ≈ 0.9189892552937267
        @test lat1 ≈ 0.4719777676633856
        @test baz ≈ mod(-1.9047280216911358, 2π) atol=1e-6
    end

    # Make sure if we don't move, we end up in the same place
    let lon0 = 0.1, lat0 = 0.2, azimuth = 1, distance = 0, a = 1, f = 0
        lon1, lat1, baz = Geodesics.forward(lon0, lat0, azimuth, distance, a, f)
        @test lon1 == lon0
        @test lat1 == lat0
        @test baz ≈ mod(azimuth + π, 2π)
    end
end

@testset "Inverse" begin
    @test_throws ArgumentError Geodesics.inverse(0, 0, 1, 1, 0, 0)
    @test_throws ArgumentError Geodesics.inverse(0, 2π, 0, -3π, 1, 1)
    @test_throws ArgumentError Geodesics.inverse(0, 0, 1, 1, 1, 1)

    # Compare to Haversine computation on the unit sphere
    let lon0 = 0, lat0 = 0, lon1 = 1, lat1 = 1, a = 1, f = 0, atol = 1e-6
        distance, azimuth, backazimuth = Geodesics.inverse(lon0, lat0, lon1, lat1, a, f)
        @test distance ≈ 1.2745557823062943 atol=atol
        @test azimuth ≈ 0.49536728921867335 atol=atol
        @test backazimuth ≈ mod(-2.0661636160135703, 2π) atol=atol
    end
end

@testset "Ang dist" begin
    # Degrees
    let lon0 = 360*rand() - 180, lat0 = 180*rand() - 90, lon1 = 360*rand() - 180,
            lat1 = 180*rand() - 90, atol = 1e-6
        @test Geodesics.angular_distance(lon0, lat0, lon1, lat1, f=0) ≈
            rad2deg(Geodesics.inverse(deg2rad.((lon0, lat0, lon1, lat1))..., 1, 0)[1]) atol=atol
        @test Geodesics.angular_distance(lon0, lat0, lon1, lat1) ≈
            rad2deg(Geodesics.inverse(deg2rad.((lon0, lat0, lon1, lat1))..., 1, Geodesics.F_WGS84)[1]) atol=atol
    end

    # Radians
    let lon0 = 2π*(rand() - 0.5), lat0 = π*(rand() - 0.5), lon1 = 2π*(rand() - 0.5),
            lat1 = π*(rand() - 0.5), atol = 1e-6
        @test Geodesics.angular_distance(lon0, lat0, lon1, lat1, false, f=0) ≈
            Geodesics.inverse(lon0, lat0, lon1, lat1, 1, 0)[1] atol=atol
        @test Geodesics.angular_distance(lon0, lat0, lon1, lat1, false) ≈
            Geodesics.inverse(lon0, lat0, lon1, lat1, 1, Geodesics.F_WGS84)[1] atol=atol
    end
end

@testset "Azimuth" begin
    # Degrees
    let lon0 = 360*rand() - 180, lat0 = 180*rand() - 90, lon1 = 360*rand() - 180,
            lat1 = 180*rand() - 90, atol = 1e-6
        @test Geodesics.azimuth(lon0, lat0, lon1, lat1, f=0) ≈
            rad2deg(Geodesics.inverse(deg2rad.((lon0, lat0, lon1, lat1))..., 1, 0)[2]) atol=atol
        @test Geodesics.azimuth(lon0, lat0, lon1, lat1) ≈
            rad2deg(Geodesics.inverse(deg2rad.((lon0, lat0, lon1, lat1))..., 1, Geodesics.F_WGS84)[2]) atol=atol
    end

    # Radians
    let lon0 = 2π*(rand() - 0.5), lat0 = π*(rand() - 0.5), lon1 = 2π*(rand() - 0.5),
            lat1 = π*(rand() - 0.5), atol = 1e-6
        @test Geodesics.azimuth(lon0, lat0, lon1, lat1, false, f=0) ≈
            Geodesics.inverse(lon0, lat0, lon1, lat1, 1, 0)[2] atol=atol
        @test Geodesics.azimuth(lon0, lat0, lon1, lat1, false) ≈
            Geodesics.inverse(lon0, lat0, lon1, lat1, 1, Geodesics.F_WGS84)[2] atol=atol
    end

    # One point at pole
    let lon0 = 0, lat0 = -90, lon1 = 0, lat1 = 0
        for T in (Int16, Int32, Int64, Float16, Float32, Float64)
            @test Geodesics.azimuth(T(lon0), T(lat0), T(lon1), T(lat1), true) == 0.0
        end
    end
    let lon0 = 90, lat0 = 0, lon1 = 90, lat1 = 90
        for T in (Int16, Int32, Int64, Float16, Float32, Float64)
            @test Geodesics.azimuth(T(lon0), T(lat0), T(lon1), T(lat1), true) == 0.0
        end
    end
end

@testset "Ang step" begin
    # Degrees
    let lon0 = 360*rand() - 180, lat0 = 180*rand() - 90, dist = 360*rand(),
            az = 360*rand() - 90, atol = 1e-6
        @test all(isapprox.(Geodesics.angular_step(lon0, lat0, az, dist, f=0),
            rad2deg.(Geodesics.forward(deg2rad.((lon0, lat0, az, dist))..., 1, 0)), atol=atol))
        @test all(isapprox.(Geodesics.angular_step(lon0, lat0, az, dist),
            rad2deg.(Geodesics.forward(deg2rad.((lon0, lat0, az, dist))..., 1, Geodesics.F_WGS84)), atol=atol))
    end

    # Radians
    let lon0 = 2π*(rand() - 0.5), lat0 = π*(rand() - 0.5), dist = 2π*rand(),
            az = 2π*(rand() - 0.5), atol = 1e-6
        @test all(isapprox.(Geodesics.angular_step(lon0, lat0, az, dist, false, f=0),
            Geodesics.forward(lon0, lat0, az, dist, 1, 0), atol=atol))
        @test all(isapprox.(Geodesics.angular_step(lon0, lat0, az, dist, false),
            Geodesics.forward(lon0, lat0, az, dist, 1, Geodesics.F_WGS84), atol=atol))
    end
end

@testset "Surf dist" begin
    # Degrees
    let lon0 = 360*rand() - 180, lat0 = 180*rand() - 90, lon1 = 360*rand() - 180,
            lat1 = 180*rand() - 90, a = 10_000_000*rand(), atol = 1e-6
        @test Geodesics.surface_distance(lon0, lat0, lon1, lat1, a, f=0) ≈
            Geodesics.inverse(deg2rad.((lon0, lat0, lon1, lat1))..., a, 0)[1] atol=atol
        @test Geodesics.surface_distance(lon0, lat0, lon1, lat1, a) ≈
            Geodesics.inverse(deg2rad.((lon0, lat0, lon1, lat1))..., a, Geodesics.F_WGS84)[1] atol=atol
    end

    # Radians
    let lon0 = 2π*(rand() - 0.5), lat0 = π*(rand() - 0.5), lon1 = 2π*(rand() - 0.5),
            lat1 = π*(rand() - 0.5), a = 10_000_000*rand(), atol = 1e-6
        @test Geodesics.surface_distance(lon0, lat0, lon1, lat1, a, false, f=0) ≈
            Geodesics.inverse(lon0, lat0, lon1, lat1, a, 0)[1] atol=atol
        @test Geodesics.surface_distance(lon0, lat0, lon1, lat1, a, false) ≈
            Geodesics.inverse(lon0, lat0, lon1, lat1, a, Geodesics.F_WGS84)[1] atol=atol
    end
end


# Part 1
program = map(i -> [i[1], parse(Int32, i[2])], split.(readlines("day08_input.txt")))

# Execute an instruction in prog, returning next instruction index and acc
function exec_inst(prog, instr, acc)
    instruction = prog[instr]
    handlers = Dict(
        "acc" => (instr, acc, arg) -> (instr + 1, acc + arg),
        "jmp" => (instr, acc, arg) -> (instr + arg, acc),
        "nop" => (instr, acc, arg) -> (instr + 1, acc)
    )
    handlers[instruction[1]](instr, acc, instruction[2])
end

# Execute instructions in prog until an instruction is executed twice
# If loop detected, return acc, false
# If about to exec 1 past the instruction array, return acc, true
function exec_until_loop(prog, acc=0)
    next_instr = 1
    executed = Set()
    while next_instr ∉ executed
        push!(executed, next_instr)
        next_instr, acc = exec_inst(prog, next_instr, acc)
        if next_instr == length(prog) + 1 return acc, true end
    end
    acc, false
end

println("Part 1: accumulator=$(first(exec_until_loop(program)))")

# Part 2
# Patch each instruction in the program (swapping jmp <-> nop) until it exits normally
function try_patches(prog)
    for i in 1:length(prog)
        patched = deepcopy(prog)
        patched[i][1] = patched[i][1] |> x -> (x == "jmp" ? "nop" : (x == "nop" ? "jmp" : x))
        acc, good = exec_until_loop(patched)
        if good return acc end
    end
end

println("Part 2: accumulator=$(try_patches(program))")


function δ(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E<:AllowedSpeciesTypes}
    for s in species(N)
        @assert haskey(L, s)
    end

    this_l = filter(p -> p.first in species(N), L)
    tl = [this_l[s] for s in species(N; dims=1)]
    bl = [this_l[s] for s in species(N; dims=2)]
    D = tl .== bl'
    return D
end

"""
    Q(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E<:AllowedSpeciesTypes}

Modularity of a network and its partition. The second argument is a dictionary
where every species of `N` is associated to an `Int64` value representing the
identity of the module. This function returns the same value of bipartite
networks and their unipartite projection.
"""
function Q(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E<:AllowedSpeciesTypes}
    for s in species(N)
        @assert haskey(L, s)
    end

    # Degrees
    dkin, dkout = degree_in(N), degree_out(N)
    kin = map(x -> dkin[x], species(N; dims=2))
    kout = map(x -> dkout[x], species(N; dims=1))

    # Value of m -- sum of weights, total number of int, ...
    m = links(N)

    # Null model
    kikj = (kout .* kin')
    Pij = kikj ./ m

    # Difference
    diff = N.A .- Pij

    # Diff × delta
    dd = diff .* δ(N,L)

    return sum(dd)/m
end

"""
    Qr(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E<:AllowedSpeciesTypes}

Realized modularity -- this function returns a value giving the proportion of
all links that are within the same module. Higher values reflect a more strongly
modular partition (whereas `Q` represents the deviation of modularity from the
random expectation).
"""
function Qr(N::T, L::Dict{E,Int64}) where {T<:AbstractEcologicalNetwork,E<:AllowedSpeciesTypes}
    for s in species(N)
        @assert haskey(L, s)
    end
    W = sum(N.A .* δ(N, L))
    B = links(N)
    return 2.0 * (W/B) - 1.0
end

function tidy_modules!(L::Dict{E,Int64}) where {E<:AllowedSpeciesTypes}
  l_values = sort(unique(collect(values(L))))
  for i in 1:length(l_values)
    for (k,v) in L
      if v == l_values[i]
        L[k] = i
      end
    end
  end
end


immutable Point
  x::Int
  y::Int
end

t = Point(10, 10)
println(t)

function norm(p::Point)
  return sqrt(p.x^2+p.y^2)
end

println(norm(Point(10, 20)))


using Test
using BranchAndBound
using JuMP, GLPK
using MathOptInterface

const BB = BranchAndBound
const MOI = MathOptInterface

@testset "Abstract Node" begin
    struct MyNode <: BB.AbstractNode end
    node = MyNode()
    @test BB.branch(node) == []
    BB.bound!(node)
    BB.heuristics!(node)
    BB.apply_changes!(node)
end

model = Model(GLPK.Optimizer)
@variable(model, 0 <= x <= 1)
@objective(model, Min, x)

tree = BB.AbstractTree()
node = BB.JuMPNode{BB.VariableBranch}(model)

@testset "Initial Tree" begin
    @test tree.node_counter == 0
    @test length(tree.nodes) == 0
    @test length(tree.processed) == 0
end

@testset "Initial Node" begin
    @test node.id == -1
    @test isnothing(node.parent)
    @test node.depth == 0
    @test isnothing(node.branch)
    @test node.solution_status == MOI.OPTIMIZE_NOT_CALLED
    @test isinf(-node.bound)
    @test !isnothing(node.model)
end

@testset "Tree with root" begin
    BB.push!(tree, node)
    @test tree.node_counter == 1
    @test length(tree.nodes) == 1
    @test length(tree.processed) == 0
    @test tree.nodes[1] == node

    tree1 = BB.initialize_tree(model)
    @test tree1.node_counter == 1
    @test length(tree1.nodes) == 1
    @test length(tree1.processed) == 0
end

current_node = BB.next_node(tree)

@testset "Root node with model" begin
    @test !isnothing(current_node.model)
    @test JuMP.lower_bound(x) == 0.
    @test JuMP.upper_bound(x) == 1.
end

@testset "Bounding root node" begin
    BB.bound!(current_node)
    @test current_node.solution_status == MOI.OPTIMAL
    @test current_node.bound == 0.0
    @test JuMP.value(x) == 0.0
end

@testset "Marked root node as processed" begin
    BB.processed!(tree, current_node)
    @test BB.termination(tree)
    @test length(tree.processed) == 1
    @test tree.processed[1] == current_node
end

@testset "Empty branching" begin
    BB.branch!(tree, current_node)
    @test tree.node_counter == 1
    @test length(tree.nodes) == 0
    @test length(tree.processed) == 1
end

@testset "Branching at root" begin
    child1 = BB.create_child_node_with_lb(current_node, x, 0.5)
    child2 = BB.create_child_node_with_ub(current_node, x, 0.5)
    BB.push!(tree, [child1,child2])
    @test tree.node_counter == 3
    @test length(tree.nodes) == 2
    @test length(tree.processed) == 1
end

next_node = BB.next_node(tree)

@testset "Child 1 with model" begin
    @test !isnothing(next_node.model)
    print(next_node.model)
    @test JuMP.lower_bound(x) == 0.
    @test JuMP.upper_bound(x) == 0.5
    @test next_node.bound == 0.0
end

@testset "Bounding Child 1 node" begin
    BB.bound!(next_node)
    @test next_node.solution_status == MOI.OPTIMAL
    @test next_node.bound == 0.0
    @test JuMP.value(x) == 0.0
end


# Tests for Profit DEA Models
@testset "ProfitDEAModel" begin

    ## Test Profit DEA Model
    # Test with Zofio, Pastor and Aparicio (2013) data
    X = [1 1; 1 1; 0.75 1.5; 0.5 2; 0.5 2; 2 2; 2.75 3.5; 1.375 1.75]
    Y = [1 11; 5 3; 5 5; 2 9; 4 5; 4 2; 3 3; 4.5 3.5]
    P = [2 1; 2 1; 2 1; 2 1; 2 1; 2 1; 2 1; 2 1]
    W = [2 1; 2 1; 2 1; 2 1; 2 1; 2 1; 2 1; 2 1]

    GxGydollar = 1 ./ (sum(P, dims = 2) + sum(W, dims = 2))
    GxGydollar = repeat(GxGydollar, 1, 2)

    deaprofitdollar = deaprofit(X, Y, W, P, GxGydollar, GxGydollar)
    @test efficiency(deaprofitdollar, :Economic)   ≈ [2; 2; 0; 2; 2; 8; 12; 4] atol = 1e-3
    @test efficiency(deaprofitdollar, :Technical)  ≈ [0; 0; 0; 0; 0; 6; 12; 3] atol = 1e-3
    @test efficiency(deaprofitdollar, :Allocative) ≈ [2; 2; 0; 2; 2; 2; 0; 1] atol = 1e-3

    # Print
    show(IOBuffer(), deaprofitdollar)

    # Test errors
    @test_throws ErrorException deaprofit([1; 2 ; 3], [4 ; 5], [1; 1; 1], [4; 5], [1; 2 ; 3], [4 ; 5]) #  Different number of observations
    @test_throws ErrorException deaprofit([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], [4; 5; 6], [1; 2; 3], [4; 5; 6]) # Different number of observation in input prices
    @test_throws ErrorException deaprofit([1; 2; 3], [4; 5; 6], [1; 2; 3], [4; 5; 6; 7], [1; 2; 3], [4; 5; 6]) # Different number of observation in output prices
    @test_throws ErrorException deaprofit([1 1; 2 2; 3 3], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], [4; 5; 6], [1 1; 2 2; 3 3], [4; 5; 6]) # Different number of input prices and inputs
    @test_throws ErrorException deaprofit([1; 2; 3], [4 4; 5 5; 6 6], [1; 2; 3], [4 4 4; 5 5 5; 6 6 6], [1; 2; 3], [4 4; 5 5; 6 6]) # Different number of oputput prices and outputs
    @test_throws ErrorException deaprofit([1 1; 2 2; 3 3], [4; 5; 6], [1 1; 2 2; 3 3], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], [4; 5; 6]) # Different size of inputs direction
    @test_throws ErrorException deaprofit([1; 2; 3], [4 4; 5 5; 6 6], [1; 2; 3], [4 4; 5 5; 6 6], [1; 2; 3], [4 4 4; 5 5 5; 6 6 6]) # Different size of inputs direction

end


using SolitonBEC
using Printf
using Plots
using Elliptic
hbar = 6.62607015e-34 / (2 * pi)
width = 1.42e-6
mass = 6.941 * 1.660539e-27
omega_perp = 2 * pi * 710
l_perp = sqrt(hbar / (mass * omega_perp))
N = 4e3
as = -0.21e-9

interaction_g = abs(2*hbar^2 * as / mass / l_perp^2)
ggg = 2*hbar*omega_perp*abs(as)

print("\ninteraction_g: ", interaction_g)

# mass as total mass or mass per atom??
time_unit = hbar^2/(mass * ggg * N)
space_unit = hbar^3/(mass * ggg^2 * N^2)
velocity_unit = ggg * N / hbar

# --------- Simulation ---------
khaykovich_gpe = Simulation(
  "GPE",
  "barrier",
)

# --------- Apparata ---------
std_apparatus = Apparatus(
  mass, #m (conversion AMU -> kg)
  as, # as
  omega_perp, # ω_perp
  N, #N
  1.77e-11, #γ
  2 * pi * 4,# ω_z
)

# Energy = 1 / 2 * (std_apparatus.as^2 / (width^2) + std_apparatus.as^2 / (2 * l_perp^2) + width^2 / std_apparatus.as^2 - interaction_g * std_apparatus.as^3 / (l_perp * width^2)) * hbar * omega_perp
L = 2 * space_unit

velocity_unit = ggg*N/hbar
print("\n\tvelocity unit: ", velocity_unit, " m/s")
v0 = 1*velocity_unit
x0 =-L/4
# --------- InitialStates ---------
InitialState1 = InitialState(
  "sech", #type
  width*1.8, # width
  v0, #v0
  x0,
)

# --------- Numerics ---------
T = abs(InitialState1.x0)/ v0 * 2
print("\n\ttime_unit: ", time_unit)
print("\n\tspace_unit: ", space_unit)
num = Numerics(
  T, #T
  T * 1e-3, #dt
  L, #S
  L * 1e-3, #ds
)

#--------- Potential ---------

function barrier_height(energy::Float64)
  r = Potential(
    "barrier", # type
    L * 1e-4, #width
    0, #position
    energy, #energy
    0, #ϵ
    10e-6 #a
  )
  return r
end

energy_unit = mass * N^2 * ggg^2 / hbar^2
print("\n\tenergy unit: ", energy_unit, " J")
print("\n==> desired group velocity: ", v0, " m/s\n\n")

GSEnergy = energy_unit * (-N/24)
#print("\nground state energy: ", GSEnergy/hbar, " hbar\n")
phys_vel = 4.918e-3
normd_vel = phys_vel*hbar/ggg/N
Energy = GSEnergy + energy_unit * normd_vel^2*N/2
#print("\ntraveling state energy: ", Energy/hbar, " hbar\n")

## Configurations
configs = []
for barrier_energy in LinRange(0, 1, 3)
  potential = barrier_height(barrier_energy * energy_unit)
  push!(configs, (num, khaykovich_gpe, std_apparatus, potential, InitialState1))
end

mem_limit = 15000000000 #byte
cnt = 1

pyplot()
p = Plots.palette(:rainbow_bgyr_35_85_c72_n256, length(configs) + 3)

## Soliton - barrier collision --------------------------------------------
run_dynamics(configs[2]...)

fig1 = plot(title="|ψ|^2 after collision with barrier",
  xlabel="space [mm]",
  ylabel="|ψ|^2",
  reuse=false,
  size=(800, 400),
  legend=:topleft)

for (num, sim, app, pot, state) in configs

  global cnt
  @printf("\n------Running simulation # %3i \n", cnt)
  estimate = mem_estimate(num)

  if estimate < mem_limit
    #@printf("Estimated memory consumption: %4.1f MiB\n", estimate / (1024^2))

    coeffs = get_coefficients(sim, app, pot, state)

    time, space, ψ, ψ_spect = @time ssfm_solve(num, coeffs)
    #plot_dynamics(time, space, ψ, ψ_spect)
    peak, idx = findmax(abs.(ψ[:, end]).^2)

    # plot!(space*1e3, abs.(ψ[:, end]).^2)

    print("peak: ", peak)
    group_velocity = (space[idx]-state.x0)/time[length(time)]
    print("\n==> measured group velocity: ", group_velocity, " m/s\n")

    print("\n ratio group velocity: ", state.v0/group_velocity, "\n")

    plot!(space * 1e3,
      abs.(ψ[:, end]) .^ 2,
      label="peak potential = $(pot.energy / (sqrt(2 * pi) * pot.width) /1.602176634e-19)   [eV]",
      lw=1,
      color=p[cnt])

    # plot!(space * 1e3,
    #   abs.(map(s -> potential(sim, app, pot, s), space)) / hbar,
    #   lw=1,
    #   ls=:dot,
    #   label=nothing,
    #   color=p[cnt])

    else
    @printf("Estimated memory consumption (%4.1f MiB) exceed maximum!\n", estimate / (1024^2))
  end

  cnt += 1

end
display(fig1)


include("EllCrv/EllCrv.jl")
include("EllCrv/Misc.jl")
include("EllCrv/QQ.jl")
include("EllCrv/Finite.jl")


module SimradEK60

using SimradRaw

export Sv, TS, pings, power, powerdb,  athwartshipangle, alongshipangle, R, filetime, power2db, octet2deg

# Float32 is okay for most practical purposes and within 10^-5 of EchoView
# Float64 gets to 10^-7 of Echoview
FLOAT_TYPE = Float32

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

# We split the RAW file into a series of EK60 specific pings.

# Whilst it is tempting to include latitude and longitude in the ping,
# this requires NMEA parsing and decisions about interpolation. Some
# people use external GPSs. Since the processing approach could be a
# user preference and potential source of errors, we choose not to
# include them. A caller can easily access NMEA datagrams produced by
# SimradRaw.jl and call PyNmea.

struct EK60Ping
    power::Vector{Int16}
    athwartshipangle::Vector{Int8}
    alongshipangle::Vector{Int8}
    filetime::UInt64
    frequency::FLOAT_TYPE
    soundvelocity::FLOAT_TYPE
    sampleinterval::FLOAT_TYPE
    absorptioncoefficient::FLOAT_TYPE
    transmitpower::FLOAT_TYPE
    pulselength::FLOAT_TYPE
    gain::FLOAT_TYPE
    equivalentbeamangle::FLOAT_TYPE
    sacorrection::FLOAT_TYPE
    sacorrectiontable::Vector{FLOAT_TYPE}
    pulselengthtable::Vector{FLOAT_TYPE}
end

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

function load(filename::AbstractString)
    ps = collect(pings(filename))
end

"""
    R(r, s, T)

TVG range correction for Ex60

R is the corrected range (m)

r the uncorrected range (m)

s is the TvgRangeCorrectionOffset value, See Echoview documentation,
Simrad Time Varied Gain (TVG) range correction, http://bit.ly/2pqzS2D

T is the sample thickness (m)
"""
R(r, s, T) = max.(0f0, r .- s*T)

"""
    R(ping::EK60Ping; soundvelocity = nothing, rangecorrectionoffset=2)

Returns the corrected range (depth) of samples in `ping`.
"""
function R(ping::EK60Ping; soundvelocity = nothing, rangecorrectionoffset=2)

    if soundvelocity == nothing
        soundvelocity = ping.soundvelocity
    end

    p = ping.power

    samplethickness = soundvelocity .* ping.sampleinterval / 2 # in metres of range

    r = [x* samplethickness for x in 0:length(p)-1]
    R(r, rangecorrectionoffset, samplethickness)

end

function R(pings::Vector{EK60Ping}; soundvelocity = nothing, rangecorrectionoffset=2)

    r = [R(ping,
           soundvelocity=soundvelocity,
           rangecorrectionoffset=rangecorrectionoffset) for ping in pings]

    myhcat(r)

end

R(pings::Channel{EK60Ping}; soundvelocity = nothing, rangecorrectionoffset=2) =
    R(collect(pings),
      soundvelocity=soundvelocity,
      rangecorrectionoffset=rangecorrectionoffset)

###

"""
    filetime(ping::EK60Ping)

Returns the FILETIME timestamp for a ping
"""
filetime(ping::EK60Ping) = ping.filetime

function filetime(pings::Vector{EK60Ping})
    [ping.filetime for ping in pings]
end

filetime(pings::Channel{EK60Ping}) = filetime(collect(pings))

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

# We must apply the SONAR equation and instrument corrections as
# described in http://bit.ly/2o1oOrq

## % An earlier MATLAB implementation
## % Sv = recvPower + 20 log10(Range) + (2 *  alpha * Range) - (10 * ...
## %           log10((xmitPower * (10^(gain/10))^2 * lambda^2 * ...
## %           c * tau * 10^(psi/10)) / (32 * pi^2)) - (2 * SaCorrection)

"""
    Sv(Pr, λ, G, Ψ, c, α, Pt, τ, Sa, R)
where:

R = the corrected range (m) - see TVG Range Correction

Pr = received power (dB re 1 W) - see Simrad EK numbers to Power

Pt = transmitted power (W)

α = absorption coefficient (dB/m)

G0 = transducer peak gain (non-dimensional) calculated as 10^(G/10)

G is the Transducer gain (dB re 1)

λ = wavelength (m) = c/f

f = frequency (Hz)

c = sound speed (m/s)

τ = transmit pulse duration (s) - also known as the
TransmittedPulseLength

ψ = Equivalent Two-way beam angle (Steradians) calculated as 10^(Ψ/10)

Ψ is the Two-way beam angle (dB re 1 Steradian)

Sa = Simrad correction factor (dB re 1m−1) determined during
calibration. This represents the correction required to the Sv
constant to harmonize the TS and NASC measurements.

"""
function Sv(Pr, λ, G, Ψ, c, α, Pt, τ, Sa, R)

    # TVG is applied to samples with ranges greater than 1 meter.
    tvg =  20log10.(max.(1, R))

    csv = 10log10.((Pt * (10^(G/10))^2 *  λ^2 * c * τ * 10^(Ψ/10)) /
                   (32 * FLOAT_TYPE(pi)^2))

    # Ignore absorption for R < 1m
    r =  [(x < 1) ? 0 : x for x in R]

    Pr .+ tvg + (2 * α * r) .- csv .- 2Sa
end

"""
     TS(Pr, λ, G, α, Pt, R)

Target strength

Pr = received power (dB re 1 W) - see Simrad EK numbers to Power

G is the Transducer gain (dB re 1)

α = absorption coefficient (dB/m)

Pt = transmitted power (W)

R = the corrected range (m) - see TVG Range Correction

"""
function TS(Pr, λ, G, α, Pt, R)

    # TVG is applied to samples with ranges greater than 1 meter.
    tvg =  40log10.(max.(1, R))

    csv = 10log10.((Pt * (10^(G/10))^2 *  λ^2) /
                   (16 * FLOAT_TYPE(pi)^2))

    # Ignore absorption for R < 1m
    r =  [(x < 1) ? 0 : x for x in R]

    Pr .+ tvg + (2 * α * r) .- csv
end

"""
    pings(datagrams::Channel{SimradRaw.Datagram})

Returns a `Channel` of `Ping`s from a `Channel` of `Datagram`s.

"""
function pings(datagrams::Channel{SimradRaw.Datagram})

    function _it(chn1)
        config = nothing
        for datagram in datagrams
            if datagram.dgheader.datagramtype == "RAW0"

                transducer = config.configurationtransducers[datagram.channel]

                if transducer.frequency != datagram.frequency
                    @error("$(config.configurationtransducer)  $datagram")
                end

                idx = findfirst(datagram.pulselength .== transducer.pulselengthtable)

                sacorrection = transducer.sacorrectiontable[idx]
                gain = transducer.gaintable[idx]

                ping = EK60Ping(datagram.power,
                                datagram.athwartshipangle,
                                datagram.alongshipangle,
                                SimradRaw.filetime(datagram.dgheader.datetime),
                                datagram.frequency,
                                datagram.soundvelocity,
                                datagram.sampleinterval,
                                datagram.absorptioncoefficient,
                                datagram.transmitpower,
                                datagram.pulselength,
                                gain,
                                transducer.equivalentbeamangle,
                                sacorrection,
                                transducer.sacorrectiontable,
                                transducer.pulselengthtable)

                put!(chn1, ping)

            elseif datagram.dgheader.datagramtype == "CON0"
                config = datagram
            end
        end
    end

    return Channel(_it, ctype=EK60Ping)
end

function pings(datagrams::Vector{SimradRaw.Datagram})
    # Is there a neater method here?
    chnl = Channel(c->foreach(i->put!(c,i), datagrams), ctype=SimradRaw.Datagram)
    pings(chnl)
end

"""
    pings(filename::AbstractString)

Returns a `Channel` over `Ping`s from the RAW file designated by
`filename`.

"""
function pings(filename::AbstractString)
    pings([filename])
end

"""
    pings(filenames::Vector{AbstractString}})

Returns a `Channel` of `Pings` from the RAW files designated by
`filenames`.

"""
function pings(filenames::Vector{String})
    pings(datagrams(filenames))

end

"""
    Sv(ping::EK60Ping;
            frequency=nothing,
            gain=nothing,
            equivalentbeamangle=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            pulselength=nothing,
            sacorrection=nothing,
            rangecorrectionoffset=2)

Returns a `Vector` of Sv, the (Mean) Volume backscattering strength (MVBS) in (dB re
1 m-1) for a given `ping`.

The function accepts a number of optional arguments which, if
specified, override the ping's own settings. This facilitates external
calibration.

"""
function Sv(ping::EK60Ping;
            frequency=nothing,
            gain=nothing,
            equivalentbeamangle=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            pulselength=nothing,
            sacorrection=nothing,
            rangecorrectionoffset=2)

    if frequency == nothing
        frequency = ping.frequency
    end

    if gain == nothing
        gain = ping.gain
    end

    if equivalentbeamangle == nothing
        equivalentbeamangle = ping.equivalentbeamangle
    end

    if soundvelocity == nothing
        soundvelocity = ping.soundvelocity
    end

    if absorptioncoefficient == nothing
        absorptioncoefficient = ping.absorptioncoefficient
    end

    if transmitpower == nothing
        transmitpower= ping.transmitpower
    end

    if pulselength == nothing
        pulselength = ping.pulselength
    end

    if sacorrection == nothing
        sacorrection = ping.sacorrection
    end

    pdb = powerdb(ping)

    rangecorrected = R(ping,
                       soundvelocity = soundvelocity,
                       rangecorrectionoffset=rangecorrectionoffset)

    λ =  soundvelocity / frequency # calculate wavelength

    Sv(pdb, λ, gain, equivalentbeamangle,
            soundvelocity, absorptioncoefficient,
            transmitpower, pulselength, sacorrection,
            rangecorrected)

end

"""
    TS(ping::EK60Ping;
            frequency=nothing,
            gain=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            rangecorrectionoffset=0)

Target strength

"""
function TS(ping::EK60Ping;
            frequency=nothing,
            gain=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            rangecorrectionoffset=0)

    if frequency == nothing
        frequency = ping.frequency
    end

    if gain == nothing
        gain = ping.gain
    end

    if soundvelocity == nothing
        soundvelocity = ping.soundvelocity
    end

    if absorptioncoefficient == nothing
        absorptioncoefficient = ping.absorptioncoefficient
    end

    if transmitpower == nothing
        transmitpower= ping.transmitpower
    end

    pdb = powerdb(ping)

    rangecorrected = R(ping,
                       soundvelocity = soundvelocity,
                       rangecorrectionoffset = rangecorrectionoffset)

    λ =  soundvelocity / frequency # calculate wavelength

    TS(pdb, λ, gain, absorptioncoefficient, transmitpower, rangecorrected)

end

function myhcat(s; missingvalue=NaN32)
    ls = map(length,s)
    minlength =minimum(ls)
    maxlength =maximum(ls)
    if minlength == maxlength
        return hcat(s...)
    end

    etype = typeof(missingvalue)

    array = Array{etype}(undef, maxlength, length(s))

    array .= missingvalue
    for i in 1:length(s)
        r = 1:length(s[i])
        array[r, i] .= s[i][r]
    end
    return array
end

"""
    Sv(pings::Vector{EK60Ping};
        frequency=nothing,
        gain=nothing,
        equivalentbeamangle=nothing,
        soundvelocity=nothing,
        absorptioncoefficient=nothing,
        transmitpower=nothing,
        pulselength=nothing,
        sacorrection=nothing,
        rangecorrectionoffset=2)

Returns an `Array` of Sv, the Volume backscattering strength in (dB re
1 m-1) for a set of contiguous `pings`.

The function accepts a number of optional arguments which, if
specified, override the pings' own settings. This facilitates external
calibration.

"""
function Sv(pings::Vector{EK60Ping};
            frequency=nothing,
            gain=nothing,
            equivalentbeamangle=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            pulselength=nothing,
            sacorrection=nothing,
            rangecorrectionoffset=2)
    s = [Sv(ping,
            frequency=frequency,
            gain=gain,
            equivalentbeamangle=equivalentbeamangle,
            soundvelocity=soundvelocity,
            absorptioncoefficient=absorptioncoefficient,
            transmitpower=transmitpower,
            pulselength=pulselength,
            sacorrection=sacorrection,
            rangecorrectionoffset=rangecorrectionoffset) for ping in pings]

    myhcat(s)
end

"""
    TS(pings::Vector{EK60Ping};
            frequency=nothing,
            gain=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            rangecorrectionoffset=0)

Target strength
"""
function TS(pings::Vector{EK60Ping};
            frequency=nothing,
            gain=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            rangecorrectionoffset=0)
    s = [TS(ping,
            frequency=frequency,
            gain=gain,
            soundvelocity=soundvelocity,
            absorptioncoefficient=absorptioncoefficient,
            transmitpower=transmitpower,
            rangecorrectionoffset=rangecorrectionoffset) for ping in pings]

    myhcat(s)
end

function Sv(pings::Channel{EK60Ping};
            frequency=nothing,
            gain=nothing,
            equivalentbeamangle=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            pulselength=nothing,
            sacorrection=nothing,
            rangecorrectionoffset=2)
    Sv(collect(pings),
       frequency=frequency,
       gain=gain,
       equivalentbeamangle=equivalentbeamangle,
       soundvelocity=soundvelocity,
       absorptioncoefficient=absorptioncoefficient,
       transmitpower=transmitpower,
       pulselength=pulselength,
       sacorrection=sacorrection,
       rangecorrectionoffset=rangecorrectionoffset)
end

function TS(pings::Channel{EK60Ping};
            frequency=nothing,
            gain=nothing,
            soundvelocity=nothing,
            absorptioncoefficient=nothing,
            transmitpower=nothing,
            rangecorrectionoffset=0)
    TS(collect(pings),
       frequency=frequency,
       gain=gain,
       soundvelocity=soundvelocity,
       absorptioncoefficient=absorptioncoefficient,
       transmitpower=transmitpower,
       rangecorrectionoffset=rangecorrectionoffset)
end

#

"""
    athwartshipangle(ping::EK60Ping)

Returns the athwartship phase angle difference vector for the given
ping.

"""
athwartshipangle(ping::EK60Ping) = ping.athwartshipangle

function athwartshipangle(pings::Vector{EK60Ping})
    s = [ping.athwartshipangle for ping in pings]
    myhcat(s, missingvalue=Int8(0))
end

athwartshipangle(pings::Channel{EK60Ping}) = athwartshipangle(collect(pings))

@deprecate(athwartshipanglematrix, athwartshipangle)

#

"""
    alongshipangle(ping::EK60Ping)

Returns the alongship phase angle difference vector for the given
ping.

"""
alongshipangle(ping::EK60Ping) = ping.alongshipangle

function alongshipangle(pings::Vector{EK60Ping})
    s = [ping.alongshipangle for ping in pings]
    myhcat(s, missingvalue=Int8(0))
end

alongshipangle(pings::Channel{EK60Ping}) = alongshipangle(collect(pings))

@deprecate(alongshipanglematrix, alongshipangle)

#

"""
    power(ping::EK60Ping)

Returns power data for a ping in manufacturer units. For decibels,
see `powerdb`.

"""
power(ping::EK60Ping) = ping.power

function power(pings::Vector{EK60Ping})
    s = [ping.power for ping in pings]
    myhcat(s, missingvalue=Int16(0))
end

power(pings::Channel{EK60Ping}) = power(collect(pings))

#

const POWER_MULTIPLIER = FLOAT_TYPE(10 * log10(2) / 256)

"""
    power2db(x)

EK60 power data is stored in a "compressed", manufacturer specific
format. This function converts to decibels.

"""
power2db(x) = x * POWER_MULTIPLIER

"""
    powerdb(x)

Returns power data in decibels. x is one or more `Ping`s.

"""
powerdb(x) = power2db(power(x))

"""
    octet2deg(x)

EK60 split beam angles are stored as signed octets with a resolution
of 180/128.  Convert from octet to degrees.
"""
octet2deg(x) = x * 180/128

end # module


# sqfChol.jl

export sqfChol_to_Σ, PDMat_adj;

## Note this follows my unconventional definition of the square-root-free
# Cholesky decomposition with construction from the back of the vector.
function sqfChol_to_Σ(β::Array{Array{T, 1}, 1}, δ::Array{T, 1},
    pdthrowerror::Bool=true) where T <: Real

    sqrtΔ = Diagonal(sqrt.(δ))
    βmat = BayesInference.xpnd_tri(vcat(β...), false)'

    βinvSqrtΔ = Matrix(βmat \ sqrtΔ)
    Σ = βinvSqrtΔ * βinvSqrtΔ'

    PDMat_adj(Σ, throwerror=pdthrowerror)
end

function lNX_sqfChol(X::Union{Array{T, 1}, Array{T, 2}}, μ::Array{T, 1},
    β::Array{Array{T, 1}, 1}, δ::Array{T, 1}, pdthrowerror::Bool=true) where T <: Real

    size(X,1) == length(μ) || throw("In lNX_sqfChol, the columns of X are the observations.")

    Σ = sqfChol_to_Σ(β, δ, pdthrowerror)

    if isnothing(Σ)
        return nothing
    else
        d = MultivariateNormal(μ, Σ)
        if pdthrowerror
            return logpdf(d, X)
        else
            try logpdf(d, X)
            catch excep
                if isa(excep, SingularException)
                    println("Singular exception in lNX_sqfChol.\n")
                    return nothing
                else
                    return logpdf(d, X)
                end
            end
        end
    end
end

# function lNX_seq(x::Array{T, 1}, μ::Array{T, 1},
#     β::Array{Array{T, 1}, 1}, δ::Array{T, 1}) where T <: Real
#
#     K = length(μ)
#     dev = x .- μ
#
#     out = ldnorm(x[K], μ[K], δ[K])
#
#     for k = (K-1):-1:1
#         out += ldnorm(x[k], μ[k] - sum(β[k] .* dev[(k+1):K]) , δ[k])
#     end
#
#     out
# end

### test
#
# using BayesInference
# using Distributions
# using PDMats
# using LinearAlgebra
#
# n = 1000
# K = 5
# X = randn(n, K)
#
# μ_x = randn(K)
# δ_x = exp.(randn(K))
# β_x = [ randn(k) for k = (K-1):-1:1 ]
#
#
# Σ = sqfChol_to_Σ(β_x, δ_x)
# MvNormal(μ_x, Σ)
#
# using BenchmarkTools
#
# @time out1 = lNX(Array(X'), μ_x, β_x, δ_x)
# @btime lNX(Array(X'), μ_x, β_x, δ_x)
# # lNX(X[1,:], μ_x, β_x, δ_x)
#
# @time out2 = [ lNX_seq(X[i,:], μ_x, β_x, δ_x) for i = 1:size(X, 1) ]
# @btime [ lNX_seq(X[i,:], μ_x, β_x, δ_x) for i = 1:size(X, 1) ]
#
# out1 ≈ out2

# naive implemtation (non-sequential) is far superior

function PDMat_adj(A::Matrix{Float64}; maxadd::Float64=1.0e-6,
        epsfact::Float64=100.0, cumadd::Float64=0.0, throwerror::Bool=true)

    try PDMat(A)
    catch excep
        if isa(excep, PosDefException) && cumadd <= maxadd
            a = epsfact * eps(Float64)
            A += a * I
            cumadd += a
            epsfactnext = 10.0 * epsfact
            return PDMat_adj(A, maxadd=maxadd, epsfact=epsfactnext,
                    cumadd=cumadd, throwerror=throwerror)
        else
            println("Failure to adjust to positive definiteness.\n A = ", A, "\n")
            if throwerror
                return PDMat(A) # just trigger original error
            else
                return nothing
            end
        end
    end
end


using DataStructures

struct Graph
    vertices::Vector{Int}
    edges::Set{Tuple{Tuple, Int}}
end

function getNeighbors(arr, (i, j))
    maxrow = size(arr, 1)
    maxcol = size(arr, 2)
    neighboridxs = Tuple[]
    if i == 1 || i == maxrow || j == 1 || j == maxcol
        if (i,j) == (1,1)
            neighboridxs = push!(neighboridxs, (1,2))
            neighboridxs = push!(neighboridxs, (2,1))
        elseif (i,j) == (1, maxcol)
            neighboridxs = push!(neighboridxs, (1,maxcol-1))
            neighboridxs = push!(neighboridxs, (2,maxcol))
        elseif (i,j) == (maxrow, 1)
            neighboridxs = push!(neighboridxs, (maxrow-1,1))
            neighboridxs = push!(neighboridxs, (maxrow,2))
        elseif (i,j) == (maxrow, maxcol)
            neighboridxs = push!(neighboridxs, (maxrow-1,maxcol))
            neighboridxs = push!(neighboridxs, (maxrow,maxcol-1))
        elseif i == 1
            neighboridxs = push!(neighboridxs, (i, j-1))
            neighboridxs = push!(neighboridxs, (i, j+1))
            neighboridxs = push!(neighboridxs, (i+1, j))
        elseif j == 1
            neighboridxs = push!(neighboridxs, (i-1, j))
            neighboridxs = push!(neighboridxs, (i, j+1))
            neighboridxs = push!(neighboridxs, (i+1, j))
        elseif j == maxcol
            neighboridxs = push!(neighboridxs, (i-1, maxcol))
            neighboridxs = push!(neighboridxs, (i+1, maxcol))
            neighboridxs = push!(neighboridxs, (i, maxcol - 1))
        elseif i == maxrow
            neighboridxs = push!(neighboridxs, (maxrow, j-1))
            neighboridxs = push!(neighboridxs, (maxrow, j+1))
            neighboridxs = push!(neighboridxs, (maxrow-1, j))
        end
    else
        neighboridxs = push!(neighboridxs, (i, j-1))
        neighboridxs = push!(neighboridxs, (i, j+1))
        neighboridxs = push!(neighboridxs, (i+1, j))
        neighboridxs = push!(neighboridxs, (i-1, j))
    end
    return neighboridxs
end

function toGraph(filename, part2 = false)
    values = readlines(filename*".txt")
    values = split.(values, "")
    intvalues = Vector{Int}[]
    for line in values
        convline = parse.(Int, line)
        intvalues = push!(intvalues, convline)
    end
    valarr = zeros(Int, size(intvalues)[1],size(intvalues[1])[1])
    for i in 1:length(intvalues)
        for j in 1:length(intvalues[1])
            valarr[i,j] = intvalues[i][j]
        end
    end
    if part2
        origvalarr = copy(valarr)
        arrs = []
        for i in 1:5
            line = []
            for j in 1:5
                arr = copy(origvalarr)
                toadd_x = i - 1
                toadd_y = j - 1
                toadd = toadd_x + toadd_y
                for k in 1:size(origvalarr,1)
                    for l in 1:size(origvalarr,2)
                        if origvalarr[k,l] + toadd != 9
                            arr[k,l] = (origvalarr[k,l] + toadd) % 9
                            if arr[k,l] == 0
                                arr[k,l] = 1
                            end
                        else
                            arr[k,l] = 9
                        end
                    end
                end
                line = push!(line, arr)
            end
            arrs = push!(arrs, line)
        end
        larrs = []
        for line in arrs
            larr = line[1]
            for i in 1:4
                larr = hcat(larr,line[i+1])
            end
            larrs = push!(larrs, larr)
        end
        valarr = larrs[1]
        for i in 1:4
            valarr = vcat(valarr,larrs[i+1])
        end
    end

    vertices = Int[]
    edges = Set{Tuple{Tuple,Int}}()
    for i in 1:size(valarr)[1]
        for j in 1:size(valarr)[2]
            vidx = ((i-1) * size(valarr)[2]) + j
            vertices = push!(vertices,vidx)
            ngb = getNeighbors(valarr, (i,j))
            for n in ngb
                org = vidx
                nidx = ((n[1]-1) * size(valarr)[2]) + n[2]
                weight = valarr[n[1], n[2]]
                edges = push!(edges, ((org,nidx),weight))
            end
        end
    end
    graph = Graph(vertices, edges)
    return graph, valarr
end

function get2Dcoords(x,valarr)
    row = Int(floor(x/size(valarr, 2))) + 1
    col = Int(x % (size(valarr, 2)))
    if col == 0
        col = size(valarr, 2)
        row = row - 1
    end
    return (row,col)
end

function dijkstra(graph, start, stop, valarr)
    dist = zeros(size(graph.vertices))
    previous = zeros(Int, size(graph.vertices))
    for (i,el) in enumerate(dist)
        dist[i] = Inf
        previous[i] = -1
    end
    dist[start] = 0

    G = PriorityQueue()
    for v in graph.vertices
        G = enqueue!(G, v, dist[v])
    end

    while isempty(G) == false
        u = dequeue!(G)
        if u == stop
            break
        end
        ngb = getNeighbors(valarr, get2Dcoords(u, valarr))
        for n in ngb
            d = dist[u] + valarr[n[1], n[2]]
            n = (n[1]-1) * size(valarr, 2) + n[2]
            if d < dist[n]
                dist[n] = d
                G[n] = d
                previous[n] = u
            end
        end
    end
    return dist, previous
end

function solve(filename, part2 = false)
    graph, valarr = toGraph(filename, part2)
    stopcoord = length(graph.vertices)
    dist, prev = dijkstra(graph, 1, stopcoord, valarr)

    println(dist[stopcoord])
end

solve("example")
solve("example", true)
solve("input")
solve("input", true)
#solve("test")
#solve("test2")


using NeuralVerification
using LazySets
using DataStructures
using LinearAlgebra

function priority_optimization(network, lbs, ubs, optimize_reach, evaluate_objective;
        n_steps = 1000, solver=Ai2z(), early_stop=true, stop_freq=200, stop_gap=1e-4, initial_splits = 0)

    start_cell = Hyperrectangle(low = lbs, high = ubs)
    initial_cells = split_multiple_times(start_cell, initial_splits)

    # Create your queue, then add your original new_cells
    cells = PriorityQueue(Base.Order.Reverse) # pop off largest first
    for cell in initial_cells
        enqueue!(cells, cell, optimize_reach(forward_network(solver, network, cell)))
    end

    best_lower_bound = -Inf
    best_x = nothing

    # For n_steps dequeue a cell, split it, and then
    for i = 1:n_steps
        cell, value = peek(cells) # peek instead of dequeue to get value, is there a better way?
        dequeue!(cells)

        # Early stopping
        if early_stop
            if i % stop_freq == 0
                lower_bound = evaluate_objective(network, cell.center)
                if lower_bound > best_lower_bound
                    best_lower_bound = lower_bound
                    best_x = cell.center
                end
                # println("i: ", i)
                # println("lower bound: ", lower_bound)
                # println("best lower bound: ", best_lower_bound)
                # println("value: ", value)
                if (value .- lower_bound) ≤ stop_gap
                    return best_x, best_lower_bound, value
                end
                # println("max radius: ", max(radius(cell)))
            end
        end

        new_cells = split_cell(cell)

        # Enqueue each of the new cells
        for new_cell in new_cells
            # If you've made the max objective cell tiny
            # break (otherwise we end up with zero radius cells)
            if max(radius(new_cell) < sqrt(eps()))
                # Return a concrete value and the upper bound from the parent cell
                # that was just dequeued, as it must have higher value than all other cells
                # that were on the queue, and they constitute a tiling of the space
                lower_bound = evaluate_objective(network, cell.center)
                if lower_bound > best_lower_bound
                    best_lower_bound = lower_bound
                    best_x = cell.center
                end
                return best_x, best_lower_bound, value
            end

            new_value = optimize_reach(forward_network(solver, network, new_cell))
            enqueue!(cells, new_cell, new_value)
        end
    end

    # The largest value in our queue is the approximate optimum
    cell, value = peek(cells)
    lower_bound = evaluate_objective(network, cell.center)

    if lower_bound > best_lower_bound
        best_lower_bound = lower_bound
        best_x = cell.center
    end
    return best_x, best_lower_bound, value
end

""" Helpers """

elem_basis(i, n) = [k == i ? 1.0 : 0.0 for k in 1:n]

function split_multiple_times(cell, n)
    q = Queue{Hyperrectangle}()
    enqueue!(q, cell)
    for i = 1:n
        new_cells = split_cell(dequeue!(q))
        enqueue!(q, new_cells[1])
        enqueue!(q, new_cells[2])
    end
    return q
end

function split_cell(cell::Hyperrectangle)
    lbs, ubs = low(cell), high(cell)
    largest_dimension = argmax(ubs .- lbs)
    # have a vector [0, 0, ..., 1/2 largest gap at largest dimension, 0, 0, ..., 0]
    delta = elem_basis(largest_dimension, length(lbs)) * 0.5 * (ubs[largest_dimension] - lbs[largest_dimension])
    cell_one = Hyperrectangle(low=lbs, high=(ubs .- delta))
    cell_two = Hyperrectangle(low=(lbs .+ delta), high=ubs)
    return [cell_one, cell_two]
end


using Revise
using AdFem
using PyCall
using LinearAlgebra
using ADCME
using MAT
using JLD2
using DelimitedFiles
using PyPlot
np = pyimport("numpy")
mpl = pyimport("tikzplotlib")

mode = "training"
# Domain information
NT = 20
Δt = 1/NT
n = 10
m = 2*n
h = 1.0/n
bdnode = bcnode("lower", m, n, h)

b = 1.0
# pl = placeholder([2.0;0.5;1.0])
pl = Variable(0.8*ones(3))
λ, μ, invη = pl[1], pl[2], pl[3]

iS = tensor(
        [1+2/3*μ*Δt*invη -1/3*μ*Δt*invη 0.0
        -1/3*μ*Δt*invη 1+2/3*μ*Δt*invη 0.0
        0.0 0.0 1+μ*Δt*invη]
    )
S = inv(iS)
H = S * tensor([
    2μ+λ λ 0.0
    λ 2μ+λ 0.0
    0.0 0.0 μ
])

bd = bcedge("upper", m, n, h)
Q, Prhs = compute_fvm_tpfa_matrix(ones(4*m*n), bd, zeros(size(bd,1)),m, n, h)
Q = SparseTensor(Q)
K = compute_fem_stiffness_matrix(H, m, n, h)
L = SparseTensor(compute_interaction_matrix(m, n, h))
M = SparseTensor(compute_fvm_mass_matrix(m, n, h))
A = [K -b*L'
b*L/Δt 1/Δt*M-Q]
A, Abd = fem_impose_coupled_Dirichlet_boundary_condition(A, bdnode, m, n, h)
# error()
U = zeros(m*n+2(m+1)*(n+1), NT+1)
x = Float64[]; y = Float64[]
for j = 1:n+1
    for i = 1:m+1
        push!(x, (i-1)*h)
        push!(y, (j-1)*h)
    end
end

injection = (div(n,2)-1)*m + 3
production = (div(n,2)-1)*m + m-3

function get_disp(ipval)

    function condition(i, tas...)
        i<=NT
    end

    function body(i, tas...)
        ta_u, ta_ε, ta_σ = tas
        u = read(ta_u, i)
        σ0 = read(ta_σ, i)
        ε0 = read(ta_ε, i)
        rhs1 = compute_fem_viscoelasticity_strain_energy_term(ε0, σ0, S, H, m, n, h)
        rhs1 = scatter_update(rhs1, [bdnode; bdnode .+ (m+1)*(n+1)], zeros(2length(bdnode)))
        rhs2 = zeros(m*n)
        rhs2[injection] += 1.0
        rhs2[production] -= 1.0
        rhs2 = rhs2 * ipval + b*L*u[1:2(m+1)*(n+1)]/Δt +
                M * u[2(m+1)*(n+1)+1:end]/Δt + Prhs

        rhs = [rhs1;rhs2]
        o = A\rhs
        ε = eval_strain_on_gauss_pts(o, m, n, h)
        σ = σ0*S + (ε - ε0)*H
        ta_u = write(ta_u, i+1, o)
        ta_ε = write(ta_ε, i+1, ε)
        ta_σ = write(ta_σ, i+1, σ)
        i+1, ta_u, ta_ε, ta_σ
    end

    i = constant(1, dtype=Int32)
    ta_u = TensorArray(NT+1); ta_u = write(ta_u, 1, constant(zeros(2(m+1)*(n+1)+m*n)))
    ta_ε = TensorArray(NT+1); ta_ε = write(ta_ε, 1, constant(zeros(4*m*n, 3)))
    ta_σ = TensorArray(NT+1); ta_σ = write(ta_σ, 1, constant(zeros(4*m*n, 3)))
    _, u_out, ε_out, σ_out = while_loop(condition, body, [i, ta_u, ta_ε, ta_σ])
    u_out = stack(u_out)
    u_out.set_shape((NT+1, size(u_out,2)))
    σ_out = stack(σ_out)
    ε_out = stack(ε_out)

    upper_idx = Int64[]
    for i = 1:m+1
        push!(upper_idx, (div(n,3)-1)*(m+1)+i)
        push!(upper_idx, (div(n,3)-1)*(m+1)+i + (m+1)*(n+1))
    end
    for i = 1:m
        push!(upper_idx, (div(n,3)-1)*m+i+2(m+1)*(n+1))
    end

    u_out, σ_out
end

Us = Array{PyObject}(undef, 6)
Ss = Array{PyObject}(undef, 6)

for i = 1:5
    Us[i], Ss[i] = get_disp(0.2*i)
end
Us[6], Ss[6] = get_disp(0.5)

if isfile("invdata.mat")
    d = matread("invdata.mat")
    Us_ = d["U"]
    Ss_ = d["S"]
    global Sigma0 = Ss_[6]
    global U0 = Us_[6]
    global loss = sum([sum((Us[i][:,1:m+1] - Us_[i][:,1:m+1])^2) for i = 1:5])
end
sess = Session(); init(sess)

function visualize(i)
    visualize_von_mises_stress(Sigma_, m, n, h, name="_nn$i")
    visualize_displacement(Array(U_'), m, n, h, name="_nn0$i",
                    xlim_=[-3h, m*h+2h], ylim_=[-2h, n*h+2h])
    close("all")
    figure(figsize=(13,4))
    subplot(121)
    plot(LinRange(0, 20, NT+1), U0[:,1], "r--")
    plot(LinRange(0, 20, NT+1), U_[:,1], "r")
    plot(LinRange(0, 20, NT+1), U0[:,1+(n+1)*(m+1)], "g--")
    plot(LinRange(0, 20, NT+1), U_[:,1+(n+1)*(m+1)], "g")
    xlabel("Time")
    ylabel("Displacement")
    subplot(122)
    plot(LinRange(0, 20, NT+1), mean(Sigma0[:,1:4,1], dims=2)[:],"r--", label="\$\\sigma_{xx}\$")
    plot(LinRange(0, 20, NT+1), mean(Sigma0[:,1:4,2], dims=2)[:],"b--", label="\$\\sigma_{yy}\$")
    plot(LinRange(0, 20, NT+1), mean(Sigma0[:,1:4,3], dims=2)[:],"g--", label="\$\\sigma_{xy}\$")
    plot(LinRange(0, 20, NT+1), mean(Sigma_[:,1:4,1], dims=2)[:],"r-")
    plot(LinRange(0, 20, NT+1), mean(Sigma_[:,1:4,2], dims=2)[:],"b-")
    plot(LinRange(0, 20, NT+1), mean(Sigma_[:,1:4,3], dims=2)[:],"g-")
    legend()
    legend()
    xlabel("Time")
    ylabel("Stress")
    savefig("disp$i.png")
    savefig("disp$i.pdf")
    matwrite("nn$i.mat", Dict("U"=>U_, "S"=>Sigma_))
end

if mode=="data"
    U = run(sess, Us)
    S = run(sess, Ss)
    matwrite("invdata.mat", Dict("U"=>U, "S"=>S))
    global Sigma_ = S[6]
    global U_ = U[6]
    global Sigma0 = Sigma_
    global U0 = U_
    visualize("true")
    error("Stop!")
end

# @show run(sess, loss)
# lineview(sess, pl, loss, [2.0;0.5;1.0], [5.0;5.0;5.0])
# savefig("line.png")

# gradview(sess, pl, loss, [5.0;5.0;5.0])
# savefig("line.png")

# meshview(sess, pl, loss, [2.0;0.5;1.0], 0.1, 0.1)
# savefig("line.png")

loss_ = BFGS!(sess, loss)
p = run(sess, pl)
matwrite("param.mat", Dict("loss"=>loss_, "p"=>p))


#=
Similar to Batch Normalization, except online and without the rescaling/skew
    y = (a .- μ) ./ σ

    TODO: This is currently broken because OnlineStats.Variances no longer
    exists.
=#
type InputNorm{T,W<:Weight} <: Transformation
    n::Int  # maps n --> n
    vars::Variances{W}
    input::SumNode{T,1}    # a
    output::OutputNode{T,1}  # y
end

function InputNorm{T}(::Type{T}, n::Int,
                    lookback::Int = 100,
                    α::Float64 = NaN,
                    wgt::Weight =  BoundedEqualWeight(isnan(α) ? lookback : α)
                    )
    InputNorm(n,
        Variances(n, wgt),
        InputNode(T, n),
        OutputNode(T, n)
    )
end
InputNorm(n::Int, args...; kw...) = InputNorm(Float64, n, args...; kw...)

function Base.show(io::IO, t::InputNorm)
    print(io, "InputNorm{n=$(t.n)}")
end

function transform!{T}(layer::InputNorm{T})
    a = layer.input.val
    y = layer.output.val

    OnlineStats.fit!(layer.vars, a)
    μ = mean(layer.vars)
    σ = std(layer.vars)
    for i=1:layer.n
        y[i] = (a[i] - μ[i]) / max(σ[i], T(1e-10))
    end
    y
end

function grad!{T}(layer::InputNorm{T})
    ∇a = layer.input.∇
    ∇y = layer.output.∇

    σ = std(layer.vars)
    for i=1:layer.n
        ∇a[i] = ∇y[i] / max(σ[i], T(1e-10))
    end
    ∇a
end


using NiLang, NiLang.AD
using Test, Random
using ForwardDiff
using FixedPointNumbers
using NiLangCore: default_constructor

@testset "ULogarithmic" begin
	@test check_grad(PlusEq(gaussian_log), (1.0, 2.0); iloss=1)
	function muleq(f, x, y, z)
		x = default_constructor(ULogarithmic, x)
		y = default_constructor(ULogarithmic, y)
		z = default_constructor(ULogarithmic, z)
		x *= f(y, z)
		x.log
	end
	g1 = ForwardDiff.gradient(arr->muleq(+, arr...), [7.0, 5.0, 3.0])
	x, y, z = default_constructor(ULogarithmic, 7.0),
		default_constructor(ULogarithmic, 5.0),
		default_constructor(ULogarithmic, 3.0)
	@instr (MulEq(+))(x, y, z)
	@instr GVar(x)
	@instr GVar(y)
	@instr GVar(z)
	@instr x.log.g += 1
	@instr (~MulEq(+))(x, y, z)
	@test grad(x.log) ≈ g1[1]
	@test grad(y.log) ≈ g1[2]
	@test grad(z.log) ≈ g1[3]

	g2 = ForwardDiff.gradient(arr->muleq(-, arr...), [7.0, 5.0, 3.0])
	x, y, z = default_constructor(ULogarithmic, 2.0),
		default_constructor(ULogarithmic, 5.0),
		default_constructor(ULogarithmic, 3.0)
	@instr (MulEq(-))(x, y, z)
	@instr GVar(x)
	@instr GVar(y)
	@instr GVar(z)
	@instr x.log.g += 1
	@instr (~MulEq(-))(x, y, z)
	@test grad(x.log) ≈ g2[1]
	@test grad(y.log) ≈ g2[2]
	@test grad(z.log) ≈ g2[3]
end

@testset "iexp" begin
	@i function i_exp(y!::T, x::T) where T<:Union{Fixed, GVar{<:Fixed}}
	    @invcheckoff begin
	        @routine begin
	            s ← one(ULogarithmic{T})
	            lx ← one(ULogarithmic{T})
	            k ← 0
	        end
	        lx *= convert(x)
	        y! += convert(s)
	        while (s.log > -20, k != 0)
	            k += 1
	            s *= lx / k
	            y! += convert(s)
	        end
	        ~(while (s.log > -20, k != 0)
	            k += 1
	            s *= x / k
	        end)
	        lx /= convert(x)
	        ~@routine
	    end
	end

	x = Fixed43(3.5)
	res = i_exp(Fixed43(0.0), x)[1]
	gx = grad(Grad(i_exp)(Val(1), Fixed43(0.0), x)[3])
	@test res ≈ exp(3.5)
	@test gx ≈ exp(3.5)
end


# Construct a sample using Reservoir sampling

type Sample{T} <: StreamStat
    sample::Vector{T}
    n::Int
end

Sample(T::DataType, k::Integer) = Sample(Array(T, k), 0)

function update!(stat::Sample, x::Any)
    k = length(stat.sample)

    stat.n += 1

    if stat.n <= k
        stat.sample[stat.n] = x
    else
       i = rand(1:stat.n)
       if i <= k
           stat.sample[i] = x
       end
    end

    return
end

# This should really return a Nullable
function state{T}(stat::Sample{T})
   if stat.n < length(stat.sample)
       return T[]
   else
       return stat.sample
   end
end

StatsBase.sample{T}(stat::Sample{T}) = state(stat)

nobs(stat::Sample) = stat.n

Base.copy(stat::Sample) = Sample(stat.sample, stat.n)

# TODO: Implement random mixture
# Base.merge(a::Std, b::Std) = Std(merge(a.var, b.var))

function Base.empty!(stat::Sample)
    empty!(stat.sample)
    stat.n = 0
    return
end

function Base.show(io::IO, stat::Sample)
    @printf(io, "Online Random Sample\n")
    @printf(io, " * K: %d\n", length(stat.sample))
    @printf(io, " * N: %d\n", stat.n)
    return
end


using Dates, Gadfly

set_default_plot_size(25cm, 35cm)

# simple Integrate and Fire Synapse Mode
type Synapse
    τ::Real     # membrane time constant [ms]
    dt::DateTime # time stamp of last update
    v::Real # tracks current voltage
    function Synapse(v::Real;τ::Real=10., dt::DateTime=now())
        new(τ, dt, v)
    end
end

nextDt(dt::DateTime, addMs::Int64) = dt + Millisecond(addMs)

function update!(s::Synapse, iV::Real, dt::DateTime)
    Δdt = int(dt - s.dt)
    s.v += Δdt*(1.0/s.τ)*(-s.v + iV)
    return s.v
end

function getResponse(V::Vector{Float64}, Δτ::Vector{Int64})
    n = length(V)
    data = zeros(n)
    tm = zeros(Int64,n)
    s = Synapse(0,τ=10)
#     tm = ones(n)
    for i = 1:n
        v = update!(s, V[i], nextDt(s.dt, Δτ[i]))
        data[i] = v
        tm[i] = i == 1? 0 : tm[i-1] + Δτ[i]
    end
    return data,tm
end

n = 20
Δτ = rand(1:1,n)
# V = ones(n) - .0
# V = max(randn(n) * 2 .+ 0.,0)
V = randn(n) * 2
response, tm = getResponse(V, Δτ)
vstack(plot(x = tm, y=V, Geom.line), plot(x=tm, y=response, Geom.point))

response[1]
tm[1]


__precompile__()

module QuantumWalk

using Expokit
using LightGraphs
using Optim

using LinearAlgebra
using SparseArrays
using Arpack

export
	AbstractCTQW,
	AbstractSzegedy,
	CTQW,
    CTQWDense,
	QSearchState,
	QWDynamics,
	QWEvolution,
	QWModel,
	QWModelCont,
	QWModelDiscr,
	QWSearch,
	QuantumWalk,
	Szegedy,
	check_qwdynamics,
	evolve,
	execute,
	execute_all,
	execute_all_measured,
	execute_single,
	execute_single_measured,
	expected_runtime,
	graph,
	initial_state,
	marked,
	matrix,
	maximize_quantum_search,
	measure,
	model,
	parameters,
	penalty,
	probability,
	runtime,
	sqrtstochastic,
	state

"""
    QuantumWalk

An package for quantum walk and quantum search simulation.

The packaged provides hierarchy of Quantum walk models and their numerical description together
with general function for quantum walk simulation and quantum search evolution. In
particular `CTQW` and `Szegedy` walk models are implemented as an example of continuous
and discrete walk models.

For both continuous and discrete model quantum walk simulator is implemented and
function finding optimal measurement time for quantum search function. For discrete
walk evolution various optimization may be chosen.

Full documentation is available in [GitHub pages](https://quantumwalks.github.io/QuantumWalk.jl/latest/).

"""
QuantumWalk

include("type_hierarchy.jl")
include("dynamics.jl")

include("qwdynamics/qwdynamics_import.jl")
include("qwmodels/qwmodels_import.jl")

end


using OrdinaryDiffEq
using Trixi

###############################################################################
# semidiscretization of the compressible Euler equations

equations = CompressibleEulerEquations3D(1.4)

"""
    initial_condition_taylor_green_vortex(x, t, equations::CompressibleEulerEquations3D)

The classical inviscid Taylor-Green vortex.
"""
function initial_condition_taylor_green_vortex(x, t, equations::CompressibleEulerEquations3D)
  A  = 1.0 # magnitude of speed
  Ms = 0.1 # maximum Mach number

  rho = 1.0
  v1  =  A * sin(x[1]) * cos(x[2]) * cos(x[3])
  v2  = -A * cos(x[1]) * sin(x[2]) * cos(x[3])
  v3  = 0.0
  p   = (A / Ms)^2 * rho / equations.gamma # scaling to get Ms
  p   = p + 1.0/16.0 * A^2 * rho * (cos(2*x[1])*cos(2*x[3]) + 2*cos(2*x[2]) + 2*cos(2*x[1]) + cos(2*x[2])*cos(2*x[3]))

  return prim2cons(SVector(rho, v1, v2, v3, p), equations)
end
initial_condition = initial_condition_taylor_green_vortex

volume_flux = flux_ranocha
solver = DGSEM(polydeg=3, surface_flux=flux_lax_friedrichs,
               volume_integral=VolumeIntegralFluxDifferencing(volume_flux))

coordinates_min = (-1.0, -1.0, -1.0) .* pi
coordinates_max = ( 1.0,  1.0,  1.0) .* pi
mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=3,
                n_cells_max=100_000)

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver)

###############################################################################
# ODE solvers, callbacks etc.

tspan = (0.0, 10.0)
ode = semidiscretize(semi, tspan)

summary_callback = SummaryCallback()

analysis_interval = 100
analysis_callback = AnalysisCallback(semi, interval=analysis_interval)

alive_callback = AliveCallback(analysis_interval=analysis_interval)

save_solution = SaveSolutionCallback(interval=100,
                                     save_initial_solution=true,
                                     save_final_solution=true,
                                     solution_variables=cons2prim)

stepsize_callback = StepsizeCallback(cfl=1.4)

callbacks = CallbackSet(summary_callback,
                        analysis_callback, alive_callback,
                        save_solution,
                        stepsize_callback)

###############################################################################
# run the simulation

sol = solve(ode, CarpenterKennedy2N54(williamson_condition=false),
            dt=1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
            save_everystep=false, callback=callbacks);
summary_callback() # print the timer summary


include("kern_proc.jl")

d = 10
n = [50, 100, 500, 1000]

srand(1)

cov_procs = Array(Proc, 0)
push!(cov_procs, KernelTest(LinIso(0.0), d, GaussianProcesses.cov))
push!(cov_procs, KernelTest(LinArd(rand(d)), d, GaussianProcesses.cov))
push!(cov_procs, KernelTest(Noise(0.0), d, GaussianProcesses.cov))
push!(cov_procs, KernelTest(SEIso(1.0, 1.0) + Noise(0.0), d, GaussianProcesses.cov))
push!(cov_procs, KernelTest(SEArd(rand(d), 1.0) + Noise(0.0), d, GaussianProcesses.cov))
cov_table = run(cov_procs, n)

grad_stack_procs = Array(Proc, 0)
push!(grad_stack_procs, KernelTest(LinIso(0.0), d, GaussianProcesses.grad_stack))
push!(grad_stack_procs, KernelTest(LinArd(rand(d)), d, GaussianProcesses.grad_stack))
push!(grad_stack_procs, KernelTest(Noise(0.0), d, GaussianProcesses.grad_stack))
push!(grad_stack_procs, KernelTest(SEIso(1.0, 1.0) + Noise(0.0), d, GaussianProcesses.grad_stack))
push!(grad_stack_procs, KernelTest(SEIso(1.0, 1.0) + Noise(0.0), d, GaussianProcesses.grad_stack))
push!(grad_stack_procs, KernelTest(SEArd(rand(d), 1.0) + Noise(0.0), d, GaussianProcesses.grad_stack))
grad_stack_table = run(grad_stack_procs, n)

println("\ncov calculations")
show(cov_table)
println("\grad_stack calculations")
show(grad_stack_table)


include("kaiming_init.jl")
using Flux
using Flux:@treelike
using NNlib: @fix, σ_stable

################################### NAC ########################################

struct NAC{F,S,T}
  Ŵ::S
  M̂::F
  b::T
end

NAC(in::Integer, out::Integer;
               initŴ = kaiming_uniform, initM̂ = kaiming_uniform, initb = zeros) =
    NAC(param(initŴ(out, in; a=√5.0)), param(initM̂(out, in; a=√5.0)), param(initb(out)))

@treelike NAC

function (nac::NAC)(x)
  Ŵ, M̂, b = nac.Ŵ, nac.M̂, nac.b
  W = tanh.(Ŵ) .* σ_stable.(M̂)
  @fix W*x .+ b
end

Base.show(io::IO, l::NAC) = print(io, "NAC(", size(l.Ŵ, 2), ", ", size(l.Ŵ, 1), ")")

################################## NALU ########################################

struct NALU{N,S,T,E}
  nac::N
  G::S
  b::T
  ϵ::E
end

NALU(in::Integer, out::Integer;
               initĜ = kaiming_uniform, initb = zeros) =
  NALU(NAC(in, out), param(initĜ(out, in; a=√5.0)), param(initb(out)), 1e-10)

@treelike NALU

function (nalu::NALU)(x)
  nac, G, b, ϵ = nalu.nac, nalu.G, nalu.b, nalu.ϵ
  a = nac(x)
  log_inp = log.(abs.(x) .+ ϵ)
  m = exp.(nac(log_inp))
  g = σ_stable.(G*x .+ b)
  @fix g .* a .+ (1.0 .- g) .* m
end

Base.show(io::IO, l::NALU) = print(io, "NALU(", size(l.G, 2), ", ", size(l.G, 1), ")")


"""
      parsimony(tree::N, char::Dict{String,String}, gap::String="-")::Float64 where N<: GeneralNode

Do parsimony reconstruction for a tree and a set of characters. The characters are
supplied as the char dictionary with the key being the name of a leave and the value
string representation of the character.
"""
function parsimony(tree::N, char::Dict{String,String}, gap::String="-")::Float64 where N<: GeneralNode
      po = post_order(tree)
      states = filter(x -> x != "-", unique(collect(values(char))))
      nStates = length(states)
      if nStates == 0
            return 0.0
      end
      nNodes = length(po)
      costMatrix = zeros(nNodes, nStates) .- 1
      for node in po
            if node.nchild == 0
                  # i am a leave node
                  s = char[node.name]
                  costMatrix[node.num, :] .= s == gap ? 0 : abs.(log.(states .== s))
            else
                  # i am an internal node
                  daughters = node.children
                  localCosts = zeros(nStates, nStates, length(daughters))
                  for (j, s1) in enumerate(states),
                        (k, s2) in enumerate(states),
                        (l, d) in enumerate(daughters)

                        localCosts[j, k, l] = costMatrix[d.num, k] + Int(s1 != s2)
                  end
                  costMatrix[i, :] = vec(
                        mapslices(
                              sum,
                              mapslices(minimum, localCosts, dims = 2),
                              dims = 3,
                        ),
                  )
            end
      end
      minimum(costMatrix[tree.num, :])
end


#-------------------------------------------------------------------------#
#                                                                         #
#                               Simulating_nets                           #
#                                                                         #
#-------------------------------------------------------------------------#

#-------------------------------------------------------------------------#
# OBSERVATIONS:

# 1. This code is for simulating phylogenetic networks in a Newick format
# 2. This code is a work in progress
# 3. The functions are in the file Fn_PhyloSim
# 4. The principal example is on "Example of a random tree generation" section
#
#-------------------------------------------------------------------------#

## Packages
# using Plots
using PhyloNetworks
using PhyloPlots

## Example of a random tree generation

# Including functions
include("Fn_PhyloSim.jl")
# Fixing seed
# Random.seed!(42)

# Different taxa examples
taxa = ["1ba","b","c","d"]
taxa = ["A","B"]
taxa = collect('a':'d')
taxa = collect('a':'z')
taxa = string.(1:40)

# Total branch lengths
total_bl = 10

# Return a network of a tree
is_network = true

# Generating a random tree topology and branch length in Newikc format
x = Sim_tree(taxa,total_bl,is_network)
print(join(x))

# Transformation to a "HybridNetwork" variable of the package PhyloNetworks
net = readTopology(join(x))
tipLabels(net)

# Plot of the network
plot(net,:R,showEdgeLength=true,showGamma=true)
# plot(net,:R,showEdgeLength=true)
# plot(net,:R)
net
## Some info about the network
printEdges(net)


using ACEds

using StaticArrays
using ACEds.equivRotations3D

"""
Given an l-vector `ll` iterate over all combinations of `mm` vectors  of
the same length such that `sum(mm) == 1`
"""
struct MRange2{N, T2}
   ll::SVector{N, Int}
   cartrg::T2
end

Base.length(mr::MRange2) = sum(_->1, _mrange2(mr.ll))

_mrange2(ll) = MRange2(ll, Iterators.Stateful(
               filter((x) -> abs(sum(x))<= 1, Tuple.(CartesianIndices(ntuple(i -> -ll[i]:ll[i], length(ll)))))
                     ))

function Base.iterate(mr::MRange2, args...)
   while true
      if isempty(mr.cartrg)
         return nothing
      end
      mpre = popfirst!(mr.cartrg)
      return SVector(mpre), nothing
   end
end

"""
Compare behaviour of iterate(mr::MRange) with iterate(mr::MRange2)
"""

function mytest(mr)
   for (i,val) in enumerate(mr)
      print("index:", i, " m-vector: ", val, "\n")
   end
end

ll = @SVector([10,10,10])
mr1 = ACEds.equivRotations3D._mrange(ll)
mr2 = _mrange2(ll)
mytest(mr2)
print("-----------")