#
# Copyright 2024-2026 Universidad Complutense de Madrid
#
# This file is part of Numina
#
# SPDX-License-Identifier: GPL-3.0-or-later
# License-Filename: LICENSE.txt
#
import logging

from astropy.io import fits
import numpy as np
from pathlib import Path

from .define_3d_wcs import wcs_to_header_using_cd_keywords


def generate_image3d_method0_ifu(
    wcs3d,
    header_keys,
    simulated_x_ifu_all,
    simulated_y_ifu_all,
    simulated_wave_all,
    bins_x_ifu,
    bins_y_ifu,
    bins_wave,
    prefix_intermediate_fits,
    logger=None,
    output_dir=".",
):
    """Compute 3D image3 IFU, method 0

    The image is calculated by performing a three-dimensional histogram
    over the X, Y coordinates in the IFU and the wavelength of the
    simulated photons.

    The result is saved as a FITS file.

    Parameters
    ----------
    wcs3d : `~astropy.wcs.wcs.WCS`
        WCS of the data cube.
    header_keys : `~astropy.io.fits.header.Header`
        FITS header with additional keywords to be merged together with
        the WCS information.
    simulated_x_ifu_all : `~astropy.units.Quantity`
        Simulated X coordinates of the photons in the IFU.
    simulated_y_ifu_all : `~astropy.units.Quantity`
        Simulated Y coordinates of the photons in the IFU.
    simulated_wave_all : `~astropy.units.Quantity`
        Simulated wavelengths of the photons in the IFU.
    bins_x_ifu : `~numpy.ndarray`
        Bin edges in the naxis1_ifu direction
        (along the slice).
    bins_y_ifu : `~numpy.ndarray`
        Bin edges in the naxis2_ifu direction
        (perpendicular to the slice).
    bins_wave : `~numpy.ndarray`
        Bin edges in the wavelength direction.
    prefix_intermediate_fits : str
        Prefix for output intermediate FITS files. If the length of
        this string is 0, no output is generated.
    logger : `~logging.Logger`, optional
        Logger for logging messages. If `None`, the root logger is used.
    output_dir : str or `~pathlib.Path`, optional
        Output directory to store results. Default is the current directory.
    """
    if logger is None:
        logger = logging.getLogger()

    if simulated_x_ifu_all is None or simulated_y_ifu_all is None or simulated_wave_all is None:
        nphotons_all = 0
    else:
        nphotons_all = len(simulated_x_ifu_all)
        if nphotons_all != len(simulated_y_ifu_all) or nphotons_all != len(simulated_wave_all):
            raise ValueError(
                "The length of simulated_x_ifu_all, simulated_y_ifu_all, and simulated_wave_all must be the same."
            )

    # generate image
    if nphotons_all > 0:
        image3d_method0_ifu, edges = np.histogramdd(
            sample=(simulated_wave_all.value, simulated_y_ifu_all.value, simulated_x_ifu_all.value),
            bins=(bins_wave.value, bins_y_ifu.value, bins_x_ifu.value),
        )
    else:
        image3d_method0_ifu = np.zeros((len(bins_wave) - 1, len(bins_y_ifu) - 1, len(bins_x_ifu) - 1), dtype=int)

    # save FITS file
    if len(prefix_intermediate_fits) > 0:
        outfile = f"{prefix_intermediate_fits}_ifu_3D_method0.fits"
        if image3d_method0_ifu.max() <= 65535:
            bitpix = 16
            hdu = fits.PrimaryHDU(image3d_method0_ifu.astype(np.uint16))
        else:
            # use float to avoid saturation problem
            logger.warning(f"The maximum value in {outfile} is greater than 65535.")
            logger.warning("Saving the image using float32 to avoid saturation.")
            bitpix = -32
            hdu = fits.PrimaryHDU(image3d_method0_ifu.astype(np.float32))
        pos0 = len(hdu.header) - 1
        hdu.header.extend(wcs_to_header_using_cd_keywords(wcs3d), update=True)
        hdu.header.update(header_keys)
        hdu.header.insert(pos0, ("COMMENT", "FITS (Flexible Image Transport System) format is defined in 'Astronomy"))
        hdu.header.insert(
            pos0 + 1, ("COMMENT", "and Astrophysics', volume 376, page 359; bibcode: 2001A&A...376..359H")
        )
        hdul = fits.HDUList([hdu])
        logger.info(f"Saving file: {outfile} (BITPIX={bitpix})")
        hdul.writeto(f"{Path(output_dir) / outfile}", overwrite=True)
