/** Boost */
#include <boost/serialization/serialization.hpp>
#include <boost/serialization/unordered_map.hpp>
#include <boost/serialization/unordered_set.hpp>
#include <boost/serialization/utility.hpp>
#include <boost/serialization/vector.hpp>

#include <boost/archive/binary_iarchive.hpp>
#include <boost/archive/binary_oarchive.hpp>

#include <boost/program_options.hpp>
namespace po = boost::program_options;

/** STL */
#include <chrono>
#include <condition_variable>
#include <future>
#include <mutex>
#include <thread>
#include <vector>

/** Ligero */
#include "Common.hpp"

#include "EncryptedClient.hpp"
#include "EncryptedCoordinator.hpp"

#include "ZkArgument.hpp"
#include "protocol/ExpressNPStatement.hpp"

#include "FiniteFields.hpp"
#include "Hash.hpp"
#include "SecretSharingNTT.hpp"

using namespace ligero;
using namespace std::chrono_literals;

/** ===== Thread synchronization ===== */
std::mutex mutex;
std::condition_variable cv;
bool thread_aborted = false;

/** Parameters RSA Ceremony */
/** ======================= */

constexpr auto degree = 1 << 16;
constexpr auto p = 9;
constexpr auto q = 21;
constexpr auto sigma = 8;
constexpr auto lambda = 128;
constexpr auto tau_limit_bit = 175;
constexpr auto tau = 1000;
constexpr auto numCandidates = 2048;

/** throughput parties */
constexpr auto wait_timeout = 300 * 1000; /* wait 5 min when start */
constexpr auto timeout = 25 * 1000;       /* a little longer than test */
constexpr auto data_size = 32 * 1024;     /* 32 KB */
constexpr auto nb_max_send = 1024;
constexpr auto parties = 1;

/** throughtput coordinator */
constexpr auto pbs = 1000;
constexpr auto duration = 20s;
constexpr auto wait_time = 120s;
constexpr auto cleanup_time = 5s;
constexpr auto throughput_cutoff = 100 * 1024; /* 200 KB/s */

constexpr auto receive_timeout = 60min;

/** Parameters Zero-Knowledge */
/** ========================= */

/** Settings Zero-Knowledge */
RegularInteger p_zk(10514644122014433281ULL);
typedef Fpp_Fixed<p_zk> FieldT;

zksnark::InterleavedRSCode irs;        /** Reed-Solomon Interleaved Code */
ligero::MTParameters<FieldT> mtParams; /** Merkle Tree Parameters */
constexpr size_t l = 8192;             // 4096;             /** Bloc size */
constexpr size_t k = 16384;            // 8192;             /** Degree */
constexpr size_t n = 65536;            // 16384;            /** # of Shares */

/** Columns to Open Up (By default half the columns) */
constexpr size_t t = 189;
constexpr auto ptNumberEvaluation = 3;

/** Random Generator */
const hash::HashIntegerSeedGeneratorF<FieldT>& fieldGenerator =
    hash::blake2bIntegerSeedGenerator<FieldT>;
const hash::HashStringSeedGeneratorF& stringGenerator = hash::blake2bStringHash;

/** Threading */
constexpr size_t maxNumberThreads = 50;
/** [m5.metal has 96 threads over 48 cores] */

/** Selecting a port for distributed verifying */
constexpr char* port = "5556";

/** Depositing the data here */
PublicData pdata;
SigmaProtocolPublicData spdata;
std::vector<ligero::zksnark::FullFSTranscript<FieldT>> transcripts;
size_t availableTranscriptsIndex = 0;
bool firstIteration = true;

std::string localIpAddr = "UNKNOWN";

/** Stitching Proofs */
constexpr int nbIterations = 3;
constexpr int nbVarsStitched = 4;
constexpr int nbBlocksPerVar = degree / l;

/** Data */
std::string hashedPublicData;

void resetGlobalVariables() {
  pdata = PublicData{};
  transcripts.clear();
  availableTranscriptsIndex = 0;
  firstIteration = true;
  thread_aborted = false;
  hashedPublicData.clear();
}

/** This separate function handles the data coming in from the coordinator, as
 * well as sigma protocol verification. The concern is that otherwise we might
 * fall behind on receiving public data & proof because of processing delays.
 * @param transport the transport object handling communication with the
 * coordinator
 */
