#!/usr/bin/env python
#-*- coding: utf-8 -*-
#{{{
"""
This is a simulation of a terahertz pulse propagating through a TiO2 sphere. As the
sphere should exhibit dielectric resonances, also R and T spectra are calculated.

About this script:
 * Written in 2012 by Filip Dominec (dominecf at the server of fzu.cz).
 * Being distributed under BSD license, this script is free as speech after five beers. (First 
   it cost me something and now it is free, but nobody guarantees it will make any sense.)
 * You are encouraged to use and modify it as you need. Feel free to write me.
 * Hereby I thank to the MEEP/python_meep authors and people of meep mailing list who helped me a lot.

Used features: 
  * 3D simulation, 
  * Bloch-periodic walls, 
  * simulation output restricted to slices,
  * structure defined by a (slow) Python callback,
  * lossy medium using polarizability (todo: use measured, not fitted values),
  * complex amplitude r(f), t(f) retrieval (instead of energies)
    '-> compare with T(f), R(f) calculation for a full plane (i. e. for the 0th diffraction order)
  * permittivity and permeability calculation, see http://www.eecs.wsu.edu/~schneidj/ufdtd/
Conventions: 
  * all units in SI (native time unit is then (1 m/c) = 3.33 ns)
  * meep remains in its own namespace ("import meep" instead of "from meep import *")
Geometry:
  * Simulation volume 100x100x200 um, structure: PML, source, monitor_t, air, object, air, monitor_t, PML
  * Sphere (eps=94) with diameter 39 um
    (Comsol, HFSS and experiment show first resonance at 750-800 GHz, second at ca. 1600 GHz.
    The nodal plane is in the middle of the sphere, planparallel to the wavefront.
    The second nodal "plane" is a nearly concentric sphere.)
TODO features: 
    * solve numerical instability (?) for timespan >~ 84 
      (manifests as transversal standing wave)
    * (double axis symmetry)
    * (variable angle of incidence)
    * (resonant modes extraction via HarmInv, done in a branched file) 
    * play http://www.youtube.com/watch?v=xpcUxwpOQ_A during computing
    * 
"""
#}}}
import numpy, subprocess, time, matplotlib
import matplotlib.pyplot as plt
#import meep_mpi as meep
import  meep

def run_bash(cmd): #{{{
    #print cmd
    p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
    out = p.stdout.read().strip()
    return out
#}}}

class epsilon_sphere(meep.Callback):#{{{ ## define the structure point-by-point
    def __init__(self, ref_free_space=False, epsilon=1, size_z=200e-6):
        meep.Callback.__init__(self)
        self.ref_free_space = ref_free_space
        self.epsilon = epsilon
        self.size_z = size_z
    def double_vec(self, r):
        ## Reference option (wave passing through air)
        if self.ref_free_space: return 1. 

        ## Define the structure
        (cx, cy, cz) = (50e-6, 50e-6, self.size_z/2)
        radius = 19.5e-6
        if  ((cx-r.x())**2 + (cy-r.y())**2 + (cz-r.z())**2)**.5 > radius:
            return 1.
        else:
            return self.epsilon
#}}}
class epsilon_slab(meep.Callback):#{{{ ## define the structure point-by-point
    def __init__(self, ref_free_space=False, epsilon=1):
        meep.Callback.__init__(self)
        self.ref_free_space = ref_free_space
        self.epsilon = epsilon
    def double_vec(self, r):
        ## Reference option (wave passing through air)
        if self.ref_free_space: return 1. 

        ## Define the structure
        (cx, cy, cz) = (50e-6, 50e-6, 100e-6)
        half_z_thickness = 25e-6
        if  abs(cz-r.z()) > half_z_thickness:
            return 1.
        else:
            return self.epsilon
#}}}

