from PyQt5.QtCore import QObject
from PyQt5.QtGui import QStandardItemModel

from xplordb.datamodel.survey import Survey


class SurveyCreationLayerModel(QStandardItemModel):
    DEPTH_COL = 0
    DIP_COL = 1
    AZIMUTH_COL = 2

    def __init__(self, parent: QObject = None):
        """
        QStandardItemModel for survey creation.

        Args:
            parent: QObject parent
        """
        super(SurveyCreationLayerModel, self).__init__(parent)
        self.setHorizontalHeaderLabels([self.tr("Depth"), self.tr("Dip"), self.tr("Azimuth")])

    def get_survey_list(self) -> [Survey]:
        """
        Get survey list (warning no person and dataset are defined)

        Returns: survey list (warning no person and dataset are defined)

        """
        result = []
        for i in range(0, self.rowCount()):
            result.append(self._survey(i))
        return result

    def add_survey(self, survey: Survey = None):
        """
        Add row in model from Survey

        Args:
            survey: Survey
        """
        self.insertRow(self.rowCount())
        if survey:
            self._set_survey_data(self.rowCount()-1, survey)

    def _survey(self, row: int) -> Survey:
        """
        Get survey definition for specific row

        Args:
            row: model row

        Returns: Survey for specific row

        """
        survey = Survey(hole_id="",  # Hole ID must be defined by widget using model
                        data_set="",  # Data set must be defined by widget using model
                        loaded_by="",  # Loaded by person must be defined by widget using model
                        dip=self.data(self.index(row, self.DIP_COL)),
                        azimuth=self.data(self.index(row, self.AZIMUTH_COL)),
                        depth=self.data(self.index(row, self.DEPTH_COL)))
        return survey

    def _set_survey_data(self, row: int, survey: Survey) -> None:
        """
        Define QStandardItemModel data from Survey

        Args:
            row: model row
            survey: Survey
        """
        self.setData(self.index(row, self.DIP_COL), survey.dip)
        self.setData(self.index(row, self.AZIMUTH_COL), survey.azimuth)
        self.setData(self.index(row, self.DEPTH_COL), survey.depth)
