#!/usr/bin/env python
# coding: utf-8

# # MatRepr TensorFlow
# 
# Neatly format dense `tf.Tensor` and sparse `tf.SparseTensor`.
# 
# 1D dense and sparse tensors are formatted as row vectors.
# 
# 2D dense and sparse tensors are formatted as matrices.
# 
# 3D+ sparse tensors are formatted as a list of index/value tuples. Dense 3D+ tensors rendered with native TensorFlow formatter.

# In[1]:


import tensorflow as tf

tf.random.set_seed(1234)

# so matrepr can be imported from the source tree.
import sys
sys.path.insert(0, '..')

from matrepr import mdisplay, mprint


# In[2]:


scalar = tf.constant(5)

dense1D = tf.random.uniform(shape=(5000,))

rand2D = tf.random.uniform(shape=(128, 64)).numpy()
rand2D[rand2D < 0.6] = 0
rand2D_sparse = tf.sparse.from_dense(tf.convert_to_tensor(rand2D))

small3D = tf.constant([[[1., 0], [2., 3.]], [[4., 0], [5., 6.]]])
coo3D = tf.sparse.from_dense(small3D)


# ## TensorFlow default formatting

# In[3]:


dense1D


# In[4]:


rand2D_sparse


# In[5]:


coo3D


# In[6]:


scalar


# ## MatRepr default
# 
# Load the MatRepr Jupyter extension with `%load_ext matrepr` to render tensors with MatRepr by default.
# 
# A single-use alternative is to use `matrepr.mdisplay()`. For console use `matrepr.mprint()`. 

# In[7]:


get_ipython().run_line_magic('load_ext', 'matrepr')


# In[8]:


dense1D


# In[9]:


rand2D_sparse


# In[10]:


coo3D


# In[11]:


scalar


# # Labels
# 
# Specify title, row and/or column labels to help the reader quickly understand what they are looking at.

# In[12]:


dense = tf.random.uniform(shape=(64, 32)).numpy()
dense[dense < 0.3] = 0
tensor = tf.sparse.from_dense(tf.convert_to_tensor(dense))

obs_labels = [f"observation {i}" for i in range(tensor.shape[0])]  # list of labels
feature_labels = {i: f"feature {i+1}" for i in range(tensor.shape[1])}  # map index to label works too


# In[13]:


mdisplay(tensor, row_labels=obs_labels, col_labels=feature_labels, title="Random Dataset")


# ## LaTeX

# In[14]:


mdisplay(rand2D_sparse, "latex")


# ## String

# In[15]:


mprint(rand2D_sparse)


# ## More compact size

# In[16]:


import matrepr
matrepr.params.max_rows = 10
matrepr.params.max_cols = 7
matrepr.params.num_after_dots = 0


# In[17]:


rand2D_sparse


# In[17]: