"""
This scripts generates ascii grids of selected parameters in nominated areas
extracted from the specified sww results file
Input: sww results file from model_run.py
       area(s) for grids to be created specified in this script
Outputs: ascii grids for each specified area and specified variable(s)
         stored in the 'results/spatial' folder for the .sww file

Note:
If producing a grid for the enitre extent of the SETAS model, cellsize
should be greater than 50m. If producing grids for local inundation areas
cell size should be be greater than the mesh size in that area(ie ~20m min)
"""

##### Import the required system files
from os import remove
import sys
from os.path import join, sep

###### Import anuga
import anuga

# Other python libraries
import shapefile              # the pyshp library - need to have shapefile.py accessible
import shutil                 # high level file copy etc functions

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


# sww file of interest located in the directory of interest is

sww_directory   = model_ini.results_dir+sep
sww_pathname    = join(model_ini.results_dir,model_ini.domain_name ) +'.sww'


#Now set the timestep at which you want the raster generated.
#Either set the actual timestep required or use 'None' to indicate that
#you want the maximum values in the raster over all timesteps
timestep = None   # plot max


# 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()
    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


# Set the special areas of interest to be gridded.
# Areas identified below are specified by window coordinates as follows

areas = ['South', 'Central', 'North']

# ASCII export grid window for South Communities
xminSouth = 567000
xmaxSouth = 634000
yminSouth = 5282000
ymaxSouth = 5361000
cellsizeSouth   = 20        # cell size m - needs vfine res in area in model mesh to make sense
if model_ini.model_checks :
    # create a check plot of the  area
    print model_ini.prefix+ ' >>>> Check plotting export grid window polygon for' , areas[0] # plot it out in checks directory
    plot_polys   = []
    plot_window  = []
    plot_window  = [[xminSouth,yminSouth],[xminSouth,ymaxSouth],[ xmaxSouth,ymaxSouth],[xmaxSouth,yminSouth]]
    plot_polys   = coastal_polygons
    plot_polys.append(plot_window)
    poly_window = anuga.plot_polygons(plot_polys,  style = None, figname=join(model_ini.checks_dir, 'Hobart Communities'), label='Hobart Communities')

# ASCII export grid window for Central Communities
xminCentral = 593000
xmaxCentral = 628000
yminCentral = 5358000
ymaxCentral = 5436000
cellsizeCentral   = 20        # cell size m - needs vfine res in area in model mesh to make sense
if model_ini.model_checks :
    # create a check plot of the  area
    print model_ini.prefix+ ' >>>> Check plotting export grid window polygon for ' , areas[1] # plot it out in checks directory
    plot_polys   = []
    plot_window  = []
    plot_window  = [[xminCentral,yminCentral],[xminCentral,ymaxCentral],[ xmaxCentral,ymaxCentral],[xmaxCentral,yminCentral]]
    plot_polys   = coastal_polygons
    plot_polys.append(plot_window)
    poly_window = anuga.plot_polygons(plot_polys,  style = None, figname=join(model_ini.checks_dir, 'Northern Communities'), label='Nothern Communities')


# ASCII export grid window for North Communities
xminNorth = 597000
xmaxNorth = 649000
yminNorth = 5433000
ymaxNorth = 5502000
cellsizeNorth   = 20        # cell size m - needs vfine res in area in model mesh to make sense
if model_ini.model_checks :
    # create a check plot of the  area
    print model_ini.prefix+ ' >>>> Check plotting export grid window polygon for ' , areas[2]  # plot it out in checks directory
    plot_polys   = []
    plot_window  = []
    plot_window  = [[xminNorth,yminNorth],[xminNorth,ymaxNorth],[ xmaxNorth,ymaxNorth],[xmaxNorth,yminNorth]]
    plot_polys   = coastal_polygons
    plot_polys.append(plot_window)
    poly_window = anuga.plot_polygons(plot_polys,  style = None, figname=join(model_ini.checks_dir, 'Southern Communities'), label='Southern Communities')



###### Set or more key strings from var_equations below

vars = ['stage', 'depth', 'speed']

################################################################################
#                                                                              #
# Define allowed variable names and associated equations to generate values.   #
# This should not normally change.                                             #
#                                                                              #
################################################################################

var_equations = {'stage':     'stage',
                 'momentum':  '(xmomentum**2 + ymomentum**2)**0.5',
                 'depth':     'stage-elevation',
                 'speed':     '(xmomentum**2 + ymomentum**2)**0.5/(stage-elevation+1.e-6)',
                 'elevation': 'elevation' }


################################################################################
#                                                                              #
# Start  sww2dem looping through variables, area  ( as set by user)            #
#                                                                              #
################################################################################

for which_var in vars:
    if which_var not in var_equations:    # loop over variables
        print model_ini.prefix+ ' ERROR > Unrecognized variable name: %s' % which_var
        break

    for which_area in areas:              # loop over areas
        try:
            easting_min = eval('xmin%s' % which_area)
            easting_max = eval('xmax%s' % which_area)
            northing_min = eval('ymin%s' % which_area)
            northing_max = eval('ymax%s' % which_area)
            cellsize     = eval('cellsize%s' % which_area)
        except AttributeError:
            print model_ini.prefix+ ' ERROR > Unrecognized area name: %s' % which_area
            break

        outname = model_ini.spatial_dir+sep+'_' + which_area + '_' + which_var+'.asc'
        quantityname = var_equations[which_var]

        print model_ini.prefix+ ' >>>> start sww2dem using sww_file=%s' % sww_pathname
        print model_ini.prefix+ ' >>>> output grid file name = %s' % outname

        anuga.sww2dem(name_in  = sww_pathname,
                      name_out = outname,
                  quantity     = quantityname,
                  cellsize     = cellsize,
                  easting_min  = easting_min,
                  easting_max  = easting_max,
                  northing_min = northing_min,
                  northing_max = northing_max,
                  reduction    = max,
                  verbose      = True)




        ########################################################################
        # Note the following code to rewrite the *.prj file is a crude fix     #
        # to rewrite the prj file to reflect the desired projection as the     #
        # prj file created by sww2dem  has no ability to reflect the           #
        # projection the model was built in - Oct 2015/ER.                     #
        ########################################################################

        proj_pathname =  model_ini.spatial_dir+sep+'_' + which_area + '_' + which_var+'.prj'
        print model_ini.prefix + ' >>>>> Reseting the projection file for  ', proj_pathname
        remove(proj_pathname)
        shutil.copy(model_ini.spatial_dir+sep+'projectionSETAS.prj', proj_pathname)

################ END OF TEMPORARY CODE BLOCK ##########################


