import json

from typing import Optional

from docdeid import Annotation, AnnotationSet

from deduce import Deduce

def regression_test(

    model: Deduce,

    examples_file: str,

    enabled: set[str],

    known_failures: Optional[set[int]] = None,

):

    if known_failures is None:

        known_failures = set()

    with open(examples_file, "rb") as file:

        examples = json.load(file)["examples"]

    failures = set()

    for example in examples:

        trues = AnnotationSet(

            Annotation(**annotation) for annotation in example["annotations"]

        )

        preds = model.deidentify(text=example["text"], enabled=enabled).annotations

        try:

            assert trues == preds

        except AssertionError:

            failures.add(example["id"])

    assert failures == known_failures

def annotators_from_group(model: Deduce, group: str) -> set[str]:

    return {name for name, _ in model.processors[group]}.union({group})

class TestRegression:

    def test_regression_name(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/names.json",

            enabled=annotators_from_group(model, "names"),

        )

    def test_regression_location(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/locations.json",

            enabled=annotators_from_group(model, "locations"),

        )

    def test_regression_institution(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/institutions.json",

            enabled=annotators_from_group(model, "institutions"),

        )

    def test_regression_date(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/dates.json",

            enabled=annotators_from_group(model, "dates"),

        )

    def test_regression_age(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/ages.json",

            enabled=annotators_from_group(model, "ages"),

        )

    def test_regression_identifier(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/identifiers.json",

            enabled=annotators_from_group(model, "identifiers"),

        )

    def test_regression_phone(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/phone_numbers.json",

            enabled=annotators_from_group(model, "phone_numbers"),

        )

    def test_regression_email(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/emails.json",

            enabled=annotators_from_group(model, "email_addresses"),

        )

    def test_regression_url(self, model):

        regression_test(

            model=model,

            examples_file="tests/data/regression_cases/urls.json",

            enabled=annotators_from_group(model, "urls"),

        )