import unittest
from unittest import TestCase

from shelxfile.misc.misc import frac_to_cart
from shelxfile.shelx.shelx import Shelxfile


class TestAtoms(TestCase):

    def setUp(self) -> None:
        self.shx = Shelxfile()
        self.shx.read_file('tests/resources/p21c.res')

    def test_u_eq_of_hydrogen_atom_in_single_afix(self):
        h = self.shx.atoms.get_atom_by_name('H34')
        self.assertEqual('C34', h.pivot.name)
        self.assertAlmostEqual(0.02956060, round(h.Uiso, 8))
        self.assertAlmostEqual(h.pivot.Uiso * 1.2, h.Uiso)

    def test_u_eq_of_hydrogen_atom_in_afix_137(self):
        h1 = self.shx.atoms.get_atom_by_name('H36A')
        h2 = self.shx.atoms.get_atom_by_name('H36B')
        h3 = self.shx.atoms.get_atom_by_name('H36C')
        self.assertEqual('C36', h1.pivot.name)
        self.assertEqual('C36', h2.pivot.name)
        self.assertEqual('C36', h3.pivot.name)
        self.assertAlmostEqual(0.05117007, round(h1.Uiso, 8))
        self.assertAlmostEqual(0.05117007, round(h2.Uiso, 8))
        self.assertAlmostEqual(0.05117007, round(h3.Uiso, 8))

    def test_uiso_of_pivot_of_H36x(self):
        c36 = self.shx.atoms.get_atom_by_name('C36')
        h36a = self.shx.atoms.get_atom_by_name('H36A')
        self.assertIsNone(c36.pivot)
        self.assertAlmostEqual(0.03411338, round(c36.Uiso, 8))
        # H36A/B/C Uiso = 1.5 × C36 Ueq
        self.assertAlmostEqual(c36.Uiso * 1.5, h36a.Uiso, places=12)

    def test_number(self):
        self.assertEqual(148, self.shx.atoms.number)

    def test_has_atom_al1(self):
        self.assertEqual(True, self.shx.atoms.has_atom('Al1'))

    def test_has_atom_al1_lower(self):
        self.assertEqual(True, self.shx.atoms.has_atom('al1'))

    def test_has_atom_H37A(self):
        self.assertEqual(True, self.shx.atoms.has_atom('H37A'))
        self.assertEqual(True, self.shx.atoms.has_atom('H37a_0'))

    def test_has_atom_al1_0(self):
        self.assertEqual(True, self.shx.atoms.has_atom('Al1_0'))

    def test_has_atom_al2_0(self):
        self.assertEqual(False, self.shx.atoms.has_atom('Al2_0'))

    def test_get_atom_by_name(self):
        self.assertEqual("Atom ID: 73", self.shx.atoms.get_atom_by_name('Al1').__repr__())

    def test_get_all_atomcoordinates(self):
        coordinate = {'O1_4': (0.074835, 0.238436, 0.402457), 'C1_4': (0.028576, 0.234542, 0.337234)}
        self.assertEqual(coordinate['O1_4'], self.shx.atoms.get_all_atomcoordinates()['O1_4'])
        self.assertEqual(coordinate['C1_4'], self.shx.atoms.get_all_atomcoordinates()['C1_4'])

    def test_residues(self):
        self.assertEqual([0, 1, 2, 3, 4], self.shx.atoms.residues)

    def test_q_peaks(self):
        self.assertEqual(['Atom ID: 328', 'Atom ID: 329', 'Atom ID: 330', 'Atom ID: 331', 'Atom ID: 332'],
                         [x.__repr__() for x in self.shx.atoms.q_peaks[:5]])

    def test_distance(self):
        self.assertEqual(2.154399, round(self.shx.atoms.distance('F1_2', 'F2_2'), 6))

    def test_distance2(self):
        self.assertEqual(1.332854, round(self.shx.atoms.distance('C2_2', 'F1_2'), 6))

    def test_angle(self):
        at1 = self.shx.atoms.get_atom_by_name('O1_4')
        at2 = self.shx.atoms.get_atom_by_name('C1_4')
        at3 = self.shx.atoms.get_atom_by_name('C2_4')
        # The angle between the three atoms is ... degree:
        self.assertEqual(109.688123, round(self.shx.atoms.angle(at1, at2, at3), 6))

    def test_torsion_angle(self):
        at1 = self.shx.atoms.get_atom_by_name('O1')
        at2 = self.shx.atoms.get_atom_by_name('C1')
        at3 = self.shx.atoms.get_atom_by_name('C2')
        at4 = self.shx.atoms.get_atom_by_name('F1')
        # The torsion angle between the four atoms is ... degree:
        self.assertEqual(74.095731, round(self.shx.atoms.torsion_angle(at1, at2, at3, at4), 6))

        # Thisis the other way round:
        at4 = self.shx.atoms.get_atom_by_name('F2')
        self.assertEqual(-44.467358, round(self.shx.atoms.torsion_angle(at1, at2, at3, at4), 6))

    def test_atoms_in_class(self):
        atoms = ['O1', 'C1', 'C2', 'F1', 'F2', 'F3', 'C3', 'F4', 'F5', 'F6', 'C4', 'F7', 'F8', 'F9']
        self.assertEqual(atoms, self.shx.atoms.atoms_in_class('CCF3'))

    def test_sfac_num(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        self.assertEqual(1, at.sfac_num)

    def test_element_c(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        self.assertEqual('C', at.element)

    def test_element_o(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        at.element = 'O'
        self.assertEqual('O', at.element)
        self.assertEqual(3, at.sfac_num)

    def test_new_element(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        at.element = 'Na'
        self.assertEqual('Na', at.element)

    def test_new_element_sfacnum(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        self.assertEqual('SFAC C  H  O  F  Al  Ga', str(at.shx.sfac_table))
        at.element = 'Na'
        self.assertEqual(7, at.sfac_num)

    def test_new_element_sfacnum_atom_line(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        at.element = 'Na'
        self.assertEqual(('C1    7    0.028576    0.234542    0.337234   -31.00000    0.02311    '
                          '0.03617      0.01096   -0.01000    0.00201    0.00356'), str(at))

    def test_new_element_sfac_unit(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        at.element = 'Na'
        self.assertEqual('SFAC C  H  O  F  Al  Ga  Na', self.shx.sfac_table.__str__())
        self.assertEqual('UNIT 1  2  3  4  5  6  1', self.shx.unit.__str__())

    def test_an(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        self.assertEqual(6, at.an)

    def test_delete_show(self):
        # We have an atom with ID = 40
        at = self.shx.atoms.get_atom_by_id(40)
        self.assertEqual('Atom ID: 40', at.__repr__())

    def test_delete_show_list(self):
        # Before and after atom 40 are 38 and 42
        self.assertEqual(['Atom ID: 38', 'Atom ID: 40', 'Atom ID: 42'],
                         [x.__repr__() for x in self.shx.atoms.all_atoms[:3]])

    def test_c1_4(self):
        # The name of atom 40 is C1_4
        self.assertEqual('C1_4', self.shx.atoms.get_atom_by_id(40).fullname)

    def test_rename_c1_4(self):
        # The name of atom 40 is C1_4
        atom = self.shx.atoms.get_atom_by_id(40)
        self.assertEqual('C1', atom.name)
        atom.name = 'C99'
        self.assertEqual('C99', atom.name)
        self.assertEqual(('C99   1    0.028576    0.234542    0.337234   -31.00000    0.02311    '
                          '0.03617      0.01096   -0.01000    0.00201    0.00356'), atom.__str__())

    def test_rename_fails(self):
        # The name of atom 40 is C1_4
        atom = self.shx.atoms.get_atom_by_id(40)
        atom.name = 'C99_3'
        # Atom name stays, because 'C99_3' is illegal.
        self.assertEqual('C1', atom.name)

    def test_delete_atom40(self):
        # get atom 40
        at = self.shx.atoms.get_atom_by_id(40)
        # Now delete atom 40
        at.delete()
        # The resulting atoms list does not contain atom 40
        new_atoms = ['Atom ID: 38', 'Atom ID: 41', 'Atom ID: 43']
        self.assertEqual(new_atoms, [x.__repr__() for x in self.shx.atoms.all_atoms[:3]])

    def test_delete2(self):
        self.shx.atoms.get_atom_by_id(40).delete()
        self.assertEqual(None, self.shx.atoms.get_atom_by_id(40))

    def test_delete3(self):
        self.shx.atoms.get_atom_by_id(40).delete()
        self.assertEqual(None, self.shx.atoms.get_atom_by_name('C1_4'))

    def test_equal(self):
        self.assertEqual(True, self.shx.atoms.get_atom_by_id(40) == self.shx.atoms.get_atom_by_name('C1_4'))

    def test_find_atoms_around(self):
        at = self.shx.atoms.get_atom_by_name('C1_4')
        found = [x.__repr__() for x in at.find_atoms_around(dist=2, only_part=2)]
        shouldfind = ['Atom ID: 38', 'Atom ID: 42', 'Atom ID: 50', 'Atom ID: 58']
        self.assertEqual(shouldfind, found)

    def test_get_coordinates(self):
        c = self.shx.atoms.get_atom_by_name('C1_4').cart_coords
        self.assertEqual([-0.1977746458215, 4.9027486970000, 6.8977664006568], [round(x, 13) for x in c])

    def test_to_isotropic(self):
        # We have regular u values of atom 40
        self.assertEqual([0.02311, 0.03617, 0.01096, -0.01, 0.00201, 0.00356], self.shx.atoms.get_atom_by_id(40).uvals)
        # We make them isotropic
        self.shx.atoms.get_atom_by_id(40).to_isotropic()
        # And the result is only one atomic u value:
        self.assertEqual([0.04, 0.0, 0.0, 0.0, 0.0, 0.0], self.shx.atoms.get_atom_by_id(40).uvals)

    def test_cart_coords(self):
        self.assertEqual([-0.1977746458215, 4.9027486970000, 6.8977664006568],
                         [round(x, 13) for x in self.shx.atoms.get_atom_by_id(40).cart_coords])

    def test_cartesian_from_method(self):
        self.assertEqual(
            [round(x, 13) for x in frac_to_cart(self.shx.atoms.get_atom_by_id(40).frac_coords, list(self.shx.cell))],
            [round(x, 13) for x in self.shx.atoms.get_atom_by_id(40).cart_coords])

    def test_frac_coords(self):
        self.assertEqual((0.028576, 0.234542, 0.337234), self.shx.atoms.get_atom_by_id(40).frac_coords)

    def test_frac_x(self):
        self.assertEqual(0.028576, self.shx.atoms.get_atom_by_id(40).x)

    def test_cart_x(self):
        self.assertEqual(-0.1977746458, round(self.shx.atoms.get_atom_by_id(40).xc, 10))

    def test_radius(self):
        self.assertEqual(0.77, self.shx.atoms.get_atom_by_id(40).radius)

    def test_resinum(self):
        self.assertEqual(4, self.shx.atoms.get_atom_by_id(40).resinum)

    def test_resiclass(self):
        self.assertEqual('CCF3', self.shx.atoms.get_atom_by_id(40).resiclass)

    def test_fullname(self):
        self.assertEqual('C1_4', self.shx.atoms.get_atom_by_id(40).fullname)

    def test_part(self):
        self.assertEqual('PART 2 -31', self.shx.atoms.get_atom_by_id(40).part.__repr__())

    def test_part_number(self):
        self.assertEqual(2, self.shx.atoms.get_atom_by_id(40).part.n)

    def test_part_occupancy(self):
        self.assertEqual(0.44236, self.shx.atoms.get_atom_by_id(40).occupancy)

    def test_part_sof(self):
        self.assertEqual(-31.0, self.shx.atoms.get_atom_by_id(40).part.sof)

    def test_get_pivot_atom(self):
        self.assertEqual('C34', self.shx.atoms.all_atoms[45].pivot.name)

    def test_get_pivot_atom_of_fluorine(self):
        # This atom should not have a pivot atom, because it is heavy and not in an AFIX:
        self.assertEqual('F2', self.shx.atoms.all_atoms[4].name)
        self.assertEqual(None, self.shx.atoms.all_atoms[4].pivot)

    def test_hydrogen_atoms(self):
        self.assertEqual('[Atom ID: 134, Atom ID: 141, Atom ID: 148]', str(self.shx.atoms.hydrogen_atoms[:3]))

    def test_riding_atoms(self):
        self.assertEqual('[Atom ID: 134, Atom ID: 141, Atom ID: 148]', str(self.shx.atoms.riding_atoms[:3]))


class TestOccupancyOnSpecialPositions(TestCase):
    """
    Atoms on special positions with a site occupancy factor != 1 combined with a
    free variable reference. FVAR 0.22604 0.76052 0.85152
    """

    def setUp(self) -> None:
        self.shx = Shelxfile()
        self.shx.read_file('tests/resources/p-31c.res')

    def test_positive_fvar_with_site_occupancy_factor(self):
        # 30.33333 -> 0.33333 * FVAR3
        self.assertAlmostEqual(0.2838372, self.shx.atoms.get_atom_by_name('C1').occupancy, places=6)
        # 20.33333 -> 0.33333 * FVAR2
        self.assertAlmostEqual(0.2535041, self.shx.atoms.get_atom_by_name('C12').occupancy, places=6)
        # 10.33333 -> 0.33333 (FVAR1 is 1.0)
        self.assertAlmostEqual(0.33333, self.shx.atoms.get_atom_by_name('N3').occupancy, places=6)

    def test_negative_fvar_with_site_occupancy_factor(self):
        # -30.33333 -> 0.33333 * (1 - FVAR3)
        self.assertAlmostEqual(0.049493, self.shx.atoms.get_atom_by_name("C1'").occupancy, places=6)
        # -20.33333 -> 0.33333 * (1 - FVAR2)
        self.assertAlmostEqual(0.0798259, self.shx.atoms.get_atom_by_name("C12'").occupancy, places=6)

    def test_negative_fvar_without_site_occupancy_factor(self):
        # -31.00000 -> 1.0 * (1 - FVAR3)
        self.assertAlmostEqual(0.14848, self.shx.atoms.get_atom_by_name("N1'").occupancy, places=6)

    def test_afix_without_sof_keeps_atom_sof(self):
        # AFIX 137 without own sof must not overwrite the sof of the atom line:
        h1d = self.shx.atoms.get_atom_by_name('H1D')
        self.assertEqual(-30.33333, h1d.sof)
        self.assertAlmostEqual(0.049493, h1d.occupancy, places=6)
        h23a = self.shx.atoms.get_atom_by_name('H23A')
        self.assertEqual(10.33333, h23a.sof)
        self.assertAlmostEqual(0.33333, h23a.occupancy, places=6)


class TestNegativeUisoResolution(TestCase):
    """
    SHELXL encodes the Uiso of riding atoms as a negative multiplier of the
    Ueq of their pivot atom. Both ueq and Uiso must resolve this encoding.
    """

    def setUp(self) -> None:
        self.shx = Shelxfile()
        self.shx.read_file('tests/resources/p-31c.res')

    def test_ueq_of_riding_hydrogen_is_resolved(self):
        h23a = self.shx.atoms.get_atom_by_name('H23A')
        c23 = self.shx.atoms.get_atom_by_name('C23')
        self.assertEqual(-1.5, h23a.uvals[0])
        self.assertEqual('C23', h23a.pivot.name)
        self.assertAlmostEqual(0.078595, h23a.ueq, places=6)
        self.assertAlmostEqual(c23.ueq * 1.5, h23a.ueq, places=12)

    def test_ueq_of_riding_hydrogen_on_nitrogen(self):
        h1 = self.shx.atoms.get_atom_by_name('H1')
        n1 = self.shx.atoms.get_atom_by_name('N1')
        self.assertAlmostEqual(n1.ueq * 1.3, h1.ueq, places=12)

    def test_uiso_is_alias_of_ueq(self):
        for name in ('H23A', 'C23', 'H1D', "C1'"):
            atom = self.shx.atoms.get_atom_by_name(name)
            self.assertEqual(atom.ueq, atom.Uiso)

    def test_ueq_of_anisotropic_atom_is_unchanged(self):
        c23 = self.shx.atoms.get_atom_by_name('C23')
        self.assertAlmostEqual(0.052397, c23.ueq, places=6)


class TestRidingAtoms(TestCase):
    def setUp(self) -> None:
        self.shx = Shelxfile()
        self.shx.read_file('tests/resources/sad-final.res')

    def test_hydrogen_atoms(self):
        # Returns the list of hydrogen atoms, regardless of their model
        self.assertEqual('[Atom ID: 71, Atom ID: 75, Atom ID: 76]', str(self.shx.atoms.hydrogen_atoms[3:6]))

    def test_riding_atoms(self):
        # Returns only riding atoms, therefore different IDs than above
        self.assertEqual('[Atom ID: 75, Atom ID: 76, Atom ID: 77]', str(self.shx.atoms.riding_atoms[3:6]))

    def test_number_of_anisotropic_atoms(self):
        self.assertEqual(79, self.shx.atoms.n_anisotropic_atoms)

    def test_number_of_isotropic_atoms(self):
        self.assertEqual(49, self.shx.atoms.n_isotropic_atoms)

    def test_number_of_hydrogen_atoms(self):
        self.assertEqual(49, self.shx.atoms.n_hydrogen_atoms)

    def test_number_of_anisotropic_hydrogen_atoms(self):
        self.assertEqual(0, self.shx.atoms.n_anisotropic_hydrogen_atoms)

    def test_number_of_hydrogen_adtoms_with_constrained_u_values(self):
        self.assertEqual(49, self.shx.atoms.n_hydrogen_atoms_with_constr_u_val)


class TestUisoOfFreeRefine(TestCase):
    def setUp(self) -> None:
        self.shx = Shelxfile()
        self.shx.read_file('tests/resources/2240189.res')

    def test_free_refined_hydrogen(self):
        h1a = self.shx.atoms.get_atom_by_name('H1A')
        # TODO: I am usure about the values:
        # self.assertEqual(0.04654, h1a.Uiso)
        # self.assertEqual([0.04654, 0.0, 0.0, 0.0, 0.0, 0.0], h1a.uvals)
        self.assertEqual("O3'", h1a.pivot.name)


class TestBedeLoneFile(TestCase):
    def setUp(self) -> None:
        self.shx = Shelxfile(debug=True)
        self.shx.read_file('tests/resources/test_bedelone.res')

    def test_lone_pairs(self):
        self.assertEqual(150, len(self.shx.atoms))
        l50 = self.shx.atoms.get_atom_by_name('L50')
        self.assertEqual('L50  2   0.822400    0.641000   0.461520   11.00000    0.00000',
                         str(l50))
        self.assertEqual(0, l50.part.n)
        self.assertEqual(11.0, l50.part.sof)
        self.assertEqual([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], l50.uvals)
        self.assertEqual(0.8224, l50.x)
