pedigree.h

/* Spell-QTL  Software suite for the QTL analysis of modern datasets.
 * Copyright (C) 2016,2017  Damien Leroux <damien.leroux@inra.fr>, Sylvain Jasson <sylvain.jasson@inra.fr>
 *
 * 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/>.
 */

#ifndef _SPELL_BAYES_CSV_H_
#define _SPELL_BAYES_CSV_H_

#include <iostream>
#include "file.h"
#include <string>
#include <sstream>
#include <stdexcept>
#include <vector>
#include <utility>
#include <unordered_set>


/*#include "permutation.h"*/
/*#include "symmetry.h"*/
#include "geno_matrix.h"
#include "linear_combination.h"
#include "pedigree_tree.h"
#include "bayes/output.h"


struct bn_label_type {
    char first;
    char second;
    char first_allele;
    char second_allele;
    bn_label_type() : first(0), second(0), first_allele(0), second_allele(0) {}
    bn_label_type(int) : first(0), second(0), first_allele(0), second_allele(0) {}
    bn_label_type(char f, char s, char fa, char sa)
        : first(f), second(s), first_allele(fa), second_allele(sa)
    {}

    int compact() const { return *(int*) this; }
    int& compact() { return *(int*) this; }

    friend std::ostream& operator << (std::ostream& os, const bn_label_type& bl)
    {
        if (bl.second != GAMETE_EMPTY) {
            return os << bl.first << ((int) bl.first_allele) << bl.second << ((int) bl.second_allele);
        } else {
            return os << bl.first << ((int) bl.first_allele);
        }
    }

    bool operator < (const bn_label_type& other) const
    {
        /*return first < other.first || (first == other.first*/
            /*&& (second < other.second || (second == other.second*/
            /*&& (first_allele < other.first_allele || (first_allele == other.first_allele*/
            /*&& second_allele < other.second_allele)))));*/
        return (*(int*)this) < (*(int*) &other);
    }

    bool operator == (const bn_label_type& other) const
    {
        return (*(int*)this) == (*(int*) &other);
    }

    bool operator != (const bn_label_type& other) const
    {
        return (*(int*)this) != (*(int*) &other);
    }
};

typedef combination_type<int, bn_label_type> genotype_comb_type;


template <typename Arg>
int read_field(std::stringstream& s, char sep, Arg& arg)
{
    std::string field;
    std::getline(s, field, sep);
    /*MSG_DEBUG("CSV FIELD |" << field << "|");*/
    std::stringstream ss(field);
    ss >> arg;
    return 0;
}


#define do_with_arg_pack(_expr) do { using _ = int[]; (void)_{0, ((_expr), void(), 0)...}; } while(0)

template <typename... Args>
void read_csv_line(std::istream& is, char sep, Args&... args)
{
    std::string line;
    std::getline(is, line);
    /*MSG_DEBUG("CSV LINE |" << line << "|");*/
    std::stringstream ss(line);
    do_with_arg_pack(read_field(ss, sep, args));
}





struct pedigree_item {
    std::string gen_name;
    int id;
    int p1;
    int p2;

    pedigree_item() : gen_name(), id(0), p1(0), p2(0) {}

    pedigree_item(const char* gn, int i, int a, int b)
        : gen_name(gn), id(i), p1(a), p2(b)
    {}

    pedigree_item(std::istream& is, char field_sep=';')
        : gen_name()
    {
        id = p1 = p2 = 0;
        read_csv_line(is, field_sep, gen_name, id, p1, p2);
        if (id == p1 && id == p2 && id == 0) {
            return;
        }
        if (id <= p1 || id <= p2) {
            throw std::runtime_error("Bad ID! ID must be greater than p1 AND p2");
            /*MSG_DEBUG("BAD ID!! " << id << " must be greater than " << p1 << " AND " << p2);*/
        }
    }

    bool is_ancestor() const { return p1 == 0 && p2 == 0; }
    bool is_self() const { return p1 > 0 && p1 == p2; }
    bool is_cross() const { return p1 > 0 && p2 > 0 && p1 != p2; }
    bool is_dh() const { return p1 != p2 && p1 >= 0 && p2 >= 0 && (p1 * p2) == 0; }
    bool is_bullshit() const { return !(is_ancestor() || is_self() || is_cross() || is_dh()); }
};


std::vector<pedigree_item>
read_csv(const std::string& pedigree_file, char field_sep=';');



typedef std::map<size_t, size_t> ancestor_list_type;


inline
ancestor_list_type reentrants(const ancestor_list_type& a)
{
    ancestor_list_type ret;
    for (const auto& kv: a) {
        if (kv.second > 1) {
            ret.emplace(kv);
        }
    }
    return ret;
}


inline
ancestor_list_type operator + (const ancestor_list_type& a1, const ancestor_list_type& a2)
{
    ancestor_list_type ret(a1);
    for (const auto& kv: a2) {
        ret[kv.first] += kv.second;
    }
    return ret;
}


inline
ancestor_list_type operator / (const ancestor_list_type& a, const ancestor_list_type& restr)
{
    ancestor_list_type ret;
    for (const auto& kv: a) {
        auto i = restr.find(kv.first);
        if (i != restr.end()) {
            ret.emplace(kv.first, std::min(kv.second, i->second));
        }
    }
    return ret;
}


inline
ancestor_list_type operator % (const ancestor_list_type& a, const ancestor_list_type& restr)
{
    ancestor_list_type ret;
    for (const auto& kv: a) {
        if (restr.find(kv.first) != restr.end()) {
            ret.emplace(kv);
        }
    }
    return ret;
}


inline
ancestor_list_type operator - (const ancestor_list_type& a, const ancestor_list_type& restr)
{
    ancestor_list_type ret;
    for (const auto& kv: a) {
        auto it = restr.find(kv.first);
        if (it == restr.end()) {
            ret.emplace(kv);
        } else if (kv.second > it->second) {
            ret.emplace(kv.first, kv.second - it->second);
        }
    }
    return ret;
}


inline
ancestor_list_type operator * (const ancestor_list_type& a, size_t weight)
{
    ancestor_list_type ret;
    for (const auto& kv: a) {
        ret.emplace(kv.first, kv.second * weight);
    }
    return ret;
}


