"""
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 = ['Hobart', 'NW', 'South']  

# ASCII export grid window for Hobart Communities
xminHobart = 520000
xmaxHobart = 545000
yminHobart = 5230000
ymaxHobart = 5260000
cellsizeHobart   = 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  = [[xminHobart,yminHobart],[xminHobart,ymaxHobart],[ xmaxHobart,ymaxHobart],[xmaxHobart,yminHobart]]
    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 North Communities
xminNW = 548000
xmaxNW = 561000
yminNW = 5250000
ymaxNW = 5258000
cellsizeNW   = 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  = [[xminNW,yminNW],[xminNW,ymaxNW],[ xmaxNW,ymaxNW],[xmaxNW,yminNW]]
    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 South Communities
xminSouth = 523000
xmaxSouth = 533300
yminSouth = 5197000
ymaxSouth = 5214000   
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[2]  # 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, '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 ##########################