from PIL import Image
import imageio


# Function to map (x, y) to (u, v)
def map_function(x, y, scale_s, scale_t, mapparam):
    # You can define your mapping function here.
    # For example, this function scales and shifts the coordinates:
    u = mapparam * (scale_s*scale_t) * x / (x*x+y*y)
    v = mapparam * (scale_s*scale_t) * y / (x*x+y*y)
    return u, v

def create_frames(stepsize,N_pics,w,h,wt,ht,mapparam):
    default_color = (0, 0, 0)
    output_images = []
    scale_s = min(w, h)
    scale_t = min(wt, ht)

    for iter in range(N_pics):
        print("%d\r"%iter,end="")
        offset = -iter*stepsize
        def wrap(int_u):
            return (int_u+offset) % (wt - 1)

        output_image = Image.new("RGB", (w, h))
        for x in range(w):
            for y in range(h):
                realx, realy = x-w//2, y-h//2
                if realx==0 and realy==0:
                    output_image.putpixel((x,y), default_color)
                    continue
                realu, realv = map_function(realx, realy, scale_s, scale_t, mapparam)
                u, v = realu + wt//2, realv + ht//2
                
                # Interpolation (you can use different methods, such as bilinear)
                u_floor, v_floor = int(u), int(v)
                u_frac, v_frac = u - u_floor, v - v_floor

                # Ensure u and v are within the bounds of the texture image
                if u_floor < 0 or u_floor > wt - 2 or v_floor < 0 or v_floor > ht - 2:
                    output_image.putpixel((x,y), default_color)
                    continue

                color_top_left = texture_image.getpixel((wrap(u_floor), v_floor))
                color_top_right = texture_image.getpixel((wrap(u_floor + 1), v_floor))
                color_bottom_left = texture_image.getpixel((wrap(u_floor), v_floor + 1))
                color_bottom_right = texture_image.getpixel((wrap(u_floor + 1), v_floor + 1))

                # Interpolate colors
                color = (
                    int((1 - u_frac) * (1 - v_frac) * color_top_left[0] +
                        u_frac * (1 - v_frac) * color_top_right[0] +
                        (1 - u_frac) * v_frac * color_bottom_left[0] +
                        u_frac * v_frac * color_bottom_right[0]),
                    int((1 - u_frac) * (1 - v_frac) * color_top_left[1] +
                        u_frac * (1 - v_frac) * color_top_right[1] +
                        (1 - u_frac) * v_frac * color_bottom_left[1] +
                        u_frac * v_frac * color_bottom_right[1]),
                    int((1 - u_frac) * (1 - v_frac) * color_top_left[2] +
                        u_frac * (1 - v_frac) * color_top_right[2] +
                        (1 - u_frac) * v_frac * color_bottom_left[2] +
                        u_frac * v_frac * color_bottom_right[2])
                )

                output_image.putpixel((x, y), color)
        output_images.append(output_image)

    return output_images

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Turn an image into a stylized gif where every frame of the gif is\
                    a conformal mapping of the original image shifted and wrapped around in the width direction')
    parser.add_argument('input_file')
    parser.add_argument('output_file')
    parser.add_argument('--output_width', type=int, default=400)
    parser.add_argument('--output_height', type=int, default=400)
    parser.add_argument('--num_frames', type=int, default=60)
    parser.add_argument('--x_shift_per_frame', type=int, nargs="?")
    parser.add_argument('--mapping_param', type=float, default=0.05)
    args = parser.parse_args()

    # Load your texture image
    texture_image = Image.open(args.input_file)

    # Set the width and height of the output image
    w, h = args.output_width, args.output_height
    # Set the width and height of the texture image
    wt, ht = texture_image.width, texture_image.height

    mapparam = args.mapping_param
    N_pics = args.num_frames
    if args.x_shift_per_frame==None:
        stepsize = wt / N_pics # one whole period
    else:
        stepsize = args.x_shift_per_frame

    output_images = create_frames(stepsize,N_pics,w,h,wt,ht,mapparam)
    imageio.mimsave(args.output_file, output_images)