using Pkg
Pkg.activate("/home/Morgan/FrankWolfe.jl")

using FrankWolfe
using Test
using LinearAlgebra
using DoubleFloats

# include("/home/Morgan/FrankWolfe.jl/src/FrankWolfe.jl")

include("lmo.jl")
include("function_gradient.jl")
include("active_set.jl")
include("utils.jl")
include("pairwise.jl")

@testset "Testing vanilla Frank-Wolfe with various step size and momentum strategies" begin
    f(x) = norm(x)^2
    function grad!(storage, x)
        return storage .= 2x
    end
    lmo_prob = FrankWolfe.ProbabilitySimplexOracle(1)
    x0 = FrankWolfe.compute_extreme_point(lmo_prob, zeros(5))
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Agnostic(),
            verbose=true,
        )[3] - 0.2,
    ) < 1.0e-5
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Agnostic(),
            verbose=false,
            gradient=collect(similar(x0)),
        )[3] - 0.2,
    ) < 1.0e-5
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Goldenratio(),
            verbose=true,
        )[3] - 0.2,
    ) < 1.0e-5
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Backtracking(),
            verbose=false,
        )[3] - 0.2,
    ) < 1.0e-5
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Nonconvex(),
            verbose=false,
        )[3] - 0.2,
    ) < 1.0e-2
    @test FrankWolfe.frank_wolfe(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=1000,
        line_search=FrankWolfe.Shortstep(),
        L=2,
        verbose=false,
    )[3] ≈ 0.2
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Nonconvex(),
            verbose=false,
        )[3] - 0.2,
    ) < 1.0e-2
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Agnostic(),
            verbose=false,
            momentum=0.9,
        )[3] - 0.2,
    ) < 1.0e-3
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Agnostic(),
            verbose=false,
            momentum=0.5,
        )[3] - 0.2,
    ) < 1.0e-3
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Agnostic(),
            verbose=false,
            momentum=0.9,
            emphasis=FrankWolfe.memory,
        )[3] - 0.2,
    ) < 1.0e-3
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Adaptive(),
            L=100,
            verbose=false,
            momentum=0.9,
        )[3] - 0.2,
    ) < 1.0e-3
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Adaptive(),
            L=100,
            verbose=false,
            momentum=0.5,
        )[3] - 0.2,
    ) < 1.0e-3
    @test abs(
        FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Adaptive(),
            L=100,
            verbose=false,
            momentum=0.9,
            emphasis=FrankWolfe.memory,
        )[3] - 0.2,
    ) < 1.0e-3
end

@testset "Gradient with momentum correctly updated" begin
    # fixing https://github.com/ZIB-IOL/FrankWolfe.jl/issues/47
    include("momentum_memory.jl")
end

@testset "Testing Lazified Conditional Gradients with various step size strategies" begin
    f(x) = norm(x)^2
    function grad!(storage, x)
        @. storage = 2x
        return nothing
    end
    lmo_prob = FrankWolfe.ProbabilitySimplexOracle(1)
    x0 = FrankWolfe.compute_extreme_point(lmo_prob, zeros(5))
    @test abs(
        FrankWolfe.lazified_conditional_gradient(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Goldenratio(),
            verbose=true,
        )[3] - 0.2,
    ) < 1.0e-5
    @test abs(
        FrankWolfe.lazified_conditional_gradient(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Backtracking(),
            verbose=true,
        )[3] - 0.2,
    ) < 1.0e-5
    @test abs(
        FrankWolfe.lazified_conditional_gradient(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=1000,
            line_search=FrankWolfe.Shortstep(),
            L=2,
            verbose=true,
        )[3] - 0.2,
    ) < 1.0e-5
end

@testset "Testing Lazified Conditional Gradients with cache strategies" begin
    n = Int(1e5)
    L = 2
    k = 1000
    bound = 16 * L * 2 / (k + 2)

    f(x) = norm(x)^2
    function grad!(storage, x)
        @. storage = 2 * x
        return nothing
    end
    lmo_prob = FrankWolfe.ProbabilitySimplexOracle(1)
    x0 = FrankWolfe.compute_extreme_point(lmo_prob, zeros(n))

    x, v, primal, dual_gap, trajectory = FrankWolfe.lazified_conditional_gradient(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Shortstep(),
        L=2,
        verbose=true,
    )

    @test primal - 1 / n <= bound

    x, v, primal, dual_gap, trajectory = FrankWolfe.lazified_conditional_gradient(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Shortstep(),
        L=2,
        cache_size=100,
        verbose=false,
    )

    @test primal - 1 / n <= bound

    x, v, primal, dual_gap, trajectory = FrankWolfe.lazified_conditional_gradient(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Shortstep(),
        L=2,
        cache_size=100,
        greedy_lazy=true,
        verbose=false,
    )

    @test primal - 1 / n <= bound
end

@testset "Testing emphasis blas vs memory" begin
    n = Int(1e5)
    k = 100
    xpi = rand(n)
    total = sum(xpi)
    xp = xpi ./ total
    f(x) = norm(x - xp)^2
    function grad!(storage, x)
        @. storage = 2 * (x - xp)
        return nothing
    end
    @testset "Using sparse structure" begin
        lmo_prob = FrankWolfe.ProbabilitySimplexOracle(1.0)
        x0 = FrankWolfe.compute_extreme_point(lmo_prob, zeros(n))

        x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=k,
            line_search=FrankWolfe.Backtracking(),
            print_iter=k / 10,
            verbose=false,
            emphasis=FrankWolfe.blas,
        )

        @test x !== nothing

        x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            x0,
            max_iteration=k,
            line_search=FrankWolfe.Backtracking(),
            print_iter=k / 10,
            verbose=false,
            emphasis=FrankWolfe.memory,
        )

        @test x !== nothing
    end
    @testset "Using dense structure" begin
        lmo_prob = FrankWolfe.L1ballDense{Float64}(1)
        x0 = FrankWolfe.compute_extreme_point(lmo_prob, zeros(n))

        x, _ = FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            copy(x0),
            max_iteration=k,
            line_search=FrankWolfe.Backtracking(),
            print_iter=k / 10,
            verbose=false,
            emphasis=FrankWolfe.blas,
        )

        @test x !== nothing

        x, _ = FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            copy(x0),
            max_iteration=k,
            line_search=FrankWolfe.Backtracking(),
            print_iter=k / 10,
            verbose=false,
            emphasis=FrankWolfe.memory,
        )

        @test x !== nothing

        line_search = FrankWolfe.MonotonousStepSize()
        x, _, primal_conv, _ = FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            copy(x0),
            max_iteration=k,
            line_search=line_search,
            print_iter=k / 10,
            verbose=false,
            emphasis=FrankWolfe.memory,
        )
        @test line_search.factor < 20

        line_search = FrankWolfe.MonotonousNonConvexStepSize()
        x, _, primal_nonconv, _ = FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo_prob,
            copy(x0),
            max_iteration=k,
            line_search=line_search,
            print_iter=k / 10,
            verbose=false,
            emphasis=FrankWolfe.memory,
        )
        @test line_search.factor < 20
    end
end
@testset "Testing rational variant" begin
    rhs = 1
    n = 40
    k = 1000

    xpi = rand(big(1):big(100), n)
    total = sum(xpi)
    xp = xpi .// total

    f(x) = norm(x - xp)^2
    function grad!(storage, x)
        @. storage = 2 * (x - xp)
    end

    lmo = FrankWolfe.ProbabilitySimplexOracle{Rational{BigInt}}(rhs)
    direction = rand(n)
    x0 = FrankWolfe.compute_extreme_point(lmo, direction)
    @test eltype(x0) == Rational{BigInt}

    x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
        f,
        grad!,
        lmo,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Agnostic(),
        print_iter=k / 10,
        emphasis=FrankWolfe.blas,
        verbose=false,
    )

    @test eltype(x0) == Rational{BigInt}

    x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
        f,
        grad!,
        lmo,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Agnostic(),
        print_iter=k / 10,
        emphasis=FrankWolfe.memory,
        verbose=true,
    )
    @test eltype(x0) == eltype(x) == Rational{BigInt}
    @test f(x) <= 1e-4

    # very slow computation, explodes quickly
    x0 = collect(FrankWolfe.compute_extreme_point(lmo, direction))
    x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
        f,
        grad!,
        lmo,
        x0,
        max_iteration=15,
        line_search=FrankWolfe.RationalShortstep(),
        L=2,
        print_iter=k / 100,
        emphasis=FrankWolfe.memory,
        verbose=true,
    )

    x0 = FrankWolfe.compute_extreme_point(lmo, direction)
    x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
        f,
        grad!,
        lmo,
        x0,
        max_iteration=15,
        line_search=FrankWolfe.RationalShortstep(),
        L=2,
        print_iter=k / 10,
        emphasis=FrankWolfe.memory,
        verbose=true,
    )
    @test eltype(x) == Rational{BigInt}
end
@testset "Multi-precision tests" begin
    rhs = 1
    n = 100
    k = 1000

    xp = zeros(n)

    L = 2
    bound = 2 * L * 2 / (k + 2)

    f(x) = norm(x - xp)^2
    function grad!(storage, x)
        @. storage = 2 * (x - xp)
    end
    test_types = (Float16, Float32, Float64, Double64, BigFloat, Rational{BigInt})

    @testset "Multi-precision test for $T" for T in test_types
        println("\nTesting precision for type: ", T)
        lmo = FrankWolfe.ProbabilitySimplexOracle{T}(rhs)
        direction = rand(n)
        x0 = FrankWolfe.compute_extreme_point(lmo, direction)

        x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo,
            x0,
            max_iteration=k,
            line_search=FrankWolfe.Agnostic(),
            print_iter=k / 10,
            emphasis=FrankWolfe.blas,
            verbose=true,
        )

        @test eltype(x0) == T
        @test primal - 1 / n <= bound

        x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
            f,
            grad!,
            lmo,
            x0,
            max_iteration=k,
            line_search=FrankWolfe.Agnostic(),
            print_iter=k / 10,
            emphasis=FrankWolfe.memory,
            verbose=true,
        )

        @test eltype(x0) == T
        @test primal - 1 // n <= bound

        x, v, primal, dual_gap, trajectory = FrankWolfe.away_frank_wolfe(
            f,
            grad!,
            lmo,
            x0,
            max_iteration=k,
            line_search=FrankWolfe.Adaptive(),
            print_iter=k / 10,
            emphasis=FrankWolfe.memory,
            verbose=true,
        )

        @test eltype(x0) == T
        @test primal - 1 // n <= bound

        x, v, primal, dual_gap, trajectory = FrankWolfe.blended_conditional_gradient(
            f,
            grad!,
            lmo,
            x0,
            max_iteration=k,
            line_search=FrankWolfe.Adaptive(),
            print_iter=k / 10,
            emphasis=FrankWolfe.memory,
            verbose=true,
        )

        @test eltype(x0) == T
        @test primal - 1 // n <= bound

    end
end

@testset "Stochastic FW linear regression" begin
    function simple_reg_loss(θ, data_point)
        (xi, yi) = data_point
        (a, b) = (θ[1:end-1], θ[end])
        pred = a ⋅ xi + b
        return (pred - yi)^2 / 2
    end

    function ∇simple_reg_loss(storage, θ, data_point)
        (xi, yi) = data_point
        (a, b) = (θ[1:end-1], θ[end])
        pred = a ⋅ xi + b
        storage[1:end-1] .+= xi * (pred - yi)
        storage[end] += pred - yi
        return storage
    end

    xs = [10 * randn(5) for i in 1:20000]
    params = rand(6) .- 1 # start params in (-1,0)
    bias = 2π
    params_perfect = [1:5; bias]

    params = rand(6) .- 1 # start params in (-1,0)

    data_perfect = [(x, x ⋅ (1:5) + bias) for x in xs]
    f_stoch = FrankWolfe.StochasticObjective(simple_reg_loss, ∇simple_reg_loss, data_perfect, similar(params))
    lmo = FrankWolfe.LpNormLMO{2}(1.1 * norm(params_perfect))

    θ, _, _, _, _ = FrankWolfe.stochastic_frank_wolfe(
        f_stoch,
        lmo,
        copy(params),
        momentum=0.95,
        verbose=false,
        line_search=FrankWolfe.Nonconvex(),
        max_iteration=100_000,
        batch_size=length(f_stoch.xs) ÷ 100,
        trajectory=false,
    )
    @test norm(θ - params_perfect) ≤ 0.05 * length(θ)

    # SFW with incrementing batch size
    batch_iterator = FrankWolfe.IncrementBatchIterator(
        length(f_stoch.xs) ÷ 1000,
        length(f_stoch.xs) ÷ 10,
        2,
    )
    θ, _, _, _, _ = FrankWolfe.stochastic_frank_wolfe(
        f_stoch,
        lmo,
        copy(params),
        momentum=0.95,
        verbose=false,
        line_search=FrankWolfe.Nonconvex(),
        max_iteration=5000,
        batch_iterator=batch_iterator,
        trajectory=false,
    )
    @test batch_iterator.maxreached
    # SFW damped momentum
    momentum_iterator = FrankWolfe.ExpMomentumIterator()
    θ, _, _, _, _ = FrankWolfe.stochastic_frank_wolfe(
        f_stoch,
        lmo,
        copy(params),
        verbose=false,
        line_search=FrankWolfe.Nonconvex(),
        max_iteration=5000,
        batch_size=1,
        trajectory=false,
        momentum_iterator=momentum_iterator,
    )
    θ, _, _, _, _ = FrankWolfe.stochastic_frank_wolfe(
        f_stoch,
        lmo,
        copy(params),
        line_search=FrankWolfe.Nonconvex(),
        max_iteration=5000,
        batch_size=1,
        verbose=false,
        trajectory=false,
        momentum_iterator=nothing,
    )
