# This file is part of Kpax3. License is MIT.

function test_mcmc_gibbs_init()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  settings = Kpax3.KSettings(ifile, ofile, maxclust=3, maxunit=1)

  data = UInt8[0 0 0 0 0 1;
               1 1 1 1 1 0;
               0 0 1 0 1 1;
               1 1 0 1 0 0;
               1 1 0 0 0 0;
               0 0 0 1 1 0;
               1 1 1 0 0 0;
               0 0 0 1 1 1;
               0 0 1 0 0 0;
               1 0 0 1 0 1;
               0 1 0 0 1 0;
               0 0 0 0 0 1;
               1 1 1 0 0 0;
               0 0 0 1 1 0;
               1 1 0 0 1 1;
               0 0 1 1 0 0;
               1 1 0 1 0 0;
               0 0 1 0 1 1]

  (m, n) = size(data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(data, settings.γ, settings.r)

  # remove unit 5 from cluster 2 and put it into another cluster
  R = [1; 1; 1; 2; 2; 3]

  i = 5
  hi = 2

  state = Kpax3.AminoAcidState(data, R, priorR, priorC, settings)
  support = Kpax3.MCMCSupport(state, priorC)

  Kpax3.initsupportgibbs!(hi, priorR, state, support)

  @test support.k == 2
  @test support.cl[1:support.k] == [1; 3]
  @test support.t[1:support.k] == [Kpax3.clusterweight(state.v[1], priorR);
                                   Kpax3.clusterweight(state.v[3], priorR)]

  nothing
end

test_mcmc_gibbs_init()

function test_mcmc_gibbs_move()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  settings = Kpax3.KSettings(ifile, ofile, maxclust=4, maxunit=1)

  data = UInt8[0 0 0 0 0 1;
               1 1 1 1 1 0;
               0 0 1 0 1 1;
               1 1 0 1 0 0;
               1 1 0 0 0 0;
               0 0 0 1 1 0;
               1 1 1 0 0 0;
               0 0 0 1 1 1;
               0 0 1 0 0 0;
               1 0 0 1 0 1;
               0 1 0 0 1 0;
               0 0 0 0 0 1;
               1 1 1 0 0 0;
               0 0 0 1 1 0;
               1 1 0 0 1 1;
               0 0 1 1 0 0;
               1 1 0 1 0 0;
               0 0 1 0 1 1]

  (m, n) = size(data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(data, settings.γ, settings.r)

  # remove unit 5 from cluster 2 and put it into another cluster
  R = [1; 1; 1; 2; 2; 3]

  i = 5
  hi = 2
  li = 3
  lj = 4

  state = Kpax3.AminoAcidState(data, R, priorR, priorC, settings)
  support = Kpax3.MCMCSupport(state, priorC)

  support.ni = vec(float(data[:, 4]))
  support.vi = 1

  Kpax3.initsupportgibbs!(hi, priorR, state, support)

  support.t[li] = Kpax3.clusterweight(support.vi, priorR)
  support.t[lj] = Kpax3.clusterweight(1, state.k, priorR)

  g = 0
  l = 0
  y = 0.0
  v = 0
  z = zeros(Float64, 2)
  for b in 1:m
    support.lpi[1, b] = priorC.logγ[1]
    support.lpi[2, b] = priorC.logγ[2]
    support.lpi[3, b] = priorC.logγ[3]
    support.lpi[4, b] = priorC.logγ[3]

    g = 1
    l = 1
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmove!(y, v, b, l, g, state.k, priorC, support)

    @test support.lq[1, l, b] == Kpax3.logmarglik(y, v, priorC.A[1, b], priorC.B[1, b])
    @test support.lq[2, l, b] == Kpax3.logmarglik(y, v, priorC.A[2, b], priorC.B[2, b])
    @test support.lq[3, l, b] == Kpax3.logmarglik(y, v, priorC.A[3, b], priorC.B[3, b])
    @test support.lq[4, l, b] == Kpax3.logmarglik(y, v, priorC.A[4, b], priorC.B[4, b])

    g = 3
    l = 2
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmove!(y, v, b, l, g, state.k, priorC, support)

    @test support.lq[1, l, b] == Kpax3.logmarglik(y, v, priorC.A[1, b], priorC.B[1, b])
    @test support.lq[2, l, b] == Kpax3.logmarglik(y, v, priorC.A[2, b], priorC.B[2, b])
    @test support.lq[3, l, b] == Kpax3.logmarglik(y, v, priorC.A[3, b], priorC.B[3, b])
    @test support.lq[4, l, b] == Kpax3.logmarglik(y, v, priorC.A[4, b], priorC.B[4, b])

    @test isapprox(support.lpi[1, b], priorC.logγ[1] + support.lp[1, 1, b] + support.lp[1, 3, b], atol=ε)
    @test isapprox(support.lpi[2, b], priorC.logγ[2] + support.lp[2, 1, b] + support.lp[2, 3, b], atol=ε)

    z[1] = log(exp(priorC.logω[state.k][1] + support.lp[3, 1, b]) + exp(priorC.logω[state.k][2] + support.lp[4, 1, b]))
    z[2] = log(exp(priorC.logω[state.k][1] + support.lp[3, 3, b]) + exp(priorC.logω[state.k][2] + support.lp[4, 3, b]))

    @test isapprox(support.lpi[3, b], priorC.logγ[3] + z[1] + z[2], atol=ε)

    z[1] = support.lq[1, 1, b] + support.lq[1, li, b] - support.lp[1, 1, b]
    z[2] = support.lq[1, 2, b] + support.lq[1, li, b] - support.lp[1, 3, b]
    @test isapprox(support.lr[1, 1, b], z[1], atol=ε)
    @test isapprox(support.lr[1, 2, b], z[2], atol=ε)

    z[1] = support.lq[2, 1, b] + support.lq[2, li, b] - support.lp[2, 1, b]
    z[2] = support.lq[2, 2, b] + support.lq[2, li, b] - support.lp[2, 3, b]
    @test isapprox(support.lr[2, 1, b], z[1], atol=ε)
    @test isapprox(support.lr[2, 2, b], z[2], atol=ε)

    z[1] = log(exp(priorC.logω[state.k][1] + support.lq[3, 1, b]) + exp(priorC.logω[state.k][2] + support.lq[4, 1, b])) - log(exp(priorC.logω[state.k][1] + support.lp[3, 1, b]) + exp(priorC.logω[state.k][2] + support.lp[4, 1, b]))
    z[2] = log(exp(priorC.logω[state.k][1] + support.lq[3, 2, b]) + exp(priorC.logω[state.k][2] + support.lq[4, 2, b])) - log(exp(priorC.logω[state.k][1] + support.lp[3, 3, b]) + exp(priorC.logω[state.k][2] + support.lp[4, 3, b]))

    @test isapprox(support.lr[3, 1, b], z[1], atol=ε)
    @test isapprox(support.lr[3, 2, b], z[2], atol=ε)

    Kpax3.gibbsupdateclusteri!(state.k, b, li, priorC, support)

    y = support.ni[b]
    v = support.vi

    @test support.lq[1, li, b] == Kpax3.logmarglik(y, v, priorC.A[1, b], priorC.B[1, b])
    @test support.lq[2, li, b] == Kpax3.logmarglik(y, v, priorC.A[2, b], priorC.B[2, b])
    @test support.lq[3, li, b] == Kpax3.logmarglik(y, v, priorC.A[3, b], priorC.B[3, b])
    @test support.lq[4, li, b] == Kpax3.logmarglik(y, v, priorC.A[4, b], priorC.B[4, b])

    z[1] = priorC.logω[state.k][1] + support.lq[3, li, b]
    z[2] = priorC.logω[state.k][2] + support.lq[4, li, b]

    @test isapprox(support.lpj[1, b], log(exp(z[1]) + exp(z[2])),  atol=ε)

    Kpax3.gibbsupdateclusterj!(data[b, i], state.k + 1, b, li, lj, priorC, support)

    y = support.ni[b]
    v = 1

    z[1] = priorC.logω[state.k + 1][1] + Kpax3.logmarglik(y, v, priorC.A[3, b], priorC.B[3, b])
    z[2] = priorC.logω[state.k + 1][2] + Kpax3.logmarglik(y, v, priorC.A[4, b], priorC.B[4, b])

    @test isapprox(support.lpj[2, b], log(exp(z[1]) + exp(z[2])), atol=ε)

    @test support.lq[1, lj, b] == Kpax3.logmarglik(data[b, i], 1, priorC.A[1, b], priorC.B[1, b])
    @test support.lq[2, lj, b] == Kpax3.logmarglik(data[b, i], 1, priorC.A[2, b], priorC.B[2, b])
    @test support.lq[3, lj, b] == Kpax3.logmarglik(data[b, i], 1, priorC.A[3, b], priorC.B[3, b])
    @test support.lq[4, lj, b] == Kpax3.logmarglik(data[b, i], 1, priorC.A[4, b], priorC.B[4, b])

    z[1] = priorC.logω[state.k + 1][1] + support.lq[3, lj, b]
    z[2] = priorC.logω[state.k + 1][2] + support.lq[4, lj, b]

    @test isapprox(support.lpj[3, b], log(exp(z[1]) + exp(z[2])), atol=ε)
  end

  Kpax3.gibbscomputeprobi!(li, support)

  for b in 1:support.m
    Kpax3.gibbscomputeprobg!(b, 1, li, support)
    Kpax3.gibbscomputeprobg!(b, 2, li, support)
    Kpax3.gibbscomputeprobj!(b, li, lj, support)
  end

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 1; 3], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)
  t = Kpax3.clusterweight(state.v[1], priorR) + su.logmlik

  @test isapprox(support.t[1], t, atol=ε)

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 3; 3], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)
  t = Kpax3.clusterweight(state.v[3], priorR) + su.logmlik

  @test isapprox(support.t[2], t, atol=ε)

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 2; 3], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)
  t = Kpax3.clusterweight(state.v[2] - 1, priorR) + su.logmlik

  @test isapprox(support.t[li], t, atol=ε)

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 4; 3], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)
  t = Kpax3.clusterweight(1, state.k, priorR) + su.logmlik

  @test isapprox(support.t[lj], t, atol=ε)

  logc = Kpax3.gibbscomputenormconst(lj, support)

  @test isapprox(logc, log(sum(exp.(support.t[1:4]))), atol=ε)

  nothing