void collectData(ZeroMQClientTransport& transport) {
  DBG("Verifier " << transport.getSocketId() << " begining data acquisition.");

  /** Tell the gathering coodrinator we are ready */
  if (firstIteration) transport.send(MessageType::VERIFIER_READY, localIpAddr, false);

  /** Receive Sigma Protocol Public Data */
  auto maybe_sigma = transport.awaitReply<SigmaProtocolPublicData>(
      MessageType::GATHER_SIGMA_DATA, true);

  if (hasError(maybe_sigma)) {
    LOG(ERROR) << "Kill/Restart received";
    thread_aborted = true;
    cv.notify_one();
    return;
  } else {
    spdata = getResult(maybe_sigma);
  }

  /** Receive public data */
  auto maybe_pdata =
      transport.awaitReply<PublicData>(MessageType::GATHER_PUBLIC_DATA, true);

  if (hasError(maybe_pdata)) {
    LOG(ERROR) << "Kill/Restart received";
    thread_aborted = true;
    cv.notify_one();
    return;
  } else {
    pdata = getResult(maybe_pdata);
    DBG("Received and deserialized public data.");

    /** Hashing Public Data */
    pdata.modulusIdx = 0;
    pdata.roots.clear();
    pdata.ptNumberEvaluation = ptNumberEvaluation;
    hashedPublicData = hash::hashPublicData(stringGenerator, pdata, spdata);
  }

  /** Verify Sigma Protocol */
  LOG(INFO) << "Verifying Sigma protocol proof";

  auto [alphasGCD, _drop_val] =
      math::fixed_bucket_n_primes(3210, primesGCD, 175);
  std::vector<mpz_class> coefs;
  std::vector<mpz_class> nfl_primes;
  std::vector<mpz_class> nfl_primes_full;
  for (int i = 0; i < 9; i++)
    nfl_primes.emplace_back(mpz_class(nfl::params<uint64_t>::P[i]));
  for (int i = 0; i < NbPrimesQ; i++)
    nfl_primes_full.emplace_back(mpz_class(nfl::params<uint64_t>::P[i]));

  /** Recomputing the final candidate from the CRT factors */
  std::vector<mpz_class> finalModuli_GCD;
  finalModuli_GCD.resize(alphasGCD.size());
  for (int alphaGCDidx = 0; alphaGCDidx < alphasGCD.size(); alphaGCDidx++) {
    std::vector<mpz_class> prime_decomp_f;
    for (int i = 0; i < NbPrimesQ; i++) {
      prime_decomp_f.emplace_back(
          mpz_class(pdata.finalModuli_GCD[alphaGCDidx * NbPrimesQ + i]));
    }
    finalModuli_GCD[alphaGCDidx] =
        math::crt_reconstruct(prime_decomp_f, coefs, nfl_primes_full);
  }
  std::vector<mpz_class> candidates;
  candidates.emplace_back(
      math::crt_reconstruct(finalModuli_GCD, coefs, alphasGCD));

  /** Recomputing the gamma values */
  mpz_class gammaSeed = pdata.gammaSeed;
  std::vector<mpz_class> gammaValues;
  gammaValues.resize(candidates.size() * 128);
  std::vector<mpz_class> engine =
      math::generateRandomVector(gammaSeed, kJacobiNumberOfRandomValues, 2048);
  mpz_t one;
  mpz_init(one);
  mpz_set_ui(one, 1);
  size_t engineCount = 0;

  {
    mpz_t gcdResult;
    mpz_init(gcdResult);

    for (int i = 0; i < candidates.size(); ++i) {
      const mpz_class& N = candidates[i];

      DBG(" N % 2 = 1; N = " << N);
      /** We expect that N is 3 (mod 4), and therefore must be odd. */
      assert(N % 2 == 1);

      int foundGamma = 0;
      while (foundGamma < 128) {
        mpz_class g = engine[engineCount];
        ++engineCount;
        if (engineCount >= engine.size()) {
          engineCount = 0;
          DBG("Regenerating random values for Jacobi (GCD) gamma values");
          engine = math::generateRandomVector(
              engine[0], kJacobiNumberOfRandomValues, 2048);
        }
        g = g % N;
        assert(g != 0);

        DBG("g = " << g);
        if (mpz_jacobi(g.get_mpz_t(), N.get_mpz_t()) == 1) {
          gammaValues[i * 128 + foundGamma] = g;
          foundGamma++;
        }
      }

      /** Sanity check, assert if gcd(gammaValues[i], N) == 1 */
      for (int j = 0; j < 128; ++j) {
        mpz_gcd(gcdResult, gammaValues[i * 128 + j].get_mpz_t(), N.get_mpz_t());
        assert(mpz_cmp(gcdResult, one) == 0);
      }
    }

    mpz_clear(one);
    mpz_clear(gcdResult);
  }
  for (int j = 0; j < 128; j++) {
    std::vector<mpz_class> sigma_a_GCD;
    sigma_a_GCD.resize(alphasGCD.size());
    std::vector<mpz_class> sigma_e_GCD;
    sigma_e_GCD.resize(alphasGCD.size());
    std::vector<mpz_class> sigma_z_GCD;
    sigma_z_GCD.resize(alphasGCD.size());
    std::vector<mpz_class> sigma_g_GCD;
    sigma_g_GCD.resize(alphasGCD.size());
    for (int alphaGCDidx = 0; alphaGCDidx < alphasGCD.size(); alphaGCDidx++) {
      std::vector<mpz_class> prime_decomp_a;
      std::vector<mpz_class> prime_decomp_e;
      std::vector<mpz_class> prime_decomp_z;
      std::vector<mpz_class> prime_decomp_g;
      for (int i = 0; i < 9; i++) {
        prime_decomp_a.emplace_back(
            mpz_class(spdata.sigmaaGCD[i * 128 * alphasGCD.size() +
                                       j * alphasGCD.size() + alphaGCDidx]));
        prime_decomp_e.emplace_back(
            mpz_class(spdata.sigmaeGCD[i * 128 * alphasGCD.size() +
                                       j * alphasGCD.size() + alphaGCDidx]));
        prime_decomp_z.emplace_back(
            mpz_class(spdata.sigmazGCD[i * 128 * alphasGCD.size() +
                                       j * alphasGCD.size() + alphaGCDidx]));
        prime_decomp_g.emplace_back(
            mpz_class(spdata.sigmagGCD[i * 128 * alphasGCD.size() +
                                       j * alphasGCD.size() + alphaGCDidx]));
      }
      sigma_a_GCD[alphaGCDidx] =
          math::crt_reconstruct(prime_decomp_a, coefs, nfl_primes);
      sigma_e_GCD[alphaGCDidx] =
          math::crt_reconstruct(prime_decomp_e, coefs, nfl_primes);
      sigma_z_GCD[alphaGCDidx] =
          math::crt_reconstruct(prime_decomp_z, coefs, nfl_primes);
      sigma_g_GCD[alphaGCDidx] =
          math::crt_reconstruct(prime_decomp_g, coefs, nfl_primes);
    }
    mpz_class sigmaa = math::crt_reconstruct(sigma_a_GCD, coefs, alphasGCD);
    mpz_class sigmae = math::crt_reconstruct(sigma_e_GCD, coefs, alphasGCD);
    mpz_class sigmaz = math::crt_reconstruct(sigma_z_GCD, coefs, alphasGCD);
    mpz_class sigmag = math::crt_reconstruct(sigma_g_GCD, coefs, alphasGCD);
    mpz_class sigma_lhs = math::powm(gammaValues[j], sigmaz, candidates[0]);
    mpz_class sigma_rhs = math::mod(
        sigmaa * math::powm(sigmag, sigmae, candidates[0]), candidates[0]);

    assert(sigma_lhs == sigma_rhs);
  }

  LOG(INFO) << "Verified Sigma protocol proof";

  for (size_t modulusIdx = 0; modulusIdx < q; modulusIdx++) {
    /** Receive proofs */
    auto maybe_proof =
        transport.awaitReply<ligero::zksnark::FullFSTranscript<FieldT>>(
            MessageType::GATHER_PROOFS, true);

    if (hasError(maybe_proof)) {
      LOG(ERROR) << "Kill/Restart received";
      thread_aborted = true;
      cv.notify_one();
      return;
    } else {
      /** get lock and notify verifier */
      std::unique_lock<std::mutex> lock(mutex);

      transcripts.push_back(getResult(maybe_proof));
      availableTranscriptsIndex++;

      lock.unlock();
      cv.notify_one();
    }
  }

  LOG(INFO) << "Got all proofs";

  return;
}

/** This function verifies the proofs as they are gathered as well as
 * cross-moduli integrity constraints ("stitching")
 * @param transport the transport object handling communication with the
 * coordinator
 * @param myTimers reference to the timers to use for this run
 * @return Unit if successful
 */
expected<Unit> verify(ZeroMQClientTransport& transport, timers& myTimers) {
  /** Reset all global variables */
  resetGlobalVariables();

  /** Verifying */
  /** ================================================================ */
  FieldT::modulusIdx_ = 0;
  size_t aggregateSuccess = 0;

  Statements statement = Statements::RSACeremony;
  Statements localStatement = statement;

  /** Primitive Roots of Unity for All Moduli */
  nfl::poly_p<uint64_t, degree, q> roots({0});
  for (size_t idx = 0; idx < q; idx++) {
    roots(idx, 1) = 1ull;
  }
  roots.ntt_pow_phi();

  ligero::zksnark::FullFSTranscript<FieldT> transcript;

  SecretData sdata;

  /** Dimensioning ai scalars */
  sdata.ai.resize(nbIterations);
  for (size_t iter = 0; iter < nbIterations; iter++) {
    sdata.ai[iter].resize(nbVarsStitched);
    for (size_t var = 0; var < nbVarsStitched; var++) {
      sdata.ai[iter][var].resize(nbBlocksPerVar);
      for (size_t block = 0; block < nbBlocksPerVar; block++) {
        sdata.ai[iter][var][block].resize(l);
      }
    }
  }

  /** Storing the results of the verification per modulus */
  std::vector<int> report(q, 0);

  /** Import Public Data */
  std::thread* dataCollection =
      new std::thread(collectData, std::ref(transport));

  /** CRT representation to check integrity for the proofs */
  std::vector<mpz_class> stitchingCRT;
  mpz_class stitching0;

  for (size_t modulusIdx = 0; modulusIdx < q; modulusIdx++) {
    std::string timerVerificationIdle =
        std::string("6.a. verification, idle, modulus ") +
        std::to_string(modulusIdx);
    myTimers.begin(1, timerVerificationIdle.c_str(), 0);

    {
      /** get lock and wait for complete */
      std::unique_lock<std::mutex> lock(mutex);
      auto received = cv.wait_for(lock, receive_timeout, [&] {
        return thread_aborted || (availableTranscriptsIndex > modulusIdx);
      });

      /** check if thread failed or timeout */
      /** we won't restart this verifier because the corresponding party will be
       */
      /** kicked out */
      if (thread_aborted || !received) {
        LOG(INFO) << "Working thread failed or timed out, quit verifier";
        dataCollection->join();
        delete dataCollection;
        return Error::TIMED_OUT;
      }

      lock.unlock();
    }

    myTimers.end(1, timerVerificationIdle.c_str(), 0, true);

    std::string timerVerificationCompute =
        std::string("6.a. verification, compute, modulus ") +
        std::to_string(modulusIdx);
    myTimers.begin(1, timerVerificationCompute.c_str(), 0);

    auto& transcript = transcripts[modulusIdx];
    bool success = true;

    LOG(INFO) << "Verifying modulus idx:" << modulusIdx;

    /** Zero-Knowledge */
    p_zk = nfl::params<uint64_t>::P[modulusIdx];
    FieldT::modulusIdx_ = modulusIdx;
    FieldT::preprocessing(log2(n) + 1);

    pdata.modulusIdx = modulusIdx;

    /** For portions 7 to 12, only run the first 9 moduli */
    if ((localStatement == Statements::RSACeremony) && (pdata.modulusIdx > 8))
      localStatement = Statements::RSACeremony_Rounds3to6;

    /** Roots of Unity */
    pdata.roots.assign(
        roots.poly_obj().data() + pdata.modulusIdx * (size_t)roots.degree,
        roots.poly_obj().data() +
            (pdata.modulusIdx + 1) * (size_t)roots.degree);
    pdata.ptNumberEvaluation = ptNumberEvaluation;

    zksnark::Verifier<FieldT> verifier(
        localStatement, irs, t, l, fieldGenerator, stringGenerator, mtParams,
        pdata, spdata, sdata, permut<k>::compute, permut<n>::compute,
        hashedPublicData);

    try {
      verifier.verifyProof(transcript);

      LOG(INFO) << "Verifier mod " << modulusIdx << ":" << verifier._stitching;
    } catch (failedTest& e) {
      DBG("Failed Test: " << e.what());
      success = false;
    }

    if (modulusIdx == 0) {
      stitching0 = verifier._stitching;
    } else {
      stitchingCRT.push_back(verifier._stitching);
    }

    if (success) {
      LOG(INFO) << "Verification for modulus idx:" << modulusIdx
                << " succeeded";
      report[modulusIdx] = 1;
      aggregateSuccess++;
    } else {
      LOG(INFO) << "Verification for modulus idx:" << modulusIdx << " failed";
      report[modulusIdx] = 0;
    }

    myTimers.end(1, timerVerificationCompute.c_str(), 0, true);
  }

  /** Cross-Moduli Verification */
  {
    bool success = false;
    std::vector<mpz_class> mods, coefs;
    for (size_t i = 1; i < q; i++) {
      mods.emplace_back(params<uint64_t>::P[i]);
    }
    success = (math::mod(stitching0 -
                             math::crt_reconstruct(stitchingCRT, coefs, mods),
                         mpz_class(params<uint64_t>::P[0])) == mpz_class(0));

    if (success) {
      LOG(INFO) << "Verification for cross-moduli integrity succeeded";
    } else {
      LOG(INFO) << "Verification for cross-moduli integrity failed";
      for (size_t idx = 0; idx < q; idx++) {
        report[idx] = 0;
      }
      aggregateSuccess = 0;
    }
  }

  dataCollection->join();
  delete dataCollection;

  /** Wait for the signal from the gathering coordinator */
  auto maybe_response_reports = transport.awaitReply();

  if (hasError(maybe_response_reports)) {
    if (getError(maybe_response_reports) == Error::RESTART) {
      LOG(INFO) << "Restart verifier";
      myTimers.reset();
      return verify(transport, myTimers);
    } else {
      LOG(INFO) << "Timed out";
      return Unit{};
    }
  } else {
    if (getResult(maybe_response_reports) != MessageType::GATHER_REPORTS) {
      LOG(ERROR) << "Out of sequence message received when expecting gathering "
                    "reports, aborting.";
      return Error::OUT_OF_SYNC;
    }
  }

  if (aggregateSuccess == q) {
    LOG(INFO) << "All proofs successfully verified.";
  } else {
    LOG(INFO) << "Failed to verify some of the proofs.";
  }

  /** Send the verifier report to the gathering coordinator */
  transport.send<std::vector<int>>(MessageType::VERIFIER_REPORT, report, false);

  /** Finally wait for restart message if any verification failed */
  auto maybeRestart = transport.awaitReply();

  if (hasError(maybeRestart) && getError(maybeRestart) == Error::RESTART) {
    LOG(INFO) << "Restart message received, restarting verifier";
    myTimers.reset();
    return verify(transport, myTimers);
  } else if (getResult(maybeRestart) == MessageType::END) {
    LOG(INFO) << "All done.";
    return Unit{};
  } else {
    LOG(INFO) << "Unknown error happened!";
    return Error::UNKNOWN_ERROR;
  }
}

