import numpy as np
import cv2
import argparse
import time
from numpy import random

def load_image(image_path):
    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
    if image is None:
        raise ValueError(f"Image at path {image_path} could not be loaded.")
    return image.astype(np.float32)

def draw_circle(image, reference_image, center, radius):
    mask = np.zeros(image.shape[:2], dtype=np.uint8)
    cv2.circle(mask, center, int(radius), 255, -1)

    mean_color = cv2.mean(reference_image, mask=mask)[:3]
    mean_color = tuple(map(int, mean_color))

    tmp_image = image.copy()
    cv2.circle(tmp_image, center, int(radius), mean_color, -1)
    return tmp_image

def calculate_vector_distance(image1, image2):
    return np.linalg.norm(image1 - image2)

def main(in_path, out_path, iterations, seed):
    random.seed(seed)
    image1 = load_image(in_path)
    height, width, _ = image1.shape
    image2 = np.zeros((height,width,3), np.uint8)

    last_time = time.time()

    for i in range(iterations):
        print(f"Image {in_path} pass {i}", end="\r")

        center = (random.randint(width), random.randint(height))
        radius = max(width, height) / 2
        step = radius

        for _ in range(int(np.log2(max(width, height))) + 1):
            candidate_up = draw_circle(image2, image1, center, radius + step)
            candidate_up_distance = calculate_vector_distance(image1, candidate_up)

            candidate_down = draw_circle(image2, image1, center, max(0, radius - step))
            candidate_down_distance = calculate_vector_distance(image1, candidate_down)

            if candidate_down_distance <= candidate_up_distance:
                radius = max(0, radius - step)
            else:
                radius += step

            step /= 2

        image2 = draw_circle(image2, image1, center, radius)

    print(f"{out_path} took {-last_time + time.time()} seconds")
    cv2.imwrite(f"{out_path}.png", np.uint8(image2))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Draw and optimize a circle on an image using gradient ascent.")
    parser.add_argument("input", help="Path to the reference image.")
    parser.add_argument("output", help="Path to save the output image.")
    parser.add_argument("circles", help="The amount of circles to draw.")
    parser.add_argument("seed", help="The seed to use.")

    args = parser.parse_args()
    main(args.input, args.output, int(args.circles), int(args.seed))