#ifndef BRAIN_H
#define BRAIN_H
// #pragma once

#include <vector>
#include <cmath>
using namespace std;

#define nn_eta 0.15 // learning rate
#define nn_alpha 0.5 // momentum
#define nn_weight_size 0.1 // random weight size multiplier

namespace nn
{
	namespace nt
	{
		enum
		{
			hidden = 1,
			input,
			output,
			bias,
			memory,
		};
	}

	class neuron;

	// neuron(output) ---> axon(weight) ---> target neuron( sum += output * weight )
	class axon
	{
	public:
		neuron *target;
		virtual double GetWeight() = 0;
		virtual double GetDetla() = 0;
		virtual void SetWeight(double NewWeight) = 0;
		virtual void SetDelta(double NewDelta) = 0;
		virtual bool IsMem() = 0; // memory axon doesn't need weight or delta, only a target neuron what is always memory neuron.

		axon(neuron * Target) {
			target = Target;
		}
		virtual ~axon() = default;
	};

	class normal_axon : public axon
	{
	public:
		double weight;
		double delta;
		double GetWeight() { return weight; }
		double GetDetla() { return delta; }
		void SetWeight(double NewWeight) { weight = NewWeight; }
		void SetDelta(double NewDelta){ delta = NewDelta; }
		bool IsMem() { return false; }
		normal_axon(neuron * Target) : axon(Target) {
			this->SetWeight(((double)rand() / double(RAND_MAX)) * nn_weight_size);
			this->SetDelta(0.0);
		};
	};

	class memory_axon : public axon
	{
	public:
		double GetWeight() { return 0.0; }
		double GetDetla() { return 0.0; }
		void SetWeight(double NewWeight) { }
		void SetDelta(double NewDelta) { }
		bool IsMem() { return true; }
		memory_axon(neuron * Target) : axon(Target) {};
	};

	class neuron {
	public:

		virtual ~neuron() = default;

		virtual double GetOutput() = 0;
		virtual void setOutput(double NewOutput) = 0;
		virtual double GetGradient() = 0;
		virtual void SetGradient(double NewGradient) = 0;
		virtual void SetTarget(double NewTarget) = 0;
		virtual double GetTarget() = 0;
		virtual bool HaveMemoryNeuron() = 0;
		virtual neuron *GetMemoryNeuron() = 0;
		virtual void PrintType() = 0;
		virtual int GetType() = 0;
		virtual void UpdateWeights() = 0;
		virtual void CalculateGradient() = 0;
		virtual void Activate() = 0;
		virtual void MemorizeOutput( double output ) = 0;
		virtual void Fire(double signal) = 0;
		virtual vector<axon *> GetAxons() = 0;
		virtual void AddConnection(neuron *Target) = 0;
		virtual bool CanBeConnected() = 0;
	};

	class ninput : public neuron {
	public:

		// output is the value it feeds to other neurons, since it input neuron it gets its output value from outside world, training data for example.
		double output;

		// connections to other neurons
		vector<axon *> axons;

		double GetOutput() { return this->output; }
		void setOutput(double NewOutput)
		{
			this->MemorizeOutput(this->output); // In case this neuron is connected to memory neuron, lets have that memory neuron remember the old output value
			this->output = NewOutput;
		}
		double GetGradient() { return 0.0; }
		void SetGradient(double NewGradient) {}
		void SetTarget(double NewTarget) {}
		double GetTarget() { return 0.0; }
		bool HaveMemoryNeuron() { if (this->axons.size() && this->axons[0]->IsMem()) return true; return false; }
		neuron *GetMemoryNeuron() { if (this->HaveMemoryNeuron()) return this->axons[0]->target; return 0; }
		void PrintType() { cout << "Input"; }
		int GetType() { return nt::input; }

		void UpdateWeights() {
			for (int a = this->HaveMemoryNeuron() == true ? 1 : 0; a < this->axons.size(); a++) {
				this->axons[a]->SetDelta(nn_eta * this->output * this->axons[a]->target->GetGradient() + this->axons[a]->GetDetla() * nn_alpha);
				this->axons[a]->SetWeight(this->axons[a]->GetWeight() + this->axons[a]->GetDetla());
			}
		}

