Index

def with_alpha(self, alpha: float) -> Vec4:
    """Returns the color with a custom transparency level."""
    return Color.hex_to_rgba(self.value, alpha=alpha)

@classmethod
def hex_to_rgba(cls, hex_str: str, alpha: float = 1.0) -> Vec4:
    """Converts '#RRGGBB' or 'RRGGBB' to normalized [0, 1] RGBA."""
    hex_str = hex_str.lstrip("#")
    # Convert hex to integers using bit-shifting
    lv = len(hex_str)
    if lv == 6:
        rgb = tuple(int(hex_str[i : i + 2], 16) for i in (0, 2, 4))
    else:
        raise ValueError(f"Invalid hex color: {hex_str}. Expected 6 characters.")

    return np.array([*(np.array(rgb) / 255.0), alpha])

@classmethod
def rgba255_to_rgba(cls, rgba255: Vec4) -> Vec4:
    """Converts [0-255, 0-255, 0-255, 0-1] to normalized [0, 1] RGBA."""
    res = np.array(rgba255, dtype=float)
    res[:3] /= 255.0
    return res

@classmethod
def rgba_to_hex(cls, rgba: Vec4) -> str:
    """Converts normalized RGBA to '#RRGGBB' (alpha is discarded)."""
    rgba = np.asarray(rgba, dtype=float)
    rgb255 = (rgba[:3] * 255).astype(int)
    return "#{:02x}{:02x}{:02x}".format(*rgb255)

@classmethod
def rgba_to_rgba255(cls, rgba: Vec4) -> Vec4:
    """Converts normalized RGBA to [0-255, 0-255, 0-255, 0-1]."""
    res = np.array(rgba, dtype=float)
    res[:3] *= 255.0
    return res

@classmethod
def from_series(cls, series: pd.Series, kind: InterpOptions = "linear") -> Self:
    """Create an interpolator directly from a Pandas Series."""
    return cls(x=np.asarray(series.index), y=np.asarray(series.values), kind=kind)

def get_sphere_to_sphere_proximity(
    self,
    state: MjState,
) -> tuple[float, Vec3, Vec3, bool]:
    """
    Calculates the shortest distance between two geometries using their bounding spheres.

    Args:
        state: The paired MuJoCo model and data instance.

    Returns:
        tuple[float, Vec3, Vec3, bool]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and if the estimated distance exceeds dist_max.

    """
    # get world orientations and origins
    origin_geom_1 = self.geom_1.rt_pos(state)
    mat_geom_1 = self.geom_1.rt_xmat(state)

    origin_geom_2 = self.geom_2.rt_pos(state)
    mat_geom_2 = self.geom_2.rt_xmat(state)

    if np.isnan(self.geom_1._rad):
        self.geom_1._rad, self.geom_1._local_centroid = self.geom_1.vertex_max_norm(
            state.model
        )

    if np.isnan(self.geom_2._rad):
        self.geom_2._rad, self.geom_2._local_centroid = self.geom_2.vertex_max_norm(
            state.model
        )

    # shift centers to pre-calculated centroids
    pos_geom_1 = origin_geom_1 + (mat_geom_1 @ self.geom_1._local_centroid)
    pos_geom_2 = origin_geom_2 + (mat_geom_2 @ self.geom_2._local_centroid)

    rad_geom_1 = self.geom_1._rad
    rad_geom_2 = self.geom_2._rad

    vec_geom_1_to_geom_2 = pos_geom_2 - pos_geom_1
    d_centers = float(np.linalg.norm(vec_geom_1_to_geom_2))
    dist = d_centers - (rad_geom_1 + rad_geom_2)

    dist = max(0.0, dist)  # clip to zero
    exceeds_dist_max = dist > self.dist_max

    if d_centers > 1e-9:
        unit_vec = vec_geom_1_to_geom_2 / d_centers
        p1 = pos_geom_1 + (unit_vec * rad_geom_1)
        p2 = pos_geom_2 - (unit_vec * rad_geom_2)
    else:
        p1 = pos_geom_1
        p2 = pos_geom_2

    self.update_last(p1, p2, state)
    return dist, p1, p2, exceeds_dist_max

