rivet/renderer.py

from __future__ import annotations

import os
from typing import TYPE_CHECKING

import pygame

from structure import Structure
from vec import Vec

if TYPE_CHECKING:
    from config import Config
    from range import Range
    from schema import InstructionSetSchema


class Renderer:
    WIDTH = 1200
    HEIGHT = 800
    
    def __init__(self, config: Config, display: bool = False) -> None:
        self.config = config
        self.display = display
        pygame.init()
        if self.display:
            self.win = pygame.display.set_mode([Renderer.WIDTH, Renderer.HEIGHT])
        
        self.surf = pygame.Surface([Renderer.WIDTH, Renderer.HEIGHT], pygame.SRCALPHA)
        self.font = pygame.font.SysFont(self.config.DEFAULT_FONT_FAMILY, self.config.DEFAULT_FONT_SIZE)
        self.italicFont = pygame.font.SysFont(self.config.ITALIC_FONT_FAMILY, self.config.ITALIC_FONT_SIZE, italic=True)
        
        self.margins = self.config.MARGINS
    
    def render(self, schema: InstructionSetSchema) -> None:
        self.surf.fill(self.config.BACKGROUND_COLOR)
        
        self.drawStructure(schema.structures["main"], schema.structures, self.margins[3], self.margins[0])
        
        if self.display:
            name = os.path.basename(schema.path)
            pygame.display.set_caption(f"Rivet - {name}")
            self.win.fill(self.config.BACKGROUND_COLOR)
            self.win.blit(self.surf, [0, 0])
            pygame.display.flip()
    
    def save(self, path: str) -> None:
        pygame.image.save(self.surf, path)
    
    def drawStructure(self,
                      struct: Structure,
                      structures: dict[str, Structure],
                      ox: float = 0,
                      oy: float = 0) -> float:
        
        bgCol = self.config.BACKGROUND_COLOR
        txtCol = self.config.TEXT_COLOR
        borderCol = self.config.BORDER_COLOR
        bitW = self.config.BIT_WIDTH
        bitH = self.config.BIT_HEIGHT
        
        bitsX, bitsY = ox, oy + bitH
        bitsWidth = struct.bits * bitW
        startBit = struct.start
        
        # Draw rectangle around structure
        pygame.draw.rect(self.surf, borderCol, [bitsX, bitsY, bitsWidth, bitH], 2)
        
        for i in range(struct.bits):
            bitX = ox + i * bitW
            
            bitITxt = self.font.render(str(struct.bits - i - 1 + startBit), True, txtCol)
            self.surf.blit(bitITxt, [
                bitX + (bitW - bitITxt.get_width())/2,
                oy + (bitH - bitITxt.get_height())/2
            ])
            
            # Draw separator
            if i != 0:
                pygame.draw.line(self.surf, borderCol, [bitX, bitsY], [bitX, bitsY + bitH])
        
        ranges = struct.getSortedRanges()
        descX = ox + max(0, (struct.bits-12) * bitW)
        descY = bitsY + bitH * 2
        
        # Names + simple descriptions
        for range_ in ranges:
            rStartI = struct.bits - range_.end + startBit - 1
            
            rStartX = bitsX + rStartI * bitW
            rWidth = range_.bits * bitW
            
            nameTxt = self.font.render(range_.name, True, txtCol)
            nameX = rStartX + (rWidth - nameTxt.get_width())/2
            nameY = bitsY + (bitH - nameTxt.get_height())/2
            pygame.draw.rect(self.surf, bgCol, [rStartX + bitW/2, nameY, rWidth - bitW, nameTxt.get_height()], 0)
            self.surf.blit(nameTxt, [nameX, nameY])

            if range_.description:
                descX, descY = self.drawDescription(range_, rStartX, bitsY, rWidth, descX, descY)
        
        # Dependencies
        for range_ in ranges:
            if range_.values is not None and range_.dependsOn is not None:
                descX, descY = self.drawDependency(struct, structures, bitsX, bitsY, range_, descX, descY)
        
        return descY
    
    def drawUnderbracket(self, start: float, end: float, bitsY: float) -> None:
        bitW = self.config.BIT_WIDTH
        bitH = self.config.BIT_HEIGHT
        
        x0 = start + bitW/2
        x1 = end - bitW/2
        y0 = bitsY + bitH * 1.25
        y1 = bitsY + bitH * 1.5
        
        pygame.draw.lines(self.surf, self.config.LINK_COLOR, False, [
            [x0, y0],
            [x0, y1],
            [x1, y1],
            [x1, y0]
        ])
    
    def drawLink(self, startX: float, startY: float, endX: float, endY: float) -> None:
        bitH = self.config.BIT_HEIGHT
        arrowMargin = self.config.ARROW_MARGIN
        
        pygame.draw.lines(self.surf, self.config.LINK_COLOR, False, [
            [startX, startY + bitH*1.5],
            [startX, endY + bitH/2],
            [endX - arrowMargin, endY + bitH/2]
        ])
    
    def drawDescription(self,
                        range_: Range,
                        rStartX: float,
                        rStartY: float,
                        rWidth: float,
                        descX: float,
                        descY: float) -> tuple[float, float]:

        bitW = self.config.BIT_WIDTH
        bitH = self.config.BIT_HEIGHT
        
        descX = max(descX, rStartX + rWidth/2 + bitW)
        
        self.drawUnderbracket(rStartX, rStartX + rWidth, rStartY)
        
        midX = rStartX + rWidth/2
        self.drawLink(midX, rStartY, descX, descY)
        
        descTxt = self.font.render(range_.description, True, self.config.TEXT_COLOR)
        self.surf.blit(descTxt, [descX, descY + (bitH - descTxt.get_height())/2])
        
        descY += descTxt.get_height()
        
        if range_.values is not None and range_.dependsOn is None:
            descX, descY = self.drawValues(range_.values, descX, descY)
        
        descY += self.config.DESCRIPTION_MARGIN
        
        return (descX, descY)

    def drawValues(self, values: dict[str, str], descX: float, descY: float) -> tuple[float, float]:
        textCol = self.config.TEXT_COLOR
        bitW = self.config.BIT_HEIGHT
        gap = self.config.VALUES_GAP
        
        for val, desc in sorted(values.items(), key=lambda vd: vd[0]):
            descY += gap
            valTxt = self.italicFont.render(f"{val} = {desc}", True, textCol)
            self.surf.blit(valTxt, [descX + bitW/2, descY])
            descY += valTxt.get_height()
        
        return (descX, descY)

    def drawDependency(self,
                       struct: Structure,
                       structures: dict[str,
                       Structure],
                       bitsX: float,
                       bitsY: float,
                       range_: Range,
                       descX: float,
                       descY: float) -> tuple[float, float]:
        
        bitW = self.config.BIT_WIDTH
        bitH = self.config.BIT_HEIGHT
        arrowMargin = self.config.ARROW_MARGIN
        
        rStartI = struct.bits - range_.end - 1
        rStartX = bitsX + rStartI * bitW
        rWidth = range_.bits * bitW
        
        self.drawUnderbracket(rStartX, rStartX + rWidth, bitsY)
        prevY = bitsY + bitH * 1.5
        dependRange = struct.ranges[range_.dependsOn]
        dependStartI = struct.bits - dependRange.end - 1
        dependStartX = bitsX + dependStartI * bitW
        dependWidth = dependRange.bits * bitW
        dependMid = dependStartX + dependWidth/2
        self.drawUnderbracket(dependStartX, dependStartX + dependWidth, bitsY)
        
        for val, data in sorted(range_.values.items(), key=lambda vd: vd[0]):
            self.drawArrow(dependMid, prevY, dependMid, descY - arrowMargin)
            
            valRanges = {}
            for i in range(dependRange.bits):
                valRanges[dependRange.end - i] = {
                    "name": val[i]
                }
            
            valStruct = {
                "bits": dependRange.bits,
                "start": dependRange.start,
                "ranges": valRanges
            }
            valStruct = Structure.load("", valStruct)
            
            self.drawStructure(valStruct, structures, dependStartX, descY)
            
            y = descY + bitH * 1.5
            
            # Arrow from left to right
            if dependRange.end > range_.start:
                x1 = dependStartX + dependWidth + arrowMargin
                x2 = rStartX - arrowMargin
            
            # Arrow from right to left
            else:
                x1 = dependStartX - arrowMargin
                x2 = rStartX + rWidth + arrowMargin
            
            self.drawArrow(x1, y, x2, y, data["description"])
            
            self.drawArrow(rStartX + rWidth - bitW,
                           prevY,
                           rStartX + rWidth - bitW,
                           descY + bitH - arrowMargin)

            prevY = descY + bitH*2 + arrowMargin
            descY = self.drawStructure(structures[data["structure"]], structures, rStartX, descY)
        
        return (descX, descY)

    def drawArrow(self, sx: float, sy: float, ex: float, ey: float, label: str = "") -> None:
        dashLen = self.config.DASH_LENGTH
        dashSpace = self.config.DASH_SPACE
        arrowSize = self.config.ARROW_SIZE
        linkCol = self.config.LINK_COLOR
        textCol = self.config.TEXT_COLOR
        arrowLabelDist = self.config.ARROW_LABEL_DISTANCE
        
        start = Vec(sx, sy)
        end = Vec(ex, ey)
        startEnd = end - start
        d = startEnd.norm()
        
        dashes = int(startEnd.mag() / (dashLen + dashSpace))
        
        for i in range(dashes):
            a = start + d * i * (dashLen + dashSpace)
            b = a + d*dashLen
            pygame.draw.line(self.surf, linkCol, [a.x, a.y], [b.x, b.y])
        
        n = Vec(d.y, -d.x)
        
        width = arrowSize / 1.5
        p1 = end - d * arrowSize + n * width
        p2 = end - d * arrowSize - n * width
        pygame.draw.polygon(self.surf, linkCol, [
            [end.x, end.y],
            [p1.x, p1.y],
            [p2.x, p2.y]], 0)
        
        if label:
            txt = self.italicFont.render(label, True, textCol)
            self.surf.blit(txt, [
                (start.x + end.x - txt.get_width())/2,
                (start.y + end.y)/2 + arrowLabelDist
            ])