import glob
import shutil
from pathlib import Path

from check import create_rectangle


def get_exps(pattern: str, splitter: str = "_", dry_run: bool = True):
    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 + "*")
        logs = glob.glob(g[0] + "/events*")[-1]
        source_path = Path(logs)
        print(logs)
        if not dry_run:
            c = g[0] + "/checkpoints"
            latest_checkpoint = glob.glob(c + "/*")[-1]
            print(latest_checkpoint)
        if not dry_run:
            shutil.copy(source_path, dir_path)
            create_rectangle(latest_checkpoint, f"out/version_{i}")
            # create_rectangle(latest_checkpoint, f"out/version_{i}b", color=False)
        else:
            print("Would copy", source_path, dir_path)
    return H


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

    import numpy as np

    D = np.array(D)
    # print(len(D), "\n", D)
    import pandas as pd

    df = pd.DataFrame(D)
    df.columns = ["batch_size", "alpha", "learning_rate"]
    df.to_csv("experiments.csv")
    print(df)