%%time
# Solve the systems for each polarization
Ainv = simpeg.SolverLU(A)

%%time
e_x = Ainv*rhs_x
e_y = Ainv*rhs_y

"""
Explanation: We are using splu in scipy package. This is bit slow, but on the cluster you can use mumps, which might a lot faster. We can think about having better iterative solver.
End of explanation
"""


Meinv = M.getEdgeInnerProduct(np.ones_like(sig), invMat=True)

j_x = Meinv*Msig*e_x
j_y = Meinv*Msig*e_x

e_x_CC = M.aveE2CCV*e_x
e_y_CC = M.aveE2CCV*e_y
j_x_CC = M.aveE2CCV*j_x
j_y_CC = M.aveE2CCV*j_y
# j_x_CC = Utils.sdiag(np.r_[sig, sig, sig])*e_x_CC

"""
Explanation: I want to visualize electrical field, which is a vector, so I average them on to cell center. Also I want to see current density ($\vec{j} = \sigma \vec{e}$).
End of explanation
"""


fig, ax = plt.subplots(1,2, figsize = (12, 5))
dat0 = M.plotSlice(abs(e_x_CC), vType='CCv', view='vec', streamOpts={'color': 'k'}, normal='X', ax = ax[0])
cb0 = plt.colorbar(dat0[0], ax = ax[0])
dat1 = M.plotSlice(abs(j_x_CC), vType='CCv', view='vec', streamOpts={'color': 'k'}, normal='X', ax = ax[1])
cb1 = plt.colorbar(dat1[0], ax = ax[1])

"""
Explanation: Then use "plotSlice" function, to visualize 2D sections
End of explanation
"""


# Calculate the data
rx_x, rx_y = np.meshgrid(np.arange(-500,501,50),np.arange(-500,501,50))
rx_loc = np.hstack((simpeg.Utils.mkvc(rx_x,2),simpeg.Utils.mkvc(rx_y,2),np.zeros((np.prod(rx_x.shape),1))))
# Get the projection matrices
Qex = M.getInterpolationMat(rx_loc,'Ex')
Qey = M.getInterpolationMat(rx_loc,'Ey')
Qez = M.getInterpolationMat(rx_loc,'Ez')
Qfx = M.getInterpolationMat(rx_loc,'Fx')
Qfy = M.getInterpolationMat(rx_loc,'Fy')
Qfz = M.getInterpolationMat(rx_loc,'Fz')

e_x_loc = np.hstack([simpeg.Utils.mkvc(Qex*e_x,2),simpeg.Utils.mkvc(Qey*e_x,2),simpeg.Utils.mkvc(Qez*e_x,2)])
e_y_loc = np.hstack([simpeg.Utils.mkvc(Qex*e_y,2),simpeg.Utils.mkvc(Qey*e_y,2),simpeg.Utils.mkvc(Qez*e_y,2)])
h_x_loc = np.hstack([simpeg.Utils.mkvc(Qfx*C*e_x,2),simpeg.Utils.mkvc(Qfy*C*e_x,2),simpeg.Utils.mkvc(Qfz*C*e_x,2)])
h_y_loc = np.hstack([simpeg.Utils.mkvc(Qfx*C*e_y,2),simpeg.Utils.mkvc(Qfy*C*e_y,2),simpeg.Utils.mkvc(Qfz*C*e_y,2)])

# Make a combined matrix
dt = np.dtype([('ex1',complex),('ey1',complex),('ez1',complex),('hx1',complex),('hy1',complex),('hz1',complex),('ex2',complex),('ey2',complex),('ez2',complex),('hx2',complex),('hy2',complex),('hz2',complex)])
combMat = np.empty((len(e_x_loc)),dtype=dt)
combMat['ex1'] = e_x_loc[:,0]
combMat['ey1'] = e_x_loc[:,1]
combMat['ez1'] = e_x_loc[:,2]
combMat['ex2'] = e_y_loc[:,0]
combMat['ey2'] = e_y_loc[:,1]
combMat['ez2'] = e_y_loc[:,2]
combMat['hx1'] = h_x_loc[:,0]
combMat['hy1'] = h_x_loc[:,1]
combMat['hz1'] = h_x_loc[:,2]
combMat['hx2'] = h_y_loc[:,0]
combMat['hy2'] = h_y_loc[:,1]
combMat['hz2'] = h_y_loc[:,2]


def calculateImpedance(fieldsData):
    ''' 
    Function that calculates MT impedance data from a rec array with E and H field data from both polarizations
    '''
    zxx = (fieldsData['ex1']*fieldsData['hy2'] - fieldsData['ex2']*fieldsData['hy1'])/(fieldsData['hx1']*fieldsData['hy2'] - fieldsData['hx2']*fieldsData['hy1'])
    zxy = (-fieldsData['ex1']*fieldsData['hx2'] + fieldsData['ex2']*fieldsData['hx1'])/(fieldsData['hx1']*fieldsData['hy2'] - fieldsData['hx2']*fieldsData['hy1'])
    zyx = (fieldsData['ey1']*fieldsData['hy2'] - fieldsData['ey2']*fieldsData['hy1'])/(fieldsData['hx1']*fieldsData['hy2'] - fieldsData['hx2']*fieldsData['hy1'])
    zyy = (-fieldsData['ey1']*fieldsData['hx2'] + fieldsData['ey2']*fieldsData['hx1'])/(fieldsData['hx1']*fieldsData['hy2'] - fieldsData['hx2']*fieldsData['hy1'])
    return zxx, zxy, zyx, zyy

zxx, zxy, zyx, zyy = calculateImpedance(combMat)

zxy

zyx


zxy + zyx


"""
Explanation: Is it reasonable?: Based on that you put resistive target that makes sense to me; current does not want to flow on resistive target so they just do roundabout:). And see air interface. It is continuous on current but not on electric field, which looks reasonable.
End of explanation
"""


