package nn

import (
	"gonum.org/v1/gonum/mat"
	"math/rand"
	"time"
)

type Network struct {
	layers []uint // number of neurons of each layer

	fed bool // whether the network has the state fed by a forward pass

	neurons []*mat.Dense // neuron value vector for each layer (size=L)
	biases  []*mat.Dense // neuron bias vector for each layer (size=L-1)
	weights []*mat.Dense // weights between each 2-layers (size=L-1)
}

const MaxLayerCount = 255
const LearningRate = 1

// Empty creates a new network where each argument is the number of neurons of a layer
// each param represents a layer, including input/output
func Empty(_layers ...uint) (*Network, error) {

	var L = len(_layers)
	// check size
	if L < 2 {
		return nil, ErrMissingLayers
	} else if L > MaxLayerCount {
		return nil, ErrTooMuchLayers
	}

	net := &Network{
		layers:  _layers,
		fed:     false,
		neurons: make([]*mat.Dense, 0),
		biases:  make([]*mat.Dense, 0),
		weights: make([]*mat.Dense, 0),
	}

	for i, layer := range _layers {

		// check neuron count
		if layer < 1 {
			return nil, ErrEmptyLayer
		}

		// create neurons
		net.neurons = append(net.neurons, mat.NewDense(int(layer), 1, nil))

		// do not create weights nor biases for first layer
		// (no previous layer to bound to)
		if i == 0 {
			continue
		}

		// create biases
		biases := make([]float64, 0, layer+1)
		for b := uint(0); b < layer; b++ {
			rand.Seed(time.Now().UnixNano())
			biases = append(biases, rand.Float64())
		}
		biasesVec := mat.NewDense(int(layer), 1, biases)
		net.biases = append(net.biases, biasesVec)

		rows, cols := int(layer), int(_layers[i-1])
		weights := make([]float64, 0, rows*cols+1)
		for v := 0; v < rows*cols; v++ {
			rand.Seed(time.Now().UnixNano())
			weights = append(weights, rand.Float64())
		}
		weightsMat := mat.NewDense(rows, cols, weights)
		net.weights = append(net.weights, weightsMat)

	}

	return net, nil

}

// reset neuron values
func (net *Network) reset() {
	net.fed = false

	for i, _ := range net.neurons {
		net.neurons[i] = mat.NewDense(int(net.layers[i]), 1, nil)
	}
}

// forward processes a forward propagation from an input vector
// and lets the network in the final processing state
func (net *Network) forward(_input ...float64) error {

	// check input size
	if len(_input) < net.neurons[0].ColView(0).Len() {
		return ErrMissingInput
	}
	// reset neuron values
	net.reset()

	// forward input to first layer
	for n, l := 0, net.neurons[0].ColView(0).Len(); n < l; n++ {
		net.neurons[0].Set(n, 0, _input[n])
	}

	// process each layer from the previous one
	for l, ll := 1, len(net.layers); l < ll; l++ {

		// Z = w^l . a^(l-1) + b^l
		z := net.neurons[l]

		a := net.neurons[l-1] // neurons of previous layer
		w := net.weights[l-1] // shifted by 1 because no weights between layers -1 and 0
		b := net.biases[l-1]  // same shift as weights

		z.Mul(w, a)
		z.Add(z, b)
		z.Apply(sigmoid, z)
	}

	net.fed = true

	return nil

}

// Cost returns the cost from the given output
func (net *Network) Cost(_expect ...float64) (float64, error) {

	costVec, err := net.costVec(_expect...)
	if err != nil {
		return 0, err
	}

	return mat.Sum(costVec), nil
}

// costVec returns the cost derivative for each output (as a vector)
// from the given _expect data
func (net *Network) costVec(_expect ...float64) (*mat.Dense, error) {

	if !net.fed {
		return nil, ErrNoState
	}

	out := net.neurons[len(net.neurons)-1]

	// check output size
	if len(_expect) < out.ColView(0).Len() {
		return nil, ErrMissingOutput
	}

	// build expect vector
	expect := mat.NewDense(len(_expect), 1, _expect)

	// process cost = 1/2 * learningRate * (out - expect)^2
	cost := new(mat.Dense)
	cost.Sub(out, expect)                                             // out - expect
	cost.MulElem(cost, cost)                                          // (out - expect)^2
	cost.Mul(cost, mat.NewDense(1, 1, []float64{0.5 * LearningRate})) // 1/2 *learningRate * (out - expect)^2

	return cost, nil
}

