#!/usr/bin/env python

# System modules
import struct
import time

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

KONTINUE = b'K'
RESEND = b'R'
STOP = b'S'
BAUDRATE = 115200
PIXEL_SIZE = 2


image_type = {
    "dark":  0,
    "light": 1,
    "auto":  2  # Light-Dark (only binned).
}

frame_type = {
    "full":   {"width": 640,  "height": 480,  "blocks": 4096},
    "crop":   {"width": 512,  "height": 480,  "blocks": 4096},
    "binned": {"width": 320,  "height": 240,  "blocks": 1024},  # only auto
    #    "custom": {"width": size, "height": size, "blocks": size }
}

# if params.imagetyp == 'auto':
#     params.frametyp='binned';


class AllSky340:
    def __init__(self, device, timeout=0.160):
        self.ser = serial.serial_for_url(device)
        self.ser.baudrate = BAUDRATE
        self.ser.bytesize = serial.EIGHTBITS
        self.ser.timeout = timeout
        self.width = 512
        self.height = 480
        self.expected = 4096
        self.status = False
        for i in range(3):
            self.ping()

    def wait(self):
        overhead = 1.0
        return (self.expected*PIXEL_SIZE) / (BAUDRATE/10) * overhead

    def checksum(self, c):
        inv = ~c[0] & 0xFF
        mask = ~(1 << 7)
        checksum = inv & mask
        if len(c) > 1:
            for i in range(1, len(c)):
                inv = ~c[i] & 0xFF
                xor = checksum ^ inv
                checksum = xor & mask
        log.info(f"Checksum: {chr(checksum).encode()}")
        return chr(checksum).encode()

    def command(self, cmd, response_bytes=None):

        check = self.checksum(cmd)
        to_send = cmd + check
        log.info(f"Sending:  {to_send}")
        self.ser.write(to_send)
        time.sleep(self.ser.timeout)
        resp = self.ser.read(1)
        if resp == check:
            log.info(f"Response ok: '{resp}' == '{check}'")
            self.status = True
        else:
            log.warning(f"Wired response: '{resp}' == '{check}'")
            self.status = False

        if response_bytes:
            resp = self.ser.read(response_bytes)
            log.info(f"Received: {resp}")
            return resp
        else:
            return self.status

    def ping(self):
        log.info("Pinging Camera")
        resp = self.command(b"E", 1)
        if resp == b"O":
            log.info(f"Response ok: {resp}")
            return True
        else:
            log.warning(f"Response not ok: {resp}")
            return False

    def open_shutter(self):
        log.info("Opening camera shutter")
        return self.command(b"O")

    def close_shutter(self):
        log.info("Closing camera shutter")
        return self.command(b"C")

    def shutter_off(self):
        log.info("De-energizing camera shutter")
        return self.command(b"K")

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

    def activate_guide_relays(self):
        log.error("Not implemented")
        pass

    def force_guide_relays(self):
        log.error("Not implemented")
        pass

    def abort(self):
        log.info("Aborting current exposure")
        return self.command(b"A")

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

    def set_max_move(self):
        data = b"M" + struct.pack("<2B", *[15, 0])
        resp = self.command(data, 1)
        return True if resp == b"K" else False

    def get_max_move(self):
        resp = self.command(b"m", 2)
        log.info(f"Camera max move: {resp[0]} {resp[1]}")
        return f"{resp[0]}{resp[1]}"

    def set_min_move(self):
        data = b"N" + struct.pack("<2B", *[15, 0])
        resp = self.command(data, 1)
        return True if resp == b"K" else False

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

    def set_y_aggressiveness(self):
        # 90 / 255*100
        data = b"Y" + struct.pack("<1B", *[90])
        resp = self.command(data, 1)
        return True if resp == b"K" else False

    def get_y_aggressiveness(self):
        resp = self.command(b"y", 1)
        value = ord(resp)/255*100
        log.info(f"Camera y aggressiveness: {value:.2f}")
        return round(value, 2)

    def set_x_aggressiveness(self):
        # 90 / 255*100
        data = b"Z" + struct.pack("<1B", *[90])
        resp = self.command(data, 1)
        return True if resp == b"K" else False

    def get_x_aggressiveness(self):
        resp = self.command(b"z", 1)
        value = ord(resp)/255*100
        log.info(f"Camera x aggressiveness: {value:.2f}")
        return round(value, 2)

    def heater_on(self):
        cmd = b'g' + struct.pack("B", 1)
        log.info("Turning camera heater on")
        self.command(cmd)
        return True

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

    def shutter_chop_on(self):
        cmd = b'U' + struct.pack("B", 1)
        log.info("Enabling chopping shutter mode.")
        self.command(cmd)
        return True

    def shutter_chop_off(self):
        cmd = b'U' + struct.pack("B", 0)
        log.info("Disabling chopping shutter mode.")
        self.command(cmd)
        return True

    def define_subframe(self, frametyp="custom", cx=123, cy=234, size=127):

        params = frame_type[frametyp]
        if frametype == "custom":
            x = struct.pack('<h', cx)
            y = struct.pack('<h', cy)

        combuf = bytearray(6)
        combuf[0] = ord("S")
        combuf[1] = x[1]
        combuf[2] = x[0]
        combuf[3] = y[1]
        combuf[4] = y[0]
        combuf[5] = self.expected

        self.command(combuf)
        return True

    def take_image(self, exposure, light=True, binning=1):

        if binning == 1:
            print(binning)
            if self.height == 480:
                self.expected = 4096
                print(self.height)
                if self.width == 640:
                    log.info("Full frame")
                    frame = 0x00
                elif self.width == 512:
                    log.info("Crop frame")
                    frame = 0x01
            else:
                log.info("Windowing!")
                frame = 0xff

        else:
            log.warning("Binning 2")
            self.expected = 1024
            self.width = 320
            self.height = 240
            frame = 0x02

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

        #exp = int(exptime).to_bytes(3, byteorder='little')
        exp = struct.pack("<I", int(exptime))[:3]

        combuf = bytearray(6)
        combuf[0] = ord("T")
        combuf[1] = exp[2]
        combuf[2] = exp[1]
        combuf[3] = exp[0]
        combuf[4] = frame  # 0x01 cropped, 0x00 full, 0x02 binning 2
        combuf[5] = 0x01 if light else 0x00

        log.info("Starting exposure")
        self.command(combuf)

        state = {
            b"E": "Expose",
            b"R": "Readout",
            b"": "Reading",
            b"D": "Done",
        }

        out = b"E"
        while out != b"D":
            time.sleep(0.1)
            out = self.ser.read(1)
            try:
                log.info(state[out])
            except:
                log.warning(f"Wired state not in dict: {out}")

    def transfer_block(self, tries=3, cs=1):
        """
        Get one 'block' of image data. At fullframe 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)
        """

        data = self.ser.read(self.expected*PIXEL_SIZE)
        ndata = len(data)
        log.info(f"{ndata}")
        csum_byte = self.ser.read(1)

        try:
            csum_byte = ord(csum_byte)
        except TypeError as e:
            log.error(f"Empty csum_byte: {csum_byte}")
            self.status = False
            return data

        # two ways to calculate checksum byte
        if cs == 1:
            csum = 0
            for byte in data:
                csum = csum ^ byte
        elif cs == 2:
            csum = 0
            for byte in data:
                csum ^= (byte & 0xff)
                csum &= 0xff

        if csum == csum_byte and ndata == self.expected*PIXEL_SIZE:
            log.info(f"Read to Kontinue after {ndata}")
            self.status = True
        else:
            log.warning(
                f"Read to Resend: checksum {csum == csum_byte}, data {ndata}")
            self.status = False

        return data

    def download(self, max_tries=5):

        # Download Image
        # blocks_expected = int((320 * 240) / expected )
        log.info(f"{self.width}, {self.height}, {self.expected}")
        blocks_expected = (self.width * self.height) // self.expected

        # Download Image
        data = b""

        self.command(b"X")

        done = False
        blocks_complete = 0
        tries = 0
        while not blocks_complete == blocks_expected and tries < max_tries:
            time.sleep(self.wait())
            datum = self.transfer_block()

            if self.status:
                data += datum
                self.ser.write(b"K")
                blocks_complete += 1
                tries = 0
            else:
                tries += 1
                self.ser.write(b"R")

            log.info(f"*** {blocks_complete} of {blocks_expected}")
            if blocks_complete == blocks_expected:
                log.info("Done")
            if tries >= max_tries:
                log.error("Max tries exceeded")

        self.ser.write(b"S")
        return data

    def create_fits(self, data):

        header = fits.Header()

        # Now make into a fits image
        data = np.frombuffer(data, dtype=np.uint16)
        data = data.reshape((self.height, self.width))

        fits.writeto("ciccio.fits", data, header=header, overwrite=True)


if __name__ == '__main__':
    #    c = AllSky340(device="socket://151.3.45.60:8899")

    c = AllSky340(device="/dev/ttyr00")

    c.abort()
    c.abort()
    c.abort()

    c.open_shutter()
    c.take_image(1.2345)
    c.close_shutter()
    c.shutter_off()
    c.download()
    c.create_fits()

    # head = fits.Header()

    # # Now make into a fits image
    # #data = np.fromstring(data, dtype=np.uint16)
    # data = np.frombuffer(data, dtype=np.uint16)
    # data = data.reshape((480, 640))
