/*
 * control_algorithm.c
 *
 *  Created on: Feb 20, 2016
 *      Author: ucart
 */

// This implemented modular quadrotor software implements a PID control algorithm

#include "control_algorithm.h"
#include "graph_blocks.h"
#include "PID.h"
#include "util.h"
#include "timer.h"

#define ROLL_PITCH_MAX_ANGLE 0.35 // 20 degrees
#define PWM_DIFF_BOUNDS 20000
#define VRPN_REFRESH_TIME 0.01f // 10ms

void connect_autonomous(parameter_t* ps) {
	struct computation_graph* graph = ps->graph;
	graph_set_source(graph, ps->pitch_pid, PID_SETPOINT, ps->x_pos_pid, PID_CORRECTION);
	graph_set_source(graph, ps->roll_pid, PID_SETPOINT, ps->y_pos_pid, PID_CORRECTION);
	//graph_set_source(graph, ps->mixer, MIXER_THROTTLE, ps->throttle_trim_add, ADD_SUM);
	graph_set_source(graph, ps->yaw_r_pid, PID_SETPOINT, ps->yaw_pid, PID_CORRECTION);
}

void connect_manual(parameter_t* ps) {
	struct computation_graph* graph = ps->graph;
	graph_set_source(graph, ps->pitch_pid, PID_SETPOINT, ps->rc_pitch, CONST_VAL);
	graph_set_source(graph, ps->roll_pid, PID_SETPOINT, ps->rc_roll, CONST_VAL);
	graph_set_source(graph, ps->mixer, MIXER_THROTTLE, ps->rc_throttle, CONST_VAL);
	graph_set_source(graph, ps->yaw_r_pid, PID_SETPOINT, ps->rc_yaw, CONST_VAL);
}

