/*
 * Copyright (c) 2014-2015, Hewlett-Packard Development Company, LP.
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation; either version 2 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details. You should have received a copy of the GNU General Public
 * License along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 *
 * HP designates this particular file as subject to the "Classpath" exception
 * as provided by HP in the LICENSE.txt file that accompanied this code.
 */
/**
 * @file foedus/graphlda/lda.cpp
 * @brief Latent-Dirichlet Allocation experiment for Topic Modeling
 * @author kimurhid
 * @date 2014/11/22
 * @details
 * This is a port of Fei Chen's LDA topic modeling program to FOEDUS.
 */
#include <stdint.h>
#include <boost/math/special_functions/gamma.hpp>
#include <gflags/gflags.h>
#include <glog/logging.h>

#include <algorithm>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

#include "foedus/engine.hpp"
#include "foedus/engine_options.hpp"
#include "foedus/error_stack.hpp"
#include "foedus/debugging/debugging_supports.hpp"
#include "foedus/debugging/stop_watch.hpp"
#include "foedus/fs/filesystem.hpp"
#include "foedus/memory/engine_memory.hpp"
#include "foedus/proc/proc_manager.hpp"
#include "foedus/soc/shared_memory_repo.hpp"
#include "foedus/soc/soc_manager.hpp"
#include "foedus/storage/storage_id.hpp"
#include "foedus/storage/storage_manager.hpp"
#include "foedus/storage/array/array_storage.hpp"
#include "foedus/thread/numa_thread_scope.hpp"
#include "foedus/thread/thread.hpp"
#include "foedus/thread/thread_pool.hpp"
#include "foedus/thread/thread_pool_pimpl.hpp"
#include "foedus/xct/xct_id.hpp"
#include "foedus/xct/xct_manager.hpp"

namespace foedus {
namespace graphlda {

DEFINE_string(dictionary_fname, "dictionary.txt", "Dictionary file");
DEFINE_string(counts_fname, "counts.tsv", "Counts file");
DEFINE_int32(ntopics, 50, "Number of topics");
DEFINE_int32(nburnin, 50, "Number of burnin iterations");
DEFINE_int32(nsamples, 10, "Number of sampling iterations");
DEFINE_int32(nthreads, 16, "Total number of threads");
DEFINE_double(alpha, 1.0, "Alpha prior");
DEFINE_double(beta, 0.1, "Beta prior");
DEFINE_bool(fork_workers, false, "Whether to fork(2) worker threads in child processes rather"
    " than threads in the same process. This is required to scale up to 100+ cores.");
DEFINE_bool(profile, false, "Whether to profile the execution with gperftools.");

typedef uint32_t WordId;
typedef uint32_t DocId;
typedef uint16_t TopicId;
typedef uint32_t Count;
const TopicId kNullTopicId = -1;

/** Used for overflow check */
const Count kUnreasonablyLargeUint32 = (0xFFFF0000U);
inline Count check_overflow(Count count) {
  if (count > kUnreasonablyLargeUint32) {
    return 0;
  } else {
    return count;
  }
}

// storage names
/**
 * n_td(t,d) [1d: d * ntopics + t] Number of occurences of topic t in document d.
 * As an array storage, it's d records of ntopics Count.
 */
const char* kNtd = "n_td";
/**
 * n_tw(t,w) [1d: w * ntopics + t] Number of occurences of word w in topic t.
 * As an array storage, it's w records of ntopics Count.
 */
const char* kNtw = "n_tw";
/**
 * n_t(t) [1d: t] The total number of words assigned to topic t.
 * As an array storage, it's 1 record of ntopics Count.
 */
const char* kNt = "n_t";

struct Token {
  WordId word;
  DocId doc;
  Token(WordId word = 0, DocId doc = 0) : word(word), doc(doc) { }
};
std::ostream& operator<<(std::ostream& out, const Token& tok) {
  return out << "(" << tok.word << ", " << tok.doc << ")";
}

/** This object is held only by the driver thread. Individual worker uses SharedInputs. */
struct Corpus {
  size_t nwords, ndocs, ntokens;
  std::vector< Token > tokens;
  std::vector<std::string> dictionary;

  Corpus(const fs::Path& dictionary_fname, const fs::Path& counts_fname)
    : nwords(0), ndocs(0), ntokens(0) {
    dictionary.reserve(20000);
    tokens.reserve(100000);
    load_dictionary(dictionary_fname);
    load_counts(counts_fname);
  }
  void load_dictionary(const fs::Path& fname) {
    std::ifstream fin(fname.c_str());
    std::string str;
    while (fin.good()) {
      std::getline(fin, str);
      if (fin.good()) {
        dictionary.push_back(str);
        nwords++;
      }
    }
    fin.close();
  }

  void load_counts(const fs::Path& fname)  {
    std::ifstream fin(fname.c_str());
    while (fin.good()) {
      // Read a collection of tokens
      const size_t NULL_VALUE(-1);
      size_t word = NULL_VALUE, doc = NULL_VALUE, count = NULL_VALUE;
      fin >> doc >> word >> count;
      if (fin.good()) {
        ASSERT_ND(word != NULL_VALUE && doc != NULL_VALUE && count != NULL_VALUE);
        // update the doc counter
        ndocs = std::max(ndocs, doc + 1);
        // Assert valid word
        ASSERT_ND(word < nwords);
        // Update the words in document counter
        // Add all the tokens
        Token tok; tok.word = word; tok.doc = doc;
        for (size_t i = 0; i < count; ++i) {
            tokens.push_back(tok);
        }
      }
    }
    fin.close();
    ntokens = tokens.size();
  }
};

/** A copy of essential parts of Corpus, which can be placed in shared memory. No vector/string */
struct SharedInputs {
  uint32_t  nwords;
  uint32_t  ndocs;
  uint32_t  ntokens;
  uint32_t  ntopics;
  uint32_t  niterations;
  double    alpha;
  double    beta;
  /**
   * Actually "Token tokens[ntokens]" and then "TopicId topics[ntokens]".
   * Note that you can't do sizeof(SharedCorpus).
   */
  char      dynamic_data[8];

  static uint64_t get_object_size(uint32_t ntokens) {
    return sizeof(SharedInputs) + (sizeof(Token) + sizeof(TopicId)) * ntokens;
  }

  void init(const Corpus& original) {
    nwords = original.nwords;
    ndocs = original.ndocs;
    ntokens = original.ntokens;
    ntopics = FLAGS_ntopics;
    niterations = 0;  // set separately
    alpha = FLAGS_alpha;
    beta = FLAGS_beta;
    std::memcpy(get_tokens_data(), &(original.tokens[0]), sizeof(Token) * ntokens);
    std::memset(get_topics_data(), 0xFFU, sizeof(TopicId) * ntokens);  // null topic ID
  }

  Token* get_tokens_data() {
    return reinterpret_cast<Token*>(dynamic_data);
  }
  TopicId* get_topics_data() {
    return reinterpret_cast<TopicId*>(dynamic_data + sizeof(Token) * ntokens);
  }
};

/** Core implementation of LDA experiment. */
class LdaWorker {
 public:
  enum Constant {
    kIntNormalizer = 1 << 24,
    kNtApplyBatch = 1 << 8,
  };
  explicit LdaWorker(const proc::ProcArguments& args) {
    engine = args.engine_;
    context = args.context_;
    shared_inputs = reinterpret_cast<SharedInputs*>(
      engine->get_soc_manager()->get_shared_memory_repo()->get_global_user_memory());
    *args.output_used_ = 0;
  }
  ErrorStack on_lda_worker_task() {
    uint16_t thread_id = context->get_thread_global_ordinal();
    LOG(INFO) << "Thread-" << thread_id << " started";
    unirand.set_current_seed(thread_id);

    Token* shared_tokens = shared_inputs->get_tokens_data();
    ntopics = shared_inputs->ntopics;
    ntopics_aligned = assorted::align8(ntopics);
    numwords = shared_inputs->nwords;

    TopicId* shared_topics = shared_inputs->get_topics_data();
    const uint32_t numtokens = shared_inputs->ntokens;
    const uint32_t niterations = shared_inputs->niterations;
    const uint16_t all_threads = engine->get_options().thread_.get_total_thread_count();
    tokens_per_core = numtokens / all_threads;
    int_a = static_cast<uint64_t>(shared_inputs->alpha * kIntNormalizer);
    int_b = static_cast<uint64_t>(shared_inputs->beta * kIntNormalizer);

    // allocate tmp memory on local NUMA node.
    // These are the only memory that need dynamic allocation in main_loop
    uint64_t tmp_memory_size = ntopics_aligned * sizeof(uint64_t)  // conditional
        + tokens_per_core * sizeof(Token)  // assigned_tokens
        + ntopics_aligned * sizeof(Count)  // record_td
        + ntopics_aligned * sizeof(Count)  // record_tw
        + ntopics_aligned * sizeof(Count)  // record_t
        + ntopics_aligned * sizeof(Count)  // record_t_diff
        + tokens_per_core * sizeof(TopicId);  // topics_tmp
    tmp_memory.alloc(
      tmp_memory_size,
      1 << 21,
      memory::AlignedMemory::kNumaAllocOnnode,
      context->get_numa_node());
    char* tmp_block = reinterpret_cast<char*>(tmp_memory.get_block());

    int_conditional = reinterpret_cast<uint64_t*>(tmp_block);
    tmp_block += ntopics_aligned * sizeof(uint64_t);

    assigned_tokens = reinterpret_cast<Token*>(tmp_block);
    tmp_block += tokens_per_core * sizeof(Token);

    record_td = reinterpret_cast<Count*>(tmp_block);
    tmp_block += ntopics_aligned * sizeof(Count);

    record_tw = reinterpret_cast<Count*>(tmp_block);
    tmp_block += ntopics_aligned * sizeof(Count);

    record_t = reinterpret_cast<Count*>(tmp_block);
    tmp_block += ntopics_aligned * sizeof(Count);

    record_t_diff = reinterpret_cast<Count*>(tmp_block);
    tmp_block += ntopics_aligned * sizeof(Count);

    topics_tmp = reinterpret_cast<TopicId*>(tmp_block);
    tmp_block += tokens_per_core * sizeof(TopicId);

    uint64_t size_check = tmp_block - reinterpret_cast<char*>(tmp_memory.get_block());
    ASSERT_ND(size_check == tmp_memory_size);

    // initialize with previous result (if this is burnin, all null-topic)
    uint64_t assign_pos = tokens_per_core * thread_id;
    std::memcpy(topics_tmp, shared_topics + assign_pos, tokens_per_core * sizeof(TopicId));
    std::memcpy(assigned_tokens, shared_tokens + assign_pos, tokens_per_core * sizeof(Token));

    n_td = storage::array::ArrayStorage(engine, kNtd);
    n_tw = storage::array::ArrayStorage(engine, kNtw);
    n_t = storage::array::ArrayStorage(engine, kNt);
    ASSERT_ND(n_td.exists());
    ASSERT_ND(n_tw.exists());
    ASSERT_ND(n_t.exists());

    nchanges = 0;
    std::memset(int_conditional, 0, sizeof(uint64_t) * ntopics_aligned);
    std::memset(record_td, 0, sizeof(Count) * ntopics_aligned);
    std::memset(record_tw, 0, sizeof(Count) * ntopics_aligned);
    std::memset(record_t, 0, sizeof(Count) * ntopics_aligned);
    std::memset(record_t_diff, 0, sizeof(Count) * ntopics_aligned);

    // Loop over all the tokens
    for (size_t gn = 0; gn < niterations; ++gn) {
      LOG(INFO) << "Thread-" << thread_id << " iteration " << gn << "/" << niterations;
      WRAP_ERROR_CODE(main_loop());
    }

    // copy back the topics to shared. it's per-core, so this is the only thread that modifies it.
    std::memcpy(shared_topics + assign_pos, topics_tmp, tokens_per_core * sizeof(TopicId));

    const uint64_t ntotal = niterations * tokens_per_core;
    LOG(INFO) << "Thread-" << thread_id << " done. nchanges/ntotal="
      << nchanges << "/" << ntotal << "=" << (static_cast<double>(nchanges) / ntotal);
    return kRetOk;
  }

