# -*- coding: utf-8 -*-

"""
/***************************************************************************
 OpenHLZ
                                 A QGIS plugin
 Open-source HLZ Identification Plugin
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2023-08-13
        copyright            : (C) 2026 by CustomCartographix
        email                : john@customcartographix.com
        version              : 2.0.1
        version date         : 2026-07-04
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/

 /***************************************************************************
 *                                                                         *
 *   This script contains the necessary functions for the OpenHLZ          *
 *   processing algorithm to work.                                         *
 *                                                                         *
 ***************************************************************************/
"""

# Import statements
import os

from qgis.core import (QgsVectorLayer,
                       QgsVectorFileWriter,
                       QgsNetworkAccessManager,
                       QgsProject,
                       QgsProcessingException,
                       QgsProcessingLayerPostProcessorInterface)
from qgis import processing
from osgeo import gdal
from pyproj import Proj, transform
from PyQt5.QtNetwork import QNetworkRequest, QNetworkReply
from PyQt5.QtCore import QUrl
from json import load


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _blockingDownload(url, output_filename, description, feedback=None):
    """
    Download the contents of a URL to a file using QgsNetworkAccessManager.

    Surfaces network failures as a QgsProcessingException (instead of silently
    writing an error page into the output file) and closes the file handle
    deterministically via a context manager.
    """
    request = QNetworkRequest(QUrl(url))
    network_access_manager = QgsNetworkAccessManager()
    response = network_access_manager.blockingGet(request)

    if response.error() != QNetworkReply.NoError:
        raise QgsProcessingException(
            'Network error while downloading ' + description + ': '
            + response.errorString() + ' (' + url + ')')

    content = response.content()
    with open(output_filename, 'wb') as output_file:
        output_file.write(content)

    if feedback is not None:
        feedback.pushInfo('Downloaded ' + description)

    return output_filename


def writeToShapefile(layer, output_path, only_selected=False):
    """
    Write a vector layer to an ESRI Shapefile using the non-deprecated
    writeAsVectorFormatV3 API.
    """
    options = QgsVectorFileWriter.SaveVectorOptions()
    options.driverName = 'ESRI Shapefile'
    options.fileEncoding = 'utf-8'
    options.onlySelectedFeatures = only_selected

    error = QgsVectorFileWriter.writeAsVectorFormatV3(
        layer,
        output_path,
        QgsProject.instance().transformContext(),
        options)

    if error[0] != QgsVectorFileWriter.NoError:
        raise QgsProcessingException(
            'Failed to write vector file ' + output_path + ': ' + str(error[1]))

    return output_path


# ---------------------------------------------------------------------------
# Downloads
# ---------------------------------------------------------------------------

def downloadCopernicusGlo30Data(aoi_layer, scratch_folder, project_coordinate_system_name, feedback=None):
    """
    Function to download Copernicus GLO-30 data based on input AOI using OpenTopography API.

    Returns output filename
    """
    # Get bounding box
    extent_var = aoi_layer.extent()
    extent_array = [extent_var.xMinimum(), extent_var.yMinimum(), extent_var.xMaximum(), extent_var.yMaximum()]
    input_projection = Proj(project_coordinate_system_name)
    output_projection = Proj('epsg:4326')
    min_x_proj, min_y_proj = transform(input_projection, output_projection, extent_array[0], extent_array[1])
    max_x_proj, max_y_proj = transform(input_projection, output_projection, extent_array[2], extent_array[3])
    url_bbox = ('south=' + str(min_x_proj) + '&north=' + str(max_x_proj) + '&west=' + str(min_y_proj) +
                '&east=' + str(max_y_proj))

    # Download Data
    download_url = ('https://portal.opentopography.org/API/globaldem?demtype=COP30&' + url_bbox +
                    '&outputFormat=GTiff&API_Key=8422c95f5e3490ecca69be478123007c')
    output_filename = os.path.join(scratch_folder, 'srtm_dem.tif')
    _blockingDownload(download_url, output_filename, 'Copernicus GLO-30 elevation data', feedback)

    # Return output filename
    return output_filename


