Source code for suboptimumg.log_analysis.yaw_artifact

from __future__ import annotations

import pickle
from pathlib import Path

from perda.core_data_structures import SingleRunData
from pydantic import BaseModel, ConfigDict, Field

from .yaw_fit_models import YawFit
from .yaw_segmentation_models import MaskReport




class SegmentationConfig(BaseModel):
    """A record of how one log was cut into sub-windows.

    Recorded alongside the sub-windows so a saved artifact answers "where did
    these come from" without re-reading the notebook that made it.
    """

    logfile: str = Field(description="Source log file path")
    trim_range_s: tuple[float, float] | None = Field(
        default=None,
        description="Time range the log was trimmed to (s), or None if untrimmed",
    )
    manual_windows: list[tuple[float, float]] = Field(
        default_factory=list,
        description="Operator-picked coarse (start, end) windows (s)",
    )
    gps_lag_s: float = Field(default=0.0, description="GPS channel lag correction applied (s)")
    min_speed_mps: float = Field(default=4.0, description="Speed floor a sample had to clear (m/s)")
    max_slip_deg: float = Field(
        default=7.0, description="Body-slip ceiling a sample had to stay under (deg)"
    )
    min_window_s: float = Field(default=8.0, description="Shortest sub-window kept (s)")
    target_hz: float = Field(
        default=100.0,
        description="Uniform grid the sub-windows were resampled onto (Hz)",
    )
    mask_reports: list[MaskReport] = Field(
        default_factory=list, description="Per-manual-window masking diagnostics"
    )





class YawArtifact(BaseModel):
    """One log's yaw-response work: its sub-windows, how they were made, and the fits.

    Bundles everything a later session needs to re-evaluate or compare against
    this log without re-running preparation. ``save`` / ``load`` round-trip it
    through pickle, because the PERDA ``SingleRunData`` sub-windows have no
    JSON form. Treat the file as a local cache, not an interchange format: it
    is tied to the installed PERDA/pydantic versions, and loading one executes
    arbitrary code, so only load artifacts you produced.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    log_name: str = Field(description="Identifying name for the source log")
    config: SegmentationConfig = Field(description="How the sub-windows were produced")
    subwindows: list[SingleRunData] = Field(description="The clean, resampled sub-windows")
    fits: dict[str, YawFit] = Field(
        default_factory=dict, description="Fitted models, keyed by spec name"
    )



    def add_fit(self, fit: YawFit) -> None:
        """Record ``fit`` under its spec name, replacing any fit of that name."""
        self.fits[fit.spec.name] = fit




    def save(self, path: str | Path) -> Path:
        """Pickle this artifact to ``path``, creating parent directories.

        Returns
        -------
        Path
            The resolved output path.
        """
        out = Path(path)
        out.parent.mkdir(parents=True, exist_ok=True)
        with open(out, "wb") as f:
            pickle.dump(self, f)
        return out




    @staticmethod
    def load(path: str | Path) -> YawArtifact:
        """Inverse of ``save``. Only load artifacts you produced -- see the class docstring.

        Raises
        ------
        TypeError
            If ``path`` does not hold a ``YawArtifact``.
        """
        with open(Path(path), "rb") as f:
            obj = pickle.load(f)
        if not isinstance(obj, YawArtifact):
            raise TypeError(f"{path} did not contain a YawArtifact")
        return obj