"""
Analyze a PX4 log to perform sysid and control design.
"""
from __future__ import print_function

from collections import OrderedDict

import control
import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize
import scipy.signal


# pylint: disable=invalid-name, no-member, too-many-locals, too-many-arguments

# noinspection PyUnusedLocal,PyUnusedLocal,PyUnusedLocal
def lqr_ofb_con(K, R, Q, X, ss_o):
    # type: (np.array, np.array, np.array, np.array, control.ss) -> np.array
    """
    Constraint for LQR output feedback optimization.
    This asserts that all eigenvalues are negative so that
    the system is stable.
    @K gain matrix
    @Q process noise covariance matrix
    @X initial state covariance matrix
    @ss_o open loop state space system
    @return constraint
    """
    # pylint: disable=unused-argument
    K = np.matrix(K).T
    A = np.matrix(ss_o.A)
    B = np.matrix(ss_o.B)
    C = np.matrix(ss_o.C)
    A_c = A - B * K * C
    return -np.real(np.linalg.eig(A_c)[0])


def lqr_ofb_cost(K, R, Q, X, ss_o):
    # type: (np.array, np.array, np.array, np.array, control.ss) -> np.array
    """
    Cost for LQR output feedback optimization.
    @K gain matrix
    @Q process noise covariance matrix
    @X initial state covariance matrix
    @ss_o open loop state space system
    @return cost
    """
    K = np.matrix(K).T
    A = np.matrix(ss_o.A)
    B = np.matrix(ss_o.B)
    C = np.matrix(ss_o.C)
    A_c = A - B * K * C
    Q_c = C.T * K.T * R * K * C + Q
    P = scipy.linalg.solve_lyapunov(A_c.T, -Q_c)
    J = np.trace(P * X)
    return J


def lqr_ofb_jac(K, R, Q, X, ss_o):
    # type: (np.array, np.array, np.array, np.array, control.ss) -> np.array
    """
    Jacobian for LQR Output feedback optimization.
    TODO: might be an error here, doesn't not help optim
    """
    K = np.matrix(K).T
    A = np.matrix(ss_o.A)
    B = np.matrix(ss_o.B)
    C = np.matrix(ss_o.C)
    A_c = A - B * K * C
    Q_c = C.T * K.T * R * K * C + Q
    P = scipy.linalg.solve_lyapunov(A_c.T, -Q_c)
    S = scipy.linalg.solve_lyapunov(A_c, -X)
    J = 2 * (R * K * C * S * C.T - B.T * P * S * C.T)
    J = np.array(J)[:, 0]
    return J


def lqr_ofb_design(K_guess, ss_o, verbose=False):
    """
    LQR output feedback controller design.
    @K_guess initial stabilizing gains
    @ss_o open loop state space system
    @return gain matrix
    """
    n_x = ss_o.A.shape[0]
    n_u = ss_o.B.shape[1]
    R = 1e-6 * np.eye(n_u)
    Q = np.eye(n_x)
    X = 1e-3 * np.eye(n_x)

    constraints = [
        {'type': 'ineq',
         'fun': lqr_ofb_con,
         'args': (R, Q, X, ss_o),
         }
    ]

    res = scipy.optimize.minimize(
        fun=lqr_ofb_cost,
        method='SLSQP',
        args=(R, Q, X, ss_o),
        x0=K_guess,
        # jac=lqr_ofb_jac,
        bounds=len(K_guess) * [[1e-6, 100]],
        constraints=constraints
    )
    K = np.matrix(res['x'])

    if verbose:
        print(res)
    if not res['success']:
        print('cost', lqr_ofb_cost(K, R, Q, X, ss_o))
        print('jac', lqr_ofb_jac(K, R, Q, X, ss_o))
        print('constraint', lqr_ofb_con(K, R, Q, X, ss_o))
        raise RuntimeError('optimization failed')

    return K.T


def attitude_sysid(y_acc, u_mix, verbose=False):
    # type: (np.array, np.array, bool) -> (control.tf, float, float)
    """
    roll/pitch system id assuming delay/gain model
    :param y_acc: (roll or pitch acceleration)
    :param u_mix: (roll or pitch acceleration command)
    :param verbose: show debug output
    :return: (G_ol, delay, k)
        G_ol: open loop plant
        delay: time delay (sec)
        k: gain
    """
    if verbose:
        print('solving for plant model ', end='')
    k, delay = delay_and_gain_sysid(y_acc, u_mix, verbose)
    if verbose:
        print('done')

    # open loop, rate output, mix input plant
    tf_acc = k * control.tf(*control.pade(delay, 1))  # order 1 approx
    # can neglect motor, pole too fast to matter compared to delay, not used in sysid either
    tf_integrator = control.tf((1), (1, 0))
    G_ol = tf_acc * tf_integrator

    return G_ol, delay, k


