import numpy as np
import matplotlib.pyplot as plt
import torch

class 더하기:
    def __init__(self):
        self.x = None
        self.y = None
        
    def __call__(self, x, y):
        """순전파"""
        self.x = x
        self.y = y
        z = x + y
        return z
    
    def backward(self, dz):
        """역전파"""
        dx = dz * 1
        dy = dz * 1
        return dx, dy


class 곱하기:
    def __init__(self):
        self.x = None
        self.y = None
        
    def __call__(self, x, y):
        """순전파"""
        self.x = x
        self.y = y
        z = x * y
        return z
    
    def backward(self, dz):
        """역전파"""
        dx = dz * self.y
        dy = dz * self.x
        return dx, dy
    
class Sum:
    def __init__(self):
        self.input_shape = None
        self.axis = None

    def __call__(self, xs, axis=None):
        """순전파"""
        self.input_shape = xs.shape
        self.axis = axis
        z = np.sum(xs, axis=axis)
        return z

    def backward(self, dz):
        """역전파"""
        # dz를 배열로 변환 (스칼라 처리 포함)
        dz = np.asarray(dz)

        # 출력 형상 계산
        if self.axis is not None:
            if isinstance(self.axis, int):
                output_shape = tuple(dim for i, dim in enumerate(self.input_shape) if i != self.axis)
            else:  # 다중 축 처리
                reduced_shape = list(self.input_shape)
                for ax in sorted(self.axis):
                    reduced_shape[ax] = 1
                output_shape = tuple(dim for dim in reduced_shape if dim != 1)
        else:
            # axis=None (전체 합산)
            output_shape = ()

        # dz의 형상 검증
        if dz.shape != output_shape and dz.shape != ():
            raise ValueError(f"dz의 크기 {dz.shape}가 출력 형상 {output_shape}와 일치하지 않습니다.")

        # dz를 입력 형상으로 확장
        if self.axis is not None:
            dz_broadcasted = np.expand_dims(dz, axis=self.axis)
        else:
            dz_broadcasted = dz
        dxs = dz_broadcasted * np.ones(self.input_shape)
        return dxs


# 구현 확인
x = np.array([[1, 2, 3], [4, 5, 6]])
sum = Sum()

# 전체 합산 (axis=None)
assert sum(x) == 21
assert np.all(sum.backward(0.1) == 0.1 * np.ones_like(x))

# 축=0
assert np.all(sum(x, axis=0) == [5, 7, 9])
assert np.all(sum.backward([0.1, 0.2, 0.3]) == [[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]])

# 축=1
assert np.all(sum(x, axis=1) == [6, 15])
assert np.all(sum.backward([0.1, 0.2]) == [[0.1, 0.1, 0.1], [0.2, 0.2, 0.2]])

# 다중 축 (axis=(0, 1))
assert sum(x, axis=(0, 1)) == 21
assert np.all(sum.backward(0.1) == 0.1 * np.ones_like(x))

x = torch.tensor(x.astype('float32'), requires_grad=True)
y = torch.sum(x)
print(y)
y.backward(torch.ones_like(y))
print(x.grad)
x.grad.zero_()

y = torch.sum(x, axis=0)
y.backward(torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32))
print(x.grad.numpy().round(1))
x.grad.zero_()

y = torch.sum(x, axis=1)
y.backward(torch.tensor([0.1, 0.2], dtype=torch.float32))
print(x.grad.numpy().round(1))

곱1 = 곱하기()
곱2 = 곱하기()

사과단가, 사과수량, 부가세 = 100, 2, 1.1

사과가격 = 곱1.forward(사과단가, 사과수량)
사과가격

가격 = 곱2.forward(사과가격, 부가세)
가격

d사과가격, d부가세 = 곱2.backward(1)
d사과가격, d부가세

d사과단가, d사과수량 = 곱1.backward(d사과가격)
d사과단가, d사과수량

class 더하기:
    def __call__(self, x, y):
        return x + y

    def backward(self, dout):
        """역전파
        dout: 상류에서 전해지는 미분값
        """
        dx = dout * 1
        dy = dout * 1
        return dx, dy
    
