# Revisions copyright 2020 Joao Rodrigues. All rights reserved.
#
# This file is part of the Biopython distribution and governed by your
# choice of the "Biopython License Agreement" or the "BSD 3-Clause License".
# Please see the LICENSE file that should have been included as part of this
# package.

"""Unit tests for disordered atoms in the Bio.PDB module."""

import os
import sys
import tempfile
import unittest
import warnings

import numpy as np

from Bio.PDB import PDBIO
from Bio.PDB import PDBParser
from Bio.PDB.Atom import Atom
from Bio.PDB.Atom import DisorderedAtom


class TestDisordered(unittest.TestCase):
    """Tests for operations on DisorderedEntities."""

    @classmethod
    def setUpClass(cls):
        cls.parser = parser = PDBParser(QUIET=1)
        cls.structure = parser.get_structure("x", "PDB/disordered.pdb")

    def unpack_all_atoms(self, structure):
        """Return a list of all atoms in the structure."""
        return [a for r in structure.get_residues() for a in r.get_unpacked_list()]

    # Copying
    def test_copy_disordered_atom(self):
        """Copies disordered atoms and all their children."""
        resi27 = self.structure[0]["A"][27]
        resi27_copy = resi27.copy()

        self.assertNotEqual(id(resi27), id(resi27_copy))  # did we really copy

        resi27_atoms = resi27.get_unpacked_list()
        resi27_copy_atoms = resi27_copy.get_unpacked_list()
        self.assertEqual(len(resi27_atoms), len(resi27_copy_atoms))

        for ai, aj in zip(resi27_atoms, resi27_copy_atoms):
            self.assertEqual(ai.name, aj.name)

    def test_copy_selected_child(self):
        """Checks if the DisorderedAtom.selected_child reference is it's own child."""
        resi27 = self.structure[0]["A"][27]
        disordered_nh1 = resi27.child_dict["NH1"]
        self.assertTrue(
            disordered_nh1.is_disordered()
        )  # Check if actually DisorderedAtom

        disordered_copy_nh1 = disordered_nh1.copy()

        atom_possibilities = disordered_copy_nh1.disordered_get_list()
        # Re-calculate the proper selected child from the Atom children
        best_atom_of_copy = max(
            atom_possibilities, key=lambda pos_atom: pos_atom.occupancy
        )
        self.assertEqual(best_atom_of_copy, disordered_copy_nh1.selected_child)

    def test_copy_selected_child_is_not_shared_with_original(self):
        """Copy of a DisorderedAtom selects one of its own children, not the original's."""
        original = self.structure[0]["A"][27].child_dict["NH1"]
        self.assertTrue(original.is_disordered())

        clone = original.copy()

        # The bug: the shallow copy kept the original's selected_child, so the
        # clone forwarded every attribute access to an atom of the original.
        self.assertIsNot(clone.selected_child, original.selected_child)
        self.assertTrue(
            any(child is clone.selected_child for child in clone.disordered_get_list())
        )
        self.assertEqual(
            clone.selected_child.get_altloc(), original.selected_child.get_altloc()
        )

        # The visible consequence: transforming the clone appeared to do
        # nothing, because clone.coord read the original's coordinates.
        before = original.coord.copy()
        clone.transform(np.eye(3), np.array([10.0, 0.0, 0.0]))
        self.assertTrue(np.allclose(clone.coord, before + [10.0, 0.0, 0.0]))
        self.assertTrue(np.allclose(original.coord, before))

    def test_copy_zero_occupancy(self):
        """Copy of an all-zero-occupancy DisorderedAtom still selects a child."""
        # disordered_add only selects a child whose occupancy is strictly
        # greater than last_occupancy, so copy() must reset last_occupancy to
        # the same sentinel __init__ uses. Resetting it to 0 would leave
        # selected_child as None here, which a fresh parse never produces.
        original = DisorderedAtom("CZ")
        for altloc, x in (("A", 1.0), ("B", 2.0)):
            original.disordered_add(
                Atom(
                    "CZ",
                    np.array([x, 0.0, 0.0], dtype="f"),
                    26.0,
                    0.00,  # occupancy
                    altloc,
                    " CZ ",
                    1,
                    element="C",
                )
            )
        self.assertEqual(original.selected_child.get_altloc(), "A")

        clone = original.copy()

        self.assertIsNotNone(clone.selected_child)
        self.assertEqual(clone.selected_child.get_altloc(), "A")
        self.assertTrue(np.allclose(clone.coord, [1.0, 0.0, 0.0]))

    def test_copy_preserves_explicit_selection(self):
        """Copy keeps a non-default altloc chosen with disordered_select().

        Re-adding the children selects by occupancy, so a lower-occupancy
        altloc the caller had selected would silently be swapped for the
        highest-occupancy one, changing the copy's coordinates.
        """
        original = DisorderedAtom("CZ")
        for altloc, x, occ in (("A", 1.0, 0.7), ("B", 2.0, 0.3)):
            original.disordered_add(
                Atom(
                    "CZ",
                    np.array([x, 0.0, 0.0], dtype="f"),
                    26.0,
                    occ,
                    altloc,
                    " CZ ",
                    1,
                    element="C",
                )
            )
        # Occupancy would pick A; the caller overrides to B.
        original.disordered_select("B")
        self.assertTrue(np.allclose(original.coord, [2.0, 0.0, 0.0]))

        clone = original.copy()

        self.assertEqual(clone.selected_child.get_altloc(), "B")
        self.assertTrue(np.allclose(clone.coord, [2.0, 0.0, 0.0]))

    def test_copy_entire_chain(self):
        """Copy propagates throughout SMCRA object."""
        s = self.structure
        s_copy = s.copy()

        self.assertNotEqual(id(s), id(s_copy))  # did we really copy

        atoms = self.unpack_all_atoms(s)
        copy_atoms = self.unpack_all_atoms(s_copy)
        self.assertEqual(len(atoms), len(copy_atoms))

        for ai, aj in zip(atoms, copy_atoms):
            self.assertEqual(ai.name, aj.name)

    # Transforming
    def test_transform_disordered(self):
        """Transform propagates through disordered atoms."""
        # This test relates to issue #455 where applying a transformation
        # to a copied structure did not work for disordered atoms.
        s = self.structure
        s_copy = s.copy()

        mtx = ((1, 0, 0), (0, 1, 0), (0, 0, 1))
        tr_vec = (20.0, 0.0, 0.0)

        s_copy.transform(mtx, tr_vec)  # transform copy

        atoms = self.unpack_all_atoms(s)
        copy_atoms = self.unpack_all_atoms(s_copy)
        self.assertEqual(len(atoms), len(copy_atoms))
        for ai, aj in zip(atoms, copy_atoms):
            self.assertEqual(ai - aj, 20.0)  # check distance == 20.0

    # Extract and write
    def test_copy_and_write_disordered(self):
        """Extract, save, and parse again disordered atoms."""
        writer = PDBIO()

        s = self.structure

        # Extract the chain object
        chain = s[0]["A"]

        writer.set_structure(chain)

        filenumber, filename = tempfile.mkstemp()  # save to temp file
        os.close(filenumber)
        try:
            writer.save(filename)

            # Parse again
            s2 = self.parser.get_structure("x_copy", filename)

            # Do we have the same stuff?
            atoms1 = self.unpack_all_atoms(s)
            atoms2 = self.unpack_all_atoms(s2)
            self.assertEqual(len(atoms1), len(atoms2))
            for ai, aj in zip(atoms1, atoms2):
                self.assertEqual(ai.name, aj.name)

        finally:
            os.remove(filename)

    # DisorderedAtom.center_of_mass
    def test_structure_w_disordered_com(self):
        """Calculate center of mass of structure including DisorderedAtoms."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            s = self.parser.get_structure("b", "PDB/disordered.pdb")

        com = s.center_of_mass()

        self.assertTrue(np.allclose(com, [54.545, 19.868, 31.212], atol=1e-3))

    def test_disordered_cog(self):
        """Calculate DisorderedAtom center of geometry."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            s = self.parser.get_structure("b", "PDB/disordered.pdb")

        arg27 = s[0]["A"][27]

        # Detach all children but NH1
        for atom in list(arg27):
            if atom.name != "NH1":
                arg27.detach_child(atom.name)

        res_cog = arg27.center_of_mass()
        self.assertTrue(np.allclose(res_cog, [59.555, 21.033, 25.954], atol=1e-3))

        # Now compare to DisorderedAtom.center_of_mass
        da_cog = arg27["NH1"].center_of_mass()
        self.assertTrue(np.allclose(res_cog, da_cog, atol=1e-3))

    def test_empty_disordered(self):
        """Raise ValueError on center of mass calculation of empty DisorderedAtom."""
        da = DisorderedAtom("dummy")
        with self.assertRaises(ValueError):
            da.center_of_mass()

    # disordered_remove
    def test_remove_disordered_residue(self):
        """Remove residues from DisorderedResidue entities."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            s = self.parser.get_structure("a", "PDB/a_structure.pdb")

        # Residue 10 of chain A is disordered
        disres = s[1]["A"][(" ", 10, " ")]
        self.assertEqual(disres.is_disordered(), 2)
        self.assertEqual(len(disres.child_dict), 2)  # GLY and SER

        disres.disordered_remove("GLY")
        self.assertEqual(len(disres.child_dict), 1)
        self.assertEqual(disres.resname, "SER")  # selects new child?

        disres.disordered_remove("SER")
        self.assertEqual(len(disres.child_dict), 0)
        self.assertIsNone(disres.selected_child)
        with self.assertRaises(AttributeError):
            _ = disres.resname

        self.assertEqual(str(disres), "<Empty DisorderedResidue>")

        # Should still be in chain with the same id though
        # Up to the user to detach the DisorderedResidue from its parent.
        disres = s[1]["A"][(" ", 10, " ")]

    def test_remove_disordered_atom(self):
        """Remove altlocs from DisorderedAtom entities."""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            s = self.parser.get_structure("a", "PDB/a_structure.pdb")

        # Residue 3 of chain A is disordered
        disres = s[1]["A"][(" ", 3, " ")]
        self.assertEqual(disres.is_disordered(), 1)

        disatom = disres["N"]
        self.assertEqual(len(disatom.child_dict), 2)  # "A" and " "

        # Add new atom with bogus occupancy to test selection on removal
        atom = disatom.child_dict["A"].copy()
        atom.altloc = "B"
        atom.occupancy = 2.0
        disatom.disordered_add(atom)

        disatom.disordered_remove(" ")
        self.assertEqual(len(disatom.child_dict), 2)
        self.assertEqual(disatom.altloc, "B")

        # Remove all children
        disatom.disordered_remove("A")
        disatom.disordered_remove("B")

        self.assertEqual(len(disatom.child_dict), 0)
        self.assertIsNone(disatom.selected_child)
        self.assertEqual(disatom.last_occupancy, -sys.maxsize)
        with self.assertRaises(AttributeError):
            _ = disatom.altloc

        self.assertEqual(str(disatom), "<Empty DisorderedAtom N>")

        disatm = s[1]["A"][(" ", 3, " ")]["N"]


if __name__ == "__main__":
    runner = unittest.TextTestRunner(verbosity=2)
    unittest.main(testRunner=runner)
