# ---------------------------------------------------------------------------
# leaf_disease.py
#
# AI project on Computer Vision: spotting a diseased crop leaf from a photo,
# and saying how much of the leaf is affected.
#
# CBSE Class X, Artificial Intelligence (417), Part C project work.
# SDG 2 (Zero Hunger) and SDG 12 (Responsible Consumption and Production).
#
# The leaf is photographed on a plain sheet of paper. The program then:
#   1. finds which pixels are leaf and which are paper (segmentation),
#   2. counts how many leaf pixels are brown or yellow instead of green,
#   3. turns that into one FEATURE: the percentage of the leaf affected, and
#   4. learns, from the training photos, the cut-off percentage that best
#      separates healthy leaves from diseased ones.
#
# Libraries: numpy (pixel arithmetic), matplotlib (reading and showing
# images, and the charts), pandas (the table of features), os (folders).
# ---------------------------------------------------------------------------

# %% Stage 1 - Problem scoping
print("=" * 60)
print("IS THIS LEAF DISEASED?")
print("=" * 60)
print("Who   : farmers, especially those far from an expert")
print("What  : leaf disease noticed late spreads through the field")
print("Where : vegetable and crop fields, in the growing season")
print("Why   : early action saves the harvest and avoids extra spraying")
print()

# %% Stage 2 - Data acquisition
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

pd.set_option("display.width", 120)

classes = ["healthy", "diseased"]
paths = []
labels = []
for c in classes:
    for name in sorted(os.listdir(os.path.join("images", c))):
        paths.append(os.path.join("images", c, name))
        labels.append(c)

print("--- Stage 2: Data acquisition ---")
print("Healthy leaves  :", labels.count("healthy"))
print("Diseased leaves :", labels.count("diseased"))
print("Total photos    :", len(paths))
print("Size of a photo :", plt.imread(paths[0]).shape)
print()

healthy_paths = []
diseased_paths = []
for i in range(len(paths)):
    if labels[i] == "healthy":
        healthy_paths.append(paths[i])
    else:
        diseased_paths.append(paths[i])

