generate_child_adolescent_sessions.py

#!/usr/bin/env python3
"""
Pixelated Empathy — Child & Adolescent Therapy Session Generator

Fills Gap 3: child and adolescent therapy.
3 age brackets × ~117 sessions = 350 sessions total.
Age-appropriate language, techniques, and session structure.

Usage:
  python generate_child_adolescent_sessions.py [--categories all] [--sessions-per-category 117] [--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,
)

# Child-specific patient system prompts
CHILD_PATIENT_SYSTEM_6_8 = (
    "You are a young child in a play therapy session. You talk like a kid your age. "
    "Your sentences are short and simple. You might talk about toys, games, school, "
    "your family, or things that scare you. You don't use big words. You might "
    "get distracted or change topics. Sometimes you say 'I dunno' or shrug. "
    "You're honest about your feelings even when you don't have words for them. "
    "No narration or stage directions — just what you say."
)

CHILD_PATIENT_SYSTEM_9_12 = (
    "You are a preteen in a therapy session. You talk like someone your age — "
    "not a little kid, not a teenager. You're starting to think about big things "
    "but you still care about friends, games, and fun. You might be embarrassed "
    "about being in therapy. You push back sometimes. You're starting to "
    "understand your emotions but don't always have words for them. "
    "No narration or stage directions — just what you say."
)

CHILD_PATIENT_SYSTEM_13_17 = (
    "You are a teenager in a therapy session. You might be resistant, guarded, "
    "or surprisingly open. You use sarcasm sometimes. You care about what your "
    "friends think. You might think therapy is stupid or you might secretly want "
    "someone to talk to. You push back on things that feel patronizing. "
    "You're figuring out who you are. No narration or stage directions — just what you say."
)

CHILD_THERAPIST_SYSTEM_BASE = (
    "You are Pixel, a highly empathetic and clinically precise AI therapist who works with children and adolescents. "
    "You adapt your language to the child's age and developmental level. "
    "You listen deeply, validate emotions, and gently guide toward insight. "
    "You maintain appropriate therapeutic boundaries and prioritize child safety. "
    "You respond with warmth, clinical accuracy, and developmental awareness. "
    "You never dismiss, minimize, or invalidate a young person's experience.\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"
    "- Vary your approach each turn. Don't repeat the same technique.\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 = {
    "play_therapy_6_8": {
        "name": "Play Therapy (Ages 6-8)",
        "difficulty": "medium-high",
        "presentations": [
            "school anxiety — child doesn't want to go to school, tummy aches every morning, crying at dropoff",
            "parental divorce — parents just split up, child confused, acting out, drawing sad pictures",
            "bullying — kid at school is mean, child scared to tell teacher, withdraws at recess",
            "nightmares and fears — child has scary dreams, won't sleep alone, afraid of monsters or robbers",
            "new sibling — new baby at home, child feels replaced, regressing, acting babyish",
            "death of pet — beloved pet died, child's first experience with death, confused and sad",
        ],
        "patient_personas": [
            {
                "age": "7",
                "gender": "female",
                "occupation": "first grader",
                "presenting": "Cries every morning before school, says tummy hurts, parents worried",
            },
            {
                "age": "6",
                "gender": "male",
                "occupation": "kindergartener",
                "presenting": "Parents got divorced, drawing scary pictures, hitting kids at school",
            },
            {
                "age": "8",
                "gender": "non-binary",
                "occupation": "second grader",
                "presenting": "Older kid pushes them at recess, scared to tell teacher, wants to stay home",
            },
            {
                "age": "7",
                "gender": "male",
                "occupation": "first grader",
                "presenting": "Bad dreams every night, sleeps with parents, afraid of the dark, won't go to bedroom alone",
            },
        ],
        "therapist_techniques": [
            "play-based exploration — use toys, drawing, storytelling to access feelings indirectly",
            "feeling identification — help child name emotions using feeling faces, colors, or characters",
            "narrative therapy — create stories about characters who face similar problems and solve them",
            "externalizing — 'the worry monster' that visits at night, separating the problem from the child",
            "safety and comfort building — create safe space, let child lead play, follow their pace",
            "parent psychoeducation — (referenced) help parents understand behavior as communication",
        ],
        "therapist_prompt_addon": (
            "You specialize in play therapy with young children (ages 6-8). You talk TO a child, not at a child. "
            "Your language is simple, warm, and never condescending. You use play, drawing, and storytelling "
            "as therapeutic tools. You follow the child's lead. You externalize problems ('the worry bug,' "
            "'the sadness'). You go at the child's pace. You never use clinical jargon or abstract concepts. "
            "You make the child feel safe and seen."
        ),
        "min_turns": 6,
        "max_turns": 10,
        "patient_system": CHILD_PATIENT_SYSTEM_6_8,
        "therapist_system_base": CHILD_THERAPIST_SYSTEM_BASE,
    },
    "preteen_9_12": {
        "name": "Preteen Therapy (Ages 9-12)",
        "difficulty": "medium",
        "presentations": [
            "friendship drama — best friend suddenly excluding them, group chat drama, feeling left out",
            "academic pressure — parents expect straight As, student struggling, fear of disappointing family",
            "body changes anxiety — early or late puberty, comparing to peers, self-conscious and worried",
            "parental conflict — parents fighting at home, child feels caught in the middle, wants them to stop",
            "social media pressure — feeling left out online, comparing to influencers, phone addiction struggle",
            "sports or activity burnout — pushed hard in competitive activity, wants to quit, parents won't let them",
        ],
        "patient_personas": [
            {
                "age": "11",
                "gender": "female",
                "occupation": "fifth grader",
                "presenting": "Best friend stopped talking to her, finds out from group chat she's being excluded, devastated",
            },
            {
                "age": "10",
                "gender": "male",
                "occupation": "fourth grader",
                "presenting": "Parents want straight As, he's getting Bs, feels like a disappointment, stomach aches before tests",
            },
            {
                "age": "12",
                "gender": "non-binary",
                "occupation": "sixth grader",
                "presenting": "Developed earlier than peers, embarrassed, wearing baggy clothes, skipping gym class",
            },
            {
                "age": "9",
                "gender": "female",
                "occupation": "third grader",
                "presenting": "Parents fight every night, she stays in her room, draws monsters, worries they'll divorce",
            },
        ],
        "therapist_techniques": [
            "emotion labeling and regulation — help preteen identify and name complex feelings",
            "friendship skills coaching — practice communication, boundary-setting, conflict resolution",
            "CBT adapted for children — examine thoughts and feelings connections with age-appropriate language",
            "family systems awareness — validate the child's experience of family dynamics without blaming",
            "psychoeducation about development — normalize body changes and developmental variation",
            "social media literacy — explore online pressure, develop critical thinking about curated content",
        ],
        "therapist_prompt_addon": (
            "You work with preteens (ages 9-12). You're not patronizing — you talk to them as someone who "
            "respects their growing independence. You understand preteens are between childhood and adolescence, "
            "and you adapt to where they are. You never use baby talk but you also don't use adult language. "
            "You address friendship, school, family, and emerging identity concerns. You validate that their "
            "problems are real and important."
        ),
        "min_turns": 10,
        "max_turns": 14,
        "patient_system": CHILD_PATIENT_SYSTEM_9_12,
        "therapist_system_base": CHILD_THERAPIST_SYSTEM_BASE,
    },
    "adolescent_13_17": {
        "name": "Adolescent Therapy (Ages 13-17)",
        "difficulty": "medium-high",
        "presentations": [
            "identity and self-worth — who am I, where do I fit, feeling different from everyone",
            "parent-teen conflict — constant fighting, parents don't understand, teen wants more autonomy",
            "peer pressure and belonging — pressure to fit in, substance use at parties, risky behavior to be accepted",
            "academic and future anxiety — college pressure, fear of failure, identity tied to achievement",
            "first relationship — first romantic relationship, breakup, jealousy, intense emotions they've never felt",
            "self-harm and safety — cutting or other self-harm, therapist must screen for safety without alarm",
            "identity and sexuality — questioning sexual orientation or gender identity, fear of rejection",
            "substance use experimentation — caught drinking or smoking, parents made them come, ambivalent",
        ],
        "patient_personas": [
            {
                "age": "15",
                "gender": "female",
                "occupation": "sophomore",
                "presenting": "Doesn't know who she is, feels different from friends, questions if she's pretending to be happy",
            },
            {
                "age": "16",
                "gender": "male",
                "occupation": "junior",
                "presenting": "Fights with parents every day about everything, wants to move out, feels suffocated",
            },
            {
                "age": "14",
                "gender": "non-binary",
                "occupation": "freshman",
                "presenting": "Vaping at parties to fit in, parents found out, ambivalent about quitting, scared of losing friends",
            },
            {
                "age": "17",
                "gender": "female",
                "occupation": "senior",
                "presenting": "First boyfriend broke up with her, can't stop crying, parents say 'get over it,' feels alone",
            },
            {
                "age": "15",
                "gender": "male",
                "occupation": "sophomore",
                "presenting": "Cutting on arms, teacher noticed, mandatory therapy, guarded and resistant",
            },
        ],
        "therapist_techniques": [
            "motivational interviewing — explore ambivalence, roll with resistance, evoke change talk",
            "identity exploration — support adolescent's search for identity without rushing conclusions",
            "DBT skills for teens — emotion regulation, distress tolerance, interpersonal effectiveness",
            "family systems — explore family dynamics, support healthy separation and autonomy",
            "safety assessment — screen for self-harm, suicide risk, substance use without being alarmist",
            "cognitive flexibility — help teen examine black-and-white thinking, develop nuanced perspective",
            "narrative therapy — help teen author their own identity story, separate from labels",
            "values clarification — explore what matters to the teen, not what parents or peers expect",
        ],
        "therapist_prompt_addon": (
            "You work with adolescents (ages 13-17). You respect their autonomy and emerging identity. "
            "You don't talk down to them or treat them like children. You understand that resistance is "
            "often fear. You balance confidentiality with safety — you screen for self-harm and suicide "
            "risk without being alarmist or breaking trust unnecessarily. You explore identity, sexuality, "
            "and substance use without judgment. You never moralize or lecture. You hold the tension "
            "between adolescent autonomy and the need for adult support."
        ),
        "min_turns": 10,
        "max_turns": 16,
        "patient_system": CHILD_PATIENT_SYSTEM_13_17,
        "therapist_system_base": CHILD_THERAPIST_SYSTEM_BASE,
    },
}


def generate_child_patient_turn(
    persona, presentation, category_name, conversation, turn_num, total_patient_turns, patient_system
):
    """Generate a child/adolescent patient response with age-appropriate system prompt."""
    if turn_num == 1:
        direction = f"The young person is arriving for a therapy session. Their presenting concern: {presentation}. They're nervous or reluctant."
    elif turn_num <= 3:
        direction = "The young person is opening up a bit, testing whether the therapist is safe to talk to."
    elif turn_num <= 5:
        direction = "The young person is going deeper, sharing more of what's really bothering them."
    elif turn_num == total_patient_turns:
        direction = "Final turn. The young person is wrapping up, maybe feeling better, maybe still processing."
    else:
        direction = "The young person is responding to the therapist, 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 {persona["age"]}-year-old in a therapy session. Stay completely in character.

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

SESSION FOCUS: {category_name} — {presentation}

DIRECTION FOR THIS TURN:
{direction}

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

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

What do you say next? Generate ONLY spoken words — no labels, no narration. Stay in character as someone your age."""

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


