import parmepy
#parmepy.__VERBOSE__ = True
from parmepy.domain.box import Box
# import parmepy.gpu
# parmepy.gpu.CL_PROFILE = True
from parmepy.fields.continuous import Field
from parmepy.operator.advection import Advection
from parmepy.problem.transport import TransportProblem
from parmepy.operator.monitors.printer import Printer
from parmepy.problem.simulation import Simulation
from parmepy.operator.analytic import Analytic
from parmepy.operator.redistribute import Redistribute
from parmepy.methods_keys import TimeIntegrator, Interpolation, Remesh,\
    Support, Splitting, MultiScale
#from parmepy.numerics.integrators.runge_kutta2 import RK2 as RK
from parmepy.numerics.integrators.runge_kutta4 import RK4 as RK
from parmepy.numerics.interpolation import Linear
from parmepy.numerics.remeshing import L6_6 as rmsh
from parmepy.mpi.main_var import main_size
from parmepy.constants import np, HDF5


def vitesse(res, x, y, t=0.):
    res[0][...] = -np.sin(x * np.pi) ** 2 * np.sin(y * np.pi * 2) * \
        np.cos(t * np.pi / 3.)
    res[1][...] = np.sin(y * np.pi) ** 2 * np.sin(x * np.pi * 2) * \
        np.cos(t * np.pi / 3.)
    return res


def scalaire(res, x, y, t=0.):
    rr = np.sqrt((x - 0.5) ** 2 + (y - 0.75) ** 2)
    res[0][...] = 0.
    res[0][rr < 0.15] = 1.
    return res


dim = 2
boxLength = [1., 1.]
boxMin = [0., 0.]
nbElem = [513, ] * 2

timeStep = 0.025
finalTime = 3.
outputFilePrefix = 'levelSet_2D_'
outputModulo = 10
simu = Simulation(tinit=0.0, tend=finalTime, timeStep=timeStep, iterMax=120)

## Domain
box = Box(dim, length=boxLength, origin=boxMin)

## Fields
scal = Field(domain=box, name='Scalar', formula=scalaire
             )
velo = Field(domain=box, name='Velocity', isVector=True, formula=vitesse
             )

## Operators
# By default, with a 'gpu' support, operator is going to use the defaut device
# given at cmake. To specify an other device, user must set the proper
# parameters: platform_id and device_id.
# parameter device_type can be used to get a specific device type
advec = Advection(velo, scal,
                  resolutions={velo: nbElem,
                               scal: nbElem},
                  method={TimeIntegrator: RK,
                          Interpolation: Linear,
                          Remesh: rmsh,
                          Support: 'gpu_1k',
                          Splitting: 'o2'},
                  # platform_id=0,
                  # device_id=0,
                  # device_type='cpu'
                  )
advec.discretize()
velocity = Analytic(velo,
                    resolutions={velo: nbElem},
                    topo=advec.advecDir[0].discreteFields[velo].topology
                    )

if main_size > 1:
    distrForAdvecY = Redistribute([velo], velocity, advec.advecDir[1],
                                  component=1)

##Problem
# Sequential : no need of redistribute
if main_size == 1:
    pb = TransportProblem([velocity, advec], simu, dumpFreq=-1)
else:
    pb = TransportProblem([velocity, distrForAdvecY, advec],
                          simu, dumpFreq=-1)

## Setting solver to Problem
pb.setUp()

print scal.topoInit
p = Printer(variables=[scal],
            topo=scal.topoInit,
            frequency=outputModulo,
            prefix=outputFilePrefix,
            formattype=HDF5)
pb.addMonitors([p])
p.apply(simu)

## Solve problem
pb.solve()

p.apply(simu)

pb.finalize()
print pb.timer