"""
-----------------------------------------------------------------model_data.py
ANUGA TSUNAMI MODELLING TEMPLATE (Data Assembly)
This script extracts/converts  client model-scenario-event data
for use in the subsequent Anuga mesh-domain build process described 
in model_build.py.


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 THE REQUIRED PYTHON MODULES                        #
#                                                                               #
#################################################################################

# Standard python modules
import os
import sys
from os import sep
from os.path import join


# Related ANUGA modules
import anuga   

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

# Other python libraries
import shapefile              # the pyshp library - need to have shapefile.py accessible

# Simulation specific imports
import model_ini              # Provides the simulation specific variables/files

barrier()
if myid == 0:
    print model_ini.prefix+' >>>> Imported required modules for the model_data assembly process '

#################################################################################
#                                                                               #
#                    ASSEMBLE THE BACKDROP COASTAL POLYGON                      #
#                                                                               #
#################################################################################

if myid == 0:    
    print model_ini.prefix+' >>>> Commencing data assembly for this simulation '

# load the coasline poly as backdrop for check plots if check true
if model_ini.model_checks :
    sf = shapefile.Reader(model_ini.coastal_poly_pathname)            # power up pyshp reader and get shape geom
    shapes =sf.shapes()
    if myid == 0:
        print  model_ini.prefix+' >>>> Reading in the backdrop coastline for check plots if checks true'
        print  model_ini.prefix+' >>>> %i polygons in coastal shapefile ' %(len(shapes))
    coastal_polygons =[]
    for i in range(len(shapes)) :                                     # loop for each poly in shp file
        coastal_polygon_points = []
        coastal_polygon_points = shapes[i].points                     # list of of xy coords for poly less closing point
        coastal_polygons.append(coastal_polygon_points)               # append poly points list to polygons list

    if model_ini.model_verbose and myid == 0 :
        for i in range(len(shapes)) :
            print model_ini.prefix+' ----- dump of first 2 coords of each coastal poly (i,[x,y]) - ', i,coastal_polygons[i][0:2]
 
    if myid == 0:   
        poly_window = anuga.plot_polygons(coastal_polygons,  style = None, figname=join(model_ini.checks_dir,'Tasmanian Coastline'), label='Tasmanian Coastline')

  
################################################################################
#                                                                              #
#           ASSEMBLE THE BOUNDING POLYGON & TAGGED BOUNDARY STRINGS            #
#                                                                              #
################################################################################

#------------------------------------------------------------------------------
# Define bounding polygon and default res to be used later in model_build.py to 
# build the mesh and boundary strings with their tags
#
# OUTPUTS OF THIS CODE BLOCK ARE;
# bounding_polygon[]  -- The bounding polygon within which the mesh 
#                        will be created (from the bounding_poly shapefile)
# boundary_strings{}  -- The tagged boundary strings dictionary
#
#------------------------------------------------------------------------------

##### Read in bounding poly  and check plot if checks set to true
##### Note that the GA sts BC data contains boundary points in a counterclockwise direction
##### around the bounding poly starting from the most southerly ocean point.
##### As a consequence any attempt to create the bounding poly as a conventional GIS (shp) 
##### file fails as the shape file automatically reorders in a clockwise manner around the bounding poly
##### as that is the required order for an outer ring/boundary in GIS. (counter clockwise denotes
##### an inner ring/hole poly). Must therefor read the bounding poly from a csv points list directly
##### into Anuga so that the non GIS standard of the sts oredering can be maintained.

barrier()
bounding_polygon=[]
bounding_polygon = anuga.read_polygon(model_ini.bounding_poly_pathname)

if myid == 0:
    print  model_ini.prefix+' >>>> Reading in the bounding polygon from ', model_ini.bounding_poly_pathname  # single poly

    if model_ini.model_verbose :
        print model_ini.prefix+' ----- There were %i unique points read in for the bounding polygon   ' %(len(bounding_polygon))   
        for i in range(len(bounding_polygon)) :
            print model_ini.prefix+' ----- dump of the coords of the bounding polygon ([point][x,y]) - ', i,bounding_polygon[i][0:2]

    if model_ini.model_checks :
        print model_ini.prefix+ ' >>>> Check plotting bounding polygon'  # plot it out in checks directory
        plot_polys = []
        for i in range(len(coastal_polygons)) :
            plot_polys.append(coastal_polygons[i])
        plot_polys.append(bounding_polygon)
        poly_window = anuga.plot_polygons(plot_polys,  figname=join(model_ini.checks_dir,'bounding_polygon'), label='Bounding Polygon')

    print model_ini.prefix+ ' >>>> Area of bounding polygon in Km2', anuga.polygon_area(bounding_polygon)/1000000.0
    def_triangle_side = ((model_ini.default_res)* 2.0)**0.5
    print model_ini.prefix+ ' >>>> Default resolution within bounding polygon is ',model_ini.default_res,' m2 ie ',def_triangle_side,' m triangle side'


##### Assign the required tagged boundary segments (verticies are those of bounding_poly)

if myid == 0:
    print model_ini.prefix+ ' >>>> Assembling the Anuga tagged boundary strings' 

# Note ANUGA will auto assign 'exterior' to any remaining segments of the bounding_poly
# vertex numbers/segnents are derived from display of the bounding poly vertex numbers
# Three explict and one implicit ('exterior') boundary strings are used to reflect the different BCs in the model
# ANUGA will auto assign 'exterior' to any  segments of the bounding_poly not assigned here.
# BCs will be assigned to these boundary types in model_build.py
# Note there are 32 nodes in the NETAS boundary file with 16 sts nodes so 15 ocean_segments
#      the sts bdry pts are indexed from the south to the north requiring a non (GIS)
#      standard (counter clockwise) numbering sequence around the Anuga bounding_polygon
#      The domain bdry pts are exactly those of the STS, with top-land-bottom added
#
#                   top (2)
# viz          17*----<----*15 (model bdry node no)
#                |         |
#          Land  V         ^
#          (13)  |         | Ocean (15 segments)
#              30*----->---*0
#                  bottom (2)


top_segments     = []
bottom_segments  = []
ocean_segments   = []
top_segments     = [15,16]
bottom_segments  = [30,31]
ocean_segments   = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14]
boundary_strings = {'top'    : top_segments,
                    'bottom' : bottom_segments,
                    'ocean'  : ocean_segments}

if myid == 0:
    print model_ini.prefix+ ' >>>> %i tagged boundary strings assigned by user to model' %(len(boundary_strings))

if model_ini.model_verbose :
    for key in boundary_strings :
        if myid == 0:
            print model_ini.prefix+ ' ----- tagged boundary ', key, boundary_strings[key]



#################################################################################
#                                                                               #
#                   ASSEMBLE INTERIOR RESOLUTION POLYGONS                       #
#                                                                               #
#################################################################################

#-------------------------------------------------------------------------------
# Define Interior region polygons for later use in model_run.py to define zones 
# of different mesh resolution
#
# OUTPUTS OF THIS CODE BLOCK ARE;
# interior_res_regions[]-- List of interior res regions containing [res_poly, res]
#                         to be used in model_run to create the mesh
#-------------------------------------------------------------------------------
# Note coarse_polygon not needed as  set as default for bounding_polgon

# Interior_res_regions and assoc resolutions list are created here for use in model_build.py
# order so coarsest interior res is first and finest last as each overstamps when assigned
# note res polys must be simple, single part with vertices spaced out further than res.
# A later res polys can be used when it lies  fully inside a prior loaded  poly but poly 
# boundaries canot touch or intersect. There should also be sufficient space between the 
# inner and outer polys for decently shaped triangles to form.
# 
# the foillowing copde  can probably be compacted into some sort of loop so can have variable 
# numbers of int polys ???????????ER

barrier()
interior_res_polygons = []
interior_res_regions = []

if myid == 0:
    print model_ini.prefix+ ' >>>> Assembling any inner region polygons and setting associated resolutions'
    print model_ini.prefix+' >>>>> ',"about to load internal res1 medium ocean10to30m polygons ", #model_ini.int1_res_poly_pathname

int1_polygons = []
sf = shapefile.Reader(model_ini.int1_res_poly_pathname)   # power up pyshp reader
shapes =sf.shapes()
region_data = []
# read basic geometry of shapefile
for i in range(len(shapes)) :                             # loop for each poly in shp file
   int1_polygon_points = []
   int1_polygon_points = shapes[i].points                # list of polys xy coords less closing point
   int1_polygons.append(int1_polygon_points)             # append the points for that poly
   interior_res_polygons.append(int1_polygon_points)     # keep polygons seperate for check plots
   region_data=[int1_polygon_points,model_ini.int1_res]  # Add the int res data for this poly
   interior_res_regions.append(region_data)              # Append these region_data polys to int_res_regions
if model_ini.model_verbose and myid == 0:
   print model_ini.prefix+' -----  ', model_ini.int1_res_poly_pathname,' read in containing ', len(shapes), ' polygons'
 
if myid == 0:
    print model_ini.prefix+' >>>>> ',"about to load internal res2 fine coastal strip polygons ", model_ini.int2_res_poly_pathname

int2_polygons = []
sf = shapefile.Reader(model_ini.int2_res_poly_pathname)   
shapes =sf.shapes()
region_data = []
for i in range(len(shapes)) :
   int2_polygon_points = []
   int2_polygon_points = shapes[i].points
   int2_polygons.append(int2_polygon_points)
   interior_res_polygons.append(int2_polygon_points)
   region_data = [int2_polygon_points, model_ini.int2_res]
   interior_res_regions.append(region_data)
if model_ini.model_verbose and myid == 0:
    print model_ini.prefix+' -----  ', model_ini.int2_res_poly_pathname,' read in containing ', len(shapes), ' polygons'

if myid == 0:
    print model_ini.prefix+' >>>>> ',"about to load internal res3 coarse coastal land above 10m polygons",model_ini.int3_res_poly_pathname
int3_polygons = []
sf = shapefile.Reader(model_ini.int3_res_poly_pathname)   
shapes =sf.shapes()
region_data = []
for i in range(len(shapes)) :
    int3_polygon_points = []
    int3_polygon_points = shapes[i].points
    int3_polygons.append(int3_polygon_points)
    interior_res_polygons.append(int3_polygon_points)
    region_data = [int3_polygon_points, model_ini.int3_res]
    interior_res_regions.append(region_data)
if model_ini.model_verbose and myid == 0:
    print model_ini.prefix+' -----  ', model_ini.int3_res_poly_pathname,' read in containing ', len(shapes), ' polygons'
if myid == 0:
    print model_ini.prefix+' >>>>> ',"about to load internal res4  vfine plot area polygons ", model_ini.int4_res_poly_pathname
int4_polygons = []
sf = shapefile.Reader(model_ini.int4_res_poly_pathname)   
shapes =sf.shapes()
region_data = []
for i in range(len(shapes)) :
    int4_polygon_points = []
    int4_polygon_points = shapes[i].points
    int4_polygons.append(int4_polygon_points)
    interior_res_polygons.append(int4_polygon_points)
    region_data = [int4_polygon_points, model_ini.int4_res]
    interior_res_regions.append(region_data)
if model_ini.model_verbose and myid == 0:
    print model_ini.prefix+' -----  ', model_ini.int4_res_poly_pathname,' read in containing ', len(shapes), ' polygons'
if myid == 0:  
   print model_ini.prefix+' >>>>> ',"about to load internal res5  xfine study area polygons ", model_ini.int5_res_poly_pathname
int5_polygons = []
sf = shapefile.Reader(model_ini.int5_res_poly_pathname)   
shapes =sf.shapes()
region_data = []
for i in range(len(shapes)) :
    int5_polygon_points = []
    int5_polygon_points = shapes[i].points
    int5_polygons.append(int5_polygon_points)
    interior_res_polygons.append(int5_polygon_points)
    region_data = [int5_polygon_points, model_ini.int5_res]
    interior_res_regions.append(region_data)
if model_ini.model_verbose and myid == 0:
    print model_ini.prefix+' -----  ', model_ini.int5_res_poly_pathname,' read in containing ', len(shapes), ' polygons'



# Do some final sanity checks on the poly data read in
if (len(interior_res_regions) != len(interior_res_polygons)) :  
    if myid == 0:  
        print model_ini.prefix+ ' ERROR - Length of int_respolys and int_res_regions not the same ', len(interior_res_polygons), len(interior_res_regions)
        sys.exit(0)

if myid == 0:
    print model_ini.prefix+ ' >>>> %i interior resolution polygons created and processed' %(len(interior_res_regions))

if model_ini.model_verbose and myid == 0: 
    for i in range(len(interior_res_regions)) :
        print model_ini.prefix+' ----- dump of internal resolution m2 for poly ', i, interior_res_regions[i][1]
    for i in range(len(interior_res_regions)) :
        print model_ini.prefix+' ----- dump of coords for poly ', i, interior_res_polygons[i][0],'...',interior_res_polygons[i][-1]

 
if model_ini.model_checks :
    if myid == 0:
        print model_ini.prefix+ ' >>>> Check plotting interior resolution polygons'  # plot it out in data directory
        plot_polys = []
        plot_polys= interior_res_polygons                                            
#       plot_polys.append(bounding_polygon)            # seems to create a problem if appended???ER??
        poly_window = anuga.plot_polygons(plot_polys,  figname=join(model_ini.checks_dir,'Interior Resolution Polygons'), label='Interior Resolution Polygons')



#################################################################################
#                                                                               #
#                       ASSEMBLE THE ELEVATION DATA                             #
#                                                                               #
#################################################################################

#------------------------------------------------------------------------------
# Assign reference to file with elevation data for later use in model_build.py
# the following assumes this  elevation data file is stored with the other
# datafiles in the model/scenario/event/data  directory
#
# OUTPUTS OF THIS CODE BLOCK ARE;
# elevation_pathname -- full path name of file containing ALS data to be used to set mesh elevations
#------------------------------------------------------------------------------

# Note must be either 1 or 2
#elev_option = 1 # assign elevations to the mesh using the fast netcdf elevation dataset
elev_option = 2 # assign elevations to the mesh using the original (slower) yxz/txt dataset
if elev_option != 1 and elev_option != 2 :
    print model_ini.prefix+ ' ERROR - elev_option must be 1 or 2 - correct and rerun - was ', elev_option
    sys.exit(0)


# NOTES #####
# OPTION 1 (Preferred but limited to smaller ALS datasets)
#           This option converts the raw xyz ALS data into NETCDF (pts) format
#           The pts file is then used to assign mesh elevations in model_build.py. 
#           While interpolating and assigning elevations to the mesh from the pts file, 
#           this option loads all data into memory (no blocking) so is limited in
#           regard to the number of points it can assign to the mesh.
#           If there are more than about 6 million points in the ALS dataset
#           this option will fail with a memory over-run error.
#           In such a case the elevation data should be left in csv (xyz/txt) format.
#
# OPTION 2 (Leave elevation data in original (csv)(xyz) format and block read when elevating mesh)
#

barrier()
if elev_option == 1 :   #   Convert xyz/txt data into NETCDF pts format to speed later use in setting mesh elevations

    if myid == 0:
        print model_ini.prefix+' >>>> Elevation x,y,z scatterpoint dataset %s being used to create NETCDF pts file' %(model_ini.elevation_points_pathname)

    G = anuga.Geospatial_data(file_name      = model_ini.elevation_points_pathname, 
	                      max_read_lines = 500, 
			      load_file_now  = True, 
			      verbose        = model_ini.model_verbose)
    G.export_points_file(model_ini.elevation_points_pathname[:-4] + '.pts')  # export geospatial data to  NETCDF pts formatted file
    elevation_pathname = model_ini.elevation_points_pathname[:-4] + '.pts'   # switch elevation_filename to now point to the NETCDF pts file
    if myid == 0:
        print model_ini.prefix+ ' >>>> Mesh elevation data will be assigned from %s' %(elevation_pathname)
	
else :                 # OPTION 2 (Leave elevation data in original (csv)(xyz) format and block read when elevating mesh in model_build.py)
    elevation_pathname = model_ini.elevation_points_pathname
    if myid == 0:
        print model_ini.prefix+ ' >>>> Mesh elevation data will be assigned from %s' %(elevation_pathname)

#################################################################################
#                                                                               #
#             ASSEMBLE SURFACE ROUGHNESS POLYGONS AND ASSOCIATED DATA           #
#                                                                               #
#################################################################################


# This block would typically read in the roughness data from a shapefile but for 
# more complex roughness dtaasets the roughness will be directly assigned to the 
# mesh in model_build from an asc grid file

barrier()
if myid == 0:
    print model_ini.prefix+ ' >>>> The surface roughness data will be assembled from ' , model_ini.mannings_grid_pathname


#################################################################################
#                                                                               #
#            ASSEMBLE THE DUNE EROSION ZONE POLYGONS (IF INCLUDED)              #
#                                                                               #
#################################################################################


# This block conditionally reads in the dune erosion zone polygons from a shapefile  
# if the simulation is flaged in model_ini to include dune erosion 
# Inclded if model_ini.dune_erosion = True

barrier()
if model_ini.dune_erosion :
    if myid == 0 :
        print model_ini.prefix+ ' >>>> Dune erosion is to be simulated within the polygons contained in ' , model_ini.erosion_zones_pathname
    
    erosion_polygons = []
    sf = shapefile.Reader(model_ini.erosion_zones_pathname)   
    shapes  = sf.shapes()
    records = sf.records()
    for i in range(len(shapes)) :                       # setup to read multiple erosion polys 
        erosion_polygon_points = []
        erosion_polygon_points = shapes[i].points
        erosion_polygon_name   = records[i][0]          # shp first attrib field
        erosion_polygon_nsbase = records[i][3]          # shp fourth attrib field
        erosion_polygons.append([erosion_polygon_points,erosion_polygon_nsbase,erosion_polygon_name])  
        if myid == 0 :
            print model_ini.prefix+ ' >>>> Dune erosion polys read in for ' , erosion_polygon_name, ' with no scour base at RL ', erosion_polygon_nsbase
    
    if (model_ini.model_verbose and myid ==0) :
         print model_ini.prefix+' -----  ', model_ini.erosion_zones_pathname,' read in containing ', len(shapes), ' shapes',len(erosion_polygons), ' polygons'
         for i in range(len(shapes)) :
             print model_ini.prefix+' -----   dump of name, coords and nsbase for erosion_polygon',  i, erosion_polygons[i][2],erosion_polygons[i][0],erosion_polygons[i][1] 
    
    if len(erosion_polygons) != len(shapes) :
         print model_ini.prefix+' ERROR - len(erosion_polygons) %i does not equal no of polys %i read in from shpfile ' %(len(erosion_polygons),len(shapes))  
         sys.exit(0)
    
    if (model_ini.model_checks and myid == 0) :
        print model_ini.prefix+ ' >>>> Check plotting erosion zone polygons'  # plot zones out in checks directory
        plot_polys = []
        for i in range(len(coastal_polygons)) :
            plot_polys.append(coastal_polygons[i])
        for i in range(len(erosion_polygons)) :
            plot_polys.append(erosion_polygons[i][0])  
        poly_window = anuga.plot_polygons(plot_polys,  figname=join(model_ini.checks_dir,'Dune Erosion Polygons'), label='Dune Erosion Polygons')
else :
    if myid == 0 :
        print model_ini.prefix+ ' >>>> No Dune erosion zones in this scenario'  
          