inline
std::ostream& operator << (std::ostream& os, const ancestor_list_type& a)
{
    auto i = a.begin();
    auto j = a.end();
    if (i != j) {
        os << i->first << ':' << i->second;
        for (++i; i != j; ++i) {
            os << ' ' << i->first << ':' << i->second;
        }
    } else {
        os << "empty";
    }
    return os;
}


#if 0
label_type operator * (label_type a, label_type b)
{
    label_type ret;
    if (a.second == GAMETE_EMPTY) {
        ret = {a.first, b.first};
    } else {
        if (b.second == GAMETE_EMPTY) {
            ret = {SELECT(a, b.first), GAMETE_EMPTY};
        } else {
            ret = {SELECT(a, b.first), SELECT(a, b.second)};
        }
    }
    /*MSG_DEBUG("" << a << " * " << b << " = " << ret);*/
    return ret;
}
#endif


#define SELECT_A(__p, __b) ((__b) == GAMETE_R ? (__p).second_allele : (__p).first_allele)

inline
bn_label_type operator * (bn_label_type a, bn_label_type b)
{
    bn_label_type ret;
    if (a.second == GAMETE_EMPTY) {
        ret = {a.first, b.first, a.first_allele, b.first_allele};
    } else {
        if (b.second == GAMETE_EMPTY) {
            ret = {SELECT(a, b.first), GAMETE_EMPTY, SELECT_A(a, b.first), 0};
        } else {
            ret = {SELECT(a, b.first), SELECT(a, b.second), SELECT_A(a, b.first), SELECT_A(a, b.second)};
        }
    }
    /*MSG_DEBUG("" << a << " * " << b << " = " << ret);*/
    return ret;
}


template <typename F> struct make_one;
template <> struct make_one<MatrixXd> {
    static MatrixXd& _(bool der) {
        static MatrixXd one = MatrixXd::Ones(1, 1);
        static MatrixXd zero = MatrixXd::Zero(1, 1);
        return der ? zero : one;
    }
};
template <> struct make_one<VectorXd> {
    static VectorXd& _(bool der) {
        static VectorXd one = VectorXd::Ones(1);
        static VectorXd zero = VectorXd::Zero(1);
        return der ? zero : one;
    }
};


template <typename PARENT_TYPE, typename STATE_TYPE>
struct rw_comb : public rw_any<rw_comb<PARENT_TYPE, STATE_TYPE>> {
    typedef combination_type<PARENT_TYPE, STATE_TYPE> comb_type;

    virtual ~rw_comb() {}

    using rw_any<rw_comb<PARENT_TYPE, STATE_TYPE>>::fourcc;
    using rw_any<rw_comb<PARENT_TYPE, STATE_TYPE>>::ref;
    using rw_any<rw_comb<PARENT_TYPE, STATE_TYPE>>::operator ();

    void operator () (ifile& ifs, bn_label_type& l) { l.compact() = read_int(ifs); }
    void operator () (ofile& ofs, bn_label_type& l) { write_int(ofs, l.compact()); }

    void operator () (ifile& fs, typename comb_type::key_type& key) { ref() (fs, key.parent); ref() (fs, key.state); }
    void operator () (ofile& fs, typename comb_type::key_type& key) { ref() (fs, key.parent); ref() (fs, key.state); }

    void operator () (ifile& fs, typename comb_type::key_list& keys) { ref() (fs, keys.keys); }
    void operator () (ofile& fs, typename comb_type::key_list& keys) { ref() (fs, keys.keys); }

    void operator () (ifile& fs, typename comb_type::element_type& elem) { ref() (fs, elem.keys); ref() (fs, elem.coef); }
    void operator () (ofile& fs, typename comb_type::element_type& elem) { ref() (fs, elem.keys); ref() (fs, elem.coef); }

    void operator () (ifile& fs, comb_type& comb)
    {
        if (fourcc(fs, "COMB")) { return; }
        ref() (fs, comb.m_combination);
    }

    void operator () (ofile& fs, comb_type& comb)
    {
        if (fourcc(fs, "COMB")) { return; }
        ref() (fs, comb.m_combination);
    }
};


struct rw_tree : public rw_any<rw_tree> {
    virtual ~rw_tree() {}

    using rw_any<rw_tree>::fourcc;
    using rw_any<rw_tree>::ref;
    using rw_any<rw_tree>::operator ();

    void operator () (ifile& fs, pedigree_tree_type::node_type& node) { ref() (fs, node.p1); ref() (fs, node.p2); }
    void operator () (ofile& fs, pedigree_tree_type::node_type& node) { ref() (fs, node.p1); ref() (fs, node.p2); }

    template <typename STREAM_TYPE, typename TREE_TYPE>
        void tree_io_impl(STREAM_TYPE& fs, TREE_TYPE&& tree)
        {
            ref() (fs, tree.m_leaves);
            ref() (fs, tree.m_roots);
            ref() (fs, tree.m_nodes);
            ref() (fs, tree.m_must_recompute);
            ref() (fs, tree.m_node_number_to_ind_number);
            ref() (fs, tree.m_ind_number_to_node_number);
            ref() (fs, tree.m_original_ordering);
        }

    void operator () (ifile& fs, pedigree_tree_type& tree)
    {
        tree_io_impl(fs, tree);
    }

    void operator () (ofile& fs, pedigree_tree_type& tree)
    {
        tree_io_impl(fs, tree);
    }

    void operator () (ifile& fs, pedigree_item& pi)
    {
        (*this)(fs, pi.gen_name);
        (*this)(fs, pi.id);
        (*this)(fs, pi.p1);
        (*this)(fs, pi.p2);
    }

    void operator () (ofile& fs, pedigree_item& pi)
    {
        (*this)(fs, pi.gen_name);
        (*this)(fs, pi.id);
        (*this)(fs, pi.p1);
        (*this)(fs, pi.p2);
    }
};


/* TODO extraire l'arbre du pedigree
 * TODO opérations sur l'arbre :
 * TODO - insérer un nouveau noeud étant donné {P1, P2} (Pi étant soit néant soit un noeud existant)
 * TODO - extraire sous-arbre étant donné {RACINE, {FEUILLES}}
 * TODO - comparer deux arbres
 * TODO - pour deux arbres comparables, déterminer la rotation du second pour matcher le premier
 */

/*
 * pedigree_type: implements all facilities to compute proper geno_matrices for any pedigree, including reentrant individuals.
 */