// errorVec returns the cost derivative  (also called ERROR) for each
// output (as a vector) from the given _expect data
func (net *Network) errorVec(_expect ...float64) (*mat.Dense, error) {

	if !net.fed {
		return nil, ErrNoState
	}

	outLayer := net.neurons[len(net.neurons)-1]

	// check output size
	if len(_expect) < outLayer.ColView(0).Len() {
		return nil, ErrMissingOutput
	}

	// build expect vector
	expect := mat.NewDense(len(_expect), 1, _expect)

	// calc cost derivative = 2 * learningRate * (expect - out)
	cost := new(mat.Dense)
	cost.Sub(expect, outLayer)
	cost.Mul(cost, mat.NewDense(1, 1, []float64{2 * LearningRate}))

	// return diff (derivative of cost)
	return cost, nil
}

// backward processes the backpropagation from the current network state
// and the expected data : _expect
func (net *Network) backward(_expect ...float64) error {

	out := net.neurons[len(net.neurons)-1]

	// 1. fail on no state (no forward pass applied first)
	if !net.fed {
		return ErrNoState
	}

	// fail on invalid _expect size
	if len(_expect) != out.ColView(0).Len() {
		return ErrMissingOutput
	}

	// calc ERROR = 0.5 * learningRate * (expect - out)
	// *it is in fact the cost derivative
	errors, err := net.errorVec(_expect...)
	if err != nil {
		return err
	}

	// FOR EACH LAYER (from last to 1)
	for l := len(net.layers) - 1; l > 0; l-- {

		neurons := net.neurons[l]
		previous := net.neurons[l-1]
		weights := net.weights[l-1] // from l-1 to l
		biases := net.biases[l-1]   // at l

		// calc GRADIENTS = sigmoid'( neuron[l-1] )
		gradients := new(mat.Dense)
		gradients.Apply(derivateSigmoid, neurons)
		gradients.MulElem(gradients, errors)
		gradients.Mul(gradients, mat.NewDense(1, 1, []float64{LearningRate}))

		// calc WEIGHTS DELTAS = gradients . previous^T
		wdeltas := new(mat.Dense)
		wdeltas.Mul(gradients, previous.T())

		// update weights
		weights.Add(weights, wdeltas)

		// adjust biases
		biases.Add(biases, gradients)

		// update ERRORS
		previousErrors := new(mat.Dense)
		previousErrors.Clone(errors)
		errors.Reset()
		errors.Mul(weights.T(), previousErrors)
	}

	return nil

}

// Guess uses the trained network to guess output from an input
func (net *Network) Guess(_input ...float64) ([]float64, error) {

	// process feed forward
	err := net.forward(_input...)
	if err != nil {
		return nil, err
	}

	// extract output
	return net.Output()

}

// Train uses the trained network to train with the _input and tries to learn
// to guess the _expect instead
func (net *Network) Train(_input []float64, _expect []float64) error {

	out := net.neurons[len(net.neurons)-1]

	// check output size
	if len(_expect) != out.ColView(0).Len() {
		return ErrMissingOutput
	}

	// process guess subroutine
	_, err := net.Guess(_input...)
	if err != nil {
		return err
	}

	// process backward propagation
	return net.backward(_expect...)

}

// Output returns the output data (only if the network has been fed)
func (net Network) Output() ([]float64, error) {

	if !net.fed {
		return nil, ErrNoState
	}

	out := net.neurons[len(net.neurons)-1]
	output := make([]float64, 0, net.layers[len(net.layers)-1])
	for n, l := 0, out.ColView(0).Len(); n < l; n++ {
		output = append(output, out.At(n, 0))
	}

	return output, nil

}