import torch
import torchvision
import matplotlib.pyplot as plt
import numpy as np
from torchvision import transforms
from PIL import Image
from sklearn.decomposition import PCA

# Load the Caltech101 dataset
caltech101_data = torchvision.datasets.Caltech101('/content/', download=True)

def extract_random_patch(image, patch_size=8):
    """Extract a random patch from an image."""
    # Convert PIL Image to tensor and handle grayscale conversion
    if isinstance(image, Image.Image):
        # Ensure the image is large enough
        if image.size[0] < patch_size or image.size[1] < patch_size:
            image = image.resize((patch_size*2, patch_size*2))
        
        # Convert to tensor
        to_tensor = transforms.ToTensor()
        image = to_tensor(image)

    # Convert to grayscale if it's RGB
    if image.shape[0] == 3:
        image = 0.299 * image[0] + 0.587 * image[1] + 0.114 * image[2]
    elif image.shape[0] == 1:
        image = image.squeeze(0)

    # Ensure we have valid dimensions
    assert image.dim() == 2, f"Expected 2D tensor, got shape {image.shape}"
    h, w = image.shape
    assert h >= patch_size and w >= patch_size, f"Image too small: {h}x{w}, need at least {patch_size}x{patch_size}"

    # Get valid patch coordinates
    i = np.random.randint(0, h - patch_size + 1)
    j = np.random.randint(0, w - patch_size + 1)
    
    # Extract and flatten patch
    patch = image[i:i+patch_size, j:j+patch_size].reshape(-1)
    
    # Normalize patch
    patch_mean = patch.mean()
    patch_std = patch.std()
    if patch_std == 0:
        patch_std = 1e-8
    patch = (patch - patch_mean) / patch_std
    
    return patch.numpy()

# Collect patches
n_patches = 10000
patch_size = 8
patches = []

print("Collecting patches...")
for _ in range(n_patches):
    idx = np.random.randint(len(caltech101_data))
    image, _ = caltech101_data[idx]
    patch = extract_random_patch(image)
    patches.append(patch)

# Convert to numpy array
patches = np.array(patches)
print(f"Collected {patches.shape[0]} patches of size {patches.shape[1]}")

# Perform PCA
n_components = patch_size * patch_size  # Same as GHA output size
print("Performing PCA...")
pca = PCA(n_components=n_components)
pca.fit(patches)

# Plot the principal components
def plot_principal_components(components):
    """Plot the principal components in an 8x8 grid."""
    fig, axes = plt.subplots(8, 8, figsize=(10, 10))
    for i in range(8):
        for j in range(8):
            idx = i * 8 + j
            pc = components[idx].reshape(8, 8)
            axes[i, j].imshow(pc, cmap='gray')
            axes[i, j].axis('off')
    plt.tight_layout()
    return fig

# Plot results
print("Plotting results...")
fig = plot_principal_components(pca.components_)
fig.savefig("PCA_Caltech101.png")
plt.show()

# Print explained variance ratios for the first few components
print("\nExplained variance ratios for first 10 components:")
for i, ratio in enumerate(pca.explained_variance_ratio_[:10]):
    print(f"PC {i+1}: {ratio:.3f}")

# Save the PCA components for comparison
pca_components = torch.from_numpy(pca.components_)