text:
importall MathProgBase.SolverInterface
import Base.convert

# --------------------------------------------------------------------------------
# Begin implementation of the MPB low-level interface
# Implements
# - ConicModel
# - loadproblem!
# - optimize!
# - status
# - numvar
# - numconstr
# http://mathprogbasejl.readthedocs.org/en/latest/solverinterface.html
#

immutable ConicIPSolver <: AbstractMathProgSolver
  preprocess
  options
end

ConicIPSolver(;equalities_as_double_inequalities = false,
                preprocess = true,
                kwargs...) =
                ConicIPSolver( preprocess,
                                kwargs )

"""
Type which encapsulates a ConicIP model, and a way to convert
its output into a MPB model

minimize    ½yᵀQy - cᵀy
s.t         Ay >= b
            Gy  = d
"""
type ConicIPModel <: AbstractConicModel

    # Parameters to be passed into ConicIP

    Q           ::AbstractMatrix
    c           ::AbstractMatrix
    A           ::AbstractMatrix
    b           ::AbstractMatrix
    cone_dims
    G           ::AbstractMatrix
    d           ::AbstractMatrix

    # Solutions

    solve_stat
    sol
    obj_val     ::Float64
    primal_sol  ::Vector{Float64}
    dual_sol    ::Vector{Float64}
    vardual_sol ::Vector{Float64}
    slack       ::Float64

    # Bookkeeping for indicies, so we can figure out what
    # the dual variables are.

    I_Al        ::Vector{Int64}
    I_Gl        ::Vector{Int64}
    I_A         ::Vector{Int64}
    I_G         ::Vector{Int64}

    I_vAl       ::Vector{Int64}
    I_vGl       ::Vector{Int64}
    I_vA        ::Vector{Int64}
    I_vG        ::Vector{Int64}

    # Number of constraints and varconstraints
    n_constr
    n_varconstr

    preprocess  ::Bool
    options

    nA          ::Float64

end

"""
  signed_assign!(x, y, I, J)

signed assignment, a hack to pass in sign information through
indices.

x[I] =  y[J]  when J .> 0
x[I] = -y[-J] when J .< 0
"""
function signed_assign!(x, y, I, J)
  I = abs(I)
  P = J .> 0
  x[I[P]] = y[J[P]]
  x[I[!P]] = -y[-J[!P]]
end

ConicModel(s::ConicIPSolver) = ConicIPModel(
  zeros(0,0),
  zeros(0,0),
  zeros(0,0),
  zeros(0,0),
  zeros(0,0),
  zeros(0,0),
  zeros(0,0),
  nothing,
  nothing,
  0,
  zeros(0),
  zeros(0),
  zeros(0),
  0,
  round(Int, zeros(0)),
  round(Int, zeros(0)),
  round(Int, zeros(0)),
  round(Int, zeros(0)),
  round(Int, zeros(0)),
  round(Int, zeros(0)),
  round(Int, zeros(0)),
  round(Int, zeros(0)),
  0,
  0,
  s.preprocess,
  s.options,
  NaN)

LinearQuadraticModel(s::ConicIPSolver)  = ConicToLPQPBridge(ConicModel(s))
status(m::ConicIPModel)                = m.solve_stat
getobjval(m::ConicIPModel)             = m.obj_val
getsolution(m::ConicIPModel)           = copy(m.primal_sol)
numvar(m::ConicIPModel)                = length(m.c)
numconstr(m::ConicIPModel)             = m.n_constr + m.n_varconstr
supportedcones(s::ConicIPSolver)        = [:Free,
                                           :Zero,
                                           :NonNeg,
                                           :NonPos,
                                           :SOC,
                                           :SDP]

function show(m::ConicIPModel)

  println("Q = ",full(m.Q))
  println("c = ",m.c)
  println("A = ",full(m.A))
  println("b = ",m.b)
  println("cone_dims = ", m.cone_dims)
  println("G = ",full(m.G))
  println("d = ",m.d)

end

function loadproblem!(m::ConicIPModel, c,
  A_, b_, constr_cones, var_cones)

  n = size(A_, 2)

  A_ = sparse(A_)

  A = spzeros(0,n)
  b = zeros(0,1)
  G = spzeros(0,n)
  d = zeros(0,1)

  I_A   = round(Int, zeros(0))
  I_G   = round(Int, zeros(0))
  I_Al  = round(Int, zeros(0))
  I_Gl  = round(Int, zeros(0))

  I_vA  = round(Int, zeros(0))
  I_vG  = round(Int, zeros(0))
  I_vAl = round(Int, zeros(0))
  I_vGl = round(Int, zeros(0))

  n_constr    = sum([length(ind) for (cone, ind) = constr_cones])
  n_varconstr = sum([length(ind) for (cone, ind) = var_cones])

  cone_dims = []

  Gs = 0
  As = 0

  for (cone, ind) = constr_cones

    if typeof(ind) <: Integer
      ind = ind:ind
    end

    if cone == :Zero
      i = Gs; Gs = Gs + length(ind); j = Gs
      d = [ d ; b_[ind,:] ]
      append!( I_G  , ind   )
      append!( I_Gl , i+1:j )
    end

    if cone == :NonPos
      i = As; As = As + length(ind); j = As
      push!( cone_dims, ("R", length(ind)) )
      b = [b ; b_[ind,:] ]
      append!( I_A  , ind      )
      append!( I_Al , -(i+1:j) )
    end

    if cone == :NonNeg
      i = As; As = As + length(ind); j = As
      push!( cone_dims, ("R", length(ind)) )
      b = [ b ; -b_[ind,:] ]
      append!( I_A  , -ind   )
      append!( I_Al , i+1:j )
    end

    if cone == :SOC
      i = As; As = As + length(ind); j = As
      push!( cone_dims, ("Q", length(ind)) )
      b = [ b ; -b_[ind,:] ]
      append!( I_A  , -ind   )
      append!( I_Al , i+1:j )
    end

    if cone == :SDP
      i = As; As = As + length(ind); j = As
      push!( cone_dims, ("S", length(ind)) )
      b = [ b ; -b_[ind,:] ]
      append!( I_A  , -ind   )
      append!( I_Al , i+1:j )
    end

  end

  I = speye(n,n)

  cum_free_dim = 0

  base_sise_A = size(A,1)
  base_size_G = size(G,1)

  for (cone, ind) = var_cones

    if typeof(ind) <: Integer
      ind = ind:ind
    end

    if cone == :Free
      cum_free_dim = cum_free_dim + length(ind)
    end

    if cone == :Zero
      i = Gs; Gs = Gs + length(ind); j = Gs
      d = [d ; zeros(length(ind), 1) ]
      append!( I_vG  , ind     )
      append!( I_vGl , (i+1:j) )
    end

    if cone == :NonPos
      i = As; As = As + length(ind); j = As
      push!( cone_dims, ("R", length(ind)) )
      b = [ b ; zeros(length(ind), 1) ]
      append!( I_vA  , -ind      )
      append!( I_vAl , -(i+1:j) )
    end

    if cone == :NonNeg
      i = As; As = As + length(ind); j = As
      push!( cone_dims, ("R", length(ind)) )
      b = [ b ; zeros(length(ind), 1) ]
      append!( I_vA  , ind     )
      append!( I_vAl , (i+1:j) )
    end

    if cone == :SOC
      i = As; As = As + length(ind); j = As
      push!( cone_dims, ("Q", length(ind)) )
      b = [ b ; zeros(length(ind), 1) ]
      append!( I_vA  , ind     )
      append!( I_vAl , (i+1:j) )
    end

    if cone == :SDP
      i = As; As = As + length(ind); j = As
      push!( cone_dims, ("S", length(ind)) )
      b = [ b ; zeros(length(ind), 1) ]
      append!( I_vA  , ind     )
      append!( I_vAl , (i+1:j) )
    end

  end

  nA = vecnorm(A_)

  Λ = sign([I_A;I_vA])
  A = Λ.*[A_[abs(I_A),:]; nA*I[abs(I_vA),:]]

  Λ = sign([I_G;I_vG])
  G = Λ.*[A_[abs(I_G),:]; nA*I[abs(I_vG),:]]

  m.Q           = spzeros(n,n)
  m.c           = -c''
  m.A           = A
  m.b           = b
  m.cone_dims   = cone_dims
  m.G           = G
  m.d           = d
  m.I_A         = I_A
  m.I_Al        = I_Al
  m.I_G         = I_G
  m.I_Gl        = I_Gl
  m.I_vA        = I_vA
  m.I_vAl       = I_vAl
  m.I_vG        = I_vG
  m.I_vGl       = I_vGl
  m.n_constr    = n_constr
  m.n_varconstr = n_varconstr
  m.nA          = nA

