# Copyright (C) 2021 - 2025 ANSYS, Inc. and/or its affiliates.
# SPDX-License-Identifier: MIT
#
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""The lxp module contains classes and functions to simplify the interaction with ray data.

Ray data is provided as lpf file.
LPF files contain a set of simulated rays with all their intersections and properties.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Optional, Union

import ansys.api.speos.lpf.v2.lpf_file_reader_pb2 as lpf_file_reader__v2__pb2
import ansys.api.speos.lpf.v2.lpf_file_reader_pb2_grpc as lpf_file_reader__v2__pb2_grpc
from ansys.speos.core.generic.general_methods import graphics_required
from ansys.speos.core.project import Project, Speos

if TYPE_CHECKING:  # pragma: no cover
    from ansys.tools.visualization_interface import Plotter
try:
    from ansys.speos.core.generic.general_methods import run_if_graphics_required

    run_if_graphics_required(warning=True)
except ImportError as err:  # pragma: no cover
    raise err

ERROR_IDS = [7, 8, 9, 10, 11, 12, 13, 14, 15]
"""Intersection types indicating an error state."""

NO_ERROR_IDS = [0, 1, 2, 3, 4, 5, 6, 16, -7, -6, -5, -5, -4, -3, -2, -1]
"""Intersection types indicating a correct ray state."""


class RayPath:
    """Framework representing a singular ray path.

    Parameters
    ----------
    raypath : ansys.api.speos.lpf.v2.lpf_file_reader__v2__pb2.RayPath
        RayPath object
    sensor_contribution : bool
        Defines if sensor contributions are stored within the data.
        By default ``False``.
    """

    def __init__(
        self,
        raypath: lpf_file_reader__v2__pb2.RayPath,
        sensor_contribution: bool = False,
    ):
        self._nb_impacts = len(raypath.impacts)
        self._impacts = [[inter.x, inter.y, inter.z] for inter in raypath.impacts]
        self._wl = raypath.wavelengths[0]
        self._body_ids = raypath.body_context_ids
        self._face_ids = raypath.unique_face_ids
        self._last_direction = [
            raypath.lastDirection.x,
            raypath.lastDirection.y,
            raypath.lastDirection.z,
        ]
        self._intersection_type = raypath.interaction_statuses
        if sensor_contribution:
            self._sensor_contribution = [
                {
                    "sensor_id": sc.sensor_id,
                    "position": [sc.coordinates.x, sc.coordinates.y],
                }
                for sc in raypath.sensor_contributions
            ]
        else:
            self._sensor_contribution = None

    @property
    def nb_impacts(self) -> int:
        """Number of impacts contained in ray path.

        Returns
        -------
        int
            Number of impacts
        """
        return self._nb_impacts

    @property
    def impacts(self) -> list[list[float]]:
        """XYZ coordinates for each impact.

        Returns
        -------
        list[list[float]]
            list containing the impact coordinates [[x0,y0,z0],[x1,y1,z1],...]
        """
        return self._impacts

    @property
    def wl(self) -> float:
        """Wavelength of the ray.

        Returns
        -------
        float
            Wavelength in nm
        """
        return self._wl

    @property
    def body_ids(self) -> list[int]:
        """Body ID of interacted body for each impact.

        Returns
        -------
        list[int]
            List of body IDs for each impact.
        """
        return self._body_ids

    @property
    def face_ids(self) -> list[int]:
        """Face ID of interacted body for each impact.

        Returns
        -------
        list[int]
            List of face IDs for each impact.
        """
        return self._face_ids

    @property
    def last_direction(self) -> list[float]:
        """Last direction of the ray.

        Returns
        -------
        list[float]
            Last direction of the rays as list[x,y,z].
        """
        return self._last_direction

    @property
    def intersection_type(self) -> list[int]:
        """Intersection type of the ray for each impact.

        Returns
        -------
        list[int]
            Intersection type at each impact.

        Notes
        -----
        Available intersection types:

        - StatusAbsorbed = 0
        - StatusSpecularTransmitted = 1
        - StatusGaussianTransmitted = 2
        - StatusLambertianTransmitted = 3
        - StatusVolumicDiffused = 4
        - StatusJustEmitted = 5
        - StatusDiracTransmitted = 6
        - StatusError = 7
        - StatusErrorVolumicBodyNotClosed = 8
        - StatusErrorVolumeConflict = 9
        - StatusError2DTangency = 10
        - StatusError2DIntersect3DWarning = 11
        - StatusErrorNonOpticalMaterial = 12
        - StatusErrorIntersection = 13
        - StatusErrorNonOpticalMaterialAtEmission = 14
        - StatusError3DTextureSupportTangency = 15
        - StatusLast = 16
        - StatusFirst = -7
        - StatusDiracReflected = -6
        - StatusReserved = -5
        - StatusGrinStep = -4
        - StatusLambertianReflected = -3
        - StatusGaussianReflected = -2
        - StatusSpecularReflected = -1
        """
        return self._intersection_type

    @property
    def sensor_contribution(self) -> Union[None, list[dict]]:
        """Provide the sensor contribution information for each sensor.

        Returns
        -------
        Union[None, list[dict]]
            If no sensor contribution, None will be returned. If there is sensor contribution, \
            a dictionary with the following information is returned:\
            {“sensor_id”: sc.sensor_id,
            “position”: [sc.coordinates.x, sc.coordinates.y]}
        """
        return self._sensor_contribution

    def get(self, key=""):
        """Retrieve any information from the RayPath object.

        Parameters
        ----------
        key : str
            Name of the property.

        Returns
        -------
        property
            Values/content of the associated property.
        """
        data = {k: v.fget(self) for k, v in RayPath.__dict__.items() if isinstance(v, property)}
        if key == "":
            return data
        elif data.get(key):
            return data.get(key)
        else:
            print("Used key: {} not found in key list: {}.".format(key, data.keys()))

    def __str__(self):
        """Create string representation of a RayPath."""
        return str(self.get())


