import numpy as np

class SirenLayer:
    """
    SIREN의 단일 레이어.
    Forward: 선형 변환 -> Sine 활성화
    Backward: Chain Rule을 이용한 기울기 계산
    """
    def __init__(self, in_features, out_features, w0=1.0, is_first=False):
        self.in_features = in_features
        self.out_features = out_features
        self.w0 = w0
        
        # 1. 초기화 (Sitzmann's Initialization)
        if is_first:
            limit = 1.0 / in_features
            self.W = np.random.uniform(-limit, limit, (in_features, out_features))
        else:
            limit = np.sqrt(6.0 / in_features) / w0
            self.W = np.random.uniform(-limit, limit, (in_features, out_features))
            
        self.b = np.zeros((1, out_features))
        
        # 역전파를 위한 캐시 (입력값, 선형변환 결과값 저장)
        self.x_cache = None
        self.z_cache = None # z = Wx + b

    def forward(self, x):
        self.x_cache = x
        # 선형 변환: Z = XW + b
        self.z_cache = np.dot(x, self.W) + self.b
        # 활성화: A = sin(w0 * Z)
        return np.sin(self.w0 * self.z_cache)

    def backward(self, grad_output, learning_rate):
        """
        grad_output: 다음 레이어에서 넘어온 기울기 (dL/dA)
        """
        # 1. 활성화 함수 미분 (Chain Rule)
        # y = sin(w0 * z)  -->  dy/dz = w0 * cos(w0 * z)
        # dL/dz = dL/dy * dy/dz
        grad_activation = self.w0 * np.cos(self.w0 * self.z_cache)
        delta = grad_output * grad_activation # Element-wise multiplication
        
        # 2. 가중치(W) 기울기 계산
        # z = xW + b  -->  dz/dW = x
        # dL/dW = x.T * delta
        grad_W = np.dot(self.x_cache.T, delta)
        
        # 3. 편향(b) 기울기 계산
        grad_b = np.sum(delta, axis=0, keepdims=True)
        
        # 4. 입력(x)에 대한 기울기 계산 (이전 레이어로 전파)
        # z = xW + b  -->  dz/dx = W
        # dL/dx = delta * W.T
        grad_input = np.dot(delta, self.W.T)
        
        # 5. 가중치 업데이트 (Gradient Descent)
        self.W -= learning_rate * grad_W
        self.b -= learning_rate * grad_b
        
        return grad_input

class SirenNetwork:
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers=3, w0=30.0):
        self.layers = []
        # Input Layer (w0=30, First)
        self.layers.append(SirenLayer(input_dim, hidden_dim, w0=w0, is_first=True))
        
        # Hidden Layers (w0=1, Not First)
        for _ in range(num_layers - 2):
            self.layers.append(SirenLayer(hidden_dim, hidden_dim, w0=1.0, is_first=False))
            
        # Output Layer (Linear - 활성화 함수 처리는 backward 구조상 SirenLayer를 변형해 사용하거나 Linear만 따로 구현)
        # 편의상 여기서는 마지막 층도 SirenLayer를 쓰되 w0를 작게 하여 선형에 가깝게 하거나, 
        # 실제 구현에선 활성화 함수 없는 Linear Layer를 쓴다. (여기선 예시 단순화를 위해 생략하고 개념 전달에 집중)
        self.last_layer = SirenLayer(hidden_dim, output_dim, w0=1.0, is_first=False) 

    def forward(self, x):
        out = x
        for layer in self.layers:
            out = layer.forward(out)
        out = self.last_layer.forward(out) 
        return out

    def backward(self, grad_loss):
        # 역전파: 마지막 층 -> 첫 번째 층
        grad = self.last_layer.backward(grad_loss, learning_rate=0.001)
        for layer in reversed(self.layers):
            grad = layer.backward(grad, learning_rate=0.001)

# --- 실행 예제 ---
# 간단한 목표: f(x) = sin(x) 를 학습
x_train = np.random.uniform(-np.pi, np.pi, (64, 1))
y_target = np.sin(x_train)

model = SirenNetwork(input_dim=1, hidden_dim=32, output_dim=1)

# 학습 루프 (1회만 예시)
y_pred = model.forward(x_train)
loss = np.mean((y_pred - y_target)**2) # MSE Loss

# dL/dy_pred 계산 (Loss 미분)
grad_loss = 2 * (y_pred - y_target) / x_train.shape[0]

# 역전파 수행
model.backward(grad_loss)

print(f"Loss: {loss:.6f}")
print("역전파 완료: 가중치 업데이트 됨")