__doc__ = """
Actuation tools
"""
__all__ = ['ApplyActuations', 'ContinuousActuation']


import numpy as np
from numba import njit

from elastica._linalg import _batch_cross
from elastica._calculus import quadrature_kernel, difference_kernel
from elastica.external_forces import inplace_addition
from elastica.external_forces import NoForces

from coomm._rod_tool import _lab_to_material, _material_to_lab, average2D

@njit(cache=True)
def _internal_to_external_load(
    director_collection, kappa, tangents,
    rest_lengths, rest_voronoi_lengths,
    dilatation, voronoi_dilatation,
    internal_force, internal_couple,
    external_force, external_couple
    ):

    external_force[:, :] = difference_kernel(
        _material_to_lab(director_collection, internal_force)
        )

    external_couple[:, :] = (
        difference_kernel(internal_couple) +
        quadrature_kernel(
            _batch_cross(kappa, internal_couple) * rest_voronoi_lengths
        ) +
        _batch_cross(
                _lab_to_material(director_collection, tangents * dilatation),
                internal_force
            )  * rest_lengths
    )

@njit(cache=True)
def _force_induced_couple(internal_forces, distance, internal_couples):
    internal_couples[:, :] = average2D(
        _batch_cross(distance, internal_forces)
    )

class ContinuousActuation:
    """ 
    Classes inherited from this base class should contain parameters and
    functions that describe and calculate forces / couples generated by
    actuators.
    """

    def __init__(self, n_elements: int, **kwargs):
        super().__init__()
        self.n_elements = n_elements
        self.internal_force = np.zeros((3, n_elements))       # material frame
        self.external_force = np.zeros((3, n_elements+1))     # lab frame
        self.internal_couple = np.zeros((3, n_elements-1))    # material frame
        self.external_couple = np.zeros((3, n_elements))      # material frame

    def reset_actuation(self,):
        """
        Set actuation force/torques to be zero.
        """
        self.internal_force[:, :] *= 0
        self.external_force[:, :] *= 0
        self.internal_couple[:, :] *= 0
        self.external_couple[:, :] *= 0

class ApplyActuations(NoForces):
    """ApplyActuations
    """

    def __init__(self, actuations, step_skip: int, callback_params_list: list):
        """
        TODO : need documentation on how to initialize

        Parameters
        ----------
        actuations :
        step_skip : int
        callback_params_list : Dictionary[list]
        """
        self.current_step = 0
        self.actuations = actuations
        self.every = step_skip
        self.callback_params_list = callback_params_list

    def apply_torques(self, system, time: np.float64 = 0.0):
        """apply_torques.

        Parameters
        ----------
        system :
        time : np.float64
        """
        for actuation in self.actuations:
            actuation(system)
            inplace_addition(
                system.external_forces, actuation.external_force
            )
            inplace_addition(
                system.external_torques, actuation.external_couple
            )

        self.make_callback()

    def make_callback(self):
        """make_callback.
        """
        if self.current_step % self.every == 0:
            self.callback_func(
                self.actuations,
                self.callback_params_list
            )
        self.current_step += 1

    def callback_func(self, actuations, callback_params_list):
        """callback_func.

        Parameters
        ----------
        actuations :
        callback_params_list :
        """
        for actuation, callback_params in zip(
            actuations, callback_params_list
        ):
            callback_params['internal_force'].append(
                actuation.internal_force.copy()
            )
            callback_params['internal_couple'].append(
                actuation.internal_couple.copy()
            )
            callback_params['external_force'].append(
                actuation.external_forces.copy()
            )
            callback_params['external_couple'].append(
                actuation.external_couple.copy()
            )
