from __future__ import division

import unittest
import meep as mp

@unittest.skipIf(mp.count_processors() < 2, "MPI specific test")
class TestDivideParallelProcesses(unittest.TestCase):

    def test_divide_parallel_processes(self):        
        resolution = 20

        sxy = 4
        dpml = 1
        cell = mp.Vector3(sxy+2*dpml,sxy+2*dpml)

        pml_layers = [mp.PML(dpml)]

        n = mp.divide_parallel_processes(2)
        fcen = 1.0/(n+1)

        sources = [mp.Source(src=mp.GaussianSource(fcen,fwidth=0.2*fcen),
                             center=mp.Vector3(),
                             component=mp.Ez)]
        
        symmetries = [mp.Mirror(mp.X),
                      mp.Mirror(mp.Y)]

        self.sim = mp.Simulation(cell_size=cell,
                                 resolution=resolution,
                                 sources=sources,
                                 symmetries=symmetries,
                                 boundary_layers=pml_layers)

    
        flux_box = self.sim.add_flux(fcen, 0, 1,
                                     mp.FluxRegion(mp.Vector3(y=0.5*sxy), size=mp.Vector3(sxy)),
                                     mp.FluxRegion(mp.Vector3(y=-0.5*sxy), size=mp.Vector3(sxy), weight=-1),
                                     mp.FluxRegion(mp.Vector3(0.5*sxy), size=mp.Vector3(y=sxy)),
                                     mp.FluxRegion(mp.Vector3(-0.5*sxy), size=mp.Vector3(y=sxy), weight=-1))

        self.sim.run(until_after_sources=30)

        tot_flux = mp.get_fluxes(flux_box)[0]
        
        tot_fluxes = mp.merge_subgroup_data(tot_flux)
        fcens = mp.merge_subgroup_data(fcen)            

        self.assertEqual(fcens[0],1)
        self.assertEqual(fcens[1],0.5)
        self.assertAlmostEqual(tot_fluxes[0], 9.8628728533)
        self.assertAlmostEqual(tot_fluxes[1], 19.6537275387)
        
if __name__ == '__main__':
    unittest.main()