# Copyright 2021 AlQuraishi Laboratory
# 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.

from functools import partial
import logging
import ml_collections
import numpy as np
import torch
import torch.nn as nn
from torch.distributions.bernoulli import Bernoulli
from typing import Dict, Optional, Tuple

from myopenfold.np import residue_constants
from myopenfold.utils import feats
from myopenfold.utils.rigid_utils import Rotation, Rigid
from myopenfold.utils.tensor_utils import (
    tree_map,
    tensor_tree_map,
    masked_mean,
    permute_final_dims,
    batched_gather,
)

from myopenfold.utils import debug

####################################################################### utils #######################################################################

def softmax_cross_entropy(logits, labels):
    loss = -1 * torch.sum(
        labels * torch.nn.functional.log_softmax(logits, dim=-1),
        dim=-1,
    )
    return loss


def sigmoid_cross_entropy(logits, labels):
    logits_dtype = logits.dtype
    logits = logits.double()
    labels = labels.double()
    log_p = torch.nn.functional.logsigmoid(logits)
    # log_p = torch.log(torch.sigmoid(logits))
    log_not_p = torch.nn.functional.logsigmoid(-1 * logits)
    # log_not_p = torch.log(torch.sigmoid(-logits))
    loss = (-1. * labels) * log_p - (1. - labels) * log_not_p
    loss = loss.to(dtype=logits_dtype)
    return loss

####################################################################### metrics #######################################################################

def lddt_ca(
    all_atom_pred_pos: torch.Tensor,
    all_atom_positions: torch.Tensor,
    all_atom_mask: torch.Tensor,
    cutoff: float = 15.0,
    eps: float = 1e-10,
    per_residue: bool = True,
) -> torch.Tensor:
    ca_pos = residue_constants.atom_order["CA"]
    all_atom_pred_pos = all_atom_pred_pos[..., ca_pos, :]
    all_atom_positions = all_atom_positions[..., ca_pos, :]
    all_atom_mask = all_atom_mask[..., ca_pos : (ca_pos + 1)]  # keep dim

    return lddt(
        all_atom_pred_pos,
        all_atom_positions,
        all_atom_mask,
        cutoff=cutoff,
        eps=eps,
        per_residue=per_residue,
    )

####################################################################### prepare #######################################################################

