

from typing import Dict, Tuple, overload

import torch
import torch.types
from torch import nn

from . import residue_constants as rc
from .rigid_utils import Rigid, Rotation
from .tensor_utils import batched_gather
@overload
def 		UpperCamelCase_(    snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor						,					snake_case__:						None     ) ->				torch.Tensor:
 ...
@overload
def 		UpperCamelCase_(    snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor     ) ->				Tuple[torch.Tensor, torch.Tensor]:
 ...
def 		UpperCamelCase_(    snake_case__:						Optional[Any]						,					snake_case__:						Optional[int]						,					snake_case__:						Union[str, Any]     ) ->				List[Any]:
 UpperCAmelCase__           =       aatype == rc.restype_order['G']
 UpperCAmelCase__           =       rc.atom_order['CA']
 UpperCAmelCase__           =       rc.atom_order['CB']
 UpperCAmelCase__           =       torch.where(
     is_gly[..., None].expand(*((-1,) * len(is_gly.shape     ))						,					3     )						,					all_atom_positions[..., ca_idx, :]						,					all_atom_positions[..., cb_idx, :]						,					)

 if all_atom_masks is not None:
  UpperCAmelCase__           =       torch.where(
      snake_case__						,					all_atom_masks[..., ca_idx]						,					all_atom_masks[..., cb_idx]						,					)
  return pseudo_beta, pseudo_beta_mask
 else:
  return pseudo_beta
def 		UpperCamelCase_(    snake_case__:						torch.Tensor						,					snake_case__:						Dict[str, torch.Tensor]     ) ->				torch.Tensor:
 UpperCAmelCase__           =       batched_gather(
     snake_case__						,					batch['residx_atom37_to_atom14']						,					dim=-2						,					no_batch_dims=len(atomaa.shape[:-2]     )						,					)

 UpperCAmelCase__           =       atomaa_data * batch['atom37_atom_exists'][..., None]

 return atomaa_data
def 		UpperCamelCase_(    snake_case__:						Dict[str, torch.Tensor]     ) ->				torch.Tensor:
 UpperCAmelCase__           =       template_feats['template_aatype']
 UpperCAmelCase__           =       template_feats['template_torsion_angles_sin_cos']
 UpperCAmelCase__           =       template_feats['template_alt_torsion_angles_sin_cos']
 UpperCAmelCase__           =       template_feats['template_torsion_angles_mask']
 UpperCAmelCase__           =       torch.cat(
     [
         nn.functional.one_hot(snake_case__						,					22     ),
         torsion_angles_sin_cos.reshape(*torsion_angles_sin_cos.shape[:-2]						,					14     ),
         alt_torsion_angles_sin_cos.reshape(*alt_torsion_angles_sin_cos.shape[:-2]						,					14     ),
         torsion_angles_mask,
     ]						,					dim=-1						,					)

 return template_angle_feat
def 		UpperCamelCase_(    snake_case__:						Dict[str, torch.Tensor]						,					snake_case__:						torch.types.Number						,					snake_case__:						torch.types.Number						,					snake_case__:						int						,					snake_case__:						bool = False						,					snake_case__:						float = 1e-20						,					snake_case__:						float = 1e8						,					) ->				torch.Tensor:
 UpperCAmelCase__           =       batch['template_pseudo_beta_mask']
 UpperCAmelCase__           =       template_mask[..., None] * template_mask[..., None, :]

 # Compute distogram (this seems to differ slightly from Alg. 5)
 UpperCAmelCase__           =       batch['template_pseudo_beta']
 UpperCAmelCase__           =       torch.sum((tpb[..., None, :] - tpb[..., None, :, :]) ** 2						,					dim=-1						,					keepdim=snake_case__     )
 UpperCAmelCase__           =       torch.linspace(snake_case__						,					snake_case__						,					snake_case__						,					device=tpb.device     ) ** 2
 UpperCAmelCase__           =       torch.cat([lower[1:], lower.new_tensor([inf]     )]						,					dim=-1     )
 UpperCAmelCase__           =       ((dgram > lower) * (dgram < upper)).type(dgram.dtype     )

 UpperCAmelCase__           =       [dgram, template_mask_ad[..., None]]

 UpperCAmelCase__           =       nn.functional.one_hot(
     batch['template_aatype']						,					rc.restype_num + 2						,					)

 UpperCAmelCase__           =       batch['template_aatype'].shape[-1]
 to_concat.append(aatype_one_hot[..., None, :, :].expand(*aatype_one_hot.shape[:-2]						,					snake_case__						,					-1						,					-1     )     )
 to_concat.append(aatype_one_hot[..., None, :].expand(*aatype_one_hot.shape[:-2]						,					-1						,					snake_case__						,					-1     )     )

 UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__           =       [rc.atom_order[a] for a in ['N', 'CA', 'C']]
 UpperCAmelCase__           =       Rigid.make_transform_from_reference(
     n_xyz=batch['template_all_atom_positions'][..., n, :]						,					ca_xyz=batch['template_all_atom_positions'][..., ca, :]						,					c_xyz=batch['template_all_atom_positions'][..., c, :]						,					eps=snake_case__						,					)
 UpperCAmelCase__           =       rigids.get_trans()[..., None, :, :]
 UpperCAmelCase__           =       rigids[..., None].invert_apply(snake_case__     )

 UpperCAmelCase__           =       torch.rsqrt(eps + torch.sum(rigid_vec**2						,					dim=-1     )     )

 UpperCAmelCase__           =       batch['template_all_atom_mask']
 UpperCAmelCase__           =       t_aa_masks[..., n] * t_aa_masks[..., ca] * t_aa_masks[..., c]
 UpperCAmelCase__           =       template_mask[..., None] * template_mask[..., None, :]

 UpperCAmelCase__           =       inv_distance_scalar * template_mask_ad
 UpperCAmelCase__           =       rigid_vec * inv_distance_scalar[..., None]

 if not use_unit_vector:
  UpperCAmelCase__           =       unit_vector * 0.0

 to_concat.extend(torch.unbind(unit_vector[..., None, :]						,					dim=-1     )     )
 to_concat.append(template_mask_ad[..., None]     )

 UpperCAmelCase__           =       torch.cat(snake_case__						,					dim=-1     )
 UpperCAmelCase__           =       act * template_mask_ad[..., None]

 return act
