#include "optimize.hpp"
#include <boost/program_options.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include "boost/filesystem.hpp"
#include <time.h>
#include <set>

Timer timer;

po::variables_map check_options(int argc, char** argv) {
    uint32_t num_cores = tbb::task_scheduler_init::default_num_threads();
    std::string num_threads_message = "Number of threads to use when possible [DEFAULT uses all available cores, " + std::to_string(num_cores) + " detected on this machine]";

    po::options_description desc("optimize options");
    desc.add_options()
        ("input-mat,i", po::value<std::string>()->required(),
         "Input mutation-annotated tree file to optimize [REQUIRED].")
        ("output-mat,o", po::value<std::string>()->required(),
         "Output optimized mutation-annotated tree file [REQUIRED].")
        ("sample-names,f", po::value<std::string>()->default_value(""),
         "File containing sample names whose ancestral nodes would be considered for optimization.")
        ("input-vcf,v", po::value<std::string>()->default_value(""),
         "Input vcf filename to reassign resolved bases using Fitch-Sankoff initially.")
        ("radius,r", po::value<uint32_t>()->default_value(10), \
         "Radius in which to restrict the SPR moves.")
        ("optimization-seconds,s", po::value<uint32_t>()->default_value(3600), \
         "Approximate number of seconds to run the tree optimization stage. The stage terminates as soon as the elapsed time exceeds this value.")
        ("save-every-seconds,S", po::value<uint32_t>()->default_value(300), \
         "Periodically save the optimized tree after every specified number of seconds have elapsed since the last save.")
        ("threads,T", po::value<uint32_t>()->default_value(num_cores), num_threads_message.c_str())
        ("help,h", "Print help messages");

    po::options_description all_options;
    all_options.add(desc);
    po::positional_options_description p;
    po::variables_map vm;
    try{
        po::store(po::command_line_parser(argc, argv)
                .options(all_options)
                .positional(p)
                .run(), vm);
        po::notify(vm);
    }
    catch(std::exception &e){
        std::cerr << desc << std::endl;
        // Return with error code 1 unless
        // the user specifies help
            if (vm.count("help"))
                exit(0);
            else
                exit(1);
    }
    return vm;
}

struct Pruned_Sample {
    std::string sample_name;
    std::vector<MAT::Mutation> sample_mutations;
    std::unordered_set<uint32_t> positions;

    // Assumes mutations are added in reverse chrono order
    void add_mutation (MAT::Mutation mut) {
        // If not reversal to reference allele
        if ((mut.ref_nuc != mut.mut_nuc) && (positions.find(mut.position) == positions.end())) {
            auto iter = std::lower_bound(sample_mutations.begin(), sample_mutations.end(), mut);
            mut.par_nuc = mut.ref_nuc;
            sample_mutations.insert(iter, mut);
        }
        positions.insert(mut.position);
    }

    Pruned_Sample (std::string name) {
        sample_name = name;
        sample_mutations.clear();
        positions.clear();
    }
};

inline std::string get_reversal_mutation(const std::string& mut_string) {
    std::vector<std::string> words;
    MAT::string_split(mut_string, words);
    assert(words.size() == 4);
    return words[0] + "\t" + words[3] + "\t" + words[2] + "\t" + words[1];
}

size_t get_node_distance (const MAT::Tree& T, MAT::Node* source, MAT::Node* dest) {
    auto s_anc = T.rsearch(source->identifier, true);
    auto d_anc = T.rsearch(dest->identifier, true);

    size_t s_off=0, d_off=0, num_iter=0;
    if (s_anc.size() <= d_anc.size()) {
        d_off = d_anc.size() - s_anc.size();
        num_iter = s_anc.size();
    }
    else {
        s_off = s_anc.size() - d_anc.size();
        num_iter = d_anc.size();
    }

    auto distance = s_off+d_off;
    for (size_t i=0; i<num_iter; i++) {
        if (s_anc[s_off+i] == d_anc[d_off+i]) {
            break;
        }
        else {
            distance+=2;
        }
    }

    return distance;
}

/*
void optimize_main_old(po::parsed_options parsed) {
    po::variables_map vm = parse_optimize_command(parsed);
    std::string input_mat_filename = vm["input-mat"].as<std::string>();
    std::string output_mat_filename = vm["output-mat"].as<std::string>();
    uint32_t max_seconds = vm["optimization-seconds"].as<uint32_t>();
    uint32_t num_threads = vm["threads"].as<uint32_t>();
    tbb::task_scheduler_init init(num_threads);
    srand (time(NULL));
    timer.Start();
    fprintf(stderr, "Loading input MAT file %s.\n", input_mat_filename.c_str());
    // Load input MAT and uncondense tree
    MAT::Tree T = MAT::load_mutation_annotated_tree(input_mat_filename);
    T.uncondense_leaves();
    fprintf(stderr, "The parsimony score for this MAT is %zu\n", T.get_parsimony_score());
    fprintf(stderr, "Completed in %ld msec \n\n", timer.Stop());
    timer.Start();
    fprintf(stderr, "Starting tree optimization.\n\n");

    fprintf(stderr, "Finding the nodes with recurrent or reversal mutations.\n");
    auto dfs = T.depth_first_expansion();

    std::map<std::string, int> mutation_counts;
    tbb::mutex tbb_lock;
    static tbb::affinity_partitioner ap;
    tbb::parallel_for(tbb::blocked_range<size_t>(0, dfs.size()),
            [&](const tbb::blocked_range<size_t> r) {
            for (size_t i=r.begin(); i<r.end(); ++i){
            auto n = dfs[i];
            n->clear_annotations();

            for (auto m: n->mutations) {
               std::string mut_string = m.chrom + "\t" + std::to_string(m.par_nuc) + "\t" +
                   std::to_string(m.position) + "\t" + std::to_string(m.mut_nuc);
               tbb_lock.lock();
               if (mutation_counts.find(mut_string) == mutation_counts.end()) {
                   mutation_counts[mut_string] = 1;
               }
               else {
                   mutation_counts[mut_string] += 1;
               }
               tbb_lock.unlock();
           }
        }
    }, ap);
    std::vector<std::string> nodes_to_prune;
    for (auto n: dfs) {
        if ((n == T.root) || (T.get_num_leaves(n) > 1000)) {
            continue;
        }
        for (auto m: n->mutations) {
            std::string mut_string = m.chrom + "\t" + std::to_string(m.par_nuc) + "\t" +
                std::to_string(m.position) + "\t" + std::to_string(m.mut_nuc);
            // Contains recurrent mutation
            if (mutation_counts[mut_string] > 1) {
                nodes_to_prune.emplace_back(n->identifier);
                break;
            }
            auto reversal_string = get_reversal_mutation(mut_string);
            // reversal mutation found
            if (mutation_counts.find(reversal_string) != mutation_counts.end()) {
                nodes_to_prune.emplace_back(n->identifier);
                break;
            }
        }
    }
    // Shuffle the nodes to be pruned
    std::random_shuffle(nodes_to_prune.begin(), nodes_to_prune.end());
    fprintf(stderr, "%zu nodes found with recurrent or reversal mutations.\n\n", nodes_to_prune.size());
    auto best_parsimony_score = T.get_parsimony_score();
    for (auto node_to_prune: nodes_to_prune) {
        if (T.get_node(node_to_prune) == NULL) {
            continue;
        }
        auto copy = MAT::get_tree_copy(T);
        fprintf(stderr, "Pruning subtree at node %s.\n", node_to_prune.c_str());
        auto leaves_to_prune = T.get_leaves(node_to_prune);
        std::vector<Pruned_Sample> pruned_samples;
        for (auto l: leaves_to_prune) {
            Pruned_Sample to_prune(l->identifier);
            auto root_to_node = T.rsearch(l->identifier, true);
            std::reverse(root_to_node.begin(), root_to_node.end());
            for (auto curr: root_to_node) {
                for (auto m: curr->mutations) {
                    to_prune.add_mutation(m);
                }
            }
            // move_level set to false as moving it would be an unnecessary overhead
            T.remove_node(l->identifier, false);
            pruned_samples.emplace_back(to_prune);
        }
        fprintf(stderr, "Pruned %zu samples at node %s.\n", pruned_samples.size(), node_to_prune.c_str());
        fprintf(stderr, "Placing %zu samples at node %s in random order.\n", pruned_samples.size(), node_to_prune.c_str());
        std::random_shuffle(pruned_samples.begin(), pruned_samples.end());
        for (auto s: pruned_samples) {
            dfs = T.depth_first_expansion();
            size_t total_nodes = dfs.size();
            // Stores the excess mutations to place the sample at each
            // node of the tree in DFS order. When placement is as a
            // child, it only contains parsimony-increasing mutations in
            // the sample. When placement is as a sibling, it contains
            // parsimony-increasing mutations as well as the mutations
            // on the placed node in common with the new sample. Note
            // guaranteed to be corrrect only for optimal nodes since
            // the mapper can terminate the search early for non-optimal
            // nodes
            std::vector<std::vector<MAT::Mutation>> node_excess_mutations(total_nodes);
            // Stores the imputed mutations for ambiguous bases in the
            // sampled in order to place the sample at each node of the
            // tree in DFS order. Again, guaranteed to be corrrect only
            // for pasrimony-optimal nodes
            std::vector<std::vector<MAT::Mutation>> node_imputed_mutations(total_nodes);
            std::vector<int> node_set_difference;
            size_t best_node_num_leaves = 0;
            // The maximum number of mutations is bound by the number
            // of mutations in the missing sample (place at root)
            //int best_set_difference = 1e9;
            // TODO: currently number of root mutations is also added to
            // this value since it forces placement as child but this
            // could be changed later
            int best_set_difference = s.sample_mutations.size() + T.root->mutations.size() + 1;
            size_t best_j = 0;
            bool best_node_has_unique = false;
            std::vector<bool> node_has_unique(total_nodes, false);
            std::vector<size_t> best_j_vec;
            size_t num_best = 1;
            MAT::Node* best_node = T.root;
            best_j_vec.emplace_back(0);
            // Parallel for loop to search for most parsimonious
            // placements. Real action happens within mapper2_body
            static tbb::affinity_partitioner ap;
            tbb::parallel_for( tbb::blocked_range<size_t>(0, total_nodes),
                    [&](tbb::blocked_range<size_t> r) {
                    for (size_t k=r.begin(); k<r.end(); ++k){
                    mapper2_input inp;
                    inp.T = &T;
                    inp.node = dfs[k];
                    inp.missing_sample_mutations = &s.sample_mutations;
                    inp.excess_mutations = &node_excess_mutations[k];
                    inp.imputed_mutations = &node_imputed_mutations[k];
                    inp.best_node_num_leaves = &best_node_num_leaves;
                    inp.best_set_difference = &best_set_difference;
                    inp.best_node = &best_node;
                    inp.best_j =  &best_j;
                    inp.num_best = &num_best;
                    inp.j = k;
                    inp.has_unique = &best_node_has_unique;
                    inp.best_j_vec = &best_j_vec;
                    inp.node_has_unique = &(node_has_unique);
                    mapper2_body(inp, false);
                    }
                    }, ap);
            // Ensure sample not already in the tree
            if (T.get_node(s.sample_name) == NULL) {
                // Is placement as sibling
                if (best_node->is_leaf() || best_node_has_unique) {
                    std::string nid = std::to_string(++T.curr_internal_node);
                    T.create_node(nid, best_node->parent->identifier);
                    T.create_node(s.sample_name, nid);
                    T.move_node(best_node->identifier, nid);
                    // common_mut stores mutations common to the
                    // best node branch and the sample, l1_mut
                    // stores mutations unique to best node branch
                    // and l2_mut stores mutations unique to the
                    // sample not in best node branch
                    std::vector<MAT::Mutation> common_mut, l1_mut, l2_mut;
                    std::vector<MAT::Mutation> curr_l1_mut;
                    // Compute current best node branch mutations
                    for (auto m1: best_node->mutations) {
                        MAT::Mutation m = m1.copy();
                        curr_l1_mut.emplace_back(m);
                    }
                    // Clear mutations on the best node branch which
                    // will be later replaced by l1_mut
                    best_node->clear_mutations();
                    // Compute l1_mut
                    for (auto m1: curr_l1_mut) {
                        bool found = false;
                        for (auto m2: node_excess_mutations[best_j]) {
                            if (m1.is_masked()) {
                                break;
                            }
                            if (m1.position == m2.position) {
                                if (m1.mut_nuc == m2.mut_nuc) {
                                    found = true;
                                    break;
                                }
                            }
                        }
                        if (!found) {
                            MAT::Mutation m = m1.copy();
                            l1_mut.emplace_back(m);
                        }
                    }
                    // Compute l2_mut
                    for (auto m1: node_excess_mutations[best_j]) {
                        bool found = false;
                        for (auto m2: curr_l1_mut) {
                            if (m1.is_masked()) {
                                break;
                            }
                            if (m1.position == m2.position) {
                                if (m1.mut_nuc == m2.mut_nuc) {
                                    found = true;
                                    MAT::Mutation m = m1.copy();
                                    common_mut.emplace_back(m);
                                    break;
                                }
                            }
                        }
                        if (!found) {
                            MAT::Mutation m = m1.copy();
                            l2_mut.emplace_back(m);
                        }
                    }
                    // Add mutations to new node using common_mut
                    for (auto m: common_mut) {
                        T.get_node(nid)->add_mutation(m);
                    }
                    // Add mutations to best node using l1_mut
                    for (auto m: l1_mut) {
                        T.get_node(best_node->identifier)->add_mutation(m);
                    }
                    // Add new sample mutations using l2_mut
                    for (auto m: l2_mut) {
                        T.get_node(s.sample_name)->add_mutation(m);
                    }
                }
                // Else placement as child
                else {
                    MAT::Node* node = T.create_node(s.sample_name, best_node->identifier);
                    std::vector<MAT::Mutation> node_mut;
                    std::vector<MAT::Mutation> curr_l1_mut;
                    for (auto m1: best_node->mutations) {
                        MAT::Mutation m = m1.copy();
                        curr_l1_mut.emplace_back(m);
                    }
                    for (auto m1: node_excess_mutations[best_j]) {
                        bool found = false;
                        for (auto m2: curr_l1_mut) {
                            if (m1.is_masked()) {
                                break;
                            }
                            if (m1.position == m2.position) {
                                if (m1.mut_nuc == m2.mut_nuc) {
                                    found = true;
                                    break;
                                }
                            }
                        }
                        if (!found) {
                            MAT::Mutation m = m1.copy();
                            node_mut.emplace_back(m);
                        }
                    }
                    for (auto m: node_mut) {
                        node->add_mutation(m);
                    }
                }
            }
        }
        fprintf(stderr, "Done pruning and placing samples at node %s\n", node_to_prune.c_str());
        auto new_parsimony_score = T.get_parsimony_score();
        if (new_parsimony_score >= best_parsimony_score) {
            fprintf(stderr, "Placement resulted in a parsimony score of %zu. Retaining old tree.\n", new_parsimony_score);
            MAT::clear_tree(T);
            T = copy;
        }
        else {
            MAT::clear_tree(copy);
            best_parsimony_score = new_parsimony_score;
            fprintf(stderr, "Placement lowered parsimony score to %zu!\n", best_parsimony_score);
        }
        auto elapsed_time = timer.Stop()/1000;
        fprintf(stderr, "Elapsed time (tree optimization stage): %ld seconds.\n\n", elapsed_time);
        if (elapsed_time >= max_seconds) {
            break;
        }
    }

    fprintf(stderr, "Optimization complete. Saving the final MAT with a parsimony score of %zu\n", T.get_parsimony_score());
    T.collapse_tree();
    T.condense_leaves();
    MAT::save_mutation_annotated_tree(T, output_mat_filename);

    fprintf(stderr, "Completed in %ld msec \n\n", timer.Stop());
}
*/

