# Path: tensorflow_decision_forests/component/inspector/blob_sequence.py
# class Reader(object):
# class Writer(object):
#   def __init__(self, path: str):
#   def open(self, path: str):
#   def close(self):
#   def read(self) -> Optional[bytes]:
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_val, exc_tb):
#   def __iter__(self) -> Iterator[bytes]:
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_val, exc_tb):
#   def __init__(self, path: Optional[str] = None):
#   def open(self, path: str):
#   def write(self, blob: bytes):
#   def close(self):


import os
import tensorflow.compat.v1 as tf
from absl import logging
from tensorflow_decision_forests.component.inspector import blob_sequence
from yggdrasil_decision_forests.dataset import data_spec_pb2
from yggdrasil_decision_forests.model import abstract_model_pb2
from yggdrasil_decision_forests.model.decision_tree import decision_tree_pb2
from yggdrasil_decision_forests.model.gradient_boosted_trees import gradient_boosted_trees_pb2
from yggdrasil_decision_forests.model.random_forest import random_forest_pb2
from yggdrasil_decision_forests.utils import distribution_pb2
  logging.info("Create toy model in %s", path)

  tf.io.gfile.makedirs(path)

  with tf.io.gfile.GFile(os.path.join(path, "done"), "w") as f:
    f.write("Something")

  data_spec = build_toy_data_spec(
      add_boolean_features=add_boolean_features, has_catset=has_catset)
  with tf.io.gfile.GFile(os.path.join(path, "data_spec.pb"), "w") as f:
    f.write(data_spec.SerializeToString())

  header = abstract_model_pb2.AbstractModel(
      name="RANDOM_FOREST",
      task=abstract_model_pb2.CLASSIFICATION,
      label_col_idx=3,
      input_features=[0, 1, 2] + ([5] if add_boolean_features else []) +
      ([5, 6] if has_catset else []))
  with tf.io.gfile.GFile(os.path.join(path, "header.pb"), "w") as f:
    f.write(header.SerializeToString())

  rf_header = random_forest_pb2.Header(
      num_node_shards=1,
      num_trees=num_trees,
      winner_take_all_inference=winner_take_all_inference,
      node_format="BLOB_SEQUENCE")
  with tf.io.gfile.GFile(os.path.join(path, "random_forest_header.pb"),
                         "w") as f:
    f.write(rf_header.SerializeToString())


# Path: tensorflow_decision_forests/component/inspector/blob_sequence.py
# class Reader(object):
# class Writer(object):
#   def __init__(self, path: str):
#   def open(self, path: str):
#   def close(self):
#   def read(self) -> Optional[bytes]:
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_val, exc_tb):
#   def __iter__(self) -> Iterator[bytes]:
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_val, exc_tb):
#   def __init__(self, path: Optional[str] = None):
#   def open(self, path: str):
#   def write(self, blob: bytes):
#   def close(self):


# Path: tensorflow_decision_forests/tensorflow/ops/inference/test_utils.py
import os
import tensorflow.compat.v1 as tf
from absl import logging
from tensorflow_decision_forests.component.inspector import blob_sequence
from yggdrasil_decision_forests.dataset import data_spec_pb2
from yggdrasil_decision_forests.model import abstract_model_pb2
from yggdrasil_decision_forests.model.decision_tree import decision_tree_pb2
from yggdrasil_decision_forests.model.gradient_boosted_trees import gradient_boosted_trees_pb2
from yggdrasil_decision_forests.model.random_forest import random_forest_pb2
from yggdrasil_decision_forests.utils import distribution_pb2

  logging.info("Create toy model in %s", path)

  tf.io.gfile.makedirs(path)

  with tf.io.gfile.GFile(os.path.join(path, "done"), "w") as f:
    f.write("Something")

  data_spec = build_toy_data_spec(
      add_boolean_features=add_boolean_features, has_catset=has_catset)
  with tf.io.gfile.GFile(os.path.join(path, "data_spec.pb"), "w") as f:
    f.write(data_spec.SerializeToString())

  header = abstract_model_pb2.AbstractModel(
      name="RANDOM_FOREST",
      task=abstract_model_pb2.CLASSIFICATION,
      label_col_idx=3,
      input_features=[0, 1, 2] + ([5] if add_boolean_features else []) +
      ([5, 6] if has_catset else []))
  with tf.io.gfile.GFile(os.path.join(path, "header.pb"), "w") as f:
    f.write(header.SerializeToString())

  rf_header = random_forest_pb2.Header(
      num_node_shards=1,
      num_trees=num_trees,
      winner_take_all_inference=winner_take_all_inference,
      node_format="BLOB_SEQUENCE")
  with tf.io.gfile.GFile(os.path.join(path, "random_forest_header.pb"),
                         "w") as f:
    f.write(rf_header.SerializeToString())


  with blob_sequence.Writer(os.path.join(