from gusto import *
from firedrake import PeriodicSquareMesh, exp, SpatialCoordinate, Constant


def setup_gaussian(dirname):
    n = 16
    L = 1.
    mesh = PeriodicSquareMesh(n, n, L)

    parameters = ShallowWaterParameters(H=1.0, g=1.0)
    dt = 0.1
    output = OutputParameters(dirname=dirname+'/sw_plane_gaussian_subcycled')
    diagnostic_fields = [CourantNumber()]

    state = State(mesh,
                  dt=dt,
                  output=output,
                  parameters=parameters,
                  diagnostic_fields=diagnostic_fields)

    eqns = ShallowWaterEquations(state, family="BDM", degree=1,
                                 fexpr=Constant(1.))

    D0 = state.fields("D")
    x, y = SpatialCoordinate(mesh)
    H = Constant(state.parameters.H)
    D0.interpolate(H + exp(-50*((x-0.5)**2 + (y-0.5)**2)))

    advected_fields = []
    advected_fields.append((SSPRK3(state, "u", options=EmbeddedDGOptions(), subcycles=2)))
    advected_fields.append((SSPRK3(state, "D", subcycles=2)))

    linear_solver = ShallowWaterSolver(state, eqns)

    # build time stepper
    stepper = CrankNicolson(state, eqns, advected_fields, linear_solver)

    return stepper


def run(dirname):
    stepper = setup_gaussian(dirname)
    stepper.run(t=0, tmax=0.3)


def test_subcycling(tmpdir):
    dirname = str(tmpdir)
    run(dirname)

@@ -246,17 +246,18 @@
  dt = 0.
-1
+08
 %0A    out
@@ -330,50 +330,8 @@
 ed')
-%0A    diagnostic_fields = %5BCourantNumber()%5D
 %0A%0A  
@@ -410,16 +410,16 @@
 output,%0A
+
         
@@ -453,63 +453,8 @@
 ters
-,%0A                  diagnostic_fields=diagnostic_fields
 )%0A%0A 
@@ -912,61 +912,8 @@
 ))%0A%0A
-    linear_solver = ShallowWaterSolver(state, eqns)%0A%0A
     
@@ -993,23 +993,8 @@
 elds
-, linear_solver
 )%0A%0A 

fix test_subcycling - this involves slightly reducing the timestep which is ok because the test is not doing the right thing on master (the timestep is also reduced in the forcing there which is wrong)