# SemanticCitation — Databricks Free Personal Account POC
# Maps to the 5-stage architecture. Zero/minimal pip beyond stdlib + display.
# Production path: replace in-memory docs with Delta tables + Delta Sharing → edge worker.

# ── Cell 0 · Setup ──────────────────────────────────────────────
from __future__ import annotations
import json, hashlib, re, time
from dataclasses import dataclass, asdict
from typing import List, Optional, Tuple
from datetime import datetime, timezone

# ── Cell 1 · Massive Ingestion (mock Unity Catalog rows) ────────
# Stage 01: in production these are Delta table rows shared via R2.

CORPUS = [
  {
    "doc_id": "NCT04829123",
    "catalog_row": "uc.pharma.trials.results.row_1842",
    "title": "Phase III mAb-X NSCLC",
    "text": (
      "In the intention-to-treat population (n=842), mAb-X plus platinum-doublet "
      "chemotherapy achieved a median progression-free survival of 11.4 months "
      "(95% CI 9.8–13.1) versus 6.9 months (95% CI 5.7–8.2) for chemotherapy alone "
      "(HR 0.62, p<0.001). The objective response rate was 48.2% in the experimental arm. "
      "Grade ≥3 treatment-related adverse events occurred in 34% of patients receiving mAb-X."
    ),
  },
  {
    "doc_id": "FDA-label-mAbX-2025",
    "catalog_row": "uc.pharma.labels.fda.row_2201",
    "title": "FDA label mAb-X",
    "text": (
      "mAb-X is a humanized IgG1 monoclonal antibody directed against PD-L1. "
      "Recommended dosage is 1200 mg administered as an intravenous infusion over 60 minutes "
      "every 3 weeks. Immune-mediated pneumonitis occurred in 3.1% of patients."
    ),
  },
  {
    "doc_id": "US11234567B2",
    "catalog_row": "uc.ip.patents.uspto.row_441",
    "title": "Patent dual-payload ADC",
    "text": (
      "An antibody-drug conjugate comprising a monoclonal antibody that binds HER2, "
      "conjugated via a cleavable valine-citrulline linker to a maytansinoid payload at a "
      "drug-to-antibody ratio of 3.5 to 4.0. ADC-14 exhibited IC50 of 0.8 nM against SK-BR-3 "
      "cells and reduced mean tumor volume by 78% at day 28."
    ),
  },
  {
    "doc_id": "EMA-CHMP-2025-44",
    "catalog_row": "uc.reg.ema.chmp.row_44",
    "title": "EMA CK-771",
    "text": (
      "CK-771 is a selective oral KRAS G12C inhibitor. Confirmed overall response rate was "
      "37.5% (48/128) with median duration of response 8.3 months. Plasma half-life averaged "
      "18.2 hours under fed conditions."
    ),
  },
  {
    "doc_id": "WIPO-WO2024-991122",
    "catalog_row": "uc.ip.patents.wipo.row_991",
    "title": "Continuous-flow pyrrolotriazine",
    "text": (
      "The continuous-flow process produces the pyrrolotriazine intermediate at 94.6% purity "
      "with a space-time yield of 2.1 kg L⁻¹ h⁻¹ at 80 °C. Residual palladium content was "
      "below 2 ppm in all pilot batches (n=14)."
    ),
  },
  {
    "doc_id": "NEJM-2024-881",
    "catalog_row": "uc.lit.nejm.row_881",
    "title": "HER2 ADC cardiotoxicity meta-analysis",
    "text": (
      "Across twelve randomized trials (n=6,412), symptomatic left-ventricular dysfunction "
      "was reported in 2.4% of patients receiving HER2-directed ADCs versus 0.9% with control "
      "(RR 2.67, 95% CI 1.41–5.04)."
    ),
  },
]

def chunk_docs(docs, max_chars=400):
    """Stage 02 helper: naive chunker before embedding / hybrid retrieval."""
    chunks = []
    for d in docs:
        text = d["text"]
        for i in range(0, len(text), max_chars):
            chunks.append({
                "doc_id": d["doc_id"],
                "catalog_row": d["catalog_row"],
                "chunk_id": f"{d['doc_id']}::c{i // max_chars}",
                "text": text[i:i + max_chars],
            })
    return chunks

CHUNKS = chunk_docs(CORPUS)
print(f"Ingested {len(CORPUS)} docs → {len(CHUNKS)} chunks")

