# Functions for Parametric Bipartite Correlation Clustering
using JuMP
using Gurobi
using SparseArrays
using LinearAlgebra

gurobi_env = Gurobi.Env()

"""
Returns the linear programming relaxation of the Parametric Bipartite
Correlation Clustering objective on a bipartite graph H with parameters
mu and beta.
"""
function PBCC_LP(H::SparseMatrixCSC,mu::Float64,beta::Float64,outputflag::Bool=true,optimalflag::Bool=false)

    m = Model(with_optimizer(Gurobi.Optimizer,OutputFlag=outputflag, gurobi_env))

    l,r = size(H)           # number of nodes on each side
    n = l+r                 # total number of nodes in graph
    L = collect(1:l)        # Left indices
    R = collect(l+1:(n))    # Right indices, R[1] is the global index of 1st node on right side

    Hbeta = Matrix(H)-beta*ones(l,r)
    A = [-mu*ones(l,l) Hbeta; Matrix(Hbeta') -mu*ones(r,r)]

    Al = H*H'       # Gives neighbors of left side: Al[i,j] = 1 if nodes i and j, on same side, share a neighbor on the other side
    Ar = H'*H       # Gives neighbors or right side

    if optimalflag
        @variable(m, x[1:n,1:n],Bin)
    else
        @variable(m, x[1:n,1:n])
        @constraint(m,x .<= ones(n,n))
        @constraint(m,x .>= zeros(n,n))
    end

    @objective(m, Min, sum((A[i,j])*x[i,j] for i=1:n-1 for j = i+1:n))

    # println("Symmetry constraints")
    for i = 1:n
        for j = 1:n
            @constraint(m, x[i,j] == x[j,i])
        end
    end
    # println("First solve")
    # One run of method
    JuMP.optimize!(m)

    while true
         D = JuMP.value.(x)

         # Store violating tuples in a vector
         violations = Vector{Tuple{Int,Int,Int}}()
         find_violations!(D,violations)

         # Iterate through and add constraints
         numvi = size(violations,1)
         # println("\t\t\t\t\tViolations = $numvi")
         for v in violations
             @constraint(m,x[v[1],v[2]] - x[min(v[1],v[3]),max(v[1],v[3])] - x[min(v[2],v[3]),max(v[2],v[3])] <= 0)
         end

         if numvi == 0
             break
         end
         JuMP.optimize!(m)
     end

     ## Output the answer
     X = JuMP.value.(x)
     LPval= JuMP.objective_value(m)

     LPbound = 0.0
     for i = 1:n
         for j = i+1:n
            aij = A[i,j]
            if aij < 0
                LPbound += aij*(X[j,i]-1)
            else
                LPbound += aij*X[j,i]
            end
         end
     end

    return X, LPbound
end

## Round vector X. There are far more efficient ways, but this won't
# be the bottleneck in the code anyways.
function RoundX(X::Matrix{Float64}, delta::Float64)
    Clusters = Vector{Vector{Int64}}()
    n = size(X,1)
    clustered = 0
    Unclustered = collect(1:n)
    label = 1
    Labels = zeros(Int64,n)
    while clustered < n
        pivot = rand(Unclustered)
        neighbs = findall(x->x<delta, X[:,pivot])
        cluster = intersect(neighbs,Unclustered)
        Unclustered = setdiff(Unclustered,cluster)
        push!(Clusters,cluster)
        clustered += length(cluster)
        Labels[cluster] .= label
        label += 1
    end

    C = spzeros(n,length(Clusters))
    for j = 1:length(Clusters)
        C[Clusters[j],j] .= 1
    end
    return C, Labels
end

## Round vector X. There are far more efficient ways, but this won't
# be the bottleneck in the code anyways.
function ManyRoundX(H::SparseMatrixCSC{Float64,Int64},X::Matrix{Float64}, delta::Float64,mu::Float64,beta::Float64,k::Int64)

    best = Inf
    bestC = spzeros(1,1)
    bestL = Vector{Int64}()

    for i = 1:k
        C, Labels = RoundX(X,delta)
        obj = PBCC_Objective(H,C,mu,beta)
        if obj < best
            best = obj
            bestC = C
            bestL = Labels
        end
    end

    return bestC, bestL, best
end

## Round vector X with a bunch of delta values and take the best
function RangeRoundX(H::SparseMatrixCSC{Float64,Int64},X::Matrix{Float64},mu::Float64,beta::Float64,k::Int64)

    best = Inf
    bestC = spzeros(1,1)
    bestL = Vector{Int64}()
    bestD = 0

    for delta = 0.05:0.05:0.95
        C, Labels, obj = ManyRoundX(H,X,delta,mu,beta,k)
        if obj < best
            best = obj
            bestC = C
            bestL = Labels
            bestD = delta
        end
    end

    return bestC, bestL, best, bestD
end

# Given a clustering of the matrix A = [0 H, H' 0], output a clustering
# of just H, i.e. a set of two vectors, x and y, so that x gives the cluster
# labels for nodes on one side of bipartite matrix H, and y gives the cluster
# labels for the other side
function  Biclustering(H::SparseMatrixCSC{Float64,Int64},L::Vector{Int64})

    l,r = size(H)
    x = L[1:l]
    y = L[l+1:l+r]

    return x,y
end

# Objective for a clustering, a super inefficient way to do things
function PBCC_Objective(H::SparseMatrixCSC{Float64,Int64},
    C::SparseMatrixCSC{Float64,Int64},mu::Float64,beta::Float64)

    n = size(C,1)

    l,r = size(H)
    Hbeta = Matrix(H)-beta*ones(l,r)
    A = [-mu*ones(l,l) Hbeta; Matrix(Hbeta') -mu*ones(r,r)]

    D = ones(n,n)  # dense, inefficient
    for i = 1:n
        cluster_ind = findall(x->x==1,C[i,:])   # cluster index
        cluster = findall(x->x==1,C[:,cluster_ind])
        D[i,cluster] .= 0.0
    end

    obj = 0.0
    for i = 1:n
        for j = i+1:n
           aij = A[i,j]
           if aij < 0
               obj += aij*(D[j,i]-1)
           else
               obj += aij*D[j,i]
           end
        end
    end
    return obj
end

"""
Returns the linear programming relaxation of the Bicluster Deletion Problem.
"""
function BiclusterDeletion_LP(H::SparseMatrixCSC,outputflag::Bool=true,optimalflag::Bool=false)

    m = Model(with_optimizer(Gurobi.Optimizer,OutputFlag=outputflag, gurobi_env))

    l,r = size(H)           # number of nodes on each side
    n = l+r                 # total number of nodes in graph
    L = collect(1:l)        # Left indices
    R = collect(l+1:(n))    # Right indices, R[1] is the global index of 1st node on right side

    A = [spzeros(l,l) H; sparse(H') spzeros(r,r)]

    # Al = Project(H,1)       # Gives neighbors of left side: Al[i,j] = 1 if nodes i and j, on same side, share a neighbor on the other side
    # Ar = Project(H,2)       # Gives neighbors or right side

    Al = H*H'       # Gives neighbors of left side: Al[i,j] = 1 if nodes i and j, on same side, share a neighbor on the other side
    Ar = H'*H       # Gives neighbors or right side

    if optimalflag
        @variable(m, x[1:n,1:n],Bin)
    else
        @variable(m, x[1:n,1:n])
        @constraint(m,x .<= ones(n,n))
        @constraint(m,x .>= zeros(n,n))
    end

    @objective(m, Min, sum((A[i,j])*x[i,j] for i=1:n-1 for j = i+1:n))

    # Make the matrix symmetric
    println("First constraints")
    for i = 1:n
        for j = 1:n
            @constraint(m, x[i,j] == x[j,i])
        end
    end

    println("Second constraints")
    # Two nodes on opposite sides can't be together if they don't share an edge
    for i = 1:l
        for j = 1:r
            if H[i,j] == 0
                @constraint(m,x[i,R[j]] == 1)
            end
        end
    end

    # println("Third constraints")
    # Two nodes on the same side can't be together if they don't share a neighbor
    for i = 1:l-1
        for j = i+1:l
            if Al[i,j] == 0
                # No shared neighbor means can't be clustered together
                @constraint(m, x[i,j] == 1)
            end
        end
    end

    # println("Fourth constraints")
    # Do the same for the right side, careful to adjust indices using R
    for i = 1:r-1
        for j = i+1:r
            if Ar[i,j] == 0
                # No shared neighbor means can't be clustered together
                @constraint(m, x[R[i],R[j]] == 1)
            end
        end
    end

    # println("First solve")
    # One run of method
    JuMP.optimize!(m)

    while true
         D = JuMP.value.(x)

         # Store violating tuples in a vector
         violations = Vector{Tuple{Int,Int,Int}}()
         find_violations!(D,violations)

         # Iterate through and add constraints
         numvi = size(violations,1)
         # println("\t\t\t\t\tViolations = $numvi")
         for v in violations
             @constraint(m,x[v[1],v[2]] - x[min(v[1],v[3]),max(v[1],v[3])] - x[min(v[2],v[3]),max(v[2],v[3])] <= 0)
         end

         if numvi == 0
             break
         end
         JuMP.optimize!(m)
     end

     ## Output the answer
     X = JuMP.value.(x)
     LPval= JuMP.objective_value(m)

    return X, LPval
end

# find_violations
# Given a candidate distance matrix D, iterate through all 3-tuples of nodes
# and return the tuples where triangle inequality constraints have been violated.
#
# Output is stored in vector 'violations'
#
# Heuristic speed ups:
#       - Only grab Dij, Dik, Djk once per tuple
#       - if a - b - c > 0, then a>b and a>c. By checking these first
#           we rule out a lot of situations where constraints do not need to be
#           be added, so in all test cases this gave a speed up
#
# Note that we want D to be lower triangular here, though if it is symmetric
# this will work fine as well. We just need to make sure the distance information
# in D is not just stored in the upper triangular portion
function find_violations!(D::Matrix{Float64}, violations::Vector{Tuple{Int,Int,Int}},epsi::Float64=1e-8)
  n = size(D,1)

  @inbounds for i = 1:n-2
       for j = i+1:n-1
          a = D[j,i]
           for k = j+1:n
              b = D[k,i]
              c = D[k,j]
        if a - b > epsi && a - c > epsi && a-b-c > epsi
            push!(violations, (i,j,k))
            # @constraint(m, x[i,j] - x[i,k] - x[j,k] <= 0)
        end
        if b - a > epsi && b - c > epsi && b-a-c > epsi
            push!(violations, (i,k,j))
            # @constraint(m, x[i,k] - x[i,j] - x[j,k] <= 0)
        end

        if c - a > epsi && c-b>epsi && c-a-b > epsi
            push!(violations, (j,k,i))
            # @constraint(m, x[j,k] - x[i,k] - x[i,j] <= 0)
        end
      end
    end
  end
end


"""
    $FUNCTIONNAME(pc::PlayerCost, x::SVector{nx}, u::SVector{nu}, t::AbstractFloat)

A convencience implementation of the cost quadraticization.
"""
function _quadraticize_ad(pc::PlayerCost, g::GeneralGame, x::SVector, u::SVector,
                          t::AbstractFloat)
    @warn "You are using the fallback quadraticization using ForwardDiff.
    Consider implementing a custom `quadraticize` for your `ControlSystem`
    type." maxlog=1

    x_cost = x->pc(g, x, u, t)
    u_cost = u->pc(g, x, u, t)
    # we can compute the gradient and the hessian in one go
    diff_x = DiffResults.HessianResult(x)
    diff_x = ForwardDiff.hessian!(diff_x, x_cost, x)
    # the linear state component of the cost is the gradient in x
    l = DiffResults.gradient(diff_x)
    # the quadratic state component of the cost is the hessian in x
    Q = DiffResults.hessian(diff_x)

    diff_u = DiffResults.HessianResult(u)
    diff_u = ForwardDiff.hessian!(diff_u, u_cost, u)
    # the linear control cost is the gradient in u
    r = DiffResults.gradient(diff_u)
    # the quadratic control component o the cost is the hessian in u
    R = DiffResults.hessian(diff_u)

    return QuadraticPlayerCost(l, Q, r, R)
end

quadraticize!(qcache::QuadCache, pc::PlayerCost, g::GeneralGame, x::SVector,
              u::SVector, t::AbstractFloat) = _quadraticize_ad(pc, g, x, u, t)


function f_model(
    f_x::R, grad_f_x::AbstractArray{C}, res::AbstractArray{C}, gamma::R
) where {R <: Real, C <: RealOrComplex{R}}
    return f_x - real(dot(grad_f_x, res)) + (0.5/gamma)*norm(res)^2
end


"""
    ArrayIndex{N}

Subtypes of `ArrayIndex` represent series of transformations for a provided index to some
buffer which is typically accomplished with square brackets (e.g., `buffer[index[inds...]]`).
The only behavior that is required of a subtype of `ArrayIndex` is the ability to transform
individual index elements (i.e. not collections). This does not guarantee bounds checking or
the ability to iterate (although additional functionallity may be provided for specific
types).
"""
abstract type ArrayIndex{N} end

const MatrixIndex = ArrayIndex{2}

const VectorIndex = ArrayIndex{1}

Base.ndims(::Type{<:ArrayIndex{N}}) where {N} = N

struct BidiagonalIndex <: MatrixIndex
    count::Int
    isup::Bool
end

struct TridiagonalIndex <: MatrixIndex
    count::Int# count==nsize+nsize-1+nsize-1
    nsize::Int
    isrow::Bool
end

struct BandedMatrixIndex <: MatrixIndex
    count::Int
    rowsize::Int
    colsize::Int
    bandinds::Array{Int,1}
    bandsizes::Array{Int,1}
    isrow::Bool
end

function _bandsize(bandind, rowsize, colsize)
    -(rowsize - 1) <= bandind <= colsize - 1 || throw(ErrorException("Invalid Bandind"))
    if (bandind * (colsize - rowsize) > 0) & (abs(bandind) <= abs(colsize - rowsize))
        return min(rowsize, colsize)
    elseif bandind * (colsize - rowsize) <= 0
        return min(rowsize, colsize) - abs(bandind)
    else
        return min(rowsize, colsize) - abs(bandind) + abs(colsize - rowsize)
    end
end

function BandedMatrixIndex(rowsize, colsize, lowerbandwidth, upperbandwidth, isrow)
    upperbandwidth > -lowerbandwidth || throw(ErrorException("Invalid Bandwidths"))
    bandinds = upperbandwidth:-1:-lowerbandwidth
    bandsizes = [_bandsize(band, rowsize, colsize) for band in bandinds]
    BandedMatrixIndex(sum(bandsizes), rowsize, colsize, bandinds, bandsizes, isrow)
end

struct BlockBandedMatrixIndex <: MatrixIndex
    count::Int
    refinds::Array{Int,1}
    refcoords::Array{Int,1}# storing col or row inds at ref points
    isrow::Bool
end

function BlockBandedMatrixIndex(nrowblock, ncolblock, rowsizes, colsizes, l, u)
    blockrowind = BandedMatrixIndex(nrowblock, ncolblock, l, u, true)
    blockcolind = BandedMatrixIndex(nrowblock, ncolblock, l, u, false)
    sortedinds = sort(
        [(blockrowind[i], blockcolind[i]) for i = 1:length(blockrowind)],
        by = x -> x[1],
    )
    sort!(sortedinds, by = x -> x[2], alg = InsertionSort)# stable sort keeps the second index in order
    refinds = Array{Int,1}()
    refrowcoords = Array{Int,1}()
    refcolcoords = Array{Int,1}()
    rowheights = pushfirst!(copy(rowsizes), 1)
    cumsum!(rowheights, rowheights)
    blockheight = 0
    blockrow = 1
    blockcol = 1
    currenti = 1
    lastrowind = sortedinds[1][1] - 1
    lastcolind = sortedinds[1][2]
    for ind in sortedinds
        rowind, colind = ind
        if colind == lastcolind
            if rowind > lastrowind
                blockheight += rowsizes[rowind]
            end
        else
            for j = blockcol:blockcol+colsizes[lastcolind]-1
                push!(refinds, currenti)
                push!(refrowcoords, blockrow)
                push!(refcolcoords, j)
                currenti += blockheight
            end
            blockcol += colsizes[lastcolind]
            blockrow = rowheights[rowind]
            blockheight = rowsizes[rowind]
        end
        lastcolind = colind
        lastrowind = rowind
    end
    for j = blockcol:blockcol+colsizes[lastcolind]-1
        push!(refinds, currenti)
        push!(refrowcoords, blockrow)
        push!(refcolcoords, j)
        currenti += blockheight
    end
    push!(refinds, currenti)# guard
    push!(refrowcoords, -1)
    push!(refcolcoords, -1)
    rowindobj = BlockBandedMatrixIndex(currenti - 1, refinds, refrowcoords, true)
    colindobj = BlockBandedMatrixIndex(currenti - 1, refinds, refcolcoords, false)
    rowindobj, colindobj
end

struct BandedBlockBandedMatrixIndex <: MatrixIndex
    count::Int
    refinds::Array{Int,1}
    refcoords::Array{Int,1}# storing col or row inds at ref points
    reflocalinds::Array{BandedMatrixIndex,1}
    isrow::Bool
end

function BandedBlockBandedMatrixIndex(
    nrowblock,
    ncolblock,
    rowsizes,
    colsizes,
    l,
    u,
    lambda,
    mu,
)
    blockrowind = BandedMatrixIndex(nrowblock, ncolblock, l, u, true)
    blockcolind = BandedMatrixIndex(nrowblock, ncolblock, l, u, false)
    sortedinds = sort(
        [(blockrowind[i], blockcolind[i]) for i = 1:length(blockrowind)],
        by = x -> x[1],
    )
    sort!(sortedinds, by = x -> x[2], alg = InsertionSort)# stable sort keeps the second index in order
    rowheights = pushfirst!(copy(rowsizes), 1)
    cumsum!(rowheights, rowheights)
    colwidths = pushfirst!(copy(colsizes), 1)
    cumsum!(colwidths, colwidths)
    currenti = 1
    refinds = Array{Int,1}()
    refrowcoords = Array{Int,1}()
    refcolcoords = Array{Int,1}()
    reflocalrowinds = Array{BandedMatrixIndex,1}()
    reflocalcolinds = Array{BandedMatrixIndex,1}()
    for ind in sortedinds
        rowind, colind = ind
        localrowind =
            BandedMatrixIndex(rowsizes[rowind], colsizes[colind], lambda, mu, true)
        localcolind =
            BandedMatrixIndex(rowsizes[rowind], colsizes[colind], lambda, mu, false)
        push!(refinds, currenti)
        push!(refrowcoords, rowheights[rowind])
        push!(refcolcoords, colwidths[colind])
        push!(reflocalrowinds, localrowind)
        push!(reflocalcolinds, localcolind)
        currenti += localrowind.count
    end
    push!(refinds, currenti)
    push!(refrowcoords, -1)
    push!(refcolcoords, -1)
    rowindobj = BandedBlockBandedMatrixIndex(
        currenti - 1,
        refinds,
        refrowcoords,
        reflocalrowinds,
        true,
    )
    colindobj = BandedBlockBandedMatrixIndex(
        currenti - 1,
        refinds,
        refcolcoords,
        reflocalcolinds,
        false,
    )
    rowindobj, colindobj
end

struct StrideIndex{N,R,C,S,O,O1} <: ArrayIndex{N}
    strides::S
    offsets::O
    offset1::O1

    function StrideIndex{N,R,C}(s::S, o::O, o1::O1) where {N,R,C,S<:Tuple{Vararg{<:CanonicalInt,N}},O<:Tuple{Vararg{<:CanonicalInt,N}},O1}
        return new{N,R::NTuple{N,Int},C::Int,S,O,O1}(s, o, o1)
    end
    function StrideIndex{N,R,C}(a::A) where {N,R,C,A}
        return StrideIndex{N,R,C}(strides(a),offsets(a),offset1(a))
    end
    function StrideIndex(a::A) where {A}
        return StrideIndex{ndims(A),known(stride_rank(A)), known(contiguous_axis(A))}(a)
    end
end

Base.firstindex(i::Union{TridiagonalIndex,BandedBlockBandedMatrixIndex,BandedMatrixIndex,BidiagonalIndex,BlockBandedMatrixIndex}) = 1
Base.lastindex(i::Union{TridiagonalIndex,BandedBlockBandedMatrixIndex,BandedMatrixIndex,BidiagonalIndex,BlockBandedMatrixIndex}) = i.count
Base.length(i::Union{TridiagonalIndex,BandedBlockBandedMatrixIndex,BandedMatrixIndex,BidiagonalIndex,BlockBandedMatrixIndex}) = i.count
function Base.getindex(ind::BidiagonalIndex, i::Int)
    1 <= i <= ind.count || throw(BoundsError(ind, i))
    if ind.isup
        ii = i + 1
    else
        ii = i + 1 + 1
    end
    convert(Int, floor(ii / 2))
end

function Base.getindex(ind::TridiagonalIndex, i::Int)
    1 <= i <= ind.count || throw(BoundsError(ind, i))
    offsetu = ind.isrow ? 0 : 1
    offsetl = ind.isrow ? 1 : 0
    if 1 <= i <= ind.nsize
        return i
    elseif ind.nsize < i <= ind.nsize + ind.nsize - 1
        return i - ind.nsize + offsetu
    else
        return i - (ind.nsize + ind.nsize - 1) + offsetl
    end
end

function Base.getindex(ind::BandedMatrixIndex, i::Int)
    1 <= i <= ind.count || throw(BoundsError(ind, i))
    _i = i
    p = 1
    while _i - ind.bandsizes[p] > 0
        _i -= ind.bandsizes[p]
        p += 1
    end
    bandind = ind.bandinds[p]
    startfromone = !xor(ind.isrow, (bandind > 0))
    if startfromone
        return _i
    else
        return _i + abs(bandind)
    end
end

function Base.getindex(ind::BlockBandedMatrixIndex, i::Int)
    1 <= i <= ind.count || throw(BoundsError(ind, i))
    p = 1
    while i - ind.refinds[p] >= 0
        p += 1
    end
    p -= 1
    _i = i - ind.refinds[p]
    if ind.isrow
        return ind.refcoords[p] + _i
    else
        return ind.refcoords[p]
    end
end

function Base.getindex(ind::BandedBlockBandedMatrixIndex, i::Int)
    1 <= i <= ind.count || throw(BoundsError(ind, i))
    p = 1
    while i - ind.refinds[p] >= 0
        p += 1
    end
    p -= 1
    _i = i - ind.refinds[p] + 1
    ind.reflocalinds[p][_i] + ind.refcoords[p] - 1
end

@inline function Base.getindex(x::StrideIndex{N}, i::CanonicalInt...) where {N}
    return _strides2int(x.offsets, x.strides, Tuple(i)) + x.offset1
end
@inline function _strides2int(o::Tuple, s::Tuple, i::Tuple)
    return ((first(i) - first(o)) * first(s)) + _strides2int(tail(o), tail(s), tail(i))
end
_strides2int(::Tuple{}, ::Tuple{}, ::Tuple{}) = static(0)


# Tests for module InterestRates

using Test
using BusinessDays
using InterestRates
using Dates

BusinessDays.initcache(BusinessDays.Brazil())

vert_x = [11, 15, 19, 23]
vert_y = [0.10, 0.15, 0.20, 0.19]

dt_curve = Date(2015,08,03)

# Error testing
@test_throws AssertionError InterestRates.IRCurve("", InterestRates.Actual365(), InterestRates.ExponentialCompounding(),
    InterestRates.Linear(), dt_curve, [1, 2, 6, 5], [.1, .1, .1, .1])

@test_throws AssertionError InterestRates.IRCurve("", InterestRates.Actual365(), InterestRates.ExponentialCompounding(),
    InterestRates.Linear(), dt_curve, [1, 2, 3, 4], [.1, .1, .1])

@test_throws AssertionError InterestRates.IRCurve("", InterestRates.Actual365(), InterestRates.ExponentialCompounding(),
    InterestRates.Linear(), dt_curve, Vector{Int}() , [.1, .1, .1])

@test_throws AssertionError InterestRates.IRCurve("", InterestRates.Actual365(), InterestRates.ExponentialCompounding(),
    InterestRates.Linear(), dt_curve, Vector{Int}(undef, 1) , Vector{Float64}())

@test_throws AssertionError InterestRates.IRCurve("", InterestRates.Actual365(), InterestRates.ExponentialCompounding(),
    InterestRates.Linear(), dt_curve, Vector{Int}() , Vector{Float64}())

curve_b252_ec_lin = InterestRates.IRCurve("dummy-linear", InterestRates.BDays252(BusinessDays.Brazil()),
    InterestRates.ExponentialCompounding(), InterestRates.Linear(), dt_curve,
    vert_x, vert_y)

@test_throws AssertionError zero_rate(curve_b252_ec_lin, dt_curve - Dates.Day(10))

@test InterestRates.curve_get_name(curve_b252_ec_lin) == "dummy-linear"
@test InterestRates.curve_get_date(curve_b252_ec_lin) == dt_curve

# daycount equality
let
    x = InterestRates.BDays252(BusinessDays.Brazil())
    y = InterestRates.BDays252(BusinessDays.Brazil())
    @test x == y

    z = InterestRates.Actual360()
    @test x != z
    @test y != z

    @test hash(x) == hash(y)
    @test hash(x) != hash(z)
    @test hash(y) != hash(z)
end

# method equality
let
    @test InterestRates.Linear() == InterestRates.Linear()
    @test InterestRates.Linear() != InterestRates.FlatForward()
    @test InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.FlatForward()) == InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.FlatForward())
    @test InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.FlatForward()) != InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.StepFunction())
    @test InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.FlatForward()) != InterestRates.Linear()

    @test hash(InterestRates.Linear()) == hash(InterestRates.Linear())
    @test hash(InterestRates.Linear()) != hash(InterestRates.FlatForward())
    @test hash(InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.FlatForward())) == hash(InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.FlatForward()))
    @test hash(InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.FlatForward())) != hash(InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.StepFunction()))
    @test hash(InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.FlatForward())) != hash(InterestRates.Linear())
