import os
import cv2

import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)  # all RuntimeWarnings

import numpy as np
import matplotlib.pyplot as plt

# Set random seed for NumPy’s random number generator
#     Conceptually: makes random operations (e.g., random initialization in KMeans, shuffling data, or selecting random samples) produce the same results every time the code is run
#     aka your code becomes reporducable
np.random.seed(42)

# If running on colab and you don't already have the necessary HW data, reclone it
REPO_URL = "https://github.com/nd-cse-30124-fa25/cse-30124-homeworks.git"
REPO_NAME = "cse-30124-homeworks"
HW_FOLDER = "homework03"

try:
    import google.colab

    # Clone repo if not already present
    if not os.path.exists(REPO_NAME):
        !git clone {REPO_URL}

    # cd into the homework folder
    %cd {REPO_NAME}/{HW_FOLDER}

    !git clone https://github.com/rapidsai/rapidsai-csp-utils.git
    !python rapidsai-csp-utils/colab/pip-install.py

except ImportError:
    print('Assignment does not appear to be running on colab, not cloning data repo nor installing CUML')
    print('\t[WARNING]: This means that your SVM will run very, very slowly')
    print('\t[WARNING]: It took me about 25 minutes to run it on my mac mini')
    print('\t[WARNING]: You could install CuML manually')

# Load letters (modify these according to your actual data paths if needed)
image_paths = [f'pages/note_page_{i}.png' for i in range(1, 5)]
mask_paths = [f'pages/mask_page_{i}.png' for i in range(1, 5)]

images = [cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB) for image_path in image_paths]
masks = [cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) for mask_path in mask_paths]

print('Image Dimensions:', images[0].shape)
print('Mask Dimensions:', masks[0].shape)

class PCA:
    def __init__(self, n_components=1.0):
        """
        Basic PCA class

        Parameters
        ----------
        n_components:
          - float in (0, 1] -> keep the minimum # of components to reach this
                               fraction of explained variance
        """
        self._expl_var = n_components

    def fit(self, X):
        """
        Fits class to input data

        Parameters
        ----------

        X : array-like, shape (n_samples, n_features)
            Data matrix.
        """

        X = np.asarray(X, dtype=float)
        n_samples, _ = X.shape

        # TODO: Center the data

        # TODO: Use Singular Value Decomposition to get the Singular Values and Explained Variance
        # Hint: numpy has a built in function for calculating SVD
        
        # TODO: Decide how many components to keep based on n_components
        
        # Retain only the resolved number of components and related statistics
        self.components_ = Vt[:self._n_components]

        return self

    def transform(self, X):
        """
        Projects the input data into the space defined by the principal components

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Data matrix.

        Returns
        -------
        Z : array-like, shape (n_samples, self._n_components)
            Data matrix.
        """

        # TODO: Transform the dataset into the new space described by the principal components

        return Z

    def fit_transform(self, X):
        return self.fit(X).transform(X)

class kmeans():
    def __init__(self, n_clusters, max_iters=300, tol=1e-4, random_state=None):
        """
        Basic k-means clustering class

        Parameters
        ---------
        n_clusters : int
            Number of clusters.
        max_iters : int
            Maximum iterations per run.
        tol : float
            Convergence tolerance on centroid movement (L2 norm).
        random_state : int or None
            Seed for reproducibility.
        """

        self._n_clusters = n_clusters
        self._max_iters = max_iters
        self._tol = tol
        self._random_state = random_state

    def _init_centroids(self, X):
        """
        Randomly initializes centroids

        Parameters
        ----------

        X : array-like, shape (n_samples, n_features)
            Data matrix.

        Returns
        -------
        X[idx].copy : ndarray, shape (n_clusters,)
            Randomly selected centroids
        """

        rng = np.random.default_rng(self._random_state)
        n_samples, _ = X.shape

        # TODO: Randomly select n_clusters centroids

        return X[idx].copy()

    def fit_predict(self, X):
        """
        Basic k-means clustering.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Data matrix.

        Returns
        -------
        np.logical_not(labels.astype(bool)) : ndarray, shape (n_samples,)
            Cluster assignment for each sample
        """
        X = np.asarray(X, dtype=float)

        labels = None
        centroids = None

        # TODO: Initialize centroids

        for _ in range(self._max_iters):
            # TODO: Calculate euclidian distances for each point from the centroids
            # TODO: Assign each sample to the closest centroid

            # TODO: For each cluster that has members, calculate the average position of members
            # TODO: For each cluster that doesn't have members, reseed it to the furthest point
            # TODO: Replace the original Centroids with these new average positions

            # TODO: Calculate the centroid movement and compare to tol

            if shift < self._tol:
                break

        # Flip the foreground and background labels to match the sklearn output
        return np.logical_not(labels.astype(bool))

