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

"""
Rast_SplitRaster.py
***************************************************************************
*                                                                         *
*   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__ = 'Leandro França'
__date__ = '2023-06-18'
__copyright__ = '(C) 2023, Leandro França'

from qgis.core import (Qgis,
                       QgsProject,
                       QgsProcessingException,
                       QgsProcessingAlgorithm,
                       QgsProcessingParameterBoolean,
                       QgsProcessingParameterFile,
                       QgsProcessingParameterNumber,
                       QgsProcessingParameterRasterLayer,
                       QgsApplication,
                       QgsProject,
                       QgsRasterLayer)

from math import floor, ceil
from osgeo import osr, gdal_array, gdal #https://gdal.org/python/
import numpy as np
from lftools.geocapt.imgs import Imgs
from lftools.translations.translate import translate
import os
from qgis.PyQt.QtGui import QIcon

class SplitRaster(QgsProcessingAlgorithm):

    LOC = QgsApplication.locale()[:2]

    def tr(self, *string):
        return translate(string, self.LOC)

    def createInstance(self):
        return SplitRaster()

    def name(self):
        return 'splitraster'

    def displayName(self):
        return self.tr('Split raster', 'Dividir raster')

    def group(self):
        return self.tr('Raster')

    def groupId(self):
        return 'raster'

    def tags(self):
        return 'GeoOne,splitting,dividir,separar,quebrar,break,retalhar,slices'.split(',')

    def icon(self):
        return QIcon(os.path.join(os.path.dirname(os.path.dirname(__file__)), 'images/raster.png'))

    txt_en = '''Splits a raster dataset into smaller pieces, by horizontal and vertical tiles.'''
    txt_pt = '''Divide um raster em pedaços menores, por blocos horizontais e verticais'''
    figure = 'images/tutorial/raster_split.jpg'

    def shortHelpString(self):
        social_BW = Imgs().social_BW
        footer = '''<div align="center">
                      <img src="'''+ os.path.join(os.path.dirname(os.path.dirname(__file__)), self.figure) +'''">
                      </div>
                      <div align="right">
                      <p align="right">
                      <b>'''+self.tr('Author: Leandro Franca', 'Autor: Leandro França')+'''</b>
                      </p>'''+ social_BW + '''</div>
                    </div>'''
        return self.tr(self.txt_en, self.txt_pt) + footer

    INPUT ='INPUT'
    VERTICAL = 'VERTICAL'
    HORIZONTAL = 'HORIZONTAL'
    OUTPUT_FOLDER = 'OUTPUT_FOLDER'
    OPEN = 'OPEN'

    def initAlgorithm(self, config=None):
        # INPUT
        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.INPUT,
                self.tr('Input Raster', 'Raster de entrada'),
                [Qgis.ProcessingSourceType.TypeRaster]
            )
        )

        self.addParameter(
            QgsProcessingParameterNumber(
                self.HORIZONTAL,
                self.tr("X axis split", 'Eixo X'),
                type = QgsProcessingParameterNumber.Type.Integer,
                minValue = 1,
                defaultValue = 2
                )
            )

        self.addParameter(
            QgsProcessingParameterNumber(
                self.VERTICAL,
                self.tr("Y axis split", 'Eixo Y'),
                type = QgsProcessingParameterNumber.Type.Integer,
                minValue = 1,
                defaultValue = 2
                )
            )

        # OUTPUT
        self.addParameter(
            QgsProcessingParameterFile(
                self.OUTPUT_FOLDER,
                self.tr('Folder for split rasters', 'Pasta para rasters divididos'),
                behavior = Qgis.ProcessingFileParameterBehavior.Folder,
                defaultValue = None
            )
        )

        self.addParameter(
            QgsProcessingParameterBoolean(
                self.OPEN,
                self.tr('Load output raster', 'Carregar imagens de saída'),
                defaultValue = True
            )
        )

    def processAlgorithm(self, parameters, context, feedback):

        RasterIN = self.parameterAsRasterLayer(
            parameters,
            self.INPUT,
            context
        )
        if RasterIN is None:
            raise QgsProcessingException(self.invalidSourceError(parameters, self.INPUT))
        nome = RasterIN.name()
        RasterIN = RasterIN.dataProvider().dataSourceUri()

        pasta_out = self.parameterAsFile(
            parameters,
            self.OUTPUT_FOLDER,
            context
        )
        if not pasta_out:
            raise QgsProcessingException(self.invalidSourceError(parameters, self.OUTPUT_FOLDER))

        n_clip_lin = self.parameterAsInt(
            parameters,
            self.VERTICAL,
            context
        )

        n_clip_col = self.parameterAsInt(
            parameters,
            self.HORIZONTAL,
            context
        )

        Carregar = self.parameterAsBool(
            parameters,
            self.OPEN,
            context
        )

        # Abrir Raster layer como array
        feedback.pushInfo(self.tr('Opening raster...', 'Abrindo raster...'))
        image = gdal.Open(RasterIN)
        prj = image.GetProjection()
        banda = image.GetRasterBand(1).ReadAsArray()
        Pixel_Nulo = image.GetRasterBand(1).GetNoDataValue()
        n_bands = image.RasterCount
        # Origem e resolucao da imagem
        geotransform  = image.GetGeoTransform()
        ulx, xres, xskew, uly, yskew, yres  = geotransform
        origem = [ulx, uly]
        # Pegar tipo de dado do array
        GDT = gdal_array.NumericTypeCodeToGDALTypeCode(banda.dtype)


        # Retalhar imagem
        n_lin = image.RasterYSize # numero de linhas
        n_col = image.RasterXSize # numero de colunas
        linhas = np.round(np.linspace(0, n_lin, n_clip_lin+1))
        colunas = np.round(np.linspace(0, n_col, n_clip_col+1))
        lista = []
        for i in range(n_clip_lin):
            for j in range(n_clip_col):
                new_name = nome + '_' + '{:02}'.format(j+1) +'_' + '{:02}'.format(i+1) + '.tif'
                feedback.pushInfo(self.tr('Exporting file {} ...', 'Exportando arquivo {} ...').format(new_name))
                Output = os.path.join(pasta_out, new_name)
                cols = colunas[j+1] - colunas[j]
                rows = linhas[i+1] - linhas[i]
                ulx = origem[0] + abs(xres)*colunas[j]
                uly = origem[1] - abs(yres)*linhas[i]
                geotransform = [ulx, xres, 0, uly, 0, yres]
                # Arquivo de saída
                Driver = gdal.GetDriverByName('GTiff').Create(Output, int(cols), int(rows), n_bands, GDT)
                Driver.SetGeoTransform(geotransform)
                Driver.SetProjection(prj)

                for k in range(n_bands):
                    banda = image.GetRasterBand(k+1).ReadAsArray()
                    split_banda = banda[int(linhas[i]):int(linhas[i+1]), int(colunas[j]):int(colunas[j+1])]
                    outband = Driver.GetRasterBand(k+1)
                    feedback.pushInfo(self.tr('Writing Band {}...'.format(k+1), 'Escrevendo Banda {}...'.format(k+1)))
                    outband.WriteArray(split_banda)
                    if Pixel_Nulo or Pixel_Nulo == 0:
                        outband.SetNoDataValue(Pixel_Nulo)
                    del banda, split_banda
                lista += [Output]
                Driver.FlushCache() # write to disk
                Driver = None  # save, close
        image = None # Close dataset

        feedback.pushInfo(self.tr('Operation completed successfully!', 'Operação finalizada com sucesso!'))
        feedback.pushInfo(self.tr('Leandro Franca - Cartographic Engineer', 'Leandro França - Eng Cart'))
        self.CAMINHO = lista
        self.CARREGAR = Carregar
        return {self.OUTPUT_FOLDER: pasta_out}

    def postProcessAlgorithm(self, context, feedback):
        if self.CARREGAR:
            for file_path in self.CAMINHO:
                layer_name = os.path.basename(file_path)[:-4]
                rlayer = QgsRasterLayer(file_path, layer_name)
                QgsProject.instance().addMapLayer(rlayer)
        return {}
