// $Id$
/*
 * This is an extensive modification by Greg Landrum of
 * pieces from several files in the vflib-2.0 distribution
 *
 * The initial version of the modifications was completed
 *   in April 2009.
 *
 * the original author of the vflib files is:
 *    Author: P. Foggia
 *  http://amalfi.dis.unina.it/graph/db/vflib-2.0/doc/vflib.html
 *
 */
#include <boost/graph/adjacency_list.hpp>
#include <vector>
#include <algorithm>
#include <cstring>

#ifndef __BGL_VF2_SUB_STATE_H__
#define __BGL_VF2_SUB_STATE_H__

namespace boost{
  namespace detail {
    typedef unsigned short node_id;
    const node_id NULL_NODE=0xFFFF;
    struct NodeInfo {
      node_id id;
      node_id in;
      node_id out;
    };

    /**
     * The ordering by in/out degree
     */
    static bool nodeInfoComp1(const NodeInfo &a, const NodeInfo &b) {
      if(a.out < b.out) return true;
      if(a.out > b.out) return false;
      if(a.in < b.in) return true;
      if(a.in > b.in) return false;
      return false;
    }

    /**
     * The ordering by frequency/valence.
     * The frequency is in the out field, the valence in `in'.
     */
    static int nodeInfoComp2(const NodeInfo &a, const NodeInfo &b) {
      if (!a.in && b.in ) return 1;
      if (a.in && !b.in) return -1;
      if (a.out < b.out) return -1;
      if (a.out > b.out) return 1;
      if( a.in < b.in ) return -1;
      if (a.in > b.in) return 1;
      return 0;
    }

    template <class Graph,class VertexDescr,class EdgeDescr>
    VertexDescr getOtherIdx(const Graph &g,const EdgeDescr &edge,const VertexDescr &vertex) {
      VertexDescr tmp=boost::source(edge,g);
      if(tmp==vertex){
        tmp=boost::target(edge,g);
      }
      return tmp;
    }

    /*----------------------------------------------------
     * Sorts the nodes of a graphs, returning a
     * heap-allocated vector (using new) with the node ids
     * in the proper orders.
     * The sorting criterion takes into account:
     *    1 - The number of nodes with the same in/out
     *        degree.
     *    2 - The valence of the nodes.
     * The nodes at the beginning of the vector are
     * the most singular, from which the matching should
     * start.
     *--------------------------------------------------*/
    template <class Graph>
    node_id* SortNodesByFrequency(const Graph *g) {
      std::vector<NodeInfo> vect;
      vect.reserve(boost::num_vertices(*g));
      typename Graph::vertex_iterator bNode,eNode;
      boost::tie(bNode,eNode) = boost::vertices(*g);
      while(bNode!=eNode){
        NodeInfo t;
        t.id=vect.size();
        t.in=boost::out_degree(*bNode,*g);// <- assuming undirected graph
        t.out=boost::out_degree(*bNode,*g);
        vect.push_back(t);
        ++bNode;
      }
      std::sort(vect.begin(),vect.end(),nodeInfoComp1);

      unsigned int run=1;
      for(unsigned int i=0; i<vect.size(); i+=run){
        for(run=1; i+run<vect.size() &&
              vect[i+run].in==vect[i].in &&
              vect[i+run].out==vect[i].out;
            ++run)
          ;
        for(unsigned int j=0; j<run; ++j) {
          vect[i+j].in += vect[i+j].out;
          vect[i+j].out=run;
        }
      }
      std::sort(vect.begin(),vect.end(),nodeInfoComp2);

      node_id *nodes=new node_id[vect.size()];
      for(unsigned int i=0; i<vect.size(); ++i){
        nodes[i]=vect[i].id;
      }

      return nodes;
    }

    /*----------------------------------------------------------
     * class VF2SubState
     * A representation of the SSS current state
     ---------------------------------------------------------*/
    template <class Graph,class VertexCompatible,class EdgeCompatible,class MatchChecking >
    class VF2SubState
    {
    private:
      Graph *g1, *g2;
      VertexCompatible &vc;
      EdgeCompatible &ec;
      MatchChecking &mc;
      unsigned int n1, n2;

      unsigned int core_len, orig_core_len;
      unsigned int added_node1;
      unsigned int t1both_len, t2both_len;
      unsigned int t1in_len, t1out_len;
      unsigned int t2in_len, t2out_len; // Core nodes are also counted by these...
      node_id *core_1;
      node_id *core_2;
      node_id *in_1;
      node_id *in_2;
      node_id *out_1;
      node_id *out_2;

      node_id *order;

      long *share_count;
      int *vs_compared;

    public:
      VF2SubState(Graph *ag1, Graph *ag2,
                  VertexCompatible &avc,
                  EdgeCompatible &aec,
                  MatchChecking &amc,
                  bool sortNodes=false) : g1(ag1), g2(ag2), vc(avc), ec(aec), mc(amc),
                                          n1(num_vertices(*ag1)),n2(num_vertices(*ag2)) {
        if (sortNodes){
          order = SortNodesByFrequency(ag1);
        } else {
          order = NULL;
        }

        core_len=orig_core_len=0;
        t1both_len=t1in_len=t1out_len=0;
        t2both_len=t2in_len=t2out_len=0;

        added_node1=NULL_NODE;

        core_1=new node_id[n1];
        core_2=new node_id[n2];
        in_1=new node_id[n1];
        in_2=new node_id[n2];
        out_1=new node_id[n1];
        out_2=new node_id[n2];
        share_count = new long;

        for(unsigned int i=0; i<n1; i++){
          core_1[i]=NULL_NODE;
          in_1[i]=0;
          out_1[i]=0;
        }
        for(unsigned int i=0; i<n2; i++){
          core_2[i]=NULL_NODE;
          in_2[i]=0;
          out_2[i]=0;
        }
        vs_compared=0;
        //vs_compared = new int[n1*n2];
        //memset((void *)vs_compared,0,n1*n2*sizeof(int));

        //es_compared = new std::map<unsigned int,bool>();
        *share_count = 1;
      };

      VF2SubState(const VF2SubState &state) :
        g1(state.g1), g2(state.g2), vc(state.vc), ec(state.ec), mc(state.mc),
        n1(state.n1),n2(state.n2), order(state.order),vs_compared(state.vs_compared)
        //es_compared(state.es_compared)
      {

        core_len=orig_core_len=state.core_len;
        t1in_len=state.t1in_len;
        t1out_len=state.t1out_len;
        t1both_len=state.t1both_len;
        t2in_len=state.t2in_len;
        t2out_len=state.t2out_len;
        t2both_len=state.t2both_len;

        added_node1=NULL_NODE;

        core_1=state.core_1;
        core_2=state.core_2;
        in_1=state.in_1;
        in_2=state.in_2;
        out_1=state.out_1;
        out_2=state.out_2;
        share_count=state.share_count;

        ++(*share_count);
      };

      ~VF2SubState(){
        if (-- *share_count == 0) {
          delete [] core_1;
          delete [] core_2;
          delete [] in_1;
          delete [] out_1;
          delete [] in_2;
          delete [] out_2;
          delete share_count;
          delete [] order;
          //delete [] vs_compared;
          //delete es_compared;
        }
      };

      bool IsGoal() { return core_len==n1 ; };
      bool MatchChecks(const node_id c1[],const node_id c2[]){
        return mc(c1,c2);
      };
      bool IsDead() { return n1>n2  ||
          t1both_len>t2both_len ||
          t1out_len>t2out_len ||
          t1in_len>t2in_len;
      };
      unsigned int CoreLen() { return core_len; }
      Graph *GetGraph1() { return g1; }
      Graph *GetGraph2() { return g2; }

