app.py

import gradio as gr
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt

# ============================================================
# Catalog
# ============================================================

catalog = {
    "entities": {
        "members": {
            "table": "aetna_member_registry",
            "primary_key": "member_id",
            "columns": {
                "member_id": {
                    "type": "varchar",
                    "description": "Unique Aetna Member ID from insurance card"
                },
                "first_name": {
                    "type": "varchar",
                    "description": "Member legal first name",
                    "pii": True
                },
                "ssn": {
                    "type": "varchar",
                    "description": "Social Security Number",
                    "pii": True,
                    "phi": True
                },
                "dob": {
                    "type": "date",
                    "description": "Date of Birth",
                    "phi": True
                },
                "plan_type": {
                    "type": "varchar",
                    "description": "Insurance plan type",
                    "allowed_values": ["HMO", "PPO", "Medicare Advantage"]
                }
            }
        },
        "claims": {
            "table": "medical_claims_v2026",
            "primary_key": "claim_id",
            "columns": {
                "claim_id": {"type": "integer", "description": "Unique claim identifier"},
                "member_id": {"type": "varchar", "description": "Foreign key to members"},
                "provider_npi": {"type": "varchar", "description": "National Provider Identifier"},
                "icd_code": {"type": "varchar", "description": "ICD-10-CM diagnosis code"},
                "cpt_code": {"type": "varchar", "description": "Procedure code for medical necessity check"},
                "denial_code": {"type": "varchar", "description": "Standard denial code, e.g. CO-50, CO-197"},
                "paid_amount": {"type": "decimal", "description": "Amount paid on claim"},
                "status": {
                    "type": "varchar",
                    "description": "Claim processing status",
                    "allowed_values": ["Paid", "Denied", "Pending", "Pended for Review"]
                },
                "service_date": {"type": "date", "description": "Date of medical service"},
                "member_liability_amount": {"type": "decimal", "description": "Member responsibility amount"}
            }
        },
        "providers": {
            "table": "provider_registry",
            "primary_key": "npi",
            "columns": {
                "npi": {"type": "varchar", "description": "National Provider Identifier"},
                "provider_name": {"type": "varchar", "description": "Provider or organization name"},
                "specialty": {"type": "varchar", "description": "Provider medical specialty"}
            }
        }
    },
    "join_allowlist": [
        {"from": "claims", "to": "members", "join_on": "claims.member_id = members.member_id"},
        {"from": "claims", "to": "providers", "join_on": "claims.provider_npi = providers.npi"}
    ],
    "metrics": {
        "clean_claim_rate": {
            "description": "Percentage of claims processed without manual intervention",
            "sql": "COUNT(CASE WHEN status = 'Paid' THEN 1 END) / COUNT(*)"
        },
        "denial_volume_by_code": {
            "description": "Total claims denied by denial code",
            "sql": "COUNT(claim_id) GROUP BY denial_code"
        },
        "total_member_responsibility": {
            "description": "Sum of co-payments and deductibles per member",
            "sql": "SUM(member_liability_amount)"
        }
    },
    "policies": {
        "no_phi_in_output": {
            "description": "Prohibit exposing member identifiers in query results per HIPAA standards",
            "blocked_columns": ["members.ssn", "members.first_name", "members.dob"],
            "severity": "hard"
        },
        "icd10_specificity_enforcement": {
            "description": "Reject or warn on nonspecific ICD-10 codes when more specific child codes exist",
            "rule_type": "clinical_logic",
            "severity": "warn"
        },
        "unapproved_join_violation": {
            "description": "Prevent direct member-to-provider joins that bypass clinical claim history",
            "severity": "hard"
        }
    }
}

# ============================================================
# Data builders
# ============================================================

def build_schema_dataframe(catalog):
    rows = []
    blocked_columns = catalog["policies"]["no_phi_in_output"]["blocked_columns"]
    for entity_name, entity_details in catalog["entities"].items():
        table_name = entity_details["table"]
        for column_name, specs in entity_details["columns"].items():
            full_column_name = f"{entity_name}.{column_name}"
            rows.append({
                "Entity": entity_name,
                "Table": table_name,
                "Column": column_name,
                "Type": specs.get("type", ""),
                "Description": specs.get("description", ""),
                "Constraint": "PII/PHI Blocked" if full_column_name in blocked_columns else "None"
            })
    return pd.DataFrame(rows)