def pid_design(G, K_guess, d_tc, verbose=False, use_P=True, use_I=True, use_D=True):
    # type: (control.tf, np.array, float, bool, bool, bool, bool) -> (np.array, control.tf, control.tf)
    """
    :param G: transfer function
    :param K_guess: gain matrix guess
    :param d_tc: time constant for derivative
    :param verbose: show debug output
    :param use_P: use p gain in design
    :param use_I: use i gain in design
    :param use_D: use d gain in design
    :return: (K, G_comp, Gc_comp)
        K: gain matrix
        G_comp: open loop compensated plant
        Gc_comp: closed loop compensated plant
    """
    # compensator transfer function
    H = []
    if use_P:
        H += [control.tf(1, 1)]
    if use_I:
        H += [control.tf((1), (1, 0))]
    if use_D:
        H += [control.tf((1, 0), (d_tc, 1))]
    H = np.array([H]).T
    H_num = [[H[i][j].num[0][0] for i in range(H.shape[0])] for j in range(H.shape[1])]
    H_den = [[H[i][j].den[0][0] for i in range(H.shape[0])] for j in range(H.shape[1])]
    H = control.tf(H_num, H_den)

    # print('G', G)
    # print('H', H)

    ss_open = control.tf2ss(G * H)

    if verbose:
        print('optimizing controller')
    K = lqr_ofb_design(K_guess, ss_open, verbose)
    if verbose:
        print('done')

    # print('K', K)
    # print('H', H)
    G_comp = control.series(G, H * K)
    Gc_comp = control.feedback(G_comp, 1)

    return K, G_comp, Gc_comp


def plot_loops(name, G_ol, G_cl):
    # type: (str, control.tf, control.tf) -> None
    """
    Plot loops
    :param name: Name of axis
    :param G_ol: open loop transfer function
    :param G_cl: closed loop transfer function
    """
    plt.figure()
    plt.plot(*control.step_response(G_cl, np.linspace(0, 1, 1000)))
    plt.title(name + ' step response')
    plt.grid()

    plt.figure()
    control.bode(G_ol)
    print('margins', control.margin(G_ol))
    plt.subplot(211)
    plt.title(name + ' open loop bode plot')

    plt.figure()
    control.rlocus(G_ol, np.logspace(-2, 0, 1000))
    for pole in G_cl.pole():
        plt.plot(np.real(pole), np.imag(pole), 'rs')
    plt.title(name + ' root locus')
    plt.grid()


def attitude_control_design(
        name, y, u, rolling_mean_window=100,
        do_plot=False, verbose=False, d_tc=1.0 / 125):
    # type: (str, np.array, np.array, int, bool, bool, float) -> (np.array, control.tf)
    """
    Do sysid and control design for roll/pitch rate loops.
    :param name: 
    :param y: output
    :param u: input
    :param rolling_mean_window: number of samples in rolling mean
    :param do_plot: controls plotting
    :param verbose: show debug output
    :param d_tc: derivative time constant
    :return: (gain, closed loop transfer function)
    """
    K_guess = np.matrix([[0.1, 0.1, 0.001]]).T

    # remove bias
    y_bias = y.rolling(rolling_mean_window).mean()
    y_debiased = y - y_bias
    u_bias = u.rolling(rolling_mean_window).mean()
    u_debiased = u - u_bias

    G_ol, delay, k = attitude_sysid(
        y_debiased, u_debiased, verbose)

    K, G_ol_rate, G_cl_rate = pid_design(
        G_ol, K_guess, d_tc, verbose)

    if do_plot:
        plt.figure()
        u_debiased.plot(label='debiased input')
        u.plot(label='input')
        plt.legend()
        plt.xlabel('t, sec')
        plt.title(name + ' input')
        plt.grid()

        plt.figure()
        y_debiased.plot(label='debiased output')
        y.plot(label='output')
        plt.legend()
        plt.xlabel('t, sec')
        plt.title(name + ' output')
        plt.grid()

        plt.figure()
        plot_delay_and_gain_fit(k, delay, y_debiased, u_debiased)
        plt.title(name + ' fit')
        plt.grid()

        plt.figure()
        plot_loops(name, G_ol_rate, G_cl_rate)

    return K, G_cl_rate


