/* -------------------------------------------------------------------------- *
 *                      SimTK Core: SimTK Simbody(tm)                         *
 * -------------------------------------------------------------------------- *
 * This is part of the SimTK Core biosimulation toolkit originating from      *
 * Simbios, the NIH National Center for Physics-Based Simulation of           *
 * Biological Structures at Stanford, funded under the NIH Roadmap for        *
 * Medical Research, grant U54 GM072970. See https://simtk.org.               *
 *                                                                            *
 * Portions copyright (c) 2007 Stanford University and the Authors.           *
 * Authors: Peter Eastman                                                     *
 * Contributors: Christopher Bruns                                            *
 *                                                                            *
 * Permission is hereby granted, free of charge, to any person obtaining a    *
 * copy of this software and associated documentation files (the "Software"), *
 * to deal in the Software without restriction, including without limitation  *
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
 * and/or sell copies of the Software, and to permit persons to whom the      *
 * Software is furnished to do so, subject to the following conditions:       *
 *                                                                            *
 * The above copyright notice and this permission notice shall be included in *
 * all copies or substantial portions of the Software.                        *
 *                                                                            *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
 * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
 * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
 * -------------------------------------------------------------------------- */

#include "SimTKmolmodel.h"
#include <iostream>

// define SHOW_VTK for visualisation (for debugging)
// #define SHOW_VTK 1

#ifdef SHOW_VTK
#include "SimTKsimbody_aux.h"
#endif

using namespace SimTK;
using namespace SimTK::units::md;
using namespace std;

#define ASSERT(cond) {SimTK_ASSERT_ALWAYS(cond, "Assertion failed");}

const int NUM_MOLECULES = 10;
const MDLength BOND_LENGTH = 0.5 * nanometers;

typedef units::md::multiply_typeof_helper<MDEnergy, MDEnergy>::type
    SquareEnergy;

/**
 * Create a force between every pair of bodies (including ground) with the potential k*(1/x^2 - 1/x).
 */

class BetweenBodyForce : public Force::Custom::Implementation {
public:
    BetweenBodyForce(const SimbodyMatterSubsystem& matter) : matter(matter), k(20.0) {
    }
    bool dependsOnlyOnPositions() {
        return true;
    }
    void calcForce(const State& state, Vector_<SpatialVec>& bodyForces, Vector_<Vec3>& particleForces, Vector& mobilityForces) const {
        for (int i = 0; i < matter.getNumBodies(); ++i) {
            const MobilizedBody& body1 = matter.getMobilizedBody(MobilizedBodyIndex(i));
            const MDDisplacement pos1 = body1.getBodyOriginLocation(state) * nanometers;
            for (int j = i+1; j < matter.getNumBodies(); ++j) {
                const MobilizedBody& body2 = matter.getMobilizedBody(MobilizedBodyIndex(j));
                const MDDisplacement pos2 = body2.getBodyOriginLocation(state) * nanometers;
                const MDLength dist = (pos2-pos1).norm();
                const MDInverseLength invDist = 1.0/dist;

                typedef units::md::multiply_typeof_helper<MDInverseLength, MDInverseLength>::type 
                    OOLength2;
                const OOLength2 invDist2 = invDist*invDist;

                // This potential does not lend itself well to units
                const MDForce f(k*(dist-2.0*nanometers)*(pos2-pos1)*invDist2*invDist2 * kilojoules_per_mole * nanometers);

                const kilojoule_per_mole_nanometer_unit& fu = kilojoules_per_mole_nanometer;
                body1.applyBodyForce(state, SpatialVec(Vec3(0),  f/fu), bodyForces);
                body2.applyBodyForce(state, SpatialVec(Vec3(0), -f/fu), bodyForces);
            }
        }
    }
    Real calcPotentialEnergy(const State& state) const {
        MDEnergy pe = 0 * kilojoules_per_mole;
        for (int i = 0; i < matter.getNumBodies(); ++i) {
            const MobilizedBody& body1 = matter.getMobilizedBody(MobilizedBodyIndex(i));
            const MDDisplacement pos1 = body1.getBodyOriginLocation(state) * nanometers;
            for (int j = i+1; j < matter.getNumBodies(); ++j) {
                const MobilizedBody& body2 = matter.getMobilizedBody(MobilizedBodyIndex(j));
                const MDDisplacement pos2 = body2.getBodyOriginLocation(state) * nanometers;
                const MDLength dist = (pos2-pos1).norm();
                const MDInverseLength invDist = 1.0/dist;

                // This potential does not lend itself well to units
                pe += MDEnergy(k*(invDist*nanometers-1)*invDist*nanometers * kilojoules_per_mole);
            }
        }
        return pe/kilojoules_per_mole;
    }
    const SimbodyMatterSubsystem& matter;

