import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path


def get_quadrant(x, y):
    """Return the quadrant of a given point."""
    if x >= 0 and y >= 0:
        return 1
    elif x < 0 and y >= 0:
        return 2
    elif x < 0 and y < 0:
        return 3
    else:
        return 4


def is_overlapping(x, y, existing_points, radius):
    """Check if a point overlaps with existing points, crosses a quadrant or goes beyond the quadrant lines."""
    current_quadrant = get_quadrant(x, y)
    buffer = 0.25
    # Check if circle touches/crosses the x=0 or y=0 lines.
    if (
        abs(x) + radius > buffer
        and abs(x) - radius < buffer
        or abs(y) + radius > buffer
        and abs(y) - radius < buffer
    ):
        return True

    for ex, ey in existing_points:
        # Check overlap with existing points
        if np.sqrt((x - ex) ** 2 + (y - ey) ** 2) <= 2 * radius + 0.5:
            return True
        # Check if the circle touches another quadrant
        if (
            get_quadrant(ex, ey) != current_quadrant
            and np.sqrt((x - ex) ** 2 + (y - ey) ** 2) <= radius
        ):
            return True
    return False


DIR = "/teamspace/studios/this_studio/out_sortcolors"

N = 100
DPI = 300
SIZE = 72  # canvas size?
radius = 3
variance = 72  # Adjust variance as needed

Path(DIR).mkdir(exist_ok=True, parents=True)


with open(f"{DIR}/arrangement_grid.txt", "w") as f:
    radius = 3
    half = SIZE / 2 - 4
    # make points an equispaced grid of 10 x 10 ranging from [-35, 35]
    interval = (half - (-half)) / 9  # 10 points, so 9 intervals

    # Generate the grid points
    points = [
        (x, y)
        for x in np.arange(-half, half + 0.1, interval)
        for y in np.arange(-half, half + 0.1, interval)
    ]

    for x, y in points:
        f.write(f"{x-radius}, {y+radius}\n")
        # f.write(f"{x}, {y}\n")

print("wrote grid")


for seed in range(11, 22):
    try:
        np.random.seed(seed)

        # To store plotted points
        points = []
        max_iterations = int(1e7)
        iterations = 0
        # Generate points
        for k in range(N):
            while True:
                # x, y = np.random.normal(0, variance), np.random.normal(0, variance)
                random_angle = np.random.uniform(0, 2 * np.pi)
                # random_radius = np.random.uniform(0.25+radius, SIZE/2 - radius - 0.25)
                random_radius = abs(np.random.normal(0, variance))
                x = random_radius * np.cos(random_angle)
                y = random_radius * np.sin(random_angle)
                iterations += 1
                if not is_overlapping(x, y, points, radius):
                    if max(abs(x), abs(y)) + radius < SIZE / 2 - 0.25:
                        points.append((x, y))
                        break
                if iterations > max_iterations:
                    raise ValueError(f"Too many iterations: {k} points")
            print(f"{k}: ({x}, {y}) @ {iterations:09d}")

        # Create plot with circles
        fig, ax = plt.subplots(1, 1, figsize=(SIZE, SIZE))
        for x, y in points:
            circle = patches.Circle((x, y), radius, color="black")
            ax.add_patch(circle)

        # Draw the standard quadrants
        ax.axhline(0, color="black", linewidth=0.5)
        ax.axvline(0, color="black", linewidth=0.5)

        # Use square axis and set limits

        lim_x = (
            max(
                abs(max(points, key=lambda t: t[0])[0]),
                abs(min(points, key=lambda t: t[0])[0]),
            )
            + radius
        )
        lim_y = (
            max(
                abs(max(points, key=lambda t: t[1])[1]),
                abs(min(points, key=lambda t: t[1])[1]),
            )
            + radius
        )
        lim_x = lim_y = SIZE / 2
        ax.axis("off")
        ax.set_aspect("equal")
        ax.set_xlim(-lim_x, lim_x)
        ax.set_ylim(-lim_y, lim_y)

        # Save and show
        fig.tight_layout(pad=0)
        plt.savefig(
            f"{DIR}/arrangement_{seed}.png", dpi=DPI, bbox_inches="tight", pad_inches=0
        )
        # plt.show()
        # also save x/y coords as text file
        with open(f"{DIR}/arrangement_{seed}.txt", "w") as f:
            for x, y in points:
                f.write(f"{x-radius}, {y+radius}\n")
                # f.write(f"{x}, {y}\n")
    except ValueError as e:
        print(f"{seed}: {e}")
    except AssertionError as e:
        print(f"{seed}: {e}")