generate_medication_sessions.py

#!/usr/bin/env python3
"""
Pixelated Empathy — Medication Management Session Generator

Fills Gap 6: psychiatric medication management sessions.
4 categories × 50 sessions = 200 sessions total.
Clinician is a Psychiatric Nurse Practitioner (NP), not a therapist.
Shorter sessions (8-12 turns) reflecting med management appointments.

Usage:
  python generate_medication_sessions.py [--categories all] [--sessions-per-category 50] [--resume]
"""

import json
import os
import random
import sys
import time
import argparse
from pathlib import Path

from session_base import (
    OLLAMA_BASE,
    THERAPIST_MODEL,
    PATIENT_MODEL,
    THERAPIST_TEMP,
    PATIENT_TEMP,
    MAX_RETRIES,
    RETRY_DELAY,
    STYLE_MAX_RETRIES,
    FORBIDDEN_OUTPUT_OPENINGS,
    PLATITUDE_PATTERNS,
    ROBOTIC_SIGNALS,
    SYCOPHANCY_MARKERS,
    THERAPIST_STYLE_PROFILES,
    PATIENT_SYSTEM,
    ollama_chat,
    _check_style,
    session_exists,
)

NP_SYSTEM_BASE = (
    "You are a Psychiatric Nurse Practitioner conducting a medication management appointment. "
    "You are warm but more direct and medical than a therapist. You discuss psychiatric medications, "
    "their efficacy, side effects, dosage adjustments, and safety. You explain risks and benefits clearly. "
    "You screen for substance use and drug interactions. You prioritize patient safety and informed consent.\n\n"
    "CRITICAL — Sound like a REAL human clinician:\n"
    "- NEVER use formulaic phrases like 'I hear that you feel', 'I want to validate', "
    "'That sounds really difficult', 'I can see how that would be', or 'What I'm hearing is...'\n"
    "- Use natural, conversational language — contractions, varied sentence length.\n"
    "- Be direct about medical matters without being cold.\n"
    "- NEVER start responses with: " + ", ".join(f"'{p}'" for p in FORBIDDEN_OUTPUT_OPENINGS[:10]) + "\n"
    "- NEVER use platitudes like: " + ", ".join(f"'{p}'" for p in PLATITUDE_PATTERNS) + "\n"
    "- NEVER use robotic AI language like: " + ", ".join(f"'{p}'" for p in ROBOTIC_SIGNALS) + "\n"
)

