# script gets all Australian BoM weather station observations
# ... and applies an interpolated temperature to all GNAF points in a 100m grid

# TODO:
#  1. remove temperature biases due to altitude differences
#       a. Add SRTM altitudes to GNAF
#       b. Add interpolated altitude from weather stations to GNAF
#       c. adjust where the difference is > 100m
#  2. generate temps outside the weather station network to catch the ~3,100 GNAF points outside the interpolated area
#

import geopandas
import io
import json
import logging
# import matplotlib.pyplot as plt
import multiprocessing
import numpy
import os
import pandas
import psycopg
import requests
import scipy.interpolate
import sqlalchemy
import struct
import urllib.request
import zipfile

# from bs4 import BeautifulSoup
from datetime import datetime
from osgeo import gdal

# where to save the files
output_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "data")

# states to include (note: no "ACT" or "OT" state, Antarctica is part of TAS in BoM observations)
state_list = [{"state": "NSW", "product": "IDN60801"},
             {"state": "NT", "product": "IDD60801"},
             {"state": "QLD", "product": "IDQ60801"},
             {"state": "SA", "product": "IDS60801"},
             {"state": "TAS", "product": "IDT60801"},
             {"state": "VIC", "product": "IDV60801"},
             {"state": "WA", "product": "IDW60801"},
             {"state": "ANT", "product": "IDT60801"}]

# urls for each state's weather observations
# http://www.bom.gov.au/fwo/IDW60801/IDW60801.95214.json
# base_url = "http://www.bom.gov.au/{0}/observations/{0}all.shtml"
base_url = "http://www.bom.gov.au/fwo/{0}/{0}.{1}.json"

# postgres connect strings
pg_connect_string = "dbname='geo' host='localhost' port='5432' user='postgres' password='password'"
sql_alchemy_engine_string = "postgresql+psycopg://postgres:password@localhost/geo"


