from PIL import Image
import imageio
import math


# Function to map (x, y) to (u, v)
def map_function(x, y, map_type):
    if map_type == "holo":
        u = x / (x*x+y*y)
        v = -y / (x*x+y*y)
    elif map_type == "antiholo":
        u = x / (x*x+y*y)
        v = y / (x*x+y*y)
    elif map_type == "exp":
        u, v = math.exp(x)*math.cos(y), math.exp(x)*math.sin(y)
    elif map_type == "log":
        u, v = math.log(x*x+y*y), math.atan(y/(x+1e-9))
    return u, v

def create_frames(stepsize,N_pics,w,h,wt,ht,s_in,s_out,map_type):
    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*s_in/scale_s, realy*s_in/scale_s, map_type)
                realu, realv = realu*s_out*scale_t, realv*s_out*scale_t
                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('--scale_in', type=float, default=1, help="multiply to the input to the mapping,\
                        default range [-0.5, 0.5] x [-0.5, 0.5]")
    parser.add_argument('--scale_out', type=float, default=0.05, help="multiply to the output of the mapping,\
                        those out of range [-0.5, 0.5] x [-0.5, 0.5] is discarded")
    parser.add_argument('--map_type',type=str,choices=["holo","antiholo","exp","log"],default="antiholo",\
                        help="holo: holomorphic mapping for f(z) = 1/z; antiholo: mapping for f(z) = 1/conj(z);\
                            exp: holomorphic mapping for f(z) = exp(z); log: f(z) = log(z)")
    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

    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,args.scale_in,args.scale_out,args.map_type)
    imageio.mimsave(args.output_file, output_images)