#!/usr/bin/env python
# -*- coding: utf-8 -*-

# From the Meep tutorial: plotting permittivity and fields of a bent waveguide
from __future__ import division

import meep as mp
#import matplotlib.pyplot as plt
import os
import sys
import numpy as np
from matplotlib import cm
import pygame

#mp.quiet(True)

Nx = int(sys.argv[1])
Ny = int(sys.argv[2])
try:
    fps = float(sys.argv[3])
except (IndexError, ValueError):
    fps = 40



cell = mp.Vector3(8,4,0)
geometry = [mp.Block(mp.Vector3(12,1,mp.inf),
                     center=mp.Vector3(0,0),
                     material=mp.Medium(epsilon=24))
            ]#,
            #mp.Block(mp.Vector3(1,5,mp.inf),
            #         center=mp.Vector3(1,1),
            #         material=mp.Medium(epsilon=12))]
#pml_layers = [mp.PML(2.0)]
resolution = 10

sources = [mp.Source(mp.ContinuousSource(frequency=0.15, width=20), # wavelength=2*(11**0.5)
                     component=mp.Ez,
                     center=mp.Vector3(0,0),
                     size=mp.Vector3(0,1))
]

sim = mp.Simulation(cell_size=cell,
                    geometry=geometry,
                    sources=sources,
                    resolution=resolution)


buffer = np.empty([80,40])
clock = pygame.time.Clock()

def get_slice(sim):
    val = sim.get_array(center=mp.Vector3(), size=cell, component=mp.Ez, arr=buffer)
    #plt.figure()
    #plt.imshow(val, interpolation='none', cmap='RdBu')
    #plt.axis('off')
    #plt.show()
    out = (val - np.amin(val)) / (np.amax(val) - np.amin(val))
    out = np.transpose(out) # or switch x and y
    out = out.ravel()
    out = 250*cm.RdBu(out)
    out = out[:,0:3]
    out = out.ravel()
    out = out.astype(np.uint8).tobytes()
    os.write(1, out)
    clock.tick_busy_loop(fps)

def plot_dielectric(sim):
    val = sim.get_array(center=mp.Vector3(), size=cell, component=mp.Dielectric, arr=buffer)
    #plt.figure()
    #plt.imshow(val, interpolation='none', cmap='RdBu')
    #plt.axis('off')
    #plt.show()
    #print()


#sim.run(mp.at_beginning(plot_dielectric), mp.at_every(0.6, get_slice), until=200)
sim.run(mp.at_every(0.00001, get_slice), until=1000)