import numpy as np
import random
import tensorflow as tf
import inspect
from typing import Any, Callable, Dict, Optional, Text, Union, Iterable
import os

#隨機生成131072的dna序列
length = 131072
def randomSeq(length):
  return ''.join([random.choice('ACGT') for i in range(length)])

#DNA序列轉(zhuǎn)為one-hot編碼，可以直接拿來用
def one_hot_encode(sequence: str,
                   alphabet: str = 'ACGT',
                   neutral_alphabet: str = 'N',
                   neutral_value: Any = 0,
                   dtype=np.float32) -> np.ndarray:
  """One-hot encode sequence."""
  def to_uint8(string):
    return np.frombuffer(string.encode('ascii'), dtype=np.uint8)
  hash_table = np.zeros((np.iinfo(np.uint8).max, len(alphabet)), dtype=dtype)
  hash_table[to_uint8(alphabet)] = np.eye(len(alphabet), dtype=dtype)
  hash_table[to_uint8(neutral_alphabet)] = neutral_value
  hash_table = hash_table.astype(dtype)
  return hash_table[to_uint8(sequence)]

#給出結(jié)果的tfr文件的路徑
path = '/content/drive/MyDrive/test_Enformer/result.tfr'
#dna的numpy數(shù)組轉(zhuǎn)成字節(jié)流，這樣才能存儲
dnaVec = dnaVec.tobytes()

from tensorflow.core.example.feature_pb2 import BytesList
with tf.io.TFRecordWriter(path) as writer:
  feature = {
      #序列使用的是tf.train.BytesList類型
      'sequence':tf.train.Feature(bytes_list=tf.train.BytesList(value=[dnaVec])),
      #label是隨機生成的0，或者1
      'label':tf.train.Feature(int64_list=tf.train.Int64List(value=[np.random.choice([0,1])]))
  }
  example = tf.train.Example(features=tf.train.Features(feature=feature))
  writer.write(example.SerializeToString())

#從剛才的路徑中加載數(shù)據(jù)集
dataset = tf.data.TFRecordDataset(path)

#定義Feature結(jié)構(gòu)，告訴解碼器每個Feature的類型是什么
feature_description = {"sequence": tf.io.FixedLenFeature((), tf.string),
              "label": tf.io.FixedLenFeature((), tf.int64)}

 #將 TFRecord 文件中的每一個序列化的 tf.train.Example 解碼
def parse_example(example_string):
  #解析之后得到的example
  example = tf.io.parse_single_example(example_string,feature_description)
  #example['sequence']還是字節(jié)流的形式，重新轉(zhuǎn)為數(shù)字向量
  sequence = tf.io.decode_raw(example['sequence'], tf.float32)
  sequence = tf.reshape(sequence,(length,4))  #形狀需要重塑，不然就是一個長向量
  label = tf.cast(example['label'],tf.int64)  #標(biāo)簽對應(yīng)的類型轉(zhuǎn)換
#每一天example解析后返回對應(yīng)的一個字典
  return {
      'sequence':sequence,
      'label': label
  }

#把parse_example函數(shù)映射到dataset中的每個example,
#這里的dataset中只有一個example
dataset = dataset.map(parse_example)