#!/usr/bin/env python
#
# file: autoencoder.py
#
# revision history:
#
# 20260408 (SP): implement deep autoencoder for self-supervised learning
#------------------------------------------------------------------------------

# import system modules
#
import os
import sys
import random
import math

# import third-party visualization and numerical libraries
#
import matplotlib.pyplot as plt
import numpy as np

#------------------------------------------------------------------------------
#
# global variables are listed here
#
#------------------------------------------------------------------------------

# set the filename using basename
#
__FILE__ = os.path.basename(__file__)

# define default hyperparameters
#
DEF_EPOCHS          = int(1500)
DEF_NUM_SAMPLES     = 300
DEF_LEARNING_RATE   = 0.01
DEF_LAMBDA          = 0.0005
DEF_RANDOM_SEED     = 42
DEF_OUTPUT_FILE     = 'autoencoder_reconstruction.png'

# define network architecture (number of neurons per layer)
# autoencoder: input -> encoder -> latent -> decoder -> output
#
DEF_INPUT_DIM      = 2

# encoder
#
DEF_HIDDEN_1_DIM   = 8  

# latent layer (compressed representation)
#
DEF_HIDDEN_2_DIM   = 2  

# reconstructed output
# 
DEF_OUTPUT_DIM     = 2   

# set the numpy random seed for reproducible data generation
#
np.random.seed(DEF_RANDOM_SEED)

# initialize empty lists for our dataset
#
def_data_x = []
def_data_y = []

# generate a 2D circular dataset to evaluate the self-supervised autoencoder
#
for _ in range(DEF_NUM_SAMPLES):
    
    # generate points roughly along a circle with some noise
    #
    angle = random.uniform(0, 2 * math.pi)
    noise_x = np.random.normal(0, 0.1)
    noise_y = np.random.normal(0, 0.1)
    
    px = math.cos(angle) * 2.0 + noise_x
    py = math.sin(angle) * 2.0 + noise_y
    
    # in self-supervised learning, the input data acts as its own label
    #
    def_data_x.append([px, py])
    def_data_y.append([px, py])

# expose the combined and fully generated standard python lists
#
DEF_DATA_X = def_data_x
DEF_DATA_Y = def_data_y

#------------------------------------------------------------------------------
#
# functions are listed here
#
#------------------------------------------------------------------------------

# apply the Tanh activation function
#
def tanh_act(x):
    """
    method: tanh_act
    arguments:
        x: float input value
    return:
        result: the non-linear hyperbolic tangent activation
    description:
        Calculates the Tanh activation function, mapping inputs 
        to a range between -1.0 and 1.0.
    """
    
    # apply math hyperbolic tangent
    #
    val = math.tanh(x)
    
    # exit gracefully
    #
    return val

# calculate the derivative of the Tanh activation
#
def tanh_deriv(x):
    """
    method: tanh_deriv
    arguments:
        x: float input value representing the pre-activation value
    return:
        result: the gradient of the tanh function
    description:
        Calculates the derivative of the Tanh function for backpropagation.
    """

    # derivative of tanh(z) is 1 - tanh^2(z)
    #
    t = math.tanh(x)
    val = 1.0 - (t * t)
    
    # exit gracefully
    #
    return val

# initialize a weight matrix with random values
#
def init_weights(rows, cols, seed=DEF_RANDOM_SEED):
    """
    method: init_weights
    arguments:
        rows: integer representing the number of neurons in the current layer
        cols: integer representing the number of neurons in the previous layer
        seed: random seed for reproducibility
    return:
        matrix: a 2D list of randomly initialized weights
    description:
        Creates and populates a 2D python list with small random floats 
        to serve as initial weights for a neural network layer.
    """
    
    # set the seed
    #
    random.seed(seed)
    
    # generate a matrix with random values centered around 0
    #
    matrix = [[random.uniform(-0.5, 0.5) for _ in range(cols)] for _ in range(rows)]
    
    # exit gracefully
    #
    return matrix