def get_convex_hull_proximity(
    self,
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the shortest distance between two geometries using their convex hull.

    Args:
        state: The paired MuJoCo model and data instance.

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    # ========== BROADPHASE ==========
    if self.geom_1._proximity_configured_for != ProximityType.CONVEX_HULL:
        self.geom_1.bake_proximity(state.model, ProximityType.CONVEX_HULL)

    if self.geom_2._proximity_configured_for != ProximityType.CONVEX_HULL:
        self.geom_2.bake_proximity(state.model, ProximityType.CONVEX_HULL)

    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(state)

    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== NARROWPHASE ==========
    # temp buffer for MuJoCo's 6-element output [x1,y1,z1, x2,y2,z2]
    mj_fromto = np.zeros(6)
    min_dist = mujoco.mj_geomDistance(
        m=state.model,
        d=state.data,
        geom1=self.geom_1.get_id(state.model),
        geom2=self.geom_2.get_id(state.model),
        distmax=self.dist_max,
        fromto=mj_fromto,
    )

    min_dist = max(0.0, min_dist)  # clip from below to zero

    p1 = mj_fromto[:3].copy()
    p2 = mj_fromto[3:6].copy()
    self.update_last(p1, p2, state)
    return min_dist, p1, p2, ProximityType.CONVEX_HULL

def get_vertex_to_face_proximity(
    self,
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the vertex to face distance using a multi-phase Bounding Volume Hierarchy (BVH) query.

    Phases:
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Mid Phase: BVH Traversal (eliminating triangle groups). No exit here.
        3. Narrow Phase: Point-to-Face proximity.

    Args:
        state: The paired MuJoCo model and data instance.

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    if self.geom_1._proximity_configured_for != ProximityType.VERTEX_TO_FACE:
        self.geom_1.bake_proximity(state.model, ProximityType.VERTEX_TO_FACE)

    if self.geom_2._proximity_configured_for != ProximityType.VERTEX_TO_FACE:
        self.geom_2.bake_proximity(state.model, ProximityType.VERTEX_TO_FACE)

    assert self.geom_1._baked_query and self.geom_2._baked_query
    assert (
        self.geom_2._local_verts is not None
        and self.geom_2._local_verts is not None
    )

    # ========== BROADPHASE: Sphere-Sphere check ==========
    # find center to center to center distance and return early if broad phase
    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(state)
    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== COORDINATE TRANSFORMATION ==========
    pos_geom_1 = self.geom_1.rt_pos(state)
    pos_geom_2 = self.geom_2.rt_pos(state)

    mat_geom_1 = self.geom_1.rt_xmat(state)  # already Mat3 (3x3)
    mat_geom_2 = self.geom_2.rt_xmat(state)
    rel_pos = pos_geom_2 - pos_geom_1

    # ========== NARROWPHASE A: Geom_1-Surface vs. Geom_2-Vertices ==========
    # trimesh uses a BVH internall here (Mid-phase) to find closest triangles
    # combine transforms from geom_1 to geom_2: V_local_geom_1 = R_geom_1.T @ (R_geom_2 @ V_local_geom_2 + p_geom_2 - p_geom_1)
    geom_2_v_in_geom_1 = (
        self.geom_2._local_verts @ mat_geom_2.T + rel_pos
    ) @ mat_geom_1
    pts_on_geom_1, dist_a, _ = self.geom_1._baked_query.on_surface(
        geom_2_v_in_geom_1
    )
    idx_a = np.argmin(dist_a)
    min_a = dist_a[idx_a]

    # ========== NARROWPHASE B: Geom_1-Vertices vs. Geom_2-Surface  ==========
    # transform geom_1 vertices into geom_2's local frame
    geom_1_v_in_geom_2 = (
        self.geom_1._local_verts @ mat_geom_1.T - rel_pos
    ) @ mat_geom_2
    pts_on_geom_2, dist_b, _ = self.geom_2._baked_query.on_surface(
        geom_1_v_in_geom_2
    )
    idx_b = np.argmin(dist_b)
    min_b = dist_b[idx_b]

    # ========== CLEANUP ==========
    # find global min
    if min_a < min_b:
        min_dist = float(min_a)
        p1 = (pts_on_geom_1[idx_a] @ mat_geom_1.T) + pos_geom_1
        p2 = (geom_2_v_in_geom_1[idx_a] @ mat_geom_1.T) + pos_geom_1

        self.update_last(p1, p2, state)
        return min_dist, p1, p2, ProximityType.VERTEX_TO_FACE
    else:
        min_dist = float(min_b)

        # pt_on_geom_2 was calculated in geom_2's local frame
        p2 = (pts_on_geom_2[idx_b] @ mat_geom_2.T) + pos_geom_2
        p1 = (geom_1_v_in_geom_2[idx_b] @ mat_geom_2.T) + pos_geom_2

        self.update_last(p1, p2, state)
        return min_dist, p1, p2, ProximityType.VERTEX_TO_FACE

def get_face_to_face_proximity(
    self,
    state: MjState,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the face to face distance using a multi-phase Bounding Volume Hierarchy (BVH) query.

    This is more accurate than the vertex to face method, but comes at higher computational cost.

    Phases:
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Mid Phase: BVH Traversal (eliminating triangle groups). No exit here.
        3. Narrow Phase: Face-to-Face proximity.

    Args:
        state: The paired MuJoCo model and data instance.

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    if self.geom_1._proximity_configured_for != ProximityType.FACE_TO_FACE:
        self.geom_1.bake_proximity(state.model, ProximityType.FACE_TO_FACE)

    if self.geom_2._proximity_configured_for != ProximityType.FACE_TO_FACE:
        self.geom_2.bake_proximity(state.model, ProximityType.FACE_TO_FACE)

    assert self.geom_1._baked_manager and self.geom_2._baked_manager

    # ========== BROADPHASE: Sphere-Sphere check ==========

    # find center to center to center distance and return early if broad phase
    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(state)
    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== NARROWPHASE ==========
    # set the other transformation relative to geom_1's local frame
    t_geom_1 = np.eye(4)
    t_geom_1[:3, :3] = self.geom_1.rt_xmat(state)
    t_geom_1[:3, 3] = self.geom_1.rt_pos(state)

    t_geom_2 = np.eye(4)
    t_geom_2[:3, :3] = self.geom_2.rt_xmat(state)
    t_geom_2[:3, 3] = self.geom_2.rt_pos(state)

    self.geom_1._baked_manager.set_transform(self.geom_1.name, t_geom_1)
    self.geom_2._baked_manager.set_transform(self.geom_2.name, t_geom_2)

    # CollisionManager returns distance and the two closest points
    result = self.geom_1._baked_manager.min_distance_other(
        self.geom_2._baked_manager, return_data=True
    )
    min_dist = float(result[0])  # pyright: ignore[reportIndexIssue]
    min_dist = max(0.0, min_dist)  # clip to zero
    data = result[1]  # pyright: ignore[reportIndexIssue]

    assert data
    p1 = data.point(self.geom_1.name)  # pyright: ignore[reportAttributeAccessIssue]
    p2 = data.point(self.geom_2.name)  # pyright: ignore[reportAttributeAccessIssue]

    self.update_last(p1, p2, state)
    return min_dist, p1, p2, ProximityType.FACE_TO_FACE

def get_proximity(self, state: MjState) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the shortest distance between two geometries using the specified proximity algorithm.

    This is a general dispatcher method that routes to different proximity calculation algorithms based on the `algorithm` parameter. Each mode offers different tradeoffs between speed and precision:

    **Modes:**
        - `SPHERE_TO_SPHERE`: Fastest. Uses bounding sphere distance only (broadphase).
        - `CONVEX_HULL`: Fast & accurate. Uses MuJoCo's convex hull-based distance (default).
        - `VERTEX_TO_FACE`: Accurate. Multi-phase BVH with vertex-to-surface queries.
        - `FACE_TO_FACE`: Most accurate but slowest. Full mesh-to-mesh distance calculation.

    **Phases (for non-sphere modes):**
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Narrow Phase: Algorithm-specific distance calculation.

    Args:
        state: The paired MuJoCo model and data instance.

    Returns:
        tuple[float, ProximityType]: If fromto=False, returns the unsigned (`>= 0`) minimum distance and which algorithm produced the result.

        tuple[tuple[float, Vec3, Vec3], ProximityType]: If fromto=True, returns the minimum distance, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which algorithm produced the result.

    The result is cached per-timestep (keyed on `state.data.time`), so calling this more than once during the same step (e.g. once from `request()`'s telemetry sampler and again from a user-defined runtime input/Load that reads the same `Proximity` instance) only pays for the underlying calculation once.

    """
    is_cached = (
        self._last_prox_type is not None
        and not np.isnan(self._last_dist)
        and state.data.time == self._last_t
    )
    if is_cached:
        assert self._last_prox_type is not None  # narrows for the type checker
        return self._last_dist, self._last_p1, self._last_p2, self._last_prox_type

    match self.algorithm:
        case ProximityType.SPHERE_TO_SPHERE:
            d_est, p1, p2, _skip = self.get_sphere_to_sphere_proximity(state)
            result = (d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE)
        case ProximityType.CONVEX_HULL:
            result = self.get_convex_hull_proximity(state)
        case ProximityType.VERTEX_TO_FACE:
            result = self.get_vertex_to_face_proximity(state)
        case ProximityType.FACE_TO_FACE:
            result = self.get_face_to_face_proximity(state)
        case _:
            msg = f"Method for {self.algorithm.name} not implemented."
            logger.error(msg)
            raise NotImplementedError(msg)

    # the configured algorithm's broadphase sphere-to-sphere check can exit early
    # (objects farther apart than dist_max), so the phase that actually produced
    # this result can flip between SPHERE_TO_SPHERE and the configured narrowphase
    # algorithm from one call to the next; surface that for performance debugging
    prox_type = result[3]
    if self._last_prox_type is not None and prox_type != self._last_prox_type:
        logger.debug(
            f"Proximity {self.pair_name} switched phase: "
            f"{self._last_prox_type.name} -> {prox_type.name} (t={state.data.time:.4f})"
        )
    self._last_prox_type = prox_type
    self._last_dist = result[0]
    # _last_t/_last_p1/_last_p2 were already set by update_last() inside whichever
    # method above produced `result`

    return result

def request(
    self,
    signal_manager: SignalManager | None = None,
    channels: list[Literal["dist", "fromto", "prox_type"]]
    | dict[Literal["dist", "fromto", "prox_type"], dict[str, Any] | None] = [
        "dist",
        "prox_type",
    ],
):
    """
    Registers specific channels for logging.

    | Channel     | Description                                                            | Type   |
    |:------------|:-----------------------------------------------------------------------|:-------|
    | `dist`      | minimum distance between the geom pair, per the proximity algorithm    | scalar |
    | `fromto`    | world coordinates of the nearest point on each geom                    | xyz    |
    | `prox_type` | the `ProximityType` used to compute `dist` and `fromto`, as an integer | scalar |

    Each channel is posted under `subgroups=(pair_name,)`, where `pair_name` is `f"{geom_1.name}_to_{geom_2.name}"`.

    * A `scalar` is posted as a single value with `attr=channel`.
    * An `xyz` is a cartesian vector without a magnitude component, posted as 3 values (`x`, `y`, `z`). `fromto` posts one `xyz` for each geom in the pair, under `subgroups=(pair_name, "fromto", geom_name)`.

    Only the computations required by the requested channels are performed each timestep.

    Each signal is tagged with built-in `dimension`/`unit` metadata for its channel (`dist`/`fromto` as a length, `prox_type` as dimensionless).

    If `signal_manager` is omitted, the `SignalManager` of the active `RuntimeManager` `with` block is used. If that `RuntimeManager` has no `SignalManager` configured, this is a no-op.

    Args:
        signal_manager: The signal manager to register the sampler with.
        channels: The proximity data channels to log. Pass a list to select channels, or a dict mapping channel name to metadata overrides (or `None`) to select channels and attach per-channel metadata in one step.

    """
    from mujoco_mojo.runtime.signal_manager import resolve_signal_manager

    signal_manager = resolve_signal_manager(signal_manager)
    if signal_manager is None:
        return

    if isinstance(channels, dict):
        _meta = cast("dict[str, dict[str, Any] | None]", channels)
        channels = list(channels.keys())
    else:
        _meta = {}

    pair_name = self.pair_name
    _prox_attrs = {"dist", "fromto", "prox_type"}
    needs_proximity = bool(set(channels) & _prox_attrs)

    if needs_proximity:
        self._requested = True

    def sample(state: MjState):
        dist: float = np.nan
        p1: Vec3 = np.zeros(3)
        p2: Vec3 = np.zeros(3)
        prox_type = ProximityType.SPHERE_TO_SPHERE

        if needs_proximity:
            dist, p1, p2, prox_type = self.get_proximity(state)

        for channel in channels:
            meta = merge_signal_metadata(
                _REQUEST_CHANNEL_METADATA.get(channel),
                channel,
                _meta,
                unit_system=state.us,
            )

            match channel:
                case "dist":
                    signal_manager.post(
                        value=dist,
                        category=SignalCategory.PROXIMITIES,
                        subgroups=(pair_name,),
                        attr=channel,
                        metadata=meta,
                    )
                case "fromto":
                    for i, attr in enumerate("xyz"):
                        signal_manager.post(
                            value=float(p1[i]),
                            category=SignalCategory.PROXIMITIES,
                            subgroups=(pair_name, channel, str(self.geom_1.name)),
                            attr=attr,
                            metadata=meta,
                        )
                    for i, attr in enumerate("xyz"):
                        signal_manager.post(
                            value=float(p2[i]),
                            category=SignalCategory.PROXIMITIES,
                            subgroups=(pair_name, channel, str(self.geom_2.name)),
                            attr=attr,
                            metadata=meta,
                        )
                case "prox_type":
                    signal_manager.post(
                        value=float(prox_type.value),
                        category=SignalCategory.PROXIMITIES,
                        subgroups=(pair_name,),
                        attr=channel,
                        metadata=meta,
                    )
                case _:
                    continue

    signal_manager.register_sampler(sample)

def run(
    self,
    global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
    clean_workdir: bool = False,
    cleanup_delay: int = 10,
    execution_mode: ExecutionMode = ExecutionMode.LOCAL,
    trial_nums: list[int] | None = None,
) -> bool:
    """Vectors a job to be either computed locally or to be orchestrated by SLURM."""
    funcs = {
        "generator": self.generator,
        "runtime": self.runtime,
        "objective": self.objective,
    }
    non_none = {k: v for k, v in funcs.items() if v is not None}
    seen: dict[int, str] = {}
    for role, func in non_none.items():
        fid = id(func)
        if fid in seen:
            msg = f"'{getattr(func, '__qualname__', func)}' is assigned to both '{seen[fid]}' and '{role}'. Each function must be unique."
            logger.error(msg)
            raise ValueError(msg)
        seen[fid] = role

    if clean_workdir:
        if self.config.resume:
            msg = "clean_workdir and resume are mutually exclusive with one another. Use one or the other."
            logger.error(msg)
            raise ValueError(msg)

        self.force_remove_dir(countdown_from=cleanup_delay, path=self.workdir)

    self.workdir.mkdir(parents=True, exist_ok=True)
    if not (self.workdir / ".gitignore").exists():
        (self.workdir / ".gitignore").write_text("*", encoding="utf-8")

    # SLURM workers re-enter run() with execution_mode=LOCAL; these only need
    # to happen once, during orchestration, to avoid every worker racing to
    # write the same shared files
    if self.slurm_trial_id is None:
        write_dojo_script(self.workdir)
        self.capture_environment()

    match execution_mode:
        case ExecutionMode.LOCAL:
            return self.run_local(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        case ExecutionMode.SLURM:
            return self.orchestrate_slurm(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        case _:
            msg = f"No run command has been configured for execution mode {execution_mode}"
            logger.error(msg)
            raise NotImplementedError(msg)

def orchestrate_slurm(
    self,
    global_overrides: NamedValueDict[NDArray],
    trial_nums: list[int] | None = None,
) -> bool:
    """Generates an sbatch script and submits the job array to SLURM for a given config."""
    (self.workdir / "logs").mkdir(exist_ok=True)

    from rich.console import Console

    Console().print(
        "[dim]Tip: pass --slurm-config/-sc with a flat JSON file to auto-inject extra "
        "#SBATCH lines (prefix a key with 'sbatch.', e.g. 'sbatch.account') and/or "
        "environment variables (any other key) into the generated submission script.[/dim]"
    )

    if isinstance(self.config, MonteCarloConfig):
        try:
            had_fails = self.orchestrate_slurm_monte_carlo(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        except (BdbQuit, KeyboardInterrupt):
            print("\n")
            logger.error("Aborted SLURM orchestration!")
            had_fails = True
    else:
        msg = f"A SLURM configuration for {self.config.__class__.__name__} has not been implemented."
        logger.error(msg)
        raise NotImplementedError(msg)
    return had_fails

def execute_single_trial(
    self, trial_num: int, seed: int | None, overrides_payload: dict
) -> tuple[MojoModel | None, TrialStatus, MjState | None]:
    """
    Helper to package a Trial and run it.

    The model and state are live in-process objects; callers that need to cross a ProcessPoolExecutor boundary (i.e. parallel MC) must use `_execute_trial_subprocess` instead, which drops the non-picklable objects before the result is serialized back.

    Returns `(mojo_model, trial_status, state)`.
    """
    overrides = NamedValueDict[NDArray].model_validate(overrides_payload)

    trial = Trial(
        trial_num=trial_num,
        base_dir=self.workdir,
        xml_name=self.xml_name,
        model_config_name=self.model_config_name,
        padding_style=self.config.padding_style,
    )

    mojo_model, trial_status, state = trial.run(
        generator=self.generator,
        runtime=self.runtime,
        seed=seed,
        overrides=overrides,
        gen_args=self.gen_args,
        gen_kwargs=self.gen_kwargs,
        run_args=self.run_args,
        run_kwargs=self.run_kwargs,
    )

    # write distribution tables from this process so they never need to
    # cross the pickle boundary; FileLock serializes concurrent workers,
    # and the inner exists() check ensures only the first one writes
    if mojo_model is not None and mojo_model.dists:
        stochas_dir = self.workdir / STOCHAS_DIR_NAME
        stochas_dir.mkdir(exist_ok=True)
        dists_json_path = stochas_dir / STOCHAS_DISTS_FNAME
        with FileLock(stochas_dir / ".dists.lock"):
            if not dists_json_path.exists():
                mojo_model.dists.to_tables(stochas_dir)
                tmp = stochas_dir / "dists.tmp.json"
                tmp.write_text(mojo_model.dists.model_dump_json(), encoding="utf-8")
                tmp.replace(dists_json_path)

    return mojo_model, trial_status, state

def run_monte_carlo(
    self,
    global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
    trial_nums: list[int] | None = None,
) -> bool:
    """Orchestrates a Monte Carlo job."""
    if self.slurm_trial_id is not None:
        tn = self.slurm_trial_id
        logger.info(f"SLURM Worker detected. Executing Trial {tn}")
        _mojo_model, trial_status, _ = self.execute_single_trial(
            trial_num=tn,
            seed=self.seed,
            overrides_payload=global_overrides.model_dump(),
        )

        return trial_status.completion == Completion.ERROR

    job_trial_nums = trial_nums if trial_nums else self.config.trial_nums

    # initialize the status tracker
    status_tracker = JobStatus(
        workdir=self.workdir.resolve(),
        job_type=JobType.MONTE_CARLO,
        execution_mode=ExecutionMode.LOCAL,
        n_proc=self.config.n_proc,
        seed=self.seed,
        padding_style=self.config.padding_style,
        generator=MojoRunner.inspect_protocol(self.generator),
        runtime=MojoRunner.inspect_protocol(self.runtime),
        objective=MojoRunner.inspect_protocol(self.objective),
        gen_args_used=bool(self.gen_args),
        gen_kwargs_used=bool(self.gen_kwargs),
        run_args_used=bool(self.run_args),
        run_kwargs_used=bool(self.run_kwargs),
        trial_nums=job_trial_nums,
    )

    # decide which trials to execute
    if self.config.resume:
        self._renumber_trial_folders(
            self.workdir.resolve(), self.config.padding_style
        )
        status_tracker.refresh_from_disk(n_proc=self.config.n_proc)
    status_tracker.dump_to_path(self.workdir / JOB_STATUS_FNAME)

    to_run = status_tracker.pending_trial_nums

    if not to_run:
        logger.info("All trials were already completed. Nothing to do.")
        return bool(status_tracker.unsuccessful_trial_nums)

    if self.config.is_parallel:
        logger.info(
            f"Running {len(to_run)} trials with {self.config.n_proc} processors. {status_tracker.n_done}/{self.config.n_trial} ({status_tracker.progress:.2%}) trials completed."
        )
        # needed for logging on Windows
        with multiprocessing.Manager() as m:
            log_queue = m.Queue()

            parent_log_level = logging.getLogger().getEffectiveLevel()
            listener = QueueListener(log_queue, *logging.getLogger().handlers)
            listener.start()

            executor = ProcessPoolExecutor(
                max_workers=self.config.n_proc,
                initializer=worker_init,
                initargs=(log_queue, parent_log_level),
            )
            try:
                future_to_tn = {
                    executor.submit(
                        self._execute_trial_subprocess,
                        tn,
                        self.seed,
                        global_overrides.model_dump(),
                    ): tn
                    for tn in to_run
                }
                for f in as_completed(future_to_tn):
                    tn = future_to_tn[f]
                    try:
                        trial_status = f.result()
                        status_tracker.update_trial(status=trial_status)
                    except (BdbQuit, KeyboardInterrupt):
                        # user is quitting from breakpoint() or CTRL+C
                        raise
                    except Exception as e:
                        logger.exception(f"Trial {tn} failed: {e}")
                        status_tracker.update_trial(
                            status=TrialStatus(
                                trial_num=tn, completion=Completion.ERROR
                            )
                        )
                    status_tracker.generate_report()
            except (BdbQuit, KeyboardInterrupt):
                # allows killing the job with one CTRL+C
                logger.warning("Interrupt recieved. Stopping all trials.")
                executor.shutdown(wait=False, cancel_futures=True)
                raise
            finally:
                listener.stop()
                executor.shutdown(wait=True)
    else:
        for tn in to_run:
            try:
                _, trial_status, _ = self.execute_single_trial(
                    trial_num=tn,
                    seed=self.seed,
                    overrides_payload=global_overrides.model_dump(),
                )
                status_tracker.update_trial(
                    status=trial_status,
                )
            except (BdbQuit, KeyboardInterrupt):
                # user is quitting from breakpoint() or CTRL+C
                raise
            except Exception as e:
                logger.exception(f"A trial failed with error: {e}")
                status_tracker.update_trial(
                    status=TrialStatus(trial_num=tn, completion=Completion.ERROR)
                )
            status_tracker.generate_report()

    status_tracker.generate_report(alert_generation=True)
    return bool(status_tracker.unsuccessful_trial_nums)