end

test_mcmc_gibbs_move()

function test_mcmc_gibbs_merge()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  settings = Kpax3.KSettings(ifile, ofile, maxclust=3, maxunit=1)

  data = UInt8[0 0 0 0 0 1;
               1 1 1 1 1 0;
               0 0 1 0 1 1;
               1 1 0 1 0 0;
               1 1 0 0 0 0;
               0 0 0 1 1 0;
               1 1 1 0 0 0;
               0 0 0 1 1 1;
               0 0 1 0 0 0;
               1 0 0 1 0 1;
               0 1 0 0 1 0;
               0 0 0 0 0 1;
               1 1 1 0 0 0;
               0 0 0 1 1 0;
               1 1 0 0 1 1;
               0 0 1 1 0 0;
               1 1 0 1 0 0;
               0 0 1 0 1 1]

  (m, n) = size(data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(data, settings.γ, settings.r)

  # remove unit 6 from cluster 3 and put it into another cluster
  R = [1; 1; 1; 2; 2; 3]

  i = 6
  hi = 3
  li = 3
  lj = 3

  state = Kpax3.AminoAcidState(data, R, priorR, priorC, settings)
  support = Kpax3.MCMCSupport(state, priorC)

  Kpax3.initsupportgibbs!(hi, priorR, state, support)

  support.t[lj] = Kpax3.clusterweight(1, state.k - 1, priorR)

  g = 0
  l = 0
  y = 0.0
  v = 0
  z = zeros(Float64, 2)
  for b in 1:m
    support.lpi[1, b] = priorC.logγ[1]
    support.lpi[2, b] = priorC.logγ[2]
    support.lpi[3, b] = priorC.logγ[3]

    g = 1
    l = 1
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmerge!(y, v, b, l, g, support.k, priorC, support)

    @test support.lq[1, l, b] == Kpax3.logmarglik(y, v, priorC.A[1, b], priorC.B[1, b])
    @test support.lq[2, l, b] == Kpax3.logmarglik(y, v, priorC.A[2, b], priorC.B[2, b])
    @test support.lq[3, l, b] == Kpax3.logmarglik(y, v, priorC.A[3, b], priorC.B[3, b])
    @test support.lq[4, l, b] == Kpax3.logmarglik(y, v, priorC.A[4, b], priorC.B[4, b])

    g = 2
    l = 2
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmerge!(y, v, b, l, g, support.k, priorC, support)

    @test support.lq[1, l, b] == Kpax3.logmarglik(y, v, priorC.A[1, b], priorC.B[1, b])
    @test support.lq[2, l, b] == Kpax3.logmarglik(y, v, priorC.A[2, b], priorC.B[2, b])
    @test support.lq[3, l, b] == Kpax3.logmarglik(y, v, priorC.A[3, b], priorC.B[3, b])
    @test support.lq[4, l, b] == Kpax3.logmarglik(y, v, priorC.A[4, b], priorC.B[4, b])

    @test isapprox(support.lpi[1, b], priorC.logγ[1] + support.lp[1, 1, b] + support.lp[1, 2, b], atol=ε)
    @test isapprox(support.lpi[2, b], priorC.logγ[2] + support.lp[2, 1, b] + support.lp[2, 2, b], atol=ε)

    z[1] = log(exp(priorC.logω[support.k][1] + support.lp[3, 1, b]) + exp(priorC.logω[support.k][2] + support.lp[4, 1, b]))
    z[2] = log(exp(priorC.logω[support.k][1] + support.lp[3, 2, b]) + exp(priorC.logω[support.k][2] + support.lp[4, 2, b]))

    @test isapprox(support.lpi[3, b], priorC.logγ[3] + z[1] + z[2], atol=ε)

    z[1] = support.lq[1, 1, b] + support.lq[1, li, b] - support.lp[1, 1, b]
    z[2] = support.lq[1, 2, b] + support.lq[1, li, b] - support.lp[1, 2, b]
    @test isapprox(support.lr[1, 1, b], z[1], atol=ε)
    @test isapprox(support.lr[1, 2, b], z[2], atol=ε)

    z[1] = support.lq[2, 1, b] + support.lq[2, li, b] - support.lp[2, 1, b]
    z[2] = support.lq[2, 2, b] + support.lq[2, li, b] - support.lp[2, 2, b]
    @test isapprox(support.lr[2, 1, b], z[1], atol=ε)
    @test isapprox(support.lr[2, 2, b], z[2], atol=ε)

    z[1] = log(exp(priorC.logω[support.k][1] + support.lq[3, 1, b]) + exp(priorC.logω[support.k][2] + support.lq[4, 1, b])) + support.lpj[1, b] - log(exp(priorC.logω[support.k][1] + support.lp[3, 1, b]) + exp(priorC.logω[support.k][2] + support.lp[4, 1, b]))
    z[2] = log(exp(priorC.logω[support.k][1] + support.lq[3, 2, b]) + exp(priorC.logω[support.k][2] + support.lq[4, 2, b])) + support.lpj[1, b] - log(exp(priorC.logω[support.k][1] + support.lp[3, 2, b]) + exp(priorC.logω[support.k][2] + support.lp[4, 2, b]))

    @test isapprox(support.lr[3, 1, b], z[1], atol=ε)
    @test isapprox(support.lr[3, 2, b], z[2], atol=ε)
  end

  Kpax3.gibbscomputeprobi!(lj, support)

  for b in 1:support.m
    Kpax3.gibbscomputeprobg!(b, 1, support)
    Kpax3.gibbscomputeprobg!(b, 2, support)
  end

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 2; 1], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)
  t = Kpax3.clusterweight(state.v[1], priorR) + su.logmlik

  @test isapprox(support.t[1], t, atol=ε)

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 2; 2], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)
  t = Kpax3.clusterweight(state.v[2], priorR) + su.logmlik

  @test isapprox(support.t[2], t, atol=ε)

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 2; 3], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)
  t = Kpax3.clusterweight(1, state.k - 1, priorR) + su.logmlik

  @test isapprox(support.t[lj], t, atol=ε)

  logc = Kpax3.gibbscomputenormconst(lj, support)

  @test isapprox(logc, log(sum(exp.(support.t[1:3]))), atol=ε)

  nothing
