#####
##	NOTE: In order for this script to run as both a stand-alone Python
##		geoprocessing script, and within ArcCatalog/ArcMap via "Load..."
##		in the Python interpreter window I added arcpy.MakeFeatureLayer...
##		calls after the creation of temporary datasets.	This unfortunately
##		breaks when you're running in ArcMap which by default adds the
##		temporary datasets to the current document.	To fix this I added
##		checks at each step to see if that exists in order for it to process.
#####

#####
##	Population/census data shapefile should be updated with columns
##		'ADMINID', 'ADMINPOP', 'YEARPOP' and 'ISO'
##	(NOTE: takes place of admin.shp in Benin example)

##	GRUMP shapefile is set below, e.g. af_as_lac_urban.shp, a subset of
##		GRUMP world shapefile with columns 'ISOURBID' and 'POP' used in code
##		(already in shapefile)

##	Other notes:
##		Check in initial population census shapefile for a RECNO and POP fields - if there are, delete
##	Geometry of census shp should be checked. if issues, topology may need to be corrected/adjusted (ex. VNM)
#####

#####
##  PREREQUISITES:
##      prediction_density feature in specific UTM with cell size 100 in the designated output folder.
##      prediction_density_prj feature in WGS 1984 with cell size .0008333 in the tmp of output folder.
##      This does not produce your prediction_density surface, only distributes populations onto them.

#####
##	BEGIN:	Set per-country configuration options

##	Parse main configuration file, which will set the country and root_path
##		variables:
##	This is the primary per-configuration file, parsed by each script,
##		in this folder, including both R and Python scripts.  There should
##		be fewer per-script configuration items now as these will control
##		the country-name and root path (path to the folder containing the RF
##		/src, /data, /output, etc. folders.

##	NOTE: Make sure that you don't use the R assignment operator <- as
##		this file is also parsed by Python, so use = instead:


##	Configure the country abbreviation and name:
country = "BWA"

##	The version of the scripts used to produce the mapping products, and
##		which will match the "_v" portion of the filename outputs:
rf_version = "2c"

##	This should be set to the folder *containing* the "RF" folder structure,
##		but make sure that you use forward slashes instead of back slashes or
##		double back slashes:
root_path = "D:/Research/Population/Data/RF/"



##	Round to whole population counts?
round_counts = False

##	END:	Set per-country configuration options
#####



#####
##	NOTICE:  In practice nothing below this line should need to be regularly
##		edited unless very specific details need to be changed about the
##		modeling process (e.g. speciyfing an extent or something detailed
##		about the Python processing).
#####



#####
##	BEGIN:	Set general configuration options

data_path = root_path + "/data/" + country + "/"
output_path = root_path + "/output/" + country + "/"
tmp_path = output_path + "tmp/"


##	Input projection is configured for each country based on the most
##		appropriate output for distance/area considerations... Change to the local UTM
##	Final projection for population maps:
final_prj = "GEOGCS['GCS_WGS_1984',DATUM['D_WGS_1984',SPHEROID['WGS_1984',6378137.0,298.257223563]],PRIMEM['Greenwich',0.0],UNIT['Degree',0.0174532925199433]]"

#final_prj = "PROJCS['WGS_1984_UTM_Zone_45N',GEOGCS['GCS_WGS_1984',DATUM['D_WGS_1984',SPHEROID['WGS_1984',6378137.0,298.257223563]],PRIMEM['Greenwich',0.0],UNIT['Degree',0.0174532925199433]],PROJECTION['Transverse_Mercator'],PARAMETER['False_Easting',500000.0],PARAMETER['False_Northing',0.0],PARAMETER['Central_Meridian',87.0],PARAMETER['Scale_Factor',0.9996],PARAMETER['Latitude_Of_Origin',0.0],UNIT['Meter',1.0],AUTHORITY['EPSG',32645]]"


##	END:	Set general configuration options
#####



#####
##	BEGIN: Import packages and set GeoProcessing environment

##	Import the Spatial Analyst extension for slope and raster algebra
##		functionality:
import os
import glob
import json
#from pprint import pprint

import arcpy
arcpy.CheckOutExtension("Spatial")
from arcpy.sa import *


##	Configure arcpy:

# Set the compression environment to LZ77:
arcpy.env.compression = "LZ77"

