# -*- mode: python-mode; python-indent-offset: 4 -*-
# coding: utf-8

"""
  Select snaphots with a criterion
"""

from utils.runselector import RunSelector
from plotter import Plotter, U
import os

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, time_min=0):
    """  
    Once other filter are applied, select one output based on the time given by time function

    Args:
        prep_function (study:studyProcessor -> None): prepare a study object
        time_function (study:studyProcessor, run:str -> time:float): compute selected time from the object 
    """
    nums = {}
    prep_function(study)
    for run in study.runs:
        time_target = max(time_min, time_function(study, run))
        rs = RunSelector(path_in=study.path, in_runs=run, time_min=time_min, 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()


if __name__ == '__main__':
    path_from_home = "simus/ismfeed/allmode"
    names = "n6_st_2e5_seed3_T5Myr_nsink1e3_comp*"

    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, ".")