import numpy as np
import parmepy as pp
import math
from parmepy.fields.continuous import Field
from parmepy.operator.analytic import Analytic
from parmepy.domain.obstacle.sphere import Sphere
pi = math.pi
from parmepy.operator.penalization import Penalization


def vitesse(x, y, z):
    return x, y, z

nb = 9

Lx = Ly = Lz = 2
dom = pp.Box(dimension=3, length=[Lx, Ly, Lz], origin=[-1., -1., -1.])
resol3D = [nb, nb, nb]

# Fields
scal = Field(domain=dom, name='Scalar')
op = Analytic(scal, formula=vitesse, resolutions={scal: resol3D})

op.setUp()

topo = dom.topologies[0]
coords = topo.mesh.coords

ff = scal.discreteFields[topo].data

ff[...]=128

sphere = Sphere(dom, position=[0., 0., 0.], radius=0.5)
ind = sphere.discretize(topo)

penal = Penalization(scal, sphere, coeff=1e6,
                     resolutions={scal: resol3D})

penal.setUp()

print scal.norm()

penal.apply(0.1)

print scal.norm()