"""
    This code simulates light being absorbed by the Lyman series of Hydrogen.
    It produces a graph of an initially flat spectra with absorption lines in. The shape (width and curve) of the
    absorption lines are determined by the line profiles, which are modelled using Voigt profiles.
        The Voigt profiles depend on the column density of the absorber, the temperature of the gas
        and the wavelength of the absorbing transition.
    The depth of the absorption lines is primarily determined by the column density of teh absorber.
        The column densities are randomly selected from a provided grid of column densities (such as the one
        generated by CreateGrid.py).

    I used this to demonstrate the difference in absorption spectra between Lyman-alpha forest, Lyman Limit system
    and Damped Lyman Alpha gas.

    Requirements:
        numpy
        matplotlib.pyplot
        scipy.special
        a grid of column density in cgs units (as produced by CreateGrid.py)

    Use:
        Set all variables under the "Set me" header.
        Run this file.

"""

from __future__ import division
from scipy import integrate
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
import scipy.special

# Set me ============================================================================================================= #
BIG           = 1                               # 1: Make plots big enough for spanning two columns. 0: Just span 1 column
save_folder   = "FoF/snap_012_z003p017/"        # Folder to save output (plots) to
grid_folder   = save_folder                     # Folder to retrieve grid from
grid_name     = "3NGrids3RMin128RMax2048D0.1Grid.npy"

T             = 1e4                             # Temperature of the gas            :: K (also use 13600#1e4)
I0            = 1                               # Initial intensity of the light    :: Same units as I
isolate       = 0                               # Which column density systems to isolate
                                                #   0 : Lyman alpha forest
                                                #   1 : Lyman limit systems
                                                #   2 : Damped Lyman alpha

small_fig_width_pt = 223.51561 + 18+12.5        # Get this from LaTeX using \showthe\columnwidth
big_fig_width_pt   = 526.72534 + 25 + 18*2

if BIG:
    fig_width_pt    = big_fig_width_pt
    save_name_add1  = "Big"

else:
    save_name_add1  = "Small"
    fig_width_pt    = small_fig_width_pt

inches_per_pt   = 1.0/72.27                     # Convert pt to inches
golden_mean     = 0.18                          # Aesthetic ratio (=height/width) : 772 or 5**0.5-1.0)/2.0 or custom
fig_width       = fig_width_pt*inches_per_pt    # width in inches
fig_height      = fig_width*golden_mean         # height in inches
fig_size        = [fig_width,fig_height]

params = {                                      # Pyplot plotting parameters (tweak as desired)
        'backend'               : 'ps',
        'mathtext.default'      : 'regular',
        'font.family'           : 'Times',
        'font.size'             : 11,
        'font.weight'           : 400,          # 'bold'
        'axes.labelsize'        : 11,
        'legend.fontsize'       : 8,
        'xtick.labelsize'       : 8,
        'ytick.labelsize'       : 8,
        'xtick.major.size'      : 1,            # major tick size in points
        'xtick.minor.size'      : .5,
        'ytick.major.size'      : 1,            # major tick size in points
        'ytick.minor.size'      : .5,
        'text.usetex'           : False,
        'text.latex.unicode'    : True,
        'figure.figsize'        : fig_size
        }
plt.rcParams.update(params)


# Constants ========================================================================================================== #
pi      = np.pi
e       = 4.80320425e-10        # Elementary charge         :: cgs statocoloumbs
c       = 2.99792458e10         # Speed of light            :: cm/s
k_B     = 1.3806488e-16         # Boltzmann Constant        :: cm^2 g s^-1 K^-1 (erg K^-1)
m_e     = 9.10938291e-28        # Mass of electron          :: g
m_HI    = 1.6737236e-24         # Mass of neutral Hydrogen  :: g
redshift= 3.0165046357126       # Redshift                  :: unitless
Hz      = 9.964708881888158E-18 # Hubble parameter          :: s^-1     (np.sqrt( ( (H_0/Mpc )**2 * (Omega0/(a**3) + OmegaLambda)) )

lyman       = np.array([121.6, 102.6, 97.3, 95.0, 93.8, 93.1, 92.6, 92.3, 92.1, 91.9, 91.18]) * 1e-7                                                    # Wavelengths of Lyman series   :: cm
lyman_rates = np.array([4.6986, 5.5751e-1, 1.2785e-1, 4.1250e-2, 1.6440-2, 7.5684-3, 3.8694-3, 2.1425e-3, 1.2631e-3, 7.8340e-4, 1.2258e-6]) * (10**8)   # Rates of transitions          :: s^-1
osc_str     = np.array([4.1641e-1, 7.9142e-2, 2.9006e-2, 1.3945e-2, 7.8035e-3, 4.8164e-3, 3.1850e-3, 2.2172e-3,  1.6062e-3,  1.2011e-3, 2.4474e-5])     # Oscillator strengths          ::
lyman_rates_cumulative = np.cumsum(lyman_rates)
lyman_names = ["Ly-2","Ly-3","Ly-4","Ly-5","Ly-6","Ly-7","Ly-8","Ly-9","Ly-10","Ly-11","Ly-Limit"]

# Line Profile Functions ============================================================================================= #
def voigt_faddeva(f, f0, lorentz_width, doppler_width):
    z = ( (f-f0) + 1j * lorentz_width ) / ( doppler_width * np.sqrt(2) )
    V = scipy.special.wofz(z)
    V = V.real / ( doppler_width * np.sqrt(2*pi) )
    return V

def voigt_faddeva_wavelength(x, x0, lorentz_width, doppler_width):
    z = ( (x-x0) + 1j * lorentz_width ) / ( doppler_width * np.sqrt(2) )
    V = scipy.special.wofz(z)
    V = V.real / ( doppler_width * np.sqrt(2*pi) )
    return V

