"""
----------------------------------------------------------------- model_build.py
ANUGA MODELLING TEMPLATE (Domain Construction )
Script for building the model domain and setting initial conditions and
(temporally varying)  boundary conditions appropriate to the particular
model-scenario-event to be modelled. (Parallel processor version)

Most model data for the run is read in or otherwise pre-processed in
model_data.py and imported into model_build.py to create the model mesh and
domain before running the simulation.

Version        1.00 June 2008 Initial release for tuflow flood simulation
               2.00 August 2015 modified for tsunami modelling in SETAS
Author         : E Rigby (0437 250 500) ted.rigby@rienco.com.au
-------------------------------------------------------------------------------
"""

#------------------------------------------------------------------------------
# Import necessary modules
#------------------------------------------------------------------------------

# Standard python modules
import os
import time
import sys


# ANUGA modules
import anuga
from anuga import Region

# Model specific imports
import model_ini        # Defines and provides the scenario-event specific data
import model_data       # Reads in and pre-processes external data required by build

# Import the tools to manage the parallel interface
from anuga import distribute, myid, numprocs, finalize, barrier

# Import numpy
import numpy as num

barrier()
t0 = time.time()        # record start time of this build (secs)
if myid == 0:
    print model_ini.prefix+' >>>> Comencing model_build.py code execution at t= %.2f hours'  %(float(time.time()-model_ini.t0)/3600.0)

