#!/usr/bin/env python3
"""Generate JSON index for experiment viewer.

Scans the figures directory, parses each filename into structured metadata
(model, component, plot_type, metric) following the convention in
docs/FIGURE_NAMING.md, and writes a JSON index that viewer.html consumes.

Usage:
    python generate_viewer_index.py --out corr_out
    python generate_viewer_index.py --out corr_out --serve
"""

import argparse
import json
import os
import re
from pathlib import Path
import http.server
import socketserver
import webbrowser
from threading import Timer


# ── Filename parsing ──────────────────────────────────────────────────

COMPONENTS = ["W_QK", "W_OV", "b_Q", "b_K", "b_V", "b_O"]
_COMP_RE = "|".join(re.escape(c) for c in COMPONENTS)

PLOT_TYPES = [
    "corr_vs_layer_distance", "dominant_eigenvectors", "MP_overlay",
    "Q_heatmap", "Q_eigenvalues", "block_means", "P_Q",
    "scalars_vs_layer", "scalars_heatmap", "scalar_correlations",
]

METRICS = [
    "hist_symmetric_kl", "hist_jensen_shannon",
    "symmetric_kl", "jensen_shannon",
    "frob_cosine", "pearson_corr", "two_point", "connected_corr",
]

# Display labels for the viewer
PLOT_TYPE_DISPLAY = {
    "Q_heatmap":                "Correlation heatmap",
    "block_means":              "Block means",
    "P_Q":                      "Overlap distribution P(Q)",
    "Q_eigenvalues":            "Eigenvalue spectrum",
    "corr_vs_layer_distance":   "Correlation vs distance",
    "MP_overlay":               "Marchenko-Pastur",
    "dominant_eigenvectors":    "Dominant eigenvectors",
    "cross_heatmap":            "Cross-correlation heatmap",
    "cross_diagonal":           "Cross-correlation diagonal",
    "scalars_vs_layer":         "Scalars vs layer",
    "scalars_heatmap":          "Scalar heatmaps",
    "scalar_correlations":      "Scalar correlations",
}

METRIC_DISPLAY = {
    "frob_cosine":          "Frobenius cosine",
    "pearson_corr":         "Pearson correlation",
    "two_point":            "Two-point function",
    "connected_corr":       "Connected correlation",
    "hist_symmetric_kl":    "Symmetric KL (hist)",
    "hist_jensen_shannon":  "Jensen-Shannon (hist)",
    "symmetric_kl":         "Symmetric KL (KDE)",
    "jensen_shannon":       "Jensen-Shannon (KDE)",
}

# Compiled regexes for cross-correlation filenames
_CROSS_DIAG_RE = re.compile(
    rf"^(.+?)_cross_diagonal_({_COMP_RE})_vs_({_COMP_RE})$"
)
_CROSS_RE = re.compile(
    rf"^(.+?)_cross_({_COMP_RE})_vs_({_COMP_RE})_(.+)$"
)
# Standard self-correlation: {model}_{component}_{plot_type}[_{metric}]
_SELF_RE = re.compile(
    rf"^(.+?)_({_COMP_RE})_(.+)$"
)


def parse_filename(filename):
    """Parse a figure filename into structured metadata.

    Returns dict with: model, component, plot_type, metric, cross_pair.
    """
    name = filename.removesuffix(".png")
    result = {
        "model": "", "component": "", "plot_type": "", "metric": "",
        "cross_pair": "",
    }

    # 1. Cross-diagonal: {model}_cross_diagonal_{compA}_vs_{compB}
    m = _CROSS_DIAG_RE.match(name)
    if m:
        result["model"] = m.group(1)
        result["component"] = f"{m.group(2)} vs {m.group(3)}"
        result["cross_pair"] = result["component"]
        result["plot_type"] = "cross_diagonal"
        return result

    # 2. Cross-heatmap: {model}_cross_{compA}_vs_{compB}_{metric}
    m = _CROSS_RE.match(name)
    if m:
        result["model"] = m.group(1)
        result["component"] = f"{m.group(2)} vs {m.group(3)}"
        result["cross_pair"] = result["component"]
        result["plot_type"] = "cross_heatmap"
        result["metric"] = m.group(4)
        return result

    # 3. Scalars: {model}_scalars_vs_layer / _scalars_heatmap / _scalar_correlations
    for pt in ("scalars_vs_layer", "scalars_heatmap", "scalar_correlations"):
        tag = f"_{pt}"
        if name.endswith(tag):
            result["model"] = name[:-len(tag)]
            result["component"] = "scalars"
            result["plot_type"] = pt
            return result

    # 4. Standard: {model}_{component}_{plot_type}[_{metric}]
    m = _SELF_RE.match(name)
    if m:
        result["model"] = m.group(1)
        result["component"] = m.group(2)
        rest = m.group(3)

        # Match known plot types (longest first is handled by ordering)
        for pt in PLOT_TYPES:
            if rest == pt:
                result["plot_type"] = pt
                return result
            if rest.startswith(pt + "_"):
                result["plot_type"] = pt
                result["metric"] = rest[len(pt) + 1:]
                return result

        # Fallback: entire rest is plot_type
        result["plot_type"] = rest
        return result

    # 5. Total fallback
    result["model"] = name
    return result


