


'''simple docstring'''
import dataclasses
import re
import string
from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple

import numpy as np

from . import residue_constants


UpperCamelCase__    					=      Mapping[str, np.ndarray]
UpperCamelCase__    					=      Mapping[str, Any]  # Is a nested dict.
UpperCamelCase__    					=      0.01




@dataclasses.dataclass(frozen=__a      )
class 			lowerCamelCase_   :
      lowerCAmelCase__		     =     42  # [num_res, num_atom_type, 3]

      # Amino-acid type for each residue represented as an integer between 0 and
      # 20, where 20 is 'X'.
      lowerCAmelCase__		     =     42  # [num_res]

      # Binary float mask to indicate presence of a particular atom. 1.0 if an atom
      # is present and 0.0 if not. This should be used for loss masking.
      lowerCAmelCase__		     =     42  # [num_res, num_atom_type]

      # Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
      lowerCAmelCase__		     =     42  # [num_res]

      # B-factors, or temperature factors, of each residue (in sq. angstroms units),
      # representing the displacement of the residue from its ground truth mean
      # value.
      lowerCAmelCase__		     =     42  # [num_res, num_atom_type]

      # Chain indices for multi-chain predictions
      lowerCAmelCase__		     =     None

      # Optional remark about the protein. Included as a comment in output PDB
      # files
      lowerCAmelCase__		     =     None

      # Templates used to generate this protein (prediction-only)
      lowerCAmelCase__		     =     None

      # Chain corresponding to each parent
      lowerCAmelCase__		     =     None





def 					a__       (					lowerCAmelCase__						)    ->	Protein:
    UpperCAmelCase__			:		Union[str, Any]													=				R'''(\[[A-Z]+\]\n)'''
    UpperCAmelCase__			:		List[str]													=				[tag.strip() for tag in re.split(lowerCAmelCase__  ,					lowerCAmelCase__						) if len(lowerCAmelCase__						) > 0]
    UpperCAmelCase__			:		Iterator[Tuple[str, List[str]]]													=				zip(tags[0::2]  ,					[l.split('''\n'''						) for l in tags[1::2]]						)

    UpperCAmelCase__			:		List[str]													=				["N", "CA", "C"]
    UpperCAmelCase__			:		str													=				None
    UpperCAmelCase__			:		List[Any]													=				None
    UpperCAmelCase__			:		List[Any]													=				None
    for g in groups:
        if "[PRIMARY]" == g[0]:
            UpperCAmelCase__			:		str													=				g[1][0].strip()
            for i in range(len(lowerCAmelCase__						)						):
                if seq[i] not in residue_constants.restypes:
                    UpperCAmelCase__			:		Union[str, Any]													=				'''X'''  # FIXME: strings are immutable
            UpperCAmelCase__			:		str													=				np.array(
                [residue_constants.restype_order.get(lowerCAmelCase__  ,					residue_constants.restype_num						) for res_symbol in seq]						)
        elif "[TERTIARY]" == g[0]:
            UpperCAmelCase__			:		List[List[float]]													=				[]
            for axis in range(3						):
                tertiary.append(list(map(lowerCAmelCase__  ,					g[1][axis].split()						)						)						)
            UpperCAmelCase__			:		List[Any]													=				np.array(lowerCAmelCase__						)
            UpperCAmelCase__			:		Optional[int]													=				np.zeros((len(tertiary[0]						) // 3, residue_constants.atom_type_num, 3)						).astype(np.floataa						)
            for i, atom in enumerate(lowerCAmelCase__						):
                UpperCAmelCase__			:		Any													=				np.transpose(tertiary_np[:, i::3]						)
            atom_positions *= PICO_TO_ANGSTROM
        elif "[MASK]" == g[0]:
            UpperCAmelCase__			:		Optional[Any]													=				np.array(list(map({'''-''': 0, '''+''': 1}.get  ,					g[1][0].strip()						)						)						)
            UpperCAmelCase__			:		str													=				np.zeros(
                (
                    len(lowerCAmelCase__						),
                    residue_constants.atom_type_num,
                )						).astype(np.floataa						)
            for i, atom in enumerate(lowerCAmelCase__						):
                UpperCAmelCase__			:		int													=				1
            atom_mask *= mask[..., None]

    assert aatype is not None

    return Protein(
        atom_positions=lowerCAmelCase__  ,					atom_mask=lowerCAmelCase__  ,					aatype=lowerCAmelCase__  ,					residue_index=np.arange(len(lowerCAmelCase__						)						)  ,					b_factors=lowerCAmelCase__  ,					)





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__ = 0						)    ->	List[str]:
    UpperCAmelCase__			:		List[str]													=				[]

    UpperCAmelCase__			:		Tuple													=				prot.remark
    if remark is not None:
        pdb_headers.append(F"""REMARK {remark}"""						)

    UpperCAmelCase__			:		str													=				prot.parents
    UpperCAmelCase__			:		Tuple													=				prot.parents_chain_index
    if parents is not None and parents_chain_index is not None:
        UpperCAmelCase__			:		Tuple													=				[p for i, p in zip(lowerCAmelCase__  ,					lowerCAmelCase__						) if i == chain_id]

    if parents is None or len(lowerCAmelCase__						) == 0:
        UpperCAmelCase__			:		int													=				['''N/A''']

    pdb_headers.append(F"""PARENT {" ".join(lowerCAmelCase__						)}"""						)

    return pdb_headers





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	str:
    UpperCAmelCase__			:		List[str]													=				[]
    UpperCAmelCase__			:		str													=				pdb_str.split('''\n'''						)

    UpperCAmelCase__			:		int													=				prot.remark
    if remark is not None:
        out_pdb_lines.append(F"""REMARK {remark}"""						)

    UpperCAmelCase__			:		List[List[str]]
    if prot.parents is not None and len(prot.parents						) > 0:
        UpperCAmelCase__			:		int													=				[]
        if prot.parents_chain_index is not None:
            UpperCAmelCase__			:		Dict[str, List[str]]													=				{}
            for p, i in zip(prot.parents  ,					prot.parents_chain_index						):
                parent_dict.setdefault(str(lowerCAmelCase__						)  ,					[]						)
                parent_dict[str(lowerCAmelCase__						)].append(lowerCAmelCase__						)

            UpperCAmelCase__			:		Tuple													=				max([int(lowerCAmelCase__						) for chain_idx in parent_dict]						)
            for i in range(max_idx + 1						):
                UpperCAmelCase__			:		int													=				parent_dict.get(str(lowerCAmelCase__						)  ,					['''N/A''']						)
                parents_per_chain.append(lowerCAmelCase__						)
        else:
            parents_per_chain.append(list(prot.parents						)						)
    else:
        UpperCAmelCase__			:		Union[str, Any]													=				[['''N/A''']]

    def make_parent_line(lowerCAmelCase__						) -> str:
        return F"""PARENT {" ".join(lowerCAmelCase__						)}"""

    out_pdb_lines.append(make_parent_line(parents_per_chain[0]						)						)

    UpperCAmelCase__			:		List[str]													=				0
    for i, l in enumerate(lowerCAmelCase__						):
        if "PARENT" not in l and "REMARK" not in l:
            out_pdb_lines.append(lowerCAmelCase__						)
        if "TER" in l and "END" not in lines[i + 1]:
            chain_counter += 1
            if not chain_counter >= len(lowerCAmelCase__						):
                UpperCAmelCase__			:		Any													=				parents_per_chain[chain_counter]
            else:
                UpperCAmelCase__			:		int													=				['''N/A''']

            out_pdb_lines.append(make_parent_line(lowerCAmelCase__						)						)

    return "\n".join(lowerCAmelCase__						)





def 					a__       (					lowerCAmelCase__						)    ->	str:
    UpperCAmelCase__			:		Union[str, Any]													=				residue_constants.restypes + ['''X''']

    def res_atoa(lowerCAmelCase__						) -> str:
        return residue_constants.restype_atoa.get(restypes[r]  ,					'''UNK'''						)

    UpperCAmelCase__			:		Dict													=				residue_constants.atom_types

    UpperCAmelCase__			:		List[str]													=				[]

    UpperCAmelCase__			:		Optional[Any]													=				prot.atom_mask
    UpperCAmelCase__			:		Tuple													=				prot.aatype
    UpperCAmelCase__			:		Any													=				prot.atom_positions
    UpperCAmelCase__			:		List[Any]													=				prot.residue_index.astype(np.intaa						)
    UpperCAmelCase__			:		Any													=				prot.b_factors
    UpperCAmelCase__			:		List[Any]													=				prot.chain_index

    if np.any(aatype > residue_constants.restype_num						):
        raise ValueError('''Invalid aatypes.'''						)

    UpperCAmelCase__			:		str													=				get_pdb_headers(lowerCAmelCase__						)
    if len(lowerCAmelCase__						) > 0:
        pdb_lines.extend(lowerCAmelCase__						)

    UpperCAmelCase__			:		int													=				aatype.shape[0]
    UpperCAmelCase__			:		Tuple													=				1
    UpperCAmelCase__			:		str													=				0
    UpperCAmelCase__			:		Any													=				string.ascii_uppercase
    UpperCAmelCase__			:		List[str]													=				None
    # Add all atom sites.
    for i in range(lowerCAmelCase__						):
        UpperCAmelCase__			:		Tuple													=				res_atoa(aatype[i]						)
        for atom_name, pos, mask, b_factor in zip(lowerCAmelCase__  ,					atom_positions[i]  ,					atom_mask[i]  ,					b_factors[i]						):
            if mask < 0.5:
                continue

            UpperCAmelCase__			:		Any													=				'''ATOM'''
            UpperCAmelCase__			:		Tuple													=				atom_name if len(lowerCAmelCase__						) == 4 else F""" {atom_name}"""
            UpperCAmelCase__			:		Any													=				''''''
            UpperCAmelCase__			:		Union[str, Any]													=				''''''
            UpperCAmelCase__			:		int													=				1.0_0
            UpperCAmelCase__			:		Optional[int]													=				atom_name[0]  # Protein supports only C, N, O, S, this works.
            UpperCAmelCase__			:		List[str]													=				''''''

            UpperCAmelCase__			:		List[Any]													=				'''A'''
            if chain_index is not None:
                UpperCAmelCase__			:		Union[str, Any]													=				chain_tags[chain_index[i]]

            # PDB is a columnar format, every space matters here!
            UpperCAmelCase__			:		Dict													=				(
                F"""{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"""
                F"""{res_name_a:>3} {chain_tag:>1}"""
                F"""{residue_index[i]:>4}{insertion_code:>1}   """
                F"""{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"""
                F"""{occupancy:>6.2f}{b_factor:>6.2f}          """
                F"""{element:>2}{charge:>2}"""
            )
            pdb_lines.append(lowerCAmelCase__						)
            atom_index += 1

        UpperCAmelCase__			:		Optional[int]													=				i == n - 1
        if chain_index is not None:
            if i != n - 1 and chain_index[i + 1] != prev_chain_index:
                UpperCAmelCase__			:		str													=				True
                UpperCAmelCase__			:		Dict													=				chain_index[i + 1]

        if should_terminate:
            # Close the chain.
            UpperCAmelCase__			:		Any													=				'''TER'''
            UpperCAmelCase__			:		str													=				(
                F"""{chain_end:<6}{atom_index:>5}      {res_atoa(aatype[i]						):>3} {chain_tag:>1}{residue_index[i]:>4}"""
            )
            pdb_lines.append(lowerCAmelCase__						)
            atom_index += 1

            if i != n - 1:
                # "prev" is a misnomer here. This happens at the beginning of
                # each new chain.
                pdb_lines.extend(get_pdb_headers(lowerCAmelCase__  ,					lowerCAmelCase__						)						)

    pdb_lines.append('''END'''						)
    pdb_lines.append(''''''						)
    return "\n".join(lowerCAmelCase__						)





def 					a__       (					lowerCAmelCase__						)    ->	np.ndarray:
    return residue_constants.STANDARD_ATOM_MASK[prot.aatype]







def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__ = None  ,					lowerCAmelCase__ = None  ,					lowerCAmelCase__ = None  ,					lowerCAmelCase__ = None  ,					lowerCAmelCase__ = None  ,					)    ->	Protein:
    return Protein(
        aatype=features['''aatype''']  ,					atom_positions=result['''final_atom_positions''']  ,					atom_mask=result['''final_atom_mask''']  ,					residue_index=features['''residue_index'''] + 1  ,					b_factors=b_factors if b_factors is not None else np.zeros_like(result['''final_atom_mask''']						)  ,					chain_index=lowerCAmelCase__  ,					remark=lowerCAmelCase__  ,					parents=lowerCAmelCase__  ,					parents_chain_index=lowerCAmelCase__  ,					)




