# ---------------------------------------------------------------------------
# pin_reader.py
#
# Capstone project: reading handwritten PIN codes on envelopes, so that
# letters and parcels can be sorted by machine.
#
# CBSE Class XII, Artificial Intelligence (843), Part C - Capstone Project.
# SDG 9 (Industry, Innovation and Infrastructure).
#
# Data Science Methodology, 10 steps. COMPUTER VISION with a NEURAL
# NETWORK. The data is real: the UCI handwritten digits dataset (8 x 8
# images). A multi-layer perceptron is compared with Logistic Regression,
# with stratified 5-fold cross validation; then OpenCV loads, splits and
# resizes whole PIN-code images for the model to read.
#
# Libraries: pandas, numpy, matplotlib, scikit-learn, OpenCV (cv2), joblib.
# ---------------------------------------------------------------------------

# %% Step 1 - Business Understanding
print("=" * 64)
print("READING HANDWRITTEN PIN CODES")
print("=" * 64)
print("Problem : millions of envelopes carry handwritten PIN codes;")
print("          reading them by hand is slow and error-prone")
print("Goal    : recognise each handwritten digit, and so the whole PIN")
print()

# %% Step 2 - Analytic Approach
print("Analytic approach : predictive - classification into 10 classes (0-9)")
print()

# %% Step 3 - Data Requirements
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import cv2
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
from sklearn.linear_model import LogisticRegression
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import accuracy_score, confusion_matrix, recall_score
import joblib

# %% Step 4 - Data Collection
data = pd.read_csv("digits.csv")
pixels = ["px" + str(i) for i in range(64)]
print("Images collected :", len(data))
print("Values per image :", len(pixels), "(an 8 x 8 grid)")
print("Pixel values run from", data[pixels].values.min(), "to", data[pixels].values.max())
print()

# %% Step 5 - Data Understanding
print("--- Step 5: Data Understanding ---")
print("Missing values in the whole table :", int(data.isnull().sum().sum()))
counts = data["digit"].value_counts().sort_index()
print("Images of each digit:")
print(counts.to_string())
print()

# Chart 1: sixteen of the images, as the computer stores them.
plt.figure(figsize=(8, 4.2))
for i in range(16):
    plt.subplot(2, 8, i + 1)
    plt.imshow(data.loc[i, pixels].values.astype(float).reshape(8, 8), cmap="gray_r")
    plt.title(str(data.loc[i, "digit"]), fontsize=10)
    plt.axis("off")
plt.suptitle("Sixteen handwritten digits - each is just 64 numbers")
plt.tight_layout()
plt.savefig("chart1_samples.png")
plt.close()

# Chart 2: the average image of each digit.
plt.figure(figsize=(9, 2.4))
for d in range(10):
    plt.subplot(1, 10, d + 1)
    plt.imshow(data[data["digit"] == d][pixels].mean().values.reshape(8, 8), cmap="magma")
    plt.title(str(d))
    plt.axis("off")
plt.suptitle("The average of every image of each digit")
plt.tight_layout()
plt.savefig("chart2_average_digits.png")
plt.close()

# %% Step 6 - Data Preparation
print("--- Step 6: Data Preparation ---")
before = len(data)
data = data.drop_duplicates(subset=pixels + ["digit"])
print("Duplicate images removed :", before - len(data))
# The pixel values run 0-16. Dividing by 16 puts them on a 0-1 scale,
# which helps a neural network learn.
X = data[pixels] / 16.0
y = data["digit"]
print("Pixels rescaled from 0-16 to 0-1")
print()

# %% Step 7 - Modelling
print("--- Step 7: Modelling ---")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=3)
print("Training images :", len(X_train), " Testing images :", len(X_test))
logistic = LogisticRegression(max_iter=2000)
network = MLPClassifier(hidden_layer_sizes=(64,), max_iter=1000, random_state=0)
logistic.fit(X_train, y_train)
network.fit(X_train, y_train)
print("Neural network : 64 inputs -> 64 hidden neurons -> 10 outputs")
print("Weights learned by the network :", sum(w.size for w in network.coefs_) + sum(b.size for b in network.intercepts_))
print()

# Chart 3: the network's training - its loss falling with each pass.
plt.figure(figsize=(7, 4))
plt.plot(network.loss_curve_, color="#b07aa1")
plt.title("The neural network learning: loss after each pass")
plt.xlabel("Pass through the training images (epoch)")
plt.ylabel("Loss")
plt.grid(True, linestyle="--", alpha=0.4)
plt.tight_layout()
plt.savefig("chart3_loss_curve.png")
plt.close()

# %% Step 8 - Evaluation
print("--- Step 8: Evaluation ---")
pred_log = logistic.predict(X_test)
pred_net = network.predict(X_test)
print("Accuracy, Logistic Regression :", round(accuracy_score(y_test, pred_log), 4))
print("Accuracy, neural network      :", round(accuracy_score(y_test, pred_net), 4))
folds = StratifiedKFold(n_splits=5, shuffle=True, random_state=1)
cv_log = cross_val_score(LogisticRegression(max_iter=2000), X, y, cv=folds)
cv_net = cross_val_score(MLPClassifier(hidden_layer_sizes=(64,), max_iter=1000, random_state=0), X, y, cv=folds)
print("Accuracy in each fold, Logistic Regression :", np.round(cv_log, 3))
print("Accuracy in each fold, neural network      :", np.round(cv_net, 3))
print("Mean CV accuracy, Logistic Regression :", round(cv_log.mean(), 4))
print("Mean CV accuracy, neural network      :", round(cv_net.mean(), 4))
print()

