#-*- coding: utf-8 -*-

import numpy as np
from matplotlib import pyplot as plt

wp = 1      # plasma frequency
c0 = 1      # speed of light
g = 0.2     # damping
ediel = 1   # permittivity of dielectric
w = np.linspace(0, 1.5, 100) # frequency

drude_real = 1 - wp**2 / (w**2 + g**2)          # real part of the Drude formula
drude_imag = (wp**2 * g)/(w * (w**2 + g**2))    # imaginary part

# real, imaginary part and absolute value of k_x
kr = w/c0 * np.sqrt((ediel * drude_real)/(ediel + drude_real))
ki = (((drude_real * ediel)/(drude_real + ediel))**(3.0/2.0) * (drude_imag/(2*(drude_real**2))))
k = np.sqrt(kr**2 + ki**2)

# Plot statements
fig = plt.figure(figsize=(12, 6))
ax1 = fig.add_subplot(111)
ax1.set_xlabel(r"wave vector $k_x c / omega_p$")
ax1.set_ylabel(r"frequency $omega / omega_p$")
ax1.plot(k, w, c="g", label="|k_x|")
ax1.plot(kr, w, c="r", label="Re k_x")
ax1.plot(ki, w, c="r", label="Im k_x", linestyle="--")
ax1.plot(w, w, c="b", label="light line", linestyle="--")
leg = ax1.legend(loc="best")
ax1.grid()
plt.ylim((0, 1.5))
plt.xlim((0, 3))
plt.axhspan(1.0/np.sqrt(2), 1.0, alpha=0.3, color='green')      # leaky mode region
plt.show()