OpticsTools/OpticsTools.py

import sys
import json
import webbrowser
import subprocess
from argparse import ArgumentParser

from PyQt5 import QtWidgets,QtGui

from ui.OpticsToolsGui import Ui_OpticsGUI
from plot import OpticsPlot
from model import Model
from machine import Machine
from reference import ReferenceManager
from sandbox import Sandbox
from matchmaker import MatchMaker


class OpticsTools(QtWidgets.QMainWindow, Ui_OpticsGUI):
    def __init__(self,phase=0, office= 1):
        super(OpticsTools, self).__init__()
        self.setupUi(self)

        office = office == 1
        if phase > 0:
            office = True

        self.version = '1.1.1'
        self.setWindowIcon(QtGui.QIcon("rsc/iconoptics.png"))

        self.plot = OpticsPlot(parent=self)
        self.plot.show()

        # initialize online model
        self.model = Model(phase=phase,parent=self)

        # initialize modeling
        self.match = MatchMaker()
        self.UIMatchOpticsSelect.clear()
        for key in self.match.matchlist.keys():
            self.UIMatchOpticsSelect.addItem(key)
        self.UIMatchOpticsSelect.setCurrentIndex(0)
        self.updateMatchingCase()


        self.machine = Machine(parent = True, office = office)
        self.machine.initPVs(self.model.getElements())
        self.sandbox = Sandbox(parent = self, machine = self.machine)
        self.sandbox.updateSandbox()
        self.reference = ReferenceManager(parent=self)
        self.reference.initReferencePoints(self.match)

        title = "SwissFEL Optics Tools - Lattice %s (Phase %d)" % (self.model.getLatticeVersion(),phase)
        if office:
            title += " - offline"
        self.setWindowTitle(title)


#        events handling
        self.UIMatchOpticsSelect.currentIndexChanged.connect(self.updateMatchingCase)
        self.UIMatchSelected.clicked.connect(self.doMatch)
        self.actionOpen_2.triggered.connect(self.loadSettings)
        self.actionSave.triggered.connect(self.saveSettings)
        self.UIUpdateFromMachine.clicked.connect(self.fullUpdate)
        self.actionHelp.triggered.connect(self.openGit)
        self.actionAbout.triggered.connect(self.about)
        self.actionOpenScriptEditor.triggered.connect(self.editMatchingScript)


    def doMatch(self):
        """
        match the lattice for the given matching scripts. These can be selected by the individual check boxed.
        The matching is in the order: Injector -> Athos -> Porthos -> Aramis.
        The online model is updated with the updated match values.
        :return: None
        """
        injector = self.UIMatchInjector.isChecked()
        aramis = self.UIMatchAramis.isChecked()
        athos = self.UIMatchAthos.isChecked()
        porthos = self.UIMatchPorthos.isChecked()
        twiss = self.match.match(self.model.om,Injector = injector,Athos = athos, Aramis = aramis, Porthos = porthos)
        energy = self.model.calcEnergyProfile(twiss)
        #enfore the writing of a new lattice so that the magnet settings are in the new lattice
        self.model.forceLat = True
        self.plot.newData(twiss, energy)


    def updateMatchingCase(self):
        """
        Update the check box for selecting the different matching steps and initial settings if reference file is present
        :return: None
        """
        target = self.UIMatchOpticsSelect.currentText()
        self.match.initScripts(target)
        self.updateMatchingCaseScript(self.UIMatchInjector,self.match.scriptInjector)
        self.updateMatchingCaseScript(self.UIMatchAthos, self.match.scriptAthos)
        self.updateMatchingCaseScript(self.UIMatchAramis, self.match.scriptAramis)
        self.updateMatchingCaseScript(self.UIMatchPorthos, self.match.scriptPorthos)

    def updateMatchingCaseScript(self,widget,state):
        """
        Generalized routine to select and enable checkbox widgets
        :param widget: checkbox widget
        :param state: True or False
        :return: None
        """
        widget.setChecked(state)
        widget.setEnabled(state)

    def closeEvent(self, event):
        self.plot.close()
        event.accept()

    def about(self):
        QtWidgets.QMessageBox.about(self, "Optics Tool",
                                    "Version:%s\nContact: Sven Reiche\nEmail: sven.reiche@psi.ch" % self.version)

    def openGit(self):
        webbrowser.open("https://gitea.psi.ch/reiche/opticstool")

    def editMatchingScript(self):
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.DontUseNativeDialog
        fileName, _ = QtWidgets.QFileDialog.getOpenFileName(self, "Open Matching Script File",
                                                            "Scripts/switchyard.madx",
                                                            "MadX Files (*.madx)", options=options)
        if not fileName:
            return
        subprocess.Popen(["emacs", fileName])

##################################################
    def saveSettings(self):
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.DontUseNativeDialog
        fileName, _ = QtWidgets.QFileDialog.getSaveFileName(self, "Save Settings",
                                                            "Settings/newSetting.json",
                                                            "Json Files (*.json)", options=options)
        if not fileName:
            return
        self.saveSettingDirect(fileName)


    def loadSettings(self):
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.DontUseNativeDialog
        fileName, _ = QtWidgets.QFileDialog.getOpenFileName(self, "Open Settings",
                                                            "Settings/ReferenceSetting.json",
                                                            "Json Files (*.json)", options=options)
        if not fileName:
            return
        self.loadSettingsdirect(fileName)

    def loadSettingsdirect(self,fileName):
        with open(fileName, 'r', encoding='utf-8') as f:
            settings = json.load(f)
        self.model.loadSettings(settings)
        self.status('Reference loaded')

    def saveSettingsdirect(self,fileName):
        settings = self.model.getSettings()
        with open(fileName, 'w', encoding='utf-8') as f:
            json.dump(settings, f, ensure_ascii=False, indent=4)

####################################################
    def fullUpdate(self):
        machine = self.machine.getMachineStatus()
        self.model.updateFromMachine(machine)
        self.sandbox.updateSandbox()
        self.status('Machine Settings')

    def status(self,msg=''):
        self.UIStatus.setText(msg)
    # --------------------------------
    # Main routine


if __name__ == '__main__':
    QtWidgets.QApplication.setStyle(QtWidgets.QStyleFactory.create("plastique"))
    parser = ArgumentParser()
    parser.add_argument('-phase', type=int, help='Select Phase of the Lattice', default=0)
    parser.add_argument('-offline', type=int, help='Excludes any connection to control system', default=1)
    args = parser.parse_args()
    app = QtWidgets.QApplication(sys.argv)
    
    main = OpticsTools(phase = args.phase, office = args.offline)
    if main:
        main.show()
    sys.exit(app.exec_())