plotter.py

'''
Turn a config file to a 2-D array of pixels with colors, which can be visualized by rendering or plotting later.
'''

import json
from essentials import *

class elements:
    def __init__(self, config: dict, default_color: str | tuple[int, int, int] | None = None):
        '''
        Translate a config to a text element
        '''
        if not isinstance(config, dict): raise TypeError("Config is not dictionary.")
        # Read config
        # Read and check text
        self.text       = config.get("value")
        if not isinstance(self.text, str): raise TypeError("Text will be plotted is not in str format or not given.")
        if self.text == "": raise ValueError("Text is empty.")

        # Read and check font name
        self.font_name  = config.get("font")  or None
        if not isinstance(self.font_name, str): raise KeyError(f"Font name is not in str format or not given. (Occured when plotting {self.text})")
        if self.font_name not in get_font_list(): raise ValueError(f"Font {self.font_name} not found. (Occured when plotting {self.text})")
        
        # Read font
        with open(f'{FONT_DIR}/{self.font_name}.json', 'r') as file:
            font = json.load(file)

        # Read type, colour, distance and caps setting
        self.type       = bool(font.get("is_symbol") or False) # True is symbol, False is text
        if self.type: self.text=[(self.text)]
        self.color      = (normalize_color(default_color),
                           normalize_color(config.get("color") or None))
        
        # char_dist and space_dist can be empty to use data provided from the font
        self.char_dist  = config.get("char_dist")
        if not isinstance(self.char_dist, int): self.char_dist = font.get("char_dist")
        self.space_dist = config.get("space_dist")
        if not isinstance(self.space_dist, int): self.space_dist = font.get("space_dist")
        self.force_caps = config.get("force_caps",True)

        if not isinstance(self.char_dist , int) and len(self.text) > 1: raise KeyError(f"Distance between characters is invalid. (Occured when plotting {self.text})")
        if not isinstance(self.space_dist, int) and " " in list(self.text): raise KeyError(f"Space span between characters is invalid. (Occured when plotting {self.text})")

        self.offset_x = config.get("offset_x") or 0
        self.offset_y = config.get("offset_y") or 0

        for _ in (self.offset_x, self.offset_y):
            if not isinstance(_, int): raise TypeError(f"Offset parameter must be integer. (Occured when plotting {self.text})")
            
        
    def plot(self):
        '''
        Plot text
        '''
        with open(f'{FONT_DIR}/{self.font_name}.json', 'r') as file:
            font = json.load(file)
        mask_height = font.get("height")
        char_lib = font.get("character")

        if mask_height == None: raise KeyError(f"Character height of {self.font_name} not found.")

        mask_width = 0

        mask=[[] for _ in range(mask_height)]

        is_first_digit=True
        # Plot text
        for char in self.text:
            # Space
            if char==" ": 
                for row in range(mask_height): mask[row] += [self.color[0]] * (self.space_dist - self.char_dist)
                mask_width += self.space_dist - self.char_dist
                continue

            # Retrieve character list
            if not self.type:
                if self.force_caps: char = char.upper()
            char_matrix = char_lib.get(char)
            if char_matrix == None: raise KeyError(f"{self.font_name} does not have key {char}")
            # Fill 
            shape = char_matrix.get("shape")
            char_width = char_matrix.get("width") or len(shape[0])
            ret = [[self.color[shape[row][col]] 
                    for col in range(char_width)] 
                    for row in range(mask_height)]
            # Add to mask
            if not is_first_digit: 
                mask_width += self.char_dist
            mask_width += char_width
            for row in range(mask_height):
                if not is_first_digit: 
                    mask[row]+=[self.color[0]]*self.char_dist
                mask[row]+=ret[row]
            if is_first_digit: is_first_digit=False
        return mask, mask_height, mask_width

class components:
    def __init__(self, config: dict, default_color):
        self.align_type    = config.get("align_type")
        raw_elements       = config.get("elements") or []
        self.elements      = []
        self.default_color = default_color
        for element in raw_elements: 
            try:
                cur = elements(element,default_color)
            except ValueError as e:
                if str(e)=="Text is empty.":
                    pass
                else:
                    raise
            except Exception:
                raise
            else:
                self.elements.append(cur)

    def plot(self):
        ret      = [] # will be resized when first element being plotted
        ret_w    = 0
        ret_h    = 0
        offset_x = 0
        offset_y = 0
        is_first = True
        curpos_x = 0
        curpos_y = 0
        for element in self.elements:
            mask, mask_height, mask_width = element.plot()

            if is_first: 
                offset_x = element.offset_x
                offset_y = element.offset_y
                ret = [[] for _ in range(mask_height)]
                ret_h = mask_height
            else: 
                curpos_x += element.offset_x
                curpos_y  = element.offset_y
            is_first = False

            mask_hdiff = ret_h - mask_height
            
            align_offset = 0
            if self.align_type == "middle": curpos_y += mask_hdiff //2
            if self.align_type == "bottom": curpos_y += mask_hdiff

            # allocate space
            if curpos_x<0: 
                adj = -curpos_x
                offset_x -= adj
                for i in range(ret_h): ret[i] = [self.default_color for _ in range(adj)] + ret[i]
                curpos_x = 0
                ret_w += adj
            if curpos_y <0:
                adj = -curpos_y
                offset_y -= adj
                ret = [[self.default_color for _ in range(ret_w)] for _ in range(adj)] + ret
                curpos_y = 0
                ret_h += adj
            
            end_x = curpos_x + mask_width
            end_y = curpos_y + mask_height

            if end_x >= ret_w:
                diff = end_x - ret_w
                for i in range(ret_h): ret[i] += [self.default_color for _ in range(diff)]
                ret_w = end_x
            
            if end_y >= ret_h:
                diff = end_y - ret_h
                ret += [[self.default_color for _ in range(ret_w)] for _ in range(diff)]
                ret_h = end_y
            
            # Plot
            for i in range(mask_height):
                ret[i+curpos_y+align_offset] = ret[i+curpos_y+align_offset][:ret_w-mask_width] + mask[i]

            curpos_x += mask_width

        return ret, offset_x, offset_y

