#include "drill_smo.h"

#define DS_PI 3.14159265358979323846f
#define DS_TWO_PI 6.28318530717958647692f
#define DS_HALF_PI 1.57079632679489661923f
#define DS_CORDIC_K 0.60725293500888125617f
#define DS_FLOAT_MAX 3.402823466e+38f

static const float cordic_angles[18] = {
    0.78539816339744830962f, 0.46364760900080611621f,
    0.24497866312686415417f, 0.12435499454676143503f,
    0.06241880999595734847f, 0.03123983343026827625f,
    0.01562372862047683129f, 0.00781234106010111130f,
    0.00390623013196697183f, 0.00195312251647881869f,
    0.00097656218955931943f, 0.00048828121119489829f,
    0.00024414062014936177f, 0.00012207031189367021f,
    0.00006103515617420877f, 0.00003051757811552610f,
    0.00001525878906131576f, 0.00000762939453110197f
};

static bool finite_value(float x) {
    return x == x && x <= DS_FLOAT_MAX && x >= -DS_FLOAT_MAX;
}

static float absf_local(float x) { return x < 0.0f ? -x : x; }

static float clamp(float x, float low, float high) {
    if (x < low) return low;
    if (x > high) return high;
    return x;
}

static float wrap_pi(float angle) {
    while (angle > DS_PI) angle -= DS_TWO_PI;
    while (angle < -DS_PI) angle += DS_TWO_PI;
    return angle;
}

static float wrap_two_pi(float angle) {
    while (angle >= DS_TWO_PI) angle -= DS_TWO_PI;
    while (angle < 0.0f) angle += DS_TWO_PI;
    return angle;
}

static float atan2_cordic(float y, float x) {
    if (x == 0.0f && y == 0.0f) return 0.0f;
    float angle = 0.0f;
    if (x < 0.0f) {
        if (y >= 0.0f) angle = DS_PI;
        else angle = -DS_PI;
        x = -x;
        y = -y;
    }
    for (unsigned i = 0; i < 18; ++i) {
        const float scale = 1.0f / (float)(1u << i);
        const float old_x = x;
        if (y > 0.0f) {
            x += y * scale;
            y -= old_x * scale;
            angle += cordic_angles[i];
        } else {
            x -= y * scale;
            y += old_x * scale;
            angle -= cordic_angles[i];
        }
    }
    return wrap_pi(angle);
}

void drill_smo_sincos(float angle_rad, float *sine, float *cosine) {
    if (!sine || !cosine) return;
    if (!finite_value(angle_rad)) {
        *sine = 0.0f;
        *cosine = 0.0f;
        return;
    }
    if (absf_local(angle_rad) > 1000000.0f) {
        *sine = 0.0f;
        *cosine = 0.0f;
        return;
    }
    float z = angle_rad;
    if (z > DS_PI || z < -DS_PI) {
        const int32_t turns = (int32_t)(z / DS_TWO_PI);
        z -= (float)turns * DS_TWO_PI;
        if (z > DS_PI) z -= DS_TWO_PI;
        if (z < -DS_PI) z += DS_TWO_PI;
    }
    float sign = 1.0f;
    if (z > DS_HALF_PI) {
        z -= DS_PI;
        sign = -1.0f;
    } else if (z < -DS_HALF_PI) {
        z += DS_PI;
        sign = -1.0f;
    }
    float x = DS_CORDIC_K;
    float y = 0.0f;
    for (unsigned i = 0; i < 18; ++i) {
        const float scale = 1.0f / (float)(1u << i);
        const float old_x = x;
        if (z >= 0.0f) {
            x -= y * scale;
            y += old_x * scale;
            z -= cordic_angles[i];
        } else {
            x += y * scale;
            y -= old_x * scale;
            z += cordic_angles[i];
        }
    }
    *sine = y * sign;
    *cosine = x * sign;
}

static bool config_valid(const drill_smo_config *c) {
    if (!c || !finite_value(c->phase_resistance_ohm) ||
        c->phase_resistance_ohm <= 0.0f || c->phase_resistance_ohm > 100.0f ||
        !finite_value(c->phase_inductance_h) || c->phase_inductance_h <= 0.0f ||
        c->phase_inductance_h > 1.0f ||
        !finite_value(c->sliding_gain_v) || c->sliding_gain_v <= 0.0f ||
        c->sliding_gain_v > 1000.0f ||
        !finite_value(c->current_error_sat_a) || c->current_error_sat_a <= 0.0f ||
        !finite_value(c->bemf_filter_tau_s) || c->bemf_filter_tau_s <= 0.0f ||
        c->bemf_filter_tau_s > 0.01f ||
        !finite_value(c->pll_kp_per_s) || c->pll_kp_per_s <= 0.0f ||
        c->pll_kp_per_s > 50000.0f ||
        !finite_value(c->pll_ki_per_s2) || c->pll_ki_per_s2 <= 0.0f ||
        c->pll_ki_per_s2 > 1000000000.0f ||
        !finite_value(c->min_bemf_v) || c->min_bemf_v <= 0.0f ||
        !finite_value(c->max_phase_voltage_v) || c->max_phase_voltage_v <= 0.0f ||
        c->max_phase_voltage_v > 100.0f ||
        c->min_bemf_v >= c->max_phase_voltage_v ||
        !finite_value(c->max_phase_current_a) || c->max_phase_current_a <= 0.0f ||
        c->max_phase_current_a > 8.0f ||
        c->current_error_sat_a > c->max_phase_current_a ||
        !finite_value(c->max_mechanical_rpm) || c->max_mechanical_rpm <= 0.0f ||
        c->max_mechanical_rpm > 200000.0f ||
        !finite_value(c->max_dt_s) || c->max_dt_s <= 0.0f || c->max_dt_s > 0.0001f ||
        !finite_value(c->lock_phase_error_rad) || c->lock_phase_error_rad <= 0.0f ||
        c->lock_phase_error_rad >= DS_HALF_PI || c->pole_pairs == 0u ||
        c->pole_pairs > 32u || c->lock_samples < 4u) return false;
    return true;
}

