#include "mutation_annotated_tree.hpp"
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <iomanip>
#include <iostream>

// Uses one-hot encoding if base is unambiguous
// A:1,C:2,G:4,T:8
int8_t Mutation_Annotated_Tree::get_nuc_id (char nuc) {
    int8_t ret = 0b1111;
    switch(nuc) {
        case 'a':
        case 'A': ret = 0b1;
                  break;
        case 'c':
        case 'C': ret = 0b10;
                  break;
        case 'g':
        case 'G': ret = 0b100;
                  break;
        case 't':
        case 'T': ret = 0b1000;
                  break;
        case 'R': ret = 0b101;
                  break;
        case 'Y': ret = 0b1010;
                  break;
        case 'S': ret = 0b110;
                  break;
        case 'W': ret = 0b1001;
                  break;
        case 'K': ret = 0b1100;
                  break;
        case 'M': ret = 0b11;
                  break;
        case 'B': ret = 0b1110;
                  break;
        case 'D': ret = 0b1101;
                  break;
        case 'H': ret = 0b1011;
                  break;
        case 'V': ret = 0b111;
        case 'n':
        case 'N':
        default: ret = 0b1111;
                 break;
    }
    return ret;
}

// Sets bits at positions specified by nuc_vec to 1 in int8
int8_t Mutation_Annotated_Tree::get_nuc_id (std::vector<int8_t> nuc_vec) {
    int8_t ret = 0;
    int8_t one = 1;
    for (auto nuc: nuc_vec) {
        assert((nuc >= 0) && (nuc <=3));
        ret += (one << nuc);
    }
    return ret;
}

// Convert nuc_id back to IUPAC base
char Mutation_Annotated_Tree::get_nuc (int8_t nuc_id) {
    char ret = 'N';
    //assert ((nuc_id >= 1) && (nuc_id <= 15));
    switch(nuc_id) {
        case 1: ret = 'A';
                break;
        case 2: ret = 'C';
                break;
        case 3: ret = 'M';
                break;
        case 4: ret = 'G';
                break;
        case 5: ret = 'R';
                break;
        case 6: ret = 'S';
                break;
        case 7: ret = 'V';
                break;
        case 8: ret = 'T';
                break;
        case 9: ret = 'W';
                break;
        case 10: ret = 'Y';
                break;
        case 11: ret = 'H';
                break;
        case 12: ret = 'K';
                break;
        case 13: ret = 'D';
                break;
        case 14: ret = 'B';
                break;
        default: ret = 'N';
                 break;
    }
    return ret;
}

// A:0, C:1, G:2, T:3
int8_t Mutation_Annotated_Tree::get_nt (int8_t nuc_id) {
    int8_t ret = 0;
    switch(nuc_id) {
        case 1: ret = 0;
                break;
        case 2: ret = 1;
                break;
        case 4: ret = 2;
                break;
        case 8: ret = 3;
                break;
        default: ret = -1;
                 break;
    }
    return ret;
}

std::vector<int8_t> Mutation_Annotated_Tree::get_nuc_vec (char c) {
    switch (c) {
        case 'a':
        case 'A': return std::vector<int8_t>{0};
        case 'c':
        case 'C': return std::vector<int8_t>{1};
        case 'g':
        case 'G': return std::vector<int8_t>{2};
        case 't':
        case 'T': return std::vector<int8_t>{3};
        case 'R': return std::vector<int8_t>{0,2};
        case 'Y': return std::vector<int8_t>{1,3};
        case 'S': return std::vector<int8_t>{1,2};
        case 'W': return std::vector<int8_t>{0,3};
        case 'K': return std::vector<int8_t>{2,3};
        case 'M': return std::vector<int8_t>{0,1};
        case 'B': return std::vector<int8_t>{1,2,3};
        case 'D': return std::vector<int8_t>{0,2,3};
        case 'H': return std::vector<int8_t>{0,1,3};
        case 'V': return std::vector<int8_t>{0,1,2};
        case 'n':
        case 'N': return std::vector<int8_t>{0,1,2,3};
        default: return std::vector<int8_t>{0,1,2,3};
    }
}

std::vector<int8_t> Mutation_Annotated_Tree::get_nuc_vec_from_id (int8_t nuc_id) {
    return get_nuc_vec(get_nuc(nuc_id));
}

// Get newick stringstream for the input subtree rooted at some node (node) in
// the input tree T. Boolean arguments decide whether
// internal node ids and branch lengths are printed. If last boolean argument is
// set, branch lengths from input tree are retained, otherwise, branch length
// for a branch is equal to the number of mutations annotated on that branch
void Mutation_Annotated_Tree::write_newick_string (std::stringstream& ss, const Mutation_Annotated_Tree::Tree& T, Mutation_Annotated_Tree::Node* node, bool print_internal, bool print_branch_len, bool retain_original_branch_len, bool uncondense_leaves) {
    TIMEIT();

    std::vector<Node*> traversal = T.depth_first_expansion(node);
    size_t level_offset = node->level-1;
    size_t curr_level = 0;
    bool prev_open = true;

    std::stack<std::string> node_stack;
    std::stack<float> branch_length_stack;

    for (auto n: traversal) {
        size_t level = n->level-level_offset;
        float branch_length = n->branch_length;
        if (!retain_original_branch_len) {
            branch_length = static_cast<float>(n->mutations.size());
        }
        if (curr_level < level) {
            if (!prev_open) {
                ss << ',';
            }
            size_t l = level - 1;
            if (curr_level > 1) {
                l = level - curr_level;
            }
            for (size_t i=0; i < l; i++) {
                ss << '(';
                prev_open = true;
            }
            if (n->is_leaf()) {
                if (uncondense_leaves && (T.condensed_nodes.find(n->identifier) != T.condensed_nodes.end())) {
                    auto cn = T.condensed_nodes.at(n->identifier);
                    auto cn_size = cn.size();
                    for (size_t idx = 0; idx < cn_size; idx++) {
                        ss << cn[idx];
                        if (idx+1 < cn_size) {
                            ss << ',';
                        }
                    }
                }
                else {
                    ss << n->identifier;
                }
                if ((print_branch_len) && (branch_length >= 0)) {
                    ss << ':';
                    ss << branch_length;
                }
                prev_open = false;
            }
            else {
                node_stack.push(n->identifier);
                branch_length_stack.push(branch_length);
            }
        }
        else if (curr_level > level) {
            prev_open = false;
            for (size_t i = level; i < curr_level; i++) {
                ss << ')';
                if (print_internal){
                    ss << node_stack.top();
                }
                if ((print_branch_len) && (branch_length_stack.top() >= 0)) {
                    ss << ':';
                    ss << branch_length_stack.top();
                }
                node_stack.pop();
                branch_length_stack.pop();
            }
            if (n->is_leaf()) {
                if (uncondense_leaves && (T.condensed_nodes.find(n->identifier) != T.condensed_nodes.end())) {
                    auto cn = T.condensed_nodes.at(n->identifier);
                    ss << ',';
                    auto cn_size = cn.size();
                    for (size_t idx = 0; idx < cn_size; idx++) {
                        ss << cn[idx];
                        if (idx+1 < cn_size) {
                            ss << ',';
                        }
                    }
                }
                else {
                    ss << ',';
                    ss << n->identifier;
                }
                if ((print_branch_len) && (branch_length >= 0)) {
                    ss << ':';
                    ss << branch_length;
                }
            }
            else {
                node_stack.push(n->identifier);
                branch_length_stack.push(branch_length);
            }
        }
        else {
            prev_open = false;
            if (n->is_leaf()) {
                if (uncondense_leaves && (T.condensed_nodes.find(n->identifier) != T.condensed_nodes.end())) {
                    auto cn = T.condensed_nodes.at(n->identifier);
                    ss << ',';
                    auto cn_size = cn.size();
                    for (size_t idx = 0; idx < cn_size; idx++) {
                        ss << cn[idx];
                        if (idx+1 < cn_size) {
                            ss << ',';
                        }
                    }
                }
                else {
                    ss << ',';
                    ss << n->identifier;
                }
                if ((print_branch_len) && (branch_length >= 0)) {
                    ss << ':';
                    ss << branch_length;
                }
            }
            else {
                node_stack.push(n->identifier);
                branch_length_stack.push(branch_length);
            }
        }
        curr_level = level;
    }
    size_t remaining = node_stack.size();
    for (size_t i = 0; i < remaining; i++) {
        ss << ')';
        if (print_internal) {
            ss << node_stack.top();
        }
        if ((print_branch_len) && (branch_length_stack.top() >= 0)) {
            ss << ':';
            ss << branch_length_stack.top();
        }
        node_stack.pop();
        branch_length_stack.pop();
    }

    ss << ';';
}

std::string Mutation_Annotated_Tree::get_newick_string (const Mutation_Annotated_Tree::Tree& T, Mutation_Annotated_Tree::Node* node, bool print_internal, bool print_branch_len, bool retain_original_branch_len, bool uncondense_leaves) {
    std::stringstream newick_ss;
    write_newick_string(newick_ss, T, node, print_internal, print_branch_len, retain_original_branch_len, uncondense_leaves);
    return newick_ss.str();
}

std::string Mutation_Annotated_Tree::get_newick_string (const Tree& T, bool print_internal, bool print_branch_len, bool retain_original_branch_len, bool uncondense_leaves) {
    return get_newick_string(T, T.root, print_internal, print_branch_len, retain_original_branch_len, uncondense_leaves);
}

// Split string into words for a specific delimiter delim
void Mutation_Annotated_Tree::string_split (std::string const& s, char delim, std::vector<std::string>& words) {
    TIMEIT();
    size_t start_pos = 0, end_pos = 0;
    while ((end_pos = s.find(delim, start_pos)) != std::string::npos) {
        if ((end_pos == start_pos) || end_pos >= s.length()) {
            break;
        }
        words.emplace_back(s.substr(start_pos, end_pos-start_pos));
        start_pos = end_pos+1;
    }
    auto last = s.substr(start_pos, s.size()-start_pos);
    if (last != "") {
        words.push_back(std::move(last));
    }

}

