"""Interaction vertex reconstruction module."""

import numpy as np

from spine.post.base import PostBase
from spine.utils.globals import SHOWR_SHP, TRACK_SHP
from spine.utils.vertex import get_vertex

__all__ = ["VertexProcessor"]


class VertexProcessor(PostBase):
    """Reconstruct one vertex for each interaction in the provided list."""

    # Name of the post-processor (as specified in the configuration)
    name = "vertex"

    # Alternative allowed names of the post-processor
    aliases = ("reconstruct_vertex",)

    # Set of post-processors which must be run before this one is
    _upstream = ("direction",)

    def __init__(
        self,
        include_shapes=(SHOWR_SHP, TRACK_SHP),
        use_primaries=True,
        update_orientations=False,
        update_primaries=False,
        anchor_vertex=True,
        touching_threshold=2.0,
        angle_threshold=0.3,
        run_mode="both",
        truth_point_mode="points",
    ):
        """Initialize the vertex finder properties.

        Parameters
        ----------
        include_shapes : List[int], default [0, 1]
            List of semantic classes to consider for vertex reconstruction
        use_primaries : bool, default True
            If true, only considers primary particles to reconstruct the vertex
        update_orientations : bool, default False
            Use the reconstructed vertex to update track orientations
        update_primaries : bool, default False
            Use the reconstructed vertex to update primaries
        anchor_vertex : bool, default True
            If true, anchor the candidate vertex to particle objects,
            with the expection of interactions only composed of showers
        touching_threshold : float, default 2.0
            Maximum distance for two track points to be considered touching (cm)
        angle_threshold : float, default 0.3 radians
            Maximum angle between the vertex-to-start-point vector and a shower
            direction to consider that a shower originated from the vertex
        """
        # Initialize the parent class
        super().__init__("interaction", run_mode, truth_point_mode)

        # Store the relevant parameters
        self.include_shapes = include_shapes
        self.use_primaries = use_primaries
        self.update_orientations = update_orientations
        self.update_primaries = update_primaries
        self.anchor_vertex = anchor_vertex
        self.touching_threshold = touching_threshold
        self.angle_threshold = angle_threshold

    def process(self, data):
        """Reconstruct the vertex position for each interaction in one entry.

        Parameters
        ----------
        data : dict
            Dictionary of data products
        """
        # Loop over interaction objects
        for k in self.interaction_keys:
            for inter in data[k]:
                self.reconstruct_vertex_single(inter)

    def reconstruct_vertex_single(self, inter):
        """Post-processor which reconstructs one vertex for each interaction
        in the provided list.

        Parameters
        ----------
        inter : List[RecoInteraction, TruthInteraction]
            Reconstructed/truth interaction object
        """
        # Selected the set of particles to use as a basis for vertex prediction
        particles = [
            part
            for part in inter.particles
            if part.size > 0 and part.shape in self.include_shapes
        ]

        if self.use_primaries and any(part.is_primary for part in particles):
            particles = [part for part in particles if part.is_primary]

        if len(particles) == 0:
            particles = [part for part in inter.particles if part.size > 0]

        # Update the vertex if we have at least one particle to work with
        if len(particles) > 0:
            # Collapse particle objects to start, end points and directions
            start_points = np.vstack([part.start_point for part in particles])
            end_points = np.vstack([part.end_point for part in particles])
            directions = np.vstack([part.start_dir for part in particles])
            shapes = np.array([part.shape for part in particles])

            # Reconstruct the vertex for this interaction
            vtx, _ = get_vertex(
                start_points,
                end_points,
                directions,
                shapes,
                self.anchor_vertex,
                self.touching_threshold,
                return_mode=True,
            )

            # Assign it to the appropriate interaction attribute
            if not inter.is_truth:
                inter.vertex = vtx
            else:
                inter.reco_vertex = vtx

            # If requested, update track orientations on the basis of the vertex
            if not inter.is_truth and self.update_orientations:
                for part in inter.particles:
                    if part.shape != TRACK_SHP:
                        continue
                    elif np.linalg.norm(
                        inter.vertex - part.start_point
                    ) > np.linalg.norm(inter.vertex - part.end_point):
                        part.start_point, part.end_point = (
                            part.end_point,
                            part.start_point,
                        )
                        part.start_dir, part.end_dir = (-part.end_dir, -part.start_dir)

            # If requested, update primaries on the basis of the vertex
            if not inter.is_truth and self.update_primaries:
                for part in inter.particles:
                    if part.shape not in [SHOWR_SHP, TRACK_SHP]:
                        part.is_primary = False
                    elif (
                        np.linalg.norm(part.start_point - inter.vertex)
                        < self.touching_threshold
                    ):
                        part.is_primary = True
                    elif part.shape == SHOWR_SHP:
                        vec = part.start_point - inter.vertex
                        cosang = np.dot(vec / np.linalg.norm(vec), part.start_dir)
                        if np.arccos(cosang) < self.angle_threshold:
                            part.is_primary = True