##	Should we overwrite any already existing derived data:
overwrite = True
arcpy.env.overwriteOutput = overwrite

##	Set the output coordinate system and extent to that of our buffer:
arcpy.env.outputCoordinateSystem = final_prj


##	END: Import packages and set GeoProcessing environment
#####



#####
##	BEGIN: Define utility functions

##	Create a directory function to check for and create data directories if they don't exist:
def ensure_dir(d):
	if not os.path.isdir(d):
		os.makedirs(d)
	return d


##	END: Define utility functions
#####



#####
##	BEGIN: Data pre-processing for needed datasets


##	Get a list of all existing folders in the country's data folder:
dataset_folders = os.walk(data_path).next()[1]


##	Check to see if an alternative census folder exists in the folder list
##		and use it for processing our maps.
##	TODO:  Eventually this needs to be pulled from the Metadata.R file
##		instead of a hard coded directory presence.
if ("! New Census" in dataset_folders):
	census_folder = "! New Census"
else:
	census_folder = ""


##	Datasets:

##	Estimated population density from randomForest to be used as density
##		weightings (gridx) from the old model, a file that should be in the
##		directory (we changed formats between version numbers so check for both:
dataset_path = glob.glob( output_path + "predict_density.*" )
if dataset_path:
	popdensity_weighting = dataset_path[0]
else:
	print("ERROR:  No \"predict_density\" TIF or IMG found in the output folder!  You first need to run the 1.3 R script!")
	exit()

##	Population data
if census_folder != "":
	census_path = data_path + census_folder + "/"
else:
	census_path = data_path + "Census/"


##	Instead of hardcoding it above we'll just pull in the only shapefile
##		that should be in the directory:
dataset_name = os.path.basename( glob.glob( census_path + "*.shp" )[0] )
adminpop = census_path + dataset_name


##	NOTE: To properly set the extent we need to query the description of
##		the shapefile by hand, for some reason setting the extent directly
##		to the name of a dataset doesn't work...
desc = arcpy.Describe(adminpop)
arcpy.env.extent = desc.Extent


##	Set the workspace for data creation and processing:
arcpy.env.workspace = ensure_dir(tmp_path)


##	BEGIN: Data pre-processing for needed datasets
#####



#####
##	BEGIN: Population map creation

print("PREPROCESS: Finalized density weights...")

##	Project our predicted density weighting file back to GCS1984:
in_path = popdensity_weighting
out_name = "predict_density_prj.tif"

out_path = output_path + "tmp/" + out_name

data_desc = arcpy.Describe(in_path)
##input_prj = data_desc.SpatialReference.exportToString()

## You may have to manually create the projected raster object in order to continue if it chokes here.


tmp_extent = arcpy.env.extent
##	There may be an inconsistency in the way that 10.3 handles extents
##		during reprojection.  To fix this the following line may be
##		uncommented:
#arcpy.env.extent = ""

##arcpy.ProjectRaster_management(in_path,out_path,final_prj,"BILINEAR","0.0008333","#","#",input_prj)


arcpy.env.extent = tmp_extent

##	Set our snapping environment to our new, reprojected and
##		majority resampled land cover file:

#popdensity_weighting_final = in_path
#arcpy.env.snapRaster = in_path
popdensity_weighting_final = out_path
arcpy.env.snapRaster = out_path



##	Create a temporary copy of our census data with an appropriate
##		population field to sum and distribute:
if (arcpy.Exists("admin") == False): arcpy.MakeFeatureLayer_management(adminpop, "admin")
arcpy.FeatureClassToFeatureClass_conversion('admin', arcpy.env.workspace, 'admin_Union')
if (arcpy.Exists("admin_Union") == False): arcpy.MakeFeatureLayer_management("/admin_Union.shp", "admin_Union")

if (len(arcpy.ListFields('admin_Union', "POP")) == 0):
	arcpy.AddField_management('admin_Union', 'POP', 'Double')
arcpy.CalculateField_management('admin_Union', 'POP', ' [ADMINPOP] ', '')



##	Population redistribution procedure using population from "admin_Union.shp"
##		converted to a raster and redistributed according to our weights:
print("POPMAP: Begin creating population redistribution...")
print("PPP: Using census file - " + dataset_name)


