from PIL import Image, ImageDraw, ImageFilter, ImageFont
import numpy as np
from transformers.tools import Tool

class PlaceholderImageTool(Tool):
    """Replacement 'image_generator'
    - transformers HfAgent likes to use the image_generator tool.
    - I don't have the disk space for that.
    - This replacement provides a random placeholder image instead.
    Install it thusly: `agent.toolbox['image_generator']=PlaceholderImageTool()`
    """
    name = "image_generator"
    description = "I will generate an image based on a prompt."
    inputs = ["text"]
    outputs = ["image"]
    def sepia(self, dim:int=64) -> Image:
        "produce a sepia image, size=(dim,dim,3)"
        data = np.random.randint(0,255,(dim,dim,3), dtype=np.uint8)
        image = Image.fromarray(data)
        sepia_matrix = [
            0.393, 0.769, 0.189, 0,
            0.349, 0.686, 0.168, 0,
            0.272, 0.534, 0.131, 0
        ]
        applied = image.convert('RGB', sepia_matrix)
        return applied
    def __call__(self, prompt):
        "returns a sepia image overlaid with a random digit (1,6 inclusive)"
        DIM=64
        image = self.sepia(DIM)
        #image = image.filter(ImageFilter.SMOOTH)
        draw = ImageDraw.Draw(image)
        digit = str(np.random.randint(1,7))
        font = ImageFont.truetype("./Lora-Bold.ttf", size=48)
        draw.text((DIM//2,DIM//2), digit, fill=(0,0,0), font=font, anchor="mm")
        return image