# Program 15: Confusion matrix, accuracy, precision, recall and F1 score
# The activity from the CBSE AI handbook: a model is tested on 1000
# students - 900 actually passed (Yes) and 100 did not (No). A faulty
# model predicts "Yes" for everybody. Then a better model is tested.

import matplotlib.pyplot as plt

actual = ["Yes"] * 900 + ["No"] * 100
faulty = ["Yes"] * 1000
# A better model: right for most students, wrong for a few.
better = ["Yes"] * 870 + ["No"] * 30 + ["No"] * 85 + ["Yes"] * 15

names = ["Faulty model (always Yes)", "Better model"]
for k in range(2):
    predicted = [faulty, better][k]
    tp = tn = fp = fn = 0
    for i in range(len(actual)):
        if actual[i] == "Yes" and predicted[i] == "Yes":
            tp = tp + 1
        elif actual[i] == "No" and predicted[i] == "No":
            tn = tn + 1
        elif actual[i] == "No" and predicted[i] == "Yes":
            fp = fp + 1
        else:
            fn = fn + 1
    accuracy = (tp + tn) / (tp + tn + fp + fn)
    precision = tp / (tp + fp)
    recall = tp / (tp + fn)
    f1 = 2 * precision * recall / (precision + recall)
    # For the "No" class: how many of the students who failed were found?
    no_found = tn / (tn + fp)

    print(names[k])
    print("  TP =", tp, " FN =", fn, " FP =", fp, " TN =", tn)
    print("  Accuracy  =", round(accuracy, 3), "=", round(accuracy * 100, 1), "%")
    print("  Precision =", round(precision, 3))
    print("  Recall    =", round(recall, 3))
    print("  F1 score  =", round(f1, 3))
    print("  Students who did not pass, found by the model:", tn, "of", tn + fp)
    print()

    plt.subplot(1, 2, k + 1)
    grid = [[tp, fn], [fp, tn]]
    plt.imshow(grid, cmap="Blues")
    labels = [["TP", "FN"], ["FP", "TN"]]
    for r in range(2):
        for c in range(2):
            if grid[r][c] > 450:
                ink = "white"
            else:
                ink = "black"
            plt.text(c, r, labels[r][c] + "\n" + str(grid[r][c]), ha="center", va="center",
                     color=ink, fontsize=11)
    plt.xticks([0, 1], ["Pred Yes", "Pred No"])
    plt.yticks([0, 1], ["Actual Yes", "Actual No"])
    plt.title(names[k], fontsize=10)
plt.gcf().set_size_inches(9, 4)
plt.tight_layout()
plt.show()