int main(int argc, char** argv) {
  std::string coordinatorIP("127.0.0.1");

  po::options_description desc("Allowed options");
  desc.add_options()("help", "produce help message")(
      "localIPaddress", po::value<std::string>(), "sets the local IP address")(
      "ip", po::value<std::string>(), "coordinator IP address");

  po::variables_map vm;
  po::positional_options_description p;
  p.add("coordinator-ip", -1);

  po::store(
      po::command_line_parser(argc, argv).options(desc).positional(p).run(),
      vm);

  po::notify(vm);

  if (vm.count("help")) {
    LOG(INFO) << desc << "\n";
    LOG(INFO)
        << "USAGE: --localIPaddress <IP Distributed Verifier> <IP Coordinator>";
    return 0;
  }

  if (vm.count("localIPaddress")) {
    localIpAddr = vm["localIPaddress"].as<std::string>();
    LOG(INFO) << "Local IP Address " << localIpAddr;
  }

  if (vm.count("ip")) {
    coordinatorIP = vm["ip"].as<std::string>();
    LOG(INFO) << "Gathering Coordinator IP " << coordinatorIP;
  }

  SocketId socketId = boost::uuids::random_generator()();
  ZeroMQClientTransport transport(std::string("tcp://") + coordinatorIP +
                                      std::string(":") + std::string(port),
                                  socketId, receive_timeout);

  /** Set Zero-Knowledge Parameters */
  irs.k = k; /** Degree of the Polynomial */
  irs.n = n; /** Number of Shares */

  mtParams.digestLength = 32;
  mtParams.zkByteSize = 8;
  mtParams.hashInnerNodeF = hash::blake2bTwoToOneHash;
  mtParams.hashLeafContentF = hash::blake2bFieldElementHash<FieldT>;
  mtParams.hashZKRandomnessF = hash::blake2bZKElementHash;
  mtParams.leavesNumber = irs.n;

  /** Configure easylogging */
  configureEasyLogging("distributed_verifier");

  /** Add timers */
  timers myTimers;
  myTimers.initialize_timer();

  auto result = verify(transport, myTimers);

  if (hasError(result)) {
    return EXIT_FAILURE;
  }

  return 0;
}


#pragma once

#include <map>
#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/range/adaptor/reversed.hpp>
#include "libhmsbeagle/beagle.h"
#include "tree.hpp"
#include "tree_manip.hpp"
#include "data.hpp"
#include "model.hpp"
#include "xlorad.hpp"

#include "output_manager.hpp"
extern lorad::OutputManager om;

namespace lorad {

    class Likelihood {
        public:
                                                    Likelihood();
                                                    ~Likelihood();

            void                                    setRooted(bool is_rooted);
            void                                    setPreferGPU(bool prefer_gpu);
            void                                    setAmbiguityEqualsMissing(bool ambig_equals_missing);

            bool                                    usingStoredData() const;
            void                                    useStoredData(bool using_data);
            void                                    useUnderflowScaling(bool do_scaling);

            std::string                             beagleLibVersion() const;
            std::string                             availableResources() const;
            std::string                             usedResources() const;

            void                                    initBeagleLib();
            void                                    finalizeBeagleLib(bool use_exceptions);

            double                                  calcLogLikelihood(Tree::SharedPtr t);

            Data::SharedPtr                         getData();
            void                                    setData(Data::SharedPtr d);

            Model::SharedPtr                        getModel();
            void                                    setModel(Model::SharedPtr m);

            void                                    clear();

            unsigned                                calcNumEdgesInFullyResolvedTree() const;
            unsigned                                calcNumInternalsInFullyResolvedTree() const;

        private:

            struct InstanceInfo {
                int handle;
                int resourcenumber;
                std::string resourcename;
                unsigned nstates;
                unsigned nratecateg;
                unsigned npatterns;
                unsigned partial_offset;
                unsigned tmatrix_offset;
                bool invarmodel;
                std::vector<unsigned> subsets;

                InstanceInfo() : handle(-1), resourcenumber(-1), resourcename(""), nstates(0), nratecateg(0), npatterns(0), partial_offset(0), tmatrix_offset(0), invarmodel(false) {}
            };

            typedef std::pair<unsigned, int>        instance_pair_t;

            unsigned                                getScalerIndex(Node * nd, InstanceInfo & info) const;
            unsigned                                getPartialIndex(Node * nd, InstanceInfo & info) const;
            unsigned                                getTMatrixIndex(Node * nd, InstanceInfo & info, unsigned subset_index) const;
            void                                    updateInstanceMap(instance_pair_t & p, unsigned subset);
            void                                    newInstance(unsigned nstates, int nrates, std::vector<unsigned> & subset_indices);
            void                                    setTipStates();
            void                                    setTipPartials();
            void                                    setPatternPartitionAssignments();
            void                                    setPatternWeights();
            void                                    setAmongSiteRateHeterogenetity();
            void                                    setModelRateMatrix();
            void                                    addOperation(InstanceInfo & info, Node * nd, Node * lchild, Node * rchild, unsigned subset_index);
            void                                    queuePartialsRecalculation(Node * nd, Node * lchild, Node * rchild, Node * polytomy = 0);
            void                                    queueTMatrixRecalculation(Node * nd);
            void                                    defineOperations(Tree::SharedPtr t);
            void                                    updateTransitionMatrices();
            void                                    calculatePartials();
            double                                  calcInstanceLogLikelihood(InstanceInfo & inst, Tree::SharedPtr t);

