from __future__ import annotations

import textwrap
from abc import ABC, abstractmethod
from collections import ChainMap
from typing import TYPE_CHECKING, Iterable, Iterator, Optional

import numpy as np
import pint
from attrs import Factory, field, frozen, validators
from frozendict import frozendict
from intervaltree import Interval, IntervalTree
from plotly import graph_objects as go
from typing_extensions import assert_never, override

from xq_pulse.pulse import expression
from xq_pulse.pulse.channel import Channel
from xq_pulse.pulse.envelope import Envelope, SquareEnvelope
from xq_pulse.pulse.expression import (
    Expression,
    Literal,
    Parameter,
    Parameterized,
    bind,
    eval_expression,
    simplify,
)
from xq_pulse.pulse.validators import ExpressionOrQuantityValidator
from xq_pulse.util import unit, units

if TYPE_CHECKING:
    from xq_pulse.pulse.parameter import Binding


def normalize_phase(phase: pint.Quantity) -> pint.Quantity:
    """Normalize a literal phase into [0, 2π) with canonical radian units."""
    phase_rad = phase.to(unit.rad)
    return phase_rad % (2 * np.pi * unit.rad)


def evaluate_phase(phase: Expression) -> pint.Quantity:
    """Evaluate and normalize a phase expression for execution/visualization."""
    return normalize_phase(eval_expression(phase))


class ActivePulseIntervals:
    """Split unrolled pulses into contiguous regions with stable active-pulse sets."""

    _intervals: IntervalTree
    _active_intervals: tuple[ActiveInterval, ...]
    _bounds_by_pulse_id: dict[int, tuple[tuple[pint.Quantity, pint.Quantity], ...]]

    def __init__(self, unrolled_pulse: UnrolledPulse):
        """
        Break the sequence up at points where the active pulses change.

        """
        self._bounds_by_pulse_id = self._collect_bounds_by_pulse(unrolled_pulse)
        intervals = unrolled_pulse.intervals.copy()
        intervals.split_overlaps()
        intervals.merge_overlaps(
            data_reducer=lambda old, new: (*old, new),
            data_initializer=tuple(),
        )
        self._intervals = intervals
        self._active_intervals = tuple(
            self._build_active_interval(interval) for interval in sorted(intervals, key=lambda x: x.begin)
        )

    @staticmethod
    def _collect_bounds_by_pulse(
        unrolled_pulse: UnrolledPulse,
    ) -> dict[int, tuple[tuple[pint.Quantity, pint.Quantity], ...]]:
        """Index original unrolled interval bounds by pulse identity."""
        bounds_by_pulse_id: dict[int, list[tuple[pint.Quantity, pint.Quantity]]] = {}
        for interval in unrolled_pulse.intervals:
            pulse = interval.data
            pulse_id = id(pulse)
            bounds_by_pulse_id.setdefault(pulse_id, []).append((interval.begin, interval.end))
        return {pulse_id: tuple(bounds) for pulse_id, bounds in bounds_by_pulse_id.items()}

    def _resolve_pulse_bounds(
        self, pulse: LeafPulse, interval_begin: pint.Quantity, interval_end: pint.Quantity
    ) -> tuple[pint.Quantity, pint.Quantity]:
        """Resolve the source pulse bounds that contain the given split interval."""
        candidates = self._bounds_by_pulse_id.get(id(pulse))
        assert candidates, f"Could not resolve pulse bounds for pulse {pulse!r}"
        matching = [bounds for bounds in candidates if bounds[0] <= interval_begin and interval_end <= bounds[1]]
        assert len(matching) == 1, (
            f"Expected exactly one matching source interval for pulse {pulse!r} in [{interval_begin}, {interval_end}), "
            f"got {len(matching)} candidates."
        )
        return matching[0]

    def _build_active_interval(self, interval: Interval) -> ActiveInterval:
        """Build one enriched active interval with per-pulse slice metadata."""
        slices: list[PulseIntervalSlice] = []
        for pulse in interval.data:
            if isinstance(pulse, DelayPulse):
                continue
            pulse_begin, pulse_end = self._resolve_pulse_bounds(pulse, interval.begin, interval.end)
            slices.append(
                PulseIntervalSlice(
                    pulse=pulse,
                    pulse_start=pulse_begin,
                    pulse_end=pulse_end,
                    interval_start=interval.begin,
                    interval_end=interval.end,
                )
            )

        return ActiveInterval(
            start=interval.begin,
            end=interval.end,
            slices=tuple(slices),
        )

    def __getitem__(self, key: pint.Quantity) -> set[LeafPulse]:
        """
        Get a tuple of the active pulses at a given time.
        """
        assert isinstance(key, pint.Quantity), f"Key {key} is not a Quantity"
        assert key.is_compatible_with(unit.s), f"Key {key} is not a time unit but {key.u}"
        intervals_at_time: Iterator[Interval] = self._intervals[key]
        assert len(intervals_at_time) == 1, f"Expected 1 interval at time {key} but got {len(intervals_at_time)}"
        single_interval = next(iter(intervals_at_time))
        assert isinstance(single_interval, Interval), f"Expected an Interval but got {type(single_interval)}"
        assert isinstance(single_interval.data, tuple), (
            f"Expected a tuple of pulses but got {type(single_interval.data)}"
        )

        return set(iter(single_interval.data))

    def __len__(self) -> int:
        """Return the number of split intervals."""
        return len(self._intervals)

    def __iter__(self) -> Iterator[tuple[pint.Quantity, set[LeafPulse]]]:
        """Iterate split intervals as (duration, active non-delay pulses)."""
        for interval in self._active_intervals:
            yield (
                interval.duration,
                set(interval.pulses),
            )

    def iter_active_intervals(self) -> Iterator[ActiveInterval]:
        """Iterate split intervals with per-pulse slice metadata."""
        return iter(self._active_intervals)

    def __repr__(self) -> str:
        body = "\n".join(f"  {duration}: {pulses}" for duration, pulses in self)
        return f"ActivePulseIntervals(\n{body}\n)"