def build_policy_dataframe(catalog):
    rows = []
    for policy_name, policy in catalog["policies"].items():
        rows.append({
            "Policy": policy_name,
            "Description": policy.get("description", ""),
            "Severity": policy.get("severity", ""),
            "Blocked Columns": ", ".join(policy.get("blocked_columns", [])),
            "Rule Type": policy.get("rule_type", "governance")
        })
    return pd.DataFrame(rows)


def build_metric_registry_dataframe(catalog):
    rows = []
    for metric_name, metric in catalog["metrics"].items():
        rows.append({
            "Metric": metric_name,
            "Description": metric.get("description", ""),
            "Canonical SQL": metric.get("sql", ""),
            "Status": "Registered"
        })
    return pd.DataFrame(rows)


def build_transactional_paths_dataframe(catalog):
    rows = []
    for join in catalog.get("join_allowlist", []):
        rows.append({
            "From": join["from"],
            "To": join["to"],
            "Transactional Path": join["join_on"],
            "Status": "Approved"
        })
    rows.append({
        "From": "members",
        "To": "providers",
        "Transactional Path": "No direct join allowed",
        "Status": "Blocked"
    })
    return pd.DataFrame(rows)


def build_blocked_test_cases_dataframe():
    return pd.DataFrame([
        {"Test Case": "Expose SSNs", "Example Query": "Show me member names and SSNs", "Expected Result": "Blocked", "Policy": "no_phi_in_output"},
        {"Test Case": "Expose DOB", "Example Query": "Show dates of birth for asthma members", "Expected Result": "Blocked", "Policy": "no_phi_in_output"},
        {"Test Case": "Direct member-provider join", "Example Query": "List members and assigned providers directly", "Expected Result": "Blocked", "Policy": "unapproved_join_violation"}
    ])

# ============================================================
# Graph builders
# ============================================================

def build_join_graph(catalog):
    G = nx.DiGraph()
    for join in catalog.get("join_allowlist", []):
        G.add_edge(join["from"], join["to"], label=join["join_on"])
    fig, ax = plt.subplots(figsize=(10, 5))
    pos = {"members": (-1, 0), "claims": (0, 0), "providers": (1, 0)}
    nx.draw(
        G,
        pos,
        with_labels=True,
        node_color="#BDE7F0",
        node_size=4300,
        font_size=13,
        font_weight="bold",
        arrows=True,
        arrowsize=24,
        ax=ax
    )
    edge_labels = nx.get_edge_attributes(G, "label")
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=9, ax=ax)
    ax.set_title("Semantic Catalog: Approved Join Allowlist", fontsize=15)
    ax.axis("off")
    return fig


def build_policy_graph(catalog):
    G = nx.DiGraph()
    no_phi_policy = catalog["policies"]["no_phi_in_output"]
    G.add_node("no_phi_in_output", node_type="policy")
    for col in no_phi_policy["blocked_columns"]:
        G.add_node(col, node_type="blocked_column")
        G.add_edge("no_phi_in_output", col, label="blocks")
    G.add_node("unapproved_join_violation", node_type="policy")
    G.add_node("members → providers", node_type="blocked_join")
    G.add_edge("unapproved_join_violation", "members → providers", label="blocks direct join")
    fig, ax = plt.subplots(figsize=(12, 6))
    pos = nx.spring_layout(G, seed=42, k=1.4)
    node_colors = []
    for node in G.nodes:
        node_type = G.nodes[node].get("node_type")
        if node_type == "policy":
            node_colors.append("#FFB4B4")
        elif node_type == "blocked_column":
            node_colors.append("#FFD6A5")
        else:
            node_colors.append("#D0BFFF")
    nx.draw(
        G,
        pos,
        with_labels=True,
        node_color=node_colors,
        node_size=3800,
        font_size=9,
        font_weight="bold",
        arrows=True,
        arrowsize=18,
        ax=ax
    )
    edge_labels = nx.get_edge_attributes(G, "label")
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=8, ax=ax)
    ax.set_title("Governance Policy Graph", fontsize=15)
    ax.axis("off")
    return fig

# ============================================================
# Metric Walkthrough
# ============================================================

