from . import laser as las
from .agilent_lightwave_connection import AgilentLightWaveConnection
from ..utils import gnuplot as gp
import time
import struct
import numpy as np
from scipy import interpolate
import os
import serial as ser
import warnings

class LaserAgilent8164B(AgilentLightWaveConnection, las.LaserTunable):
    '''
    Controls the laser module in the Agilent 8164B.

    Args:
        gpib_num (int): The GPIB bus number the laser is on.
        gpib_dev_num (int): The device number the laser in on
            the bus.
        power_unit (str): Either \'W\' or \'dBm\' depending on what
            units the Agilent 8164B should use for the laser power.
        output_mode (str):
            'HIGH' -> The High Power output is regulated.
            'LOWS' -> The Low SSE output is regulated.
            'BHR' -> Both outputs are active but only the High Power output is Regulated.
            'BLR' -> Both outputs are active but only the Low SSE output is Regulated.
    '''
    def __init__(self, gpib_num=None, gpib_dev_num=None, power_unit='W',
                 output_mode='high'):
        super().__init__(gpib_num, gpib_dev_num)
        self._power_unit = self.set_unit(power_unit)
        self._set_output_mode(output_mode)

    def turn_on(self):
        # Turn laser on.
        self._write('outp1')
        # Open output shutter.
        self._write('sour0:pow:stat 1')
        return self.get_on_or_off()

    def turn_off(self):
        # Turn laser off.
        self._write('outp0')
        # Open output shutter.
        self._write('sour0:pow:stat 0')
        return self.get_on_or_off()

    def get_on_or_off(self):
        on_off = self._query('sour0:pow:stat?')
        return bool(int(on_off))

    def set_power_W(self, power_W):
        assert power_W < 5.e-3, 'Can\'t set power that high.'
        if self._power_unit == 'dBm':
            power_dbm = LaserAgilent8164B.watts_to_dbm(power_W)
            cmd = 'sour0:pow %fdbm' % power_dbm
        elif self._power_unit == 'W':
            cmd = 'sour0:pow %fw' % power_W
        self._write(cmd)
        return self.get_power_W()

    def get_power_W(self):
        inString = 'sour0:pow?'
        data = self._query(inString)
        return float(data)

    def set_wavelength_m(self, wavelength_m):
        cmd = 'sour0:wav %sM' % wavelength_m
        self._write(cmd)
        r = self.get_wavelength_m()
        return r

    def get_wavelength_m(self):
        cmd = 'sour0:wav?'
        r = float(self._query(cmd))
        return r

    def set_unit(self, unit):
        '''
        Sets the units the laser should operate in.

        This will affect in what units the laser power
        is displayed on the screen.

        Args:
            unit (str): Either \'dBm\' or \'W\'.

        Returns:
            str: The units the laser is oeprating in.
        '''
        assert unit.lower() in ('dbm', 'w')
        cmd = 'sour0:pow:unit %s' % unit
        self._write(cmd)
        self._power_unit = self.get_unit()
        return self._power_unit

    def get_unit(self):
        '''
        Gets the units the laser is operating in.

        The units affect in what units the laser power
        is displayed on the screen.

        Returns:
            str: The units the laser is oeprating in.
        '''
        inString = 'sour0:pow:unit?'
        data = self._query(inString)
        unit = int(data)
        if unit == 0:
            unit_str = 'dBm'
        else:
            unit_str = 'W'
        return unit_str

    def last_operation_completed(self):
        return self._query('*OPC?')

    def wait_for_last_operation_completed(self):
        finish = False
        while not finish:
            finish = self.last_operation_completed()
        return finish

    def _get_output_mode(self):
        m = self._query('outp0:path?')
        return m

    def _set_output_mode(self, mode):
        '''
        Sets the mode the power meter will operate in.

        Args:
            mode (str):
                HIGH -> The High Power output is regulated.
                LOWS -> The Low SSE output is regulated.
                BHR -> Both outputs are active but only the High Power output is Regulated.
                BLR -> Both outputs are active but only the Low SSE output is Regulated.

        Returns:
            str: The mode the laser is operating in.
        '''
        m = mode.lower()
        assert m in ('high', 'lows', 'bhr', 'blr')
        self._write('outp0:path %s' % m)
        return self._get_output_mode()

    def get_velocity_nm_s(self):
        return float(self._query('sour0:wav:swe:spe?'))

    def set_velocity_nm_s(self, sweep_speed_nm_per_sec):
        assert float(sweep_speed_nm_per_sec) in (0.5, 1., 2., 5., 10., 20., 40., 80.), \
            'Invalid sweep speed choice; laser only supports certain sweep speeds nm/s.'
        self._write('sour0:wav:swe:spe %fnm/s' % sweep_speed_nm_per_sec)
        return self.get_velocity_nm_s()

    def wavelength_sweep(self,
                         start_wavelength_nm,
                         stop_wavelength_nm,
                         step_wavelength_nm,
                         max_power_mW,
                         sweep_speed_nm_per_sec=5.,
                         power_metre_slot=1,
                         int_time_ms=20.,
                         filename=None,
                         preserve_settings=True):
        assert step_wavelength_nm >= 0.6e-3, '`step_wavelength_nm` to small.'

        max_power_W = max_power_mW / 1000.
        int_time_us = int_time_ms * 1000.
        int_time_s = int_time_ms / 1000.

        # Save current settings
        if preserve_settings:
            cmds = ['trig:conf?',
                    'sour0:pow:stat?',
                    'trig0:outp?',
                    'trig0:inp?',
                    'sour0:wav:swe:star?',
                    'sour0:wav:swe:stop?',
                    'sour0:wav:swe:step?',
                    'sour0:wav:swe:mode?',
                    'sour0:wav:swe:llog?',
                    'trig1:inp?',
                    'sens'+str(power_metre_slot)+':pow:unit?',
                    'sens'+str(power_metre_slot)+':pow:rang:auto?',
                    'sens'+str(power_metre_slot)+':pow:rang?',
                    'sens'+str(power_metre_slot)+':pow:wav?',
                    'sens'+str(power_metre_slot)+':func:par:logg?']
            temps = [self._query(cmd) for cmd in cmds]

        max_power_dbm = LaserAgilent8164B.watts_to_dbm(max_power_W)

        # Setup lasers and triggers.
        self._write('trig:conf loop')
        self._write('sour0:pow:stat 1')
        self._query('sour0:pow:stat?')
        self._write('trig0:outp stf')
        self._write('trig0:inp sws')
        self._write('sour0:wav:swe:star %fnm' % start_wavelength_nm)
        self._write('sour0:wav:swe:stop %fnm' % stop_wavelength_nm)
        self._write('sour0:wav:swe:step %fnm' % step_wavelength_nm)
        self.set_velocity_nm_s(sweep_speed_nm_per_sec)
        self._write('sour0:wav:swe:mode cont')
        self._write('sour0:wav:swe:llog 1')

        param_check = self._query('sour0:wav:swe:chec?').strip()
        if param_check != '0,OK':
            raise RuntimeError('Invalid sweep parameters: `%s`.' % param_check)

        max_selectable_power_W = float(self._query('sour0:wav:swe:pmax? %fnm,%fnm' %
                                                  (start_wavelength_nm, stop_wavelength_nm)))
        if max_power_W > max_selectable_power_W:
            raise RuntimeError('Invalid `max_power` choice.')
        self._write('sour0:pow %fW' % max_power_W)

        # This command doesn't work for our laser module: int(self._query('sour0:wav:swe:exp?'))
        # so it must be manually calculated.
        expected_triggers = int(self._query('sour0:wav:swe:exp?'))
        power_meas_time = int_time_s*expected_triggers
        sweep_time = (stop_wavelength_nm-start_wavelength_nm)/sweep_speed_nm_per_sec
        if power_meas_time <= sweep_time:
            warnings.warn('The time needed for power readings %.2fsec is longer than the sweep %.2fsec.  Either reduce the integration time of the power meter, reduce the sweep speed, or reduce the sampling rate (`step_wavelength_nm`).' % (power_meas_time, sweep_time))
        self._write('sour0:wav:swe 1')

        # Setup power meter.
        self._write('trig1:inp sme')
        self._write('sens'+str(power_metre_slot)+':pow:unit 1')
        self._write('sens'+str(power_metre_slot)+':pow:rang:auto 0')
        self._write('sens'+str(power_metre_slot)+':pow:rang %fdbm' % max_power_dbm)
        self._query('sens'+str(power_metre_slot)+':pow:rang?')
        centre_wavelength_nm = 0.5*(start_wavelength_nm+stop_wavelength_nm)
        self._write('sens'+str(power_metre_slot)+':pow:wav %fnm' % centre_wavelength_nm)
        self._write('sens'+str(power_metre_slot)+':func:par:logg %i,%ius' % (expected_triggers, int_time_us))

        # Set to one trigger per measurment.
        self._write('sens'+str(power_metre_slot)+':func:stat logg,star')

        finished = False
        while not finished:
            flag = int(self._query('source0:wav:sweep:flag?').strip())
            finished = bool(flag)

        self._write('sour0:wav:swe:soft')

        # Wait for log to complete.
        while self._query('sour0:wav:sweep:state?').strip() != '+0':
            time.sleep(0.1)

        # Version from manual.  Doesn't work.
        #while self._query('sens'+str(power_metre_slot)+':func:stat?').strip() != 'LOGGING_STABILITY,COMPLETE':
        #    self._query('sens'+str(power_metre_slot)+':func:stat?')
        #    time.sleep(50.e-3)

        # Get logged data.
        self._write('sens'+str(power_metre_slot)+':func:res?')
        r = self._read(2)
        assert r[0] == '#', 'A \'#\' should have been read first.'
        size_of_num_bytes = int(r[1])
        num_bytes = int(self._read(size_of_num_bytes))
        data = self._read_raw(num_bytes)
        _ = self._read(1) # Read the remaining '\n' character.
        data = [data[i:i+4] for i in range(0, len(data), 4)]
        powers = [struct.unpack('f', d)[0] for d in data]

        while int(self._query('sour0:wav:swe:flag?')) not in (flag+1, flag+2):
            time.sleep(50.e-3)

        self._write('sour0:read:data? llog')
        r = self._read(2)
        assert r[0] == '#', 'A \'#\' should have been read first.'
        size_of_num_bytes = int(r[1])
        num_bytes = int(self._read(size_of_num_bytes))
        data = self._read_raw(num_bytes)
        _ = self._read(1) # Read the remaining '\n' character.
        data = [data[i:i+8] for i in range(0, len(data), 8)]
        wavelengths = [struct.unpack('d', d)[0] for d in data]

        # Wait for sweep to finish incase it is still running.
        while int(self._query('sour0:wav:swe?')) != 0:
            time.sleep(50.e-3)

        # Reload settings
        if preserve_settings:
            for cmd, temp in zip(cmds, temps):
                self._write(cmd[:-1] + ' ' + temp)

        # Save data to file if necessary.
        try:
            with open(filename, 'w') as fs:
                for w, p in zip(wavelengths, powers):
                    fs.write('%e,%e\n' % (w, p))
            script_dir = os.path.dirname(os.path.abspath(__file__))
            gp.Gnuplot(script_dir+'/agilent_laser_sweep.gpi', {'filename':filename})
        except TypeError:
            pass

        assert expected_triggers == len(wavelengths) == len(powers), 'Invalid data.'

        return wavelengths, powers

    def wavelength_sweep_interp(self,
                                start_wavelength_nm,
                                stop_wavelength_nm,
                                step_wavelength_nm,
                                max_power_W,
                                sweep_speed_nm_per_sec=5.,
                                filename=None):
        wavelengths, powers = self.wavelength_sweep(start_wavelength_nm,
                                                    stop_wavelength_nm,
                                                    step_wavelength_nm,
                                                    max_power_W,
                                                    sweep_speed_nm_per_sec,
                                                    filename)
        powers_interp = interpolate.interp1d(wavelengths, powers)
        return powers_interp
