/**
 * MCU Mattress Controller - C Implementation
 * 
 * This program implements sleep posture detection and mattress control for MCUs.
 * It uses TensorFlow Lite for Microcontrollers for inference.
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <pthread.h>
#include "tensorflow/lite/micro/all_ops_resolver.h"
#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/schema/schema_generated.h"

// If model data is included as a C array
#include "model_tflite.h" // Replace with your model header

// Constants for mode definitions
#define MODE_MANUAL     0
#define MODE_ADAPTIVE   1
#define MODE_PREVENTION 2
#define MODE_ABAB       3
#define MODE_TURNING    4
#define MODE_FLOATING   5

// Command codes from remote controller
#define CMD_STOP           0x00
#define CMD_ADAPTIVE_MODE  0x01
#define CMD_PREVENTION_MODE 0x02
#define CMD_ABAB_MODE      0x03
#define CMD_TURNING_MODE   0x04
#define CMD_FLOATING_MODE  0x05
#define CMD_DETECT_POSTURE 0x06

// Posture definitions
#define POSTURE_SUPINE 0  // Supine (back)
#define POSTURE_LEFT   1  // Left side
#define POSTURE_RIGHT  2  // Right side

// Serial port defines - platform specific
#define MATTRESS_PORT "/dev/ttyS1"
#define SENSOR_PORT   "/dev/ttyS2"
#define REMOTE_PORT   "/dev/ttyS0"

// Data structures
typedef struct {
    int posture_index;
    float confidence;
    bool success;
} PostureResult;

typedef struct {
    // Serial handles
    void* mattress_ser;
    void* sensor_ser;
    void* remote_ser;
    
    // Thread management
    pthread_t running_threads[5]; // Support up to 5 concurrent threads
    int thread_count;
    pthread_mutex_t thread_mutex;
    bool stop_event;
    
    // TFLite model
    tflite::MicroErrorReporter error_reporter;
    tflite::AllOpsResolver resolver;
    const tflite::Model* model;
    tflite::MicroInterpreter* interpreter;
    TfLiteTensor* input_tensor;
    TfLiteTensor* output_tensor;
    bool model_loaded;
    bool is_quantized;
    
    // Device ID bytes for mattress commands
    uint8_t id_bytes[4];
    
    // Current posture tracking
    int current_posture;
    float confidence;
    
    // Memory buffer for TensorFlow Lite arena
    uint8_t tensor_arena[128 * 1024]; // 128KB tensor arena
} MCUController;

// Forward declarations for threading
void* adaptive_thread_func(void* arg);
void* prevention_thread_func(void* arg);
void* abab_thread_func(void* arg);
void* turning_thread_func(void* arg);
void* floating_thread_func(void* arg);

// Serial communication stubs (platform-specific implementations needed)
bool serial_open(void** handle, const char* port, int baudrate);
void serial_close(void* handle);
bool serial_write(void* handle, const uint8_t* data, size_t length);
int serial_read(void* handle, uint8_t* buffer, size_t max_length);
int serial_available(void* handle);

/**
 * Initialize the controller
 */
bool controller_init(MCUController* controller) {
    memset(controller, 0, sizeof(MCUController));
    
    // Initialize mutex
    pthread_mutex_init(&controller->thread_mutex, NULL);
    
    // Initialize device ID bytes
    controller->id_bytes[0] = 0x01;
    controller->id_bytes[1] = 0x02;
    controller->id_bytes[2] = 0x03;
    controller->id_bytes[3] = 0x04;
    
    // Load TFLite model
    controller->model = tflite::GetModel(g_model_tflite); // Use the model data from header
    if (controller->model->version() != TFLITE_SCHEMA_VERSION) {
        printf("Model schema version mismatch!\n");
        return false;
    }
    
    // Create interpreter
    static tflite::MicroInterpreter static_interpreter(
        controller->model, 
        controller->resolver,
        controller->tensor_arena, 
        sizeof(controller->tensor_arena),
        &controller->error_reporter);
    
    controller->interpreter = &static_interpreter;
    
    // Allocate tensors
    TfLiteStatus allocate_status = controller->interpreter->AllocateTensors();
    if (allocate_status != kTfLiteOk) {
        printf("Failed to allocate tensors!\n");
        return false;
    }
    
    // Get input and output tensors
    controller->input_tensor = controller->interpreter->input(0);
    controller->output_tensor = controller->interpreter->output(0);
    
    // Check if model is quantized
    controller->is_quantized = controller->input_tensor->type == kTfLiteInt8;
    
    printf("Model loaded successfully\n");
    printf("Model is %s\n", controller->is_quantized ? "quantized" : "floating-point");
    
    controller->model_loaded = true;
    return true;
}

