import numpy as np
import sympy as sp


def eval_func(func: sp.Expr, sub_1: sp.Expr, sub_2: sp.Expr, grid: np.ndarray) -> np.ndarray:
    # return np.array([[float(func.subs(sub_1, x_).subs(sub_2, y_)) for x_, y_ in line] for line in grid])
    func = sp.lambdify([sub_1, sub_2], func, 'numpy')
    return func(grid[:, :, 0], grid[:, :, 1])


def get_orientation_phase_grid(step_phase: float, step_orientation: float) -> np.ndarray:
    """
    Returns a grid of x and y values for plotting.
    :param step_phase: step for the phase (phi) - in degrees
    :param step_orientation: step for the orientation (theta) - in degrees
    :return: numpy array of shape (n_orientation, n_phase). Each element is a tuple (theta, phi)
    """
    # phase <-> phi
    # orientation <-> theta
    step_phase *= np.pi / 180
    step_orientation *= np.pi / 180
    phi = np.arange(0, 2 * np.pi, step_phase)
    theta = np.arange(0, np.pi, step_orientation)
    return np.array(np.meshgrid(theta, phi)).T.reshape(-1, len(phi), 2)


def get_spatial_grid(step_x: float, step_y: float, size: float = 1) -> np.ndarray:
    """
    Returns a grid of x and y values for plotting.
    :param step_x: step for the x-coordinate
    :param step_y: step for the y-coordinate
    :param size: size of the grid
    :return: numpy array of shape (2 * size / step_x, 2 * size / step_y). Each element is a tuple (x, y)
    """
    x = np.arange(-size, size, step_x)
    y = np.arange(-size, size, step_y)
    return np.array(np.meshgrid(x, y)).T.reshape(-1, len(x), 2)