server/simulator/latent_state.py

"""Latent (hidden) state of the LHC simulator.

The agent never sees these structures. They define the ground-truth particle
properties, detector imperfections, experiment progress flags, and the live
resource budget.
"""

from __future__ import annotations

from typing import Dict, List, Optional

from pydantic import BaseModel, Field


# ── Particle truth ────────────────────────────────────────────────────────


class LatentParticle(BaseModel):
    """The hidden mystery particle that the agent must discover.

    Defines the true mass, width, decay branching ratios, spin, parity,
    production cross-section, and dominant decay channel. The agent has to
    recover these values from noisy observations.
    """

    name: str = "X"
    mass_gev: float = 125.0
    width_gev: float = 0.004
    spin: int = 0  # 0, 1, or 2
    parity: str = "+"  # "+" or "-"
    cross_section_fb: float = 50.0  # signal cross-section in femtobarns
    decay_branching: Dict[str, float] = Field(
        default_factory=lambda: {
            "diphoton": 0.0023,
            "dilepton_ee": 0.00003,
            "dilepton_mumu": 0.00022,
            "four_lepton": 0.000125,
            "bb": 0.58,
            "dijet": 0.30,
        },
        description="Branching ratio (BR) per decay channel, sums to ~1.",
    )
    primary_channel: str = "diphoton"


# ── Detector & accelerator state ─────────────────────────────────────────


class DetectorState(BaseModel):
    """Hidden detector and accelerator parameters that shape noise.

    These influence resolution, trigger efficiency, pileup, and systematic
    uncertainties applied to every observation.
    """

    detector_resolution_gev: float = 1.5  # absolute mass resolution σ_m
    pileup_mu: float = 30.0               # average pileup interactions per crossing
    trigger_efficiency: float = 0.85
    luminosity_uncertainty: float = 0.025  # 2.5% relative uncertainty
    energy_scale_offset: float = 0.0       # systematic shift in GeV
    energy_scale_uncertainty: float = 0.3  # σ on the scale
    background_shape_alpha: float = -2.5   # exponent of background ~ 1/m^|α|
    qcd_background_strength: float = 1.0   # scale factor for hadronic background
    detector_calibrated: bool = False
    tracker_aligned: bool = False
    # Channel-dependent reconstruction efficiency
    channel_efficiency: Dict[str, float] = Field(
        default_factory=lambda: {
            "diphoton": 0.45,
            "dilepton_ee": 0.55,
            "dilepton_mumu": 0.70,
            "four_lepton": 0.40,
            "dijet": 0.80,
            "bb": 0.50,
        }
    )


# ── Experiment progress flags ────────────────────────────────────────────


class ExperimentProgress(BaseModel):
    """Boolean milestones used by rules and reward shaping."""

    beam_configured: bool = False
    luminosity_allocated: bool = False
    trigger_set: bool = False
    collisions_collected: bool = False
    detector_calibrated: bool = False
    tracks_reconstructed: bool = False
    channel_selected: bool = False
    invariant_mass_built: bool = False
    background_subtracted: bool = False
    resonance_fitted: bool = False
    bump_scanned: bool = False
    angular_measured: bool = False
    significance_estimated: bool = False
    systematics_requested: bool = False
    theory_review_requested: bool = False
    claim_submitted: bool = False

    n_events_collected: int = 0
    n_signal_candidates: int = 0
    n_background_estimate: int = 0
    best_fit_mass_gev: Optional[float] = None
    best_fit_width_gev: Optional[float] = None
    best_significance_sigma: Optional[float] = None
    best_channel: Optional[str] = None
    best_beam_energy: Optional[str] = None


# ── Resources ─────────────────────────────────────────────────────────────


class ResourceState(BaseModel):
    """Live resource accounting (superset of the agent-visible ResourceUsage)."""

    budget_total_musd: float = 100.0
    budget_used_musd: float = 0.0
    luminosity_total_fb: float = 300.0
    luminosity_used_fb: float = 0.0
    time_limit_days: float = 365.0
    time_used_days: float = 0.0
    compute_hours_used: float = 0.0

    @property
    def budget_remaining(self) -> float:
        return max(0.0, self.budget_total_musd - self.budget_used_musd)

    @property
    def luminosity_remaining(self) -> float:
        return max(0.0, self.luminosity_total_fb - self.luminosity_used_fb)

    @property
    def time_remaining(self) -> float:
        return max(0.0, self.time_limit_days - self.time_used_days)

    @property
    def budget_exhausted(self) -> bool:
        return self.budget_remaining <= 0

    @property
    def luminosity_exhausted(self) -> bool:
        return self.luminosity_remaining <= 0

    @property
    def time_exhausted(self) -> bool:
        return self.time_remaining <= 0


# ── Aggregate hidden state ───────────────────────────────────────────────


class FullLatentState(BaseModel):
    """Complete hidden state of the simulated LHC analysis world."""

    particle: LatentParticle = Field(default_factory=LatentParticle)
    detector: DetectorState = Field(default_factory=DetectorState)
    progress: ExperimentProgress = Field(default_factory=ExperimentProgress)
    resources: ResourceState = Field(default_factory=ResourceState)

    selected_channel: Optional[str] = None
    selected_beam_energy: Optional[str] = None
    selected_trigger: Optional[str] = None

    candidate_masses_gev: List[float] = Field(default_factory=list)
    candidate_significances: List[float] = Field(default_factory=list)

    hidden_failure_conditions: List[str] = Field(default_factory=list)
    rng_seed: int = 42
    step_count: int = 0