# -*- coding: utf-8 -*-
"""
/***************************************************************************
 ExtractToPointsDialog
                                 A QGIS plugin
 extract
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                             -------------------
        begin                : 2019-06-06
        git sha              : $Format:%H$
        copyright            : (C) 2019 by zxx
        email                : zhengxinxin@whu.edu.cn
 ***************************************************************************/

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

import os

from PyQt5 import uic
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import QDialog, QTableWidgetItem, QMessageBox, QFileDialog, QInputDialog
from qgis.core import QgsWkbTypes, QgsVectorLayer, QgsProject, QgsFields, QgsField, QgsFeature, QgsRectangle, QgsGeometry, QgsFeatureRequest, QgsRaster, QgsVectorFileWriter

# This loads your .ui file so that PyQt can populate your plugin with the elements from Qt Designer
UI_Dialog = uic.loadUiType(os.path.join(os.path.dirname(__file__), 'ExtractToPoints_dialog_base.ui'))[0]


class ExtractToPointsDialog(QDialog, UI_Dialog):
    
    targetItems = {}
    polygonItems = {}
    rasterItems = {}
    fields = []
    
    def __init__(self, iface):
        QDialog.__init__(self)
        self.iface = iface
        self.setupUi(self)
        
        #connect
        self.BrowseButton.clicked.connect(self.browseFile)
        self.TargetLayer.currentIndexChanged.connect(self.updateFieldsList)
        self.AddToMap.setCheckState(Qt.Checked)
        self.SourceField.itemSelectionChanged.connect(self.updateFieldsTable)
        self.FieldsTable.cellChanged.connect(self.fieldNameChanged)
        self.buttonBox.accepted.connect(self.accept)
        self.buttonBox.rejected.connect(self.reject)
        
        mapCanvas = self.iface.mapCanvas()
        
        # initial dictionaries of items:
        self.targetItems = {}
        self.polygonItems = {}
        self.rasterItems = {}
        
        for i in range(mapCanvas.layerCount()):
            layer = mapCanvas.layer(i)
            
            if (layer.type() == layer.VectorLayer) and (layer.geometryType() == QgsWkbTypes.PointGeometry):
                # identify and get point layer
                provider = layer.dataProvider()
                fields = provider.fields()
                allItem = [layer]
                for j in fields:
                    allItem += [[str(j.name()), str(j.name()), False]]
                self.targetItems[str(layer.name())] = allItem
                self.TargetLayer.addItem(layer.name())
                
            elif (layer.type() == layer.VectorLayer) and (layer.geometryType() == QgsWkbTypes.PolygonGeometry):
                # identify and get polygon layer
                provider = layer.dataProvider()
                fields = provider.fields()
                allItem = [layer]
                for j in fields:
                    allItem += [[str(j.name()), str(j.name()), False]]
                self.polygonItems[str(layer.name())] = allItem
                
            elif layer.type() == layer.RasterLayer:
                # identify and get raster layer
                allItem = [layer]
                for j in range(layer.bandCount()):
                    if layer.bandCount() == 1:
                        name1 = layer.bandName(j+1) 
                        name2 = layer.name()[:10]
                    else:
                        name1 = layer.bandName(j+1)
                        name2 = layer.name()[:8] + "_" + str(j+1)
                    allItem += [[name1, name2, False]]
                self.rasterItems[str(layer.name())] = allItem
        
        self.updateFieldsList()
        
    def updateFieldsList(self):
        
        self.SourceField.clear()
        if not self.TargetLayer.count():
            return
        # i:layer j:fields
        i = self.TargetLayer.currentText()
        for j in range(1, len(self.targetItems[i])):
            #clear previously enabled fields (as they aren't selected in the widget)
            self.targetItems[i][j][2] = False
            self.SourceField.addItem(self.targetItems[i][0].name() + " : " + self.targetItems[i][j][0] + "  (source layer field)")
            
        for i in self.polygonItems:
            for j in range(1, len(self.polygonItems[i])):
                self.SourceField.addItem(str(self.polygonItems[i][0].name()) + " : " + str(self.polygonItems[i][i][0]) + "  (polygon)")
            
        for i in self.rasterItems:
            for j in range(1, len(self.rasterItems[i])):
                self.SourceField.addItem(str(self.rasterItems[i][0].name()) + " : " + str(self.rasterItems[i][j][0]) + "  (raster)")
        
        self.updateFieldsTable()
        self.repaint()
        
        
    def updateFieldsTable(self):
        # after selection, update the table
        
        # mark selected point items
        n = 0
        i = self.TargetLayer.currentText()
        for j in range(1, len(self.targetItems[i])):
            if self.SourceField.item(n) and self.SourceField.item(n).isSelected():
                self.targetItems[i][j][2] = True
            else:
                self.targetItems[i][j][2] = False
            n += 1
            
        # mark selected polygon items
        for i in self.polygonItems:
            for j in range(1, len(self.polygonItems[i])):
                if self.SourceField.item(n) and self.SourceField.item(n).isSelected():
                    self.polygonItems[i][j][2] = True
                else:
                    self.polygonItems[i][j][2] = False
                n += 1
        
        # mark selected raster items
        for i in self.rasterItems:
            for j in range(1, len(self.rasterItems[i])):
                if self.SourceField.item(n) and self.SourceField.item(n).isSelected():
                    self.rasterItems[i][j][2] = True
                else:
                    self.rasterItems[i][j][2] = False
                n += 1
        
        # fill FieldsTable with point, polygon and raster items
        self.fields = []
        n = 0
        self.FieldsTable.setRowCount(0)
        # point
        i = self.TargetLayer.currentText()
        for j in range(1, len(self.targetItems[i])):
            if self.targetItems[i][j][2]:
                self.fields += [["point",i,j]]
                self.FieldsTable.setRowCount(n+1)
                cell = QTableWidgetItem(str(self.targetItems[i][0].name()) + " : " + str(self.targetItems[i][j][0]))
                cell.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
                self.FieldsTable.setItem(n, 0, cell)
                self.FieldsTable.setItem(n, 1, QTableWidgetItem(str(self.targetItems[i][j][1])))
                n += 1
        #polygon
        for i in self.polygonItems:
            for j in range(1, len(self.polygonItems[i])):
                if self.polygonItems[i][j][2]:
                    self.fields += [["polygon",i,j]]
                    self.FieldsTable.setRowCount(n+1)
                    cell = QTableWidgetItem(str(self.polygonItems[i][0].name()) + " : " +str(self.polygonItems[i][j][0]))
                    cell.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
                    self.FieldsTable.setItem(n, 0, cell)
                    self.FieldsTable.setItem(n, 1, QTableWidgetItem(str(self.polygonItems[i][j][1])))
                    n += 1
        #raster
        for i in self.rasterItems:
            for j in range(1, len(self.rasterItems[i])):
                if self.rasterItems[i][j][2]:
                    self.fields += [["raster",i,j]]
                    self.FieldsTable.setRowCount(n+1)
                    cell = QTableWidgetItem(str(self.rasterItems[i][0].name()) + " : " +str(self.rasterItems[i][j][0]))
                    cell.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
                    self.FieldsTable.setItem(n, 0, cell)
                    self.FieldsTable.setItem(n, 1, QTableWidgetItem(str(self.rasterItems[i][j][1])))
                    n += 1
        
        self.FieldsTable.resizeColumnsToContents()
        
    def fieldNameChanged(self, n):
        # when cell of the FieldsTable was modified
        
        # exit when false alarm
        if len(self.fields) == 0:
            return 0
        if self.FieldsTable.rowCount() == 0:
            return 0
        
        updatedItem = self.FieldsTable.item(n,1)
        if updatedItem == None:
            return 0
        
        # update
        updatedText = str(updatedItem.text())
        if self.fields[n][0] == "point":
            self.targetItems[self.fields[n][1]][self.fields[n][2]][1] = updatedText[:10]
        
        elif self.fields[n][0] == "polygon":
            self.polygonItems[self.fields[n][1]][self.fields[n][2]][1] = updatedText[:10]
        
        else:
            self.rasterItems[self.fields[n][1]][self.fields[n][2]][1] = updatedText[:10]
        
        # limit 10 characters
        if len(updatedText) > 10:
            self.updateFieldsTable()
            QMessageBox.information(self, self.tr("Extract Tool"), self.tr("Limit 10 Characters."))
            
        self.FieldsTable.resizeColumnsToContents()
        
    def browseFile(self):
        # call dialog for output file
        self.OutLayer.clear()
        outName, _ = QFileDialog().getSaveFileName(self, self.tr("Output file"), '.', 
                                                   self.tr("Shapefiles (*.shp);;Comma separated values (*.csv);;GeoPackages(*.gpkg)"), 
                                                   options = QFileDialog.DontConfirmOverwrite)
        outPath = QFileInfo(outName).absoluteFilePath()
        if not outPath.upper().endswith('.SHP') and not outPath.upper().endswith('.CSV') and not outPath.upper().endswith('.GPKG'):
            outPath += '.gpkg'
        if outName:
            self.OutLayer.clear()
            self.OutLayer.insert(outPath)
            
    def accept(self):
        # "OK" button
        
        # check if all fields are filled up
        nothingSelected = True
        for i in self.polygonItems:
            for j in range(1, len(self.polygonItems[i])):
                if self.polygonItems[i][j][2]:
                    nothingSelected = False
        
        for i in self.rasterItems:
            for j in range(1, len(self.rasterItems[i])):
                if self.rasterItems[i][j][2]:
                    nothingSelected = False
                    
        if self.TargetLayer.currentText() == "":
            QMessageBox.information(self, self.tr("Extract Tool"), self.tr("Please select target point layer"))
            return
        if nothingSelected:
            QMessageBox.information(self, self.tr("Extract Tool"), self.tr("Please select at least one field"))
            return 
        if self.OutLayer.text() == "":
            QMessageBox.information(self, self.tr("Extract Tool"), self.tr("Please specify output file name"))
            return
        
        # check if target field names are unique
        if not self.testFieldNames(self.fields):
            self.updateFieldsTable()
            QMessageBox.warning(self, self.tr("Extract Tool"), self.tr("At least two field names are the same!\nPlease type unique names."))
            return
        
        # check if there a CRS mismatch
        targetLayerSrid = list(self.targetItems.values())[0][0].crs().postgisSrid()
        msg = self.tr('''<html>All layers must have the same coordinate reference system. The <b>%s</b> layer seems to have different CRS id (<b>%d</b>)
                   than the point layer (<b>%d</b>). If they are two different CRSes, you need to reproject one of the layers first,
                   otherwise results will be wrong.<br/>
                   However, if you are sure both CRSes are the same, and they are just improperly recognized, you can safely continue.
                   Do you want to continue?</html>''')
        
        for i in self.polygonItems:
            for j in range(1, len(self.polygonItems[i])):
                if self.polygonItems[i][j][2]:
                    layerSrid = self.polygonItems[i][0].crs().postgisSrid()
                    if layerSrid != targetLayerSrid:
                        if QMessageBox.question(self, self.tr("Extract Tool: layer CRS mismatch!"), msg % (i, layerSrid, targetLayerSrid), QMessageBox.Yes | QMessageBox.No) != QMessageBox.Yes:
                            return
        
        for i in self.rasterItems:
            for j in range(1, len(self.rasterItems[i])):
                if self.rasterItems[i][j][2]:
                    layerSrid = self.rasterItems[i][0].crs().postgisSrid()
                    if layerSrid != targetLayerSrid:
                        if QMessageBox.question(self, self.tr("Extract Tool: layer CRS mismatch!"), msg % (i, layerSrid, targetLayerSrid), QMessageBox.Yes | QMessageBox.No) != QMessageBox.Yes:
                            return
        
        if True:
            # all tests passed, go on
            
            self.repaint()
            outPath = self.OutLayer.text()
            outPath = outPath.replace("\\","/")
            if not outPath.upper().endswith('.SHP') and not outPath.upper().endswith('.CSV') and not outPath.upper().endswith('.GPKG'):
                outPath += '.gpkg'
            outName = QFileInfo(outPath).fileName()
            tableName = None
            oldFile = QFile(outPath)
            if oldFile.exists():
                if not outPath.upper().endswith('.GPKG'):
                    if QMessageBox.question(self, self.tr("Extract Tool"), self.tr("File %s already exists. Do you want to overwrite?") % outName) == QMessageBox.No:
                        self.OutLayer.clear()
                        self.repaint()
                        return
                    
                else:
                    msg = self.tr("""Please provide <b>table name</b> for your layer.<br/>
                      <b>WARNING: </b>Database %s already exists. If you select a table existing in it, the table will be overwritten.""") % outName
                    tableName, result = QInputDialog.getText(self, "Extract Tool", msg, text = outName[:-5])
                    
                    if not result:
                        self.OutLayer.clear()
                        self.repaint()
                        return
            
            self.repaint()
            # execute main function
            if not self.extract(outPath, tableName):
                return 
        
            self.OutLayer.clear()
            if self.AddToMap.checkState() == Qt.Checked:
                uri = outPath
                layerName = outName
                if tableName:
                    uri += "|layername=%s" % tableName
                    layerName += ": %s" % tableName
                self.vlayer = QgsVectorLayer(uri, layerName, "ogr")
                
                if self.vlayer.isValid():
                    # Add layer to the map, but first remove it if already present
                    for l in QgsProject.instance().mapLayers().values():
                        if hasattr(l, 'source') and l.source() == self.vlayer.source():
                            QgsProject.instance().removeMapLayer(l)
                            
                    QgsProject.instance().addMapLayer(self.vlayer)
                
                else:
                    QMessageBox.warning(self, self.tr("Extract Tool"), self.tr("The new layer seems to be created, but is invalid.\nIt won't be loaded."))
                
    def extract(self, outPath, tableName):
        # main process
        
        # open target point layer
        targetLayer = self.targetItems[str(self.TargetLayer.currentText())][0]
        targetprovider = targetLayer.dataProvider()
        allAttrs = targetprovider.attributeIndexes()
        sRs = targetLayer.crs()
        
        # create output layer: first create list of selected fields
        fieldList = QgsFields()
        for i in range(len(self.fields)):
            if self.fields[i][0] == "point":
                # copying fields from source layer
                field = targetprovider.fields()[targetprovider.fieldNameIndex(self.targetItems[self.fields[i][1]][self.fields[i][2]][0])]
                field.setName(self.targetItems[self.fields[i][1]][self.fields[i][2]][1])
                
            elif self.fields[i][0] == "polygon": 
                # copying fields from polygon layers
                polyLayer = self.polygonItems[self.fields[i][1]][0]
                polyProvider = polyLayer.dataProvider()
                field = polyProvider.fields()[polyProvider.fieldNameIndex(self.polygonItems[self.fields[i][1]][self.fields[i][2]][0])]
                field.setName(self.polygonItems[self.fields[i][1]][self.fields[i][2]][1])
            
            else: 
                # creating fields for raster layers
                field = QgsField(self.rasterItems[self.fields[i][1]][self.fields[i][2]][1], QVariant.Double, "real", 20, 5, "")
            
            fieldList.append(field)
            
        # create temporary memory layer 
        memLayer = QgsVectorLayer("Point?crs=epsg:%d" % sRs.postgisSrid(), 'temp layer', 'memory')
        memLayer.startEditing()
        for field in fieldList:
            memLayer.addAttribute(field)
        memLayer.commitChanges()
        
        self.repaint()
        # process point after point
        pointFeat = QgsFeature()
        np = 0
        snp = targetprovider.featureCount()
        for pointFeat in targetprovider.getFeatures():
            np += 1
            # convert multipoint[0] to point
            pointGeom = pointFeat.geometry()
            if pointGeom.wkbType() == QgsWkbTypes.MultiPoint:
                pointPoint = pointGeom.asMultiPoint()[0]
            else:
                pointPoint = pointGeom.asPoint()
            outFeat = QgsFeature()
            outFeat.setGeometry(pointGeom)
            
            # and next loop inside: field after field
            bBox = QgsRectangle(pointPoint.x()-0.001,pointPoint.y()-0.001,pointPoint.x()+0.001,pointPoint.y()+0.001) 
            # reusable rectangle buffer around the point feature
            previousPolyLayer = None  # reuse previous feature if it's still the same layer
            previousPolyFeat = None   # reuse previous feature if it's still the same layer
            previousRastLayer = None  # reuse previous raster multichannel sample if it's still the same layer
            previousRastSample = None # reuse previous raster multichannel sample if it's still the same layer
            attrs = []
            for i in range(len(self.fields)):
                field = self.fields[i]
                
                if field[0] == "point":
                    attr = pointFeat.attributes()[targetprovider.fieldNameIndex(self.targetItems[field[1]][field[2]][0])]
                    attrs += [attr]
                    
                elif field[0] == "polygon":
                    polyLayer = self.polygonItems[field[1]][0]
                    polyProvider = polyLayer.dataProvider()
                    
                    if polyLayer == previousPolyLayer:
                        polyFeat = previousPolyFeat
                        
                    else:
                        polyFeat = None
                        pointGeom = QgsGeometry().fromPointXY(pointPoint)
                        for iFeat in polyProvider.getFeatures(QgsFeatureRequest().setFilterRect(bBox)):
                            if pointGeom.intersects(iFeat.geometry()):
                                polyFeat = iFeat
                                
                    if polyFeat:
                        attr = polyFeat.attributes()[polyProvider.fieldNameIndex(self.polygonItems[field[1]][field[2]][0])]
                    
                    else:
                        attr = None
                    attrs += [attr] 
                    #only last one if more polygons overlaps
                    previousPolyLayer = polyLayer
                    previousPolyFeat = polyFeat
                    
                else:
                    # raster
                    rastLayer = self.rasterItems[field[1]][0]
                    if rastLayer == previousRastLayer:
                        rastSample = previousRastSample
                        
                    else:
                        rastSample = rastLayer.dataProvider().identify(pointPoint, QgsRaster.IdentifyFormatValue).results()
                    
                    try:
                        bandNo = field[2]
                        attr = float(rastSample[bandNo])
                    except: # point is out of raster extent
                        attr = None
                    
                    attrs += [attr]
                    previousRastLayer = rastLayer
                    previousRastSample = rastSample
                    
            outFeat.initAttributes(len(attrs))
            outFeat.setAttributes(attrs)
            memLayer.dataProvider().addFeature(outFeat)
        
        # write memlayer to the output file
        so = QgsVectorFileWriter.SaveVectorOptions()
        so.fileEncoding = 'UTF-8'
        
        if outPath.upper().endswith('SHP'):
            so.driverName = "ESRI Shapefile"
        elif outPath.upper().endswith('CSV'):
            so.driverName = "CSV"
        else:
            so.driverName = "GPKG"
            if tableName:
                so.actionOnExistingFile = QgsVectorFileWriter.CreateOrOverwriteLayer
                so.layerName = tableName
        result, errMsg = QgsVectorFileWriter.writeAsVectorFormat(memLayer, outPath, so)
        if result:
            QMessageBox.critical(self, self.tr("Extract tool"), errMsg)
            return False
        else:
            del memLayer
            return True
        
    def testFieldNames(self, fields):
        #tests uniqueness of field names
        ok = True
        if len(fields) > 1:
            for field1 in fields:
                for field2 in fields:
                    if field1[0] == "point": 
                        name1 = self.targetItems[field1[1]][field1[2]][1]
                    elif field1[0] == "polygon": 
                        name1 = self.polygonItems[field1[1]][field1[2]][1]
                    else: 
                        name1 = self.rasterItems[field1[1]][field1[2]][1]
                        
                    if field2[0] == "point": 
                        name2 = self.targetItems[field2[1]][field2[2]][1]
                    elif field2[0] == "polygon": 
                        name2 = self.polygonItems[field2[1]][field2[2]][1]
                    else: 
                        name2 = self.rasterItems[field2[1]][field2[2]][1]
                        
                    if (name1 == name2) and (field1 != field2):
                        ok = False
                        
        return ok