/**
 * Connect to mattress
 */
bool connect_mattress(MCUController* controller, const char* port, int baudrate) {
    // Close existing connection if any
    if (controller->mattress_ser) {
        serial_close(controller->mattress_ser);
        controller->mattress_ser = NULL;
    }
    
    // Open new connection
    if (serial_open(&controller->mattress_ser, port, baudrate)) {
        printf("Connected to mattress on %s\n", port);
        return true;
    } else {
        printf("Failed to connect to mattress on %s\n", port);
        return false;
    }
}

/**
 * Connect to sensor
 */
bool connect_sensor(MCUController* controller, const char* port, int baudrate) {
    // Close existing connection if any
    if (controller->sensor_ser) {
        serial_close(controller->sensor_ser);
        controller->sensor_ser = NULL;
    }
    
    // Open new connection
    if (serial_open(&controller->sensor_ser, port, baudrate)) {
        printf("Connected to sensor on %s\n", port);
        return true;
    } else {
        printf("Failed to connect to sensor on %s\n", port);
        return false;
    }
}

/**
 * Disconnect from all devices
 */
void disconnect_all(MCUController* controller) {
    // Stop all running threads
    stop_current_mode(controller);
    
    // Close serial connections
    if (controller->mattress_ser) {
        serial_close(controller->mattress_ser);
        controller->mattress_ser = NULL;
    }
    
    if (controller->sensor_ser) {
        serial_close(controller->sensor_ser);
        controller->sensor_ser = NULL;
    }
    
    if (controller->remote_ser) {
        serial_close(controller->remote_ser);
        controller->remote_ser = NULL;
    }
}

/**
 * Read pressure data from sensor
 */
bool read_pressure_data(MCUController* controller, float* data, size_t data_size) {
    if (!controller->sensor_ser) {
        printf("Sensor not connected\n");
        return false;
    }
    
    // Send command to request data
    uint8_t cmd[] = {0xAA, 0x01, 0xBB};
    serial_write(controller->sensor_ser, cmd, sizeof(cmd));
    
    // Wait for response
    struct timespec ts;
    ts.tv_sec = 0;
    ts.tv_nsec = 100000000; // 100ms
    nanosleep(&ts, NULL);
    
    // Calculate expected data size (1024 float values = 4096 bytes)
    size_t expected_bytes = data_size * sizeof(float);
    
    // Check if enough data is available
    if (serial_available(controller->sensor_ser) >= expected_bytes) {
        // Read raw bytes
        uint8_t raw_data[expected_bytes];
        size_t bytes_read = serial_read(controller->sensor_ser, raw_data, expected_bytes);
        
        if (bytes_read == expected_bytes) {
            // Convert raw bytes to float array
            memcpy(data, raw_data, expected_bytes);
            return true;
        } else {
            printf("Incomplete data: got %zu bytes, expected %zu\n", bytes_read, expected_bytes);
        }
    } else {
        printf("Insufficient data: %d bytes available\n", serial_available(controller->sensor_ser));
    }
    
    return false;
}

/**
 * Preprocess the raw sensor data
 */