# perform a single forward pass through the autoencoder network
#
def forward_pass(x, W1, b1, W2, b2, W3, b3):
    """
    method: forward_pass
    arguments:
        x: list representing a single input feature vector
        W1, b1: weights and biases for the encoder layer
        W2, b2: weights and biases for the latent layer
        W3, b3: weights and biases for the decoder (output) layer
    return:
        z1, a1, z2, a2, z3, y_hat: intermediate pre-activations and activations
    description:
        Pushes a single input vector forward through the autoencoder architecture.
    """
    
    # calculate the encoder layer: z1 = W1*x + b1, a1 = tanh_act(z1)
    #
    z1 = []
    for i in range(len(W1)):
        z1_sum = 0.0
        for j in range(len(x)):
            z1_sum += W1[i][j] * x[j]
        z1.append(z1_sum + b1[i])
        
    # get the activations for the encoder layer
    #
    a1 = []
    for z in z1:
        a1.append(tanh_act(z))
    
    # calculate the latent layer: z2 = W2*a1 + b2, a2 = tanh_act(z2)
    #
    z2 = []
    for i in range(len(W2)):
        z2_sum = 0.0
        for j in range(len(a1)):
            z2_sum += W2[i][j] * a1[j]
        z2.append(z2_sum + b2[i])
    
    # get the activations for the latent layer
    #
    a2 = []
    for z in z2:
        a2.append(tanh_act(z))
    
    # calculate the output decoder layer: z3 = W3*a2 + b3
    #
    z3 = []
    y_hat = []
    
    for i in range(len(W3)):
        z3_sum = 0.0
        for j in range(len(a2)):
            z3_sum += W3[i][j] * a2[j]
        z3.append(z3_sum + b3[i])
        
        # apply linear activation for regression (reconstruction)
        #
        y_hat.append(z3[-1])

    # exit gracefully
    #
    return z1, a1, z2, a2, z3, y_hat