end

test_mcmc_gibbs_merge()

function test_mcmc_gibbs_sample_cluster()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  settings = Kpax3.KSettings(ifile, ofile, maxclust=4, maxunit=1)

  data = UInt8[0 0 0 0 0 1;
               1 1 1 1 1 0;
               0 0 1 0 1 1;
               1 1 0 1 0 0;
               1 1 0 0 0 0;
               0 0 0 1 1 0;
               1 1 1 0 0 0;
               0 0 0 1 1 1;
               0 0 1 0 0 0;
               1 0 0 1 0 1;
               0 1 0 0 1 0;
               0 0 0 0 0 1;
               1 1 1 0 0 0;
               0 0 0 1 1 0;
               1 1 0 0 1 1;
               0 0 1 1 0 0;
               1 1 0 1 0 0;
               0 0 1 0 1 1]

  (m, n) = size(data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(data, settings.γ, settings.r)

  # remove unit 5 from cluster 2 and put it into another cluster
  R = [1; 1; 1; 2; 2; 3]

  state = Kpax3.AminoAcidState(data, R, priorR, priorC, settings)
  support = Kpax3.MCMCSupport(state, priorC)

  support.t[1] = log(0.001) + log(0.4)
  support.t[2] = log(0.001) + log(0.3)
  support.t[3] = log(0.001) + log(0.2)
  support.t[4] = log(0.001) + log(0.1)

  logc = log(0.001)

  N = 1000000
  g = 0
  counter = zeros(Float64, 6)
  for t in 1:N
    g = Kpax3.gibbssamplecluster(logc, 4, support)
    counter[g] += 1
  end

  counter /= N

  @test maximum(abs.(counter - [0.4; 0.3; 0.2; 0.1; 0.0; 0.0])) <= 0.005

  nothing
end

test_mcmc_gibbs_sample_cluster()

function test_mcmc_gibbs_perform_move()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  settings = Kpax3.KSettings(ifile, ofile, maxclust=4, maxunit=1)

  data = UInt8[0 0 0 0 0 1;
               1 1 1 1 1 0;
               0 0 1 0 1 1;
               1 1 0 1 0 0;
               1 1 0 0 0 0;
               0 0 0 1 1 0;
               1 1 1 0 0 0;
               0 0 0 1 1 1;
               0 0 1 0 0 0;
               1 0 0 1 0 1;
               0 1 0 0 1 0;
               0 0 0 0 0 1;
               1 1 1 0 0 0;
               0 0 0 1 1 0;
               1 1 0 0 1 1;
               0 0 1 1 0 0;
               1 1 0 1 0 0;
               0 0 1 0 1 1]

  (m, n) = size(data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(data, settings.γ, settings.r)

  # remove unit 5 from cluster 2 and put it into another cluster
  R = [1; 1; 1; 2; 2; 3]

  i = 5
  hi = 2
  li = 3
  lj = 4

  state = Kpax3.AminoAcidState(data, R, priorR, priorC, settings)
  support = Kpax3.MCMCSupport(state, priorC)

  support.ni = vec(float(data[:, 4]))
  support.vi = 1

  Kpax3.initsupportgibbs!(hi, priorR, state, support)

  support.t[li] = Kpax3.clusterweight(support.vi, priorR)
  support.t[lj] = Kpax3.clusterweight(1, state.k, priorR)

  g = 0
  l = 0
  y = 0.0
  v = 0
  z = zeros(Float64, 2)
  for b in 1:m
    support.lpi[1, b] = priorC.logγ[1]
    support.lpi[2, b] = priorC.logγ[2]
    support.lpi[3, b] = priorC.logγ[3]
    support.lpi[4, b] = priorC.logγ[3]

    g = 1
    l = 1
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmove!(y, v, b, l, g, state.k, priorC, support)

    g = 3
    l = 2
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmove!(y, v, b, l, g, state.k, priorC, support)

    Kpax3.gibbsupdateclusteri!(state.k, b, li, priorC, support)
    Kpax3.gibbsupdateclusterj!(data[b, i], state.k + 1, b, li, lj, priorC, support)
  end

  Kpax3.gibbscomputeprobi!(li, support)

  for b in 1:support.m
    Kpax3.gibbscomputeprobg!(b, 1, li, support)
    Kpax3.gibbscomputeprobg!(b, 2, li, support)
    Kpax3.gibbscomputeprobj!(b, li, lj, support)
  end

  # R = [1; 1; 1; 2; 3; 3]
  l = 2

  Kpax3.gibbsmove!(i, hi, support.cl[l], li, l, data, priorR, support, state)

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 3; 3], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)

  @test state.R == st.R
  @test state.k == st.k

  @test !state.emptycluster[1]
  @test state.cl[1] == st.cl[1]
  @test state.v[1] == st.v[1]
  @test state.n1s[1, :] == st.n1s[1, :]
  @test state.unit[1][1:state.v[1]] == [1; 2; 3]

  @test !state.emptycluster[2]
  @test state.cl[2] == st.cl[2]
  @test state.v[2] == st.v[2]
  @test state.n1s[2, :] == st.n1s[2, :]
  @test state.unit[2][1:state.v[2]] == [4]

  @test !state.emptycluster[3]
  @test state.cl[3] == st.cl[3]
  @test state.v[3] == st.v[3]
  @test state.n1s[3, :] == st.n1s[3, :]
  @test state.unit[3][1:state.v[3]] == [6; 5]

  @test isapprox(state.logpR, st.logpR, atol=ε)

  @test isapprox(support.logmlik, su.logmlik, atol=ε)

  @test support.lp == su.lp

  nothing