class sigma_sphere(meep.Callback):#{{{ ## define the structure point-by-point
    def __init__(self, ref_free_space=False, sigma=0., size_z=200e-6):
        meep.Callback.__init__(self)
        self.ref_free_space = ref_free_space
        self.sigma = sigma
        self.size_z = size_z
    def double_vec(self, r):
        ## Reference option (wave passing through air)
        if self.ref_free_space: return 0. 

        ## Define the structure
        (cx, cy, cz) = (50e-6, 50e-6, self.size_z/2)
        radius = 19.5e-6
        if  ((cx-r.x())**2 + (cy-r.y())**2 + (cz-r.z())**2)**.5 > radius:
            return 0.
        else:
            return self.sigma ##

            ## experiments with "s.add_polarizability(meep.DBL1, 9e5/c, 0./c)"
            #return 94.*1e-3 ## resonance shifted by -5 GHz same as gamma=.1e9/c

            ## experiments with "s.add_polarizability(meep.DBL1, 9e11/c, 0.1e9/c)"
            #return 94.*1e-01 ## resonance shifted by -15 GHz
            #return 94.*1e-02 ## resonance shifted by -15 GHz
            #return 94.*1e-03 ## resonance shifted by -5 GHz
            #return 94.*1e-10 ## res. at orig. place
            #return 94.*1e5 ## returned NaN
            #return 94.*1e2 ## returned T=1, R=0 (in the observed region)
#}}}


print "=== User defitinions ===" 
## Structure setup
## Note: The resonant structure must be spaced by at least 250 um from the PML to
## ensure numerical stability for simulation time > 1000 ps. Higher size_z may be
## needed for even longer time.
size_x, size_y, size_z = (100e-6, 100e-6, 400e-6)
resolution = 3e-6
pml_thickness = 30e-6
two_pass = True

## Following constants define the TiO2 material model for 0 - 2 THz frequencies.
## One resonance peak at 5.6 THz is defined for absorption [Baumard1977], the 
## epsilon=22. constant is the remaining high-frequency part of permittivity. 
## Real part of epsilon at 1 THz is ~94.
## Note: To simulate TiO2 sphere without absorption, use: epsilon = 94.; 
##   sigma = 0.000 instead.
epsilon = 22.  # shifts the dispersion curve to obtain realistic value 
sigma = epsilon*3.
omega = 5.6e12      # [Hz]
gamma = 8.1e11      # [Hz]

## Source setup for broadband gaussian pulse
srcFreq = 1207e9    # [Hz]
srcWidth = 2000e9

## FFT analysis
# for more info on instability, see: http://permalink.gmane.org/gmane.comp.science.electromagnetism.meep.general/238
# http://www.mail-archive.com/meep-discuss@ab-initio.mit.edu/msg02450.html
timespan = 50
freq_zoom = 2
freq_points = 10#  int(1*timespan*20./freq_zoom)
print "freq_points =", freq_points

#meep.use_Courant(.01)

simulation_name = "eps94pml030_Z400_r%04.1f_t%04.1f_sigma%.2e_gamma%.2e_omega%.2e.png" % (int(resolution*1e6), timespan, sigma, gamma, omega) 
#simulation_name = "Sa_r%04.1f_t%04.1f" % (int(resolution*1e6), timespan)
## GIF output slices
slices=[]
## full Y-Z slice
#slices.append({ "name":"Ex_yz_1396GHz-zoom", "component":meep.Ex, 
            #"geom":meep.volume(meep.vec(size_x/2, 0., 0.), 
                #meep.vec(size_x/2, size_y, size_z)) })
## Y-Z slice around the sphere only
#slices.append({ "name":"Ex_yz_1396GHz-zoom", "component":meep.Ex, 
            #"geom":meep.volume(meep.vec(size_x/2, size_y/2-30e-6, size_z/2-30e-6), 
                #meep.vec(size_x/2, size_y/2+30e-6, size_z/2+30e-6)) })
print meep.my_rank()


