"""Pulse width and glitch triggering for Oscura.

Provides pulse width triggering, glitch detection, and runt pulse
detection for signal integrity analysis.

Example:
    >>> from oscura.triggering.pulse import PulseWidthTrigger, find_glitches
    >>> # Find pulses between 100ns and 200ns
    >>> trigger = PulseWidthTrigger(level=1.5, min_width=100e-9, max_width=200e-9)
    >>> events = trigger.find_events(trace)
    >>> # Find glitches shorter than 50ns
    >>> glitches = find_glitches(trace, max_width=50e-9)
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np

from oscura.core.exceptions import AnalysisError
from oscura.core.types import DigitalTrace, WaveformTrace
from oscura.triggering.base import (
    Trigger,
    TriggerEvent,
    TriggerType,
    interpolate_crossing,
)


@dataclass
class PulseInfo:
    """Information about a detected pulse.

    Attributes:
        start_time: Start time of pulse in seconds.
        end_time: End time of pulse in seconds.
        width: Pulse width in seconds.
        polarity: "positive" or "negative".
        start_index: Sample index at pulse start.
        end_index: Sample index at pulse end.
        amplitude: Peak amplitude during pulse.
    """

    start_time: float
    end_time: float
    width: float
    polarity: Literal["positive", "negative"]
    start_index: int
    end_index: int
    amplitude: float


class PulseWidthTrigger(Trigger):
    """Pulse width trigger for detecting pulses in a width range.

    Triggers on pulses that fall within the specified width range.

    Attributes:
        level: Threshold level for pulse detection.
        polarity: Pulse polarity - "positive", "negative", or "either".
        min_width: Minimum pulse width (None for no minimum).
        max_width: Maximum pulse width (None for no maximum).
    """

    def __init__(
        self,
        level: float,
        polarity: Literal["positive", "negative", "either"] = "positive",
        min_width: float | None = None,
        max_width: float | None = None,
    ) -> None:
        """Initialize pulse width trigger.

        Args:
            level: Threshold level for pulse detection.
            polarity: Pulse polarity to detect.
            min_width: Minimum pulse width in seconds.
            max_width: Maximum pulse width in seconds.

        Raises:
            AnalysisError: If min_width is greater than max_width.
        """
        self.level = level
        self.polarity = polarity
        self.min_width = min_width
        self.max_width = max_width

        if min_width is not None and max_width is not None and min_width > max_width:
            raise AnalysisError("min_width cannot be greater than max_width")

    def find_events(
        self,
        trace: WaveformTrace | DigitalTrace,
    ) -> list[TriggerEvent]:
        """Find pulses matching the width criteria.

        Args:
            trace: Input trace.

        Returns:
            List of trigger events for matching pulses.
        """
        pulses = self._find_all_pulses(trace)

        # Filter by width
        events: list[TriggerEvent] = []
        for pulse in pulses:
            if self.min_width is not None and pulse.width < self.min_width:
                continue
            if self.max_width is not None and pulse.width > self.max_width:
                continue

            events.append(
                TriggerEvent(
                    timestamp=pulse.start_time,
                    sample_index=pulse.start_index,
                    event_type=TriggerType.PULSE_WIDTH,
                    level=pulse.amplitude,
                    duration=pulse.width,
                    data={
                        "polarity": pulse.polarity,
                        "end_time": pulse.end_time,
                        "end_index": pulse.end_index,
                    },
                )
            )

        return events

    def _find_all_pulses(
        self,
        trace: WaveformTrace | DigitalTrace,
    ) -> list[PulseInfo]:
        """Find all pulses in the trace."""
        if isinstance(trace, DigitalTrace):
            data = trace.data.astype(np.float64)
            level = 0.5
        else:
            data = trace.data
            level = self.level

        sample_period = trace.metadata.time_base
        pulses: list[PulseInfo] = []

        # Find all threshold crossings
        above = data >= level
        below = data < level

        # Rising edges: transition from below to above
        rising = np.where(below[:-1] & above[1:])[0]
        # Falling edges: transition from above to below
        falling = np.where(above[:-1] & below[1:])[0]

        if self.polarity in ("positive", "either"):
            # Positive pulses: rising -> falling
            for r_idx in rising:
                # Find next falling edge
                next_falling = falling[falling > r_idx]
                if len(next_falling) == 0:
                    continue
                f_idx = next_falling[0]

                start_time = interpolate_crossing(data, r_idx, level, sample_period, True)
                end_time = interpolate_crossing(data, f_idx, level, sample_period, False)
                width = end_time - start_time

                # Get peak amplitude
                pulse_data = data[r_idx : f_idx + 1]
                amplitude = float(np.max(pulse_data)) if len(pulse_data) > 0 else level

                pulses.append(
                    PulseInfo(
                        start_time=start_time,
                        end_time=end_time,
                        width=width,
                        polarity="positive",
                        start_index=int(r_idx),
                        end_index=int(f_idx),
                        amplitude=amplitude,
                    )
                )

        if self.polarity in ("negative", "either"):
            # Negative pulses: falling -> rising
            for f_idx in falling:
                # Find next rising edge
                next_rising = rising[rising > f_idx]
                if len(next_rising) == 0:
                    continue
                r_idx = next_rising[0]

                start_time = interpolate_crossing(data, f_idx, level, sample_period, False)
                end_time = interpolate_crossing(data, r_idx, level, sample_period, True)
                width = end_time - start_time

                # Get peak (minimum) amplitude
                pulse_data = data[f_idx : r_idx + 1]
                amplitude = float(np.min(pulse_data)) if len(pulse_data) > 0 else level

                pulses.append(
                    PulseInfo(
                        start_time=start_time,
                        end_time=end_time,
                        width=width,
                        polarity="negative",
                        start_index=int(f_idx),
                        end_index=int(r_idx),
                        amplitude=amplitude,
                    )
                )

        # Sort by start time
        pulses.sort(key=lambda p: p.start_time)
        return pulses


class GlitchTrigger(Trigger):
    """Glitch trigger for detecting narrow pulses.

    Glitches are pulses shorter than a maximum width threshold.

    Attributes:
        level: Threshold level.
        max_width: Maximum pulse width to be considered a glitch.
        polarity: Glitch polarity - "positive", "negative", or "either".
    """

    def __init__(
        self,
        level: float,
        max_width: float = 100e-9,
        polarity: Literal["positive", "negative", "either"] = "either",
    ) -> None:
        """Initialize glitch trigger.

        Args:
            level: Threshold level.
            max_width: Maximum pulse width to trigger (in seconds).
            polarity: Glitch polarity to detect.
        """
        self.level = level
        self.max_width = max_width
        self.polarity = polarity

    def find_events(
        self,
        trace: WaveformTrace | DigitalTrace,
    ) -> list[TriggerEvent]:
        """Find all glitches in the trace.

        Args:
            trace: Input trace.

        Returns:
            List of trigger events for each glitch.
        """
        pulse_trigger = PulseWidthTrigger(
            level=self.level,
            polarity=self.polarity,
            min_width=None,
            max_width=self.max_width,
        )

        events = pulse_trigger.find_events(trace)

        # Reclassify as glitch events
        for event in events:
            event.event_type = TriggerType.GLITCH

        return events


class RuntTrigger(Trigger):
    """Runt pulse trigger for detecting incomplete transitions.

    Runt pulses cross one threshold but not the other, indicating
    incomplete signal transitions.

    Attributes:
        low_threshold: Lower threshold level.
        high_threshold: Upper threshold level.
        polarity: Runt polarity - "positive", "negative", or "either".
    """

    def __init__(
        self,
        low_threshold: float,
        high_threshold: float,
        polarity: Literal["positive", "negative", "either"] = "either",
    ) -> None:
        """Initialize runt trigger.

        Args:
            low_threshold: Lower threshold (e.g., logic low).
            high_threshold: Upper threshold (e.g., logic high).
            polarity: "positive" for rising runts, "negative" for falling.

        Raises:
            AnalysisError: If low_threshold is not less than high_threshold.
        """
        if low_threshold >= high_threshold:
            raise AnalysisError("low_threshold must be less than high_threshold")

        self.low_threshold = low_threshold
        self.high_threshold = high_threshold
        self.polarity = polarity

    def find_events(
        self,
        trace: WaveformTrace | DigitalTrace,
    ) -> list[TriggerEvent]:
        """Find all runt pulses in the trace.

        Args:
            trace: Input trace.

        Returns:
            List of trigger events for each runt pulse.
        """
        if isinstance(trace, DigitalTrace):
            # Digital traces don't have runts
            return []

        data = trace.data
        sample_period = trace.metadata.time_base
        events: list[TriggerEvent] = []

        # Track signal zones
        # Zone 0: below low_threshold
        # Zone 1: between thresholds
        # Zone 2: above high_threshold
        def get_zone(value: float) -> int:
            if value < self.low_threshold:
                return 0
            elif value > self.high_threshold:
                return 2
            else:
                return 1

        zones = np.array([get_zone(v) for v in data])

        # Find runt pulses: transitions that enter zone 1 but don't reach the other side
        i = 0
        while i < len(zones) - 1:
            curr_zone = zones[i]

            if curr_zone == 0:
                # Starting low - look for positive runt
                if self.polarity in ("positive", "either"):
                    # Find transition to zone 1
                    if zones[i + 1] == 1:
                        start_idx = i
                        # Track through zone 1
                        j = i + 1
                        while j < len(zones) and zones[j] == 1:
                            j += 1
                        if j < len(zones) and zones[j] == 0:
                            # Returned to low without reaching high - RUNT
                            peak = float(np.max(data[start_idx : j + 1]))
                            events.append(
                                TriggerEvent(
                                    timestamp=start_idx * sample_period,
                                    sample_index=start_idx,
                                    event_type=TriggerType.RUNT,
                                    level=peak,
                                    duration=(j - start_idx) * sample_period,
                                    data={
                                        "polarity": "positive",
                                        "expected_high": self.high_threshold,
                                        "actual_peak": peak,
                                    },
                                )
                            )
                        i = j
                        continue

            elif curr_zone == 2:
                # Starting high - look for negative runt
                if self.polarity in ("negative", "either") and zones[i + 1] == 1:
                    start_idx = i
                    j = i + 1
                    while j < len(zones) and zones[j] == 1:
                        j += 1
                    if j < len(zones) and zones[j] == 2:
                        # Returned to high without reaching low - RUNT
                        trough = float(np.min(data[start_idx : j + 1]))
                        events.append(
                            TriggerEvent(
                                timestamp=start_idx * sample_period,
                                sample_index=start_idx,
                                event_type=TriggerType.RUNT,
                                level=trough,
                                duration=(j - start_idx) * sample_period,
                                data={
                                    "polarity": "negative",
                                    "expected_low": self.low_threshold,
                                    "actual_trough": trough,
                                },
                            )
                        )
                    i = j
                    continue

            i += 1

        return events


def find_pulses(
    trace: WaveformTrace,
    *,
    level: float | None = None,
    polarity: Literal["positive", "negative", "either"] = "positive",
    min_width: float | None = None,
    max_width: float | None = None,
) -> list[TriggerEvent]:
    """Find pulses matching width criteria.

    Args:
        trace: Input waveform trace.
        level: Threshold level. If None, uses 50% of amplitude.
        polarity: Pulse polarity to find.
        min_width: Minimum pulse width in seconds.
        max_width: Maximum pulse width in seconds.

    Returns:
        List of trigger events for matching pulses.

    Example:
        >>> # Find all positive pulses between 1us and 10us
        >>> pulses = find_pulses(trace, min_width=1e-6, max_width=10e-6)
    """
    if level is None:
        level = (np.min(trace.data) + np.max(trace.data)) / 2

    trigger = PulseWidthTrigger(
        level=level,
        polarity=polarity,
        min_width=min_width,
        max_width=max_width,
    )
    return trigger.find_events(trace)


def find_glitches(
    trace: WaveformTrace,
    max_width: float = 100e-9,
    *,
    level: float | None = None,
    polarity: Literal["positive", "negative", "either"] = "either",
) -> list[TriggerEvent]:
    """Find glitches (narrow pulses) in a trace.

    Args:
        trace: Input waveform trace.
        max_width: Maximum width to be considered a glitch (default 100ns).
        level: Threshold level. If None, uses 50% of amplitude.
        polarity: Glitch polarity to find.

    Returns:
        List of trigger events for each glitch.

    Example:
        >>> # Find all glitches shorter than 50ns
        >>> glitches = find_glitches(trace, max_width=50e-9)
        >>> print(f"Found {len(glitches)} glitches")
    """
    if level is None:
        level = (np.min(trace.data) + np.max(trace.data)) / 2

    trigger = GlitchTrigger(
        level=level,
        max_width=max_width,
        polarity=polarity,
    )
    return trigger.find_events(trace)


def find_runt_pulses(
    trace: WaveformTrace,
    low_threshold: float | None = None,
    high_threshold: float | None = None,
    *,
    polarity: Literal["positive", "negative", "either"] = "either",
) -> list[TriggerEvent]:
    """Find runt pulses (incomplete transitions) in a trace.

    Args:
        trace: Input waveform trace.
        low_threshold: Lower threshold. If None, uses 20% of amplitude.
        high_threshold: Upper threshold. If None, uses 80% of amplitude.
        polarity: Runt polarity to find.

    Returns:
        List of trigger events for each runt pulse.

    Example:
        >>> # Find runts using standard 20%/80% thresholds
        >>> runts = find_runt_pulses(trace)
        >>> for runt in runts:
        ...     print(f"Runt at {runt.timestamp*1e6:.2f} us")
    """
    if low_threshold is None:
        amplitude = np.max(trace.data) - np.min(trace.data)
        low_threshold = np.min(trace.data) + 0.2 * amplitude

    if high_threshold is None:
        amplitude = np.max(trace.data) - np.min(trace.data)
        high_threshold = np.min(trace.data) + 0.8 * amplitude

    trigger = RuntTrigger(
        low_threshold=low_threshold,
        high_threshold=high_threshold,
        polarity=polarity,
    )
    return trigger.find_events(trace)


def pulse_statistics(
    trace: WaveformTrace,
    *,
    level: float | None = None,
    polarity: Literal["positive", "negative"] = "positive",
) -> dict[str, float]:
    """Calculate pulse width statistics.

    Args:
        trace: Input waveform trace.
        level: Threshold level.
        polarity: Pulse polarity to analyze.

    Returns:
        Dictionary with pulse statistics:
        - count: Number of pulses
        - min_width: Minimum pulse width
        - max_width: Maximum pulse width
        - mean_width: Mean pulse width
        - std_width: Standard deviation of pulse widths

    Example:
        >>> stats = pulse_statistics(trace)
        >>> print(f"Mean pulse width: {stats['mean_width']*1e6:.2f} us")
    """
    if level is None:
        level = (np.min(trace.data) + np.max(trace.data)) / 2

    trigger = PulseWidthTrigger(level=level, polarity=polarity)
    events = trigger.find_events(trace)

    if len(events) == 0:
        return {
            "count": 0,
            "min_width": np.nan,
            "max_width": np.nan,
            "mean_width": np.nan,
            "std_width": np.nan,
        }

    widths = np.array([e.duration for e in events if e.duration is not None])

    return {
        "count": len(widths),
        "min_width": float(np.min(widths)),
        "max_width": float(np.max(widths)),
        "mean_width": float(np.mean(widths)),
        "std_width": float(np.std(widths)),
    }


__all__ = [
    "GlitchTrigger",
    "PulseInfo",
    "PulseWidthTrigger",
    "RuntTrigger",
    "find_glitches",
    "find_pulses",
    "find_runt_pulses",
    "pulse_statistics",
]
