decision/engine/training_data_generator.py

"""
Training Data Generator (Phase 3C).

Generates synthetic training samples from taxonomy trigger definitions
for model retraining and augmentation.

Each domain's triggers.yaml contains carefully curated phrases at
multiple confidence tiers. These can be converted into labeled training
examples to improve DriveHealthBERT's domain coverage.

Capabilities:
  - Generate ESCALATION samples from R2/R3 domain triggers
  - Generate domain-specific samples with label mapping
  - Generate negation examples (negative samples for safety)
  - Template-based augmentation with variations
  - JSONL output compatible with train.py

Usage:
    generator = TrainingDataGenerator(config)
    samples = generator.generate_all()
    generator.write_jsonl(samples, "data/taxonomy_augmented.jsonl")
"""

from __future__ import annotations

import json
import logging
import random
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple

from decision.engine.config_loader import DecisionConfigLoader

logger = logging.getLogger("decision.training_data_generator")

# Map domains to DriveHealthBERT labels
_DOMAIN_TO_LABEL = {
    # R3 domains → ESCALATION
    "chest_pain": "ESCALATION",
    "suicidal_ideation": "ESCALATION",
    "homicidal_ideation": "ESCALATION",
    "stroke_symptoms": "ESCALATION",
    "seizure": "ESCALATION",
    "severe_bleeding": "ESCALATION",
    "anaphylaxis": "ESCALATION",
    "airway_compromise": "ESCALATION",
    "overdose": "ESCALATION",
    "pregnancy_emergency": "ESCALATION",
    "ectopic_signal": "ESCALATION",
    "febrile_neutropenia": "ESCALATION",
    "meningitis_signal": "ESCALATION",
    "dka_hhs": "ESCALATION",
    "severe_hypoglycemia": "ESCALATION",
    "syncope": "ESCALATION",
    "thunderclap_headache": "ESCALATION",
    "sudden_vision_loss": "ESCALATION",
    "acute_limb_deficit": "ESCALATION",

    # R2 domains → ESCALATION (moderate, but still escalation)
    "shortness_of_breath_rest": "ESCALATION",
    "palpitations_dizziness": "ESCALATION",
    "heart_failure_worsening": "ESCALATION",
    "gi_bleed": "ESCALATION",
    "severe_abdominal_pain": "ESCALATION",
    "wound_infection": "ESCALATION",
    "fall_high_risk": "ESCALATION",
    "behavioral_crisis": "ESCALATION",
    "acute_confusion": "ESCALATION",
    "persistent_vomiting": "ESCALATION",
    "post_discharge_fever": "ESCALATION",
    "dvt_signal": "ESCALATION",
    "npo_violation": "ESCALATION",
    "anticoagulant_hold_failure": "ESCALATION",
    "acute_illness_preop": "ESCALATION",
    "worsening_depression_si": "ESCALATION",
    "procedure_anxiety_severe": "ESCALATION",
    "domestic_safety_concern": "ESCALATION",

    # R1/R0 symptom domains → SYMPTOM_CHECK
    "exertional_dyspnea": "SYMPTOM_CHECK",
    "persistent_dizziness": "SYMPTOM_CHECK",
    "abnormal_bp": "SYMPTOM_CHECK",
    "reproducible_chest_discomfort": "SYMPTOM_CHECK",
    "wound_concern": "SYMPTOM_CHECK",
    "uncontrolled_pain": "SYMPTOM_CHECK",
    "worsening_chronic_pain": "SYMPTOM_CHECK",
    "progressive_weakness": "SYMPTOM_CHECK",
    "urinary_retention": "SYMPTOM_CHECK",
    "severe_diarrhea": "SYMPTOM_CHECK",
    "poor_oral_intake": "SYMPTOM_CHECK",
    "new_jaundice": "SYMPTOM_CHECK",
    "glucose_out_of_range": "SYMPTOM_CHECK",
    "recurrent_falls": "SYMPTOM_CHECK",
    "post_procedure_incontinence": "SYMPTOM_CHECK",

    # Medication domains → MEDICATION
    "medication": "MEDICATION",
    "medication_nonadherence": "MEDICATION",

    # Scheduling domains → APPOINTMENT
    "scheduling": "APPOINTMENT",
    "scheduling_barrier": "APPOINTMENT",
    "missed_followup": "APPOINTMENT",
    "post_discharge_engagement": "APPOINTMENT",

    # Screening domains → SYMPTOM_CHECK
    "phq2_positive": "SYMPTOM_CHECK",
    "audit_c_positive": "SYMPTOM_CHECK",
    "screening_gap_identified": "SYMPTOM_CHECK",
    "screening_refusal": "GENERAL_INQUIRY",
    "cognitive_concern": "SYMPTOM_CHECK",

    # Social/safety domains
    "caregiver_safety_concern": "ESCALATION",
    "food_insecurity": "GENERAL_INQUIRY",
    "social_isolation": "GENERAL_INQUIRY",
    "tobacco_use_active": "GENERAL_INQUIRY",
}