int main(int argc, char** argv) {
    po::variables_map vm = check_options(argc, argv);
    std::string input_mat_filename = vm["input-mat"].as<std::string>();
    std::string input_vcf_filename = vm["input-vcf"].as<std::string>();
    std::string output_mat_filename = vm["output-mat"].as<std::string>();
    std::string sample_name_filename = vm["sample-names"].as<std::string>();
    uint32_t max_seconds = vm["optimization-seconds"].as<uint32_t>();
    uint32_t save_every = vm["save-every-seconds"].as<uint32_t>();
    uint32_t radius = vm["radius"].as<uint32_t>();

    uint32_t num_threads = vm["threads"].as<uint32_t>();

    tbb::task_scheduler_init init(num_threads);
    srand (time(NULL));

    static tbb::affinity_partitioner ap;

    timer.Start();
    fprintf(stderr, "Loading input MAT file %s.\n", input_mat_filename.c_str());

    // Load input MAT and uncondense tree
    MAT::Tree T = MAT::load_mutation_annotated_tree(input_mat_filename);
    T.uncondense_leaves();

    // If VCF specified, re-compute the assignments using Fitch-Sankoff
    if (input_vcf_filename != "") {
        fprintf(stderr, "Computing parsimonious assignments for input variants.\n");
        timer.Start();

        // Variables below used to store the different fields of the input VCF file
        bool header_found = false;
        std::vector<std::string> variant_ids;
        std::vector<Missing_Sample> missing_samples;

        // Vector used to store all tree nodes in breadth-first search (BFS) order
        std::vector<MAT::Node*> bfs;
        // Map the node identifier string to index in the BFS traversal
        std::unordered_map<std::string, size_t> bfs_idx;

        // Breadth-first expansion to populate bfs and bfs_idx
        bfs = T.breadth_first_expansion();
        for (size_t idx = 0; idx < bfs.size(); idx++) {
            bfs_idx[bfs[idx]->identifier] = idx;
            // clear the node mutations
            bfs[idx]->clear_mutations();
        }

        // Boost library used to stream the contents of the input VCF file in
        // uncompressed or compressed .gz format
        std::ifstream infile(input_vcf_filename, std::ios_base::in | std::ios_base::binary);
        if (!infile) {
            fprintf(stderr, "ERROR: Could not open the VCF file: %s!\n", input_vcf_filename.c_str());
            exit(1);
        }
        boost::iostreams::filtering_istream instream;
        try {
            if (input_vcf_filename.find(".gz\0") != std::string::npos) {
                instream.push(boost::iostreams::gzip_decompressor());
            }
            instream.push(infile);
        }
        catch(const boost::iostreams::gzip_error& e) {
            std::cout << e.what() << '\n';
        }

        // A TBB flow graph containing a single source_node (reader) connected
        // to several mappers. The source_node sequentially reads in the different
        // lines of the input VCF file and constructs a mapper task for each
        // VCF line. Each mapper task takes a mapper_input as input, which stores
        // the alternate alleles, ambiguous bases and missing data (Ns) for
        // different tree samples at the corresponding VCF line/position. The
        // mappers use Fitch-Sankoff algorithm to assign mutations at different
        // branches of the tree and update the mutation-annotated tree (T)
        // accordingly.
        tbb::flow::graph mapper_graph;

        tbb::flow::function_node<mapper_input, int> mapper(mapper_graph, tbb::flow::unlimited, mapper_body());
        tbb::flow::source_node <mapper_input> reader (mapper_graph,
                [&] (mapper_input &inp) -> bool {

                //check if reached end-of-file
                int curr_char = instream.peek();
                if(curr_char == EOF)
                    return false;

                std::string s;
                std::getline(instream, s);
                std::vector<std::string> words;
                MAT::string_split(s, words);
                inp.variant_pos = -1;

                // Header is found when "POS" is the second word in the line
                if ((not header_found) && (words.size() > 1)) {
                  if (words[1] == "POS") {
                  // Sample names start from the 10th word in the header
                    for (size_t j=9; j < words.size(); j++) {
                      variant_ids.emplace_back(words[j]);
                      // If sample name not in tree, add it to missing_samples
                      if (bfs_idx.find(words[j]) == bfs_idx.end()) {
                        missing_samples.emplace_back(Missing_Sample(words[j]));
                      }
                    }
                    header_found = true;
                  }
                }
                else if (header_found) {
                    if (words.size() != 9+variant_ids.size()) {
                        fprintf(stderr, "ERROR! Incorrect VCF format.\n");
                        exit(1);
                    }
                    std::vector<std::string> alleles;
                    alleles.clear();
                    inp.variant_pos = std::stoi(words[1]);
                    MAT::string_split(words[4], ',', alleles);
                    // T will be modified by the mapper with mutation
                    // annotations
                    inp.T = &T;
                    inp.chrom = words[0];
                    inp.bfs = &bfs;
                    inp.bfs_idx = &bfs_idx;
                    inp.variant_ids = &variant_ids;
                    inp.missing_samples = &missing_samples;
                    // Ref nuc id uses one-hot encoding (A:0b1, C:0b10, G:0b100,
                    // T:0b1000)
                    inp.ref_nuc = MAT::get_nuc_id(words[3][0]);
                    assert((inp.ref_nuc & (inp.ref_nuc-1)) == 0); //check if it is power of 2
                    inp.variants.clear();
                    for (size_t j=9; j < words.size(); j++) {
                        if (isdigit(words[j][0])) {
                            int allele_id = std::stoi(words[j]);
                            if (allele_id > 0) {
                                std::string allele = alleles[allele_id-1];
                                inp.variants.emplace_back(std::make_tuple(j-9, MAT::get_nuc_id(allele[0])));
                            }
                        }
                        else {
                            inp.variants.emplace_back(std::make_tuple(j-9, MAT::get_nuc_id('N')));
                        }
                    }
                }
                return true;
                }, true );
        tbb::flow::make_edge(reader, mapper);
        mapper_graph.wait_for_all();

        fprintf(stderr, "Completed in %ld msec \n\n", timer.Stop());
    }

    // Collapse tree for optimal performance and results
    T.collapse_tree();

    fprintf(stderr, "The parsimony score for this MAT is %zu\n", T.get_parsimony_score());
    fprintf(stderr, "Completed in %ld msec \n\n", timer.Stop());

    float last_update = 0;

    timer.Start();
    fprintf(stderr, "Finding the nodes to prune.\n");
    auto bfs = T.breadth_first_expansion();

    std::set<std::string> nodes_to_prune;

    if (sample_name_filename != "") {
        std::ifstream infile(sample_name_filename);
        if (!infile) {
            fprintf(stderr, "ERROR: Could not open the sample name file: %s!\n", sample_name_filename.c_str());
            exit(1);
        }
        std::string line;

        fprintf(stderr, "Reading sample name file.\n");
        timer.Start();
        while (std::getline(infile, line)) {
            std::vector<std::string> words;
            MAT::string_split(line, words);
            if (words.size() != 1) {
                fprintf(stderr, "ERROR: Incorrect format for sample name file: %s!\n", sample_name_filename.c_str());
                exit(1);
            }
            auto n = T.get_node(words[0]);
            if (n == NULL) {
                fprintf(stderr, "WARNING: Node id %s not found in the tree!\n", words[0].c_str());
                exit(1);
            }
            else {
                for (auto anc: T.rsearch(words[0], true)) {
                    if (anc->level >= 4) {
                        nodes_to_prune.insert(anc->identifier);
                    }
                }
            }
        }
        infile.close();
    }

    else {
        fprintf(stderr, "\n");
        std::unordered_map<int, std::vector<std::string>> pos_to_nid;
        for (auto n: bfs) {
            for (auto m: n->mutations) {
                if (pos_to_nid.find(m.position) != pos_to_nid.end()) {
                    pos_to_nid[m.position].emplace_back(n->identifier);
                }
                else {
                    pos_to_nid[m.position] = std::vector<std::string>();
                    pos_to_nid[m.position].emplace_back(n->identifier);
                }
            }
        }

        size_t curr=0;
        tbb::mutex tbb_lock;
        tbb::parallel_for(tbb::blocked_range<size_t>(0, pos_to_nid.size()),
                [&](tbb::blocked_range<size_t> r) {
                for (size_t it = r.begin(); it < r.end(); it++) {
                  auto at  = __sync_fetch_and_add(&curr, 1);
                  if (at % 10 == 0) {
                     tbb_lock.lock();
                     fprintf(stderr, "\rAt %zu of %zu", at, pos_to_nid.size());
                     fflush(stderr);
                     tbb_lock.unlock();
                  }
                  auto cn = pos_to_nid.begin();
                  std::advance(cn, it);
                  size_t num_elem = cn->second.size();
                  for (size_t i=0; i<num_elem; i++) {
                      tbb_lock.lock();
                      if (nodes_to_prune.find(cn->second[i]) != nodes_to_prune.end()) {
                          tbb_lock.unlock();
                          continue;
                      }
                      tbb_lock.unlock();
                      auto n = T.get_node(cn->second[i]);
                      for (size_t j=i+1; j<num_elem; j++) {
                          auto n2 = T.get_node(cn->second[j]);
                          if (get_node_distance(T, n, n2) <= radius) {
                              tbb_lock.lock();
                              nodes_to_prune.insert(n->identifier);
                              nodes_to_prune.insert(n2->identifier);
                              tbb_lock.unlock();
                              break;
                          }
                      }
                  }
                }
            }, ap);

        fprintf(stderr, "\n");
    }

    fprintf(stderr, "%zu nodes found to prune.\n", nodes_to_prune.size());
    fprintf(stderr, "Completed in %ld msec \n\n", timer.Stop());

    timer.Start();
    fprintf(stderr, "Starting tree optimization.\n\n");

    auto best_parsimony_score = T.get_parsimony_score();
    size_t curr=0;
    for (auto nid_to_prune: nodes_to_prune) {
        fprintf(stderr, "At %zu of %zu\n", ++curr, nodes_to_prune.size());

        auto node_to_prune = T.get_node(nid_to_prune);
        if (node_to_prune == NULL) {
            continue;
        }

        fprintf(stderr, "Pruning subtree at node %s.\n", nid_to_prune.c_str());

        // Find mutations on the node to prune
        Pruned_Sample pruned_sample(nid_to_prune);

        auto node_to_root = T.rsearch(nid_to_prune, true);
        //std::reverse(root_to_node.begin(), root_to_node.end());
        for (auto curr: node_to_root) {
            for (auto m: curr->mutations) {
                pruned_sample.add_mutation(m);
            }
        }

        // Now prune the node_to_prune from the tree
        auto curr_parent = node_to_prune->parent;
        auto iter = std::find(curr_parent->children.begin(), curr_parent->children.end(), node_to_prune);
        assert (iter != curr_parent->children.end());
        curr_parent->children.erase(iter);

        // Set source to current parent
        auto source = curr_parent;

        // Remove curr_parent if it has no children
        if (curr_parent->children.size() == 0) {
            auto ancestors = T.rsearch(curr_parent->identifier, true);
            T.remove_node(curr_parent->identifier, true);

            // Since remove_node can remove multiple levels of ancestors,
            // reassign source to nearest ancestor currently in the tree
            for (auto anc: ancestors) {
                if (T.get_node(anc->identifier) != NULL) {
                    source = anc;
                    break;
                }
            }
        }

        // Move the remaining child one level up if it is the only child of its parent
        if (curr_parent->children.size() == 1) {
            auto child = curr_parent->children[0];
            if (curr_parent->parent != NULL) {
                child->parent = curr_parent->parent;
                child->level = curr_parent->parent->level + 1;

                std::vector<MAT::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);

                curr_parent->children.clear();
                T.remove_node(curr_parent->identifier, false);

                source = child;
            }
        }

        bfs = T.breadth_first_expansion();
        size_t total_nodes = bfs.size();

        std::vector<std::vector<MAT::Mutation>> node_excess_mutations(total_nodes);
        std::vector<std::vector<MAT::Mutation>> imputed_mutations(total_nodes);

        std::vector<int> node_set_difference;

        size_t best_node_num_leaves = 0;
        int best_set_difference = 1e9;

        std::vector<bool> node_has_unique(total_nodes);
        size_t best_j = 0;
        bool best_node_has_unique = false;

        size_t best_distance = 1e9;

        std::vector<size_t> best_j_vec;

        size_t num_best = 1;
        MAT::Node* best_node = T.root;
        best_j_vec.emplace_back(0);

        std::vector<size_t> node_distance(total_nodes);

        static tbb::affinity_partitioner ap;
        tbb::parallel_for( tbb::blocked_range<size_t>(0, total_nodes),
                [&](tbb::blocked_range<size_t> r) {
                for (size_t k=r.begin(); k<r.end(); ++k){
                node_distance[k] = get_node_distance(T, source, bfs[k]);
                if (node_distance[k] > radius) {
                   continue;
                }

                node_has_unique[k] = false;

                mapper2_input inp;
                inp.T = &T;
                inp.node = bfs[k];
                inp.missing_sample_mutations = &pruned_sample.sample_mutations;
                inp.excess_mutations = &node_excess_mutations[k];
                inp.imputed_mutations = &imputed_mutations[k];
                inp.best_node_num_leaves = &best_node_num_leaves;
                inp.best_set_difference = &best_set_difference;
                inp.best_node = &best_node;
                inp.best_j =  &best_j;
                inp.num_best = &num_best;
                inp.j = k;
                inp.has_unique = &best_node_has_unique;

                inp.distance = node_distance[k];
                inp.best_distance = &best_distance;

                inp.best_j_vec = &best_j_vec;
                inp.node_has_unique = &(node_has_unique);

                mapper2_body(inp, false);
                }
                }, ap);

        auto distance = node_distance[best_j];

        // Clear current mutations of the node to prune
        node_to_prune->clear_mutations();

        // Is placement as sibling
        if (best_node->is_leaf() || best_node_has_unique) {
            std::string nid = std::to_string(++T.curr_internal_node);
            auto new_node = T.create_node(nid, best_node->parent->identifier);
            node_to_prune->parent = new_node;
            new_node->children.emplace_back(node_to_prune);
            T.move_node(best_node->identifier, nid);

            std::vector<MAT::Mutation> common_mut, l1_mut, l2_mut;
            std::vector<MAT::Mutation> curr_l1_mut;

            // Compute current best node branch mutations
            for (auto m1: best_node->mutations) {
                MAT::Mutation m = m1.copy();
                curr_l1_mut.emplace_back(m);
            }
            // Clear mutations on the best node branch which
            // will be later replaced by l1_mut
            best_node->clear_mutations();

            // Compute l1_mut
            for (auto m1: curr_l1_mut) {
                bool found = false;
                for (auto m2: node_excess_mutations[best_j]) {
                    if (m1.is_masked()) {
                        break;
                    }
                    if (m1.position == m2.position) {
                        if (m1.mut_nuc == m2.mut_nuc) {
                            found = true;
                            break;
                        }
                    }
                }
                if (!found) {
                    MAT::Mutation m = m1.copy();
                    l1_mut.emplace_back(m);
                }
            }
            // Compute l2_mut
            for (auto m1: node_excess_mutations[best_j]) {
                bool found = false;
                for (auto m2: curr_l1_mut) {
                    if (m1.is_masked()) {
                        break;
                    }
                    if (m1.position == m2.position) {
                        if (m1.mut_nuc == m2.mut_nuc) {
                            found = true;
                            MAT::Mutation m = m1.copy();
                            common_mut.emplace_back(m);
                            break;
                        }
                    }
                }
                if (!found) {
                    MAT::Mutation m = m1.copy();
                    l2_mut.emplace_back(m);
                }
            }

            // Add mutations to new node using common_mut
            for (auto m: common_mut) {
                new_node->add_mutation(m);
            }
            // Add mutations to best node using l1_mut
            for (auto m: l1_mut) {
                best_node->add_mutation(m);
            }
            // Add new sample mutations using l2_mut
            for (auto m: l2_mut) {
                node_to_prune->add_mutation(m);
            }
        }
        // Else placement as child
        else {
            best_node->children.emplace_back(node_to_prune);
            node_to_prune->parent = best_node;

            std::vector<MAT::Mutation> node_mut;

            std::vector<MAT::Mutation> curr_l1_mut;

            for (auto m1: best_node->mutations) {
                MAT::Mutation m = m1.copy();
                curr_l1_mut.emplace_back(m);
            }

            for (auto m1: node_excess_mutations[best_j]) {
                bool found = false;
                for (auto m2: curr_l1_mut) {
                    if (m1.is_masked()) {
                        break;
                    }
                    if (m1.position == m2.position) {
                        if (m1.mut_nuc == m2.mut_nuc) {
                            found = true;
                            break;
                        }
                    }
                }
                if (!found) {
                    MAT::Mutation m = m1.copy();
                    node_mut.emplace_back(m);
                }
            }
            for (auto m: node_mut) {
                node_to_prune->add_mutation(m);
            }
        }

        fprintf(stderr, "Done pruning and placing samples at node %s\n", best_node->identifier.c_str());
        fprintf(stderr, "Distance from original placement: %zu\n", distance);

        auto new_parsimony_score = T.get_parsimony_score();
        assert(new_parsimony_score <= best_parsimony_score);
        if (new_parsimony_score == best_parsimony_score) {
            fprintf(stderr, "Parsimony score of %zu unchanged after placement.\n", new_parsimony_score);
        }
        else {
            best_parsimony_score = new_parsimony_score;
            fprintf(stderr, "Placement lowered parsimony score to %zu!\n", best_parsimony_score);
        }

        float elapsed_sec = static_cast<float>(timer.Stop())/1000;
        fprintf(stderr, "Elapsed time (tree optimization stage): %.3f sec\n\n", elapsed_sec);

        if (elapsed_sec >= max_seconds) {
            break;
        }

        if (elapsed_sec-last_update >= save_every) {
            // Levels need to be adjusted after placement
            bfs = T.depth_first_expansion();
            for (auto n: bfs) {
                if (n == T.root) {
                    n->level = 1;
                }
                else {
                    n->level = n->parent->level+1;
                }
            }

            fprintf(stderr, "Saving the current MAT with a parsimony score of %zu\n\n", T.get_parsimony_score());
            T.condense_leaves();
            MAT::save_mutation_annotated_tree(T, output_mat_filename);
            last_update = elapsed_sec;

            // Uncondense for the next iteration
            T.uncondense_leaves();
        }
    }

    // Levels need to be adjusted after placement
    bfs = T.breadth_first_expansion();
    for (auto n: bfs) {
        if (n == T.root) {
            n->level = 1;
        }
        else {
            n->level = n->parent->level+1;
        }
    }

    fprintf(stderr, "Optimization complete. Saving the final MAT with a parsimony score of %zu\n", T.get_parsimony_score());
    T.condense_leaves();
    MAT::save_mutation_annotated_tree(T, output_mat_filename);

    fprintf(stderr, "Completed in %ld msec \n\n", timer.Stop());
}


