from abc import ABC, abstractmethod
from enum import Enum
import numpy as np
from typing import Tuple, Optional

class ArrayType(Enum):
    """Enumeration of supported antenna array types."""
    ULA = 'ula'
    URA = 'ura'
    UCA = 'uca'
    NESTED = 'nested'
    COPRIME = 'coprime'
    MRA = 'mra'
    CYLINDRICAL = "cylindrical"
    CUSTOM = 'custom'

class BaseArray(ABC):
    """Abstract base class for antenna arrays.

    Attributes:
        positions (np.ndarray): Physical coordinates of each element, shape (dim, M).
            - dim: Spatial dimension of the array (1 for linear, 2 for planar, etc.)
            - M: Number of antenna elements
        type (ArrayType): Array geometry type (e.g., ULA, URA, UCA, etc.).
        spacing (Optional[float]): Inter-element spacing, if uniformly defined.
        dim (int): Spatial dimension, inferred from positions.shape[0].
        count (int): Total number of elements, inferred from positions.shape[1].
        indices (Optional[np.ndarray]): Structured grid indices, shape (dim, M), used in structured arrays like URA.
        position_errors (np.ndarray): Perturbations in element positions, same shape as `positions`.
        gain_errors (np.ndarray): Per-element gain variation factors (linear scale), same shape as `positions`.
        phase_errors (np.ndarray): Per-element phase offsets (in radians), same shape as `positions`.
    """

    def __init__(
        self,
        positions: np.ndarray,
        type: ArrayType,
        spacing: Optional[float] = None,
        indices: Optional[np.ndarray] = None,
        position_errors: Optional[np.ndarray] = None,
        gain_errors: Optional[np.ndarray] = None,
        phase_errors: Optional[np.ndarray] = None
    ):
        """Initialize a base antenna array.

        Args:
            positions (np.ndarray): Element positions, shape (dim, M), where:
                - dim is the spatial dimension (1, 2, or 3),
                - M is the number of elements.
            type (ArrayType): Type of the array structure (e.g., ArrayType.ULA, ArrayType.URA).
            spacing (Optional[float]): Inter-element spacing, relevant only for structured arrays.
            indices (Optional[np.ndarray]): Grid-based indices, same shape as `positions`.
            position_errors (Optional[np.ndarray]): Position perturbations, same shape as `positions`.
            gain_errors (Optional[np.ndarray]): Gain perturbations (linear), same shape as `positions`.
            phase_errors (Optional[np.ndarray]): Phase perturbations (in radians), same shape as `positions`.

        Raises:
            ValueError: If input arrays have mismatched shapes or invalid dimensions.
        """

        # element positions
        positions = np.array(positions, dtype=float)
        if positions.ndim == 1:
            positions = positions.reshape(1, -1) # reshape to row vector
        if positions.shape[0] not in (1, 2, 3):
            raise ValueError(f"positions must be in 1D, 2D, or 3D space. Got dim={positions.shape[0]}.")
        self.positions = positions
        self.dim = positions.shape[0]
        self.count = positions.shape[1]

        # index
        if indices is not None:
            indices = np.array(indices)
            if indices.shape != positions.shape:
                raise ValueError(f"Shape mismatch: indices {indices.shape} vs positions {positions.shape}")
        self.indices = indices

        # errors and perturbations
        # TODO: mutual coupling needs to be modeled
        # TODO: if position error is provided, the dimension of the array should change w.s.t. the dimension of the error
        if position_errors is not None:
            position_errors = np.array(position_errors, dtype=float)
            if position_errors.shape != positions.shape:
                raise ValueError(f"Shape mismatch: position_errors {position_errors.shape} vs positions {positions.shape}")
            self.position_errors = position_errors
        else:
            self.position_errors = position_errors # None

        if gain_errors is not None:
            gain_errors = np.array(gain_errors, dtype=float)
            if gain_errors.shape != positions.shape:
                raise ValueError(f"Shape mismatch: gain_errors {gain_errors.shape} vs positions {positions.shape}")
            self.gain_errors = gain_errors
        else:
            self.gain_errors = gain_errors # None

        if phase_errors is not None:
            phase_errors = np.array(phase_errors, dtype=float)
            if phase_errors.shape != positions.shape:
                raise ValueError(f"Shape mismatch: phase_errors {phase_errors.shape} vs positions {positions.shape}")
            self.phase_errors = phase_errors
        else:
            self.phase_errors = phase_errors # None

        # other attributes
        self.spacing = spacing
        self.type = type

    @abstractmethod
    def steering_matrix(
        self,
        wavelength: float,
        directions: np.ndarray
    ) -> np.ndarray:
        """Computes the array steering matrix for specified direction(s).

        Args:
            wavelength (float): Wavelength of the incoming signal (same unit as element positions).
            directions (np.ndarray): Incident direction(s) in radians. Shape should be:
                - (dim,) for a single direction
                - (dim, D) for D directions

        Returns:
            np.ndarray: Complex-valued steering matrix of shape (M, D), where M is the number of elements
                and D is the number of directions. For a single direction, shape is (M,).

        Raises:
            ValueError: If `wavelength` is non-positive.
            ValueError: If `directions` shape is incompatible with array dimension.
        """
        pass

    @abstractmethod
    def rotate(
        self,
        target_normal: np.ndarray,
        current_normal: Optional[np.ndarray] = None,
        center: Optional[np.ndarray] = None
    ) -> "BaseArray":
        """Rotate the array such that its current normal vector aligns with a target direction.
        This method performs a rigid-body rotation around a given center (default: centroid):
        1. Translates the array such that the rotation center lies at the origin.
        2. Applies the rotation that maps `current_normal` to `target_normal`.
        3. Translates the array back to its centroid or specified center.

        Args:
            target_normal (np.ndarray): Target normal vector, shape (dim,). Must be non-zero.
            current_normal (Optional[np.ndarray]): Current normal vector, (shape (dim,).
                If None, defaults to the vector from centroid to origin.
            center (Optional[np.ndarray]): Optional rotation center, shape (dim,).
                If None, uses the centroid of the array.

        Returns:
            BaseArray: A new array instance with rotated element positions (and other attributes if applicable).

        Raises:
            ValueError: If `target_normal` or `current_normal` has incorrect shape or is near-zero in norm.
            NotImplementedError: If the method is not implemented in a concrete subclass.
        """
        pass

    def translate(self, offset: np.ndarray):
        """Translate the array to a new center by applying the given offset.

        Args:
            offset (np.ndarray): Translation vector of shape (dim,) or (1, dim).

        Raises:
            ValueError: If the dimension of the offset does not match the array's spatial dimension.
        """
        if offset.ndim == 1 and offset.shape[0] == self.dim:
            offset = np.reshape(offset, (-1, 1))
        elif offset.ndim == 2:
            if offset.shape == (1, self.dim):
                offset = np.reshape(offset, (self.dim, 1))
            elif offset.shape == (self.dim, 1):
                pass  # already correct shape
            else:
                raise ValueError(
                    f"Offset must be convertible to shape ({self.dim}, 1), "
                    f"but got shape {offset.shape}."
                )
        else:
            raise ValueError(
                f"Offset must be a 1D or 2D array convertible to shape ({self.dim}, 1), "
                f"but got shape {offset.shape}."
            )
        self.positions += offset

    def center(self):
        """Center the array at the origin."""
        centroid = np.mean(self.positions, axis=1)
        self.translate(-centroid)