end

maturity_2_days = advancebdays(BusinessDays.Brazil(), dt_curve, vert_x[1] + 2)
yrs = (vert_x[1] + 2) / 252.0
zero_rate_2_days = 0.125
disc_2_days = 1.0 / ( (1.0 + zero_rate_2_days)^yrs)
@test zero_rate_2_days ≈ zero_rate(curve_b252_ec_lin, maturity_2_days) # Linear interpolation
@test disc_2_days ≈ discountfactor(curve_b252_ec_lin, maturity_2_days)
@test zero_rate(curve_b252_ec_lin, advancebdays(BusinessDays.Brazil(), dt_curve, 11)) ≈ 0.10
@test_throws AssertionError zero_rate(curve_b252_ec_lin, advancebdays(BusinessDays.Brazil(), dt_curve, -4)) # maturity before curve date
@test zero_rate(curve_b252_ec_lin, advancebdays(BusinessDays.Brazil(), dt_curve, 11-4)) ≈ 0.05 # extrapolate before first vertice
@test zero_rate(curve_b252_ec_lin, advancebdays(BusinessDays.Brazil(), dt_curve, 23+4)) ≈ 0.18 # extrapolate after last vertice

dt_maturity = dt_curve+Dates.Day(30)
@test ERF_to_rate(curve_b252_ec_lin, ERF(curve_b252_ec_lin, dt_maturity), InterestRates.yearfraction(curve_b252_ec_lin, dt_maturity)) ≈ zero_rate(curve_b252_ec_lin, dt_maturity)

