from PIL import Image, ImageDraw, ImageFont
import numpy as np
import cv2
import os

FONTS_FOLDER = "fonts"
NUM_IMAGES_PER_FONT = 5

def generate_images(text):
    images = []
    for font_file in os.listdir(FONTS_FOLDER):
        font_path = os.path.join(FONTS_FOLDER, font_file)
        for i in range(NUM_IMAGES_PER_FONT):
            img = generate_text_image(text, font_path)
            images.append((img, font_file))
    return images

def generate_text_image(text, font_path, fontsize=None):
    if not fontsize:
        fontsize = int(np.random.normal(loc=50, scale=10))
    font = ImageFont.truetype(font_path, fontsize)
    text_size = font.getsize(text)
    img = Image.new('RGB', text_size, color='black')
    draw = ImageDraw.Draw(img)
    draw.text((0, 0), text, font=font, fill='white')
    noise = np.random.normal(loc=0, scale=10, size=(img.size[1], img.size[0]))[..., np.newaxis]
    noise = np.tile(noise, [1, 1, 3])
    img = Image.fromarray(np.clip(np.array(img) + noise, 0, 255).astype(np.uint8), 'RGB')
    return np.array(img)

def flann_matching_alt(generated_images, query_image, num_trees=5, num_checks=50):
    query_image_gray = cv2.cvtColor(query_image, cv2.COLOR_BGR2GRAY)
    generated_images_gray = []
    for img, _ in generated_images:
        generated_images_gray.append(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))

    sift = cv2.SIFT_create()
    index_params = dict(algorithm=0, trees=num_trees)
    search_params = dict(checks=num_checks)
    flann = cv2.FlannBasedMatcher(index_params, search_params)

    query_kp, query_desc = sift.detectAndCompute(query_image_gray, None)

    if query_desc is None:
        return []

    matching_results = []
    for i, (img, font_file) in enumerate(generated_images):
        kp, desc = sift.detectAndCompute(generated_images_gray[i], None)
        if desc is not None:
            matches = flann.knnMatch(query_desc, desc, k=2)
            good_matches = []
            for m, n in matches:
                if m.distance < 0.75 * n.distance:
                    good_matches.append(m)
            matching_img = cv2.drawMatches(query_image_gray, query_kp, generated_images_gray[i], kp, good_matches, None, flags=2)
            # Calculate percentage match
            num_query_kp = len(query_kp)
            num_matches = len(good_matches)
            match_percent = 100 * num_matches / num_query_kp
            matching_results.append((matching_img, font_file, match_percent))

    return matching_results

def flann_matching(generated_images, query_image, num_trees=5, num_checks=50):
    query_image_gray = cv2.cvtColor(query_image, cv2.COLOR_BGR2GRAY)
    generated_images_gray = []
    for img, _ in generated_images:
        generated_images_gray.append(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))
    
    sift = cv2.SIFT_create()
    index_params = dict(algorithm=0, trees=num_trees)
    search_params = dict(checks=num_checks)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    
    query_kp, query_desc = sift.detectAndCompute(query_image_gray, None)
    generated_kp = []
    generated_desc = []
    for img in generated_images_gray:
        kp, desc = sift.detectAndCompute(img, None)
        generated_kp.append(kp)
        generated_desc.append(desc)
    
    matching_results = []
    for i, (img, font_file) in enumerate(generated_images):
        matches = flann.knnMatch(query_desc, generated_desc[i], k=2)
        good_matches = []
        for m, n in matches:
            if m.distance < 0.75*n.distance:
                good_matches.append([m])
        matching_img = cv2.drawMatchesKnn(query_image_gray, query_kp, img, generated_kp[i], good_matches, None, flags=2)
        # Calculate percentage match
        num_query_kp = len(query_kp)
        num_matches = len(good_matches)
        match_percent = 100*num_matches/num_query_kp
        matching_results.append((matching_img, font_file, match_percent))
        
    return matching_results

def flann_matching_3(generated_images, query_image, num_trees=5, num_checks=50):
    query_image_gray = cv2.cvtColor(query_image, cv2.COLOR_BGR2GRAY)
    generated_images_gray = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img, _ in generated_images]
    
    sift = cv2.SIFT_create()
    index_params = dict(algorithm=0, trees=num_trees)
    search_params = dict(checks=num_checks)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    
    query_kp, query_desc = sift.detectAndCompute(query_image_gray, None)
    
    if query_desc is None:
        return []
    
    matching_results = []
    for i, (img, font_file) in enumerate(generated_images):
        kp, desc = sift.detectAndCompute(generated_images_gray[i], None)
        
        if desc is None:
            continue
        
        matches = flann.knnMatch(query_desc, desc, k=2)
        good_matches = []
        for m, n in matches:
            if m.distance < 0.75 * n.distance:
                good_matches.append(m)
        
        if len(good_matches) < 10:
            continue
        
        src_pts = np.float32([query_kp[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        M, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
        
        if M is None:
            continue
        
        h, w = query_image_gray.shape
        dst_img = cv2.warpPerspective(img, M, (w, h))
        dst_gray = cv2.cvtColor(dst_img, cv2.COLOR_BGR2GRAY)
        
        orb = cv2.ORB_create()
        kp1, desc1 = sift.detectAndCompute(query_image_gray, None)
        kp2, desc2 = sift.detectAndCompute(dst_gray, None)
        
        if desc1 is None or desc2 is None:
            continue
        
        matches = flann.knnMatch(desc1, desc2, k=2)
        good_matches = []
        for m, n in matches:
            if m.distance < 0.75 * n.distance:
                good_matches.append(m)
        
        if len(good_matches) < 10:
            continue
        
        matching_img = cv2.drawMatches(query_image_gray, kp1, dst_gray, kp2, good_matches, None, flags=2)
        # Calculate percentage match
        num_query_kp = len(kp1)
        num_matches = len(good_matches)
        match_percent = 100 * num_matches / num_query_kp
        matching_results.append((matching_img, font_file, match_percent))
        
    return matching_results

def flann_matching_4(generated_images, query_image, num_trees=5, num_checks=50):
    query_image_gray = cv2.cvtColor(query_image, cv2.COLOR_BGR2GRAY)
    generated_images_gray = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img, _ in generated_images]

    sift = cv2.SIFT_create()
    index_params = dict(algorithm=0, trees=num_trees)
    search_params = dict(checks=num_checks)
    flann = cv2.FlannBasedMatcher(index_params, search_params)

    query_kp, query_desc = sift.detectAndCompute(query_image_gray, None)

    if query_desc is None:
        return []

    matching_results = []
    for img, font_file in generated_images:
        generated_image_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        kp, desc = sift.detectAndCompute(generated_image_gray, None)

        if desc is None:
            continue

        matches = flann.knnMatch(query_desc, desc, k=2)
        good_matches = []
        for m, n in matches:
            if m.distance < 0.75 * n.distance:
                good_matches.append(m)
        matching_img = cv2.drawMatches(query_image_gray, query_kp, generated_image_gray, kp, good_matches, None, flags=2)
        
        # Calculate percentage match
        num_query_kp = len(query_kp)
        num_matches = len(good_matches)
        match_percent = 100 * num_matches / num_query_kp
        matching_results.append((matching_img, font_file, match_percent))

    return matching_results