#!/usr/bin/env python3
"""Refine a saved STROID mesh once, retaining the original and provenance.

Run with a Python environment containing the multiblock-capable STROID build.
The output is a complete, already-refined mesh: pass it to the validation
experiment without any further refinement, including during saved-field replay.
"""

import argparse
import hashlib
import json
from pathlib import Path
import sys

import stroid
from stroid import _stroid
from stroid.IO import LoadStroidMesh, SaveStroidMesh
from stroid.refinement import UniformRefinement


def digest(path):
    with Path(path).open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def counts(mesh):
    result = stroid.stats.ComputeMeshStats(
        mesh, stroid.stats.MeshStatFeatures.ELEMENT_COUNT
    ).element_counts
    return {name: getattr(result, name) for name in
            ("total", "core", "envelope", "vacuum", "other")}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path)
    parser.add_argument("output_directory", type=Path,
                        help="New directory; existing directories are refused")
    args = parser.parse_args()
    source = args.input.resolve(strict=True)
    original_digest = digest(source)
    mesh = LoadStroidMesh(str(source))
    if not hasattr(mesh.config, "core_mapping"):
        raise RuntimeError("STROID is too old: no core_mapping support")
    if not mesh.has_mesh() or not mesh.has_rmesh():
        raise RuntimeError("Both physical and logical reference meshes are required")
    before = counts(mesh)
    initial_level = mesh.refinement_levels
    mapping = mesh.config.core_mapping
    order = mesh.config.order
    args.output_directory.mkdir(exist_ok=False)
    print(f"Loaded {mesh}; mapping={mapping}, geometry order={order}, "
          f"stored refinement level={initial_level}", flush=True)

    # This is an ADDITIONAL level on the loaded mesh, not regeneration from
    # MeshConfig.refinement_levels. STROID rebuilds the high-order geometry.
    UniformRefinement(mesh, 1)

    after = counts(mesh)
    if mesh.refinement_levels != initial_level + 1:
        raise RuntimeError("STROID did not advance exactly one refinement level")
    if any(after[name] != 8 * count for name, count in before.items()):
        raise RuntimeError(f"Expected eight children per hex: {before} -> {after}")
    if mesh.config.core_mapping != mapping or mesh.config.order != order:
        raise RuntimeError("Refinement changed mapping strategy or geometry order")
    target = args.output_directory / "refined.smesh"
    SaveStroidMesh(mesh, str(target), "One additional UniformRefinement of " + str(source))
    reloaded = LoadStroidMesh(str(target))
    if counts(reloaded) != after or reloaded.refinement_levels != initial_level + 1:
        raise RuntimeError("Saved refinement failed its load/metadata round trip")
    if reloaded.config.core_mapping != mapping or reloaded.config.order != order:
        raise RuntimeError("Saved refinement lost its mapping/order configuration")
    if digest(source) != original_digest:
        raise RuntimeError("Input mesh changed during the operation")
    provenance = {
        "source": str(source), "source_sha256": original_digest,
        "refined_mesh": str(target.resolve()), "refined_sha256": digest(target),
        "operation": "stroid.refinement.UniformRefinement(loaded_mesh, 1)",
        "additional_levels": 1, "initial_level": initial_level,
        "final_level": mesh.refinement_levels,
        "geometry_order": order, "core_mapping": mapping,
        "counts_before": before, "counts_after": after,
        "python": sys.executable, "stroid_version": stroid.__version__,
        "stroid_extension": _stroid.__file__,
        "stroid_extension_sha256": digest(_stroid.__file__),
        "geometry_policy": "STROID reprojects refined logical geometry; not fixed physical-polynomial subdivision",
        "downstream_extra_refinements": 0,
    }
    with (args.output_directory / "refinement.json").open("x") as stream:
        json.dump(provenance, stream, indent=2)
        stream.write("\n")
    print(f"Saved {mesh} to {target}; load with extra refinement=0", flush=True)


if __name__ == "__main__":
    main()