###############################################################################
#                                                                             #
#                    ONLY CREATE ONE DOMAIN (ON PROCESSOR 0)                  #
#                                                                             #
###############################################################################
if myid == 0:

    ################################################################################
    #                                                                              #
    #              CREATE THE TRIANGULAR MESH AND SINGLE DOMAIN                    #
    #                                                                              #
    ################################################################################

    # Note: The domain and mesh are created based on overall clipping polygon with a tagged
    # boundary and interior regions as defined in model_data.py along with
    # resolutions (maximal area m2 per triangle) for each interior polygon

    print model_ini.prefix+' >>>> Creating domain within bounding_poly reflecting regional resolutions'
    # Create the domain reflecting the bounding (default) and internal region mesh resolutions set in model_data.py
    # with boundary tags as defined in model_data.py


    domain = anuga.create_domain_from_regions(bounding_polygon  = model_data.bounding_polygon,
                                         boundary_tags          = model_data.boundary_strings,
                                         maximum_triangle_area  = model_ini.default_res,
                                         minimum_triangle_angle = 30.0,
                                         interior_regions       = model_data.interior_res_regions,
                                         mesh_filename          = model_ini.mesh_pathname,
                                         use_cache              = True,
                                         verbose                = model_ini.model_verbose)


    print model_ini.prefix+' >>>>> Completed  mesh %s and domain creation at t= %.2f hours ' %(model_ini.mesh_filename,float(time.time()-model_ini.t0)/3600.0)


    # set the basic domain parameters
    domain.set_name(model_ini.domain_name)                                # already created in model_ini.py
    domain.set_datadir(model_ini.results_dir)                             # already created and/or checked in model_ini.py
    domain.set_minimum_storable_height(model_ini.minimum_storable_height) # Min water depths stored in sww
    domain.set_maximum_allowed_speed (model_ini.maximum_allowed_speed )   # Max particle speed allowed in calcs.
    domain.set_minimum_allowed_height(model_ini.minimum_allowed_height)   # Min depth used in calcs.
    domain.set_flow_algorithm(model_ini.flow_algorithm)                   # refer anuga config.py for options
    domain.set_starttime(model_ini.starttime)                             # Set the simulation starttime if not at t=0

    # Display some basic domain info to confirm what has been done
    if model_ini.model_verbose  :
        print model_ini.prefix+' ----- Domain Statistics ','\n',  domain.statistics()

    print model_ini.prefix+' >>>> Completed basic domain construction at t= %.2f hours'  %(float(time.time()-model_ini.t0)/3600.0)


    ################################################################################
    #                                                                              #
    #                ASSIGN THE SINGLE DOMAIN ELEVATIONS                           #
    #                                                                              #
    ################################################################################

    #------------------------------------------------------------------------------
    # This section assigns the initial (t=0) conditions for the domain
    # Domain initial conditions can be fixed, read from a file, or from a function
    # Note also that anuga permits almost all initial conditions to subsequently vary
    # with time, so while the following conditions apply to the whole domain at t=0
    # all can be modified within the evolve time loop. In this way scour, deposition,
    # blockage, changes in roughness etc with time can be explicitly included in a
    # simulation
    #------------------------------------------------------------------------------

    print model_ini.prefix+' >>>> Assigning initial conditions to %s ' %(model_ini.domain_name)

    # Elevation data may be assigned to the domain from a csv or (NetCDF) pts file.

    # For a smaller csv that will be read in repeatedly, best to use geospatial_data to readin and
    # export_point_file to save the points permanently in a NetCDF *.pts file. This runs faster
    # than the block read of a csv/txt - but - is limited to about 6 million points.
    # Alternatively - by directly assigning elevation from a csv/txt file can take advantage of the block
    # read in domain.set_quantity so the elevation dataset can be virtually unlimited in size.

    # In this simulation the size of the elevation dataset is too large for conversion to a (NetCDF)
    # pts format so elevation data is read in from the raw ALS csv file.
    #


    domain.set_quantity('elevation',
                    filename  = model_ini.elevation_points_pathname,
                    use_cache = True,
                    verbose   = model_ini.model_verbose,
                    alpha     = 0.1)

    print model_ini.prefix+' >>>> Finished assigning initial elevation conditions to %s ' %(model_ini.domain_name)

    # Note that different 'scenarios' of elevation can be stamped over the base dataset by adding the
    # new elevation dataset after the base elevation has been read in (like TUFLOW). This can be repeated
    # as needed to add scenario fragments to build up the new final landform.
    # Format is: domain.add_quantity('elevation',filename= new_elevation_data)
    # As with TUFLOW need to be careful this overstamping does not create  discontinuities at
    # the boundary between datasets!



    ################################################################################
    #                                                                              #
    #                ASSIGN THE SINGLE DOMAIN MANNINGS ROUGHNESS                   #
    #                                                                              #
    ################################################################################
    # Assign uniform initial friction value to all domain mesh centroids initially
    # Note: Friction will then be re-assigned based on spatial roughness data read
    # in from a shapefile or grid
    # NOTE : In the GA2009 model friction was globally set at 0.010 within the domain poly
    #        In MRT2015 models this will be set spatially by an asc grid



    domain.set_quantity('friction',
                     filename=model_ini.mannings_grid_pathname,
                    use_cache = True,
                     location = 'centroids',
                      verbose = model_ini.model_verbose)

    print model_ini.prefix+' >>>> Finished assigning initial friction conditions to %s ' %(model_ini.domain_name)


    ################################################################################
    #                                                                              #
    #                ASSIGN THE SINGLE DOMAIN INITIAL CONDITIONS                   #
    #                                                                              #
    ################################################################################
    # Set the intial water level (stage) across the domain

    initial_water_level = model_ini.initial_water_level            # check  IWL set to boundary water level at t=starttime
    if domain.get_quantity('elevation') < initial_water_level:     # This can create ponds if land is below initial tailwater level
        domain.set_quantity('stage', initial_water_level)          # set stage at initial sea level if grd elev below IWL
    else:
        domain.set_quantity('stage', expression='elevation')       # else set initial stage at land level (dry) if grd elev above IWL

    print model_ini.prefix+' >>>> Finished assigning initial water elevation conditions to %s ' %(model_ini.domain_name)

else:
    domain = None

barrier()

################################################################################
#                                                                              #
#          IF DUNE EROSION is TRUE - Construct no scour base elevation         #
#                                                                              #
#                                                                              #
################################################################################