 private:
  Engine* engine;
  thread::Thread* context;
  SharedInputs* shared_inputs;
  uint16_t ntopics;
  uint16_t ntopics_aligned;
  uint32_t numwords;
  uint32_t tokens_per_core;
  uint64_t nchanges;
  uint64_t int_a;
  uint64_t int_b;
  storage::array::ArrayStorage n_td;  // d records of ntopics Count
  storage::array::ArrayStorage n_tw;  // w records of ntopics Count
  storage::array::ArrayStorage n_t;  // 1 record of ntopics Count
  assorted::UniformRandom unirand;

  /** local memory to back the followings */
  memory::AlignedMemory tmp_memory;
  uint64_t* int_conditional;  // [ntopics]
  Token* assigned_tokens;  // [tokens_per_core]
  Count* record_td;  // [ntopics]
  Count* record_tw;  // [ntopics]
  Count* record_t;  // [ntopics]
  Count* record_t_diff;  // [ntopics]. This batches changes to n_t, which is so contentious.
  TopicId* topics_tmp;  // [tokens_per_core]

  ErrorCode main_loop() {
    xct::XctManager* xct_manager = engine->get_xct_manager();
    uint16_t batched_t = 0;
    for (size_t i = 0; i < tokens_per_core; ++i) {
      // Get the word and document for the ith token
      const WordId w = assigned_tokens[i].word;
      const DocId d = assigned_tokens[i].doc;
      const TopicId old_topic = topics_tmp[i];

      CHECK_ERROR_CODE(xct_manager->begin_xct(context, xct::kDirtyRead));

      // First, read from the global arrays
      CHECK_ERROR_CODE(n_td.get_record(context, d, record_td));
      CHECK_ERROR_CODE(n_tw.get_record(context, w, record_tw));

      // Construct the conditional
      uint64_t normalizer = 0;
      // We do the following calculation without double/float to speed up.
      // Notation: int_x := x * N (eg, int_ntd := ntd * N, int_alpha = alpha * N)
      //           where N is some big number, such as 2^24
      //      conditional[t] = (a + ntd) * (b + ntw) / (b * nwords + nt)
      // <=>  conditional[t] = (int_a + int_ntd) * (b + ntw) / (int_b * nwords + int_nt)
      // <=>  int_conditional[t] = (int_a + int_ntd) * (int_b + int_ntw) / (int_b * nwords + int_nt)
      for (TopicId t = 0; t < ntopics; ++t) {
        uint64_t int_ntd = record_td[t], int_ntw = record_tw[t], int_nt = record_t[t];
        if (t == old_topic) {
          // locally apply the decrements. we apply it to global array
          // when we are sure that new_topic != old_topic
          --int_ntd;
          --int_ntw;
          --int_nt;
        }
        int_ntd = check_overflow(int_ntd) * kIntNormalizer;
        int_ntw = check_overflow(int_ntw) * kIntNormalizer;
        int_nt = check_overflow(int_nt) * kIntNormalizer;

        int_conditional[t] = (int_a + int_ntd) * (int_b + int_ntw) / (int_b * numwords + int_nt);
        normalizer += int_conditional[t];
      }

      // Draw a new valuez
      TopicId new_topic = topic_multinomial(normalizer);

      // Update the topic assignment and counters
      if (new_topic != old_topic) {
        nchanges++;
        topics_tmp[i] = new_topic;

        // We apply the inc/dec to global array only in this case.
        if (old_topic != kNullTopicId) {
          // Remove the word from the old counters
          uint16_t offset = old_topic * sizeof(Count);
          CHECK_ERROR_CODE(n_td.increment_record_oneshot<Count>(context, d, -1, offset));
          CHECK_ERROR_CODE(n_tw.increment_record_oneshot<Count>(context, w, -1, offset));
          // change to t are batched. will be applied later
          --record_t[old_topic];
          --record_t_diff[old_topic];
        }

        // Add the word to the new counters
        uint16_t offset = new_topic * sizeof(Count);
        CHECK_ERROR_CODE(n_td.increment_record_oneshot<Count>(context, d, 1, offset));
        CHECK_ERROR_CODE(n_tw.increment_record_oneshot<Count>(context, w, 1, offset));
        ++record_t[new_topic];
        ++record_t_diff[new_topic];

        Epoch ep;
        // because this is not serializable level and the only update is one-shot increment,
        // no race is possible.
        CHECK_ERROR_CODE(xct_manager->precommit_xct(context, &ep));

        ++batched_t;
        if (batched_t >= kNtApplyBatch) {
          CHECK_ERROR_CODE(sync_record_t());
          batched_t = 0;
        }
      } else {
        CHECK_ERROR_CODE(xct_manager->abort_xct(context));
      }
    }

    CHECK_ERROR_CODE(sync_record_t());

    return kErrorCodeOk;
  }

  TopicId topic_multinomial(uint64_t normalizer) {
    uint64_t sample = unirand.next_uint64() % normalizer;
    uint64_t cur_sum = 0;
    for (TopicId t = 0; t < ntopics; ++t) {
      cur_sum += int_conditional[t];
      if (sample <= cur_sum) {
        return t;
      }
    }
    return ntopics - 1U;
  }

  ErrorCode sync_record_t() {
    xct::XctManager* xct_manager = engine->get_xct_manager();
    Epoch ep;
    // first, "flush" batched increments/decrements to the global n_t
    // this transactionally applies all the inc/dec.
    CHECK_ERROR_CODE(xct_manager->begin_xct(context, xct::kDirtyRead));
    for (TopicId t = 0; t < ntopics_aligned; t += 4) {
      // for each uint64_t. we reserved additional size (align8) for simplifying this.
      uint64_t diff = *reinterpret_cast<uint64_t*>(&record_t_diff[t]);
      if (diff != 0) {
        uint16_t offset = t * sizeof(Count);
        CHECK_ERROR_CODE(n_t.increment_record_oneshot<uint64_t>(context, 0, diff, offset));
      }
    }
    CHECK_ERROR_CODE(xct_manager->precommit_xct(context, &ep));

    // then, retrive a fresh new copy of record_t, which contains updates from other threads.
    std::memset(record_t, 0, sizeof(Count) * ntopics_aligned);
    std::memset(record_t_diff, 0, sizeof(Count) * ntopics_aligned);
    CHECK_ERROR_CODE(xct_manager->begin_xct(context, xct::kDirtyRead));
    CHECK_ERROR_CODE(n_t.get_record(context, 0, record_t));
    CHECK_ERROR_CODE(xct_manager->abort_xct(context));
    return kErrorCodeOk;
  }
};

ErrorStack lda_worker_task(const proc::ProcArguments& args) {
  LdaWorker worker(args);
  return worker.on_lda_worker_task();
}

void run_workers(Engine* engine, SharedInputs* shared_inputs, uint32_t niterations) {
  const EngineOptions& options = engine->get_options();

  auto* thread_pool = engine->get_thread_pool();
  std::vector< thread::ImpersonateSession > sessions;

  shared_inputs->niterations = niterations;
  for (uint16_t node = 0; node < options.thread_.group_count_; ++node) {
    for (uint16_t ordinal = 0; ordinal < options.thread_.thread_count_per_group_; ++ordinal) {
      thread::ImpersonateSession session;
      bool ret = thread_pool->impersonate_on_numa_node(
        node,
        "worker_task",
        nullptr,
        0,
        &session);
      if (!ret) {
        LOG(FATAL) << "Couldn't impersonate";
      }
      sessions.emplace_back(std::move(session));
    }
  }
  LOG(INFO) << "Started running! #iter=" << niterations;
  for (uint32_t i = 0; i < sessions.size(); ++i) {
    LOG(INFO) << "result[" << i << "]=" << sessions[i].get_result();
    COERCE_ERROR(sessions[i].get_result());
    sessions[i].release();
  }
  LOG(INFO) << "Completed";
}

ErrorStack lda_populate_array_task(const proc::ProcArguments& args) {
  ASSERT_ND(args.input_len_ == sizeof(storage::StorageId));
  storage::StorageId storage_id = *reinterpret_cast<const storage::StorageId*>(args.input_buffer_);
  storage::array::ArrayStorage array(args.engine_, storage_id);
  ASSERT_ND(array.exists());
  thread::Thread* context = args.context_;

  uint16_t nodes = args.engine_->get_options().thread_.group_count_;
  uint16_t node = context->get_numa_node();
  uint64_t records_per_node = array.get_array_size() / nodes;
  storage::array::ArrayOffset begin = records_per_node * node;
  storage::array::ArrayOffset end = records_per_node * (node + 1U);
  WRAP_ERROR_CODE(array.prefetch_pages(context, true, false, begin, end));

  LOG(INFO) << "Population of " << array.get_name() << " done in node-" << node;
  return kRetOk;
}

void create_count_array(
  Engine* engine,
  const char* name,
  uint64_t records,
  uint64_t counts) {
  Epoch ep;
  uint64_t counts_aligned = assorted::align8(counts);
  storage::array::ArrayMetadata meta(name, counts_aligned * sizeof(Count), records);
  COERCE_ERROR(engine->get_storage_manager()->create_storage(&meta, &ep));
  LOG(INFO) << "Populating empty " << name << "(id=" << meta.id_ << ")...";

  // populate it with one thread per node.
  const EngineOptions& options = engine->get_options();
  auto* thread_pool = engine->get_thread_pool();
  std::vector< thread::ImpersonateSession > sessions;
  for (uint16_t node = 0; node < options.thread_.group_count_; ++node) {
    thread::ImpersonateSession session;
    bool ret = thread_pool->impersonate_on_numa_node(
      node,
      "populate_array_task",
      &meta.id_,
      sizeof(meta.id_),
      &session);
    if (!ret) {
      LOG(FATAL) << "Couldn't impersonate";
    }
    sessions.emplace_back(std::move(session));
  }
  for (uint32_t i = 0; i < sessions.size(); ++i) {
    LOG(INFO) << "populate result[" << i << "]=" << sessions[i].get_result();
    COERCE_ERROR(sessions[i].get_result());
    sessions[i].release();
  }
  LOG(INFO) << "Populated empty " << name;
}

double run_experiment(Engine* engine, const Corpus &corpus) {
  const EngineOptions& options = engine->get_options();

  // Input information placed in shared memory
  SharedInputs* shared_inputs = reinterpret_cast<SharedInputs*>(
    engine->get_soc_manager()->get_shared_memory_repo()->get_global_user_memory());
  ASSERT_ND(options.soc_.shared_user_memory_size_kb_ * (1ULL << 10)
    >= SharedInputs::get_object_size(corpus.ntokens));
  shared_inputs->init(corpus);

  // create empty tables.
  create_count_array(engine, kNtd, corpus.ndocs, FLAGS_ntopics);
  create_count_array(engine, kNtw, corpus.nwords, FLAGS_ntopics);
  create_count_array(engine, kNt, 1, FLAGS_ntopics);

  LOG(INFO) << "Created arrays";

  if (FLAGS_profile) {
    COERCE_ERROR(engine->get_debug()->start_profile("lda.prof"));
  }

  debugging::StopWatch watch;
  run_workers(engine, shared_inputs, FLAGS_nburnin);
  LOG(INFO) << "Completed burnin!";
  run_workers(engine, shared_inputs, FLAGS_nsamples);
  LOG(INFO) << "Completed final sampling!";

  watch.stop();

  LOG(INFO) << "Experiment ended. Elapsed time=" << watch.elapsed_sec() << "sec";
  if (FLAGS_profile) {
    engine->get_debug()->stop_profile();
    LOG(INFO) << "Check out the profile result: pprof --pdf lda lda.prof > lda.pdf; okular lda.pdf";
  }

  LOG(INFO) << "Shutting down...";

  LOG(INFO) << engine->get_memory_manager()->dump_free_memory_stat();
  return watch.elapsed_sec();
}

int driver_main(int argc, char** argv) {
  gflags::SetUsageMessage("LDA sampler code");
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  fs::Path dictionary_fname(FLAGS_dictionary_fname);
  fs::Path counts_fname(FLAGS_counts_fname);

  if (!fs::exists(dictionary_fname)) {
    std::cerr << "The dictionary file doesn't exist: " << dictionary_fname;
    return 1;
  } else if (!fs::exists(counts_fname)) {
    std::cerr << "The count file doesn't exist: " << counts_fname;
    return 1;
  }

  fs::Path folder("/dev/shm/foedus_lda");
  if (fs::exists(folder)) {
    fs::remove_all(folder);
  }
  if (!fs::create_directories(folder)) {
    std::cerr << "Couldn't create " << folder << ". err=" << assorted::os_error();
    return 1;
  }

  EngineOptions options;

  fs::Path savepoint_path(folder);
  savepoint_path /= "savepoint.xml";
  options.savepoint_.savepoint_path_.assign(savepoint_path.string());
  ASSERT_ND(!fs::exists(savepoint_path));

  std::cout << "NUMA node count=" << static_cast<int>(options.thread_.group_count_) << std::endl;
  uint16_t thread_count = FLAGS_nthreads;
  if (thread_count < options.thread_.group_count_) {
    std::cout << "nthreads less than socket count. Using subset of sockets" << std::endl;
    options.thread_.group_count_ = thread_count;
    options.thread_.thread_count_per_group_ = 1;
  } else if (thread_count % options.thread_.group_count_ != 0) {
    std::cout << "nthreads not multiply of #sockets. adjusting" << std::endl;
    options.thread_.thread_count_per_group_ = thread_count / options.thread_.group_count_;
  } else {
    options.thread_.thread_count_per_group_ = thread_count / options.thread_.group_count_;
  }

  options.snapshot_.folder_path_pattern_ = "/dev/shm/foedus_lda/snapshot/node_$NODE$";
  options.snapshot_.snapshot_interval_milliseconds_ = 100000000U;

  options.log_.folder_path_pattern_ = "/dev/shm/foedus_lda/log/node_$NODE$/logger_$LOGGER$";
  options.log_.loggers_per_node_ = 1;
  options.log_.flush_at_shutdown_ = false;
  options.log_.log_file_size_mb_ = 1 << 16;
  options.log_.emulation_.null_device_ = true;

  options.memory_.page_pool_size_mb_per_node_ = 1 << 8;
  options.cache_.snapshot_cache_size_mb_per_node_ = 1 << 6;

  if (FLAGS_fork_workers) {
    std::cout << "Will fork workers in child processes" << std::endl;
    options.soc_.soc_type_ = kChildForked;
  }

  options.debugging_.debug_log_min_threshold_
    = debugging::DebuggingOptions::kDebugLogInfo;
    // = debugging::DebuggingOptions::kDebugLogWarning;
  options.debugging_.verbose_modules_ = "";
  options.debugging_.verbose_log_level_ = -1;

  std::cout << "Loading the corpus." << std::endl;
  Corpus corpus(dictionary_fname, counts_fname);

  std::cout << "Number of words:   " << corpus.nwords << std::endl
            << "Number of docs:    " << corpus.ndocs << std::endl
            << "Number of tokens:  " << corpus.ntokens << std::endl
            << "Number of topics:  " << FLAGS_ntopics << std::endl
            << "Number of threads: " << options.thread_.get_total_thread_count() << std::endl
            << "Alpha:             " << FLAGS_alpha   << std::endl
            << "Beta:              " << FLAGS_beta    << std::endl;

  options.soc_.shared_user_memory_size_kb_
    = (SharedInputs::get_object_size(corpus.ntokens) >> 10) + 64;

  double elapsed;
  {
    Engine engine(options);
    engine.get_proc_manager()->pre_register("worker_task", lda_worker_task);
    engine.get_proc_manager()->pre_register("populate_array_task", lda_populate_array_task);
    COERCE_ERROR(engine.initialize());
    {
      UninitializeGuard guard(&engine);
      elapsed = run_experiment(&engine, corpus);
      COERCE_ERROR(engine.uninitialize());
    }
  }

  std::cout << "All done. elapsed time=" << elapsed << "sec" << std::endl;

  return 0;
}

}  // namespace graphlda
}  // namespace foedus

