1"""writes indexes to the model directory for the frontend to use or for record keeping"""
  2
  3import importlib.metadata
  4import importlib.resources
  5import inspect
  6import itertools
  7import json
  8from collections.abc import Callable
  9from pathlib import Path
 10from typing import Literal
 11
 12import pattern_lens
 13from pattern_lens.attn_figure_funcs import (
 14	_FIGURE_NAMES_KEY,
 15	ATTENTION_MATRIX_FIGURE_FUNCS,
 16)
 17
 18
 19def generate_prompts_jsonl(model_dir: Path) -> None:
 20	"""creates a `prompts.jsonl` file with all the prompts in the model directory
 21
 22	looks in all directories in `{model_dir}/prompts` for a `prompt.json` file
 23	"""
 24	prompts: list[dict] = list()
 25	for prompt_dir in (model_dir / "prompts").iterdir():
 26		prompt_file: Path = prompt_dir / "prompt.json"
 27		if prompt_file.exists():
 28			with open(prompt_file, "r") as f:
 29				prompt_data: dict = json.load(f)
 30				prompts.append(prompt_data)
 31
 32	with open(model_dir / "prompts.jsonl", "w") as f:
 33		for prompt in prompts:
 34			f.write(json.dumps(prompt))
 35			f.write("\n")
 36
 37
 38def generate_models_jsonl(path: Path) -> None:
 39	"""creates a `models.jsonl` file with all the models"""
 40	models: list[dict] = list()
 41	for model_dir in (path).iterdir():
 42		model_cfg_path: Path = model_dir / "model_cfg.json"
 43		if model_cfg_path.exists():
 44			with open(model_cfg_path, "r") as f:
 45				model_cfg: dict = json.load(f)
 46				models.append(model_cfg)
 47
 48	with open(path / "models.jsonl", "w") as f:
 49		for model in models:
 50			f.write(json.dumps(model))
 51			f.write("\n")
 52
 53
 54def get_func_metadata(func: Callable) -> list[dict[str, str | None]]:
 55	"""get metadata for a function
 56
 57	# Parameters:
 58	- `func : Callable` which has a `_FIGURE_NAMES_KEY` (by default `_figure_names`) attribute
 59
 60	# Returns:
 61
 62	`list[dict[str, str | None]]`
 63	each dictionary is for a function, containing:
 64
 65	- `name : str` : the name of the figure
 66	- `func_name : str`
 67		the name of the function. if not a multi-figure function, this is identical to `name`
 68		if it is a multi-figure function, then `name` is `{func_name}.{figure_name}`
 69	- `doc : str` : the docstring of the function
 70	- `figure_save_fmt : str | None` : the format of the figure that the function saves, using the `figure_save_fmt` attribute of the function. `None` if the attribute does not exist
 71	- `source : str | None` : the source file of the function
 72	- `code : str | None` : the source code of the function, split by line. `None` if the source file cannot be read
 73
 74	"""
 75	source_file: str | None = inspect.getsourcefile(func)
 76	output: dict[str, str | None] = dict(
 77		func_name=func.__name__,
 78		doc=func.__doc__,
 79		figure_save_fmt=getattr(func, "figure_save_fmt", None),
 80		source=Path(source_file).as_posix() if source_file else None,
 81	)
 82
 83	try:
 84		output["code"] = inspect.getsource(func)
 85	except OSError:
 86		output["code"] = None
 87
 88	fig_names: list[str] | None = getattr(func, _FIGURE_NAMES_KEY, None)
 89	if fig_names:
 90		return [
 91			{
 92				"name": func_name,
 93				**output,
 94			}
 95			for func_name in fig_names
 96		]
 97	else:
 98		return [
 99			{
100				"name": func.__name__,
101				**output,
102			},
103		]
104
105
106def generate_functions_jsonl(path: Path) -> None:
107	"unions all functions from `figures.jsonl` and `ATTENTION_MATRIX_FIGURE_FUNCS` into the file"
108	figures_file: Path = path / "figures.jsonl"
109	existing_figures: dict[str, dict] = dict()
110
111	if figures_file.exists():
112		with open(figures_file, "r") as f:
113			for line in f:
114				func_data: dict = json.loads(line)
115				existing_figures[func_data["name"]] = func_data
116
117	# Add any new functions from ALL_FUNCTIONS
118	new_functions_lst: list[dict] = list(
119		itertools.chain.from_iterable(
120			get_func_metadata(func) for func in ATTENTION_MATRIX_FIGURE_FUNCS
121		),
122	)
123	new_functions: dict[str, dict] = {func["name"]: func for func in new_functions_lst}
124
125	all_functions: list[dict] = list(
126		{
127			**existing_figures,
128			**new_functions,
129		}.values(),
130	)
131
132	with open(figures_file, "w") as f:
133		for func_meta in sorted(all_functions, key=lambda x: x["name"]):
134			json.dump(func_meta, f)
135			f.write("\n")
136
137
138def inline_assets(
139	html: str,
140	assets: list[tuple[Literal["script", "style"], str]],
141	base_path: Path,
142) -> str:
143	"""Inline specified local CSS/JS files into an HTML document.
144
145	Each entry in `assets` should be a tuple like `("script", "app.js")` or `("style", "style.css")`.
146
147	# Parameters:
148	- `html : str`
149		input HTML content.
150	- `assets : list[tuple[Literal["script", "style"], str]]`
151		List of (tag_type, filename) tuples to inline.
152
153	# Returns:
154	`str` : Modified HTML content with inlined assets.
155	"""
156	for tag_type, filename in assets:
157		if tag_type not in ("style", "script"):
158			err_msg: str = f"Unsupported tag type: {tag_type}"
159			raise ValueError(err_msg)
160
161		# Dynamically create the pattern for the given tag and filename
162		pattern: str = rf'<{tag_type} src="{filename}"></{tag_type}>'
163		# assert it's in the text exactly once
164		assert html.count(pattern) == 1, (
165			f"Pattern {pattern} should be in the html exactly once, found {html.count(pattern) = }"
166		)
167		# read the content and create the replacement
168		content: str = (base_path / filename).read_text()
169		replacement: str = f"<{tag_type}>\n{content}\n</{tag_type}>"
170		# perform the replacement
171		html = html.replace(pattern, replacement)
172
173	return html
174
175
176def write_html_index(path: Path) -> None:
177	"""writes an index.html file to the path"""
178	# TYPING: error: Argument 1 to "Path" has incompatible type "Traversable"; expected "str | PathLike[str]"  [arg-type]
179	frontend_resources_path: Path = Path(
180		importlib.resources.files(pattern_lens).joinpath("frontend"),  # type: ignore[arg-type]
181	)
182	html_index: str = (frontend_resources_path / "index.template.html").read_text(
183		encoding="utf-8",
184	)
185	# inline assets
186	html_index = inline_assets(
187		html_index,
188		[
189			("style", "style.css"),
190			("script", "util.js"),
191			("script", "app.js"),
192		],
193		base_path=frontend_resources_path,
194	)
195
196	# add version
197	pattern_lens_version: str = importlib.metadata.version("pattern-lens")
198	html_index = html_index.replace("$$PATTERN_LENS_VERSION$$", pattern_lens_version)
199	# write the index.html file
200	with open(path / "index.html", "w", encoding="utf-8") as f:
201		f.write(html_index)