#include <ScrimpTrivialParVec.hpp>
#include <logging.hpp>
#include <partitioning_1d.h>
#include <binproffile.h>
#include <settings.h>

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <limits>
#include <algorithm>
#include <string>
#include <sstream>
#include <chrono>
#include <cassert>
#include <memory>
#include <cstddef>

#include <boost/assert.hpp>
#include <boost/align.hpp>
#include <boost/filesystem.hpp>

#include <mpi.h>

#include <papiwrapper.hpp>

using DiagPartition = Partition1D;

// easy aligned allocation.
// code from https://www.boost.org/doc/libs/1_57_0/doc/html/align/examples.html#align.examples.aligned_ptr, licensed under the Boost Software License - Version 1.0
// see copy of the license terms provided under /licenses/boostSoftwareLicense.txt
template<class T>
using aligned_ptr = std::unique_ptr<T,
  boost::alignment::aligned_delete>;

template<class T, class... Args>
inline aligned_ptr<T> make_aligned(Args&&... args)
{
  auto p = boost::alignment::
    aligned_alloc(alignof(T), sizeof(T));
  if (!p) {
	throw std::bad_alloc();
  }
  try {
	auto q = ::new(p) T(std::forward<Args>(args)...);
	return aligned_ptr<T>(q);
  } catch (...) {
	boost::alignment::aligned_free(p);
	throw;
  }
}
// condinue own code

