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  , 0  , 0)  , # Recovered
	'@' : (0  , 255, 255), # Own position
	None: (0  , 0  , 0)  , # Not on disk
}

BORDERCOLOR = (64,64,64)

W = 1920 // 2 - 2
W = 1920 - 2
H = 1040
CELLSIZE = 8
CELLBORDER = 0
CELLMARGIN = 3

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

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()
		self.histograms = None
		self.squares = None

		# Order of colors to show...
		self.order = ['@', '-', '/', '*', '?', '+']

	def gen_histograms(self):
		self.histograms = []
		sz = self.mapfile.size
		sqsz = max(sz / CELLS, 4096)
		for cid in range(CELLS):
			start = cid * sqsz
			end = (cid + 1) * sqsz
			hist = self.mapfile.get_hist(start, end)
			self.histograms.append(hist)

	
	def gen_squares(self):
		if self.histograms is None: self.gen_histograms()
		self.squares = []
		for hist in self.histograms:
			symb = self.get_symb(hist)
			color = colors[symb]
			self.squares.append(color)

	def draw(self):
		self.disp.fill((0,0,0)) # Dark theme :-)
		if self.squares is None: self.gen_squares()
		for cell, color in enumerate(self.squares):
			line = cell // COLUMNS
			col = cell % COLUMNS
			pos = (line, col)
#			DARKEN = 0.7
#			if line % 2 == 0: color = (color[0] * DARKEN, color[1] * DARKEN, color[2] * DARKEN)

			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.histograms = None
		self.squares = None
		self.draw()

	def get_symb(self, hist):
		## Majority:
		#symb, _count = max(hist.items(), key=lambda x: x[1])
		#return symb
		
		for ch in self.order:
			if ch in hist:
				return ch

	
	def run(self, *, refresh=1800):
		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])
#		pygame.event.set_allowed([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]:
				self.draw()
			if ev.type == pygame.KEYDOWN and ev.key == pygame.K_r:
				self.reload()
			if ev.type == pygame.KEYDOWN and ev.key == pygame.K_a:
				self.order = ['@', '-', '/', '*', '?', '+']
				self.squares = None
				self.draw()
			if ev.type == pygame.KEYDOWN and ev.key == pygame.K_m:
				self.order = ['@', '?', '-', '/', '*', '+']
				self.squares = None
				self.draw()
			if ev.type == pygame.KEYDOWN and ev.key == pygame.K_g:
				colors['+'] = (0, 128, 0)
				self.squares = None
				self.draw()
			if ev.type == pygame.KEYDOWN and ev.key == pygame.K_b:
				colors['+'] = (0, 0, 0)
				self.squares = None
				self.draw()
			if ev.type == pygame.KEYDOWN and ev.key == pygame.K_x:
				self.order = ['@', '+', '?', '-', '/', '*']
				self.squares = None
				self.draw()
			if ev.type == pygame.KEYDOWN and ev.key == pygame.K_p:
				sz = self.mapfile.size
				sqsz = max(sz / CELLS, 4096)
				print(f"Each square represents {sqsz} bytes.")
			if ev.type == pygame.USEREVENT:
				self.reload()