struct pedigree_type {
    /*
     * original data
     */
    std::vector<pedigree_item> items;

    /*
     * pedigree tree implementation
     */
    pedigree_tree_type tree;

    /*
     * pedigree tree metadata
     */
    typedef size_t geno_matrix_index_type;
    typedef int individual_index_type;
    std::vector<std::shared_ptr<geno_matrix>> generations;
    std::vector<geno_matrix_index_type> node_generations;
    std::map<individual_index_type, char> ancestor_letters;
    std::map<char, std::string> ancestor_names;
    std::map<geno_matrix_index_type, std::string> generation_names;
    std::map<int, int> m_id;
    bool with_LC;
    /*std::vector<VectorLC> LC;*/

    /*
     * geno_matrix cache to avoid recomputing identical generations
     */
    std::map<geno_matrix_index_type, geno_matrix_index_type> cache_gamete;
    std::map<std::pair<geno_matrix_index_type, geno_matrix_index_type>, geno_matrix_index_type> cache_geno;

    /*
     * geno_matrix database
     */

    std::map<std::string, std::set<geno_matrix_index_type>> geno_matrix_by_generation_name;
    std::map<std::string, std::vector<individual_index_type>> individuals_by_generation_name;
    std::map<individual_index_type, const std::string*> generation_name_by_individual;

    /*
     * overlump control
     */

    size_t max_states;

    /*
     * BN metadata
     */
    size_t n_alleles;

    /*
     * Metadata for XML output and recreating command line
     */
    std::string filename;

    /*
     * Actual output: LC and factors for bayesian network
     */
    std::vector<std::vector<gencomb_type>> LC;
    std::vector<std::vector<std::map<bn_label_type, genotype_comb_type>>> factor_messages;
    std::vector<std::vector<size_t>> individuals_in_factors;
    /* i-th element means the i-th variable receives a message through this factor from variables in genotype_comb_type:keys */

    /*
     * default ctor
     */
    pedigree_type()
        : tree(), node_generations(), ancestor_letters(), ancestor_names(), generation_names(), m_id(),
          with_LC(true),
          cache_gamete(), cache_geno(),
          max_states(NONE),
          n_alleles(1),
          filename("<none>"),
          LC(),
          factor_messages()
    {
        __init();
    }

    void __init()
    {
        generations.emplace_back();
    }

#if 0
    /*
     * prealloc ctor
     */
    pedigree_type(size_t n_ind)
    {
        n_alleles = 1;
        max_states = NONE;
        nodes.reserve(3 * n_ind);
        /*ind_generations.reserve(n_ind);*/
        ind_number_to_node_number.reserve(n_ind);
        __init();
    }

    size_t last_node_index() const { return tree.size() - 1; }
#endif

    individual_index_type spawn_gamete(const std::string&, int parent_node)
    {
        int n = tree.add_node(parent_node);
        node_generations.emplace_back(node_generations[parent_node]);
        /*MSG_DEBUG_INDENT_EXPR("[compute " << gamete_name << " gamete] ");*/
        /*compute_generation(n);*/
        /*compute_LC(n);*/
        /*MSG_DEBUG_DEDENT;*/
        return n;
    }

    individual_index_type spawn(const std::string& generation_name, std::initializer_list<individual_index_type> parents)
    {
        individual_index_type ind = tree.next_ind_idx();
        switch (parents.size()) {
            case 0: /* ancestor */
                {
                    MSG_DEBUG_INDENT_EXPR("[compute gen #" << ind << "] ");
                    MSG_DEBUG("ANCESTOR");
                    int n = tree.add_node();
                    MSG_DEBUG("node=" << n << " ind=" << ind);
                    compute_generation(generation_name, n);
                    compute_LC(n);
                    MSG_DEBUG_DEDENT;
                }
                break;
            case 1: /* doubling */
                {
                    MSG_DEBUG_INDENT_EXPR("[compute gen #" << ind << "] ");
                    MSG_DEBUG("DOUBLING");
                    individual_index_type p1 = *parents.begin();
                    int g1 = spawn_gamete("M", tree.ind2node(p1));
                    int n = tree.add_node(g1, g1);
                    MSG_DEBUG("node=" << n << " ind=" << ind);
                    compute_generation(generation_name, n);
                    compute_LC(n);
                    /*compute_data_for_bn(n);*/
                    MSG_DEBUG_DEDENT;
                }
                break;
            case 2: /* crossing & selfing */
                {
                    MSG_DEBUG_INDENT_EXPR("[compute gen #" << ind << "] ");
                    MSG_DEBUG("CROSSING/SELFING");
                    auto i = parents.begin();
                    individual_index_type p1 = *i++;
                    int n1 = tree.ind2node(p1);
                    individual_index_type p2 = *i;
                    int n2 = tree.ind2node(p2);
                    /*MSG_DEBUG("p1=" << p1 << " p2=" << p2 << " n1=" << n1 << " n2=" << n2);*/
                    int g1 = spawn_gamete("M", n1);
                    int g2 = spawn_gamete("F", n2);
                    int n = tree.add_node(g1, g2);
                    MSG_DEBUG("node=" << n << " ind=" << ind);
                    compute_generation(generation_name, n);
                    compute_LC(n);
                    /*compute_data_for_bn(n);*/
                    MSG_DEBUG_DEDENT;
                }
                break;
            default:;
        };
        return ind;
    }

    /*individual_index_type crossing(std::string& generation_name, individual_index_type p1, individual_index_type p2) { return spawn(generation_name, {p1, p2}); }*/
    /*individual_index_type selfing(std::string& generation_name, individual_index_type p1) { return spawn(generation_name, {p1, p1}); }*/
    /*individual_index_type dh(std::string& generation_name, individual_index_type p1) { return spawn(generation_name, {p1}); }*/
    /*individual_index_type ancestor(std::string& generation_name) { return spawn(generation_name, {}); }*/

    individual_index_type fill_db(const std::string& name, individual_index_type ind)
    {
        geno_matrix_by_generation_name[name].insert(get_gen_index(ind));
        individuals_by_generation_name[name].push_back(ind);
        auto it = individuals_by_generation_name.find(name);
        generation_name_by_individual[ind] = &it->first;
        return ind;
    }

