import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import tensorflow as tf

tf.keras.utils.set_random_seed(7)

train_features = tf.constant(
    [
        [510.0, 6.5, 1.2],
        [620.0, 7.1, 0.8],
        [470.0, 5.9, 1.7],
        [730.0, 8.0, 0.4],
        [690.0, 7.8, 0.6],
        [540.0, 6.8, 1.0],
        [455.0, 5.7, 1.9],
        [760.0, 8.3, 0.3],
    ],
    dtype=tf.float32,
)
train_labels = tf.constant([0, 1, 0, 1, 1, 0, 0, 1], dtype=tf.float32)

validation_features = tf.constant(
    [
        [500.0, 6.4, 1.1],
        [710.0, 7.9, 0.5],
        [480.0, 6.0, 1.6],
        [745.0, 8.1, 0.4],
    ],
    dtype=tf.float32,
)
validation_labels = tf.constant([0, 1, 0, 1], dtype=tf.float32)

feature_scale = tf.constant([1000.0, 10.0, 10.0], dtype=tf.float32)


def prepare_example(features, label):
    features = tf.cast(features, tf.float32) / feature_scale
    label = tf.cast(label, tf.float32)
    return features, label


def build_dataset(features, labels, batch_size, training):
    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
    dataset = dataset.map(prepare_example, num_parallel_calls=tf.data.AUTOTUNE)
    if training:
        dataset = dataset.shuffle(
            buffer_size=int(features.shape[0]),
            seed=7,
            reshuffle_each_iteration=True,
        )
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset


train_dataset = build_dataset(
    train_features,
    train_labels,
    batch_size=4,
    training=True,
)
validation_dataset = build_dataset(
    validation_features,
    validation_labels,
    batch_size=2,
    training=False,
)

model = tf.keras.Sequential(
    [
        tf.keras.layers.Input(shape=(3,)),
        tf.keras.layers.Dense(4, activation="relu"),
        tf.keras.layers.Dense(1, activation="sigmoid"),
    ]
)
model.compile(optimizer="adam", loss="binary_crossentropy")

history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=1,
    shuffle=False,
    verbose=0,
)

train_batch_features, train_batch_labels = next(iter(train_dataset))
validation_batch_features, validation_batch_labels = next(iter(validation_dataset))
train_feature_spec, train_label_spec = train_dataset.element_spec
validation_feature_spec, validation_label_spec = validation_dataset.element_spec

print(f"tensorflow={tf.__version__}")
print(
    "train_feature_spec="
    f"shape={train_feature_spec.shape}, dtype={train_feature_spec.dtype.name}"
)
print(
    "train_label_spec="
    f"shape={train_label_spec.shape}, dtype={train_label_spec.dtype.name}"
)
print(
    "validation_feature_spec="
    f"shape={validation_feature_spec.shape}, dtype={validation_feature_spec.dtype.name}"
)
print(
    "validation_label_spec="
    f"shape={validation_label_spec.shape}, dtype={validation_label_spec.dtype.name}"
)
print(f"train_batches={int(train_dataset.cardinality())}")
print(f"validation_batches={int(validation_dataset.cardinality())}")
print(f"train_batch_shape={tuple(train_batch_features.shape)}")
print(f"train_label_shape={tuple(train_batch_labels.shape)}")
print(f"validation_batch_shape={tuple(validation_batch_features.shape)}")
print(f"validation_label_shape={tuple(validation_batch_labels.shape)}")
print(f"history_keys={sorted(history.history.keys())}")
print(f"fit_epochs={len(history.epoch)}")