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

import unittest

import torch
import torch.nn.functional as F

from openfold3.core.data.primitives.featurization.structure import (
    create_atom_to_token_index,
)
from openfold3.core.utils.atomize_utils import (
    aggregate_atom_feat_to_tokens,
    broadcast_token_feat_to_atoms,
    get_token_atom_index_offset,
    get_token_center_atoms,
    get_token_frame_atoms,
    get_token_representative_atoms,
)


def example1():
    # Standard amino acid residues
    # Protein: ALA GLY
    # NumAtoms: 5 4
    restype = F.one_hot(torch.Tensor([[0, 7]]).long(), num_classes=32).float()

    batch = {
        "restype": restype,
        "asym_id": torch.Tensor([[0, 0]]),
        "token_mask": torch.ones((1, 2)),
        "is_protein": torch.ones((1, 2)),
        "is_dna": torch.zeros((1, 2)),
        "is_rna": torch.zeros((1, 2)),
        "is_atomized": torch.zeros((1, 2)),
        "start_atom_index": torch.Tensor([[0, 5]]),
        "num_atoms_per_token": torch.Tensor([5, 4]),
    }

    x = torch.randn((1, 9, 3))
    atom_mask = torch.ones((1, 9))

    return batch, x, atom_mask


def example2():
    # Modified amino acid residues
    # Protein: ALA GLY/A
    # NumAtoms: 5 4
    restype = F.one_hot(torch.Tensor([[0, 7, 7, 7, 7]]).long(), num_classes=32).float()

    batch = {
        "restype": restype,
        "token_mask": torch.ones((1, 5)),
        "is_protein": torch.ones((1, 5)),
        "is_dna": torch.zeros((1, 5)),
        "is_rna": torch.zeros((1, 5)),
        "is_atomized": torch.Tensor([[0, 1, 1, 1, 1]]),
        "start_atom_index": torch.Tensor([[0, 5, 6, 7, 8]]),
    }

    x = torch.randn((1, 9, 3))
    atom_mask = torch.ones((1, 9))

    return batch, x, atom_mask


def example3():
    # Standard nucleotide residues
    # Protein 1: A U
    # Protein 2: DG DC
    # NumAtoms 1: 22 20
    # NumAtoms 2: 22 19
    restype = F.one_hot(
        torch.Tensor([[21, 24], [27, 28]]).long(), num_classes=32
    ).float()

    batch = {
        "restype": restype,
        "asym_id": torch.Tensor([[0, 0], [0, 1]]),
        "token_mask": torch.ones((2, 2)),
        "is_protein": torch.zeros((2, 2)),
        "is_dna": torch.Tensor([[0, 0], [1, 1]]),
        "is_rna": torch.Tensor([[1, 1], [0, 0]]),
        "is_atomized": torch.zeros((2, 2)),
        "start_atom_index": torch.Tensor([[0, 22], [0, 22]]),
        "num_atoms_per_token": torch.Tensor([[22, 20], [22, 19]]),
    }

    x = torch.randn((2, 42, 3))
    atom_mask = torch.ones((2, 42))
    atom_mask[1, -1] = 0

    return batch, x, atom_mask


def example4():
    # Modified nucleotide residues
    # Protein 1: A U/A
    # Protein 2: DG/A DC
    # NumAtoms 1: 22 20
    # NumAtoms 2: 22 19
    token_mask = torch.Tensor([[1] * 21 + [0] * 2, [1] * 23])
    restype = (
        F.one_hot(
            torch.Tensor([[21] + [24] * 20 + [31] * 2, [27] * 22 + [28]]).long(),
            num_classes=32,
        ).float()
        * token_mask[..., None]
    )

    batch = {
        "restype": restype,
        "token_mask": token_mask,
        "is_protein": torch.zeros((2, 23)),
        "is_rna": torch.Tensor([[1] * 21 + [0] * 2, [0] * 23]),
        "is_dna": torch.Tensor([[0] * 23, [1] * 23]),
        "is_atomized": torch.Tensor([[0] + [1] * 20 + [0] * 2, [1] * 22 + [0]]),
        "start_atom_index": torch.Tensor(
            [[0] + [i for i in range(22, 42)] + [0] * 2, [i for i in range(23)]]
        ),
    }

    x = torch.randn((2, 42, 3))
    atom_mask = torch.ones((2, 42))
    atom_mask[1, -1] = 0

    return batch, x, atom_mask


