from PIL import Image, ImageFilter
import random
import math

SQUARE_SIZE = 32
SQUARES = 2**8
ITERATIONS_PER_PIXEL = 1000
COORDS = [(1, 0), (1, 1), (0, 1), (-1, 1), (-1, 0), (-1, -1), (0, -1), (1, -1)]


def random_point_on_sphere():
    theta = random.uniform(0, 2 * math.pi)
    phi = math.acos(random.uniform(-1, 1))
    x = math.sin(phi) * math.cos(theta)
    y = math.sin(phi) * math.sin(theta)
    z = math.cos(phi)
    return [x, y, z]


calculated_squares = dict()
with Image.new("RGBA", (SQUARE_SIZE*SQUARES, SQUARE_SIZE), "black") as im:
    px = im.load()

    for square_id in range(SQUARES):
        print(f"square {square_id}: ", end="")
        square = square_id
        
        # if 2 side touching tiles are obscuring, then we can assume the corner is obscuring as well
        if square & (1 << 6) > 0 and square & (1 << 4) > 0:
            square = square | (1 << 5)
        if square & (1 << 4) > 0 and square & (1 << 2) > 0:
            square = square | (1 << 3)
        if square & (1 << 2) > 0 and square & (1 << 0) > 0:
            square = square | (1 << 1)
        if square & (1 << 0) > 0 and square & (1 << 6) > 0:
            square = square | (1 << 7)
        
        # check if we've already done the task in a rotated variation
        instance_found = False
        for rotation in range(4):  # clockwise rotations by pi/2
            rotated_square = (square >> (rotation * 2)) | ((square & ((1 << rotation*2)-1)) << 2*(4 - rotation))
            if rotated_square in calculated_squares:
                square_id_from = calculated_squares[rotated_square]
                print(f"already calculated at {square_id_from}")

                box_from = (SQUARE_SIZE*square_id_from, 0, SQUARE_SIZE*(square_id_from+1), SQUARE_SIZE)
                box_to = (SQUARE_SIZE*square_id, 0, SQUARE_SIZE*(square_id+1), SQUARE_SIZE)
                region = im.crop(box_from)
                if rotation == 1:
                    region = region.transpose(Image.Transpose.ROTATE_270)
                if rotation == 2:
                    region = region.transpose(Image.Transpose.ROTATE_180)
                if rotation == 3:
                    region = region.transpose(Image.Transpose.ROTATE_90)
                im.paste(region, box_to)

                instance_found = True
                break
        if instance_found:
            continue

        # actually calculate the occlusion
        obscured = [i for c, i in enumerate(COORDS) if (square >> c) & 1 == 1]
        for x in range(SQUARE_SIZE):
            for y in range(SQUARE_SIZE):
                point_coord = ((x+0.5)/SQUARE_SIZE-0.5, (y+0.5)/SQUARE_SIZE-0.5, 0)
                count = 0
                for i in range(ITERATIONS_PER_PIXEL):
                    point = random_point_on_sphere()
                    point[0] += point_coord[0]
                    point[1] += point_coord[1]
                    rounded_point = (round(point[0]), round(point[1]))
                    if rounded_point in obscured:
                        count += 1
                result_color = 255-255*count//ITERATIONS_PER_PIXEL
                px[x+square_id*SQUARE_SIZE, y] = (result_color, result_color, result_color, 255)
        calculated_squares[square] = square_id
        
        # smooth it
        box = (SQUARE_SIZE*square_id, 0, SQUARE_SIZE*(square_id+1), SQUARE_SIZE)
        region = im.crop(box)
        region = region.filter(filter=ImageFilter.GaussianBlur(2))
        im.paste(region, box)

        print(f"calculated")

    print(f"total squares calculated: {len(calculated_squares)}")
    
    # result_side_length = int(SQUARES**0.5)
    # with Image.new("RGBA", (SQUARE_SIZE*result_side_length, SQUARE_SIZE*result_side_length), "black") as new_im:
    #     for i in range(result_side_length):
    #         box_from = (SQUARE_SIZE*result_side_length*i, 0, SQUARE_SIZE*result_side_length*(i+1), SQUARE_SIZE)
    #         box_to = (0, SQUARE_SIZE*i, SQUARE_SIZE*result_side_length, SQUARE_SIZE*(i+1))
    #         new_im.paste(im.crop(box_from), box_to)
    #     new_im.show()
    #     new_im.save("../images/ambient_occlusion.png")

    im.show()
    im.save("../images/ambient_occlusion.png")