t = []
E_r0 = []
E_t0 = []
E_r = []
E_t = []
for ref_free_space in [True, False] if two_pass else [False]:
    print "=== Initialize a computational volume ==="
    c=2.997e8
    #vol = meep.vol3d(size_x, size_y, size_z, int(1./resolution))
    vol = meep.vol3d(size_x, size_y, size_z, 1./resolution)

    print "=== Initialize the structure and fields ==="
    ## Averaging (extremely slow!)
    #meep.use_averaging(True)

    ## Define the refractive index (for dielectric objects).
    ## Use the two sides at Z axis as 100% absorbing (Perfectly Matching Layers),
    ## the remaining sides at X, Y axes are borders to another similar cell (Bloch-periodic).
    m = epsilon_sphere(ref_free_space=ref_free_space, epsilon=epsilon, size_z=size_z) 
    #m = epsilon_slab(ref_free_space=ref_free_space, epsilon=epsilon)
    meep.set_EPS_Callback(m.__disown__())


    perfectly_matched_layers = meep.pml(pml_thickness, meep.Z) # note: if 2nd par. omitted, all faces become PML
    s = meep.structure(vol, meep.EPS, perfectly_matched_layers)

    ## Define the polarizability (-> conductivity, for metallic objects)
    # from https://lists.launchpad.net/python-meep/msg00040.html
    pol1 = sigma_sphere(ref_free_space=ref_free_space, sigma=sigma, size_z=size_z) 
    meep.set_DBL1_Callback(pol1.__disown__())
    s.add_polarizability(meep.DBL1, omega/c, gamma/c)

    f = meep.fields(s)
    f.use_bloch(meep.vec(0.0, size_y, 0.0))
    f.use_bloch(meep.vec(size_x, 0, 0.0))

    ## Export the dielectric structure so that we can visually verify it
    run_bash("mkdir -p output/")
    eps_file =  meep.prepareHDF5File("output/structure.h5")
    f.output_hdf5(meep.Dielectric, vol.surroundings(), eps_file)   
    del(eps_file)

    print "=== Define source ==="
    ## Define source at one of the Z-faces (padded from the absorbing layer)
    f.add_volume_source(
            meep.Ex, 
            meep.gaussian_src_time(srcFreq/c, srcWidth/c), 
            meep.volume(meep.vec(0., 0, pml_thickness), meep.vec(size_x, size_y, pml_thickness)),
            True)

    print "=== Add flux plane to 'measure' the fields ==="
    ## Define the frequencies at which the R, T spectra shall be calculated 
    freq1, freq2 = (srcFreq-srcWidth/freq_zoom, srcFreq+srcWidth/freq_zoom)   
    #freq_points = 2

    ## Define the reflection flux plane (near source)
    flux_0 = f.add_dft_flux_plane(
            meep.volume(meep.vec(0., 0, size_z/2-40e-6), meep.vec(size_x, size_y, size_z/2-40e-6)),
            freq1/c, 
            freq2/c, 
            freq_points)

    ## Define the transmission flux plane (at the side far from source)
    flux_t = f.add_dft_flux_plane(
            meep.volume(meep.vec(0., 0,  size_z/2+70e-6), meep.vec(size_x, size_y, size_z/2+70e-6)),
            freq1/c, 
            freq2/c, 
            freq_points)

    print "=== Simulate the system and save the slices ==="

    if not ref_free_space and two_pass:
        flux_0.load_hdf5(f, "minusflux")

    for slice in slices:
        slice['file'] = meep.prepareHDF5File("output/%s.h5" % slice['name'])
    steps_number = 0 
    images_number = 0 
    print "    Source frequency, width:", srcFreq, srcWidth
    print "    Last source time:", f.last_source_time()/c, "s (virtual:", f.last_source_time(), ")"
    print "    Simulation time:", f.last_source_time()/c*timespan, "s (virtual:", f.last_source_time()*timespan, ")"
    def get_avg_Ex_field_xy_plane(field, sizex, sizey, zpos):
        ## 5x5 grid is optimal (no visible difference between 10x10 grid and 5x5 grid)
        ## 1x1 grid (single on-axis point) gives ca. 2x stronger epsilon in resonance and other errors
        count = 3
        field_sum = 0 
        for x in [x0*sizex/count+(sizex/2/count) for x0 in range(count)]:
            for y in [y0*sizey/count+(sizey/2/count) for y0 in range(count)]:
                field_sum += field.get_field(meep.Ex, meep.vec(x, y, zpos))
        return field_sum/(count**2)


    while (f.time() < (f.last_source_time()*timespan)):
        steps_number += 1 
        f.step()
        if ref_free_space:
            t.append(f.time()/c)
            E_r0.append(get_avg_Ex_field_xy_plane(f, size_x, size_y, size_z/2-40e-6))
            E_t0.append(get_avg_Ex_field_xy_plane(f, size_x, size_y, size_z/2+60e-6))
        else:
            E_r.append(get_avg_Ex_field_xy_plane(f, size_x, size_y, size_z/2-40e-6))
            E_t.append(get_avg_Ex_field_xy_plane(f, size_x, size_y, size_z/2+60e-6))

        if not ref_free_space and steps_number%1==0 and \
                (f.time() > (f.last_source_time()*7.)) and (f.time() < (f.last_source_time()*7.05)):
           images_number += 1 
           for slice in slices:
                f.output_hdf5(slice['component'], slice['geom'], slice['file'], 1) 

    print "=== Retrieve data from the flux plane and plot the R, T graph ==="

    if ref_free_space:
        fluxdata_in_ref = numpy.array(meep.getFluxData(flux_0))
        fluxdata_out_ref = numpy.array(meep.getFluxData(flux_t))


        h5file = open("minusflux.h5", 'w') #truncate file if already exists
        h5file.close()
        flux_0.scale_dfts(-1);
        flux_0.save_hdf5(f, "minusflux")
        flux_0.scale_dfts(-1);
    else:
        fluxdata_in = numpy.array(meep.getFluxData(flux_0))
        fluxdata_out = numpy.array(meep.getFluxData(flux_t))