@frozen
class PulseIntervalSlice:
    """
    Metadata describing which part of a pulse is active within a split interval.
    """

    pulse: LeafPulse
    pulse_start: pint.Quantity
    pulse_end: pint.Quantity
    interval_start: pint.Quantity
    interval_end: pint.Quantity

    def __attrs_post_init__(self) -> None:
        assert self.pulse_start <= self.interval_start <= self.interval_end <= self.pulse_end, (
            f"Interval [{self.interval_start}, {self.interval_end}) must be inside pulse "
            f"[{self.pulse_start}, {self.pulse_end})."
        )

    @property
    def pulse_duration(self) -> pint.Quantity:
        """Return full duration of the source pulse for this slice."""
        return self.pulse_end - self.pulse_start

    @property
    def interval_duration(self) -> pint.Quantity:
        """Return duration of this split interval."""
        return self.interval_end - self.interval_start

    @property
    def local_start(self) -> pint.Quantity:
        """Return pulse-local slice start time."""
        return self.interval_start - self.pulse_start

    @property
    def local_end(self) -> pint.Quantity:
        """Return pulse-local slice end time."""
        return self.interval_end - self.pulse_start

    @property
    def tau_start(self) -> float:
        """Return normalized pulse-local slice start in [0, 1]."""
        tau = self.local_start.to(unit.s).magnitude / self.pulse_duration.to(unit.s).magnitude
        return float(np.clip(tau, 0.0, 1.0))

    @property
    def tau_end(self) -> float:
        """Return normalized pulse-local slice end in [0, 1]."""
        tau = self.local_end.to(unit.s).magnitude / self.pulse_duration.to(unit.s).magnitude
        return float(np.clip(tau, 0.0, 1.0))


@frozen
class ActiveInterval:
    """
    Split interval with enriched pulse-slice metadata.
    """

    start: pint.Quantity
    end: pint.Quantity
    slices: tuple[PulseIntervalSlice, ...]

    def __attrs_post_init__(self) -> None:
        assert self.start <= self.end, f"Active interval start ({self.start}) must be <= end ({self.end})"

    @property
    def duration(self) -> pint.Quantity:
        """Return interval duration."""
        return self.end - self.start

    @property
    def pulses(self) -> tuple[LeafPulse, ...]:
        """Return pulses active in this interval, preserving slice order."""
        return tuple(pulse_slice.pulse for pulse_slice in self.slices)


