# /// script
# requires-python = ">=3.9,<3.14"
# dependencies = [
#     "anonymacy",
#     "datasets",
#     "gliner==0.2.24",
#     "gliner[tokenizers]",
#     "hf_xet",
#     "ipython==9.8.0",
# ]
#
# [tool.uv.sources]
# anonymacy = { path = "../", editable = true }
# ///

import marimo

__generated_with = "0.18.4"
app = marimo.App(width="medium")


@app.cell
def _():
    from datasets import load_dataset
    from gliner import GLiNER
    import pandas as pd
    import json
    from tqdm import tqdm
    from sklearn.metrics import precision_recall_fscore_support

    dataset = load_dataset("E3-JSI/synthetic-multi-pii-ner-v1")
    model = GLiNER.from_pretrained("knowledgator/gliner-pii-large-v1.0")

    SUPPORTED_LABELS = {
        "name", "first name", "last name", "name medical professional",
        "dob", "age", "gender", "marital status", "email address",
        "phone number", "ip address", "url", "location address",
        "location street", "location city", "location state",
        "location country", "location zip", "account number",
        "bank account", "routing number", "credit card",
        "credit card expiration", "cvv", "ssn", "money", "condition",
        "medical process", "drug", "dose", "blood type", "injury",
        "organization medical facility", "healthcare number",
        "medical code", "passport number", "driver license",
        "username", "password", "vehicle id"
    }

    DUTCH_MAPPING = {
        "naam": "name", "voornaam": "first name", "achternaam": "last name",
        "volledige naam": "name", "persoon": "name",
        "patiënt": "name medical professional", "arts": "name medical professional",
        "dokter": "name medical professional", "huisarts": "name medical professional",
        "geboortedatum": "dob", "geboortejaar": "dob", "leeftijd": "age",
        "geboorteplaats": "location city", "nationaliteit": "location country",
        "geslacht": "gender", "burgerlijke staat": "marital status",
        "e-mail": "email address", "emailadres": "email address",
        "telefoonnummer": "phone number", "faxnummer": "phone number",
        "ip adres": "ip address", "url": "url",
        "adres": "location address", "straat": "location street",
        "plaats": "location city", "stad": "location city",
        "postcode": "location zip", "land": "location country",
        "account": "account number", "bankrekening": "bank account",
        "creditcardnummer": "credit card", "bsn": "ssn",
        "bedrag": "money", "aandoening": "condition",
        "behandeling": "medical process", "medicatie": "drug",
        "dosis": "dose", "bloedgroep": "blood type",
        "ziekenhuis": "organization medical facility",
    }

    def create_token_labels(text, tokens, entities):
        """Create binary labels for tokens using actual text positions."""
        labels = [0] * len(tokens)
        char_pos = 0

        for i, token in enumerate(tokens):
            # Find token's actual position in text
            char_pos = text.find(token, char_pos)
            if char_pos == -1:
                continue

            token_start = char_pos
            token_end = char_pos + len(token)

            # Check overlap with entities
            for ent_start, ent_end, _ in entities:
                if not (token_end <= ent_start or token_start >= ent_end):
                    labels[i] = 1
                    break

            char_pos = token_end

        return labels

    def benchmark_language(dataset, language):
        """Benchmark model on a specific language."""
        all_true, all_pred = [], []

        samples = [s for split in dataset.values() for s in split
                   if s["language"] == language]

        for sample in tqdm(samples, desc=f"Benchmarking {language}"):
            text = sample["text"]
            tokens = list(sample["gliner_tokenized_text"])
            true_entities = json.loads(sample["gliner_entities"])

            # Map Dutch entities
            if language == "Dutch":
                true_entities = [(s, e, DUTCH_MAPPING.get(t, t))
                                for s, e, t in true_entities]
                true_entities = [(s, e, t) for s, e, t in true_entities
                                if t in SUPPORTED_LABELS]

            if not true_entities:
                continue

            # Get predictions
            entity_types = list(set(ent[2] for ent in true_entities))
            predictions = model.predict_entities(text, entity_types)
            pred_entities = [(p["start"], p["end"], p["label"])
                            for p in predictions]

            # Create labels
            all_true.extend(create_token_labels(text, tokens, true_entities))
            all_pred.extend(create_token_labels(text, tokens, pred_entities))

        if not all_true:
            return {"Precision": "0.0000", "Recall": "0.0000",
                    "F1": "0.0000", "Samples": 0}

        p, r, f1, _ = precision_recall_fscore_support(
            all_true, all_pred, average="binary", zero_division=0
        )

        return {"Precision": f"{p:.4f}", "Recall": f"{r:.4f}",
                "F1": f"{f1:.4f}", "Samples": len(samples)}

    results = {lang: benchmark_language(dataset, lang)
               for lang in ["Dutch", "English"]}

    print("=" * 60)
    print("GLINER-PII-LARGE BENCHMARK")
    print("=" * 60)
    print(pd.DataFrame(results).T.to_string())
    return


@app.cell
def _():
    return


@app.cell
def _():
    #         Precision  Recall      F1 Samples
    # nvidia/gliner-pii
    # Dutch      0.7873  0.7520  0.7693     446
    # English    0.7343  0.6875  0.7101     452

    # knowledgator/gliner-x-large
    # Dutch      0.7749  0.8829  0.8253     446
    # English    0.7203  0.7859  0.7517     452

    # knowledgator/gliner-pii-large-v1.0
    #Dutch      0.7856  0.8303  0.8074     446
    #English    0.7269  0.7620  0.7440     452

    # E3-JSI/gliner-multi-pii-domains-v1
    #Dutch      0.7810  0.8753  0.8255     446
    #English    0.7151  0.7895  0.7504     452
    return


@app.cell
def _():
    return


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()
