crime_env/case_generator.py

"""Component 1 — Case Generator.

Generates randomized crime scenarios for the AI Crime Investigation World.
Criminal is randomly assigned to either Suspect_A or Suspect_B each episode.
"""

import random

from crime_env.constants import (
    SUSPECTS,
    SUSPECT_A,
    SUSPECT_B,
    SUSPECT_A_KEY,
    SUSPECT_B_KEY,
    WITNESS_1,
    WITNESS_1_KEY,
)


# ── Pools for random sampling ───────────────────────────────────────────────

CRIMES = [
    "art theft",
    "corporate fraud",
    "warehouse robbery",
    "arson",
]

LOCATIONS = [
    "east wing",
    "server room",
    "vault",
    "parking lot",
    "basement archive",
]

TIMES = [
    "8:45 PM",
    "9:15 PM",
    "10:00 PM",
]

METHODS = [
    "lock picking with custom tools",
    "insider access using stolen credentials",
    "brute-force entry through side door",
    "social engineering the night guard",
    "disabling the alarm system remotely",
    "cloning a supervisor's keycard",
    "crawling through the ventilation shaft",
    "bribing a maintenance worker for access",
    "exploiting a gap in the security patrol rotation",
    "cutting the perimeter fence near the north wall",
]

CLOTHING_OPTIONS = [
    "dark hoodie and jeans",
    "grey trench coat",
    "black leather jacket",
    "navy blue overalls",
    "security guard uniform",
    "dark windbreaker and cargo pants",
    "green army surplus jacket",
    "plain black tracksuit",
    "brown corduroy jacket and khakis",
    "dark bomber jacket",
]

MOTIVES = [
    "financial pressure",
    "personal grudge",
    "professional rivalry",
    "jealousy",
    "revenge",
    "desperation",
]

# ── Type-constrained evidence templates ─────────────────────────────────────
# Each evidence type has its own pool to prevent mismatches.

KEYCARD_LOG_TEMPLATES = [
    {
        "description": "Keycard access log shows {suspect}'s badge was used at the {location} door at {time}",
        "points_to": "{suspect}",
    },
    {
        "description": "Keycard system recorded an unauthorized swipe at {time} — badge ID unregistered",
        "points_to": "unknown",
    },
    {
        "description": "Keycard records show two swipes at the {location} within 3 minutes of each other at {time} — one belonged to {suspect}",
        "points_to": "{suspect}",
    },
    {
        "description": "Keycard system was offline for several hours due to a power glitch — no access logs available for the {location} until after {time}",
        "points_to": "unknown",
    },
]

CCTV_FOOTAGE_TEMPLATES = [
    {
        "description": "CCTV footage shows a figure resembling {suspect} entering the {location} at {time}",
        "points_to": "{suspect}",
    },
    {
        "description": "CCTV footage captured a person in dark clothing near the {location} at {time} — face obscured",
        "points_to": "unknown",
    },
    {
        "description": "CCTV camera at {location} was disabled until shortly before {time} — no footage available for the relevant window",
        "points_to": "unknown",
    },
    {
        "description": "CCTV from the adjacent hallway shows {suspect} walking briskly toward the {location} at {time}",
        "points_to": "{suspect}",
    },
]

FORENSIC_REPORT_TEMPLATES = [
    {
        "description": "Forensic report: fingerprints matching {suspect} found on the tool bag recovered at the scene",
        "points_to": "{suspect}",
    },
    {
        "description": "Forensic report: torn fabric consistent with a grey trench coat found near the {location}",
        "points_to": "unknown",
    },
    {
        "description": "Forensic report: DNA trace on a dropped glove near the {location} matches {suspect}",
        "points_to": "{suspect}",
    },
    {
        "description": "Forensic report: muddy boot prints (size 10) found at the scene — owner not yet identified",
        "points_to": "unknown",
    },
]

# ── Specific alibi pools (no overlap between fake and real) ─────────────────

FAKE_ALIBI_POOL = [
    "I was at a family dinner at my uncle's house on Oak Street from 7 PM to 11 PM",
    "I was home alone watching a movie all evening — I didn't leave the house",
    "I was volunteering at the Riverside Community Centre until about 10:30 PM",
    "I was at my cousin's birthday party across town — there were at least 20 people there",
    "I was at the public library studying until they closed at 9 PM, then I walked home",
    "I was helping a friend move furniture into their new apartment on Elm Avenue",
]

