ddrescue_visualisation/ddresc_visu/draw.py

from .mapfile import MapFile

import pygame

colors = {
	'?' : (64 , 64 , 64) , # Non-tried
	'*' : (0  , 0  , 255), # Non-trimmed
	'/' : (255, 255, 0)  , # Non-scraped
	'-' : (255, 0  , 0)  , # Bad
	'+' : (0  , 128, 0)  , # Recovered
	'@' : (0  , 255, 255), # Own position
	None: (0  , 0  , 0)  , # Not on disk
}

BORDERCOLOR = (64,64,64)

W = 1920
H = 1040
CELLSIZE = 8
CELLBORDER = 0
CELLMARGIN = 3

CELLOFFSET = CELLSIZE + 2*CELLBORDER + CELLMARGIN
COLUMNS = W // CELLOFFSET
ROWS = H // CELLOFFSET
CELLS = COLUMNS * ROWS

def get_symb(hist):
	## Majority:
	#symb, _count = max(hist.items(), key=lambda x: x[1])
	#return symb
	
	# The worst:
	order = ['@', '-', '/', '*', '?', '+']
	for ch in order:
		if ch in hist:
			return ch

class Visualisation:
	def __init__(self, fn, **kwargs):
		pygame.init()
		# https://stackoverflow.com/questions/46160971/pygame-full-core-usage-in-simple-loop
		pygame.mixer.quit()
		self.disp = pygame.display.set_mode((W, H))

		self.mapfile = MapFile(fn, **kwargs)
		self.mapfile.load()

	def draw(self):
		self.disp.fill((0,0,0)) # Dark theme :-)

		sz = self.mapfile.size
		sqsz = max(sz / CELLS, 4096)
		print(f"INFO: Each square represents {sqsz} bytes")
		for cid in range(CELLS):
			start = cid * sqsz
			end = (cid + 1) * sqsz
			hist = self.mapfile.get_hist(start, end)
			symb = get_symb(hist)
			color = colors[symb]

			line = cid // COLUMNS
			col = cid % COLUMNS
			pos = (line, col)

			rect = pygame.Rect(col * CELLOFFSET + CELLBORDER, line * CELLOFFSET + CELLBORDER, CELLSIZE+2*CELLBORDER, CELLSIZE+2*CELLBORDER)
			pygame.draw.rect(self.disp, color, rect)
			if CELLBORDER != 0: pygame.draw.rect(self.disp, BORDERCOLOR, rect, width=CELLBORDER)

		pygame.display.flip()

	def reload(self):
		self.mapfile.load()
		self.draw()
	
	def run(self, *, refresh=2):
		pygame.time.set_timer(pygame.USEREVENT, refresh * 1000)
		pygame.event.set_blocked(None)
		pygame.event.set_allowed([pygame.USEREVENT, pygame.VIDEORESIZE, pygame.VIDEOEXPOSE, pygame.QUIT, pygame.KEYDOWN])
		self.draw()
		while True:
			ev = pygame.event.wait()
			if ev.type == pygame.QUIT or ev.type == pygame.KEYDOWN and ev.key == pygame.K_q:
				return
			if ev.type in [pygame.VIDEOEXPOSE, pygame.VIDEORESIZE, pygame.USEREVENT]:
				self.reload()