import json
import os
from pathlib import Path

import numpy as np

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1")

import tensorflow as tf

tf.get_logger().setLevel("ERROR")

tf_config = json.loads(os.environ["TF_CONFIG"])
task = tf_config["task"]
workers = tf_config["cluster"]["worker"]

per_worker_batch = 16
global_batch = per_worker_batch * len(workers)
steps_per_epoch = 4
epochs = 3

rng = np.random.default_rng(7)
features = rng.normal(size=(256, 8)).astype("float32")
score = (
    features[:, 0] * 0.8
    + features[:, 1] * 0.6
    - features[:, 2] * 0.4
    + features[:, 3] * 0.2
)
labels = (score > 0).astype("float32")[:, None]

strategy = tf.distribute.MultiWorkerMirroredStrategy()

print(
    f"worker={task['index']} "
    f"cluster_workers={len(workers)} "
    f"replicas={strategy.num_replicas_in_sync} "
    f"global_batch={global_batch}",
    flush=True,
)


def dataset_fn(input_context):
    batch_size = input_context.get_per_replica_batch_size(global_batch)
    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
    dataset = dataset.shard(
        input_context.num_input_pipelines,
        input_context.input_pipeline_id,
    )
    dataset = dataset.shuffle(128, seed=7, reshuffle_each_iteration=False)
    return dataset.repeat().batch(batch_size).prefetch(tf.data.AUTOTUNE)


with strategy.scope():
    model = tf.keras.Sequential(
        [
            tf.keras.layers.Input(shape=(8,)),
            tf.keras.layers.Dense(16, activation="relu"),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ]
    )
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.03)
    loss_object = tf.keras.losses.BinaryCrossentropy(
        reduction=tf.keras.losses.Reduction.NONE
    )
    train_accuracy = tf.keras.metrics.BinaryAccuracy(name="binary_accuracy")

dist_dataset = strategy.distribute_datasets_from_function(dataset_fn)
dist_iterator = iter(dist_dataset)


@tf.function
def train_step(iterator):
    def step_fn(inputs):
        batch_features, batch_labels = inputs
        with tf.GradientTape() as tape:
            predictions = model(batch_features, training=True)
            per_example_loss = loss_object(batch_labels, predictions)
            loss = tf.nn.compute_average_loss(
                per_example_loss,
                global_batch_size=global_batch,
            )
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        train_accuracy.update_state(batch_labels, predictions)
        return loss

    per_replica_losses = strategy.run(step_fn, args=(next(iterator),))
    return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)


for epoch in range(1, epochs + 1):
    train_accuracy.reset_state()
    total_loss = 0.0
    for _ in range(steps_per_epoch):
        total_loss += train_step(dist_iterator)
    print(
        f"worker={task['index']} "
        f"epoch={epoch} "
        f"loss={float(total_loss / steps_per_epoch):.4f} "
        f"binary_accuracy={float(train_accuracy.result()):.4f}",
        flush=True,
    )

if task["type"] == "worker" and task["index"] == 0:
    Path("multiworker-complete.txt").write_text("chief worker completed training\n")

print(f"worker={task['index']} training_complete", flush=True)