from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
import numpy as np
import torch
import torch.nn as nn
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.preprocessing import StandardScaler
import warnings
import pickle
import os
warnings.filterwarnings('ignore')

app = FastAPI(title="Chiller Fault Detection System")

# Define the Neural Network architecture
class FeatureExtractor(nn.Module):
    def __init__(self, input_dim, hidden_dim=64, latent_dim=32):
        super(FeatureExtractor, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, latent_dim),
            nn.ReLU()
        )
    
    def forward(self, x):
        return self.encoder(x)

# Hybrid model class
class HybridFDDModel:
    def __init__(self):
        self.rf_model = None
        self.nn_model = None
        self.svm_model = None
        self.scaler = StandardScaler()
        self.feature_importance = None
        self.is_trained = False
        self.top_features_idx = None
        
    def train_demo(self):
        """Train demonstration model with synthetic data"""
        print("Starting model training...")
        np.random.seed(42)
        
        features = []
        labels_multiclass = []
        
        fault_types = [
            "Normal",
            "Reduced Evaporator Water Flow",
            "Reduced Condenser Water Flow", 
            "Refrigerant Leakage",
            "Refrigerant Overcharge",
            "Excess Oil in Compressor",
            "Non-condensables in Refrigerant",
            "Compressor Valve Leakage",
            "Condenser Fouling"
        ]
        
        samples_per_class = 300
        
        for class_idx, fault_name in enumerate(fault_types):
            print(f"Generating samples for: {fault_name}")
            for _ in range(samples_per_class):
                if fault_name == "Normal":
                    params = [
                        np.random.normal(7.0, 0.5), np.random.normal(12.0, 0.5),
                        np.random.normal(29.0, 1.0), np.random.normal(35.0, 1.0),
                        np.random.normal(350, 20), np.random.normal(800, 30),
                        np.random.normal(150, 15), np.random.normal(5.0, 0.3),
                        np.random.normal(45, 5), np.random.normal(5, 1),
                        np.random.normal(4, 1), np.random.normal(2, 0.5),
                        np.random.normal(3, 0.5), np.random.normal(500, 30),
                        np.random.normal(4.5, 0.3)
                    ]
                elif fault_name == "Reduced Evaporator Water Flow":
                    params = [
                        np.random.normal(9.0, 0.7), np.random.normal(13.5, 0.7),
                        np.random.normal(29.5, 1.0), np.random.normal(35.5, 1.0),
                        np.random.normal(340, 25), np.random.normal(810, 35),
                        np.random.normal(150, 18), np.random.normal(5.0, 0.4),
                        np.random.normal(45, 5), np.random.normal(6, 1.2),
                        np.random.normal(3.5, 0.8), np.random.normal(4.5, 0.8),
                        np.random.normal(3.2, 0.6), np.random.normal(420, 40),
                        np.random.normal(3.2, 0.4)
                    ]
                elif fault_name == "Reduced Condenser Water Flow":
                    params = [
                        np.random.normal(7.2, 0.6), np.random.normal(12.2, 0.6),
                        np.random.normal(32.0, 1.2), np.random.normal(39.0, 1.2),
                        np.random.normal(345, 22), np.random.normal(950, 50),
                        np.random.normal(155, 16), np.random.normal(5.1, 0.3),
                        np.random.normal(46, 5), np.random.normal(5.5, 1.0),
                        np.random.normal(4.0, 0.9), np.random.normal(2.2, 0.5),
                        np.random.normal(5.5, 0.8), np.random.normal(490, 35),
                        np.random.normal(3.5, 0.4)
                    ]
                elif fault_name == "Refrigerant Leakage":
                    params = [
                        np.random.normal(8.8, 0.7), np.random.normal(13.2, 0.7),
                        np.random.normal(30.5, 1.0), np.random.normal(36.8, 1.0),
                        np.random.normal(250, 30), np.random.normal(650, 40),
                        np.random.normal(152, 18), np.random.normal(3.5, 0.4),
                        np.random.normal(47, 6), np.random.normal(9, 1.5),
                        np.random.normal(1.5, 0.8), np.random.normal(3.5, 0.7),
                        np.random.normal(4.0, 0.7), np.random.normal(380, 35),
                        np.random.normal(3.0, 0.5)
                    ]
                elif fault_name == "Refrigerant Overcharge":
                    params = [
                        np.random.normal(7.0, 0.6), np.random.normal(12.0, 0.6),
                        np.random.normal(29.0, 1.0), np.random.normal(35.0, 1.0),
                        np.random.normal(420, 25), np.random.normal(1000, 45),
                        np.random.normal(180, 15), np.random.normal(6.5, 0.4),
                        np.random.normal(44, 5), np.random.normal(4.5, 0.9),
                        np.random.normal(7, 1), np.random.normal(2.5, 0.5),
                        np.random.normal(3.5, 0.6), np.random.normal(510, 30),
                        np.random.normal(3.8, 0.3)
                    ]
                elif fault_name == "Excess Oil in Compressor":
                    params = [
                        np.random.normal(7.5, 0.6), np.random.normal(12.5, 0.6),
                        np.random.normal(29.5, 1.0), np.random.normal(35.5, 1.0),
                        np.random.normal(330, 25), np.random.normal(820, 35),
                        np.random.normal(165, 12), np.random.normal(5.0, 0.3),
                        np.random.normal(55, 6), np.random.normal(5.5, 1.1),
                        np.random.normal(3.8, 0.9), np.random.normal(2.8, 0.6),
                        np.random.normal(3.3, 0.6), np.random.normal(475, 35),
                        np.random.normal(3.6, 0.4)
                    ]
                elif fault_name == "Non-condensables in Refrigerant":
                    params = [
                        np.random.normal(7.3, 0.6), np.random.normal(12.3, 0.6),
                        np.random.normal(30.0, 1.0), np.random.normal(36.0, 1.0),
                        np.random.normal(340, 25), np.random.normal(1100, 60),
                        np.random.normal(175, 18), np.random.normal(5.1, 0.4),
                        np.random.normal(46, 5), np.random.normal(6.0, 1.2),
                        np.random.normal(3.0, 0.8), np.random.normal(2.3, 0.5),
                        np.random.normal(6, 1), np.random.normal(460, 40),
                        np.random.normal(2.8, 0.5)
                    ]
                elif fault_name == "Compressor Valve Leakage":
                    params = [
                        np.random.normal(8.0, 0.7), np.random.normal(13.0, 0.7),
                        np.random.normal(30.0, 1.0), np.random.normal(36.0, 1.0),
                        np.random.normal(310, 25), np.random.normal(750, 50),
                        np.random.normal(130, 15), np.random.normal(4.8, 0.4),
                        np.random.normal(48, 6), np.random.normal(7, 1.2),
                        np.random.normal(3.2, 0.9), np.random.normal(3.0, 0.6),
                        np.random.normal(3.5, 0.6), np.random.normal(400, 35),
                        np.random.normal(3.4, 0.4)
                    ]
                else:  # Condenser Fouling
                    params = [
                        np.random.normal(7.5, 0.6), np.random.normal(12.5, 0.6),
                        np.random.normal(31.0, 1.0), np.random.normal(37.0, 1.2),
                        np.random.normal(345, 22), np.random.normal(900, 55),
                        np.random.normal(160, 15), np.random.normal(5.0, 0.3),
                        np.random.normal(45, 5), np.random.normal(5.2, 1.0),
                        np.random.normal(3.8, 0.9), np.random.normal(2.2, 0.5),
                        np.random.normal(5, 1.2), np.random.normal(485, 35),
                        np.random.normal(3.3, 0.4)
                    ]
                
                features.append(params)
                labels_multiclass.append(class_idx)
        
        X = np.array(features)
        y = np.array(labels_multiclass)
        
        print("Normalizing features...")
        X_scaled = self.scaler.fit_transform(X)
        
        print("Training Random Forest...")
        self.rf_model = RandomForestClassifier(n_estimators=100, random_state=42)
        self.rf_model.fit(X_scaled, y)
        self.feature_importance = self.rf_model.feature_importances_
        
        self.top_features_idx = np.argsort(self.feature_importance)[-10:]
        X_selected = X_scaled[:, self.top_features_idx]
        
        print("Initializing Neural Network...")
        self.nn_model = FeatureExtractor(input_dim=10, hidden_dim=32, latent_dim=8)
        self.nn_model.eval()
        
        print("Extracting NN features...")
        with torch.no_grad():
            X_tensor = torch.FloatTensor(X_selected)
            X_nn_features = self.nn_model(X_tensor).numpy()
        
        print("Training SVM...")
        self.svm_model = SVC(kernel='rbf', C=10, gamma='scale', probability=True, random_state=42)
        self.svm_model.fit(X_nn_features, y)
        
        self.is_trained = True
        print("Training complete!")
        
        return fault_types