@frozen
class UnrolledPulse:
    """
    Aggregates the result of unrolling a pulse tree.
    """

    intervals: IntervalTree
    source_by_unrolled: frozendict[int, Pulse]
    unrolled_by_source: frozendict[int, tuple[LeafPulse]]

    @classmethod
    def from_pulse(cls, pulse: LeafPulse, source: Pulse) -> UnrolledPulse:
        """Create an UnrolledPulse from a single pulse and its source."""
        duration: pint.Quantity = units.round_time(eval_expression(pulse.duration))
        assert isinstance(duration, pint.Quantity), (
            f"Evaluating the pulse duration did not yield a Quantity but {type(duration).__name__}: {duration}"
        )
        assert duration.is_compatible_with(unit.s), f"Pulse duration is not a time unit but {duration.u}"
        if isinstance(pulse, DrivePulse):
            pulse = DrivePulse(
                duration=pulse.duration,
                amplitude=pulse.amplitude,
                frequency=pulse.frequency,
                phase=evaluate_phase(pulse.phase),
                envelope=pulse.envelope,
            )
        return cls(
            intervals=IntervalTree([Interval(0 * unit.s, duration, pulse)]),
            source_by_unrolled={id(pulse): source},
            unrolled_by_source={id(source): (pulse,)},
        )

    @classmethod
    def empty(cls) -> UnrolledPulse:
        return cls(
            intervals=IntervalTree(),
            source_by_unrolled=frozendict(),
            unrolled_by_source=frozendict(),
        )

    @staticmethod
    def intervals_duration(intervals: IntervalTree) -> pint.Quantity:
        duration = intervals.end()
        if isinstance(duration, int) and duration == 0:
            return 0 * unit.s
        return units.round_time(duration)

    @classmethod
    def combine(cls, pulses: Iterable[UnrolledPulse], *, sequential: bool) -> UnrolledPulse:
        intervals: list[Interval] = []
        source_by_unrolled: dict[int, Pulse] = {}
        unrolled_by_source: dict[int, list[LeafPulse]] = {}
        offset = 0 * unit.s
        num_pulses = 0

        for pulse in pulses:
            num_pulses += 1
            pulse_offset = offset if sequential else 0 * unit.s
            for interval in pulse.intervals:
                intervals.append(
                    Interval(
                        begin=units.round_time(interval.begin + pulse_offset),
                        end=units.round_time(interval.end + pulse_offset),
                        data=interval.data,
                    )
                )
            source_by_unrolled.update(pulse.source_by_unrolled)
            for source_id, source_pulses in pulse.unrolled_by_source.items():
                unrolled_by_source.setdefault(source_id, []).extend(source_pulses)

            if sequential:
                offset = units.round_time(offset + cls.intervals_duration(pulse.intervals))

        if num_pulses == 0:
            return cls.empty()

        return cls(
            intervals=IntervalTree(intervals),
            source_by_unrolled=frozendict(source_by_unrolled),
            unrolled_by_source=frozendict(
                {source_id: tuple(source_pulses) for source_id, source_pulses in unrolled_by_source.items()}
            ),
        )

    @classmethod
    def concatenate(cls, pulses: Iterable[UnrolledPulse]) -> UnrolledPulse:
        """Append unrolled pulses in sequence."""
        return cls.combine(pulses, sequential=True)

    @classmethod
    def overlay(cls, pulses: Iterable[UnrolledPulse]) -> UnrolledPulse:
        """Overlay unrolled pulses in parallel from a shared time origin."""
        return cls.combine(pulses, sequential=False)

    @property
    def pulse_starts(self) -> list[pint.Quantity]:
        return sorted(set([interval.begin for interval in self.intervals]))

    @property
    def pulse_ends(self) -> list[pint.Quantity]:
        return [interval.end for interval in self.intervals]

    @property
    def unrolled_pulses(self) -> list[LeafPulse]:
        return [interval.data for interval in self.intervals]

    @property
    def source_pulses(self) -> list[Pulse]:
        return list(self.source_by_unrolled.values())

    def __getitem__(self, key: pint.Quantity) -> set[LeafPulse]:
        return set([interval.data for interval in self.intervals[key]])

    def append(self, other: UnrolledPulse) -> UnrolledPulse:
        """Append two UnrolledPulses where the other occurs after this one (self)."""
        return UnrolledPulse.concatenate((self, other))

    def merge(self, other: UnrolledPulse) -> UnrolledPulse:
        """Merge two UnrolledPulses which start at the same time."""
        return UnrolledPulse.overlay((self, other))