#define TIMEPOINT( varname ) const std::chrono::steady_clock::time_point varname = std::chrono::steady_clock::now();

using namespace matrix_profile;

static FactoryRegistration<ScrimpTrivialParVec> s_trivParReg("triv_par_vec");
static const int notification_interval_iter = 10000;

void commit_mpi_SOA_prof_type(MPI_Datatype* ptr_mpi_type, tsa_dtype* profile_buf, idx_dtype* idx_buf, int profile_length)
{
	const int blocklen[] = {profile_length,profile_length};
	MPI_Aint profBufAddr, idxBufAddr;
	MPI_Aint disps[2];
	//assertion for DOUBLE usage!!! TODO employ some preprocessor or template magic to automatically use a diffferent MPI Datatype...
	static_assert(sizeof(tsa_dtype) == sizeof(double), "NO FLOAT POSSIBLE without modifing MPI Datatype!" );
	static_assert(std::is_same<idx_dtype, long>::value, "Index datatype does not match the MPI Datatype!");
	const MPI_Datatype dtypes[] = {MPI_DOUBLE, MPI_LONG};
	// official way of handling addresses
	MPI_Get_address(profile_buf, &profBufAddr);
	MPI_Get_address(idx_buf, &idxBufAddr);
	disps[0] = 0;
	disps[1] = MPI_Aint_diff(idxBufAddr, profBufAddr);

	MPI_Type_create_struct(2, blocklen, disps, dtypes, ptr_mpi_type);
	MPI_Type_commit(ptr_mpi_type);
}

void decompose_datatype(
        const void * const invec,
        const void * const inoutvec,
        MPI_Datatype *datatype,
        tsa_dtype** inoutprof, idx_dtype** inout_idx,
        tsa_dtype** in2_prof, idx_dtype** in2_idx,
        idx_dtype& blocklen
        )
{
	EXEC_TRACE("reducing a MatProfSOA!");
	// retrieve the profile length: This is equal to the block length of the datastructure, not the length specified in the reduction!
	int num_int, num_add, num_dtype, combiner;
	MPI_Aint inout_addr, in2_addr;
	MPI_Get_address(invec, &in2_addr);
	MPI_Get_address(inoutvec, &inout_addr);
	MPI_Type_get_envelope(*datatype, &num_int, &num_add ,&num_dtype, &combiner);

	if (num_int>3 || num_add>2 || num_dtype >2){
		EXEC_ERROR("Invalid dataype sizes retrieved: num_i " << num_int
		           << " num_add " << num_add
		           << " num_dtype " << num_dtype
		           << "aborting as buffers are too small");
		MPI_Abort(MPI_COMM_WORLD, 1);
	}

	int blocklens[3];
	MPI_Aint addresses[2];
	MPI_Datatype dtypes[2];
	MPI_Type_get_contents(*datatype, num_int, num_add, num_dtype, blocklens, addresses, dtypes);

	assert(blocklens[1] == blocklens[2]);

	//"shortcuts" for coding conveniencebb
	*inoutprof = reinterpret_cast<tsa_dtype*>(MPI_Aint_add(inout_addr, addresses[0]) );;
	*inout_idx = reinterpret_cast<idx_dtype*>(MPI_Aint_add(inout_addr, addresses[1]) );
	*in2_prof = reinterpret_cast<tsa_dtype*>(MPI_Aint_add(in2_addr, addresses[0]) );
	*in2_idx = reinterpret_cast<idx_dtype*>(MPI_Aint_add(in2_addr, addresses[1]) );
	blocklen = blocklens[1];
}

void ScrimpTrivialParVec::MPI_MatProfSOA_reduction(void * invec, void *inoutvec, int *len, MPI_Datatype *datatype) {
	tsa_dtype *inout_prof, *in2_prof;
	idx_dtype *inout_idx, *in2_idx;
	idx_dtype blocklen;

	assert(*len == 1); // only reduction with size 1 is supported: Otherwise strides between SOAs needed to be considered....
	decompose_datatype(invec, inoutvec, datatype, &inout_prof, &inout_idx, &in2_prof, &in2_idx, blocklen);
	EXEC_TRACE("reduction with len " << *len << " and blocklen " << blocklen);

	//finally the matrix profile reduction
	const int limit = blocklen; // avoid unnecessary reads
	for (int i = 0; i < limit; ++i) {
		if (in2_prof[i] < inout_prof[i]) {
			inout_prof[i] = in2_prof[i];
			inout_idx[i] = in2_idx[i];
		}
	}
}

void ScrimpTrivialParVec::compute_matrix_profile(const Scrimppp_params& params) {
	TIMEPOINT( tstart );
	aligned_tsdtype_vec A = fetch_time_series<aligned_tsdtype_vec::allocator_type>(params); //load the time series data
	TIMEPOINT( tfileread );
	aligned_tsdtype_vec AMean(A.size());
	aligned_tsdtype_vec ASigma(A.size());
	aligned_tsdtype_vec dotproducts(A.size());
	int windowSize = params.query_window_len;
	int exclusionZone = windowSize / 4;
	int timeSeriesLength = A.size();
	int profile_length = timeSeriesLength - windowSize + 1;
	aligned_tsdtype_vec profile(profile_length, -1.0);//overallocation to enable usage of aligned avx-writes up to the end
	aligned_int_vec profileIndex(profile_length, 0);  //overallocation to enable usage of aligned avx-writes up to the end
	const int BLOCKING_SIZE = params.block_length;
	const bool distrib_io = (params.filetype == Scrimppp_params::BIN) && params.use_distributed_io;

	if (params.filetype != Scrimppp_params::BIN && params.use_distributed_io) {
		EXEC_ERROR("distributed I/O is NOT supported with ASCII files. Falling back to master I/O!");
	}

	// code instrumentaion (if enabled)
	PerfCounters ctrs_diaginit("diagonal initialization");
	PerfCounters ctrs_eval("block evaluations");
	PerfCounters ctrs_mpi_comm("MPI communication");

	// retrieve basic MPI information
	int world_size;
	int world_rank;
	MPI_Datatype mpi_result_type;
	MPI_Op reduction_op;
	MPI_Comm_size(MPI_COMM_WORLD, &world_size);
	MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);

	EXEC_INFO( "MPI INFO: comm size: " << world_size << " world rank: " << world_rank);
	EXEC_INFO( "Blocking size: " << BLOCKING_SIZE)

	if (profile_length > RESULT_ALLOC_LEN) { // check, whether the data fit into the preallocated result storage
		throw std::runtime_error("resulting profile length EXCEEDS MAXIMUM LENGTH. Recompile with different allocation size or choose a smaller problem");
	}
	else if(profile_length<=0) {
		throw std::runtime_error("parameters chosen badly: resulting matrix profile of length 0. done.");
	}

	//create the MPI dataype for the result
	commit_mpi_SOA_prof_type(&mpi_result_type, profile.data(), profileIndex.data(), profile_length );

	//create the reduction operation for the custom datatype
	MPI_Op_create(MPI_MatProfSOA_reduction, 1, &reduction_op);

	// values regarding work partitioning
	const int num_partitions = 2*world_size; //number of partitions. 2 partitions per process for the sake of balancing
	const int diags_to_process_world = profile_length-exclusionZone-1; //total number of diagonals in the "adjacency" matrix to process
	// likely the partitions can not be of exactly even size. We assume the first ones to be of partition_size_full_load and the remaining ones to be one element smaller
	const PartitioningInfo partinfo = {
	    ._num_partitions = 2*world_size,
	    ._num_partitions_full_load = diags_to_process_world - ((diags_to_process_world/num_partitions) * num_partitions),
	    ._size_full_load = (diags_to_process_world/num_partitions) + 1
	};

	/*const bool proc_has_reduced_load = world_rank<num_reduced_load;
	const int diags_to_process_proc = (proc_has_reduced_load?diags_to_process_reduced_load:diags_to_process_full_load);
	const int proc_offset = (proc_has_reduced_load? world_rank*diags_to_process_reduced_load
													: diags_to_process_reduced_load*num_reduced_load
														+diags_to_process_full_load*(world_rank-num_reduced_load));*/

	// retrieve the local partitions: each process gets 2. A "long" one in the upper half of the triangle and a "short" one in the lower half of the triangle
	std::vector<DiagPartition> local_partitions;
	    // the upper partition is the "process-rank"-th one
	DiagPartition part_offsets = get_partition(world_rank, partinfo);
	local_partitions.push_back( {exclusionZone+1+part_offsets._first_id, std::min(profile_length-1, exclusionZone+1+part_offsets._last_id)} );
	part_offsets = get_partition(num_partitions-world_rank-1, partinfo);
	local_partitions.push_back( {exclusionZone+1+part_offsets._first_id, std::min(profile_length-1, exclusionZone+1+part_offsets._last_id)} );

	//sum up, how many diags the local proc has to evaluate
	int diags_to_process_proc = 0;
	for (auto partit = local_partitions.begin(); partit < local_partitions.end(); ++partit) {
		diags_to_process_proc += partit->_last_id - partit->_first_id +1;
	}