    individual_index_type crossing(const std::string& name, individual_index_type p1, individual_index_type p2) { return fill_db(name, spawn(name, {p1, p2})); }
    individual_index_type selfing(const std::string& name, individual_index_type p1) { return fill_db(name, spawn(name, {p1, p1})); }
    individual_index_type dh(const std::string& name, individual_index_type p1) { return fill_db(name, spawn(name, {p1})); }
    individual_index_type ancestor(const std::string& name) { return fill_db(name, spawn(name, {})); }

#if 0
    void propagate_symmetries(int n, geno_matrix& gen)
    {
        MSG_DEBUG_INDENT_EXPR("[propagate symmetries #" << n << "] ");
        std::vector<int> in, out;
        auto expr = tree.extract_expression(n, in, out);
        std::vector<pedigree_tree_type> input_trees;
        input_trees.reserve(in.size());
        for (int t: in) {
            input_trees.emplace_back(tree.extract_subtree(t));
        }
        /*auto recompute = tree.get_deep_recompute_vec(n);*/
        /*MSG_DEBUG("RECOMPUTE: " << recompute);*/
        auto get_lumper
            = [&, this] (int node) -> MatrixXb
            {
                /*if (tree[node].is_gamete() || recompute[node]) {*/
                    /*MSG_DEBUG("NIL lumper for node #" << node << " because" << (recompute[node] && tree[node].is_gamete() ? " recompute flag is set and it is a gamete" : tree[node].is_gamete() ? " it is a gamete" : " recompute flag is set"));*/
                    /*return {};*/
                /*}*/
                return get_node_gen(node)->collect.cast<bool>();
            };
        symmetry_propagator sp(expr);
        MSG_DEBUG_INDENT_EXPR("[SYMMETRIES] ");
        gen.symmetries = sp.compute_propagated_symmetries(
                [&] (int i) -> const pedigree_tree_type& { MSG_DEBUG("request subtree for input #" << i); return input_trees[i]; },
                [&] (int node) -> const symmetry_group_type& { MSG_DEBUG("request symmetries for node #" << node); return get_node_gen(node)->symmetries; },
                [&] (int node) -> const symmetry_group_type& { MSG_DEBUG("request latent symmetries for node #" << node); return get_node_gen(node)->latent_symmetries; },
                gen.labels, gen.inf_mat,
                get_lumper
                );
        MSG_DEBUG_DEDENT;
        MSG_DEBUG_INDENT_EXPR("[LATENT SYMMETRIES] ");
        auto temp = sp.compute_propagated_latent_symmetries(
                [&] (int i) -> const pedigree_tree_type& { MSG_DEBUG("request subtree for input #" << i); return input_trees[i]; },
                [&] (int node) -> const symmetry_group_type& { MSG_DEBUG("request symmetries for node #" << node); return get_node_gen(node)->symmetries; },
                [&] (int node) -> const symmetry_group_type& { MSG_DEBUG("request latent symmetries for node #" << node); return get_node_gen(node)->latent_symmetries; },
                gen.labels, gen.inf_mat,
                get_lumper
                );
        MSG_DEBUG_DEDENT;
        MSG_DEBUG(temp);
        gen.latent_symmetries = temp - gen.symmetries;
        MSG_DEBUG_INDENT_EXPR("[AFTER SYMMETRY PROPAGATION] ");
        MSG_DEBUG(gen);
        MSG_DEBUG_DEDENT;
        MSG_DEBUG_DEDENT;
    }
#endif

    void compute_generation(const std::string& generation_name, int n)
    {
        MSG_DEBUG("Computing generation for node " << tree.make_node_label(n));
        /*MSG_DEBUG(render_tree());*/

        int np1 = tree.get_p1(n);
        int np2 = tree.get_p2(n);
        node_generations.emplace_back(generations.size());
        geno_matrix new_gen;
        geno_matrix_index_type* cached_gen = NULL;
        if (np1 == NONE && np2 == NONE) {
            /* ancestor */
            char l = 'a' + ancestor_letters.size();
            ancestor_letters[n] = l;
            ancestor_names[l] = generation_name;
            new_gen = ancestor_matrix(generation_name, l);
        } else if (np2 == NONE) {
            /* gamete */
            auto gp = node_generations[np1];
            auto& cache = cache_gamete;
            cached_gen = &cache[gp];
            if (*cached_gen) {
                MSG_DEBUG("GENERATION HAS ALREADY BEEN COMPUTED");
                node_generations[n] = *cached_gen;
                return;
            }
            new_gen = kronecker(*generations[gp], gamete);
            new_gen.name = generation_name;
            /*MSG_DEBUG("TMP GAMETE GEN");*/
            /*MSG_DEBUG(new_gen);*/
        } else {
            /*auto ngp1 = node_generations[np1];*/
            /*auto ngp2 = node_generations[np2];*/
            /*auto gp1 = generations[ngp1];*/
            /*auto gp2 = generations[ngp2];*/
            MSG_DEBUG("Child of " << tree.node2ind(tree.get_p1(np1)) << " and " << tree.node2ind(tree.get_p1(np2)));

            /* use node's grandparents generations, not parents (no gamete generation) */
            size_t g1 = node_generations[tree.get_p1(np1)];
            size_t g2 = node_generations[tree.get_p1(np2)];
            cached_gen = &cache_geno[{g1, g2}];
            if (*cached_gen) {
                MSG_DEBUG("GENERATION HAS ALREADY BEEN COMPUTED");
                node_generations[n] = *cached_gen;
                return;
            }

            const auto& recompute = tree.get_recompute_vec(n);

            std::vector<bool> visited_clear(recompute.size(), false);
            std::vector<bool> visited;

            new_gen.name = generation_name;

            visited = visited_clear;
            MSG_DEBUG("COMPUTING INF_MAT");
            new_gen.inf_mat = eval(n, &geno_matrix::inf_mat, &pedigree_type::kron_d, recompute, visited);
            MSG_DEBUG(MATRIX_SIZE(new_gen.inf_mat));
            visited = visited_clear;
            MSG_DEBUG("COMPUTING DIAG");
            new_gen.diag = eval(n, &geno_matrix::diag, &pedigree_type::kron_d_diag, recompute, visited);
            visited = visited_clear;
            MSG_DEBUG("COMPUTING STAT_DIST");
            new_gen.stat_dist = eval(n, &geno_matrix::stat_dist, &pedigree_type::kron, recompute, visited);
            visited = visited_clear;
            MSG_DEBUG("COMPUTING P");
            new_gen.p = eval(n, &geno_matrix::p, &pedigree_type::kron, recompute, visited);
            visited = visited_clear;
            MSG_DEBUG("COMPUTING P_INV");
            new_gen.p_inv = eval(n, &geno_matrix::p_inv, &pedigree_type::kron_d, recompute, visited);
            /*new_gen.labels = eval_labels(n, recompute, visited_clear);*/
            new_gen.labels = eval_vector(n, recompute, &pedigree_type::get_labels, reentrant_label);
            new_gen.variant = (tree[n].is_genotype()
                               ? Geno
                               : Gamete);
            new_gen.dispatch = new_gen.collect = MatrixXd::Identity(new_gen.diag.size(), new_gen.diag.size());
            MSG_DEBUG(MATRIX_SIZE(new_gen.inf_mat));
            MSG_DEBUG(MATRIX_SIZE(new_gen.p));
            MSG_DEBUG(MATRIX_SIZE(new_gen.p_inv));
            MSG_DEBUG(MATRIX_SIZE(new_gen.diag));
            MSG_DEBUG("new_gen.labels.size()=" << new_gen.labels.size());
            MSG_QUEUE_FLUSH();

            /*if (!(ind_number_to_node_number.size() == 9 && tree.size() == 23)) {*/
                /*MSG_DEBUG("PROPAGATING SYMMETRIES");*/
                /*propagate_symmetries(new_gen, recompute, n);*/
                /*study_expression_symmetries(new_gen);*/
                /*complete_symmetries(new_gen);*/
                /*MSG_DEBUG("COMPUTING LATENT SYMMETRY");*/
            /*}*/
        }
        node_generations[n] = generations.size();
        generations.emplace_back(new geno_matrix());
        /*MSG_DEBUG("BEFORE LUMPING");*/
        /*MSG_DEBUG(new_gen);*/
        *generations.back() = lump(new_gen, max_states);
        /*if (tree[n].is_crossing()) {*/
            /*propagate_symmetries(n, *generations.back());*/
        /*} else if (tree[n].is_ancestor()) {*/
//             generations.back()->symmetries = symmetry_group_type(generations.back()->labels);
        /*}*/
        /**node_generations[n] = lump(new_gen);*/
        if (cached_gen) {
            *cached_gen = node_generations[n];
        }
        MSG_DEBUG("DONE COMPUTING GENERATION FOR NODE #" << n);
        MSG_DEBUG_INDENT_EXPR("[RESULT " << tree.make_node_label(n) << " gen#" << node_generations[n] << "] ");
        MSG_DEBUG((*generations.back()));
        MSG_DEBUG_DEDENT;
        /*MSG_DEBUG((*generations[node_generations[n]]));*/
        MSG_DEBUG("=========================================================================");
    }

    std::map<genotype_comb_type::key_list, double> GLC_norm_factors(const std::vector<genotype_comb_type>& expanded)
    {
        std::map<genotype_comb_type::key_list, double> ret;
        for (const auto& e: expanded) {
            const auto& elem = e.m_combination.front();  /* all LC are singletons in the expanded vector */
            for (const auto& k: elem.keys) {
                auto sub_k = elem.keys - k;
                ret[sub_k] += elem.coef;
            }
        }
        return ret;
    }

    void compute_data_for_bn(int n)
    {
        compute_LC(n);

        factor_messages.resize(n + 1);
        for (auto& dest_f: compute_factors(n, true)) {
            factor_messages[dest_f.first].emplace_back(std::move(dest_f.second));
        }
        for (auto& dest_f: compute_factors(n, false)) {
            factor_messages[dest_f.first].emplace_back(std::move(dest_f.second));
        }
        individuals_in_factors.resize(n + 1);
        std::vector<size_t>& iif = individuals_in_factors.back();
        std::vector<int> in;
        std::vector<int> out;
        auto expr = tree.extract_expression(n, in, out);
        iif.reserve(expr.m_nodes.size());
        for (size_t i = 0; i < expr.m_nodes.size(); ++i) {
            if (expr.m_nodes[i].is_genotype()) {
                iif.push_back(expr.original_node_number(i));
            }
        }
        MSG_DEBUG("INDIVIDUALS IN FACTOR: " << iif);
    }

    void
        compute_LC(int n)
        {
            if (!with_LC) {
                return;
            }
            int p1 = tree.get_p1(n);
            std::vector<gencomb_type> lc;
            const geno_matrix& m = *generations[node_generations[n]];
            if (p1 == NONE) {
                lc.emplace_back(1.);
            } else {
                lc = eval_vector(tree.size() - 1, tree.get_recompute_vec(n), &pedigree_type::get_LC, reentrant_LC);
            }
            VectorLC tmp(lc.size()), lumped;
            for (size_t i = 0; i < lc.size(); ++i) {
                tmp(i) = lc[i];
            }
            lumped = m.collect.cast<gencomb_type>() * tmp;
            LC.resize(n + 1);
            LC.back().assign(lumped.data(), lumped.data() + lumped.size());
            MSG_DEBUG("Computed new LC:");
            MSG_DEBUG("" << LC.back());
        }