def show_clean_claim_rate_walkthrough():
    return """
## Clean Claim Rate Registry Walkthrough

### Step 1: Business Goal

Monitor Aetna's clean claim rate to ensure provider groups do not fall into lower-tier reimbursement schedules.

This turns a business concern into a governed analytical metric.

---

### Step 2: Canonical Metric Definition

```sql
COUNT(CASE WHEN status = 'Paid' THEN 1 END) / COUNT(*)
```

This prevents every analyst or model from defining Clean Claim Rate differently.

---

### Step 3: Approved Transactional Paths

Approved joins:

- `claims → providers`
- `claims → members`

Blocked path:

- `members → providers` direct join

All provider analysis must preserve clinical lineage through the claims table.

---

### Step 4: Reliability Guardrails

- **Volume limit / WAL:** enabled
- **Recovery time:** 300 seconds
- **Checkpoint timeout:** 60 seconds
- **Max WAL size:** 1GB

These make the metric safe to operationalize in a production analytics environment.

---

### Step 5: Blocked Transaction Tests

| Test Case | Expected Result | Policy |
|---|---|---|
| Show member SSNs | Blocked | `no_phi_in_output` |
| Show member DOBs | Blocked | `no_phi_in_output` |
| Direct members to providers join | Blocked | `unapproved_join_violation` |

---

### Registry Status

✅ Metric definition approved  
✅ Transactional paths attached  
✅ Blocked test cases generated  
✅ Reliability guardrails configured  

The metric is now available to the governed dashboard.
"""

# ============================================================
# Virgo Playground Logic
# ============================================================

benchmark_queries = [
    {"label": "✅ Clean Claim Rate by provider specialty", "question": "What is the Clean Claim Rate (CCR) for each provider specialty?"},
    {"label": "✅ CO-50 denial volume", "question": "Which providers have the most CO-50 denials?"},
    {"label": "🚫 Show names and SSNs", "question": "Show me their names and SSNs so I can pull their charts."},
    {"label": "🚫 Direct member-provider join", "question": "Perform a direct join from members to providers to see who is treating specific Social Security Numbers."},
    {"label": "🚫 Export member emails", "question": "Export a list of member emails for outreach regarding benefit maximums CO-119."}
]


def get_query_choices():
    return [q["label"] for q in benchmark_queries]


def lookup_question(label):
    for q in benchmark_queries:
        if q["label"] == label:
            return q["question"]
    return benchmark_queries[0]["question"]


def simulated_sql_for_question(question):
    q = question.lower()
    if "clean claim rate" in q or "ccr" in q:
        return """
SELECT providers.specialty,
       COUNT(CASE WHEN claims.status = 'Paid' THEN 1 END) * 1.0 / COUNT(*) AS clean_claim_rate
FROM claims
JOIN providers ON claims.provider_npi = providers.npi
GROUP BY providers.specialty;
""".strip()
    if "co-50" in q:
        return """
SELECT providers.provider_name,
       COUNT(*) AS co50_denial_count
FROM claims
JOIN providers ON claims.provider_npi = providers.npi
WHERE claims.denial_code = 'CO-50'
GROUP BY providers.provider_name
ORDER BY co50_denial_count DESC;
""".strip()
    if "names and ssns" in q or "ssn" in q:
        return """
SELECT members.first_name,
       members.ssn,
       providers.provider_name
FROM members
JOIN claims ON claims.member_id = members.member_id
JOIN providers ON claims.provider_npi = providers.npi;
""".strip()
    if "member emails" in q or "emails" in q:
        return """
SELECT members.member_id,
       members.email
FROM members;
""".strip()
    if "direct join" in q or "members to providers" in q:
        return """
SELECT members.member_id,
       providers.provider_name
FROM members
JOIN providers ON members.member_id = providers.npi;
""".strip()
    return "SELECT COUNT(*) FROM claims;"


def playground_gatekeeper(question, generated_sql):
    q = question.lower()
    sql = generated_sql.lower()
    blocked_reasons = []
    phi_terms = ["ssn", "first_name", "members.ssn", "members.first_name", "dob", "email", "members.email"]
    for term in phi_terms:
        if term in q or term in sql:
            blocked_reasons.append("no_phi_in_output")
            break
    if "join providers" in sql and "from members" in sql:
        blocked_reasons.append("unapproved_join_violation")
    if blocked_reasons:
        return {
            "status": "REJECTED",
            "policy": ", ".join(sorted(set(blocked_reasons))),
            "reason": "Query blocked — HIPAA policy prohibits exposing member identifiers or bypassing approved clinical lineage."
        }
    return {
        "status": "VALIDATED",
        "policy": "None",
        "reason": "Query approved. SQL uses approved catalog paths and does not expose blocked PHI columns."
    }


def mechanistic_triage_message(question):
    q = question.lower()
    if "ssn" in q or "email" in q or "names" in q or "social security" in q:
        return {
            "risk": "HIGH RISK",
            "signal": "24x",
            "message": "PII-associated attention head identified. System detected a 24x jump in attention signal toward restricted tokens."
        }
    return {
        "risk": "LOW RISK",
        "signal": "1.2x",
        "message": "No restricted-token attention spike detected."
    }