// Split string into words (delimited by space, tabs etc.)
void Mutation_Annotated_Tree::string_split (std::string s, std::vector<std::string>& words) {
    std::string curr = "";
    std::vector<std::string> ret;

    // Used to split string around spaces.
    std::istringstream ss(s);

    std::string word;
    // Traverse through all words
    while (ss >> word) {
        words.push_back(std::move(word));
    };
}

Mutation_Annotated_Tree::Tree Mutation_Annotated_Tree::create_tree_from_newick_string (std::string newick_string) {
    TIMEIT();
    Tree T;

    std::vector<std::string> leaves;
    std::vector<size_t> num_open;
    std::vector<size_t> num_close;
    std::vector<std::queue<float>> branch_len (128);  // will be resized later if needed
    size_t level = 0;

    std::vector<std::string> s1;
    string_split(newick_string, ',', s1);

    num_open.reserve(s1.size());
    num_close.reserve(s1.size());

    for (auto s: s1) {
        size_t no = 0;
        size_t nc = 0;
        bool stop = false;
        bool branch_start = false;
        std::string leaf = "";
        std::string branch = "";
        for (auto c: s) {
            if (c == ':') {
                stop = true;
                branch = "";
                branch_start = true;
            }
            else if (c == '(') {
                no++;
                level++;
                if (branch_len.size() <= level) {
                  branch_len.resize(level*2);
                }
            }
            else if (c == ')') {
                stop = true;
                nc++;
                float len = (branch.size() > 0) ? std::stof(branch) : -1.0;
                branch_len[level].push(len);
                level--;
                branch_start = false;
            }
            else if (!stop) {
                leaf += c;
                branch_start = false;
            }
            else if (branch_start) {
                if (isdigit(c)  || c == '.' || c == 'e' || c == 'E' || c == '-' || c == '+') {
                    branch += c;
                }
            }
        }
        leaves.push_back(std::move(leaf));
        num_open.push_back(no);
        num_close.push_back(nc);
        float len = (branch.size() > 0) ? std::stof(branch) : -1.0;
        branch_len[level].push(len);
    }

    if (level != 0) {
        fprintf(stderr, "ERROR: incorrect Newick format!\n");
        exit(1);
    }

    T.curr_internal_node = 0;
    std::stack<Node*> parent_stack;

    for (size_t i=0; i<leaves.size(); i++) {
        auto leaf = leaves[i];
        auto no = num_open[i];
        auto nc = num_close[i];
        for (size_t j=0; j<no; j++) {
            std::string nid = std::to_string(++T.curr_internal_node);
            Node* new_node = NULL;
            if (parent_stack.size() == 0) {
                new_node = T.create_node(nid, branch_len[level].front());
            }
            else {
                new_node = T.create_node(nid, parent_stack.top(), branch_len[level].front());
            }
            branch_len[level].pop();
            level++;
            parent_stack.push(new_node);
        }
        T.create_node(leaf, parent_stack.top(), branch_len[level].front());
        branch_len[level].pop();
        for (size_t j=0; j<nc; j++) {
            parent_stack.pop();
            level--;
        }
    }

    if (T.root == NULL) {
        fprintf(stderr, "WARNING: Tree found empty!\n");
    }

    return T;
}

Mutation_Annotated_Tree::Tree Mutation_Annotated_Tree::create_tree_from_newick (std::string filename) {
    std::ifstream infile(filename);
    if (!infile) {
        fprintf(stderr, "ERROR: Could not open the tree file: %s!\n", filename.c_str());
        exit(1);
    }
    std::string newick_string;
    std::getline(infile, newick_string);

    return create_tree_from_newick_string(newick_string);
}

Mutation_Annotated_Tree::Tree Mutation_Annotated_Tree::load_mutation_annotated_tree (std::string filename) {
    TIMEIT();
    Tree tree;

    Parsimony::data data;

    std::ifstream inpfile(filename, std::ios::in | std::ios::binary);
    if (!inpfile) {
        fprintf(stderr, "ERROR: Could not load the mutation-annotated tree object from file: %s!\n", filename.c_str());
        exit(1);
    }

    // Boost library used to stream the contents of the input protobuf file in
    // uncompressed or compressed .gz format
    if (filename.find(".gz\0") != std::string::npos) {
        boost::iostreams::filtering_istream instream;
        try {
            instream.push(boost::iostreams::gzip_decompressor());
            instream.push(inpfile);
        }
        catch(const boost::iostreams::gzip_error& e) {
            std::cout << e.what() << '\n';
        }

        data.ParseFromIstream(&instream);
        inpfile.close();
    }
    else {
        data.ParseFromIstream(&inpfile);
        inpfile.close();
    }

    //check if the pb has a metadata field
    bool hasmeta = (data.metadata_size()>0);
    if (!hasmeta) {
        fprintf(stderr, "WARNING: This pb does not include any metadata. Filling in default values\n");
    }
    tree = create_tree_from_newick_string(data.newick());
    auto dfs = tree.depth_first_expansion();
    static tbb::affinity_partitioner ap;
    tbb::parallel_for( tbb::blocked_range<size_t>(0, dfs.size()),
            [&](tbb::blocked_range<size_t> r) {
            for (size_t idx = r.begin(); idx < r.end(); idx++) {
               auto node = dfs[idx];
               auto mutation_list = data.node_mutations(idx);
               if (hasmeta) {
                   for (int k = 0; k < data.metadata(idx).clade_annotations_size(); k++) {
                       node->clade_annotations.emplace_back(data.metadata(idx).clade_annotations(k));
                   }
               }
               for (int k = 0; k < mutation_list.mutation_size(); k++) {
                  auto mut = mutation_list.mutation(k);
                  Mutation m;
                  m.chrom = mut.chromosome();
                  m.position = mut.position();
                  if (!m.is_masked()) {
                     m.ref_nuc = (1 << mut.ref_nuc());
                     m.par_nuc = (1 << mut.par_nuc());
                     m.is_missing = false;
                     std::vector<int8_t> nuc_vec(mut.mut_nuc_size());
                     for (int n = 0; n < mut.mut_nuc_size(); n++) {
                        nuc_vec[n] = mut.mut_nuc(n);
                     }
                     m.mut_nuc = get_nuc_id(nuc_vec);
                     if (m.mut_nuc != m.par_nuc) {
                         node->add_mutation(m);
                     }
                  }
                  else {
                      // Mutation masked
                      m.ref_nuc = 0;
                      m.par_nuc = 0;
                      m.mut_nuc = 0;
                      node->add_mutation(m);
                  }
               }
               if (!std::is_sorted(node->mutations.begin(), node->mutations.end())) {
                   fprintf(stderr, "WARNING: Mutations not sorted!\n");
                   std::sort(node->mutations.begin(), node->mutations.end());
               }
            }
        }, ap);

    size_t num_condensed_nodes = static_cast<size_t>(data.condensed_nodes_size());
    tbb::parallel_for( tbb::blocked_range<size_t>(0, num_condensed_nodes),
            [&](tbb::blocked_range<size_t> r) {
            for (size_t idx = r.begin(); idx < r.end(); idx++) {
               auto cn = data.condensed_nodes(idx);
               tree.condensed_nodes.emplace(std::pair<std::string, std::vector<std::string>>(cn.node_name(), std::vector<std::string>(cn.condensed_leaves_size())));
               for (int k = 0; k < cn.condensed_leaves_size(); k++) {
                  tree.condensed_nodes[cn.node_name()][k] = cn.condensed_leaves(k);
                  tree.condensed_leaves.emplace(cn.condensed_leaves(k));
               }
            }
    }, ap);

    return tree;
}

void Mutation_Annotated_Tree::save_mutation_annotated_tree (Mutation_Annotated_Tree::Tree tree, std::string filename) {
    TIMEIT();
    Parsimony::data data;
    data.set_newick(get_newick_string(tree, false, true, true));

    auto dfs = tree.depth_first_expansion();

    for (size_t idx = 0; idx < dfs.size(); idx++) {
        auto meta = data.add_metadata();
        for (size_t k = 0; k < dfs[idx]->clade_annotations.size(); k++) {
            meta->add_clade_annotations(dfs[idx]->clade_annotations[k]);
        }
        auto mutation_list = data.add_node_mutations();
        for (auto m: dfs[idx]->mutations) {
            auto mut = mutation_list->add_mutation();
            mut->set_chromosome(m.chrom);
            mut->set_position(m.position);

            if (m.is_masked()) {
                mut->set_ref_nuc(-1);
                mut->set_par_nuc(-1);
            }
            else {
                int8_t j = get_nt(m.ref_nuc);
                assert (j >= 0);
                mut->set_ref_nuc(j);

                j = get_nt(m.par_nuc);
                assert(j >= 0);
                mut->set_par_nuc(j);

                mut->clear_mut_nuc();
                for (auto nuc: get_nuc_vec_from_id(m.mut_nuc)) {
                    mut->add_mut_nuc(nuc);
                }
            }
        }
    }

    // Add condensed nodes
    for (auto cn: tree.condensed_nodes) {
        auto cn_ptr = data.add_condensed_nodes();
        cn_ptr->set_node_name(cn.first);
        for (auto lid: cn.second) {
            cn_ptr->add_condensed_leaves(lid);
        }
    }

    // Boost library used to stream the contents to the output protobuf file in
    // uncompressed or compressed .gz format
    std::ofstream outfile(filename, std::ios::out | std::ios::binary);
    boost::iostreams::filtering_streambuf< boost::iostreams::output> outbuf;

    if (filename.find(".gz\0") != std::string::npos) {
        try {
            outbuf.push(boost::iostreams::gzip_compressor());
            outbuf.push(outfile);
            std::ostream outstream(&outbuf);
            data.SerializeToOstream(&outstream);
            boost::iostreams::close(outbuf);
            outfile.close();
        }
        catch(const boost::iostreams::gzip_error& e) {
            std::cout << e.what() << '\n';
        }
    }
    else {
        data.SerializeToOstream(&outfile);
        outfile.close();
    }
}

/* === Node === */
bool Mutation_Annotated_Tree::Node::is_leaf () {
    return (children.size() == 0);
}

