"""
APEX-VW Propagation Demo -- Gradio app
======================================

Interactive demonstration of *correction propagation* for English->Spanish
post-editing, built around the APEX-VW corpus workflow.

Flow
----
1. Load the built-in NHS virtual-ward sample, or paste your own English text.
2. The app segments it and machine-translates one paragraph at a time
   (source / MT / post-edit, mirroring apex_vw_src_mt_pe.csv).
3. You post-edit the current paragraph. The app diffs raw MT against your edit
   and mines phrase-substitution rules (the edit-memory engine).
4. On every later paragraph, rules that match its MT appear as checkboxes.
   Tick = accept, untick = reject. Accepted rules are applied on top of the raw
   MT before you finalise that paragraph.

MT backends
-----------
- "Demo (offline)"                   -> no download; canned / rule-based pseudo-MT.
- "Helsinki-NLP/opus-mt-en-es"       -> open-source MarianMT (recommended).
- "facebook/nllb-200-distilled-600M" -> open-source multilingual NLLB.

Run:  python app.py       (local)
Deploy: Hugging Face Spaces, SDK = Gradio, app_file = app.py
"""

from __future__ import annotations

import difflib
import html
import re
from dataclasses import dataclass

import gradio as gr
import spaces  # ZeroGPU: enables @spaces.GPU; a no-op off ZeroGPU

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MODELS = {
    "Demo (offline, no download)": {"kind": "mock"},
    "Helsinki-NLP/opus-mt-en-es": {"kind": "marian", "name": "Helsinki-NLP/opus-mt-en-es"},
    "facebook/nllb-200-distilled-600M": {
        "kind": "nllb", "name": "facebook/nllb-200-distilled-600M",
        "src": "eng_Latn", "tgt": "spa_Latn",
    },
}

_NO_SPACE_BEFORE = set(".,;:!?%)]}\u00bb")
_NO_SPACE_AFTER = set("([{\u00bf\u00a1\u00ab")
MAX_RULE_SPAN = 6

# ---------------------------------------------------------------------------
# Built-in sample (source + canned "MT" so the offline demo is realistic)
# ---------------------------------------------------------------------------

SAMPLE_SOURCE = """A virtual ward allows patients to receive hospital-level care in their own home. Each virtual ward is led by a consultant and supported by a multidisciplinary team.

Referral to a virtual ward should follow the agreed clinical criteria. Before admission to the virtual ward, the team confirms that remote monitoring is appropriate.

Staff on the virtual ward review patient observations at least once a day. If a patient deteriorates, the virtual ward escalates care according to the local pathway.

Technology-enabled virtual wards use apps and wearable devices to monitor patients. A technology-enabled virtual ward can safely support a larger number of patients.

Discharge from the virtual ward is agreed jointly by the clinical team and the patient. After discharge, the virtual ward shares a summary with the patient's GP."""

SAMPLE_MT = {
    "A virtual ward allows patients to receive hospital-level care in their own home. Each virtual ward is led by a consultant and supported by a multidisciplinary team.":
        "Una sala virtual permite a los pacientes recibir atenci\u00f3n de nivel hospitalario en su propio domicilio. Cada sala virtual est\u00e1 dirigida por un m\u00e9dico consultor y cuenta con el apoyo de un equipo multidisciplinar.",
    "Referral to a virtual ward should follow the agreed clinical criteria. Before admission to the virtual ward, the team confirms that remote monitoring is appropriate.":
        "La derivaci\u00f3n a una sala virtual debe seguir los criterios cl\u00ednicos acordados. Antes del ingreso en la sala virtual, el equipo confirma que la monitorizaci\u00f3n remota es apropiada.",
    "Staff on the virtual ward review patient observations at least once a day. If a patient deteriorates, the virtual ward escalates care according to the local pathway.":
        "El personal de la sala virtual revisa las observaciones del paciente al menos una vez al d\u00eda. Si un paciente se deteriora, la sala virtual escala la atenci\u00f3n seg\u00fan la v\u00eda local.",
    "Technology-enabled virtual wards use apps and wearable devices to monitor patients. A technology-enabled virtual ward can safely support a larger number of patients.":
        "Las salas virtuales habilitadas por tecnolog\u00eda utilizan aplicaciones y dispositivos ponibles para monitorizar a los pacientes. Una sala virtual habilitada por tecnolog\u00eda puede atender de forma segura a un mayor n\u00famero de pacientes.",
    "Discharge from the virtual ward is agreed jointly by the clinical team and the patient. After discharge, the virtual ward shares a summary with the patient's GP.":
        "El alta de la sala virtual se acuerda de forma conjunta entre el equipo cl\u00ednico y el paciente. Tras el alta, la sala virtual comparte un resumen con el m\u00e9dico de cabecera del paciente.",
}

