import argparse
import numpy as np
from PIL import Image

'''
    ULTRA ZAJEBISTY LABELTAG IMG GENERATOR
    by masuo
    v1.0
'''

OUTER_RADIUS = 9.4
INNER_RADIUS = 4.2
ANGLE_PER_HEIGHT = 1.26

def load_image(path):
    return Image.open(path)

def save_image(img, path):
    img.save(path)

def get_disc_coords(label_coords, half_disc_size, label_height):
    x, y = label_coords
    angle_rad = x * ANGLE_PER_HEIGHT / label_height
    radius_mm = OUTER_RADIUS - (y / label_height) * (OUTER_RADIUS - INNER_RADIUS)
    radius_px = (radius_mm / OUTER_RADIUS) * half_disc_size
    disc_x = half_disc_size + radius_px * np.sin(angle_rad)
    disc_y = half_disc_size - radius_px * np.cos(angle_rad) - 1

    return int(np.round(disc_x)), int(np.round(disc_y))

def get_sizes(img):
    width, height = img.size
    if width != height:
        raise ValueError("Disc image must be square for proper mapping.")
    
    disc_size = width
    if disc_size % 2 != 0:
        raise ValueError("Disc image dimensions must be even for proper mapping.")
    
    half_disc_size = disc_size // 2
    
    label_height_float = disc_size * (OUTER_RADIUS - INNER_RADIUS) / OUTER_RADIUS
    label_width_float = label_height_float * np.pi * 2 / (ANGLE_PER_HEIGHT)
    
    label_height = int(np.round(label_height_float))
    label_width = int(np.round(label_width_float))

    return disc_size, half_disc_size, label_height, label_width    

def apply_mapping(disc_img):
    disc_size, half_disc_size, label_height, label_width = get_sizes(disc_img)

    disc_img = disc_img.convert('RGBA')
    data = np.array(disc_img)
    label_data = np.zeros((label_height, label_width, 4), dtype=np.uint8)

    for y in range(label_height):
        for x in range(label_width):
            disc_x, disc_y = get_disc_coords((x, y), half_disc_size, label_height)
            if 0 <= disc_x < disc_size and 0 <= disc_y < disc_size:
                label_data[y, x] = data[disc_y, disc_x]

    return Image.fromarray(label_data, 'RGBA')

def main():
    p = argparse.ArgumentParser()
    p.add_argument('input', help='input PNG file')
    p.add_argument('output', help='output PNG file')
    args = p.parse_args()
    img = load_image(args.input)
    out = apply_mapping(img)
    save_image(out, args.output)

if __name__ == '__main__':
    main()