baseline/policy.py

import os, sys, json, re
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from openai import OpenAI
from models import CityObservation, ContainmentAction


def get_client() -> OpenAI:
    return OpenAI(
        api_key  = os.environ.get("API_KEY") or os.environ.get("HF_TOKEN", ""),
        base_url = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1"),
    )


def build_prompt(obs: CityObservation) -> str:
    num_districts = len(obs.districts)
    has_data_lag  = obs.data_lag_days > 0

    sorted_districts = sorted(
        obs.districts,
        key=lambda d: d.reported_infection_rate,
        reverse=True
    )

    lines = [
        "You are an epidemic response coordinator making life-or-death resource decisions.",
        "Your goal: keep infection BELOW 0.40 in all districts and prevent hospital collapse.",
        "",
        f"Step {obs.current_step}/{obs.max_steps} | Resources remaining: {obs.available_resources}",
        "ℹ️  Resources replenish by 1 each step, capped at starting pool. Spend carefully.",
        "",
        "Districts (sorted by CURRENT infection, highest threat first):",
    ]

    for d in sorted_districts:
        if d.reported_infection_rate > 0.4:
            status = "🔴 CRITICAL"
        elif d.reported_infection_rate > 0.2:
            if d.growth_rate_hint > 0.06:
                status = "🟡 WARNING→CRITICAL SOON"
            else:
                status = "🟡 WARNING"
        else:
            status = "🟢 SAFE"

        if d.hospital_capacity_remaining < 0.3:
            hosp_status = "⚠️ HOSPITAL DANGER"
        elif d.hospital_capacity_remaining < 0.6:
            hosp_status = "hospital LOW"
        else:
            hosp_status = "hospital OK"

        if has_data_lag:
            # pre-compute estimated current infection so the LLM doesn't have to do the math
            estimated = round(min(1.0, d.reported_infection_rate + 3 * d.growth_rate_hint), 2)
            if estimated > 0.4:
                est_status = "🔴 EST.CRITICAL"
            elif estimated > 0.2:
                est_status = "🟡 EST.WARNING"
            else:
                est_status = "🟢 EST.SAFE"
            lines.append(
                f"  D{d.district_id}: reported={d.reported_infection_rate:.2f} [3 days old] "
                f"growth={d.growth_rate_hint:.2f} → ESTIMATED NOW={estimated:.2f} {est_status} "
                f"{hosp_status}({d.hospital_capacity_remaining:.2f})"
            )
        else:
            lines.append(
                f"  D{d.district_id}: {status} infection={d.reported_infection_rate:.2f} "
                f"growth={d.growth_rate_hint:.2f} {hosp_status}({d.hospital_capacity_remaining:.2f})"
            )

    lines += [""]

    if not has_data_lag:
        lines += [
            "HOW ACTIONS WORK:",
            "  - 'allocate': costs 1 resource. REDUCES existing infection by 5% AND slows spread.",
            "    Infections naturally recover 1%/day but spread (3-8%/day) dominates without action.",
            "    You need SUSTAINED allocation (multiple steps) to drive a district below safe level.",
            "  - 'restrict': FREE. Slows future spread only. Does NOT reduce existing infection.",
            "    Use only when you have no resources OR for districts already below 0.20.",
            "  - 'test': wastes 1 resource. Data is already real-time. NEVER use this.",
            "",
            "DECISION RULES — follow this priority order every step:",
            "1. HOSPITAL EMERGENCY: If ANY hospital < 0.30 capacity → allocate on that district NOW.",
            "2. TRIAGE: Look at ALL districts. Find the one with the HIGHEST infection rate right now.",
            "   That is your target this step. Do not stick to the same district if another is worse.",
            "3. CRITICAL DISTRICT (above 0.40 and have resources): allocate on the highest.",
            "4. WARNING DISTRICT (0.20-0.40) with growth > 0.06 AND resources available:",
            "   allocate on it NOW to prevent it from becoming CRITICAL next step.",
            "5. If resources = 0: restrict on the highest infected district.",
            "6. NEVER restrict a district below 0.20 — you will be penalised.",
            "7. NEVER use 'test' — it wastes a resource you cannot afford.",
            "",
            "KEY INSIGHT: Infection spreads 3-8% per day. A WARNING district at 0.38 with",
            "growth=0.07 will be CRITICAL next step. Act before it escalates, not after.",
        ]
    else:
        lines += [
            "HOW ACTIONS WORK:",
            "  - 'allocate': costs 1 resource. Reduces infection AND slows spread.",
            "  - 'restrict': FREE. Slows future spread only.",
            "  - Data is 3 DAYS OLD — you cannot see current true infection rates.",
            "    Use growth_hint to estimate which districts are getting worse fastest.",
            "",
            "DECISION RULES for delayed-information scenario:",
            "1. HOSPITAL EMERGENCY: If ANY hospital < 0.30 → allocate on that district NOW.",
            "2. Look at ESTIMATED NOW column — this is your best estimate of current infection.",
            "   The estimation accounts for 3 days of growth automatically.",
            "3. Allocate on the district with HIGHEST 'ESTIMATED NOW' value.",
            "4. If resources = 0: restrict on the district with highest ESTIMATED NOW.",
            "5. NEVER use 'test' — data lag is structural, testing does not help.",
        ]

    lines += [
        "",
        "Think briefly (1 sentence): Which district is most dangerous RIGHT NOW and why?",
        "Then give your JSON decision.",
        "",
        "Example response:",
        'District 0 is critical at 0.65 and growing fastest.',
        '{"action_type": "allocate", "district_id": 0}',
        "",
        "Your response:",
    ]

    return "\n".join(lines)


def call_llm(prompt: str, client: OpenAI) -> str:
    response = client.chat.completions.create(
        model    = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct"),
        messages = [
            {
                "role":    "system",
                "content": (
                    "You are an epidemic response AI. "
                    "First write one sentence of reasoning, then a JSON action on the next line. "
                    "JSON must be valid and contain action_type and district_id."
                )
            },
            {"role": "user", "content": prompt}
        ],
        max_tokens  = 60,
        temperature = 0.1,
    )
    return (response.choices[0].message.content or "").strip()


def parse_action(response: str, num_districts: int) -> ContainmentAction:
    valid_types = {"test", "restrict", "allocate"}
    try:
        cleaned = re.sub(r"```(?:json)?|```", "", response).strip()
        match   = re.search(r"\{.*?\}", cleaned, re.DOTALL)
        if match:
            cleaned = match.group()
        data        = json.loads(cleaned)
        action_type = str(data.get("action_type", "allocate")).lower().strip()
        district_id = int(data.get("district_id", 0))
        if action_type not in valid_types:
            action_type = "allocate"
        district_id = max(0, min(district_id, num_districts - 1))
        return ContainmentAction(action_type=action_type, district_id=district_id)
    except Exception:
        return ContainmentAction(action_type="allocate", district_id=0)


def get_action(obs: CityObservation, client: OpenAI) -> ContainmentAction:
    prompt   = build_prompt(obs)
    response = call_llm(prompt, client)
    return parse_action(response, len(obs.districts))


def build_prompt_with_memory(obs: CityObservation, memory) -> str:
    """
    Builds the LLM prompt augmented with relevant past decisions from episodic memory.
    Injects memory block just before the 'Your response:' line so the model sees
    prior high-reward decisions as concrete examples before making its choice.
    """
    base         = build_prompt(obs)
    memory_block = memory.retrieve(obs)
    if not memory_block:
        return base
    injection = "\n" + memory_block + "\nApply these lessons to your current decision.\n"
    return base.replace("Your response:", injection + "Your response:")