      bool NextPair(node_id *pn1, node_id *pn2,
                    node_id prev_n1=NULL_NODE, node_id prev_n2=NULL_NODE){
        if (prev_n1==NULL_NODE)
          prev_n1=0;
        if (prev_n2==NULL_NODE)
          prev_n2=0;
        else
          prev_n2++;

#if 0
    std::cerr<<" **** np: "<< prev_n1<<","<<prev_n2<<std::endl;
    std::cerr<<"in_1 ";
    for(unsigned int i=0;i<n1;++i){
      std::cerr<<"("<<in_1[i]<<","<<out_1[i]<<"), ";
    }
    std::cerr<<std::endl;
    std::cerr<<"in_2 ";
    for(unsigned int i=0;i<n2;++i){
      std::cerr<<"("<<in_2[i]<<","<<out_2[i]<<"), ";
    }
    std::cerr<<std::endl;
#endif
        if (t1both_len>core_len && t2both_len>core_len) {
          while (prev_n1<n1 &&
                 (core_1[prev_n1]!=NULL_NODE || out_1[prev_n1]==0
                  || in_1[prev_n1]==0) ) {
            prev_n1++;
            prev_n2=0;
          }
        }
        else if (t1out_len>core_len && t2out_len>core_len) {
          while (prev_n1<n1 &&
                 (core_1[prev_n1]!=NULL_NODE || out_1[prev_n1]==0) ){
            prev_n1++;
            prev_n2=0;
          }
        }
        else if (t1in_len>core_len && t2in_len>core_len) {
          while (prev_n1<n1 &&
                 (core_1[prev_n1]!=NULL_NODE || in_1[prev_n1]==0) ) {
            prev_n1++;
            prev_n2=0;
          }
        }
        else if (prev_n1==0 && order!=NULL) {
          unsigned int i=0;
          while (i<n1 && core_1[prev_n1=order[i]] != NULL_NODE)
            i++;
          if (i==n1)
            prev_n1=n1;
        }
        else {
          while (prev_n1<n1 && core_1[prev_n1]!=NULL_NODE ){
            prev_n1++;
            prev_n2=0;
          }
        }

        if (t1both_len>core_len && t2both_len>core_len) {
          while (prev_n2<n2 &&
                 (core_2[prev_n2]!=NULL_NODE || out_2[prev_n2]==0
                  || in_2[prev_n2]==0) ) {
            prev_n2++;
          }
        }
        else if (t1out_len>core_len && t2out_len>core_len) {
          while (prev_n2<n2 &&
                 (core_2[prev_n2]!=NULL_NODE || out_2[prev_n2]==0) ) {
            prev_n2++;
          }
        }
        else if (t1in_len>core_len && t2in_len>core_len) {
          while (prev_n2<n2 &&
                 (core_2[prev_n2]!=NULL_NODE || in_2[prev_n2]==0) ) {
            prev_n2++;
          }
        }
        else {
          while (prev_n2<n2 && core_2[prev_n2]!=NULL_NODE ){
            prev_n2++;
          }
        }
        //std::cerr<<" "<< prev_n1<<"<"<<n1<<" "<<prev_n2<<"<"<<n2;
        if (prev_n1<n1 && prev_n2<n2) {
          *pn1=prev_n1;
          *pn2=prev_n2;
          //std::cerr<<"  Found"<<std::endl;
          return true;
        }
        //std::cerr<<"  nope"<< std::endl;
        return false;
      };
      bool IsFeasiblePair(node_id node1, node_id node2){
        assert(node1 < n1);
        assert(node2 < n2);
        assert(core_1[node1] == NULL_NODE);
        assert(core_2[node2] == NULL_NODE);

        //std::cerr<<"  ifp:"<<node1<<"-"<<node2<<" "<<vs_compared->size()<<std::endl;
        // int &isCompat=vs_compared[node1*n2+node2];
        // if(isCompat==0){
        //   isCompat=vc(node1,node2)?1:-1;
        // }
        // if( isCompat<0 ){
        //   //std::cerr<<"  short1"<<std::endl;
        //   return false;
        // }
        if(!vc(node1,node2)) return false;

        unsigned int other1, other2;
        unsigned int termout1 = 0, termout2 = 0, termin1 = 0, termin2 = 0;
        unsigned int new1 = 0, new2 = 0;

        // Check the out edges of node1
        typename Graph::out_edge_iterator bNbrs,eNbrs;
        boost::tie(bNbrs,eNbrs) = boost::out_edges(node1,*g1);
        while(bNbrs!=eNbrs){
          other1=getOtherIdx(*g1,*bNbrs,node1);
          if (core_1[other1] != NULL_NODE) {
            other2 = core_1[other1];
            typename Graph::edge_descriptor oEdge;
            bool found;
            boost::tie(oEdge,found) = boost::edge(node2,other2,*g2);
            if(!found || !ec(*bNbrs,oEdge) ){
              //std::cerr<<"  short2"<<std::endl;
              return false;
            }
          } else {
            if (in_1[other1]) ++termin1;
            if (out_1[other1]) ++termout1;
            if (!in_1[other1] && !out_1[other1]) ++new1;
          }
          ++bNbrs;
        }

        // Check the out edges of node2
        boost::tie(bNbrs,eNbrs) = boost::out_edges(node2,*g2);
        while(bNbrs!=eNbrs){
          other2=getOtherIdx(*g2,*bNbrs,node2);
          if (core_2[other2] != NULL_NODE) {
            // do nothing
          } else {
            if (in_2[other2]) ++termin2;
            if (out_2[other2]) ++termout2;
            if (!in_2[other2] && !out_2[other2]) ++new2;
          }
          ++bNbrs;
        }
        //std::cerr<<(termin1 <= termin2 && termout1 <= termout2 && (termin1+termout1+new1)<=(termin2+termout2+new2))<<std::endl;

        return termin1 <= termin2 && termout1 <= termout2 && (termin1+termout1+new1)<=(termin2+termout2+new2);
      };
      void AddPair(node_id node1, node_id node2){
        assert(node1 < n1);
        assert(node2 < n2);
        assert(core_len < n1);
        assert(core_len < n2);

        ++core_len;
        added_node1 = node1;

        if (!in_1[node1]) {
          in_1[node1] = core_len;
          ++t1in_len;
          if (out_1[node1]) ++t1both_len;
        }
        if (!out_1[node1]) {
          out_1[node1] = core_len;
          ++t1out_len;
          if (in_1[node1]) ++t1both_len;
        }

        if (!in_2[node2]) {
          in_2[node2] = core_len;
          ++t2in_len;
          if (out_2[node2]) ++t2both_len;
        }
        if (!out_2[node2]) {
          out_2[node2] = core_len;
          ++t2out_len;
          if (in_2[node2]) ++t2both_len;
        }

        core_1[node1] = node2;
        core_2[node2] = node1;

        typename Graph::out_edge_iterator bNbrs,eNbrs;
        // FIX: this is explicitly ignoring directionality
        boost::tie(bNbrs,eNbrs) = boost::out_edges(node1,*g1);
        while(bNbrs!=eNbrs){
          unsigned int other = getOtherIdx(*g1,*bNbrs,node1);
          if (!in_1[other]) {
            in_1[other] = core_len;
            ++t1in_len;
            if (out_1[other])  ++t1both_len;
          }
          if (!out_1[other]) {
            out_1[other] = core_len;
            ++t1out_len;
            if (in_1[other])  ++t1both_len;
          }
          ++bNbrs;
        }

        // FIX: this is explicitly ignoring directionality
        boost::tie(bNbrs,eNbrs) = boost::out_edges(node2,*g2);
        while(bNbrs!=eNbrs){
          unsigned int other = getOtherIdx(*g2,*bNbrs,node2);
          if (!in_2[other]) {
            in_2[other] = core_len;
            ++t2in_len;
            if (out_2[other]) ++t2both_len;
          }
          if (!out_2[other]) {
            out_2[other] = core_len;
            ++t2out_len;
            if (in_2[other]) ++t2both_len;
          }
          ++bNbrs;
        }
      };
      void GetCoreSet(node_id c1[], node_id c2[]){
        unsigned int i, j;
        for (i = 0, j = 0; i < n1; ++i){
          if (core_1[i] != NULL_NODE) {
            c1[j] = i;
            c2[j] = core_1[i];
            ++j;
          }
        }

      };
      VF2SubState *Clone(){
        return new VF2SubState(*this);
      };
      void BackTrack(){
        assert(core_len - orig_core_len <= 1);
        assert(added_node1 != NULL_NODE);

        if (orig_core_len < core_len) {

          if (in_1[added_node1] == core_len) in_1[added_node1] = 0;
          if (out_1[added_node1] == core_len)  out_1[added_node1] = 0;

          typename Graph::out_edge_iterator bNbrs,eNbrs;
          boost::tie(bNbrs,eNbrs) = boost::out_edges(added_node1,*g1);
          while(bNbrs!=eNbrs){
            unsigned int other = getOtherIdx(*g1,*bNbrs,added_node1);
            if (out_1[other] == core_len) out_1[other] = 0;
            if (in_1[other] == core_len) in_1[other] = 0;
            ++bNbrs;
          }

          unsigned int node2 = core_1[added_node1];
          if (in_2[node2] == core_len) in_2[node2] = 0;
          if (out_2[node2] == core_len) out_2[node2] = 0;

          boost::tie(bNbrs,eNbrs) = boost::out_edges(node2,*g2);
          while(bNbrs!=eNbrs){
            unsigned int other = getOtherIdx(*g2,*bNbrs,node2);
            if (out_2[other] == core_len) out_2[other] = 0;
            if (in_2[other] == core_len) in_2[other] = 0;
            ++bNbrs;
          }

          core_1[added_node1] = NULL_NODE;
          core_2[node2] = NULL_NODE;

          core_len = orig_core_len;
          added_node1 = NULL_NODE;
        }
      };
    };

    /*-------------------------------------------------------------
     * static bool match(pn, c1, c2, s)
     * Finds a matching between two graphs, if it exists, starting
     * from state s.
     * Returns true a match has been found.
     * *pn is assigned the numbero of matched nodes, and
     * c1 and c2 will contain the ids of the corresponding nodes
     * in the two graphs.
     ------------------------------------------------------------*/
    template <class SubState>
    bool match(int *pn, node_id c1[], node_id c2[], SubState &s)
    {
      if (s.IsGoal() ) {
        s.GetCoreSet(c1, c2);
        if(s.MatchChecks(c1,c2)) {
          *pn=s.CoreLen();
          return true;
        }
      }

      if (s.IsDead())
        return false;
      //std::cerr<<"  > match: "<<*pn<<" "<<&s<<std::endl;
      node_id n1=NULL_NODE, n2=NULL_NODE;
      bool found=false;
      while (!found && s.NextPair(&n1, &n2, n1, n2)) {
        //std::cerr<<"           "<<n1<<","<<n2<<std::endl;
        if (s.IsFeasiblePair(n1, n2)){
          SubState *s1=s.Clone();
          s1->AddPair(n1, n2);
          found=match(pn, c1, c2, *s1);
          s1->BackTrack();
          delete s1;
        }
      }
      //std::cerr<<"  < returning: "<<found<<" "<<*pn<<" "<<&s<<std::endl;
      return found;
    }

    /*-------------------------------------------------------------
     * static bool match(c1, c2, vis, usr_data, pcount)
     * Visits all the matchings between two graphs,  starting
     * from state s.
     * Returns true if the caller must stop the visit.
     * Stops when there are no more matches
     *
     ------------------------------------------------------------*/
    template <class SubState,class DoubleBackInsertionSequence>
    bool match(node_id c1[], node_id c2[], SubState &s, DoubleBackInsertionSequence &res,
               unsigned int max_results) {
      if (s.IsGoal()){
        s.GetCoreSet(c1, c2);
        if(s.MatchChecks(c1,c2)) {
          typename DoubleBackInsertionSequence::value_type newSeq;
          for(unsigned int i=0;i<s.CoreLen();++i){
            newSeq.push_back(std::pair<int,int>(c1[i],c2[i]));
          }
          res.push_back(newSeq);
          if(res.size()>=max_results) return true;
        }
        return false;
      }

      if (s.IsDead())
        return false;

      node_id n1=NULL_NODE, n2=NULL_NODE;
      while (s.NextPair(&n1, &n2, n1, n2)) {
        if (s.IsFeasiblePair(n1, n2)){
          SubState *s1=s.Clone();
          s1->AddPair(n1, n2);
          if (match(c1, c2, *s1,res,max_results)){
            s1->BackTrack();
            delete s1;
            return true;
          }
          else {
            s1->BackTrack();
            delete s1;
          }
        }
      }
      return false;
    }
  }; //end of namespace detail

  template <  class Graph
              , class VertexLabeling    // binary predicate
              , class EdgeLabeling      // binary predicate
              , class MatchChecking      // binary predicate
              , class BackInsertionSequence   // contains std::pair<vertex_descriptor,vertex_descriptor>
              >
  bool vf2(const Graph &g1,const Graph &g2,
           VertexLabeling& vertex_labeling,
           EdgeLabeling& edge_labeling,
           MatchChecking& match_checking,
           BackInsertionSequence& F){
    detail::VF2SubState<const Graph,VertexLabeling,EdgeLabeling,MatchChecking> s0(&g1,&g2,vertex_labeling,
                                                                                edge_labeling,match_checking,false);
    detail::node_id *ni1 = new detail::node_id[num_vertices(g1)];
    detail::node_id *ni2 = new detail::node_id[num_vertices(g2)];
    int n=0;

    F.clear();
    F.resize(0);
    if(match(&n,ni1,ni2,s0)){
      for(unsigned int i=0;i<num_vertices(g1);i++){
        F.push_back(std::pair<int,int>(ni1[i],ni2[i]));
      }
    }
    delete [] ni1;
    delete [] ni2;

    return !F.empty();
  };
  template <  class Graph
              , class VertexLabeling    // binary predicate
              , class EdgeLabeling      // binary predicate
              , class MatchChecking      // binary predicate
              , class DoubleBackInsertionSequence   // contains a back insertion sequence
              >
  bool vf2_all(const Graph& g1, const Graph& g2,
               VertexLabeling& vertex_labeling,
               EdgeLabeling& edge_labeling,
               MatchChecking& match_checking,
               DoubleBackInsertionSequence& F,
               unsigned int max_results=1000) {
    detail::VF2SubState<const Graph,VertexLabeling,EdgeLabeling,MatchChecking> s0(&g1,&g2,vertex_labeling,
                                                                                edge_labeling,match_checking,false);
    detail::node_id *ni1 = new detail::node_id[num_vertices(g1)];
    detail::node_id *ni2 = new detail::node_id[num_vertices(g2)];

    F.clear();
    F.resize(0);

    match(ni1,ni2,s0,F,max_results);

    delete [] ni1;
    delete [] ni2;

    return !F.empty();
  };
} // end of namespace boost
#endif


// Ultra fast anagram generator

#include <algorithm>
#include <bitset>
#include <cassert>
#include <codecvt>
#include <fstream>
#include <ios>
#include <iostream>
#include <limits>
#include <locale>
#include <numeric>
#include <optional>
#include <sstream>
#include <string>
#include <tuple>
#include <unordered_map>
#include <utility>
#include <vector>

#include <boost/functional/hash.hpp>
#include <boost/program_options.hpp>

#include <unicode/stringoptions.h>
#include <unicode/uchar.h>
#include <unicode/unistr.h>
#include <unicode/ustream.h>
#include <unicode/utypes.h>

using std::bitset;
using std::cerr;
using std::cout;
using std::endl;
using std::forward;
using std::ifstream;
using std::ios;
using std::nullopt;
using std::optional;
using std::ostream;
using std::ostringstream;
using std::pair;
using std::string;
using std::tie;
using std::unordered_map;
using std::vector;

using boost::hash_range;
using icu::UnicodeString;

typedef unsigned __int128 bigint;

namespace po = boost::program_options;

// total maximum of letters in input
constexpr int MAX_LETTERS = 128;

typedef uint8_t CharIdx;
constexpr int MAX_CHARIDX = std::numeric_limits<CharIdx>::max();

// first 168 primes
constexpr int PRIMES[] = {
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71,
    73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151,
    157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223, 227, 229, 233,
    239, 241, 251, 257, 263, 269, 271, 277, 281, 283, 293, 307, 311, 313, 317,
    331, 337, 347, 349, 353, 359, 367, 373, 379, 383, 389, 397, 401, 409, 419,
    421, 431, 433, 439, 443, 449, 457, 461, 463, 467, 479, 487, 491, 499, 503,
    509, 521, 523, 541, 547, 557, 563, 569, 571, 577, 587, 593, 599, 601, 607,
    613, 617, 619, 631, 641, 643, 647, 653, 659, 661, 673, 677, 683, 691, 701,
    709, 719, 727, 733, 739, 743, 751, 757, 761, 769, 773, 787, 797, 809, 811,
    821, 823, 827, 829, 839, 853, 857, 859, 863, 877, 881, 883, 887, 907, 911,
    919, 929, 937, 941, 947, 953, 967, 971, 977, 983, 991, 997
};

