import numpy as np
from skimage import io, color, transform
from skimage.feature import hog

image_path = "C:/Users/AlpOzzy/Desktop/Thesis/Segmentation_400webpages/output_folder/"
# train_image_path = "C:/Users/AlpOzzy/Desktop/Thesis/Segmentation_400webpages/train_images/"
image1_name = "f9cb7acd-8d39-4819-b9b9-ecef3c9ae2b0.png"
image2_name = "f9f607be-23e9-4cae-856c-16b341b9a9c2.png"


def extract_hog_features(image, orientations=9, pixels_per_cell=(8, 8),
                         cells_per_block=(2, 2)):
    """
    Extract HOG features from an image.
    """
    features = hog(image, orientations=orientations,
                   pixels_per_cell=pixels_per_cell,
                   cells_per_block=cells_per_block,
                   block_norm='L2-Hys',
                   feature_vector=True)
    return features


def get_pyramid_features(image, levels=3):
    """
    Extract features using spatial pyramid matching.
    """
    height, width = image.shape
    features = []

    for level in range(levels):
        num_cells = 2 ** level
        h_step = height // num_cells
        w_step = width // num_cells

        for i in range(num_cells):
            for j in range(num_cells):
                h_start = i * h_step
                h_end = (i + 1) * h_step
                w_start = j * w_step
                w_end = (j + 1) * w_step
                region = image[h_start:h_end, w_start:w_end]

                hog_features = extract_hog_features(region)
                # Normalize the features
                hog_features = hog_features / (np.sum(hog_features) + 1e-10)
                features.append(hog_features)

    # Weight features by level
    weighted_features = []
    start_idx = 0
    for level in range(levels):
        num_regions = 4 ** level
        level_features = features[start_idx:start_idx + num_regions]
        weight = 2 ** (levels - level - 1) if level != levels - 1 else 2 ** (levels - 1)
        weighted_features.extend([feat * weight for feat in level_features])
        start_idx += num_regions

    return np.concatenate(weighted_features)


def histogram_intersection(hist1, hist2):
    """
    Compute normalized histogram intersection similarity.
    Returns a value between 0 and 1.
    """
    # Ensure features are normalized
    hist1_norm = hist1 / (np.sum(hist1) + 1e-10)
    hist2_norm = hist2 / (np.sum(hist2) + 1e-10)

    intersection = np.sum(np.minimum(hist1_norm, hist2_norm))
    # Normalize by the total possible intersection
    max_intersection = np.minimum(np.sum(hist1_norm), np.sum(hist2_norm))

    return intersection / (max_intersection + 1e-10)


def compute_similarity(features1, features2, metric='intersection'):
    """
    Compute similarity between two feature vectors using different metrics.

    Args:
        features1: First feature vector
        features2: Second feature vector
        metric: Similarity metric ('intersection', 'l1', 'l2', or 'cosine')

    Returns:
        Similarity score
    """
    if metric == 'intersection':
        return histogram_intersection(features1, features2)

    elif metric == 'cosine':
        dot_product = np.dot(features1, features2)
        norm1 = np.linalg.norm(features1)
        norm2 = np.linalg.norm(features2)
        return dot_product / (norm1 * norm2 + 1e-10)

    elif metric == 'l1':
        distance = np.sum(np.abs(features1 - features2))
        return 1.0 / (1.0 + distance)  # Convert distance to similarity

    elif metric == 'l2':
        distance = np.sqrt(np.sum((features1 - features2) ** 2))
        return 1.0 / (1.0 + distance)  # Convert distance to similarity

    else:
        raise ValueError(f"Unknown metric: {metric}")


def match_images(image1, image2, levels=3, metric='intersection'):
    """
    Match two images using Spatial Pyramid Matching with HOG features.

    Args:
        image1: First image (grayscale)
        image2: Second image (grayscale)
        levels: Number of pyramid levels
        metric: Similarity metric ('intersection', 'l1', 'l2', or 'cosine')

    Returns:
        Similarity score between the images
    """
    features1 = get_pyramid_features(image1, levels)
    features2 = get_pyramid_features(image2, levels)

    return compute_similarity(features1, features2, metric)


# Example usage:
if __name__ == "__main__":
    # Load and preprocess images
    img1 = io.imread(image_path + image1_name)
    img2 = io.imread(image_path + image2_name)

    # Convert to grayscale if needed
    if len(img1.shape) == 3:
        img1 = color.rgb2gray(img1)
    if len(img2.shape) == 3:
        img2 = color.rgb2gray(img2)

    # Resize to same dimensions if needed
    img2 = transform.resize(img2, img1.shape)

    # Compare using different metrics
    metrics = ['intersection', 'cosine', 'l1', 'l2']
    for metric in metrics:
        similarity = match_images(img1, img2, metric=metric)
        print(f"{metric} similarity: {similarity:.4f}")