#include "chassis_power.h"

#include <math.h>
#include <string.h>

ChassisPower_State_s chassis_power_debug_state;

static float ChassisPowerLimitValue(float value, float min, float max)
{
    if (value < min) {
        return min;
    }
    if (value > max) {
        return max;
    }
    return value;
}

ChassisPowerMotorModel_s ChassisPowerGetDefault3508Model(void)
{
    const float it_to_ireal_m3508 = 20.0f / 16384.0f;
    const float i_to_t_m3508 = 0.3f / (3591.0f / 187.0f);
    const float it_to_t_m3508 = it_to_ireal_m3508 * i_to_t_m3508;

    ChassisPowerMotorModel_s model = {
        .k_current_speed = it_to_t_m3508 / 9.55f,
        .k_speed2 = 1.453e-7f,
        .k_current2 = 2.53e-7f,
        .k0 = 4.081f / 4.0f,
    };
    return model;
}

void ChassisPowerInit(ChassisPowerInstance *instance, const ChassisPower_Init_Config_s *config)
{
    if (instance == 0) {
        return;
    }

    memset(instance, 0, sizeof(ChassisPowerInstance));

    instance->wheel_model = ChassisPowerGetDefault3508Model();
    instance->default_power_limit = 60.0f;
    instance->min_power_limit = 20.0f;
    instance->buffer_low = 30.0f;
    instance->buffer_high = 50.0f;
    instance->buffer_gain = 0.25f;
    instance->steering_power_reserve = 6.0f;
    instance->min_scale = 0.20f;
    instance->rls_enable = 0u;

    float lambda = 0.995f;
    float init_covariance = 1000.0f;

    if (config != 0) {
        instance->wheel_model = config->wheel_model;
        if (config->default_power_limit > 0.0f) {
            instance->default_power_limit = config->default_power_limit;
        }
        if (config->min_power_limit > 0.0f) {
            instance->min_power_limit = config->min_power_limit;
        }
        if (config->buffer_low > 0.0f) {
            instance->buffer_low = config->buffer_low;
        }
        if (config->buffer_high > instance->buffer_low) {
            instance->buffer_high = config->buffer_high;
        }
        if (config->buffer_gain > 0.0f) {
            instance->buffer_gain = config->buffer_gain;
        }
        if (config->steering_power_reserve >= 0.0f) {
            instance->steering_power_reserve = config->steering_power_reserve;
        }
        if (config->min_scale > 0.0f && config->min_scale <= 1.0f) {
            instance->min_scale = config->min_scale;
        }
        if (config->rls_lambda > 0.0f) {
            lambda = config->rls_lambda;
        }
        if (config->rls_init_covariance > 0.0f) {
            init_covariance = config->rls_init_covariance;
        }
        instance->rls_enable = config->rls_enable;
    }

    float theta[CHASSIS_POWER_RLS_ORDER] = {
        instance->wheel_model.k_current_speed,
        instance->wheel_model.k_speed2,
        instance->wheel_model.k_current2,
        instance->wheel_model.k0 * 2.0f,
    };
    RLSInit(&instance->rls, CHASSIS_POWER_RLS_ORDER, lambda, init_covariance, theta);
    memcpy(instance->state.rls_theta, theta, sizeof(theta));
}

float ChassisPowerEstimateMotor(const ChassisPowerMotorModel_s *model, float speed, float current_cmd)
{
    if (model == 0) {
        return 0.0f;
    }

    float power = model->k_current_speed * speed * current_cmd +
                  model->k_speed2 * speed * speed +
                  model->k_current2 * current_cmd * current_cmd +
                  model->k0;

    if (power < 0.0f) {
        power = 0.0f;
    }
    return power;
}

static float ChassisPowerGetAllowedPower(ChassisPowerInstance *instance, const ChassisPower_Limit_Input_s *input)
{
    float allowed = input->referee_limit > 1.0f ? input->referee_limit : instance->default_power_limit;

    if (input->referee_limit > 1.0f) {
        if (input->power_buffer < instance->buffer_low) {
            allowed -= (instance->buffer_low - input->power_buffer) * instance->buffer_gain;
        } else if (input->power_buffer > instance->buffer_high) {
            allowed += (input->power_buffer - instance->buffer_high) * instance->buffer_gain;
        }
    }

    return ChassisPowerLimitValue(allowed, instance->min_power_limit, instance->default_power_limit + 20.0f);
}

void ChassisPowerUpdateMeasureState(ChassisPowerInstance *instance, const ChassisPower_Limit_Input_s *input)
{
    if (instance == 0 || input == 0) {
        return;
    }

    ChassisPower_State_s *state = &instance->state;
    const float rf_power = ChassisPowerEstimateMotor(&instance->wheel_model, input->rf_speed, input->rf_current_cmd);
    const float lb_power = ChassisPowerEstimateMotor(&instance->wheel_model, input->lb_speed, input->lb_current_cmd);

    state->rf_ref_raw = input->rf_speed_ref;
    state->lb_ref_raw = input->lb_speed_ref;
    state->rf_ref_limited = input->rf_speed_ref;
    state->lb_ref_limited = input->lb_speed_ref;
    state->rf_power = rf_power;
    state->lb_power = lb_power;
    state->estimated_wheel_power = rf_power + lb_power;
    state->estimated_total_power = state->estimated_wheel_power;
    state->allowed_power = ChassisPowerGetAllowedPower(instance, input);
    state->wheel_budget = state->allowed_power;
    state->scale = 1.0f;
    state->measured_power = input->measured_power;
    state->referee_power = input->referee_power;
    state->power_buffer = input->power_buffer;
    state->rf_speed = input->rf_speed;
    state->lb_speed = input->lb_speed;
    state->rf_current_cmd = input->rf_current_cmd;
    state->lb_current_cmd = input->lb_current_cmd;
    state->rf_output_cmd = input->rf_output_cmd;
    state->lb_output_cmd = input->lb_output_cmd;
    state->rf_feedback_current = input->rf_feedback_current;
    state->lb_feedback_current = input->lb_feedback_current;
    state->power_meter_online = input->power_meter_online;
    state->limited = 0u;

    state->features[0] = input->rf_speed * input->rf_current_cmd + input->lb_speed * input->lb_current_cmd;
    state->features[1] = input->rf_speed * input->rf_speed + input->lb_speed * input->lb_speed;
    state->features[2] = input->rf_current_cmd * input->rf_current_cmd + input->lb_current_cmd * input->lb_current_cmd;
    state->features[3] = 1.0f;
    chassis_power_debug_state = *state;
}

