generate_termination_sessions.py

#!/usr/bin/env python3
"""
Pixelated Empathy — Therapy Termination Session Generator

Fills Gap 7: therapy termination sessions.
4 categories × 50 sessions = 200 sessions total.
These are FINAL therapy sessions — the goodbye.

Usage:
  python generate_termination_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,
    PIXEL_SYSTEM_BASE,
    PATIENT_SYSTEM,
    ollama_chat,
    _check_style,
    session_exists,
)

TERMINATION_PATIENT_SYSTEM = (
    PATIENT_SYSTEM + "\n\n"
    "This is the FINAL therapy session. You know this is the last appointment. "
    "You may feel relief, grief, gratitude, ambivalence, or all of these at once. "
    "The relationship with this therapist has mattered to you. "
    "Express what ending feels like honestly."
)

CATEGORIES = {
    "successful_termination": {
        "name": "Successful Termination",
        "difficulty": "low",
        "presentations": [
            "anxiety treatment completed — patient came for anxiety, learned CBT skills, symptoms manageable, both agree it's time",
            "depression remission — depression lifted months ago, maintenance solid, patient feels ready to fly solo",
            "relationship work done — couple came for communication, made real progress, both feel equipped to continue alone",
            "grief processed — came after a loss, grief has integrated, patient feels ready to move forward",
            "skills mastery — patient came for DBT skills, mastered them, using them independently, graduation natural",
        ],
        "patient_personas": [
            {
                "age": "32",
                "gender": "female",
                "occupation": "teacher",
                "presenting": "Came for anxiety 6 months ago, CBT skills working, hasn't had a panic attack in 2 months, feels ready",
            },
            {
                "age": "28",
                "gender": "male",
                "occupation": "engineer",
                "presenting": "Depression in remission, exercising again, socializing, feels like himself, wants to stop therapy",
            },
            {
                "age": "35",
                "gender": "non-binary",
                "occupation": "designer",
                "presenting": "Grief over parent's death 1 year ago, has processed, still misses them but functioning, ready to end",
            },
            {
                "age": "41",
                "gender": "female",
                "occupation": "nurse manager",
                "presenting": "Learned DBT skills, using them daily, hasn't had an emotional crisis in months, confident in skills",
            },
        ],
        "therapist_techniques": [
            "progress review — concretely review what brought them in and what changed, celebrate specific gains",
            "relapse prevention plan — identify triggers, early warning signs, coping strategies, when to return",
            "confidence building — reflect patient's growth back to them, reinforce self-efficacy",
            "relationship closure — acknowledge the therapeutic relationship, say goodbye genuinely",
            "booster session planning — offer a check-in in 1 month, leave the door open without creating dependency",
            "future orientation — help patient envision continued growth without therapy, normalize future struggles",
        ],
        "therapist_prompt_addon": (
            "This is a successful termination. The patient has made real progress and you both agree it's time "
            "to end. You celebrate specific, concrete gains — not vague 'you've grown so much.' You review "
            "what brought them in and what changed. You build a relapse prevention plan together. You say "
            "goodbye with genuine warmth, not clinical detachment. You leave the door open for future "
            "check-ins without creating dependency. This is a graduation, not a loss."
        ),
    },
    "relapse_prevention_planning": {
        "name": "Relapse Prevention Planning",
        "difficulty": "medium",
        "presentations": [
            "anxiety relapse prevention — patient improved but anxiety has a pattern of returning, plan for future",
            "depression maintenance — depression is episodic, building a plan for when the next episode comes",
            "substance use recovery plan — early recovery, building a detailed plan for triggers and urges",
            "trauma recovery maintenance — trauma processing complete, but triggers will return, plan for future",
            "relationship patterns — old patterns may resurface, build a plan for recognizing and interrupting",
        ],
        "patient_personas": [
            {
                "age": "30",
                "gender": "female",
                "occupation": "marketing manager",
                "presenting": "Anxiety improved but came back before after stopping therapy, wants a plan this time, cautiously hopeful",
            },
            {
                "age": "38",
                "gender": "male",
                "occupation": "teacher",
                "presenting": "Third depressive episode, each time stopped therapy too early, this time wants a real plan",
            },
            {
                "age": "25",
                "gender": "non-binary",
                "occupation": "student",
                "presenting": "6 months sober, therapy ending, nervous about maintaining recovery without the anchor of weekly sessions",
            },
            {
                "age": "44",
                "gender": "female",
                "occupation": "social worker",
                "presenting": "Processed childhood trauma, but knows triggers will come back, wants a concrete plan for next time",
            },
        ],
        "therapist_techniques": [
            "trigger identification — collaboratively map specific triggers and early warning signs",
            "coping strategy rehearsal — review and practice skills learned in therapy, ensure they're internalized",
            "relapse signature — identify patient's unique pattern of relapse (sleep change first? isolation? irritability?)",
            "cris resource list — concrete: who to call, what to do, when to return to therapy, crisis line numbers",
            "self-monitoring plan — what to watch for, how often to check in with self, journaling or tracking tools",
            "compassionate self-talk plan — prepare the internal voice for when things get hard, not harsh self-criticism",
        ],
        "therapist_prompt_addon": (
            "This session focuses on building a concrete relapse prevention plan. You're practical and "
            "collaborative. You help the patient identify their unique early warning signs — not generic "
            "ones. You review the specific skills they learned and ensure they can use them independently. "
            "You create a written plan they can take with them. You normalize that relapse is part of "
            "recovery, not failure. You build the patient's confidence in self-management."
        ),
    },
    "referral_handoff": {
        "name": "Referral & Transition Handoff",
        "difficulty": "medium-high",
        "presentations": [
            "therapist relocation — therapist is moving, patient must transfer to someone new, attachment rupture",
            "insurance change — patient's insurance changed, therapist is out of network, forced to switch",
            "specialty referral — patient needs specialized care therapist can't provide, referral to specialist",
            "therapist leave — therapist going on extended leave (medical, sabbatical), patient feels abandoned",
            "trainee transition — patient was seeing a trainee who is graduating, must transition to permanent therapist",
        ],
        "patient_personas": [
            {
                "age": "29",
                "gender": "female",
                "occupation": "nurse",
                "presenting": "Therapist is moving across country, she's been seeing her for 3 years, feels like being abandoned all over again",
            },
            {
                "age": "35",
                "gender": "male",
                "occupation": "electrician",
                "presenting": "Insurance changed, can't afford out of pocket, feels like the system is forcing him to start over",
            },
            {
                "age": "26",
                "gender": "non-binary",
                "occupation": "student",
                "presenting": "Therapist says they need specialized trauma care she can't provide, feels rejected even though it's the right call",
            },
            {
                "age": "33",
                "gender": "female",
                "occupation": "teacher",
                "presenting": "Therapist going on medical leave for 4 months, devastated, working through what it brings up from her past",
            },
        ],
        "therapist_techniques": [
            "attachment processing — validate that this transition triggers attachment wounds, it's not just logistics",
            "grief work — this is a real loss, name it, grieve it, don't minimize it as 'just a transfer'",
            "continuity reassurance — explain what will transfer (records, treatment plan) and what won't (the relationship)",
            "transition preparation — help patient articulate what they need from the next therapist",
            "saying goodbye — acknowledge the relationship mattered, don't pretend it didn't",
            "self-advocacy coaching — prepare patient to advocate for themselves with new provider",
        ],
        "therapist_prompt_addon": (
            "This is a forced termination — you're referring or transitioning the patient, not ending because "
            "treatment is complete. You validate that this is a loss. You don't minimize it as 'just "
            "a transfer.' You process the attachment rupture honestly. You acknowledge that the relationship "
            "mattered and can't simply be replaced. You provide practical continuity (records, treatment "
            "plan, what to tell the new therapist) while holding space for the grief. You're not defensive "
            "about the patient's feelings about the transition."
        ),
    },
    "attachment_disruption": {
        "name": "Attachment Disruption Processing",
        "difficulty": "high",
        "presentations": [
            "patient-initiated premature termination — patient wants to stop suddenly, therapist suspects avoidance",
            "frustrated and leaving — patient angry at therapist, feels therapy isn't working, wants to quit",
            "ambivalent departure — patient says they want to stop but clearly has more work to do, therapist must respect autonomy",
            "attachment-triggered departure — patient leaving because the therapeutic relationship is getting close and that's terrifying",
            "devaluation and exit — patient devaluing therapy and therapist as a defense against ending, classic borderline pattern",
        ],
        "patient_personas": [
            {
                "age": "27",
                "gender": "female",
                "occupation": "barista",
                "presenting": "Wants to stop after 2 months, says she's fine, therapist knows she's not, feels like she's running",
            },
            {
                "age": "40",
                "gender": "male",
                "occupation": "business owner",
                "presenting": "Frustrated, says therapy is a waste of time, actually therapy is touching real issues and he's scared",
            },
            {
                "age": "24",
                "gender": "non-binary",
                "occupation": "student",
                "presenting": "Says they want to stop but cries when they talk about it, clearly ambivalent, therapist must hold both",
            },
            {
                "age": "31",
                "gender": "female",
                "occupation": "therapist",
                "presenting": "Getting close to core material, wants to quit, the closeness is triggering her own attachment wounds",
            },
            {
                "age": "35",
                "gender": "male",
                "occupation": "lawyer",
                "presenting": "Started devaluing therapy, 'this isn't working,' projection of critical father onto therapist, edge of personality patterns",
            },
        ],
        "therapist_techniques": [
            "non-defensive processing — don't take the devaluation personally, explore it curiously",
            "autonomy respect — patient has the right to leave, name it without guilt-tripping",
            "ambivalence exploration — hold both 'I want to go' and 'I'm not ready' simultaneously",
            "attachment interpretation — gently name the pattern: 'leaving before being left?' — without forcing insight",
            "door left open — make clear they can return, no judgment, no 'I told you so'",
            "meaning-making — what is this departure communicating, what does it protect against",
        ],
        "therapist_prompt_addon": (
            "This is a difficult termination — the patient is leaving prematurely, possibly as an avoidance "
            "or attachment defense. You do NOT guilt-trip or convince them to stay. You respect their "
            "autonomy absolutely. But you also name what you see: the pattern, the timing, what it might "
            "be protecting against. You're not defensive if they're angry or devaluing. You hold both: "
            "'this is your choice' AND 'I think there's more work here.' You leave the door open without "
            "conditions. You process your own feelings about the rupture without burdening the patient."
        ),
    },
}


def generate_termination_patient_turn(
    persona, presentation, category_name, conversation, turn_num, total_patient_turns
):
    """Generate patient turn in a termination session."""
    if turn_num == 1:
        direction = (
            f"This is the final therapy session. The patient knows this is the last appointment. "
            f"Their presenting concern: {presentation}. They're arriving with awareness that this is ending."
        )
    elif turn_num <= 3:
        direction = "The patient is processing the ending, maybe sharing what the therapy has meant, or struggling with the goodbye."
    elif turn_num <= 5:
        direction = (
            "The patient is going deeper — expressing what ending brings up, maybe grief, maybe relief, maybe both."
        )
    elif turn_num == total_patient_turns:
        direction = "Final turn. The patient is saying goodbye — in whatever way feels real to them."
    else:
        direction = "The patient is responding to the therapist, processing the ending, maybe pushing back or having a realization."

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

    prompt = f"""You are playing a therapy patient in their FINAL session. Stay completely in character.

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

