server/tasks/scenarios.py

"""Built-in physics scenarios + procedural sampling.

Each scenario binds a hidden ``LatentParticle`` truth and a public
``TaskSpec`` (search window, available channels, resource budgets, expected
findings, paper references). Curated scenarios are inspired by famous LHC
discoveries; procedural ones randomise mass, channel, width and budgets to
build a curriculum.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional

import numpy as np

from models import (
    DetectorChannel,
    ExpectedFinding,
    PaperReference,
    TOOL_REGISTRY,
    TaskSpec,
)

from server.simulator.latent_state import (
    DetectorState,
    FullLatentState,
    LatentParticle,
    ResourceState,
)


@dataclass
class Scenario:
    name: str
    difficulty: str
    task: TaskSpec
    latent: FullLatentState

    def fresh_latent(self) -> FullLatentState:
        # Pydantic deep-copy so the env can mutate freely
        return self.latent.model_copy(deep=True)


# ── Curated, story-driven scenarios ──────────────────────────────────────


def _higgs_like_scenario() -> Scenario:
    particle = LatentParticle(
        name="HiggsLike",
        mass_gev=125.0,
        width_gev=0.004,
        spin=0,
        parity="+",
        cross_section_fb=55.0,
        decay_branching={
            "diphoton": 0.0023,
            "dilepton_ee": 0.00003,
            "dilepton_mumu": 0.00022,
            "four_lepton": 0.000125,
            "bb": 0.58,
            "dijet": 0.30,
        },
        primary_channel="diphoton",
    )
    detector = DetectorState(
        detector_resolution_gev=1.5,
        pileup_mu=30.0,
        trigger_efficiency=0.85,
    )
    resources = ResourceState(
        budget_total_musd=120.0,
        luminosity_total_fb=300.0,
        time_limit_days=365.0,
    )
    latent = FullLatentState(
        particle=particle,
        detector=detector,
        resources=resources,
        rng_seed=125,
    )
    task = TaskSpec(
        problem_statement=(
            "An anomalous excess at ~125 GeV is rumoured in early 13 TeV runs. "
            "Plan a campaign to confirm or refute a Standard-Model Higgs-like scalar. "
            "Pick channels, allocate luminosity, fit, and submit a calibrated discovery claim."
        ),
        target_collider="LHC",
        mass_search_window_gev=[100.0, 200.0],
        budget_limit_musd=120.0,
        luminosity_budget_fb=300.0,
        time_limit_days=365.0,
        prior_observations=[
            "Earlier Tevatron data shows a mild diphoton excess near 125 GeV.",
            "ATLAS/CMS rumour mills suggest a 4ℓ excess at low mass.",
        ],
        success_criteria=[
            "Identify a resonance within 1 GeV of the truth.",
            "Reach ≥5σ local significance.",
            "Submit confidence consistent with calibration.",
        ],
        paper_references=[
            PaperReference(
                title="Observation of a new particle in the search for the SM Higgs boson",
                arxiv_id="1207.7214",
                doi="10.1016/j.physletb.2012.08.020",
            ),
        ],
        expected_findings=[
            ExpectedFinding(finding="Diphoton resonance at ~125 GeV", category="discovery"),
            ExpectedFinding(finding="Spin-0, even parity", category="property"),
        ],
        difficulty="medium",
        available_tools=list(TOOL_REGISTRY.keys()),
    )
    return Scenario(name="higgs_like_125", difficulty="medium", task=task, latent=latent)


def _hidden_zprime_scenario() -> Scenario:
    particle = LatentParticle(
        name="ZPrime",
        mass_gev=600.0,
        width_gev=18.0,
        spin=1,
        parity="-",
        cross_section_fb=12.0,
        decay_branching={
            "diphoton": 0.0,
            "dilepton_ee": 0.04,
            "dilepton_mumu": 0.04,
            "four_lepton": 0.0,
            "bb": 0.20,
            "dijet": 0.70,
        },
        primary_channel="dilepton_mumu",
    )
    detector = DetectorState(
        detector_resolution_gev=8.0,
        pileup_mu=45.0,
        trigger_efficiency=0.78,
        qcd_background_strength=1.2,
    )
    resources = ResourceState(
        budget_total_musd=140.0,
        luminosity_total_fb=200.0,
        time_limit_days=400.0,
    )
    latent = FullLatentState(
        particle=particle, detector=detector, resources=resources, rng_seed=600,
    )
    task = TaskSpec(
        problem_statement=(
            "Run-2 dilepton spectra hint at a high-mass excess. Hunt for a heavy "
            "Z'-like vector resonance and characterise spin-1, parity-odd hypothesis."
        ),
        mass_search_window_gev=[300.0, 1500.0],
        budget_limit_musd=140.0,
        luminosity_budget_fb=200.0,
        time_limit_days=400.0,
        prior_observations=[
            "High-pT dilepton tail shows a 2.7σ shoulder near 600 GeV.",
            "Dijet smooth-fit residuals consistent with the same window.",
        ],
        success_criteria=[
            "Identify a high-mass dilepton/dijet resonance.",
            "Constrain spin to be vector (1).",
            "Report calibrated mass within 5% and ≥4σ significance.",
        ],
        paper_references=[
            PaperReference(
                title="Search for high-mass dilepton resonances at the LHC",
                arxiv_id="1903.06248",
            ),
        ],
        expected_findings=[
            ExpectedFinding(finding="Heavy Z'-like dilepton resonance", category="discovery"),
            ExpectedFinding(finding="Spin-1, parity-odd", category="property"),
        ],
        difficulty="hard",
        available_tools=list(TOOL_REGISTRY.keys()),
    )
    return Scenario(name="hidden_zprime_600", difficulty="hard", task=task, latent=latent)


def _diboson_resonance_scenario() -> Scenario:
    particle = LatentParticle(
        name="Graviton",
        mass_gev=750.0,
        width_gev=45.0,
        spin=2,
        parity="+",
        cross_section_fb=6.0,
        decay_branching={
            "diphoton": 0.06,
            "dilepton_ee": 0.005,
            "dilepton_mumu": 0.005,
            "four_lepton": 0.001,
            "bb": 0.15,
            "dijet": 0.70,
        },
        primary_channel="diphoton",
    )
    detector = DetectorState(
        detector_resolution_gev=12.0,
        pileup_mu=50.0,
        trigger_efficiency=0.80,
    )
    resources = ResourceState(
        budget_total_musd=110.0,
        luminosity_total_fb=180.0,
        time_limit_days=350.0,
    )
    latent = FullLatentState(
        particle=particle, detector=detector, resources=resources, rng_seed=750,
    )
    task = TaskSpec(
        problem_statement=(
            "A faint γγ excess at 750 GeV stirred the field briefly in 2015-2016. "
            "Re-investigate with the modern luminosity budget and decide if it is "
            "real or a fluctuation."
        ),
        mass_search_window_gev=[400.0, 1200.0],
        budget_limit_musd=110.0,
        luminosity_budget_fb=180.0,
        time_limit_days=350.0,
        prior_observations=[
            "Public CMS/ATLAS data show a 2-3σ diphoton bump near 750 GeV.",
            "Theory papers proposed graviton, scalar singlet, and SM-fluctuation explanations.",
        ],
        success_criteria=[
            "Decide between discovery and fluctuation with calibrated confidence.",
        ],
        paper_references=[
            PaperReference(
                title="Search for resonant production of high-mass diphoton pairs",
                arxiv_id="1606.04093",
            ),
        ],
        expected_findings=[
            ExpectedFinding(finding="Possible diphoton resonance near 750 GeV", category="discovery"),
        ],
        difficulty="hard",
        available_tools=list(TOOL_REGISTRY.keys()),
    )
    return Scenario(name="diphoton_750", difficulty="hard", task=task, latent=latent)


def _easy_diphoton_scenario() -> Scenario:
    """Generous budgets, narrow scalar, single obvious channel."""
    particle = LatentParticle(
        name="EasyScalar",
        mass_gev=160.0,
        width_gev=0.5,
        spin=0,
        parity="+",
        cross_section_fb=120.0,
        decay_branching={
            "diphoton": 0.05,
            "dilepton_ee": 0.001,
            "dilepton_mumu": 0.005,
            "four_lepton": 0.0001,
            "bb": 0.50,
            "dijet": 0.30,
        },
        primary_channel="diphoton",
    )
    detector = DetectorState(
        detector_resolution_gev=2.0,
        pileup_mu=20.0,
        trigger_efficiency=0.9,
    )
    resources = ResourceState(
        budget_total_musd=200.0,
        luminosity_total_fb=400.0,
        time_limit_days=500.0,
    )
    latent = FullLatentState(
        particle=particle, detector=detector, resources=resources, rng_seed=160,
    )
    task = TaskSpec(
        problem_statement=(
            "Tutorial scenario: discover a narrow scalar that decays cleanly to "
            "two photons. Resources are abundant; focus on running a clean pipeline."
        ),
        mass_search_window_gev=[80.0, 300.0],
        budget_limit_musd=200.0,
        luminosity_budget_fb=400.0,
        time_limit_days=500.0,
        success_criteria=[
            "Identify the diphoton peak and submit a calibrated 5σ claim.",
        ],
        expected_findings=[
            ExpectedFinding(finding="Diphoton scalar near 160 GeV", category="discovery"),
        ],
        difficulty="easy",
        available_tools=list(TOOL_REGISTRY.keys()),
    )
    return Scenario(name="easy_diphoton_160", difficulty="easy", task=task, latent=latent)


CURATED_SCENARIOS: List[Scenario] = [
    _easy_diphoton_scenario(),
    _higgs_like_scenario(),
    _hidden_zprime_scenario(),
    _diboson_resonance_scenario(),
]


# ── Procedural sampler ───────────────────────────────────────────────────


_DIFFICULTY_TIERS = {
    "easy":   {"mass_lo": 90.0,  "mass_hi": 250.0,  "xsec_lo": 80.0, "xsec_hi": 150.0, "res": 1.5,  "budget": 200.0, "lumi": 400.0},
    "medium": {"mass_lo": 100.0, "mass_hi": 600.0,  "xsec_lo": 25.0, "xsec_hi": 80.0,  "res": 3.0,  "budget": 150.0, "lumi": 300.0},
    "hard":   {"mass_lo": 250.0, "mass_hi": 1500.0, "xsec_lo": 5.0,  "xsec_hi": 25.0,  "res": 8.0,  "budget": 110.0, "lumi": 200.0},
}


def _procedural_scenario(difficulty: str, rng: np.random.Generator) -> Scenario:
    tier = _DIFFICULTY_TIERS.get(difficulty, _DIFFICULTY_TIERS["medium"])
    mass = float(rng.uniform(tier["mass_lo"], tier["mass_hi"]))
    xsec = float(rng.uniform(tier["xsec_lo"], tier["xsec_hi"]))
    spin = int(rng.choice([0, 1, 2]))
    parity = str(rng.choice(["+", "-"]))
    primary = str(rng.choice([c.value for c in DetectorChannel]))

    branching = {c.value: 0.001 for c in DetectorChannel}
    branching[primary] = float(rng.uniform(0.02, 0.6))
    # normalise so it sums to ~1
    total = sum(branching.values())
    branching = {k: v / total for k, v in branching.items()}

    particle = LatentParticle(
        name=f"Mystery_{int(mass)}GeV",
        mass_gev=mass,
        width_gev=float(rng.uniform(0.5, 30.0) if difficulty != "easy" else rng.uniform(0.05, 2.0)),
        spin=spin,
        parity=parity,
        cross_section_fb=xsec,
        decay_branching=branching,
        primary_channel=primary,
    )
    detector = DetectorState(
        detector_resolution_gev=tier["res"],
        pileup_mu=float(rng.uniform(20.0, 60.0)),
        trigger_efficiency=float(rng.uniform(0.7, 0.92)),
        qcd_background_strength=float(rng.uniform(0.8, 1.3)),
    )
    resources = ResourceState(
        budget_total_musd=tier["budget"],
        luminosity_total_fb=tier["lumi"],
        time_limit_days=float(rng.uniform(300.0, 500.0)),
    )
    latent = FullLatentState(
        particle=particle, detector=detector, resources=resources,
        rng_seed=int(rng.integers(1, 1_000_000)),
    )
    window_lo = max(50.0, mass - 200.0)
    window_hi = mass + 300.0
    task = TaskSpec(
        problem_statement=(
            f"Procedural ({difficulty}): a hidden resonance lives somewhere in "
            f"[{window_lo:.0f}, {window_hi:.0f}] GeV. Discover and characterise it."
        ),
        mass_search_window_gev=[window_lo, window_hi],
        budget_limit_musd=tier["budget"],
        luminosity_budget_fb=tier["lumi"],
        time_limit_days=resources.time_limit_days,
        difficulty=difficulty,
        available_tools=list(TOOL_REGISTRY.keys()),
        success_criteria=[
            "Discover the hidden resonance with a calibrated mass and channel.",
        ],
    )
    return Scenario(
        name=f"procedural_{difficulty}_{int(mass)}",
        difficulty=difficulty,
        task=task,
        latent=latent,
    )


def sample_scenario(
    *,
    difficulty: Optional[str] = None,
    name: Optional[str] = None,
    seed: Optional[int] = None,
) -> Scenario:
    rng = np.random.default_rng(seed)

    if name:
        for s in CURATED_SCENARIOS:
            if s.name == name:
                fresh = Scenario(
                    name=s.name,
                    difficulty=s.difficulty,
                    task=s.task,
                    latent=s.fresh_latent(),
                )
                if seed is not None:
                    fresh.latent.rng_seed = int(seed)
                return fresh

    if difficulty in {"easy", "medium", "hard"}:
        # mix curated + procedural
        curated_pool = [s for s in CURATED_SCENARIOS if s.difficulty == difficulty]
        if curated_pool and rng.random() < 0.4:
            picked = curated_pool[int(rng.integers(0, len(curated_pool)))]
            return Scenario(
                name=picked.name,
                difficulty=picked.difficulty,
                task=picked.task,
                latent=picked.fresh_latent(),
            )
        return _procedural_scenario(difficulty, rng)

    # default: random difficulty
    diff = str(rng.choice(["easy", "medium", "hard"]))
    return _procedural_scenario(diff, rng)


__all__ = ["CURATED_SCENARIOS", "Scenario", "sample_scenario"]