dockipelago/worlds/dk64/randomizer/Patching/Library/Image.py

"""Image modification library functions."""

import zlib
import random
import gzip
import math
from enum import IntEnum, auto
from PIL import Image, ImageEnhance
from randomizer.Patching.Patcher import ROM, LocalROM
from randomizer.Settings import ColorblindMode
from randomizer.Enums.Kongs import Kongs
from randomizer.Patching.Library.Assets import getPointerLocation, TableNames
from typing import Tuple, Union


class TextureFormat(IntEnum):
    """Texture Format Enum."""

    Null = auto()
    RGBA5551 = auto()
    RGBA32 = auto()
    I8 = auto()
    I4 = auto()
    IA8 = auto()
    IA4 = auto()


class ExtraTextures(IntEnum):
    """Extra Textures in Table 25 after the bonus skins."""

    FakeGBShine = 0
    RainbowCoin0 = auto()
    RainbowCoin1 = auto()
    RainbowCoin2 = auto()
    MelonSurface = auto()
    BonusShell = auto()
    OSprintLogoLeft = auto()
    OSprintLogoRight = auto()
    BLockerItemMove = auto()
    BLockerItemBlueprint = auto()
    BLockerItemFairy = auto()
    BLockerItemBean = auto()
    BLockerItemPearl = auto()
    BLockerItemRainbowCoin = auto()
    BLockerItemIceTrap = auto()
    BLockerItemPercentage = auto()
    BLockerItemBalloon = auto()
    BLockerItemCompanyCoin = auto()
    BLockerItemKong = auto()
    BeetleTex0 = auto()
    BeetleTex1 = auto()
    BeetleTex2 = auto()
    BeetleTex3 = auto()
    BeetleTex4 = auto()
    BeetleTex5 = auto()
    BeetleTex6 = auto()
    Feather0 = auto()
    Feather1 = auto()
    Feather2 = auto()
    Feather3 = auto()
    Feather4 = auto()
    Feather5 = auto()
    Feather6 = auto()
    Feather7 = auto()
    FoolOverlay = auto()
    MedalRim = auto()
    MushTop0 = auto()
    MushTop1 = auto()
    ShellWood = auto()
    ShellMetal = auto()
    ShellQMark = auto()
    RocketTop = auto()
    BlastTop = auto()
    Anniv25Sticker = auto()
    Anniv25Barrel = auto()
    BeanSpin01 = auto()
    BeanSpin02 = auto()
    BeanSpin03 = auto()
    BeanSpin04 = auto()
    BeanSpin05 = auto()
    BeanSpin06 = auto()
    BeanSpin07 = auto()
    BeanSpin08 = auto()
    BeanSpin09 = auto()
    BeanSpin10 = auto()
    BeanSpin11 = auto()
    BeanSpin12 = auto()
    KrushaFace1 = auto()
    KrushaFace2 = auto()
    KrushaFace3 = auto()
    KrushaFace4 = auto()
    KrushaFace5 = auto()
    KrushaFace321 = auto()
    KrushaFace322 = auto()
    KrushaFace323 = auto()
    KrushaFace324 = auto()
    KrushaFace325 = auto()
    APPearl0 = auto()
    APPearl1 = auto()
    APPearl2 = auto()
    APPearl3 = auto()
    APPearl4 = auto()
    APPearl5 = auto()


barrel_skins = (
    "gb",
    "dk",
    "diddy",
    "lanky",
    "tiny",
    "chunky",
    "bp",
    "nin_coin",
    "rw_coin",
    "key",
    "crown",
    "medal",
    "potion",
    "bean",
    "pearl",
    "fairy",
    "rainbow",
    "fakegb",
    "melon",
    "cranky",
    "funky",
    "candy",
    "snide",
    "hint",
    "ap",
)


def getBonusSkinOffset(offset: int):
    """Get texture index after the barrel skins."""
    return 6026 + (3 * len(barrel_skins)) + offset


