B n_mutations(const alphabet& a, const vector<int>& letters, const SequenceTree& T,const ublas::matrix<B>& cost,
	      ublas::matrix<B>& n_muts, const vector<const_branchview>& branches)
{
  int root = T.directed_branch(0).target();

  peel_n_mutations(a,letters,T,cost,n_muts,branches);

  return row_min(n_muts,root);
}

/// Does any branch in T imply the partition p?
bool implies(const SequenceTree& T,const Partition& p) {
  bool result = false;
  for(int b=0;b<T.n_branches() and not result;b++) {
    dynamic_bitset<> bp = branch_partition(T,b);

    if (implies(bp,p)) return true;
  }
  return false;
}

vector<Partition> all_partitions_from_tree(const SequenceTree& T) 
{
  vector<Partition> partitions;

  for(int b=0;b<T.n_branches();b++)
    partitions.push_back(partition_from_branch(T,b));

  return partitions;
}

void remap_T_indices(SequenceTree& T,const vector<string>& names)
{
  //----- Remap leaf indices for T onto A's leaf sequence indices -----//
  try {
    vector<int> mapping = compute_mapping(T.get_sequences(),names);

    T.standardize(mapping);
  }
  catch(const bad_mapping<string>& b)
  {
    bad_mapping<string> b2(b.missing,b.from);
    if (b.from == 0)
      b2<<"Couldn't find leaf sequence \""<<b2.missing<<"\" in names.";
    else
      b2<<"Sequence '"<<b2.missing<<"' not found in the tree.";
    throw b2;
  }
}

int which_branch(const SequenceTree& T, const Partition& p) 
{
  for(int b=0; b<2*T.n_branches(); b++) {
    dynamic_bitset<> bp = branch_partition(T,b);
    if( directed_implies(bp,p) )
      return b;
  }
  return -1;
}

void write_partitions(std::ostream& o,const vector<Partition>& partitions)
{
  vector<Partition> full;
  vector<Partition> sub;
  for(int i=0;i<partitions.size();i++)
    if (partitions[i].full())
      full.push_back(partitions[i]);
    else
      sub.push_back(partitions[i]);

  if (full.size()) {
    SequenceTree consensus = get_mf_tree(partitions[0].names,full);
    o<<consensus.write(false)<<endl;
  }

  for(int i=0;i<sub.size();i++)
    o<<sub[i]<<endl;
}

/// \brief  Remap the leaf indices of tree \a T to match the alignment \a A: check the result
///
/// \param A The alignment.
/// \param T The tree.
/// \param internal_sequences Should the resulting alignment have sequences for internal nodes on the tree?
///
void link(alignment& A,SequenceTree& T,bool internal_sequences) 
{
  check_names_unique(A);

  // Later, might we WANT sub-branches???
  if (has_sub_branches(T))
    remove_sub_branches(T);

  if (internal_sequences and not is_Cayley(T)) {
    assert(has_polytomy(T));
    throw myexception()<<"Cannot link a multifurcating tree to an alignment with internal sequences.";
  }

  //------ IF sequences < leaf nodes THEN complain ---------//
  if (A.n_sequences() < T.n_leaves())
    throw myexception()<<"Tree has "<<T.n_leaves()<<" leaves but Alignment only has "
		       <<A.n_sequences()<<" sequences.";

  //----- IF sequences = leaf nodes THEN maybe add internal sequences.
  else if (A.n_sequences() == T.n_leaves()) {
    if (internal_sequences)
      A = add_internal(A,T);
  }
  //----- IF sequences > leaf nodes THEN maybe complain -------//
  else {
    if (not internal_sequences) {
      alignment A2 = chop_internal(A);
      if (A2.n_sequences() == T.n_leaves()) {
	A = A2;
      }
      else
	throw myexception()<<"More alignment sequences than leaf nodes!";
    } 
    else if (A.n_sequences() > T.n_nodes())
      throw myexception()<<"More alignment sequences than tree nodes!";
    else if (A.n_sequences() < T.n_nodes())
      throw myexception()<<"Fewer alignment sequences than tree nodes!";
  }
  
  //---------- double-check that we have the right number of sequences ---------//
  if (internal_sequences)
    assert(A.n_sequences() == T.n_nodes());
  else
    assert(A.n_sequences() == T.n_leaves());

  //----- Remap leaf indices for T onto A's leaf sequence indices -----//
  remap_T_indices(T,A);

  if (internal_sequences)
    connect_leaf_characters(A,T);

  //---- Check to see that internal nodes satisfy constraints ----//
  check_alignment(A,T,internal_sequences);
}

