# Path: alphafold/model/tf/protein_features.py
# class FeatureType(enum.Enum):
#   ZERO_DIM = 0  # Shape [x]
#   ONE_DIM = 1  # Shape [num_res, x]
#   TWO_DIM = 2  # Shape [num_res, num_res, x]
#   MSA = 3  # Shape [msa_length, num_res, x]
# NUM_RES = "num residues placeholder"
# NUM_SEQ = "length msa placeholder"
# NUM_TEMPLATES = "num templates placeholder"
# FEATURES = {
#     #### Static features of a protein sequence ####
#     "aatype": (tf.float32, [NUM_RES, 21]),
#     "between_segment_residues": (tf.int64, [NUM_RES, 1]),
#     "deletion_matrix": (tf.float32, [NUM_SEQ, NUM_RES, 1]),
#     "domain_name": (tf.string, [1]),
#     "msa": (tf.int64, [NUM_SEQ, NUM_RES, 1]),
#     "num_alignments": (tf.int64, [NUM_RES, 1]),
#     "residue_index": (tf.int64, [NUM_RES, 1]),
#     "seq_length": (tf.int64, [NUM_RES, 1]),
#     "sequence": (tf.string, [1]),
#     "all_atom_positions": (tf.float32,
#                            [NUM_RES, residue_constants.atom_type_num, 3]),
#     "all_atom_mask": (tf.int64, [NUM_RES, residue_constants.atom_type_num]),
#     "resolution": (tf.float32, [1]),
#     "template_domain_names": (tf.string, [NUM_TEMPLATES]),
#     "template_sum_probs": (tf.float32, [NUM_TEMPLATES, 1]),
#     "template_aatype": (tf.float32, [NUM_TEMPLATES, NUM_RES, 22]),
#     "template_all_atom_positions": (tf.float32, [
#         NUM_TEMPLATES, NUM_RES, residue_constants.atom_type_num, 3
#     ]),
#     "template_all_atom_masks": (tf.float32, [
#         NUM_TEMPLATES, NUM_RES, residue_constants.atom_type_num, 1
#     ]),
# }
# FEATURE_TYPES = {k: v[0] for k, v in FEATURES.items()}
# FEATURE_SIZES = {k: v[1] for k, v in FEATURES.items()}
# def register_feature(name: str,
#                      type_: tf.dtypes.DType,
#                      shape_: Tuple[Union[str, int]]):
# def shape(feature_name: str,
#           num_residues: int,
#           msa_length: int,
#           num_templates: Optional[int] = None,
#           features: Optional[FeaturesMetadata] = None):


from typing import Dict, Mapping, Optional, Sequence
from alphafold.model.tf import protein_features
import numpy as np
import tensorflow.compat.v1 as tf
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Datasets consisting of proteins."""

TensorDict = Dict[str, tf.Tensor]


def parse_tfexample(
    raw_data: bytes,

# Path: alphafold/model/tf/protein_features.py
# class FeatureType(enum.Enum):
#   ZERO_DIM = 0  # Shape [x]
#   ONE_DIM = 1  # Shape [num_res, x]
#   TWO_DIM = 2  # Shape [num_res, num_res, x]
#   MSA = 3  # Shape [msa_length, num_res, x]
# NUM_RES = "num residues placeholder"
# NUM_SEQ = "length msa placeholder"
# NUM_TEMPLATES = "num templates placeholder"
# FEATURES = {
#     #### Static features of a protein sequence ####
#     "aatype": (tf.float32, [NUM_RES, 21]),
#     "between_segment_residues": (tf.int64, [NUM_RES, 1]),
#     "deletion_matrix": (tf.float32, [NUM_SEQ, NUM_RES, 1]),
#     "domain_name": (tf.string, [1]),
#     "msa": (tf.int64, [NUM_SEQ, NUM_RES, 1]),
#     "num_alignments": (tf.int64, [NUM_RES, 1]),
#     "residue_index": (tf.int64, [NUM_RES, 1]),
#     "seq_length": (tf.int64, [NUM_RES, 1]),
#     "sequence": (tf.string, [1]),
#     "all_atom_positions": (tf.float32,
#                            [NUM_RES, residue_constants.atom_type_num, 3]),
#     "all_atom_mask": (tf.int64, [NUM_RES, residue_constants.atom_type_num]),
#     "resolution": (tf.float32, [1]),
#     "template_domain_names": (tf.string, [NUM_TEMPLATES]),
#     "template_sum_probs": (tf.float32, [NUM_TEMPLATES, 1]),
#     "template_aatype": (tf.float32, [NUM_TEMPLATES, NUM_RES, 22]),
#     "template_all_atom_positions": (tf.float32, [
#         NUM_TEMPLATES, NUM_RES, residue_constants.atom_type_num, 3
#     ]),
#     "template_all_atom_masks": (tf.float32, [
#         NUM_TEMPLATES, NUM_RES, residue_constants.atom_type_num, 1
#     ]),
# }
# FEATURE_TYPES = {k: v[0] for k, v in FEATURES.items()}
# FEATURE_SIZES = {k: v[1] for k, v in FEATURES.items()}
# def register_feature(name: str,
#                      type_: tf.dtypes.DType,
#                      shape_: Tuple[Union[str, int]]):
# def shape(feature_name: str,
#           num_residues: int,
#           msa_length: int,
#           num_templates: Optional[int] = None,
#           features: Optional[FeaturesMetadata] = None):


# Path: alphafold/model/tf/proteins_dataset.py
from typing import Dict, Mapping, Optional, Sequence
from alphafold.model.tf import protein_features
import numpy as np
import tensorflow.compat.v1 as tf

# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Datasets consisting of proteins."""

TensorDict = Dict[str, tf.Tensor]


def parse_tfexample(
    raw_data: bytes,

    features: protein_features.FeaturesMetadata,