end

function optimize!(m::ConicIPModel)

  if m.preprocess
    m.sol = preprocess_conicIP(m.Q,m.c,m.A,m.b,m.cone_dims,m.G,m.d; m.options...)
  else
    m.sol = conicIP(m.Q,m.c,m.A,m.b,m.cone_dims,m.G,m.d; m.options...)
  end

  m.solve_stat = m.sol.status

  m.obj_val = -vecdot(m.sol.y,m.c)
  m.primal_sol = m.sol.y[:]

  m.dual_sol = zeros(m.n_constr)
  signed_assign!(m.dual_sol, m.sol.v, m.I_A, m.I_Al)
  signed_assign!(m.dual_sol, m.sol.w, m.I_G, m.I_Gl)

  m.vardual_sol = zeros(m.n_varconstr)
  signed_assign!(m.vardual_sol, m.sol.v*m.nA, m.I_vA, m.I_vAl)
  signed_assign!(m.vardual_sol, m.sol.w*m.nA, m.I_vG, m.I_vGl)

end

getdual(m::ConicIPModel)    = m.dual_sol
getvardual(m::ConicIPModel) = m.vardual_sol


text:
using ITensors
include("../networks/lattices.jl")

struct LocalMPO
  mpo::MPO
  coord1::Tuple{<:Integer,<:Integer}
  coord2::Tuple{<:Integer,<:Integer}
end

struct LineMPO
  mpo::MPO
  coord::Union{Tuple{Colon,<:Integer},Tuple{<:Integer,Colon}}
end

# Transverse field
# The critical point is h = 1.0
# This is the most challenging part of the model for DMRG
function mpo(::Model"tfim", sites::Matrix{<:Index}; h::Float64)
  Ny, Nx = size(sites)
  sites_vec = reshape(sites, Nx * Ny)
  lattice = square_lattice(Nx, Ny; yperiodic=false)

  opsum = OpSum()
  for b in lattice
    opsum += -1, "X", b.s1, "X", b.s2
  end
  for i in 1:(Nx * Ny)
    opsum += h, "Z", i
  end
  return MPO(opsum, sites_vec)
end

function localham_term(
  ::Model"tfim",
  sites::Matrix{<:Index},
  bond::Tuple{Tuple{<:Integer,<:Integer},Tuple{<:Integer,<:Integer}};
  h::Float64,
)
  Ny, Nx = size(sites)
  coord1, coord2 = bond
  opsum = OpSum()
  opsum += -1, "X", 1, "X", 2
  if coord2[1] == coord1[1] + 1
    opsum += h, "Z", 1
  end
  if coord2[1] == coord1[1] + 1 && coord2[1] == Ny
    opsum += h, "Z", 2
  end
  mpo = MPO(opsum, [sites[coord1...], sites[coord2...]])
  return LocalMPO(mpo, coord1, coord2)
end

function localham(m::Model, sites; kwargs...)
  Ny, Nx = size(sites)
  lattice = Square((Ny, Nx))
  bds = bonds(lattice; periodic=false)
  return [localham_term(m, sites, bond; kwargs...) for bond in bds]
end

function localham_term(
  ::Model"tfim", sites::Matrix{<:Index}, bond::Tuple{Colon,<:Integer}; h::Float64
)
  Ny, Nx = size(sites)
  opsum = OpSum()
  for i in 1:(Ny - 1)
    opsum += -1, "X", i, "X", i + 1
    opsum += h, "Z", i
  end
  opsum += h, "Z", Ny
  return LineMPO(MPO(opsum, sites[bond...]), bond)
end

function localham_term(
  ::Model"tfim", sites::Matrix{<:Index}, bond::Tuple{<:Integer,Colon}; h::Float64
)
  Ny, Nx = size(sites)
  opsum = OpSum()
  for i in 1:(Nx - 1)
    opsum += -1, "X", i, "X", i + 1
  end
  return LineMPO(MPO(opsum, sites[bond...]), bond)
end

function lineham(m::Model, sites; kwargs...)
  Ny, Nx = size(sites)
  lattice = Square((Ny, Nx))
  bonds_row = [(i, :) for i in 1:Ny]
  bonds_column = [(:, i) for i in 1:Nx]
  bds = vcat(bonds_row, bonds_column)
  return [localham_term(m, sites, bond; kwargs...) for bond in bds]
end

# Check that the local Hamiltonian is the same as the MPO
function checkham(Hlocal::Array{LocalMPO}, H, sites)
  @disable_warn_order begin
    Ny, Nx = size(sites)
    lattice = Square((Ny, Nx))
    bds = bonds(lattice; periodic=false)
    Hlocal_full = ITensor()
    for (i, bond) in enumerate(bds)
      Hlocalterm_full = prod(Hlocal[i].mpo)
      for y in 1:Ny
        for x in 1:Nx
          if !((y, x) in bond)
            Hlocalterm_full *= op("Id", vec(sites), (x - 1) * Ny + y)
          end
        end
      end
      Hlocal_full += Hlocalterm_full
    end
    @show norm(Hlocal_full - prod(H))
  end
  return isapprox(norm(Hlocal_full), norm(prod(H)))
end

function checkham(Hline::Array{LineMPO}, H, sites)
  @disable_warn_order begin
    Ny, Nx = size(sites)
    Hlocal_full = ITensor()
    for h in Hline
      h_full = prod(h.mpo)
      if h.coord[1] isa Colon
        for y in 1:Ny
          for x in 1:Nx
            if x != h.coord[2]
              h_full *= op("Id", vec(sites), (x - 1) * Ny + y)
            end
          end
        end
      else
        for y in 1:Ny
          for x in 1:Nx
            if y != h.coord[1]
              h_full *= op("Id", vec(sites), (x - 1) * Ny + y)
            end
          end
        end
      end
      Hlocal_full += h_full
    end
    @show norm(Hlocal_full - prod(H))
  end
  return isapprox(norm(Hlocal_full), norm(prod(H)))