def example5():
    # Ligands
    # Ligand 1 + GLY (4 atoms)
    # Ligand 2 + A (22 atoms)
    # Ligand 3 + DG (22 atoms)
    # Ligand 4
    token_mask = torch.ones((4, 4))
    token_mask[-1, -1] = 0
    restype = (
        F.one_hot(
            torch.Tensor(
                [
                    [20, 20, 20, 7],
                    [20, 20, 20, 21],
                    [20, 20, 20, 27],
                    [20, 20, 20, 31],
                ]
            ).long(),
            num_classes=32,
        ).float()
        * token_mask[..., None]
    )

    is_protein = torch.zeros((4, 4))
    is_protein[0, -1] = 1
    is_rna = torch.zeros((4, 4))
    is_rna[1, -1] = 1
    is_dna = torch.zeros((4, 4))
    is_dna[2, -1] = 1

    is_atomized = torch.concat([torch.ones((4, 3)), torch.zeros((4, 1))], dim=-1)

    start_atom_index = torch.arange(4).unsqueeze(0).repeat((4, 1))
    start_atom_index = start_atom_index * token_mask

    batch = {
        "restype": restype,
        "token_mask": token_mask,
        "is_protein": is_protein,
        "is_dna": is_dna,
        "is_rna": is_rna,
        "is_atomized": is_atomized,
        "start_atom_index": start_atom_index,
        "num_atoms_per_token": torch.Tensor(
            [[1, 1, 1, 4], [1, 1, 1, 22], [1, 1, 1, 22], [1, 1, 1, 0]]
        ),
        "asym_id": torch.Tensor(
            [[0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 1, 0]]
        ),
    }

    # Example 1: Ligand atom with valid frames
    ligand1 = torch.Tensor([[0, 0, 0], [1, 0, 0], [0, 2, 0]])

    # Example 2: Ligand atom with failed angle constraint (outward colinear)
    ligand2 = torch.Tensor([[0, 0, 0], [1, 0, 0], [-1, 0.1, 0]])

    # Example 3: Ligand atom with failed angle constraint (inward colinear)
    ligand3 = torch.Tensor([[0, 0, 0], [1, 0, 0], [0.9, 0.1, 0]])

    # Example 4: Ligand atom with failed chain constraint
    ligand4 = torch.Tensor([[0, 0, 0], [1, 0, 0], [0, 2, 0]])

    ligands = torch.stack([ligand1, ligand2, ligand3, ligand4], dim=0)
    x = torch.concat([ligands, torch.randn((4, 22, 3))], dim=1)

    atom_mask = torch.Tensor(
        [[1] * 7 + [0] * 18, [1] * 25, [1] * 25, [1] * 3 + [0] * 22]
    )

    return batch, x, atom_mask


def example6():
    # Unknown residues
    # Protein: UNK N DN
    # NumAtoms: 7 12 11
    token_mask = torch.Tensor([[1, 1, 1, 0]])
    restype = (
        F.one_hot(
            torch.Tensor([[20, 25, 30, 31]]).long(),
            num_classes=32,
        ).float()
        * token_mask[..., None]
    )

    batch = {
        "restype": restype,
        "token_mask": token_mask,
        "asym_id": torch.Tensor([[0, 1, 1, 2]]),
        "is_protein": torch.Tensor([[1, 0, 0, 0]]),
        "is_rna": torch.Tensor([[0, 1, 0, 0]]),
        "is_dna": torch.Tensor([[0, 0, 1, 0]]),
        "is_atomized": torch.Tensor([[0, 0, 0, 0]]),
        "start_atom_index": torch.Tensor([[0, 7, 19, 0]]),
        "num_atoms_per_token": torch.Tensor([[7, 12, 11, 0]]),
    }

    x = torch.randn((1, 30, 3))
    atom_mask = torch.ones((1, 30))

    return batch, x, atom_mask


