#from IPython.display import Latex

"""
Explanation: Simple notebook on performing a 1D MT problem.
End of explanation
"""


import sys
sys.path.append('C:/GudniWork/Codes/python/simpeg')
import SimPEG as simpeg, numpy as np, scipy, scipy.sparse as sp

"""
Explanation: Maxwell's equations in 1D are as follows
$i \omega b = - \partial_z e $ 
$  s =  \partial_z(\mu^{-1} b) - \sigma(z) e$
$b(0) = 1 \hspace{1cm} ;\hspace{1cm} b(-\infty) = 0$
where $e = \widehat{\overrightarrow{E}}_x $ and $b = \widehat{\overrightarrow{B}}_y$
In weak form the equations become
$i \omega(b,f) = - (\partial_ze,f)$ 
$(s,w) = -(\mu^{-1} b, \partial_z w) - (\sigma(z) e, w)$
where f and w are abritrary functions living in same discritizational space as b and e, respectivily.
We consider e on nodes and b on cell centers. This way the derivative of any nodal function becomes
$ (\partial_z u)_k \approx $
$h_{k}^{-1} ( u_{k+\frac{1}{2}} - u_{k-\frac{1}{2}})  + O(h^2) $
Matrix form
$ e_z \approx \textbf{L}^{-1} \textbf{G} e = \begin{bmatrix} h_1^{-1} & & & \\ & h_2^{-1} & & \\ & &  \ddots & \\ & & & h_n^{-1} \end{bmatrix}^{(n,n)}
\begin{bmatrix} -1 & 1 & & & & \\ & -1 & 1 & & & \\ & & \ddots & \ddots & \\ & & & -1 & 1 \end{bmatrix}^{(n,n+1)} 
\begin{bmatrix} e_1 \\ \\ \vdots \\ \\ e_{n+1} \end{bmatrix}^{(n+1,1)} $
where $ \textbf{L} = diag(h) $ is the cell size and $ \textbf{G}$ is the gradient operator with -1,1 representing the topology of the mesh, taking the difference between adjoint cells.
We need to compute 2 inner products, on cell centers and from nodes to cell centers.
Cell centers inner product is
$ (b,f) \approx \sum\limits_k h_k \textbf{b}_k \textbf{f}_k + O(h^2) $
and in matrix from
$ (b,) \approx \textbf{b}^T \textbf{M}^f \textbf{f}$  and $ (\mu^{-1} b,f) \approx \textbf{b}^T \textbf{M}_{\mu}^f \textbf{f}$ 
where $ \textbf{M}_{\mu}^f = diag(\textbf{h} \odot \mu^{-1}) $ and $ \textbf{M}^f = diag(\textbf{h}) $ are the matrices.
Nodes to cell centers inner product is
$ (\sigma e, w) \approx \sum\limits_k \frac{h_k \sigma_k}{4} ( e_{k+\frac{1}{2}} w_{k+\frac{1}{2}} + e_{k+\frac{1}{2}} w_{k+\frac{1}{2}} )$
and in matrix from
$ (\sigma e, w ) \approx (\textbf{h} \odot \sigma )^T ( \textbf{A}_v (\textbf{e} \odot \textbf{w} = \textbf{w}^T diag(\textbf{A}_v^T (\textbf{h} \odot \sigma)) \textbf{e} $
Here $\odot$ is a point wise Hadamard product and $\textbf{A}_v$ is the averaging operator/matrix from nodes to cell centers
$ \textbf{A}_v = \begin{bmatrix} \frac{1}{2} & \frac{1}{2} &  & \\ & \ddots & \ddots & \\ &  &  \frac{1}{2} & \frac{1}{2} \end{bmatrix} ^{(n+1 , n)} $
The sigma mass matrix is defined as 
$ \textbf{M}_{\sigma}^{e} = diag(\textbf{A}_v^T (\textbf{h} \odot \sigma)) $
In the MT problem there is no source in the domain, so $ s = 0 $. How ever the boundary conditions provide the right hand side where 
$ (\partial_z \mu^{-1} b, w ) = - (\mu^{-1} b, \partial_z w ) + (\mu^{-1} b w )|_0^{end} $
where 
$ (\mu^{-1} b w )|_0^{end} = \textbf{bc}^T (\textbf{BC w}) $
here $\textbf{BC}$ is an matrix operator that extracts the boundary elements from $ \textbf{w}$ and $\textbf{bc} $ are the known boundary condintions. For the 1D case with homogenous boundary conditions we have 
$ \textbf{B} = \begin{bmatrix} -1 & 0 \\ \vdots & \vdots \\ 0 & 1 \end{bmatrix} $ 
The weak form is 
$ (\mu^{-1} b, \partial_z w) + (\sigma e, w) = (\mu^{-1} b w )|_0^{end} $
$ (i \omega b,f) + (\partial_z e , f) = 0 $
Using the above matrix represntation we get the Maxwells equations in following form
$ i \omega \textbf{f}^T \textbf{M}^f \textbf{b} + \textbf{f}^T \textbf{M}^f \textbf{L}^{-1} \textbf{G} \textbf{e} = 0 $ 
$ \textbf{w}^T \textbf{G}^T \textbf{L}^{-1}  \textbf{M}^f_{\mu} \textbf{b} + \textbf{w}^T \textbf{M}_{\sigma}^e \textbf{e} = \textbf{w}^T \textbf{bc}^T \textbf{BC} $ 
Here we use that 
$ (\textbf{b},\textbf{f}) \approx \textbf{b}^T \textbf{M}^f \textbf{f} = \textbf{f}^T  \textbf{M}^f \textbf{b} $ 
since $\textbf{M}^f$ is a symmetric diaoganal matrix of size n by n and $\textbf{b}$ and $\textbf{f}$ are vectors of length n.
We eliminate testing vectors and get system of equations to solve
$ i \omega \textbf{b} + \textbf{L}^{-1} \textbf{G} \textbf{e} = 0 $ 
$ \textbf{G}^T \textbf{L}^{-1}  \textbf{M}^f_{\mu} \textbf{b} +  \textbf{M}_{\sigma}^e \textbf{e} = \textbf{bc}^T \textbf{BC} $ 
and as $ \textbf{A} \textbf{x} = \textbf{bc} $ system that we will solve, where 
$ \textbf{A} = \begin{bmatrix} \textbf{G}^T \textbf{L}^{-1}  \textbf{M}^f_{\mu}  & \textbf{M}_{\sigma}^e  \\ i \omega  &  \textbf{L}^{-1} \textbf{G}  \end{bmatrix} $
$ \textbf{x} = \begin{bmatrix} \textbf{b} \\ \textbf{e} \end{bmatrix} $
$ \textbf{bc} = \begin{bmatrix} \textbf{bc}^T \textbf{BC} \\ \textbf{0} \end{bmatrix} $
End of explanation
"""


 # Set up  the problem