end

ext: jl


text:
# Performance evaluation

## correctrate & errorrate

correctrate(gt::IntegerVector, r::IntegerVector) = counteq(gt, r) / length(gt)
errorrate(gt::IntegerVector, r::IntegerVector) = countne(gt, r) / length(gt)

## confusion matrix

function confusmat(k::Integer, gts::IntegerVector, preds::IntegerVector)
    n = length(gts)
    length(preds) == n || throw(DimensionMismatch("Inconsistent lengths."))
    R = zeros(Int, k, k)
    for i = 1:n
        @inbounds g = gts[i]
        @inbounds p = preds[i]
        R[g, p] += 1
    end
    return R
end

## counthits & hitrate

function counthits(gt::IntegerVector, rklst::IntegerMatrix, k::Integer)
    n = length(gt)
    size(rklst, 2) == n || throw(DimensionMismatch("Input dimensions mismatch."))
    m = min(size(rklst, 1), @compat(Int(k)))

    cnt = 0
    @inbounds for j = 1:n
        rj = view(rklst, :, j)
        gj = gt[j]
        for i = 1:m
            if rj[i] == gj
                cnt += 1
                break
            end
        end
    end
    return cnt::Int
end

function counthits(gt::IntegerVector, rklst::IntegerMatrix, ks::IntegerVector)
    n = length(gt)
    size(rklst, 2) == n || throw(DimensionMismatch("Input dimensions mismatch."))
    issorted(ks) || throw(DimensionMismatch("ks must be sorted."))

    m = min(size(rklst, 1), ks[end])
    nk = length(ks)
    cnts = zeros(Int, nk)
    @inbounds for j = 1:n
        rj = view(rklst, :, j)
        gj = gt[j]
        for i = 1:m
            if rj[i] == gj
                ik = 1
                while ks[ik] < i; ik += 1; end
                while ik <= nk
                    cnts[ik] += 1
                    ik += 1
                end
                break
            end
        end
    end
    return cnts
end

hitrate(gt::IntegerVector, rklst::IntegerMatrix, k::Integer) =
    (counthits(gt, rklst, k) / length(gt))::Float64

function hitrates(gt::IntegerVector, rklst::IntegerMatrix, ks::IntegerVector)
    n = length(gt)
    h = counthits(gt, rklst, ks)
    nk = length(ks)
    r = Array(Float64, nk)
    for i = 1:nk
        r[i] = h[i] / n
    end
    return r
end

## ROC

immutable ROCNums{T<:Real}
    p::T    # positive in ground-truth
    n::T    # negative in ground-truth
    tp::T   # correct positive prediction
    tn::T   # correct negative prediction
    fp::T   # (incorrect) positive prediction when ground-truth is negative
    fn::T   # (incorrect) negative prediction when ground-truth is positive
end

function show(io::IO, x::ROCNums)
    println(io, "$(typeof(x))")
    println(io, "  p = $(x.p)")
    println(io, "  n = $(x.n)")
    println(io, "  tp = $(x.tp)")
    println(io, "  tn = $(x.tn)")
    println(io, "  fp = $(x.fp)")
    println(io, "  fn = $(x.fn)")
end

true_positive(x::ROCNums) = x.tp
true_negative(x::ROCNums) = x.tn
false_positive(x::ROCNums) = x.fp
false_negative(x::ROCNums) = x.fn

true_positive_rate(x::ROCNums) = x.tp / x.p
true_negative_rate(x::ROCNums) = x.tn / x.n
false_positive_rate(x::ROCNums) = x.fp / x.n
false_negative_rate(x::ROCNums) = x.fn / x.p

recall(x::ROCNums) = true_positive_rate(x)
sensitivity(x::ROCNums) = true_positive_rate(x)
specificity(x::ROCNums) = true_negative_rate(x)
precision(x::ROCNums) = x.tp / (x.tp + x.fp)
positive_predictive_value(x::ROCNums) = precision(x)
negative_predictive_value(x::ROCNums) = x.tn / (x.tn + x.fn)
false_discovery_rate(x::ROCNums) = x.fp / (x.tp + x.fp)

accuracy(x::ROCNums) = (x.tp + x.tn) / (x.p + x.n)
f1score(x::ROCNums) = (tp2 = x.tp + x.tp; tp2 / (tp2 + x.fp + x.fn) )

## roc: produce a single instance of ROCNums from predictions

_ispos(x::Bool) = x
_ispos(x::Real) = x > zero(x)

function _roc(gt::IntegerVector, pr)
    len = length(gt)
    length(pr) == len || throw(DimensionMismatch("Inconsistent lengths."))

    p = 0
    n = 0
    tp = 0
    tn = 0
    fp = 0
    fn = 0

    for i = 1:len
        @inbounds gi = gt[i]
        @inbounds ri = pr[i]
        if _ispos(gi)   # gt = true
            p += 1
            if ri == gi
                tp += 1
            elseif ri == 0
                fn += 1
            end
        else        # gt = false
            n += 1
            if _ispos(ri)
                fp += 1
            else
                tn += 1
            end
        end
    end

    return ROCNums{Int}(p, n, tp, tn, fp, fn)
end

# compute roc numbers based on prediction
roc(gt::IntegerVector, pr::IntegerVector) = _roc(gt, pr)

##
#   BinaryThresPredVec immutates a vector:
#
#   v[i] := scores[i] < thres ? 0 : 1
#
immutable BinaryThresPredVec{ScoreVec <: RealVector,
                             T <: Real,
                             Ord <: Ordering}
    scores::ScoreVec
    thres::T
    ord::Ord
end

length(v::BinaryThresPredVec) = length(v.scores)
getindex(v::BinaryThresPredVec, i::Integer) = !lt(v.ord, v.scores[i], v.thres)

# compute roc numbers based on scores & threshold
roc(gt::IntegerVector, scores::RealVector, t::Real, ord::Ordering) =
    _roc(gt, BinaryThresPredVec(scores, t, ord))

roc(gt::IntegerVector, scores::RealVector, thres::Real) =
    roc(gt, scores, thres, Forward)

##
#   ThresPredVec immutates a vector:
#
#   v[i] := scores[i] < thres ? 0 : preds[i]
#
immutable ThresPredVec{PredVec <: IntegerVector,
                       ScoreVec <: RealVector,
                       T <: Real,
                       Ord <: Ordering}

    preds::PredVec
    scores::ScoreVec
    thres::T
    ord::Ordering
end

function ThresPredVec{PVec<:IntegerVector,SVec<:RealVector,T<:Real,Ord<:Ordering}(
    preds::PVec, scores::SVec, thres::T, ord::Ord)
    n = length(preds)
    length(scores) == n || throw(DimensionMismatch("Inconsistent lengths."))
    ThresPredVec{PVec,SVec,T,Ord}(preds, scores, thres, ord)
