#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
This modules defines utility functions for the pose extractor node.


"""

import copy
import numpy as np
import geometry_msgs.msg
import udom_modeling_msgs.msg
import tf
import tf.transformations as transformations


def extract_points(mesh, indices, zero_based=True):
    """
    Extracts the nodes' positions based on the mesh and indices.

    :param mesh: The mesh from which the pose will be extracted.
    :type mesh: udom_modeling_msgs.msg.Mesh

    :param indices: Indices of the node to return.
    :type indices: list of int

    :param zero_based: If true, it assumes the indices are zero-based.
        Otherwise, it assumes the indices start from one.
    :type zero_based: bool

    :return: The extracted nodes' positions, if there is an error it returns None.
    :rtype: list of geometry_msgs.msg.Point or None

    """
    try:
        if zero_based:
            return [mesh.vertices[index] for index in indices]
        else:
            return [mesh.vertices[index - 1] for index in indices]
    except IndexError:
        return None


def list_to_mesh(mesh, points, nodes=None):
    """
    Updates the nodes in a mesh based on a list which follows this order:
    [p1_x, p1_y, p_1z, p2_x, p2_y, p2_z....]

    :param mesh: The mesh to update.
    :type mesh: udom_modeling_msgs.msg.Mesh

    :param points: The points representing the mesh's nodes position.
    :type points: list

    :param nodes: The indices of the nodes (in order) that are represented by the points.
    :type nodes: list

    :return: The mesh with the updated position of its nodes.
    :rtype: udom_modeling_msgs.msg.Mesh

    """
    if nodes:
        mesh_out = udom_modeling_msgs.msg.Mesh()
        for ii, _ in enumerate(nodes):
            mesh_out.vertices.append(geometry_msgs.msg.Point(
                x=points[ii * 3], y=points[ii * 3 + 1], z=points[ii * 3 + 2]))
    else:
        mesh_out = copy.deepcopy(mesh)
        for ii, _ in enumerate(mesh_out.vertices):
            mesh_out.vertices[ii] = geometry_msgs.msg.Point(
                x=points[ii * 3], y=points[ii * 3 + 1], z=points[ii * 3 + 2])
    return mesh_out


def mesh_to_list(mesh, indices=None):
    """
    Converts the nodes in a mesh into a list.
    The order follows: [p1_x, p1_y, p_1z, p2_x, p2_y, p2_z....]

    :param mesh: The mesh from which the points will be converted to a list.
    :type mesh: udom_modeling_msgs.msg.Mesh

    :param indices: If specified, it returns only the position for the
        nodes represented by the indices.
    :type indices: list

    :return: The extracted nodes' positions.
    :rtype: list

    """
    points = []
    for ii, vertex in enumerate(mesh.vertices):
        if indices:
            if ii in indices:
                p_x = vertex.x
                p_y = vertex.y
                p_z = vertex.z
                points.extend([p_x, p_y, p_z])
        else:
            p_x = vertex.x
            p_y = vertex.y
            p_z = vertex.z
            points.extend([p_x, p_y, p_z])
    return points


def extract_forces(force_array, indices):
    """
    Extracts the force values on the nodes as represented by the indices.

    :param force_array: An array representing the three-dimensional forces applied
        to each node on a mesh.
    :type force_array: std_msgs.msg.Float32MultiArray

    :param indices: Indices of the nodes where the force is applied.
    :type indices: list

    :return: The extracted forces.
    :rtype: list

    """
    forces = []
    try:
        for ii in indices:
            f_x = force_array.data[ii * 3]
            f_y = force_array.data[ii * 3 + 1]
            f_z = force_array.data[ii * 3 + 2]
            forces.extend([f_x, f_y, f_z])
        return forces
    except IndexError:
        return None


def fill_forces(forces, indices, number_of_nodes):
    """
    Fills a force array of length number_of_nodes*3 with the forces,
    where the indices represent the nodes where the forces are applied.
    The forces are assumed to have the same order as the indices and
    they are arranged as 'fx, fy, fz'.

    :param forces: An array representing the three-dimensional forces applied
        on the nodes (as specified by the indices) of the mesh.
    :type forces: list

    :param indices: Indices of the nodes where the force is applied.
    :type indices: list

    :param number_of_nodes: Number of nodes of the mesh.
    :type number_of_nodes: int

    :return: The filled force array.
    :rtype: list

    """
    force_array = np.zeros(number_of_nodes * 3).tolist()

    try:
        for ii, index in enumerate(indices):
            force_array[index * 3] = forces[ii * 3]
            force_array[index * 3 + 1] = forces[ii * 3 + 1]
            force_array[index * 3 + 2] = forces[ii * 3 + 2]
        return force_array
    except IndexError:
        return None


def compute_centroid(points):
    """
    Computes the centroid of the points based on the arithmetic mean.

    :param points: The points on which the centroid will be computed.
    :type points: list of geometry_msgs.msg.Point

    :return: The centroid of the points.
    :rtype: geometry_msgs.msg.Point

    """
    centroid = geometry_msgs.msg.Point()
    points = [[point.x, point.y, point.z] for point in points]

    c = np.mean(points, axis=0)
    centroid.x = c[0]
    centroid.y = c[1]
    centroid.z = c[2]

    return centroid


def compute_normal_as_quaternion(points, rotation=90.0):
    """
    Computes the normal, as a quaternion, based on three points.

    :param points: The points on which the normal will be computed.
    :type points: list of geometry_msgs.msg.Point

    :param rotation: Rotates the pose around the normal axis by the specified angle (in degrees).
    :type rotation: float

    :return: The normal of the three points as a quaternion.
    :rtype: geometry_msgs.msg.Quaternion

    """
    quaternion = geometry_msgs.msg.Quaternion()

    p_1 = np.array([points[0].x, points[0].y, points[0].z])
    p_2 = np.array([points[1].x, points[1].y, points[1].z])
    p_3 = np.array([points[2].x, points[2].y, points[2].z])
    normal = np.cross((p_2 - p_1), (p_3 - p_1))

    quat = tf.transformations.quaternion_about_axis(np.radians(rotation), normal)
    quaternion.x = quat[0]
    quaternion.y = quat[1]
    quaternion.z = quat[2]
    quaternion.w = quat[3]

    return quaternion


def rotate_pose(pose, angle=0.0, reference_axis='z'):
    """
    Rotates the orientation of a pose, for a single rotation axis (reference_axis),
    by the specified angle.

    :param pose: The pose to be modified.
    :type pose: geometry_msgs.msg.PoseStamped

    :param angle: The angle to rotate the pose (in degrees).
    :type angle: float

    :param reference_axis: The rotation axis of the pose to be modified (e.g. x, y, z).
    :type reference_axis: str

    :return: The modified pose.
    :rtype: geometry_msgs.msg.PoseStamped

    """
    assert reference_axis.lower() in ['x', 'y', 'z'], \
        "'reference_axis' must be 'x', 'y', or 'z'."

    pose_out = copy.deepcopy(pose)
    orientation_in = (
        pose_out.pose.orientation.x, pose_out.pose.orientation.y,
        pose_out.pose.orientation.z, pose_out.pose.orientation.w)

    angles_out = np.array(transformations.euler_from_quaternion(orientation_in))

    # Ensure the offset is not more than 360 degrees
    angle %= 360

    if reference_axis.lower() == 'x':
        angles_out[0] += np.radians(angle)
    elif reference_axis.lower() == 'y':
        angles_out[1] += np.radians(angle)
    elif reference_axis.lower() == 'z':
        angles_out[2] += np.radians(angle)

    orientation_out = transformations.quaternion_from_euler(
        angles_out[0], angles_out[1], angles_out[2])

    pose_out.pose.orientation.x = orientation_out[0]
    pose_out.pose.orientation.y = orientation_out[1]
    pose_out.pose.orientation.z = orientation_out[2]
    pose_out.pose.orientation.w = orientation_out[3]

    return pose_out