// If the given function returns false, terminate.
template <typename Fn>
bool forAllAlpha(const UnicodeString &s, Fn &&f) {
    for (int j = 0, je = s.length(); j<je; j++) {
        UChar c = s[j];
        if (u_isUAlphabetic(c)) {
            if (!f(c))
                return false;
        }
    }
    return true;
}

// We map the characters in the input to integers starting from 0 (and
// that must fit in CharIdx).
typedef unordered_map<UChar, int> CharMap;

static size_t hash_bigint(const bigint &m) {
    // Let's decide the lowest bits are good enough, even though it's
    // guaranteed we have count(most_common_letter) trailing zero
    // bits.
    return static_cast<size_t>(m);
}

// A multiset of characters.
//
// The allowable characters are passed in the CharMap, which maps
// unicode characters to integers starting from 0.
//
// To facilitate the required multiset operations, namely subtraction
// and subset testing, we map a multiset to an integer (which may not
// fit in a fixed size integer type)
//
//    product(nth_prime(i)**char_count(i)).
//
// Now,
//
// * x is a subset of y iff x.num % y.num == 0
// * x-y maps to integer division.
//
// The order of characters is chosen such that the most common
// characters in the input map to smallest primes, which results in
// smallest numbers.
class CharBag {
public:
    static optional<CharBag> fromUString(const UnicodeString &str, const CharMap &charmap);
    static optional<CharBag> fromLowerUString(const UnicodeString &str, const CharMap &charmap);
    static optional<CharBag> fromNativeString(const string &str, const CharMap &charmap) {
        return fromUString(UnicodeString(str.c_str()), charmap);
    }
    const bigint &num() const { return m_num; }
    const int size() const { return m_size; }
    size_t hash() const { return m_hash; }

    bool empty() const { return m_size == 0; }

    bool isSubsetOf(const CharBag &other) const {
        auto cs_opt = other-*this;
        return (bool)cs_opt;
    }

    optional<CharBag> operator-(const CharBag &rhs) const;
    bool operator==(const CharBag &rhs) const;
private:
    CharBag(bigint num, int size) :
        m_num(num), m_size(size), m_hash(compute_hash(m_size, m_num)) {}

    static size_t compute_hash(const int &size, const bigint &num) {
        // As a hack to find words faster, we want longer words to map
        // to smaller hashes. Do this by mapping negated size to the
        // high bits of the hash.
        constexpr int bits = std::numeric_limits<size_t>::digits;
        size_t hi = (~static_cast<size_t>(size)) << (bits-7);
        size_t lo = hash_bigint(num) & ((size_t(1) << (bits-7)) - 1);
        return hi | lo;
    }

    bigint m_num;
    int m_size;
    size_t m_hash;
};

namespace std {
  template <> struct hash<CharBag> {
    size_t operator()(const CharBag &cs) const {
        return cs.hash();
    }
  };
}

bool CharBag::operator==(const CharBag &rhs) const {
    return m_hash == rhs.m_hash && m_num == rhs.m_num;
}

optional<CharBag> CharBag::operator-(const CharBag &rhs) const {
    if (rhs.m_size > m_size)
        return nullopt;

    if (m_size == rhs.m_size && m_num == rhs.m_num)
        return CharBag{bigint(1), 0};

    // if (!mpz_divisible_p(m_num.get_mpz_t(), rhs.m_num.get_mpz_t()))
    //  return nullopt;
    if (m_num % rhs.m_num != 0)
        return nullopt;

    return CharBag{m_num/rhs.m_num, m_size-rhs.m_size};
}

optional<CharBag> CharBag::fromLowerUString(const UnicodeString &str, const CharMap &charmap) {
    bigint n(1);
    int size = 0;

    static_assert(sizeof(PRIMES)/sizeof(PRIMES[0]) >= MAX_LETTERS);

    if (forAllAlpha(str, [&n, &size, &charmap](UChar c) {
            auto it = charmap.find(c);
            if (it == charmap.end())
                return false;
            int idx = it->second;
            assert(idx < MAX_LETTERS);
            n *= PRIMES[idx];
            ++size;
            return true;
            })) {
        return CharBag(n, size);
    } else
        return nullopt;
}

optional<CharBag> CharBag::fromUString(const UnicodeString &str_, const CharMap &charmap) {
    UnicodeString str(str_);
    str.toLower();
    return fromLowerUString(str, charmap);
}

// [[maybe_unused]]
// static ostream &operator<<(ostream &os, const CharBag &cs) {
//     return os << "CharBag{" << cs.size() << ", " << cs.num() << "}";
// }

// [[maybe_unused]]
// static ostream &operator<<(ostream &os, const optional<CharBag> &cs) {
//     if (cs)
//      os << *cs;
//     else
//      os << "nil";
//     return os;
// }

// from https://stackoverflow.com/questions/17074324/
template <typename T>
void apply_permutation_in_place(std::vector<T>& vec,
                                const std::vector<std::size_t>& p) {
    std::vector<bool> done(vec.size());
    for (std::size_t i = 0; i < vec.size(); ++i) {
        if (done[i])
            continue;
        done[i] = true;
        std::size_t prev_j = i;
        std::size_t j = p[i];
        while (i != j) {
            std::swap(vec[prev_j], vec[j]);
            done[j] = true;
            prev_j = j;
            j = p[j];
        }
    }
}

static vector<char> slurp(const string &fileName) {
    ifstream ifs(fileName.c_str(), ios::binary | ios::ate);

    ifstream::pos_type size = ifs.tellg();
    ifs.seekg(0, ios::beg);

    vector<char> bytes(size);
    ifs.read(bytes.data(), size);

    return bytes;
}

// The dictionary words are sorted by length (number of alphabetic
// characters). This is important, it's used in
// forAllAnagrams_iter_last().
static pair<vector<vector<string>>, vector<CharBag>> loadDictionary(
    const string &fname, const CharBag &cset, const CharMap &cmap) {
    const auto raw_contents = slurp(fname);
    UnicodeString contents = UnicodeString(raw_contents.data(), raw_contents.size());
    contents.toLower();
    size_t contents_len = contents.length();

    vector<vector<string>> words;
    vector<CharBag> charbags;
    unordered_map<CharBag, size_t> charbag_map;

    size_t count = 0;

    size_t line_start = 0;
    while (true) {
        auto newline = contents.indexOf(UChar('\n'), line_start);

        size_t endpos;
        if (newline == -1)
            endpos = contents_len;
        else
            endpos = newline;

        auto str_len = endpos - line_start;
        auto line = contents.tempSubString(line_start, str_len);
        if (str_len > 0) {
            optional<CharBag> cs = CharBag::fromLowerUString(line, cmap);
            if (cs && cset - *cs) {
                if (cs->empty())
                    continue;
                ++count;

                ostringstream line_sstream;
                line_sstream << line;

                auto it = charbag_map.find(*cs);
                if (it == charbag_map.end()) {
                    words.emplace_back(vector<string>{{line_sstream.str()}});
                    charbags.emplace_back(*cs);
                    charbag_map[*cs] = words.size()-1;
                } else
                    words[it->second].emplace_back(line_sstream.str());
            }
        }
        if (newline == -1)
            break;
        else
            line_start = newline+1;
    }

    // Sort the words by charbag hash.

    vector<size_t> hash_sort_order(charbags.size());
    std::iota(hash_sort_order.begin(), hash_sort_order.end(), 0);

    std::sort(hash_sort_order.begin(), hash_sort_order.end(),
              [&charbags](int a, int b) { return charbags[a].hash() < charbags[b].hash(); });

    apply_permutation_in_place(words, hash_sort_order);
    apply_permutation_in_place(charbags, hash_sort_order);

    //cerr << "Loaded " << count << " dictionary words, " << words.size() << " distinct." << endl;
    return {words, charbags};
}

template <typename Fn>
void forAllAnagrams_iter_last(const vector<CharBag> &dict_charbags,
                              const vector<int> &possible_charbags,
                              const CharBag &charbag, Fn &&f,
                              vector<size_t> &words, size_t start_idx) {
    size_t required_hash = charbag.hash();

    auto compare_charbags = [&dict_charbags, required_hash](const int a, const int b) {
        size_t a_hash, b_hash;

        // FIXME how to do this properly? This is really ugly.
        if (a == -1)
            a_hash = required_hash;
        else
            a_hash = dict_charbags[a].hash();
        if (b == -1)
            b_hash = required_hash;
        else
            b_hash = dict_charbags[b].hash();

        return a_hash < b_hash;
    };

    {
        auto it = std::lower_bound(possible_charbags.begin() + start_idx,
                                   possible_charbags.end(), -1, compare_charbags);

        for (auto ie = possible_charbags.end();
             it != ie && dict_charbags[*it].hash() == required_hash; ++it)
            if (charbag == dict_charbags[*it]) {
                words.emplace_back(*it);
                f(words);
                words.pop_back();
                // There's at most one hit, and we found one.
                return;
            }
    }
}

template <typename Fn>
void forAllAnagrams_iter(const vector<CharBag> &dict_charbags,
                         const vector<int> &old_possible_charbags,
                         const CharBag &charbag, Fn &&f,
                         vector<size_t> &words, size_t start_idx, int curr_len, int max_len) {
    if (curr_len+1 >= max_len) {
        forAllAnagrams_iter_last(dict_charbags, old_possible_charbags, charbag, f,
                                 words, start_idx);
        return;
    }
    vector<int> possible_charbags;
    for (int i = start_idx, ie = old_possible_charbags.size(); i<ie; i++)
        if (dict_charbags[old_possible_charbags[i]].isSubsetOf(charbag))
            possible_charbags.emplace_back(old_possible_charbags[i]);

    for (int i = 0, ie = possible_charbags.size(); i<ie; i++) {
        auto cs = (charbag - dict_charbags[possible_charbags[i]]).value();
        words.emplace_back(possible_charbags[i]);
        if (cs.empty())
            f(words);
        else
            forAllAnagrams_iter(dict_charbags, possible_charbags,
                                cs, forward<Fn>(f), words, i, curr_len+1, max_len);
        words.pop_back();
    }
}

template <typename Fn>
void forAllAnagrams(const vector<CharBag> &dict_charbags, const CharBag &charbag,
                    int max_len, Fn &&f) {
    vector<size_t> words;
    vector<int> possible_charbags(dict_charbags.size());
    std::iota(possible_charbags.begin(), possible_charbags.end(), 0);
    forAllAnagrams_iter(dict_charbags, possible_charbags,
                        charbag, forward<Fn>(f), words, 0, 0, max_len);
}

// The words vector contains vectors of anagram-equivalent words.
// Output all possible combinations of them.
static void outputWords(ostream &stream, const vector<size_t> &word_idxs,
                        const vector<vector<string>> &words) {
    int size = word_idxs.size();
    vector<size_t> idxs(size);

    while (true) {
        stream << words[word_idxs[0]][idxs[0]];
        for (int i = 1; i<size; i++)
            stream << " " << words[word_idxs[i]][idxs[i]];
        stream << endl;

        // increment
        int curr = size-1;
        while (curr >= 0 && ++idxs[curr] == words[word_idxs[curr]].size()) {
            idxs[curr] = 0;
            curr--;
        }
        if (curr < 0)
            break;
    }
}

static void usage(char * const *argv, po::options_description &visible) {
    cout << "Usage: " << argv[0] << " [options] sentence" << endl;
    cout << visible << endl;
    exit(0);
}

static po::variables_map parse_args(int argc, char * const *argv) {
    bool help = false;

    po::options_description visible("Allowed options"), cmdline_opt;
    visible.add_options()
        ("help,h", po::bool_switch(&help)->default_value(false), "show this help")
        ("dict,d", po::value<string>()->default_value("words.txt"),
         "dictionary (word list) to use")
        ("len,l", po::value<int>()->default_value(3), "maximum anagram length in words")
        ("remove,r", po::value<string>()->default_value(""),
         "remove characters from input (find anagrams containing given string)");

    po::options_description hidden;
    hidden.add_options()
        ("sentence", po::value<string>()->required(), "sentence");

    cmdline_opt.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("sentence", 1);

    po::variables_map vm;
    try {
        po::store(po::command_line_parser(argc, argv).options(cmdline_opt).positional(p).run(), vm);
        po::notify(vm);
    } catch (po::error &) {
        help = true;
    }

    if (help)
        usage(argv, visible);

    return vm;
}

