kernel =  lambda m: V.dot(m)
derChk = lambda m: (kernel(m), lambda v: kernel(v) )
Tests.checkDerivative(derChk, mest, dx = mest*0.1)

"""
Explanation: \begin{eqnarray}
    m = e^{m_1} + (e^{m_2}-e^{m_1})(\frac{tan^{-1}(\alpha \ f(\vec{r},c))}{\pi}+0.5) 
\end{eqnarray}
\begin{align}
        \frac{\partial m}{\partial m_1} = e^{m_1}-e^{m_1}tan^{-1}(\frac{\alpha \ f}{\pi}+0.5)
\end{align}
\begin{align}
        \frac{\partial m}{\partial m_2} = e^{m_2}tan^{-1}(\frac{\alpha \ f}{\pi}+0.5)
\end{align}
\begin{align}
        \frac{\partial m}{\partial c} = (e^{m_2}-e^{m_1})\frac{\partial f}{\partial c}\frac{\alpha}{1+(\alpha f)^2}\frac{1}{\pi}
\end{align}
Step1: Check derivative for "c"
End of explanation
"""


def Poly2Dmap(m1, m2, c, xy):
    alpha = 1e4
    f = polynomial.polyval(xy[:,0], c) - xy[:,1]
    m = m1+(m2-m1)*(np.arctan(alpha*f)/np.pi+0.5)
    return m

"""
Explanation: Step2: Check $\frac{\partial m}{\partial c}$
End of explanation
"""


def Poly2DmapDeriv(m1, m2, c, xy):
    alpha = 1e4
    V = polynomial.polyvander(xy[:,0], len(c)-1)
    f = polynomial.polyval(xy[:,0], c) - xy[:,1]    
    deriv = Utils.sdiag(alpha*(m2-m1)/(1.+(alpha*f)**2)/np.pi)*V
    return deriv

kernel = lambda x: Poly2Dmap(1., 2., x, mesh.gridCC)
kernel_deriv = lambda x: Poly2DmapDeriv(1., 2., x, mesh.gridCC)

derChk = lambda m: (kernel(m), lambda v: np.dot(kernel_deriv(m), v) )
Tests.checkDerivative(derChk, mest, dx = mest*0.1)

from SimPEG.Maps import IdentityMap, PolyMap, ActiveCells

actind = (mesh.gridCC[:,0]>-15.) & (mesh.gridCC[:,0] <15.)

meshact = Mesh.TensorMesh([mesh.hx[actx],mesh.hy], x0 = np.r_[ -15., mesh.x0[1]])

sigma0 = np.ones(mesh.nC)*1e0
sigma0[mesh.gridCC[:,1]>-5] = 1e-2
mesh.plotImage(sigma0)

actmap = ActiveCells(mesh, actind, sigma0)

weight = (1./(V**2).sum(axis=0))**0.5
weight = weight / weight.max()

m1D = Mesh.TensorMesh([V.shape[1]+2])
weightmap = Maps.Weighting(m1D, weights=np.r_[1., 1., weight])

polymap = PolyMap(mesh, V.shape[1]-1, logSigma=True)

mapping = polymap*weightmap
# mapping = polymap

polymap.test()

mtrue = np.r_[np.log(1e-2), np.log(1e0), mest / weight]
m0 = np.r_[np.log(1e-2), np.log(1e-2), np.r_[-3., np.zeros(V.shape[1]-1)] / weight]

# mtrue = np.r_[np.log(1e-2), np.log(1e0), mest]
# m0 = np.r_[np.log(1e-2), np.log(1e0), np.r_[-3., np.zeros(V.shape[1]-1)]]

sigma = mapping*mtrue

# actind = (mesh.gridCC[:,0] > -15) & (mesh.gridCC[:,0] < 15)
# meshact = Mesh.TensorMesh([mesh.hx, mesh.hy], x0=mesh.x0)
# actmap = Maps.ActiveCells(mesh, actind, 2e-3)

xr = np.linspace(-15, 15, 20)
xz_A = Utils.ndgrid(xr, np.r_[-0.25])
xz_B = Utils.ndgrid(np.ones_like(xr)*22, np.r_[-0.25])
xz_M = Utils.ndgrid(xr, np.r_[-0.25])
xz_N = Utils.ndgrid(np.ones_like(xr)*-22, np.r_[-0.25])