class LightPathFinder:
    """Define an interface to read LPF files.

    LPF files contain a set of simulated rays including their intersections and properties.

    Parameters
    ----------
    speos : ansys.speos.core.speos.Speos
        Speos Session (connected to Speos gRPC server).
    path : str
        Path to the LPF file to be opened.

    """

    def __init__(self, speos: Speos, path: str):
        self.client = speos.client
        """Speos instance client"""
        self._stub = lpf_file_reader__v2__pb2_grpc.LpfFileReader_MonoStub(self.client.channel)
        self.__open(path)
        self._data = self._stub.GetInformation(
            lpf_file_reader__v2__pb2.GetInformation_Request_Mono()
        )
        self._nb_traces = self._data.nb_of_traces
        self._nb_xmps = self._data.nb_of_xmps
        self._has_sensor_contributions = self._data.has_sensor_contributions
        self._sensor_names = self._data.sensor_names
        self._rays = self.__parse_traces()
        self._filtered_rays = []

    @property
    def nb_traces(self) -> int:
        """Number of light paths within LPF data set."""
        return self._nb_traces

    @property
    def nb_xmps(self) -> int:
        """Number of sensors involved within LPF data set."""
        return self._nb_xmps

    @property
    def has_sensor_contributions(self) -> bool:
        """Define if a LPF file contains information regarding the sensor contribution."""
        return self._has_sensor_contributions

    @property
    def sensor_names(self) -> list[str]:
        """List of involved sensor names."""
        return self._sensor_names

    @property
    def rays(self) -> list[RayPath]:
        """List ray paths within LPF file."""
        return self._rays

    @property
    def filtered_rays(self) -> list[RayPath]:
        """List of filtered ray paths."""
        return self._filtered_rays

    def __str__(self):
        """Create string representation of LightPathFinder."""
        return str(
            {
                k: v.fget(self)
                for k, v in LightPathFinder.__dict__.items()
                if isinstance(v, property) and "rays" not in k
            }
        )

    def __open(self, path: str):
        """Open LPF file.

        Parameters
        ----------
        path : str
            Path to file
        """
        self._stub.InitLpfFileName(
            lpf_file_reader__v2__pb2.InitLpfFileName_Request_Mono(lpf_file_uri=path)
        )

    def __parse_traces(self) -> list[RayPath]:
        """Read all ray paths from lpf dataset.

        Returns
        -------
            list[script.RayPath]

        """
        raypaths = []
        for rp in self._stub.Read(lpf_file_reader__v2__pb2.Read_Request_Mono()):
            raypaths.append(RayPath(rp, self._has_sensor_contributions))
        return raypaths

    def __filter_by_last_intersection_types(self, options: list[int], new=True):
        """Filter ray paths based on last intersection types.

        Populate filtered_rays property.
        """
        if new:
            self._filtered_rays = []
            for ray in self._rays:
                if int(ray.intersection_type[-1]) in options:
                    self._filtered_rays.append(ray)
        else:
            temp_rays = self._filtered_rays
            self._filtered_rays = []
            for ray in temp_rays:
                if int(ray.intersection_type[-1]) in options:
                    self._filtered_rays.append(ray)

    def filter_by_face_ids(self, options: list[int], new=True) -> LightPathFinder:
        """Filter ray paths based on face IDs and populates filtered_rays property.

        Parameters
        ----------
        options : list[int]
            List of face IDs.
        new : bool
            Define if a new filter is created or an existing filter is filtered.

        Returns
        -------
        ansys.speos.core.lxp.LightPathFinder
            LightPathFinder Instance.
        """
        if new:
            self._filtered_rays = []
            for ray in self._rays:
                if any([faceid in options for faceid in ray.face_ids]):
                    self._filtered_rays.append(ray)
        else:
            temp_rays = self._filtered_rays
            self._filtered_rays = []
            for ray in temp_rays:
                if any([faceid in options for faceid in ray.face_ids]):
                    self._filtered_rays.append(ray)
        return self

    def filter_by_body_ids(self, options: list[int], new=True) -> LightPathFinder:
        """Filter ray paths based on body IDs and populates filtered_rays property.

        Parameters
        ----------
        options : list[int]
            List of body IDs.
        new : bool
            Define if a new filter is created or an existing filter is filtered.

        Returns
        -------
        ansys.speos.core.lxp.LightPathFinder
            LightPathFinder Instance.
        """
        if new:
            self._filtered_rays = []
            for ray in self._rays:
                if any([body_id in options for body_id in ray.body_ids]):
                    self._filtered_rays.append(ray)
        else:
            temp_rays = self._filtered_rays
            self._filtered_rays = []
            for ray in temp_rays:
                if any([body_id in options for body_id in ray.body_ids]):
                    self._filtered_rays.append(ray)
        return self

    def filter_error_rays(self) -> LightPathFinder:
        """Filter ray paths and only shows rays in error.

        Returns
        -------
        ansys.speos.core.lxp.LightPathFinder
            LightPathFinder Instance.
        """
        self.__filter_by_last_intersection_types(options=ERROR_IDS)
        return self

    def remove_error_rays(self) -> LightPathFinder:
        """Filter rays and only shows rays not in error.

        Returns
        -------
        ansys.speos.core.lxp.LightPathFinder
            LightPathFinder Instance.
        """
        self.__filter_by_last_intersection_types(options=NO_ERROR_IDS)
        return self

    @staticmethod
    @graphics_required
    def __add_ray_to_pv(plotter: Plotter, ray: RayPath, max_ray_length: float):
        """Add a ray to pyvista plotter.

        Parameters
        ----------
        plotter : Plotter
            Ansys plotter object to which rays should be added.
        ray : script.RayPath
            RayPath object which contains ray information to be added.
        max_ray_length : float
            Length of the last ray.
        """
        import pyvista as pv

        temp = ray.impacts.copy()
        if not 7 <= ray.intersection_type[-1] <= 15:
            temp.append(
                [
                    ray.impacts[-1][0] + max_ray_length * ray.last_direction[0],
                    ray.impacts[-1][1] + max_ray_length * ray.last_direction[1],
                    ray.impacts[-1][2] + max_ray_length * ray.last_direction[2],
                ]
            )
        if len(ray.impacts) > 2:
            mesh = pv.MultipleLines(temp)
        else:
            mesh = pv.Line(temp[0], temp[1])
        plotter.plot(mesh, color=wavelength_to_rgb(ray.wl), line_width=2)

    @graphics_required
    def preview(
        self,
        nb_ray: int = 100,
        max_ray_length: float = 50.0,
        ray_filter: bool = False,
        project: Project = None,
        screenshot: Optional[Union[str, Path]] = None,
    ) -> LightPathFinder:
        """Preview LPF file with pyvista.

        Parameters
        ----------
        nb_ray : int
            Number of rays to be visualized.
        max_ray_length : float
            Length of last ray.
        ray_filter : bool
            Boolean to decide if filtered rays or all rays should be shown.
        project : ansys.speos.core.project.Project
            Speos Project/Geometry to be added to pyvista visualisation.
        screenshot : str or Path or ``None``
            Path to save a screenshot of the plotter. If defined Plotter will only create the
            screenshot

        Returns
        -------
        ansys.speos.core.lxp.LightPathFinder
            LightPathFinder Instance.

        Notes
        -----
        Please use the ``q``-key to close the plotter as some
        operating systems (namely Windows) will experience issues
        saving a screenshot if the exit button in the GUI is pressed.
        """
        from ansys.tools.visualization_interface import Plotter

        if ray_filter:
            if len(self._filtered_rays) > 0:
                temp_rays = self._filtered_rays
            else:
                print("no filtered rays")
                temp_rays = self._rays
        else:
            temp_rays = self._rays
        if not project:
            plotter = Plotter()
            if nb_ray > len(temp_rays):
                for ray in temp_rays:
                    self.__add_ray_to_pv(plotter, ray, max_ray_length)
            else:
                for i in range(nb_ray):
                    self.__add_ray_to_pv(plotter, temp_rays[i], max_ray_length)
        else:
            plotter = project._create_preview(viz_args={"opacity": 0.5})
            if nb_ray > len(temp_rays):
                for ray in temp_rays:
                    self.__add_ray_to_pv(plotter, ray, max_ray_length)
            else:
                for i in range(nb_ray):
                    self.__add_ray_to_pv(plotter, temp_rays[i], max_ray_length)
        plotter.show(screenshot=screenshot)
        return self