            std::vector<InstanceInfo>               _instances;
            std::map<int, std::string>              _beagle_error;
            std::map<int, std::vector<int> >        _operations;
            std::map<int, std::vector<int> >        _pmatrix_index;
            std::map<int, std::vector<double> >     _edge_lengths;
            std::map<int, std::vector<int> >        _eigen_indices;
            std::map<int, std::vector<int> >        _category_rate_indices;
            double                                  _relrate_normalizing_constant;

            std::vector<int>                        _subset_indices;
            std::vector<int>                        _parent_indices;
            std::vector<int>                        _child_indices;
            std::vector<int>                        _tmatrix_indices;
            std::vector<int>                        _weights_indices;
            std::vector<int>                        _freqs_indices;
            std::vector<int>                        _scaling_indices;

            Model::SharedPtr                        _model;

            Data::SharedPtr                         _data;
            unsigned                                _ntaxa;
            bool                                    _rooted;
            bool                                    _prefer_gpu;
            bool                                    _ambiguity_equals_missing;
            bool                                    _underflow_scaling;
            bool                                    _using_data;

            std::vector<Node *>                     _polytomy_helpers;
            std::map<int, std::vector<int> >        _polytomy_map;
            std::vector<double>                     _identity_matrix;

        public:
            typedef std::shared_ptr< Likelihood >   SharedPtr;
    };

    inline Likelihood::Likelihood() {
        clear();
    }

    inline Likelihood::~Likelihood() {
        finalizeBeagleLib(false);
        clear();
    }

    inline unsigned Likelihood::calcNumEdgesInFullyResolvedTree() const {
        assert(_ntaxa > 0);
        return (_rooted ? (2*_ntaxa - 2) : (2*_ntaxa - 3));
    }

    inline unsigned Likelihood::calcNumInternalsInFullyResolvedTree() const {
        assert(_ntaxa > 0);
        return (_rooted ? (_ntaxa - 1) : (_ntaxa - 2));
    }

    inline void Likelihood::finalizeBeagleLib(bool use_exceptions) {
        // Close down all BeagleLib instances if active
        for (auto info : _instances) {
            if (info.handle >= 0) {
                int code = beagleFinalizeInstance(info.handle);
                if (code != 0) {
                    if (use_exceptions)
                        throw XLorad(boost::format("Likelihood failed to finalize BeagleLib instance. BeagleLib error code was %d (%s).") % code % _beagle_error[code]);
                    else
                        std::cerr << boost::format("Likelihood destructor failed to finalize BeagleLib instance. BeagleLib error code was %d (%s).") % code % _beagle_error[code] << std::endl;
                }
            }
        }
        _instances.clear();
    }

    inline void Likelihood::clear() {
        finalizeBeagleLib(true);

        _ntaxa                      = 0;
        _rooted                     = false;
        _prefer_gpu                 = false;
        _ambiguity_equals_missing   = true;
        _underflow_scaling          = false;
        _using_data                 = true;
        _data                       = nullptr;

        _operations.clear();
        _pmatrix_index.clear();
        _edge_lengths.clear();
        _eigen_indices.clear();
        _category_rate_indices.clear();
        _relrate_normalizing_constant = 1.0;
        _subset_indices.assign(1, 0);
        _parent_indices.assign(1, 0);
        _child_indices.assign(1, 0);
        _tmatrix_indices.assign(1, 0);
        _weights_indices.assign(1, 0);
        _freqs_indices.assign(1, 0);
        _scaling_indices.assign(1, 0);
        _identity_matrix.assign(1, 0.0);

        _model = Model::SharedPtr(new Model());

        // Store BeagleLib error codes so that useful
        // error messages may be provided to the user
        _beagle_error.clear();
        _beagle_error[0]  = std::string("success");
        _beagle_error[-1] = std::string("unspecified error");
        _beagle_error[-2] = std::string("not enough memory could be allocated");
        _beagle_error[-3] = std::string("unspecified exception");
        _beagle_error[-4] = std::string("the instance index is out of range, or the instance has not been created");
        _beagle_error[-5] = std::string("one of the indices specified exceeded the range of the array");
        _beagle_error[-6] = std::string("no resource matches requirements");
        _beagle_error[-7] = std::string("no implementation matches requirements");
        _beagle_error[-8] = std::string("floating-point range exceeded");
    }

    inline std::string Likelihood::beagleLibVersion() const {
        return std::string(beagleGetVersion());
    }

    inline std::string Likelihood::availableResources() const {
        BeagleResourceList * rsrcList = beagleGetResourceList();
        std::string s;
        for (int i = 0; i < rsrcList->length; ++i) {
            std::string desc = rsrcList->list[i].description;
            boost::trim(desc);
            if (desc.size() > 0)
                s += boost::str(boost::format("  resource %d: %s (%s)\n") % i % rsrcList->list[i].name % desc);
            else
                s += boost::str(boost::format("  resource %d: %s\n") % i % rsrcList->list[i].name);
        }
        boost::trim_right(s);
        return s;
    }

    inline std::string Likelihood::usedResources() const {
        std::string s;
        for (unsigned i = 0; i < _instances.size(); i++) {
            s += boost::str(boost::format("  instance %d: %s (resource %d)\n") % _instances[i].handle % _instances[i].resourcename % _instances[i].resourcenumber);
        }
        return s;
    }

    inline Data::SharedPtr Likelihood::getData() {
        return _data;
    }

    inline void Likelihood::setData(Data::SharedPtr data) {
        assert(_instances.size() == 0);
        assert(!data->getDataMatrix().empty());
        _data = data;
    }

    inline Model::SharedPtr Likelihood::getModel() {
        return _model;
    }

    inline void Likelihood::setModel(Model::SharedPtr m) {
        assert(_instances.size() == 0); // can't change model after initBeagleLib called
        _model = m;
    }

    inline void Likelihood::setRooted(bool is_rooted) {
        assert(_instances.size() == 0 || _rooted == is_rooted); // can't change rooting status after initBeagleLib called
        _rooted = is_rooted;
    }

    inline void Likelihood::setAmbiguityEqualsMissing(bool ambig_equals_missing) {
        // Can't change GPU preference status after initBeagleLib called
        assert(_instances.size() == 0 || _ambiguity_equals_missing == ambig_equals_missing);
        _ambiguity_equals_missing = ambig_equals_missing;
    }

    inline void Likelihood::setPreferGPU(bool prefer_gpu) {
        // Can't change GPU preference status after initBeagleLib called
        assert(_instances.size() == 0 || _prefer_gpu == prefer_gpu);
        _prefer_gpu = prefer_gpu;
    }

    inline bool Likelihood::usingStoredData() const {
        return _using_data;
    }

    inline void Likelihood::useStoredData(bool using_data) {
        _using_data = using_data;
    }

    inline void Likelihood::useUnderflowScaling(bool do_scaling) {
        _underflow_scaling = do_scaling;
    }

    inline void Likelihood::initBeagleLib() {
        assert(_data);
        assert(_model);

        // Close down any existing BeagleLib instances
        finalizeBeagleLib(true);

        _ntaxa = _data->getNumTaxa();

        unsigned nsubsets = _data->getNumSubsets();
        std::set<instance_pair_t> nstates_ncateg_combinations;
        std::map<instance_pair_t, std::vector<unsigned> > subsets_for_pair;
        for (unsigned subset = 0; subset < nsubsets; subset++) {
            // Create a pair comprising number of states and number of rate categories
            unsigned nstates = _data->getNumStatesForSubset(subset);
            bool invar_model = _model->getSubsetIsInvarModel(subset);
            int nrates = (invar_model ? -1 : 1)*_model->getSubsetNumCateg(subset);
            instance_pair_t p = std::make_pair(nstates, nrates);

            // Add combo to set
            nstates_ncateg_combinations.insert(p);
            subsets_for_pair[p].push_back(subset);
        }

        // Create one instance for each distinct nstates-nrates combination
        _instances.clear();
        for (auto p : nstates_ncateg_combinations) {
            newInstance(p.first, p.second, subsets_for_pair[p]);

            InstanceInfo & info = *_instances.rbegin();
            ::om.outputConsole(boost::format("Created BeagleLib instance %d (%d states, %d rate%s, %d subset%s, %s invar. sites model)\n") % info.handle % info.nstates % info.nratecateg % (info.nratecateg == 1 ? "" : "s") % info.subsets.size() % (info.subsets.size() == 1 ? "" : "s") % (info.invarmodel ? "is" : "not"));
        }

        if (_ambiguity_equals_missing)
            setTipStates();
        else
            setTipPartials();
        setPatternWeights();
        setPatternPartitionAssignments();
    }

