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

"""
/***************************************************************************
 GreenViewIndex
                                 A QGIS plugin
 A plugin for Green View Index (GVI) operations
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2023-04-21
        copyright            : (C) 2023 by Alexandros Voukenas
        email                : avoukenas@gmail.com
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   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.                                   *
 *                                                                         *
 ***************************************************************************/
"""

__author__ = 'Alexandros Voukenas'
__date__ = '2023-04-21'
__copyright__ = '(C) 2023 by Alexandros Voukenas'

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

__revision__ = '$Format:%H$'
import os
import inspect
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtCore import QCoreApplication, QVariant
from qgis.core import (QgsProcessing,
                       QgsFeatureSink,
                       QgsVectorLayer,
                       QgsProcessingException,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterFeatureSource,
                       QgsProcessingParameterFeatureSink,
                       QgsProcessingParameterField,
                       QgsProcessingParameterEnum,
                       QgsProcessingParameterNumber,
                       QgsProcessingParameterString,
                       QgsProcessingParameterBoolean,
                       QgsProcessingParameterFolderDestination,
                       QgsField,
                       QgsFields,
                       QgsFeature)
from qgis import processing
import requests,csv,os
from skimage.io import imread,imsave
from skimage.filters.rank import modal
from skimage.color import rgb2hsv
import numpy as np

