from __future__ import annotations

import json
import sys

from dagplanner import ExposureGuardDAGPlanner


def demo_single():
    planner = ExposureGuardDAGPlanner(risk_threshold=0.20, budget=1.0)
    result = planner.plan(
        risk_score=0.74,
        retok_prob=0.68,
        active_modalities=["text", "image", "ehr"],
        patient_id="P-0042",
    )
    print(planner.summary(result))
    print(json.dumps(result, indent=2))
    return result


def demo_batch():
    planner = ExposureGuardDAGPlanner(
        risk_threshold=0.20,
        budget=1.2,
        modality_weights={"image": 1.4, "audio": 1.2},
    )
    records = [
        {"patient_id": "P-001", "risk_score": 0.82, "retok_prob": 0.71,
         "active_modalities": ["text", "image", "audio", "ehr"]},
        {"patient_id": "P-002", "risk_score": 0.45, "retok_prob": 0.30,
         "active_modalities": ["text", "ehr"]},
        {"patient_id": "P-003", "risk_score": 0.15, "retok_prob": 0.10,
         "active_modalities": ["ehr"]},
        {"patient_id": "P-004", "risk_score": 0.61, "retok_prob": 0.58,
         "active_modalities": ["text", "audio"]},
    ]
    results = planner.batch_plan(records)
    for r in results:
        print(planner.summary(r))
    return results


def demo_pipeline():
    fedcrdt_out = {"patient_id": "P-0099", "risk_score": 0.67, "retok_prob": 0.61}
    encoder_out = {"risk_score": 0.70}
    merged_risk = 0.5 * fedcrdt_out["risk_score"] + 0.5 * encoder_out["risk_score"]

    planner = ExposureGuardDAGPlanner(risk_threshold=0.20, budget=1.0)
    result = planner.plan(
        risk_score=merged_risk,
        retok_prob=fedcrdt_out["retok_prob"],
        active_modalities=["text", "image", "ehr"],
        patient_id=fedcrdt_out["patient_id"],
    )
    print("\n--- Pipeline demo ---")
    print(planner.summary(result))
    return result


if __name__ == "__main__":
    print("=== Single ===")
    demo_single()
    print("\n=== Batch ===")
    demo_batch()
    demo_pipeline()