# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_repr_rgb.ipynb.

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

# %% ../nbs/01_repr_rgb.ipynb 5
from typing import Any, Optional as O

from matplotlib import axes, figure
from IPython.core.pylabtools import print_figure
from PIL import Image
import torch

from lovely_numpy.utils.utils import cached_property
from lovely_numpy.utils.pad import pad_frame_gutters
from lovely_numpy.utils.tile2d import hypertile
from lovely_numpy.repr_rgb import fig_rgb
from lovely_numpy import config as np_config

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


# %% ../nbs/01_repr_rgb.ipynb 6
# This is here for the monkey-patched tensor use case.

# I want to be able to call both `tensor.rgb` and `tensor.rgb(stats)`. For the
# first case, the class defines `_repr_png_` to send the image to Jupyter. For
# the later case, it defines __call__, which accps the argument.

class RGBProxy():
    """Flexible `PIL.Image.Image` wrapper"""
    
    def __init__(self, t:torch.Tensor):
        assert t.ndim >= 3, f"Expecting at least 3 dimensions, got shape{t.shape}={t.dim()}"
        self.t =t
        self.params = dict(denorm   = None,
                            cl         = False,
                            gutter_px  = 3,     
                            frame_px   = 1,
                            scale      = 1,
                            view_width = 966,
                            ax = None)


    def __call__(self,
                denorm      :Any    =None,
                cl          :Any    =False,
                gutter_px   :O[int] =None,
                frame_px    :O[int] =None,
                scale       :O[int] =None,
                view_width  :O[int] =None,
                ax          :O[axes.Axes]=None):
        
        self.params.update( { k:v for
                            k,v in locals().items()
                            if k != "self" and v is not None } )
        _ = self.fig # Trigger figure generation
        return self

    @cached_property
    def fig(self) -> figure.Figure:
        cfg = get_config()
        with np_config(fig_close=cfg.fig_close, fig_show=cfg.fig_show):
            return fig_rgb(to_numpy(self.t), **self.params)

    def _repr_png_(self):
        return print_figure(self.fig, fmt="png", pad_inches=0,
            metadata={"Software": "Matplotlib, https://matplotlib.org/"})


# %% ../nbs/01_repr_rgb.ipynb 7
def rgb(x           :torch.Tensor,  # Tensor to display. [[...], C,H,W] or [[...], H,W,C]
        denorm      :Any =None,     # Reverse per-channel normalizatoin
        cl          :Any =False,    # Channel-last
        gutter_px   :int =3,        # If more than one tensor -> tile with this gutter width
        frame_px    :int =1,        # If more than one tensor -> tile with this frame width
        scale       :int =1,        # Scale up. Can't scale down.
        view_width  :int =966,      # target width of the image
        ax          :O[axes.Axes] =None # Use this Axes
        ) -> RGBProxy:

    args = locals()
    del args["x"]
    
    return RGBProxy(x)(**args)