end

test_mcmc_gibbs_perform_move()

function test_mcmc_gibbs_perform_merge()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  settings = Kpax3.KSettings(ifile, ofile, maxclust=3, maxunit=1)

  data = UInt8[0 0 0 0 0 1;
               1 1 1 1 1 0;
               0 0 1 0 1 1;
               1 1 0 1 0 0;
               1 1 0 0 0 0;
               0 0 0 1 1 0;
               1 1 1 0 0 0;
               0 0 0 1 1 1;
               0 0 1 0 0 0;
               1 0 0 1 0 1;
               0 1 0 0 1 0;
               0 0 0 0 0 1;
               1 1 1 0 0 0;
               0 0 0 1 1 0;
               1 1 0 0 1 1;
               0 0 1 1 0 0;
               1 1 0 1 0 0;
               0 0 1 0 1 1]

  (m, n) = size(data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(data, settings.γ, settings.r)

  # remove unit 6 from cluster 3 and put it into another cluster
  R = [1; 1; 1; 2; 2; 3]

  i = 6
  hi = 3
  li = 3
  lj = 3

  state = Kpax3.AminoAcidState(data, R, priorR, priorC, settings)
  support = Kpax3.MCMCSupport(state, priorC)

  Kpax3.initsupportgibbs!(hi, priorR, state, support)

  support.t[lj] = Kpax3.clusterweight(1, state.k - 1, priorR)

  g = 0
  l = 0
  y = 0.0
  v = 0
  z = zeros(Float64, 2)
  for b in 1:m
    support.lpi[1, b] = priorC.logγ[1]
    support.lpi[2, b] = priorC.logγ[2]
    support.lpi[3, b] = priorC.logγ[3]

    g = 1
    l = 1
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmerge!(y, v, b, l, g, support.k, priorC, support)

    g = 2
    l = 2
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmerge!(y, v, b, l, g, support.k, priorC, support)
  end

  Kpax3.gibbscomputeprobi!(lj, support)

  for b in 1:support.m
    Kpax3.gibbscomputeprobg!(b, 1, support)
    Kpax3.gibbscomputeprobg!(b, 2, support)
  end

  # R = [1; 1; 1; 2; 2; 1]
  l = 1

  Kpax3.gibbsmerge!(i, hi, support.cl[l], l, data, priorR, support, state)

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 2; 1], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)

  @test state.R == st.R
  @test state.k == st.k

  @test !state.emptycluster[1]
  @test state.cl[1] == st.cl[1]
  @test state.v[1] == st.v[1]
  @test state.n1s[1, :] == st.n1s[1, :]
  @test state.unit[1][1:state.v[1]] == [1; 2; 3; 6]

  @test !state.emptycluster[2]
  @test state.cl[2] == st.cl[2]
  @test state.v[2] == st.v[2]
  @test state.n1s[2, :] == st.n1s[2, :]
  @test state.unit[2][1:state.v[2]] == [4; 5]

  @test state.emptycluster[3]

  @test isapprox(state.logpR, st.logpR, atol=ε)

  @test isapprox(support.logmlik, su.logmlik, atol=ε)

  @test support.lp[:, 1:2, :] == su.lp[:, 1:2, :]

  nothing
end

test_mcmc_gibbs_perform_merge()

function test_mcmc_gibbs_perform_split()
  ifile = "data/read_proper_aa.fasta"
  ofile = "../build/test"

  # state should be resized
  settings = Kpax3.KSettings(ifile, ofile, maxclust=3, maxunit=1)

  data = UInt8[0 0 0 0 0 1;
               1 1 1 1 1 0;
               0 0 1 0 1 1;
               1 1 0 1 0 0;
               1 1 0 0 0 0;
               0 0 0 1 1 0;
               1 1 1 0 0 0;
               0 0 0 1 1 1;
               0 0 1 0 0 0;
               1 0 0 1 0 1;
               0 1 0 0 1 0;
               0 0 0 0 0 1;
               1 1 1 0 0 0;
               0 0 0 1 1 0;
               1 1 0 0 1 1;
               0 0 1 1 0 0;
               1 1 0 1 0 0;
               0 0 1 0 1 1]

  (m, n) = size(data)

  priorR = Kpax3.EwensPitman(settings.α, settings.θ)
  priorC = Kpax3.AminoAcidPriorCol(data, settings.γ, settings.r)

  # remove unit 5 from cluster 2 and put it into another cluster
  R = [1; 1; 1; 2; 2; 3]

  i = 5
  hi = 2
  li = 3
  lj = 4

  state = Kpax3.AminoAcidState(data, R, priorR, priorC, settings)
  support = Kpax3.MCMCSupport(state, priorC)

  Kpax3.resizesupport!(support, state.k + 1)

  support.ni = vec(float(data[:, 4]))
  support.vi = 1

  Kpax3.initsupportgibbs!(hi, priorR, state, support)

  support.t[li] = Kpax3.clusterweight(support.vi, priorR)
  support.t[lj] = Kpax3.clusterweight(1, state.k, priorR)

  g = 0
  l = 0
  y = 0.0
  v = 0
  z = zeros(Float64, 2)
  for b in 1:m
    support.lpi[1, b] = priorC.logγ[1]
    support.lpi[2, b] = priorC.logγ[2]
    support.lpi[3, b] = priorC.logγ[3]
    support.lpi[4, b] = priorC.logγ[3]

    g = 1
    l = 1
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmove!(y, v, b, l, g, state.k, priorC, support)

    g = 3
    l = 2
    y = state.n1s[g, b] + float(data[b, i])
    v = state.v[g] + 1

    Kpax3.gibbsupdateclustergmove!(y, v, b, l, g, state.k, priorC, support)

    Kpax3.gibbsupdateclusteri!(state.k, b, li, priorC, support)
    Kpax3.gibbsupdateclusterj!(data[b, i], state.k + 1, b, li, lj, priorC,
                             support)
  end

  Kpax3.gibbscomputeprobi!(li, support)

  for b in 1:support.m
    Kpax3.gibbscomputeprobg!(b, 1, li, support)
    Kpax3.gibbscomputeprobg!(b, 2, li, support)
    Kpax3.gibbscomputeprobj!(b, li, lj, support)
  end

  # R = [1; 1; 1; 2; 4; 3]
  l = 4

  Kpax3.gibbssplit!(i, hi, li, lj, data, priorR, settings, support, state)

  st = Kpax3.AminoAcidState(data, [1; 1; 1; 2; 4; 3], priorR, priorC, settings)
  su = Kpax3.MCMCSupport(st, priorC)

  @test state.R == st.R
  @test state.k == st.k

  @test !state.emptycluster[1]
  @test state.cl[1] == st.cl[1]
  @test state.v[1] == st.v[1]
  @test state.n1s[1, :] == st.n1s[1, :]
  @test state.unit[1][1:state.v[1]] == [1; 2; 3]

  @test !state.emptycluster[2]
  @test state.cl[2] == st.cl[2]
  @test state.v[2] == st.v[2]
  @test state.n1s[2, :] == st.n1s[2, :]
  @test state.unit[2][1:state.v[2]] == [4]

  @test !state.emptycluster[3]
  @test state.cl[3] == st.cl[3]
  @test state.v[3] == st.v[3]
  @test state.n1s[3, :] == st.n1s[3, :]
  @test state.unit[3][1:state.v[3]] == [6]

  @test !state.emptycluster[4]
  @test state.cl[4] == st.cl[4]
  @test state.v[4] == st.v[4]
  @test state.n1s[4, :] == st.n1s[4, :]
  @test state.unit[4][1:state.v[4]] == [i]

  @test isapprox(state.logpR, st.logpR, atol=ε)

  @test isapprox(support.logmlik, su.logmlik, atol=ε)

  @test support.lp[:, 1:4, :] == su.lp

  nothing
