from PIL import Image
import math

v = 139
r = 5

# Define the smoothstep function for vignette effect
def smoothstep(edge0, edge1, x):
    """Smoothly interpolate between edge0 and edge1 based on x."""
    if edge0 == edge1:
        return 0.0 if x < edge0 else 1.0
    t = min(max((x - edge0) / (edge1 - edge0), 0.0), 1.0)
    return t * t * (3 - 2 * t)

# Define the hexagon function to compute coordinates
def hexagon(p):
    """
    Compute hexagon coordinates and distances for a given point p.
    Returns (h_x, h_y, e, f) where h_x and h_y are integer coordinates.
    """
    # Transform to hexagonal coordinate system
    q = (p[0] * 2.0 * 0.5773503, p[1] + p[0] * 0.5773503)
    pi = (math.floor(q[0]), math.floor(q[1]))
    pf = (q[0] - pi[0], q[1] - pi[1])
    v = (pi[0] + pi[1]) % 3.0
    ca = 1.0 if v >= 1.0 else 0.0
    cb = 1.0 if v >= 2.0 else 0.0
    ma = (1.0 if pf[1] >= pf[0] else 0.0, 1.0 if pf[0] >= pf[1] else 0.0)
    temp = (
        1.0 - pf[1] + ca * (pf[0] + pf[1] - 1.0) + cb * (pf[1] - 2.0 * pf[0]),
        1.0 - pf[0] + ca * (pf[0] + pf[1] - 1.0) + cb * (pf[0] - 2.0 * pf[1])
    )
    e = ma[0] * temp[0] + ma[1] * temp[1]
    p2_x = (q[0] + math.floor(0.5 + p[1] / 1.5)) * 0.5 + 0.5
    p2_y = (4.0 * p[1] / 3.0) * 0.5 + 0.5
    fract_p2 = (p2_x - math.floor(p2_x), p2_y - math.floor(p2_y))
    f = math.sqrt((fract_p2[0] - 0.5)**2 + ((fract_p2[1] - 0.5) * 0.85)**2)
    h_xy = (pi[0] + ca - cb * ma[0], pi[1] + ca - cb * ma[1])
    return (h_xy[0], h_xy[1], e, f)

def ura(p, r=1.0, v=10.0):
    """
    Generate a grayscale value based on hexagonal coordinates, emulating the HLSL URA function.
    
    Parameters:
        p (tuple): A 2-tuple (x, y) of coordinates.
        r (float): Multiplier for p[0]. Default is 1.0.
        v (float): Modulus value controlling the pattern frequency. Default is 10.0.
    
    Returns:
        float: 1.0 if no pattern match is found, otherwise 0.0.
    """
    import math
    l = math.fmod(p[1] + r * p[0], v)
    rz = 1.0
    for i in range(1, int(v/2)):
        if math.isclose(math.fmod(i * i, v), l, abs_tol=1e-6):
            rz = 0.0
            break
    return rz

# Define the color palette
default_colors = [
    (255, 0, 0),   # Red
    (0, 255, 0),   # Green
    (0, 0, 255)    # Blue
]

# Generate the image with colorful_hexagonal pattern
def generate_image_color(width, height, colors=default_colors):
    """Generate an RGB image with a colorful hexagonal pattern."""
    img = Image.new('RGB', (width, height))
    aspect = width / height
    for j in range(height):
        for i in range(width):
            # Normalize pixel coordinates to [0, 1]
            q_x = i / width
            q_y = j / height
            # Transform to centered coordinates with aspect ratio
            p_x = (q_x * 2.0 - 1.0) * aspect
            p_y = q_y * 2.0 - 1.0
            p = (p_x, p_y)
            # Scale coordinates for pattern frequency
            h = hexagon((p[0] * 21.0, p[1] * 21.0))
            h_xy = (int(h[0]), int(h[1]))
            # Assign color based on hexagon coordinates
            color_index = (h_xy[0] + h_xy[1]) % len(colors)
            col = colors[color_index]
            # Apply vignette effect
            q = (q_x * 2.0 - 1.0, q_y * 2.0 - 1.0)
            vignette = smoothstep(1.01, 0.97, max(abs(q[0]), abs(q[1])))
            col = tuple(int(c * vignette) for c in col)
            # Set the pixel color
            img.putpixel((i, j), col)
    return img

def generate_image_grayscale(width, height):
    img = Image.new('RGB', (width, height))
    aspect = width / height
    for j in range(height):
        for i in range(width):
            q_x = i / width
            q_y = j / height
            p_x = (q_x * 2.0 - 1.0) * aspect
            p_y = q_y * 2.0 - 1.0
            p = (p_x, p_y)
            h = hexagon((p[0] * 21.0, p[1] * 21.0))
            rz = ura(h[:2])
            smooth = smoothstep(-0.2, 0.13, h[2])
            if rz > 0.5:
                col = smooth
            else:
                col = 1.0 - smooth
            q = (q_x * 2.0 - 1.0, q_y * 2.0 - 1.0)
            vignette = smoothstep(1.01, 0.97, max(abs(q[0]), abs(q[1])))
            col *= vignette
            color = int(col * 255)
            img.putpixel((i, j), (color, color, color))
    return img

# Example usage
#if __name__ == "__main__":
#    # Set image dimensions
#    width, height = 800, 600
#    img = generate_image(width, height)
#    img.save('colorful_hexagon_pattern.png')
#    print("Image saved as 'colorful_hexagon_pattern.png'")