void preprocess_data(const float* raw_data, float* processed_data, size_t data_size) {
    const float clip_min = 0.0f;
    const float clip_max = 300.0f;
    const float epsilon = 1e-6f;
    
    // Calculate mean and std_dev
    float sum = 0.0f;
    float sum_sq = 0.0f;
    
    // First pass: calculate mean
    for (size_t i = 0; i < data_size; i++) {
        // Apply clipping
        float value = raw_data[i];
        if (value < clip_min) value = clip_min;
        if (value > clip_max) value = clip_max;
        
        // Sqrt(x+1)
        value = sqrtf(value + 1.0f);
        
        // Accumulate for mean calculation
        sum += value;
        sum_sq += value * value;
    }
    
    float mean = sum / data_size;
    float variance = (sum_sq / data_size) - (mean * mean);
    float std_dev = sqrtf(variance) + epsilon;
    
    // Second pass: apply normalization, tanh, and scaling
    for (size_t i = 0; i < data_size; i++) {
        // Apply clipping
        float value = raw_data[i];
        if (value < clip_min) value = clip_min;
        if (value > clip_max) value = clip_max;
        
        // Sqrt(x+1)
        value = sqrtf(value + 1.0f);
        
        // Normalize
        value = (value - mean) / std_dev;
        
        // Apply tanh
        value = tanhf(value);
        
        // Scale to [0,1]
        value = (value + 1.0f) / 2.0f;
        
        processed_data[i] = value;
    }
}

/**
 * Detect sleep posture from sensor data
 */
PostureResult detect_posture(MCUController* controller) {
    PostureResult result = {0};
    result.success = false;
    
    if (!controller->model_loaded) {
        printf("Model not loaded\n");
        return result;
    }
    
    // Read pressure data
    float raw_data[1024];
    if (!read_pressure_data(controller, raw_data, 1024)) {
        printf("Failed to read pressure data\n");
        return result;
    }
    
    // Preprocess data
    float processed_data[1024];
    preprocess_data(raw_data, processed_data, 1024);
    
    // Copy data to input tensor
    if (controller->is_quantized) {
        // For quantized model, quantize the input
        int8_t* input_data = controller->input_tensor->data.int8;
        float input_scale = controller->input_tensor->params.scale;
        int input_zero_point = controller->input_tensor->params.zero_point;
        
        for (int i = 0; i < 1024; i++) {
            input_data[i] = (int8_t)(processed_data[i] / input_scale + input_zero_point);
        }
    } else {
        // For float model, copy directly
        float* input_data = controller->input_tensor->data.f;
        memcpy(input_data, processed_data, 1024 * sizeof(float));
    }
    
    // Run inference
    TfLiteStatus invoke_status = controller->interpreter->Invoke();
    if (invoke_status != kTfLiteOk) {
        printf("Invoke failed\n");
        return result;
    }
    
    // Get output
    float output_values[3]; // Assuming 3 classes
    
    if (controller->is_quantized) {
        // Dequantize output for quantized model
        int8_t* output_data = controller->output_tensor->data.int8;
        float output_scale = controller->output_tensor->params.scale;
        int output_zero_point = controller->output_tensor->params.zero_point;
        
        for (int i = 0; i < 3; i++) {
            output_values[i] = (output_data[i] - output_zero_point) * output_scale;
        }
    } else {
        // Copy output for float model
        float* output_data = controller->output_tensor->data.f;
        memcpy(output_values, output_data, 3 * sizeof(float));
    }
    
    // Find predicted class and confidence
    int predicted_class = 0;
    float max_confidence = output_values[0];
    
    for (int i = 1; i < 3; i++) {
        if (output_values[i] > max_confidence) {
            max_confidence = output_values[i];
            predicted_class = i;
        }
    }
    
    // Set result
    result.posture_index = predicted_class;
    result.confidence = max_confidence * 100.0f; // Convert to percentage
    result.success = true;
    
    // Store current posture and confidence
    controller->current_posture = predicted_class;
    controller->confidence = result.confidence;
    
    // Map class to posture name for debugging
    const char* posture_names[] = {"Supine", "Left Side", "Right Side"};
    printf("Detected posture: %s (confidence: %.1f%%)\n", 
           posture_names[predicted_class], result.confidence);
    
    return result;
}

/**
 * Send command to the mattress controller
 */