end

test_mcmc_gibbs_perform_split()


"""
    DecisionVariable <: AbstractVariableRef

Holds a reference to the stochastic program and the corresponding MOI.VariableIndex.
"""
struct DecisionVariable <: JuMP.AbstractVariableRef
    stochasticprogram::StochasticProgram
    index::MOI.VariableIndex
end

# Getters #
# ========================== #
function get_decisions(dvar::DecisionVariable)
    return decisions(structure(owner_model(dvar)))
end
"""
    decision(dvar::DecisionVariable)

Return the internal `Decision` associated with `dvar`.
"""
function decision(dvar::DecisionVariable)
    decisions = get_decisions(dvar)::Decisions
    return decision(decisions, index(dvar))
end
"""
    state(dvar::DecisionVariable)

Return the `DecisionState` of `dvar`.
"""
function state(dvar::DecisionVariable)
    return decision(dvar).state
end

# Setters #
# ========================== #
function take_decisions!(stochasticprogram::StochasticProgram, dvars::Vector{DecisionVariable}, vals::AbstractVector)
    # Check that all given decisions are in model
    map(dvar -> check_belongs_to_model(dvar, stochasticprogram), dvars)
    # Check decision length
    length(dvars) == length(vals) || error("Given decision of length $(length(vals)) not compatible with number of decision variables $(length(dvars)).")
    # Update decisions
    for (dvar, val) in zip(dvars, vals)
        d = decision(dvar)
        # Update state
        d.state = Taken
        # Update value
        d.value = val
    end
    # Update objective and constraints in model
    update_decisions!(stochasticprogram, DecisionsStateChange())
    return nothing
end

function untake_decisions!(model::JuMP.Model, dvars::Vector{DecisionVariable})
    # Check that all given decisions are in model
    map(dvar -> check_belongs_to_model(dvar, model), dvars)
    # Update decisions
    need_update = false
    for dvar in dvars
        d = decision(dvar)
        if state(d) == Taken
            need_update |= true
            # Update state
            d.state = NotTaken
        end
    end
    # Update objective and constraints in model (if needed)
    need_update && update_decisions!(stochasticprogram, DecisionsStateChange())
    return nothing
end

function update_decisions!(stochasticprogram::StochasticProgram, change::DecisionModification)
    update_decisions!(structure(stochasticprogram), change)
end

# JuMP overloads #
# ========================== #
function JuMP.delete(stochasticprogram::StochasticProgram, dvar::DecisionVariable)
    if stochasticprogram !== owner_model(dvar)
        error("The decision variable you are trying to delete does not " *
              "belong to the stochastic program.")
    end
    MOI.delete(structure(stochasticprogram), index(dvar))
end

function delete(stochasticprogram::StochasticProgram, con_ref::ConstraintRef{<:StochasticProgram})
    if stochasticprogram !== con_ref.model
        error("The constraint reference you are trying to delete does not " *
              "belong to the stochasticprogram.")
    end
    MOI.delete(structure(sto), index(con_ref))
end

