import os
from pathlib import Path

os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")

import tensorflow as tf


def float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))


def int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


rows = [
    {"feature_a": 0.10, "feature_b": 1.20, "label": 0},
    {"feature_a": 0.40, "feature_b": 0.70, "label": 1},
    {"feature_a": 0.30, "feature_b": 1.10, "label": 0},
    {"feature_a": 0.90, "feature_b": 0.20, "label": 1},
]

record_path = Path("data/training.tfrecord")
record_path.parent.mkdir(parents=True, exist_ok=True)

with tf.io.TFRecordWriter(str(record_path)) as writer:
    for row in rows:
        example = tf.train.Example(
            features=tf.train.Features(
                feature={
                    "feature_a": float_feature(row["feature_a"]),
                    "feature_b": float_feature(row["feature_b"]),
                    "label": int64_feature(row["label"]),
                }
            )
        )
        writer.write(example.SerializeToString())

feature_description = {
    "feature_a": tf.io.FixedLenFeature([], tf.float32),
    "feature_b": tf.io.FixedLenFeature([], tf.float32),
    "label": tf.io.FixedLenFeature([], tf.int64),
}


def parse_record(serialized_example):
    parsed = tf.io.parse_single_example(serialized_example, feature_description)
    label = parsed.pop("label")
    return parsed, label


def read_records():
    return tf.data.TFRecordDataset(str(record_path))


parsed_dataset = read_records().map(parse_record)

record_count = sum(1 for _ in read_records())
first_features, first_label = next(iter(parsed_dataset))
batch_features, batch_labels = next(iter(parsed_dataset.batch(2)))

first_record = {
    "feature_a": round(float(first_features["feature_a"].numpy()), 2),
    "feature_b": round(float(first_features["feature_b"].numpy()), 2),
    "label": int(first_label.numpy()),
}
first_batch_features = {
    name: [round(float(value), 2) for value in tensor.numpy().tolist()]
    for name, tensor in batch_features.items()
}

print(f"tfrecord_path={record_path}")
print(f"file_exists={record_path.exists()}")
print(f"record_count={record_count}")
print(f"first_record={first_record}")
print(f"first_batch_features={first_batch_features}")
print(f"first_batch_labels={batch_labels.numpy().tolist()}")