mu = 4*np.pi*1e-7
eps0 = 8.85e-12
# Frequency
fr = np.array([1e1]) #np.logspace(0,5,200) #np.array([2000]) #np.logspace(-4,5,82)
omega = 2*np.pi*fr
# Mesh
sig0 = 1e-2
#L = 3*np.sqrt(2/(mu*omega[0]*sig0))
#nn=np.ceil(np.log(0.3*L + 1)/np.log(1.3))
#h = 5*(1.3**(np.arange(nn+1)))

h = np.ones(18)
x0 = np.array([0])
#sig = sig0*np.ones((len(h),1)) 
#sig[0:50] = 0.1
#sig[50:100] = 1
# Make the mesh
mesh = simpeg.Mesh.TensorMesh([h],x0)
sig = np.zeros(mesh.nC) + 1e-8
sig[mesh.vectorCCx<=0] = 1e-2

fr,omega

# Make the operators
G = mesh.nodalGrad
Av = mesh.aveN2CC
Li = scipy.sparse.spdiags(1/mesh.hx,0,mesh.nNx,mesh.nNx)
Mmu = scipy.sparse.spdiags(mesh.hx/mu,0,mesh.nCx,mesh.nCx)
Msig = scipy.sparse.spdiags(Av.T.dot(mesh.hx*sig.ravel()),0,mesh.nNx,mesh.nNx)
# The boundaries
bc_b = np.zeros((mesh.nCx,1))
bc_b[0] = -1 # Set the top b field to 1
bc_e = np.zeros((mesh.nNx,1))
# Make the sparse matrix
bc = sp.vstack((bc_b,bc_e))


b = np.empty((mesh.nCx,len(omega)),dtype=np.complex64)
e = np.empty((mesh.nNx,len(omega)),dtype=np.complex64)
# Loop all the frequencies
for nrOm, om in enumerate(omega):
    # Left hand side
    A = sp.vstack((sp.hstack(( -G.conj().T.dot(Mmu), - Msig)), sp.hstack((1j*om*scipy.sparse.identity(mesh.nCx) , G))))
    #A = A.tocsr
    # Solve the system
    bef = scipy.sparse.linalg.spsolve(A,bc)
    # Sort the output
    b[:,nrOm] = bef[0:mesh.nCx]
    e[:,nrOm] = bef[mesh.nCx::]


import matplotlib.pyplot as plt
# Plot the solution
z=e[0,:]/(b[0,:]/mu)
app_res = ((1./(8e-7*np.pi**2))/fr)*np.abs(z)**2
app_phs = np.arctan(z.imag/z.real)*(180/np.pi)
ax_res = plt.subplot(2,1,1)
ax_res.loglog(fr,app_res)
ax_phs = plt.subplot(2,1,2)
ax_phs.semilogx(fr,app_phs)

plt.show()

# Calculate the impedance
z = e[0,:]/(b[0,:]/mu)
z

app_res

"""
Explanation: We have
$ i \omega b = - \partial_z e \hspace{1cm} ; \hspace{1cm} \partial_z(\mu^{-1} b) - \sigma(z) e \hspace{1cm} ;\hspace{1cm} b(0) = 1 \hspace{1cm} ;\hspace{1cm} b(-\infty) = 0 $
To deal with boundary: we assume that below depth L both $ \sigma $ and $ \mu $ are constants ($ z < - L $). At the boundary we have that 
$ e = c \exp(ikz) \hspace{0.2cm} where \hspace{0.2cm} k = \sqrt{i\omega\mu\sigma} $. 
Therefore for $ z < - L $ we have that
$\omega b - k e = 0 $.
We discretize the e field on the nodes and b field at the cell centers. The system we want to solve is 
$\begin{bmatrix} i  \omega & \frac{\partial}{\partial z} \\ \frac{1}{\mu} \frac{\partial}{\partial z} & -\sigma \end{bmatrix} 
\begin{bmatrix} b \\ e \end{bmatrix}  = \begin{bmatrix} s1 \\ s2 \end{bmatrix}$
End of explanation
"""


