from __future__ import division

import math
import unittest

import meep as mp


class TestMultiLevelAtom(unittest.TestCase):

    def test_multilevel_atom(self):
        resolution = 40
        ncav = 1.5
        Lcav = 1
        dpad = 1
        dpml = 1
        sz = Lcav + dpad + dpml

        cell_size = mp.Vector3(z=sz)
        dimensions = 1
        pml_layers = [mp.PML(dpml, side=mp.High)]

        omega_a = 40
        freq_21 = omega_a / (2 * math.pi)

        gamma_perp = 4
        gamma_21 = (2 * gamma_perp) / (2 * math.pi)

        theta = 1
        sigma_21 = 2 * theta * theta * omega_a

        rate_21 = 0.005
        N0 = 28
        Rp = 0.0051

        t1 = mp.Transition(
            1,
            2,
            pumping_rate=Rp,
            frequency=freq_21,
            gamma=gamma_21,
            sigma_diag=mp.Vector3(sigma_21, sigma_21, sigma_21)
        )
        t2 = mp.Transition(2, 1, transition_rate=rate_21)
        ml_atom = mp.MultilevelAtom(sigma=1, transitions=[t1, t2], initial_populations=[N0])
        two_level = mp.Medium(index=ncav, E_susceptibilities=[ml_atom])

        geometry = [mp.Block(center=mp.Vector3(z=(-0.5 * sz) + (0.5 * Lcav)),
                             size=mp.Vector3(mp.inf, mp.inf, Lcav), material=two_level)]

        sim = mp.Simulation(cell_size=cell_size,
                            resolution=resolution,
                            boundary_layers=pml_layers,
                            geometry=geometry,
                            dimensions=dimensions)

        def field_func(p):
            return 1 if p.z == (-0.5 * sz) + (0.5 * Lcav) else 0

        def check_field(sim):
            fp = sim.get_field_point(mp.Ex, mp.Vector3(z=(-0.5 * sz) + Lcav + (0.5 * dpad))).real
            self.assertAlmostEqual(fp, -2.7110969214986387)

        sim.init_sim()
        sim.fields.initialize_field(mp.Ex, field_func)
        sim.run(mp.at_end(check_field), until=7000)


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