SystemSimulationApp/app/simulation/examples/testmodel/run.py

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path

from app.simulation.examples.testmodel.closure import TestModelSolveDiagnostics
from app.simulation.examples.testmodel.system import (
    InitializationDiagnostics,
    TestModelConfig,
    TestModelSystem,
)
from app.simulation.paths import (
    MODELICA_TESTMODEL_RESULT_PATH,
    PROJECT_ROOT,
    SIMULATION_RUNS_DIR,
)
from app.simulation.reporting import (
    COMPARISON_KEYS,
    PRIMARY_KEYS,
    TestModelArtifacts,
    export_testmodel_artifacts,
    format_testmodel_run_report,
    load_modelica_series,
    write_testmodel_run_report,
)
from app.simulation.solvers.solver import SolveIVPConfig


@dataclass(frozen=True)
class TestModelSamplingConfig:
    step: float = 0.1


@dataclass(frozen=True)
class TestModelPathConfig:
    output_dir: Path | None = None
    modelica_result_path: Path | None = None


@dataclass(frozen=True)
class TestModelExecutionConfig:
    use_modelica_reference_if_available: bool = True


@dataclass(frozen=True)
class TestModelRunConfig:
    model: TestModelConfig = field(default_factory=TestModelConfig)
    solver: SolveIVPConfig = field(default_factory=SolveIVPConfig)
    sampling: TestModelSamplingConfig = field(default_factory=TestModelSamplingConfig)
    paths: TestModelPathConfig = field(default_factory=TestModelPathConfig)
    execution: TestModelExecutionConfig = field(default_factory=TestModelExecutionConfig)

    @property
    def sample_step(self) -> float:
        return self.sampling.step

    def sample_times(self) -> list[float]:
        return _sample_times(
            self.solver.t_start,
            self.solver.t_stop,
            step=self.sampling.step,
        )


@dataclass(frozen=True)
class PreparedTestModelRun:
    run_config: TestModelRunConfig
    repo_root: Path
    output_dir: Path
    modelica_result_path: Path
    t_eval: tuple[float, ...]
    use_modelica_reference_if_available: bool
    modelica_reference_exists: bool


@dataclass(frozen=True)
class TestModelRunResult:
    run_config: TestModelRunConfig
    prepared_run: PreparedTestModelRun
    system: TestModelSystem
    initialization: InitializationDiagnostics
    raw_initial_state: tuple[float, ...]
    consistent_initial_state: tuple[float, ...]
    solution: object
    series: dict[str, list[float]]
    solve_diagnostics: TestModelSolveDiagnostics | None
    artifacts: TestModelArtifacts
    comparison_summary: dict[str, tuple[float, float]] | None
    used_modelica_reference: bool


def _sample_times(t_start: float, t_stop: float, step: float) -> list[float]:
    point_count = int(round((t_stop - t_start) / step))
    return [t_start + index * step for index in range(point_count + 1)]


def _default_run_output_dir() -> Path:
    timestamp = datetime.now(UTC).strftime("testmodel_%Y%m%d_%H%M%S_%f")
    return SIMULATION_RUNS_DIR / timestamp


def prepare_testmodel_run(
    *,
    run_config: TestModelRunConfig | None = None,
    output_dir: Path | None = None,
    modelica_result_path: Path | None = None,
) -> PreparedTestModelRun:
    run_config = run_config or TestModelRunConfig()
    resolved_output_dir = (
        output_dir
        or run_config.paths.output_dir
        or _default_run_output_dir()
    )
    resolved_modelica_result_path = (
        modelica_result_path
        or run_config.paths.modelica_result_path
        or MODELICA_TESTMODEL_RESULT_PATH
    )
    t_eval = tuple(run_config.sample_times())
    return PreparedTestModelRun(
        run_config=run_config,
        repo_root=PROJECT_ROOT,
        output_dir=resolved_output_dir,
        modelica_result_path=resolved_modelica_result_path,
        t_eval=t_eval,
        use_modelica_reference_if_available=run_config.execution.use_modelica_reference_if_available,
        modelica_reference_exists=resolved_modelica_result_path.exists(),
    )


def run_prepared_testmodel(prepared_run: PreparedTestModelRun) -> TestModelRunResult:
    run_config = prepared_run.run_config
    system = TestModelSystem(config=run_config.model)
    raw_initial_state = tuple(system.initial_state_vector())
    initialization = system.initialize_consistent_state()
    consistent_initial_state = tuple(initialization.state_vector)
    solution = system.simulate(config=run_config.solver, t_eval=list(prepared_run.t_eval))
    series = system.evaluate_solution(solution)
    solve_diagnostics = system.last_solve_diagnostics

    modelica_series = None
    used_modelica_reference = False
    if (
        prepared_run.use_modelica_reference_if_available
        and prepared_run.modelica_reference_exists
    ):
        modelica_series = load_modelica_series(
            prepared_run.modelica_result_path,
            COMPARISON_KEYS,
        )
        used_modelica_reference = True

    artifacts, comparison_summary = export_testmodel_artifacts(
        output_dir=prepared_run.output_dir,
        series=series,
        modelica_series=modelica_series,
    )
    report_text = format_testmodel_run_report(
        network_summary=system.network.summary(),
        initialization=initialization,
        raw_initial_state=raw_initial_state,
        consistent_initial_state=consistent_initial_state,
        solution=solution,
        series=series,
        solve_diagnostics=solve_diagnostics,
        artifacts=artifacts,
        comparison_summary=comparison_summary,
    )
    write_testmodel_run_report(prepared_run.output_dir, report_text)

    return TestModelRunResult(
        run_config=run_config,
        prepared_run=prepared_run,
        system=system,
        initialization=initialization,
        raw_initial_state=raw_initial_state,
        consistent_initial_state=consistent_initial_state,
        solution=solution,
        series=series,
        solve_diagnostics=solve_diagnostics,
        artifacts=artifacts,
        comparison_summary=comparison_summary,
        used_modelica_reference=used_modelica_reference,
    )


def run_testmodel(
    *,
    run_config: TestModelRunConfig | None = None,
    output_dir: Path | None = None,
    modelica_result_path: Path | None = None,
) -> TestModelRunResult:
    prepared_run = prepare_testmodel_run(
        run_config=run_config,
        output_dir=output_dir,
        modelica_result_path=modelica_result_path,
    )
    return run_prepared_testmodel(prepared_run)


def main() -> None:
    run_config = TestModelRunConfig()
    result = run_testmodel(run_config=run_config)
    print(
        format_testmodel_run_report(
            network_summary=result.system.network.summary(),
            initialization=result.initialization,
            raw_initial_state=result.raw_initial_state,
            consistent_initial_state=result.consistent_initial_state,
            solution=result.solution,
            series=result.series,
            solve_diagnostics=result.solve_diagnostics,
            artifacts=result.artifacts,
            comparison_summary=result.comparison_summary,
        ),
        end="",
    )


if __name__ == "__main__":
    main()