bool Mutation_Annotated_Tree::Node::is_root() {
    return (parent == NULL);
}

Mutation_Annotated_Tree::Node::Node() {
    level = 0;
    identifier = "";
    parent = NULL;
    branch_length = -1.0;
    clade_annotations.clear();
    mutations.clear();
}

Mutation_Annotated_Tree::Node::Node (std::string id, float len) {
    identifier = id;
    parent = NULL;
    level = 1;
    branch_length = len;
    mutations.clear();
}

Mutation_Annotated_Tree::Node::Node (std::string id, Node* p, float len) {
    identifier = id;
    parent = p;
    level = p->level + 1;
    branch_length = len;
    mutations.clear();
}

// Assumes mutations are added in chronological order. If a new mutation occurs
// at the same position, it should either be updated to the new allele or
// removed entirely (in case of reversal mutation)
void Mutation_Annotated_Tree::Node::add_mutation (Mutation mut) {
    auto iter = std::lower_bound(mutations.begin(), mutations.end(), mut);
    // check if mutation at the same position has occured before
    if ((iter != mutations.end()) && (iter->position == mut.position)) {
        // update to new allele
        if (iter->par_nuc != mut.mut_nuc) {
            iter->mut_nuc = mut.mut_nuc;
        }
        //reversal mutation
        else {
            std::vector<Mutation> tmp;
            for (auto m: mutations) {
                if (m.position != iter->position) {
                    tmp.emplace_back(m.copy());
                }
            }
            mutations.clear();
            for (auto m: tmp) {
                mutations.emplace_back(m.copy());
            }
        }
    }
    // new mutation
    else {
        mutations.insert(iter, mut);
    }
}

void Mutation_Annotated_Tree::Node::clear_mutations() {
    mutations.clear();
}

void Mutation_Annotated_Tree::Node::clear_annotations() {
    clade_annotations.clear();
}

/* === Tree === */
size_t Mutation_Annotated_Tree::Tree::get_max_level () const {
    size_t max_level = 0;
    for (auto x: all_nodes) {
        if (x.second->level > max_level) {
            max_level = x.second->level;
        }
    }
    return max_level;
}

size_t Mutation_Annotated_Tree::Tree::get_num_annotations () const {
    size_t ret = 0;
    if (root != NULL) {
        ret = root->clade_annotations.size();
    }
    return ret;
}

void Mutation_Annotated_Tree::Tree::rename_node(std::string old_nid, std::string new_nid) {
    auto n = get_node(old_nid);
    if (n != NULL) {
        n->identifier = new_nid;
        all_nodes.erase(old_nid);
        all_nodes[new_nid] = n;
    }
    else {
        fprintf(stderr, "ERROR: %s not found in the Tree!\n", old_nid.c_str());
        exit(1);
    }
}

std::vector<Mutation_Annotated_Tree::Node*> Mutation_Annotated_Tree::Tree::get_leaves(std::string nid) {
    std::vector<Node*> leaves;
    if (nid == "") {
        if (root == NULL) {
            return leaves;
        }
        nid = root->identifier;
    }
    Node* node = all_nodes[nid];

    std::queue<Node*> remaining_nodes;
    remaining_nodes.push(node);
    while (remaining_nodes.size() > 0) {
        Node* curr_node = remaining_nodes.front();
        if (curr_node->children.size() == 0)
            leaves.push_back(curr_node);
        remaining_nodes.pop();
        for (auto c: curr_node->children) {
            remaining_nodes.push(c);
        }
    }
    return leaves;
}

std::vector<std::string> Mutation_Annotated_Tree::Tree::get_leaves_ids(std::string nid) {
    std::vector<std::string> leaves_ids;
    if (nid == "") {
        if (root == NULL) {
            return leaves_ids;
        }
        nid = root->identifier;
    }
    Node* node = all_nodes[nid];

    std::queue<Node*> remaining_nodes;
    remaining_nodes.push(node);
    while (remaining_nodes.size() > 0) {
        Node* curr_node = remaining_nodes.front();
        if (curr_node->children.size() == 0)
            leaves_ids.push_back(curr_node->identifier);
        remaining_nodes.pop();
        for (auto c: curr_node->children) {
            remaining_nodes.push(c);
        }
    }
    return leaves_ids;
}

size_t Mutation_Annotated_Tree::Tree::get_num_leaves(Node* node) {
    if (node == NULL) {
        node = root;
    }

    if (node->is_leaf()) {
        return 1;
    }
    size_t num_leaves = 0;
    for (auto c: node->children) {
        num_leaves += get_num_leaves(c);
    }
    return num_leaves;
}

Mutation_Annotated_Tree::Node* Mutation_Annotated_Tree::Tree::create_node (std::string const& identifier, float branch_len, size_t num_annotations) {
    all_nodes.clear();
    Node* n = new Node(identifier, branch_len);
    for (size_t k=0; k < num_annotations; k++) {
        n->clade_annotations.emplace_back("");
    }
    root = n;
    all_nodes[identifier] = root;
    return n;
}

Mutation_Annotated_Tree::Node* Mutation_Annotated_Tree::Tree::create_node (std::string const& identifier, Node* par, float branch_len) {
    if (all_nodes.find(identifier) != all_nodes.end()) {
        fprintf(stderr, "Error: %s already in the tree!\n", identifier.c_str());
        exit(1);
    }
    Node* n = new Node(identifier, par, branch_len);
    size_t num_annotations = get_num_annotations();
    for (size_t k=0; k < num_annotations; k++) {
        n->clade_annotations.emplace_back("");
    }
    all_nodes[identifier] = n;
    par->children.push_back(n);
    return n;
}

Mutation_Annotated_Tree::Node* Mutation_Annotated_Tree::Tree::create_node (std::string const& identifier, std::string const& parent_id, float branch_len) {
    Node* par = all_nodes[parent_id];
    return create_node(identifier, par, branch_len);
}

Mutation_Annotated_Tree::Node* Mutation_Annotated_Tree::Tree::get_node (std::string nid) const {
    if (all_nodes.find(nid) != all_nodes.end()) {
        return all_nodes.at(nid);
    }
    return NULL;

}

bool Mutation_Annotated_Tree::Tree::is_ancestor (std::string anc_id, std::string nid) const {
    Node* node = get_node(nid);
    while (node->parent != NULL) {
        node = node->parent;
        if (node->identifier == anc_id) {
            return true;
        }
    }
    return false;
}

std::vector<Mutation_Annotated_Tree::Node*> Mutation_Annotated_Tree::Tree::rsearch (const std::string& nid, bool include_self) const {
    std::vector<Node*> ancestors;
    Node* node = get_node(nid);
    if (node==NULL) {
        return ancestors;
    }
    if (include_self) {
        ancestors.push_back(node);
    }
    while (node->parent != NULL) {
        ancestors.push_back(node->parent);
        node = node->parent;
    }
    return ancestors;
}

void Mutation_Annotated_Tree::Tree::remove_node_helper (std::string nid, bool move_level) {
    auto it = all_nodes.find(nid);
    if (it == all_nodes.end()) {
        fprintf(stderr, "ERROR: Tried to remove node identifier %s but it was not found!\n", nid.c_str());
        exit(1);
    }
    Node* source = it->second;
    Node* curr_parent = source->parent;

    if (curr_parent != NULL) {
        // Remove source from curr_parent
        auto iter = std::find(curr_parent->children.begin(), curr_parent->children.end(), source);
        assert (iter != curr_parent->children.end());
        curr_parent->children.erase(iter);

        // Remove parent if it no longer has any children
        if (curr_parent->children.size() == 0) {
            if (curr_parent == root) {
                fprintf(stderr, "ERROR: Tree empty!\n");
                exit(1);
            }
            remove_node_helper (curr_parent->identifier, move_level);
        }
        // Move the remaining child one level up if it is the only child of its parent
        else if (move_level && (curr_parent->children.size() == 1)) {
            auto child = curr_parent->children[0];
            if (curr_parent->parent != NULL) {
                for (size_t k=0; k < curr_parent->clade_annotations.size(); k++) {
                    if (child->clade_annotations[k] == "") {
                        child->clade_annotations[k] = curr_parent->clade_annotations[k];
                    }
                }
                child->parent = curr_parent->parent;
                child->level = curr_parent->parent->level + 1;
                child->branch_length += curr_parent->branch_length;

                std::vector<Mutation> tmp;
                for (auto m: child->mutations) {
                    tmp.emplace_back(m.copy());
                }

                //Clear and add back mutations in chrono order
                child->clear_mutations();
                for (auto m: curr_parent->mutations) {
                    child->add_mutation(m.copy());
                }
                for (auto m: tmp) {
                    child->add_mutation(m.copy());
                }

                curr_parent->parent->children.push_back(child);

                iter = std::find(curr_parent->parent->children.begin(), curr_parent->parent->children.end(), curr_parent);
                assert(iter != curr_parent->parent->children.end());
                curr_parent->parent->children.erase(iter);

                // Update levels of source descendants
                std::queue<Node*> remaining_nodes;
                remaining_nodes.push(child);
                while (remaining_nodes.size() > 0) {
                    Node* curr_node = remaining_nodes.front();
                    remaining_nodes.pop();
                    curr_node->level = curr_node->parent->level + 1;
                    for (auto c: curr_node->children) {
                        remaining_nodes.push(c);
                    }
                }
            }

            auto par_it = all_nodes.find(curr_parent->identifier);
            assert (par_it != all_nodes.end());
            all_nodes.erase(par_it);
            delete curr_parent;
        }
    }

    //Remove source and descendants from all_nodes
    std::queue<Node*> desc;
    desc.push(source);
    while (desc.size() > 0) {
        Node* curr_node = desc.front();
        desc.pop();
        for (auto c: curr_node->children) {
            desc.push(c);
        }
        it = all_nodes.find(curr_node->identifier);
        all_nodes.erase(it);
        delete curr_node;
    }
}

void Mutation_Annotated_Tree::Tree::remove_node (std::string nid, bool move_level) {
    TIMEIT();
    remove_node_helper (nid, move_level);
}

