from SimPEG.mesh import TensorMesh
from SimPEG.forward import Problem, SyntheticProblem
from SimPEG.tests import checkDerivative
from SimPEG.utils import ModelBuilder, sdiag
import numpy as np
import scipy as sp
import scipy.sparse.linalg as linalg

class DCProblem(Problem):
    """
        **DCProblem**

        Geophysical DC resistivity problem.

    """
    def __init__(self, mesh):
        super(DCProblem, self).__init__(mesh)
        self.mesh.setCellGradBC('neumann')

    def createMatrix(self, m):
        """
            Makes the matrix A(m) for the DC resistivity problem.

            :param numpy.array m: model
            :rtype: scipy.csc_matrix
            :return: A(m)

            .. math::
                c(m,u) = A(m)u - q = G\\text{sdiag}(M(mT(m)))Du - q = 0

            Where M() is the mass matrix and mT is the model transform.
        """
        D = self.mesh.faceDiv
        G = self.mesh.cellGrad
        sigma = self.modelTransform(m)
        Msig = self.mesh.getFaceMass(sigma)
        A = D*Msig*G
        return A.tocsc()

    def field(self, m):
        A = self.createMatrix(m)
        solve = linalg.factorized(A)

        nRHSs = self.RHS.shape[1]  # Number of RHSs
        phi = np.zeros((self.mesh.nC, nRHSs)) + np.nan
        for ii in range(nRHSs):
            phi[:,ii] = solve(self.RHS[:,ii])

        return phi

    def J(self, m, v, u=None, solve=None):
        """
            :param numpy.array m: model
            :param numpy.array v: vector to multiply
            :param numpy.array u: fields
            :rtype: numpy.array
            :return: Jv

            .. math::
                c(m,u) = A(m)u - q = G\\text{sdiag}(M(mT(m)))Du - q = 0

                \\nabla_u (A(m)u - q) = A(m)

                \\nabla_m (A(m)u - q) = G\\text{sdiag}(Du)\\nabla_m(M(mT(m)))

            Where M() is the mass matrix and mT is the model transform.

            .. math::
                J = - P \left( \\nabla_u c(m, u) \\right)^{-1} \\nabla_m c(m, u)

                J(v) = - P ( A(m)^{-1} ( G\\text{sdiag}(Du)\\nabla_m(M(mT(m))) v ) )
        """
        P = self.P
        D = self.mesh.faceDiv
        G = self.mesh.cellGrad
        A = self.createMatrix(m)
        Av_dm = self.mesh.getFaceMassDeriv()
        mT_dm = self.modelTransformDeriv(m)

        dCdu = A
        dCdm = D * ( sdiag( G * u ) * ( Av_dm * ( mT_dm * v ) ) )

        if solve is None:
            solve = linalg.factorized(dCdu)

        Jv = - P * solve(dCdm)
        return Jv

    def Jt(self, m, v, u=None, solve=None):
        P = self.P
        D = self.mesh.faceDiv
        G = self.mesh.cellGrad
        A = self.createMatrix(m)
        Av_dm = self.mesh.getFaceMassDeriv()
        mT_dm = self.modelTransformDeriv(m)

        dCdu = A.T

        if solve is None:
            solve = linalg.factorized(dCdu.tocsc())
        w = solve(P.T*v)

        Jtv = - mT_dm.T * ( Av_dm.T * ( sdiag( G * u ) * ( D.T * w ) ) )
        return Jtv


if __name__ == '__main__':

    from SimPEG.regularization import Regularization
    from SimPEG import inverse

    # Create the mesh
    h1 = np.ones(100)
    h2 = np.ones(100)
    mesh = TensorMesh([h1,h2])

    # Create some parameters for the model
    sig1 = 1
    sig2 = 0.01

    # Create a synthetic model from a block in a half-space
    p0 = [20, 20]
    p1 = [50, 50]
    condVals = [sig1, sig2]
    mSynth = ModelBuilder.defineBlockConductivity(p0,p1,mesh.gridCC,condVals)
    mesh.plotImage(mSynth, showIt=False)


    # Set up the projection
    nelec = 50
    spacelec = 2
    surfloc = 0.5
    elecini = 0.5
    elecend = 0.5+spacelec*(nelec-1)
    elecLocR = np.linspace(elecini, elecend, nelec)
    rxmidLoc = (elecLocR[0:nelec-1]+elecLocR[1:nelec])*0.5
    q, Q, rxmidloc = genTxRxmat(nelec, spacelec, surfloc, elecini, mesh)
    P = Q.T

    # Create some data
    class syntheticDCProblem(DCProblem, SyntheticProblem):
        pass

    synthetic = syntheticDCProblem(mesh);
    synthetic.P = P
    synthetic.RHS = q
    dobs, Wd = synthetic.createData(mSynth, std=0.05)

    u = synthetic.field(mSynth)
    mesh.plotImage(u[:,10], showIt=False)

    # Now set up the problem to do some minimization
    problem = DCProblem(mesh)
    problem.P = P
    problem.RHS = q
    problem.W = Wd
    problem.dobs = dobs
    m0 = mesh.gridCC[:,0]*0+sig1

    # print problem.misfit(m0)
    # print problem.misfit(mSynth)

    opt = inverse.InexactGaussNewton()
    reg = Regularization(mesh)

    inv = inverse.Inversion(problem, reg, opt)

    inv.run(m0)

    # Check Derivative
    derChk = lambda m: [problem.misfit(m), problem.misfitDeriv(m)]
    checkDerivative(derChk, mSynth)

    # Adjoint Test
    u = np.random.rand(mesh.nC)
    v = np.random.rand(mesh.nC)
    w = np.random.rand(dobs.shape[0])
    print w.dot(problem.J(mSynth, v, u=u))
    print v.dot(problem.Jt(mSynth, w, u=u))





def genTxRxmat(nelec, spacelec, surfloc, elecini, mesh):
    """ Generate projection matrix (Q) and """
    elecend = 0.5+spacelec*(nelec-1)
    elecLocR = np.linspace(elecini, elecend, nelec)
    elecLocT = elecLocR+1
    nrx = nelec-1
    ntx = nelec-1
    q = np.zeros((mesh.nC, ntx))
    Q = np.zeros((mesh.nC, nrx))

    for i in range(nrx):

        rxind1 = np.argwhere((mesh.gridCC[:,0]==surfloc) & (mesh.gridCC[:,1]==elecLocR[i]))
        rxind2 = np.argwhere((mesh.gridCC[:,0]==surfloc) & (mesh.gridCC[:,1]==elecLocR[i+1]))

        txind1 = np.argwhere((mesh.gridCC[:,0]==surfloc) & (mesh.gridCC[:,1]==elecLocT[i]))
        txind2 = np.argwhere((mesh.gridCC[:,0]==surfloc) & (mesh.gridCC[:,1]==elecLocT[i+1]))

        q[txind1,i] = 1
        q[txind2,i] = -1
        Q[rxind1,i] = 1
        Q[rxind2,i] = -1

    Q = sp.csr_matrix(Q)
    rxmidLoc = (elecLocR[0:nelec-1]+elecLocR[1:nelec])*0.5
    return q, Q, rxmidLoc