CharMap generateCharMap(const UnicodeString &input) {
    CharMap charmap;

    unordered_map<UChar, int> char_counts;

    forAllAlpha(input, [&char_counts](UChar c) {
            ++char_counts[c];
            return true;
        });

    if (char_counts.size() > MAX_LETTERS) {
        cerr << "Error: More than " << MAX_CHARIDX+1
             << " different characters in input." << endl;
        exit(1);
    }

    vector<pair<int, UChar>> charmap_with_counts;
    for (const auto &it : char_counts)
        charmap_with_counts.emplace_back(-it.second, it.first);

    std::sort(charmap_with_counts.begin(), charmap_with_counts.end());

    int i = 0;
    for (const auto &p : charmap_with_counts)
        charmap[p.second] = i++;

    return charmap;
}

int main(int argc, char **argv) {
    std::locale::global(std::locale(""));

    auto vm = parse_args(argc, argv);

    UnicodeString input = UnicodeString(vm["sentence"].as<string>().c_str()).toLower();
    UnicodeString remove = UnicodeString(vm["remove"].as<string>().c_str()).toLower();

    CharMap charmap;
    vector<UChar> reverse_charmap;

    charmap = generateCharMap(input);

    auto input_charbag = CharBag::fromUString(input, charmap).value();
    auto remove_charbag_opt = CharBag::fromUString(remove, charmap);

    if (remove_charbag_opt == nullopt || !remove_charbag_opt.value().isSubsetOf(input_charbag)) {
        cerr << "Error: characters for substring \"" << remove
                << "\" not found in string \"" << input
                << "\"!" << endl;
        exit(2);
    }

    auto stripped_input_charbag = (input_charbag - remove_charbag_opt.value()).value();

    vector<vector<string>> dict_words;
    vector<CharBag> dict_charbags;
    tie(dict_words, dict_charbags) = loadDictionary(
        vm["dict"].as<string>(), stripped_input_charbag, charmap);

    forAllAnagrams(dict_charbags, stripped_input_charbag, vm["len"].as<int>(),
                   [&dict_words](const vector<size_t> &word_idxs) {
            assert(!word_idxs.empty());
            outputWords(cout, word_idxs, dict_words);
        });
}


// -*- mode: C++; c-indent-level: 4; c-basic-offset: 4; indent-tabs-mode: nil; -*-
//   Learn more about Rcpp at:
//
//   http://www.rcpp.org/
//   http://adv-r.had.co.nz/Rcpp.html
//   http://gallery.rcpp.org/
//

// we only include RcppArmadillo.h which pulls Rcpp.h in for us
#include "RcppArmadillo.h"
using namespace Rcpp;

// via the depends attribute we tell Rcpp to create hooks for
// RcppArmadillo so that the build process will know what to do
//
// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::depends(BH)]]

// one include file from Boost to access the digamma function
#include <boost/math/special_functions/digamma.hpp>

/*----------------------------------------------------------------------------*/
// SlalomModel class definition and module //

//' SlalomModel C++ class
//'
//' @description
//' A C++ class for SlalomModel models.
//'
//' @param Y_init matrix of expression values
//' @param pi_init G x K matrix with each entry being the prior
//' probability for a gene g being active for factor k.
//' @param X_init matrix of initial factor states (N x K)
//' @param W_init G x K matrix of initial weights
//' @param prior_alpha numeric vector of length two giving prior values
//' for the gamma hyperparameters of the precisions
//' @param prior_epsilon numeric vector of length two giving prior values
//' for the gamma hyperparameters of the residual variances
//'
//' @return
//' an object of the SlalomModel class
//'
//' @useDynLib slalom, .registration=TRUE, .fixes="Rcpp_"
//' @importFrom Rcpp evalCpp
//' @exportClass Rcpp_SlalomModel
//' @name SlalomModel
//' @aliases SlalomModel
class SlalomModel {
    // base class for slalom holding models
public:
    // declare necessary variables for alpha
    double alpha_pa;
    double alpha_pb;
    arma::vec alpha_a;
    arma::vec alpha_b;
    arma::vec alpha_E1;
    arma::vec alpha_lnE1;
    // declare necessary variables for epsilon
    double epsilon_pa;
    double epsilon_pb;
    arma::vec epsilon_a;
    arma::vec epsilon_b;
    arma::vec epsilon_E1;
    arma::vec epsilon_lnE1;
    arma::vec epsilon_diagSigmaS;
    // declare necessary variables for X
    arma::mat X_E1;
    arma::mat X_diagSigmaS;
    arma::mat X_init;
    // declare necessary variables for W
    arma::mat W_E1;
    arma::mat W_sigma2;
    arma::mat W_E2diag;
    arma::mat W_gamma0;
    arma::mat W_gamma1;
    // declare necessary variables for Z
    arma::mat Z_E1;
    // declare other variables
    arma::mat I;
    arma::mat Pi_a;
    arma::mat Pi_b;
    arma::mat Pi_pa;
    arma::mat Pi_E1;
    arma::mat Y;
    arma::mat pseudo_Y;
    arma::vec YY;
    arma::mat Known;
    int K;
    int N;
    int G;
    int nAnnotated;
    int nHidden;
    int nKnown;
    int nIterations;
    int minIterations;
    int iterationCount;
    double tolerance;
    bool forceIterations;
    bool shuffle;
    bool converged;
    double nScale;
    std::vector< std::string > noiseModel;
    double onF;
    arma::vec nOn;
    arma::vec iUnannotatedDense;
    arma::vec iUnannotatedSparse;
    arma::uvec doUpdate;
    arma::uvec pretrain_order;
    bool dropFactors;
    bool learnPi;
    // names
    Rcpp::StringVector termNames;
    Rcpp::StringVector cellNames;
    Rcpp::StringVector geneNames;
    // methods
    void train(void);
    void update(void);
    void updateEpsilon(void);
    void updateAlpha(const int);
    void updateX(const int);
    void updateW(const int);
    void updatePi(const int);
    // constructor
    SlalomModel() {}
    SlalomModel(arma::mat Y_init, arma::mat pi_init, arma::mat X_init,
                arma::mat W_init, arma::vec prior_alpha, arma::vec prior_epsilon) :
        Y(Y_init), Pi_E1(pi_init), I(pi_init), W_gamma0(pi_init), X_E1(X_init),
        W_E1(W_init) {
        K = pi_init.n_cols;
        N = Y_init.n_rows;
        G = Y_init.n_cols;
        nScale = 100.0;
        // initialise squared column sums
        YY = arma::zeros(G);
        for (int g=0; g < G; g++) {
            YY(g) = arma::accu(Y.col(g) % Y.col(g));
        }
        // initialise alpha variables
        alpha_pa = prior_alpha[0];
        alpha_pb = prior_alpha[1];
        alpha_a = arma::ones(K) * alpha_pa;
        alpha_b = arma::ones(K) * alpha_pb;
        alpha_E1 = alpha_b / alpha_a;
        alpha_lnE1 = arma::ones(K);
        for (int i = 0; i < K; i++) {
            alpha_lnE1(i) = boost::math::digamma(alpha_a(i));
        }
        alpha_lnE1 = alpha_lnE1 - arma::log(alpha_b);
        // initialise epsilon variables
        epsilon_pa = prior_epsilon[0];
        epsilon_pb = prior_epsilon[1];
        epsilon_a = arma::ones(G) * epsilon_pa;
        epsilon_b = arma::ones(G) * epsilon_pb;
        epsilon_E1 = epsilon_b / epsilon_a;
        epsilon_lnE1 = arma::ones(G);
        for (int i = 0; i < G; i++) {
            epsilon_lnE1(i) = boost::math::digamma(epsilon_a(i));
        }
        epsilon_lnE1 = epsilon_lnE1 - arma::log(epsilon_b);
        epsilon_diagSigmaS = arma::zeros(G);
        // initialise X variables
        X_diagSigmaS = arma::ones(N, K);
        // initialise W variables
        W_gamma1 = 1.0 - W_gamma0;
        W_sigma2 = arma::ones(G, K);
        W_E2diag = arma::zeros(G, K);
        // initialise iterations
        converged = false;
        iterationCount = 0;
        tolerance = 1e-08;
        learnPi = true;
        doUpdate = arma::ones<arma::uvec>(K);
        minIterations = 1;
        nIterations = 1;
    }
};

/*----------------------------------------------------------------------------*/
// SlalomModel class train method

// train method
void SlalomModel::train(void) {
    /*
     Iterate updates of weights (with spike-and-slab prior), ARD parameters, factors, annd noise parameters.

     No arguments, but utilises the following parameters defined in the object:
     nIterations          (int): Number of iterations.
     forceIterations     (bool): Force the model to update `nIteration` times.
     tolerance          (float): Tolerance to monitor convergence of reconstruction error
     minIterations        (int): Minimum number of iterations the model should perform.
     */
    arma::umat Ion = (this->W_gamma0 > .5);
    arma::mat Zr = this->X_E1 * (this->W_E1.t() % Ion.t());
    arma::mat Zd = this->Z_E1 - Zr;
    arma::rowvec error(1);
    arma::rowvec error_old(1);
    error.fill(arma::mean(arma::mean((arma::abs(Zd)))));
    double meanerr = arma::mean(error);
    // mean absolute error
    error_old.fill(100);
    this->converged = false;
    // iterate model to train it
    this->iterationCount = 0;
    for (int iter = 0; iter <= this->nIterations; iter++) {
        // t = time.time();
        this->update();
        this->iterationCount++;
        if (iter % 100 == 0) {
            Rcout << "iteration " << iter << std::endl;
        }
        if (iter % 50 == 0) {
            error_old = error;
            Zr = this->X_E1 * this->W_E1.t();
            Zd = this->Z_E1 - Zr;
            error.fill(arma::mean(arma::mean((arma::abs(Zd)))));
            // mean absolute error
            this->converged = arma::approx_equal(error_old, error, "absdiff",
                                           this->tolerance);
            double meanerr = arma::mean(error);
        }
        if ( this->converged && !(this->forceIterations) &&
             (iter > this->minIterations) ) {
            Rcout << "Model converged after " << iter << " iterations." << std::endl;
            break;
        }
        // this->Z_E1 = this->X_E1 * this->W_E1.t();
    }
    if (!converged) {
        Rcout << "Model not converged after " <<  this->nIterations << " iterations." << std::endl;
    }
}

/*----------------------------------------------------------------------------*/
// SlalomModel class update methods //

// wrapper around R's RNG such that we get a uniform distribution over
// [0,n) as required by the STL algorithm
inline int randWrapper(const int n) { return floor(unif_rand() * n); }

void SlalomModel::update(void) {
    /*
     Do one update of weights (with spike-and-slab prior), ARD parameters, factors, annd noise parameters.

     */
    this->epsilon_diagSigmaS = arma::zeros(this->K);
    // arma::umat Ion = (this->W_gamma0 > .5);
    // check above: set parameter to zero with every update?
    arma::uvec kRange;
    if (this->iterationCount == 1) {
        kRange = this->pretrain_order;
    } else {
        kRange = arma::regspace<arma::uvec>(0,  (this->K - 1));
    }
    if (this->shuffle == true && this->iterationCount > 1) {
        arma::uvec kunfix = arma::regspace<arma::uvec>(this->nKnown,
                                                       (this->K - 1));
        std::random_shuffle(kunfix.begin(), kunfix.end(), randWrapper);
        for (int i = 0; i < (this->K - this->nKnown); ++i) {
            int ii = i + this->nKnown;
            kRange(ii) = kunfix(i);
        }
    }
    // switch factors off such that they can't be turned back on
    // Regular iterator, non-C++11
    for (int m = 0; m < this->K; m++) {
        int k = kRange[m];
        if (this->doUpdate[k]) {
            if (this->dropFactors == false || this->iterationCount < 10 ||
                (this->alpha_E1[k] / arma::var(this->X_E1.col(k))) < 1e10) {
                this->updateW(k);
                if (this->learnPi) {
                    if (arma::any((this->iUnannotatedSparse - 1) == k)) {
                        this->updatePi(k);
                    }
                }
                this->updateAlpha(k);
                this->updateX(k);
            } else {
                this->doUpdate[k] = 0;
                Rcout << "Switched off factor " << k << std::endl;
            }
        }
    }
    // const char noise_gauss = "gauss";
    // if (std::strcmp(this->noiseModel, noise_gauss) == 0) {
    //     // this->updateEpsilon();
    // };
    this->updateEpsilon();
}

