from pylablib.devices import NI

class Shutter():
    def __init__(self):
        self.daq_name = 'Dev1'


class ShutterPump(Shutter):
    """
    Controls shutter between opo and sample. 
    Intended to be used instead of switching laser on and off.
    Not yet in use.
    """
    def __init__(self):
        super().__init__()    
    
    def open_(self):
        with NI.NIDAQ(self.daq_name) as shutter:
            shutter.add_voltage_output('vout', 'ao1', (0,5))
            shutter.set_voltage_outputs('vout', 5)
        print('Pump shutter opened')

    def close_(self):
        with NI.NIDAQ(self.daq_name) as shutter:
            shutter.add_voltage_output('vout', 'ao1', (0,5))
            shutter.set_voltage_outputs('vout', 0)
        print('Pump shutter closed')


class ShutterProbe(Shutter):
    """
    Controls shutter between bentham TMC300 and sample. 
    Blocks probe light when not needed.
    """
    def __init__(self):
        super().__init__()    
    
    def open_(self):
        with NI.NIDAQ(self.daq_name) as shutter:
            shutter.add_voltage_output('vout', 'ao0', (0,5))
            shutter.set_voltage_outputs('vout', 5)
        print('Probe shutter opened')

    def close_(self):
        with NI.NIDAQ(self.daq_name) as shutter:
            shutter.add_voltage_output('vout', 'ao0', (0,5))
            shutter.set_voltage_outputs('vout', 0)
        print('Probe shutter closed')

            
if __name__ == '__main__':
    from DeviceManager import DeviceManager
    dm = DeviceManager()
    
    #dm.shutter_probe.open_()
    dm.shutter_probe.close_()
    
    #dm.shutter_pump.open_()
    #dm.shutter_pump.close_()