def main():
    start_time = datetime.now()

    # connect to Postgres
    try:
        pg_conn = psycopg.connect(pg_connect_string)
        pg_conn.autocommit = True
        pg_cur = pg_conn.cursor()
    except psycopg.Error:
        logger.fatal("Unable to connect to database\nACTION: Check your Postgres parameters and/or database security")
        return False

    # create df of state info
    state_df = pandas.DataFrame(state_list)

    # download weather stations
    station_list = get_weather_stations()
    logger.info("Downloaded {:,} weather stations : {}".format(len(station_list), datetime.now() - start_time))

    # create dataframe of weather stations
    temp_station_df = pandas.DataFrame(station_list)

    # merge with state info to get state product name
    station_df = temp_station_df.merge(state_df, on="state")

    # download weather obs JSON files into a list
    obs_list = get_weather_observations(station_df.to_dict("records"))
    logger.info("Downloaded {:,} latest observations : {}".format(len(obs_list), datetime.now() - start_time))
    start_time = datetime.now()

    # create dataframe of weather obs
    obs_df = pandas.DataFrame(obs_list).drop_duplicates()

    # merge data and add points to dataframe
    df = (obs_df.merge(station_df, on="wmo")
          .drop(["lat", "lon"], axis=1)
          )
    # gdf = geopandas.GeoDataFrame(df, geometry=geopandas.points_from_xy(df.longitude, df.latitude), crs="EPSG:4283")

    # select rows from the last hour with air temps
    air_temp_df = df[(df["utc_time_diff"] < 3600.0) & (df["air_temp"].notna())
                     & (df["longitude"] > 112.0) & (df["longitude"] < 162.0)
                     & (df["latitude"] > -45.0) & (df["latitude"] < -8.0)]

    # # testing - get histogram of observation time
    # air_temp_df.hist("utc_time")
    # plt.savefig(os.path.join(output_path, "histogram.png"), dpi=300, facecolor="w", pad_inches=0.0, metadata=None)

    # export dataframe to PostGIS
    export_dataframe(pg_cur, air_temp_df, "testing", "weather_stations", "replace")
    logger.info("Exported weather station dataframe to PostGIS: {}".format(datetime.now() - start_time))
    start_time = datetime.now()

    # # save to disk for debugging
    # air_temp_df.to_feather(os.path.join(output_path "temp_df.ipc"))

    # # load from disk if debugging
    # temp_df = pandas.read_feather(os.path.join(output_path "temp_df.ipc"))

    # extract lat, long and air temp as arrays
    x = air_temp_df["longitude"].to_numpy()
    y = air_temp_df["latitude"].to_numpy()
    z = air_temp_df["air_temp"].to_numpy()
    h = air_temp_df["altitude"].to_numpy()

    logger.info("Filtered observations dataframe with weather station coordinates : {} rows : {}"
                .format(len(air_temp_df.index), datetime.now() - start_time))
    start_time = datetime.now()

    # open SRTM 3 second DEM of Australia (ESRI Binary Grid format)
    dem_file_name = "/Users/hugh.saalmans/Downloads/3secSRTM_DEM/DEM_ESRI_GRID_16bit_Integer/dem3s_int/hdr.adf"
    dem_dataset = gdal.Open(dem_file_name, gdal.GA_ReadOnly)
    dem_geotransform = dem_dataset.GetGeoTransform()

    # get DEM origin point and pixel size to create numpy arrays from
    dem_num_x, dem_num_y = dem_dataset.RasterXSize, dem_dataset.RasterYSize
    dem_origin_x, dem_origin_y = dem_geotransform[0], dem_geotransform[3]
    dem_origin_delta_x, dem_origin_delta_y = dem_geotransform[1], dem_geotransform[5]

    # select GNAF coordinates - group by 3 decimal places to create a ~100m grid of addresses
    # sql = """SELECT latitude::numeric(5,3) as latitude, longitude::numeric(6,3) as longitude, count(*) as address_count
    #          FROM gnaf_202408.address_principals
    #          GROUP BY latitude::numeric(5,3), longitude::numeric(6,3)"""
    sql = """SELECT * FROM testing.gnaf_points_with_pop_and_height"""
    gnaf_df = pandas.read_sql_query(sql, pg_conn)

    # save to feather file for future use (GNAF only changes once every 3 months)
    gnaf_df.to_feather(os.path.join(output_path, "gnaf.ipc"))

    # # load from feather file
    # gnaf_df = pandas.read_feather(os.path.join(output_path, "gnaf.ipc"))

    gnaf_x = gnaf_df["longitude"].to_numpy()
    gnaf_y = gnaf_df["latitude"].to_numpy()
    gnaf_counts = gnaf_df["count"].to_numpy()
    gnaf_dem_elevation = gnaf_df["elevation"].to_numpy()

    logger.info("Loaded {:,} GNAF points : {}".format(len(gnaf_df.index), datetime.now() - start_time))
    start_time = datetime.now()

    # # interpolate temperatures for GNAF coordinates
    gnaf_points = numpy.array((gnaf_x.flatten(), gnaf_y.flatten())).T
    gnaf_temps = scipy.interpolate.griddata((x, y), z, gnaf_points, method="linear")
    gnaf_weather_elevation = scipy.interpolate.griddata((x, y), h, gnaf_points, method="linear")

    # create results dataframe
    temperature_df = pandas.DataFrame({"latitude": gnaf_y, "longitude": gnaf_x,
                                       "count": gnaf_counts, "dem_elevation": gnaf_dem_elevation,
                                       "weather_elevation": gnaf_weather_elevation, "air_temp": gnaf_temps})

    # add temperatures adjusted for altitude differences between GNAF point and nearby weather stations
    temperature_df["adjusted_temp"] = temperature_df["air_temp"] + \
                                      (temperature_df["weather_elevation"] - temperature_df["dem_elevation"]) / 150.0

    # print(temperature_df)

    # get count of rows with a temperature
    row_count = len(temperature_df[temperature_df["air_temp"].notna()].index)

    logger.info("Got {:,} interpolated temperatures and elevations for GNAF points : {}"
                .format(row_count, datetime.now() - start_time))
    start_time = datetime.now()

    # # plot a map of gnaf points by temperature
    # temperature_df.plot.scatter("longitude", "latitude", c="air_temp", colormap="jet")
    # plt.axis("off")
    # plt.savefig(os.path.join(output_path, "interpolated.png"), dpi=300, facecolor="w", pad_inches=0.0, metadata=None)
    #
    # logger.info("Plotted points to PNG file : {}".format(datetime.now() - start_time))
    # start_time = datetime.now()

    # export dataframe to PostGIS
    export_dataframe(pg_cur, temperature_df, "testing", "gnaf_temperature", "replace")
    logger.info("Exported GNAF temperature dataframe to PostGIS: {}".format(datetime.now() - start_time))
    # start_time = datetime.now()

    return True


def export_dataframe(pg_cur, df, schema_name, table_name, export_mode):
    # create geodataframe
    gdf = geopandas.GeoDataFrame(df, geometry=geopandas.points_from_xy(df.longitude, df.latitude), crs="EPSG:4283")

    # export to GeoPackage
    # gdf.to_file(os.path.join(output_path, "{}.gpkg".format(table_name)), driver="GPKG")
    #
    # logger.info("Exported points to GeoPackage : {}".format(datetime.now() - start_time))
    # start_time = datetime.now()

    # export to PostGIS
    engine = sqlalchemy.create_engine(sql_alchemy_engine_string)
    gdf.to_postgis(table_name, engine, schema=schema_name, if_exists=export_mode)

    pg_cur.execute("ANALYSE {}.{}".format(schema_name, table_name))
    # pg_cur.execute("ALTER TABLE testing.weather_stations ADD CONSTRAINT weather_stations_pkey PRIMARY KEY (wmo)"
    #                .format(schema_name, table_name))
    pg_cur.execute("ALTER TABLE {0}.{1} CLUSTER ON idx_{1}_geometry".format(schema_name, table_name))
    pg_cur.execute("ALTER TABLE {}.{} RENAME COLUMN geometry TO geom".format(schema_name, table_name))


