# 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 os as os
from pathlib import Path as Path
import cupy as cp
import numpy as np
import struct
from typing import Tuple, Optional, Union

from ...utils._kernels import get_kernel

path = Path(__file__).parent / "poisson_flash.cu"
with open(path) as _source_file:
    source_code = _source_file.read()

include_path = ''.join([os.environ['CPATH']])
module = cp.RawModule(code=source_code, backend='nvcc', options=(''.join(['-I',include_path]),'--use_fast_math', '-std=c++11', '--extra-device-vectorization', '-O3'))
generate_map = module.get_function("generate_map")
generate_arrival_times = module.get_function("generate_arrival_times")

def generate_poisson_flash(
        n_pixels: int,
        mean: Union[cp.ndarray, np.ndarray],
        t_start: float,
        duration: float,
        seed: Optional[int] = None) -> Tuple[cp.ndarray, cp.ndarray]:
    """Generate Poisson arrival times in a time window equal for all pixels.

    This function simulates the arrival times of photo-electrons
    generated by a Poisson process, representing a light flash.  It assumes
    the same time window (start time and duration) for all pixels, but allows
    for a different mean number of photo-electrons for each pixel.

    Parameters
    ----------
    n_pixels : int
        Total number of pixels.
    mean : Union[cp.ndarray, numpy.ndarray]
        Mean number of photo-electrons for each pixel. This should be a
        1D array of length `n_pixels`.
    t_start : float
        Start time of the flash (in ns).
    duration : float
        Duration of the flash (in ns).
    seed : int, optional
        Random seed for reproducibility. If None, a random seed is generated.

    Returns
    -------
    Tuple[cp.ndarray, cp.ndarray]
        A tuple containing:

        - **pes**: A 1D CuPy array of shape (total_n_pes,) containing the
          arrival times of all photo-electrons.
        
        - **pes_map**: A 1D CuPy array of shape (n_pixels + 1,) representing
          a mapping to access the arrival times for each pixel.  The arrival
          times for pixel `i` are given by ```pes[pes_map[i]:pes_map[i+1]]```.

    Examples
    --------
    .. code-block:: python

        import cupy as cp
        n_pixels = 10
        mean = cp.array([2.0, 3.5, 1.0, 4.2, 0.8, 2.7, 1.9, 3.1, 2.5, 1.3])
        t_start = 0.0
        duration = 10.0
        pes, pes_map = generate_poisson_flash(n_pixels, mean, t_start, duration)
        # Access arrival times for pixel 3:
        pixel_3_times = pes[pes_map[3]:pes_map[4]]

    """
    if seed is None:
        seed = struct.unpack('>I', os.urandom(4))[0]

    mean = cp.asarray(mean, dtype=cp.float32)
    
    pe_per_pixel = cp.empty((n_pixels+1,), dtype=cp.int32)
    
    blocksize = 128
    nblock = int(np.ceil(n_pixels/blocksize))
    generate_map((nblock,), (blocksize,), (pe_per_pixel,mean, cp.uint64(seed), cp.int32(n_pixels)))

    pe_map = cp.cumsum(pe_per_pixel, dtype=cp.int32)

    n_pes = pe_map[-1].get()
    pe_arrival_times = cp.empty((n_pes,), dtype=cp.float32)

    t_start = cp.float32(t_start)
    duration = cp.float32(duration)
    
    blocksize = 128
    nblock = int(np.ceil(n_pes/blocksize))
    generate_arrival_times((nblock,), (blocksize,), (pe_arrival_times, t_start, duration, cp.uint64(seed), cp.int32(n_pes)))

    return pe_arrival_times, pe_map