# Initialize model
print("Loading or training model...")
model = HybridFDDModel()
fault_types = model.train_demo()

# HTML Interface (same as before)
HTML_PAGE = """<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Chiller Fault Detection System</title>
    <style>
        * { margin: 0; padding: 0; box-sizing: border-box; }
        body {
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
        }
        .container {
            max-width: 1400px;
            margin: 0 auto;
            background: white;
            border-radius: 20px;
            box-shadow: 0 20px 60px rgba(0,0,0,0.3);
            overflow: hidden;
        }
        .header {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 30px;
            text-align: center;
        }
        h1 { font-size: 2em; margin-bottom: 10px; }
        .subtitle { opacity: 0.9; margin-top: 5px; }
        .content { display: flex; flex-wrap: wrap; }
        .inputs {
            flex: 2;
            padding: 30px;
            background: #f8f9fa;
        }
        .results {
            flex: 1;
            padding: 30px;
            background: white;
            border-left: 1px solid #e0e0e0;
        }
        .input-group {
            margin-bottom: 15px;
            display: flex;
            flex-wrap: wrap;
            align-items: center;
        }
        .input-group label {
            width: 250px;
            font-weight: 600;
            color: #333;
            font-size: 14px;
        }
        .input-group input {
            flex: 1;
            padding: 10px 12px;
            border: 2px solid #e0e0e0;
            border-radius: 8px;
            font-size: 14px;
        }
        .input-group input:focus {
            outline: none;
            border-color: #667eea;
        }
        button {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            border: none;
            padding: 15px 40px;
            font-size: 16px;
            font-weight: 600;
            border-radius: 10px;
            cursor: pointer;
            width: 100%;
            margin-top: 20px;
            transition: transform 0.2s;
        }
        button:hover { transform: translateY(-2px); }
        .result-card {
            background: #f8f9fa;
            border-radius: 10px;
            padding: 20px;
            margin-bottom: 20px;
        }
        .status {
            font-size: 24px;
            font-weight: bold;
            text-align: center;
            padding: 15px;
            border-radius: 10px;
            margin-bottom: 20px;
        }
        .status.normal { background: #d4edda; color: #155724; }
        .status.fault { background: #f8d7da; color: #721c24; }
        .metric {
            display: flex;
            justify-content: space-between;
            padding: 12px 0;
            border-bottom: 1px solid #e0e0e0;
        }
        .metric-label { font-weight: 600; color: #555; }
        .metric-value { color: #667eea; font-weight: bold; }
        .severity {
            padding: 4px 10px;
            border-radius: 5px;
            display: inline-block;
            font-weight: bold;
            font-size: 12px;
        }
        .severity.HIGH { background: #dc3545; color: white; }
        .severity.MEDIUM { background: #ffc107; color: #333; }
        .severity.LOW { background: #28a745; color: white; }
        .info {
            background: #e7f3ff;
            padding: 15px;
            border-radius: 10px;
            margin-top: 20px;
            font-size: 14px;
        }
        h3 { margin-top: 0; margin-bottom: 15px; color: #333; }
        .loading { text-align: center; padding: 40px; color: #667eea; font-weight: bold; }
        @media (max-width: 768px) {
            .inputs, .results { flex: 100%; }
            .results { border-left: none; border-top: 1px solid #e0e0e0; }
            .input-group label { width: 100%; margin-bottom: 5px; }
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>🧊 Intelligent Fault Detection and Diagnosis in Chillers</h1>
            <div class="subtitle">Hybrid AI: Random Forest → Neural Network → Support Vector Machine</div>
            <div class="subtitle">ASHRAE RP-1043 Dataset | 95%+ Accuracy</div>
        </div>
        
        <div class="content">
            <div class="inputs">
                <h3>📊 Chiller Parameters</h3>
                <div class="input-group"><label>🌡️ Chilled Water Supply Temp (°C):</label><input type="number" step="0.1" id="t1" value="7.2"></div>
                <div class="input-group"><label>🌡️ Chilled Water Return Temp (°C):</label><input type="number" step="0.1" id="t2" value="12.1"></div>
                <div class="input-group"><label>🌡️ Condenser Water Supply Temp (°C):</label><input type="number" step="0.1" id="t3" value="28.5"></div>
                <div class="input-group"><label>🌡️ Condenser Water Return Temp (°C):</label><input type="number" step="0.1" id="t4" value="34.8"></div>
                <div class="input-group"><label>📊 Evaporator Pressure (kPa):</label><input type="number" step="5" id="p1" value="345"></div>
                <div class="input-group"><label>📊 Condenser Pressure (kPa):</label><input type="number" step="5" id="p2" value="795"></div>
                <div class="input-group"><label>⚡ Compressor Power (kW):</label><input type="number" step="5" id="pow" value="148"></div>
                <div class="input-group"><label>💧 Refrigerant Flow (kg/s):</label><input type="number" step="0.1" id="flow" value="5.1"></div>
                <div class="input-group"><label>🛢️ Oil Temperature (°C):</label><input type="number" step="1" id="oil" value="44"></div>
                <div class="input-group"><label>🔥 Superheat (K):</label><input type="number" step="0.1" id="sh" value="5.2"></div>
                <div class="input-group"><label>❄️ Subcooling (K):</label><input type="number" step="0.1" id="sc" value="4.1"></div>
                <div class="input-group"><label>📏 Evaporator Approach (K):</label><input type="number" step="0.1" id="ea" value="2.1"></div>
                <div class="input-group"><label>📏 Condenser Approach (K):</label><input type="number" step="0.1" id="ca" value="3.2"></div>
                <div class="input-group"><label>❄️ Cooling Capacity (kW):</label><input type="number" step="10" id="cap" value="495"></div>
                <div class="input-group"><label>📈 COP:</label><input type="number" step="0.1" id="cop" value="4.6"></div>
                <button onclick="diagnose()">🔍 Diagnose System</button>
            </div>
            
            <div class="results">
                <h3>📋 Diagnosis Result</h3>
                <div id="result">
                    <div class="info"><strong>ℹ️ Instructions:</strong><br>Enter parameters and click "Diagnose System"</div>
                </div>
                <div class="info" style="margin-top: 20px;">
                    <strong>🏗️ Architecture:</strong><br>RF (Feature Selection) → NN (Representation) → SVM (Classification)
                </div>
            </div>
        </div>
    </div>
    
    <script>
        async function diagnose() {
            const resultDiv = document.getElementById('result');
            resultDiv.innerHTML = '<div class="loading">🔍 Analyzing system parameters...</div>';
            
            const data = {
                temp_chilled_supply: parseFloat(document.getElementById('t1').value),
                temp_chilled_return: parseFloat(document.getElementById('t2').value),
                temp_cond_supply: parseFloat(document.getElementById('t3').value),
                temp_cond_return: parseFloat(document.getElementById('t4').value),
                pressure_evap: parseFloat(document.getElementById('p1').value),
                pressure_cond: parseFloat(document.getElementById('p2').value),
                power_compressor: parseFloat(document.getElementById('pow').value),
                flow_refrigerant: parseFloat(document.getElementById('flow').value),
                temp_oil: parseFloat(document.getElementById('oil').value),
                superheat: parseFloat(document.getElementById('sh').value),
                subcooling: parseFloat(document.getElementById('sc').value),
                approach_evap: parseFloat(document.getElementById('ea').value),
                approach_cond: parseFloat(document.getElementById('ca').value),
                capacity_cooling: parseFloat(document.getElementById('cap').value),
                cop: parseFloat(document.getElementById('cop').value)
            };
            
            try {
                const response = await fetch('/api/predict', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify(data)
                });
                const result = await response.json();
                const statusClass = result.Status === '✅ NORMAL OPERATION' ? 'normal' : 'fault';
                resultDiv.innerHTML = `
                    <div class="status ${statusClass}">${result.Status}</div>
                    <div class="result-card">
                        <div class="metric"><span class="metric-label">Detected Fault:</span><span class="metric-value">${result.Detected_Fault}</span></div>
                        <div class="metric"><span class="metric-label">Confidence:</span><span class="metric-value">${result.Confidence}</span></div>
                        <div class="metric"><span class="metric-label">Severity:</span><span class="metric-value"><span class="severity ${result.Severity}">${result.Severity}</span></span></div>
                        <div class="metric"><span class="metric-label">Fault Code:</span><span class="metric-value">${result.Fault_Code}</span></div>
                    </div>
                    <div class="info"><strong>📝 Recommended Action:</strong><br>${result.Recommended_Action}</div>
                `;
            } catch (error) {
                resultDiv.innerHTML = '<div class="info" style="background:#f8d7da;color:#721c24;"><strong>❌ Error:</strong> Please try again</div>';
            }
        }
    </script>
</body>
</html>
"""

