#!/usr/bin/env python

import logging
import os
import sys

LOG_FMT = "%(asctime)s %(levelname).1s - %(message)s"
logging.basicConfig(level=logging.INFO, format=LOG_FMT, datefmt="%Y/%m/%d %H:%M:%S")
log = logging.getLogger()

strategy = None
strategy_desc = ""
if "TF_CONFIG" in os.environ:
    tf_config = os.environ["TF_CONFIG"]
    log.info(f"TF_CONFIG detected: {tf_config}")
    log.info("Switching to MultiWorker distributed strategy")

    import tensorflow as tf

    strategy = tf.distribute.MultiWorkerMirroredStrategy()
    strategy_desc = "Multi Host / Multi GPU"

from tensorpotential.cli.gracemaker import main

if __name__ == "__main__":
    main(sys.argv[1:], strategy=strategy, strategy_desc=strategy_desc)
