Train¶

class LaTexTableLogger(Logger):
    # TODO: settings has to be configurable in Trainer

    """A Lightning Logger that accumulates metrics and writes a LaTex comparison table.

    Multiple instances (one per model) share a class-level store keyed by experiment_name.
    Every time finalize() is called (after fit() or test() for each model), the current
    state of all accumulated metrics is written to a LaTex file. The last model to
    finalize produces the most complete table.

    This means the file is progressively updated as models finish training/testing,
    so you can open it mid-run to see partial results.

    Args:
        save_dir: Base directory where the comparison/ subfolder will be created.
        model_name: The model's full name (e.g., "mlp:mean").
        experiment_name: Shared key that groups all models in the same experiment.
        precision: Decimal places for metric values in the table.
    """

    # Class-level shared store: {experiment_name: {model_name: {metric_name: value}}}
    __shared_stores: ClassVar[dict[str, dict[str, dict[str, Any]]]] = {}

    def __init__(
        self,
        save_dir: str | Path,
        model_name: str,
        experiment_name: str,
        precision: int = 4,
        options: LaTexTableConfig | None = None,
    ) -> None:
        super().__init__()
        self.__save_dir = save_dir
        self.__model_name = model_name
        self.__experiment_name = experiment_name
        self.__precision = precision

        default_empty_options: LaTexTableConfig = {}
        self.__options = options if options is not None else default_empty_options

        if experiment_name not in self.__shared_stores:
            self.__shared_stores[experiment_name] = {}

    @property
    def name(self) -> str:
        return "LaTexTableLogger"

    @property
    def version(self) -> str | int:
        return self.__model_name

    @property
    def store(self) -> dict[str, dict[str, Any]]:
        """Access the shared store for the current experiment."""
        return dict(self.__shared_stores.get(self.__experiment_name, {}))

    @property
    def save_dir(self) -> str | Path:
        return self.__save_dir

    @property
    def experiment_name(self) -> str | Path:
        return self.__experiment_name

    def log_metrics(self, metrics: dict[str, Any], step: int | None = None) -> None:
        """Accumulate metrics for this model. Called by Lightning on every log step.

        Keeps only the latest value for each metric name. For example, if
        "val_auc" is logged at step 10 and step 20, only the step 20 value is kept.
        """
        store = self.__shared_stores[self.__experiment_name]
        if self.__model_name not in store:
            store[self.__model_name] = {}
        store[self.__model_name].update(metrics)

    def log_hyperparams(self, params: Any) -> None:
        pass

    def finalize(self, status: str) -> None:
        """Write the LaTex comparison table with all accumulated metrics so far.

        Called by Lightning after fit() and after test() for each model. Since models
        train/test sequentially, each finalize() overwrites the file with all data
        accumulated up to that point. The file grows more complete over time.
        """
        test_results, train_results, val_results = self.__split_results()
        if not test_results and not train_results and not val_results:
            return

        comparison_dir = Path(self.__save_dir) / "comparison"
        table_caption_opt = self.__options.get("table_caption")
        sort_by_opt = self.__options.get("sort_by")
        border_opt = self.__options.get("border")

        table_caption = table_caption_opt if isinstance(table_caption_opt, str) else None
        sort_by = sort_by_opt if isinstance(sort_by_opt, list) and sort_by_opt else ["asc"]
        border = border_opt if isinstance(border_opt, bool) else True

        self.__save_comparison_tables(
            test_results=test_results,
            save_dir=comparison_dir,
            train_results=train_results or None,
            val_results=val_results or None,
            precision=self.__precision,
            table_caption=table_caption,
            sort_by=sort_by,
            border=border,
        )
        self.__save_comparison_tables(
            test_results=test_results,
            save_dir=comparison_dir,
            train_results=None,
            val_results=None,
            precision=self.__precision,
            filename="test.tex",
            table_caption=table_caption,
            sort_by=sort_by,
            border=border,
        )

    def __split_results(
        self,
    ) -> tuple[dict[str, dict[str, Any]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]]:
        """
        Split all accumulated metrics into test vs train/val groups.

        Metrics are classified by their name prefix:
        - "test_*"  --> test_results
        - "train_*" --> train_results
        - "val_*" --> val_results
        - anything else (e.g., "epoch") --> ignored
        """
        store = self.__shared_stores.get(self.__experiment_name, {})
        test_results: dict[str, dict[str, Any]] = {}
        train_results: dict[str, dict[str, Any]] = {}
        val_results: dict[str, dict[str, Any]] = {}

        for model_name, metrics in store.items():
            test_metrics: dict[str, Any] = {}
            train_metrics: dict[str, Any] = {}
            val_metrics: dict[str, Any] = {}

            for metric_name, value in metrics.items():
                if metric_name.startswith("test"):
                    test_metrics[metric_name] = value
                elif metric_name.startswith("train"):
                    train_metrics[metric_name] = value
                elif metric_name.startswith("val"):
                    val_metrics[metric_name] = value

            if test_metrics:
                test_results[model_name] = test_metrics
            if train_metrics:
                train_results[model_name] = train_metrics
            if val_metrics:
                val_results[model_name] = val_metrics

        return test_results, train_results, val_results

    def clear(self, experiment_name: str) -> None:
        """Remove accumulated data for an experiment."""
        self.__shared_stores.pop(experiment_name, None)

    def __build_comparison_table(
        self,
        sections_data: list[tuple[str, Mapping[str, Mapping[str, Any]]]],
        precision: int = 4,
        table_caption: str | None = None,
        sort_by: list[str] | None = None,
        border: bool = True,
    ) -> str:
        if not sections_data:
            return ""

        # One tabular must have fixed column count; use max needed across sections.
        max_metrics = max(len({m for mm in rs.values() for m in mm}) for _, rs in sections_data)
        total_cols = 1 + max_metrics
        if border:
            col_spec = "|".join(["l", *(["c"] * (total_cols - 1))])
        else:
            col_spec = "l" + "c" * (total_cols - 1)

        lines: list[str] = [
            rf"\begin{{tabular}}{{{col_spec}}}",
            r"\hline" if border else r"\toprule",
        ]

        for title, results in sections_data:
            lines.extend(
                self.__get_section_lines(title, results, total_cols, precision, sort_by, border)
            )

        # Replace the last section-ending rule with a final closing \hline.
        last_rule = r"\hline" if border else r"\midrule"
        final_rule = r"\hline"
        lines and lines[-1] == last_rule and lines.pop()
        (lines and lines[-1] == final_rule) or lines.append(final_rule)

        lines.append(r"\end{tabular}")
        table_lines: list[str] = [r"\begin{table}[htbp]", r"\centering"]

        if table_caption:
            table_lines.append(rf"\caption{{{self.__escape(table_caption)}}}")

        table_lines.extend(lines)
        table_lines.append(r"\end{table}")
        return "\n".join(table_lines) + "\n"

    def __get_section_lines(
        self,
        title: str,
        results: Mapping[str, Mapping[str, Any]],
        total_cols: int,
        precision: int,
        sort_by: list[str] | None,
        border: bool,
    ) -> list[str]:
        metrics = sorted({m for mm in results.values() for m in mm})
        sort_orders = sort_by or ["asc"]

        normalized_orders: list[str] = []
        for order in sort_orders:
            normalized = order.lower()
            if normalized not in ("asc", "des"):
                raise ValueError(f"Invalid sort_by value: {order}. Use 'asc' or 'des'.")
            normalized_orders.append(normalized)

        metric_sort: dict[str, str] = {}
        for idx, metric in enumerate(metrics):
            metric_sort[metric] = (
                normalized_orders[idx] if idx < len(normalized_orders) else normalized_orders[-1]
            )

        metric_bounds = collect_metric_bounds(results, metrics)

        best_by_metric: dict[str, float] = {}

        for metric in metrics:
            vals = [
                metric_value
                for model_metrics in results.values()
                for metric_name, metric_value in model_metrics.items()
                if metric_name == metric and isinstance(metric_value, (int, float))
            ]
            if vals:
                best_by_metric[metric] = (
                    min(vals) if metric_sort.get(metric, "asc") == "asc" else max(vals)
                )

        header_cells = ["Model", *[self.__escape(metric) for metric in metrics]]
        while len(header_cells) < total_cols:
            header_cells.append("")

        lines = [
            r"\addlinespace[3pt]",
            rf"\multicolumn{{{total_cols}}}{{c}}{{\textbf{{{self.__escape(title)}}}}} \\",
            r"\midrule",
            " & ".join(header_cells) + r" \\",
        ]

        for model_name in sorted(results):
            model_metrics = results[model_name]
            row = [self.__escape(model_name)]

            for metric in metrics:
                value = model_metrics.get(metric)
                if isinstance(value, (int, float)):
                    formatted = f"{value:.{precision}f}"
                    best = best_by_metric.get(metric)
                    if best is not None and value == best:
                        formatted = rf"\underline{{{formatted}}}"

                    row.append(
                        colorize_metric_value(
                            metric=metric,
                            value=float(value),
                            text=formatted,
                            metric_bounds=metric_bounds,
                            sort_order=metric_sort.get(metric, "asc"),
                        )
                    )
                else:
                    row.append("-")

            while len(row) < total_cols:
                row.append("")
            lines.append(" & ".join(row) + r" \\")

        lines.append(r"\hline" if border else r"\midrule")
        return lines

    def __escape(self, value: str) -> str:
        return (
            value.replace("\\", "\\textbackslash{}")
            .replace("&", "\\&")
            .replace("%", "\\%")
            .replace("$", "\\$")
            .replace("#", "\\#")
            .replace("_", "\\_")
            .replace("{", "\\{")
            .replace("}", "\\}")
            .replace("~", "\\textasciitilde{}")
            .replace("^", "\\textasciicircum{}")
        )

    def __save_comparison_tables(
        self,
        test_results: Mapping[str, Mapping[str, Any]],
        save_dir: str | Path,
        train_results: Mapping[str, Mapping[str, Any]] | None = None,
        val_results: Mapping[str, Mapping[str, Any]] | None = None,
        filename: str = "overall.tex",
        precision: int = 4,
        table_caption: str | None = None,
        sort_by: list[str] | None = None,
        border: bool = True,
    ) -> Path:
        sections_data: list[tuple[str, Mapping[str, Mapping[str, Any]]]] = []
        if test_results:
            sections_data.append(("Test Results", test_results))
        if train_results:
            sections_data.append(("Train Results", train_results))
        if val_results:
            sections_data.append(("Val Results", val_results))

        content = self.__build_comparison_table(
            sections_data=sections_data,
            precision=precision,
            border=border,
            table_caption=table_caption,
            sort_by=sort_by,
        )

        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)

        file_path = save_path / filename
        if content != "":
            content = (
                "% Requires: \\usepackage{booktabs}\n"
                "% Requires: \\usepackage[table]{xcolor}\n" + content
            )
        file_path.write_text(content)
        return file_path

