Main.java

package Final;

import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.Random;
import java.util.Scanner;

class ActivationFunction {
    // possibly swith to rmsprop for better op
    // switch to leaky or parametric relu
    public static double[][] sigmoid(double[][] matrix) {
        double[][] newMatrix = new double[matrix.length][matrix[0].length];
    
        for (int i = 0; i < newMatrix.length; i++) {
            for (int j = 0; j < newMatrix[i].length; j++) {
                newMatrix[i][j] = 1 / (1 + Math.exp(-matrix[i][j]));
            }
        }
    
        return newMatrix;
    }
    
    public static double[][] sigmoidDerivative(double[][] matrix) {
        double[][] sigmoidMatrix = sigmoid(matrix);  
        double[][] derivative = new double[matrix.length][matrix[0].length];
        
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[i].length; j++) {
                double sigmoidValue = sigmoidMatrix[i][j];
                derivative[i][j] = sigmoidValue * (1 - sigmoidValue);
            }
        }
        
        return derivative;
    }
    

    // public static double sigmoidDerivative(double x) {
    //     return (Math.exp(-x) / Math.pow(1 + Math.exp(-x)))
    // }

    public static double[][] rectifiedLinearUnit(double[][] matrix) {
        double[][] newMatrix = new double[matrix.length][matrix[0].length];

        for (int i = 0; i < newMatrix.length; i++) {
            for (int j = 0; j < newMatrix[i].length; j++) {
                if (matrix[i][j] >= 0.0) {
                    newMatrix[i][j] = matrix[i][j];
                }
                else {
                    newMatrix[i][j] = 0.01 * matrix[i][j];
                }
            }
        }

        return newMatrix;
    }

    public static double[][] deravtiveRectifiedLinearUnit(double[][] matrix) {
        double[][] derivative = new double[matrix.length][matrix[0].length];
        
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[i].length; j++) {
                if (matrix[i][j] > 0) {
                    derivative[i][j] = 1;
                } else {
                    derivative[i][j] = 0.01;
                }
            }
        }
        
        return derivative;
    }
    
    
    public static double[][] softMax(double[][] matrix) {
        // matrix shape is [numClasses][numExamples]; normalize per example (column).
        int numClasses = matrix.length;
        int numExamples = matrix[0].length;
        double[][] softmax = new double[numClasses][numExamples];

        for (int ex = 0; ex < numExamples; ex++) {
            double max = matrix[0][ex];
            for (int c = 1; c < numClasses; c++) {
                if (matrix[c][ex] > max) max = matrix[c][ex];
            }

            double sum = 0.0;
            for (int c = 0; c < numClasses; c++) {
                double e = Math.exp(matrix[c][ex] - max); // stability
                softmax[c][ex] = e;
                sum += e;
            }
            for (int c = 0; c < numClasses; c++) {
                softmax[c][ex] /= sum;
            }
        }

        return softmax;
    }    

    public static double[][] sigmoidDerivativeFromActivated(double[][] sigmoidActivated) {
        double[][] derivative = new double[sigmoidActivated.length][sigmoidActivated[0].length];
        for (int i = 0; i < sigmoidActivated.length; i++) {
            for (int j = 0; j < sigmoidActivated[i].length; j++) {
                double s = sigmoidActivated[i][j];
                derivative[i][j] = s * (1.0 - s);
            }
        }
        return derivative;
    }
}

class LossFunction {
    public static double meanSquaredError(double[][] actual, double[][] predictions) {
        int numClasses = actual.length;
        int numExamples = actual[0].length;
        double sumSquaredError = 0;
    
        for (int i = 0; i < numExamples; i++) {
            for (int j = 0; j < numClasses; j++) {
                double error = actual[j][i] - predictions[j][i]; 
                sumSquaredError += Math.pow(error, 2); 
            }
        }
    
        double mse = sumSquaredError / (numClasses * numExamples);
        return mse;
    }
    
