"""Gear Element module.

This module defines the GearElement classes which can be used to represent
gears or gearboxes used to couple different shafts in the MultiRotor class.
"""

import numpy as np
from plotly import graph_objects as go
from ross.units import Q_

from ross.disk_element import DiskElement
from ross.materials import steel


__all__ = ["GearElement"]


class GearElement(DiskElement):
    """A gear element.

    This class creates a gear element from input data of inertia and mass.

    Parameters
    ----------
    n: int
        Node in which the gear will be inserted.
    m : float, pint.Quantity
        Mass of the gear element (kg).
    Id : float, pint.Quantity
        Diametral moment of inertia.
    Ip : float, pint.Quantity
        Polar moment of inertia.
    width: float, pint.Quantity
        width of the gear (mm).
    n_teth: int, pint.Quantity
        number of teeth from the gear
    base_diameter : float, pint.Quantity
        Base diameter of the gear (m).
        If given pitch_diameter is not necessary.
    pitch_diameter : float, pint.Quantity
        Pitch diameter of the gear (m).
        If given base_diameter is not necessary.
    pressure_angle : float, pint.Quantity, optional
        The normal pressure angle of the gear (rad).
        Default is 20 deg (converted to rad).
    material: ross.material, optional
        material of the gear. Default is steel.
    tag : str, optional
        A tag to name the element.
        Default is None.
    helix_angle: float, pint.Quantity, optional
        value of helix angle for helical gears. Default is 0 representing spur gear.
    scale_factor: float or str, optional
        The scale factor is used to scale the gear drawing.
        For gears it is also possible to provide 'mass' as the scale factor.
        In this case the code will calculate scale factors for each gear based
        on the gear with the higher mass. Notice that in this case you have to
        create all gears with the scale_factor='mass'.
        Default is 1.
    color : str, optional
        A color to be used when the element is represented.
        Default is 'Goldenrod'.

    Examples
    --------
    >>> gear = GearElement(
    ...        n=0, m=4.67, Id=0.015, Ip=0.030, width=0.07, n_teeth=50,
    ...        pitch_diameter=0.187,
    ...        pressure_angle=Q_(22.5, "deg"),
    ...        helix_angle=0,
    ... )
    >>> gear.pressure_angle # doctest: +ELLIPSIS
    0.392699...
    """

    def __init__(
        self,
        n,
        m,
        Id,
        Ip,
        width,
        n_teeth,
        pitch_diameter=None,
        base_diameter=None,
        pressure_angle=None,
        material=steel,
        tag=None,
        helix_angle=0,
        scale_factor=1.0,
        color="Goldenrod",
    ):
        if pressure_angle is None:
            pressure_angle = Q_(20, "deg")

        self.pressure_angle = float(pressure_angle)

        if base_diameter:
            self.base_radius = float(base_diameter) / 2
        elif pitch_diameter:
            self.base_radius = float(pitch_diameter) * np.cos(self.pressure_angle) / 2
        else:
            raise TypeError(
                "At least one of the following must be informed for GearElement: base_diameter or pitch_diameter"
            )

        if not pitch_diameter:
            self.pitch_diameter = (self.base_radius * 2) / np.cos(self.pressure_angle)
        else:
            self.pitch_diameter = float(pitch_diameter)

        self.helix_angle = float(helix_angle)
        self.width = float(width)

        self.material = material

        self.n_teeth = n_teeth

        super().__init__(n, m, Id, Ip, tag, scale_factor, color)

    @classmethod
    def from_geometry(
        cls,
        n,
        material,
        width,
        i_d,
        o_d,
        n_teeth,
        pressure_angle=None,
        tag=None,
        helix_angle=0,
        scale_factor=1.0,
        color="Goldenrod",
    ):
        """Create a gear element from geometry properties.

        This class method will create a gear element from geometry data.
        Properties are calculated as per :cite:`friswell2010dynamics`, appendix 1
        for a hollow cylinder:

        Mass:

        :math:`m = \\rho \\pi w (d_o^2 - d_i^2) / 4`

        Polar moment of inertia:

        :math:`I_p = m (d_o^2 + d_i^2) / 8`

        Diametral moment of inertia:

        :math:`I_d = \\frac{1}{2} I_p + \\frac{1}{12} m w^2`

        Where :math:`\\rho` is the material density, :math:`w` is the gear width,
        :math:`d_o` is the outer diameter and :math:`d_i` is the inner diameter.

        Parameters
        ----------
        n : int
            Node in which the gear will be inserted.
        material: ross.Material
            Gear material.
        width : float, pint.Quantity
            The face width of the gear considering that the gear body has the
            same thickness (m).
        i_d : float, pint.Quantity
            Inner diameter (the diameter of the shaft on which the gear is mounted).
        o_d : float, pint.Quantity
            Outer pitch diameter (m).
        n_teeth: int, pint.Quantity
            Number of teeth of the gear.
        pressure_angle : float, pint.Quantity, optional
            The normal pressure angle of the gear (rad).
            Default is 20 deg (converted to rad).
        tag : str, optional
            A tag to name the element.
            Default is None.
        helix_angle: float, pint.Quantity, optional
            Value of helix angle for helical gears. Default is 0 representing spur gear.
        scale_factor: float, optional
            The scale factor is used to scale the gear drawing.
            Default is 1.
        color : str, optional
            A color to be used when the element is represented.
            Default is 'Goldenrod'.

        Attributes
        ----------
        m : float
            Mass of the gear element.
        Id : float
            Diametral moment of inertia.
        Ip : float
            Polar moment of inertia.

        Examples
        --------
        >>> from ross.materials import steel
        >>> gear = GearElement.from_geometry(0, steel, 0.07, 0.05, 0.28, 50)
        >>> gear.base_radius # doctest: +ELLIPSIS
        0.131556...
        >>>
        """
        m = material.rho * np.pi * width * (o_d**2 - i_d**2) / 4
        Ip = m * (o_d**2 + i_d**2) / 8
        Id = 1 / 2 * Ip + 1 / 12 * m * width**2

        if pressure_angle is None:
            pressure_angle = Q_(20, "deg")

        return cls(
            n,
            m,
            Id,
            Ip,
            width=width,
            n_teeth=n_teeth,
            pitch_diameter=o_d,
            pressure_angle=pressure_angle,
            material=material,
            tag=tag,
            helix_angle=helix_angle,
            scale_factor=scale_factor,
            color=color,
        )

    def _patch(self, position, fig):
        """Gear element patch.

        Patch that will be used to draw the gear element using plotly library.

        Parameters
        ----------
        position : float
            Position in which the patch will be drawn.
        fig : plotly.graph_objects.Figure
            The figure object which traces are added on.

        Returns
        -------
        fig : plotly.graph_objects.Figure
            The figure object which traces are added on.
        """

        zpos, ypos, yc_pos, scale_factor = position
        scale_factor *= 2
        radius = min(self.base_radius * 1.1 + 0.05, 1)

        z_upper = [
            zpos + scale_factor / 25,
            zpos + scale_factor / 25,
            zpos - scale_factor / 25,
            zpos - scale_factor / 25,
        ]
        y_upper = [ypos, ypos + radius, ypos + radius, ypos]

        z_lower = [
            zpos + scale_factor / 25,
            zpos + scale_factor / 25,
            zpos - scale_factor / 25,
            zpos - scale_factor / 25,
        ]
        y_lower = [-ypos, -ypos - radius, -ypos - radius, -ypos]

        z_pos = z_upper
        z_pos.append(None)
        z_pos.extend(z_lower)

        y_pos = y_upper
        y_upper.append(None)
        y_pos.extend(y_lower)

        customdata = [self.n, self.Ip, self.Id, self.m, self.base_radius * 2]
        hovertemplate = (
            f"Gear Node: {customdata[0]}<br>"
            + f"Polar Inertia: {customdata[1]:.3e}<br>"
            + f"Diametral Inertia: {customdata[2]:.3e}<br>"
            + f"Gear Mass: {customdata[3]:.3f}<br>"
            + f"Gear Base Diam.: {customdata[4]:.3f}<br>"
        )

        fig.add_trace(
            go.Scatter(
                x=z_pos,
                y=[y + yc_pos if y is not None else None for y in y_pos],
                customdata=[customdata] * len(z_pos),
                text=hovertemplate,
                mode="lines",
                fill="toself",
                fillcolor=self.color,
                fillpattern=dict(
                    shape="/", fgcolor="rgba(0, 0, 0, 0.2)", bgcolor=self.color
                ),
                opacity=0.8,
                line=dict(width=2.0, color="rgba(0, 0, 0, 0.2)"),
                showlegend=False,
                name=self.tag,
                legendgroup="gears",
                hoveron="points+fills",
                hoverinfo="text",
                hovertemplate=hovertemplate,
                hoverlabel=dict(bgcolor=self.color),
            )
        )

        return fig