class Array1D(BaseArray):
    """
    One-dimensional linear array.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if self.dim != 1:
            raise ValueError(f"Array1D requires 1D array. Got dimension = {self.dim}.")

    def steering_matrix(
        self,
        wavelength: float,
        directions: np.ndarray,
        compute_derivative: bool = False
    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        """Compute the steering matrix of 1D array.

        Args:
            wavelength (float): Signal wavelength.
            directions (np.ndarray): Input directions, must be in one of the following formats:
                - shape (D,) or (1, D): D broadside angles (in radians), valid for 1D array
                - shape (2, D): D pairs of azimuth and elevation angles (in radians), used for projection to broadside
                The final broadside angle is computed as asin(sin(az) * cos(el)).
            compute_derivative (bool): Whether to compute the derivative of the steering matrix w.r.t. each DOA (only valid for 1D DOAs).

        Returns:
            A (np.ndarray): Steering matrix of shape (M, D)
            DA (np.ndarray or None): Derivative matrix (optional, only valid for 1D DOAs)

        Raises:
            ValueError: If dimension of directions is not 1 or 2.
        """
        directions = np.array(directions)
        if directions.ndim in (0, 1):
            directions = np.reshape(directions, (1, -1))
        elif directions.ndim == 2:
            if directions.shape[0] != 2:
                raise ValueError(
                    f"If directions is 2D, its shape must be (2, D), "
                    f"but got shape {directions.shape}."
                )
            az, el = directions[0, :], directions[1, :]
            directions = np.arcsin(np.sin(az) * np.cos(el))
        else:
            raise ValueError(
                f"Invalid shape for directions ({directions.ndim}). "
                "Expected shape (D,) or (1, D) for 1D broadside angles or shape (2, D) for azimuth/elevation."
            )

        DA = None
        positions = self.positions

        k = 2 * np.pi / wavelength
        phase = k * (np.outer(positions, np.sin(directions)))
        A = np.exp(1j * phase)
        if compute_derivative:
            if directions.ndim != 1:
                raise ValueError("compute_derivative is only valid for 1D DOAs (broadside angles).")
            DA = 1j * 2 * np.pi / wavelength * np.outer(positions, np.cos(directions)) * A

        # apply gain/phase errors
        if self.gain_errors is not None:
            A_gain = self.gain_errors.flatten()[:, None]
            A *= A_gain
            if DA is not None:
                DA *= A_gain
        if self.phase_errors is not None:
            A_phase = np.exp(1j * self.phase_errors.flatten())[:, None]
            A *= A_phase
            if DA is not None:
                DA *= A_phase

        return A, DA

    def rotate(
        self,
        target_normal: np.ndarray,
        current_normal: Optional[np.ndarray] = None,
        center: Optional[np.ndarray] = None
    ) -> Tuple[BaseArray, np.ndarray]:
        """
        Rotate a 1D array to align its normal vector with a target direction.

        The array is assumed to be initially aligned along the y-axis (in 2D) or y-direction
        in the xOy plane (in 3D). The method performs a rigid-body rotation to align the
        array's normal with the specified `target_normal`.

        The rotation is performed in three steps:
            1. Translate the centroid of the array to the origin.
            2. Apply the rotation based on the angle between current and target normal vectors.
            3. Translate the array back to the specified or original centroid.

        Args:
            target_normal (np.ndarray): Target direction to align the array's normal to.
                Must be a 2D or 3D vector of shape (2,) or (3,).
            current_normal (Optional[np.ndarray]): Original normal direction of the array.
                Defaults to [1, 0] in 2D and [0, 0, 1] in 3D.
            center (Optional[np.ndarray]): Rotation center, defaults to the array's centroid.
                Must be the same dimension as `target_normal`.

        Returns:
            Tuple[BaseArray, np.ndarray]:
                - A new rotated array object (Array2D or Array3D) with updated element positions.
                - The rotation matrix used in the transformation (2x2 or 3x3).

        Raises:
            ValueError: If `target_normal` is not 2D or 3D.
            ValueError: If `center` is provided but its dimension does not match `target_normal`.
        """

        target_normal = (np.array(target_normal, dtype=float)).flatten()
        target_dim = target_normal.shape[0]
        if target_dim not in (2, 3):
            raise ValueError(f"Target normal must be a 2D or 3D vector, got shape {target_normal.shape}.")

        positions = self.positions
        elem_num = self.count

        # Determine rotation center
        if center is None:
            center = np.mean(positions, axis=1)
        center = np.reshape(np.array(center, dtype=float), (-1, 1))
        if center.shape[0] != positions.shape[0]:
            raise ValueError(
                f"Array center and element positions must have the same dimension. "
                f"Got center dim = {center.shape[0]}, positions dim = {positions.shape[0]}."
            )

        # Centered positions aligned along y-axis
        if target_dim == 2:
            # Arrange elements along Y-axis and center at origin
            centered_pos = np.vstack((np.zeros((1, elem_num)), positions - center))
            # Get normal vector
            if current_normal is None:
                current_normal = np.array([1.0, 0.0]) # assume
            else:
                current_normal = (np.array(current_normal, dtype=float)).flatten()
                current_dim = current_normal.shape[0]
                if current_dim != target_dim:
                    raise ValueError(f"Dimension of target normal ({target_dim}) and current normal ({current_dim}) must match.")
            current_normal = current_normal / np.linalg.norm(current_normal)
            target_normal = target_normal / np.linalg.norm(target_normal)
            # Get the angle between current and target normal vector
            angle = np.arctan2(target_normal[1], target_normal[0]) - np.arctan2(current_normal[1], current_normal[0])
            rot_mat = np.array([
                [np.cos(angle), -np.sin(angle)],
                [np.sin(angle),  np.cos(angle)]
            ])
            # expand the dimension of rotation center
            center_exp = np.vstack([
                np.zeros((1, 1)),
                center.reshape(1, 1)
            ])
            # perform rotation and translation
            positions_rot = rot_mat @ centered_pos + center_exp
            return Array2D(
                positions = positions_rot,
                type = self.type,
                spacing = self.spacing,
                indices = self.indices,
                position_errors = self.position_errors,
                gain_errors = self.gain_errors,
                phase_errors = self.phase_errors
            ), rot_mat
        elif target_dim == 3:
            # Arrange elements along y-axis and center at origin
            centered_pos = np.vstack((
                np.zeros((1, elem_num)),
                positions - center,
                np.zeros((1, elem_num))
            ))
            # Get normal vector
            if current_normal is None:
                current_normal = np.array([0, 0, 1])
            else:
                current_normal = (np.array(current_normal, dtype=float)).flatten()
                current_dim = current_normal.shape[0]
                if current_dim != target_dim:
                    raise ValueError(f"Dimension of target normal ({target_dim}) and current normal ({current_dim}) must match.")
            current_normal = current_normal / np.linalg.norm(current_normal)
            target_normal = target_normal / np.linalg.norm(target_normal)
            # Rotate to a 3D direction using axis-angle method
            axis = np.cross(current_normal, target_normal)
            norm_axis = np.linalg.norm(axis)
            if norm_axis < 1e-8:
                # If vectors are parallel or anti-parallel
                if np.dot(current_normal, target_normal) > 0:
                    rot_mat = np.eye(3) # No rotation
                else:
                    rot_mat = np.diag([-1, -1, 1]) # 180-degree rotation
            else:
                axis = axis / norm_axis
                angle = np.arccos(np.dot(current_normal, target_normal))
                K = np.array([
                    [0, -axis[2], axis[1]],
                    [axis[2], 0, -axis[0]],
                    [-axis[1], axis[0], 0]
                ])
                rot_mat = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * K @ K

            # expand the dimension of rotation center
            center_exp = np.vstack([
                np.zeros((1, 1)),
                center.reshape(1, 1),
                np.zeros((1, 1))
            ])
            # perform rotation and translation
            positions_rot = rot_mat @ centered_pos + center
            return Array3D(
                positions = positions_rot,
                type = self.type,
                spacing = self.spacing,
                indices = self.indices,
                position_errors = self.position_errors,
                gain_errors = self.gain_errors,
                phase_errors = self.phase_errors
            ), rot_mat
        else:
            raise ValueError("Invalid dimension")

class Array2D(BaseArray):
    """
    Two-dimensional planar array.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if self.dim != 2:
            raise ValueError(f"Array2D requires 2D array. Got dimension = {self.dim}.")

    def steering_matrix(
        self,
        wavelength: float,
        directions: np.ndarray,
    ) -> np.ndarray:
        """Compute the steering matrix of 2D array.

        Args:
            wavelength (float): Signal wavelength.
            directions (np.ndarray): Input directions, must be in one of the following formats:
                - shape (2, D): D pairs of azimuth and elevation angles (in radians), used for projection to broadside

        Returns:
            A (np.ndarray): Steering matrix of shape (M, D)

        Raises:
            ValueError: If dimension of directions is 2.
        """
        directions = np.array(directions)

        if directions.ndim == 2:
            if directions.shape[0] != 2:
                raise ValueError(
                    f"If directions is 2D, its shape must be (2, D), "
                    f"but got shape {directions.shape}."
                )
            az, el = directions[0, :], directions[1, :]
            cos_el = np.cos(el)
            x_proj = cos_el * np.cos(az)
            y_proj = cos_el * np.sin(az)
        else:
            raise ValueError(
                f"Invalid shape for directions ({directions.ndim}). "
                "Expected shape (2, D) for azimuth/elevation."
            )

        positions = self.positions

        # Compute phase shifts
        k = 2 * np.pi / wavelength
        phase = k * (positions[0, :, None] * x_proj + positions[1, :, None] * y_proj)
        A = np.exp(1j * phase)

        # apply gain/phase errors
        if self.gain_errors is not None:
            A *= self.gain_errors.flatten()
        if self.phase_errors is not None:
            A *= np.exp(1j * self.phase_errors.flatten())

        return A

    def rotate(
        self,
        target_normal: np.ndarray,
        current_normal: Optional[np.ndarray] = None,
        center: Optional[np.ndarray] = None
    ) -> BaseArray:
        """
        Rotate a 2D array to align its normal vector with a target direction.

        The array is assumed to be initially aligned along xOy plane (in 3D). The method performs a rigid-body rotation to align the
        array's normal with the specified `target_normal`.

        The rotation is performed in three steps:
            1. Translate the centroid of the array to the origin.
            2. Apply the rotation based on the angle between current and target normal vectors.
            3. Translate the array back to the specified or original centroid.

        Args:
            target_normal (np.ndarray): Target direction to align the array's normal to.
                Must be a 3D vector of shape (3,).
            current_normal (Optional[np.ndarray]): Original normal direction of the array.
                Defaults to [0, 0, 1] in 3D.
            center (Optional[np.ndarray]): Rotation center, defaults to the array's centroid.
                Must be the same dimension as `target_normal`.

        Returns:
            Tuple[BaseArray, np.ndarray]:
                - A new rotated array object (Array3D) with updated element positions.
                - The rotation matrix used in the transformation (3x3).

        Raises:
            ValueError: If `target_normal` is not 3D.
            ValueError: If `center` is provided but its dimension does not match `target_normal`.
        """

        target_normal = (np.array(target_normal, dtype=float)).flatten()
        target_dim = target_normal.shape[0]
        if target_dim != 3:
            raise ValueError(f"Target normal must be a 3D vector for 2D array, got shape {target_normal.shape}.")

        positions = self.positions
        elem_num = self.count

        # Determine rotation center
        if center is None:
            center = np.mean(positions, axis=1)
        center = np.reshape(np.array(center, dtype=float), (-1, 1))
        if center.shape[0] != positions.shape[0]:
            raise ValueError(
                f"Array center and element positions must have the same dimension. "
                f"Got center dim = {center.shape[0]}, positions dim = {positions.shape[0]}."
            )

        # Arrange elements along xOy plane and center at origin
        centered_pos = np.vstack((
            positions - center,
            np.zeros((1, elem_num))
        ))
        # Get normal vector
        if current_normal is None:
            current_normal = np.array([0, 0, 1])
        else:
            current_normal = (np.array(current_normal, dtype=float)).flatten()
            current_dim = current_normal.shape[0]
            if current_dim != target_dim:
                raise ValueError(f"Dimension of target normal ({target_dim}) and current normal ({current_dim}) must match.")
        current_normal = current_normal / np.linalg.norm(current_normal)
        target_normal = target_normal / np.linalg.norm(target_normal)
        # Rotate to a 3D direction using axis-angle method
        axis = np.cross(current_normal, target_normal)
        norm_axis = np.linalg.norm(axis)
        if norm_axis < 1e-8:
            # If vectors are parallel or anti-parallel
            if np.dot(current_normal, target_normal) > 0:
                rot_mat = np.eye(3) # No rotation
            else:
                rot_mat = np.diag([-1, -1, 1]) # 180-degree rotation
        else:
            axis = axis / norm_axis
            angle = np.arccos(np.dot(current_normal, target_normal))
            K = np.array([
                [0, -axis[2], axis[1]],
                [axis[2], 0, -axis[0]],
                [-axis[1], axis[0], 0]
            ])
            rot_mat = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * K @ K
        # expand the dimension of rotation center
        center_exp = np.vstack([
            center.reshape(2, 1),
            np.zeros((1, 1))
        ])
        # perform rotation and translation
        positions_rot = rot_mat @ centered_pos + center_exp
        return Array3D(
            positions = positions_rot,
            type = self.type,
            spacing = self.spacing,
            indices = self.indices,
            position_errors = self.position_errors,
            gain_errors = self.gain_errors,
            phase_errors = self.phase_errors
        ), rot_mat

