# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/03_utils.utils.ipynb.

# %% auto 0
__all__ = ['sci_mode', 'pretty_str', 'plain_repr', 'plain_str', 'sparse_join', 'ansi_color', 'np_to_str_common',
           'history_warning']

# %% ../../nbs/03_utils.utils.ipynb 3
from collections import defaultdict
import warnings
import numpy as np
from typing import Optional, Union
from .config import get_config


# %% ../../nbs/03_utils.utils.ipynb 4
# Do we want this float in decimal or scientific mode?
def sci_mode(f: float):
    config = get_config()
    return ((abs(f) < 10**config.threshold_min) or
            (abs(f) > 10**config.threshold_max))

# %% ../../nbs/03_utils.utils.ipynb 7
# Convert an ndarray or scalar into a string.
# This only looks good for small tensors, which is how it's intended to be used.
def pretty_str(x):
    """A slightly better way to print `float`-y values.
    Works for `np.ndarray`, `torch.Tensor`, `jax.DeviceArray`, and scalars."""

    if isinstance(x, int):
        return '{}'.format(x)
    elif isinstance(x, float):
        if x == 0.:
            return "0."

        sci = sci_mode(x) if get_config().sci_mode is None else get_config().sci_mode
        
        fmt = f"{{:.{get_config().precision}{'e' if sci else 'f'}}}"

        return fmt.format(x)
    elif x.ndim == 0:
            return pretty_str(x.item())
    else:
        slices = [pretty_str(x[i]) for i in range(0, x.shape[0])]
        return '[' + ", ".join(slices) + ']'

# %% ../../nbs/03_utils.utils.ipynb 12
def plain_repr(x):
    "Pick the right function to get a plain repr. `x` is not limited to numpy"
    # assert isinstance(x, np.ndarray), f"expected np.ndarray but got {type(x)}" # Could be a sub-class.
    return x._plain_repr() if hasattr(type(x), "_plain_repr") else repr(x)

def plain_str(x):
    "Pick the right function to get a plain str. `x` is not limited to numpy"
    # assert isinstance(x, np.ndarray), f"expected np.ndarray but got {type(x)}"
    return x._plain_str() if hasattr(type(x), "_plain_str") else str(x)

# %% ../../nbs/03_utils.utils.ipynb 13
def sparse_join(lst, sep=" "):
    # Join non-empty list elements into a space-sepaeated string
    return sep.join( [ l for l in lst if l] )

# %% ../../nbs/03_utils.utils.ipynb 15
def ansi_color(s: str, col: str, use_color=True):
        "Very minimal ANSI color support"
        style = defaultdict(str)
        style["grey"] = "\x1b[38;2;127;127;127m"
        style["red"] = "\x1b[31m"
        end_style = "\x1b[0m"
       
        return style[col]+s+end_style if use_color else s

# %% ../../nbs/03_utils.utils.ipynb 18
def np_to_str_common(x: Union[np.ndarray, np.generic],  # Input
                        color=True,                     # ANSI color highlighting
                        ddof=0):                        # For "std" unbiasing

    if x.size == 0:
        return ansi_color("empty", "grey", color)

    zeros = ansi_color("all_zeros", "grey", color) if np.equal(x, 0.).all() and x.size > 1 else None
    pinf = ansi_color("+Inf!", "red", color) if np.isposinf(x).any() else None
    ninf = ansi_color("-Inf!", "red", color) if np.isneginf(x).any() else None
    nan = ansi_color("NaN!", "red", color) if np.isnan(x).any() else None

    attention = sparse_join([zeros,pinf,ninf,nan])
    numel = f"n={x.size}" if x.size > 5 and max(x.shape) != x.size else None

    summary = None
    if not zeros and isinstance(x, np.ndarray):
        # Calculate stats on good values only.
        gx = x[ np.isfinite(x) ]

        minmax = f"x∈[{pretty_str(gx.min())}, {pretty_str(gx.max())}]" if gx.size > 2 else None
        meanstd = f"μ={pretty_str(gx.mean())} σ={pretty_str(gx.std(ddof=ddof))}" if gx.size >= 2 else None
        summary = sparse_join([numel, minmax, meanstd])


    return sparse_join([ summary, attention])

# %% ../../nbs/03_utils.utils.ipynb 22
def history_warning():
    "Issue a warning (once) ifw e are running in IPYthon with output cache enabled"
    if "get_ipython" in globals() and get_ipython().cache_size > 0:
        warnings.warn("IPYthon has its output cache enabled. See https://xl0.github.io/lovely-tensors/history.html")