curve_ac360_cont_ff = InterestRates.IRCurve("dummy-cont-flatforward", InterestRates.Actual360(),
    InterestRates.ContinuousCompounding(), InterestRates.FlatForward(), dt_curve,
    vert_x, vert_y)

@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(11)) ≈ 0.1
@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(15)) ≈ 0.15
@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(19)) ≈ 0.20
@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(23)) ≈ 0.19
@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(16)) > 0.15
@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(17)) < 0.20
@test forward_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(11), dt_curve + Dates.Day(15)) ≈ 0.2875 # forward_rate calculation on vertices
@test forward_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(11), dt_curve + Dates.Day(13)) ≈ 0.2875 # forward_rate calculation on interpolated maturity
@test ERF(curve_ac360_cont_ff, dt_curve + Dates.Day(13)) ≈ 1.00466361875533 # ffwd interp on ERF

dt_maturity = dt_curve+Dates.Day(30)
@test ERF_to_rate(curve_ac360_cont_ff, ERF(curve_ac360_cont_ff, dt_maturity), InterestRates.yearfraction(curve_ac360_cont_ff, dt_maturity)) ≈ zero_rate(curve_ac360_cont_ff, dt_maturity)
@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(13)) ≈ 0.128846153846152 # ffwd interp as zero_rate
@test ERF(curve_ac360_cont_ff, dt_curve + Dates.Day(19), dt_curve + Dates.Day(23)) ≈ 1.00158458746737
@test forward_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(19), dt_curve + Dates.Day(23)) ≈ 0.1425000000000040
@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(30)) ≈ 0.1789166666666680 # ffwd extrap after last vertice
@test forward_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(19), dt_curve + Dates.Day(23)) ≈ forward_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(50), dt_curve + Dates.Day(51))
@test forward_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(19), dt_curve + Dates.Day(23)) ≈ forward_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(50), dt_curve + Dates.Day(100))

@test forward_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(11), dt_curve + Dates.Day(15)) ≈ 0.2875
@test zero_rate(curve_ac360_cont_ff, dt_curve + Dates.Day(9)) ≈ 0.05833333333333 # ffwd extrap before first vertice

@test discountfactor(curve_ac360_cont_ff, dt_curve) == 1
@test isnan(ERF_to_rate(curve_ac360_cont_ff, 1.0, 0.0))

# Null curve tests
n_curve = InterestRates.NullIRCurve()
@test n_curve == InterestRates.NullIRCurve() # Should be singleton
@test ERF(n_curve, Date(2000,1,1)) == 1.0
@test ER(n_curve, Date(2000,1,1)) == 0.0
@test discountfactor(n_curve, Date(2000,1,1)) == 1.0
@test InterestRates.curve_get_name(n_curve) == "NullCurve"
@test isnullcurve(n_curve) == true
@test isnullcurve(curve_ac360_cont_ff) == false
@test forward_rate(n_curve, Date(2000,1,1), Date(2000,1,2)) == 0.0
@test zero_rate(n_curve, Date(2000,1,1)) == 0.0
@test zero_rate(n_curve, [Date(2000,1,1), Date(2000,1,2)]) == [ 0.0, 0.0 ]

# Tests for vector functions
dt_curve = Date(2015, 08, 07)
curve_ac365_simple_linear = InterestRates.IRCurve("dummy-simple-linear", InterestRates.Actual365(),
    InterestRates.SimpleCompounding(), InterestRates.Linear(), dt_curve,
    vert_x, vert_y)
mat_vec = [ Date(2015,08,17), Date(2015,08,18), Date(2015,08,19), Date(2015,08,20), Date(2015,08,21), Date(2015,08,22)]
@test zero_rate(curve_ac365_simple_linear, mat_vec) ≈ [0.0875,0.1,0.1125,0.1250,0.1375,0.15]
@test ERF(curve_ac365_simple_linear, mat_vec) ≈ [1.00239726027397, 1.00301369863014, 1.00369863013699, 1.00445205479452, 1.00527397260274, 1.00616438356164]
@test discountfactor(curve_ac365_simple_linear, mat_vec) ≈ [0.997608472839084, 0.996995356459984, 0.996314999317592, 0.995567678145244, 0.994753696259454, 0.993873383253914]

dt_maturity = dt_curve+Dates.Day(30)
@test ERF_to_rate(curve_ac365_simple_linear, ERF(curve_ac365_simple_linear, dt_maturity), InterestRates.yearfraction(curve_ac365_simple_linear, dt_maturity)) ≈ zero_rate(curve_ac365_simple_linear, dt_maturity)

comp = InterestRates.CompositeInterpolation(InterestRates.StepFunction(), InterestRates.Linear(), InterestRates.StepFunction())
curve_ac365_simple_comp = InterestRates.IRCurve("dummy-simple-linear", InterestRates.Actual365(),
    InterestRates.SimpleCompounding(), comp, dt_curve,
    vert_x, vert_y)
@test zero_rate(curve_ac365_simple_comp, mat_vec) ≈ [0.1,0.1,0.1125,0.1250,0.1375,0.15]
@test zero_rate(curve_ac365_simple_comp, Date(2100,2,2)) ≈ 0.19

curve_step = InterestRates.IRCurve("step-curve", InterestRates.Actual365(),
    InterestRates.SimpleCompounding(), InterestRates.StepFunction(), dt_curve,
    vert_x, vert_y)
mat_vec = [ Date(2015,08,08), Date(2015,08,12), Date(2015,08,17), Date(2015,08,18), Date(2015,08,19),
    Date(2015,08,21), Date(2015,08,22), Date(2015,08,23), Date(2015,08,25), Date(2015,08,26),
    Date(2015,08,27), Date(2015,08,29), Date(2015,08,30), Date(2015,08,31), Date(2015,09,26)]

@test zero_rate(curve_step, mat_vec) ≈ [0.10, 0.10, 0.10, 0.10, 0.10, 0.10, 0.15, 0.15, 0.15, 0.20, 0.20, 0.20, 0.19, 0.19, 0.19]

dt_maturity = dt_curve+Dates.Day(30)
@test ERF_to_rate(curve_step, ERF(curve_step, dt_maturity), InterestRates.yearfraction(curve_step, dt_maturity)) ≈ zero_rate(curve_step, dt_maturity)

# Nelson Siegel
dt_curve = Date(2015, 08, 11)
curve_NS = InterestRates.IRCurve("dummy-continuous-nelsonsiegel", InterestRates.Actual360(),
    InterestRates.ContinuousCompounding(), InterestRates.NelsonSiegel(), dt_curve,
    [0.1, 0.2, 0.3, 0.5])

mat_vec = [Date(2015,8,12), Date(2016,8,12)]
@test zero_rate(curve_NS, mat_vec) ≈ [0.300069315921728, 0.311522078457982]
@test discountfactor(curve_NS, mat_vec) ≈ [0.999166821408637, 0.727908844513432]

dt_maturity = dt_curve+Dates.Day(30)
@test ERF_to_rate(curve_NS, ERF(curve_NS, dt_maturity), InterestRates.yearfraction(curve_NS, dt_maturity)) ≈ zero_rate(curve_NS, dt_maturity)

# Svensson
dt_curve = Date(2015, 08, 11)
curve_sven = InterestRates.IRCurve("dummy-continuous-svensson", InterestRates.Actual360(),
    InterestRates.ContinuousCompounding(), InterestRates.Svensson(), dt_curve,
    [0.1, 0.2, 0.3, 0.4, 0.5, 0.8])

mat_vec = [Date(2015,8,12), Date(2016,8,12)]
@test zero_rate(curve_sven, mat_vec) ≈ [0.300513102478340, 0.408050168725566]
@test discountfactor(curve_sven, mat_vec) ≈ [0.999165589696054, 0.659690510410030]

dt_maturity = dt_curve+Dates.Day(30)
@test ERF_to_rate(curve_NS, ERF(curve_NS, dt_maturity), InterestRates.yearfraction(curve_NS, dt_maturity)) ≈ zero_rate(curve_NS, dt_maturity)

# Splines

vert_x = [11, 15, 19, 23, 25]
vert_y = [0.10, 0.12, 0.20, 0.22, 0.2]

sp = InterestRates.splinefit([1, 2, 3, 4], [1.0, 2.0, 3.0, 4.0])
InterestRates.splineint(sp, [5, 6])
sp = InterestRates.splinefit([1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0])
InterestRates.splineint(sp, [1.0, 1.5])

sp = InterestRates.splinefit(vert_x, vert_y)
y = InterestRates.splineint(sp, convert(Vector{Int}, 1:30))
@test y[vert_x] ≈ vert_y

y_benchmark = [0.09756098, 0.09780488, 0.09804878, 0.09829268, 0.09853659, 0.09878049, 0.09902439, 0.09926829, 0.09951220, 0.09975610, 0.10000000, 0.10054116, 0.10286585,
0.10875762, 0.12000000, 0.13753049, 0.15890244, 0.18082317, 0.20000000, 0.21371189, 0.22152439, 0.22357470, 0.22000000, 0.21137195, 0.20000000, 0.18817073,
0.17634146, 0.16451220, 0.15268293, 0.14085366]

@test all(round.(y; digits=8) .== round.(y_benchmark; digits=8))

curve_spline_rates = InterestRates.IRCurve("dummy-SplineOnRates", InterestRates.Actual360(),
    InterestRates.ContinuousCompounding(), InterestRates.CubicSplineOnRates(), dt_curve,
    vert_x, vert_y)

@test zero_rate(curve_spline_rates, dt_curve + Dates.Day(11)) ≈ 0.1
@test zero_rate(curve_spline_rates, [dt_curve+Dates.Day(11), dt_curve+Dates.Day(15)]) ≈ vert_y[1:2]

dt_maturity = dt_curve+Dates.Day(30)
@test ERF_to_rate(curve_spline_rates, ERF(curve_spline_rates, dt_maturity), InterestRates.yearfraction(curve_spline_rates, dt_maturity)) ≈ zero_rate(curve_spline_rates, dt_maturity)

mat_vec = [ Dates.Day(i) for i in 1:30 ] + dt_curve
@test zero_rate(curve_spline_rates, mat_vec) ≈ y

curve_spline_discount = InterestRates.IRCurve("dummy-SplineOnDiscountFactors", InterestRates.Actual360(),
    InterestRates.ContinuousCompounding(), InterestRates.CubicSplineOnDiscountFactors(), dt_curve,
    vert_x, vert_y)

@test zero_rate(curve_spline_discount, dt_curve + Dates.Day(11)) ≈ 0.1
@test zero_rate(curve_spline_discount, [dt_curve+Dates.Day(11), dt_curve+Dates.Day(15)]) ≈ vert_y[1:2]