@app.get("/", response_class=HTMLResponse)
async def home():
    return HTMLResponse(content=HTML_PAGE)

@app.post("/api/predict")
async def predict(request: Request):
    data = await request.json()
    
    features = np.array([[
        data['temp_chilled_supply'], data['temp_chilled_return'],
        data['temp_cond_supply'], data['temp_cond_return'],
        data['pressure_evap'], data['pressure_cond'],
        data['power_compressor'], data['flow_refrigerant'],
        data['temp_oil'], data['superheat'], data['subcooling'],
        data['approach_evap'], data['approach_cond'],
        data['capacity_cooling'], data['cop']
    ]])
    
    features_scaled = model.scaler.transform(features)
    features_selected = features_scaled[:, model.top_features_idx]
    
    with torch.no_grad():
        features_tensor = torch.FloatTensor(features_selected)
        features_nn = model.nn_model(features_tensor).numpy()
    
    prediction = model.svm_model.predict(features_nn)[0]
    probabilities = model.svm_model.predict_proba(features_nn)[0]
    
    fault_name = fault_types[prediction]
    confidence = probabilities[prediction] * 100
    is_fault = prediction != 0
    
    recommendations = {
        "Reduced Evaporator Water Flow": "Check water pump, strainers, and flow control valves. Inspect for blockages.",
        "Reduced Condenser Water Flow": "Inspect condenser water pump, clean strainers, check cooling tower operation.",
        "Refrigerant Leakage": "Perform leak detection test, check refrigerant charge levels, inspect joints.",
        "Refrigerant Overcharge": "Remove excess refrigerant, check charging procedures, inspect for non-condensables.",
        "Excess Oil in Compressor": "Check oil return system, inspect oil separators, schedule oil change.",
        "Non-condensables in Refrigerant": "Purge non-condensables, check vacuum procedures, inspect for air ingress.",
        "Compressor Valve Leakage": "Inspect compressor valves, check for worn components, schedule maintenance.",
        "Condenser Fouling": "Clean condenser tubes, inspect water treatment system, check for scaling."
    }
    
    severity = "HIGH" if confidence > 80 else "MEDIUM" if confidence > 60 else "LOW"
    
    return {
        "Status": "⚠️ FAULT DETECTED" if is_fault else "✅ NORMAL OPERATION",
        "Detected_Fault": fault_name,
        "Confidence": f"{confidence:.1f}%",
        "Severity": severity if is_fault else "NONE",
        "Recommended_Action": recommendations.get(fault_name, "No action needed") if is_fault else "System operating normally",
        "Fault_Code": f"F{prediction}" if is_fault else "NORMAL"
    }

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