def getImageFromAddress(ROM_COPY: Union[LocalROM, ROM], rom_address: int, width: int, height: int, compressed: bool, file_size: int, format: TextureFormat):
    """Get image from a ROM address."""
    ROM_COPY.seek(rom_address)
    data = ROM_COPY.readBytes(file_size)
    if compressed:
        data = zlib.decompress(data, (15 + 32))
    im_f = Image.new(mode="RGBA", size=(width, height))
    pix = im_f.load()
    for y in range(height):
        for x in range(width):
            if format == TextureFormat.RGBA32:
                offset = ((y * width) + x) * 4
                pix_data = int.from_bytes(data[offset : offset + 4], "big")
                red = (pix_data >> 24) & 0xFF
                green = (pix_data >> 16) & 0xFF
                blue = (pix_data >> 8) & 0xFF
                alpha = pix_data & 0xFF
            elif format == TextureFormat.RGBA5551:
                offset = ((y * width) + x) * 2
                pix_data = int.from_bytes(data[offset : offset + 2], "big")
                red = ((pix_data >> 11) & 31) << 3
                green = ((pix_data >> 6) & 31) << 3
                blue = ((pix_data >> 1) & 31) << 3
                alpha = (pix_data & 1) * 255
            elif format == TextureFormat.IA8:
                offset = (y * width) + x
                pix_data = int.from_bytes(data[offset : offset + 1], "big")
                intensity = int(((pix_data >> 4) / 0xF) * 255)
                alpha = int(((pix_data & 0xF) / 0xF) * 255)
                red = intensity
                green = intensity
                blue = intensity
            else:
                raise Exception(f"Unhandled Codec: {format}")
            pix[x, y] = (red, green, blue, alpha)
    return im_f


def getImageFile(ROM_COPY: Union[LocalROM, ROM], table_index: TableNames, file_index: int, compressed: bool, width: int, height: int, format: TextureFormat):
    """Grab image from file."""
    file_start = getPointerLocation(table_index, file_index)
    file_end = getPointerLocation(table_index, file_index + 1)
    file_size = file_end - file_start
    return getImageFromAddress(ROM_COPY, file_start, width, height, compressed, file_size, format)


def getRandomHueShift(min: int = -359, max: int = 359) -> int:
    """Get random hue shift."""
    return random.randint(min, max)


def hueShift(im, amount: int):
    """Apply a hue shift on an image."""
    hsv_im = im.convert("HSV")
    im_px = im.load()
    w, h = hsv_im.size
    hsv_px = hsv_im.load()
    amount = int(amount * (256 / 360))  # Truncate to within 256
    for y in range(h):
        for x in range(w):
            old = list(hsv_px[x, y]).copy()
            old[0] = (old[0] + amount) % 256
            hsv_px[x, y] = (old[0], old[1], old[2])
    rgb_im = hsv_im.convert("RGB")
    rgb_px = rgb_im.load()
    for y in range(h):
        for x in range(w):
            new = list(rgb_px[x, y])
            new.append(list(im_px[x, y])[3])
            im_px[x, y] = (new[0], new[1], new[2], new[3])
    return im


def hueShiftImageFromAddress(ROM_COPY: Union[ROM, LocalROM], address: int, width: int, height: int, format: TextureFormat, shift: int):
    """Hue shift image located at a certain ROM address."""
    size_per_px = {
        TextureFormat.RGBA5551: 2,
        TextureFormat.RGBA32: 4,
    }
    data_size_per_px = size_per_px.get(format, None)
    if data_size_per_px is None:
        raise Exception(f"Texture Format unsupported by this function. Let the devs know if you see this. Attempted format: {format.name}")
    loaded_im = getImageFromAddress(ROM_COPY, address, width, height, False, data_size_per_px * width * height, format)
    loaded_im = hueShift(loaded_im, shift)
    loaded_px = loaded_im.load()
    bytes_array = []
    for y in range(height):
        for x in range(width):
            pix_data = list(loaded_px[x, y])
            if format == TextureFormat.RGBA32:
                bytes_array.extend(pix_data)
            elif format == TextureFormat.RGBA5551:
                red = int((pix_data[0] >> 3) << 11)
                green = int((pix_data[1] >> 3) << 6)
                blue = int((pix_data[2] >> 3) << 1)
                alpha = int(pix_data[3] != 0)
                value = red | green | blue | alpha
                bytes_array.extend([(value >> 8) & 0xFF, value & 0xFF])
    px_data = bytearray(bytes_array)
    ROM_COPY.seek(address)
    ROM_COPY.writeBytes(px_data)


def clampRGBA(n):
    """Restricts input to integer value between 0 and 255."""
    return math.floor(max(0, min(n, 255)))


def convertRGBAToBytearray(rgba_lst):
    """Convert RGBA list with 4 items (r,g,b,a) to a two-byte array in RGBA5551 format."""
    twobyte = (rgba_lst[0] << 11) | (rgba_lst[1] << 6) | (rgba_lst[2] << 1) | (rgba_lst[3] & 1)
    lower = twobyte % 256
    upper = int(twobyte / 256) % 256
    return [upper, lower]


