# Source code for svir.dialogs.load_gmf_data_as_layer_dialog

    # -*- coding: utf-8 -*-
    # /***************************************************************************
    # Irmt
    #                                 A QGIS plugin
    # OpenQuake Integrated Risk Modelling Toolkit
    #                              -------------------
    #        begin                : 2013-10-24
    #        copyright            : (C) 2014 by GEM Foundation
    #        email                : devops@openquake.org
    # ***************************************************************************/
    #
    # OpenQuake is free software: you can redistribute it and/or modify it
    # under the terms of the GNU Affero General Public License as published
    # by the Free Software Foundation, either version 3 of the License, or
    # (at your option) any later version.
    #
    # OpenQuake is distributed in the hope that it will be useful,
    # but WITHOUT ANY WARRANTY; without even the implied warranty of
    # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    # GNU General Public License for more details.
    #
    # You should have received a copy of the GNU Affero General Public License
    # along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

    from qgis.core import (
        QgsFeature, QgsGeometry, QgsPointXY, edit, QgsTask, QgsApplication)
    from svir.dialogs.load_output_as_layer_dialog import LoadOutputAsLayerDialog
    from svir.calculations.calculate_utils import add_numeric_attribute
    from svir.utilities.utils import WaitCursorManager, log_msg, extract_npz
    from svir.tasks.extract_npz_task import ExtractNpzTask

    class LoadGmfDataAsLayerDialog(LoadOutputAsLayerDialog):
        """
        Dialog to load gmf_data from an oq-engine output, as layer
        """

        def __init__(self, drive_engine_dlg, iface, viewer_dock, session, hostname,
                     calc_id, output_type='gmf_data', path=None, mode=None,
                     engine_version=None):
            assert output_type == 'gmf_data'
            LoadOutputAsLayerDialog.__init__(
                self, drive_engine_dlg, iface, viewer_dock, session, hostname,
                calc_id, output_type=output_type, path=path, mode=mode,
                engine_version=engine_version)

            self.setWindowTitle(
                'Load ground motion fields as layer')
            self.create_num_sites_indicator()
            # NOTE: gmpe and gsim are synonyms
            self.create_rlz_or_stat_selector('Ground Motion Prediction Equation')
            self.create_imt_selector()
            self.create_eid_selector()

            self.extract_npz_task = ExtractNpzTask(
                'Extract ground motion fields', QgsTask.CanCancel, self.session,
                self.hostname, self.calc_id, self.output_type, self.finalize_init,
                self.on_extract_error)
            QgsApplication.taskManager().addTask(self.extract_npz_task)

        def set_ok_button(self):
            self.ok_button.setEnabled(self.imt_cbx.currentIndex() != -1)

        def show_num_sites(self):
            # NOTE: we are assuming all realizations have the same number of sites,
            #       which is always true for scenario calculations.
            #       If different realizations have a different number of sites, we
            #       need to move this block of code inside on_rlz_or_stat_changed()
            rlz = self.rlz_or_stat_cbx.itemData(
                self.rlz_or_stat_cbx.currentIndex())
            gmf_data = self.npz_file[rlz]
            self.num_sites_lbl.setText(
                self.num_sites_msg % gmf_data.shape)

        def populate_rlz_or_stat_cbx(self):
            self.rlzs_or_stats = [key for key in sorted(self.npz_file)
                                  if key not in ('imtls', 'array')]
            with WaitCursorManager(
                    'Extracting...', message_bar=self.iface.messageBar()):
                self.rlzs_npz = extract_npz(
                    self.session, self.hostname, self.calc_id, 'realizations',
                    message_bar=self.iface.messageBar(), params=None)
            # rlz[1] is the branch_path field
            self.gsims = [rlz[1].decode('utf8').strip('"')
                          for rlz in self.rlzs_npz['array']]
            self.rlz_or_stat_cbx.clear()
            self.rlz_or_stat_cbx.setEnabled(True)
            for gsim, rlz in zip(self.gsims, self.rlzs_or_stats):
                # storing gsim as text, rlz as hidden data
                self.rlz_or_stat_cbx.addItem(gsim, userData=rlz)

        def on_rlz_or_stat_changed(self):
            rlz = self.rlz_or_stat_cbx.itemData(
                self.rlz_or_stat_cbx.currentIndex())
            self.dataset = self.npz_file[rlz]
            imts = self.dataset.dtype.names[2:]
            self.imt_cbx.clear()
            self.imt_cbx.setEnabled(True)
            self.imt_cbx.addItems(imts)
            self.set_ok_button()

        def on_imt_changed(self):
            imt = self.imt_cbx.currentText()
            if imt:
                min_eid = 0
                max_eid = (self.dataset[imt].shape[1] - 1)
                self.eid_sbx.cleanText()
                self.eid_sbx.setEnabled(True)
                self.eid_lbl.setText(
                    'Event ID (used for default styling) (range %d-%d)' % (
                        min_eid, max_eid))
                self.eid_sbx.setRange(min_eid, max_eid)
            self.set_ok_button()

        def load_from_npz(self):
            for rlz, gsim in zip(self.rlzs_or_stats, self.gsims):
                if (not self.load_all_rlzs_or_stats_chk.isChecked()
                        and gsim != self.rlz_or_stat_cbx.currentText()):
                    continue
                with WaitCursorManager('Creating layer for "%s"...'
                                       % gsim, self.iface.messageBar()):
                    self.build_layer(rlz_or_stat=rlz, gsim=gsim)
                    self.style_maps(self.layer, self.default_field_name,
                                    self.iface, self.output_type)
            if self.npz_file is not None:
                self.npz_file.close()

        def build_layer_name(self, gsim=None, **kwargs):
            self.imt = self.imt_cbx.currentText()
            self.eid = self.eid_sbx.value()
            self.default_field_name = '%s-%s' % (self.imt, self.eid)
            # NOTE: assuming it's a scenario calculation
            layer_name = "scenario_gmfs_%s_eid-%s" % (gsim, self.eid)
            return layer_name

        def get_field_names(self, **kwargs):
            # NOTE: we need a list instead of a tuple, because we want to be able
            #       to modify the list afterwards, to keep track of the actual
            #       field names created in the layer, that might be laundered to be
            #       compliant with shapefiles constraints
            field_names = list(self.dataset.dtype.names)
            return field_names

        def add_field_to_layer(self, field_name):
            field_name = "%s-%s" % (field_name, self.eid)
            added_field_name = add_numeric_attribute(field_name, self.layer)
            return added_field_name

        def read_npz_into_layer(self, field_names, rlz_or_stat, **kwargs):
            with edit(self.layer):
                feats = []
                fields = self.layer.fields()
                layer_field_names = [field.name() for field in fields]
                dataset_field_names = self.get_field_names()
                d2l_field_names = dict(
                    list(zip(dataset_field_names[2:], layer_field_names)))
                for row in self.npz_file[rlz_or_stat]:
                    # add a feature
                    feat = QgsFeature(fields)
                    for field_name in dataset_field_names:
                        if field_name in ['lon', 'lat']:
                            continue
                        layer_field_name = d2l_field_names[field_name]
                        value = float(row[field_name][self.eid])
                        feat.setAttribute(layer_field_name, value)
                    feat.setGeometry(QgsGeometry.fromPointXY(
                        QgsPointXY(row[0], row[1])))
                    feats.append(feat)
                added_ok = self.layer.addFeatures(feats)
                if not added_ok:
                    msg = 'There was a problem adding features to the layer.'
                    log_msg(msg, level='C', message_bar=self.iface.messageBar())