#include "general_quantum_operator.hpp"

#include <Eigen/Dense>
#include <csim/stat_ops.hpp>
#include <csim/update_ops.hpp>
#include <csim/update_ops_dm.hpp>
#include <cstring>
#include <fstream>
#include <numeric>

#include "gate_factory.hpp"
#include "pauli_operator.hpp"
#include "state.hpp"
#include "type.hpp"
#include "utility.hpp"

#ifdef _USE_GPU
#include <gpusim/update_ops_cuda.h>
#endif

GeneralQuantumOperator::GeneralQuantumOperator(const UINT qubit_count)
    : _qubit_count(qubit_count), _is_hermitian(true) {}

GeneralQuantumOperator::~GeneralQuantumOperator() {
    for (auto& term : this->_operator_list) {
        delete term;
    }
}

void GeneralQuantumOperator::add_operator(const PauliOperator* mpt) {
    PauliOperator* _mpt = mpt->copy();
    if (!check_Pauli_operator(this, _mpt)) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: GeneralQuantumOperator::add_operator(const "
               "PauliOperator*): pauli_operator applies target qubit of "
               "which the index is larger than qubit_count";
        throw std::invalid_argument(error_message_stream.str());
    }
    if (this->_is_hermitian && std::abs(_mpt->get_coef().imag()) > 0) {
        this->_is_hermitian = false;
    }
    this->_operator_list.push_back(_mpt);
}

void GeneralQuantumOperator::add_operator(
    CPPCTYPE coef, std::string pauli_string) {
    PauliOperator* _mpt = new PauliOperator(pauli_string, coef);
    if (!check_Pauli_operator(this, _mpt)) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: "
               "GeneralQuantumOperator::add_operator(double,std::string):"
               " pauli_operator applies target qubit of which the index "
               "is larger than qubit_count";
        throw std::invalid_argument(error_message_stream.str());
    }
    if (this->_is_hermitian && std::abs(coef.imag()) > 0) {
        this->_is_hermitian = false;
    }
    this->add_operator(_mpt);
    delete _mpt;
}

CPPCTYPE GeneralQuantumOperator::get_expectation_value(
    const QuantumStateBase* state) const {
    if (this->_qubit_count > state->qubit_count) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: GeneralQuantumOperator::get_expectation_value(const "
               "QuantumStateBase*): invalid qubit count";
        throw std::invalid_argument(error_message_stream.str());
    }
    auto sum = std::accumulate(this->_operator_list.cbegin(),
        this->_operator_list.cend(), (CPPCTYPE)0.0,
        [&](CPPCTYPE acc, PauliOperator* pauli) {
            return acc + pauli->get_expectation_value(state);
        });
    return sum;
}

CPPCTYPE GeneralQuantumOperator::get_expectation_value_single_thread(
    const QuantumStateBase* state) const {
    if (this->_qubit_count > state->qubit_count) {
        std::cerr
            << "Error: GeneralQuantumOperator::get_expectation_value(const "
               "QuantumStateBase*): invalid qubit count"
            << std::endl;
        return 0.;
    }
    auto sum = std::accumulate(this->_operator_list.cbegin(),
        this->_operator_list.cend(), (CPPCTYPE)0.0,
        [&](CPPCTYPE acc, PauliOperator* pauli) {
            return acc + pauli->get_expectation_value_single_thread(state);
        });
    return sum;
}

CPPCTYPE GeneralQuantumOperator::get_transition_amplitude(
    const QuantumStateBase* state_bra,
    const QuantumStateBase* state_ket) const {
    if (this->_qubit_count > state_bra->qubit_count ||
        state_bra->qubit_count != state_ket->qubit_count) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: GeneralQuantumOperator::get_transition_amplitude(const "
               "QuantumStateBase*, const QuantumStateBase*): invalid qubit "
               "count";
        throw std::invalid_argument(error_message_stream.str());
    }

    auto sum = std::accumulate(this->_operator_list.cbegin(),
        this->_operator_list.cend(), (CPPCTYPE)0.0,
        [&](CPPCTYPE acc, PauliOperator* pauli) {
            return acc + pauli->get_transition_amplitude(state_bra, state_ket);
        });
    return sum;
}

void GeneralQuantumOperator::add_random_operator(const UINT operator_count) {
    const auto qubit_count = this->get_qubit_count();
    Random random;
    for (UINT operator_index = 0; operator_index < operator_count;
         operator_index++) {
        auto target_qubit_index_list = std::vector<UINT>(qubit_count, 0);
        auto target_qubit_pauli_list = std::vector<UINT>(qubit_count, 0);
        for (UINT qubit_index = 0; qubit_index < qubit_count; qubit_index++) {
            const UINT pauli_id = random.int32() % 4;
            target_qubit_index_list.at(qubit_index) = qubit_index;
            target_qubit_pauli_list.at(qubit_index) = pauli_id;
        }
        // -1.0 <= coef <= 1.0
        const CPPCTYPE coef = random.uniform() * 2 - 1.0;
        auto pauli_operator = PauliOperator(
            target_qubit_index_list, target_qubit_pauli_list, coef);
        this->add_operator(&pauli_operator);
    }
}

CPPCTYPE
GeneralQuantumOperator::solve_ground_state_eigenvalue_by_arnoldi_method(
    QuantumStateBase* state, const UINT iter_count, const CPPCTYPE mu) const {
    if (this->get_term_count() == 0) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: "
               "GeneralQuantumOperator::solve_ground_state_eigenvalue_by_"
               "arnoldi_method("
               "QuantumStateBase * state, const UINT iter_count, const "
               "CPPCTYPE mu): At least one PauliOperator is required.";
        throw std::invalid_argument(error_message_stream.str());
    }

    // Implemented based on
    // https://files.transtutors.com/cdn/uploadassignments/472339_1_-numerical-linear-aljebra.pdf
    const auto qubit_count = this->get_qubit_count();
    auto present_state = QuantumState(qubit_count);
    auto tmp_state = QuantumState(qubit_count);
    auto multiplied_state = QuantumState(qubit_count);
    auto mu_timed_state = QuantumState(qubit_count);

    // Vectors composing Krylov subspace.
    std::vector<QuantumStateBase*> state_list;
    state_list.reserve(iter_count + 1);
    state->normalize(state->get_squared_norm());
    state_list.push_back(state->copy());

    CPPCTYPE mu_;
    if (mu == 0.0) {
        // mu is not changed from default value.
        mu_ = this->calculate_default_mu();
    } else {
        mu_ = mu;
    }

    ComplexMatrix hessenberg_matrix =
        ComplexMatrix::Zero(iter_count, iter_count);
    for (UINT i = 0; i < iter_count; i++) {
        mu_timed_state.load(state_list[i]);
        mu_timed_state.multiply_coef(-mu_);
        this->apply_to_state(&tmp_state, *state_list[i], &multiplied_state);
        multiplied_state.add_state(&mu_timed_state);

        for (UINT j = 0; j < i + 1; j++) {
            const auto coef = state::inner_product(
                static_cast<QuantumState*>(state_list[j]), &multiplied_state);
            hessenberg_matrix(j, i) = coef;
            tmp_state.load(state_list[j]);
            tmp_state.multiply_coef(-coef);
            multiplied_state.add_state(&tmp_state);
        }

        const auto norm = multiplied_state.get_squared_norm();
        if (i != iter_count - 1) {
            hessenberg_matrix(i + 1, i) = std::sqrt(norm);
        }
        multiplied_state.normalize(norm);
        state_list.push_back(multiplied_state.copy());
    }

    Eigen::ComplexEigenSolver<ComplexMatrix> eigen_solver(hessenberg_matrix);
    const auto eigenvalues = eigen_solver.eigenvalues();
    const auto eigenvectors = eigen_solver.eigenvectors();

    // Find ground state eigenvalue and eigenvector.
    UINT minimum_eigenvalue_index = 0;
    auto minimum_eigenvalue = eigenvalues[0];
    for (UINT i = 0; i < eigenvalues.size(); i++) {
        if (eigenvalues[i].real() < minimum_eigenvalue.real()) {
            minimum_eigenvalue_index = i;
            minimum_eigenvalue = eigenvalues[i];
        }
    }

    // Compose ground state vector and store it to `state`.
    present_state.multiply_coef(0.0);
    for (UINT i = 0; i < state_list.size() - 1; i++) {
        tmp_state.load(state_list[i]);
        tmp_state.multiply_coef(eigenvectors(i, minimum_eigenvalue_index));
        present_state.add_state(&tmp_state);
    }
    state->load(&present_state);

    // Free states allocated by `QuantumState::copy()`.
    for (auto used_state : state_list) {
        delete used_state;
    }
    return minimum_eigenvalue + mu_;
}

CPPCTYPE GeneralQuantumOperator::solve_ground_state_eigenvalue_by_power_method(
    QuantumStateBase* state, const UINT iter_count, const CPPCTYPE mu) const {
    if (this->get_term_count() == 0) {
        std::stringstream error_message_stream;
        error_message_stream
            << "Error: "
               "GeneralQuantumOperator::solve_ground_state_eigenvalue_by_"
               "power_method("
               "QuantumStateBase * state, const UINT iter_count, const "
               "CPPCTYPE mu): At least one PauliOperator is required.";
        throw std::invalid_argument(error_message_stream.str());
    }

    CPPCTYPE mu_;
    if (mu == 0.0) {
        // mu is not changed from default value.
        mu_ = this->calculate_default_mu();
    } else {
        mu_ = mu;
    }

    // Stores a result of A|a>
    auto multiplied_state = QuantumState(state->qubit_count);
    // Stores a result of -\mu|a>
    auto mu_timed_state = QuantumState(state->qubit_count);
    auto work_state = QuantumState(state->qubit_count);
    for (UINT i = 0; i < iter_count; i++) {
        mu_timed_state.load(state);
        mu_timed_state.multiply_coef(-mu_);

        multiplied_state.multiply_coef(0.0);
        this->apply_to_state(&work_state, *state, &multiplied_state);
        state->load(&multiplied_state);
        state->add_state(&mu_timed_state);
        state->normalize(state->get_squared_norm());
    }
    return this->get_expectation_value(state) + mu;
}

void GeneralQuantumOperator::apply_to_state(QuantumStateBase* work_state,
    const QuantumStateBase& state_to_be_multiplied,
    QuantumStateBase* dst_state) const {
    if (state_to_be_multiplied.qubit_count != dst_state->qubit_count) {
        throw std::invalid_argument(
            "Qubit count of state_to_be_multiplied and dst_state must be the "
            "same");
    }

    dst_state->multiply_coef(0.0);
    const auto term_count = this->get_term_count();
    for (UINT i = 0; i < term_count; i++) {
        work_state->load(&state_to_be_multiplied);
        const auto term = this->get_term(i);
        auto pauli_operator =
            gate::Pauli(term->get_index_list(), term->get_pauli_id_list());
        pauli_operator->update_quantum_state(work_state);
        work_state->multiply_coef(term->get_coef());
        dst_state->add_state(work_state);
        delete pauli_operator;
    }
}