    inline void Likelihood::newInstance(unsigned nstates, int nrates, std::vector<unsigned> & subset_indices) {
        unsigned num_subsets = (unsigned)subset_indices.size();

        bool is_invar_model = (nrates < 0 ? true : false);
        unsigned ngammacat = (unsigned)(is_invar_model ? -nrates : nrates);

        // Create an identity matrix used for computing partials
        // for polytomies (represents the transition matrix
        // for the zero-length edges inserted to arbitrarily
        // resolve each polytomy)
        _identity_matrix.assign(nstates*nstates*ngammacat, 0.0);
        for (unsigned k = 0; k < ngammacat; k++) {
            unsigned offset = k*nstates*nstates;
            _identity_matrix[0+offset]  = 1.0;
            _identity_matrix[5+offset]  = 1.0;
            _identity_matrix[10+offset] = 1.0;
            _identity_matrix[15+offset] = 1.0;
        }

        //...

        unsigned num_patterns = 0;
        for (auto s : subset_indices) {
            num_patterns += _data->getNumPatternsInSubset(s);
        }

        unsigned num_internals = calcNumInternalsInFullyResolvedTree();

        // add 1 to num_edges so that subroot node will have a tmatrix, root tip's tmatrix is never used
        unsigned num_edges = calcNumEdgesInFullyResolvedTree();
        unsigned num_nodes = num_edges + 1;
        unsigned num_transition_probs = num_nodes*num_subsets;

        long requirementFlags = 0;

        long preferenceFlags = BEAGLE_FLAG_PRECISION_SINGLE | BEAGLE_FLAG_THREADING_CPP;
        if (_underflow_scaling) {
            preferenceFlags |= BEAGLE_FLAG_SCALING_MANUAL;
            preferenceFlags |= BEAGLE_FLAG_SCALERS_LOG;
        }
        if (_prefer_gpu)
            preferenceFlags |= BEAGLE_FLAG_PROCESSOR_GPU;
        else
            preferenceFlags |= BEAGLE_FLAG_PROCESSOR_CPU;

        BeagleInstanceDetails instance_details;
        unsigned npartials = num_internals + _ntaxa;
        unsigned nscalers = num_internals;  // one scale buffer for every internal node
        unsigned nsequences = 0;
        if (_ambiguity_equals_missing) {
            npartials -= _ntaxa;
            nsequences += _ntaxa;
        }

        int inst = beagleCreateInstance(
             _ntaxa,                        // tips
             2*npartials,                   // partials
             nsequences,                    // sequences
             nstates,                       // states
             num_patterns,                  // patterns (total across all subsets that use this instance)
             num_subsets,                   // models (one for each distinct eigen decomposition)
             2*num_subsets*num_transition_probs, // transition matrices (one for each edge in each subset)
             ngammacat,                     // rate categories
             (_underflow_scaling ? 2*nscalers + 1 : 0),  // scale buffers (+1 is for the cumulative scaler at index 0)
             NULL,                          // resource restrictions
             0,                             // length of resource list
             preferenceFlags,               // preferred flags
             requirementFlags,              // required flags
             &instance_details);            // pointer for details

        if (inst < 0) {
            // beagleCreateInstance returns one of the following:
            //   valid instance (0, 1, 2, ...)
            //   error code (negative integer)
            throw XLorad(boost::str(boost::format("Likelihood init function failed to create BeagleLib instance (BeagleLib error code was %d)") % _beagle_error[inst]));
        }

        InstanceInfo info;
        info.handle         = inst;
        info.resourcenumber = instance_details.resourceNumber;
        info.resourcename   = instance_details.resourceName;
        info.nstates        = nstates;
        info.nratecateg     = ngammacat;
        info.invarmodel     = is_invar_model;
        info.subsets        = subset_indices;
        info.npatterns      = num_patterns;
        info.partial_offset = num_internals;
        info.tmatrix_offset = num_nodes;
        _instances.push_back(info);
    }

    inline void Likelihood::setTipStates() {
        assert(_instances.size() > 0);
        assert(_data);
        Data::state_t one = 1;

        for (auto & info : _instances) {
            std::vector<int> states(info.nstates*info.npatterns);

            // Loop through all rows of the data matrix, setting the tip states for one taxon each row
            unsigned t = 0;
            for (auto & row : _data->getDataMatrix()) {

                // Loop through all subsets assigned to this instance
                unsigned k = 0;
                for (unsigned s : info.subsets) {

                    // Loop through all patterns in this subset
                    auto interval = _data->getSubsetBeginEnd(s);
                    for (unsigned p = interval.first; p < interval.second; p++) {

                        // d is the state for taxon t, pattern p (in subset s)
                        // d is stored as a bit field (e.g., for nucleotide data, A=1, C=2, G=4, T=8, ?=15),
                        // but BeagleLib expects states to be integers (e.g. for nucleotide data,
                        // A=0, C=1, G=2, T=3, ?=4).
                        Data::state_t d = row[p];

                        // Handle common nucleotide case separately
                        if (info.nstates == 4) {
                            if (d == 1)
                                states[k++] = 0;
                            else if (d == 2)
                                states[k++] = 1;
                            else if (d == 4)
                                states[k++] = 2;
                            else if (d == 8)
                                states[k++] = 3;
                            else
                                states[k++] = 4;
                        }
                        else {
                            // This case is for any other data type except nucleotide
                            int s = -1;
                            for (unsigned b = 0; b < info.nstates; b++) {
                                if (d == one << b) {
                                    s = b;
                                    break;
                                }
                            }
                            if (s == -1)
                                states[k++] = info.nstates;
                            else
                                states[k++] = s;
                        }
                    } // pattern loop
                }   // subset loop

            int code = beagleSetTipStates(
                info.handle,    // Instance number
                t,              // Index of destination compactBuffer
                &states[0]);    //  Pointer to compact states vector

            if (code != 0)
                throw XLorad(boost::format("failed to set tip state for taxon %d (\"%s\"; BeagleLib error code was %d)") % (t+1) % _data->getTaxonNames()[t] % code % _beagle_error[code]);
            ++t;
            }
        }
    }

    inline void Likelihood::setTipPartials() {
        assert(_instances.size() > 0);
        assert(_data);
        Data::state_t one = 1;

        for (auto & info : _instances) {
            std::vector<double> partials(info.nstates*info.npatterns);

            // Loop through all rows of data matrix, setting the tip states for one taxon each row
            unsigned t = 0;
            for (auto & row : _data->getDataMatrix()) {

                // Loop through all subsets assigned to this instance
                unsigned k = 0;
                for (unsigned s : info.subsets) {

                    // Loop through all patterns in this subset
                    auto interval = _data->getSubsetBeginEnd(s);
                    for (unsigned p = interval.first; p < interval.second; p++) {

                        // d is the state for taxon t, pattern p (in subset s)
                        Data::state_t d = row[p];

                        // Handle common nucleotide case separately
                        if (info.nstates == 4) {
                            partials[k++] = d & 1 ? 1.0 : 0.0;
                            partials[k++] = d & 2 ? 1.0 : 0.0;
                            partials[k++] = d & 4 ? 1.0 : 0.0;
                            partials[k++] = d & 8 ? 1.0 : 0.0;
                        }
                        else {
                            // This case is for any other data type except nucleotide
                            for (unsigned b = 0; b < info.nstates; b++) {
                                partials[k++] = d & (one << b) ? 1.0 : 0.0;
                            }
                        }
                    }
                }

            int code = beagleSetTipPartials(
                info.handle,    // Instance number
                t,              // Index of destination compactBuffer
                &partials[0]);  // Pointer to compact states vector

            if (code != 0)
                throw XLorad(boost::format("failed to set tip state for taxon %d (\"%s\"; BeagleLib error code was %d)") % (t+1) % _data->getTaxonNames()[t] % code % _beagle_error[code]);
            ++t;
            }
        }
    }