		void CalculateGradient() {}
		void Activate() {
			// Lets loop this neuron all connections and fire our neuron's (output * current connection weight) to another neuron
			// If we have memory axon then ignore it. Its function isin't fire neuron. it's just a pointer to memory neuron
			for (int a = this->HaveMemoryNeuron() == true ? 1 : 0; a < this->axons.size(); a++)
			{
				this->axons[a]->target->Fire( this->output * this->axons[a]->GetWeight() );
			}
			// since input neuron doesn't use activation function, we're done here
		}
		void MemorizeOutput(double output) {
			if (this->HaveMemoryNeuron()) this->GetMemoryNeuron()->setOutput(output);
		}

		bool CanBeConnected() { return false; }

		void Fire(double signal) {}
		vector<axon *> GetAxons() { return this->axons; }

		void AddConnection( neuron *Target ) {
			if (Target->GetType() == nt::memory) {
				if (this->HaveMemoryNeuron()) return;

				// Insert memory axon as first so it would be easier to find it later on
				this->axons.insert(this->axons.begin(), new memory_axon(Target));
			}
			else
			{
				if (!Target->CanBeConnected()) return;
				this->axons.push_back(new normal_axon(Target));
			}
		}

		ninput() {
			output = 0.0;
		}

		~ninput() {
			output = 0.0;
		}
	};

	class nhidden : public neuron {
	public:
		double output;
		double gradient;
		double sum;
		vector<axon *> axons;

		double GetOutput() { return this->output; }
		void setOutput(double NewOutput)
		{
			this->MemorizeOutput(this->output);
			this->output = NewOutput;
		}
		double GetGradient() { return this->gradient; }
		void SetGradient(double NewGradient) { this->gradient = NewGradient; }
		void SetTarget(double NewTarget) {}
		double GetTarget() { return 0.0; }
		bool HaveMemoryNeuron() { if (this->axons.size() && this->axons[0]->IsMem()) return true; return false; }
		neuron *GetMemoryNeuron() { if (this->HaveMemoryNeuron()) return this->axons[0]->target; return 0; }
		void PrintType() { cout << "hidden"; }
		int GetType() { return nt::hidden; }

		void UpdateWeights() {
			for (int a = this->HaveMemoryNeuron() == true ? 1 : 0; a < this->axons.size(); a++) {
				this->axons[a]->SetDelta( nn_eta * this->output * this->axons[a]->target->GetGradient() + this->axons[a]->GetDetla() * nn_alpha);
				this->axons[a]->SetWeight(this->axons[a]->GetWeight() + this->axons[a]->GetDetla());
			}
		}

		void CalculateGradient() {
			double Gsum = 0.0;
			for (int a = this->HaveMemoryNeuron() == true ? 1 : 0; a < this->axons.size(); a++) {
				Gsum += this->axons[a]->GetWeight() * this->axons[a]->target->GetGradient();
			}
			this->gradient = Gsum * (1.0 - this->output * this->output);
		}

		void Activate() {
			this->setOutput(tanh(this->sum));
			this->sum = 0;
			for (int a = this->HaveMemoryNeuron() == true ? 1 : 0; a < this->axons.size(); a++) {
				this->axons[a]->target->Fire(this->output * this->axons[a]->GetWeight());
			}
		}
		void MemorizeOutput(double output) {
			if (this->HaveMemoryNeuron()) this->GetMemoryNeuron()->setOutput(output);
		}

		bool CanBeConnected() { return true; }

		void Fire(double signal) {
			this->sum += signal;
		}
		vector<axon *> GetAxons() { return this->axons; }

		void AddConnection(neuron *Target) {
			if (Target->GetType() == nt::memory) {
				if (this->HaveMemoryNeuron()) return;

				this->axons.insert(this->axons.begin(), new memory_axon(Target));
			}
			else
			{
				if (!Target->CanBeConnected()) return;
				this->axons.push_back(new normal_axon(Target));
			}
		}

		nhidden() {
			output = 0.0;
			gradient = 0.0;
			sum = 0.0;
		}

		~nhidden() {
		}
	};


	class nbias : public neuron {
	public:
		vector<axon *> axons;

		double GetOutput() { return 1.0; }
		void setOutput(double NewOutput) { }
		double GetGradient() { return 0.0; }
		void SetGradient(double NewGradient) { }
		void SetTarget(double NewTarget) {}
		double GetTarget() { return 1.0; }
		bool HaveMemoryNeuron() { return false; }
		neuron *GetMemoryNeuron() { return 0; }
		void PrintType() { cout << "bias"; }
		int GetType() { return nt::bias; }