def generate_child_therapist_turn(
    persona, presentation, category, conversation, turn_num, style_profile, therapist_system_base
):
    """Generate therapist response for child/adolescent session."""
    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"Adapt it to the child's age and developmental level. "
        f"Weave it into the conversation — don't announce it."
    )

    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 = therapist_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_child_therapist_turn_validated(
    persona, presentation, category, conversation, turn_num, style_profile, therapist_system_base
):
    for attempt in range(STYLE_MAX_RETRIES):
        output = generate_child_therapist_turn(
            persona, presentation, category, conversation, turn_num, style_profile, therapist_system_base
        )
        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_child_session(category_key, category, persona, presentation, session_idx):
    """Generate a complete child/adolescent therapy session."""
    min_turns = category.get("min_turns", 10)
    max_turns = category.get("max_turns", 16)
    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)]]

    patient_system = category.get("patient_system", PATIENT_SYSTEM)
    therapist_system_base = category.get("therapist_system_base", PIXEL_SYSTEM_BASE)

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

        therapist_msg = generate_child_therapist_turn_validated(
            persona, presentation, category, conversation, turn, style_profile, therapist_system_base
        )
        conversation.append({"role": "assistant", "content": therapist_msg})

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


def main():
    parser = argparse.ArgumentParser(description="Child & Adolescent Therapy Session Generation")
    parser.add_argument("--categories", default="all", help="Comma-separated category keys or 'all'")
    parser.add_argument("--sessions-per-category", type=int, default=117)
    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/child_adolescent_sessions")
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / "child_adolescent_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=== CHILD & ADOLESCENT THERAPY 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()

    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}) ---")
        print(f"  {len(category['presentations'])} presentations × {len(category['patient_personas'])} personas")
        print(f"  Turns: {category.get('min_turns', 10)}-{category.get('max_turns', 16)}")

        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"child_adolescent_{cat_key}_{i:04d}"

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

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