# download weather observations for all 'wmo' stations
def get_weather_observations(station_list):
    obs_urls = list()
    obs_list = list()

    for station in station_list:
        obs_url = base_url.format(station["product"], station["wmo"])
        obs_urls.append(obs_url)

        # with open(os.path.join(output_path, "weather_observations_urls.txt"), "w", newline="") as output_file:
        #     output_file.write("\n".join(obs_urls))

    # download each obs file using multiprocessing
    pool = multiprocessing.Pool(processes=1)
    results = pool.imap_unordered(run_multiprocessing, obs_urls)

    pool.close()
    pool.join()

    for result in list(results):
        if result.get("error") is not None:
            logger.warning("\t- Failed to parse {}".format(result["error"]))
        else:
            obs_list.append(result)

    return obs_list


def get_weather_stations():
    # get weather stations - obs have poor coordinates
    response = urllib.request.urlopen("ftp://ftp.bom.gov.au/anon2/home/ncc/metadata/sitelists/stations.zip")
    data = io.BytesIO(response.read())
    station_file = zipfile.ZipFile(data, "r", zipfile.ZIP_DEFLATED).read("stations.txt").decode("utf-8")
    stations = station_file.split("\r\n")

    station_list = list()

    # split fixed width file and get the fields we want
    field_widths = (-8, -6, 41, -8, -7, 9, 10, -15, 4, 11, -9, 7)  # negative widths represent ignored fields
    format_string = " ".join("{}{}".format(abs(fw), "x" if fw < 0 else "s") for fw in field_widths)
    field_struct = struct.Struct(format_string)
    parser = field_struct.unpack_from
    # print("fmtstring: {!r}, recsize: {} chars".format(fmtstring, fieldstruct.size))

    # skip first 5 rows (lazy coding!)
    stations.pop(0)
    stations.pop(0)
    stations.pop(0)
    stations.pop(0)
    stations.pop(0)

    # add each station to a list of dictionaries
    for station in stations:
        if len(station) > 128:
            fields = parser(bytes(station, "utf-8"))

            # convert to list
            field_list = list()

            for field in fields:
                field_list.append(field.decode("utf-8").lstrip().rstrip())

            if field_list[5] != "..":
                station_dict = dict()
                station_dict["name"] = field_list[0]
                station_dict["latitude"] = float(field_list[1])
                station_dict["longitude"] = float(field_list[2])
                station_dict["state"] = field_list[3]
                if field_list[4] != "..":
                    station_dict["altitude"] = float(field_list[4])
                station_dict["wmo"] = int(field_list[5])

                station_list.append(station_dict)

    return station_list


def run_multiprocessing(url):
    # file_path = os.path.join(output_path, "obs", url.split("/")[-1])

    # try:
    headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36'}
    obs_text = requests.get(url, headers=headers).text
    # obs_text = requests.get(url).text

    # with open(file_path, "w", newline="") as output_file:
    #     output_file.write(obs_text)

    # print(obs_text)

    obs_json = json.loads(obs_text)
    obs_list = obs_json["observations"]["data"]

    try:
        # default is an error for when there are no observations
        result = dict()
        result["error"] = "{} : No observations".format(url)

        for obs in obs_list:
            if obs["sort_order"] == 0:
                result = obs

                # add utc time
                obs["utc_time"] = datetime.strptime(obs["aifstime_utc"], "%Y%m%d%H%M%S")
                obs["utc_time_diff"] = (datetime.utcnow() - obs["utc_time"]).total_seconds()

    except Exception as ex:
        result = dict()
        result["error"] = "{} : {}".format(url, ex)
        # print(result)

    return result


if __name__ == "__main__":
    full_start_time = datetime.now()

    logger = logging.getLogger()

    # set logger
    log_file = os.path.abspath(__file__).replace(".py", ".log")
    logging.basicConfig(filename=log_file, level=logging.DEBUG, format="%(asctime)s %(message)s",
                        datefmt="%m/%d/%Y %I:%M:%S %p")

    # setup logger to write to screen as well as writing to log file
    # define a Handler which writes INFO messages or higher to the sys.stderr
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    # set a format which is simpler for console use
    formatter = logging.Formatter("%(name)-12s: %(levelname)-8s %(message)s")
    # tell the handler to use this format
    console.setFormatter(formatter)
    # add the handler to the root logger
    logging.getLogger("").addHandler(console)

    logger.info("")
    logger.info("Start weather obs download")
    # geoscape.check_python_version(logger)

    if main():
        logger.info("Finished successfully! : {}".format(datetime.now() - full_start_time))
    else:
        logger.fatal("Something bad happened!")

    logger.info("")
    logger.info("-------------------------------------------------------------------------------")
