demeter

  1import base64
  2import os
  3import uvicorn
  4import numpy as np
  5import torch
  6from io import BytesIO
  7from collections import deque
  8from fastapi import FastAPI
  9from pydantic import BaseModel
 10from typing import List
 11from PIL import Image
 12from dotenv import load_dotenv
 13from pymongo import MongoClient, ReturnDocument
 14from pymongo.errors import ConfigurationError
 15from datetime import datetime
 16
 17# Load env from root
 18current_dir = os.path.dirname(os.path.abspath(__file__))
 19env_path = os.path.join(current_dir, "..", ".env")
 20load_dotenv(env_path)
 21
 22MONGO_URI = os.environ.get("MONGODB_URI")
 23mongo_client = MongoClient(MONGO_URI)
 24try:
 25    db = mongo_client.get_default_database()
 26except ConfigurationError:
 27    db = mongo_client["test"]
 28crops_collection = db["cropstates"]
 29sim_state_collection = db["simulator_state"]
 30
 31print(f"[DEBUG] MongoDB connected to DB: '{db.name}'")
 32print(f"[DEBUG] MONGO_URI = {MONGO_URI[:40] if MONGO_URI else 'NOT SET'}...")
 33
 34MODEL_PATH = "models/PPO/lettuce_brain_v1.zip"
 35HISTORY_LEN = 20
 36
 37CROP_LIFECYCLES = {
 38    "lettuce": {
 39        "stages": [
 40            {"name": "seedling", "end_hour": 168},
 41            {"name": "vegetative", "end_hour": 504},
 42            {"name": "harvest", "end_hour": 999999},
 43        ]
 44    },
 45    "tomato": {
 46        "stages": [
 47            {"name": "seedling", "end_hour": 336},
 48            {"name": "vegetative", "end_hour": 1008},
 49            {"name": "flowering", "end_hour": 1680},
 50            {"name": "fruiting", "end_hour": 999999},
 51        ]
 52    },
 53    "basil": {
 54        "stages": [
 55            {"name": "seedling", "end_hour": 168},
 56            {"name": "vegetative", "end_hour": 672},
 57            {"name": "harvest", "end_hour": 999999},
 58        ]
 59    },
 60    "strawberry": {
 61        "stages": [
 62            {"name": "seedling", "end_hour": 336},
 63            {"name": "vegetative", "end_hour": 1008},
 64            {"name": "flowering", "end_hour": 1512},
 65            {"name": "fruiting", "end_hour": 999999},
 66        ]
 67    },
 68}
 69
 70
 71class FarmAction(BaseModel):
 72    acid_dosage_ml: float = 0.0
 73    base_dosage_ml: float = 0.0
 74    nutrient_dosage_ml: float = 0.0
 75    fan_speed_pct: float = 0.0
 76    water_refill_l: float = 0.0
 77    debug_force_ph: float | None = None
 78
 79
 80class BatchActionRequest(BaseModel):
 81    crop_id: str
 82    action: FarmAction
 83
 84
 85class ResidualPhysicsNet(torch.nn.Module):
 86    def __init__(self, state_dim, action_dim):
 87        super().__init__()
 88        self.net = torch.nn.Sequential(
 89            torch.nn.Linear(state_dim + action_dim, 64),
 90            torch.nn.Tanh(),
 91            torch.nn.Linear(64, 64),
 92            torch.nn.ReLU(),
 93            torch.nn.Linear(64, state_dim),
 94        )
 95
 96    def forward(self, state, action):
 97        x = torch.cat([state, action], dim=-1)
 98        return self.net(x)
 99
