"""This is a module for the SBIG ST-5 camera.

This module provides functions to configure and operate the SBIG ST-5 Camera.

(c) 2017 László Dobos, József Stéger
         Eötvös Loránd University
         Department of Physics of Complex Systems
"""

import array
import struct
import math
import numpy as np
import socket
import time
import sys
import logging
import urllib.parse

import sbigpy

try:
    import serial
except ImportError:
    serial = None

class ST5():
    CTR_START = 0xA5  # Start of Packet
    CTR_ACK = 0x06  # Packet received ok
    CTR_NAK = 0x15  # Bad checksum
    CTR_CAN = 0x18  # Bad command, length, parameter

    CMD_TAKE_IMAGE = 0x01
    CMD_END_EXPOSURE = 0x02
    CMD_GET_READOUT_PEAK = 0x03
    CMD_GET_ACTIVITY_STATUS = 0x05
    CMD_CLR_BUF = 0x06
    CMD_READ_BLANK_VIDEO = 0x12
    CMD_GET_ROM_VERSION = 0x19
    CMD_SET_COM_BAUD = 0x1A
    CMD_GET_UNCOMPRESSED_LINE = 0x1F
    CMD_GET_CPU_INFO = 0x25
    CMD_ACCUM_IMAGE = 0x0A
    CMD_REGULATE_TEMP = 0x0E
    CMD_OUTPUT_TEMP = 0x10
    CMD_RESET = 0x1B
    CMD_READ_THERMISTOR = 0x1D
    CMD_CAL_WITDH = 0x1E
    CMD_GET_TEMP_STATUS = 0x20
    CMD_FLUSH_CCD = 0x27

    BUFFER_DARK = 0
    BUFFER_LIGHT = 1
    BUFFER_ACCU = 2

    ABG_STATE_LOW = 0
    ABG_STATE_CLOCKED = 1
    ABG_STATE_MID = 2

    READOUT_MODE_HIGH = 0
    READOUT_MODE_LOW = 1

    STATUS_IDLE = 0
    STATUS_SENT_TO_FOREGROUND = 1
    STATUS_WAITING_FOR_SHUTTER = 2
    STATUS_FLUSHING_CCD = 3
    STATUS_TIMING_EXPOSURE = 4
    STATUS_WAITING_FOR_END_EXPOSURE = 5
    STATUS_TRANSFERING_CCD = 6
    STATUS_WAITING_FOR_READOUT = 7
    STATUS_READING_CCD = 8
    STATUS_POST_PROCESSING = 9

    TEMP_T_0 = 25.0
    TEMP_R_0 = 3.0
    TEMP_DT = 50.0
    TEMP_R_RATIO = 9.1
    TEMP_R_BRIDGE = 9.09
    TEMP_MAX_AD = 8192
    TEMP_MIN_CELSIUS = -10
    TEMP_MAX_CELSIUS = 15

    DEFAULT_SERIAL_URI = 'pty://dev/ttyS0&b=9600'
    DEFAULT_TCP_URI = 'tcp://localhost:2002'
    DEFAULT_BAUDRATE = 9600
    DEFAULT_TEMP_SAMP_RATE = 10
    DEFAULT_TEMP_P_GAIN = 1000
    DEFAULT_TEMP_I_GAIN = 164

    def __init__(self, uri=DEFAULT_SERIAL_URI):
        self.uri = uri
        self.socket = None
        self.last_command = None
        self.rom_version = None

    def open(self):
        uri = urllib.parse.urlparse(self.uri)
        if uri.scheme == 'pty':
            params = urllib.parse.parse_qs(uri.query)
            port = '/' + uri.hostname + uri.path
            baudrate = int(params['b'][0])
            # If baudrate is not default, set it up first
            if baudrate != ST5.DEFAULT_BAUDRATE:
                self.open_serial(port, ST5.DEFAULT_BAUDRATE)
                self.get_rom_version()
                self.set_com_baud(baudrate)
                time.sleep(0.8)
                self.close_serial()
            self.open_serial(port, baudrate)
        elif uri.scheme == 'tcp':
            self.open_tcp(uri.hostname, uri.port)
        else:
            raise sbigpy.ST5Error("Wrong URI scheme.")
        self.get_rom_version()

    def close(self):
        if (isinstance(self.socket, socket.socket)):
            self.close_tcp()
        elif (isinstance(self.socket, serial.Serial)):
            self.close_serial()
        self.socket = None

    """
    Serial and TCP communication
    """

    def open_serial(self, port, baudrate):
        # user must be in dialout group
        # sudo adduser USERNAME dialout
        # sudo chmod a+rw /dev/ttyUSB0
        logging.debug("Opening port %s at baud rate %d" % (port, baudrate))
        self.socket = serial.Serial(
            port=port,
            baudrate=baudrate,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=3)
        self.rom_version = self.get_rom_version()

    def close_serial(self):
        if self.socket and self.socket.baudrate != ST5.DEFAULT_BAUDRATE:
            self.set_com_baud(ST5.DEFAULT_BAUDRATE)
            time.sleep(0.8)
        if self.socket:
            self.socket.close()

    def open_tcp(self, hostname, port):
        logging.debug("Opening tcp connection to %s:%d" % (hostname, port))
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.connect((hostname, port))

    def close_tcp(self):
        if self.socket:
            self.socket.close()

    def int8(value):
        return struct.pack('B', value)

    def int16(value):
        return struct.pack('<H', value)

    def int32(value):
        return struct.pack('<I', value)

    def build_packet(self, cmd, data):
        size = len(data)
        checksum = 0
        buffer = [ST5.CTR_START, cmd]
        buffer += ST5.int16(size)
        buffer += data
        for d in buffer:
            checksum += d
        checksum = checksum & 0xFFFF
        buffer += ST5.int16(checksum)
        return buffer

    def send_packet(self, buffer):
        buffer = array.array('B', buffer).tobytes()
        logging.debug("Sending:")
        logging.debug(buffer)
        self.socket_send(buffer)

    def socket_send(self, buffer):
        if (isinstance(self.socket, socket.socket)):
            self.socket.sendall(buffer)
        elif (isinstance(self.socket, serial.Serial)):
            self.socket.write(buffer)
            self.socket.flush()

    def read_packet(self):
        (res,) = struct.unpack('<B', self.socket_read())
        if res == ST5.CTR_ACK:
            logging.debug("Received: ACK")
            return res, None, None, None
        elif res == ST5.CTR_CAN:
            logging.debug("Received: CAN")
            raise sbigpy.ST5Error('Bad command, length or parameter.')
        else:
            buffer = self.socket_read(3)
            cmd, size = struct.unpack('<BH', buffer)
            logging.debug("Receiving %d bytes", size)
            data = self.socket_read(size)
            (checksum,) = struct.unpack('<H', self.socket_read(2))
            checksum -= res
            for d in buffer:
                checksum -= d
            for d in data:
                checksum -= d
            checksum = checksum & 0xFFFF
            if checksum != 0:
                raise sbigpy.ST5Error('Checksum error')
            return res, cmd, size, data

    def socket_read(self, size=1):
        if (isinstance(self.socket, socket.socket)):
            buffer = b''
            while len(buffer) < size:
                buffer += self.socket.recv(size - len(buffer))
            return buffer
        elif (isinstance(self.socket, serial.Serial)):
            return self.socket.read(size)

    def command(self, cmd, data=[]):
        self.last_command = cmd
        buffer = self.build_packet(cmd, data)
        self.send_packet(buffer)
        return self.read_packet()

    """
    Utility and high level functions
    """

    def print_progress(self, total, value, caption='', ln=40.0):
        s = int(round(ln * value / total))
        sys.stdout.write('\r' + caption + ': ' + '+' * s + '-' * (int(ln) - s))

    def wait_for_idle(self):
        while True:
            status = self.get_activity_status(self.last_command)
            if status == ST5.STATUS_IDLE:
                break
            logging.debug("Status: %d, sleeping..." % status)
            time.sleep(0.3)

    def print_status(self):
        print("ROM version: %s" % self.get_rom_version())
        info = self.get_cpu_info()
        info.prettyprint()
        temp_status = self.get_temp_status()
        temp_status.prettyprint()
        print("Temperature: %f °C" % self.get_temperature())
        if serial and isinstance(self.socket, serial.Serial):
            print("Baud rate: %d" % self.socket.baudrate)

    def set_baudrate(self, baudrate):
        # Only set baud rate if connected on serial port
        if serial and isinstance(self.socket, serial.Serial):
            self.set_com_baud(baudrate)
            time.sleep(0.8)
            self.close_serial()
            self.open_serial(self.port, self.baudrate)
            self.get_rom_version()

    def read_image(self, buffer = BUFFER_LIGHT,
                  line_start = 0, line_len = 240,
                  pixel_start = 0, pixel_len = 320):

        # Switch to higher baud rate
        # self.set_baudrate(ST5.DEFAULT_READ_IMAGE_BAUDRATE)

        image = np.empty([pixel_len, line_len], dtype=np.uint16)
        for line in range(line_len):
            self.print_progress(line_len, line, 'Downloading image')
            data = self.get_uncompressed_line(buffer, line_start + line, pixel_start, pixel_len)
            image[:,line] = np.fromstring(data, dtype=np.uint16)
        print()

        # Switch back to standard baud rate
        # self.set_baudrate(ST5.DEFAULT_BAUDRATE)

        return image


    def get_setpoint(self, celsius):
        r = ST5.TEMP_R_0 * math.exp(math.log(ST5.TEMP_R_RATIO) * (ST5.TEMP_T_0 - celsius) / ST5.TEMP_DT)
        setpoint = ST5.TEMP_MAX_AD / (ST5.TEMP_R_BRIDGE / r + 1.0)
        setpoint = int(setpoint)
        return setpoint

    def get_celsius(self, setpoint):
        r = ST5.TEMP_R_BRIDGE / (ST5.TEMP_MAX_AD / setpoint - 1.0)
        celsius = ST5.TEMP_T_0 - ST5.TEMP_DT * math.log(r / ST5.TEMP_R_0) / math.log(ST5.TEMP_R_RATIO)
        return celsius

    def get_temperature(self):
        setpoint = self.read_thermistor()
        celsius = self.get_celsius(setpoint)
        return celsius

    def set_temperature(self, celsius):
        if celsius < ST5.TEMP_MIN_CELSIUS or celsius > ST5.TEMP_MAX_CELSIUS:
            raise sbigpy.ST5Error("Invalid temperature.")
        setpoint = self.get_setpoint(celsius)
        self.regulate_temp(1, setpoint)

    def reset_temperature(self):
        self.regulate_temp(0)

    """
    ST-5 CPU commands in order of CMD codes
    """

    def take_image(self, exposure_time = 1,
                   line_start = 0, line_len = 240,
                   pixel_start = 0, pixel_len = 320,
                   enable_dcs = 0, dc_restore = 0,
                   abg_state = ABG_STATE_CLOCKED, abg_period = 6000,
                   buffer = BUFFER_LIGHT,
                   auto_dark = 0,
                   readout_mode = READOUT_MODE_HIGH,
                   open_shutter = 0):
        """
        Purpose: Control shutter, Clear CCD, Time exposure and Readout CCD.
                 (see docs for details)
        Response: ACK
        Status Values:
          0=Idle, 1=Sent to foreground, 2=Waiting for shutter, 3=Flushing CCD,
          4=Timing exposure, 5=Waiting for end_exposure, 6=Transferring CCD,
          7=Waiting for readout, 8=Reading CCD, 9=Post processing,
          100-341=Digitizing line n where n=status-100
        """
        logging.debug("Command: CMD_TAKE_IMAGE")
        data = struct.pack('<I12H', exposure_time, line_start, line_len, pixel_start, pixel_len,
                           enable_dcs, dc_restore, abg_state, abg_period, buffer, auto_dark, readout_mode,
                           open_shutter)
        res, cmd, size, data = self.command(ST5.CMD_TAKE_IMAGE, data)
        self.wait_for_idle()

    def end_exposure(self, skip_transfer):
        """
        Purpose: Transfers and reads out CCD.
        Parameters:
          abort (boolean)-Skip transfer and readout of CCD if TRUE
        Response: ACK
        Status Values:
          0=Idle, 6=Transferring CCD, 7=Waiting for readout, 8=Reading CCD,
          9=Post processing, 100-341=Digitizing line n where n=status-100
        Notes:
        * Aborts exposure in progress if not waiting for end_exposure command.
        * Use this command after issuing a take_image command with the exposure_time
          set to zero.
        """
        logging.debug("Command: CMD_END_EXPOSURE")
        data = struct.pack('<H', skip_transfer)
        res, cmd, size, data = self.command(ST5.CMD_END_EXPOSURE, data)

    def get_readout_peak(self):
        """
        Purpose: Get the peak pixel value and location found during the previous readout.
        Response: (immediate packet)
          peak_value (int)-Peak pixel value found during readout
          peak_x (int)-Location of peak in x
          peak_y (int)-Location of peak in y
        Status Values:
          0=Idle, 2=In progress
        Notes:
        * If this command is issued after the take_image command with the auto_dark
          parameter TRUE it will report the peak value after the dark frame subtraction.
        """
        logging.debug("Command: CMD_GET_READOUT_PEAK")
        res, cmd, size, data = self.command(ST5.CMD_GET_READOUT_PEAK)
        (peak_value, peak_x, peak_y) = struct.unpack('<HHH', data)
        return peak_value, peak_x, peak_y


    def get_activity_status(self, command):
        """
        Purpose: Report activity status of any command to the Host.
        Parameters:
            command (int)-command to get status word for
        Response: (immediate packet)
            command (int)-command for which status is being reported
            status (int)-status of above command
        Status Values:
            0=Idle, 2=In progress
        Notes:
        • Use this command to monitor the progress of any command that has been sent to
          the foreground for processing.
        """
        lastcmd = self.last_command
        logging.debug("Command: CMD_GET_ACTIVITY_STATUS")
        data = struct.pack('<H', command)
        res, cmd, size, data = self.command(ST5.CMD_GET_ACTIVITY_STATUS, data)
        (_, status) = struct.unpack('<HH', data)
        self.last_command = lastcmd
        return status

    def clr_buf(self, buffer = BUFFER_LIGHT):
        """
        Purpose: Clears one of the image buffers by filling it with zeros.
        Parameters:
          buf (buffer)-Destination buffer to clear
        Response:ACK
        Status Values:
          0=Idle, 1=Sent to foreground, 2=In progress
        Notes:
        * You do not need to call this function prior to using the take_image command
          since the buffers are overwritten by that command.
        """
        logging.debug("Command: CMD_CLR_BUF")
        data = struct.pack('<H', buffer)
        res, cmd, size, data = self.command(ST5.CMD_CLR_BUF, data)
        self.wait_for_idle()


    def read_blank_video(self, enable_dcs=0, head_offset=0):
        """
        Purpose: Read the output level of the CCD at the black level as digitized by the A/D.
        Parameters:
          enable_dcs (boolean)-For ST-6 cameras enables DCS when TRUE. Other cameras
          ignore the setting of this parameter.
          head_offset (int)-For ST-6 cameras this is the setting of the offset DAC in the head
          with allowed values of 0-255. Other cameras ignore the setting of this parameter.
        Response: (immediate packet)
          video (int)-CCD black level A/D value
        Status Values:
          0=Idle, 2=In progress
        Notes:
        * Physical activation of this command only occurs if an exposure is not
        """
        logging.debug("Command: CMD_READ_BLANK_VIDEO")
        data = struct.pack('<HH', enable_dcs, head_offset)
        res, cmd, size, data = self.command(ST5.CMD_READ_BLANK_VIDEO, data)
        (value,) = struct.unpack('<H', data)
        return value

    def get_rom_version(self):
        """
        Purpose: Report CPU internal firmware version.
        Response: firmware_version (int)-this is the firmware version
        Status Values: 0=Idle, 2=In progress
        Notes:
        * Interpret the firmware version as a 4 digit BCD number with a two digit fraction
          (for example version 257 decimal = 0201H would be version 2.01)
        * You should check the ROM version prior to issuing commands that are not
          supported by all ROMS. For example before you use the binning modes
          (readout_mode parameter) of the take_image command other than 0 and 1 you
          should check that the ROM supports those modes.
        """
        logging.debug("Command: CMD_GET_ROM_VERSION")
        res, cmd, size, data = self.command(ST5.CMD_GET_ROM_VERSION)
        version = "v%d.%d" % (data[0], data[1])
        logging.debug("ROM version is %s" % version)
        return version

    def set_com_baud(self, baudrate):
        """
        Purpose: Set the baud rate of the COM port for communications with the Host.
        Parameters: baud (long)-baud rate to set COM port to
        Response:ACK
        Status Values: 0=Idle, 2=In progress
        Notes:
        * The CPU sends the ACK at the old baud rate then switches to the new rate. You
          must then send and the CPU must receive a get_rom_version command within
          1.0 second or the CPU will switch back down to 9600 baud.
        """
        logging.debug("Command: CMD_SET_COM_BAUD")
        logging.debug("Setting baud rate to %d" % baudrate)
        data = struct.pack('<L', baudrate)
        res, cmd, size, data = self.command(ST5.CMD_SET_COM_BAUD, data)
        self.baudrate = baudrate

    def get_uncompressed_line(self, buffer=BUFFER_LIGHT, line_start=0, pixel_start=0, pixel_len=240):
        """
        Purpose: Transmit an uncompressed line of image data.
        Parameters:
          buf (buffer)-Source image buffer from which data will be transmitted
          line_start (int)-line to transmit (value depends on the size of the image buffers, 0 thru
          max height - 1)
          pixel_start (int)-leftmost pixel to transmit (value depends on the size of the image
          buffers, 0 thru max width - 1)
          pixel_len (int)-number of pixels to transmit (value depends on the size of the image
          buffers, 1 thru max width)
        Response: (immediate packet)
          line_start (int)-line being sent
          data_1 (int)-1st pixel
          .
          .
          data_N (int)-last pixel
        Status Values:
          0=Idle, 2=In progress
        Notes: (see docs)
        """
        logging.debug("Command: CMD_GET_UNCOMPRESSED_LINE")
        data = struct.pack('<HHHH', buffer, line_start, pixel_start, pixel_len)
        res, cmd, size, data = self.command(ST5.CMD_GET_UNCOMPRESSED_LINE, data)
        return data[2:]

    def get_cpu_info(self):
        """
        Purpose:Return the camera model and capabilities.
        Response: (immediate packet) -- see docs
        Status Values: 0=Idle, 2=In progress
        Notes:
        * The packet length of this response will vary from camera model to camera model
          as different models support different numbers of readout modes.
        """
        logging.debug("Command: CMD_GET_CPU_INFO")
        res, cmd, size, data = self.command(ST5.CMD_GET_CPU_INFO)
        info = sbigpy.Cpu_Info()
        info.unpack(data)
        return info

    def accum_image(self, x_offset=0, y_offset=0):
        """
        Purpose: Add the light buffer to the accumulation buffer with X and Y offsets.
        Parameters:
          x_offset (signed int)-x offset of light buffer relative to accumulation buffer. Positive
          values shifts light buffer to right
          y_offset (signed int)-y offset of light buffer relative to accumulation buffer. Positive
          values shifts light buffer down
        Response:ACK
        Status Values:
          0=Idle, 1=Sent to foreground, 2=In progress
        """
        logging.debug("Command: CMD_ACCUM_IMAGE")
        data = struct.pack('<HH', x_offset, y_offset)
        res, cmd, size, data = self.command(ST5.CMD_ACCUM_IMAGE, data)
        self.wait_for_idle()


    def regulate_temp(self, enable, setpoint = 0,
                      samp_rate=DEFAULT_TEMP_SAMP_RATE,
                      p_gain = DEFAULT_TEMP_P_GAIN,
                      i_gain = DEFAULT_TEMP_I_GAIN,
                      reset_brownout = 1):
        """
        Purpose: Enable or disable the CPU temperature regulation.
        Parameters:
          enable (boolean)-enable temperature regulation when TRUE
          setpoint (int)-temperature or thermistor setpoint in A/D units
          samp_rate (int)-temperature sampling rate in hundredths of a second
          p_gain (int)-proportional gain term
          i_gain (int)-integral gain term
          reset_brownout (boolean)-reset brownout detector when TRUE
        Response:ACK
        Status Values:
          0=Idle, temperature regulation off, 1=Enabled
        Notes:
        * Since ST-4X cameras do not have temperature regulation, the settings of all
          parameters in this command except the reset_brownout parameter are ignored by
          ST-4X cameras.
        * On power up the ST-5 and ST-6 CPUs read the current CCD temperature and start
          regulating the temperature at that setpoint. The ST-4X CPU ramps the TE cooler
          drive to the maximum drive level.
        """
        logging.debug("Command: CMD_REGULATE_TEMP")
        data = struct.pack('<6H', enable, setpoint, samp_rate, p_gain, i_gain, reset_brownout)
        res, cmd, size, data = self.command(ST5.CMD_REGULATE_TEMP, data)

    def output_temp(self, value):
        logging.debug("Command: CMD_OUTPUT_TEMP")
        data = struct.pack('<H', value)
        res, cmd, size, data = self.command(ST5.CMD_OUTPUT_TEMP, data)

    def reset(self):
        """
        Purpose: Perform a cold reset of the CPU.
        Response:ACK
        Notes:
        * The CPU will restart communicating at 9600 baud just as if freshly powered up.
        """
        logging.debug("Command: CMD_RESET")
        res, cmd, size, data = self.command(ST5.CMD_RESET)

    def read_thermistor(self):
        logging.debug("Command: CMD_READ_THERMISTOR")
        res, cmd, size, data = self.command(ST5.CMD_READ_THERMISTOR)
        (t,) = struct.unpack('<H', data)
        return t

    def cal_width(self, buffer, x_offset, y_offset, x_length, y_length):
        """
        Purpose: Calculates the peak and the average pixel value in the box, then calculates the
          height and width at half the amplitude of the peak over the average for the
          data contained within the box, placing the results into the global result buffer.
        Parameters:
          buf (buffer)-source buffer to use in calculating the width
          x_offset (int)-leftmost pixel of box (value depends on the size of the image buffers, 0
          thru max width -1 )
          y_offset (int)-top line of box (value depends on the size of the image buffers, 0 thru
          max height - 1)
          x_length (int)-width of box (value depends on the size of the image buffers, 1 thru
          max width)
          y_length (int)-height of box (value depends on the size of the image buffers, 1 thru
          max height)
        Response:ACK immediately plus buffered results
        Results: (placed into global results buffer)
          width (int)-full width at half amplitude of star in box
          height (int)-full height at half amplitude of star in box
        Status Values:
          0=Idle, 1=Sent to foreground, 2=In progress
        Notes:
        * You must use the get_result_buf command once this command has finished
          executing to get the results of this command.
        * Use the get_cpu_info command to determine the maximum size of the image
          buffers.
        """
        logging.debug("Command: CMD_CAL_WITDH")
        data = struct.pack('<H5', buffer, x_offset, y_offset, x_length, y_length)
        res, cmd, size, data = self.command(ST5.CMD_CAL_WITDH)
        (width, height) = struct.unpack('<HH', data)
        return width, height

    def get_temp_status(self):
        logging.debug("Command: CMD_GET_TEMP_STATUS")
        res, cmd, size, data = self.command(ST5.CMD_GET_TEMP_STATUS)
        status = sbigpy.Temp_Status()
        status.unpack(data)
        return status

    def flush_ccd(self, cycles):
        logging.debug("Command: CMD_FLUSH_CCD")
        data = struct.pack('<H', cycles)
        res, cmd, size, data = self.command(ST5.CMD_FLUSH_CCD, data)