    std::map<size_t, std::map<bn_label_type, genotype_comb_type>>
        compute_factors(int n, bool up)
        {
            std::vector<genotype_comb_type> glc;
            std::vector<bn_label_type> bn_labels;

            if (tree.get_p1(n) == NONE) {
                return {};
            }

            if (up) {
                std::vector<bool> recompute(n + 1, false);
                recompute[n] = true;
                recompute[tree.get_p1(n)] = true;
                recompute[tree.get_p2(n)] = true;
                glc = eval_vector(tree.size() - 1, recompute, &pedigree_type::get_GLC, reentrant_GLC);
                bn_labels = eval_vector(tree.size() - 1, recompute, &pedigree_type::get_bn_labels, reentrant_bn_label);
            } else {
                glc = eval_vector(tree.size() - 1, tree.get_recompute_vec(n), &pedigree_type::get_GLC, reentrant_GLC);
                bn_labels = eval_vector(tree.size() - 1, tree.get_recompute_vec(n), &pedigree_type::get_bn_labels, reentrant_bn_label);
            }

            MSG_DEBUG("glc.size = " << glc.size());
            MSG_DEBUG("bn_labels.size = " << bn_labels.size());

            std::map<bn_label_type, genotype_comb_type> glc_map;
            for (size_t i = 0; i < bn_labels.size(); ++i) {
                glc_map[bn_labels[i]] += glc[i];
            }

            MSG_DEBUG("Corresponding GLC:");
            std::vector<genotype_comb_type> expanded_glc;
            for (const auto& kv: glc_map) {
                MSG_DEBUG("" << kv.first << " = " << kv.second);
                for (const auto& elem: kv.second.m_combination) {
                    expanded_glc.emplace_back();
                    expanded_glc.back().m_combination.emplace_back(elem);
                    auto& keys = expanded_glc.back().m_combination.back().keys.keys;
                    keys.emplace_back(tree.size() - 1, kv.first);
                    std::sort(keys.begin(), keys.end());
                }
            }

            MSG_DEBUG("Sparse GLC coefs:");
            for (const auto& g: expanded_glc) {
                MSG_DEBUG("" << g);
            }

            auto glc_norm_factors = GLC_norm_factors(expanded_glc);
            for (const auto& kn: glc_norm_factors) {
                MSG_DEBUG('|' << kn.first << "| = " << kn.second);
            }

            std::map<size_t, std::map<bn_label_type, genotype_comb_type>> messages;
            size_t n_nodes = expanded_glc.front().m_combination.front().keys.keys.size();
            size_t first, last;
            if (up) {
                first = 0;
                last = n_nodes - 2;
            } else {
                first = 0;
                last = n_nodes - 1;
            }
            for (const auto& lc: expanded_glc) {
                const auto& elems = lc.m_combination.front();
                const auto& keys = elems.keys.keys;
                for (size_t ni = first; ni <= last; ++ni) {
                    genotype_comb_type tmp;
                    tmp.m_combination.emplace_back(elems.coef);
                    auto& msg_keys = tmp.m_combination.back().keys.keys;
                    size_t i;
                    for (i = 0; i < ni; ++i) {
                        msg_keys.emplace_back(keys[i]);
                    }
                    for (++i; i < n_nodes; ++i) {
                        msg_keys.emplace_back(keys[i]);
                    }
                    tmp.m_combination.front().coef /= glc_norm_factors[tmp.m_combination.back().keys];
                    messages[keys[ni].parent][keys[ni].state] += tmp;
                }
            }

            for (const auto& kv1: messages) {
                MSG_DEBUG("MESSAGE TOWARDS #" << kv1.first);
                for (const auto& kv2: kv1.second) {
                    MSG_DEBUG("" << kv2.first << " = " << kv2.second);
                }
            }

            return messages;
        }


    static MatrixXd kron(const MatrixXd& m1, const MatrixXd& m2)
    {
        return kroneckerProduct(m1, m2);
    }

    static MatrixXd kron_d(const MatrixXd& m1, const MatrixXd& m2)
    {
        return kroneckerProduct(m1, MatrixXd::Identity(m2.rows(), m2.cols()))
             + kroneckerProduct(MatrixXd::Identity(m1.rows(), m1.cols()), m2);
    }

    static MatrixXd kron_d_diag(const MatrixXd& m1, const MatrixXd& m2)
    {
        return kroneckerProduct(m1, MatrixXd::Ones(m2.rows(), m2.cols()))
             + kroneckerProduct(MatrixXd::Ones(m1.rows(), m1.cols()), m2);
    }

    char ancestor_letter(size_t a) const
    {
        auto i = ancestor_letters.find(tree.ind2node(a));
        return i == ancestor_letters.end() ? 0 : i->second;
    }

    geno_matrix_index_type get_gen_index(size_t ind) const
    {
        return node_generations[tree.ind2node(ind)];
    }

    const std::shared_ptr<geno_matrix> get_gen(size_t ind) const
    {
        return generations[get_gen_index(ind)];
    }

    int
        ind2id(int i) const
        {
            return items[i - 1].id;
        }

    const geno_matrix&
        get_geno_matrix_by_individual(size_t ind) const { return *generations[node_generations[tree.ind2node(ind)]]; }
    const std::set<geno_matrix_index_type>&
        get_geno_matrices_by_name(const std::string& name) const { return geno_matrix_by_generation_name.find(name)->second; }
    const std::string&
        get_generation_name_by_individual(size_t ind) const { return *generation_name_by_individual.find(ind)->second; }

    const std::shared_ptr<geno_matrix> get_node_gen(size_t node) const
    {
        return generations[node_generations[node]];
    }

    typedef MatrixXd (* KronFunc) (const MatrixXd&, const MatrixXd&);

    template <typename FIELD_TYPE>
    FIELD_TYPE eval(size_t node, FIELD_TYPE geno_matrix::* field, KronFunc func, const std::vector<bool>& recompute, std::vector<bool>& visited) const
    {
        scoped_indent _;
        /*MSG_DEBUG("eval node " << node);*/
        if (visited[node]) {
            /*MSG_DEBUG("already visited => 1");*/
            return make_one<FIELD_TYPE>::_(func == kron_d || func == kron_d_diag);
        }
        visited[node] = true;
        if (recompute[node] || tree[node].is_gamete()) {
            FIELD_TYPE ret, m1, m2;
            if (tree[node].is_genotype()) {
                m1 = eval(tree.get_p1(node), field, func, recompute, visited);
                m2 = eval(tree.get_p2(node), field, func, recompute, visited);
                ret = func(m1, m2);
            } else {
                m1 = eval(tree.get_p1(node), field, func, recompute, visited);
                ret = func(m1, gamete.*field);
            }
            return ret;
        } else {
            /*MSG_DEBUG("using actual matrix");*/
            /*MSG_DEBUG((*generations[node_generations[node]]).*field);*/
            return (*generations[node_generations[node]]).*field;
        }
        return {};
    }

    template <typename VALUE_TYPE>
    struct vector_iterator {
        typedef std::vector<VALUE_TYPE> vector_type;
        vector_type data;
        typename vector_type::const_iterator begin, end, cur;
        vector_iterator() : data(), begin(data.begin()), end(data.end()), cur(data.begin()) {}
        vector_iterator(const vector_type& l) : data(l), begin(data.begin()), end(data.end()), cur(data.begin()) {}
        void reset() { cur = begin = data.begin(); end = data.end(); }
        void start() { cur = begin; }
        bool next() { if (at_end()) return true; return ++cur == end; }
        bool at_end() const { return cur == end; }
        size_t size() const { return end - begin; }
        const VALUE_TYPE& operator * () const { return *cur; }
        const VALUE_TYPE* operator -> () const { return &*cur; }
    };

