# ---------------------------------------------------------------------------
# coral_bleaching.py
#
# AI project on Computer Vision: telling healthy coral from bleached coral in
# underwater photographs, so that bleaching can be spotted early.
#
# CBSE Class X, Artificial Intelligence (417), Part C project work.
# SDG 14 (Life below Water) and SDG 13 (Climate Action).
#
# A computer sees a photo as a grid of pixels, and each pixel as three
# numbers from 0 to 255 - how much red, green and blue it holds. The program:
#   1. reads every photo into those numbers,
#   2. finds which pixels are coral and which are water (segmentation),
#   3. measures two FEATURES of the coral: how bright it is, and how
#      colourful it is, and
#   4. learns the average healthy coral and the average bleached coral from
#      the training photos, and labels a new photo by whichever it is closer to.
#
# Libraries: numpy (the pixel arithmetic), matplotlib (reading and showing
# images, and the charts), pandas (the table of features), os (the folders).
# ---------------------------------------------------------------------------

# %% Stage 1 - Problem scoping
print("=" * 60)
print("SPOTTING CORAL BLEACHING")
print("=" * 60)
print("Who   : marine scientists and reef divers")
print("What  : thousands of reef photos, too many to check by eye")
print("Where : coral reefs, such as those around Lakshadweep")
print("Why   : bleaching caught early shows where reefs need protection")
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", "bleached"]
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 photos  :", labels.count("healthy"))
print("Bleached photos :", labels.count("bleached"))
print("Total photos    :", len(paths))

# Read one photo and look at what the computer actually sees.
first = plt.imread(paths[0])
print("One photo is stored as an array of shape :", first.shape)
print("Height x width x colours                 : 64 x 64 x 3 = 12288 numbers")
pixel = (first[40, 32] * 255).round()
print("The pixel at row 40, column 32 (R, G, B) :", int(pixel[0]), int(pixel[1]), int(pixel[2]))
print()

# Chart 1: eight photos of each kind, spread through each folder.
healthy_paths = []
bleached_paths = []
for i in range(len(paths)):
    if labels[i] == "healthy":
        healthy_paths.append(paths[i])
    else:
        bleached_paths.append(paths[i])
