#!/usr/bin/env python2
# -*- coding: utf-8 -*-


import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
from scipy.integrate import odeint

# SymPi Preamble
from sympy import *

s = symbols('s')
init_printing(use_unicode=False, wrap_line=False)


def extract_numerator_and_denominator_from_sym_py_poly_fraction(polyFraction):
    sym_numerator, sym_denominator = fraction(polyFraction)

    try:
        sym_numerator = Poly(sym_numerator).all_coeffs()
    except:
        sym_numerator = [sym_numerator]

    try:
        sym_denominator = Poly(sym_denominator).all_coeffs()
    except:
        sym_denominator = [sym_denominator]

    numerator = [float(x) for x in sym_numerator]
    denominator = [float(x) for x in sym_denominator]

    return numerator, denominator


# Define input signal
signalTime = np.arange(0, 3, 0.001)
wo = 1 * np.pi
signalDemand = np.zeros(signalTime.shape, dtype=float)
signalDemand[int(len(signalDemand)/4):] = 1.0

# Define controller parameters
Kp = 10.0
Ki = 10.0
Kd = 5.00

# PID Controller
Cs = (Kd * s ** 2 + Kp * s + Ki) / (
            s * (s / 1000000 + 1))  # Had to add a very quick pole to make the controller propper

# Plant model - Motor controller (request a velocity, get a position)
motorControllerBandwidth = 1000 * 2 * np.pi
Ps = 1 / (s * (s / motorControllerBandwidth + 1))

# Closed Loop Dynamics - Closed loop transfer function T
Ts = simplify(Cs * Ps / (1 + Cs * Ps))
Ts_numerator, Ts_denominator = extract_numerator_and_denominator_from_sym_py_poly_fraction(Ts)
Ts_transferFunction = signal.TransferFunction(Ts_numerator, Ts_denominator);

# Transfer fucntion for Demand sent to Motor controller Demand (MD)
MDs = simplify(Cs / (1 + Cs * Ps))
MDs_numerator, MDs_denominator = extract_numerator_and_denominator_from_sym_py_poly_fraction(MDs)
MDs_transferFunction = signal.TransferFunction(MDs_numerator, MDs_denominator);

# Open Loop Transfer Function G
Gs = simplify(Cs * Ps)
Gs_numerator, Gs_denominator = extract_numerator_and_denominator_from_sym_py_poly_fraction(Gs)
Gs_transferFunction = signal.TransferFunction(Gs_numerator, Gs_denominator);

# Simulate Responses
tout, signalSystemOutput, x = signal.lsim(Ts_transferFunction, signalDemand, signalTime)
tout, signalMotorDemand, x = signal.lsim(MDs_transferFunction, signalDemand, signalTime)

wRange = np.arange(1, 100000, 1)
w, magTs, phaseTs = signal.bode(Ts_transferFunction, wRange)
w, magMDs, phaseMDs = signal.bode(MDs_transferFunction, wRange)
w, magGs, phaseGs = signal.bode(Gs_transferFunction, wRange)

plt.figure(1)
plt.clf()
ax1 = plt.subplot(2, 1, 1)
plt.title("Demand and Response")
plt.plot(tout, signalDemand, label="Demand")
plt.plot(tout, signalSystemOutput, label="Response")
plt.xlabel("Time Seconds")
plt.ylabel("Position (Radians)")
plt.legend()
plt.grid()
ax1 = plt.subplot(2, 1, 2)
plt.title("Plant/Motor Velocity Demands")
plt.plot(tout, signalMotorDemand, label="Motor Demand")
plt.xlabel("Time Seconds")
plt.ylabel("Velocity (Radians Per Second)")
plt.grid()
plt.show()

# Bode plot closed loop response
plt.figure(2)
plt.clf()
ax1 = plt.subplot(2, 1, 1)
plt.title("Closed Loop Gains")
plt.semilogx(wRange, magTs, label="Closed Loop Transfer Function")
plt.semilogx(wRange, magMDs, label="Closed Loop Plant/Motor Demand")
plt.ylabel("-20log(|T(s)|) (dB)")
plt.legend()
plt.grid()
ax1 = plt.subplot(2, 1, 2)
plt.title("Closed Loop Phases")
plt.semilogx(wRange, phaseTs, label="Closed Loop Transfer Function")
plt.semilogx(wRange, phaseMDs, label="Closed Loop Plant/Motor Demand")
plt.xlabel("W (rads per second)")
plt.ylabel("phase T(s) (Degrees)")
plt.legend()
plt.grid()
plt.show()

# Nyquist Diagram for Open Loop System
plt.figure(3)
plt.clf()
plt.title("Open Loop Nyquist")
plt.polar(phaseGs / 180 * np.pi, np.power(10, -magGs / 20), label="Open Loop Nyquist")  # Convert from decibels
plt.show()