class TestBroadcastTokenFeatToAtoms(unittest.TestCase):
    def test_with_one_batch_dim(self):
        num_atoms_per_token = torch.Tensor([[3, 6, 2, 5, 1], [4, 7, 1, 3, 5]])

        token_mask = torch.Tensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 0]])

        atom_mask = broadcast_token_feat_to_atoms(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
            token_feat=token_mask,
        )

        num_atoms = torch.sum(atom_mask, dim=-1).int()
        gt_num_atoms = torch.sum(num_atoms_per_token * token_mask, dim=-1).int()

        self.assertTrue(atom_mask.shape == (2, gt_num_atoms[0]))
        self.assertTrue((num_atoms == gt_num_atoms).all())

    def test_with_two_batch_dim(self):
        num_atoms_per_token = torch.Tensor(
            [[3, 6, 2, 5, 1], [4, 7, 1, 3, 5]]
        ).unsqueeze(1)

        token_mask = torch.Tensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 0]]).unsqueeze(1)

        atom_mask = broadcast_token_feat_to_atoms(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
            token_feat=token_mask,
        )

        num_atoms = torch.sum(atom_mask, dim=-1).int()
        gt_num_atoms = torch.sum(num_atoms_per_token * token_mask, dim=-1).int()

        self.assertTrue(atom_mask.shape == (2, 1, gt_num_atoms[0]))
        self.assertTrue((num_atoms == gt_num_atoms).all())

    def test_with_two_feat_dim(self):
        c_s = 10

        num_atoms_per_token = torch.Tensor([[2, 2]])
        token_mask = torch.Tensor([[1, 1]])
        si = torch.randn((1, 2, c_s))

        sl = broadcast_token_feat_to_atoms(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
            token_feat=si,
            token_dim=-2,
        )

        sl_gt = torch.repeat_interleave(si, repeats=2, dim=1)

        self.assertTrue(sl.shape == (1, 4, c_s))
        self.assertTrue((torch.abs(sl - sl_gt) < 1e-5).all())

    def test_with_max_num_atoms_per_token(self):
        num_atoms_per_token = torch.Tensor([[3, 6]])

        token_mask = torch.Tensor([[1, 1]])

        atom_mask = broadcast_token_feat_to_atoms(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
            token_feat=token_mask,
            max_num_atoms_per_token=10,
        )

        gt_atom_mask = torch.Tensor(
            [[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]]
        )

        self.assertTrue((atom_mask == gt_atom_mask).all())