def imageToCI(ROM_COPY: ROM, im_f, ci_index: int, tex_index: int, pal_index: int):
    """Change image to a CI texture."""
    if ci_index not in (4, 8):
        return
    color_count = 1 << ci_index
    if color_count < 32:
        im_f = im_f.quantize(colors=color_count, method=Image.MAXCOVERAGE)
    else:
        im_f = im_f.convert("P", palette=Image.ADAPTIVE, colors=color_count)
    palette_indexes = list(im_f.getdata())
    palette = im_f.getpalette()
    palette_colors = [tuple(palette[i : i + 3]) for i in range(0, len(palette), 3)]
    rgba5551_values = []
    for color in palette_colors:
        colv = 0
        for channel_value in color:
            val = channel_value & 0x1F
            colv <<= 5
            colv |= val
        colv |= 1
        rgba5551_values.append(colv)
    tex_bin = []
    if ci_index == 8:
        tex_bin = palette_indexes.copy()
    else:
        output_value = 0
        for index, value in enumerate(palette_indexes):
            if (index & 1) == 0:
                output_value = (value & 0xF) << 4
            else:
                output_value |= value & 0xF
                tex_bin.append(output_value)
    pal_bin = []
    for half in rgba5551_values:
        upper = (half >> 8) & 0xFF
        lower = half & 0xFF
        pal_bin.extend([upper, lower])
    tex_bin_file = gzip.compress(bytearray(tex_bin), compresslevel=9)
    pal_bin_file = gzip.compress(bytearray(pal_bin), compresslevel=9)
    tex_start = getPointerLocation(TableNames.TexturesGeometry, tex_index)
    tex_end = getPointerLocation(TableNames.TexturesGeometry, tex_index + 1)
    pal_start = getPointerLocation(TableNames.TexturesGeometry, pal_index)
    pal_end = getPointerLocation(TableNames.TexturesGeometry, pal_index + 1)
    if (tex_end - tex_start) < len(tex_bin_file):
        return
    if (pal_end - pal_start) < len(pal_bin_file):
        return
    ROM_COPY.seek(tex_start)
    ROM_COPY.write(tex_bin_file)
    ROM_COPY.seek(pal_start)
    ROM_COPY.write(pal_bin_file)


def writeColorImageToROM(im_f, table_index: TableNames, file_index: int, width: int, height: int, transparent_border: bool, format: TextureFormat, ROM_COPY: Union[LocalROM, ROM]) -> None:
    """Write texture to ROM."""
    file_start = getPointerLocation(table_index, file_index)
    file_end = getPointerLocation(table_index, file_index + 1)
    file_size = file_end - file_start
    ROM_COPY.seek(file_start)
    pix = im_f.load()
    width, height = im_f.size
    bytes_array = []
    border = 1
    right_border = 3
    for y in range(height):
        for x in range(width):
            if transparent_border:
                if ((x < border) or (y < border) or (x >= (width - border)) or (y >= (height - border))) or (x == (width - right_border)):
                    pix_data = [0, 0, 0, 0]
                else:
                    pix_data = list(pix[x, y])
            else:
                pix_data = list(pix[x, y])
            if format == TextureFormat.RGBA32:
                bytes_array.extend(pix_data)
            elif format == TextureFormat.RGBA5551:
                red = int((pix_data[0] >> 3) << 11)
                green = int((pix_data[1] >> 3) << 6)
                blue = int((pix_data[2] >> 3) << 1)
                alpha = int(pix_data[3] != 0)
                value = red | green | blue | alpha
                bytes_array.extend([(value >> 8) & 0xFF, value & 0xFF])
            elif format == TextureFormat.IA4:
                intensity = pix_data[0] >> 5
                alpha = 0 if pix_data[3] == 0 else 1
                data = ((intensity << 1) | alpha) & 0xF
                bytes_array.append(data)
    bytes_per_px = 2
    if format == TextureFormat.IA4:
        temp_ba = bytes_array.copy()
        bytes_array = []
        value_storage = 0
        bytes_per_px = 0.5
        for idx, val in enumerate(temp_ba):
            polarity = idx % 2
            if polarity == 0:
                value_storage = val << 4
            else:
                value_storage |= val
                bytes_array.append(value_storage)
    data = bytearray(bytes_array)
    if format == TextureFormat.RGBA32:
        bytes_per_px = 4
    if len(data) > (bytes_per_px * width * height):
        print(f"Image too big error: {table_index} > {file_index}")
    if table_index in (14, 25):
        data = gzip.compress(data, compresslevel=9)
    if len(data) > file_size:
        print(f"File too big error: {table_index} > {file_index}")
    ROM_COPY.writeBytes(data)