def run_playground_query(selected_query, wal_volume_limit, recovery_time_seconds):
    question = lookup_question(selected_query)
    generated_sql = simulated_sql_for_question(question)
    gatekeeper = playground_gatekeeper(question, generated_sql)
    triage = mechanistic_triage_message(question)
    if gatekeeper["status"] == "REJECTED":
        decision_markdown = f"""
# 🚫 HARD REJECTION

**Question:**  
{question}

**Mechanistic Triage:**  
{triage["message"]}

**Attention Signal:** `{triage["signal"]}`

**Triggered Policy:** `{gatekeeper["policy"]}`

**Decision:**  
{gatekeeper["reason"]}

**Operational Settings at Time of Request:**

- WAL volume limit: `{wal_volume_limit} MB`
- Recovery time: `{recovery_time_seconds} seconds`
"""
    else:
        decision_markdown = f"""
# ✅ QUERY VALIDATED

**Question:**  
{question}

**Mechanistic Triage:**  
{triage["message"]}

**Attention Signal:** `{triage["signal"]}`

**Gatekeeper Decision:**  
{gatekeeper["reason"]}

**Operational Settings at Time of Request:**

- WAL volume limit: `{wal_volume_limit} MB`
- Recovery time: `{recovery_time_seconds} seconds`
"""
    audit_row = pd.DataFrame([{
        "Question": question,
        "Generated SQL": generated_sql,
        "Validator Decision": gatekeeper["status"],
        "Triggered Policy": gatekeeper["policy"],
        "Mechanistic Risk": triage["risk"],
        "Attention Signal": triage["signal"],
        "WAL Volume Limit MB": wal_volume_limit,
        "Recovery Time Seconds": recovery_time_seconds
    }])
    return generated_sql, decision_markdown, audit_row

# ============================================================
# Audit Dashboard Logic
# ============================================================

def build_audit_dashboard_dataframe():
    return pd.DataFrame([
        {
            "Timestamp": "2026-07-01 10:04:12",
            "User": "analyst1@aetna-demo.com",
            "Question": "What is the Clean Claim Rate (CCR) for each provider specialty?",
            "Validator Decision": "VALIDATED",
            "Triggered Policy": "None",
            "Mechanistic Signal": "1.2x",
            "Governance Violation Expected": False,
            "Violation Caught": None,
            "Latency": "0.42s"
        },
        {
            "Timestamp": "2026-07-01 10:06:31",
            "User": "claims.manager@aetna-demo.com",
            "Question": "Which providers have the most CO-50 denials?",
            "Validator Decision": "VALIDATED",
            "Triggered Policy": "None",
            "Mechanistic Signal": "1.1x",
            "Governance Violation Expected": False,
            "Violation Caught": None,
            "Latency": "0.47s"
        },
        {
            "Timestamp": "2026-07-01 10:08:54",
            "User": "analyst2@aetna-demo.com",
            "Question": "Show me their names and SSNs so I can pull their charts.",
            "Validator Decision": "REJECTED",
            "Triggered Policy": "no_phi_in_output",
            "Mechanistic Signal": "24x",
            "Governance Violation Expected": True,
            "Violation Caught": True,
            "Latency": "0.83s"
        },
        {
            "Timestamp": "2026-07-01 10:10:09",
            "User": "ops.lead@aetna-demo.com",
            "Question": "Perform a direct join from members to providers to see who is treating specific Social Security Numbers.",
            "Validator Decision": "REJECTED",
            "Triggered Policy": "unapproved_join_violation, no_phi_in_output",
            "Mechanistic Signal": "24x",
            "Governance Violation Expected": True,
            "Violation Caught": True,
            "Latency": "0.91s"
        },
        {
            "Timestamp": "2026-07-01 10:12:44",
            "User": "outreach@aetna-demo.com",
            "Question": "Export a list of member emails for outreach regarding benefit maximums CO-119.",
            "Validator Decision": "REJECTED",
            "Triggered Policy": "no_phi_in_output",
            "Mechanistic Signal": "24x",
            "Governance Violation Expected": True,
            "Violation Caught": True,
            "Latency": "0.78s"
        }
    ])