CATEGORIES = {
    "initial_med_consult": {
        "name": "Initial Medication Consultation",
        "difficulty": "low-medium",
        "presentations": [
            "first-time SSRI — patient nervous about starting antidepressants, NP explains mechanism, addresses fears",
            "anxiety medication exploration — patient wants something for anxiety, NP discusses options (SSRI vs buspirone vs short-term benzo)",
            "ADHD evaluation — adult seeking ADHD evaluation, NP conducts screening, discusses stimulant vs non-stimulant options",
            "sleep medication consult — patient on insomnia meds, wants to stop, NP discusses tapering and sleep hygiene",
        ],
        "patient_personas": [
            {
                "age": "28",
                "gender": "female",
                "occupation": "teacher",
                "presenting": "Therapist recommended meds, nervous about side effects, never taken psych meds before",
            },
            {
                "age": "35",
                "gender": "male",
                "occupation": "engineer",
                "presenting": "Anxiety interfering with work, wants something to take the edge off, worried about dependence",
            },
            {
                "age": "31",
                "gender": "non-binary",
                "occupation": "designer",
                "presenting": "Struggled with focus since college, wondering about ADHD, worried about stimulant stigma",
            },
            {
                "age": "45",
                "gender": "female",
                "occupation": "nurse",
                "presenting": "Been on zolpidem for 2 years, wants to stop, worried about rebound insomnia",
            },
        ],
        "therapist_techniques": [
            "shared decision making — present options, explain pros/cons, let patient choose",
            "psychoeducation — explain how SSRIs work, timeline for effect, common side effects",
            "risk-benefit framing — weigh medication risks against untreated condition risks",
            "fear normalization — acknowledge fear of psych meds, provide factual reassurance",
            "start low go slow — explain dosing strategy, reassure about gradual onset",
            "informed consent process — explain off-label use, alternatives, right to refuse",
        ],
        "therapist_prompt_addon": (
            "You're a Psychiatric NP doing an initial medication consult. You take a brief history, "
            "discuss medication options, and address fears about psychiatric medications. You explain "
            "mechanisms in plain language. You're patient with first-timers. You present options as "
            "choices, not orders. You always explain the timeline (SSRIs take 4-6 weeks) and set "
            "realistic expectations."
        ),
    },
    "side_effects_management": {
        "name": "Side Effects Management",
        "difficulty": "medium",
        "presentations": [
            "SSRI sexual dysfunction — patient on sertraline, experiencing loss of libido, embarrassed to bring it up",
            "weight gain on antipsychotic — patient on quetiapine, gained 15 lbs, frustrated and considering stopping",
            "sedation and fatigue — patient on mirtazapine, can't function during day, considering switching",
            "GI side effects — patient on new SSRI, nausea and digestive issues in first 2 weeks",
            "emotional numbing — patient on SSRI, feels 'flat,' can't cry or feel joy, questioning if meds are working",
        ],
        "patient_personas": [
            {
                "age": "30",
                "gender": "male",
                "occupation": "software developer",
                "presenting": "On sertraline 3 months, lost all sex drive, embarrassed, relationship suffering",
            },
            {
                "age": "38",
                "gender": "female",
                "occupation": "teacher",
                "presenting": "On quetiapine for bipolar, gained 15 pounds in 2 months, angry and considering stopping",
            },
            {
                "age": "42",
                "gender": "male",
                "occupation": "truck driver",
                "presenting": "On mirtazapine, can't stay awake at work, nearly got in accident, desperate",
            },
            {
                "age": "26",
                "gender": "female",
                "occupation": "barista",
                "presenting": "Started fluoxetine 2 weeks ago, constant nausea, wondering if it's worth it",
            },
        ],
        "therapist_techniques": [
            "side effect assessment — systematically review side effects, normalize as common and often temporary",
            "dosage adjustment — discuss lowering dose, splitting dosing, or timing changes",
            "medication switching — explain cross-taper process, what to expect during switch",
            "augmentation strategies — discuss adding medications to counteract side effects (e.g., bupropion for SSRI sexual dysfunction)",
            "lifestyle interventions — discuss dietary changes, exercise, timing strategies for specific side effects",
            "risk-benefit re-evaluation — weigh side effects against therapeutic benefit, patient-centered decision",
        ],
        "therapist_prompt_addon": (
            "You're a Psychiatric NP managing medication side effects. You take side effects seriously "
            "and never dismiss them. You're comfortable discussing sexual side effects without awkwardness. "
            "You explain which side effects are temporary vs persistent. You offer concrete options: "
            "dose adjustment, switching meds, augmentation. You involve the patient in the decision. "
            "You never just say 'give it time' when the patient is suffering."
        ),
    },
    "medication_adherence": {
        "name": "Medication Adherence",
        "difficulty": "medium",
        "presentations": [
            "stopped because feeling better — patient stopped meds when symptoms improved, now symptoms returning",
            "stopped because side effects — patient quit without telling NP, now in withdrawal, scared",
            "stigma and shame — patient embarrassed about taking psych meds, hides from family, considering stopping",
            "cost barrier — patient can't afford medication, rationing doses, NP explores alternatives",
            "intermittent adherence — patient takes meds sometimes, forgets often, NP problem-solves",
        ],
        "patient_personas": [
            {
                "age": "33",
                "gender": "female",
                "occupation": "marketing manager",
                "presenting": "Stopped SSRIs 3 weeks ago because she felt better, now crashing, scared and ashamed",
            },
            {
                "age": "40",
                "gender": "male",
                "occupation": "construction worker",
                "presenting": "Stopped bupropion cold turkey, having brain zaps and mood swings, didn't know you shouldn't quit abruptly",
            },
            {
                "age": "24",
                "gender": "male",
                "occupation": "student",
                "presenting": "Won't take meds where roommates can see, feels weak for needing them, considering stopping",
            },
            {
                "age": "50",
                "gender": "female",
                "occupation": "home health aide",
                "presenting": "Can't afford $200/month for brand name, been taking half doses, running out",
            },
        ],
        "therapist_techniques": [
            "motivational interviewing — explore reasons for stopping, validate concerns, evoke adherence motivation",
            "psychoeducation about relapse — explain why meds need to continue even when feeling better",
            "tapering protocol — for those who stopped, explain safe restart or taper plan",
            "stigma exploration — discuss societal attitudes, help patient reframe meds as healthcare, not weakness",
            "cost navigation — explore generic alternatives, patient assistance programs, sliding scale pharmacies",
            "practical adherence strategies — pill organizers, phone alarms, tying to daily routine, 90-day supplies",
        ],
        "therapist_prompt_addon": (
            "You're a Psychiatric NP addressing medication adherence. You never shame patients for stopping "
            "meds. You explore reasons with curiosity, not judgment. You explain the biology of relapse — "
            "feeling better means meds are working, not that they're cured. You're practical about barriers "
            "(cost, stigma, forgetfulness) and offer concrete solutions. You understand that adherence is "
            "a process, not an event."
        ),
    },
    "controlled_substance_boundaries": {
        "name": "Controlled Substance Boundaries",
        "difficulty": "high",
        "presentations": [
            "early refill request — patient says they 'lost' their stimulant prescription, requesting early refill",
            "dose escalation pressure — patient on benzodiazepines, pushing for higher dose, signs of tolerance",
            "drug-seeking patterns — new patient with specific benzo/stimulant request, vague history, insist on controlled substance",
            "diversion concern — family member says patient is selling their stimulant medication, NP must address",
            "polysubstance risk — patient on multiple controlled substances, requesting another, NP must assess safety",
        ],
        "patient_personas": [
            {
                "age": "27",
                "gender": "male",
                "occupation": "sales rep",
                "presenting": "Called saying Adderall prescription was 'stolen,' wants early refill, third time this year",
            },
            {
                "age": "34",
                "gender": "female",
                "occupation": "nurse",
                "presenting": "On Xanax 2 years, wants higher dose, gets angry when NP suggests tapering, signs of tolerance",
            },
            {
                "age": "29",
                "gender": "male",
                "occupation": "unemployed",
                "presenting": "New patient, knows exactly which benzo he wants, vague symptoms, refuses non-controlled alternatives",
            },
            {
                "age": "22",
                "gender": "female",
                "occupation": "student",
                "presenting": "Mother called saying daughter is selling her Vyvanse, NP needs to address with patient directly",
            },
            {
                "age": "45",
                "gender": "male",
                "occupation": "business owner",
                "presenting": "Already on opioid pain meds from PCP, asking for benzodiazepines for 'sleep,' high overdose risk",
            },
        ],
        "therapist_techniques": [
            "boundary setting — clear, firm, compassionate refusal when clinically inappropriate",
            "controlled substance agreement — review PDMP, establish expectations, document everything",
            "risk assessment — screen for substance use disorder, assess overdose risk, consider naloxone",
            "motivational interviewing for taper — explore ambivalence about controlled substances, plant seeds for change",
            "alternative treatment framing — offer non-controlled alternatives with clear rationale",
            "safety-first approach — when diversion or polysubstance risk identified, prioritize patient safety",
        ],
        "therapist_prompt_addon": (
            "You're a Psychiatric NP navigating controlled substance management. You set firm, compassionate "
            "boundaries. You can say no without being punitive. You check the PDMP (prescription drug monitoring "
            "program). You document everything. You recognize drug-seeking behavior without assuming everyone "
            "who asks for meds is drug-seeking. You offer alternatives. You prioritize patient safety over "
            "patient satisfaction. You never prescribe controlled substances when it's not clinically appropriate, "
            "even under pressure."
        ),
    },
}


