/**
 * \file internal/creating_tamd_particles.h
 * \brief creating TAMD chains internal methods
 *
 *  Copyright 2007-2021 IMP Inventors. All rights reserved.
 */

#ifndef IMPNPCTRANSPORT_INTERNAL_TAMD_CHAIN_H
#define IMPNPCTRANSPORT_INTERNAL_TAMD_CHAIN_H

#include "../npctransport_config.h"
#include <IMP/npctransport/ParticleFactory.h>
#include <IMP/npctransport/FGChain.h>
#include <IMP/atom/Hierarchy.h>
#include <IMP/core/DistancePairScore.h>
#include <string>
#include <vector>

IMPNPCTRANSPORT_BEGIN_INTERNAL_NAMESPACE

/**
    A hierarchy of centroids to represent a polymer chain with
    centroids in the non-leaf nodes, singleton particles in the
    leaves, and TAMD images attached to the centroids by spring.

    Due to its complexity, it cannot be constructed directly,
    but through factory methods

    \see create_tamd_chain()
    \see create_singleton_tamd_chain()

*/
class TAMDChain : public npctransport::FGChain{
  /*
    TAMD hierarchy:
    root - root of hierarchy
    beads - fine chain particles
    centroids - the list of centroids in tree(p)
    images - corresponding TAMD images for each centroid in centroids
    tamd_restraints - corresponding springs that attch each TAMD
                       image to each centroid
    tamd_springs - springs stored in R (in same order)
  */
  typedef npctransport::FGChain P;
 private:
  // TODO: keep those in the hierarchy?
  IMP::Particles centroids_;
  IMP::Particles images_;
  IMP::Restraints tamd_restraints_;
  core::HarmonicDistancePairScores tamd_springs_;
  bool is_initialized;

 private:
  /** Initialize a TAMD Chain that still doesn't have a root or beads,
      with backbone k and rest length factor as specified, to be utilized
      once an actual root with bead leafs is added. The user should not
      use this constructor - it is expected that the chain will be generated
      only via a factor method, causing is_initialized to become true

      \see create_tamd_chain()
  */
  TAMDChain(
           double backbone_k = 0.0,
           double rest_length_factor = 1.0,
           std::string name = "tamd_chain %1%")
    : P(nullptr, backbone_k, rest_length_factor, name),
    is_initialized(false)
    {}

  /**
     - Add d children and n beads generated by pf to chain root;
     - T_factors, F_factors and Ks are TAMD params as in
       create_tamd_chain()
     - Set root as centroid of its children (with updated list);
     - Assumes root is not null
  */
  void add_children
    ( ParticleFactory* pf,
      unsigned int n,
      unsigned int d,
      std::vector<double> T_factors,
      std::vector<double> F_factors,
      std::vector<double> Ks );

 public:
  /**
      return the restraint accosciated with internal interactions by this chain.
      In particular for a TAMD chain, this includes both the bonds between beads,
      and the TAMD spring restraints.

      \see FGChain::get_chain_restraints
  */
  virtual Restraints get_chain_restraints
    (Scoring const* scoring_manager) IMP_OVERRIDE;

  //! get TAMD images of centroids in the hierarchy
  /**
     Get TAMD images of centroids in the hierarchy, in corresponding order to centroids in
     get_tamd_centroids()

     @note In future extensions, it is expected to included images of any collective
           variables in the simulation, not only centroids)

      \see get_tamd_centroids()
  */
  Particles get_tamd_images() const
    { return images_; }

  //! get particles that are centroids in the hierarchy
  /** Get particles that are centroids in the hierarchy, in the same order as
      their corresponding images in get_tamd_images()

      \see get_tamd_images()
  */
  Particles get_tamd_centroids() const
    { return centroids_; }

  /** returns a list of all tamd springs associated with
      the tamd restraints (between a diffusing tamd particle and its
      associated collective variable, eg centroid)
  */
  core::HarmonicDistancePairScores& get_tamd_springs_byref() {
    return tamd_springs_;
  }


  IMP_OBJECT_METHODS(TAMDChain);

 /******************* factory methods are friends ********************/

  friend TAMDChain*
    create_tamd_chain( ParticleFactory* pf,
                       unsigned int nlevels,
                       unsigned int d,
                       std::vector<double> T_factors,
                       std::vector<double> F_factors,
                       std::vector<double> Ks );

  friend TAMDChain*
    create_singleton_tamd_chain( ParticleFactory* pf);

};


/************* factory methods ***************/

/**
     Create a TAMD hierarchy of nlevels depth with a core for
     each d centroids in a lower level, with a real centroid and
     restrained centroid realization

     @param pf      A factory for producing singleton particles
                   (the leaves of the chain)
     @param nlevels Number of tamd levels in the hierarchy. If 0 then return a
                    singleton particle.
     @param d       The out degree of each non-leaf node (# of children)
     @param T_factors A list of length nlevels with temperature at each level
                    from top to bottom
     @param G_factors A list of length nlevels with friction factor (G for gamma)
                    at each level from top to bottom
     @param Ks      Spring constants at each level between a particle and its TAMD
                    image (a list of length nlevels)

     @return a tuple with <root particle, centroids, images, restraints>
*/
TAMDChain*
create_tamd_chain( ParticleFactory* pf,
                   unsigned int nlevels,
                   unsigned int d,
                   std::vector<double> T_factors,
                   std::vector<double> F_factors,
                   std::vector<double> Ks );


//! create a degenerate TAMDChain with a single bead (that is also the root),
//! which is produced by pf
TAMDChain* create_singleton_tamd_chain( ParticleFactory* pf);



IMPNPCTRANSPORT_END_INTERNAL_NAMESPACE


#endif /* IMPNPCTRANSPORT_INTERNAL_TAMD_CHAIN_H */