MOCK_PHRASES = [
    ("virtual wards", "salas virtuales"), ("virtual ward", "sala virtual"),
    ("hospital-level care", "atenci\u00f3n de nivel hospitalario"),
    ("multidisciplinary team", "equipo multidisciplinar"),
    ("remote monitoring", "monitorizaci\u00f3n remota"),
    ("technology-enabled", "habilitada por tecnolog\u00eda"),
    ("clinical team", "equipo cl\u00ednico"),
]
MOCK_WORDS = {
    "a": "una", "an": "una", "the": "el", "patients": "pacientes", "patient": "paciente",
    "care": "atenci\u00f3n", "home": "domicilio", "team": "equipo", "staff": "personal",
    "referral": "derivaci\u00f3n", "admission": "ingreso", "discharge": "alta",
    "consultant": "m\u00e9dico consultor", "apps": "aplicaciones", "devices": "dispositivos",
    "monitor": "monitorizar", "summary": "resumen", "day": "d\u00eda", "pathway": "v\u00eda",
    "and": "y", "in": "en", "to": "a", "of": "de", "with": "con", "is": "es",
}

# ---------------------------------------------------------------------------
# Pure helpers (tokenisation, rule mining, rule application)
# ---------------------------------------------------------------------------

def tokenize(s: str) -> list[str]:
    return re.findall(r"\w+|[^\w\s]", s, re.UNICODE)


def detok(tokens: list[str]) -> str:
    out = ""
    for i, tok in enumerate(tokens):
        if i == 0:
            out = tok
            continue
        prev = tokens[i - 1]
        if tok and tok[0] in _NO_SPACE_BEFORE:
            out += tok
        elif prev and prev[-1] in _NO_SPACE_AFTER:
            out += tok
        else:
            out += " " + tok
    return out


def _join_pairs_html(pairs: list[tuple[str, bool]]) -> str:
    out = ""
    for i, (tok, hl) in enumerate(pairs):
        piece = html.escape(tok)
        if hl:
            piece = f'<span class="prop-hit">{piece}</span>'
        if i == 0:
            out = piece
            continue
        prev = pairs[i - 1][0]
        if tok and tok[0] in _NO_SPACE_BEFORE:
            out += piece
        elif prev and prev[-1] in _NO_SPACE_AFTER:
            out += piece
        else:
            out += " " + piece
    return out


def mine_rules(mt: str, pe: str, max_span: int = MAX_RULE_SPAN) -> list[tuple[str, str]]:
    a, b = tokenize(mt), tokenize(pe)
    sm = difflib.SequenceMatcher(a=a, b=b, autojunk=False)
    rules: list[tuple[str, str]] = []
    for tag, i1, i2, j1, j2 in sm.get_opcodes():
        if tag != "replace":
            continue
        if (i2 - i1) > max_span or (j2 - j1) > max_span:
            continue
        src, tgt = detok(a[i1:i2]), detok(b[j1:j2])
        if not src.strip() or not tgt.strip() or src == tgt:
            continue
        if re.fullmatch(r"[\W_]+", src):
            continue
        rules.append((src, tgt))
    return rules


def rule_applies(src: str, mt_tokens: list[str]) -> bool:
    st_toks = tokenize(src)
    L = len(st_toks)
    if L == 0:
        return False
    return any(mt_tokens[i:i + L] == st_toks for i in range(len(mt_tokens) - L + 1))


