#!/usr/bin/env python

## Import statements
#from fipy.meshes.grid2D import Grid2D
from fipy.meshes.periodicGrid2D import PeriodicGrid2DLeftRight
from fipy.variables.cellVariable import CellVariable
from fipy.viewers import MatplotlibViewer
from fipy.solvers.trilinos.linearGMRESSolver import LinearGMRESSolver
from fipy.terms.implicitDiffusionTerm import ImplicitDiffusionTerm
from fipy.terms.transientTerm import TransientTerm
from fipy import raw_input
from fipy import Grid2D
from fipy import parallel
#from fipy import *

# from mpi4py import MPI
# comm = MPI.COMM_WORLD
# rank = comm.Get_rank()

## Discretization
dx = 0.035e-6
dy = 0.035e-6

nx = 300
ny = 300

## Mesh construction
# mesh = Grid2D(dx=dx, dy=dy, nx=nx, ny=ny)
mesh = PeriodicGrid2DLeftRight(dx=dx, dy=dy, nx=nx, ny=ny)

## Define cell variable
phi = CellVariable(name="phi",
                     mesh=mesh,
                     value=0.0,
                     hasOld=1)

# set phi variable
center = (0.5 * nx * dx, 0.0)
radius = 1.25e-6

x, y = mesh.cellCenters[0], mesh.cellCenters[1]

print min(mesh.y)
print min(y)
print 300 * 300
print len(y)
phi.setValue(1.0,
               where=((x - center[0])**2 \
                          + (y - center[1])**2) < radius**2)

# View phi variable
viewer = MatplotlibViewer(phi, datamax=1., datamin=0.)
viewer.plot()

raw_input("Press Enter to continue...")

# Setup a simple diffusion equation
equation = TransientTerm() == ImplicitDiffusionTerm(coeff=1e-9)

solver = LinearGMRESSolver(tolerance=1e-6)

for step in range(5):
    phi.updateOld()

    equation.solve(phi, solver=solver, dt=1e-5)
    
    #residual = 1e10
    #while residual > 1e-6:
    #    residual = equation.sweep(phi, solver=solver, dt=1e-5)

    viewer.plot()
    if parallel.procID == 0:
        print step

raw_input("Press Enter to continue...")