def calculate_segmentation_metrics(pred, truth, visualize=False):
    """
    Calculate Intersection over Union (IoU) to measure segmentation accuracy

    Args:
        pred: Binary prediction mask (595, 420)
        truth: Binary ground truth mask (595, 420)


    Returns:
        iou: Intersection over Union (IoU) metric
    """
    if visualize:
      # Visualize the masks
      plt.figure(figsize=(10, 5))

      plt.subplot(1, 2, 1)
      plt.title('Prediction Mask')
      plt.imshow(pred, cmap='gray')
      plt.axis('off')

      plt.subplot(1, 2, 2)
      plt.title('Ground Truth Mask')
      plt.imshow(truth, cmap='gray')
      plt.axis('off')

      plt.show()

    # TODO: Calculate intersection and union

    # TODO: calculate IoU (Intersection over Union)
    #    hint: this is also known as the Jaccard index

    return iou

def header_printer(i):
    print(f"Image {i+1}:")
    print('#' * 80)
    print()
    print(' ' * 25 + 'EV: 0.75' + ' ' * 40 + 'EV: 0.9' + ' ' * 40 + 'EV: 0.99')
    headers = ' ' * 14 + "Student kmeans" + ' ' * 5 + "sklearn kmeans"
    print(headers * 3)

def row_formatter(ious):
    return f'{ious[0]:.4f}' + ' ' * 12 + f'{ious[1]:.4f}' + ' ' * 23

from sklearn.preprocessing import StandardScaler

from sklearn.decomposition import PCA as sklearn_pca
from sklearn.cluster import KMeans as sklearn_kmeans

print("\nEvaluating Unsupervised Segmentation Methods...")
print('#' * 80)
print()

# Initialize pca methods and kmeans methods
pca_implementations = {'Student PCA': PCA, 'sklearn PCA': sklearn_pca}
kmeans_implementations = {'Student kmeans': kmeans(n_clusters=2, random_state=42), 'sklearn kmeans': sklearn_kmeans(n_clusters=2, init='random', random_state=42)}

for i, image in enumerate(images):
    header_printer(i)

    for pca_creator, pca_method in pca_implementations.items():
        results = ' ' * 8

        for offset, ev in enumerate([0.75, 0.9, .99]):
            ious = []
            for kmeans_creator, kmeans_method in kmeans_implementations.items():
                # Reshape the image into a 2D array of pixels and 3 color values (RGB)
                pixels = image.reshape(-1, 3)          # reshaped into a 2D array where each row represents a pixel and the columns represent the RGB color channels
                                                    # Note: The -1 means "reshape it into as many rows as needed" (i.e., one row per pixel), and 3 is the number of columns corresponding to the RGB channels

                # TODO: Apply PCA to reduce dimensionality

                # TODO: Perform K-Means clustering, fit Kmeans to data to get cluster labels for each pixel

                # Reshape the labels back to the image shape
                segmented_image = labels.reshape(image.shape[:2])

                # Calculate IoU to evaluate clustering
                ious.append(calculate_segmentation_metrics(segmented_image.astype(bool), masks[i].astype(bool)))

            results += row_formatter(ious)
        print(f"{pca_creator + results}")
    print()