int main(int argc, char **argv) {
  return foedus::graphlda::driver_main(argc, argv);
}


// Boost.Units - A C++ library for zero-overhead dimensional analysis and
// unit/quantity manipulation and conversion
//
// Copyright (C) 2003-2008 Matthias Christian Schabel
// Copyright (C) 2008 Steven Watanabe
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_UNITS_CODATA_PROTON_CONSTANTS_HPP
#define BOOST_UNITS_CODATA_PROTON_CONSTANTS_HPP

#include <boost/units/quantity.hpp>
#include <boost/units/static_constant.hpp>

#include <boost/units/systems/detail/constants.hpp>
#include <boost/units/systems/si/amount.hpp>
#include <boost/units/systems/si/area.hpp>
#include <boost/units/systems/si/electric_charge.hpp>
#include <boost/units/systems/si/energy.hpp>
#include <boost/units/systems/si/frequency.hpp>
#include <boost/units/systems/si/length.hpp>
#include <boost/units/systems/si/mass.hpp>
#include <boost/units/systems/si/magnetic_flux_density.hpp>
#include <boost/units/systems/si/time.hpp>
#include <boost/units/systems/si/wavenumber.hpp>

#include <boost/units/systems/si/codata/typedefs.hpp>

/// \file
/// CODATA recommended values of fundamental atomic and nuclear constants
/// CODATA 2006 values as of 2007/03/30

namespace boost {

namespace units {

namespace si {

namespace constants {

namespace codata {

/// CODATA recommended values of the fundamental physical constants: NIST SP 961

/// proton mass
BOOST_UNITS_PHYSICAL_CONSTANT(m_p,quantity<mass>,1.672621637e-27*kilograms,8.3e-35*kilograms);
/// proton-electron mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(m_p_over_m_e,quantity<dimensionless>,1836.15267247*dimensionless(),8.0e-7*dimensionless());
/// proton-muon mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(m_p_over_m_mu,quantity<dimensionless>,8.88024339*dimensionless(),2.3e-7*dimensionless());
/// proton-tau mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(m_p_over_m_tau,quantity<dimensionless>,0.528012*dimensionless(),8.6e-5*dimensionless());
/// proton-neutron mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(m_p_over_m_n,quantity<dimensionless>,0.99862347824*dimensionless(),4.6e-10*dimensionless());
/// proton charge to mass ratio
BOOST_UNITS_PHYSICAL_CONSTANT(e_over_m_p,quantity<electric_charge_over_mass>,9.57883392e7*coulombs/kilogram,2.4e0*coulombs/kilogram);
/// proton molar mass
BOOST_UNITS_PHYSICAL_CONSTANT(M_p,quantity<mass_over_amount>,1.00727646677e-3*kilograms/mole,1.0e-13*kilograms/mole);
/// proton Compton wavelength
BOOST_UNITS_PHYSICAL_CONSTANT(lambda_C_p,quantity<length>,1.3214098446e-15*meters,1.9e-24*meters);
/// proton rms charge radius
BOOST_UNITS_PHYSICAL_CONSTANT(R_p,quantity<length>,0.8768e-15*meters,6.9e-18*meters);
/// proton magnetic moment
BOOST_UNITS_PHYSICAL_CONSTANT(mu_p,quantity<energy_over_magnetic_flux_density>,1.410606662e-26*joules/tesla,3.7e-34*joules/tesla);
/// proton-Bohr magneton ratio
BOOST_UNITS_PHYSICAL_CONSTANT(mu_p_over_mu_B,quantity<dimensionless>,1.521032209e-3*dimensionless(),1.2e-11*dimensionless());
/// proton-nuclear magneton ratio
BOOST_UNITS_PHYSICAL_CONSTANT(mu_p_over_mu_N,quantity<dimensionless>,2.792847356*dimensionless(),2.3e-8*dimensionless());
/// proton g-factor
BOOST_UNITS_PHYSICAL_CONSTANT(g_p,quantity<dimensionless>,5.585694713*dimensionless(),4.6e-8*dimensionless());
/// proton-neutron magnetic moment ratio
BOOST_UNITS_PHYSICAL_CONSTANT(mu_p_over_mu_n,quantity<dimensionless>,-1.45989806*dimensionless(),3.4e-7*dimensionless());
/// shielded proton magnetic moment
BOOST_UNITS_PHYSICAL_CONSTANT(mu_p_prime,quantity<energy_over_magnetic_flux_density>,1.410570419e-26*joules/tesla,3.8e-34*joules/tesla);
/// shielded proton-Bohr magneton ratio
BOOST_UNITS_PHYSICAL_CONSTANT(mu_p_prime_over_mu_B,quantity<dimensionless>,1.520993128e-3*dimensionless(),1.7e-11*dimensionless());
/// shielded proton-nuclear magneton ratio
BOOST_UNITS_PHYSICAL_CONSTANT(mu_p_prime_over_mu_N,quantity<dimensionless>,2.792775598*dimensionless(),3.0e-8*dimensionless());
/// proton magnetic shielding correction
BOOST_UNITS_PHYSICAL_CONSTANT(sigma_p_prime,quantity<dimensionless>,25.694e-6*dimensionless(),1.4e-8*dimensionless());
/// proton gyromagnetic ratio
BOOST_UNITS_PHYSICAL_CONSTANT(gamma_p,quantity<frequency_over_magnetic_flux_density>,2.675222099e8/second/tesla,7.0e0/second/tesla);
/// shielded proton gyromagnetic ratio
BOOST_UNITS_PHYSICAL_CONSTANT(gamma_p_prime,quantity<frequency_over_magnetic_flux_density>,2.675153362e8/second/tesla,7.3e0/second/tesla);

} // namespace codata

} // namespace constants

} // namespace si

} // namespace units

} // namespace boost

#endif // BOOST_UNITS_CODATA_PROTON_CONSTANTS_HPP


/* Siconos is a program dedicated to modeling, simulation and control
 * of non smooth dynamical systems.
 *
 * Copyright 2021 INRIA.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
*/