@frozen
class AnyPulse(ABC, Parameterized):
    """
    Abstract base class for all pulses.
    Only use this for inheritance.
    For type annotations, use Pulse, which is a union type and supports compile time exhaustiveness checking.
    """

    duration: Expression = field(
        repr=str,
        validator=ExpressionOrQuantityValidator(
            unit.seconds,
            min_value=1 * unit.ns,
            max_value=1000 * unit.us,
            min_inclusive=True,
            max_inclusive=True,
        ),
    )

    @property
    def channels(self) -> frozenset[Channel]:
        """
        Recursively collect all hardware channels used by this pulse and its descendants.
        Default for leaves without explicit channel assignment is empty.
        """
        return frozenset()

    @abstractmethod
    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> Pulse:
        """
        Recursively bind Parameters in the pulse tree to literal values.
        Returns a new pulse with all bound Parameters replaced by their respective values.

        Optionally, a source_map can be provided in which id(bound pulse) -> source pulse is is stored.
        """
        ...

    @abstractmethod
    def unroll(self) -> UnrolledPulse:
        """
        Recursively unroll the pulse tree into individual leaf pulses.
        Returns an UnrolledPulse object containing all unrolling information.
        """
        ...

    @abstractmethod
    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        """
        Return a new pulse with channel mappings applied according to mapping from abstract leaf id to Channel.
        Implemented by each concrete pulse type.
        """
        ...

    def plot(self, show=True) -> go.Figure | None:
        container_padding = 0.01
        label_height = 0.015
        height_px = 400
        fig = go.Figure()

        def evaluate_quantity(expression_or_quantity: Expression | pint.Quantity) -> pint.Quantity:
            value = (
                expression_or_quantity
                if isinstance(expression_or_quantity, pint.Quantity)
                else eval_expression(expression_or_quantity)
            )
            assert isinstance(value, pint.Quantity), f"Expected Quantity but got {type(value).__name__}: {value}"
            return value

        def convert_quantity(
            expression_or_quantity: Expression | pint.Quantity, display_unit: pint.Unit, *, compact: bool = True
        ) -> pint.Quantity:
            value = evaluate_quantity(expression_or_quantity)
            converted = value.to(display_unit)
            if compact:
                converted = converted.to_compact()
            return converted

        def format_quantity(
            expression_or_quantity: Expression | pint.Quantity, display_unit: pint.Unit, *, compact: bool = True
        ) -> str:
            return f"{convert_quantity(expression_or_quantity, display_unit, compact=compact):.3g~#C}"

        def duration_us(expression_or_quantity: Expression | pint.Quantity) -> float:
            return convert_quantity(expression_or_quantity, unit.us, compact=False).magnitude

        def height(pulse):
            match pulse:
                case DelayPulse():
                    return 0.1
                case ChannelMappedPulse():
                    return height(pulse.pulse) + label_height + container_padding
                case LaserPulse() | DrivePulse() | AcquisitionPulse():
                    return 0.1
                case SequencePulse():
                    return max((height(sp) for sp in pulse.pulses), default=0.0) + label_height + container_padding
                case ParallelPulse():
                    return sum(height(sp) for sp in pulse.pulses) + label_height + container_padding * len(pulse.pulses)
                case ForLoopPulse():
                    return (
                        max((height(iterp) for iterp in pulse.iterations()), default=0.0)
                        + label_height
                        + container_padding
                    )
                case _:
                    raise AssertionError(f"Got unknown pulse type {type(pulse)}")

        def add_outline(x0, x1, y0, y1, line):
            fig.add_shape(
                type="rect",
                x0=x0,
                x1=x1,
                y0=y0,
                y1=y1,
                fillcolor="rgba(0,0,0,0)",
                line=dict(color=line, dash="dot", width=1),
            )

        def add_label(x, y, text):
            fig.add_annotation(
                x=x,
                y=y,
                xanchor="left",
                yanchor="top",
                text=text.replace("\n", "<br>"),
                showarrow=False,
            )

        def add_bar(x0, width, y0, h, fill, line, text, insidetextanchor: str = "middle"):
            fig.add_trace(
                go.Bar(
                    orientation="h",
                    x=[width],
                    base=[x0],
                    y=[y0 + h / 2.0],
                    width=[h],
                    marker=dict(color=fill, line=dict(color=line, width=1)),
                    text=[text],
                    textposition="inside",
                    insidetextanchor=insidetextanchor,
                    textfont=dict(color="black"),
                    hoverinfo="skip",
                    cliponaxis=True,
                )
            )

        def add_envelope_trace(pulse: Pulse, x0: float, y0: float, width: float, height: float, color: str):
            envelope = getattr(pulse, "envelope", None)
            if not isinstance(envelope, Envelope):
                return
            tau, values = envelope.sample(num_samples=121)
            if np.allclose(values, values[0]):
                return
            abs_max = np.max(np.abs(values))
            if abs_max == 0:
                return
            y_center = y0 + height / 2.0
            y_points = y_center + (values / abs_max) * (height / 2.0)
            x_points = x0 + tau * width
            fig.add_trace(
                go.Scatter(
                    x=x_points,
                    y=y_points,
                    mode="lines",
                    line=dict(color=color, width=2),
                    hoverinfo="skip",
                    showlegend=False,
                )
            )

        def draw(pulse: Pulse, x: float, y: float):
            match pulse:
                case DelayPulse():
                    w = duration_us(pulse.duration)
                    h = height(pulse)
                    add_bar(
                        x,
                        w,
                        y,
                        h,
                        "lightgray",
                        "gray",
                        f"Delay ({format_quantity(pulse.duration, unit.s)})",
                    )
                case ChannelMappedPulse():
                    # Draw wrapped pulse inside a dotted container with a label
                    h = height(pulse)
                    total = convert_quantity(pulse.duration, unit.s)
                    total_us = duration_us(total)
                    # Draw child
                    draw(pulse.pulse, x, y)
                    # Outline and in-box label
                    add_outline(x, x + total_us, y, y + h, "gray")
                    label_h = label_height
                    add_bar(
                        x,
                        total_us,
                        y + h - label_h,
                        label_h,
                        "rgba(0,0,0,0)",
                        "rgba(0,0,0,0)",
                        f"Channel: {pulse.channel.name}",
                        insidetextanchor="start",
                    )
                    return
                case LaserPulse():
                    w = duration_us(pulse.duration)
                    h = height(pulse)
                    add_bar(
                        x,
                        w,
                        y,
                        h,
                        "lightgreen",
                        "green",
                        f"Laser ({format_quantity(pulse.duration, unit.s)})<br>"
                        f"λ = {format_quantity(pulse.wavelength, unit.meter)}<br>"
                        f"power = {format_quantity(pulse.power, unit.watt)}",
                    )
                    add_envelope_trace(pulse, x, y, w, h, "green")
                case DrivePulse():
                    w = duration_us(pulse.duration)
                    h = height(pulse)
                    add_bar(
                        x,
                        w,
                        y,
                        h,
                        "pink",
                        "red",
                        f"Drive ({format_quantity(pulse.duration, unit.s)})<br>"
                        f"freq = {format_quantity(pulse.frequency, unit.Hz)}<br>"
                        f"amp = {format_quantity(pulse.amplitude, unit.T)}<br>"
                        f"phase = {format_quantity(evaluate_phase(pulse.phase), unit.rad, compact=False)}",
                    )
                    add_envelope_trace(pulse, x, y, w, h, "red")
                case AcquisitionPulse():
                    w = duration_us(pulse.duration)
                    h = height(pulse)
                    add_bar(
                        x,
                        w,
                        y,
                        h,
                        "lightblue",
                        "blue",
                        f"Acq ({format_quantity(pulse.duration, unit.s)})<br>"
                        f"target = {pulse.target}<br>"
                        f"bin = {format_quantity(pulse.bin, unit.dimensionless, compact=False)}",
                    )
                case SequencePulse():
                    h = height(pulse)
                    total = convert_quantity(pulse.duration, unit.s)
                    total_us = duration_us(total)
                    # Draw children
                    t = x
                    for sub in pulse.pulses:
                        sub_w = duration_us(sub.duration)
                        draw(sub, t, y)
                        t += sub_w
                    # Outline and in-box label
                    add_outline(x, x + total_us, y, y + h, "gray")
                    label_h = label_height
                    add_bar(
                        x,
                        total_us,
                        y + h - label_h,
                        label_h,
                        "rgba(0,0,0,0)",
                        "rgba(0,0,0,0)",
                        f"Sequence ({format_quantity(total, unit.s)})<br>",
                        insidetextanchor="start",
                    )
                case ParallelPulse():
                    h = height(pulse)
                    total = convert_quantity(pulse.duration, unit.s)
                    total_us = duration_us(total)
                    add_outline(x, x + total_us, y, y + h, "gray")
                    label_h = label_height
                    add_bar(
                        x,
                        total_us,
                        y + h - label_h,
                        label_h,
                        "rgba(0,0,0,0)",
                        "rgba(0,0,0,0)",
                        f"Parallel ({format_quantity(total, unit.s)})",
                        insidetextanchor="start",
                    )
                    yy = y + container_padding
                    for sub in pulse.pulses:
                        draw(sub, x, yy)
                        yy += height(sub) + container_padding
                case ForLoopPulse():
                    h = height(pulse)
                    total = convert_quantity(pulse.duration, unit.s)
                    total_us = duration_us(total)
                    add_outline(x, x + total_us, y, y + h, "gray")
                    label_h = label_height
                    add_bar(
                        x,
                        total_us,
                        y + h - label_h,
                        label_h,
                        "rgba(0,0,0,0)",
                        "rgba(0,0,0,0)",
                        f"ForLoop ({format_quantity(total, unit.s)})",
                        insidetextanchor="start",
                    )
                    t = x
                    for iteration in pulse.iterations():
                        w = duration_us(iteration.duration)
                        draw(iteration, t, y)
                        t += w
                case _:
                    assert_never(pulse)

        draw(self, 0.0, 0.0)

        total_us = duration_us(self.duration)
        total_h = height(self)

        fig.update_xaxes(
            title_text="Time [µs]",
            range=[0, total_us],
            showgrid=True,
            zeroline=False,
            showticklabels=True,
            tick0=0,
            dtick=1,
        )
        fig.update_yaxes(
            type="linear",
            range=[0, total_h + 0.05],
            showgrid=False,
            zeroline=True,
            showticklabels=False,
            fixedrange=True,
        )
        fig.update_layout(
            height=height_px,
            margin=dict(l=30, r=10, t=10, b=30),
            template="plotly_white",
            dragmode="pan",
            barmode="overlay",
            transition=dict(duration=0),
            uniformtext=dict(mode="hide", minsize=9),
            showlegend=False,
        )
        if show:
            fig.show(config={"scrollZoom": True})
        else:
            return fig