def getNumberImage(number: int, ROM_COPY: Union[LocalROM, ROM]):
    """Get Number Image from number."""
    if number < 5:
        num_0_bounds = [0, 20, 30, 45, 58, 76]
        x = number
        return getImageFile(ROM_COPY, 14, 15, True, 76, 24, TextureFormat.RGBA5551).crop((num_0_bounds[x], 0, num_0_bounds[x + 1], 24))
    num_1_bounds = [0, 15, 28, 43, 58, 76]
    x = number - 5
    return getImageFile(ROM_COPY, 14, 16, True, 76, 24, TextureFormat.RGBA5551).crop((num_1_bounds[x], 0, num_1_bounds[x + 1], 24))


def numberToImage(number: int, dim: Tuple[int, int], ROM_COPY: Union[LocalROM, ROM]):
    """Convert multi-digit number to image."""
    digits = 1
    if number < 10:
        digits = 1
    elif number < 100:
        digits = 2
    else:
        digits = 3
    current = number
    nums = []
    total_width = 0
    max_height = 0
    sep_dist = 1
    for _ in range(digits):
        base = getNumberImage(current % 10, ROM_COPY)
        bbox = base.getbbox()
        base = base.crop(bbox)
        nums.append(base)
        base_w, base_h = base.size
        max_height = max(max_height, base_h)
        total_width += base_w
        current = int(current / 10)
    nums.reverse()
    total_width += (digits - 1) * sep_dist
    base = Image.new(mode="RGBA", size=(total_width, max_height))
    pos = 0
    for num in nums:
        base.paste(num, (pos, 0), num)
        num_w, num_h = num.size
        pos += num_w + sep_dist
    output = Image.new(mode="RGBA", size=dim)
    xScale = dim[0] / total_width
    yScale = dim[1] / max_height
    scale = xScale
    if yScale < xScale:
        scale = yScale
    new_w = int(total_width * scale)
    new_h = int(max_height * scale)
    x_offset = int((dim[0] - new_w) / 2)
    y_offset = int((dim[1] - new_h) / 2)
    new_dim = (new_w, new_h)
    base = base.resize(new_dim)
    output.paste(base, (x_offset, y_offset), base)
    return output


def getRGBFromHash(hash: str):
    """Convert hash RGB code to rgb array."""
    red = int(hash[1:3], 16)
    green = int(hash[3:5], 16)
    blue = int(hash[5:7], 16)
    return [red, green, blue]


def maskImageWithColor(im_f: Image, mask: tuple):
    """Apply rgb mask to image using a rgb color tuple."""
    w, h = im_f.size
    converter = ImageEnhance.Color(im_f)
    im_f = converter.enhance(0)
    im_dupe = im_f.copy()
    brightener = ImageEnhance.Brightness(im_dupe)
    im_dupe = brightener.enhance(2)
    im_f.paste(im_dupe, (0, 0), im_dupe)
    pix = im_f.load()
    w, h = im_f.size
    for x in range(w):
        for y in range(h):
            base = list(pix[x, y])
            if base[3] > 0:
                for channel in range(3):
                    base[channel] = int(mask[channel] * (base[channel] / 255))
                pix[x, y] = (base[0], base[1], base[2], base[3])
    return im_f


def getColorBase(mode: ColorblindMode) -> list[str]:
    """Get the color base array."""
    if mode == ColorblindMode.prot:
        return ["#000000", "#0072FF", "#766D5A", "#FFFFFF", "#FDE400"]
    elif mode == ColorblindMode.deut:
        return ["#000000", "#318DFF", "#7F6D59", "#FFFFFF", "#E3A900"]
    elif mode == ColorblindMode.trit:
        return ["#000000", "#C72020", "#13C4D8", "#FFFFFF", "#FFA4A4"]
    return ["#FFD700", "#FF0000", "#1699FF", "#B045FF", "#41FF25"]


def getKongItemColor(mode: ColorblindMode, kong: Kongs, output_as_list: bool = False) -> str:
    """Get the color assigned to a kong."""
    hash_str = getColorBase(mode)[kong]
    if output_as_list:
        return getRGBFromHash(hash_str)
    return hash_str