#include "SiconosMatrix.hpp"
#include <assert.h>                                   // for assert
#include <math.h>                                     // for fabs
#include <float.h>                                     // for DBL_EPSILON
#include <algorithm>                                  // for max, min, lower...
#include <boost/numeric/ublas/detail/config.hpp>      // for noalias, noalia...
#include <boost/numeric/ublas/detail/iterator.hpp>    // for bidirectional_i...
#include <boost/numeric/ublas/matrix_expression.hpp>  // for matrix_vector_b...
#include <boost/numeric/ublas/matrix_proxy.hpp>       // for matrix_range
#include <boost/numeric/ublas/matrix_sparse.hpp>      // for compressed_matr...
#include <boost/numeric/ublas/storage.hpp>            // for unbounded_array
#include <boost/numeric/ublas/vector.hpp>             // for vector
#include <memory>                                     // for __shared_ptr_ac...
#include <utility>                                    // for swap
#include <vector>                                     // for vector, operator==
#include "SiconosAlgebra.hpp"
#include "BlockMatrix.hpp"                            // for BlockMatrix
#include "SiconosVector.hpp"                          // for SiconosVector
#include "SimpleMatrixFriends.hpp"                    // for isComparableTo

#include "CSparseMatrix_internal.h"                   // for CSparseMatrix
#include "NumericsSparseMatrix.h"                     // for NSM_fix_csc
#include "SiconosException.hpp"

// Constructor with the type-number
SiconosMatrix::SiconosMatrix(Siconos::UBLAS_TYPE type): _num(type), _isSymmetric(false), _isPositiveDefinite(false)
{}

const SP::Index SiconosMatrix::tabRow() const
{
  THROW_EXCEPTION("not implemented for this type of matrix (Simple?) reserved to BlockMatrix.");
}

const SP::Index SiconosMatrix::tabCol() const
{
  THROW_EXCEPTION("not implemented for this type of matrix (Simple?) reserved to BlockMatrix.");
}

//=====================
// matrices comparison
//=====================
SiconosMatrix& operator *=(SiconosMatrix& m, const double& s)
{
  if(m._num == Siconos::BLOCK)  // BlockMatrix
  {
    BlockMatrix& mB = static_cast<BlockMatrix&>(m);
    BlocksMat::iterator1 it;
    BlocksMat::iterator2 it2;
    for(it = mB._mat->begin1(); it != mB._mat->end1(); ++it)
    {
      for(it2 = it.begin(); it2 != it.end(); ++it2)
        (**it2) *= s;
    }
  }
  else if(m._num == Siconos::DENSE)
    *m.dense() *= s;
  else if(m._num == Siconos::TRIANGULAR)
    *m.triang() *= s;
  else if(m._num == Siconos::SYMMETRIC)
    *m.sym() *= s;
  else if(m._num == Siconos::SPARSE)
    *m.sparse() *= s;
  else if(m._num == Siconos::SPARSE_COORDINATE)
    *m.sparseCoordinate() *= s;
  else if(m._num == Siconos::BANDED)
    *m.banded() *= s;
  else if(m._num == Siconos::ZERO) {}  // nothing!
  else //if(_num == 7)
    THROW_EXCEPTION("invalid type of matrix");

  return m;
}

SiconosMatrix& operator /=(SiconosMatrix& m, const double& s)
{
  if(m._num == Siconos::BLOCK)  // BlockMatrix
  {
    BlockMatrix& mB = static_cast<BlockMatrix&>(m);
    BlocksMat::iterator1 it;
    BlocksMat::iterator2 it2;
    for(it = mB._mat->begin1(); it != mB._mat->end1(); ++it)
    {
      for(it2 = it.begin(); it2 != it.end(); ++it2)
        (**it2) /= s;
    }
  }
  else if(m._num == Siconos::DENSE)
    *m.dense() /= s;
  else if(m._num == Siconos::TRIANGULAR)
    *m.triang() /= s;
  else if(m._num == Siconos::SYMMETRIC)
    *m.sym() /= s;
  else if(m._num == Siconos::SPARSE)
    *m.sparse() /= s;
  else if(m._num == Siconos::SPARSE_COORDINATE)
    *m.sparseCoordinate() /= s;
  else if(m._num == Siconos::BANDED)
    *m.banded() /= s;
  else if(m._num == Siconos::ZERO) {}  // nothing!
  else //if(_num == 7)
    THROW_EXCEPTION("invalid type of matrix");

  return m;
}

size_t SiconosMatrix::nnz(double tol)
{
  size_t nnz = 0;
  if(_num == Siconos::DENSE)  //dense
  {
    double* arr = getArray();
    for(size_t i = 0; i < size(0)*size(1); ++i)
    {
      if(fabs(arr[i]) > tol)
      {
        nnz++;
      }
    }
  }
  else if(_num == Siconos::SPARSE)
  {
    nnz = sparse()->nnz();
  }
  else
    THROW_EXCEPTION("not implemented for the given matrix type");

  return nnz;

}

bool SiconosMatrix::fillCSC(CSparseMatrix* csc, size_t row_off, size_t col_off, double tol)
{
  assert(csc);
  double* Mx = csc->x; // data
  CS_INT* Mi = csc->i; // row indx
  CS_INT* Mp = csc->p; // column pointers

  assert(Mp[col_off] >= 0);
  size_t nz = csc->p[col_off];

  size_t nrow = size(0);
  size_t ncol = size(1);

  CS_INT pval = Mp[col_off];

  if(_num == Siconos::DENSE)  //dense
  {
    double* arr = getArray();
    for(size_t j = 0, joff = col_off; j < ncol; ++j)
    {
      for(size_t i = 0; i < nrow; ++i)
      {
        // col-major
        double elt_val = arr[i + j*nrow];
        // std::cout << " a(i=" << i << ",j=" << j << ") = "<< elt_val << std::endl;
        if(fabs(elt_val) > tol)
        {
          Mx[pval] = elt_val;
          Mi[pval] = i + row_off;
          // std::cout << "Mx[" <<pval <<"] = " << Mx[pval]<<   std::endl;
          // std::cout << "Mp[" <<pval <<"] = " << Mi[pval]<<   std::endl;
          ++pval;
        }
      }
      // std::cout << "joff" << joff << std::endl;
      Mp[++joff] = pval;

    }
  }
  else if(_num == Siconos::SPARSE)
  {
    const Index& ptr = sparse()->index1_data();
    const Index& indx = sparse()->index2_data();
    const ublas::unbounded_array<double>& vals = sparse()->value_data();

    size_t nnz =  sparse()->nnz();

    assert(ptr.size() == ncol + 1);
    assert(indx.size() == nnz);
    assert(vals.size() == nnz);

    for(size_t i = 0; i < nnz; ++i)
    {
      Mx[pval] = vals[i];
      Mi[pval++] = row_off + indx[i];
    }
    for(size_t j = 1, joff = col_off + 1; j < ncol+1; ++j, ++joff)
    {
      Mp[joff] = nz + ptr[j];
    }
  }
  else
  {
    THROW_EXCEPTION("not implemented for the given matrix type");
  }

  return true;
}

bool SiconosMatrix::fillCSC(CSparseMatrix* csc, double tol)
{
  assert(csc);
  double* Mx = csc->x; // data
  CS_INT* Mi = csc->i; // row indx
  CS_INT* Mp = csc->p; // column pointers

  size_t nrow = size(0);
  size_t ncol = size(1);

  CS_INT pval = 0;

  if(_num == Siconos::DENSE)  //dense
  {
    double* arr = getArray();
    for(size_t j = 0, joff = 0; j < ncol; ++j)
    {
      for(size_t i = 0; i < nrow; ++i)
      {
        // col-major
        double elt_val = arr[i + j*nrow];
        // std::cout << " a(i=" << i << ",j=" << j << ") = "<< elt_val << std::endl;
        if(fabs(elt_val) > tol)
        {
          Mx[pval] = elt_val;
          Mi[pval] = i;
          // std::cout << "Mx[" <<pval <<"] = " << Mx[pval]<<   std::endl;
          // std::cout << "Mp[" <<pval <<"] = " << Mi[pval]<<   std::endl;
          ++pval;
        }
      }
      // std::cout << "joff" << joff << std::endl;
      Mp[++joff] = pval;

    }
  }
  else if(_num == Siconos::SPARSE)
  {
    const Index& ptr = sparse()->index1_data();
    const Index& indx = sparse()->index2_data();
    const ublas::unbounded_array<double>& vals = sparse()->value_data();

    size_t nnz =  sparse()->nnz();

    assert(ptr.size() == ncol + 1);
    assert(indx.size() >= nnz);
    assert(vals.size() >= nnz);

    for(size_t i = 0; i < nnz; ++i)
    {
      Mx[pval] = vals[i];
      Mi[pval++] = indx[i];
    }
    for(size_t j = 0; j < ncol+1; ++j)
    {
      Mp[j] = ptr[j];
    }
  }
  else
  {
    THROW_EXCEPTION("not implemented for the given matrix type");
  }

  return true;
}

bool SiconosMatrix::fromCSC(CSparseMatrix* csc)
{
  assert(csc);

  NSM_sort_csc(csc);

  double* Mx = csc->x; // data
  CS_INT* Mi = csc->i; // row indx
  CS_INT* Mp = csc->p; // column pointers
  CS_INT n = csc->n;
  // CS_INT m = csc->m;

  // size_t nnz = csc->p[n];

  if(_num == Siconos::SPARSE)
  {
    sparse()->clear();
    CS_INT pval=0;
    // push_back in order should be in constant time
    // http://www.guwi17.de/ublas/matrix_sparse_usage.html
    for (CS_INT col =0; col < n; col++)
    {
      for (CS_INT p = Mp[col]; p < Mp[col+1];  p++)
      {
        sparse()->push_back(Mi[pval], col, Mx[pval]);
        pval++;
      }
    }

    // not able to work directly on the contents of the sparse matrix.
    // sparse()->resize(m, nnz, false);
    // //sparse()->set_filled(nnz, n+1);
    // std::cout << sparse()->size1() << std::endl;
    // std::cout << sparse()->size2() << std::endl;

    // Index& ptr = sparse()->index1_data();
    // Index& indx = sparse()->index2_data();
    // ublas::unbounded_array<double>& vals = sparse()->value_data();

    // // ptr.resize(n+1);
    // // indx.resize(nnz);
    // // vals.resize(nnz);

    // std::cout << sparse()->filled1() << std::endl;
    // std::cout << sparse()->filled2() << std::endl;

    // // CS_INT pval=0;

    // // for(size_t i = 0; i < nnz; ++i)
    // // {
    // //   vals[i] = Mx[pval];
    // //   indx[i] = Mi[pval++];
    // // }
    // for(size_t i = 0; i < nnz; ++i)
    // {
    //   printf("vals[%i] = %e\t", i, vals[i]);
    //   printf("indx[%i] = %i\t", i, indx[i]);
    // }

    // // for(size_t j = 0 ; j < n+1; ++j)
    // // {
    // //   ptr[j] = Mp[j]  ;
    // // }
    // for(size_t j = 0 ; j < n+1; ++j) printf("ptr[%i] = %i\t", j, ptr[j]);
  }
  else if(_num == Siconos::DENSE)
  {
    CS_INT pval=0;
    // push_back in order should be in constant time
    // http://www.guwi17.de/ublas/matrix_sparse_usage.html
    for (CS_INT col =0; col < n; col++)
    {
      for (CS_INT p = Mp[col]; p < Mp[col+1];  p++)
      {
        setValue(Mi[pval], col, Mx[pval]);
        pval++;
      }
    }
  }
  else
  {
    THROW_EXCEPTION("not implemented for the given matrix type");
  }
  return true;
}

