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

def mse(img1, img2):
   h, w, _= img1.shape
   diff = cv2.subtract(img1, img2)
   err = np.sum(diff**2)
   mse = err/(float(h*w))
   return mse

def seconds_to_hms(seconds):
    hours = seconds // 3600
    minutes = (seconds % 3600) // 60
    secs = seconds % 60
    return f'{int(hours)}:{int(minutes):02}:{int(secs):02}'

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)) 
    h, w, _ = img.shape

    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: {seconds_to_hms(total_frames / video_fps)}")

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

    prev = time.perf_counter()
    best_match = float('inf')
    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, (w, h)) 
            
            error = mse(img, frame)
            
            if error < best_match:
                best_match = error
                best_match_frame = frameCouter
                cv2.imwrite(f"./output/matchFrame_{frameCouter}.bmp", frame)
        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 {seconds_to_hms(best_match_frame / video_fps)}")
            pbar.update(1)
        
    pbar.close()
    cap.release()
    
    end_time = time.time()
    print(f"{seconds_to_hms(end_time - start_time)} seconds")