		void UpdateWeights() {
			for (int a = 0; a < this->axons.size(); a++) {
				this->axons[a]->SetDelta(nn_eta * this->axons[a]->target->GetGradient() + this->axons[a]->GetDetla() * nn_alpha);
				this->axons[a]->SetWeight(this->axons[a]->GetWeight() + this->axons[a]->GetDetla());
			}
		}

		void CalculateGradient() { } // bias doesn't need gradient right?

		void Activate() {
			for (int a = 0; a < this->axons.size(); a++) {
				this->axons[a]->target->Fire(this->axons[a]->GetWeight());
			}
		}
		void MemorizeOutput(double output) { }

		bool CanBeConnected() { return false; }

		void Fire(double signal) { }
		vector<axon *> GetAxons() { return this->axons; }

		void AddConnection(neuron *Target) {
			if (Target->GetType() != nt::memory) {
				if (!Target->CanBeConnected()) return;
				this->axons.push_back(new normal_axon(Target));
			}
		}

		nbias() {
		}

		~nbias() {
		}
	};

	class nmemory : public neuron {
	public:
		double output;
		vector<axon *> axons;

		double GetOutput() { return this->output; }
		void setOutput(double NewOutput)
		{
			this->MemorizeOutput(this->output);
			this->output = NewOutput;
		}
		double GetGradient() { return 0.0; }
		void SetGradient(double NewGradient) { }
		void SetTarget(double NewTarget) {}
		double GetTarget() { return 0.0; }
		bool HaveMemoryNeuron() { if (this->axons.size() && this->axons[0]->IsMem()) return true; return false; }
		neuron *GetMemoryNeuron() { if (this->HaveMemoryNeuron()) return this->axons[0]->target; return 0; }
		void PrintType() { cout << "memory"; }
		int GetType() { return nt::memory; }

		void UpdateWeights() {
			for (int a = this->HaveMemoryNeuron() == true ? 1 : 0; a < this->axons.size(); a++) {
				this->axons[a]->SetDelta(nn_eta * this->output * this->axons[a]->target->GetGradient() + this->axons[a]->GetDetla() * nn_alpha);
				this->axons[a]->SetWeight(this->axons[a]->GetWeight() + this->axons[a]->GetDetla());
			}
		}

		void CalculateGradient() { }

		void Activate() {
			for (int a = this->HaveMemoryNeuron() == true ? 1 : 0; a < this->axons.size(); a++) {
				this->axons[a]->target->Fire(this->output * this->axons[a]->GetWeight());
			}
		}
		void MemorizeOutput(double output) {
			if (this->HaveMemoryNeuron()) this->GetMemoryNeuron()->setOutput(output);
		}

		bool CanBeConnected() { return true; }

		void Fire(double signal) {}
		vector<axon *> GetAxons() { return this->axons; }

		void AddConnection(neuron *Target) {
			if (Target->GetType() == nt::memory) {
				if (this->HaveMemoryNeuron()) return;
				this->axons.insert(this->axons.begin(), new memory_axon(Target));
			}
			else
			{
				if (!Target->CanBeConnected()) return;
				this->axons.push_back(new normal_axon(Target));
			}
		}

		nmemory() {
			output = 0.0;
		}

		~nmemory() {
		}
	};

	class noutput : public neuron {
	public:
		double output;
		double target;
		double gradient;
		double sum;
		vector<axon *> axons;

		double GetOutput() { return this->output; }
		void setOutput(double NewOutput) { this->output = NewOutput; }
		double GetGradient() { return this->gradient; }
		void SetGradient(double NewGradient) { this->gradient = NewGradient; }
		void SetTarget(double NewTarget)
		{
			this->MemorizeOutput(this->target);
			this->target = NewTarget;
		}
		double GetTarget() { return this->target; }
		bool HaveMemoryNeuron() { if (this->axons.size() && this->axons[0]->IsMem()) return true; return false; }
		neuron *GetMemoryNeuron() { if (this->HaveMemoryNeuron()) return this->axons[0]->target; return 0; }
		void PrintType() { cout << "output"; }
		int GetType() { return nt::output; }

		void UpdateWeights() {}

		void CalculateGradient() {
			this->gradient = ( this->target - this->output ) * (1.0 - this->output * this->output);
		}