class MarkdownTableLogger(Logger):
    """A Lightning Logger that accumulates metrics and writes a markdown comparison table.

    Multiple instances (one per model) share a class-level store keyed by experiment_name.
    Every time finalize() is called (after fit() or test() for each model), the current
    state of all accumulated metrics is written to a markdown file. The last model to
    finalize produces the most complete table.

    This means the file is progressively updated as models finish training/testing,
    so partial results are available while running.

    Args:
        save_dir: Base directory where the comparison/ subfolder will be created.
        model_name: The model's full name (e.g., "mlp:mean").
        experiment_name: Shared key that groups all models in the same experiment.
        precision: Decimal places for metric values in the table.
    """

    # Class-level shared store: {experiment_name: {model_name: {metric_name: value}}}
    __shared_stores: ClassVar[dict[str, dict[str, dict[str, float]]]] = {}

    def __init__(
        self,
        save_dir: str | Path,
        model_name: str,
        experiment_name: str,
        precision: int = 4,
    ) -> None:
        super().__init__()
        self.__save_dir = save_dir
        self.__model_name = model_name
        self.__experiment_name = experiment_name
        self.__precision = precision

        if experiment_name not in self.__shared_stores:
            self.__shared_stores[experiment_name] = {}

    @property
    def name(self) -> str:
        return "MarkdownTableLogger"

    @property
    def version(self) -> str | int:
        return self.__model_name

    @property
    def store(self) -> dict[str, dict[str, float]]:
        """Access the shared store for the current experiment."""
        return copy.deepcopy(self.__shared_stores.get(self.__experiment_name, {}))

    @property
    def save_dir(self) -> str | Path:
        return self.__save_dir

    @property
    def experiment_name(self) -> str | Path:
        return self.__experiment_name

    def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
        """Accumulate metrics for this model. Called by Lightning on every log step.

        Keeps only the latest value for each metric name. For example, if
        "val_auc" is logged at step 10 and step 20, only the step 20 value is kept.
        """
        store = self.__shared_stores[self.__experiment_name]
        if self.__model_name not in store:
            store[self.__model_name] = {}
        store[self.__model_name].update(metrics)

    def log_hyperparams(self, params: Any) -> None:
        pass

    def finalize(self, status: str) -> None:
        """Write the markdown comparison table with all accumulated metrics so far.

        Called by Lightning after fit() and after test() for each model. Since models
        train/test sequentially, each finalize() overwrites the file with all data
        accumulated up to that point. The file grows more complete over time.

        Args:
            status: The stage that just completed, e.g., "fit" or "test".
        """
        test_results, train_results, val_results = self.__split_results()

        if not test_results and not train_results and not val_results:
            return

        comparison_dir = Path(self.__save_dir) / "comparison"
        self.__save_comparison_tables(
            test_results=test_results,
            save_dir=comparison_dir,
            train_results=train_results or None,
            val_results=val_results or None,
            precision=self.__precision,
        )

    def __split_results(
        self,
    ) -> tuple[
        dict[str, dict[str, float]], dict[str, dict[str, float]], dict[str, dict[str, float]]
    ]:
        """Split all accumulated metrics into test vs train/val groups.

        Metrics are classified by their name prefix:
        - "test*"  --> test_results
        - "train*" --> train_results
        - "val*" --> val_results
        - anything else (e.g., "epoch") --> ignored

        Returns:
            Tuple of (test_results, train_results, val_results), where each is a dict
            mapping model names to their respective metric dicts. Models with no metrics
            in a category are excluded from that category's dict.
        """
        store = self.__shared_stores.get(self.__experiment_name, {})
        test_results: dict[str, dict[str, float]] = {}
        train_results: dict[str, dict[str, float]] = {}
        val_results: dict[str, dict[str, float]] = {}

        for model_name, metrics in store.items():
            test_metrics: dict[str, float] = {}
            train_metrics: dict[str, float] = {}
            val_metrics: dict[str, float] = {}

            for metric_name, value in metrics.items():
                if metric_name.startswith("test"):
                    test_metrics[metric_name] = value
                elif metric_name.startswith("train"):
                    train_metrics[metric_name] = value
                elif metric_name.startswith("val"):
                    val_metrics[metric_name] = value

            if test_metrics:
                test_results[model_name] = test_metrics
            if train_metrics:
                train_results[model_name] = train_metrics
            if val_metrics:
                val_results[model_name] = val_metrics

        return test_results, train_results, val_results

    def clear(self, experiment_name: str) -> None:
        """Remove accumulated data for an experiment.
        Args:
            experiment_name: The experiment name whose data should be cleared.

        """
        self.__shared_stores.pop(experiment_name, None)

    def __build_comparison_table(
        self,
        results: Mapping[str, Mapping[str, float]],
        precision: int = 4,
    ) -> str:
        """Build a markdown comparison table from model results.

        Examples:
            Input:

            ```python
            {
            "mlp:mean": {"test_auc": 0.85, "test_loss": 0.32},
            "gat:default": {"test_auc": 0.82},
            }
            ```

            Output:

            ```md
            | Model | test_auc | test_loss |
            | --- | --- | --- |
            | gat:default | 0.8200 | - |
            | mlp:mean | 0.8500 | 0.3200 |
            ```

        Args:
            results: Mapping of model names to metric dictionaries.
            precision: Number of decimal places for numeric metric values.

        Returns:
            Markdown table string. Returns an empty string if ``results`` is empty.


        """
        if not results:
            return ""

        # Collect all unique metric names across all models, sorted for determinism
        all_metrics = sorted(
            {metric for model_metrics in results.values() for metric in model_metrics}
        )

        # Build header row
        header = "| Model | " + " | ".join(all_metrics) + " |"
        separator = "| --- | " + " | ".join("---" for _ in all_metrics) + " |"

        # Build one row per model, sorted by model name for determinism
        rows = []
        for model_name in sorted(results):
            model_metrics = results[model_name]
            cells = []
            for metric in all_metrics:
                value = model_metrics.get(metric)
                if isinstance(value, (int, float)):
                    cells.append(f"{value:.{precision}f}")
                else:
                    cells.append("-")
            rows.append(f"| {model_name} | " + " | ".join(cells) + " |")

        return "\n".join([header, separator, *rows])

    def __save_comparison_tables(
        self,
        test_results: Mapping[str, Mapping[str, float]],
        save_dir: str | Path,
        train_results: Mapping[str, Mapping[str, float]] | None = None,
        val_results: Mapping[str, Mapping[str, float]] | None = None,
        filename: str = "results.md",
        precision: int = 4,
    ) -> Path:
        """Build and save markdown comparison tables to a file.

        Writes two sections:
        - "## Test Results" with the test metrics table
        - "## Train/Val Results" with the train/val metrics table (if provided)

        Args:
            test_results: Dict from test_all(), mapping model names to test metric dicts.
            save_dir: Directory where the markdown file will be written.
            train_results: Optional dict mapping model names to train metric dicts.
            val_results: Optional dict mapping model names to val metric dicts.
            filename: Name of the output file.
            precision: Decimal places for metric values.

        Returns:
            Path to the written file.
        """
        sections = []

        # Test results table
        test_table = self.__build_comparison_table(test_results, precision)
        if test_table:
            sections.append(f"## Test Results\n\n{test_table}")

        # Train/val results table
        if train_results or val_results:
            if train_results:
                train_table = self.__build_comparison_table(train_results, precision)
                sections.append(f"## Train Results\n\n{train_table}")
            if val_results:
                val_table = self.__build_comparison_table(val_results, precision)
                sections.append(f"## Val Results\n\n{val_table}")

        content = "\n\n".join(sections) + "\n" if sections else ""

        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)

        file_path = save_path / filename
        file_path.write_text(content)

        return file_path