    inline void Likelihood::setPatternPartitionAssignments() {
        assert(_instances.size() > 0);
        assert(_data);

        // beagleSetPatternPartitions does not need to be called if data are unpartitioned
        // (and, in fact, BeagleLib only supports partitioning for 4-state instances if GPU is used,
        // so not calling beagleSetPatternPartitions allows unpartitioned codon model analyses)
        if (_instances.size() == 1 && _instances[0].subsets.size() == 1)
            return;

        Data::partition_key_t v;

        // Loop through all instances
        for (auto & info : _instances) {
            unsigned nsubsets = (unsigned)info.subsets.size();
            v.resize(info.npatterns);
            unsigned pattern_index = 0;

            // Loop through all subsets assigned to this instance
            unsigned instance_specific_subset_index = 0;
            for (unsigned s : info.subsets) {
                // Loop through all patterns in this subset
                auto interval = _data->getSubsetBeginEnd(s);
                for (unsigned p = interval.first; p < interval.second; p++) {
                    v[pattern_index++] = instance_specific_subset_index;
                }
                ++instance_specific_subset_index;
            }

            int code = beagleSetPatternPartitions(
               info.handle, // instance number
               nsubsets,    // number of data subsets (equals 1 if data are unpartitioned)
               &v[0]);      // vector of subset indices: v[i] = 0 means pattern i is in subset 0

            if (code != 0) {
                throw XLorad(boost::format("failed to set pattern partition. BeagleLib error code was %d (%s)") % code % _beagle_error[code]);
            }
        }
    }

    inline void Likelihood::setPatternWeights() {
        assert(_instances.size() > 0);
        assert(_data);
        Data::pattern_counts_t v;
        auto pattern_counts = _data->getPatternCounts();
        assert(pattern_counts.size() > 0);

        // Loop through all instances
        for (auto & info : _instances) {
            v.resize(info.npatterns);
            unsigned pattern_index = 0;

            // Loop through all subsets assigned to this instance
            for (unsigned s : info.subsets) {

                // Loop through all patterns in this subset
                auto interval = _data->getSubsetBeginEnd(s);
                for (unsigned p = interval.first; p < interval.second; p++) {
                    v[pattern_index++] = pattern_counts[p];
                }
            }

            int code = beagleSetPatternWeights(
               info.handle,   // instance number
               &v[0]);        // vector of pattern counts: v[i] = 123 means pattern i was encountered 123 times

            if (code != 0)
                throw XLorad(boost::format("Failed to set pattern weights for instance %d. BeagleLib error code was %d (%s)") % info.handle % code % _beagle_error[code]);
        }
    }

    inline void Likelihood::setAmongSiteRateHeterogenetity() {
        assert(_instances.size() > 0);
        int code = 0;

        // Loop through all instances
        for (auto & info : _instances) {

            // Loop through all subsets assigned to this instance
            unsigned instance_specific_subset_index = 0;
            for (unsigned s : info.subsets) {
                code = _model->setBeagleAmongSiteRateVariationRates(info.handle, s, instance_specific_subset_index);
                if (code != 0)
                    throw XLorad(boost::str(boost::format("Failed to set category rates for BeagleLib instance %d. BeagleLib error code was %d (%s)") % info.handle % code % _beagle_error[code]));

                code = _model->setBeagleAmongSiteRateVariationProbs(info.handle, s, instance_specific_subset_index);
                if (code != 0)
                    throw XLorad(boost::str(boost::format("Failed to set category probabilities for BeagleLib instance %d. BeagleLib error code was %d (%s)") % info.handle % code % _beagle_error[code]));

                ++instance_specific_subset_index;
            }
        }
    }

    inline void Likelihood::setModelRateMatrix() {
        // Loop through all instances
        for (auto & info : _instances) {

            // Loop through all subsets assigned to this instance
            unsigned instance_specific_subset_index = 0;
            for (unsigned s : info.subsets) {
                int code = _model->setBeagleStateFrequencies(info.handle, s, instance_specific_subset_index);
                if (code != 0)
                    throw XLorad(boost::str(boost::format("Failed to set state frequencies for BeagleLib instance %d. BeagleLib error code was %d (%s)") % info.handle % code % _beagle_error[code]));

                code = _model->setBeagleEigenDecomposition(info.handle, s, instance_specific_subset_index);
                if (code != 0)
                    throw XLorad(boost::str(boost::format("Failed to set eigen decomposition for BeagleLib instance %d. BeagleLib error code was %d (%s)") % info.handle % code % _beagle_error[code]));

                ++instance_specific_subset_index;
            }
        }
    }

    inline unsigned Likelihood::getScalerIndex(Node * nd, InstanceInfo & info) const {
        unsigned sindex = BEAGLE_OP_NONE;
        if (_underflow_scaling) {
            sindex = nd->_number - _ntaxa + 1; // +1 to skip the cumulative scaler vector
            if (nd->isAltPartial())
                sindex += info.partial_offset;
        }
        return sindex;
    }

    inline unsigned Likelihood::getPartialIndex(Node * nd, InstanceInfo & info) const {
        // Note: do not be tempted to subtract _ntaxa from pindex: BeagleLib does this itself
        assert(nd->_number >= 0);
        unsigned pindex = nd->_number;
        if (pindex >= _ntaxa) {
            if (nd->isAltPartial())
                pindex += info.partial_offset;
        }
        return pindex;
    }

    inline unsigned Likelihood::getTMatrixIndex(Node * nd, InstanceInfo & info, unsigned subset_index) const {
        unsigned tindex = 2*subset_index*info.tmatrix_offset + nd->_number;
        if (nd->isAltTMatrix())
            tindex += info.tmatrix_offset;
        return tindex;
    }

    inline void Likelihood::defineOperations(Tree::SharedPtr t) {
        assert(_instances.size() > 0);
        assert(t);
        assert(t->isRooted() == _rooted);
        assert(_polytomy_helpers.empty());
        assert(_polytomy_map.empty());

        _relrate_normalizing_constant = _model->calcNormalizingConstantForSubsetRelRates();

        // Start with a clean slate
        for (auto & info : _instances) {
            _operations[info.handle].clear();
            _pmatrix_index[info.handle].clear();
            _edge_lengths[info.handle].clear();
            _eigen_indices[info.handle].clear();
            _category_rate_indices[info.handle].clear();
        }

        // Loop through all nodes in reverse level order
        for (auto nd : boost::adaptors::reverse(t->_levelorder)) {
            assert(nd->_number >= 0);
            if (!nd->_left_child) {
                // This is a leaf
                if (nd->isSelTMatrix())
                    queueTMatrixRecalculation(nd);
            }
            else {
                // This is an internal node
                if (nd->isSelTMatrix())
                    queueTMatrixRecalculation(nd);

                // Internal nodes have partials to be calculated, so define
                // an operation to compute the partials for this node
                if (nd->isSelPartial()) {
                    TreeManip tm(t);
                    if (tm.isPolytomy(nd)) {
                        // Internal node is a polytomy
                        unsigned nchildren = tm.countChildren(nd);
                        Node * a = nd->_left_child;
                        Node * b = a->_right_sib;
                        Node * c = 0;
                        for (unsigned k = 0; k < nchildren - 2; k++) {
                            c = tm.getUnusedNode();
                            c->_left_child = a;
                            _polytomy_helpers.push_back(c);

                            queuePartialsRecalculation(c, a, b, nd);

                            // Tackle next arm of the polytomy
                            b = b->_right_sib;
                            a = c;
                        }

                        // Now add operation to compute the partial for the real internal node
                        queuePartialsRecalculation(nd, a, b);
                    }
                    else {
                        // Internal node is not a polytomy
                        Node * lchild = nd->_left_child;
                        assert(lchild);
                        Node * rchild = lchild->_right_sib;
                        assert(rchild);
                        queuePartialsRecalculation(nd, lchild, rchild);
                    }
                }
            }
        }
    }