bool SiconosMatrix::fillTriplet(CSparseMatrix* triplet, size_t row_off, size_t col_off, double tol)
{
  assert(triplet);
  size_t nrow = size(0);
  size_t ncol = size(1);

  if(_num == Siconos::DENSE)  //dense
  {
    double* arr = getArray();
    for(size_t j = 0; j < ncol; ++j)
    {
      for(size_t i = 0; i < nrow; ++i)
      {
        // col-major

        CSparseMatrix_zentry(triplet, i + row_off, j + col_off, arr[i + j*nrow], DBL_EPSILON);
      }
    }
  }
  else
  {
    THROW_EXCEPTION("not implemented for the given matrix type");
  }

  return true;
}

std::ostream& operator<<(std::ostream& os, const SiconosMatrix& sm)
{
  os << sm.toString();
  return os;
}

void SiconosMatrix::private_prod(unsigned int startRow, const SiconosVector& x, SiconosVector& y, bool init) const
{
  assert(!(isFactorized()) && "A is Factorized in prod !!");

  // Computes y = subA *x (or += if init = false), subA being a sub-matrix of A, between el. of index (row) startRow and startRow + sizeY

  if(init)  // y = subA * x , else y += subA * x
    y.zero();
  private_addprod(startRow, 0, x, y);
}

/** Computation of y = subA.x

    where subA is a sub-matrix of A, subA[0,0] = A[startRow, startCol]
 */
void SiconosMatrix::private_addprod(unsigned startRow, unsigned int startCol, const SiconosVector& x, SiconosVector& y) const
{
  assert(!(isFactorized()) && "A is Factorized in prod !!");
  assert(!isBlock() && "private_addprod(start,x,y) error: not yet implemented for block matrix.");

  // we take a submatrix subA of A, starting from row startRow to row (startRow+sizeY) and between columns startCol and (startCol+sizeX).
  // Then computation of y = subA*x + y.
  Siconos::UBLAS_TYPE numA = num();
  Siconos::UBLAS_TYPE numY = y.num();
  Siconos::UBLAS_TYPE numX = x.num();
  unsigned int sizeX = x.size();
  unsigned int sizeY = y.size();

  assert(numX == numY && "private_addprod(A,start,x,y) error: not yet implemented for x and y of different types.");

  if(numY == Siconos::DENSE && numX == Siconos::DENSE)
  {

    assert(y.dense() != x.dense());

    if(numA == Siconos::DENSE)
      noalias(*y.dense()) += prod(ublas::subrange(*dense(), startRow, startRow + sizeY, startCol, startCol + sizeX), *x.dense());
    else if(numA == Siconos::TRIANGULAR)
      noalias(*y.dense()) += prod(ublas::subrange(*triang(), startRow, startRow + sizeY, startCol, startCol + sizeX), *x.dense());
    else if(numA == Siconos::SYMMETRIC)
      noalias(*y.dense()) += prod(ublas::subrange(*sym(), startRow, startRow + sizeY, startCol, startCol + sizeX), *x.dense());
    else if(numA == Siconos::SPARSE)
      noalias(*y.dense()) += prod(ublas::subrange(*sparse(), startRow, startRow + sizeY, startCol, startCol + sizeX), *x.dense());
    else //if(numA==Siconos::BANDED)
      noalias(*y.dense()) += prod(ublas::subrange(*banded(), startRow, startRow + sizeY, startCol, startCol + sizeX), *x.dense());
  }
  else // x and y sparse
  {
    if(numA == Siconos::SPARSE)
      *y.sparse() += prod(ublas::subrange(*sparse(), startRow, startRow + sizeY, startCol, startCol + sizeX), *x.sparse());
    else
      THROW_EXCEPTION("not yet implemented for x, y  sparse and A not sparse.");
  }
}


//=======================================================================
// Copyright 2014-2015 David Simmons-Duffin.
// Distributed under the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#include <algorithm>
#include <iostream>
#include <vector>
#include "omp.h"
//Tweak to allow Ubuntu-14.04/gcc-4.8.4 and similar environments to compile
#define BOOST_NO_CXX11_SCOPED_ENUMS
#include <boost/filesystem.hpp>
#undef BOOST_NO_CXX11_SCOPED_ENUMS
#include "SDPSolver.h"
#include "Timers.h"

using boost::filesystem::path;
using boost::timer::nanosecond_type;
using std::cout;

// A note on function conventions: For functions which return large
// objects, we usually use references to the returned objects as
// arguments to a void function, which are then modified in place:
//
// void myFunction(const InputType1 &input1,
//                 const InputType2 &input2,
//                 ...
//                 OutputType1 &output1,
//                 OutputType2 &output2,
//                 ...) {
//   use input1, etc. to modify output1, etc.
// }
//
// We try to mark large inputs as `const' references.  However,
// several MBLAS and MLAPACK functions are not correctly annotated
// with `const's, so we occasionally have input arguments which cannot
// be marked as const because they're used in MBLAS/MLAPACK functions.
// We try to use comments to distinguish which arguments should be
// considered inputs and which should be considered outputs.

/***********************************************************************/
// Create and initialize an SDPSolver for the given SDP and
// SDPSolverParameters

SDPSolver::SDPSolver(const SDP &sdp, const SDPSolverParameters &parameters):
  sdp(sdp),
  parameters(parameters),
  x(sdp.primalObjective.size(), 0),
  X(sdp.psdMatrixBlockDims()),
  y(sdp.dualObjective.size(), 0),
  Y(X),
  dx(x),
  dX(X),
  dy(y),
  dY(Y),
  PrimalResidues(X),
  dualResidues(x),
  XCholesky(X),
  YCholesky(X),
  Z(X),
  R(X),
  BilinearPairingsXInv(sdp.bilinearPairingBlockDims()),
  BilinearPairingsY(BilinearPairingsXInv),
  SchurComplement(sdp.schurBlockDims()),
  SchurComplementCholesky(SchurComplement),
  SchurOffDiagonal(sdp.FreeVarMatrix),
  schurStabilizeIndices(SchurComplement.blocks.size()),
  schurStabilizeLambdas(SchurComplement.blocks.size()),
  stabilizeBlocks(SchurComplement.blocks.size()),
  Q(sdp.FreeVarMatrix.cols, sdp.FreeVarMatrix.cols),
  Qpivots(sdp.FreeVarMatrix.cols),
  dyExtended(Q.rows),
  StepMatrixWorkspace(X)
{
  // initialize bilinearPairingsWorkspace, eigenvaluesWorkspace, QRWorkspace
  for (unsigned int b = 0; b < sdp.bilinearBases.size(); b++) {
    bilinearPairingsWorkspace.push_back(Matrix(X.blocks[b].rows,
                                               BilinearPairingsXInv.blocks[b].cols));
    eigenvaluesWorkspace.push_back(Vector(X.blocks[b].rows));
    QRWorkspace.push_back(Vector(3*X.blocks[b].rows - 1));
  }

  // X = \Omega_p I
  X.addDiagonal(parameters.initialMatrixScalePrimal);
  // Y = \Omega_d I
  Y.addDiagonal(parameters.initialMatrixScaleDual);
}

/***********************************************************************/
// Specialized linear algebra helper functions needed for the solver.

// Result = Q'^T A Q', where Q' = Q \otimes 1, where \otimes denotes
// tensor product and 1 is an mxm identity matrix.
// Inputs:
// - A      : l*m x l*m symmetric matrix
// - Q      : l   x n   matrix
// Workspace:
// - Work   : l*m x n*m matrix, intermediate workspace (overwritten)
// Output:
// - Result : n*m x n*m symmetric matrix (overwritten)
//
// An explanation of the name: a 'congruence' refers to the action M
// -> A M A^T.  We use 'transpose congruence' to refer to a congruence
// with the transposed matrix M -> A^T M A.  'tensor' here refers to
// the fact that we're performing a congruence with the tensor product
// Q \otimes 1.
//
void tensorTransposeCongruence(const Matrix &A,
                               const Matrix &Q,
                               Matrix &Work,
                               Matrix &Result) {
  int m = A.rows / Q.rows;

  assert(Result.rows == Q.cols * m);
  assert(Result.cols == Q.cols * m);

  assert(Work.rows == A.rows);
  assert(Work.cols == Result.cols);

  // Work = A Q'
  for (int c = 0; c < Work.cols; c++) {
    int qCol       = c % Q.cols;
    int aColOffset = (c / Q.cols) * Q.rows;

    for (int r = 0; r < Work.rows; r++) {
      Real tmp = 0;
      for (int k = 0; k < Q.rows; k++) {
        tmp += A.elt(r, aColOffset + k) * Q.elt(k, qCol);
      }

      Work.elt(r, c) = tmp;
    }
  }

  // Result = Q'^T Work
  for (int c = 0; c < Result.cols; c++) {
    // since Result is symmetric, only compute its upper triangle
    for (int r = 0; r <= c; r++) {
      int qCol          = r % Q.cols;
      int workRowOffset = (r / Q.cols) * Q.rows;

      Real tmp = 0;
      for (int k = 0; k < Q.rows; k++) {
        tmp += Q.elt(k, qCol) * Work.elt(workRowOffset + k, c);
      }

      Result.elt(r, c) = tmp;

      // lower triangle is the same as upper triangle
      if (c != r) {
        Result.elt(c, r) = tmp;
      }
    }
  }
}

// Result_b = Q[b]'^T A_b Q[b]' for each block 0 <= b < Q.size()
// - Result_b, A_b denote the b-th blocks of Result, A, resp.
// - Q[b]' = Q[b] \otimes 1, where \otimes denotes tensor product
// - for each b, L.blocks[b], Q[b], Work[b], and Result.blocks[b] must
//   have the structure described above for
//   `tensorTransposeCongruence'
//
void blockTensorTransposeCongruence(const BlockDiagonalMatrix &A,
                                    const vector<Matrix> &Q,
                                    vector<Matrix> &Work,
                                    BlockDiagonalMatrix &Result) {
  #pragma omp parallel for schedule(dynamic)
  for (unsigned int b = 0; b < Q.size(); b++)
    tensorTransposeCongruence(A.blocks[b], Q[b], Work[b], Result.blocks[b]);
}

// Result = Q'^T A^{-1} Q', where Q' = Q \otimes 1, where \otimes
// denotes tensor product and 1 is an mxm idenity matrix.
// Inputs:
// - L      : l*m x l*m cholesky decomposition of A
// - Q      : l   x n   matrix
// Workspace:
// - Work   : l*m x n*m matrix, intermediate workspace (overwritten)
// Output:
// - Result : n*m x n*m symmetric matrix (overwritten)
//
void tensorInvTransposeCongruenceWithCholesky(const Matrix &L,
                                              const Matrix &Q,
                                              Matrix &Work,
                                              Matrix &Result) {
  // Work = L^{-1} (Q \otimes 1);
  for (int cw = 0; cw < Work.cols; cw++) {
    int mc  = cw / Q.cols;

    for (int rw = mc*Q.rows; rw < Work.rows; rw++) {
      int mr = rw / Q.rows;

      Real tmp = (mr != mc) ? Real(0) : Q.elt(rw % Q.rows, cw % Q.cols);
      for (int cl = mc*Q.rows; cl < rw; cl++)
        tmp -= L.elt(rw, cl)*Work.elt(cl, cw);

      Work.elt(rw, cw) = tmp/L.elt(rw, rw);
    }
  }

  // Result = Work^T Work
  for (int cr = 0; cr < Result.cols; cr++) {
    int mc = cr / Q.cols;

    for (int rr = 0; rr <= cr; rr++) {
      int mr = rr / Q.cols;

      Real tmp = 0;
      for (int rw = max(mr, mc)*Q.rows; rw < Work.rows; rw++)
        tmp += Work.elt(rw, cr)*Work.elt(rw, rr);

      Result.elt(rr, cr) = tmp;
      if (rr != cr)
        Result.elt(cr, rr) = tmp;
    }
  }
}

