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

# In[1]:


import numpy as np


# # softmax 
# - $softmax(x_i) = \frac{e^{x_i}}{\sum_j{e^{x_j}}}$
# 
# - $\begin{align} (softmax(x + c))_{i}= \frac{e^{x_{i} + c}}{\sum_{j} e^{x_{j} + c}} = \frac{e^{x_{i}} \times e^{c}}{e^{c} \times \sum_{j} e^{x_{j}}} \\ = \frac{e^{x_{i}} \times {e^{c}}}{{e^{c}} \times \sum_{j} e^{x_{j}}} = (softmax(x))_{i}  \end{align}$
# 
# so:
# 
# - $softmax(x) = softmax(x + c)$

# In[2]:


def softmax(x):
    orig_shape = x.shape

    if len(x.shape) > 1:
        # Matrix
        ### YOUR CODE HERE
        x_max = np.max(x, axis=1).reshape(x.shape[0], 1)
        x -= x_max
        exp_sum = np.sum(np.exp(x), axis=1).reshape(x.shape[0], 1)
        x = np.exp(x) / exp_sum 
        ### END YOUR CODE
    else:
        # Vector
        ### YOUR CODE HERE
        x_max = np.max(x)
        x -= x_max
        exp_sum = np.sum(np.exp(x))
        x = np.exp(x) / exp_sum
        ### END YOUR CODE
        #or:  x = (np.exp(x)/sum(np.exp(x)))   

    assert x.shape == orig_shape
    return x


# In[3]:


def test_softmax_basic():
    """
    Some simple tests to get you started.
    Warning: these are not exhaustive.
    """
    print("Running basic tests...")
    test1 = softmax(np.array([1,2]))
    print(test1)
    ans1 = np.array([0.26894142,  0.73105858])
    assert np.allclose(test1, ans1, rtol=1e-05, atol=1e-06)

    test2 = softmax(np.array([[1001,1002],[3,4]]))
    print(test2)
    ans2 = np.array([
        [0.26894142, 0.73105858],
        [0.26894142, 0.73105858]])
    assert np.allclose(test2, ans2, rtol=1e-05, atol=1e-06)

    test3 = softmax(np.array([[-1001,-1002]]))
    print(test3)
    ans3 = np.array([0.73105858, 0.26894142])
    assert np.allclose(test3, ans3, rtol=1e-05, atol=1e-06)

    print("You should be able to verify these results by hand!\n")

if __name__ == "__main__":
    test_softmax_basic()


# In[4]:


x = np.array([[1,2],[4,3]])


# In[5]:


np.max(x, axis=1)


# In[ ]:





# In[ ]:





# In[ ]:





# In[ ]: