import numpy as np
from typing import Tuple
from __env__ import H, H_Z, SIGMA_0, GAIT_SETTINGS
from numba import jit
import numba


f_0 = GAIT_SETTINGS["FREQ_BASE"] 
SIGMA_0 = SIGMA_0
H_Z = H_Z


#@jit
def FTG(sigma_i_0: float, t: float, f_i: float) -> Tuple[np.array, float]:
    """
        Generates a vector in R^3 representing the desired foot position 
        (end efector) in the H_i frame corresponding to the robots i-th leg 
        horizontal frame below its hip.
        
        Arguments:
        ----------
            sigma_i_0 : float 
                Contact phase.

            t  : float 
                Timestep.

            f_i: float 
                i-th leg frequency offset (from NN policy).
            
        Returns:
        -------
            numpy.ndarray, shape (3,) 
                Vector expresed in the i-th leg horizontal frame Hi,
                representing the target foot position.

            float
                FTG frequency. This output is used as an input of the neural 
                network
    """

    sigma_i = (sigma_i_0 + t * (f_0 + f_i)) % (2 * np.pi)
    k = 2 * (sigma_i - np.pi) / np.pi
    param = 0 # 0.5
    h= H*0.8

    position = (h * (-2 * k ** 3 + 3 * k ** 2) - param) * (0 < k < 1) + \
                (h * (2 * k ** 3 - 9 * k ** 2 + 12 * k - 4) - param) *\
                (1 < k < 2) - param * (k > 2 or k < 0)
    
    return position * H_Z, sigma_i

#@jit
def compute_foot_trajectories(t, frequencies)-> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    
    """
    Compute the foot trajectories for the given frequencies, for all the
    four legs.
    """

    target_foot_positions = np.zeros((4, 3))
    FTG_frequencies = np.zeros(4)
    FTG_phases = np.zeros((4,2))

    for i in range(4):
        r, sigma_i = FTG(SIGMA_0[i], t, frequencies[i])
    
        FTG_frequencies[i] = f_0 + frequencies[i]
        FTG_phases[i] = [np.sin(sigma_i), np.cos(sigma_i)]
        target_foot_positions[i] = r

    return target_foot_positions, FTG_frequencies, FTG_phases