
from pycamia import info_manager

__info__ = info_manager(
    project = "PyCAMIA",
    package = "batorch",
    fileinfo = "Tensor functions for `batorch`.",
    requires = ["numpy", "torch", "pycamia", "pyoverload", "matplotlib"]
)

__all__ = """
    batorch_wrapper
    conv
    crop_as
    decimal
    divide
    dot
    equals
    gaussian_kernel
    image_grid               imagegrid
    mask_center
    one_hot
    pad
    permute_space
    normalize01

    grad_image               Jacobian
    up_scale                 down_scale
    
    upper_matrix             upper_triangular_matrix
    lower_matrix             lower_triangular_matrix
    skew_symmetric           cross_matrix
    skew_symmetric_params    uncross_matrix
    
    norm                     norm2
    Fnorm                    Fnorm2
    frobenius_norm
    meannorm                 meannorm2
    mean_norm                mean_norm2
    
    summary                  display
""".split()

import math, re

with __info__:
    import numpy as np
    import torch
    import batorch as bt
    from functools import wraps
    from pyoverload import *
    from pycamia import avouch, touch, alias, decorator
    from pycamia import get_environ_vars, get_snakenames, get_args_expression
    from pycamia import SPrint, Version, argmin, tokenize, to_tuple
    from matplotlib import pyplot as plt
    
@decorator
def batorch_wrapper(func: callable=None, **options):
    """
    The wrapper for batorch functions, ensuring the inputs of the function are batorch 
        tensors, and the return type will be casted as the input type for non-batorch inputs. 
    """
    if func is None: return lambda f: batorch_wrapper(f, **options)
    i_arg = options.pop('i_arg', None)
    default_aux = options
    @wraps(func)
    def wrapper_func(*args, **kwargs):
        values = get_arg_values(get_func_info(func), *args, **kwargs)
        f_varnames = func.__code__.co_varnames
        new_args = []
        new_kwargs = {}
        input_class = None
        for i, (n, v) in enumerate(zip(f_varnames, values)):
            if func.__annotations__.get(n, None) in (array, bt.Tensor):
                if i_arg is None and input_class is None or i == i_arg: input_class = v.__class__
                if v is None: ...
                elif not isinstance(v, torch.Tensor): v = bt.tensor(v)
                elif not isinstance(v, bt.Tensor): v = v.as_subclass(bt.Tensor).init_special()
            if (func.__code__.co_flags & 0x04) and i == func.__code__.co_argcount + func.__code__.co_kwonlyargcount: new_args.extend(v)
            elif (func.__code__.co_flags & 0x08) and i == func.__code__.co_argcount + func.__code__.co_kwonlyargcount + ((func.__code__.co_flags & 0x04) >> 2): new_kwargs.update(v)
            elif i < len(args): new_args.append(v)
            else: new_kwargs[n] = v
        ret = func(*new_args, **new_kwargs)
        if not isinstance(ret, bt.Tensor): return ret
        ret_obj = None
        if input_class is not None:
            if issubclass(input_class, torch.Tensor) and input_class != bt.Tensor: ret_obj = ret.as_subclass(input_class)
            elif issubclass(input_class, np.ndarray): ret_obj = ret.detach().cpu().numpy()
            elif issubclass(input_class, list): ret_obj = ret.detach().cpu().numpy().tolist()
            elif issubclass(input_class, tuple): ret_obj = tuple(ret.detach().cpu().numpy().tolist())
        if ret_obj is None: ret_obj = ret
        for k, default in default_aux.items():
            aux = getattr(ret, k, default)
            if issubclass(default.__class__, torch.Tensor):
                if input_class != bt.Tensor: aux = (aux if isinstance(aux, torch.Tensor) else torch.tensor(aux)).as_subclass(default.__class__)
                else: aux = aux if isinstance(aux, bt.Tensor) else bt.tensor(aux)
            elif issubclass(default.__class__, np.ndarray): aux = aux.detach().cpu().numpy() if isinstance(aux, torch.Tensor) else np.array(aux)
            elif issubclass(default.__class__, list): aux = aux.detach().cpu().numpy().tolist() if isinstance(aux, torch.Tensor) else list(aux)
            elif issubclass(default.__class__, tuple): aux = tuple(aux.detach().cpu().numpy().tolist()) if isinstance(aux, torch.Tensor) else tuple(aux)
            try: setattr(ret_obj, k, aux)
            except AttributeError: pass
        return ret_obj
    return wrapper_func

# @overload
# def batorch_wrapper(func: callable, i_arg=None, **default_aux):
#     @wraps
#     def wrapper_func(*args, **kwargs):
#         values = get_arg_values(get_func_info(func), *args, **kwargs)
#         f_varnames = func.__code__.co_varnames
#         n_posarg = func.__code__.co_posonlyargcount
#         new_args = []
#         new_kwargs = {}
#         input_class = None
#         for i, (n, v) in enumerate(zip(f_varnames, values)):
#             if func.__annotations__.get(n, None) == bt.Tensor:
#                 if i_arg is None and input_class is None or i == i_arg: input_class = v.__class__
#                 if not isinstance(v, torch.Tensor): v = bt.tensor(v)
#                 else: v = v.as_subclass(bt.Tensor)
#             if i < n_posarg: new_args.append(v)
#             else: new_kwargs[n] = v
#         ret = func(*new_args, **new_kwargs)
#         if input_class is not None:
#             if issubclass(input_class, torch.Tensor): ret_obj = ret.as_subclass(input_class)
#             elif issubclass(input_class, np.ndarray): ret_obj = ret.detach().cpu().numpy()
#             elif issubclass(input_class, list): ret_obj = ret.detach().cpu().numpy().tolist()
#             elif issubclass(input_class, tuple): ret_obj = tuple(ret.detach().cpu().numpy().tolist())
#         for k, default in default_aux:
#             aux = getattr(ret, k, default)
#             if issubclass(default.__class__, torch.Tensor): aux = (aux if isinstance(aux, torch.Tensor) else torch.tensor(aux)).as_subclass(default.__class__)
#             elif issubclass(default.__class__, np.ndarray): aux = aux.detach().cpu().numpy() if isinstance(aux, torch.Tensor) else np.array(aux)
#             elif issubclass(default.__class__, list): aux = aux.detach().cpu().numpy().tolist() if isinstance(aux, torch.Tensor) else list(aux)
#             elif issubclass(default.__class__, tuple): aux = tuple(aux.detach().cpu().numpy().tolist()) if isinstance(aux, torch.Tensor) else tuple(aux)
#             setattr(ret_obj, k, aux)
#         return ret_obj
#     return wrapper_func

# @overload
# def batorch_wrapper(i_arg: int | null =None, **default_aux):
#     return wraps(lambda f: batorch_wrapper(f, i_arg, **default_aux))

