#include "Initialization_stage.h"

#include <sstream>
#include <math.h>

#include "System_of_Units.h"
#include "Atom.h"
#include "Atom_kit.h"
#include "Model_proxy.h"
#include "Protocol.h"
#include "User.h"

namespace MM
{

//T_Base_random    Initialization_stage::rand_(0., 1.);
MTRand          Initialization_stage::distr_ (156416);

Initialization_stage::
Initialization_stage (Protocol &protocol, Text const & init)
:
    Protocol_stage_frame (protocol)//,
    //rand_(0., 1.)
{
    using namespace std;
    stringstream str(init.c_str());
    string intention, T, P, comment_first_word, comment_end;
    str >> intention >> T >> P >> comment_first_word;

    temperature_.from_text (T.c_str());
    pressure_.   from_text (P.c_str());

    getline (str, comment_end);
    comment_ = (comment_first_word + comment_end).c_str();
}

void Initialization_stage::
execute ()
{
    setup_velocities (protocol_.model(), temperature_);
}

void Initialization_stage::
setup_velocities (Atom_group & model, Temperature_unit const& temperature)
{
    double T0 =
        //System_of_Units::singleton().temperature (temperature_) *
        System_of_Units::singleton().temperature (temperature) *
        System_of_Units::singleton().gas_constant();

    //Model & model = protocol_.model();

    int i;
    int count = model.atom_count();

    for (i=0;  i<count;  ++i)
    {
        Atom & atom  = model.atom (i);
        double mass  = atom.mass();
        double sigma = sqrt (T0 / mass);
        double x     = sigma * gauss ();
        double y     = sigma * gauss ();
        double z     = sigma * gauss ();
        atom.kit().set_velocity (Vector_3D_impl (x, y, z));
    }

    // remove mass center movement
    Vector_3D_impl sum (0., 0., 0.);
    double sum_mass = 0.;

    for (i=0;  i<count;  ++i)
    {
        Atom & atom  = model.atom (i);
        double mass  = atom.mass();
        sum.add_scaled (mass, atom.kit().velocity ());
        sum_mass += mass;
    }

    sum *= 1. / sum_mass;
    sum.negate ();

    for (i=0;  i<count;  ++i)
    {
        Atom & atom  = model.atom (i);
        atom.kit().add_velocity (sum);
    }


    //fix remove moments
}

double Initialization_stage::
gauss ()
{
    return distr_.randNorm ();

    //double R_2 = 2., v1, v2;

    //while (R_2 >= 1.)
    //{
    //    v1 = 2. * rand_.next() - 1.;
    //    v2 = 2. * rand_.next() - 1.;
    //    R_2 = v1*v1 + v2*v2;
    //}

    //double result = v1 * sqrt (-2. * log(R_2) / R_2);
    //return result;
}

Text Initialization_stage::
option (Text const & key) const
{
    if (key == "T")                 return temperature_.text ();
    else if (key == "P")            return pressure_.text ();
    else if (key == "Comment")      return comment_;
    return "";
}

void Initialization_stage::
set_option (Text const & key, Text const & value)
{
    if (key=="T")
        temperature_.from_text (value);

    else if (key=="P")
        pressure_.from_text (value);

    else if (key=="Comment")
        comment_ = value;

    else
        to_user().status (Text("'") + key + "' - unknown key.");
}

void Initialization_stage::
get_option (int i, Text & key, Text & value) const
{
    switch (i)
    {
    case 0: key = "T";          value = "298 K";        return;
    case 1: key = "P";          value = "1 atm";        return;
    case 2: key = "Comment";    value = "No comments";  return;
    default:
        FLAW (Text(i) + " is out of range. Options count is " 
                      + option_count() + ".");
    }
}

int Initialization_stage::
option_count () const
{
    return 3;
}
}//MM


