#!/bin/bash

# Script to geoscrub-enable a postgres database, developed in the

# context of BIEN3 geoscrubbing but with intentions of generality.

#

# Basic workflow:

#  1. Import table dumps obtained from geonames.org

#  2. Add additional custom mapping info, including:

#     a. Madagascar provinces used in GADM2 but not geonames

#     b. extra country name mappings/fixes specific to BIEN3 data

#     c. extra state/prov name mappings/fixes specific to BIEN3 data

#     d. custom state/prov name mappings binding GADM2 to geonames

#

# Jim Regetz

# NCEAS

# Created Nov 2012

#

# refactored and reorganized by

# Paul Sarando

# iPlant Collaborative

# Updated Oct 2013

DB_NAME="geoscrub"

DB_USER="bien"

DB_HOST_OPT=""

SCRIPT_DIR="$(dirname $0)"

DATADIR="${HOME}/geonames"

GEONAMES_DUMP_URL="http://download.geonames.org/export/dump"

ALL_COUNTRIES_ZIP=allCountries.zip

ALL_COUNTRIES_TXT=allCountries.txt

ALTERNATE_NAMES_ZIP=alternateNames.zip

ALTERNATE_NAMES_TXT=alternateNames.txt

COUNTRY_INFO_TXT=countryInfo.txt

HIERARCHY_ZIP=hierarchy.zip

HIERARCHY_TXT=hierarchy.txt

function usage {

    echo "Usage: $0 [OPTIONS]" >&2

    echo "Valid Options:" >&2

    echo "-d, --dbname=DBNAME      database name psql commands will connect to" >&2

    echo "-h, --host=HOSTNAME      database server host or socket directory" >&2

    echo "-U, --username=USERNAME  database user name" >&2

    echo ""

    echo "Input Data Options:" >&2

    echo "-m, --geonames-data      Geonames.org data directory (default: ${HOME}/geonames)" >&2

    echo "                         Delete this directory, or the geonames.org data in it," >&2

    echo "                         to re-download the data." >&2

    exit 1;

}

while [[ $# -gt 0  ]]; do

    case "$1" in

        -\? | --help)

            usage

            ;;

        -h)

            if [[ -z $2  ]];  then

                echo "Option $1 requires an argument." >&2

                usage

            fi

            DB_HOST_OPT="-h $2"

            shift 2

            ;;

        --host=*)

            DB_HOST_OPT="-h ${1#*=}"

            shift

            ;;

        -U)

            if [[ -z $2  ]];  then

                echo "Option $1 requires an argument." >&2

                usage

            fi

            DB_USER="$2"

            shift 2

            ;;

        --username=*)

            DB_USER="${1#*=}"

            shift

            ;;

        -d)

            if [[ -z $2  ]];  then

                echo "Option $1 requires an argument." >&2

                usage

            fi

            DB_NAME="$2"

            shift 2

            ;;

        --dbname=*)

            DB_NAME="${1#*=}"

            shift

            ;;

        -m)

            if [[ -z $2  ]];  then

                echo "Option $1 requires an argument." >&2

                usage

            fi

            DATADIR="$2"

            shift 2

            ;;

        --geonames-data=*)

            DATADIR="${1#*=}"

            shift

            ;;

        *)

            echo "Invalid option: $1" >&2

            usage

            ;;

    esac

done

function run_sql_cmd {

    local SQL_CMD=$1

    psql -e -U "$DB_USER" $DB_HOST_OPT -d "$DB_NAME" -c "$SQL_CMD"

    if [[ $? != 0 ]]; then

        echo "Error while executing SQL command."

        exit 1

    fi

}

function run_sql_script {

    local SCRIPT=$1

    psql -e --set ON_ERROR_STOP=1 -U "$DB_USER" $DB_HOST_OPT -d "$DB_NAME" < "$SCRIPT"

    if [[ $? != 0 ]]; then

        echo "Error while executing SQL script ${SCRIPT}"

        exit 1

    fi

}

echo "Updating geonames.org tables..."

# Check for data downloaded into geonames under the current directory.

