import os
import unittest

import numpy as np
from spglib import (
    SpglibDataset,
    get_symmetry,
    get_symmetry_dataset,
    get_symmetry_from_database,
)
from vasp import read_vasp

data_dir = os.path.dirname(os.path.abspath(__file__))

dirnames = (
    "cubic",
    "hexagonal",
    "monoclinic",
    "orthorhombic",
    "tetragonal",
    "triclinic",
    "trigonal",
    "distorted",
    "virtual_structure",
)


class TestChangeOfBasis(unittest.TestCase):
    def setUp(self):
        self._filenames = []
        self._datasets: list[SpglibDataset] = []
        self._cells = []
        self._symprecs = []
        for d in dirnames:
            dirname = os.path.join(data_dir, "data", d)
            filenames = os.listdir(dirname)
            self._filenames += [os.path.join(dirname, fname) for fname in filenames]

        for fname in self._filenames:
            cell = read_vasp(fname)
            if "distorted" in fname:
                symprec = 1e-1
            else:
                symprec = 1e-5
            self._datasets.append(get_symmetry_dataset(cell, symprec=symprec))
            self._cells.append(cell)
            self._symprecs.append(symprec)

    def tearDown(self):
        pass

    def test_change_of_basis(self):
        for fname, dataset, cell, symprec in zip(
            self._filenames,
            self._datasets,
            self._cells,
            self._symprecs,
        ):
            std_pos = dataset.std_positions
            tmat = dataset.transformation_matrix
            orig_shift = dataset.origin_shift
            lat = np.dot(cell[0].T, np.linalg.inv(tmat))
            pos = np.dot(cell[1], tmat.T) + orig_shift
            for p in pos:
                diff = std_pos - p
                diff -= np.rint(diff)
                diff = np.dot(diff, lat.T)
                delta = np.sqrt((diff**2).sum(axis=1))
                indices = np.where(delta < symprec)[0]
                self.assertEqual(len(indices), 1, msg=("multi: %s %s" % (p, indices)))

    def test_std_symmetry(self):
        for fname, dataset, cell, symprec in zip(
            self._filenames,
            self._datasets,
            self._cells,
            self._symprecs,
        ):
            symmetry = get_symmetry_from_database(dataset.hall_number)
            std_pos = dataset.std_positions

            # for r, t in zip(symmetry['rotations'], symmetry['translations']):
            #     for rp in (np.dot(std_pos, r.T) + t):
            #         diff = std_pos - rp
            #         diff -= np.rint(diff)
            #         num_match = len(np.where(abs(diff).sum(axis=1) < 1e-3)[0])
            #         self.assertEqual(num_match, 1, msg="%s" % fname)

            # Equivalent above by numpy hack
            # 15 sec on macOS 2.3 GHz Intel Core i5 (4times faster than above)
            rot = symmetry["rotations"]
            trans = symmetry["translations"]
            # (n_sym, 3, n_atom)
            rot_pos = np.dot(rot, std_pos.T) + trans[:, :, None]
            for p in std_pos:
                diff = rot_pos - p[None, :, None]
                diff -= np.rint(diff)
                num_match = (abs(diff).sum(axis=1) < 1e-3).sum(axis=1)
                self.assertTrue((num_match == 1).all(), msg="%s" % fname)

    def test_std_rotation(self):
        for fname, dataset, cell, symprec in zip(
            self._filenames,
            self._datasets,
            self._cells,
            self._symprecs,
        ):
            std_lat = dataset.std_lattice
            tmat = dataset.transformation_matrix
            lat = np.dot(cell[0].T, np.linalg.inv(tmat))
            lat_rot = np.dot(dataset.std_rotation_matrix, lat)
            np.testing.assert_allclose(
                std_lat,
                lat_rot.T,
                atol=symprec,
                err_msg="%s" % fname,
            )

    def _test_get_symmetry(self):
        """Caution!: This test must be executed with spglib compiled with -DSPGTEST."""
        for fname, dataset, cell, symprec in zip(
            self._filenames,
            self._datasets,
            self._cells,
            self._symprecs,
        ):
            cell_spin = cell + (
                [
                    1,
                ]
                * len(cell[2]),
            )
            symmetry = get_symmetry(cell_spin, symprec=symprec)

            self.assertEqual(
                len(dataset.rotations),
                len(symmetry["rotations"]),
                msg=("%s" % fname),
            )

            for r, t in zip(dataset.rotations, dataset.translations):
                found = False
                for r_, t_ in zip(symmetry["rotations"], symmetry["translations"]):
                    if (r == r_).all():
                        diff = t - t_
                        diff -= np.rint(diff)
                        if (abs(diff) < symprec).all():
                            found = True
                            break
                self.assertTrue(found, msg="%s" % fname)


if __name__ == "__main__":
    suite = unittest.TestLoader().loadTestsFromTestCase(TestChangeOfBasis)
    unittest.TextTestRunner(verbosity=2).run(suite)
    # unittest.main()
