import json


# We have four types of outputs to leverage from radgraph
# 1. Entity token
# 2. Entity label
# 3. Relation tokens
# 4. Relation label

# Question for Jean Benoit: good idea to be based on f1 score ?
# We are using sets: so we do not account for repetitions of a word ! Except through relations ?

# We do three rewards for each type, so a total of six rewards
# +1 reward that does everything at the same time (maybe not that great)
# We combine these rewards with chosen coefficients (hyperparameters)

# To test our reward, we want that no token < token with wrong label < token with correct label
# Compares important clinical tokens of both texts
# Ex: is "pleural effusion" in both texts ?

# Liste de nos modifications
# 1. on met des lower partout, just in case
# 2. on met des return 0 si aucune entité ou relation n'est detectée
# 3. enlever les prints


def exact_entity_token_if_all_match_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    candidates = []
    for annotation_list in [hypothesis_annotation_list, reference_annotation_list]:
        candidate = []
        for entity in annotation_list["entities"].values():
            if not entity["relations"]:
                candidate.append((entity["tokens"], entity["label"]))
            if entity["relations"]:
                candidate.extend([(entity["tokens"].lower(),
                                   entity["label"],
                                   r[0],
                                   annotation_list["entities"][r[1]]["tokens"].lower())
                                  for r in entity["relations"]]
                                 )
               
        candidate = set(candidate)
        candidates.append(candidate)

    hypothesis_relation_token_list, reference_relation_token_list = candidates
    precision = (
        sum(
            [
                1
                for x in hypothesis_relation_token_list
                if (x in reference_relation_token_list)
            ]
        )
        / len(hypothesis_relation_token_list)
        if len(hypothesis_relation_token_list) > 0
        else 0.0
    )
    recall = (
        sum(
            [
                1
                for x in reference_relation_token_list
                if (x in hypothesis_relation_token_list)
            ]
        )
        / len(reference_relation_token_list)
        if len(reference_relation_token_list) > 0
        else 0.0
    )
    f1_score = (
        (2 * precision * recall / (precision + recall))
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


def exact_entity_token_if_rel_match_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    candidates = []
    for annotation_list in [hypothesis_annotation_list, reference_annotation_list]:
        candidate = []
        for entity in annotation_list["entities"].values():
            if not entity["relations"]:
                candidate.append((entity["tokens"], entity["label"]))
            if entity["relations"]:
                candidate.append(tuple((entity["tokens"].lower(),
                                        annotation_list["entities"][r[1]]["tokens"].lower())
                                       for r in entity["relations"])
                                 )
        candidate = set(candidate)
        candidates.append(candidate)

    hypothesis_relation_token_list, reference_relation_token_list = candidates

    precision = (
        sum(
            [
                1
                for x in hypothesis_relation_token_list
                if (x in reference_relation_token_list)
            ]
        )
        / len(hypothesis_relation_token_list)
        if len(hypothesis_relation_token_list) > 0
        else 0.0
    )
    recall = (
        sum(
            [
                1
                for x in reference_relation_token_list
                if (x in hypothesis_relation_token_list)
            ]
        )
        / len(reference_relation_token_list)
        if len(reference_relation_token_list) > 0
        else 0.0
    )
    f1_score = (
        (2 * precision * recall / (precision + recall))
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


def exact_entity_token_if_rel_exists_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    candidates = []
    for annotation_list in [hypothesis_annotation_list, reference_annotation_list]:
        candidate = []
        for entity in annotation_list["entities"].values():
            if not entity["relations"]:
                candidate.append((entity["tokens"], entity["label"]))
            if entity["relations"]:
                candidate.append((entity["tokens"], entity["label"], True))

        candidate = set(candidate)
        candidates.append(candidate)

    hypothesis_relation_token_list, reference_relation_token_list = candidates

    precision = (
        sum(
            [
                1
                for x in hypothesis_relation_token_list
                if (x in reference_relation_token_list)
            ]
        )
        / len(hypothesis_relation_token_list)
        if len(hypothesis_relation_token_list) > 0
        else 0.0
    )
    recall = (
        sum(
            [
                1
                for x in reference_relation_token_list
                if (x in hypothesis_relation_token_list)
            ]
        )
        / len(reference_relation_token_list)
        if len(reference_relation_token_list) > 0
        else 0.0
    )
    f1_score = (
        (2 * precision * recall / (precision + recall))
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


def exact_entity_token_match_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    candidates = []
    for annotation_list in [hypothesis_annotation_list, reference_annotation_list]:
        candidate = []
        for entity in annotation_list["entities"].values():
            candidate.append((entity["tokens"], entity["label"]))

        candidate = set(candidate)
        candidates.append(candidate)

    hypothesis_relation_token_list, reference_relation_token_list = candidates

    precision = (
        sum(
            [
                1
                for x in hypothesis_relation_token_list
                if (x in reference_relation_token_list)
            ]
        )
        / len(hypothesis_relation_token_list)
        if len(hypothesis_relation_token_list) > 0
        else 0.0
    )
    recall = (
        sum(
            [
                1
                for x in reference_relation_token_list
                if (x in hypothesis_relation_token_list)
            ]
        )
        / len(reference_relation_token_list)
        if len(reference_relation_token_list) > 0
        else 0.0
    )

    f1_score = (
        (2 * precision * recall / (precision + recall))
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


# To be tested alone
def exact_entity_token_match_reward_old(
        hypothesis_annotation_list, reference_annotation_list
):
    hypothesis_entity_token_list = set(
        map(
            lambda x: x["tokens"].lower(),
            hypothesis_annotation_list["entities"].values(),
        )
    )
    reference_entity_token_list = set(
        map(
            lambda x: x["tokens"].lower(),
            reference_annotation_list["entities"].values(),
        )
    )

    precision = (
        sum(
            [
                1
                for x in hypothesis_entity_token_list
                if x in reference_entity_token_list
            ]
        )
        / len(hypothesis_entity_token_list)
        if len(hypothesis_entity_token_list) > 0
        else 0.0
    )
    recall = (
        sum(
            [
                1
                for x in reference_entity_token_list
                if x in hypothesis_entity_token_list
            ]
        )
        / len(reference_entity_token_list)
        if len(reference_entity_token_list) > 0
        else 0.0
    )
    f1_score = (
        2 * precision * recall / (precision + recall)
        if (precision + recall) > 0
        else 0.0
    )
    return f1_score


# Compares important clinical entities and their syntaxtic use in both text
# Ex: is "pleural effusion" as a Definitively Absent Observation in both texts ?
# To be tested alone
def exact_entity_token_and_label_match_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    hypothesis_entity_token_and_label_list = set(
        map(
            lambda x: (x["tokens"].lower(), x["label"]),
            hypothesis_annotation_list["entities"].values(),
        )
    )
    reference_entity_token_and_label_list = set(
        map(
            lambda x: (x["tokens"].lower(), x["label"]),
            reference_annotation_list["entities"].values(),
        )
    )

    precision = (
        sum(
            [
                1
                for x in hypothesis_entity_token_and_label_list
                if x in reference_entity_token_and_label_list
            ]
        )
        / len(hypothesis_entity_token_and_label_list)
        if len(hypothesis_entity_token_and_label_list) > 0
        else 0.0
    )

    recall = (
        sum(
            [
                1
                for x in reference_entity_token_and_label_list
                if x in hypothesis_entity_token_and_label_list
            ]
        )
        / len(reference_entity_token_and_label_list)
        if len(reference_entity_token_and_label_list) > 0
        else 0.0
    )

    f1_score = (
        2 * precision * recall / (precision + recall)
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


# Ex: if "pleural effusion" is in both texts, is it as a Definitively Absent Observation in both texts?
# Not to be tested
def exact_entity_label_if_correct_token_match_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    raise Exception("not supposed to be used")
    hypothesis_entity_token_and_label_list = set(
        map(
            lambda x: (x["tokens"], x["label"]),
            hypothesis_annotation_list["entities"].values(),
        )
    )
    reference_entity_token_and_label_list = set(
        map(
            lambda x: (x["tokens"], x["label"]),
            reference_annotation_list["entities"].values(),
        )
    )

    hypothesis_entity_token_list = set(
        map(lambda x: x["tokens"], hypothesis_annotations["entities"].values())
    )
    reference_entity_token_list = set(
        map(lambda x: x["tokens"], reference_annotations["entities"].values())
    )

    hypothesis_entity_token_and_label_list = [
        (tokens, label)
        for (tokens, label) in hypothesis_entity_token_and_label_list
        if tokens in reference_entity_token_list
    ]
    reference_entity_token_and_label_list = [
        (tokens, label)
        for (tokens, label) in reference_entity_token_and_label_list
        if tokens in hypothesis_entity_token_list
    ]

    precision = (
        sum(
            [
                1
                for x in hypothesis_entity_token_and_label_list
                if x in reference_entity_token_and_label_list
            ]
        )
        / len(hypothesis_entity_token_and_label_list)
        if len(hypothesis_entity_token_and_label_list) > 0
        else 0.0
    )

    recall = (
        sum(
            [
                1
                for x in reference_entity_token_and_label_list
                if x in hypothesis_entity_token_and_label_list
            ]
        )
        / len(reference_entity_token_and_label_list)
        if len(reference_entity_token_and_label_list) > 0
        else 0.0
    )

    f1_score = (
        2 * precision * recall / (precision + recall)
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


# Test alone, maybe best one ?
def partially_exact_entity_token_and_label_match_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    hypothesis_entity_token_and_label_list = set(
        map(
            lambda x: (x["tokens"].lower(), x["label"]),
            hypothesis_annotation_list["entities"].values(),
        )
    )
    reference_entity_token_and_label_list = set(
        map(
            lambda x: (x["tokens"].lower(), x["label"]),
            reference_annotation_list["entities"].values(),
        )
    )

    hypothesis_entity_token_list = set(
        map(
            lambda x: x["tokens"].lower(),
            hypothesis_annotation_list["entities"].values(),
        )
    )
    reference_entity_token_list = set(
        map(
            lambda x: x["tokens"].lower(),
            reference_annotation_list["entities"].values(),
        )
    )

    precision = (
        sum(
            [
                1
                for (x, y) in hypothesis_entity_token_and_label_list
                if (x, y) in reference_entity_token_and_label_list
            ]
            + [
                0.5
                for (x, y) in hypothesis_entity_token_and_label_list
                if (
                        ((x, y) not in reference_entity_token_and_label_list)
                        and (x in reference_entity_token_list)
                )
            ]
        )
        / len(hypothesis_entity_token_and_label_list)
        if len(hypothesis_entity_token_and_label_list) > 0
        else 0.0
    )

    recall = (
        sum(
            [
                1
                for (x, y) in reference_entity_token_and_label_list
                if (x, y) in hypothesis_entity_token_and_label_list
            ]
            + [
                0.5
                for (x, y) in reference_entity_token_and_label_list
                if (
                        ((x, y) not in hypothesis_entity_token_and_label_list)
                        and (x in hypothesis_entity_token_list)
                )
            ]
        )
        / len(reference_entity_token_and_label_list)
        if len(reference_entity_token_and_label_list) > 0
        else 0.0
    )

    f1_score = (
        2 * precision * recall / (precision + recall)
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


# We do not do the partially exact loss for the relations, because does not make sense to do not have the label
# Just an exact match one that we combine with the other reward, and we are good

# Compares the relations, directly both the start and end tokens AS WELL AS the label
# Because a relation without the label does not really make sense
# keep a list or a set ?


def exact_relation_token_and_label_match_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    # print(json.dumps(hypothesis_annotation_list, indent=4))
    hypothesis_relation_token_and_label_list = list(
        map(
            lambda x: [
                (x["tokens"].lower(), relation[0], relation[1])
                for relation in x["relations"]
            ],
            hypothesis_annotation_list["entities"].values(),
        )
    )
    # print(hypothesis_relation_token_and_label_list)
    # troll

    hypothesis_relation_token_and_label_list = set(
        [
            (
                tokens_1,
                label,
                hypothesis_annotation_list["entities"][tokens_2]["tokens"].lower(),
            )
            for relations in hypothesis_relation_token_and_label_list
            for (tokens_1, label, tokens_2) in relations
        ]
    )

    reference_relation_token_and_label_list = list(
        map(
            lambda x: [
                (x["tokens"].lower(), relation[0], relation[1])
                for relation in x["relations"]
            ],
            reference_annotation_list["entities"].values(),
        )
    )

    reference_relation_token_and_label_list = set(
        [
            (
                tokens_1,
                label,
                reference_annotation_list["entities"][tokens_2]["tokens"].lower(),
            )
            for relations in reference_relation_token_and_label_list
            for (tokens_1, label, tokens_2) in relations
        ]
    )

    precision = (
        sum(
            [
                1
                for x in hypothesis_relation_token_and_label_list
                if x in reference_relation_token_and_label_list
            ]
        )
        / len(hypothesis_relation_token_and_label_list)
        if len(hypothesis_relation_token_and_label_list) > 0
        else 0.0
    )

    recall = (
        sum(
            [
                1
                for x in reference_relation_token_and_label_list
                if x in hypothesis_relation_token_and_label_list
            ]
        )
        / len(reference_relation_token_and_label_list)
        if len(reference_relation_token_and_label_list) > 0
        else 0.0
    )

    f1_score = (
        2 * precision * recall / (precision + recall)
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


def partially_exact_relation_token_and_label_match_reward(
        hypothesis_annotation_list, reference_annotation_list
):
    hypothesis_relation_token_and_label_list = list(
        map(
            lambda x: [
                (x["tokens"].lower(), relation[0], relation[1])
                for relation in x["relations"]
            ],
            hypothesis_annotation_list["entities"].values(),
        )
    )

    hypothesis_relation_token_and_label_list = set(
        [
            (
                tokens_1,
                label,
                hypothesis_annotation_list["entities"][tokens_2]["tokens"].lower(),
            )
            for relations in hypothesis_relation_token_and_label_list
            for (tokens_1, label, tokens_2) in relations
        ]
    )

    hypothesis_relation_token_list = set(
        [
            (token_1, token_2)
            for (token_1, label, token_2) in hypothesis_relation_token_and_label_list
        ]
    )

    reference_relation_token_and_label_list = list(
        map(
            lambda x: [
                (x["tokens"].lower(), relation[0], relation[1])
                for relation in x["relations"]
            ],
            reference_annotation_list["entities"].values(),
        )
    )

    reference_relation_token_and_label_list = set(
        [
            (
                tokens_1,
                label,
                reference_annotation_list["entities"][tokens_2]["tokens"].lower(),
            )
            for relations in reference_relation_token_and_label_list
            for (tokens_1, label, tokens_2) in relations
        ]
    )

    reference_relation_token_list = set(
        [
            (token_1, token_2)
            for (token_1, label, token_2) in reference_relation_token_and_label_list
        ]
    )

    precision = (
        sum(
            [
                1
                for (x, y, z) in hypothesis_relation_token_and_label_list
                if (x, y, z) in reference_relation_token_and_label_list
            ]
            + [
                0.5
                for (x, y, z) in hypothesis_relation_token_and_label_list
                if (
                        ((x, y, z) not in reference_relation_token_and_label_list)
                        and (
                                ((x, z) in reference_relation_token_list)
                                or ((z, x) in reference_relation_token_list)
                        )
                )
            ]
        )
        / len(hypothesis_relation_token_and_label_list)
        if len(hypothesis_relation_token_and_label_list) > 0
        else 0.0
    )

    recall = (
        sum(
            [
                1
                for (x, y, z) in reference_relation_token_and_label_list
                if (x, y, z) in hypothesis_relation_token_and_label_list
            ]
            + [
                0.5
                for (x, y, z) in reference_relation_token_and_label_list
                if (
                        ((x, y, z) not in hypothesis_relation_token_and_label_list)
                        and (
                                ((x, z) in hypothesis_relation_token_list)
                                or ((z, x) in hypothesis_relation_token_list)
                        )
                )
            ]
        )
        / len(reference_relation_token_and_label_list)
        if len(reference_relation_token_and_label_list) > 0
        else 0.0
    )

    f1_score = (
        2 * precision * recall / (precision + recall)
        if (precision + recall) > 0
        else 0.0
    )

    return f1_score


hypothesis_annotations = {
    "text": "FINAL REPORT INDICATION : ___ F with cough / / Cough TECHNIQUE : PA and lateral views of the chest . COMPARISON : None . FINDINGS : The lungs are clear without focal consolidation , , or edema . The cardiomediastinal silhouette is within normal limits . No acute osseous abnormalities . IMPRESSION : No acute cardiopulmonary process .",
    "entities": {
        "1": {
            "tokens": "lungs",
            "label": "ANAT-DP",
            "start_ix": 28,
            "end_ix": 28,
            "relations": [],
        },
        "2": {
            "tokens": "clear",
            "label": "OBS-DP",
            "start_ix": 30,
            "end_ix": 30,
            "relations": [["located_at", "1"]],
        },
        "3": {
            "tokens": "focal",
            "label": "OBS-DA",
            "start_ix": 32,
            "end_ix": 32,
            "relations": [["modify", "4"]],
        },
        "4": {
            "tokens": "consolidation",
            "label": "OBS-DA",
            "start_ix": 33,
            "end_ix": 33,
            "relations": [["located_at", "1"]],
        },
        "5": {
            "tokens": "edema",
            "label": "OBS-DA",
            "start_ix": 37,
            "end_ix": 37,
            "relations": [],
        },
        "6": {
            "tokens": "cardiomediastinal",
            "label": "ANAT-DP",
            "start_ix": 40,
            "end_ix": 40,
            "relations": [],
        },
        "7": {
            "tokens": "silhouette",
            "label": "ANAT-DP",
            "start_ix": 41,
            "end_ix": 41,
            "relations": [["modify", "6"]],
        },
        "8": {
            "tokens": "within",
            "label": "OBS-DP",
            "start_ix": 43,
            "end_ix": 43,
            "relations": [],
        },
        "9": {
            "tokens": "normal",
            "label": "OBS-DP",
            "start_ix": 44,
            "end_ix": 44,
            "relations": [["located_at", "6"]],
        },
        "10": {
            "tokens": "limits",
            "label": "OBS-DP",
            "start_ix": 45,
            "end_ix": 45,
            "relations": [["modify", "9"]],
        },
        "11": {
            "tokens": "acute",
            "label": "OBS-DA",
            "start_ix": 48,
            "end_ix": 48,
            "relations": [["modify", "13"]],
        },
        "12": {
            "tokens": "osseous",
            "label": "ANAT-DP",
            "start_ix": 49,
            "end_ix": 49,
            "relations": [],
        },
        "13": {
            "tokens": "abnormalities",
            "label": "OBS-DA",
            "start_ix": 50,
            "end_ix": 50,
            "relations": [["located_at", "12"]],
        },
        "14": {
            "tokens": "acute",
            "label": "OBS-DA",
            "start_ix": 55,
            "end_ix": 55,
            "relations": [["modify", "16"]],
        },
        "15": {
            "tokens": "cardiopulmonary",
            "label": "ANAT-DP",
            "start_ix": 56,
            "end_ix": 56,
            "relations": [],
        },
        "16": {
            "tokens": "process",
            "label": "OBS-DA",
            "start_ix": 57,
            "end_ix": 57,
            "relations": [["located_at", "15"]],
        },
    },
    "data_source": None,
    "data_split": "inference",
}
reference_annotations = {
    "text": "FINAL REPORT INDICATION : ___ F with cough / / Cough TECHNIQUE : PA and lateral views of the chest . COMPARISON : None . FINDINGS : The lungs are clear without focal consolidation , , or edema . The cardiomediastinal silhouette is within normal limits . No acute osseous abnormalities . IMPRESSION : No acute cardiopulmonary process .",
    "entities": {
        "1": {
            "tokens": "lungs",
            "label": "ANAT-DP",
            "start_ix": 28,
            "end_ix": 28,
            "relations": [],
        },
        "2": {
            "tokens": "clear",
            "label": "OBS-DP",
            "start_ix": 30,
            "end_ix": 30,
            "relations": [["located_at", "1"]],
        },
        "3": {
            "tokens": "focal",
            "label": "OBS-DA",
            "start_ix": 32,
            "end_ix": 32,
            "relations": [["modify", "4"]],
        },
        "4": {
            "tokens": "consolidation",
            "label": "OBS-DA",
            "start_ix": 33,
            "end_ix": 33,
            "relations": [["located_at", "1"]],
        },
        "5": {
            "tokens": "edema",
            "label": "OBS-DA",
            "start_ix": 37,
            "end_ix": 37,
            "relations": [],
        },
        "6": {
            "tokens": "cardiomediastinal",
            "label": "ANAT-DP",
            "start_ix": 40,
            "end_ix": 40,
            "relations": [],
        },
        "7": {
            "tokens": "silhouette",
            "label": "ANAT-DP",
            "start_ix": 41,
            "end_ix": 41,
            "relations": [["modify", "6"]],
        },
        "8": {
            "tokens": "within",
            "label": "OBS-DP",
            "start_ix": 43,
            "end_ix": 43,
            "relations": [],
        },
        "9": {
            "tokens": "normal",
            "label": "OBS-DP",
            "start_ix": 44,
            "end_ix": 44,
            "relations": [["located_at", "6"]],
        },
        "10": {
            "tokens": "limits",
            "label": "OBS-DP",
            "start_ix": 45,
            "end_ix": 45,
            "relations": [["modify", "9"]],
        },
        "11": {
            "tokens": "acute",
            "label": "OBS-DA",
            "start_ix": 48,
            "end_ix": 48,
            "relations": [["modify", "13"]],
        },
        "12": {
            "tokens": "osseous",
            "label": "ANAT-DP",
            "start_ix": 49,
            "end_ix": 49,
            "relations": [],
        },
        "13": {
            "tokens": "abnormalities",
            "label": "OBS-DA",
            "start_ix": 50,
            "end_ix": 50,
            "relations": [["located_at", "12"]],
        },
        "14": {
            "tokens": "acute",
            "label": "OBS-DA",
            "start_ix": 55,
            "end_ix": 55,
            "relations": [["modify", "16"]],
        },
        "15": {
            "tokens": "cardiopulmonary",
            "label": "ANAT-DP",
            "start_ix": 56,
            "end_ix": 56,
            "relations": [["located_at", "16"]],
        },
        "16": {
            "tokens": "process",
            "label": "OBS-DA",
            "start_ix": 57,
            "end_ix": 57,
            "relations": [],
        },
    },
    "data_source": None,
    "data_split": "inference",
}

# score = partially_exact_relation_token_and_label_match_reward(
#    hypothesis_annotations, reference_annotations
# )

# print(score)
# print(2 * (9.5 / 10) * (9.5 / 10) / (9.5 / 10 + 9.5 / 10))
# print(2 * (13 / 15) * (13 / 14) / (13 / 15 + 13 / 14))
