#!/usr/bin/python3
# Script to have a better analysis on exercise 2 session 5 of Basics of
# Robotics for Manipulation
#
# Louis Munier - 211026
# Evgenia Roussinova and Louis Munier - 221026
from typing import List
from math import radians
from sys import platform

import numpy as np
import matplotlib.pyplot as plt

if platform == "linux" or platform == "linux2":
    from mpl_toolkits import mplot3d


def compute_grid_search_res(delta_alpha: float, d: float, range_q: List[float],
                            resolution: int, plot: bool = True) -> np.float:
    """
        Compute the different resolution values to plot them
    """
    q_1 = np.tile(np.linspace(
        range_q[0], range_q[1], num=resolution), (resolution, 1)).T
    q_2 = np.tile(np.linspace(
        range_q[2], range_q[3], num=resolution), (resolution, 1))

    um_scale = np.power(10, 6)
    d_delta_a = d * delta_alpha

    delta_x = -d_delta_a * (np.sin(q_1) + 2 * np.sin(q_1 + q_2)) * um_scale
    delta_y = d_delta_a * (np.cos(q_1) + 2 * np.cos(q_1 + q_2)) * um_scale

    delta_s = np.sqrt(np.power(delta_x, 2) + np.power(delta_y, 2))

    if plot:
        max_delta_x = round(np.max(delta_x), 2)
        max_rel_delta_x = round(np.max(delta_x) / (um_scale * d_delta_a), 2)
        print(
            f"Max delta X : {max_delta_x} [um] => max delta X/(d*delta alpha) : {max_rel_delta_x}")

        max_delta_y = round(np.max(delta_y), 2)
        max_rel_delta_y = round(np.max(delta_y) / (um_scale * d_delta_a), 2)
        print(
            f"Max delta Y: {max_delta_y} [um] => max delta Y/(d*delta alpha) : {max_rel_delta_y}")

        max_delta_s = round(np.max(delta_s), 2)
        max_rel_delta_ys = round(np.max(delta_s) / (um_scale * d_delta_a), 2)
        print(
            f"Max delta S : {max_delta_s} [um] => max delta S/(d*delta alpha) : {max_rel_delta_ys}")

        labels = {
            "title": "Delta X resolution with respect to q1 and q2.",
            "xlabel": "q1 [°]",
            "ylabel": "q2 [°]",
            "zlabel": "Delta X [um]"
        }
        plot_resolutions(q_1, q_2, delta_x, labels)

        labels = {
            "title": "Delta Y resolution with respect to q1 and q2.",
            "xlabel": "q1 [°]",
            "ylabel": "q2 [°]",
            "zlabel": "Delta Y [um]"
        }
        plot_resolutions(q_1, q_2, delta_y, labels)

        labels = {
            "title": "Delta S resolution with respect to q1 and q2.",
            "xlabel": "q1 [°]",
            "ylabel": "q2 [°]",
            "zlabel": "Delta S [um]"
        }
        plot_resolutions(q_1, q_2, delta_s, labels)
        plt.show()


def plot_resolutions(data_x: np.float, data_y: np.float, data_z: np.float, labels: List[str]):
    """
        Plot the different resolutions on 3D graphs
    """
    plt.figure()
    ax = plt.axes(projection='3d')
    ax.plot_surface(np.degrees(data_x), np.degrees(data_y), data_z, rstride=1, cstride=1,
                    cmap='viridis', edgecolor='none')
    ax.set_title(labels['title'])
    ax.set_xlabel(labels['xlabel'])
    ax.set_ylabel(labels['ylabel'])
    ax.set_zlabel(labels['zlabel'])


if __name__ == '__main__':
    D = 0.350
    RESOLUTION = 50
    DELTA_ALPHA = radians(0.001)
    RANGE_Q = [0, 2*np.pi, -np.pi, np.pi]

    compute_grid_search_res(DELTA_ALPHA, D, RANGE_Q, RESOLUTION)