##	Read the first row of the attribute table from our census data to
##		determine the year of our census data:
rows = arcpy.SearchCursor(adminpop)
row = rows.next()
census_year = row.YEARPOP

## create an alternate popdensity_weighting_final feature to feed into the script, that allows for the inclusion of
## census unit counts that may not have an included settlement area.

## Declare the root path: MUST BE CHANGED FOR EACH RUN

arcpy.CreateFileGDB_management(output_path+"tmp", "Prediction_Density_2")
path = output_path + "tmp\\Prediction_Density_2.gdb"
#"C:\\Users\\preed\\Desktop\\Thesis\\Comparisons\\Scripts\\01.Opt.1\\Prediction_Density.gdb"

gridp = arcpy.FeatureToRaster_conversion('admin_Union', 'ADMINPOP', 'gridp', '0.0008333')
# may need to be specified as gridp.tif depending on where your default directory is. (gdb vs folder)

ZonalSt = arcpy.sa.ZonalStatisticsAsTable('admin_Union', 'ADMINID', Raster(popdensity_weighting_final),(path + "\\ZonalSt"), "DATA", "SUM")




## INSTRUCTIONS:
    ## The following steps are run  in ArcMap's python visualizer.
    ## There is some visual Q/A that has to occur at the end of each section.

    ## If you finish STEP 1: and all records are selected from admin_union, then run the commented code at the end of the section. popmap_a is your product!
    ## If you finish STEP 1: and not every record is selected from admin_union, then proceed to STEP 2.

    ## If in STEP 2: there are no remaining gaps in popmap_ab, then run the last few lines of commented code. popmap_ab is your product!
    ## If in STEP 2: there are still some areas in popmap_ab that were not populated, then proceed to run STEP 3: to fill in the gaps.

    ## If STEP 3: was run, then your output is popmap_final.

    ## There has been a continuous save error when running through the python visualizer, so outputs are not automatically saved in the correct format.
    ## For the time, manual export of the final product is the way to go.


## SELECTION 1: SETTLED CENSUS UNITS

arcpy.AddJoin_management('admin_Union', 'ADMINID', ZonalSt, 'ADMINID', 'KEEP_ALL')
arcpy.SelectLayerByAttribute_management('admin_Union', 'NEW_SELECTION', '\"ZonalSt.Sum\" >0')

Predict_Density_a = arcpy.sa.ExtractByMask(Raster(popdensity_weighting_final),'admin_Union')
Predict_Density_a.save(path + "\\Predict_Density_a")

gridy_a =arcpy.sa.ZonalStatistics('admin_Union', 'admin_Union.ADMINID', Predict_Density_a, 'SUM', 'DATA')
gridy_a.save(path + "\\gridy_a")

popmap_a = gridp * Predict_Density_a / gridy_a
popmap_a.save(path + "\\popmap_a")


## If working with the traditional WorldPop model, go ahead and run these last few steps:
##      OR
## If all records are selected in the admin_uion feature, then run this commented section and be done:
##arcpy.RemoveJoin_management('admin_Union')
##arcpy.SelectLayerByAttribute_management('admin_Union', 'CLEAR_SELECTION')
##arcpy.sa.ZonalStatisticsAsTable('admin_Union', 'ADMINID', popmap_a,(path + "\\PopmapSUM_a"), "DATA", "SUM")


## SELECTION 2: NON-SETTLED CENSUS UNITS

arcpy.SelectLayerByAttribute_management('admin_Union', 'SWITCH_SELECTION')

EvenDensity = arcpy.sa.Con(Raster(popdensity_weighting_final)>= 0,1,0)
EvenDensity.save(path + "\\EvenDensity")

Predict_Density_b = arcpy.sa.ExtractByMask(EvenDensity,'admin_Union')
Predict_Density_b.save(path + "\\Predict_Density_b")

gridy_b =arcpy.sa.ZonalStatistics('admin_Union', 'admin_Union.ADMINID', Predict_Density_b, 'SUM', 'DATA')
gridy_b.save(path + "\\gridy_b")


popmap_b = gridp * Predict_Density_b / gridy_b
popmap_b.save(path + "\\popmap_b")


popmap_ab = arcpy.sa.Con(IsNull(popmap_a),popmap_b,popmap_a)

