# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
"""
This module is for color formats, units and whatnot. This is exposed
automatically with::

   from k1lib.imports import *
   fmt.txt # exposed
"""
import k1lib, math, re; from k1lib import cli
from typing import Dict, Iterator, Tuple
pygments = k1lib.dep("pygments")
__all__ = ["generic", "metricPrefixes", "size", "fromSize", "sizeOf",
           "comp", "compRate", "time", "item", "throughput", "txt", "code", "h", "pre", "row", "col", "colors", "rmAnsi"]
k1lib.settings.add("fmt", k1lib.Settings().add("separator", True, "whether to have a space between the number and the unit"), "from k1lib.fmt module");
settings = k1lib.settings.fmt
metricPrefixes = {-8:"y",-7:"z",-6:"a",-5:"f",-4:"p",-3:"n",-2:"u",-1:"m",0:"",1:"k",2:"M",3:"G",4:"T",5:"P",6:"E",7:"Z",8:"Y"}
#metricPrefixes = ["", "k", "M", "G", "T", "P", "E", "Z", "Y"]
def generic(x, units:Dict[int, str]):                                            # generic
    c = " " if settings.separator else ""                                        # generic
    for i, unit in units.items():                                                # generic
        upperBound = 1000 * 1000**i                                              # generic
        if abs(x) < upperBound:                                                  # generic
            return f"{round(1e3*x/upperBound, 2)}{c}{unit}"                      # generic
    return (f"{round(1e3*x/upperBound, 2)}{c}{unit}").strip()                    # generic
sizes = {i: f"{p}B" for i, p in metricPrefixes.items() if i >= 0}; #sizes[0] = "bytes" # generic
def size(_bytes=0):                                                              # size
    """Formats disk size.
Example::

    # returns "50.0 bytes"
    fmt.size(50)
    # returns "12.0 MB"
    fmt.size(1.2e7)
"""                                                                              # size
    return generic(_bytes, sizes)                                                # size
sizeInv = {"k": 1e3, "m": 1e6, "g": 1e9, "t": 1e12, "p": 1e15, "e": 1e18, "z": 1e21, "y": 1e24} # size
def fromSize(s:str) -> int: # this is ugly I know, doesn't fit well into others. But making a generalized version seems hard, and I just need this right now # fromSize
    """Grabs size from string representation.
Example::

    fromSize("31.5k") # returns 31500
    fromSize("31.5kB") # also returns 31500
"""                                                                              # fromSize
    s = s.lower().replace(" ", "").rstrip("b"); ch = s[-1]                       # fromSize
    if ch in sizeInv: return int(float(s[:-1])*sizeInv[ch])                      # fromSize
    return int(s)                                                                # fromSize
def sizeOf(l:Iterator[float]) -> Tuple[str, Iterator[float]]:                    # sizeOf
    """Figures out appropriate scale, scales back the Iterator, and return both.
Example::

    x = torch.abs(torch.randn(2)) * 1e4 + 1e5
    label, t = fmt.sizeOf(x) # label is "kB"
    (t | toTensor()).min() # min value should be close to 100"""                 # sizeOf
    l = list(l | cli.apply(lambda n: abs(n)))                                    # sizeOf
    v = l | cli.toMax()                                                          # sizeOf
    v = math.log10(v) if v > 0 else -math.log10(-v)                              # sizeOf
    idx = math.floor(v/3)                                                        # sizeOf
    coef = 1.0/1000**idx                                                         # sizeOf
    return sizes[idx], l | cli.apply(lambda x: x * coef) | cli.deref()           # sizeOf
computations = {i: f"{p}FLOPs" for i, p in metricPrefixes.items() if i >= 0}     # sizeOf
def comp(flop=0):                                                                # comp
    """Formats computation amount.
Example::

    # returns "50.0 FLOPs"
    fmt.computation(50)
    # returns "50.0 MFLOPs"
    fmt.computation(5e7)
"""                                                                              # comp
    return generic(flop, computations)                                           # comp
computationRates = {i: f"{p}FLOPS" for i, p in metricPrefixes.items() if i >= 0} # comp
def compRate(flops=0):                                                           # compRate
    """Formats computation rate.
Example::

    # returns "50.0 FLOPS"
    fmt.computationRate(50)
    # returns "50.0 MFLOPS"
    fmt.computationRate(5e7)
"""                                                                              # compRate
    return generic(flops, computationRates)                                      # compRate
