# Copyright (C) 2024- Davide Mollica <davide.mollica@inaf.it>
# SPDX-License-Identifier: GPL-3.0-or-later
#
# This file is part of iactsim.
#
# iactsim is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# iactsim is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with iactsim.  If not, see <https://www.gnu.org/licenses/>.

import numpy as np

def pointing_dir(altitude: float, azimuth: float) -> np.ndarray:
    """Telescope pointing direction in the corsika-like local reference frame.

    *   North -> (0, 0) : x-axis
    *   East -> (0, 90)
    *   South -> (0, 180)
    *   West -> (0, 270) : y-axis
    *   Zenith -> (90, any) : z-axis

    Parameters
    ----------
    altitude : float
        Altitude angle in degrees.
    azimuth : float
        Azimuth angle in degrees. Measured from North, increasing towards East (clockwise).

    Returns
    -------
    direction : ndarray, shape (3,)
        Pointing unit vector in the local reference frame (North-East-Up).
    """
    azimuth_rad = np.deg2rad(azimuth)
    altitude_rad = np.deg2rad(altitude)
    return np.array([
        np.cos(altitude_rad) * np.cos(azimuth_rad),
        -np.cos(altitude_rad) * np.sin(azimuth_rad),
        np.sin(altitude_rad),
    ])

def local_to_pointing_rotation(altitude: float, azimuth: float) -> np.ndarray:
    """Compute the rotation matrix to transform from the local reference frame (North-West-Up) to the "pointing reference frame".

    This reference frame is defined such that when the telescope is
    pointing at the horizon towards the North (i.e. alt: 0, az: 0), the axes are aligned as follows:

    *   z-axis: aligned with the telescope optical axis, pointing North.
    *   x-axis: perpendicular to the optical axis, pointing downward.
    *   y-axis: perpendicular to the optical axis, pointing west.

    Parameters
    ----------
    altitude : float
        Altitude angle of the telescope pointing, in degrees.
    azimuth : float
        Azimuth angle of the telescope pointing, in degrees. Measured from North, increasing towards East.

    Returns
    -------
    R : ndarray, shape (3, 3)
        Rotation matrix. When a vector in the local reference frame is multiplied by this matrix
        it is transformed into the telescope reference frame.

    Notes
    -----
    This is a simple rotation of the CORSIKA reference frame that move the up-vector (z) into the telescope direction.

    """
    theta = np.radians(90-altitude)
    phi = -np.radians(azimuth)

    ct = np.cos(theta)
    st = np.sin(theta)
    cp = np.cos(phi)
    sp = np.sin(phi)

    # pointing north (0,0)
    #  z axis is pointing north (local x axis)
    #  x axis is pointing down (opposite to local z axis)
    #  y axis is pointing west (local y axis)
    R = np.array([
        [ct*cp, ct*sp, -st],
        [  -sp,    cp,   0],
        [st*cp, st*sp,  ct]
    ])
    return R

def local_to_telescope_rotation(altitude: float, azimuth: float, custom_rotation=None) -> np.ndarray:
    """Compute the rotation matrix to transform from the local reference frame (North-West-Up) to the "telescope reference frame".
    By deault the telescope refrence frame is the pointing reference frame (see :py:func:`local_to_pointing_rotation`).

    Parameters
    ----------
    altitude : float
        Altitude angle of the telescope pointing, in degrees.
    azimuth : float
        Azimuth angle of the telescope pointing, in degrees. Measured from North, increasing towards East.

    Returns
    -------
    R : ndarray, shape (3, 3)
        Rotation matrix. When a vector in the local reference frame is multiplied by this matrix
        it is transformed into the telescope reference frame.

    Examples
    --------
    To rotate the telescope frame around the optical axis (Z-axis), while keeping Z pointing in the same direction:

        .. code-block:: python

            import numpy as np

            # Rotation around Z-axis
            rotation = np.array([
                [ 0, 1, 0], # new X is old Y
                [ 1, 0, 0], # new Y is old X
                [ 0, 0, 1] 
            ])

            # Get the telescope transformation matrix
            R_tel = local_to_telescope_rotation(
                altitude=70, 
                azimuth=180, 
                custom_rotation=rotation
            )
    
    Now when the telescope is pointing North:

    *   z-axis: aligned with the telescope optical axis, pointing North;
    *   y-axis: perpendicular to the optical axis, pointing downward;
    *   x-axis: perpendicular to the optical axis, pointing west.

    """
    # Rotation into the local reference frame
    R = local_to_pointing_rotation(altitude, azimuth)
    
    if custom_rotation is None:
        custom_rotation = np.eye(3)

    return custom_rotation @ R