end

@testset "Away-step FW" begin
    n = 50
    lmo_prob = FrankWolfe.ProbabilitySimplexOracle(1.0)
    x0 = FrankWolfe.compute_extreme_point(lmo_prob, rand(n))
    f(x) = norm(x)^2
    function grad!(storage, x)
        @. storage = 2x
    end
    k = 1000
    active_set = ActiveSet([(1.0, x0)])

    # compute reference from vanilla FW
    xref, _ = FrankWolfe.frank_wolfe(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Backtracking(),
        verbose=false,
        emphasis=FrankWolfe.blas,
    )

    x, v, primal, dual_gap, trajectory = FrankWolfe.away_frank_wolfe(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Backtracking(),
        print_iter=k / 10,
        verbose=true,
        emphasis=FrankWolfe.blas,
    )

    @test x !== nothing
    @test xref ≈ x atol = (1e-3 / length(x))

    x, v, primal, dual_gap, trajectory = FrankWolfe.away_frank_wolfe(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        away_steps = false,
        line_search=FrankWolfe.Backtracking(),
        print_iter=k / 10,
        verbose=true,
        emphasis=FrankWolfe.blas,
    )

    @test x !== nothing
    @test xref ≈ x atol = (1e-3 / length(x))

    xs, v, primal, dual_gap, trajectory = FrankWolfe.away_frank_wolfe(
        f,
        grad!,
        lmo_prob,
        active_set,
        max_iteration=k,
        line_search=FrankWolfe.Backtracking(),
        print_iter=k / 10,
        verbose=true,
        emphasis=FrankWolfe.blas,
    )

    @test xs !== nothing
    @test xref ≈ xs atol = (1e-3 / length(x))

    x, v, primal, dual_gap, trajectory = FrankWolfe.away_frank_wolfe(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Backtracking(),
        print_iter=k / 10,
        verbose=true,
        emphasis=FrankWolfe.memory,
    )

    @test x !== nothing
    @test xref ≈ x atol = (1e-3 / length(x))

    x, v, primal, dual_gap, trajectory = FrankWolfe.away_frank_wolfe(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        away_steps=false,
        line_search=FrankWolfe.Backtracking(),
        print_iter=k / 10,
        verbose=true,
        emphasis=FrankWolfe.memory,
    )

    @test x !== nothing
    @test xref ≈ x atol = (1e-3 / length(x))

    xs, v, primal, dual_gap, trajectory = FrankWolfe.away_frank_wolfe(
        f,
        grad!,
        lmo_prob,
        active_set,
        max_iteration=k,
        line_search=FrankWolfe.Backtracking(),
        print_iter=k / 10,
        verbose=true,
        emphasis=FrankWolfe.memory,
    )

    @test xs !== nothing
    @test xref ≈ xs atol = (1e-3 / length(x))

    empty!(active_set)
    @test_throws ArgumentError("Empty active set") FrankWolfe.away_frank_wolfe(
        f,
        grad!,
        lmo_prob,
        active_set,
        max_iteration=k,
        line_search=FrankWolfe.Backtracking(),
        print_iter=k / 10,
        verbose=true,
        emphasis=FrankWolfe.blas,
    )

end

@testset "Blended conditional gradient" begin
    n = 50
    lmo_prob = FrankWolfe.ProbabilitySimplexOracle(1.0)
    x0 = FrankWolfe.compute_extreme_point(lmo_prob, randn(n))
    f(x) = norm(x)^2
    function grad!(storage, x)
        @. storage = 2x
    end
    k = 1000

    # compute reference from vanilla FW
    xref, _ = FrankWolfe.frank_wolfe(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=k,
        line_search=FrankWolfe.Backtracking(),
        verbose=true,
        emphasis=FrankWolfe.blas,
    )

    x, v, primal, dual_gap, trajectory = FrankWolfe.blended_conditional_gradient(
        f,
        grad!,
        lmo_prob,
        x0;
        line_search=FrankWolfe.Backtracking(),
        L=Inf,
        epsilon=1e-9,
        max_iteration=k,
        print_iter=1,
        trajectory=false,
        verbose=false,
        linesearch_tol=1e-10,
    )

    @test x !== nothing
    @test f(x) ≈ f(xref)

end

include("oddities.jl")

# in separate module for name space issues
module BCGDirectionError
using Test
@testset "BCG direction accuracy" begin
    include("bcg_direction_error.jl")
end
end

module RationalTest
using Test
@testset "Rational test and shortstep" begin
    include("rational_test.jl")
end
end


function gradcheckCPU(net;showgrad=false,returngrad=false)
    #ONLY WORKS ON THE CPU AT THE MOMENT
    # (c) David Barber, University College London 2015
    println("GRADCHECK CPU")
    value=net.value
    node=net.node
    net2=deepcopy(net) ## CAREFUL -- DEEPCOPY IS BUGGY FOR CUDAARRAY

    ADforward!(net)
    ADbackward!(net)
    g=net.gradient

    epsilon=1e-7;
    tol=1e-6
    gemp=deepcopy(g)
    #for n=1:length(net.node) # loop over the nodes in the network
    for n in net.validnodes # loop over the nodes in the network
        if node[n].returnderivative
            println("------------------------------------------------------------------------------")
            println("node $n:")
            for linearidx in 1:prod(size(value[n]))
                net2.value[n][linearidx]=value[n][linearidx]+epsilon
                fplus=ADforward!(net2,returnf=true)
                net2.value[n][linearidx]=value[n][linearidx]-epsilon
                fminus=ADforward!(net2,returnf=true)
                net2.value[n][linearidx]=value[n][linearidx]
                gemp[n][linearidx]=0.5*(fplus-fminus)/epsilon
            end
            diff=mean(abs(g[n]-gemp[n]))
            reldiff=mean(abs((g[n]-gemp[n])./(realmin()+g[n])))
            println("absolute difference between analytic and empirical gradient=$diff")
            println("relative difference between analytic and empirical gradient=$reldiff")

                #if showgrad
                #println("analytic gradient:"); println(g[n])
                #println("empirical gradient:"); println(gemp[n])
                #end

            if diff>tol
                if showgrad
                println("analytic gradient:"); println(g[n])
                println("empirical gradient:"); println(gemp[n])
                end
                print_with_color(:red,"failed: analytic and empiricial gradient mismatch more than $tol\n")
            else
                print_with_color(:green,"passed\n")
            end
        end
    end
    if returngrad
        return gemp
    end
end


### Notation
#   small letters => vectors
# capital letters => matrices

function build_A(k, u, η, D, y, params)

    N = params.Ny
    g = params.g

    k2 = k^2
    ik = 1im*k

    # This removed the free-surface deformation in the basic state,
    # which is more like what is done in the Hydrostatic case

    h = 0*η .+ params.H

     U = diagm(0 => u)
     H = diagm(0 => h)
    dU = diagm(0 => D*u)
    dH = diagm(0 => D*h)
     F = diagm(0 => params.f₀ .+ 0*y)
     I = diagm(0 => ones(N+1))

    # Form 1L-RSW Matrix
    #   [u1, v1, h1]

    A = [    k*U         1im*(F-dU)[:, 2:N]    g*k*I;
          -1im*F[2:N,:]     k*U[2:N,2:N]    -1im*g*D[2:N,:];
             k*H         -1im*D*H[:,  2:N]       k*U];

    return A
end


# __author__:aliyousuuf@gmail.com

# This file contains all the codes from lecture, pactise quizzes, and assesment
# of week 02.

# Assume that we've downloaded the EBola data from wiki's website,
# and save it as csv format under file name epidata.csv.

epidata = readdlm("epidata.csv", ',')

col1 = epidata[:,1]
for date = 1:54
    col1[date] = DateTime(col1[date], "d u y")
end

#println(col1)
#println(Dates.datetime2rata(col1[1]))

dayssincemar22(x) = Dates.datetime2rata(x)
epidays = Array(Int64, 54)
for i = 1:54
    epidays[i] = dayssincemar22(col1[i])
end

epidata[:,1] = epidays
writedlm("epidata_converted.csv", epidata, ',')
epidata_converted = readdlm("epidata_converted.csv",',')

# Now we're going to plot epidemic days vs all the cases.

epidays = epidata_converted[:,1]
allcases = epidata_converted[:,2]

# akkcases contains two types of data--string and Int--we've to convert them
# into one type.

for item = 1:54
    if typeof(allcases[item,1]) == SubString{String}
        allcases[item,1] = parse.(Int64,replace(allcases[item,1], ",", ""))
    end
end

using Plots
#Plots.PyPlotBackend()
#pyplot()
#using Gadfly
#plot(epidays, allcases)

# Week_02_Practice Quiz: Q.2

f(x) = 3 * x^2 + 6 * x - 9
for x = -5:5
  #println((,x, ", ", f(x), ")")
end
using Plots
#plot(f, -4, 3) # plot f over [-4,4]
#plot!(zero, -4, 3)

#Week_02_Practice Quiz: Q.3

data = [1.6800483  -1.641695388;
        0.501309281 -0.977697538;
        1.528012113 0.52771122;
        1.70012253 1.711524991;
        1.992493625 1.891000015;
        2.706075824 -0.463427794;
        2.994931927 -0.443566619;
        3.491852811 -1.275179133;
        3.501191722 -0.690499597;
        4.459924502 -5.516130799;
        4.936965851 -6.001703074;
        5.023289852 -8.36416901;
        5.04233698 -7.924477517;
        5.50739285 -10.77482371;
        5.568665171 -10.9171878]

x,y = data[:,1] , data[:,2]
Plots.PyPlotBackend()
#plot(x, y, linetype = :scatter, leg = false)

#Week_02_Practice Quiz: Q.7
n = 20
x = sort(rand(20)); y = rand(20)
#plot!(x,y, leg=false, title = "A sample plot")

# gui()

# replace all the instances of "-" and "--" into 0, and "," into "".
row,col = size(epidata_converted)
for r = 1:row
    for c = 1: col
        if typeof(epidata_converted[r,c]) == SubString{String}
            epidata_converted[r,c] = replace(epidata_converted[r,c], ",", "")
            if isnumber(string(epidata_converted[r,c])) == false
                epidata_converted[r,c] = 0
            end
        end
    end
end

# Convert all the string into integer
for r = 1:row
    for c = 1: col
       if isnumber(string(epidata_converted[r,c])) == true
           epidata_converted[r,c] = parse(Int64,string(epidata_converted[r,c]))
       end
   end
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.4
#     language: julia
#     name: julia-1.6
# ---

# %%
using Distributions
using Roots
using Plots
using StatsFuns: logistic, logit

# \approx TAB → ≈
# \lessapprox TAB → ⪅
x ⪅ y = x < y || x ≈ y

# %%
# Bernoulli分布モデルでの通常の信頼区間

"""P値函数(exact版)"""
function pvalue_exact(n, p, k)
    bin = Binomial(n, p)
    p0 = pdf(bin, k)
    s = sum(pdf(bin, j) for j in support(bin) if pdf(bin, j) ⪅ p0)
    min(1, s)
end

"""P値函数(正規分布近似)"""
function pvalue_normal(n, p, k)
    bin = Binomial(n, p)
    normal = Normal(mean(bin), std(bin))
    min(1, 2cdf(normal, k), 2ccdf(normal, k))
end

"""信頼区間函数"""
function confidence_interval(pvalue_func, n, k; α = 0.05)
    f(t) = pvalue_func(n, logistic(t), k) - α
    CI = logistic.(find_zeros(f, -10.0, 10.0))
    if length(CI) < 2
        return 2k ≤ n ? [0.0, first(CI)] : [first(CI), 1.0]
    else
        return [first(CI), last(CI)]
    end
end