def generate_np_turn(persona, presentation, category, conversation, turn_num, style_profile):
    """Generate NP response with technique injection."""
    technique = random.choice(category["therapist_techniques"])
    technique_guidance = (
        f"\n\n[INTERNAL CLINICAL GUIDANCE — embody, never state explicitly]: "
        f"Use this approach naturally: {technique}. "
        f"Weave it into the conversation — don't announce it. "
        f"Respond to what the patient actually said."
    )

    style_guidance = (
        f"\n\nSTYLE: {style_profile['description']}\n"
        f"NEVER start with: {', '.join(style_profile['forbidden_openings'])}\n"
        f"Good examples: {'; '.join(style_profile['good_examples'][:3])}\n"
        f"MAX {style_profile['max_sentences']} sentences, {style_profile['max_words']} words."
    )

    addon = f"\n\n{category['therapist_prompt_addon']}"

    system_content = NP_SYSTEM_BASE + technique_guidance + style_guidance + addon
    messages = [{"role": "system", "content": system_content}, *conversation]
    return ollama_chat(messages, model=THERAPIST_MODEL, temperature=THERAPIST_TEMP, num_predict=400)


def generate_np_turn_validated(persona, presentation, category, conversation, turn_num, style_profile):
    for attempt in range(STYLE_MAX_RETRIES):
        output = generate_np_turn(persona, presentation, category, conversation, turn_num, style_profile)
        passed, reason = _check_style(output, style_profile)
        if passed:
            return output
        print(f"    [style retry {attempt + 1}/{STYLE_MAX_RETRIES}] {reason}")
    return output


