import os
from glob import glob
from tqdm import tqdm
from slic.utils import json_load, json_save


# select pgroup here
pgroup = "p21623"

# set the output folder here
output = "./new_jsons"

# select run range here
run_min = 40
run_max = 69

# add more channels here, if needed
new_bs_chans = []
new_cam_chans =[
    "SATES21-CAMS-PATT1:FPICTURE"
]

os.makedirs(output, exist_ok=True)

pattern = f"/sf/maloja/data/{pgroup}/raw/run0*"
runs = sorted(glob(pattern))
for run in tqdm(runs):
    nr = run.split("/")[-1][3:7]
    nr = int(nr)
    if not (run_min <= nr < run_max):
        continue

    acqs = sorted(glob(f"{run}/meta/acq*.json"))
    for acq in tqdm(acqs):
        data = json_load(acq)

        # update the channels:
        bs_chans  = data.get("channels_list", [])
        cam_chans = data.get("camera_list", [])

        bs_chans.extend(new_bs_chans)
        cam_chans.extend(new_cam_chans)

        data["channels_list"] = bs_chans
        data["camera_list"]   = cam_chans

        na = acq.split("/")[-1][3:7]
        na = int(na)

        new = f"{output}/run{nr:04}_acq{na:04}.json"
        json_save(data, new)