import petsc4py
import sys

petsc4py.init(sys.argv)

from petsc4py import PETSc
import mshr
from dolfin import *
import sympy as sy
import numpy as np
import ExactSol
import MatrixOperations as MO
import CheckPetsc4py as CP
from  dolfin import __version__

def Domain(n):

    # defining the L-shaped domain
    # domain = mshr.Rectangle(Point(-1., -1.), Point(1., 1.)) - mshr.Rectangle(Point(0., -1.), Point(1., 0.) )
    # mesh = mshr.generate_mesh(domain, n)
    if __version__ == '1.6.0':
        mesh = RectangleMesh(Point(-1., -1.), Point(1., 1.),n,n)
    else:
        mesh = RectangleMesh(-1,-1,1,1,n,n, 'left')
    cell_f = CellFunction('size_t', mesh, 0)
    for cell in cells(mesh):
        v = cell.get_vertex_coordinates()
        y = v[np.arange(0,6,2)]
        x = v[np.arange(1,6,2)]
        xone = np.ones(3)
        xone[x > 0] = 0
        yone = np.ones(3)
        yone[y < 0] = 0
        if np.sum(xone)+ np.sum(yone)>5.5:
            cell_f[cell] = 1
    mesh = SubMesh(mesh, cell_f, 0)


    # cell_markers = CellFunction("bool", mesh)
    # cell_markers.set_all(False)
    # origin = Point(0., 0.)
    # for cell in cells(mesh):
    #     p = cell.midpoint()
    #     if abs(p.distance(origin)) < 0.6:
    #         cell_markers[cell] = True

    # mesh = refine(mesh, cell_markers)


    # cell_markers = CellFunction("bool", mesh)
    # cell_markers.set_all(False)
    # origin = Point(0., 0.)
    # for cell in cells(mesh):
    #     p = cell.midpoint()
    #     if abs(p.distance(origin)) < 0.4:
    #         cell_markers[cell] = True

    # mesh = refine(mesh, cell_markers)


    # cell_markers = CellFunction("bool", mesh)
    # cell_markers.set_all(False)
    # origin = Point(0., 0.)
    # for cell in cells(mesh):
    #     p = cell.midpoint()
    #     if abs(p.distance(origin)) < 0.2:
    #         cell_markers[cell] = True

    # mesh = refine(mesh, cell_markers)



    # Creating classes that define the boundary of the domain
    class Left(SubDomain):
        def inside(self, x, on_boundary):
            return near(x[0], -1.0)

    class Right(SubDomain):
        def inside(self, x, on_boundary):
            return near(x[0], 1.0)

    class Bottom(SubDomain):
        def inside(self, x, on_boundary):
            return near(x[1], -1.0)

    class Top(SubDomain):
        def inside(self, x, on_boundary):
            return near(x[1], 1.0)

    class CornerTop(SubDomain):
        def inside(self, x, on_boundary):
            return near(x[1], 0.0) and between(x[0], (0.0,1.0))

    class CornerLeft(SubDomain):
        def inside(self, x, on_boundary):
            return near(x[0], 0.0) and between(x[1], (-1.0,0.0))


    left = Left()
    top = Top()
    right = Right()
    bottom = Bottom()
    cleft = CornerLeft()
    ctop = CornerTop()

    # Initialize mesh function for the domain
    domains = CellFunction("size_t", mesh)
    domains.set_all(0)

    # Initialize mesh function for boundary domains
    boundaries = FacetFunction("size_t", mesh)
    boundaries.set_all(0)

    left.mark(boundaries, 1)
    top.mark(boundaries, 1)
    bottom.mark(boundaries, 1)
    right.mark(boundaries, 1)
    cleft.mark(boundaries, 2)
    ctop.mark(boundaries, 2)
    return mesh, boundaries, domains


# functions that perform partial derivatives of x and y with respect to polar coordinates
def polarx(u, rho, phi):
    return sy.cos(phi)*sy.diff(u, rho) - (1./rho)*sy.sin(phi)*sy.diff(u, phi)

def polary(u, rho, phi):
    return sy.sin(phi)*sy.diff(u, rho) + (1./rho)*sy.cos(phi)*sy.diff(u, phi)

def polarr(u, x, y):
    return (1./sqrt(x**2 + y**2))*(x*sy.diff(u,x)+y*sy.diff(u,y))

def polart(u, x, y):
    return -y*sy.diff(u,x)+x*sy.diff(u,y)

def SolutionPolar(mesh, params):

    l = 0.54448373678246
    omega = (3./2)*np.pi

    phi = sy.symbols('x[1]')
    rho = sy.symbols('x[0]')
    z = sy.symbols('z')

    # looked at all the exact solutions and they seems to be the same as the paper.....
    psi = (sy.sin((1+l)*phi)*sy.cos(l*omega))/(1+l) - sy.cos((1+l)*phi) - (sy.sin((1-l)*phi)*sy.cos(l*omega))/(1-l) + sy.cos((1-l)*phi)
    psi_prime = sy.diff(psi, phi)
    psi_3prime = sy.diff(psi, phi, phi, phi)

    u = rho**l*((1+l)*sy.sin(phi)*psi + sy.cos(phi)*psi_prime)
    v = rho**l*(-(1+l)*sy.cos(phi)*psi + sy.sin(phi)*psi_prime)

    uu0 = Expression((sy.ccode(u),sy.ccode(v)))

    # sssss
    p = -rho**(l-1)*((1+l)**2*psi_prime + psi_3prime)/(1-l)
    pu0 = Expression(sy.ccode(p))

    f = rho**(2./3)*sy.sin((2./3)*phi)
    # b = sy.diff(f,rho)
    b = polarx(f, rho, phi)
    # d = (1./rho)*sy.diff(f,phi)
    d = polary(f, rho, phi)

    bu0 = Expression((sy.ccode(b),sy.ccode(d)))

    r = sy.diff(phi,rho)
    ru0 = Expression(sy.ccode(r))


    # Defining polarx and polary as the x and y derivatives with respect to polar coordinates (rho, phi). Writing the right handside with respect to cartesian coords

    #Laplacian
    L1 = polarx(polarx(u, rho, phi), rho, phi) + polary(polary(u, rho, phi), rho, phi)
    L2 = polarx(polarx(v, rho, phi), rho, phi) + polary(polary(v, rho, phi), rho, phi)

    # Advection
    A1 = u*polarx(u, rho, phi)+v*polary(u, rho, phi)
    A2 = u*polarx(v, rho, phi)+v*polary(v, rho, phi)

    # Pressure gradient
    P1 = polarx(p, rho, phi)
    P2 = polary(p, rho, phi)

    # Curl-curl
    C1 = polarx(polary(d, rho, phi), rho, phi) - polary(polary(b, rho, phi), rho, phi)
    C2 = polarx(polary(b, rho, phi), rho, phi) - polary(polary(d, rho, phi), rho, phi)

    # Multiplier gradient
    R1 = sy.diff(r, rho)
    R2 = sy.diff(r, rho)

    # Coupling term for fluid variables
    NS1 = -d*(polarx(d, rho, phi)-polary(b, rho, phi))
    NS2 = b*(polarx(d, rho, phi)-polary(b, rho, phi))

    # Coupling term for Magnetic variables
    M1 = polary(u*d-v*b, rho, phi)
    M2 = -polarx(u*d-v*b, rho, phi)

    # Using https://en.wikipedia.org/wiki/Del_in_cylindrical_and_spherical_coordinates defintitions of the derivative operators      (sy.diff(u,rho) means partial derivative of u with respect to rho)

    # Laplacian
    L11 = (1./rho)*sy.diff(rho*sy.diff(u,rho),rho) + (1./(rho**2))*sy.diff(sy.diff(u,phi),phi) - (1./rho**2)*u - (2./rho**2)*sy.diff(v, phi)
    L22 = (1./rho)*sy.diff(rho*sy.diff(v,rho),rho) + (1./(rho**2))*sy.diff(sy.diff(v,phi),phi) - (1./rho**2)*v + (2./rho**2)*sy.diff(u, phi)


    # Advection
    A11 = u*sy.diff(u, rho) + (1./rho)*v*sy.diff(u, phi) - u**2/rho
    A22 = u*sy.diff(v, rho) + (1./rho)*v*sy.diff(v, phi) + v*u/rho

    # Pressure gradient
    P11 = sy.diff(p, rho)
    P22 = (1./rho)*sy.diff(p, phi)

    # Curl-curl
    c = (1./rho)*(sy.diff(rho*d, rho) - sy.diff(b, phi))
    C11 = (1./rho)*sy.diff(c, phi)
    C22 = -sy.diff(c, rho)

    # Multiplier gradient
    R11 = sy.diff(r, rho)
    R22 = sy.diff(r, rho)

    # Coupling term for fluid variables
    NS11 = -c*d
    NS22 = c*b

    # Coupling term for Magnetic variables
    c = u*d-v*b
    M11 = (1./rho)*sy.diff(c, phi)
    M22 = -sy.diff(c, rho)
    FF = sy.diff(u, rho) + (1./rho)*sy.diff(v, phi)

    # print "\n\n\nL limits \n\n"
    # print sy.limit(L1, rho,0), sy.limit(sy.limit(L1, phi,0),rho,0)
    # print sy.limit(L11, rho,0), sy.limit(sy.limit(L11, phi,0),rho,0)
    # print "\n", sy.limit(L2, rho,0), sy.limit(sy.limit(L2, phi,0),rho,0)
    # print sy.limit(L22, rho,0), sy.limit(sy.limit(L22, phi,0),rho,0)

    # print "\n\n\nA limits \n\n"
    # print sy.limit(A1, rho,0), sy.limit(sy.limit(A1, phi,0),rho,0)
    # print sy.limit(A11, rho,0), sy.limit(sy.limit(A11, phi,0),rho,0)
    # print "\n", sy.limit(A2, rho,0), sy.limit(sy.limit(A2, phi,0),rho,0)
    # print sy.limit(A22, rho,0), sy.limit(sy.limit(A22, phi,0),rho,0)

    # print "\n\n\nP limits \n\n"
    # print sy.limit(P1, rho,0), sy.limit(sy.limit(P1, phi,0),rho,0)
    # print sy.limit(P11, rho,0), sy.limit(sy.limit(P11, phi,0),rho,0)
    # print "\n", sy.limit(P2, rho,0), sy.limit(sy.limit(P2, phi,0),rho,0)
    # print sy.limit(P22, rho,0), sy.limit(sy.limit(P22, phi,0),rho,0)

    # print "\n\n\nC limits \n\n"
    # print sy.limit(C1, rho,0), sy.limit(sy.limit(C1, phi,0),rho,0)
    # print sy.limit(C11, rho,0), sy.limit(sy.limit(C11, phi,0),rho,0)
    # print "\n", sy.limit(C2, rho,0), sy.limit(sy.limit(C2, phi,0),rho,0)
    # print sy.limit(C22, rho,0), sy.limit(sy.limit(C22, phi,0),rho,0)

    # print "\n\n\nR limits \n\n"
    # print sy.limit(R1, rho,0), sy.limit(sy.limit(R1, phi,0),rho,0)
    # print sy.limit(R11, rho,0), sy.limit(sy.limit(R11, phi,0),rho,0)
    # print "\n", sy.limit(R2, rho,0), sy.limit(sy.limit(R2, phi,0),rho,0)
    # print sy.limit(R22, rho,0), sy.limit(sy.limit(R22, phi,0),rho,0)

    # print "N\n\n\nS limits \n\n"
    # print sy.limit(NS1, rho,0), sy.limit(sy.limit(NS1, phi,0),rho,0)
    # print sy.limit(NS11, rho,0), sy.limit(sy.limit(NS11, phi,0),rho,0)
    # print "\n", sy.limit(NS2, rho,0), sy.limit(sy.limit(NS2, phi,0),rho,0)
    # print sy.limit(NS22, rho,0), sy.limit(sy.limit(NS22, phi,0),rho,0)

    # print "\n\n\nM limits \n\n"
    # print sy.limit(M1, rho,0), sy.limit(sy.limit(M1, phi,0),rho,0)
    # print sy.limit(M11, rho,0), sy.limit(sy.limit(M11, phi,0),rho,0)
    # print "\n", sy.limit(M2, rho,0), sy.limit(sy.limit(M2, phi,0),rho,0)
    # print sy.limit(M22, rho,0), sy.limit(sy.limit(M22, phi,0),rho,0)

    # print "\n\n\Fluid limits \n\n"
    # print sy.limit(u, rho,0), sy.limit(sy.limit(u, phi,0),rho,0)
    # print sy.limit(v, rho,0), sy.limit(sy.limit(v, phi,0),rho,0)
    # print sy.limit(p, rho,0), sy.limit(sy.limit(p, phi,0),rho,0)

    # print "\n\n\Magnetic limits \n\n"
    # print sy.limit(b, rho,0), sy.limit(sy.limit(b, phi,0),rho,0)
    # print sy.limit(d, rho,0), sy.limit(sy.limit(d, phi,0),rho,0)
    # print sy.limit(r, rho,0), sy.limit(sy.limit(r, phi,0),rho,0)



    # ssss
    # graduu0 = Expression(sy.ccode(sy.diff(u, rho) + (1./rho)*sy.diff(u, phi)))
    graduu0 = Expression((sy.ccode(sy.diff(u, rho)),sy.ccode(sy.diff(v, rho))))
    Laplacian = Expression((sy.ccode(L11),sy.ccode(L22)))
    Advection = Expression((sy.ccode(A11),sy.ccode(A22)))
    gradPres = Expression((sy.ccode(P11),sy.ccode(P22)))
    CurlCurl = Expression((sy.ccode(C11),sy.ccode(C22)))
    gradR = Expression((sy.ccode(R11).replace('M_PI','pi'),sy.ccode(R22).replace('M_PI','pi')))
    NS_Couple = Expression((sy.ccode(NS11),sy.ccode(NS22)))
    M_Couple = Expression((sy.ccode(M11),sy.ccode(M22)))

    # ignore this! Just removes the singularity (atan2(0,0) = NaN) and makes all functions zero at the origin
    class u0(Expression):
        def __init__(self, mesh, uu0):
            self.mesh = mesh
            self.u0 = uu0
        def eval_cell(self, values, x, ufc_cell):
            if abs(x[0]) < 1e-3 and abs(x[1]) < 1e-3:
                values[0] = 0.0
                values[1] = 0.0
            else:
                r = sqrt(x[0]**2 + x[1]**2)
                theta = np.arctan2(x[1],x[0])
                # print theta
                if theta < 0:
                    theta += 2*np.pi
                values[0] = self.u0(r, theta)[0]
                values[1] = self.u0(r,theta)[1]
        def value_shape(self):
            return (2,)

    class gradu0(Expression):
        def __init__(self, mesh, graduu0):
            self.mesh = mesh
            self.gradu0 = graduu0
        def eval_cell(self, values, x, ufc_cell):
            # if abs(x[0]) < 1e-3 and abs(x[1]) < 1e-3:
            #     values =  0.0
            # else:
                r = sqrt(x[0]**2 + x[1]**2)
                theta = np.arctan2(x[1],x[0])
                if theta < 0:
                    theta += 2*np.pi
                values = self.gradu0(r,theta)
    class p0(Expression):
        def __init__(self, mesh, pu0):
            self.mesh = mesh
            self.p0 = pu0
        def eval_cell(self, values, x, ufc_cell):
            if abs(x[0]) < 1e-3 and abs(x[1]) < 1e-3:
                values[0] = 0.0
            else:
                r = sqrt(x[0]**2 + x[1]**2)
                theta = np.arctan2(x[1],x[0])
                if theta < 0:
                    theta += 2*np.pi
                values[0] = self.p0(r,theta)
        def eval(self, values, x):
            # if abs(x[0]) < 1e-3 and abs(x[1]) < 1e-3:
            #     values = 0.0
            # else:
                r = sqrt(x[0]**2 + x[1]**2)
                theta = np.arctan2(x[1],x[0])
                if theta < 0:
                    theta += 2*np.pi
                values[0] = self.p0(r,theta)
    class p0Vec(Expression):
        def __init__(self, mesh, pu0):
            self.mesh = mesh
            self.p0 = pu0
        def eval_cell(self, values, x, ufc_cell):
            # if abs(x[0]) < 1e-3 and abs(x[1]) < 1e-3:
            #     values[0] = 0.0
            # else:
                r = sqrt(x[0]**2 + x[1]**2)
                theta = np.arctan2(x[1],x[0])
                if theta < 0:
                    theta += 2*np.pi
                values[0] = self.p0(r,theta)
                values[1] = self.p0(r,theta)
        def eval(self, values, x):
            # if abs(x[0]) < 1e-3 and abs(x[1]) < 1e-3:
            #     values = 0.0
            # else:
                r = sqrt(x[0]**2 + x[1]**2)
                theta = np.arctan2(x[1],x[0])
                if theta < 0:
                    theta += 2*np.pi
                values[0] = self.p0(r,theta)
                values[1] = self.p0(r,theta)