widget_ao.py

#!/usr/bin/env python
 
"""
Widget to simulate a closed loop
 
Usage:
  widget_ao.py [<parameters_filename>] [options]
 
with 'parameters_filename' the path to the parameters file
 
Options:
  -h --help          Show this help message and exit
  --cacao            Distribute data with cacao
  --expert           Display expert panel
  -d, --devices devices      Specify the devices
  -i, --interactive  keep the script interactive
"""
 
import os, sys
 
import numpy as np
import time
 
import pyqtgraph as pg
 
from shesha.util.tools import plsh, plpyr
from shesha.config import ParamConfig
 
try:
    from PyQt5 import QtWidgets
    from PyQt5.QtCore import Qt
except ModuleNotFoundError as e:
    try:    
        from PySide2 import QtWidgets
        from PySide2.QtCore import Qt
    except ModuleNotFoundError as e:
        raise ModuleNotFoundError("No module named 'PyQt5' or PySide2', please install one of them\nException raised: "+e.msg)
 
from typing import Any, Dict, Tuple
 
from docopt import docopt
from collections import deque
 
from shesha.widgets.widget_base import WidgetBase, uiLoader
 
AOWindowTemplate, AOClassTemplate = uiLoader('widget_ao')
 
from shesha.supervisor.compassSupervisor import CompassSupervisor, scons
 
