# -*- coding: utf-8 -*-
import unittest
from copy import deepcopy
from math import degrees, radians

import numpy as np
from numpy.linalg import norm

from crystals import Crystal, Lattice, LatticeSystem
from crystals.lattice import lattice_vectors_from_parameters
from crystals.affine import rotation_matrix

np.random.seed(23)


def angle_between(v1, v2):
    """ Returns the angle in degrees between vectors `v1` and `v2`"""
    v1 /= norm(v1)
    v2 /= norm(v2)
    rad = np.arccos(np.clip(np.dot(v1, v2), -1.0, 1.0))
    return degrees(rad)


class TestEuclidianLattice(unittest.TestCase):
    def setUp(self):
        self.lattice = Lattice(np.eye(3))

    def test_euclidian_lattice(self):
        self.assertTrue(np.array_equal(self.lattice.a1, [1, 0, 0]))
        self.assertTrue(np.array_equal(self.lattice.a2, [0, 1, 0]))
        self.assertTrue(np.array_equal(self.lattice.a3, [0, 0, 1]))

    def test_volume(self):
        self.assertEqual(self.lattice.volume, 1)

    def test_equality(self):
        """ Test equality between identical Lattice instances and copies """
        self.assertEqual(self.lattice, self.lattice)
        self.assertEqual(self.lattice, deepcopy(self.lattice))

        self.assertNotEqual(self.lattice, Lattice(2 * np.eye(3)))


class TestLatticeArray(unittest.TestCase):
    def test_shape(self):
        """ Test that array(Lattice(...)) is always 3x3 """
        arr = np.random.random(size=(3, 3))
        lattice = Lattice(arr)
        self.assertTrue(np.allclose(arr, np.array(lattice)))

    def test_dtype(self):
        """ Test that the data-type of array(Lattice(...)) is respected """
        arr = np.random.random(size=(3, 3))
        lattice = Lattice(arr)
        self.assertTrue(np.array(lattice, dtype=np.int).dtype, np.int)


class TestLatticeMeshes(unittest.TestCase):
    def test_frac_mesh(self):
        """ Test that Lattice.frac_mesh is working as expected compared to numpy.meshgrid """
        lattice = Lattice(np.eye(3))
        x = np.linspace(0, 1, num=8)

        for out, n in zip(lattice.frac_mesh(x), np.meshgrid(x, x, x)):
            self.assertTrue(np.allclose(out, n))

    def test_frac_mesh_two_arr(self):
        """ Test that Lattice.frac_mesh is raising an exception for two input vectors """
        lattice = Lattice(np.eye(3))
        x = np.linspace(0, 1, num=2)

        with self.assertRaises(ValueError):
            lattice.frac_mesh(x, x)

    def test_real_mesh_trivial(self):
        """ Test that Lattice.mesh works identically to Lattice.frac_mesh for trivial lattice """
        lattice = Lattice(np.eye(3))
        x = np.linspace(0, 1, num=8)

        for frac, real in zip(lattice.frac_mesh(x), lattice.mesh(x)):
            self.assertTrue(np.allclose(frac, real))

    def test_real_mesh(self):
        """ Test that Lattice.mesh works as expected """
        lattice = Lattice(2 * np.eye(3))
        x = np.linspace(0, 1, num=8)

        # since lattice is a stretched euclidian lattice, we expect
        # a maximum length of 2
        self.assertTrue(np.max(lattice.mesh(x)[0]) == 2)


