package neuralnetwork

import (
	"testing"

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

func TestNewNeuralNetwork(t *testing.T) {
	nn, err := NewNeuralNetwork([]int{}, nil)
	if nn != nil || err == nil {
		t.Error("nn initialized, but shouldn't ", err)
	}

	nn, err = NewNeuralNetwork([]int{0, 0, 0, 0}, nil)
	if nn != nil || err == nil {
		t.Error("nn initialized, but shouldn't ", err)
	}

	nn, err = NewNeuralNetwork([]int{1, 1, 1, 1}, nil)
	if nn != nil || err == nil {
		t.Error("nn initialized, but shouldn't ", err)
	}

	nn, err = NewNeuralNetwork([]int{5, 5}, nil)
	if nn != nil || err == nil {
		t.Error("nn initialized, but shouldn't ", err)
	}

	nn, err = NewNeuralNetwork([]int{5, 1, 5, 5}, nil)
	if nn != nil || err == nil {
		t.Error("nn initialized, but shouldn't ", err)
	}

	nn, err = NewNeuralNetwork([]int{5, 4, 4, 5}, nil)
	if nn == nil || err != nil {
		t.Error("nn is not initialized, but should be ", err)
	}
}

func TestNeuralNetworkPredict(t *testing.T) {
	nn, _ := NewNeuralNetwork([]int{3, 4, 4, 2}, nil)

	aIn := &mat.Dense{}
	index, max := nn.Predict(aIn)
	if index != -1 || max != 0.0 {
		t.Error("Prediction when empty aIn shouldn't be possibe but predicted", index, max)
	}

	aIn = mat.NewDense(2, 1, []float64{0.1, 0.2})
	index, max = nn.Predict(aIn)
	if index != -1 || max != 0.0 {
		t.Error("Prediction aIn has invalid size shouldn't be possibe but predicted", index, max)
	}

	aIn = mat.NewDense(3, 1, []float64{0.1, 0.2, 0.3})
	index, max = nn.Predict(aIn)
	if index == -1 || max == 0.0 {
		t.Error("Prediction of aIn valid size should be predicted", index, max)
	}

	aIn = mat.NewDense(4, 1, []float64{0.1, 0.2, 0.3, 0.4})
	index, max = nn.Predict(aIn)
	if index != -1 || max != 0.0 {
		t.Error("Prediction aIn has invalid size shouldn't be possibe but predicted", index, max)
	}
}