def control_design(raw_data, do_plot=False, rolling_mean_window=100, verbose=False):
    # type: (pd.DataFrame, bool, int, bool) -> (OrderedDict, dict)
    """
    Design a PID controller from log file.
    :param raw_data: data frame
    :param do_plot: controls plotting
    :param rolling_mean_window: number of rolling mean samples
    :param verbose: show debug output
    :return: (gain ordered dict, local variables)
    """
    data, dt = setup_data(raw_data)
    d_tc = 1.0 / 125  # nyquist frequency of derivative in PID, (250 Hz/2)
    K_guess_att = np.matrix([[1.0]]).T

    roll_acc = data.ATT_RollRate.diff() / dt
    K_rollrate, G_cl_rollrate = attitude_control_design(
        'roll rate', roll_acc, data.ATTC_Roll,
        rolling_mean_window=rolling_mean_window,
        do_plot=do_plot, verbose=verbose, d_tc=d_tc)

    tf_integrator = control.tf((1), (1, 0))

    K_roll, G_roll, G_cl_roll = pid_design(
        G_cl_rollrate * tf_integrator, K_guess_att, d_tc,
        verbose=verbose, use_I=False, use_D=False)

    if do_plot:
        plt.figure()
        plot_loops('roll', G_roll, G_cl_roll)

    pitch_acc = data.ATT_PitchRate.diff() / dt
    K_pitchrate, G_cl_pitchrate = attitude_control_design(
        'pitch rate', pitch_acc, data.ATTC_Pitch,
        rolling_mean_window=rolling_mean_window,
        do_plot=do_plot, d_tc=d_tc, verbose=verbose)

    K_pitch, G_pitch, G_cl_pitch = pid_design(
        G_cl_pitchrate * tf_integrator, K_guess_att, d_tc,
        verbose=verbose, use_I=False, use_D=False)

    if do_plot:
        plt.figure()
        plot_loops('pitch', G_pitch, G_cl_pitch)

    if verbose:
        print('G_roll', G_roll)
        print('G_cl_roll', G_cl_roll)
        print('G_cl_rollrate', G_cl_rollrate)
        print('G_pitch', G_pitch)
        print('G_cl_pitch', G_cl_pitch)
        print('G_cl_pitchrate', G_cl_pitchrate)

    return OrderedDict([
        ('MC_ROLL_P', round(K_roll[0, 0], 3)),
        ('MC_ROLLRATE_P', round(K_rollrate[0, 0], 3)),
        ('MC_ROLLRATE_I', round(K_rollrate[1, 0], 3)),
        ('MC_ROLLRATE_D', round(K_rollrate[2, 0], 3)),
        ('MC_PITCH_P', round(K_pitch[0, 0], 3)),
        ('MC_PITCHRATE_P', round(K_pitchrate[0, 0], 3)),
        ('MC_PITCHRATE_I', round(K_pitchrate[1, 0], 3)),
        ('MC_PITCHRATE_D', round(K_pitchrate[2, 0], 3)),
    ]), locals()


def control_design_ulog(raw_data, do_plot=False, rolling_mean_window=100, verbose=False):
    # type: (pd.DataFrame, bool, int, bool) -> (OrderedDict, dict)

    """
    Design a PID controller from log file.
    TODO, debug this
    :param raw_data: data frame
    :param do_plot: controls plotting
    :param rolling_mean_window: number of rolling mean samples
    :param verbose: show debug output
    :return: (gain ordered dict, local variables)
    """
    data, dt = setup_data(raw_data)
    d_tc = 1.0 / 125  # nyquist frequency of derivative in PID, (250 Hz/2)
    K_guess_att = np.matrix([[1.0]]).T

    roll_acc = data.ATT_RollRate.diff() / dt
    K_rollrate, G_cl_rollrate = attitude_control_design(
        'roll rate', roll_acc, data.ATTC_Roll,
        rolling_mean_window=rolling_mean_window,
        do_plot=do_plot, verbose=verbose, d_tc=d_tc)

    tf_integrator = control.tf((1), (1, 0))

    K_roll, G_roll, G_cl_roll = pid_design(
        G_cl_rollrate * tf_integrator, K_guess_att, d_tc,
        verbose=verbose, use_I=False, use_D=False)

    if do_plot:
        plt.figure()
        plot_loops('roll', G_roll, G_cl_roll)

    pitch_acc = data.ATT_PitchRate.diff() / dt
    K_pitchrate, G_cl_pitchrate = attitude_control_design(
        'pitch rate', pitch_acc, data.ATTC_Pitch,
        rolling_mean_window=rolling_mean_window,
        do_plot=do_plot, d_tc=d_tc, verbose=verbose)

    K_pitch, G_pitch, G_cl_pitch = pid_design(
        G_cl_pitchrate * tf_integrator, K_guess_att, d_tc,
        verbose=verbose, use_I=False, use_D=False)

    if do_plot:
        plt.figure()
        plot_loops('pitch', G_pitch, G_cl_pitch)

    if verbose:
        print('G_roll', G_roll)
        print('G_cl_roll', G_cl_roll)
        print('G_cl_rollrate', G_cl_rollrate)
        print('G_pitch', G_pitch)
        print('G_cl_pitch', G_cl_pitch)
        print('G_cl_pitchrate', G_cl_pitchrate)

    return OrderedDict([
        ('MC_ROLL_P', round(K_roll[0, 0], 3)),
        ('MC_ROLLRATE_P', round(K_rollrate[0, 0], 3)),
        ('MC_ROLLRATE_I', round(K_rollrate[1, 0], 3)),
        ('MC_ROLLRATE_D', round(K_rollrate[2, 0], 3)),
        ('MC_PITCH_P', round(K_pitch[0, 0], 3)),
        ('MC_PITCHRATE_P', round(K_pitchrate[0, 0], 3)),
        ('MC_PITCHRATE_I', round(K_pitchrate[1, 0], 3)),
        ('MC_PITCHRATE_D', round(K_pitchrate[2, 0], 3)),
    ]), locals()
