if __name__=='__main__':
    import sys, os, yaml, re, pickle
    # print()
    print(os.path.dirname(os.path.abspath(__file__)))
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from mi_boiler.sigma_conf import load_master_config, getenv_roundup
    from mi_boiler.compression_tools import iterate_over_compact_chunks
    from tqdm import tqdm
    import numpy as np
    import torch
    from collections import defaultdict
    from .ptrue import PTrueEstimator

    # can configure the run here if needed
    conf, printable_conf, oconf = load_master_config('uncertainty/uncertainties_computation.yaml', abspath=True, return_original_conf=True)
    # unpack into globals
    globals().update(conf)
    # print the used environment variables
    print(getenv_roundup())
    print(printable_conf)

    # get all the directories
    dirs = sorted([os.path.join(prefix_dir, d) for d in os.listdir(prefix_dir) if re.fullmatch(subdir_select_regex, d)])
    print(dirs)

    # select ones that are in the current grid
    dirs = [d for i,d in enumerate(dirs) if (i%grid_size)==my_grid_id]
    print(dirs)

    # set the name for the range stretch that will be saved
    range_suffix = f'.{starting_range:05d}_{ending_range:05d}_' if starting_range>0 or ending_range!=-1 else ''

    # now can go over every each applicable dir
    for rdir in dirs:
        results = defaultdict(list) # {k: [] for k in partial_uncert_fns.keys()}

        for records in iterate_over_compact_chunks(
            rdir,
            range_start=starting_range, 
            range_end=ending_range, 
            restrict_sets=None,
            pigz_fast_temp_path=pigz_temp_path
        ):
            # small type conversion
            print(f"Normalizing the chunk: ")
            for subset in ['ms', 'bs']:
                for r in records[subset]:
                    for k in r.keys():
                        if k.startswith('txt') and isinstance(r[k], np.ndarray):
                            r[k] = list(r[k])
                    for k, v in r.items():
                        if isinstance(v, np.ndarray):
                            r[k] = torch.from_numpy(r[k])
            
            # compute corresponding uncertainty functions
            for uname, uncert_fn in partial_uncert_fns.items():
                if overwrite_existing or not os.path.exists(os.path.join(rdir, 'uncert_'+uname+'.pkl')):
                    print(f"Computing : {uname}")
                    if isinstance(uncert_fn, PTrueEstimator) and "SPAR" in rdir:
                        print("Ptrue not implemented for vlms! Skipping..")
                        continue
                    uncert_res = uncert_fn.compute_for_records(records)
                    results[uname].extend(uncert_res)

        # save the results
        for name, res in results.items():
            with open(os.path.join(rdir, f'{range_suffix}uncert_{name}.pkl'), 'wb') as f:
                pickle.dump(res, f)
        del results