void Mutation_Annotated_Tree::Tree::move_node (std::string source_id, std::string dest_id, bool move_level) {
    Node* source = all_nodes[source_id];
    Node* destination = all_nodes[dest_id];
    Node* curr_parent = source->parent;

    source->parent = destination;
    source->branch_length = -1.0; // Invalidate source branch length

    destination->children.push_back(source);

    // Remove source from curr_parent
    auto iter = std::find(curr_parent->children.begin(), curr_parent->children.end(), source);
    curr_parent->children.erase(iter);
    if (curr_parent->children.size() == 0) {
        remove_node(curr_parent->identifier, move_level);
    }

    // Update levels of source descendants
    std::queue<Node*> remaining_nodes;
    remaining_nodes.push(source);
    while (remaining_nodes.size() > 0) {
        Node* curr_node = remaining_nodes.front();
        remaining_nodes.pop();
        curr_node->level = curr_node->parent->level + 1;
        for (auto c: curr_node->children) {
            remaining_nodes.push(c);
        }
    }
}

std::vector<Mutation_Annotated_Tree::Node*> Mutation_Annotated_Tree::Tree::breadth_first_expansion(std::string nid) {
    std::vector<Node*> traversal;

    if (nid == "") {
        if (root == NULL) {
            return traversal;
        }
        nid = root->identifier;
    }

    Node* node = all_nodes[nid];

    std::queue<Node*> remaining_nodes;
    remaining_nodes.push(node);
    while (remaining_nodes.size() > 0) {
        Node* curr_node = remaining_nodes.front();
        traversal.push_back(curr_node);
        remaining_nodes.pop();
        for (auto c: curr_node->children) {
            remaining_nodes.push(c);
        }
    }

    return traversal;
}

void Mutation_Annotated_Tree::Tree::depth_first_expansion_helper(Mutation_Annotated_Tree::Node* node, std::vector<Mutation_Annotated_Tree::Node*>& vec) const {
    vec.push_back(node);
    for (auto c: node->children) {
        depth_first_expansion_helper(c, vec);
    }
}

std::vector<Mutation_Annotated_Tree::Node*> Mutation_Annotated_Tree::Tree::depth_first_expansion(Mutation_Annotated_Tree::Node* node) const {
    TIMEIT();
    std::vector<Node*> traversal;
    if (node == NULL) {
        node = root;
    }
    if (node == NULL) {
        return traversal;
    }
    depth_first_expansion_helper(node, traversal);
    return traversal;
}

size_t Mutation_Annotated_Tree::Tree::get_parsimony_score() {
    size_t score = 0;
    auto dfs = depth_first_expansion();
    for (auto n: dfs) {
        score += n->mutations.size();
    }
    return score;
}

void Mutation_Annotated_Tree::Tree::condense_leaves(std::vector<std::string> missing_samples) {
    if (condensed_nodes.size() > 0) {
        fprintf(stderr, "WARNING: tree contains condensed nodes. It may be condensed already!\n");
    }

    auto tree_leaves = get_leaves_ids();
    for (auto l1_id: tree_leaves) {
        std::vector<Node*> polytomy_nodes;

        auto l1 = get_node(l1_id);
        if (l1 == NULL) {
            continue;
        }
        if (std::find(missing_samples.begin(), missing_samples.end(), l1->identifier) != missing_samples.end()) {
            continue;
        }
        if (l1->mutations.size() > 0) {
            continue;
        }

        for (auto l2: l1->parent->children) {
                if (std::find(missing_samples.begin(), missing_samples.end(), l2->identifier) != missing_samples.end()) {
                    continue;
                }
            if (l2->is_leaf() && (get_node(l2->identifier) != NULL) && (l2->mutations.size() == 0)) {
                polytomy_nodes.push_back(l2);
            }
        }
        if (polytomy_nodes.size() > 1) {
            std::string new_node_name = "node_" + std::to_string(1+condensed_nodes.size()) + "_condensed_" + std::to_string(polytomy_nodes.size()) + "_leaves";

            auto curr_node = get_node(l1->identifier);
            auto new_node = create_node(new_node_name, curr_node->parent, l1->branch_length);

            new_node->clear_mutations();

            condensed_nodes[new_node_name] = std::vector<std::string>(polytomy_nodes.size());

            for (size_t it = 0; it < polytomy_nodes.size(); it++) {
                condensed_nodes[new_node_name][it] = polytomy_nodes[it]->identifier;
                remove_node(polytomy_nodes[it]->identifier, false);
            }
        }
    }
}

void Mutation_Annotated_Tree::Tree::uncondense_leaves() {
    tbb::mutex tbb_lock;
    static tbb::affinity_partitioner ap;
    tbb::parallel_for(tbb::blocked_range<size_t>(0, condensed_nodes.size()),
            [&](tbb::blocked_range<size_t> r) {
            for (size_t it = r.begin(); it < r.end(); it++) {
                auto cn = condensed_nodes.begin();
                std::advance(cn, it);

                tbb_lock.lock();
                auto n = get_node(cn->first);
                tbb_lock.unlock();
                auto par = (n->parent != NULL) ? n->parent : n;

                size_t num_samples = cn->second.size();

                if (num_samples > 0) {
                    tbb_lock.lock();
                    rename_node(n->identifier, cn->second[0]);
                    tbb_lock.unlock();
                }

                for (size_t s = 1; s < num_samples; s++) {
                    tbb_lock.lock();
                    auto new_n = create_node(cn->second[s], par, n->branch_length);
                    tbb_lock.unlock();
                    for (auto m: n->mutations) {
                        new_n->add_mutation(m.copy());
                    }
                }
            }
        }, ap);
    condensed_nodes.clear();
    condensed_leaves.clear();
}

void Mutation_Annotated_Tree::Tree::collapse_tree() {
    auto bfs = breadth_first_expansion();

    for (size_t idx = 1; idx < bfs.size(); idx++) {
        auto node = bfs[idx];
        auto mutations = node->mutations;
        if (mutations.size() == 0) {
            auto parent = node->parent;
            auto children = node->children;
            for (auto child: children) {
                move_node(child->identifier, parent->identifier, false);
            }
        }
        //If internal node has one child, the child can be moved up one level
        else if (node->children.size() == 1) {
            auto child = node->children.front();
            auto parent = node->parent;
            for (auto m: mutations) {
                child->add_mutation(m.copy());
            }
            move_node(child->identifier, parent->identifier, false);
        }
    }
}

Mutation_Annotated_Tree::Tree Mutation_Annotated_Tree::get_tree_copy(const Mutation_Annotated_Tree::Tree& tree, const std::string& identifier) {
    TIMEIT();
    auto root = tree.root;
    if (identifier != "") {
        root = tree.get_node(identifier);
    }

    Tree copy = create_tree_from_newick_string (get_newick_string(tree, root, true, true));

    std::vector<Node*> dfs1;
    std::vector<Node*> dfs2;

    static tbb::affinity_partitioner ap;
    tbb::parallel_for(tbb::blocked_range<size_t>(0, 2),
            [&](tbb::blocked_range<size_t> r) {
            for (size_t k=r.begin(); k<r.end(); ++k){
              if (k==0) {
                dfs1 = tree.depth_first_expansion(root);
              }
              else {
                dfs2 = copy.depth_first_expansion();
              }
            }
            }, ap);

    tbb::parallel_for(tbb::blocked_range<size_t>(0, dfs1.size()),
            [&](tbb::blocked_range<size_t> r) {
            for (size_t k=r.begin(); k<r.end(); ++k){
              auto n1 = dfs1[k];
              auto n2 = dfs2[k];
              n2->clade_annotations.resize(n1->clade_annotations.size());
              for (size_t i=0; i<n1->clade_annotations.size(); i++) {
                 n2->clade_annotations[i] = n1->clade_annotations[i];
              }
              for (auto m: n1->mutations) {
                Mutation m2 = m.copy();
                n2->add_mutation(m2);
                }
              }
            }, ap);

    size_t num_condensed_nodes = static_cast<size_t>(tree.condensed_nodes.size());
    tbb::parallel_for( tbb::blocked_range<size_t>(0, num_condensed_nodes),
            [&](tbb::blocked_range<size_t> r) {
            for (size_t idx = r.begin(); idx < r.end(); idx++) {
               auto cn = tree.condensed_nodes.begin();
               std::advance(cn, idx);
               copy.condensed_nodes.insert(std::pair<std::string, std::vector<std::string>>(cn->first, std::vector<std::string>(cn->second.size())));
               for (size_t k = 0; k < cn->second.size(); k++) {
                  copy.condensed_nodes[cn->first][k] = cn->second[k];
                  copy.condensed_leaves.insert(cn->second[k]);
               }
            }
    }, ap);

    return copy;
}

// Get the last common ancestor of two node identifiers. Return NULL if does not
// exist
Mutation_Annotated_Tree::Node* Mutation_Annotated_Tree::LCA (const Mutation_Annotated_Tree::Tree& tree, const std::string& nid1, const std::string& nid2) {
    TIMEIT();

    if ((tree.get_node(nid1) == NULL) || (tree.get_node(nid2) == NULL)) {
        return NULL;
    }

    auto n2_ancestors = tree.rsearch(nid2, true);

    for (auto anc1: tree.rsearch(nid1, true)) {
        for (auto anc2: n2_ancestors) {
            if (anc1 == anc2) {
                return anc1;
            }
        }
    }

    return NULL;
}