# ── Cell 2 · Semantic Retrieval (mock RRF) ──────────────────────
# Stage 02: production uses voyage-3-large + sqlite-vec + BM25 fusion.

def retrieve(query: str, k: int = 4):
    q_tokens = set(re.findall(r"[a-z0-9]{4,}", query.lower()))
    scored = []
    for c in CHUNKS:
        t = set(re.findall(r"[a-z0-9]{4,}", c["text"].lower()))
        score = len(q_tokens & t)
        scored.append((score, c))
    scored.sort(key=lambda x: -x[0])
    return [c for s, c in scored[:k] if s > 0] or [c for _, c in scored[:k]]

# ── Cell 3 · Claim extraction + Verbatim Enforcer ───────────────
# Stage 04: deterministic exact substring. No model in the gate.

def normalize(s: str) -> str:
    s = s.lower().replace("–", "-").replace("—", "-")
    return re.sub(r"\s+", " ", s).strip()

def extract_sentences(draft: str) -> List[str]:
    parts = re.split(r"(?<=[.!?])\s+|\n+", draft.strip())
    return [p.strip().rstrip(".!?") for p in parts if len(p.strip()) > 20]

@dataclass
class ClaimResult:
    claim: str
    status: str  # verified | flagged
    doc_id: Optional[str]
    catalog_row: Optional[str]
    snippet: Optional[str]

def enforce_verbatim(draft: str, corpus=CORPUS) -> List[ClaimResult]:
    results = []
    for sent in extract_sentences(draft):
        nclaim = normalize(sent)
        hit = None
        for d in corpus:
            if nclaim in normalize(d["text"]):
                hit = d
                break
        if hit:
            results.append(ClaimResult(
                claim=sent, status="verified",
                doc_id=hit["doc_id"], catalog_row=hit["catalog_row"],
                snippet=sent[:160],
            ))
        else:
            results.append(ClaimResult(
                claim=sent, status="flagged",
                doc_id=None, catalog_row=None, snippet=None,
            ))
    return results

# ── Cell 4 · Proof certificate ──────────────────────────────────
# Stage 05: XFPC-style JSON artifact.

def proof_certificate(draft: str, results: List[ClaimResult]) -> dict:
    verified = [r for r in results if r.status == "verified"]
    flagged = [r for r in results if r.status == "flagged"]
    payload = {
        "certificate_id": "XFPC-SC-" + datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S"),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "input_hash": hashlib.sha256(draft.encode()).hexdigest(),
        "verified_count": len(verified),
        "flagged_count": len(flagged),
        "match_rate": len(verified) / max(1, len(results)),
        "determinism_score": 1.0,
        "sentence_map": [asdict(r) for r in results],
        "signature": "ed25519-mock:" + hashlib.sha256(
            json.dumps([asdict(r) for r in results], sort_keys=True).encode()
        ).hexdigest()[:24],
    }
    return payload

# ── Cell 5 · End-to-end demo ────────────────────────────────────
DRAFT = (
    "In NCT04829123, mAb-X plus platinum-doublet chemotherapy achieved a median "
    "progression-free survival of 11.4 months versus 6.9 months for chemotherapy alone "
    "(HR 0.62, p<0.001). The objective response rate was 48.2% in the experimental arm. "
    "Overall survival improved to 24.8 months, establishing a new standard of care."
)

t0 = time.time()
hits = retrieve(DRAFT)
print("Top retrieval:", [h["doc_id"] for h in hits])
# Stage 03 would call Gemma 4 here with <|think|> — skipped in free notebook.
results = enforce_verbatim(DRAFT)
cert = proof_certificate(DRAFT, results)
cert["processing_ms"] = int((time.time() - t0) * 1000)

# Dashboard-style summary
print("\n=== VERBATIM ENFORCER DASHBOARD ===")
print(f"Match rate:   {cert['match_rate']:.0%}")
print(f"Verified:     {cert['verified_count']}")
print(f"Flagged:      {cert['flagged_count']}")
print(f"Certificate:  {cert['certificate_id']}")
for r in results:
    mark = "✓" if r.status == "verified" else "✗"
    print(f"  {mark} [{r.status}] {r.claim[:80]}...")
    if r.catalog_row:
        print(f"      → {r.catalog_row}")

# Optional: display as table in Databricks
try:
    import pandas as pd
    display(pd.DataFrame([asdict(r) for r in results]))  # noqa: F821
    display(cert)  # noqa: F821
except Exception:
    print(json.dumps(cert, indent=2)[:1200])