bool send_mattress_command(MCUController* controller, uint8_t mode_byte, uint8_t* args, size_t args_length) {
    if (!controller->mattress_ser) {
        printf("Mattress not connected\n");
        return false;
    }
    
    // Create command packet
    uint8_t packet[20]; // Allocate enough space for command
    size_t packet_index = 0;
    
    // Start byte
    packet[packet_index++] = 0xAA;
    
    // Mode byte
    packet[packet_index++] = mode_byte;
    
    // Parameter bytes
    for (size_t i = 0; i < args_length && packet_index < 14; i++) {
        packet[packet_index++] = args[i];
    }
    
    // Pad to fixed length if needed
    while (packet_index < 14) {
        packet[packet_index++] = 0x00;
    }
    
    // Add ID bytes
    for (int i = 0; i < 4; i++) {
        packet[packet_index++] = controller->id_bytes[i];
    }
    
    // End byte
    packet[packet_index++] = 0xBB;
    
    // Send packet
    if (serial_write(controller->mattress_ser, packet, packet_index)) {
        printf("Sent mattress command: mode=%d\n", mode_byte);
        return true;
    } else {
        printf("Failed to send mattress command\n");
        return false;
    }
}

/**
 * Stop the current operating mode
 */
bool stop_current_mode(MCUController* controller) {
    // Set stop event
    controller->stop_event = true;
    
    // Wait for threads to finish (simple implementation)
    // In a real-time system, you would use proper thread joining with timeout
    struct timespec ts;
    ts.tv_sec = 1;
    ts.tv_nsec = 0;
    nanosleep(&ts, NULL);
    
    // Reset thread count and stop event
    pthread_mutex_lock(&controller->thread_mutex);
    controller->thread_count = 0;
    controller->stop_event = false;
    pthread_mutex_unlock(&controller->thread_mutex);
    
    // Send stop command to mattress
    if (controller->mattress_ser) {
        uint8_t args[1] = {0};
        send_mattress_command(controller, MODE_MANUAL, args, 0);
    }
    
    return true;
}

/**
 * Run adaptive mode based on posture detection
 */
bool run_adaptive_mode(MCUController* controller, int detection_interval, int runtime) {
    if (!controller->mattress_ser || !controller->sensor_ser) {
        printf("Both mattress and sensor must be connected\n");
        return false;
    }
    
    // Stop any current mode
    stop_current_mode(controller);
    
    // Create thread arguments (allocated on heap to avoid stack issues)
    typedef struct {
        MCUController* controller;
        int detection_interval;
        int runtime;
    } AdaptiveArgs;
    
    AdaptiveArgs* args = (AdaptiveArgs*)malloc(sizeof(AdaptiveArgs));
    args->controller = controller;
    args->detection_interval = detection_interval;
    args->runtime = runtime;
    
    // Start adaptive thread
    pthread_mutex_lock(&controller->thread_mutex);
    
    if (controller->thread_count >= 5) {
        printf("Too many threads running\n");
        pthread_mutex_unlock(&controller->thread_mutex);
        free(args);
        return false;
    }
    
    pthread_t thread;
    if (pthread_create(&thread, NULL, adaptive_thread_func, args) != 0) {
        printf("Failed to create adaptive thread\n");
        pthread_mutex_unlock(&controller->thread_mutex);
        free(args);
        return false;
    }
    
    controller->running_threads[controller->thread_count++] = thread;
    pthread_mutex_unlock(&controller->thread_mutex);
    
    printf("Started adaptive mode\n");
    return true;
}

