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

"""
/***************************************************************************
 DemShading - ambient occlusion
 This algorithm simulates ambiental occlusion effect over an elevation model (DEM)
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2020-02-20
        copyright            : (C) 2020 by Zoran Čučković
 ***************************************************************************/
/***************************************************************************
 *                                                                         *
 *   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.                                   *
 *                                                                         *
 ***************************************************************************/
"""

# ========= TODO : BETTER HANDLE BORDERS

__author__ = 'Zoran Čučković'
__date__ = '2020-02-05'
__copyright__ = '(C) 2020 by Zoran Čučković'

# This will get replaced with a git SHA1 when you do a git archive

__revision__ = '$Format:%H$'

from os import sys, path

from PyQt5.QtCore import QCoreApplication
from qgis.core import (QgsProcessing,
                       QgsProcessingException,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterRasterLayer,
                       QgsProcessingParameterRasterDestination,
                        QgsProcessingParameterBoolean,
                      QgsProcessingParameterNumber,
                       QgsProcessingParameterEnum,


                       QgsProcessingUtils,
                        QgsRasterBandStats,
                       QgsSingleBandGrayRenderer,
                       QgsContrastEnhancement
                        )

from processing.core.ProcessingConfig import ProcessingConfig

import gdal
import numpy as np
from .modules.helpers import view, window_loop

from qgis.core import QgsMessageLog # for testing

