/*
 * Luviner - Predictive Maintenance Demo
 * ======================================
 * Industrial motor vibration monitoring on ESP32.
 *
 * Detects 4 states:
 *   0 = Normal              -> GREEN LED
 *   1 = Rotor Imbalance     -> YELLOW LED
 *   2 = Bearing Wear        -> RED LED
 *   3 = Shaft Misalignment  -> RED LED (blinking)
 *
 * Hardware:
 *   Green LED  -> GPIO 13
 *   Yellow LED -> GPIO 12
 *   Red LED    -> GPIO 14
 */

#include <stdio.h>
#include <stdint.h>
#include <math.h>

extern "C" {
  #include "motor_model.h"
}

/* LED Pins */
#define LED_GREEN  13
#define LED_YELLOW 12
#define LED_RED    14

/* Class names */
const char* CLASS_NAMES[] = {
    "NORMAL",
    "IMBALANCE",
    "BEARING WEAR",
    "MISALIGNMENT"
};

/* Normalization parameters (from training set) */
const float norm_mean[32] = {
    0.104083, 1.143582, 3.978953, 1.402556, 0.109119, 1.419285, 5.222976, 1.820687,
    0.111078, 1.425596, 5.222961, 1.864608, 0.114237, 1.434222, 5.207515, 1.875538,
    0.108678, 1.276182, 4.788409, 1.664082, 0.250746, 1.279924, 4.748700, 1.666998,
    0.086222, 1.992232, 7.596256, 2.720997, 0.105403, 2.000372, 7.583913, 2.752669
};
const float norm_std[32] = {
    0.190031, 0.907843, 2.553461, 1.004364, 0.202944, 0.931424, 2.950773, 1.218039,
    0.193831, 0.936839, 2.964299, 1.289780, 0.195806, 0.942552, 2.946763, 1.351781,
    0.194676, 0.752271, 2.640957, 1.175996, 0.457545, 0.747722, 2.610656, 1.191325,
    0.200552, 1.791024, 7.071104, 2.654848, 0.197748, 1.797600, 7.058487, 2.570005
};

uint32_t sim_seed = 12345;
float sim_rand() {
    sim_seed = sim_seed * 1103515245 + 12345;
    return ((float)((sim_seed >> 16) & 0x7FFF)) / 32768.0f - 0.5f;
}

/* Simulate reading from 8 accelerometers */
void simulate_sensor_reading(float features[32], int scenario) {
    int i = 0;
    float signal_buf[200];

    for (int sensor = 0; sensor < 8; sensor++) {
        float base_freq = 25.0f + sensor * 2.0f;

        for (int k = 0; k < 200; k++) {
            float t = (float)k / 200.0f * 10.0f * 3.14159f * 2.0f;
            signal_buf[k] = 0.5f * sinf(t * base_freq / (2.0f * 3.14159f));

            if (scenario == 1) {
                /* Imbalance: strong 1x frequency */
                signal_buf[k] = 2.5f * sinf(t * base_freq / (2.0f * 3.14159f));
            } else if (scenario == 2) {
                /* Bearing wear: impulses + broadband noise */
                float bpfo = base_freq * 3.58f;
                float impulse = sinf(t * bpfo / (2.0f * 3.14159f));
                if (impulse > 0) impulse = impulse * impulse * impulse;
                else impulse = 0;
                signal_buf[k] += 2.0f * impulse;
                signal_buf[k] += sim_rand() * 0.4f;
            } else if (scenario == 3) {
                /* Misalignment: strong 2x and 3x harmonics */
                signal_buf[k] += 2.0f * sinf(t * base_freq * 2.0f / (2.0f * 3.14159f));
                signal_buf[k] += 1.5f * sinf(t * base_freq * 3.0f / (2.0f * 3.14159f));
                if (sensor >= 6) signal_buf[k] *= 2.5f;
            }
            signal_buf[k] += sim_rand() * 0.1f;
        }

        /* 4 statistical features per sensor */
        float s_mean = 0, s_min = signal_buf[0], s_max = signal_buf[0];
        float s_sum_sq = 0, s_rough = 0;
        for (int k = 0; k < 200; k++) {
            s_mean += signal_buf[k];
            if (signal_buf[k] > s_max) s_max = signal_buf[k];
            if (signal_buf[k] < s_min) s_min = signal_buf[k];
            s_sum_sq += signal_buf[k] * signal_buf[k];
            if (k > 0) s_rough += fabsf(signal_buf[k] - signal_buf[k-1]);
        }
        s_mean /= 200.0f;
        float s_std = sqrtf(s_sum_sq / 200.0f - s_mean * s_mean);
        float s_ptp = s_max - s_min;
        s_rough /= 199.0f;

        features[i++] = s_mean;
        features[i++] = s_std;
        features[i++] = s_ptp;
        features[i++] = s_rough;
    }
}

