SystemSimulationApp/app/simulation/reporting/test_mql_comparison.py

from __future__ import annotations

from bisect import bisect_left
import csv
from dataclasses import dataclass
from pathlib import Path

from app.simulation.reporting.amesim_results import AmesimResults


DEFAULT_TEST_MQL_ALIGNMENT_PATHS = (
    "temp3@pn_c1_8",
    "press3@pn_c1_8",
    "vvol1@pn_brp2_8",
    "vol1@pn_brp2_8",
)


@dataclass(frozen=True)
class TestMqlComparisonMetric:
    data_path: str
    sample_count: int
    max_abs_error: float
    mean_abs_error: float
    max_rel_error: float
    undefined_rel_error_count: int
    near_zero_baseline_count: int
    final_abs_error: float


@dataclass(frozen=True)
class TestMqlComparisonResult:
    metrics: tuple[TestMqlComparisonMetric, ...]

    def metric(self, data_path: str) -> TestMqlComparisonMetric:
        for metric in self.metrics:
            if metric.data_path == data_path:
                return metric
        raise KeyError(data_path)

    @property
    def max_abs_error(self) -> float:
        return max((metric.max_abs_error for metric in self.metrics), default=0.0)

    @property
    def max_rel_error(self) -> float:
        return max((metric.max_rel_error for metric in self.metrics), default=0.0)

    @property
    def undefined_rel_error_count(self) -> int:
        return sum(metric.undefined_rel_error_count for metric in self.metrics)


class TestMqlComparisonError(ValueError):
    """Raised when Python and AMESim series cannot be aligned."""


def compare_test_mql_series(
    *,
    python_times: tuple[float, ...] | list[float],
    python_series_by_data_path: dict[str, tuple[float, ...] | list[float]],
    amesim_results: AmesimResults,
    data_paths: tuple[str, ...] | list[str] | None = None,
    relative_floor: float = 1.0e-12,
) -> TestMqlComparisonResult:
    """Compare current values directly with AMESim simulation values.

    ``relative_floor`` only identifies near-zero baselines for reporting.  It is
    never substituted into the relative-error denominator.  An exact zero
    AMESim baseline has undefined relative error and is counted separately;
    absolute error remains available for judgement.
    """

    if relative_floor < 0.0:
        raise TestMqlComparisonError("relative_floor cannot be negative.")
    _validate_time_axis(python_times)
    selected_paths = _select_data_paths(python_series_by_data_path, amesim_results, data_paths)
    metrics = []
    for data_path in selected_paths:
        python_values = tuple(float(value) for value in python_series_by_data_path[data_path])
        if len(python_values) != len(python_times):
            raise TestMqlComparisonError(
                f"Python series length mismatch for {data_path!r}: "
                f"{len(python_values)} values for {len(python_times)} time samples."
            )
        amesim_values = amesim_results.series(data_path)
        abs_errors = []
        rel_errors = []
        undefined_rel_error_count = 0
        near_zero_baseline_count = 0
        for time_value, python_value in zip(python_times, python_values):
            amesim_value = interpolate_series_value(amesim_results.times, amesim_values, time_value)
            abs_error = abs(python_value - amesim_value)
            abs_errors.append(abs_error)
            if abs(amesim_value) <= relative_floor:
                near_zero_baseline_count += 1
            if amesim_value == 0.0:
                undefined_rel_error_count += 1
            else:
                rel_errors.append(abs_error / abs(amesim_value))
        final_amesim_value = interpolate_series_value(
            amesim_results.times,
            amesim_values,
            float(python_times[-1]),
        )
        metrics.append(
            TestMqlComparisonMetric(
                data_path=data_path,
                sample_count=len(python_times),
                max_abs_error=max(abs_errors, default=0.0),
                mean_abs_error=sum(abs_errors) / max(len(abs_errors), 1),
                max_rel_error=max(rel_errors, default=0.0),
                undefined_rel_error_count=undefined_rel_error_count,
                near_zero_baseline_count=near_zero_baseline_count,
                final_abs_error=abs(python_values[-1] - final_amesim_value),
            )
        )
    return TestMqlComparisonResult(metrics=tuple(metrics))


def write_test_mql_amesim_baseline_csv(
    output_dir: Path,
    amesim_results: AmesimResults,
    data_paths: tuple[str, ...] | list[str] = DEFAULT_TEST_MQL_ALIGNMENT_PATHS,
) -> Path:
    output_dir.mkdir(parents=True, exist_ok=True)
    csv_path = output_dir / "test_mql_amesim_baseline.csv"
    _validate_amesim_data_paths(amesim_results, data_paths)
    with csv_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["time_s", *data_paths])
        for index, time_value in enumerate(amesim_results.times):
            writer.writerow(
                [time_value, *(amesim_results.series(data_path)[index] for data_path in data_paths)]
            )
    return csv_path


