""" Correctly fill location and state_label in state nodes based on semantic context. """ import os, psycopg, re from datetime import datetime, timezone from dotenv import load_dotenv from pathlib import Path load_dotenv(Path(__file__).resolve().parents[1] / ".env") conn = psycopg.connect( host=os.getenv("PGHOST"), port=os.getenv("PGPORT"), dbname=os.getenv("PGDATABASE"), user=os.getenv("PGUSER"), password=os.getenv("PGPASSWORD")) cur = conn.cursor() NOW = datetime.now(timezone.utc) # ── Location inference ── # Map body name → orbit location, surface location (if applicable) BODY_LOC = { "mercury": ("Mercury Orbit", "Mercury Surface"), "venus": ("Venus Orbit", None), # no crewed surface ops "earth": ("LEO", "Earth Surface"), "moon": ("Lunar Orbit", "Lunar Surface"), "lunar": ("Lunar Orbit", "Lunar Surface"), "mars": ("Mars Orbit", "Mars Surface"), "phobos": ("Phobos Orbit", None), "deimos": ("Deimos Orbit", None), "ceres": ("Ceres Orbit", "Ceres Surface"), "vesta": ("Vesta Orbit", "Vesta Surface"), "jupiter": ("Jupiter System", None), "io": ("Io Orbit", "Io Surface"), "europa": ("Europa Orbit", "Europa Surface"), "ganymede": ("Ganymede Orbit", "Ganymede Surface"), "callisto": ("Callisto Orbit", "Callisto Surface"), "saturn": ("Saturn System", None), "titan": ("Titan Orbit", "Titan Surface"), "iapetus": ("Iapetus Orbit", "Iapetus Surface"), "enceladus": ("Enceladus Orbit", None), "dione": ("Dione Orbit", None), "rhea": ("Rhea Orbit", None), "tethys": ("Tethys Orbit", None), "mimas": ("Mimas Orbit", None), "uranus": ("Uranus System", None), "miranda": ("Miranda Orbit", None), "ariel": ("Ariel Orbit", None), "umbriel": ("Umbriel Orbit", None), "titania": ("Titania Orbit", None), "oberon": ("Oberon Orbit", None), "neptune": ("Neptune System", None), "triton": ("Triton Orbit", "Triton Surface"), "pluto": ("Pluto Orbit", "Pluto Surface"), "charon": ("Charon Orbit", None), "sun": ("Solar Orbit", None), "solar": ("Solar Orbit", None), "star port": ("LEO", None), "star ring": ("LEO", None), "leo": ("LEO", None), "geo": ("GEO", None), } # Keywords that indicate surface ops vs orbital SURFACE_WORDS = {"surface", "landing", "descend", "edl", "eva", "base", "outpost", "crewed landing", "ground", "walk", "construction"} ORBIT_WORDS = {"orbit", "orbital", "arrival", "departure", "soi", "flyby", "survey", "parking", "dock", "rendezvous", "approach", "braking burn", "course correction", "insertion"} TRANSFER_WORDS = {"transfer", "transit", "outbound", "inbound", "return", "departure", "en route", "heading", "burn", "leave", "left", "sail"} # State label inference STATE_KEYWORDS = [ (["construction", "assembly", "assembl", "fitting", "install", "module delivery", "first module", "construction complete", "assembly halt", "assembly complete", "resume construction"], "Construction"), (["trial run", "testing", "test mission", "commissioning", "shakedown", "deep space test"], "Testing"), (["maintenance", "upgrade", "refit", "repair", "servicing", "parked to", "upgrade to"], "Maintenance"), (["departure", "depart", "leave", "outbound", "earth departure", "mars departure", "europa departure", "left", "en route to", "heading to"], "start-transit"), (["arrival", "arrive", "arrived", "return", "inbound", "braking burn", "course correction", "soi", "crew arrival", "crew return", "earth arrival", "mars arrival", "arriving"], "transit-complete"), (["exploration", "expedition", "surface operation", "surface expedition", "descend", "ascend", "eva operation", "orbital survey", "probe deployment", "atmosphere survey", "crewed landing", "lander", "surface ops", "rover", "observation"], "Exploration"), (["docked", "docking", "dock", "berthing", "mooring"], "Docked"), (["autonomous standby", "standby"], "Autonomous Standby"), (["crew handoff", "handoff", "handover"], "Crew Handoff"), ] def find_body(text): """Find the primary celestial body mentioned in text (whole-word match only).""" if not text: return None text_lower = text.lower() # Sort by key length (longest first) to match "star port" before "star" for body in sorted(BODY_LOC.keys(), key=len, reverse=True): pattern = r'\b' + re.escape(body) + r'\b' if re.search(pattern, text_lower): return body return None def infer_location(node_title, node_detail, entry_title, mission_label, asset_name, asset_type, home_region): """Infer the correct location for a state node.""" node_text = f"{node_title or ''} {node_detail or ''}".lower() full_text = f"{entry_title or ''} {mission_label or ''}".lower() # Priority 1: body mentioned in node title/detail body = find_body(node_text) if body and body in BODY_LOC: orbit_loc, surf_loc = BODY_LOC[body] has_surface = any(w in node_text for w in SURFACE_WORDS) has_orbit = any(w in node_text for w in ORBIT_WORDS) if surf_loc and has_surface and not has_orbit: return surf_loc return orbit_loc # Priority 2: body in entry title / mission label body = find_body(full_text) if body and body in BODY_LOC: orbit_loc, surf_loc = BODY_LOC[body] has_surface = any(w in node_text for w in SURFACE_WORDS) has_orbit = any(w in node_text for w in ORBIT_WORDS) if surf_loc and has_surface: return surf_loc return orbit_loc # Priority 3: transfer keywords all_text = f"{node_text} {full_text}" if any(w in all_text for w in TRANSFER_WORDS): if body: return BODY_LOC.get(body, (None, None))[0] or "Transfer" return "Transfer" # Priority 4: home region or asset type default if home_region: hr = str(home_region).lower() for b in BODY_LOC: if b in hr: return BODY_LOC[b][0] return "LEO" def infer_state(node_title, node_detail, entry_title, mission_label): """Infer the correct state_label for a state node.""" context = f"{node_title or ''} {node_detail or ''} {entry_title or ''} {mission_label or ''}".lower() for keywords, label in STATE_KEYWORDS: for kw in keywords: if kw in context: return label return None # ── Main ── cur.execute(""" SELECT sn.id::text, sn.title, sn.detail, sn.at_time, le.title as entry_title, le.mission_label, a.name as asset_name, a.asset_type, a.home_region FROM asset_state_nodes sn JOIN asset_log_entries le ON sn.entry_id = le.id JOIN assets a ON le.asset_id = a.id ORDER BY a.name, sn.at_time """) rows = cur.fetchall() updated = 0 for r in rows: nid, ntitle, ndetail, atime, entry_title, mission, asset_name, asset_type, home_region = r new_loc = infer_location(ntitle, ndetail, entry_title, mission, asset_name, asset_type, home_region) new_st = infer_state(ntitle, ndetail, entry_title, mission) # Only update if different from current cur.execute("SELECT location, state_label FROM asset_state_nodes WHERE id = %s", (nid,)) cur_loc, cur_st = cur.fetchone() if (new_loc and new_loc != cur_loc) or (new_st and new_st != cur_st): cur.execute(""" UPDATE asset_state_nodes SET location = COALESCE(%s, location), state_label = COALESCE(%s, state_label), updated_at = %s WHERE id = %s """, (new_loc, new_st, NOW, nid)) updated += 1 print(f" {asset_name[:20]:20s} | {str(ntitle)[:45]:45s} | {cur_loc or 'NULL':20s} → {new_loc or 'NULL':20s} | {cur_st or 'NULL':15s} → {new_st or 'NULL':15s}") conn.commit() cur.execute(""" SELECT COUNT(*), COUNT(CASE WHEN location IS NOT NULL AND location!='' THEN 1 END), COUNT(CASE WHEN state_label IS NOT NULL AND state_label!='' THEN 1 END) FROM asset_state_nodes """) r = cur.fetchone() print(f"\nFinal: {r[0]} total, {r[1]} with location, {r[2]} with state_label") print(f"Updated: {updated} nodes") conn.close()