// Extract the subtree consisting of the specified set of samples. This routine
// maintains the internal node names of the input tree. Mutations are copied
// from the tree such that the path of mutations from root to the sample is
// same as the original tree.
Mutation_Annotated_Tree::Tree Mutation_Annotated_Tree::get_subtree (const Mutation_Annotated_Tree::Tree& tree, const std::vector<std::string>& samples) {
    TIMEIT();
    Tree subtree;

    // Set of leaf and internal nodes corresponding to the subtree
    tbb::concurrent_unordered_set<Node*> subtree_nodes;
    // Maintain a set of all ancestors of a sample for each sample
    std::vector<tbb::concurrent_unordered_set<Node*>> all_ancestors(samples.size());

    static tbb::affinity_partitioner ap;
    tbb::parallel_for(tbb::blocked_range<size_t>(0, samples.size()),
            [&](tbb::blocked_range<size_t> r) {
            for (size_t k=r.begin(); k<r.end(); ++k){
               subtree_nodes.insert(tree.get_node(samples[k]));
               for (auto anc: tree.rsearch(samples[k], true)) {
                   all_ancestors[k].insert(anc);
               }
            }
    }, ap);

    tbb::parallel_for(tbb::blocked_range<size_t>(0, samples.size()),
            [&](tbb::blocked_range<size_t> r) {
            for (size_t i=r.begin(); i<r.end(); ++i){
               for (size_t j=i+1; j<samples.size(); ++j){
                   for (auto anc: tree.rsearch(samples[i], true)) {
                      if (all_ancestors[j].find(anc) != all_ancestors[j].end()) {
                         subtree_nodes.insert(anc);
                         break;
                      }
                   }
               }
            }
    }, ap);

    auto dfs = tree.depth_first_expansion();
    size_t num_annotations = tree.get_num_annotations();

    std::stack<Node*> last_subtree_node;
    for (auto n: dfs) {
        // If the node is in subtree_nodes, it should be added to the subtree
        if (subtree_nodes.find(n) != subtree_nodes.end()) {
            Node* subtree_parent = NULL;
            if (last_subtree_node.size() > 0) {
                while (!tree.is_ancestor(last_subtree_node.top()->identifier, n->identifier)) {
                    last_subtree_node.pop();
                }
                subtree_parent = last_subtree_node.top();
            }
            // Add as root of the subtree
            if (subtree_parent == NULL) {
                // for root node, need to size the annotations vector
                Node* new_node = subtree.create_node(n->identifier, -1.0, num_annotations);
                // need to assign any clade annotations which would belong to that root as well
                for (size_t k = 0; k < num_annotations; k++) {
                    if (n->clade_annotations[k] != "") {
                        new_node->clade_annotations[k] = n->clade_annotations[k];
                    }
                }
                std::vector<Node*> root_to_node = tree.rsearch(n->identifier, true);
                std::reverse(root_to_node.begin(), root_to_node.end());
                root_to_node.emplace_back(n);

                for (auto curr: root_to_node) {
                    for (auto m: curr->mutations) {
                        new_node->add_mutation(m);
                    }
                }
            }
            // Add to the parent identified
            else {
                Node* new_node = subtree.create_node(n->identifier, subtree_parent->identifier);

                auto par_to_node = tree.rsearch(n->identifier, true);
                std::reverse(par_to_node.begin(), par_to_node.end());
                par_to_node.erase(par_to_node.begin(), std::find(par_to_node.begin(), par_to_node.end(), subtree_parent)+1);

                for (auto curr: par_to_node) {
                    for (size_t k = 0; k < num_annotations; k++) {
                        if (curr->clade_annotations[k] != "") {
                            new_node->clade_annotations[k] = curr->clade_annotations[k];
                        }
                    }
                    for (auto m: curr->mutations) {
                        new_node->add_mutation(m);
                    }
                }
            }
            last_subtree_node.push(n);
        }
    }

    subtree.curr_internal_node = tree.curr_internal_node;

    return subtree;
}

void Mutation_Annotated_Tree::clear_tree(Mutation_Annotated_Tree::Tree& T) {
    for (auto n: T.depth_first_expansion()) {
        delete(n);
    }
}


#include <boost/algorithm/hex.hpp>
#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>

#include <nil/crypto3/zk/components/blueprint.hpp>
#include <nil/crypto3/zk/components/blueprint_variable.hpp>

#include <nil/crypto3/zk/snark/schemes/ppzksnark/r1cs_gg_ppzksnark.hpp>
#include <nil/crypto3/zk/snark/schemes/ppzksnark/r1cs_gg_ppzksnark/marshalling.hpp>

#include <nil/crypto3/zk/snark/algorithms/generate.hpp>
#include <nil/crypto3/zk/snark/algorithms/verify.hpp>
#include <nil/crypto3/zk/snark/algorithms/prove.hpp>

#include "./comparerLogic.h"

using namespace std;

using namespace nil::crypto3::zk::components;
using namespace nil::crypto3::zk::snark;

ushort INVALID_PROOF_RETURN_CODE = 200;

std::string convert_byteblob_to_hex_string(std::vector<std::uint8_t> blob) {
    // convert byte_blob to hex string and print it to output
    std::string hex;
    hex.reserve(blob.size() * 2);
    boost::algorithm::hex(blob.begin(), blob.end(), back_inserter(hex));
    return hex;
}

void save_byteblob(std::vector<std::uint8_t> byteblob, boost::filesystem::path fname) {
    boost::filesystem::ofstream out(fname);
    for (const auto &v : byteblob) {
        out << v;
    }
    out.close();
}

std::vector<std::uint8_t> load_byteblob(boost::filesystem::path fname) {
    boost::filesystem::ifstream stream(fname, std::ios::in | std::ios::binary);
    std::vector<std::uint8_t> contents((std::istreambuf_iterator<char>(stream)), std::istreambuf_iterator<char>());
    if (contents.size() == 0) {
        throw std::ios_base::failure("Empty file");
    }
    return contents;
}

// proving key

void save_proving_key(scheme_type::proving_key_type pk, boost::filesystem::path fname) {
    std::vector<std::uint8_t> byteblob = nil::marshalling::verifier_input_serializer_tvm<scheme_type>::process(pk);
    save_byteblob(byteblob, fname);
}

scheme_type::proving_key_type load_proving_key(boost::filesystem::path fname) {
    std::vector<std::uint8_t> byteblob = load_byteblob(fname);
    nil::marshalling::status_type processingStatus = nil::marshalling::status_type::success;
    return nil::marshalling::verifier_input_deserializer_tvm<scheme_type>::proving_key_process(
        byteblob.cbegin(),
        byteblob.cend(),
        processingStatus);
}

// verification key

void save_verification_key(scheme_type::verification_key_type vk, boost::filesystem::path fname) {
    std::vector<std::uint8_t> byteblob = nil::marshalling::verifier_input_serializer_tvm<scheme_type>::process(vk);
    save_byteblob(byteblob, fname);
}

scheme_type::verification_key_type load_verification_key(boost::filesystem::path fname) {
    std::vector<std::uint8_t> byteblob = load_byteblob(fname);
    nil::marshalling::status_type processingStatus = nil::marshalling::status_type::success;
    return nil::marshalling::verifier_input_deserializer_tvm<scheme_type>::verification_key_process(
        byteblob.cbegin(),
        byteblob.cend(),
        processingStatus);
}

// proof

void save_proof(scheme_type::proof_type proof, boost::filesystem::path fname) {
    std::vector<std::uint8_t> byteblob = nil::marshalling::verifier_input_serializer_tvm<scheme_type>::process(proof);
    save_byteblob(byteblob, fname);
}

scheme_type::proof_type load_proof(boost::filesystem::path fname) {
    std::vector<std::uint8_t> byteblob = load_byteblob(fname);
    nil::marshalling::status_type processingStatus = nil::marshalling::status_type::success;
    return nil::marshalling::verifier_input_deserializer_tvm<scheme_type>::proof_process(
        byteblob.cbegin(),
        byteblob.cend(),
        processingStatus);
}

//  primary input

void save_primary_input(zk::snark::r1cs_primary_input<field_type> primary_input, boost::filesystem::path fname) {
    std::vector<std::uint8_t> byteblob = nil::marshalling::verifier_input_serializer_tvm<scheme_type>::process(primary_input);
    save_byteblob(byteblob, fname);
}

zk::snark::r1cs_primary_input<field_type> load_primary_input(boost::filesystem::path fname) {
    std::vector<std::uint8_t> byteblob = load_byteblob(fname);
    nil::marshalling::status_type processingStatus = nil::marshalling::status_type::success;
    return nil::marshalling::verifier_input_deserializer_tvm<scheme_type>::primary_input_process(
        byteblob.cbegin(),
        byteblob.cend(),
        processingStatus);
}

int setup_keys(boost::filesystem::path pk_path, boost::filesystem::path vk_path) {

    blueprint<field_type> bp;
    ComparerLogic comparerLogic(bp);
    comparerLogic.generate_r1cs_constraints(bp);

    cout << "Blueprint size: " << bp.num_variables() << endl;
    cout << "Generating constraint system..." << endl;
    const r1cs_constraint_system<field_type> constraint_system = bp.get_constraint_system();
    cout << "Number of R1CS constraints: " << constraint_system.num_constraints() << endl;

    cout << "Generating keypair..." << endl;
    scheme_type::keypair_type keypair = generate<scheme_type>(constraint_system);

    cout << "Saving proving key to a file " << pk_path<< endl;
    save_proving_key(keypair.first, pk_path);

    cout << "Saving verification key to a file " << vk_path << endl;
    save_verification_key(keypair.second, vk_path);

    return 0;
}

int create_proof(boost::filesystem::path pk_path, boost::filesystem::path proof_path, boost::filesystem::path pi_path, int minYear, int maxYear, int year) {

    cout << "Loading proving key from a file " << pk_path << endl;
    typename scheme_type::proving_key_type pk = load_proving_key(pk_path);

    blueprint<field_type> bp;
    ComparerLogic comparerLogic(bp);

    cout << "Generating constraint system..." << endl;
    comparerLogic.generate_r1cs_constraints(bp);

    cout << "Generating witness..." << endl;
    comparerLogic.generate_r1cs_witness(bp, minYear, maxYear, year);

    cout << "Blueprint is satisfied: " << bp.is_satisfied() << endl;

    if (!bp.is_satisfied()) {
        return INVALID_PROOF_RETURN_CODE;
    }

    cout << "Generating proof..." << endl;
    const scheme_type::proof_type proof = prove<scheme_type>(pk, bp.primary_input(), bp.auxiliary_input());

    cout << "Saving proof to file " << proof_path << endl;
    save_proof(proof, proof_path);

    cout << "Saving primary input to file " << pi_path << endl;
    save_primary_input(bp.primary_input(), pi_path);
    return 0;
}