def apply_rules(mt: str, accepted: list[tuple[str, str]]):
    tokens = tokenize(mt)
    compiled = sorted(
        ((tokenize(s), tokenize(t), s, t) for s, t in accepted),
        key=lambda r: len(r[0]), reverse=True,
    )
    out_pairs: list[tuple[str, bool]] = []
    applied: dict[tuple[str, str], int] = {}
    i, n = 0, len(tokens)
    while i < n:
        matched = False
        for src_toks, tgt_toks, s, t in compiled:
            L = len(src_toks)
            if L and tokens[i:i + L] == src_toks:
                out_pairs.extend((tok, True) for tok in tgt_toks)
                applied[(s, t)] = applied.get((s, t), 0) + 1
                i += L
                matched = True
                break
        if not matched:
            out_pairs.append((tokens[i], False))
            i += 1
    plain = detok([tok for tok, _ in out_pairs])
    html_text = _join_pairs_html(out_pairs)
    return plain, html_text, [(s, t, c) for (s, t), c in applied.items()]


def segment_document(text: str, mode: str = "paragraph") -> list[str]:
    text = text.strip()
    if not text:
        return []
    if mode == "paragraph":
        segs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
        if len(segs) > 1:
            return segs
    sents = re.split(r"(?<=[.!?])\s+", text.replace("\n", " "))
    return [s.strip() for s in sents if s.strip()]


def mock_translate(text: str) -> str:
    if text in SAMPLE_MT:
        return SAMPLE_MT[text]
    out = text
    for en, es in MOCK_PHRASES:
        out = re.sub(re.escape(en), es, out, flags=re.IGNORECASE)
    return detok([MOCK_WORDS.get(t.lower(), t) for t in tokenize(out)])


@dataclass(frozen=True)
class Rule:
    src: str
    tgt: str
    learned_from: int


# ---------------------------------------------------------------------------
# MT backends
# ---------------------------------------------------------------------------

_PIPELINES: dict[str, object] = {}


def _get_pipeline(cfg):
    key = cfg["name"]
    if key not in _PIPELINES:
        from transformers import pipeline
        try:
            import torch
            device = 0 if torch.cuda.is_available() else -1
        except Exception:
            device = -1
        if cfg["kind"] == "nllb":
            _PIPELINES[key] = pipeline("translation", model=key,
                                       src_lang=cfg["src"], tgt_lang=cfg["tgt"], device=device)
        else:
            _PIPELINES[key] = pipeline("translation", model=key, device=device)
    return _PIPELINES[key]


@spaces.GPU(duration=30)
def _gpu_translate(model_key: str, sentences: list[str]) -> list[str]:
    """Real-model inference. On ZeroGPU this runs inside an allocated GPU
    context (where torch.cuda becomes available); off ZeroGPU it runs on CPU.
    Its presence is also what satisfies the ZeroGPU startup check."""
    pipe = _get_pipeline(MODELS[model_key])
    return [pipe(s, max_length=512)[0]["translation_text"].strip() for s in sentences]


def translate_segment(text: str, model_key: str) -> str:
    cfg = MODELS[model_key]
    if cfg["kind"] == "mock":
        return mock_translate(text)
    sentences = [s for s in re.split(r"(?<=[.!?])\s+", text) if s.strip()]
    try:
        pieces = _gpu_translate(model_key, sentences)
    except Exception as exc:
        gr.Warning(f"Could not load {model_key} ({exc}). Using offline demo MT.")
        return mock_translate(text)
    return " ".join(pieces)


# ---------------------------------------------------------------------------
# State + derived data
# ---------------------------------------------------------------------------

def default_state() -> dict:
    return {"segments": [], "idx": 0, "committed": {}, "mt_cache": {},
            "model_key": list(MODELS)[0], "applic": []}


def current_mt(state: dict, idx: int) -> str:
    key = (state["model_key"], idx)
    if key not in state["mt_cache"]:
        state["mt_cache"][key] = translate_segment(state["segments"][idx], state["model_key"])
    return state["mt_cache"][key]


def learned_rules(state: dict, exclude_idx: int | None = None) -> list[Rule]:
    seen: set[tuple[str, str]] = set()
    rules: list[Rule] = []
    for i in sorted(state["committed"]):
        if i == exclude_idx:
            continue
        for s, t in mine_rules(current_mt(state, i), state["committed"][i]):
            if (s, t) not in seen:
                seen.add((s, t))
                rules.append(Rule(s, t, i))
    return rules


