"""
03_tune_controller
------------------ 
Load a turbine model and tune the controller
In this example:

* Read a .yaml file
* Create a ROSCO turbine object fromt he OpenFAST model
* Tune a ROSCO controller object
* Write a controller input file (``'03_DISCON.IN'``)
* Plot gain schedule (``PC_GS_KP`` and ``PC_GS_KI`` versus ``PS_GS_angles``):

.. figure:: /images/examples/03_GainSched.png
   :align: center
   :width: 70%

"""

# Python modules
import matplotlib.pyplot as plt 
import os 
# ROSCO toolbox modules 
from rosco.toolbox import controller as ROSCO_controller
from rosco.toolbox import turbine as ROSCO_turbine
from rosco.toolbox.utilities import write_DISCON
from rosco.toolbox.inputs.validation import load_rosco_yaml

this_dir = os.path.dirname(os.path.abspath(__file__))
tune_dir =  os.path.join(this_dir,'Tune_Cases')
example_out_dir = os.path.join(this_dir,'examples_out')
os.makedirs(example_out_dir, exist_ok=True)

def main():
    # Inputs: tuning yaml (will point to openfast model file)
    parameter_filename = os.path.join(tune_dir,'IEA15MW.yaml')
    
    # Output: controller DISCON input file
    param_file = os.path.join(example_out_dir,'IEA15MW_DISCON.IN')
    
    # Load yaml file 
    inps = load_rosco_yaml(parameter_filename)
    path_params         = inps['path_params']
    turbine_params      = inps['turbine_params']
    controller_params   = inps['controller_params']
    yaml_dir = os.path.dirname(parameter_filename)

    # Instantiate turbine, controller, and file processing classes
    turbine         = ROSCO_turbine.Turbine(turbine_params)
    controller      = ROSCO_controller.Controller(controller_params)

    # Load turbine data from OpenFAST and rotor performance text file
    cp_filename = os.path.join(yaml_dir,path_params['rotor_performance_filename'])
    turbine.load_from_fast(
        path_params['FAST_InputFile'],
        os.path.join(yaml_dir,path_params['FAST_directory']),
        rot_source='txt',txt_filename= cp_filename
        )

    # Tune controller 
    controller.tune_controller(turbine)

    # Write parameter input file
    write_DISCON(
        turbine,controller,
        param_file=param_file, 
        txt_filename=cp_filename
        )

    # Plot gain schedule
    fig, ax = plt.subplots(2,2,constrained_layout=True,sharex=True)
    ax = ax.flatten()
    ax[0].plot(controller.v[len(controller.v_below_rated)+1:], controller.omega_pc_U)
    ax[0].set_ylabel('omega_pc')

    ax[1].plot(controller.v[len(controller.v_below_rated)+1:], controller.zeta_pc_U)
    ax[1].set_ylabel('zeta_pc')

    ax[2].plot(controller.v[len(controller.v_below_rated)+1:], controller.pc_gain_schedule.Kp)
    ax[2].set_xlabel('Wind Speed')
    ax[2].set_ylabel('Proportional Gain')

    ax[3].plot(controller.v[len(controller.v_below_rated)+1:], controller.pc_gain_schedule.Ki)
    ax[3].set_xlabel('Wind Speed')
    ax[3].set_ylabel('Integral Gain')

    plt.suptitle('Pitch Controller Gains')
        

    if False:
        plt.show()
    else:
        plt.savefig(os.path.join(example_out_dir,'03_GainSched.png'))

if __name__ == "__main__":
    main()