class TestLatticeParameters(unittest.TestCase):
    def test_orthorombic(self):
        """ alpha = beta = gamma = 90"""
        a1, a2, a3 = Lattice.from_parameters(2, 1, 5, 90, 90, 90).lattice_vectors
        self.assertTrue(np.allclose(a1, [2, 0, 0]))
        self.assertTrue(np.allclose(a2, [0, 1, 0]))
        self.assertTrue(np.allclose(a3, [0, 0, 5]))

    def test_monoclinic(self):
        """ beta =/= 90 """
        a1, a2, a3 = Lattice.from_parameters(1, 2, 3, 90, 120, 90).lattice_vectors

        with self.subTest("Lengths"):
            self.assertAlmostEqual(norm(a1), 1)
            self.assertAlmostEqual(norm(a2), 2)
            self.assertAlmostEqual(norm(a3), 3)

        with self.subTest("Angles"):
            self.assertAlmostEqual(angle_between(a2, a3), 90)
            self.assertAlmostEqual(angle_between(a3, a1), 120)
            self.assertAlmostEqual(angle_between(a1, a2), 90)

    def test_triclinic(self):
        """ alpha, beta, gama =/= 90 """
        a1, a2, a3 = Lattice.from_parameters(1, 2, 3, 75, 40, 81).lattice_vectors

        with self.subTest("Lengths"):
            self.assertAlmostEqual(norm(a1), 1)
            self.assertAlmostEqual(norm(a2), 2)
            self.assertAlmostEqual(norm(a3), 3)

        with self.subTest("Angles"):
            self.assertAlmostEqual(angle_between(a2, a3), 75)
            self.assertAlmostEqual(angle_between(a3, a1), 40)
            self.assertAlmostEqual(angle_between(a1, a2), 81)

    def test_reciprocal_and_back(self):
        """ Create lattice from parameters, take reciprocal twice, 
        and see if the parameters have changed. """
        triclinic = (3, 4, 20, 45, 90, 126)
        triclinic2 = Lattice.from_parameters(
            *triclinic
        ).reciprocal.reciprocal.lattice_parameters
        self.assertTrue(np.allclose(triclinic, triclinic2))

    def test_lattice_parameters_back_and_forth(self):
        """ Test that the conversion between lattice vectors and lattice parameters
        is working """
        for name in Crystal.builtins:
            with self.subTest(name):
                c = Crystal.from_database(name)
                lv1 = c.lattice_vectors
                params = c.lattice_parameters
                lv2 = lattice_vectors_from_parameters(*params)

                self.assertTrue(np.allclose(lv1, lv2))


class TestLatticeSystems(unittest.TestCase):
    def test_cubic_lattice_system(self):
        """ Test that the Lattice.lattice_system attribute is working properly for cubic lattice """
        self.assertEqual(Lattice(2 * np.eye(3)).lattice_system, LatticeSystem.cubic)

    def test_tetragonal_lattice_system(self):
        """ Test that the Lattice.lattice_system attribute is working properly for tetragonal lattice """
        parameters = (2, 2, 3, 90, 90, 90)
        l = Lattice.from_parameters(*parameters)
        self.assertEqual(l.lattice_system, LatticeSystem.tetragonal)

    def test_rhombohedral_lattice_system(self):
        """ Test that the Lattice.lattice_system attribute is working properly for rhombohedral lattice """
        parameters = (1, 1, 1, 87, 87, 87)
        l = Lattice.from_parameters(*parameters)
        self.assertEqual(l.lattice_system, LatticeSystem.rhombohedral)

    def test_monoclinic_lattice_system(self):
        """ Test that the Lattice.lattice_system attribute is working properly for monoclinic lattice
        including all possible permutations. """
        with self.subTest("permutation 1"):
            parameters = (1, 2, 3, 90, 115, 90)
            l = Lattice.from_parameters(*parameters)
            self.assertEqual(l.lattice_system, LatticeSystem.monoclinic)

        with self.subTest("permutation 2"):
            parameters = (2, 3, 1, 115, 90, 90)
            l = Lattice.from_parameters(*parameters)
            self.assertEqual(l.lattice_system, LatticeSystem.monoclinic)

        with self.subTest("permutation 3"):
            parameters = (3, 1, 2, 90, 90, 115)
            l = Lattice.from_parameters(*parameters)
            self.assertEqual(l.lattice_system, LatticeSystem.monoclinic)

    def test_hexagonal_lattice_system(self):
        """ Test that the Lattice.lattice_system attribute is working properly for hexagonal lattice,
        including all possible permutations of lattice parameters. """
        with self.subTest("Gamma = 120deg"):
            parameters = (2, 2, 3, 90, 90, 120)
            l = Lattice.from_parameters(*parameters)
            self.assertEqual(l.lattice_system, LatticeSystem.hexagonal)

        with self.subTest("alpha = 120deg"):
            parameters = (3, 2, 2, 120, 90, 90)
            l = Lattice.from_parameters(*parameters)
            self.assertEqual(l.lattice_system, LatticeSystem.hexagonal)

        with self.subTest("beta = 120deg"):
            parameters = (2, 3, 2, 90, 120, 90)
            l = Lattice.from_parameters(*parameters)
            self.assertEqual(l.lattice_system, LatticeSystem.hexagonal)

        with self.subTest("Equal lengths"):
            parameters = (2, 2, 2, 90, 120, 90)
            l = Lattice.from_parameters(*parameters)
            self.assertEqual(l.lattice_system, LatticeSystem.hexagonal)

    def test_triclinic_lattice_system(self):
        """ Test that the Lattice.lattice_system attribute is working properly for triclinic lattice """
        l = Lattice.from_parameters(1, 2, 3, 75, 40, 81)
        self.assertEqual(l.lattice_system, LatticeSystem.triclinic)

    def test_graphite(self):
        """ Test that the builtin Crystal for graphite has a hexagonal lattice system """
        graphite = Crystal.from_database("C")
        self.assertEqual(graphite.lattice_system, LatticeSystem.hexagonal)

    def test_lead(self):
        """ Test that the builtin Crystal for lead has a cubic lattice system """
        pb = Crystal.from_database("Pb")
        self.assertEqual(pb.lattice_system, LatticeSystem.cubic)

    def test_vo2(self):
        """ Test that the builtin Crystal for monoclinic M1 VO2 has a monoclinic lattice system """
        vo2 = Crystal.from_database("vo2-m1")
        self.assertEqual(vo2.lattice_system, LatticeSystem.monoclinic)