@test InterestRates.advancedays(InterestRates.BDays252(BusinessDays.Brazil()), Date(2015,9,1), [0, 1, 3, 4, 5]) == [Date(2015,9,1),Date(2015,9,2),Date(2015,9,4),Date(2015,9,8),Date(2015,9,9)]
@test InterestRates.advancedays(InterestRates.Actual360(), Date(2015,9,1), [0, 1, 3, 4, 5]) == [Date(2015,9,1),Date(2015,9,2),Date(2015,9,4),Date(2015,9,5),Date(2015,9,6)]
@test InterestRates.advancedays(InterestRates.Actual365(), Date(2015,9,1), [0, 1, 3, 4, 5]) == [Date(2015,9,1),Date(2015,9,2),Date(2015,9,4),Date(2015,9,5),Date(2015,9,6)]

param = 10
InterestRates.curve_set_dict_parameter!(curve_spline_discount, :custom_parameter, param)
@test InterestRates.curve_get_dict_parameter(curve_spline_discount, :custom_parameter) == param

# CompositeIRCurve
composite_curve = InterestRates.CompositeIRCurve(curve_spline_rates, curve_NS)
@test_throws AssertionError composite_curve_2 = InterestRates.CompositeIRCurve(curve_NS, curve_b252_ec_lin)
@test InterestRates.curve_get_date(composite_curve) == InterestRates.curve_get_date(curve_NS)
@test InterestRates.discountfactor(composite_curve, Date(2015,10,1)) ≈ InterestRates.discountfactor(curve_NS, Date(2015,10,1)) * InterestRates.discountfactor(curve_spline_rates, Date(2015,10,1))
@test InterestRates.ERF(composite_curve, Date(2015,10,1)) == InterestRates.ERF(curve_NS, Date(2015,10,1)) * InterestRates.ERF(curve_spline_rates, Date(2015,10,1))

# BufferedIRCurve
buffered_curve_ac360_cont_ff = InterestRates.BufferedIRCurve(curve_ac360_cont_ff)
dt_curve = InterestRates.curve_get_date(buffered_curve_ac360_cont_ff)

@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(11)) ≈ 0.1
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(11)) ≈ 0.1
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(15)) ≈ 0.15
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(15)) ≈ 0.15
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(19)) ≈ 0.20
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(19)) ≈ 0.20
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(23)) ≈ 0.19
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(23)) ≈ 0.19
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(16)) > 0.15
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(16)) > 0.15
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(17)) < 0.20
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(17)) < 0.20
@test forward_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(11), dt_curve + Dates.Day(15)) ≈ 0.2875 # forward_rate calculation on vertices
@test forward_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(11), dt_curve + Dates.Day(13)) ≈ 0.2875 # forward_rate calculation on interpolated maturity
@test ERF(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(13)) ≈ 1.00466361875533 # ffwd interp on ERF
@test ERF(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(13)) ≈ 1.00466361875533

dt_maturity = dt_curve+Dates.Day(30)
@test ERF_to_rate(buffered_curve_ac360_cont_ff, ERF(buffered_curve_ac360_cont_ff, dt_maturity), InterestRates.yearfraction(buffered_curve_ac360_cont_ff, dt_maturity)) ≈ zero_rate(buffered_curve_ac360_cont_ff, dt_maturity)
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(13)) ≈ 0.128846153846152 # ffwd interp as zero_rate
@test ERF(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(19), dt_curve + Dates.Day(23)) ≈ 1.00158458746737
@test forward_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(19), dt_curve + Dates.Day(23)) ≈ 0.1425000000000040
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(30)) ≈ 0.1789166666666680 # ffwd extrap after last vertice
@test forward_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(19), dt_curve + Dates.Day(23)) ≈ forward_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(50), dt_curve + Dates.Day(51))
@test forward_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(19), dt_curve + Dates.Day(23)) ≈ forward_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(50), dt_curve + Dates.Day(100))

@test forward_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(11), dt_curve + Dates.Day(15)) ≈ 0.2875
@test zero_rate(buffered_curve_ac360_cont_ff, dt_curve + Dates.Day(9)) ≈ 0.05833333333333 # ffwd extrap before first vertice

@test discountfactor(buffered_curve_ac360_cont_ff, dt_curve) == 1
@test isnan(ERF_to_rate(buffered_curve_ac360_cont_ff, 1.0, 0.0))

include("usage.jl")
include("perftests.jl")


export gradiente

function gradiente(
  nlp::AbstractNLPModel;
  atol = 1e-6,
  rtol = 1e-6,
  max_iter = 1000,
  max_time = 10.0
)
  x = copy(nlp.meta.x0)

  f(x) = obj(nlp, x)
  ∇f(x) = grad(nlp, x)

  ϵ = atol + rtol * norm(∇f(x))
  t₀ = time()

  iter = 0
  Δt = time() - t₀
  resolvido = norm(∇f(x)) < ϵ
  cansado = iter > max_iter || Δt > max_time

  status = :unknown

  α = 1.0
  while !(resolvido || cansado)

    x -= α * ∇f(x)
    α *= 0.99

    iter += 1
    Δt = time() - t₀
    resolvido = norm(∇f(x)) < ϵ
    cansado = iter > max_iter || Δt > max_time
  end

  if resolvido
    status = :first_order
  elseif cansado
    if iter > max_iter
      status = :max_iter
    elseif Δt > max_time
      status = :max_time
    end
  end

  return GenericExecutionStats(
    status,
    nlp,
    solution=x,
    objective=f(x),
    dual_feas=norm(∇f(x)),
    elapsed_time=Δt,
    iter=iter
  )
end


julia> meandegrees([350, 10])
0.0

julia> meandegrees([90, 180, 270, 360]])
0.0

julia> meandegrees([10, 20, 30]])
19.999999999999996


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
module O

using Plots

struct Wireframe{F, X, Y}
    f::F
    x::X
    y::Y
end

vecwithnan(X) = [
    vec([X ; fill(NaN, 1, size(X,  2))])
    vec([X'; fill(NaN, 1, size(X', 2))])
]