if [[ ! -d "${DATADIR}" ]]; then

    echo "making directory ${DATADIR}"

    mkdir -p "${DATADIR}"

    if [[ $? != 0 ]]; then

        echo "Could not create directory ${DATADIR}"

        exit 1

    fi

fi

run_sql_script "${SCRIPT_DIR}/truncate.geonames.sql"

#

# load geonames.org geonames

#

if [[ ! -r "${DATADIR}/${ALL_COUNTRIES_TXT}" ]]; then

    pushd "${DATADIR}"

    wget "${GEONAMES_DUMP_URL}/${ALL_COUNTRIES_ZIP}"

    if [[ $? != 0 ]]; then

        echo "Could not download ${ALL_COUNTRIES_ZIP}"

        exit 1

    fi

    unzip "$ALL_COUNTRIES_ZIP"

    rm -f "$ALL_COUNTRIES_ZIP"

    popd

fi

# import data

echo "Copying geonames from ${DATADIR}/${ALL_COUNTRIES_TXT}"

run_sql_cmd "\COPY geonames FROM '${DATADIR}/${ALL_COUNTRIES_TXT}' DELIMITER E'\t' NULL ''"

#

# load geonames.org alternate names

#

if [[ ! -r "${DATADIR}/${ALTERNATE_NAMES_TXT}" ]]; then

    pushd "${DATADIR}"

    wget "${GEONAMES_DUMP_URL}/${ALTERNATE_NAMES_ZIP}"

    if [[ $? != 0 ]]; then

        echo "Could not download ${ALTERNATE_NAMES_ZIP}"

        exit 1

    fi

    unzip "$ALTERNATE_NAMES_ZIP"

    rm -f "$ALTERNATE_NAMES_ZIP"

    popd

fi

# import data

echo "Copying alternateNames from ${DATADIR}/${ALTERNATE_NAMES_TXT}"

run_sql_cmd "\COPY alternateNames FROM '${DATADIR}/${ALTERNATE_NAMES_TXT}' DELIMITER E'\t' NULL ''"

run_sql_script "${SCRIPT_DIR}/update.alternateNames.sql"

#

# load geonames.org country info

#

if [[ ! -r "${DATADIR}/${COUNTRY_INFO_TXT}" ]]; then

    pushd "${DATADIR}"

    wget "${GEONAMES_DUMP_URL}/${COUNTRY_INFO_TXT}"

    if [[ $? != 0 ]]; then

        echo "Could not download ${COUNTRY_INFO_TXT}"

        exit 1

    fi

    popd

fi

# import geonames country table

echo "Copying countries from ${DATADIR}/${COUNTRY_INFO_TXT}"

grep -ve '^#' "${DATADIR}/${COUNTRY_INFO_TXT}" | \

    psql -U "$DB_USER" $DB_HOST_OPT -d "$DB_NAME" -c \

    "\COPY countries FROM STDIN DELIMITER E'\t' NULL ''"

if [[ $? != 0 ]]; then

    echo "Could not copy countries table data."

    exit 1

fi

run_sql_script "${SCRIPT_DIR}/update.alt_country.sql"

run_sql_script "${SCRIPT_DIR}/update.alt_stateprovince.sql"

#

# load geonames.org hierarchy info

#

if [[ ! -r "${DATADIR}/${HIERARCHY_TXT}" ]]; then

    pushd "${DATADIR}"

    wget "${GEONAMES_DUMP_URL}/${HIERARCHY_ZIP}"

    if [[ $? != 0 ]]; then

        echo "Could not download ${HIERARCHY_ZIP}"

        exit 1

    fi

    unzip "$HIERARCHY_ZIP"

    rm -f "$HIERARCHY_ZIP"

    popd

fi

# import geonames country table

echo "Copying hierarchy from ${DATADIR}/${HIERARCHY_TXT}"

run_sql_cmd "\COPY hierarchy FROM '${DATADIR}/${HIERARCHY_TXT}' DELIMITER E'\t' NULL ''"

run_sql_script "${SCRIPT_DIR}/update.hierarchy.sql"

# augment geonames with some additional places we need, and update

# hierarchy table accordingly

echo "Updating geonames and hierarchy tables..."

run_sql_script "${SCRIPT_DIR}/update.geonames.sql"

echo "Update geonames.org tables successfully completed."