    typedef vector_iterator<label_type> label_iterator;
//     typedef vector_iterator<symmetry_table_type> symmetry_iterator;

    template <typename VALUE_TYPE>
        VALUE_TYPE eval_one(size_t node, const std::vector<bool>& recompute, const std::vector<vector_iterator<VALUE_TYPE>>& iterators,
                            const std::vector<size_t>& node_to_iterator, std::vector<bool>& visited,
                            VALUE_TYPE (&eval_reent)(size_t, const VALUE_TYPE&),
                            std::vector<VALUE_TYPE>& reent) const
        {
            /*scoped_indent _(SPELL_STRING("[eval #" << node << "] "));*/
            if (visited[node]) {
                auto ret = eval_reent(node, reent[node]);
                /*ret = eval_reent(node, iterators[node_to_iterator[node]]);*/
                /*MSG_DEBUG("reentrant; " << ret);*/
                return ret;
            } else {
                visited[node] = true;
                if (recompute[node] || tree[node].is_gamete()) {
                    if (tree[node].is_genotype()) {
                        /*MSG_DEBUG("eval p1");*/
                        auto p1 = eval_one(tree.get_p1(node), recompute, iterators, node_to_iterator, visited, eval_reent, reent);
                        /*MSG_DEBUG("p1 = " << p1);*/
                        /*MSG_DEBUG("eval p2");*/
                        auto p2 = eval_one(tree.get_p2(node), recompute, iterators, node_to_iterator, visited, eval_reent, reent);
                        /*MSG_DEBUG("p2 = " << p2);*/
                        reent[node] = p1 * p2;
                    } else {
                        reent[node] = eval_one(tree.get_p1(node), recompute, iterators, node_to_iterator, visited, eval_reent, reent)
                                    * *iterators[node_to_iterator[node]];
                    }
                } else {
                    reent[node] = *iterators[node_to_iterator[node]];
                }
            }
            /*MSG_DEBUG("ret = " << reent[node]);*/
            /*MSG_QUEUE_FLUSH();*/
            return reent[node];
        }

    struct skip_eval_exception {};

    template <typename VALUE_TYPE>
        VALUE_TYPE eval_one(size_t node, const std::vector<bool>& recompute, const std::vector<vector_iterator<VALUE_TYPE>>& iterators,
                            const std::vector<size_t>& node_to_iterator, std::vector<bool>& visited,
                            VALUE_TYPE (&eval_reent)(size_t, const VALUE_TYPE&),
                            std::vector<VALUE_TYPE>& reent,
                            bool (&skip_predicate)(const VALUE_TYPE&)) const
        {
            /*scoped_indent _(SPELL_STRING("[eval #" << node << "] "));*/
            if (visited[node]) {
                auto ret = eval_reent(node, reent[node]);
                /*ret = eval_reent(node, iterators[node_to_iterator[node]]);*/
                /*MSG_DEBUG("reentrant; " << ret);*/
                return ret;
            } else {
                visited[node] = true;
                if (recompute[node] || tree[node].is_gamete()) {
                    if (tree[node].is_genotype()) {
                        /*MSG_DEBUG("eval p1");*/
                        auto p1 = eval_one(tree.get_p1(node), recompute, iterators, node_to_iterator, visited, eval_reent, reent);
                        /*MSG_DEBUG("p1 = " << p1);*/
                        /*MSG_DEBUG("eval p2");*/
                        auto p2 = eval_one(tree.get_p2(node), recompute, iterators, node_to_iterator, visited, eval_reent, reent);
                        /*MSG_DEBUG("p2 = " << p2);*/
                        reent[node] = p1 * p2;
                    } else {
                        reent[node] = eval_one(tree.get_p1(node), recompute, iterators, node_to_iterator, visited, eval_reent, reent)
                                    * *iterators[node_to_iterator[node]];
                    }
                } else {
                    reent[node] = *iterators[node_to_iterator[node]];
                }
            }
            /*MSG_DEBUG("ret = " << reent[node]);*/
            /*if (skip_predicate(reent[node])) {*/
                /*throw skip_eval_exception();*/
            /*}*/
            return reent[node];
            (void)skip_predicate;
        }

    static label_type reentrant_label(size_t, const label_type& l) { return l; }
    static bn_label_type reentrant_bn_label(size_t, const bn_label_type& l) { return l; }
    static gencomb_type reentrant_LC(size_t, const gencomb_type&) { return {1.}; }
    static genotype_comb_type reentrant_GLC(size_t, const genotype_comb_type&) { return {1.}; }
//     static symmetry_table_type reentrant_sym(size_t, const symmetry_table_type& S) { return {permutation_type::identity(1), S.letters}; }

    template <typename VALUE_TYPE>
        void init_iterators_rec(size_t node, const std::vector<bool>& recompute, std::vector<bool>& visited,
                                std::vector<vector_iterator<VALUE_TYPE>>& iterators, std::vector<size_t>& node_to_iterator,
                                std::vector<VALUE_TYPE> (pedigree_type:: *accessor)(bool, size_t, const geno_matrix&) const) const
        {
            if (visited[node]) {
                return;
            }
            visited[node] = true;
            if (recompute[node] || tree[node].is_gamete()) {
                init_iterators_rec(tree.get_p1(node), recompute, visited, iterators, node_to_iterator, accessor);
                if (tree[node].is_genotype()) {
                    init_iterators_rec(tree.get_p2(node), recompute, visited, iterators, node_to_iterator, accessor);
                } else {
                    node_to_iterator[node] = iterators.size();
                    iterators.emplace_back((this->*accessor)(true, node, gamete));
                    /*MSG_DEBUG("GAMETE");*/
                    /*MSG_DEBUG("" << (this->*accessor)(true, node, gamete));*/
                }
#if 0
                switch (nodes[node].type) {
                    case NTGenotype:
                        break;
                    case NTGamete:
                        break;
                    case NTDoubling:
                        node_to_iterator[node] = iterators.size();
                        iterators.emplace_back((this->*accessor)(true, node, doubling_gamete));
                        MSG_DEBUG("DOUBLING GAMETE");
                        MSG_DEBUG("" << (this->*accessor)(true, node, doubling_gamete));
                };
#endif
            } else {
                node_to_iterator[node] = iterators.size();
                iterators.emplace_back((this->*accessor)(false, node, *generations[node_generations[node]]));
                /*MSG_DEBUG("GENERATION");*/
                /*MSG_DEBUG("" << (this->*accessor)(false, node, *generations[node_generations[node]]));*/
            }
        }