class TestAggregateAtomFeatToTokens(unittest.TestCase):
    def test_with_one_batch_dim(self):
        num_atoms_per_token = torch.Tensor([[3, 6, 2, 5, 1], [4, 7, 1, 3, 5]])

        token_mask = torch.Tensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 0]])

        atom_mask = broadcast_token_feat_to_atoms(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
            token_feat=token_mask,
        )
        atom_feat = atom_mask.clone()

        atom_to_token_index = create_atom_to_token_index(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
        )

        token_feat = aggregate_atom_feat_to_tokens(
            token_mask=token_mask,
            atom_to_token_index=atom_to_token_index,
            atom_mask=atom_mask,
            atom_feat=atom_feat,
            atom_dim=-1,
            aggregate_fn="mean",
            eps=1e-9,
        )

        self.assertTrue((token_feat == token_mask).all())

        atom_mask[0, :4] = 0

        token_feat = aggregate_atom_feat_to_tokens(
            token_mask=token_mask,
            atom_to_token_index=atom_to_token_index,
            atom_mask=atom_mask,
            atom_feat=atom_feat,
            atom_dim=-1,
            aggregate_fn="mean",
            eps=1e-9,
        )

        token_mask[0, 0] = 0

        self.assertTrue((token_feat == token_mask).all())

    def test_with_two_batch_dim(self):
        num_atoms_per_token = torch.Tensor(
            [[3, 6, 2, 5, 1], [4, 7, 1, 3, 5]]
        ).unsqueeze(1)

        token_mask = torch.Tensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 0]]).unsqueeze(1)

        atom_mask = broadcast_token_feat_to_atoms(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
            token_feat=token_mask,
        )
        atom_feat = atom_mask.clone()

        atom_to_token_index = create_atom_to_token_index(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
        )

        token_feat = aggregate_atom_feat_to_tokens(
            token_mask=token_mask,
            atom_to_token_index=atom_to_token_index,
            atom_mask=atom_mask,
            atom_feat=atom_feat,
            atom_dim=-1,
            aggregate_fn="mean",
            eps=1e-9,
        )

        self.assertTrue((token_feat == token_mask).all())

    def test_with_two_feat_dim(self):
        num_atoms_per_token = torch.Tensor([[3, 6]])
        token_mask = torch.Tensor([[1, 1]])

        atom_mask = torch.Tensor([[1, 1, 1, 1, 0, 1, 1, 1, 1]])
        atom_feat = torch.randn((1, 9, 5))

        atom_to_token_index = create_atom_to_token_index(
            token_mask=token_mask,
            num_atoms_per_token=num_atoms_per_token,
        )

        token_feat = aggregate_atom_feat_to_tokens(
            token_mask=token_mask,
            atom_to_token_index=atom_to_token_index,
            atom_mask=atom_mask,
            atom_feat=atom_feat,
            atom_dim=-2,
            aggregate_fn="mean",
            eps=1e-9,
        )

        gt_token_feat = torch.concat(
            [
                torch.sum(atom_feat[:, :3], dim=-2, keepdim=True) / 3,
                torch.sum(atom_feat[:, [3, 5, 6, 7, 8]], dim=-2, keepdim=True) / 5,
            ],
            dim=-2,
        )

        self.assertTrue((torch.abs(token_feat - gt_token_feat) < 1e-5).all())


class TestGetTokenAtomIndex(unittest.TestCase):
    def test_with_amino_acid_backbone_residue(self):
        restype = F.one_hot(
            torch.Tensor([[0, 7, 13, 20], [21, 26, 24, 30]]).long(), num_classes=32
        )
        token_atom_index_offset, token_atom_mask = get_token_atom_index_offset(
            atom_name="CA", restype=restype
        )
        gt_token_atom_index_offset = torch.Tensor([[1, 1, 1, 1], [-1, -1, -1, -1]])
        gt_token_atom_mask = torch.Tensor([[1, 1, 1, 1], [0, 0, 0, 0]])
        self.assertTrue((token_atom_index_offset == gt_token_atom_index_offset).all())
        self.assertTrue((token_atom_mask == gt_token_atom_mask).all())

    def test_with_amino_acid_sidechain_residue(self):
        restype = F.one_hot(
            torch.Tensor([[0, 7, 13, 20], [21, 26, 24, 30]]).long(), num_classes=32
        )
        token_atom_index_offset, token_atom_mask = get_token_atom_index_offset(
            atom_name="CB", restype=restype
        )
        gt_token_atom_index_offset = torch.Tensor([[4, -1, 4, 4], [-1, -1, -1, -1]])
        gt_token_atom_mask = torch.Tensor([[1, 0, 1, 1], [0, 0, 0, 0]])
        self.assertTrue((token_atom_index_offset == gt_token_atom_index_offset).all())
        self.assertTrue((token_atom_mask == gt_token_atom_mask).all())

    def test_with_nucleotide_backbone_residue(self):
        restype = F.one_hot(
            torch.Tensor([[0, 7, 20, 31], [21, 25, 26, 30]]).long(), num_classes=32
        )
        token_atom_index_offset, token_atom_mask = get_token_atom_index_offset(
            atom_name="C3'", restype=restype
        )
        gt_token_atom_index_offset = torch.Tensor([[-1, -1, -1, -1], [7, 7, 7, 7]])
        gt_token_atom_mask = torch.Tensor([[0, 0, 0, 0], [1, 1, 1, 1]])
        self.assertTrue((token_atom_index_offset == gt_token_atom_index_offset).all())
        self.assertTrue((token_atom_mask == gt_token_atom_mask).all())

    def test_with_nucleotide_sidechain_residue(self):
        restype = F.one_hot(
            torch.Tensor([[0, 7, 13, 20], [21, 25, 26, 30]]).long(), num_classes=32
        )
        token_atom_index_offset, token_atom_mask = get_token_atom_index_offset(
            atom_name="C4", restype=restype
        )
        gt_token_atom_index_offset = torch.Tensor([[-1, -1, -1, -1], [21, -1, 20, -1]])
        gt_token_atom_mask = torch.Tensor([[0, 0, 0, 0], [1, 0, 1, 0]])
        self.assertTrue((token_atom_index_offset == gt_token_atom_index_offset).all())
        self.assertTrue((token_atom_mask == gt_token_atom_mask).all())