class 곱하기:
    def __init__(self):
        self.x = None
        self.y = None

    def __call__(self, x, y):
        self.x = x
        self.y = y
        return x * y
    
    def backward(self, dout):
        """역전파
        dout: 상류에서 전해지는 미분값
        """
        dx = dout * self.y
        dy = dout * self.x
        return dx, dy

class Inverse:
    def __init__(self):
        self.x = None

    def __call__(self, x):
        self.x = x
        return 1 / x
    
    def backward(self, dout):
        dx = dout * (-1 / self.x ** 2)
        return dx
    
class 나누기:
    def __init__(self):
        self.곱하기 = 곱하기()
        self.역수 = Inverse()

    def __call__(self, x, y):
        x = self.역수(x)
        out = self.곱하기(x, y)
        return out
    
    def backward(self, dout):
        dx, dy = self.곱하기.backward(dout)
        dx = self.역수.backward(dx)
        return dx, dy
    
class Exponetial:
    def __init__(self):
        self.out = None

    def __call__(self, x):
        self.out = np.exp(x)
        return self.out

    def backward(self, dout):
        dx = dout * self.out
        return dx

class Sigmoid:
    def __init__(self):
        # y = 1 / (1 + exp(-x))
        self.mul = 곱하기()
        self.exp = Exponetial()
        self.add = 더하기()
        self.inv = Inverse()

    def __call__(self, x):
        x = self.mul(x, -1) # -x
        x = self.exp(x) # exp(-x)
        x = self.add(x, 1) # 1 + exp(-x)
        y = self.inv(x) # 1 / (1 + exp(-x))
        return y
    
    def backward(self, dout):
        dx = self.inv.backward(dout)
        dx, _ = self.add.backward(dx)
        dx = self.exp.backward(dx)
        dx, _ = self.mul.backward(dx)
        return dx

sigmoid = Sigmoid()
x = 0.0
y = sigmoid(x)
dydx = sigmoid.backward(1.0)
# 순전파 검증
assert np.isclose(y, 1 / (1 + np.exp(-x)))
# 해석적 미분과 역전파가 일치하는지 확인
assert np.isclose(dydx, y * (1 - y))

class Sigmoid:
    def __init__(self):
        self.y = None

    def __call__(self, x):
        return self.forward(x)

    def forward(self, x):
        y = 1 / (1 + np.exp(-x))
        self.y = y
        return y

    def backward(self, dout):
        dx = dout * self.y * (1 - self.y)
        return dx

class 완전연결:
    def __init__(self, 입력수, 출력수, 활성화=None):
        self.W = np.random.randn(입력수, 출력수)
        self.b = np.zeros(출력수)
        self.activation = 활성화
        self.x = None
        self.dW = None
        self.db = None

    def __call__(self, x):
        return self.forward(x)

    def forward(self, x):
        self.x = x
        z = np.dot(x, self.W) + self.b
        if self.activation:
            return self.activation(z)
        return z

    def backward(self, dout):
        if self.activation:
            dout = self.activation.backward(dout)
            
        self.dW = np.dot(self.x.T, dout)
        # 배치 단위 연산 시, 미분값을 모든 표본에 대해 더합니다.
        self.db = np.sum(dout, axis=0)
        
        dx = np.dot(dout, self.W.T)
        return dx

def softmax(z):
    if z.ndim == 1:
        z = z.reshape(1, -1)
    exp_z = np.exp(z - np.max(z, axis=1).reshape(-1, 1))
    return exp_z / np.sum(exp_z, axis=1).reshape(-1, 1)

def 교차엔트로피오차(y, y_pred):
    delta = 1e-7
    배치크기 = y.shape[0]
    return -np.sum(y * np.log(y_pred + delta)) / 배치크기