# %%
n = 20
x = 0:n
θ = 0:0.002:1
z_exact = pvalue_exact.(n, θ, x')
z_normal = pvalue_normal.(n, θ, x')

P = plot(; size=(500, 500), colorbar=false)
heatmap!(x, θ, z_exact; clim=(-0.1, 1.0))
plot!(; xlabel="x", ylabel="θ")
plot!(; xtick=0:n, ytick=0:0.05:1)
title!("pvalue_exact(n=$n, θ, x)"; titlefontsize=12)

Q = plot(; size=(500, 500), colorbar=false)
heatmap!(x, θ, z_normal; clim=(-0.1, 1.0))
plot!(; xlabel="x", ylabel="θ")
plot!(; xtick=0:n, ytick=0:0.05:1)
title!("pvalue_normal(n=$n, θ, x)"; titlefontsize=12)

plot(P, Q; size=(1000, 500))
plot!(; leftmargin=3Plots.mm, bottommargin=3Plots.mm)

# %%


# Fizz BUzz with Julia

for num in 1:100
    if num % 15 == 0
        println("FizzBuzz")
    elseif num % 3 == 0
        println("Fizz")
    elseif num % 5 == 0
        println("Buzz")
    else
        println(num)
    end
end


using FirstOrderSolvers: KKTMatrix, AffinePlusLinear
using Convex

Random.seed!(10)

#Test KKTMatrix
A = randn(10,20)
M1 = [I A';A -I]
M2 = KKTMatrix(A)
x = randn(30)
y1 = randn(30)
y2 = randn(30)
mul!(y1, M1, x)
mul!(y2, M2, x)
@test y1 ≈ y2

mul!(y1, transpose(M1), x)
mul!(y2, transpose(M2), x)
@test y1 ≈ y2

#Test AffinePlusLinear

x0 = randn(20)
z0 = randn(10)
q = randn(20)
b = randn(10)

# β = 1
β = 1
S2 = AffinePlusLinear(A, b, q, β)
y2 = Array{Float64,1}(undef, 30)
FirstOrderSolvers.prox!(y2, S2, [x0;z0])
x2 = view(y2,1:20)
z2 = view(y2,21:30)

# #This test is equivalent to y3
# x = Variable(20)
# z = Variable(10)
# p = minimize(1/2*norm(x-x0)^2+1/2*norm(z-z0)^2+dot(q,x), [A*x-z==b])
# solve!(p, FirstOrderSolvers.DR(eps=1e-9))
# x1 = x.value
# z1 = z.value
# y1 = [x1; z1]
# @test y1 ≈ y2

y3 = [I A'; A -I]\[x0-q+A'z0; b]
@test y3 ≈ y2

β = -1
S2 = AffinePlusLinear(A, b, q, β)
y2 = Array{Float64,1}(undef, 30)
FirstOrderSolvers.prox!(y2, S2, [x0;z0])
x2 = view(y2,1:20)
z2 = view(y2,21:30)

# #This test is equivalent to y3
# x = Variable(20)
# z = Variable(10)
# p = minimize(1/2*norm(x-x0)^2+1/2*norm(z-z0)^2+dot(q,x), [A*x-β*z==b])
# solve!(p, FirstOrderSolvers.DR(eps=1e-9))
# x1 = x.value
# z1 = z.value
# y1 = [x1; z1]
# @test y1 ≈ y2

y3 = [I -A'; A I]\[x0-q-A'z0; b]
@test y3 ≈ y2


# TODO: I cound that "utils.jl" has been `include`ed in the main `VectorSphericalWaves` module. Is there a more neat way? when I `include("utils.jl")` here, I get lots of warnings.

export M_mn_wave
export N_mn_wave
export M_mn_wave_SVector
export N_mn_wave_SVector

export B_mn_of_θ
export C_mn_of_θ
export P_mn_of_θ

export B_mn_of_θ_ϕ
export C_mn_of_θ_ϕ
export P_mn_of_θ_ϕ

#############################################################################################
# calculate B(θ), C(θ), P(θ)
"""
    Calculation of B(θ), equation C.19, returns array
I assume each of m, n, θ is a single number
"""
function B_mn_of_θ(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    return vcat(
        zero(θ),                  # r-component
        τₘₙ(m, n, θ),      # θ-component
        im * πₘₙ(m, n, θ)  # ϕ-component
    ) # equation C.19
end

"""
    Calculation of B(θ), equation C.19, returns SVector
I assume each of m, n, θ is a single number
"""
function B_mn_of_θ_SVector(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    """
    I assume each of m, n, θ is a single number
    """
    return SVector(
        zero(θ),                  # r-component
        τₘₙ(m, n, θ),      # θ-component
        im * πₘₙ(m, n, θ)  # ϕ-component
    ) # equation C.19
end

"""
    Calculation of C(θ), equation C.20, returns array
I assume each of m, n, θ is a single number
"""
function C_mn_of_θ(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    return vcat(
        zero(θ),                  # r-component
        im * πₘₙ(m, n, θ), # θ-component
        -1 * τₘₙ(m, n, θ),    # ϕ-component
    ) # equation C.20
end

"""
    Calculation of C(θ), equation C.20, returns SVector
I assume each of m, n, θ is a single number
"""
function C_mn_of_θ_SVector(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    return SVector(
        zero(θ),                  # r-component
        im * πₘₙ(m, n, θ), # θ-component
        -1 * τₘₙ(m, n, θ),    # ϕ-component
    ) # equation C.20
end

# TODO use partial application for M,N, fold them into the type sig of struct P
"""
    equation C.21, returns array
"""
function P_mn_of_θ(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    # TODO @Alok, should I replace arrays with SMatrix? what are the drawbacks?
    # TODO: replace "0" with zero(type)
    return vcat(
        wignerdjmn(n, 0, m, θ), # r-component
        zero(θ),                # θ-component
        zero(θ),                # ϕ-component
    ) # equation C.21
end

"""
    equation C.21, returns SVector
"""
function P_mn_of_θ_SVector(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    # TODO @Alok, should I replace arrays with SMatrix? what are the drawbacks?
    # TODO: replace "0" with zero(type)
    return SVector(
        wignerdjmn(n, 0, m, θ), # r-component
        zero(θ),                # θ-component
        zero(θ),                # ϕ-component
    ) # equation C.21
end

#############################################################################################
# calculate B(θ,ϕ), C(θ,ϕ), P(θ,ϕ), returns Array
function B_mn_of_θ_ϕ(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    return (-1)^m * convert(R, sqrt_factorial_n_plus_m_over_factorial_n_minus_m(m, n)) * B_mn_of_θ(m, n, θ) * exp(im * m * ϕ) # equation C.16
end

function C_mn_of_θ_ϕ(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    return (-1)^m * convert(R, sqrt_factorial_n_plus_m_over_factorial_n_minus_m(m, n)) * C_mn_of_θ(m, n, θ) * exp(im * m * ϕ) # equation C.17
end

function P_mn_of_θ_ϕ(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    return (-1)^m * convert(R, sqrt_factorial_n_plus_m_over_factorial_n_minus_m(m, n)) * P_mn_of_θ(m, n, θ) * exp(im * m * ϕ) # equation C.18
end

# calculate B(θ,ϕ), C(θ,ϕ), P(θ,ϕ), returns SVector
function B_mn_of_θ_ϕ_SVector(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    return (-1)^m * convert(R, sqrt_factorial_n_plus_m_over_factorial_n_minus_m(m, n)) * B_mn_of_θ_SVector(m, n, θ) * exp(im * m * ϕ) # equation C.16
end

function C_mn_of_θ_ϕ_SVector(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    return (-1)^m * convert(R, sqrt_factorial_n_plus_m_over_factorial_n_minus_m(m, n)) * C_mn_of_θ_SVector(m, n, θ) * exp(im * m * ϕ) # equation C.17
end

function P_mn_of_θ_ϕ_SVector(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    return (-1)^m * convert(R, sqrt_factorial_n_plus_m_over_factorial_n_minus_m(m, n)) * P_mn_of_θ_SVector(m, n, θ) * exp(im * m * ϕ) # equation C.18
end

#############################################################################################
# calculate spherical Bessel and Hankel functions and their derivarive
#function spherical_Bessel_j_n(n::I, x::NN) where {NN <: Number, I <: Integer}
    """
    Spherical Bessel function of the first kind.
    It can be calculated from ordinary Bessel function of the first kind "besselj" as in the code.
    check https://en.wikipedia.org/wiki/Bessel_function#Spherical_Bessel_functions
    """
#    return √(π / 2x) * besselj(n + 1 / 2, x)
#end

#function spherical_Bessel_y_n(n::I, x::NN) where {NN <: Number, I <: Integer}
    """
    Spherical Bessel function of the second kind.
    It can be calculated from ordinary Bessel function of the second kind "bessely" as in the code.
    check https://en.wikipedia.org/wiki/Bessel_function#Spherical_Bessel_functions
    """
#    return √(π / 2x) * bessely(n + 1 / 2, x)
#end

# overriding Bessel functions, not using SpecialFunctions
include("bessel.jl")
spherical_Bessel_j_n = spherical_Bessel_j_n_ELZOUKA
spherical_Bessel_y_n = spherical_Bessel_y_n_ELZOUKA

function spherical_Hankel_h1_n(n::I, x::NN) where {NN <: Number, I <: Integer}
    """
    Spherical Hankel function of the first kind. It can be calculated from spherical Bessel functions of the first and second kinds as in the code:
    """
    return spherical_Bessel_j_n(n, x) + im * spherical_Bessel_y_n(n, x)
end

function one_over_x_by_∂_x_j_n_by_∂x(n::I, x::NN) where {NN <: Number, I <: Integer}
    """
    Derivative of (spherical Bessel of first kind * x) divided by x
    """
    return (spherical_Bessel_j_n(n - 1, x) - n / x * spherical_Bessel_j_n(n, x))
end

function one_over_x_by_∂_x_y_n_by_∂x(n::I, x::NN) where {NN <: Number, I <: Integer}
    """
    Derivative of (spherical Bessel of second kind * x) divided by x
    """
    return (spherical_Bessel_y_n(n - 1, x) - n / x * spherical_Bessel_y_n(n, x))
end

function one_over_x_by_∂_x_h_n_by_∂x(n::I, x::NN) where {NN <: Number, I <: Integer}
    """
    Derivative of (spherical Hankel of first kind * x) divided by x
    """
    return one_over_x_by_∂_x_j_n_by_∂x(n, x) + im * one_over_x_by_∂_x_y_n_by_∂x(n, x)
end

#############################################################################################
# calculate (Rg)M(kr,θ,ϕ), (Rg)N(kr,θ,ϕ), returns Array
"""
    Parameters
    ==========
    kind: string, either ["regular" or "incoming"] or ["irregular" or "outgoing"]
"""
function M_mn_wave(m::I, n::I, kr::NN, θ::R, ϕ::R; kind="regular") where {R <: Real, I <: Integer, NN <: Number}
    radial_function, _ = get_radial_function_and_special_derivative_given_kind(kind)
    return convert(R, γ_mn(m, n)) * radial_function(n, kr) * C_mn_of_θ_ϕ(m, n, θ, ϕ)
end

"""
    Parameters
    ==========
    kind: string, either ["regular" or "incoming"] or ["irregular" or "outgoing"]
"""
function N_mn_wave(m::I, n::I, kr::NN, θ::R, ϕ::R; kind="regular") where {R <: Real, I <: Integer, NN <: Number}
    radial_function, radial_function_special_derivative  = get_radial_function_and_special_derivative_given_kind(kind)
    return convert(R, γ_mn(m, n)) * (
        n * (n + 1) / kr * radial_function(n, kr) * P_mn_of_θ_ϕ(m, n, θ, ϕ)
        + (radial_function_special_derivative(n, kr) * B_mn_of_θ_ϕ(m, n, θ, ϕ))
    )
end

# calculate (Rg)M(kr,θ,ϕ), (Rg)N(kr,θ,ϕ), returns SVector
"""
    Parameters
    ==========
    kind: string, either ["regular" or "incoming"] or ["irregular" or "outgoing"]
"""
function M_mn_wave_SVector(m::I, n::I, kr::NN, θ::R, ϕ::R; kind="regular") where {R <: Real, I <: Integer, NN <: Number}
    radial_function, _ = get_radial_function_and_special_derivative_given_kind(kind)
    return convert.(typeof(Complex(θ,θ)),
        convert(R, γ_mn(m, n)) * radial_function(n, kr) * C_mn_of_θ_ϕ_SVector(m, n, θ, ϕ)
    ) # make sure the output is of the same type as the input. # TODO: find a better way
end

"""
    Parameters
    ==========
    kind: string, either ["regular" or "incoming"] or ["irregular" or "outgoing"]
"""
function N_mn_wave_SVector(m::I, n::I, kr::NN, θ::R, ϕ::R; kind="regular") where {R <: Real, I <: Integer, NN <: Number}
    radial_function, radial_function_special_derivative  = get_radial_function_and_special_derivative_given_kind(kind)
    return convert.(typeof(Complex(θ,θ)),
            convert(R, γ_mn(m, n)) * (
            n * (n + 1) / kr * radial_function(n, kr) * P_mn_of_θ_ϕ_SVector(m, n, θ, ϕ)
            + (radial_function_special_derivative(n, kr) * B_mn_of_θ_ϕ_SVector(m, n, θ, ϕ))
        )
    )
end


"""
    _SOCtoPSDaff{T}(f::MOI.VectorAffineFunction{T}, g::MOI.ScalarAffineFunction{T})

Builds a VectorAffineFunction representing the upper (or lower) triangular part of the matrix
[ f[1]     f[2:end]' ]
[ f[2:end] g * I     ]
"""
function _SOCtoPSDaff(f::MOI.VectorAffineFunction{T}, g::MOI.ScalarAffineFunction{T}) where T
    dim = MOI.output_dimension(f)
    n = div(dim * (dim+1), 2)
    # Needs to add t*I
    N0 = length(f.terms)
    Ni = length(g.terms)
    N = N0 + (dim-1) * Ni
    terms = Vector{MOI.VectorAffineTerm{T}}(undef, N)
    terms[1:N0] = MOI.VectorAffineTerm.(map(t -> trimap.(t.output_index, 1), f.terms),
                                        MOI.ScalarAffineTerm.(map(t -> t.scalar_term.coefficient, f.terms),
                                                              map(t -> t.scalar_term.variable_index, f.terms)))
    constant = [f.constants; zeros(T, n - dim)]
    cur = N0
    for i in 2:dim
        k = trimap(i, i)
        terms[cur.+(1:Ni)]  = MOI.VectorAffineTerm.(k, MOI.ScalarAffineTerm.(map(t -> t.coefficient, g.terms),
                                                                             map(t -> t.variable_index, g.terms)))
        constant[k] = g.constant
        cur += Ni
    end
    MOI.VectorAffineFunction(terms, constant)
end

"""
The `SOCtoPSDBridge` transforms the second order cone constraint ``\\lVert x \\rVert \\le t`` into the semidefinite cone constraints
```math
\\begin{pmatrix}
  t & x^\\top\\\\
  x & tI
\\end{pmatrix} \\succeq 0
```
Indeed by the Schur Complement, it is positive definite iff
```math
\\begin{align*}
  tI & \\succ 0\\\\
  t - x^\\top (tI)^{-1} x & \\succ 0
\\end{align*}
```
which is equivalent to
```math
\\begin{align*}
  t & > 0\\\\
  t^2 & > x^\\top x
\\end{align*}
```
"""
struct SOCtoPSDBridge{T} <: AbstractBridge
    dim::Int
    cr::CI{MOI.VectorAffineFunction{T}, MOI.PositiveSemidefiniteConeTriangle}
end
function bridge_constraint(::Type{SOCtoPSDBridge{T}}, model, f,
                           s::MOI.SecondOrderCone) where T
    d = MOI.dimension(s)
    cr = MOI.add_constraint(model, _SOCtoPSDaff(f, T), MOI.PositiveSemidefiniteConeTriangle(d))
    SOCtoPSDBridge(d, cr)
end

_SOCtoPSDaff(f::MOI.VectorOfVariables, ::Type{T}) where T = _SOCtoPSDaff(MOI.VectorAffineFunction{T}(f), T)
_SOCtoPSDaff(f::MOI.VectorAffineFunction, ::Type) = _SOCtoPSDaff(f, MOIU.eachscalar(f)[1])

MOI.supports_constraint(::Type{SOCtoPSDBridge{T}}, ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}}, ::Type{MOI.SecondOrderCone}) where T = true
added_constraint_types(::Type{SOCtoPSDBridge{T}}, ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}}, ::Type{MOI.SecondOrderCone}) where T = [(MOI.VectorAffineFunction{T}, MOI.PositiveSemidefiniteConeTriangle)]

function MOI.get(model::MOI.AbstractOptimizer, a::MOI.ConstraintPrimal, c::SOCtoPSDBridge)
    MOI.get(model, a, c.cr)[trimap.(1:c.dim, 1)]
end
function MOI.get(model::MOI.AbstractOptimizer, a::MOI.ConstraintDual, c::SOCtoPSDBridge)
    dual = MOI.get(model, a, c.cr)
    tdual = sum(i -> dual[trimap(i, i)], 1:c.dim)
    [tdual; dual[trimap.(2:c.dim, 1)]*2]
end

MOI.get(::SOCtoPSDBridge{T}, ::MOI.NumberOfConstraints{MOI.VectorAffineFunction{T}, MOI.PositiveSemidefiniteConeTriangle}) where T = 1
MOI.get(b::SOCtoPSDBridge{T}, ::MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T}, MOI.PositiveSemidefiniteConeTriangle}) where T = [b.cr]

function MOI.delete(model::MOI.AbstractOptimizer, c::SOCtoPSDBridge)
    MOI.delete(model, c.cr)
end

"""
The `RSOCtoPSDBridge` transforms the second order cone constraint ``\\lVert x \\rVert \\le 2tu`` with ``u \\ge 0`` into the semidefinite cone constraints
```math
\\begin{pmatrix}
  t & x^\\top\\\\
  x & 2uI
\\end{pmatrix} \\succeq 0
```
Indeed by the Schur Complement, it is positive definite iff
```math
\\begin{align*}
  uI & \\succ 0\\\\
  t - x^\\top (2uI)^{-1} x & \\succ 0
\\end{align*}
```
which is equivalent to
```math
\\begin{align*}
  u & > 0\\\\
  2tu & > x^\\top x
\\end{align*}
```
"""
struct RSOCtoPSDBridge{T} <: AbstractBridge
    dim::Int
    cr::CI{MOI.VectorAffineFunction{T}, MOI.PositiveSemidefiniteConeTriangle}
end

MOI.supports_constraint(::Type{RSOCtoPSDBridge{T}}, ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}}, ::Type{MOI.RotatedSecondOrderCone}) where T = true
added_constraint_types(::Type{RSOCtoPSDBridge{T}}, ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}}, ::Type{MOI.RotatedSecondOrderCone}) where T = [(MOI.VectorAffineFunction{T}, MOI.PositiveSemidefiniteConeTriangle)]

function bridge_constraint(::Type{RSOCtoPSDBridge{T}}, model, f,
                           s::MOI.RotatedSecondOrderCone) where T
    d = MOI.dimension(s)-1
    cr = MOI.add_constraint(model, _RSOCtoPSDaff(f, T), MOI.PositiveSemidefiniteConeTriangle(d))
    RSOCtoPSDBridge(d, cr)
end

_RSOCtoPSDaff(f::MOI.VectorOfVariables, ::Type{T}) where T = _RSOCtoPSDaff(MOI.VectorAffineFunction{T}(f), T)
function _RSOCtoPSDaff(f::MOI.VectorAffineFunction, ::Type{T}) where T
    n = MOI.output_dimension(f)
    f_scalars = MOIU.eachscalar(f)
    g = MOIU.operate!(*, T, f_scalars[2], convert(T, 2))
    _SOCtoPSDaff(f_scalars[[1; 3:n]], g)
end

function MOI.get(model::MOI.AbstractOptimizer, a::MOI.ConstraintPrimal, c::RSOCtoPSDBridge)
    x = MOI.get(model, MOI.ConstraintPrimal(), c.cr)[[trimap(1, 1); trimap(2, 2); trimap.(2:c.dim, 1)]]
    x[2] /= 2 # It is (2u*I)[1,1] so it needs to be divided by 2 to get u
    x
end
function MOI.get(model::MOI.AbstractOptimizer, a::MOI.ConstraintDual, c::RSOCtoPSDBridge)
    dual = MOI.get(model, MOI.ConstraintDual(), c.cr)
    udual = sum(i -> dual[trimap(i, i)], 2:c.dim)
    [dual[1]; 2udual; dual[trimap.(2:c.dim, 1)]*2]
end

MOI.get(::RSOCtoPSDBridge{T}, ::MOI.NumberOfConstraints{MOI.VectorAffineFunction{T}, MOI.PositiveSemidefiniteConeTriangle}) where T = 1
MOI.get(b::RSOCtoPSDBridge{T}, ::MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T}, MOI.PositiveSemidefiniteConeTriangle}) where T = [b.cr]

function MOI.delete(model::MOI.AbstractOptimizer, c::RSOCtoPSDBridge)
    MOI.delete(model, c.cr)
end


"""
    latlong_to_xy(latitude, longitude, utm_zone; isnorth=true, units="m")

Converts arrays of points from latitude and longitude to x and y in meters in a local
coordinate frame based on the point with the lowest magnitude latitude,

# Arguments
- `latitude::Array{Float,N}`
- `longitude::Array{Float,N}`
- `isnorth::Float`: specifies if the point is in the northern hemisphere (defaul: true)
"""
function latlong_to_xy(latitude, longitude, utm_zone; isnorth=true)

    # get number of points
    npoints = length(latitude)

    # set up utm region
    utmregion = gd.UTMfromLLA(utm_zone, isnorth, wgs84)

    # get zero point index
    zp = argmin(latitude)

    # get zero point lat long
    minlla = gd.LLA(latitude[zp], longitude[zp])

    # get zero point utm
    minutm = utmregion(minlla)

    # get x and y for all points
    x = zeros(npoints)
    y = zeros(npoints)
    for i in 1:npoints
        # lat and long for current point
        lla = gd.LLA(latitude[i], longitude[i])

        # convert to utm coordinates
        utm = utmregion(lla)

        # get x location
        x[i] = utm.x - minutm.x

        # get y location
        y[i] = utm.y - minutm.y

    end

    return x, y
end

"""
    hermite_spline(x, x0, x1, y0, dy0, y1, dy1)

Produces the y and (optionally) dy values for a hermite cubic spline
interpolating between two end points with known slopes

# Arguments
- `x::Float`: x position of output y
- `x0::Float`: x position of upwind endpoint of spline
- `x1::Float`: x position of downwind endpoint of spline
- `y0::Float`: y position of upwind endpoint of spline
- `dy0::Float`: slope at upwind endpoint of spline
- `y1::Float`: y position of downwind endpoint of spline
- `dy1::Float`: slope at downwind endpoint of spline
"""
function hermite_spline(x, x0, x1, y0, dy0, y1, dy1; return_deriv=false)

    # initialize coefficients for parametric Hermite cubic spline
    c3 = (2.0*(y1))/(x0^3 - 3.0*x0^2*x1 + 3.0*x0*x1^2 - x1^3) -
        (2.0*(y0))/(x0^3 - 3.0*x0^2*x1 + 3.0*x0*x1^2 - x1^3) +
        (dy0)/(x0^2 - 2.0*x0*x1 + x1^2) +
        (dy1)/(x0^2 - 2.0*x0*x1 + x1^2)

    c2 = (3.0*(y0)*(x0 + x1))/(x0^3 - 3.0*x0^2*x1 + 3.0*x0*x1^2 - x1^3) -
        ((dy1)*(2.0*x0 + x1))/(x0^2 - 2.0*x0*x1 + x1^2) - ((dy0)*(x0 +
        2.0*x1))/(x0^2 - 2.0*x0*x1 + x1^2) - (3.0*(y1)*(x0 + x1))/(x0^3 -
        3.0*x0^2*x1 + 3.0*x0*x1^2 - x1^3)

    c1 = ((dy0)*(x1^2 + 2.0*x0*x1))/(x0^2 - 2.0*x0*x1 + x1^2) + ((dy1)*(x0^2 +
        2.0*x1*x0))/(x0^2 - 2.0*x0*x1 + x1^2) - (6.0*x0*x1*(y0))/(x0^3 -
        3.0*x0^2*x1 + 3.0*x0*x1^2 - x1^3) + (6.0*x0*x1*(y1))/(x0^3 -
        3.0*x0^2*x1 + 3.0*x0*x1^2 - x1^3)

    c0 = ((y0)*(- x1^3 + 3.0*x0*x1^2))/(x0^3 - 3.0*x0^2*x1 + 3.0*x0*x1^2 -
        x1^3) - ((y1)*(- x0^3 + 3.0*x1*x0^2))/(x0^3 - 3.0*x0^2*x1 +
        3.0*x0*x1^2 - x1^3) - (x0*x1^2*(dy0))/(x0^2 - 2.0*x0*x1 + x1^2) -
        (x0^2*x1*(dy1))/(x0^2 - 2.0*x0*x1 + x1^2)

    # Solve for y and dy values at the given point
    y = c3*x^3 + c2*x^2 + c1*x + c0
    dy_dx = c3*3*x^2 + c2*2*x + c1

    if return_deriv
        return y, dy_dx
    else
        return y
    end
end

"""
    overlap_area_func(turbine_y, turbine_z, rotor_diameter, wake_center_y,
    wake_center_z, wake_diameter; tol=1E-6)

Produces the y and (optionally) dy values for a hermite cubic spline
interpolating between two end points with known slopes

All dimensions should be in meters

# Arguments
- `turbine_y::Float`: cross wind location of turbine hub
- `turbine_z::Float`: vertical location of turbine hub
- `rotor_diameter::Float`
- `wake_center_y::Float`: cross wind location of wake center
- `wake_center_z::Float`: vertical location of wake center
- `wake_diameter::Float`
- `tol::Float`: extra distance for comparisons. Default 1E-6
"""
# calculates the overlap area between a given wake and a rotor area
function overlap_area_func(turbine_y, turbine_z, rotor_diameter, wake_center_y,
    wake_center_z, wake_diameter; tol=1E-6)

    # distance between wake center and rotor center
    if (wake_center_z > (turbine_z + tol)) || (wake_center_z < (turbine_z - tol))
        OVdYd = sqrt((wake_center_y-turbine_y)^2 + (wake_center_z - turbine_z)^2)
    elseif (wake_center_y > (turbine_y + tol))
        OVdYd = wake_center_y - turbine_y
    elseif (turbine_y > (wake_center_y + tol))
        OVdYd = turbine_y - wake_center_y
    else
        OVdYd = 0.0
    end

    # find rotor radius
    OVr = rotor_diameter/2.0

    # find wake radius
    OVRR = wake_diameter/2.0

    # make sure the distance from wake center to turbine hub is positive
    OVdYd = abs(OVdYd)

    # determine if there is overlap
    if (OVdYd < (OVr+OVRR)) # if the rotor overlaps the wake zone

        # check that turbine and wake centers are not perfectly aligned
        if (OVdYd > (0.0 + tol))

            # check if the rotor is wholly contained in the wake
            if ((OVdYd + OVr) < OVRR + tol)
                # wake_overlap = pi*OVr*OVr
                wake_overlap = 3.1415926535897*OVr*OVr
            elseif ((OVdYd + OVRR) < OVr + tol)
                # wake_overlap = pi*OVRR*OVRR
                wake_overlap = 3.1415926535897*OVRR*OVRR
            else
                # calculate the distance from the wake center to the chord connecting the lens cusps
                OVL = (-OVr*OVr+OVRR*OVRR+OVdYd*OVdYd)/(2.0*OVdYd)

                OVz = sqrt(OVRR*OVRR-OVL*OVL)
                OVz2 = sqrt(OVr*OVr-(OVdYd-OVL)*(OVdYd-OVL))

                wake_overlap = OVRR*OVRR*acos(OVL/OVRR) + OVr*OVr*acos((OVdYd-OVL)/OVr) - OVL*OVz - (OVdYd-OVL)*OVz2
            end

        # perfect overlap case where the wake is larger than the rotor
        elseif (OVRR > OVr)
            wake_overlap = pi*OVr*OVr
        # perfect overlap case where the rotor is larger than the wake
        else
            wake_overlap = pi*OVRR*OVRR
        end

    # case with no overlap
    else
        wake_overlap = 0.0
    end

end

"""
    smooth_max_ndim(x; s=100.0)

Calculate the smoothmax (a.k.a. softmax or LogSumExponential) of the elements in x.

Based on John D. Cook's writings at
(1) https://www.johndcook.com/blog/2010/01/13/soft-maximum/
and
(2) https://www.johndcook.com/blog/2010/01/20/how-to-compute-the-soft-maximum/

# Arguments
- `x::Float`: first value for comparison
- `y::Float`: second value for comparison
- `s::Float` : controls the level of smoothing used in the smooth max
"""
function smooth_max(x, y; s=10.0)

    # LogSumExponential Method - used this in the past
    # g = (x*exp(s*x)+y*exp(s*y))/(exp(s*x)+exp(s*y))

    # non-overflowing version of Smooth Max function (see ref 2 above)
    max_val = max(x, y)
    min_val = min(x, y)
    r = (log(1.0 + exp(s*(min_val - max_val))) + s*max_val)/s

    return r

end

"""
    smooth_max(x; s=10.0)

Calculate the smoothmax (a.k.a. softmax or LogSumExponential) of the elements in x.

Based on John D. Cook's writings at
(1) https://www.johndcook.com/blog/2010/01/13/soft-maximum/
and
(2) https://www.johndcook.com/blog/2010/01/20/how-to-compute-the-soft-maximum/

And based on article in FeedlyBlog
(3) https://blog.feedly.com/tricks-of-the-trade-logsumexp/

# Arguments
- `x::Array{Float,1}` : vector with all the input values
- `s::Float` : controls the level of smoothing used in the smooth max
"""
function smooth_max(x; s=10.0)

    # non-overflowing version of Smooth Max function (see ref 2 and 3 above)

    # get the maximum value and the index of maximum value
    max_val, max_ind = findmax(x)

    # get the indices of x
    indices = collect(1:length(x))

    # remove the index of the maximum value
    splice!(indices, max_ind)

    # LogSumExp with smoothing factor s
    r = (log(sum([1.0; exp.(s*(x[indices].-max_val))])) + s*max_val)/s

    return r

end

"""
    closeBndryLists(bndryPts_x, bndryPts_y)

Appends the 1st element to the end of each array for a closed boundary.
Note, this will not function properly if there is only one region.
For only one region, use `closeBndryList(bndryPts_x, bndryPts_y)` (note the
singular, not plural 'List' in the function title)

# Arguments
- `bndryPts_x::Array{Float,1}` : N-D array of x-coordinates for the vertices
        around N-many closed boundaries
- `bndryPts_y::Array{Float,1}` : N-D array of y-coordinates for the vertices
        around N-many closed boundaries
"""
function closeBndryLists(region_bndry_x, region_bndry_y)
    # Determine how many regions and points per region were passed
    nRegions = length(region_bndry_x)
    # For every region
    for i in 1:nRegions
        # Append the initial points to the end of that row (if needed)
        region_bndry_x[i], region_bndry_y[i] = closeBndryList(region_bndry_x[i], region_bndry_y[i])
    end

    return region_bndry_x, region_bndry_y
end

"""
    closeBndryList(bndryPts_x, bndryPts_y)

Appends the 1st element to the end of the coordinate arrays if it is not already
repeated. Note, this will only work on 1-D arrays. For an array of 1-D arrays,
use `closeBndryLists(bndryPts_x, bndryPts_y)` (note the
plural, not singular 'Lists' in the function title)

# Arguments
- `bndryPts_x::Array{Float,1}` : 1-D array of x-coordinates for the vertices
        around a singlar closed boundary
- `bndryPts_y::Array{Float,1}` : 1-D array of y-coordinates for the vertices
        around a singlar closed boundary
"""
function closeBndryList(bndry_x, bndry_y)
    # If the end point isn't already a repeat of the staring point
    if !((bndry_x[1] == bndry_x[end]) && (bndry_y[1] == bndry_y[end]))
        # Append the first point to the end of the array
        bndryPts_x_clsd = push!(bndry_x, bndry_x[1])
        bndryPts_y_clsd = push!(bndry_y, bndry_y[1])
    end

    return bndryPts_x_clsd, bndryPts_y_clsd
end

"""
    PointsOnCircum(center_x, center_y, r, n = 100)

Given a circle center, radius, and number of discrete points, returns an array
of discrete points along the circle's circumference

# Arguments
- `center_x::Float64` : cartesian x-coordinate for the center of the circle
- `center_y::Float64` : cartesian y-coordinate for the center of the circle
- `r::Float64` : distance from circle's center to the circumference points
- `n::Float64` : defaults to 100, is the number of discrete evenly-spaced points
        that will be returned along the circle's circumference
"""
function DiscreteCircum(center_x, center_y, r, n = 100)
    bndry_x = zeros(n+1)
    bndry_y = zeros(n+1)
    for i in 0:n
        bndry_x[i+1] = center_x + cos(2*pi/n*i)*r
        bndry_y[i+1] = center_y + sin(2*pi/n*i)*r
    end
    return bndry_x, bndry_y
end

"""
    calcMinorAngle(bndry_x, bndry_y, bndry_z=[0,0,0])

Given three points in space, calculates the magnitude of the non-reflex angle
formed at the center point. Created to be used in VR_bounary_startup()

# Arguments
- `bndry_x::Array{Float,1}` : 1-D array of x-coordinates for the vertices
        around a singlar closed boundary
- `bndry_y::Array{Float,1}` : 1-D array of y-coordinates for the vertices
        around a singlar closed boundary
- `bndry_z::Array{Float,1}` : 1-D array of z-coordinates for the vertices
        around a singlar closed boundary. Default to [0,0,0] for X-Y plane
"""
function calcMinorAngle(bndry_x, bndry_y, bndry_z=[0,0,0])
    # Calculates the magnitude of the non-reflex angle formed at the center point
    ABx = bndry_x[2]-bndry_x[1]
    ABy = bndry_y[2]-bndry_y[1]
    ABz = bndry_z[2]-bndry_z[1]

    BCx = bndry_x[2]-bndry_x[3]
    BCy = bndry_y[2]-bndry_y[3]
    BCz = bndry_z[2]-bndry_z[3]

    Num = (ABx*BCx) + (ABy*BCy) + (ABz*BCz) # Dot Product

    Denom = sqrt(ABx^2 + ABy^2 + ABz^2) * sqrt(BCx^2 + BCy^2 + BCz^2) # Multiplication of magnitudes
    Theta = acosd(Num/Denom) # Get the angle formed

    # If it's greater than 180, get the
    if (Theta > 180.0)
        Theta = 360.0 - Theta    # Get the explementary angle
    end

    return Theta
end

"""
    calcSmallestAngle(bndry_x_clsd, bndry_y_clsd)

Given a 1-D closed array of boundary verticies (with first point repeated at the
end) it determines the smallest non-reflex angle created by any three
consecutive verticies along the boundary. Created to be used in
VR_bounary_startup()

# Arguments
- `bndry_x::Array{Float,1}` : 1-D array of x-coordinates for the vertices
        around a singlar closed boundary
- `bndry_y::Array{Float,1}` : 1-D array of y-coordinates for the vertices
        around a singlar closed boundary
"""
function calcSmallestAngle(bndry_x_clsd, bndry_y_clsd)
    num_angles = length(bndry_x_clsd)-1
    # Loop the second point back on so we get the angle at the beginning
    if !(((bndry_x_clsd[1] == bndry_x_clsd[end-1]) && (bndry_y_clsd[1] == bndry_y_clsd[end-1])) && ((bndry_x_clsd[2] == bndry_x_clsd[end]) && (bndry_y_clsd[2] == bndry_y_clsd[end])))
        bndryPts_x_loopd = vcat(bndry_x_clsd, bndry_x_clsd[2])
        bndryPts_y_loopd = vcat(bndry_y_clsd, bndry_y_clsd[2])
    end

    #- Calculate the smallest angle -#
    smallest_angle = 360
    for i in 1:num_angles
        temp_angle = calcMinorAngle(bndryPts_x_loopd[i:i+2], bndryPts_y_loopd[i:i+2])

        if (temp_angle < smallest_angle)
            smallest_angle = temp_angle
        end
    end

    return smallest_angle
end

"""
    getPerimeterLength(bndry_x_clsd, bndry_y_clsd)

Given a 1-D closed array of boundary verticies (with first point repeated at the
end) returns the length along the perimeter. Created to be used in
VR_bounary_startup()

# Arguments
- `bndry_x::Array{Float,1}` : 1-D array of x-coordinates for the vertices
        around a singlar closed boundary
- `bndry_y::Array{Float,1}` : 1-D array of y-coordinates for the vertices
        around a singlar closed boundary
"""
function getPerimeterLength(bndry_x_clsd, bndry_y_clsd)
    num_bndry_pts = length(bndry_x_clsd)-1
    nLength = zeros(num_bndry_pts)
    for i in 1:num_bndry_pts
        nLength[i] = sqrt(
                 (bndry_x_clsd[i+1]-bndry_x_clsd[i])^2
                +(bndry_y_clsd[i+1]-bndry_y_clsd[i])^2)
    end

    return nLength
end

"""
    coordDist(x1, y1, x2, y2)

Given a two points (x1, y1) and (x2, y2), returns the euclidean distance between
them

# Arguments
- `x1::Float64` : x-coord of the first point
- `y1::Float64` : y-coord of the first point
- `x2::Float64` : x-coord of the second point
- `y2::Float64` : y-coord of the second point
"""
function coordDist(x1, y1, x2, y2)
    xDiff = x1 - x2
    yDiff = y1 - y2
    return sqrt(xDiff^2 + yDiff^2)
end


# Tests for mutual information function
# (NB These tests will fail with the shrinkage estimator, because the 1D,
# 2D and 3D distributions on the right will shrink towards different
# targets by default, whereas the mutual information formula overcomes
# by estimating the probabilities of the 3D distribution, then projecting)

# Test mutual information formula
@test get_mutual_information(arr1, arr2) ≈ get_entropy(arr1) + get_entropy(arr2) - get_entropy(arr1, arr2)
println("Mutual information formula passed.")


abstract type AbstractBroadcastedNormal{T} <: ContinuousBatchDistribution end

function rand(rng::AbstractRNG, bd::AbstractBroadcastedNormal, dims::Int...)
    return bd.m .+ std(bd) .* randnsimilar(rng, bd.m, dims...)
end

function logpdf(bd::AbstractBroadcastedNormal{T}, x) where {T}
    diff = x .- bd.m
    return -(log(2 * T(pi)) .+ logvar(bd) .+ diff .* diff ./ var(bd)) / 2
end

mean(bd::AbstractBroadcastedNormal) = bd.m
mode(bd::AbstractBroadcastedNormal) = bd.m

_vlogv(bd::AbstractBroadcastedNormal) = (var(bd), logvar(bd))

function kldiv(bd1::AbstractBroadcastedNormal{T}, bd2::AbstractBroadcastedNormal{T}) where {T}
    diff = bd2.m .- bd1.m
    v1, logv1 = _vlogv(bd1)
    v2, logv2 = _vlogv(bd2)
    return (logv2 .- logv1 .- 1 .+ v1 ./ v2 .+ diff .* diff ./ v2) / 2
end

"""
Broadcasted Normal distribution with standard deviation.
"""
struct BroadcastedNormalStd{T<:Number,Tm<:AbstractArray{T},Ts<:AbstractArray{T}} <: AbstractBroadcastedNormal{T}
    m::Tm
    s::Ts
end

Broadcast.broadcastable(bd::BroadcastedNormalStd) = Ref(bd)

Normal(m::AbstractArray, s::AbstractArray) = BroadcastedNormalStd(m, s)
NormalStd(m::AbstractArray, s::AbstractArray) = BroadcastedNormalStd(m, s)

   std(bd::BroadcastedNormalStd) = bd.s
   var(bd::BroadcastedNormalStd) = bd.s.^2
logvar(bd::BroadcastedNormalStd) = 2log.(bd.s)

"""
Broadcasted Normal distribution with variance.
"""
struct BroadcastedNormalVar{T<:Number,Tm<:AbstractArray{T},Tv<:AbstractArray{T}} <: AbstractBroadcastedNormal{T}
    m::Tm
    v::Tv
end

Broadcast.broadcastable(bd::BroadcastedNormalVar) = Ref(bd)

NormalVar(m::AbstractArray, v::AbstractArray) = BroadcastedNormalVar(m, v)

   std(bd::BroadcastedNormalVar) = sqrt.(bd.v)
   var(bd::BroadcastedNormalVar) = bd.v
logvar(bd::BroadcastedNormalVar) = log.(bd.v)

"""
Broadcasted Normal distribution with log standard deviation.
"""
struct BroadcastedNormalLogStd{T<:Number,Tm<:AbstractArray{T},Ts<:AbstractArray{T}} <: AbstractBroadcastedNormal{T}
       m::Tm    # mean
    logs::Ts    # log std
end

Broadcast.broadcastable(bd::BroadcastedNormalLogStd) = Ref(bd)

NormalLogStd(m::AbstractArray, logs::AbstractArray) = BroadcastedNormalLogStd(m, logs)

   std(bd::BroadcastedNormalLogStd) = exp.(bd.logs)
   var(bd::BroadcastedNormalLogStd) = exp.(2bd.logs)
logvar(bd::BroadcastedNormalLogStd) = 2bd.logs

function _vlogv(bd::BroadcastedNormalLogStd)
    logv = logvar(bd)
    v = exp.(logv)
    return v, logv
end

"""
Broadcasted Normal distribution with log variance.
"""
struct BroadcastedNormalLogVar{T<:Number,Tm<:AbstractArray{T},Tv<:AbstractArray{T}} <: AbstractBroadcastedNormal{T}
       m::Tm    # mean
    logv::Tv    # log std
end

Broadcast.broadcastable(bd::BroadcastedNormalLogVar) = Ref(bd)

NormalLogVar(m::AbstractArray, logv::AbstractArray) = BroadcastedNormalLogVar(m, logv)

   std(bd::BroadcastedNormalLogVar) = exp.(bd.logv ./ 2)
   var(bd::BroadcastedNormalLogVar) = exp.(bd.logv)
logvar(bd::BroadcastedNormalLogVar) = bd.logv

function _vlogv(bd::BroadcastedNormalLogVar)
    logv = logvar(bd)
    v = exp.(logv)
    return v, logv
end

###

# using LinearAlgebra: det, tr, inv
# using Distributions: MvNormal

# """
#     kldiv(mvn1::MvNormal, mvn2::MvNormal)

# Compute ``KL(MvNormal_1||MvNormal_2)``.
# """
# function kldiv(mvn1::MvNormal, mvn2::MvNormal)
#     d = length(mvn1.μ)
#     diff = mvn2.μ .- mvn1.μ
#     Σ1 = Matrix(mvn1.Σ)
#     Σ2 = Matrix(mvn2.Σ)
#     return (log(det(Σ2)) - log(det(Σ1)) - d + tr(inv(Σ2) * Σ1) + diff' * inv(Σ2) * diff) ./ 2
# end


using DifferentialEvolutionMCMC, Test, Random, Turing, Distributions
Random.seed!(973536)

bounds = ((-Inf,Inf),(0.0,Inf))

data = rand(Normal(0.0, 1.0), 100)

function prior_loglike(μ, σ)
    LL = 0.0
    LL += logpdf(Normal(0, 10), μ)
    LL += logpdf(truncated(Cauchy(0, 1), 0, Inf), σ)
    return LL
end

function sample_prior()
    μ = rand(Normal(0, 10))
    σ = rand(truncated(Cauchy(0, 1), 0, Inf))
    return [μ,σ]
end

bounds = ((-Inf,Inf),(0.0,Inf))

data = rand(Normal(0.0, 1.0), 50)

function loglike(data, μ, σ)
    return sum(logpdf.(Normal(μ, σ), data))
end

names = (:μ,:σ)

model = DEModel(;
    sample_prior,
    prior_loglike,
    loglike,
    data,
    names
)

de = DE(;sample_prior, bounds, burnin=1500, Np=6)
n_iter = 3000
chains = sample(model, de, n_iter, progress=true)
μ_de = describe(chains)[1][:,:mean]
σ_de = describe(chains)[1][:,:std]
rhat = describe(chains)[1][:,:rhat]

@model model1(data) = begin
    μ ~ Normal(0, 10)
    σ ~ truncated(Cauchy(0, 1), 0.0, Inf)
    for i in 1:length(data)
        data[i] ~ Normal(μ, σ)
    end
end
chn = sample(model1(data), NUTS(1000, .85), 2000)
μ_nuts = describe(chn)[1][:,:mean]
σ_nuts = describe(chn)[1][:,:std]

@test all(isapprox.(rhat, fill(1.0, 2), atol=.05))
@test all(isapprox.(μ_nuts, μ_de, atol=.01))
@test all(isapprox.(σ_nuts, σ_de, atol=.01))


function extrapolate_left!(f, r, rₘᵢₙ = 1e-7)
  m = 2
  for i in eachindex(r)
      if r[i] ≥ rₘᵢₙ
          m = i
          break
      end
  end
  tg = ((f[m-1] - f[m])*(r[m+1] - r[m])^2 - (f[m+1] - f[m])*(r[m-1] - r[m])^2) /
       ((r[m+1] - r[m])*(r[m-1] - r[m])*(r[m+1] - r[m-1]))
  f0 = f[m] - tg*r[m]
  for i=1:m-1
      f[i] = tg*r[i] + f0
  end
  return m, tg, f0
end

"""
    radial_shr_eq(V::Function = r -> -1/r, x::AbstractRange{T} = -30.0:0.1:5.0;
    conf = 1, μ = 1.0, α = 1e5) where {T}

    radial_shr_eq(V::Array{S}, x::AbstractRange{T};
    conf = 1, μ = 1.0, α = 1e5) where {S, T}

Solve Schrödinger equation ([atomic units](https://en.wikipedia.org/wiki/Hartree_atomic_units) are assumed)
```math
-\\frac{1}{2\\mu}\\left[\\frac{1}{r^{2}}\\frac{\\partial}{\\partial r}\\left(r^{2}\\frac{\\partial R_{nl}}{\\partial r}\\right)-\\frac{l(l+1)}{r^{2}}R_{nl}(r)\\right]+V(r)R_{nl}(r)=E_{nl}R_{nl}(r)
```
for [the particle in a spherically symmetric potential.](https://en.wikipedia.org/wiki/Particle_in_a_spherically_symmetric_potential)

## On input:

* `V` - potential, provided as an explicit function ``V(r)`` (by default `r -> -1/r`), or an array of values on a radial grid.
* `x` - uniform grid (`AbstractRange` object), such as ``r_i = \\exp(x_i)``. Default grid is `-30.0:0.1:5.0`
* `conf` - electronic configuration of interest, by default: ground state `conf = ((1.0),)`
* `μ` - [reduced mass](https://en.wikipedia.org/wiki/Reduced_mass), by default `μ = 1.0` (electron mass in a.u.)
* `α` - parameter ``α = 10^5`` is used in the matrix pencil ``\\mathrm{H} + α\\mathrm{S}`` for the generalized eigenproblem.

## On output function returns:

* sum of one-particle energies: ``E = \\sum_{i=1}^{N} n_i \\epsilon_i``
* particles density: ``\\rho(x) = \\frac{1}{4\\pi} \\sum_{i=1}^{N} n_i y_{i}^{2}(x)``
* orbitals ``y_i(x)``, corresponding eigenvalues ``\\epsilon_i`` (energy levels), azimuthal and radial quantum numbers ``l``, ``n_r``, and level populations ``n_i`` (as listed in `conf`).

## Examples
[Hydrogen atom](https://en.wikipedia.org/wiki/Hydrogen_atom)

```jldoctest
julia> radial_shr_eq(r -> -1/r, conf = conf_enc("1s1")).energy ≈ -1/2
true

julia> radial_shr_eq(r -> -1/r, conf = conf_enc("2s1")).energy ≈ -1/8
true

julia> radial_shr_eq(r -> -1/r, conf = conf_enc("2p1")).energy ≈ -1/8
true

julia> radial_shr_eq(r -> -1/r, conf = conf_enc("3s1")).energy ≈ -1/18
true

julia> radial_shr_eq(r -> -1/r, conf = conf_enc("3p1")).energy ≈ -1/18
true

julia> radial_shr_eq(r -> -1/r, conf = conf_enc("3d1")).energy ≈ -1/18
true
```
"""
function radial_shr_eq(V::AbstractArray,
                       x::AbstractRange = -35.0:0.1:5.0;
                    conf::Array{Array{T,1},1} = [[1]],
                       μ::Real = 1,
                       Α::Real = 1e5) where T <: Real

  if length(V) != length(x)
    throw(DimensionMismatch("V and r have different length"))
  end

  r, r², n, dx = exp.(x), exp.(2x), length(x), step(x)

  D = 1:n+1:n*n

  H = -1/2μ*laplacian(x); HD = H[D] .+ V .* r²
  S = copy(H)

  ε = similar(r); ρ = zero(r); ∑nᵢεᵢ = 0.0

  ψ = Dict()
  # the equation is solved separately for each subshell: s, p, d, f
  for l=0:length(conf)-1
    @. H[D] = HD + 1/2μ * (l + 1/2)^2
    @. @views S[D] = H[D] + Α * r²
    θ, y = eigen(Symmetric(H), Symmetric(S))
    @. ε = Α * θ / (1 - θ)
    for (nᵣ, nᵢ) in enumerate(conf[l+1])
      @views y[:, nᵣ] ./= sqrt(∫(dx, y[:, nᵣ] .^ 2 .* r²))
      @views ρ .+= nᵢ / 4π * y[:, nᵣ] .^ 2
      ∑nᵢεᵢ += nᵢ * ε[nᵣ]
      ψ[(nᵣ = nᵣ - 1, l = l)] = (ϵᵢ = ε[nᵣ], nᵢ = nᵢ, ψᵢ = y[:, nᵣ])
    end
  end
  @. ρ /= r
  extrapolate_left!(ρ, r)
  @. ρ *= r
  return (energy = ∑nᵢεᵢ, density = ρ, orbitals = ψ)
end
radial_shr_eq(V::Function = r -> -1/r, x::AbstractRange = -35.0:0.1:5.0; conf = 1, μ = 1, Α = 1e5) =
              radial_shr_eq(V.(exp.(x)), x, conf = conf, μ = μ, Α = Α)


using FastTransforms, FillArrays, Test
import FastTransforms: clenshaw, clenshaw!, forwardrecurrence!, forwardrecurrence

@testset "clenshaw" begin
    @testset "Chebyshev T" begin
        for elty in (Float64, Float32)
            c = [1,2,3]
            cf = elty.(c)
            @test @inferred(clenshaw(c,1)) ≡ 1 + 2 + 3
            @test @inferred(clenshaw(c,0)) ≡ 1 + 0 - 3
            @test @inferred(clenshaw(c,0.1)) == 1 + 2*0.1 + 3*cos(2acos(0.1))
            @test @inferred(clenshaw(c,[-1,0,1])) == clenshaw!(c,[-1,0,1]) == [2,-2,6]
            @test clenshaw(c,[-1,0,1]) isa Vector{Int}
            @test @inferred(clenshaw(elty[],1)) ≡ zero(elty)

            x = elty[1,0,0.1]
            @test @inferred(clenshaw(c,x)) ≈ @inferred(clenshaw!(c,copy(x))) ≈
                @inferred(clenshaw!(c,x,similar(x))) ≈
                @inferred(clenshaw(cf,x)) ≈ @inferred(clenshaw!(cf,copy(x))) ≈
                @inferred(clenshaw!(cf,x,similar(x))) ≈ elty[6,-2,-1.74]

            @testset "Strided" begin
                cv = view(cf,:)
                xv = view(x,:)
                @test clenshaw!(cv, xv, similar(xv)) == clenshaw!(cf,x,similar(x))

                cv2 = view(cf,1:2:3)
                @test clenshaw!(cv2, xv, similar(xv)) == clenshaw([1,3], x)

                # modifies x and xv
                @test clenshaw!(cv2, xv) == xv == x == clenshaw([1,3], elty[1,0,0.1])
            end

            @testset "matrix coefficients" begin
                c = [1 2; 3 4; 5 6]
                @test clenshaw(c,0.1) ≈ [clenshaw(c[:,1],0.1), clenshaw(c[:,2],0.1)]
                @test clenshaw(c,[0.1,0.2]) ≈ [clenshaw(c[:,1], 0.1) clenshaw(c[:,2], 0.1); clenshaw(c[:,1], 0.2) clenshaw(c[:,2], 0.2)]
            end
        end
    end

    @testset "Chebyshev U" begin
        N = 5
        A, B, C = Fill(2,N-1), Zeros{Int}(N-1), Ones{Int}(N)
        @testset "forwardrecurrence!" begin
            @test @inferred(forwardrecurrence(N, A, B, C, 1)) == @inferred(forwardrecurrence!(Vector{Int}(undef,N), A, B, C, 1)) == 1:N
            @test forwardrecurrence!(Vector{Int}(undef,N), A, B, C, -1) == (-1) .^ (0:N-1) .* (1:N)
            @test forwardrecurrence(N, A, B, C, 0.1) ≈ forwardrecurrence!(Vector{Float64}(undef,N), A, B, C, 0.1) ≈
                    sin.((1:N) .* acos(0.1)) ./ sqrt(1-0.1^2)
        end

        c = [1,2,3]
        @test c'forwardrecurrence(3, A, B, C, 0.1) ≈ clenshaw([1,2,3], A, B, C, 0.1) ≈
            1 + (2sin(2acos(0.1)) + 3sin(3acos(0.1)))/sqrt(1-0.1^2)

        @testset "matrix coefficients" begin
            c = [1 2; 3 4; 5 6]
            @test clenshaw(c,A,B,C,0.1) ≈ [clenshaw(c[:,1],A,B,C,0.1), clenshaw(c[:,2],A,B,C,0.1)]
            @test clenshaw(c,A,B,C,[0.1,0.2]) ≈ [clenshaw(c[:,1], A,B,C,0.1) clenshaw(c[:,2], A,B,C,0.1); clenshaw(c[:,1], A,B,C,0.2) clenshaw(c[:,2], A,B,C,0.2)]
        end
    end

    @testset "Chebyshev-as-general" begin
        @testset "forwardrecurrence!" begin
            N = 5
            A, B, C = [1; fill(2,N-2)], fill(0,N-1), fill(1,N)
            Af, Bf, Cf = float(A), float(B), float(C)
            @test forwardrecurrence(N, A, B, C, 1) == forwardrecurrence!(Vector{Int}(undef,N), A, B, C, 1) == ones(Int,N)
            @test forwardrecurrence!(Vector{Int}(undef,N), A, B, C, -1) == (-1) .^ (0:N-1)
            @test forwardrecurrence(N, A, B, C, 0.1) ≈ forwardrecurrence!(Vector{Float64}(undef,N), A, B, C, 0.1) ≈ cos.((0:N-1) .* acos(0.1))
        end

        c, A, B, C = [1,2,3], [1,2,2], fill(0,3), fill(1,4)
        cf, Af, Bf, Cf = float(c), float(A), float(B), float(C)
        @test @inferred(clenshaw(c, A, B, C, 1)) ≡ 6
        @test @inferred(clenshaw(c, A, B, C, 0.1)) ≡ -1.74
        @test @inferred(clenshaw([1,2,3], A, B, C, [-1,0,1])) == clenshaw!([1,2,3],A, B, C, [-1,0,1]) == [2,-2,6]
        @test clenshaw(c, A, B, C, [-1,0,1]) isa Vector{Int}
        @test @inferred(clenshaw(Float64[], A, B, C, 1)) ≡ 0.0

        x = [1,0,0.1]
        @test @inferred(clenshaw(c, A, B, C, x)) ≈ @inferred(clenshaw!(c, A, B, C, copy(x))) ≈
            @inferred(clenshaw!(c, A, B, C, x, one.(x), similar(x))) ≈
            @inferred(clenshaw!(cf, Af, Bf, Cf, x, one.(x),similar(x))) ≈
            @inferred(clenshaw([1.,2,3], A, B, C, x)) ≈
            @inferred(clenshaw!([1.,2,3], A, B, C, copy(x))) ≈ [6,-2,-1.74]
    end

    @testset "Legendre" begin
        @testset "Float64" begin
            N = 5
            n = 0:N-1
            A = (2n .+ 1) ./ (n .+ 1)
            B = zeros(N)
            C = n ./ (n .+ 1)
            v_1 = forwardrecurrence(N, A, B, C, 1)
            v_f = forwardrecurrence(N, A, B, C, 0.1)
            @test v_1 ≈ ones(N)
            @test forwardrecurrence(N, A, B, C, -1) ≈ (-1) .^ (0:N-1)
            @test v_f ≈ [1,0.1,-0.485,-0.1475,0.3379375]

            n = 0:N # need extra entry for C in Clenshaw
            C = n ./ (n .+ 1)
            for j = 1:N
                c = [zeros(j-1); 1]
                @test clenshaw(c, A, B, C, 1) ≈ v_1[j] # Julia code
                @test clenshaw(c, A, B, C, 0.1) ≈  v_f[j] # Julia code
                @test clenshaw!(c, A, B, C, [1.0,0.1], [1.0,1.0], [0.0,0.0])  ≈ [v_1[j],v_f[j]] # libfasttransforms
            end
        end

        @testset "BigFloat" begin
            N = 5
            n = BigFloat(0):N-1
            A = (2n .+ 1) ./ (n .+ 1)
            B = zeros(N)
            C = n ./ (n .+ 1)
            @test forwardrecurrence(N, A, B, C, parse(BigFloat,"0.1")) ≈ [1,big"0.1",big"-0.485",big"-0.1475",big"0.3379375"]
        end
    end

    @testset "Int" begin
        N = 10; A = 1:10; B = 2:11; C = range(3; step=2, length=N+1)
        v_i = forwardrecurrence(N, A, B, C, 1)
        v_f = forwardrecurrence(N, A, B, C, 0.1)
        @test v_i isa Vector{Int}
        @test v_f isa Vector{Float64}

        j = 3
        @test clenshaw([zeros(Int,j-1); 1; zeros(Int,N-j)], A, B, C, 1) == v_i[j]
    end

    @testset "Zeros diagonal" begin
        N = 10; A = randn(N); B = Zeros{Int}(N); C = randn(N+1)
        @test forwardrecurrence(N, A, B, C, 0.1) == forwardrecurrence(N, A, Vector(B), C, 0.1)
        c = randn(N)
        @test clenshaw(c, A, B, C, 0.1) == clenshaw(c, A, Vector(B), C, 0.1)
    end
end


module TestParametricBootstrap

using Bootstrap
using Test

using DataFrames
using Distributions
using GLM
using RDatasets
using StatsBase

const city = dataset("boot", "city")
const aircondit = convert(Vector{Float64}, dataset("boot", "aircondit")[!, :Hours])

@testset "Parametric bootstraps" begin

    function test_bootsample(bs, ref, raw_data, n)

        show(IOBuffer(), bs)
        @test typeof(bs) <: ParametricBootstrapSample
        t0 = original(bs)
        @test length(t0) == length(ref)
        [@test t ≈ r for (t, r) in zip(t0, ref)]

        t1 = straps(bs)
        @test length(t1) == length(t0)
        [@test length(t) == n for t in t1]
        [@test (minimum(t) <= tr && maximum(t) >= tr) for (t, tr) in zip(t1, t0)]
        [@test eltype(t) == eltype(tr) for (t, tr) in zip(t1, t0)]

        @test Bootstrap.data(bs) == raw_data

        @test nrun(sampling(bs)) == n
        @test nrun(bs) == n

        @test length(bias(bs)) == length(ref)
        [@test (b[1] > -Inf && b[1] < Inf) for b in bias(bs)]
        @test length(stderror(bs)) == length(ref)
        [@test s >= 0 for s in stderror(bs)]

        [@test original(bs, i) == original(bs)[i]  for i in 1:nvar(bs)]
        [@test straps(bs, i) == straps(bs)[i]  for i in 1:nvar(bs)]
        [@test bias(bs, i) == bias(bs)[i]  for i in 1:nvar(bs)]
        [@test stderror(bs, i) == stderror(bs)[i]  for i in 1:nvar(bs)]

        m = model(bs)
        @test typeof(m) <: Bootstrap.Model

        return Nothing
    end

    function test_confint(bs)

        cim_all = (BasicConfInt(), PercentileConfInt(), NormalConfInt(), BCaConfInt())
        for cim in cim_all
            c = confint(bs, cim)
            [@test (x[1] >= x[2]  && x[1] <= x[3]) for x in c]
            [@test x[1] ≈ t0 for (x, t0) in zip(c, original(bs))]
        end

        return Nothing
    end

    n = 100

    @testset "Parametric distributions" begin

        r = randn(20)
        ref = mean(r)

        @testset "Basic resampling: Normal distribution" begin
            bs = bootstrap(mean, r, Model(Normal), BasicSampling(n))
            test_bootsample(bs, ref, r, n)
            test_confint(bs)
        end

        @testset "Balanced resampling: Normal distribution" begin
            bs = bootstrap(mean, r, Model(Normal), BalancedSampling(n))
            test_bootsample(bs, ref, r, n)
            test_confint(bs)
        end

        ref = mean(fit(Exponential, aircondit))

        @testset "Basic resampling: Exponential distribution" begin
            bs = bootstrap(mean, aircondit, Model(Exponential), BasicSampling(n))
            test_bootsample(bs, ref, aircondit, n)
            test_confint(bs)
        end

        @testset "Balanced resampling: Exponential distribution" begin
            bs = bootstrap(mean, aircondit, Model(Exponential), BalancedSampling(n))
            test_bootsample(bs, ref, aircondit, n)
            test_confint(bs)
        end

    end

    @testset "Linear regression models" begin

        ref = coef(fit(LinearModel, @formula(log10(X) ~ log10(U)), city))

        @testset "Residual resampling" begin
            bs = bootstrap(coef, city, Model(LinearModel, @formula(log10(X) ~ log10(U))), ResidualSampling(n))
            test_bootsample(bs, ref, city, n)
        end

        @testset "Wild resampling: Rademacher" begin
            bs = bootstrap(coef, city, Model(LinearModel, @formula(log10(X) ~ log10(U))), WildSampling(n, rademacher))
            test_bootsample(bs, ref, city, n)
            @test isa(noise(sampling(bs)), Function)
        end

        @testset "Wild resampling: Mammen" begin
            bs = bootstrap(coef, city, Model(LinearModel, @formula(log10(X) ~ log10(U))), WildSampling(n, mammen))
            test_bootsample(bs, ref, city, n)
            @test isa(noise(sampling(bs)), Function)
        end

    end

    @testset "Generalized linear regression models" begin

        ref = coef(fit(GeneralizedLinearModel, @formula(log10(X) ~ log10(U)), city, Normal()))
        maxiter = 200
        atol = 1e-3

        @testset "Residual resampling" begin
            bs = bootstrap(coef, city,
                           Model(GeneralizedLinearModel, @formula(log10(X) ~ log10(U)), Normal(), maxiter = maxiter, atol = atol),
                           ResidualSampling(n))
            test_bootsample(bs, ref, city, n)
        end

        @testset "Residual resampling with link function" begin
            bs = bootstrap(coef, city,
                           Model(GeneralizedLinearModel, @formula(log10(X) ~ log10(U)), Normal(), IdentityLink(), maxiter = maxiter, atol = atol),
                           ResidualSampling(n))
            test_bootsample(bs, ref, city, n)
        end

        @testset "Wild resampling: Rademacher" begin
            bs = bootstrap(coef, city,
                           Model(GeneralizedLinearModel, @formula(log10(X) ~ log10(U)), Normal(), maxiter = maxiter, atol = atol),
                           WildSampling(n, rademacher))
            test_bootsample(bs, ref, city, n)
            @test isa(noise(sampling(bs)), Function)
        end

        @testset "Wild resampling with link function: Mammen" begin
            bs = bootstrap(coef, city,
                           Model(GeneralizedLinearModel, @formula(log10(X) ~ log10(U)), Normal(), IdentityLink(), maxiter = maxiter, atol = atol),
                           WildSampling(n, mammen))
            test_bootsample(bs, ref, city, n)
            @test isa(noise(sampling(bs)), Function)
        end

    end

end

end


@testset "modifiers.jl" begin
    @testset "reparametrized" begin
        # Interval by x->x²
        interval = Cube(1)
        ri = reparametrized(interval, x->x^2)
        @test @capture_out(print(ri)) == @capture_err(print(stderr, ri))
        @test ri(0.0) == SVector(0.00)
        @test ri(0.5) == SVector(0.25)
        @test ri(1.0) == SVector(1.00)
        r = rand()
        @test offsetframe(ri, r)[1] ≈ offsetframe(interval, r^2)[1]
        @test offsetframe(ri, r)[2] ≈ offsetframe(interval, r^2)[2]

        # Sphere by 2
        circle = Sphere(1)
        rc = reparametrized(circle, 2)
        @test @capture_out(print(rc)) == @capture_err(print(stderr, rc))
        @test rc(0.0) == rc(0.5) == rc(1.0)
        r = rand()
        @test offsetframe(rc, r)[1] ≈ offsetframe(circle, 2r)[1]
        @test offsetframe(rc, r)[2] ≈ offsetframe(circle, 2r)[2]
        @test norm(rand(rc)) ≈ 1

        # Identity reprametrization does not change the space.
        for i in 1:5
            @test reparametrized(Sphere(i), (args...) -> args) ≅ Sphere(i)
        end
    end

    @testset "scaled" begin
        #
        interval = Cube(1)
        si = scaled(interval, 2)
        @test @capture_out(print(si)) == @capture_err(print(stderr, si))
        @test si(0.0) == SVector(0.0)
        @test si(0.5) == SVector(1.0)
        @test si(1.0) == SVector(2.0)
        @test 2interval == si
        r = rand()
        @test offsetframe(si, r) == offsetframe(interval, 2r)
        @test scaled(interval, 2) == 2interval == interval * 2

        # Radius of sphere should match scaling.
        for i in 1:5
            r = rand()
            sc = r * Sphere(i)
            @test norm(rand(sc)) ≈ r
        end

        # Mulitplying by 1 doesn't change the space.
        for i in 1:5
            @test 1 * Sphere(i) ≅ Sphere(i)
        end
    end

    @testset "translated" begin
    end

    @testset "flattened" begin
        fs = flattened(Sphere(1))
        @test @capture_out(print(fs)) == @capture_err(print(stderr, fs))
        @test offsetframe(fs, 0.0)[1] ≈ SVector(1.0, 0.0)
        @test offsetframe(fs, 0.0)[2] ≈ reshape([0.0, 0.0, 1.0], (3, 1))
        @test fs ≅ Sphere(1)

        # Noop for spaces higher dimensions.
        for i in 2:5
            fi = flattened(Sphere(i))
            r = rand(i)
            @test offsetframe(fi, r...)[1] ≈ offsetframe(Sphere(i), r...)[1]
            @test offsetframe(fi, r...)[2] ≈ offsetframe(Sphere(i), r...)[2]
            @test fi ≅ Sphere(i)
        end
    end
end


using AnalyticalEOR
using Test

@testset "AnalyticalEOR.jl" begin

    # Test relative permeability functions
    @test water_rel_perm(0.5, 0.2, 0.2, 0.3, 3.0) ≈ 0.0374999999
    @test water_rel_perm([0.5, 0.5], 0.2, 0.2, 0.3, 3.0) ≈ Vector([0.0375, 0.0375])

    @test oil_rel_perm(0.5, 0.2, 0.2, 0.3, 3.0) ≈ 0.0375
    @test oil_rel_perm([0.5, 0.5], 0.2, 0.2, 0.3, 3.0) ≈ Vector([0.0375, 0.0375])

    @test fractional_flow([0.5, 0.6], 0.2, 0.2, 0.3, 1.0, 3.0, 3.0, 1.0, 1.0) ≈ Vector([ 0.2307692307692305,  0.7058823529411761])
end

@testset "Derivative Tests" begin

    @test krw_derivative([0.5, 0.5], 0.2, 0.2, 0.3, 3.0) ≈ [0.375, 0.375]
    @test kro_derivative([0.5, 0.5], 0.2, 0.2, 0.3, 3.0) ≈ [-0.375, -0.375]

    # test derivative of fractional flow
end


function Burgers(x,t;α::Real=1.0,β::Real=0.0)::Float64
    y=mod(x-β*t,2π)
    if y≤π
        flag=2
    else
        flag=1
        y=2π-y
    end
    N=100
    y0=range(0,π,length=N)
    y1=similar(y0)
    @.y1=y0+α*sin(y0)*t
    ℓ= α≥0 ? findlast(y1.≤y) : findfirst(y1.≥y)
    v=α*sin(y0[ℓ])
    err=abs(v-α*sin(y-t*v))
    while err > 1e-13
        v=v-(v-α*sin(y-t*v))/(1+α*t*cos(y-t*v))
        err=abs(v-α*sin(y-t*v))
    end
    u=(-1)^flag*v+β
end
function Burgers(x::AbstractArray,t;α::Real=1.0,β::Real=0.0)
    N=100
    y0=range(0,π,length=N)
    y1=y0+α*sin.(y0)*t
    y=mod.(x.-β*t,2π)
    u=similar(y)
    @inbounds Threads.@threads for k=1:length(u)
        if y[k]≤π
            flag=2
        else
            flag=1
            y[k]=2π-y[k]
        end
        ℓ= α≥0 ? findlast(y1.≤y[k]) : findfirst(y1.≥y[k])
        v=α*sin(y0[ℓ])
        err=abs(v-α*sin(y[k]-t*v))
        while err > 1e-13
            v=v-(v-α*sin(y[k]-t*v))/(1+α*t*cos(y[k]-t*v))
            err=abs(v-α*sin(y[k]-t*v))
        end
        u[k]=(-1)^flag*v+β
    end
    u
end

function Burgers(x,y,t;α::Real=1.0,β::Real=0.0)
    Burgers(x+y,2t;α=α,β=β)
end

function errLp(u::Function,mesh::Grid2DRect,U::AbstractMatrix,
    t,quad1D_num::Integer=10;p=2)::Float64
    P_dof=size(U,1)
    P_order=isqrt(P_dof)-1
    quad1D_p=range(-1,1,length=quad1D_num)
    quad_num=quad1D_num^2
    quad_p=Array{Float64,2}(undef,quad_num,2)
    quad_w=1/quad_num
    for ℓ₁=1:quad1D_num, ℓ₂=1:quad1D_num
        ℓ=(ℓ₂-1)*quad1D_num+ℓ₁
        quad_p[ℓ,1]=quad1D_p[ℓ₁]
        quad_p[ℓ,2]=quad1D_p[ℓ₂]
    end
    basisP = Array{Float64,2}(undef,P_dof,quad_num)
    for i=0:P_order, j=0:P_order, ℓ=1:quad_num
        basisP[i*(P_order+1)+j+1,ℓ]=JacobiP(i,0,0,quad_p[ℓ,1])*JacobiP(j,0,0,quad_p[ℓ,2])
    end
    quad_p = ( quad_p .+1 ) / 2 # uniform_p now is in [0,1].
    err = 0.0
    uₕ=basisP'*U
    if p<Inf
        @inbounds for j=1:mesh.NumofE
            x₁=mesh.Coordinates[mesh.Elements[1,j],1]
            x₂=mesh.Coordinates[mesh.Elements[2,j],1]
            y₁=mesh.Coordinates[mesh.Elements[1,j],2]
            y₂=mesh.Coordinates[mesh.Elements[4,j],2]
            for ℓ=1:quad_num
                x_ℓj=(1-quad_p[ℓ,1])*x₁+quad_p[ℓ,1]*x₂
                y_ℓj=(1-quad_p[ℓ,2])*y₁+quad_p[ℓ,2]*y₂
                u_ℓj=u(x_ℓj,y_ℓj,t)
                err+=mesh.Tarea[j]*quad_w*abs(u_ℓj-uₕ[ℓ,j])^p
            end
        end
        err=err^(1/p)
    else
        @inbounds for j=1:mesh.NumofE
            x₁=mesh.Coordinates[mesh.Elements[1,j],1]
            x₂=mesh.Coordinates[mesh.Elements[2,j],1]
            y₁=mesh.Coordinates[mesh.Elements[1,j],2]
            y₂=mesh.Coordinates[mesh.Elements[4,j],2]
            for ℓ=1:quad_num
                x_ℓj=(1-quad_p[ℓ,1])*x₁+quad_p[ℓ,1]*x₂
                y_ℓj=(1-quad_p[ℓ,2])*y₁+quad_p[ℓ,2]*y₂
                u_ℓj=u(x_ℓj,y_ℓj,t)
                Δu_ℓj=abs(u_ℓj-uₕ[ℓ,j])
                err=max(err,Δu_ℓj)
            end
        end
    end
    err
end


function Obj_crit_mix(model::EoSModel,F,z,V,T)
    f(x) = eos(model,V,T,x)
    H(x) = ForwardDiff.hessian(f,x)/8.134/T
    L(x) = det(H(x))
    dL(x) = ForwardDiff.gradient(L,x)
    M(x) = [H(x)[1:end-1,:];transpose(dL(x))]
    F[1] = L(z)
    F[2] = det(M(z))
    return F
end

function crit_mix(model::EoSModel,z;v0=nothing)
    if v0 === nothing
        v0 = x0_crit_mix(model,z)
    end
    f! = (F,x) -> Obj_crit_mix(model, F, z, exp10(x[1]), x[2])
    r  = Solvers.nlsolve(f!,v0[1:end],LineSearch(Newton()))
    sol = Solvers.x_sol(r)
    T_c = sol[2]
    V_c = exp10(sol[1])
    p_c = pressure(model, V_c, T_c, z)
    return (T_c, p_c, V_c)
end

function x0_crit_mix(model::EoSModel,z)
    pure = split_model(model)
    crit = crit_pure.(pure)
    n_c  = length(pure)
    V_c  = sum(z[i]*crit[i][3] for i in 1:n_c)
    T_c  = prod(crit[i][1]^z[i] for i ∈ 1:n_c)
    return [log10(V_c),T_c]
end


function makescatteringnoisetrans(M,ω=1.,basis="Hermite",Nk=M)
#Noise transform: 2-field quadrature to go from basis to k-space
kn,wkn,Tkn = nfieldtrans(Nk,3,ω=1/ω,basis=basis)
Tkn = eigmat(M,kn,ω=1/ω,basis=basis)*diagm((-im).^(0:M-1))

#Noise: position space roots and weights for 3-field quadrature
xn,wxn,Txn = nfieldtrans(M,3,ω=ω,basis=basis)
Txkn = conj(Tkn'*Txn)
return xn,wxn,Txn,kn,wkn,Tkn
end


"""
(1-exp(-Λ*x)/x
"""
function kernel(ω::Float)
    if ω < 1e-5
        return 1 - ω / 2 + ω^2 / 6 - ω^3 / 24 + ω^4 / 120 - ω^5 / 720
    else
        # return (1 - exp(-ω)) / ω
        return -expm1(-ω) / ω
    end
end

##################### Particle-hole symmetric kernel #############################

"""
particle-hole symmetric kernel: K(ω, τ)=e^{-ω*τ}+e^{-ω*(β-τ)}

KK=int_0^{1/2} dτ K(ω1,τ)*K(ω2,τ)=(1-e^{ω1+ω2})/(ω1+ω2)+(e^{-ω2}-e^{-ω1})/(ω1-ω2)
"""
function projPH_ω(Λ::Float, ω1::Float, ω2::Float)
        if ω1 > ω2
        return kernel(ω1 + ω2) + exp(-ω2) * kernel(ω1 - ω2)
    else
        return kernel(ω1 + ω2) + exp(-ω1) * kernel(ω2 - ω1)
    end
end

"""
particle-hole symmetric kernel: K(ω, τ)=e^{-ω*τ}+e^{-ω*(β-τ)}

KK=int_0^{Λ} dτ K(ω,t1)*K(ω2,t2)=(1-e^{t1+t2})/(t1+t2)+(1-e^{2β-t1-t2})/(2β-t1-t2)+(1-e^{β+t1-t2})/(β+t1-t2)+(1-e^{β-t1+t2})/(β-t1+t2)
"""
function projPH_τ(Λ::Float, t1::Float, t2::Float)
    # note that here Λ = \beta/2
    return kernel(t1 + t2) + kernel(4 * Λ - t1 - t2) + kernel(2 * Λ - t1 + t2) + kernel(2 * Λ + t1 - t2)
end

##################### Particle-hole asymmetric kernel #############################
"""
particle=hole asymmetric kernel: K(ω, τ)=e^{-ω*τ}-e^{-ω*(β-τ)}

KK=int_0^{1/2} dτ K(ω1,τ)*K(ω2,τ)=(1-e^{ω1+ω2})/(ω1+ω2)-(e^{-ω2}-e^{-ω1})/(ω1-ω2)
"""
function projPHA_ω(Λ::Float, ω1::Float, ω2::Float)
    if ω1 > ω2
    return kernel(ω1 + ω2) - exp(-ω2) * kernel(ω1 - ω2)
    else
    return kernel(ω1 + ω2) - exp(-ω1) * kernel(ω2 - ω1)
    end
end

"""
particle-hole asymmetric kernel: K(ω, τ)=e^{-ω*τ}-e^{-ω*(β-τ)}

KK=int_0^{Λ} dτ K(ω,t1)*K(ω2,t2)=(1-e^{t1+t2})/(t1+t2)+(1-e^{2β-t1-t2})/(2β-t1-t2)-(1-e^{β+t1-t2})/(β+t1-t2)-(1-e^{β-t1+t2})/(β-t1+t2)
"""
function projPHA_τ(Λ::Float, t1::Float, t2::Float)
    return kernel(t1 + t2) + kernel(4 * Λ - t1 - t2) - kernel(2 * Λ - t1 + t2) - kernel(2 * Λ + t1 - t2)
end


module TestBasic

using Test
using TreeParzen

# Quadratic TPE
objective(params) = (params[:x] - 3)^2
best = fmin(objective, Dict(:x => HP.Uniform(:x, -5.0, 5.0)), 50)
@test abs(best[:x] - 3.0) < .25

# Generate trials to calculate
objective(params) = params[:x]^2 + params[:y]^2
space = Dict(
    :x => HP.Uniform(:x, -10.0, 10.0),
    :y => HP.Uniform(:y, -10.0, 10.0),
)
points = [Dict(:x => 0.0, :y => 0.0), Dict(:x => 1.0, :y => 1.0)]
best = fmin(objective, space, 10, points)
@test best == Dict(:x => 0.0, :y => 0.0)

end
true


function nlp_007_010(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    # Test Goals:
    # - infeasible model

    model = Model(optimizer)

    @variable(model, x)
    @variable(model, y)

    @NLconstraint(model, y == exp(x))
    @constraint(model, x == y^2)

    optimize!(model)

    return check_status(
        model,
        INFEASIBLE_PROBLEM,
        termination_target,
        primal_target,
    )
end


using Distributions, ForwardDiff
using Random: seed!
seed!(1)
Y = rand(Normal(4.0, 1.0), 100)
epochs = 200
cost(Y, μ) = loglikelihood(Normal(μ, 1.0), Y)

# adagrad
θ_adagrad = rand(Normal(), 1)
opt = Adagrad(η=1.0)

for i in 1:epochs
    g = ForwardDiff.gradient(μ -> cost(Y, μ[1]), θ_adagrad)

    δ = update(opt, g)
    θ_adagrad .+= δ
end

# rmsprop
θ_rmsprop = rand(Normal(), 1)
opt = RMSprop(η=0.1)

for i in 1:(epochs*2)
    g = ForwardDiff.gradient(μ -> cost(Y, μ[1]), θ_rmsprop)

    δ = update(opt, g)
    θ_rmsprop .+= δ
end

# adam
θ_adam = rand(Normal(), 1)
opt = Adam(α=1.0)

for i in 1:epochs
    g = ForwardDiff.gradient(μ -> cost(Y, μ[1]), θ_adam)

    δ = update(opt, g)
    θ_adam .+= δ
end

# adamax
θ_adamax = rand(Normal(), 1)
opt = Adamax(α=1.0)

for i in 1:epochs
    g = ForwardDiff.gradient(μ -> cost(Y, μ[1]), θ_adamax)

    δ = update(opt, g)
    θ_adamax .+= δ
end

## nadam
θ_nadam = rand(Normal(), 1)
opt = Nadam(η=1.0)

for i in 1:epochs
    g = ForwardDiff.gradient(μ -> cost(Y, μ[1]), θ_nadam)

    δ = update(opt, g)
    θ_nadam .+= δ
end

## vanilla
θ_vanilla = rand(Normal(), 1)
opt = VanillaGradDescent(η=0.001)

for i in 1:epochs
    g = ForwardDiff.gradient(μ -> cost(Y, μ[1]), θ_vanilla)

    δ = update(opt, g)
    θ_vanilla .+= δ
end

## inversedecay
θ_invdec = rand(Normal(), 1)
opt = Inversedecay(κ=0.9)

for i in 1:epochs
    g = ForwardDiff.gradient(μ -> cost(Y, μ[1]), θ_invdec)

    δ = update(opt, g)
    θ_invdec .+= δ
end

@testset "Normal MLE" begin
    @test mean(Y) ≈ θ_adagrad[1] atol=1e-3
    @test mean(Y) ≈ θ_rmsprop[1] atol=1e-1
    @test mean(Y) ≈ θ_adam[1] atol=1e-3
    @test mean(Y) ≈ θ_adamax[1] atol=1e-3
    @test mean(Y) ≈ θ_nadam[1] atol=1e-3
    @test mean(Y) ≈ θ_vanilla[1] atol=1e-3
    @test mean(Y) ≈ θ_invdec[1] atol=1e-3
end