# ── Index generation ──────────────────────────────────────────────────

def find_images(figures_dir):
    figures_path = Path(figures_dir)
    if not figures_path.exists():
        print(f"Warning: Figures directory not found: {figures_dir}")
        return []

    images = []
    for img_path in sorted(figures_path.glob("*.png")):
        stat = img_path.stat()
        meta = parse_filename(img_path.name)
        images.append({
            "filename": img_path.name,
            "mtime": stat.st_mtime,
            "size": stat.st_size,
            **meta,
        })
    return images


def generate_index(output_dir, figures_subdir="figures", relative_paths=True):
    output_path = Path(output_dir)
    figures_dir = output_path / figures_subdir
    images = find_images(figures_dir)

    base_path = figures_subdir if relative_paths else str(figures_dir.absolute())

    # Collect unique values for each filter dimension
    models = sorted({img["model"] for img in images if img["model"]})
    components = sorted({img["component"] for img in images if img["component"]})
    plot_types = sorted({img["plot_type"] for img in images if img["plot_type"]})
    metrics = sorted({img["metric"] for img in images if img["metric"]})

    index_data = {
        "basePath": base_path,
        "imageCount": len(images),
        "images": images,
        "filters": {
            "models": models,
            "components": components,
            "plot_types": plot_types,
            "metrics": metrics,
        },
        "display": {
            "plot_types": PLOT_TYPE_DISPLAY,
            "metrics": METRIC_DISPLAY,
        },
    }

    index_path = output_path / "image_index.json"
    with open(index_path, "w") as f:
        json.dump(index_data, f, indent=2)

    print(f"Generated index: {len(images)} images")
    print(f"  Models:     {models}")
    print(f"  Components: {components}")
    print(f"  Plot types: {plot_types}")
    print(f"  Metrics:    {metrics}")
    print(f"  Saved to:   {index_path}")

    return index_path, len(images)


def copy_viewer(output_dir, script_dir, force=False):
    import shutil
    viewer_source = Path(script_dir) / "viewer.html"
    viewer_dest = Path(output_dir) / "viewer.html"
    if not viewer_source.exists():
        print(f"Warning: viewer.html not found at {viewer_source}")
        return None
    if not viewer_dest.exists() or force:
        shutil.copy(viewer_source, viewer_dest)
        print(f"Copied viewer.html to {viewer_dest}")
    else:
        print(f"viewer.html already exists (use --force-copy to update)")
    return viewer_dest


def serve_viewer(output_dir, port=8000):
    os.chdir(output_dir)
    Handler = http.server.SimpleHTTPRequestHandler

    def open_browser():
        webbrowser.open(f"http://localhost:{port}/viewer.html")

    Timer(1.5, open_browser).start()

    with socketserver.TCPServer(("", port), Handler) as httpd:
        print(f"\nServing viewer at: http://localhost:{port}/viewer.html")
        print(f"Press Ctrl+C to stop\n")
        try:
            httpd.serve_forever()
        except KeyboardInterrupt:
            print("\nShutting down server...")
            httpd.shutdown()


# ── CLI ───────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Generate index for experiment viewer")
    parser.add_argument("--out", type=str, default="corr_out",
                        help="Output directory containing experiment results")
    parser.add_argument("--figures-dir", type=str, default="figures",
                        help="Subdirectory containing figures (default: figures)")
    parser.add_argument("--serve", action="store_true",
                        help="Start HTTP server to view results")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--absolute-paths", action="store_true",
                        help="Use absolute paths in index (for file:// viewing)")
    parser.add_argument("--no-copy-viewer", action="store_true")
    parser.add_argument("--force-copy", action="store_true",
                        help="Force update viewer.html")
    args = parser.parse_args()

    script_dir = Path(__file__).parent

    print(f"Scanning: {args.out}/{args.figures_dir}")
    index_path, image_count = generate_index(
        args.out, args.figures_dir,
        relative_paths=not args.absolute_paths,
    )

    if image_count == 0:
        print(f"\nNo images found in {args.out}/{args.figures_dir}")
        return

    if not args.no_copy_viewer:
        copy_viewer(args.out, script_dir, force=args.force_copy)

    if args.serve:
        serve_viewer(args.out, args.port)
    else:
        print(f"\nTo view: run with --serve, or:")
        print(f"  cd {args.out} && python -m http.server {args.port}")
        print(f"  open http://localhost:{args.port}/viewer.html")


if __name__ == "__main__":
    main()