import cv2 
import numpy as np
import matplotlib.pyplot as plt
import pprint
import math


# /opt/homebrew/bin/ffmpeg -framerate 30 -pattern_type glob -i '*_EurchinsHg.jpg' -c:v libx264 -crf 20 EurchinsHg.mp4
# /opt/homebrew/bin/ffmpeg -i EurchinsHg.mp4 -vf "vidstabtransform=input=transforms.trf" -crf 20 EurchinsStabilized.mp4


def replace_color_with_image(background, overlay, target_color, tolerance, output_path):
    """
    Replaces a specific color in an image with another image.

    Args:
        background_path (str): Path to the background image.
        overlay_path (str): Path to the overlay image (can be the same as background if no overlay needed).
        target_color (tuple): RGB color to replace (e.g., (255, 0, 0) for red).
        tolerance (int): Tolerance range for color matching.
        replacement_image_path (str): Path to the replacement image.
        output_path (str): Path to save the output image.
    """
    if background is None or overlay is None:
        raise FileNotFoundError("One or more images not found.")

    # Create a mask for the target color
    lower_bound = np.array([max(0, c - tolerance) for c in target_color], dtype="uint8")
    upper_bound = np.array([min(255, c + tolerance) for c in target_color], dtype="uint8")
    mask = cv2.inRange(overlay, lower_bound, upper_bound)

    kernel = np.ones((3, 3), np.uint8) 
    dialated_mask = cv2.dilate(mask, kernel, iterations=1) 

    mask_rgb = cv2.merge((dialated_mask,dialated_mask,dialated_mask))

    mask_inv = cv2.bitwise_not(dialated_mask)
    mask_inv_rgb = cv2.merge((mask_inv,mask_inv,mask_inv))

    masked_background = cv2.bitwise_and(mask_rgb, background) # pixels from background where target_color is in the overlay.
    masked_overlay = cv2.bitwise_and(mask_inv_rgb, overlay) # pixels from background where target_color is in the overlay.
    combined = cv2.add(masked_background, masked_overlay)


    cv2.imwrite(output_path, combined)



f = open("transforms.trf", "w")
f.write("VID.STAB 1")
f.write("#      accuracy = 15")
f.write("#     shakiness = 3")
f.write("#      stepsize = 4")
f.write("#   mincontrast = 0.200000")
f.write("Frame 1 (List 0 [])
")

img1 = cv2.imread('0012_EurchinsHg.jpg')  
mask = cv2.imread('mask.png', cv2.IMREAD_GRAYSCALE)

currentMask = mask

#sift
sift = cv2.SIFT_create()

keypoints_1, descriptors_1 = sift.detectAndCompute(img1,mask)

for n in range(1,67): # Remember to +1 the end number.
    print("Frame " + str(n))
    img2 = cv2.imread("{:04d}_EurchinsHg.jpg".format(n)) 
    #img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

    keypoints_2, descriptors_2 = sift.detectAndCompute(img2,currentMask)

    #feature matching
    bf = cv2.BFMatcher(cv2.NORM_L1, crossCheck=True)

    matches = bf.match(descriptors_1,descriptors_2)
    matches = sorted(matches, key = lambda x:x.distance)

    # Get key points
    points1 = np.float32([keypoints_1[m.queryIdx].pt for m in matches]).reshape(-1, 1, 2)
    points2 = np.float32([keypoints_2[m.trainIdx].pt for m in matches]).reshape(-1, 1, 2)

    # Affine matrix
    matrix, _ = cv2.findHomography(points2, points1, cv2.RANSAC, ransacReprojThreshold=5.0)


    # Align
    new_background_color = (0, 255, 0)  # Green

    aligned_image = cv2.warpPerspective(img2, matrix, (img1.shape[1], img1.shape[0]), borderValue=new_background_color)

    replace_color_with_image(img1, aligned_image, new_background_color, 1, "aligned_{:04d}_EurchinsHg.jpg".format(n))


f.close()

/opt/homebrew/bin/ffmpeg -framerate 30 -pattern_type glob -i 'aligned_*_EurchinsHg.jpg' -c:v libx264 -crf 20 AlignedEurchinsHg.mp4