def getJson(aoi_layer, scratch_folder, project_coordinate_system_name, feedback=None):
    """
    Function to get JSON file return for either DEM or LULC

    Returns JSON files for 1m, 10m, and 1/9 arc-second data
    """

    # Get bounding box
    extent_var = aoi_layer.extent()
    extent_array = [extent_var.xMinimum(), extent_var.yMinimum(), extent_var.xMaximum(), extent_var.yMaximum()]
    input_projection = Proj(project_coordinate_system_name)
    output_projection = Proj('epsg:4326')
    min_x_proj, min_y_proj = transform(input_projection, output_projection, extent_array[0], extent_array[1])
    max_x_proj, max_y_proj = transform(input_projection, output_projection, extent_array[2], extent_array[3])
    url_bbox = str(min_y_proj) + ',' + str(min_x_proj) + ',' + str(max_y_proj) + ',' + str(max_x_proj)

    # Get json return for 1m DEM
    url = 'https://tnmaccess.nationalmap.gov/api/v1/products?bbox=' + url_bbox + \
          '&datasets=Digital+Elevation+Model+(DEM)+1+meter&prodFormats=GeoTIFF&outputFormat=json'
    output_json_filename_1m = os.path.join(scratch_folder, 'json_1m.json')
    _blockingDownload(url, output_json_filename_1m, '1m DEM index', feedback)
    with open(output_json_filename_1m) as temp_file:
        json_1m = load(temp_file)

    # Get json return for 1/3 arc-second (10m) data
    url = 'https://tnmaccess.nationalmap.gov/api/v1/products?bbox=' + url_bbox + \
          ('&datasets=National%20Elevation%20Dataset%20('
           'NED)%201%20arc-second&prodFormats=GeoTIFF&outputFormat=json')
    output_json_filename_10m = os.path.join(scratch_folder, 'json_10m.json')
    _blockingDownload(url, output_json_filename_10m, '10m DEM index', feedback)
    with open(output_json_filename_10m) as temp_file_10m:
        json_10m = load(temp_file_10m)

    # Get json return for 1/9 arc-second (1m) DEM
    url = 'https://tnmaccess.nationalmap.gov/api/v1/products?bbox=' + url_bbox + \
          ('&datasets=National%20Elevation%20Dataset%20(NED)%201/9%20arc-second&prodFormats=GeoTIFF&'
           'outputFormat=json')
    output_json_filename_19as = os.path.join(scratch_folder, 'json_19as.json')
    _blockingDownload(url, output_json_filename_19as, '1/9 arc-second DEM index', feedback)
    with open(output_json_filename_19as) as temp_file:
        json_19as = load(temp_file)

    return json_1m, json_10m, json_19as


def downloadDem(i, json_file, scratch_folder, input_data, feedback=None):
    """
    Function to download DEM data based on input location. Will first try 1m, then will try 10m if no 1m

    Returns filename of saved DEM raster
    """

    # Download 1m data
    if input_data == '1m':
        temp_dict = json_file['items'][i]
        output_filename = os.path.join(scratch_folder, 'dem_1m_' + str(i) + '.tif')
        return _blockingDownload(temp_dict['downloadURL'], output_filename, '1m DEM tile', feedback)

    # Download 1/9m data
    if input_data == '19m':
        temp_dict = json_file['items'][i]
        output_filename = os.path.join(scratch_folder, 'dem_19m_' + str(i) + '.tif')
        return _blockingDownload(temp_dict['downloadURL'], output_filename, '1/9 arc-second DEM tile', feedback)

    # Else download 10m data
    else:
        temp_dict = json_file['items'][i]
        output_filename = os.path.join(scratch_folder, 'dem_10m_' + str(i) + '.tif')
        return _blockingDownload(temp_dict['downloadURL'], output_filename, '10m DEM tile', feedback)


def downloadLandCoverRaster(i, utm_zones, scratch_folder, feedback=None):
    """
    Function to download land cover data based on current location

    Returns filename of saved land cover raster
    """

    # Download data based on whichever UTM zone the aoi is within
    # Get UTM zone and create output filename
    zone = utm_zones[i]
    zone_id = str(zone[0]) + zone[1]
    output_filename = os.path.join(scratch_folder, zone_id + '.tif')

    # Download data from Living Atlas and write to scratch folder (2025 collection)
    url = ("https://lulctimeseries.blob.core.windows.net/lulctimeseriesv003/lc2025/" + zone_id +
           "_20250101-20251231.tif")
    _blockingDownload(url, output_filename, 'land cover tile ' + zone_id, feedback)

    # Return filename of downloaded land cover raster
    return output_filename


# ---------------------------------------------------------------------------
# Mosaic / clip
# ---------------------------------------------------------------------------

def _writeAoiCutline(aoi_layer, scratch_folder):
    """
    Write the AOI polygon out to a shapefile so it can be used as a gdal.Warp
    cutline (gdal.Warp needs a data-source path, not an in-memory layer).
    """
    cutline_path = os.path.join(scratch_folder, 'aoi_cutline.shp')
    writeToShapefile(aoi_layer, cutline_path, only_selected=False)
    return cutline_path


def _makeWarpProgressCallback(feedback):
    """
    Build a GDAL progress callback that forwards progress to the Processing
    feedback panel and honours cancellation (returning 0 aborts the warp).
    """
    def _callback(complete, message, unknown):
        if feedback is not None:
            if feedback.isCanceled():
                return 0
            feedback.setProgress(complete * 100)
        return 1
    return _callback


