text: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  """
  Utility code for converting between protein representations.
  
  Note that several methods here are no longer used in any of the training routines.
  However, they were quite useful to us during the course of research,
  so we are releasing them here in case they help others.
  """
  import collections
  import os
  import os.path as osp
  import pickle
  import random
  from itertools import product
  from multiprocessing import Pool
  
  import numpy as np
  import pandas as pd
  
  import gemmi
  from amino_acid_config import kvs, res_atoms, res_children, res_chis, res_parents
  from config import MMCIF_PATH, ROTAMER_LIBRARY_PATH
  from constants import atom_names, residue_names
  from math_utils import rotate_v1_v2, rotate_v1_v2_vec
  
  
  def parse_dense_format(node_embed):
      """
      In protein-ebm, we represent amino acids in two different formats.
      This method converts from the dense format to a sparse format.
      
      ===============
      ==== Dense ====
      ===============
      The dense format represents a protein using a is a D x 6 dimensional represention.
      Each 6 dimensional vector represents an atom, following this scheme:
          [1]: amino acid identity of the amino acid the atom is part of (residue_idx)
          [2]: element identity of the amino acid the atom is part of (atom_idx)
          [3]: positional location of atom in the amino acid (atom_num)
          [4..6]: x,y,z coordinates
      The dense format is useful for feeding data into a neural network.
  
      ===============
      ==== Sparse ===
      ===============
      The sparse format represents a data based on its topology (parent/child/etc).
      It follows this scheme:
          amino_name: amino acid to substitue
          par: A N x 20 encoding of the relative offset of the parent of each atom. For example,
                  the amino acid glycine would be represented as [-18 -1 -1 -1 0, ...]
          child: A N x 20 encoding of the child of each atom. For example, the amino acid glycine
                  would be represented as [1 1 18 0 0 0 ..]
          pos_exist: A N x 20 mask encoding of which atoms are valid for each amino acid so for
          example the amino acid glycine would be represented as [1 1 1 1 0 0 ...]
          chi_valid: A N x 5 mask encoding which chi angles are valid, so for example glycine would
          be represented as [0 0 0 0 0]
          pos: A N x 20 x 3 encoding the (x, y, z) coordinates of each atom per amino acid in a protein
          i: amino acid position to substitute
          sequence_map: map from amino acid to structure
          rotate_matrix: matrix of rotation to amino acid position
      This format is easier for manipulating the proteins, e.g changing the rotamers
      during negative sampling. 
  
      See comments in the implementation below for more details.
      """
  
      # The input is a list of atoms. We keep track of how many we have processed.
      start = 0
  
      # Construct amino acid-level information from the atomic inputs
      # Each amino acid is described on the atomic-level by 20-dim lists
      pars = []  # ordinal distance of parent atoms
      childs = []  # ordinal distance of cildren atoms
      pos = []  # 3d translations of each atom
      pos_exists = []  # whether a position exists or not
      residues = []  # the name of the amino acid
      chis_valid = []  # a 20-dim list describing which atoms are part of the chi angle
  
      # consume all of the atoms in the input
      while start < node_embed.shape[0]:
          idx = int(node_embed[start, 0])
          residue = residue_names[idx]
  
          # Get the parent and child representation (see amino_acid_config.py)
          par = res_parents[residue].copy()
          child = res_children[residue].copy()
          n = len(par)
  
          # 20-dim mask of which positions encode meaningful values
          pos_exist = [1] * n + [0] * (20 - n)  # this is the mask
  
          # pad up to 20-dim with 0s
          par = par + [0] * (20 - n)
          child = child + [0] * (20 - len(child))
  
          # x,y,z coordinates for each of the atoms in the amino acid, padded to 20-dim
          pos_temp = np.concatenate(
              [node_embed[start : start + n, -3:], np.zeros((20 - n, 3))], axis=0
          )
  
          # If we can fit these n atom in, then record the information
          if start + n <= node_embed.shape[0]:
              pars.append(par)
              childs.append(child)
              pos.append(pos_temp)
              pos_exists.append(pos_exist)
              chis = res_chis[residue]
              chis_valid.append([1] * len(chis) + [0] * (20 - len(chis)))
              residues.append(residue.lower())
  
          # All atoms from start <-> start+n should belong to the same amino acid
          if not (node_embed[start : start + n, 0] == idx).all():
              return None, None, None, None, None, None
  
          # keep track of number of atoms consumeed
          start = start + n
  
      # Don't proceess single amino acid prorteins
      if len(pos) < 2:
          return None, None, None, None, None, None
  
      # Wrap the results in numpy arrays
      pars, childs, pos, pos_exists, chis_valid = (
          np.array(pars),
          np.array(childs),
          np.stack(pos, axis=0),
          np.array(pos_exists),
          np.array(chis_valid),
      )
  
      # The code above assumes that each nitrogen is connected to previous carbon
      # and each carbon is connected to the next nitrogen. This is not the case
      # for the N-terminus and C-terminus, so we need to override those cases.
      pars[0, 0] = 0
      childs[-1, 2] = 0
  
      # return the new encoding in amino acid form
      return pars, childs, pos, pos_exists, residues, chis_valid
  
  
  def reencode_dense_format(node_embed, pos_new, pos_exist):
      """Updates x,y,z positions in dense encoding with new positions"""
      node_embed_new = node_embed.copy()
      pos_mask = pos_exist.astype(np.bool)
      elem_num = pos_mask.sum()
      node_embed_new[:elem_num, -3:] = pos_new[pos_mask]
  
      return node_embed_new
  
  
  def cif_to_embed(cif_file, ix=None, parse_skip=False):
      """
      Parses a CIF file into a more convenient representation.
  
      # Embedding format for nodes:
      # 'one hot amino acid' amino type of molecule
      # 'x, y, z' positional encoding
      # 'one hot representation of atom type', either C, CA, N, O,
  
      """
      st = gemmi.read_structure(cif_file)
  
      results = []
      skips = []
      for model in st:
          for i, chain in enumerate(model):
  
              if (ix is not None) and (ix != i):
                  continue
  
              atoms = []
              node_embeddings = []
              for j, residue in enumerate(chain):
                  translation = []
  
                  if residue.name not in residue_names:
                      # Skip over any structure that contains nucleotides
                      if residue.name in ["DA", "DC", "DG", "DT"]:
                          return None, None
                      else:
                          continue
  
                  residue_counter = 0
                  namino_elements = len(res_parents[residue.name])
                  amino_atoms = res_atoms[residue.name]
  
                  residue_atoms = []
                  residue_embed = []
  
                  # reisdue object contains information about the residue, including identity
                  # and spatial coordiantes for atoms in the residue. We parse this into a
                  # dense encoding, for feeding into a neural network.
                  node_embed = parse_residue_embed(residue)
  
                  if len(node_embed) == 0:
                      skips.append(j)
  
                  node_embeddings.extend(node_embed)
  
              node_embeddings = np.array(node_embeddings)
  
              result = (node_embeddings,)
              results.append(result)
  
      if parse_skip:
          return st, results, skips
      else:
          return st, results
  
  
  def vis_cif(cif_path, im_path):
      import pymol
      from pymol import cmd
  
      cmd.load(cif_path, "mov")
      cmd.zoom()
      cmd.png(im_path, 300, 200)
  
  
  def compute_chi_angle_st(st, ix):
      angles = []
  
      num = int(ix)
      chain_counter = 0
      for model in st:
          for chain in model:
              if num != chain_counter:
                  chain_counter += 1
                  continue
              else:
                  for residue in chain:
                      if residue.name in residue_names:
                          chi_angles = compute_chi_angle_residue(residue)
                          if chi_angles is not None:
                              angles.append(chi_angles)
                  return angles
  
  
  def compute_chi_angle_residue(residue):
      # look up the atoms that are used for computing the chi angles.
      chi_angles_atoms = kvs[residue.name]
  
      angles = []
  
      try:
          for chi_angles_atom in chi_angles_atoms:
              atoms = chi_angles_atom.split("-")
              pos = []
  
              for atom in atoms:
                  # In some cases, amino acid side chains are listed with CD1 instead of CD
                  if atom == "CD":
                      if "CD" not in residue:
                          atom = residue["CD1"]
                      else:
                          atom = residue[atom]
                  else:
                      atom = residue[atom]
  
                  pos.append((atom.pos.x, atom.pos.y, atom.pos.z))
  
              pos = np.array(pos)
              diff_vec = pos[2] - pos[1]
              # Compute the axis in which we are computing the dihedral angle
              diff_vec_normalize = diff_vec / np.linalg.norm(diff_vec)
  
              diff_bot = pos[0] - pos[1]
              diff_top = pos[3] - pos[2]
  
              # Now project diff_bot and diff_top to be on the plane
              diff_bot = diff_bot - diff_bot.dot(diff_vec_normalize) * diff_vec_normalize
              diff_top = diff_top - diff_top.dot(diff_vec_normalize) * diff_vec_normalize
  
              diff_bot_normalize = diff_bot / np.linalg.norm(diff_bot)
              diff_top_normalize = diff_top / np.linalg.norm(diff_top)
  
              # Compute the dot product for cos and cross product for sin
              sin = (np.cross(diff_bot_normalize, diff_top_normalize) * diff_vec_normalize).sum(
                  axis=1
              )
              cos = diff_bot_normalize.dot(diff_top_normalize)
  
              # print("trig value ", sin, cos, np.linalg.norm([sin, cos]))
              angle = np.arctan2(sin, cos)
              # print("angle ", angle)
  
              angles.append(angle)
      except Exception as e:
          return None
  
      return angles
  
  
  def parse_cif(path):
      base_folder, f = osp.split(path)
      base_name, *junk = f.split(".")
  
      st, infos = cif_to_embed(path)
  
      if infos is not None:
          for i, info in enumerate(infos):
              pickle_file = osp.join(base_folder, "{}.{}.p".format(base_name, i))
              pickle.dump(info, open(pickle_file, "wb"))
  
      return None
  
  
  def script_parse_cif():
      mmcif_path = osp.join(MMCIF_PATH, "mmCIF")
      files = []
      dirs = os.listdir(mmcif_path)
      pool = Pool()
  
      for d in dirs:
          directory = osp.join(mmcif_path, d)
          d_files = os.listdir(directory)
          files_tmp = [osp.join(directory, d_file) for d_file in d_files if ".cif" in d_file]
          files.extend(files_tmp)
  
      pool.map(parse_cif, files)
  
  
  def clean_cif():
      mmcif_path = osp.join(MMCIF_PATH, mmCIF)
      dirs = os.listdir(mmcif_path)
  
      for d in dirs:
          directory = osp.join(mmcif_path, d)
          d_files = os.listdir(directory)
          files_tmp = [osp.join(directory, d_file) for d_file in d_files if ".p" in d_file]
  
          for f in files_tmp:
              os.remove(f)
  
  
  def recorrect_name(name):
      if (name[-1]).isdigit() and name[-1] == "1":
          return name[:-1]
      elif not (name[-1].isdigit()):
          return name + "1"
      else:
          return name
  
  
  def _parse_residue(residue):
      """Obtains a sparse representation of residue from gemmi"""
  
      # list of atoms in the residue (e.g. N-CA-C-O)
      atoms = res_atoms[residue.name]
  
      # ordinal encoding of how far away the parents are
      parents = res_parents[residue.name]
  
      # ordinal encoding of how far away the children are
      children = res_children[residue.name]
  
      # atoms belonging to chi anglse
      chis = res_chis[residue.name]
  
      # accumulate the xyz postions of the atoms, and node_embed encodings
      pos, node_embeds = [], []
  
      residue_counter = 0
      for atom in atoms:
          if atom in residue:
              atom = residue[atom]
          elif recorrect_name(atom) in residue:
              atom = residue[recorrect_name(atom)]
          else:
              return None
  
          pos.append((atom.pos.x, atom.pos.y, atom.pos.z))
          node_embeds.append(
              (
                  residue_names.index(residue.name),
                  atom_names.index(atom.element.name),
                  residue_counter,
                  atom.pos.x,
                  atom.pos.y,
                  atom.pos.z,
              )
          )
          residue_counter = residue_counter + 1
  
      # 20-dim mask for each residue for atom existence
      exist = [1] * len(parents) + [0] * (20 - len(parents))
  
      # pad the parents and children to 20-dim
      parents = parents + [0] * (20 - len(parents))
      children = children + [0] * (20 - len(children))
  
      # place the x,y,z coordinates into a numpy array
      pos_fill = np.zeros((20, 3))
      pos_fill[: len(pos)] = pos
  
      # pad the chi angles
      chis = [1] * len(chis) + [0] * (5 - len(chis))
  
      # return the new representation
      return parents, children, pos_fill, exist, chis, node_embeds
  
  
  # shorthand methods for the above, since logic is the same
  def parse_residue(residue):
      ret = _parse_residue(residue, 0)
  
      if ret:
          parents, children, pos_fill, exist, chis, _, _ = ret
          return parents, children, pos_fill, exist, chis
      else:
          return None, None, None, None, None
  
  
  def parse_residue_embed(residue):
      ret = _parse_residue(residue)
  
      if ret:
          _, _, _, _, _, node_embeds = ret
          return node_embeds
      else:
          return []
  
  
  def flatten(arr):
      return arr.reshape((-1, *arr.shape[2:]))
  
  
  def rotate_dihedral_fast(a, p, c, pos, pos_e, ch, chv, idx):
      """
      Where as rotate_dihedral(...) rotates all amino acids in the batch by some angle,
      this function just rotates a single amino acid in a protein.
      """
      pos = pos.copy()
      ai, pi, ci, pos_i, pos_ei, chi, chvi = (
          a[idx - 1 : idx + 1],
          p[idx - 1 : idx + 1],
          c[idx - 1 : idx + 1],
          pos[idx - 1 : idx + 1],
          pos_e[idx - 1 : idx + 1],
          ch[idx - 1 : idx + 1],
          chv[idx - 1 : idx + 1],
      )
      pnew = rotate_dihedral(ai, pi, ci, pos_i, pos_ei, chi, chvi)
      pos[idx - 1 : idx + 1] = pnew
  
      return pos
  
  
  def rotate_dihedral(angles, par, child, pos, pos_exist, chis, chi_valid):
      """Rotate a protein representation by a set of dihedral angles:
          N represents the number of amino acids in the batch, 20 is the number of atoms.
          angles: N x 20 set of angles to rotate each atom by
          par: A N x 20 encoding of the relative offset of the parent of each atom. For example,
              the amino acid glycine would be represented at [-18 -1 -1 -1 0, ...]
          child: A N x 20 encoding of the child of each atom. For example, the amino acid glycine
              would be represented as [1 1 18 0 0 0 ..]
          pos_exist: A N x 20 mask encoding of which atoms are valid for each amino acid so for
              example the amino acid glycine would be represented as [1 1 1 1 0 0 ...]
          chis: A N x 20 representation of the existing chi angles
          chi_valid: A N x 5 mask encoding which chi angles are valid, so for example glycine would
          be represented as [0 0 0 0 0]
          """
  
      angles = angles / 180 * np.pi
      chis = chis / 180 * np.pi
      pos_orig = pos
      pos = pos.copy()
  
      for i in range(4):
          # There are a maximum of 5 chi angles
          p2 = pos[:, 4 + i]
          index = np.tile(4 + i, (pos.shape[0], 1)) + par[:, 4 + i : 5 + i]
          # print("index, pos shape ", index.shape, pos.shape)
          p1 = np.take_along_axis(pos, index[:, :, None], axis=1)[:, 0, :]
  
          rot_angle = chis[:, i] - angles[:, 4 + i]
  
          diff_vec = p2 - p1
          diff_vec_normalize = diff_vec / (np.linalg.norm(diff_vec, axis=1, keepdims=True) + 1e-10)
  
          # Rotate all subsequent points by the rotamor angle with the defined line where normalize on the origin
          rot_points = pos[:, 5 + i :].copy() - p1[:, None, :]
  
          par_points = (rot_points * diff_vec_normalize[:, None, :]).sum(
              axis=2, keepdims=True
          ) * diff_vec_normalize[:, None, :]
          perp_points = rot_points - par_points
  
          perp_points_norm = np.linalg.norm(perp_points, axis=2, keepdims=True) + 1e-10
          perp_points_normalize = perp_points / perp_points_norm
  
          a3 = np.cross(diff_vec_normalize[:, None, :], perp_points_normalize)
  
          rot_points = (
              perp_points * np.cos(rot_angle)[:, None, None]
              + np.sin(rot_angle)[:, None, None] * a3 * perp_points_norm
              + par_points
              + p1[:, None, :]
          )
  
          rot_points[np.isnan(rot_points)] = 10000
  
          # Only set the points that vald chi angles
          first_term = rot_points * chi_valid[:, i : i + 1, None]
          second_term = pos[:, 5 + i :] * (1 - chi_valid[:, i : i + 1, None])
  
          pos[:, 5 + i :] = first_term + second_term
  
      return pos
  
  
  def compute_dihedral(par, child, pos, pos_exist, reshape=True):
      """Compute the dihedral angles of all atoms in a structure
          par: A N x 20 encoding of the relative offset of the parent of each atom. For example,
               the amino acid glycine would be represented at [-18 -1 -1 -1 0, ...]
          child: A N x 20 encoding of the child of each atom. For example, the amino acid glycine
                  would be represented as [1 1 18 0 0 0 ..]
          pos_exist: A N x 20 mask encoding of which atoms are valid for each amino acid so for
          pos: A N x 20 x 3 encoding the (x, y, z) coordinates of each atom per amino acid in a protein
      """
      par, child, pos, pos_exist = flatten(par), flatten(child), flatten(pos), flatten(pos_exist)
      # pos[~pos_exist] = 0.1
      idx = np.arange(par.shape[0])
  
      child_idx = idx + child
      child_pos = pos[child_idx, :].copy()
      up_edge_idx = idx + par
      up_edge_pos = pos[up_edge_idx, :].copy()
      parent_idx = up_edge_idx + par[up_edge_idx]
      parent_pos = pos[parent_idx, :].copy()
  
      # The dihedral angle is given by parent_pos -> up_edge_pos -> pos -> child_pos
      p0, p1, p2, p3 = parent_pos, up_edge_pos, pos, child_pos
  
      p23 = p3 - p2
      p12 = p2 - p1
      p01 = p1 - p0
  
      n1 = np.cross(p01, p12)
      n2 = np.cross(p12, p23)
  
      n1 = n1 / (np.linalg.norm(n1, axis=1, keepdims=True) + 1e-10)
      n2 = n2 / (np.linalg.norm(n2, axis=1, keepdims=True) + 1e-10)
  
      sin = (np.cross(n1, n2) * p12 / (np.linalg.norm(p12, axis=1, keepdims=True) + 1e-10)).sum(
          axis=1
      )
      cos = (n1 * n2).sum(axis=1)
  
      angle = np.arctan2(sin, cos)
  
      # Convert the angles to -180 / 180
      angle = angle / np.pi * 180
  
      if reshape:
          angle = angle.reshape((-1, 20))
  
      return angle
  
  
  # The code below does sampling from the dunbrack library
  def sample_df(df, uniform=False, sample=1):
      """Sample from rotamer library based off gaussian on nearby slots"""
      cum_prob = df["Probabil"].cumsum()
      cutoff = np.random.uniform(0, cum_prob.max(), (sample,))
      ixs = cum_prob.searchsorted(cutoff)
  
      if uniform:
          ix = cum_prob.searchsorted(0.99)
          if ix == 0:
              ixs = [0] * sample
          else:
              ixs = np.random.randint(ix, size=(sample,))
  
      chis_list = []
      for ix in ixs:
          means = df[["chi{}Val".format(i) for i in range(1, 5)]][ix : ix + 1].to_numpy()
          std = df[["chi{}Sig".format(i) for i in range(1, 5)]][ix : ix + 1].to_numpy()
          chis = std[0] * np.random.normal(0, 1, (4,)) + means[0]
          chis[chis > 180] = chis[chis > 180] - 360
          chis[chis < -180] = chis[chis < -180] + 360
          chis_list.append(chis)
  
      if sample == 1:
          chis_list = chis_list[0]
  
      return chis_list
  
  
  def sample_weighted_df(dfs, weights_array, uniform=False):
      """sample from rotamer library based off a weighted gaussian from nearby slots"""
      n = min(df["Probabil"].to_numpy().shape[0] for df in dfs)
      dfs = [df[:n].sort_values("chi1Val") for df in dfs]
  
      probs = []
      for weight, df in zip(weights_array, dfs):
          probs.append(df["Probabil"].to_numpy()[:n] * weight)
  
      probs = np.sum(np.array(probs), axis=0) / 100
      cum_prob = np.cumsum(probs)
      cutoff = np.random.uniform(0, cum_prob.max())
      ix = np.searchsorted(cum_prob, cutoff)
  
      if uniform:
          # ix = np.searchsorted(cum_prob, 0.99)
          if ix == 0:
              ix = 0
          else:
              ix = np.random.randint(ix)
  
      means = [
          weight * df[["chi{}Val".format(i) for i in range(1, 5)]].to_numpy()[:n]
          for weight, df in zip(weights_array, dfs)
      ]
      std = [
          weight * df[["chi{}Sig".format(i) for i in range(1, 5)]].to_numpy()[:n]
          for weight, df in zip(weights_array, dfs)
      ]
  
      means = np.sum(means, axis=0) / 100.0
      std = np.sum(std, axis=0) / 100
  
      chis = std[ix] * np.random.normal(0, 1, (4,)) + means[ix]
  
      # chis = (360 - chis) % 360
  
      chis[chis > 180] = chis[chis > 180] - 360
      chis[chis < -180] = chis[chis < -180] + 360
      return chis
  
  
  def discrete_angle_to_bucket(ang):
      assert isinstance(ang, int)
      assert ang % 10 == 0
      assert -180 <= ang < 180
      return (ang + 180) // 10
  
  
  def get_rotind(r1, r2, r3, r4):
      return 1000000 * r1 + 10000 * r2 + 100 * r3 + r4
  
  
  QuadrantData = collections.namedtuple(
      "QuadrantData",
      ["chimeans", "chisigmas", "probs", "meanprobs", "cumprobs", "exists", "rotinds"],
  )
  
  
  def _preprocess_db(db, name):
      df = db[name]
  
      bucketed_data = [[{} for _1 in range(36)] for _2 in range(36)]
  
      df_rows = df.to_dict("records")
      for row in df_rows:
          phi, psi = row["Phi"], row["Psi"]
  
          wraparound = False
          if phi == 180:
              wraparound = True
              phi = -180
          if psi == 180:
              wraparound = True
              psi = -180
  
          phi_bucket, psi_bucket = discrete_angle_to_bucket(phi), discrete_angle_to_bucket(psi)
  
          rotind = get_rotind(row["r1"], row["r2"], row["r3"], row["r4"])
          chimeans = np.array([row[f"chi{i}Val"] for i in range(1, 5)])
          chisigmas = np.array([row[f"chi{i}Sig"] for i in range(1, 5)])
          prob = row["Probabil"]
  
          bucket = bucketed_data[phi_bucket][psi_bucket]
          bucket_data = (chimeans, chisigmas, prob)
  
          if wraparound:
              assert (
                  (bucket[rotind][0] == bucket_data[0]).all()
                  and (bucket[rotind][1] == bucket_data[1]).all()
                  and (bucket[rotind][2] == bucket_data[2])
              )
  
          else:
              bucket[rotind] = bucket_data
  
      quadrant_data = [[None for _1 in range(36)] for _2 in range(36)]
  
      for lower_phi_bucket in range(36):
          for lower_psi_bucket in range(36):
              upper_phi_bucket = (lower_phi_bucket + 1) % 36
              upper_psi_bucket = (lower_psi_bucket + 1) % 36
  
              quadrants = [
                  bucketed_data[lower_phi_bucket][lower_psi_bucket],
                  bucketed_data[upper_phi_bucket][lower_psi_bucket],
                  bucketed_data[lower_phi_bucket][upper_psi_bucket],
                  bucketed_data[upper_phi_bucket][upper_psi_bucket],
              ]
  
              rotinds = np.array(
                  sorted(set().union(*[set(quadrant.keys()) for quadrant in quadrants])),
                  dtype=np.int,
              )
  
              assert len(rotinds) > 0
  
              exists = np.zeros((len(rotinds), 4), dtype=np.bool)
              probs = np.zeros((len(rotinds), 4), dtype=np.float64)
              chimeans = np.zeros((len(rotinds), 4, 4), dtype=np.float64)
              chisigmas = np.zeros((len(rotinds), 4, 4), dtype=np.float64)
  
              for i, rotind in enumerate(rotinds):
                  for qid, quadrant in enumerate(quadrants):
                      if rotind not in quadrant:
                          continue
  
                      quadrant_chimeans, quadrant_chisigmas, quadrant_prob = quadrant[rotind]
  
                      exists[i, qid] = True
                      probs[i, qid] = quadrant_prob
                      chimeans[i, qid] = quadrant_chimeans
                      chisigmas[i, qid] = quadrant_chisigmas
  
              meanprobs = probs.mean(1)
              order = np.argsort(-meanprobs, kind="stable")
              meanprobs = meanprobs[order]
              cumprobs = np.cumsum(meanprobs)
  
              assert np.abs(cumprobs[-1] - 1) < 1e-5
  
              quadrant_data[lower_phi_bucket][lower_psi_bucket] = QuadrantData(
                  chimeans=chimeans[order],
                  chisigmas=chisigmas[order],
                  probs=probs[order],
                  exists=exists[order],
                  rotinds=rotinds[order],
                  meanprobs=meanprobs,
                  cumprobs=cumprobs,
              )
  
      return quadrant_data
  
  
  _PREPROCESS_DB_CACHE = {}
  
  
  def preprocess_db(db, name):
      key = (id(db), name)
      val = _PREPROCESS_DB_CACHE.get(key)
  
      if val is None:
          val = _preprocess_db(db, name)
          _PREPROCESS_DB_CACHE[key] = val
  
      return val
  
  
  def get_quadrant_data_with_interpolated_weights(db, name, phi, psi):
      lower_phi, lower_psi = int(phi // 10) * 10, int(psi // 10) * 10
      upper_phi, upper_psi = lower_phi + 10, lower_psi + 10
      lower_phi_bucket, lower_psi_bucket = (
          discrete_angle_to_bucket(lower_phi),
          discrete_angle_to_bucket(lower_psi),
      )
  
      quadrant_data = preprocess_db(db, name)[lower_phi_bucket][lower_psi_bucket]
  
      weights = np.array(
          [
              (10 - (phi - lower_phi)) * (10 - (psi - lower_psi)),
              (10 - (upper_phi - phi)) * (10 - (psi - lower_psi)),
              (10 - (phi - lower_phi)) * (10 - (upper_psi - psi)),
              (10 - (upper_phi - phi)) * (10 - (upper_psi - psi)),
          ]
      )
  
      sum_existing_weights = (weights[np.newaxis, :] * quadrant_data.exists).sum(1)
      effective_weights = weights[np.newaxis, :] / sum_existing_weights[:, np.newaxis]
  
      return quadrant_data, effective_weights
  
  
  def exhaustive_sample(db, phi, psi, name, tresh=0.99, chi_mean=False):
      """sample  a set of discrete possibilitys for rotamers following protocol used in Rosetta"""
      quadrant_data, weights = get_quadrant_data_with_interpolated_weights(db, name, phi, psi)
  
      chimeans = (quadrant_data.chimeans * weights[:, :, np.newaxis]).sum(1)
      chisigmas = (quadrant_data.chisigmas * weights[:, :, np.newaxis]).sum(1)
  
      cumprobs = quadrant_data.cumprobs
      search_limit = (np.searchsorted(cumprobs, tresh) + 1) if tresh < (1 - 1e-6) else len(cumprobs)
      assert search_limit <= len(cumprobs)
  
      chimeans = chimeans[:search_limit]
      chisigmas = chisigmas[:search_limit]
  
      sigma_masks = np.array(list(product([-1, 0, 1], [-1, 0, 1], [0], [0])), dtype=np.float64)
  
      if chi_mean:
          return list(chimeans)
  
      angles = chimeans[:, np.newaxis, :] + (
          chisigmas[:, np.newaxis, :] * sigma_masks[np.newaxis, :, :]
      )
      angles = angles.reshape(-1, 4)
  
      for _ in range(2):
          angles[angles >= 180] = angles[angles >= 180] - 360
          angles[angles < -180] = angles[angles < -180] + 360
  
      return list(angles)
  
  
  def _sample_from_cumprobs(cumprobs, n, uniform):
      if uniform:
          return np.random.randint(len(cumprobs), size=n)
      else:
          searchvals = np.random.uniform(low=0.0, high=cumprobs[-1], size=n)
          indices = np.searchsorted(cumprobs, searchvals)
          assert (indices < len(cumprobs)).all()
          return indices
  
  
  def interpolated_sample_normal(db, phi, psi, name, n, uniform=False):
      quadrant_data, weights = get_quadrant_data_with_interpolated_weights(db, name, phi, psi)
  
      chimeans = (quadrant_data.chimeans * weights[:, :, np.newaxis]).sum(1)
      chisigmas = (quadrant_data.chisigmas * weights[:, :, np.newaxis]).sum(1)
  
      cumprobs = quadrant_data.cumprobs
      sample_indices = _sample_from_cumprobs(cumprobs=cumprobs, n=n, uniform=uniform)
      assert sample_indices.shape == (n,)
  
      chimeans = chimeans[sample_indices]
      chisigmas = chisigmas[sample_indices]
  
      angles = chimeans + np.random.randn(n, 4) * chisigmas
  
      for _ in range(2):
          angles[angles >= 180] = angles[angles >= 180] - 360
          angles[angles < -180] = angles[angles < -180] + 360
  
      return list(angles)
  
  
  def mixture_sample_normal(db, phi, psi, name, n, uniform=False):
      quadrant_data, weights = get_quadrant_data_with_interpolated_weights(db, name, phi, psi)
  
      chimeans = quadrant_data.chimeans
      chisigmas = quadrant_data.chisigmas
  
      cumprobs = quadrant_data.cumprobs
      sample_indices = _sample_from_cumprobs(cumprobs=cumprobs, n=n, uniform=uniform)
      assert sample_indices.shape == (n,)
  
      angles = np.zeros((n, 4))
  
      for aidx in range(n):
          i = sample_indices[aidx]
          quadrant = np.random.choice(4, p=weights[i])
          chimean = chimeans[i, quadrant]
          chisigma = chisigmas[i, quadrant]
          angles[aidx] = chimean + np.random.randn(4) * chisigma
  
      for _ in range(2):
          angles[angles >= 180] = angles[angles >= 180] - 360
          angles[angles < -180] = angles[angles < -180] + 360
  
      return list(angles)
  
  
  def sample_rotomor_angle_db(db, phi, psi, name, uniform=False, n=1):
      df = db[name]
      lower_phi = (phi // 10) * 10
      upper_phi = lower_phi + 10
  
      lower_psi = (psi // 10) * 10
      upper_psi = lower_psi + 10
  
      weights = [
          (10 - (phi - lower_phi)) * (10 - (psi - lower_psi)),
          (10 - (upper_phi - phi)) * (10 - (psi - lower_psi)),
          (10 - (phi - lower_phi)) * (10 - (upper_psi - psi)),
          (10 - (upper_phi - phi)) * (10 - (upper_psi - psi)),
      ]
      weights_array = weights
  
      weights = np.cumsum(weights)
  
      dfs = [
          df[(df.Phi == lower_phi) & (df.Psi == lower_psi)],
          df[(df.Phi == upper_phi) & (df.Psi == lower_psi)],
          df[(df.Phi == lower_phi) & (df.Psi == upper_psi)],
          df[(df.Phi == upper_phi) & (df.Psi == upper_psi)],
      ]
  
      calc = np.random.uniform(0, 100, (n,))
  
      if n == 1:
          idxs = np.searchsorted(weights, calc)
          chis = sample_df(dfs[idxs[0]], uniform=uniform)
          return chis
      else:
          idxs = np.searchsorted(weights, calc)
          chis = []
  
          for i in range(4):
              count = (idxs == i).sum()
              if count > 0:
                  chi = sample_df(dfs[i], uniform=uniform, sample=count)
  
                  if count > 1:
                      chis.extend(chi)
                  else:
                      chis.append(chi)
  
          return chis
  
  
  def load_rotamor_library():
      # Loads the rotamor library
      amino_acids = [
          "arg",
          "asp",
          "asn",
          "cys",
          "glu",
          "gln",
          "his",
          "ile",
          "leu",
          "lys",
          "met",
          "phe",
          "pro",
          "ser",
          "thr",
          "trp",
          "tyr",
          "val",
      ]
      db = {}
  
      columns = collections.OrderedDict()
      columns["T"] = np.str
      columns["Phi"] = np.int64
      columns["Psi"] = np.int64
      columns["Count"] = np.int64
      columns["r1"] = np.int64
      columns["r2"] = np.int64
      columns["r3"] = np.int64
      columns["r4"] = np.int64
      columns["Probabil"] = np.float64
      columns["chi1Val"] = np.float64
      columns["chi2Val"] = np.float64
      columns["chi3Val"] = np.float64
      columns["chi4Val"] = np.float64
      columns["chi1Sig"] = np.float64
      columns["chi2Sig"] = np.float64
      columns["chi3Sig"] = np.float64
      columns["chi4Sig"] = np.float64
  
      for amino_acid in amino_acids:
          db[amino_acid] = pd.read_csv(
              osp.join(ROTAMER_LIBRARY_PATH, f"ExtendedOpt1-5/{amino_acid}.bbdep.rotamers.lib"),
              names=list(columns.keys()),
              dtype=columns,
              comment="#",
              delim_whitespace=True,
              engine="c",
          )
  
      return db
  
  
  def compute_rotamer_score_planar(gt_chi, neg_chi, chi_valid, res_name):
      select_res = {"phe": 1, "tyr": 1, "asp": 1, "glu": 2}
  
      if res_name in select_res.keys():
          n = select_res[res_name]
          chi_val = (
              np.minimum(
                  np.minimum(
                      np.abs(neg_chi[:n] - gt_chi[:n]), np.abs(neg_chi[:n] - gt_chi[:n] - 360)
                  ),
                  np.abs(neg_chi[:n] - gt_chi[:n] + 360),
              )
          ) * chi_valid[:n]
          chi_bool_i = chi_val < 20
  
          c1, c2 = neg_chi[n], gt_chi[n]
          c1, c2 = c1 % 180, c2 % 180
          min_dist = min(min(abs(c1 - c2), abs(c1 - c2 - 180)), abs(c1 - c2 + 180)) * chi_valid[n]
          chi_bool_last = min_dist < 20
  
          max_dist = chi_val + [min_dist]
          chi_bool = chi_bool_i.all() & chi_bool_last
      else:
          angle_diff = np.minimum(
              np.minimum(np.abs(neg_chi - gt_chi), np.abs(neg_chi - gt_chi - 360)),
              np.abs(neg_chi - gt_chi + 360),
          )
  
          angle_diff[np.isnan(angle_diff)] = 0
          chi_bool = (angle_diff * chi_valid) < 20
          max_dist = angle_diff * chi_valid
  
      return chi_bool.all(), max_dist
  
  
  def compute_dihedral_test():
      pos = np.array([[-1, 0, 0], [0, 0, 0], [0, 0, 1], [0, 1, 1]])[:, None, :]
      par = np.array([0, -1, -1, -1])
      child = np.array([1, 1, 1, 0])
      pos_exist = np.zeros(3)
      angle = compute_dihedral(par, child, pos, pos_exist, reshape=False)
      assert angle[2] == -90
  
  
  def compute_dihedral_test_2():
      # Load a random CIF file
      _, dense_embed = cif_to_embed("tests/101m.cif.gz")
      dense_embed = dense_embed[0][0]
      par, child, pos, pos_exist, res, chis_valid = parse_dense_format(dense_embed)
  
      angles = compute_dihedral(par, child, pos, pos_exist, reshape=True)
      perturb_angles = np.random.uniform(-180, 180, (angles.shape[0], 5))
  
      pos_new = rotate_dihedral(angles, par, child, pos, pos_exist, perturb_angles, chis_valid)
      new_angles = compute_dihedral(par, child, pos_new, pos_exist, reshape=True)
  
      chi_vals = new_angles[:, 4:9]
      error = (chi_vals - perturb_angles) * chis_valid[:, 4:9]
      zeros = np.zeros_like(error)
      assert np.isclose(error, zeros, rtol=1e-6).all()
  
  
  if __name__ == "__main__":
      print("Conducting unit tests on an example pdb file...")
  
      # Unit test for computing dihedral angles for simple perpendicular angles
      compute_dihedral_test()
  
      # Test compute dihedral angles via rotations from rotate_dihedral_angle_fast
      compute_dihedral_test_2()
  
      print("Unit tests passed.")
  
      # Preprocess dataset
      print("Preprocessing MMCIF files to construct the dataset...")
      script_parse_cif()
      print("Operation completed.")
  