#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Nov  6 19:13:12 2024

@author: amriksen
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
from mpl_toolkits.mplot3d import Axes3D

# Generate sample data for linear regression
np.random.seed(0)
X = 2 * np.random.rand(50, 1)  # Reduced data points to 50
y = 4 + 7 * X + np.random.randn(50, 1)*3 # higher variance  = 9

# Define prediction and cost functions
def predict(X, m, c):
    return m * X + c

def compute_cost(X, y, m, c):
    return np.mean((predict(X, m, c) - y) ** 2)

# Define gradients
def compute_gradients(X_i, y_i, m, c):
    error = (m * X_i + c) - y_i
    d_m = 2 * error * X_i
    d_c = 2 * error
    return d_m.item(), d_c.item()  # Ensure gradients are scalars

# Implement Stochastic Gradient Descent
def stochastic_gradient_descent(X, y, learning_rate=0.05, n_epochs=15):  # Increased learning rate, reduced epochs
    m, c = 0, 0  # Initial guesses for m and c
    m_history = [m]
    c_history = [c]
    cost_history = [compute_cost(X, y, m, c)]
    
    for epoch in range(n_epochs):
        # for i in range(len(X)):
        #     X_i = X[i]
        #     y_i = y[i]
        for _ in range(len(X)):
            # Randomly select an index for stochastic gradient computation
            random_index = np.random.randint(len(X))
            X_i = X[random_index]
            y_i = y[random_index]    
            d_m, d_c = compute_gradients(X_i, y_i, m, c)
            m = m - learning_rate * d_m
            c = c - learning_rate * d_c
            m_history.append(m)
            c_history.append(c)
            cost_history.append(compute_cost(X, y, m, c))
    
    return np.array(m_history), np.array(c_history), np.array(cost_history)

# Run SGD
learning_rate = 0.005
n_epochs = 15
m_history, c_history, cost_history = stochastic_gradient_descent(X, y, learning_rate, n_epochs)

# Create a grid of m and c values for plotting the error surface
m_vals = np.linspace(-2, 15, 100)
c_vals = np.linspace(-2, 15, 100)
M, C = np.meshgrid(m_vals, c_vals)
Z = np.array([[compute_cost(X, y, m, c) for m, c in zip(row_m, row_c)] for row_m, row_c in zip(M, C)])

# Prepare regression line data
y_pred_history = [predict(X, m, c) for m, c in zip(m_history, c_history)]

# Initialize the plot with two subplots
fig = plt.figure(figsize=(14, 6))
ax1 = fig.add_subplot(121, projection='3d')
ax2 = fig.add_subplot(122)

# Plot the MSE surface on the first subplot
ax1.plot_surface(M, C, Z, cmap="viridis", alpha=0.6, edgecolor='none')
ax1.contour(M, C, Z, levels=30, offset=-1, cmap="viridis", linestyles="dashed")  # Contour plot at the base
path_line, = ax1.plot([], [], [], 'r-', markersize=5, label="Path of SGD")

# Initialize the scatter plot and line for regression
ax2.scatter(X, y, color='blue', label="Data Points")
line, = ax2.plot(X, y_pred_history[0], color='red', label="Regression Line")
cost_text = ax2.text(0.02, 0.95, '', transform=ax2.transAxes, fontsize=12, color="purple")

# Set labels and legends
ax1.set_xlabel("m")
ax1.set_ylabel("c")
ax1.set_zlabel("MSE")
ax1.set_title("3D MSE Surface with SGD Path")
ax1.legend()

ax2.set_xlabel("X")
ax2.set_ylabel("y")
ax2.legend()

# Sample every 5 frames for animation
sample_rate = 5
m_history_sampled = m_history[::sample_rate]
c_history_sampled = c_history[::sample_rate]
cost_history_sampled = cost_history[::sample_rate]
y_pred_history_sampled = y_pred_history[::sample_rate]

# Update function for animation
def animate(i):
    # Update MSE surface path
    path_line.set_data(m_history_sampled[:i+1], c_history_sampled[:i+1])
    path_line.set_3d_properties(cost_history_sampled[:i+1])
    
    # Update regression line
    line.set_ydata(y_pred_history_sampled[i])  # Update regression line
    cost_text.set_text(f"Cost: {cost_history_sampled[i]:.2f}")
    ax1.view_init(elev=30, azim=3 * i)  # Rotate view for better visualization
    
    return path_line, line, cost_text

# Create animation
anim = FuncAnimation(fig, animate, frames=len(m_history_sampled), interval=200, blit=True)

# Save the animation as a GIF with reduced frame rate
anim.save("sgd_mse_and_regression.gif", writer=PillowWriter(fps=5))

plt.show()