"""Simulation of the nonsymmetrical dam dreak problem.

   Copyright 2005
   Christopher Zoppou, Stephen Roberts
   Australian National University

"""

#---------------------------------------
# Setup Path and import modules
import sys
from os import sep, path
sys.path.append('..'+sep+'pyvolution')

from shallow_water import Domain, Transmissive_boundary, Reflective_boundary,\
     Dirichlet_boundary
#from pmesh2domain import pmesh_to_domain_instance
from anuga.pyvolution.util import Polygon_function
from mesh_factory import rectangular_cross


def cut_out_region(domain):
    """
    To do: make better comments!
    Deal with passing the boundary info as well
    """

    points = domain.coordinates
    elements = domain.triangles
    boundary = domain.boundary
    centroid_coordinates = domain.centroid_coordinates
    N = domain.number_of_elements

    elements_in  = []
    elements_out = []
    for i in range(N):
        element = elements[i]
        #print element
        [x,y] = centroid_coordinates[i]
        #print x,y
        if x>10 and y>10:
            #print i,'Out region'
            elements_out.append(i)
        else:
            #print i,'In region'
            elements_in.append(i)

    #print elements_in
    #print elements_out

    points_in = {}
    for i in elements_in:
        #print i
        [v0,v1,v2] = elements[i]
        #print v0,v1,v2
        points_in[v0] = v0
        points_in[v1] = v1
        points_in[v2] = v2


    new_index = []
    for i,value in enumerate(points_in):
        #print i , value
        points_in[value] = i


    #print points_in
    new_elements  = []
    for i in elements_in:
        #print i
        [v0,v1,v2] = elements[i]
        #print v0,v1,v2

        nv0 = points_in[v0]
        nv1 = points_in[v1]
        nv2 = points_in[v2]

        new_elements.append([nv0,nv1,nv2])

    new_points = []
    for key,value in points_in.iteritems():
        [x,y] = points[key]
        new_points.append([x,y])

    #print new_points

    return new_points, new_elements







#---------------------------------------
# Boundary conditions
h0 = 1.0
h1 = 10.0
print 'Boundary conditions'

#---------------------------------------
# Domain
n = 100
m = 100
lenx = 20.0
leny = 20.0
delta_x = lenx/n
delta_y = leny/m
origin = (0.0, 0.0)

points, elements, boundary = rectangular_cross(m, n, lenx, leny, origin)
domain = Domain(points, elements, boundary)

new_points, new_elements = cut_out_region(domain)
domain = Domain(new_points, new_elements)


R = Reflective_boundary(domain)
T = Transmissive_boundary(domain)
D = Dirichlet_boundary([h1, 0.0, 0.0])
domain.set_boundary({'exterior': R})







print "Number of triangles = ", len(domain)


#---------------------------------------
#Initial condition
p0 = [[0.0, 0.0], [10.0, 0.0], [10.0, 20.0], [0.0, 20.0]]
domain.set_quantity('stage',Polygon_function([(p0,h1)],default = h0))


#---------------------------------------
# Provide file name for storing output
domain.store = True
domain.format = 'sww'
domain.set_name('Non_symetrical_Dam_Break_second_order')


# Visualization smoothing
domain.visualise=True

#---------------------------------------
#Decide which quantities are to be stored at each timestep
domain.quantities_to_be_stored = ['stage', 'xmomentum', 'ymomentum']


#---------------------------------------
# Set bed elevation
def x_slope(x,y):
    n = x.shape[0]
    z = 0*x
    return z
domain.set_quantity('elevation', x_slope)


#----------------------------
# Friction
domain.set_quantity('friction', 0.0)


#--------------------------------------
# Evolution

yieldstep = 0.1
finaltime = 15.0

domain.CFL = 0.75

domain.default_order = 1
domain.smooth = True

import time
t0 = time.time()
for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    domain.write_time()

print 'That took %.2f seconds' %(time.time()-t0)