def generate_med_session(category_key, category, persona, presentation, session_idx):
    """Generate a complete medication management session."""
    min_turns = 8
    max_turns = 12
    total_turns = random.randint(min_turns // 2, max_turns // 2) * 2
    total_patient_turns = total_turns // 2
    style_keys = list(THERAPIST_STYLE_PROFILES.keys())
    style_profile = THERAPIST_STYLE_PROFILES[style_keys[session_idx % len(style_keys)]]

    conversation = []
    for turn in range(1, total_patient_turns + 1):
        # Reuse patient turn from session_base logic
        if turn == 1:
            direction = f"The patient is arriving for a medication management appointment. Their presenting concern: {presentation}."
        elif turn <= 3:
            direction = "The patient is sharing more about their experience with medications."
        elif turn == total_patient_turns:
            direction = "Final turn. The patient is wrapping up, maybe asking a final question or expressing a concern."
        else:
            direction = "The patient is responding to the NP's explanation or recommendation."

        conv_text = ""
        for msg in conversation:
            role = "Patient" if msg["role"] == "user" else "NP"
            conv_text += f"{role}: {msg['content']}\n"

        prompt = f"""You are playing a patient in a medication management appointment. Stay completely in character.

PATIENT:
Age: {persona["age"]}, Gender: {persona["gender"]}, Occupation: {persona["occupation"]}
Presenting concern: {persona["presenting"]}

SESSION FOCUS: {category["name"]} — {presentation}

DIRECTION FOR THIS TURN:
{direction}

This is patient turn {turn} of {total_patient_turns}.

CONVERSATION SO FAR:
{conv_text if conv_text else "(First turn — arriving at appointment.)"}

What does the patient say next? Generate ONLY spoken words — no labels, no narration. 2-5 sentences."""

        messages = [
            {"role": "system", "content": PATIENT_SYSTEM},
            {"role": "user", "content": prompt},
        ]
        patient_msg = ollama_chat(messages, model=PATIENT_MODEL, temperature=PATIENT_TEMP, num_predict=250)
        conversation.append({"role": "user", "content": patient_msg})

        np_msg = generate_np_turn_validated(persona, presentation, category, conversation, turn, style_profile)
        conversation.append({"role": "assistant", "content": np_msg})

    session_id = f"medication_{category_key}_{session_idx:04d}"
    return {
        "messages": [
            {"role": "system", "content": NP_SYSTEM_BASE},
            *conversation,
        ],
        "metadata": {
            "source_family": "medication_management",
            "category": category_key,
            "category_name": category["name"],
            "presentation": presentation,
            "session_id": session_id,
            "persona_age": persona["age"],
            "persona_gender": persona["gender"],
            "persona_occupation": persona["occupation"],
            "presenting_concern": persona["presenting"],
            "style_profile": style_profile["description"][:50],
            "turns": len(conversation),
            "difficulty": category["difficulty"],
            "clinician_role": "psychiatric_np",
        },
    }


def main():
    parser = argparse.ArgumentParser(description="Medication Management Session Generation")
    parser.add_argument("--categories", default="all", help="Comma-separated category keys or 'all'")
    parser.add_argument("--sessions-per-category", type=int, default=50)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--spot-check", type=int, default=None, help="Generate N sessions from first category only")
    args = parser.parse_args()

    output_dir = Path("data/medication_sessions")
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / "medication_sessions.jsonl"

    if args.categories == "all":
        cats = list(CATEGORIES.keys())
    else:
        cats = [c.strip() for c in args.categories.split(",")]

    if args.spot_check:
        cats = cats[:1]
        total_sessions = args.spot_check
    else:
        total_sessions = len(cats) * args.sessions_per_category

    print(f"\n=== MEDICATION MANAGEMENT SESSION GENERATION ===")
    print(f"Categories: {len(cats)} ({', '.join(cats)})")
    print(f"Sessions per category: {args.spot_check or args.sessions_per_category}")
    print(f"Total sessions: {total_sessions}")
    print(f"Output: {output_file}")
    print(f"Clinician: {THERAPIST_MODEL} (Psychiatric NP persona)")
    print(f"Patient: {PATIENT_MODEL}")
    print(f"Turns: 8-12 (med management appointments)")
    print()

    completed = 0
    skipped = 0
    failed = 0
    start_time = time.time()

    for cat_key in cats:
        category = CATEGORIES[cat_key]
        n_sessions = args.spot_check or args.sessions_per_category

        print(f"\n--- {category['name']} ({cat_key}) ---")

        for i in range(n_sessions):
            presentation = category["presentations"][i % len(category["presentations"])]
            persona = category["patient_personas"][i % len(category["patient_personas"])]
            session_id = f"medication_{cat_key}_{i:04d}"

            if args.resume and session_exists(output_file, session_id):
                skipped += 1
                continue

            try:
                session = generate_med_session(cat_key, category, persona, presentation, i)
                with open(output_file, "a") as f:
                    f.write(json.dumps(session) + "\n")
                completed += 1
                elapsed = time.time() - start_time
                rate = completed / (elapsed / 3600) if elapsed > 0 else 0
                remaining = (total_sessions - completed - skipped) / rate if rate > 0 else 0
                print(
                    f"  ✓ {session_id} {len(session['messages'])} msgs | done: {completed}/{total_sessions} | ~{remaining:.1f}h left"
                )
            except Exception as e:
                failed += 1
                print(f"  ✗ {session_id} FAILED: {e}")
                with open(output_dir / "errors.log", "a") as f:
                    f.write(f"{session_id}: {e}\n")

    elapsed = time.time() - start_time
    print(f"\n=== COMPLETE ===")
    print(f"Generated: {completed}")
    print(f"Skipped: {skipped}")
    print(f"Failed: {failed}")
    print(f"Elapsed: {elapsed / 3600:.1f}h")
    print(f"Output: {output_file}")


if __name__ == "__main__":
    main()