void SlalomModel::updateW(const int k) {
    // update the factor weights
    // define logPi values
    arma::vec logPi;
    int Muse = arma::accu(this->doUpdate);
    if (k < this->nKnown || arma::any((this->iUnannotatedSparse - 1) == k) ||
        arma::any((this->iUnannotatedDense - 1) == k)) {
        logPi = arma::log(this->Pi_E1.col(k) / (1.0 - this->Pi_E1.col(k)));
        // careful of divide-by-zero errors here
    } else if (this->nScale > 0 && this->nScale < this->N) {
        logPi = arma::log(this->Pi_E1.col(k) / (1.0 - this->Pi_E1.col(k)));
        // careful of divide-by-zero errors here
        arma::uvec isOFF_ = arma::find(this->Pi_E1.col(k) < 0.5);
        arma::uvec kvec(1);
        kvec(0) = k;
        logPi(isOFF_) = ((this->N / this->nScale) *
            arma::log(this->Pi_E1(isOFF_, kvec) /
                (1 - this->Pi_E1(isOFF_, kvec))));
        // careful of divide-by-zero errors here
        arma::uvec isON_ = arma::find(this->Pi_E1.col(k) > 0.5);
        if (this->onF > 1.0) {
            logPi(isON_) = this->onF * arma::log(
                this->Pi_E1(isON_, kvec) / (1 - this->Pi_E1(isON_, kvec)));
        }
    } else {
        logPi = arma::log(this->Pi_E1.col(k) / (1 - this->Pi_E1.col(k)));
    }
    arma::vec sigma2Sigmaw;
    sigma2Sigmaw = (1.0 / this->epsilon_E1) * this->alpha_E1(k);
    arma::uvec set1 = arma::regspace<arma::uvec>(0, 1, k - 1);    // zero-indexing
    arma::uvec set2 = arma::regspace<arma::uvec>(k + 1, 1, this->K - 1);
    arma::uvec setMinus = arma::join_cols(set1, set2);
    arma::uvec idx = arma::find(this->doUpdate(setMinus) == 1);
    setMinus = setMinus(idx);

    arma::vec SmTSk = arma::sum(
        (arma::repmat(this->X_E1.col(k), 1, Muse - 1) %
             this->X_E1.cols(setMinus)), 0).t();
    double tmp = arma::as_scalar(this->X_E1.col(k).t() * this->X_E1.col(k));
    double SmTSm = (tmp + arma::sum(this->X_diagSigmaS.col(k)));

    arma::vec b;
    arma::vec diff;
    b = (this->W_gamma0.cols(setMinus) % this->W_E1.cols(setMinus)) * SmTSk;
    diff = (this->X_E1.col(k).t() * this->Z_E1).t() - b;
    arma::vec diff2 = diff % diff;
    arma::vec SmTSmSig = SmTSm + sigma2Sigmaw; // same as Py

    // update gamma and W
    arma::vec u_qm = (logPi + 0.5 * arma::log(sigma2Sigmaw) - 0.5 *
        arma::log(SmTSmSig) + (0.5 * this->epsilon_E1) %
        (diff2 / SmTSmSig));
    this->W_gamma0.col(k) = 1.0 / (1 + arma::exp(-u_qm));
    this->W_gamma1.col(k) = 1.0 - this->W_gamma0.col(k);
    this->W_E1.col(k) = (diff / SmTSmSig);
    this->W_sigma2.col(k) = (1.0 / this->epsilon_E1) / SmTSmSig;
    this->W_E2diag.col(k) = ((this->W_E1.col(k) % this->W_E1.col(k)) +
        this->W_sigma2.col(k));
}

void SlalomModel::updatePi(const int k) {
    // update Pi in a SlalomModel class object
    this->Pi_a.col(k) = this->Pi_pa.col(k) + arma::accu(this->W_gamma0.col(k));
    this->Pi_b.col(k) = (this->Pi_pa.col(k) + this->G -
        arma::accu(this->W_gamma0.col(k)));
    this->Pi_E1.col(k) = this->Pi_a.col(k) / this->Pi_a.col(k) +
        this->Pi_b.col(k);
}

void SlalomModel::updateAlpha(const int k) {
    // pdate alpha in a SlalomModel class object - precisions
    double Ewdwd = arma::accu(this->W_gamma0.col(k) % this->W_E2diag.col(k));
    // elementwise mult.
    this->alpha_a(k) = this->alpha_pa + 0.5 * Ewdwd;
    this->alpha_b(k) = this->alpha_pb + arma::accu(this->W_gamma0.col(k)) / 2.0;
    this->alpha_E1(k) = this->alpha_b(k) / this->alpha_a(k);
}

void SlalomModel::updateX(const int k) {
    // update the factor states
    arma::uvec set1 = arma::regspace<arma::uvec>(0, 1, k - 1);
    arma::uvec set2 = arma::regspace<arma::uvec>(k + 1, 1, this->K - 1);
    arma::uvec setMinus = arma::join_cols(set1, set2);
    arma::uvec idx = arma::find(this->doUpdate(setMinus) == 1);
    setMinus = setMinus(idx);

    arma::vec SW_sigma;
    arma::vec SW2_sigma;

    SW2_sigma = ((this->W_gamma0.col(k) % this->W_E2diag.col(k)) %
                     this->epsilon_E1);
    double alphaSm = arma::sum(SW2_sigma);

    for (int i = 0; i < this->N; i++) {
        this->X_diagSigmaS(i, k) = 1.0 / (1.0 + alphaSm);
    }

    if (k >= this->nKnown) {
        SW_sigma = (this->W_gamma0.col(k) %
                        this->W_E1.col(k)) % this->epsilon_E1;     // Gx1 vec
        arma::mat b0 = (this->X_E1.cols(setMinus) *
            (this->W_gamma0.cols(setMinus) % this->W_E1.cols(setMinus)).t());  // NxG mat.
        arma::vec b = b0 * SW_sigma;    // NxG x Gx1 -> Nx1 vec
        arma::vec barmuX = (this->Z_E1 * SW_sigma) - b;
        // update X
        this->X_E1.col(k) = barmuX / (1.0 + alphaSm);
        // keep diagSigmaS
        this->epsilon_diagSigmaS(k) = arma::accu(this->X_diagSigmaS.col(k));
    }
}

void SlalomModel::updateEpsilon(void) {
    // update Epsilon (vectorised) - noise parameters
    arma::uvec update_cols = arma::find(this->doUpdate == 1);
    arma::mat SW_sigma = (this->W_gamma0.cols(update_cols) %
                              this->W_E1.cols(update_cols));  // elementwise mult.; GxK mat
    arma::mat SW2_sigma = (this->W_gamma0.cols(update_cols) %
                               this->W_E2diag.cols(update_cols));  // elementwise mult.; GxK mat
    arma::mat muSTmuS = (this->X_E1.cols(update_cols).t() *
        this->X_E1.cols(update_cols));     // KxK matrix
    arma::mat newmat = this->Z_E1.t() * this->X_E1.cols(update_cols);
    arma::vec t1 = arma::sum(SW_sigma % newmat, 1);  // K length vec
    arma::vec t2 = arma::sum(SW2_sigma % arma::repmat(muSTmuS.diag().t() +
        this->epsilon_diagSigmaS(update_cols).t(), this->G, 1), 1);
    // set diagonals to zeros in muSTmuS for next calculation
    muSTmuS.diag().zeros();
    arma::vec t3 = arma::sum((SW_sigma * muSTmuS) % SW_sigma, 1);
    this->epsilon_E1 = 1.0 / ((0.5 * (this->YY  + (-2 * t1  + t2 + t3))) /
        (0.5 * this->N));
    this->epsilon_a.fill(0.5 * this->N + this->epsilon_pa);
    this->epsilon_b = this->epsilon_pb + 0.5 * (this->YY + (-2 * t1 + t2 + t3));
    for (int i = 0; i < this->epsilon_E1.n_elem; i++) {
        if (this->epsilon_E1(i) > 1e6) {
            this->epsilon_E1(i) = 1e6;
        }
    }
}

/*----------------------------------------------------------------------------*/
// SlalomModel Rcpp module

// Define a module to make the C++ class available to R
RCPP_MODULE(SlalomModel) {
    using namespace Rcpp;

    class_<SlalomModel>("SlalomModel")
        // expose the default constructor
        .constructor()
        .constructor<arma::mat, arma::mat, arma::mat, arma::mat, arma::vec,
        arma::vec>()
        //
        // fields
        .field("K", &SlalomModel::K)
        .field("N", &SlalomModel::N)
        .field("G", &SlalomModel::G)
        .field("nScale", &SlalomModel::nScale)
        .field("nAnnotated", &SlalomModel::nAnnotated)
        .field("nHidden", &SlalomModel::nHidden)
        .field("nKnown", &SlalomModel::nKnown)
        .field("nIterations", &SlalomModel::nIterations)
        .field("minIterations", &SlalomModel::minIterations)
        .field("iterationCount", &SlalomModel::iterationCount)
        .field("forceIterations", &SlalomModel::forceIterations)
        .field("tolerance", &SlalomModel::tolerance)
        .field("shuffle", &SlalomModel::shuffle)
        .field("converged", &SlalomModel::converged)
        .field("noiseModel", &SlalomModel::noiseModel)
        .field("onF", &SlalomModel::onF)
        // alpha
        .field("alpha_pa", &SlalomModel::alpha_pa)
        .field("alpha_pb", &SlalomModel::alpha_pb)
        .field("alpha_a", &SlalomModel::alpha_a)
        .field("alpha_b", &SlalomModel::alpha_b)
        .field("alpha_E1", &SlalomModel::alpha_E1)
        .field("alpha_lnE1", &SlalomModel::alpha_lnE1)
        // epsilon
        .field("epsilon_pa", &SlalomModel::epsilon_pa)
        .field("epsilon_pb", &SlalomModel::epsilon_pb)
        .field("epsilon_a", &SlalomModel::epsilon_a)
        .field("epsilon_b", &SlalomModel::epsilon_b)
        .field("epsilon_E1", &SlalomModel::epsilon_E1)
        .field("epsilon_lnE1", &SlalomModel::epsilon_lnE1)
        .field("epsilon_diagSigmaS", &SlalomModel::epsilon_diagSigmaS)
        // X
        .field("X_E1", &SlalomModel::X_E1)
        .field("X_diagSigmaS", &SlalomModel::X_diagSigmaS)
        .field("X_init", &SlalomModel::X_init)
        // W
        .field("W_E1", &SlalomModel::W_E1)
        .field("W_sigma2", &SlalomModel::W_sigma2)
        .field("W_E2diag", &SlalomModel::W_E2diag)
        .field("W_gamma0", &SlalomModel::W_gamma0)
        .field("W_gamma1", &SlalomModel::W_gamma1)
        // Z
        .field("Z_E1", &SlalomModel::Z_E1)
        // Pi
        .field("Pi_a", &SlalomModel::Pi_a)
        .field("Pi_pa", &SlalomModel::Pi_pa)
        .field("Pi_b", &SlalomModel::Pi_b)
        .field("Pi_E1", &SlalomModel::Pi_E1)
        .field("I", &SlalomModel::I)
        // other variables
        .field("Known", &SlalomModel::Known)
        .field("Y", &SlalomModel::Y)
        .field("pseudo_Y", &SlalomModel::pseudo_Y)
        .field("YY", &SlalomModel::YY)
        .field("iUnannotatedDense", &SlalomModel::iUnannotatedDense)
        .field("iUnannotatedSparse", &SlalomModel::iUnannotatedSparse)
        .field("nOn", &SlalomModel::nOn)
        .field("doUpdate", &SlalomModel::doUpdate)
        .field("pretrain_order", &SlalomModel::pretrain_order)
        .field("learnPi", &SlalomModel::learnPi)
        .field("dropFactors", &SlalomModel::dropFactors)
        // names
        .field("termNames", &SlalomModel::termNames)
        .field("cellNames", &SlalomModel::cellNames)
        .field("geneNames", &SlalomModel::geneNames)
        // methods
        .method("train", &SlalomModel::train , "Train the SlalomModel")
        .method("update", &SlalomModel::update , "Update the SlalomModel")
        .method("updateW", &SlalomModel::updateW , "Update W")
        .method("updateX", &SlalomModel::updateX , "Update X")
        .method("updatePi", &SlalomModel::updatePi , "Update Pi")
        .method("updateEpsilon", &SlalomModel::updateEpsilon , "Update Epsilon")
        .method("updateAlpha", &SlalomModel::updateAlpha , "Update alpha")
        ;
}