bool drill_smo_init(drill_smo_state *s, const drill_smo_config *c) {
    if (!s) return false;
    drill_smo_reset(s);
    if (!config_valid(c)) return false;
    s->config = *c;
    s->initialized = true;
    return true;
}

void drill_smo_reset(drill_smo_state *s) {
    if (!s) return;
    s->config = (drill_smo_config){0};
    s->estimated_current_alpha_a = 0.0f;
    s->estimated_current_beta_a = 0.0f;
    s->estimated_bemf_alpha_v = 0.0f;
    s->estimated_bemf_beta_v = 0.0f;
    s->electrical_angle_rad = 0.0f;
    s->electrical_speed_rad_s = 0.0f;
    s->lock_count = 0u;
    s->unlock_count = 0u;
    s->initialized = false;
    s->angle_valid = false;
    s->observer_locked = false;
}

bool drill_smo_svpwm_to_ab(float duty_a, float duty_b, float duty_c,
                           float bus_voltage_v, float *voltage_alpha_v,
                           float *voltage_beta_v) {
    if (!voltage_alpha_v || !voltage_beta_v) return false;
    *voltage_alpha_v = 0.0f;
    *voltage_beta_v = 0.0f;
    if (!finite_value(duty_a) || !finite_value(duty_b) ||
        !finite_value(duty_c) || !finite_value(bus_voltage_v) ||
        duty_a < 0.0f || duty_a > 1.0f || duty_b < 0.0f || duty_b > 1.0f ||
        duty_c < 0.0f || duty_c > 1.0f || bus_voltage_v <= 0.0f ||
        bus_voltage_v > 100.0f) return false;
    const float common_duty = (duty_a + duty_b + duty_c) / 3.0f;
    const float phase_a_v = bus_voltage_v * (duty_a - common_duty);
    const float phase_b_v = bus_voltage_v * (duty_b - common_duty);
    const float phase_c_v = bus_voltage_v * (duty_c - common_duty);
    *voltage_alpha_v = (2.0f/3.0f) *
                       (phase_a_v - 0.5f*phase_b_v - 0.5f*phase_c_v);
    *voltage_beta_v = (phase_b_v - phase_c_v) * 0.5773502691896258f;
    return true;
}

static bool input_valid(const drill_smo_state *s, const drill_smo_input *i) {
    const drill_smo_config *c = &s->config;
    return i && finite_value(i->voltage_alpha_v) && finite_value(i->voltage_beta_v) &&
           finite_value(i->current_alpha_a) && finite_value(i->current_beta_a) &&
           finite_value(i->dt_s) && i->dt_s > 0.0f && i->dt_s <= c->max_dt_s &&
           (i->commanded_direction == 1 || i->commanded_direction == -1) &&
           absf_local(i->current_alpha_a) <= c->max_phase_current_a &&
           absf_local(i->current_beta_a) <= c->max_phase_current_a &&
           i->current_alpha_a*i->current_alpha_a +
               i->current_beta_a*i->current_beta_a <=
               c->max_phase_current_a*c->max_phase_current_a &&
           absf_local(i->voltage_alpha_v) <= c->max_phase_voltage_v &&
           absf_local(i->voltage_beta_v) <= c->max_phase_voltage_v &&
           i->voltage_alpha_v*i->voltage_alpha_v +
               i->voltage_beta_v*i->voltage_beta_v <=
               c->max_phase_voltage_v*c->max_phase_voltage_v;
}

static float sliding_injection(float error_a, float limit_a, float gain_v) {
    return gain_v * clamp(error_a / limit_a, -1.0f, 1.0f);
}