    public static double crossEntropy(double[][] probabilities, double[][] oneHotEncoded) {
        // Both matrices are expected as [numClasses][numExamples]
        int numClasses = oneHotEncoded.length;
        int numExamples = oneHotEncoded[0].length;
        double loss = 0.0;
        double epsilon = 1e-15; // Small value to avoid log(0)

        // Cross-entropy: -1/N * sum_{ex} sum_{c} y[c,ex] * log(p[c,ex])
        for (int ex = 0; ex < numExamples; ex++) {
            for (int c = 0; c < numClasses; c++) {
                double y = oneHotEncoded[c][ex];
                if (y != 0.0) {
                    double p = probabilities[c][ex];
                    if (p < epsilon) p = epsilon;
                    loss += -Math.log(p);
                }
            }
        }

        return loss / numExamples;
    } 
}

class Neuron {
    private double[] weights;
    private double bias;

    public Neuron(int numInputs, int numOutputs) {
        Random rand = new Random();
        double initWeightRange = Math.sqrt(6.0 / (numInputs + numOutputs)); // Xavier initialization
        weights = new double[numInputs];
        for (int i = 0; i < numInputs; i++) {
            weights[i] = rand.nextDouble() * 2 * initWeightRange - initWeightRange;
        }
        bias = 0.0; // Initialize bias to zero
    }

    public double[] getWeights() {
        return weights;
    }

    public double getBias() {
        return bias;
    }

    public void updateWeights(double[] weights) {
        this.weights = weights;
    }

    public void updateBias(double bias) {
        this.bias = bias;
    }
}

class Layer {
    private Neuron[] neurons;

    public Layer(int numNeurons, int numInputsPerNeuron, int numOutputNeurons) {
        neurons = new Neuron[numNeurons];
        for (int i = 0; i < numNeurons; i++) {
            neurons[i] = new Neuron(numInputsPerNeuron, numOutputNeurons);
        }
    }

    // public double[][] calculateOutputs(double[] inputs, int numberOfImages) {
    //     double[][] outputs = new double[numberOfImages][neurons.length];

    //     for (int r = 0; r < neurons.length; r++) {
    //         for (int c = 0; c < neurons.length; c++) {
    //             outputs[r][c] = neurons[c].calculateOutput(inputs);
    //         }
    //     }
        
    //     return outputs;
    // }

    // public double[] calculateOutputs(double[] inputs, double[] target) {
    //     double[] classProb = new double[inputs.length];
    //     for (int i = 0; i < inputs.length; i++) {
    //         classProb[i] = ActivationFunction.softMax(inputs[i], inputs);
    //     }
    //     return classProb;
    // }

    public static int classify(double[] classProb) {
        double highest = classProb[0];
        int highestIndex = 0;

        for (int j = 0; j < classProb.length; j++) {
            if (highest < classProb[j]) {
                highest = classProb[j];
                highestIndex = j;
            }
        }
        
        return highestIndex;
    }

    public Neuron[] getNeurons() {
        return neurons;
    }

    public void updateNeurons(double[][] weights, double[][] bias) {
        for (int i = 0; i < neurons.length; i++) {
            neurons[i].updateWeights(weights[i]);
            neurons[i].updateBias(bias[i][0]);
        }
    }
}

class NeuralNetwork {
    private Layer hiddenLayer;
    private Layer outputLayer;

    // Momentum buffers (persist across epochs).
    private double[][] hiddenWeightVelocity;
    private double[][] hiddenBiasVelocity;
    private double[][] outputWeightVelocity;
    private double[][] outputBiasVelocity;
    private boolean velocitiesInitialized = false;

    public NeuralNetwork(int numInputNeurons, int numHiddenNeurons, int numOutputNeurons) {
        hiddenLayer = new Layer(numHiddenNeurons, numInputNeurons, numOutputNeurons);
        outputLayer = new Layer(numOutputNeurons, numHiddenNeurons, numOutputNeurons);
    }

    // public double[] forwardPass(double[] inputs) {
    //     double[] hiddenOutputs = hiddenLayer.calculateOutputs(inputs);
    //     return outputLayer.calculateOutputs(hiddenOutputs); //fix
    // }

    public int[] getPredctions(double[][] percentageMatrix) {
        int[] predictions = new int[percentageMatrix[0].length];
        double[][] tempMatrix = PreProcess.transpose(percentageMatrix);

        for (int i = 0; i < tempMatrix.length; i++) {
            double highest = tempMatrix[i][0];
            int highestIndex = 0;

            for (int j = 0; j < tempMatrix[i].length; j++) {
                if (tempMatrix[i][j] > highest) {
                    highest = tempMatrix[i][j];
                    highestIndex = j;
                }
            }

            predictions[i] = highestIndex;
        }
        
        return predictions;
    }