int verify_proof(boost::filesystem::path proof_path, boost::filesystem::path vk_path, boost::filesystem::path pi_path) {

    cout << "Loading proof from a file " << proof_path << endl;
    typename scheme_type::proof_type proof = load_proof(proof_path);

    cout << "Loading primary input from a file " << pi_path << endl;
    r1cs_primary_input<field_type> input = load_primary_input(pi_path);

    cout << "Loading verification key from a file " << vk_path << endl;
    typename scheme_type::verification_key_type vk = load_verification_key(vk_path);

    // verify
    using basic_proof_system = r1cs_gg_ppzksnark<curve_type>;
    const bool verified = verify<basic_proof_system>(vk, input, proof);
    cout << "Verification status: " << verified << endl;

    return verified ? 0 : INVALID_PROOF_RETURN_CODE;
}

int main(int argc, char *argv[]) {
    int maxYear, minYear, year;
    boost::filesystem::path pk_path, vk_path, proof_path, pi_path;
    // bool hexFlag;

    boost::program_options::options_description options("CLI Proof Generator");
    options.add_options()
    // ("hex,h", boost::program_options::bool_switch(&hexFlag), "print only hex proof to output")
    ("minYear,minYear", boost::program_options::value<int>(&minYear)->default_value(0))
    ("maxYear,maxYear", boost::program_options::value<int>(&maxYear)->default_value(100))
    ("year,year", boost::program_options::value<int>(&year)->default_value(18))
    ("proving-key-path,pk", boost::program_options::value<boost::filesystem::path>(&pk_path)->default_value("proving.key"))
    ("verification-key-path,vk", boost::program_options::value<boost::filesystem::path>(&vk_path)->default_value("verification.key"))
    ("proof-path,p", boost::program_options::value<boost::filesystem::path>(&proof_path)->default_value("proof"))
    ("primary-input-path,pi", boost::program_options::value<boost::filesystem::path>(&pi_path)->default_value("primary.input"));

    boost::program_options::variables_map vm;
    boost::program_options::store(boost::program_options::command_line_parser(argc, argv).options(options).run(), vm);
    boost::program_options::notify(vm);

    cout << setprecision(16);

    if (!argv[1]) {
        cout << "Please select a command: [setup/prove/verify]" << endl;
        return 0;
    }
    else if (string(argv[1]) == "setup") {
        // Generate proving.key & verification.key
        return setup_keys(pk_path, vk_path);
    } else if (string(argv[1]) == "prove") {
        return create_proof(pk_path, proof_path, pi_path, minYear, maxYear, year);
    } else if (string(argv[1]) == "verify") {
        return verify_proof(proof_path, pi_path, vk_path);
    }

    return 0;
}


#define _LIBCPP_ENABLE_THREAD_SAFETY_ANNOTATIONS
#include "matching.h"

#include <boost/histogram/ostream.hpp>
#include <cstdint>
#include <fstream>
#include <gsl/gsl>
#include <iterator>
#include <opencv2/core/base.hpp>
#include <opencv2/core/mat.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/features2d.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <sens_loc/analysis/distance.h>
#include <sens_loc/io/histogram.h>
#include <sens_loc/io/image.h>
#include <sens_loc/util/console.h>
#include <sens_loc/util/thread_analysis.h>
#include <util/batch_visitor.h>
#include <util/statistic_visitor.h>

using namespace cv;
using namespace std;
using namespace gsl;

namespace {

struct descriptor_stat_data {
    descriptor_stat_data() = default;

    void insert_matches(gsl::span<DMatch> matches,
                        int               descriptor_count) noexcept {
        lock_guard l{_mutex};
        transform(begin(matches), end(matches),
                  back_inserter(_global_minimal_distances),
                  [](const DMatch& m) { return m.distance; });
        _total_descriptors += descriptor_count;
    }

    pair<vector<float>, int64_t> extract() noexcept {
        lock_guard                   l{_mutex};
        pair<vector<float>, int64_t> p{move(_global_minimal_distances),
                                       _total_descriptors};
        _global_minimal_distances = vector<float>();
        _total_descriptors        = 0UL;
        return p;
    }

  private:
    mutex                                   _mutex;
    vector<float> _global_minimal_distances GUARDED_BY(_mutex);
    int64_t _total_descriptors              GUARDED_BY(_mutex) = 0L;
};

class matching {
  public:
    matching(descriptor_stat_data& accumulated_data,
             NormTypes             norm_to_use,
             bool                  crosscheck,
             string_view           input_pattern,
             optional<string_view> output_pattern,
             optional<string_view> original_files) noexcept
        : accumulated_data{accumulated_data}
        , matcher{BFMatcher::create(norm_to_use, crosscheck)}
        , input_pattern{input_pattern}
        , output_pattern{output_pattern}
        , original_images{original_files} {
        // XOR is true if both operands have the same value.
        Expects(!(output_pattern.has_value() ^ original_images.has_value()) &&
                "Either both or none are set");
    }

    void operator()(int idx,
                    optional<vector<KeyPoint>> /*keypoints*/,  // NOLINT
                    optional<Mat> descriptors) noexcept {
        Expects(descriptors.has_value());
        if (descriptors->rows == 0)
            return;

        try {
            const int         previous_idx = idx - 1;
            const FileStorage previous_img = sens_loc::io::open_feature_file(
                fmt::format(input_pattern, previous_idx));
            Mat previous_descriptors =
                sens_loc::io::load_descriptors(previous_img);

            vector<DMatch> matches;
            matcher->match(*descriptors, previous_descriptors, matches);
            accumulated_data.insert_matches(matches, descriptors->rows);

            // Plot the matching between the descriptors of the previous and the
            // current frame.
            if (output_pattern) {
                const vector<KeyPoint> previous_keypoints =
                    sens_loc::io::load_keypoints(previous_img);

                const FileStorage this_feature =
                    sens_loc::io::open_feature_file(
                        fmt::format(input_pattern, idx));
                const vector<KeyPoint> this_keypoints =
                    sens_loc::io::load_keypoints(this_feature);

                const string img_p1 = fmt::format(*original_images, idx - 1);
                const string img_p2 = fmt::format(*original_images, idx);
                auto         img1   = sens_loc::io::load_as_8bit_gray(img_p1);
                auto         img2   = sens_loc::io::load_as_8bit_gray(img_p2);

                if (!img1 || !img2)
                    return;

                Mat out_img;
                drawMatches(img2->data(), this_keypoints, img1->data(),
                            previous_keypoints, matches, out_img,
                            Scalar(0, 0, 255), Scalar(255, 0, 0));

                const string output = fmt::format(*output_pattern, idx);
                imwrite(output, out_img);
            }
        } catch (...) {
            std::cerr << sens_loc::util::err{}
                      << "Could not initialize data for idx: " << idx << "\n";
            return;
        }
    }

    size_t postprocess(const optional<string>& stat_file,
                       const optional<string>& matched_distance_histo) {
        auto [distances, total_descriptors] = accumulated_data.extract();
        if (distances.empty())
            return 0UL;

        sort(begin(distances), end(distances));
        const auto                   dist_bins = 25;
        sens_loc::analysis::distance distance_stat{distances, dist_bins};

        if (stat_file) {
            cv::FileStorage stat_out{*stat_file,
                                     cv::FileStorage::WRITE |
                                         cv::FileStorage::FORMAT_YAML};
            stat_out.writeComment(
                "The following values contain the results of the statistical "
                "analysis for descriptor distance to the closest descriptor "
                "after matching");
            write(stat_out, "match_distance", distance_stat.get_statistic());
            stat_out.release();
        } else {
            cout << "==== Match Distances\n"
                 << "total count:    " << total_descriptors << "\n"
                 << "matched count:  " << distances.size() << "\n"
                 << "matched/total:  "
                 << narrow_cast<double>(distances.size()) /
                        narrow_cast<double>(total_descriptors)
                 << "\n"
                 << "min:            " << distance_stat.min() << "\n"
                 << "max:            " << distance_stat.max() << "\n"
                 << "median:         " << distance_stat.median() << "\n"
                 << "mean:           " << distance_stat.mean() << "\n"
                 << "Variance:       " << distance_stat.variance() << "\n"
                 << "StdDev:         " << distance_stat.stddev() << "\n"
                 << "Skewness:       " << distance_stat.skewness() << "\n";
        }
        if (matched_distance_histo) {
            std::ofstream gnuplot_data{*matched_distance_histo};
            gnuplot_data << sens_loc::io::to_gnuplot(distance_stat.histogram())
                         << std::endl;
        } else {
            cout << distance_stat.histogram() << "\n";
        }
        return distances.size();
    }

  private:
    descriptor_stat_data& accumulated_data;

    Ptr<BFMatcher>        matcher;
    string_view           input_pattern;
    optional<string_view> output_pattern;
    optional<string_view> original_images;
};
}  // namespace

namespace sens_loc::apps {
int analyze_matching(util::processing_input       in,
                     NormTypes                    norm_to_use,
                     bool                         crosscheck,
                     const optional<string>&      stat_file,
                     const optional<string>&      matched_distance_histo,
                     const optional<string_view>& output_pattern,
                     const optional<string_view>& original_files) {
    Expects(in.start < in.end && "Matching requires at least 2 images");
    using visitor = statistic_visitor<matching, required_data::descriptors>;
    descriptor_stat_data data;
    auto analysis_v = visitor{/*input_pattern=*/in.input_pattern,
                              /*accumulated_data=*/data,
                              /*norm_to_use=*/norm_to_use,
                              /*crosscheck=*/crosscheck,
                              /*input_pattern=*/in.input_pattern,
                              /*output_pattern=*/output_pattern,
                              /*original_files=*/original_files};

    auto f = parallel_visitation(
        in.start + 1,  // Because two consecutive images are matched, the first
                       // index is skipped. This requires "backwards" matching.
        in.end, analysis_v);
    size_t n_elements = f.postprocess(stat_file, matched_distance_histo);

    return n_elements > 0UL ? 0 : 1;
}
}  // namespace sens_loc::apps


/**
 *****************************************************************************
 * @author     This file is part of tinyram_snark, developed by SCIPR Lab
 *             and contributors (see AUTHORS).
 * @copyright  MIT license (see LICENSE file)
 *****************************************************************************/
