#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Allows to tune PID values for closed loop moves of Physik Instrumente controllers.
# It allows to do a move while recording commanded and actual position, and
# update the PID parameter values.
# Example way to start:
# ./util/pituner --port autoip --controller 1
# To test:
# ./util/pituner --port /dev/fake --controller 1

# Note for tuning:
#  * Start tuning with P, with I & D = 0, then tune I, and then D.
#  * (Temporarily) Reduce the settle window/increase the settle time
#  * Try with various magnitudes of move distances, in both directions
#  * Check on different parts of the axis.
'''
Created on February 2016

@author: Éric Piel

Copyright © 2016 Éric Piel, Delmic

pituner is free software: you can redistribute it and/or modify it under the terms
of the GNU General Public License version 2 as published by the Free Software
Foundation.

pituner is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
PURPOSE. See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with
pituner. If not, see http://www.gnu.org/licenses/.
'''

from __future__ import division, absolute_import, print_function

import argparse
from builtins import input
import logging
import math
import numpy
from odemis.driver import pigcs
from odemis.util import driver
import sys
import time

import matplotlib.pyplot as plt


class PITuner(object):
    def __init__(self, cont, axis):
        """
        cont: the PIGCS controller
        axis (str): axis number
        """
        self.cont = cont
        self.axis = axis

        # print(cont._sendQueryCommand("HDR?\n"))

        # TODO: set different values based on the type of controller
        self._tablelen = 1024  # True on E-861 (but could be read using HDR?)
        # Doc says 20µs but actually it's  50µs
        self._cycledur = 50e-6  # s, for the E-861
        self._upm = cont._upm[axis]  # ratio : user unit -> m

        # Hack to force the servo to be always on.
        # We could try to init with auto_suspend = False, but that only works with
        # some type of controllers.
        if hasattr(cont, "_acquireAxis"):
            cont._acquireAxis(axis)

        # Configure the recording options
        # 1 = commanded position
        # 2 = actual position
        # 3 = position error
        # 73 = motor output
        cont.SetRecordConfig(1, axis, 1)
        cont.SetRecordConfig(2, axis, 2)
        try:
            self.cont.checkError()
        except Exception:
            logging.warning("Failed to set record config", exc_info=True)

        # distance to move (negative to move backward)
        self._distm = 10e-6  # m

        # Get speed and acceleration
        try:
            self._vel = cont.GetCLVelocity(axis) * self._upm
            self._acc = cont.GetCLAcceleration(axis) * self._upm
            # Set the deceleration same as the acceleration
            self.cont.SetCLDeceleration(self.axis, self._acc / self._upm)
            self.cont.checkError()
        except Exception:
            logging.warning("Failed to read velocity/acceleration", exc_info=True)
            self._vel = None
            self._acc = None

    def tune(self, show_vel=False):
        """
        Main loop for the tuning process: asks the user, does one move, show
          the result... rinse and repeat.
        show_vel (bool): True to also display the measured velocity
        """

        while self.ask_next_move():
            if self._vel is not None:
                movedur = driver.estimateMoveDuration(abs(self._distm), self._vel, self._acc)
            else:
                # Fill up by not too crazy velocity/acceleration
                movedur = driver.estimateMoveDuration(abs(self._distm), 2e-3, 3e-3)
            logging.info("Expected move time = %g s", movedur)

            # Adjust record rate to fit the move duration
            rrate = max(1, int(math.ceil((movedur + min(movedur, 2) + 0.5) / (self._tablelen * self._cycledur))))
            self.cont.SetRecordRate(rrate)
            recdur = self._tablelen * self._cycledur * rrate
            logging.info("Recording every %d cycle => duration = %g s", rrate, recdur)

            # Start moving and recording
            self.cont.MoveRelRecorded(self.axis, self._distm / self._upm)
            tstart = time.time()
            tendmax = tstart + 2 * recdur + 1

            # Wait until the move is done
            while True:
                time.sleep(0.01)

                try:
                    if self.cont.IsOnTarget(self.axis, check=True):
                        logging.debug("Move finished after %g s", time.time() - tstart)
                        break
                except pigcs.PIGCSError as exp:
                    logging.warning("Controller reports error %s", exp)
                    break

                if time.time() > tendmax:
                    logging.warning("Controller still moving after %g s, stopping it", tendmax - tstart)
                    self.cont.Stop()
                    break

            # Wait until all the data is recorded
            left = tstart + recdur - time.time()
            if left > 0:
                time.sleep(left)

            # Read the recorded data from the move
            data = self.cont.GetRecordedData()

            # plot data using matplotlib
            idx = numpy.linspace(0, (len(data) - 1) * self._cycledur * rrate, len(data))
            cp, ap = numpy.array(zip(*data)) * self._upm * 1e6  # -> µm
            ep = ap - cp
            nb_plots = 2
            if show_vel:
                # Compute actual velocity
                av = (ap[1:] - ap[:-1]) / (self._cycledur * rrate)
                nb_plots += 1

            f, ax = plt.subplots(nb_plots, sharex=True)
            ax[0].plot(idx, cp, 'b-', idx, ap, 'r-')
            ax[0].set_title(u'Commanded vs. Actual position (µm)')
            ax[1].plot(idx, ep, 'b-')
            ax[1].set_title(u'Position Error (µm)')
            if show_vel:
                ax[2].plot(idx[1:], av, 'b-')
                ax[2].set_title(u'Actual velocity (µm/s)')

            plt.show()

            # print("# Commanded pos\tActual pos")
            # for cp, ap in data:
            #    print("%g\t%g" % (cp, ap))

    def ask_next_move(self):
        """
        Ask the user for the size of the next move, or to update the settings of
        the controller.
        return (bool): False if needs to stop, True if a move is requested
        """
        P = self.cont.GetParameter(self.axis, 1)
        I = self.cont.GetParameter(self.axis, 2)
        D = self.cont.GetParameter(self.axis, 3)
        print("P=%s, I=%s, D=%s" % (P, I, D))
        if self._vel is not None:
            print(u"Vel=%g µm/s, Acc/Dec=%g µm/s²" % (self._vel * 1e6, self._acc * 1e6))
        print(u"Current move distance: %g µm" % (self._distm * 1e6,))

        while True:
            print("Change P, I, D, (M)ove distance, (V)elocity, (A)cceleration, (Q)uit or press Enter to start next move: ", end="")
            choice = input().upper()
            if choice in ("P", "I", "D"):
                val = input("Enter new value for %s: " % (choice,))
                param = {"P": 1, "I": 2, "D": 3}[choice]
                try:
                    self.cont.SetParameter(self.axis, param, val)
                except Exception:
                    logging.exception("Failed to write %s to parameter 0x%x", val, param)
            elif choice == "V":
                print(u"Enter new velocity in µm/s: ", end="")
                val = input()
                self._vel = float(val) * 1e-6
                self.cont.SetCLVelocity(self.axis, self._vel / self._upm)
                self.cont.checkError()
            elif choice == "A":
                print(u"Enter new acceleration/deceleration in µm/s²: ", end="")
                val = input()
                self._acc = float(val) * 1e-6
                self.cont.SetCLAcceleration(self.axis, self._acc / self._upm)
                self.cont.SetCLDeceleration(self.axis, self._acc / self._upm)
                self.cont.checkError()
            elif choice == "M":
                print(u"Enter new distance in µm (can be negative): ", end="")
                val = input()
                self._distm = float(val) * 1e-6
            elif choice == "":
                return True
            elif choice == "Q":
                return False
            else:
                # Just keep asking
                print("Command '%s' not understood" % (choice,))


