import pickle
import numpy as np
import os

def save_snapshot(results, filename):
    """
    Save results (e.g., dict of outputs, predictions, metrics) to a file.
    """
    with open(filename, "wb") as f:
        pickle.dump(results, f)

def load_snapshot(filename):
    """
    Load previously saved results from a file.
    Raises FileNotFoundError if the file does not exist.
    """
    if not os.path.isfile(filename):
        raise FileNotFoundError(f"Snapshot file '{filename}' does not exist.")
    with open(filename, "rb") as f:
        return pickle.load(f)

def compare_snapshots(results, snapshot, atol=1e-6):
    """
    Compare current results to a snapshot.
    Returns True if all values are close, False otherwise.
    """
    for key in results:
        if key not in snapshot:
            print(f"Key {key} missing in snapshot.")
            return False
        if isinstance(results[key], np.ndarray):
            if not np.allclose(results[key], snapshot[key], atol=atol):
                print(f"Mismatch in {key}")
                return False
        elif isinstance(results[key], float):
            if not np.isclose(results[key], snapshot[key], atol=atol):
                print(f"Mismatch in {key}")
                return False
        # Add more types as needed
    return True