from dolfin import *
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter



#cf 'Contenu des fichiers' à propos de 'reference'
def reference(tabC,res1,c,ec,fr):
    #res=256
    k,ll,mm,nn=tabC.shape
    mesh = UnitSquareMesh(res1, res1, "crossed")

    class LocalC(UserExpression):
        def __init__(self, Cij, mesh):
            self.Cij = Cij
            self.Nx, self.Ny, a, b = Cij.shape
            assert a==3 and b==3, "Wrong shape"
            self.mesh = mesh
            UserExpression.__init__(self)
        def eval_cell(self, value, x, ufc_cell):
            p = Cell(mesh, ufc_cell.index).midpoint()
            i, j = int(p[0]*(self.Nx)), int(p[1]*(self.Ny))
            value[:] = self.Cij[j, i,:,:].flatten()

        def value_shape(self):
            return (3, 3)



#Ici on crée, à partir du tableau des tenseurs d'élasticité correspondant au domaine dont on veut calculer le Chom, un tableau Cv représentant le même Cn mais adapté à une résolution plus fine 'res'
    # Cv = np.array([[np.zeros((3,3)) for i in range(res)] for j in range(res)])
    # for i0 in range(k):
    #     for j0 in range(k):
    #         for i in range(int(i0*res/k),int((i0+1)*res/k)):
    #             for j in range (int(j0*res/k),int((j0+1)*res/k)):
    #                 Cv[i,j]=tabC[i0,j0]

    # plt.imshow(tabC[:,:,0,0])
    # plt.show()
    # plt.close()
    # plt.imshow(Cv[:,:,0,0])
    # plt.show()
    # plt.close()

#Cvar est la transformation du tableau Cv en une fonction définie sur 'mesh'
    Cvar= LocalC(tabC,mesh)

    # class used to define the periodic boundary map
    class PeriodicBoundary(SubDomain):
        def __init__(self, vertices, tolerance=DOLFIN_EPS):
            """ vertices stores the coordinates of the 4 unit cell corners"""
            SubDomain.__init__(self, tolerance)
            self.tol = tolerance
            self.vv = vertices
            self.a1 = self.vv[1,:]-self.vv[0,:] # first vector generating periodicity
            self.a2 = self.vv[3,:]-self.vv[0,:] # second vector generating periodicity
            # check if UC vertices form indeed a parallelogram
            assert np.linalg.norm(self.vv[2, :]-self.vv[3, :] - self.a1) <= self.tol
            assert np.linalg.norm(self.vv[2, :]-self.vv[1, :] - self.a2) <= self.tol

        def inside(self, x, on_boundary):
            # return True if on left or bottom boundary AND NOT on one of the
            # bottom-right or top-left vertices
            return bool((near(x[0], self.vv[0,0] + x[1]*self.a2[0]/self.vv[3,1], self.tol) or
                        near(x[1], self.vv[0,1] + x[0]*self.a1[1]/self.vv[1,0], self.tol)) and
                        (not ((near(x[0], self.vv[1,0], self.tol) and near(x[1], self.vv[1,1], self.tol)) or
                        (near(x[0], self.vv[3,0], self.tol) and near(x[1], self.vv[3,1], self.tol)))) and on_boundary)

        def map(self, x, y):
            if near(x[0], self.vv[2,0], self.tol) and near(x[1], self.vv[2,1], self.tol): # if on top-right corner
                y[0] = x[0] - (self.a1[0]+self.a2[0])
                y[1] = x[1] - (self.a1[1]+self.a2[1])
            elif near(x[0], self.vv[1,0] + x[1]*self.a2[0]/self.vv[2,1], self.tol): # if on right boundary
                y[0] = x[0] - self.a1[0]
                y[1] = x[1] - self.a1[1]
            else:   # should be on top boundary
                y[0] = x[0] - self.a2[0]
                y[1] = x[1] - self.a2[1]



#la fonction reference_problem prend en argument C, donnant le tenseur d'élasticité en un point de 'mesh', et renvoie le Chom du domaine. k étant le nombre de phases différentes sur le domaine, qui correspond à la taille du tableau Cn pris en argument par la fonction 'reference'
    def reference_problem(C, mesh,k):

        def eps(v):
            return as_vector([v[0].dx(0),
                            v[1].dx(1),
                            (v[0].dx(1)+v[1].dx(0))/sqrt(2)])

        def sigma(v, Eps):
            return dot(C, eps(v)+Eps)


        Ve = VectorElement("CG", mesh.ufl_cell(), 2)
        Re = VectorElement("R", mesh.ufl_cell(), 0)
        vertices = np.array([[0, 0], [1, 0], [1, 1], [0, 1]])
        W = FunctionSpace(mesh, MixedElement([Ve, Re]), constrained_domain=PeriodicBoundary(vertices, tolerance=1e-10))
        V = VectorFunctionSpace(mesh,'CG',degree=2,dim=3)

        v_,lamb_ = TestFunctions(W)
        dv, dlamb = TrialFunctions(W)
        w = Function(W)

        Eps = Constant((0,)*3)
        F = inner(sigma(dv, Eps), eps(v_))*dx
        a, L = lhs(F), rhs(F)
        a += dot(lamb_,dv)*dx + dot(dlamb,v_)*dx

        Chom = np.zeros((3, 3))
        for (j, case) in enumerate(["Exx", "Eyy", "Exy"]):
            ee = np.zeros((3,))
            ee[j] = 1
            Eps.assign(Constant(ee))
            solve(a == L, w, solver_parameters={"linear_solver": "mumps"})#, solver_parameters={"linear_solver": "bicgstab", "preconditioner":"ilu"})
            (v, lamb) = split(w)
            ffile = XDMFFile('c='+str(c)+'/ech='+str(ec)+'/frac='+str(fr)+'/Strain EF '+case+'.xdmf')
            ffile.parameters["functions_share_mesh"] = True
            vf = Function(V, name='Strain-'+case)
            vf.assign(project(eps(v),V))
            ffile.write(vf,0.)

            plt.figure()
            pp=plot(eps(v)[0])
            plt.colorbar(pp)
            plt.savefig('c='+str(c)+'/ech='+str(ec)+'/frac='+str(fr)+'/Aloc EF 1'+str(j+1)+'.pdf')
            #plt.show()
            plt.close()
            plt.figure()
            pp=plot(eps(v)[1])
            plt.colorbar(pp)
            plt.savefig('c='+str(c)+'/ech='+str(ec)+'/frac='+str(fr)+'/Aloc EF 2'+str(j+1)+'.pdf')
            #plt.show()
            plt.close()
            plt.figure()
            pp=plot(eps(v)[2])
            plt.colorbar(pp)
            plt.savefig('c='+str(c)+'/ech='+str(ec)+'/frac='+str(fr)+'/Aloc EF 3'+str(j+1)+'.pdf')
            #plt.show()
            plt.close()

            Sigma = np.zeros((3,))
            for kk in range(3):
                Sigma[kk] = assemble(sigma(v, Eps)[kk]*dx)
            Chom[j, :] = Sigma

        ffile.close()
        return Chom

    Chom = reference_problem(Cvar, mesh,k)
    CVoigt = np.array([[0., 0., 0.],
                [0., 0., 0.],
                [0., 0., 0.]])
    for i in range(k):
        for j in range(k):
            CVoigt=CVoigt+tabC[i,j]

    CVoigt=CVoigt/k/k

    print(Chom)
    return Chom,CVoigt