def main(args):
    """
    Handles the command line arguments
    args is the list of arguments passed
    return (int): value to return to the OS as program exit code
    """

    # arguments handling
    parser = argparse.ArgumentParser(prog="pituner",
                                     description="PI PID tuner")

    parser.add_argument("--log-level", dest="loglev", metavar="<level>", type=int,
                        default=1, help="set verbosity level (0-2, default = 1)")

    parser.add_argument('--port', dest="port", required=True,
                        help="Port name (ex: /dev/ttyUSB0, autoip, or 192.168.95.5)")
    parser.add_argument('--controller', dest="addr", type=int,
                        help="Controller address (if controller needs it)")
    parser.add_argument("--velocity", "-s", dest="vel", action="store_true", default=False,
                        help="Also shows velocity based on the measured position.")

    options = parser.parse_args(args[1:])

    # Set up logging before everything else
    if options.loglev < 0:
        logging.error("Log-level must be positive.")
        return 127
    loglev_names = (logging.WARNING, logging.INFO, logging.DEBUG)
    loglev = loglev_names[min(len(loglev_names) - 1, options.loglev)]
    logging.getLogger().setLevel(loglev)

    try:
        kwargs = {}
        if options.addr is None:
            # If no address, there is also no master (for IP)
            kwargs["master"] = None

        if options.port == "/dev/fake":
            kwargs["_addresses"] = {options.addr: True}
            acc = pigcs.FakeBus._openPort(options.port, **kwargs)
        else:
            acc = pigcs.Bus._openPort(options.port, **kwargs)

        # TODO: allow to specify the axis
        cont = pigcs.Controller(acc, address=options.addr, axes={"1": True})
        tuner = PITuner(cont, "1")
        tuner.tune(options.vel)

        cont.terminate()
        acc.terminate()
    except ValueError as exp:
        logging.error("%s", exp)
        return 127
    except IOError as exp:
        logging.error("%s", exp)
        return 129
    except Exception:
        logging.exception("Unexpected error while performing action.")
        return 130

    return 0


if __name__ == '__main__':
    ret = main(sys.argv)
    exit(ret)