"""
Discovers pre-exported steering-trajectory frame PNGs (one per alpha value)
and the corresponding assembled GIF, per case:
  assets/trajectory_frames/<case>/frame_NNN_alphaX.XX.png
  assets/trajectory_gifs/<case>.gif
Both produced by the physics_gif() function in
steering_paper_v2/notebooks/v3/02_paper_results.ipynb. No Gradio dependency.
"""

import os
import re

ASSETS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
FRAMES_DIR = os.path.join(ASSETS_DIR, "trajectory_frames")
GIFS_DIR = os.path.join(ASSETS_DIR, "trajectory_gifs")

_FRAME_RE = re.compile(r"alpha([\d.]+)\.png$")

# case key -> human-readable label (from PHYSICS_GIF_LABELS in
# 02_paper_results.ipynb, LaTeX stripped down to plain text since Gradio
# Dropdown choices don't render LaTeX).
CASE_LABELS = {
    "s2r": "phi -> psi_R (scalar to right-handed fermion)",
    "s2l": "phi -> psi_L (scalar to left-handed fermion)",
    "r2l": "psi_R -> psi_L",
    "su3_1to3": "SU(3): 1 -> 3",
    "su3_1to3bar": "SU(3): 1 -> 3-bar",
    "su3_3_to_3bar": "SU(3): 3 -> 3-bar",
    "su2_1to2": "SU(2): 1 -> 2",
    "su2_1to3": "SU(2): 1 -> 3",
    "su2_1toany": "SU(2) extension (rep. size ambiguous)",
    "modintq_div2": "U(1): Q_int -> Q_int / 2",
    "u1_b2c": "Enforce U(1) charge conservation",
    "susy_all": "SUSY: 1-3 fields (unrestricted)",
    "susy_1f": "SUSY: 1 field",
    "susy_2f": "SUSY: 2 fields",
    "add_s2r_r2l_e1": "(phi -> psi_R) + (psi_R -> psi_L), additive",
    "susy_1f_s": "SUSY: 1 scalar field (scalar source)",
    "acc_rel": "Gauge anomaly-free: by field extension",
    "acc_abs": "Gauge anomaly-free: by field extension (absolute scaling)",
    "acc_3f": "Gauge anomaly-free: by delta-Q",
    "wit_uniq": "Witten anomaly-free",
}


# Cases present on disk (assets/trajectory_frames/<case>/) but deliberately
# left out of the dropdown — data stays put, just not shown.
EXCLUDED_CASES = {"acc_abs"}

# Dropdown order, matching HEADLINE_CASES in 02_paper_results.ipynb (the
# table behind the paper's combined headline plots) — grouped by family:
# Lorentz -> SU(3) -> SU(2) -> U(1) -> Theory extension -> Anomaly. Any
# case not listed here (e.g. a new one added before this list is updated)
# falls back to alphabetical, appended after the ordered ones.
CASE_ORDER = [
    "s2r", "s2l", "r2l", "add_s2r_r2l_e1",
    "su3_1to3", "su3_1to3bar", "su3_3_to_3bar",
    "su2_1to2", "su2_1to3",
    "modintq_div2", "u1_b2c",
    "susy_1f", "susy_1f_s", "susy_2f", "susy_all", "su2_1toany",
    "acc_rel", "acc_3f", "wit_uniq",
]


def _discover_cases():
    if not os.path.isdir(FRAMES_DIR):
        return []
    found = {
        d for d in os.listdir(FRAMES_DIR)
        if os.path.isdir(os.path.join(FRAMES_DIR, d)) and d not in EXCLUDED_CASES
    }
    ordered = [c for c in CASE_ORDER if c in found]
    leftover = sorted(found - set(CASE_ORDER))
    return ordered + leftover


def _discover_frames(case_dir):
    alphas_to_path = {}
    for fname in os.listdir(case_dir):
        m = _FRAME_RE.search(fname)
        if m:
            alphas_to_path[float(m.group(1))] = os.path.join(case_dir, fname)
    if not alphas_to_path:
        raise ValueError(f"No frame_*_alpha*.png files found in {case_dir}")
    return dict(sorted(alphas_to_path.items()))


CASES = _discover_cases()
# case key -> {alpha: frame_path}, sorted by alpha
CASE_FRAMES = {case: _discover_frames(os.path.join(FRAMES_DIR, case)) for case in CASES}
# case key -> [(alpha, path), ...] sorted by alpha, for index-based access —
# most cases have a uniform alpha step, but at least one (acc_abs) doesn't,
# so the slider indexes into this list by position rather than by alpha
# value directly.
CASE_FRAME_LIST = {case: sorted(frames.items()) for case, frames in CASE_FRAMES.items()}
# dropdown choices as (label, value) pairs, falling back to the raw key
CASE_CHOICES = [(CASE_LABELS.get(case, case), case) for case in CASES]


def n_frames(case):
    return len(CASE_FRAME_LIST[case])


def frame_by_index(case, index):
    frames = CASE_FRAME_LIST[case]
    index = max(0, min(index, len(frames) - 1))
    alpha, path = frames[index]
    return alpha, path


def gif_path(case):
    path = os.path.join(GIFS_DIR, f"{case}.gif")
    if not os.path.isfile(path):
        raise ValueError(f"No assembled GIF found for case {case!r} at {path}")
    return path