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
}

func (nn *NeuralNetwork) Result() mat.Matrix {
	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)
	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) 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)
	}
}