# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.


import unittest

import numpy as np

from pymatgen.core import Lattice
from pymatgen.electronic_structure.core import Magmom, Orbital, Spin


class SpinTest(unittest.TestCase):
    def test_init(self):
        self.assertEqual(int(Spin.up), 1)
        self.assertEqual(int(Spin.down), -1)

    def test_from_int(self):
        self.assertEqual(Spin(1), Spin.up)
        self.assertEqual(Spin(-1), Spin.down)
        self.assertRaises(ValueError, Spin, 0)

    def test_cached(self):
        self.assertEqual(id(Spin(1)), id(Spin.up))


class OrbitalTest(unittest.TestCase):
    def test_init(self):
        for orb in Orbital:
            self.assertEqual(Orbital(orb.value), orb)
        self.assertRaises(ValueError, Orbital, 100)

    def test_cached(self):
        self.assertEqual(id(Orbital(0)), id(Orbital.s))


class MagmomTest(unittest.TestCase):
    def test_init(self):
        # backwards compatibility for scalar-like magmoms
        magmom = Magmom(2.0)
        self.assertEqual(float(magmom), 2.0)
        # backwards compatibility for list-like magmoms
        magmom2 = Magmom([1, 2, 3])
        self.assertEqual(list(magmom2), [1, 2, 3])
        self.assertEqual(magmom2.global_moment.tolist(), [1, 2, 3])
        # non-default saxis, normalized internally
        magmom3 = Magmom([1, 2, 3], saxis=[1, 1, 1])
        self.assertTrue(np.allclose(magmom3.saxis, [np.sqrt(1 / 3.0)] * 3))
        # test construction from known global moment and desired, non-default saxis
        magmom4 = Magmom.from_global_moment_and_saxis([1, 2, 3], saxis=[1, 0, 0])
        self.assertTrue(np.allclose(magmom4.moment, [-3, 2, 1]))
        # test global moments with non-default saxis
        magmom5 = Magmom([-3, 2, 1], saxis=[1, 0, 0])
        self.assertTrue(np.allclose(magmom5.global_moment, [1, 2, 3]))

    def test_get_moments(self):

        # simple cases
        magmom_along_x = Magmom([1, 0, 0])
        self.assertTrue(np.allclose(magmom_along_x.get_moment(saxis=[1, 0, 0]), [0, 0, 1]))

        magmom_along_y = Magmom([0, 1, 0])
        self.assertTrue(np.allclose(magmom_along_y.get_moment(saxis=[0, 1, 0]), [0, 0, 1]))

        # test transformations
        magmoms = [
            [0, 0, 0],
            [0, 0, 1],
            [0, 0, -1],
            [1, 2, 3],
            [-1, 2, 3],
            [-1, -2, -3],
        ]

        for magmom in magmoms:
            magmom1 = Magmom(magmom)
            # transform to non-default saxis
            magmom2 = magmom1.get_00t_magmom_with_xyz_saxis()
            # and back to default saxis
            magmom3 = magmom2.get_xyz_magmom_with_001_saxis()
            self.assertTrue(np.allclose(magmom1.moment, magmom))
            self.assertTrue(np.allclose(magmom1.saxis, [0, 0, 1]))
            self.assertTrue(np.allclose(magmom1.get_moment(saxis=magmom1.saxis), magmom1.moment))
            self.assertTrue(np.allclose(magmom1.get_moment(saxis=magmom2.saxis), magmom2.moment))
            self.assertTrue(np.allclose(magmom2.get_moment(saxis=[0, 0, 1]), magmom1.moment))
            self.assertTrue(np.allclose(magmom2.get_moment(saxis=magmom2.saxis), magmom2.moment))
            self.assertTrue(np.allclose(magmom3.moment, magmom1.moment))

    def test_is_collinear(self):
        magmoms_list = [
            [0, 0, 0],
            [1, 1, 1],
            [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
            [[0, 0, 1], [0, 0, 1], [0, 0, 1]],
            [[0, 0, -1], [0, 0, 1], [0, 0, 1]],
            [[2, 2, 2], [-2, -2, -2], [2, 2, 2]],
        ]
        for magmoms in magmoms_list:
            self.assertEqual(Magmom.are_collinear(magmoms), True)
        ncl_magmoms = [[[0, 0, 1], [0, 0, 1], [1, 2, 3]]]
        self.assertEqual(Magmom.are_collinear(ncl_magmoms), False)

    def test_have_consistent_saxis(self):
        magmom1 = Magmom([1, 2, 3])
        magmom2 = Magmom([1, 2, 3])
        magmom3 = Magmom([1, 2, 3], saxis=[0, 0, -1])
        magmom4 = Magmom([1, 2, 3], saxis=[1, 2, 3])
        self.assertTrue(Magmom.have_consistent_saxis([magmom1, magmom2]))
        self.assertFalse(Magmom.have_consistent_saxis([magmom1, magmom3]))
        self.assertFalse(Magmom.have_consistent_saxis([magmom1, magmom4]))

    def test_get_consistent_set_and_saxis(self):
        magmoms = [1, 1, 2, 2, 0, 0, 2]
        magmoms, saxis = Magmom.get_consistent_set_and_saxis(magmoms)
        self.assertTrue(np.allclose(saxis, [0, 0, 1]))

        magmoms = [[0, 0, 0], [1, 1, 1], [2, 2, 2]]
        magmoms, saxis = Magmom.get_consistent_set_and_saxis(magmoms)
        self.assertTrue(np.allclose(saxis, [np.sqrt(1 / 3.0)] * 3))

    def test_relative_to_crystal_axes(self):
        lattice = Lattice.from_parameters(5, 10, 5, 90, 110, 90)
        moment = [1, 0, 2]
        magmom = Magmom.from_moment_relative_to_crystal_axes(moment, lattice)
        self.assertTrue(np.allclose(magmom.moment, [0.93969262, 0.0, 1.65797986]))
        self.assertTrue(np.allclose(magmom.get_moment_relative_to_crystal_axes(lattice), moment))

    def test_equality(self):
        self.assertTrue(Magmom([1, 1, 1]) == Magmom([1, 1, 1]))
        self.assertFalse(Magmom([1, 1, 2]) == Magmom([1, 1, 1]))
        self.assertTrue(Magmom([0, 0, 10]) == 10)

    def test_negative(self):
        self.assertEqual(-Magmom([1, 2, 3]), Magmom([-1, -2, -3]))


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