    std::vector<label_type> get_labels(bool, size_t, const geno_matrix& m) const { return m.labels; }
    std::vector<bn_label_type> get_bn_labels(bool, size_t, const geno_matrix& m) const
    {
        if (m.variant == Geno) {
            std::vector<bn_label_type> ret;
            ret.reserve(m.labels.size() * n_alleles * n_alleles);
            for (const auto& l: m.labels) {
                for (int a1 = 0; a1 < (int) n_alleles; ++a1) {
                    for (int a2 = 0; a2 < (int) n_alleles; ++a2) {
                        ret.emplace_back(l.first(), l.second(), a1, a2);
                    }
                }
            }
            return ret;
        } else if (m.variant == Gamete) {
            return {{GAMETE_L, GAMETE_EMPTY, 0, 0}, {GAMETE_R, GAMETE_EMPTY, 0, 0}};
        } else if (m.variant == DoublingGamete) {
            return {{GAMETE_L, GAMETE_L, 0, 0}, {GAMETE_R, GAMETE_R, 0, 0}};
        } else {
            std::vector<bn_label_type> ret;
            ret.reserve(m.labels.size() * n_alleles);
            for (const auto& l: m.labels) {
                for (int a1 = 0; a1 < (int) n_alleles; ++a1) {
                    ret.emplace_back(l.first(), l.second(), a1, 0);
                }
            }
            return ret;
        }
    }

    std::vector<gencomb_type> get_LC(bool recompute, size_t node, const geno_matrix& m) const
    {
        /*return {LC[node].data(), LC[node].data() + LC[node].size()};*/
        if (recompute && !tree[node].is_genotype()) {
            return {.5, .5};
        } else {
            std::vector<gencomb_type> ret(m.cols());
            for (size_t i = 0; i < ret.size(); ++i) {
                ret[i].m_combination.emplace_back(gencomb_type::key_type{node, i}, 1.);
            }
            return ret;
        }
        return {};
    }

    std::vector<genotype_comb_type> get_GLC(bool recompute, size_t node, const geno_matrix& m) const
    {
        /*return {LC[node].data(), LC[node].data() + LC[node].size()};*/
        if (recompute && !tree[node].is_genotype()) {
            return {.5, .5};
        } else {
            if (tree[node].is_genotype()) {
                std::vector<genotype_comb_type> ret(m.cols() * n_alleles * n_alleles);
                size_t n = 0;
                double norm = 1.; // / (n_alleles * n_alleles);
                for (size_t i = 0; i < m.labels.size(); ++i) {
                    for (char a1 = 0; a1 < (char) n_alleles; ++a1) {
                        for (char a2 = 0; a2 < (char) n_alleles; ++a2) {
                            ret[n++].m_combination.emplace_back(genotype_comb_type::key_type{(int) node, {m.labels[i].first(), m.labels[i].second(), a1, a2}}, norm);
                        }
                    }
                }
                return ret;
            } else {
                std::vector<genotype_comb_type> ret(m.cols() * n_alleles);
                size_t n = 0;
                double norm = 1.; // / n_alleles;
                for (size_t i = 0; i < m.labels.size(); ++i) {
                    for (char a1 = 0; a1 < (char) n_alleles; ++a1) {
                        ret[n++].m_combination.emplace_back(genotype_comb_type::key_type{(int) node, {m.labels[i].first(), m.labels[i].second(), a1, 0}}, norm);
                    }
                }
                return ret;
            }
        }
        return {};
    }

    template <typename VALUE_TYPE>
        bool next(std::vector<vector_iterator<VALUE_TYPE>>& iterators) const
        {
            auto i = iterators.rbegin(), j = iterators.rend();
            for (; i != j; ++i) {
                if (i->next()) {
                    i->start();
                } else {
                    break;
                }
            }
            return i != j;
        }

    template <typename VALUE_TYPE>
        void reset(std::vector<vector_iterator<VALUE_TYPE>>& iterators) const
        {
            for (auto& vi: iterators) {
                vi.reset();
            }
        }

    template <typename VALUE_TYPE>
        std::vector<VALUE_TYPE>
        eval_vector(size_t node, const std::vector<bool>& recompute,
                    std::vector<VALUE_TYPE> (pedigree_type:: *accessor)(bool, size_t, const geno_matrix&) const,
                    VALUE_TYPE (&eval_reent)(size_t, const VALUE_TYPE&)) const
        {
            std::vector<VALUE_TYPE> reent(tree.size());
            /*MSG_DEBUG_INDENT_EXPR("[eval_vector #" << node << "] ");*/
            std::vector<bool> visited(tree.size(), false);
            std::vector<vector_iterator<VALUE_TYPE>> iterators;
            std::vector<size_t> node_to_iterator;
            std::vector<VALUE_TYPE> ret;

            iterators.reserve(tree.size());
            node_to_iterator.reserve(tree.size());

            /*MSG_DEBUG("init iterators");*/
            init_iterators_rec(node, recompute, visited, iterators, node_to_iterator, accessor);

            size_t total_size = 1;

            for (const auto& i: iterators) {
                total_size *= i.size();
            }
            /*MSG_DEBUG("total size = " << total_size);*/

            if (!total_size) {
                /*MSG_DEBUG("EMPTY!");*/
                /*MSG_DEBUG_DEDENT;*/
                return {};
            }

            ret.reserve(total_size);

            /*size_t i = 0;*/
            do {
                /*MSG_DEBUG((++i) << "/" << total_size);*/
                visited.assign(tree.size(), false);
                ret.push_back(eval_one(node, recompute, iterators, node_to_iterator, visited, eval_reent, reent));
            } while (next(iterators));

            /*MSG_DEBUG("DONE");*/
            /*MSG_DEBUG_DEDENT;*/