KSP_project/scripts/fill_location_state.py

"""
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()