    public double train(double[][] inputs, double[][] targets, double learningRate, int numberOfImages, int epoch) {
        int batchSize = 64;
        double momentum = 0.9;

        int numExamples = inputs[0].length;
        int actualNumImages = Math.min(numberOfImages, numExamples);

        // Current parameters
        Neuron[] hiddenNeurons = hiddenLayer.getNeurons();
        double[][] hiddenWeights = new double[hiddenNeurons.length][];
        double[][] hiddenBiases = new double[hiddenNeurons.length][1];
        for (int i = 0; i < hiddenNeurons.length; i++) {
            hiddenWeights[i] = hiddenNeurons[i].getWeights();
            hiddenBiases[i][0] = hiddenNeurons[i].getBias();
        }

        Neuron[] outputNeurons = outputLayer.getNeurons();
        double[][] outputWeights = new double[outputNeurons.length][];
        double[][] outputBiases = new double[outputNeurons.length][1];
        for (int i = 0; i < outputNeurons.length; i++) {
            outputWeights[i] = outputNeurons[i].getWeights();
            outputBiases[i][0] = outputNeurons[i].getBias();
        }

        // Initialize momentum buffers once (shape depends on layer sizes).
        if (!velocitiesInitialized) {
            hiddenWeightVelocity = new double[hiddenWeights.length][];
            for (int i = 0; i < hiddenWeights.length; i++) {
                hiddenWeightVelocity[i] = new double[hiddenWeights[i].length];
            }
            hiddenBiasVelocity = new double[hiddenWeights.length][1];

            outputWeightVelocity = new double[outputWeights.length][];
            for (int i = 0; i < outputWeights.length; i++) {
                outputWeightVelocity[i] = new double[outputWeights[i].length];
            }
            outputBiasVelocity = new double[outputWeights.length][1];

            velocitiesInitialized = true;
        }

        // Shuffle training examples.
        int[] indices = new int[actualNumImages];
        for (int i = 0; i < actualNumImages; i++) {
            indices[i] = i;
        }
        Random rng = new Random(epoch + 1234);
        for (int i = actualNumImages - 1; i > 0; i--) {
            int j = rng.nextInt(i + 1);
            int tmp = indices[i];
            indices[i] = indices[j];
            indices[j] = tmp;
        }

        // Mini-batch SGD with momentum.
        for (int start = 0; start < actualNumImages; start += batchSize) {
            int end = Math.min(actualNumImages, start + batchSize);
            int currentBatchSize = end - start;

            // Slice batch columns from inputs/targets.
            double[][] batchInputs = new double[inputs.length][currentBatchSize];
            for (int r = 0; r < inputs.length; r++) {
                for (int b = 0; b < currentBatchSize; b++) {
                    batchInputs[r][b] = inputs[r][indices[start + b]];
                }
            }

            double[][] batchTargets = new double[targets.length][currentBatchSize];
            for (int c = 0; c < targets.length; c++) {
                for (int b = 0; b < currentBatchSize; b++) {
                    batchTargets[c][b] = targets[c][indices[start + b]];
                }
            }

            // Forward pass.
            double[][] hiddenInputs = PreProcess.matrixOperations(
                PreProcess.dot(hiddenWeights, batchInputs),
                PreProcess.copyAcross(hiddenBiases, currentBatchSize),
                false
            );
            double[][] hiddenOutputs = ActivationFunction.rectifiedLinearUnit(hiddenInputs);

            double[][] outputInputs = PreProcess.matrixOperations(
                PreProcess.dot(outputWeights, hiddenOutputs),
                PreProcess.copyAcross(outputBiases, currentBatchSize),
                false
            );
            double[][] actualOutputs = ActivationFunction.softMax(outputInputs);

            // Backprop for softmax + cross entropy: dZ = (P - Y) / N
            double[][] outputErrors = PreProcess.matrixCoefficientMultiplication(
                PreProcess.matrixOperations(actualOutputs, batchTargets, true),
                (1.0 / currentBatchSize)
            );

            double[][] deravtiveOutputWeigths = PreProcess.dot(outputErrors, PreProcess.transpose(hiddenOutputs));
            double[][] deravtiveOutputBiases = PreProcess.sumSecondAxis(outputErrors);

            double[][] hiddenErrors = PreProcess.elementWiseMult(
                PreProcess.dot(PreProcess.transpose(outputWeights), outputErrors),
                ActivationFunction.deravtiveRectifiedLinearUnit(hiddenInputs)
            );

            double[][] deravtiveHiddenWeigths = PreProcess.dot(hiddenErrors, PreProcess.transpose(batchInputs));
            double[][] deravtiveHiddenBiases = PreProcess.sumSecondAxis(hiddenErrors);

            // Momentum updates.
            for (int i = 0; i < outputWeights.length; i++) {
                for (int j = 0; j < outputWeights[i].length; j++) {
                    outputWeightVelocity[i][j] = momentum * outputWeightVelocity[i][j] + deravtiveOutputWeigths[i][j];
                    outputWeights[i][j] -= learningRate * outputWeightVelocity[i][j];
                }
                outputBiasVelocity[i][0] = momentum * outputBiasVelocity[i][0] + deravtiveOutputBiases[i][0];
                outputBiases[i][0] -= learningRate * outputBiasVelocity[i][0];
            }

            for (int i = 0; i < hiddenWeights.length; i++) {
                for (int j = 0; j < hiddenWeights[i].length; j++) {
                    hiddenWeightVelocity[i][j] = momentum * hiddenWeightVelocity[i][j] + deravtiveHiddenWeigths[i][j];
                    hiddenWeights[i][j] -= learningRate * hiddenWeightVelocity[i][j];
                }
                hiddenBiasVelocity[i][0] = momentum * hiddenBiasVelocity[i][0] + deravtiveHiddenBiases[i][0];
                hiddenBiases[i][0] -= learningRate * hiddenBiasVelocity[i][0];
            }

            // Commit updated parameters.
            outputLayer.updateNeurons(outputWeights, outputBiases);
            hiddenLayer.updateNeurons(hiddenWeights, hiddenBiases);
        }

        // Accuracy/loss reporting only every 50 epochs (keeps runtime reasonable).
        if (epoch % 50 == 0) {
            double[][] probs = forward(inputs);
            int[] predictions = getPredctions(probs);

            int correct = 0;
            for (int k = 0; k < predictions.length; k++) {
                if (targets[predictions[k]][k] == 1.0) {
                    correct++;
                }
            }

            System.out.println("Loss: " + LossFunction.crossEntropy(probs, targets));
            return ((double) correct / actualNumImages) * 100.0;
        }

        return 0;
    }