# Chart 1: six leaves of each kind.
plt.figure(figsize=(9, 3.2))
for i in range(6):
    plt.subplot(2, 6, i + 1)
    plt.imshow(plt.imread(healthy_paths[i * len(healthy_paths) // 6]))
    plt.axis("off")
    if i == 0:
        plt.title("healthy", fontsize=9, loc="left")
    plt.subplot(2, 6, i + 7)
    plt.imshow(plt.imread(diseased_paths[i * len(diseased_paths) // 6]))
    plt.axis("off")
    if i == 0:
        plt.title("diseased", fontsize=9, loc="left")
plt.tight_layout()
plt.savefig("chart1_sample_leaves.png")
plt.close()

# %% Stage 3 - Data exploration
print("--- Stage 3: Data exploration ---")

# Feature extraction, done the same way for every photo:
#   - Paper is grey: its red, green and blue are almost equal. A pixel whose
#     largest and smallest colour differ by more than 35 is leaf.
#   - Healthy leaf is green: much more green than red. A leaf pixel whose
#     red is at least 0.75 of its green is brown or yellow - affected.
rows = []
for i in range(len(paths)):
    img = plt.imread(paths[i])[:, :, :3].astype(float) * 255
    spread = img.max(axis=2) - img.min(axis=2)
    leaf = spread > 35
    affected = leaf & (img[:, :, 0] >= 0.75 * img[:, :, 1])
    percent = affected.sum() / leaf.sum() * 100
    rows.append([paths[i], labels[i], leaf.sum(), affected.sum(), round(percent, 1)])

data = pd.DataFrame(rows, columns=["photo", "label", "leaf_pixels", "affected_pixels", "affected_pct"])
print("Average of each kind:")
print(data.groupby("label")[["leaf_pixels", "affected_pixels", "affected_pct"]].mean().round(1))
print()
print("Healthy leaves  - most affected :", data[data["label"] == "healthy"]["affected_pct"].max(), "%")
print("Diseased leaves - least affected:", data[data["label"] == "diseased"]["affected_pct"].min(), "%")
print()

# Chart 2: what the program sees on one diseased leaf.
example = plt.imread(diseased_paths[0])[:, :, :3]
img = example.astype(float) * 255
leaf = (img.max(axis=2) - img.min(axis=2)) > 35
affected = leaf & (img[:, :, 0] >= 0.75 * img[:, :, 1])
marked = example.copy()
marked[affected] = [1.0, 0.0, 0.0]
plt.figure(figsize=(8, 3))
plt.subplot(1, 3, 1)
plt.imshow(example)
plt.title("Photo")
plt.axis("off")
plt.subplot(1, 3, 2)
plt.imshow(leaf, cmap="gray")
plt.title("Leaf pixels (white)")
plt.axis("off")
plt.subplot(1, 3, 3)
plt.imshow(marked)
plt.title("Affected pixels (red)")
plt.axis("off")
plt.tight_layout()
plt.savefig("chart2_pixel_masks.png")
plt.close()

# %% Stage 4 - Modelling
print("--- Stage 4: Modelling ---")
data = data.sample(frac=1, random_state=4)
cut = int(len(data) * 0.8)
train = data[:cut]
test = data[cut:]
print("Training photos :", len(train))
print("Testing photos  :", len(test))

# Learn the cut-off: try every value from 0% to 40% in steps of 0.5, count
# the training mistakes of "diseased if affected_pct is at least this", and
# keep the cut-off with the fewest.
truth = train["label"] == "diseased"
best_cutoff = 0
best_mistakes = len(train) + 1
cutoffs = np.arange(0, 40.5, 0.5)
for c in cutoffs:
    mistakes = ((train["affected_pct"] >= c) != truth).sum()
    if mistakes < best_mistakes:
        best_mistakes = mistakes
        best_cutoff = c
print("Learned cut-off            :", best_cutoff, "% of the leaf affected")
print("Mistakes on training data  :", best_mistakes, "of", len(train))
print()

# Chart 3: the feature for both kinds, with the learned cut-off.
plt.figure(figsize=(7, 4.3))
plt.hist(train[train["label"] == "healthy"]["affected_pct"], bins=20, alpha=0.65, color="#59a14f", label="healthy")
plt.hist(train[train["label"] == "diseased"]["affected_pct"], bins=20, alpha=0.65, color="#b07aa1", label="diseased")
plt.axvline(best_cutoff, color="#e15759", linestyle="--", label="learned cut-off: " + str(best_cutoff) + "%")
plt.title("Share of the leaf affected (training photos)")
plt.xlabel("Affected pixels (% of the leaf)")
plt.ylabel("Number of leaves")
plt.legend()
plt.tight_layout()
plt.savefig("chart3_affected_share.png")
plt.close()

# %% Stage 5 - Evaluation
print("--- Stage 5: Evaluation ---")
tp = fn = fp = tn = 0
wrong = []
for i in range(len(test)):
    row = test.iloc[i]
    said_diseased = row["affected_pct"] >= best_cutoff
    if row["label"] == "diseased" and said_diseased:
        tp = tp + 1
    elif row["label"] == "diseased":
        fn = fn + 1
        wrong.append(i)
    elif said_diseased:
        fp = fp + 1
        wrong.append(i)
    else:
        tn = tn + 1
accuracy = (tp + tn) / len(test) * 100
precision = 0
if tp + fp > 0:
    precision = tp / (tp + fp) * 100
recall = 0
if tp + fn > 0:
    recall = tp / (tp + fn) * 100
f1 = 0
if precision + recall > 0:
    f1 = 2 * precision * recall / (precision + recall)
print("TP, FN, FP, TN :", tp, fn, fp, tn)
print("Accuracy       :", round(accuracy, 1), "%")
print("Precision      :", round(precision, 1), "%")
print("Recall         :", round(recall, 1), "%")
print("F1 score       :", round(f1, 1), "%")
print("Photos wrong   :", len(wrong))
for i in wrong:
    row = test.iloc[i]
    print("  " + os.path.basename(row["photo"]), "really", row["label"], "-", row["affected_pct"], "% affected")
print()

# Chart 4: confusion matrix.
grid = [[tp, fn], [fp, tn]]
plt.figure(figsize=(5, 4.3))
plt.imshow(grid, cmap="Greens")
names = [["TP", "FN"], ["FP", "TN"]]
for r in range(2):
    for c in range(2):
        if grid[r][c] > max(max(grid)) / 2:
            ink = "white"
        else:
            ink = "black"
        plt.text(c, r, names[r][c] + "\n" + str(grid[r][c]), ha="center", va="center", fontsize=14, color=ink)
plt.xticks([0, 1], ["Said diseased", "Said healthy"])
plt.yticks([0, 1], ["Diseased", "Healthy"])
plt.xlabel("What the model said")
plt.ylabel("The truth")
plt.title("Confusion matrix (testing photos)")
plt.tight_layout()
plt.savefig("chart4_matrix.png")
plt.close()

# %% Stage 6 - Deployment
print("--- Stage 6: Deployment ---")
new_names = sorted(os.listdir("new_photos"))
plt.figure(figsize=(8, 2.8))
for i in range(len(new_names)):
    photo = plt.imread(os.path.join("new_photos", new_names[i]))[:, :, :3]
    img = photo.astype(float) * 255
    leaf = (img.max(axis=2) - img.min(axis=2)) > 35
    affected = leaf & (img[:, :, 0] >= 0.75 * img[:, :, 1])
    percent = affected.sum() / leaf.sum() * 100
    if percent < best_cutoff:
        verdict = "healthy"
    elif percent < 15:
        verdict = "mild"
    elif percent < 30:
        verdict = "moderate"
    else:
        verdict = "severe"
    print(new_names[i], "->", round(percent, 1), "% affected ->", verdict.upper())
    plt.subplot(1, len(new_names), i + 1)
    plt.imshow(photo)
    plt.title(verdict + " (" + str(round(percent)) + "%)", fontsize=10)
    plt.axis("off")
plt.tight_layout()
plt.savefig("chart5_new_leaves.png")
plt.close()
print("Mild, moderate or severe: remove affected leaves and show the photo to")
print("the nearest Krishi Vigyan Kendra before spraying anything.")