def compute_renamed_ground_truth(
    batch: Dict[str, torch.Tensor],
    atom14_pred_positions: torch.Tensor,
    eps=1e-10,
) -> Dict[str, torch.Tensor]:
    """
    Find optimal renaming of ground truth based on the predicted positions.

    Alg. 26 "renameSymmetricGroundTruthAtoms"

    This renamed ground truth is then used for all losses,
    such that each loss moves the atoms in the same direction.

    Args:
      batch: Dictionary containing:
        * atom14_gt_positions: Ground truth positions.
        * atom14_alt_gt_positions: Ground truth positions with renaming swaps.
        * atom14_atom_is_ambiguous: 1.0 for atoms that are affected by
            renaming swaps.
        * atom14_gt_exists: Mask for which atoms exist in ground truth.
        * atom14_alt_gt_exists: Mask for which atoms exist in ground truth
            after renaming.
        * atom14_atom_exists: Mask for whether each atom is part of the given
            amino acid type.
      atom14_pred_positions: Array of atom positions in global frame with shape
    Returns:
      Dictionary containing:
        alt_naming_is_better: Array with 1.0 where alternative swap is better.
        renamed_atom14_gt_positions: Array of optimal ground truth positions
          after renaming swaps are performed.
        renamed_atom14_gt_exists: Mask after renaming swap is performed.
    """

    pred_dists = torch.sqrt(
        eps
        + torch.sum(
            (
                atom14_pred_positions[..., None, :, None, :]
                - atom14_pred_positions[..., None, :, None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )

    atom14_gt_positions = batch["atom14_gt_positions"]
    gt_dists = torch.sqrt(
        eps
        + torch.sum(
            (
                atom14_gt_positions[..., None, :, None, :]
                - atom14_gt_positions[..., None, :, None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )

    atom14_alt_gt_positions = batch["atom14_alt_gt_positions"]
    alt_gt_dists = torch.sqrt(
        eps
        + torch.sum(
            (
                atom14_alt_gt_positions[..., None, :, None, :]
                - atom14_alt_gt_positions[..., None, :, None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )

    lddt = torch.sqrt(eps + (pred_dists - gt_dists) ** 2)
    alt_lddt = torch.sqrt(eps + (pred_dists - alt_gt_dists) ** 2)

    atom14_gt_exists = batch["atom14_gt_exists"]
    atom14_atom_is_ambiguous = batch["atom14_atom_is_ambiguous"]
    mask = (
        atom14_gt_exists[..., None, :, None]
        * atom14_atom_is_ambiguous[..., None, :, None]
        * atom14_gt_exists[..., None, :, None, :]
        * (1.0 - atom14_atom_is_ambiguous[..., None, :, None, :])
    )

    per_res_lddt = torch.sum(mask * lddt, dim=(-1, -2, -3))
    alt_per_res_lddt = torch.sum(mask * alt_lddt, dim=(-1, -2, -3))

    fp_type = atom14_pred_positions.dtype
    alt_naming_is_better = (alt_per_res_lddt < per_res_lddt).type(fp_type)

    renamed_atom14_gt_positions = (
        1.0 - alt_naming_is_better[..., None, None]
    ) * atom14_gt_positions + alt_naming_is_better[
        ..., None, None
    ] * atom14_alt_gt_positions

    renamed_atom14_gt_mask = (
        1.0 - alt_naming_is_better[..., None]
    ) * atom14_gt_exists + alt_naming_is_better[..., None] * batch[
        "atom14_alt_gt_exists"
    ]

    return {
        "alt_naming_is_better": alt_naming_is_better,
        "renamed_atom14_gt_positions": renamed_atom14_gt_positions,
        "renamed_atom14_gt_exists": renamed_atom14_gt_mask,
    }


def compute_plddt(logits: torch.Tensor) -> torch.Tensor:
    num_bins = logits.shape[-1]
    bin_width = 1.0 / num_bins
    bounds = torch.arange(
        start=0.5 * bin_width, end=1.0, step=bin_width, device=logits.device
    )
    probs = torch.nn.functional.softmax(logits, dim=-1)
    pred_lddt_ca = torch.sum(
        probs * bounds.view(*((1,) * len(probs.shape[:-1])), *bounds.shape),
        dim=-1,
    )
    return pred_lddt_ca * 100

####################################################################### fape loss #######################################################################

def compute_fape(
    pred_frames: Rigid,
    target_frames: Rigid,
    frames_mask: torch.Tensor,
    pred_positions: torch.Tensor,
    target_positions: torch.Tensor,
    positions_mask: torch.Tensor,
    length_scale: float,
    l1_clamp_distance: Optional[float] = None,
    eps=1e-8,
) -> torch.Tensor:
    """
        Computes FAPE loss.

        Args:
            pred_frames:
                [*, N_frames] Rigid object of predicted frames
            target_frames:
                [*, N_frames] Rigid object of ground truth frames
            frames_mask:
                [*, N_frames] binary mask for the frames
            pred_positions:
                [*, N_pts, 3] predicted atom positions
            target_positions:
                [*, N_pts, 3] ground truth positions
            positions_mask:
                [*, N_pts] positions mask
            length_scale:
                Length scale by which the loss is divided
            l1_clamp_distance:
                Cutoff above which distance errors are disregarded
            eps:
                Small value used to regularize denominators
        Returns:
            [*] loss tensor
    """
    #. pred_frames [N_sm, *, N_frames]
    #. target_frames [1, *, N_frames]
    #. frames_mask [1, *, N_frames]
    #. pred_positions [N_sm, *, N_pts, 3]
    #. target_positions [1, *, N_pts, 3]
    #. positions_mask [1, *, N_pts]
#    print("pred_frames", pred_frames.shape)  # DEBUG
#    print("target_frames", target_frames.shape)  # DEBUG
#    print("frames_mask", frames_mask.shape)  # DEBUG
#    print("pred_positions", pred_positions.shape)  # DEBUG
#    print("target_positions", target_positions.shape)  # DEBUG
#    print("positions_mask", positions_mask.shape)  # DEBUG

    # if pred_positions.shape[0] == 8:
    #     for i in range(8):
    #         print("difference between 0 and", i, torch.sum(torch.abs(pred_positions[0] - pred_positions[i])))

    # [*, N_frames, N_pts, 3]
    #. pred_frames.invert()[..., None] [N_sm, *, N_frames, 1]
    #. local_pred_pos [N_sm, *, N_frames, N_pts, 3]
    local_pred_pos = pred_frames.invert()[..., None].apply(
        pred_positions[..., None, :, :],  #. [N_sm, *, 1, N_pts, 3]
    )
    #. local_target_pos [1, *, N_frames, N_pts, 3]
    local_target_pos = target_frames.invert()[..., None].apply(
        target_positions[..., None, :, :],
    )
#    print("local_pred_pos", local_pred_pos.shape)  # DEBUG
#    print("local_target_pos", local_target_pos.shape)  # DEBUG

    #. error_dist [N_sm, *, N_frames, N_pts]
    error_dist = torch.sqrt(
        torch.sum((local_pred_pos - local_target_pos) ** 2, dim=-1) + eps
    )
#    print("error_dist", error_dist.shape)  # DEBUG

    if l1_clamp_distance is not None:
        error_dist = torch.clamp(error_dist, min=0, max=l1_clamp_distance)

    normed_error = error_dist / length_scale
#    print("normed_error", normed_error.shape)  # DEBUG
    normed_error = normed_error * frames_mask[..., None]
#    print("normed_error", normed_error.shape)  # DEBUG
    normed_error = normed_error * positions_mask[..., None, :]
#    print("normed_error", normed_error.shape)  # DEBUG

    # FP16-friendly averaging. Roughly equivalent to:
    #
    # norm_factor = (
    #     torch.sum(frames_mask, dim=-1) *
    #     torch.sum(positions_mask, dim=-1)
    # )
    # normed_error = torch.sum(normed_error, dim=(-1, -2)) / (eps + norm_factor)
    #
    # ("roughly" because eps is necessarily duplicated in the latter)
    #. normed_error [N_sm, *, N_frames]
    normed_error = torch.sum(normed_error, dim=-1)
#    print("normed_error", normed_error.shape)  # DEBUG
    #. normed_error [N_sm, *, N_frames]
    normed_error = (
        normed_error / (eps + torch.sum(frames_mask, dim=-1))[..., None]
    )
#    print("normed_error", normed_error.shape)  # DEBUG
    #. normed_error [N_sm, *, N_frames]
    normed_error = torch.sum(normed_error, dim=-1)
#    print("normed_error", normed_error.shape)  # DEBUG
    #. normed_error [N_sm, *]
    normed_error = normed_error / (eps + torch.sum(positions_mask, dim=-1))
#    print("normed_error", normed_error.shape)  # DEBUG

    return normed_error


def backbone_loss(
    backbone_rigid_tensor: torch.Tensor,
    backbone_rigid_mask: torch.Tensor,
    traj: torch.Tensor,
    use_clamped_fape: Optional[torch.Tensor] = None,
    clamp_distance: float = 10.0,
    loss_unit_distance: float = 10.0,
    eps: float = 1e-4,
    **kwargs,
) -> torch.Tensor:
    pred_aff = Rigid.from_tensor_7(traj)
    pred_aff = Rigid(
        Rotation(rot_mats=pred_aff.get_rots().get_rot_mats(), quats=None),
        pred_aff.get_trans(),
    )

    # DISCREPANCY: DeepMind somehow gets a hold of a tensor_7 version of
    # backbone tensor, normalizes it, and then turns it back to a rotation
    # matrix. To avoid a potentially numerically unstable rotation matrix
    # to quaternion conversion, we just use the original rotation matrix
    # outright. This one hasn't been composed a bunch of times, though, so
    # it might be fine.
    gt_aff = Rigid.from_tensor_4x4(backbone_rigid_tensor)

    fape_loss = compute_fape(
        pred_aff,
        gt_aff[None],
        backbone_rigid_mask[None],
        pred_aff.get_trans(),
        gt_aff[None].get_trans(),
        backbone_rigid_mask[None],
        l1_clamp_distance=clamp_distance,
        length_scale=loss_unit_distance,
        eps=eps,
    )
    if use_clamped_fape is not None:
        unclamped_fape_loss = compute_fape(
            pred_aff,
            gt_aff[None],
            backbone_rigid_mask[None],
            pred_aff.get_trans(),
            gt_aff[None].get_trans(),
            backbone_rigid_mask[None],
            l1_clamp_distance=None,
            length_scale=loss_unit_distance,
            eps=eps,
        )

        fape_loss = fape_loss * use_clamped_fape + unclamped_fape_loss * (
            1 - use_clamped_fape
        )

#    print("fape_loss", fape_loss)  # DEBUG
    #. average over iteration dim
    fape_loss = torch.mean(fape_loss, dim = 0)
    # TODO average over the batch dimension outside of the loss function
    # Average over the batch dimension
    # fape_loss = torch.mean(fape_loss)

    return fape_loss


def sidechain_loss(
    sidechain_frames: torch.Tensor,
    sidechain_atom_pos: torch.Tensor,
    rigidgroups_gt_frames: torch.Tensor,
    rigidgroups_alt_gt_frames: torch.Tensor,
    rigidgroups_gt_exists: torch.Tensor,
    renamed_atom14_gt_positions: torch.Tensor,
    renamed_atom14_gt_exists: torch.Tensor,
    alt_naming_is_better: torch.Tensor,
    clamp_distance: float = 10.0,
    length_scale: float = 10.0,
    eps: float = 1e-4,
    **kwargs,
) -> torch.Tensor:
    renamed_gt_frames = (
        1.0 - alt_naming_is_better[..., None, None, None]
    ) * rigidgroups_gt_frames + alt_naming_is_better[
        ..., None, None, None
    ] * rigidgroups_alt_gt_frames

    # Steamroll the inputs
    sidechain_frames = sidechain_frames[-1]
    batch_dims = sidechain_frames.shape[:-4]
    sidechain_frames = sidechain_frames.view(*batch_dims, -1, 4, 4)
    sidechain_frames = Rigid.from_tensor_4x4(sidechain_frames)
    renamed_gt_frames = renamed_gt_frames.view(*batch_dims, -1, 4, 4)
    renamed_gt_frames = Rigid.from_tensor_4x4(renamed_gt_frames)
    rigidgroups_gt_exists = rigidgroups_gt_exists.reshape(*batch_dims, -1)
    sidechain_atom_pos = sidechain_atom_pos[-1]
    sidechain_atom_pos = sidechain_atom_pos.view(*batch_dims, -1, 3)
    renamed_atom14_gt_positions = renamed_atom14_gt_positions.view(
        *batch_dims, -1, 3
    )
    renamed_atom14_gt_exists = renamed_atom14_gt_exists.view(*batch_dims, -1)

    fape = compute_fape(
        sidechain_frames,
        renamed_gt_frames,
        rigidgroups_gt_exists,
        sidechain_atom_pos,
        renamed_atom14_gt_positions,
        renamed_atom14_gt_exists,
        l1_clamp_distance=clamp_distance,
        length_scale=length_scale,
        eps=eps,
    )

    return fape


def fape_loss(
    out: Dict[str, torch.Tensor],
    batch: Dict[str, torch.Tensor],
    config: ml_collections.ConfigDict,
) -> torch.Tensor:
#    print("In fape_loss..")  # DEBUG

    #. bb_loss [*,]
    bb_loss = backbone_loss(
        traj=out["sm"]["frames"],
        **{**batch, **config.backbone},
    )
#    print("bb_loss", bb_loss.shape)  # DEBUG

    #. sc_loss [*,]
    sc_loss = sidechain_loss(
        out["sm"]["sidechain_frames"],
        out["sm"]["positions"],
        **{**batch, **config.sidechain},
    )
#    print("sc_loss", sc_loss.shape)  # DEBUG

#    print("bb_loss", bb_loss)  # DEBUG
#    print("sc_loss", sc_loss)  # DEBUG

    #. loss [*,]
    loss = config.backbone.weight * bb_loss + config.sidechain.weight * sc_loss
#    print("loss", loss.shape)  # DEBUG
    
    # TODO average over the batch dimension outside of the loss function
    # Average over the batch dimension
    # loss = torch.mean(loss)

    return loss


####################################################################### chi loss #######################################################################

def supervised_chi_loss(
    angles_sin_cos: torch.Tensor,
    unnormalized_angles_sin_cos: torch.Tensor,
    aatype: torch.Tensor,
    seq_mask: torch.Tensor,
    chi_mask: torch.Tensor,
    chi_angles_sin_cos: torch.Tensor,
    chi_weight: float,
    angle_norm_weight: float,
    eps=1e-6,
    **kwargs,
) -> torch.Tensor:
    """
        Implements Algorithm 27 (torsionAngleLoss)
        NOTE only compute supervised loss for chi angles, 
        NOTE  this is slightly different from the paper but same as the src code

        Args:
            angles_sin_cos:
                [*, N, 7, 2] predicted angles
            unnormalized_angles_sin_cos:
                The same angles, but unnormalized
            aatype:
                [*, N] residue indices
            seq_mask:
                [*, N] sequence mask
            chi_mask:
                [*, N, 7] angle mask
            chi_angles_sin_cos:
                [*, N, 7, 2] ground truth angles
            chi_weight:
                Weight for the angle component of the loss
            angle_norm_weight:
                Weight for the normalization component of the loss
        Returns:
            [*] loss tensor
    """
#    print("In supervised_chi_loss..")  # DEBUG
#    print("angles_sin_cos", angles_sin_cos.shape)  # DEBUG
#    print("unnormalized_angles_sin_cos", unnormalized_angles_sin_cos.shape)  # DEBUG
#    print("aatype", aatype.shape)  # DEBUG
#    print("seq_mask", seq_mask.shape)  # DEBUG
#    print("chi_mask", chi_mask.shape)  # DEBUG
#    print("chi_angles_sin_cos", chi_angles_sin_cos.shape)  # DEBUG
#    print("chi_weight", chi_weight)  # DEBUG
#    print("angle_norm_weight", angle_norm_weight)  # DEBUG
    #. angles_sin_cos [N_sm, *, N_res, 7, 2]
    #. unnormalized_angles_sin_cos [N_sm, *, N_res, 7, 2]
    #. aatype [*, N_res]
    #. seq_mask [*, N_res]
    #. chi_mask [*, N_res, 4]
    #. chi_angles_sin_cos [*, N_res, 4, 2]

    #. pred_angles [N_sm, *, N_res, 7, 2]
    pred_angles = angles_sin_cos[..., 3:, :]
#    print("pred_angles", pred_angles.shape)  # DEBUG

    #. residue_type_one_hot [*, N_res, 21]
    residue_type_one_hot = torch.nn.functional.one_hot(
        aatype,
        residue_constants.restype_num + 1,
    )
#    print("residue_type_one_hot", residue_type_one_hot.shape)  # DEBUG

    # NOTE why not use matrix multiplication?
    #. which chi angles are pi periodic
    #. chi_pi_periodic [*, N_res, 4]
    chi_pi_periodic = torch.einsum(
        "...ij,jk->ik",
        residue_type_one_hot.type(angles_sin_cos.dtype),  #. [*, N_res, 21]
        angles_sin_cos.new_tensor(residue_constants.chi_pi_periodic),  #. [21, 4]
    )
#    print("chi_pi_periodic", chi_pi_periodic.shape)  # DEBUG

    #. true_chi [1, *, N_res, 4, 2]
    true_chi = chi_angles_sin_cos[None]
#    print("true_chi", true_chi.shape)  # DEBUG

    #. shifted_mask [*, N_res, 4, 1]
    shifted_mask = (1 - 2 * chi_pi_periodic).unsqueeze(-1)  #. pi_periodic 1 -> -1, 0 -> 1
#    print("shifted_mask", shifted_mask.shape)  # DEBUG

    #. true_chi_shifted [1, *, N_res, 4, 2]
    true_chi_shifted = shifted_mask * true_chi
#    print("true_chi_shifted", true_chi_shifted.shape)  # DEBUG

    #. sq_chi_error [N_sm, *, N_res, 4]
    sq_chi_error = torch.sum((true_chi - pred_angles) ** 2, dim=-1)
#    print("sq_chi_error", sq_chi_error.shape)  # DEBUG

    #. sq_chi_error_shifted [N_sm, *, N_res, 4]
    sq_chi_error_shifted = torch.sum(
        (true_chi_shifted - pred_angles) ** 2, dim=-1
    )
#    print("sq_chi_error_shifted", sq_chi_error_shifted.shape)  # DEBUG

    #. sq_chi_error [N_sm, *, N_res, 4]
    sq_chi_error = torch.minimum(sq_chi_error, sq_chi_error_shifted)
#    print("sq_chi_error", sq_chi_error.shape)  # DEBUG
    
    #. important permute oper. Let batch dims be first, followed by [N_sm, N_res, 4]
    #. sq_chi_error [*, N_sm, N_res, 4]
    # The ol' switcheroo
    sq_chi_error = sq_chi_error.permute(
        *range(len(sq_chi_error.shape))[1:-2], 0, -2, -1
    )
#    print("sq_chi_error", sq_chi_error.shape)  # DEBUG

    #. sq_chi_loss [*,]
    sq_chi_loss = masked_mean(
        chi_mask[..., None, :, :], sq_chi_error, dim=(-1, -2, -3)
    )
#    print("sq_chi_loss", sq_chi_loss.shape)  # DEBUG

    #. loss [*,]
    loss = chi_weight * sq_chi_loss
#    print("loss", loss.shape)  # DEBUG

    #. angle_norm [N_sm, *, N_res, 7]
    angle_norm = torch.sqrt(
        torch.sum(unnormalized_angles_sin_cos ** 2, dim=-1) + eps
    )
#    print("angle_norm", angle_norm.shape)  # DEBUG

    #. norm_error [N_sm, *, N_res, 7] -> [*, N_sm, N_res, 7]
    norm_error = torch.abs(angle_norm - 1.0)
    norm_error = norm_error.permute(
        *range(len(norm_error.shape))[1:-2], 0, -2, -1
    )
#    print("norm_error", norm_error.shape)  # DEBUG

    #. angle_norm_loss [*,]
    angle_norm_loss = masked_mean(
        seq_mask[..., None, :, None], norm_error, dim=(-1, -2, -3)
    )
#    print("angle_norm_loss", angle_norm_loss.shape)  # DEBUG

    loss = loss + angle_norm_weight * angle_norm_loss

    # TODO average over the batch dimension outside of the loss function
    # Average over the batch dimension
    # loss = torch.mean(loss)

    return loss


####################################################################### plddt loss #######################################################################

def lddt(
    all_atom_pred_pos: torch.Tensor,
    all_atom_positions: torch.Tensor,
    all_atom_mask = None,   #. torch.Tensor
    cutoff: float = 15.0,
    eps: float = 1e-10,
    per_residue: bool = True,
) -> torch.Tensor:
    #. all_atom_pred_pos [*, N_atom, 3]
    #. all_atom_positions [*, N_atom, 3]
    #. all_atom_mask [*, N_atom]

    # NOTE add this in case we pre-select the non-nan atoms
    if all_atom_mask is None:
        dims = all_atom_pred_pos.shape[:-1] + (1,)
        all_atom_mask = torch.ones(dims, dtype=torch.bool, device=all_atom_pred_pos.device)

    #. n N_atom
    n = all_atom_mask.shape[-2]

    #. dmat_true [*, N_atom, N_atom]
    dmat_true = torch.sqrt(
        eps
        + torch.sum(
            (
                all_atom_positions[..., None, :]
                - all_atom_positions[..., None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )

    #. dmat_pred [*, N_atom, N_atom]
    dmat_pred = torch.sqrt(
        eps
        + torch.sum(
            (
                all_atom_pred_pos[..., None, :]
                - all_atom_pred_pos[..., None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )

    #. mask those:
    #.  1. dist > cutoff
    #.  2. not in all_atom_mask
    #.  3. diagonal
    #. dists_to_score [*, N_atom, N_atom]
    dists_to_score = (
        (dmat_true < cutoff)
        * all_atom_mask
        * permute_final_dims(all_atom_mask, (1, 0))
        * (1.0 - torch.eye(n, device=all_atom_mask.device))
    )

    #. dist_l1 [*, N_atom, N_atom]
    dist_l1 = torch.abs(dmat_true - dmat_pred)

    #. score [*, N_atom, N_atom]
    score = (
        (dist_l1 < 0.5).type(dist_l1.dtype)
        + (dist_l1 < 1.0).type(dist_l1.dtype)
        + (dist_l1 < 2.0).type(dist_l1.dtype)
        + (dist_l1 < 4.0).type(dist_l1.dtype)
    )
    score = score * 0.25

    dims = (-1,) if per_residue else (-2, -1)

    #. norm [*,] or [*, N_atom]
    norm = 1.0 / (eps + torch.sum(dists_to_score, dim=dims))

    #. score [*,] or [*, N_atom]
    score = norm * (eps + torch.sum(dists_to_score * score, dim=dims))

    return score


def lddt_loss(
    logits: torch.Tensor,
    all_atom_pred_pos: torch.Tensor,
    all_atom_positions: torch.Tensor,
    all_atom_mask: torch.Tensor,
    resolution: torch.Tensor,
    cutoff: float = 15.0,
    no_bins: int = 50,
    min_resolution: float = 0.1,
    max_resolution: float = 3.0,
    eps: float = 1e-10,
    **kwargs,
) -> torch.Tensor:
#    print("In lddt_loss..")  # DEBUG
#    print("logits.shape: ", logits.shape)  # DEBUG
#    print("all_atom_pred_pos.shape: ", all_atom_pred_pos.shape)  # DEBUG
#    print("all_atom_positions.shape: ", all_atom_positions.shape)  # DEBUG
#    print("all_atom_mask.shape: ", all_atom_mask.shape)  # DEBUG
#    print("resolution.shape: ", resolution.shape)  # DEBUG
    #. logits [*, N_res, 50]
    #. all_atom_pred_pos [*, N_res, 37, 3]
    #. all_atom_positions [*, N_res, 37, 3]
    #. all_atom_mask [*, N_res, 37]
    #. resolution [*,]

    #. N_atom = 37
    n = all_atom_mask.shape[-2]

    #. extract CA positions
    #. all_atom_pred_pos [*, N_res, 3]
    #. all_atom_positions [*, N_res, 3]
    ca_pos = residue_constants.atom_order["CA"]
    all_atom_pred_pos = all_atom_pred_pos[..., ca_pos, :]
    all_atom_positions = all_atom_positions[..., ca_pos, :]
#    print("all_atom_pred_pos.shape: ", all_atom_pred_pos.shape)  # DEBUG
#    print("all_atom_positions.shape: ", all_atom_positions.shape)  # DEBUG

    #. all_atom_mask [*, N_res, 1]
    all_atom_mask = all_atom_mask[..., ca_pos : (ca_pos + 1)]  # keep dim
#    print("all_atom_mask.shape: ", all_atom_mask.shape)  # DEBUG

    #. compute ground truth plddt score
    #. resolution [*,]
    score = lddt(
        all_atom_pred_pos, 
        all_atom_positions, 
        all_atom_mask, 
        cutoff=cutoff, 
        eps=eps
    )
#    print("score.shape: ", score.shape)  # DEBUG

    score = score.detach()

    #. ground truth bin index
    #. bin_index [*, N_res]
    bin_index = torch.floor(score * no_bins).long()
    bin_index = torch.clamp(bin_index, max=(no_bins - 1))
#    print("bin_index.shape: ", bin_index.shape)  # DEBUG

    #. one-hot encoding of ground truth bin index
    #. lddt_ca_one_hot [*, N_res, 50]
    lddt_ca_one_hot = torch.nn.functional.one_hot(
        bin_index, num_classes=no_bins
    )
#    print("lddt_ca_one_hot.shape: ", lddt_ca_one_hot.shape)  # DEBUG

    #. compute cross entropy loss
    #. errors [*, N_res]
    errors = softmax_cross_entropy(logits, lddt_ca_one_hot)
#    print("errors.shape: ", errors.shape)  # DEBUG

    #. all_atom_mask [*, N_res, 1]
    all_atom_mask = all_atom_mask.squeeze(-1)
#    print("all_atom_mask.shape: ", all_atom_mask.shape)  # DEBUG

    #. loss per sample, average over all residues
    #. loss [*,]
    loss = torch.sum(errors * all_atom_mask, dim=-1) / (
        eps + torch.sum(all_atom_mask, dim=-1)
    )
#    print("loss.shape: ", loss.shape)  # DEBUG

    #. mask out resolution out of range
    loss = loss * (
        (resolution >= min_resolution) & (resolution <= max_resolution)
    )

    # TODO average over the batch dimension outside of the loss function
    # Average over the batch dimension
    # loss = torch.mean(loss)

    return loss


####################################################################### distogram loss #######################################################################

def distogram_loss(
    logits,
    pseudo_beta,
    pseudo_beta_mask,
    min_bin=2.3125,
    max_bin=21.6875,
    no_bins=64,
    eps=1e-6,
    **kwargs,
):
#    print("In distogram_loss..")  # DEBUG
#    print("logits.shape: ", logits.shape)  # DEBUG
#    print("pseudo_beta.shape: ", pseudo_beta.shape)  # DEBUG
#    print("pseudo_beta_mask.shape: ", pseudo_beta_mask.shape)  # DEBUG
    #. logits [*, N_res, N_res, 64]
    #. pseudo_beta [*, N_res, 3]
    #. pseudo_beta_mask [*, N_res]

    #. bin partition boundaries
    #. boundaries [64]
    boundaries = torch.linspace(
        min_bin,
        max_bin,
        no_bins - 1,
        device=logits.device,
    )
    boundaries = boundaries ** 2
#    print("boundaries.shape: ", boundaries.shape)  # DEBUG

    #. compute distance matrix
    #. dists [*, N_res, N_res, 1]
    dists = torch.sum(
        (pseudo_beta[..., None, :] - pseudo_beta[..., None, :, :]) ** 2,
        dim=-1,
        keepdims=True,
    )
#    print("dists.shape: ", dists.shape)  # DEBUG

    #. compute true bins
    #. true_bins [*, N_res, N_res]
    true_bins = torch.sum(dists > boundaries, dim=-1)  # NOTE low readability, but seems right
#    print("true_bins.shape: ", true_bins.shape)  # DEBUG

    #. compute loss
    #. errors [*, N_res, N_res]
    errors = softmax_cross_entropy(
        logits,
        torch.nn.functional.one_hot(true_bins, no_bins),  #. [*, N_res, N_res, 64]
    )
#    print("torch.nn.functional.one_hot(true_bins, no_bins)", torch.nn.functional.one_hot(true_bins, no_bins).shape)  # DEBUG
#    print("errors.shape: ", errors.shape)  # DEBUG

    #. 2D mask
    #. square_mask [*, N_res, N_res]
    square_mask = pseudo_beta_mask[..., None] * pseudo_beta_mask[..., None, :]
#    print("square_mask.shape: ", square_mask.shape)  # DEBUG

    # FP16-friendly sum. Equivalent to:
    # mean = (torch.sum(errors * square_mask, dim=(-1, -2)) /
    #         (eps + torch.sum(square_mask, dim=(-1, -2))))
    scale = 0.5  # NOTE add a scale factor to avoid overflow, same as original implementation in masked_msa_loss
    #. denom[..., None] [*]
    denom = eps + torch.sum(scale * square_mask, dim=(-1, -2))  # NOTE added scale
#    print("denom.shape: ", denom.shape)  # DEBUG

    #. mean [*, N_res, N_res] -> [*, N_res] -> [*]
    mean = errors * square_mask
#    print("mean.shape: ", mean.shape)  # DEBUG
    mean = torch.sum(mean, dim=-1)
    mean = mean / denom[..., None]
    mean = torch.sum(mean, dim=-1)
    mean = mean * scale  # NOTE added scale

    # TODO average over the batch dimension outside of the loss function
    # Average over the batch dimensions
    # mean = torch.mean(mean)

    return mean


####################################################################### violation loss #######################################################################

def between_residue_bond_loss(
    pred_atom_positions: torch.Tensor,  # (*, N, 37/14, 3)
    pred_atom_mask: torch.Tensor,  # (*, N, 37/14)
    residue_index: torch.Tensor,  # (*, N)
    aatype: torch.Tensor,  # (*, N)
    tolerance_factor_soft=12.0,
    tolerance_factor_hard=12.0,
    eps=1e-6,
) -> Dict[str, torch.Tensor]:
    """Flat-bottom loss to penalize structural violations between residues.

    This is a loss penalizing any violation of the geometry around the peptide
    bond between consecutive amino acids. This loss corresponds to
    Jumper et al. (2021) Suppl. Sec. 1.9.11, eq 44, 45.

    Args:
      pred_atom_positions: Atom positions in atom37/14 representation
      pred_atom_mask: Atom mask in atom37/14 representation
      residue_index: Residue index for given amino acid, this is assumed to be
        monotonically increasing.
      aatype: Amino acid type of given residue
      tolerance_factor_soft: soft tolerance factor measured in standard deviations
        of pdb distributions
      tolerance_factor_hard: hard tolerance factor measured in standard deviations
        of pdb distributions

    Returns:
      Dict containing:
        * 'c_n_loss_mean': Loss for peptide bond length violations
        * 'ca_c_n_loss_mean': Loss for violations of bond angle around C spanned
            by CA, C, N
        * 'c_n_ca_loss_mean': Loss for violations of bond angle around N spanned
            by C, N, CA
        * 'per_residue_loss_sum': sum of all losses for each residue
        * 'per_residue_violation_mask': mask denoting all residues with violation
            present.
    """
    # Get the positions of the relevant backbone atoms.
    this_ca_pos = pred_atom_positions[..., :-1, 1, :]
    this_ca_mask = pred_atom_mask[..., :-1, 1]
    this_c_pos = pred_atom_positions[..., :-1, 2, :]
    this_c_mask = pred_atom_mask[..., :-1, 2]
    next_n_pos = pred_atom_positions[..., 1:, 0, :]
    next_n_mask = pred_atom_mask[..., 1:, 0]
    next_ca_pos = pred_atom_positions[..., 1:, 1, :]
    next_ca_mask = pred_atom_mask[..., 1:, 1]
    has_no_gap_mask = (residue_index[..., 1:] - residue_index[..., :-1]) == 1.0

    # Compute loss for the C--N bond.
    c_n_bond_length = torch.sqrt(
        eps + torch.sum((this_c_pos - next_n_pos) ** 2, dim=-1)
    )

    # The C-N bond to proline has slightly different length because of the ring.
    next_is_proline = aatype[..., 1:] == residue_constants.resname_to_idx["PRO"]
    gt_length = (
        ~next_is_proline
    ) * residue_constants.between_res_bond_length_c_n[
        0
    ] + next_is_proline * residue_constants.between_res_bond_length_c_n[
        1
    ]
    gt_stddev = (
        ~next_is_proline
    ) * residue_constants.between_res_bond_length_stddev_c_n[
        0
    ] + next_is_proline * residue_constants.between_res_bond_length_stddev_c_n[
        1
    ]
    c_n_bond_length_error = torch.sqrt(eps + (c_n_bond_length - gt_length) ** 2)
    c_n_loss_per_residue = torch.nn.functional.relu(
        c_n_bond_length_error - tolerance_factor_soft * gt_stddev
    )
    mask = this_c_mask * next_n_mask * has_no_gap_mask
    c_n_loss = torch.sum(mask * c_n_loss_per_residue, dim=-1) / (
        torch.sum(mask, dim=-1) + eps
    )
    c_n_violation_mask = mask * (
        c_n_bond_length_error > (tolerance_factor_hard * gt_stddev)
    )

    # Compute loss for the angles.
    ca_c_bond_length = torch.sqrt(
        eps + torch.sum((this_ca_pos - this_c_pos) ** 2, dim=-1)
    )
    n_ca_bond_length = torch.sqrt(
        eps + torch.sum((next_n_pos - next_ca_pos) ** 2, dim=-1)
    )

    c_ca_unit_vec = (this_ca_pos - this_c_pos) / ca_c_bond_length[..., None]
    c_n_unit_vec = (next_n_pos - this_c_pos) / c_n_bond_length[..., None]
    n_ca_unit_vec = (next_ca_pos - next_n_pos) / n_ca_bond_length[..., None]

    ca_c_n_cos_angle = torch.sum(c_ca_unit_vec * c_n_unit_vec, dim=-1)
    gt_angle = residue_constants.between_res_cos_angles_ca_c_n[0]
    gt_stddev = residue_constants.between_res_bond_length_stddev_c_n[0]
    ca_c_n_cos_angle_error = torch.sqrt(
        eps + (ca_c_n_cos_angle - gt_angle) ** 2
    )
    ca_c_n_loss_per_residue = torch.nn.functional.relu(
        ca_c_n_cos_angle_error - tolerance_factor_soft * gt_stddev
    )
    mask = this_ca_mask * this_c_mask * next_n_mask * has_no_gap_mask
    ca_c_n_loss = torch.sum(mask * ca_c_n_loss_per_residue, dim=-1) / (
        torch.sum(mask, dim=-1) + eps
    )
    ca_c_n_violation_mask = mask * (
        ca_c_n_cos_angle_error > (tolerance_factor_hard * gt_stddev)
    )

    c_n_ca_cos_angle = torch.sum((-c_n_unit_vec) * n_ca_unit_vec, dim=-1)
    gt_angle = residue_constants.between_res_cos_angles_c_n_ca[0]
    gt_stddev = residue_constants.between_res_cos_angles_c_n_ca[1]
    c_n_ca_cos_angle_error = torch.sqrt(
        eps + torch.square(c_n_ca_cos_angle - gt_angle)
    )
    c_n_ca_loss_per_residue = torch.nn.functional.relu(
        c_n_ca_cos_angle_error - tolerance_factor_soft * gt_stddev
    )
    mask = this_c_mask * next_n_mask * next_ca_mask * has_no_gap_mask
    c_n_ca_loss = torch.sum(mask * c_n_ca_loss_per_residue, dim=-1) / (
        torch.sum(mask, dim=-1) + eps
    )
    c_n_ca_violation_mask = mask * (
        c_n_ca_cos_angle_error > (tolerance_factor_hard * gt_stddev)
    )

    # Compute a per residue loss (equally distribute the loss to both
    # neighbouring residues).
    per_residue_loss_sum = (
        c_n_loss_per_residue + ca_c_n_loss_per_residue + c_n_ca_loss_per_residue
    )
    per_residue_loss_sum = 0.5 * (
        torch.nn.functional.pad(per_residue_loss_sum, (0, 1))
        + torch.nn.functional.pad(per_residue_loss_sum, (1, 0))
    )

    # Compute hard violations.
    violation_mask = torch.max(
        torch.stack(
            [c_n_violation_mask, ca_c_n_violation_mask, c_n_ca_violation_mask],
            dim=-2,
        ),
        dim=-2,
    )[0]
    violation_mask = torch.maximum(
        torch.nn.functional.pad(violation_mask, (0, 1)),
        torch.nn.functional.pad(violation_mask, (1, 0)),
    )

    return {
        "c_n_loss_mean": c_n_loss,
        "ca_c_n_loss_mean": ca_c_n_loss,
        "c_n_ca_loss_mean": c_n_ca_loss,
        "per_residue_loss_sum": per_residue_loss_sum,
        "per_residue_violation_mask": violation_mask,
    }


def between_residue_clash_loss(
    atom14_pred_positions: torch.Tensor,
    atom14_atom_exists: torch.Tensor,
    atom14_atom_radius: torch.Tensor,
    residue_index: torch.Tensor,
    overlap_tolerance_soft=1.5,
    overlap_tolerance_hard=1.5,
    eps=1e-10,
) -> Dict[str, torch.Tensor]:
    """Loss to penalize steric clashes between residues.

    This is a loss penalizing any steric clashes due to non bonded atoms in
    different peptides coming too close. This loss corresponds to the part with
    different residues of
    Jumper et al. (2021) Suppl. Sec. 1.9.11, eq 46.

    Args:
      atom14_pred_positions: Predicted positions of atoms in
        global prediction frame
      atom14_atom_exists: Mask denoting whether atom at positions exists for given
        amino acid type
      atom14_atom_radius: Van der Waals radius for each atom.
      residue_index: Residue index for given amino acid.
      overlap_tolerance_soft: Soft tolerance factor.
      overlap_tolerance_hard: Hard tolerance factor.

    Returns:
      Dict containing:
        * 'mean_loss': average clash loss
        * 'per_atom_loss_sum': sum of all clash losses per atom, shape (N, 14)
        * 'per_atom_clash_mask': mask whether atom clashes with any other atom
            shape (N, 14)
    """
    fp_type = atom14_pred_positions.dtype

    # Create the distance matrix.
    # (N, N, 14, 14)
    dists = torch.sqrt(
        eps
        + torch.sum(
            (
                atom14_pred_positions[..., :, None, :, None, :]
                - atom14_pred_positions[..., None, :, None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )

    # Create the mask for valid distances.
    # shape (N, N, 14, 14)
    dists_mask = (
        atom14_atom_exists[..., :, None, :, None]
        * atom14_atom_exists[..., None, :, None, :]
    ).type(fp_type)

    # Mask out all the duplicate entries in the lower triangular matrix.
    # Also mask out the diagonal (atom-pairs from the same residue) -- these atoms
    # are handled separately.
    dists_mask = dists_mask * (
        residue_index[..., :, None, None, None]
        < residue_index[..., None, :, None, None]
    )

    # Backbone C--N bond between subsequent residues is no clash.
    c_one_hot = torch.nn.functional.one_hot(
        residue_index.new_tensor(2), num_classes=14
    )
    c_one_hot = c_one_hot.reshape(
        *((1,) * len(residue_index.shape[:-1])), *c_one_hot.shape
    )
    c_one_hot = c_one_hot.type(fp_type)
    n_one_hot = torch.nn.functional.one_hot(
        residue_index.new_tensor(0), num_classes=14
    )
    n_one_hot = n_one_hot.reshape(
        *((1,) * len(residue_index.shape[:-1])), *n_one_hot.shape
    )
    n_one_hot = n_one_hot.type(fp_type)

    neighbour_mask = (
        residue_index[..., :, None, None, None] + 1
    ) == residue_index[..., None, :, None, None]
    c_n_bonds = (
        neighbour_mask
        * c_one_hot[..., None, None, :, None]
        * n_one_hot[..., None, None, None, :]
    )
    dists_mask = dists_mask * (1.0 - c_n_bonds)

    # Disulfide bridge between two cysteines is no clash.
    cys = residue_constants.restype_name_to_atom14_names["CYS"]
    cys_sg_idx = cys.index("SG")
    cys_sg_idx = residue_index.new_tensor(cys_sg_idx)
    cys_sg_idx = cys_sg_idx.reshape(
        *((1,) * len(residue_index.shape[:-1])), 1
    ).squeeze(-1)
    cys_sg_one_hot = torch.nn.functional.one_hot(cys_sg_idx, num_classes=14)
    disulfide_bonds = (
        cys_sg_one_hot[..., None, None, :, None]
        * cys_sg_one_hot[..., None, None, None, :]
    )
    dists_mask = dists_mask * (1.0 - disulfide_bonds)

    # Compute the lower bound for the allowed distances.
    # shape (N, N, 14, 14)
    dists_lower_bound = dists_mask * (
        atom14_atom_radius[..., :, None, :, None]
        + atom14_atom_radius[..., None, :, None, :]
    )

    # Compute the error.
    # shape (N, N, 14, 14)
    dists_to_low_error = dists_mask * torch.nn.functional.relu(
        dists_lower_bound - overlap_tolerance_soft - dists
    )

    # Compute the mean loss.
    # shape ()
    mean_loss = torch.sum(dists_to_low_error, dim=(-1, -2, -3, -4)) / (1e-6 + torch.sum(dists_mask, dim=(-1, -2, -3, -4)))

    # Compute the per atom loss sum.
    # shape (N, 14)
    per_atom_loss_sum = torch.sum(dists_to_low_error, dim=(-4, -2)) + torch.sum(
        dists_to_low_error, axis=(-3, -1)
    )

    # Compute the hard clash mask.
    # shape (N, N, 14, 14)
    clash_mask = dists_mask * (
        dists < (dists_lower_bound - overlap_tolerance_hard)
    )

    # Compute the per atom clash.
    # shape (N, 14)
    per_atom_clash_mask = torch.maximum(
        torch.amax(clash_mask, axis=(-4, -2)),
        torch.amax(clash_mask, axis=(-3, -1)),
    )

    return {
        "mean_loss": mean_loss,  # shape ()
        "per_atom_loss_sum": per_atom_loss_sum,  # shape (N, 14)
        "per_atom_clash_mask": per_atom_clash_mask,  # shape (N, 14)
    }


def within_residue_violations(
    atom14_pred_positions: torch.Tensor,
    atom14_atom_exists: torch.Tensor,
    atom14_dists_lower_bound: torch.Tensor,
    atom14_dists_upper_bound: torch.Tensor,
    tighten_bounds_for_loss=0.0,
    eps=1e-10,
) -> Dict[str, torch.Tensor]:
    """Loss to penalize steric clashes within residues.

    This is a loss penalizing any steric violations or clashes of non-bonded atoms
    in a given peptide. This loss corresponds to the part with
    the same residues of
    Jumper et al. (2021) Suppl. Sec. 1.9.11, eq 46.

    Args:
        atom14_pred_positions ([*, N, 14, 3]):
            Predicted positions of atoms in global prediction frame.
        atom14_atom_exists ([*, N, 14]):
            Mask denoting whether atom at positions exists for given
            amino acid type
        atom14_dists_lower_bound ([*, N, 14]):
            Lower bound on allowed distances.
        atom14_dists_upper_bound ([*, N, 14]):
            Upper bound on allowed distances
        tighten_bounds_for_loss ([*, N]):
            Extra factor to tighten loss

    Returns:
      Dict containing:
        * 'per_atom_loss_sum' ([*, N, 14]):
              sum of all clash losses per atom, shape
        * 'per_atom_clash_mask' ([*, N, 14]):
              mask whether atom clashes with any other atom shape
    """
    # Compute the mask for each residue.
    dists_masks = 1.0 - torch.eye(14, device=atom14_atom_exists.device)[None]
    dists_masks = dists_masks.reshape(
        *((1,) * len(atom14_atom_exists.shape[:-2])), *dists_masks.shape
    )
    dists_masks = (
        atom14_atom_exists[..., :, :, None]
        * atom14_atom_exists[..., :, None, :]
        * dists_masks
    )

    # Distance matrix
    dists = torch.sqrt(
        eps
        + torch.sum(
            (
                atom14_pred_positions[..., :, :, None, :]
                - atom14_pred_positions[..., :, None, :, :]
            )
            ** 2,
            dim=-1,
        )
    )

    # Compute the loss.
    dists_to_low_error = torch.nn.functional.relu(
        atom14_dists_lower_bound + tighten_bounds_for_loss - dists
    )
    dists_to_high_error = torch.nn.functional.relu(
        dists - (atom14_dists_upper_bound - tighten_bounds_for_loss)
    )
    loss = dists_masks * (dists_to_low_error + dists_to_high_error)

    # Compute the per atom loss sum.
    per_atom_loss_sum = torch.sum(loss, dim=-2) + torch.sum(loss, dim=-1)

    # Compute the violations mask.
    violations = dists_masks * (
        (dists < atom14_dists_lower_bound) | (dists > atom14_dists_upper_bound)
    )

    # Compute the per atom violations.
    per_atom_violations = torch.maximum(
        torch.max(violations, dim=-2)[0], torch.max(violations, axis=-1)[0]
    )

    return {
        "per_atom_loss_sum": per_atom_loss_sum,
        "per_atom_violations": per_atom_violations,
    }


def find_structural_violations(
    batch: Dict[str, torch.Tensor],
    atom14_pred_positions: torch.Tensor,
    violation_tolerance_factor: float,
    clash_overlap_tolerance: float,
    **kwargs,
) -> Dict[str, torch.Tensor]:
    """Computes several checks for structural violations."""

    # Compute between residue backbone violations of bonds and angles.
#    print("Computing structural violation loss..")  # DEBUG
#    print("atom14_pred_positions", atom14_pred_positions.shape)  # DEBUG
#    print("batch['atom14_atom_exists']", batch["atom14_atom_exists"].shape)  # DEBUG
#    print("batch['residue_index']", batch["residue_index"].shape)  # DEBUG
#    print("batch['aatype']", batch["aatype"].shape)  # DEBUG
    #. atom14_pred_positions [*, N_res, 14, 3]
    #. atom14_atom_exists [*, N_res, 14]
    #. residue_index [*, N_res,]
    #. aatype [*, N_res,]
    
    connection_violations = between_residue_bond_loss(
        pred_atom_positions=atom14_pred_positions,
        pred_atom_mask=batch["atom14_atom_exists"],
        residue_index=batch["residue_index"],
        aatype=batch["aatype"],
        tolerance_factor_soft=violation_tolerance_factor,
        tolerance_factor_hard=violation_tolerance_factor,
    )
#    print("connection_violations", debug.print_shape(connection_violations))  # DEBUG
    #. c_n_loss_mean []
    #. ca_c_n_loss_mean []
    #. c_n_ca_loss_mean []
    #. per_residue_loss_sum [*, N_res]
    #. per_residue_violation_mask [*, N_res]

    # Compute the Van der Waals radius for every atom
    # (the first letter of the atom name is the element type).
    # Shape: (N, 14).
    #. atomtype_radius [*, N_res, 14]
    atomtype_radius = [
        residue_constants.van_der_waals_radius[name[0]]
        for name in residue_constants.atom_types
    ]
    atomtype_radius = atom14_pred_positions.new_tensor(atomtype_radius)
#    print("atomtype_radius", debug.print_shape(atomtype_radius))  # DEBUG

    
    atom14_atom_radius = (
        batch["atom14_atom_exists"]
        * atomtype_radius[batch["residx_atom14_to_atom37"]]
    )
#    print("atom14_atom_radius", atom14_atom_radius.shape)  # DEBUG

    # Compute the between residue clash loss.
    #. mean_loss []
    #. per_atom_loss_sum [*, N_res, 14]
    #. per_atom_clash_mask [*, N_res, 14]
    between_residue_clashes = between_residue_clash_loss(
        atom14_pred_positions=atom14_pred_positions,
        atom14_atom_exists=batch["atom14_atom_exists"],
        atom14_atom_radius=atom14_atom_radius,
        residue_index=batch["residue_index"],
        overlap_tolerance_soft=clash_overlap_tolerance,
        overlap_tolerance_hard=clash_overlap_tolerance,
    )
#    print("between_residue_clashes", debug.print_shape(between_residue_clashes))  # DEBUG

    #. lower_bound [21, 14, 14]
    #. upper_bound [21, 14, 14]
    #. stddev [21, 14, 14]
    # Compute all within-residue violations (clashes,
    # bond length and angle violations).
    restype_atom14_bounds = residue_constants.make_atom14_dists_bounds(
        overlap_tolerance=clash_overlap_tolerance,
        bond_length_tolerance_factor=violation_tolerance_factor,
    )
#    print("restype_atom14_bounds", debug.print_shape(restype_atom14_bounds))  # DEBUG

    #. atom14_atom_exists [*, N_res, 14]
    #. atom14_dists_lower_bound [*, N_res, 14, 14]
    #. atom14_dists_upper_bound [*, N_res, 14, 14]
    #. residue_violations["per_atom_loss_sum"] [*, N_res, 14]
    #. residue_violations["per_atom_violations"] [*, N_res, 14]
    atom14_atom_exists = batch["atom14_atom_exists"]
    atom14_dists_lower_bound = atom14_pred_positions.new_tensor(
        restype_atom14_bounds["lower_bound"]
    )[batch["aatype"]]
    atom14_dists_upper_bound = atom14_pred_positions.new_tensor(
        restype_atom14_bounds["upper_bound"]
    )[batch["aatype"]]
    residue_violations = within_residue_violations(
        atom14_pred_positions=atom14_pred_positions,
        atom14_atom_exists=batch["atom14_atom_exists"],
        atom14_dists_lower_bound=atom14_dists_lower_bound,
        atom14_dists_upper_bound=atom14_dists_upper_bound,
        tighten_bounds_for_loss=0.0,
    )
#    print("atom14_dists_lower_bound", atom14_dists_lower_bound.shape)  # DEBUG
#    print("atom14_dists_upper_bound", atom14_dists_upper_bound.shape)  # DEBUG
#    print("residue_violations", debug.print_shape(residue_violations))  # DEBUG

    # Combine them to a single per-residue violation mask (used later for LDDT).
    #. connection_violations["per_residue_violation_mask"] [*, N_res]
    #. between_residue_clashes["per_atom_clash_mask"] [*, N_res, 14]
    #. residue_violations["per_atom_violations"] [*, N_res, 14]
    #. per_residue_violations_mask [*, N_res]
    per_residue_violations_mask = torch.max(
        torch.stack(
            [
                connection_violations["per_residue_violation_mask"],
                torch.max(
                    between_residue_clashes["per_atom_clash_mask"], dim=-1
                )[0],
                torch.max(residue_violations["per_atom_violations"], dim=-1)[0],
            ],
            dim=-1,
        ),
        dim=-1,
    )[0]
#    print('connection_violations["per_residue_violation_mask"]', connection_violations["per_residue_violation_mask"].shape)  # DEBUG
#    print('between_residue_clashes["per_atom_clash_mask"]', between_residue_clashes["per_atom_clash_mask"].shape)  # DEBUG
#    print('residue_violations["per_atom_violations"]', residue_violations["per_atom_violations"].shape)  # DEBUG
#    print("per_residue_violations_mask", per_residue_violations_mask.shape)  # DEBUG

    return {
        "between_residues": {
            "bonds_c_n_loss_mean": connection_violations["c_n_loss_mean"],  # ()
            "angles_ca_c_n_loss_mean": connection_violations[
                "ca_c_n_loss_mean"
            ],  # ()
            "angles_c_n_ca_loss_mean": connection_violations[
                "c_n_ca_loss_mean"
            ],  # ()
            "connections_per_residue_loss_sum": connection_violations[
                "per_residue_loss_sum"
            ],  # (N)
            "connections_per_residue_violation_mask": connection_violations[
                "per_residue_violation_mask"
            ],  # (N)
            "clashes_mean_loss": between_residue_clashes["mean_loss"],  # ()
            "clashes_per_atom_loss_sum": between_residue_clashes[
                "per_atom_loss_sum"
            ],  # (N, 14)
            "clashes_per_atom_clash_mask": between_residue_clashes[
                "per_atom_clash_mask"
            ],  # (N, 14)
        },
        "within_residues": {
            "per_atom_loss_sum": residue_violations[
                "per_atom_loss_sum"
            ],  # (N, 14)
            "per_atom_violations": residue_violations[
                "per_atom_violations"
            ],  # (N, 14),
        },
        "total_per_residue_violations_mask": per_residue_violations_mask,  # (N)
    }


def violation_loss(
    violations: Dict[str, torch.Tensor],
    atom14_atom_exists: torch.Tensor,
    eps=1e-6,
    **kwargs,
) -> torch.Tensor:
#    print("In violation_loss..")  # DEBUG
#    print("atom14_atom_exists", atom14_atom_exists.shape)  # DEBUG
#    print("violations['between_residues']", debug.print_shape(violations["between_residues"]))  # DEBUG
#    print("violations['within_residues']", debug.print_shape(violations["within_residues"]))  # DEBUG
    #. atom14_atom_exists [*, N_res, 14]

    #. violations['between_residues']
    #. bonds_c_n_loss_mean [*,]
    #. angles_ca_c_n_loss_mean [*,]
    #. angles_c_n_ca_loss_mean [*,]
    #. connections_per_residue_loss_sum [*, N_res]
    #. connections_per_residue_violation_mask [*, N_res]
    #. clashes_mean_loss [*,]
    #. clashes_per_atom_loss_sum [*, N_res, 14]
    #. clashes_per_atom_clash_mask [*, N_res, 14]

    #. violations['within_residues']
    #. per_atom_loss_sum [*, N_res, 14]
    #. per_atom_violations [*, N_res, 14]

    # NOTE sum over 'num_atoms' and 'l_clash' only over dim -1 and -2, quick fix for multi-sample mini-batch
    num_atoms = torch.sum(atom14_atom_exists, dim=(-1, -2))  #. [*,]
    l_clash = torch.sum(
        violations["between_residues"]["clashes_per_atom_loss_sum"]
        + violations["within_residues"]["per_atom_loss_sum"],
        dim=(-1, -2)
    )
    l_clash = l_clash / (eps + num_atoms)
    loss = (
        violations["between_residues"]["bonds_c_n_loss_mean"]  #. [*,]
        + violations["between_residues"]["angles_ca_c_n_loss_mean"]  #. [*,]
        + violations["between_residues"]["angles_c_n_ca_loss_mean"]  #. [*,]
        + l_clash
    )

    # NOTE quick fix
    # loss = torch.mean(loss)

    return loss

####################################################################### experimentally-resolved loss #######################################################################

def experimentally_resolved_loss(
    logits: torch.Tensor,
    atom37_atom_exists: torch.Tensor,
    all_atom_mask: torch.Tensor,
    resolution: torch.Tensor,
    min_resolution: float,
    max_resolution: float,
    eps: float = 1e-8,
    **kwargs,
) -> torch.Tensor:
#    print("In experimentally_resolved_loss..")  # DEBUG
#    print("logits: ", logits.shape)  # DEBUG
#    print("atom37_atom_exists: ", atom37_atom_exists.shape)  # DEBUG
#    print("all_atom_mask: ", all_atom_mask.shape)  # DEBUG
#    print("resolution: ", resolution, resolution.shape)  # DEBUG
#    # print("min_resolution: ", min_resolution)  # DEBUG
#    # print("max_resolution: ", max_resolution)  # DEBUG
    #. logits [*, N_res, 37]
    #. atom37_atom_exists [*, N_res, 37]
    #. all_atom_mask [*, N_res, 37]
    #. resolution [*,]

    #. errors [*, N_res, 37]
    errors = sigmoid_cross_entropy(logits, all_atom_mask)
#    print("errors: ", errors.shape)  # DEBUG

    #. loss [*, N_res]
    loss = torch.sum(errors * atom37_atom_exists, dim=-1)
#    print("loss: ", loss.shape)  # DEBUG

    # NOTE added an unsqueeze to make it work for multi-sample minibatch
    #. torch.sum(atom37_atom_exists, dim=(-1, -2)) [*,] -> [*, 1]
    #. loss [*, N_res]
#    print("torch.sum(atom37_atom_exists, dim=(-1, -2)).unsqueeze(-1)", torch.sum(atom37_atom_exists, dim=(-1, -2)).unsqueeze(-1).shape)  # DEBUG
    loss = loss / (eps + torch.sum(atom37_atom_exists, dim=(-1, -2))[..., None])
#    print("loss: ", loss.shape)  # DEBUG

    #. sum over residues
    #. loss [*,]
    loss = torch.sum(loss, dim=-1)
#    print("loss: ", loss.shape)  # DEBUG
    
    #. mask out resolutions outside of range
    #. loss [*,]
    loss = loss * (
        (resolution >= min_resolution) & (resolution <= max_resolution)
    )
#    print("loss: ", loss.shape)  # DEBUG

    #. average over batch
    #. loss []
    # loss = torch.mean(loss)
#    print("loss: ", loss.shape)  # DEBUG
 
    return loss

####################################################################### masked MSA loss #######################################################################

def masked_msa_loss(logits, true_msa, bert_mask, eps=1e-8, **kwargs):
    """
    Computes BERT-style masked MSA loss. Implements subsection 1.9.9.

    Args:
        logits: [*, N_seq, N_res, 23] predicted residue distribution
        true_msa: [*, N_seq, N_res] true MSA
        bert_mask: [*, N_seq, N_res] MSA mask
    Returns:
        Masked MSA loss
    """
#    print("In masked_msa_loss..")  # DEBUG
#    print("logits: ", logits.shape)  # DEBUG
#    print("true_msa: ", true_msa.shape)  # DEBUG
#    print("bert_mask: ", bert_mask.shape)  # DEBUG
    #. logits [*, N_seq, N_res, 23]
    #. true_msa [*, N_seq, N_res]
    #. bert_mask [*, N_seq, N_res]

    #. errors [*, N_seq, N_res]
    errors = softmax_cross_entropy(
        logits, torch.nn.functional.one_hot(true_msa, num_classes=23)
    )
#    print("torch.nn.functional.one_hot(true_msa, num_classes=23)", torch.nn.functional.one_hot(true_msa, num_classes=23).shape)  # DEBUG
#    print("errors: ", errors.shape)  # DEBUG

    # FP16-friendly averaging. Equivalent to:
    # loss = (
    #     torch.sum(errors * bert_mask, dim=(-1, -2)) /
    #     (eps + torch.sum(bert_mask, dim=(-1, -2)))
    # )
    #. loss [*, N_seq, N_res] -> [*, N_seq] -> [*,]
    #. denom [*, N_seq]
    loss = errors * bert_mask
    loss = torch.sum(loss, dim=-1)
    scale = 0.5  #. seems to be used for fp16 numerical stability
    denom = eps + torch.sum(scale * bert_mask, dim=(-1, -2))
    loss = loss / denom[..., None]
    loss = torch.sum(loss, dim=-1)
    loss = loss * scale

    #. average over batch
    # loss = torch.mean(loss)

    return loss

####################################################################### Main class #######################################################################

class AlphaFoldLoss(nn.Module):
    """Aggregation of the various losses described in the supplement"""
    def __init__(self, config):
        super(AlphaFoldLoss, self).__init__()
        self.config = config
        print("In AlphaFoldLoss..")  # DEBUG
        print("Using settings..")  # DEBUG
        print("violation:", self.config.violation)  # DEBUG
        print("distogram:", self.config.distogram)  # DEBUG
        print("experimentally_resolved:", self.config.experimentally_resolved)  # DEBUG
        print("masked_msa:", self.config.masked_msa)  # DEBUG
        print("fape:", self.config.fape)  # DEBU
        print("plddt:", self.config.plddt_loss)  # DEBUG
        print("supervised chi:", self.config.supervised_chi)  # DEBUG
        print("Will scale cum_loss with", self.config.cum_loss_scale)  # DEBUG

    def forward(self, out, batch, _return_breakdown=False):
        if "violation" not in out.keys():
            out["violation"] = find_structural_violations(
                batch,
                out["sm"]["positions"][-1],
                **self.config.violation,
            )

        if "renamed_atom14_gt_positions" not in out.keys():
            batch.update(
                compute_renamed_ground_truth(
                    batch,
                    out["sm"]["positions"][-1],
                )
            )

        loss_fns = {
            # NOTE understanding level: 3/5
            # TODO average over batch outside: trivial
            "distogram": lambda: distogram_loss(
                logits=out["distogram_logits"],
                **{**batch, **self.config.distogram},
            ),

            # NOTE understanding level: 3/5
            # TODO average over batch outside: trivial
            "experimentally_resolved": lambda: experimentally_resolved_loss(
                logits=out["experimentally_resolved_logits"],
                **{**batch, **self.config.experimentally_resolved},
            ),

            # NOTE understanding level: 1/5         remain to go into detail
            # TODO average over batch outside: trivial
            "fape": lambda: fape_loss(
                out,
                batch,
                self.config.fape,
            ),

            # NOTE understanding level: 1/5         remain to go into detail
            # TODO average over batch outside: trivial
            "plddt_loss": lambda: lddt_loss(
                logits=out["lddt_logits"],
                all_atom_pred_pos=out["final_atom_positions"],
                **{**batch, **self.config.plddt_loss},
            ),

            # NOTE understanding level: 3/5
            # TODO average over batch outside: trivial
            "masked_msa": lambda: masked_msa_loss(
                logits=out["masked_msa_logits"],
                **{**batch, **self.config.masked_msa},
            ),

            # NOTE understanding level: 3/5
            # TODO average over batch outside: trivial
            "supervised_chi": lambda: supervised_chi_loss(
                out["sm"]["angles"],
                out["sm"]["unnormalized_angles"],
                **{**batch, **self.config.supervised_chi},
            ),

            # NOTE understanding level: 1/5         remain to go into detail
            # TODO average over batch outside: seemingly trivial after a quick fix
            "violation": lambda: violation_loss(
                out["violation"],
                **batch,
            ),
        }

        if(self.config.tm.enabled):
            loss_fns["tm"] = lambda: tm_loss(
                logits=out["tm_logits"],
                **{**batch, **out, **self.config.tm},
            )

        cum_loss = 0.
        losses = {}
        for loss_name, loss_fn in loss_fns.items():
#            print('Computing loss: ', loss_name)  # DEBUG
            weight = self.config[loss_name].weight
            loss = loss_fn()
#            print(loss_name, loss)  # DEBUG
            if(torch.isnan(loss).any() or torch.isinf(loss)).any():
                #for k,v in batch.items():
                #    if(torch.any(torch.isnan(v)) or torch.any(torch.isinf(v))):
                #        logging.warning(f"{k}: is nan")
                #logging.warning(f"{loss_name}: {loss}")
                logging.warning(f"{loss_name} loss is NaN. Skipping...")
                loss = loss.new_tensor(0., requires_grad=True)
            cum_loss = cum_loss + weight * loss
            losses[loss_name] = loss.detach().clone().cpu().numpy()

        losses["unscaled_loss"] = cum_loss.detach().clone().cpu().numpy()

        # NOTE could be improved by using the mean seq len of the batch
        #. subsection 1.9.1, this version scales the loss by the mean seq len of the batch
        #. while the original paper scales independently by each sequence length
        # Scale the loss by the square root of the minimum of the crop size and
        # the (average) sequence length. See subsection 1.9.
        # seq_len = torch.mean(batch["seq_length"].float())
        if cum_loss.ndim == 0:
            cum_loss = cum_loss.unsqueeze(0)
        
        n_sample = cum_loss.shape[0]
        crop_len = batch["aatype"].shape[-1]

        seq_len = torch.where(batch["seq_length"].float() < crop_len, batch["seq_length"].float(), torch.Tensor([crop_len]).float().to(batch["seq_length"].device).expand(n_sample))

        # print(cum_loss, cum_loss.shape)
        # print("crop_len", crop_len)
        # print(batch["seq_length"], batch["seq_length"].shape)
        # print("\n\n", seq_len, "\n\n")
        
        cum_loss = torch.mean(cum_loss * torch.sqrt(seq_len))   # min(seq_len, crop_len)
        cum_loss = cum_loss * self.config.cum_loss_scale

        losses["loss"] = cum_loss.detach().clone().cpu().numpy()

        if(not _return_breakdown):
            return cum_loss
        
        return cum_loss, losses

####################################################################### experiment #######################################################################

class EvoformerEmbeddingLoss(nn.Module):
    """Aggregation of the various losses described in the supplement"""
    def __init__(self, config):
        super(EvoformerEmbeddingLoss, self).__init__()
        self.config = config
        print("In EvoformerEmbeddingLoss..")  # DEBUG
        print("Using settings..")  # DEBUG
        print("distogram:", self.config)  # DEBUG

    def forward(self, distogram_logits, batch, _return_breakdown=False):
        no_heads = len(distogram_logits)
        # debug.print_shape(distogram_logits)
        # debug.print_shape(batch)

        cum_loss = 0.
        losses = {}

        for i in range(no_heads):
            loss = distogram_loss(
                logits=distogram_logits[i],
                **{**batch, **self.config},
            )
            cum_loss = cum_loss + self.config.weight * loss
            losses["distogram_%d" % i] = loss.detach().clone().cpu().numpy()

        losses["unscaled_loss"] = cum_loss.detach().clone().cpu().numpy()

        # NOTE could be improved by using the mean seq len of the batch
        #. subsection 1.9.1, this version scales the loss by the mean seq len of the batch
        #. while the original paper scales independently by each sequence length
        # Scale the loss by the square root of the minimum of the crop size and
        # the (average) sequence length. See subsection 1.9.
        # seq_len = torch.mean(batch["seq_length"].float())
        if cum_loss.ndim == 0:
            cum_loss = cum_loss.unsqueeze(0)
        
        n_sample = cum_loss.shape[0]
        crop_len = batch["aatype"].shape[-1]

        seq_len = torch.where(batch["seq_length"].float() < crop_len, batch["seq_length"].float(), torch.Tensor([crop_len]).float().to(batch["seq_length"].device).expand(n_sample))
        
        # NOTE compute mean later
        cum_loss = cum_loss * torch.sqrt(seq_len)  # min(seq_len, crop_len)

        losses["loss"] = cum_loss.detach().clone().cpu().numpy()

        if(not _return_breakdown):
            return cum_loss
        
        return cum_loss, losses


####################################################################### tm-related functions #######################################################################

def _calculate_bin_centers(boundaries: torch.Tensor):
    step = boundaries[1] - boundaries[0]
    bin_centers = boundaries + step / 2
    bin_centers = torch.cat(
        [bin_centers, (bin_centers[-1] + step).unsqueeze(-1)], dim=0
    )
    return bin_centers


def _calculate_expected_aligned_error(
    alignment_confidence_breaks: torch.Tensor,
    aligned_distance_error_probs: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    bin_centers = _calculate_bin_centers(alignment_confidence_breaks)
    return (
        torch.sum(aligned_distance_error_probs * bin_centers, dim=-1),
        bin_centers[-1],
    )


def compute_predicted_aligned_error(
    logits: torch.Tensor,
    max_bin: int = 31,
    no_bins: int = 64,
    **kwargs,
) -> Dict[str, torch.Tensor]:
    """Computes aligned confidence metrics from logits.

    Args:
      logits: [*, num_res, num_res, num_bins] the logits output from
        PredictedAlignedErrorHead.
      max_bin: Maximum bin value
      no_bins: Number of bins
    Returns:
      aligned_confidence_probs: [*, num_res, num_res, num_bins] the predicted
        aligned error probabilities over bins for each residue pair.
      predicted_aligned_error: [*, num_res, num_res] the expected aligned distance
        error for each pair of residues.
      max_predicted_aligned_error: [*] the maximum predicted error possible.
    """
    boundaries = torch.linspace(
        0, max_bin, steps=(no_bins - 1), device=logits.device
    )

    aligned_confidence_probs = torch.nn.functional.softmax(logits, dim=-1)
    (
        predicted_aligned_error,
        max_predicted_aligned_error,
    ) = _calculate_expected_aligned_error(
        alignment_confidence_breaks=boundaries,
        aligned_distance_error_probs=aligned_confidence_probs,
    )

    return {
        "aligned_confidence_probs": aligned_confidence_probs,
        "predicted_aligned_error": predicted_aligned_error,
        "max_predicted_aligned_error": max_predicted_aligned_error,
    }


def compute_tm(
    logits: torch.Tensor,
    residue_weights: Optional[torch.Tensor] = None,
    max_bin: int = 31,
    no_bins: int = 64,
    eps: float = 1e-8,
    **kwargs,
) -> torch.Tensor:
    if residue_weights is None:
        residue_weights = logits.new_ones(logits.shape[-2])

    boundaries = torch.linspace(
        0, max_bin, steps=(no_bins - 1), device=logits.device
    )

    bin_centers = _calculate_bin_centers(boundaries)
    clipped_n = max(torch.sum(residue_weights), 19)

    d0 = 1.24 * (clipped_n - 15) ** (1.0 / 3) - 1.8

    probs = torch.nn.functional.softmax(logits, dim=-1)

    tm_per_bin = 1.0 / (1 + (bin_centers ** 2) / (d0 ** 2))
    predicted_tm_term = torch.sum(probs * tm_per_bin, dim=-1)

    normed_residue_mask = residue_weights / (eps + residue_weights.sum())
    per_alignment = torch.sum(predicted_tm_term * normed_residue_mask, dim=-1)

    weighted = per_alignment * residue_weights
     
    argmax = (weighted == torch.max(weighted)).nonzero()[0]
    return per_alignment[tuple(argmax)]


def tm_loss(
    logits,
    final_affine_tensor,
    backbone_rigid_tensor,
    backbone_rigid_mask,
    resolution,
    max_bin=31,
    no_bins=64,
    min_resolution: float = 0.1,
    max_resolution: float = 3.0,
    eps=1e-8,
    **kwargs,
):
    pred_affine = Rigid.from_tensor_7(final_affine_tensor)
    backbone_rigid = Rigid.from_tensor_4x4(backbone_rigid_tensor)

    def _points(affine):
        pts = affine.get_trans()[..., None, :, :]
        return affine.invert()[..., None].apply(pts)

    sq_diff = torch.sum(
        (_points(pred_affine) - _points(backbone_rigid)) ** 2, dim=-1
    )

    sq_diff = sq_diff.detach()

    boundaries = torch.linspace(
        0, max_bin, steps=(no_bins - 1), device=logits.device
    )
    boundaries = boundaries ** 2
    true_bins = torch.sum(sq_diff[..., None] > boundaries, dim=-1)

    errors = softmax_cross_entropy(
        logits, torch.nn.functional.one_hot(true_bins, no_bins)
    )

    square_mask = (
        backbone_rigid_mask[..., None] * backbone_rigid_mask[..., None, :]
    )

    loss = torch.sum(errors * square_mask, dim=-1)
    scale = 0.5  # hack to help FP16 training along
    denom = eps + torch.sum(scale * square_mask, dim=(-1, -2))
    loss = loss / denom[..., None]
    loss = torch.sum(loss, dim=-1)
    loss = loss * scale

    loss = loss * (
        (resolution >= min_resolution) & (resolution <= max_resolution)
    )

    # Average over the loss dimension
    # loss = torch.mean(loss)

    return loss


####################################################################### amber relaxation-related functions #######################################################################

def find_structural_violations_np(
    batch: Dict[str, np.ndarray],
    atom14_pred_positions: np.ndarray,
    config: ml_collections.ConfigDict,
) -> Dict[str, np.ndarray]:
    to_tensor = lambda x: torch.tensor(x)
    batch = tree_map(to_tensor, batch, np.ndarray)
    atom14_pred_positions = to_tensor(atom14_pred_positions)

    out = find_structural_violations(batch, atom14_pred_positions, **config)

    to_np = lambda x: np.array(x)
    np_out = tensor_tree_map(to_np, out)

    return np_out


def extreme_ca_ca_distance_violations(
    pred_atom_positions: torch.Tensor,  # (N, 37(14), 3)
    pred_atom_mask: torch.Tensor,  # (N, 37(14))
    residue_index: torch.Tensor,  # (N)
    max_angstrom_tolerance=1.5,
    eps=1e-6,
) -> torch.Tensor:
    """Counts residues whose Ca is a large distance from its neighbour.

    Measures the fraction of CA-CA pairs between consecutive amino acids that are
    more than 'max_angstrom_tolerance' apart.

    Args:
      pred_atom_positions: Atom positions in atom37/14 representation
      pred_atom_mask: Atom mask in atom37/14 representation
      residue_index: Residue index for given amino acid, this is assumed to be
        monotonically increasing.
      max_angstrom_tolerance: Maximum distance allowed to not count as violation.
    Returns:
      Fraction of consecutive CA-CA pairs with violation.
    """
    this_ca_pos = pred_atom_positions[..., :-1, 1, :]
    this_ca_mask = pred_atom_mask[..., :-1, 1]
    next_ca_pos = pred_atom_positions[..., 1:, 1, :]
    next_ca_mask = pred_atom_mask[..., 1:, 1]
    has_no_gap_mask = (residue_index[..., 1:] - residue_index[..., :-1]) == 1.0
    ca_ca_distance = torch.sqrt(
        eps + torch.sum((this_ca_pos - next_ca_pos) ** 2, dim=-1)
    )
    violations = (
        ca_ca_distance - residue_constants.ca_ca
    ) > max_angstrom_tolerance
    mask = this_ca_mask * next_ca_mask * has_no_gap_mask
    mean = masked_mean(mask, violations, -1)
    return mean


def compute_violation_metrics(
    batch: Dict[str, torch.Tensor],
    atom14_pred_positions: torch.Tensor,  # (N, 14, 3)
    violations: Dict[str, torch.Tensor],
) -> Dict[str, torch.Tensor]:
    """Compute several metrics to assess the structural violations."""
    ret = {}
    extreme_ca_ca_violations = extreme_ca_ca_distance_violations(
        pred_atom_positions=atom14_pred_positions,
        pred_atom_mask=batch["atom14_atom_exists"],
        residue_index=batch["residue_index"],
    )
    ret["violations_extreme_ca_ca_distance"] = extreme_ca_ca_violations
    ret["violations_between_residue_bond"] = masked_mean(
        batch["seq_mask"],
        violations["between_residues"][
            "connections_per_residue_violation_mask"
        ],
        dim=-1,
    )
    ret["violations_between_residue_clash"] = masked_mean(
        mask=batch["seq_mask"],
        value=torch.max(
            violations["between_residues"]["clashes_per_atom_clash_mask"],
            dim=-1,
        )[0],
        dim=-1,
    )
    ret["violations_within_residue"] = masked_mean(
        mask=batch["seq_mask"],
        value=torch.max(
            violations["within_residues"]["per_atom_violations"], dim=-1
        )[0],
        dim=-1,
    )
    ret["violations_per_residue"] = masked_mean(
        mask=batch["seq_mask"],
        value=violations["total_per_residue_violations_mask"],
        dim=-1,
    )
    return ret


def compute_violation_metrics_np(
    batch: Dict[str, np.ndarray],
    atom14_pred_positions: np.ndarray,
    violations: Dict[str, np.ndarray],
) -> Dict[str, np.ndarray]:
    to_tensor = lambda x: torch.tensor(x)
    batch = tree_map(to_tensor, batch, np.ndarray)
    atom14_pred_positions = to_tensor(atom14_pred_positions)
    violations = tree_map(to_tensor, violations, np.ndarray)

    out = compute_violation_metrics(batch, atom14_pred_positions, violations)

    to_np = lambda x: np.array(x)
    return tree_map(to_np, out, torch.Tensor)

####################################################################### not used in pivotal workflow #######################################################################

#. not used in final-version loss
def torsion_angle_loss(
    a,  # [*, N, 7, 2]
    a_gt,  # [*, N, 7, 2]
    a_alt_gt,  # [*, N, 7, 2]
):
    # [*, N, 7]
    norm = torch.norm(a, dim=-1)

    # [*, N, 7, 2]
    a = a / norm.unsqueeze(-1)

    # [*, N, 7]
    diff_norm_gt = torch.norm(a - a_gt, dim=-1)
    diff_norm_alt_gt = torch.norm(a - a_alt_gt, dim=-1)
    min_diff = torch.minimum(diff_norm_gt ** 2, diff_norm_alt_gt ** 2)

    # [*]
    l_torsion = torch.mean(min_diff, dim=(-1, -2))
    l_angle_norm = torch.mean(torch.abs(norm - 1), dim=(-1, -2))

    an_weight = 0.02
    return l_torsion + an_weight * l_angle_norm