def compute_applicable(state: dict, idx: int) -> list[tuple[str, str, str]]:
    """Return (label, src, tgt) for rules from other paragraphs that match this MT."""
    mt_toks = tokenize(current_mt(state, idx))
    out = []
    for r in learned_rules(state, exclude_idx=idx):
        if rule_applies(r.src, mt_toks):
            label = f"{r.src}  \u2192  {r.tgt}"
            out.append((label, r.src, r.tgt))
    return out


def wrap_html(inner: str, applied: list[tuple[str, str, int]]) -> str:
    note = ""
    if applied:
        note = ('<div class="applied-note">Applied: '
                + "; ".join(f"{s} \u2192 {t} (\u00d7{c})" for s, t, c in applied) + "</div>")
    return f'<div class="mtbox">{inner or "&nbsp;"}</div>{note}'


# ---------------------------------------------------------------------------
# Render + event handlers
# ---------------------------------------------------------------------------

def render(state: dict):
    if not state["segments"]:
        return (
            "### Load a document to begin \u2192",
            "", "",
            gr.update(choices=[], value=[]),
            wrap_html("", []),
            gr.update(value=""),
            [], [],
        )
    idx, segs = state["idx"], state["segments"]
    n = len(segs)
    mt = current_mt(state, idx)
    applic = compute_applicable(state, idx)
    state["applic"] = applic
    labels = [a[0] for a in applic]
    accepted_pairs = [(s, t) for (_, s, t) in applic]
    _, html_s, applied = apply_rules(mt, accepted_pairs)

    flag = " \u00b7 committed \u2713" if idx in state["committed"] else ""
    progress = (f"### Paragraph {idx + 1} of {n}{flag}\n"
                f"{len(state['committed'])}/{n} paragraphs post-edited")
    editor_val = state["committed"].get(idx, mt)
    learned = [[r.src, r.tgt, r.learned_from + 1] for r in learned_rules(state)]
    committed = [[i + 1, segs[i], current_mt(state, i), state["committed"][i]]
                 for i in sorted(state["committed"])]

    return (progress, segs[idx], mt,
            gr.update(choices=labels, value=labels),
            wrap_html(html_s, applied),
            gr.update(value=editor_val),
            learned, committed)


def on_load_sample(state):
    new = default_state()
    new["model_key"] = state["model_key"]
    new["segments"] = segment_document(SAMPLE_SOURCE, "paragraph")
    return (new, *render(new))


def on_load_doc(state, text, mode):
    if not (text or "").strip():
        gr.Warning("Paste some English text first, or load the sample.")
        return (state, *render(state))
    new = default_state()
    new["model_key"] = state["model_key"]
    new["segments"] = segment_document(text, mode)
    return (new, *render(new))


def on_model_change(state, model_key):
    state["model_key"] = model_key
    return (state, *render(state))


def on_rules_change(state, accepted_labels):
    if not state["segments"]:
        return wrap_html("", [])
    mt = current_mt(state, state["idx"])
    pairs = [(s, t) for (lab, s, t) in state.get("applic", []) if lab in accepted_labels]
    _, html_s, applied = apply_rules(mt, pairs)
    return wrap_html(html_s, applied)


def on_use_suggestion(state, accepted_labels):
    if not state["segments"]:
        return gr.update()
    mt = current_mt(state, state["idx"])
    pairs = [(s, t) for (lab, s, t) in state.get("applic", []) if lab in accepted_labels]
    plain, _, _ = apply_rules(mt, pairs)
    return gr.update(value=plain)


def on_reset_editor(state):
    if not state["segments"]:
        return gr.update()
    return gr.update(value=current_mt(state, state["idx"]))


def on_commit(state, editor_text):
    if state["segments"]:
        idx = state["idx"]
        state["committed"][idx] = editor_text
        if idx + 1 < len(state["segments"]):
            state["idx"] = idx + 1
    return (state, *render(state))


def on_nav(state, delta):
    if state["segments"]:
        state["idx"] = max(0, min(len(state["segments"]) - 1, state["idx"] + delta))
    return (state, *render(state))


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