#ifndef NDEBUG
	// sum up over all processes
	int num_processed_diags = 0;
	MPI_Reduce(&diags_to_process_proc, &num_processed_diags, 1, MPI_INTEGER, MPI_SUM, 0, MPI_COMM_WORLD);
	if (world_rank == 0) {
		BOOST_ASSERT_MSG(num_processed_diags==diags_to_process_world, "work partitioning of the diagonals is ivalid");
	}
	EXEC_INFO("partitioning is fine!");
#endif

//	EXEC_INFO("Processing all the matrix profile!");
//	local_partitions.clear();
//	local_partitions.push_back( {exclusionZone+1, ProfileLength-1} );
/*
	const int tmprank=0;
	EXEC_INFO("Processing the rank" << tmprank << " partition of the matrix profile!");
	local_partitions.clear();
	part_offsets = get_partition(tmprank, partinfo);
	local_partitions.push_back( {exclusionZone+1+part_offsets._first_id, std::min(ProfileLength-1, exclusionZone+1+part_offsets._last_id)} );
	part_offsets = get_partition(num_partitions-tmprank-1, partinfo);
	local_partitions.push_back( {exclusionZone+1+part_offsets._first_id, std::min(ProfileLength-1, exclusionZone+1+part_offsets._last_id)} );
*/

	//validation of parameters
	if (timeSeriesLength < windowSize) {
		throw std::invalid_argument("ERROR: Time series is shorter than the window length, can not proceed");
	}

	TIMEPOINT( tsetup );

	//precompute the mean and standard deviations of the sliding windows along the time series
	precompute_window_statistics(windowSize, A, profile_length, AMean, ASigma);

	//start time measurment
	TIMEPOINT( tprecomputations );

	// randomly shuffle the comutation order of diagonal blocks (according to the blocking size).
	// In order to do so split the local partition into blocks and shuffle them.
	std::vector<DiagPartition> eval_blocks;
	eval_blocks.reserve((diags_to_process_proc/BLOCKING_SIZE) +1);

	for (auto partit = local_partitions.begin(); partit < local_partitions.end(); ++partit) {
		    // split the partitions into block which preserve cache locality
		const int partlen = partit->_last_id - partit->_first_id;
		const int num_prolonged_blocks = std::max(partlen % BLOCKING_SIZE, partlen/BLOCKING_SIZE) ;
		int prev_block_end = partit->_first_id-1;

		EXEC_TRACE("splitting into blocks a partition with start: " << partit->_first_id << " and last: " << partit->_last_id);

		for ( int blocki = 0; prev_block_end < partit->_last_id; ++blocki) {
			DiagPartition part;
			part._first_id = prev_block_end+1;
			if (blocki < num_prolonged_blocks) {
				part._last_id = std::min(part._first_id + BLOCKING_SIZE-1, partit->_last_id);
			}
			else {
				part._last_id = std::min(part._first_id+BLOCKING_SIZE-1, partit->_last_id);
			}
			eval_blocks.push_back( part );
			prev_block_end = part._last_id;
		}
	}
	//std::random_shuffle(eval_blocks.begin(), eval_blocks.end());

	TIMEPOINT( tpartitioning )
	EXEC_DEBUG( "Proc " << world_rank << " shall process " << diags_to_process_proc << " diagonals");

	//init_all_diagonals(dotproducts, A, windowSize);

	//iteratively evaluate the diagonals of the distance matrix
	int progressctr=0;
	int next_notification_thresh = notification_interval_iter;
	for (auto blockiter = eval_blocks.begin(); blockiter < eval_blocks.end(); ++blockiter) {
		const int blocklen = blockiter->_last_id +1 - blockiter->_first_id;

		// compute distance of the first entry for each entry in the block

		    // might be more efficient, if performed with MASS for ech partition (when splitting them up into blocks)
		{
			const ScopedPerfAccumulator monitor(ctrs_diaginit);
			init_diagonals(blockiter->_first_id, blockiter->_last_id, dotproducts, A, windowSize);
		}

		EXEC_DEBUG("Rank " << world_rank << " evaluating diags " << blockiter->_first_id << " to " << blockiter->_last_id);
		{
			const ScopedPerfAccumulator monitor(ctrs_eval);
			eval_diagonal_block(blocklen, profileIndex, A, dotproducts, windowSize, ASigma, AMean, blockiter->_first_id, profile);

			progressctr+=blocklen;
		}
		//Show time per 10000 iterations
		if (progressctr > next_notification_thresh && false) //disable in or der for the logging not to disturb the timing
		{
			const auto tcur = std::chrono::steady_clock::now();
			std::chrono::duration<float> time_elapsed = tcur - tpartitioning;
			EXEC_INFO ("finished " << progressctr << " iterations in: " << std::setprecision(4) << time_elapsed.count() << " seconds.");
			next_notification_thresh += notification_interval_iter;
		}
	}

	TIMEPOINT( tevaluations)
	EXEC_DEBUG("apply transformation from score to distance values")
	// apply a correction of the distance values, as we dropped a factor of 2 to avoid unnecessary computations
	tsa_dtype twice_m = 2.0*static_cast<tsa_dtype>(windowSize);
	for (int i = 0; i<profile_length; ++i) {
		profile[i] = twice_m - 2.0 * profile[i];
	}
	TIMEPOINT( tpostprocessing)
	{
		const ScopedPerfAccumulator monitor(ctrs_mpi_comm);
		EXEC_INFO("reduce the processes individual results")
		//update matrix profile and matrix profile index if the current distance value is smaller
		if (distrib_io) { //distributed I/O => everyone needs the result => Allreduce
			// actually only a fraction is required: the partition which will be written. That would require several reductions (one for each partition), where the receiver is the process responsible for the partition...
			MPI_Allreduce(MPI_IN_PLACE, profile.data(), 1, mpi_result_type, reduction_op, MPI_COMM_WORLD); // in case of distributed binary I/O everyone needs the result
		}
		else { // reduction into master only, in case it is the only one writing...
			if (world_rank != 0) {
				MPI_Reduce(profile.data(), profile.data(), 1, mpi_result_type, reduction_op, 0, MPI_COMM_WORLD); // Reduction of length 1, as the datatype is defined as ProfileLength!
			}
			else {
				MPI_Reduce(MPI_IN_PLACE, profile.data(), 1, mpi_result_type, reduction_op, 0, MPI_COMM_WORLD); // Reduction of length 1, as the datatype is defined as ProfileLength!
			}
		}
		EXEC_INFO( "rank " << world_rank << " done with reduction");
	}

	TIMEPOINT( tcommunication)
	// MPI_Barrier(MPI_COMM_WORLD); //just for "debugging"

	//store the result
	if (world_rank == 0 || distrib_io ){ // storing in parallel if BIN type specified...
		EXEC_DEBUG("Rank " << world_rank << "storing the result");
		store_matrix_profile(profile.data(), profileIndex.data(), profile_length, params, distrib_io);
	}

	TIMEPOINT( tfilewrite )

	// MPI cleanup
	MPI_Op_free(&reduction_op);
	MPI_Type_free(&mpi_result_type);

	TIMEPOINT( tmpiclean)

	std::chrono::duration<double> setup_time, comp_time, comm_time, evaluation_time, io_time, work_time, precomp_time;
	setup_time = tsetup-tfileread + tmpiclean-tfilewrite + tpartitioning-tprecomputations;
	comm_time = tcommunication-tpostprocessing ;
	evaluation_time = tevaluations-tpartitioning;
	precomp_time = tprecomputations-tsetup;
	comp_time = evaluation_time + precomp_time + tpostprocessing-tevaluations;
	work_time = comp_time + comm_time + setup_time;
	io_time = tfileread - tstart + tfilewrite-tcommunication;
	// log computation performance
	const double triang_len = profile_length-exclusionZone;
	auto timerprecision = std::numeric_limits<tsa_dtype>::digits10 + 2;
	PERF_LOG ( "evaluation time: " << std::setprecision(timerprecision) << evaluation_time.count() << " seconds." );
	PERF_LOG ( "local comp time (eval+precomp): " << std::setprecision(timerprecision) << comp_time.count() << " seconds." );
	PERF_LOG ( "communication time: " << std::setprecision(timerprecision) << comm_time.count() << " seconds" );
	PERF_LOG ( "local working time: " << std::setprecision(timerprecision) << work_time.count() << " seconds" );
	PERF_LOG ( "I/O time: " << std::setprecision(timerprecision) << io_time.count() << " seconds" );

	PERF_LOG ( "local throughput evaluations: " << std::setprecision(6) << diags_to_process_proc * triang_len / evaluation_time.count() << " matrix entries/second" );
	PERF_LOG ( "throughput evaluations: " << std::setprecision(6) << triang_len * triang_len / evaluation_time.count() << " matrix entries/second (estimate based on local)" );
	PERF_LOG ( "local throughput computations: " << std::setprecision(6) << diags_to_process_proc * triang_len / comp_time.count() << " matrix entries/second" );
	PERF_LOG ( "throughput computations: " << std::setprecision(6) << triang_len * triang_len / comp_time.count() << " matrix entries/second (estimate based on local)" );

	using unit_seconds = std::chrono::duration<double>;
	// log exhaustive timing information