@batorch_wrapper
def conv(tensor: bt.Tensor, kernel: bt.Tensor, padding='SAME'):
    """
    Convolution for a batorch tensor: only in space dimensions. 
    
    Args:
        tensor (bt.Tensor): With shape ({n_batch: optional}, [n_channel_in: optional], *n_space).
        kernel (bt.Tensor): With shape ([n_channel_out: optional], [n_channel_in: optional], *n_kernel_size).
        padding (str, int, tuple): Convolution padding: string 'SAME' stands for not changing the size of 'tensor'. 
    
    Returns:
        output (bt.Tensor): With shape ({n_batch: optional}, [n_channel_out: optional], *(n_space - n_kernel_size // 2)).
    """
    input_shape = tensor.shape
    if not tensor.has_batch: tensor = tensor.unsqueeze({})
    if not tensor.has_feature: tensor = tensor.unsqueeze([])
    n_dim = tensor.n_space_dim
    avouch(kernel.n_space_dim == n_dim, f"Kernel for {n_dim}D convolution should be of size ([channel_out, channel_in], n_kernel_1, ..., n_kernel_{n_dim}) with a dimension of {n_dim}. ")
    if kernel.has_feature:
        avouch(kernel.n_feature_dim == 2, f"Kernel for {n_dim}D convolution should have exactly 2 channel dimensions, only size ([channel_out, channel_in], n_kernel_1, ..., n_kernel_{n_dim}) is allowed, not {kernel.shape}.")
        avouch(tensor.has_channel, f"Tensor input for convolution should have channel dimension (one feature dimension): not tensor size {tensor.shape} and kernel size {kernel.shape}. Use kernel without special dimension to perform spatial-dimension-only convolution. ")
        avouch(tensor.n_feature_dim == 1 and tensor.n_channel == kernel.feature[1], f"Channels of convolution should match: tensor size {str(tensor.shape).replace('[', '[>').replace(']', '<]')} and kernel size {str(kernel.shape).replace(f'{kernel.feature[-1]}]', f'>{kernel.feature[-1]}<]')}. ")
        space_conv = False
    else: space_conv = True

    is_same = False
    if isinstance(padding, str) and padding == 'SAME':
        is_same = True
        padding = tuple(x // 2 for x in kernel.space)
    if isinstance(padding, int):
        padding = (padding,) * n_dim
    avouch(isinstance(padding, tuple), f"Unrecognized argument 'padding' for convolution: {padding}. ")
    conv = eval("bt.nn.functional.conv%dd"%n_dim)
    
    if space_conv:
        prev_batch_feature = bt.Size({tensor.n_batch}) + tensor.feature
        tensor = tensor.merge_dims([], {}).unsqueeze([])
        kernel = kernel.unsqueeze(0, 0)
    result = conv(tensor, kernel.type(tensor.dtype), padding=padding)
    if space_conv: result = result.squeeze([]).splitdim({}, prev_batch_feature)
    if not input_shape.has_batch: result.squeeze_({})
    if not input_shape.has_feature: result.squeeze_([])
    if is_same: result = crop_as(result, input_shape.space)
    
    # if kernel.size(0) > 1:
    #     if input_shape.has_channel:
    #         input_shape = input_shape.with_dim_size(input_shape.channel_dim, kernel.size(0))
    
    # if is_same: result = crop_as(result, tensor).view(*input_shape)
    # else: result = result.view(input_shape.with_space(tuple(t - k + 1 + 2 * p for t, k, p in zip(tensor.space, kernel.space, padding))))
    return result

@alias("meannorm2", "mean_norm2", root = False, mean = True)
@alias("meannorm", "mean_norm", mean = True)
@alias("Fnorm2", "norm2", root = False)
@alias("Fnorm", "frobenius_norm")
@batorch_wrapper
def norm(tensor: bt.Tensor, p = 2, root = True, mean = False, dim=null):
    """
    The norm of a tensor in (the first available condition):
    (1) all feature dimensions (if n_feature_dim >= 1);
    (2) all space dimensions (if n_space_dim >= 1); 
    (3) all sequence dimensions (if n_sequence_dim >= 1). 
    Use 'dim=None' to compute for all dimensions. 

    Args:
        tensor (bt.Tensor): The tensor to compute norm, commonly in shape ({n_batch}, [n_channel], n_1, ..., n_{n_dim})
        p (int, optional): The order of norm. Defaults to 2.
            p = 1 for L1-norm/ LASSO;
            p = 2 for L2-norm or Frobenius;
            p = math.inf / float('inf') for L∞-norm / maximum. 
        root (bool, optional): Whether to calculate the root. Defaults to True for common norms.
        mean (bool, optional): Whether to use mean across the dimensions. Defaults to False for common norms.
        dim (_type_, optional): The dimension to perform norm. Defaults to null, which auto selects the dimensions.
            One need to manually set 'dim=None' to perform norm in all dimensions. 
    """
    if dim is None: dim = bt.exist_dim(tensor, dim)
    else: dim = bt.linalg_dim[1:](tensor, None if dim is null else dim)
    tensor = tensor.abs()
    if p == float('inf'): return tensor.max(*dim)
    if mean: return (tensor ** p).mean(*dim) ** (1 / p) if root else (tensor ** p).mean(*dim)
    else: return (tensor ** p).sum(*dim) ** (1 / p) if root else (tensor ** p).sum(*dim)

@batorch_wrapper
def decimal(tensor: bt.Tensor):
    """
    The decimal part of tensor: {x} = x-[x]. 
    """
    return tensor - bt.floor(tensor)

@batorch_wrapper
def divide(a: bt.Tensor, b: bt.Tensor, limit=1, tol=1e-6):
    """
    def divide(a: bt.Tensor, b: bt.Tensor, limit=1, tol=1e-6)
    Element-wise division `a / b`, except it equals to `limit` if the divider `b` is smaller than `tol` (which is regarded as the division of zero)
    """
    a_s, b_s = a.shape ^ b.shape
    a = a.view(a_s)
    b = b.view(b_s)
    shape = bt.Size(max(x, y) for x, y in zip(a_s, b_s)).special_from(a_s)
    return bt.where(b.abs() < tol, bt.where(a.abs() < tol, bt.zeros(shape), limit * bt.ones(shape)), a / bt.where(b.abs() < tol, tol * bt.ones(shape), b))

@batorch_wrapper
def equals(x: bt.Tensor, y: bt.Tensor, tol=1e-6):
    """
    Determine whether the corresponding elements in two tensors are equal to each other. 
    The region of equivalence is a union of a diamond area at the origin point (⟐) 
        and a cross area in shape 'x' (with borders not passing the origin: '⪥').
    The cross area includes points near lines y=x and y=-x with an angle dilation of approximately arctan(tol). 
    The diamond area has a side length of √2*tol. 
    The reason why a diamond area is added at the origin instead of square (⊡) or circular region (⊙) is that,
        for inputs of difference storage accuracy (e.g. int & float), higher tolerance should be placed for errors 
        in the lower accuracy (e.g. float) tensor than in the higher accuracy (e.g. int) tensor, causing peaks in tolerance map. 
    """
    return ((bt.abs(x - y) / (bt.abs(x) + bt.abs(y) + tol)) < tol) | (bt.abs(x) + bt.abs(y) < tol)

def gaussian_kernel(n_dim = 2, kernel_size = 3, sigma = 0, normalize = True):
    """
    Create a bt.Tensor Gaussian kernel of size (kernel_size, ..., kernel_size). 
                                                \---{n_dim} kernel_size's---/

    Args:
        n_dim (int, optional): The dimension of the kernel. Defaults to 2.
        kernel_size (int, optional): The side length of the kernel. Defaults to 3.
        sigma (int or float, optional): The standard deviation sigma of Gaussian distribution, 
                                            in unit 'pixel'. Defaults to 0.
        normalize (bool, optional): Whether the sum of kernel should be 1. Defaults to True.
    """
    if not isinstance(kernel_size, bt.torch.Tensor):
        if not isinstance(kernel_size, (tuple, list)): kernel_size = (kernel_size,)
        kernel_size = list(kernel_size)
        if len(kernel_size) == 1: kernel_size *= n_dim
        kernel_size = bt.Tensor(kernel_size)
    elif not isinstance(kernel_size, bt.Tensor): kernel_size = kernel_size.as_subclass(bt.Tensor).init_special()
    kernel_size.n_feature_dim = 1
    radius = (kernel_size - 1) / 2

    if sigma == 0: sigma = radius * 0.6
    elif not isinstance(sigma, bt.torch.Tensor):
        if not isinstance(sigma, (tuple, list)): sigma = (sigma,)
        sigma = list(sigma)
        if len(sigma) == 1: sigma *= n_dim
        sigma = bt.Tensor(sigma)
    elif not isinstance(sigma, bt.Tensor): sigma = sigma.as_subclass(bt.Tensor).init_special()
    sigma.n_feature_dim = 1

    grid = bt.image_grid(*kernel_size.tolist()).float()
    kernel = bt.exp(- (((grid - radius) / sigma) ** 2).sum(0) / 2)
    return (kernel / kernel.sum()) if normalize else kernel

@batorch_wrapper
def dot(g1: bt.Tensor, g2: bt.Tensor, dim=null):
    """
    The dot product of two tensors in (the first available condition):
    (1) all feature dimensions (if n_feature_dim >= 1);
    (2) all space dimensions (if n_space_dim >= 1); 
    (3) all sequence dimensions (if n_sequence_dim >= 1). 
    Use 'dim=None' to compute for all dimensions. 
    """
    if dim is None: dim = bt.exist_dim(g1, dim)
    else: dim = bt.linalg_dim[1:](g1, None if dim is null else dim)
    avouch(g1.shape == g2.shape, "Please make sure the dot product recieve two tensor of a same shape. Use 'bt.expand_to' to adjust if necessary. ")
    avouch(g1.n_feature_dim == g2.n_feature_dim, "Please make sure the inputs of 'dot' have a same feature dimension. ")
    return (g1 * g2).sum(*dim)

@batorch_wrapper
def Jacobian(X, Y, dt=1, pad=False):
    """
    The Jacobian matrix; Note that it is a transpose of grad_image if X is standard grid as it follows the orientation of Jacobian.
    
    Args:
        X (bt.Tensor): in shape ({n_batch}, [n_dim], n_1, ..., n_{n_dim})
        Y (bt.Tensor): in shape ({n_batch}, [n_feature], n_1, ..., n_{n_dim})
        
    Returns:
        Jacobian(output, pad = True): in shape ({n_batch}, [n_feature, n_dim], n_1, ..., n_{n_dim})
        Jacobian(output, pad = False): in shape ({n_batch}, [n_feature, n_dim], n_1 - dt, ..., n_{n_dim} - dt)
    """
    dX = grad_image(X, dx=dt, pad=pad) # ({n_batch}, [n_dim, n_dim], n_1, ..., n_{n_dim})
    dY = grad_image(Y, dx=dt, pad=pad) # ({n_batch}, [n_dim, n_feature], n_1, ..., n_{n_dim})
    # dy_k/dx_i = sum_j [(dy_k/dt_j) / (dx_i/dt_j)] => Jac[{b}, [k, i], ...] = sum_j Jac_j[{b}, [j, k, i], ...] = dY[{b}, [j, k, *], ...] / dX[{b}, [j, *, i], ...]
    return bt.divide(dY.unsqueeze([3]), dX.unsqueeze([2]), 0.0, tol=1e-4).sum([0])

@batorch_wrapper
def grad_image(array: bt.Tensor, dx=1, pad=False):
    '''
        Gradient image / tensor of array (dx is the displacement for finite difference). 
        If pad is not True, the image will trim. The sizes should be: 
        array: ({n_batch}, [n_feature], n_1, ..., n_{n_dim})
        output (pad = True): ({n_batch}, [n_dim, n_feature], n_1, ..., n_{n_dim})
        output (pad = False): ({n_batch}, [n_dim, n_feature], n_1 - dx, ..., n_{n_dim} - dx)
        OR:
        array: ({n_batch}, n_1, ..., n_{n_dim})
        output (pad = True): ({n_batch}, [n_dim], n_1, ..., n_{n_dim})
        output (pad = False): ({n_batch}, [n_dim], n_1 - dx, ..., n_{n_dim} - dx)
    '''
    if not array.has_batch: array = array.unsqueeze({})
    output = []
    size = array.space
    if not pad: size = tuple(x - dx for x in size)
    for d in range(*array.space_range):
        b = (slice(None, None),) * d + (slice(dx, None),) + (slice(None, None),) * (array.ndim - d - 1)
        a = (slice(None, None),) * d + (slice(None, -dx),) + (slice(None, None),) * (array.ndim - d - 1)
        output.append(bt.crop_as((array[b] - array[a]) / dx, size))
    return bt.stack(output, [int(array.has_batch)])

@overload
@batorch_wrapper(roi=tuple())
def crop_as(x: array, y: tuple, center: tuple, fill: number=0) -> array:
    """
    crop a Tensor `x` as the shape given by `y`. 

    Args:
        x (bt.Tensor): The data to crop (or pad if the target shape is larger), 
            Shape: ({n_batch}, [*n_feature], *n_space). 
            Note that only the space dimensions of the tensor are cropped/padded by default.
        y (bt.Tensor or tuple): The target shape in tuple or another tensor to provide the target shape.
        center (tuple, optional): The center of the target box. Defaults to the center of x's shape. 
            Note: Do calculate the center w.r.t. input `x` if one is expanding the tensor, 
                as `x`'s center coordinates in y-space is different from `y`'s center 
                coordinates in x-space (which is correct).
        fill (number, optional): The number to fill for paddings. Defaults to 0.
    """
    size_x = x.shape
    size_y = bt.Size(y)

    if isinstance(size_y, bt.Size) and size_x.n_space_dim == size_y.n_space_dim:
        size_y = size_y.space
    exp_size_x, exp_size_y = size_x ^ bt.Size(-1 if s == 1 else s for s in size_y).special_from(size_y)
    avouch(len(exp_size_x) == len(size_x), f"Cannot match crop size {size_y} with tensor size {size_x} in 'crop_as'. ")
    exp_size_y = bt.Size(-s if s in (1, -1) else s for s in exp_size_y).special_from(exp_size_y)
    size_x, size_y = exp_size_x, exp_size_y

    center = bt.Size(center)
    if len(center) == len(size_x): center = center.special_from(size_x)
    elif len(center) == size_x.n_space_dim: center = (center ^ size_x)[0]
    elif len(x for x in center if x >= 0) == len(x for x in size_y if x >= 0):
        center = bt.Size(a if b >= 0 else -1 for a, b in zip(center, size_y)).special_from(size_x)
    avouch(len(center) == len(size_x), TypeError("Mismatch dimensions for the center in 'crop_as', please use -1 if the dimension that is centered or doesn't need cropping. "))
    center = bt.Size(-1 if sx == sy or sy < 0 else (sx / 2 if c < 0 else c) for c, sx, sy in zip(center, size_x, size_y)).special_from(center)
    size_y = bt.Size(sx if sy < 0 else sy for sx, sy in zip(size_x, size_y)).special_from(size_y)

    z = fill * bt.ones(*size_y).to(x.device).type_as(x)
    """
    We then calculate the intersected boxes in source 'x' and target 'z', 
    to manage the occasions when the two regions are not containing each other (like: ⧉). 
    """
    def intersect(u, v): return max(u[0], v[0]), min(u[1], v[1])
    z_box = [(0, lx) if m < 0 or ly < 0 else intersect((0, ly), (- round(float(m - float(ly) / 2)), - round(float(m - float(ly) / 2)) + lx)) for m, lx, ly in zip(center, size_x, size_y)]
    x_box = [(0, lx) if m < 0 or ly < 0 else intersect((0, lx), (+ round(float(m - float(ly) / 2)), + round(float(m - float(ly) / 2)) + ly)) for m, lx, ly in zip(center, size_x, size_y)]
    # if the two boxes are seperated
    if any([r[0] >= r[1] for r in z_box]) or any([r[0] >= r[1] for r in x_box]): z.roi = None; return z
    region_z = tuple(slice(u, v) for u, v in z_box)
    region_x = tuple(slice(u, v) for u, v in x_box)
    z[region_z] = x[region_x]
    z.roi = region_x
    z.special_from(size_x)
    return z

@overload
def crop_as(x: array, y: array, center: tuple, fill: number=0) -> array:
    return crop_as(x, y.shape, center, fill)

@overload
def crop_as(x: array, y: union(tuple, array), fill: number=0) -> array:
    center = tuple(-1 for m in x.shape)
    return crop_as(x, y, center, fill)

@overload
def crop_as(x: array, *y: int) -> array:
    center = tuple(-1 for m in x.shape)
    return crop_as(x, y, center)

@batorch_wrapper
def pad(x: bt.Tensor, p = 1, fill = 0):
    """
    Pad the tensor by `p` pixels with values `fill`. 
    p can be a tuple of len `n_dim` to pad differently for each dimension. 
    """
    p = to_tuple(p)
    if len(p) == 1: p *= x.n_space_dim
    return crop_as(x, tuple(s + 2 * q for s, q in zip(x.space, p)), fill = fill)

def add_special(size, special, fill=1):
    s = special
    if len(s) == 0: pass
    elif len(s) == 1: size = size[:s[0]] + (fill,) + size[s[0]:]
    else: size = size[:s[0]] + (fill,) + size[s[0]:s[1]-1] + (fill,) + size[s[1]-1:]
    return size

@batorch_wrapper
def up_scale(image: bt.Tensor, *scaling:int):
    """
    up_scale(tensor, s_1, ..., s_{n_dim}): equivalent to tensor.amplify(s_i, i)
    """
    if len(scaling) == 0:
        scaling = (1,)
    elif len(scaling) == 1 and iterable(scaling[0]):
        scaling = scaling[0]
    if len(scaling) == 1:
        if isinstance(scaling[0], int):
            scaling *= image.n_space_dim
            scaling = add_special(scaling, image.special_dims, 1)
        else: raise TypeError("Unknown scaling type for 'up_scale'. ")
    elif len(scaling) < image.n_dim and len(scaling) == image.n_space_dim:
        scaling = add_special(scaling, image.special_dims, 1)
    for i, s in enumerate(scaling):
        image = (
            image
            .transpose(i, -1)
            .unsqueeze(-1)
            .repeat((1,) * image.ndim + (int(s),))
            .flatten(-2)
            .transpose(i, -1)
        )
    return image

@batorch_wrapper
def down_scale(image: bt.Tensor, *scaling:int):
    """
    down_scale(tensor, s_1, ..., s_{n_dim}): equivalent to tensor[::s_i, ..., ::s_{n_dim}]
    """
    if len(scaling) == 0:
        scaling = (1,)
    elif len(scaling) == 1 and iterable(scaling[0]):
        scaling = scaling[0]
    if len(scaling) == 1:
        if isinstance(scaling[0], int):
            scaling *= image.n_space_dim
            scaling = add_special(scaling, image.special_dims, 1)
        else: raise TypeError("Unknown scaling type for 'down_scale'. ")
    elif len(scaling) < image.n_dim and len(scaling) == image.n_space_dim:
        scaling = add_special(scaling, image.special_dims, 1)
    return image[tuple(slice(None, None, s) for s in scaling)]

@alias("imagegrid")
def image_grid(*shape):
    """
    Create the indices for a space represented by tuple, or the space-dimension space of a tensor. 
    Note that the indices are distributed as indexing = 'ij' in the representation of keyword argument for torch.meshgrid.
    
    Outputs:
        (1) For input of space shape: (n_1, n_2, ..., n_{n_dim}), 
            output a Tensor in shape ([n_dim], n_1, n_2, ..., n_{n_dim}).
        (2) For input of tensor shape: ({n_batch}, n_1, ..., n_{n_dim}, 'n_sequence'), 
            output a Tensor in shape ([n_dim], n_1, ..., n_{n_dim}). 
        (3) ** WARNING: Please be careful with this usage as the feature dimension of the output is irrelevant to the input. **
            For input of tensor shape with feature: ({n_batch}, [n_feature], n_1, ..., n_{n_dim}, 'n_sequence'), 
            output a Tensor in shape ([n_dim], n_1, ..., n_{n_dim}). 
    """
    if len(shape) == 1 and isinstance(shape[0], (list, tuple)): shape = shape[0]
    if len(shape) == 1 and hasattr(shape[0], 'space'): shape = shape[0].space
    if len(shape) == 1 and hasattr(shape[0], 'shape'): shape = shape[0].shape
    shape = tuple(bt.Size(*shape).space)
    kwargs = {'indexing': 'ij'} if bt.__torchversion__ >= Version('1.10') else {}
    if any(x < 0 for x in shape): raise TypeError(f"Negative size in shape {shape} for image grid.")
    if len(shape) == 0: return bt.zeros([0])
    return bt.stack(bt.meshgrid(*[bt.arange(x) for x in shape], **kwargs), [])

@batorch_wrapper
def mask_center(mask: bt.Tensor):
    """
    Compute the center of gravity for a mask.
    """
    avouch(issubclass(mask.dtype, dtype(bool)), TypeError("Only bool mask is allowed in function 'bt.mask_center'. "))
    coords = bt.image_grid(mask).change_special_dim([], bt.Size(0).with_func_dim(True))
    return (coords * mask).float().sum(...) / mask.sum(...)

# def linear(input, weight, bias):
#     result = input @ weight.T
#     if bias is not None:
#         if bias.dim() == 2:
#             return result + bias
#         return result + bias.unsqueeze(0)
#     return result

@batorch_wrapper
def one_hot(k: bt.Tensor, num_classes=None):
    """
    one_hot(k, n) -> bt.Tensor
    Create a one-hot vector as a `bt.Tensor` with length `n` with all elements 0 except the `k`-th element being 1. 
    e.g. k == -1 gives tensor([0, 0, ..., 0, 1])
    """
    avouch(not k.is_floating_point(), "'one_hot' requires the index argument of integers. ")
    if num_classes is None: num_classes = k.max().item() + 1
    l = bt.zeros(k.shape + bt.func_dim_size(num_classes)).int(); l[*image_grid(*k.shape), k] = 1
    return l

@batorch_wrapper
def permute_space(data: bt.Tensor, *dims):
    """
    permute the space section in data. len(dims) should be data.n_space_dim. 
    """
    if len(dims) == 1 and isinstance(dims[0], (list, tuple)): dims = dims[0]
    avouch(len(dims) == data.n_space_dim, TypeError(f"'permute_space' only takes 'dims' of length data.n_space_dim={data.n_space_dim}, currently {dims}."))
    dimensions = list(range(data.space_start)) + [d + data.space_start for d in dims] + list(range(data.space_stop, data.n_dim))
    return data.permute(*dimensions)

@batorch_wrapper
def normalize01(data: bt.Tensor, func='linear'):
    """
    Normalize the data to range [0, 1]. Outliers are automatically excluded. 
        The result is continuous, even at the two ends.

    Args:
        data (bt.Tensor): The input data of size ({n_batch}, [n_feature], n_1, ..., n_r) [r=n_dim]
        func (str): The basic mapping function. Defaults to 'linear'.
            'linear': x => (x - x_min) / (x_max - x_min)
            'sigmoid': x => 1 / (1 + 19 ^ (2 * x - x_max - x_min) / (x_max - x_min))
    """
    if data.has_feature:
        norm = (data ** 2).sum([], keepdim=True).clamp(bt.eps).sqrt()
        m = norm.quantile(0.05, ..., keepdim=True)
        M = norm.quantile(0.95, ..., keepdim=True)
        new_norm = normalize01(norm.squeeze([]), func=func).view_as(norm)
        non_zero_norm = bt.where(norm == 0, bt.ones_like(norm), norm)
        return data * bt.where(new_norm == 0, bt.ones_like(new_norm), new_norm / non_zero_norm)
    else:
        m = data.quantile(0.05, ..., keepdim=True)
        M = data.quantile(0.95, ..., keepdim=True)
        clamped = bt.min(bt.max(data, m), M)
        if func.lower() == 'linear':
            return (clamped - m) / (M - m).clamp(bt.eps)
        elif func.lower() == 'sigmoid':
            return 1 / (bt.exp(- math.log(19) * (2 * clamped - M - m) / (M - m).clamp(bt.eps)) + 1)
        else: raise TypeError("Only 'linear' and 'sigmoid' functions are available for argument 'func' in 'normalize01'. ")

# # class Input:
    
# #     def __init__(self, **kwargs):
# #         """
# #         Reshape the tensors to the shape_exprs.

# #         Params:
# #             shape_expr (str): a string containing a shape expression. The expression is in shape format, 
# #                 except unknown number of dimensions can be specified with '...' and variable can be specified with words.
                
# #         Examples:
# #         ----------
# #         >>> a, b = bt.Input(n_dim=3)\
# #         ...          .set_shape(bt.zeros(3, 4, 5), "({n_batch}, [n_dim], n@1, ..., n@n_dim)")\
# #         ...          .set_shape(bt.zeros([2], 34), "({n_batch}, m_1)")\
# #         ...          .init()
# #         ...
# #         >>> n_batch, n_dim
# #         (2, 3)
# #         >>> a
# #         Tensor([[[[[...]]]]], shape=batorch.Size({2}, [3], 3, 4, 5))
# #         >>> b
# #         Tensor([[...]], shape=batorch.Size({2}, 34))
# #         >>> m_1, n_1, n_n_dim
# #         34, 3, 5
# #         """
# #         self.env_vars = get_environ_vars()
# #         self.vars = self.env_vars.locals
# #         self.vars.update(kwargs)
# #         self.env_vars.update(kwargs)
# #         self.data = []
# #         self.shape_exprs = []

# #     def set_shape(self, data, shape_expr: str):
# #         self.data.append(data)
# #         self.shape_exprs.append(shape_expr)
# #         return self

# #     def init(self):
# #         print(sum([get_snakenames(x) for x in self.shape_exprs], []))
# #         # for data, shape_list in zip(self.data, self.shape_exprs):
# #         #     old_shape = data.shape
# #         #     dim_pos = [0] * len(shape_list)
# #         #     pointer = 0
# #         #     for i, x in enumerate(shape_list):
# #         #         if x == '...': break
# #         #         if x.startswith('[') and x.endswith(']'): dim_pos[i] = old_shape.batch_dimension; x = x.strip('[]')
# #         #         if x.startswith('{') and x.endswith('}'): dim_pos[i] = old_shape.channel_dimension; x = x.strip('{}')
# #         #         if x.endswith(":optional"): x = x[:-len(":optional")]
# #         #         dim_pos[i] = pointer; pointer += 1
# #         #         if pointer in old_shape._special: pointer += 1
# #         #         if pointer in old_shape._special: pointer += 1 # two lines needed
# #         #         if is_dig
# #         #         x = get_snakename(x)
# #         #         old_value = self.vars.get(x)
# #         #         new_value = old_shape[dim_pos[i]]
            
# #         # return tuple(data_list)
# #         return tuple(self.data)
        
# # def as_shape(self, shape_expr: str):
# #     env_vars = get_environ_vars()
# #     loc_vars = env_vars.locals
# #     shape_list = [x.strip() for x in shape_expr.strip("()").split(',')]
# #     ibatch, ichannel = None, None
# #     old_shape = self.shape
# #     dim_pos = [0] * len(shape_list)
# #     pointer = 0
# #     for i, x in enumerate(shape_list):
# #         if x == '...': break
# #         if x.startswith('[') and x.endswith(']'): dim_pos[i] = old_shape.batch_dimension; x = x.strip('[]')
# #         if x.startswith('{') and x.endswith('}'): dim_pos[i] = old_shape.channel_dimension; x = x.strip('{}')
# #         if x.endswith(":optional"): x = x[:-len(":optional")]
# #         x = get_snakename(x)
# #         if x in loc_vars: x = loc_vars[x]

@alias("upper_triangular_matrix")
@batorch_wrapper
def upper_matrix(elements: bt.Tensor):
    '''
    elements: ({n_batch}, [n_dim * (n_dim + 1) / 2])
    output: ({n_batch}, [n_dim, n_dim])
    Compute the matrix with upper triangular elements of `elements`:
        For nxn matrix space (n_dim = n), the input is 𝜔:({n_batch}, [m = n(n+1)/2])
            The output is in shape ({n_batch}, [n, n]) of matrices:
                [[𝜔ₘ 𝜔ₘ₋₁ ... 𝜔ₘ₋ₙ₊₂  𝜔ₘ₋ₙ₊₁ ]
                 [0  𝜔ₘ₋ₙ ... 𝜔ₘ₋₂ₙ₊₄ 𝜔ₘ₋₂ₙ₊₃]
                 [︙  ︙   ⋱    ︙       ︙   ]
                 [0   0   ...   𝜔₃      𝜔₂   ]
                 [0   0   ...   0       𝜔₁   ]]
        ...
    '''
    avouch(elements.n_feature_dim == 1, TypeError("elements for 'upper_matrix' should be a tensor of shape ({n_batch}, [n_dim * (n_dim + 1) / 2])"))
    if not elements.has_batch: elements = elements.unsqueeze({})
    n_batch = elements.n_batch
    n_dim = int(math.sqrt(2 * elements.n_channel))
    elements = elements.flip([])
    output = bt.zeros({n_batch}, [n_dim, n_dim], dtype=elements.dtype, device=elements.device)
    row_len = bt.arange(n_dim, 0, -1)
    for i in range(n_dim): output[:, i, i:] = elements[:, row_len[:i].sum():row_len[:(i+1)].sum()]
    return output

@alias("lower_triangular_matrix")
@batorch_wrapper
def lower_matrix(elements: bt.Tensor):
    '''
    elements: ({n_batch}, [n_dim * (n_dim + 1) / 2])
    output: ({n_batch}, [n_dim, n_dim])
    Compute the matrix with lower triangular elements of `elements`:
        For nxn matrix space (n_dim = n), the input is 𝜔:({n_batch}, [m = n(n+1)/2])
            The output is in shape ({n_batch}, [n, n]) of matrices:
                [[ 𝜔ₘ       0     ...    0   0 ]
                 [𝜔ₘ₋₁    𝜔ₘ₋ₙ    ...    0   0 ]
                 [ ︙       ︙      ⋱   ︙   ︙ ]
                 [𝜔ₘ₋ₙ₊₂  𝜔ₘ₋₂ₙ₊₄  ...   𝜔₃  0  ]
                 [𝜔ₘ₋ₙ₊₁  𝜔ₘ₋₂ₙ₊₃  ...   𝜔₂  𝜔₁ ]]
        ...
    '''
    avouch(elements.n_feature_dim == 1, TypeError("elements for 'upper_matrix' should be a tensor of shape ({n_batch}, [n_dim * (n_dim + 1) / 2])"))
    if not elements.has_batch: elements = elements.unsqueeze({})
    n_batch = elements.n_batch
    n_dim = int(math.sqrt(2 * elements.n_channel))
    elements = elements.flip([])
    output = bt.zeros({n_batch}, [n_dim, n_dim], dtype=elements.dtype, device=elements.device)
    row_len = bt.arange(n_dim, 0, -1)
    for i in range(n_dim): output[:, i, i:] = elements[:, row_len[:i].sum():row_len[:(i+1)].sum()]
    return output.T

@alias("skew_symmetric")
@batorch_wrapper
def cross_matrix(axis: bt.Tensor):
    '''
    axis: ({n_batch}, [n_dim * (n_dim - 1) / 2])
    output: ({n_batch}, [n_dim, n_dim])
    Compute the skew symmetric matrix:
        For 2x2 matrix space (n_dim = 2), the input is 𝜗:({n_batch}, [1])
            The output is in shape ({n_batch}, [2, 2]) of matrices:
                [[0  𝜗]
                 [-𝜗 0]]
        For 3x3 matrix space (n_dim = 3), the input is 𝜔:({n_batch}, [3])
            The output is in shape ({n_batch}, [3, 3]) of matrices:
                [[ 0   𝜔₃ -𝜔₂]
                 [-𝜔₃  0   𝜔₁]
                 [ 𝜔₂ -𝜔₁  0 ]]
        For 4x4 matrix space (n_dim = 4), the input is 𝜔:({n_batch}, [6])
            The output is in shape ({n_batch}, [4, 4]) of matrices:
                [[ 0   𝜔₆ -𝜔₅  𝜔₄]
                 [-𝜔₆  0   𝜔₃ -𝜔₂]
                 [ 𝜔₅ -𝜔₃  0   𝜔₁]
                 [-𝜔₄  𝜔₂ -𝜔₁  0 ]]
        ...
    '''
    avouch(axis.n_feature_dim == 1, TypeError("axis for 'cross_matrix' should be a tensor of shape ({n_batch}, [n_dim * (n_dim - 1) / 2])"))
    if not axis.has_batch: axis = axis.unsqueeze({})
    n_batch = axis.n_batch
    n_dim = int(math.sqrt(2 * axis.n_channel)) + 1
    axis = axis.flip([])
    upper = bt.zeros({n_batch}, [n_dim, n_dim], dtype=axis.dtype, device=axis.device)
    row_len = bt.arange(n_dim - 1, 0, -1)
    row_sign = bt.where(row_len % 2 == 0, bt.ones_like(row_len), -bt.ones_like(row_len)).as_feature_tensor()
    for i in range(n_dim):
        upper[:, i, i+1:] = axis[:, row_len[:i].sum():row_len[:i+1].sum()] * row_sign[:(n_dim - i - 1)]
    return upper - upper.T

@alias("skew_symmetric_params")
@batorch_wrapper
def uncross_matrix(cross_matrix: bt.Tensor):
    '''
    cross_matrix: ({n_batch}, [n_dim, n_dim])
    output: ({n_batch}, [n_dim * (n_dim - 1) / 2])
    The inverse operation of cross_matrix. 
    '''
    avouch(cross_matrix.n_feature_dim == 2 or not cross_matrix.has_feature and cross_matrix.n_space_dim == 2, 
           TypeError("cross_matrix for 'cross_matrix' should be a tensor of shape ({n_batch}, [n_dim, n_dim])"))
    if not cross_matrix.has_batch: cross_matrix = cross_matrix.unsqueeze({})
    if not cross_matrix.has_feature:
        cross_matrix.sz_feature_dim_(cross_matrix.n_space_dim if cross_matrix.sz_sequence_dim < 0 else -cross_matrix.n_space_dim)
    avouch(cross_matrix.feature[0] == cross_matrix.feature[1], 
           TypeError("cross_matrix for 'cross_matrix' should be a tensor of shape ({n_batch}, [n_dim, n_dim])"))
    n_batch = cross_matrix.n_batch
    n_dim = cross_matrix.feature[0]
    axis = bt.zeros({n_batch}, [n_dim * (n_dim - 1) // 2], dtype=cross_matrix.dtype, device=cross_matrix.device)
    row_len = bt.arange(n_dim - 1, 0, -1)
    row_sign = bt.where(row_len % 2 == 0, bt.ones_like(row_len), -bt.ones_like(row_len)).as_feature_tensor()
    for i in range(n_dim):
        axis[:, row_len[:i].sum():row_len[:i+1].sum()] = cross_matrix[:, i, i+1:] * row_sign[:(n_dim - i - 1)]
    axis = axis.flip([])
    return axis

class summary:
    def __init__(self, *Xs, show_thumb = False, thumb_size = None, plot = False):
        self.smys = []
        self.show_as_plot = plot
        if plot: self.canvases = []
        if not hasattr(self, 'var_names'):
            args_exp = get_args_expression("summary")
            tokens = tokenize(args_exp, sep=', ')
            self.var_names = [vn.strip() for vn in tokens if len(tokenize(vn, sep='=')) == 1]
        avouch(len(self.var_names) <= len(Xs), f"Too many variable names recognized: {self.var_names}")
        starts_with_star = [i for i, x in enumerate(self.var_names) if x.startswith('*')]
        n_starts_with_star = 0
        if len(starts_with_star) > 1:
            avouch(None, f"Too much expandable variable in {self.var_names}: more than one. ")
        elif len(starts_with_star) == 0:
            avouch(len(self.var_names) == len(Xs), f"No expandable variable in {self.var_names} but detected more input arguments. ")
        else:
            starts_with_star = starts_with_star[0]
            n_starts_with_star = len(Xs) - len(self.var_names) + 1
        
        if thumb_size is None: thumb_size = 4
        else: show_thumb = True
        self.show_thumb = show_thumb
        self.thumb_size = thumb_size
        
        for i, x in enumerate(Xs):
            if n_starts_with_star > 0:
                if 0 <= i - starts_with_star < n_starts_with_star: j = i - starts_with_star; i = starts_with_star
                elif i >= starts_with_star + n_starts_with_star: i = i - n_starts_with_star
            var_name = self.var_names[i] + (f'[{j}]' if i == starts_with_star else '')
            class_name = x.__class__.__name__
            g = None
            if not isinstance(x, bt.torch.Tensor):
                l = touch(lambda: len(x), default=0)
                if l <= 10:
                    self.smys.append(f">>{var_name}<< " + str(x) + '\n')
                    continue
                else: x = bt.tensor(x)
            elif not isinstance(x, bt.Tensor):
                if x.is_leaf and x.grad is not None: g = bt.tensor(x.grad.clone().detach())
                x = bt.tensor(x)
            special_dim_str = f"[{x.n_special_dim}]+" if x.n_special_dim > 0 else ''
            self.out = SPrint(f">>Summary of {var_name}<<\n")
            def size_str(Size=None, Dimension=None, Class=None):
                return f"{f'[{Size.n_special_dim}]+' if Size.n_special_dim > 0 else ''}{Size.n_space_dim}D {Class.__name__} of size {self.__raw_str__(Size)}"
            self.__add_prop__('Size', value2str_func = size_str, Size=x.shape, Dimension=x.n_dim, Class=x.__class__)
            self.__add_prop__('Device', x.device, lambda d: 'cpu' if d.type == 'cpu' else f"{d.type}:{d.index}")
            self.__add_prop__('Dtype', x.dtype, self.__raw_str__)
            self.__add_prop__('Requires_Gradient', x.requires_grad)
            if x.requires_grad: self.__add_prop__('Gradient_Func', x.grad_fn)

            self.show_grad = True
            self.__display_values__(x)
            if self.show_grad and g is not None:
                self.__add_prop__('Gradient_Size', g.shape, self.__raw_str__)
                self.__display_values__(x, is_grad=True)
            self.smys.append(self.out.text)
            if plot:
                while x.n_dim > 2:
                    if x.has_special:
                        d_sample = x.special_dims[0]
                    else: d_sample = argmin(list(x.shape))[0]
                    x = x.pick(x.size(d_sample) // 2, d_sample)
                self.canvases.append(x.detach())
                
    def __add_prop__(self, name, value=null, value2str_func=None, prompt=None, display_only=False, **assign_values):
        if prompt is not None: prompt = f" ({prompt})"
        else: prompt = ''
        if value is not null:
            if value2str_func is not None: value_str = value2str_func(value)
            else: value_str = str(value)
            if not display_only: setattr(self, name, value)
        else:
            avouch(value2str_func is not None, TypeError("'value2str_func' should be specified when multiple values are given for 'summary.__add_prop__'. "))
            value_str = value2str_func(**assign_values)
            if not display_only:
                for n, v in assign_values.items(): setattr(self, n, v)
        tag = f"【{name}{prompt}】"
        n_align = math.ceil((len(tag) + 2) / 25)
        self.out(tag + ' ' * (n_align * 25 - len(tag) - 2) + value_str)
             
    def __raw_str__(self, x):
        sx = str(x)
        if 'Size(' in sx: return sx.split('Size')[-1]
        return tokenize(tokenize(sx.split('(', 1)[-1], sep=[',', ')'])[0], sep='.')[-1]

    def __approx_quantile__(self, x):
        valid_dims = bt.tensor(x.shape).cumprod(0) >> 24 < 1
        if all(valid_dims): x = x.flatten()
        elif all(~valid_dims):
            n_replicate = (x.n_ele >> 24) + 1
            n_sample = x.n_ele // n_replicate
            x = x[:n_replicate * n_sample].reshape(n_sample, n_replicate)
        else: x = x.flatten(0, valid_dims.argmin().item()-1)
        x = x.float().quantile(bt.arange(6).float() / 5, dim=0)
        if x.n_dim > 1: x = x.mean(*range(1, x.n_dim))
        return x
    
    def __display_values__(self, x, is_grad=False):
        nan_mask = bt.isnan(x)
        num_nans = nan_mask.sum()
        inf_mask = bt.isinf(x)
        num_infs = inf_mask.sum()
        nan_inf_mask = nan_mask | inf_mask
        valid_nums = x[~nan_inf_mask]
        if num_nans + num_infs == 0: show_error = False; show_valid = True
        elif (num_nans + num_infs) / x.n_ele < 0.5: show_error = True; show_valid = True
        else: show_error = True; show_valid = False; self.show_grad = False
        if show_valid:
            values_str = lambda vs: f"(min:{vs[0]}, med:{vs[1]}, max:{vs[2]})"
            if x.is_complex():
                vre_name, vim_name, name = 'Values_Re', 'Values_Im', 'Values_Re/Im'
                if is_grad: vre_name, vim_name, name = ['Gradient_' + x for x in (vre_name, vim_name, name)]
                vre_val = (valid_nums.real.min(), valid_nums.real.median(), valid_nums.real.max())
                vim_val = (valid_nums.imag.min(), valid_nums.imag.median(), valid_nums.imag.max())
                self.__add_prop__(name, value2str_func=lambda **k: f"Re{values_str(k[vre_name])}, Im{values_str(k[vim_name])}", **{vre_name: vre_val, vim_name: vim_val})
            else:
                name = 'Gradient_Values' if is_grad else 'Values'
                val = (valid_nums.min(), valid_nums.median(), valid_nums.max())
                self.__add_prop__(name, val, values_str)
        if show_error:
            if num_nans > 0:
                nan_name = f"{num_nans} NaN"
                if num_nans > 1: nan_name += 's'
                has_nan_maps = nan_mask.sum(...) > 0
                num_nan_maps = has_nan_maps.sum()
                nan_mask = nan_mask[has_nan_maps, ...][0]
                nan_center = bt.mask_center(nan_mask)
                if is_grad:
                    self.Gradient_NaN = self.Gradient_NaNs = num_nans
                else:
                    self.NaN = self.NaNs = num_nans
                    self.NaN_maps = num_nan_maps
                    self.NaN_center = tuple(nan_center)
                self.__add_prop__(nan_name, value2str_func=lambda **_: f"{num_nan_maps}/{has_nan_maps.n_ele} maps with nan, e.g. center: {self.__raw_str__(nan_center)}", display_only=True)
            if num_infs > 0:
                inf_name = f"{num_infs} Inf"
                if num_infs > 1: inf_name += 's'
                has_inf_maps = inf_mask.sum(...) > 0
                num_inf_maps = has_inf_maps.sum()
                inf_mask = inf_mask[has_inf_maps, ...][0]
                inf_center = bt.mask_center(inf_mask)
                if is_grad:
                    self.Gradient_Inf = self.Gradient_Infs = num_infs
                else:
                    self.Inf = self.Infs = num_infs
                    self.Inf_maps = num_inf_maps
                    self.Inf_center = tuple(inf_center)
                self.__add_prop__(inf_name, value2str_func=lambda **_: f"{num_inf_maps}/{has_inf_maps.n_ele} maps with inf, e.g. center: {self.__raw_str__(inf_center)}", display_only=True)
        if show_valid:
            approx_q = x.n_ele >> 24 > 0
            self.__add_prop__('Gradient_Quantiles' if is_grad else 'Quantiles', self.__approx_quantile__(valid_nums), 
                              lambda q: ('\n >> ' if approx_q else '') + f"([0%]{q[0]}, [20%]{q[1]}, [40%]{q[2]}, [60%]{q[3]}, [80%]{q[4]}, [100%]{q[5]})", 
                              prompt='approx. due to too many elements' if approx_q else None)
            values = valid_nums.unique()
            def range_str(values):
                if values.n_ele <= 10:
                    if values.n_ele > 6: return f"\n >> [{', '.join(str(x) for x in values)}]"
                    else:
                        num_pairs = [(v, (x == v).sum().item()) for v in values]
                        return f"\n >> [{', '.join(f'{v}({n})' for v, n in num_pairs)}]"
                else: return f"\n >> [{', '.join(str(x) for x in values[:5])}, ..., {', '.join(str(x) for x in values[-5:])}]"
            self.__add_prop__('Gradient_Range' if is_grad else 'Range', values, value2str_func=range_str, 
                              prompt=f"{values.n_ele} unique element" + ('s' if values.n_ele > 1 else ''))
        if self.show_thumb:
            gaps = bt.floor((bt.channel_tensor(list(x.space)) - 1) / (self.thumb_size - 1)).clamp(min=1).long()
            down_sample = x[(slice(None),) * x.space_start + (image_grid(*(self.thumb_size if s > self.thumb_size else max(s // g, 1) for s, g in zip(x.space, gaps))) * gaps).split(1, dim=[], squeezedim=True)]
            for i in down_sample.special_dims[::-1]: down_sample = down_sample.pick(0, dim=i)
            self.__add_prop__('Gradient_Thumb' if is_grad else 'Thumb', down_sample, 
                              value2str_func=lambda x: '\n >> ' + bt.Tensor.__shift_repr__(x.__raw_repr__(), 4), 
                              prompt=f"downsampled to {self.thumb_size}{f' x {self.thumb_size}'*(x.n_space_dim-1)}")
    
    def show(self):
        if not self.show_as_plot:
            print(self)
            return
        plt.figure(figsize=(15, 7))
        
        n_plots = len(self.canvases)
        self.subplot_rows = math.floor(math.sqrt(n_plots))
        self.subplot_cols = math.ceil(n_plots / self.subplot_rows)

        old_minus = plt.rcParams['axes.unicode_minus']
        old_font = plt.rcParams['font.sans-serif']
        font = 'Heiti TC'
        if font not in plt.rcParams['font.sans-serif']:
            plt.rcParams['font.sans-serif'].insert(0, font)
        plt.rcParams['axes.unicode_minus'] = False
        for i, x in enumerate(self.canvases):
            ax = plt.subplot(self.subplot_rows, self.subplot_cols, i+1)
            old_size = plt.rcParams['font.size']
            plt.rcParams['font.size'] = 6
            if x.n_dim == 2:
                img = plt.imshow(x, cmap=plt.cm.gray)
                plt.colorbar(mappable=img, ax=ax, orientation='vertical')
                plt.text(-0.5, -0.5, self.smys[i])
            elif x.n_dim == 1:
                plt.plot(bt.arange(len(x)), x)
                plt.text(0, plt.ylim()[1], self.smys[i])
            else: raise TypeError("Failed to display 0-dimensional data in plot.")
            plt.rcParams['font.size'] = old_size
        plt.rcParams['font.sans-serif'] = old_font
        plt.rcParams['axes.unicode_minus'] = old_minus 
        plt.show()

    @alias("__repr__")
    def __str__(self): return ''.join(self.smys)

class display(summary):
    def __init__(self, *args, **kwargs):
        kwargs['plot'] = True
        args_exp = get_args_expression("display")
        tokens = tokenize(args_exp, sep=', ')
        self.var_names = [vn.strip() for vn in tokens if len(tokenize(vn, sep='=')) == 1]
        return super().__init__(*args, **kwargs)