@recipe function f(wf::Wireframe)
    f, x, y = wf.f, wf.x, wf.y
    X, Y = reim(complex.(x', y))
    XX, YY = vecwithnan(X), vecwithnan(Y)
    ZZ = f.(XX, YY)
    XX, YY, ZZ
end

@recipe function f(wf::Wireframe{<:AbstractMatrix})
    Z, x, y = wf.f, wf.x, wf.y
    X, Y = reim(complex.(x', y))
    XX, YY, ZZ = vecwithnan(X), vecwithnan(Y), vecwithnan(Z)
    XX, YY, ZZ
end

end

# %%
using Plots
x = y = range(-1, 1, 51)
f(x, y) = x == y == 0 ? zero(x) : x^2*y/(x^2 + y^2)
plot(O.Wireframe(f, x, y); label="")

# %%
Z = f.(x', y)
plot(O.Wireframe(Z, x, y); label="")

# %%
wireframe(x, y, f)

# %%


using Distributions
using Evolutionary

"""

Return the expected fraction of the dataset processed per iteration. Here, `θs` is the fraction of
the dataset stored, `ps` is the number of partitions, and `Ls` is the probability of finishing
among the `nwait` fastest workers, for each worker.
"""
function fraction_processed(ps::AbstractVector, θs::AbstractVector, Ls::AbstractVector)
    rv = 0.0
    for i in 1:length(Ls)
        rv += Ls[i] * θs[i] / ps[i]
    end
    rv
end

"""Helper function for running event-drive simulations"""
function simulate(;θs, ps, cms_comp, cvs_comp, ds_comm, nwait, niterations=100, nsamples=10, update_latency=0)
    comp_means = cms_comp .* θs ./ ps
    comp_vars = cvs_comp .* θs ./ ps
    ds_comp = DSAGAnalysis.distribution_from_mean_variance.(Gamma, comp_means, comp_vars)
    nworkers = length(ps)
    latency = 0.0
    Ls = zeros(nworkers)
    for _ in 1:nsamples
        dfi, Lsi = simulate_iterations(;nwait, niterations, ds_comm, ds_comp, update_latency)
        latency += mean(dfi.latency)
        Ls .+= Lsi
    end
    latency /= nsamples
    Ls ./= nsamples
    latency, Ls
end

"""

Minimize the variance of the vector `Ls .* θs ./ ps` using a genetic optimization algorithm.
"""
function balance_contribution(ps0::AbstractVector{<:Real}, min_processed_fraction::Real; θs::AbstractVector{<:Real}, comp_means, comp_vars, comm_means, comm_vars, nwait::Integer, populationSize::Integer=100, tournamentSize::Integer=10, mutationRate::Real=1.0, time_limit::Real=10.0)
    nworkers = length(ps0)
    length(θs) == nworkers || throw(DimensionMismatch("θs has dimension $(length(θs)), but nworkers is $nworkers"))
    length(comp_means) == nworkers || throw(DimensionMismatch("comp_means has dimension $(length(comp_means)), but nworkers is $nworkers"))
    length(comp_vars) == nworkers || throw(DimensionMismatch("comp_vars has dimension $(length(comp_vars)), but nworkers is $nworkers"))
    length(comm_means) == nworkers || throw(DimensionMismatch("comm_means has dimension $(length(comp_means)), but nworkers is $nworkers"))
    length(comm_vars) == nworkers || throw(DimensionMismatch("comm_vars has dimension $(length(comp_vars)), but nworkers is $nworkers"))
    all((x)->0<x<=1, θs) || throw(ArgumentError("The entries of θs must be in (0, 1], but got $θs"))
    all((x)->0<x, ps0) || throw(ArgumentError("The entries of ps must be positive, but got $ps0"))

    # initialization
    @info "comp" comp_means comp_vars
    @info "comm" comm_means comm_vars
    ds_comm = DSAGAnalysis.distribution_from_mean_variance.(Gamma, comm_means, comm_vars)
    cms_comp = comp_means ./ (θs ./ ps0)
    cvs_comp = comp_vars ./ (θs ./ ps0)

    # constraint
    # (the total expected conbtribution must be above some threshold)
    function g(ps, Ls)
        fraction_processed(ps, θs, Ls) - min_processed_fraction
    end

    # objective function
    # (the variance of the contribution between workers)
    fworst = 0.0
    function f(ps)
        _, Ls = simulate(;θs, ps, cms_comp, cvs_comp, ds_comm, nwait)
        c = g(ps, Ls)
        if c < 0
            rv = fworst + abs(c)
        else
            rv = var(Ls .* θs ./ ps)
            fworst = max(fworst, rv)
        end
        rv
    end

    # evolutionary algorithm setup
    nworkers = length(ps0)
    selection = Evolutionary.tournament(tournamentSize)
    crossover = Evolutionary.LX()
    lower = max.(ones(Int, nworkers), ceil.(Int, ps0 ./ 2))
    upper = 2 .* ps0
    mutation = Evolutionary.domainrange((lower .- upper) ./ 10) # as recommended in the BGA paper

    # integer-output mutation that wraps another mutation
    function integer_mutation(x)
        if isone(mutationRate) || rand() < mutationRate
            mutation(x)
        end
        for i in 1:length(x)
            if rand() < 0.5
                x[i] = floor(x[i])
            else
                x[i] = ceil(x[i])
            end
        end
        x
    end

    # optimization algorithm
    opt = Evolutionary.GA(;populationSize, mutationRate=1.0, selection, crossover, mutation=integer_mutation)
    options = Evolutionary.Options(;time_limit, Evolutionary.default_options(opt)...)
    Evolutionary.optimize(f, lower, upper, float.(ps0), opt, options)
end

function load_balancer(chin::Channel; chout::Channel, ps0::AbstractVector, θs, min_processed_fraction::Real, nwait::Integer)
    @info "load_balancer task started"
    nworkers = length(ps0)
    0 < nworkers || throw(ArgumentError("nworkers must be positive, but is $nworkers"))
    length(θs) == nworkers || throw(DimensionMismatch("θs has dimension $(length(θs)), but nworkers is $nworkers"))
    ps = Vector{Int}(ps0)
    comp_means = fill(NaN, nworkers)
    comp_vars = fill(NaN, nworkers)
    comm_means = fill(NaN, nworkers)
    comm_vars = fill(NaN, nworkers)

    function process_sample(v)
        0 < v.worker <= nworkers || throw(ArgumentError("v.worker is $(v.worker), but nworkers is $nworkers"))
        isnan(v.comp_mean) || (comp_means[v.worker] = v.comp_mean)
        isnan(v.comp_var) || (comp_vars[v.worker] = v.comp_var)
        isnan(v.comm_mean) || (comm_means[v.worker] = v.comm_mean)
        isnan(v.comm_var) || (comm_vars[v.worker] = v.comm_var)
        return
    end

    # helper to check if there is any missing latency data
    all_populated = false
    function check_populated()
        if all_populated
            return all_populated
        end
        all_populated = iszero(count(isnan, comp_means))
        if all_populated
            all_populated = all_populated && iszero(count(isnan, comp_vars))
        end
        if all_populated
            all_populated = all_populated && iszero(count(isnan, comm_means))
        end
        if all_populated
            all_populated = all_populated && iszero(count(isnan, comm_vars))
        end
        all_populated
    end

    while isopen(chin)

        # consume all values currently in the channel
        try
            vin = take!(chin)
            process_sample(vin)
        catch e
            if e isa InvalidStateException
                break
            else
                rethrow()
            end
        end
        while isready(chin)
            try
                vin = take!(chin)
                process_sample(vin)
            catch e
                if e isa InvalidStateException
                    break
                else
                    rethrow()
                end
            end
        end

        # verify that we have complete latency information for all workers
        if !check_populated()
            continue
        end

        # new = balance_contribution(ps, min_processed_fraction; θs, ds_comm, cms_comp, cvs_comp, nwait)
        @info "started load-balancing optimization"
        result = balance_contribution(ps, min_processed_fraction; θs, comp_means, comp_vars, comm_means, comm_vars, nwait)
        new = result.minimizer

        # push any changes into the output channel
        for i in 1:nworkers
            if new[i] != ps[i]
                vout = @NamedTuple{worker::Int,p::Int}((i, new[i]))
                @info vout
                push!(chout, vout)
                ps[i] = new[i]
            end
        end
    end
    @info "load_balancer task finished"
    return
end


include("GB100.jl")
using DelimitedFiles
import GLM: fit, predict, LinearModel
using Distributions,SpecialFunctions
using LinearAlgebra, MultivariateStats
using StatsBase, Random, Distributed
using DelimitedFiles, DataStructures
using CSV, MLBase
using DataFrames

rng                     = Random.seed!(2019)

######## σ² = 1 #########
Pval = []
for i in 1:1000
    if (10 <= i <= 19)
        push!(Pval, i)
    elseif (20 <= i <= 99)
        if (i % 10 == 0)
            push!(Pval, i)
        end
    else
        if (i % 100 == 0)
            push!(Pval, i)
        end

    end
end

Pval
temp_storedim = size(collect(Pval),1)

rmsefitted2 = []
HyperS = Array{Union{Missing, Nothing, Int, Float64}, 2}(nothing, temp_storedim, 1000)

STORE = [] # storing udpated σ

Pval
for pval in Pval

    # reproducibility
    rng = Random.seed!(2019)
    gbdt = GB100.GBDT(;
      sampling_rate         = 0.75,                # Sampling rate
      learning_rate         = 0.1,                 # Learning rate
      num_iterations        = 100,                # NuGBer of iterations
      tree_options          = Dict(
          :nsubfeatures     => 0,
          :min_samples_leaf => 1,
          :max_depth        => 5))

    x, y                    = GB100.test_data(1000; p = pval, sigma = sqrt(1))
    x_test, y_test          = GB100.test_data(500; p = pval, sigma = sqrt(1))
    mu                      = GB100.fried(x)
    mu_test                 = GB100.fried(x_test)

    # initialize the hypers
    hypers = GB100.Hypers(x, y, 1, 1, 0.8)
    hypers.σ = sqrt(1)
    hypers.τ = 3 / 10 * var(y) / gbdt.num_iterations
    hypers.α = 1
    hypers.λ = nothing  # hypers.λ = 2 or 20 for Poisson(λ); hypers.λ = 1.0 or 2.0 for expDecay(λ = 1),vor hypers.λ = nothing for no priors
    hypers.s =  ones(size(x, 2)) / size(x, 2)

    # initialize opt
    opts = GB100.Opts()
    opts.num_burn = 5

    rmse2 = GB100.TBD_predict_and_rmse2!(
        x,
        y,
        mu,
        x_test,
        y_test,
        mu_test,
        gbdt,
        hypers,
        opts,
        rng)
    println("The rmse2 b.w. μ_test and pred for p = $(pval)  is $(rmse2).")

    push!(rmsefitted2, rmse2)
    push!(STORE, hypers.σ)

    # store the hyper.s for the ensembel tree
    j = findall(x -> x == pval, Pval)
    HyperS[j, 1:size(hypers.s, 1)] = hypers.s
    #println("After $(opts.num_burn) times of burning, the final sum of the first 5 elements of the prob vec is:\n$(sum(hypers.s[1:5])) ")
end

rmsefitted2
HyperS
STORE

HyperS_copy = copy(HyperS)
for i in 1 : size(HyperS, 1)
    for j in 1 : size(HyperS, 2)
        if HyperS_copy[i, j] == nothing
            HyperS_copy[i, j] = 0.0
        end
    end
end
HyperS_copy

rmsefitted2 = DataFrame(rmsefitted2 = rmsefitted2)
CSV.write("_rmse2_sigma1.csv", rmsefitted2)
HyperS = DataFrame(HyperS)
HyperS_copy = DataFrame(HyperS_copy)
CSV.write("_HyperS_copy_sigma1.csv", HyperS_copy)
STORE

for i in 1 : size(STORE, 1)
    temp = DataFrame(STORE[i,])
    CSV.write("STORE.csv", temp, append = true)
end
##### not runned yet

######## σ² = 10 #########
Pval = []
for i in 1:1000
    if (10 <= i <= 19)
        push!(Pval, i)
    elseif (20 <= i <= 99)
        if (i % 10 == 0)
            push!(Pval, i)
        end
    else
        if (i % 100 == 0)
            push!(Pval, i)
        end

    end
end

Pval
temp_storedim = size(collect(Pval),1)
rmsefitted2_sigma10 = []
HyperS_sigma10 = Array{Union{Missing, Nothing, Int, Float64}, 2}(nothing, temp_storedim, 1000)
STORE_sigma10 = []
Pval

for pval in Pval

    # reproducibility
    rng = Random.seed!(2019)

    gbdt = GB100.GBDT(;
      sampling_rate         = 0.75,                # Sampling rate
      learning_rate         = 0.1,                 # Learning rate
      num_iterations        = 100,                # NuGBer of iterations
      tree_options          = Dict(
          :nsubfeatures     => 0,
          :min_samples_leaf => 1,
          :max_depth        => 5))

    ### simulate the training and testing data
    X, Y                    = GB100.test_data(1000; p = pval, sigma = sqrt(10.0))
    X_test, Y_test          = GB100.test_data(500; p = pval, sigma = sqrt(10.0))
    μ                       = GB100.fried(X)
    μ_test                  = GB100.fried(X_test)

    # initialize the hypers
    hypers                  = GB100.Hypers(X, Y, 1, 1, 0.8)
    hypers.σ                = sqrt(1)
    hypers.τ                = 3 / 10 * var(Y) / gbdt.num_iterations
    hypers.α                = 1
    hypers.λ = nothing  # hypers.λ = 2 or 20 for Poisson(λ); hypers.λ = 1.0 or 2.0 for expDecay(λ = 1), vor hypers.λ = nothing for no priors
    hypers.s                =  ones(size(X, 2)) / size(X, 2)

    # initialize opt
    opts = GB100.Opts()
    opts.num_burn = 5

    # train the model and get the final model as well as rmse2
    rmse2 = GB100.TBD_predict_and_rmse2!(
        X,
        Y,
        μ,
        X_test,
        Y_test,
        μ_test,
        gbdt,
        hypers,
        opts,
        rng)

    println("The rmse2 b.w. μ_test and pred for p = $(pval)  is $(rmse2).")
    push!(rmsefitted2_sigma10, rmse2)
    push!(STORE_sigma10, hypers.σ)
    # store the hyper.s for the ensembel tree
    j = findall(x -> x == pval, Pval)
    HyperS_sigma10[j, 1:size(hypers.s, 1)] = hypers.s
    #println("After $(opts.num_burn) times of burning, the sum of the first 5 elements in the prob vec is: $(sum(hypers.s[1:5])).")
end

rmsefitted2_sigma10
HyperS_sigma10
STORE_sigma10

rmsefitted2_sigma10 = DataFrame(rmsefitted2_sigma10 = rmsefitted2_sigma10)
CSV.write("_rmse2_sigma10.csv", rmsefitted2_sigma10)
HyperS_sigma10 = DataFrame(HyperS_sigma10)
HyperS_copy_sigma10 = copy(HyperS_sigma10)

for i in 1 : size(HyperS_sigma10, 1)

    for j in 1 : size(HyperS_sigma10, 2)

        if HyperS_copy_sigma10[i, j] == nothing

            HyperS_copy_sigma10[i, j] = 0.0

        end

    end

end

HyperS_copy_sigma10
CSV.write("_HyperS_sigma10.csv", HyperS_sigma10)

STORE_sigma10

for i in 1 : size(STORE_sigma10, 1)
    temp = DataFrame(STORE_sigma10[i,])
    CSV.write("STORE_sigma10.csv", temp, append = true)
end

############# Visualizing the result ################
using DataFrames, Gadfly
ind = collect(1:28)
mydata = DataFrame(hcat(ind, rmsefitted2, rmsefitted2_sigma10))
names!(mydata, [:indSeq, :σ²1, :σ²10])
mystack = stack(mydata, [:σ²1, :σ²10])

myplot = plot(stack(mydata, [:σ²1, :σ²10]),
    x = :indSeq,
    y = :value,
    color = :variable,
    shape = :variable,
    Geom.point,
    Geom.line)

using Cairo
#draw(PNG("c.png", 7inch, 4.88inch), myplot)  # expDecay(c = 1.0)
#draw(PNG("c2.png", 7inch, 4.88inch), myplot) # expDecay(c = 2.0) this one is not good at all
#draw(PNG("lambda2.png", 7inch, 4.88inch), myplot) # Poisson(λ = 2)
#draw(PNG("lambda20.png", 7inch, 4.88inch), myplot) # Poisson(λ = 20)
#draw(PNG("updateSigma.png", 7inch, 4.88inch), myplot) # Poisson(λ = 20)
#draw(PNG("updateSigmaPoisson2.png", 7inch, 4.88inch), myplot) # Poisson(λ = 20)
#draw(PNG("updateSigmaPoisson20.png", 7inch, 4.88inch), myplot) # Poisson(λ = 20)
draw(PNG("tuneSigmabyMU_test.png", 7inch, 4.88inch), myplot) # Poisson(λ = 20)


#===============================================================================
This file contains an independent implementation of the excellent  line search
algorithm described in:

	"Line Search Algorithms with Guaranteed Sufficient Decrease"
	Jorge J. Moré and David J. Thuente
    ACM Transactions on Mathematical Software (TOMS) (1994)
	doi: 10.1145/192115.192132

Features:
  - Finite termination
  - Wolfe conditions
  - α stays in minstep/maxstep
Rwquires:
  - c1 ≤ c2
  - ϕ is C1
===============================================================================#

struct MoreThuente{T1, T2} <: LineSearch
    ratio::T1
	c::T1
	maxiter::T2
	verbose::Bool
end

#===============================================================================
                            Trial value selection
	Notation:
		αc is the minimizer of the cubic that interpolates fl, ft, gl, gt
		αq is the minimizer of the quadratic that interpolates fl, ft, gl
		αs is the minimizer of the quadratic that interpolates fl, gl, gt
===============================================================================#
#===============================================================================
	CASE 1: a) ft > fl

	Compute αc and αq and set αt⁺ = ifelse(|αc-αl|<|αq-αl, αc, (αq+αc)/2).
	This choice places us closer to αl - the step length that gives us the
	lowest function value so far.
===============================================================================#
αt⁺ = abs(αc - αl) < abs(αq - αl) ? αc : (αq + αc)/2
#===============================================================================
	CASE 2: a) ft ≤ fl
	        b) gt×gl < 0
===============================================================================#
αt⁺ = abs(αc - αt) < abs(αs - αt) ? αc : αs
#===============================================================================
	CASE 3: a) ft ≤ fl
            b) gt×gl ≥ 0
            c) |gt| ≤ |gl|

	αt⁺ :  cubic that interpolates fl, ft, gl, gt may not have a minimizer
===============================================================================#
if cubic tend to inf for alpha inc and is αc > αt?
	αt⁺ = ifelse(abs(αc - αt) < abs(αs - αt)), αc, αs
else
	αt⁺ = αs
end
#===============================================================================
	CASE 4: a) ft ≤ fl
	        b) gt×gl ≥ 0
	        c) |gt| < |gl|

	αt⁺ : minimizer of cubic that interpolates fu, ft, gu, gt
===============================================================================#

function find_steplength(ls::Union{Backtracking, TwoPointQuadratic}, φ::T, d, x,
	                     λ, #
						 ϕ_0::Tf,
						 ∇f_x,
						 dϕ_0=dot(d, ∇f_x);
						 minlength,
						 maxlength) where {T, Tf}
	ratio, c, maxiter, verbose = Tf(ls.ratio), Tf(ls.c), ls.maxiter, ls.verbose

    if verbose
        println("Entering line search with step size: ", λ)
        println("Initial value: ", ϕ_0)
        println("Value at first step: ", φ(λ))
    end

    t = -c*dϕ_0

    iter, α, β = 0, λ, λ # iteration variables
	f_α = φ(α)   # initial function value

	is_solved = isfinite(f_α) && f_α <= ϕ_0 + c*α*t

    while !is_solved && iter <= maxiter
        iter += 1
        β, α, f_α = _solve_pol(ls, φ, ϕ_0, dϕ_0, α, f_α, ratio)
		is_solved = isfinite(f_α) && f_α <= ϕ_0 + c*α*t
    end

	ls_success = iter >= maxiter ? false : true

    if verbose
		!ls_success && println("maxiter exceeded in backtracking")
        println("Exiting line search with step size: ", α)
        println("Exiting line search with value: ", f_α)
    end
    return α, f_α, ls_success