'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__    					=      {
    '''configuration_distilbert''': [
        '''DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''DistilBertConfig''',
        '''DistilBertOnnxConfig''',
    ],
    '''tokenization_distilbert''': ['''DistilBertTokenizer'''],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      ['''DistilBertTokenizerFast''']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''DistilBertForMaskedLM''',
      '''DistilBertForMultipleChoice''',
      '''DistilBertForQuestionAnswering''',
      '''DistilBertForSequenceClassification''',
      '''DistilBertForTokenClassification''',
      '''DistilBertModel''',
      '''DistilBertPreTrainedModel''',
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''TFDistilBertForMaskedLM''',
      '''TFDistilBertForMultipleChoice''',
      '''TFDistilBertForQuestionAnswering''',
      '''TFDistilBertForSequenceClassification''',
      '''TFDistilBertForTokenClassification''',
      '''TFDistilBertMainLayer''',
      '''TFDistilBertModel''',
      '''TFDistilBertPreTrainedModel''',
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''FlaxDistilBertForMaskedLM''',
      '''FlaxDistilBertForMultipleChoice''',
      '''FlaxDistilBertForQuestionAnswering''',
      '''FlaxDistilBertForSequenceClassification''',
      '''FlaxDistilBertForTokenClassification''',
      '''FlaxDistilBertModel''',
      '''FlaxDistilBertPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_distilbert import (
      DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      DistilBertConfig,
      DistilBertOnnxConfig,
  )
  from .tokenization_distilbert import DistilBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_distilbert_fast import DistilBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_distilbert import (
        DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        DistilBertForMaskedLM,
        DistilBertForMultipleChoice,
        DistilBertForQuestionAnswering,
        DistilBertForSequenceClassification,
        DistilBertForTokenClassification,
        DistilBertModel,
        DistilBertPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_distilbert import (
        TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDistilBertForMaskedLM,
        TFDistilBertForMultipleChoice,
        TFDistilBertForQuestionAnswering,
        TFDistilBertForSequenceClassification,
        TFDistilBertForTokenClassification,
        TFDistilBertMainLayer,
        TFDistilBertModel,
        TFDistilBertPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_distilbert import (
        FlaxDistilBertForMaskedLM,
        FlaxDistilBertForMultipleChoice,
        FlaxDistilBertForQuestionAnswering,
        FlaxDistilBertForSequenceClassification,
        FlaxDistilBertForTokenClassification,
        FlaxDistilBertModel,
        FlaxDistilBertPreTrainedModel,
    )

else:
  import sys

  UpperCamelCase__    					=      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
