import numpy as np
class SoftmaxCrossEntropy:
     def __init__ (self):
        self.loss = None # Loss
        self.y = None # Output of softmax
        self.t = None # Label data (one-hot vector)
        
        
     def softmax(x):
        x = x - np.max(x, axis=-1, keepdims=True)
        return np.exp(x) / np.sum(np.exp(x), axis=-1, keepdims=True)
    
     def cross_entropy_error(self,y, t):
         if y.ndim == 1:
            t = t.reshape(1, t.size)
            y = y.reshape(1, y.size)
         if t.size == y.size:
            t = t.argmax(axis=1)
         batch_size = y.shape[0]
         return -np.sum(np.log(y[np.arange(batch_size), t] + 1e-7)) / batch_size
        
    
     def forward(self, x, t):
        self.t = t
        self.y = self.softmax(x)
        self.loss = self.cross_entropy_error(self.y, self.t)
        return self.loss

     def backward(self, dout=1):
        batch_size = self.t.shape[0]
        dx = (self.y - self.t) / batch_size
        return dx
    
     