def maskImage(im_f, base_index, min_y, keep_dark=False, mode=ColorblindMode.off):
    """Apply RGB mask to image."""
    w, h = im_f.size
    converter = ImageEnhance.Color(im_f)
    im_f = converter.enhance(0)
    im_dupe = im_f.crop((0, min_y, w, h))
    if keep_dark is False:
        brightener = ImageEnhance.Brightness(im_dupe)
        im_dupe = brightener.enhance(2)
    im_f.paste(im_dupe, (0, min_y), im_dupe)
    pix = im_f.load()
    mask = getKongItemColor(mode, base_index, True)
    w, h = im_f.size
    for x in range(w):
        for y in range(min_y, h):
            base = list(pix[x, y])
            if base[3] > 0:
                for channel in range(3):
                    base[channel] = int(mask[channel] * (base[channel] / 255))
                pix[x, y] = (base[0], base[1], base[2], base[3])
    return im_f


def hueShiftImageContainer(table: int, image: int, width: int, height: int, format: TextureFormat, shift: int, ROM_COPY: ROM):
    """Load an image, shift the hue and rewrite it back to ROM."""
    loaded_im = getImageFile(ROM_COPY, table, image, table != 7, width, height, format)
    loaded_im = hueShift(loaded_im, shift)
    loaded_px = loaded_im.load()
    bytes_array = []
    for y in range(height):
        for x in range(width):
            pix_data = list(loaded_px[x, y])
            if format == TextureFormat.RGBA32:
                bytes_array.extend(pix_data)
            elif format == TextureFormat.RGBA5551:
                red = int((pix_data[0] >> 3) << 11)
                green = int((pix_data[1] >> 3) << 6)
                blue = int((pix_data[2] >> 3) << 1)
                alpha = int(pix_data[3] != 0)
                value = red | green | blue | alpha
                bytes_array.extend([(value >> 8) & 0xFF, value & 0xFF])
    px_data = bytearray(bytes_array)
    if table != 7:
        px_data = gzip.compress(px_data, compresslevel=9)
    ROM_COPY.seek(getPointerLocation(table, image))
    ROM_COPY.writeBytes(px_data)


def getLuma(color: tuple) -> float:
    """Get the luma value of a color."""
    return (0.299 * color[0]) + (0.587 * color[1]) + (0.114 * color[2])


def hueShiftColor(color: tuple, amount: int, head_ratio: int = None) -> tuple:
    """Apply a hue shift to a color."""
    # RGB -> HSV Conversion
    red_ratio = color[0] / 255
    green_ratio = color[1] / 255
    blue_ratio = color[2] / 255
    color_max = max(red_ratio, green_ratio, blue_ratio)
    color_min = min(red_ratio, green_ratio, blue_ratio)
    color_delta = color_max - color_min
    hue = 0
    if color_delta != 0:
        if color_max == red_ratio:
            hue = 60 * (((green_ratio - blue_ratio) / color_delta) % 6)
        elif color_max == green_ratio:
            hue = 60 * (((blue_ratio - red_ratio) / color_delta) + 2)
        else:
            hue = 60 * (((red_ratio - green_ratio) / color_delta) + 4)
    sat = 0 if color_max == 0 else color_delta / color_max
    val = color_max
    # Adjust Hue
    if head_ratio is not None and sat != 0:
        amount = head_ratio / (sat * 100)
    hue = (hue + amount) % 360
    # HSV -> RGB Conversion
    c = val * sat
    x = c * (1 - abs(((hue / 60) % 2) - 1))
    m = val - c
    if hue < 60:
        red_ratio = c
        green_ratio = x
        blue_ratio = 0
    elif hue < 120:
        red_ratio = x
        green_ratio = c
        blue_ratio = 0
    elif hue < 180:
        red_ratio = 0
        green_ratio = c
        blue_ratio = x
    elif hue < 240:
        red_ratio = 0
        green_ratio = x
        blue_ratio = c
    elif hue < 300:
        red_ratio = x
        green_ratio = 0
        blue_ratio = c
    else:
        red_ratio = c
        green_ratio = 0
        blue_ratio = x
    return (int((red_ratio + m) * 255), int((green_ratio + m) * 255), int((blue_ratio + m) * 255))


def rgba32to5551(rgba_32: list[int]) -> list[int]:
    """Convert list as RGBA32 bytes with no alpha to list of RGBA5551 bytes."""
    val_r = int((rgba_32[0] >> 3) << 11)
    val_g = int((rgba_32[1] >> 3) << 6)
    val_b = int((rgba_32[2] >> 3) << 1)
    rgba_val = val_r | val_g | val_b | 1
    return [(rgba_val >> 8) & 0xFF, rgba_val & 0xFF]