import cv2
import numpy as np
from pathlib import Path
import time
from tqdm import tqdm
import argparse

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    
    parser.add_argument("-i", "--Input", type=str)
    parser.add_argument("-o", "--Output", type=str, default="output.mp4")
    parser.add_argument("-t", "--Target", type=str)
    
    args = parser.parse_args()

    start_time = time.time()
    img = cv2.imread(str(Path(args.Target)), cv2.IMREAD_COLOR)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.resize(img, (224, 224))
    
    ssim_algo = cv2.quality.QualitySSIM_create(img)

    cap = cv2.VideoCapture(str(Path(args.Input)))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    video_fps = cap.get(cv2.CAP_PROP_FPS)
    print(f"Video FPS: {video_fps:.2f}, Total frames: {total_frames}, Video length: {(total_frames / video_fps):.2f} seg")

    pbar = tqdm(
        total=total_frames,
        bar_format="{desc}",      # only render the description
        desc="aFPS: 0.00"          # initial text
    )

    prev = time.perf_counter()
    best_match = -1
    frameCouter = 0
    best_match_frame = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        frameCouter += 1
        try:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            frame = cv2.resize(frame, (224, 224))
            
            sim = ssim_algo.compute(frame)[0]
            
            if sim > best_match:
                best_match = sim
                best_match_frame = frameCouter
        except Exception as e:
            continue
        finally:
            now = time.perf_counter()
            fps = 1.0 / (now - prev)
            prev = now

            pbar.set_description(f"aFPS: {fps:.2f}, best Match. {best_match_frame} / {best_match:.2f}; Aprox timestamp {(best_match_frame / video_fps):.4f}")
            pbar.update(1)
        
    pbar.close()
    cap.release()
    
    end_time = time.time()
    print(f"{end_time - start_time:.4f} seconds")