# # Snapshot checkpoint (non-persistent)
# The checkpoint method described in [Writing function checkpoints](./writing_functions_checkpoint)
# are *N-to-M*, meaning that you can write them out on N-processes and read them in on M processes.
#
# As discussed in that chapter, these checkpoints need to be associated with a mesh.
# This is because the function is defined on a specific function space, which in turn is
# defined on a specific mesh.
#
# However, there are certain scenarios where you simply want to store a checkpoint associated
# with the current mesh, that should only be possible to use during this simulation.
# An example use-case is when running an iterative solver, and wanting a fall-back mechanism that
# does not require extra RAM.

# In this example, we will demonstrate how to write a
# {py:func}`snapshot checkpoint<io4dolfinx.snapshot_checkpoint>` to disk.

# First we define a {py:class}`function<dolfinx.fem.Function>` `f` that we want to represent in
# the {py:class}`function space<dolfinx.fem.FunctionSpace>`.

# +
import logging
from pathlib import Path

import ipyparallel as ipp


def f(x):
    import numpy as np

    return np.sin(x[0]) + 0.1 * x[1]


# -

# Next, we create a mesh and an appropriate function space and read and write from file.
# Note that for both these operations, we use {py:func}`io4dolfinx.snapshot_checkpoint`,
# with different read and write modes.


def read_write_snapshot(filename: Path):
    from mpi4py import MPI

    import dolfinx
    import numpy as np

    import io4dolfinx

    mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 3, 7, 4)
    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 5))
    u = dolfinx.fem.Function(V)
    u.interpolate(f)
    u.name = "Current_solution"
    # Next, we store the solution to file
    io4dolfinx.snapshot_checkpoint(u, filename, io4dolfinx.FileMode.write)

    # Next, we create a new function and load the solution into it
    u_new = dolfinx.fem.Function(V)
    u_new.name = "Read_solution"
    io4dolfinx.snapshot_checkpoint(u_new, filename, io4dolfinx.FileMode.read)

    # Next, we verify that the solution is correct
    np.testing.assert_allclose(u_new.x.array, u.x.array, atol=np.finfo(float).eps)

    print(f"{MPI.COMM_WORLD.rank + 1}/{MPI.COMM_WORLD.size}: Successfully wrote and read snapshot")


# +

mesh_file = Path("snapshot.bp")

with ipp.Cluster(engines="mpi", n=3, log_level=logging.ERROR) as cluster:
    cluster[:].push({"f": f})
    query = cluster[:].apply_async(
        read_write_snapshot,
        mesh_file,
    )
    query.wait()
    assert query.successful(), query.error
    print("".join(query.stdout))
# -