end


module SingleFloats

export Single32

using Base: IEEEFloat
import Base: promote_rule, promote_type, convert, hash, show, string, repr, parse, tryparse,
    cmp, ==, !=, >=, >, <=, <, isless, isequal, iseven, isfinite, isnan, isodd, issubnormal,
    +, -, *, /, \, %, √, ÷, inv,
    prevfloat, nextfloat, eps, exponent, significand, zero, one, iszero, isone, isinteger, decompose,
    typemin, typemax, floatmin, floatmax, maxintfloat, widen, precision,
    ^, sqrt, cbrt, ∛, div, rem, fld, mod, cld, hypot, min, max, minmax, maximum, minimum, rad2deg, deg2rad, mod2pi, rem2pi,
    abs, abs2, signbit, sign, copysign, flipsign, floor, ceil, trunc, round, divrem, fldmod, clamp

import Base.Math: exp, exp2, exp10, expm1, log, log2, log10, log1p,
    sin, cos, tan, csc, sec, cot, asin, acos, atan, acsc, asec, acot,
    sinh, cosh, tanh, csch, sech, coth, asinh, acosh, atanh, acsch, asech, acoth

using LinearAlgebra
import LinearAlgebra: rank, cond, norm, opnorm, det, tr, inv, pinv,
                      transpose, adjoint, eigvals, eigvals!, svdvals, svdvals!, svd,
                      iszero, isone, isdiag, issymmetric, ishermitian,
                      isposdef, isposdef!, istril, istriu,
                      lu, lu!, qr, qr!,
                      schur, schur!, hessenberg, hessenberg!,
                      cholesky, cholesky!, factorize
using Random

include("type/Single32.jl")
include("math/arith.jl")
include("math/elementary.jl")
include("math/linearalgebra.jl")

end # SingleFloats


####
#### Network
####

"""
    Network

A graph along with strings labeling the vertices.
"""
struct Network
    graph::Graph
    persons::Vector{String}
end

"""
    Network()

Create a network with no persons.
"""
Network() = Network(Graph(), String[])

"""
    Network(persons)

Create a network initialized with the iterator `persons` of `String`s.
The network contains no friend relations.
"""
Network(persons) = Network(Graph(length(persons)), collect(persons))

## Make Network behave as a scalar in broadcasting expressions. That is, when
## automatically broadcasting a scalar function over arrays, make Network behave
## like a scalar would.
Base.Broadcast.broadcastable(net::Network) = Ref(net)

# Show a compact representation.
function Base.show(io::IO, net::Network)
    print(io, "Network(num_persons=$(nv(net.graph)), num_friendships=$(ne(net.graph)))")
    return nothing
end

# Forward some functions to the graph.
for func in (:adj_list, :add_vertex!)
    @eval ($func)(net::Network, args...) = ($func)(net.graph, args...)
end

"""
    persons(net::Network)

Return the list of all persons in `net`.
"""
persons(net::Network) = net.persons

"""
    person(net::Network, i::Integer) = (net.persons)[i]

Return the name (`String`) of the `i`th person in `net`.
"""
person(net::Network, i::Integer) = (net.persons)[i]

## Index notation looks up a person's index from the name. Eg.
## net["Alice"] -> 10
Base.getindex(net::Network, person::AbstractString) = findfirst(isequal(person), net.persons)

"""
    ispresent(net::Network, person)

Return `true` if `person` is in `net`.
"""
ispresent(net::Network, person) = (findfirst(isequal(person), net.persons) != nothing)

"""
    friends(net::Network, person)

Return the list of friends of `person`.
"""
function friends(net::Network, person)
    iperson = net[person]
    friend_inds::Vector{Int} = adj_list(net)[iperson]
    return net.persons[friend_inds]
end

"""
    add_person!(net::Network, person::AbstractString)

Add `person` to `net`. No friend relations are added.
"""
function add_person!(net::Network, person::AbstractString)
    persons = net.persons
    ispresent(net, person) && error("Person $person already present.")
    add_vertex!(net)
    push!(persons, person)
    return net
end

"""
    friend_pair!(net::Network, person1, person2)

Cause `person1` to "friend" `person2` (and vice-versa) in the network `net`.
"""
function friend_pair!(net::Network, person1, person2)
    add_edge!(net.graph, net[person1], net[person2])
    return net
end

"""
    defriend_pair!(net::Network, person1, person2)

Cause `person1` to "defriend" `person2` in the network `net`.
This just removes the symmetric friend relation between the
two persons.
"""
function defriend_pair!(net::Network, person1, person2)
    del_edge!(net.graph, net[person1], net[person2])
    return net
end

"""
    degree_of_separation(net::Network, person1, person2)

Return the degree of separation between `person1` and `person2`.
If the persons are equal, the distance is zero. If they are nearest
neighbors, it is one. In general, it is the length of a shortest
friendship path between the persons.
"""
function degree_of_separation(net::Network, person1, person2)
    return shortest_path_length(net.graph, net[person1], net[person2])
end


module TestModels

include("util.jl")

using QuantEcon, QuantEcon.Models
using Distributions
using Base.Test
using FactCheck

function set_up_data(cp::ConsumerProblem)
    f = get_data_file()

    # read in or compute v_star
    if exists(f, "ifp/v_star")
        v_star = read(f, "ifp/v_star")
    else
        v_init, c_init = init_values(cp)
        v_star = compute_fixed_point(x-> bellman_operator(cp, x), v_init,
                                     max_iter=5000, err_tol=1e-10,
                                     verbose=true, print_skip=10)
        write(f, "ifp/v_star", v_star)
    end

    # read in or compute c_star_vfi
    if exists(f, "ifp/c_star_vfi")
        c_star_vfi = read(f, "ifp/c_star_vfi")
    else
        c_star_vfi = get_greedy(cp, v_star)
        write(f, "ifp/c_star_vfi", c_star_vfi)
    end

    # read in or compute c_star
    if exists(f, "ifp/c_star")
        c_star = read(f, "ifp/c_star")
    else
        v_init, c_init = init_values(cp)
        c_star = compute_fixed_point(x-> coleman_operator(cp, x), c_init,
                                     max_iter=5000, err_tol=1e-10,
                                     verbose=true, print_skip=10)
        write(f, "ifp/c_star", c_star)
    end

    close(f)

    return v_star, c_star_vfi, c_star
end

function set_up_data(jv::JvWorker)
    f = get_data_file()

    # read in or compute v_star
    if exists(f, "jv/v_star")
        v_star = read(f, "jv/v_star")
    else
        v_init = [jv.x_grid .* 0.5]
        v_star = compute_fixed_point(x-> bellman_operator(jv, x), v_init,
                                     max_iter=5000, err_tol=1e-10,
                                     verbose=true, print_skip=10)
        write(f, "jv/v_star", v_star)
    end

    # read in or compute policies
    if exists(f, "jv/s_star")
        s_star = read(f, "jv/s_star")
    else
        v_init = [jv.x_grid .* 0.5]
        s_star, phi_star = get_greedy(jv, v_star)
        write(f, "jv/s_star", s_star)
    end

    if exists(f, "jv/phi_star")
        phi_star = read(f, "jv/phi_star")
    else
        if !isdefined(:phi_star)  # otherwise computed above
            _, phi_star = get_greedy(jv, v_star)
        end
        write(f, "jv/phi_star", phi_star)
    end

    close(f)

    return v_star, s_star, phi_star
end

function set_up_data(lt::LucasTree)
    f = get_data_file()

    # read in or compute prices
    if exists(f, "lt/prices")
        prices = read(f, "lt/prices")
    else
        prices = compute_lt_price(lt, err_tol=1e-10, max_iter=5000)
        write(f, "lt/prices", prices)
    end

    close(f)

    return prices
end

function set_up_data(sp::SearchProblem)
    f = get_data_file()

    # read in or compute v_star
    if exists(f, "odu/v_star")
        v_star = read(f, "odu/v_star")
    else
        v_init = fill(sp.c / (1 - sp.bet), sp.n_w, sp.n_pi)
        v_star = compute_fixed_point(x-> bellman_operator(sp, x), v_init,
                                     print_skip=1, err_tol=1e-10, max_iter=5000)
        write(f, "odu/v_star", v_star)
    end

    # read in or compute phi_vfi
    if exists(f, "odu/phi_vfi")
        phi_vfi = read(f, "odu/phi_vfi")
    else
        phi_vfi = get_greedy(sp, v_star)
        write(f, "odu/phi_vfi", phi_vfi)
    end

    # read in or compute phi_pfi
    if exists(f, "odu/phi_pfi")
        phi_pfi = read(f, "odu/phi_pfi")
    else
        phi_init = ones(sp.n_pi)
        phi_pfi = compute_fixed_point(x-> res_wage_operator(sp, x), phi_init,
                                      print_skip=1, err_tol=1e-10, max_iter=5000)
        write(f, "odu/phi_pfi", phi_pfi)
    end

    close(f)

    return v_star, phi_vfi, phi_pfi
end

function set_up_data(gm::GrowthModel)
    f = get_data_file()

    # read in or compute v_star
    if exists(f, "gm/v_star")
        v_star = read(f, "gm/v_star")
    else
        v_init = 5 .* gm.u([gm.grid]) .- 25
        v_star = compute_fixed_point(x->bellman_operator(gm, x), v_init,
                                     err_tol=1e-10, max_iter=5000, print_skip=1)
        write(f, "gm/v_star", v_star)
    end

    close(f)

    return v_star
end

facts("Testing asset_pricing.jl") do
    n = 5
    P = 0.0125 .* ones(n, n)
    P .+= diagm(0.95 .- 0.0125 .* ones(5))
    s = [1.05, 1.025, 1.0, 0.975, 0.95]
    γ = 2.0
    β = 0.94
    ζ = 1.0
    p_s = 150.0
    ap = AssetPrices(β, P, s, γ)

    tree_p = tree_price(ap)
    consol_p = consol_price(ap, ζ)
    call_option_p = call_option(ap, ζ, p_s)

    context("test shapes") do
        @fact size(ap.P, 1) --> size(ap.P, 2)
        @fact size(ap.P, 1) --> ap.n
        @fact length(tree_p) --> ap.n
        @fact length(consol_p) --> ap.n
        @fact length(call_option_p[1]) --> ap.n
    end

    context("test P_tilde") do
        P_tilde = ap.P_tilde
        P_tilde2 = similar(P_tilde)
        for i=1:ap.n, j=1:ap.n
            P_tilde2[i, j] = ap.P[i, j] * ap.s[j]^(1.0-ap.gamm)
        end
        @fact P_tilde --> roughly(P_tilde2)
    end

    context("test P_check") do
        P_check = ap.P_check
        P_check2 = similar(P_check)
        for i=1:ap.n, j=1:ap.n
            P_check2[i, j] = ap.P[i, j] * ap.s[j]^(- ap.gamm)
        end
        @fact P_check --> roughly(P_check2)
    end

    context("test multiple periods call option") do
        w_bars = call_option(ap, ζ, p_s, [5, 7])[2]
        @fact length(w_bars) --> 2
    end
end  # facts

facts("Testing career.jl") do
    cp = CareerWorkerProblem()
    v_init = rand(cp.N, cp.N)
    v_prime = bellman_operator(cp, v_init)
    greedy = get_greedy(cp, v_init)

    context("test shapes") do
        @fact size(v_init) --> size(v_prime)
        @fact size(v_init) --> size(greedy)
    end

    context("test model intuition") do
        # new life with worst job and career
        if any(greedy .== 3)
            @fact greedy[1, 1] --> 3
        end

        # new job with best career and worst job
        @fact greedy[end, 1] --> 2

        # say put with best job and career
        if any(greedy .== 1)
            @fact greedy[end, end] --> 1
        end
    end

    context("test mutating functions") do
        # new array to update inplace in bellman_operator! and get_greedy!
        out = copy(v_init)
        get_greedy!(cp, v_init, out)
        @fact out --> roughly(greedy, atol=1e-12)

        # now test bellman_operator!
        bellman_operator!(cp, v_init, out)
        @fact out --> roughly(v_prime)
    end
end  # facts