/// \brief Re-index the leaves of tree \a T1 so that the labels have the same ordering as in \a T2.
///
/// \param T1 The leaf-labelled tree to re-index.
/// \param T2 The leaf-labelled tree to match.
///
void remap_T_indices(SequenceTree& T1,const SequenceTree& T2)
{
  if (T1.n_leaves() != T2.n_leaves())
    throw myexception()<<"Trees do not correspond: different numbers of leaves.";

  //----- Remap leaf indices for T onto A's leaf sequence indices -----//
  try {
    remap_T_indices(T1,T2.get_sequences());
  }
  catch(const bad_mapping<string>& b)
  {
    bad_mapping<string> b2(b.missing,b.from);
    if (b.from == 0)
      b2<<"Couldn't find leaf sequence \""<<b2.missing<<"\" in second tree.";
    else
      b2<<"Couldn't find leaf sequence \""<<b2.missing<<"\" in first tree.";
    throw b2;
  }
}

SequenceTree get_mf_tree(const std::vector<std::string>& names,
			 const std::vector<Partition>& partitions) 
{
  SequenceTree T = star_tree(names);

  int i=0;
  try {
    for(;i<partitions.size();i++)
      T.induce_partition(partitions[i].group1);
  }
  catch(...) {
    throw myexception()<<"Partition ("<<partitions[i]<<") conflicts with tree "<<T;
  }

  for(int i=0;i<T.n_branches();i++)
    T.branch(i).set_length(1.0);

  return T;
}

int choose_SPR_target(SequenceTree& T1, int b1_) 
{
  const_branchview b1 = T1.directed_branch(b1_);

  //----- Select the branch to move to ------//
  dynamic_bitset<> subtree_nodes = T1.partition(b1.reverse());
  subtree_nodes[b1.target()] = true;

  vector<int> branches;
  vector<double> lengths;

  for(int i=0;i<T1.n_branches();i++) 
  {
    const_branchview bi = T1.branch(i);

    // skip branch if its contained in the subtree
    if (subtree_nodes[bi.target()] and 
	subtree_nodes[bi.source()])
      continue;

    double L = 1.0;

    // down-weight branch if it is one of the subtree's 2 neighbors
    if (subtree_nodes[bi.target()] or 
	subtree_nodes[bi.source()])
      L = 0.5;

    branches.push_back(i);
    lengths.push_back(L);
  }

  try {
    int b2 = branches[ choose(lengths) ];

    return b2;
  }
  catch (choose_exception<efloat_t>& c)
  {
    c.prepend(__PRETTY_FUNCTION__);
    throw c;
  }
}

/// \brief Re-index the leaves of tree \a T so that the labels have the same ordering as in \a names.
///
/// \param T The leaf-labelled tree.
/// \param names The ordered leaf labels.
///
void remap_T_leaf_indices(SequenceTree& T,const vector<string>& names)
{
  assert(names.size() == T.n_leaves());
  //----- Remap leaf indices for T onto A's leaf sequence indices -----//
  try {
    vector<int> mapping = compute_mapping(T.get_leaf_labels(), names);

    T.standardize(mapping);
  }
  catch(const bad_mapping<string>& b)
  {
    bad_mapping<string> b2 = b;
    b2.clear();
    if (b2.from == 0)
      b2<<"Couldn't find leaf sequence \""<<b2.missing<<"\" in names.";
    else
      b2<<"Sequence '"<<b2.missing<<"' not found in the tree.";
    throw b2;
  }
}

data_partition::data_partition(const string& n, const alignment& a,const SequenceTree& t,
			       const substitution::MultiModel& SM)
  :SModel_(SM),
   partition_name(n),
   cached_alignment_prior_for_branch(t.n_branches()),
   cached_alignment_counts_for_branch(t.n_branches(),ublas::matrix<int>(5,5)),
   cached_sequence_lengths(a.n_sequences()),
   branch_mean_(1.0),
   smodel_full_tree(true),
   A(a),
   T(t),
   MC(t,SM),
   LC(t,SModel()),
   branch_HMMs(t.n_branches()),
   branch_HMM_type(t.n_branches(),0),
   beta(2, 1.0)
{
  for(int b=0;b<cached_alignment_counts_for_branch.size();b++)
    cached_alignment_counts_for_branch[b].invalidate();
}