#ifndef TCM_BLAS_WRAPPER_HPP
#define TCM_BLAS_WRAPPER_HPP

#include <type_traits>
#include <complex>
#include <cassert>

#include <boost/numeric/conversion/cast.hpp>

#include <detail/config.hpp>
#include <detail/utils.hpp>
#include <detail/iterator.hpp>

///////////////////////////////////////////////////////////////////////////////
/// \file blas_wrapper.hpp
/// \brief Defines template versions of some BLAS routines.
///////////////////////////////////////////////////////////////////////////////

#define REGISTER_BLAS(general_f, s_f, d_f, c_f, z_f)                           \
	namespace {                                                                \
		template<class... Args>                                                \
        __attribute__((always_inline))                                         \
		inline                                                                 \
		auto general_f##_impl( utils::Type2Type<float>                         \
		                     , Args&&... args) noexcept                        \
		{ return s_f(std::forward<Args>(args)...); }                           \
		                                                                       \
		template<class... Args>                                                \
        __attribute__((always_inline))                                         \
		inline                                                                 \
		auto general_f##_impl( utils::Type2Type<double>                        \
		                     , Args&&... args) noexcept                        \
		{ return d_f(std::forward<Args>(args)...); }                           \
		                                                                       \
		template<class... Args>                                                \
        __attribute__((always_inline))                                         \
		inline                                                                 \
		auto general_f##_impl( utils::Type2Type<std::complex<float>>           \
		                     , Args&&... args) noexcept                        \
		{ return c_f(std::forward<Args>(args)...); }                           \
		                                                                       \
		template<class... Args>                                                \
        __attribute__((always_inline))                                         \
		inline                                                                 \
		auto general_f##_impl( utils::Type2Type<std::complex<double>>          \
		                     , Args&&... args) noexcept                        \
		{ return z_f(std::forward<Args>(args)...); }                           \
	}                                                                          \
                                                                               \
	template<class T, class... Args>                                           \
	__attribute__((always_inline))                                             \
	inline                                                                     \
	auto general_f(Args&&... args) noexcept                                    \
	{ return general_f##_impl( utils::Type2Type<T>{}                           \
	                         , std::forward<Args>(args)...);                   \
	}

namespace tcm {

namespace import {

enum class Operator : char {None = 'N', T = 'T', H = 'C'};

#ifdef USING_INTEL_MKL
	using blas_int = long long int;
#else
#	ifdef USING_ATLAS
		using blas_int = int;
#	else
#		error "Need BLAS"
#	endif
#endif

// ============================================================================
//                                DOT PRODUCT
// ============================================================================

extern "C" {

float sdot_
( blas_int const* N
, float const* X, blas_int const* INCX
, float const* Y, blas_int const* INCY );

double ddot_
( blas_int const* N
, double const* X, blas_int const* INCX
, double const* Y, blas_int const* INCY );

std::complex<float> cdotc_
( blas_int const* N
, std::complex<float> const* X, blas_int const* INCX
, std::complex<float> const* Y, blas_int const* INCY );

std::complex<double> zdotc_
( blas_int const* N
, std::complex<double> const* X, blas_int const* inc_X
, std::complex<double> const* Y, blas_int const* inc_Y );

} // extern "C"

REGISTER_BLAS(dotc, sdot_, ddot_, cdotc_, zdotc_)

template< class _T
        , class = std::enable_if_t
                  <   std::is_same<_T, float>()
                   or std::is_same<_T, double>()
                   or std::is_same<_T, std::complex<float>>()
                   or std::is_same<_T, std::complex<double>>()
                  >
        >
inline
auto dot( std::size_t const n
        , _T const* X, blas_int const INCX
        , _T const* Y, blas_int const INCY
        ) -> _T
{
	if(n == 0) return _T{0.0};

	assert(X != nullptr and INCX != 0);
	assert(Y != nullptr and INCY != 0);
	/*
	utils::assert_valid( const_blas_iterator<_T>{X           , INCX}
	                   , const_blas_iterator<_T>{X + n * INCX, INCX} );
	utils::assert_valid( const_blas_iterator<_T>{Y           , INCY}
	                   , const_blas_iterator<_T>{Y + n * INCY, INCY} );
	*/

	auto const N = boost::numeric_cast<blas_int>(n);
	auto const result = dotc<_T>(&N, X, &INCX, Y, &INCY);

	utils::assert_valid(result);
	return result;
}

// ============================================================================
//                                   AXPY
// ============================================================================

extern "C" {

void saxpy_
( blas_int const* N, float const* A
, float const* X, blas_int const* INCX
, float      * Y, blas_int const* INCY );

void daxpy_
( blas_int const* N, double const* A
, double const* X, blas_int const* INCX
, double      * Y, blas_int const* INCY );

void caxpy_
( blas_int const* N, std::complex<float> const* A
, std::complex<float> const* X, blas_int const* INCX
, std::complex<float>      * Y, blas_int const* INCY );

void zaxpy_
( blas_int const* N, std::complex<double> const* A
, std::complex<double> const* X, blas_int const* INCX
, std::complex<double>      * Y, blas_int const* INCY );

} // extern "C"

REGISTER_BLAS(axpy, saxpy_, daxpy_, caxpy_, zaxpy_)

template< class _T
        , class = std::enable_if_t
                  <   std::is_same<_T, float>()
                   or std::is_same<_T, double>()
                   or std::is_same<_T, std::complex<float>>()
                   or std::is_same<_T, std::complex<double>>()
                  >
        >
inline
auto axpy( std::size_t const n, _T const a
         , _T const* X, blas_int const INCX
         , _T      * Y, blas_int const INCY ) -> void
{
	if(n == 0) return;

	assert(X != nullptr and INCX != 0);
	assert(Y != nullptr and INCY != 0);
	/*
	assert_valid( const_blas_iterator<_T>{X           , INCX}
	            , const_blas_iterator<_T>{X + n * INCX, INCX} );
	assert_valid( const_blas_iterator<_T>{Y           , INCY}
	            , const_blas_iterator<_T>{Y + n * INCY, INCY} );
	*/
	auto const N = boost::numeric_cast<blas_int>(n);

	axpy<_T>(&N, &a, X, &INCX, Y, &INCY);

	/*
	assert_valid( const_blas_iterator<_T>{Y           , INCY}
	            , const_blas_iterator<_T>{Y + n * INCY, INCY} );
	*/
}

// ============================================================================
//                         MATRIX-VECTOR MULTIPLICATION
// ============================================================================

extern "C" {

void sgemv_
( char const* TRANSA
, blas_int const* M, blas_int const* N
, float const* ALPHA, float const* A, blas_int const* LDA
, float const* X, blas_int const* INCX
, float const* BETA, float* Y, blas_int const* INCY );

void dgemv_
( char const* TRANSA
, blas_int const* M, blas_int const* N
, double const* ALPHA, double const* A, blas_int const* LDA
, double const* X, blas_int const* INCX
, double const* BETA, double* Y, blas_int const* INCY );

void cgemv_
( char const* TRANSA
, blas_int const* M, blas_int const* N
, std::complex<float> const* ALPHA
, std::complex<float> const* A, blas_int const* LDA
, std::complex<float> const* X, blas_int const* INCX
, std::complex<float> const* BETA
, std::complex<float>* Y, blas_int const* INCY );

void zgemv_
( char const* TRANSA
, blas_int const* M, blas_int const* N
, std::complex<double> const* ALPHA
, std::complex<double> const* A, blas_int const* LDA
, std::complex<double> const* X, blas_int const* INCX
, std::complex<double> const* BETA
, std::complex<double>* Y, blas_int const* INCY );

} // extern "C"

REGISTER_BLAS(gemv, sgemv_, dgemv_, cgemv_, zgemv_)

template< class _T
        , class = std::enable_if_t
                  <   std::is_same<_T, float>()
                   or std::is_same<_T, double>()
                   or std::is_same<_T, std::complex<float>>()
                   or std::is_same<_T, std::complex<double>>()
                  >
        >
auto gemv( Operator const op_A
         , std::size_t const m, std::size_t const n
         , _T const ALPHA, _T const* A, blas_int const LDA
         , _T const* X, blas_int const INCX
         , _T const BETA, _T* Y, blas_int const INCY ) -> void
{
	if(n == 0 or m == 0) return;

	assert(A != nullptr and LDA  != 0);
	assert(X != nullptr and INCX != 0);
	assert(Y != nullptr and INCY != 0);
	/*
	utils::assert_valid(A, A + LDA * n);
	utils::assert_valid
		( const_blas_iterator<_T>{X, INCX}
		, const_blas_iterator<_T>{ X + (op_A == Operator::None ? n : m) * INCX
		                         , INCX }
		);
	utils::assert_valid
		( const_blas_iterator<_T>{Y, INCY}
		, const_blas_iterator<_T>{ Y + (op_A == Operator::None ? m : n) * INCY
		                         , INCY }
		);
	*/

	auto const TRANS = static_cast<char>(op_A);
	auto const M     = boost::numeric_cast<blas_int>(m);
	auto const N     = boost::numeric_cast<blas_int>(n);

	assert(LDA >= M);

	gemv<_T>( &TRANS, &M, &N
	        , &ALPHA, A, &LDA, X, &INCX
	        , &BETA, Y, &INCY );

	/*
	utils::assert_valid
		( const_blas_iterator<_T>{Y, INCY}
		, const_blas_iterator<_T>{ Y + (op_A == Operator::None ? m : n) * INCY
		                         , INCY }
		);
	*/
}

// ============================================================================
//                           MATRIX-MATRIX MULTIPLICATION
// ============================================================================

extern "C" {

void sgemm_
( char const* TRANSA, char const* TRANSB
, blas_int const* M, blas_int const* N, blas_int const* K
, float const* ALPHA
, float const* A, blas_int const* LDA
, float const* B, blas_int const* LDB
, float const* BETA
, float* C, blas_int const* LDC );

void dgemm_
( char const* TRANSA, char const* TRANSB
, blas_int const* M, blas_int const* N, blas_int const* K
, double const* ALPHA
, double const* A, blas_int const* LDA
, double const* B, blas_int const* LDB
, double const* BETA
, double* C, blas_int const* LDC );

void cgemm_
( char const* TRANSA, char const* TRANSB
, blas_int const* M, blas_int const* N, blas_int const* K
, std::complex<float> const* ALPHA
, std::complex<float> const* A, blas_int const* LDA
, std::complex<float> const* B, blas_int const* LDB
, std::complex<float> const* BETA
, std::complex<float>* C, blas_int const* LDC );

void zgemm_
( char const* TRANSA, char const* TRANSB
, blas_int const* M, blas_int const* N, blas_int const* K
, std::complex<double> const* ALPHA
, std::complex<double> const* A, blas_int const* LDA
, std::complex<double> const* B, blas_int const* LDB
, std::complex<double> const* BETA
, std::complex<double>* C, blas_int const* LDC );

} // extern "C"

REGISTER_BLAS(gemm, sgemm_, dgemm_, cgemm_, zgemm_)

template< class _T
        , class = std::enable_if_t
                  <   std::is_same<_T, float>()
                   or std::is_same<_T, double>()
                   or std::is_same<_T, std::complex<float>>()
                   or std::is_same<_T, std::complex<double>>()
                  >
        >
inline
auto gemm( Operator const op_A, Operator const op_B
         , std::size_t const m, std::size_t const n, std::size_t const k
         , _T const ALPHA, _T const* A, blas_int const LDA
         , _T const* B, blas_int const LDB
         , _T const BETA, _T* C, blas_int const LDC ) -> void
{
	if(m == 0 or n == 0 or k == 0) return;

	assert(A != nullptr and LDA != 0);
	assert(B != nullptr and LDB != 0);
	assert(C != nullptr and LDC != 0);
	/*
	utils::assert_valid(A, A + LDA * (op_A == Operator::None ? k : m));
	utils::assert_valid(B, B + LDB * (op_B == Operator::None ? n : k));
	utils::assert_valid(C, C + LDC * n);
	*/

	auto const M      = boost::numeric_cast<blas_int>(m);
	auto const N      = boost::numeric_cast<blas_int>(n);
	auto const K      = boost::numeric_cast<blas_int>(k);
	auto const TRANSA = static_cast<char>(op_A);
	auto const TRANSB = static_cast<char>(op_B);

	assert( LDA >= (op_A == Operator::None ? M : K) );
	assert( LDB >= (op_B == Operator::None ? K : N) );
	assert( LDC >= M );

	gemm<_T>( &TRANSA, &TRANSB, &M, &N, &K
	        , &ALPHA, A, &LDA, B, &LDB
	        , &BETA, C, &LDC );

	/*
	utils::assert_valid(C, C + LDC * n);
	*/
}

} // namespace import

} // namespace tcm