@frozen
class DelayPulse(AnyPulse):
    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> DelayPulse:
        result = DelayPulse(duration=bind(self.duration, binding))
        if source_map is not None:
            source_map[id(result)] = self
        return result

    def unroll(self) -> UnrolledPulse:
        return UnrolledPulse.from_pulse(self, self)

    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        return self

    @property
    def parameters(self) -> frozenset[Parameter]:
        return frozenset(
            {
                *({self.duration} if isinstance(self.duration, Parameter) else {}),
            }
        )

    @override
    def simplify(self) -> LeafPulse:
        return DelayPulse(duration=simplify(self.duration))


@frozen
class LaserPulse(AnyPulse):
    wavelength: Expression = field(
        repr=str,
        validator=ExpressionOrQuantityValidator(
            unit.meter,
            min_value=400 * unit.nm,
            max_value=1000 * unit.nm,
            min_inclusive=False,
            max_inclusive=False,
        ),
    )
    power: Expression = field(
        repr=str,
        validator=ExpressionOrQuantityValidator(
            unit.watt,
            min_value=0 * unit.mW,
            max_value=500 * unit.mW,
            min_inclusive=False,
            max_inclusive=False,
        ),
    )

    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> LaserPulse:
        result = LaserPulse(
            duration=bind(self.duration, binding),
            wavelength=bind(self.wavelength, binding),
            power=bind(self.power, binding),
        )
        if source_map is not None:
            source_map[id(result)] = self
        return result

    def unroll(self) -> UnrolledPulse:
        return UnrolledPulse.from_pulse(self, self)

    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        ch = mapping.get(id(self))
        if ch is None:
            return self
        return ChannelMappedPulse(pulse=self, channel=ch)

    @property
    def parameters(self) -> frozenset[Parameter]:
        params: set[Parameter] = set()
        for expr in (self.duration, self.wavelength, self.power):
            if isinstance(expr, Parameter):
                params.add(expr)
            elif isinstance(expr, Parameterized):
                params.update(expr.parameters)
        return frozenset(params)

    @override
    def simplify(self) -> LeafPulse:
        return LaserPulse(
            duration=simplify(self.duration),
            wavelength=simplify(self.wavelength),
            power=simplify(self.power),
        )


