#pragma once

/*
 * Nom: William Nolin
 * Code permanent : NOLW76060101
 * Email : william.nolin.1@ens.etsmtl.ca
 */

#include "IKSolver.h"
#include "Armature.h"
#include "SVD.h"

using namespace gti320;

namespace {
}

IKSolver::IKSolver(Armature *_armature, Vector3f &_targetPos) : m_armature(_armature), m_targetPos(_targetPos), m_J() {
}


float IKSolver::getError(Vector3f &dx) const {
    // Compute the error between the current end effector
    //      position and the target position
    dx.setZero();

    const int numLinks = m_armature->links.size();
    Link *endEffector = m_armature->links[numLinks - 1];

    Vector3f f_theta = endEffector->globalPosition();
    Vector3f ddx = m_targetPos - f_theta;
    dx(0) = ddx(0);
    dx(1) = ddx(1);
    dx(2) = ddx(2);

    return ddx.norm();
}

// NOUVELLE FONCTION
// Produit scalaire entre une sous-matrice et un vecteur, permettant d'éviter
//   la copie de données dans un nouveau vecteur pour utiliser l'opérateur fait dans le premier laboratoire.
// La sous matrice doit avoir au moins une colonne, et toutes les autres seront ignorées.
template<typename _Scalar, int _Rows, int _Storage>
float dotP(SubMatrix<_Scalar, _Rows, Dynamic, _Storage> &a, Vector<_Scalar, _Rows> &b) {
    assert(a.cols() >= 1);
    assert(a.rows() == b.rows());

    _Scalar product = 0;
    for (int i = 0; i < b.size(); i++) {
        product += a(i, 0) * b(i);
    }
    return product;
}

// NOUVELLE FONCTION
// Produit vectoriel entre une sous-matrice et un vecteur.
// La sous-matrice doit comporter au moins une colonne et trois rangées, et la sous-matrice résultante doit contenir au moins une colonne.
template<int _Rows, int _Columns, int _Storage>
void crossP(SubMatrix<float, 3, Dynamic, _Storage> &result,
            SubMatrix<float, _Rows, _Columns, _Storage> &a, Vector3f &b) {
    assert(result.cols() >= 1);
    assert(a.rows() >= 3);
    assert(a.cols() >= 1);

    result(0, 0) = a(1, 0) * b(2) - a(2, 0) * b(1);
    result(1, 0) = a(2, 0) * b(0) - a(0, 0) * b(2);
    result(2, 0) = a(0, 0) * b(1) - a(1, 0) * b(0);
}

void IKSolver::solve() {
    const int numLinks = m_armature->links.size();
    const int dim = 3 * (numLinks);
    m_J.resize(3, dim);

    // We assume that the last link is the "end effector"
    //
    Link *endEffector = m_armature->links[numLinks - 1];

    // Build the Jacobian matrix m_J.
    //      Each column corresponds to a separate link
    for (int i = 0; i < numLinks; i++) {
        Link *link = m_armature->links[i];
        Vector3f ri = endEffector->globalPosition() - link->globalPosition();

        for (int j = 0; j < 3; j++) {
            auto ji = m_J.block(0, i * 3 + j, 3, 1);
            auto oij = link->M.block(0, j, 3, 1);
            crossP(ji, oij, ri);
        }
    }

    // Compute the error between the current end effector
    //      position and the target position by calling getError()
    Vector3f dx;
    float error = getError(dx);

    // Compute the change in the joint angles by solving:
    //    df/dtheta * delta_theta = delta_x
    //  where df/dtheta is the Jacobian matrix.
    auto svd = SVD(m_J);
    svd.decompose();

    int m = svd.getSigma().size();
    int rank = 0;
    for (int i = 0; i < m; i++) {
        if (svd.getSigma()(i) == 0) {
            break;
        }

        rank++;
    }

    auto z = Vector<float, Dynamic>(m);
    for (int i = 0; i < rank; i++) {
        auto ui = svd.getU().block(0, i, 3, 1);
        auto si = svd.getSigma()(i);
        z(i) = dotP(ui, dx) / si;
    }
    Vector<float, Dynamic> d_theta = svd.getV() * z;

    // Perform gradient descent method with line search
    //      to move the end effector toward the target position.
    //
    //   Hint: use the Armature::unpack() and Armature::pack() functions
    //   to set and get the joint angles of the armature.
    // 
    //   Hint: whenever you change the joint angles, you must also call
    //   armature->updateKinematics() to compute the global position.
    float alpha = 1;
    Vector<float, Dynamic> initial_theta;
    m_armature->pack(initial_theta);
    do {
        m_armature->unpack(initial_theta + alpha * d_theta);
        m_armature->updateKinematics();

        alpha /= 2;
    } while (getError(dx) >= error && alpha > 0);
}