#include "rls.h"

#include <string.h>

void RLSInit(RLSInstance *rls, uint8_t order, float lambda, float init_covariance, const float *initial_theta)
{
    if (rls == 0) {
        return;
    }

    memset(rls, 0, sizeof(RLSInstance));

    if (order == 0u) {
        return;
    }
    if (order > RLS_MAX_ORDER) {
        order = RLS_MAX_ORDER;
    }
    if (lambda < 0.01f) {
        lambda = 0.01f;
    }
    if (lambda > 1.0f) {
        lambda = 1.0f;
    }
    if (init_covariance <= 0.0f) {
        init_covariance = 1000.0f;
    }

    rls->order = order;
    rls->lambda = lambda;
    rls->initialized = 1u;

    for (uint8_t i = 0u; i < order; i++) {
        rls->theta[i] = initial_theta != 0 ? initial_theta[i] : 0.0f;
        rls->covariance[i][i] = init_covariance;
    }
}

float RLSPredict(const RLSInstance *rls, const float *x)
{
    if (rls == 0 || x == 0 || rls->initialized == 0u) {
        return 0.0f;
    }

    float prediction = 0.0f;
    for (uint8_t i = 0u; i < rls->order; i++) {
        prediction += rls->theta[i] * x[i];
    }
    return prediction;
}

float RLSUpdate(RLSInstance *rls, const float *x, float y)
{
    float px[RLS_MAX_ORDER] = {0.0f};
    float xtp[RLS_MAX_ORDER] = {0.0f};

    if (rls == 0 || x == 0 || rls->initialized == 0u) {
        return 0.0f;
    }

    const uint8_t order = rls->order;
    const float prediction = RLSPredict(rls, x);
    const float error = y - prediction;

    for (uint8_t i = 0u; i < order; i++) {
        for (uint8_t j = 0u; j < order; j++) {
            px[i] += rls->covariance[i][j] * x[j];
            xtp[i] += x[j] * rls->covariance[j][i];
        }
    }

    float denominator = rls->lambda;
    for (uint8_t i = 0u; i < order; i++) {
        denominator += x[i] * px[i];
    }
    if (denominator <= 1.0e-6f) {
        denominator = 1.0e-6f;
    }

    for (uint8_t i = 0u; i < order; i++) {
        rls->gain[i] = px[i] / denominator;
        rls->theta[i] += rls->gain[i] * error;
    }

    for (uint8_t i = 0u; i < order; i++) {
        for (uint8_t j = 0u; j < order; j++) {
            rls->covariance[i][j] = (rls->covariance[i][j] - rls->gain[i] * xtp[j]) / rls->lambda;
        }
    }

    rls->last_prediction = prediction;
    rls->last_error = error;
    return prediction;
}

const float *RLSGetTheta(const RLSInstance *rls)
{
    return rls != 0 ? rls->theta : 0;
}