times = {i:f"{p}s" for i, p in metricPrefixes.items() if i <= 0}                 # compRate
def time(seconds=0):                                                             # time
    """Formats small times.
Example::

    fmt.time(50) # returns "50.0 s"
    fmt.time(4000) # returns "4000.0 s"
    fmt.time(0.02) # returns "20.0 ms"
    fmt.time(1e-5) # returns "10.0 us"
"""                                                                              # time
    return generic(seconds, times)                                               # time
items = {0: "", 1: "k", 2: "M", 3: "B", 4: "T"}                                  # time
def item(n=0):                                                                   # item
    """Formats generic item.
Example::

    # returns "50.0"
    fmt.item(50)
    # returns "500.0 k"
    fmt.item(5e5)
"""                                                                              # item
    return generic(n, items)                                                     # item
def throughput(n, unit=""):                                                      # throughput
    """Formats item throughput.
Example::

    # returns "3.16/year"
    fmt.throughput(1e-7)
    # returns "2.63/month"
    fmt.throughput(1e-6)
    # returns "3.6/hour"
    fmt.throughput(1e-3)
    # returns "100.0 k/s"
    throughput(1e5)

    # returns "100.0 k epochs/s"
    throughput(1e5, " epochs")

:param n: items per second
:param unit: optional item unit"""                                               # throughput
    if n < 10/(365.25*86400): return item(n*(365.25*86400)) + f"{unit}/year"     # throughput
    if n < 10/(30.4375*86400): return item(n*(30.4375*86400)) + f"{unit}/month"  # throughput
    if n < 10/86400: return item(n*86400) + f"{unit}/day"                        # throughput
    if n < 10/3600: return item(n*3600) + f"{unit}/hour"                         # throughput
    if n < 10/60: return item(n*60) + f"{unit}/minute"                           # throughput
    return item(n) + f"{unit}/s"                                                 # throughput
_esc = '\033['                                                                   # throughput
_end = f'{_esc}0m'                                                               # throughput
class txt:                                                                       # txt
    """Text formatting.
Example::

    # will print out red text
    print(fmt.txt.red("some text"))"""                                           # txt
    @staticmethod                                                                # txt
    def darkcyan(s:str):  return f"{_esc}36m{s}{_end}"                           # txt
    @staticmethod                                                                # txt
    def red(s:str):       return f"{_esc}91m{s}{_end}"                           # txt
    @staticmethod                                                                # txt
    def green(s:str):     return f"{_esc}92m{s}{_end}"                           # txt
    @staticmethod                                                                # txt
    def yellow(s:str):    return f"{_esc}93m{s}{_end}"                           # txt
    @staticmethod                                                                # txt
    def blue(s:str):      return f"{_esc}94m{s}{_end}"                           # txt
    @staticmethod                                                                # txt
    def purple(s:str):    return f"{_esc}95m{s}{_end}"                           # txt
    @staticmethod                                                                # txt
    def cyan(s:str):      return f"{_esc}96m{s}{_end}"                           # txt
    @staticmethod                                                                # txt
    def bold(s:str):      return f"{_esc}1m{s}{_end}"                            # txt
    @staticmethod                                                                # txt
    def grey(s:str):      return f"{_esc}38;2;150;150;150m{s}{_end}"             # txt
    @staticmethod                                                                # txt
    def darkgrey(s:str):  return f"{_esc}38;2;100;100;100m{s}{_end}"             # txt
    @staticmethod                                                                # txt
    def underline(s:str): return f"{_esc}4m{s}{_end}"                            # txt
    @staticmethod                                                                # txt
    def identity(s:str):  return f"{s}"                                          # txt
class code:                                                                      # code
    def python(code:str):                                                        # code
        """Formats Python code.
Example::

    fmt.code.python(\"\"\"
    def f(x:int):
        return x + 3
    \"\"\") | aS(IPython.display.HTML)
"""                                                                              # code
        a = "<style>" + pygments.formatters.HtmlFormatter().get_style_defs(".highlight") + "</style>" # code
        b = pygments.highlight(code, pygments.lexers.PythonLexer(), pygments.formatters.HtmlFormatter()) # code
        return a + b                                                             # code
def h(code:str, level:int=1) -> str:                                             # h
    """Wraps content inside a 'h' html tag.
Example::

    fmt.h("abc", 2) # returns "<h2>abc</h2>"

:param level: what's the header level?"""                                        # h
    return f"<h{level}>{code}</h{level}>"                                        # h
