import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
import psutil


NUM_ITERATIONS = 1000


def plot_memory_usage(memory_usage):
    plt.figure(figsize=(10, 5))
    x = np.arange(len(memory_usage))
    plt.plot(x, memory_usage, label="Memory Usage (KB)", color="blue")
    plt.title("Memory Usage Over Time")
    plt.xlabel("Iterations")
    plt.ylabel("Memory Usage (KB)")
    plt.legend()
    plt.grid()
    plt.savefig("test_opencv_memory_leak.png")


def main():
    process = psutil.Process(os.getpid())

    shape = (800, 600, 3)
    image = np.random.randint(0, 255 + 1, np.prod(shape), np.uint8).reshape(shape)
    memory_usage = np.zeros(NUM_ITERATIONS, dtype=np.float32)

    cv2.namedWindow("image", cv2.WINDOW_NORMAL)
    for i in range(NUM_ITERATIONS):
        rect = cv2.getWindowImageRect("image")

        cv2.imshow("image", image)

        memory_usage[i] = process.memory_info().rss / 1024

        if cv2.waitKey(1) == ord("q"):
            break

    print("Current memory:", memory_usage[-1], "KB")
    print("Peak memory usage:", np.max(memory_usage), "KB")

    plot_memory_usage(memory_usage)


if __name__ == "__main__":
    main()