// Implementation of adaptive thread function
void* adaptive_thread_func(void* arg) {
    typedef struct {
        MCUController* controller;
        int detection_interval;
        int runtime;
    } AdaptiveArgs;
    
    AdaptiveArgs* args = (AdaptiveArgs*)arg;
    MCUController* controller = args->controller;
    int detection_interval = args->detection_interval;
    int runtime = args->runtime;
    
    // Free args struct - we've copied the data
    free(args);
    
    printf("Adaptive thread starting (interval=%d, runtime=%d)\n", detection_interval, runtime);
    
    int adjustment_count = 0;
    time_t start_time = time(NULL);
    
    while (!controller->stop_event) {
        // Check if runtime exceeded
        time_t current_time = time(NULL);
        int elapsed_time = (int)(current_time - start_time);
        
        if (elapsed_time >= runtime) {
            printf("Adaptive mode completed after %d seconds\n", elapsed_time);
            break;
        }
        
        // Calculate progress percentage
        int progress = (elapsed_time * 100) / runtime;
        printf("Adaptive mode progress: %d%%\n", progress);
        
        // Detect current posture
        PostureResult result = detect_posture(controller);
        
        if (result.success) {
            int posture_index = result.posture_index;
            float confidence = result.confidence;
            
            // Only adjust if confidence is high enough
            if (confidence >= 70.0f) {
                uint8_t args[4];
                
                // Adjust mattress based on posture
                if (posture_index == POSTURE_SUPINE) {
                    // Supine: neutral position
                    args[0] = 50;
                    args[1] = 50;
                    args[2] = 50;
                    args[3] = 50;
                } else if (posture_index == POSTURE_LEFT) {
                    // Left side: adjust for comfort
                    args[0] = 70;
                    args[1] = 40;
                    args[2] = 70;
                    args[3] = 40;
                } else if (posture_index == POSTURE_RIGHT) {
                    // Right side: adjust for comfort
                    args[0] = 40;
                    args[1] = 70;
                    args[2] = 40;
                    args[3] = 70;
                }
                
                send_mattress_command(controller, MODE_MANUAL, args, 4);
                adjustment_count++;
                
                const char* posture_names[] = {"Supine", "Left Side", "Right Side"};
                printf("Adjusted mattress for %s (adjustment #%d)\n", 
                       posture_names[posture_index], adjustment_count);
            } else {
                printf("Low confidence (%.1f%%), no adjustment made\n", confidence);
            }
        } else {
            printf("Posture detection failed, no adjustment made\n");
        }
        
        // Wait for next detection interval or until stopped
        for (int i = 0; i < detection_interval && !controller->stop_event; i++) {
            struct timespec ts;
            ts.tv_sec = 1;
            ts.tv_nsec = 0;
            nanosleep(&ts, NULL);
        }
    }
    
    // Reset to manual mode when finished
    if (!controller->stop_event) {
        uint8_t args[1] = {0};
        send_mattress_command(controller, MODE_MANUAL, args, 0);
        printf("Adaptive mode completed\n");
    } else {
        printf("Adaptive mode stopped\n");
    }
    
    return NULL;
}

/**
 * Process a command from the remote controller
 */
bool process_remote_command(MCUController* controller, uint8_t command_byte, uint8_t param1, uint8_t param2) {
    printf("Processing remote command: cmd=%d, param1=%d, param2=%d\n", command_byte, param1, param2);
    
    switch (command_byte) {
        case CMD_STOP:
            return stop_current_mode(controller);
            
        case CMD_ADAPTIVE_MODE: {
            // param1: interval in minutes, param2: runtime in hours
            int interval = param1 > 0 ? param1 * 60 : 300;
            int runtime = param2 > 0 ? param2 * 3600 : 3600;
            return run_adaptive_mode(controller, interval, runtime);
        }
            
        case CMD_PREVENTION_MODE: {
            // Implementation would be similar to adaptive mode
            printf("Prevention mode not fully implemented\n");
            return false;
        }
            
        case CMD_ABAB_MODE: {
            // Implementation would be similar to adaptive mode
            printf("ABAB mode not fully implemented\n");
            return false;
        }
            
        case CMD_TURNING_MODE: {
            // Implementation would be similar to adaptive mode
            printf("Turning mode not fully implemented\n");
            return false;
        }
            
        case CMD_FLOATING_MODE: {
            // Implementation would be similar to adaptive mode
            printf("Floating mode not fully implemented\n");
            return false;
        }
            
        case CMD_DETECT_POSTURE: {
            // Just detect and report posture
            PostureResult result = detect_posture(controller);
            return result.success;
        }
            
        default:
            printf("Unknown command: %d\n", command_byte);
            return false;
    }
}