end

length(v::ThresPredVec) = length(v.preds)
getindex(v::ThresPredVec, i::Integer) = ifelse(lt(v.ord, v.scores[i], v.thres), 0, v.preds[i])

# compute roc numbers based on predictions & scores & threshold
roc{PV<:IntegerVector,SV<:RealVector}(gt::IntegerVector, preds::@compat(Tuple{PV,SV}), t::Real, ord::Ordering) =
    _roc(gt, ThresPredVec(preds..., t, ord))

roc{PV<:IntegerVector,SV<:RealVector}(gt::IntegerVector, preds::@compat(Tuple{PV,SV}), thres::Real) =
    roc(gt, preds, thres, Forward)

## roc: produces a series of ROCNums instances with multiple thresholds

# find_thresbin
#
#  x < threshold[1] --> 1
#  threshold[i] <= x < threshold[i+1] --> i+1
#  x >= threshold[n] --> n+1
#
function find_thresbin(x::Real, thresholds::RealVector, ord::Ordering)
    n = length(thresholds)
    r = 1
    if !lt(ord, x, thresholds[1])
        l = 1
        r = n + 1
        while l + 1 < r
            m = (l + r) >> 1
            if lt(ord, x, thresholds[m])
                r = m
            else
                l = m
            end
        end
    end
    return r::Int
end

find_thresbin(x::Real, thresholds::RealVector) = find_thresbin(x, thresholds, Forward)

lin_thresholds(scores::RealArray, n::Integer, ord::ForwardOrdering) =
    ((s0, s1) = extrema(scores); intv = (s1 - s0) / (n-1); s0:intv:s1)

lin_thresholds(scores::RealArray, n::Integer, ord::ReverseOrdering{ForwardOrdering}) =
    ((s0, s1) = extrema(scores); intv = (s0 - s1) / (n-1); s1:intv:s0)

# roc for binary predictions
function roc(gt::IntegerVector, scores::RealVector, thresholds::RealVector, ord::Ordering)
    issorted(thresholds, ord) || error("thresholds must be sorted w.r.t. the given ordering.")

    ns = length(scores)
    nt = length(thresholds)

    # scan scores and classify them into bins
    hp = zeros(Int, nt + 1)
    hn = zeros(Int, nt + 1)
    p = 0
    n = 0
    for i = 1:ns
        @inbounds s = scores[i]
        @inbounds g = gt[i]
        k = find_thresbin(s, thresholds, ord)
        if _ispos(g)
            hp[k] += 1
            p += 1
        else
            hn[k] += 1
            n += 1
        end
    end

    # produce results
    r = Array(ROCNums{Int}, nt)
    fn = 0
    tn = 0
    @inbounds for i = 1:nt
        fn += hp[i]
        tn += hn[i]
        tp = p - fn
        fp = n - tn
        r[i] = ROCNums{Int}(p, n, tp, tn, fp, fn)
    end
    return r
end

roc(gt::IntegerVector, scores::RealVector, thresholds::RealVector) = roc(gt, scores, thresholds, Forward)

roc(gt::IntegerVector, scores::RealVector, n::Integer, ord::Ordering) =
    roc(gt, scores, lin_thresholds(scores, n, ord), ord)

roc(gt::IntegerVector, scores::RealVector, n::Integer) = roc(gt, scores, n, Forward)

roc(gt::IntegerVector, scores::RealVector, ord::Ordering) = roc(gt, scores, 100, ord)
roc(gt::IntegerVector, scores::RealVector) = roc(gt, scores, Forward)

# roc for multi-way predictions
function roc{PV<:IntegerVector,SV<:RealVector}(
    gt::IntegerVector, preds::@compat(Tuple{PV,SV}), thresholds::RealVector, ord::Ordering)

    issorted(thresholds, ord) || error("thresholds must be sorted w.r.t. the given ordering.")
    pr::PV = preds[1]
    scores::SV = preds[2]

    ns = length(scores)
    nt = length(thresholds)

    # scan scores and classify them into bins
    hp = zeros(Int, nt + 1)
    hn = zeros(Int, nt + 1)
    htp = zeros(Int, nt + 1)

    p = 0
    n = 0
    tp = 0

    for i = 1:ns
        @inbounds pi = pr[i]
        @inbounds si = scores[i]
        @inbounds gi = gt[i]
        k = find_thresbin(si, thresholds, ord)
        if _ispos(gi)
            hp[k] += 1
            p += 1

            if pi == gi
                htp[k] += 1
                tp += 1
            end
        else
            hn[k] += 1
            n += 1
        end
    end

    # produce results
    r = Array(ROCNums{Int}, nt)
    fn = 0
    tn = 0
    @inbounds for i = 1:nt
        fn += hp[i]
        tn += hn[i]
        tp -= htp[i]
        fp = n - tn
        r[i] = ROCNums{Int}(p, n, tp, tn, fp, fn)
    end
    return r
end

roc{PV<:IntegerVector, SV<:RealVector}(gt::IntegerVector, preds::@compat(Tuple{PV,SV}), thresholds::RealVector) =
    roc(gt, preds, thresholds, Forward)

roc{PV<:IntegerVector, SV<:RealVector}(gt::IntegerVector, preds::@compat(Tuple{PV,SV}), n::Integer, ord::Ordering) =
    roc(gt, preds, lin_thresholds(preds[2],n,ord), ord)

roc{PV<:IntegerVector, SV<:RealVector}(gt::IntegerVector, preds::@compat(Tuple{PV,SV}), n::Integer) =
    roc(gt, preds, n, Forward)

roc{PV<:IntegerVector, SV<:RealVector}(gt::IntegerVector, preds::@compat(Tuple{PV,SV}), ord::Ordering) =
    roc(gt, preds, 100, ord)

roc{PV<:IntegerVector, SV<:RealVector}(gt::IntegerVector, preds::@compat(Tuple{PV,SV})) =
    roc(gt, preds, Forward)


text:
module PkgTemplates

using DocStringExtensions
using Pkg

export Template,
    generate

include("gitignore.jl")
include("runtests.jl")
include("runtests_gpu.jl")
include("formatter.jl")

include("azure.jl")
include("travis.jl")
include("appveyor.jl")
include("github_docker_ci.jl")
include("github_ci.jl")
include("ci.jl")

include("github_actions.jl")

include("bors.jl")