#define TIME_TRACE( timer, description ) PERF_TRACE ( "timepoint after " << description << ": " << \
	    std::setprecision(timerprecision) << std::chrono::duration_cast< unit_seconds >(timer.time_since_epoch()).count() << " seconds");
	TIME_TRACE( tstart, "start");
	TIME_TRACE( tfileread, "fileread");
	TIME_TRACE( tsetup, "setup");
	TIME_TRACE( tprecomputations, "precomputations");
	TIME_TRACE( tpartitioning, "partitioning");
	TIME_TRACE( tevaluations, "matrix_eval");
	TIME_TRACE( tpostprocessing, "postprocessing");
	TIME_TRACE( tcommunication, "communication");
	TIME_TRACE( tfilewrite, "postprocessing");
	TIME_TRACE( tmpiclean, "mpi_cleanup")

    #ifdef PROFILING
	    PERF_LOG("number of matrix profile updates: " << _profile_updates);
	    PERF_LOG("number of evaluations: " << _eval_ctr);
    #endif
	ctrs_diaginit.log_perf();
	ctrs_eval.log_perf();
	ctrs_mpi_comm.log_perf();
}


//=============================================================================================================
/**
 * @file     mne_cov_matrix.cpp
 * @author   Lorenz Esch <lesch@mgh.harvard.edu>;
 *           Matti Hamalainen <msh@nmr.mgh.harvard.edu>;
 *           Christoph Dinh <chdinh@nmr.mgh.harvard.edu>
 * @version  dev
 * @date     January, 2017
 *
 * @section  LICENSE
 *
 * Copyright (C) 2017, Lorenz Esch, Matti Hamalainen, Christoph Dinh. All rights reserved.
 *
 * 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 MNE-CPP authors 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 OWNER 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.
 *
 *
 * @brief    Definition of the MneCovMatrix Class.
 *
 */

//*************************************************************************************************************
//=============================================================================================================
// INCLUDES
//=============================================================================================================

#include "mne_cov_matrix.h"
#include "mne_sss_data.h"
#include "mne_proj_item.h"
#include "mne_proj_op.h"

#include <Eigen/Core>
#include <Eigen/Eigenvalues>

//*************************************************************************************************************
//=============================================================================================================
// USED NAMESPACES
//=============================================================================================================

using namespace Eigen;
using namespace FIFFLIB;
using namespace MNELIB;

#ifndef FAIL
#define FAIL -1
#endif

#ifndef OK
#define OK 0
#endif

#define MALLOC_30(x,t) (t *)malloc((x)*sizeof(t))
#define REALLOC_30(x,y,t) (t *)((x == NULL) ? malloc((y)*sizeof(t)) : realloc((x),(y)*sizeof(t)))

#define FREE_30(x) if ((char *)(x) != NULL) free((char *)(x))

#define FREE_CMATRIX_30(m) mne_free_cmatrix_30((m))
#define FREE_DCMATRIX_30(m) mne_free_dcmatrix_30((m))

void mne_free_cmatrix_30 (float **m)
{
    if (m) {
        FREE_30(*m);
        FREE_30(m);
    }
}

void mne_free_dcmatrix_30 (double **m)

{
    if (m) {
        FREE_30(*m);
        FREE_30(m);
    }
}

#define ALLOC_CMATRIX_30(x,y) mne_cmatrix_30((x),(y))

#define ALLOC_DCMATRIX_30(x,y) mne_dmatrix_30((x),(y))

static void matrix_error_30(int kind, int nr, int nc)

{
    if (kind == 1)
        printf("Failed to allocate memory pointers for a %d x %d matrix\n",nr,nc);
    else if (kind == 2)
        printf("Failed to allocate memory for a %d x %d matrix\n",nr,nc);
    else
        printf("Allocation error for a %d x %d matrix\n",nr,nc);
    if (sizeof(void *) == 4) {
        printf("This is probably because you seem to be using a computer with 32-bit architecture.\n");
        printf("Please consider moving to a 64-bit platform.");
    }
    printf("Cannot continue. Sorry.\n");
    exit(1);
}

float **mne_cmatrix_30(int nr,int nc)

{
    int i;
    float **m;
    float *whole;

    m = MALLOC_30(nr,float *);
    if (!m) matrix_error_30(1,nr,nc);
    whole = MALLOC_30(nr*nc,float);
    if (!whole) matrix_error_30(2,nr,nc);

    for(i=0;i<nr;i++)
        m[i] = whole + i*nc;
    return m;
}

double **mne_dmatrix_30(int nr, int nc)

{
    int i;
    double **m;
    double *whole;

    m = MALLOC_30(nr,double *);
    if (!m) matrix_error_30(1,nr,nc);
    whole = MALLOC_30(nr*nc,double);
    if (!whole) matrix_error_30(2,nr,nc);

    for(i=0;i<nr;i++)
        m[i] = whole + i*nc;
    return m;
}

//============================= mne_decompose.c =============================

int mne_decompose_eigen (double *mat,
                         double *lambda,
                         float  **vectors, /* Eigenvectors fit into floats easily */
                         int    dim)
/*
      * Compute the eigenvalue decomposition of
      * a symmetric matrix using the LAPACK routines
      *
      * 'mat' contains the lower triangle of the matrix
      */
{
    int    np  =   dim*(dim+1)/2;
    double *w    = MALLOC_30(dim,double);
    double *z    = MALLOC_30(dim*dim,double);
    double *work = MALLOC_30(3*dim,double);
    double *dmat = MALLOC_30(np,double);
    float  *vecp = vectors[0];

    const char   *uplo  = "U";
    const char   *compz = "V";
    int    info,k;
    int    one = 1;
    int    maxi;
    double scale;

//    maxi = idamax(&np,mat,&one);
// idamax workaround begin
    maxi = 0;
    for(int i = 0; i < np; ++i)
        if (std::fabs(mat[i]) > std::fabs(mat[maxi]))
            maxi = i;
// idamax workaround end

    scale = 1.0/mat[maxi];//scale = 1.0/mat[maxi-1];

    for (k = 0; k < np; k++)
        dmat[k] = mat[k]*scale;
//    dspev(compz,uplo,&dim,dmat,w,z,&dim,work,&info);

// dspev workaround begin
    MatrixXd dmat_tmp = MatrixXd::Zero(dim,dim);
    int idx = 0;
    for (int i = 0; i < dim; ++i) {
        for(int j = 0; j <= i; ++j) {
            dmat_tmp(i,j) = dmat[idx];
            dmat_tmp(j,i) = dmat[idx];
            ++idx;
        }
    }
    SelfAdjointEigenSolver<MatrixXd> es;
    es.compute(dmat_tmp);
    for ( int i = 0; i < dim; ++i )
        w[i] = es.eigenvalues()[i];

    idx = 0;
    for ( int j = 0; j < dim; ++j ) {
        for( int i = 0; i < dim; ++i ) {
            z[idx] = es.eigenvectors()(i,j);// Column Major
            ++idx;
        }
    }
// dspev workaround end

    info = 0;

    qDebug() << "!!!DEBUG ToDo: dspev(compz,uplo,&dim,dmat,w,z,&dim,work,&info);";

    FREE_30(work);
    if (info != 0)
        printf("Eigenvalue decomposition failed (LAPACK info = %d)",info);
    else {
        scale = 1.0/scale;
        for (k = 0; k < dim; k++)
            lambda[k] = scale*w[k];
        for (k = 0; k < dim*dim; k++)
            vecp[k] = z[k];
    }
    FREE_30(w);
    FREE_30(z);
    if (info == 0)
        return 0;
    else
        return -1;
}

double **mne_dmatt_dmat_mult2 (double **m1,double **m2, int d1,int d2,int d3)
/* Matrix multiplication
      * result(d1 x d3) = m1(d2 x d1)^T * m2(d2 x d3) */

{
    double **result = ALLOC_DCMATRIX_30(d1,d3);
#ifdef BLAS
    char  *transa = "N";
    char  *transb = "T";
    double zero = 0.0;
    double one  = 1.0;

    dgemm (transa,transb,&d3,&d1,&d2,
           &one,m2[0],&d3,m1[0],&d1,&zero,result[0],&d3);

    return result;
#else
    int j,k,p;
    double sum;

    for (j = 0; j < d1; j++)
        for (k = 0; k < d3; k++) {
            sum = 0.0;
            for (p = 0; p < d2; p++)
                sum = sum + m1[p][j]*m2[p][k];
            result[j][k] = sum;
        }
    return result;
#endif
}

//*************************************************************************************************************
//=============================================================================================================
// DEFINE MEMBER METHODS
//=============================================================================================================

MneCovMatrix::MneCovMatrix(int p_kind,
                           int p_ncov,
                           const QStringList& p_names,
                           double *p_cov,
                           double *p_cov_diag,
                           FiffSparseMatrix* p_cov_sparse)
:kind(p_kind)
,ncov(p_ncov)
,nproj(0)
,nzero(0)
,names(p_names)
,cov(p_cov)
,cov_diag(p_cov_diag)
,cov_sparse(p_cov_sparse)
,eigen(NULL)
,lambda(NULL)
,chol(NULL)
,inv_lambda(NULL)
,nfree(1)
,ch_class(NULL)
,proj(NULL)
,sss(NULL)
,nbad(0)
{

}

//*************************************************************************************************************

MneCovMatrix::~MneCovMatrix()
{
    FREE_30(cov);
    FREE_30(cov_diag);
    if(cov_sparse)
        delete cov_sparse;
    names.clear();
    FREE_CMATRIX_30(eigen);
    FREE_30(lambda);
    FREE_30(inv_lambda);
    FREE_30(chol);
    FREE_30(ch_class);
    if(proj)
        delete proj;
    if (sss)
        delete sss;
    bads.clear();
}

//*************************************************************************************************************

MneCovMatrix *MneCovMatrix::mne_dup_cov(MneCovMatrix *c)
{
    double       *vals;
    int          nval;
    int          k;
    MneCovMatrix* res;

    if (c->cov_diag)
        nval = c->ncov;
    else
        nval = (c->ncov*(c->ncov+1))/2;

    vals = MALLOC_30(nval,double);
    if (c->cov_diag) {
        for (k = 0; k < nval; k++)
            vals[k] = c->cov_diag[k];
        res = mne_new_cov(c->kind,c->ncov,c->names,NULL,vals);
    }
    else {
        for (k = 0; k < nval; k++)
            vals[k] = c->cov[k];
        res = mne_new_cov(c->kind,c->ncov,c->names,vals,NULL);
    }
    /*
        * Duplicate additional items
        */
    if (c->ch_class) {
        res->ch_class = MALLOC_30(c->ncov,int);
        for (k = 0; k < c->ncov; k++)
            res->ch_class[k] = c->ch_class[k];
    }
    res->bads = c->bads;
    res->nbad = c->nbad;
    res->proj = MneProjOp::mne_dup_proj_op(c->proj);
    res->sss  = new MneSssData(*(c->sss));

    return res;
}

//*************************************************************************************************************

int MneCovMatrix::mne_is_diag_cov(MneCovMatrix *c)
{
    return c->cov_diag != NULL;
}

//*************************************************************************************************************