/* Set LED based on prediction */
void set_leds(int prediction) {
    digitalWrite(LED_GREEN, LOW);
    digitalWrite(LED_YELLOW, LOW);
    digitalWrite(LED_RED, LOW);

    switch (prediction) {
        case 0:  /* Normal */
            digitalWrite(LED_GREEN, HIGH);
            break;
        case 1:  /* Imbalance */
            digitalWrite(LED_YELLOW, HIGH);
            break;
        case 2:  /* Bearing wear */
            digitalWrite(LED_RED, HIGH);
            break;
        case 3:  /* Misalignment */
            /* Blinking red */
            for (int b = 0; b < 3; b++) {
                digitalWrite(LED_RED, HIGH);
                delay(150);
                digitalWrite(LED_RED, LOW);
                delay(150);
            }
            digitalWrite(LED_RED, HIGH);
            break;
    }
}

void setup() {
    Serial.begin(115200);

    /* Configure LED pins */
    pinMode(LED_GREEN, OUTPUT);
    pinMode(LED_YELLOW, OUTPUT);
    pinMode(LED_RED, OUTPUT);

    /* LED startup test */
    digitalWrite(LED_GREEN, HIGH); delay(300);
    digitalWrite(LED_GREEN, LOW);
    digitalWrite(LED_YELLOW, HIGH); delay(300);
    digitalWrite(LED_YELLOW, LOW);
    digitalWrite(LED_RED, HIGH); delay(300);
    digitalWrite(LED_RED, LOW);

    delay(500);

    Serial.println();
    Serial.println("==============================================");
    Serial.println("  LUVINER - Predictive Maintenance Demo");
    Serial.println("  Industrial motor vibration monitoring");
    Serial.println("==============================================");
    Serial.println();

    nc3_init();
    Serial.println("AI model loaded. Starting monitoring...");
    Serial.println();
}

int reading_num = 0;
int scenarios[] = {0, 0, 0, 1, 1, 2, 2, 3, 3, 0, 0, 0, 1, 2, 3, 0};
int n_scenarios = 16;

void loop() {
    int scenario = scenarios[reading_num % n_scenarios];

    float features[32];
    int16_t input_q8[32];

    /* 1. Read sensors */
    simulate_sensor_reading(features, scenario);

    /* 2. Normalize -> Q8 fixed-point */
    for (int j = 0; j < 32; j++) {
        float norm = (features[j] - norm_mean[j]) / (norm_std[j] + 1e-8f);
        if (norm > 4.0f) norm = 4.0f;
        if (norm < -4.0f) norm = -4.0f;
        input_q8[j] = (int16_t)(norm * 256.0f);
    }

    /* 3. AI inference */
    int pred = nc3_predict(input_q8);

    /* 4. Update LEDs */
    set_leds(pred);

    /* 5. Serial output */
    Serial.printf("[%3d] %-20s -> AI: %-20s %s\n",
                   reading_num + 1,
                   CLASS_NAMES[scenario],
                   CLASS_NAMES[pred],
                   (pred == scenario) ? "OK" : "!!");

    if (pred > 0) {
        Serial.printf("      ALERT: %s detected — schedule maintenance\n",
                      CLASS_NAMES[pred]);
    }

    reading_num++;

    /* Reading every 2 seconds (configurable in production) */
    delay(2000);
}