// Result_b = Q[b]'^T A_b^{-1} Q[b]' for each block 0 <= b < Q.size()
// - Result_b, A_b denote the b-th blocks of Result, A, resp.
// - Q[b]' = Q[b] \otimes 1, where \otimes denotes tensor product
// - for each b, L.blocks[b], Q[b], Work[b], and Result.blocks[b] must
//   have the structure described above for
//   `tensorInvTransposeCongruenceWithCholesky'
//
void blockTensorInvTransposeCongruenceWithCholesky(const BlockDiagonalMatrix &L,
                                                   const vector<Matrix> &Q,
                                                   vector<Matrix> &Work,
                                                   BlockDiagonalMatrix &Result) {
  #pragma omp parallel for schedule(dynamic)
  for (unsigned int b = 0; b < Q.size(); b++)
    tensorInvTransposeCongruenceWithCholesky(L.blocks[b], Q[b],
                                             Work[b], Result.blocks[b]);
}

// Result^(blockRow,blockCol) = V D V^T, where D=diag(d) is a diagonal
// matrix.
//
// Here, we view Result as a matrix on the tensor product
// R^V.rows \otimes R^k.  Result^(blockRow,blockCol) refers to the
// block submatrix of size V.rows x V.rows at position (blockRow,
// blockCol) in the second tensor factor.
//
// Inputs:
// - d        : pointer to beginning of a length-V.cols vector
// - V        : V.rows x V.cols Matrix
// - blockRow : integer < k
// - blockCol : integer < k
// Output:
// - Result   : (k*V.rows) x (k*V.rows) square Matrix (overwritten)
//
void diagonalCongruence(Real const *d,
                        const Matrix &V,
                        const int blockRow,
                        const int blockCol,
                        Matrix &Result) {
  for (int p = 0; p < V.rows; p++) {
    for (int q = 0; q <= p; q++) {
      Real tmp = 0;

      for (int n = 0; n < V.cols; n++)
        tmp += *(d+n) * V.elt(p, n)*V.elt(q, n);

      Result.elt(blockRow*V.rows + p, blockCol*V.rows + q) = tmp;
      if (p != q)
        Result.elt(blockRow*V.rows + q, blockCol*V.rows + p) = tmp;
    }
  }
}

// v^T A^(blockRow, blockCol) v, where A^(r,s) is the (r,s)-th dim x
// dim block inside A.
//
// Input:
// - v        : pointer to the beginning of a vector of length dim
// - dim      : length of the vector v
// - A        : (k*dim) x (k*dim) matrix, where k > blockRow, blockCol
// - blockRow : integer labeling block of A
// - blockCol : integer labeling block of A
// Output: v^T A^(blockRow, blockCol) v (returned)
//
Real bilinearBlockPairing(const Real *v,
                          const int dim,
                          const Matrix &A,
                          const int blockRow,
                          const int blockCol) {
  Real result = 0;

  for (int r = 0; r < dim; r++) {
    Real tmp = 0;

    for (int c = 0; c < dim; c++)
      tmp += *(v+c) * A.elt(blockRow*dim + r, blockCol*dim + c);
    result += *(v+r) * tmp;
  }
  return result;
}

/***********************************************************************/
// Subroutines needed for each solver iteration

// Compute the SchurComplement matrix using BilinearPairingsXInv and
// BilinearPairingsY and the formula
//
//   S_{(j,r1,s1,k1), (j,r2,s2,k2)} = \sum_{b \in blocks[j]}
//          (1/4) (BilinearPairingsXInv_{ej s1 + k1, ej r2 + k2}*
//                 BilinearPairingsY_{ej s2 + k2, ej r1 + k1} +
//                 swaps (r1 <-> s1) and (r2 <-> s2))
//
// where ej = d_j + 1.
//
// Inputs: sdp, BilinearPairingsXInv, BilinearPairingsY
// Output: SchurComplement (overwritten)
//
void computeSchurComplement(const SDP &sdp,
                            const BlockDiagonalMatrix &BilinearPairingsXInv,
                            const BlockDiagonalMatrix &BilinearPairingsY,
                            BlockDiagonalMatrix &SchurComplement) {
  #pragma omp parallel for schedule(dynamic)
  for (unsigned int j = 0; j < sdp.dimensions.size(); j++) {
    const int ej = sdp.degrees[j] + 1;

    for (unsigned int u1 = 0; u1 < sdp.constraintIndices[j].size(); u1++) {
      const int ej_r1 = sdp.constraintIndices[j][u1].r * ej;
      const int ej_s1 = sdp.constraintIndices[j][u1].s * ej;
      const int k1    = sdp.constraintIndices[j][u1].k;

      for (unsigned int u2 = 0; u2 <= u1; u2++) {
        const int ej_r2 = sdp.constraintIndices[j][u2].r * ej;
        const int ej_s2 = sdp.constraintIndices[j][u2].s * ej;
        const int k2    = sdp.constraintIndices[j][u2].k;

        Real tmp = 0;
        for (vector<int>::const_iterator b = sdp.blocks[j].begin();
             b != sdp.blocks[j].end(); b++) {
          tmp += (BilinearPairingsXInv.blocks[*b].elt(ej_s1 + k1, ej_r2 + k2) *
                  BilinearPairingsY   .blocks[*b].elt(ej_s2 + k2, ej_r1 + k1) +
                  BilinearPairingsXInv.blocks[*b].elt(ej_r1 + k1, ej_r2 + k2) *
                  BilinearPairingsY   .blocks[*b].elt(ej_s2 + k2, ej_s1 + k1) +
                  BilinearPairingsXInv.blocks[*b].elt(ej_s1 + k1, ej_s2 + k2) *
                  BilinearPairingsY   .blocks[*b].elt(ej_r2 + k2, ej_r1 + k1) +
                  BilinearPairingsXInv.blocks[*b].elt(ej_r1 + k1, ej_s2 + k2) *
                  BilinearPairingsY   .blocks[*b].elt(ej_r2 + k2, ej_s1 + k1))/4;
        }
        SchurComplement.blocks[j].elt(u1, u2) = tmp;
        if (u2 != u1)
          SchurComplement.blocks[j].elt(u2, u1) = tmp;
      }
    }
  }
}

// dualResidues[p] = primalObjective[p] - Tr(A_p Y) - (FreeVarMatrix y)_p,
// for 0 <= p < primalObjective.size()
//
// The pairings Tr(A_p Y) can be written in terms of BilinearPairingsY:
//
//   Tr(A_(j,r,s,k) Y) = \sum_{b \in blocks[j]}
//                       (1/2) (BilinearPairingsY_{ej r + k, ej s + k} +
//                              swap (r <-> s))
// where ej = d_j + 1.
//
// Inputs: sdp, y, BilinearPairingsY
// Output: dualResidues (overwriten)
//
void computeDualResidues(const SDP &sdp,
                         const Vector &y,
                         const BlockDiagonalMatrix &BilinearPairingsY,
                         Vector &dualResidues) {
  #pragma omp parallel for schedule(dynamic)
  for (unsigned int j = 0; j < sdp.dimensions.size(); j++) {
    const int ej = sdp.degrees[j] + 1;

    for (vector<IndexTuple>::const_iterator t = sdp.constraintIndices[j].begin();
         t != sdp.constraintIndices[j].end();
         t++) {
      const int p    = t->p;
      const int ej_r = t->r * ej;
      const int ej_s = t->s * ej;
      const int k    = t->k;

      // dualResidues[p] = -Tr(A_p Y)
      dualResidues[p] = 0;
      for (vector<int>::const_iterator b = sdp.blocks[j].begin();
           b != sdp.blocks[j].end(); b++) {
        dualResidues[p] -= BilinearPairingsY.blocks[*b].elt(ej_r+k, ej_s+k);
        dualResidues[p] -= BilinearPairingsY.blocks[*b].elt(ej_s+k, ej_r+k);
      }
      dualResidues[p] /= 2;

      // dualResidues[p] = -Tr(A_p Y) - (FreeVarMatrix y)_p
      for (int n = 0; n < sdp.FreeVarMatrix.cols; n++)
        dualResidues[p] -= sdp.FreeVarMatrix.elt(p, n)*y[n];

      // dualResidues[p] = primalObjective[p] - Tr(A_p Y) - (FreeVarMatrix y)_p
      dualResidues[p] += sdp.primalObjective[p];
    }
  }
}

// Result = \sum_p a[p] A_p,
//
// where a[p] is a vector of length primalObjective.size() and the
// constraint matrices A_p are given by
//
//   A_(j,r,s,k) = \sum_{b \in blocks[j]}
//                     Block_b(v_{b,k} v_{b,k}^T \otimes E^{rs}),
//
// where v_{b,k} is the k-th column of bilinearBases[b], as described
// in SDP.h.
//
// Inputs: sdp, a
// Output: Result (overwritten)
//
void constraintMatrixWeightedSum(const SDP &sdp,
                                 const Vector a,
                                 BlockDiagonalMatrix &Result)  {
  #pragma omp parallel for schedule(dynamic)
  for (unsigned int j = 0; j < sdp.dimensions.size(); j++) {
    const int ej = sdp.degrees[j] + 1;

    // For each j, t points to the first IndexTuple corresponding to j
    for (vector<IndexTuple>::const_iterator t = sdp.constraintIndices[j].begin();
         t != sdp.constraintIndices[j].end();
         t += ej) {
      const int p = t->p;
      const int r = t->r;
      const int s = t->s;
      assert(t->k == 0);

      for (vector<int>::const_iterator b = sdp.blocks[j].begin();
           b != sdp.blocks[j].end(); b++) {

        // Result.blocks[b]^(r,s) = V diag(a') V^T, where
        // V=sdp.bilinearBases[b], a' denotes the subvector of a
        // corresponding to j, and M^(r,s) denotes the (r,s)-th block
        // of M.
        diagonalCongruence(&a[p], sdp.bilinearBases[*b], r, s, Result.blocks[*b]);

        // Result should be symmetric, so if r != s, we must divide
        // the (r,s)-th block of Result.blocks[b] by 2 and copy its
        // transpose to the (s,r)-th block.
        if (r != s) {
          const int u = sdp.bilinearBases[*b].rows;
          for (int m = r*u; m < (r+1)*u; m++) {
            for (int n = s*u; n < (s+1)*u; n++) {
              Result.blocks[*b].elt(m, n) /= 2;
              Result.blocks[*b].elt(n, m) = Result.blocks[*b].elt(m, n);
            }
          }
        }
      }
    }
  }
}