    inline void Likelihood::queuePartialsRecalculation(Node * nd, Node * lchild, Node * rchild, Node * polytomy) {
        for (auto & info : _instances) {
            unsigned instance_specific_subset_index = 0;
            // line below produces warning about s not being used, but this is unavoidable
            // because values in vector info.subsets may not be sequential
            for (unsigned s : info.subsets) {
                (void)s;    // to prevent compiler warning about s not being used

                if (polytomy) {
                    // nd has been pulled out of tree's _unused_nodes vector to break up the polytomy
                    // Note that the parameter "polytomy" is the polytomous node itself

                    // First get the transition matrix index
                    unsigned tindex = getTMatrixIndex(nd, info, instance_specific_subset_index);

                    // Set the transition matrix for nd to the identity matrix
                    // note: last argument 1 is the value used for ambiguous states (should be 1 for transition matrices)
                    int code = beagleSetTransitionMatrix(info.handle, tindex, &_identity_matrix[0], 1);
                    if (code != 0)
                        throw XLorad(boost::str(boost::format("Failed to set transition matrix for instance %d. BeagleLib error code was %d (%s)") % info.handle % code % _beagle_error[code]));

                    // Set the edgelength to 0.0 to maintain consistency with the transition matrix
                    nd->setEdgeLength(0.0);

                    // If employing underflow scaling, the scaling factors for these fake nodes need to be
                    // transferred to the polytomous node, as that will be the only node remaining after the
                    // likelihood has been calculated. Save the scaling factor index, associating it with
                    // the scaling factor index of the polytomy node.
                    if (_underflow_scaling) {
                        // Get the polytomy's scaling factor index
                        int spolytomy = getScalerIndex(polytomy, info);

                        // Get nd's scaling factor index
                        int snd = getScalerIndex(nd, info);

                        // Save nd's index in the vector associated with polytomy's index
                        _polytomy_map[spolytomy].push_back(snd);
                    }

                }

                addOperation(info, nd, lchild, rchild, instance_specific_subset_index);
                ++instance_specific_subset_index;
            }
        }
    }

    inline void Likelihood::queueTMatrixRecalculation(Node * nd) {
        Model::subset_relrate_vect_t & subset_relrates = _model->getSubsetRelRates();

        for (auto & info : _instances) {
            unsigned instance_specific_subset_index = 0;
            for (unsigned s : info.subsets) {
                double subset_relative_rate = subset_relrates[s]/_relrate_normalizing_constant;

                unsigned tindex = getTMatrixIndex(nd, info, instance_specific_subset_index);
                _pmatrix_index[info.handle].push_back(tindex);
                _edge_lengths[info.handle].push_back(nd->_edge_length*subset_relative_rate);
                _eigen_indices[info.handle].push_back(s);
                _category_rate_indices[info.handle].push_back(s);

                ++instance_specific_subset_index;
            }
        }
    }

    inline void Likelihood::addOperation(InstanceInfo & info, Node * nd, Node * lchild, Node * rchild, unsigned subset_index) {
        assert(nd);
        assert(lchild);
        assert(rchild);

        // 1. destination partial to be calculated
        int partial_dest = getPartialIndex(nd, info);
        _operations[info.handle].push_back(partial_dest);

        // 2. destination scaling buffer index to write to
        int scaler_index = getScalerIndex(nd, info);
        _operations[info.handle].push_back(scaler_index);

        // 3. destination scaling buffer index to read from
        _operations[info.handle].push_back(BEAGLE_OP_NONE);

        // 4. left child partial index
        int partial_lchild = getPartialIndex(lchild, info);
        _operations[info.handle].push_back(partial_lchild);

        // 5. left child transition matrix index
        unsigned tindex_lchild = getTMatrixIndex(lchild, info, subset_index);
        _operations[info.handle].push_back(tindex_lchild);

        // 6. right child partial index
        int partial_rchild = getPartialIndex(rchild, info);
        _operations[info.handle].push_back(partial_rchild);

        // 7. right child transition matrix index
        unsigned tindex_rchild = getTMatrixIndex(rchild, info, subset_index);
        _operations[info.handle].push_back(tindex_rchild);

        if (info.subsets.size() > 1) {
            // 8. index of partition subset
            _operations[info.handle].push_back(subset_index);

            // 9. cumulative scale index
            _operations[info.handle].push_back(BEAGLE_OP_NONE); // accumulate in calcInstanceLogLikelihood
        }
    }

    inline void Likelihood::updateTransitionMatrices() {
        assert(_instances.size() > 0);
        if (_pmatrix_index.size() == 0)
            return;

        // Loop through all instances
        for (auto & info : _instances) {
            int code = 0;

            unsigned nsubsets = (unsigned)info.subsets.size();
            if (nsubsets > 1) {
                code = beagleUpdateTransitionMatricesWithMultipleModels(
                    info.handle,                                // Instance number
                    &_eigen_indices[info.handle][0],            // Index of eigen-decomposition buffer
                    &_category_rate_indices[info.handle][0],    // category rate indices
                    &_pmatrix_index[info.handle][0],            // transition probability matrices to update
                    NULL,                                       // first derivative matrices to update
                    NULL,                                       // second derivative matrices to update
                    &_edge_lengths[info.handle][0],             // List of edge lengths
                    (int)_pmatrix_index[info.handle].size());   // Length of lists
            }
            else {
                code = beagleUpdateTransitionMatrices(
                    info.handle,                                // Instance number
                    0,                                          // Index of eigen-decomposition buffer
                    &_pmatrix_index[info.handle][0],            // transition probability matrices to update
                    NULL,                                       // first derivative matrices to update
                    NULL,                                       // second derivative matrices to update
                    &_edge_lengths[info.handle][0],             // List of edge lengths
                    (int)_pmatrix_index[info.handle].size());   // Length of lists
            }

            if (code != 0)
                throw XLorad(boost::str(boost::format("Failed to update transition matrices for instance %d. BeagleLib error code was %d (%s)") % info.handle % code % _beagle_error[code]));

        }
    }

    inline void Likelihood::calculatePartials() {
        assert(_instances.size() > 0);
        if (_operations.size() == 0)
            return;
        int code = 0;

        // Loop through all instances
        for (auto & info : _instances) {
            unsigned nsubsets = (unsigned)info.subsets.size();

            if (nsubsets > 1) {
                code = beagleUpdatePartialsByPartition(
                    info.handle,                                                    // Instance number
                    (BeagleOperationByPartition *) &_operations[info.handle][0],    // BeagleOperation list specifying operations
                    (int)(_operations[info.handle].size()/9));                      // Number of operations
                if (code != 0)
                    throw XLorad(boost::format("failed to update partials. BeagleLib error code was %d (%s)") % code % _beagle_error[code]);

                if (_underflow_scaling) {
                    // Accumulate scaling factors across polytomy helpers and assign them to their parent node
                    for (auto & m : _polytomy_map) {
                        for (unsigned subset = 0; subset < nsubsets; subset++) {
                            code = beagleAccumulateScaleFactorsByPartition(info.handle, &m.second[0], (int)m.second.size(), m.first, subset);
                            if (code != 0) {
                                throw XLorad(boost::format("failed to transfer scaling factors to polytomous node. BeagleLib error code was %d (%s)") % code % _beagle_error[code]);
                            }
                        }
                    }
                }
            }
            else {
                // no partitioning, just one data subset
                code = beagleUpdatePartials(
                    info.handle,                                        // Instance number
                    (BeagleOperation *) &_operations[info.handle][0],   // BeagleOperation list specifying operations
                    (int)(_operations[info.handle].size()/7),           // Number of operations
                    BEAGLE_OP_NONE);                                    // Index number of scaleBuffer to store accumulated factors
                if (code != 0)
                    throw XLorad(boost::format("failed to update partials. BeagleLib error code was %d (%s)") % code % _beagle_error[code]);

                if (_underflow_scaling) {
                    // Accumulate scaling factors across polytomy helpers and assign them to their parent node
                    for (auto & m : _polytomy_map) {
                        code = beagleAccumulateScaleFactors(info.handle, &m.second[0], (int)m.second.size(), m.first);
                        if (code != 0) {
                            throw XLorad(boost::format("failed to transfer scaling factors to polytomous node. BeagleLib error code was %d (%s)") % code % _beagle_error[code]);
                        }
                    }
                }
            }
        }
    }