def extract_features(image, patch_size=5):
    """
    Extract features for each pixel using RGB values and local neighborhood.

    Args:
        image: RGB image (595, 420, 3)
        patch_size: Size of the local neighborhood patch (must be odd)

    Returns:
        features: Array of shape (n_pixels, n_features) (249900, 78) (595 x 420, 5 x 5 x 3 + 1 x 3)
    """
    height, width, channels = image.shape

    # Add padding to image
    padding = patch_size // 2

    padded = cv2.copyMakeBorder(
        image,
        padding, padding, padding, padding,
        cv2.BORDER_REFLECT
    )

    # TODO: Calculate total number of pixels

    # TODO: Calculate total number of features per pixel

    # TODO: Initialize data structue to store features for all pixels

    # TODO: Loop through pixels and extracts features for each pixel
    #     hint: the feature vector will be all RGB values for every surrounding pixel concatenated with the pixel itself
    #     hint: this results in a 78 dimensional feature vector, for each pixel!
    #     hint: 5 x 5 patch = 25, 3 colors per pixel gives 3 x 25 = 75 + 3 for the pixel itself again gives 78
    #     hint: use current pixel RGB values and the surrounding local patch
    pixel_idx = 0

    # Return feature matrix
    return features

try:
    import cudf
    from cuml.svm import SVC, LinearSVC
except:
    from sklearn.svm import SVC, LinearSVC

# -------------------------------------
# Task 5: SVM Segmentation
# -------------------------------------
def svm_segmentation(model, pca, test_image):
    # TODO: Extract features of testing image

    # TODO: Apply transformations to testing image if PCA != None
    #     hint: transformations means scaling and dimensionality reduction

    # Attempt to move test features to GPU
    try:
        X_test = cudf.DataFrame.from_records(X_test)
    except:
        pass

    # TODO: Predict on test image

    # Attempt to retrieve results from GPU
    try:
        predictions = predictions.to_cupy().get()  # get() brings it to CPU as np.ndarray
    except:
        pass

    # Reshape predictions to match image matrix shape
    segmented_image = predictions.reshape(test_image.shape[:2])

    return segmented_image

print("Evaluating Supervised Segmentation Methods...")
print('#' * 80)
print()

# We only use the first 1000 pixels from the training image and mask, otherwise it takes too long to train
# We'll train on only the first 1000 pixels from the first image, and test only on the second (again for sake of runtimes)
# Really we should be doing some 4 fold cross validation with full images here but I don't want to wait 9 hours for it to run
train_image = images[0][:1000]
train_mask = masks[0][:1000]

test_image = images[1]
test_mask = masks[1]

# TODO: Extract features of training image

y_train = (train_mask > 0).reshape(-1)

# TODO: Declare the kernels we want to test
for kernel in []: 
    print(f"Testing an SVM with a {kernel} kernel:")

    # TODO: Declare the explained variances we want to test
    for ev in []:
        # TODO: Apply PCA

        # Attempt to move to GPU (as cuDF)
        try:
            X_train_pca = cudf.DataFrame.from_records(X_train_pca)
            y_train = cudf.Series(y_train)
        except:
            pass

        # TODO: Initialize the SVM using LinearSVC for the linear kernel and SVC for everything else

        # TODO: Fit the SVM to the training data

        # TODO: segment the test image and calculate the IoU for it compared to the test mask
        
        print(f"\tAn explained variance of {ev} results in an IoU of {iou:.2f}")

    print()
    print('#' * 80)
    print()

from skimage.measure import label, regionprops
import os