#include <fstream>
#include <iostream>
#ifndef MINDEPS
#include <boost/program_options.hpp>
#endif

#include <tinyram_snark/common/default_types/tinyram_ppzksnark_pp.hpp>
#include <tinyram_snark/relations/ram_computations/rams/tinyram/tinyram_params.hpp>
#include <tinyram_snark/zk_proof_systems/ppzksnark/ram_ppzksnark/ram_ppzksnark.hpp>

#ifndef MINDEPS
namespace po = boost::program_options;

bool process_prover_command_line(const int argc, const char** argv,
                                 std::string &processed_assembly_fn,
                                 std::string &proving_key_fn,
                                 std::string &primary_input_fn,
                                 std::string &auxiliary_input_fn,
                                 std::string &proof_fn)
{
    try
    {
        po::options_description desc("Usage");
        desc.add_options()
            ("help", "print this help message")
            ("processed_assembly", po::value<std::string>(&processed_assembly_fn)->required())
            ("proving_key", po::value<std::string>(&proving_key_fn)->required())
            ("primary_input", po::value<std::string>(&primary_input_fn)->required())
            ("auxiliary_input", po::value<std::string>(&auxiliary_input_fn)->required())
            ("proof", po::value<std::string>(&proof_fn)->required());

        po::variables_map vm;
        po::store(po::parse_command_line(argc, argv, desc), vm);

        if (vm.count("help"))
        {
            std::cout << desc << "\n";
            return false;
        }

        po::notify(vm);
    }
    catch(std::exception& e)
    {
        std::cerr << "Error: " << e.what() << "\n";
        return false;
    }

    return true;
}
#endif

using namespace tinyram_snark;

int main(int argc, const char * argv[])
{
    default_tinyram_ppzksnark_pp::init_public_params();

#ifdef MINDEPS
    std::string processed_assembly_fn = "processed_assembly.txt";
    std::string proving_key_fn = "proving_key.txt";
    std::string primary_input_fn = "primary_input.txt";
    std::string auxiliary_input_fn = "auxiliary_input.txt";
    std::string proof_fn = "proof.txt";
#else
    std::string processed_assembly_fn;
    std::string proving_key_fn;
    std::string primary_input_fn;
    std::string auxiliary_input_fn;
    std::string proof_fn;

    if (!process_prover_command_line(argc, argv, processed_assembly_fn,
                                     proving_key_fn, primary_input_fn, auxiliary_input_fn, proof_fn))
    {
        return 1;
    }
#endif
    libff::start_profiling();

    /* load everything */
    libff::enter_block("Deserialize proving key");
    ram_ppzksnark_proving_key<default_tinyram_ppzksnark_pp> pk;
    std::ifstream pk_file(proving_key_fn);
    pk_file >> pk;
    pk_file.close();
    libff::leave_block("Deserialize proving key");

    std::ifstream processed(processed_assembly_fn);
    tinyram_program program = load_preprocessed_program(pk.ap, processed);

    std::ifstream f_primary_input(primary_input_fn);
    std::ifstream f_auxiliary_input(auxiliary_input_fn);
    tinyram_input_tape primary_input = load_tape(f_primary_input);
    tinyram_input_tape auxiliary_input = load_tape(f_auxiliary_input);

    const ram_boot_trace<default_tinyram_ppzksnark_pp> boot_trace = tinyram_boot_trace_from_program_and_input(pk.ap, pk.primary_input_size_bound, program, primary_input);
    const ram_ppzksnark_proof<default_tinyram_ppzksnark_pp> proof = ram_ppzksnark_prover<default_tinyram_ppzksnark_pp>(pk, boot_trace,  auxiliary_input);

    libff::enter_block("Serialize proof");
    std::ofstream proof_file(proof_fn);
    proof_file << proof;
    proof_file.close();
    libff::leave_block("Serialize proof");
}


#ifndef DQN_HPP_
#define DQN_HPP_

#include <memory>
#include <random>
#include <tuple>
#include <unordered_map>
#include <vector>
#include <ale_interface.hpp>
#include <caffe/caffe.hpp>
#include <boost/functional/hash.hpp>
#include <boost/optional.hpp>

namespace dqn
{
constexpr auto kRawFrameHeight       = 210;
constexpr auto kRawFrameWidth        = 160;
constexpr auto kCroppedFrameSize     = 84;
constexpr auto kCroppedFrameDataSize = kCroppedFrameSize * kCroppedFrameSize;
constexpr auto kOutputCount          = 18;

constexpr auto last_action_layer_name = "last_action_input_layer";

constexpr auto frames_layer_name = "frames_input_layer";
constexpr auto cont_layer_name   = "cont_input_layer";
constexpr auto target_layer_name = "target_input_layer";
constexpr auto filter_layer_name = "filter_input_layer";

constexpr auto train_last_action_blob_name  = "actions";
constexpr auto test_last_action_blob_name  = "last_actions";

constexpr auto train_frames_blob_name = "frames";
constexpr auto test_frames_blob_name  = "all_frames";

constexpr auto target_blob_name       = "target";
constexpr auto filter_blob_name       = "filter";
constexpr auto cont_blob_name         = "cont";
constexpr auto q_values_blob_name     = "q_values";

constexpr auto ip1Size = 512;
constexpr auto lstmSize = 512;

using LastAction       = std::vector<Action>;
using LastActionBatch  = std::vector<LastAction>;

using FrameData        = std::array<uint8_t, kCroppedFrameDataSize>;
using FrameDataSp      = std::shared_ptr<FrameData>;
using InputFrames      = std::vector<FrameDataSp>;
using InputFramesBatch = std::vector<InputFrames>;
using Transition       = std::tuple<Action, FrameDataSp, Action, float,
      boost::optional<FrameDataSp> >;
using Episode          = std::vector<Transition>;
using ReplayMemory     = std::deque<Episode>;
using MemoryLayer      = caffe::MemoryDataLayer<float>;
using FrameVec         = std::vector<FrameDataSp>;

using ActionValue = std::pair<Action, float>;
using SolverSp = std::shared_ptr<caffe::Solver<float>>;
using NetSp = boost::shared_ptr<caffe::Net<float>>;

/**
 * Deep Q-Network
 */
class DQN
{
public:
    DQN(const ActionVect& legal_actions,
        const int replay_memory_capacity,
        const double gamma,
        const int clone_frequency,
        const int unroll,
        const int minibatch_size,
        const int frames_per_timestep);

    // Initialize DQN. Must be called before calling any other method.
    void Initialize(caffe::SolverParameter& solver_param);

    // Create the caffe net .prototxt
    caffe::NetParameter CreateNet(bool unroll1_is_lstm);

    // Load a trained model from a file.
    void LoadTrainedModel(const std::string& model_file);

    // Restore solving from a solver file.
    void RestoreSolver(const std::string& solver_file);

    // Snapshot the model/solver/replay memory. Produces files:
    // snapshot_prefix_iter_N.[caffemodel|solverstate|replaymem]. Optionally
    // removes snapshots that share the same prefix but have a lower
    // iteration number.
    void Snapshot(const std::string& snapshot_prefix, bool remove_old=false,
                  bool snapshot_memory=true);

    // A specialized method for producing a high-score
    // snapshot. Optionally remove older HiScore snapshots
    void SnapshotHiScore(const std::string& snapshot_prefix,
                         double avg_score, double std_dev,
                         bool remove_old=true);

    // Select an action by epsilon-greedy. If cont is false, LSTM state
    // will be reset. cont should be true only at start of new episodes.
    Action SelectAction(const InputFrames& frames, const LastAction& last_action, double epsilon, bool cont);

    // Select a batch of actions by epsilon-greedy.
    ActionVect SelectActions(const InputFramesBatch& frames_batch,
							 const LastActionBatch& last_action_batch,
                             double epsilon, bool cont);

    // Add an episode to the replay memory
    void RememberEpisode(const Episode& episode);

    // Update DQN. Returns the number of solver steps executed.
///    int UpdateSequential();
    // Updates from a random minibatch of experiences
    int UpdateRandom();

    // Clear the replay memory
    void ClearReplayMemory();

    // Save the replay memory to a gzipped compressed file
    void SnapshotReplayMemory(const std::string& filename);

    // Load the replay memory from a gzipped compressed file
    void LoadReplayMemory(const std::string& filename);

    // Get the number of episodes stored in the replay memory
    int memory_episodes() const
    {
        return replay_memory_.size();
    }

    // Get the number of transitions store in the replay memory
    int memory_size() const
    {
        return replay_memory_size_;
    }

    // Return the current iteration of the solver
    int current_iteration() const
    {
        return solver_->iter();
    }

    void CloneTestNet()
    {
        CloneNet(*test_net_);
    }

    // Benchmark the speed of the learning by doing some number of
    // iterations of updates and selects. random_updates toggles
    // random/sequential updating.
    //void Benchmark(int iterations, bool random_updates);

    // Obscures the screen by zeroing everything with a given probability.
    void ObscureScreen(FrameDataSp& screen, double obscure_prob);
    // Re-Display the last seen screen with probability prob
    void RedisplayScreen(FrameDataSp& screen, double prob);

    // Returns the number of transitions in the last episode added to
    // the memory or 0 if the memory is empty.
    int GetLastEpisodeSize();

protected:
    // Clone the given net and store the result in clone_net_
    void CloneNet(caffe::Net<float>& net);

    // Given a set of input frames and a network, select an
    // action. Returns the action and the estimated Q-Value.
	ActionValue SelectActionGreedily(caffe::Net<float>& net,
                                     const InputFrames& last_frames,
                                     const Action& last_action,
                                     bool cont);

    // Given a vector of frames, return a batch of selected actions + values.
	std::vector<ActionValue> SelectActionGreedily(caffe::Net<float>& net,
												   const InputFramesBatch& frames_batch,
                                                   const LastActionBatch& last_action_batch,
                                                   bool cont);

    // Input data into the Frames/Target/Filter layers of the given
    // net. This must be done before forward is called.
    void InputDataIntoLayers(caffe::Net<float>& net,
                             float* frames_input,
							 float* last_action_input,
                             float* cont_input,
                             float* target_input,
                             float* filter_input);

protected:
    int unroll_; // Number of steps to unroll recurrent layers
    int minibatch_size_; // Size of each minibatch
    int frames_per_timestep_; // History of frames given at each timestep
    int frames_per_forward_; // Number of frames needed by each forward