def write_test_mql_comparison_csv(
    *,
    output_dir: Path,
    python_times: tuple[float, ...] | list[float],
    python_series_by_data_path: dict[str, tuple[float, ...] | list[float]],
    amesim_results: AmesimResults,
    data_paths: tuple[str, ...] | list[str] | None = None,
) -> tuple[Path, Path, TestMqlComparisonResult]:
    output_dir.mkdir(parents=True, exist_ok=True)
    selected_paths = _select_data_paths(python_series_by_data_path, amesim_results, data_paths)
    comparison = compare_test_mql_series(
        python_times=python_times,
        python_series_by_data_path=python_series_by_data_path,
        amesim_results=amesim_results,
        data_paths=selected_paths,
    )
    csv_path = output_dir / "test_mql_amesim_comparison.csv"
    summary_path = output_dir / "test_mql_amesim_comparison_summary.txt"

    with csv_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        header = ["time_s"]
        for data_path in selected_paths:
            header.extend(
                [
                    f"python.{data_path}",
                    f"amesim.{data_path}",
                    f"abs_error.{data_path}",
                    f"rel_error.{data_path}",
                ]
            )
        writer.writerow(header)
        for index, time_value in enumerate(python_times):
            row = [time_value]
            for data_path in selected_paths:
                python_value = float(python_series_by_data_path[data_path][index])
                amesim_value = interpolate_series_value(
                    amesim_results.times,
                    amesim_results.series(data_path),
                    float(time_value),
                )
                abs_error = abs(python_value - amesim_value)
                rel_error = (
                    None
                    if amesim_value == 0.0
                    else abs_error / abs(amesim_value)
                )
                row.extend(
                    [
                        python_value,
                        amesim_value,
                        abs_error,
                        "" if rel_error is None else rel_error,
                    ]
                )
            writer.writerow(row)

    summary_lines = [
        (
            f"{metric.data_path}: samples={metric.sample_count}, "
            f"max_abs_error={metric.max_abs_error:.12g}, "
            f"mean_abs_error={metric.mean_abs_error:.12g}, "
            f"max_rel_error={metric.max_rel_error:.12%}, "
            f"undefined_rel_error_count={metric.undefined_rel_error_count}, "
            f"final_abs_error={metric.final_abs_error:.12g}"
        )
        for metric in comparison.metrics
    ]
    summary_path.write_text("\n".join(summary_lines) + "\n", encoding="utf-8")
    return csv_path, summary_path, comparison


def interpolate_series_value(
    time_values: tuple[float, ...] | list[float],
    values: tuple[float, ...] | list[float],
    target_time: float,
) -> float:
    if len(time_values) != len(values):
        raise TestMqlComparisonError("time and value series lengths differ.")
    if not time_values:
        raise TestMqlComparisonError("cannot interpolate an empty series.")
    if target_time <= time_values[0]:
        return float(values[0])
    if target_time >= time_values[-1]:
        return float(values[-1])

    right_index = bisect_left(time_values, target_time)
    if right_index < len(time_values) and abs(time_values[right_index] - target_time) <= 1.0e-12:
        return float(values[right_index])

    left_index = right_index - 1
    left_time = float(time_values[left_index])
    right_time = float(time_values[right_index])
    fraction = (target_time - left_time) / (right_time - left_time)
    return float(values[left_index]) + fraction * (float(values[right_index]) - float(values[left_index]))


def _select_data_paths(
    python_series_by_data_path: dict[str, tuple[float, ...] | list[float]],
    amesim_results: AmesimResults,
    data_paths: tuple[str, ...] | list[str] | None,
) -> tuple[str, ...]:
    if data_paths is None:
        data_paths = tuple(
            data_path
            for data_path in python_series_by_data_path
            if data_path in amesim_results.series_by_data_path
        )
    selected_paths = tuple(data_paths)
    if not selected_paths:
        raise TestMqlComparisonError("no common Data_Path values are available for comparison.")
    missing_python = [data_path for data_path in selected_paths if data_path not in python_series_by_data_path]
    if missing_python:
        raise TestMqlComparisonError(f"Python series missing Data_Path values: {missing_python}")
    _validate_amesim_data_paths(amesim_results, selected_paths)
    return selected_paths


def _validate_amesim_data_paths(
    amesim_results: AmesimResults,
    data_paths: tuple[str, ...] | list[str],
) -> None:
    missing_amesim = [data_path for data_path in data_paths if data_path not in amesim_results.series_by_data_path]
    if missing_amesim:
        raise TestMqlComparisonError(f"AMESim results missing Data_Path values: {missing_amesim}")


def _validate_time_axis(time_values: tuple[float, ...] | list[float]) -> None:
    if not time_values:
        raise TestMqlComparisonError("Python time axis is empty.")
    previous = float(time_values[0])
    for value in time_values[1:]:
        value = float(value)
        if value < previous:
            raise TestMqlComparisonError("Python time axis must be monotonically increasing.")
        previous = value