#!/usr/bin/python

#
#   Get sequence of focussing images and
#   compute the best focus value
#
#   This variant uses Fatkhullin's AndorAPI2 library
#   to control Andor IXon Ultra EMCCD camera
#   The script uses TCP-socket Python API to control focus
#   of the SAO RAS Zeiss-1000 telescope
#


from OBSUTILS import focussing_app
import sys
import argparse as ap
import subprocess as sp
import numpy as np
import socket
import pathlib as pl
import time


def init_seq(seq_kwds):
    pt = pl.Path(seq_kwds["root_filename"])

    foc_files = list(
        pt.absolute().parent.glob("{}_[0-9][0-9][0-9][0-9].fit".format(pt.stem))
    )

    seq_kwds["start_idx"] = 1
    N = len(foc_files)
    if N > 0:
        max = 0
        for fname in foc_files:
            inum = int(fname.stem.split("_")[1])
            if inum > max:
                max = inum

        seq_kwds["start_idx"] = max + 1

    return 0


def set_focus(foc_val):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect(("ztcs.sao.ru", 4444))

    bt = bytearray("goto={:g}".format(foc_val), 'utf8')

    sock.send(bt)
    resp = sock.recv(20).decode('utf8')

    if resp != "OK":
        sock.close()
        return 1

    tp_start = time.monotonic()

    while True:
        time.sleep(0.5)
        sock.send(b'status')
        resp = sock.recv(20).decode('utf8')
        if resp == 'OK':
            break;

        if (time.monotonic()-tp_start) > 30.0:
             sock.close() # timeout!
             return 2

    sock.close()

    return 0


#
#   Assumes Andor IXon EMCCD server is listenning UNIX socket on localhost
#
def get_image(filename, exp_time):
    cmd = ["ixonultra_cmdclient", "-A", "CCD", "-p", "1", "--hs", "0", "-T", "{:g}".format(exp_time), filename]

    ret = sp.run(cmd, stdout=sp.PIPE, stderr=sp.PIPE)
    return ret.returncode


if __name__ == "__main__":
    parser = ap.ArgumentParser(
        prog="{}".format(pl.Path(sys.argv[0]).name),
        description="Zeiss-1000 and Andor Ixon EMCCD hardware: focussing sequence implementation.",
    )

    ret = focussing_app(parser, init_seq, set_focus, get_image)

    sys.exit(ret)