#!/usr/bin/env python
"""AllSky340

Base class and utilities for operating an SBIG AllSky340 allsky
camera.  Tested and used with the monochrome model only.  Color
version will require some modifications. The camera is based on the
SG-4 autoguider so this could easily be adapted to that as well.  The
information used here was gleaned from:

http://sbig.impulse.net/pdffiles/SBIG%20SG-4%20&%20AllSky-340(C)%20Serial%20Interface%20Specification.pdf

Currently only supporting 1x1 binning and either full (640x480) or
cropped (512x480).

Author                     Version             Date
--------------------------------------------------------
TE Pickering                 0.1             20130114

TODO
--------------------------------------------------------
- add support for SG-4 and color version.  split into basic class for SG-4
  that the AllSky340 classes can inherit from.

- figure out how to do simultaneous expose and readout...

Updates
--------------------------------------------------------

"""

# System modules
import logging
import os
import struct
import sys
import time

# Third-party modules
import numpy as np
import serial
from astropy.io import fits as pyfits


def emit_colored_ansi(fn):
    def new(*args):
        levelno = args[1].levelno
        if(levelno >= 50):
            color = '\x1b[31m'  # red
        elif(levelno >= 40):
            color = '\x1b[31m'  # red
        elif(levelno >= 30):
            color = '\x1b[33m'  # yellow
        elif(levelno >= 20):
            color = '\x1b[32m'  # green
        elif(levelno >= 10):
            color = '\x1b[35m'  # pink
        else:
            color = '\x1b[0m'  # normal
        args[1].levelname = color + args[1].levelname + '\x1b[0m'  # normal
        return fn(*args)
    return new


# Initialize logger
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
cam_log = logging.getLogger()
fh = logging.FileHandler("%s/skycam.log" % os.environ['HOME'])
fh.setLevel(logging.DEBUG)
fh.setFormatter(logging.Formatter("%(asctime)s: %(levelname)s - %(message)s"))
cam_log.addHandler(fh)
ch = logging.StreamHandler()
f = logging.Formatter("%(levelname)s: %(message)s")
ch.setFormatter(f)


logging.StreamHandler.emit = emit_colored_ansi(logging.StreamHandler.emit)