#endif // TCM_BLAS_WRAPPER_HPP


#include <cstdio>
#include <cmath>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <map>
#include <cfloat>

// Boost
#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>

// Eigen
#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>

// HDF5
#include <H5Cpp.h>

// OpenMP
#include <omp.h>

// Point-triangle distance and ray-triangle intersection.
#include "triangle_point/poitri.h"
#include "triangle_ray/raytri.h"
#include "box_triangle/aabb_triangle_overlap.h"

/** \brief Compute triangle point distance and corresponding closest point.
 * \param[in] point point
 * \param[in] v1 first vertex
 * \param[in] v2 second vertex
 * \param[in] v3 third vertex
 * \param[out] ray corresponding closest point
 * \return distance
 */
float triangle_point_distance(const Eigen::Vector3f &point, const Eigen::Vector3f &v1, const Eigen::Vector3f &v2, const Eigen::Vector3f &v3,
    Eigen::Vector3f &closest_point) {

  Vec3f x0(point.data());
  Vec3f x1(v1.data());
  Vec3f x2(v2.data());
  Vec3f x3(v3.data());

  Vec3f r(0);
  float distance = point_triangle_distance(x0, x1, x2, x3, r);

  for (int d = 0; d < 3; d++) {
    closest_point(d) = r[d];
  }

  return distance;
}

/** \brief Test triangle ray intersection.
 * \param[in] origin origin of ray
 * \param[in] dest destination of ray
 * \param[in] v1 first vertex
 * \param[in] v2 second vertex
 * \param[in] v3 third vertex
 * \return intersects
 */
bool triangle_ray_intersection(const Eigen::Vector3f &origin, const Eigen::Vector3f &dest,
    const Eigen::Vector3f &v1, const Eigen::Vector3f &v2, const Eigen::Vector3f &v3, float &t) {

  double _origin[3] = {origin(0), origin(1), origin(2)};
  double _dir[3] = {dest(0) - origin(0), dest(1) - origin(1), dest(2) - origin(2)};
  double _v1[3] = {v1(0), v1(1), v1(2)};
  double _v2[3] = {v2(0), v2(1), v2(2)};
  double _v3[3] = {v3(0), v3(1), v3(2)};

  // t is the distance, u and v are barycentric coordinates
  // http://fileadmin.cs.lth.se/cs/personal/tomas_akenine-moller/code/raytri_tam.pdf
  double _t, u, v;
  int success = intersect_triangle(_origin, _dir, _v1, _v2, _v3, &_t, &u, &v);
  t = _t;

  if (success) {
    return true;
  }

  return false;
}

/** \brief Compute triangle box intersection.
 * \param[in] min defining voxel
 * \param[in] max defining voxel
 * \param[in] v1 first vertex
 * \param[in] v2 second vertex
 * \param[in] v3 third vertex
 * \return intersects
 */
bool triangle_box_intersection(const Eigen::Vector3f &min, Eigen::Vector3f &max, const Eigen::Vector3f &v1, const Eigen::Vector3f &v2, const Eigen::Vector3f &v3) {
  float half_size[3] = {
    (max(0) - min(0))/2.,
    (max(1) - min(1))/2.,
    (max(2) - min(2))/2.
  };

  float center[3] = {
    max(0) - half_size[0],
    max(1) - half_size[1],
    max(2) - half_size[2]
  };

  float vertices[3][3] = {{v1(0), v1(1), v1(2)}, {v2(0), v2(1), v2(2)}, {v3(0), v3(1), v3(2)}};
  return triBoxOverlap(center, half_size, vertices);
}

/** \brief Specifies the voxelization mode, i.e. which point of a voxel to use for SDF computation. */
enum VoxelizationMode {
  CENTER = 0,
  CORNER = 1
};

/** \brief Just encapsulating vertices and faces. */
class Mesh {
public:
  /** \brief Empty constructor. */
  Mesh() {

  }

  /** \brief Reading an off file and returning the vertices x, y, z coordinates and the
   * face indices.
   * \param[in] filepath path to the OFF file
   * \param[out] mesh read mesh with vertices and faces
   * \return success
   */
  static bool from_off(const std::string filepath, Mesh& mesh) {

    std::ifstream* file = new std::ifstream(filepath.c_str());
    std::string line;
    std::stringstream ss;
    int line_nb = 0;

    std::getline(*file, line);
    ++line_nb;

    if (line != "off" && line != "OFF") {
      std::cout << "[Error] Invalid header: \"" << line << "\", " << filepath << std::endl;
      return false;
    }

    size_t n_edges;
    std::getline(*file, line);
    ++line_nb;

    int n_vertices;
    int n_faces;
    ss << line;
    ss >> n_vertices;
    ss >> n_faces;
    ss >> n_edges;

    for (size_t v = 0; v < n_vertices; ++v) {
      std::getline(*file, line);
      ++line_nb;

      ss.clear();
      ss.str("");

      Eigen::Vector3f vertex;
      ss << line;
      ss >> vertex(0);
      ss >> vertex(1);
      ss >> vertex(2);

      mesh.add_vertex(vertex);
    }

    size_t n;
    for (size_t f = 0; f < n_faces; ++f) {
      std::getline(*file, line);
      ++line_nb;

      ss.clear();
      ss.str("");

      size_t n;
      ss << line;
      ss >> n;

      if(n != 3) {
        std::cout << "[Error] Not a triangle (" << n << " points) at " << (line_nb - 1) << std::endl;
        return false;
      }

      Eigen::Vector3i face;
      ss >> face(0);
      ss >> face(1);
      ss >> face(2);

      mesh.add_face(face);
    }

    if (n_vertices != mesh.num_vertices()) {
      std::cout << "[Error] Number of vertices in header differs from actual number of vertices." << std::endl;
      return false;
    }

    if (n_faces != mesh.num_faces()) {
      std::cout << "[Error] Number of faces in header differs from actual number of faces." << std::endl;
      return false;
    }

    file->close();
    delete file;

    return true;
  }

  /** \brief Write mesh to OFF file.
   * \param[in] filepath path to OFF file to write
   * \return success
   */
  bool to_off(const std::string filepath) {
    std::ofstream* out = new std::ofstream(filepath, std::ofstream::out);
    if (!static_cast<bool>(out)) {
      return false;
    }

    (*out) << "OFF" << std::endl;
    (*out) << this->num_vertices() << " " << this->num_faces() << " 0" << std::endl;

    for (unsigned int v = 0; v < this->num_vertices(); v++) {
      (*out) << this->vertices[v](0) << " " << this->vertices[v](1) << " " << this->vertices[v](2) << std::endl;
    }

    for (unsigned int f = 0; f < this->num_faces(); f++) {
      (*out) << "3 " << this->faces[f](0) << " " << this->faces[f](1) << " " << this->faces[f](2) << std::endl;
    }

    out->close();
    delete out;

    return true;
  }

  /** \brief Add a vertex.
   * \param[in] vertex vertex to add
   */
  void add_vertex(Eigen::Vector3f& vertex) {
    this->vertices.push_back(vertex);
  }

  /** \brief Get the number of vertices.
   * \return number of vertices
   */
  int num_vertices() {
    return static_cast<int>(this->vertices.size());
  }

  /** \brief Add a face.
   * \param[in] face face to add
   */
  void add_face(Eigen::Vector3i& face) {
    this->faces.push_back(face);
  }

  /** \brief Get the number of faces.
   * \return number of faces
   */
  int num_faces() {
    return static_cast<int>(this->faces.size());
  }

  /** \brief Translate the mesh.
   * \param[in] translation translation vector
   */
  void translate(const Eigen::Vector3f& translation) {
    for (int v = 0; v < this->num_vertices(); ++v) {
      for (int i = 0; i < 3; ++i) {
        this->vertices[v](i) += translation(i);
      }
    }
  }

  /** \brief Scale the mesh.
   * \param[in] scale scale vector
   */
  void scale(const Eigen::Vector3f& scale) {
    for (int v = 0; v < this->num_vertices(); ++v) {
      for (int i = 0; i < 3; ++i) {
        this->vertices[v](i) *= scale(i);
      }
    }
  }

  /** \brief Voxelize the given mesh into a SDF.
   * \param[out] sdf volume to fill with sdf values
   */
  void voxelize_sdf(Eigen::Tensor<float, 3, Eigen::RowMajor>& sdf, const VoxelizationMode &mode) {

    int height = sdf.dimension(0);
    int width = sdf.dimension(1);
    int depth = sdf.dimension(2);

    #pragma omp parallel
    {
      #pragma omp for
      for (int i = 0; i < height*width*depth; i++) {
        int d = i%depth;
        int w = (i/depth)%width;
        int h = (i/depth)/width;

        sdf(h, w, d) = FLT_MAX;

        // the box corresponding to this voxel
        Eigen::Vector3f min(w, h, d);
        Eigen::Vector3f max(w + 1, h + 1, d + 1);

        Eigen::Vector3f center(w + 0.5f, h + 0.5f, d + 0.5f);
        if (mode == VoxelizationMode::CORNER) {
          center = Eigen::Vector3f(w, h, d);
        }

        // count number of intersections.
        int num_intersect = 0;
        for (unsigned int f = 0; f < this->num_faces(); ++f) {

          Eigen::Vector3f v1 = this->vertices[this->faces[f](0)];
          Eigen::Vector3f v2 = this->vertices[this->faces[f](1)];
          Eigen::Vector3f v3 = this->vertices[this->faces[f](2)];

          Eigen::Vector3f closest_point;
          triangle_point_distance(center, v1, v2, v3, closest_point);
          float distance = (center - closest_point).norm();

          if (distance < sdf(h, w, d)) {
            sdf(h, w, d) = distance;
          }

          bool intersect = triangle_ray_intersection(center, Eigen::Vector3f(0, 0, 0), v1, v2, v3, distance);

          if (intersect && distance >= 0) {
            num_intersect++;
          }
        }

        if (num_intersect%2 == 1) {
          sdf(h, w, d) *= -1;
        }
      }
    }
  }

  /** \brief Voxelize the given mesh into an occupancy grid.
   * \param[out] occ volume to fill
   */
  void voxelize_occ(Eigen::Tensor<int, 3, Eigen::RowMajor>& occ, const VoxelizationMode &mode) {

    int height = occ.dimension(0);
    int width = occ.dimension(1);
    int depth = occ.dimension(2);

    #pragma omp parallel
    {
      #pragma omp for
      for (int i = 0; i < height*width*depth; i++) {
        int d = i%depth;
        int w = (i/depth)%width;
        int h = (i/depth)/width;

        Eigen::Vector3f min(w, h, d);
        Eigen::Vector3f max(w + 1, h + 1, d + 1);

        for (unsigned int f = 0; f < this->num_faces(); ++f) {

          Eigen::Vector3f v1 = this->vertices[this->faces[f](0)];
          Eigen::Vector3f v2 = this->vertices[this->faces[f](1)];
          Eigen::Vector3f v3 = this->vertices[this->faces[f](2)];

          bool overlap = triangle_box_intersection(min, max, v1, v2, v3);
          if (overlap) {
            occ(h, w, d) = 1;
            break;
          }
        }
      }
    }
  }

private:

