#!/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
#


import OBSUTILS as obsutil
import argparse as ap
import sys
import numpy as np
import subprocess as sp
import socket


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

    sock.send("goto={:g}".format(foc_val))
    resp = sock.recv(20)

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

    sock.close()

    return 0


def get_image(filename, exp_time):
    cmd = ["ixonultra_cmdclient", "-A", "CCD", "-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=sys.argv[0])

    parser.add_argument("focus_start", help="focus start value", type=float)
    parser.add_argument("focus_stop", help="focus stop value", type=float)
    parser.add_argument(
        "focus_step", help="focus step value", type=float, nargs="?", default=np.nan
    )

    parser.add_argument(
        "--guess",
        help="Use of two-steps focussing strategy. Input focus start and stop values are interpretated"
        + "as a guess range while step value is one for precise measurements.",
        action="store_true",
        default=False,
    )

    parser.add_argument(
        "-N",
        "--num",
        help="Number of points for precise stage measurement. Default is 7. Ignored if no '--guess'.",
        type=np.uint64,
        default=7,
    )

    parser.add_argument(
        "-v", "--verbose", help="Verbose output", action="store_true", default=False
    )

    parser = obsutil.getFocussingSequenceCmdlinePars(parser)

    args = parser.parse_args()

    if (args.num < 3) and (args.guess):
        if args.verbose:
            print("Number of focus values must be greater than 2! Abort!")
        sys.exit(-1)

    focus_start = float(args.focus_start)
    focus_stop = float(args.focus_stop)
    focus_step = float(args.focus_step)

    if np.isfinite(focus_step):
        if np.isclose(focus_step, 0.0):
            if args.verbose:
                print("Invalid (zero) focus step value! Abort!")
            sys.exit(1)

        if np.isclose(focus_start, focus_stop):
            if args.verbose:
                print("Invalid focus range parameters (an empty range)! Abort!")
            sys.exit(1)

        if (focus_stop - focus_start) * focus_step < 0:
            if args.verbose:
                print("Invalid focus range parameters! Abort!")
            sys.exit(1)

        # if (focus_start > focus_stop) and (focus_step > 0):
        #     if args.verbose:
        #         print("Invalid focus range parameters! Abort!")
        #     sys.exit(1)

        # if (focus_start < focus_stop) and (focus_step < 0):
        #     if args.verbose:
        #         print("Invalid focus range parameters! Abort!")
        #     sys.exit(1)
    else:
        if args.guess:
            if args.verbose:
                print("If '--guess' is set then step value must be given! Abort!")
            exit(1)

    log = []

    seq_kwds = obsutil.parsFocussingSequenceCmdlinePars(args, log)

    if args.verbose:
        if len(log):
            for log_str in log:
                print(log_str)

    if args.verbose:
        print("START FOCUSSING SEQUENCE ...")

    if args.guess:  # rough focus estimation
        if args.verbose:
            print("\tStart rough focus estimation:")
        result = obsutil.focussingSequence(
            [focus_start, focus_stop],
            set_focus,
            get_image,
            None,
            sys.stdout,
            **seq_kwds
        )

        if result["ret_code"] != 0:
            print("\n\tERROR: Cannot perform rough focus estimation! Abort!")
            print("\nFAIL!")
            sys.exit(result["ret_code"])

        # compute precise range start and stop values
        r = args.num * focus_step / 2.0
        focus_start = result["focus_value"] - r
        focus_stop = result["focus_value"] + r

    if args.guess:  # rough focus estimation
        if args.verbose:
            print("\n\tStart precise focus measurements:")

    if np.isfinite(args.focus_step):
        result = obsutil.focussingSequence(
            [focus_start, focus_stop, focus_step],
            set_focus,
            get_image,
            None,
            sys.stdout,
            **seq_kwds
        )
    else:
        result = obsutil.focussingSequence(
            [focus_start, focus_stop],
            set_focus,
            get_image,
            None,
            sys.stdout,
            **seq_kwds
        )

    if args.verbose:
        print("\tThe best focus value: {:g}".format(result["focus_value"]))

    if not args.do_no_set:
        if args.verbose:
            print("\n\tSet focus to the best value ...", end="")
            set_focus(result["focus_value"])
            print("\tOK")

    if args.verbose:
        print("DONE.")

    sys.exit(0)