import json from typing import Iterable from dataclasses import dataclass, asdict from src.evilnkode import Observation from src.utils import observations, passcode_generator from pathlib import Path from PIL import Image, ImageDraw, ImageFont from src.keypad.keypad import ( BaseKeypad, SlidingSplitShuffleKeypad, SlidingTowerShuffleKeypad, RandomShuffleKeypad, RandomSplitShuffleKeypad, ) import argparse @dataclass class ObservationSequence: target_passcode: list[int] observations: list[Observation] def new_observation_sequence( keypad: BaseKeypad, passcode_len: int, complexity: int, disparity: int, numb_runs: int, ) -> ObservationSequence: passcode = passcode_generator(keypad.k, keypad.p, passcode_len, complexity, disparity) obs_gen = observations( keypad=keypad, target_passcode=passcode, number_of_observations=numb_runs, ) return ObservationSequence(target_passcode=passcode, observations=[obs for obs in obs_gen]) def _next_json_filename(base_dir: Path) -> Path: """Find the next available observation_X.json file in base_dir.""" counter = 1 while True: candidate = base_dir / f"observation_{counter}.json" if not candidate.exists(): return candidate counter += 1 def save_observation_sequence_to_json(seq: ObservationSequence, filename: Path) -> None: filename.parent.mkdir(parents=True, exist_ok=True) with filename.open("w", encoding="utf-8") as f: json.dump(asdict(seq), f, indent=4) # ---------- Helpers ---------- def _load_font(preferred: str, size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: """Try a preferred TTF, fall back to common monospace, then PIL default.""" candidates = [ preferred, "DejaVuSansMono.ttf", # common on Linux "Consolas.ttf", # Windows "Menlo.ttc", "Menlo.ttf", # macOS "Courier New.ttf", ] for c in candidates: try: return ImageFont.truetype(c, size) except Exception: continue return ImageFont.load_default() def _text_size(draw: ImageDraw.ImageDraw, text: str, font: ImageFont.ImageFont) -> tuple[int, int]: """Get (w, h) using font bbox for accurate layout.""" left, top, right, bottom = draw.textbbox((0, 0), text, font=font) return int(right - left), int(bottom - top) def _join_nums(nums: Iterable[int]) -> str: return " ".join(str(n) for n in nums) def _next_available_path(path: Path) -> Path: """If path exists, append _1, _2, ...""" if not path.exists(): return path base, suffix = path.stem, path.suffix or ".png" i = 1 while True: candidate = path.with_name(f"{base}_{i}{suffix}") if not candidate.exists(): return candidate i += 1 # ---------- Core rendering ---------- def render_observation_to_png( target_passcode: list[int], obs: Observation, out_path: Path, *, header_font_name: str = "DejaVuSans.ttf", body_font_name: str = "DejaVuSans.ttf", header_size: int = 28, body_size: int = 24, margin: int = 32, row_padding_xy: tuple[int, int] = (16, 12), # (x, y) padding inside row box row_spacing: int = 14, header_spacing: int = 10, section_spacing: int = 18, bg_color: str = "white", fg_color: str = "black", row_fill: str = "#f7f7f7", row_outline: str = "#222222", ): """ Render a single observation: - Top lines: Target Passcode: {target} Selected Keys: {selected keys} - Then a stack of row boxes representing the keypad rows. """ out_path.parent.mkdir(parents=True, exist_ok=True) out_path = _next_available_path(out_path) # Fonts header_font = _load_font(header_font_name, header_size) body_font = _load_font(body_font_name, body_size) # Prepare strings header1 = f"Target Passcode: {_join_nums(target_passcode)}" header2 = f"Selected Keys: {_join_nums(obs.key_selection)}" row_texts = [_join_nums(row) for row in obs.keypad] # Measure to compute canvas size # Provisional image for measurement temp_img = Image.new("RGB", (1, 1), bg_color) d = ImageDraw.Draw(temp_img) h1_w, h1_h = _text_size(d, header1, header_font) h2_w, h2_h = _text_size(d, header2, header_font) row_text_sizes = [_text_size(d, t, body_font) for t in row_texts] row_box_widths = [tw + 2 * row_padding_xy[0] for (tw, th) in row_text_sizes] row_box_heights = [th + 2 * row_padding_xy[1] for (tw, th) in row_text_sizes] content_width = max([h1_w, h2_w] + (row_box_widths or [0])) total_rows_height = sum(row_box_heights) + row_spacing * max(0, len(row_box_heights) - 1) width = content_width + 2 * margin height = ( margin + h1_h + header_spacing + h2_h + section_spacing + total_rows_height + margin ) # Create final image img = Image.new("RGB", (max(width, 300), max(height, 200)), bg_color) draw = ImageDraw.Draw(img) # Draw headers x = margin y = margin draw.text((x, y), header1, font=header_font, fill=fg_color) y += h1_h + header_spacing draw.text((x, y), header2, font=header_font, fill=fg_color) y += h2_h + section_spacing # Draw row boxes with evenly spaced numbers max_box_width = max(row_box_widths) if row_box_widths else 0 for row, box_h in zip(obs.keypad, row_box_heights): box_left = x box_top = y box_right = x + max_box_width box_bottom = y + box_h # draw row rectangle draw.rectangle( [box_left, box_top, box_right, box_bottom], fill=row_fill, outline=row_outline, width=2 ) # evenly spaced numbers n = len(row) if n > 0: available_width = max_box_width - 2 * row_padding_xy[0] spacing = available_width / (n + 1) for idx, num in enumerate(row, start=1): num_text = str(num) num_w, num_h = _text_size(draw, num_text, body_font) num_x = box_left + row_padding_xy[0] + spacing * idx - num_w / 2 num_y = box_top + (box_h - num_h) // 2 draw.text((num_x, num_y), num_text, font=body_font, fill=fg_color) y = box_bottom + row_spacing img.save(out_path, format="PNG") def _next_run_dir(base_dir: Path) -> Path: """Find the next available run directory under base_dir (run_001, run_002, ...).""" counter = 1 while True: run_dir = base_dir / f"run_{counter:03d}" if not run_dir.exists(): run_dir.mkdir(parents=True) return run_dir counter += 1 def render_sequence_to_pngs(seq: ObservationSequence, out_dir: Path) -> None: out_dir.mkdir(parents=True, exist_ok=True) run_dir = _next_run_dir(out_dir) for i, obs in enumerate(seq.observations, start=1): filename = run_dir / f"observation_{i:03d}.png" render_observation_to_png(seq.target_passcode, obs, filename) if __name__ == "__main__": shuffle_classes = { 'RandomSplitShuffle': RandomSplitShuffleKeypad, 'RandomShuffle': RandomShuffleKeypad, 'SlidingSplitShuffle': SlidingSplitShuffleKeypad, 'SlidingTowerShuffle': SlidingTowerShuffleKeypad } parser = argparse.ArgumentParser(description="Generate and save observation sequences with optional PNG rendering.") parser.add_argument("--number-of-keys", type=int, default=6, help="Number of keys in the keypad (default: 6)") parser.add_argument("--properties-per-key", type=int, default=9, help="Properties per key (default: 9)") parser.add_argument("--passcode-length", type=int, default=4, help="Length of the passcode (default: 4)") parser.add_argument("--complexity", type=int, default=0, help="Complexity of the passcode (default: 0)") parser.add_argument("--disparity", type=int, default=0, help="Disparity of the passcode (default: 0)") parser.add_argument("--num-runs", type=int, default=50, help="Number of observations to generate (default: 50)") parser.add_argument("--shuffle-type", type=str, default="SlidingTowerShuffle", choices=list(shuffle_classes.keys()), help="Keypad shuffle type: 'RandomShuffle' or 'SlidingTowerShuffle' (default: SlidingTowerShuffle)") parser.add_argument("--output-dir", type=str, default="./output", help="Custom output directory for JSON and PNG files") args = parser.parse_args() keypad = shuffle_classes[args.shuffle_type].new_keypad(6, 9) obs_seq = new_observation_sequence(keypad, 4, 0, 0, numb_runs=50) shuffle_type = str(type(keypad)).lower().split('.')[-1].replace("'>", "") output_dir = Path(args.output_dir) save_observation_sequence_to_json(obs_seq, output_dir / "obs.json") render_sequence_to_pngs(obs_seq, output_dir / "obs_png")