from __future__ import division

import unittest
import numpy as np
import meep as mp

## Farady rotation rate for gyrotropic Lorentzian medium
def kgyro_lorentzian(freq, epsn, f0, gamma, sigma, b0):
    dfsq = (f0**2 - 1j*freq*gamma - freq**2)
    eperp = epsn + sigma * f0**2 * dfsq / (dfsq**2 - (freq*b0)**2)
    eta = sigma * f0**2 * freq * b0 / (dfsq**2 - (freq*b0)**2)
    return 2*np.pi*freq * np.sqrt(0.5*(eperp - np.sqrt(eperp**2 - eta**2)))

## Farady rotation rate for gyrotropic Drude medium
def kgyro_drude(freq, epsn, f0, gamma, sigma, b0):
    dfsq = - 1j*freq*gamma - freq**2
    eperp = epsn + sigma * f0**2 * dfsq / (dfsq**2 - (freq*b0)**2)
    eta = sigma * f0**2 * freq * b0 / (dfsq**2 - (freq*b0)**2)
    return 2*np.pi*freq * np.sqrt(0.5*(eperp - np.sqrt(eperp**2 - eta**2)))

## Farady rotation rate for Landau-Lifshitz-Gilbert medium
def kgyro_llg(freq, epsn, f0, gamma, sigma, alpha):
    df1 = f0 - 1j*freq*alpha
    df2 = freq + 1j*gamma
    eperp = epsn + sigma * df1/(df1**2 - df2**2)
    eta = sigma * df2 / (df1**2 - df2**2)
    return 2*np.pi*freq * np.sqrt(0.5*(eperp - np.sqrt(eperp**2 - eta**2)))

class TestFaradayRotation(unittest.TestCase):
    ## Simulate a linearly polarized plane wave traveling along the gyrotropy axis.
    ## Extract Faraday rotation angle by comparing the Ex and Ey amplitudes, and
    ## compare to a theoretical result corresponding to rotation rate KPRED.
    ## The default acceptable tolerance TOL is 1.5 degrees.
    def check_rotation(self, mat, L, fsrc, zsrc, resolution, tmax, zout, kpred, tol=1.5):
        cell = mp.Vector3(0, 0, L)
        pml_layers = [mp.PML(thickness=1.0, direction=mp.Z)]
        sources = [mp.Source(mp.ContinuousSource(frequency=fsrc),
                             component=mp.Ex, center=mp.Vector3(0, 0, zsrc))]

        self.sim = mp.Simulation(cell_size=cell, geometry=[], sources=sources,
                                 boundary_layers=pml_layers,
                                 default_material=mat, resolution=resolution)

        record_Ex, record_Ey, record_t = [], [], []

        def record_ex_ey(sim):
            record_Ex.append(sim.get_field_point(mp.Ex, mp.Vector3(0, 0, zout)))
            record_Ey.append(sim.get_field_point(mp.Ey, mp.Vector3(0, 0, zout)))
            record_t.append(sim.meep_time())

        self.sim.run(mp.after_time(0.5*tmax, mp.at_every(1e-6, record_ex_ey)), until=tmax)

        ex_rel = np.amax(abs(np.fft.fft(record_Ex)))
        ey_rel = np.amax(abs(np.fft.fft(record_Ey)))
        result = np.arctan2(ey_rel, ex_rel) * 180/np.pi

        Ex_theory = np.abs(np.cos(kpred * (zout - zsrc)).real)
        Ey_theory = np.abs(np.sin(kpred * (zout - zsrc)).real)
        expected = np.arctan2(Ey_theory, Ex_theory) * 180/np.pi

        print("Rotation angle (in degrees): {}, expected {}\n".format(result, expected))
        np.testing.assert_allclose(expected, result, atol=tol)

    def test_faraday_rotation(self):
        L, zsrc, zout = 12.0, -4.5, 4.0
        freq, tmax = 0.8, 100.0
        resolution = 24

        ## Test gyrotropic Lorentzian medium
        epsn, f0, gamma, sn, b0  = 1.5, 1.0, 1e-3, 0.1, 0.15
        susc = [mp.GyrotropicLorentzianSusceptibility(frequency=f0, gamma=gamma, sigma=sn,
                                                      bias=mp.Vector3(0, 0, b0))]
        mat = mp.Medium(epsilon=epsn, mu=1, E_susceptibilities=susc)
        k = kgyro_lorentzian(freq, epsn, f0, gamma, sn, b0)
        print('=' * 24)
        print("Testing Faraday rotation for gyrotropic Lorentzian model...")
        self.check_rotation(mat, L, freq, zsrc, resolution, tmax, zout, k)

        ## Test gyrotropic Drude medium
        susc = [mp.GyrotropicDrudeSusceptibility(frequency=f0, gamma=gamma, sigma=sn,
                                                 bias=mp.Vector3(0, 0, b0))]
        mat = mp.Medium(epsilon=epsn, mu=1, E_susceptibilities=susc)
        k = kgyro_drude(freq, epsn, f0, gamma, sn, b0)
        print('=' * 24)
        print("Testing Faraday rotation for gyrotropic Drude model...")
        self.check_rotation(mat, L, freq, zsrc, resolution, tmax, zout, k)

        ## Test Landau-Lifshitz-Gilbert medium
        alpha = 1e-5
        susc = [mp.GyrotropicSaturatedSusceptibility(frequency=f0, gamma=gamma, sigma=sn,
                                                     alpha=alpha,
                                                     bias=mp.Vector3(0, 0, 1.0))]
        mat = mp.Medium(epsilon=epsn, mu=1, E_susceptibilities=susc)
        k = kgyro_llg(freq, epsn, f0, gamma, sn, alpha)
        print('=' * 24)
        print("Testing Faraday rotation for Landau-Lifshitz-Gilbert model...")
        self.check_rotation(mat, L, freq, zsrc, resolution, tmax, zout, k)

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