# MOI #
# ========================== #
function MOI.get(stochasticprogram::StochasticProgram, attr::MOI.AbstractVariableAttribute,
                 dvar::DecisionVariable)
    check_belongs_to_model(dvar, stochasticprogram)
    if MOI.is_set_by_optimize(attr)
        # Check if there is a cached solution
        cache = solutioncache(stochasticprogram)
        if haskey(cache, :solution)
            # Returned cached solution
            return MOI.get(cache[:solution], attr, index(dvar))
        end
        check_provided_optimizer(stochasticprogram.optimizer)
        if MOI.get(stochasticprogram, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED
            throw(OptimizeNotCalled())
        end
        return MOI.get(optimizer(stochasticprogram), attr, index(dvar))
    end
    return MOI.get(structure(stochasticprogram), attr, index(dvar))
end
function MOI.get(stochasticprogram::StochasticProgram, attr::MOI.AbstractConstraintAttribute,
                 cr::ConstraintRef{<:StochasticProgram})
    check_belongs_to_model(cr, stochasticprogram)
    if MOI.is_set_by_optimize(attr)
        # Check if there is a cached solution
        cache = solutioncache(stochasticprogram)
        if haskey(cache, :solution)
            # Returned cached solution
            return MOI.get(cache[:solution], attr, index(cr))
        end
        check_provided_optimizer(stochasticprogram.optimizer)
        if MOI.get(stochasticprogram, MOI.TerminationStatus()) == MOI.OPTIMIZE_NOT_CALLED
            throw(OptimizeNotCalled())
        end
        return MOI.get(optimizer(stochasticprogram), attr, index(cr))
    else
        return MOI.get(structure(stochasticprogram), attr, index(cr))
    end
end

function MOI.set(stochasticprogram::StochasticProgram, attr::MOI.AbstractVariableAttribute,
                 dvar::DecisionVariable, value)
    check_belongs_to_model(dvar, stochasticprogram)
    MOI.set(structure(stochasticprogram), attr, index(dvar), value)
    return nothing
end
function MOI.set(stochasticprogram::StochasticProgram, attr::MOI.AbstractConstraintAttribute,
                 cr::ConstraintRef, value)
    check_belongs_to_model(cr, stochasticprogram)
    MOI.set(structure(stochasticprogram), attr, index(cr), value)
end

# JuMP variable interface #
# ========================== #
JuMP.name(dvar::DecisionVariable) = MOI.get(owner_model(dvar), MOI.VariableName(), dvar)::String
function JuMP.set_name(dvar::DecisionVariable, name::String)
    return MOI.set(owner_model(dvar), MOI.VariableName(), dvar, name)
end

function decision_by_name(stochasticprogram::StochasticProgram, name::String)
    index = MOI.get(structure(stochasticprogram), MOI.VariableIndex, name)
    if index isa Nothing
        return nothing
    else
        return DecisionVariable(stochasticprogram, index)
    end
end

JuMP.index(dvar::DecisionVariable) = dvar.index

function JuMP.value(dvar::DecisionVariable; result::Int = 1)::Float64
    d = decision(dvar)
    if d.state == Taken
        # If decision has been fixed the value can be fetched
        # directly
        return d.value
    end
    return MOI.get(owner_model(dvar), MOI.VariablePrimal(result), dvar)
end

function JuMP.is_fixed(dvar::DecisionVariable)
    if state(dvar) == Taken
        return true
    end
    return false
end
"""
    unfix(dvar::DecisionVariable)

Unfix the decision associated with `dvar`. If the decision is already in a `NotTaken` state, this does nothing.

See also [`fix`](@ref).
"""
function JuMP.unfix(dvar::DecisionVariable)
    if state(dvar) == NotTaken
        # Nothing to do, just return
        return nothing
    end
    d = decision(dvar)
    # Update state
    d.state = NotTaken
    # Prepare modification
    change = DecisionStateChange(index(dvar), NotTaken, -value(dvar))
    # Update objective and constraints
    update_decisions!(JuMP.owner_model(dvar), change)
    return nothing
end
"""
    fix(dvar::DecisionVariable, val::Number)

Fix the decision associated with `dvar` to `val`. In contexts where `dvar` is a variable, the variable is fixed to the value. In contexts where `dvar` is a known parameter value, the value is updated.

See also [`unfix`](@ref).
"""
function JuMP.fix(dvar::DecisionVariable, val::Number)
    d = decision(dvar)
    if state(dvar) == NotTaken
        # Prepare modification
        change = DecisionStateChange(index(dvar), Taken, val)
        # Update state
        d.state = Taken
        # Update value
        d.value = val
    else
        # Prepare modification
        change = DecisionStateChange(index(dvar), Taken, val - d.value)
        # Just update value
        d.value = val
    end
    # Update objective and constraints
    update_decisions!(JuMP.owner_model(dvar), change)
    return nothing
end

JuMP.owner_model(dvar::DecisionVariable) = dvar.stochasticprogram

struct DecisionVariableNotOwned <: Exception
    dvar::DecisionVariable
end

function JuMP.check_belongs_to_model(dvar::DecisionVariable, stochasticprogram::StochasticProgram)
    if owner_model(dvar) !== stochasticprogram
        throw(DecisionNotOwned(dvar))
    end
end

function JuMP.is_valid(stochasticprogram::StochasticProgram, dvar::DecisionVariable)
    return stochasticprogram === owner_model(dvar)
end

function JuMP.has_lower_bound(dvar::DecisionVariable)
    index = MOI.ConstraintIndex{SingleDecision, MOI.GreaterThan{Float64}}(index(dvar).value)
    return MOI.is_valid(structure(owner_model(dvar)), index)
end
function JuMP.LowerBoundRef(dvar::DecisionVariable)
    moi_lb =  MOI.ConstraintIndex{SingleDecision, MOI.GreaterThan{Float64}}
    index = moi_lb(index(dvar).value)
    sp = owner_model(dvar)
    SP = typeof(sp)
    return ConstraintRef{SP, moi_lb, ScalarShape}(sp,
                                                  index,
                                                  ScalarShape())
end
function JuMP.set_lower_bound(dvar::DecisionVariable, lower::Number)
    new_set = MOI.GreaterThan(convert(Float64, lower))
    if has_lower_bound(dvar)
        cindex = MOI.ConstraintIndex{SingleDecision, MOI.GreaterThan{Float64}}(index(dvar).value)
        MOI.set(structure(owner_model(dvar)), MOI.ConstraintSet(), cindex, new_set)
    else
        MOI.add_constraint(structure(owner_model(dvar)), SingleDecision(index(dvar)), new_set)
    end
    return nothing
end
function JuMP.delete_lower_bound(dvar::DecisionVariable)
    JuMP.delete(owner_model(dvar), LowerBoundRef(dvar))
end
function JuMP.lower_bound(dvar::DecisionVariable)
    if !has_lower_bound(dvar)
        error("Decision variable $(dvar) does not have a lower bound.")
    end
    cset = MOI.get(structure(owner_model(dvar)), MOI.ConstraintSet(),
                   LowerBoundRef(dvar))::MOI.GreaterThan{Float64}
    return cset.lower
end

function JuMP.has_upper_bound(dvar::DecisionVariable)
    index = MOI.ConstraintIndex{SingleDecision, MOI.LessThan{Float64}}(index(dvar).value)
    return MOI.is_valid(structure(owner_model(dvar)), index)
end
function JuMP.UpperBoundRef(dvar::DecisionVariable)
    moi_ub =  MOI.ConstraintIndex{SingleDecision, MOI.LessThan{Float64}}
    index = moi_ub(index(dvar).value)
    sp = owner_model(dvar)
    SP = typeof(sp)
    return ConstraintRef{SP, moi_ub, ScalarShape}(sp,
                                                  index,
                                                  ScalarShape())
end
function JuMP.set_upper_bound(dvar::DecisionVariable, lower::Number)
    new_set = MOI.LessThan(convert(Float64, lower))
    if has_upper_bound(dvar)
        cindex = MOI.ConstraintIndex{SingleDecision, MOI.LessThan{Float64}}(index(dvar).value)
        MOI.set(structure(owner_model(dvar)), MOI.ConstraintSet(), cindex, new_set)
    else
        MOI.add_constraint(structure(owner_model(dvar)), SingleDecision(index(dvar)), new_set)
    end
    return nothing
end
function JuMP.delete_upper_bound(dvar::DecisionVariable)
    JuMP.delete(owner_model(dvar), UpperBoundRef(dvar))
end
function JuMP.upper_bound(dvar::DecisionVariable)
    if !has_upper_bound(dvar)
        error("Decision $(dvar) does not have a upper bound.")
    end
    cset = MOI.get(structure(owner_model(dvar)), MOI.ConstraintSet(),
                   UpperBoundRef(dvar))::MOI.LessThan{Float64}
    return cset.upper
end

function JuMP.is_integer(dvar::DecisionVariable)
    index = MOI.ConstraintIndex{SingleDecision, MOI.Integer}(index(dvar).value)
    return MOI.is_valid(structure(owner_model(dvar)), index)
end
function JuMP.set_integer(dvar::DecisionVariable)
    if is_integer(dvar)
        return nothing
    elseif is_binary(dvar)
        error("Cannot set the decision $(dvar) to integer as it is already binary.")
    else
        MOI.add_constraint(structure(owner_model(dvar)), SingleDecision(index(dvar)), MOI.Integer())
    end
end
function unset_integer(dvar::DecisionVariable)
    JuMP.delete(owner_model(dvar), IntegerRef(dvar))
    return nothing
end
function JuMP.IntegerRef(dvar::DecisionVariable)
    moi_int =  MOI.ConstraintIndex{SingleDecision, MOI.Integer}
    index = moi_int(index(dvar).value)
    sp = owner_model(dvar)
    SP = typeof(sp)
    return ConstraintRef{SP, moi_int, ScalarShape}(sp,
                                                   index,
                                                   ScalarShape())
end

function JuMP.is_binary(dvar::DecisionVariable)
    index = MOI.ConstraintIndex{SingleDecision, MOI.ZeroOne}(index(dvar).value)
    return MOI.is_valid(structure(owner_model(dvar)), index)
end
function JuMP.set_binary(dvar::DecisionVariable)
    if is_binary(dvar)
        return nothing
    elseif is_integer(dvar)
        error("Cannot set the decision $(dvar) to binary as it is already integer.")
    else
        MOI.add_constraint(structure(owner_model(dvar)), SingleDecision(index(dvar)), MOI.ZeroOne())
    end
end
function unset_binary(dvar::DecisionVariable)
    JuMP.delete(owner_model(dvar), BinaryRef(dvar))
    return nothing
end
function JuMP.BinaryRef(dvar::DecisionVariable)
    moi_bin =  MOI.ConstraintIndex{SingleDecision, MOI.ZeroOne}
    index = moi_bin(index(dvar).value)
    sp = owner_model(dvar)
    SP = typeof(sp)
    return ConstraintRef{SP, moi_bin, ScalarShape}(sp,
                                                   index,
                                                   ScalarShape())
end

function JuMP.start_value(dvar::DecisionVariable)
    return MOI.get(owner_model(dvar), MOI.VariablePrimalStart(), dvar)
end
function JuMP.set_start_value(dvar::DecisionVariable, value::Number)
    MOI.set(owner_model(dvar), MOI.VariablePrimalStart(), dvar, Float64(value))
end

function Base.hash(dvar::DecisionVariable, h::UInt)
    return hash(objectid(owner_model(dvar)), hash(dvar.index, h))
end

JuMP.isequal_canonical(d::DecisionVariable, other::DecisionVariable) = isequal(d, other)
function Base.isequal(dvar::DecisionVariable, other::DecisionVariable)
    return owner_model(dvar) === owner_model(other) && dvar.index == other.index
end

Base.iszero(::DecisionVariable) = false
Base.copy(dvar::DecisionVariable) = DecisionVariable(dvar.stochasticprogram, dvar.index)
Base.broadcastable(dvar::DecisionVariable) = Ref(dvar)


function has_quadratic(model::LinQuadOptimizer)
    return model.obj_type == QuadraticObjective ||
        length(cmap(model).q_less_than) > 0 ||
        length(cmap(model).q_greater_than) > 0 ||
        length(cmap(model).q_equal_to) > 0
end

#=
    Optimize the model
=#

function MOI.optimize!(model::LinQuadOptimizer)
    # reset storage
    fill!(model.variable_primal_solution, NaN)
    fill!(model.variable_dual_solution, NaN)
    fill!(model.constraint_primal_solution, NaN)
    fill!(model.constraint_dual_solution, NaN)
    model.primal_status = MOI.UnknownResultStatus
    model.dual_status   = MOI.UnknownResultStatus
    model.primal_result_count = 0
    model.dual_result_count = 0

    start_time = time()
    if has_integer(model)
        solve_mip_problem!(model)
    elseif has_quadratic(model)
        solve_quadratic_problem!(model)
    else
        solve_linear_problem!(model)
    end
    model.solvetime = time() - start_time

    # termination_status
    model.termination_status = get_termination_status(model)
    model.primal_status = get_primal_status(model)
    model.dual_status = get_dual_status(model)

    if model.primal_status in [MOI.FeasiblePoint, MOI.InfeasiblePoint]
        get_variable_primal_solution!(model, model.variable_primal_solution)
        get_linear_primal_solution!(model, model.constraint_primal_solution)
        if has_quadratic(model)
            get_quadratic_primal_solution!(model, model.qconstraint_primal_solution)
        end
        model.primal_result_count = 1
    elseif model.primal_status == MOI.InfeasibilityCertificate
        get_unbounded_ray!(model, model.variable_primal_solution)
        model.primal_result_count = 1
    end
    if model.dual_status in [MOI.FeasiblePoint, MOI.InfeasiblePoint]
        get_variable_dual_solution!(model, model.variable_dual_solution)
        get_linear_dual_solution!(model, model.constraint_dual_solution)
        if has_quadratic(model)
            get_quadratic_dual_solution!(model, model.qconstraint_dual_solution)
        end
        model.dual_result_count = 1
    elseif model.dual_status == MOI.InfeasibilityCertificate
        get_farkas_dual!(model, model.constraint_dual_solution)
        model.dual_result_count = 1
    end

    if MOI.get(model, MOI.ObjectiveSense()) == MOI.MaxSense
        model.constraint_dual_solution *= -1
        model.variable_dual_solution *= -1
    end
end

#=
    Result Count
=#
MOI.canget(::LinQuadOptimizer, ::MOI.ResultCount) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.ResultCount)
    return max(model.primal_result_count, model.dual_result_count)
