"""Module that contains parsers that do not fit in other categories.

Contains the following parsers:
- :class:`Meta2DParser`
- :class:`Meta3DParser`
- :class:`RunInfoParser`
- :class:`CRTHitParser`
- :class:`TriggerParser`
"""

import numpy as np

from spine.data import CRTHit, Flash, Meta, RunInfo, Trigger
from spine.geo import GeoManager
from spine.utils.conditional import larcv
from spine.utils.optical import FlashMerger

from .base import ParserBase
from .data import ParserObjectList

__all__ = [
    "MetaParser",
    "RunInfoParser",
    "FlashParser",
    "CRTHitParser",
    "TriggerParser",
]


class MetaParser(ParserBase):
    """Get the metadata information to translate into real world coordinates.

    Each entry in a dataset is a cube, where pixel/voxel coordinates typically
    go from 0 to some integer N in each dimension. If you wish to translate
    these pixel/voxel coordinates back into real world coordinates, you can use
    the output of this parser to compute it.

    .. code-block. yaml

        schema:
          meta:
            parser: meta
            sparse_event: sparse3d_pcluster
    """

    # Name of the parser (as specified in the configuration)
    name = "meta"

    # Alternative allowed names of the parser
    aliases = ("meta2d", "meta3d")

    # Type of object(s) returned by the parser
    returns = "object"

    # Overlay strategy for the objects returned by the parser
    overlay = "match"

    def __call__(self, trees):
        """Parse one entry.

        Parameters
        ----------
        trees : dict
            Dictionary which maps each data product name to a LArCV object
        """
        return self.process(**self.get_input_data(trees))

    def __init__(self, projection_id=None, **kwargs):
        """Initialize the parser.

        Parameters
        ----------
        projection_id : int, optional
            Projection ID to get the 2D image from (if fetching from 2D)
        **kwargs : dict, optional
            data product arguments to be passed to the `process` function
        """
        # Initialize the parent class
        super().__init__(**kwargs)

        # Store the revelant attributes
        self.projection_id = projection_id

    def process(self, sparse_event=None, cluster_event=None):
        """Fetches the metadata from one object that has it.

        Parameters
        ----------
        sparse_event : Union[larcv.EventSparseTensor2D
                             larcv.EventSparseTensor3D], optional
            Tensor which contains the metadata information as an attribute
        cluster_event : Union[larcv.EventClusterPixel2D,
                              larcv.EventClusterVoxel3D], optional
            Cluster which contains the metadata information as an attribute

        Returns
        -------
        Meta
            Metadata information for one image
        """
        # Check on the input, pick a source for the metadata
        if (sparse_event is None) == (cluster_event is None):
            raise ValueError(
                "Must specify exactly one of `sparse_event` or `cluster_event`."
            )

        ref_event = sparse_event if sparse_event is not None else cluster_event
        assert ref_event is not None  # For the type checker's benefitV

        # Fetch a specific projection, if needed
        if isinstance(
            ref_event, (larcv.EventSparseTensor2D, larcv.EventClusterPixel2D)
        ):
            ref_event = ref_event.sparse_tensor_2d(self.projection_id)

        return Meta.from_larcv(ref_event.meta())


class RunInfoParser(ParserBase):
    """Parse run information (run, subrun, event number).

    .. code-block. yaml

        schema:
          run_info:
            parser: run_info
            sparse_event: sparse3d_pcluster
    """

    # Name of the parser (as specified in the configuration)
    name = "run_info"

    # Type of object(s) returned by the parser
    returns = "object"

    # Overlay strategy for the objects returned by the parser
    overlay = "cat"

    def __call__(self, trees):
        """Parse one entry.

        Parameters
        ----------
        trees : dict
            Dictionary which maps each data product name to a LArCV object
        """
        return self.process(**self.get_input_data(trees))

    def process(self, sparse_event=None, cluster_event=None):
        """Fetches the run information from one object that has it.

        Parameters
        ----------
        sparse_event : Union[larcv.EventSparseTensor2D
                             larcv.EventSparseTensor3D], optional
            Tensor which contains the run information as an attribute
        cluster_event : Union[larcv.EventClusterPixel2D,
                              larcv.EventClusterVoxel3D], optional
            Cluster which contains the run information as an attribute

        Returns
        -------
        RunInfo
            Run information object
        """
        # Check on the input, pick a source for the run information
        assert (sparse_event is not None) ^ (
            cluster_event is not None
        ), "Must specify either `sparse_event` or `cluster_event`."
        ref_event = sparse_event if sparse_event is not None else cluster_event

        return RunInfo.from_larcv(ref_event)