@frozen
class DrivePulse(AnyPulse):
    amplitude: Expression = field(
        repr=str,
        validator=ExpressionOrQuantityValidator(
            unit.T,
            min_value=0 * unit.T,
            max_value=10 * unit.mT,
            min_inclusive=False,
            max_inclusive=False,
        ),
    )
    frequency: Expression = field(
        repr=str,
        validator=ExpressionOrQuantityValidator(
            unit.Hz,
            min_value=0 * unit.Hz,
            max_value=20 * unit.GHz,
            min_inclusive=False,
            max_inclusive=False,
        ),
    )
    phase: Expression = field(
        repr=str,
        validator=ExpressionOrQuantityValidator(
            unit.rad,
        ),
    )
    envelope: Envelope = field(
        factory=SquareEnvelope,
        validator=validators.instance_of(Envelope),
        repr=False,
    )

    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> DrivePulse:
        result = DrivePulse(
            duration=bind(self.duration, binding),
            amplitude=bind(self.amplitude, binding),
            frequency=bind(self.frequency, binding),
            phase=bind(self.phase, binding),
            envelope=self.envelope,
        )
        if source_map is not None:
            source_map[id(result)] = self
        return result

    def unroll(self) -> UnrolledPulse:
        return UnrolledPulse.from_pulse(self, self)

    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        ch = mapping.get(id(self))
        if ch is None:
            return self
        return ChannelMappedPulse(pulse=self, channel=ch)

    @property
    def parameters(self) -> frozenset[Parameter]:
        params: set[Parameter] = set()
        for expr in (self.duration, self.amplitude, self.frequency, self.phase):
            if isinstance(expr, Parameter):
                params.add(expr)
            elif isinstance(expr, Parameterized):
                params.update(expr.parameters)
        return frozenset(params)

    @override
    def simplify(self) -> LeafPulse:
        return DrivePulse(
            duration=simplify(self.duration),
            amplitude=simplify(self.amplitude),
            frequency=simplify(self.frequency),
            phase=simplify(self.phase),
            envelope=self.envelope,
        )


@frozen
class AcquisitionTarget:
    name: str
    bins: int


@frozen
class AcquisitionPulse(AnyPulse):
    target: AcquisitionTarget
    bin: Expression = field(
        repr=str,
        validator=ExpressionOrQuantityValidator(
            unit.dimensionless,
            min_value=0 * unit.dimensionless,
            min_inclusive=True,
        ),
    )

    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> AcquisitionPulse:
        result = AcquisitionPulse(
            duration=bind(self.duration, binding),
            target=self.target,
            bin=bind(self.bin, binding),
        )
        if source_map is not None:
            source_map[id(result)] = self
        return result

    def unroll(self) -> UnrolledPulse:
        return UnrolledPulse.from_pulse(self, self)

    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        ch = mapping.get(id(self))
        if ch is None:
            return self
        return ChannelMappedPulse(pulse=self, channel=ch)

    @property
    def parameters(self) -> frozenset[Parameter]:
        return frozenset(
            {
                *({self.duration} if isinstance(self.duration, Parameter) else {}),
                *({self.bin} if isinstance(self.bin, Parameter) else {}),
            }
        )

    @override
    def simplify(self) -> LeafPulse:
        return AcquisitionPulse(
            duration=simplify(self.duration),
            target=self.target,
            bin=simplify(self.bin),
        )