end

#=
    Termination status
=#
MOI.canget(::LinQuadOptimizer, ::MOI.TerminationStatus) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.TerminationStatus)
    return model.termination_status
end

#=
    Primal status
=#
function MOI.canget(model::LinQuadOptimizer,  status::MOI.PrimalStatus)
    if status.N != 1
        error("Multiple primal statuses are not supported.")
    end
    return model.primal_result_count >= status.N
end
function MOI.get(model::LinQuadOptimizer, ::MOI.PrimalStatus)
    return model.primal_status
end

#=
    Dual status
=#
function MOI.canget(model::LinQuadOptimizer,  status::MOI.DualStatus)
    if status.N != 1
        error("Multiple dual statuses are not supported.")
    end
    return model.dual_result_count >= status.N
end
function MOI.get(model::LinQuadOptimizer, ::MOI.DualStatus)
    return model.dual_status
end

#=
    Objective Value
=#
function MOI.canget(::LinQuadOptimizer, attr::MOI.ObjectiveValue)
    return attr.resultindex == 1
end
function MOI.get(model::LinQuadOptimizer, attr::MOI.ObjectiveValue)
    if attr.resultindex == 1
        # Note: we add m.objective_constant here to account for any constant
        # term which was not passed to the solver itself (and which therefore
        # would not be accounted for in `get_objective_value(m)`. We do *not*
        # call `get_constant_objective(m)` because that would also pull any
        # constants which were passed to the solver, resulting those constants
        # being counted twice. This confusion will be alleviated when all LQOI
        # solvers implement `get_constant_objective()` and
        # `set_constant_objective!()` by actually passing the relevant constants
        # to the solvers, at which point we can just get rid of
        # m.objective_constant entirely.
        return get_objective_value(model) + model.objective_constant
    else
        error("Unable to access multiple objective values")
    end