"""
    Template

# Fields
$(DocStringExtensions.FIELDS)
"""
Base.@kwdef struct Template
    "Organization name"
    org::String="charleskawczynski"
    "Directory where package is located"
    directory::String=""
    "Description"
    description::String="Package description goes here"
    "Add GPU test environment and gpu test runtests"
    gpu::Bool=true
    "Package name"
    pkg_name::String=""
    "Julia version"
    julia_version::String
    "Package name without the extension"
    pkg_name_no_ext::String=replace(pkg_name, ".jl" => "")
    "Package root directory"
    pkg_root::String=joinpath(directory, pkg_name)

    # CI
    "Continuous Integration"
    ci::ContinuousIntegration = ContinuousIntegration(
        azure=Azure(;org=org, pkg_name=pkg_name, julia_version=julia_version),
        travis=Travis(;julia_version=julia_version),
        appveyor=AppVeyor(;julia_version=julia_version),
        # github_ci=GitHubCI(;org=org, pkg_name=pkg_name, julia_version=julia_version),
        # github_docker_ci=GitHubDockerCI(;org=org, pkg_name=pkg_name, julia_version=julia_version),
    )

    # GitHub Actions
    "Doc build"
    doc_build::DocBuild = DocBuild("docs-build", julia_version)
    "TagBot"
    tag_bot::TagBot = TagBot()
    "CompatHelper"
    compat_helper::CompatHelper = CompatHelper(julia_version)
    "Code-coverage configuration"
    code_cov::CodeCov = CodeCov(pkg_name, julia_version)
    "Formatter Check"
    formatter_check::FormatterCheck = FormatterCheck(julia_version)

    "Bors configuration"
    bors::Bors = Bors(;checks=[ci_checks(ci)..., doc_build.build_name])
end

include("readme.jl")
include("docs_make.jl")
include("index_md.jl")

function write_contents_to_file(file, contents)
  open(file, "w") do io
    print(io, contents)
  end
end

function generate(t)

    # Use Pkg.generate
    cd(t.directory) do
      Pkg.generate(t.pkg_name_no_ext)
      # rename folder if ".jl" is in package name
      mv(t.pkg_name_no_ext, t.pkg_name)
    end

    mkpath(joinpath(t.pkg_root, "test"))

    mkpath(joinpath(t.pkg_root, "docs"))
    mkpath(joinpath(t.pkg_root, "docs", "src"))

    mkpath(joinpath(t.pkg_root, "env"))
    mkpath(joinpath(t.pkg_root, "env", "test"))
    mkpath(joinpath(t.pkg_root, "env", "viz"))

    if t.gpu
        env_gpu = joinpath(@__DIR__, "env_gpu.jl")
        cd(t.pkg_root) do
            run(`julia --project=env/gpu/ $(env_gpu)`)
        end
    end
    env_test = joinpath(@__DIR__, "env_test.jl")
    cd(t.pkg_root) do
        run(`julia --project=env/test/ $(env_test)`)
    end

    env_docs = joinpath(@__DIR__, "env_docs.jl")
    cd(t.pkg_root) do
        run(`julia --project=docs/ $(env_docs)`)
    end

    env_viz = joinpath(@__DIR__, "env_viz.jl")
    cd(t.pkg_root) do
        run(`julia --project=env/viz/ $(env_viz)`)
    end

    mkpath(joinpath(t.pkg_root, ".github"))
    mkpath(joinpath(t.pkg_root, ".github", "workflows"))

    mkpath(joinpath(t.pkg_root, ".dev"))

    contents = gitignore()
    file = joinpath(t.pkg_root, ".gitignore")
    write_contents_to_file(file, contents)

    contents = docs_make(t)
    file = joinpath(t.pkg_root, "docs", "make.jl")
    write_contents_to_file(file, contents)

    contents = index_md(t)
    file = joinpath(t.pkg_root, "docs", "src", "index.md")
    write_contents_to_file(file, contents)

    contents = formatter()
    file = joinpath(t.pkg_root, ".dev", "format.jl")
    write_contents_to_file(file, contents)

    contents = formatter_Project_toml()
    file = joinpath(t.pkg_root, ".dev", "Project.toml")
    write_contents_to_file(file, contents)

    contents = yaml(t.doc_build)
    file = joinpath(t.pkg_root, ".github", "workflows", "Docs.yml")
    write_contents_to_file(file, contents)

    contents = yaml(t.tag_bot)
    file = joinpath(t.pkg_root, ".github", "workflows", "TagBot.yml")
    write_contents_to_file(file, contents)

    contents = yaml(t.compat_helper)
    file = joinpath(t.pkg_root, ".github", "workflows", "CompatHelper.yml")
    write_contents_to_file(file, contents)

    contents = yaml(t.code_cov)
    file = joinpath(t.pkg_root, ".github", "workflows", "CodeCov.yml")
    write_contents_to_file(file, contents)

    contents = yaml(t.formatter_check)
    file = joinpath(t.pkg_root, ".github", "workflows", "Formatter.yml")
    write_contents_to_file(file, contents)

    contents = runtests(t.pkg_name_no_ext)
    file = joinpath(t.pkg_root, "test", "runtests.jl")
    write_contents_to_file(file, contents)

    if t.gpu
        contents = runtests_gpu(t.pkg_name_no_ext)
        file = joinpath(t.pkg_root, "test", "runtests_gpu.jl")
        write_contents_to_file(file, contents)
    end

    contents = toml(t.bors)
    file = joinpath(t.pkg_root, "bors.toml")
    write_contents_to_file(file, contents)

    contents = readme(t)
    file = joinpath(t.pkg_root, "README.md")
    write_contents_to_file(file, contents)

end

end # module

max_stars_repo_name: charleskawczynski/PkgTemplates

max_issues_repo_name: charleskawczynski/PkgTemplates

max_forks_repo_name: charleskawczynski/PkgTemplates


text:
ENV["GENIE_ENV"] = "dev"
push!(LOAD_PATH, "app/resources/insurancecontracts")
using Pkg
Pkg.add("Revise")
using Revise
run(```psql -f sqlsnippets/droptables.sql```)
include("test/testsCreateContract.jl")

max_stars_repo_name: michaelfliegner/BitempWebApp

max_issues_repo_name: michaelfliegner/BitempWebApp


text:
# Support functionality for amend_coverage_from_src!

isevaldef(x) = Base.Meta.isexpr(x, :(=)) && Base.Meta.isexpr(x.args[1], :call) &&
               x.args[1].args[1] == :eval

# This detects different types of function declarations:
# - those that use the  `function` keyword;
# - those that are defined as `f = x -> expr`
# - those that are defined as `f() = expr`
# - those that are defined as `f() where {T} = expr`
# - ... and various more
isfuncexpr(ex::Expr) =
    ex.head == :function || ex.head == :-> ||
        Base.is_short_function_def(ex)
isfuncexpr(arg) = false

function_body_lines(ast, coverage, lineoffset) = function_body_lines!(Int[], ast, coverage, lineoffset, false)

function_body_lines!(flines, arg, coverage, lineoffset, infunction) = flines

function function_body_lines!(flines, node::LineNumberNode, coverage, lineoffset, infunction)
    line = node.line
    if infunction
        push!(flines, line)
    end
    flines
end