class TestGetTokenCenterAtom(unittest.TestCase):
    def test_with_standard_amino_acid_residues(self):
        batch, x, atom_mask = example1()

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_x = x[:, [1, 6], :]
        gt_center_atom_mask = torch.Tensor([1, 1])
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

        atom_mask[0, 6] = 0

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_x = x[:, [1, 6], :]
        gt_center_atom_mask = torch.Tensor([1, 0])
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

    def test_with_modified_amino_acid_residues(self):
        batch, x, atom_mask = example2()

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_x = x[:, [1, 5, 6, 7, 8], :]

        gt_center_atom_mask = torch.Tensor([1, 1, 1, 1, 1])
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

        atom_mask[0, 8] = 0

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_atom_mask = torch.Tensor([1, 1, 1, 1, 0])
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

    def test_with_standard_nucleotide_residues(self):
        batch, x, atom_mask = example3()

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_x = torch.stack([x[0, [11, 33], :], x[1, [10, 32], :]], dim=0)

        gt_center_atom_mask = torch.Tensor([[1, 1], [1, 1]])
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

        atom_mask[0, 11] = 0

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_atom_mask = torch.Tensor([[0, 1], [1, 1]])
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

    def test_with_modified_nucleotide_residues(self):
        batch, x, atom_mask = example4()

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_x = torch.stack(
            [
                x[0, [11] + [i for i in range(22, 42)] + [0] * 2, :],
                x[1, [i for i in range(22)] + [32], :],
            ],
            dim=0,
        )

        gt_center_atom_mask = torch.Tensor([[1] * 21 + [0] * 2, [1] * 23])
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

        atom_mask[1, 10] = 0

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_atom_mask[1, 10] = 0
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

    def test_with_ligands(self):
        batch, x, atom_mask = example5()

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_x = torch.stack(
            [
                x[0, [0, 1, 2, 4], :],
                x[1, [0, 1, 2, 14], :],
                x[2, [0, 1, 2, 13], :],
                x[3, [0, 1, 2, 0], :],
            ],
            dim=0,
        )

        gt_center_atom_mask = torch.ones((4, 4))
        gt_center_atom_mask[-1, -1] = 0
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

        atom_mask[0, 4] = 0

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_atom_mask[0, -1] = 0
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

    def test_with_unknown_residues(self):
        batch, x, atom_mask = example6()

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_x = x[:, [1, 18, 29, 0], :]

        gt_center_atom_mask = torch.Tensor([[1, 1, 1, 0]])

        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())

        atom_mask[0, 18] = 0

        center_x, center_atom_mask = get_token_center_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_center_atom_mask[0, 1] = 0
        self.assertTrue((torch.abs(center_x - gt_center_x) < 1e-5).all())
        self.assertTrue((center_atom_mask == gt_center_atom_mask).all())


