import subprocess
import numpy as np
import cv2
import torch
import torch.nn.functional as F
from scipy import signal

TARGET_FPS = 30

def get_device():

    if torch.cuda.is_available():
        device = torch.device("cuda")

        print("GPU detected")
        print("Device:", torch.cuda.get_device_name(0))
        print("ROCm build:", torch.version.hip is not None)

        return device

    print("GPU not available, falling back to CPU")
    return torch.device("cpu")


device = get_device()


def ffmpeg_frame_stream(path):

    cmd = [
        "ffmpeg",
        "-i", path,
        "-vf", "fps=30,scale=64:36,format=bgr24",
        "-f", "rawvideo",
        "-pix_fmt", "bgr24",
        "-"
    ]

    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        bufsize=10**8
    )

    frame_size = 64 * 36 * 3

    while True:
        raw = proc.stdout.read(frame_size)
        if len(raw) < frame_size:
            break

        frame = np.frombuffer(raw, np.uint8).reshape((36, 64, 3))
        yield frame

    proc.stdout.close()
    proc.wait()


def compute_signal(video_path):

    prev_blocks = None
    signal_out = []

    for frame in ffmpeg_frame_stream(video_path):

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        t = torch.from_numpy(gray).float().to(device) / 255.0

        # reshape for conv
        t = t.unsqueeze(0).unsqueeze(0)

        # block averages using avg pool
        blocks = F.avg_pool2d(t, kernel_size=8, stride=8)

        blocks = blocks.squeeze()

        if prev_blocks is None:
            prev_blocks = blocks
            signal_out.append(0)
            continue

        # block difference energy
        diff = torch.abs(blocks - prev_blocks).sum()

        # edge energy
        sobel_x = torch.tensor(
            [[-1,0,1],[-2,0,2],[-1,0,1]],
            dtype=torch.float32, device=device
        ).view(1,1,3,3)

        sobel_y = torch.tensor(
            [[-1,-2,-1],[0,0,0],[1,2,1]],
            dtype=torch.float32, device=device
        ).view(1,1,3,3)

        gx = F.conv2d(t, sobel_x, padding=1)
        gy = F.conv2d(t, sobel_y, padding=1)

        edge_energy = torch.mean(torch.sqrt(gx**2 + gy**2))

        value = diff + 0.5 * edge_energy

        signal_out.append(value.item())

        prev_blocks = blocks

    signal_out = np.array(signal_out)

    # temporal derivative
    signal_out = np.diff(signal_out, prepend=signal_out[0])

    return signal_out


def normalize(sig):
    return (sig - np.mean(sig)) / (np.std(sig) + 1e-6)


def find_offset(video_a, video_b):

    sig_a = compute_signal(video_a)
    sig_b = compute_signal(video_b)

    sig_a = normalize(sig_a)
    sig_b = normalize(sig_b)

    corr = signal.correlate(sig_a, sig_b, mode="full")
    lags = signal.correlation_lags(len(sig_a), len(sig_b), mode="full")

    best = lags[np.argmax(corr)]

    seconds = best / TARGET_FPS

    return best, seconds


def main():

    video_a = "/home/honney/mount/Media/tmp/serien/ger_low_quality/Black Bullet/Black Bullet Season 01 E03 Kinder des Schicksals.mp4"
    video_b = "/home/honney/mount/Media/tmp/serien/high_quality/Black Bullet (Season 1)/Season 01/Black Bullet - S01E03.mkv"

    frames, seconds = find_offset(video_a, video_b)

    print("offset frames:", frames)
    print("offset seconds:", seconds)


if __name__ == "__main__":
    main()