function function_body_lines!(flines, ast::Expr, coverage::Vector{CovCount}, lineoffset, infunction)
    if ast.head == :line
        line = ast.args[1]
        if infunction
            push!(flines, line)
        end
        return flines
    elseif ast.head == :module
        # Ignore automatically added eval definitions
        args = ast.args[end].args
        if length(args) >= 2 && isevaldef(args[1]) && isevaldef(args[2])
            args = args[3:end]
        end
    else
        args = ast.args
    end

    # Check whether we are looking at a function declaration.
    # Then also make sure we are not looking at a function declaration
    # that does not also define a method by checking the length of
    # ast.args and making sure it is at least length==2. The test for
    # Expr might not be necessary but also can't harm. Note that this works
    # for both function declarations that use the `function` keyword, and
    # declarations of the form `foo() = expr`, in both cases the method
    # body ends up being `ast.args[2]` (if it exists).
    if isfuncexpr(ast) && length(ast.args)>=2 && ast.args[2] isa Expr
        # Only look in function body and ignore the function signature
        # itself. Sometimes function signatures have line nodes inside
        # and we don't want those lines to be identified as runnable code.
        # In this context, ast.args[1] is the function signature and
        # ast.args[2] is the method body
        #
        # now compute all lines in the body of this function, but for now,
        # track them separately from flines
        flines_new = Int[]
        for arg in ast.args[2].args
            function_body_lines!(flines_new, arg, coverage, lineoffset, true)
        end

        # if any of the lines in the body of the function already has
        # coverage, we assume that the function was executed, generally
        # speaking, and so we should *not* apply our heuristic (which is mean
        # to mark functions which were never executed as code, which the
        # coverage data returned by Julia normally does not)
        not_covered(l) = (l + lineoffset > length(coverage) || coverage[l + lineoffset] === nothing)
        if all(not_covered, flines_new)
            append!(flines, flines_new)
        end
    else
        for arg in args
            function_body_lines!(flines, arg, coverage, lineoffset, infunction)
        end
    end

    flines
end


text:
module q1

# function for calculating temperature and pressure
function tandp(ke, pe, vir, N)
    # Take the time average of the virial, ⟨∑∑r⃗∂V/∂r⟩
    virialavg = means(vir)

    # Find the temperatures
    T = 2ke/3N

    # Set Boltsmann constant
    kB = 1.3806e-23

    kBTs = kB*T

    p = 1 - virialavg./(6N*kBTs)

    return T, p
end


text:
################################################################################
# Comparison of Zig-Zag for diffusion bridges with tailored Poisson rates
# and with adapted Poisson rate. Reference: https://arxiv.org/abs/2001.05889.
################################################################################
using Pkg
Pkg.activate(@__DIR__)
cd(@__DIR__)
using ZigZagBoomerang, SparseArrays, LinearAlgebra, Revise
#using CairoMakie
include("faberschauder.jl")
const ZZB = ZigZagBoomerang

α = 1.5
# Number of levels of FS basis functions (level == first index)
L = 7
#Length diffusion brifges
T = 50.0
# Drift
b(x) = α * sin(x)
# First derivative
b′(x) = α * cos(x)
# Second derivative
b″(x) = -α * sin(x)

n = (2 << L) + 1
ξ0 = 0randn(n)
u, v = -π, 3π  # initial and fianl point
ξ0[1] = u / sqrt(T)
ξ0[end] = v / sqrt(T)
c = ones(n)
c[end] = c[1] = 0.0
θ0 = rand((-1.0, 1.0), n)
θ0[end] = θ0[1] = 0.0 # fix final point
T′ = 1000.0 # final clock of the pdmp

Γ = sparse(1.0I, n, n)
#trace, (t, ξ, θ), (acc, num) = @time pdmp(∇ϕ!, 0.0, ξ0, θ0, T, 10.0, Boomerang(Γ, ξ0*0, 0.1; ρ=0.9), 1, L, adapt=false);
#trace, (t, ξ, θ), (acc, num) = @time pdmp(∇ϕ, 0.0, ξ0, rand((-1.0, 1.0), n), T, 40.0*ones(n), ZigZag(Γ, ξ0*0), 5, L, adapt=false);
trace, (t, ξ, θ), (acc, num), c = @time spdmp(∇ϕmoving, 0.0, ξ0, θ0, T′, c, ZigZag(Γ, ξ0 * 0),
                        SelfMoving(), L, T, adapt = true);

################################################################################
# Overloafing Poisson times in order to have tighter upperbounds
################################################################################
struct MyBound
    c::Float64
end
function ZZB.adapt!(b::Vector{MyBound}, i, x)
    b[i] = MyBound(b[i].c * x)
end

"""
    poisson_time(a, b, c, u)
Obtaining waiting time for inhomogeneous Poisson Process
with rate of the form λ(t) = a + (b + c*t)^+, where `c`,`a`> 0 ,`b` ∈ R, `u` uniform random variable
"""
function ZZB.poisson_time((a, b, c)::NTuple{3}, u = rand()) # formula (22)
    if b > 0
        return (-(b + a) + sqrt((b + a)^2 - 2.0 * c * log(u))) / c # positive solution of quadratic equation c*0.5 x^2 + (b + a) x + log(u) = 0
    elseif a * b / c <= log(u)
        return -log(u) / a
    else
        return (-(a + b) + sqrt((a + b)^2 - 2.0 * c * (b * b * 0.5 / c + log(u)))) / c    # positive solution of quadratic equation c*0.5 x^2 + (b + a) x + log(u) + b*b*0.5/c = 0
    end
end

ZZB.sλ̄((a, b, c)::NTuple{3}, Δt) = a + ZZB.pos(b + c * Δt)
"""
    abc(G, i, x, θ, c, Flow)

Returns the constant term `a` and linear term `b` when computing the Poisson times
from the upper upper bounding rates λᵢ(t) = max(a + b*t)^2. The factors `a` and `b`
can be function of the current position `x`, velocity `θ`, tuning parameter `c` and
the Graph `G`
"""
function ZZB.ab(G, i, x, θ, c::Vector{MyBound}, F::ZigZag)
    if i == 1
        a = c[i].c + T^(1.5)*0.5*(α^2 + α) * abs(θ[i])  # initial point
        b1 = θ[i]*(x[i] - x[end])
        b2 = θ[i]*(θ[i] - θ[end])
    elseif i == (2 << L) + 1
        a = c[i].c + T^(1.5)*0.5*(α^2 + α) * abs(θ[i])  # final point
        b1 = θ[i]*(x[i] - x[1])
        b2 = θ[i]*(θ[i] - θ[1])
    else
        l = lvl(i, L)
        a = c[i].c + T^(1.5) / 2^((L - l) * 1.5 + 2) * (α^2 + α) * abs(θ[i]) # formula (22)
        b1 = x[i] * θ[i]
        b2 = θ[i] * θ[i]
    end
    return a, b1, b2
end

c = [MyBound(0.0) for i in 1:n]
trace, (t, ξ, θ), num, c = @time spdmp(∇ϕmoving, 0.0, ξ0, θ0, T′, c, ZigZag(Γ, ξ0 * 0),
                        SelfMoving(), L, T, adapt = false);

ts, ξs = splitpairs(discretize(trace, T′/n))
S = T*(0:n)/(n+1)

### change to false is you want to Plot the results
if true
    error("Stop before plotting the results...")
end

using CairoMakie
p1 = lines(S, [dotψ(ξ, s, L, T) for s in S], linewidth=0.3)
for ξ in ξs[1:5:end]
    lines!(p1, S, [dotψ(ξ, s, L, T) for s in S], linewidth=0.3)
end
display(p1)


text:
using MatrixCompletion
import Distributions

