"""
===================
A sequence function
===================

The previous lesson,
:doc:`/generated/gallery/05-sequence-modules/01_sequence_modules`, designed an
excitation and a readout with modules. This lesson assembles them into a
complete acquisition. A sequence function
takes the system limits and a protocol: it designs the modules, plays one
repetition per line of k-space and returns the sequence. The function records
the encoding labels with :class:`~pypulseqpp.sequences.Labels`.
:func:`~pypulseqpp.sequences.parameters` reads its protocol, and
:func:`~pypulseqpp.sequences.write` writes the sequence it returns.

The sequence is the RF-spoiled slice-selective gradient echo of
:doc:`/generated/gallery/02-spoiling/01_spoiling`, expressed as a sequence
function rather than as a loop over events. Each sequence in :doc:`/sequences` is
such a function, in its own module of ``pypulseqpp.sequences``. The structure is
described in :doc:`/explanations/sequence-design`.

Learning objectives
-------------------

After this lesson, you should be able to:

- state what a sequence function designs before its loop, plays in the loop
  and records after it;
- write a function of the system limits and a protocol of keyword parameters
  that designs and plays a sequence, recording its encoding indices with
  :class:`~pypulseqpp.sequences.Labels`;
- read the acquisition order back from the ``LIN`` labels of the sequence;
- read the protocol of a sequence function with
  :func:`~pypulseqpp.sequences.parameters`, derive its command-line options
  from the same signature, and write what it returns with
  :func:`~pypulseqpp.sequences.write`.
"""

# %%
# Sequence function
# -----------------
#
# ``system`` is the first parameter and the keyword parameters after it are the
# protocol. Before the loop, the function lowers the limits of ``system`` to the
# gradient and slew limits it is designed for with
# :func:`~pypulseqpp.cap_system`, which never raises a limit, and designs the
# excitation and readout modules; the readout module solves the echo and
# repetition times that are ``None``. The loop plays one repetition per line of
# k-space. A label keeps its value until an event changes it, so
# :class:`~pypulseqpp.sequences.Labels` returns only the ``LIN`` events a block
# needs. After the loop, the geometry a reconstruction reads is recorded as
# definitions and the function returns the sequence.

import numpy as np

import pypulseqpp as pp
import pypulseqpp.sequences as design


def spoiled_gradient_echo(
    system: pp.Opts,
    *,
    fov: float = 220e-3,
    matrix: int = 128,
    slice_thickness: float = 5e-3,
    flip_angle_deg: float = 12.0,
    te: float | None = None,
    tr: float | None = None,
) -> pp.Sequence:
    """RF-spoiled 2D Cartesian gradient echo, one line per repetition.

    Parameters
    ----------
    fov : float, default=0.22
        Isotropic field of view (m).
    matrix : int, default=128
        Matrix size along both encoded axes.
    slice_thickness : float, default=0.005
        Slice thickness (m).
    flip_angle_deg : float, default=12.0
        Excitation flip angle (degrees).
    te : float | None, default=None
        Echo time (s). ``None`` is as short as the readout admits.
    tr : float | None, default=None
        Repetition time (s). ``None`` is as short as possible.
    """
    system = pp.cap_system(system, max_grad=40.0, max_slew=150.0)
    excitation = design.SpatialSelectiveExcitation(
        system,
        flip_angle_deg=flip_angle_deg,
        thickness_m=slice_thickness,
        duration_s=3e-3,
    )
    readout = design.LineReadout2D(
        system,
        excitation.rf,
        excitation.gz,
        excitation.gz_reph,
        fov=(fov, fov),
        matrix=(matrix, matrix),
        te=te,
        tr=tr,
    )

    # Quadratic RF spoiling: the phase increment grows by 117 degrees per
    # repetition.
    phases = np.deg2rad(117.0) * np.cumsum(np.arange(matrix))

    seq, labels = pp.Sequence(system), design.Labels()
    for line in range(matrix):
        phase = phases[line] % (2 * np.pi)
        readout.rf.phase_offset = readout.adc.phase_offset = phase
        step = (line - matrix // 2) / (matrix / 2)
        seq.add_block(readout.rf, readout.gz, *labels(LIN=line))
        seq.add_block(
            readout.gx_pre, pp.scale_grad(readout.gy_pre, step), readout.gz_reph
        )
        seq.add_block(readout.gx, readout.adc)
        seq.add_block(readout.gx_spoil, pp.scale_grad(readout.gy_rew, step))
        if getattr(readout, "wait_tr", None) is not None:
            seq.add_block(readout.wait_tr)

    seq.set_definition("FOV", [fov, fov, slice_thickness])
    seq.set_definition("Matrix", [matrix, matrix, 1])
    seq.set_definition("Name", "tour_gre_2d")
    return seq


# %%
# Sequence construction
# ---------------------
#
# Calling the function designs the whole scan. The system is the default
# :class:`~pypulseqpp.Opts`, whose limits the function lowers to at most
# 40 mT/m and 150 T/m/s.

seq = spoiled_gradient_echo(pp.Opts(), matrix=128)
print(f"{seq.num_blocks} blocks, {seq.duration()[0]:.2f} s")
print("timing:", seq.check_timing()[0])
print("definitions:", sorted(seq.definitions))

# %%
# Sequence diagram
# ----------------

seq.paper_plot(tr=32)

# %%
# Acquisition order
# -----------------
#
# The first change of a label is a SET, a change that repeats the previous
# change is an INC, and an unchanged value writes no event.

labels = design.Labels()
for line in (0, 1, 2, 3, 3):
    print(f"LIN={line}:", [(e.type, e.value) for e in labels(LIN=line)])

# %%
# Every acquisition of the sequence carries its ``LIN`` label, so the order the
# loop played is read back from the sequence rather than reconstructed.

pp.plot.plot_kspace(seq, color_by="order", plane="xy", show_trajectory=False)

# %%
# Protocol
# --------
#
# The type and default of a protocol parameter are read from the signature, and
# its unit from the first parenthesised group of its description in the
# ``Parameters`` section. A protocol editor reads this record, and the command
# line derives its options from the same signature and ``Parameters`` section.

for name, parameter in design.parameters(spoiled_gradient_echo).items():
    print(
        f"{name:16}{parameter.type.__name__:7}{parameter.unit:9}"
        f"default {parameter.default}"
    )

# %%
# Command-line interface
# ----------------------
#
# :func:`pypulseqpp.cli.run` derives one option from each keyword parameter,
# with the help text from the ``Parameters`` section, builds ``system`` from
# the limit options and writes the result with
# :func:`~pypulseqpp.sequences.write`. A script that ends with
#
# .. code-block:: python
#
#    if __name__ == "__main__":
#        raise SystemExit(cli.run(spoiled_gradient_echo, sys.argv[1:]))
#
# is run as ``python gre.py --matrix 64 --te 5e-3 -o gre.seq``. A function that
# returns a list of sequences is written as a chain of files, each naming the
# next with ``NextSequence``.

from pathlib import Path
from tempfile import mkdtemp

path = Path(mkdtemp()) / "tour_gre_2d.seq"
print([Path(written).name for written in design.write(path, seq)])