from pathlib import Path
from muscriptor import TranscriptionModel

# Downloads the default "medium" variant from HuggingFace (cached under
# ~/.cache/muscriptor/). Also accepts "small"/"large", a local safetensors
# path, or an hf:// / http(s):// URL.
model = TranscriptionModel.load_model()

# Stream events as they're transcribed. Optionally tell the model which
# instruments to expect — run `muscriptor list-instruments` for the names.
for event in model.transcribe("audio.wav", instruments=["acoustic_piano", "drums"]):
    print(event)

# Or get a MIDI file directly
midi_bytes = model.transcribe_to_midi("audio.wav")
Path("out.mid").write_bytes(midi_bytes)

from dataclasses import dataclass
from typing import Generator


@dataclass
class NoteStartEvent:
    pitch: int
    # Start of the note in seconds, from the beginning of the audio.
    start_time: float
    # A unique index for this note, used to match the corresponding
    # NoteEndEvent.
    index: int
    instrument: str


@dataclass
class NoteEndEvent:
    end_time: float
    # The NoteStartEvent this end matches. Convenient for consumers —
    # when serializing (e.g. to JSON) drop this field and rely on
    # `start_event_index` to refer back to the start by id.
    start_event: NoteStartEvent

    @property
    def start_event_index(self) -> int:
        return self.start_event.index


@dataclass
class ProgressEvent:
    # A coarse progress anchor woven into the event stream: `completed` of
    # `total` 5-second chunks have been transcribed. One is emitted up front
    # with completed == 0 (so consumers learn `total`), then one per finished
    # chunk. Advisory only — consumers that just build notes can ignore them.
    completed: int
    total: int


class TranscriptionModel:
    ...
    def transcribe(
            self,
            audio: str | Path | tuple[torch.Tensor, int],
            use_sampling: bool = False,
            temperature: float = 1.0,
            cfg_coef: float = 1.0,
            instruments: list[str] | None = None,
            batch_size: int | None = None,
            no_eos_is_ok: bool = True,
            beam_size: int = 1,
            prelude_forcing: bool = True,
        ) -> Generator[NoteStartEvent | NoteEndEvent | ProgressEvent, None, None]:
        """Transcribe audio into a stream of note events.

        Args:
            audio: Path to an audio file, or a tuple `(tensor, sample_rate)`
                with a float tensor of shape [T] or [1, T] at `sample_rate`
                Hz. The tuple form is useful when the audio is already
                loaded in memory.
            use_sampling: Use temperature sampling instead of greedy decoding.
            temperature: Sampling temperature (only used when use_sampling=True).
            cfg_coef: Classifier-free guidance coefficient. Keep to 1 for the released models (they are post-RL)
            instruments: Optional list of instrument group names to
                restrict decoding to (exact names, e.g.
                ["acoustic_piano", "drums"]). Run `muscriptor
                list-instruments` (or GET /instruments on the server)
                for the full list of valid names. When given, this is a
                hard constraint: every program/drum token outside the
                list is masked out during generation, so no unlisted
                instrument can appear in the output. Leave unset to let
                the model decode whatever instruments it detects.
            batch_size: Number of 5-second chunks processed per forward
                pass. `None` (default) resolves to 1 while prelude forcing
                is on; with `prelude_forcing=False` it picks a value based
                on the device: 1 on CPU, 4 on GPU. Values > 1 trade
                chunk-boundary quality for throughput and therefore require
                `prelude_forcing=False` (ValueError otherwise). Batching
                also delays streaming: several chunks generate together, so
                events belonging to later chunks of a batch won't arrive
                until the whole batch finishes. Within a batch, events are
                always yielded in temporal order; all events from chunk N
                are emitted before any event from chunk N+1.
            no_eos_is_ok: If True, a chunk that doesn't emit EOS within
                the generation budget produces a warning instead of raising.
            beam_size: Beam search width. 1 (default) uses greedy decoding
                (or sampling, with use_sampling=True); >= 2 enables beam
                search, which is slower but can be more accurate.
            prelude_forcing: If True (default), every chunk after the first
                has its tie prologue — the tokens declaring which notes are
                sustained from the previous chunk — teacher-forced from the
                previous chunk's actually-unfinished notes, instead of
                letting the model guess (which occasionally makes a chunk
                restart with the wrong instruments). Requires chunks to be
                generated strictly in order, so it only works with
                batch_size=1 (the default); set prelude_forcing=False
                explicitly to use larger batches.

        Returns:
            Generator of NoteStartEvent, NoteEndEvent and ProgressEvent
            objects. Every
            NoteStartEvent is guaranteed to be followed by exactly one
            matching NoteEndEvent later in the stream (with the same
            `index`). Drum hits appear as a NoteStartEvent immediately
            followed by its matching NoteEndEvent at the same start time
            plus a tiny duration. Note: this tokenizer does not preserve
            velocity (loudness) — only onset/offset timing, pitch, and
            instrument are recovered.
        """

    def transcribe_to_midi(
            self,
            audio: str | Path | tuple[torch.Tensor, int],
            use_sampling: bool = False,
            temperature: float = 1.0,
            cfg_coef: float = 1.0,
            instruments: list[str] | None = None,
            batch_size: int | None = None,
            no_eos_is_ok: bool = True,
            beam_size: int = 1,
            prelude_forcing: bool = True,
        ) -> bytes:
        """Same as `transcribe`, but returns a MIDI file as bytes instead
        of a generator of events. Useful when you want to save the MIDI
        to disk or send it over a network without going through the
        event stream.
        """