void GeneralQuantumOperator::apply_to_state(
    QuantumStateBase* state, QuantumStateBase* dst_state) const {
    if (state->qubit_count != dst_state->qubit_count) {
        throw std::invalid_argument(
            "Qubit count of state_to_be_multiplied and dst_state must be the "
            "same");
    }

    dst_state->set_zero_norm_state();
    const auto term_count = this->get_term_count();
    for (UINT i = 0; i < term_count; i++) {
        const auto term = this->get_term(i);
        _apply_pauli_to_state(
            term->get_pauli_id_list(), term->get_index_list(), state);
        dst_state->add_state_with_coef(term->get_coef(), state);
        _apply_pauli_to_state(
            term->get_pauli_id_list(), term->get_index_list(), state);
    }
}

void GeneralQuantumOperator::apply_to_state_single_thread(
    QuantumStateBase* state, QuantumStateBase* dst_state) const {
    if (state->qubit_count != dst_state->qubit_count) {
        throw std::invalid_argument(
            "Qubit count of state_to_be_multiplied and dst_state must be the "
            "same");
    }

    dst_state->set_zero_norm_state();
    const auto term_count = this->get_term_count();
    for (UINT i = 0; i < term_count; i++) {
        const auto term = this->get_term(i);
        _apply_pauli_to_state_single_thread(
            term->get_pauli_id_list(), term->get_index_list(), state);
        dst_state->add_state_with_coef_single_thread(term->get_coef(), state);
        _apply_pauli_to_state_single_thread(
            term->get_pauli_id_list(), term->get_index_list(), state);
    }
}

void GeneralQuantumOperator::_apply_pauli_to_state(
    std::vector<UINT> pauli_id_list, std::vector<UINT> target_index_list,
    QuantumStateBase* state) const {
    // this function is same as the gate::Pauli update quantum state
    if (state->is_state_vector()) {
#ifdef _USE_GPU
        if (state->get_device_name() == "gpu") {
            multi_qubit_Pauli_gate_partial_list_host(target_index_list.data(),
                pauli_id_list.data(), (UINT)target_index_list.size(),
                state->data(), state->dim, state->get_cuda_stream(),
                state->device_number);
            // _update_func_gpu(this->_target_qubit_list[0].index(), _angle,
            // state->data(), state->dim);
            return;
        }
#endif
        multi_qubit_Pauli_gate_partial_list(target_index_list.data(),
            pauli_id_list.data(), (UINT)target_index_list.size(),
            state->data_c(), state->dim);
    } else {
        dm_multi_qubit_Pauli_gate_partial_list(target_index_list.data(),
            pauli_id_list.data(), (UINT)target_index_list.size(),
            state->data_c(), state->dim);
    }
}

void GeneralQuantumOperator::_apply_pauli_to_state_single_thread(
    std::vector<UINT> pauli_id_list, std::vector<UINT> target_index_list,
    QuantumStateBase* state) const {
    // this function is same as the gate::Pauli update quantum state
    if (state->is_state_vector()) {
#ifdef _USE_GPU
        if (state->get_device_name() == "gpu") {
            // TODO: make it single thread for this function
            multi_qubit_Pauli_gate_partial_list_host(target_index_list.data(),
                pauli_id_list.data(), (UINT)target_index_list.size(),
                state->data(), state->dim, state->get_cuda_stream(),
                state->device_number);
            // _update_func_gpu(this->_target_qubit_list[0].index(), _angle,
            // state->data(), state->dim);
            return;
        }
#endif
        multi_qubit_Pauli_gate_partial_list_single_thread(
            target_index_list.data(), pauli_id_list.data(),
            (UINT)target_index_list.size(), state->data_c(), state->dim);
    } else {
        // TODO: make it single thread for this function
        dm_multi_qubit_Pauli_gate_partial_list(target_index_list.data(),
            pauli_id_list.data(), (UINT)target_index_list.size(),
            state->data_c(), state->dim);
    }
}

CPPCTYPE GeneralQuantumOperator::calculate_default_mu() const {
    double mu = 0.0;
    const auto term_count = this->get_term_count();
    for (UINT i = 0; i < term_count; i++) {
        const auto term = this->get_term(i);
        mu += std::abs(term->get_coef().real());
    }
    return static_cast<CPPCTYPE>(mu);
}

GeneralQuantumOperator* GeneralQuantumOperator::copy() const {
    auto quantum_operator = new GeneralQuantumOperator(_qubit_count);
    for (auto pauli : this->_operator_list) {
        quantum_operator->add_operator(pauli->copy());
    }
    return quantum_operator;
}

GeneralQuantumOperator GeneralQuantumOperator::operator+(
    const GeneralQuantumOperator& target) const {
    auto res = this->copy();
    *res += target;
    return *res;
}

GeneralQuantumOperator* GeneralQuantumOperator::get_dagger() const {
    auto quantum_operator = new GeneralQuantumOperator(_qubit_count);
    for (auto pauli : this->_operator_list) {
        quantum_operator->add_operator(
            std::conj(pauli->get_coef()), pauli->get_pauli_string());
    }
    return quantum_operator;
}

