import subprocess
import numpy as np
import matplotlib.pyplot as plt
import sys

# --- Setup ---
video_a = "/home/honney/projects/mkv/Saga of Tanya the Evil/Season 1/Saga Of Tanya The Evil - S01E01 - The Devil Of The Rhine.mkv"
video_b = "/home/honney/projects/mkv/Saga of Tanya the Evil (2017-2017) [imdbid-tt6455986]/Season 01/Saga of Tanya the Evil S01E001.mkv"

WIDTH = 160
HEIGHT = int((WIDTH/16)*9)
FPS = 24000/1001
MAX_OFFSET_SECONDS = 2  # ±2 Sekunden Suche


def extract_hist_signal(video):
    """Extrahiert pro Frame einen Histogrammvektor unter Nutzung von VA-API (AMD GPU)."""
    # Wir nutzen vaapi für Hardware-Dekodierung und Scaling
    proc = subprocess.Popen([
        "ffmpeg",
        "-loglevel", "error",
        "-hwaccel", "vaapi",             # Nutze VA-API Hardware-Beschleunigung
        "-hwaccel_device", "/dev/dri/renderD128", # Pfad zur RX 7900 XTX
        "-hwaccel_output_format", "vaapi",
        "-i", video,
        "-vf", (
            f"fps={FPS},"                # Framerate Anpassung
            f"scale_vaapi=w={WIDTH}:h={HEIGHT}:format=gray8," # GPU-Scaling + Grayscale
            "hwdownload,format=gray8"    # Zurück in den RAM schieben
        ),
        "-f", "rawvideo",
        "-"
    ], stdout=subprocess.PIPE)

    frame_size = WIDTH * HEIGHT
    hist_signal = []

    # Der restliche Loop bleibt gleich, da die Histogramm-Berechnung 
    # (np.histogram) weiterhin auf der CPU via NumPy erfolgt.
    while True:
        data = proc.stdout.read(frame_size)
        if not data or len(data) < frame_size:
            break
        frame = np.frombuffer(data, dtype=np.uint8)
        hist, _ = np.histogram(frame, bins=32, range=(0, 255))
        hist_signal.append(hist / np.sum(hist))

    proc.wait()
    return np.array(hist_signal)


print("Extracting signal of video A")
a = extract_hist_signal(video_a)

print("Extracting signal of video B")
b = extract_hist_signal(video_b)

print("Calculating cross correlation")

max_lag_frames = int(MAX_OFFSET_SECONDS * FPS)
n_a, n_bins = a.shape
n_b, _ = b.shape

lags = np.arange(-max_lag_frames, max_lag_frames + 1)
corr = []

for lag in lags:
    if lag < 0:
        # Offset negativ → Video B vor Video A
        sub_a = a[:lag]           # bis lag Frames von Video A
        sub_b = b[-lag:]          # von -lag bis Ende von Video B
    else:
        sub_a = a[lag:]           # ab lag bis Ende von Video A
        sub_b = b[:n_b - lag]     # bis n_b - lag von Video B

    # Frames angleichen
    min_len = min(len(sub_a), len(sub_b))
    if min_len == 0:
        corr.append(0)
        continue
    sub_a = sub_a[:min_len]
    sub_b = sub_b[:min_len]

    corr.append(np.sum(sub_a * sub_b))

corr = np.array(corr)
best_index = np.argmax(corr)
best_lag = lags[best_index]

print("Best offset frames:", best_lag)
print("Best offset seconds:", best_lag / FPS)

sys.exit()

# Plot
plt.figure(figsize=(10, 5))
plt.plot(lags / FPS, corr, marker='o')
plt.axvline(best_lag / FPS, color='r', linestyle='--', label=f"Best offset = {best_lag / FPS:.3f}s")
plt.xlabel("Offset (seconds)")
plt.ylabel("Cross correlation")
plt.title("Video cross correlation (Histogram-based)")
plt.legend()
plt.grid(True)
plt.show()