def wavelength_to_rgb(wavelength: float, gamma: float = 0.8) -> [int, int, int, int]:
    """Convert a given wavelength of light to an approximate RGB color value.

    The wavelength must be given in nanometers in the range from 380 nm to 750 nm.
    Based on the code from http://www.physics.sfasu.edu/astro/color/spectra.html

    Parameters
    ----------
    wavelength : float
        Wavelength in nanometer between 380-750 nm
    gamma : float
        Gamma value.
        By default : ``0.8``
    """
    wavelength = float(wavelength)
    if 380 <= wavelength <= 440:
        attenuation = 0.3 + 0.7 * (wavelength - 380) / (440 - 380)
        r = ((-(wavelength - 440) / (440 - 380)) * attenuation) ** gamma
        g = 0.0
        b = (1.0 * attenuation) ** gamma
    elif 440 <= wavelength <= 490:
        r = 0.0
        g = ((wavelength - 440) / (490 - 440)) ** gamma
        b = 1.0
    elif 490 <= wavelength <= 510:
        r = 0.0
        g = 1.0
        b = (-(wavelength - 510) / (510 - 490)) ** gamma
    elif 510 <= wavelength <= 580:
        r = ((wavelength - 510) / (580 - 510)) ** gamma
        g = 1.0
        b = 0.0
    elif 580 <= wavelength <= 645:
        r = 1.0
        g = (-(wavelength - 645) / (645 - 580)) ** gamma
        b = 0.0
    elif 645 <= wavelength <= 750:
        attenuation = 0.3 + 0.7 * (750 - wavelength) / (750 - 645)
        r = (1.0 * attenuation) ** gamma
        g = 0.0
        b = 0.0
    else:
        r = 0.0
        g = 0.0
        b = 0.0
    r *= 255
    g *= 255
    b *= 255
    return [int(r), int(g), int(b), 255]