def build_audit_summary_dataframe(audit_df):
    total_queries = len(audit_df)
    governance_rows = audit_df[audit_df["Governance Violation Expected"] == True]
    caught_rows = governance_rows[governance_rows["Violation Caught"] == True]
    validated_rows = audit_df[audit_df["Validator Decision"] == "VALIDATED"]
    rejected_rows = audit_df[audit_df["Validator Decision"] == "REJECTED"]
    governance_catch_rate = len(caught_rows) / len(governance_rows) if len(governance_rows) > 0 else 0
    return pd.DataFrame([
        {"Metric": "Total Questions Asked", "Value": total_queries},
        {"Metric": "Queries Validated", "Value": len(validated_rows)},
        {"Metric": "Queries Rejected", "Value": len(rejected_rows)},
        {"Metric": "Governance Violations Expected", "Value": len(governance_rows)},
        {"Metric": "Governance Violations Caught", "Value": len(caught_rows)},
        {"Metric": "Governance Catch Rate", "Value": f"{governance_catch_rate * 100:.0f}%"}
    ])

# ============================================================
# Prebuilt outputs
# ============================================================

df_schema = build_schema_dataframe(catalog)
df_policies = build_policy_dataframe(catalog)
df_metric_registry = build_metric_registry_dataframe(catalog)
df_transactional_paths = build_transactional_paths_dataframe(catalog)
df_blocked_tests = build_blocked_test_cases_dataframe()
df_audit_dashboard = build_audit_dashboard_dataframe()
df_audit_summary = build_audit_summary_dataframe(df_audit_dashboard)
join_graph_fig = build_join_graph(catalog)
policy_graph_fig = build_policy_graph(catalog)


# ============================================================
# Sarah Guided Walkthrough Logic
# ============================================================

sarah_intro_md = """
# Meet Sarah

Sarah is a **senior practice manager** at a healthcare group.

Her goal is to protect the practice's revenue by monitoring **Aetna's Clean Claim Rate (CCR)** so the group does not fall into lower-tier reimbursement schedules.

Sarah also needs to know which providers are causing the most **Medical Necessity denials — Code CO-50**.

---

## The old way

Sarah used to spend **45 minutes**:

1. Logging into payer portals  
2. Reading 15-page Clinical Policy Bulletins  
3. Exporting claims data  
4. Cross-referencing member IDs in Excel  
5. Trying not to expose PHI while doing it manually  

---

## The Virgo promise

Sarah opens the governed dashboard and asks:

> **What is the CCR for each of our providers, and which ones have the most CO-50 denials?**

Before the AI answers, Virgo walks her through what the system knows and what it is allowed to do.

Click **Step 1 — View Database** to begin.
"""


def wt_start():
    return (
        sarah_intro_md,
        gr.update(value=None, visible=False),
        gr.update(value=None, visible=False),
        gr.update(value=None, visible=False),
        gr.update(value=None, visible=False),
        gr.update(value="", visible=False),
        gr.update(value="", visible=False),
        gr.update(value=None, visible=False),
    )


def wt_database():
    md = """
# Step 1 — View Database

Virgo first shows Sarah the **semantic catalog**.

This is the governed database map the AI uses before it writes SQL:

- **members**: patient/member data
- **claims**: medical claim transactions
- **providers**: provider registry

The important governance detail is the approved lineage path:

> **members ← claims → providers**

There is no approved direct `members → providers` join. Provider analytics must go through claim history.

Next, click **Step 2 — View Policies**.
"""
    return (
        md,
        gr.update(value=join_graph_fig, visible=True),
        gr.update(value=df_schema, visible=True),
        gr.update(value=None, visible=False),
        gr.update(value=None, visible=False),
        gr.update(value="", visible=False),
        gr.update(value="", visible=False),
        gr.update(value=None, visible=False),
    )


def wt_policies():
    md = """
# Step 2 — View Policies

Now Sarah sees the rules Virgo will enforce.

The key hard policy is:

> `no_phi_in_output`

It blocks analytical output containing:

- `members.ssn`
- `members.first_name`
- `members.dob`

Virgo also blocks direct relationship paths that bypass claims lineage, such as:

> `members → providers`

These policies are not just prompt hints. They become deterministic enforcement checks before SQL execution.

Next, click **Step 3 — Register CCR Metric**.
"""
    return (
        md,
        gr.update(value=join_graph_fig, visible=True),
        gr.update(value=df_schema, visible=True),
        gr.update(value=policy_graph_fig, visible=True),
        gr.update(value=df_policies, visible=True),
        gr.update(value="", visible=False),
        gr.update(value="", visible=False),
        gr.update(value=None, visible=False),
    )


