package neuralnetworkbase

import (
	mat "gonum.org/v1/gonum/mat"
)

type NeuralNetwork struct {
	Count   int
	Sizes   []int
	Biases  []*mat.Dense
	Weights []*mat.Dense
	A       []*mat.Dense
	Z       []*mat.Dense
	alpha   float64
}

func (nn *NeuralNetwork) Result() *mat.Dense {
	return nn.A[nn.Count-1]
}

func NewNeuralNetwork(Sizes []int) (nn *NeuralNetwork) {
	nn = &NeuralNetwork{}
	nn.Sizes = Sizes
	nn.Count = len(Sizes)
	nn.Weights = make([]*mat.Dense, nn.Count)
	nn.Biases = make([]*mat.Dense, nn.Count)
	nn.A = make([]*mat.Dense, nn.Count)
	nn.Z = make([]*mat.Dense, nn.Count)
	nn.alpha = 0.2 / float64(nn.Sizes[0])

	for i := 1; i < nn.Count; i++ {
		nn.Weights[i] = generateRandomDense(nn.Sizes[i], nn.Sizes[i-1])
		nn.Biases[i] = generateRandomDense(nn.Sizes[i], 1)
	}
	return
}

func (nn *NeuralNetwork) Predict(aIn mat.Matrix) (maxIndex int, max float64) {
	nn.Forward(aIn)
	result := nn.Result()
	r, _ := result.Dims()
	max = 0.0
	maxIndex = 0
	for i := 0; i < r; i++ {
		if result.At(i, 0) > max {
			max = result.At(i, 0)
			maxIndex = i
		}
	}
	return
}

func (nn *NeuralNetwork) Forward(aIn mat.Matrix) {
	nn.A[0] = mat.DenseCopyOf(aIn)

	for i := 1; i < nn.Count; i++ {
		nn.A[i] = mat.NewDense(nn.Sizes[i], 1, nil)
		aSrc := nn.A[i-1]
		aDst := nn.A[i]
		aDst.Mul(nn.Weights[i], aSrc)
		aDst.Add(aDst, nn.Biases[i])
		nn.Z[i] = mat.DenseCopyOf(aDst)
		aDst.Apply(applySigmoid, aDst)
	}
}

func (nn *NeuralNetwork) Backward(aIn, aOut mat.Matrix) {
	nn.Forward(aOut)

	//Initial error
	err := &mat.Dense{}
	err.Sub(aOut, nn.Result())

	for i := nn.Count - 1; i > 0; i-- {
		sigmoidsPrime := &mat.Dense{}
		sigmoidsPrime.Apply(applySigmoidPrime, nn.Z[i])

		delta := &mat.Dense{}
		delta.MulElem(err, sigmoidsPrime)
		err = delta

		biases := mat.DenseCopyOf(delta)
		weights := &mat.Dense{}
		weights.Mul(nn.A[i-1], delta)

		// Scale down
		nn.Weights[i] = makeBackGradien(weights, nn.Weights[i], nn.alpha)
		nn.Biases[i] = makeBackGradien(biases, nn.Biases[i], nn.alpha)
	}
}

func makeBackGradien(in mat.Matrix, actual mat.Matrix, alpha float64) *mat.Dense {
	scaled := &mat.Dense{}
	result := &mat.Dense{}

	scaled.Scale(alpha, in)
	result.Sub(actual, scaled)
	return result
}