package net.woodyfolsom.msproj.ann2;

import java.util.List;

import net.woodyfolsom.msproj.ann2.math.ErrorFunction;
import net.woodyfolsom.msproj.ann2.math.MSSE;

public class BackPropagation implements TrainingMethod {
	private final ErrorFunction errorFunction;
	private final double learningRate;
	private final double momentum;

	public BackPropagation(double learningRate, double momentum) {
		this.errorFunction = MSSE.function;
		this.learningRate = learningRate;
		this.momentum = momentum;
	}

	@Override
	public void iterate(FeedforwardNetwork neuralNetwork,
			List<NNDataPair> trainingSet) {
		System.out.println("Learningrate: " + learningRate);
		System.out.println("Momentum: " + momentum);

		//zeroErrors(neuralNetwork);
		
		for (NNDataPair trainingPair : trainingSet) {
			zeroErrors(neuralNetwork);
			
			System.out.println("Training with: " + trainingPair.getInput());

			NNData ideal = trainingPair.getIdeal();
			NNData actual = neuralNetwork.compute(trainingPair);
			
			System.out.println("Updating weights.  Ideal Output: " + ideal);
			System.out.println("Actual Output: " + actual);

			updateErrors(neuralNetwork, ideal);
			
			updateWeights(neuralNetwork);
		}
		
		//updateWeights(neuralNetwork);
	}

	@Override
	public double computeError(FeedforwardNetwork neuralNetwork,
			List<NNDataPair> trainingSet) {
		int numDataPairs = trainingSet.size();
		int outputSize = neuralNetwork.getOutput().length;
		int totalOutputSize = outputSize * numDataPairs;

		double[] actuals = new double[totalOutputSize];
		double[] ideals = new double[totalOutputSize];
		for (int dataPair = 0; dataPair < numDataPairs; dataPair++) {
			NNDataPair nnDataPair = trainingSet.get(dataPair);
			double[] actual = neuralNetwork.compute(nnDataPair.getInput()
					.getValues());
			double[] ideal = nnDataPair.getIdeal().getValues();
			int offset = dataPair * outputSize;

			System.arraycopy(actual, 0, actuals, offset, outputSize);
			System.arraycopy(ideal, 0, ideals, offset, outputSize);
		}

		double MSSE = errorFunction.compute(ideals, actuals);
		return MSSE;
	}

	private void updateErrors(FeedforwardNetwork neuralNetwork, NNData ideal) {
		Neuron[] outputNeurons = neuralNetwork.getOutputNeurons();
		double[] idealValues = ideal.getValues();
		
		for (int i = 0; i < idealValues.length; i++) {
			double output = outputNeurons[i].getOutput();
			double derivative = outputNeurons[i].getActivationFunction()
					.derivative(output);
			outputNeurons[i].setError(outputNeurons[i].getError() + derivative * (idealValues[i] - output));
		}
		// walking down the list of Neurons in reverse order, propagate the
		// error
		Neuron[] neurons = neuralNetwork.getNeurons();

		for (int n = neurons.length - 1; n >= 0; n--) {
			
			Neuron neuron = neurons[n];
			double error = neuron.getError();
			
			Connection[] connectionsFromN = neuralNetwork
					.getConnectionsFrom(neuron.getId());
			if (connectionsFromN.length > 0) {

				double derivative = neuron.getActivationFunction().derivative(
						neuron.getOutput());
				for (Connection connection : connectionsFromN) {
					error += derivative * connection.getWeight() * neuralNetwork.getNeuron(connection.getDest()).getError();
				}
			}
			neuron.setError(error);
		}
	}
	
	private void updateWeights(FeedforwardNetwork neuralNetwork) {
		for (Connection connection : neuralNetwork.getConnections()) {
			Neuron srcNeuron = neuralNetwork.getNeuron(connection.getSrc());
			Neuron destNeuron = neuralNetwork.getNeuron(connection.getDest());
			double delta = learningRate * srcNeuron.getOutput() * destNeuron.getError();
			//TODO allow for momentum
			//double lastDelta = connection.getLastDelta();
			connection.addDelta(delta); 
		}
	}
	
	private void zeroErrors(FeedforwardNetwork neuralNetwork) {
		// Set output errors relative to ideals, all other errors to 0.
		for (Neuron neuron : neuralNetwork.getNeurons()) {
			neuron.setError(0.0);
		}
	}
}