import os
import sys
sys.path.insert(0, os.path.abspath('../../'))

import numpy as np
from matplotlib import pyplot as plt
import arrayfire as af

from dg_maxwell import params
from dg_maxwell import lagrange
from dg_maxwell import wave_equation as w1d
from dg_maxwell import utils

af.set_backend('opencl')
af.set_device(1)
af.info()

plt.rcParams['figure.figsize']     = 12, 7.5
plt.rcParams['lines.linewidth']    = 1.5
plt.rcParams['font.family']        = 'serif'
plt.rcParams['font.weight']        = 'bold'
plt.rcParams['font.size']          = 20  
plt.rcParams['font.sans-serif']    = 'serif'
plt.rcParams['text.usetex']        = True
plt.rcParams['axes.linewidth']     = 1.5
plt.rcParams['axes.titlesize']     = 'medium'
plt.rcParams['axes.labelsize']     = 'medium'

plt.rcParams['xtick.major.size']   = 8
plt.rcParams['xtick.minor.size']   = 4
plt.rcParams['xtick.major.pad']    = 8
plt.rcParams['xtick.minor.pad']    = 8
plt.rcParams['xtick.color']        = 'k'
plt.rcParams['xtick.labelsize']    = 'medium'
plt.rcParams['xtick.direction']    = 'in'    

plt.rcParams['ytick.major.size']   = 8
plt.rcParams['ytick.minor.size']   = 4
plt.rcParams['ytick.major.pad']    = 8
plt.rcParams['ytick.minor.pad']    = 8
plt.rcParams['ytick.color']        = 'k'
plt.rcParams['ytick.labelsize']    = 'medium'
plt.rcParams['ytick.direction']    = 'in'
plt.rcParams['text.usetex']        = True
plt.rcParams['text.latex.unicode'] = True

# 1. Set the initial conditions

E_00 = 1.
E_01 = 1.

B_00 = 0.2
B_01 = 0.5

E_z_init = E_00 * af.sin(2 * np.pi * params.element_LGL) \
         + E_01 * af.cos(2 * np.pi * params.element_LGL)

B_y_init = B_00 * af.sin(2 * np.pi * params.element_LGL) \
         + B_01 * af.cos(2 * np.pi * params.element_LGL)

u_init = af.constant(0., d0 = params.N_LGL, d1 = params.N_Elements, d2 = 2)
u_init[:, :, 0] = E_z_init
u_init[:, :, 1] = B_y_init

element_LGL_flat = af.flat(params.element_LGL)
E_z_init_flat    = af.flat(u_init[:, :, 0])
B_y_init_flat    = af.flat(u_init[:, :, 1])

plt.plot(element_LGL_flat, E_z_init_flat, label = r'$E_z$')
plt.plot(element_LGL_flat, B_y_init_flat, label = r'$B_y$')

plt.title(r'Plot of $E_z(t = 0)$ and $B_y(t = 0)$')
plt.xlabel(r'$x$')
plt.ylabel(r'$y$')

plt.legend(prop={'size': 14})

plt.show()

"""
Explanation: Developmental file for modifying the 1D advection solver to work for multiple wave equations
End of explanation
"""


# Older LF flux code
u_n = u_init[:, :, :]

u_iplus1_0    = af.shift(u_n[0, :], 0, -1)
u_i_N_LGL     = u_n[-1, :]
flux_iplus1_0 = w1d.flux_x(u_iplus1_0)
flux_i_N_LGL  = w1d.flux_x(u_i_N_LGL)

boundary_flux = (flux_iplus1_0 + flux_i_N_LGL) / 2 \
              - params.c_lax * (u_iplus1_0 - u_i_N_LGL) / 2


print(boundary_flux)


"""
Explanation: Prototype implementation of LF flux for multiple-u's
End of explanation
"""