def gaussian(samples, transitions, D_width, wavelength ='y'):
    if wavelength == 'n':
        return 1/np.sqrt(pi) * ( 1 / D_width ) * np.exp( - ( (samples - transitions)/D_width ) **2 )
    if wavelength == 'y':
        return 1/np.sqrt(pi) * ( 1 / D_width ) * np.exp( - ( ( (c/samples) - (c/transitions) )/D_width ) **2 )

print "Gaussian integral: ", integrate.quad(gaussian, a=-1000, b=1000, args=(0,2,'n'))[0]

# Real data ========================================================================================================== #

# Load the grid (real data)
grid = np.load(grid_folder + grid_name)

# Isolate systems
if isolate == 0:
    grid = ma.masked_where(np.log10(grid) < 17.2, grid)            # Isolate LLs
    grid = ma.masked_where(np.log10(grid) > 20.3, grid)
elif isolate == 1:
    grid = ma.masked_where(np.log10(grid) > 17.2, grid)            # Isolate Lyman - alpha forest
elif isolate == 2:
    grid = ma.masked_where(np.log10(grid) < 21, grid)              # Isolate DLAs
else:
    print "Not isolating any systems, because isolate = " + str(isolate)

# Randomly select a cell from the grid.
r = np.random.randint(0,len(grid.nonzero()[0]))
cellX=grid.nonzero()[0][r]
cellY=grid.nonzero()[1][r]

N_HI = grid[cellX, cellY]
print "Column Density:", N_HI


def intensity_spectrum(transitions, osc_str, col_density, I0=1, T=1e4, m=m_HI):
    """
    Calculates an intensity spectra.
    :param transitions  : Wavelengths of the electronic transitions at line center      :: cm
    :param osc_str      : The oscillator strengths that correpsond to the transitions   :: cgs
    :param I0           : Initial intensity                                             :: up to user
    :param col_density  : The column density                                            :: cm^-2
    :param T (optional) : Temperature of absorbing gas                                  :: K
    :param m (optional) : Mass of atoms doing the absorbing                             :: g
    :return             : Sampled wavelengths, Intensity                                :: same units as I0
    """

    # Define the range either side of the transition in which to sample (Wider lines require larger ratio)
    if   col_density > 1*(10**20.3):
        ratio = 0.2
    elif col_density > 1*(10**17.2):
        ratio = 0.05
    else:
        ratio = 0.01

    # Calculate the transition widths
    L_width = (transitions**2/c)*lyman_rates_cumulative / (2*pi)
    D_width = (1/np.sqrt(2))*(transitions / c) * np.sqrt(2 * k_B * T / m)

    # Calculate the absorption cross-section.
    total_samples = []
    total_absorption_cross_section = []
    for i in range(len(transitions)):
        samples = np.linspace( (1-ratio)*transitions[i], (1+ratio)*transitions[i], 10000)
        line_profile = voigt_faddeva_wavelength(samples, transitions[i], lorentz_width=L_width[i], doppler_width=D_width[i])
        print "Transition", i, "line profile at center", max(line_profile)
        absorption_cross_section = (transitions[i]**2/c)*(pi*e**2/(m_e*c))*osc_str[i] * np.sqrt(pi) * line_profile *2

        total_absorption_cross_section.extend(absorption_cross_section)
        total_samples.extend(samples)

    # Convert to arrays
    total_absorption_cross_section = np.array(total_absorption_cross_section)
    total_samples = np.array(total_samples)

    # Find intensity
    optical_depth = total_absorption_cross_section * col_density
    I = I0 * np.exp(-optical_depth)
    return total_samples, I


# Plotting =========================================================================================================== #

# Find wavelength of transitions at corresponding redshift.
transitions  = np.array(lyman)

# Calculate the intensity of the resulting spectrum
total_samples, I = intensity_spectrum(transitions, osc_str, N_HI, I0, T)

# Convert to angstrom for plotting
transitions   *= 1e8
total_samples *= 1e8

total_samples *= (1+redshift)
transitions   *= (1+redshift)

# Plot intensity vs. wavelength
plt.figure()
plt.plot(total_samples, I)
plt.axhline(I0, color='g', alpha = 0.2)
for i in range(len(transitions)):
    plt.axvline(transitions[i],ymin=0, ymax=1, color='r', alpha=0.2)
# Plot labels of line centers
for i in range(len(transitions)-8):
    plt.text(x=transitions[i],y=plt.ylim()[1],s=lyman_names[i], verticalalignment='bottom', horizontalalignment='center')
plt.text(x=transitions[len(transitions)-1],y=plt.ylim()[1],s=lyman_names[-1], verticalalignment='bottom', horizontalalignment='center')
plt.figtext(x=0.72, y=-0.12,s=(r"$ \rm N_{HI}=$"+'{0:.3g}'.format(N_HI)+r"$ \rm cm^{-2}$"))
plt.xlabel(r"$ \rm Wavelength (\AA)$", labelpad=3)
plt.ylabel(r"$ \rm I (erg s^{-1} cm^{-2} \AA^{-1})$", labelpad=2)
plt.xlim((3500,5500))
plt.savefig("NewSpectra/Spectra"+str(N_HI)+".pdf", dpi=1000, bbox_inches='tight', pad_inches=0.04)
plt.savefig("NewSpectra/Spectra"+str(N_HI)+".png", dpi=300, bbox_inches='tight', pad_inches=0.04)
plt.savefig("NewSpectra/Spectra"+str(N_HI)+".jpg", dpi=300, bbox_inches='tight', pad_inches=0.04)
plt.show()

