from run_selector import RunSelector
from plotter import Plotter, U
import os

path_from_home = "simus/ismfeed/allmode"
names = "n6_st_2e5_seed3_T5Myr_nsink1e3_comp*"


def prep_mcons(study):
    study.coarse_step_from_log()


def time_mcons(study, run, target=0.2):
    mcons = study.get_value(f"/series/coarse_step_from_log/mcons/{run}")
    time = study.get_value(f"/series/coarse_step_from_log/time/{run}")
    idx = -mcons[::-1].searchsorted(-target)
    time_target = time[idx] * study.info["unit_time"].express(U.Myr)
    return time_target


def find_nums(study, prep_function, time_function):
    nums = {}
    prep_function(study)
    for run in study.runs:
        time_target = time_function(study, run)
        rs = RunSelector(path_in=in_dir, in_runs=run, time=time_target, unit_time=U.Myr)
        nums.update(rs.nums)
    return nums


def write_paths(nums, path_from_home, filename="~/list_file"):
    paths = []
    for key in nums:
        for num in nums[key]:
            paths.append(f"{path_from_home}/{key}/output_{num:05}\n")
    f = open(os.path.expanduser(filename), "w")
    f.writelines(paths)
    f.close()


in_dir = os.path.expanduser(f"~/{path_from_home}")
study = Plotter(
    in_dir,
    filter_name=names,
    nums="first",
    tag="select",
).study


nums = find_nums(study, prep_mcons, time_mcons)
write_paths(nums, ".")

pl = Plotter(in_dir, runs=nums.keys(), nums=nums)

pl.rho_pdf()