import glob
from pathlib import Path
import shutil
from check import make_image


def get_exps(pattern: str, splitter: str = "_"):
    basedir = "/teamspace/jobs/"
    chkpt_basedir = "/work/colors/lightning_logs/"
    location = basedir + pattern
    res = glob.glob(location)
    location = location.replace('*', '')
    H = []  # hyperparams used
    # print(res)
    for r in res:
        d = r.replace(location, '').split(splitter)
        d = list(float(_d) for _d in d)
        d[0] = int(d[0])
        H.append(d)
    for i, r in enumerate(res):
        dir_path = Path(f"/teamspace/studios/this_studio/colors/lightning_logs/version_{i}/")
        dir_path.mkdir(parents=True, exist_ok=True)
        g = glob.glob(r + chkpt_basedir + "*")
        c = g[0] + "/checkpoints"
        latest_checkpoint = glob.glob(c + "/*")[-1]
        # print(latest_checkpoint)
        logs = glob.glob(g[0] + "/events*")[-1]
        print(logs)
        source_path = Path(logs)
        # print("Would copy", source_path, dir_path)
        # shutil.copy(source_path, dir_path)
        make_image(latest_checkpoint, f"out/version_{i}")
        make_image(latest_checkpoint, f"out/version_{i}b", color=False)

    return H


if __name__ == "__main__":
    D = get_exps("color_*", "_")
    print(len(D), D)