  /** \brief Vertices as (x,y,z)-vectors. */
  std::vector<Eigen::Vector3f> vertices;

  /** \brief Faces as list of vertex indices. */
  std::vector<Eigen::Vector3i> faces;
};

/** \brief Write the given set of volumes to h5 file.
 * \param[in] filepath h5 file to write
 * \param[in] n number of volumes
 * \param[in] height height of volumes
 * \param[in] width width of volumes
 * \param[in] depth depth of volumes
 * \param[in] dense volume data
 */
template<int RANK>
bool write_float_hdf5(const std::string filepath, Eigen::Tensor<float, RANK, Eigen::RowMajor>& tensor) {

  try {

    /*
     * Turn off the auto-printing when failure occurs so that we can
     * handle the errors appropriately
     */
    H5::Exception::dontPrint();

    /*
     * Create a new file using H5F_ACC_TRUNC access,
     * default file creation properties, and default file
     * access properties.
     */
    H5::H5File file(filepath, H5F_ACC_TRUNC);

    /*
     * Define the size of the array and create the data space for fixed
     * size dataset.
     */
    hsize_t rank = RANK;
    hsize_t dimsf[rank];
    for (int i = 0; i < rank; i++) {
      dimsf[i] = tensor.dimension(i);

    }
    H5::DataSpace dataspace(rank, dimsf);

    /*
     * Define datatype for the data in the file.
     * We will store little endian INT numbers.
     */
    H5::IntType datatype(H5::PredType::NATIVE_FLOAT);
    datatype.setOrder(H5T_ORDER_LE);

    /*
     * Create a new dataset within the file using defined dataspace and
     * datatype and default dataset creation properties.
     */
    H5::DataSet dataset = file.createDataSet("tensor", datatype, dataspace);

    /*
     * Write the data to the dataset using default memory space, file
     * space, and transfer properties.
     */
    float* data = static_cast<float*>(tensor.data());
    dataset.write(data, H5::PredType::NATIVE_FLOAT);
  }  // end of try block

  // catch failure caused by the H5File operations
  catch(H5::FileIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSet operations
  catch(H5::DataSetIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataSpaceIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataTypeIException error) {
    error.printError();
    return false;
  }

  return true;
}

/** \brief Write the given set of volumes to h5 file.
 * \param[in] filepath h5 file to write
 * \param[in] n number of volumes
 * \param[in] height height of volumes
 * \param[in] width width of volumes
 * \param[in] depth depth of volumes
 * \param[in] dense volume data
 */
template<int RANK>
bool write_int_hdf5(const std::string filepath, Eigen::Tensor<int, RANK, Eigen::RowMajor>& tensor) {

  try {

    /*
     * Turn off the auto-printing when failure occurs so that we can
     * handle the errors appropriately
     */
    H5::Exception::dontPrint();

    /*
     * Create a new file using H5F_ACC_TRUNC access,
     * default file creation properties, and default file
     * access properties.
     */
    H5::H5File file(filepath, H5F_ACC_TRUNC);

    /*
     * Define the size of the array and create the data space for fixed
     * size dataset.
     */
    hsize_t rank = RANK;
    hsize_t dimsf[rank];
    for (int i = 0; i < rank; i++) {
      dimsf[i] = tensor.dimension(i);

    }
    H5::DataSpace dataspace(rank, dimsf);

    /*
     * Define datatype for the data in the file.
     * We will store little endian INT numbers.
     */
    H5::IntType datatype(H5::PredType::NATIVE_INT);
    datatype.setOrder(H5T_ORDER_LE);

    /*
     * Create a new dataset within the file using defined dataspace and
     * datatype and default dataset creation properties.
     */
    H5::DataSet dataset = file.createDataSet("tensor", datatype, dataspace);

    /*
     * Write the data to the dataset using default memory space, file
     * space, and transfer properties.
     */
    int* data = static_cast<int*>(tensor.data());
    dataset.write(data, H5::PredType::NATIVE_INT);
  }  // end of try block

  // catch failure caused by the H5File operations
  catch(H5::FileIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSet operations
  catch(H5::DataSetIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataSpaceIException error) {
    error.printError();
    return false;
  }

  // catch failure caused by the DataSpace operations
  catch(H5::DataTypeIException error) {
    error.printError();
    return false;
  }

  return true;
}

/** \brief Read all files in a directory matching the given extension.
 * \param[in] directory path to directory
 * \param[out] files read file paths
 * \param[in] extension extension to filter for
 */
void read_directory(const boost::filesystem::path directory, std::map<int, boost::filesystem::path>& files, const std::string extension = ".off") {

  files.clear();
  boost::filesystem::directory_iterator end;

  for (boost::filesystem::directory_iterator it(directory); it != end; ++it) {
    if (it->path().extension().string() == extension) {
      int number = std::stoi(it->path().filename().string());
      files.insert(std::pair<int, boost::filesystem::path>(number, it->path()));
    }
  }
}

/** \brief Main entrance point of the script.
 * Expects one parameter, the path to the corresponding config file in config/.
 */
int main(int argc, char** argv) {
  boost::program_options::options_description desc("Allowed options");
  desc.add_options()
      ("help", "produce help message")
      ("mode", boost::program_options::value<std::string>()->default_value("occ"), "operation mode, 'occ' or 'sdf'")
      ("input", boost::program_options::value<std::string>(), "input, either single OFF file or directory containing OFF files where the names correspond to integers (zero padding allowed) and are consecutively numbered starting with zero")
      ("height", boost::program_options::value<int>()->default_value(32), "height of volume, corresponding to y-axis (=up)")
      ("width", boost::program_options::value<int>()->default_value(32), "width of volume, corresponding to x-axis (=right")
      ("depth", boost::program_options::value<int>()->default_value(32), "depth of volume, corresponding to z-axis (=forward)")
      ("center", boost::program_options::bool_switch()->default_value(false), "by default, the top-left-front corner is used for SDF computation; if instead the voxel centers should be used, set this flag")
      ("output", boost::program_options::value<std::string>(), "output file, will be a HDF5 file containing either a N x C x height x width x depth tensor or a C x height x width x depth tensor, where N is the number of files and C=2 the number of channels, N is discarded if only a single file is processed; should have the .h5 extension");

  boost::program_options::positional_options_description positionals;
  positionals.add("mode", 1);
  positionals.add("input", 1);
  positionals.add("output", 1);

  boost::program_options::variables_map parameters;
  boost::program_options::store(boost::program_options::command_line_parser(argc, argv).options(desc).positional(positionals).run(), parameters);
  boost::program_options::notify(parameters);

  if (parameters.find("help") != parameters.end()) {
    std::cout << desc << std::endl;
    return 0;
  }

  std::string mode = parameters["mode"].as<std::string>();
  if (mode == "occ") {
    std::cout << "Voxelizing occupancy grids." << std::endl;
  }
  else if (mode == "sdf") {
    std::cout << "Voxelizing SDFs." << std::endl;
  }
  else {
    std::cout << "Invalid mode, choose from occ or sdf." << std::endl;
    return 1;
  }

  boost::filesystem::path input(parameters["input"].as<std::string>());
  if (!boost::filesystem::is_directory(input) && !boost::filesystem::is_regular_file(input)) {
    std::cout << "Input is neither directory nor file." << std::endl;
    return 1;
  }

  boost::filesystem::path output(parameters["output"].as<std::string>());
  if (boost::filesystem::is_regular_file(output)) {
    std::cout << "Output file already exists; overwriting." << std::endl;
  }

  VoxelizationMode voxelization_mode;
  if (parameters["center"].as<bool>()) {
    voxelization_mode = VoxelizationMode::CENTER;
    std::cout << "Using the top-left-front voxel corner for voxelization." << std::endl;
  }
  else {
    voxelization_mode = VoxelizationMode::CORNER;
    std::cout << "Using the voxel center for voxelization." << std::endl;
  }

  int height = parameters["height"].as<int>();
  int width = parameters["width"].as<int>();
  int depth = parameters["depth"].as<int>();

  std::cout << "Voxelizing into " << height << " x " << width << " x " << depth << " (height x width x depth)." << std::endl;

  if (boost::filesystem::is_regular_file(input)) {
    Mesh mesh;
    bool success = Mesh::from_off(input.string(), mesh);

    if (!success) {
      std::cout << "Could not read " << input << "." << std::endl;
      return 1;
    }

    std::cout << "Read " << input << "." << std::endl;

    if (mode == "sdf") {
      Eigen::Tensor<float, 3, Eigen::RowMajor> tensor(height, width, depth);

      mesh.voxelize_sdf(tensor, voxelization_mode);
      std::cout << "Voxelized " << input << "." << std::endl;

      bool success = write_float_hdf5<3>(output.string(), tensor);

      if (!success) {
        std::cout << "Could not write " << output << "." << std::endl;
        return 1;
      }
    }
    if (mode == "occ") {
      Eigen::Tensor<int, 3, Eigen::RowMajor> tensor(height, width, depth);
      tensor.setZero();

      mesh.voxelize_occ(tensor, voxelization_mode);
      std::cout << "Voxelized " << input << "." << std::endl;

      bool success = write_int_hdf5<3>(output.string(), tensor);

      if (!success) {
        std::cout << "Could not write " << output << "." << std::endl;
        return 1;
      }
    }

    std::cout << "Wrote " << output << "." << std::endl;
    std::cout << "The output is a " << height << " x " << width << " x " << depth << " tensor." << std::endl;
  }
  else {
    std::map<int, boost::filesystem::path> input_files;
    read_directory(input, input_files);

    if (input_files.size() <= 0) {
      std::cout << "Could not find any OFF files in the input directory." << std::endl;
      return 1;
    }

    std::cout << "Read " << input_files.size() << " files." << std::endl;

    if (mode == "sdf") {
      Eigen::Tensor<float, 4, Eigen::RowMajor> tensor(input_files.size(), height, width, depth);

      int i = 0;
      for (std::map<int, boost::filesystem::path>::iterator it = input_files.begin(); it != input_files.end(); it++) {
        Mesh mesh;
        bool success = Mesh::from_off(it->second.string(), mesh);

        if (!success) {
          std::cout << "Could not read " << it->second << "." << std::endl;
          return 1;
        }

        Eigen::Tensor<float, 3, Eigen::RowMajor> slice(height, width, depth);
        mesh.voxelize_sdf(slice, voxelization_mode);
        tensor.chip(i, 0) = slice;
        std::cout << "Voxelized " << it->second << " (" << (i + 1) << " of " << input_files.size() << ")." << std::endl;

        i++;
      }

      bool success = write_float_hdf5<4>(output.string(), tensor);

      if (!success) {
        std::cout << "Could not write " << output << "." << std::endl;
        return 1;
      }
    }
    if (mode == "occ") {
      Eigen::Tensor<int, 4, Eigen::RowMajor> tensor(input_files.size(), height, width, depth);
      tensor.setZero();

      int i = 0;
      for (std::map<int, boost::filesystem::path>::iterator it = input_files.begin(); it != input_files.end(); it++) {
        Mesh mesh;
        bool success = Mesh::from_off(it->second.string(), mesh);

        if (!success) {
          std::cout << "Could not read " << it->second << "." << std::endl;
          return 1;
        }

        Eigen::Tensor<int, 3, Eigen::RowMajor> slice(height, width, depth);
        slice.setZero();

        mesh.voxelize_occ(slice, voxelization_mode);
        tensor.chip(i, 0) = slice;
        std::cout << "Voxelized " << it->second << " (" << (i + 1) << " of " << input_files.size() << ")." << std::endl;

        i++;
      }

      bool success = write_int_hdf5<4>(output.string(), tensor);

      if (!success) {
        std::cout << "Could not write " << output << "." << std::endl;
        return 1;
      }
    }

    std::cout << "Wrote " << output << "." << std::endl;
    std::cout << "The output is a " << input_files.size() << " x " << height << " x " << width << " x " << depth << " tensor." << std::endl;
  }

  return 0;
}