# Plot the MEEP's internally computed R(f), T(f) (fluxdata) #{{{
#print meep.my_rank(), meep.my_rank() == 0 TODO
if meep.my_rank() == 0:
    ## Save the data
    norm = fluxdata_in_ref[1] if two_pass else fluxdata_in[1]

    refl = (fluxdata_in[1])/norm
    transm = fluxdata_out[1]/norm
    numpy.savetxt(simulation_name+".dat", (refl, transm))

    plt.figure(figsize=(9,9))
    print "Plotting %d points for reflection"
    plt.plot(fluxdata_in[0]*c, -refl, 
            color="#AA4A00", label=u'R', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(fluxdata_in[0]*c, transm, 
            color="#004AAA", label=u'T', linewidth=1, linestyle=('-','--','-.',':')[0])
    print (transm, refl)
    plt.plot(fluxdata_in[0]*c, transm-refl, 
            color="#000000", label=u'Energy', linewidth=1, linestyle=('-','--','-.',':')[0])

    plt.xlabel(u"Frequency [Hz]") 
    plt.ylabel(u"Intensity")
    plt.ylim((-0.1,1.1))
    #plt.xlim((6.2e11,8.5e11))
    plt.legend()
    plt.grid()
    plt.savefig(simulation_name, bbox_inches='tight')
#}}}
#  Plot the time-domain data#{{{
if meep.my_rank() == 0:
    plt.figure(figsize=(9,9))
    print "Plotting %d points for reflection"
    plt.plot(t, E_r0, color="#AA004A", label=u'r0', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(t, E_r, color="#AA4A00", label=u'r', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(t, [x-y for (x,y) in zip(E_r,E_r0)], color="#4AAA00", label=u'r-r0', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(t, E_t0, color="#004AAA", label=u't0', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(t, E_t, color="#4a00AA", label=u't', linewidth=1, linestyle=('-','--','-.',':')[0])

    plt.xlabel(u"Time [s]") 
    plt.ylabel(u"E field")
    plt.legend()
    plt.grid()
    plt.savefig("E_t_"+simulation_name, bbox_inches='tight')
#}}}
#  Plot the FFTs of time-domain data #{{{
if meep.my_rank() == 0:
    (E_r, E_t, E_r0, E_t0) = map(lambda x: numpy.real(numpy.array(x)), (E_r, E_t, E_r0, E_t0))

    n = len(t)
    timestep = (t[1]-t[0])
    freq = (numpy.arange(0., int(n/2)))/(n/2) / timestep / 2 # positive frequency range => divided by two 

    # calculate the raw values of electric field spectra at monitor points:
    (rfrefl, rftransm, rfrefl0, rftransm0) = map(lambda x: numpy.fft.fft(x)[0:int(n/2)], (E_r, E_t, E_r0, E_t0))
    frefl = (rfrefl-rfrefl0) / rfrefl0
    ftransm = rftransm / rfrefl0
    numpy.savetxt(simulation_name+".dat", zip(freq, frefl, ftransm))


    plt.figure(figsize=(11,11))
    xlim = (7e11, 9e11)
    plt.subplot(411)
    plt.plot(freq, abs(frefl), 
            color="#AA4A00", label=u'$|r|$', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(freq, abs(ftransm), 
            color="#004AAA", label=u'$|t|$', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(freq, abs(ftransm)**2+abs(frefl)**2, 
            color="#000000", label=u'$W$', linewidth=1, linestyle=('-','--','-.',':')[1])
    plt.plot(fluxdata_in[0]*c, transm-refl, 
            color="#666666", label=u'$W_M$', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.ylabel(u"Amplitude")
    plt.xlim(xlim)
    plt.ylim((-0.1,1.1))
    plt.legend()
    plt.grid()

    plt.subplot(412)
    plt.plot(freq, numpy.unwrap(numpy.angle(frefl)), 
            color="#AA4A00", label=u'$\\phi(r)$', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(freq, numpy.unwrap(numpy.angle(ftransm)), 
            color="#004AAA", label=u'$\\phi(t)$', linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.ylabel(u"Phase")
    plt.xlim(xlim)
    plt.ylim((-5.,5))
    plt.legend()
    plt.grid()


    d=100e-6        
    branch = 0
    k = 2*numpy.pi* freq/c                      # wave vector
    ## complex refractive index and impedance calculated according to Smith2002
    N = (numpy.arccos(((1-frefl**2+ftransm**2)/2/ftransm))+branch*numpy.pi*2) / (k*d)                       
    print N[1:5]
    Z = ((1+frefl)**2-ftransm**2) / ((1-frefl)**2 - ftransm**2)     # impedance (complex)
    plt.subplot(413)
    plt.plot(freq, numpy.real(N), 
            color="#99AA00", label=u"$N$'", linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(freq, numpy.imag(N), 
            color="#99AA66", label=u'$N$"', linewidth=1, linestyle=('-','--','-.',':')[1])
    plt.plot(freq, numpy.real(Z), 
            color="#00AADD", label=u"$Z$'", linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(freq, numpy.imag(Z), 
            color="#00AADD", label=u'$Z$"', linewidth=1, linestyle=('-','--','-.',':')[1])
    plt.ylabel(u"Value")
    plt.xlim(xlim)
    plt.ylim((-3.,3.))
    plt.legend()
    plt.grid()


    eps = N/Z
    mu = N*Z
    plt.subplot(414)
    plt.plot(freq, numpy.real(eps), 
            color="#AA00DD", label=u"$\\epsilon$'", linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(freq, numpy.imag(eps), 
            color="#AA66DD", label=u'$\\epsilon$"', linewidth=1, linestyle=('-','--','-.',':')[1])
    plt.plot(freq, numpy.real(mu), 
            color="#BB9900", label=u"$\\mu$'", linewidth=1, linestyle=('-','--','-.',':')[0])
    plt.plot(freq, numpy.imag(mu), 
            color="#BBAA44", label=u'$\\mu$"', linewidth=1, linestyle=('-','--','-.',':')[1])
    plt.xlabel(u"Frequency [THz]") 
    plt.ylabel(u"Value")
    plt.xlim(xlim)
    plt.ylim((-5.,5.))
    plt.legend()
    plt.grid()
    plt.savefig("FFT"+simulation_name, bbox_inches='tight')
#}}}

print "=== Convert field data to image ==="#{{{
for slice in slices:
    del(slice['file'])
    run_bash("cd output; rm *png")
    h5time = time.time()
    run_bash("cd output; h5topng -t 0:%d -R -Zc dkbluered -a yarg %s.h5 -S 1" % (images_number-1, slice['name']))
    pngtime = time.time()
    print "\n    Images to gif"
    run_bash("cd output; convert -compress None -delay 10 *png %s.gif" % slice['name'])
    giftime = time.time()
    #run_bash("cd output; chromium-browser *gif")
    print "Field time usage:"
    print "    h5topng: "+ `pngtime-h5time` +" s (per slice)"
    print "    convert png to gif: "+ `giftime-pngtime` +" s (per slice)"
#}}}
