// SPDX-License-Identifier: LGPL-3.0-or-later
// Author: Kristian Lytje

#pragma once

#include <rigidbody/constraints/IDistanceConstraint.h>
#include <data/DataFwd.h>
#include <utility/observer_ptr.h>

#include <string>

namespace ausaxs::rigidbody::constraints {
    /**
     * This class is the glue that keeps separate bodies together during the optimization. Each constraint is between two individual atoms of two different bodies, and works by 
     * adding a new term to the chi-square which is a function of the distance between the two atoms. Thus the optimizer is penalized depending on how much it separates the two 
     * constrained atoms. 
     */
    class DistanceConstraint : public IDistanceConstraint {
        public: 
            /**
             * @brief Create a new constraint between a pair of atoms.
             * 
             * We use indexes since the bodies and atoms may change during the optimization.
             * Thus to use a constraint, the order of the bodies and atoms cannot change. 
             * 
             * Complexity: O(1)
             * 
             * @param protein The protein this constraint belongs to.
             * @param ibody1 The index of the first body.
             * @param ibody2 The index of the second body.
             * @param iatom1 The index of the first atom in the first body.
             * @param iatom2 The index of the second atom in the second body.
             */
            DistanceConstraint(observer_ptr<const data::Molecule> molecule, int ibody1, int ibody2, int iatom1, int iatom2);

            /**
             * @brief Create a new constraint between a pair of atoms in the two bodies.
             * 
             * Complexity: O(n)
             * 
             * @param body1 The first body.
             * @param body2 The second body.
             * @param center_mass Create a constraint from the center-masses of the two bodies. If false, the two closest carbon atoms are used instead. 
             */
            DistanceConstraint(observer_ptr<const data::Molecule> molecule, int ibody1, int ibody2, bool center_mass = false);

            /**
             * @brief Create a new constraint between a pair of atoms.
             * 
             * Complexity: O(n) where n is the number of atoms in the protein.
             * 
             * @param protein The protein this constraint belongs to.
             * @param atom1 The first atom.
             * @param atom2 The second atom.
             */
            DistanceConstraint(observer_ptr<const data::Molecule> molecule, const data::AtomFF& atom1, const data::AtomFF& atom2);

            virtual ~DistanceConstraint() override = default;

            /**
             * @brief Evaluate this constraint for the current positions. 
             * 
             * @return The chi2 contribution of this constraint.
             */
            double evaluate() const override;

            /**
             * @brief Check if a constraint is identical to this object. 
             * 
             * @param constraint The constraint to be checked for equality. 
             */
            bool operator==(const DistanceConstraint& constraint) const;

            /**
             * @brief Generate a string representation of this constraint.
             */
            std::string to_string() const;

        private: 
            /**
             * @brief Transforms a distance into a proper constraint for least-squares fitting. 
             * 
             * @param offset The radial offset between the new and original positions. 
             */
            static double transform(double offset);
    };
}