REAL_ALIBI_POOL = [
    "I was working late at the Morgan & Associates office until 10 PM — my manager Sarah Chen can confirm, and I paid for takeout at 8:50 PM with my debit card",
    "I was at O'Brien's Bar on 5th Street from 7 PM to 11 PM — I paid my tab by card at 10:45 PM and the bartender knows me by name",
    "I was at the Cineplex on Main Street watching a 7:30 PM showing — I still have my ticket stub and the cashier can verify",
    "I was at the downtown gym from 6:30 PM to 9:30 PM — I scanned my membership card on entry and exit, which is logged electronically",
    "I was attending a night class at the community college from 7 PM to 9:30 PM — the instructor marked attendance electronically",
    "I was at St. Mary's Hospital visiting my grandmother from 6 PM to 9 PM — the visitor log has my name and sign-in time",
]

BEHAVIORAL_PATTERNS = [
    "uses family alibis",
    "deflects blame to associates",
    "becomes aggressive under pressure",
    "provides overly detailed timelines",
    "claims memory loss for key periods",
    "feigns cooperation while withholding details",
]

PAST_CRIME_POOL = [
    "petty theft",
    "fraud",
    "trespassing",
    "embezzlement",
    "forgery",
    "breaking and entering",
]

LIE_TOPICS = ["location", "alibi", "knows_associate", "was_at_scene", "time"]


# ── Helpers ─────────────────────────────────────────────────────────────────


def _random_past_lies(count: int) -> list[dict]:
    """Generate a list of documented past lies."""
    lies = []
    for _ in range(count):
        topic = random.choice(LIE_TOPICS)
        lies.append(
            {
                "topic": topic,
                "claimed": _fake_claim(topic),
                "disproved_by": _disprove_source(topic),
            }
        )
    return lies


def _fake_claim(topic: str) -> str:
    claims = {
        "location": "claimed to be at home",
        "knows_associate": "denied knowing the accomplice",
        "was_at_scene": "denied being near the scene",
        "time": "claimed to be elsewhere during the incident window",
        "alibi": "said they were with a friend who later denied it",
    }
    return claims.get(topic, "made an unverifiable statement")


def _disprove_source(topic: str) -> str:
    sources = {
        "location": "GPS records showed otherwise",
        "knows_associate": "phone records confirmed contact",
        "was_at_scene": "camera logs confirmed presence",
        "time": "timestamped records contradicted the timeline",
        "alibi": "friend's testimony contradicted the claim",
    }
    return sources.get(topic, "documentary evidence")


def _build_evidence(criminal: str, location: str, time: str) -> list[dict]:
    """Build 3 pieces of physical evidence — one per type.

    Index 0 → keycard_log
    Index 1 → cctv_footage
    Index 2 → forensic_report

    Each type draws from its own template pool, ensuring descriptions
    always match the evidence type name.
    """
    other_suspect = SUSPECT_B if criminal == SUSPECT_A else SUSPECT_A

    evidence_types = [
        ("keycard_log", KEYCARD_LOG_TEMPLATES),
        ("cctv_footage", CCTV_FOOTAGE_TEMPLATES),
        ("forensic_report", FORENSIC_REPORT_TEMPLATES),
    ]

    evidence = []

    for ev_name, pool in evidence_types:
        template = random.choice(pool)
        points_to_raw = template["points_to"]

        if points_to_raw == "{suspect}":
            # Bias toward implicating the actual criminal
            suspect = random.choices(
                [criminal, other_suspect],
                weights=[0.7, 0.3],
            )[0]
        else:
            suspect = None

        desc = template["description"].format(
            suspect=suspect or "", location=location, time=time
        )
        points_to = suspect if suspect else "unknown"

        evidence.append(
            {
                "name": ev_name,
                "description": desc.strip(),
                "points_to": points_to,
                "visible_to_detective": False,
            }
        )

    return evidence


