import os
import cv2
import numpy as np
import pandas as pd

# Read in the video (our video has a frequency of 30 frames/second)
video_file = "world.mp4"
cap = cv2.VideoCapture(video_file)

# Read in the CSV file (this file was generated by the idt.py from the gaze positions csv file "4bwCZ9awAx_unfamiliar.csv")
csv_file = "fixation_4bwCZ9awAx_unfamiliar.csv"
df = pd.read_csv(csv_file)

# Define a function to apply the homography matrix to the x and y coordinates
def apply_homography(x, y, H):
    # Convert x and y to homogeneous coordinates
    points = np.array([[x, y, 1]], dtype=np.float32).T

    # Apply the homography matrix
    new_points = np.dot(H, points)

    # Convert back to cartesian coordinates
    new_x = new_points[0, 0] / new_points[2, 0]
    new_y = new_points[1, 0] / new_points[2, 0]

    return new_x, new_y

# Loop through each pair of consecutive rows in the CSV file
for index, row in df.iterrows():
    if index == len(df) - 1:
        break  # Skip the last row

    # Get the start and end frame numbers from the CSV file
    start_frame = int(row["start_frame"])
    end_frame = int(df.loc[index+1, "end_frame"])

    # Stitch all frames between the two frame numbers
    frames = []
    for frame_num in range(start_frame, end_frame+1):
        # Set the video capture to the correct frame
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)

        # Capture the frame
        ret, frame = cap.read()

        # Check if the frame was successfully captured
        if not ret:
            continue

        frames.append(frame)

    # Apply image stitching algorithm to get the homography matrix (You can replace this with your own image stitching algorithm)
    stitcher = cv2.Stitcher.create()
    (status, H) = stitcher.stitch(frames)

    # Save the frames to the folder if the homography matrix is not None
    if H is not None:
        # Create a folder for the frames
        folder_name = "{}-{}".format(start_frame, end_frame)
        if not os.path.exists(folder_name):
            os.makedirs(folder_name)

        # Save the frames to the folder
        for i, frame in enumerate(frames):
            filename = os.path.join(folder_name, "{}.jpg".format(start_frame + i))
            cv2.imwrite(filename, frame)

        # Output the homography matrix
        print("Homography Matrix for frames {} to {}: {}".format(start_frame, end_frame, H))

        # Apply the homography matrix to the x and y coordinates
        # (you can multiply the "x_mean" and "y_mean" with the number of pixels on the x and y axis to have the absolute pixel index of every fixation)
        x1 = row["x_mean"]
        y1 = row["y_mean"]
        x2 = df.loc[index+1, "x_mean"]
        y2 = df.loc[index+1, "y_mean"]
        x1_corrected, y1_corrected = apply_homography(x1, y1, H)
        x2_corrected, y2_corrected = apply_homography(x2, y2, H)

        # Add the corrected x and y values to the CSV file
        df.loc[index, "corrected-x"] = x1_corrected
        df.loc[index, "corrected-y"] = y1_corrected
        df.loc[index+1, "corrected-x"] = x2_corrected
        df.loc[index+1, "corrected-y"] = y2_corrected
    else:
        df.loc[index, "corrected-x"] = x1
        df.loc[index, "corrected-y"] = y1
        df.loc[index + 1, "corrected-x"] = x2
        df.loc[index + 1, "corrected-y"] = y2

# Save the modified CSV file. This file contains the corrected fixation positions.
df.to_csv("corrected_"+csv_file, index=False)