class CrossEntropy:
    def __init__(self, from_logits=False):
        self.y = None
        self.from_logits = from_logits
        self.proba = None

    def __call__(self, z, y):
        return self.forward(z, y)

    def forward(self, outputs, y):
        self.y = y
        # z -> softmax -> proba
        if self.from_logits:            
            self.proba = outputs
        else:
            self.proba = softmax(outputs)
        # CEE(y, proba)
        손실 = 교차엔트로피오차(y, self.proba)
        return 손실

    def backward(self, dout=1):
        배치크기 = len(self.y)
        dz = self.proba - self.y
        return dz / 배치크기

class 역전파신경망:
    def __init__(self, 손실함수):
        self.layers = []
        self.loss_func = 손실함수

    def add(self, layer):
        self.layers.append(layer)

    def __call__(self, x):
        """순전파"""
        outputs = x
        for layer in self.layers:
            outputs = layer(outputs)
        return outputs # z_last

    def 손실산출(self, x, y):
        outputs = self(x)
        손실 = self.loss_func(outputs, y)
        return 손실

    def fit(self, x, y, 배치크기, 에폭수, 학습률):
        에폭당_배치수 = len(x) // 배치크기
        학습횟수 = 에폭당_배치수 * 에폭수
        print(f'배치크기={배치크기}, 에폭수={에폭수}, 학습횟수={학습횟수}({에폭당_배치수}/에폭)')
        손실변화 = []
        for 학습 in range(학습횟수):
            # 1. 미니 배치
            표본수 = S = len(x)
            배치색인 = np.random.choice(표본수, 배치크기)
            x_batch = x[배치색인]
            y_batch = y[배치색인]
            # 2. 경사 산출 (역전파)
            # 1) 순전파
            손실 = self.손실산출(x_batch, y_batch)
            손실변화.append(손실)
            # 2) 역전파            
            dout = self.loss_func.backward(1)
            for layer in reversed(self.layers):
                dout = layer.backward(dout)
            # 3. 매개변수 갱신 (경사 하강)
            for layer in self.layers:
                if isinstance(layer, 완전연결):
                    layer.W -= layer.dW * 학습률
                    layer.b -= layer.db * 학습률

            if 학습 == 0 or (학습 + 1) % 100 == 0:
                print(f'[학습 {학습 + 1}] Loss: {손실:.3f}')
        return 손실변화

import pickle

파일경로 = '../data/mnist/mnist_ndarray.pkl'
with open(파일경로, 'rb') as 파일:
    mnist_data = pickle.load(파일)

(train_images, train_labels), (test_images, test_labels) = mnist_data

print(train_images.shape, test_images.shape)
assert train_images.shape[1:] == test_images.shape[1:]
assert len(train_images) == len(train_labels)
assert len(test_images) == len(test_labels)

def 전처리(images):
    X = images.reshape(-1, 28 * 28)
    X = X.astype('float32')
    X /= 255
    return X

X_train = 전처리(train_images)
X_test = 전처리(test_images)
print(X_train.shape, X_test.shape)
print(X_train.max(), X_test.max())

Y_train = np.eye(10)[train_labels]
Y_train.shape

FC = 완전연결

model = 역전파신경망(손실함수=CrossEntropy())
model.add(FC(784, 50, Sigmoid()))
# model.add(Sigmoid())
model.add(FC(50, 100, Sigmoid()))
# model.add(Sigmoid())
model.add(FC(100, 10))

outputs = model(X_test)
print(f'훈련 전 성능: {np.mean(test_labels == np.argmax(outputs, 1)):.2%}')

손실변화 = model.fit(X_train, Y_train, 배치크기=100, 에폭수=10, 학습률=1.0)

outputs = model(X_test)
print(outputs.shape)
예측 = np.argmax(outputs, axis=1)
Y_test = np.eye(10)[test_labels]
손실 = CrossEntropy()(outputs, Y_test)
정확도 = np.mean(test_labels == 예측)
print(f'[시험] 손실: {손실:.3f}, 정확도: {정확도:.2%}')

# 에폭마다 손실 표시
에폭별_손실 = 손실변화[::600]
plt.plot(에폭별_손실, 'go--')
plt.show()