SESSION FOCUS: {category_name} — {presentation}
This is the LAST therapy session.

DIRECTION FOR THIS TURN:
{direction}

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

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

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

    messages = [
        {"role": "system", "content": TERMINATION_PATIENT_SYSTEM},
        {"role": "user", "content": prompt},
    ]
    return ollama_chat(messages, model=PATIENT_MODEL, temperature=PATIENT_TEMP, num_predict=250)


def generate_termination_session(category_key, category, persona, presentation, session_idx):
    """Generate a complete termination therapy session."""
    min_turns = 8
    max_turns = 14
    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)]]

    # Use termination-specific system prompt for therapist too
    termination_addon = category["therapist_prompt_addon"]

    conversation = []
    for turn in range(1, total_patient_turns + 1):
        patient_msg = generate_termination_patient_turn(
            persona, presentation, category["name"], conversation, turn, total_patient_turns
        )
        conversation.append({"role": "user", "content": patient_msg})

        # Therapist turn with termination-specific guidance
        technique = random.choice(category["therapist_techniques"])
        technique_guidance = (
            f"\n\n[INTERNAL CLINICAL GUIDANCE — embody, never state explicitly]: "
            f"Use this technique naturally: {technique}. "
            f"This is a termination session — the ending IS the work."
        )

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

        system_content = PIXEL_SYSTEM_BASE + technique_guidance + style_guidance + f"\n\n{termination_addon}"
        messages = [{"role": "system", "content": system_content}, *conversation]
        therapist_msg = ollama_chat(messages, model=THERAPIST_MODEL, temperature=THERAPIST_TEMP, num_predict=400)

        # Style validation
        passed, reason = _check_style(therapist_msg, style_profile)
        if not passed:
            for attempt in range(STYLE_MAX_RETRIES - 1):
                print(f"    [style retry {attempt + 1}/{STYLE_MAX_RETRIES}] {reason}")
                therapist_msg = ollama_chat(
                    messages, model=THERAPIST_MODEL, temperature=THERAPIST_TEMP, num_predict=400
                )
                passed, reason = _check_style(therapist_msg, style_profile)
                if passed:
                    break

        conversation.append({"role": "assistant", "content": therapist_msg})

    session_id = f"termination_{category_key}_{session_idx:04d}"
    return {
        "messages": [
            {"role": "system", "content": PIXEL_SYSTEM_BASE},
            *conversation,
        ],
        "metadata": {
            "source_family": "therapy_termination",
            "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"],
            "session_type": "termination",
        },
    }


def main():
    parser = argparse.ArgumentParser(description="Therapy Termination 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/termination_sessions")
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / "termination_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=== THERAPY TERMINATION 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"Therapist: {THERAPIST_MODEL}")
    print(f"Patient: {PATIENT_MODEL}")
    print(f"Turns: 8-14 (termination sessions)")
    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"termination_{cat_key}_{i:04d}"

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

            try:
                session = generate_termination_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()