from anuga.abstract_2d_finite_volumes.generic_domain import Generic_Domain
import anuga.abstract_2d_finite_volumes.mesh_factory as mesh_factory
import anuga.abstract_2d_finite_volumes.neighbour_mesh as neighbour_mesh
from anuga.abstract_2d_finite_volumes.generic_boundary_conditions import Dirichlet_boundary
from kinematic_viscosity import Kinematic_Viscosity_Operator
import numpy as num
import anuga.utilities.log as log


# FIXME SR: Seems a little strange that the timing for the elliptic and parabolic solves
# are the same. should investigate

print "Setting up"
grid_rows = 40
grid_cols = 40
#Set up a rectangular grid
points,elements,boundary = mesh_factory.rectangular_cross(grid_rows,grid_cols)
mesh = neighbour_mesh.Mesh(points,elements)
n = len(mesh)
print "n =",n
boundary_map = {}
for (vol_id,edge) in boundary.keys():
    x = mesh.get_edge_midpoint_coordinate(vol_id,edge)[0]
    y = mesh.get_edge_midpoint_coordinate(vol_id,edge)[1]
    #Define boundary conditions: u=x^2-y^2, v=x-y+2
    boundary_map[(vol_id,edge)] = Dirichlet_boundary([1,x*x-y*y,x-y+2])
#Now define an operator
domain = Generic_Domain(source=points,triangles=elements,boundary=boundary_map)
KV_Operator = Kinematic_Viscosity_Operator(domain)
print "Applying stage heights"
#Set up the operator and RHS
h = num.ones((n,),num.float)
    
def solve():   
    KV_Operator.apply_stage_heights(h)
    rhs = num.zeros((n,2),num.float) #we will solve for u and v
    print "Solving"
    #Solve
    velocities = KV_Operator.cg_solve(rhs)
    new_velocities = KV_Operator.parabolic_solver(velocities)

#Do the profiling
import profile, pstats
FN = 'kv_profile.dat'
print "Starting solve..."
profile.run('solve()', FN)
S = pstats.Stats(FN)
s = S.sort_stats('cumulative').print_stats(30)