def local_to_photon_rotation(u: np.ndarray) -> np.ndarray:
    """Compute the rotation matrix to transform vector `u` to the local vertical Z-axis (0,0,1).
    
    This uses the Rodrigues rotation formula.

    Parameters
    ----------
    u : ndarray, shape (3,)
        The source unit vector (e.g., photon direction).
    
    Returns
    -------
    R : ndarray, shape (3, 3)
        Rotation matrix such that R @ u = [0, 0, 1].
    """
    # Normalize input
    u = u / np.linalg.norm(u)
    
    # Define target (transform in a reference system where u = v)
    v = np.array([0.0, 0.0, 1.0])

    # Cosine of the angle between u and v
    c = np.dot(u, v)
    
    # Vectors are already parallel
    if c > 0.999999:
        return np.eye(3)
        
    # Vectors are anti-parallel
    # Rotate 180 degrees around any axis perpendicular to u.
    if c < -0.999999:
        # Pick a perpendicular axis
        # If u is not parallel to X, use cross(u, X). Else use cross(u, Y).
        if np.abs(u[0]) < 0.9:
            axis = np.cross(u, np.array([1, 0, 0]))
        else:
            axis = np.cross(u, np.array([0, 1, 0]))
        axis = axis / np.linalg.norm(axis)
        
        # Rodrigues formula for 180 degrees
        #   R = 2(axis * axis.T) - I
        return 2 * np.outer(axis, axis) - np.eye(3)

    ### Rodrigues rotation formula
    #   R = I + K + K^2 * (1 / (1 + c))
    # See for example https://math.stackexchange.com/a/476311
    
    # Rotation axis k = u x v
    k = np.cross(u, v)
    
    # Skew-symmetric cross-product matrix of k
    # [ 0  -k.z  k.y]
    # [ k.z  0  -k.x]
    # [-k.y  k.x  0 ]
    K = np.array([
        [   0., -k[2],  k[1]],
        [ k[2],    0., -k[0]],
        [-k[1],  k[0],    0.]
    ])
    
    R = np.eye(3) + K + (K @ K) * (1. / (1. + c))
    
    return R

def create_surface_rotation(normal: np.ndarray, up: np.ndarray = np.array([1, 0, 0])) -> np.ndarray:
    """
    Computes a rotation matrix to transform from the telescope reference frame to the surface reference frame.

    In the surface reference system:

    * the Z-axis (0, 0, 1) is aligned with the input ``normal``;
    * the Y-axis (0, 1, 0) is aligned as close as possible to the input ``up`` vector;
    * the X-axis (1, 0, 0) is mutually orthogonal.

    Parameters
    ----------
    normal : np.ndarray
        The vector in the telescope reference system that defines where the surface is facing. This will become the local Z-axis.
    up : np.ndarray, optional
        The vector in the telescope reference system that defines the "top" of the surface. 
        The local Y-axis will be constructed by projecting this vector onto the plane 
        perpendicular to the normal. Defaults to the X axis of telescope reference frame ([1, 0, 0]).

    Returns
    -------
    R : ndarray, shape (3, 3)
        The 3x3 rotation matrix. 
        
        * Usage: 
          
           - ``v_surf = R @ v_telescope`` if ``v_telescope`` has shape (3,)
           - ``v_surf = v_telescope @ R.T`` if ``v_telescope`` has shape (N,3)

    Notes
    -----

    The ID of each pixel of a SiPM tile is defined in the tile reference system as follows::

        ^ y
        |     _____   _____   _____
        |    |     | |     | |     |
        |    |  6  | |  7  | |  8  |
        |    |_____| |_____| |_____|
        |     _____   _____   _____
        |    |     | |     | |     |
        |    |  3  | |  4  | |  5  |
        |    |_____| |_____| |_____|
        |     _____   _____   _____
        |    |     | |     | |     |
        |    |  0  | |  1  | |  2  |
        |    |_____| |_____| |_____|
        |
        ------------------------------> x

    So, the SiPM tile rotation matrix must be carefully set to match the desired ID convention.
    
    Examples
    --------

        .. code-block:: python

            # Opposite to the pointing direction
            sensor_normal = np.array([0., 0., -1.])

            # Align the surface y axis with the telescope y axis
            R = create_surface_rotation(normal=sensor_normal, up=np.array([0., 1., 0.]))
        
    """
    # Surface Z axis is the normal
    z_axis = normal / np.linalg.norm(normal)
    
    # New X axis is perpendicular to both up and normal vectors
    x_axis = np.cross(up, z_axis)
    
    # Handle collinearity (if normal is parallel to up)
    if np.linalg.norm(x_axis) < 1e-6:
        raise(ValueError("Up vector and normal vector are parallel."))
    
    x_axis = x_axis / np.linalg.norm(x_axis)
    
    # New Y axis is perpendicular to Z and X
    y_axis = np.cross(z_axis, x_axis)
    
    # Construct rotation matrix
    R = np.array([
        x_axis,
        y_axis,
        z_axis
    ])
    
    return R