"""Descriptors for computing statistics on chunked data, compatible with ctapipe aggregators."""

import numpy as np
from ctapipe.containers import EventType
from ctapipe.core import Component
from ctapipe.core.traits import ComponentName
from ctapipe.monitoring.aggregator import BaseAggregator, BaseChunking
from traitlets import Enum

__all__ = [
    "TrivialDescriptor",
    "PercentageDescriptor",
    "RateCounter",
    "EventRateCounter",
]


class TrivialDescriptor(Component):
    """
    A descriptor that returns data without performing aggregation.

    This handles two use cases based on ``data_category``. For ``DL1_MONITORING``,
    pre-aggregated monitoring data with ``time_start``/``time_end`` and
    ``event_id_start``/``event_id_end`` is passed through unchanged. For
    ``DL0_MONITORING``, raw monitoring data with a ``time`` column is converted to
    ``time_start``/``time_end``.

    Parameters
    ----------
    chunk_size : None
        Must be None for TrivialDescriptor (no chunking)
    parent : Metric
        Parent metric instance
    """

    def __call__(self, table, col_name):
        """
        Return the input data, adapting format as needed.

        Parameters
        ----------
        table : astropy.table.Table
            Input table, either:
            - DL1_MONITORING: has time_start, time_end, event_id_start, event_id_end
            - DL0_MONITORING: has 'time' column
        col_name : str
            Name of the data column to preserve

        Returns
        -------
        astropy.table.Table
            Table with columns:
            - time_start, time_end: Time window boundaries
            - event_id_start, event_id_end: Event ID boundaries (NaN for DL0_MONITORING)
            - {col_name}: The data value
        """
        from astropy.table import Table

        if col_name not in table.colnames:
            raise ValueError(f"Input table must have a '{col_name}' column")

        # Case 1: Pre-aggregated data (already has time_start/time_end)
        if "time_start" in table.colnames and "time_end" in table.colnames:
            # Return as-is - data is already properly formatted
            return table

        # Case 2: Raw monitoring data with timestamps
        if "time" not in table.colnames:
            raise ValueError(
                "Input table must have either (time_start, time_end) for DL1_MONITORING "
                "or 'time' column for DL0_MONITORING"
            )

        # Create result table for raw monitoring data
        result = Table()

        # Convert 'time' to time_start and time_end (both equal for point measurements)
        result["time_start"] = table["time"]
        result["time_end"] = table["time"]

        # Preserve the original data column
        result[col_name] = table[col_name]

        return result


class PercentageDescriptor(BaseAggregator):
    """Computes percentage of non-zero values per chunk."""

    def _add_result_columns(self, data, masked_elements_of_sample, results_dict):
        value = np.array(data)

        # Apply mask if provided
        if masked_elements_of_sample is not None:
            masked_data = np.ma.array(value, mask=masked_elements_of_sample)
            valid_data = masked_data.compressed()
        else:
            valid_data = value.flatten()

        # Compute percentage of non-zero values
        if len(valid_data) > 0:
            percentage = (valid_data != 0).sum() / len(valid_data) * 100
        else:
            percentage = np.nan

        results_dict["percentage"].append(percentage)

    def _set_result_units(self, table, unit):
        # Percentage is dimensionless
        pass


class RateCounter(BaseAggregator):
    """Computes measurement rate in Hz for each time-based chunk."""

    # Override chunking_type to be fixed to TimeChunking (non-configurable)
    chunking_type = ComponentName(
        BaseChunking,
        default_value="TimeChunking",
        help="Fixed to TimeChunking for rate calculations (non-configurable).",
    )  # No .tag(config=True) - makes it non-configurable

    def __init__(self, config=None, parent=None, **kwargs):
        """Initialize RateCounter with fixed TimeChunking configuration."""
        super().__init__(config=config, parent=parent, **kwargs)

        # Check and enforce default values for consistent rate calculations
        if self.chunking.allow_undersized_tables:
            self.log.warning(
                "RateCounter requires allow_undersized_tables=False for consistent rate calculations, "
                "but got %s. Resetting to False.",
                self.chunking.allow_undersized_tables,
            )
            self.chunking.allow_undersized_tables = False

        if self.chunking.last_chunk_policy != "overlap":
            self.log.warning(
                "RateCounter requires last_chunk_policy='overlap' for consistent rate calculations, "
                "but got '%s'. Resetting to 'overlap'.",
                self.chunking.last_chunk_policy,
            )
            self.chunking.last_chunk_policy = "overlap"

    def _add_result_columns(self, data, masked_elements_of_sample, results_dict):
        # Count measurements in this chunk
        measurement_count = len(data)

        # Use configured chunk duration for consistent rate calculation
        # Since we enforce TimeChunking, chunk_duration is guaranteed to exist
        duration = self.chunking.chunk_duration.to_value("s")

        # Compute rate in Hz (measurements/second)
        rate = measurement_count / duration

        results_dict["rate"].append(rate)

    def _set_result_units(self, table, unit):
        # Rate has units of Hz (1/s)
        import astropy.units as u

        if "rate" in table.colnames:
            table["rate"].unit = u.Hz


class EventRateCounter(RateCounter):
    """Computes event rate in Hz for each time-based chunk, with event type filtering."""

    # Override to require event_type as data source and add event type filtering
    event_type_filter = Enum(
        values=[event_type.name for event_type in EventType],
        default_value="SUBARRAY",
        help="Event type to count for rate calculation. Based on ctapipe EventType enum.",
    ).tag(config=True)

    def __call__(self, table, col_name="event_type", **kwargs):
        """
        Compute event rate filtering by event type.

        Parameters
        ----------
        table : astropy.table.Table
            Input table with 'event_type' column
        col_name : str
            Must be "event_type" for EventRateCounter
        **kwargs
            Additional arguments passed to parent

        Returns
        -------
        astropy.table.Table
            Rate table with event type filtering applied
        """
        if col_name != "event_type":
            self.log.warning(
                "EventRateCounter requires col_name='event_type', got '%s'. "
                "Automatically using 'event_type' for event type counting.",
                col_name,
            )
            col_name = "event_type"

        if "event_type" not in table.colnames:
            raise ValueError(
                "Input table must have 'event_type' column for EventRateCounter"
            )

        return super().__call__(table, col_name=col_name, **kwargs)

    def _add_result_columns(self, data, masked_elements_of_sample, results_dict):
        # Get EventType enum value directly by name
        target_event_type = EventType[self.event_type_filter]

        # Filter events by type - data contains event_type values
        matching_events = np.sum(data == target_event_type.value)

        # Use configured chunk duration for consistent rate calculation
        duration = self.chunking.chunk_duration.to_value("s")

        # Compute rate in Hz (matching events/second)
        rate = matching_events / duration

        results_dict["rate"].append(rate)
