# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_repr_str.ipynb.

# %% auto 0
__all__ = ['lovely']

# %% ../nbs/00_repr_str.ipynb 3
import warnings

import torch

from lovely_numpy import np_to_str_common, pretty_str, sparse_join, ansi_color, in_debugger, bytes_to_human
from lovely_numpy import config as lnp_config

from .utils.config import get_config, config
from .utils.misc import to_numpy

# %% ../nbs/00_repr_str.ipynb 6
def type_to_dtype(t: str) -> torch.dtype:
    "Convert str, e.g. 'float32' to torch.dtype e.g torch.float32"

    dtp = vars(torch)[t]
    assert isinstance(dtp, torch.dtype)
    return dtp

# %% ../nbs/00_repr_str.ipynb 8
dtnames = { type_to_dtype(k): v
                for k,v in {"float32": "",
                            "float16": "f16",
                            "float64": "f64",
                            "bfloat16": "bf16",
                            "uint8": "u8", # torch does not have uint16/32/64
                            "int8": "i8",
                            "int16": "i16",
                            "int32": "i32",
                            "int64": "i64",
                        }.items()
}

def short_dtype(x):
    # split(".") to torch.uint4 -> uint4
    return dtnames.get(x.dtype, str(x.dtype).split(".")[-1])

# %% ../nbs/00_repr_str.ipynb 10
def plain_repr(x: torch.Tensor):
    "Pick the right function to get a plain repr"
    return x._plain_repr() if hasattr(x, "_plain_repr") else repr(x)

def plain_str(x: torch.Tensor):
    "Pick the right function to get a plain str."
    return x._plain_str() if hasattr(x, "_plain_str") else str(x)

# %% ../nbs/00_repr_str.ipynb 11
def is_nasty(t: torch.Tensor):
    """Return true of any `t` values are inf or nan"""

    if t.numel() == 0: return False # amin/amax don't like zero-lenght tensors
    if (t.device.type == "meta"): return False
    # Unlike .min()/.max(), amin/amax do not allocate extra GPU memory.

    t_min = t.amin().cpu()
    t_max = t.amax().cpu()

    return (t_min.isnan() or t_min.isinf() or t_max.isinf()).item()

# %% ../nbs/00_repr_str.ipynb 13
def torch_to_str_common(t: torch.Tensor,  # Input
                        color=True,       # ANSI color highlighting
                        ) -> str:

    if t.numel() == 0: return ansi_color("empty", "grey", color)
    if t.device.type == "meta": return ansi_color("meta", "grey", color)

    # Unlike .min()/.max(), amin/amax do not allocate extra GPU memory.
    amin, amax = t.amin().cpu(), t.amax().cpu()

    zeros = ansi_color("all_zeros", "grey", color) if amin.eq(0) and amax.eq(0) and t.numel() > 1 else None

    summary = None
    if not zeros:
        minmax = f"x∈[{pretty_str(amin)}, {pretty_str(amax)}]" if t.numel() > 2 else None
        meanstd = f"μ={pretty_str(t.mean())} σ={pretty_str(t.std())}" if t.numel() >= 2 else None
        summary = sparse_join([minmax, meanstd])

    return sparse_join([ summary, zeros ])

# %% ../nbs/00_repr_str.ipynb 14
# tensor.is_cpu was only introduced in 1.13.
def is_cpu(t: torch.Tensor):
    return t.device == torch.device("cpu")

# %% ../nbs/00_repr_str.ipynb 15
# Fake tensors are used by torch.compile when tracing the graph, or something.
try:
    from torch._subclasses.fake_tensor import is_fake
except ImportError:
    is_fake = lambda t: False