cm = confusion_matrix(y_test, pred_net)
recall_each = recall_score(y_test, pred_net, average=None)
print("Recall for each digit (neural network):")
for d in range(10):
    print("  digit", d, ":", round(recall_each[d], 3))
pairs = []
for a in range(10):
    for b in range(10):
        if a != b and cm[a, b] > 0:
            pairs.append([cm[a, b], a, b])
pairs.sort(reverse=True)
print("Most common mistakes (true -> predicted, count):")
for count, a, b in pairs[:4]:
    print("  ", a, "->", b, ":", count)
print()

# Chart 4: the 10 x 10 confusion matrix.
plt.figure(figsize=(6.5, 5.5))
plt.imshow(cm, cmap="Blues")
for a in range(10):
    for b in range(10):
        if cm[a, b] > 0:
            plt.text(b, a, cm[a, b], ha="center", va="center", fontsize=9,
                     color="white" if cm[a, b] > cm.max() / 2 else "black")
plt.xticks(range(10))
plt.yticks(range(10))
plt.xlabel("Predicted digit")
plt.ylabel("True digit")
plt.title("Neural network: confusion matrix (test images)")
plt.tight_layout()
plt.savefig("chart4_confusion.png")
plt.close()

# Chart 5: the test images the network got wrong.
wrong = np.where(pred_net != y_test.values)[0]
plt.figure(figsize=(1.3 * max(len(wrong), 1) + 0.5, 2.0))
for k, i in enumerate(wrong):
    plt.subplot(1, len(wrong), k + 1)
    plt.imshow(X_test.iloc[i].values.reshape(8, 8), cmap="gray_r")
    plt.title(str(y_test.iloc[i]) + " -> " + str(pred_net[i]), fontsize=9)
    plt.axis("off")
plt.tight_layout()
plt.savefig("chart5_mistakes.png")
plt.close()

# Chart 6: accuracy in each fold, both models.
x = np.arange(5)
plt.figure(figsize=(7.5, 4))
plt.bar(x - 0.2, cv_log, width=0.4, label="Logistic Regression", color="#9c9c9c")
plt.bar(x + 0.2, cv_net, width=0.4, label="Neural network", color="#b07aa1")
plt.xticks(x, ["Fold 1", "Fold 2", "Fold 3", "Fold 4", "Fold 5"])
plt.ylim(0.9, 1.0)
plt.ylabel("Accuracy")
plt.title("Stratified 5-fold cross validation")
plt.legend()
plt.tight_layout()
plt.savefig("chart6_kfold.png")
plt.close()

# %% Step 9 - Deployment
print("--- Step 9: Deployment ---")
joblib.dump(network, "digit_model.pkl")
model = joblib.load("digit_model.pkl")
truth = pd.read_csv("pin_truth.csv", dtype={"true_pin": str})
results = []
for k in range(len(truth)):
    name = truth.loc[k, "file"]
    strip = cv2.imread(os.path.join("pin_codes", name), cv2.IMREAD_GRAYSCALE)   # load with OpenCV
    if strip is None:                         # imread gives None, not an error, if the file is missing
        raise FileNotFoundError("Could not read pin_codes/" + name)
    strip = 255 - strip                                                          # ink becomes bright
    read = ""
    for j in range(6):
        digit_img = strip[:, j * 64:(j + 1) * 64]                                 # cut out one digit
        small = cv2.resize(digit_img, (8, 8), interpolation=cv2.INTER_AREA)       # resize to 8 x 8
        vector = (small / 255.0 * 16).round() / 16.0                              # same 0-1 scale as training
        read = read + str(model.predict(vector.reshape(1, 64))[0])
    results.append([name, truth.loc[k, "true_pin"], read])
    status = "correct" if read == truth.loc[k, "true_pin"] else "WRONG"
    print(name, " true", truth.loc[k, "true_pin"], " read", read, " ", status)
print()

# Chart 7: four of the PIN images, with what the model read.
plt.figure(figsize=(9, 4.5))
for n, k in enumerate([0, 1, 7, 9]):
    img = cv2.imread(os.path.join("pin_codes", results[k][0]), cv2.IMREAD_GRAYSCALE)
    plt.subplot(4, 1, n + 1)
    plt.imshow(img, cmap="gray")
    plt.title("read as " + results[k][2] + "   (true " + results[k][1] + ")", fontsize=10)
    plt.axis("off")
plt.tight_layout()
plt.savefig("chart7_pin_codes.png")
plt.close()

# %% Step 10 - Feedback
print("--- Step 10: Feedback ---")
digits_right = 0
pins_right = 0
clean_right = 0
messy_right = 0
for k in range(len(results)):
    true_pin, read = results[k][1], results[k][2]
    digits_right = digits_right + sum(1 for a, b in zip(true_pin, read) if a == b)
    if true_pin == read:
        pins_right = pins_right + 1
        if k < 6:
            clean_right = clean_right + 1
        else:
            messy_right = messy_right + 1
print("Digits read correctly on envelopes :", digits_right, "of", 6 * len(results))
print("Whole PIN codes read correctly     :", pins_right, "of", len(results))
print("Clean envelopes, PINs correct      :", clean_right, "of 6")
print("Messy envelopes, PINs correct      :", messy_right, "of 4")
per_digit = digits_right / (6 * len(results))
print("If each digit is right with probability", round(per_digit, 3), ", a 6-digit PIN is all right")
print("with probability about", round(per_digit ** 6, 3))