lambda_function.py

import json
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
import joblib
import os
import boto3

# --- 1. CONFIGURATION ---
WINDOW_SIZE = 20  # Must match the value used in dataset_lstm.py
FEATURES = ['eda', 'temp', 'hr', 'hrv']
CLASSES = ['happy', 'neutral', 'relaxed', 'stressed'] # NOTE: Check the alphabetical order from LabelEncoder!
CONTROL_TOPIC = 'tracker/control/commands' # Topic to publish the final command

# Global variables for model and scaler (Cached loading)
global model
global scaler
global mqtt_client # Use a client for publishing back to IoT Core

# --- 2. CACHED INITIALIZATION ---
# This runs only once when the Lambda container starts
def init_resources():
    global model
    global scaler
    global mqtt_client

    # 1. Load the Model and Scaler from a persistent location (e.g., /tmp, S3)
    #    In AWS, you typically zip these files with your Lambda code or load them from S3.
    #    We assume the model/scaler files are packaged with the deployment.
    try:
        # Load Keras Model (must be in the deployment package)
        model = load_model('wesad_lstm_model.h5')
        
        # Load StandardScaler object (joblib is commonly used for this)
        scaler = joblib.load('scaler_params.joblib') # You need to save this object during training!
        
        # Initialize IoT Core client for publishing control commands
        mqtt_client = boto3.client('iot-data', region_name=os.environ['AWS_REGION'])
        
        print("Resources loaded successfully.")

    except Exception as e:
        print(f"Error loading resources: {e}")
        model = None
        scaler = None
        mqtt_client = None

# Run initialization when the module is imported
init_resources()


# --- 3. INFERENCE HELPER FUNCTION (Simulates Time Series Logic) ---
# NOTE: This is the simplified version. A production system needs to store
# the previous 19 windows in DynamoDB or ElastiCache to form the full 20-timestep window.
# For now, we simulate a single window prediction.

def process_and_predict(feature_values):
    """Takes a single window of features, scales it, and returns the prediction."""
    
    if model is None or scaler is None:
        raise Exception("Model or Scaler failed to load.")
    
    # 1. Convert incoming single row (4 features) to a NumPy array
    #    The input structure must be (1, 1, 4) for a single timestep prediction
    #    For the full 20-timestep window, you would load 19 prior windows here.
    
    # *** CRITICAL ASSUMPTION: ***
    # Since we can't maintain state easily in this template, we assume the Lambda 
    # receives a full (20, 4) array OR we use the single point to predict the immediate state.
    # We will format the single point (1, 1, 4) for basic testing.
    
    # Convert list to array, and then reshape to (1, 1, 4)
    # NOTE: You must adjust the window to (1, 20, 4) if you pass a full window.
    X_raw = np.array([feature_values], dtype=np.float32)
    
    # 2. Reshape and Scale (simulating the flattening done during training)
    # The scaler was fitted on a flattened (N*T, F) array.
    X_flat = X_raw.reshape(-1, len(FEATURES)) # Reshape to (Timesteps, Features)
    X_scaled = scaler.transform(X_flat)
    
    # 3. Reshape back to LSTM input (1 sample, 1 timestep, 4 features)
    # NOTE: In a production system, this would be (1, 20, 4).
    X_input = X_scaled.reshape(1, X_raw.shape[0], len(FEATURES))
    
    # 4. Predict
    predictions = model.predict(X_input)
    
    # 5. Get the result
    predicted_index = np.argmax(predictions[0])
    return CLASSES[predicted_index]


# --- 4. ROOM CONTROL DECISION ---
def generate_command(emotion):
    """Maps the predicted emotion to a room control command."""
    if emotion == 'stressed':
        return {"command": "DIM_LIGHTS", "level": 30, "music": "calm"}
    elif emotion == 'happy':
        return {"command": "BRIGHTEN_LIGHTS", "level": 80, "music": "upbeat"}
    elif emotion == 'relaxed':
        return {"command": "MAINTAIN", "level": 50, "music": "ambient"}
    else: # neutral
        return {"command": "DEFAULT", "level": 60, "music": "none"}


# --- 5. LAMBDA HANDLER ---
def lambda_handler(event, context):
    try:
        # The IoT Rule passes the MQTT payload as the event body
        # The event is usually a dictionary containing the MQTT fields.
        
        # 1. Extract Features from MQTT Payload
        # NOTE: Keys are lowercase as sent by ESP32 JSON
        features = [
            event.get('eda'), 
            event.get('temp'), 
            event.get('hr'), 
            event.get('hrv')
        ]
        
        # Check for missing/invalid data before proceeding
        if any(f is None for f in features):
            print("Missing feature data in payload. Skipping.")
            return {"status": "error", "message": "Missing features"}

        # 2. Get Prediction
        # NOTE: This call is simplified to use only the current point.
        # It needs to be wrapped in 20 timesteps of historic data for a true LSTM prediction.
        predicted_emotion = process_and_predict(features)

        # 3. Determine Command
        room_command = generate_command(predicted_emotion)
        
        # 4. Publish Command back to IoT Core
        # The response topic is used by a hypothetical device/service running the room controls
        mqtt_client.publish(
            topic=CONTROL_TOPIC,
            qos=1,
            payload=json.dumps({"emotion": predicted_emotion, "action": room_command})
        )
        
        print(f"Prediction: {predicted_emotion}, Action: {room_command['command']}")
        
        return {
            "statusCode": 200,
            "body": predicted_emotion,
            "command": room_command['command']
        }

    except Exception as e:
        print(f"Fatal error during Lambda execution: {e}")
        return {"status": "failed", "error": str(e)}