    private double[][] forward(double[][] inputs) {
        int numberOfImages = inputs[0].length;

        // Hidden layer
        Neuron[] hiddenNeurons = hiddenLayer.getNeurons();
        double[][] hiddenWeights = new double[hiddenNeurons.length][];
        double[][] hiddenBiases = new double[hiddenNeurons.length][1];
        for (int i = 0; i < hiddenNeurons.length; i++) {
            hiddenWeights[i] = hiddenNeurons[i].getWeights();
            hiddenBiases[i][0] = hiddenNeurons[i].getBias();
        }

        double[][] hiddenInputs = PreProcess.matrixOperations(
            PreProcess.dot(hiddenWeights, inputs),
            PreProcess.copyAcross(hiddenBiases, numberOfImages),
            false
        );
        double[][] hiddenOutputs = ActivationFunction.rectifiedLinearUnit(hiddenInputs);

        // Output layer
        Neuron[] outputNeurons = outputLayer.getNeurons();
        double[][] outputWeights = new double[outputNeurons.length][];
        double[][] outputBiases = new double[outputNeurons.length][1];
        for (int i = 0; i < outputNeurons.length; i++) {
            outputWeights[i] = outputNeurons[i].getWeights();
            outputBiases[i][0] = outputNeurons[i].getBias();
        }

        double[][] outputInputs = PreProcess.matrixOperations(
            PreProcess.dot(outputWeights, hiddenOutputs),
            PreProcess.copyAcross(outputBiases, numberOfImages),
            false
        );

        // Probabilities for each class (shape [numClasses][numExamples])
        return ActivationFunction.softMax(outputInputs);
    }