int MneCovMatrix::mne_add_inv_cov(MneCovMatrix *c)
/*
          * Calculate the inverse square roots for whitening
          */
{
    double *src = c->lambda ? c->lambda : c->cov_diag;
    int k;

    if (src == NULL) {
        qCritical("Covariance matrix is not diagonal or not decomposed.");
        return FAIL;
    }
    c->inv_lambda = REALLOC_30(c->inv_lambda,c->ncov,double);
    for (k = 0; k < c->ncov; k++) {
        if (src[k] <= 0.0)
            c->inv_lambda[k] = 0.0;
        else
            c->inv_lambda[k] = 1.0/sqrt(src[k]);
    }
    return OK;
}

//*************************************************************************************************************

int MneCovMatrix::condition_cov(MneCovMatrix *c, float rank_threshold, int use_rank)
{
    double *scale  = NULL;
    double *cov    = NULL;
    double *lambda = NULL;
    float  **eigen = NULL;
    double **data1 = NULL;
    double **data2 = NULL;
    double magscale,gradscale,eegscale;
    int    nmag,ngrad,neeg,nok;
    int    j,k;
    int    res = FAIL;

    if (c->cov_diag)
        return OK;
    if (!c->ch_class) {
        qCritical("Channels not classified. Rank cannot be determined.");
        return FAIL;
    }
    magscale = gradscale = eegscale = 0.0;
    nmag = ngrad = neeg = 0;
    for (k = 0; k < c->ncov; k++) {
        if (c->ch_class[k] == MNE_COV_CH_MEG_MAG) {
            magscale += c->cov[mne_lt_packed_index(k,k)]; nmag++;
        }
        else if (c->ch_class[k] == MNE_COV_CH_MEG_GRAD) {
            gradscale += c->cov[mne_lt_packed_index(k,k)]; ngrad++;
        }
        else if (c->ch_class[k] == MNE_COV_CH_EEG) {
            eegscale += c->cov[mne_lt_packed_index(k,k)]; neeg++;
        }
#ifdef DEBUG
        fprintf(stdout,"%d ",c->ch_class[k]);
#endif
    }
#ifdef DEBUG
    fprintf(stdout,"\n");
#endif
    if (nmag > 0)
        magscale = magscale > 0.0 ? sqrt(nmag/magscale) : 0.0;
    if (ngrad > 0)
        gradscale = gradscale > 0.0 ? sqrt(ngrad/gradscale) : 0.0;
    if (neeg > 0)
        eegscale = eegscale > 0.0 ? sqrt(neeg/eegscale) : 0.0;
#ifdef DEBUG
    fprintf(stdout,"%d %g\n",nmag,magscale);
    fprintf(stdout,"%d %g\n",ngrad,gradscale);
    fprintf(stdout,"%d %g\n",neeg,eegscale);
#endif
    scale = MALLOC_30(c->ncov,double);
    for (k = 0; k < c->ncov; k++) {
        if (c->ch_class[k] == MNE_COV_CH_MEG_MAG)
            scale[k] = magscale;
        else if (c->ch_class[k] == MNE_COV_CH_MEG_GRAD)
            scale[k] = gradscale;
        else if (c->ch_class[k] == MNE_COV_CH_EEG)
            scale[k] = eegscale;
        else
            scale[k] = 1.0;
    }
    cov    = MALLOC_30(c->ncov*(c->ncov+1)/2.0,double);
    lambda = MALLOC_30(c->ncov,double);
    eigen  = ALLOC_CMATRIX_30(c->ncov,c->ncov);
    for (j = 0; j < c->ncov; j++)
        for (k = 0; k <= j; k++)
            cov[mne_lt_packed_index(j,k)] = c->cov[mne_lt_packed_index(j,k)]*scale[j]*scale[k];
    if (mne_decompose_eigen(cov,lambda,eigen,c->ncov) == 0) {
#ifdef DEBUG
        for (k = 0; k < c->ncov; k++)
            fprintf(stdout,"%g ",lambda[k]/lambda[c->ncov-1]);
        fprintf(stdout,"\n");
#endif
        nok = 0;
        for (k = c->ncov-1; k >= 0; k--) {
            if (lambda[k] >= rank_threshold*lambda[c->ncov-1])
                nok++;
            else
                break;
        }
        printf("\n\tEstimated covariance matrix rank = %d (%g)\n",nok,lambda[c->ncov-nok]/lambda[c->ncov-1]);
        if (use_rank > 0 && use_rank < nok) {
            nok = use_rank;
            fprintf(stderr,"\tUser-selected covariance matrix rank = %d (%g)\n",nok,lambda[c->ncov-nok]/lambda[c->ncov-1]);
        }
        /*
         * Put it back together
         */
        for (j = 0; j < c->ncov-nok; j++)
            lambda[j] = 0.0;
        data1 = ALLOC_DCMATRIX_30(c->ncov,c->ncov);
        for (j = 0; j < c->ncov; j++) {
#ifdef DEBUG
            mne_print_vector(stdout,NULL,eigen[j],c->ncov);
#endif
            for (k = 0; k < c->ncov; k++)
                data1[j][k] = sqrt(lambda[j])*eigen[j][k];
        }
        data2 = mne_dmatt_dmat_mult2(data1,data1,c->ncov,c->ncov,c->ncov);
#ifdef DEBUG
        printf(">>>\n");
        for (j = 0; j < c->ncov; j++)
            mne_print_dvector(stdout,NULL,data2[j],c->ncov);
        printf(">>>\n");
#endif
        /*
         * Scale back
         */
        for (k = 0; k < c->ncov; k++)
            if (scale[k] > 0.0)
                scale[k] = 1.0/scale[k];
        for (j = 0; j < c->ncov; j++)
            for (k = 0; k <= j; k++)
                if (c->cov[mne_lt_packed_index(j,k)] != 0.0)
                    c->cov[mne_lt_packed_index(j,k)] = scale[j]*scale[k]*data2[j][k];
        res = nok;
    }
    FREE_30(cov);
    FREE_30(lambda);
    FREE_CMATRIX_30(eigen);
    FREE_DCMATRIX_30(data1);
    FREE_DCMATRIX_30(data2);
    return res;
}

//*************************************************************************************************************

int MneCovMatrix::mne_decompose_eigen_cov_small(MneCovMatrix *c, float p_small, int use_rank)
/*
          * Do the eigenvalue decomposition
          */
{
    int   np,k,p,rank;
    float rank_threshold = 1e-6;

    if (p_small < 0)
        p_small = 1.0;

    if (!c)
        return OK;
    if (c->cov_diag)
        return mne_add_inv_cov(c);
    if (c->lambda && c->eigen) {
        fprintf(stderr,"\n\tEigenvalue decomposition had been precomputed.\n");
        c->nzero = 0;
        for (k = 0; k < c->ncov; k++, c->nzero++)
            if (c->lambda[k] > 0)
                break;
    }
    else {
        FREE_30(c->lambda); c->lambda = NULL;
        FREE_CMATRIX_30(c->eigen); c->eigen = NULL;

        if ((rank = condition_cov(c,rank_threshold,use_rank)) < 0)
            return FAIL;

        np = c->ncov*(c->ncov+1)/2;
        c->lambda = MALLOC_30(c->ncov,double);
        c->eigen  = ALLOC_CMATRIX_30(c->ncov,c->ncov);
        if (mne_decompose_eigen (c->cov,c->lambda,c->eigen,c->ncov) != 0)
            goto bad;
        c->nzero = c->ncov - rank;
        for (k = 0; k < c->nzero; k++)
            c->lambda[k] = 0.0;
        /*
         * Find which eigenvectors correspond to EEG/MEG
         */
        {
            float meglike,eeglike;
            int   nmeg,neeg;

            nmeg = neeg = 0;
            for (k = c->nzero; k < c->ncov; k++) {
                meglike = eeglike = 0.0;
                for (p = 0; p < c->ncov; p++)  {
                    if (c->ch_class[p] == MNE_COV_CH_EEG)
                        eeglike += std::fabs(c->eigen[k][p]);
                    else if (c->ch_class[p] == MNE_COV_CH_MEG_MAG || c->ch_class[p] == MNE_COV_CH_MEG_GRAD)
                        meglike += std::fabs(c->eigen[k][p]);
                }
                if (meglike > eeglike)
                    nmeg++;
                else
                    neeg++;
            }
            printf("\t%d MEG and %d EEG-like channels remain in the whitened data\n",nmeg,neeg);
        }
    }
    return mne_add_inv_cov(c);

bad : {
        FREE_30(c->lambda); c->lambda = NULL;
        FREE_CMATRIX_30(c->eigen); c->eigen = NULL;
        return FAIL;
    }
}

//*************************************************************************************************************

int MneCovMatrix::mne_decompose_eigen_cov(MneCovMatrix *c)

{
    return mne_decompose_eigen_cov_small(c,-1.0,-1);
}

//*************************************************************************************************************

int MneCovMatrix::mne_lt_packed_index(int j, int k)

{
    if (j >= k)
        return k + j*(j+1)/2;
    else
        return j + k*(k+1)/2;
}


/***********************************************************************
*
* OpenSkeletonFitting
* Skeleton fitting by the use of energy minimization
* Copyright (C) 2012 Norman Link <norman.link@gmx.net>
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program.  If not, see <http://www.gnu.org/licenses/>.
*
***********************************************************************/

#include "precompiled.h"
#include <iostream>
#include <stdio.h>
#include <boost/math/quaternion.hpp>
#include "Utils.h"

using namespace std;

namespace osf
{
	bool loadCvMat(const char* filename, cv::Mat& image, float* timestamp)
	{
		if (!filename || strlen(filename) == 0)
					throw Exception("invalid input filename");

		FILE* file = NULL;

#ifdef _WIN32
		errno_t err = fopen_s(&file, filename, "rb");
		if(!file || err) {
			cerr << "could not open file: " << filename << endl;
			return false;
		}
#elif __APPLE__ & __MACH__
		file = fopen(filename, "rb");
		if(!file || ferror(file)) {
			cerr << "could not open file: " << filename << endl;
			return false;
		}
#endif

		// read header
		int flags = 0;
		int headerSize = 0;
		cv::Size imgSize(0, 0);
		int dims = 0;
		cv::Size origSize(0, 0);
		cv::Point startPoint(0, 0);
		int size = 0;
		float ts = 0;

		char buffer[4];
		fread(&buffer, sizeof(char), 3, file);
		buffer[3] = '\0';

		if (strcmp(buffer, "CVM") != 0) {
			cerr << "file does not have cvm format: " << filename << endl;
			fclose(file);
			return 0;
		}

		// read header
		fread(&headerSize, sizeof(int), 1, file);
		fread(&flags, sizeof(int), 1, file);
		fread(&imgSize.width, sizeof(int), 1, file);
		fread(&imgSize.height, sizeof(int), 1, file);
		fread(&dims, sizeof(int), 1, file);
		fread(&origSize.width, sizeof(int), 1, file);
		fread(&origSize.height, sizeof(int), 1, file);
		fread(&startPoint.x, sizeof(int), 1, file);
		fread(&startPoint.y, sizeof(int), 1, file);
		fread(&size, sizeof(int), 1, file);

		// read timestamp
		fread(&ts, sizeof(float), 1, file);

		if (timestamp)
			*timestamp = ts;

		image = cv::Mat(origSize.height, origSize.width, CV_MAT_TYPE(flags));

		// set file pointer
		fseek(file, headerSize, SEEK_SET);

		// read actual data
		fread(image.data, sizeof(unsigned char), size, file);

		// NOTE: Can be speeded up: don't create new image with the specified roi
		// like here, but only set image.dataStart and image.dataEnd accordingly
		// to point to the image region specified by the roi. [9/15/2011 Norman]
		image = image(cv::Rect(startPoint.x, startPoint.y, imgSize.width, imgSize.height));

		fclose(file);

		return true;
	}