class FlashParser(ParserBase):
    """Copy construct Flash and return an array of `Flash`.

    This parser also takes care of flashes that have been split between their
    respective optical volumes, provided a `flash_event_list`. This parser
    assumes that the trees are provided in order of the volume ID they
    correspond to.

    .. code-block. yaml
        schema:
          flashes:
            parser: flash
            flash_event_list:
              - flash_cryoE
              - flash_cryoW
    """

    # Name of the parser (as specified in the configuration)
    name = "flash"

    # Alternative allowed names of the parser
    aliases = ("opflash",)

    # Type of object(s) returned by the parser
    returns = "object_list"

    def __init__(self, merge=None, **kwargs):
        """Initialize the flash parser.

        Parameters
        ----------
        merge : dict, optional
            Flash merging configuration
        **kwargs : dict, optional
            data product arguments to be passed to the `process` function
        """
        # Initialize the parent class
        super().__init__(**kwargs)

        # Initialize the flash merging class, if needed
        self.merger = None
        if merge is not None:
            self.merger = FlashMerger(**merge)

        # Fetch the geometry information to resize the flash objects
        geo = GeoManager.get_instance()
        if geo.optical is not None:
            self.num_channels_per_volume = geo.optical.num_channels_per_volume
        else:
            raise ValueError(
                "Optical geometry not found, required to resize the flash objects."
            )

    def __call__(self, trees):
        """Parse one entry.

        Parameters
        ----------
        trees : dict
            Dictionary which maps each data product name to a LArCV object
        """
        return self.process(**self.get_input_data(trees))

    def process(self, flash_event=None, flash_event_list=None):
        """Fetches the list of optical flashes.

        Parameters
        -------------
        flash_event : larcv.EventFlash, optional
            Optical flash event which contains a list of flash objects
        flash_event_list : List[larcv.EventFlash], optional
            List of optical flash events, each a list of flash objects

        Returns
        -------
        List[Flash]
            List of optical flash objects
        """
        # Check on the input
        assert (flash_event is not None) ^ (
            flash_event_list is not None
        ), "Must specify either `flash_event` or `flash_event_list`."

        # Parse flash objects
        flashes = []
        if flash_event is not None:
            # If there is a single flash event, parse it as is
            flash_list = flash_event.as_vector()
            flashes = [Flash.from_larcv(larcv.Flash(f)) for f in flash_list]

        elif flash_event_list is not None:
            # Otherwise, set the volume ID of the flash to the source index
            # and count the flash index from 0 to the largest number
            idx = 0
            for volume_id, flash_event in enumerate(flash_event_list):
                for f in flash_event.as_vector():
                    # Cast and update attributes
                    flash = Flash.from_larcv(f)
                    flash.id = idx
                    flash.volume_id = volume_id

                    # Append, increment counter
                    flashes.append(flash)
                    idx += 1

        # Resize the flash objects to match the geometry
        for flash in flashes:
            pe_per_ch = np.zeros(
                self.num_channels_per_volume, dtype=flash.pe_per_ch.dtype
            )
            if len(flash.pe_per_ch) > len(pe_per_ch):
                # If the flash spans > 1 optical volume, reshape
                lower = flash.volume_id * len(pe_per_ch)
                upper = (flash.volume_id + 1) * len(pe_per_ch)
                pe_per_ch = flash.pe_per_ch[lower:upper]

            elif len(flash.pe_per_ch) < len(pe_per_ch):
                # Pad if it does not fill the full length
                pe_per_ch[: len(flash.pe_per_ch)] = flash.pe_per_ch

            else:
                pe_per_ch = flash.pe_per_ch

            flash.pe_per_ch = pe_per_ch

        # If requested, merge flashes which match in time
        if self.merger is not None:
            flashes, _ = self.merger(flashes)

        return ParserObjectList(flashes, Flash())


class CRTHitParser(ParserBase):
    """Copy construct CRTHit and return an array of `CRTHit`.

    .. code-block. yaml
        schema:
          crthits:
            parser: crthit
            crthit_event: crthit_crthit
    """

    # Name of the parser (as specified in the configuration)
    name = "crthit"

    # Type of object(s) returned by the parser
    returns = "object_list"

    def __call__(self, trees):
        """Parse one entry.

        Parameters
        ----------
        trees : dict
            Dictionary which maps each data product name to a LArCV object
        """
        return self.process(**self.get_input_data(trees))

    def process(self, crthit_event):
        """Fetches the list of CRT hits.

        Parameters
        ----------
        crthit_event : larcv.CRTHitEvent

        Returns
        -------
        List[CRTHit]
            List of CRT hit objects
        """
        # Output as a list of LArCV CRT hit objects
        crthit_list = crthit_event.as_vector()
        crthits = [CRTHit.from_larcv(larcv.CRTHit(c)) for c in crthit_list]

        return ParserObjectList(crthits, CRTHit())


class TriggerParser(ParserBase):
    """Copy construct Trigger and return a `Trigger`.

    .. code-block. yaml
        schema:
          trigger:
            parser: trigger
            trigger_event: trigger_base
    """

    # Name of the parser (as specified in the configuration)
    name = "trigger"

    # Type of object(s) returned by the parser
    returns = "object"

    # Overlay strategy for the objects returned by the parser
    overlay = "cat"

    def __call__(self, trees):
        """Parse one entry.

        Parameters
        ----------
        trees : dict
            Dictionary which maps each data product name to a LArCV object
        """
        return self.process(**self.get_input_data(trees))

    def process(self, trigger_event):
        """Fetches the trigger information.

        Parameters
        ----------
        trigger_event : larcv.TriggerEvent

        Returns
        -------
        Trigger
            Trigger object
        """
        # Output as a trigger objects
        trigger = Trigger.from_larcv(larcv.Trigger(trigger_event))

        return trigger
