import comfy.utils
import math

MAX_RESOLUTION=8192


def ratiocalc(width, height, multiplier):
    aspectratio = width / height

    #Comfy needs images that are divisible by 64 below only works for a1111
    #newwidth = width * multiplier
    #newheight = math.floor(newwidth / aspectratio)
    #newwidth = math.floor(newwidth)

    newwidth = int((width/64)*multiplier)*64
    newheight = int(((width/64)*multiplier)/aspectratio)*64

    print(f"Multiplier:{multiplier} Ratio{aspectratio}. Res: {newwidth}x{newheight}")
    return newwidth, newheight

class LatentUpscaleByRatio:
    upscale_methods = ["nearest-exact", "bilinear", "area"]
    crop_methods = ["disabled", "center"]

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
                              "original_width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "original_height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "upscale_ratio": ("FLOAT", {"default": 2.00, "min": 1.00, "max": 6.00, "step": 0.01}),
                              "crop": (s.crop_methods,)}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "upscale"
    CATEGORY = "latent"

    def upscale(self, samples, upscale_method, original_width, original_height, crop, upscale_ratio):

        original_width, original_height = ratiocalc(original_width, original_height, upscale_ratio)

        s = samples.copy()
        s["samples"] = comfy.utils.common_upscale(samples["samples"], original_width // 8, original_height // 8, upscale_method, crop)
        return (s,)


class ImageScaleByRatio:
    upscale_methods = ["nearest-exact", "bilinear", "area"]
    crop_methods = ["disabled", "center"]

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
                              "original_width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "original_height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
                              "upscale_ratio": ("FLOAT", {"default": 2.00, "min": 1.00, "max": 6.00, "step": 0.01}),
                              "crop": (s.crop_methods,)}}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"

    CATEGORY = "image/upscaling"
    def upscale(self, image, upscale_method, original_width, original_height, crop, upscale_ratio):

        original_width, original_height = ratiocalc(original_width, original_height, upscale_ratio)

        samples = image.movedim(-1,1)
        s = comfy.utils.common_upscale(samples, original_width, original_height, upscale_method, crop)
        s = s.movedim(1,-1)
        return (s,)


# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
    "LatentUpscaleByRatio": LatentUpscaleByRatio,
    "ImageScaleByRatio": ImageScaleByRatio
}

# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
    "LatentUpscaleByRatio": "Latent Upscale By Ratio",
    "ImageScaleByRatio": "Image Scale By Ratio"
}