	bool saveCvMat(const char* filename, const cv::Mat& image)
	{
		if (!filename || strlen(filename) == 0)
			throw Exception("invalid input filename");

		FILE* file = NULL;

#ifdef _WIN32
		errno_t err = fopen_s(&file, filename, "wb");
		if(!file || err) {
			cerr << "could not create file: " << filename << endl;
			return false;
		}
#elif __APPLE__ & __MACH__
		file = fopen(filename, "wb");
		if(!file || ferror(file)) {
			cerr << "could not create file: " << filename << endl;
			return false;
		}
#endif

		// process roi information if available
		size_t step = image.step;
		cv::Size origSize = cv::Size(0, 0);
		cv::Point startPoint = cv::Point(0, 0);

		image.locateROI(origSize, startPoint);
		int size = image.elemSize() * origSize.width * origSize.height;

		// header size
		const int headerSize = 3 * sizeof(char) +
			10 * sizeof(int) +
			sizeof(float);

		// write identification string
		fwrite("CVM", sizeof(char), 3, file);

		// write header
		fwrite(&headerSize, sizeof(int), 1, file);
		fwrite(&image.flags, sizeof(int), 1, file);
		fwrite(&image.cols, sizeof(int), 1, file);
		fwrite(&image.rows, sizeof(int), 1, file);
		fwrite(&image.dims, sizeof(int), 1, file);
		fwrite(&origSize.width, sizeof(int), 1, file);
		fwrite(&origSize.height, sizeof(int), 1, file);
		fwrite(&startPoint.x, sizeof(int), 1, file);
		fwrite(&startPoint.y, sizeof(int), 1, file);
		fwrite(&size, sizeof(int), 1, file);

		// write timestamp
		float time = Timer::getTimestamp();
		fwrite(&time, sizeof(float), 1, file);

		// write data
		fwrite(image.data, sizeof(uchar), size, file);

		fclose(file);

		return true;
	}

	void matrix2Quat(const double* rot, double* quat)
	{
		// convert rotation matrix to quaternion (adapted from OgreQuaternion.cpp)

		double matrix[3][3] = {
			{rot[0], rot[1], rot[2]},
			{rot[3], rot[4], rot[5]},
			{rot[6], rot[7], rot[8]}
		};

		double trace = matrix[0][0] + matrix[1][1] + matrix[2][2];
		double root;

		if (trace > 0.0) {
			root = sqrt(trace + 1.0);
			quat[3] = 0.5 * root;
			root = 0.5 / root;

			quat[0] = (matrix[2][1] - matrix[1][2]) * root;
			quat[1] = (matrix[0][2] - matrix[2][0]) * root;
			quat[2] = (matrix[1][0] - matrix[0][1]) * root;
		}
		else {
			static size_t next[3] = {1, 2, 0};
			size_t i = 0;
			if (matrix[1][1] > matrix[0][0])
				i = 1;
			if (matrix[2][2] > matrix[i][i])
				i = 2;
			size_t j = next[i];
			size_t k = next[j];

			root = sqrt(matrix[i][i] - matrix[j][j] - matrix[k][k] + 1.0);
			double* apkQuat[3] = {&quat[0], &quat[1], &quat[2]};
			*apkQuat[i] = 0.5 * root;
			root = 0.5 / root;
			quat[3] = (matrix[k][j] - matrix[j][k]) * root;
			*apkQuat[j] = (matrix[j][i] + matrix[i][j]) * root;
			*apkQuat[k] = (matrix[k][i] + matrix[i][k]) * root;
		}
	}

	void matrix2Quat(const cv::Mat& rot, quaternion<double>& quat)
	{
		if (rot.type() != CV_64F)
			throw Exception("invalid type");

		double q[4] = { quat.R_component_2(),
			quat.R_component_3(),
			quat.R_component_4(),
			quat.R_component_1() };

		matrix2Quat((double*)&rot.data[0], q);

		quat = boost::math::quaternion<double>(q[3], q[0], q[1], q[2]);
	}

	void quat2Matrix(const double* quat, double* rot)
	{
		// convert quaternion to rotation matrix (adapted from OgreQuaternion.cpp)

		double tx = quat[0] + quat[0];
		double ty = quat[1] + quat[1];
		double tz = quat[2] + quat[2];
		double twx = tx * quat[3];
		double twy = ty * quat[3];
		double twz = tz * quat[3];
		double txx = tx * quat[0];
		double txy = ty * quat[0];
		double txz = tz * quat[0];
		double tyy = ty * quat[1];
		double tyz = tz * quat[1];
		double tzz = tz * quat[2];

		double matrix[3][3];

		matrix[0][0] = 1.0 - (tyy + tzz);
		matrix[0][1] = txy - twz;
		matrix[0][2] = txz + twy;
		matrix[1][0] = txy + twz;
		matrix[1][1] = 1.0 - (txx + tzz);
		matrix[1][2] = tyz - twx;
		matrix[2][0] = txz - twy;
		matrix[2][1] = tyz + twx;
		matrix[2][2] = 1.0 - (txx + tyy);

		rot[0] = matrix[0][0];
		rot[1] = matrix[0][1];
		rot[2] = matrix[0][2];
		rot[3] = matrix[1][0];
		rot[4] = matrix[1][1];
		rot[5] = matrix[1][2];
		rot[6] = matrix[2][0];
		rot[7] = matrix[2][1];
		rot[8] = matrix[2][2];
	}

	void quat2Matrix(const quaternion<double>& quat, cv::Mat& rot)
	{
		if (!rot.empty() && rot.type() != CV_64F)
			throw Exception("invalid type");
		else
			rot = cv::Mat(3, 3, CV_64FC1);

		double q[4] = { quat.R_component_2(),
			quat.R_component_3(),
			quat.R_component_4(),
			quat.R_component_1() };

		quat2Matrix(q, (double*)&rot.data[0]);
	}

	void euler2Quat(quaternion<double>& quat, double angleX, double angleY, double angleZ)
	{
		double r = DEG2RAD(angleX / 2.0);
		double p = DEG2RAD(angleY / 2.0);
		double y = DEG2RAD(angleZ / 2.0);

		double sinp = sin(p);
		double siny = sin(y);
		double sinr = sin(r);
		double cosp = cos(p);
		double cosy = cos(y);
		double cosr = cos(r);

		quat = quaternion<double>(
			cosr * cosp * cosy + sinr * sinp * siny,
			sinr * cosp * cosy - cosr * sinp * siny,
			cosr * sinp * cosy + sinr * cosp * siny,
			cosr * cosp * siny - sinr * sinp * cosy);

		// normalize
		quat /= norm(quat);

		/*quaternion<double> quatX(0, 0, 0, 0);
		axis2Quat(quatX, cv::Point3d(1, 0, 0), angleX);

		quaternion<double> quatY(0, 0, 0, 0);
		axis2Quat(quatY, cv::Point3d(0, 1, 0), angleY);

		quaternion<double> quatZ(0, 0, 0, 0);
		axis2Quat(quatZ, cv::Point3d(0, 0, 1), angleZ);

		quat = quatX * quatY * quatZ;
		quat /= norm(quat);*/
	}

	void quat2Euler(const quaternion<double>& quat, double& angleX, double& angleY, double& angleZ)
	{
		quaternion<double> myQuat = quat;
		myQuat /= norm(myQuat);

		cv::Point3d euler(0, 0, 0);

		double qW = myQuat.R_component_1();
		double qX = myQuat.R_component_2();
		double qY = myQuat.R_component_3();
		double qZ = myQuat.R_component_4();

		double test = (qW * qY - qZ * qX);
		double unit = qX * qX + qY * qY + qZ * qZ + qW * qW;

		// handle singularities
		if (test > 0.4999999 * unit) {
			euler.x = 2 * atan2(qX, qW);
			euler.y = M_PI / 2.0;
			euler.z = 0;
		}
		else if (test < -0.4999999 * unit) {
			euler.x = 2 * atan2(qX, qW);
			euler.y = -M_PI / 2.0;
			euler.z = 0;
		}
		else {
			euler.x = atan2(2 * (qW * qX + qY * qZ), 1 - 2 * (qX * qX + qY * qY));
			euler.y = asin(2 * test);
			euler.z = atan2(2 * (qW * qZ + qX * qY), 1 - 2 * (qY * qY + qZ * qZ));
		}

		angleX = RAD2DEG(euler.x);
		angleY = RAD2DEG(euler.y);
		angleZ = RAD2DEG(euler.z);
	}

	void axis2Quat(quaternion<double>& quat, const cv::Point3d& axis, double angle)
	{
		cv::Point3d myAxis = axis;
		myAxis *= (1.0 / cv::norm(myAxis));

		double halfAngle = DEG2RAD(angle / 2.0);
		double sinAngle = sin(halfAngle);

		quat = quaternion<double>(cos(halfAngle),
			myAxis.x * sinAngle, myAxis.y * sinAngle, myAxis.z * sinAngle);
	}

	void quat2Axis(const quaternion<double>& quat, cv::Point3d& axis, double& angle)
	{
		quaternion<double> myQuat = quat;
		double qw = myQuat.R_component_1();

		// normalize
		if (qw > 1.0)
			myQuat /= norm(myQuat);

		angle = RAD2DEG(2.0 * acos(qw));
		double s = sqrt(1.0 - qw * qw);

		if (s < 0.0001) {
			// avoid divbyzero, any arbitrary axis is valid
			axis.x = 0;
			axis.y = 1;
			axis.z = 0;
		}
		else {
			axis.x = myQuat.R_component_2() / s;
			axis.y = myQuat.R_component_3() / s;
			axis.z = myQuat.R_component_4() / s;
		}
	}

	void quatRotate(const quaternion<double>& quat, cv::Point3d& point)
	{
		quaternion<double> pointTemp(0, point.x, point.y, point.z);
		pointTemp = conj(quat) * pointTemp * quat;
		point = cv::Point3d(pointTemp.R_component_2(), pointTemp.R_component_3(),
			pointTemp.R_component_4());
	}

	void quatRotateInv(const quaternion<double>& quat, cv::Point3d& point)
	{
		quaternion<double> pointTemp(0, point.x, point.y, point.z);
		pointTemp = quat * pointTemp * conj(quat);
		point = cv::Point3d(pointTemp.R_component_2(), pointTemp.R_component_3(),
			pointTemp.R_component_4());
	}
}