import unittest
import numpy as np
from xrd_simulator.detector import Detector
from xrd_simulator.phase import Phase
from xrd_simulator.scatterer import Scatterer
from scipy.spatial import ConvexHull
from xrd_simulator.beam import Beam
import os
import matplotlib.pyplot as plt

class TestDetector(unittest.TestCase):

    def setUp(self):
        self.pixel_size_z = 50.
        self.pixel_size_y = 40.
        self.detector_size = 10000.
        self.d0 = np.array([1,0,0])*self.detector_size
        self.d1 = np.array([1,1,0])*self.detector_size
        self.d2 = np.array([1,0,1])*self.detector_size
        self.detector = Detector( self.pixel_size_z, self.pixel_size_y, self.d0, self.d1, self.d2 )

    def test_init(self):

        for o,otrue in zip(self.detector.d0, np.array([1,0,0])*self.detector_size):
            self.assertAlmostEqual(o,otrue ,msg="detector origin is incorrect")

        for z,ztrue in zip(self.detector.zdhat, np.array([0,0,1])):
            self.assertAlmostEqual(z, ztrue ,msg="zdhat is incorrect")

        for y,ytrue in zip(self.detector.ydhat, np.array([0,1,0])):
            self.assertAlmostEqual(y, ytrue ,msg="ydhat is incorrect")

        self.assertAlmostEqual(self.detector.zmax, self.detector_size, msg="Bad detector dimensions in zmax")
        self.assertAlmostEqual(self.detector.ymax, self.detector_size, msg="Bad detector dimensions in ymax")

        for n,ntrue in zip(self.detector.normal, np.array([-1,0,0])):
            self.assertAlmostEqual(n, ntrue ,msg="Bad detector normal")

    def test_centroid_render(self):
        v = self.detector.ydhat + self.detector.zdhat
        v = v/np.linalg.norm(v)
        verts1 = np.array([[0,0,0],
                          [0,0,1],
                          [0,1,0],
                          [1,0,0]]) + v*np.sqrt(2)*self.detector_size/2. # tetra at detector centre
        ch1 = ConvexHull( verts1 )
        verts2 = np.array([[0,0,0],
                          [0,0,1],
                          [0,1,0],
                          [1,0,0]]) + 2*v*np.sqrt(2)*self.detector_size # tetra out of detector bounds
        ch2 = ConvexHull( verts2 )
        wavelength = 1.0

        incident_wave_vector = 2*np.pi* np.array([1,0,0])/(wavelength)
        scattered_wave_vector = self.d0 + self.pixel_size_y*3*self.detector.ydhat + self.pixel_size_z*2*self.detector.zdhat
        scattered_wave_vector = 2*np.pi*scattered_wave_vector/(np.linalg.norm(scattered_wave_vector)*wavelength)

        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
        phase   = Phase(unit_cell, sgname, path_to_cif_file=data)
        phase.setup_diffracting_planes(wavelength, 0, 20*np.pi/180)

        scatterer1 = Scatterer( ch1, 
                                scattered_wave_vector=scattered_wave_vector, 
                                incident_wave_vector=incident_wave_vector, 
                                wavelength=wavelength,
                                incident_polarization_vector=np.array([0,1,0]), 
                                rotation_axis=np.array([0,0,1]),
                                time=0, 
                                phase=phase, 
                                hkl_indx=0 )
        scatterer2 = Scatterer( ch2, 
                                scattered_wave_vector=scattered_wave_vector, 
                                incident_wave_vector=incident_wave_vector, 
                                wavelength=wavelength,
                                incident_polarization_vector=np.array([0,1,0]), 
                                rotation_axis=np.array([0,0,1]),
                                time=0, 
                                phase=phase, 
                                hkl_indx=0 )

        self.detector.frames.append([scatterer1, scatterer2])
        diffraction_pattern = self.detector.render(frame_number=0, lorentz=False, polarization=False, structure_factor=False, method="centroid")
        
        # the sample sits at the centre of the detector.
        expected_z_pixel = int( self.detector_size / (2*self.pixel_size_z) ) + 2
        expected_y_pixel = int( self.detector_size / (2*self.pixel_size_y) ) + 3

        self.assertAlmostEqual(diffraction_pattern[expected_z_pixel,expected_y_pixel], ch1.volume, msg="detector rendering did not capture scatterer")
        self.assertAlmostEqual(np.sum(diffraction_pattern), ch1.volume, msg="detector rendering captured out of bounds scatterer")

        # Try rendering with advanced intensity model
        diffraction_pattern = self.detector.render(frame_number=0, lorentz=True, polarization=False, structure_factor=False)
        self.assertTrue(diffraction_pattern[expected_z_pixel,expected_y_pixel]!=ch1.volume, msg="detector rendering did not use lorentz factor")
        
        diffraction_pattern = self.detector.render(frame_number=0, lorentz=False, polarization=True, structure_factor=False)
        self.assertTrue(diffraction_pattern[expected_z_pixel,expected_y_pixel]!=ch1.volume, msg="detector rendering did not use polarization factor")
        
        diffraction_pattern = self.detector.render(frame_number=0, lorentz=False, polarization=False, structure_factor=True)
        self.assertTrue(diffraction_pattern[expected_z_pixel,expected_y_pixel]!=ch1.volume, msg="detector rendering did not use structure_factor factor")

    def test_projection_render(self):
        
        # Convex hull of a sphere placed at the centre of the detector
        phi, theta  = np.meshgrid( np.linspace( 0, 2*np.pi, 25 ), np.linspace( 0, 2*np.pi, 25 ), indexing='ij' )
        r     = 1.0*self.detector_size/4.
        x,y,z = r*np.cos( phi )*np.sin( theta ),r*np.sin( phi )*np.sin( theta ),r*np.cos( theta )
        hull_points = np.array( [x.flatten(),y.flatten(),z.flatten()] ).T
        v = self.detector.ydhat + self.detector.zdhat
        v = v/np.linalg.norm(v)
        sphere_hull = ConvexHull( hull_points + v*np.sqrt(2)*self.detector_size/2. )

        # The spherical scatterer forward scatters.
        wavelength = 1.0
        incident_wave_vector = 2*np.pi* np.array([1,0,0])/wavelength
        scattered_wave_vector = 2*np.pi* np.array([1,0,0])/wavelength
        scattered_wave_vector = 2*np.pi*scattered_wave_vector/(np.linalg.norm(scattered_wave_vector)*wavelength)

        # The spherical scatterer is composed of Fe_mp-150 (a pure iron crystal)
        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
        phase   = Phase(unit_cell, sgname, path_to_cif_file=data)
        phase.setup_diffracting_planes(wavelength, 0, 20*np.pi/180)

        scatterer = Scatterer(  sphere_hull, 
                                scattered_wave_vector=scattered_wave_vector, 
                                incident_wave_vector=incident_wave_vector, 
                                wavelength=wavelength,
                                incident_polarization_vector=np.array([0,1,0]), 
                                rotation_axis=np.array([0,0,1]),
                                time=0, 
                                phase=phase,
                                hkl_indx=0 )

        self.detector.frames.append([scatterer])
        diffraction_pattern = self.detector.render(frame_number=0, lorentz=False, polarization=False, structure_factor=False, method="project")

        projected_summed_intensity = np.sum(diffraction_pattern)
        relative_error             = np.abs( scatterer.volume - projected_summed_intensity ) / scatterer.volume
        self.assertLessEqual( relative_error, 1e-4, msg="Projected mass does not match the hull volume" )

        index = np.where(diffraction_pattern==np.max(diffraction_pattern))
        self.assertEqual( index[0][0], self.detector_size//(2*self.pixel_size_z), msg="Projected mass does not match the hull volume" )
        self.assertEqual( index[1][0], self.detector_size//(2*self.pixel_size_y), msg="Projected mass does not match the hull volume" )

        no_pixels_z = int( self.detector_size//self.pixel_size_z )
        no_pixels_y = int( self.detector_size//self.pixel_size_y )


        for i in range(no_pixels_z):
            for j in range(no_pixels_y):
                zd = i*self.pixel_size_z
                yd = j*self.pixel_size_y
                if (zd - self.detector_size/2.)**2 + (yd - self.detector_size/2.)**2 > (1.01*r)**2:
                    self.assertAlmostEqual( diffraction_pattern[i,j], 0 )
                elif (zd - self.detector_size/2.)**2 + (yd - self.detector_size/2.)**2 < (0.99*r)**2:
                    self.assertGreater( diffraction_pattern[i,j], 0 )

    def test_get_intersection(self):

        # central normal algined ray
        ray_direction = np.array([2.23,0.,0.])
        source_point  = np.array([0.,0.,0.])
        z, y  = self.detector.get_intersection(ray_direction, source_point)
        self.assertAlmostEqual(z,0,msg="central detector-normal algined ray does not intersect at 0")
        self.assertAlmostEqual(y,0,msg="central detector-normal algined ray does not intersect at 0")

        # translate the ray
        source_point  += self.detector.ydhat * self.pixel_size_y
        source_point  -= self.detector.zdhat * 2 * self.pixel_size_z
        z, y  = self.detector.get_intersection(ray_direction, source_point)
        self.assertAlmostEqual(z, -2*self.pixel_size_z, msg="translated detector-normal algined ray does not intersect properly")
        self.assertAlmostEqual(y, self.pixel_size_y, msg="translated detector-normal algined ray does not intersect properly")

        # tilt the ray
        ang = np.arctan(self.pixel_size_y/self.detector_size)
        frac = np.tan(ang)*np.linalg.norm(ray_direction)
        ray_direction += self.detector.ydhat * frac * 3
        z, y  = self.detector.get_intersection(ray_direction, source_point)
        self.assertAlmostEqual(z, -2*self.pixel_size_z, msg="translated and tilted ray does not intersect properly")
        self.assertAlmostEqual(y, 4*self.pixel_size_y, msg="translated and tilted ray does not intersect properly")

    def test_contains(self):
        c1 = self.detector.contains(self.detector_size/10., self.detector_size/5.)
        self.assertTrue(c1, msg="detector does no contain included point")
        c2= self.detector.contains(-self.detector_size/8., self.detector_size/3.)
        self.assertTrue(not c2, msg="detector contain negative points")
        c4= self.detector.contains(self.detector_size*2*self.pixel_size_z, self.detector_size/374.)
        self.assertTrue(not c4, msg="detector contain out of bounds points")

    def test_get_wrapping_cone(self):
        wavelength = 1.0
        k = 2 * np.pi * np.array([1,0,0]) / wavelength
        source_point  = (self.detector.zdhat + self.detector.ydhat) * self.detector_size/2.
        opening_angle = self.detector.get_wrapping_cone(k, source_point)

        normed_det_center = (source_point + self.d0)/np.linalg.norm(source_point + self.d0)
        normed_det_origin = self.d0/np.linalg.norm(self.d0)
        expected_angle = np.arccos( np.dot( normed_det_center, normed_det_origin ) ) / 2.

        self.assertAlmostEqual(opening_angle, expected_angle, msg="detector centered wrapping cone has faulty opening angle")

        source_point  = (self.detector.zdhat + self.detector.ydhat) * self.detector_size/2.
        source_point -= self.detector.zdhat*10*self.pixel_size_z
        source_point -= self.detector.ydhat*10*self.pixel_size_y
        opening_angle = self.detector.get_wrapping_cone(k, source_point)
        self.assertGreaterEqual(opening_angle, expected_angle, msg="detector off centered wrapping cone has opening angle")

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