end

#=
    Variable Primal solution
=#
MOI.canget(::LinQuadOptimizer, ::MOI.VariablePrimal, ::Type{VarInd}) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.VariablePrimal, index::VarInd)
    column = get_column(model, index)
    return model.variable_primal_solution[column]
end

MOI.canget(::LinQuadOptimizer, ::MOI.VariablePrimal, ::Type{Vector{VarInd}}) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.VariablePrimal, indices::Vector{VarInd})
    MOI.get.(Ref(model), Ref(MOI.VariablePrimal()), indices)
end

#=
    Variable Dual solution
=#
"""
    is_binding(set, value::Float64)

Return true if `value` is an extreme point of the set `set`.
"""
is_binding(set::LE, value::Float64) = isapprox(set.upper, value)
is_binding(set::GE, value::Float64) = isapprox(set.lower, value)
is_binding(set::EQ, value::Float64) = isapprox(set.value, value)
is_binding(set::IV, value::Float64) = isapprox(set.lower, value) || isapprox(set.upper, value)

function MOI.canget(::LinQuadOptimizer, ::MOI.ConstraintDual, ::Type{SVCI{S}}) where S <: LinSets
    return true
end

function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintDual, index::SVCI{<: LinSets})
    column = get_column(model, model[index])
    # the variable reduced cost is only the constriant dual if the bound is active.
    set = MOI.get(model, MOI.ConstraintSet(), index)
    primal_value = model.variable_primal_solution[column]
    if is_binding(set, primal_value)
        return model.variable_dual_solution[column]
    else
        return 0.0
    end
end

function MOI.canget(s::LinQuadOptimizer, ::MOI.ConstraintDual,
                    ::Type{VVCI{S}}) where S <: Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives}
    return true
end
function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintDual, index::VVCI{<: Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives}})
    return [model.constraint_dual_solution[row] for row in model[index]]
end

#=
    Variable Bound Primal solution
=#

MOI.canget(::LinQuadOptimizer, ::MOI.ConstraintPrimal, ::Type{SVCI{S}}) where S <: LinSets = true

function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintPrimal, index::SVCI{<: LinSets})
    column = get_column(model, model[index])
    return model.variable_primal_solution[column]
end

function MOI.canget(::LinQuadOptimizer, ::MOI.ConstraintPrimal,
                   ::Type{VVCI{S}}) where S <: Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives}
    return true
end

function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintPrimal,
                 index::VVCI{<: Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives}})
    return [model.constraint_primal_solution[row] for row in model[index]]
end

#=
    Constraint Primal solution
=#

MOI.canget(m::LinQuadOptimizer, ::MOI.ConstraintPrimal, ::Type{LCI{S}}) where S <: LinSets = true

function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintPrimal, index::LCI{<: LinSets})
    row = model[index]
    return model.constraint_primal_solution[row] + model.constraint_constant[row]
end

MOI.canget(::LinQuadOptimizer, ::MOI.ConstraintPrimal, ::Type{QCI{S}}) where S <: LinSets = true

function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintPrimal, index::QCI{<: LinSets})
    row = model[index]
    return model.qconstraint_primal_solution[row]
end

MOI.canget(::LinQuadOptimizer, ::MOI.ConstraintPrimal, ::Type{VLCI{S}}) where S <: Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives} = true

# vector valued constraint duals
function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintPrimal, index::VLCI{<: Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives}})
    row = model[index]
    return model.constraint_primal_solution[row] + model.constraint_constant[row]
end

#=
    Constraint Dual solution
=#

__assert_dual_sense__(::LCI{LE}, dual) = @assert dual <= 0.0
__assert_dual_sense__(::LCI{GE}, dual) = @assert dual >= 0.0
__assert_dual_sense__(::LCI{IV}, dual) = nothing
__assert_dual_sense__(::LCI{EQ}, dual) = nothing

MOI.canget(::LinQuadOptimizer, ::MOI.ConstraintDual, ::Type{LCI{S}}) where S <: LinSets = true

function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintDual, index::LCI{<: LinSets})
    row = model[index]
    dual = model.constraint_dual_solution[row]
    __assert_dual_sense__(index, dual)
    return dual
end

MOI.canget(model::LinQuadOptimizer, ::MOI.ConstraintDual, ::Type{QCI{S}}) where S <: LinSets = true

function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintDual, index::QCI{<: LinSets})
    row = model[index]
    return model.qconstraint_dual_solution[row]
end

# vector valued constraint duals
MOI.canget(::LinQuadOptimizer, ::MOI.ConstraintDual, ::Type{VLCI{S}}) where S <: Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives} = true

function MOI.get(model::LinQuadOptimizer, ::MOI.ConstraintDual, index::VLCI{<: Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives}})
    rows = model[index]
    return model.constraint_dual_solution[rows]
end

#=
    Solution Attributes
=#

MOI.supports(::LinQuadOptimizer, ::MOI.ObjectiveBound) = true
MOI.canget(::LinQuadOptimizer, ::MOI.ObjectiveBound) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.ObjectiveBound)
    return get_objective_bound(model)
end

MOI.supports(::LinQuadOptimizer, ::MOI.RelativeGap) = true
MOI.canget(::LinQuadOptimizer, ::MOI.RelativeGap) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.RelativeGap)
    return get_relative_mip_gap(model)
end

MOI.supports(::LinQuadOptimizer, ::MOI.SolveTime) = true
MOI.canget(::LinQuadOptimizer, ::MOI.SolveTime) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.SolveTime)
    return model.solvetime
end

MOI.supports(::LinQuadOptimizer, ::MOI.SimplexIterations) = true
MOI.canget(::LinQuadOptimizer, ::MOI.SimplexIterations) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.SimplexIterations)
    return get_iteration_count(model)
end

MOI.supports(::LinQuadOptimizer, ::MOI.BarrierIterations) = true
MOI.canget(::LinQuadOptimizer, ::MOI.BarrierIterations) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.BarrierIterations)
    return get_barrier_iterations(model)
end

MOI.supports(::LinQuadOptimizer, ::MOI.NodeCount) = true
MOI.canget(::LinQuadOptimizer, ::MOI.NodeCount) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.NodeCount)
    return get_node_count(model)
end

MOI.supports(::LinQuadOptimizer, ::MOI.RawSolver) = true
MOI.canget(::LinQuadOptimizer, ::MOI.RawSolver) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.RawSolver)
    return model
end