using Random, LinearAlgebra, StatsFuns
#Random.seed!(23)

# Activation function: Sigmoid function
Sigmoid(x) = 1.0 ./ (1.0 .+ exp(-x))

∇Sigmoid(x) = Sigmoid(x)*(1 - Sigmoid(x))

#Softmax(T, k) = exp(T[k])/sum([exp(T[l]) for l=1:length(T)])

Tanh(x) = tanh(x)
∇Tanh(x) = sech(x)

# Matrix rows sum
RowSum(A) = sum(A, dims=1)[1, :]

# Training function
function TrainMLN(Xin, Y, M, γ)
	# First step: add 1s to each row of X
	K = length(Y[1])
	N, P = size(Xin)
	P = P + 1
	X = ones(N, P)
	X[:, 2:P] = Xin
	# Initial random weigths
	α = 0.3*rand(M, P)
	β = 0.4*rand(K, M)	
	# Training loop
	Z = zeros(N,M)
	Yb = zeros(N,K)
	for it in 1:2000
		# Feedforward
		#println("a=", α)
		#println("b=", β)

		δ = zeros(K,N)
		s = zeros(M,N)
		
		# Loop over the Training set
		for i=1:N
			# Feedforward batch i-th
			for m=1:M
				Z[i,m] = Sigmoid(α[m,:]'*X[i,:])
			end

			for k=1:K
				Yb[i,k] = Tanh(β[k,:]'*Z[i,:])
			end

			# Compute deltas for gradiente descent
			for k = 1:K
				δ[k,i] = - 2*(Y[i,k] - Yb[i,k])*∇Sigmoid(β[k,:]'*Z[i,:])
			end			
		
			for m = 1:M
				for k = 1:K
					s[m,i] = s[m,i] + δ[k,i]*β[k,m]*∇Tanh(α[m,:]'*X[i,:])
				end				
			end
		end
		
		# Update weights with gradient descent
		for k = 1:K
			for m = 1:M
				d = zero(0.0)
				for i = 1:N			
					d = d + δ[k,i]*Z[i,m]				
				end
				β[k,m] = β[k,m] - γ * d
			end
		end
	
		for m = 1:M
			for l = 1:P
				d = zero(0.0)
				for i = 1:N			
					d = d + s[m,i]*X[i,l]
				end
				α[m,l] = α[m,l] - γ * d				
			end
		end

		# Logging loss function values
		
		# Stopping if error goes to zero
		E = Y - Yb
		println(it, "-->  Loss(w): ", norm(E), " - Weights: ", β)
		if norm(E) < 0.1			
			break
		end
	end
	
	println("Yb=", Yb)
end

#Xxor = [0 0; 0 1; 1 0; 1 1]
#Yxor = [-1.0, 1, 1, -1.0]
#println("START LEARNING OR")
#TrainMLN(Xxor, Yxor, 3, 0.9)

L = [1 1 1; 1 0 1; 1 1 0; 1 0 0; 0 1 1; 0 0 1; 0 1 0; 0 0 0]
Y = [1, 1, 1, 1, -1, 1, -1, -1]
TrainMLN(L, Y, 5, 0.1)