from PIL import Image, ImageDraw, ImageFont
import numpy as np

class FrameRenderer:
    """
    Helper class for drawing headers, text, and observation images onto a PIL.Image frame.
    """

    def __init__(
        self,
        frame_width,
        frame_height,
        header_rect,
        obs_rect,
        obs_low_res_rect=None,
        font_size=16,
        font_path=None,
    ):
        """
        frame_width, frame_height: output frame size
        header_rect: (x, y, w, h) region for header
        obs_rect: (x, y, w, h) region for main observation
        obs_low_res_rect: (x, y, w, h) region for low-res obs (optional)
        font_size: font size for text
        font_path: path to .ttf font file (optional)
        """
        self.frame_width = frame_width
        self.frame_height = frame_height
        self.header_rect = header_rect
        self.obs_rect = obs_rect
        self.obs_low_res_rect = obs_low_res_rect
        self.font_size = font_size

        if font_path:
            self.font = ImageFont.truetype(font_path, font_size)
        else:
            self.font = ImageFont.load_default()

    def clear_header(self, image):
        """
        Draw a black rectangle with white outline at the header area of image.
        """
        draw = ImageDraw.Draw(image)
        x, y, w, h = self.header_rect
        draw.rectangle([x, y, x + w, y + h], fill="black", outline="white", width=1)

    def draw_text(self, image, text, idx_line, idx_column, num_cols):
        """
        Draw text at (row, column) in the header area of image.
        """
        x_header, y_header, header_width, header_height = self.header_rect
        x_pos = 5 + idx_column * int(header_width // num_cols)
        y_pos = 5 + idx_line * self.font_size
        assert (0 <= x_pos <= header_width) and (0 <= y_pos <= header_height)
        draw = ImageDraw.Draw(image)
        draw.text((x_header + x_pos, y_header + y_pos), text, font=self.font, fill="white")

    def draw_obs(self, image, obs, obs_low_res=None):
        """
        Draw the main observation (and optional low-res obs) onto the image.
        """
        # Main observation
        assert obs.ndim == 4 and obs.size(0) == 1
        obs_img = Image.fromarray(
            obs[0].add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy()
        )
        x, y, w, h = self.obs_rect
        obs_img = obs_img.resize((w, h), resample=Image.BICUBIC)
        image.paste(obs_img, (x, y))

        # Small resolution observation (if any)
        if obs_low_res is not None and self.obs_low_res_rect is not None:
            assert obs_low_res.ndim == 4 and obs_low_res.size(0) == 1
            obs_lr_img = Image.fromarray(
                obs_low_res[0].add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy()
            )
            x_lr, y_lr, w_lr, h_lr = self.obs_low_res_rect
            obs_lr_img = obs_lr_img.resize((w_lr, h_lr), resample=Image.BICUBIC)
            image.paste(obs_lr_img, (x_lr, y_lr))

    def draw_header_and_text(self, image, header):
        """
        Draw header grid and all text in the header area.
        header: list of columns, each column is a list of text (rows)
        """
        self.clear_header(image)
        num_cols = len(header)
        for j, col in enumerate(header):
            for i, row in enumerate(col):
                self.draw_text(image, row, idx_line=i, idx_column=j, num_cols=num_cols)