def merge_mask(original_mask: list, new_mask: list, pos_w: int, pos_h: int, default_color: tuple[int,int,int] | None):
    ''' Override a mask onto portion of another mask
    '''
    px_count_height=len(original_mask)
    px_count_width=len(original_mask[0])

    for i in range(len(new_mask)):
        if not 0<=pos_h+i<px_count_height: continue
        for j in range(len(new_mask[i])):
            if not 0<=pos_w+j<px_count_width: continue
            if original_mask[pos_h+i][pos_w+j] == default_color:
                original_mask[pos_h+i][pos_w+j]=new_mask[i][j]
    return original_mask

def plot(config: dict) -> list:
    '''Plot the headsign as a matrix of colors accroding to the config'''
    # Read config
    file_name         = config.get("file_name") or "LED.svg"
    svg_style         = config.get("svg_style")
    if not isinstance(svg_style, dict): raise TypeError("Invalid SVG style type")
    default_color = normalize_color(svg_style.get("px_default_color") )

    dimensions        = config.get("dimensions")
    if not isinstance(dimensions, dict): raise TypeError("Invalid dimension type")

    width_px       = dimensions.get("width_px")
    height_px      = dimensions.get("height_px")
    for _ in (width_px, height_px):
        if not isinstance(_, int): raise TypeError("Dimension parameter must be integer.")
        if _ <= 0: raise ValueError("Dimension parameter must be positive.")
    
    components_left   = components(config.get("components_left") or DEFAULT_COMPONENTS, default_color)
    components_middle_all = config.get("components_middle") or None
    components_middle = components_middle_all.get("middle") or None
    if components_middle == None:
        components_up   = components(components_middle_all.get("up")   or DEFAULT_COMPONENTS, default_color)
        components_down = components(components_middle_all.get("down") or DEFAULT_COMPONENTS, default_color)
    else:
        components_middle = components(components_middle_all.get("middle"), default_color)
    components_right  = components(config.get("components_right")  or DEFAULT_COMPONENTS, default_color)
    
    # Initialize headsign
    sign = [[default_color for _ in range(width_px)]
                          for _ in range(height_px)]

    # Plot left part
    mask, offset_x, offset_y = components_left.plot()
    base_x = 0
    base_y = 0
    sign = merge_mask(sign, mask, base_x + offset_x, base_y + offset_y, default_color)
    width_L = max(get_mask_width(mask) + offset_x, 0)

    # Plot right part
    mask, offset_x, offset_y = components_right.plot()
    base_x = width_px - get_mask_width(mask)
    base_y = 0
    sign = merge_mask(sign, mask, base_x + offset_x, base_y + offset_y, default_color)
    width_R = max(get_mask_width(mask) + (-offset_x), 0)

    width_mid = width_px - width_L - width_R
    mid_px = width_L + width_mid//2

    # Plot middle part
    if components_middle == None:
        # up
        mask, offset_x, offset_y = components_up.plot()
        base_x = mid_px - get_mask_width(mask) //2
        base_y = 0
        sign = merge_mask(sign, mask, base_x + offset_x, base_y + offset_y, default_color)

        # down
        mask, offset_x, offset_y = components_down.plot()
        base_x = mid_px - get_mask_width(mask) //2
        base_y = height_px - len(mask)
        sign = merge_mask(sign, mask, base_x + offset_x, base_y + offset_y, default_color)
    else:
        mask, offset_x, offset_y = components_middle.plot()
        base_x = mid_px - get_mask_width(mask) //2
        base_y = (height_px - len(mask)) //2
        sign = merge_mask(sign, mask, base_x + offset_x, base_y + offset_y, default_color)

    return sign

if __name__=="__main__":
    import json
    with open("config.json", "r") as f:
        config = json.load(f)
    from mask_to_svg import plot_svg
    mask = plot(config)
    plot_svg(mask, config.get("svg_style"))