def 		UpperCamelCase_(    snake_case__:						Dict[str, torch.Tensor]     ) ->				torch.Tensor:
 UpperCAmelCase__           =       nn.functional.one_hot(batch['extra_msa']						,					23     )
 UpperCAmelCase__           =       [
     msa_ahot,
     batch['extra_has_deletion'].unsqueeze(-1     ),
     batch['extra_deletion_value'].unsqueeze(-1     ),
 ]
 return torch.cat(snake_case__						,					dim=-1     )
def 		UpperCamelCase_(    snake_case__:						Rigid						,					snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor						,					) ->				Rigid:
 # [*, N, 8, 4, 4]
 UpperCAmelCase__           =       rrgdf[aatype, ...]

 # [*, N, 8] transformations, i.e.
 #   One [*, N, 8, 3, 3] rotation matrix and
 #   One [*, N, 8, 3]    translation matrix
 UpperCAmelCase__           =       r.from_tensor_axa(snake_case__     )

 UpperCAmelCase__           =       alpha.new_zeros((*((1,) * len(alpha.shape[:-1]     )), 2)     )
 UpperCAmelCase__           =       1

 # [*, N, 8, 2]
 UpperCAmelCase__           =       torch.cat([bb_rot.expand(*alpha.shape[:-2]						,					-1						,					-1     ), alpha]						,					dim=-2     )

 # [*, N, 8, 3, 3]
 # Produces rotation matrices of the form:
 # [
 #   [1, 0  , 0  ],
 #   [0, a_2,-a_1],
 #   [0, a_1, a_2]
 # ]
 # This follows the original code rather than the supplement, which uses
 # different indices.

 UpperCAmelCase__           =       alpha.new_zeros(default_r.get_rots().get_rot_mats().shape     )
 UpperCAmelCase__           =       1
 UpperCAmelCase__           =       alpha[..., 1]
 UpperCAmelCase__           =       -alpha[..., 0]
 UpperCAmelCase__           =       alpha

 UpperCAmelCase__           =       default_r.compose(Rigid(Rotation(rot_mats=snake_case__     )						,					snake_case__     )     )

 UpperCAmelCase__           =       all_frames[..., 5]
 UpperCAmelCase__           =       all_frames[..., 6]
 UpperCAmelCase__           =       all_frames[..., 7]

 UpperCAmelCase__           =       all_frames[..., 4]
 UpperCAmelCase__           =       chia_frame_to_bb.compose(snake_case__     )
 UpperCAmelCase__           =       chia_frame_to_bb.compose(snake_case__     )
 UpperCAmelCase__           =       chia_frame_to_bb.compose(snake_case__     )

 UpperCAmelCase__           =       Rigid.cat(
     [
         all_frames[..., :5],
         chia_frame_to_bb.unsqueeze(-1     ),
         chia_frame_to_bb.unsqueeze(-1     ),
         chia_frame_to_bb.unsqueeze(-1     ),
     ]						,					dim=-1						,					)

 UpperCAmelCase__           =       r[..., None].compose(snake_case__     )

 return all_frames_to_global







def 		UpperCamelCase_(    snake_case__:						Rigid						,					snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor						,					snake_case__:						torch.Tensor						,					) ->				torch.Tensor:
 # [*, N, 14]
 UpperCAmelCase__           =       group_idx[aatype, ...]

 # [*, N, 14, 8]
 UpperCAmelCase__           =       nn.functional.one_hot(
     snake_case__						,					num_classes=default_frames.shape[-3]						,					)

 # [*, N, 14, 8]
 UpperCAmelCase__           =       r[..., None, :] * group_mask_one_hot

 # [*, N, 14]
 UpperCAmelCase__           =       t_atoms_to_global.map_tensor_fn(lambda snake_case__     : torch.sum(snake_case__						,					dim=-1     )     )

 # [*, N, 14, 1]
 UpperCAmelCase__           =       atom_mask[aatype, ...].unsqueeze(-1     )

 # [*, N, 14, 3]
 UpperCAmelCase__           =       lit_positions[aatype, ...]
 UpperCAmelCase__           =       t_atoms_to_global.apply(snake_case__     )
 UpperCAmelCase__           =       pred_positions * atom_mask

 return pred_positions






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_UpperCamelCase     =    {
    '''configuration_fsmt''': ['''FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''FSMTConfig'''],
    '''tokenization_fsmt''': ['''FSMTTokenizer'''],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _UpperCamelCase     =    ['''FSMTForConditionalGeneration''', '''FSMTModel''', '''PretrainedFSMTModel''']


if TYPE_CHECKING:
 from .configuration_fsmt import FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP, FSMTConfig
 from .tokenization_fsmt import FSMTTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_fsmt import FSMTForConditionalGeneration, FSMTModel, PretrainedFSMTModel

else:
 import sys

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



