# -*- coding: utf-8 -*-
"""
/***************************************************************************
 ClipRaster
                                 A QGIS plugin
 This plugin allows you to split raster files based on given polygon data.
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2020-12-13
        git sha              : $Format:%H$
        copyright            : (C) 2020 by Murat ÇALIŞKAN
        email                : caliskan.murat.20@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.                                   *
 *                                                                         *
 ***************************************************************************/
"""
from qgis.PyQt.QtCore import QSettings, QTranslator, QCoreApplication
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtWidgets import QAction, QFileDialog
from qgis.utils import iface

from qgis.core import Qgis, QgsProject, QgsGeometry
import processing

# Initialize Qt resources from file resources.py
from .resources import *
# Import the code for the dialog
from .easy_raster_splitter_dialog import ClipRasterDialog
import os.path

import numpy as np
from osgeo import gdal
import os

class ClipRaster:
    """QGIS Plugin Implementation."""

    def __init__(self, iface):
        """Constructor.

        :param iface: An interface instance that will be passed to this class
            which provides the hook by which you can manipulate the QGIS
            application at run time.
        :type iface: QgsInterface
        """
        # Save reference to the QGIS interface
        self.iface = iface
        # initialize plugin directory
        self.plugin_dir = os.path.dirname(__file__)
        # initialize locale
        locale = QSettings().value('locale/userLocale')[0:2]
        locale_path = os.path.join(
            self.plugin_dir,
            'i18n',
            'ClipRaster_{}.qm'.format(locale))

        if os.path.exists(locale_path):
            self.translator = QTranslator()
            self.translator.load(locale_path)
            QCoreApplication.installTranslator(self.translator)

        # Declare instance attributes
        self.actions = []
        self.menu = self.tr(u'&Easy Raster Splitter')

        # Check if plugin was started the first time in current QGIS session
        # Must be set in initGui() to survive plugin reloads
        self.first_start = None

    # noinspection PyMethodMayBeStatic
    def tr(self, message):
        """Get the translation for a string using Qt translation API.

        We implement this ourselves since we do not inherit QObject.

        :param message: String for translation.
        :type message: str, QString

        :returns: Translated version of message.
        :rtype: QString
        """
        # noinspection PyTypeChecker,PyArgumentList,PyCallByClass
        return QCoreApplication.translate('ClipRaster', message)


    def add_action(
        self,
        icon_path,
        text,
        callback,
        enabled_flag=True,
        add_to_menu=True,
        add_to_toolbar=True,
        status_tip=None,
        whats_this=None,
        parent=None):
        """Add a toolbar icon to the toolbar.

        :param icon_path: Path to the icon for this action. Can be a resource
            path (e.g. ':/plugins/foo/bar.png') or a normal file system path.
        :type icon_path: str

        :param text: Text that should be shown in menu items for this action.
        :type text: str

        :param callback: Function to be called when the action is triggered.
        :type callback: function

        :param enabled_flag: A flag indicating if the action should be enabled
            by default. Defaults to True.
        :type enabled_flag: bool

        :param add_to_menu: Flag indicating whether the action should also
            be added to the menu. Defaults to True.
        :type add_to_menu: bool

        :param add_to_toolbar: Flag indicating whether the action should also
            be added to the toolbar. Defaults to True.
        :type add_to_toolbar: bool

        :param status_tip: Optional text to show in a popup when mouse pointer
            hovers over the action.
        :type status_tip: str

        :param parent: Parent widget for the new action. Defaults None.
        :type parent: QWidget

        :param whats_this: Optional text to show in the status bar when the
            mouse pointer hovers over the action.

        :returns: The action that was created. Note that the action is also
            added to self.actions list.
        :rtype: QAction
        """

        icon = QIcon(icon_path)
        action = QAction(icon, text, parent)
        action.triggered.connect(callback)
        action.setEnabled(enabled_flag)

        if status_tip is not None:
            action.setStatusTip(status_tip)

        if whats_this is not None:
            action.setWhatsThis(whats_this)

        if add_to_toolbar:
            # Adds plugin icon to Plugins toolbar
            self.iface.addToolBarIcon(action)

        if add_to_menu:
            self.iface.addPluginToRasterMenu(
                self.menu,
                action)

        self.actions.append(action)

        return action

    def initGui(self):
        """Create the menu entries and toolbar icons inside the QGIS GUI."""

        icon_path = ':/plugins/easy_raster_splitter/icon.png'
        self.add_action(
            icon_path,
            text=self.tr(u'Easy Raster Splitter'),
            callback=self.run,
            parent=self.iface.mainWindow())

        # will be set False in run()
        self.first_start = True


    def unload(self):
        """Removes the plugin menu item and icon from QGIS GUI."""
        for action in self.actions:
            self.iface.removePluginRasterMenu(
                self.tr(u'&Easy Raster Splitter'),
                action)
            self.iface.removeToolBarIcon(action)

    def selectOutput(self):
        self.dlg.lineEdit_outRaster.setText("")
        self.OutputDir = QFileDialog.getExistingDirectory(None, 'Open working directory', "", QFileDialog.ShowDirsOnly)
        self.dlg.lineEdit_outRaster.setText(self.OutputDir)
        self.dlg.button_box.setEnabled(True)
                
    def getLayers(self):
        self.dlg.comboBox_in_vector.clear()
        self.dlg.comboBox_in_raster.clear()
        
        self.dlg.comboBox_in_vector.addItems(self.vector_layers.keys())
        self.dlg.comboBox_in_raster.addItems(self.raster_layers.keys())
        
    def fillFields(self):
        self.dlg.comboBox_names.clear()
        self.selectedLayer = self.vector_layers[self.dlg.comboBox_in_vector.currentText()]
        self.fields = ["fid"] + [field.name() for field in self.selectedLayer.fields() if (field.typeName().lower()=='string') or (field.typeName().lower().startswith("int"))]
        self.dlg.comboBox_names.addItems(self.fields)        
        
    def toggleCheckBox(self, name):
        if name == "full":
            if self.dlg.checkBox_full.isChecked():
                self.dlg.lineEdit_name.setEnabled(True)
                self.dlg.lineEdit_name.setText("")
            else:
                self.dlg.lineEdit_name.setEnabled(False)
                self.dlg.lineEdit_name.setText("")
                
        elif name == "feature":
            if self.dlg.checkBox_feature.isChecked():
                self.dlg.comboBox_names.setEnabled(True)
                self.fillFields()
            else:
                self.dlg.comboBox_names.setEnabled(False)
                self.dlg.comboBox_names.clear()

    def getResults(self, layer, raster, out_folder, geot, nodata, outputNameFeatureField, outputNameFull, erased=False, split_features=False, split_full=True):
        
        if split_full:
            if not outputNameFull:
                outputNameFull = "outRaster_full"
                
            self.out_file = os.path.join(out_folder, outputNameFull + ".tif")
            
            if not erased:
                params = {
                      'INPUT': self.dlg.comboBox_in_raster.currentText(),
                      'MASK': self.dlg.comboBox_in_vector.currentText(),
                      'NODATA': None,
                      'ALPHA_BAND': False,
                      'CROP_TO_CUTLINE': True,
                      'KEEP_RESOLUTION': False,
                      'OPTIONS': 'COMPRESS=LZW',
                      'DATA_TYPE': 0,
                      'OUTPUT': self.out_file,
                }
                  
                self.result = processing.run('gdal:cliprasterbymasklayer', params)
                
                if self.dlg.checkBox_add.isChecked():
                    iface.addRasterLayer(self.out_file, outputNameFull)
                
            else:
                self.raster_layer_name = self.dlg.comboBox_in_raster.currentText()
                self.vector_layer_name = self.dlg.comboBox_in_vector.currentText()
                
                self.vectorLayer = QgsProject.instance().mapLayersByName(self.vector_layer_name)[0]
                
                self.paramsExtent = {
                    'INPUT':self.raster_layer_name,
                    'ROUND_TO':0,
                    'OUTPUT':'TEMPORARY_OUTPUT'
                }
                
                self.rasterExtent = processing.run("native:polygonfromlayerextent", self.paramsExtent)
                self.rasterExtent = self.rasterExtent["OUTPUT"]
                
                self.extentFeatures = {i.id():i.geometry() for i in self.rasterExtent.getFeatures()}
                self.vectorFeatures = [i.geometry() for i in self.vectorLayer.getFeatures()]
                
                for geomLayer in self.vectorFeatures:
                    for k, geomExtent in self.extentFeatures.items():
                        self.extentFeatures[k] = geomExtent.difference(geomLayer)
                
                self.rasterExtent.startEditing()
                for fid,geom in self.extentFeatures.items():
                    _ = self.rasterExtent.changeGeometry(fid,geom)
                self.rasterExtent.commitChanges()
                
                self.paramsClip = {
                    'INPUT': self.raster_layer_name,
                    'MASK': self.rasterExtent,
                    'ALPHA_BAND': False,
                    'CROP_TO_CUTLINE': True,
                    'KEEP_RESOLUTION': False,
                    'OPTIONS': 'COMPRESS=LZW',
                    'DATA_TYPE': 0,
                    'OUTPUT': self.out_file,
                }
                
                _ = processing.run('gdal:cliprasterbymasklayer', self.paramsClip)
                
                if self.dlg.checkBox_add.isChecked():
                    iface.addRasterLayer(self.out_file, outputNameFull)
        
        if split_features:
            if erased:
                self.raster_layer_name = self.dlg.comboBox_in_raster.currentText()
                self.vector_layer_name = self.dlg.comboBox_in_vector.currentText()
                
                self.vectorLayer = QgsProject.instance().mapLayersByName(self.vector_layer_name)[0]
                
                self.paramsExtent = {
                    'INPUT':self.raster_layer_name,
                    'ROUND_TO':0,
                    'OUTPUT':'TEMPORARY_OUTPUT'
                }
                
                self.rasterExtent = processing.run("native:polygonfromlayerextent", self.paramsExtent)
                self.rasterExtent = self.rasterExtent["OUTPUT"]
                self.extentFeatures = {i.id():i.geometry() for i in self.rasterExtent.getFeatures()}
                
                self.currentFilterText = self.vectorLayer.subsetString()
                                                
                if outputNameFeatureField != "fid":
                    self.features = [str(i[outputNameFeatureField]) for i in self.vectorLayer.getFeatures()]    
                else:
                    self.features = [str(i.id()) for i in self.vectorLayer.getFeatures()]
                    
                for name in set(self.features):
                    self.out_file = os.path.join(out_folder, name+".tif")
                    
                    if outputNameFeatureField =="fid":
                        if self.vectorLayer.storageType()=='Memory storage':
                            self.q = '$id = \'{}\''.format(name)
                        else:
                            self.q = '"fid" = \'{}\''.format(name)
                    else:
                        self.q = '"{}" = \'{}\''.format(outputNameFeatureField, name)
                        
                    self.vectorLayer.setSubsetString(self.q)
                
                    self.vectorFeatures = [i.geometry() for i in self.vectorLayer.getFeatures()]    
                    for k, geomExtent in self.extentFeatures.items():
                        self.geomExtent = QgsGeometry.fromRect(geomExtent.boundingBox())
                        for geomLayer in self.vectorFeatures:
                            self.extentFeatures[k] = self.geomExtent.difference(geomLayer)
                
                    self.rasterExtent.startEditing()
                    for fid,geom in self.extentFeatures.items():
                        _ = self.rasterExtent.changeGeometry(fid,geom)
                    self.rasterExtent.commitChanges()
                        
                    self.params = {
                        'INPUT': self.raster_layer_name,
                        'MASK': self.rasterExtent,
                        'ALPHA_BAND': False,
                        'CROP_TO_CUTLINE': True,
                        'KEEP_RESOLUTION': False,
                        'OPTIONS': 'COMPRESS=LZW',
                        'DATA_TYPE': 0,
                        'OUTPUT': self.out_file,
                    }
                    self.result = processing.run('gdal:cliprasterbymasklayer', self.params)
                
                    if self.dlg.checkBox_add.isChecked():
                        iface.addRasterLayer(self.out_file, str(name))
                
                self.vectorLayer.setSubsetString(self.currentFilterText)
            
            else:
                
                self.rasterLayer = QgsProject.instance().mapLayersByName(self.dlg.comboBox_in_raster.currentText())[0]
                self.vectorLayer = QgsProject.instance().mapLayersByName(self.dlg.comboBox_in_vector.currentText())[0]

                self.currentFilterText = self.vectorLayer.subsetString()
                                
                if outputNameFeatureField != "fid":
                    self.features = [str(i[outputNameFeatureField]) for i in self.vectorLayer.getFeatures()]    
                else:
                    self.features = [str(i.id()) for i in self.vectorLayer.getFeatures()]

                
                for name in set(self.features):
                    self.out_file = os.path.join(out_folder, name+".tif")
                    if outputNameFeatureField =="fid":
                        if self.vectorLayer.storageType()=='Memory storage':
                            self.q = '$id = \'{}\''.format(name)
                        else:
                            self.q = '"fid" = \'{}\''.format(name)
                    else:
                        self.q = '"{}" = \'{}\''.format(outputNameFeatureField, name)
                        
                    self.vectorLayer.setSubsetString(self.q)
                    
                    self.params = {
                        'INPUT': self.rasterLayer,
                        'MASK': self.vectorLayer,
                        'ALPHA_BAND': False,
                        'CROP_TO_CUTLINE': True,
                        'KEEP_RESOLUTION': False,
                        'OPTIONS': 'COMPRESS=LZW',
                        'DATA_TYPE': 0,
                        'OUTPUT': self.out_file,
                    }
                    self.result = processing.run('gdal:cliprasterbymasklayer', self.params)
                    
                    if self.dlg.checkBox_add.isChecked():
                        iface.addRasterLayer(self.out_file, str(name))
                                    
                self.vectorLayer.setSubsetString(self.currentFilterText)
                
               
    def run(self):
        """Run method that performs all the real work"""             
        # Create the dialog with elements (after translation) and keep reference
        # Only create GUI ONCE in callback, so that it will only load when the plugin is started
        if self.first_start == True:
            self.dlg = ClipRasterDialog()
            self.vector_layers = {layer.name():layer for layer in QgsProject.instance().mapLayers().values() if (layer.type()== 0 and layer.wkbType() in (3,6))}
            self.raster_layers = {layer.name():layer for layer in QgsProject.instance().mapLayers().values() if layer.type()== 1}
            self.getLayers()
            
            self.dlg.checkBox_full.toggled.connect(lambda x:self.toggleCheckBox("full"))
            self.dlg.checkBox_feature.toggled.connect(lambda x:self.toggleCheckBox("feature"))
            self.dlg.comboBox_in_vector.currentTextChanged.connect(self.fillFields)
            self.dlg.toolButton_selectOutFolder.clicked.connect(self.selectOutput)

        # show the dialog
        self.dlg.show()
        
        # Run the dialog event loop
        result = self.dlg.exec_()
        # See if OK was pressed
        if result:
            # Do something useful here - delete the line containing pass and
            # substitute with your code.
            
            self.rasterName = self.dlg.comboBox_in_raster.currentText()
            self.vectorName = self.dlg.comboBox_in_vector.currentText()
            
            self.raster_path = self.raster_layers[self.rasterName].source()
            self.vector_layer = self.vector_layers[self.vectorName]
            
            self.out_folder = self.dlg.lineEdit_outRaster.text()
            self.erased = False if self.dlg.radioButton_clip.isChecked() else True
            self.split_features = True if self.dlg.checkBox_feature.isChecked() else False
            self.split_full = True if self.dlg.checkBox_full.isChecked() else False
            self.outputNameFeatureField = self.dlg.comboBox_names.currentText()
            self.outputNameFull = self.dlg.lineEdit_name.text()             
            
            self.raster = gdal.Open(self.raster_path)
            self.geot = self.raster.GetGeoTransform()
            self.bands = [self.raster.GetRasterBand(band_no) for band_no in range(1, self.raster.RasterCount+1)]
            self.dtype = self.bands[0].DataType
            self.nodata = self.bands[0].GetNoDataValue()
            
            if self.nodata is None:
                self.arrays = [band.ReadAsArray(0,0,0,0) for band in self.bands]
                try:
                    self.nanofbands = [np.iinfo(array.dtype).min for array in self.arrays]
                    self.nodata = min(self.nanofbands)
                except:
                    self.nanofbands = [np.finfo(array.dtype).min for array in self.arrays]
                    self.nodata = min(self.nanofbands)
                        
            if self.out_folder:
                self.getResults(self.vector_layer, self.raster, self.out_folder, self.geot, self.nodata, self.outputNameFeatureField, self.outputNameFull, self.erased, self.split_features, self.split_full)                    
                self.iface.messageBar().pushMessage("Success", "Raster was clipped/erased successfully!" , level=Qgis.Success, duration=5)