facts("Testing ifp.jl") do
    cp = ConsumerProblem()
    v_star, c_star_vfi, c_star = set_up_data(cp)

    context("test bellman and coleman policies agree") do
        @fact Base.maxabs(c_star_vfi - c_star) <= 0.2 --> true
    end

    context("test bellman solution is fixed point") do
        @fact v_star --> roughly(bellman_operator(cp, v_star); atol=1e-6)
    end

    context("test coleman solution is fixed point") do
        @fact c_star --> roughly(coleman_operator(cp, c_star); atol=1e-6)
    end

    context("test shape of init_values") do
        v_init, c_init = init_values(cp)
        shapes = (length(cp.asset_grid), length(cp.z_vals))
        @fact size(v_init) --> shapes
        @fact size(c_init) --> shapes
    end

    context("test mutating functions") do
        # new array to update inplace in bellman_operator! and get_greedy!
        out = copy(v_star)
        get_greedy!(cp, v_star, out)
        @fact out --> roughly(c_star_vfi, atol=1e-12)

        # now test bellman_operator!
        bellman_operator!(cp, v_star, out)
        @fact out --> roughly(v_star, atol=1e-12)

        # now test coleman_operator!
        coleman_operator!(cp, c_star, out)
        @fact out --> roughly(c_star; atol=1e-6)
    end
end  # facts

facts("Testing jv.jl") do
    A = 1.4
    α = 0.6
    β = 0.96
    grid_size = 50

    jv = JvWorker(A=A, alpha=α, bet=β, grid_size=grid_size)
    v_star, s_star, phi_star = set_up_data(jv)
    n = length(jv.x_grid)

    context("test model intuition") do
        # s preferred to phi with low x?
        # low x is an early index
        @fact s_star[1] > phi_star[1] --> true

        # phi preferred to s with high x?
        # high x is a late index
        @fact phi_star[end] >  s_star[end] --> true
    end

    context("policies correct size") do
        @fact length(s_star) --> n
        @fact length(phi_star) --> n
    end

    context("solution to bellman is fixed point") do
        @fact v_star --> roughly(bellman_operator(jv, v_star); atol=1e-6)
    end

    context("test mutating functions") do
        # new array to update inplace in bellman_operator! and get_greedy!
        out1 = copy(v_star)
        out2 = copy(v_star)
        get_greedy!(jv, v_star, (out1, out2))
        @fact out1 --> roughly(s_star, atol=1e-12)
        @fact out2 --> roughly(phi_star, atol=1e-12)

        # now test bellman_operator!
        bellman_operator!(jv, v_star, out1)
        @fact out1 --> roughly(v_star, atol=1e-12)
    end
end  # facts

facts("Testing lucastree.jl") do
    # model parameters
    γ = 2.0
    β  = 0.95
    α = 0.90
    σ = 0.1

    lt = LucasTree(γ, β, α, σ)
    prices = set_up_data(lt)

    @fact size(prices) --> size(lt.grid)

    context("test integrate") do
        g(x) = x .* 0.0 + 1.0
        est = Models.integrate(lt, g)

        # Same as bounds in lucastree.jl
        _int_min, _int_max = exp(-4 * σ), exp(4 * σ)
        exact = cdf(lt.phi, _int_max) - cdf(lt.phi, _int_min)

        @fact abs(est - exact) <= 0.1 --> true
    end

    context("test lucas_operator fp") do
        old_f = prices ./ (lt.grid.^γ)
        new_f = lucas_operator(lt, old_f)
        new_p = new_f .* lt.grid.^γ

        @fact prices --> roughly(new_p; atol=1e-3)
    end

    context("test prices increasing in y") do
        @fact prices --> sort(prices)
    end
end  # facts

facts("Testing odu.jl") do
    β = 0.95
    c = 0.6
    F_a = 1
    F_b = 1
    G_a = 3
    G_b = 1.2
    w_max = 2
    w_grid_size = 40
    pi_grid_size = 40

    sp = SearchProblem(β, c, F_a, F_b, G_a, G_b, w_max, w_grid_size,
                       pi_grid_size)

    v_star, phi_vfi, phi_pfi = set_up_data(sp)

    context("tests shapes of vfi outputs") do
        @fact size(v_star) --> size(phi_vfi)
    end

    context("phi_vfi increasing?") do
        phi_vfi_sorted = true
        for col=1:sp.n_pi
            if !(issorted(phi_vfi[:, col]))
                phi_vfi_sorted = false
            end
        end
        @fact phi_vfi_sorted --> true
    end

    context("v_star increasing?") do
        v_star_sorted = true
        for col=1:sp.n_pi
            if !(issorted(v_star[:, col]))
                v_star_sorted = false
            end
        end
        @fact v_star_sorted --> true
    end

    context("test model intuition") do
        # phi_pfi increasing?
        @fact issorted(phi_pfi, rev=true) --> true

        # v_star fixed_point?
        @fact v_star --> roughly(bellman_operator(sp, v_star); atol=1e-5)
    end

    context("phi_pfi fixed_point?") do
        @fact phi_pfi --> roughly(res_wage_operator(sp, phi_pfi); atol=1e-5)
    end

    context("test mutating functions") do
        # new array to update inplace in bellman_operator! and get_greedy!
        out = copy(v_star)
        get_greedy!(sp, v_star, out)
        @fact out --> roughly(phi_vfi, atol=1e-12)

        # now test bellman_operator!
        bellman_operator!(sp, v_star, out)
        @fact out --> roughly(v_star, atol=1e-12)

        # now test res_wage_operator!
        out2 = copy(phi_pfi)
        res_wage_operator!(sp, phi_pfi, out2)
        @fact out2 --> roughly(phi_pfi, atol=1e-12)
    end
end  # facts

facts("Testing optgrowth.jl") do
    # model parameters
    α = 0.65
    func(k) = k.^α
    β = 0.95
    u = log
    grid_max = 2
    grid_size = 150
    gm = GrowthModel(func, β, u, grid_max, grid_size)

    # get value/policy functions we solved for
    v_star = set_up_data(gm)
    sigma = get_greedy(gm, v_star)

    # compute true policy rule
    true_sigma = (1 - α*β) .* gm.grid.^α

    # compute true value function
    ab = α * β
    c1 = (log(1 - ab) + log(ab) * ab / (1 - ab)) / (1 - β)
    c2 = α / (1 - ab)
    true_v_star(k) = c1 + c2.*log(k)

    context("test computed/analytical policies are close") do
        @fact sigma --> roughly(true_sigma; atol=1e-2)
    end

    context("test computed/analytical values are close") do
        # Note first point is garbage b/c of interp
        @fact v_star[2:end] --> roughly(true_v_star(gm.grid)[2:end]; atol=5e-2)
    end

    context("test v_star fixed point.") do
        @fact v_star[2:end] --> roughly(bellman_operator(gm, v_star)[2:end];
                                       atol=5e-2)
    end

    context("test mutating functions") do
        # new array to update inplace in bellman_operator! and get_greedy!
        out = copy(v_star)
        get_greedy!(gm, v_star, out)
        @fact out --> roughly(sigma, atol=1e-12)

        # now test bellman_operator!
        bellman_operator!(gm, v_star, out)
        @fact out --> roughly(v_star, atol=1e-12)
    end
end  # facts

end  # module


################################
####### Define parameters ######
################################
import Parameters: @with_kw

"""
Mutable structure containing the variables required to solve the analytical approach. All the functions are solve using the internal values of the structure. You should declare a mutable structure to the perform the analytical estimations.

# Parameters
 - `gamNeg::Int64`: Selection coefficient for deleterious alleles
 - `gL::Int64`: Selection coefficient for weakly benefical alleles
 - `gH::Int64`: Selection coefficient for strongly benefical alleles
 - `alLow::Float64`: Proportion of α due to weak selection
 - `alTot::Float64`: α
 - `thetaF::Float64`: Mutation rate defining BGS strength
 - `thetaMidNeutral::Float64`: Mutation rate on coding region
 - `al::Float64`: DFE shape parameter
 - `be::Float64`: DFE scale parameter
 - `B::Float64`: BGS strength
 - `B_bins::Array{Float64,1}`: BGS values to simulate
 - `pposL::Float64`: Fixation probabily of weakly beneficial alleles
 - `pposH::Float64`: Fixation probabily of strongly beneficial alleles
 - `N::Int64`: Population size
 - `n::Int64`: Sample size
 - `Lf::Int64`: Flanking region length
 - `rho::Float64`: Recombination rate
 - `TE::Float64`

"""
@with_kw mutable struct parameters
	gamNeg::Int64              = -457
	gL::Int64                  = 10
	gH::Int64                  = 500
	alLow::Float64             = 0.2
	alTot::Float64             = 0.4
	thetaF::Float64            = 1e-3
	thetaMidNeutral::Float64   = 1e-3
	al::Float64                = 0.184
	be::Float64                = 0.000402
	B::Float64                 = 0.999
	B_bins::Array{Float64,1}   = push!(collect(0.1:0.025:0.975),0.999)
	pposL::Float64             = 0
	pposH::Float64             = 0.001
	N::Int64                   = 1000
	n::Int64                   = 500
	Lf::Int64                  = 2*10^5
	rho::Float64               = 0.001
	TE::Float64                = 5.0
	diploid::Bool              = false

	NN::Int64 = 2*N
	nn::Int64 = 2*n
	θᵣ::Array{Float64,1} = fill(thetaMidNeutral,nn-1)
	dac::Array{Int64,1} = [2,4,5,10,20,50,200,500,700]
end

"""
Mutable structure containing the downsampled SFS.

# Returns
 - `bn::Dict`: SparseMatrixCSC containing the binomial convolution

"""
@with_kw mutable struct binomialDict
	bn::Dict = Dict{Float64,SparseMatrixCSC{Float64,Int64}}()
end
#=const binomialDict = Dict{Float64,SparseMatrixCSC{Float64,Int64}}()=#

################################
###### Solving parameters ######
################################
"""
	Br(Lmax,theta)

Expected reduction in nucleotide diversity. Explored at [Charlesworth B., 1994](https://doi.org/10.1017/S0016672300032365):

```math
\\frac{\\pi}{\\pi_{0}} = e^{\\frac{4\\muL}{2rL+t}}
```

# Arguments
 - `param::parameters`
 - `theta::Float64`
# Returns
 - `Float64`: expected reduction in diversity given a non-coding length, mutation rate and defined recombination.
"""
function Br(param::parameters,theta::Float64)

	ρ::Float64     = param.rho
	t::Float64     = -1.0*param.gamNeg/(param.NN+0.0)
	μ::Float64     = theta/(2.0*param.NN)
	r::Float64     = ρ/(2.0*param.NN)

	out::Float64  = ℯ^(-4*μ*param.Lf/(2*param.Lf*r+t))
	return out
end

# Set mutation rate given the expected reduction in nucleotide diversity (B value ) in a locus.
"""
	setThetaF!(param)

Find the optimum mutation given the expected reduction in nucleotide diversity (B value) in a locus.

# Returns
 - `adap.thetaF::Float64`: changes adap.thetaF value.
"""
function setThetaF!(param::parameters)

	i(θ,p=param) = Br(p,θ)-p.B
	thetaF       = find_zero(i,0.0)
	param.thetaF = thetaF
end

function alphaExpSimLow(param::parameters,pposL::Float64,pposH::Float64)
	return fixPosSim(param,param.gL,0.5*pposL)/(fixPosSim(param,param.gL,0.5*pposL)+fixPosSim(param,param.gH,0.5*pposH) + fixNegB(param,0.5*pposL+0.5*pposH))
end

function alphaExpSimTot(param::parameters,pposL::Float64,pposH::Float64)
	return (fixPosSim(param,param.gL,0.5*pposL)+fixPosSim(param,param.gH,0.5*pposH))/(fixPosSim(param,param.gL,0.5*pposL)+fixPosSim(param,param.gH,0.5*pposH)+fixNegB(param,0.5*pposL+0.5*pposH))
end

function solvEqns(param,config)

	pposL,pposH = config
	return (alphaExpSimTot(param,pposL,pposH)-param.alTot,alphaExpSimLow(param,pposL,pposH)-param.alLow)
end

"""
	setPpos!(param)

Find the probabilty of positive selected alleles given the model. It solves a equation system taking into account fixations probabilities of weakly and strong beneficial alleles.

# Returns
 - `Tuple{Float64,Float64}`: weakly and strong beneficial alleles probabilites.
"""
function setPpos!(param::parameters)

	function f!(F,x,param=param)
		F[1] = alphaExpSimTot(param,x[1],x[2])-param.alTot
		F[2] = alphaExpSimLow(param,x[1],x[2])-param.alLow
	end

	pposL,pposH = nlsolve(f!,[0.0; 0.0]).zero

	if pposL < 0.0
		 pposL = 0.0
	end
	if pposH < 0.0
		 pposH = 0.0
	end

	param.pposL,param.pposH = pposL, pposH

 end