def _build_prior_history(criminal: str = SUSPECT_A) -> dict:
    """Build prior history for both suspects.

    The criminal is biased toward having a more checkered past
    (more prior crimes, lower trust), which improves realism.
    """
    history = {}
    for suspect in SUSPECTS:
        is_criminal = suspect == criminal
        # Criminal: 1-2 priors, low trust; Innocent: 0-1, higher trust
        n_crimes = random.randint(1, 2) if is_criminal else random.randint(0, 1)
        n_lies = random.randint(1, 2) if is_criminal else random.randint(0, 1)
        on_parole = random.choices([True, False], weights=[0.8, 0.2])[0] if is_criminal else random.choice([True, False])
        trust = round(random.uniform(0.0, 0.4), 2) if is_criminal else round(random.uniform(0.4, 1.0), 2)

        history[suspect] = {
            "previous_crimes": random.sample(
                PAST_CRIME_POOL, k=min(n_crimes, len(PAST_CRIME_POOL))
            ),
            "past_lies_on_record": _random_past_lies(n_lies),
            "behavioral_pattern": random.choice(BEHAVIORAL_PATTERNS),
            "on_parole": on_parole,
            "parole_condition": (
                random.choice(
                    [
                        "must not leave the city",
                        "mandatory weekly check-ins",
                        "no contact with known criminals",
                        "curfew between 9 PM and 6 AM",
                    ]
                )
                if on_parole
                else None
            ),
            "trust_score": trust,
        }
    return history


def _build_agent_knowledge(
    criminal: str, location: str, time: str
) -> dict:
    """Build what each agent privately knows.

    Args:
        criminal: Which suspect is the criminal.
        location: Crime location.
        time: Crime time.
    """
    innocent = SUSPECT_B if criminal == SUSPECT_A else SUSPECT_A

    fake_alibi = random.choice(FAKE_ALIBI_POOL)
    real_alibi = random.choice(REAL_ALIBI_POOL)

    clothing = random.choice(CLOTHING_OPTIONS)

    # Enhancement 2: 30% chance the witness got the clothing wrong
    witness_clothing = clothing
    if random.random() < 0.3:
        # Pick a different clothing item
        other_options = [c for c in CLOTHING_OPTIONS if c != clothing]
        witness_clothing = random.choice(other_options)

    knowledge = {
        SUSPECT_A_KEY if criminal == SUSPECT_A else SUSPECT_B_KEY: {
            "knows_own_guilt": True,
            "fake_alibi": fake_alibi,
            "knows_cctv_status": random.choice([True, False]),
            "knows_keycard_log_exists": random.choice([True, False]),
        },
        SUSPECT_A_KEY if innocent == SUSPECT_A else SUSPECT_B_KEY: {
            "knows_own_guilt": False,
            "real_alibi": real_alibi,
            "knows_suspect_A": random.choice([True, False]),
        },
        WITNESS_1_KEY: {
            "saw": f"a person in {witness_clothing} near the {location} around {time}",
            "actual_criminal_clothing": clothing,
            "time_seen": time,
            "bias_target": random.choice([SUSPECT_A, SUSPECT_B, None]),
            "bias_strength": round(random.uniform(0.0, 0.5), 2),
        },
    }
    return knowledge


def _build_detective_briefing(prior_history: dict) -> str:
    """Create a 2-3 sentence briefing for the detective (no criminal identity)."""
    parts = []
    for suspect in SUSPECTS:
        h = prior_history[suspect]
        crimes = h["previous_crimes"]
        crime_str = (
            f"has prior convictions for {', '.join(crimes)}"
            if crimes
            else "has no prior convictions on record"
        )

        lies = h["past_lies_on_record"]
        lie_str = (
            f"and has been caught lying about {lies[0]['topic']} in the past"
            if lies
            else "with no documented dishonesty"
        )

        parole_str = ""
        if h["on_parole"]:
            parole_str = f" Currently on parole ({h['parole_condition']})."

        parts.append(
            f"{suspect} {crime_str} {lie_str}. "
            f"Known behavioral pattern: {h['behavioral_pattern']}.{parole_str}"
        )

    return " ".join(parts)


# ── Public API ──────────────────────────────────────────────────────────────


def generate_case() -> dict:
    """Generate a complete randomized crime case.

    Criminal is randomly assigned to either Suspect_A or Suspect_B.

    Returns a dictionary containing all case information needed for one
    episode of the Crime Investigation environment.
    """
    crime = random.choice(CRIMES)
    criminal = random.choice(list(SUSPECTS))
    time = random.choice(TIMES)
    location = random.choice(LOCATIONS)
    method = random.choice(METHODS)
    motive = random.choice(MOTIVES)

    physical_evidence = _build_evidence(criminal, location, time)
    prior_history = _build_prior_history(criminal)
    agent_knowledge = _build_agent_knowledge(criminal, location, time)
    detective_briefing = _build_detective_briefing(prior_history)

    return {
        "crime": crime,
        "criminal": criminal,
        "time": time,
        "location": location,
        "method": method,
        "motive": motive,
        "physical_evidence": physical_evidence,
        "prior_history": prior_history,
        "agent_knowledge": agent_knowledge,
        "detective_briefing": detective_briefing,
    }