"""
Created on Fri Apr 21 2023
@name:   position_manipulation.py
@desc:   This module contains some utility functions for position manipulation.
@auth:   Djalim Simaila
@e-mail: djalim.simaila@inrae.fr
"""
import numpy as np
from utils.files.input import ScannedObject
from utils.math.data_extraction import get_mean


def get_mass_properties(obj:ScannedObject)->tuple:
    '''
    Evaluate and return a tuple with the following elements:
        - the volume
        - the position of the center of gravity (COG)
        - the inertia matrix expressed at the COG
    :param obj: Object to analyse
    :return: tuple(float, numpy.array, numpy.array)

    From numpy-stl:(https://pypi.org/project/numpy-stl/)
        Documentation can be found here:
        http://www.geometrictools.com/Documentation/PolyhedralMassProperties.pdf

    '''

    verts = np.asarray(obj.get_vertices())
    faces = np.asarray(obj.get_faces())
    faces = verts[faces]
    x = np.asarray([[point[0] for point in faces] for faces in faces])
    y = np.asarray([[point[1] for point in faces] for faces in faces])
    z = np.asarray([[point[2] for point in faces] for faces in faces])

    def subexpression(x):
        w0, w1, w2 = x[:, 0], x[:, 1], x[:, 2]
        temp0 = w0 + w1
        f1 = temp0 + w2
        temp1 = w0 * w0
        temp2 = temp1 + w1 * temp0
        f2 = temp2 + w2 * f1
        f3 = w0 * temp1 + w1 * temp2 + w2 * f2
        g0 = f2 + w0 * (f1 + w0)
        g1 = f2 + w1 * (f1 + w1)
        g2 = f2 + w2 * (f1 + w2)
        return f1, f2, f3, g0, g1, g2

    x0, x1, x2 = x[:, 0], x[:, 1], x[:, 2]
    y0, y1, y2 = y[:, 0], y[:, 1], y[:, 2]
    z0, z1, z2 = z[:, 0], z[:, 1], z[:, 2]
    a1, b1, c1 = x1 - x0, y1 - y0, z1 - z0
    a2, b2, c2 = x2 - x0, y2 - y0, z2 - z0
    d0, d1, d2 = b1 * c2 - b2 * c1, a2 * c1 - a1 * c2, a1 * b2 - a2 * b1

    f1x, f2x, f3x, g0x, g1x, g2x = subexpression(x)
    f1y, f2y, f3y, g0y, g1y, g2y = subexpression(y)
    f1z, f2z, f3z, g0z, g1z, g2z = subexpression(z)

    intg = np.zeros((10))
    intg[0] = sum(d0 * f1x)
    intg[1:4] = sum(d0 * f2x), sum(d1 * f2y), sum(d2 * f2z)
    intg[4:7] = sum(d0 * f3x), sum(d1 * f3y), sum(d2 * f3z)
    intg[7] = sum(d0 * (y0 * g0x + y1 * g1x + y2 * g2x))
    intg[8] = sum(d1 * (z0 * g0y + z1 * g1y + z2 * g2y))
    intg[9] = sum(d2 * (x0 * g0z + x1 * g1z + x2 * g2z))
    intg /= np.array([6, 24, 24, 24, 60, 60, 60, 120, 120, 120])
    volume = intg[0]
    cog = intg[1:4] / volume
    cogsq = cog ** 2
    inertia = np.zeros((3, 3))
    inertia[0, 0] = intg[5] + intg[6] - volume * (cogsq[1] + cogsq[2])
    inertia[1, 1] = intg[4] + intg[6] - volume * (cogsq[2] + cogsq[0])
    inertia[2, 2] = intg[4] + intg[5] - volume * (cogsq[0] + cogsq[1])
    inertia[0, 1] = inertia[1, 0] = -(intg[7] - volume * cog[0] * cog[1])
    inertia[1, 2] = inertia[2, 1] = -(intg[8] - volume * cog[1] * cog[2])
    inertia[0, 2] = inertia[2, 0] = -(intg[9] - volume * cog[2] * cog[0])
    return volume, cog, inertia

def verticalise(obj:ScannedObject):
    """
    Rotate the object so that the principal axis of inertia is vertical
    :param obj: Object to analyse
    """
    cog = get_mass_properties(obj)
    cog, inertia = get_mass_properties(obj)[1:]
    [val,vect] = np.linalg.eig(inertia)

    if np.linalg.det(vect) < 0:
        vect[:,2] = - vect[:,2]

    rot = vect.T
    faces = obj.get_faces(True)

    theta = -np.pi/2

    R_y = np.array([[np.cos(theta), 0, np.sin(theta)],
                    [0, 1, 0],
                    [-np.sin(theta), 0, np.cos(theta)]])
    
    for face,_ in enumerate(faces):
        for vertex in range(3):
            faces[face][vertex] = rot@(faces[face][vertex] - cog)
            faces[face][vertex] = R_y@(faces[face][vertex])

    obj.update_from_faces(faces)