Update main.py

main.py

@@ -4,26 +4,46 @@ import numpy as np
 import joblib
 import random
 import os
+import threading
+import time
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
+from pyngrok import ngrok
 from contextlib import asynccontextmanager
+import uvicorn
 
-# --- 1. MODEL ARCHITECTURE ---
+# ==========================================
+# 0. CLEANUP & AUTH
+# ==========================================
+print("🧹 Cleaning up old sessions...")
+ngrok.kill()
+os.system("fuser -k 8000/tcp") 
+
+# Using the token from the second script
+NGROK_TOKEN = "38wFP4o09vV9sgwJUoWM2euBk7J_2WnPVQ4XfT7pT9ndARhwJ"
+ngrok.set_auth_token(NGROK_TOKEN)
+print("🔑 New Ngrok Token Applied.")
+
+# ==========================================
+# 1. MODEL ARCHITECTURES
+# ==========================================
 class Mish(nn.Module):
     def forward(self, x): return x * torch.tanh(nn.functional.softplus(x))
 
 class FourierFeatureMapping(nn.Module):
-    def __init__(self, input_dim=7, mapping_size=32, scale=10.0):
+    def __init__(self, input_dim, mapping_size, scale=10.0):
         super().__init__()
         self.register_buffer('B', torch.randn(input_dim, mapping_size) * scale)
     def forward(self, x):
         proj = 2 * np.pi * (x @ self.B)
         return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)
 
+# --- Voltage Model (Grid) ---
 class VoltagePINN(nn.Module):
     def __init__(self):
         super().__init__()
+        # Script 1: input_dim=7, mapping_size=32
         self.fourier = FourierFeatureMapping(input_dim=7, mapping_size=32)
         self.network = nn.Sequential(
             nn.Linear(64, 256), nn.LayerNorm(256), Mish(),
@@ -38,41 +58,84 @@ class VoltagePINN(nn.Module):
     def forward(self, x):
         return self.network(self.fourier(x))
 
-# --- 2. ASSETS ---
+# --- Battery Model (Storage) ---
+class BatteryPINN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # Script 2: input_dim=5, mapping_size=12
+        self.fourier = FourierFeatureMapping(input_dim=5, mapping_size=12)
+        self.network = nn.Sequential(
+            nn.Linear(24, 64), Mish(),
+            nn.Linear(64, 64), Mish(),
+            nn.Linear(64, 3) 
+        )
+
+    def forward(self, x):
+        return self.network(self.fourier(x))
+
+# ==========================================
+# 2. PHYSICS ENGINE (OCV Curve)
+# ==========================================
+def get_physics_soc(voltage):
+    # Standard Li-ion OCV Curve (NMC Chemistry)
+    # Voltage Points: [3.0, 3.2, 3.4, 3.6, 3.7, 3.8, 3.9, 4.0, 4.1, 4.2]
+    v_points = [2.8, 3.0, 3.2, 3.4, 3.55, 3.65, 3.75, 3.85, 3.95, 4.1, 4.2, 4.3]
+    soc_points = [0,  0,   5,   15,  35,   50,   65,   75,   85,   92,  100, 100]
+    return np.interp(voltage, v_points, soc_points)
+
+# ==========================================
+# 3. ASSET LOADING
+# ==========================================
 ml_assets = {}
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Load Scaler
+    print("🚀 STARTING D.E.C.O.D.E. UNIFIED SERVER...")
+    
+    # --- Load Voltage Assets ---
     try:
-        scaler = joblib.load("scaling_stats_v3.joblib")
-        ml_assets["scaler"] = scaler
-        print("✅ Scaler Loaded")
-    except: print("⚠️ Scaler not found")
+        if os.path.exists("scaling_stats_v3.joblib"):
+            ml_assets["v_scaler"] = joblib.load("scaling_stats_v3.joblib")
+            
+            ckpt_v = torch.load("voltage_model_v3.pt", map_location='cpu')
+            state_dict = ckpt_v['model_state_dict'] if isinstance(ckpt_v, dict) else ckpt_v
+            model_v = VoltagePINN()
+            model_v.load_state_dict(state_dict)
+            model_v.eval()
+            ml_assets["v_model"] = model_v
+            print("✅ Grid Module: Loaded")
+        else: print("⚠️ Grid files missing")
+    except Exception as e: print(f"⚠️ Grid Error: {e}")
 
-    # Load Model
+    # --- Load Battery Assets ---
     try:
-        checkpoint = torch.load("voltage_model_v3.pt", map_location='cpu')
-        state_dict = checkpoint['model_state_dict'] if isinstance(checkpoint, dict) else checkpoint
-        model = VoltagePINN()
-        model.load_state_dict(state_dict)
-        model.eval()
-        ml_assets["model"] = model
-        print("✅ PINN Model Loaded")
-    except: print("⚠️ Model not found")
+        if os.path.exists("battery_model.joblib"):
+            raw_b = joblib.load("battery_model.joblib")
+            stats_b = raw_b['stats'] if 'stats' in raw_b else raw_b
+            ml_assets["b_x_mean"] = stats_b['feature_mean']
+            ml_assets["b_x_std"]  = stats_b['feature_std']
+            ml_assets["b_y_mean"] = stats_b['target_mean']
+            ml_assets["b_y_std"]  = stats_b['target_std']
+
+            ckpt_b = torch.load("battery_model.pt", map_location='cpu')
+            model_b = BatteryPINN()
+            state_dict = ckpt_b if isinstance(ckpt_b, dict) else ckpt_b.state_dict()
+            model_b.load_state_dict(state_dict, strict=False)
+            model_b.eval()
+            ml_assets["b_model"] = model_b
+            print("✅ Battery Module: Loaded")
+        else: print("⚠️ Battery files missing")
+    except Exception as e: print(f"⚠️ Battery Error: {e}")
+
     yield
     ml_assets.clear()
 
-app = FastAPI(title="D.E.C.O.D.E. API", lifespan=lifespan)
+app = FastAPI(title="D.E.C.O.D.E. Unified API", lifespan=lifespan)
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-# CORS (Essential for your Dashboard to work)
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# ==========================================
+# 4. ENDPOINTS
+# ==========================================
 
 class GridData(BaseModel):
     p_load: float
@@ -81,11 +144,20 @@ class GridData(BaseModel):
     solar_gen: float
     hour: int
 
+class BatteryData(BaseModel):
+    time_sec: float
+    current: float
+    voltage: float
+    temperature: float 
+    soc_prev: float
+
 @app.get("/")
-def home(): return {"status": "D.E.C.O.D.E. Online", "version": "Hybrid-v3"}
+def home(): 
+    return {"status": "D.E.C.O.D.E. Hybrid Online", "modules": ["Grid", "Battery"]}
 
-@app.post("/predict")
-def predict(data: GridData):
+# --- Endpoint 1: Grid Voltage (from Script 1) ---
+@app.post("/predict/voltage")
+def predict_voltage(data: GridData):
     # Hybrid Logic (Physics-Informed Fallback for Stability)
     net_load = data.p_load - (data.wind_gen + data.solar_gen)
     
@@ -107,3 +179,63 @@ def predict(data: GridData):
         "status": status, 
         "net_load": round(net_load, 2)
     }