## If in STEP 2: there are no remaining gaps in popmap_ab, then run the last few lines of commented code. popmap_ab is your product!
arcpy.RemoveJoin_management('admin_Union')
arcpy.SelectLayerByAttribute_management('admin_Union', 'CLEAR_SELECTION')
arcpy.sa.ZonalStatisticsAsTable('admin_Union', 'ADMINID', popmap_ab,(path + "\\PopmapSUM_ab"), "DATA", "SUM")



## SELECTION 3: REMAINING EXCLUDED AREA
## Sometimes there are census units that have settlement pixels right on the border between units,
## which for whatever reason are included in the initial settlement unit selection, but do not produce a surface.
## The prediction_density_a and gridy_a parameter are both set to 0, and are therefore excluded from the final.

## We can create a predict_density_c based on only the excluded area = 1. Clip according to even density.
predict_density_c =arcpy.sa.Con(IsNull(popmap_ab),1)
predict_density_c =arcpy.sa.ExtractByMask(predict_density_c,EvenDensity)
predict_density_c.save(path + "\\predict_density_c")

##      Run zonal statistics as table on predict_density_c by admin_Union.AdminID as SUM.
##      Calculate zonal stats for the EvenDensity layer, so that we can produce a difference
##      Join back to admin_Union by adminID.
arcpy.RemoveJoin_management('admin_Union')
arcpy.SelectLayerByAttribute_management('admin_Union', 'CLEAR_SELECTION')

PreDen_C_Tbl = arcpy.sa.ZonalStatisticsAsTable('admin_Union', 'ADMINID', predict_density_c,(path + "\\PreDen_C_Tbl"), "DATA", "SUM")
EvenDensity_tbl = arcpy.sa.ZonalStatisticsAsTable('admin_Union', 'ADMINID', EvenDensity,(path + "\\EvenDensity_tbl"), "DATA", "SUM")

##     Add a difference field to admin_Union
arcpy.AddField_management('admin_Union', "Diff", "DOUBLE")

arcpy.AddJoin_management('admin_Union', 'ADMINID', PreDen_C_Tbl, 'ADMINID', 'KEEP_ALL')
arcpy.AddJoin_management('admin_Union', 'admin_Union.ADMINID', EvenDensity_tbl, 'ADMINID', 'KEEP_ALL')

##     Calculate the difference field. And select those areas that have a similarity if 50 or greater.
arcpy.SelectLayerByAttribute_management('admin_Union', 'NEW_SELECTION', '\"PreDen_C_Tbl.Sum\" >=1')
arcpy.CalculateField_management('admin_Union', "Diff", "([PreDen_C_Tbl.SUM] / [EvenDensity_tbl.SUM])*100")
arcpy.SelectLayerByAttribute_management('admin_Union', 'NEW_SELECTION', '\"admin_Union.Diff\" >=50')

##     Remove existing joins.
arcpy.RemoveJoin_management('admin_Union')

## produce gridy_c using the selection.
gridy_c =arcpy.sa.ZonalStatistics('admin_Union', 'ADMINID', predict_density_c, 'SUM', 'DATA')
gridy_c.save(path + "\\gridy_c")

##     Produce popmap_c
popmap_c = gridp * predict_density_c / gridy_c
popmap_c.save(path + "\\popmap_c")



## Concatonate popmap_ab and popmap_c
popmap_abc = arcpy.sa.Con(IsNull(popmap_ab),popmap_c,popmap_ab)

## Fill in the gaps with values of 0 based on predict_density_c
gaps = arcpy.sa.Con(predict_density_c==1, 0, predict_density_c)
popmap_final = arcpy.sa.Con(IsNull(popmap_abc), gaps, popmap_abc)

arcpy.sa.ZonalStatisticsAsTable('admin_Union', 'ADMINID', popmap_final,(path + "\\PopmapSUM_Final_2"), "DATA", "SUM")
arcpy.sa.ZonalStatisticsAsTable('admin_Union', 'ADMINID', popmap_abc,(path + "\\PopmapSUM_abc"), "DATA", "SUM")

popmap_final.save(path + "\\popmap_final")




## Dubious save function.
outPath = output_path + country + "_ppp_v" + rf_version + "_" + str(census_year) + ".tif"
popmap_final.save(outPath)








print("POPMAP: Completed!")


print("COMPLETED:  Succesffully created population map outputs!")

##	BEGIN: Population map creation
#####
