from . import healthbench_meta_eval

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))


def test_compute_agreement_for_rater_by_class():
    self_pred_list = [True, False, True]
    other_preds_list = [[True, True, False], [True, False], [False]]
    cluster_list = ["a", "a", "b"]
    model_or_physician = "model"
    metrics = healthbench_meta_eval.compute_metrics_for_rater_by_class(
        self_pred_list, other_preds_list, cluster_list, model_or_physician
    )

    # precision overall
    index_str_pos_precision = healthbench_meta_eval.INDEX_STR_TEMPLATE.format(
        model_or_physician=model_or_physician, metric="precision", pred_str="pos"
    )
    index_str_neg_precision = healthbench_meta_eval.INDEX_STR_TEMPLATE.format(
        model_or_physician=model_or_physician, metric="precision", pred_str="neg"
    )
    overall_pos_precision = metrics[index_str_pos_precision]
    overall_neg_precision = metrics[index_str_neg_precision]
    expected_overall_pos_precision = (2 + 0 + 0) / (3 + 0 + 1)
    expected_overall_neg_precision = (0 + 1 + 0) / (0 + 2 + 0)
    assert overall_pos_precision["value"] == expected_overall_pos_precision
    assert overall_neg_precision["value"] == expected_overall_neg_precision
    assert overall_pos_precision["n"] == 4
    assert overall_neg_precision["n"] == 2

    # recall overall
    index_str_pos_recall = healthbench_meta_eval.INDEX_STR_TEMPLATE.format(
        model_or_physician=model_or_physician, metric="recall", pred_str="pos"
    )
    index_str_neg_recall = healthbench_meta_eval.INDEX_STR_TEMPLATE.format(
        model_or_physician=model_or_physician, metric="recall", pred_str="neg"
    )
    overall_pos_recall = metrics[index_str_pos_recall]
    overall_neg_recall = metrics[index_str_neg_recall]
    expected_overall_pos_recall = (2 + 0 + 0) / (2 + 1 + 0)
    expected_overall_neg_recall = (0 + 1 + 0) / (1 + 1 + 1)
    assert overall_pos_recall["value"] == expected_overall_pos_recall
    assert overall_neg_recall["value"] == expected_overall_neg_recall
    assert overall_pos_recall["n"] == 3
    assert overall_neg_recall["n"] == 3

    # f1 overall
    index_str_pos_f1 = healthbench_meta_eval.INDEX_STR_TEMPLATE.format(
        model_or_physician=model_or_physician, metric="f1", pred_str="pos"
    )
    index_str_neg_f1 = healthbench_meta_eval.INDEX_STR_TEMPLATE.format(
        model_or_physician=model_or_physician, metric="f1", pred_str="neg"
    )
    overall_pos_f1 = metrics[index_str_pos_f1]
    overall_neg_f1 = metrics[index_str_neg_f1]
    expected_overall_pos_f1 = (
        2
        * expected_overall_pos_precision
        * expected_overall_pos_recall
        / (expected_overall_pos_precision + expected_overall_pos_recall)
    )
    expected_overall_neg_f1 = (
        2
        * expected_overall_neg_precision
        * expected_overall_neg_recall
        / (expected_overall_neg_precision + expected_overall_neg_recall)
    )
    assert overall_pos_f1["value"] == expected_overall_pos_f1
    assert overall_neg_f1["value"] == expected_overall_neg_f1

    # balanced f1
    index_str_balanced_f1 = healthbench_meta_eval.INDEX_STR_TEMPLATE.format(
        model_or_physician=model_or_physician, metric="f1", pred_str="balanced"
    )
    balanced_f1 = metrics[index_str_balanced_f1]
    expected_balanced_f1 = (expected_overall_pos_f1 + expected_overall_neg_f1) / 2
    assert balanced_f1["value"] == expected_balanced_f1

    # by cluster
    # precision
    cluster_a_str_pos_precision = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="a", index_str=index_str_pos_precision
    )
    cluster_a_str_neg_precision = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="a", index_str=index_str_neg_precision
    )
    cluster_a_pos_precision = metrics[cluster_a_str_pos_precision]
    cluster_a_neg_precision = metrics[cluster_a_str_neg_precision]
    assert cluster_a_pos_precision["value"] == (
        # example 1, 2 in order
        (2 + 0) / (3 + 0)
    )
    assert cluster_a_neg_precision["value"] == (
        # example 1, 2 in order
        (0 + 1) / (0 + 2)
    )
    assert cluster_a_pos_precision["n"] == 3
    assert cluster_a_neg_precision["n"] == 2

    # recall
    cluster_a_str_pos_recall = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="a", index_str=index_str_pos_recall
    )
    cluster_a_str_neg_recall = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="a", index_str=index_str_neg_recall
    )
    cluster_a_pos_recall = metrics[cluster_a_str_pos_recall]
    cluster_a_neg_recall = metrics[cluster_a_str_neg_recall]
    assert cluster_a_pos_recall["value"] == (
        # example 1, 2 in order
        (2 + 0) / (2 + 1)
    )
    assert cluster_a_neg_recall["value"] == (
        # example 1, 2 in order
        (0 + 1) / (1 + 1)
    )
    assert cluster_a_pos_recall["n"] == 3
    assert cluster_a_neg_recall["n"] == 2

    # cluster B
    # precision
    cluster_b_str_pos_precision = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="b", index_str=index_str_pos_precision
    )
    cluster_b_str_neg_precision = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="b", index_str=index_str_neg_precision
    )
    cluster_b_str_pos_precision = metrics[cluster_b_str_pos_precision]
    assert cluster_b_str_neg_precision not in metrics
    assert cluster_b_str_pos_precision["value"] == (
        # example 3 only
        0 / 1
    )
    assert cluster_b_str_pos_precision["n"] == 1

    # recall
    cluster_b_str_pos_recall = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="b", index_str=index_str_pos_recall
    )
    cluster_b_str_neg_recall = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="b", index_str=index_str_neg_recall
    )
    assert cluster_b_str_pos_recall not in metrics
    cluster_b_neg_recall = metrics[cluster_b_str_neg_recall]
    assert cluster_b_neg_recall["value"] == (
        # example 3 only
        0 / 1
    )
    assert cluster_b_neg_recall["n"] == 1

    # f1
    index_str_pos_f1 = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="b", index_str=index_str_pos_f1
    )
    index_str_neg_f1 = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="b", index_str=index_str_neg_f1
    )
    index_str_balanced_f1 = healthbench_meta_eval.CLUSTER_STR_TEMPLATE.format(
        cluster="b", index_str=index_str_balanced_f1
    )
    assert index_str_pos_f1 not in metrics
    assert index_str_neg_f1 not in metrics
    assert index_str_balanced_f1 not in metrics


if __name__ == "__main__":
    test_compute_agreement_for_rater_by_class()