CSS = """
.mtbox { background:#f6f8fa; border:1px solid #e1e4e8; border-radius:6px;
         padding:10px 12px; line-height:1.6; }
.prop-hit { background:#c8f7c5; border-radius:3px; padding:0 2px; }
.applied-note { color:#6a737d; font-size:0.85rem; margin-top:6px; }
.srcbox textarea { background:#fffdf5 !important; }
"""


def build_demo():
    with gr.Blocks() as demo:
        gr.HTML(f"<style>{CSS}</style>")
        state = gr.State(default_state())

        gr.Markdown(
            "# APEX-VW \u00b7 Correction Propagation Demo\n"
            "Post-edit EN\u2192ES MT one paragraph at a time. Your edits become rules "
            "(edit-memory) that are proposed on later paragraphs for you to **accept or reject**."
        )

        with gr.Row():
            # ---- Left: controls + source + MT --------------------------------
            with gr.Column(scale=3):
                with gr.Row():
                    model_dd = gr.Dropdown(list(MODELS), value=list(MODELS)[0],
                                           label="MT model", scale=3)
                    seg_mode = gr.Radio(["paragraph", "sentence"], value="paragraph",
                                        label="Segment by", scale=2)
                with gr.Accordion("Input a document", open=True):
                    doc_in = gr.Textbox(lines=4, label="Paste English text",
                                        placeholder="Paste a document, or use the sample \u2192")
                    with gr.Row():
                        load_doc_btn = gr.Button("Load document")
                        sample_btn = gr.Button("Load NHS sample", variant="secondary")

                progress_md = gr.Markdown("### Load a document to begin \u2192")
                with gr.Row():
                    prev_btn = gr.Button("\u25c0 Prev")
                    next_btn = gr.Button("Next \u25b6")

                source_box = gr.Textbox(label="Source (English)", interactive=False,
                                        lines=3, elem_classes="srcbox")
                mt_box = gr.Textbox(label="Raw MT (Spanish)", interactive=False, lines=3)

            # ---- Right: rules panel + suggestion + editor --------------------
            with gr.Column(scale=2):
                gr.Markdown("### Propagation rules")
                gr.Markdown("<span style='color:#6a737d;font-size:0.85rem'>Corrections mined "
                            "from your earlier edits. Tick the ones to propagate onto this "
                            "paragraph.</span>")
                rules_group = gr.CheckboxGroup(choices=[], value=[],
                                               label="Rules (tick to propagate)")
                gr.Markdown("**Suggested translation** (accepted rules applied)")
                suggested_html = gr.HTML(wrap_html("", []))
                with gr.Row():
                    use_btn = gr.Button("\u2b07 Use suggestion", variant="primary")
                    reset_btn = gr.Button("\u21ba Reset to raw MT")
                editor = gr.Textbox(label="Your post-edit (final Spanish)", lines=4)
                commit_btn = gr.Button("Commit & continue \u25b6", variant="primary")

        with gr.Accordion("Learned rules (edit-memory)", open=False):
            learned_df = gr.Dataframe(headers=["MT phrase", "Correction", "From \u00b6"],
                                      interactive=False, wrap=True)
        with gr.Accordion("Post-edited segments (source / targetmt / targetpe)", open=False):
            committed_df = gr.Dataframe(headers=["#", "source", "targetmt", "targetpe"],
                                        interactive=False, wrap=True)

        # ---- Wiring ------------------------------------------------------------
        seg_out = [progress_md, source_box, mt_box, rules_group, suggested_html,
                   editor, learned_df, committed_df]

        sample_btn.click(on_load_sample, [state], [state, *seg_out])
        load_doc_btn.click(on_load_doc, [state, doc_in, seg_mode], [state, *seg_out])
        model_dd.change(on_model_change, [state, model_dd], [state, *seg_out])
        rules_group.change(on_rules_change, [state, rules_group], suggested_html)
        use_btn.click(on_use_suggestion, [state, rules_group], editor)
        reset_btn.click(on_reset_editor, [state], editor)
        commit_btn.click(on_commit, [state, editor], [state, *seg_out])
        prev_btn.click(lambda s: on_nav(s, -1), [state], [state, *seg_out])
        next_btn.click(lambda s: on_nav(s, +1), [state], [state, *seg_out])

    return demo


demo = build_demo()

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