# SPDX-FileCopyrightText: 2026 Jukka Aho
# SPDX-License-Identifier: AGPL-3.0-or-later
"""Check the unmodified OpenPFC diffusion starter against its discrete mode.

Usage: python3 check-results.py input.json ATTEMPT/results [OTHER_ATTEMPT/results]
Only the Python standard library is required. This checks values and ownership,
not performance or arbitrary edited models. Raises an error on a failed check.
"""

import argparse
import csv
import json
import math
from pathlib import Path

TOLERANCE = 2e-11


def require(condition, message):
    if not condition:
        raise ValueError(message)


def close(actual, expected, label):
    require(math.isfinite(actual), f"{label}: non-finite value")
    require(abs(actual - expected) <= TOLERANCE, f"{label}: {actual} != {expected}")


def check(input_path, results):
    settings = json.loads(Path(input_path).read_text())
    cells, steps = settings["cells"], settings["steps"]
    dt, diffusivity = settings["dt"], settings["diffusivity"]
    require(type(cells) is int and 4 <= cells <= 256, "invalid cells")
    require(type(steps) is int and 1 <= steps <= 1000000, "invalid steps")
    require(
        math.isfinite(dt)
        and math.isfinite(diffusivity)
        and dt > 0
        and diffusivity > 0
        and dt * diffusivity <= 1 / 6,
        "invalid time step/diffusivity",
    )
    gain = 1 - 4 * diffusivity * dt * math.sin(math.pi / cells) ** 2
    amplitude = 0.1 * gain**steps
    summary = json.loads((results / "summary.json").read_text())
    ranks = summary["ranks"]
    require(type(ranks) is int and ranks >= 1, "invalid ranks")
    require(summary["steps"] == steps, "wrong step count")
    close(summary["time"], steps * dt, "final time")
    close(summary["mean"], 1, "summary mean")
    close(summary["initial_l2"], math.sqrt(1.005), "initial RMS")
    close(summary["final_l2"], math.sqrt(1 + amplitude**2 / 2), "final RMS")
    with (results / "diagnostics.csv").open(newline="") as stream:
        rows = list(csv.DictReader(stream))
    require(len(rows) == steps + 1, "missing/extra diagnostics rows")
    for step, row in enumerate(rows):
        require(int(row["step"]) == step, "wrong diagnostics step")
        close(float(row["time"]), step * dt, "diagnostics time")
        close(float(row["mean"]), 1, "diagnostics mean")
        close(
            float(row["l2"]),
            math.sqrt(1 + (0.1 * gain**step) ** 2 / 2),
            "diagnostics RMS",
        )
    files = list(results.glob("field_rank_*.csv"))
    require(
        {file.name for file in files}
        == {f"field_rank_{rank}.csv" for rank in range(ranks)},
        "missing/extra rank files",
    )
    field = {}
    for file in files:
        with file.open(newline="") as stream:
            for row in csv.DictReader(stream):
                point = tuple(float(row[axis]) for axis in ["x", "y", "z"])
                require(
                    all(
                        math.isfinite(x) and x.is_integer() and 0 <= x < cells
                        for x in point
                    ),
                    f"invalid owned coordinate: {point}",
                )
                require(point not in field, f"duplicate owned cell: {point}")
                value = float(row["u"])
                expected = 1 + amplitude * math.cos(2 * math.pi * point[0] / cells)
                close(value, expected, f"field at {point}")
                field[point] = value
    require(len(field) == cells**3, "incomplete global owned field")
    close(math.fsum(field.values()) / cells**3, 1, "field mean")
    measured_amplitude = (
        2
        * math.fsum(
            (value - 1) * math.cos(2 * math.pi * point[0] / cells)
            for point, value in field.items()
        )
        / cells**3
    )
    close(measured_amplitude, amplitude, "mode amplitude")
    max_error = max(
        abs(value - (1 + amplitude * math.cos(2 * math.pi * point[0] / cells)))
        for point, value in field.items()
    )
    report = {
        "ranks": ranks,
        "owned_cells": len(field),
        "mean": math.fsum(field.values()) / cells**3,
        "measured_amplitude": measured_amplitude,
        "discrete_amplitude": amplitude,
        "pde_amplitude": 0.1
        * math.exp(-diffusivity * (2 * math.pi / cells) ** 2 * dt * steps),
        "max_discrete_field_error": max_error,
    }
    return field, report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path)
    parser.add_argument(
        "results", type=Path, nargs="+", help="One or two results directories"
    )
    args = parser.parse_args()
    require(len(args.results) <= 2, "supply one or two results directories")
    checked = [check(args.input, directory) for directory in args.results]
    report = {"tolerance": TOLERANCE, "runs": [result[1] for result in checked]}
    if len(checked) == 2:
        left, right = checked[0][0], checked[1][0]
        require(left.keys() == right.keys(), "global ownership differs")
        error = max(abs(left[point] - right[point]) for point in left)
        require(error <= TOLERANCE, "global fields differ across rank counts")
        report["max_rank_count_difference"] = error
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()