config.py

# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2025 SWGY, Inc
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
Configuration loading for the blast-pressure simulation.

The YAML file is the primary interface; individual keys can be overridden on
the command line with repeatable dotted ``--set key.sub=value`` flags. Values
are parsed with the YAML scalar rules, so ``--set blast.cube_segments=12``
yields an int and ``--set blast.cube_center_m=[0,0,0]`` yields a list.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, List, Optional, Tuple

import yaml


@dataclass(frozen=True)
class BlastCfg:
    cube_center_m: Tuple[float, float, float]
    cube_extents_m: float
    cube_segments: int
    min_dist_m: float


@dataclass(frozen=True)
class RaysCfg:
    resolution_arcmin: float
    max_bounces: int
    max_length_m: float
    fan_mode: str  # "sensor_cone" | "geometry"


@dataclass(frozen=True)
class PhysicsCfg:
    initial_peak_pressure_kpa: float
    positive_phase_duration_s: float
    decay_alpha: float
    reflection_loss: float


@dataclass(frozen=True)
class SensorCfg:
    diameter_m: float

    @property
    def radius_m(self) -> float:
        """True sphere radius (half the diameter)."""
        return self.diameter_m / 2.0


@dataclass(frozen=True)
class GeometryCfg:
    helmet: str
    vest: str


@dataclass(frozen=True)
class ComputeCfg:
    device: str           # "cuda" | "cpu" (Warp)


@dataclass(frozen=True)
class OutputCfg:
    file: Optional[str]


@dataclass(frozen=True)
class DiffractionCfg:
    enabled: bool                 # add first-order BTM edge diffraction
    edge_mode: str                # "rim_plane" | "feature_edges" | "polyline"
    rim_plane_normal: Tuple[float, float, float]
    rim_plane_offset_m: float
    feature_angle_deg: float
    segment_length_m: float
    wedge_model: str              # "screen" (nu=1/2) | "dihedral" (nu from faces)
    polyline_file: Optional[str]


@dataclass(frozen=True)
class Config:
    blast: BlastCfg
    rays: RaysCfg
    physics: PhysicsCfg
    sensor: SensorCfg
    geometry: GeometryCfg
    compute: ComputeCfg
    output: OutputCfg
    diffraction: DiffractionCfg


_FAN_MODES = ("sensor_cone", "geometry")
_EDGE_MODES = ("rim_plane", "feature_edges", "polyline")
_WEDGE_MODELS = ("screen", "dihedral")


def _apply_override(tree: dict, dotted: str) -> None:
    """Apply a single ``key.sub=value`` override to a nested dict in place."""
    if "=" not in dotted:
        raise ValueError(f"--set expects key.path=value, got: {dotted!r}")
    path, raw = dotted.split("=", 1)
    keys = path.strip().split(".")
    value = yaml.safe_load(raw)  # YAML scalar parsing: ints, floats, bools, lists

    node = tree
    for k in keys[:-1]:
        if k not in node or not isinstance(node[k], dict):
            raise KeyError(f"--set path not found in config: {path!r}")
        node = node[k]
    if keys[-1] not in node:
        raise KeyError(f"--set path not found in config: {path!r}")
    node[keys[-1]] = value


def _require(cond: bool, msg: str) -> None:
    if not cond:
        raise ValueError(f"Invalid configuration: {msg}")


def _build(tree: dict) -> Config:
    b = tree["blast"]
    r = tree["rays"]
    p = tree["physics"]
    s = tree["sensor"]
    g = tree["geometry"]
    c = tree["compute"]
    o = tree["output"]
    d = tree.get("diffraction", {})   # optional; absent -> disabled defaults

    center = tuple(float(v) for v in b["cube_center_m"])
    _require(len(center) == 3, "blast.cube_center_m must have three components")
    _require(b["cube_segments"] > 1, "blast.cube_segments must be > 1")
    _require(b["cube_extents_m"] > 0, "blast.cube_extents_m must be > 0")
    _require(b["min_dist_m"] >= 0, "blast.min_dist_m must be >= 0")
    _require(r["resolution_arcmin"] > 0, "rays.resolution_arcmin must be > 0")
    _require(r["max_bounces"] >= 0, "rays.max_bounces must be >= 0")
    _require(r["max_length_m"] > 0, "rays.max_length_m must be > 0")
    _require(r["fan_mode"] in _FAN_MODES, f"rays.fan_mode must be one of {_FAN_MODES}")
    _require(p["positive_phase_duration_s"] > 0, "physics.positive_phase_duration_s must be > 0")
    _require(s["diameter_m"] > 0, "sensor.diameter_m must be > 0")
    _require(c["device"] in ("cuda", "cpu"), "compute.device must be 'cuda' or 'cpu'")

    edge_mode = str(d.get("edge_mode", "rim_plane"))
    wedge_model = str(d.get("wedge_model", "screen"))
    seg_len = float(d.get("segment_length_m", 0.01))
    feat_angle = float(d.get("feature_angle_deg", 30.0))
    normal = tuple(float(v) for v in d.get("rim_plane_normal", (0.0, 0.0, 1.0)))
    poly = d.get("polyline_file", None)
    _require(edge_mode in _EDGE_MODES, f"diffraction.edge_mode must be one of {_EDGE_MODES}")
    _require(wedge_model in _WEDGE_MODELS, f"diffraction.wedge_model must be one of {_WEDGE_MODELS}")
    _require(seg_len > 0, "diffraction.segment_length_m must be > 0")
    _require(0.0 < feat_angle < 180.0, "diffraction.feature_angle_deg must be in (0, 180)")
    _require(len(normal) == 3, "diffraction.rim_plane_normal must have three components")
    _require(edge_mode != "polyline" or poly,
             "diffraction.polyline_file is required when edge_mode is 'polyline'")

    return Config(
        blast=BlastCfg(center, float(b["cube_extents_m"]), int(b["cube_segments"]),
                       float(b["min_dist_m"])),
        rays=RaysCfg(float(r["resolution_arcmin"]), int(r["max_bounces"]),
                     float(r["max_length_m"]), str(r["fan_mode"])),
        physics=PhysicsCfg(float(p["initial_peak_pressure_kpa"]),
                           float(p["positive_phase_duration_s"]),
                           float(p["decay_alpha"]), float(p["reflection_loss"])),
        sensor=SensorCfg(float(s["diameter_m"])),
        geometry=GeometryCfg(str(g["helmet"]), str(g["vest"])),
        compute=ComputeCfg(str(c["device"])),
        output=OutputCfg(o["file"] if o["file"] else None),
        diffraction=DiffractionCfg(
            bool(d.get("enabled", False)), edge_mode, normal,
            float(d.get("rim_plane_offset_m", -0.09)), feat_angle, seg_len,
            wedge_model, str(poly) if poly else None),
    )


def load_config(path: str, overrides: Optional[List[str]] = None) -> Config:
    """Load a YAML config file and apply any ``--set`` overrides."""
    with open(path, "r") as fh:
        tree: dict[str, Any] = yaml.safe_load(fh)

    for dotted in overrides or []:
        _apply_override(tree, dotted)

    return _build(tree)