@testset "$(format("GD Optimizer: Negative Binomial Loss [Small][Forgiving][Native]"))" begin
  let
    input_size = 1600*400
    y = rand(Distributions.NegativeBinomial(10, 0.6), input_size) * 1.0
    mle_x = MatrixCompletion.Losses.negative_binomial_train(fx      = rand(input_size),
                                                            y       = y,
                                                            c       = zeros(input_size),
                                                            ρ       = 0,
                                                            γ       = 0.2,
                                                            iter    = 200,
                                                            verbose = true,
                                                            r_estimate = 10)
    @show(mle_x[1:20])
    prediction = predict(NegativeBinomial(),forward_map(NegativeBinomial(),mle_x,r_estimate=10))
    tc = provide(Diagnostics{Poisson()}(),
                 input_data=prediction, reference=y)
    display(tc)
  end
end


text:
export VisitationFrequency, probabilities

"""
    VisitationFrequency(r::RectangularBinning) <: BinningProbabilitiesEstimator

A probability estimator based on binning data into rectangular boxes dictated by
the binning scheme `r`.

## Example

```julia
# Construct boxes by dividing each coordinate axis into 5 equal-length chunks.
b = RectangularBinning(5)

# A probabilities estimator that, when applied a dataset, computes visitation frequencies
# over the boxes of the binning
est = VisitationFrequency(b)
```

See also: [`RectangularBinning`](@ref).
"""
struct VisitationFrequency{RB<:RectangularBinning} <: BinningProbabilitiesEstimator
    binning::RB
end

function probabilities(x::AbstractDataset, est::VisitationFrequency)
    _non0hist(x, est.binning)[1]
end


text:
# NOTE: these currently just test that the inference methods do not crash
# put longer-running statistical tests of inference methods elsewhere

@testset "importance sampling" begin

    @gen function model()
        x = @addr(normal(0, 1), :x)
        @addr(normal(x, 1), :y)
    end

    @gen function proposal()
        @addr(normal(0, 2), :x)
    end

    y = 2.
    observations = DynamicAssignment()
    set_leaf_node!(observations, :y, y)

    n = 4

    (traces, log_weights, lml_est) = importance_sampling(model, (), observations, n)
    @test length(traces) == n
    @test length(log_weights) == n
    @test isapprox(logsumexp(log_weights), 0., atol=1e-14)
    @test !isnan(lml_est)
    for trace in traces
        @test get_assignment(trace)[:y] == y
    end

    (traces, log_weights, lml_est) = importance_sampling(model, (), observations, proposal, (), n)
    @test length(traces) == n
    @test length(log_weights) == n
    @test isapprox(logsumexp(log_weights), 0., atol=1e-14)
    @test !isnan(lml_est)
    for trace in traces
        @test get_assignment(trace)[:y] == y
    end

    (trace, lml_est) = importance_resampling(model, (), observations, n)
    @test isapprox(logsumexp(log_weights), 0., atol=1e-14)
    @test !isnan(lml_est)
    @test get_assignment(trace)[:y] == y

    (trace, lml_est) = importance_resampling(model, (), observations, proposal, (), n)
    @test isapprox(logsumexp(log_weights), 0., atol=1e-14)
    @test !isnan(lml_est)
    @test get_assignment(trace)[:y] == y
end

@testset "SGD training" begin

    Random.seed!(1)

    # x y   z

    # 0 0   0
    # 0 1   1
    # 1 0   1
    # 1 1   0

    # x y   p(x, y | z=0)
    # 0 0   0.5
    # 0 1   0
    # 1 0   0
    # 1 1   0.5

    # x y   p(x, y | z=1)
    # 0 0   0
    # 0 1   0.5
    # 1 0   0.5
    # 1 1   0

    @gen function teacher()
        x = @addr(bernoulli(0.5), :x)
        y = @addr(bernoulli(0.5), :y)
        z::Bool = xor(x, y)
        @addr(bernoulli(z ? 1.0 : 0.0), :z)
    end

    sigmoid(val) = 1.0 / (1.0 + exp(-val))

    @gen function batch_student(zs::Vector{Bool})
        @param theta1::Float64
        @param theta2::Float64
        @param theta3::Float64
        @param theta4::Float64
        @param theta5::Float64
        for (i, z) in enumerate(zs)
            x = @addr(bernoulli(sigmoid(theta1)), i => :x)
            if z && x
                prob_y = sigmoid(theta2) # should be near 0
            elseif z && !x
                prob_y = sigmoid(theta3) # should be near 1
            elseif !z && x
                prob_y = sigmoid(theta4) # should be near 1
            elseif !z && !x
                prob_y = sigmoid(theta5) # should be near 0
            end
            @addr(bernoulli(prob_y), i => :y)
        end
    end

    # theta1 = 0.0
    # theta2 -> inf (prob_y -> 1)
    # theta3 -> -inf (prob_y -> 0)

    input_extractor = (samples::Vector) -> (Bool[s[:z] for s in samples],)

    function constraint_extractor(samples::Vector)
        constraints = DynamicAssignment()
        for (i, s) in enumerate(samples)
            constraints[i => :x] = s[:x]
            constraints[i => :y] = s[:y]
        end
        constraints
    end

    function minibatch_callback(batch, minibatch, avg_score, verbose)
        for name in [:theta1, :theta2, :theta3, :theta4, :theta5]
            grad = get_param_grad(batch_student, name)
            value = get_param(batch_student, name)
            set_param!(batch_student, name, value + grad * 0.01)
            zero_param_grad!(batch_student, name)
        end
    end

    batch_callback = (batch, verbose) -> nothing

    init_param!(batch_student, :theta1, 0.)
    init_param!(batch_student, :theta2, 0.)
    init_param!(batch_student, :theta3, 0.)
    init_param!(batch_student, :theta4, 0.)
    init_param!(batch_student, :theta5, 0.)

    # check gradients using finite differences on a simulated batch
    assignments = Vector{Any}(undef, 100)
    for i=1:100
        assignments[i] = get_assignment(simulate(teacher, ()))
    end
    input = input_extractor(assignments)
    constraints = constraint_extractor(assignments)
    student_trace = assess(batch_student, input, constraints)
    backprop_params(batch_student, student_trace, nothing)
    for name in [:theta1, :theta2, :theta3, :theta4, :theta5]
        actual = get_param_grad(batch_student, name)
        dx = 1e-6
        value = get_param(batch_student, name)
        set_param!(batch_student, name, value + dx)
        lpdf_pos = get_call_record(assess(batch_student, input, constraints)).score
        set_param!(batch_student, name, value - dx)
        lpdf_neg = get_call_record(assess(batch_student, input, constraints)).score
        set_param!(batch_student, name, value)
        expected = (lpdf_pos - lpdf_neg) / (2 * dx)
        @test isapprox(actual, expected, atol=1e-4)
    end

    conf = SGDTrainConf(2000, 50, 1, 50,
        input_extractor, constraint_extractor,
        minibatch_callback, batch_callback)
    sgd_train_batch(teacher, (), batch_student, conf, false)

    # p(x | z=0) = p(x | z=1) = 0.5
    @test isapprox(get_param(batch_student, :theta1), 0., atol=0.2)

    # y | z, x = xor(x, z)
    @test get_param(batch_student, :theta2) < -5
    @test get_param(batch_student, :theta3) > 5
    @test get_param(batch_student, :theta4) > 5
    @test get_param(batch_student, :theta5) < -5