# For debug
# from IPython.core.debugger import Pdb
# then add this line to create a breakpoint
# Pdb().set_trace()
 
 
class widgetAOWindow(AOClassTemplate, WidgetBase):
 
    def __init__(self, config_file: Any = None, cacao: bool = False,
                 expert: bool = False, devices: str = None,
                 hide_histograms: bool = False, twoStages: bool = False) -> None:
        WidgetBase.__init__(self, hide_histograms=hide_histograms)
        AOClassTemplate.__init__(self)
        self.twoStagestwoStages = twoStages
        self.cacaocacao = cacao
        self.rollingWindowrollingWindow = 100
        self.SRLESRLE = deque(maxlen=self.rollingWindowrollingWindow)
        self.SRSESRSE = deque(maxlen=self.rollingWindowrollingWindow)
        self.numiternumiter = deque(maxlen=self.rollingWindowrollingWindow)
        self.expertexpert = expert
        self.devicesdevices = devices
 
        self.uiAOuiAO = AOWindowTemplate()
        self.uiAOuiAO.setupUi(self)
 
        
 
        self.supervisorsupervisor = None
        self.configconfig = None
        self.stopstop = False  # type: bool  # Request quit
 
        self.uiAOuiAO.wao_nbiters.setValue(1000)  # Default GUI nIter box value
        self.nbiternbiter = self.uiAOuiAO.wao_nbiters.value()
        self.refreshTimerefreshTime = 0  # type: float  # System time at last display refresh
        self.assistantassistant = None  # type: Any
 
        
 
        # Default path for config files
        self.defaultParPathdefaultParPathdefaultParPath = os.environ["SHESHA_ROOT"] + "/data/par/par4bench"
        self.defaultAreaPathdefaultAreaPathdefaultAreaPath = os.environ["SHESHA_ROOT"] + "/data/layouts"
        self.loadDefaultConfigloadDefaultConfig()
 
        self.uiAOuiAO.wao_run.setCheckable(True)
        self.uiAOuiAO.wao_run.clicked[bool].connect(self.aoLoopClickedaoLoopClicked)
        self.uiAOuiAO.wao_open_loop.setCheckable(True)
        self.uiAOuiAO.wao_open_loop.clicked[bool].connect(self.aoLoopOpenaoLoopOpen)
        self.uiAOuiAO.wao_next.clicked.connect(self.loop_onceloop_once)
        self.uiAOuiAO.wao_resetSR.clicked.connect(self.resetSRresetSR)
        self.uiAOuiAO.wao_resetCoro.clicked.connect(self.resetCororesetCoro)
        # self.uiAO.wao_actionHelp_Contents.triggered.connect(self.on_help_triggered)
 
        self.uiAOuiAO.wao_allTarget.stateChanged.connect(self.updateAllTargetupdateAllTarget)
        self.uiAOuiAO.wao_allCoro.stateChanged.connect(self.updateAllCoroupdateAllCoro)
        self.uiAOuiAO.wao_forever.stateChanged.connect(self.updateForeverupdateForever)
 
        self.uiAOuiAO.wao_atmosphere.clicked[bool].connect(self.enable_atmosenable_atmos)
        self.dispStatsInTerminaldispStatsInTerminal = False
        self.uiAOuiAO.wao_clearSR.clicked.connect(self.clearSRclearSR)
        # self.uiAO.actionStats_in_Terminal.toggled.connect(self.updateStatsInTerminal)
 
        self.uiAOuiAO.wao_run.setDisabled(True)
        self.uiAOuiAO.wao_next.setDisabled(True)
        self.uiAOuiAO.wao_unzoom.setDisabled(True)
        self.uiAOuiAO.wao_resetSR.setDisabled(True)
        self.uiAOuiAO.wao_resetCoro.setDisabled(True)
        self.uiAOuiAO.wao_allCoro.setDisabled(True)
 
        p1 = self.uiAOuiAO.wao_SRPlotWindow.addPlot(title='SR evolution')
        self.curveSRSEcurveSRSE = p1.plot(pen=(255, 0, 0), symbolBrush=(255, 0, 0), name="SR SE")
        self.curveSRLEcurveSRLE = p1.plot(pen=(0, 0, 255), symbolBrush=(0, 0, 255), name="SR LE")
 
        self.SRCrossXSRCrossX = {}  # type: Dict[str, pg.ScatterPlotItem]
        self.SRCrossYSRCrossY = {}  # type: Dict[str, pg.ScatterPlotItem]
        self.SRcirclesSRcircles = {}  # type: Dict[str, pg.ScatterPlotItem]
        self.PyrEdgeXPyrEdgeX = {}  # type: Dict[str, pg.ScatterPlotItem]
        self.PyrEdgeYPyrEdgeY = {}  # type: Dict[str, pg.ScatterPlotItem]
 
        self.natmnatm = 0
        self.nwfsnwfs = 0
        self.ndmndm = 0
        self.ntarntar = 0
        self.PSFzoomPSFzoom = 50
        self.firstTimefirstTime = 1
        self.uiAOuiAO.wao_SRDock.setVisible(False)
 
        self.addDockWidget(Qt.DockWidgetArea(1), self.uiBaseuiBase.wao_ConfigDock)
        self.addDockWidget(Qt.DockWidgetArea(1), self.uiBaseuiBase.wao_DisplayDock)
        self.uiBaseuiBase.wao_ConfigDock.setFloating(False)
        self.uiBaseuiBase.wao_DisplayDock.setFloating(False)
 
        if expert:
            from shesha.widgets.widget_ao_expert import WidgetAOExpert
            self.expertWidgetexpertWidget = WidgetAOExpert()
            # self.expertWidget.setupUi(self)
            self.addDockWidget(
                    Qt.DockWidgetArea(1), self.expertWidgetexpertWidget.uiExpert.wao_expertDock)
            self.expertWidgetexpertWidget.uiExpert.wao_expertDock.setFloating(False)
 
        self.adjustSize()
 
        if config_file is not None:
            self.uiBaseuiBase.wao_selectConfig.clear()
            self.uiBaseuiBase.wao_selectConfig.addItem(config_file)
            self.load_configload_configload_config(config_file=config_file)
            self.init_configinit_configinit_config()
 
    # def on_help_triggered(self, i: Any=None) -> None:
    #     if i is None:
    #         return
    #     if not self.assistant or \
    #        not self.assistant.poll():
 
    #         helpcoll = os.environ["COMPASS_ROOT"] + "/doc/COMPASS.qhc"
    #         cmd = "assistant -enableRemoteControl -collectionFile %s" % helpcoll
    #         self.assistant = Popen(cmd, shell=True, stdin=PIPE)
 
    
 
    # def updateStatsInTerminal(self, state):
    #     self.dispStatsInTerminal = state
 
    def updateAllTarget(self, state):
        self.uiAOuiAO.wao_resetSR_tarNum.setDisabled(state)
 
    def updateAllCoro(self, state):
        self.uiAOuiAO.wao_resetCoro_coroNum.setDisabled(state)
 
    def updateForever(self, state):
        self.uiAOuiAO.wao_nbiters.setDisabled(state)
 
    def enable_atmos(self, atmos):
        self.supervisorsupervisor.atmos.enable_atmos(atmos)
 
    def resetSR(self) -> None:
        if self.uiAOuiAO.wao_allTarget.isChecked():
            for t in range(len(self.configconfig.p_targets)):
                self.supervisorsupervisor.target.reset_strehl(t)
        else:
            tarnum = self.uiAOuiAO.wao_resetSR_tarNum.value()
            print("Reset SR on target %d" % tarnum)
            self.supervisorsupervisor.target.reset_strehl(tarnum)
 
    def resetCoro(self) -> None:
        # TODO Adapt for multiple corono 
        if self.uiAOuiAO.wao_allCoro.isChecked():
            for c in range(self.ncoroncoro):
                self.supervisorsupervisor.corono.reset()
        else:
            coroNum = self.uiAOuiAO.wao_resetCoro_coroNum.value()
            print("Reset Coro %d" % coroNum)
            self.supervisorsupervisor.corono.reset()
 
    def add_dispDock(self, name: str, parent, type: str = "pg_image") -> None:
        d = WidgetBase.add_dispDock(self, name, parent, type)
        if type == "SR":
            d.addWidget(self.uiAOuiAO.wao_Strehl)
 
    def load_config(self, *args, config_file=None, supervisor=None, **kwargs) -> None:
        '''
            Callback when 'LOAD' button is hit
            * required to catch positionals, as by default
            if a positional is allowed the QPushButton will send a boolean value
            and hence overwrite supervisor...
        '''
 
        WidgetBase.load_config(self)
        for key, pgpl in self.SRcirclesSRcircles.items():
            self.viewboxesviewboxes[key].removeItem(pgpl)
 
        for key, pgpl in self.SRCrossXSRCrossX.items():
            self.viewboxesviewboxes[key].removeItem(pgpl)
 
        for key, pgpl in self.SRCrossYSRCrossY.items():
            self.viewboxesviewboxes[key].removeItem(pgpl)
 
        for key, pgpl in self.PyrEdgeXPyrEdgeX.items():
            self.viewboxesviewboxes[key].removeItem(pgpl)
 
        for key, pgpl in self.PyrEdgeYPyrEdgeY.items():
            self.viewboxesviewboxes[key].removeItem(pgpl)
 
        if config_file is None:
            config_file = str(self.uiBaseuiBase.wao_selectConfig.currentText())
            sys.path.insert(0, self.defaultParPathdefaultParPathdefaultParPath)
 
        if supervisor is None:
            self.configconfig = ParamConfig(config_file)
        else:
            self.configconfig = supervisor.get_config()
 
        if self.devicesdevices:
            self.configconfig.p_loop.set_devices([
                    int(device) for device in self.devicesdevices.split(",")
            ])
 
        try:
            sys.path.remove(self.defaultParPathdefaultParPathdefaultParPath)
        except:
            pass
 
        self.SRcirclesSRcircles.clear()
        self.SRCrossXSRCrossX.clear()
        self.SRCrossYSRCrossY.clear()
        self.PyrEdgeXPyrEdgeX.clear()
        self.PyrEdgeYPyrEdgeY.clear()
        self.nctrlnctrl = len(self.configconfig.p_controllers)
        for ctrl in range(self.nctrlnctrl):
            self.add_dispDockadd_dispDockadd_dispDock("modalGains_"+str(ctrl), self.wao_graphgroup_cbwao_graphgroup_cb, "MPL")
 
        self.natmnatm = len(self.configconfig.p_atmos.alt)
        for atm in range(self.natmnatm):
            name = 'atm_%d' % atm
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_phasesgroup_cbwao_phasesgroup_cb)
 
        self.nwfsnwfs = len(self.configconfig.p_wfss)
        for wfs in range(self.nwfsnwfs):
            name = 'wfs_%d' % wfs
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_phasesgroup_cbwao_phasesgroup_cb)
            name = 'slpComp_%d' % wfs
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_graphgroup_cbwao_graphgroup_cb, "MPL")
            name = 'slpGeom_%d' % wfs
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_graphgroup_cbwao_graphgroup_cb, "MPL")
            if self.configconfig.p_wfss[wfs].type == scons.WFSType.SH:
                name = 'SH_%d' % wfs
                self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
            elif self.configconfig.p_wfss[
                    wfs].type == scons.WFSType.PYRHR or self.configconfig.p_wfss[
                            wfs].type == scons.WFSType.PYRLR:
                name = 'pyrFocalPlane_%d' % wfs
                self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
                name = 'pyrHR_%d' % wfs
                self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
                name = 'pyrLR_%d' % wfs
                self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
            else:
                raise "Analyser unknown"
 
        self.ndmndm = len(self.configconfig.p_dms)
        for dm in range(self.ndmndm):
            name = 'dm_%d' % dm
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_phasesgroup_cbwao_phasesgroup_cb)
 
        self.ntarntar = len(self.configconfig.p_targets)
        for tar in range(self.ntarntar):
            name = 'tar_%d' % tar
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_phasesgroup_cbwao_phasesgroup_cb)
        for tar in range(self.ntarntar):
            name = 'psfSE_%d' % tar
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
        for tar in range(self.ntarntar):
            name = 'psfLE_%d' % tar
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
        if(self.configconfig.p_coronos) is not None:
            self.ncoroncoro = len(self.configconfig.p_coronos)
        else:
            self.ncoroncoro = 0
        for coro in range(self.ncoroncoro):
            name = 'coroImageLE_%d' % coro
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
        for coro in range(self.ncoroncoro):
            name = 'coroImageSE_%d' % coro
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
        for coro in range(self.ncoroncoro):
            name = 'coroPSFLE_%d' % coro
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
        for coro in range(self.ncoroncoro):
            name = 'coroPSFSE_%d' % coro
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
            
        self.add_dispDockadd_dispDockadd_dispDock("Strehl", self.wao_graphgroup_cbwao_graphgroup_cb, "SR")
        for coro in range(self.ncoroncoro):
            self.add_dispDockadd_dispDockadd_dispDock(f"ContrastLE_{coro}", self.wao_graphgroup_cbwao_graphgroup_cb, "MPL")
            self.add_dispDockadd_dispDockadd_dispDock(f"ContrastSE_{coro}", self.wao_graphgroup_cbwao_graphgroup_cb, "MPL")
 
        self.uiAOuiAO.wao_resetSR_tarNum.setValue(0)
        self.uiAOuiAO.wao_resetSR_tarNum.setMaximum(len(self.configconfig.p_targets) - 1)
 
        self.uiAOuiAO.wao_resetCoro_coroNum.setValue(0)
        self.uiAOuiAO.wao_resetCoro_coroNum.setMaximum(self.ncoroncoro - 1)
 
        self.uiAOuiAO.wao_dispSR_tar.setValue(0)
        self.uiAOuiAO.wao_dispSR_tar.setMaximum(len(self.configconfig.p_targets) - 1)
 
        self.uiAOuiAO.wao_run.setDisabled(True)
        self.uiAOuiAO.wao_next.setDisabled(True)
        self.uiAOuiAO.wao_unzoom.setDisabled(True)
        self.uiAOuiAO.wao_resetSR.setDisabled(True)
        self.uiAOuiAO.wao_resetCoro.setDisabled(True)
        self.uiAOuiAO.wao_allCoro.setDisabled(True)
 
        self.uiBaseuiBase.wao_init.setDisabled(False)
 
        if self.expertexpert:
            self.expertWidgetexpertWidget.setSupervisor(self.supervisorsupervisor)
            self.expertWidgetexpertWidget.updatePanels()
 
        if (hasattr(self.configconfig._config, "layout")):
            area_filename = self.defaultAreaPathdefaultAreaPathdefaultAreaPath + "/" + self.configconfig._config.layout + ".area"
            self.loadArealoadArea(filename=area_filename)
 
        self.adjustSize()
 
    def aoLoopClicked(self, pressed: bool) -> None:
        if pressed:
            self.stopstop = False
            self.refreshTimerefreshTime = time.time()
            self.nbiternbiter = self.uiAOuiAO.wao_nbiters.value()
            if self.dispStatsInTerminaldispStatsInTerminal:
                if self.uiAOuiAO.wao_forever.isChecked():
                    print("LOOP STARTED")
                else:
                    print("LOOP STARTED FOR %d iterations" % self.nbiternbiter)
            self.runrunrun()
        else:
            self.stopstop = True
 
    def aoLoopOpen(self, pressed: bool) -> None:
        if (pressed):
            self.supervisorsupervisor.rtc.close_loop()
            self.uiAOuiAO.wao_open_loop.setText("Open Loop")
        else:
            self.supervisorsupervisor.rtc.open_loop()
            self.uiAOuiAO.wao_open_loop.setText("Close Loop")
 
    def init_config(self) -> None:
        if(self.twoStagestwoStages):
            from shesha.supervisor.stageSupervisor import StageSupervisor, scons
            self.supervisorsupervisor = StageSupervisor(self.configconfig, cacao=self.cacaocacao)
        else:
            self.supervisorsupervisor = CompassSupervisor(self.configconfig, cacao=self.cacaocacao)
        WidgetBase.init_config(self)
 
    def init_configThread(self) -> None:
        self.uiAOuiAO.wao_deviceNumber.setDisabled(True)
 
    def init_configFinished(self) -> None:
        # Thread carmaWrap context reload:
        self.supervisorsupervisor.force_context()
 
        for i in range(self.natmnatm):
            key = "atm_%d" % i
            data = self.supervisorsupervisor.atmos.get_atmos_layer(i)
            cx, cy = self.circleCoordscircleCoords(self.configconfig.p_geom.pupdiam / 2, 1000,
                                       data.shape[0], data.shape[1])
            self.SRcirclesSRcircles[key] = pg.ScatterPlotItem(cx, cy, pen='r', size=1)
            self.viewboxesviewboxes[key].addItem(self.SRcirclesSRcircles[key])
            self.SRcirclesSRcircles[key].setData(cx, cy)
 
        for i in range(self.nwfsnwfs):
            key = "wfs_%d" % i
            data = self.supervisorsupervisor.wfs.get_wfs_phase(i)
            cx, cy = self.circleCoordscircleCoords(self.configconfig.p_geom.pupdiam / 2, 1000,
                                       data.shape[0], data.shape[1])
            self.SRcirclesSRcircles[key] = pg.ScatterPlotItem(cx, cy, pen='r', size=1)
            self.viewboxesviewboxes[key].addItem(self.SRcirclesSRcircles[key])
            self.SRcirclesSRcircles[key].setData(cx, cy)
            key = 'slpComp_%d' % i
            key = 'slpGeom_%d' % i
 
            # if self.config.p_wfss[i].type == scons.WFSType.SH:
            #     key = "SH_%d" % i
            #     self.addSHGrid(self.docks[key].widgets[0],
            #                    self.config.p_wfss[i].get_validsub(), 8, 8)
 
        for i in range(self.ndmndm):
            key = "dm_%d" % i
            dm_type = self.configconfig.p_dms[i].type
            alt = self.configconfig.p_dms[i].alt
            data = self.supervisorsupervisor.dms.get_dm_shape(i)
            cx, cy = self.circleCoordscircleCoords(self.configconfig.p_geom.pupdiam / 2, 1000,
                                       data.shape[0], data.shape[1])
            self.SRcirclesSRcircles[key] = pg.ScatterPlotItem(cx, cy, pen='r', size=1)
            self.viewboxesviewboxes[key].addItem(self.SRcirclesSRcircles[key])
            self.SRcirclesSRcircles[key].setData(cx, cy)
 
        for i in range(len(self.configconfig.p_targets)):
            key = "tar_%d" % i
            data = self.supervisorsupervisor.target.get_tar_phase(i)
            cx, cy = self.circleCoordscircleCoords(self.configconfig.p_geom.pupdiam / 2, 1000,
                                       data.shape[0], data.shape[1])
            self.SRcirclesSRcircles[key] = pg.ScatterPlotItem(cx, cy, pen='r', size=1)
            self.viewboxesviewboxes[key].addItem(self.SRcirclesSRcircles[key])
            self.SRcirclesSRcircles[key].setData(cx, cy)
 
            data = self.supervisorsupervisor.target.get_tar_image(i)
            for psf in ["psfSE_", "psfLE_"]:
                key = psf + str(i)
                Delta = 5
                self.SRCrossXSRCrossX[key] = pg.PlotCurveItem(
                        np.array([
                                data.shape[0] / 2 + 0.5 - Delta,
                                data.shape[0] / 2 + 0.5 + Delta
                        ]), np.array([data.shape[1] / 2 + 0.5, data.shape[1] / 2 + 0.5]),
                        pen='r')
                self.SRCrossYSRCrossY[key] = pg.PlotCurveItem(
                        np.array([data.shape[0] / 2 + 0.5, data.shape[0] / 2 + 0.5]),
                        np.array([
                                data.shape[1] / 2 + 0.5 - Delta,
                                data.shape[1] / 2 + 0.5 + Delta
                        ]), pen='r')
                # Put image in plot area
                self.viewboxesviewboxes[key].addItem(self.SRCrossXSRCrossX[key])
                # Put image in plot area
                self.viewboxesviewboxes[key].addItem(self.SRCrossYSRCrossY[key])
 
 
        for i in range(self.ncoroncoro):
            data = self.supervisorsupervisor.corono.get_image(i, expo_type="se")
            for psf in ["coroImageSE_", "coroImageLE_", "coroPSFSE_", "coroPSFLE_"]:
                key = psf + str(i)
                Delta = 5
                if("Image" in key):
                    center = 0.
                else:
                    center = 0.5
                self.SRCrossXSRCrossX[key] = pg.PlotCurveItem(
                        np.array([
                                data.shape[0] / 2 + center - Delta,
                                data.shape[0] / 2 + center + Delta
                        ]), np.array([data.shape[1] / 2 + center, data.shape[1] / 2 + center]),
                        pen='r')
                self.SRCrossYSRCrossY[key] = pg.PlotCurveItem(
                        np.array([data.shape[0] / 2 + center, data.shape[0] / 2 + center]),
                        np.array([
                                data.shape[1] / 2 + center - Delta,
                                data.shape[1] / 2 + center + Delta
                        ]), pen='r')
                # Put image in plot area
                self.viewboxesviewboxes[key].addItem(self.SRCrossXSRCrossX[key])
                # Put image in plot area
                self.viewboxesviewboxes[key].addItem(self.SRCrossYSRCrossY[key])
 
 
 
        for i in range(len(self.configconfig.p_wfss)):
            if (self.configconfig.p_wfss[i].type == scons.WFSType.PYRHR or
                        self.configconfig.p_wfss[i].type == scons.WFSType.PYRLR):
                key = "pyrFocalPlane_%d" % i
                data = self.supervisorsupervisor.wfs.get_pyr_focal_plane(i)
                Delta = len(data) / 2
                self.PyrEdgeXPyrEdgeX[key] = pg.PlotCurveItem(
                        np.array([
                                data.shape[0] / 2 + 0.5 - Delta,
                                data.shape[0] / 2 + 0.5 + Delta
                        ]), np.array([data.shape[1] / 2 + 0.5, data.shape[1] / 2 + 0.5]),
                        pen='b')
                self.PyrEdgeYPyrEdgeY[key] = pg.PlotCurveItem(
                        np.array([data.shape[0] / 2 + 0.5, data.shape[0] / 2 + 0.5]),
                        np.array([
                                data.shape[1] / 2 + 0.5 - Delta,
                                data.shape[1] / 2 + 0.5 + Delta
                        ]), pen='b')
                # Put image in plot area
                self.viewboxesviewboxes[key].addItem(self.PyrEdgeXPyrEdgeX[key])
                # Put image in plot area
                self.viewboxesviewboxes[key].addItem(self.PyrEdgeYPyrEdgeY[key])
 
        print(self.supervisorsupervisor)
 
        if self.expertexpert:
            self.expertWidgetexpertWidget.displayRtcMatrix()
 
        self.updateDisplayupdateDisplayupdateDisplay()
 
        self.uiAOuiAO.wao_run.setDisabled(False)
        self.uiAOuiAO.wao_next.setDisabled(False)
        self.uiAOuiAO.wao_open_loop.setDisabled(False)
        self.uiAOuiAO.wao_unzoom.setDisabled(False)
        self.uiAOuiAO.wao_resetSR.setDisabled(False)
        if(self.ncoroncoro):
            self.uiAOuiAO.wao_resetCoro.setDisabled(False)
            self.uiAOuiAO.wao_allCoro.setDisabled(False)
 
        WidgetBase.init_configFinished(self)
 
    def circleCoords(self, ampli: float, npts: int, datashape0: int,
                     datashape1: int) -> Tuple[float, float]:
        cx = ampli * np.sin((np.arange(npts) + 1) * 2. * np.pi / npts) + datashape0 / 2
        cy = ampli * np.cos((np.arange(npts) + 1) * 2. * np.pi / npts) + datashape1 / 2
        return cx, cy
 
    def clearSR(self):
        self.SRLESRLE = deque(maxlen=20)
        self.SRSESRSE = deque(maxlen=20)
        self.numiternumiter = deque(maxlen=20)
 
    def updateSRDisplay(self, SRLE, SRSE, numiter):
        self.SRLESRLE.append(SRLE)
        self.SRSESRSE.append(SRSE)
        self.numiternumiter.append(numiter)
        self.curveSRSEcurveSRSE.setData(self.numiternumiter, self.SRSESRSE)
        self.curveSRLEcurveSRLE.setData(self.numiternumiter, self.SRLESRLE)
 
    def updateDisplay(self) -> None:
        if (self.supervisorsupervisor is None or self.supervisorsupervisor.is_init is False):
            # print("Widget not fully initialized")
            return
        if not self.loopLockloopLock.acquire(False):
            return
        else:
            try:
                for key, dock in self.docksdocks.items():
                    if key in ["Strehl"]:
                        continue
                    elif dock.isVisible():
                        index = int(key.split("_")[-1])
                        data = None
                        if "atm" in key:
                            data = self.supervisorsupervisor.atmos.get_atmos_layer(index)
                        if "wfs" in key:
                            data = self.supervisorsupervisor.wfs.get_wfs_phase(index)
                        if "dm" in key:
                            dm_type = self.configconfig.p_dms[index].type
                            alt = self.configconfig.p_dms[index].alt
                            data = self.supervisorsupervisor.dms.get_dm_shape(index)
                        if "tar" in key:
                            data = self.supervisorsupervisor.target.get_tar_phase(index)
                        if "psfLE" in key:
                            data = self.supervisorsupervisor.target.get_tar_image(
                                    index, expo_type="le")
                        if "psfSE" in key:
                            data = self.supervisorsupervisor.target.get_tar_image(
                                    index, expo_type="se")
                        if "coroImageLE" in key:
                            data = self.supervisorsupervisor.corono.get_image(index, expo_type="le")
                        if "coroImageSE" in key:
                            data = self.supervisorsupervisor.corono.get_image(index, expo_type="se")
                        if "coroPSFLE" in key:
                            data = self.supervisorsupervisor.corono.get_psf(index, expo_type="le")
                        if "coroPSFSE" in key:
                            data = self.supervisorsupervisor.corono.get_psf(index, expo_type="se")
 
                        if "psf" in key or "coro" in key:
                            if (self.uiAOuiAO.actionPSF_Log_Scale.isChecked()):
                                if np.any(data <= 0):
                                    # warnings.warn("\nZeros founds, filling with min nonzero value.\n")
                                    data[data <= 0] = np.min(data[data > 0])
                                data = np.log10(data)
                            if (self.supervisorsupervisor.get_frame_counter() < 10):
                                self.viewboxesviewboxes[key].setRange(
                                        xRange=(data.shape[0] / 2 + 0.5 - self.PSFzoomPSFzoom,
                                                data.shape[0] / 2 + 0.5 + self.PSFzoomPSFzoom),
                                        yRange=(data.shape[1] / 2 + 0.5 - self.PSFzoomPSFzoom,
                                                data.shape[1] / 2 + 0.5 + self.PSFzoomPSFzoom),
                                )
 
                        if "SH" in key:
                            data = self.supervisorsupervisor.wfs.get_wfs_image(index)
                        if "pyrLR" in key:
                            data = self.supervisorsupervisor.wfs.get_wfs_image(index)
                        if "pyrHR" in key:
                            data = self.supervisorsupervisor.wfs.get_pyrhr_image(index)
                        if "pyrFocalPlane" in key:
                            data = self.supervisorsupervisor.wfs.get_pyr_focal_plane(index)
 
                        if (data is not None):
                            autoscale = True  # self.uiAO.actionAuto_Scale.isChecked()
                            # if (autoscale):
                            #     # inits levels
                            #     self.hist.setLevels(data.min(), data.max())
                            self.imgsimgs[key].setImage(data, autoLevels=autoscale)
                            # self.p1.autoRange()
                        elif "slp" in key:  # Slope display
                            self.imgsimgs[key].canvas.axes.clear()
                            if "Geom" in key:
                                slopes = self.supervisorsupervisor.rtc.get_slopes_geom()
                                x, y, vx, vy = plsh(
                                        slopes, self.configconfig.p_wfss[index].nxsub,
                                        self.configconfig.p_tel.cobs, returnquiver=True
                                )  # Preparing mesh and vector for display
                            if "Comp" in key:
                                centroids = self.supervisorsupervisor.rtc.get_slopes(index)
                                nmes = [
                                        2 * p_wfs._nvalid for p_wfs in self.configconfig.p_wfss
                                ]
                                first_ind = np.sum(nmes[:index], dtype=np.int32)
                                if (self.configconfig.p_wfss[index].type == scons.WFSType.PYRHR
                                            or self.configconfig.p_wfss[index].type == scons.
                                            WFSType.PYRLR):
                                    #TODO: DEBUG... with full pixels (only works with classic slopes method)
                                    plpyr(
                                            centroids[first_ind:first_ind + nmes[index]],
                                            np.stack([
                                                    self.configconfig.p_wfss[index]._validsubsx,
                                                    self.configconfig.p_wfss[index]._validsubsy
                                            ]))
                                else:
                                    x, y, vx, vy = plsh(
                                            centroids[first_ind:first_ind + nmes[index]],
                                            self.configconfig.p_wfss[index].nxsub,
                                            self.configconfig.p_tel.cobs, returnquiver=True
                                    )  # Preparing mesh and vector for display
                                self.imgsimgs[key].canvas.axes.quiver(x, y, vx, vy)
                            self.imgsimgs[key].canvas.draw()
                        elif "modalGains" in key:
                            data = self.supervisorsupervisor.rtc.get_modal_gains(index)
                            self.imgsimgs[key].canvas.axes.clear()
                            self.imgsimgs[key].canvas.axes.plot(data)
                            self.imgsimgs[key].canvas.draw()
                        elif "ContrastLE" in key:     
                            distances, mean, std, mini, maxi = self.supervisorsupervisor.corono.get_contrast(index, expo_type='le')
                            if(np.all(mean)): 
                                self.imgsimgs[key].canvas.axes.clear()
                                self.imgsimgs[key].canvas.axes.plot(distances, mean)
                                self.imgsimgs[key].canvas.axes.set_yscale('log')
                                self.imgsimgs[key].canvas.axes.set_xlabel("angular distance (Lambda/D)")
                                self.imgsimgs[key].canvas.axes.set_ylabel("Raw contrast")
 
                                self.imgsimgs[key].canvas.axes.grid()
                                self.imgsimgs[key].canvas.draw()
                        elif "ContrastSE" in key:
                            distances, mean, std, mini, maxi = self.supervisorsupervisor.corono.get_contrast(index, expo_type='se')
                            if(np.all(mean)): 
                                self.imgsimgs[key].canvas.axes.clear()
                                self.imgsimgs[key].canvas.axes.plot(distances, mean)
                                self.imgsimgs[key].canvas.axes.set_yscale('log')
                                self.imgsimgs[key].canvas.axes.set_xlabel("angular distance (Lambda/D)")
                                self.imgsimgs[key].canvas.axes.set_ylabel("Raw contrast")
                                self.imgsimgs[key].canvas.axes.grid()
                                self.imgsimgs[key].canvas.draw()
                self.firstTimefirstTime = 1
 
            finally:
                self.loopLockloopLock.release()
 
    def updateSRSE(self, SRSE):
        self.uiAOuiAO.wao_strehlSE.setText(SRSE)
 
    def updateSRLE(self, SRLE):
        self.uiAOuiAO.wao_strehlLE.setText(SRLE)
 
    def updateCurrentLoopFrequency(self, freq):
        self.uiAOuiAO.wao_currentFreq.setValue(freq)
 
    def loop_once(self) -> None:
        if not self.loopLockloopLock.acquire(False):
            print("Display locked")
            return
        else:
            try:
                start = time.time()
                self.supervisorsupervisor.next()
                for t in range(len(self.supervisorsupervisor.config.p_targets)):
                    self.supervisorsupervisor.target.comp_tar_image(t)
                loopTime = time.time() - start
 
                refreshDisplayTime = 1. / self.uiBaseuiBase.wao_frameRate.value()
 
                if (time.time() - self.refreshTimerefreshTime > refreshDisplayTime):
                    signal_le = ""
                    signal_se = ""
                    for t in range(len(self.configconfig.p_targets)):
                        SR = self.supervisorsupervisor.target.get_strehl(t)
                        # TODO: handle that !
                        if (t == self.uiAOuiAO.wao_dispSR_tar.value()
                            ):  # Plot on the wfs selected
                            self.updateSRDisplayupdateSRDisplay(SR[1], SR[0],
                                                 self.supervisorsupervisor.get_frame_counter())
                        signal_se += "%1.2f   " % SR[0]
                        signal_le += "%1.2f   " % SR[1]
 
                    currentFreq = 1 / loopTime
                    refreshFreq = 1 / (time.time() - self.refreshTimerefreshTime)
 
                    self.updateSRSEupdateSRSE(signal_se)
                    self.updateSRLEupdateSRLE(signal_le)
                    self.updateCurrentLoopFrequencyupdateCurrentLoopFrequency(currentFreq)
 
                    if (self.dispStatsInTerminaldispStatsInTerminal):
                        self.printInPlaceprintInPlace(
                                "iter #%d SR: (L.E, S.E.)= (%s, %s) running at %4.1fHz (real %4.1fHz)"
                                % (self.supervisorsupervisor.get_frame_counter(), signal_le,
                                   signal_se, refreshFreq, currentFreq))
 
                    self.refreshTimerefreshTime = start
            except Exception as e:
                print(e)
                print("error!!")
                raise e
            finally:
                self.loopLockloopLock.release()
 
    def run(self):
        WidgetBase.run(self)
        if not self.uiAOuiAO.wao_forever.isChecked():
            self.nbiternbiter -= 1
 
        if self.nbiternbiter <= 0:
            self.stopstop = True
            self.uiAOuiAO.wao_run.setChecked(False)
 