    inline double Likelihood::calcInstanceLogLikelihood(InstanceInfo & info, Tree::SharedPtr t) {
        int code = 0;
        unsigned nsubsets = (unsigned)info.subsets.size();
        assert(nsubsets > 0);

        // Assuming there are as many transition matrices as there are edge lengths
        assert(_pmatrix_index[info.handle].size() == _edge_lengths[info.handle].size());

        int state_frequency_index  = 0;
        int category_weights_index = 0;
        int cumulative_scale_index = (_underflow_scaling ? 0 : BEAGLE_OP_NONE);
        int child_partials_index   = getPartialIndex(t->_root, info);
        int parent_partials_index = getPartialIndex(t->_preorder[0], info);
        int parent_tmatrix_index = getTMatrixIndex(t->_preorder[0], info, 0);

        // storage for results of the likelihood calculation
        std::vector<double> subset_log_likelihoods(nsubsets, 0.0);
        double log_likelihood = 0.0;

        if (_underflow_scaling) {
            // Create vector of all scaling vector indices in current use
            std::vector<int> internal_node_scaler_indices;
            for (auto nd : t->_preorder) {
                if (nd->_left_child) {
                    unsigned s = getScalerIndex(nd, info);
                    internal_node_scaler_indices.push_back(s);
                }
            }

            if (nsubsets == 1) {
                code = beagleResetScaleFactors(info.handle, cumulative_scale_index);
                if (code != 0)
                    throw XLorad(boost::str(boost::format("failed to reset scale factors in calcInstanceLogLikelihood. BeagleLib error code was %d (%s)") % code % _beagle_error[code]));

                code = beagleAccumulateScaleFactors(
                     info.handle,
                     &internal_node_scaler_indices[0],
                     (int)internal_node_scaler_indices.size(),
                     cumulative_scale_index);
                if (code != 0)
                    throw XLorad(boost::str(boost::format("failed to accumulate scale factors in calcInstanceLogLikelihood. BeagleLib error code was %d (%s)") % code % _beagle_error[code]));
            }
            else {
                for (unsigned s = 0; s < nsubsets; ++s) {
                    code = beagleResetScaleFactorsByPartition(info.handle, cumulative_scale_index, s);
                    if (code != 0)
                        throw XLorad(boost::str(boost::format("failed to reset scale factors for subset %d in calcInstanceLogLikelihood. BeagleLib error code was %d (%s)") % s % code % _beagle_error[code]));

                    code = beagleAccumulateScaleFactorsByPartition(
                        info.handle,
                        &internal_node_scaler_indices[0],
                        (int)internal_node_scaler_indices.size(),
                        cumulative_scale_index,
                        s);
                    if (code != 0)
                        throw XLorad(boost::str(boost::format("failed to acccumulate scale factors for subset %d in calcInstanceLogLikelihood. BeagleLib error code was %d (%s)") % s % code % _beagle_error[code]));
                }
            }
        }

        if (nsubsets > 1) {
            _parent_indices.assign(nsubsets, parent_partials_index);
            _child_indices.assign(nsubsets, child_partials_index);
            _weights_indices.assign(nsubsets, category_weights_index);
            _scaling_indices.resize(nsubsets);
            _subset_indices.resize(nsubsets);
            _freqs_indices.resize(nsubsets);
            _tmatrix_indices.resize(nsubsets);

            for (unsigned s = 0; s < nsubsets; s++) {
                _scaling_indices[s]  = (_underflow_scaling ? 0 : BEAGLE_OP_NONE);
                _subset_indices[s]  = s;
                _freqs_indices[s]   = s;
                _tmatrix_indices[s] = getTMatrixIndex(t->_preorder[0], info, s); //index_focal_child + s*tmatrix_skip;
            }

            code = beagleCalculateEdgeLogLikelihoodsByPartition(
                info.handle,                 // instance number
                &_parent_indices[0],         // indices of parent partialsBuffers
                &_child_indices[0],          // indices of child partialsBuffers
                &_tmatrix_indices[0],        // transition probability matrices for this edge
                NULL,                        // first derivative matrices
                NULL,                        // second derivative matrices
                &_weights_indices[0],        // weights to apply to each partialsBuffer
                &_freqs_indices[0],          // state frequencies for each partialsBuffer
                &_scaling_indices[0],        // scaleBuffers containing accumulated factors
                &_subset_indices[0],         // indices of subsets
                nsubsets,                    // partition subset count
                1,                           // number of distinct eigen decompositions
                &subset_log_likelihoods[0],  // address of vector of log likelihoods (one for each subset)
                &log_likelihood,             // destination for resulting log likelihood
                NULL,                        // destination for vector of first derivatives (one for each subset)
                NULL,                        // destination for first derivative
                NULL,                        // destination for vector of second derivatives (one for each subset)
                NULL);                       // destination for second derivative
        }
        else {
            code = beagleCalculateEdgeLogLikelihoods(
                info.handle,                 // instance number
                &parent_partials_index,      // indices of parent partialsBuffers
                &child_partials_index,       // indices of child partialsBuffers
                &parent_tmatrix_index,       // transition probability matrices for this edge
                NULL,                        // first derivative matrices
                NULL,                        // second derivative matrices
                &category_weights_index,     // weights to apply to each partialsBuffer
                &state_frequency_index,      // state frequencies for each partialsBuffer
                &cumulative_scale_index,     // scaleBuffers containing accumulated factors
                1,                           // Number of partialsBuffer
                &log_likelihood,             // destination for log likelihood
                NULL,                        // destination for first derivative
                NULL);                       // destination for second derivative
        }

        // ...

        if (code != 0) {
            std::cerr << "Problem computing likelihood for this tree:\n";
            std::cerr << TreeManip(t).makeNewick(9, true) << std::endl;
            throw XLorad(boost::str(boost::format("failed to calculate edge log-likelihoods in calcInstanceLogLikelihood. BeagleLib error code was %d (%s)") % code % _beagle_error[code]));
        }

        if (info.invarmodel) {
            auto monomorphic = _data->getMonomorphic();
            auto counts = _data->getPatternCounts();
            std::vector<double> site_log_likelihoods(info.npatterns, 0.0);
            double * siteLogLs = &site_log_likelihoods[0];

            beagleGetSiteLogLikelihoods(info.handle, siteLogLs);

            // Loop through all subsets assigned to this instance
            double lnL = 0.0;
            unsigned i = 0;
            for (unsigned s : info.subsets) {
                const ASRV & asrv = _model->getASRV(s);
                const QMatrix & qmatrix = _model->getQMatrix(s);
                const double * freq = qmatrix.getStateFreqs();

                double pinvar = *(asrv.getPinvarSharedPtr());
                assert(pinvar >= 0.0 && pinvar <= 1.0);

                if (pinvar == 0.0) {
                    // log likelihood for this subset is equal to the sum of site log-likelihoods
                    auto interval = _data->getSubsetBeginEnd(s);
                    for (unsigned p = interval.first; p < interval.second; p++) {
                        lnL += counts[p]*site_log_likelihoods[i++];
                    }
                }
                else {
                    // Loop through all patterns in this subset
                    double log_pinvar = log(pinvar);
                    double log_one_minus_pinvar = log(1.0 - pinvar);
                    auto interval = _data->getSubsetBeginEnd(s);
                    for (unsigned p = interval.first; p < interval.second; p++) {
                        // Loop through all states for this pattern
                        double invar_like = 0.0;
                        if (monomorphic[p] > 0) {
                            for (unsigned k = 0; k < info.nstates; ++k) {
                                Data::state_t x = (Data::state_t)1 << k;
                                double condlike = (x & monomorphic[p] ? 1.0 : 0.0);
                                double basefreq = freq[k];
                                invar_like += condlike*basefreq;
                            }
                        }
                        double site_lnL = site_log_likelihoods[i++];
                        double log_like_term = log_one_minus_pinvar + site_lnL;
                        if (invar_like > 0.0) {
                            double log_invar_term = log_pinvar + log(invar_like);
                            double site_log_like = (log_like_term + log(1.0 + exp(log_invar_term - log_like_term)));
                            lnL += counts[p]*site_log_like;
                        }
                        else {
                            lnL += counts[p]*log_like_term;
                        }
                    }
                }
            }
            log_likelihood = lnL;
        }

        return log_likelihood;
    }

    inline double Likelihood::calcLogLikelihood(Tree::SharedPtr t) {
        assert(_instances.size() > 0);

        if (!_using_data)
            return 0.0;

        // Must call setData and setModel before calcLogLikelihood
        assert(_data);
        assert(_model);

        if (t->_is_rooted)
            throw XLorad("This version of the program can only compute likelihoods for unrooted trees");

        // Assuming "root" is leaf 0
        assert(t->_root->_number == 0 && t->_root->_left_child == t->_preorder[0] && !t->_preorder[0]->_right_sib);

        setModelRateMatrix();
        setAmongSiteRateHeterogenetity();
        defineOperations(t);
        updateTransitionMatrices();
        calculatePartials();

        double log_likelihood = 0.0;
        for (auto & info : _instances) {
            log_likelihood += calcInstanceLogLikelihood(info, t);
        }

        // We no longer need the internal nodes brought out of storage
        // and used to compute partials for polytomies
        TreeManip tm(t);
        for (Node * h : _polytomy_helpers) {
            tm.putUnusedNode(h);
        }
        _polytomy_helpers.clear();
        _polytomy_map.clear();

        return log_likelihood;
    }

}