		void Activate() {
			this->setOutput(tanh(this->sum));
			this->sum = 0;
		}
		void MemorizeOutput(double output) {
			if (this->HaveMemoryNeuron()) this->GetMemoryNeuron()->setOutput(output);
		}

		bool CanBeConnected() { return true; }

		void Fire(double signal) {
			this->sum += signal;
		}
		vector<axon *> GetAxons() { return this->axons; }

		void AddConnection(neuron *Target) {
			if (Target->GetType() == nt::memory) {
				if (this->HaveMemoryNeuron()) return;
				this->axons.insert(this->axons.begin(), new memory_axon(Target));
			}
		}

		noutput() {
			output = 0.0;
			gradient = 0.0;
			sum = 0.0;
			target = 0.0;
		}

		~noutput() {
		}
	};


	class brain
	{
	public:
		vector<vector<neuron *>> map;
		void FeedForward()
		{
			for (int y = 0; y < map.size(); y++) {
				for (int x = 0; x < map[y].size(); x++) {
					map[y][x]->Activate();
				}
			}
		}

		void BackPropagation()
		{
			for (int y = map.size() - 1; y >= 0; y--) {
				for (int x = 0; x < map[y].size(); x++) {
					map[y][x]->CalculateGradient();
				}
			}

			for (int y = map.size() - 1; y >= 0; y--) {
				for (int x = 0; x < map[y].size(); x++) {
					map[y][x]->UpdateWeights();
				}
			}
		}

		void AddNeuron(int layer, int NeuronType, int count)
		{
			for (int a = 0; a < count; a++) AddNeuron(layer, NeuronType);
		}

		neuron *AddNeuron( int layer, int NeuronType )
		{
			_MakeLayer(layer);
			neuron *ret = 0;
			switch (NeuronType)
			{
			case nt::bias:
				ret = new nbias();
				map[layer].push_back(ret);
				break;
			case nt::hidden:
				ret = new nhidden();
				map[layer].push_back(ret);
				break;
			case nt::input:
				ret = new ninput();
				map[layer].push_back(ret);
				break;
			case nt::memory:
				ret = new nmemory();
				map[layer].push_back(ret);
				break;
			case nt::output:
				ret = new noutput();
				map[layer].push_back(ret);
				break;
			}
			return ret;
		}

		neuron *GetNeuron(int y, int x) {
			return map[y][x];
		}

		void ConnectLayers(int LayerOne, int LayerTwo)
		{
			for (int x = 0; x < map[LayerOne].size(); x++)  _ConnectWithLayer(map[LayerOne][x], LayerTwo);
		}

		void _MakeLayer(int layer)
		{
			if (layer >= map.size()) {
				for (int a = 0; a < (layer - map.size()) + 1; a++) map.push_back(vector<neuron *>());
			}

		}

		void _ConnectWithLayer(neuron *Neuron, int layer)
		{
			for (int x = 0; x < map[layer].size(); x++){
				Neuron->AddConnection(map[layer][x]);
			}
		}

		brain(){}
		~brain(){}
	};


	class input
	{
	public:
		vector<double> inputs;
		input(int InputCount, ...)
		{
			va_list vl;
			va_start(vl, InputCount);
			for (int i = 0; i<InputCount; i++) {
				inputs.push_back(va_arg(vl, double));
			}
			va_end(vl);
		}
	};

	class target
	{
	public:
		vector<double> targets;
		target(int ResultCount, ...)
		{
			va_list vl;
			va_start(vl, ResultCount);
			for (int i = 0; i<ResultCount; i++)
			{
				targets.push_back(va_arg(vl, double));
			}
			va_end(vl);
		}
	};

	class TrainingData
	{
	public:
		vector<input> Inputs;
		vector<target> Targets;
		int InputCount() { return Inputs[0].inputs.size(); }
		int TargetCount() { return Targets[0].targets.size(); }
		int DataCount() { return Inputs.size(); }

		double GetInput(int index, int index2) {
			return Inputs[index].inputs[index2];
		}

		double GetTarget(int index, int index2) {
			return Targets[index].targets[index2];
		}

		TrainingData(int TrainingDataCount, ...)
		{
			va_list vl;
			va_start(vl, TrainingDataCount);
			for (int i = 0; i<TrainingDataCount; i++)
			{
				Inputs.push_back(va_arg(vl, input));
				Targets.push_back(va_arg(vl, target));
			}
			va_end(vl);
		}
	};
}
#endif