# -*- 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'
    #DIRECTION= 'DIRECTION'
    RADIUS= 'RADIUS'
    DENOISE = 'DENOISE'
    ANALYSIS_TYPE='ANALYSIS_TYPE'
    OUTPUT = 'OUTPUT'

    ANALYSIS_TYPES = ['Sky-view', '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
            5, 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)

        openness = self.parameterAsInt(parameters,self.ANALYSIS_TYPE, context)
        
        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)
        mx_b = np.zeros(mx_z.shape)
        mx_c = np.zeros(mx_z.shape)
                          
        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)

            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
                                    
            mx_b[:]=0
            # 8 standard lines, for more : step = 0.5; 0.25
            for dy in [-1,0,1]:
                for dx in [-1,0,1]:

                    if dx == dy == 0 : continue

                    mx_a [:]= -9999 if openness else 0
                    mx_c [:]= -9999 if openness else 0   
                    
                    for r in range (1, radius + 1):

                        view_in, view_out = view(r * dy, r * dx, mx_z.shape)

                        angles = mx_z[view_in] - mx_z[view_out]
                        angles /= r * pixel_size 
			
                        np.maximum(mx_a[view_out], angles, mx_a[view_out] )
                        np.maximum(mx_c[view_in], -angles, mx_c[view_in] )
                        # arg 3 = result keeping array
                    
                    mx_b += mx_a + mx_c

                    counter += 1
                    feedback.setProgress(100 * chunk * (counter/8) /  xsize)
                    if feedback.isCanceled(): sys.exit()
        
            mx_b /= 8

            #For speed we should do without tangent calcualtion, but it's not crucial... (and not evident)
            if openness:
                out = 1 - (np.arctan(mx_b) / np.pi + 0.5)
            else:
                out = 1 - (np.arctan(mx_b)* 2 / np.pi )

            ds.GetRasterBand(1).WriteArray(out[mx_view_out], * gdal_put[:2])

        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())

    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:
                 - Analysis type: Sky-view option will model the light comming from directions above individual pixels (i.e. from the sky). Openness option 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()