def mosaicAndClipRasters(input_raster_filenames, aoi_layer, scratch_folder, output_raster, feedback=None):
    """
    Reproject, mosaic, and clip the input rasters to the AOI in a single
    streaming pass with osgeo.gdal.Warp.

    Warp reads block-by-block, so it never allocates the full output in RAM the
    way gdal:merge did (which blew up on large AOIs). It also natively accepts
    inputs in different UTM CRSs, applies the AOI cutline, preserves the source
    colour table, and writes an LZW-compressed, tiled GeoTIFF. Nearest-neighbour
    resampling is used so categorical land-cover class values are never blended
    into nonsense classes.
    """
    if not input_raster_filenames:
        raise QgsProcessingException('No input rasters were provided to mosaic and clip.')

    cutline_path = _writeAoiCutline(aoi_layer, scratch_folder)
    target_crs = aoi_layer.crs().authid()

    warp_options = gdal.WarpOptions(
        format='GTiff',
        dstSRS=target_crs,
        cutlineDSName=cutline_path,
        cropToCutline=True,
        resampleAlg='near',
        multithread=True,
        creationOptions=['COMPRESS=LZW', 'TILED=YES'],
        callback=_makeWarpProgressCallback(feedback))

    result = gdal.Warp(output_raster, input_raster_filenames, options=warp_options)
    if result is None:
        raise QgsProcessingException('GDAL Warp failed to mosaic and clip the input rasters.')

    # Flush and release the dataset / GDAL block cache
    result = None


# ---------------------------------------------------------------------------
# HLS / HLZ processing
# ---------------------------------------------------------------------------

def generateHlsRaster(output, projected_dem, projected_land_cover, slope_caution, slope_limit,
                      model_coordinate_system, scratch_folder, suitable_land_cover_classes='5,8,11'):
    """
    Function to generate HLS raster based on slope and LULC constraints (see README.md on GitHub for more info on
    methodology).
    """
    # Calculate slope raster
    processing.run("gdal:slope", {'INPUT': projected_dem, 'BAND': 1, 'SCALE': 1, 'AS_PERCENT': False,
                                  'COMPUTE_EDGES': False, 'ZEVENBERGEN': False, 'OPTIONS': '', 'EXTRA': '',
                                  'OUTPUT': scratch_folder + '/slope_raster.tif'})

    # Reclassify slope and LULC rasters
    processing.run("qgis:rastercalculator", {'EXPRESSION': '"slope_raster@1" < ' + str(slope_caution),
                                             'LAYERS': [scratch_folder + '/slope_raster.tif'],
                                             'CELLSIZE': 0,
                                             'EXTENT': None,
                                             'CRS': model_coordinate_system,
                                             'OUTPUT': scratch_folder + '/slope_caution.tif'})
    processing.run("qgis:rastercalculator", {'EXPRESSION': '"slope_raster@1" < ' + str(slope_limit),
                                             'LAYERS': [scratch_folder + '/slope_raster.tif'],
                                             'CELLSIZE': 0,
                                             'EXTENT': None,
                                             'CRS': model_coordinate_system,
                                             'OUTPUT': scratch_folder + '/slope_limit.tif'})
    processing.run("qgis:rastercalculator", {'EXPRESSION': '"slope_limit@1" + "slope_caution@1"',
                                             'LAYERS': [scratch_folder + '/slope_limit.tif',
                                                        scratch_folder + '/slope_caution.tif'],
                                             'CELLSIZE': 0,
                                             'EXTENT': None,
                                             'CRS': model_coordinate_system,
                                             'OUTPUT': scratch_folder + '/slope_reclass.tif'})
    os.remove(projected_dem)
    os.remove(scratch_folder + '/slope_limit.tif')
    os.remove(scratch_folder + '/slope_caution.tif')

    # Reclassify land cover based on possible input classes
    input_classes = suitable_land_cover_classes.split(',')
    input_classes_len = len(input_classes)
    expression = ''
    for i in range(input_classes_len):
        value = str(input_classes[i])
        if i == (input_classes_len - 1):
            expression += '("clipped_lc@1" = ' + value + ')'
        else:
            expression += '("clipped_lc@1" = ' + value + ') or '
    processing.run("qgis:rastercalculator", {'EXPRESSION': expression,
                                             'LAYERS': [projected_land_cover],
                                             'CELLSIZE': 0,
                                             'EXTENT': None,
                                             'CRS': model_coordinate_system,
                                             'OUTPUT': scratch_folder + '/lc_reclass.tif'})
    os.remove(projected_land_cover)

    # Generate HLS
    hls_raster = processing.run("qgis:rastercalculator", {'EXPRESSION': '"lc_reclass@1" * "slope_reclass@1"',
                                                          'LAYERS': [scratch_folder + '/lc_reclass.tif',
                                                                     scratch_folder + '/slope_reclass.tif'],
                                                          'CELLSIZE': 0,
                                                          'EXTENT': None,
                                                          'CRS': model_coordinate_system,
                                                          'OUTPUT': output})['OUTPUT']
    os.remove(scratch_folder + '/lc_reclass.tif')
    os.remove(scratch_folder + '/slope_reclass.tif')

    return hls_raster