@frozen
class ChannelMappedPulse(AnyPulse):
    pulse: LeafPulse
    channel: Channel
    duration: Expression = field(
        repr=str,
        default=Factory(lambda self: self.pulse.duration, takes_self=True),
    )

    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> ChannelMappedPulse:
        result = ChannelMappedPulse(
            pulse=self.pulse.bind(binding),
            channel=self.channel,
        )
        if source_map is not None:
            source_map[id(result)] = self
        return result

    def unroll(self) -> UnrolledPulse:
        result = UnrolledPulse.from_pulse(self, self)
        return result

    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        base = self.pulse
        ch = mapping.get(id(base))
        if ch is None:
            # Drop old wrapper if any and return the base leaf
            return base
        return ChannelMappedPulse(pulse=base, channel=ch)

    @property
    def parameters(self) -> frozenset[Parameter]:
        return self.pulse.parameters

    @property
    def channels(self) -> frozenset[Channel]:
        # Include this pulse's channel and any channels from the inner pulse (if present)
        return frozenset({self.channel, *self.pulse.channels})

    @override
    def simplify(self) -> LeafPulse:
        return ChannelMappedPulse(
            pulse=simplify(self.pulse),
            channel=self.channel,
        )


@frozen
class SequencePulse(AnyPulse):
    pulses: tuple[Pulse, ...]
    duration: Expression = field(
        repr=str,
        default=Factory(
            lambda self: (
                0 * unit.s if len(self.pulses) == 0 else expression.sum(*(pulse.duration for pulse in self.pulses))
            ),
            takes_self=True,
        ),
    )

    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> SequencePulse:
        result = SequencePulse(
            pulses=tuple(pulse.bind(binding, source_map=source_map) for pulse in self.pulses),
        )
        if source_map is not None:
            source_map[id(result)] = self
        return result

    def unroll(self) -> UnrolledPulse:
        assert self.is_bound, "Cannot unroll a pulse with free parameters"
        return UnrolledPulse.concatenate(subpulse.unroll() for subpulse in self.pulses)

    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        return SequencePulse(pulses=tuple(p.apply_channel_mapping(mapping) for p in self.pulses))

    @property
    def parameters(self) -> frozenset[Parameter]:
        return frozenset(ChainMap(*(pulse.parameters for pulse in self.pulses)))

    @property
    def channels(self) -> frozenset[Channel]:
        return frozenset().union(*(p.channels for p in self.pulses))

    @override
    def simplify(self) -> Pulse:
        if len(self.pulses) == 1:
            return self.pulses[0].simplify()

        simplified_children: list[Pulse] = [pulse.simplify() for pulse in self.pulses]
        flattened_children: list[Pulse] = []

        for child in simplified_children:
            if isinstance(child, SequencePulse):
                flattened_children.extend(child.pulses)
            else:
                flattened_children.append(child)

        if len(flattened_children) == 1:
            return flattened_children[0]

        return SequencePulse(pulses=tuple(flattened_children))

    def __repr__(self) -> str:
        body = ",\n".join(repr(pulse) for pulse in self.pulses)
        body = textwrap.indent(body, " " * 4)
        return f"""{self.__class__.__name__}(
{body},
)"""


@frozen
class ParallelPulse(AnyPulse):
    pulses: tuple[Pulse, ...]
    duration: Expression = field(
        repr=str,
        default=Factory(
            lambda self: (
                0 * unit.s if len(self.pulses) == 0 else expression.max(*(pulse.duration for pulse in self.pulses))
            ),
            takes_self=True,
        ),
    )

    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> ParallelPulse:
        result = ParallelPulse(
            pulses=tuple(pulse.bind(binding, source_map=source_map) for pulse in self.pulses),
        )
        if source_map is not None:
            source_map[id(result)] = self
        return result

    @property
    def parameters(self) -> frozenset[Parameter]:
        return frozenset(ChainMap(*(pulse.parameters for pulse in self.pulses)))

    def unroll(self) -> UnrolledPulse:
        return UnrolledPulse.overlay(subpulse.unroll() for subpulse in self.pulses)

    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        return ParallelPulse(pulses=tuple(p.apply_channel_mapping(mapping) for p in self.pulses))

    @override
    def simplify(self) -> Pulse:
        if len(self.pulses) == 1:
            return self.pulses[0].simplify()
        return ParallelPulse(pulses=tuple(pulse.simplify() for pulse in self.pulses))

    def __repr__(self) -> str:
        body = ",\n".join(repr(pulse) for pulse in self.pulses)
        body = textwrap.indent(body, " " * 4)
        return f"""{self.__class__.__name__}(
{body},
)"""

    @property
    def channels(self) -> frozenset[Channel]:
        return frozenset().union(*(p.channels for p in self.pulses))