/// \brief Re-index the leaves of tree \a T so that the labels have the same ordering as in \a A.
///
/// \param T The leaf-labelled tree.
/// \param A A multiple sequence alignment.
///
alignment remap_A_indices(alignment& A, const SequenceTree& T)
{
  vector<string> labels = T.get_labels();

  if (A.n_sequences() == T.n_leaves())
  {
    labels.resize(T.n_leaves());

  }
  else if (A.n_sequences() != T.n_nodes())
    throw myexception()<<"Cannot map alignment onto tree:\n  Alignment has "<<A.n_sequences()<<" sequences.\n  Tree has "<<T.n_leaves()<<" leaves and "<<T.n_nodes()<<" nodes.";
      

  for(int i=0;i<labels.size();i++)
    if (labels[i] == "")
    {
      if (i<T.n_leaves())
	throw myexception()<<"Tree has empty label for a leaf node: not allowed!";
      else
	throw myexception()<<"Alignment has internal node information, but tree has empty label for an internal node: not allowed!";
    }

  assert(A.n_sequences() == labels.size());

  //----- Remap leaf indices for T onto A's leaf sequence indices -----//
  try {
    vector<int> mapping = compute_mapping(labels, sequence_names(A));

    return reorder_sequences(A,mapping);
  }
  catch(const bad_mapping<string>& b)
  {
    bad_mapping<string> b2 = b;
    b2.clear();
    if (b.from == 0)
      b2<<"Couldn't find sequence \""<<b2.missing<<"\" in alignment.";
    else
      b2<<"Alignment sequence '"<<b2.missing<<"' not found in the tree.";
    throw b2;
  }
}

bool update_lengths(const SequenceTree& Q,const SequenceTree& T,
		    valarray<double>& branch_lengths, 
		    valarray<double>& branch_lengths_squared, 
		    valarray<double>& node_lengths)
{
  // map branches from Q -> T
  vector<int> branches_map = extends_map(T,Q);
  if (not branches_map.size())
    return false;

  // incorporate lengths of branches that map to Q
  for(int b=0;b<Q.n_branches();b++)
  {
    int b2 = branches_map[b];
    double L = T.directed_branch(b2).length();
    branch_lengths[b] += L;
    branch_lengths_squared[b] += L*L;
  }

  // map nodes from T -> Q
  vector<int> nodes_map = get_nodes_map(Q,T,branches_map);

  // incorprate lengths of branches that map to nodes in Q
  for(int i=T.n_leafbranches();i<T.n_branches();i++) 
  {
    const_branchview b = T.branch(i);
    int n1 = nodes_map[b.source()];
    int n2 = nodes_map[b.target()];

    if (n1 == n2)
      node_lengths[n1] += T.branch(i).length();
  }

  return true;
}

void load_As_and_random_T(const variables_map& args,vector<alignment>& alignments,SequenceTree& T,const vector<bool>& internal_sequences)
{
  //align - filenames
  vector<string> filenames = args["align"].as<vector<string> >();

  // load the alignments
  alignments = load_alignments(filenames,load_alphabets(args));

  //------------- Load random tree ------------------------//
  SequenceTree TC = star_tree(sequence_names(alignments[0]));
  if (args.count("t-constraint"))
    TC = load_constraint_tree(args["t-constraint"].as<string>(),sequence_names(alignments[0]));

  T = TC;
  RandomTree(T,1.0);

  //-------------- Link --------------------------------//
  link(alignments,T,internal_sequences);

  //---------------process----------------//
  for(int i=0;i<alignments.size();i++) 
  {
    
    //---------------- Randomize alignment? -----------------//
    if (args.count("randomize-alignment"))
      alignments[i] = randomize(alignments[i],T.n_leaves());
  
    //------------------ Analyze 'internal'------------------//
    if ((args.count("internal") and args["internal"].as<string>() == "+")
	or args.count("randomize-alignment"))
      for(int column=0;column< alignments[i].length();column++) {
	for(int j=T.n_leaves();j<alignments[i].n_sequences();j++) 
	  alignments[i](column,j) = alphabet::not_gap;
      }

    //---- Check that internal sequence satisfy constraints ----//
    check_alignment(alignments[i],T,internal_sequences[i]);
  }
}

