#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Nov 12 00:02:15 2024

@author: amriksen
"""

# Import necessary libraries
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# Set a random seed for reproducibility
np.random.seed(42)

# Data Generation
# Create a simple regression dataset with one feature and one target
N = 100  # Number of data points
X = np.linspace(-1, 1, N).reshape(-1, 1)  # Feature
y = 2 * X.squeeze() + np.random.randn(N) * 0.2  # Target with some noise

# Normalize the target values to fall between 0 and 1 (suitable for sigmoid activation)
y_min, y_max = y.min(), y.max()
y = (y - y_min) / (y_max - y_min)

# Neural Network Parameters
input_dim = 1
hidden_dim = 1
output_dim = 1
learning_rate = 0.1
epochs = 50

# Initialize weights and biases
W1 = np.random.randn(input_dim, hidden_dim)
b1 = np.zeros((1, hidden_dim))
W2 = np.random.randn(hidden_dim, output_dim)
b2 = np.zeros((1, output_dim))

# Activation Function: Sigmoid
def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# Derivative of Sigmoid
def sigmoid_derivative(x):
    sig = sigmoid(x)
    return sig * (1 - sig)

# Mean Squared Error Loss
def compute_loss(y_true, y_pred):
    return np.mean((y_true - y_pred) ** 2)

# Stochastic Gradient Descent (SGD) Optimization
loss_history = []
weight_updates = []

for epoch in range(epochs):
    for i in range(N):
        # Select a random sample
        idx = np.random.randint(N)
        x_sample = X[idx].reshape(1, -1)
        y_sample = y[idx].reshape(1, -1)

        # Forward Pass
        hidden = sigmoid(np.dot(x_sample, W1) + b1)
        output = np.dot(hidden, W2) + b2

        # Compute Loss
        loss = compute_loss(y_sample, output)

        # Backward Pass (Gradients)
        d_loss = 2 * (output - y_sample)  # Loss gradient
        d_W2 = np.dot(hidden.T, d_loss)
        d_b2 = d_loss
        d_hidden = np.dot(d_loss, W2.T) * sigmoid_derivative(np.dot(x_sample, W1) + b1)
        d_W1 = np.dot(x_sample.T, d_hidden)
        d_b1 = d_hidden

        # Update Weights and Biases
        W1 -= learning_rate * d_W1
        b1 -= learning_rate * d_b1
        W2 -= learning_rate * d_W2
        b2 -= learning_rate * d_b2

        # Record Loss and Weights
        loss_history.append(loss)
        weight_updates.append((W1.copy(), W2.copy()))

# Animation Setup
fig, ax = plt.subplots()
ax.set_xlim(-1, 1)
ax.set_ylim(0, 1)
line, = ax.plot([], [], 'r-', lw=2)
scatter = ax.scatter(X, y, color='blue', label='Data')
loss_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)

# Initialize Animation
def init():
    line.set_data([], [])
    loss_text.set_text('')
    return line, loss_text

# Update Function for Animation
def update(frame):
    W1, W2 = weight_updates[frame]
    hidden = sigmoid(np.dot(X, W1) + b1)
    y_pred = np.dot(hidden, W2) + b2
    line.set_data(X, y_pred)
    loss_text.set_text(f'Loss: {loss_history[frame]:.4f}')
    return line, loss_text, scatter

# Create Animation
ani = animation.FuncAnimation(fig, update, frames=len(weight_updates), init_func=init, blit=True, interval=50)

# Save as GIF
gif_path = "sgd_animation_with_sigmoid.gif"
ani.save(gif_path, writer='imagemagick', fps=30)

# Display the last plot frame for verification
hidden = sigmoid(np.dot(X, W1) + b1)
y_pred = np.dot(hidden, W2) + b2
plt.scatter(X, y, color='blue', label='Data')
plt.plot(X, y_pred, color='red', label='Predicted')
plt.title('Two-Layer Neural Network with Sigmoid Activation (Regression)')
plt.xlabel('X')
plt.ylabel('y (normalized)')
plt.legend()
plt.show()

# Return the path of the generated GIF
gif_path