#include "particlesys.h"
#include <cassert>
using namespace std;

ParticleFunc::ParticleFunc(const std::vector<RcPointer<Force> > &forces,
                           const Vecd &mass) : _forces(forces), _mass(mass)
{
}

Vecd ParticleFunc::operator()(const Vecd &x, const Vecd &v) const
{
    int i;
    int n = x.Elts();
    Vecd ret(n,vl_zero);

    for (i=0; i<(int) _forces.size(); i++)
    {
        // MSVC, please do return value optimization
        SparseVecd sv(_forces[i]->eval(x,v,_mass));
        SVIterd j;
        
        for (j.Begin(sv); !j.AtEnd(); j.Inc())
            ret[j.Index()] += j.Data();
    }

    for (i=0; i<_mass.Elts(); i++)
    {
        if (_mass[i] != HUGE_VAL)
        {
            ret[O(i,0)] /= _mass[i];
            ret[O(i,1)] /= _mass[i];
            ret[O(i,2)] /= _mass[i];
        }
        else
        {
            ret[O(i,0)] = 0;
            ret[O(i,1)] = 0;
            ret[O(i,2)] = 0;
        }
    }

    return ret;
}

SparseMatd  ParticleFunc::dx(const Vecd &x, const Vecd &v) const
{
    int i,j;
    int n = x.Elts();

    SparseMatd ret(n,n);

    for (i=0; i<(int)_forces.size(); i++)
    {
        Force::DblSparse_t M(_forces[i]->evalDx(x,v,_mass));

        for (j=0; j<(int)M.size(); j++)
        {
            int row = M[j].first;
            const SparseVecd &sv = M[j].second;
            ret[row] += sv;
        }
    }

    for (i=0; i<_mass.Elts(); i++)
    {
        if (_mass[i] != HUGE_VAL)
        {
            ret[O(i,0)] /= _mass[i];
            ret[O(i,1)] /= _mass[i];
            ret[O(i,2)] /= _mass[i];
        }
        else
        {
            ret[O(i,0)].MakeZero();
            ret[O(i,1)].MakeZero();
            ret[O(i,2)].MakeZero();
        }
    }

    return ret;
}

SparseMatd  ParticleFunc::dv(const Vecd &x, const Vecd &v) const
{
    int i,j;
    int n = x.Elts();

    SparseMatd ret(n,n);

    for (i=0; i<(int)_forces.size(); i++)
    {
        Force::DblSparse_t M(_forces[i]->evalDv(x,v,_mass));

        for (j=0; j<(int)M.size(); j++)
        {
            int row = M[j].first;
            const SparseVecd &sv = M[j].second;
            ret[row] += sv;
        }
    }

    for (i=0; i<_mass.Elts(); i++)
    {
        if (_mass[i] != HUGE_VAL)
        {
            ret[O(i,0)] /= _mass[i];
            ret[O(i,1)] /= _mass[i];
            ret[O(i,2)] /= _mass[i];
        }
        else
        {
            ret[O(i,0)].MakeZero();
            ret[O(i,1)].MakeZero();
            ret[O(i,2)].MakeZero();
        }
    }

    return ret;
}

double  ParticleFunc::max_displacement(const Vecd &x) const
{
  double max = 0.0;
  for (int i=0; i<(int)_forces.size(); i++)
  {
    double temp = _forces[i]->Displacement(x);
    if (temp > max) max = temp;
  }
  return (max);
}


///////////////////////////////////////////////////////////////////////
ParticleSys::ParticleSys(const vector<Vec3d> &positions, 
                         const vector<Vec3d> &velocities,
                         const vector<double> &masses,
                         const vector<RcPointer<Force> > &forces,
                         const RcPointer<OdeSolver2> &solver)
                         : _solver(solver) 
{
    assert(positions.size() == velocities.size() && velocities.size() == masses.size());
    int n = masses.size();
    
    Vecd mass(n);
    _x.SetSize(3*n);
    _v.SetSize(3*n);
    
    for (int i=0; i<n; i++)
    {
        _x[O(i,0)] = positions[i][0];
        _x[O(i,1)] = positions[i][1];
        _x[O(i,2)] = positions[i][2];
        _v[O(i,0)] = velocities[i][0];
        _v[O(i,1)] = velocities[i][1];
        _v[O(i,2)] = velocities[i][2];
        mass[i] = masses[i];
    }
    
    _pf = new ParticleFunc(forces, mass);
    
    
    
}

vector<Vec3d>   ParticleSys::getPositions() const
{
    int n= _pf->_mass.Elts();
    vector<Vec3d> ret(n);
    
    for (int i=0; i<n; i++)
    {
        ret[i][0] = _x[O(i,0)];
        ret[i][1] = _x[O(i,1)];
        ret[i][2] = _x[O(i,2)];
    }
    
    return ret;
}

vector<Vec3d>   ParticleSys::getVelocities() const
{
    int n= _pf->_mass.Elts();
    vector<Vec3d> ret(n);
    
    for (int i=0; i<n; i++)
    {
        ret[i][0] = _v[O(i,0)];
        ret[i][1] = _v[O(i,1)];
        ret[i][2] = _v[O(i,2)];
    }
    
    return ret;
}

vector<double>  ParticleSys::getMasses() const
{
    int n= _pf->_mass.Elts();
    return vector<double>(_pf->_mass.Ref(), _pf->_mass.Ref()+n);
}

void ParticleSys::setPosition(int p, const Vec3d &v)
{
    unsigned int n= _pf->_mass.Elts();
    assert((unsigned)p < n);
    
    _x[O(p,0)] = v[0];
    _x[O(p,1)] = v[1];
    _x[O(p,2)] = v[2];
}

void ParticleSys::setVelocity(int p, const Vec3d &v)
{
    unsigned int n= _pf->_mass.Elts();
    assert((unsigned)p < n);
    
    _v[O(p,0)] = v[0];
    _v[O(p,1)] = v[1];
    _v[O(p,2)] = v[2];
}

void ParticleSys::setMass(int p, double m)
{
    unsigned int n= _pf->_mass.Elts();
    assert((unsigned)p < n);

    _pf->_mass[p] = m;
}

void ParticleSys::step(double h, double max_allowed_displacement)
{
    //_q = (*_solver)(*_pf, _q, h);
    Vecd x,v;
    static int count = 0;

    count++;

    (*_solver).step(x,v, *_pf, _x, _v, h);

    if ((count < 4) && (max_allowed_displacement > 0.0)) {
      if (_pf->max_displacement(x) > max_allowed_displacement) {
        std::cout << "halving step to " << h / 2.0 << " because max displacement is " << _pf->max_displacement(x) << "\n";
        step(h / 2.0, max_allowed_displacement);
        step(h / 2.0, max_allowed_displacement);

        count--;
        return;
      }
    }

    _x = x;
    _v = v;
    count--;

}
