"""
PSF image analysis for ray tracing results.

Processes photon position lists to compute image centroids,
PSF containment diameters (D80 and other fractions),
effective area, and cumulative radial distributions.
"""

import gzip
import logging
import tempfile
from math import fabs, pi, sqrt
from pathlib import Path

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np

from simtools import settings
from simtools.job_execution import job_manager
from simtools.utils.general import collect_kwargs, set_default_kwargs


class PSFImage:
    """
    Image composed of list of photon positions (2D).

    Load photon list from sim_telarray file and compute centroids, psf containers, effective area,
    as well as plot the image as a 2D histogram.
    Internal units: photon positions in cm internally.

    Parameters
    ----------
    focal_length: float
        Focal length of the system in cm. If not given, PSF can only be computed in cm.
    total_scattered_area: float
        Scatter area of all photons in cm^2. If not given, effective area cannot be computed.
    containment_fraction: float
        Containment fraction for PSF calculation.
    """

    __PSF_RADIUS = "Radius [cm]"
    __PSF_CUMULATIVE = "Cumulative PSF"

    def __init__(
        self,
        focal_length=None,
        total_scattered_area=None,
        containment_fraction=None,
    ):
        """Initialize PSFImage class."""
        self._logger = logging.getLogger(__name__)

        self._total_photons = None
        self._number_of_detected_photons = None
        self._effective_area = None
        self.photon_pos_x = []
        self.photon_pos_y = []
        self.photon_r = []
        self.centroid_x = None
        self.centroid_y = None
        self._total_area = total_scattered_area
        self.cos_camera_rotation = 1.0
        self.sin_camera_rotation = 0.0
        self._stored_psf = {}
        try:
            self._cm_to_deg = 180.0 / pi / focal_length if focal_length is not None else None
        except ZeroDivisionError:
            self._cm_to_deg = None
            self._logger.warning("Focal length is zero; no conversion from cm to deg possible.")
        self._containment_fraction = containment_fraction

    def process_photon_list(self, photon_file, use_rx):
        """
        Read and process a photon list file generated by sim_telarray.

        Extracts photon positions, computes centroids and effective area.
        Can use either sim_telarray's rx binary or internal computation.

        Parameters
        ----------
        photons_file: str
            Name of sim_telarray file with photon list.
        use_rx: bool
            Use the RX method for analysis.
        """
        if use_rx:
            self._process_simtel_file_using_rx(photon_file)
        else:
            self.read_photon_list_from_simtel_file(photon_file)

    def _process_simtel_file_using_rx(self, photon_file):
        """
        Process a sim_telarray file with photon lists using the RX method.

        Parameters
        ----------
        photons_file: str
            Name of sim_telarray file with photon list.
        """
        rx_command = [
            f"{settings.config.sim_telarray_path}/bin/rx",
            "-f",
            f"{self._containment_fraction:.2f}",
            "-v",
        ]

        try:
            with tempfile.NamedTemporaryFile(mode="wb", delete=False) as tmp_file:
                with gzip.open(photon_file, "rb") as gz:
                    tmp_file.write(gz.read())
                tmp_file_path = tmp_file.name

            try:
                with open(tmp_file_path, "rb") as tmp_f:
                    result = job_manager.submit(rx_command, stdin=tmp_f)

                data_lines = [
                    line.strip()
                    for line in result.stdout.strip().split("\n")
                    if line.strip() and not line.startswith("#")
                ]

                if not data_lines:
                    raise IndexError("No data line found in RX output")

                rx_output = data_lines[-1].split()
                self.set_psf(
                    2 * float(rx_output[0]), fraction=self._containment_fraction, unit="cm"
                )
                self.centroid_x = float(rx_output[1])
                self.centroid_y = float(rx_output[2])
                self._effective_area = float(rx_output[5])

            finally:
                Path(tmp_file_path).unlink(missing_ok=True)

        except FileNotFoundError as e:
            raise FileNotFoundError(f"Photon list file not found: {photon_file}") from e
        except (IndexError, ValueError) as e:
            raise type(e)(
                f"Invalid RX output format: {locals().get('rx_output', 'unknown')}"
            ) from e

    def read_photon_list_from_simtel_file(self, photons_file):
        """
        Read photon list file generated by sim_telarray and store the photon positions (2D).

        Computes image centroid as mean of photon positions,
        effective area as (detected_photons * total_scattered_area) / total_photons,
        and sorts radial distances from centroid for efficient PSF calculation.

        Parameters
        ----------
        photons_file: str
            Name of sim_telarray file with photon list.

        Raises
        ------
        RuntimeError
            If photon positions X and Y are not compatible or are empty.

        """
        self._logger.info(f"Reading sim_telarray photon file {photons_file}")
        self._total_photons = 0
        if Path(photons_file).suffix == ".gz":
            file_open_function = gzip.open
        else:
            file_open_function = open
        with file_open_function(photons_file, "rb") as f:
            for line in f:
                self._process_simtel_line(line)

        if not self._is_photon_positions_ok():
            raise RuntimeError(
                f"Problems reading sim_telarray photon file {photons_file} - invalid data"
            )

        self.centroid_x = np.mean(self.photon_pos_x)
        self.centroid_y = np.mean(self.photon_pos_y)
        self._number_of_detected_photons = len(self.photon_pos_x)
        self._effective_area = (
            self._number_of_detected_photons * self._total_area / self._total_photons
        )
        self.photon_r = np.sort(
            np.sqrt(
                (self.photon_pos_x - self.centroid_x) ** 2
                + (self.photon_pos_y - self.centroid_y) ** 2
            )
        )

    def _is_photon_positions_ok(self):
        """
        Verify if the photon positions are ok.

        Returns
        -------
        bool
            True if photon positions are ok, False if they are not.
        """
        return len(self.photon_pos_x) > 0 and len(self.photon_pos_x) == len(self.photon_pos_y)

    def _process_simtel_line(self, line):
        """
        Supporting function to read_photon_list_from_simtel_file.

        Parameters
        ----------
        line: str
            A line from the photon list file generated by sim_telarray.
        """
        words = line.split()

        if b"falling on an area of" in line:
            self._total_photons += int(words[4])
            total_area_in_file = float(words[14])
            if self._total_area is None:
                self._total_area = total_area_in_file
            elif total_area_in_file != self._total_area:
                self._logger.warning(
                    "Conflicting value of the total area found"
                    f" {self._total_area} != {total_area_in_file}"
                    " - Keeping the original value"
                )
        elif b"Camera rotation angle" in line:
            camera_rotation_angle = np.deg2rad(float(words[5]))
            self.cos_camera_rotation = np.cos(camera_rotation_angle)
            self.sin_camera_rotation = np.sin(camera_rotation_angle)
        elif b"#" in line or len(words) == 0:
            return
        else:
            # Photon positions from cols 2 and 3; apply camera rotation
            x, y = float(words[2]), float(words[3])
            self.photon_pos_x.append(x * self.cos_camera_rotation - y * self.sin_camera_rotation)
            self.photon_pos_y.append(y * self.cos_camera_rotation + x * self.sin_camera_rotation)

    def get_effective_area(self, tel_transmission=1.0):
        """
        Return effective area pre calculated.

        Parameters
        ----------
        telescope_transmission : float
            Telescope transmission parameter.

        Returns
        -------
        float
            Pre-calculated effective area. None if it could not be calculated (e.g because the total
            scattering area was not set).
        """
        if "_effective_area" in self.__dict__ and self._effective_area is not None:
            return self._effective_area * tel_transmission

        self._logger.error("Effective Area could not be calculated")
        return None

    def set_effective_area(self, value):
        """
        Set effective area.

        Parameters
        ----------
        value: float
            Effective area

        """
        self._effective_area = value

    def get_psf(self, fraction=0.8, unit="cm"):
        """
        Return PSF containment diameter.

        Computes (if not cached) the diameter containing the specified fraction
        of photons using an iterative algorithm or brute-force scanning.

        Parameters
        ----------
        fraction: float
            Fraction of photons within the containing radius.
        unit: str
            'cm' or 'deg'. 'deg' will not work if focal length was not set.

        Returns
        -------
        float:
            Containing diameter for a certain intensity fraction (PSF).

        """
        if unit == "deg" and self._cm_to_deg is None:
            self._logger.error("PSF cannot be computed in deg because focal length is not set")
            return None
        if fraction not in self._stored_psf:
            self._compute_psf(fraction)
        unit_factor = 1 if unit == "cm" else self._cm_to_deg
        return self._stored_psf[fraction] * unit_factor

    def set_psf(self, value, fraction=0.8, unit="cm"):
        """
        Set PSF calculated from other methods.

        Parameters
        ----------
        value: float
            PSF value to be set
        fraction: float
            Fraction of photons within the containing radius.
        unit: str
            'cm' or 'deg'. 'deg' will not work if focal length was not set.
        """
        if unit == "deg" and self._cm_to_deg is None:
            self._logger.error("PSF cannot be set in deg because focal length is not set")
            return
        unit_factor = 1 if unit == "cm" else 1.0 / self._cm_to_deg
        self._stored_psf[fraction] = value * unit_factor

    def _compute_psf(self, fraction):
        """
        Compute and store PSF.

        Parameters
        ----------
        fraction: float
            Fraction of photons within the containing radius
        """
        self._stored_psf[fraction] = self._find_psf(fraction)

    def _find_psf(self, fraction):
        """
        Find PSF containment diameter using iterative gradient-based algorithm.

        Computes the radius containing a specified fraction of photons by iterating
        from an initial estimate (1.5 * radius_sigma) with adaptive step size.
        Falls back to brute-force scanning if convergence fails after 1000 iterations.

        Parameters
        ----------
        fraction: float
            Fraction of photons within the containing radius.

        Returns
        -------
        float:
            Diameter of the circular container with a certain fraction of the photons.

        """
        self._logger.debug(f"Finding PSF for fraction = {fraction}")

        x_pos_sq = [i**2 for i in self.photon_pos_x]
        y_pos_sq = [i**2 for i in self.photon_pos_y]
        x_pos_sig = sqrt(np.mean(x_pos_sq) - self.centroid_x**2)
        y_pos_sig = sqrt(np.mean(y_pos_sq) - self.centroid_y**2)
        radius_sig = sqrt(x_pos_sig**2 + y_pos_sig**2)

        target_number = fraction * self._number_of_detected_photons
        current_radius = 1.5 * radius_sig
        start_number = self._sum_photons_in_radius(current_radius)
        scale = 0.5 * sqrt(current_radius * current_radius / start_number)
        delta_number = start_number - target_number
        n_iter = 0
        max_iter = 1000
        tolerance = self._number_of_detected_photons / 1000.0
        found_radius = False
        while not found_radius and n_iter < max_iter:
            n_iter += 1
            dr = -delta_number * scale / sqrt(target_number)
            while current_radius + dr < 0:
                dr *= 0.5
            current_radius += dr
            current_number = self._sum_photons_in_radius(current_radius)
            delta_number = current_number - target_number
            found_radius = fabs(delta_number) < tolerance

        if found_radius:
            return 2 * current_radius

        self._logger.warning("Could not find PSF efficiently - trying by scanning")
        return 2 * self._find_radius_by_scanning(target_number, radius_sig)

    def _find_radius_by_scanning(self, target_number, radius_sig):
        """
        Find PSF containment radius by brute-force scanning.

        Performs two-pass scan with progressively finer step sizes
        (0.1*radius_sig, then 0.005*radius_sig) to find the radius
        containing the target number of photons.

        Parameters
        ----------
        target_number: float
            Number of photons inside the diameter to be found.
        radius_sig: float
            Sigma of the radius to be used as scale.

        Returns
        -------
        float:
            Radius of the circle with target_number photons inside.
        """
        self._logger.debug("Finding PSF by scanning")

        def scan(dr, rad_min, rad_max):
            """
            Scan the image from rad_min to rad_max until it finds target_number photons inside.

            Scanning is done in steps of dr.

            Returns
            -------
            (float, float, float):
                Average radius, min radius, max radius of the interval where target_number photons
                are inside.

            Raises
            ------
            RuntimeError
                if radius is not found (found_radius is False)
            """
            r0, r1 = rad_min, rad_min + dr
            found_radius = False
            while not found_radius:
                s0, s1 = self._sum_photons_in_radius(r0), self._sum_photons_in_radius(r1)
                if s0 < target_number <= s1:
                    found_radius = True
                    break
                if r1 > rad_max:
                    break
                r0 += dr
                r1 += dr
            if found_radius:
                return (r0 + r1) / 2, r0, r1

            self._logger.error("Could not find PSF by scanning")
            raise RuntimeError

        # Run scan few times with smaller dr to optimize search.
        # Step 0
        radius, rad_min, rad_max = scan(0.1 * radius_sig, 0, 4 * radius_sig)
        # Step 1
        radius, rad_min, rad_max = scan(0.005 * radius_sig, rad_min, rad_max)
        return radius

    def _sum_photons_in_radius(self, radius):
        """Return number of photons inside radius using binary search on sorted radial distances."""
        return np.searchsorted(self.photon_r, radius)

    def get_image_data(self, centralized=True):
        """
        Provide image data (2D photon positions in cm) as lists.

        Parameters
        ----------
        centralized: bool
            Centroid of the image is set to (0, 0) if True.

        Returns
        -------
        (x, y), the photons positions in cm.
        """
        if centralized:
            x_pos_data = np.array(self.photon_pos_x) - self.centroid_x
            y_pos_data = np.array(self.photon_pos_y) - self.centroid_y
        else:
            x_pos_data = np.array(self.photon_pos_x)
            y_pos_data = np.array(self.photon_pos_y)
        d_type = {"names": ("X", "Y"), "formats": ("f8", "f8")}
        result = np.recarray((len(x_pos_data),), dtype=d_type)
        result.X = x_pos_data
        result.Y = y_pos_data
        return result

    def plot_image(self, centralized=True, file_name=None, **kwargs):
        """
        Plot 2D image as histogram (in cm).

        Parameters
        ----------
        centralized: bool
            Centroid of the image is set to (0, 0) if True.
        **kwargs:
            image_* for the histogram plot and psf_* for the psf circle.
        """
        data = self.get_image_data(centralized)

        kwargs = set_default_kwargs(
            kwargs,
            image_bins=150,
            image_cmap=plt.cm.gist_heat_r,
            psf_color="k",
            psf_fill=False,
            psf_lw=2,
            psf_ls="--",
        )
        kwargs_for_image = collect_kwargs("image", kwargs)
        kwargs_for_psf = collect_kwargs("psf", kwargs)

        ax = plt.gca()
        ax.set_xlabel("X Position (cm)")
        ax.set_ylabel("Y Position (cm)")
        ax.hist2d(data["X"], data["Y"], **kwargs_for_image)
        ax.set_aspect("equal", "box")

        # PSF circle (containment fraction)
        fraction = self._containment_fraction if self._containment_fraction is not None else 0.8
        center = (0, 0) if centralized else (self.centroid_x, self.centroid_y)
        circle = plt.Circle(center, self.get_psf(fraction) / 2, **kwargs_for_psf)
        ax.add_artist(circle)

        ax.axhline(0, color="k", linestyle="--", zorder=3, linewidth=0.5)
        ax.axvline(0, color="k", linestyle="--", zorder=3, linewidth=0.5)

        if file_name is not None:
            plt.savefig(file_name)
            plt.close()

    def get_cumulative_data(self, radius=None):
        """
        Compute cumulative intensity as function of radius from centroid.

        For each radius, counts the fraction of total photons within that distance
        using binary search on pre-sorted radial distances.

        Parameters
        ----------
        radius: array
            Array with radius calculate the cumulative PSF in distance units.

        Returns
        -------
        (radius, intensity)
        """
        if radius is not None:
            radius_all = radius.to(u.cm).value if isinstance(radius, u.Quantity) else radius
        else:
            fraction = self._containment_fraction if self._containment_fraction is not None else 0.8
            radius_all = list(np.linspace(0, 1.6 * self.get_psf(fraction), 30))
        intensity = [
            self._sum_photons_in_radius(rad) / self._number_of_detected_photons
            for rad in radius_all
        ]
        d_type = {
            "names": (self.__PSF_RADIUS, self.__PSF_CUMULATIVE),
            "formats": ("f8", "f8"),
        }
        result = np.recarray((len(radius_all),), dtype=d_type)
        result[self.__PSF_RADIUS] = radius_all
        result[self.__PSF_CUMULATIVE] = intensity

        return result

    def plot_cumulative(self, file_name=None, psf_diameter_cm=None, **kwargs):
        """Plot cumulative data (intensity vs radius).

        Parameters
        ----------
        file_name: str
            Name of the file to save the plot to.
        psf_diameter_cm: float
            PSF diameter value to be marked in the plot (in cm).
        **kwargs:
            Customization of line plot (e.g., color, linestyle, linewidth).
        """
        data = self.get_cumulative_data()
        fig, ax = plt.subplots(constrained_layout=True)
        ax.set_xlabel("Radius (cm)")
        ax.set_ylabel("Contained light %")
        ax.plot(data[self.__PSF_RADIUS], data[self.__PSF_CUMULATIVE], **kwargs)
        fraction = self._containment_fraction if self._containment_fraction is not None else 0.8
        ax.axvline(x=self.get_psf(fraction) / 2, color="b", linestyle="--", linewidth=1)
        if psf_diameter_cm is not None:
            ax.axvline(x=psf_diameter_cm / 2.0, color="r", linestyle="--", linewidth=1)
        if file_name is not None:
            fig.savefig(file_name)
            plt.close(fig)