int control_algorithm_init(parameter_t * ps)
{
    struct computation_graph* graph = create_graph();
    ps->graph = graph;

    // Create all the PID blocks
    ps->roll_pid = graph_add_defined_block(graph, BLOCK_PID, "Roll PID");
    ps->pitch_pid = graph_add_defined_block(graph, BLOCK_PID, "Pitch PID");
    ps->yaw_pid = graph_add_defined_block(graph, BLOCK_PID, "Yaw PID");
    ps->roll_r_pid = graph_add_defined_block(graph, BLOCK_PID, "Roll Rate PID");
    ps->pitch_r_pid = graph_add_defined_block(graph, BLOCK_PID, "Pitch Rate PID");
    ps->yaw_r_pid = graph_add_defined_block(graph, BLOCK_PID, "Yaw Rate PID");
    ps->x_pos_pid = graph_add_defined_block(graph, BLOCK_PID, "X pos PID");
    ps->y_pos_pid = graph_add_defined_block(graph, BLOCK_PID, "Y pos PID");
    ps->alt_pid = graph_add_defined_block(graph, BLOCK_PID, "Altitude PID");
    ps->x_set = graph_add_defined_block(graph, BLOCK_CONSTANT, "X Setpoint"); // ID 9
    ps->y_set = graph_add_defined_block(graph, BLOCK_CONSTANT, "Y Setpoint");
    ps->alt_set = graph_add_defined_block(graph, BLOCK_CONSTANT, "Alt Setpoint");
    ps->yaw_set = graph_add_defined_block(graph, BLOCK_CONSTANT, "Yaw Setpoint");
    ps->throttle_trim = graph_add_defined_block(graph, BLOCK_CONSTANT, "Throttle trim");
    ps->throttle_trim_add = graph_add_defined_block(graph, BLOCK_ADD, "T trim add");

    // Create blocks for sensor inputs
    ps->cur_pitch = graph_add_defined_block(graph, BLOCK_CONSTANT, "Pitch"); // ID 15
    ps->cur_roll = graph_add_defined_block(graph, BLOCK_CONSTANT, "Roll");
    ps->cur_yaw = graph_add_defined_block(graph, BLOCK_CONSTANT, "Yaw");
    // Sensor trims
    ps->pitch_trim = graph_add_defined_block(graph, BLOCK_CONSTANT, "Pitch trim");
    ps->pitch_trim_add = graph_add_defined_block(graph, BLOCK_ADD, "Pitch trim add");

	// Yaw angular velocity PID
    // theta_dot is the angular velocity about the y-axis
    // phi_dot is the angular velocity about the x-axis
	// psi_dot is the angular velocity about the z-axis
	// These are calculated from using the gimbal equations
    ps->theta_dot = graph_add_defined_block(graph, BLOCK_CONSTANT, "dTheta");
    ps->phi_dot = graph_add_defined_block(graph, BLOCK_CONSTANT, "dPhi");
    ps->psi_dot = graph_add_defined_block(graph, BLOCK_CONSTANT, "dPsi");
    ps->clamp_d_pwmP = graph_add_defined_block(graph, BLOCK_BOUNDS, "P PWM Clamp");
    ps->clamp_d_pwmR = graph_add_defined_block(graph, BLOCK_BOUNDS, "R PWM Clamp");
    ps->clamp_d_pwmY = graph_add_defined_block(graph, BLOCK_BOUNDS, "Y PWM Clamp");

    // Create blocks for VRPN data
    ps->vrpn_x = graph_add_defined_block(graph, BLOCK_CONSTANT, "VRPN X");
    ps->vrpn_y = graph_add_defined_block(graph, BLOCK_CONSTANT, "VRPN Y");
    ps->vrpn_alt = graph_add_defined_block(graph, BLOCK_CONSTANT, "VRPN Alt");
    ps->vrpn_pitch = graph_add_defined_block(graph, BLOCK_CONSTANT, "VRPN Pitch");
    ps->vrpn_roll = graph_add_defined_block(graph, BLOCK_CONSTANT, "VRPN Roll");

    // Create blocks for RC controller
    ps->rc_pitch = graph_add_defined_block(graph, BLOCK_CONSTANT, "RC Pitch");
    ps->rc_roll = graph_add_defined_block(graph, BLOCK_CONSTANT, "RC Roll");
    ps->rc_yaw = graph_add_defined_block(graph, BLOCK_CONSTANT, "RC Yaw");
    ps->rc_throttle = graph_add_defined_block(graph, BLOCK_CONSTANT, "RC Throttle");

    ps->mixer = graph_add_defined_block(graph, BLOCK_MIXER, "Signal Mixer");

    ps->angle_time = graph_add_defined_block(graph, BLOCK_CONSTANT, "Ts_IMU");
    ps->pos_time = graph_add_defined_block(graph, BLOCK_CONSTANT, "Ts_VRPN");

    // Connect pitch PID chain
    // Trims
    graph_set_source(graph, ps->pitch_trim_add, ADD_SUMMAND1, ps->pitch_trim, CONST_VAL);
    graph_set_source(graph, ps->pitch_trim_add, ADD_SUMMAND2, ps->cur_pitch, CONST_VAL);
    // Controllers
    graph_set_source(graph, ps->pitch_r_pid, PID_SETPOINT, ps->pitch_pid, PID_CORRECTION);
    graph_set_source(graph, ps->pitch_r_pid, PID_CUR_POINT, ps->theta_dot, CONST_VAL);
    graph_set_source(graph, ps->pitch_r_pid, PID_DT, ps->angle_time, CONST_VAL);
    graph_set_source(graph, ps->pitch_pid, PID_CUR_POINT, ps->pitch_trim_add, ADD_SUM);
    //graph_set_source(graph, ps->pitch_pid, PID_CUR_POINT, ps->vrpn_pitch, CONST_VAL);
    graph_set_source(graph, ps->pitch_pid, PID_DT, ps->angle_time, CONST_VAL);

     // Connect roll PID chain
    graph_set_source(graph, ps->roll_r_pid, PID_SETPOINT, ps->roll_pid, PID_CORRECTION);
    graph_set_source(graph, ps->roll_r_pid, PID_CUR_POINT, ps->phi_dot, CONST_VAL);
    graph_set_source(graph, ps->roll_r_pid, PID_DT, ps->angle_time, CONST_VAL);
    graph_set_source(graph, ps->roll_pid, PID_CUR_POINT, ps->cur_roll, CONST_VAL);
    //graph_set_source(graph, ps->roll_pid, PID_CUR_POINT, ps->vrpn_roll, CONST_VAL);
    graph_set_source(graph, ps->roll_pid, PID_DT, ps->angle_time, CONST_VAL);

    // Connect yaw PID chain
    graph_set_source(graph, ps->yaw_r_pid, PID_SETPOINT, ps->rc_yaw, PID_CORRECTION);
    graph_set_source(graph, ps->yaw_r_pid, PID_CUR_POINT, ps->psi_dot, CONST_VAL);
    graph_set_source(graph, ps->yaw_r_pid, PID_DT, ps->angle_time, CONST_VAL);


    // X autonomous
    graph_set_source(graph, ps->x_pos_pid, PID_DT, ps->pos_time, CONST_VAL);
    graph_set_source(graph, ps->x_pos_pid, PID_CUR_POINT, ps->vrpn_x, CONST_VAL);
    graph_set_source(graph, ps->x_pos_pid, PID_SETPOINT, ps->x_set, CONST_VAL);
    // Y autonomous
    graph_set_source(graph, ps->y_pos_pid, PID_DT, ps->pos_time, CONST_VAL);
    graph_set_source(graph, ps->y_pos_pid, PID_CUR_POINT, ps->vrpn_y, CONST_VAL);
    graph_set_source(graph, ps->y_pos_pid, PID_SETPOINT, ps->y_set, CONST_VAL);
    // Alt autonomous
    graph_set_source(graph, ps->alt_pid, PID_DT, ps->pos_time, CONST_VAL);
    graph_set_source(graph, ps->alt_pid, PID_CUR_POINT, ps->vrpn_alt, CONST_VAL);
    graph_set_source(graph, ps->alt_pid, PID_SETPOINT, ps->alt_set, CONST_VAL);
    graph_set_source(graph, ps->alt_set, CONST_VAL, ps->vrpn_alt, CONST_VAL);
    graph_set_source(graph, ps->throttle_trim_add, ADD_SUMMAND1, ps->alt_pid, PID_CORRECTION);
    graph_set_source(graph, ps->throttle_trim_add, ADD_SUMMAND2, ps->throttle_trim, CONST_VAL);
    // Yaw autonomous
    graph_set_source(graph, ps->yaw_pid, PID_DT, ps->pos_time, CONST_VAL);
    graph_set_source(graph, ps->yaw_pid, PID_SETPOINT, ps->yaw_set, CONST_VAL);
    graph_set_source(graph, ps->yaw_pid, PID_CUR_POINT, ps->cur_yaw, CONST_VAL);

    // Connect PWM clamping blocks
    graph_set_source(graph, ps->clamp_d_pwmP, BOUNDS_IN, ps->pitch_r_pid, PID_CORRECTION);
    graph_set_source(graph, ps->clamp_d_pwmR, BOUNDS_IN, ps->roll_r_pid, PID_CORRECTION);
    graph_set_source(graph, ps->clamp_d_pwmY, BOUNDS_IN, ps->yaw_r_pid, PID_CORRECTION);

    // Connect signal mixer
    //graph_set_source(graph, ps->mixer, MIXER_THROTTLE, ps->rc_throttle, CONST_VAL);
    graph_set_source(graph, ps->mixer, MIXER_PITCH, ps->clamp_d_pwmP, BOUNDS_OUT);
    graph_set_source(graph, ps->mixer, MIXER_ROLL, ps->clamp_d_pwmR, BOUNDS_OUT);
    graph_set_source(graph, ps->mixer, MIXER_YAW, ps->clamp_d_pwmY, BOUNDS_OUT);

    // Set pitch PID constants
    graph_set_param_val(graph, ps->pitch_pid, PID_KP, PITCH_ANGLE_KP);
    graph_set_param_val(graph, ps->pitch_pid, PID_KI, PITCH_ANGLE_KI);
    graph_set_param_val(graph, ps->pitch_pid, PID_KD, PITCH_ANGLE_KD);
    graph_set_param_val(graph, ps->pitch_r_pid, PID_KP, PITCH_ANGULAR_VELOCITY_KP);
    graph_set_param_val(graph, ps->pitch_r_pid, PID_KI, PITCH_ANGULAR_VELOCITY_KI);
    graph_set_param_val(graph, ps->pitch_r_pid, PID_KD, PITCH_ANGULAR_VELOCITY_KD);

    // Set roll PID constants
    graph_set_param_val(graph, ps->roll_pid, PID_KP, ROLL_ANGLE_KP);
    graph_set_param_val(graph, ps->roll_pid, PID_KI, ROLL_ANGLE_KI);
    graph_set_param_val(graph, ps->roll_pid, PID_KD, ROLL_ANGLE_KD);
    graph_set_param_val(graph, ps->roll_r_pid, PID_KP, ROLL_ANGULAR_VELOCITY_KP);
    graph_set_param_val(graph, ps->roll_r_pid, PID_KI, ROLL_ANGULAR_VELOCITY_KI);
    graph_set_param_val(graph, ps->roll_r_pid, PID_KD, ROLL_ANGULAR_VELOCITY_KD);

    // Set yaw PID constants
    graph_set_param_val(graph, ps->yaw_pid, PID_KP, YAW_ANGLE_KP);
    graph_set_param_val(graph, ps->yaw_pid, PID_KI, YAW_ANGLE_KI);
    graph_set_param_val(graph, ps->yaw_pid, PID_KD, YAW_ANGLE_KD);
    graph_set_param_val(graph, ps->yaw_r_pid, PID_KP, YAW_ANGULAR_VELOCITY_KP);
    graph_set_param_val(graph, ps->yaw_r_pid, PID_KI, YAW_ANGULAR_VELOCITY_KI);
    graph_set_param_val(graph, ps->yaw_r_pid, PID_KD, YAW_ANGULAR_VELOCITY_KD);

    // Set X/Y/Alt constants
    graph_set_param_val(graph, ps->x_pos_pid, PID_KP, XPOS_KP);
    graph_set_param_val(graph, ps->x_pos_pid, PID_KI, XPOS_KI);
    graph_set_param_val(graph, ps->x_pos_pid, PID_KD, XPOS_KD);
    graph_set_param_val(graph, ps->x_pos_pid, PID_ALPHA, XPOS_ALPHA);
    graph_set_param_val(graph, ps->y_pos_pid, PID_KP, YPOS_KP);
    graph_set_param_val(graph, ps->y_pos_pid, PID_KI, YPOS_KI);
    graph_set_param_val(graph, ps->y_pos_pid, PID_KD, YPOS_KD);
    graph_set_param_val(graph, ps->y_pos_pid, PID_ALPHA, YPOS_ALPHA);
    graph_set_param_val(graph, ps->alt_pid, PID_KP, ALT_ZPOS_KP);
    graph_set_param_val(graph, ps->alt_pid, PID_KI, ALT_ZPOS_KI);
    graph_set_param_val(graph, ps->alt_pid, PID_KD, ALT_ZPOS_KD);

    // Set PWM difference clamping limits
    graph_set_param_val(graph, ps->clamp_d_pwmP, BOUNDS_MIN, -PWM_DIFF_BOUNDS);
    graph_set_param_val(graph, ps->clamp_d_pwmP, BOUNDS_MAX, PWM_DIFF_BOUNDS);
    graph_set_param_val(graph, ps->clamp_d_pwmR, BOUNDS_MIN, -PWM_DIFF_BOUNDS);
    graph_set_param_val(graph, ps->clamp_d_pwmR, BOUNDS_MAX, PWM_DIFF_BOUNDS);
    graph_set_param_val(graph, ps->clamp_d_pwmY, BOUNDS_MIN, -PWM_DIFF_BOUNDS);
    graph_set_param_val(graph, ps->clamp_d_pwmY, BOUNDS_MAX, PWM_DIFF_BOUNDS);

    // Set trims
    graph_set_param_val(graph, ps->pitch_trim, CONST_SET, 0.02);

    // Initial value for sampling periods
    graph_set_param_val(graph, ps->pos_time, CONST_SET, 0.04);
    graph_set_param_val(graph, ps->angle_time, CONST_SET, 0.005);
    // Set initial mode
    connect_manual(ps);

	return 0;
 }

 int control_algorithm(log_t* log_struct, user_input_t * user_input_struct, sensor_t* sensor_struct, setpoint_t* setpoint_struct, parameter_t* ps, user_defined_t* user_defined_struct, actuator_command_t* actuator_struct, modular_structs_t* structs)
 {
    struct computation_graph* graph = ps->graph;
	// use the 'flap' switch as the flight mode selector
	int cur_fm_switch = read_flap(user_input_struct->rc_commands[FLAP]);
	static int last_fm_switch = MANUAL_FLIGHT_MODE;

	// reset flight_mode to MANUAL right away if the flap switch is in manual position
	// to engage AUTO mode the code waits for a new packet after the flap is switched to auto
	// before actually engaging AUTO mode
	if(cur_fm_switch == MANUAL_FLIGHT_MODE) {
		if (last_fm_switch == AUTO_FLIGHT_MODE) {
			connect_manual(ps);
		}
		user_defined_struct->flight_mode = MANUAL_FLIGHT_MODE;
	}

	// flap switch was just toggled to auto flight mode
	if((last_fm_switch != cur_fm_switch) && (cur_fm_switch == AUTO_FLIGHT_MODE))
	{
		user_defined_struct->engaging_auto = 1;

		graph_set_param_val(graph, ps->throttle_trim, CONST_SET, user_input_struct->rc_commands[THROTTLE]);
	}

	if(user_input_struct->locationFresh && user_defined_struct->engaging_auto == 1)
			user_defined_struct->engaging_auto = 2;


	// The last packet ID from VRPN.
	static int last_vrpn_id = 0;
	// if the flap switch was toggled to AUTO_FLIGHT_MODE and we've received a new packet
	// then record the current position as the desired position
	// also reset the previous error and accumulated error from the position PIDs
	if((cur_fm_switch == AUTO_FLIGHT_MODE) && (user_defined_struct->engaging_auto == 2))
	{
		graph_set_param_val(graph, ps->x_set, CONST_SET, sensor_struct->currentQuadPosition.x_pos);
		graph_set_param_val(graph, ps->y_set, CONST_SET, sensor_struct->currentQuadPosition.y_pos);
		graph_set_param_val(graph, ps->alt_set, CONST_SET, sensor_struct->currentQuadPosition.alt_pos);
		graph_set_param_val(graph, ps->yaw_set, CONST_SET, sensor_struct->currentQuadPosition.yaw);

		// reset the flag that engages auto mode
		user_defined_struct->engaging_auto = 0;
		// finally engage the AUTO_FLIGHT_MODE
		// this ensures that we've gotten a new update packet right after the switch was set to auto mode
		user_defined_struct->flight_mode = AUTO_FLIGHT_MODE;
		connect_autonomous(ps);
		// Reset this when autonomous is engaged, so there is not a large difference at the start of autonomous
		last_vrpn_id = sensor_struct->currentQuadPosition.packetId - 1;
	}

	//PIDS///////////////////////////////////////////////////////////////////////

	/* 					Position loop
	 * Reads current position, and outputs
	 * a pitch or roll for the angle loop PIDs
	 */

	if(user_input_struct->locationFresh)
	{
	    // VRPN values
	    graph_set_param_val(graph, ps->vrpn_x, CONST_SET, sensor_struct->currentQuadPosition.x_pos);
	    graph_set_param_val(graph, ps->vrpn_y, CONST_SET, sensor_struct->currentQuadPosition.y_pos);
	    graph_set_param_val(graph, ps->vrpn_alt, CONST_SET, sensor_struct->currentQuadPosition.alt_pos);
	    graph_set_param_val(graph, ps->vrpn_pitch, CONST_SET, sensor_struct->currentQuadPosition.pitch);
	    graph_set_param_val(graph, ps->vrpn_roll, CONST_SET, sensor_struct->currentQuadPosition.roll);
	    graph_set_param_val(graph, ps->cur_yaw, CONST_SET, sensor_struct->currentQuadPosition.yaw);

	    // Use the difference in IDs to compute the sampling time for the position PIDs
	    int current_vrpn_id = sensor_struct->currentQuadPosition.packetId;
	    graph_set_param_val(graph, ps->pos_time, CONST_SET, VRPN_REFRESH_TIME * (current_vrpn_id - last_vrpn_id));
	    last_vrpn_id = current_vrpn_id;
	}

	// Loop time
    graph_set_param_val(graph, ps->angle_time, CONST_SET, get_last_loop_time());

	// Sensor values
    graph_set_param_val(graph, ps->cur_pitch, CONST_SET, sensor_struct->pitch_angle_filtered);
    graph_set_param_val(graph, ps->cur_roll, CONST_SET, sensor_struct->roll_angle_filtered);
    //graph_set_param_val(graph, ps->theta_dot, CONST_SET, sensor_struct->theta_dot);
    //graph_set_param_val(graph, ps->phi_dot, CONST_SET, sensor_struct->phi_dot);
    //graph_set_param_val(graph, ps->psi_dot, CONST_SET, sensor_struct->psi_dot);
    graph_set_param_val(graph, ps->theta_dot, CONST_SET, sensor_struct->gyr_y);
    graph_set_param_val(graph, ps->phi_dot, CONST_SET, sensor_struct->gyr_x);
    graph_set_param_val(graph, ps->psi_dot, CONST_SET, sensor_struct->gyr_z);

    // RC values
    graph_set_param_val(graph, ps->rc_pitch, CONST_SET, user_input_struct->pitch_angle_manual_setpoint);
    graph_set_param_val(graph, ps->rc_roll, CONST_SET, user_input_struct->roll_angle_manual_setpoint);
    graph_set_param_val(graph, ps->rc_yaw, CONST_SET, user_input_struct->yaw_manual_setpoint);
    graph_set_param_val(graph, ps->rc_throttle, CONST_SET, user_input_struct->rc_commands[THROTTLE]);

    // Compute VRPN blocks so they can be logged
    int outputs[4] = {ps->mixer, ps->vrpn_alt, ps->vrpn_pitch, ps->vrpn_roll};
    graph_compute_nodes(graph, outputs, 4);

	 // here for now so in case any flight command is not PID controlled, it will default to rc_command value:
	//memcpy(raw_actuator_struct->controller_corrected_motor_commands, user_input_struct->rc_commands, sizeof(int) * 6);

	// don't use the PID corrections if the throttle is less than about 10% of its range

    // re-enable the check for AUTO_FLIGHT_MODE when autonomous is fully enabled!!!
	if((user_input_struct->rc_commands[THROTTLE] > 118000))
	{
			//THROTTLE

			actuator_struct->pwms[0] = graph_get_output(graph, ps->mixer, MIXER_PWM0);
			actuator_struct->pwms[1] = graph_get_output(graph, ps->mixer, MIXER_PWM1);
			actuator_struct->pwms[2] = graph_get_output(graph, ps->mixer, MIXER_PWM2);
			actuator_struct->pwms[3] = graph_get_output(graph, ps->mixer, MIXER_PWM3);
	}
	else
	{
		actuator_struct->pwms[0] = user_input_struct->rc_commands[THROTTLE];
		actuator_struct->pwms[1] = user_input_struct->rc_commands[THROTTLE];
		actuator_struct->pwms[2] = user_input_struct->rc_commands[THROTTLE];
		actuator_struct->pwms[3] = user_input_struct->rc_commands[THROTTLE];
	}

	last_fm_switch = cur_fm_switch;

	// Make location stale now
	user_input_struct->locationFresh = 0;

    return 0;
 }

 void setPIDCoeff(PID_t* p, float pValue, float iValue, float dValue) {

 	p->Kp = pValue;
 	p->Ki = iValue;
 	p->Kd = dValue;

 }