float ChassisPowerLimitWheelRefs(ChassisPowerInstance *instance, const ChassisPower_Limit_Input_s *input, float *rf_ref, float *lb_ref)
{
    if (instance == 0 || input == 0 || rf_ref == 0 || lb_ref == 0) {
        return 1.0f;
    }

    ChassisPower_State_s *state = &instance->state;
    const float rf_raw = *rf_ref;
    const float lb_raw = *lb_ref;
    const float allowed = ChassisPowerGetAllowedPower(instance, input);
    float steering_power = input->steering_power_estimate;

    if (steering_power <= 0.0f) {
        steering_power = instance->steering_power_reserve;
    }

    const float rf_power = ChassisPowerEstimateMotor(&instance->wheel_model, input->rf_speed, input->rf_current_cmd);
    const float lb_power = ChassisPowerEstimateMotor(&instance->wheel_model, input->lb_speed, input->lb_current_cmd);
    const float wheel_power = rf_power + lb_power;
    const float total_power = wheel_power + steering_power;
    float wheel_budget = allowed - steering_power;
    float scale = 1.0f;

    if (wheel_budget < 1.0f) {
        wheel_budget = 1.0f;
    }

    if (wheel_power > wheel_budget && wheel_power > 1.0e-3f) {
        scale = sqrtf(wheel_budget / wheel_power);
        scale = ChassisPowerLimitValue(scale, instance->min_scale, 1.0f);
        *rf_ref = rf_raw * scale;
        *lb_ref = lb_raw * scale;
        state->limited = 1u;
    } else {
        state->limited = 0u;
    }

    state->rf_ref_raw = rf_raw;
    state->lb_ref_raw = lb_raw;
    state->rf_ref_limited = *rf_ref;
    state->lb_ref_limited = *lb_ref;
    state->rf_power = rf_power;
    state->lb_power = lb_power;
    state->estimated_wheel_power = wheel_power;
    state->estimated_total_power = total_power;
    state->allowed_power = allowed;
    state->wheel_budget = wheel_budget;
    state->scale = scale;
    state->measured_power = input->measured_power;
    state->referee_power = input->referee_power;
    state->power_buffer = input->power_buffer;
    state->rf_speed = input->rf_speed;
    state->lb_speed = input->lb_speed;
    state->rf_current_cmd = input->rf_current_cmd;
    state->lb_current_cmd = input->lb_current_cmd;
    state->rf_output_cmd = input->rf_output_cmd;
    state->lb_output_cmd = input->lb_output_cmd;
    state->rf_feedback_current = input->rf_feedback_current;
    state->lb_feedback_current = input->lb_feedback_current;
    state->power_meter_online = input->power_meter_online;

    state->features[0] = input->rf_speed * input->rf_current_cmd + input->lb_speed * input->lb_current_cmd;
    state->features[1] = input->rf_speed * input->rf_speed + input->lb_speed * input->lb_speed;
    state->features[2] = input->rf_current_cmd * input->rf_current_cmd + input->lb_current_cmd * input->lb_current_cmd;
    state->features[3] = 1.0f;
    chassis_power_debug_state = *state;

    return scale;
}

void ChassisPowerUpdateRLS(ChassisPowerInstance *instance, const ChassisPower_Limit_Input_s *input, float measured_wheel_power)
{
    if (instance == 0 || input == 0 || instance->rls_enable == 0u || measured_wheel_power <= 0.0f) {
        return;
    }

    float features[CHASSIS_POWER_RLS_ORDER] = {
        input->rf_speed * input->rf_current_cmd + input->lb_speed * input->lb_current_cmd,
        input->rf_speed * input->rf_speed + input->lb_speed * input->lb_speed,
        input->rf_current_cmd * input->rf_current_cmd + input->lb_current_cmd * input->lb_current_cmd,
        1.0f,
    };

    RLSUpdate(&instance->rls, features, measured_wheel_power);
    const float *theta = RLSGetTheta(&instance->rls);
    if (theta == 0) {
        return;
    }

    instance->wheel_model.k_current_speed = theta[0];
    instance->wheel_model.k_speed2 = theta[1] > 0.0f ? theta[1] : 0.0f;
    instance->wheel_model.k_current2 = theta[2] > 0.0f ? theta[2] : 0.0f;
    instance->wheel_model.k0 = theta[3] > 0.0f ? theta[3] * 0.5f : 0.0f;
    memcpy(instance->state.rls_theta, theta, sizeof(instance->state.rls_theta));
    chassis_power_debug_state = instance->state;
}

void ChassisPowerSetWheelModel(ChassisPowerInstance *instance, const ChassisPowerMotorModel_s *model)
{
    if (instance == 0 || model == 0) {
        return;
    }
    instance->wheel_model = *model;
}

const ChassisPower_State_s *ChassisPowerGetState(const ChassisPowerInstance *instance)
{
    return instance != 0 ? &instance->state : 0;
}