1# coding: utf-8
2"""
3Collection of low-level tools that faciliate the interface with resource managers.
4
5The preferred way of importing this module is:
6
7    import qutils as qu
8"""
9from monty.string import is_string
10from pymatgen.core.units import Time, Memory
11from abipy.tools import duck
12
13import logging
14logger = logging.getLogger(__name__)
15
16
17def slurm_parse_timestr(s):
18    """
19    A slurm time parser. Accepts a string in one the following forms:
20
21        # "days-hours",
22        # "days-hours:minutes",
23        # "days-hours:minutes:seconds".
24        # "minutes",
25        # "minutes:seconds",
26        # "hours:minutes:seconds",
27
28    Returns: Time in seconds.
29
30    Raises:
31        `ValueError` if string is not valid.
32    """
33    days, hours, minutes, seconds = 0, 0, 0, 0
34
35    if duck.is_number_like(s):
36        return Time(s, "s")
37
38    if '-' in s:
39        # "days-hours",
40        # "days-hours:minutes",
41        # "days-hours:minutes:seconds".
42        days, s = s.split("-")
43        days = int(days)
44
45        if ':' not in s:
46            hours = int(float(s))
47        elif s.count(':') == 1:
48            hours, minutes = map(int, s.split(':'))
49        elif s.count(':') == 2:
50            hours, minutes, seconds = map(int, s.split(':'))
51        else:
52            raise ValueError("More that 2 ':' in string!")
53
54    else:
55        # "minutes",
56        # "minutes:seconds",
57        # "hours:minutes:seconds",
58        if ':' not in s:
59            minutes = int(float(s))
60        elif s.count(':') == 1:
61            minutes, seconds = map(int, s.split(':'))
62        elif s.count(':') == 2:
63            hours, minutes, seconds = map(int, s.split(':'))
64        else:
65            raise ValueError("More than 2 ':' in string!")
66
67    return Time((days*24 + hours)*3600 + minutes*60 + seconds, "s")
68
69
70def time2slurm(timeval, unit="s"):
71    """
72    Convert a number representing a time value in the given unit (Default: seconds)
73    to a string following the slurm convention: "days-hours:minutes:seconds".
74
75    >>> assert time2slurm(61) == '0-0:1:1' and time2slurm(60*60+1) == '0-1:0:1'
76    >>> assert time2slurm(0.5, unit="h") == '0-0:30:0'
77    """
78    d, h, m, s = 24*3600, 3600, 60, 1
79
80    timeval = Time(timeval, unit).to("s")
81    days, hours = divmod(timeval, d)
82    hours, minutes = divmod(hours, h)
83    minutes, secs = divmod(minutes, m)
84
85    return "%d-%d:%d:%d" % (days, hours, minutes, secs)
86
87
88def time2pbspro(timeval, unit="s"):
89    """
90    Convert a number representing a time value in the given unit (Default: seconds)
91    to a string following the PbsPro convention: "hours:minutes:seconds".
92
93    >>> assert time2pbspro(2, unit="d") == '48:0:0'
94    """
95    h, m, s = 3600, 60, 1
96
97    timeval = Time(timeval, unit).to("s")
98    hours, minutes = divmod(timeval, h)
99    minutes, secs = divmod(minutes, m)
100
101    return "%d:%d:%d" % (hours, minutes, secs)
102
103
104def time2loadlever(timeval, unit="s"):
105    """
106    Convert a number representing a time value in the given unit (Default: seconds)
107    to a string following the LoadLever convention. format hh:mm:ss (hours:minutes:seconds)
108
109    >>> assert time2loadlever(2, unit="d") == '48:00:00'
110    """
111    h, m, s = 3600, 60, 1
112
113    timeval = Time(timeval, unit).to("s")
114    hours, minutes = divmod(timeval, h)
115    minutes, secs = divmod(minutes, m)
116
117    return "%d:%02d:%02d" % (hours, minutes, secs)
118
119
120def timelimit_parser(s):
121    """Convert a float or a string into time in seconds."""
122    try:
123        return Time(float(s), "s")
124    except ValueError:
125        return slurm_parse_timestr(s)
126
127
128def any2mb(s):
129    """Convert string or number to memory in megabytes."""
130    if is_string(s):
131        return int(Memory.from_string(s).to("Mb"))
132    else:
133        return int(s)
134