# train the self-supervised autoencoder
#
def train_autoencoder(X, y, epochs=DEF_EPOCHS, lr=DEF_LEARNING_RATE, lambda_param=DEF_LAMBDA, seed=DEF_RANDOM_SEED):
    """
    method: train_autoencoder
    arguments:
        X: list of lists representing the input vectors
        y: list of lists representing the target vectors (identical to X)
        epochs: number of times to iterate over the dataset
        lr: learning rate for gradient descent
        lambda_param: L2 regularization strength parameter
        seed: random seed for reproducibility
    return:
        W1, b1, W2, b2, W3, b3: optimized weights and biases
    description:
        Trains a deep autoencoder using Stochastic Gradient Descent (SGD) 
        and backpropagation with pure python structures. Minimizes 
        Mean Squared Error (MSE) loss combined with L2 regularization.
    """
    
    # set the random seed
    #
    random.seed(int(seed))
    
    # initialize weight matrices and bias vectors
    #
    W1 = init_weights(DEF_HIDDEN_1_DIM, DEF_INPUT_DIM, seed)
    b1 = [0.0 for _ in range(DEF_HIDDEN_1_DIM)]
    
    W2 = init_weights(DEF_HIDDEN_2_DIM, DEF_HIDDEN_1_DIM, seed+1)
    b2 = [0.0 for _ in range(DEF_HIDDEN_2_DIM)]
    
    W3 = init_weights(DEF_OUTPUT_DIM, DEF_HIDDEN_2_DIM, seed+2)
    b3 = [0.0 for _ in range(DEF_OUTPUT_DIM)]

    # loop up to the maximum number of epochs
    #
    for epoch in range(epochs):
        
        # pair features and labels so we can shuffle them together
        #
        dataset = list(zip(X, y))
        random.shuffle(dataset)
        
        # iterate through each individual sample (Stochastic Gradient Descent)
        #
        for x_val, y_val in dataset:
            
            # forward pass
            #
            z1, a1, z2, a2, z3, y_hat = forward_pass(x_val, W1, b1, W2, b2, W3, b3)
            
            # backward pass (backpropagation)
            # calculate gradient of Mean Squared Error loss w.r.t z3 (linear activation)
            #
            dz3 = []
            for k in range(DEF_OUTPUT_DIM):
                dz3.append(y_hat[k] - y_val[k])
            
            # calculate gradients for output layer (layer 3)
            #
            dW3 = []
            for i in range(DEF_OUTPUT_DIM):
                row = []
                for j in range(DEF_HIDDEN_2_DIM):
                    row.append(dz3[i] * a2[j])
                dW3.append(row)
            db3 = dz3
            
            # backpropagate error to latent layer (layer 2)
            #
            da2 = []
            for j in range(DEF_HIDDEN_2_DIM):
                da2_sum = 0.0
                for k in range(DEF_OUTPUT_DIM):
                    da2_sum += W3[k][j] * dz3[k]
                da2.append(da2_sum)

            dz2 = []
            for i in range(DEF_HIDDEN_2_DIM):
                dz2.append(da2[i] * tanh_deriv(z2[i]))
            
            # calculate gradients for latent layer
            #
            dW2 = []
            for i in range(DEF_HIDDEN_2_DIM):
                row = []
                for j in range(DEF_HIDDEN_1_DIM):
                    row.append(dz2[i] * a1[j])
                dW2.append(row)
            db2 = dz2
            
            # backpropagate error to encoder layer (layer 1)
            #
            da1 = []
            for j in range(DEF_HIDDEN_1_DIM):
                da1_sum = 0.0
                for i in range(DEF_HIDDEN_2_DIM):
                    da1_sum += W2[i][j] * dz2[i]
                da1.append(da1_sum)

            dz1 = []
            for j in range(DEF_HIDDEN_1_DIM):
                dz1.append(da1[j] * tanh_deriv(z1[j]))
            
            # calculate gradients for encoder layer
            #
            dW1 = []
            for i in range(DEF_HIDDEN_1_DIM):
                row = []
                for j in range(DEF_INPUT_DIM):
                    row.append(dz1[i] * x_val[j])
                dW1.append(row)
            db1 = dz1
            
            # update weights and biases (with L2 Regularization applied to weights)
            #
            # update layer 3
            #
            for i in range(DEF_OUTPUT_DIM):
                for j in range(DEF_HIDDEN_2_DIM):
                    W3[i][j] -= lr * (dW3[i][j] + lambda_param * W3[i][j])
                b3[i] -= lr * db3[i]
            
            # update layer 2
            #
            for i in range(DEF_HIDDEN_2_DIM):
                for j in range(DEF_HIDDEN_1_DIM):
                    W2[i][j] -= lr * (dW2[i][j] + lambda_param * W2[i][j])
                b2[i] -= lr * db2[i]
                
            # update layer 1
            #
            for i in range(DEF_HIDDEN_1_DIM):
                for j in range(DEF_INPUT_DIM):
                    W1[i][j] -= lr * (dW1[i][j] + lambda_param * W1[i][j])
                b1[i] -= lr * db1[i]

    # exit gracefully
    #
    return W1, b1, W2, b2, W3, b3

# calculate the reconstruction error of the trained autoencoder
#
def calculate_reconstruction_error(X, W1, b1, W2, b2, W3, b3):
    """
    method: calculate_reconstruction_error
    arguments:
        X: list of lists representing the input vectors
        W1, b1, W2, b2, W3, b3: the optimized network parameters
    return:
        mse: float representing the mean squared error across all points
    description:
        Evaluates the autoencoder's accuracy by determining the mean squared
        distance between the original inputs and their reconstructed outputs.
    """
    
    # initialize a variable to accumulate error
    #
    total_error = 0.0
    total_points = len(X)

    # iterate through every point in the dataset
    #
    for i in range(total_points):
        
        # calculate the network output (reconstruction)
        #
        _, _, _, _, _, y_hat = forward_pass(X[i], W1, b1, W2, b2, W3, b3)
        
        # calculate squared euclidean distance
        #
        sq_dist = 0.0
        for dim in range(DEF_INPUT_DIM):
            sq_dist += (X[i][dim] - y_hat[dim]) ** 2
            
        total_error += sq_dist

    # calculate the mean squared error
    #
    mse = total_error / total_points

    # exit gracefully
    #
    return mse