    // Size of the input blobs to the memory layers
    int frame_input_size_TRAIN_, target_input_size_TRAIN_,
        filter_input_size_TRAIN_, cont_input_size_TRAIN_,
		last_action_input_size_TRAIN_;
    int frame_input_size_TEST_, cont_input_size_TEST_,
		last_action_input_size_TEST_;

    const ActionVect legal_actions_;
    const int replay_memory_capacity_;
    const double gamma_;
    const int clone_frequency_; // How often (steps) the clone_net is updated
    int replay_memory_size_; // Number of transitions in replay memory
    ReplayMemory replay_memory_;
    SolverSp solver_;
    NetSp net_; // The primary network used for action selection.
    NetSp test_net_; // Net used for testing
    NetSp clone_net_; // Clone used to generate targets.
    int last_clone_iter_; // Iteration in which the net was last cloned
    std::mt19937 random_engine;
    float smoothed_loss_;
    std::vector<uint8_t> last_displayed_screen_; // Used in RedisplayScreen
};

/**
 * Returns a vector of filenames matching a given regular expression.
 */
std::vector<std::string> FilesMatchingRegexp(const std::string& regexp);

/**
 * Removes snapshots starting with snapshot_prefix that have an
 * iteration less than min_iter. Does not remove high-score snapshots.
 */
void RemoveSnapshots(const std::string& snapshot_prefix, int min_iter);

/**
 * Look for the latest snapshot to resume from. Returns a string
 * containing the path to the .solverstate. Returns empty string if
 * none is found. Will only return if the snapshot contains all of:
 * .solverstate,.caffemodel,.replaymemory
 */
std::string FindLatestSnapshot(const std::string& snapshot_prefix);
/**
 * Returns a list of high score snapshots
 */
std::vector<std::string> GetHiScoreSnapshots(const std::string& snapshot_prefix);

/**
 * Look for the best HiScore matching the given snapshot prefix
 */
float FindHiScore(const std::string& snapshot_prefix);

/**
 * Remove all high-score snapshots matching the given snapshot prefix
 */
void RemoveHiScoreSnapshots(const std::string& snapshot_prefix);

/**
 * Preprocess an ALE screen (downsampling & grayscaling). Optionally
 * obscure the screen to make ALE into a POMDP.
 */
FrameDataSp PreprocessScreen(const ALEScreen& raw_screen);

}

#endif /* DQN_HPP_ */


// Boost.Geometry

// Copyright (c) 2021, Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

// Licensed under the Boost Software License version 1.0.
// http://www.boost.org/users/license.html

#ifndef BOOST_GEOMETRY_STRATEGIES_AZIMUTH_GEOGRAPHIC_HPP
#define BOOST_GEOMETRY_STRATEGIES_AZIMUTH_GEOGRAPHIC_HPP

// TODO: move this file to boost/geometry/strategy
#include <boost/geometry/strategies/geographic/azimuth.hpp>

#include <boost/geometry/strategies/azimuth/services.hpp>
#include <boost/geometry/strategies/detail.hpp>

namespace boost { namespace geometry
{

namespace strategies { namespace azimuth
{

template
<
    typename FormulaPolicy = strategy::andoyer,
    typename Spheroid = srs::spheroid<double>,
    typename CalculationType = void
>
class geographic : strategies::detail::geographic_base<Spheroid>
{
    using base_t = strategies::detail::geographic_base<Spheroid>;

public:
    geographic()
        : base_t()
    {}

    explicit geographic(Spheroid const& spheroid)
        : base_t(spheroid)
    {}

    auto azimuth() const
    {
        return strategy::azimuth::geographic
            <
                FormulaPolicy, Spheroid, CalculationType
            >(base_t::m_spheroid);
    }
};

namespace services
{

template <typename Point1, typename Point2>
struct default_strategy<Point1, Point2, geographic_tag, geographic_tag>
{
    using type = strategies::azimuth::geographic<>;
};

template <typename FP, typename S, typename CT>
struct strategy_converter<strategy::azimuth::geographic<FP, S, CT> >
{
    static auto get(strategy::azimuth::geographic<FP, S, CT> const& strategy)
    {
        return strategies::azimuth::geographic<FP, S, CT>(strategy.model());
    }
};

} // namespace services

}} // namespace strategies::azimuth

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_STRATEGIES_AZIMUTH_GEOGRAPHIC_HPP


/*

Copyright (c) 2005-2017, University of Oxford.
All rights reserved.

University of Oxford means the Chancellor, Masters and Scholars of the
University of Oxford, having an administrative office at Wellington
Square, Oxford OX1 2JD, UK.

This file is part of Chaste.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.
 * Neither the name of the University of Oxford nor the names of its
   contributors may be used to endorse or promote products derived from this
   software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

*/
#ifndef DYNAMICALLY_LOADABLE_LR91_HPP_
#define DYNAMICALLY_LOADABLE_LR91_HPP_

#include "ChasteSerialization.hpp"
#include <boost/serialization/base_object.hpp>

#include "AbstractCardiacCell.hpp"
#include "AbstractStimulusFunction.hpp"
#include "AbstractDynamicallyLoadableEntity.hpp"
#include <vector>

/**
 * This class represents the Luo-Rudy 1991 system of equations,
 * with support for being compiled into a .so and loaded at run-time.
 */
class DynamicallyLoadableLr91 : public AbstractCardiacCell, public AbstractDynamicallyLoadableEntity
{
private:
    /** Needed for serialization. */
    friend class boost::serialization::access;
    /**
     * Archive the member variables.
     *
     * @param archive
     * @param version
     */
    template<class Archive>
    void serialize(Archive & archive, const unsigned int version)
    {
        archive & boost::serialization::base_object<AbstractCardiacCell>(*this);
        archive & boost::serialization::base_object<AbstractDynamicallyLoadableEntity>(*this);
    }

    /* Constants for the model */

    /** membrane capcaitance*/
    static const double membrane_C;
    /** Faraday constant*/
    static const double membrane_F;
    /** Universal gas constant*/
    static const double membrane_R;
    /** Temeperature*/
    static const double membrane_T;
    /** Reversal potentila for background current*/
    static const double background_current_E_b;
    /** Maximal conductance for background current*/
    static const double background_current_g_b;
    /** Maximal conductance for sodium current*/
    static const double fast_sodium_current_g_Na;
    /** Intracellular potassium concentration*/
    static const double ionic_concentrations_Ki;
    /** Extracellular potassium concentration*/
    static const double ionic_concentrations_Ko;
    /** Intracellular sodium concentration*/
    static const double ionic_concentrations_Nai;
    /** Extracellular sodium concentration*/
    static const double ionic_concentrations_Nao;
    /** Maximal conductance for plateau current*/
    static const double plateau_potassium_current_g_Kp;
    /** Permeability ratio Na/K for potassium currents*/
    static const double time_dependent_potassium_current_PR_NaK;

    /** Another parameter, which is a function of the above */
    double fast_sodium_current_E_Na;

    /**
     *  Range-checking on the current values of the state variables. Make sure
     *  all gating variables have are within zero and one, and all concentrations
     *  are positive
     */
    void VerifyStateVariables();

public:
    /**
     * Constructor
     *
     * @param pSolver is a pointer to the ODE solver
     * @param pIntracellularStimulus is a pointer to the intracellular stimulus
     */
    DynamicallyLoadableLr91(boost::shared_ptr<AbstractIvpOdeSolver> pSolver,
                            boost::shared_ptr<AbstractStimulusFunction> pIntracellularStimulus);

    /**
     * Destructor
     */
    ~DynamicallyLoadableLr91();

    /**
     * Fill in a vector representing the RHS of the Luo-Rudy 1991 system
     * of Odes at each time step, y' = [y1' ... yn'].
     * Some ODE solver will call this function repeatedly to solve for y = [y1 ... yn].
     *
     * @param time  the current time, in milliseconds
     * @param rY  current values of the state variables
     * @param rDY  to be filled in with derivatives
     */
    void EvaluateYDerivatives(double time, const std::vector<double> &rY, std::vector<double> &rDY);

    /**
     * Returns the ionic current
     *
     * @param pStateVariables  optional state at which to evaluate the current
     * @return the total ionic current
     */
    double GetIIonic(const std::vector<double>* pStateVariables=NULL);

    /**
     * Get the intracellular calcium concentration
     *
     * @return the intracellular calcium concentration
     */
    double GetIntracellularCalciumConcentration();
};

#include "SerializationExportWrapper.hpp"
CHASTE_CLASS_EXPORT(DynamicallyLoadableLr91)

namespace boost
{
namespace serialization
{
/**
 * Allow us to not need a default constructor, by specifying how Boost should
 * instantiate a DynamicallyLoadableLr91 instance.
 */
template<class Archive>
inline void save_construct_data(
    Archive & ar, const DynamicallyLoadableLr91 * t, const unsigned int file_version)
{
    const boost::shared_ptr<AbstractIvpOdeSolver> p_solver = t->GetSolver();
    const boost::shared_ptr<AbstractStimulusFunction> p_stimulus = t->GetStimulusFunction();
    ar << p_solver;
    ar << p_stimulus;
}

/**
 * Allow us to not need a default constructor, by specifying how Boost should
 * instantiate a DynamicallyLoadableLr91 instance (using existing constructor).
 *
 * NB this constructor allocates memory for the other member variables too.
 */
template<class Archive>
inline void load_construct_data(
    Archive & ar, DynamicallyLoadableLr91 * t, const unsigned int file_version)
{

    boost::shared_ptr<AbstractIvpOdeSolver> p_solver;
    boost::shared_ptr<AbstractStimulusFunction> p_stimulus;
    ar >> p_solver;
    ar >> p_stimulus;
    ::new(t)DynamicallyLoadableLr91(p_solver, p_stimulus);
}
}
} // namespace ...

#endif // DYNAMICALLY_LOADABLE_LR91_HPP_