/**
 * Listen for remote controller commands
 */
void* remote_listener_thread(void* arg) {
    MCUController* controller = (MCUController*)arg;
    
    // Buffer for received bytes
    uint8_t buffer[16];
    int buffer_index = 0;
    
    while (1) {
        // Check if data is available
        if (serial_available(controller->remote_ser) > 0) {
            // Read one byte
            uint8_t byte;
            if (serial_read(controller->remote_ser, &byte, 1) == 1) {
                buffer[buffer_index++] = byte;
                
                // Check for complete command (3 bytes)
                if (buffer_index >= 3) {
                    // Process command
                    uint8_t cmd = buffer[0];
                    uint8_t param1 = buffer[1];
                    uint8_t param2 = buffer[2];
                    
                    printf("Received command: %d %d %d\n", cmd, param1, param2);
                    
                    // Process the command
                    process_remote_command(controller, cmd, param1, param2);
                    
                    // Clear buffer
                    buffer_index = 0;
                }
            }
        }
        
        // Small delay to prevent CPU overuse
        struct timespec ts;
        ts.tv_sec = 0;
        ts.tv_nsec = 10000000; // 10ms
        nanosleep(&ts, NULL);
    }
    
    return NULL;
}

/**
 * Start remote controller listener
 */
bool start_remote_listener(MCUController* controller, const char* port, int baudrate) {
    // Connect to remote controller
    if (controller->remote_ser) {
        serial_close(controller->remote_ser);
        controller->remote_ser = NULL;
    }
    
    if (!serial_open(&controller->remote_ser, port, baudrate)) {
        printf("Failed to connect to remote controller on %s\n", port);
        return false;
    }
    
    printf("Connected to remote controller on %s\n", port);
    
    // Start listener thread
    pthread_t thread;
    if (pthread_create(&thread, NULL, remote_listener_thread, controller) != 0) {
        printf("Failed to create remote listener thread\n");
        serial_close(controller->remote_ser);
        controller->remote_ser = NULL;
        return false;
    }
    
    // Make thread detached so it cleans up automatically
    pthread_detach(thread);
    
    return true;
}

/**
 * Main entry point
 */
int main() {
    // Create controller
    MCUController controller;
    
    printf("Initializing MCU Mattress Controller...\n");
    
    // Initialize controller
    if (!controller_init(&controller)) {
        printf("Failed to initialize controller\n");
        return 1;
    }
    
    // Connect to devices
    connect_mattress(&controller, MATTRESS_PORT, 115200);
    connect_sensor(&controller, SENSOR_PORT, 115200);
    
    // Start remote controller listener
    start_remote_listener(&controller, REMOTE_PORT, 9600);
    
    // Main loop
    printf("Controller running...\n");
    
    // In a real implementation, you might need some form of
    // application main loop here. For simplicity, we just sleep.
    while (1) {
        struct timespec ts;
        ts.tv_sec = 1;
        ts.tv_nsec = 0;
        nanosleep(&ts, NULL);
    }
    
    // Clean up (unreachable in this simple example)
    disconnect_all(&controller);
    
    return 0;
}

// Platform-specific serial communication implementations would go here
// These are just stubs and would need to be implemented for your specific platform

bool serial_open(void** handle, const char* port, int baudrate) {
    // Implementation depends on your platform (UART, USART, etc.)
    printf("Opening serial port %s at %d baud\n", port, baudrate);
    *handle = malloc(1); // Dummy handle
    return true;
}

void serial_close(void* handle) {
    // Implementation depends on your platform
    if (handle) free(handle);
}

bool serial_write(void* handle, const uint8_t* data, size_t length) {
    // Implementation depends on your platform
    printf("Writing %zu bytes to serial port\n", length);
    return true;
}

int serial_read(void* handle, uint8_t* buffer, size_t max_length) {
    // Implementation depends on your platform
    return 0; // Return number of bytes read
}

int serial_available(void* handle) {
    // Implementation depends on your platform
    return 0; // Return number of bytes available
}