' ==========================================================
' == Using the ADAM Optimizer in jdBasic
' == Goal: Train a neural network to learn the function y = x^3 - x^2 + 2
' ==========================================================

DLLIMPORT "aifunc"

CLS
PRINT "--- Neural Network with ADAM Optimizer ---"
PRINT

' --- 1. Generate Training Data ---
' We will create data points for the function y = x^3 - x^2 + 2
' so the network can learn to approximate it.
PRINT "Generating training data..."

' Generate 100 data points from x = -2 to x = 2
NUM_POINTS = 100
X_FLAT_DATA = []
Y_FLAT_DATA = []
FOR I = 0 TO NUM_POINTS - 1
    ' Create x values between -2 and 2
    x_val = -2 + 4 * (I / (NUM_POINTS - 1))
    ' Calculate y = x^3 - x^2 + 2, with some small random noise
    y_val = (x_val^3) - (x_val^2) + 2 + (RND(1) - 0.5) * 0.1
    
    ' Append scalar values to flat 1D arrays first.
    ' Using APPEND with an array as the second argument would flatten the result.
    X_FLAT_DATA = APPEND(X_FLAT_DATA, x_val)
    Y_FLAT_DATA = APPEND(Y_FLAT_DATA, y_val)
NEXT I

' Reshape the flat data into 2D column vectors (matrices of size N x 1).
' This is the correct way to create rank-2 tensors for MATMUL.
X_DATA = RESHAPE(X_FLAT_DATA, [NUM_POINTS, 1])
Y_DATA = RESHAPE(Y_FLAT_DATA, [NUM_POINTS, 1])

' Convert the 2D data arrays into Tensors. This is crucial for
' enabling automatic differentiation (autodiff).
INPUTS = TENSOR.FROM(X_DATA)
TARGETS = TENSOR.FROM(Y_DATA)

PRINT "Data generated."
PRINT


' --- 2. Define the Neural Network Model ---
' The model is a standard Map data structure that holds the layers.
MODEL = {}

' We'll use a simple feed-forward network with two hidden layers.
' Input -> Dense(20) -> ReLU -> Dense(20) -> ReLU -> Output
MODEL{"layer1"} = TENSOR.CREATE_LAYER("DENSE", {"input_size": 1, "units": 20})
MODEL{"layer2"} = TENSOR.CREATE_LAYER("DENSE", {"input_size": 20, "units": 20})
MODEL{"output_layer"} = TENSOR.CREATE_LAYER("DENSE", {"input_size": 20, "units": 1})

PRINT "Model created."
PRINT


' --- 3. Setup the ADAM Optimizer and Training Parameters ---
' Here we create the ADAM optimizer instead of SGD.
' We provide the learning rate and the beta parameters.
' beta1: decay rate for the first moment estimate (momentum)
' beta2: decay rate for the second moment estimate (RMSprop)
ADAM_OPTIONS = {"learning_rate": 0.01, "beta1": 0.9, "beta2": 0.999}
OPTIMIZER = TENSOR.CREATE_OPTIMIZER("ADAM", ADAM_OPTIONS)

EPOCHS = 8000 ' The number of times to iterate over the training data.

PRINT "ADAM Optimizer configured."
PRINT


' --- 4. Define Helper Functions ---

' The forward pass function takes the model and an input tensor,
' and returns the model's prediction. It uses the ReLU activation function.
FUNC MODEL_FORWARD(current_model, input_tensor)
    ' Pass through Layer 1
    l1_out = TENSOR.MATMUL(input_tensor, current_model{"layer1"}{"weights"}) + current_model{"layer1"}{"bias"}
    l1_activated = TENSOR.RELU(l1_out)
    
    ' Pass through Layer 2
    l2_out = TENSOR.MATMUL(l1_activated, current_model{"layer2"}{"weights"}) + current_model{"layer2"}{"bias"}
    l2_activated = TENSOR.RELU(l2_out)

    ' Pass through Output Layer (no activation, as it's a regression problem)
    output = TENSOR.MATMUL(l2_activated, current_model{"output_layer"}{"weights"}) + current_model{"output_layer"}{"bias"}
    
    RETURN output
ENDFUNC

' The loss function measures how wrong the model's predictions are.
' We use Mean Squared Error (MSE), which is standard for regression.
FUNC MSE_LOSS(predicted_tensor, actual_tensor)
    errors = actual_tensor - predicted_tensor
    squared_errors = errors ^ 2
    ' Return the average of the squared errors.
    RETURN SUM(squared_errors) / LEN(TENSOR.TOARRAY(errors))[0]
ENDFUNC


' --- 5. The Training Loop ---
PRINT "--- Starting Training with ADAM ---"

FOR epoch = 1 TO EPOCHS
    ' --- a) Forward Pass ---
    ' Get the model's current predictions for the training inputs.
    PREDICTIONS = MODEL_FORWARD(MODEL, INPUTS)

    ' --- b) Calculate Loss ---
    ' Calculate how far off the predictions are from the actual targets.
    ' The result is a Tensor that is the root of our computation graph.
    LOSS = MSE_LOSS(PREDICTIONS, TARGETS)
    
    ' --- c) Backpropagation ---
    ' This one command calculates the gradient (derivative) of the loss
    ' with respect to every parameter (weight and bias) in the model.
    TENSOR.BACKWARD LOSS

    ' --- d) Update Parameters with ADAM ---
    ' The TENSOR.UPDATE function will see that the optimizer is of type "ADAM"
    ' and apply the correct update logic, managing the momentum (m) and
    ' variance (v) states internally.
    'MODEL = TENSOR.UPDATE(MODEL, OPTIMIZER)

    MODEL = ADAM.UPDATE(MODEL, OPTIMIZER)

    ' --- e) Print Progress ---
    IF epoch MOD 500 = 0 THEN
        PRINT "Epoch:"; epoch; ", Loss:"; TENSOR.TOARRAY(LOSS)[0]
    ENDIF
NEXT epoch

PRINT
PRINT "--- Training Complete ---"
PRINT


' --- 6. Verification and Final Predictions ---
PRINT "--- Verifying Final Predictions ---"
final_predictions = MODEL_FORWARD(MODEL, INPUTS)
final_predictions_array = TENSOR.TOARRAY(final_predictions)

PRINT "Sample of Inputs, Targets, and Predictions:"
PRINT "-------------------------------------------"
PRINT "  Input  |  Target  | Predicted"
PRINT "-------------------------------------------"
FOR i = 0 TO 9
    ' *** FIX: Corrected array indexing from [0,i] to [i,0] ***
    x_val_str = FORMAT$("{:7.3f}", X_DATA[i,0])
    y_val_str = FORMAT$("{:8.3f}", Y_DATA[i,0])
    pred_val_str = FORMAT$("{:9.3f}", final_predictions_array[i,0])
    PRINT x_val_str; " |"; y_val_str; " |"; pred_val_str
NEXT i
PRINT "-------------------------------------------"
