"""
Experiment result tracking and comparison table generation.

Logs results from all tokenizer × domain × model × mode × split combinations
and generates CSV / LaTeX tables for the paper.
"""

import csv
import json
import os
from collections import defaultdict


class ExperimentTracker:
    """
    Track and store experimental results across all configurations.

    Results are stored in a nested dict:
        results[domain][tokenizer][model][mode][split] = metrics_dict
    """

    def __init__(self, output_dir: str = "results"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        self.results: dict = defaultdict(
            lambda: defaultdict(
                lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))
            )
        )

    def log_result(
        self,
        domain: str,
        tokenizer: str,
        model: str,
        mode: str,
        split: str,
        metrics: dict,
    ) -> None:
        """
        Log a single experimental result.

        Args:
            domain: e.g. 'blocks', 'gripper'
            tokenizer: e.g. 'wl', 'simhash', 'shortest_path', 'graphbpe', 'random'
            model: e.g. 'lstm', 'xgboost'
            mode: 'state' or 'delta'
            split: 'validation', 'test-interpolation', 'test-extrapolation'
            metrics: dict with keys like 'solved_rate', 'exec_rate', 'vocab_size'
        """
        self.results[domain][tokenizer][model][mode][split] = metrics
        print(
            f"  Logged: {domain}/{tokenizer}/{model}/{mode}/{split} → "
            f"{metrics.get('solved_rate', 'N/A')}"
        )

    def generate_comparison_table(self) -> str:
        """
        Generate a comparison table across all methods.

        Returns: CSV string
        """
        rows = []
        header = [
            "Domain",
            "Tokenizer",
            "Model",
            "Mode",
            "Val Solved",
            "Interp Solved",
            "Extrap Solved",
            "Val Exec",
            "Interp Exec",
            "Extrap Exec",
            "Vocab Size",
        ]
        rows.append(header)

        for domain in sorted(self.results.keys()):
            for tokenizer in sorted(self.results[domain].keys()):
                for model in sorted(self.results[domain][tokenizer].keys()):
                    for mode in sorted(
                        self.results[domain][tokenizer][model].keys()
                    ):
                        splits = self.results[domain][tokenizer][model][mode]
                        row = [domain, tokenizer, model, mode]

                        for split in [
                            "validation",
                            "test-interpolation",
                            "test-extrapolation",
                        ]:
                            m = splits.get(split, {})
                            row.append(f"{m.get('solved_rate', 0):.2%}")

                        for split in [
                            "validation",
                            "test-interpolation",
                            "test-extrapolation",
                        ]:
                            m = splits.get(split, {})
                            row.append(f"{m.get('exec_rate', 0):.2%}")

                        # Vocab size (same for all splits)
                        any_m = next(iter(splits.values()), {})
                        row.append(str(any_m.get("vocab_size", "N/A")))

                        rows.append(row)

        # Write CSV
        csv_path = os.path.join(self.output_dir, "tokenization_comparison.csv")
        with open(csv_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerows(rows)
        print(f"Saved comparison table to {csv_path}")

        # Generate LaTeX
        self._generate_latex_table(rows)

        return csv_path

    def _generate_latex_table(self, rows: list[list]) -> None:
        """Generate LaTeX table from rows."""
        header = rows[0]
        data = rows[1:]

        n_cols = len(header)
        col_spec = "l" * n_cols

        lines = [
            r"\begin{table}[ht]",
            r"\centering",
            r"\caption{Comparison of tokenization strategies}",
            r"\label{tab:tokenization_comparison}",
            f"\\begin{{tabular}}{{{col_spec}}}",
            r"\toprule",
            " & ".join(header) + r" \\",
            r"\midrule",
        ]

        for row in data:
            lines.append(" & ".join(str(x) for x in row) + r" \\")

        lines.extend(
            [
                r"\bottomrule",
                r"\end{tabular}",
                r"\end{table}",
            ]
        )

        tex_path = os.path.join(self.output_dir, "tokenization_comparison.tex")
        with open(tex_path, "w") as f:
            f.write("\n".join(lines))
        print(f"Saved LaTeX table to {tex_path}")

    def save_results(self, filepath: str | None = None) -> None:
        """Save all results to JSON."""
        if filepath is None:
            filepath = os.path.join(self.output_dir, "full_results.json")

        # Convert defaultdicts to regular dicts for JSON serialization
        def to_dict(d):
            if isinstance(d, defaultdict):
                return {k: to_dict(v) for k, v in d.items()}
            return d

        with open(filepath, "w") as f:
            json.dump(to_dict(self.results), f, indent=2)
        print(f"Saved full results to {filepath}")

    def load_results(self, filepath: str) -> None:
        """Load results from JSON."""
        with open(filepath, "r") as f:
            data = json.load(f)

        # Populate results
        for domain, tok_data in data.items():
            for tok, model_data in tok_data.items():
                for model, mode_data in model_data.items():
                    for mode, split_data in mode_data.items():
                        for split, metrics in split_data.items():
                            self.results[domain][tok][model][mode][split] = metrics

    def generate_vocabulary_stats(self) -> None:
        """Generate vocabulary size stats per tokenizer per domain."""
        stats = {}
        for domain in self.results:
            for tok in self.results[domain]:
                for model in self.results[domain][tok]:
                    for mode in self.results[domain][tok][model]:
                        for split, m in self.results[domain][tok][model][
                            mode
                        ].items():
                            if "vocab_size" in m:
                                stats.setdefault(tok, {})[domain] = m["vocab_size"]

        stats_path = os.path.join(self.output_dir, "vocabulary_stats.json")
        with open(stats_path, "w") as f:
            json.dump(stats, f, indent=2)
        print(f"Saved vocabulary stats to {stats_path}")