# %% ../nbs/00_repr_str.ipynb 16
@torch.no_grad()
def to_str(t: torch.Tensor,
            plain: bool=False,
            verbose: bool=False,
            depth=0,
            lvl=0,
            color=None) -> str:

    if plain or is_fake(t):
        return plain_repr(t)

    conf = get_config()

    names = getattr(t, "names", None)
    if names and not all(n is None for n in names):
        # amin/amax don't like named tensors
        t_no_names = t.rename(None)
    else:
        t_no_names = t
        names = None # If all names are None, we don't want to show them.

    tname = "tensor" if type(t) is torch.Tensor else type(t).__name__.split(".")[-1]

    if names:
        shape = [f"{n}={v}" if n else str(v) for n,v in zip(names, t.shape)]
        shape = f"[{', '.join(shape)}]"
    else:
        shape = str(list(t.shape)) if t.ndim else None
    type_str = sparse_join([tname, shape], sep="")

    dev = str(t.device) if t.device.type != "cpu" else None
    dtype = short_dtype(t)
    grad_fn = t.grad_fn.name() if t.grad_fn else None
    # PyTorch does not want you to know, but all `grad_fn``
    # tensors actuall have `requires_grad=True`` too.
    grad = "grad" if t.requires_grad else None


    # For complex tensors, just show the shape / size part for now.
    if not t.is_complex():
        if color is None: color=conf.color
        if in_debugger(): color=False
        # `lovely-numpy` is used to calculate stats when doing so on GPU would require
        # memory allocation (not float tensors, tensors with bad numbers), or if the
        # data is on CPU (because numpy is faster).
        #
        # Temporarily set the numpy config to match our config for consistency.
        with lnp_config(precision=conf.precision,
                        threshold_min=conf.threshold_min,
                        threshold_max=conf.threshold_max,
                        sci_mode=conf.sci_mode):
            if is_nasty(t_no_names) or not t.is_floating_point():
                common = np_to_str_common(to_numpy(t), color=color, ddof=1)
            else:
                common = torch_to_str_common(t_no_names, color=color)

            numel = None
            nbytes = t.numel() * t.element_size()
            if t.shape and max(t.shape) != t.numel():
                numel = f"n={t.numel()}"
                if get_config().show_mem_above <= nbytes:
                    numel = sparse_join([numel, f"({bytes_to_human(nbytes)})"])
            elif get_config().show_mem_above <= nbytes:
                numel = bytes_to_human(nbytes)

            vals = None
            if t.device.type != "meta":
                vals = pretty_str(to_numpy(t)) if 0 < t.numel() <= 10 else None
            res = sparse_join([type_str, dtype, numel, common, grad, grad_fn, dev, vals])
    else:
        res = plain_repr(t)


    if verbose:
        res += "\n" + plain_repr(t)

    if depth and t.dim() > 1:
        with config(show_mem_above=torch.inf):
            deep_width = min((t.shape[0]), conf.deeper_width) # Print at most this many lines
            deep_lines = [ " "*conf.indent*(lvl+1) + to_str(t[i,:], depth=depth-1, lvl=lvl+1)
                                for i in range(deep_width)]

            # If we were limited by width, print ...
            if deep_width < t.shape[0]: deep_lines.append(" "*conf.indent*(lvl+1) + "...")

            res += "\n" + "\n".join(deep_lines)

    return res

# %% ../nbs/00_repr_str.ipynb 17
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")

# %% ../nbs/00_repr_str.ipynb 20
class StrProxy():
    def __init__(self, t: torch.Tensor, plain=False, verbose=False, depth=0, lvl=0, color=None):
        self.t = t
        self.plain = plain
        self.verbose = verbose
        self.depth=depth
        self.lvl=lvl
        self.color=color
        history_warning()

    def __repr__(self):
        return to_str(self.t, plain=self.plain, verbose=self.verbose,
                      depth=self.depth, lvl=self.lvl, color=self.color)

    # This is used for .deeper attribute and .deeper(depth=...).
    # The second onthe results in a __call__.
    def __call__(self, depth=1):
        return StrProxy(self.t, depth=depth)

# %% ../nbs/00_repr_str.ipynb 21
def lovely(t: torch.Tensor, # Tensor of interest
            verbose=False,  # Whether to show the full tensor
            plain=False,    # Just print if exactly as before
            depth=0,        # Show stats in depth
            color=None):    # Force color (True/False) or auto.
    return StrProxy(t, verbose=verbose, plain=plain, depth=depth, color=color)