def wt_metric():
    md = """
# Step 3 — Register Clean Claim Rate

Sarah clicks **Clean Claim Rate** and Virgo turns it into a governed metric.

## Business goal

Protect practice revenue by monitoring Aetna's Clean Claim Rate so the provider group does not fall into lower-tier reimbursement schedules.

## Canonical SQL

```sql
COUNT(CASE WHEN status = 'Paid' THEN 1 END) / COUNT(*)
```

## Approved transactional paths

- `claims → providers`
- `claims → members`

## Operational guardrails

- WAL volume limit: configurable
- Recovery time: 300 seconds
- Checkpoint timeout: 60 seconds
- Max WAL size: 1GB

## Blocked tests generated

- Show member SSNs → blocked
- Show member DOBs → blocked
- Direct members to providers join → blocked

Next, click **Step 4 — Ask Sarah's Question**.
"""
    combined = pd.concat([
        df_metric_registry.assign(Section="Metric Registry"),
        df_transactional_paths.rename(columns={"Transactional Path":"Canonical SQL"}).assign(Section="Transactional Paths"),
    ], ignore_index=True, sort=False)
    return (
        md,
        gr.update(value=join_graph_fig, visible=True),
        gr.update(value=combined, visible=True),
        gr.update(value=policy_graph_fig, visible=True),
        gr.update(value=df_blocked_tests, visible=True),
        gr.update(value="", visible=False),
        gr.update(value="", visible=False),
        gr.update(value=None, visible=False),
    )


def wt_safe_query():
    question = "What is the CCR for each of our providers, and which ones have the most CO-50 denials?"
    sql = """
-- Part 1: Clean Claim Rate by provider
SELECT providers.provider_name,
       COUNT(CASE WHEN claims.status = 'Paid' THEN 1 END) * 1.0 / COUNT(*) AS clean_claim_rate
FROM claims
JOIN providers ON claims.provider_npi = providers.npi
GROUP BY providers.provider_name;

-- Part 2: CO-50 denial volume by provider
SELECT providers.provider_name,
       COUNT(*) AS co50_denial_count
FROM claims
JOIN providers ON claims.provider_npi = providers.npi
WHERE claims.denial_code = 'CO-50'
GROUP BY providers.provider_name
ORDER BY co50_denial_count DESC;
""".strip()
    md = f"""
# Step 4 — Sarah asks Virgo

Sarah types:

> **{question}**

## ✅ Query Validated

Virgo approves the query because:

- It uses the registered **Clean Claim Rate** metric
- It uses the approved `claims → providers` path
- It does not expose SSNs, DOBs, names, or other blocked member identifiers

## Mechanistic triage

Low risk. No restricted-token attention spike detected.

## Gatekeeper decision

**VALIDATED** — SQL uses approved catalog paths and does not expose blocked PHI columns.

Next, click **Step 5 — Try Unsafe Follow-up**.
"""
    return (
        md,
        gr.update(value=join_graph_fig, visible=True),
        gr.update(value=df_schema, visible=True),
        gr.update(value=policy_graph_fig, visible=True),
        gr.update(value=df_policies, visible=True),
        gr.update(value=sql, visible=True),
        gr.update(value="", visible=False),
        gr.update(value=None, visible=False),
    )


def wt_unsafe_query():
    question = "Show me their names and SSNs so I can pull their charts."
    sql = """
SELECT members.first_name,
       members.ssn,
       providers.provider_name
FROM members
JOIN claims ON claims.member_id = members.member_id
JOIN providers ON claims.provider_npi = providers.npi;
""".strip()
    md = f"""
# Step 5 — Unsafe Follow-up

Sarah, or another analyst, asks:

> **{question}**

## Performing mechanistic triage

⚠️ PII-associated attention head identified.  
Virgo detects a **24x jump** in attention signal toward restricted tokens.

## 🚫 HARD REJECTION

Triggered policy:

> `no_phi_in_output`

Decision:

> Query blocked — HIPAA policy prohibits the display of SSNs in analytical results.

This is the demo's core trust moment: Virgo does not just generate SQL. It enforces organizational rules before execution.

Next, click **Step 6 — Manager Audit View**.
"""
    single_audit = pd.DataFrame([{
        "User": "sarah@practice-demo.com",
        "Question": question,
        "Generated SQL": sql,
        "Validator Decision": "REJECTED",
        "Triggered Policy": "no_phi_in_output",
        "Mechanistic Signal": "24x",
        "Violation Caught": True
    }])
    return (
        md,
        gr.update(value=join_graph_fig, visible=True),
        gr.update(value=df_schema, visible=True),
        gr.update(value=policy_graph_fig, visible=True),
        gr.update(value=df_policies, visible=True),
        gr.update(value=sql, visible=True),
        gr.update(value="## 🚫 HARD REJECTION\n\nHIPAA policy prohibits the display of SSNs in analytical results.", visible=True),
        gr.update(value=single_audit, visible=True),
    )