class OcclusionAlgorithm(QgsProcessingAlgorithm):
    """
    This algorithm simulates ambient lighting over a raster DEM (in input). 
    """

    # Constants used to refer to parameters and outputs. They will be
    # used when calling the algorithm from another algorithm, or when
    # calling from the QGIS console.

    INPUT = 'INPUT'
    RADIUS= 'RADIUS'
    DENOISE = 'DENOISE'
    ANALYSIS_TYPE='ANALYSIS_TYPE'
    OUTPUT = 'OUTPUT'

    ANALYSIS_TYPES = ['Sky-view','Sky-view (symmetric)', 'Openness']

    output_model = None #for post-processing

    def initAlgorithm(self, config):
        """
        Here we define the inputs and output of the algorithm, along
        with some other properties.
        """

        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.INPUT,
                self.tr('Digital elevation model')
            ) )
        
        self.addParameter(QgsProcessingParameterEnum (
            self.ANALYSIS_TYPE,
            self.tr('Analysis type'),
            self.ANALYSIS_TYPES,
            defaultValue=0))
                    
        self.addParameter(QgsProcessingParameterNumber(
            self.RADIUS,
            self.tr('Radius (pixels)'),
            0, # QgsProcessingParameterNumber.Integer = 0
            7, False, 0, 100))
        
        self.addParameter(QgsProcessingParameterBoolean(
            self.DENOISE,
            self.tr('Denoise'),
            False, False)) 
        
        self.addParameter(
            QgsProcessingParameterRasterDestination(
                self.OUTPUT,
            self.tr("Ambient occlusion")))
        
    def processAlgorithm(self, parameters, context, feedback):
        """
        Here is where the processing itself takes place.
        """

        elevation_model= self.parameterAsRasterLayer(parameters,self.INPUT, context)

        if elevation_model.crs().mapUnits() != 0 :
            err= " \n ****** \n ERROR! \n Raster data has to be projected in a metric system!"
            feedback.reportError(err, fatalError = False)
           # raise QgsProcessingException(err)
        #could also use:
        #raise QgsProcessingException(self.invalidSourceError(parameters, self.INPUT))

        if  round(abs(elevation_model.rasterUnitsPerPixelX()),
                    2) !=  round(abs(elevation_model.rasterUnitsPerPixelY()),2):
            
            err= (" \n ****** \n ERROR! \n Raster pixels are irregular in shape " +
                  "(probably due to incorrect projection)!")
            feedback.reportError(err, fatalError = False)
            #raise QgsProcessingException(err)

        self.output_model = self.parameterAsOutputLayer(parameters,self.OUTPUT,context)

        #direction = self.parameterAsDouble(parameters,self.DIRECTION, context)
        denoise = self.parameterAsInt(parameters,self.DENOISE, context)        
       
        radius =self.parameterAsInt(parameters,self.RADIUS, context)
        
        method = self.parameterAsInt(parameters,self.ANALYSIS_TYPE, context)
        symmetric = method == 1
        openness = method == 2
        
        overlap = radius if not denoise else radius +1
             
        dem = gdal.Open(elevation_model.source())
          
        # ! attention: x in gdal is y dimension un numpy (the first dimension)
        xsize, ysize = dem.RasterXSize,dem.RasterYSize
        #assuming one band dem !
        nodata = dem.GetRasterBand(1).GetNoDataValue()
        
        pixel_size = dem.GetGeoTransform()[1]
        
        chunk = int(ProcessingConfig.getSetting('DATA_CHUNK')) * 1000000
        chunk = min(chunk // xsize, xsize)
        
        # writing output to dump data chunks
        driver = gdal.GetDriverByName('GTiff')
        ds = driver.Create(self.output_model, xsize,ysize, 1, gdal.GDT_Float32)
        ds.SetProjection(dem.GetProjection())
        ds.SetGeoTransform(dem.GetGeoTransform())
        
             
        chunk_slice = (ysize, chunk + 2 * overlap) 
        
        mx_z = np.zeros(chunk_slice)
        mx_a = np.zeros(mx_z.shape)
        if symmetric: mx_b = mx_a
        else: mx_b = np.zeros(mx_z.shape)
        out =  np.zeros(mx_z.shape)
        
        
        # intialise the count of lines per pixel
        mx_cnt = np.ones(mx_z.shape)
        # set borders first  
        mx_cnt[:]= 5 if not symmetric else 3
        # main area : 8 lines per pixel (or 4 if symmetric algo)
        mx_cnt[1:-1, 1:-1] = 8 if not symmetric else 4
        # corners
        for v in [(0,0),(-1,-1),(0,-1), (-1,0)]: mx_cnt[v] = 3
                          
        counter = 0
            
        for mx_view_in, gdal_take, mx_view_out, gdal_put in window_loop ( 
            shape = (xsize, ysize), 
            chunk = chunk,
            overlap = overlap) :
     
            mx_z[mx_view_in]= dem.ReadAsArray(*gdal_take).astype(float)
            # NODATA : TODO !
            # mx_z[mx_z == nodata] = 0
            
            if denoise :
                mx_a[:]=0
                for i in range(-1,2):
                    for j in range(-1,2):
                        view_in, view_out = view(i , j ,mx_z.shape)
                        mx_a[view_out] += mx_z[view_in]
                    
                mx_z = mx_a/9
         
                # 8 standard lines, we use symetry to optimise
            for dy, dx in [(0,1), (1,0), (1, -1), (1,1)]:
                
                mx_a [:] = 0 if not openness else -9999999
                mx_b [:] = 0 if not openness else -9999999
                
                for r in range (1, radius + 1): # we could probably sample over radius, not all pixels are needed
                                  
                    view_in, view_out = view(r * dx, r * dy, mx_z[mx_view_in].shape)

                    angles = mx_z[view_in] - mx_z[view_out]
                                                   # diagonals         
                    dist = r * pixel_size * (1.4142 if dx * dy != 0 else 1) 
                    
                    angles /= dist 
                    
                    # slow functions, perhaps with cumsum would be faster (?)
                    np.maximum(mx_a[view_out], angles, mx_a[view_out] )
                    np.maximum(mx_b[view_in], -angles, mx_b[view_in] )
                
               # ugly patch, edge values not overwritten...
               # TODO : better handling of raster edges !
                if openness: mx_a[mx_a < -99] = 0; mx_b[mx_b < -99] = 0 

                # average of angles: see Kokalj et al. 2011
                # these operations are costly, however ...
                out += np.sin(np.arctan(mx_a))          
                if not symmetric : out += np.sin(np.arctan(mx_b)) 

                counter += 1
                feedback.setProgress(100 * chunk * (counter/4) /  xsize)
                if feedback.isCanceled(): sys.exit()
            
            # this is a patch : last chunk is often spilling outside raster edge 
            # so, move the edge values to match raster edge
            end = gdal_take[2]           
            if end + gdal_take[0] == xsize : 
                mx_cnt[:, end-1: end] = mx_cnt[:, -1:] 
                
            out = 1 - out/mx_cnt
            
            ds.GetRasterBand(1).WriteArray(out[mx_view_out], * gdal_put[:2])
                
            out[:]=0 # RESET

        ds = None
        
        return {self.OUTPUT: self.output_model}

    def postProcessAlgorithm(self, context, feedback):

        output = QgsProcessingUtils.mapLayerFromString(self.output_model, context)
        provider = output.dataProvider()

        stats = provider.bandStatistics(1,QgsRasterBandStats.All,output.extent(),0)
        mean, sd = stats.mean, stats.stdDev
        
        rnd = QgsSingleBandGrayRenderer(provider, 1)
        ce = QgsContrastEnhancement(provider.dataType(1))
        ce.setContrastEnhancementAlgorithm(QgsContrastEnhancement.StretchToMinimumMaximum)
       
        ce.setMinimumValue(mean-3*sd)
        ce.setMaximumValue(min(1, mean+2*sd))

        rnd.setContrastEnhancement(ce)

        output.setRenderer(rnd)
        
        output.triggerRepaint()

        return {self.OUTPUT: self.output_model}

    def name(self):
        """
        Returns the algorithm name, used for identifying the algorithm. This
        string should be fixed for the algorithm, and must not be localised.
        The name should be unique within each provider. Names should contain
        lowercase alphanumeric characters only and no spaces or other
        formatting characters.
        """
        return 'Ambient occlusion'

    def displayName(self):
        """
        Returns the translated algorithm name, which should be used for any
        user-visible display of the algorithm name.
        """
        return self.tr(self.name()+ " (sky-view)")

    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def shortHelpString(self):
        curr_dir = path.dirname(path.realpath(__file__))
        h = ( """
                Ambient occlusion of a locale is the proportion of the ambient light that it may recieve. This algorithm assumes equal light intensity from all directions (simple ambient lighting).
                Parameters:
                 - Sky-view: models the light comming from 8 directions towards individual pixels (i.e. from the sky).
                 - Symmetric sky-view: For each pair of opposite directions, take the one with higher horizon. Nice visual effect.
                 - Openness: allows for light sources situated below the horizontal plane.
                 - Radius: The ambient occlusion is caluclated within a defined radius for each raster pixel (computation time is directly dependant on the analysis radius).
                 - Denoise: Apply a simple smoothing filter.
                NB. This algorithm is made for terrain visualisation, it is not appropriate for precise calculation of solar exposition or of incident light.
                For more information see <a href="https://zoran-cuckovic.github.io/QGIS-terrain-shading/"> the manual </a> and the in-depth <a href=https://landscapearchaeology.org/2020/ambient-occlusion/"> blog post</a>.
                
                If you find this tool useful, consider to :
                 
             <a href='https://ko-fi.com/D1D41HYSW' target='_blank'><img height='30' style='border:0px;height:36px;' src='%s/help/kofi2.webp' /></a>
            """) % curr_dir
		
        return self.tr(h)

    def createInstance(self):
        return OcclusionAlgorithm()