GeneralQuantumOperator GeneralQuantumOperator::operator+(
    const PauliOperator& target) const {
    auto res = this->copy();
    *res += target;
    return *res;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator+=(
    const GeneralQuantumOperator& target) {
    ITYPE i, j;
    auto terms = target.get_terms();
#pragma omp parallel for
    for (i = 0; i < _operator_list.size(); i++) {
        auto pauli_operator = _operator_list[i];
        for (j = 0; j < terms.size(); j++) {
            auto target_operator = terms[j];
            auto pauli_x = pauli_operator->get_x_bits();
            auto pauli_z = pauli_operator->get_z_bits();
            auto target_x = target_operator->get_x_bits();
            auto target_z = target_operator->get_z_bits();
            if (pauli_x.size() != target_x.size()) {
                UINT max_size = std::max(pauli_x.size(), target_x.size());
                pauli_x.resize(max_size);
                pauli_z.resize(max_size);
                target_x.resize(max_size);
                target_z.resize(max_size);
            }
            if (pauli_x == target_x && pauli_z == target_z) {
                _operator_list[i]->change_coef(_operator_list[i]->get_coef() +
                                               target_operator->get_coef());
            }
        }
    }
    for (j = 0; j < terms.size(); j++) {
        auto target_operator = terms[j];
        bool flag = true;
        for (i = 0; i < _operator_list.size(); i++) {
            auto pauli_operator = _operator_list[i];
            auto pauli_x = pauli_operator->get_x_bits();
            auto pauli_z = pauli_operator->get_z_bits();
            auto target_x = target_operator->get_x_bits();
            auto target_z = target_operator->get_z_bits();
            if (pauli_x.size() != target_x.size()) {
                UINT max_size = std::max(pauli_x.size(), target_x.size());
                pauli_x.resize(max_size);
                pauli_z.resize(max_size);
                target_x.resize(max_size);
                target_z.resize(max_size);
            }
            if (pauli_x == target_x && pauli_z == target_z) {
                flag = false;
            }
        }
        if (flag) {
            this->add_operator(target_operator->copy());
        }
    }
    return *this;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator+=(
    const PauliOperator& target) {
    bool flag = true;
    ITYPE i;
#pragma omp parallel for
    for (i = 0; i < _operator_list.size(); i++) {
        auto pauli_operator = _operator_list[i];
        auto pauli_x = pauli_operator->get_x_bits();
        auto pauli_z = pauli_operator->get_z_bits();
        auto target_x = target.get_x_bits();
        auto target_z = target.get_z_bits();
        if (pauli_x.size() != target_x.size()) {
            UINT max_size = std::max(pauli_x.size(), target_x.size());
            pauli_x.resize(max_size);
            pauli_z.resize(max_size);
            target_x.resize(max_size);
            target_z.resize(max_size);
        }
        if (pauli_x == target_x && pauli_z == target_z) {
            _operator_list[i]->change_coef(
                _operator_list[i]->get_coef() + target.get_coef());
            flag = false;
        }
    }
    if (flag) {
        this->add_operator(target.copy());
    }
    return *this;
}

GeneralQuantumOperator GeneralQuantumOperator::operator-(
    const GeneralQuantumOperator& target) const {
    auto res = this->copy();
    *res -= target;
    return *res;
}

GeneralQuantumOperator GeneralQuantumOperator::operator-(
    const PauliOperator& target) const {
    auto res = this->copy();
    *res -= target;
    return *res;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator-=(
    const GeneralQuantumOperator& target) {
    ITYPE i, j;
    auto terms = target.get_terms();
#pragma omp parallel for
    for (i = 0; i < _operator_list.size(); i++) {
        auto pauli_operator = _operator_list[i];
        for (j = 0; j < terms.size(); j++) {
            auto target_operator = terms[j];
            auto pauli_x = pauli_operator->get_x_bits();
            auto pauli_z = pauli_operator->get_z_bits();
            auto target_x = target_operator->get_x_bits();
            auto target_z = target_operator->get_z_bits();
            if (pauli_x.size() != target_x.size()) {
                UINT max_size = std::max(pauli_x.size(), target_x.size());
                pauli_x.resize(max_size);
                pauli_z.resize(max_size);
                target_x.resize(max_size);
                target_z.resize(max_size);
            }
            if (pauli_x == target_x && pauli_z == target_z) {
                _operator_list[i]->change_coef(_operator_list[i]->get_coef() -
                                               target_operator->get_coef());
            }
        }
    }
    for (j = 0; j < terms.size(); j++) {
        auto target_operator = terms[j];
        bool flag = true;
        for (i = 0; i < _operator_list.size(); i++) {
            auto pauli_operator = _operator_list[i];
            auto pauli_x = pauli_operator->get_x_bits();
            auto pauli_z = pauli_operator->get_z_bits();
            auto target_x = target_operator->get_x_bits();
            auto target_z = target_operator->get_z_bits();
            if (pauli_x.size() != target_x.size()) {
                UINT max_size = std::max(pauli_x.size(), target_x.size());
                pauli_x.resize(max_size);
                pauli_z.resize(max_size);
                target_x.resize(max_size);
                target_z.resize(max_size);
            }
            if (pauli_x == target_x && pauli_z == target_z) {
                flag = false;
            }
        }
        if (flag) {
            auto copy = target_operator->copy();
            copy->change_coef(-copy->get_coef());
            this->add_operator(copy);
        }
    }
    return *this;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator-=(
    const PauliOperator& target) {
    bool flag = true;
    ITYPE i;
    for (i = 0; i < _operator_list.size(); i++) {
        auto pauli_operator = _operator_list[i];
        auto pauli_x = pauli_operator->get_x_bits();
        auto pauli_z = pauli_operator->get_z_bits();
        auto target_x = target.get_x_bits();
        auto target_z = target.get_z_bits();
        if (pauli_x.size() != target_x.size()) {
            UINT max_size = std::max(pauli_x.size(), target_x.size());
            pauli_x.resize(max_size);
            pauli_z.resize(max_size);
            target_x.resize(max_size);
            target_z.resize(max_size);
        }
        if (pauli_x == target_x && pauli_z == target_z) {
            _operator_list[i]->change_coef(
                _operator_list[i]->get_coef() - target.get_coef());
            flag = false;
        }
    }
    if (flag) {
        auto copy = target.copy();
        copy->change_coef(-copy->get_coef());
        this->add_operator(copy);
    }
    return *this;
}
GeneralQuantumOperator GeneralQuantumOperator::operator*(
    const GeneralQuantumOperator& target) const {
    auto res = this->copy();
    *res *= target;
    return *res;
}

GeneralQuantumOperator GeneralQuantumOperator::operator*(
    const PauliOperator& target) const {
    auto res = this->copy();
    *res *= target;
    return *res;
}
GeneralQuantumOperator GeneralQuantumOperator::operator*(
    CPPCTYPE target) const {
    auto res = this->copy();
    *res *= target;
    return *res;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator*=(
    const GeneralQuantumOperator& target) {
    auto copy = this->copy();
    _operator_list.clear();
    auto terms = copy->get_terms();
    auto target_terms = target.get_terms();
    ITYPE i, j;
#pragma omp parallel for
    for (i = 0; i < terms.size(); i++) {
        auto pauli_operator = terms[i];
        for (j = 0; j < target_terms.size(); j++) {
            auto target_operator = target_terms[j];
            PauliOperator* product = new PauliOperator;
            *product = (*pauli_operator) * (*target_operator);
            *this += *product;
        }
    }
    return *this;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator*=(
    const PauliOperator& target) {
    auto copy = this->copy();
    _operator_list.clear();
    ITYPE i;
    auto terms = copy->get_terms();
#pragma omp parallel for
    for (i = 0; i < terms.size(); i++) {
        auto pauli_operator = terms[i];
        PauliOperator* product = new PauliOperator;
        *product = (*pauli_operator) * (target);
        *this += *product;
    }
    return *this;
}

GeneralQuantumOperator& GeneralQuantumOperator::operator*=(CPPCTYPE target) {
    ITYPE i;
#pragma omp parallel for
    for (i = 0; i < _operator_list.size(); i++) {
        *_operator_list[i] *= target;
    }
    return *this;
}
namespace quantum_operator {
GeneralQuantumOperator* create_general_quantum_operator_from_openfermion_file(
    std::string file_path) {
    UINT qubit_count = 0;
    std::vector<CPPCTYPE> coefs;
    std::vector<std::string> ops;

    // loading lines and check qubit_count
    std::string str_buf;
    std::vector<std::string> index_list;

    std::ifstream ifs;
    std::string line;
    ifs.open(file_path);

    while (getline(ifs, line)) {
        std::tuple<double, double, std::string> parsed_items =
            parse_openfermion_line(line);
        const auto coef_real = std::get<0>(parsed_items);
        const auto coef_imag = std::get<1>(parsed_items);
        str_buf = std::get<2>(parsed_items);

        CPPCTYPE coef(coef_real, coef_imag);
        coefs.push_back(coef);
        ops.push_back(str_buf);
        index_list = split(str_buf, "IXYZ ");

        for (UINT i = 0; i < index_list.size(); ++i) {
            UINT n = std::stoi(index_list[i]) + 1;
            if (qubit_count < n) qubit_count = n;
        }
    }
    if (!ifs.eof()) {
        std::stringstream error_message_stream;
        error_message_stream << "ERROR: Invalid format";
        throw std::runtime_error(error_message_stream.str());
    }
    ifs.close();

    GeneralQuantumOperator* general_quantum_operator =
        new GeneralQuantumOperator(qubit_count);

    for (UINT i = 0; i < ops.size(); ++i) {
        general_quantum_operator->add_operator(
            new PauliOperator(ops[i].c_str(), coefs[i]));
    }

    return general_quantum_operator;
}

GeneralQuantumOperator* create_general_quantum_operator_from_openfermion_text(
    std::string text) {
    UINT qubit_count = 0;
    std::vector<CPPCTYPE> coefs;
    std::vector<std::string> ops;

    std::string str_buf;
    std::vector<std::string> index_list;

    std::vector<std::string> lines;
    lines = split(text, "\n");
    for (std::string line : lines) {
        std::tuple<double, double, std::string> parsed_items =
            parse_openfermion_line(line);
        const auto coef_real = std::get<0>(parsed_items);
        const auto coef_imag = std::get<1>(parsed_items);
        str_buf = std::get<2>(parsed_items);

        CPPCTYPE coef(coef_real, coef_imag);
        coefs.push_back(coef);
        ops.push_back(str_buf);
        index_list = split(str_buf, "IXYZ ");

        for (UINT i = 0; i < index_list.size(); ++i) {
            UINT n = std::stoi(index_list[i]) + 1;
            if (qubit_count < n) qubit_count = n;
        }
    }
    GeneralQuantumOperator* general_quantum_operator =
        new GeneralQuantumOperator(qubit_count);

    for (UINT i = 0; i < ops.size(); ++i) {
        general_quantum_operator->add_operator(
            new PauliOperator(ops[i].c_str(), coefs[i]));
    }

    return general_quantum_operator;
}

std::pair<GeneralQuantumOperator*, GeneralQuantumOperator*>
create_split_general_quantum_operator(std::string file_path) {
    UINT qubit_count = 0;
    std::vector<CPPCTYPE> coefs;
    std::vector<std::string> ops;

    std::ifstream ifs;
    ifs.open(file_path);

    if (!ifs) {
        std::stringstream error_message_stream;
        error_message_stream << "ERROR: Cannot open file";
        throw std::runtime_error(error_message_stream.str());
    }

    // loading lines and check qubit_count
    std::string str_buf;
    std::vector<std::string> index_list;

    std::string line;
    while (getline(ifs, line)) {
        std::tuple<double, double, std::string> parsed_items =
            parse_openfermion_line(line);
        const auto coef_real = std::get<0>(parsed_items);
        const auto coef_imag = std::get<1>(parsed_items);
        str_buf = std::get<2>(parsed_items);
        if (str_buf == (std::string)NULL) {
            continue;
        }
        CPPCTYPE coef(coef_real, coef_imag);
        coefs.push_back(coef);
        ops.push_back(str_buf);
        index_list = split(str_buf, "IXYZ ");

        for (UINT i = 0; i < index_list.size(); ++i) {
            UINT n = std::stoi(index_list[i]) + 1;
            if (qubit_count < n) qubit_count = n;
        }
    }
    if (!ifs.eof()) {
        std::stringstream error_message_stream;
        error_message_stream << "ERROR: Invalid format";
        throw std::runtime_error(error_message_stream.str());
    }
    ifs.close();

    GeneralQuantumOperator* general_quantum_operator_diag =
        new GeneralQuantumOperator(qubit_count);
    GeneralQuantumOperator* general_quantum_operator_non_diag =
        new GeneralQuantumOperator(qubit_count);

    for (UINT i = 0; i < ops.size(); ++i) {
        if (ops[i].find("X") != std::string::npos ||
            ops[i].find("Y") != std::string::npos) {
            general_quantum_operator_non_diag->add_operator(
                new PauliOperator(ops[i].c_str(), coefs[i]));
        } else {
            general_quantum_operator_diag->add_operator(
                new PauliOperator(ops[i].c_str(), coefs[i]));
        }
    }

    return std::make_pair(
        general_quantum_operator_diag, general_quantum_operator_non_diag);
}
}  // namespace quantum_operator

bool check_Pauli_operator(const GeneralQuantumOperator* quantum_operator,
    const PauliOperator* pauli_operator) {
    auto vec = pauli_operator->get_index_list();
    UINT val = 0;
    if (vec.size() > 0) {
        val = std::max(val, *std::max_element(vec.begin(), vec.end()));
    }
    return val < (quantum_operator->get_qubit_count());
}

std::string GeneralQuantumOperator::to_string() const {
    std::stringstream os;
    auto term_count = this->get_term_count();
    for (UINT index = 0; index < term_count; index++) {
        os << this->get_term(index)->get_coef() << " ";
        os << this->get_term(index)->get_pauli_string();
        if (index != term_count - 1) {
            os << " + ";
        }
    }
    return os.str();
}


#include <alien/kernels/mtl/algebra/MTLInternalLinearAlgebra.h>

#include <alien/core/backend/LinearAlgebraT.h>

#include <alien/kernels/mtl/data_structure/MTLInternal.h>
#include <alien/kernels/mtl/data_structure/MTLMatrix.h>
#include <alien/kernels/mtl/data_structure/MTLVector.h>
#include <alien/kernels/mtl/MTLBackEnd.h>
#include <boost/numeric/mtl/mtl.hpp>

#include <alien/core/impl/MultiMatrixImpl.h>
#include <alien/core/impl/MultiVectorImpl.h>

#include <arccore/base/NotImplementedException.h>

/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/

namespace Alien {

/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/

template class ALIEN_EXTERNAL_PACKAGES_EXPORT LinearAlgebra<BackEnd::tag::mtl>;

IInternalLinearAlgebra<MTLMatrix, MTLVector>*
MTLInternalLinearAlgebraFactory()
{
  return new MTLInternalLinearAlgebra();
}

/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/

MTLInternalLinearAlgebra::MTLInternalLinearAlgebra()
{
  // Devrait faire le MTLInitialize qui est actuellement dans le solveur
  // Attention, cette initialisation serait globale et non restreinte à cet objet
}

/*---------------------------------------------------------------------------*/

MTLInternalLinearAlgebra::~MTLInternalLinearAlgebra()
{
}

/*---------------------------------------------------------------------------*/

Arccore::Real
MTLInternalLinearAlgebra::norm0(const Vector& x) const
{
  return infinity_norm(x.internal()->m_internal);
}

/*---------------------------------------------------------------------------*/

Arccore::Real
MTLInternalLinearAlgebra::norm1(const Vector& x) const
{
  return one_norm(x.internal()->m_internal);
}

/*---------------------------------------------------------------------------*/

Arccore::Real
MTLInternalLinearAlgebra::norm2(const Vector& x) const
{
  return two_norm(x.internal()->m_internal);
}

/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::mult(const Matrix& a, const Vector& x, Vector& r) const
{
  r.internal()->m_internal = a.internal()->m_internal * x.internal()->m_internal;
}

/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::axpy(
    Real alpha, const Vector& x, Vector& r) const
{
  r.internal()->m_internal += alpha * x.internal()->m_internal;
}

/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::aypx(
    Real alpha, Vector& y, const Vector& x) const
{
  throw Arccore::NotImplementedException(
      A_FUNCINFO, "MTLInternalLinearAlgebra::aypx not implemented");
}

/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::copy(const Vector& x, Vector& r) const
{
  r.internal()->m_internal = x.internal()->m_internal;
}

/*---------------------------------------------------------------------------*/

Arccore::Real
MTLInternalLinearAlgebra::dot(const Vector& x, const Vector& y) const
{
  return dot_real(x.internal()->m_internal, y.internal()->m_internal);
}

/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::scal(Real alpha, Vector& x) const
{
  throw Arccore::NotImplementedException(
      A_FUNCINFO, "MTLInternalLinearAlgebra::scal not implemented");
}

/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::diagonal(const Matrix& a, Vector& x) const
{
  throw Arccore::NotImplementedException(
      A_FUNCINFO, "MTLInternalLinearAlgebra::diagonal not implemented");
}

/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::reciprocal(Vector& x) const
{
  throw Arccore::NotImplementedException(
      A_FUNCINFO, "MTLInternalLinearAlgebra::reciprocal not implemented");
}

/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::pointwiseMult(const Vector& x, const Vector& y, Vector& w) const
{
  throw Arccore::NotImplementedException(
      A_FUNCINFO, "MTLInternalLinearAlgebra::pointwiseMult not implemented");
}

/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/

void
MTLInternalLinearAlgebra::mult(
    const Matrix& a, const UniqueArray<Real>& x, UniqueArray<Real>& r) const
{
  throw NotImplementedException(A_FUNCINFO, "LinearAlgebra::mult not implemented");
}
void
MTLInternalLinearAlgebra::axpy(
    Real alpha, const UniqueArray<Real>& x, UniqueArray<Real>& r) const
{
  throw NotImplementedException(A_FUNCINFO, "LinearAlgebra::axpy not implemented");
}
void
MTLInternalLinearAlgebra::aypx(
    Real alpha, UniqueArray<Real>& y, const UniqueArray<Real>& x) const
{
  throw NotImplementedException(A_FUNCINFO, "LinearAlgebra::aypx not implemented");
}
void
MTLInternalLinearAlgebra::copy(const UniqueArray<Real>& x, UniqueArray<Real>& r) const
{
  throw NotImplementedException(A_FUNCINFO, "LinearAlgebra::copy not implemented");
}
Real
MTLInternalLinearAlgebra::dot(
    Integer local_size, const UniqueArray<Real>& x, const UniqueArray<Real>& y) const
{
  throw NotImplementedException(A_FUNCINFO, "LinearAlgebra::dot not implemented");
  return Real();
}
void
MTLInternalLinearAlgebra::scal(Real alpha, UniqueArray<Real>& x) const
{
  throw NotImplementedException(A_FUNCINFO, "HypreLinearAlgebra::scal not implemented");
}

} // END_NAMESPACE

/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/


#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <boost/format.hpp>
#include <GL/glew.h>
#include <GLFW/glfw3.h>
#include <glm/glm.hpp>
#include "util.hpp"

using namespace std;
using namespace boost;
using namespace glm;

const static int WIDTH = 512;
const static int HEIGHT = 512;
const static int COUNT = 160;
const static int DEPTH = 1024;
const static int BATCH = 4;

int main(int argc, char* argv[]) {
	if (!glfwInit()) {
		throw runtime_error("Failed to initialize GLFW");
	}

	glfwWindowHint(GLFW_SAMPLES, 4);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 4);
	glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

	GLFWwindow* window = glfwCreateWindow(WIDTH, HEIGHT, "Render", NULL, NULL);
	if (!window) {
		throw runtime_error("Failed to open GLFW window");
	}
	glfwMakeContextCurrent(window);
	glfwSetInputMode(window, GLFW_STICKY_KEYS, GL_TRUE);

	glewExperimental = GL_TRUE;
	if (glewInit() != GLEW_OK) {
		throw runtime_error("Failed to initialize GLEW");
	}
	checkError("init");

	// motion
	const unsigned char* motionData = readBMP("motion.bmp");
	GLuint motion;
	glGenTextures(1, &motion);
	glActiveTexture(GL_TEXTURE0);
	glBindTexture(GL_TEXTURE_2D, motion);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, WIDTH, HEIGHT, 0, GL_RGB, GL_UNSIGNED_BYTE, motionData);
	glBindImageTexture(0, motion, 0, GL_FALSE, 0, GL_READ_ONLY, GL_RGBA8UI);
	checkError("motion");

	// position
	/*float* positionData = new float[COUNT * COUNT * DEPTH * 4];
	for (int i = 0; i < COUNT * COUNT * DEPTH; i += 4) {
		positionData[i + 3] = 0;
	}*/

	GLuint position;
	glGenTextures(1, &position);
	glActiveTexture(GL_TEXTURE1);
	glBindTexture(GL_TEXTURE_3D, position);
	glTexImage3D(GL_TEXTURE_3D, 0, GL_RGBA32F, COUNT, COUNT, DEPTH, 0, GL_RGBA, GL_FLOAT, NULL);
	glBindImageTexture(1, position, 0, GL_FALSE, 0, GL_READ_WRITE, GL_RGBA32F);
	checkError("position");

	// velocity
	/*float* velocityData = new float[COUNT * COUNT * 4];
	for (int i = 0; i < COUNT * COUNT; i += 4) {
		velocityData[i    ] = 0.5;
		velocityData[i + 1] = 0;
	}*/

	GLuint velocity;
	glGenTextures(1, &velocity);
	glActiveTexture(GL_TEXTURE2);
	glBindTexture(GL_TEXTURE_3D, velocity);
	glTexImage3D(GL_TEXTURE_3D, 0, GL_RGBA32F, COUNT, COUNT, DEPTH, 0, GL_RGBA, GL_FLOAT, NULL);
	glBindImageTexture(2, velocity, 0, GL_FALSE, 0, GL_READ_WRITE, GL_RGBA32F);
	checkError("velocity");

	// control
	float* controlData = new float[COUNT * COUNT * DEPTH * 4];
	for (int i = 0; i < (COUNT * COUNT * DEPTH * 4); i++) {
		controlData[i] = float(rand()) / RAND_MAX;
	}
	GLuint control;
	glGenTextures(1, &control);
	glActiveTexture(GL_TEXTURE3);
	glBindTexture(GL_TEXTURE_3D, control);
	glTexImage3D(GL_TEXTURE_3D, 0, GL_RGBA32F, COUNT, COUNT, DEPTH, 0, GL_RGBA, GL_FLOAT, controlData);
	glBindImageTexture(3, control, 0, GL_FALSE, 0, GL_READ_WRITE, GL_RGBA32F);
	checkError("control");

	// display
	GLuint display;
	glGenTextures(1, &display);
	glActiveTexture(GL_TEXTURE4);
	glBindTexture(GL_TEXTURE_2D, display);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, WIDTH, HEIGHT, 0, GL_RGB, GL_UNSIGNED_BYTE, NULL);
	glBindImageTexture(4, display, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA8UI);
	glGenerateMipmap(GL_TEXTURE_2D);
	checkError("display");

	// compute shader
	const static char* computeSrc =
		#include "compute.shader"
	;
	GLenum computeType[] = { GL_COMPUTE_SHADER };
	GLuint compute = buildProgram(&computeSrc, computeType, 1);
	glUseProgram(compute);
	glUniform1i(glGetUniformLocation(compute, "depth"), DEPTH);
	glUniform1i(glGetUniformLocation(compute, "batch"), BATCH);
	glUniform2i(glGetUniformLocation(compute, "size"), WIDTH, HEIGHT);
	checkError("compute");

	// draw shader
	const static char* vertexSrc =
		#include "vertex.shader"
	;
	const static char* fragmentSrc =
		#include "fragment.shader"
	;
	const char* drawSrc[] = { vertexSrc, fragmentSrc };
	GLenum drawType[] = { GL_VERTEX_SHADER, GL_FRAGMENT_SHADER };
	GLuint draw = buildProgram(drawSrc, drawType, 2);
	glUseProgram(draw);
	glUniform2i(glGetUniformLocation(draw, "size"), WIDTH, HEIGHT);
	checkError("draw");

	// draw data
	GLuint vertexArray;
	glGenVertexArrays(1, &vertexArray);
	glBindVertexArray(vertexArray);

	const static GLfloat vertexData[] = {
		-1.0f, -1.0f,
		-1.0f,  1.0f,
		 1.0f, -1.0f,
		 1.0f,  1.0f
	};

	GLuint vertexBuffer;
	glGenBuffers(1, &vertexBuffer);
	glBindBuffer(GL_ARRAY_BUFFER, vertexBuffer);
	glBufferData(GL_ARRAY_BUFFER, sizeof(vertexData), vertexData, GL_STATIC_DRAW);

	glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 0, (void*)0);
	glEnableVertexAttribArray(0);
	checkError("vertex");

	// main loop
	double lastTime = glfwGetTime();
	double nowTime;
	char fpsTitle[32];
	do {
		glUseProgram(compute);
		glDispatchCompute(COUNT / 16, COUNT / 16, 1);
		glUseProgram(draw);
		glDrawArrays(GL_TRIANGLE_STRIP, 0, sizeof(vertexData));
		glfwSwapBuffers(window);
		glfwPollEvents();
		glClearTexImage(display, 0, GL_RGB, GL_UNSIGNED_BYTE, NULL);

		nowTime = glfwGetTime();
		sprintf(fpsTitle, "Render - FPS: %.2f", 1 / (nowTime - lastTime));
		glfwSetWindowTitle(window, fpsTitle);
		lastTime = nowTime;
	} while (glfwGetKey(window, GLFW_KEY_ESCAPE) != GLFW_PRESS && glfwWindowShouldClose(window) == 0);

	return EXIT_SUCCESS;
}


#include <cstdlib>
#include <iostream>
#include <iomanip>
#include <vector>
#include <numeric>
#include <limits>
#include <boost/timer.hpp>
#include <sys/stat.h>
#include <errno.h>
#include <string.h>
#include <omp.h>

#include "absmc3D.h"

using namespace absmc;
using namespace absmc::graphics;

using util::createFileName;
using util::parseFileName;
using util::config;

using std::cout;
using std::cerr;
using std::endl;

typedef AgentBase<3> agent_t;
typedef Point<3, double> point_t;

class FSetState : public std::binary_function<size_t, agent_t, void> {
public:
    void operator()(size_t id, agent_t* agentBase) {
        SMC3D* agent = dynamic_cast<SMC3D*>(agentBase);
        if (agent) { agent->setState(SMC3D::SG2M); }
    }
};

/// This executable is outdated. Use at your own risk
/// main
int main(int argc, char *argv[])
{
    if (argc!=3) {
        cout << "Usage: runStandAlone configDir outputRoot" << endl;
        return EXIT_FAILURE;
    }

    const std::string configDir(argv[1]);
    const std::string outputRoot(argv[2]);
    if (mkdir(outputRoot.c_str(),0755)==-1 && errno != EEXIST) {
    cout<<"miserably failed to create outputRoot <" << outputRoot << ">: "<<((std::string)strerror(errno)) << endl;
    return 1;
    }

    if (mkdir((outputRoot+"/data.runStandAlone").c_str(),0755)==-1 && errno != EEXIST) {cout<<"miserably failed to create data.runStandAlone <" << outputRoot << "/data.runStandAlone>: "<<((std::string)strerror(errno)) << endl; return 1; }
    const std::string outputDir = createFileName(outputRoot, "data.runStandAlone", "");

    const std::string configFileName  = createFileName(configDir, "absmc", "cfg");
    config().readFile(configFileName);

    // agent input files
    const std::string inFileName = config().getValue<std::string>("run_input_file");
    const std::string baseName = parseFileName(inFileName).baseName;

    // morse potential parameters (dimensionless)
    const double U1 = config().getValue<double>("U1");
    const double r1 = config().getValue<double>("r1");
    const double U2 = config().getValue<double>("U2");
    const double r2 = config().getValue<double>("r2");
    // equilibrium distance (dimensionless)
    const double r0 = r1*r2/(r2-r1)*log(U1/U2*r2/r1);
    // characteristic length scale
    const double L = config().getValue<double>("smc_mean_rad");
    // boundary range
    const double BR = config().getValue<double>("Boundary_Inactivity_Range");

    // load agents from file; includes smc+iel+obstacle agents
    AgentContainer<3> agents;
    AgentFileReader<3>::readFile(createFileName(configDir, "stage3."+baseName, "dat"), AgentFactory3D(), agents);
    cout << agents.count() << " agents loaded." << endl;

    // initialise neighbourhoods
    AccumulativeNeighbourDetector<agent_t> nbDetector(r0*L);
    const std::string nbFileName = createFileName(configDir, "stage3."+baseName+"_nb", "dat");
    NBReader<3>::readFile(nbFileName, agents);

    // set agents at longitudinal (x) domain boundaries immobile;
    const double lX = config().getValue<double>("lX");
    const double mX = config().getValue<double>("run_boundary_mobility_x");
    const double mY = config().getValue<double>("run_boundary_mobility_y");
    const double mZ = config().getValue<double>("run_boundary_mobility_z");
    agents.forAll(PSelectCloseToBoundary<agent_t>(0, 0.0, lX, BR), FSetMobility<agent_t>(point_t(mX, mY, mZ) ) );

    //ADDED set mobility of EEL AND SMCS to zero
    agents.forAll(tEEL3D, FSetMobility<agent_t>(point_t(0.0, 0.0, 0.0) ) );
    agents.forAll(tIEL3D, FSetMobility<agent_t>(point_t(0.0, 0.0, 0.0) ) );
//    agents.forAll(tSMC3D, FSetMobility<agent_t>(point_t(0.0, 0.0, 0.0) ) );

    // set SMC agents at longitudinal (x) domain boundaries and at the outer surface biologically inactive
    agents.forAll(PSelectCloseToBoundary<agent_t>(0, 0.0, lX, BR), FSetSMC3DInactive() );
    // set rules for SMC agents; no rules for IEL agents.
    const double smcMaxStress = config().getValue<double>("smc_max_stress");
    SMCNecrosisRule3D smcNecrosisRule(smcMaxStress);

    const double ciWeightSMC          = config().getValue<double>("ci_weight_smc");
    const double ciWeightIEL          = config().getValue<double>("ci_weight_iel");
    const double ciWeightEEL          = config().getValue<double>("ci_weight_eel");
    const double ciWeightObstacle     = config().getValue<double>("ci_weight_obstacle");
    const double ciRangeFactor        = config().getValue<double>("ci_range_factor");
    const double ciThresholdCount     = config().getValue<double>("ci_threshold_count");
    const double ciRange = ciRangeFactor*r0*L;
    NeighbourCache<agent_t>* nbCacheForCiRule = 0;
    SMCContactInhibitionRule3D smcContactInhibitionRule(ciWeightSMC, ciWeightIEL, ciWeightEEL, ciWeightObstacle, ciRange, ciThresholdCount, &nbCacheForCiRule);

    const double drugConcThreshold    = config().getValue<double>("smc_drug_conc_threshold");
    const double wssMaxThreshold      = config().getValue<double>("smc_wss_max_threshold");
    SMCCellCycleRule3D smcCellCycleRule(drugConcThreshold, wssMaxThreshold);

    //pretending we have a blood flow code running

    //    logger::info("Sizes nagents=%d, OSI=%d, WSS=%d, Drug=%d",nAgents,dataWssOsi.size(),dataWssMax.size(),dataDrugConc.size());
        int timeStep = 0;
    //new Nitric oxide rule
    SMCNitricOxideRule3D smcNitricOxideRule3D(0.0); //is timestep defined?

    CompositeRule<SMC3D> smcRule;
//    smcRule.add(&smcNecrosisRule);
    smcRule.add(&smcContactInhibitionRule);
    smcRule.add(&smcCellCycleRule);
    smcRule.add(&smcNitricOxideRule3D);

    // set rules for bulk smc agents only; other smc agents are left biologically inactive with nil agent rule.
    const double outerRadius = config().getValue<double>("outer_radius");
    agents.forAll(PSelectBulkAgentsOfSpecifiedType<agent_t>(0.0, lX, 0.0, 0.0, outerRadius, 4.0*0.5*r0*L, tSMC3D),
                  FSetAgentRule<SMC3D>(&smcRule) );

    //set initial number of SMCS in NO rule
    smcNitricOxideRule3D.setInitial_Num_SMCs(agents.count(tSMC3D));
/*
    //////////////// physics parameters for Morse force (dimensionless)
    // cell-cell stiffness at equilibrium distance is zero, so kEq is at 0.001
    //    const double young = 0.1; //MPa
    const double youngeq = 0.77;//0.5*young/(1.0-0.27*0.27);
    const double B = 1.0/8.3439;//20.0*youngeq/(9.0*math::pi); //MPa = N/mm2
    const double kEq = 0.00273753;//0.5*B*math::pi*0.000716223/L;// N/mm/L
    // maximal time step (from rough stability estimation)
    const double maxDt = 2.0 / kEq; // L*mm/N
    // maximal time step passed to solver
    const double solverMaxDt = maxDt;  // mm/N
    // maximal displacement
    const double solverMaxDispl = 0.05*2.0*L;// mm //0.1*0.5*r0;
*/

    //////////////// physics parameters for Morse force (dimensionless)
    // cell-cell stiffness at equilibrium distance
    const double kEq = 8.0;//U1/(r1*r1)*exp(-r0/r1) - U2/(r2*r2)*exp(-r0/r2);
    // maximal time step (from rough stability estimation)
    const double maxDt = 2.0 / kEq;;
    // maximal time step passed to solver
    const double solverMaxDt = 0.01*maxDt;;
    // maximal displacement
    const double solverMaxDispl = 0.01*0.5*r0;

/*
    //////////////// physics parameters for Morse force (dimensionless)
    // cell-cell stiffness at equilibrium distance
    const double B = 1.0/8.3439;
    const double kEq = B*0.0027/0.015/0.015;//0.01/6.0/0.015;//U1/(r1*r1)*exp(-r0/r1) - U2/(r2*r2)*exp(-r0/r2);
    // maximal time step (from rough stability estimation)
    const double maxDt = 2.0 / kEq;;
    // maximal time step passed to solver
    const double solverMaxDt = 0.1*maxDt;;
    // maximal displacement
    const double solverMaxDispl = 0.1*0.5*r0;
*/

    cout << "U1                = " << U1 << endl;
    cout << "r1                = " << r1 << endl;
    cout << "U2                = " << U2 << endl;
    cout << "r2                = " << r2 << endl;
    cout << "r0                = " << r0 << endl;
    cout << "kEq               = " << kEq << endl;;
    cout << "maxDt             = " << maxDt << endl;
    cout << "solverMaxDt       = " << solverMaxDt << endl;
    cout << "solverMaxDispl    = " << solverMaxDispl << endl;

    ZeroUnaryForce3D uForce;
    Bilinear3D bForce(U1, r1, U2, r2, L);

    std::ofstream integratorLogFile(createFileName(outputDir, "fr", "dat").c_str() );
    AdaptiveEulerIntegrator<agent_t, ZeroUnaryForce3D, Bilinear3D>
            integrator(uForce, bForce, L, solverMaxDispl, solverMaxDt, nbDetector, integratorLogFile);

    // iteration parameters
    const double iterTime = config().getValue<double>("run_iter_time");
    const int maxIter     = config().getValue<int>("run_max_iter");
    const int vtkIter     = config().getValue<int>("run_vtk_iter");
    const int datIter     = config().getValue<int>("run_dat_iter");
    FixedIntervalController fixedIntervalController(iterTime);

    // system-characteristic force
    const double charForce = kEq * r0 * 0.1;
    // desired convergence level
    const double eps      = config().getValue<double>("run_convergence_level");

    cout << "charForce         = " << charForce << endl;
    cout << "eps               = " << eps << endl;
    //cahrForce*eps = 0.000108268 for morse seems like 0.00005 for beyondhertz (already multiplied by 0.015, still need to apply B).
    ForceResidualMaxNormController frMaxNormController(std::numeric_limits<double>::infinity(), charForce, eps);

    // vtp output parameters
    const std::string vtpScalars = config().getValue<std::string>("run_vtp_scalars");
    const std::string vtpVectors = config().getValue<std::string>("run_vtp_vectors");

    boost::timer timer;
    double tIter = 0.0;
    double tTotal = 0.0;
    double omptimefinal = 0.0;

    double omptime = omp_get_wtime();

    // main loop

    timeStep = 0;
    int newSMCs = 0;
    smcNitricOxideRule3D.setNewSMCCounter(newSMCs);

    int oldNumOfSMCs = agents.count(tSMC3D);
    int newNumOfSMCs = oldNumOfSMCs;
    int selectedSMCs = 0;
    int initial_num_SMCs = agents.count(tSMC3D);
    double new_endothelial_probability = 0.0;
    util::Random::init(0); //trying to control the randomness by allowing repetition

    for (int iter=0; iter<maxIter; iter++) {

        const size_t nAgents = agents.getAgentVector().size();
        //pretending we have a blood flow code running
        for (size_t iAgent = 0; iAgent < nAgents; iAgent++)
        {
            if ((agents.getAgentVector()[iAgent]->getTypeId() == tSMC3D) || (agents.getAgentVector()[iAgent]->getTypeId() == tIEL3D))
            {
                CellBase3D * agent;
                agent = (CellBase3D*) agents.getAgentVector()[iAgent];
                //        agent->setWssOsi(dataWssOsi[iAgent]);
                agent->setWssMax(12.0);

            }
        }

        if (iter%vtkIter==0) {

            tIter = timer.elapsed();
            tTotal += tIter;
            omptimefinal = - omptime + omp_get_wtime();

            VtkWriter<agent_t>::writeVtkPolyData(agents.getAgentVector(),
                                                 createFileName(outputDir, baseName, "vtp", iter, 6),
                                                 vtpScalars, vtpVectors);

            cout << "iter=" << iter << ",  t(" << vtkIter << " it)=";
            cout << std::setprecision(4) << tIter << ",  t(total)=" << tTotal;
            cout << ",  t(omp)=" << std::setprecision(4) << omptimefinal;
            cout << ", nSMC=" << agents.count(tSMC3D) << ", nIEL=" << agents.count(tIEL3D)<< ", nEEL=" << agents.count(tEEL3D)<< ", nSEL=" << selectedSMCs
            << ", nNEW=" << newSMCs << ", prob=" << 100.0*new_endothelial_probability << std::endl;

            timer.restart();
        }

        if (datIter!=0 && iter%datIter==0) {
            NBWriter<3>::writeFile(agents, createFileName(configDir, "stage4."+baseName+"_nb", "dat") );
            agents.writeFile(createFileName(outputDir, baseName, "dat", iter, 6) );
        }
         // set SMC agents at longitudinal (x) domain boundaries and at the outer surface biologically inactive
        agents.forAll(PSelectCloseToBoundary<agent_t>(0, 0.0, lX, BR), FSetSMC3DInactive() );
        // exec agent rules
        nbCacheForCiRule = new NeighbourCache<agent_t>(agents.getAgentVector() );

        double increase_rate = 0.0205/24.0;                 // considering 59% endothelium after day 3rd and 100% after 23 days. equation y=(41/20)x + 52.85. so we just take the slope, staring point is 59%.
        //       double increase_rate = 0.41/12.0/24.0;             // considering 59% endothelium after day 3rd and 100% after 15 days. equation y=(41/12)x + 48.75. so we just take the slope, staring point is 59%.
        //double increase_rate = 1.0/23.0/24.0;         //0.181159420289855    // considering 0% endothelium after day 0 and 100% after 23 days. equation y=(100/23)x + 0. so we just take the slope, staring point is 0%.
        double old_endothelial_probability = 0.0;            // probability assuming a fixed number of cells. This is just used to calculate the probabilty defined in the next line.
        //endothelial_probability = 0.0;
        new_endothelial_probability = 0.0;            // probability that includes the newly produced SMCs, More the SMC divide into daugter cells, more the probabilty is.

        if (iter <=72) {
            increase_rate = 0.59/3.0/24.0; //going from 0 to 59% between days 0 an 3. Note: it could very well be that coverage does not end up to be exactly 59% after day 3.
        } else {
            increase_rate = 0.41/12.0/24.0;
        }

        if (iter <= 360) {                     // EC probabilty will be 100% after 15 days, 15x24=360
            //old_endothelial_probability = 100.0*increase_rate/(47.15 - iter*increase_rate); //NOW iter SHOULD BE IN HOURS, VALID FOR 3 to 23, we use 100-52.85=47.15
            //        old_endothelial_probability = 100.0*increase_rate/(51.25 - iter*increase_rate); //NOW iter SHOULD BE IN HOURS, VALID FOR 3 to 15, we use 100-48.75=51.25
            if (iter <=72) {
                old_endothelial_probability = increase_rate/(1.0 - iter*increase_rate); //0.001934235976789 //NOW iter SHOULD BE IN HOURS, VALID FOR 0 to 23
            } else {
                old_endothelial_probability = increase_rate/(1.0 - 72.0*0.59/3.0/24.0 - (iter-72.0)*increase_rate);
            }
            //double percent_cells_left = 47.15 - iter*increase_rate;        //VALID FOR 3 to 23, we use 100-52.85=47.15
            //        double percent_cells_left = 51.25 - iter*increase_rate;        //VALID FOR 3 to 15, we use 100-48.75=51.25
            double percent_cells_left = 1.0 - iter*increase_rate;    //0.936594202898551    //VALID FOR 0 to 23
            int num_cells_left = (percent_cells_left * initial_num_SMCs); //274991.550724637679798
            //        int newSMCs = 0;
            //    endothelial_probability = 100.0*(old_endothelial_probability*num_cells_left/100.0 + newSMCs)/(num_cells_left + newSMCs);
            new_endothelial_probability = old_endothelial_probability;//(old_endothelial_probability*num_cells_left + newSMCs)/(num_cells_left + newSMCs);  // this is required as the total number of SMCs are also increasing with time.

            //                if (iter == 72) { /*endothelial_probability = 59.0;*/new_endothelial_probability = 59.0;}     // first step for Nakazawa et al.
        } else { /*endothelial_probability = 100.0;*/new_endothelial_probability = 1.0; }

        //    printf("endprob = %f \n", new_endothelial_probability);

        if (new_endothelial_probability > 1.0) { /*endothelial_probability =100.0;*/new_endothelial_probability = 1.0; }

        smcNitricOxideRule3D.setnew_endothelial_probability(100.0*new_endothelial_probability);

        agents.execAgentRules();
        newNumOfSMCs = agents.count(tSMC3D);
        selectedSMCs = agents.countSelected(tSMC3D);
        newSMCs = newNumOfSMCs - oldNumOfSMCs;
        smcNitricOxideRule3D.setNewSMCCounter(newSMCs);
        delete nbCacheForCiRule;

        // ... and restore equilibrium
        nbDetector.updateSlowTimeScale(agents.getAgentVector() );
        integrator.integrate(agents.getAgentVector(), frMaxNormController);
    }

    // write the final configuration to file
    agents.writeFile(createFileName(configDir, "stage4."+baseName, "dat") );
     NBWriter<3>::writeFile(agents, createFileName(configDir, "stage4."+baseName+"_nb", "dat") );
    VtkWriter<agent_t>::writeVtkPolyData(agents.getAgentVector(), createFileName(configDir, "stage4."+baseName, "vtp"), vtpScalars, vtpVectors);

    return EXIT_SUCCESS;
}


/* Copyright © 2017 Apple Inc. All rights reserved.
 *
 * Use of this source code is governed by a BSD-3-clause license that can
 * be found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
 */
// ML Data
#include <unity/toolkits/ml_data_2/ml_data.hpp>
#include <unity/toolkits/ml_data_2/metadata.hpp>
#include <unity/toolkits/ml_data_2/ml_data_iterators.hpp>
#include <unity/toolkits/ml_data_2/sframe_index_mapping.hpp>

// Toolkits
#include <unity/toolkits/nearest_neighbors/nearest_neighbors.hpp>
#include <unity/toolkits/nearest_neighbors/ball_tree_neighbors.hpp>
#include <unity/lib/variant_deep_serialize.hpp>
#include <unity/lib/toolkit_util.hpp>
#include <logger/assertions.hpp>
#include <sframe/groupby.hpp>
#include <sframe/sframe.hpp>
#include <sframe/sarray.hpp>

// Miscellaneous
#include <timer/timer.hpp>
#include <algorithm>
#include <Eigen/SparseCore>
#include <limits>
#include <stack>
#include <table_printer/table_printer.hpp>

namespace turi {
namespace nearest_neighbors {

#define NONE_FLAG ((size_t) -1)

/**
 * Destructor. Make sure bad things don't happen
 */
ball_tree_neighbors::~ball_tree_neighbors(){

}

/**
* Set options
*/
void ball_tree_neighbors::init_options(
  const std::map<std::string, flexible_type>& _options) {

  options.create_integer_option("leaf_size",
                            "Max number of points in a leaf node of the ball tree",
                            0,
                            0,
                            std::numeric_limits<int>::max(),
                            true);

  options.create_string_option("label",
                             "Name of the reference dataset column with row labels.",
                             "",
                             false);

  // Set options and update model state with final option values
  options.set_options(_options);
  add_or_update_state(flexmap_to_varmap(options.current_option_values()));
}

/**
 * Train a ball tree nearest neighbors model.
 */
void ball_tree_neighbors::train(const sframe& X,
                                const std::vector<flexible_type>& ref_labels,
                                const std::vector<dist_component_type>& composite_distance_params,
                                const std::map<std::string, flexible_type>& opts) {

  logprogress_stream << "Starting ball tree nearest neighbors model training." << std::endl;

  timer t;
  double start_time = t.current_time();

  // Validate the inputs.
  init_options(opts);
  validate_distance_components(composite_distance_params, X);

  // Create the ml_data object for the reference data.
  initialize_model_data(X, ref_labels);

  // Initialize the distance components. NOTE: this needs data to be initialized
  // first because the row slicers need the column indices to be sorted.
  initialize_distances();

  ASSERT_FALSE(composite_distances.empty());
  dist_component c = composite_distances[0];

  if (metadata->num_dimensions() > 100) {
    logprogress_stream << "\nWARNING: The computational advantage of the "
      << "ball tree tends to diminish as the number of variables grows. With more "
      << "than 100 variables, the ball tree may not be optimal for this dataset."
      << std::endl;
  }

  // Figure out leaf size if the user didn't set it
  size_t leaf_size = (size_t)options.value("leaf_size");

  if (leaf_size == 0) {
    leaf_size = std::max((size_t)1000, (size_t)ceil((double)num_examples / 2048)); // max tree depth of 12
    options.set_option("leaf_size", leaf_size);
  }

  size_t min_leaves = ceil((double)num_examples / leaf_size);
  tree_depth = ceil(log2(min_leaves)) + 1;
  size_t num_leaves = std::max(size_t(1), size_t(std::pow(2, tree_depth - 1)));
  size_t num_nodes = 2 * num_leaves - 1;

  if (tree_depth > 12) {
    logprogress_stream << "\nWARNING: The ball tree is very large. Consider "
      << "increasing the leaf size to create a smaller tree and improve "
      << "performance." << std::endl;
  }

  // Initialize tree and loop objects
  if (is_dense) {
    pivots.resize(num_nodes);     // pivot observations
  } else {
    pivots_sp.resize(num_nodes);
  }

  node_radii.resize(num_nodes);                      // distance from pivot to furthest node member
  std::vector<double> first_child_radius (num_nodes, 0);  // distance from the pivot to the first child observation
  std::vector<double> median_dist (num_nodes);   // median of distances to the first child point
  double middle_dist;                            // the temporary container for the median distance in a node
  membership.resize(num_examples);               // point membership in nodes
  std::vector<double> pivot_dist(num_examples);  // distance from each point to its pivot (at the lowest tree level)
  std::vector<double> first_child_dist(num_examples);  // distance from each point to the first child (at the lowest tree level)

  // Set the radius and membership to 0 to start
  for (size_t i = 0; i < num_nodes; ++i) {
    node_radii[i] = 0;
  }

  for (size_t i = 0; i < num_examples; ++i) {
    membership[i] = 0;
  }

  size_t num_variables = metadata->num_dimensions();

  // Declare loop variables
  DenseVector p(num_variables);      // dense pivot observation
  DenseVector x(num_variables);      // dense generic observation
  SparseVector x_sp(num_variables);  // sparse pivot observation
  SparseVector p_sp(num_variables);  // sparse query observation

  size_t a;                          // generic row index for a point
  size_t idx_node;                   // index of the current node
  size_t idx_node_start;             // index of the first node on a level
  size_t idx_node_end;               // index of the last node on a level
  size_t num_level_nodes;            // number of nodes on a level

  // Switch for maintaining balance in the nodes. If a point is exactly on the
  // median this toggle indicates which child node to assign it to.
  bool first_child_median_flag = true;

  // Choose the first pivot
  // NOTE: for now this will be the first row of the reference data, but this
  // should probably be chosen randomly.

  if (is_dense) {
    mld_ref.get_iterator().fill_observation(x);
    pivots[0] = x;
  } else {
    mld_ref.get_iterator().fill_observation(x_sp);
    pivots_sp[0] = x_sp;
  }

  table_printer table( {{"Tree level", 0}, {"Elapsed Time", 0} });
  table.print_header();

  // The main loop over levels of the tree
  // NOTE: the second-to-last tree level creates the leaves, so the loop should
  // end at tree_depth - 1.
  for (size_t tree_level = 0; tree_level < (tree_depth - 1); ++tree_level) {

    if (cppipc::must_cancel()) {
      log_and_throw("Toolkit cancelled by user.");
    }

    // Get the node indices for nodes on the current level
    idx_node_start = std::pow(2, tree_level) - 1;
    idx_node_end = std::pow(2, (tree_level + 1)) - 2;
    num_level_nodes = idx_node_end - idx_node_start + 1;

    // First pass over the data
    for (auto it = mld_ref.get_iterator(); !it.done(); ++it) {

      // Get the required data
      a = it.row_index();
      idx_node = membership[a];

      if (is_dense) {
        p = pivots[idx_node];
        it.fill_observation(x);
        pivot_dist[a] = c.distance->distance(x, p);

        // find the largest distance to the pivot and index of the point
        if (pivot_dist[a] >= node_radii[idx_node]) {
          node_radii[idx_node] = pivot_dist[a];
          pivots[2 * idx_node + 1] = x;
        }

      } else {  // data is not dense
        p_sp = pivots_sp[idx_node];
        it.fill_observation(x_sp);
        pivot_dist[a] = c.distance->distance(x_sp, p_sp);

        // find the largest distance to the pivot and index of the point
        if (pivot_dist[a] >= node_radii[idx_node]) {
          node_radii[idx_node] = pivot_dist[a];
          pivots_sp[2 * idx_node + 1] = x_sp;
        }
      }
    }

    // Create vector of vectors to store the first child distances contiguously
    // for each node.
    std::vector<std::vector<double>> node_dists(num_level_nodes);

    // Second pass over the data
    for (auto it = mld_ref.get_iterator(); !it.done(); ++it) {

      // Get the required data
      a = it.row_index();
      idx_node = membership[a];

      if (is_dense) {
        p = pivots[2 * idx_node + 1];
        it.fill_observation(x);

        // Find all of the distances to the first child and pick the second child
        // as the point furthest away.
        first_child_dist[a] = c.distance->distance(x, p);

        if (first_child_dist[a] >= first_child_radius[idx_node]) {
          first_child_radius[idx_node] = first_child_dist[a];
          pivots[2 * idx_node + 2] = x;
        }

      } else { // data is not dense
        p_sp = pivots_sp[2 * idx_node + 1];
        it.fill_observation(x_sp);

        // Find all of the distances to the first child and pick the second child
        // as the point furthest away.
        first_child_dist[a] = c.distance->distance(x_sp, p_sp);

        if (first_child_dist[a] >= first_child_radius[idx_node]) {
          first_child_radius[idx_node] = first_child_dist[a];
          pivots_sp[2 * idx_node + 2] = x_sp;
        }
      }

      // Keep the first child distances compiled by node for median computation
      node_dists[idx_node - idx_node_start].push_back(first_child_dist[a]);
    }

    // Find the median first child distance for each node
    for (size_t j = 0; j < num_level_nodes; ++j) {
      if (node_dists[j].size() > 1) {

        std::nth_element(node_dists[j].begin(),
                         node_dists[j].begin() + node_dists[j].size()/2,
                         node_dists[j].end());
        middle_dist = node_dists[j][node_dists[j].size()/2];

        // if there are an even number of elements get the median of the middle two
        if (node_dists[j].size() % 2 == 0) {
          std::nth_element(node_dists[j].begin(),
                           node_dists[j].begin() + node_dists[j].size()/2 - 1,
                           node_dists[j].end());
          middle_dist = (middle_dist + node_dists[j][node_dists[j].size()/2 - 1]) / 2;
        }

        median_dist[j + idx_node_start] = middle_dist;

      } else {
        // set median distance to -1 so that singletons always go to second child
        median_dist[j + idx_node_start] = -1;
      }
    }

    // Third pass over the data
    // - assign each point to a child
    // - careful about maintaining balance here
    for (size_t b = 0; b < num_examples; ++b) {
      idx_node = membership[b];
      if (first_child_dist[b] < median_dist[idx_node]) {
        membership[b] = 2 * idx_node + 1;

      } else if (first_child_dist[b] > median_dist[idx_node]) {
        membership[b] = 2 * idx_node + 2;

      } else {  // the point is exactly on the median
          if (first_child_median_flag) {
            membership[b] = 2 * idx_node + 1;
            first_child_median_flag = false;

          } else {
            membership[b] = 2 * idx_node + 2;
            first_child_median_flag = true;
          }
      }
    }

    table.print_row(tree_level, progress_time());

  } // end loop over tree levels

  // Find the radii for each of the leaf nodes
  for (auto it = mld_ref.get_iterator(); !it.done(); ++it) {

    // Get the required data
    a = it.row_index();
    idx_node = membership[a];

    if (is_dense) {
      // Find the largest distance to the pivot and index of that point
      p = pivots[idx_node];
      it.fill_observation(x);
      pivot_dist[a] = c.distance->distance(x, p);

    } else { // data is not dense
      // Find the largest distance to the pivot and index of that point
      p_sp = pivots_sp[idx_node];
      it.fill_observation(x_sp);
      pivot_dist[a] = c.distance->distance(x_sp, p_sp);
    }

    if (pivot_dist[a] >= node_radii[idx_node]) {
      node_radii[idx_node] = pivot_dist[a];
    }
  }

  table.print_row(tree_depth - 1, progress_time());

  // Group the reference data by leaf node ID

  // convert the reference labels to an SArray.
  std::shared_ptr<sarray<flexible_type>> sa_ref_labels(new sarray<flexible_type>);
  sa_ref_labels->open_for_write();
  flex_type_enum ref_label_type = reference_labels[0].get_type();
  sa_ref_labels->set_type(ref_label_type);
  turi::copy(ref_labels.begin(), ref_labels.end(), *sa_ref_labels);
  sa_ref_labels->close();

  // convert membership into a shared pointer to an sarray
  std::shared_ptr<sarray<flexible_type>> member_column(new sarray<flexible_type>);
  member_column->open_for_write();
  member_column->set_type(flex_type_enum::INTEGER);
  turi::copy(membership.begin(), membership.end(), *member_column);
  member_column->close();

  // add the membership sarray as a column to the reference data and group
  sframe sf_refs = X.add_column(sa_ref_labels, "__nearest_neighbors_ref_label");
  sf_refs = sf_refs.add_column(member_column, "__nearest_neighbors_membership");
  sf_refs = turi::group(sf_refs, "__nearest_neighbors_membership");

  // extract the grouped membership vector and remove from the dataset.
  auto member_reader = sf_refs.select_column("__nearest_neighbors_membership")->get_reader();
  std::vector<flexible_type> temp(num_examples);
  member_reader->read_rows(0, num_examples, temp);
  std::copy(temp.begin(), temp.end(), membership.begin());

  size_t idx_member_column = sf_refs.column_index("__nearest_neighbors_membership");
  sf_refs = sf_refs.remove_column(idx_member_column);

  // extract the map of grouped row indices from the dataset.
  auto label_reader = sf_refs.select_column("__nearest_neighbors_ref_label")->get_reader();
  std::vector<flexible_type> temp2(num_examples);
  label_reader->read_rows(0, num_examples, temp2);

  // this modifies the model's stored reference labels, *not* the vector passed to this function.
  std::copy(temp2.begin(), temp2.end(), reference_labels.begin());

  size_t idx_label_column = sf_refs.column_index("__nearest_neighbors_ref_label");
  sf_refs = sf_refs.remove_column(idx_label_column);

  // Re-make the ML data with the row-permuted data for storage in the model
  mld_ref = v2::ml_data(metadata);
  mld_ref.fill(sf_refs);

  add_or_update_state({ {"method", "ball_tree"},
                        {"tree_depth", tree_depth},
                        {"leaf_size", leaf_size},
                        {"training_time", t.current_time() - start_time} });
  table.print_footer();
}  // end the create function

/**
 * Make predictions using an existing ball tree nearest neighbors model.
 * /note For each query point compute the distance to every reference point.
 */
sframe ball_tree_neighbors::query(const v2::ml_data& mld_queries,
                                  const std::vector<flexible_type>& query_labels,
                                  const size_t k, const double radius,
                                  const bool include_self_edges) const {

  timer t;

  size_t num_queries = mld_queries.size();
  size_t num_nodes = node_radii.size();

  // Construct the distance object pointer
  ASSERT_FALSE(composite_distances.empty());
  dist_component c = composite_distances[0];

  // Compute the actual number of nearest neighbors and construct the data
  // structures to hold candidate neighbors while reference points are searched
  size_t kstar;

  if (k == NONE_FLAG) {
    kstar = NONE_FLAG;
  } else {
    kstar = std::min(k, mld_ref.size());
  }

  std::vector<neighbor_candidates> topk (num_queries,
                    neighbor_candidates(-1, kstar, radius, include_self_edges));

  parallel_for(0, num_queries, [&](size_t i) {
    topk[i].set_label(i);
  });

  atomic<size_t> n_query_points = 0;

  table_printer table({ {"Query points", 0}, {"% Complete.", 0}, {"Elapsed Time", 0}});
  table.print_header();

  in_parallel([&](size_t thread_idx, size_t num_threads) GL_GCC_ONLY(GL_HOT) {

      // Find the nearest neighbors for each query point
      // ---------------------------------------------------------------------------

      size_t num_variables = metadata->num_dimensions();

      DenseVector x(num_variables);  // reference observation
      DenseVector q(num_variables);  // query observation
      SparseVector x_sp(num_variables);  // reference observation
      SparseVector q_sp(num_variables);  // query observation

      std::stack<size_t> node_stack; // nodes that still need to be checked
      size_t idx_node;               // node currently being checked
      size_t idx_query;              // index of the query point
      bool activate_node;            // indicates whether to traverse the node
      double dist_child1;            // distance to the first child pivot of an active node
      double dist_child2;            // distance to the second child pivot of an active node
      size_t idx_start, idx_end;     // indicate where the reference data starts and stops for leaf nodes
      double dist;                   // distance between a query point and a leaf reference point
      double min_dist_possible;      // minimum possible distance between a query and a node

      auto it_ref = mld_ref.get_iterator();

      // Iterate over query points
      for (auto it_query = mld_queries.get_iterator(thread_idx, num_threads);
           !it_query.done(); ++it_query) {

        if (cppipc::must_cancel()) {
          log_and_throw("Toolkit cancelled by user.");
        }

        ASSERT_TRUE(it_query.row_index() != NONE_FLAG);

        if (is_dense) {
          it_query.fill_observation(q);
        } else {
          it_query.fill_observation(q_sp);
        }

        idx_query = it_query.row_index();
        node_stack.push(0);
        min_dist_possible = 0;

        // Loop over nodes in the traversal queue
        while (!node_stack.empty()) {
          idx_node = node_stack.top();
          node_stack.pop();

          // Compute the minimum possible distance from the query to the node
          if (is_dense) {
            min_dist_possible =
                c.distance->distance(pivots[idx_node], q) - node_radii[idx_node];
          } else {
            min_dist_possible =
                c.distance->distance(pivots_sp[idx_node], q_sp) - node_radii[idx_node];
          }

          // Decide if the node needs to be processed
          activate_node = activate_query_node(kstar, radius, min_dist_possible,
                                              topk[idx_query].candidates.size(),
                                              topk[idx_query].get_max_dist());

          if (activate_node) {

            // The active node is internal
            if (idx_node < num_nodes / 2) {

              // find the closest child pivot to the query
              if (is_dense) {
                dist_child1 = c.distance->distance(q, pivots[2 * idx_node + 1]);
                dist_child2 = c.distance->distance(q, pivots[2 * idx_node + 2]);
              } else {
                dist_child1 = c.distance->distance(q_sp, pivots_sp[2 * idx_node + 1]);
                dist_child2 = c.distance->distance(q_sp, pivots_sp[2 * idx_node + 2]);
              }

              // add child nodes to the stack, closest on the top
              if (dist_child1 <= dist_child2) {
                node_stack.push(2 * idx_node + 2);
                node_stack.push(2 * idx_node + 1);
              } else {
                node_stack.push(2 * idx_node + 1);
                node_stack.push(2 * idx_node + 2);
              }

              // The active node is a leaf
            } else {

              // figure out where the leaf members are in the reference ml_data
              idx_start = NONE_FLAG;
              idx_end = NONE_FLAG;

              for (size_t i = 0; i < membership.size(); i++) {
                if ((idx_start == NONE_FLAG) && (membership[i] == idx_node)) {
                  idx_start = i;
                }

                if ((idx_start != NONE_FLAG) && (membership[i] == idx_node)) {
                  idx_end = i + 1;
                }
              }

              DASSERT_TRUE(idx_end >= idx_start);

              if ((idx_start == NONE_FLAG) || (idx_end == NONE_FLAG)) {
                continue;  // if the node is empty, move on to the next node in the stack
              }

              for (it_ref.seek(idx_start); it_ref.row_index() != idx_end; ++it_ref) {

                if (is_dense) {
                  it_ref.fill_observation(x);
                  dist = c.distance->distance(x, q);
                } else {
                  it_ref.fill_observation(x_sp);
                  dist = c.distance->distance(x_sp, q_sp);
                }

                DASSERT_TRUE(it_ref.row_index() != NONE_FLAG);

                topk[idx_query].evaluate_point(
                    std::pair<double, size_t>(dist, it_ref.row_index()));

              }  // end the loop over reference points in the leaf
            }  // end the leaf node processing
          } // end active node processing
        } // end tree traversal for a given query point

        size_t n_query_points_so_far = (++n_query_points);

        table.print_timed_progress_row( n_query_points_so_far,
                                        std::floor((4 * 100.0 * n_query_points_so_far) / num_queries) / 4.0,
                                        progress_time());

      }  // end the loop over query points
    });

  table.print_row("Done", " ", progress_time());
  table.print_footer();

  sframe result = write_neighbors_to_sframe(topk, reference_labels, query_labels);
  return result;
}

/**
* Turi Serialization Save
*/
void ball_tree_neighbors::save_impl(turi::oarchive& oarc) const {

  variant_deep_save(state, oarc);

  std::map<std::string, variant_type> data;

  data["membership"]         = to_variant(membership);
  data["node_radii"]         = to_variant(node_radii);
  data["tree_depth"]         = to_variant(tree_depth);
  data["is_dense"]           = to_variant(is_dense);

  variant_deep_save(data, oarc);

  // Now, a few that couldn't get saved in the above map
  oarc << pivots << pivots_sp;
  oarc << options
       << mld_ref
       << composite_params
       << untranslated_cols
       << reference_labels;
}

/**
 * Turi Serialization Load
 */
void ball_tree_neighbors::load_version(turi::iarchive& iarc, size_t version) {

  ASSERT_MSG((version == 0) || (version == 1) || (version == 2),
             "This model version cannot be loaded. Please re-save your model.");

  variant_deep_load(state, iarc);

  std::map<std::string, variant_type> data;

  variant_deep_load(data, iarc);

#define __EXTRACT(var) var = variant_get_value<decltype(var)>(data.at(#var));

  __EXTRACT(membership);
  __EXTRACT(node_radii);
  __EXTRACT(tree_depth);
  __EXTRACT(is_dense);
#undef __EXTRACT

  iarc >> pivots >> pivots_sp;
  iarc >> options;
  iarc >> mld_ref;

  metadata = mld_ref.metadata();

  // If loading from an old version, manually construct a single component that
  // assumes a single distance across all features. This is necessary because
  // the GLC v1.1 ball tree query operates on a set of distance components
  // stored in the model, rather than a distance name as in GLC v1.
  if (version == 0) {
    auto fn = function_closure_info();
    fn.native_fn_name = "_distances.";
    fn.native_fn_name += std::string(options.value("distance"));
    std::vector<std::string> features = variant_get_value<std::vector<std::string>>(state["features"]);
    dist_component_type p = std::make_tuple(features, fn, 1.0);
    composite_params = {p};

    // set empty untranslated columns for string features.
    untranslated_cols = {};
  }

  else {
    iarc >> composite_params;
    iarc >> untranslated_cols;
  }

  // construct the reference labels from the target column of the reference
  // ml_data.
  if (version < 2) {
    reference_labels.resize(mld_ref.size());

    in_parallel([&](size_t thread_idx, size_t num_threads) {
      for (auto it = mld_ref.get_iterator(); !it.done(); ++it) {
        reference_labels[it.row_index()] = metadata->target_indexer()->map_index_to_value(it.target_index());
      }
    });

    add_or_update_state({ {"num_distance_components", 1} });
  }

  else {
    iarc >> reference_labels;
  }

  initialize_distances();
}

/**
* Helper function to decide if a node should be activated for a query.
*/
bool ball_tree_neighbors::activate_query_node(size_t k, double radius,
                                              double min_poss_dist,
                                              size_t num_current_neighbors,
                                              double max_current_dist) const {
  bool activate = false;

  if (k == NONE_FLAG) {

    if (radius < 0) {         // neither k nor radius is defined
      activate = true;

    } else {                  // k is undefined, radius is defined
      if (min_poss_dist < radius)
        activate = true;
    }

  } else {
      // NOTE: in future versions, the default radius will be infinity. All
      // neighbors will be checked against radius, so the following conditional
      // won't be needed.
    if (radius < 0) {         // k is defined, radius is undefined

      // if the candidates set is empty, max_current_dist is -1.0, but
      // num_current_neighbors should be 0, so this should trigger (unless k is
      // 0). The same thing occurs when both k and radius are defined (below).
      if ((num_current_neighbors < k) || (min_poss_dist < max_current_dist))
        activate = true;

    } else {                  // both k and radius are defined
      if (min_poss_dist < radius)
       if ((min_poss_dist < max_current_dist) || (num_current_neighbors < k))
        activate = true;
    }
  }

  return activate;
}

}  // namespace nearest_neighbors
}  // namespace turi


// __BEGIN_LICENSE__
// Copyright (C) 2006-2011 United States Government as represented by
// the Administrator of the National Aeronautics and Space Administration.
// All Rights Reserved.
// __END_LICENSE__

#include <vw/Mosaic/GigapanQuadTreeConfig.h>

#include <boost/bind.hpp>

#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>
#include <boost/filesystem/fstream.hpp>
#include <boost/filesystem/convenience.hpp>
namespace fs = boost::filesystem;

namespace vw {
namespace mosaic {

  struct GigapanQuadTreeConfigData {
    BBox2 m_longlat_bbox;

    std::vector<std::pair<std::string,vw::BBox2i> > branch_func( QuadTreeGenerator const&, std::string const& name, BBox2i const& region ) const;
    void metadata_func( QuadTreeGenerator const&, QuadTreeGenerator::TileInfo const& info ) const;

    public:
      GigapanQuadTreeConfigData() : m_longlat_bbox(0, 0, 0, 0) {}
  };

  GigapanQuadTreeConfig::GigapanQuadTreeConfig()
    : m_data( new GigapanQuadTreeConfigData() )
  {}

  std::string GigapanQuadTreeConfig::image_path( QuadTreeGenerator const& qtree, std::string const& name ) {
    return QuadTreeGenerator::tiered_image_path()(qtree, name);
  }

  void GigapanQuadTreeConfig::configure( QuadTreeGenerator& qtree ) const {
    qtree.set_image_path_func( &image_path );
    qtree.set_cull_images( true );
    qtree.set_metadata_func( boost::bind(&GigapanQuadTreeConfigData::metadata_func,m_data,_1,_2) );

    if (m_data->m_longlat_bbox.width() != 0 || m_data->m_longlat_bbox.height() != 0) {
      qtree.set_branch_func( boost::bind(&GigapanQuadTreeConfigData::branch_func,m_data,_1,_2,_3) );
    }
  }

  void GigapanQuadTreeConfig::set_longlat_bbox( BBox2 const& bbox ) {
    m_data->m_longlat_bbox = bbox;
  }

  std::vector<std::pair<std::string,vw::BBox2i> > GigapanQuadTreeConfigData::branch_func( QuadTreeGenerator const& qtree, std::string const& name, BBox2i const& region ) const {
    std::vector<std::pair<std::string,vw::BBox2i> > children;
    if( region.height() > qtree.get_tile_size() ) {

      Vector2i dims = qtree.get_dimensions();
      double aspect_ratio = 2 * (region.width()/region.height()) * ( (m_longlat_bbox.width()/dims.x()) / (m_longlat_bbox.height()/dims.y()) );

      double bottom_lat = m_longlat_bbox.max().y() - region.max().y()*m_longlat_bbox.height() / dims.y();
      double top_lat = m_longlat_bbox.max().y() - region.min().y()*m_longlat_bbox.height() / dims.y();
      bool top_merge = ( bottom_lat > 0 ) && ( ( 1.0 / cos(M_PI/180 * bottom_lat) ) > aspect_ratio );
      bool bottom_merge = ( top_lat < 0 ) && ( ( 1.0 / cos(M_PI/180 * top_lat) ) > aspect_ratio );

      if( top_merge ) {
        children.push_back( std::make_pair( name + "4", BBox2i( region.min(), region.max() - Vector2i(0,region.height()/2) ) ) );
      }
      else {
        children.push_back( std::make_pair( name + "0", BBox2i( (region + region.min()) / 2 ) ) );
        children.push_back( std::make_pair( name + "1", BBox2i( (region + Vector2i(region.max().x(),region.min().y())) / 2 ) ) );
      }
      if( bottom_merge ) {
        children.push_back( std::make_pair( name + "5", BBox2i( region.min() + Vector2i(0,region.height()/2), region.max() ) ) );
      }
      else {
        children.push_back( std::make_pair( name + "2", BBox2i( (region + Vector2i(region.min().x(),region.max().y())) / 2 ) ) );
        children.push_back( std::make_pair( name + "3", BBox2i( (region + region.max()) / 2 ) ) );
      }
    }
    return children;
  }

  void GigapanQuadTreeConfigData::metadata_func( QuadTreeGenerator const& qtree, QuadTreeGenerator::TileInfo const& info ) const {
    bool root_node = ( info.name.size() == 0 );

    if ( root_node) {
      std::ostringstream json;
      fs::path file_path( info.filepath, fs::native );
      fs::path json_path = change_extension( file_path, ".json" );

      json << "{" << std::endl
           << "  \"width\": " << qtree.get_dimensions()[0] << "," << std::endl
           << "  \"height\": " << qtree.get_dimensions()[1] << "," << std::endl
           << "  \"nlevels\": " << qtree.get_tree_levels() << std::endl
           << "}" << std::endl;

      fs::ofstream jsonfs(json_path);
      jsonfs << json.str();
    }
  }

  // TODO: Is this actually the right function for Gigapan?
  cartography::GeoReference GigapanQuadTreeConfig::output_georef(uint32 xresolution, uint32 yresolution) {
    if (yresolution == 0)
      yresolution = xresolution;

    VW_ASSERT(xresolution == yresolution, LogicErr() << "TMS requires square pixels");

    cartography::GeoReference r;
    r.set_pixel_interpretation(cartography::GeoReference::PixelAsArea);

    // Note: the global TMS pixel space extends from +270 to -90
    // latitude, so that the lower-left hand corner is tile-
    // aligned, since TMS uses an origin in the lower left.
    Matrix3x3 transform;
    transform(0,0) = 360.0 / xresolution;
    transform(0,2) = -180;
    transform(1,1) = -360.0 / yresolution;
    transform(1,2) = 270;
    transform(2,2) = 1;
    r.set_transform(transform);

    return r;
  }
} // namespace mosaic
} // namespace vw