def _jsF_h(meta, level=3):                                                       # _jsF_h
    fIdx = cli.init._jsFAuto(); dataIdx = cli.init._jsDAuto()                    # _jsF_h
    return f"const {fIdx} = ({dataIdx}) => `<h{level}>${{{dataIdx}}}</h{level}>`", fIdx # _jsF_h
k1lib.settings.cli.kjs.jsF[h] = _jsF_h                                           # _jsF_h
def pre(code:str, extras:str="") -> str:                                         # pre
    """Wraps content inside a 'pre' html tag.
Example::

    fmt.pre("abc")
"""                                                                              # pre
    return f"<pre style='font-family: monospace' {extras} >{code}</pre>"         # pre
def _jsF_pre(meta):                                                              # _jsF_pre
    fIdx = cli.init._jsFAuto(); dataIdx = cli.init._jsDAuto()                    # _jsF_pre
    return f"const {fIdx} = ({dataIdx}) => `<pre style='font-family: monospace'>${{{dataIdx}}}</pre>`", fIdx # _jsF_pre
k1lib.settings.cli.kjs.jsF[pre] = _jsF_pre                                       # _jsF_pre
def col(args, margin=10):                                                        # col
    """Creates a html col of all the elements.
Example::

    fmt.col(["abc", "def"]) | aS(IPython.display.HTML)
"""                                                                              # col
    return args | cli.apply(lambda x: f"<div style='margin: {margin}px'>{x}</div>") | cli.join("") | cli.aS(lambda x: f"<div style='display: flex; flex-direction: column'>{x}</div>") # col
def _jsF_col(meta, margin=10):                                                   # _jsF_col
    fIdx = cli.init._jsFAuto(); dataIdx = cli.init._jsDAuto()                    # _jsF_col
    fIdx2 = cli.init._jsFAuto(); dataIdx2 = cli.init._jsDAuto()                  # _jsF_col
    return f"""
const {fIdx2} = ({dataIdx2}) => `<div style='margin: {margin}px'>${{{dataIdx2}}}</div>`
const {fIdx}  = ({dataIdx})  => `<div style='display: flex; flex-direction: column'>${{{dataIdx}.map({fIdx2}).join('')}}</div>`
    """, fIdx                                                                    # _jsF_col
k1lib.settings.cli.kjs.jsF[col] = _jsF_col                                       # _jsF_col
def row(args, margin=10):                                                        # row
    """Creates a html row of all the elements.
Example::

    fmt.row(["abc", "def"]) | aS(IPython.display.HTML)
"""                                                                              # row
    return args | cli.apply(lambda x: f"<div style='margin: {margin}px'>{x}</div>") | cli.join("") | cli.aS(lambda x: f"<div style='display: flex; flex-direction: row'>{x}</div>") # row
def _jsF_row(meta, margin=10):                                                   # _jsF_row
    fIdx = cli.init._jsFAuto(); dataIdx = cli.init._jsDAuto()                    # _jsF_row
    fIdx2 = cli.init._jsFAuto(); dataIdx2 = cli.init._jsDAuto()                  # _jsF_row
    return f"""
const {fIdx2} = ({dataIdx2}) => `<div style='margin: {margin}px'>${{{dataIdx2}}}</div>`
const {fIdx}  = ({dataIdx})  => `<div style='display: flex; flex-direction: row'>${{{dataIdx}.map({fIdx2}).join('')}}</div>`
    """, fIdx                                                                    # _jsF_row
k1lib.settings.cli.kjs.jsF[row] = _jsF_row                                       # _jsF_row
settings.add("colors", ["#8dd3c7", "#ffffb3", "#bebada", "#fb8072", "#80b1d3", "#fdb462", "#b3de69", "#fccde5", "#d9d9d9", "#bc80bd", "#ccebc5", "#ffed6f"], "List of colors to cycle through in fmt.colors()") # _jsF_row
def colors():                                                                    # colors
    """Returns an infinite iterator that cycles through 12 colors.
Example::

    fmt.colors() | head(3) | deref()

Color scheme taken from https://colorbrewer2.org/#type=qualitative&scheme=Set3&n=12""" # colors
    return settings.colors | cli.repeatFrom()                                    # colors
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')               # colors
def rmAnsi(text):                                                                # rmAnsi
    """Removes ansi escape characters, courtesy of https://stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python.""" # rmAnsi
    return ansi_escape.sub('', text)                                             # rmAnsi
