// Copyright (c) 2009-2016 The Regents of the University of Michigan
// This file is part of the HOOMD-blue project, released under the BSD 3-Clause License.


// Maintainer: jglaser

#include "IntegrationMethodTwoStep.h"
#include "hoomd/Variant.h"
#include "hoomd/ComputeThermo.h"

#ifndef __TWO_STEP_NPT_MTK_H__
#define __TWO_STEP_NPT_MTK_H__

/*! \file TwoStepNPTMTK.h
    \brief Declares the TwoStepNPTMTK class
*/

#ifdef NVCC
#error This header cannot be compiled by nvcc
#endif

//! Integrates part of the system forward in two steps in the NPT ensemble
/*! Implements the Martyna Tobias Klein (MTK) equations for rigorous integration in the NPT ensemble.
    The update equations are derived from a strictly measure-preserving and
    time-reversal symmetric integration scheme, closely following the one proposed by Tuckerman et al.

    Supports anisotropic (orthorhombic or tetragonal) integration modes, by implementing a special
    version of the the fully flexible cell update equations proposed in Yu et al.

    Triclinic integration for an upper triangular cell parameter matrix is supported with
    fully time-reversible and measure-preserving update equations (Glaser et al. 2013 to be published)

    \cite Martyna1994
    \cite Tuckerman2006
    \cite Yu2010
    \ingroup updaters
*/
class TwoStepNPTMTK : public IntegrationMethodTwoStep
    {
    public:
        //! Specify possible couplings between the diagonal elements of the pressure tensor
        enum couplingMode
            {
            couple_none = 0,
            couple_xy,
            couple_xz,
            couple_yz,
            couple_xyz};

        /*! Flags to indicate which degrees of freedom of the simulation box should be put under
            barostat control
         */
        enum baroFlags
            {
            baro_x = 1,
            baro_y = 2,
            baro_z = 4,
            baro_xy = 8,
            baro_xz = 16,
            baro_yz = 32
            };

        //! Constructs the integration method and associates it with the system
        TwoStepNPTMTK(boost::shared_ptr<SystemDefinition> sysdef,
                   boost::shared_ptr<ParticleGroup> group,
                   boost::shared_ptr<ComputeThermo> thermo_group,
                   boost::shared_ptr<ComputeThermo> thermo_group_t,
                   Scalar tau,
                   Scalar tauP,
                   boost::shared_ptr<Variant> T,
                   boost::python::list S,
                   couplingMode couple,
                   unsigned int flags,
                   const bool nph=false);
        virtual ~TwoStepNPTMTK();

        //! Update the temperature
        /*! \param T New temperature to set
        */
        virtual void setT(boost::shared_ptr<Variant> T)
            {
            m_T = T;
            }

    //! Update the stress components
    /*! \param S list of stress components: [xx, yy, zz, yz, xz, xy]
     */
    virtual void setS(boost::python::list S)
            {
            std::vector<boost::shared_ptr<Variant> > swapS;
            swapS.resize(0);
            for (int i = 0; i< 6; ++i)
                   {
                swapS.push_back(boost::python::extract<boost::shared_ptr<Variant>>(S[i]));
                }
            m_S.swap(swapS);
            }

        //! Update the tau value
        /*! \param tau New time constant to set
        */
        virtual void setTau(Scalar tau)
            {
            m_tau = tau;
            }

        //! Update the nuP value
        /*! \param tauP New pressure constant to set
        */
        virtual void setTauP(Scalar tauP)
            {
            m_tauP = tauP;
            }

        //! Set the scale all particles option
        /*! \param rescale_all If true, rescale all particles
        */
        void setRescaleAll(bool rescale_all)
            {
            m_rescale_all = rescale_all;
            }

        //! Performs the first step of the integration
        virtual void integrateStepOne(unsigned int timestep);

        //! Performs the second step of the integration
        virtual void integrateStepTwo(unsigned int timestep);

        //! Get needed pdata flags
        /*! TwoStepNPTMTK needs the pressure, so the isotropic_virial or pressure_tensor flag is set,
            depending on the integration mode
        */
        virtual PDataFlags getRequestedPDataFlags()
            {
            PDataFlags flags;
            flags[pdata_flag::pressure_tensor] = 1;
            if (m_aniso)
                {
                flags[pdata_flag::rotational_kinetic_energy] = 1;
//                flags[pdata_flag::rotational_virial] = 1;
                }
            flags[pdata_flag::external_field_virial]=1;
            return flags;
            }

        //! Returns a list of log quantities this compute calculates
        std::vector< std::string > getProvidedLogQuantities();

        //! Returns logged values
        Scalar getLogValue(const std::string& quantity, unsigned int timestep, bool &my_quantity_flag);

    protected:
        boost::shared_ptr<ComputeThermo> m_thermo_group;   //!< ComputeThermo operating on the integrated group at t+dt/2
        boost::shared_ptr<ComputeThermo> m_thermo_group_t; //!< ComputeThermo operating on the integrated group at t
        unsigned int m_ndof;            //!< Number of degrees of freedom from ComputeThermo

        Scalar m_tau;                   //!< tau value for Nose-Hoover
        Scalar m_tauP;                  //!< tauP value for the barostat
        boost::shared_ptr<Variant> m_T; //!< Temperature set point
        std::vector<boost::shared_ptr<Variant>> m_S;  //!< Stress matrix (upper diagonal, components [xx, yy, zz, yz, xz, xy])
        Scalar m_V;                     //!< Current volume

        couplingMode m_couple;          //!< Coupling of diagonal elements
        unsigned int m_flags;             //!< Coupling flags for barostat
        bool m_nph;                     //!< True if integrating without thermostat
        Scalar m_mat_exp_v[6];          //!< Matrix exponential for velocity update (upper triangular)
        Scalar m_mat_exp_r[6];          //!< Matrix exponential for position update (upper triangular)
        Scalar m_mat_exp_r_int[6];      //!< Integrated matrix exp. for velocity update (upper triangular)

        bool m_rescale_all;             //!< If true, rescale all particles in the system irrespective of group

        std::vector<std::string> m_log_names; //!< Name of the barostat and thermostat quantities that we log

        //! Helper function to advance the barostat parameters
        void advanceBarostat(unsigned int timestep);

        //! advance the thermostat
        /*!\param timestep The time step
         * \param broadcast True if we should broadcast the integrator variables via MPI
         */
        void advanceThermostat(unsigned int timestep);

        //! Helper function to update the propagator elements
        void updatePropagator(Scalar nuxx, Scalar nuxy, Scalar nuxz, Scalar nuyy, Scalar nuyz, Scalar nuzz);

        };

//! Exports the TwoStepNPTMTK class to python
void export_TwoStepNPTMTK();

#endif // #ifndef __TWO_STEP_NPT_MTK_H__