# Templates for wrapping trigger phrases into natural sentences
_PATIENT_TEMPLATES = [
    "{phrase}",
    "I have {phrase}",
    "I'm experiencing {phrase}",
    "I've been having {phrase}",
    "I think I have {phrase}",
    "My {phrase} is getting worse",
    "I need help with {phrase}",
    "I'm worried about {phrase}",
    "I want to talk about {phrase}",
    "Can you help me with {phrase}",
]

_NEGATION_TEMPLATES = [
    "I don't have {phrase}",
    "No {phrase}",
    "I'm not experiencing {phrase}",
    "{phrase} has gone away",
    "I used to have {phrase} but not anymore",
    "My mother had {phrase} not me",
]


class TrainingDataGenerator:
    """Generates training data from taxonomy triggers."""

    def __init__(self, config: DecisionConfigLoader):
        self._config = config
        self._triggers = config.taxonomy_triggers
        self._rules = config.taxonomy_rules

    def generate_all(
        self,
        include_negations: bool = True,
        max_per_domain: int = 50,
        augment_templates: bool = True,
        seed: int = 42,
    ) -> List[Dict[str, str]]:
        """
        Generate training samples from all taxonomy domains.

        Args:
            include_negations: Also generate negative samples from negation patterns
            max_per_domain: Maximum samples per domain
            augment_templates: Apply template-based augmentation
            seed: Random seed for reproducibility

        Returns:
            List of {"text": ..., "label": ..., "source": ...} dicts
        """
        random.seed(seed)
        all_samples: List[Dict[str, str]] = []

        for domain_name, triggers in self._triggers.items():
            label = _DOMAIN_TO_LABEL.get(domain_name)
            if not label:
                logger.debug("Skipping domain %s (no label mapping)", domain_name)
                continue

            domain_samples = self._generate_for_domain(
                domain_name, triggers, label, augment_templates, max_per_domain
            )
            all_samples.extend(domain_samples)

            # Generate negation samples
            if include_negations and label == "ESCALATION":
                neg_samples = self._generate_negations(
                    domain_name, triggers, max_per_domain // 3
                )
                all_samples.extend(neg_samples)

        random.shuffle(all_samples)

        # Summary
        label_counts = {}
        for s in all_samples:
            label_counts[s["label"]] = label_counts.get(s["label"], 0) + 1

        logger.info(
            "Generated %d training samples from %d domains: %s",
            len(all_samples),
            len(self._triggers),
            label_counts,
        )

        return all_samples

    def _generate_for_domain(
        self,
        domain_name: str,
        triggers: Dict[str, Any],
        label: str,
        augment: bool,
        max_samples: int,
    ) -> List[Dict[str, str]]:
        """Generate training samples for a single domain."""
        samples = []
        lexical = triggers.get("lexical_signals", {})

        for tier_name in ("high", "medium", "low"):
            tier = lexical.get(tier_name, {})

            # Direct phrase samples
            for phrase in tier.get("phrases", []):
                samples.append({
                    "text": phrase,
                    "label": label,
                    "source": f"taxonomy/{domain_name}/{tier_name}/phrase",
                })

                # Template augmentation
                if augment:
                    templates = random.sample(
                        _PATIENT_TEMPLATES, min(3, len(_PATIENT_TEMPLATES))
                    )
                    for template in templates:
                        try:
                            augmented = template.format(phrase=phrase.lower())
                            if augmented != phrase:
                                samples.append({
                                    "text": augmented,
                                    "label": label,
                                    "source": f"taxonomy/{domain_name}/{tier_name}/augmented",
                                })
                        except (KeyError, IndexError):
                            pass

            # Partial samples (lower confidence)
            for partial in tier.get("partials", []):
                samples.append({
                    "text": partial,
                    "label": label,
                    "source": f"taxonomy/{domain_name}/{tier_name}/partial",
                })

        # Deduplicate by text
        seen = set()
        unique = []
        for s in samples:
            text_lower = s["text"].lower().strip()
            if text_lower not in seen:
                seen.add(text_lower)
                unique.append(s)

        # Cap per domain
        if len(unique) > max_samples:
            unique = random.sample(unique, max_samples)

        return unique

    def _generate_negations(
        self,
        domain_name: str,
        triggers: Dict[str, Any],
        max_samples: int,
    ) -> List[Dict[str, str]]:
        """
        Generate negative samples from negation patterns.

        These become NON-ESCALATION training examples to reduce false positives.
        """
        samples = []
        negation = triggers.get("negation_handling", {})
        patterns = negation.get("patterns", [])

        # Use negation patterns directly as negative examples
        for pattern in patterns:
            samples.append({
                "text": pattern,
                "label": "GENERAL_INQUIRY",  # Negated safety = not escalation
                "source": f"taxonomy/{domain_name}/negation",
            })

        # Apply negation templates to domain phrases
        lexical = triggers.get("lexical_signals", {})
        high_phrases = lexical.get("high", {}).get("phrases", [])

        for phrase in high_phrases[:5]:  # Top 5 high-confidence phrases
            templates = random.sample(
                _NEGATION_TEMPLATES, min(2, len(_NEGATION_TEMPLATES))
            )
            for template in templates:
                try:
                    negated = template.format(phrase=phrase.lower())
                    samples.append({
                        "text": negated,
                        "label": "GENERAL_INQUIRY",
                        "source": f"taxonomy/{domain_name}/negation_augmented",
                    })
                except (KeyError, IndexError):
                    pass

        if len(samples) > max_samples:
            samples = random.sample(samples, max_samples)

        return samples

    def write_jsonl(
        self, samples: List[Dict[str, str]], output_path: str
    ) -> int:
        """Write samples to JSONL file compatible with train.py."""
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)

        count = 0
        with open(path, "w", encoding="utf-8") as f:
            for sample in samples:
                record = {
                    "text": sample["text"],
                    "label": sample["label"],
                }
                f.write(json.dumps(record, ensure_ascii=False) + "\n")
                count += 1

        logger.info("Wrote %d samples to %s", count, path)
        return count

    def get_label_distribution(
        self, samples: List[Dict[str, str]]
    ) -> Dict[str, int]:
        """Get label distribution of generated samples."""
        dist: Dict[str, int] = {}
        for s in samples:
            dist[s["label"]] = dist.get(s["label"], 0) + 1
        return dict(sorted(dist.items(), key=lambda x: x[1], reverse=True))

    def get_domain_coverage_report(
        self, samples: List[Dict[str, str]]
    ) -> str:
        """Generate a coverage report."""
        domain_count: Dict[str, int] = {}
        for s in samples:
            source = s.get("source", "")
            domain = source.split("/")[1] if "/" in source else "unknown"
            domain_count[domain] = domain_count.get(domain, 0) + 1

        lines = ["=== Training Data Coverage Report ===", ""]
        lines.append(f"Total samples: {len(samples)}")
        lines.append(f"Domains covered: {len(domain_count)}")
        lines.append("")

        dist = self.get_label_distribution(samples)
        lines.append("Label distribution:")
        for label, count in dist.items():
            pct = count / len(samples) * 100
            lines.append(f"  {label:20s}: {count:5d} ({pct:.1f}%)")

        lines.append("")
        lines.append("Top domains by sample count:")
        for domain, count in sorted(domain_count.items(), key=lambda x: x[1], reverse=True)[:15]:
            lines.append(f"  {domain:35s}: {count:4d}")

        return "\n".join(lines)