// Compute the vectors r_x and r_y on the right-hand side of the Schur
// complement equation:
//
// {{S, -B}, {B^T, 0}} . {dx, dy} = {r_x, r_y}
//
// where S = SchurComplement and B = FreeVarMatrix.  Specifically,
//
// r_x[p] = -dualResidues[p] - Tr(A_p Z)              for 0 <= p < P
// r_y[n] = dualObjective[n] - (FreeVarMatrix^T x)_n  for 0 <= n < N
//
// where P = primalObjective.size(), N = dualObjective.size()
//
// Inputs:
// - sdp, an SDP
// - dualResidues, a Vector of length P
// - Z = X^{-1} (PrimalResidues Y - R)
// - x, a vector of length P
// Outputs:
// - r_x, a Vector of length P
// - r_y, a Vector of length N
//
void computeSchurRHS(const SDP &sdp,
                     const Vector &dualResidues,
                     const BlockDiagonalMatrix &Z,
                     const Vector &x,
                     Vector &r_x,
                     Vector &r_y) {
  // r_x[p] = -dualResidues[p]
  for (unsigned int p = 0; p < r_x.size(); p++)
    r_x[p] = -dualResidues[p];

  // r_x[p] = -dualResidues[p] - Tr(A_p Z), where A_p are as described
  // in SDP.h.  The trace can be computed in terms of bilinearBases
  // using bilinearBlockPairing.
  #pragma omp parallel for schedule(dynamic)
  for (unsigned int j = 0; j < sdp.dimensions.size(); j++) {
    for (vector<IndexTuple>::const_iterator t = sdp.constraintIndices[j].begin();
         t != sdp.constraintIndices[j].end();
         t++) {
      for (vector<int>::const_iterator b = sdp.blocks[j].begin();
           b != sdp.blocks[j].end(); b++) {
        const int h = sdp.bilinearBases[*b].rows;
        // Pointer to the k-th column of sdp.bilinearBases[*b]
        const Real *q = &sdp.bilinearBases[*b].elements[(t->k) * h];

        r_x[t->p] -= bilinearBlockPairing(q, h, Z.blocks[*b], t->r, t->s);
      }
    }
  }

  // r_y[n] = dualObjective[n] - (FreeVarMatrix^T x)_n
  #pragma omp parallel for schedule(static)
  for (unsigned int n = 0; n < sdp.dualObjective.size(); n++) {
    r_y[n] = sdp.dualObjective[n];
    for (int p = 0; p < sdp.FreeVarMatrix.rows; p++) {
      r_y[n] -= sdp.FreeVarMatrix.elt(p, n)*x[p];
    }
  }
}

// PrimalResidues = \sum_p A_p x[p] - X
//
// Inputs: sdp, x, X
// Output: PrimalResidues (overwritten)
//
void computePrimalResidues(const SDP &sdp,
                           const Vector x,
                           const BlockDiagonalMatrix &X,
                           BlockDiagonalMatrix &PrimalResidues) {
  constraintMatrixWeightedSum(sdp, x, PrimalResidues);
  PrimalResidues -= X;
}

// Centering parameter \beta_p for the predictor step
Real predictorCenteringParameter(const SDPSolverParameters &parameters,
                                 const bool isPrimalDualFeasible) {
  return isPrimalDualFeasible ? Real(0) : parameters.infeasibleCenteringParameter;
}

// Centering parameter \beta_c for the corrector step
Real correctorCenteringParameter(const SDPSolverParameters &parameters,
                                 const BlockDiagonalMatrix &X,
                                 const BlockDiagonalMatrix &dX,
                                 const BlockDiagonalMatrix &Y,
                                 const BlockDiagonalMatrix &dY,
                                 const Real &mu,
                                 const bool isPrimalDualFeasible) {
  Real r = frobeniusProductOfSums(X, dX, Y, dY) / (mu * X.dim);
  Real beta = r < 1 ? r*r : r;

  if (isPrimalDualFeasible)
    return min(max(parameters.feasibleCenteringParameter, beta), Real(1));
  else
    return max(parameters.infeasibleCenteringParameter, beta);
}

// min(gamma \alpha(M, dM), 1), where \alpha(M, dM) denotes the
// largest positive real number such that M + \alpha dM is positive
// semidefinite.
//
// \alpha(M, dM) is computed with a Cholesky decomposition M = L L^T.
// The eigenvalues of M + \alpha dM are equal to the eigenvalues of 1
// + \alpha L^{-1} dM L^{-T}.  The correct \alpha is then -1/lambda,
// where lambda is the smallest eigenvalue of L^{-1} dM L^{-T}.
//
// Inputs:
// - MCholesky = L, the Cholesky decomposition of M (M itself is not needed)
// - dM, a BlockDiagonalMatrix with the same structure as M
// Workspace:
// - MInvDM (NB: overwritten when computing minEigenvalue)
// - eigenvalues, a Vector of eigenvalues for each block of M
// - workspace, a vector of Vectors needed by the minEigenvalue function
// Output:
// - min(\gamma \alpha(M, dM), 1) (returned)
//
Real stepLength(BlockDiagonalMatrix &MCholesky,
                BlockDiagonalMatrix &dM,
                BlockDiagonalMatrix &MInvDM,
                vector<Vector> &eigenvalues,
                vector<Vector> &workspace,
                const Real gamma) {
  // MInvDM = L^{-1} dM L^{-T}, where M = L L^T
  MInvDM.copyFrom(dM);
  lowerTriangularInverseCongruence(MInvDM, MCholesky);

  const Real lambda = minEigenvalue(MInvDM, workspace, eigenvalues);
  if (lambda > -gamma)
    return 1;
  else
    return -gamma/lambda;
}

// Compute the quantities needed to solve the Schur complement
// equation
//
// {{S, -B}, {B^T, 0}} . {dx, dy} = {r, s}
//
// (where S = SchurComplement, B = FreeVarMatrix), using the method
// described in the manual:
//
// - Compute S using BilinearPairingsXInv and BilinearPairingsY.
//
// - Stabilize S by adding a low-rank update S' = S + U U^T and
//   compute the Cholesky decomposition S' = L' L'^T.
//
// - Form B' = (B U) and compute
//
//   - SchurOffDiagonal = L'^{-1} B
//   - L'^{-1} U (this is stored implicitly in the stabilizeBlock* variables)
//   - Q = (L'^{-1} B')^T (L'^{-1} B') - {{0, 0}, {0, 1}}
//
// - Compute the LU decomposition of Q.
//
// This data is sufficient to efficiently solve the above equation for
// a given r,s.
//
// Inputs:
// - BilinearPairingsXInv, BilinearPairingsY (these are members of
//   SDPSolver, but we include them as arguments to emphasize that
//   they must be computed first)
// - choleskyStabilizeThreshold: the real constant \theta used in
//   stabilizing S
// Workspace (members of SDPSolver which are modified by this method
// and not used later):
// - SchurComplement
// - schurStabilizeIndices
// - schurStabilizeLambdas
// Outputs (members of SDPSolver which are modified by this method and
// used later):
// - SchurComplementCholesky
// - SchurOffDiagonal
// - stabilizeBlockIndices
// - stabilizeBlockUpdateRow
// - stabilizeBlockUpdateColumn
// - stabilizeBlocks
// - Q, Qpivots
//
void SDPSolver::initializeSchurComplementSolver(const BlockDiagonalMatrix &BilinearPairingsXInv,
                                                const BlockDiagonalMatrix &BilinearPairingsY,
                                                const Real &choleskyStabilizeThreshold) {
  computeSchurComplement(sdp, BilinearPairingsXInv, BilinearPairingsY, SchurComplement);

  // compute SchurComplementCholesky = L', where
  //
  //   L' L'^T = S' = SchurComplement + U U^T
  //
  // Here, the 'update' matrix U has columns given by
  //
  //   U = ( Lambda_{p_1} e_{p_1}, ..., Lambda_{p_M} e_{p_M} )
  //
  // where e_p is a unit vector in the p-th direction and the
  // Lambda_{p_m} are constants. The p_i are given by
  // schurStabilizeIndices and the corresponding Lambda_i by
  // schurStabilizeLambdas.
  //
  choleskyDecompositionStabilized(SchurComplement, SchurComplementCholesky,
                                  schurStabilizeIndices,
                                  schurStabilizeLambdas,
                                  cast2double(choleskyStabilizeThreshold));

  // SchurOffDiagonal = L'^{-1} FreeVarMatrix
  SchurOffDiagonal.copyFrom(sdp.FreeVarMatrix);
  blockMatrixLowerTriangularSolve(SchurComplementCholesky, SchurOffDiagonal);

  // Next we compute L'^{-1} U, which is stored implicitly in terms of
  // its nonzero submatrices in the stabilizeBlock* variables.

  // total number of columns in the off-diagonal part B' = (B U)
  // (currently just B; will accumulate the rest shortly)
  int offDiagonalColumns = SchurOffDiagonal.cols;

  stabilizeBlockIndices.resize(0);
  stabilizeBlockUpdateRow.resize(0);
  stabilizeBlockUpdateColumn.resize(0);

  // j runs over blocks of SchurComplement
  for (unsigned int j = 0; j < SchurComplement.blocks.size(); j++) {
    if (schurStabilizeIndices[j].size() > 0) {
      // the j-th block of S contains stabilized directions. We have a
      // block submatrix
      //
      //   U_j = stabilizeBlocks[j]
      //
      // of U for each such j.
      stabilizeBlockIndices.push_back(j);

      // index of the first stabilized direction within the j-th block
      int startIndex = schurStabilizeIndices[j][0];

      // set the dimensions of U_j
      stabilizeBlocks[j].resize(SchurComplement.blocks[j].rows - startIndex,
                                schurStabilizeIndices[j].size());
      // set U_j = 0
      stabilizeBlocks[j].setZero();
      // for each column of U_j add Lambda_p in the row (p - startIndex)
      for (unsigned int c = 0; c < schurStabilizeIndices[j].size(); c++) {
        int r = schurStabilizeIndices[j][c] - startIndex;
        stabilizeBlocks[j].elt(r, c) = schurStabilizeLambdas[j][c];
      }

      // append the row of U corresponding to the top-left of U_j
      stabilizeBlockUpdateRow.push_back(SchurComplement.blockStartIndices[j] + startIndex);
      // append the column of U corresponding to the top-left of U_j
      stabilizeBlockUpdateColumn.push_back(offDiagonalColumns);
      // update the number of off-diagonal columns
      offDiagonalColumns += stabilizeBlocks[j].cols;
    }
  }

  // Set U = L'^{-1} U
  //
  // We do this by modifying the blocks U_j = stabilizeBlocks[j]
  // in-place, multiplying by the inverse of the appropriate submatrix
  // of SchurComplementCholesky.  We henceforth refer to L'^{-1} U as
  // V to avoid confusion.
  #pragma omp parallel for schedule(dynamic)
  for (unsigned int j = 0; j < stabilizeBlockIndices.size(); j++) {
    int b = stabilizeBlockIndices[j];
    int startIndex = schurStabilizeIndices[b][0];
    Rtrsm("Left", "Lower", "NoTranspose", "NonUnitDiagonal",
          stabilizeBlocks[b].rows, stabilizeBlocks[b].cols, 1,
          &SchurComplementCholesky.blocks[b].elt(startIndex, startIndex),
          SchurComplementCholesky.blocks[b].rows,
          &stabilizeBlocks[b].elt(0, 0),
          stabilizeBlocks[b].rows);
  }

  // Next, we compute
  //
  //   Q = (L'^{-1} B')^T (L'^{-1} B') - {{0, 0}, {0, 1}}
  //
  // Where B' = (B U).  We think of Q as containing four blocks called
  // Upper/Lower-Left/Right.

  // Set the dimensions of Q
  Q.resize(offDiagonalColumns, offDiagonalColumns);
  Q.setZero();

  // Here, SchurOffDiagonal = L'^{-1} B.
  //
  // UpperLeft(Q) = SchurOffDiagonal^T SchurOffDiagonal
  matrixSquareIntoBlock(SchurOffDiagonal, Q, 0, 0);

  // Here, stabilizeBlocks contains the blocks of V = L'^{-1} U.
  //
  // LowerRight(Q) = V^T V - 1
  for (unsigned int j = 0; j < stabilizeBlockIndices.size(); j++) {
    int b = stabilizeBlockIndices[j];
    int c = stabilizeBlockUpdateColumn[j];
    matrixSquareIntoBlock(stabilizeBlocks[b], Q, c, c);
    // subtract the identity matrix from this block
    for (int i = c; i < c + stabilizeBlocks[b].cols; i++)
      Q.elt(i, i) -= 1;
  }

  // LowerLeft(Q) = V^T SchurOffDiagonal
  # pragma omp parallel for schedule(dynamic)
  for (unsigned int j = 0; j < stabilizeBlockIndices.size(); j++) {
    int b = stabilizeBlockIndices[j];
    int p = stabilizeBlockUpdateRow[j];
    int r = stabilizeBlockUpdateColumn[j];
    Rgemm("Transpose", "NoTranspose",
          stabilizeBlocks[b].cols,
          SchurOffDiagonal.cols,
          stabilizeBlocks[b].rows,
          1,
          &stabilizeBlocks[b].elements[0],
          stabilizeBlocks[b].rows,
          &SchurOffDiagonal.elt(p, 0),
          SchurOffDiagonal.rows,
          0,
          &Q.elt(r, 0),
          Q.rows);
  }

  // UpperRight(Q) = LowerLeft(Q)^T
  # pragma omp parallel for schedule(static)
  for (int c = 0; c < SchurOffDiagonal.cols; c++)
    for (int r = SchurOffDiagonal.cols; r < Q.rows; r++)
      Q.elt(c, r) = Q.elt(r, c);

  // Resize Qpivots appropriately and compute the LU decomposition of Q
  Qpivots.resize(Q.rows);
  LUDecomposition(Q, Qpivots);
}

