controller.py

import logging

from model import Model
from utils import (
    calculate_risk_score,
    convert_cyclical_to_original,
    create_timestamp_from_predictions,
    get_failure_severity,
    get_failure_type_name,
    get_risk_level,
    prepare_prediction_data,
    prepare_sensor_data_for_anomaly,
)
import numpy as np
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)


class Controller:
    def __init__(self, database_sensor):
        self.__sensor = database_sensor
        self.__model = Model()
        logger.info("Controller initialized with Model")
    
    def set_sensor_data(self, sensor):
        self.__sensor = sensor
        if isinstance(sensor, list):
            logger.debug(f"Sensor data updated with {len(sensor)} readings")
        else:
            logger.debug(f"Sensor data updated for machine_id: {sensor.get('machine_id')}")
    
    # Binary method
    def predict_binary(self):
        if self.__sensor is None:
            logger.warning("Binary prediction attempted with no sensor data")
            return {
                "success": False,
                "error": "No sensor data available from database.",
            }
        
        # Handle both list and dict cases
        if isinstance(self.__sensor, dict) and self.__sensor.get("message") == "Data already predicted":
            logger.info(f"Skipping binary prediction - data already predicted for UDI: {self.__sensor.get('udi')}")
            return {
                "success": False,
                "error": "Data already predicted",
            }
        
        if (
            self.__model.model_binary is None
            or self.__model.preprocessor_anomaly is None
        ):
            logger.error("Binary model or preprocessor not loaded - cannot perform prediction")
            return {
                "success": False,
                "error": "Binary model or preprocessor not loaded.",
            }
            
        try:
            # For binary, we always use single sensor data (dict)
            sensor_dict = self.__sensor if isinstance(self.__sensor, dict) else self.__sensor[0]
            
            X_scaled = prepare_sensor_data_for_anomaly(
                sensor_dict, self.__model.preprocessor_anomaly
            )
    
            if X_scaled is None:
                logger.error("Failed to prepare sensor data for binary prediction")
                return {"success": False, "error": "Failed to prepare sensor data."}
    
            if hasattr(self.__model.model_binary, "predict_proba"):
                probabilities = self.__model.model_binary.predict_proba(X_scaled)[0]
                confidence_normal = float(probabilities[0])
                confidence_error = float(probabilities[1])
                
                is_error = int(confidence_error > 0.5)
                confidence = confidence_error if is_error else confidence_normal
            else:
                prediction = self.__model.model_binary.predict(X_scaled)
                is_error = int(prediction[0])
                confidence = 1.0 if is_error else 0.0
    
            risk_score = confidence * 100 if is_error else (1 - confidence) * 100
    
            logger.info(f"Binary prediction complete - Failure: {bool(is_error)}, Confidence: {confidence:.2f}, Risk: {risk_score:.2f}")
            
            result = {
                "success": True,
                "failure_predicted": bool(is_error),
                "confidence": float(f"{confidence}"),
                "risk_score": float(f"{risk_score}"),
            }
            return result
        except Exception as e:
            logger.error(f"Exception in binary prediction: {str(e)}")
            return {
                "success": False,
                "message": f"Failed to predict binary data: {str(e)}"
            }

    # Classification Method
    def predict_classification(self):
        if self.__sensor is None:
            logger.warning("Classification prediction attempted with no sensor data")
            return {
                "success": False,
                "error": "No sensor data available from database.",
            }
        
        # Handle both list and dict cases
        if isinstance(self.__sensor, dict) and self.__sensor.get("message") == "Data already predicted":
            logger.info(f"Skipping classification prediction - data already predicted for UDI: {self.__sensor.get('udi')}")
            return {
                "success": False,
                "error": "Data already predicted",
            }

        if (
            self.__model.model_multiclass is None
            or self.__model.preprocessor_anomaly is None
        ):
            logger.error("Multiclass model or preprocessor not loaded - cannot perform prediction")
            return {"success": False, "error": "Model or scalers not loaded."}
        
        try:
            # For classification, we always use single sensor data (dict)
            sensor_dict = self.__sensor if isinstance(self.__sensor, dict) else self.__sensor[0]
            
            X_scaled = prepare_sensor_data_for_anomaly(
                sensor_dict, self.__model.preprocessor_anomaly
            )
    
            if X_scaled is None:
                logger.error("Failed to prepare sensor data for classification prediction")
                return {"success": False, "error": "Failed to prepare sensor data."}
    
            prediction = self.__model.model_multiclass.predict(X_scaled)
    
            if hasattr(self.__model.model_multiclass, "predict_proba"):
                probabilities = self.__model.model_multiclass.predict_proba(X_scaled)[0]
                failure_index = int(probabilities.argmax())
                confidence = float(probabilities[failure_index])
    
                all_probs = {
                    get_failure_type_name(i): float(prob)
                    for i, prob in enumerate(probabilities)
                }
            else:
                failure_index = int(prediction[0])
                confidence = 1.0
                all_probs = None
    
            failure_type = get_failure_type_name(failure_index)
    
            severity = get_failure_severity(failure_index)
            risk_score = calculate_risk_score(confidence, severity)
            risk_level = get_risk_level(risk_score)
    
            logger.info(f"Classification prediction complete - Failure Type: {failure_type}, Confidence: {confidence:.2f}, Risk Level: {risk_level}, Risk Score: {risk_score:.2f}")
            
            return {
                "success": True,
                "failure_type": failure_type,
                "confidence": float(f"{confidence}"),
                "risk_score": float(f"{risk_score}"),
                "risk_level": risk_level,
                "all_probabilities": all_probs,
            }
        except Exception as e:
            logger.error(f"Exception in classification prediction: {str(e)}")
            return {
                "success": False,
                "message": f"Failed to predict classification data: {str(e)}"
            }

    # Time series method
    def predict_time_series(self):
        # 1. Validation Checks
        if self.__sensor is None:
            logger.warning("Time series prediction attempted with no sensor data")
            return {"success": False, "error": "No sensor data available."}
        
        if isinstance(self.__sensor, dict) and self.__sensor.get("message") == "Data already predicted":
            logger.info(f"Skipping time series - data already predicted for UDI: {self.__sensor.get('udi')}")
            return {"success": False, "error": "Data already predicted"}

        if self.__model.model_lstm is None or self.__model.scaler_lstm is None:
            logger.error("LSTM model or scaler_x not loaded")
            return {"success": False, "error": "Model or scalers not loaded."}
            
        try:
            sensor_data = self.__sensor if isinstance(self.__sensor, list) else [self.__sensor]
    
            input_data = []
            for row in sensor_data:
                input_data.append([
                    row.get("air_temp"),
                    row.get("process_temp"),
                    row.get("rotational_speed"),
                    row.get("torque"),
                    row.get("tool_wear", 0) 
                ])
            
            input_array = np.array(input_data)
            if len(input_array) < 30:
                missing = 30 - len(input_array)
                padding = np.tile(input_array[0], (missing, 1))
                input_array = np.vstack([padding, input_array])
            elif len(input_array) > 30:
                input_array = input_array[-30:]
            
            input_scaled = self.__model.scaler_lstm.transform(input_array)
            input_reshaped = input_scaled.reshape(1, 30, 5)
            health_score = float(self.__model.model_lstm.predict(input_reshaped, verbose=0)[0][0])
            
            
            MAX_LIFE_DAYS = 7.0 
            MAX_LIFE_MINUTES = MAX_LIFE_DAYS * 24 * 60
            
            rul_minutes_left = health_score * MAX_LIFE_MINUTES
            days_remaining = health_score * MAX_LIFE_DAYS
            
            now_str = sensor_data[0].get("timestamp")
            if isinstance(now_str, str):
                now = datetime.fromisoformat(now_str.replace('Z', '+00:00'))
            else:
                now = datetime.now()
                
            failure_date = now + timedelta(minutes=rul_minutes_left)
            
            if days_remaining < 1.0:
                status = "Critical"
            elif days_remaining < 3.0:
                status = "Warning"
            else:
                status = "Good"

            logger.info(f"Health: {health_score*100:.1f}% -> {days_remaining:.2f} Days Left")
            
            return {
                "success": True,
                "predictions": {
                    "health_score": round(health_score, 4),           # e.g., 0.95
                    "health_percentage": round(health_score * 100, 2), # e.g., 95.0%
                    "rul_minutes": round(rul_minutes_left, 2),
                    "days_remaining": round(days_remaining, 2),
                    "predicted_failure_date": failure_date.strftime('%Y-%m-%d %H:%M:%S'),
                    "status": status
                },
                "input_timestamp": str(now_str),
            }
        except Exception as e:
            logger.error(f"Exception in time series prediction: {str(e)}")
            return {
                "success": False,
                "message": f"Failed to predict time-series data: {str(e)}"
            }