ntx = xz_A.shape[0]
txList = []
for i in range(ntx):
    offset = abs(xz_A[i,0]-xz_M[:,0])
    actrx = offset > 5.
    rx = DC.RxDipole(xz_M[actrx,:], xz_N[actrx,:])
    src = DC.SrcDipole([rx], xz_A[i,:], xz_B[i,:])
    txList.append(src)

dat = mesh.plotImage(np.log10(sigma), clim=(-2, 0))
plt.colorbar(dat[0])
plot(xz_A[:,0], xz_A[:,1], 'w.')
plot(xz_B[:,0], xz_B[:,1], 'y.')
plot(xz_N[:,0], xz_N[:,1], 'r.')

dat = mesh.plotImage(np.log10(mapping*m0), clim=(-2, 0))
plt.colorbar(dat[0])
plot(xz_A[:,0], xz_A[:,1], 'w.')
plot(xz_B[:,0], xz_B[:,1], 'y.')
plot(xz_N[:,0], xz_N[:,1], 'r.')

survey = DC.SurveyDC(txList)
problem = DC.ProblemDC_CC(mesh, mapping = mapping)
problem.pair(survey)

from pymatsolver import MumpsSolver
problem.Solver = MumpsSolver

dini = survey.dpred(m0)
dtrue  = survey.dpred(mtrue)
plot(dini)
plot(dtrue)

mtrue.shape


figsize(14*0.5,7*0.5)
circmodelest = mapping*mtrue
dat = mesh.plotImage(np.log10(circmodelest))
plot(xz_A[:,0], xz_A[:,1], 'w.')
plot(xz_B[:,0], xz_B[:,1], 'k.')
plot(xz_N[:,0], xz_N[:,1], 'r.')
# plot(temp.rxList[0].locs[0][:,0], temp.rxList[0].locs[0][:,1], 'bo')
plt.colorbar(dat[0])

noise = 0.1*abs(dtrue)*np.random.randn(dtrue.shape[0])

survey.dobs = dtrue +noise
dmis = DataMisfit.l2_DataMisfit(survey)
dmis.Wd = 1./(0.1*abs(dtrue)+abs(dini).min()*0.1)
reg = Regularization.BaseRegularization(m1D)
opt = Optimization.InexactGaussNewton(maxIter=50,tolX=1e-3, maxIterLS=20)
opt.remember('xc')
invProb = InvProblem.BaseInvProblem(dmis, reg, opt)
invProb.beta = 0.
betaSched = Directives.BetaSchedule(coolingFactor=1, coolingRate=1)
targetmis = Directives.TargetMisfit()
savemodel = Directives.SaveModelEveryIteration()
inv = Inversion.BaseInversion(invProb, directiveList=[betaSched,targetmis])
reg.mref = m0
mopt = inv.run(m0)

# invProb.beta = 0.
# mopt = inv.run(mopt)

XC = opt.recall('xc')

from ipywidgets import interact, IntSlider

def viewInv(iteration):
    fig  = plt.figure(num=0,figsize = (10,5))
    ax = plt.subplot(111)
    dat = mesh.plotImage(np.log10(mapping*XC[iteration]), grid=False, ax=ax, clim=(-2, 0))
#     ax.set_xlim(mesh.vectorNx.min(), mesh.vectorNx.max())
#     ax.set_ylim(mesh.vectorNy.min(), mesh.vectorNy.max())
    ax.plot(mesh.vectorCCx, dpred, 'w.-')
#     plt.colorbar(dat[0], ax=ax)
#     ax.set_ylim (-15, 0.)
#     ax.set_xlim (-15, 15.)
    plt.show()
    return True

interact(viewInv, iteration = IntSlider(min=0, max=opt.iter-1,step=1, value=0))


figsize(14*0.5,7*0.5)
circmodelest = mapping*mopt
dat = mesh.plotImage(np.log10(circmodelest))
plot(xz_A[:,0], xz_A[:,1], 'w.')
plot(xz_B[:,0], xz_B[:,1], 'k.')
plot(xz_N[:,0], xz_N[:,1], 'r.')
plot(mesh.vectorCCx, dpred, 'w.-')
# plot(temp.rxList[0].locs[0][:,0], temp.rxList[0].locs[0][:,1], 'bo')
plt.colorbar(dat[0])

"""
Explanation: \begin{align}
        \frac{\partial m}{\partial c} = (e^{m_2}-e^{m_1})\frac{\partial f}{\partial c}\frac{\alpha}{1+(\alpha f)^2}\frac{1}{\pi}
\end{align}
End of explanation
"""


