"""Run the original checks and full proof, limiting only scheduling and reporting progress."""
import json
import multiprocessing
import pathlib
import sys
import time
sys.path.insert(0, "/study/code")
import test_verify
import verify

class ProgressPool:
    def __init__(self, ignored):
        self.pool = multiprocessing.Pool(4)
    def __enter__(self):
        return self
    def __exit__(self, *args):
        self.pool.close()
        self.pool.join()
    def imap(self, function, chunks):
        completed = 0
        total = sum(len(chunk) for chunk in chunks)
        start = time.monotonic()
        for part in self.pool.imap(function, chunks):
            completed += len(part)
            progress = {"completed_components": completed, "total_components": total, "elapsed_seconds": round(time.monotonic() - start, 3)}
            pathlib.Path("results/progress.json").write_text(json.dumps(progress) + "\n")
            yield part

if __name__ == "__main__":
    assert test_verify.main() == 0
    verify.Pool = ProgressPool
    sys.exit(verify.main())