SwissMX/HelicalTable.py

import sys
import os
import logging
import csv
from PyQt5 import QtCore, QtGui

import numpy as np

logger = logging.getLogger("helical")

from PyQt5.QtWidgets import (
    QTableWidget,
    QApplication,
    QMainWindow,
    QTableWidgetItem,
    QFileDialog,
    QHeaderView,
    QAbstractItemView,
    QMenu)
from PyQt5.QtCore import Qt, QFileInfo, pyqtSignal, QPoint, pyqtSlot

DATA_ITEM = 1

START_OMEGA_0 = 0
START_OMEGA_120 = 1
START_OMEGA_240 = 2
STOP_OMEGA_0 = 3
STOP_OMEGA_120 = 4
STOP_OMEGA_240 = 5

ROLE_XTAL_START = 1 + Qt.UserRole
ROLE_XTAL_END = 2 + Qt.UserRole


class HelicalTableWidget(QTableWidget):
    gonioMoveRequest = pyqtSignal(float, float, float, float, float)
    def __init__(self):
        super().__init__()
        self.check_change = True
        self._scanAngularStep = 0.1
        self._scanHorizontalCount = 5
        self._scanVerticalCount = 10
        self._current_xtal = None
        self._start_angle = 0.0
        self.init_ui()

    def startAngle(self) -> float:
        return self._start_angle

    @pyqtSlot(float)
    def setStartAngle(self, angle:float):
        self._start_angle = angle

    def init_ui(self):
        self.setColumnCount(6)
        self.setEditTriggers(QAbstractItemView.NoEditTriggers)
        self.resizeColumnsToContents()
        self.setSelectionMode(QAbstractItemView.SingleSelection)
        self.setSelectionBehavior(QAbstractItemView.SelectItems)
        labels = [
            "0\nbx, bz\n(mm)",
            "120\nbx, bz\n(mm)",
            "240\nbx, bz\n(mm)",
            "0\nbx, bz\n(mm)",
            "120\nbx, bz\n(mm)",
            "240\nbx, bz\n(mm)",
        ]
        self.setHorizontalHeaderLabels(labels)

        self.verticalHeader().resizeSections(QHeaderView.ResizeToContents)
        self.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)

        self.cellChanged.connect(self.c_current)
        self.show()

    def display_coords(self, x, y):
        row = self.rowAt(y)
        col = self.columnAt(x)
        coords = self.gonio_coords_at(row, col)
        omega = col % 3 * 120
        data = coords[omega]
        logger.info("gonio at {}, {}: {}".format(row, col, data))

    def goto_gonio_position(self, x, y):
        row = self.rowAt(y)
        col = self.columnAt(x)
        coords = self.gonio_coords_at(row, col)
        omega = col % 3 * 120
        data = coords[omega]
        logger.info("move gonio to: {}".format(row, col, data))
        self.gonioMoveRequest.emit(*data)

    def remove_xtal(self, x, y):
        row = self.rowAt(y)
        col = self.columnAt(x)
        logger.info("rowCount() = {}, removing row: {}".format(self.rowCount(), row))
        self.removeRow(row)
        if self.rowCount() == 0:
            self._current_xtal = None
        logger.info("rowCount() = {}".format(self.rowCount()))

    def contextMenuEvent(self, event):
        logger.info(event.pos())
        x = event.pos().x()
        y = event.pos().y()
        menu = QMenu(self)
        gotoAction = menu.addAction("Move Gonio Here")
        gotoAction.triggered.connect(lambda b: self.goto_gonio_position(x, y))
        removeAction = menu.addAction("Remove this line")
        removeAction.triggered.connect(lambda b: self.remove_xtal(x, y))
        menu.popup(QtGui.QCursor.pos())

    def setScanHorizontalCount(self, count: int):
        logger.debug("horizontal count: {}".format(count))
        self._scanHorizontalCount = count

    def scanHorizontalCount(self) -> int:
        return self._scanHorizontalCount

    def setScanVerticalCount(self, count: int):
        logger.debug("vertical count: {}".format(count))
        self._scanVerticalCount = count

    def scanVerticalCount(self) -> int:
        return self._scanVerticalCount

    def setScanAngularStep(self, step: float):
        self._scanAngularStep = step

    def scanAngularStep(self) -> float:
        return self._scanAngularStep

    def scanTotalRange(self) -> float:
        return self._scanVerticalCount * self._scanHorizontalCount * self._scanAngularStep

    def gonio_coords_at(self, row, col):
        role = ROLE_XTAL_START if col < 3 else ROLE_XTAL_END
        coords = self.item(row, DATA_ITEM).data(role)
        return coords

    def add_xtal(self):
        row = self.rowCount()
        self.setRowCount(row + 1)
        self._current_xtal = row
        print(row)
        for n in range(self.columnCount()):
            self.setItem(row, n, QTableWidgetItem("unset"))


    def set_xtal_start(self, data: list):
        """sets data to start position
        
            data = [fx, fy, bx, bz, omega]
        
        """
        self.set_data_point(ROLE_XTAL_START, data)

    def set_xtal_end(self, data: list):
        self.set_data_point(ROLE_XTAL_END, data)

    def set_data_point(self, role: int, data: list):
        if self._current_xtal is None:
            self.add_xtal()
        row, col = self._current_xtal, DATA_ITEM

        try:
            coords = self.item(row, DATA_ITEM).data(role)
        except:
            print("empty coords; initializing...")
            coords = None

        if coords is None:
            coords = {}

        fx, fy, bx, bz, omega = data
        o = int(omega)
        coords[o] = data
        print("coords = {}".format(coords))

        if role == ROLE_XTAL_END:
            col = 3
        else:
            col = 0

        for n in range(3):
            omega = n * 120
            try:
                fx, fy, bx, bz, omega = coords[omega]
                info = "{bx:.3f} mm\n{bz:.3f} mm".format(bx=bx, bz=bz)
            except:
                info = "unset"
            # self.item(row, col + n).setData(Qt.DisplayRole, info)
            self.item(row, col + n).setText(info)
            self.item(row, col + n).setToolTip(info)

        self.item(row, DATA_ITEM).setData(role, coords)

    def c_current(self):
        if self.check_change:
            row = self.currentRow()
            col = self.currentColumn()
            value = self.item(row, col)
            # value = value.text()

    def load_datapoints(self):
        self.check_change = False
        path = QFileDialog.getOpenFileName(
            self, "Open CSV", os.getenv("HOME"), "CSV(*.csv)"
        )
        if path[0] != "":
            with open(path[0], newline="") as csv_file:
                self.setRowCount(0)
                self.setColumnCount(10)
                my_file = csv.reader(csv_file, delimiter=",", quotechar="|")
                for row_data in my_file:
                    row = self.rowCount()
                    self.insertRow(row)
                    if len(row_data) > 10:
                        self.setColumnCount(len(row_data))
                    for column, stuff in enumerate(row_data):
                        item = QTableWidgetItem(stuff)
                        self.setItem(row, column, item)
        self.check_change = True

    def get_data(self, as_numpy=False):
        """return a list of tuples with all defined data

            [
            (is_fiducial(boolean), x, y, gonio_x, gonio_y),...
            ]
        """
        data = []

        for row in range(self.rowCount()):
            start = self.item(row, DATA_ITEM).data(ROLE_XTAL_START)
            end = self.item(row, DATA_ITEM).data(ROLE_XTAL_END)

            data.append([start, end])
        # if as_numpy:
        #     data = np.asarray(data)
        return data


if __name__ == "__main__":

    class Sheet(QMainWindow):
        def __init__(self):
            super().__init__()

            self.form_widget = HelicalTableWidget()
            self.setCentralWidget(self.form_widget)
            self.show()

    app = QApplication(sys.argv)
    sheet = Sheet()
    sys.exit(app.exec_())