import unittest
import numpy as np
from xrd_simulator.scatterer import Scatterer
from scipy.spatial import ConvexHull
import pkg_resources
from xrd_simulator.phase import Phase
import os

class TestScatterer(unittest.TestCase):

    def setUp(self):
        np.random.seed(10)
        self.wavelength = 1.0
        verts = np.array([[0,0,0],
                          [0,0,1],
                          [0,1,0],
                          [1,0,0],
                          [0.1 ,0.1, 0.1]])
        self.ch = ConvexHull( verts )
        scattered_wave_vector = np.random.rand(3,)
        self.scattered_wave_vector = 2*np.pi*scattered_wave_vector/(self.wavelength * np.linalg.norm(scattered_wave_vector) )
        self.time = np.random.rand()

        data = os.path.join( os.path.join(os.path.dirname(__file__), 'data' ), 'Fe_mp-150_conventional_standard.cif' )
        unit_cell = [3.64570000, 3.64570000, 3.64570000, 90.0, 90.0, 90.0]
        sgname = 'Fm-3m' # Iron
        self.phase   = Phase(unit_cell, sgname, path_to_cif_file=data)
        self.phase.setup_diffracting_planes(self.wavelength, 0, 20*np.pi/180)

        self.incident_wave_vector = np.array([1,0.,0])
        self.incident_wave_vector = 2*np.pi*self.incident_wave_vector/(self.wavelength * np.linalg.norm(self.incident_wave_vector) )
        self.incident_polarization_vector = np.array([0.,1.,0]) 
        self.rotation_axis = np.array([0.,0, 1.])

        self.scatterer = Scatterer( self.ch, 
                                    self.scattered_wave_vector, 
                                    self.incident_wave_vector, 
                                    self.wavelength,
                                    self.incident_polarization_vector, 
                                    self.rotation_axis,
                                    self.time, 
                                    self.phase, 
                                    hkl_indx=0 )

    def test_lorentz(self):
        z = np.array([0,1,1.])
        self.scatterer.scattered_wave_vector = 2*np.pi*z/(np.linalg.norm(z)*self.wavelength)
        L = self.scatterer.lorentz_factor
        self.assertAlmostEqual( L, np.sqrt(2.) )  

        z = np.array([0,0,1.])
        self.scatterer.scattered_wave_vector = 2*np.pi*z/(self.wavelength)
        L = self.scatterer.lorentz_factor
        self.assertTrue( L is np.inf )

    def test_polarization(self):
        self.scatterer.scattered_wave_vector = np.array([0, 1.0, 0])
        P = self.scatterer.polarization_factor
        self.assertAlmostEqual( P, 0 )  

        self.scatterer.scattered_wave_vector = np.array([1./np.sqrt(2), 0, -1./np.sqrt(2)])
        P = self.scatterer.polarization_factor
        self.assertAlmostEqual( P, 1.0 ) 

    def test_hkl(self):
        hkl = self.scatterer.hkl
        for i in range(3):
            self.assertAlmostEqual( hkl[i], -1, msg="hkl is wrong" )

    def test_structure_factor(self):
        structure_factor = self.scatterer.real_structure_factor
        self.assertGreaterEqual( structure_factor, 0, msg="structure factor is wrong" )

    def test_centroid(self):
        centroid2 = self.scatterer.centroid
        for c1,c2 in zip(np.array([0.25, 0.25, 0.25]), centroid2):
            self.assertAlmostEqual( c1, c2, msg="centroid is wrong" )

    def test_volume(self):
        vol = self.scatterer.volume
        self.assertAlmostEqual( vol, 1/6., msg="volume is wrong" )

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