class TestGetTokenRepresentativeAtom(unittest.TestCase):
    def test_with_standard_amino_acid_residues(self):
        batch, x, atom_mask = example1()

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_x = x[:, [4, 6], :]
        gt_rep_atom_mask = torch.Tensor([1, 1])
        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

        atom_mask[0, 4] = 0

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_atom_mask = torch.Tensor([0, 1])
        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

    def test_with_modified_amino_acid_residues(self):
        batch, x, atom_mask = example2()

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_x = x[:, [4, 5, 6, 7, 8], :]
        gt_rep_atom_mask = torch.Tensor([1, 1, 1, 1, 1])
        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

        atom_mask[0, 7] = 0

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_atom_mask = torch.Tensor([1, 1, 1, 0, 1])
        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

    def test_with_standard_nucleotide_residues(self):
        batch, x, atom_mask = example3()

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_x = torch.stack([x[0, [21, 35], :], x[1, [21, 34], :]], dim=0)
        gt_rep_atom_mask = torch.ones((2, 2))

        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

        atom_mask[1, 33] = 0

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

    def test_with_modified_nucleotide_residues(self):
        batch, x, atom_mask = example4()

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_x = torch.stack(
            [
                x[0, [21] + [i for i in range(22, 42)] + [0] * 2, :],
                x[1, [i for i in range(22)] + [34], :],
            ],
            dim=0,
        )
        gt_rep_atom_mask = torch.Tensor([[1] * 21 + [0] * 2, [1] * 23])

        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

        atom_mask[0, 26] = 0

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_atom_mask[0, 5] = 0
        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

    def test_with_ligands(self):
        batch, x, atom_mask = example5()

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_x = torch.stack(
            [
                x[0, [0, 1, 2, 4], :],
                x[1, [0, 1, 2, 24], :],
                x[2, [0, 1, 2, 24], :],
                x[3, [0, 1, 2, 0], :],
            ],
            dim=0,
        )
        gt_rep_atom_mask = torch.ones((4, 4))
        gt_rep_atom_mask[-1, -1] = 0

        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

        atom_mask[2, 1] = 0

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_atom_mask[2, 1] = 0
        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

    def test_wtih_unknown_residues(self):
        batch, x, atom_mask = example6()

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_x = x[:, [4, 0, 0, 0], :]
        gt_rep_atom_mask = torch.Tensor([[1, 0, 0, 0]])

        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())

        atom_mask[0, 4] = 0

        rep_x, rep_atom_mask = get_token_representative_atoms(
            batch=batch, x=x, atom_mask=atom_mask
        )

        gt_rep_atom_mask[0, 0] = 0
        self.assertTrue((torch.abs(rep_x - gt_rep_x) < 1e-5).all())
        self.assertTrue((rep_atom_mask == gt_rep_atom_mask).all())