/// Reorder internal sequences of \a A to correspond to standardized node names for \a T
alignment standardize(const alignment& A, const SequenceTree& T) 
{
  SequenceTree T2 = T;

  // if we don't have any internal node sequences, then we are already standardized
  if (A.n_sequences() == T.n_leaves())
    return A;

  // standardize NON-LEAF node and branch names in T
  vector<int> mapping = T2.standardize();
  vector<int> new_order = invert(mapping);

  return reorder_sequences(A,new_order);
}

int main(int argc,char* argv[]) { 
  Arguments args;
  args.read(argc,argv);

  unsigned long seed =0;
  if (args.set("seed")) {
    seed = convertTo<unsigned long>(args["seed"]);
    myrand_init(seed);
  }
  else
    seed = myrand_init();

  assert(args.set("names"));
  vector<string> names = split(args["names"],':');

  double branch_mean = 0.1;
  if (args.set("mean"))
    branch_mean = convertTo<double>(args["mean"]);

  SequenceTree T = RandomTree(names,branch_mean);

  std::cout<<T.write()<<std::endl;
}

vector<int> get_parsimony_letters(const alphabet& a, const vector<int>& letters, const SequenceTree& T,
				  const ublas::matrix<int>& cost)
{
  int root = T.directed_branch(0).target();
  ublas::matrix<int> n_muts(T.n_nodes(),a.size());

  peel_n_mutations(a,letters,T,cost,n_muts, branches_toward_node(T,root) );

  // get an order list of branches point away from the root;
  vector<const_branchview> branches = branches_from_node(T,root);
  std::reverse(branches.begin(),branches.end());
  
  // Allocate space to store the letter for each node
  vector<int> node_letters(T.n_nodes(),-1);

  // choose the cheapest letter at the root
  node_letters[root] = row_min(n_muts,root);

  const unsigned A = a.size();
  vector<double> temp(A);

  for(int i=0;i<branches.size();i++) 
  {
    int s = branches[i].source();
    int t = branches[i].target();

    int k = node_letters[s];
    assert(k != -1);

    for(int l=0;l<A;l++)
      temp[l] = n_muts(t,l)+cost(l,k);

    node_letters[t] = argmin(temp);
  }

  return node_letters;
}

string topology(const string& t) {
    SequenceTree T = standardized(t);
    return T.write(false);
}

Parameters::Parameters(const vector<alignment>& A, const SequenceTree& t,
		       const vector<polymorphic_cow_ptr<substitution::MultiModel> >& SMs,
		       const vector<int>& s_mapping,
		       const vector<int>& scale_mapping)
  :SModels(SMs),
   smodel_for_partition(s_mapping),
   scale_for_partition(scale_mapping),
   branch_prior_type(0),
   smodel_full_tree(true),
   T(t),
   TC(star_tree(t.get_sequences())),
   branch_HMM_type(t.n_branches(),0),
   beta(2, 1.0),
   updown(-1),
   features(0)
{
  constants.push_back(-1);

  for(int i=0;i<n_scales;i++)
    add_super_parameter("mu"+convertToString(i+1),1.0);

  // check that smodel mapping has correct size.
  if (smodel_for_partition.size() != A.size())
    throw myexception()<<"There are "<<A.size()
		       <<" data partitions, but you mapped smodels onto "
		       <<smodel_for_partition.size();

  // register the substitution models as sub-models
  for(int i=0;i<SModels.size();i++) {
    string name = "S" + convertToString(i+1);
    add_submodel(name, *SModels[i]);
  }

  // NO indel model (in this constructor)

  // check that we only mapping existing smodels to data partitions
  for(int i=0;i<smodel_for_partition.size();i++) {
    int m = smodel_for_partition[i];
    if (m >= SModels.size())
      throw myexception()<<"You can't use smodel "<<m+1<<" for data partition "<<i+1
			 <<" because there are only "<<SModels.size()<<" smodels.";
  }

  // load values from sub-models (smodels/imodel)
  read();

  // don't constrain any branch lengths
  for(int b=0;b<TC->n_branches();b++)
    TC->branch(b).set_length(-1);

  // create data partitions and register as sub-models
  for(int i=0;i<A.size();i++) 
  {
    // compute name for data-partition
    string name = string("part") + convertToString(i+1);

    // get reference to smodel for data-partition
    const substitution::MultiModel& SM = SModel(smodel_for_partition[i]);

    // create data partition
    data_partitions.push_back(cow_ptr<data_partition>(data_partition(name,A[i],*T,SM)));

    // register data partition as sub-model
    add_submodel(name,*data_partitions[i]);
  }
}