# -- Markdown Cell --
# **Classical feature extractors**

# -- Code Cell --
import cv2
import numpy as np
import matplotlib.pyplot as plt

def extract_image_features(image_path):
    # Load image
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError("Image not found or path is incorrect.")

    # Resize image (standard size for consistency)
    image = cv2.resize(image, (256, 256))

    # Convert to RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 1. Color Histogram Features
    hist_r = cv2.calcHist([image_rgb], [0], None, [256], [0, 256])
    hist_g = cv2.calcHist([image_rgb], [1], None, [256], [0, 256])
    hist_b = cv2.calcHist([image_rgb], [2], None, [256], [0, 256])

    # Normalize histograms
    hist_r = hist_r.flatten() / hist_r.sum()
    hist_g = hist_g.flatten() / hist_g.sum()
    hist_b = hist_b.flatten() / hist_b.sum()

    color_features = np.concatenate([hist_r, hist_g, hist_b])

    # 2. Edge Features using Canny
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 100, 200)

    edge_density = np.sum(edges > 0) / (256 * 256)

    # 3. Statistical Features
    mean_color = np.mean(image_rgb, axis=(0, 1))
    std_color = np.std(image_rgb, axis=(0, 1))

    feature_vector = {
        "color_histogram": color_features,
        "edge_density": edge_density,
        "mean_color": mean_color,
        "std_color": std_color
    }

    return feature_vector


if __name__ == "__main__":
    image_path = "833.jpg"
    features = extract_image_features(image_path)

    print("Extracted Image Features:")
    for key, value in features.items():
        print(f"{key}: {value}")

    image_path = "834.jpg"
    features = extract_image_features(image_path)

    print("Extracted Image Features:")
    for key, value in features.items():
        print(f"{key}: {value}")

    image_path = "936.jpg"
    features = extract_image_features(image_path)

    print("Extracted Image Features:")
    for key, value in features.items():
        print(f"{key}: {value}")

    image_path = "937.jpg"
    features = extract_image_features(image_path)

    print("Extracted Image Features:")
    for key, value in features.items():
        print(f"{key}: {value}")


# -- Code Cell --
def extract_feature_vector(image_path):
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError("Image not found or path is incorrect.")

    image = cv2.resize(image, (256, 256))
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Color histograms
    hist_r = cv2.calcHist([image_rgb], [0], None, [256], [0, 256]).flatten()
    hist_g = cv2.calcHist([image_rgb], [1], None, [256], [0, 256]).flatten()
    hist_b = cv2.calcHist([image_rgb], [2], None, [256], [0, 256]).flatten()

    hist_r /= hist_r.sum()
    hist_g /= hist_g.sum()
    hist_b /= hist_b.sum()

    color_features = np.concatenate([hist_r, hist_g, hist_b])

    # Edge density
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 100, 200)
    edge_density = np.array([np.sum(edges > 0) / (256 * 256)])

    # Mean & std color
    mean_color = np.mean(image_rgb, axis=(0, 1))
    std_color = np.std(image_rgb, axis=(0, 1))

    # Final feature vector
    feature_vector = np.concatenate([
        color_features,
        edge_density,
        mean_color,
        std_color
    ])

    return feature_vector

def euclidean_distance(f1, f2):
    return np.linalg.norm(f1 - f2)

def cosine_similarity(f1, f2):
    return np.dot(f1, f2) / (np.linalg.norm(f1) * np.linalg.norm(f2))

image_paths = ["833.jpg", "834.jpg", "936.jpg", "937.jpg"]

features = {p: extract_feature_vector(p) for p in image_paths}

for i in range(len(image_paths)):
    for j in range(i + 1, len(image_paths)):
        f1 = features[image_paths[i]]
        f2 = features[image_paths[j]]

        dist = euclidean_distance(f1, f2)
        cos = cosine_similarity(f1, f2)

        print(f"{image_paths[i]} vs {image_paths[j]}")
        print(f"  Euclidean distance: {dist:.4f}")
        print(f"  Cosine similarity: {cos:.4f}")

# -- Code Cell --
query = "833.jpg"
query_feat = extract_feature_vector(query)

best_match = None
best_distance = float("inf")

for path in image_paths:
    if path == query:
        continue

    dist = euclidean_distance(query_feat, features[path])

    if dist < best_distance:
        best_distance = dist
        best_match = path

print(f"Best match for {query}: {best_match} (distance={best_distance:.4f})")

# -- Markdown Cell --
# **Deep learning feature extractors**
# 
# Closer to 1 - more similar (better match)
# 
# Around 0 - unrelated
# 
# Negative - very different

# -- Code Cell --
from torchvision.models import resnet50, ResNet50_Weights
import torch
from PIL import Image
import numpy as np

# Load weights
weights = ResNet50_Weights.DEFAULT

# Load model
model = resnet50(weights=weights)

# Remove classifier - feature extractor
model = torch.nn.Sequential(*list(model.children())[:-1])
model.eval()

# Processing
transform = weights.transforms()

def extract_resnet_features(image_path):
    image = Image.open(image_path).convert("RGB")
    tensor = transform(image).unsqueeze(0)

    with torch.no_grad():
        features = model(tensor)

    features = features.squeeze().numpy()

    # Normalize feature vector
    features = features / np.linalg.norm(features)
    return features


def cosine_similarity(a, b):
    return np.dot(a, b)


image_paths = ["833.jpg", "834.jpg", "936.jpg", "937.jpg"]
features = {p: extract_resnet_features(p) for p in image_paths}

for i in range(len(image_paths)):
    for j in range(i + 1, len(image_paths)):
        sim = cosine_similarity(features[image_paths[i]], features[image_paths[j]])
        print(f"{image_paths[i]} vs {image_paths[j]} → similarity = {sim:.4f}")


# -- Code Cell --