class AllSky340:
    def __init__(self, device, baudrate=115200, timeout=0.1):
        self.ser = serial.serial_for_url(device)
        self.ser.baudrate = baudrate
        self.ser.timeout = timeout
        for i in range(5):
            self.ping()

    def checksum(self, c):
        inv = ~ord(c[0]) & 0xFF
        mask = ~(1 << 7)
        checksum = inv & mask
        if len(c) > 1:
            for i in range(1, len(c)):
                inv = ~ord(c[i]) & 0xFF
                xor = checksum ^ inv
                checksum = xor & mask
        return checksum

    def command(self, cmd, nbytes):
        to_send = cmd.encode('utf-8') + \
            self.checksum(cmd).to_bytes(1, "little")
        self.ser.write(to_send)
        time.sleep(0.09)
        if nbytes > 0:
            return self.ser.read(nbytes)
        else:
            return True

    def ping(self):
        resp = self.command("E", 2)
        cam_log.info(f"Pinged camera and received response of {resp}")
        if resp == b":O":
            cam_log.info("Camera on-line and communicating properly.")
            return True
        else:
            cam_log.warning(f"No correct response")
            return False

    def open_shutter(self):
        cam_log.info("Opening camera shutter")
        return self.command("O", 0)

    def close_shutter(self):
        cam_log.info("Closing camera shutter")
        return self.command("C", 0)

    def shutter_off(self):
        cam_log.info("De-energizing camera shutter")
        return self.command("K", 0)

    def abort_image(self):
        cam_log.info("Aborting current exposure")
        return self.command("A", 0)

    def firmware(self):
        resp = self.command("V", 3)
        minor = resp[1]
        major = resp[0]
        if major >= 8:
            ver = "T"
            major -= 8
        else:
            ver = "V"
        version = f"{ver}{major}{minor}"
        cam_log.info(f"Camera firmware: {version}")
        return version

    def serial(self):
        resp = self.command("r", 9)
        version = resp.decode("utf-8").strip("\r")
        cam_log.info(f"Camera serial: {version}")
        return version

    def get_min_move(self):
        resp = self.command("n", 2)
        cam_log.info(f"Camera min_move: {resp[0]} {resp[1]}")
        return resp

    def set_min_move(self):
        data = "N" + struct.pack("<2B", *[0, 18]).decode()
        resp = self.command(data, 3)
        cam_log.info(f"Camera min_move: {resp}")
        return resp

    def heater_on(self):
        """per communication with SBIG, the heater is wired to the X+ guide
        relay on the camera.  use the 'force guide relays' command to
        toggle it
        """

        cmd = 'g' + struct.pack("B", 1).decode()
        cam_log.info("Turning camera heater on")
        self.command(cmd, 0)
        return True

    def heater_off(self):
        cmd = 'g' + struct.pack("B", 0).decode()
        cam_log.info("Turning camera heater off")
        self.command(cmd, 0)
        return True

    def shutter_chop_on(self):
        """
        this is an undocumented command that was provided by SBIG via
        email.  it's simply U plus 0x01 (on) or 0x00 (off).
        """
        cmd = 'U' + struct.pack("B", 1).decode("utf-8")
        cam_log.info("Enabling chopping shutter mode.")
        self.command(cmd, 0)
        return True

    def shutter_chop_off(self):
        cmd = 'U' + struct.pack("B", 0).decode("utf-8")
        cam_log.info("Disabling chopping shutter mode.")
        self.command(cmd, 0)
        return True

    def _get_image_block(self, expected=4096, ignore_cs=True):
        """
        Get one 'block' of image data. At full frame the camera returns image
        data in chunks of 4096 pixels. For different imaging modes this value
        will change, but the caller can simply change the value of expected.
        expected -- Number of pixels to retrieve
        ignore_cs -- always pass checksum without checking (for debug only)
        """

        CSUM_OK = b'K'
        CSUM_ERROR = b'R'

        valid = False
        cs_failed = 0
        while not valid:

            data = self.ser.read(expected * 2)
            cs_byte = self.ser.read(1)  # ord

            cs = 0
            for c in data:
                cs = cs ^ c
                print(cs, cs_byte)

            if cs == cs_byte:
                self.ser.write(CSUM_OK)
                valid = True
            else:
                cs_failed += 1
                self.ser.write(CSUM_ERROR)

            if ignore_cs:
                valid = True

            if cs_failed > 0:
                cam_log.error('Checksum failed')
            else:
                cam_log.info('Checksum ok')

    def block_read(self, npix, ntries=0, lrc_byte=None):
        """
        cheeky recursive function to read block of pixels and ask for repeat
        transmissions if checksums don't match
        """

        max_tries = 5

        block = self.ser.read(npix * 2 + 1)

        if lrc_byte == None:
            lrc_byte = self.ser.read(1)
        if len(lrc_byte) > 0:
            cam_lrc = lrc_byte  # ord(lrc_byte)
        else:
            cam_lrc = 0

        lrc = 0
        for b in block:
            lrc ^= b

        if lrc == cam_lrc:
            return block

        else:
            if ntries < max_tries:
                cam_log.warning(
                    f"Camera read-out error. Re-transmitting block (try #{ntries})")
                self.command('R', 0)
                # not sure why this is needed. not mentioned in document...
                # f = self.ser.read(1)
                return self.block_read(npix, ntries=ntries+1, lrc_byte=lrc_byte)
            else:
                cam_log.error("Camera read_out failed. Stopping transfer.")
                #self.command('S', 0)
                return False

    def getimage(self, exposure, light=False, cropped=False, progress_callback=None):

        exptime = exposure * 10000.0
        if exptime > 0x63FFFF:
            exptime = 0x63FFFF

        print(exptime)

        # exp = struct.pack('<' + 'I', int(exptime))[:3].decode("utf-8")
        # e = "T" + exp[::-1]

        # if light:
        #     command = e + chr(0x00) + chr(0x01)
        # else:
        #     command = e + chr(0x00) + chr(0x00)

        exp = struct.pack('<' + 'I', int(exptime))

        combuf = bytearray(6)
        combuf[0] = ord("T")
        combuf[1] = exp[2]
        combuf[2] = exp[1]
        combuf[3] = exp[0]
        combuf[4] = 0x00
        combuf[5] = 0x01 if light else 0x00

        command = combuf.decode("utf-8")

        self.command(command, 0)

        if cropped:
            crop = "cropped"
        else:
            crop = "full"

        if light:
            imtype = "light"
        else:
            imtype = "dark"

        cam_log.info(
            f"Exposing {crop} {imtype} image for {exposure} seconds....")

        out = b"E"
        #s = ""

        while out != b"D":
            time.sleep(0.2)
            out = self.ser.read(1)
            print(out)

        cam_log.info("Transferring image from camera....")

        self.command('X', 0)

        imag_str = bytearray()

        # we only use 1x1 binning and either full or cropped
        if cropped:
            nbytes = 245760
        else:
            nbytes = 307200

        # this is true for 1x1 binning.  it would be 1024 for 2x2 binning
        # or the length of one line in sub-frame mode.
        npix = 4096

        data = b''
        blocks_complete = 0

        #### new #####
        blocks_expected = nbytes // npix
        for _ in range(blocks_expected):
            time.sleep(0.2)
            data += self._get_image_block()
            blocks_complete += 1
            print("f{blocks_complete} or {_} of {blocks_expected}")
            if progress_callback is not None:
                progress_callback(float(blocks_complete) /
                                  blocks_expected * 100)

        # log.debug('Image download complete')

        time.sleep(2.2)

        # # not sure why this is needed. not mentioned in document...
        # f = self.ser.read(1)
        # print(f)

        #### original #####
        for i in range(nbytes // npix):
            time.sleep(0.160)
            block_str = self.block_read(npix)
            cam_log.warning(f"block str {block_str}")
            # if the read completely fails, just replace with zeros
            if block_str == False:
                cam_log.warning(f"block str false: {block_str}")
                l = npix*[0]
                block_str = struct.pack('B' * len(l), *l)
            imag_str.extend(block_str)
            self.command('K', 0)

        # pull these extra bytes out.  also not sure why...
        cam = self.ser.read(1)
        cam = self.ser.read(1)
        cam = self.ser.read(1)
        cam = self.ser.read(1)
        cam = self.ser.read(1)

        # this is ugly and there has to be a better way, but it works.
        earr = [z for z in imag_str[::2]]
        oarr = [z for z in imag_str[1::2]]
        arr = []

        for i in range(len(earr)):
            s = struct.pack("BB", earr[i], oarr[i])
            n = struct.unpack("H", s)[0]
            arr.append(n)

        # if cropped:
        #     imag = np.array(arr, dtype=np.int32).reshape(480, 512)
        # else:
        #     imag = np.array(arr, dtype=np.int32).reshape(480, 640)
        # return np.flipud(imag)

        return np.array(arr, dtype=np.int32)


if __name__ == '__main__':
    c = AllSky340(device="socket://151.3.45.60:8899")
    method = sys.argv[1].lower()
    args = ", ".join(sys.argv[2:])
    eval(f"c.{method}({args})")