class Array3D(BaseArray):
    """
    Three-dimensional volumetric array.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        assert self.dim == 3, "Array3D requires 3D element positions."

    def steering_matrix(
        self,
        wavelength: float,
        directions: np.ndarray,
    ) -> np.ndarray:
        """Compute the steering matrix of 2D array.

        Args:
            wavelength (float): Signal wavelength.
            directions (np.ndarray): Input directions, must be in one of the following formats:
                - shape (2, D): D pairs of azimuth and elevation angles (in radians), used for projection to broadside

        Returns:
            A (np.ndarray): Steering matrix of shape (M, D)

        Raises:
            ValueError: If dimension of directions is 2.
        """
        directions = np.array(directions)

        if directions.ndim == 2:
            if directions.shape[0] != 2:
                raise ValueError(
                    f"If directions is 2D, its shape must be (2, D), "
                    f"but got shape {directions.shape}."
                )
            az, el = directions[0, :], directions[1, :]
            cos_el = np.cos(el)
            x_proj = cos_el * np.cos(az)
            y_proj = cos_el * np.sin(az)
            z_proj = np.sin(el)
        else:
            raise ValueError(
                f"Invalid shape for directions ({directions.ndim}). "
                "Expected shape (2, D) for azimuth/elevation."
            )

        positions = self.positions

        # Compute phase shifts
        k = 2 * np.pi / wavelength
        phase = k * (positions[0, :, None] * x_proj + positions[1, :, None] * y_proj
                     + positions[2, :, None] * z_proj)
        A = np.exp(1j * phase)

        # apply gain/phase errors
        if self.gain_errors is not None:
            A *= self.gain_errors.flatten()
        if self.phase_errors is not None:
            A *= np.exp(1j * self.phase_errors.flatten())

        return A

    def rotate(
        self,
        target_normal: np.ndarray,
        current_normal: Optional[np.ndarray] = None,
        center: Optional[np.ndarray] = None
    ) -> BaseArray:
        """
        Rotate a 3D array to align its normal vector with a target direction.

        The method performs a rigid-body rotation to align the array's normal with the specified `target_normal`.

        The rotation is performed in three steps:
            1. Translate the centroid of the array to the origin.
            2. Apply the rotation based on the angle between current and target normal vectors.
            3. Translate the array back to the specified or original centroid.

        Args:
            target_normal (np.ndarray): Target direction to align the array's normal to.
                Must be a 3D vector of shape (3,).
            current_normal (Optional[np.ndarray]): Original normal direction of the array.
                Defaults to [0, 0, 1] in 3D.
            center (Optional[np.ndarray]): Rotation center, defaults to the array's centroid.
                Must be the same dimension as `target_normal`.

        Returns:
            Tuple[BaseArray, np.ndarray]:
                - A new rotated array object (Array3D) with updated element positions.
                - The rotation matrix used in the transformation (3x3).

        Raises:
            ValueError: If `target_normal` is not 3D.
            ValueError: If `center` is provided but its dimension does not match element `positions`.
        """

        target_normal = (np.array(target_normal, dtype=float)).flatten()
        target_dim = target_normal.shape[0]
        if target_dim != 3:
            raise ValueError(f"Target normal must be a 3D vector for 2D array, got shape {target_normal.shape}.")

        positions = self.positions
        elem_num = self.count

        # Determine rotation center
        if center is None:
            center = np.mean(positions, axis=1)
        center = np.reshape(np.array(center, dtype=float), (-1, 1))
        if center.shape[0] != positions.shape[0]:
            raise ValueError(
                f"Array center and element positions must have the same dimension. "
                f"Got center dim = {center.shape[0]}, positions dim = {positions.shape[0]}."
            )

        # Arrange elements along xOy plane and center at origin
        centered_pos = positions - center
        # Get normal vector
        if current_normal is None:
            current_normal = np.array([0, 0, 1])
        else:
            current_normal = (np.array(current_normal, dtype=float)).flatten()
            current_dim = current_normal.shape[0]
            if current_dim != target_dim:
                raise ValueError(f"Dimension of target normal ({target_dim}) and current normal ({current_dim}) must match.")
        current_normal = current_normal / np.linalg.norm(current_normal)
        target_normal = target_normal / np.linalg.norm(target_normal)
        # Rotate to a 3D direction using axis-angle method
        axis = np.cross(current_normal, target_normal)
        norm_axis = np.linalg.norm(axis)
        if norm_axis < 1e-8:
            # If vectors are parallel or anti-parallel
            if np.dot(current_normal, target_normal) > 0:
                rot_mat = np.eye(3) # No rotation
            else:
                rot_mat = np.diag([-1, -1, 1]) # 180-degree rotation
        else:
            axis = axis / norm_axis
            angle = np.arccos(np.dot(current_normal, target_normal))
            K = np.array([
                [0, -axis[2], axis[1]],
                [axis[2], 0, -axis[0]],
                [-axis[1], axis[0], 0]
            ])
            rot_mat = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * K @ K
        # expand the dimension of rotation center
        center_exp = center.reshape(3, 1)
        # perform rotation and translation
        positions_rot = rot_mat @ centered_pos + center_exp
        return Array3D(
            positions = positions_rot,
            type = self.type,
            spacing = self.spacing,
            indices = self.indices,
            position_errors = self.position_errors,
            gain_errors = self.gain_errors,
            phase_errors = self.phase_errors
        ), rot_mat