bool drill_smo_step(drill_smo_state *s, const drill_smo_input *i,
                    drill_smo_output *o) {
    if (o) *o = (drill_smo_output){0};
    if (!s || !s->initialized || !o || !input_valid(s, i)) {
        if (s && s->initialized) drill_smo_reset(s);
        return false;
    }
    const drill_smo_config *c = &s->config;
    const float inj_a = sliding_injection(s->estimated_current_alpha_a -
                                         i->current_alpha_a,
                                         c->current_error_sat_a,
                                         c->sliding_gain_v);
    const float inj_b = sliding_injection(s->estimated_current_beta_a -
                                         i->current_beta_a,
                                         c->current_error_sat_a,
                                         c->sliding_gain_v);
    const float inv_l = 1.0f / c->phase_inductance_h;
    s->estimated_current_alpha_a += i->dt_s * inv_l *
        (i->voltage_alpha_v - c->phase_resistance_ohm *
         s->estimated_current_alpha_a - inj_a);
    s->estimated_current_beta_a += i->dt_s * inv_l *
        (i->voltage_beta_v - c->phase_resistance_ohm *
         s->estimated_current_beta_a - inj_b);
    if (!finite_value(s->estimated_current_alpha_a) ||
        !finite_value(s->estimated_current_beta_a) ||
        absf_local(s->estimated_current_alpha_a) > 2.0f*c->max_phase_current_a ||
        absf_local(s->estimated_current_beta_a) > 2.0f*c->max_phase_current_a) {
        drill_smo_reset(s);
        return false;
    }
    const float filter_alpha = i->dt_s / (c->bemf_filter_tau_s + i->dt_s);
    s->estimated_bemf_alpha_v += filter_alpha *
                                (inj_a - s->estimated_bemf_alpha_v);
    s->estimated_bemf_beta_v += filter_alpha *
                               (inj_b - s->estimated_bemf_beta_v);
    const float bemf_sq = s->estimated_bemf_alpha_v*s->estimated_bemf_alpha_v +
                          s->estimated_bemf_beta_v*s->estimated_bemf_beta_v;
    if (!finite_value(bemf_sq)) {
        drill_smo_reset(s);
        return false;
    }
    if (bemf_sq < c->min_bemf_v*c->min_bemf_v) {
        s->lock_count = 0u;
        s->unlock_count = 0u;
        s->observer_locked = false;
        s->angle_valid = false;
        o->bemf_alpha_v = s->estimated_bemf_alpha_v;
        o->bemf_beta_v = s->estimated_bemf_beta_v;
        o->valid = true;
        return true;
    }

    const float direction = (float)i->commanded_direction;
    const float raw_angle = atan2_cordic(-direction*s->estimated_bemf_alpha_v,
                                         direction*s->estimated_bemf_beta_v);
    /* The current observer and BEMF smoother each contribute phase delay. */
    const float phase_lead = direction *
        clamp(2.0f*absf_local(s->electrical_speed_rad_s)*
              c->bemf_filter_tau_s, 0.0f, 0.65f);
    const float measured_angle = wrap_pi(raw_angle + phase_lead);
    float phase_error = 0.0f;
    if (!s->angle_valid) {
        s->electrical_angle_rad = measured_angle;
        s->electrical_speed_rad_s = 0.0f;
        s->angle_valid = true;
    } else {
        phase_error = wrap_pi(measured_angle - s->electrical_angle_rad);
        const float max_speed = c->max_mechanical_rpm *
                                ((DS_TWO_PI/60.0f)*(float)c->pole_pairs);
        s->electrical_speed_rad_s = clamp(
            s->electrical_speed_rad_s + c->pll_ki_per_s2*phase_error*i->dt_s,
            -max_speed, max_speed);
        s->electrical_angle_rad = wrap_two_pi(s->electrical_angle_rad +
            (s->electrical_speed_rad_s + c->pll_kp_per_s*phase_error)*i->dt_s);
    }
    if (!s->observer_locked) {
        if (absf_local(phase_error) <= c->lock_phase_error_rad &&
            absf_local(s->electrical_speed_rad_s) > 0.0f) {
            if (s->lock_count < c->lock_samples) ++s->lock_count;
            if (s->lock_count >= c->lock_samples) s->observer_locked = true;
        } else {
            s->lock_count = 0u;
        }
    } else {
        if (absf_local(phase_error) > 2.0f*c->lock_phase_error_rad) {
            if (s->unlock_count < c->lock_samples) ++s->unlock_count;
            if (s->unlock_count >= c->lock_samples) {
                s->observer_locked = false;
                s->lock_count = 0u;
                s->unlock_count = 0u;
            }
        } else {
            s->unlock_count = 0u;
        }
    }
    drill_smo_sincos(s->electrical_angle_rad, &o->sine, &o->cosine);
    o->electrical_angle_rad = s->electrical_angle_rad;
    o->electrical_speed_rad_s = s->electrical_speed_rad_s;
    o->mechanical_speed_rpm = s->electrical_speed_rad_s *
        (60.0f/DS_TWO_PI)/(float)c->pole_pairs;
    o->bemf_alpha_v = s->estimated_bemf_alpha_v;
    o->bemf_beta_v = s->estimated_bemf_beta_v;
    o->phase_error_rad = phase_error;
    o->angle_valid = s->angle_valid;
    o->observer_locked = s->observer_locked;
    o->valid = true;
    return true;
}
