# -*- coding: utf-8 -*-
"""
Created on Jun 2 2021

@author: Alex
"""
import numpy as np

from pyro.dynamic               import massspringdamper
from pyro.analysis.costfunction import QuadraticCostFunction
from pyro.control.lqr           import synthesize_lqr_controller


# Plant
sys = massspringdamper.TwoMass()

sys.m1 = 2
sys.m2 = 3
sys.k1 = 5
sys.k2 = 2
sys.b1 = 0.1
sys.b2 = 0.2

sys.compute_ABCD()

#Full state feedback (default of class is x2 output only)
sys.C = np.diag([1,1,1,1])
sys.p = 4 # dim of output vector
sys.output_label = sys.state_label
sys.output_units = sys.state_units

# Cost function
cf  = QuadraticCostFunction.from_sys( sys )
cf.Q[0,0] = 5
cf.Q[1,1] = 10
cf.Q[2,2] = 5
cf.Q[3,3] = 2
cf.R[0,0] = 0.01
sys.cost_function = cf

# LQR controller
ctl = synthesize_lqr_controller( sys , cf )
ctl.ref_label = ['x1d','x2d','dx1d','dx2d']

# Simulation Closed-Loop Non-linear with LQR controller
cl_sys = ctl + sys
cl_sys.x0 = np.array([0,1,0,0])
cl_sys.compute_trajectory()
cl_sys.plot_trajectory('xu')
cl_sys.plot_linearized_pz_map(1,1)
cl_sys.plot_linearized_bode(1,1)
cl_sys.plot_linearized_pz_map(0,0)
cl_sys.plot_linearized_bode(0,0)
cl_sys.animate_simulation()