"""Predictive maintenance for equipment failure prediction using Gradient Boosting."""
import logging, numpy as np
from typing import Dict, List
from datetime import datetime

logger = logging.getLogger("deltamind.predictive")

class PredictiveMaintenance:
    def __init__(self):
        self.model = None
        self._train()

    def _train(self):
        try:
            from sklearn.ensemble import GradientBoostingClassifier
            np.random.seed(42)
            n = 1000
            # Features: vibration, temperature, runtime_hours, pressure
            vibration = np.concatenate([np.random.normal(0.2, 0.05, int(n*0.8)), np.random.normal(0.7, 0.1, int(n*0.2))])
            temperature = np.concatenate([np.random.normal(65, 5, int(n*0.8)), np.random.normal(95, 8, int(n*0.2))])
            runtime = np.concatenate([np.random.normal(4000, 1000, int(n*0.8)), np.random.normal(7500, 500, int(n*0.2))])
            pressure = np.concatenate([np.random.normal(150, 10, int(n*0.8)), np.random.normal(110, 20, int(n*0.2))])
            
            X = np.column_stack([vibration, temperature, runtime, pressure])
            y = np.concatenate([np.zeros(int(n*0.8)), np.ones(int(n*0.2))])
            
            self.model = GradientBoostingClassifier(n_estimators=100, max_depth=4, random_state=42)
            self.model.fit(X, y)
            logger.info("Predictive maintenance model trained successfully")
        except Exception as e:
            logger.error(f"Predictive model training failed: {e}")

    def predict(self, equipment: Dict) -> Dict:
        if not self.model:
            return {"failure_probability": 0.5, "days_to_failure": 30, "severity": "unknown"}
            
        features = np.array([[
            equipment.get("vibration", 0.2), equipment.get("temperature", 65),
            equipment.get("runtime_hours", 4000), equipment.get("pressure", 150)
        ]])
        
        prob = float(self.model.predict_proba(features)[0][1])
        days = max(1, int(90 * (1 - prob)))
        sev = "critical" if prob > 0.8 else "high" if prob > 0.6 else "medium" if prob > 0.4 else "low"
        
        actions = {
            "critical": "Immediate shutdown and workover required.", 
            "high": "Schedule workover within 5 days. Monitor daily.",
            "medium": "Monitor closely, schedule maintenance within 30 days.", 
            "low": "Continue normal operations. Next routine check."
        }
        
        return {
            "equipment_id": equipment.get("id", "unknown"),
            "equipment_type": equipment.get("type", "ESP"),
            "failure_probability": round(prob, 3), 
            "predicted_days_to_failure": days,
            "severity": sev, 
            "recommended_action": actions[sev],
            "health_score": round((1 - prob) * 100, 1),
            "timestamp": datetime.now().isoformat()
        }

    def fleet_health(self) -> Dict:
        np.random.seed(int(datetime.now().timestamp()) % 10000)
        fleet = []
        for i in range(20):
            eq = {
                "id": f"ESP-{i+1:03d}", "type": "ESP", "oml_id": f"OML-{[14,18,22,29,58][i%5]}",
                "vibration": round(np.random.uniform(0.1, 0.8), 3),
                "temperature": round(np.random.uniform(55, 100), 1),
                "runtime_hours": int(np.random.uniform(1000, 8000)),
                "pressure": round(np.random.uniform(100, 200), 1)
            }
            fleet.append(self.predict(eq))
            
        critical = len([e for e in fleet if e["severity"] == "critical"])
        high = len([e for e in fleet if e["severity"] == "high"])
        
        return {
            "fleet_size": len(fleet), "critical": critical, "high": high,
            "avg_health": round(sum(e["health_score"] for e in fleet) / len(fleet), 1),
            "equipment": fleet
        }

predictive_model = PredictiveMaintenance()