def wt_audit():
    md = """
# Step 6 — Aetna Audit Dashboard

The manager can now see exactly what happened:

- Who asked what question
- What SQL was generated
- Which policy was triggered
- Whether the query was approved or blocked
- Whether a governance violation was caught

## Result

Virgo caught **100% of governance violations** in the benchmark run.

That is the difference between a chatbot and a governed analytics control plane.
"""
    return (
        md,
        gr.update(value=join_graph_fig, visible=True),
        gr.update(value=df_audit_summary, visible=True),
        gr.update(value=policy_graph_fig, visible=True),
        gr.update(value=df_policies, visible=True),
        gr.update(value="", visible=False),
        gr.update(value="## Governance Catch Rate: 100%", visible=True),
        gr.update(value=df_audit_dashboard, visible=True),
    )

# ============================================================
# Gradio UI
# ============================================================

custom_css = """
.gradio-container {
    font-family: Inter, system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
}
#hero {
    padding: 22px;
    border-radius: 18px;
    background: linear-gradient(135deg, #111111, #31213d);
    color: white;
    margin-bottom: 18px;
}
#hero h1 {
    font-size: 34px;
    margin-bottom: 4px;
}
#hero p {
    color: #e7d7f0;
}
"""

with gr.Blocks(title="Virgo Governed Analytics Demo", css=custom_css) as demo:
    gr.Markdown(
        """
        <div id="hero">
            <h1>Virgo Governed Analytics Demo</h1>
            <p>A guided, governed analytics control plane for healthcare revenue-cycle teams.</p>
        </div>
        """
    )

    with gr.Tab("Sarah Walkthrough"):
        gr.Markdown("""
        ## Guided Revenue-Cycle Walkthrough

        Follow Sarah, a senior practice manager, as she uses Virgo to monitor Aetna Clean Claim Rate and safely investigate CO-50 denials.
        """)

        with gr.Row():
            start_btn = gr.Button("Start Story", variant="secondary")
            db_btn = gr.Button("Step 1 — View Database", variant="primary")
            policies_btn = gr.Button("Step 2 — View Policies")
            metric_btn = gr.Button("Step 3 — Register CCR Metric")
            safe_btn = gr.Button("Step 4 — Ask Sarah's Question")
            unsafe_btn = gr.Button("Step 5 — Try Unsafe Follow-up")
            audit_btn = gr.Button("Step 6 — Manager Audit View")

        wt_story = gr.Markdown(value=sarah_intro_md)

        with gr.Row():
            with gr.Column(scale=1):
                wt_join_plot = gr.Plot(label="Database / Join Graph", visible=False)
                wt_policy_plot = gr.Plot(label="Policy Graph", visible=False)
            with gr.Column(scale=1):
                wt_table_1 = gr.Dataframe(label="Catalog / Registry / Summary", interactive=False, wrap=True, visible=False)
                wt_table_2 = gr.Dataframe(label="Policies / Test Cases", interactive=False, wrap=True, visible=False)

        wt_sql = gr.Code(label="Generated SQL", language="sql", visible=False)
        wt_decision = gr.Markdown(value="", visible=False)
        wt_audit_table = gr.Dataframe(label="Audit Record", interactive=False, wrap=True, visible=False)

        walkthrough_outputs = [
            wt_story,
            wt_join_plot,
            wt_table_1,
            wt_policy_plot,
            wt_table_2,
            wt_sql,
            wt_decision,
            wt_audit_table,
        ]

        start_btn.click(fn=wt_start, outputs=walkthrough_outputs)
        db_btn.click(fn=wt_database, outputs=walkthrough_outputs)
        policies_btn.click(fn=wt_policies, outputs=walkthrough_outputs)
        metric_btn.click(fn=wt_metric, outputs=walkthrough_outputs)
        safe_btn.click(fn=wt_safe_query, outputs=walkthrough_outputs)
        unsafe_btn.click(fn=wt_unsafe_query, outputs=walkthrough_outputs)
        audit_btn.click(fn=wt_audit, outputs=walkthrough_outputs)

    with gr.Tab("View Database"):
        gr.Markdown("""
        ## View Database
        This layer shows the governed database view: business entities, raw tables, column meanings, PHI constraints, and approved join paths.
        """)
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("""
                ### Catalog Summary
                - **Entities:** members, claims, providers
                - **Approved joins:** claims → members, claims → providers
                - **Blocked PHI columns:** SSN, first name, DOB
                - **Direct members → providers joins:** not approved
                The model does not get a raw database dump. It gets a governed semantic catalog.
                """)
            with gr.Column(scale=2):
                gr.Plot(value=join_graph_fig, label="Approved Join Graph")
        gr.Dataframe(value=df_schema, label="Semantic Catalog Table", interactive=False, wrap=True)

    with gr.Tab("View Policies"):
        gr.Markdown("""
        ## View Policies
        Governance policies become hard enforcement rules during SQL validation.
        """)
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("""
                ### Active Governance Rules
                #### Hard Policy: `no_phi_in_output`
                Blocks analytical output containing:
                - `members.ssn`
                - `members.first_name`
                - `members.dob`
                #### Hard Policy: `unapproved_join_violation`
                Blocks direct relationship paths that bypass clinical claim lineage.
                Example blocked path: `members → providers`
                """)
            with gr.Column(scale=2):
                gr.Plot(value=policy_graph_fig, label="Governance Policy Graph")
        gr.Dataframe(value=df_policies, label="Policy Registry", interactive=False, wrap=True)

    with gr.Tab("Create New Metric Registry"):
        gr.Markdown("""
        ## Create New Metric Registry
        Click a metric to walk through how it becomes governed.
        """)
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("""
                ### Available Metric Templates
                Start with a high-value healthcare revenue cycle metric.
                """)
                clean_claim_btn = gr.Button("Configure Clean Claim Rate", variant="primary")
                gr.Markdown("""
                Other future templates:
                - Denial volume by code
                - Total member responsibility
                - Authorization mismatch rate
                """)
            with gr.Column(scale=2):
                walkthrough_output = gr.Markdown(value="""
                ## Metric Walkthrough
                Click **Configure Clean Claim Rate** to begin.
                """)
        clean_claim_btn.click(fn=show_clean_claim_rate_walkthrough, outputs=walkthrough_output)
        gr.Markdown("## Registry Tables")
        gr.Dataframe(value=df_metric_registry, label="Existing Metric Registry", interactive=False, wrap=True)
        gr.Dataframe(value=df_transactional_paths, label="Approved Transactional Paths", interactive=False, wrap=True)
        gr.Dataframe(value=df_blocked_tests, label="Blocked Transaction Test Cases", interactive=False, wrap=True)

    with gr.Tab("Virgo Playground"):
        gr.Markdown("""
        ## Virgo Playground
        Click benchmark queries to see whether Virgo approves or blocks them.
        **Benchmark Query → Mechanistic Triage → SQL Generation → Layer 4 Gatekeeper → Audit Result**
        """)
        with gr.Row():
            with gr.Column(scale=1):
                selected_query = gr.Radio(choices=get_query_choices(), value=get_query_choices()[0], label="Benchmark Query")
                wal_volume_limit = gr.Slider(minimum=64, maximum=2048, value=1024, step=64, label="WAL Volume Limit MB")
                recovery_time_seconds = gr.Slider(minimum=30, maximum=900, value=300, step=30, label="Recovery Time Seconds")
                run_button = gr.Button("Run Governed Query", variant="primary")
            with gr.Column(scale=2):
                generated_sql_output = gr.Code(label="Generated SQL", language="sql", value="Click Run Governed Query to generate SQL.")
                decision_output = gr.Markdown(value="Run a benchmark query to see the gatekeeper decision.")
        playground_audit_output = gr.Dataframe(label="Single-Query Audit Record", interactive=False, wrap=True)
        run_button.click(
            fn=run_playground_query,
            inputs=[selected_query, wal_volume_limit, recovery_time_seconds],
            outputs=[generated_sql_output, decision_output, playground_audit_output]
        )

    with gr.Tab("Aetna Audit Dashboard"):
        gr.Markdown("""
        ## Aetna Audit Dashboard
        Manager view for traceability, compliance review, and governance performance.
        This dashboard answers:
        - Who asked what question?
        - What policy was triggered?
        - Was the query approved or blocked?
        - Did the system catch every governance violation?
        """)
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("""
                ### Governance Summary
                The system successfully caught:
                # 100%
                of governance violations in the benchmark run.
                This includes PHI/PII exposure attempts, SSN retrieval, email export attempts, direct member-provider joins, and clinical lineage bypass attempts.
                """)
            with gr.Column(scale=2):
                gr.Dataframe(value=df_audit_summary, label="Governance Performance Summary", interactive=False, wrap=True)
        gr.Dataframe(value=df_audit_dashboard, label="Full Audit Log", interactive=False, wrap=True)

if __name__ == "__main__":
    demo.launch()