'''
/***************************************************************************
Name	   :  import_vector_to_raster_lgp_controller.py
Description:  Import vector to raster LGP class for FEWSTools plugin
copyright  :  (C) 2019-2023 by FEWS
email      :  dhackman@contractor.usgs.gov
Created    :  3/29/2023 - dhackman
Modified   :  
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   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 osgeo import ogr

from PyQt5.QtWidgets import QMessageBox, QDialog, QFileDialog
from qgis.core import Qgis, QgsMessageLog

from fews_tools import fews_tools_config as config
from fews_tools.utilities import geoclim_gdal_utilities as g_util
from fews_tools.utilities import geoclim_qgs_utilities as qgs_util

from fews_tools.models.region_model import RegionModel
from fews_tools.models.workspace_setup_model import WorkspaceSetupModel
from fews_tools.forms.Ui_ImportVectorToRasterLGP import Ui_ImportVectorToRasterLGP


class ImportVectorToRasterControllerLGP(QDialog):
    '''
    Class for getting a LGP raster from a .shp file
    '''
    def __init__(self, reg_dic):
        QDialog.__init__(self)
        self.ui = Ui_ImportVectorToRasterLGP()
        self.ui.setupUi(self)
        self.ui.closeButton.setFocus()
        self.wrksp_setup = WorkspaceSetupModel()
        self.reg_dic = reg_dic
        self.ui.browseVectorFileButton.clicked.connect(self.browse_shp_path)
        self.ui.browseBilFileButton.clicked.connect(self.browse_outfile_path)
        self.ui.applyChangesToProgramCheckBox.setVisible(False)
        self.ui.joinVectorFiletoSecondaryDataTableCheckBox.setEnabled(False)
        self.ui.okButton.clicked.connect(self.import_to_raster)
        self.ui.vectorInLineEdit.textChanged.connect(self.enable_ok_button)
        self.ui.bilOutLineEdit.textChanged.connect(self.enable_ok_button)
        self.ui.okButton.setEnabled(False)
        self.raster_filename = ''
        self.init_value = 0


    def browse_shp_path(self):
        '''
        Browse to vector files.
        '''
        vector_file =\
            QFileDialog.getOpenFileName(self,
                                        'Please Select the Vector File',
                                        self.wrksp_setup.get_map_data_path(),
                                        config.VEC_FILE_TYPES)
        if vector_file[0]:
            self.ui.vectorInLineEdit.setText(
                self.wrksp_setup.fix_os_sep_in_path(vector_file[0]))

    def browse_outfile_path(self):
        '''
        Browse to bil output file, the name that is filled in will be the
        basename of the shp file selected as a .bil/.tif/.tiff extension.
        If no shp file has been selected the filename will default to 
        "outfile_rasterized.bil"
        '''
        if self.ui.vectorInLineEdit.text() != '':
            vector_file_basename = os.path.basename(
                self.ui.vectorInLineEdit.text().split('.')[0] + "_raster")
        else:
            vector_file_basename = 'outfile_rasterized'
        raster_file =\
            QFileDialog.getSaveFileName(
                self,
                u'Raster file location',
                os.path.join(self.wrksp_setup.get_static_data_path(),
                             vector_file_basename),
                config.RST_FILE_TYPES)
        self.ui.bilOutLineEdit.setText(raster_file[0])

    def enable_ok_button(self):
        '''
        Enables the ok button if widgets are ready
        '''
        if self.ui.vectorInLineEdit.text() and self.ui.bilOutLineEdit.text():
            self.ui.okButton.setEnabled(True)
        
        if self.ui.vectorInLineEdit.text():
            self.update_dropdown()

    def update_dropdown(self):
        '''
        Call from in the enable_ok_button
        updates the dropdown from given shapefile
        '''
        self.ui.dataFieldComboBox.clear()

        vector_file = self.ui.vectorInLineEdit.text()
        source = ogr.Open(vector_file)
        layer = source.GetLayer()
        fields = []
        ldefn = layer.GetLayerDefn()
        for n in range(ldefn.GetFieldCount()):
            fdefn = ldefn.GetFieldDefn(n)
            fields.append(fdefn.name)
        
        self.ui.dataFieldComboBox.addItems(fields)

    def get_raster_filename(self):
        '''
        Getter for raster file name
        '''
        return self.raster_filename

    def set_raster_filename(self):
        '''
        Setter for raster file name
        '''
        self.raster_filename = self.ui.bilOutLineEdit.text()

    def import_to_raster(self):
        '''
        Import the vector to the raster, using the default region extents
        and cell size.
        This will not overwrite an existing raster file.
        '''
        attribute = self.ui.dataFieldComboBox.currentText()
        vector_file = self.ui.vectorInLineEdit.text()
        self.set_raster_filename()
        reg_extent = qgs_util.get_region_extent(self.reg_dic)
        shp_extent = qgs_util.extract_shapefile_extents_rectangle(vector_file)
        # this mask extent must be made slightly larger or the mask
        # will flag an error.  Mask must be slightly larger than it's region
        fudge = self.reg_dic['Height'] / 2
        match_rect = qgs_util.rectangle_equals(reg_extent, shp_extent)
        # test if region and shapefile are the same or if shp is within region
        if reg_extent.contains(shp_extent) or match_rect:
            if not os.path.exists(self.raster_filename):
                if vector_file and self.raster_filename and\
                        os.path.exists(vector_file):
                    g_util.rasterize_vector_by_attribute(
                        vector_file, self.raster_filename, self.reg_dic['Height'],
                        reg_extent, self.init_value, attribute,
                        fudge, d_type="UInt16")

                    if os.path.exists(self.raster_filename):
                        qgs_util.display_raster_layer(self.raster_filename, 'raster')
                        QgsMessageLog.logMessage(
                            'Import LGP from Vector succeeded', level=Qgis.Info)
                    else:
                        QMessageBox.critical(
                            self,
                            'ERROR: Import LGP from Vector did not succeed',
                            'ERROR: Import LGP from Vector did not succeed',
                            QMessageBox.Ok)
                        QgsMessageLog.logMessage('Import LGP from Vector failed',
                                                level=Qgis.Critical)
                else:
                    QMessageBox.critical(
                        self,
                        'ERROR: Import LGP from Vector did not succeed',
                        'ERROR: Import LGP from Vector did not succeed',
                        QMessageBox.Ok)
                    QgsMessageLog.logMessage('Import LGP from Vector failed',
                                             level=Qgis.Critical)
            else:
                QMessageBox.information(
                    self, 'LGP file already exists',
                    'LGP file already exists!  Exiting!', QMessageBox.Ok)
        else:
            QMessageBox.information(
                self, 'Region extents do not match with shapefile',
                'Region vs shapefile mismatch!  Exiting!', QMessageBox.Ok)
        self.close()