@frozen
class ForLoopPulse(AnyPulse):
    start: Literal = field(
        repr=str,
        validator=validators.instance_of(pint.Quantity),
    )
    stop: Literal = field(
        repr=str,
        validator=validators.instance_of(pint.Quantity),
    )
    step: Literal = field(
        repr=str,
        validator=validators.instance_of(pint.Quantity),
    )
    loop_parameter: Parameter = field(repr=str)
    index_parameter: Parameter = field(repr=str)
    body: SequencePulse
    duration: Expression = field(
        repr=str,
        default=Factory(
            lambda self: expression.sum(
                *(
                    self.body.bind(
                        {
                            self.loop_parameter: value,
                            self.index_parameter: i * unit.dimensionless,
                        }
                    ).duration
                    for i, value in enumerate(self.values)
                )
            ),
            takes_self=True,
        ),
    )

    @property
    def values(self) -> np.ndarray:
        values = (
            np.arange(
                self.start.m,
                # Add a small epsilon to avoid the stop value being excluded due to floating point precision
                (self.stop + self.step * 1e-3).m_as(self.start.u),
                self.step.m_as(self.start.u),
            )
            * self.start.u
        )
        return values

    def bind(self, binding: Binding, *, source_map: Optional[dict[int, Pulse]] = None) -> ForLoopPulse:
        result = ForLoopPulse(
            start=self.start,
            stop=self.stop,
            step=self.step,
            loop_parameter=self.loop_parameter,
            index_parameter=self.index_parameter,
            body=self.body.bind(binding, source_map=source_map),
        )
        if source_map is not None:
            source_map[id(result)] = self
        return result

    @property
    def parameters(self) -> frozenset[Parameter]:
        return self.body.parameters.difference({self.loop_parameter, self.index_parameter})

    def iterations(self, *, source_map: Optional[dict[int, Pulse]] = None) -> list[SequencePulse]:
        """
        Returns a list of SequencePulses representing the iterations of the loop.
        Each iteration is a SequencePulse with the body pulse bound to the current loop value and index.
        """
        return [
            self.body.bind(
                binding={
                    self.loop_parameter: value,
                    self.index_parameter: idx * unit.dimensionless,
                },
                source_map=source_map,
            )
            for idx, value in enumerate(self.values)
        ]

    def unroll(self) -> UnrolledPulse:
        assert self.is_bound, "Cannot unroll a pulse with free parameters"

        source_map: dict[int, Pulse] = {}  # id(bound iteration) -> source iteration
        iterations = tuple(self.iterations(source_map=source_map))
        composite = SequencePulse(pulses=iterations)
        unrolled = composite.unroll()

        # The parameter binding in iterations() changes the source pulses.
        # Use the source_map to replace the bound source pulses with the original source pulses.
        source_by_unrolled = {
            unrolled_id: source_map.get(id(source)) for unrolled_id, source in unrolled.source_by_unrolled.items()
        }
        unrolled_by_source = {}
        # While source_id is unique, the source_map may map multiple source_ids to the same source pulse.
        for source_id, unrolled_pulses in unrolled.unrolled_by_source.items():
            key = id(source_map[source_id])
            unrolled_by_source.setdefault(key, []).extend(unrolled_pulses)

        return UnrolledPulse(
            intervals=unrolled.intervals,
            source_by_unrolled=source_by_unrolled,
            unrolled_by_source=frozendict(
                {source_id: tuple(unrolled_pulses) for source_id, unrolled_pulses in unrolled_by_source.items()}
            ),
        )

    def apply_channel_mapping(self, mapping: dict[int, Channel]) -> "Pulse":
        mapped_body = self.body.apply_channel_mapping(mapping)
        assert isinstance(mapped_body, SequencePulse)
        return ForLoopPulse(
            start=self.start,
            stop=self.stop,
            step=self.step,
            loop_parameter=self.loop_parameter,
            index_parameter=self.index_parameter,
            body=mapped_body,
        )

    def __repr__(self) -> str:
        body = ""
        body += f"start={self.start},"
        body += "\n"
        body += f"stop={self.stop},"
        body += "\n"
        body += f"step={self.step},"
        body += "\n"
        body += f"loop_parameter={self.loop_parameter},"
        body += "\n"
        body += f"index_parameter={self.index_parameter},"
        body += "\n"
        body += f"body={self.body},"
        body = textwrap.indent(body, " " * 4)
        return f"""{self.__class__.__name__}(
{body},
)"""

    @property
    def channels(self) -> frozenset[Channel]:
        # Channels used by the loop are those of its body
        return self.body.channels

    @override
    def simplify(self) -> Pulse:
        iterations = self.iterations()
        if len(iterations) == 1:
            return iterations[0].simplify()
        simplified_body = self.body.simplify()
        return ForLoopPulse(
            start=self.start,
            stop=self.stop,
            step=self.step,
            loop_parameter=self.loop_parameter,
            index_parameter=self.index_parameter,
            body=simplified_body
            if isinstance(simplified_body, SequencePulse)
            else SequencePulse(pulses=(simplified_body,)),
        )


LeafPulse = DelayPulse | LaserPulse | DrivePulse | AcquisitionPulse | ChannelMappedPulse
ContainerPulse = SequencePulse | ParallelPulse | ForLoopPulse
Pulse = LeafPulse | ContainerPulse