class TestGetTokenFrameAtom(unittest.TestCase):
    def test_with_standard_amino_acid_residues(self):
        angle_threshold = 25.0
        eps = 1e-8
        inf = 1e9
        batch, x, atom_mask = example1()

        (a, b, c), valid_frame_mask = get_token_frame_atoms(
            batch=batch,
            x=x,
            atom_mask=atom_mask,
            angle_threshold=angle_threshold,
            eps=eps,
            inf=inf,
        )

        gt_a = x[:, [0, 5], :]
        gt_b = x[:, [1, 6], :]
        gt_c = x[:, [2, 7], :]
        gt_valid_frame_mask = torch.Tensor([1, 1])

        self.assertTrue((torch.abs(a - gt_a) < 1e-5).all())
        self.assertTrue((torch.abs(b - gt_b) < 1e-5).all())
        self.assertTrue((torch.abs(c - gt_c) < 1e-5).all())
        self.assertTrue((valid_frame_mask == gt_valid_frame_mask).all())

        atom_mask[0, 7] = 0

        (a, b, c), valid_frame_mask = get_token_frame_atoms(
            batch=batch,
            x=x,
            atom_mask=atom_mask,
            angle_threshold=angle_threshold,
            eps=eps,
            inf=inf,
        )

        gt_valid_frame_mask = torch.Tensor([1, 0])

        self.assertTrue((torch.abs(a - gt_a) < 1e-5).all())
        self.assertTrue((torch.abs(b - gt_b) < 1e-5).all())
        self.assertTrue((torch.abs(c - gt_c) < 1e-5).all())
        self.assertTrue((valid_frame_mask == gt_valid_frame_mask).all())

    def test_with_standard_nucleotide_residues(self):
        angle_threshold = 25.0
        eps = 1e-8
        inf = 1e9
        batch, x, atom_mask = example3()

        (a, b, c), valid_frame_mask = get_token_frame_atoms(
            batch=batch,
            x=x,
            atom_mask=atom_mask,
            angle_threshold=angle_threshold,
            eps=eps,
            inf=inf,
        )

        gt_a = torch.stack(
            [
                x[0, [7, 29], :],
                x[1, [7, 29], :],
            ],
            dim=0,
        )
        gt_b = torch.stack([x[0, [11, 33], :], x[1, [10, 32], :]], dim=0)
        gt_c = torch.stack(
            [
                x[0, [5, 27], :],
                x[1, [5, 27], :],
            ],
            dim=0,
        )
        gt_valid_frame_mask = torch.Tensor([[1, 1], [1, 1]])

        self.assertTrue((torch.abs(a - gt_a) < 1e-5).all())
        self.assertTrue((torch.abs(b - gt_b) < 1e-5).all())
        self.assertTrue((torch.abs(c - gt_c) < 1e-5).all())
        self.assertTrue((valid_frame_mask == gt_valid_frame_mask).all())

        atom_mask[0, 7] = 0

        (a, b, c), valid_frame_mask = get_token_frame_atoms(
            batch=batch,
            x=x,
            atom_mask=atom_mask,
            angle_threshold=angle_threshold,
            eps=eps,
            inf=inf,
        )

        gt_valid_frame_mask = torch.Tensor([[0, 1], [1, 1]])

        self.assertTrue((torch.abs(a - gt_a) < 1e-5).all())
        self.assertTrue((torch.abs(b - gt_b) < 1e-5).all())
        self.assertTrue((torch.abs(c - gt_c) < 1e-5).all())
        self.assertTrue((valid_frame_mask == gt_valid_frame_mask).all())

    def test_with_ligands(self):
        angle_threshold = 25.0
        eps = 1e-8
        inf = 1e9
        batch, x, atom_mask = example5()

        (a, b, c), valid_frame_mask = get_token_frame_atoms(
            batch=batch,
            x=x,
            atom_mask=atom_mask,
            angle_threshold=angle_threshold,
            eps=eps,
            inf=inf,
        )

        gt_a = torch.stack(
            [
                x[0, [1, 0, 0, 3], :],
                x[1, [1, 0, 0, 10], :],
                x[2, [2, 2, 1, 10], :],
                x[3, [1, 0, 0, 0], :],
            ],
            dim=0,
        )
        gt_b = torch.stack(
            [
                x[0, [0, 1, 2, 4], :],
                x[1, [0, 1, 2, 14], :],
                x[2, [0, 1, 2, 13], :],
                x[3, [0, 1, 2, 0], :],
            ],
            dim=0,
        )
        gt_c = torch.stack(
            [
                x[0, [2, 2, 1, 5], :],
                x[1, [2, 2, 1, 8], :],
                x[2, [1, 0, 0, 8], :],
                x[3, [2, 2, 1, 0], :],
            ],
            dim=0,
        )

        # Zero out since no closest neighbors
        a[3, 2, :] = torch.zeros(3)
        c[3, 0, :] = torch.zeros(3)
        c[3, 1, :] = torch.zeros(3)
        c[3, 2, :] = torch.zeros(3)
        gt_a[3, 2, :] = torch.zeros(3)
        gt_c[3, 0, :] = torch.zeros(3)
        gt_c[3, 1, :] = torch.zeros(3)
        gt_c[3, 2, :] = torch.zeros(3)

        gt_valid_frame_mask = torch.Tensor(
            [[1, 1, 1, 1], [0, 0, 0, 1], [0, 1, 1, 1], [0, 0, 0, 0]]
        )

        self.assertTrue((torch.abs(a - gt_a) < 1e-5).all())
        self.assertTrue((torch.abs(b - gt_b) < 1e-5).all())
        self.assertTrue((torch.abs(c - gt_c) < 1e-5).all())
        self.assertTrue((valid_frame_mask == gt_valid_frame_mask).all())

        atom_mask[2, 2] = 0

        (a, b, c), valid_frame_mask = get_token_frame_atoms(
            batch=batch,
            x=x,
            atom_mask=atom_mask,
            angle_threshold=angle_threshold,
            eps=eps,
            inf=inf,
        )

        gt_valid_frame_mask = torch.Tensor(
            [[1, 1, 1, 1], [0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 0]]
        )

        # Zero out since no closest neighbors
        self.assertTrue(
            (
                torch.abs(
                    a * valid_frame_mask[..., None]
                    - gt_a * gt_valid_frame_mask[..., None]
                )
                < 1e-5
            ).all()
        )
        self.assertTrue(
            (
                torch.abs(
                    b * valid_frame_mask[..., None]
                    - gt_b * gt_valid_frame_mask[..., None]
                )
                < 1e-5
            ).all()
        )
        self.assertTrue(
            (
                torch.abs(
                    c * valid_frame_mask[..., None]
                    - gt_c * gt_valid_frame_mask[..., None]
                )
                < 1e-5
            ).all()
        )
        self.assertTrue((valid_frame_mask == gt_valid_frame_mask).all())

    def test_with_unknown_residues(self):
        angle_threshold = 25.0
        eps = 1e-8
        inf = 1e9
        batch, x, atom_mask = example6()

        (a, b, c), valid_frame_mask = get_token_frame_atoms(
            batch=batch,
            x=x,
            atom_mask=atom_mask,
            angle_threshold=angle_threshold,
            eps=eps,
            inf=inf,
        )

        gt_a = x[:, [0, 14, 26, 0], :]
        gt_b = x[:, [1, 18, 29, 0], :]
        gt_c = x[:, [2, 12, 24, 0], :]
        gt_valid_frame_mask = torch.Tensor([[1, 1, 1, 0]])

        self.assertTrue((torch.abs(a - gt_a) < 1e-5).all())
        self.assertTrue((torch.abs(b - gt_b) < 1e-5).all())
        self.assertTrue((torch.abs(c - gt_c) < 1e-5).all())
        self.assertTrue((valid_frame_mask == gt_valid_frame_mask).all())

        atom_mask[0, 12] = 0

        (a, b, c), valid_frame_mask = get_token_frame_atoms(
            batch=batch,
            x=x,
            atom_mask=atom_mask,
            angle_threshold=angle_threshold,
            eps=eps,
            inf=inf,
        )

        gt_valid_frame_mask = torch.Tensor([[1, 0, 1, 0]])

        self.assertTrue((torch.abs(a - gt_a) < 1e-5).all())
        self.assertTrue((torch.abs(b - gt_b) < 1e-5).all())
        self.assertTrue((torch.abs(c - gt_c) < 1e-5).all())
        self.assertTrue((valid_frame_mask == gt_valid_frame_mask).all())


if __name__ == "__main__":
    unittest.main()