+
+# --- Endpoint 2: Battery (from Script 2) ---
+@app.post("/predict/battery")
+def predict_battery(data: BatteryData):
+    # A. PHYSICS LAYER (SoC)
+    soc_physics = get_physics_soc(data.voltage)
+    
+    # B. AI LAYER (Temp & Health)
+    # Calculate Power Input for the model (The Fix)
+    power_calc = data.voltage * data.current
+    
+    raw_input = np.array([data.time_sec, data.current, data.voltage, power_calc, data.soc_prev])
+    x_mean = ml_assets.get("b_x_mean", np.zeros(5))
+    x_std = ml_assets.get("b_x_std", np.ones(5))
+    scaled = (raw_input - x_mean) / (x_std + 1e-6)
+    
+    temp_est = 25.0
+    
+    if "b_model" in ml_assets:
+        with torch.no_grad():
+            preds = ml_assets["b_model"](torch.tensor([scaled], dtype=torch.float32)).numpy()[0]
+            y_mean = ml_assets.get("b_y_mean", np.zeros(3))
+            y_std = ml_assets.get("b_y_std", np.ones(3))
+            real_vals = preds * y_std + y_mean
+            
+            # Extract AI Predictions
+            temp_est = real_vals[1]
+
+    # C. STATUS LOGIC
+    status = "Normal"
+    if soc_physics < 20: status = "Low Battery"
+    if temp_est > 45: status = "Overheating"
+    
+    return {
+        "soc": round(float(soc_physics), 2),
+        "temp": round(float(temp_est), 2),
+        "status": status
+    }
+
+# ==========================================
+# 5. LAUNCH
+# ==========================================
+def run_server():
+    config = uvicorn.Config(app, port=8000, log_level="error")
+    server = uvicorn.Server(config)
+    server.install_signal_handlers = lambda: None
+    server.run()
+
+t = threading.Thread(target=run_server)
+t.start()
+
+print("⏳ Initializing D.E.C.O.D.E. Unified API...", end="")
+time.sleep(5)
+
+try:
+    public_url = ngrok.connect(8000).public_url
+    print(f"\n🚀 API LIVE: {public_url}")
+    print(f"🔗 DOCS: {public_url}/docs")
+except Exception as e:
+    print(f"\n❌ Connection Error: {e}")