100
101class DigitalTwin:
102    def __init__(self, crop_id: str, initial_state: list):
103        self.crop_id = crop_id
104        self.state = np.array(initial_state, dtype=np.float32)
105        self.tank_volume = 100.0
106        self.plant_health = 100.0
107        self.residual_model = ResidualPhysicsNet(7, 4)
108        
109        # Try to load trained weights
110        trained_weights_path = os.path.join(current_dir, "residual_physics.pt")
111        if os.path.exists(trained_weights_path):
112            try:
113                checkpoint = torch.load(trained_weights_path, map_location="cpu")
114                self.residual_model.load_state_dict(checkpoint["model_state_dict"])
115                print(f"[DEBUG] Loaded trained ResidualPhysicsNet weights from {trained_weights_path}")
116            except Exception as e:
117                print(f"[WARNING] Failed to load trained weights: {e}. Using random initialization.")
118        else:
119            print(f"[DEBUG] Trained weights not found at {trained_weights_path}. Using random initialization.")
120
121        self.history = {
122            "ph": deque([float(self.state[0])] * 5, maxlen=HISTORY_LEN),
123            "ec": deque([float(self.state[1])] * 5, maxlen=HISTORY_LEN),
124            "water_temp": deque([float(self.state[2])] * 5, maxlen=HISTORY_LEN),
125            "air_temp": deque([float(self.state[3])] * 5, maxlen=HISTORY_LEN),
126            "humidity": deque([float(self.state[4])] * 5, maxlen=HISTORY_LEN),
127            "co2": deque([400.0] * 5, maxlen=HISTORY_LEN),
128            "light_intensity": deque([0.0] * 5, maxlen=HISTORY_LEN),
129            "vpd": deque([float(self.state[5])] * 5, maxlen=HISTORY_LEN),
130        }
131
132    def _calculate_vpd(self, temp, hum):
133        es = 0.61078 * np.exp((17.27 * temp) / (temp + 237.3))
134        ea = es * (hum / 100.0)
135        return max(0.0, es - ea)
136
137    def step(self, action: FarmAction = None):
138        if action is None:
139            action = FarmAction()
140
141        u = np.array(
142            [
143                action.acid_dosage_ml / 10.0,
144                action.base_dosage_ml / 10.0,
145                action.nutrient_dosage_ml / 20.0,
146                action.fan_speed_pct / 100.0,
147            ],
148            dtype=np.float32,
149        )
150
151        ph, ec, wt, at, hum, vpd, bio = self.state
152
153        d_ph = (u[1] * 0.5) - (u[0] * 0.5) + (0.001 * bio)
154        if action.debug_force_ph:
155            self.state[0] = action.debug_force_ph
156
157        uptake = 0.02 * bio * vpd
158        d_ec = (u[2] * 0.2) - (uptake / self.tank_volume)
159        d_at = 0.1 - (u[3] * 1.5)
160        d_hum = 1.0 - (u[3] * 5.0)
161
162        stress = abs(vpd - 1.0)
163        growth = 0.1 * bio * (1.0 - min(stress, 1.0))
164
165        physics_delta = np.array(
166            [d_ph, d_ec, 0, d_at, d_hum, 0, growth], dtype=np.float32
167        )
168
169        with torch.no_grad():
170            nn_delta = self.residual_model(
171                torch.tensor(self.state), torch.tensor(u)
172            ).numpy()
173
174        self.state += physics_delta + (nn_delta * 0.05)
175        self.state[3] = np.clip(self.state[3], 0, 50)
176        self.state[4] = np.clip(self.state[4], 0, 100)
177        self.state[5] = self._calculate_vpd(self.state[3], self.state[4])
178        self.state[6] = max(0.1, self.state[6])
179
180        ph_score = max(0, 1.0 - abs(self.state[0] - 6.0))
181        vpd_score = max(0, 1.0 - abs(self.state[5] - 1.0))
182        self.plant_health = max(
183            0.0, min(100.0, (float(ph_score) + float(vpd_score)) * 50.0)
184        )
185
186        self._update_history()
187        return self._generate_image()
188
189    def _update_history(self):
190        s = self.state
191        self.history["ph"].append(float(s[0]))
192        self.history["ec"].append(float(s[1]))
193        self.history["water_temp"].append(float(s[2]))
194        self.history["air_temp"].append(float(s[3]))
195        self.history["humidity"].append(float(s[4]))
196        self.history["vpd"].append(float(s[5]))
197
198    def _generate_image(self):
199        bucket = max(0, min(90, int(self.plant_health // 10) * 10))
200        filename = f"{bucket}.png"
201        if os.path.exists(filename):
202            return Image.open(filename)
203        return Image.new("RGB", (512, 512), (50, 50, 50))
204
205
206app = FastAPI()
207simulators = {}
208
209
210def sync_simulators_from_db():
211    db_crops = crops_collection.find({})
212    for crop in db_crops:
213        cid = crop.get("crop_id")
214        if not cid:
215            continue
216
217        if cid not in simulators:
218            sensors = crop.get("sensors", {})
219
220            ph_val = sensors.get("pH", [6.0])
221            ec_val = sensors.get("EC", [1.5])
222            temp_val = sensors.get("temp", [24.0])
223            hum_val = sensors.get("humidity", [60.0])
224
225            state = [
226                ph_val[-1] if isinstance(ph_val, list) else ph_val,
227                ec_val[-1] if isinstance(ec_val, list) else ec_val,
228                20.0,
229                temp_val[-1] if isinstance(temp_val, list) else temp_val,
230                hum_val[-1] if isinstance(hum_val, list) else hum_val,
231                1.0,
232                10.0,
233            ]
234            simulators[cid] = DigitalTwin(cid, state)
235
236
237@app.get("/simulation/state")
238async def get_all_states():
239    clock = sim_state_collection.find_one_and_update(
240        {"_id": "global_clock"},
241        {"$inc": {"tick_hours": 1}},
242        upsert=True,
243        return_document=ReturnDocument.AFTER,
244    )
245    current_tick = clock["tick_hours"]
246
247    crops_collection.update_many({}, {"$inc": {"simulated_age_hours": 1}})
248
249    sync_simulators_from_db()
250
251    db_crops = list(crops_collection.find({}))
252    response = []
253
254    for crop in db_crops:
255        cid = crop.get("crop_id")
256        if not cid:
257            continue
258        if cid not in simulators:
259            print(f"⚠️ Crop {cid} not in simulators after sync — skipping")
260            continue
261
262        crop_type = crop.get("crop", "lettuce").lower()
263        age_hours = crop.get("sequence_number", 0) * crop.get("cycle_duration_hours", 1)
264
265        print(f"Simulating Crop ID: {cid} | Type: {crop_type} | Age (hrs): {age_hours}")
266        cycle_duration = crop.get("cycle_duration_hours", 1)
267
268        new_stage = crop.get("stage", "seedling")
269        if crop_type in CROP_LIFECYCLES:
270            for stage_info in CROP_LIFECYCLES[crop_type]["stages"]:
271                if age_hours <= stage_info["end_hour"]:
272                    new_stage = stage_info["name"]
273                    break
274
275        if new_stage != crop.get("stage"):
276            crops_collection.update_one(
277                {"crop_id": cid}, {"$set": {"stage": new_stage}}
278            )
279
280        print(
281            f" current_tick: {current_tick} | crop_id: {cid} | age_hours: {age_hours} | stage: {new_stage} | cycle_duration: {cycle_duration}"
282        )
283
284        if current_tick % cycle_duration != 0 and current_tick != 1:
285            continue
286
287        sim = simulators[cid]
288        pil_img = sim._generate_image()
289        buf = BytesIO()
290        pil_img.save(buf, format="PNG")
291        img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
292
293        response.append(
294            {
295                "crop_id": cid,
296                "sensor_window": {k: list(v) for k, v in sim.history.items()},
297                "metadata": {
298                    "crop": crop_type,
299                    "stage": new_stage,
300                    "health": round(float(sim.plant_health), 1),
301                    "biomass_est": round(float(sim.state[6]), 2),
302                    "age_hours": age_hours,
303                    "global_tick": current_tick,
304                },
305                "image": img_b64,
306            }
307        )
308
309    return response
310
311
312@app.post("/simulation/action")
313async def take_batch_actions(payload: List[BatchActionRequest]):
314    sync_simulators_from_db()
315
316    results = []
317    for req in payload:
318        cid = req.crop_id
319        if cid not in simulators:
320            results.append(
321                {"crop_id": cid, "status": "error", "message": "Crop not found"}
322            )
323            continue
324
325        sim = simulators[cid]
326        sim.step(req.action)
327
328        push_payload = {
329            "sensors.pH": {"$each": [float(sim.state[0])], "$slice": -5},
330            "sensors.EC": {"$each": [float(sim.state[1])], "$slice": -5},
331            "sensors.temp": {"$each": [float(sim.state[3])], "$slice": -5},
332            "sensors.humidity": {"$each": [float(sim.state[4])], "$slice": -5},
333        }
334
335        set_payload = {
336            "action_taken": req.action.dict(),
337            "last_updated": datetime.utcnow(),
338        }
339
340        crops_collection.update_one(
341            {"crop_id": cid},
342            {
343                "$push": push_payload,
344                "$set": set_payload,
345                "$inc": {"sequence_number": 1},
346            },
347        )
348
349        results.append(
350            {
351                "crop_id": cid,
352                "status": "success",
353                "new_state": {
354                    "pH": float(sim.state[0]),
355                    "EC": float(sim.state[1]),
356                    "temp": float(sim.state[3]),
357                    "humidity": float(sim.state[4]),
358                },
359            }
360        )
361
362    return {"updated_crops": results}
363
364
365if __name__ == "__main__":
366    port = int(os.environ.get("SIMULATOR_PORT", 8001))
367    uvicorn.run(app, host="0.0.0.0", port=port)