' ==========================================================
' == A Complete, Trainable Neural Network in jdBasic
' == Goal: Learn the XOR logic gate
' ==========================================================

' --- 1. Setup and Helper Functions ---

' The Sigmoid activation function
FUNC SIGMOID(X)
  ' We multiply by -1 for element-wise negation of the array
  RETURN 1 / (1 + 2.71828 ^ (X * -1))
ENDFUNC

' The *derivative* of the Sigmoid function.
' This is needed for backpropagation.
FUNC SIGMOID_DERIVATIVE(X)
  RETURN X * (1 - X)
ENDFUNC

' Helper function to create an array of a given SHAPE filled with random numbers.
FUNC RAND_ARRAY(SHAPE)
  TOTAL_ELEMENTS = PRODUCT(SHAPE)
  DIM R[TOTAL_ELEMENTS]
  FOR i = 0 TO TOTAL_ELEMENTS - 1
    R[i] = RND(1)
  NEXT i
  RETURN RESHAPE(R, SHAPE)
ENDFUNC


' --- 2. Training Data and Network Architecture ---

' XOR truth table:
' Input: [0,0], Output: [0]
' Input: [0,1], Output: [1]
' Input: [1,0], Output: [1]
' Input: [1,1], Output: [0]
TRAINING_INPUTS = [[0,0], [0,1], [1,0], [1,1]]
TRAINING_OUTPUTS = [[0], [1], [1], [0]]

' Initialize weights and biases with random numbers shifted to the range [-1, 1].
' Hidden Layer: 2 input neurons, 3 hidden neurons
' Output Layer: 3 hidden neurons, 1 output neuron
WEIGHTS_HIDDEN = (RAND_ARRAY([2, 3]) - 0.5) * 2
WEIGHTS_OUTPUT = (RAND_ARRAY([3, 1]) - 0.5) * 2
BIAS_HIDDEN = (RAND_ARRAY([1, 3]) - 0.5) * 2
BIAS_OUTPUT = (RAND_ARRAY([1, 1]) - 0.5) * 2

LEARNING_RATE = 0.1
EPOCHS = 10000 ' Number of training iterations

PRINT "--- Starting Training ---"

' --- 3. The Training Loop ---

FOR i = 1 TO EPOCHS
  ' --- FORWARD PASS ---
  HIDDEN_LAYER_ACTIVATION = MATMUL(TRAINING_INPUTS, WEIGHTS_HIDDEN) + BIAS_HIDDEN
  ACTIVATED_HIDDEN = SIGMOID(HIDDEN_LAYER_ACTIVATION)
  
  OUTPUT_LAYER_ACTIVATION = MATMUL(ACTIVATED_HIDDEN, WEIGHTS_OUTPUT) + BIAS_OUTPUT
  PREDICTED_OUTPUT = SIGMOID(OUTPUT_LAYER_ACTIVATION)

  ' --- BACKPROPAGATION ---
  ' Calculate the error
  ERROR = TRAINING_OUTPUTS - PREDICTED_OUTPUT

  ' Calculate the gradients (this is the derived calculus)
  D_PREDICTED_OUTPUT = ERROR * SIGMOID_DERIVATIVE(PREDICTED_OUTPUT)
  
  ERROR_HIDDEN_LAYER = MATMUL(D_PREDICTED_OUTPUT, TRANSPOSE(WEIGHTS_OUTPUT))
  D_HIDDEN_LAYER = ERROR_HIDDEN_LAYER * SIGMOID_DERIVATIVE(ACTIVATED_HIDDEN)

  ' --- UPDATE WEIGHTS AND BIASES ---
  ' Update output layer weights
  WEIGHTS_OUTPUT = WEIGHTS_OUTPUT + MATMUL(TRANSPOSE(ACTIVATED_HIDDEN), D_PREDICTED_OUTPUT) * LEARNING_RATE
  
  ' Update hidden layer weights
  WEIGHTS_HIDDEN = WEIGHTS_HIDDEN + MATMUL(TRANSPOSE(TRAINING_INPUTS), D_HIDDEN_LAYER) * LEARNING_RATE
  
  ' Update biases
  BIAS_OUTPUT = BIAS_OUTPUT + SUM(D_PREDICTED_OUTPUT, 0) * LEARNING_RATE
  BIAS_HIDDEN = BIAS_HIDDEN + SUM(D_HIDDEN_LAYER, 0) * LEARNING_RATE

  ' Print progress every 1000 epochs
  IF i MOD 1000 = 0 THEN
    LOSS = SUM(ERROR ^ 2) / LEN(ERROR)[0]
    PRINT "Epoch:"; i; ", Loss:"; LOSS
  ENDIF
NEXT i

PRINT
PRINT "--- Training Complete ---"
PRINT "Final Predictions:"
PRINT FRMV$(PREDICTED_OUTPUT)