    public int[] predict(double[][] inputs) {
        return getPredctions(forward(inputs));
    }

    // Returns accuracy as a percentage.
    public double test(double[][] inputs, double[][] targets) {
        int[] predictions = predict(inputs);
        int numExamples = predictions.length;
        int correct = 0;

        // targets is one-hot encoded, shape [numClasses][numExamples]
        for (int k = 0; k < numExamples; k++) {
            int pred = predictions[k];
            if (targets[pred][k] == 1.0) {
                correct++;
            }
        }

        return ((double) correct / numExamples) * 100.0;
    }
}

public class Main {
    public static void main(String[] args) throws FileNotFoundException, IOException {
        NeuralNetwork neuralNetwork = new NeuralNetwork(784, 64, 10);

        System.out.println("[*] Choose an Option below and enter the number:");
        System.out.println("1. Train");
        System.out.println("2. Test");

        /*
         * TO DO:
         * Parallel Processing
         * using ByteBuffer
         * Adding print statments to show progess 
         * Allow error checking?
         * intliaze arrays outside of loops to save mem
         */

        @SuppressWarnings("resource")
        Scanner input = new Scanner(System.in);
        int option = input.nextInt();

        if (option == 1) {
            double[][] trainingInputs = PreProcess.processImages(4000);
            double[][] trainingOutputs = PreProcess.processLabels(4000);
            int epochs = 401;
            double learningRate = 0.01;
            System.out.println("Starting training");
            for (int i = 0; i < epochs; i++) {
                double accuracy = neuralNetwork.train(trainingInputs, trainingOutputs, learningRate, 4000, i);
                if (i % 50 == 0) {
                    System.out.println("Epoch: " + i + " Accuarcy: " + accuracy);
                }
            }
            
            // Evaluate on MNIST test set (t10k)
            int testSize = 2000;
            System.out.println("Loading test set (" + testSize + ")...");
            double[][] testInputs = PreProcess.processImages("t10k-images.idx3-ubyte", testSize);
            double[][] testOutputs = PreProcess.processLabels("t10k-labels.idx1-ubyte", testSize);
            double testAccuracy = neuralNetwork.test(testInputs, testOutputs);
            System.out.println("Test Accuracy: " + testAccuracy);

            // Print a few predictions
            int[] predictions = neuralNetwork.predict(testInputs);
            for (int k = 0; k < 10 && k < predictions.length; k++) {
                int pred = predictions[k];
                int trueLabel = -1;
                for (int c = 0; c < 10; c++) {
                    if (testOutputs[c][k] == 1.0) {
                        trueLabel = c;
                        break;
                    }
                }
                System.out.println("Sample " + k + ": predicted=" + pred + ", true=" + trueLabel);
            }
        } else if (option == 2) {
            // Test-only mode (weights are still random unless you trained in this run).
            int testSize = 1000;
            System.out.println("Loading test set (" + testSize + ")...");
            double[][] testInputs = PreProcess.processImages("t10k-images.idx3-ubyte", testSize);
            double[][] testOutputs = PreProcess.processLabels("t10k-labels.idx1-ubyte", testSize);
            double testAccuracy = neuralNetwork.test(testInputs, testOutputs);
            System.out.println("Test Accuracy (random weights if you didn't train): " + testAccuracy);

            int[] predictions = neuralNetwork.predict(testInputs);
            for (int k = 0; k < 10 && k < predictions.length; k++) {
                int pred = predictions[k];
                int trueLabel = -1;
                for (int c = 0; c < 10; c++) {
                    if (testOutputs[c][k] == 1.0) {
                        trueLabel = c;
                        break;
                    }
                }
                System.out.println("Sample " + k + ": predicted=" + pred + ", true=" + trueLabel);
            }
        } else {
            System.out.println("No second Chances try again next time");
        }

        // for (int i = 0; i < trainingInputs.length; i++) {
        //     double[] inputs = trainingInputs[i];
        //     double[] predictedOutputs = neuralNetwork.forwardPass(inputs);
        //     System.out.println("Input: " + inputs[0] + ", " + inputs[1] + " | Predicted Output: " + predictedOutputs[0]);
        // }
    }
}