import os
import socket
 
def tcp_connect_to_display():
        # get the display from the environment
        display_env = os.environ['DISPLAY']
 
        # parse the display string
        display_host, display_num = display_env.split(':')
        display_num_major = display_num.split('.')[0]
 
        # calculate the port number
        display_port = 6000 + int(display_num_major)
 
        # attempt a TCP connection
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
                sock.connect((display_host, display_port))
        except socket.error:
                return False
        finally:
            sock.close()
        return True
 
if __name__ == '__main__':
    # if(not tcp_connect_to_display()):
    #     raise RuntimeError("Cannot connect to display")
        
    arguments = docopt(__doc__)
    app = QtWidgets.QApplication(sys.argv)
    app.setStyle('cleanlooks')
    wao = widgetAOWindow(arguments["<parameters_filename>"], cacao=arguments["--cacao"],
                         expert=arguments["--expert"], devices=arguments["--devices"])
    wao.show()
 
    print("")
    print("If the GUI is black, you can:")
    print("    type %gui qt5 to unlock GUI")
    print(" or launch ipython with the option '--gui=qt' or '--matplotlib=qt'")
    print(" or edit ~/.ipython/profile_default/ipython_config.py to set c.TerminalIPythonApp.matplotlib = 'qt'")
 
    if arguments["--interactive"]:
        from shesha.util.ipython_embed import embed
        embed(os.path.basename(__file__), locals())