# 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 time
from typing import (
    Tuple, List, 
    Union, Optional
)

import cupy as cp
import numpy as np
from tqdm.auto import tqdm
import random

from .optics._optical_system import OpticalSystem
from .optics.utils._aspheric_coeffs import normalize_aspheric_coefficients
from .optics._cpu_transforms import local_to_telescope_rotation
from .optics._surface_misc import SurfaceShape, SensorSurfaceType
from .optics._atmosphere import AtmosphericTransmission
from .optics._materials import Materials
from .optics import ray_tracing

from .electronics import CherenkovCamera, CherenkovSiPMCamera

from .utils._timer import BenchmarkTimer

class IACT:
    """A class to represent an Imaging Atmospheric Cherenkov Telescope (IACT).

    This class serves as a wrapper for a CUDA-accelerated backend. It handles the conversion 
    of optical system and camera information into CuPy ndarrays (transferring data to the GPU), 
    and arranges this data in a format suitable for input to the CUDA kernels that perform the 
    actual simulation.

    Parameters
    ----------
    optical_system : OpticalSystem
        An instance of a class derived from :py:class:`OpticalSystem` representing the 
        telescope optical system.
    camera : CherenkovCamera
        An instance of a class derived from :py:class:`CherenkovCamera` representing the 
        telescope Cherenkov camera.
    position : Tuple[float, float, float] | List[float] | numpy.ndarray
        The Cartesian coordinates (x, y, z) representing the telescope 
        position in millimeters. Can be a tuple, list, or NumPy array.
    pointing : Tuple[float, float] | List[float] | numpy.ndarray
        The horizontal coordinates (altitude, azimuth) representing the 
        telescope pointing direction in degrees. Can be a tuple, list, 
        or NumPy array.    
    
    Raises
    ------
    TypeError
        
        - If optical_system is not an instance of :py:class:`OpticalSystem` or a derived class.
        - If camera is not an instance of :py:class:`CherenkovCamera` or a derived class.

    ValueError
        
        - If position does not have 3 elements.
        - If pointing does not have 2 elements.

    """

    def __init__(
        self,
        optical_system: OpticalSystem,
        camera: CherenkovCamera = None,
        position: Tuple[float, float, float] | List[float] | np.ndarray = (0.,0.,0.),
        pointing: Tuple[float, float] | List[float] | np.ndarray = (0,0),
    ):
        if not isinstance(optical_system, OpticalSystem):
            raise TypeError("optical_system must be an instance of OpticalSystem or a derived class")

        # Validate and convert position to NumPy array
        if not isinstance(position, (tuple, list, np.ndarray)) or len(position) != 3:
            raise ValueError("position must be a tuple, list, or NumPy array of length 3")

        self._cuda_tracing_args = [None] * 5 # ray-tracing kernel related arguments
        self._cuda_atmosphere_args = [None] * 5 # atmosphere transmission related arguments
        self._cuda_materials_args = [None] * 5 # materials refractive index args
        
        self._position = np.asarray(position)

        self._custom_rotation = np.eye(3)

        self.pointing = pointing

        self._optical_system = optical_system
        self._camera = camera

        self.atmospheric_transmission = AtmosphericTransmission()

        # Benchmark
        self.timer = BenchmarkTimer()
        self.timer.add_section('simulate_response')

        self.show_progress = False

        self._blocksize = 128 # Threads per block

    @property
    def position(self) -> np.ndarray:
        """
        numpy.ndarray: The telescope position in Cartesian coordinates (x, y, z) 
                    in millimeters.
        """
        return self._position

    @position.setter
    def position(self, value: Tuple[float, float, float] | List[float] | np.ndarray):
        if not isinstance(value, (tuple, list, np.ndarray)) or len(value) != 3:
            raise ValueError("position must be a tuple, list, or NumPy array of length 3")
        self._position = np.asarray(value)
        self._cuda_telescope_init()

    @property
    def pointing(self) -> np.ndarray:
        """
        numpy.ndarray: The telescope pointing direction in horizontal 
                    coordinates (altitude, azimuth) in degrees.
        """
        return self._pointing

    @pointing.setter
    def pointing(self, value: Tuple[float, float] | List[float] | np.ndarray):
        if not isinstance(value, (tuple, list, np.ndarray)) or len(value) != 2:
            raise ValueError("pointing must be a tuple, list, or NumPy array of length 2")
        self._pointing = np.asarray(value)
        self._altitude = value[0]
        self._azimuth = value[1]
        self._cuda_telescope_init()
    
    @property
    def altitude(self) -> float:
        """
        float: Telescope pointing altitude (0-90).
        """
        return self._altitude
    
    @altitude.setter
    def altitude(self, value: float):
        if value < 0 or value > 90:
            raise(ValueError('Telescope pointing altitude must be between 0 (horizon) and 90 (zenith) degrees.'))
        self._altitude = value
        self._pointing[0] = value
        self._cuda_telescope_init()

    @property
    def azimuth(self) -> float:
        """
        float: Telescope pointing azimuth (0-360).
        """
        return self._azimuth
    
    @azimuth.setter
    def azimuth(self, value: float):
        if value < 0 or value > 360:
            raise(ValueError('Telescope pointing azimuth must be between 0 and 360 degrees (increasing from north to east).'))
        self._azimuth = value
        self._pointing[1] = value
        self._cuda_telescope_init()
    
    @property
    def optical_system(self) -> OpticalSystem:
        """
        OpticalSystem: The telescope optical system
        """
        return self._optical_system
    
    @property
    def camera(self) -> CherenkovCamera:
        """
        CherenkovCamera: The telescope Cherenkov camera.
        """
        return self._camera

    @camera.setter
    def camera(self, a_camera):
        if a_camera is not None and not isinstance(a_camera, CherenkovCamera):
            raise TypeError("camera must be an instance of CherenkovCamera or a derived class")
        self._camera = a_camera
    
    @property
    def custom_rotation(self) -> np.ndarray:
        r"""
        The rotation matrix transforming the standard pointing frame into the telescope frame.


        .. math::
            R_{telescope} = R_{custom} \cdot R_{pointing}

        The pointing frame (when this matrix is identity) 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.

        Example
        -------
        The following rotates the telescope reference system in such a way to have, when the telescope is 
        pointing at the horizon towards the North (i.e. alt: 0, az: 0), the axes aligned as follows:

        * 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.

        .. code-block:: python
        
            rotation = np.array([
                [ 0, 1, 0], # new X is old Y
                [ 1, 0, 0], # new Y is old X
                [ 0, 0, 1] 
            ])
            telescope.custom_rotation = rotation

        Returns
        -------
        R : ndarray, shape (3, 3)
            The 3x3 rotation matrix. Defaults to the Identity matrix (no offset).
        """
        return self._custom_rotation

    @custom_rotation.setter
    def custom_rotation(self, matrix: np.ndarray):
        if matrix.shape != (3, 3):
            raise ValueError("Custom rotation must be a 3x3 matrix.")
        self._custom_rotation = matrix
    
    @property
    def local_to_telescope_rotation(self) -> np.ndarray:
        """Compute the rotation matrix to transform from the local reference frame to the telescope reference frame."""
        return local_to_telescope_rotation(*self.pointing, custom_rotation=self.custom_rotation)

    def _cuda_telescope_init(self):
        # Position and rotation to transform into the telescope reference frame
        telescope_rot = cp.asarray(self.local_to_telescope_rotation, dtype=cp.float32)
        telescope_pos = cp.asarray(self.position, dtype=cp.float32)
        
        self._cuda_tracing_args[0] = telescope_pos
        self._cuda_tracing_args[1] = telescope_rot
    
    def _cuda_surfaces_init(self):
        """Prepares optical surface data for CUDA-based ray tracing.

        This function extracts relevant properties from the optical surfaces
        defined in :py:attr:`optical_system` and converts them into CuPy arrays.
        """

        NUM_ASPHERIC_COEFFS = 10
        ROT_MATRIX_SIZE = 9

        surface_data_dtype = np.dtype([
            ('position',              np.float32, 3),                   # 12 bytes
            ('_pad1',                 np.uint8,   4),                   # ! Padding to align 'offset' (float2) to 8 bytes (12+4=16)
            ('offset',                np.float32, 2),                   # 8 bytes (starts at 16)
            ('rotation_matrix',       np.float32, ROT_MATRIX_SIZE),     # 36 bytes (starts at 24)
            ('curvature',             np.float32),                      # 4 bytes (starts at 60)
            ('conic_constant',        np.float32),                      # 4 bytes (starts at 64)
            ('aspheric_coefficients', np.float32, NUM_ASPHERIC_COEFFS), # 40 bytes (starts at 68, ends at 108)
            ('_pad2',                 np.uint8,   4),                   # ! Padding to align 'aperture_size' to 112
            ('aperture_size',         np.float32, 4),                   # 8 bytes (starts at 112)
            ('aperture_shape',        np.int32,   2),                   # 8 bytes
            ('flags',                 np.bool_,   2),                   # 2 bytes
            # align=True will automatically handle the 2-byte padding needed here for 'type' (int)
            ('type',                  np.int32),                        # 4 bytes
            ('shape',                 np.int32),                        # 4 bytes
            ('material1',             np.int32),                        # 4 bytes
            ('material2',             np.int32),                        # 4 bytes
            ('scattering_dispersion', np.float32),                      # 4 bytes
            ('sensor_info',           np.float32, 8),                   # 24 bytes
        ], align=True)

        n_surfaces = len(self.optical_system)

        h_surfaces = np.zeros(n_surfaces, dtype=surface_data_dtype)

        self._sub_pixels_per_side = []
        sensor_counter = 0
        for i, surface in enumerate(self.optical_system):
            if surface._shape not in SurfaceShape:
                raise(ValueError(f"Surface shape '{surface._shape}' ray-tracing not yet implemented."))

            h_surfaces[i]['position'] = np.asarray(surface.position, dtype=np.float32)
            h_surfaces[i]['rotation_matrix'] = np.asarray(surface.get_rotation_matrix().flatten(), dtype=np.float32)
            h_surfaces[i]['type'] = np.int32(surface.type.value)
            h_surfaces[i]['shape'] = np.int32(surface._shape.value)
            
            h_surfaces[i]['scattering_dispersion'] = np.float32(surface.scattering_dispersion/180.*np.pi)

            if surface._shape is SurfaceShape.CYLINDRICAL:
                h_surfaces[i]['material1'] = np.int32(surface.material_out.value)
                h_surfaces[i]['material2'] = np.int32(surface.material_in.value)
                h_surfaces[i]['aperture_size'][0] = np.float32(surface.radius)
                h_surfaces[i]['aperture_size'][1] = np.float32(surface.height)
                h_surfaces[i]['flags'][0] = np.bool_(surface.top)
                h_surfaces[i]['flags'][1] = np.bool_(surface.bottom)
                continue

            if surface._shape is SurfaceShape.CONICAL:
                h_surfaces[i]['material1'] = np.int32(surface.material_out.value)
                h_surfaces[i]['material2'] = np.int32(surface.material_in.value)
                h_surfaces[i]['aperture_size'][0] = np.float32(surface.radius_top)
                h_surfaces[i]['aperture_size'][1] = np.float32(surface.radius_bottom)
                h_surfaces[i]['aperture_size'][2] = np.float32(surface.height)
                h_surfaces[i]['flags'][0] = np.bool_(surface.top)
                h_surfaces[i]['flags'][1] = np.bool_(surface.bottom)
                continue

            if surface._shape is SurfaceShape.BOX:
                h_surfaces[i]['material1'] = np.int32(surface.material_out.value)
                h_surfaces[i]['material2'] = np.int32(surface.material_in.value)
                h_surfaces[i]['aperture_size'][0] = np.float32(surface.half_side_x)
                h_surfaces[i]['aperture_size'][1] = np.float32(surface.half_side_y)
                h_surfaces[i]['aperture_size'][2] = np.float32(surface.half_side_z)
                h_surfaces[i]['flags'][0] = np.bool_(surface.top)
                h_surfaces[i]['flags'][1] = np.bool_(surface.bottom)
                continue
            
            h_surfaces[i]['material1'] = np.int32(surface.material_in.value)
            h_surfaces[i]['material2'] = np.int32(surface.material_out.value)
            h_surfaces[i]['aperture_size'][0] = np.float32(surface.half_aperture)
            h_surfaces[i]['aperture_size'][1] = np.float32(surface.central_hole_half_aperture)
            h_surfaces[i]['aperture_shape'][0] = np.int32(surface.aperture_shape.value)
            h_surfaces[i]['aperture_shape'][1] = np.int32(surface.central_hole_shape.value)

            # Register pixels (in the order they have been inserted in the optical system)
            sensor_type = surface.sensor_type
            h_surfaces[i]['sensor_info'][0] = np.float32(sensor_type.value)
            if surface.sensor_type is not SensorSurfaceType.NONE:
                if sensor_type is SensorSurfaceType.SIPM_TILE:
                    h_surfaces[i]['sensor_info'][0] = np.float32(surface.sensor_type.value)
                    h_surfaces[i]['sensor_info'][1] = np.float32(sensor_counter)
                    h_surfaces[i]['sensor_info'][2] = np.float32(surface.pixels_per_side)
                    h_surfaces[i]['sensor_info'][3] = np.float32(surface.pixel_active_side)
                    h_surfaces[i]['sensor_info'][4] = np.float32(surface.pixels_separation)
                    h_surfaces[i]['sensor_info'][5] = np.float32(surface.border_to_active_area)
                    h_surfaces[i]['sensor_info'][6] = np.float32(surface._inv_ucell_size)
                    h_surfaces[i]['sensor_info'][7] = np.float32(surface._n_ucells_per_side)
                    self._sub_pixels_per_side.append(surface._n_ucells_per_side)
                    surface._first_pixel_id = sensor_counter
                sensor_counter += surface._n_pixels

            if surface._shape is SurfaceShape.FLAT:
                continue
            
            h_surfaces[i]['offset'] = np.asarray(surface.offset, dtype=np.float32)
            h_surfaces[i]['curvature'] = np.float32(surface.curvature)
            h_surfaces[i]['flags'][0] = np.bool_(surface.is_fresnel)

            if surface._shape is SurfaceShape.ASPHERICAL:
                h_surfaces[i]['conic_constant'] = np.float32(surface.conic_constant)
                surface_asph_coeffs = normalize_aspheric_coefficients(surface.aspheric_coefficients, surface.half_aperture)
                for j in range(len(surface_asph_coeffs)):
                    h_surfaces[i]['aspheric_coefficients'][j] = np.float32(surface_asph_coeffs[j])

        self._sub_pixels_per_side = cp.asarray(self._sub_pixels_per_side, dtype=cp.int32)
        self._cuda_tracing_args[2] = cp.asarray(h_surfaces.view(np.uint8), dtype=cp.uint8)

    def _cuda_surfaces_efficiency_init(self):
        optical_tables_dtype = np.dtype([
            ('transmittance_front', np.uint64),
            ('reflectance_front',   np.uint64),
            ('transmittance_back',  np.uint64),
            ('reflectance_back',    np.uint64),
            ('efficiency',          np.uint64),
            ('opt_inv_dwl',         np.float32),
            ('opt_start_wl',        np.float32),
            ('opt_inv_dang',        np.float32),
            ('opt_start_ang',       np.float32),
            ('eff_inv_dwl',         np.float32),
            ('eff_start_wl',        np.float32),
            ('eff_inv_dang',        np.float32),
            ('eff_start_ang',       np.float32)
        ], align=True)
        
        n_surfaces = len(self.optical_system)
        h_optical_tables = np.zeros(n_surfaces, dtype=optical_tables_dtype)
        self._surface_textures = [] # keep a reference to the textures to avoid garbage collection
        for k,s in enumerate(self.optical_system):
            surface_prop = s.properties
            if surface_prop is None or not surface_prop.is_defined:
                continue
            opt_textures = surface_prop.get_textures()
            self._surface_textures.append(opt_textures)
            h_optical_tables[k]['transmittance_front'] = opt_textures.transmittance.ptr
            h_optical_tables[k]['reflectance_front'] = opt_textures.reflectance.ptr
            h_optical_tables[k]['transmittance_back'] = opt_textures.transmittance_back.ptr
            h_optical_tables[k]['reflectance_back'] = opt_textures.reflectance_back.ptr
            h_optical_tables[k]['opt_inv_dwl'] = opt_textures.optical_wavelength_inv_step
            h_optical_tables[k]['opt_start_wl'] = opt_textures.optical_wavelength_start
            h_optical_tables[k]['opt_inv_dang'] = opt_textures.optical_angle_inv_step
            h_optical_tables[k]['opt_start_ang'] = opt_textures.optical_angle_start

            h_optical_tables[k]['efficiency'] = opt_textures.efficiency.ptr
            h_optical_tables[k]['eff_inv_dwl'] = opt_textures.efficiency_wavelength_inv_step
            h_optical_tables[k]['eff_start_wl'] = opt_textures.efficiency_wavelength_start
            h_optical_tables[k]['eff_inv_dang'] = opt_textures.efficiency_angle_inv_step
            h_optical_tables[k]['eff_start_ang'] = opt_textures.efficiency_angle_start
        
        self._cuda_tracing_args[3] = cp.asarray(h_optical_tables.view(np.uint8), dtype=cp.uint8)
    
    def _cuda_materials_init(self):
        """Prepares materials data for CUDA-based ray tracing.
        This function extracts relevant properties from all defined materials.
        """
        # Refractive indices
        refractive_tables_dtype = np.dtype([
            ('texture',  np.uint64),
            ('inv_dwl',  np.float32),
            ('start_wl', np.float32)
        ], align=True)

        n_materials = sum(1 for _ in Materials._members())

        h_refractive_tables = np.zeros(n_materials, dtype=refractive_tables_dtype)

        self._ri_textures = [] # keep a reference to the textures to avoid garbage collection
        for material in Materials._members():
            material_texture = material.get_refractive_index_texture()
            self._ri_textures.append(material_texture)
            h_refractive_tables[material.value]['texture'] = material_texture.texture.ptr
            h_refractive_tables[material.value]['inv_dwl'] = material_texture.wavelength_inv_step
            h_refractive_tables[material.value]['start_wl'] = material_texture.wavelength_start

        self._cuda_tracing_args[4] = cp.asarray(h_refractive_tables.view(np.uint8), dtype=cp.uint8)

    def _cuda_atmosphere_init(self):
        if self.atmospheric_transmission.value is not None:
            wl = self.atmospheric_transmission.wavelength
            zem = self.atmospheric_transmission.altitude
            self._cuda_atmosphere_args[:] = [
                cp.asarray(self.atmospheric_transmission.value, dtype=cp.float32),
                cp.asarray(wl, dtype=cp.float32),
                cp.asarray(zem, dtype=cp.float32),
                cp.asarray([len(wl), len(zem)], dtype=cp.int32),
            ]
                
    def cuda_init(self):
        """Build and copy to device all the CUDA kernels input related to the optical configuration."""

        self._cuda_surfaces_init()

        self._cuda_materials_init()

        self._cuda_telescope_init()

        self._cuda_surfaces_efficiency_init()

        if isinstance(self.camera, CherenkovSiPMCamera):
            pass

        self._cuda_atmosphere_init()

        self._d_num_surfaces = cp.int32(len(self.optical_system))

    def trace_photons(self,
            positions: Union[np.ndarray, cp.ndarray],
            directions: Union[np.ndarray, cp.ndarray],
            wavelengths: Union[np.ndarray, cp.ndarray],
            arrival_times: Union[np.ndarray, cp.ndarray],
            emission_altitudes: Optional[Union[np.ndarray, cp.ndarray]] = None,
            weights: Optional[Union[np.ndarray, cp.ndarray]] = None,
            events_mapping: Union[np.ndarray, list] = None,
            *,
            photons_per_bunch: int = 1,
            unpack_wavelength_range: Optional[tuple] = None,
            transform_to_tel = True,
            simulate_camera = False,
            min_n_pes = 1,
            get_camera_input: bool = False,
            max_bounces=100,
            save_last_bounce=False,
            seed: Optional[int] = None,
            free_memory_at_end: bool = False,
        ):
        r"""Traces photons through the telescope optical system and optionally simulates the Cherenkov camera response.

        This function simulates the propagation of photons through an optical system,
        taking into account their initial positions, directions, wavelengths, and arrival times.
        It can optionally:
        
        - model the effects of atmospheric transmission based on emission altitudes;
        - simulate the Cherenkov camera response.

        Each element in the input arrays can represent either a single photon or a "bunch" of photons.
        The ``photons_per_bunch`` parameter determines how many photons are represented by each element.
        If ``photons_per_bunch`` > 1, it is assumed that all photons within a bunch share the same
        initial position, direction and arrival time (for the wavelength see ``unpack_wavelength_range``).

        Parameters
        ----------
        positions : Union[numpy.ndarray, cp.ndarray] of shape (N, 3)
            Initial Cartesian coordinates (x, y, z) of the photon bunches in millimeters (mm).
            Must be either a NumPy (CPU) or CuPy (GPU) array.
        directions : Union[numpy.ndarray, cp.ndarray] of shape (N, 3)
            Initial direction vectors (vx, vy, vz) of the photon bunches.
            These should be normalized. Must be either a NumPy (CPU) or CuPy (GPU) array.
        wavelengths : Union[numpy.ndarray, cp.ndarray] of shape (N,)
            Wavelengths of the photon bunches in nanometers (nm).
            Must be either a NumPy (CPU) or CuPy (GPU) array.
        arrival_times : Union[numpy.ndarray, cp.ndarray] of shape (N,)
            Arrival times of the photon bunches at their respective positions in nanoseconds (ns).
            Must be either a NumPy (CPU) or CuPy (GPU) array.
        emission_altitudes : Optional[Union[numpy.ndarray, cp.ndarray]], optional
            Emission altitudes of the photon bunches in millimeters (mm).
            This is used to calculate atmospheric transmission effects if enabled.
            Must be either a NumPy (CPU) or CuPy (GPU) array.
        weights : Optional[Union[numpy.ndarray, cp.ndarray]], optional
            Weights of the photon bunches. The meaning of the weights depends on whether the bunches are unpacked or not (see Notes).
        events_mapping : numpy.ndarray or list
            A NumPy array of shape (N+1,) to split the photons in to groups corresponding to N different events when 
            generating the camera input arrays.
            ``events_mapping[i]`` indicates the starting index in ``positions`` for the i-th event.
            ``events_mapping[i+1]`` indicates the ending index (exclusive) in ``positions`` for the i-th event.
            Therefore, photons belonging to event i are located in ``positions[events_mapping[i]:events_mapping[i+1], :]``.
            The last element of events_mapping must be equal to the number of photons.
        photons_per_bunch : int, optional
            Number of photons represented by each element in the input arrays. If 1, each element represents a single photon.
            If > 1, each element represents a bunch of photons that share the same properties, by default 1. Maximum value is 65535.
        unpack_wavelength_range : tuple, optional
            Defines in which range (min,max) generate the wavelength of each unpacked photon. If None, bunches are non unpacked. By default None.
        transform_to_tel: bool, optional
            Whether to move detected photons into the telescope refrence frame or in the detector(s) reference frame. By default True.
        simulate_camera : bool, optional
            Whether to simulate the response of the telescope Cherenkov camera.
        min_n_pes : int, optional
            Minimum number of photo-electrons needed to trigger a camera simulation (if ``simulate_camera`` is True). By default 1.
        get_camera_input : bool, optional
            Whether to return the input for camera simulation (i.e. photon arrival times for each pixel) as CuPy ndarrays.
            By default False.
        max_bounces : int, optional
            Maximum number of bounces before the photon is killed. By default 100.
        save_last_bounce : float, optional
            Keep the position at the last bounce. Only for ray tracing visualization. By default False.
        seed : int, optional
            Base seed of the simulation. If None a random seed will be used. By default None.

        Returns
        -------
        ndarray
            List of triggered events index. If ``simulate_camera`` is True and ``get_camera_input`` is False.
        tuple(cp.ndarray, cp.ndarray)
            Input CuPy arrays for :py:meth:`CherenkovCamera.simulate_response()` method. If ``simulate_camera`` is False and ``get_camera_input`` is True.
        tuple(cp.ndarray, cp.ndarray, np.ndarray)
            Input CuPy arrays for :py:meth:`CherenkovCamera.simulate_response()` method and the list of triggered events index. If ``simulate_camera`` is True and ``get_camera_input`` is True.

        Warning
        -------
        Telescope configuration must be copied into the device before calling this method. Use the :py:meth:`cuda_init()` method.

        Warning
        -------

        - The ray-tracing operates with positions in mm, wavelengths in nm, and times in ns. It also handles photon bunches,
          where each thread may process multiple photons if ``photons_per_bunch`` > 1. In this case is assumed 
          that each element of the input arrays represents a bunch of photons that share the same 
          initial position, direction, wavelength, and arrival time. The arrays will maintain this
          structure, with each element representing the updated properties of a photon bunch.
        - To treat each photon in a bunch separately during the simulation, the ``unpack_wavelength_range`` 
          parameter must be provided. If ``unpack_wavelength_range`` is not None, 
          each bunch will be unpacked into its constituent photons, and the wavelength of each photon 
          will be randomly generated within the specified range following a Cherenkov-like distribution. 

        - The valid range and meaning of the ``weights`` array depends on this unpacking behavior:

          - If ``unpack_wavelength_range`` is provided, the weight of an input element represents the sum 
            of the weights of the photons within that bunch. Therefore, it must be <= ``photons_per_bunch``;
            
          - If ``unpack_wavelength_range`` is not provided, the weight applies to the bunch as a whole 
            single entity, and must be <= 1.
        
        Warning
        -------
        The input arrays can be either CuPy or NumPy ndarray. In the former case arrays are modified in place.
        In the latter case the arrays are copied to the device and can be accessed through the following attributes:

        - self._d_directions
        - self._d_positions
        - self._d_wavelengths
        - self._d_arrival_times
        - self._d_emission_altitudes
        - self._d_weights
        
        Notes
        -----
        **Memory Usage**

        The peak GPU memory (VRAM) footprint depends on the simulation flags provided and whether 
        photon bunches are unpacked (i.e., ``unpack_wavelength_range`` is not None):
        
        1. Ray tracing only (``simulate_camera=False`` and ``get_camera_input=False``):
           The method only allocates the base arrays for photon propagation.
           Peak VRAM (Bytes):

             - Packed simulation: ≈ :math:`44N`
             - Unpacked simulation: ≈ :math:`44NW + 32N`
           
        2. SiPM camera simulation without micro-cells (``camera.simulate_ucells=False``):
           Filters detected photons, sorts them, and generates the time expansion array.
           Peak VRAM (Bytes):

             - Packed simulation: ≈ :math:`45N + 64M + 4P`
             - Unpacked simulation: ≈ :math:`45NW + 64M + 4P + 32N`
           
        3. SiPM camera simulation with micro-cells (``camera.simulate_ucells=True``):
           Adds an additional output array to track the specific sub-pixel IDs during expansion.
           Peak VRAM (Bytes):

             - Packed simulation: ≈ :math:`45N + 64M + 8P`
             - Unpacked simulation: ≈ :math:`45NW + 64M + 8P + 32N`
        
        The formulas above assume the optional ``weights`` and ``emission_altitudes`` 
        parameters are not provided. If provided, each array adds 4 bytes per tracked 
        element (:math:`4N` for packed, or + :math:`4NW` for unpacked).
        
        Where:

        - :math:`N` is the total number of input photon bunches;
        - :math:`W` is the ``photons_per_bunch`` multiplier;
        - :math:`M` is the number of detected elements (photons or bunches);
        - :math:`P` is the total number of photo-electrons generated after weight expansion (:math:`P \ge M`).
        
        **Example**
        Assuming a telescope optical efficiency of 10% (:math:`M = 0.1N`) and a bunch weight of 5 photons 
        per bunch (:math:`W = 5`), and providing both ``weights`` and ``emission_altitudes``,
        the peak VRAM scales as follows:
        
        **Packed execution** (Packed into bunches, :math:`N_{sim} = N`, :math:`P = 5M`):
        
        ===================  ================  ================  ================
        Input bunches (N)    Ray tracing only  w/o microcells    w/ microcells
        ===================  ================  ================  ================
        1,000,000            ~ 52.0 MB         ~ 61.4 MB         ~ 63.4 MB
        5,000,000            ~ 260.0 MB        ~ 307.0 MB        ~ 317.0 MB
        10,000,000           ~ 520.0 MB        ~ 614.0 MB        ~ 634.0 MB
        25,000,000           ~ 1300.0 MB       ~ 1535.0 MB       ~ 1585.0 MB
        ===================  ================  ================  ================

        **Unpacked execution** (Unpacked into single photons, :math:`N_{sim} = 5N`, :math:`P = M`):
        
        ===================  ================  ================  ================
        Input bunches (N)    Ray tracing only  w/o microcells    w/ microcells
        ===================  ================  ================  ================
        1,000,000            ~ 292.0 MB        ~ 331.0 MB        ~ 333.0 MB
        5,000,000            ~ 1460.0 MB       ~ 1655.0 MB       ~ 1665.0 MB
        10,000,000           ~ 2920.0 MB       ~ 3310.0 MB       ~ 3330.0 MB
        25,000,000           ~ 7300.0 MB       ~ 8275.0 MB       ~ 8325.0 MB
        ===================  ================  ================  ================

        """
        if self.timer.active:
            cp.cuda.stream.get_current_stream().synchronize()
            t0 = time.time()
        
        if seed is None:
            seed = random.getrandbits(62)
        
        has_emission_altitudes = emission_altitudes is not None
        has_weights = weights is not None

        if photons_per_bunch > 1 and unpack_wavelength_range is not None:
            
            # Copy arrays to device
            d_in_pos = cp.asarray(positions, dtype=cp.float32)
            d_in_dir = cp.asarray(directions, dtype=cp.float32)
            d_in_time = cp.asarray(arrival_times, dtype=cp.float32)
            
            if has_emission_altitudes:
                d_in_alt = cp.asarray(emission_altitudes, dtype=cp.float32)
            else:
                d_in_alt = cp.empty(1, dtype=cp.float32) # Dummy array
            
            if has_weights:
                d_in_weights = cp.asarray(weights, dtype=cp.float32)
            else:
                d_in_weights = cp.empty(1, dtype=cp.float32) # Dummy array

            n_bunches = d_in_pos.shape[0]
            n_photons = n_bunches * photons_per_bunch

            # Allocate the memory needed for the output arrays
            self._d_positions = cp.empty((n_photons, 3), dtype=cp.float32)
            self._d_directions = cp.empty((n_photons, 3), dtype=cp.float32)
            self._d_arrival_times = cp.empty(n_photons, dtype=cp.float32)
            self._d_wavelengths = cp.empty(n_photons, dtype=cp.float32)
            
            if has_emission_altitudes:
                self._d_emission_altitudes = cp.empty(n_photons, dtype=cp.float32)
            else:
                self._d_emission_altitudes = cp.empty(1, dtype=cp.float32)
            
            if has_weights:
                self._d_weights = cp.empty(n_photons, dtype=cp.float32)
            else:
                self._d_weights = cp.empty(1, dtype=cp.float32)
            
            # Wavelength parameters
            lambda_min, lambda_max = unpack_wavelength_range
            inv_l_min = np.float32(1. / lambda_min)
            inv_l_max = np.float32(1. / lambda_max)

            # Launch unpacking kernel
            grid_size = (n_photons + self._blocksize - 1) // self._blocksize
            ray_tracing.unpack_bunches(
                (grid_size,),
                (self._blocksize,), 
                (
                    d_in_pos,
                    d_in_dir,
                    d_in_time,
                    d_in_alt,
                    d_in_weights,
                    self._d_positions,
                    self._d_directions,
                    self._d_arrival_times, 
                    self._d_emission_altitudes,
                    self._d_wavelengths,
                    self._d_weights,
                    cp.int32(photons_per_bunch),
                    cp.float32(inv_l_min),
                    cp.float32(inv_l_max),
                    cp.int32(n_photons),
                    cp.bool_(has_emission_altitudes),
                    cp.bool_(has_weights),
                    cp.uint64(seed)
                )
            )

            # Fix events mapping
            if events_mapping is None:
                events_mapping = cp.asarray([0, n_photons], dtype=cp.int32)
            else:
                events_mapping = cp.asarray(events_mapping, dtype=cp.int32) * photons_per_bunch

            # Now bunches have been unpacked
            photons_per_bunch = 1
        else:
            # Fallback logic for photons_per_bunch=1 or no wavelength extraction
            self._d_positions = cp.asarray(positions, dtype=cp.float32)
            self._d_directions = cp.asarray(directions, dtype=cp.float32)
            self._d_wavelengths = cp.asarray(wavelengths, dtype=cp.float32)
            self._d_arrival_times = cp.asarray(arrival_times, dtype=cp.float32)

            n_photons = self._d_positions.shape[0]
            
            if has_emission_altitudes:
                self._d_emission_altitudes = cp.asarray(emission_altitudes, dtype=cp.float32)
            
            if has_weights:
                self._d_weights = cp.asarray(weights, dtype=cp.float32)
            
            if events_mapping is None:
                events_mapping = cp.asarray([0, self._d_positions.shape[0]], dtype=cp.int32)
            else:
                events_mapping = cp.asarray(events_mapping, dtype=cp.int32)
        
        # ID of the pixel reached by each bunch/photon
        self._d_pixid = cp.empty((n_photons,), dtype=cp.int32)
        self._d_subpixid = cp.empty((n_photons,), dtype=cp.int32)
        self._d_surfaceid = cp.empty((n_photons,), dtype=cp.int32)

        # Calculate kernel call parameters
        num_blocks = int(np.ceil(n_photons / self._blocksize))
        d_num_photons = cp.int32(n_photons)

        if self.timer.active:
            cp.cuda.stream.get_current_stream().synchronize()
            t1 = time.time()
            self.timer.add_entry('simulate_response', 'copy to device', t1-t0)
        
        # Apply weights
        if has_weights:
            ray_tracing.reject_by_weight(
                (num_blocks,),
                (self._blocksize,),
                (
                    self._d_positions,
                    self._d_directions,
                    self._d_wavelengths,
                    self._d_arrival_times,
                    self._d_weights,
                    d_num_photons,
                    cp.uint64(seed+1),
                )
            )

        if self.timer.active:
            cp.cuda.stream.get_current_stream().synchronize()
            t1b = time.time()
            self.timer.add_entry('simulate_response', 'weight photons', t1b-t1)
        
        # Optional atmospheric transmission calculation
        if self.atmospheric_transmission.value is not None and emission_altitudes is not None:
            ray_tracing.atmospheric_transmission(
                (num_blocks,), 
                (self._blocksize,),
                (
                    self._d_positions,
                    self._d_directions,
                    self._d_wavelengths,
                    self._d_arrival_times,
                    self._d_emission_altitudes,
                    *self._cuda_atmosphere_args,
                    d_num_photons,
                    cp.uint64(seed+2),
                    cp.bool_(True) # force the use of log(z) for now
                )
            )

        if self.timer.active:
            cp.cuda.stream.get_current_stream().synchronize()
            t2 = time.time()
            self.timer.add_entry('simulate_response', 'atmospheric transmission', t2-t1)

        # Perform ray tracing
        ray_tracing.ray_tracing(
            (num_blocks,),
            (self._blocksize,),
            (
                self._d_positions,
                self._d_directions,
                self._d_wavelengths,
                self._d_arrival_times,
                self._d_pixid,
                self._d_subpixid,
                self._d_surfaceid,
                *self._cuda_tracing_args,
                self._d_num_surfaces,
                d_num_photons,
                cp.int32(max_bounces),
                cp.bool_(save_last_bounce),
                cp.bool_(transform_to_tel),
                cp.uint64(seed+3)
            ),
            shared_mem = len(self.optical_system) * 5 * 4 # Surface positions, bounding radii and shape ID
        )

        if self.timer.active:
            cp.cuda.stream.get_current_stream().synchronize()
            t3 = time.time()
            self.timer.add_entry('simulate_response', 'trace onto focal plane', t3-t2)

        # Build the input for `CherenkovCamera.simulate_response()` method
        if get_camera_input or simulate_camera:

            # Split photons into events
            if events_mapping is None:
                events_mapping = cp.asarray([0, n_photons], dtype=cp.int32)
            else:
                events_mapping = cp.asarray(events_mapping, dtype=cp.int32)

            n_pixels = self.camera.n_pixels
            n_events = len(events_mapping) - 1

            # Filter for detected photons to reduce workload
            hit_mask = cp.isfinite(self._d_arrival_times) & (self._d_pixid > -1)
            
            if not hit_mask.any():
                # No photons detected
                phs = cp.array([], dtype=cp.float32)
                phs_mapping = cp.zeros((n_events, n_pixels + 1), dtype=cp.int32)
                pe_mapping = cp.zeros((n_events + 1,), dtype=cp.int32)
                active_indices = np.array([], dtype=np.int32)
                pe_mapping_host = np.zeros((n_events + 1,), dtype=np.int32)
                if getattr(self.camera, 'simulate_ucells', False):
                    subids = cp.array([], dtype=cp.int32)
            else:
                hit_indices_global = cp.where(hit_mask)[0]
                hit_pix = self._d_pixid[hit_mask]
                hit_subpix = self._d_subpixid[hit_mask]
                hit_time = self._d_arrival_times[hit_mask]

                # Get event IDs of each detected photon bunch
                hit_event = cp.searchsorted(events_mapping, hit_indices_global, side='right').astype(cp.int32) - 1
                
                # Create composite key for the bunches to sort by event, then by pixel
                create_composite_key = cp.ElementwiseKernel(
                    'int32 event, int32 pixel, int64 n_pixels',
                    'int64 key',
                    'key = (long long)event * n_pixels + pixel',
                    'create_composite_key'
                )
                bunch_composite_key = cp.empty(len(hit_event), dtype=cp.int64)
                create_composite_key(hit_event, hit_pix, n_pixels, bunch_composite_key)
                
                # CUDA kernel that simulate ucells needs sorted arrival times per each pixel
                if getattr(self.camera, 'simulate_ucells', False) > 0:
                    # First sort key: event -> inside bunch_composite_key
                    # Second sort key: pixel -> inside bunch_composite_key
                    # Third sort key: time
                    keys = cp.stack((hit_time, bunch_composite_key))
                    sort_order = cp.lexsort(keys)
                    s_hit_subpix = hit_subpix[sort_order]
                else:
                    # First sort key: event -> inside bunch_composite_key
                    # Second sort key: pixel -> inside bunch_composite_key
                    sort_order = cp.argsort(bunch_composite_key)

                # Sort the bunches before any expansion to reduce memory footprint
                s_hit_time = hit_time[sort_order]
                s_hit_pix = hit_pix[sort_order]
                s_bunch_composite_key = bunch_composite_key[sort_order]

                # Generate mappings
                # Count bunches per bin, then multiply by the uniform bunch size to get total PEs
                flat_bunch_counts = cp.bincount(s_bunch_composite_key, minlength=n_events*n_pixels).astype(cp.int32)
                flat_counts = flat_bunch_counts * photons_per_bunch
                counts_2d = flat_counts.reshape((n_events, n_pixels))

                # phs_mapping: cumulative sum along pixels for each event
                zeros_col = cp.zeros((n_events, 1), dtype=cp.int32)
                counts_padded = cp.hstack([zeros_col, counts_2d])
                phs_mapping = cp.cumsum(counts_padded, axis=1, dtype=cp.int32)

                # pe_mapping: global offsets for slicing 'phs' by event
                event_pes = counts_2d.sum(axis=1)
                # TODO: rename this for clarity
                pe_mapping = cp.concatenate([cp.array([0], dtype=cp.int32), cp.cumsum(event_pes, dtype=cp.int32)])

                # Calculate the start index for each bunch
                starts = cp.arange(len(s_hit_time) + 1, dtype=cp.int32) * photons_per_bunch
                total_pes = int(starts[-1])
                
                # Pre-allocate the final output array
                phs = cp.empty(total_pes, dtype=cp.float32)

                if not getattr(self.camera, 'simulate_ucells', False):
                    # No micro-cells expansion: write 'time' into 'phs' for 'ppb' iterations
                    # Use an ElementwiseKernel to save GPU memory
                    expand_phs_kernel = cp.ElementwiseKernel(
                        'int32 start, int32 ppb, float32 time',
                        'raw float32 phs',
                        '''
                        for (int i = 0; i < ppb; ++i) {
                            phs[start + i] = time;
                        }
                        ''',
                        'expand_phs_kernel'
                    )
                    expand_phs_kernel(starts[:-1], cp.int32(photons_per_bunch), s_hit_time, phs)

                else:
                    if photons_per_bunch > 9:
                        raise ValueError("Microcell simulation not supported for bunch-size > 9.")
                    
                    subids = cp.empty(total_pes, dtype=cp.int32)
                    
                    # Ensure n_side is a cupy array for elementwise mapping
                    n_side_array = cp.asarray(self._sub_pixels_per_side, dtype=cp.int32)[s_hit_pix]

                    # Micro-cells expansion: spread photons of a bunch on a 9X9 microcells window 
                    #                        and write 'time' into 'phs' for 'ppb' iterations
                    # Use an ElementwiseKernel to save GPU memory
                    expand_ucells_kernel = cp.ElementwiseKernel(
                        'int32 start, int32 ppb, float32 time, int32 subpix, int32 n_side',
                        'raw float32 phs, raw int32 subids',
                        '''
                        // Define pattern offsets (9x9 window):
                        //   0: Center
                        //   1-4: cardinal (up, down, left, right)
                        //   5-8: diagonals (up-left, up-right, down-left, down-right)
                        int dr[9] = {0, -1, 1, 0, 0, -1, -1, 1, 1};
                        int dc[9] = {0, 0, 0, -1, 1, -1, 1, -1, 1};
                        
                        for (int i = 0; i < ppb; ++i) {
                            int out_idx = start + i;
                            phs[out_idx] = time;
                            
                            int pattern_idx = i % 9;
                            int r0 = subpix / n_side;
                            int c0 = subpix % n_side;
                            
                            int final_r = r0 + dr[pattern_idx];
                            if (final_r < 0) final_r = 0;
                            if (final_r > n_side - 1) final_r = n_side - 1;
                            
                            int final_c = c0 + dc[pattern_idx];
                            if (final_c < 0) final_c = 0;
                            if (final_c > n_side - 1) final_c = n_side - 1;
                            
                            subids[out_idx] = final_r * n_side + final_c;
                        }
                        ''',
                        'expand_ucells_kernel'
                    )
                    expand_ucells_kernel(
                        starts[:-1], cp.int32(photons_per_bunch), s_hit_time, s_hit_subpix, n_side_array, phs, subids
                    )
            
            if self.timer.active:
                cp.cuda.stream.get_current_stream().synchronize()
                t4 = time.time()
                self.timer.add_entry('simulate_response', 'generate camera input', t4-t3)

            output = None
            
            # Simulate camera response
            if simulate_camera:
                self.camera.base_seed = seed+4

                if phs.size > 0:
                    n_pes_per_event = pe_mapping[1:] - pe_mapping[:-1]
                    active_mask = n_pes_per_event >= min_n_pes
                    active_indices_gpu = cp.where(active_mask)[0]

                    # Transfer minimal loop control data to CPU
                    active_indices = active_indices_gpu.get()
                    pe_mapping_host = pe_mapping.get()
                else:
                    active_indices = []
                
                if self.timer.active:
                    cp.cuda.stream.get_current_stream().synchronize()
                    t5 = time.time()
                    self.timer.add_entry('simulate_response', 'filter events to be simulated', t5-t4)

                event_number = []

                for i in tqdm(active_indices, disable=not self.show_progress):
                    start_phs = pe_mapping_host[i]
                    end_phs = pe_mapping_host[i+1]
                    if getattr(self.camera, 'simulate_ucells', False):
                        source = (phs[start_phs:end_phs], phs_mapping[i], subids[start_phs:end_phs])
                    else:
                        source = (phs[start_phs:end_phs], phs_mapping[i])
                    self.camera.simulate_response(source)
                    if self.camera.triggered:
                        event_number.append(i)
                
                event_number = np.array(event_number)

                if self.timer.active:
                    cp.cuda.stream.get_current_stream().synchronize()
                    t6 = time.time()
                    self.timer.add_entry('simulate_response', 'simulate camera response', t6-t5)

                if get_camera_input:
                    if getattr(self.camera, 'simulate_ucells', False):
                        output = phs, phs_mapping, subids, event_number
                    else:
                        output = phs, phs_mapping, event_number
                else:
                    output = event_number
            
            if get_camera_input and not simulate_camera:
                if getattr(self.camera, 'simulate_ucells', False):
                    output = phs, phs_mapping, subids
                else:
                    output = phs, phs_mapping
            
            if free_memory_at_end:
                self.clean_memory()

            return output

    def clean_memory(self, reclaim: bool = False):
        """
        Clears GPU memory by explicitly removing references to CuPy arrays created during ray tracing.

        Parameters
        ----------
        reclaim : bool, optional
            If True, forces Python garbage collection and completely flushes 
            the CuPy memory pools back to the OS. This is a slow operation. Defaults to False.
        """
        # Attributes dynamically created during trace_photons
        gpu_arrays = [
            '_d_directions',
            '_d_positions',
            '_d_wavelengths',
            '_d_arrival_times',
            '_d_emission_altitudes',
            '_d_weights',
            '_d_pixid',
            '_d_subpixid',
            '_d_surfaceid'
        ]

        # Dropping references instantly returns the memory to the CuPy pool
        for attr in gpu_arrays:
            if hasattr(self, attr):
                setattr(self, attr, None)

        # Only perform the (slow) memory flush if explicitly requested
        if reclaim:
            import gc
            import cupy as cp
            gc.collect()
            cp.get_default_memory_pool().free_all_blocks()
            cp.get_default_pinned_memory_pool().free_all_blocks()

    def visualize_ray_tracing(self,
            positions: Union[np.ndarray, cp.ndarray],
            directions: Union[np.ndarray, cp.ndarray],
            wavelengths: Union[np.ndarray, cp.ndarray],
            arrival_times: Union[np.ndarray, cp.ndarray],
            emission_altitudes: Optional[Union[np.ndarray, cp.ndarray]] = None,
            weights: Optional[Union[np.ndarray, cp.ndarray]] = None,
            photons_per_bunch: int = 1,
            unpack_wavelength_range: Optional[tuple] = None,
            *,
            show_not_detected: bool = True,
            map_wavelength_color=True,
            get_renderer = False,
            opacity=None,
            show_hits=True,
            show_rays=True, 
            point_size=1,
            orthographic=False,
            focal_point=None,
            resolution=None,
            max_bounces = 100
        ):
        """
        Visualizes the ray tracing process by performing step-by-step propagation of photons.

        This method initializes a VTK renderer and iteratively calls `trace_photons` with 
        `max_bounces=1` to capture and visualize the path of photons between surfaces.

        Parameters
        ----------
        positions : Union[numpy.ndarray, cp.ndarray] of shape (N, 3)
            Initial Cartesian coordinates (x, y, z) of the photon bunches in millimeters (mm).
            Must be either a NumPy (CPU) or CuPy (GPU) array.
        directions : Union[numpy.ndarray, cp.ndarray] of shape (N, 3)
            Initial direction vectors (vx, vy, vz) of the photon bunches.
            These should be normalized. Must be either a NumPy (CPU) or CuPy (GPU) array.
        wavelengths : Union[numpy.ndarray, cp.ndarray] of shape (N,)
            Wavelengths of the photon bunches in nanometers (nm).
            Must be either a NumPy (CPU) or CuPy (GPU) array.
        arrival_times : Union[numpy.ndarray, cp.ndarray] of shape (N,)
            Arrival times of the photon bunches at their respective positions in nanoseconds (ns).
            Must be either a NumPy (CPU) or CuPy (GPU) array.
        emission_altitudes : Optional[Union[numpy.ndarray, cp.ndarray]], optional
            Emission altitudes of the photon bunches in millimeters (mm).
            This is used to calculate atmospheric transmission effects if enabled.
            Must be either a NumPy (CPU) or CuPy (GPU) array.
        weights : Optional[Union[numpy.ndarray, cp.ndarray]], optional
            Weights of the photon bunches. The meaning of the weights depends on whether the bunches are unpacked or not (see :py:meth:`trace_photons`).
        photons_per_bunch : int, optional
            Number of physical photons represented by one simulated photon. Defaults to 1.
        unpack_wavelength_range : tuple, optional
            Defines in which range (min,max) generate the wavelength of each unpacked photon. If None, bunches are non unpacked. By default None.
        show_not_detected : bool, optional
            Whether to show also photons that are not detected by the camera. Defaults to True.
        map_wavelength_color : bool, optional
            Whether to color rays based on their wavelength. Defaults to True.
        get_renderer : bool, optional
            If True, returns the `VTKOpticalSystem` instance instead of starting the render loop immediately. 
            Defaults to False.
        opacity : float, optional
            Opacity of the ray lines (0.0 to 1.0). If None, it is calculated based on the number 
            of photons. Defaults to None.
        show_hits : bool, optional
            Whether to render points at the location where rays intersect surfaces. Defaults to True.
        show_rays : bool, optional
            Whether to render lines representing the photon paths. Defaults to True.
        point_size : int, optional
            Size of the points rendered at intersection hits. Defaults to 1.
        orthographic : bool, optional
            If True, start with a parallel projection. 
            If False, start with a Perspective projection. Default is False.
        focal_point : tuple, list, or str, optional
            If tuple/list: (x, y, z) coordinates to look at.
            If str: The name of the surface in self.os to look at.
            If None: Defaults to center of system.
        resolution : float, optional
            Objects mesh resolution (in mm). By default None (depends on object extent).
        max_bounces: int, optional
            Stop the ray-tracing if the number of bounces exceeds this value. Defaults to 100.
        
        Returns
        -------
        VTKOpticalSystem or None
            Returns the renderer instance if `get_renderer` is True, otherwise returns None 
            after the render window is closed.
        """
        from .visualization._vtk_optical_system import VTKOpticalSystem
        
        # Ensure input arrays are on GPU, so any modifications is inplace
        positions = cp.asarray(positions, dtype=cp.float32)
        directions = cp.asarray(directions, dtype=cp.float32)
        wavelengths = cp.asarray(wavelengths, dtype=cp.float32)
        arrival_times = cp.asarray(arrival_times, dtype=cp.float32)
        emission_altitudes = cp.asarray(emission_altitudes, dtype=cp.float32) if emission_altitudes is not None else None
        weights = cp.asarray(weights, dtype=cp.float32) if weights is not None else None

        # Initialize VTK renderer
        renderer = VTKOpticalSystem(self, resolution=resolution)
        
        # Where to put rays if showing only detected ones at the end
        ps_start_list = []
        ps_stop_list = []
        surface_hits_list = []
        vis_wavelengths = None

        bounces = 0
        while True:
            # Store start positions
            ps_start = positions.get()

            if bounces > 0:
                unpack_wavelength_range = None
                weights = None
                emission_altitudes = None
            
            # Perform ray tracing for one bounce
            self.trace_photons(
                positions,
                directions,
                wavelengths,
                arrival_times,
                emission_altitudes,
                weights=weights,
                photons_per_bunch=photons_per_bunch,
                unpack_wavelength_range=unpack_wavelength_range,
                simulate_camera=False,
                max_bounces=1, 
                save_last_bounce=True,
                transform_to_tel=False,
            )

            # If bunches were unpacked in this bounce, duplicate the initial positions 
            # to match the expanded arrays, then set the photon bunch size to 1
            if bounces == 0 and photons_per_bunch > 1 and unpack_wavelength_range is not None:
                ps_start = np.repeat(ps_start, photons_per_bunch, axis=0)
                photons_per_bunch = 1

            # Sync local references with the internal GPU arrays that were modified
            positions = self._d_positions
            directions = self._d_directions
            wavelengths = self._d_wavelengths
            arrival_times = self._d_arrival_times
            if emission_altitudes is not None:
                emission_altitudes = self._d_emission_altitudes

            # Filter photons based on photon status
            if show_not_detected or (not show_not_detected and bounces==0):
                survivors = cp.isfinite(wavelengths)
                positions = positions[survivors]
                directions = directions[survivors]
                wavelengths = wavelengths[survivors]
                arrival_times = arrival_times[survivors]
                self._d_surfaceid = self._d_surfaceid[survivors]
                if emission_altitudes is not None:
                    emission_altitudes = emission_altitudes[survivors]
                if map_wavelength_color:
                    vis_wavelengths = wavelengths.get()
                ps_start = ps_start[survivors.get()]
            
            # Surface hit IDs and stop positions
            surface_hits = self._d_surfaceid.get()
            ps_stop = positions.get()
            
            # Stop if no more changes
            if np.array_equal(ps_start, ps_stop, equal_nan=True):
                break
            
            # Add rays to renderer
            if show_not_detected:
                renderer.add_rays(
                    ps_start, ps_stop,
                    wavelengths=vis_wavelengths,
                    surface_id=surface_hits,
                    show_hits=show_hits,
                    show_rays=show_rays,
                    point_size=point_size
                )
            # Store for later rendering
            else:
                ps_start_list.append(ps_start)
                ps_stop_list.append(ps_stop)
                surface_hits_list.append(surface_hits)
            
            # Break if max bounces reached
            if bounces > max_bounces:
                break
            
            # Update bounce counter
            bounces += 1
        
        # If showing only detected photons, add them now
        if not show_not_detected:
            for ps_start, ps_stop, surface_hits in zip(ps_start_list, ps_stop_list, surface_hits_list):
                mask = cp.isfinite(wavelengths).get()
                if map_wavelength_color:
                    vis_wavelengths = wavelengths.get()[mask]
                renderer.add_rays(
                    ps_start[mask], ps_stop[mask],
                    wavelengths=vis_wavelengths,
                    surface_id=surface_hits[mask],
                    show_hits=show_hits,
                    show_rays=show_rays,
                    point_size=point_size
                )
        
        # Set ray opacity (by default based on number of photons)
        renderer.set_ray_opacity(opacity)

        # Return the renderer if requested
        if get_renderer:
            return renderer
        # otherwise start the render loop
        else:
            renderer.start_render(focal_point=focal_point, orthographic=orthographic)