# render a plot of original points versus reconstructed points
#
def plot_reconstruction(X, W1, b1, W2, b2, W3, b3):
    """
    method: plot_reconstruction
    arguments:
        X: list of lists representing the original feature vectors
        W1, b1, W2, b2, W3, b3: the optimized network parameters
    return:
        True: indicates successful execution
    description:
        Generates a graphical plot using matplotlib to display the original
        dataset and the network's compressed-and-reconstructed version.
    """
    
    # globally set the base font size for all matplotlib elements
    #
    plt.rcParams.update({'font.size': 12})

    # gather original coordinates
    #
    orig_x = [p[0] for p in X]
    orig_y = [p[1] for p in X]
    
    # gather reconstructed coordinates via the forward pass
    #
    recon_x = []
    recon_y = []
    
    for p in X:
        _, _, _, _, _, y_hat = forward_pass(p, W1, b1, W2, b2, W3, b3)
        recon_x.append(y_hat[0])
        recon_y.append(y_hat[1])

    # plot lines connecting original points to their reconstructed counterparts
    # to visualize the individual displacement error
    #
    for i in range(len(X)):
        plt.plot([orig_x[i], recon_x[i]], [orig_y[i], recon_y[i]], color='gray', alpha=0.3)

    # plot the actual data points
    #
    plt.scatter(orig_x, orig_y, color='blue', marker='o', edgecolors='k', label='Original Data', s=40, alpha=0.7)
    plt.scatter(recon_x, recon_y, color='red', marker='X', edgecolors='k', label='Reconstructed', s=40, alpha=0.7)

    # configure the plot aesthetics and labels
    #
    plt.xlabel('Dimension 1')
    plt.ylabel('Dimension 2')
    plt.title('Self-Supervised Autoencoder Reconstruction')
    plt.legend(loc='upper right', fontsize=11)
    
    # ensure the axes are scaled equally
    #
    plt.axis('equal')
    plt.grid(True, linestyle='--', alpha=0.6)

    # save the final plot 
    #
    plt.savefig(DEF_OUTPUT_FILE)

    # exit gracefully
    #
    return True

# function: main
#
def main(argv):
    """
    method: main
    arguments:
        argv: command line arguments
    return:
        True: indicates successful execution
    description:
        Main entry point that demonstrates training a Deep Autoencoder on a 
        2D dataset using self-supervised learning principles.
    """

    # define the initial dataset and parameters
    #
    X = DEF_DATA_X
    y = DEF_DATA_Y

    # print header
    #
    print("Starting Self-Supervised Autoencoder demonstration...")
    print("Architecture: %d Inputs -> %d Encoder -> %d Latent -> %d Output" % 
          (DEF_INPUT_DIM, DEF_HIDDEN_1_DIM, DEF_HIDDEN_2_DIM, DEF_OUTPUT_DIM))
    print("Number of samples: %d" % len(X))
    print("-" * 65)

    # run the autoencoder training loop using gradient descent
    #
    print("Executing Autoencoder Stochastic Gradient Descent...")
    W1, b1, W2, b2, W3, b3 = train_autoencoder(X, y)
    
    # calculate the reconstruction error
    #
    mse = calculate_reconstruction_error(X, W1, b1, W2, b2, W3, b3)

    # print the metrics
    #
    print("Training Complete.")
    print("-" * 65)
    print("Mean Squared Error (MSE) : %.5f" % mse)
    print("-" * 65)

    # render the matplotlib plot
    #
    print("Launching Matplotlib Visualization...")
    plot_reconstruction(X, W1, b1, W2, b2, W3, b3)
    print("Plot Saved at %s. Execution finished." % DEF_OUTPUT_FILE)
    print("-" * 65)

    # exit gracefully
    #
    return True

# begin gracefully
#
if __name__ == '__main__':
    main(sys.argv[0:])

#
# end of file