<reponame>SamuelDAdams/CrypTen
import argparse
import logging
import os

from examples.multiprocess_launcher import MultiProcessLauncher


parser = argparse.ArgumentParser(description="CrypTen 1D CNN Training")
parser.add_argument(
    "--world_size",
    type=str,
    default=2,
    help="The number of parties to launch. Each party acts as its own process",
)

parser.add_argument(
    "--CNN_path",
    type=str,
    default="models/CNN.pth",
    help="The path to the pre-trained CNN model owned by Alice (variable is irrelivant to Bob)",
)

parser.add_argument(
    "--label_path",
    type=str,
    default="test/bin_labels_tensor.pth",
    help="The path to the labels of the data owned by Bob (variable is irrelivant to Alice)",
)

parser.add_argument(
    "--data_path",
    type=str,
    default="test/embeddings_tensor.pth",
    help="The path to the data owned by Bob (variable is irrelivant to Alice)",
)

parser.add_argument(
    "--multiprocess",
    default=False,
    action="store_true",
    help="Run example in multiprocess mode",
)

parser.add_argument(
    "--batch",
    type=int,
    default=0,
    help="0 means we don't want to batch classifications, 1 means we do (parties must agree on this value)",
)

parser.add_argument(
    "--count",
    type=int,
    default=0,
    help="How many samples we want to classify. So, for examples, count = 200 means we classify 200 samples (parties must agree on this value)",
)


def _run_experiment(args):
    level = logging.INFO
    if "RANK" in os.environ and os.environ["RANK"] != "0":
        level = logging.CRITICAL
    logging.getLogger().setLevel(level)
    logging.basicConfig(
        level=level,
        format="%(asctime)s - %(process)d - %(name)s - %(levelname)s - %(message)s",
    )
    from mpc_1D_CNN import run_mpc_1D_CNN

    run_mpc_1D_CNN(
        args.CNN_path, args.data_path, args.label_path, args.count, args.batch
    )


def main(run_experiment):
    args = parser.parse_args()
    if args.multiprocess:
        launcher = MultiProcessLauncher(args.world_size, run_experiment, args)
        launcher.start()
        launcher.join()
        launcher.terminate()
    else:
        run_experiment(args)


if __name__ == "__main__":
    main(_run_experiment)