    // This potential does not lend itself well to units
    const Real k;
};

class EnergyMonitor : public PeriodicEventReporter {
public:
    typedef units::md::multiply_typeof_helper<MDEnergy, MDEnergy>::type MDEnergySquared;
    mutable int eventCount;
    mutable MDEnergy sumEnergy;
    mutable MDEnergySquared sumEnergySquared;
    EnergyMonitor(MultibodySystem& system) 
        : PeriodicEventReporter(0.05), 
          system(system), 
          eventCount(0), 
          sumEnergy(0.0 * kilojoules_per_mole), 
          sumEnergySquared(0.0 * kilojoules_per_mole * kilojoules_per_mole) 
    {}
    void handleEvent(const State& state) const {
        if (state.getTime() <= 1.0)
            return;
        eventCount++;
        system.realize(state, Stage::Dynamics);
        MDEnergy energy = system.calcKineticEnergy(state) * kilojoules_per_mole;
        sumEnergy += energy;
        sumEnergySquared += energy*energy;
    }
private:
    MultibodySystem& system;
};

void testTemperature(MDTemperature temperature) {
    MultibodySystem mbs;
    SimbodyMatterSubsystem matter(mbs);
    GeneralForceSubsystem forces(mbs);
    
    // Create a gas of two atom molecules.
    
    Body::Rigid body = Body::Rigid(MassProperties(1, Vec3(0), Inertia(1)))
                                  .addDecoration(Transform(), DecorativeSphere(.1));
    Random::Uniform random(-10.0, 10.0);
    for (int i = 0; i < NUM_MOLECULES; ++i) {
        MDDisplacement pos(Vec3(random.getValue(), random.getValue(), random.getValue())*nanometers);
        MobilizedBody::Translation first(
            matter.Ground(), 
            Transform(), 
            body, 
            Transform(pos/nanometers));
        MobilizedBody::LineOrientation second(
            first, 
            Transform(), 
            body, 
            Transform(Vec3(0, 0, BOND_LENGTH/nanometers)));
    }
    Force::Custom(forces, new BetweenBodyForce(matter));
    mbs.updDefaultSubsystem().addEventHandler(new VelocityRescalingThermostat(mbs, temperature, 1.0*picoseconds));

#ifdef SHOW_VTK
    mbs.updDefaultSubsystem().addEventReporter(new VTKEventReporter(mbs, 0.1));
#endif

    EnergyMonitor* monitor = new EnergyMonitor(mbs);
    mbs.getDefaultSubsystem().addEventReporter(monitor);
    mbs.realizeTopology();
    State& state = mbs.updDefaultState();
    mbs.realizeModel(state);
    
    // Minimize energy, so low temperature version does not have to go too far
    LocalEnergyMinimizer::minimizeEnergy(mbs, state, 0.1);

    // Simulate it.
    
    RungeKuttaMersonIntegrator integ(mbs);
    TimeStepper ts(mbs, integ);
    ts.initialize(mbs.getDefaultState());
    ts.stepTo(30);
    
    // See if the temperature was being correctly maintained.
    

    ASSERT(monitor->eventCount > 100);
    MDEnergy meanEnergy = monitor->sumEnergy/monitor->eventCount;
    MDEnergy expectedEnergy = 5.0*NUM_MOLECULES*0.5*units::md::BoltzmannConstant*temperature; // kT/2 per degree of freedom
    SquareEnergy meanSqrEnergy = monitor->sumEnergySquared/monitor->eventCount;
    MDEnergy stdDevEnergy = units::md::sqrt(meanSqrEnergy-meanEnergy*meanEnergy);
    ASSERT(std::abs(meanEnergy/expectedEnergy-1.0) < 0.2);
    ASSERT(stdDevEnergy < 0.5*meanEnergy);
    std::cout << "Done" << std::endl;
}

int main() 
{
    using units::md::kelvins;

    cout << "300K" << endl;
    testTemperature(300.0 * kelvins);

    // Comment out some tests to make nightly build tests go faster

    //cout << "5000K" << endl;
    //testTemperature(5000.0);

    //cout << "10K" << endl;
    //testTemperature(10.0);

    return 0;
}