// Solve the Schur complement equation for dx, dy.
//
// - As inputs, dx and dy are the residues r_x and r_y on the
//   right-hand side of the Schur complement equation.
// - As outputs, dx and dy are overwritten with the solutions of the
//   Schur complement equation.
//
// The equation is solved using the block-decomposition described in
// the manual.
//
void SDPSolver::solveSchurComplementEquation(Vector &dx, Vector &dy) {
  // dx = SchurComplementCholesky^{-1} dx
  blockMatrixLowerTriangularSolve(SchurComplementCholesky, dx);

  // extend dy by additional coordinates z needed for stabilizing
  // dyExtended = (dy, z)
  dyExtended.resize(Q.rows);

  // dy = r_y - SchurOffDiagonal^T dx
  for (unsigned int n = 0; n < dy.size(); n++)
    dyExtended[n] = dy[n];
  vectorScaleMatrixMultiplyTransposeAdd(-1, SchurOffDiagonal, dx, 1, dyExtended);

  // z = -V^T dx
  for (unsigned int j = 0; j < stabilizeBlockIndices.size(); j++) {
    int b = stabilizeBlockIndices[j];
    int pTopLeft = stabilizeBlockUpdateRow[j];
    int cTopLeft = stabilizeBlockUpdateColumn[j];

    for (int c = 0; c < stabilizeBlocks[b].cols; c++) {
      dyExtended[cTopLeft + c] = 0;
      for (int r = 0; r < stabilizeBlocks[b].rows; r++)
        dyExtended[cTopLeft + c] -= dx[pTopLeft + r] * stabilizeBlocks[b].elt(r, c);
    }
  }

  // dyExtended = Q^{-1} dyExtended
  solveWithLUDecomposition(Q, Qpivots, dyExtended);

  // dx += SchurOffDiagonal dy
  vectorScaleMatrixMultiplyAdd(1, SchurOffDiagonal, dyExtended, 1, dx);

  // dx += V z
  for (unsigned int j = 0; j < stabilizeBlockIndices.size(); j++) {
    int b = stabilizeBlockIndices[j];
    int pTopLeft = stabilizeBlockUpdateRow[j];
    int cTopLeft = stabilizeBlockUpdateColumn[j];

    for (int c = 0; c < stabilizeBlocks[b].cols; c++)
      for (int r = 0; r < stabilizeBlocks[b].rows; r++)
        dx[pTopLeft + r] += dyExtended[cTopLeft + c] * stabilizeBlocks[b].elt(r, c);
  }

  // dx = SchurComplementCholesky^{-T} dx
  blockMatrixLowerTriangularTransposeSolve(SchurComplementCholesky, dx);

  // dy = first few entries of dyExtended
  for (unsigned int n = 0; n < dy.size(); n++)
    dy[n] = dyExtended[n];
}

// Compute the search direction (dx, dX, dy, dY) for the predictor and
// corrector phases.
//
// Inputs:
// - beta, the centering parameter
// - mu = Tr(X Y) / X.cols
// - correctorPhase: boolean indicating whether we're in the corrector
//   phase or predictor phase.
// Workspace (members of SDPSolver which are modified in-place but not
// used elsewhere):
// - Z, R
// Outputs (members of SDPSolver which are modified in-place):
// - dx, dX, dy, dY
//
void SDPSolver::computeSearchDirection(const Real &beta,
                                       const Real &mu,
                                       const bool correctorPhase) {
  // R = beta mu I - X Y (predictor phase)
  // R = beta mu I - X Y - dX dY (corrector phase)
  blockDiagonalMatrixScaleMultiplyAdd(-1, X, Y, 0, R);
  if (correctorPhase)
    blockDiagonalMatrixScaleMultiplyAdd(-1, dX, dY, 1, R);
  R.addDiagonal(beta*mu);

  // Z = Symmetrize(X^{-1} (PrimalResidues Y - R))
  blockDiagonalMatrixMultiply(PrimalResidues, Y, Z);
  Z -= R;
  blockMatrixSolveWithCholesky(XCholesky, Z);
  Z.symmetrize();

  // r_x[p] = -dualResidues[p] - Tr(A_p Z)
  // r_y[n] = dualObjective[n] - (FreeVarMatrix^T x)_n
  // Here, dx = r_x, dy = r_y.
  computeSchurRHS(sdp, dualResidues, Z, x, dx, dy);

  // Solve for dx, dy in-place
  solveSchurComplementEquation(dx, dy);

  // dX = PrimalResidues + \sum_p A_p dx[p]
  constraintMatrixWeightedSum(sdp, dx, dX);
  dX += PrimalResidues;

  // dY = Symmetrize(X^{-1} (R - dX Y))
  blockDiagonalMatrixMultiply(dX, Y, dY);
  dY -= R;
  blockMatrixSolveWithCholesky(XCholesky, dY);
  dY.symmetrize();
  dY *= -1;
}

/***********************************************************************/
// The main solver loop

SDPSolverTerminateReason SDPSolver::run(const path checkpointFile) {
  Real primalStepLength;
  Real dualStepLength;

  printHeader();

  for (int iteration = 1;; iteration++) {
    if (timers["Last checkpoint"].elapsed().wall >= parameters.checkpointInterval * 1000000000LL)
      saveCheckpoint(checkpointFile);
    if (timers["Solver runtime"].elapsed().wall >= parameters.maxRuntime * 1000000000LL)
      return MaxRuntimeExceeded;

    primalObjective = sdp.objectiveConst + dotProduct(sdp.primalObjective, x);
    dualObjective   = sdp.objectiveConst + dotProduct(sdp.dualObjective, y);
    dualityGap      = abs(primalObjective - dualObjective) /
      max(Real(abs(primalObjective) + abs(dualObjective)), Real(1));

    choleskyDecomposition(X, XCholesky);
    choleskyDecomposition(Y, YCholesky);

    // Compute the bilinear pairings BilinearPairingsXInv and
    // BilinearPairingsY needed for the dualResidues and the Schur
    // complement matrix
    blockTensorInvTransposeCongruenceWithCholesky(XCholesky, sdp.bilinearBases,
                                                  bilinearPairingsWorkspace,
                                                  BilinearPairingsXInv);
    blockTensorTransposeCongruence(Y, sdp.bilinearBases,
                                   bilinearPairingsWorkspace,
                                   BilinearPairingsY);

    // dualResidues[p] = primalObjective[p] - Tr(A_p Y) - (FreeVarMatrix y)_p,
    computeDualResidues(sdp, y, BilinearPairingsY, dualResidues);
    dualError = maxAbsVector(dualResidues);

    // PrimalResidues = \sum_p A_p x[p] - X
    computePrimalResidues(sdp, x, X, PrimalResidues);
    primalError = PrimalResidues.maxAbs();

    const bool isPrimalFeasible = primalError < parameters.primalErrorThreshold;
    const bool isDualFeasible   = dualError   < parameters.dualErrorThreshold;
    const bool isOptimal        = dualityGap  < parameters.dualityGapThreshold;

    if (isPrimalFeasible && isDualFeasible && isOptimal)                   return PrimalDualOptimal;
    else if (isPrimalFeasible && parameters.findPrimalFeasible)            return PrimalFeasible;
    else if (isDualFeasible && parameters.findDualFeasible)                return DualFeasible;
    else if (primalStepLength == 1 && parameters.detectPrimalFeasibleJump) return PrimalFeasibleJumpDetected;
    else if (dualStepLength == 1 && parameters.detectDualFeasibleJump)     return DualFeasibleJumpDetected;
    else if (iteration > parameters.maxIterations)                         return MaxIterationsExceeded;

    // Compute SchurComplement and prepare to solve the Schur
    // complement equation for dx, dy
    initializeSchurComplementSolver(BilinearPairingsXInv, BilinearPairingsY,
                                    parameters.choleskyStabilizeThreshold);

    // Compute the complementarity mu = Tr(X Y)/X.dim
    Real mu = frobeniusProductSymmetric(X, Y)/X.dim;
    if (mu > parameters.maxComplementarity)
      return MaxComplementarityExceeded;

    // Compute the predictor solution for (dx, dX, dy, dY)
    Real betaPredictor =
      predictorCenteringParameter(parameters, isPrimalFeasible && isDualFeasible);
    computeSearchDirection(betaPredictor, mu, false);

    // Compute the corrector solution for (dx, dX, dy, dY)
    Real betaCorrector =
      correctorCenteringParameter(parameters, X, dX, Y, dY, mu,
                                  isPrimalFeasible && isDualFeasible);
    computeSearchDirection(betaCorrector, mu, true);

    // Compute step-lengths that preserve positive definiteness of X, Y
    primalStepLength = stepLength(XCholesky, dX, StepMatrixWorkspace,
                                  eigenvaluesWorkspace, QRWorkspace,
                                  parameters.stepLengthReduction);
    dualStepLength   = stepLength(YCholesky, dY, StepMatrixWorkspace,
                                  eigenvaluesWorkspace, QRWorkspace,
                                  parameters.stepLengthReduction);

    // If our problem is both dual-feasible and primal-feasible,
    // ensure we're following the true Newton direction.
    if (isPrimalFeasible && isDualFeasible) {
      primalStepLength = min(primalStepLength, dualStepLength);
      dualStepLength = primalStepLength;
    }

    printIteration(iteration, mu, primalStepLength, dualStepLength, betaCorrector);

    // Update the primal point (x, X) += primalStepLength*(dx, dX)
    addScaledVector(x, primalStepLength, dx);
    dX *= primalStepLength;
    X += dX;

    // Update the dual point (y, Y) += dualStepLength*(dy, dY)
    addScaledVector(y, dualStepLength, dy);
    dY *= dualStepLength;
    Y += dY;
  }

  // Never reached
  return MaxIterationsExceeded;
}