app.py

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from contextlib import asynccontextmanager
import torch
import torch.nn as nn
import joblib
import numpy as np
import pandas as pd
import warnings
import os

warnings.filterwarnings("ignore", category=UserWarning)

class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(LSTMModel, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)
        self.fc1 = nn.Linear(hidden_size, 32)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(32, output_size)
        
    def forward(self, x):
        out, _ = self.lstm(x)
        out = self.fc1(out[:, -1, :])
        out = self.relu(out)
        out = self.fc2(out)
        return out

# Global variables for model and data
model = None
scaler = None
test_df_processed = None
feature_cols = None
SEQ_LENGTH = 30

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load artifacts on startup
    global model, scaler, test_df_processed, feature_cols
    print("===== Starting Application Startup =====")
    try:
        print("Loading model and scaler...")
        # Ensure files exist
        if not os.path.exists("scaler.pkl") or not os.path.exists("lstm_model.pth"):
            raise FileNotFoundError("Missing model or scaler artifacts!")
            
        scaler = joblib.load("scaler.pkl")
        model = LSTMModel(input_size=45, hidden_size=64, num_layers=2, output_size=1)
        
        # Explicitly load to CPU to avoid issues on CPU-only environments
        model.load_state_dict(torch.load("lstm_model.pth", map_location=torch.device('cpu')))
        model.eval()
        
        print("Loading and preprocessing test_FD001.txt...")
        if not os.path.exists("test_FD001.txt"):
            raise FileNotFoundError("test_FD001.txt not found!")
            
        op_cols = ["op_1", "op_2", "op_3"]
        sensor_cols = [f"sensor_{i}" for i in range(1, 22)]
        columns = ["engine_id", "cycle"] + op_cols + sensor_cols
        
        test_df = pd.read_csv("test_FD001.txt", sep=r"\s+", header=None)
        test_df.columns = columns
        
        drop_sensors = ["sensor_1", "sensor_5", "sensor_6", "sensor_10", "sensor_16", "sensor_18", "sensor_19"]
        test_df = test_df.drop(columns=drop_sensors)
        test_df = test_df.sort_values(["engine_id", "cycle"])
        
        top_sensors = ["sensor_11", "sensor_9", "sensor_4", "sensor_12", "sensor_14", "sensor_7", "sensor_15", "sensor_21", "sensor_2"]
        
        for sensor in top_sensors:
            test_df[f"{sensor}_rollmean"] = test_df.groupby("engine_id")[sensor].rolling(window=5, min_periods=1).mean().reset_index(level=0, drop=True)
            test_df[f"{sensor}_rollstd"] = test_df.groupby("engine_id")[sensor].rolling(window=5, min_periods=1).std().reset_index(level=0, drop=True)
            test_df[f"{sensor}_delta"] = test_df.groupby("engine_id")[sensor].diff()
            
        test_df = test_df.fillna(0)
        feature_cols = [c for c in test_df.columns if c not in ["RUL", "engine_id", "max_cycle"]]
        
        # Scale test set
        test_df[feature_cols] = scaler.transform(test_df[feature_cols].values)
        
        test_df_processed = test_df
        print("Successfully loaded artifacts and preprocessed test dataset!")
    except Exception as e:
        print(f"CRITICAL STARTUP ERROR: {str(e)}")
        # We don't raise here so the health check can still pass and we can see logs
        # But we should probably log the traceback
        import traceback
        traceback.print_exc()
        
    yield
    # Clean up on shutdown if needed
    print("===== Application Shutdown =====")

app = FastAPI(title="AeroPulse RUL API", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Artifacts are now loaded via the lifespan handler

@app.get("/engine/{engine_id}/history")
def get_engine_history(engine_id: int):
    if test_df_processed is None:
        raise HTTPException(status_code=500, detail="Data not loaded properly on server startup.")
        
    engine_data = test_df_processed[(test_df_processed["engine_id"] == engine_id)]
    if engine_data.empty:
        raise HTTPException(status_code=404, detail="Engine not found in dataset.")
        
    features = engine_data[feature_cols].values
    
    if len(features) >= SEQ_LENGTH:
        X_seq = features[-SEQ_LENGTH:]
    else:
        pad = np.zeros((SEQ_LENGTH - len(features), len(feature_cols)))
        X_seq = np.vstack([pad, features])
        
    X_test_t = torch.tensor(X_seq, dtype=torch.float32).unsqueeze(0) # shape (1, 30, 45)
    
    with torch.no_grad():
        pred = model(X_test_t)
        
    # Get last 20 cycles for the charts (unscaled versions preferably, but we scaled them inplace)
    # Let's get them from unscaled inverse transform or just return the scaled for now since 
    # the frontend draws the shape. For accuracy, let's reverse transform.
    history = engine_data.tail(20)
    unscaled_hist = scaler.inverse_transform(history[feature_cols].values)
    history_df = pd.DataFrame(unscaled_hist, columns=feature_cols)
    history_df["cycle"] = history["cycle"].values
    
    cycles = history_df["cycle"].tolist()
    sensor_11 = history_df["sensor_11"].tolist()
    sensor_14 = history_df["sensor_14"].tolist()
    
    return {
        "engine_id": engine_id,
        "current_cycle": int(cycles[-1]),
        "predicted_rul": float(pred.item()),
        "history": {
            "cycles": cycles,
            "sensor_11": sensor_11,
            "sensor_14": sensor_14
        }
    }

@app.get("/fleet")
def get_fleet_status():
    if test_df_processed is None:
        raise HTTPException(status_code=500, detail="Data not loaded.")
    
    X_test_seq = []
    engine_ids = []
    cycles = []
    
    for engine_id, group in test_df_processed.groupby("engine_id"):
        engine_ids.append(engine_id)
        cycles.append(group["cycle"].iloc[-1])
        features = group[feature_cols].values
        if len(features) >= SEQ_LENGTH:
            X_test_seq.append(features[-SEQ_LENGTH:])
        else:
            pad = np.zeros((SEQ_LENGTH - len(features), len(feature_cols)))
            X_test_seq.append(np.vstack([pad, features]))
            
    X_test_t = torch.tensor(np.array(X_test_seq), dtype=torch.float32)
    
    with torch.no_grad():
        preds = model(X_test_t).flatten().numpy()
        
    fleet = []
    for i in range(len(engine_ids)):
        fleet.append({
            "id": int(engine_ids[i]),
            "cycles": int(cycles[i]),
            "rul": float(preds[i])
        })
        
    # Sort fleet by RUL ascending
    fleet.sort(key=lambda x: x["rul"])
    return fleet

@app.get("/health")
def health_check():
    return {
        "status": "Healthy" if test_df_processed is not None else "Degraded",
        "artifacts_loaded": test_df_processed is not None
    }

@app.get("/engine_bg.png")
def serve_bg():
    return FileResponse("engine_bg.png")

@app.get("/")
def serve_ui():
    return FileResponse("index.html")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)