import os
import json
import re
from collections import defaultdict

# Absolute path to the image directory
base_dir = "/uufs/chpc.utah.edu/common/home/horel-group9/horel/uunet/uunet_figs/climo_figs/"

stations = set()
types = set()
variables = defaultdict(set)

pattern = re.compile(r"(.*?_set_\d+)_(.+)_months\.png$")

for fname in os.listdir(base_dir):
    if not fname.endswith(".png"):
        continue

    if "_" not in fname or "_months.png" not in fname:
        continue

    station_and_rest = fname.split("_", 1)
    station = station_and_rest[0]

    match = pattern.search(fname)
    if not match:
        continue  # skip if not matching expected format

    variable = match.group(1).replace(f"{station}_", "")
    typ = match.group(2)

    stations.add(station)
    types.add(typ)
    variables[station].add(variable)

# Output JSON
result = {
    "stations": sorted(stations),
    "types": sorted(types),
    "variables": {k: sorted(v) for k, v in variables.items()}
}

output_path = os.path.join(base_dir, "images_index.json")
with open(output_path, "w") as f:
    json.dump(result, f, indent=2)

print("✅ images_index.json created in", output_path)