end

function hmm_forward_alg(prior::Vector{Float64},
        emission_dists::AbstractArray{Float64,2}, transition_dists::AbstractArray{Float64,2},
        emissions::Vector{Int})
    marg_lik = 1.
    alpha = prior # p(z_1)
    for i=2:length(emissions)

        # p(z_{i-1} , y_{i-1} | y_{1:i-2}) for each z_{i-1}
        prev_posterior = alpha .* emission_dists[emissions[i-1], :]

        # p(y_{i-1} | y_{1:i-2})
        denom = sum(prev_posterior)

        # p(z_{i-1} | y_{1:i-1})
        prev_posterior = prev_posterior / denom

        # p(z_i | y_{1:i-1})
        alpha = transition_dists * prev_posterior

        # p(y_{1:i-1})
        marg_lik *= denom
    end
    prev_posterior = alpha .* emission_dists[emissions[end], :]
    denom = sum(prev_posterior)
    marg_lik *= denom
    marg_lik
end

@testset "hmm forward alg" begin

    # test the hmm_forward_alg on a hand-calculated example
    prior = [0.4, 0.6]
    emission_dists = [0.1 0.9; 0.7 0.3]'
    transition_dists = [0.5 0.5; 0.2 0.8']
    obs = [2, 1]
    expected_marg_lik = 0.
    # z = [1, 1]
    expected_marg_lik += prior[1] * transition_dists[1, 1] * emission_dists[obs[1], 1] * emission_dists[obs[2], 1]
    # z = [1, 2]
    expected_marg_lik += prior[1] * transition_dists[2, 1] * emission_dists[obs[1], 1] * emission_dists[obs[2], 2]
    # z = [2, 1]
    expected_marg_lik += prior[2] * transition_dists[1, 2] * emission_dists[obs[1], 2] * emission_dists[obs[2], 1]
    # z = [2, 2]
    expected_marg_lik += prior[2] * transition_dists[2, 2] * emission_dists[obs[1], 2] * emission_dists[obs[2], 2]
    actual_marg_lik = hmm_forward_alg(prior, emission_dists, transition_dists, obs)
    @test isapprox(actual_marg_lik, expected_marg_lik)

end

@testset "particle filtering" begin

    prior = [0.2, 0.3, 0.5]

    emission_dists = [
        0.1 0.2 0.8;
        0.2 0.8 0.1;
        0.8 0.2 0.1
    ]'

    transition_dists = [
        0.4 0.4 0.2;
        0.2 0.3 0.5;
        0.9 0.05 0.05
    ]'

    @gen function model_init()
        z = @addr(categorical(prior), :z)
        @addr(categorical(emission_dists[:,z]), :x)
        return z
    end

    @gen function model_step(t::Int, prev_z::Int, params::Nothing)
        z = @addr(categorical(transition_dists[:,prev_z]), :z)
        @addr(categorical(emission_dists[:,z]), :x)
        return z
    end

    @gen function model(num_steps::Int)
        z_init = @addr(model_init(), :init)
        change = MarkovChange(true, false, false) # we are only ever extending
        @addr(markov(model_step)(num_steps-1, z_init, nothing), :markov, change)
    end

    num_steps = 4
    obs_x = [1, 1, 2, 3]

    # latents:
    # :init => z
    # :markov => 1 => :z
    # :markov => 2 => :z
    # :markov => 3 => :z

    # observations :
    # :init => :x
    # :markov => 1 => :x
    # :markov => 2 => :x
    # :markov => 3 => :x

    @gen function proposal_init(x::Int)
        dist = prior .* emission_dists[x,:]
        @addr(categorical(dist ./ sum(dist)), :init => :z)
    end

    @gen function proposal_step(t::Int, prev_z::Int, x::Int)
        dist = transition_dists[:,prev_z] .* emission_dists[x,:]
        # NOTE: was missing :markov, this should have been an error..
        @addr(categorical(dist ./ sum(dist)), :markov => t => :z)
    end

    function get_init_observations_and_proposal_args()
        observations = DynamicAssignment()
        observations[:init => :x] = obs_x[1]
        init_proposal_args = (obs_x[1],)
        (observations, init_proposal_args)
    end

    function get_step_observations_and_proposal_args(step::Int, prev_trace)
        @assert !has_internal_node(get_assignment(prev_trace), :markov => (step - 1))
        @assert step > 1
        if step == 2
            prev_z = get_assignment(prev_trace)[:init => :z]
        else
            prev_z = get_assignment(prev_trace)[:markov => (step - 2) => :z]
        end
        observations = DynamicAssignment()
        observations[:markov => (step - 1) => :x] = obs_x[step]
        step_proposal_args = (step - 1, prev_z, obs_x[step])
        (observations, step_proposal_args)
    end

    Random.seed!(0)

    num_particles = 10000
    ess_threshold = 10000 # make sure we exercise resampling
    (_, _, log_ml_est) = particle_filter(model, (),
                    num_steps, num_particles, ess_threshold,
                    get_init_observations_and_proposal_args,
                    get_step_observations_and_proposal_args,
                    proposal_init, proposal_step; verbose=true)

    expected_log_ml = log(hmm_forward_alg(prior, emission_dists, transition_dists, obs_x))
    @assert isapprox(expected_log_ml, log_ml_est, atol=0.01)
end

max_stars_repo_name: BoltzmannBrain/Gen

max_issues_repo_name: BoltzmannBrain/Gen

max_forks_repo_name: BoltzmannBrain/Gen


text:
# define commonly used symbols
const σ0 = @SMatrix [1 + 0im 0; 0 1 + 0im]
const σx = @SMatrix [0 1 + 0im; 1 + 0im 0]
const σy = @SMatrix [0 1im; -1im 0]
const σz = @SMatrix [1 + 0im 0; 0 -1 + 0im]


text:
# Autogenerated wrapper script for Libical_jll for x86_64-w64-mingw32-cxx11
export libical, libical_cxx, libicalss, libicalss_cxx, libicalvcal

using BerkeleyDB_jll
using Glib_jll
using ICU_jll
using XML2_jll
JLLWrappers.@generate_wrapper_header("Libical")
JLLWrappers.@declare_library_product(libical, "libical.dll")
JLLWrappers.@declare_library_product(libical_cxx, "libical_cxx.dll")
JLLWrappers.@declare_library_product(libicalss, "libicalss.dll")
JLLWrappers.@declare_library_product(libicalss_cxx, "libicalss_cxx.dll")
JLLWrappers.@declare_library_product(libicalvcal, "libicalvcal.dll")
function __init__()
    JLLWrappers.@generate_init_header(BerkeleyDB_jll, Glib_jll, ICU_jll, XML2_jll)
    JLLWrappers.@init_library_product(
        libical,
        "bin\\libical.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libical_cxx,
        "bin\\libical_cxx.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libicalss,
        "bin\\libicalss.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libicalss_cxx,
        "bin\\libicalss_cxx.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libicalvcal,
        "bin\\libicalvcal.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()