class CalculateGreenViewIndex(QgsProcessingAlgorithm):

    INPUT_POINTS = 'INPUT POINTS'
    
    INPUT_FIELD="ID FIELD"
    INPUT_FOLDER = 'INPUT FOLDER'
    INPUT_WRITE_MASKED='WRITE MASKED'
    INPUT_ALGORITHM='INPUT ALGORITHM'
    
    OUTPUT = 'OUTPUT'

    def tr(self, string):
        """
        Returns a translatable string with the self.tr() function.
        """
        return QCoreApplication.translate('Processing', string)
    
    def icon(self):
        cmd_folder = os.path.split(inspect.getfile(inspect.currentframe()))[0]
        icon = QIcon(os.path.join(os.path.join(cmd_folder, 'calculate_logo.png')))
        return icon
        
    def createInstance(self):
        return CalculateGreenViewIndex()

    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 'calculate_green_view_index'

    def displayName(self):
        """
        Returns the translated algorithm name, which should be used for any
        user-visible display of the algorithm name.
        """
        return self.tr('Calculate Green View Index')

    def group(self):
        """
        Returns the name of the group this algorithm belongs to. This string
        should be localised.
        """
        return self.tr(self.groupId())

    def groupId(self):
        """
        Returns the unique ID of the group this algorithm belongs to. This
        string should be fixed for the algorithm, and must not be localised.
        The group id should be unique within each provider. Group id should
        contain lowercase alphanumeric characters only and no spaces or other
        formatting characters.
        """
        return ''
        
    def shortHelpString(self):
        """
        Returns a localised short helper string for the algorithm. This string
        should provide a basic description about what the algorithm does and the
        parameters and outputs associated with it..
        """
        return self.tr("""Given an input sample points layer and a folder containing the Google Street View images for those points, this script calculates the Green View Index for each point.
                        The Input folder is the same as the output folder from script 'Download Google Street View Images'. The Input sample points layer is the same as for that script as well.
                        The Unique ID field must also be set, in order to create correspondence between point feature and image file name.
                        The Algorithm to extract vegetation pixels can best to chose between Li et al., 2015 and Dong et al., 2018 (check documentation). Li et al. is the default one, leave it at that if you are unsure.
                        The Write green mask images option is used to create output images that mask for non vegetation (vegetation pixel=255, everything else=0).""")

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

        #add the input parameters to the algorithm
        self.addParameter(
            QgsProcessingParameterFeatureSource(
                self.INPUT_POINTS,
                self.tr('Input sample points layer'),
                [QgsProcessing.TypeVectorPoint]
            )
        )
        self.addParameter(
        QgsProcessingParameterField(
                self.INPUT_FIELD,
                'Unique ID field from sample points layer',
                '',
                self.INPUT_POINTS)
        )
        
        self.addParameter(
        QgsProcessingParameterEnum(
            self.INPUT_ALGORITHM,
            self.tr('Algorithm to extract vegetation pixels'),
            options=[self.tr('Li et al., 2015'),self.tr('Dong et al., 2018')],
            defaultValue=0,
            optional=False)
        )
        
        self.addParameter(
            QgsProcessingParameterFolderDestination(
                self.INPUT_FOLDER,
                self.tr('Input Folder containging GSV images')
            )
        )
        
        self.addParameter(
            QgsProcessingParameterBoolean(
                self.INPUT_WRITE_MASKED,
                self.tr('Write green mask images')
            )
        )
        
        self.addParameter(
            QgsProcessingParameterFeatureSink(
                self.OUTPUT,
                self.tr('Output layer')
            )
        )

    def processAlgorithm(self, parameters, context, feedback):
        
        #define the parameters for the processing algorithm
        source = self.parameterAsSource(
            parameters,
            self.INPUT_POINTS,
            context
        )
        
        id_field=self.parameterAsString(
            parameters,
            self.INPUT_FIELD,
            context
        )
        
        img_folder=self.parameterAsString(
            parameters,
            self.INPUT_FOLDER,
            context
        )
        write_masked=self.parameterAsBoolean(
            parameters,
            self.INPUT_WRITE_MASKED,
            context
        )
        
        output_path=self.parameterAsString(
            parameters,
            self.OUTPUT,
            context
        )
        
        algorithm=self.parameterAsEnum(
            parameters,
            self.INPUT_ALGORITHM,
            context
        )
        
        #set the fields for the output layer. It's one Double type field where the GVI value will be stored, and a unique ID field in order to join with the source layer.
        fields = QgsFields()
        fields.append(QgsField(id_field,QVariant.Double))
        fields.append(QgsField('GVI', QVariant.Double))
        
        (sink, dest_id) = self.parameterAsSink(
            parameters,
            self.OUTPUT,
            context, fields, source.wkbType(), source.sourceCrs()
            )
        #A list to store point IDs and corresponding GVIs
        GVIS=[["pointID","GVI"]]
        
        #Get all the image filenames from the specified directory
        filenames=os.listdir(img_folder)
        #We want to create a list, the elements of which will be lists of each filename corresponding to the same point ID, e.g.:
        #[[pointID_1_heading_0_pitch_0,pointID_1_heading_60_pitch_0],
        # [pointID_2_heading_0_pitch_0,pointID_2_heading_60_pitch_0]
        # [pointID_3_heading_0_pitch_0,pointID_3_heading_60_pitch_0]]
        filenames_new=[]
        IDs_unique=[]
        
        #look up through all the filenames and keep a list of all the pointIDs
        for filename in filenames:
            if filename.endswith(".jpg") and "GVI" not in filename:
                filenames_new.append(filename)
                IDs_unique.append(filename.split("_")[1])
                
        #We need to figure out the number of unique IDs in the filenames list. We use a set for that
        IDs_unique_set=set(IDs_unique)
        IDs_unique=list(IDs_unique_set)
        num_ids=len(IDs_unique)
        #we use numpy array split to split our list into the desired result (see above)
        filenames_splitted=np.array_split(filenames_new,num_ids)
        
        #these two lines help at communicating progress
        all_imgs=len(filenames_new)    
        p=999;k=0
        
        #for each filename group, i.e for each sublist within the mother list:
        for filename_group in filenames_splitted:
            
            d=0 #reset counter
            
            #for each filename within that group:
            for i,filename in enumerate(filename_group):
                #for that filename, get the pointID from its name, and its full path on the drive
                pointID=filename.split("_")[1]
                fullpath=os.path.join(img_folder,filename)
                
                #read the image and split rgb bands
                img=imread(fullpath)
                
                r = img[:,:,0]
                g = img[:,:,1]
                b = img[:,:,2]
                
                #create an array of 255s with the same shape as input image. In the final image, pixels with value 255 will correspond to vegetation, while value 0 is anything else
                GVI_img=np.ones(b.shape,dtype=np.uint8)*255
                
                
                #convert b g r bands to uint16
                b=np.int16(b); g=np.int16(g); r=np.int16(r)
                
                if algorithm==0:
                #Extract green pixels according to Li et. al., 2015
                    diff1=g-r
                    diff2=g-b
                    diff=diff1*diff2
                    
                    GVI_img[diff1<=0]=0
                    GVI_img[diff2<=0]=0
                    GVI_img[diff<50]=0
                    
                    
                elif algorithm==1:
                #Extract green pixels according to Dong et. al., 2018
                    hsv_img = rgb2hsv(img)
                    hue_img = hsv_img[:, :, 0]*255
                    GVI_img[(hue_img>75) & (hue_img<170)]=0
                    
                GVI_img[(b+g+r)>520]=0 #an additional threshold to filter out too bright pixels    
                GVI_img=modal(GVI_img,np.ones((11,11))) #an additional majority filter to smooth the result
                size=b.shape[0]*b.shape[1]
                
                GVIorig=round(len(GVI_img[GVI_img==255])/size,2)#calculate the GVI for that particular image            
                d=d+GVIorig #add that to the counter
                k+=1#keep track of all images that have been parsed
                if write_masked==True:
                    imsave(fullpath[0:-4]+"_GVI.jpg",GVI_img)
                
                #communicate percentage completed
                pct=int((k/all_imgs)*100)
                if pct in [10,20,25,30,33,40,50,60,66,70,75,80,90,100]:
                    if pct!=p:
                        feedback.pushInfo('Finished with {}% of images'.format(pct))
                        p=pct
                            
            #calculate the GVI for all the images corresponding to the point. Append it to the GVIS list            
            GVI_filename=d/(i+1)
            GVIS.append([int(pointID),GVI_filename])

        #write a csv file of the output. It will contain point IDs and corresponding GVI values
        csvfile=os.path.join(img_folder,"Points_GVIs.csv")
        
        with open(csvfile, 'w') as f:
            write = csv.writer(f)
            write.writerows(GVIS)
        csv_layer=QgsVectorLayer(csvfile,'gvis_csv')
        
        #an intermediate field calculator to be able to join with the initial layer
        points_str_field=processing.run("qgis:fieldcalculator", 
        {'INPUT':parameters[self.INPUT_POINTS],
        'FIELD_NAME':'id_str',
        'FIELD_TYPE':2,
        'FIELD_LENGTH':0, 
        'FIELD_PRECISION':0,
        'FORMULA':'\"{}\"'.format(parameters[self.INPUT_FIELD]),
        'OUTPUT':'TEMPORARY_OUTPUT',
        '--overwrite': True})['OUTPUT']
        
        #perform the join with the initial layer
        final_layer_joined=processing.run("qgis:joinattributestable", 
        {'INPUT':points_str_field,
        'FIELD':'id_str',
        'INPUT_2':csv_layer,
        'FIELD_2':'pointID',
        'FIELDS_TO_COPY':['GVI'],'METHOD':1,'DISCARD_NONMATCHING':False,'PREFIX':'',
        'OUTPUT':'TEMPORARY_OUTPUT',
        '--overwrite': True})['OUTPUT']
        
        #write that to feature Sink
        for i,f in enumerate(final_layer_joined.getFeatures()):
                new_feature =  QgsFeature()
                new_feature.setGeometry(f.geometry())
                new_feature.setAttributes([f[parameters[self.INPUT_FIELD]],f['GVI']])
                sink.addFeature(new_feature, QgsFeatureSink.FastInsert)

        return {self.OUTPUT: dest_id}