plt.figure(figsize=(9, 2.6))
for i in range(8):
    plt.subplot(2, 8, i + 1)
    plt.imshow(plt.imread(healthy_paths[i * len(healthy_paths) // 8]))
    plt.axis("off")
    if i == 0:
        plt.title("healthy", fontsize=9, loc="left")
    plt.subplot(2, 8, i + 9)
    plt.imshow(plt.imread(bleached_paths[i * len(bleached_paths) // 8]))
    plt.axis("off")
    if i == 0:
        plt.title("bleached", fontsize=9, loc="left")
plt.tight_layout()
plt.savefig("chart1_sample_photos.png")
plt.close()

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

# Chart 2: one photo split into its red, green and blue layers.
plt.figure(figsize=(9, 2.8))
plt.subplot(1, 4, 1)
plt.imshow(first)
plt.title("Photo")
plt.axis("off")
channel_names = ["Red", "Green", "Blue"]
channel_maps = ["Reds", "Greens", "Blues"]
for ch in range(3):
    plt.subplot(1, 4, ch + 2)
    plt.imshow(first[:, :, ch], cmap=channel_maps[ch], vmin=0, vmax=1)
    plt.title(channel_names[ch])
    plt.axis("off")
plt.tight_layout()
plt.savefig("chart2_rgb_layers.png")
plt.close()

# Feature extraction. Water is blue-green with very little red, and coral -
# healthy or bleached - has a lot of red. So a pixel with red above 0.3
# (about 77 out of 255) is counted as coral. This is segmentation: deciding
# which pixels belong to the object.
rows = []
for i in range(len(paths)):
    img = plt.imread(paths[i])[:, :, :3].astype(float) * 255
    coral = img[:, :, 0] > 77
    pixels = img[coral]                                  # only the coral pixels
    brightness = pixels.mean()                            # 0 = black, 255 = white
    colourful = (pixels.max(axis=1) - pixels.min(axis=1)).mean()   # grey/white = 0
    rows.append([paths[i], labels[i], coral.sum(), round(brightness, 1), round(colourful, 1)])

features = pd.DataFrame(rows, columns=["photo", "label", "coral_pixels", "brightness", "colourfulness"])
print("Features measured from each photo (first 5 of each kind):")
print(features[features["label"] == "healthy"].head(5).to_string(index=False))
print(features[features["label"] == "bleached"].head(5).to_string(index=False))
print()
print("Average features of each kind:")
print(features.groupby("label")[["coral_pixels", "brightness", "colourfulness"]].mean().round(1))
print()

# Chart 3: the two features, one dot per photo.
plt.figure(figsize=(7, 5))
colours = ["#b5651d", "#7fb3d5"]
for k in range(2):
    part = features[features["label"] == classes[k]]
    plt.scatter(part["brightness"], part["colourfulness"], color=colours[k], label=classes[k], alpha=0.75)
plt.title("Two features measured from every photo")
plt.xlabel("Brightness of the coral (0-255)")
plt.ylabel("Colourfulness of the coral")
plt.legend()
plt.grid(True, linestyle="--", alpha=0.4)
plt.tight_layout()
plt.savefig("chart3_features.png")
plt.close()

# %% Stage 4 - Modelling
print("--- Stage 4: Modelling ---")
use = ["brightness", "colourfulness"]

features = features.sample(frac=1, random_state=8)
cut = int(len(features) * 0.8)
train = features[:cut]
test = features[cut:]
print("Training photos :", len(train))
print("Testing photos  :", len(test))

low = train[use].min()
high = train[use].max()
averages = ((train[use] - low) / (high - low)).groupby(train["label"]).mean()
print()
print("What the model learned - the average of each kind (rescaled 0 to 1):")
print(averages.round(2))
print()

test_scaled = (test[use] - low) / (high - low)
predictions = []
for i in range(len(test)):
    distance = ((averages - test_scaled.iloc[i]) ** 2).sum(axis=1)
    predictions.append(distance.idxmin())

# %% Stage 5 - Evaluation
print("--- Stage 5: Evaluation ---")
actual = list(test["label"])

# "bleached" is the positive class: it is what the scientists are looking for.
tp = fn = fp = tn = 0
for i in range(len(actual)):
    if actual[i] == "bleached" and predictions[i] == "bleached":
        tp = tp + 1
    elif actual[i] == "bleached":
        fn = fn + 1
    elif predictions[i] == "bleached":
        fp = fp + 1
    else:
        tn = tn + 1
accuracy = (tp + tn) / len(actual) * 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()

# Looking at the mistakes is the most useful check in computer vision.
wrong = []
for i in range(len(actual)):
    if predictions[i] != actual[i]:
        wrong.append(i)
print("Photos the model got wrong :", len(wrong))
for i in wrong:
    row = test.iloc[i]
    print("  " + os.path.basename(row["photo"]).ljust(18), "really", row["label"].ljust(8),
          " brightness", row["brightness"], " colourfulness", row["colourfulness"])
print()

# Chart 5: the photos it got wrong.
if len(wrong) > 0:
    plt.figure(figsize=(1.8 * len(wrong) + 0.5, 2.4))
    for n in range(len(wrong)):
        row = test.iloc[wrong[n]]
        plt.subplot(1, len(wrong), n + 1)
        plt.imshow(plt.imread(row["photo"]))
        plt.title("really " + row["label"], fontsize=9)
        plt.axis("off")
    plt.tight_layout()
    plt.savefig("chart5_mistakes.png")
    plt.close()

# Chart 4: confusion matrix.
grid = [[tp, fn], [fp, tn]]
plt.figure(figsize=(5, 4.3))
plt.imshow(grid, cmap="Blues")
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 bleached", "Said healthy"])
plt.yticks([0, 1], ["Bleached", "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.6))
for i in range(len(new_names)):
    img = plt.imread(os.path.join("new_photos", new_names[i]))[:, :, :3].astype(float) * 255
    pixels = img[img[:, :, 0] > 77]
    point = pd.Series({
        "brightness": pixels.mean(),
        "colourfulness": (pixels.max(axis=1) - pixels.min(axis=1)).mean(),
    })
    distance = ((averages - (point - low) / (high - low)) ** 2).sum(axis=1)
    answer = distance.idxmin()
    print(new_names[i], "-> brightness", round(point["brightness"]), " colourfulness",
          round(point["colourfulness"]), "->", answer.upper())
    plt.subplot(1, len(new_names), i + 1)
    plt.imshow(img.astype("uint8"))
    plt.title(answer, fontsize=10)
    plt.axis("off")
plt.tight_layout()
plt.savefig("chart6_new_photos.png")
plt.close()