#include "ParticleSystem.h"

using namespace gti320;

/**
 * Calcule des forces qui affectent chacune des particules.
 *
 * Étant donné une particule p, la force est stockée dans le vecteur p.f
 * Les forces prisent en compte sont : la gravité et la force des ressorts.
 */
void ParticleSystem::computeForces() {
    // Calcul de la force gravitationnelle sur chacune des particules
    for (Particle &p: m_particles) {
        p.f(0) = 0; // x
        p.f(1) = -p.m * 9.81f;  // y
    }

    // Pour chaque ressort, ajouter la force exercée à chacune des extrémités sur
    // les particules appropriées. Pour un ressort s, les deux particules
    // auxquelles le ressort est attaché sont m_particles[s.index0] et
    // m_particles[s.index1]. On rappelle que les deux forces sont de même
    // magnitude mais dans des directions opposées.
    for (const Spring &s: m_springs) {
        Particle &p0 = m_particles[s.index0];
        Particle &p1 = m_particles[s.index1];

        Vector<float, 2> distance = p1.x - p0.x;
        Vector<float, 2> force = (s.k * (1 - s.l0 / distance.norm())) * distance;

        p0.f = p0.f + force;
        p1.f = p1.f - force;
    }
}

/**
 * Assemble les données du système dans les trois vecteurs d'état _pos,
 * _vel et _force.
 */
void ParticleSystem::pack(Vector<float, Dynamic> &_pos,
                          Vector<float, Dynamic> &_vel,
                          Vector<float, Dynamic> &_force) {
    // Copier les données des particules dans les vecteurs. Attention, si on a
    // changé de modèle, il est possible que les vecteurs n'aient pas la bonne
    // taille. Rappel : la taille de ces vecteurs doit être 2 fois le nombre de
    // particules.
    int required_size = (int) this->m_particles.size() * 2;
    if (_pos.size() != required_size) {
        // Si un des vecteurs n'est pas la bonne grandeur, un changement de modèle a eu lieu et on doit redimensionner tous les vecteurs.
        _pos.resize(required_size);
        _vel.resize(required_size);
        _force.resize(required_size);
    }

    for (int i = 0; i < m_particles.size(); i++) {
        Particle &p = m_particles[i];
        int ix = i * 2;
        int iy = i * 2 + 1;

        _pos(ix) = p.x(0);
        _pos(iy) = p.x(1);
        _vel(ix) = p.v(0);
        _vel(iy) = p.v(1);
        _force(ix) = p.f(0);
        _force(iy) = p.f(1);
    }
}

/**
 * Copie les données des vecteurs d'états dans le particules du système.
 */
void ParticleSystem::unpack(const Vector<float, Dynamic> &_pos,
                            const Vector<float, Dynamic> &_vel) {
    // Mise à jour des propriétés de chacune des particules à partir des valeurs
    // contenues dans le vecteur d'état.

    for (int i = 0; i < _pos.size(); i += 2) {
        int ix = i;
        int iy = i + 1;
        Particle &p = m_particles[i / 2];

        p.x(0) = _pos(ix);
        p.x(1) = _pos(iy);
        p.v(0) = _vel(ix);
        p.v(1) = _vel(iy);
    }
}


/**
 * Construction de la matrice de masses.
 */
void ParticleSystem::buildMassMatrix(Matrix<float, Dynamic, Dynamic> &M) {
    const int numParticles = m_particles.size();
    const int dim = 2 * numParticles;
    M.resize(dim, dim);
    M.setZero();

    // Inscrire la masse de chacune des particules dans la matrice de masses M
    //
    for (int i = 0; i < numParticles; ++i) {
        int j = i * 2;
        int k = j + 1;
        float m = m_particles[i].m;
        if (m_particles[i].fixed) {
            m = 10000000;
        }

        M(j, j) = m;
        M(k, k) = m;
    }
}


/**
 * Construction de la matrice de rigidité.
 */
void ParticleSystem::buildDfDx(Matrix<float, Dynamic, Dynamic> &dfdx) {
    const int numParticles = m_particles.size();
    const int numSprings = m_springs.size();
    const int dim = 2 * numParticles;
    dfdx.resize(dim, dim);
    dfdx.setZero();

    auto identity = Matrix<float, 2, 2>();
    identity.setIdentity();

    // Pour chaque ressort...
    for (const Spring &spring: m_springs) {
        // Calculer le coefficients alpha et le produit dyadique tel que décrit au cours.
        // Utiliser les indices spring.index0 et spring.index1 pour calculer les coordonnées des endroits affectés.
        //
        // Astuce: créer une matrice de taille fixe 2 par 2 puis utiliser la classe SubMatrix pour accumuler 
        // les modifications sur la diagonale (2 endroits) et pour mettre à jour les blocs non diagonale (2 endroits).

        auto p0 = m_particles[spring.index0];
        auto p1 = m_particles[spring.index1];

        float l = (p1.x - p0.x).norm();
        float a = spring.k * (1 - spring.l0 / l);
        Vector<float, 2> dij = (1 / l) * (p1.x - p0.x);
        Matrix<float, 2, 2> correction = spring.k * (spring.l0 / l) * (dij.as_matrix() * dij.transpose());

        Matrix<float, 2, 2> ndiag = a * identity + correction;
        Matrix<float, 2, 2> diag = -1.0f * ndiag;

        dfdx.block(spring.index0 * 2, spring.index0 * 2, 2, 2) += diag;
        dfdx.block(spring.index1 * 2, spring.index1 * 2, 2, 2) += diag;
        dfdx.block(spring.index0 * 2, spring.index1 * 2, 2, 2) += ndiag;
        dfdx.block(spring.index1 * 2, spring.index0 * 2, 2, 2) += ndiag;
    }
}