import numpy as np
from PIL.Image import fromarray
from IPython import get_ipython
from IPython.display import display_html


def display_np_arrays():
    def np_to_png(a):
        if 2 <= len(a.shape) <= 3:
            return fromarray(
                np.array(np.clip(a, 0, 1) * 255, dtype="uint8")
            )._repr_png_()
        else:
            return fromarray(np.zeros([1, 1], dtype="uint8"))._repr_png_()

    def np_to_text(obj, p, cycle):
        if len(obj.shape) < 2:
            print(repr(obj))
        if 2 <= len(obj.shape) <= 3:
            pass
        else:
            print("<array of shape {}>".format(obj.shape))

    get_ipython().display_formatter.formatters["image/png"].for_type(
        np.ndarray, np_to_png
    )
    get_ipython().display_formatter.formatters["text/plain"].for_type(
        np.ndarray, np_to_text
    )


_style_inline = """<style>
.sparrow-answer {
    color: transparent;
    padding: 5px 15px;
    background-color: #def;
}
.sparrow-answer:hover { color: blue; } 
</style>
"""


def guess(x):
    display_html(
        _style_inline
        + "<h4>Answer is: <span class='einops-answer'>{x}</span> (hover to see)</h4>".format(
            x=tuple(x)
        ),
        raw=True,
    )