class TestLatticeMillerScattering(unittest.TestCase):
    def setUp(self):
        self.lattice = Lattice(np.random.random((3, 3)))

    def test_scattering_vector_trivial(self):
        """ Test that Lattice.scattering_vectors is working """
        self.assertTrue(
            np.allclose(self.lattice.scattering_vector((0, 0, 0)), ((0, 0, 0)))
        )

    def test_miller_indices_trivial(self):
        """ Test that Lattice.miller_indices is working """
        self.assertTrue(
            np.allclose(self.lattice.miller_indices((0, 0, 0)), ((0, 0, 0)))
        )

    def test_back_and_forth(self):
        """ Test that Lattice.miller_indices and Lattice.scattering_vector are
        reciprocal to each other """
        h, k, l = np.random.randint(-10, 10, size=(3,))
        vector = self.lattice.scattering_vector((h, k, l))
        hp, kp, lp = self.lattice.miller_indices(vector)

        self.assertAlmostEqual(h, float(hp))
        self.assertAlmostEqual(k, float(kp))
        self.assertAlmostEqual(l, float(lp))


class TestBoundedReflections(unittest.TestCase):
    def setUp(self):
        self.crystal = Crystal.from_database(next(iter(Crystal.builtins)))

    def test_bounded_reflections_negative(self):
        """ Test that negative reflection bounds raise an Exception.
        Otherwise, an infinite number of reflections will be generated """
        with self.assertRaises(ValueError):
            next(self.crystal.bounded_reflections(-1))

    def test_bounded_reflections_zero(self):
        """ Check that bounded_reflections returns (000) for a zero bound """
        reflections = set(self.crystal.bounded_reflections(0))
        self.assertIn((0, 0, 0), reflections)
        self.assertTrue(len(reflections), 1)

    def test_bounded_reflections_all_within_bounds(self):
        """ Check that every reflection is within the bound """
        bound = 10
        vectors = (
            self.crystal.scattering_vector(refl)
            for refl in self.crystal.bounded_reflections(bound)
        )
        for vector in vectors:
            self.assertLessEqual(np.linalg.norm(vector), bound)


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