"""
	binomOp(param)

Binomial convolution to sample the allele frequencies probabilites depending on background selection values and sample size.

# Arguments
 - `param::parameters`
 - `convolutedSamples::binomialDict`

# Returns
 - `Array{Float64,2}`: convoluted SFS for each B value defined in the model (param.B_bins). The estimations are saved at *convolutedBn.bn*.
"""
function binomOp!(param::parameters,convolutedBn::Dict)

	#=bn = Dict(param.B_bins[i] => spzeros(param.nn+1,param.NN) for i in 1:length(param.B_bins))=#

	for bVal in param.B_bins

		NN2          = convert(Int64,ceil(param.NN*bVal))
		samples      = collect(1:(param.nn-1))
		pSize        = collect(0:NN2)
		samplesFreqs = permutedims(pSize/NN2)
		neutralSfs   = @. 1/pSize
		replace!(neutralSfs, Inf => 0.0)

		f(x,y=param.nn) = Binomial(y,x)
		z    = f.(samplesFreqs)

		out  = pdf.(z,samples)
		out  = round.(out,digits=10)
		outS = SparseArrays.dropzeros(SparseArrays.sparse(out))
		convolutedBn[bVal] = outS
		# param.bn[bVal] = outS
		#=param.neut[bVal] = round.(param.B*(param.thetaMidNeutral)*0.25*(outS*neutralSfs),digits=10)=#

	end
end

"""

	phiReduction(param,gamma,ppos)

Reduction in fixation probabilty due to background selection and linkage. The formulas used have been subjected to several theoretical works ([Charlesworth B., 1994](https://doi.org/10.1017/S0016672300032365), [Hudson et al., 1995](https://www.genetics.org/content/141/4/1605), [Nordborg et al. 1995](https://doi.org/10.1017/S0016672300033619), [Barton NH., 1995](https://www.genetics.org/content/140/2/821)).

The fixation probabilty of selected alleles are reduce by a factor ``\\phi``:

```math
\\phi(t,s) = e^{[\\frac{-2\\mu}{t(1+\\frac{rL}{t}+\\frac{2s}{t})}]}
```

Multiplying across all deleterious linkes sites, we find:

```math
\\Phi = \\prod_{1}^{L} = \\phi(t,s) = e^{\\frac{-2t\\mu(\\psi[1,\\frac{r+2s+L}{r}] - \\psi[1,\\frac{r(L+1)+2s+t}{r}])}{r^2}}
```

```math
\\phi(t,s) = e^{\\frac{-2t\\mu(\\psi[1,\\frac{r+2s+L}{r}] - \\psi[1,\\frac{r(L+1)+2s+t}{r}])}{r^2}}
```

# Arguments
 - `gamma::Int64`: selection coefficient.

# Returns
 - `Float64`: expected rate of positive fixations under background selection.

"""
function phiReduction(param::parameters,gammaValue::Int64)
	S::Float64 = abs(param.gamNeg/(1.0*param.NN))
	r::Float64 = param.rho/(2.0*param.NN)
	μ::Float64 = param.thetaF/(2.0*param.NN)
	s::Float64 = gammaValue/(param.NN*1.0)

	Ψ0::Float64 = polygamma(1,(s+S)/r)
	Ψ1::Float64 = polygamma(1,(r+param.Lf*r+s+S)/r)
	CC::Float64 = 1.0

	out::Float64 = (ℯ^(-2.0*S*μ*(Ψ0-Ψ1)/(r^2)))
	return out
end

"""
	analyticalAlpha(param, convolutedSamples)

Analytical α(x) estimation. Solve α(x) generally. We used the expected rates of divergence and polymorphism to approach the asympotic value accouting for background selection, weakly and strong positive selection. α(x) can be estimated taking into account the role of positive selected alleles or not. In this way we explore the role of linkage to deleterious alleles in the coding region.

```math
\\mathbb{E}[\\alpha_{x}] =  1 - \\left(\\frac{\\mathbb{E}[D_{s}]}{\\mathbb{E}[D_{N}]}\\frac{\\mathbb{E}[P_{N}]}{\\mathbb{E}[P_{S}]}\\right)
```

# Arguments
 - `param::parameters`
 - `convolutedSamples::binomialDict`
# Returns
 - `Array{Float64,1}` α(x).
"""
function analyticalAlpha(;param::parameters,convolutedSamples::binomialDict)

	################################################################
		# Solve the model similarly to original python mktest  #
	################################################################
	B = param.B

	setThetaF!(param)
	thetaF = param.thetaF
	# Solve the probabilities of fixations without BGS
	## Set non-bgs
	param.B = 0.999
	## Solve the mutation rate
	setThetaF!(param)
	## Solve the probabilities
	setPpos!(param)
	# Return to the original values
	param.thetaF = thetaF
	param.B = B

	################################################################
	 # Accounting for positive alleles segregating due to linkage #
	################################################################

	# Fixation
	fN     = param.B*fixNeut(param)
	fNeg   = param.B*fixNegB(param,0.5*param.pposH+0.5*param.pposL)
	fPosL  = fixPosSim(param,param.gL,0.5*param.pposL)
	fPosH  = fixPosSim(param,param.gH,0.5*param.pposH)

	ds = fN
	dn = fNeg + fPosL + fPosH

	## Polymorphism
	neut = DiscSFSNeutDown(param,convolutedSamples.bn[param.B])

	selH::Array{Float64,1} = if isinf(exp(param.gH * 2))
		DiscSFSSelPosDownArb(param,param.gH,param.pposH,convolutedSamples.bn[param.B])
	else
		DiscSFSSelPosDown(param,param.gH,param.pposH,convolutedSamples.bn[param.B])
	end

	selL::Array{Float64,1} = DiscSFSSelPosDown(param,param.gL,param.pposL,convolutedSamples.bn[param.B])
	selN::Array{Float64,1} = DiscSFSSelNegDown(param,param.pposH+param.pposL,convolutedSamples.bn[param.B])

	splitColumns(matrix::Array{Float64,2}) = (view(matrix, :, i) for i in 1:size(matrix, 2));
	tmp = cumulativeSfs(hcat(neut,selH,selL,selN),false)

	neut, selH, selL, selN = splitColumns(tmp)
	sel = (selH+selL)+selN

	# ps = @. neut / (sel+neut)
	# pn = @. sel / (sel+neut)

	## Outputs
	α = @. 1 - ((ds/dn) * (sel/neut))

	#############################################################
	#		Accounting for neutral and deleterious alleles		#	#############################################################
	## Fixation
	fN_nopos     = fN*(param.thetaMidNeutral/2.)*param.TE*param.NN
	fNeg_nopos   = fNeg*(param.thetaMidNeutral/2.)*param.TE*param.NN
	fPosL_nopos  = fPosL*(param.thetaMidNeutral/2.)*param.TE*param.NN
	fPosH_nopos  = fPosH*(param.thetaMidNeutral/2.)*param.TE*param.NN

	ds_nopos = fN_nopos
	dn_nopos = fNeg_nopos + fPosL_nopos + fPosH_nopos

	## Polymorphism
	sel_nopos = selN
	ps_nopos = @. neut / (sel_nopos + neut)
	pn_nopos = @. sel_nopos / (sel_nopos + neut)

	α_nopos = 1 .- (ds_nopos/dn_nopos) .* (sel_nopos./neut)

	##########
	# Output #
	##########
	return (α,α_nopos)
end


"""
    NonCommonVertices()

    Given the indices of vertices of a source simplex inside the circumsphere of a target simplex,
    find the indices of the vertices that are not shared between simplices, but still lie
    inside the circumference of the target simplex.

Input arguments
---------------
IndexIn::Array{Int, 1}              Row vector. Indices of the vertices of the source simplex inside the circumsphere of the target simplex.
InternalComIndex::Array{Int, 1}     Row vector. Indices of the vertices that both simplices share, but in the 'IndexIn' row vector of indices instead of the 'IndexComVert' vector.
Nin::Int                            Integer. Number of vertices inside the circumsphere of the target simplex. REDUNDANT? Essentially, length(IndexIn) - but remember to check for zeros
Ncomm::Int                          Integer. Number of common vertices between between source and target simplices.

Returns
-------
Res                                 Row vector. Indices of the vertices of the source simplex that are inside the circumsphere of the target simplex, but is not shared between the simplices.
"""

function NonCommonVertices(InternalComIndex,IndexIn,
                          n_vertices_in_circumsphere,
                          n_commonvertices)
  # IndexComVert: indices of the common vertices
  # IndexIn: indices of the vertices contained in the circumsphere
  # InternalComIndex: Indices of IndexComVert in IndexIn
  # Nin: number of vertices in the circumsphere

  # OUTCOME
  @show n_vertices_in_circumsphere
  @show n_commonvertices
  @show IndexIn
  @show InternalComIndex
  # Res: indices of the non common vertices still inside the circumsphere
  # if Res is empty, it returns 0
  if n_commonvertices == 0
    #println("n_commonvertices == 0")
    return IndexIn
  else
      if n_commonvertices == n_vertices_in_circumsphere
        #println("Number of vertices in circumsphere equals number of common vertices")
        return [0]
      else
        #println("Number of vertices in circumsphere does not equal number of common vertices")
        return IndexIn[complementary(InternalComIndex, n_vertices_in_circumsphere)]
     end
 end

end


instrs = split(strip(read("input.txt", String)), "\n")

function interpret(instrs)
    acc = 0
    ip = 1
    visited = Set()

    while ip <= length(instrs)
        if ip in visited
            return (acc, false)
        end

        push!(visited, ip)
        instr = instrs[ip]

        if startswith(instr, "acc")
            acc += parse(Int, instr[5:end])
        end

        if startswith(instr, "jmp")
            ip += parse(Int, instr[5:end])
        else
            ip += 1
        end
    end

    (acc, true)
end

interpret(instrs)[1]

function correct(instrs)
    for i in eachindex(instrs)
        instr = instrs[i]
        if startswith(instr, "acc")
            continue
        end
        newinstr = (startswith(instr, "jmp") ? "nop" : "jmp") * instr[4:end]
        newinstrs = deepcopy(instrs)
        newinstrs[i] = newinstr
        result = interpret(newinstrs)
        if result[2]
            return result[1]
        end
    end
end

correct(instrs)


# ## Basis Set Data
# Note use of curly braces here. Julia assumes that if you have square braces, you want
# things flattened as much as possible (to be as fast as possible, I guess). Curlys
# preserve the list structure the way I would expect from Python

# ## Basis function definitions

type PGBF
  expn::Float64
  x::Float64
  y::Float64
  z::Float64
  I::Int64
  J::Int64
  K::Int64
  norm::Float64
end

# Enable keyword-based
function pgbf(expn; x=0,y=0,z=0,I=0,J=0,K=0,norm=1)
  p = PGBF(expn,x,y,z,I,J,K,norm)
  normalize!(p)
  return p
end

function pgbf(expn, x=0,y=0,z=0,I=0,J=0,K=0,norm=1)
  p = PGBF(expn,x,y,z,I,J,K,norm)
  normalize!(p)
  return p
end

function amplitude(bf::PGBF,x,y,z)
  dx,dy,dz = x-bf.x,y-bf.y,z-bf.z
  r2 = dist2(dx,dy,dz)
  return bf.norm*(dx^bf.I)*(dy^bf.J)*(dz^bf.K)*exp(-bf.expn*r2)
end

function normalize!(pbf::PGBF)
  pbf.norm /= sqrt(overlap(pbf,pbf))
end

type CGBF
  x::Float64
  y::Float64
  z::Float64
  I::Int64
  J::Int64
  K::Int64
  norm::Float64
  pgbfs::Array{PGBF,1}
  coefs::Array{Float64,1}
end

cgbf(x=0,y=0,z=0,I=0,J=0,K=0) = CGBF(x,y,z,I,J,K,1.0,PGBF[],Float64[])

function amplitude(bf::CGBF,x,y,z)
  s = 0
  for (c,pbf) in primitives(bf)
    s += c*amplitude(pbf,x,y,z)
  end
  return bf.norm*s
end

function normalize!(bf::CGBF)
  bf.norm /= sqrt(overlap(bf,bf))
end

primitives(a::CGBF) = zip(a.coefs,a.pgbfs)

function push!(cbf::CGBF,expn,coef)
  Base.push!(cbf.pgbfs,pgbf(expn,cbf.x,cbf.y,cbf.z,cbf.I,cbf.J,cbf.K))
  Base.push!(cbf.coefs,coef)
  normalize!(cbf)
end

function contract(f,a::CGBF,b::CGBF)
  s = 0
  for (ca,abf) in primitives(a)
    for (cb,bbf) in primitives(b)
      s += ca*cb*f(abf,bbf)
    end
  end
  return a.norm*b.norm*s
end

function contract(f,a::CGBF,b::CGBF,c::CGBF,d::CGBF)
  s = 0
  for (ca,abf) in primitives(a)
    for (cb,bbf) in primitives(b)
      for (cc,cbf) in primitives(c)
        for (cd,dbf) in primitives(d)
          s += ca*cb*cc*cd*f(abf,bbf,cbf,dbf)
        end
      end
    end
  end
  return a.norm*b.norm*c.norm*d.norm*s
end