def identifyHlzs(output, hls_raster, tdp_diameter, model_coordinate_system, scratch_folder):
    """
    Function to identify possible HLZs based on criteria and HLS raster
    """
    # Convert raster (only values of 1 or 2) to polygons
    processing.run("gdal:polygonize", {'INPUT': hls_raster,
                                       'BAND': 1,
                                       'FIELD': 'DN',
                                       'EIGHT_CONNECTEDNESS': False,
                                       'EXTRA': '',
                                       'OUTPUT': scratch_folder + '/pixel_polygons.shp'})

    processing.run("native:reprojectlayer", {'INPUT': scratch_folder + '/pixel_polygons.shp',
                                             'TARGET_CRS': model_coordinate_system,
                                             'OPERATION': '+proj=noop',
                                             'OUTPUT': scratch_folder + '/hls_pixels_polygons_projected.shp'})

    selectAndExportFeatures(scratch_folder + '/hls_pixels_polygons_projected.shp',
                            '("DN" = 1) or ("DN" = 2)',
                            scratch_folder + '/selected_pixel_polygons.shp')

    # Dissolve polygons
    hls_dissolve = processing.run("native:dissolve", {'INPUT': scratch_folder + '/selected_pixel_polygons.shp',
                                                      'FIELD': [],
                                                      'SEPARATE_DISJOINT': True,
                                                      'OUTPUT': 'TEMPORARY_OUTPUT'})['OUTPUT']
    hls_dissolve_geometry = \
        processing.run("qgis:exportaddgeometrycolumns", {'INPUT': hls_dissolve,
                                                         'CALC_METHOD': 0,
                                                         'OUTPUT': scratch_folder + '/hls_dissolve_geometry.shp'})['OUTPUT']

    # Select polygons based on area
    selectAndExportFeatures(scratch_folder + '/hls_dissolve_geometry.shp',
                            '"area" > ' + str(tdp_diameter * tdp_diameter),
                            scratch_folder + '/large_enough_areas.shp')

    # Calculate centroids
    initial_hlzs = processing.run("native:centroids", {'INPUT': scratch_folder + '/large_enough_areas.shp',
                                                       'ALL_PARTS': True,
                                                       'OUTPUT': 'TEMPORARY_OUTPUT'})['OUTPUT']

    output_hlzs = processing.run("native:reprojectlayer", {'INPUT': initial_hlzs,
                                                           'TARGET_CRS': model_coordinate_system,
                                                           'OPERATION': '+proj=noop',
                                                           'OUTPUT': output})['OUTPUT']

    return output_hlzs


def selectAndExportFeatures(input_filename, input_expression, output_filename):
    """
    Function to select features within an input file based on an input expression and then export them to a new file
    """
    temp_layer = QgsVectorLayer(input_filename, 'temp_layer')
    temp_layer.selectByExpression(input_expression, QgsVectorLayer.SetSelection)
    writeToShapefile(temp_layer, output_filename, only_selected=True)


# ---------------------------------------------------------------------------
# Output styling (applied via post-processor so the styled layer is the one the
# Processing framework loads on completion, instead of manually adding a
# duplicate layer to the project).
# ---------------------------------------------------------------------------

class _StyleFromQmlPostProcessor(QgsProcessingLayerPostProcessorInterface):
    """
    Base post-processor that applies a bundled .qml style to an output layer.

    The concrete subclass holds the created instance in a class attribute so it
    is not garbage-collected before QGIS calls back into it (the documented
    PyQGIS post-processor lifetime workaround).
    """
    instance = None
    style_path = None

    def postProcessLayer(self, layer, context, feedback):
        if self.style_path and layer is not None and layer.isValid():
            layer.loadNamedStyle(self.style_path)
            layer.triggerRepaint()

    @classmethod
    def create(cls, style_path):
        cls.instance = cls()
        cls.instance.style_path = style_path
        return cls.instance


class HlsRasterStylePostProcessor(_StyleFromQmlPostProcessor):
    instance = None
    style_path = None


class HlzPointsStylePostProcessor(_StyleFromQmlPostProcessor):
    instance = None
    style_path = None


def applyStyleOnLoad(context, output_layer_id, post_processor_cls, style_path):
    """
    Register a style post-processor for an output layer, if that layer will be
    loaded into the project on completion.
    """
    if context.willLoadLayerOnCompletion(output_layer_id):
        context.layerToLoadOnCompletionDetails(output_layer_id).setPostProcessor(
            post_processor_cls.create(style_path))