def extract_and_save_letters(segmented_image, original_image, page, output_dir='letters'):
    """
    Extract each segmented letter and save as a separate PNG file.

    Args:
        segmented_image: Segmented image with unique labels for each letter
        original_image: Original RGB image
        output_dir: Directory to save the extracted letter images
    """
    # Ensure the output directory exists
    os.makedirs(output_dir, exist_ok=True)

    # Label connected components
    labeled_image = label(segmented_image)

    # Iterate over each labeled region
    for region in regionprops(labeled_image):
        # Extract the bounding box of the region
        min_row, min_col, max_row, max_col = region.bbox

        # Extract the region from the original image
        letter_image = original_image[min_row:max_row, min_col:max_col]

        # Save the extracted letter as a PNG file
        letter_filename = os.path.join(output_dir, f'note_{page}_letter_{region.label}.png')
        cv2.imwrite(letter_filename, cv2.cvtColor(letter_image, cv2.COLOR_RGB2BGR))

# Attempt to move to GPU (as cuDF)
try:
    X_train = cudf.DataFrame.from_records(X_train)
    y_train = cudf.Series(y_train)
except:
    pass

svm = SVC(kernel='rbf')
svm.fit(X_train, y_train)

for page, image in enumerate(images):
    segmented_image = svm_segmentation(svm, None, image)
    extract_and_save_letters(segmented_image, image, page)

Task ID	Description	Points
00	Load Dataset	0
01	PCA Implementation	3
02	Kmeans Clustering Implementation	4
03	Clustering Method Evaluation
03-1	- IoU Metric Function Definition	1
03-2	- Clustering Method Evaluations	2
03-3	- Clustering Short Answer Questions	2
04	Supervised Image Segmentation
04-1	- Feature Vectors Creation for SVM	2
04-2	- Support Vector Machine Function Definition	1
04-3	- Properly Set Colab Env	0
04-4	- SVM Evaluation	3
04-5	- SVM Short Answer Questions	2
05	Final Letter Data Extraction
05-1	- Computer Vision Region Cropping Function	0
05-2	- Letter Extraction	0

Introduction to Artificial Intelligence - Homework Assignment 03 (20pts.)¶

Story Progression¶

Task 00: Load Data¶

Task 00: Description (0 pts.)¶

Load image data using opencv¶

Task 00: Code (0 pts.)¶

Story Progression¶

Task 01: PCA Implementation (3 pts.)¶

Task 01: Description (0 pts.)¶

PCA from scratch¶

Task 01: Code (3 pts.)¶

Task 02: kmeans Implementation (4 pts.)¶

Task 02: Description (0 pts.)¶

kmeans from scratch¶

Task 02: Code (4 pts.)¶

Task 3: Evaluating Unsupervised Segmentation (4 pts.)¶

Task 03-1: Description (0 pts.)¶

Intersection over Union¶

Task 03-1: Code (1 pt.)¶

Task 03-2: Description (0 pts.)¶

Evaluation of Unsupervised Methods¶

Task 03-2: Helper Code (0 pts.)¶

Task 03-2: Code (1 pt.)¶

Task 03-2: Expected Output (1 pts.)¶

Task 03-3: Short Answer Questions (2 pts.)¶

Story Progression¶

Task 04: Supervised Image Segmentation¶

Task 04-1: Description (0 pts.)¶

Feature Vector Extraction¶

Task 04-1: Code (2 pts.)¶

Task 04-2: Description (0 pts.)¶

SVM Segmentation Helper Function¶

Task 04-2: Code (1 pt.)¶

Task 04-3: Description (0 pts.)¶

Properly set colab env¶

Task 04-4: Description (0 pts.)¶

Supervised Segmentation Evaluation¶

Task 04-4: Code (2 pts.)¶

Task 04-3: Expected Output (1 pt.)¶

Task 04-4: Short Answer Questions (2 pts.)¶

Task 05: Region Extraction¶

Task 05-1: Description (0 pts.)¶

Region localization from SVM image segmentation¶

Task 05-1: Code (0 pts.)¶

Task 05-2: Description (0 pts.)¶

Extract each character from our letters¶

Task 05-2: Code (0 pts.)¶

Task 05-2: Example Output (0 pts.)¶

Story Progression¶