# setup the ns base on proc 0 and save in domain for later use in mpi domains
if model_ini.dune_erosion and myid == 0 :  # Note: domain only exists on proc 0 at this time

    print model_ini.prefix+' >>>> Seting up the no scour base dataset for the %i erosion polygons using their ns limits ' %(len(model_data.erosion_polygons))


    # get the initial model surface elevation - this also needs to be run only on proc 0 as domain only exists at this time on proc 0
    nsbase_elev_c = []
    nsbase_elev_c = domain.get_quantity('elevation').get_values(location='centroids')

    if model_ini.model_verbose :
        print model_ini.prefix+' ---- nsbase_elev_c contains ',len(nsbase_elev_c) , ' triangles'

    print model_ini.prefix+' >>>> Now build the global no scour base elevation surface for the single domain'


    # build the no scour base surface for the domain by retaining initial elev where elev < nsbase
    # and using nsbase as elev where elev > nsbase
    # get the indices of triangles for each erosion polygon so can build the global ns base elevations
    # this has to be on proc 0 as domain only exists on proc 0 at this time
    poly_indices = []
    for i in range(len(model_data.erosion_polygons)) :
        print 'there are ' + str(len(model_data.erosion_polygons)) + ' dune erosion polygons'
        polynsbase = model_data.erosion_polygons[i][1]
        polyind = (Region(domain, polygon=model_data.erosion_polygons[i][0])).indices
        nsbase_elev_c[polyind] = num.minimum(nsbase_elev_c[polyind], polynsbase)
        poly_indices.append(polyind)


    if model_ini.model_verbose :
        print model_ini.prefix+' ---- There are ', len(poly_indices), '  erosion zones used to set ns base '


    # assign the nsbase_elev_c to the domain so can be later distributed to the parallel domains
    anuga.Quantity(domain,name= 'nsbase_elevation', register = True)
    domain.set_quantity('nsbase_elevation',nsbase_elev_c, location = 'centroids')


################################################################################
#                                                                              #
#            NOW PRODUCE (DISTRIBUTE)THE PARALLEL DOMAIN                       #
#                                                                              #
################################################################################

barrier()

if myid == 0 :
    print model_ini.prefix+' >>>> Distribute the single domain onto the multiple processors '

domain = distribute(domain,verbose=model_ini.model_verbose)

domain.set_store_vertices_uniquely(False)                               # true is supposedly MPI safe?? but False is default and halves sww size

if model_ini.dune_erosion :
    domain.set_quantities_to_be_stored({'elevation': 2,'stage': 2,'xmomentum': 2,'ymomentum': 2}) # 2 =store every Ts, 1 = store at t=0
    if myid == 0 and model_ini.model_verbose :
        print model_ini.prefix+' ----- Dune erosion activated, bed elevation stored as timeseries'
else :
    domain.set_quantities_to_be_stored({'elevation': 1,'stage': 2,'xmomentum': 2,'ymomentum': 2}) # 2 store every Ts, 1 at t=0
    if myid == 0 and model_ini.model_verbose :
        print model_ini.prefix+' ----- No Dune erosion included, bed elevation constant over time'


#################################################################################
#                                                                               #
#   ASSIGN THE (TEMPORAL)  BOUNDARY CONDITIONS TO THE DISTRIBUTED DOMAIN        #
#                                                                               #
#################################################################################

print model_ini.prefix+' >>>> Processing BCs for Partition %i ' %(myid)
print model_ini.prefix+' >>>> Obtaining the seaward BC from the (%s) netcdf timeseries ' %(model_ini.event_sts_pathname)
print model_ini.prefix+' >>>> Available boundary tags are ', domain.get_boundary_tags()

# Note 'exterior' are all bdry segs not explicitly set and named earlier - typically the land in a tsunami model
# As the simulation can exceed the sts time it is necessary to set a default_boundary condition to replace the
# sts boundary when t exceeds the sts time limit.
# In this nodel it is reset to equal the initial water level - this may need review in other scenarios!!!!!!!!!!

Br = anuga.Reflective_boundary(domain)
Bt = anuga.Transmissive_stage_zero_momentum_boundary(domain)
Bf = anuga.Field_boundary(model_ini.event_sts_pathname,
                    domain           = domain,
		    mean_stage       = model_ini.initial_water_level,
                    time_thinning    = 1,
                    default_boundary = anuga.Dirichlet_boundary([0, 0, 0]),
                    boundary_polygon = model_data.bounding_polygon,
                    use_cache        = True,
                    verbose          = True)

domain.set_boundary({'top':Bt,'bottom' :Bt,'ocean':Bf,'exterior':Br})
print model_ini.prefix+' >>>> Boundary tags set for the partition on ', myid, ' are ', domain.get_boundary_tags()

print model_ini.prefix+' >>>> Completed assignment of boundary conditions to partition %i at t= %.2f hours ' %(myid,float(time.time()-model_ini.t0)/3600.0)

barrier()
if myid  == 0:
    print model_ini.prefix+' >>>> Completed distributed domain BC assignment at t= %.2f hours ' %(float(time.time()-model_ini.t0)/3600.0)


